Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from diffusers import FluxDiffusionPipeline
+from huggingface_hub import hf_hub_download
 import spaces
 
 # Load the fine-tuned model
@@ -10,22 +11,37 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return model, tokenizer
 
-model, tokenizer = load_model()
+# Load the base Flux Dev model
+model_id = "black-forest-labs/FLUX.1-dev"
+pipeline = FluxDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipeline = pipeline.to("cuda")
+
+# Download and load the LoRA weights
+lora_model_path = hf_hub_download("MegaTronX/SuicideGirl-FLUX", "SuicideGirls.safetensors")
+pipeline.load_lora_weights(lora_model_path)
 
-# Function to generate an image
 @spaces.GPU
-def generate_image(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_length=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps):
+    image = pipeline(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps
+    ).images[0]
+    return image
 
-# Gradio interface
+# Create the Gradio interface
 iface = gr.Interface(
     fn=generate_image,
-    inputs="text",
-    outputs="text",
-    title="Image Generation with LoRA Fine-Tuned Model",
-    description="Enter a prompt to generate an image using the fine-tuned model."
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Textbox(label="Negative Prompt"),
+        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale"),
+        gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Number of Inference Steps")
+    ],
+    outputs=gr.Image(type="pil"),
+    title="Image Generation with Flux Dev LoRA",
+    description="Generate images using a Flux Dev model with a custom LoRA fine-tune."
 )
 
 iface.launch()