"""Resident VOSR models with batched, shared-noise latent restoration.""" import gc import json import os import sys import time from pathlib import Path from types import SimpleNamespace os.environ.setdefault("XFORMERS_DISABLED", "1") ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT / "vendor.zip" / "vosr")) sys.path.insert(0, str(ROOT / "vendor.zip" / "dino")) import numpy as np import torch import torch.nn.functional as F from diffusers import AutoencoderKL from huggingface_hub import hf_hub_download from safetensors.torch import load_file from dinov2.hub.backbones import dinov2_vitb14, dinov2_vitl14 from models.light_decoder import LightDecoder from models.lightningdit import LightningDiT from models.qwenimage_vae2d import AutoencoderKLQwenImage2D from geometry import tile_coordinates MODEL_REPO = "CSWRY/VOSR" MODEL_REVISION = "f24b3061b7f350b81e1907bdcfc27f71fb4ff3f3" MODELS = {"Fast": "VOSR_0.5B_os", "Restore": "VOSR2", "Detail": "VOSR_1.4B_ms"} MODEL_FILES = [ "VOSR_1.4B_ms/args.json", "VOSR_1.4B_ms/checkpoints/ema_model.safetensors", "VOSR_0.5B_os/args.json", "VOSR_0.5B_os/checkpoints/ema_model.safetensors", "VOSR2/args.json", "VOSR2/checkpoints/ema_model.safetensors", "Qwen-Image-vae-2d/config.json", "Qwen-Image-vae-2d/diffusion_pytorch_model.safetensors", "stable-diffusion-2-1-base/vae/config.json", "stable-diffusion-2-1-base/vae/diffusion_pytorch_model.safetensors", "sd21_lwdecoder.pth", "torch_cache/checkpoints/dinov2_vitb14_pretrain.pth", "torch_cache/checkpoints/dinov2_vitl14_pretrain.pth", ] def model_file(name): return Path(hf_hub_download(MODEL_REPO, name, revision=MODEL_REVISION)) def gaussian(height, width, device): y = torch.linspace(-1, 1, height, device=device) x = torch.linspace(-1, 1, width, device=device) return torch.exp(-(y[:, None].square() + x[None, :].square()) / 0.32).clamp_min(1e-4)[None, None] class RestorationModel: def __init__(self, mode, device="cuda", shared=None): self.device = device self.mode = mode folder = MODELS[mode] self.args = SimpleNamespace(**json.loads(model_file(f"{folder}/args.json").read_text())) a = self.args auxiliary = getattr(a, "auxiliary_time_cond", None) a.auxiliary_time_cond = (mode != "Detail") if auxiliary is None else bool(auxiliary) if shared is not None and (a.enc_type, a.ae_type) != (shared.args.enc_type, shared.args.ae_type): raise ValueError("Shared components must match encoder and VAE types") # Allocate in reduced precision from the start; no huge temporary fp32 backbone. previous = torch.get_default_dtype() try: torch.set_default_dtype(torch.bfloat16) self.model = LightningDiT( input_size=a.resolution // 8, patch_size=a.patch_size, in_channels=32 if a.ae_type == "qwen" else 8, out_channels=16 if a.ae_type == "qwen" else 4, hidden_size=a.dim, depth=a.depth, num_heads=a.num_heads, mlp_ratio=a.mlp_ratio, z_dims=a.enc_dim, encdim_ratio=a.encdim_ratio, auxiliary_time_cond=a.auxiliary_time_cond, use_qknorm=a.use_qknorm, use_swiglu=a.use_swiglu, use_rope=a.use_rope, use_rmsnorm=a.use_rmsnorm, wo_shift=a.wo_shift, num_fused_layers=len(a.layer_dinov2b_list), ) builder = dinov2_vitl14 if a.enc_type == "dinov2l" else dinov2_vitb14 self.encoder = shared.encoder if shared is not None else builder(pretrained=False) finally: torch.set_default_dtype(previous) self.model.load_state_dict(load_file(str(model_file(f"{folder}/checkpoints/ema_model.safetensors"))), strict=True) self.model.forward = self.model.forward_flexible dino_name = "dinov2_vitl14_pretrain.pth" if a.enc_type == "dinov2l" else "dinov2_vitb14_pretrain.pth" if shared is None: self.encoder.load_state_dict(torch.load(model_file(f"torch_cache/checkpoints/{dino_name}"), map_location="cpu", weights_only=True), strict=True) self.model.to(device).eval().requires_grad_(False) self.encoder.to(device).eval().requires_grad_(False) self.light_decoder = None if shared is not None: self.vae = shared.vae self.light_decoder = shared.light_decoder elif a.ae_type == "qwen": model_file("Qwen-Image-vae-2d/diffusion_pytorch_model.safetensors") vae_dir = model_file("Qwen-Image-vae-2d/config.json").parent self.vae = AutoencoderKLQwenImage2D.from_pretrained(vae_dir, local_files_only=True) else: model_file("stable-diffusion-2-1-base/vae/diffusion_pytorch_model.safetensors") vae_dir = model_file("stable-diffusion-2-1-base/vae/config.json").parent self.vae = AutoencoderKL.from_pretrained(vae_dir, local_files_only=True) state = torch.load(model_file("sd21_lwdecoder.pth"), map_location="cpu", weights_only=True) cfg = state["config"] self.light_decoder = LightDecoder( in_channels=cfg["in_channels"], out_channels=cfg["out_channels"], block_out_channels=tuple(cfg["block_out_channels"]), layers_per_block=cfg["layers_per_block"], ) self.light_decoder.load_state_dict(state["model_state_dict"], strict=True) self.light_decoder.to(device).eval().requires_grad_(False) self.vae.to(device).eval().requires_grad_(False) gc.collect() print(f"Loaded {mode}: {folder}", flush=True) def features(self, pixels): a = self.args x = F.interpolate((pixels + 1) / 2, size=(a.dinov2_size, a.dinov2_size), mode="bicubic", align_corners=False).clamp(0, 1) mean = x.new_tensor([0.485, 0.456, 0.406])[None, :, None, None] std = x.new_tensor([0.229, 0.224, 0.225])[None, :, None, None] x = self.encoder.prepare_tokens_with_masks(((x - mean) / std).to(torch.bfloat16)) found = {} wanted = a.layer_dinov2b_list for i, block in enumerate(self.encoder.blocks): x = block(x) if i in wanted: # Upstream replaces the final block feature with its normalized output. found[i] = (self.encoder.norm(x) if i == len(self.encoder.blocks) - 1 else x)[:, 1:] if i >= max(wanted): break # Unused DINO blocks need not run. return [found[i] for i in wanted] def encode(self, pixels): posterior = self.vae.encode(pixels).latent_dist z = posterior.mode().float() if self.args.ae_type == "sd2": return z * self.vae.config.scaling_factor mean = z.new_tensor(self.vae.config.latents_mean)[None, :, None, None] std = z.new_tensor(self.vae.config.latents_std)[None, :, None, None] return (z - mean) / std def decode(self, latent): if self.args.ae_type == "sd2": return self.light_decoder(latent / self.vae.config.scaling_factor).clamp(-1, 1) mean = latent.new_tensor(self.vae.config.latents_mean)[None, :, None, None] std = latent.new_tensor(self.vae.config.latents_std)[None, :, None, None] return self.vae.decode(latent * std + mean, return_dict=False)[0].clamp(-1, 1) @torch.inference_mode() def restore(self, image, seed=42, batch_size=4, tile_size=512, progress=None, overlap=64, steps=10, guidance=0.5, weak_condition=0.1): timings = {} start = time.perf_counter() x = torch.from_numpy(np.array(image, dtype=np.float32) / 127.5 - 1).permute(2, 0, 1)[None].to(self.device) _, _, h, w = x.shape lh, lw = h // 8, w // 8 # VAE tiles may be rectangular; DiT tiles remain square and stride-aligned. vh, vw = min(h, 1024), min(w, 1024) vae_coords = tile_coordinates(h, w, vh, vw, 128) latent = None weights = x.new_zeros((1, 1, lh, lw)) mask = gaussian(vh // 8, vw // 8, self.device) for n, (y, col) in enumerate(vae_coords): tile = self.encode(x[:, :, y:y + vh, col:col + vw]) if latent is None: latent = x.new_zeros((1, tile.shape[1], lh, lw)) yy, xx = y // 8, col // 8 latent[:, :, yy:yy + vh // 8, xx:xx + vw // 8] += tile * mask weights[:, :, yy:yy + vh // 8, xx:xx + vw // 8] += mask if progress: progress(0.05 + 0.15 * (n + 1) / len(vae_coords), desc="Encoding image") latent /= weights.clamp_min(1e-8) torch.cuda.synchronize() timings["encode"] = time.perf_counter() - start stage = time.perf_counter() size = tile_size // 8 if overlap not in (64, 128, 192): raise ValueError("Overlap must be 64, 128, or 192 pixels") coords = tile_coordinates(lh, lw, size, size, overlap // 8) generator = torch.Generator(device=self.device).manual_seed(int(seed)) state = torch.randn(latent.shape, device=self.device, dtype=torch.float32, generator=generator) multistep = getattr(self, "mode", "Fast") == "Detail" n_steps = int(steps) if multistep else 1 if n_steps not in (1, 5, 10, 15, 25): raise ValueError("Unsupported sampling steps") schedule = torch.linspace(1., 0., n_steps + 1, device=self.device) mask = gaussian(size, size, self.device) # DINO conditions are constant over the sampling trajectory: compute once. feature_cache = [] for offset in range(0, len(coords), batch_size): group = coords[offset:offset + batch_size] pixels = torch.cat([x[:, :, y * 8:(y + size) * 8, col * 8:(col + size) * 8] for y, col in group]) with torch.autocast(device_type=self.device, dtype=torch.bfloat16): feature_cache.append(self.features(pixels)) for step in range(n_steps): accum = torch.zeros_like(latent) weights = latent.new_zeros((1, 1, lh, lw)) for group_index, offset in enumerate(range(0, len(coords), batch_size)): group = coords[offset:offset + batch_size] condition = torch.cat([latent[:, :, y:y + size, col:col + size] for y, col in group]) z = torch.cat([state[:, :, y:y + size, col:col + size] for y, col in group]) features = feature_cache[group_index] count = len(group) inputs = torch.cat((condition, z), dim=1) # The teacher compares full visual conditioning to weak conditioning. if multistep: weak = torch.cat((condition * float(weak_condition), z), dim=1) inputs = torch.cat((inputs, weak)) features = [torch.cat((f, torch.zeros_like(f))) for f in features] count = inputs.shape[0] with torch.autocast(device_type=self.device, dtype=torch.bfloat16): velocity = self.model(inputs.to(torch.bfloat16), schedule[step].expand(count), schedule[step + 1].expand(count), features).float() if multistep: full, weak = velocity.chunk(2) velocity = weak + float(guidance) * (full - weak) for idx, (y, col) in enumerate(group): accum[:, :, y:y + size, col:col + size] += velocity[idx:idx + 1] * mask weights[:, :, y:y + size, col:col + size] += mask if progress: done = step * len(coords) + min(offset + batch_size, len(coords)) progress(0.2 + 0.6 * done / (n_steps * len(coords)), desc=f"Sampling {step + 1}/{n_steps}, tiles {min(offset + batch_size, len(coords))}/{len(coords)}") # Update the shared global latent only after every tile for this step. state -= (schedule[step] - schedule[step + 1]) * accum / weights.clamp_min(1e-8) result_latent = state del feature_cache torch.cuda.synchronize() timings["restore"] = time.perf_counter() - stage stage = time.perf_counter() out = x.new_zeros((1, 3, h, w)) weights = x.new_zeros((1, 1, h, w)) mask = gaussian(vh, vw, self.device) for n, (y, col) in enumerate(vae_coords): pixels = self.decode(result_latent[:, :, y // 8:(y + vh) // 8, col // 8:(col + vw) // 8]) out[:, :, y:y + vh, col:col + vw] += pixels * mask weights[:, :, y:y + vh, col:col + vw] += mask if progress: progress(0.8 + 0.18 * (n + 1) / len(vae_coords), desc="Decoding image") out = (out / weights.clamp_min(1e-8)).clamp(-1, 1) if not torch.isfinite(out).all(): raise RuntimeError("The model produced nonfinite pixels. Please retry with Fast mode.") result = ((out[0].permute(1, 2, 0) + 1) * 127.5).round().clamp(0, 255).to(torch.uint8).cpu().numpy() timings["decode"] = time.perf_counter() - stage timings["gpu_seconds"] = time.perf_counter() - start timings["tiles"] = len(coords) timings["batch_size"] = batch_size timings["steps"] = n_steps timings["model_evaluations"] = len(coords) * n_steps * (2 if multistep else 1) return result, timings