import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image

# Load the model
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(prompt, negative_prompt="", num_inference_steps=50, guidance_scale=7.5, height=512, width=512):
    """
    Generate an image from text prompt using Stable Diffusion
    """
    try:
        with torch.no_grad():
            image = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                height=height,
                width=width
            ).images[0]
        return image
    except Exception as e:
        return f"Error generating image: {str(e)}"

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AI Image Generator")
    gr.Markdown("Generate images from text descriptions using Stable Diffusion")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Enter a detailed description of the image you want to generate",
                lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="(Optional) Things to avoid in the image",
                lines=2
            )
            
            with gr.Row():
                steps = gr.Slider(20, 100, value=50, step=1, label="Inference Steps")
                guidance = gr.Slider(1.0, 20.0, value=7.5, step=0.5, label="Guidance Scale")
            
            with gr.Row():
                height = gr.Slider(256, 768, value=512, step=64, label="Height")
                width = gr.Slider(256, 768, value=512, step=64, label="Width")
            
            generate_btn = gr.Button("Generate Image", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil")
    
    # Connect the generate button to the function
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance, height, width],
        outputs=output_image
    )

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