tolgacangoz
/

matryoshka-diffusion-models

Model card Files Files and versions Community

tolgacangoz commited on 30 days ago

Commit

a108e65

•

1 Parent(s): c08a7a5

Upload matryoshka.py

Files changed (1) hide show

scheduler/matryoshka.py +2 -2

scheduler/matryoshka.py CHANGED Viewed

@@ -4002,7 +4002,7 @@ class MatryoshkaPipeline(
                 prompt_attention_mask = torch.cat(
                     [
                         prompt_attention_mask,
-                        torch.zeros(batch_size, max_len - len(prompt_attention_mask[0]), dtype=torch.long),
                     ],
                     dim=1,
                 )
@@ -4014,7 +4014,7 @@ class MatryoshkaPipeline(
                 negative_prompt_attention_mask = torch.cat(
                     [
                         negative_prompt_attention_mask,
-                        torch.zeros(batch_size, max_len - len(negative_prompt_attention_mask[0]), dtype=torch.long),
                     ],
                     dim=1,
                 )

                 prompt_attention_mask = torch.cat(
                     [
                         prompt_attention_mask,
+                        torch.zeros(batch_size, max_len - len(prompt_attention_mask[0]), dtype=torch.long, device=device),
                     ],
                     dim=1,
                 )
                 negative_prompt_attention_mask = torch.cat(
                     [
                         negative_prompt_attention_mask,
+                        torch.zeros(batch_size, max_len - len(negative_prompt_attention_mask[0]), dtype=torch.long, device=device),
                     ],
                     dim=1,
                 )