File size: 5,105 Bytes
93d565a 2c867cb 93d565a 48fac10 93d565a 48fac10 93d565a 2c867cb 93d565a 2c867cb 93d565a 2c867cb 5d62227 93d565a 5d62227 93d565a 5d62227 2c867cb 5d62227 2c867cb 5d62227 2c867cb 5d62227 2c867cb 5d62227 93d565a 2c867cb 93d565a | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """Gradio ZeroGPU demo for VOSR image super-resolution."""
from __future__ import annotations
import spaces
import gradio as gr
from pipeline import DEFAULT_MODEL_LABEL, MODEL_CHOICES, load_pipeline, run_sr
# Eager-load default weights onto CUDA so ZeroGPU can pack them at startup.
print("Preparing default VOSR pipeline (may download weights on first boot)...")
try:
load_pipeline(DEFAULT_MODEL_LABEL)
print("Default pipeline ready.")
except Exception as exc: # noqa: BLE001
print(f"Startup load deferred: {exc}")
def _gpu_duration(image, model_label, upscale, infer_steps, *args):
steps = int(infer_steps or 1)
mode = MODEL_CHOICES.get(model_label, {}).get("mode", "multistep")
pixels = 512 * 512
if image is not None:
w, h = image.size
u = max(int(upscale or 1), 1)
pixels = w * h * u * u
mp = max(pixels / 1e6, 0.25)
# Absolute ceiling 1800s; platform may still reject via illegal-duration / quota.
if mode == "onestep":
return int(min(1800, max(60, 45 + 25 * mp * steps)))
return int(min(1800, max(90, 60 + 8 * mp * steps)))
@spaces.GPU(duration=_gpu_duration)
def upscale_image(
image,
model_label,
upscale,
infer_steps,
cfg_scale,
weak_cond,
tile_size,
tile_overlap,
vae_tile_size,
vae_tile_overlap,
align_method,
seed,
):
return run_sr(
image=image,
model_label=model_label,
upscale=upscale,
infer_steps=infer_steps,
cfg_scale=cfg_scale,
weak_cond=weak_cond,
tile_size=tile_size,
tile_overlap=tile_overlap,
align_method=align_method,
seed=seed,
vae_tile_size=vae_tile_size,
vae_tile_overlap=vae_tile_overlap,
)
def _on_model_change(model_label):
spec = MODEL_CHOICES[model_label]
is_ms = spec["mode"] == "multistep"
return (
gr.update(value=spec["default_steps"]),
gr.update(interactive=is_ms),
gr.update(interactive=is_ms),
)
TITLE = "VOSR — Vision-Only Generative Super-Resolution"
DESCRIPTION = """
Demo of [VOSR](https://github.com/cswry/VOSR) (CVPR 2026) on **ZeroGPU**.
Upload a low-resolution image, keep the defaults (1.4B multi-step, 4×), and click **Upscale**.
For large images: set **DiT tile size** and optionally **VAE tile size** (both refer to
the upscaled result resolution; up to 8192). Prefer larger VAE tiles when VRAM allows —
tiling can still leave mild seams. Weights: [CSWRY/VOSR](https://huggingface.co/CSWRY/VOSR).
"""
with gr.Blocks(title="VOSR") as demo:
gr.Markdown(f"# {TITLE}\n{DESCRIPTION}")
with gr.Row():
with gr.Column():
inp = gr.Image(type="pil", label="Input image")
model = gr.Dropdown(
choices=list(MODEL_CHOICES.keys()),
value=DEFAULT_MODEL_LABEL,
label="Model",
)
upscale = gr.Slider(1, 8, value=4, step=1, label="Upscale factor")
btn = gr.Button("Upscale", variant="primary")
with gr.Accordion("Advanced", open=False):
infer_steps = gr.Slider(1, 50, value=25, step=1, label="Inference steps")
cfg_scale = gr.Slider(-1.0, 4.0, value=0.5, step=0.1, label="CFG scale (multi-step)")
weak_cond = gr.Slider(
0.05, 0.25, value=0.10, step=0.01, label="Weak cond strength (multi-step)"
)
tile_size = gr.Slider(
0,
8192,
value=0,
step=64,
label="DiT tile size (0 = off; pixels on result / upscaled image)",
)
tile_overlap = gr.Slider(0, 1024, value=32, step=8, label="DiT tile overlap")
vae_tile_size = gr.Slider(
0,
8192,
value=0,
step=64,
label="VAE tile size (0 = off; pixels on result / upscaled image)",
)
vae_tile_overlap = gr.Slider(
0, 1024, value=128, step=8, label="VAE tile overlap (pixels; ≥ tile/8 recommended)"
)
align_method = gr.Radio(
choices=["adain", "wavelet", "nofix"],
value="adain",
label="Color alignment",
)
seed = gr.Number(value=42, precision=0, label="Seed")
with gr.Column():
out = gr.Image(type="pil", label="Upscaled output")
model.change(_on_model_change, inputs=[model], outputs=[infer_steps, cfg_scale, weak_cond])
btn.click(
fn=upscale_image,
inputs=[
inp,
model,
upscale,
infer_steps,
cfg_scale,
weak_cond,
tile_size,
tile_overlap,
vae_tile_size,
vae_tile_overlap,
align_method,
seed,
],
outputs=[out],
)
if __name__ == "__main__":
demo.queue(max_size=4).launch()
|