Update app.py

app.py

@@ -9,6 +9,7 @@ from huggingface_hub import hf_hub_download
 import torch
 from diffusers import DiffusionPipeline
 from huggingface_hub import hf_hub_download
+# from gradio_imageslider import ImageSlider
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -26,7 +27,7 @@ pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",
 # pipe.fuse_lora(lora_scale=0.125)
 
 #pipe.enable_lora()
-pipe.to("cuda")
+pipe.to(DEVICE)
 
 def get_examples():
     case = [
@@ -38,9 +39,15 @@ def get_examples():
            [Image.open("metal.png"), "dragon.png","a dragon, in 3d melting gold metal",0.9, 0.5, 0, 4, 8, 8,  789385745, False,True, 2,  True , "text/image guided stylzation"],
         ]
     return case
-    
-def reset_do_inversion():
+
+def reset_image_input():
+    return True
+
+def reset_do_inversion(image_input):
+    if image_input:
         return True
+    else:
+        return False
     
 def resize_img(image, max_size=1024):
     width, height = image.size
@@ -100,26 +107,43 @@ def invert_and_edit(image,
                     height = 1024,
                     inverted_latent_list = None,
                     do_inversion = True,
+                    image_input = False,
                     
                    ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    if do_inversion:
-        inverted_latent_list = pipe(
-            source_prompt,
-            height=1024,
-            width=1024,
-            guidance_scale=1,
-            output_type="pil",
-            num_inference_steps=num_inversion_steps,
-            max_sequence_length=512,
-            latents=image2latent(image),
-            invert_image=True
-        )
-        do_inversion = False
+    if image_input:
+        if do_inversion:
+            inverted_latent_list = pipe(
+                source_prompt,
+                height=1024,
+                width=1024,
+                guidance_scale=1,
+                output_type="pil",
+                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]
+        
+        latents = inverted_latent_list[-1].tile(2, 1, 1)
+        guidance_scale = [1,3]
+        image_input = True
     else:
-        # move to gpu because of zero and gr.states
-        inverted_latent_list = [tensor.to(DEVICE) for tensor in inverted_latent_list]
+        latents = torch.randn(
+            (4096, 64), 
+            generator=torch.Generator(0).manual_seed(0), 
+            dtype=torch.float16,
+    device=DEVICE,
+        ).tile(2, 1, 1)
+        guidance_scale = 3.5
+        image_input = False
+        
 
     try:
         multimodal_layers = convert_string_to_list(multimodal_layers)
@@ -131,19 +155,20 @@ def invert_and_edit(image,
             [source_prompt, edit_prompt],
             height=1024,
             width=1024,
-            guidance_scale=[1,3],
+            guidance_scale=guidance_scale,
             output_type="pil",
             num_inference_steps=num_inference_steps,
             max_sequence_length=512,
-            latents=inverted_latent_list[-1].tile(2, 1, 1),
+            latents=latents,
             inverted_latent_list=inverted_latent_list,
             mm_copy_blocks=multimodal_layers,
             single_copy_blocks=single_layers,
-        ).images[1]
+        ).images
 
     # 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
+    if inverted_latent_list is not None:
+        inverted_latent_list = [tensor.cpu() for tensor in inverted_latent_list]
+    return output[0], output[1], inverted_latent_list, do_inversion, image_input, seed
 
 # UI CSS
 css = """
@@ -157,7 +182,8 @@ css = """
 with gr.Blocks(css=css) as demo:
 
     inverted_latents = gr.State()
-    do_inversion = gr.State(True)
+    do_inversion = gr.State(False)
+    image_input = gr.State(False)
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""# Stable Flow 🌊🖌️
@@ -205,6 +231,9 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
             
             with gr.Column():
                 result = gr.Image(label="Result")
+            # with gr.Column():
+            #     with gr.Group():
+            #         result = ImageSlider(position=0.5)
         
         with gr.Accordion("Advanced Settings", open=False):
             
@@ -271,10 +300,11 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
             width,
             height,
             inverted_latents,
-            do_inversion
+            do_inversion,
+            image_input
             
         ],
-        outputs=[result, inverted_latents, do_inversion, seed],
+        outputs=[input_image, result, inverted_latents, do_inversion, image_input, seed],
     )
 
     # gr.Examples(
@@ -284,18 +314,28 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                
     #         )
 
-    input_image.change(
+    input_image.input(fn=reset_image_input,
+                     outputs=[image_input]).then(
+        fn=reset_do_inversion,
+        inputs = [image_input],
+        outputs=[do_inversion]
+    )
+    
+    source_prompt.change(
         fn=reset_do_inversion,
+        inputs = [image_input],
         outputs=[do_inversion]
     )
 
     num_inversion_steps.change(
         fn=reset_do_inversion,
+        inputs = [image_input],
         outputs=[do_inversion]
     )
 
     seed.change(
         fn=reset_do_inversion,
+        inputs = [image_input],
         outputs=[do_inversion]
     )