Update diffusion/models/stable_diffusion.py

Co-authored-by: Landan Seguin <[email protected]>
mosaicml · Oct 2, 2023 · 45c0cd7 · 45c0cd7
1 parent dd35c77
commit 45c0cd7
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/diffusion/models/stable_diffusion.py b/diffusion/models/stable_diffusion.py
@@ -447,12 +447,8 @@ def generate(
                 crop_params = [0., 0.]
             if not size_params:
                 size_params = [width, height]
-            cond_original_size = torch.tensor([[width, height]]).repeat(pooled_embeddings.shape[0],
-                                                                        1).to(device).float()
-            cond_crops_coords_top_left = torch.tensor([crop_params]).repeat(pooled_embeddings.shape[0],
-                                                                            1).to(device).float()
-            cond_target_size = torch.tensor([size_params]).repeat(pooled_embeddings.shape[0], 1).to(device).float()
-            add_time_ids = torch.cat([cond_original_size, cond_crops_coords_top_left, cond_target_size], dim=1).float()
+            add_time_ids = torch.tensor([[width, height, *crop_params, *size_params]], dtype=torch.float, device=device)
+            add_time_ids = add_time_ids.repeat(pooled_embeddings.shape[0], 1)
             add_text_embeds = pooled_embeddings
 
             added_cond_kwargs = {'text_embeds': add_text_embeds, 'time_ids': add_time_ids}