# app.py
import torch
import gradio as gr
from diffusers import DiffusionPipeline

# --- MODEL LOADING ---
# Load the model outside the prediction function for faster inference.
# ⚠️ The safety_checker=None is the key to uncensored generation.
MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0" 

try:
    pipe = DiffusionPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True,
        # 🔓 DISABLING THE SAFETY FILTER IS ESSENTIAL FOR UNCENSORED OUTPUT
        safety_checker=None, 
        requires_safety_checker=False 
    ).to("cuda") # Use CUDA for fast inference on the Space GPU
    print("✅ Model Loaded: SDXL Base 1.0 (Filter DISABLED)")
except Exception as e:
    print(f"Error loading model: {e}")
    # Fallback/Error handling is crucial for deployment

# --- GENERATION FUNCTION (Called by the Gradio API) ---
def generate_image(prompt, negative_prompt, steps, cfg_scale):
    if pipe is None:
        return "Error: Model failed to load."

    params = dict(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=1024,
        height=1024,
        # Use a random seed for variety
        generator=torch.Generator("cuda").manual_seed(torch.randint(0, 100000, (1,)).item())
    )

    # Image generation
    image = pipe(**params).images[0]
    return image

# --- GRADIO INTERFACE (Creates the API Endpoint) ---

gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt (Uncensored)"),
        gr.Textbox(label="Negative Prompt"),
        gr.Slider(minimum=20, maximum=100, step=1, label="Steps", value=30),
        gr.Slider(minimum=5.0, maximum=15.0, step=0.5, label="CFG Scale", value=8.0),
    ],
    outputs="image", # Outputting the image directly
    title="SDXL Uncensored Image Generator (Hugging Face Spaces)",
    description="This Space is running a Stable Diffusion XL model with the safety checker disabled.",
    allow_flagging="never"
).launch()