import gradio as gr
import torch
from diffusers import FluxDiffusionPipeline
from huggingface_hub import hf_hub_download
import spaces

# Load the fine-tuned model
def load_model():
    model_name = "MegaTronX/SuicideGirl-FLUX"  # Replace with your model path
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return model, tokenizer

# 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)

@spaces.GPU
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

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    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()