dont pad

app.py

@@ -129,22 +129,9 @@ def cfg_wrapper(
     emb = te(**prompt_token).last_hidden_state
     neg_emb = te(**neg_prompt_token).last_hidden_state
 
-    if emb.size(1) > neg_emb.size(1):
-        pad_setting = (0, 0, 0, emb.size(1) - neg_emb.size(1))
-        neg_emb = F.pad(neg_emb, pad_setting)
-    if neg_emb.size(1) > emb.size(1):
-        pad_setting = (0, 0, 0, neg_emb.size(1) - emb.size(1))
-        emb = F.pad(emb, pad_setting)
-    text_ctx_emb = torch.concat([emb, neg_emb])
-
     def cfg_fn(x, t, cfg=cfg_scale):
-        cond, uncond = unet(
-            x.repeat(2, 1, 1, 1),
-            t.expand(x.size(0) * 2),
-            text_ctx_emb,
-        ).chunk(2)
-        cond = cond.float()
-        uncond = uncond.float()
+        cond = unet(x, t.expand(x.size(0)), emb).float()
+        uncond = unet(x, t.expand(x.size(0)), neg_emb).float()
         return uncond + (cond - uncond) * cfg
 
     return cfg_fn