Update app.py

app.py

@@ -22,8 +22,8 @@ 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"), lora_scale=0.125)
+# pipe.fuse_lora(lora_scale=0.125)
 
 #pipe.enable_lora()
 pipe.to("cuda")
@@ -83,10 +83,15 @@ def check_hyper_flux_lora(enable_hyper_flux):
         pipe.unfuse_lora()
         return 28, 28
 
+def convert_string_to_list(s):
+    return [int(x) for x in s.split(',') if x]
+
 @spaces.GPU(duration=150)
 def invert_and_edit(image,
                     source_prompt,
-                    edit_prompt, 
+                    edit_prompt,
+                    multimodal_layers,
+                    single_layers,
                     num_inversion_steps,
                     num_inference_steps,
                     seed,
@@ -115,6 +120,13 @@ 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]
+
+    try:
+        multimodal_layers = convert_string_to_list(multimodal_layers)
+        single_layers = convert_string_to_list(single_layers)
+    except:
+        multimodal_layers = MULTIMODAL_VITAL_LAYERS
+        single_layers = SINGLE_MODAL_VITAL_LAYERS
     output = pipe(
             [source_prompt, edit_prompt],
             height=1024,
@@ -125,8 +137,8 @@ def invert_and_edit(image,
             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,
+            mm_copy_blocks=multimodal_layers,
+            single_copy_blocks=single_layers,
         ).images[1]
 
     # move back to cpu because of zero and gr.states
@@ -171,6 +183,19 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
                     max_lines=1,
                     placeholder="describe the edited output",
                 )
+                with gr.Row():
+                    multimodal_layers = gr.Text(
+                        info = "the attention layers used for injection",
+                        label="vital multimodal layers",
+                        max_lines=1,
+                        placeholder="0, 1, 17, 18",
+                    )
+                    single_layers = gr.Text(
+                        info = "the attention layers used for injection",
+                        label="vital single layers",
+                        max_lines=1,
+                        placeholder="9, 34, 35, 37, 6",
+                    )
                 with gr.Row():
                     enable_hyper_flux = gr.Checkbox(label="8-step LoRA", value=False, info="may reduce edit quality", visible=False)
                     
@@ -237,6 +262,8 @@ following the algorithm proposed in [*Stable Flow: Vital Layers for Training-Fre
             input_image,
             source_prompt,
             edit_prompt,
+            multimodal_layers,
+            single_layers,
             num_inversion_steps,
             num_inference_steps,
             seed,