You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
226 lines
7.9 KiB
Python
226 lines
7.9 KiB
Python
# from github.com/AUTOMATIC1111/stable-diffusion-webui
|
|
|
|
import torch
|
|
from torch.nn.functional import silu
|
|
|
|
import ldm.modules.attention
|
|
import ldm.modules.diffusionmodules.model
|
|
|
|
|
|
|
|
module_in_gpu = None
|
|
cpu = torch.device("cpu")
|
|
device = gpu = torch.device("cuda")
|
|
|
|
|
|
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 setup_for_low_vram(sd_model, use_medvram):
|
|
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(gpu)
|
|
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
|
|
def first_stage_model_encode_wrap(self, encoder, x):
|
|
send_me_to_gpu(self, None)
|
|
return encoder(x)
|
|
|
|
def first_stage_model_decode_wrap(self, decoder, z):
|
|
send_me_to_gpu(self, None)
|
|
return decoder(z)
|
|
|
|
# remove three big modules, cond, first_stage, and unet from the model and then
|
|
# send the model to GPU. Then put modules back. the modules will be in CPU.
|
|
stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model
|
|
sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = None, None, None
|
|
sd_model.to(device)
|
|
sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.model = stored
|
|
|
|
# register hooks for those the first two models
|
|
sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)
|
|
sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)
|
|
sd_model.first_stage_model.encode = lambda x, en=sd_model.first_stage_model.encode: first_stage_model_encode_wrap(sd_model.first_stage_model, en, x)
|
|
sd_model.first_stage_model.decode = lambda z, de=sd_model.first_stage_model.decode: first_stage_model_decode_wrap(sd_model.first_stage_model, de, z)
|
|
parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model
|
|
|
|
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(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
|
|
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)
|
|
|
|
ldm.modules.diffusionmodules.model.nonlinearity = silu
|
|
|
|
try:
|
|
import xformers
|
|
except ImportError:
|
|
ldm.modules.attention.CrossAttention.forward = split_cross_attention_forward
|
|
ldm.modules.diffusionmodules.model.AttnBlock.forward = cross_attention_attnblock_forward
|
|
|
|
import math
|
|
import torch
|
|
from torch import einsum
|
|
|
|
from ldm.util import default
|
|
from einops import rearrange
|
|
|
|
|
|
|
|
# taken from https://github.com/Doggettx/stable-diffusion
|
|
def split_cross_attention_forward(self, x, context=None, mask=None):
|
|
h = self.heads
|
|
|
|
q_in = self.to_q(x)
|
|
context = default(context, x)
|
|
k_in = self.to_k(context) * self.scale
|
|
v_in = self.to_v(context)
|
|
del context, x
|
|
|
|
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
|
|
del q_in, k_in, v_in
|
|
|
|
r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
|
|
|
|
stats = torch.cuda.memory_stats(q.device)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_cuda + mem_free_torch
|
|
|
|
gb = 1024 ** 3
|
|
tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * q.element_size()
|
|
modifier = 3 if q.element_size() == 2 else 2.5
|
|
mem_required = tensor_size * modifier
|
|
steps = 1
|
|
|
|
if mem_required > mem_free_total:
|
|
steps = 2 ** (math.ceil(math.log(mem_required / mem_free_total, 2)))
|
|
# print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
|
|
# f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
|
|
|
|
if steps > 64:
|
|
max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
|
|
raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
|
|
f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
|
|
|
|
slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
|
|
for i in range(0, q.shape[1], slice_size):
|
|
end = i + slice_size
|
|
s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k)
|
|
|
|
s2 = s1.softmax(dim=-1, dtype=q.dtype)
|
|
del s1
|
|
|
|
r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
|
|
del s2
|
|
|
|
del q, k, v
|
|
|
|
r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
|
|
del r1
|
|
|
|
return self.to_out(r2)
|
|
|
|
def cross_attention_attnblock_forward(self, x):
|
|
h_ = x
|
|
h_ = self.norm(h_)
|
|
q1 = self.q(h_)
|
|
k1 = self.k(h_)
|
|
v = self.v(h_)
|
|
|
|
# compute attention
|
|
b, c, h, w = q1.shape
|
|
|
|
q2 = q1.reshape(b, c, h*w)
|
|
del q1
|
|
|
|
q = q2.permute(0, 2, 1) # b,hw,c
|
|
del q2
|
|
|
|
k = k1.reshape(b, c, h*w) # b,c,hw
|
|
del k1
|
|
|
|
h_ = torch.zeros_like(k, device=q.device)
|
|
|
|
stats = torch.cuda.memory_stats(q.device)
|
|
mem_active = stats['active_bytes.all.current']
|
|
mem_reserved = stats['reserved_bytes.all.current']
|
|
mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
|
|
mem_free_torch = mem_reserved - mem_active
|
|
mem_free_total = mem_free_cuda + mem_free_torch
|
|
|
|
tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
|
|
mem_required = tensor_size * 2.5
|
|
steps = 1
|
|
|
|
if mem_required > mem_free_total:
|
|
steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
|
|
|
|
slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
|
|
for i in range(0, q.shape[1], slice_size):
|
|
end = i + slice_size
|
|
|
|
w1 = torch.bmm(q[:, i:end], k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
|
|
w2 = w1 * (int(c)**(-0.5))
|
|
del w1
|
|
w3 = torch.nn.functional.softmax(w2, dim=2, dtype=q.dtype)
|
|
del w2
|
|
|
|
# attend to values
|
|
v1 = v.reshape(b, c, h*w)
|
|
w4 = w3.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
|
|
del w3
|
|
|
|
h_[:, :, i:end] = torch.bmm(v1, w4) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
|
|
del v1, w4
|
|
|
|
h2 = h_.reshape(b, c, h, w)
|
|
del h_
|
|
|
|
h3 = self.proj_out(h2)
|
|
del h2
|
|
|
|
h3 += x
|
|
|
|
return h3
|