Update app.py

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
 import numpy as np
-
 import spaces
 import torch
-import spaces
 import random
 
 from diffusers import FluxControlPipeline, FluxTransformer2DModel
@@ -15,8 +13,25 @@ MAX_IMAGE_SIZE = 2048
 pipe = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-Depth-dev", torch_dtype=torch.bfloat16).to("cuda")
 processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")
 
+def load_lora(lora_path):
+    if not lora_path.strip():
+        return "Please provide a valid LoRA path"
+    try:
+        pipe.load_lora_weights(lora_path)
+        return f"Successfully loaded LoRA weights from {lora_path}"
+    except Exception as e:
+        return f"Error loading LoRA weights: {str(e)}"
+
+def unload_lora():
+    try:
+        pipe.unload_lora_weights()
+        return "Successfully unloaded LoRA weights"
+    except Exception as e:
+        return f"Error unloading LoRA weights: {str(e)}"
+
 @spaces.GPU
-def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, height=1024, 
+          guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -32,12 +47,6 @@ def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, heig
         generator=torch.Generator().manual_seed(seed),
     ).images[0]
     return image, seed
-    
-examples = [
-    "a tiny astronaut hatching from an egg on the moon",
-    "a cat holding a sign that says hello world",
-    "an anime illustration of a wiener schnitzel",
-]
 
 css="""
 #col-container {
@@ -49,14 +58,24 @@ css="""
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 Depth [dev]
+        gr.Markdown(f"""# FLUX.1 Depth [dev] with LoRA Support
 12B param rectified flow transformer structural conditioning tuned, guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
 [[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
         """)
 
+        # LoRA controls
+        with gr.Row():
+            lora_path = gr.Textbox(
+                label="HuggingFace LoRA Path",
+                placeholder="e.g., Borcherding/FLUX.1-dev-LoRA-AutumnSpringTrees"
+            )
+            load_lora_btn = gr.Button("Load LoRA")
+            unload_lora_btn = gr.Button("Unload LoRA")
+        
+        lora_status = gr.Textbox(label="LoRA Status", interactive=False)
+
         control_image = gr.Image(label="Upload the image for control", type="pil")
         with gr.Row():
-            
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
@@ -64,13 +83,11 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            
             run_button = gr.Button("Run", scale=0)
         
         result = gr.Image(label="Result", show_label=False)
         
         with gr.Accordion("Advanced Settings", open=False):
-            
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -82,7 +99,6 @@ with gr.Blocks(css=css) as demo:
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
             
             with gr.Row():
-                
                 width = gr.Slider(
                     label="Width",
                     minimum=256,
@@ -100,7 +116,6 @@ with gr.Blocks(css=css) as demo:
                 )
             
             with gr.Row():
-
                 guidance_scale = gr.Slider(
                     label="Guidance Scale",
                     minimum=1,
@@ -117,11 +132,24 @@ with gr.Blocks(css=css) as demo:
                     value=28,
                 )
 
+    # Event handlers
+    load_lora_btn.click(
+        fn=load_lora,
+        inputs=[lora_path],
+        outputs=[lora_status]
+    )
+    
+    unload_lora_btn.click(
+        fn=unload_lora,
+        inputs=[],
+        outputs=[lora_status]
+    )
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
-        fn = infer,
-        inputs = [control_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-        outputs = [result, seed]
+        fn=infer,
+        inputs=[control_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, seed]
     )
 
 demo.launch()