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

# ---------------- CONFIG & LOAD ----------------
MODEL_REPO = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
MODEL_FILE = "Realistic_Vision_V6.0_NV_B1_fp16.safetensors"
hf_token = os.environ.get("HF_TOKEN")

model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, token=hf_token)
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    safety_checker=None  # Resolves the safety checker warning
).to(device)
pipe.enable_attention_slicing()

# ---------------- API FUNCTION ----------------
def generate(prompt, negative_prompt):
    # This receives the exact multi-line storybook prompt from your local engine
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=35, 
        guidance_scale=7.0
    ).images[0]
    return image

# ---------------- GRADIO INTERFACE ----------------
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"), 
        gr.Textbox(label="Negative Prompt")
    ],
    outputs=gr.Image(type="filepath")
)

# Standard launch for Hugging Face Spaces
# If you explicitly want a public share link, use: demo.launch(share=True)
# But for a Space, demo.launch() is all you need.
demo.launch()