Update app.py

app.py

@@ -104,29 +104,32 @@ def invert_and_edit(image,
             width=1024,
             guidance_scale=1,
             output_type="pil",
-            num_inference_steps=50,
+            num_inference_steps=num_inversion_steps,
             max_sequence_length=512,
             latents=image2latent(image),
             invert_image=True
         )
         do_inversion = False
-    
-    
+    else:
+        # move to gpu because of zero and gr.states
+        inverted_latent_list = [tensor.to(DEVICE) for tensor in inverted_latent_list]
     output = pipe(
             [source_prompt, edit_prompt],
             height=1024,
             width=1024,
             guidance_scale=[1,3],
             output_type="pil",
-            num_inference_steps=50,
+            num_inference_steps=num_inference_steps,
             max_sequence_length=512,
             latents=inverted_latent_list[-1].tile(2, 1, 1),
             inverted_latent_list=inverted_latent_list,
             mm_copy_blocks=MULTIMODAL_VITAL_LAYERS,
             single_copy_blocks=SINGLE_MODAL_VITAL_LAYERS,
         ).images[1]
-        
-    return output, inverted_latent_list.cpu(), do_inversion, seed
+
+    # move back to cpu because of zero and gr.states
+    inverted_latent_list = [tensor.cpu() for tensor in inverted_latent_list]
+    return output, inverted_latent_list, do_inversion, seed
 
 # UI CSS
 css = """
@@ -193,7 +196,7 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                             minimum=1,
                             maximum=50,
                             step=1,
-                            value=18,
+                            value=25,
                         )
                 
 
@@ -204,7 +207,7 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                         minimum=1,
                         maximum=50,
                         step=1,
-                        value=50,
+                        value=25,
                     )
             
             with gr.Row():