import gradio as gr
import spaces
from diffusers import DiffusionPipeline
import torch


device = "cuda" if torch.cuda.is_available() else "cpu"


pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16
).to(device)
pipe.load_lora_weights("MegaTronX/MetartLoRA", weight_name="MetartLoRA.safetensors")

@spaces.GPU(duration=75)
def generate_image(prompt, num_inference_steps=25, guidance_scale=7.5, seed=None):
    """Generates an image using the FLUX.1-dev LoRA model."""
    generator = torch.Generator("cuda").manual_seed(seed) if seed else None

    image = pipe(
        prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=generator,
    ).images[0]
    return image

# Gradio Interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(lines=3, label="Prompt"),
        gr.Slider(minimum=10, maximum=100, value=25, label="Inference Steps"),
        gr.Slider(minimum=1, maximum=15, value=7.5, label="Guidance Scale"),
        gr.Number(label="Seed (Optional)"),
    ],
    outputs=gr.Image(label="Generated Image"),
    title="FLUX.1-dev LoRA Demo",
    description="A demo of your FLUX.1-dev LoRA model.",
)

iface.launch()