r/stable-diffusion-finetune
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

add further nerf-attention

Browse source
This commit is contained in:
rromb 2022-07-06 23:44:26 +02:00
parent 0540e685e5
commit e2a6bee13c
2 changed files with 41 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
configs/stable-diffusion/txt2img-2B-clip-encoder-high-res-512-dev.yaml
View file

@ -32,11 +32,12 @@ model:
image_size: 64 # unused image_size: 64 # unused
in_channels: 4 in_channels: 4
out_channels: 4 out_channels: 4
model_channels: 352 model_channels: 384
attention_resolutions: [ 8, 4, 2 ] attention_resolutions: [ 8, 4, 2, 1 ]
num_res_blocks: [ 2, 2, 2, 6 ] num_res_blocks: [ 2, 2, 2, 5 ]
channel_mult: [ 1, 2, 4, 4 ] channel_mult: [ 1, 2, 4, 4 ]
disable_self_attentions: [ True, True, True, False ] # converts the self-attention to a cross-attention layer if true disable_self_attentions: [ False, False, False, False ] # converts the self-attention to a cross-attention layer if true
num_attention_blocks: [1, 1, 1, 3]
num_heads: 8 num_heads: 8
use_spatial_transformer: True use_spatial_transformer: True
transformer_depth: 1 transformer_depth: 1

16
ldm/modules/diffusionmodules/openaimodel.py
View file

@ -467,7 +467,8 @@ class UNetModel(nn.Module):
context_dim=None, # custom transformer support context_dim=None, # custom transformer support
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
legacy=True, legacy=True,
disable_self_attentions=None disable_self_attentions=None,
num_attention_blocks=None
): ):
super().__init__() super().__init__()
if use_spatial_transformer: if use_spatial_transformer:
@ -503,6 +504,13 @@ class UNetModel(nn.Module):
if disable_self_attentions is not None: if disable_self_attentions is not None:
# should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
assert len(disable_self_attentions) == len(channel_mult) assert len(disable_self_attentions) == len(channel_mult)
if num_attention_blocks is not None:
assert len(num_attention_blocks) == len(self.num_res_blocks)
assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
print(f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
f"attention will still not be set.") # todo: convert to warning
self.attention_resolutions = attention_resolutions self.attention_resolutions = attention_resolutions
self.dropout = dropout self.dropout = dropout
@ -538,7 +546,7 @@ class UNetModel(nn.Module):
ch = model_channels ch = model_channels
ds = 1 ds = 1
for level, mult in enumerate(channel_mult): for level, mult in enumerate(channel_mult):
for _ in range(self.num_res_blocks[level]): for nr in range(self.num_res_blocks[level]):
layers = [ layers = [
ResBlock( ResBlock(
ch, ch,
@ -564,6 +572,8 @@ class UNetModel(nn.Module):
disabled_sa = disable_self_attentions[level] disabled_sa = disable_self_attentions[level]
else: else:
disabled_sa = False disabled_sa = False
if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
layers.append( layers.append(
AttentionBlock( AttentionBlock(
ch, ch,
@ -670,6 +680,8 @@ class UNetModel(nn.Module):
disabled_sa = disable_self_attentions[level] disabled_sa = disable_self_attentions[level]
else: else:
disabled_sa = False disabled_sa = False
if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
layers.append( layers.append(
AttentionBlock( AttentionBlock(
ch, ch,