# app.py — Verified Hugging Face Space Code (CPU/GPU Safe)

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

# ------------------ CONFIG ------------------
MODEL_ID = "runwayml/stable-diffusion-v1-5"

# ------------------ SAFE MODEL LOAD ------------------
try:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if torch.cuda.is_available() else torch.float32

    pipe = StableDiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=dtype,
        safety_checker=None
    ).to(device)

except Exception as e:
    print("❌ Model Loading Error:")
    traceback.print_exc()
    pipe = None

# ------------------ GENERATION FUNCTION ------------------
def generate_image(prompt):
    if pipe is None:
        return "⚠️ Model not loaded properly. Please check logs."
    try:
        result = pipe(prompt)
        image = result.images[0]
        return image
    except Exception as e:
        print("❌ Generation Error:", e)
        return "⚠️ Error generating image."

# ------------------ GRADIO UI ------------------
title = "🪷 Shrividya Text-to-Image AI"
description = "Generate stunning images from your text prompts using Stable Diffusion."

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt", placeholder="e.g. Divine temple on riverbank at sunset"),
    outputs=gr.Image(label="Generated Image"),
    title=title,
    description=description,
    allow_flagging="never"
)

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