Update app.py

app.py

@@ -14,74 +14,9 @@ with open('loras.json', 'r') as f:
 # Initialize the base model
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-original_load_lora = copy.deepcopy(pipe.load_lora_into_transformer)
 pipe.to("cuda")
 
-def load_lora_into_transformer_patched(cls, state_dict, transformer, adapter_name=None, alpha=None, _pipeline=None):
-    from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
-
-    keys = list(state_dict.keys())
-
-    transformer_keys = [k for k in keys if k.startswith(cls.transformer_name)]
-    state_dict = {
-        k.replace(f"{cls.transformer_name}.", ""): v for k, v in state_dict.items() if k in transformer_keys
-    }
-
-    if len(state_dict.keys()) > 0:
-        # check with first key if is not in peft format
-        first_key = next(iter(state_dict.keys()))
-        if "lora_A" not in first_key:
-            state_dict = convert_unet_state_dict_to_peft(state_dict)
-
-        if adapter_name in getattr(transformer, "peft_config", {}):
-            raise ValueError(
-                f"Adapter name {adapter_name} already in use in the transformer - please select a new adapter name."
-            )
-
-        rank = {}
-        for key, val in state_dict.items():
-            if "lora_B" in key:
-                rank[key] = val.shape[1]
-
-        lora_config_kwargs = get_peft_kwargs(rank, network_alpha_dict=None, peft_state_dict=state_dict)
-        if "use_dora" in lora_config_kwargs:
-            if lora_config_kwargs["use_dora"] and is_peft_version("<", "0.9.0"):
-                raise ValueError(
-                    "You need `peft` 0.9.0 at least to use DoRA-enabled LoRAs. Please upgrade your installation of `peft`."
-                )
-            else:
-                lora_config_kwargs.pop("use_dora")
-
-        
-        lora_config_kwargs["lora_alpha"] = 42
-        lora_config = LoraConfig(**lora_config_kwargs)
-
-        # adapter_name
-        if adapter_name is None:
-            adapter_name = get_adapter_name(transformer)
-
-        # In case the pipeline has been already offloaded to CPU - temporarily remove the hooks
-        # otherwise loading LoRA weights will lead to an error
-        is_model_cpu_offload, is_sequential_cpu_offload = cls._optionally_disable_offloading(_pipeline)
-
-        inject_adapter_in_model(lora_config, transformer, adapter_name=adapter_name)
-        incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)
-
-        if incompatible_keys is not None:
-            # check only for unexpected keys
-            unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
-            if unexpected_keys:
-                logger.warning(
-                    f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
-                    f" {unexpected_keys}. "
-                )
-
-        # Offload back.
-        if is_model_cpu_offload:
-            _pipeline.enable_model_cpu_offload()
-        elif is_sequential_cpu_offload:
-            _pipeline.enable_sequential_cpu_offload()
-        # Unsafe code />
+MAX_SEED = 2**32-1
 
 def update_selection(evt: gr.SelectData):
     selected_lora = loras[evt.index]
@@ -95,7 +30,7 @@ def update_selection(evt: gr.SelectData):
     )
 
 @spaces.GPU(duration=90)
-def run_lora(prompt, cfg_scale, steps, selected_index, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
 
@@ -115,18 +50,19 @@ def run_lora(prompt, cfg_scale, steps, selected_index, seed, width, height, lora
         pipe.load_lora_into_transformer = original_load_lora
         
     # Set random seed for reproducibility
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
 
     # Generate image
     image = pipe(
         prompt=f"{prompt} {trigger_word}",
-        #negative_prompt=negative_prompt,
         num_inference_steps=steps,
         guidance_scale=cfg_scale,
         width=width,
         height=height,
         generator=generator,
-        #cross_attention_kwargs={"scale": lora_scale},
+        joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
 
     # Unload LoRA weights
@@ -159,10 +95,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
             result = gr.Image(label="Generated Image")
 
     with gr.Row():
-        #with gr.Column():
-            #prompt_title = gr.Markdown("### Click on a LoRA in the gallery to select it")
-            #negative_prompt = gr.Textbox(label="Negative Prompt", lines=2, value="low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry")
-
+        with gr.Accordion("Advanced Settings", open=False)
         with gr.Column():
             with gr.Row():
                 cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
@@ -173,7 +106,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
                 height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
             
             with gr.Row():
-                seed = gr.Slider(label="Seed", minimum=0, maximum=2**32-1, step=1, value=0, randomize=True)
+                randomize_seed = gr.Checkbox(True, label="Randomize seed")
+                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
                 lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.85)
 
     gallery.select(update_selection, outputs=[prompt, selected_info, selected_index])
@@ -181,7 +115,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
-        inputs=[prompt, cfg_scale, steps, selected_index, seed, width, height, lora_scale],
+        inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
         outputs=[result]
     )