Update app.py

app.py

@@ -9,23 +9,17 @@ import spaces
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
-
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
-
 MAX_SEED = 2**32-1
 
-# Hidden variables (you would set these based on your specific LoRA)
-LORA_PATH = "SebastianBodza/Flux_Aquarell_Watercolor_v2"
-TRIGGER_WORD = "AQUACOLTOK"
-
-# Load LoRA weights (do this once at startup)
-pipe.load_lora_weights(LORA_PATH)
-
-@spaces.GPU(duration=70)
-def generate_image(prompt, width, height):
+@spaces.GPU()
+def generate_image(prompt, width, height, lora_path, trigger_word):
+    # Load LoRA weights
+    pipe.load_lora_weights(lora_path)
+    
     # Combine prompt with trigger word
-    full_prompt = f"{TRIGGER_WORD} {prompt}"
+    full_prompt = f"{trigger_word} {prompt}"
     
     # Set up generation parameters
     seed = random.randint(0, MAX_SEED)
@@ -43,8 +37,8 @@ def generate_image(prompt, width, height):
     
     return image
 
-def run_lora(prompt, width, height):
-    return generate_image(prompt, width, height)
+def run_lora(prompt, width, height, lora_path, trigger_word):
+    return generate_image(prompt, width, height, lora_path, trigger_word)
 
 # Set up the Gradio interface
 with gr.Blocks() as app:
@@ -57,13 +51,17 @@ with gr.Blocks() as app:
         width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=512)
         height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=512)
     
+    with gr.Row():
+        lora_path = gr.Textbox(label="LoRA Path", value="SebastianBodza/Flux_Aquarell_Watercolor_v2")
+        trigger_word = gr.Textbox(label="Trigger Word", value="AQUACOLTOK")
+    
     generate_button = gr.Button("Generate Image")
     
     output_image = gr.Image(label="Generated Image")
     
     generate_button.click(
         fn=run_lora,
-        inputs=[prompt, width, height],
+        inputs=[prompt, width, height, lora_path, trigger_word],
         outputs=[output_image]
     )