"""
CCTV frame restoration demo — Stable Diffusion x4 Upscaler.
Self-contained: this is the only Python file the Space needs.
"""

import torch
import spaces
from diffusers import StableDiffusionUpscalePipeline
import gradio as gr

MODEL_ID = "stabilityai/stable-diffusion-x4-upscaler"
MAX_INPUT_SIDE = 128  # model was trained on small inputs; larger is slow and doesn't help

device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading model, this happens once at startup...")
pipe = StableDiffusionUpscalePipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
pipe.to(device)


print(f"Is CUDA available: {torch.cuda.is_available()}")
print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")

@spaces.GPU
def restore(image, prompt, steps):
    """Callback for the Gradio UI. `image` arrives as a PIL Image."""
    if image is None:
        return None

    image = image.convert("RGB")
    w, h = image.size
    scale = MAX_INPUT_SIDE / max(w, h)
    if scale < 1.0:
        image = image.resize((int(w * scale), int(h * scale)))

    result = pipe(
        prompt=prompt,
        negative_prompt="blurry, noisy, low quality, pixelated, artifacts",
        image=image,
        num_inference_steps=int(steps),
        guidance_scale=7.0,
    ).images[0]
    return result


demo = gr.Interface(
    fn=restore,
    inputs=[
        gr.Image(type="pil", label="Low-res / noisy CCTV frame"),
        gr.Textbox(
            value="a clear, sharp, well-lit security camera photograph, high detail",
            label="Restoration prompt",
        ),
        gr.Slider(10, 50, value=20, step=5, label="Diffusion steps (higher = slower, sharper)"),
    ],
    outputs=gr.Image(type="pil", label="Restored (4x upscaled)"),
    title="CCTV Frame Restoration with Stable Diffusion x4 Upscaler",
    description=(
        "Diffusion-based super-resolution for low-quality surveillance frames. "
        "The model was trained on 128x128 crops, so larger inputs are downscaled first. "
    ),
)

if __name__ == "__main__":
    demo.launch()