mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2024-06-07 21:20:49 +00:00
291 lines
11 KiB
Python
291 lines
11 KiB
Python
from contextlib import contextmanager, nullcontext
|
|
import torch
|
|
from modules import devices, shared, patches
|
|
|
|
module_in_gpu = None
|
|
cpu = torch.device("cpu")
|
|
|
|
stream_impl = devices.get_stream_impl()
|
|
stream_wrapper = devices.get_stream_wrapper()
|
|
|
|
|
|
use_streamlined_lowvram = torch.cuda.is_available() and not shared.opts.use_non_streamlined_lowvram and stream_impl is not None and stream_wrapper is not None
|
|
|
|
|
|
def is_same_device(device1, device2):
|
|
tensor1_device_type = device1.type
|
|
tensor2_device_type = device2.type
|
|
tensor1_device_index = device1.index or 0
|
|
tensor2_device_index = device2.index or 0
|
|
return (
|
|
tensor1_device_type == tensor2_device_type
|
|
and tensor1_device_index == tensor2_device_index
|
|
)
|
|
|
|
class RTTensorMoverPatches:
|
|
def __init__(self):
|
|
self.mover_stream = stream_impl(device=devices.device)
|
|
self.calc_stream = stream_impl(device=devices.device)
|
|
self.stash = {}
|
|
|
|
self.linear_original = patches.patch(
|
|
__name__,
|
|
torch.nn.functional,
|
|
"linear",
|
|
self._wrapper_default(torch.nn.functional.linear),
|
|
)
|
|
self.conv2d_original = patches.patch(
|
|
__name__,
|
|
torch.nn.functional,
|
|
"conv2d",
|
|
self._wrapper_default(torch.nn.functional.conv2d),
|
|
)
|
|
self.conv3d_original = patches.patch(
|
|
__name__,
|
|
torch.nn.functional,
|
|
"conv3d",
|
|
self._wrapper_default(torch.nn.functional.conv3d),
|
|
)
|
|
self.group_norm_original = patches.patch(
|
|
__name__,
|
|
torch.nn.functional,
|
|
"group_norm",
|
|
self._wrapper_group_norm(torch.nn.functional.group_norm),
|
|
)
|
|
self.layer_norm_original = patches.patch(
|
|
__name__,
|
|
torch.nn.functional,
|
|
"layer_norm",
|
|
self._wrapper_layer_norm(torch.nn.functional.layer_norm),
|
|
)
|
|
|
|
@contextmanager
|
|
def wrap_weight_biases(self, input, weight, bias):
|
|
if not is_same_device(input.device, devices.device):
|
|
yield (weight, bias)
|
|
return
|
|
|
|
moved = False
|
|
before_calc_event, after_calc_event = None, None
|
|
with stream_wrapper(stream=self.mover_stream):
|
|
if weight is not None and not is_same_device(weight.device, input.device):
|
|
weight = weight.to(device=input.device, copy=True, non_blocking=weight.is_pinned())
|
|
moved = True
|
|
if bias is not None and not is_same_device(bias.device, input.device):
|
|
bias = bias.to(device=input.device, copy=True, non_blocking=bias.is_pinned())
|
|
moved = True
|
|
before_calc_event = self.mover_stream.record_event()
|
|
|
|
if not moved:
|
|
yield (weight, bias)
|
|
return
|
|
|
|
with stream_wrapper(stream=self.calc_stream):
|
|
if before_calc_event is not None:
|
|
self.calc_stream.wait_event(before_calc_event)
|
|
yield (weight, bias)
|
|
after_calc_event = self.calc_stream.record_event()
|
|
self.stash[id(after_calc_event)] = (weight, bias, after_calc_event)
|
|
|
|
to_remove = []
|
|
for k, (_, _, e) in self.stash.items():
|
|
if e.query():
|
|
to_remove.append(k)
|
|
|
|
for k in to_remove:
|
|
del self.stash[k]
|
|
|
|
def _wrapper_default(self, original):
|
|
def wrapper(input, weight, bias, *args, **kwargs):
|
|
with self.wrap_weight_biases(input, weight, bias) as (w, b):
|
|
return original(input, w, b, *args, **kwargs)
|
|
return wrapper
|
|
|
|
def _wrapper_group_norm(self, original):
|
|
def wrapper(input, num_groups, weight, bias, *args, **kwargs):
|
|
with self.wrap_weight_biases(input, weight, bias) as (w, b):
|
|
return original(input, num_groups, w, b, *args, **kwargs)
|
|
return wrapper
|
|
|
|
def _wrapper_layer_norm(self, original):
|
|
def wrapper(input, normalized_shape, weight, bias, *args, **kwargs):
|
|
with self.wrap_weight_biases(input, weight, bias) as (w, b):
|
|
return original(input, normalized_shape, w, b, *args, **kwargs)
|
|
return wrapper
|
|
|
|
def close(self):
|
|
patches.undo(__name__, torch.nn.functional, "linear")
|
|
patches.undo(__name__, torch.nn.functional, "conv2d")
|
|
patches.undo(__name__, torch.nn.functional, "conv3d")
|
|
patches.undo(__name__, torch.nn.functional, "group_norm")
|
|
patches.undo(__name__, torch.nn.functional, "layer_norm")
|
|
|
|
|
|
rtmover = None
|
|
if use_streamlined_lowvram:
|
|
rtmover = RTTensorMoverPatches()
|
|
|
|
|
|
def calc_wrapper():
|
|
if rtmover is not None:
|
|
return stream_wrapper(stream=rtmover.calc_stream)
|
|
return nullcontext()
|
|
|
|
|
|
def calc_sync():
|
|
if rtmover is not None:
|
|
return rtmover.calc_stream.synchronize()
|
|
return nullcontext()
|
|
|
|
|
|
def send_everything_to_cpu():
|
|
global module_in_gpu
|
|
|
|
if module_in_gpu is not None:
|
|
module_in_gpu.to(cpu)
|
|
|
|
module_in_gpu = None
|
|
|
|
|
|
def is_needed(sd_model):
|
|
return shared.cmd_opts.lowvram or shared.cmd_opts.medvram or shared.cmd_opts.medvram_sdxl and hasattr(sd_model, 'conditioner')
|
|
|
|
|
|
def apply(sd_model):
|
|
enable = is_needed(sd_model)
|
|
shared.parallel_processing_allowed = not enable
|
|
|
|
if enable:
|
|
setup_for_low_vram(sd_model, not shared.cmd_opts.lowvram)
|
|
else:
|
|
sd_model.lowvram = False
|
|
|
|
|
|
def setup_for_low_vram(sd_model, use_medvram):
|
|
if getattr(sd_model, 'lowvram', False):
|
|
return
|
|
|
|
sd_model.lowvram = True
|
|
|
|
parents = {}
|
|
|
|
def send_me_to_gpu(module, _):
|
|
"""send this module to GPU; send whatever tracked module was previous in GPU to CPU;
|
|
we add this as forward_pre_hook to a lot of modules and this way all but one of them will
|
|
be in CPU
|
|
"""
|
|
global module_in_gpu
|
|
|
|
module = parents.get(module, module)
|
|
|
|
if module_in_gpu == module:
|
|
return
|
|
|
|
if module_in_gpu is not None:
|
|
module_in_gpu.to(cpu)
|
|
|
|
module.to(devices.device)
|
|
module_in_gpu = module
|
|
|
|
# see below for register_forward_pre_hook;
|
|
# first_stage_model does not use forward(), it uses encode/decode, so register_forward_pre_hook is
|
|
# useless here, and we just replace those methods
|
|
|
|
first_stage_model = sd_model.first_stage_model
|
|
first_stage_model_encode = sd_model.first_stage_model.encode
|
|
first_stage_model_decode = sd_model.first_stage_model.decode
|
|
|
|
def first_stage_model_encode_wrap(x):
|
|
send_me_to_gpu(first_stage_model, None)
|
|
return first_stage_model_encode(x)
|
|
|
|
def first_stage_model_decode_wrap(z):
|
|
send_me_to_gpu(first_stage_model, None)
|
|
return first_stage_model_decode(z)
|
|
|
|
to_remain_in_cpu = [
|
|
(sd_model, 'first_stage_model'),
|
|
(sd_model, 'depth_model'),
|
|
(sd_model, 'embedder'),
|
|
(sd_model, 'model'),
|
|
(sd_model, 'embedder'),
|
|
]
|
|
|
|
is_sdxl = hasattr(sd_model, 'conditioner')
|
|
is_sd2 = not is_sdxl and hasattr(sd_model.cond_stage_model, 'model')
|
|
|
|
if is_sdxl:
|
|
to_remain_in_cpu.append((sd_model, 'conditioner'))
|
|
elif is_sd2:
|
|
to_remain_in_cpu.append((sd_model.cond_stage_model, 'model'))
|
|
else:
|
|
to_remain_in_cpu.append((sd_model.cond_stage_model, 'transformer'))
|
|
|
|
# remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model
|
|
stored = []
|
|
for obj, field in to_remain_in_cpu:
|
|
module = getattr(obj, field, None)
|
|
stored.append(module)
|
|
setattr(obj, field, None)
|
|
|
|
# send the model to GPU.
|
|
sd_model.to(devices.device)
|
|
|
|
# put modules back. the modules will be in CPU.
|
|
for (obj, field), module in zip(to_remain_in_cpu, stored):
|
|
setattr(obj, field, module)
|
|
|
|
# register hooks for those the first three models
|
|
if is_sdxl:
|
|
sd_model.conditioner.register_forward_pre_hook(send_me_to_gpu)
|
|
elif is_sd2:
|
|
sd_model.cond_stage_model.model.register_forward_pre_hook(send_me_to_gpu)
|
|
sd_model.cond_stage_model.model.token_embedding.register_forward_pre_hook(send_me_to_gpu)
|
|
parents[sd_model.cond_stage_model.model] = sd_model.cond_stage_model
|
|
parents[sd_model.cond_stage_model.model.token_embedding] = sd_model.cond_stage_model
|
|
else:
|
|
sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
|
|
parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
|
|
|
|
sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
|
|
sd_model.first_stage_model.encode = first_stage_model_encode_wrap
|
|
sd_model.first_stage_model.decode = first_stage_model_decode_wrap
|
|
if sd_model.depth_model:
|
|
sd_model.depth_model.register_forward_pre_hook(send_me_to_gpu)
|
|
if sd_model.embedder:
|
|
sd_model.embedder.register_forward_pre_hook(send_me_to_gpu)
|
|
|
|
if use_medvram:
|
|
sd_model.model.register_forward_pre_hook(send_me_to_gpu)
|
|
else:
|
|
diff_model = sd_model.model.diffusion_model
|
|
|
|
# the third remaining model is still too big for 4 GB, so we also do the same for its submodules
|
|
# so that only one of them is in GPU at a time
|
|
stored = diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed
|
|
diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = None, None, None, None
|
|
sd_model.model.to(devices.device)
|
|
diff_model.input_blocks, diff_model.middle_block, diff_model.output_blocks, diff_model.time_embed = stored
|
|
|
|
# install hooks for bits of third model
|
|
|
|
if use_streamlined_lowvram:
|
|
# put it into pinned memory to achieve data transfer overlap
|
|
diff_model.time_embed._apply(lambda x: x.pin_memory())
|
|
for block in diff_model.input_blocks:
|
|
block._apply(lambda x: x.pin_memory())
|
|
diff_model.middle_block._apply(lambda x: x.pin_memory())
|
|
for block in diff_model.output_blocks:
|
|
block._apply(lambda x: x.pin_memory())
|
|
else:
|
|
diff_model.time_embed.register_forward_pre_hook(send_me_to_gpu)
|
|
for block in diff_model.input_blocks:
|
|
block.register_forward_pre_hook(send_me_to_gpu)
|
|
diff_model.middle_block.register_forward_pre_hook(send_me_to_gpu)
|
|
for block in diff_model.output_blocks:
|
|
block.register_forward_pre_hook(send_me_to_gpu)
|
|
|
|
|
|
def is_enabled(sd_model):
|
|
return sd_model.lowvram
|