import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

pipe = None

def load_pipeline():
    global pipe
    if pipe is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        revision = "fp16" if device == "cuda" else "main"
        torch_dtype = torch.float16 if device == "cuda" else torch.float32
        pipe = StableDiffusionPipeline.from_pretrained(
            "runwayml/stable-diffusion-v1-5",
            revision=revision,
            torch_dtype=torch_dtype
        ).to(device)
    return pipe

def generate_image(prompt, num_inference_steps=30, guidance_scale=7.5):
    try:
        pipeline = load_pipeline()
        with torch.autocast("cuda" if torch.cuda.is_available() else "cpu"):
            result = pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)
        image = result.images[0]
        return image
    except Exception as e:
        return f"Error: {str(e)}"

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", lines=2),
        gr.Slider(1, 100, step=1, value=30, label="Inference Steps"),
        gr.Slider(0.0, 15.0, step=0.5, value=7.5, label="Guidance Scale")
    ],
    outputs=[
        gr.Image(type="pil", label="Generated Image")
    ],
    title="Stable Diffusion Image Generator",
    description="Generate images from text prompts using Stable Diffusion."
)

if __name__ == "__main__":
    iface.launch()