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

# Load Stable Diffusion v1.5
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", 
    safety_checker=None, 
    torch_dtype=torch.float32
).to("cpu")

# Output folder
os.makedirs("outputs", exist_ok=True)

# Aspect ratio mapping
aspect_ratios = {
    "Square (512x512)": (512, 512),
    "Portrait (512x768)": (512, 768),
    "Landscape (768x512)": (768, 512)
}

def generate_image(user_prompt, ratio):
    width, height = aspect_ratios[ratio]

    # Enhance prompt internally for better quality
    enhanced_prompt = f"{user_prompt}, highly detailed, ultra realistic, cinematic lighting, sharp focus"

    image = pipe(prompt=enhanced_prompt, height=height, width=width, guidance_scale=7.5).images[0]

    # Save image
    filename = f"outputs/{uuid.uuid4().hex}.png"
    image.save(filename)

    return image

# UI
with gr.Blocks() as demo:
    gr.Markdown("# 🌌 Stable Diffusion v1.5 - Image Generator")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="A shining star in the sky")
        ratio = gr.Radio(list(aspect_ratios.keys()), label="Aspect Ratio", value="Square (512x512)")
    with gr.Row():
        generate_btn = gr.Button("Generate Image")
    with gr.Column():
        output_img = gr.Image(label="Generated Image", interactive=False)

    generate_btn.click(fn=generate_image, inputs=[prompt, ratio], outputs=output_img)

demo.launch()