Upsampler/seedvr2-3b

1# SeedVR2-3B image restoration, Upsampler v3 recipe.2#3# Ported from ByteDance-Seed/SeedVR2-3B. What changed and why:4#5# 1. Image only. The upstream Space serves video and images from one entry6#    point, carrying the sequence-parallel plumbing, frame cutting and video7#    writing along with it. This Space backs an image tool, so that is all gone.8# 2. ONE @spaces.GPU entry. Upstream decorates configure_runner,9#    generation_step AND generation_loop, so a single request booked three GPU10#    allocations of 100s each. ZeroGPU checks the requested duration against the11#    visitor's remaining quota, and an unauthenticated visitor has 120 seconds a12#    day in total, so the upstream shape cannot serve an anonymous user at all.13#    Everything now runs inside one call with a measured dynamic duration.14# 3. No apex. Upstream installs a prebuilt `apex-0.1-cp310-...whl` and selects15#    `fusedrms` / `fusedln` norms in configs_3b/main.yaml. ZeroGPU runs Python16#    3.12, where that wheel does not install, so every norm layer then failed.17#    The config now selects the `rms` / `layer` paths that the same source file18#    already implements in pure PyTorch, with identical parameter names and19#    shapes so the checkpoint loads unchanged.20# 4. No hard flash-attn dependency. See models/dit_v2/attention.py.21#22# The model restores at a fixed working resolution regardless of input size (it23# was trained at high res and NaResize scales the input to meet it), so GPU cost24# per request is essentially constant and no tiling is involved. See WORK_AREA.25 26# Must precede the torch import: the allocator reads this at initialization.27#28# Fragmentation insurance, NOT the fix for the NVML assert this Space used to29# die on. That was measured: the ZeroGPU worker printed30# 'expandable_segments:True' on the runs that still crashed, so the setting was31# applied the whole time and made no difference. What actually decides it is32# WORK_AREA below. Kept because it costs nothing and reduces fragmentation.33import os34 35os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")36 37import gc38 39import mimetypes40from pathlib import Path41 42import gradio as gr43import spaces44import torch45import torch.nn.functional as F46from einops import rearrange47from omegaconf import OmegaConf48from PIL import Image49from huggingface_hub import hf_hub_download50from torchvision.transforms import Compose, Lambda, Normalize51import torchvision.transforms as T52 53from upsampler_theme import UPSAMPLER_CSS, UPSAMPLER_THEME, footer_html, header_html54 55from data.image.transforms.divisible_crop import DivisibleCrop56from data.image.transforms.na_resize import NaResize57from data.video.transforms.rearrange import Rearrange58from common.config import load_config59from common.distributed import init_torch60from common.seed import set_seed61from projects.video_diffusion_sr.infer import VideoDiffusionInfer62 63try:64    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction65 66    USE_COLOR_FIX = True67except ImportError:68    USE_COLOR_FIX = False69    print("color fix unavailable; output will not be wavelet-reconstructed")70 71# Weights come from the official ByteDance repo rather than a re-upload. It is72# the canonical source for these files and is Apache-2.0, so there is no mirror73# in the chain that could change under us.74WEIGHTS_REPO = "ByteDance-Seed/SeedVR2-3B"75CKPT_DIR = Path("./ckpts")76CKPT_DIR.mkdir(exist_ok=True)77 78 79def _fetch(filename: str, target: Path) -> str:80    if target.exists():81        return str(target)82    path = hf_hub_download(repo_id=WEIGHTS_REPO, filename=filename)83    target.symlink_to(path)84    return str(target)85 86 87DIT_CKPT = _fetch("seedvr2_ema_3b.pth", CKPT_DIR / "seedvr2_ema_3b.pth")88VAE_CKPT = _fetch("ema_vae.pth", CKPT_DIR / "ema_vae.pth")89# The text branch is conditioned by two fixed embeddings shipped with the90# weights, so this Space runs no text encoder at all.91POS_EMB = _fetch("pos_emb.pt", Path("./pos_emb.pt"))92NEG_EMB = _fetch("neg_emb.pt", Path("./neg_emb.pt"))93 94# The resolution the model restores at, as an area. Upstream hardcodes95# 2560*1440 for images with `downsample_only=False`, meaning small inputs are96# scaled UP to it and large inputs down, because the model was only trained at97# high resolution.98#99# THIS is what decides whether the VAE decode survives on ZeroGPU. The100# decoder's peak contiguous allocation scales with it, and that allocation is101# what trips "NVML_SUCCESS == r INTERNAL ASSERT FAILED" in the caching102# allocator. Measured: 2560x1440 (the upstream value) fails, 1920x1080 passes103# in 14.7s. Left overridable so the ceiling can be probed without a code104# change, but do not raise the default without re-testing a real run.105#106# The cost is honest to state: the model was trained at high resolution, so107# restoring at 2MP rather than 3.7MP gives up some of its headroom on very108# large outputs. It still upscales ~4.9x from a 360x240 input.109WORK_AREA = int(os.environ.get("SEEDVR2_WORK_AREA", 1920 * 1080))110 111# Single-process "distributed" context. The model code routes every device112# placement through common.distributed, which reads these.113os.environ.setdefault("MASTER_ADDR", "127.0.0.1")114os.environ.setdefault("MASTER_PORT", "12355")115os.environ.setdefault("RANK", "0")116os.environ.setdefault("WORLD_SIZE", "1")117os.environ.setdefault("LOCAL_RANK", "0")118 119_runner = None120 121 122def _ensure_runner():123    """Build the runner once, inside a GPU context.124 125    Deliberately NOT done at module scope, even though ZeroGPU prefers that for126    placement: `init_torch` ends in `dist.init_process_group(backend="nccl")`127    and `torch.cuda.set_device`, which need a real device rather than the CUDA128    emulation that applies outside `@spaces.GPU`. Memoized because129    `init_process_group` raises if called twice, and because a warm worker130    should not reload 3B parameters per request.131    """132    global _runner133    if _runner is not None:134        return _runner135 136    if not torch.distributed.is_initialized():137        init_torch(cudnn_benchmark=False)138 139    runner = VideoDiffusionInfer(load_config(os.path.join("./configs_3b", "main.yaml")))140    OmegaConf.set_readonly(runner.config, False)141    runner.configure_dit_model(device="cuda", checkpoint=DIT_CKPT)142    runner.configure_vae_model()143    if hasattr(runner.vae, "set_memory_limit"):144        runner.vae.set_memory_limit(**runner.config.vae.memory_limit)145 146    _runner = runner147    return _runner148 149 150def _transform():151    return Compose(152        [153            NaResize(resolution=WORK_AREA**0.5, mode="area", downsample_only=False),154            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),155            DivisibleCrop((16, 16)),156            Normalize(0.5, 0.5),157            Rearrange("t c h w -> c t h w"),158        ]159    )160 161 162def _duration(image, steps=1, progress=None) -> int:163    """Every request restores at the same working resolution (WORK_AREA), so cost164    tracks the step count and little else. Measured on a cold worker, then165    given headroom; kept as small as honesty allows, because the request is166    checked against the visitor's remaining quota and a smaller one also ranks167    higher in the ZeroGPU queue."""168    return int(min(90, 14 + 12 * int(steps)))169 170 171@spaces.GPU(duration=_duration)172@torch.no_grad()173def upscale_image(image, steps=1, progress=gr.Progress(track_tqdm=True)):174    if image is None:175        raise gr.Error("Upload an image first.")176 177    # The allocator config has to be read by the process that actually owns the178    # CUDA context. ZeroGPU runs this function in its own worker, so an env var179    # exported at import time in the parent is not proof it applied here —180    # print what the worker sees, and set it directly as well.181    print(f"[alloc] PYTORCH_CUDA_ALLOC_CONF={os.environ.get('PYTORCH_CUDA_ALLOC_CONF')!r}", flush=True)182    try:183        torch.cuda.memory._set_allocator_settings("expandable_segments:True")184        print("[alloc] expandable_segments set at runtime", flush=True)185    except Exception as exc:  # older/newer torch may not expose this186        print(f"[alloc] runtime allocator setting unavailable: {exc}", flush=True)187    print(f"[alloc] work area {WORK_AREA} px", flush=True)188 189    runner = _ensure_runner()190 191    runner.config.diffusion.cfg.scale = 1.0192    runner.config.diffusion.cfg.rescale = 0.0193    runner.config.diffusion.timesteps.sampling.steps = int(steps)194    runner.configure_diffusion()195    # Fixed seed: a restorer should be deterministic for a given input, and the196    # knob was noise in a tool whose job is 'make this photo better'.197    set_seed(666, same_across_ranks=True)198 199    img = Image.open(image).convert("RGB") if isinstance(image, str) else image.convert("RGB")200    tensor = T.ToTensor()(img).unsqueeze(0)  # (t=1, c, h, w)201 202    cond = _transform()(tensor.to("cuda"))203    original = cond204    latents = runner.vae_encode([cond])205 206    text_embeds = {207        "texts_pos": [torch.load(POS_EMB).to("cuda")],208        "texts_neg": [torch.load(NEG_EMB).to("cuda")],209    }210 211    noise = [torch.randn_like(latent) for latent in latents]212    aug_noise = [torch.randn_like(latent) for latent in latents]213 214    def _add_noise(x, aug):215        t = torch.tensor([1000.0], device="cuda") * 0.1216        shape = torch.tensor(x.shape[1:], device="cuda")[None]217        return runner.schedule.forward(x, aug, runner.timestep_transform(t, shape))218 219    conditions = [220        runner.get_condition(n, task="sr", latent_blur=_add_noise(latent, a))221        for n, a, latent in zip(noise, aug_noise, latents)222    ]223 224    with torch.autocast("cuda", torch.bfloat16, enabled=True):225        videos = runner.inference(226            noises=noise, conditions=conditions, dit_offload=False, **text_embeds227        )228 229    sample = videos[0]230    sample = (231        rearrange(sample[:, None], "c t h w -> t c h w")232        if sample.ndim == 3233        else rearrange(sample, "c t h w -> t c h w")234    )235    reference = (236        rearrange(original[:, None], "c t h w -> t c h w")237        if original.ndim == 3238        else rearrange(original, "c t h w -> t c h w")239    )240    if USE_COLOR_FIX:241        sample = wavelet_reconstruction(sample.to("cpu"), reference[: sample.size(0)].to("cpu"))242    else:243        sample = sample.to("cpu")244 245    sample = rearrange(sample, "t c h w -> t h w c")246    sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().to(torch.uint8).numpy()247 248    del latents, conditions, videos249    gc.collect()250    torch.cuda.empty_cache()251 252    return Image.fromarray(sample[0])253 254 255with gr.Blocks(css=UPSAMPLER_CSS, theme=UPSAMPLER_THEME) as demo:256    gr.HTML(257        header_html(258            "SeedVR2 3B Image Upscaler",259            "One-step diffusion restoration that rebuilds real detail in blurry, "260            "compressed, and low-resolution photos.",261        )262    )263 264    with gr.Row():265        with gr.Column():266            image_in = gr.Image(label="Image", type="filepath")267            steps = gr.Slider(1, 4, value=1, step=1, label="Steps")268            run = gr.Button("Upscale Image", variant="primary")269        with gr.Column():270            image_out = gr.Image(label="Result", type="pil")271 272    # No leading slash: gradio prefixes it, and "/upscale_image" here would273    # publish the endpoint as "//upscale_image".274    run.click(275        upscale_image,276        inputs=[image_in, steps],277        outputs=[image_out],278        api_name="upscale_image",279    )280 281    gr.HTML(282        footer_html(283            "SeedVR2-3B is ByteDance's one-step diffusion model for image and video "284            "restoration. It rebuilds genuine texture in photos that are blurry, "285            "heavily compressed, or simply too small, restoring at high resolution "286            "rather than smoothing detail away the way a conventional upscaler does.",287            "https://upsampler.com/free-image-upscaler-no-signup",288            "free image upscaler",289        )290    )291 292demo.launch(ssr_mode=False, show_error=True)293