File size: 3,551 Bytes
fe6327d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
# NAI compatible
import torch
class HypernetworkModule(torch.nn.Module):
def __init__(self, dim, multiplier=1.0):
super().__init__()
linear1 = torch.nn.Linear(dim, dim * 2)
linear2 = torch.nn.Linear(dim * 2, dim)
linear1.weight.data.normal_(mean=0.0, std=0.01)
linear1.bias.data.zero_()
linear2.weight.data.normal_(mean=0.0, std=0.01)
linear2.bias.data.zero_()
linears = [linear1, linear2]
self.linear = torch.nn.Sequential(*linears)
self.multiplier = multiplier
def forward(self, x):
return x + self.linear(x) * self.multiplier
class Hypernetwork(torch.nn.Module):
enable_sizes = [320, 640, 768, 1280]
# return self.modules[Hypernetwork.enable_sizes.index(size)]
def __init__(self, multiplier=1.0) -> None:
super().__init__()
self.modules = []
for size in Hypernetwork.enable_sizes:
self.modules.append((HypernetworkModule(size, multiplier), HypernetworkModule(size, multiplier)))
self.register_module(f"{size}_0", self.modules[-1][0])
self.register_module(f"{size}_1", self.modules[-1][1])
def apply_to_stable_diffusion(self, text_encoder, vae, unet):
blocks = unet.input_blocks + [unet.middle_block] + unet.output_blocks
for block in blocks:
for subblk in block:
if 'SpatialTransformer' in str(type(subblk)):
for tf_block in subblk.transformer_blocks:
for attn in [tf_block.attn1, tf_block.attn2]:
size = attn.context_dim
if size in Hypernetwork.enable_sizes:
attn.hypernetwork = self
else:
attn.hypernetwork = None
def apply_to_diffusers(self, text_encoder, vae, unet):
blocks = unet.down_blocks + [unet.mid_block] + unet.up_blocks
for block in blocks:
if hasattr(block, 'attentions'):
for subblk in block.attentions:
if 'SpatialTransformer' in str(type(subblk)) or 'Transformer2DModel' in str(type(subblk)): # 0.6.0 and 0.7~
for tf_block in subblk.transformer_blocks:
for attn in [tf_block.attn1, tf_block.attn2]:
size = attn.to_k.in_features
if size in Hypernetwork.enable_sizes:
attn.hypernetwork = self
else:
attn.hypernetwork = None
return True # TODO error checking
def forward(self, x, context):
size = context.shape[-1]
assert size in Hypernetwork.enable_sizes
module = self.modules[Hypernetwork.enable_sizes.index(size)]
return module[0].forward(context), module[1].forward(context)
def load_from_state_dict(self, state_dict):
# old ver to new ver
changes = {
'linear1.bias': 'linear.0.bias',
'linear1.weight': 'linear.0.weight',
'linear2.bias': 'linear.1.bias',
'linear2.weight': 'linear.1.weight',
}
for key_from, key_to in changes.items():
if key_from in state_dict:
state_dict[key_to] = state_dict[key_from]
del state_dict[key_from]
for size, sd in state_dict.items():
if type(size) == int:
self.modules[Hypernetwork.enable_sizes.index(size)][0].load_state_dict(sd[0], strict=True)
self.modules[Hypernetwork.enable_sizes.index(size)][1].load_state_dict(sd[1], strict=True)
return True
def get_state_dict(self):
state_dict = {}
for i, size in enumerate(Hypernetwork.enable_sizes):
sd0 = self.modules[i][0].state_dict()
sd1 = self.modules[i][1].state_dict()
state_dict[size] = [sd0, sd1]
return state_dict
|