Spaces:
Running on Zero
Running on Zero
| """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) | |
| 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 | |