import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Check if a GPU is available, otherwise fall back to CPU safely
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

print(f"--- Current running device is: {device.upper()} ---")

# Load the photorealistic epiCRealism XL pipeline
pipe = DiffusionPipeline.from_pretrained(
    "glides/epicrealismxl", 
    torch_dtype=torch_dtype
)
pipe.to(device)

def generate_image(prompt, negative_prompt, guidance_scale):
    # Safe guard: SDXL models take too long on free CPUs and will timeout.
    # This prevents your space from crashing while waiting for the GPU grant.
    if device == "cpu":
        return None
        
    image = pipe(
        prompt=prompt, 
        negative_prompt=negative_prompt, 
        guidance_scale=guidance_scale
    ).images[0]
    return image

# Building a clean, lightweight Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="A stunning photorealistic portrait of..."),
        gr.Textbox(label="Negative Prompt", value="blurry, bad quality, distorted, lowres, digital art, anime"),
        gr.Slider(minimum=1, maximum=15, value=7.5, step=0.5, label="Guidance Scale")
    ],
    outputs=gr.Image(type="filepath", label="Generated Output Image")
)

demo.launch()