import gradio as gr
from diffusers import StableDiffusionPipeline
from diffusers import FluxPipeline
import torch
import spaces

@spaces.GPU
def generate_image(prompt, num_inference_steps, guidance_scale):
    # Load the base model
    base_model_id = "black-forest-labs/FLUX.1-dev"
    #pipe = StableDiffusionPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16)

    #model_path = "MegaTronX/SuicideGirl-FLUX"
    model_path = "https://huggingface.co/MegaTronX/SuicideGirl-FLUX/blob/main/SuicideGirls.safetensors"
    
    '''pipe = StableDiffusionPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16,
    use_safetensors=True
    '''

    # Load the base model and apply LoRA weights using from_single_file
    pipe = FluxPipeline.from_single_file(
        "https://huggingface.co/MegaTronX/SuicideGirl-FLUX/blob/main/SuicideGirls.safetensors",
        torch_dtype=torch.float16,
        use_safetensors=True,
        load_safety_checker=False,
        variant="fp16",
        original_config_file=base_model_id
    )


        # Load the LoRA weights
    #pipe.unet.load_attn_procs("MegaTronX/SuicideGirl-FLUX")
    
    pipe = pipe.to("cuda")
    
    # Generate the image
    image = pipe(
        prompt,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale
    ).images[0]
    
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=1, maximum=100, value=50, label="Number of Inference Steps"),
        gr.Slider(minimum=1, maximum=20, value=7.5, label="Guidance Scale")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="Image Generation with Custom LoRA",
    description="Generate images using a custom LoRA model trained on Flux Dev."
)

iface.launch()