import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# ---- Configuration ---- #
MODEL_ID = "runwayml/stable-diffusion-v1-5"
USE_CUDA = torch.cuda.is_available()
DEVICE = "cuda" if USE_CUDA else "cpu"
DTYPE = torch.float16 if USE_CUDA else torch.float32

# ---- Load Pipeline ---- #
pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE
).to(DEVICE)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_attention_slicing()

# ---- Image Generation Function ---- #
def generate_image(prompt):
    enhanced_prompt = f"{prompt}, ultra realistic, high detail, 8k resolution, DSLR photography, natural lighting"
    with torch.inference_mode():
        result = pipe(enhanced_prompt, num_inference_steps=25)
    return result.images[0]

# ---- Gradio UI ---- #
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(
        label="Enter image description",
        placeholder="e.g. A cozy cabin in a snowy forest"
    ),
    outputs=gr.Image(type="pil"),
    title="🖼️ Realistic Text-to-Image Generator",
    description="Generate high-quality, photorealistic images using Stable Diffusion v1.5 + DPM Scheduler"
)

# ---- Launch ---- #
if __name__ == "__main__":
    demo.queue().launch()