Blame for engine.py

"""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