Update src/utils.py

src/utils.py

@@ -193,7 +193,6 @@ def warpped_feature(sample, step):
 def warpped_skip_feature(block_samples, step):
     down_block_res_samples = []
     # print(block_samples.shape, step)
-    print(step)
     for sample in block_samples:
         sample_expand = warpped_feature(sample, step)
         down_block_res_samples.append(sample_expand)
@@ -205,9 +204,14 @@ def warpped_text_emb(text_emb, step):
     step: timestep span
     """
     bs, token_len, dim = text_emb.shape
-    uncond_fea, cond_fea = text_emb.chunk(2)
-    uncond_fea = uncond_fea.repeat(step,1,1) # (step * bs//2) * 77 *768
-    cond_fea = cond_fea.repeat(step,1,1) # (step * bs//2) * 77 * 768
+    if text_emb.shape[0] >= 2:
+        uncond_fea, cond_fea = text_emb.chunk(2)
+        uncond_fea = uncond_fea.repeat(step,1,1) # (step * bs//2) * 77 *768
+        cond_fea = cond_fea.repeat(step,1,1) # (step * bs//2) * 77 * 768
+    else:
+        uncond_fea, cond_fea = text_emb, text_emb
+        uncond_fea = uncond_fea.repeat(step,1,1) # (step * bs//2) * 77 *768
+        cond_fea = cond_fea.repeat(step,1,1) # (step * bs//2) * 77 * 768
     return torch.cat([uncond_fea, cond_fea]) # (step*bs) * 77 *768
 
 def warpped_timestep(timesteps, bs):