Update app.py

app.py

@@ -8,23 +8,24 @@ import huggingface_hub
 huggingface_hub.login(os.getenv('HF_TOKEN_FLUX2'))
 # Load default image from assets as an example
 default_image = Image.open("assets/1.png")
-
-@spaces.GPU()
-def generate_image(model_path, image, height, width, prompt, guidance_scale, num_steps, lora_name):
-    # Load the model
-    pipeline = FluxPipeline.from_pretrained(
-        model_path,
+pipeline = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev",
         torch_dtype=torch.bfloat16,
     ).to('cuda')
 
+
+@spaces.GPU()
+def generate_image(image, height, width, prompt, guidance_scale, num_steps, lora_name):
+    # Load the model
+    
     # Load and fuse base LoRA weights
-    pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name="pretrain.safetensors")
-    pipeline.fuse_lora()
-    pipeline.unload_lora_weights()
+    # pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name="pretrain.safetensors")
+    # pipeline.fuse_lora()
+    # pipeline.unload_lora_weights()
 
     # Load selected LoRA effect if not using the pretrained base model
-    if lora_name != 'pretrained':
-        pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name=f"{lora_name}.safetensors")
+    pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name=f"{lora_name}.safetensors")
+    pipeline.fuse_lora()
     
     # Prepare the input image
     condition_image = image.resize((height, width)).convert("RGB")
@@ -46,7 +47,6 @@ def generate_image(model_path, image, height, width, prompt, guidance_scale, num
 iface = gr.Interface(
     fn=generate_image,
     inputs=[
-        gr.Textbox(label="Model Path", value="black-forest-labs/FLUX.1-dev"),
         gr.Image(label="Input Image", type="pil", value=default_image),
         gr.Slider(label="Height", value=768, minimum=256, maximum=1024, step=64),
         gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=64),