fix the cross-attention num_head bug. · city96/ComfyUI_ExtraModels@81f58d8 · GitHub
Skip to content

Commit 81f58d8

Browse files
lawrence-cjlawrence-cj
committed
fix the cross-attention num_head bug.
1 parent 9ec31c8 commit 81f58d8

6 files changed

Lines changed: 17 additions & 22 deletions

File tree

‎Gemma/nodes.py‎

Lines changed: 2 additions & 3 deletions

‎Sana/conf.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
"depth": 28,
1515
"hidden_size": 1152,
1616
"patch_size": 1,
17-
"num_heads": 36,
17+
"num_heads": 16,
1818
"linear_head_dim": 32,
1919
"model_max_length": 300,
2020
"y_norm": True,
@@ -36,7 +36,7 @@
3636
"depth": 20,
3737
"hidden_size": 2240,
3838
"patch_size": 1,
39-
"num_heads": 70,
39+
"num_heads": 20,
4040
"linear_head_dim": 32,
4141
"model_max_length": 300,
4242
"y_norm": True,

‎Sana/loader.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def extra_conds(self, **kwargs):
4848
return out
4949

5050

51-
def load_sana(model_path, model_conf, dtype):
51+
def load_sana(model_path, model_conf):
5252
state_dict = comfy.utils.load_torch_file(model_path)
5353
state_dict = state_dict.get("model", state_dict)
5454

@@ -62,7 +62,7 @@ def load_sana(model_path, model_conf, dtype):
6262
state_dict = convert_state_dict(state_dict) # Diffusers
6363

6464
parameters = comfy.utils.calculate_parameters(state_dict)
65-
unet_dtype = dtype
65+
unet_dtype = comfy.model_management.unet_dtype()
6666
load_device = comfy.model_management.get_torch_device()
6767
offload_device = comfy.model_management.unet_offload_device()
6868

‎Sana/models/sana.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ def __init__(
159159
caption_channels=2304,
160160
pe_interpolation=1.0,
161161
config=None,
162-
model_max_length=120,
162+
model_max_length=300,
163163
qk_norm=False,
164164
y_norm=False,
165165
norm_eps=1e-5,
@@ -182,6 +182,7 @@ def __init__(
182182
self.depth = depth
183183
self.use_pe = use_pe
184184
self.y_norm = y_norm
185+
self.model_max_length = model_max_length
185186
self.fp32_attention = kwargs.get("use_fp32_attention", False)
186187

187188
kernel_size = patch_embed_kernel or patch_size

‎Sana/models/sana_multi_scale.py‎

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -296,20 +296,19 @@ def forward_raw(self, x, timestep, y, mask=None, data_info=None, **kwargs):
296296

297297
t = self.t_embedder(timestep) # (N, D)
298298

299+
y_lens = ((y != 0).sum(dim=3) > 0).sum(dim=2).squeeze().tolist()
300+
y_lens = [y_lens[1]] * bs
301+
302+
mask = torch.zeros((len(y_lens), self.model_max_length), dtype=torch.int).to(x.device)
303+
for i, count in enumerate(y_lens):
304+
mask[i, :count] = 1
305+
299306
t0 = self.t_block(t)
300307
y = self.y_embedder(y, self.training, mask=mask) # (N, D)
301308
if self.y_norm:
302309
y = self.attention_y_norm(y)
303310

304-
if mask is not None:
305-
if mask.shape[0] != y.shape[0]:
306-
mask = mask.repeat(y.shape[0] // mask.shape[0], 1)
307-
mask = mask.squeeze(1).squeeze(1)
308-
y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1])
309-
y_lens = mask.sum(dim=1).tolist()
310-
else:
311-
y_lens = [y.shape[2]] * y.shape[0]
312-
y = y.squeeze(1).view(1, -1, x.shape[-1])
311+
y = y.squeeze(1).masked_select(mask.unsqueeze(-1).bool()).view(1, -1, y.shape[-1])
313312

314313
for block in self.blocks:
315314
x = auto_grad_checkpoint(

‎Sana/nodes.py‎

Lines changed: 1 addition & 5 deletions

0 commit comments

Comments
 (0)