Update app.py

app.py

@@ -9,7 +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
+from gradio_imageslider import ImageSlider
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -23,11 +23,11 @@ SINGLE_MODAL_VITAL_LAYERS = list(np.array([28, 53, 54, 56, 25]) - 19)
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev",
                         
                                           torch_dtype=torch.bfloat16)
-# pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), lora_scale=0.125)
-# pipe.fuse_lora(lora_scale=0.125)
+pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"))
+pipe.fuse_lora(lora_scale=0.125)
 
 #pipe.enable_lora()
-pipe.to(DEVICE)
+pipe.to(DEVICE, dtype=torch.float16)
 
 def get_examples():
     case = [
@@ -103,6 +103,7 @@ def invert_and_edit(image,
                     num_inference_steps,
                     seed,
                     randomize_seed,
+                    latent_nudging_scalar,
                     width = 1024,
                     height = 1024,
                     inverted_latent_list = None,
@@ -112,7 +113,7 @@ def invert_and_edit(image,
                    ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    if image_input:
+    if image_input and (image is not None):
         if do_inversion:
             inverted_latent_list = pipe(
                 source_prompt,
@@ -122,7 +123,7 @@ def invert_and_edit(image,
                 output_type="pil",
                 num_inference_steps=num_inversion_steps,
                 max_sequence_length=512,
-                latents=image2latent(image),
+                latents=image2latent(image, latent_nudging_scalar),
                 invert_image=True
             )
             do_inversion = False
@@ -130,7 +131,7 @@ def invert_and_edit(image,
         else:
             # move to gpu because of zero and gr.states
             inverted_latent_list = [tensor.to(DEVICE) for tensor in inverted_latent_list]
-        
+        num_inference_steps = num_inversion_steps 
         latents = inverted_latent_list[-1].tile(2, 1, 1)
         guidance_scale = [1,3]
         image_input = True
@@ -168,7 +169,10 @@ def invert_and_edit(image,
     # move back to cpu because of zero and gr.states
     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
+    if image is None:
+        image = output[0]
+    
+    return image, output[1], inverted_latent_list, do_inversion, image_input, seed
 
 # UI CSS
 css = """
@@ -252,7 +256,7 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                             minimum=1,
                             maximum=50,
                             step=1,
-                            value=25,
+                            value=8,
                         )
                 
 
@@ -265,6 +269,13 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                         step=1,
                         value=25,
                     )
+                latent_nudging_scalar= gr.Slider(
+                        label="latent nudging scalar",
+                        minimum=1,
+                        maximum=5,
+                        step=0.01,
+                        value=1.15,
+                    )
             
             with gr.Row():
                 width = gr.Slider(
@@ -297,6 +308,7 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
             num_inference_steps,
             seed,
             randomize_seed,
+            latent_nudging_scalar,
             width,
             height,
             inverted_latents,