Update app.py

app.py

@@ -150,37 +150,46 @@ def call(
 
         # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
         extra_step_kwargs = pipe.prepare_extra_step_kwargs(generator, eta)
-        
+
         # 7. Prepare added time ids & embeddings
         add_text_embeds = pooled_prompt_embeds
         add_text2_embeds = pooled_prompt2_embeds
-        
-        # Default dtype if prompt_embeds or prompt2_embeds are None
-        default_dtype = torch.float32
-        
-        # Check and set dtype for add_time_ids
-        dtype_for_add_time_ids = prompt_embeds.dtype if prompt_embeds is not None else default_dtype
+
         add_time_ids = pipe._get_add_time_ids(
-            original_size, crops_coords_top_left, target_size, dtype=dtype_for_add_time_ids
+            original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype
         )
-        
-        # Check and set dtype for add_time2_ids
-        dtype_for_add_time2_ids = prompt2_embeds.dtype if prompt2_embeds is not None else default_dtype
         add_time2_ids = pipe._get_add_time_ids(
-            original_size, crops_coords_top_left, target_size, dtype=dtype_for_add_time2_ids
+            original_size, crops_coords_top_left, target_size, dtype=prompt2_embeds.dtype
         )
-        
+
         if negative_original_size is not None and negative_target_size is not None:
             negative_add_time_ids = pipe._get_add_time_ids(
                 negative_original_size,
                 negative_crops_coords_top_left,
                 negative_target_size,
-                dtype=dtype_for_add_time_ids  # Use the same default dtype for negative prompts
+                dtype=prompt_embeds.dtype,
             )
         else:
             negative_add_time_ids = add_time_ids
             negative_add_time2_ids = add_time2_ids
 
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+            prompt2_embeds = torch.cat([negative_prompt2_embeds, prompt2_embeds], dim=0)
+            add_text2_embeds = torch.cat([negative_pooled_prompt2_embeds, add_text2_embeds], dim=0)
+            add_time2_ids = torch.cat([negative_add_time2_ids, add_time2_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+        prompt2_embeds = prompt2_embeds.to(device)
+        add_text2_embeds = add_text2_embeds.to(device)
+        add_time2_ids = add_time2_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
         # 8. Denoising loop
         num_warmup_steps = max(len(timesteps) - num_inference_steps * pipe.scheduler.order, 0)