import os
import torch
from diffusers import StableDiffusionPipeline
from huggingface_hub import hf_hub_download
import gradio as gr

# ---------------- CONFIG ----------------
MODEL_REPO = "Harshavarhini/realistic-vision-v6"
MODEL_FILE = "realisticVisionV60B1_v51HyperVAE.safetensors"

# ---------------- DOWNLOAD MODEL ----------------
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE  # 🔥 REQUIRED (fixes 401 error)
)

# ---------------- DEVICE SETUP ----------------
device = "cuda" if torch.cuda.is_available() else "cpu"

# ---------------- LOAD PIPELINE ----------------
pipe = StableDiffusionPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    low_cpu_mem_usage=True
)

pipe = pipe.to(device)

# 🔥 Memory optimization (VERY IMPORTANT for Spaces)
pipe.enable_attention_slicing()

# ---------------- GENERATION FUNCTION ----------------
def generate(prompt):
    image = pipe(
        prompt,
        num_inference_steps=30,
        guidance_scale=7.5
    ).images[0]
    
    return image

# ---------------- GRADIO UI ----------------
demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter Prompt"),
    outputs=gr.Image(type="pil"),
    title="Realistic Vision v6 API",
    description="Generate realistic images using Stable Diffusion"
)

# ---------------- LAUNCH ----------------
demo.launch()