Spaces:
Sleeping
Sleeping
File size: 2,079 Bytes
a36fe68 0671f7e a36fe68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """
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()
|