import gradio as gr
import torch
import spaces
from diffusers import FluxPipeline
from safetensors.torch import load_file

# Load the Flux Dev model
model_id = "black-forest-labs/FLUX.1-dev"
pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cuda")

# Load the LoRA weights
lora_path = "MegaTronX/SuicideGirl-FLUX"
lora_weights = load_file(lora_path)

# Apply LoRA weights to the model
pipe.unet.load_attn_procs(lora_weights)

@spaces.GPU
def generate_image(prompt, negative_prompt, guidance_scale, num_inference_steps, lora_scale):
    with torch.inference_mode():
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            cross_attention_kwargs={"scale": lora_scale},
        ).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(1, 20, value=7.5, label="Guidance Scale"),
        gr.Slider(1, 100, value=50, step=1, label="Number of Inference Steps"),
        gr.Slider(0, 1, value=0.75, label="LoRA Scale"),
    ],
    outputs=gr.Image(type="pil"),
    title="Flux Dev with Custom LoRA Image Generator",
    description="Generate images using Flux Dev model with a custom LoRA trained on Civitai",
)

iface.launch()