Update app.py

app.py

@@ -150,46 +150,37 @@ 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=prompt_embeds.dtype
+            original_size, crops_coords_top_left, target_size, dtype=dtype_for_add_time_ids
         )
+        
+        # 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=prompt2_embeds.dtype
+            original_size, crops_coords_top_left, target_size, dtype=dtype_for_add_time2_ids
         )
-
+        
         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=prompt_embeds.dtype,
+                dtype=dtype_for_add_time_ids  # Use the same default dtype for negative prompts
             )
         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)