4bit build of VOSR 2.0

4bit/config.json

@@ -0,0 +1,67 @@
+{
+ "model_type": "vosr",
+ "version": "2.0",
+ "dim": 1536,
+ "depth": 36,
+ "num_heads": 24,
+ "mlp_ratio": 4,
+ "patch_size": 2,
+ "latent_channels": 16,
+ "train_tokens": 32,
+ "enc_dim": 1024,
+ "encdim_ratio": 3,
+ "dino_depth": 18,
+ "dino_heads": 16,
+ "dino_size": 448,
+ "dino_patch": 14,
+ "vae_base_dim": 96,
+ "vae_dim_mult": [
+  1,
+  2,
+  4,
+  4
+ ],
+ "vae_res_blocks": 2,
+ "latents_mean": [
+  -0.7571,
+  -0.7089,
+  -0.9113,
+  0.1075,
+  -0.1745,
+  0.9653,
+  -0.1517,
+  1.5508,
+  0.4134,
+  -0.0715,
+  0.5517,
+  -0.3632,
+  -0.1922,
+  -0.9497,
+  0.2503,
+  -0.2921
+ ],
+ "latents_std": [
+  2.8184,
+  1.4541,
+  2.3275,
+  2.6558,
+  1.2196,
+  1.7708,
+  2.6052,
+  2.0743,
+  3.2687,
+  2.1526,
+  2.8652,
+  1.5579,
+  1.6382,
+  1.1253,
+  2.8251,
+  1.916
+ ],
+ "quantization": {
+  "group_size": 64,
+  "bits": 4,
+  "mode": "affine"
+ },
+ "precision": "4bit"
+}

4bit/mlx_model.py

@@ -0,0 +1,579 @@
+"""VOSR 2.0 (cswry/VOSR, Apache-2.0) in MLX: one-step generative super-resolution.
+
+The pipeline of inference_vosr_onestep.py (tile_size 0, the default): the low-resolution picture is bicubic-upscaled
+to the output size, encoded by the 2D Qwen-Image VAE, and a 1.4B LightningDiT turns noise into the restored latent in
+one flow step (t 1 -> 0), conditioned on the VAE latent (concatenated) and on DINOv2-L block 18 features of the
+upscaled picture (cross-attention). The VAE decodes, and the wavelet color fix (cv2 GaussianBlur, sigma 5) takes the
+low frequencies from the bicubic upscale.
+
+Layout: NHWC throughout. Only the first 18 DINOv2 blocks are kept (VOSR reads block 18 and never the rest); its
+position embedding is stored already interpolated to the 448x448 input VOSR always uses. Norm statistics, the
+timestep embedding and the color fix run in float32.
+
+    import mlx_model
+    model = mlx_model.load("bf16")
+    sr = mlx_model.upscale(model, rgb_uint8, scale=4, seed=42)
+"""
+
+from __future__ import annotations
+
+import json
+import math
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+CONFIG = {"model_type": "vosr", "version": "2.0", "dim": 1536, "depth": 36, "num_heads": 24, "mlp_ratio": 4,
+          "patch_size": 2, "latent_channels": 16, "train_tokens": 32, "enc_dim": 1024, "encdim_ratio": 3,
+          "dino_depth": 18, "dino_heads": 16, "dino_size": 448, "dino_patch": 14, "vae_base_dim": 96,
+          "vae_dim_mult": [1, 2, 4, 4], "vae_res_blocks": 2,
+          "latents_mean": [-0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715,
+                           0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921],
+          "latents_std": [2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652,
+                          1.5579, 1.6382, 1.1253, 2.8251, 1.916]}
+GROUP = 64
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], np.float32)
+IMAGENET_STD = np.array([0.229, 0.224, 0.225], np.float32)
+
+
+def quantizable(path: str, module) -> bool:
+    return isinstance(module, nn.Linear) and module.weight.shape[-1] % GROUP == 0
+
+
+# ------------------------------------------------------------------------------------------------ ops
+
+def rms(x, weight, eps=1e-6):
+    y = x.astype(mx.float32)
+    y = y * mx.rsqrt(mx.mean(y * y, axis=-1, keepdims=True) + eps)
+    return y.astype(x.dtype) * weight
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+
+    def __call__(self, x):
+        return rms(x, self.weight)
+
+
+def layer_norm(ln, x):
+    return mx.fast.layer_norm(x.astype(mx.float32), ln.weight.astype(mx.float32), ln.bias.astype(mx.float32),
+                              ln.eps).astype(x.dtype)
+
+
+def bicubic_matrix(n_in: int, n_out: int) -> np.ndarray:
+    """torch F.interpolate(mode="bicubic", align_corners=False, antialias=False) along one axis, (n_out, n_in)."""
+    A = -0.75
+    m = np.zeros((n_out, n_in), np.float64)
+    scale = n_in / n_out
+    for o in range(n_out):
+        s = (o + 0.5) * scale - 0.5
+        f = int(np.floor(s))
+        t = s - f
+        w = [((A * (t + 1) - 5 * A) * (t + 1) + 8 * A) * (t + 1) - 4 * A,
+             ((A + 2) * t - (A + 3)) * t * t + 1,
+             ((A + 2) * (1 - t) - (A + 3)) * (1 - t) * (1 - t) + 1,
+             ((A * (2 - t) - 5 * A) * (2 - t) + 8 * A) * (2 - t) - 4 * A]
+        for j in range(4):
+            m[o, min(max(f - 1 + j, 0), n_in - 1)] += w[j]
+    return m.astype(np.float32)
+
+
+def bicubic(x, size):
+    """(B, H, W, C) float32 -> (B, size[0], size[1], C)."""
+    _, h, w, _ = x.shape
+    mh, mw = mx.array(bicubic_matrix(h, size[0])), mx.array(bicubic_matrix(w, size[1]))
+    y = mx.einsum("oh,bhwc->bowc", mh, x)
+    return mx.einsum("pw,bowc->bopc", mw, y)
+
+
+# ------------------------------------------------------------------------------------------------ Qwen-Image VAE (2D)
+
+class RMSNorm2D(nn.Module):
+    """F.normalize over channels * sqrt(C) * gamma, in float32."""
+
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = mx.ones((dim,))
+
+    def __call__(self, x):
+        y = x.astype(mx.float32)
+        n = mx.sqrt(mx.sum(y * y, axis=-1, keepdims=True))
+        y = y / mx.maximum(n, 1e-12) * math.sqrt(y.shape[-1])
+        return (y * self.gamma.astype(mx.float32)).astype(x.dtype)
+
+
+class ResBlock(nn.Module):
+    def __init__(self, cin, cout):
+        super().__init__()
+        self.norm1 = RMSNorm2D(cin)
+        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
+        self.norm2 = RMSNorm2D(cout)
+        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
+        if cin != cout:
+            self.conv_shortcut = nn.Conv2d(cin, cout, 1)
+
+    def __call__(self, x):
+        h = self.conv_shortcut(x) if "conv_shortcut" in self else x
+        x = self.conv1(nn.silu(self.norm1(x)))
+        x = self.conv2(nn.silu(self.norm2(x)))
+        return x + h
+
+
+class VAEAttention(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.norm = RMSNorm2D(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+    def __call__(self, x):
+        b, h, w, c = x.shape
+        q, k, v = mx.split(self.to_qkv(self.norm(x)).reshape(b, 1, h * w, 3 * c), 3, axis=-1)
+        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=c ** -0.5)
+        return self.proj(y.reshape(b, h, w, c)) + x
+
+
+class Down(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = nn.Conv2d(dim, dim, 3, stride=2)
+
+    def __call__(self, x):
+        return self.conv(mx.pad(x, [(0, 0), (0, 1), (0, 1), (0, 0)]))
+
+
+class Up(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.conv = nn.Conv2d(dim, dim // 2, 3, padding=1)
+
+    def __call__(self, x):
+        b, h, w, c = x.shape
+        x = mx.broadcast_to(x[:, :, None, :, None, :], (b, h, 2, w, 2, c)).reshape(b, h * 2, w * 2, c)
+        return self.conv(x)
+
+
+class Mid(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.resnets = [ResBlock(dim, dim), ResBlock(dim, dim)]
+        self.attentions = [VAEAttention(dim)]
+
+    def __call__(self, x):
+        return self.resnets[1](self.attentions[0](self.resnets[0](x)))
+
+
+class Encoder(nn.Module):
+    def __init__(self, dim, z_dim, mult, nres):
+        super().__init__()
+        dims = [dim * u for u in [1] + mult]
+        self.conv_in = nn.Conv2d(3, dims[0], 3, padding=1)
+        self.down_blocks = []
+        for i, (cin, cout) in enumerate(zip(dims[:-1], dims[1:])):
+            for _ in range(nres):
+                self.down_blocks.append(ResBlock(cin, cout))
+                cin = cout
+            if i != len(mult) - 1:
+                self.down_blocks.append(Down(cout))
+        self.mid_block = Mid(cout)
+        self.norm_out = RMSNorm2D(cout)
+        self.conv_out = nn.Conv2d(cout, z_dim, 3, padding=1)
+
+    def __call__(self, x):
+        x = self.conv_in(x)
+        for layer in self.down_blocks:
+            x = layer(x)
+        return self.conv_out(nn.silu(self.norm_out(self.mid_block(x))))
+
+
+class UpBlock(nn.Module):
+    def __init__(self, cin, cout, nres, up):
+        super().__init__()
+        self.resnets = [ResBlock(cin if i == 0 else cout, cout) for i in range(nres + 1)]
+        if up:
+            self.upsamplers = [Up(cout)]
+
+    def __call__(self, x):
+        for r in self.resnets:
+            x = r(x)
+        return self.upsamplers[0](x) if "upsamplers" in self else x
+
+
+class Decoder(nn.Module):
+    def __init__(self, dim, z_dim, mult, nres):
+        super().__init__()
+        dims = [dim * u for u in [mult[-1]] + mult[::-1]]
+        self.conv_in = nn.Conv2d(z_dim, dims[0], 3, padding=1)
+        self.mid_block = Mid(dims[0])
+        self.up_blocks = []
+        for i, (cin, cout) in enumerate(zip(dims[:-1], dims[1:])):
+            self.up_blocks.append(UpBlock(cin // 2 if i > 0 else cin, cout, nres, i != len(mult) - 1))
+        self.norm_out = RMSNorm2D(cout)
+        self.conv_out = nn.Conv2d(cout, 3, 3, padding=1)
+
+    def __call__(self, x):
+        x = self.mid_block(self.conv_in(x))
+        for b in self.up_blocks:
+            x = b(x)
+        return self.conv_out(nn.silu(self.norm_out(x)))
+
+
+class VAE(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        d, z, mult, nres = cfg["vae_base_dim"], cfg["latent_channels"], cfg["vae_dim_mult"], cfg["vae_res_blocks"]
+        self.encoder = Encoder(d, 2 * z, mult, nres)
+        self.quant_conv = nn.Conv2d(2 * z, 2 * z, 1)
+        self.post_quant_conv = nn.Conv2d(z, z, 1)
+        self.decoder = Decoder(d, z, mult, nres)
+        self.z = z
+
+    def encode(self, x):
+        """Pixels in [-1, 1] -> the posterior mode (the mean)."""
+        h = self.quant_conv(self.encoder(x.astype(self.quant_conv.weight.dtype)))
+        return h[..., :self.z].astype(mx.float32)
+
+    def decode(self, z):
+        x = self.decoder(self.post_quant_conv(z.astype(self.post_quant_conv.weight.dtype)))
+        return mx.clip(x.astype(mx.float32), -1, 1)
+
+
+# ------------------------------------------------------------------------------------------------ DINOv2-L
+
+class DinoBlock(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+        self.ls1 = mx.ones((dim,))
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        self.fc1 = nn.Linear(dim, dim * 4)
+        self.fc2 = nn.Linear(dim * 4, dim)
+        self.ls2 = mx.ones((dim,))
+
+    def __call__(self, x):
+        b, n, c = x.shape
+        hd = c // self.heads
+        qkv = self.qkv(layer_norm(self.norm1, x)).reshape(b, n, 3, self.heads, hd).transpose(2, 0, 3, 1, 4)
+        y = mx.fast.scaled_dot_product_attention(qkv[0], qkv[1], qkv[2], scale=hd ** -0.5)
+        x = x + self.ls1 * self.proj(y.transpose(0, 2, 1, 3).reshape(b, n, c))
+        return x + self.ls2 * self.fc2(nn.gelu(self.fc1(layer_norm(self.norm2, x))))
+
+
+class Dino(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        d, p = cfg["enc_dim"], cfg["dino_patch"]
+        g = cfg["dino_size"] // p
+        self.patch_embed = nn.Conv2d(3, d, p, stride=p)
+        self.cls_token = mx.zeros((1, 1, d))
+        self.pos_embed = mx.zeros((1, 1 + g * g, d))
+        self.blocks = [DinoBlock(d, cfg["dino_heads"]) for _ in range(cfg["dino_depth"])]
+
+    def __call__(self, x):
+        """(B, 448, 448, 3) ImageNet-normalized -> (B, 1024, 1024) patch tokens after the last kept block."""
+        x = self.patch_embed(x.astype(self.patch_embed.weight.dtype))
+        b, h, w, c = x.shape
+        x = mx.concatenate([mx.broadcast_to(self.cls_token, (b, 1, c)), x.reshape(b, h * w, c)], axis=1)
+        x = x + self.pos_embed
+        for blk in self.blocks:
+            x = blk(x)
+        return x[:, 1:]
+
+
+# ------------------------------------------------------------------------------------------------ LightningDiT
+
+def rope_tables(h, w, head_dim, train_tokens):
+    """VisionRotaryEmbeddingFast / _get_dynamic_rope: positions scaled to the 32-token training grid; the first half
+    of each head rotates with the row, the second with the column, in interleaved pairs."""
+    half = head_dim // 2
+    inv = 1.0 / (10000 ** (np.arange(0, half, 2)[:half // 2].astype(np.float32) / half))
+
+    def axis(n):
+        t = np.arange(n, dtype=np.float32) / n * train_tokens
+        return np.repeat(np.einsum("i,f->if", t, inv), 2, axis=-1)          # (n, half)
+
+    fh, fw = axis(h), axis(w)
+    f = np.concatenate([np.broadcast_to(fh[:, None], (h, w, half)), np.broadcast_to(fw[None], (h, w, half))], -1)
+    f = f.reshape(h * w, head_dim)
+    return mx.array(np.cos(f)), mx.array(np.sin(f))
+
+
+def rope(x, cos, sin):
+    s = x.shape
+    pairs = x.reshape(*s[:-1], s[-1] // 2, 2)
+    rot = mx.stack([-pairs[..., 1], pairs[..., 0]], axis=-1).reshape(s)
+    return x * cos.astype(x.dtype) + rot * sin.astype(x.dtype)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.q_norm = RMSNorm(dim // heads)
+        self.k_norm = RMSNorm(dim // heads)
+        self.proj = nn.Linear(dim, dim)
+
+    def __call__(self, x, cos, sin):
+        b, n, c = x.shape
+        hd = c // self.heads
+        qkv = self.qkv(x).reshape(b, n, 3, self.heads, hd).transpose(2, 0, 3, 1, 4)
+        q, k = rope(self.q_norm(qkv[0]), cos, sin), rope(self.k_norm(qkv[1]), cos, sin)
+        y = mx.fast.scaled_dot_product_attention(q, k, qkv[2], scale=hd ** -0.5)
+        return self.proj(y.transpose(0, 2, 1, 3).reshape(b, n, c))
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q_linear = nn.Linear(dim, dim)
+        self.k_linear = nn.Linear(dim, dim)
+        self.v_linear = nn.Linear(dim, dim)
+        self.q_norm = RMSNorm(dim // heads)
+        self.k_norm = RMSNorm(dim // heads)
+        self.proj = nn.Linear(dim, dim)
+
+    def __call__(self, x, cond):
+        b, n, c = x.shape
+        m, hd = cond.shape[1], c // self.heads
+        q = self.q_norm(self.q_linear(x).reshape(b, n, self.heads, hd).transpose(0, 2, 1, 3))
+        k = self.k_norm(self.k_linear(cond).reshape(b, m, self.heads, hd).transpose(0, 2, 1, 3))
+        v = self.v_linear(cond).reshape(b, m, self.heads, hd).transpose(0, 2, 1, 3)
+        y = mx.fast.scaled_dot_product_attention(q, k, v, scale=hd ** -0.5)
+        return self.proj(y.transpose(0, 2, 1, 3).reshape(b, n, c))
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.w12 = nn.Linear(dim, 2 * hidden)
+        self.w3 = nn.Linear(hidden, dim)
+
+    def __call__(self, x):
+        a, g = mx.split(self.w12(x), 2, axis=-1)
+        return self.w3(nn.silu(a) * g)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.norm2 = RMSNorm(dim)
+        self.attn = SelfAttention(dim, heads)
+        self.mlp = SwiGLU(dim, int(2 / 3 * dim * mlp_ratio))
+        self.scale_shift_table = mx.zeros((6, dim))
+        self.cross_attn = CrossAttention(dim, heads)
+
+    def __call__(self, x, c0, z, cos, sin):
+        b = x.shape[0]
+        mod = self.scale_shift_table[None] + c0.reshape(b, 6, -1)
+        sh1, sc1, g1, sh2, sc2, g2 = [mod[:, i:i + 1] for i in range(6)]
+        x = x + g1 * self.attn(self.norm1(x) * (1 + sc1) + sh1, cos, sin)
+        x = x + self.cross_attn(x, z)
+        return x + g2 * self.mlp(self.norm2(x) * (1 + sc2) + sh2)
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, dim, freq=256):
+        super().__init__()
+        self.freq = freq
+        self.fc1 = nn.Linear(freq, dim)
+        self.fc2 = nn.Linear(dim, dim)
+
+    def __call__(self, t):
+        half = self.freq // 2
+        freqs = mx.exp(-math.log(10000) * mx.arange(half, dtype=mx.float32) / half)
+        a = t.astype(mx.float32)[:, None] * freqs[None]
+        e = mx.concatenate([mx.cos(a), mx.sin(a)], axis=-1).astype(self.fc1.bias.dtype)
+        return self.fc2(nn.silu(self.fc1(e)))
+
+
+class DiT(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        d, p, lc = cfg["dim"], cfg["patch_size"], cfg["latent_channels"]
+        self.cfg = cfg
+        self.x_embedder = nn.Conv2d(2 * lc, d, p, stride=p)
+        self.t_embedder = TimestepEmbedder(d)
+        self.t_block = nn.Linear(d, 6 * d)
+        self.blocks = [Block(d, cfg["num_heads"], cfg["mlp_ratio"]) for _ in range(cfg["depth"])]
+        self.norm_final = RMSNorm(d)
+        self.final_linear = nn.Linear(d, p * p * lc)
+        self.final_adaln = nn.Linear(d, 2 * d)
+        self.layer_norm = nn.LayerNorm(cfg["enc_dim"], eps=1e-5)
+        self.mlp_ca_fc1 = nn.Linear(cfg["enc_dim"], d * cfg["encdim_ratio"])
+        self.mlp_ca_fc2 = nn.Linear(d * cfg["encdim_ratio"], d)
+
+    def __call__(self, x, t, feats):
+        """x: (B, H, W, 32) [lq latent, z]; t: (B,); feats: (B, M, 1024) DINOv2 tokens -> velocity (B, H, W, 16)."""
+        cfg, dt = self.cfg, self.x_embedder.weight.dtype
+        p, lc = cfg["patch_size"], cfg["latent_channels"]
+        x = self.x_embedder(x.astype(dt))
+        b, h, w, d = x.shape
+        x = x.reshape(b, h * w, d)
+        c = self.t_embedder(t)
+        c0 = self.t_block(nn.silu(c))
+        z = self.mlp_ca_fc2(nn.gelu_approx(self.mlp_ca_fc1(layer_norm(self.layer_norm, feats.astype(dt)))))
+        cos, sin = rope_tables(h, w, d // cfg["num_heads"], cfg["train_tokens"])
+        for blk in self.blocks:
+            x = blk(x, c0, z, cos, sin)
+        shift, scale = mx.split(self.final_adaln(nn.silu(c)), 2, axis=-1)
+        x = self.final_linear(self.norm_final(x) * (1 + scale[:, None]) + shift[:, None])
+        x = x.reshape(b, h, w, p, p, lc).transpose(0, 1, 3, 2, 4, 5).reshape(b, h * p, w * p, lc)
+        return x.astype(mx.float32)
+
+
+class Model(nn.Module):
+    def __init__(self, config=None):
+        super().__init__()
+        self.config = dict(config or CONFIG)
+        self.vae = VAE(self.config)
+        self.dino = Dino(self.config)
+        self.dit = DiT(self.config)
+
+
+def load(folder) -> Model:
+    folder = Path(folder)
+    config = json.loads((folder / "config.json").read_text())
+    model = Model(config)
+    q = config.get("quantization")
+    if q:
+        nn.quantize(model, group_size=q["group_size"], bits=q["bits"], class_predicate=quantizable)
+    model.load_weights(str(folder / "model.safetensors"), strict=True)
+    mx.eval(model.parameters())
+    return model
+
+
+# ------------------------------------------------------------------------------------------------ pipeline
+
+def gaussian_blur(x: np.ndarray, sigma: float = 5.0) -> np.ndarray:
+    """cv2.GaussianBlur(x, (0, 0), sigma) for float32 HxWxC: kernel 8*sigma+1, BORDER_REFLECT_101."""
+    k = int(round(sigma * 4 * 2 + 1)) | 1
+    r = k // 2
+    g = np.exp(-((np.arange(k) - r) ** 2) / (2 * sigma * sigma))
+    g = (g / g.sum()).astype(np.float32)
+    y = mx.array(x)
+    for axis in (0, 1):
+        n = y.shape[axis]
+        idx = np.abs(np.arange(-r, n + r))
+        idx = np.where(idx > n - 1, 2 * (n - 1) - idx, idx)
+        idx = mx.array(idx)
+        yp = mx.take(y, idx, axis=axis)
+        out = 0
+        for j in range(k):
+            sl = [slice(None)] * 3
+            sl[axis] = slice(j, j + n)
+            out = out + float(g[j]) * yp[tuple(sl)]
+        y = out
+    return np.array(y)
+
+
+def wavelet_color_fix(sr: np.ndarray, ref: np.ndarray) -> np.ndarray:
+    t, s = sr.astype(np.float32) / 255.0, ref.astype(np.float32) / 255.0
+    out = np.clip(gaussian_blur(s) + (t - gaussian_blur(t)), 0, 1) * 255.0
+    return out.astype(np.uint8)
+
+
+def restore(model: Model, lq: np.ndarray, noise=None, seed: int = 42) -> np.ndarray:
+    """lq: (H, W, 3) float32 in [-1, 1] at the output size, H and W multiples of 16 -> (H, W, 3) in [-1, 1]."""
+    cfg = model.config
+    x = mx.array(lq)[None]
+    mean, std = mx.array(cfg["latents_mean"]), mx.array(cfg["latents_std"])
+    lat = (model.vae.encode(x) - mean) / std
+    s = cfg["dino_size"]
+    d = mx.clip(bicubic((x + 1) * 0.5, (s, s)), 0, 1)
+    d = (d - mx.array(IMAGENET_MEAN)) / mx.array(IMAGENET_STD)
+    feats = model.dino(d)
+    z = mx.array(noise) if noise is not None else mx.random.normal(lat.shape, key=mx.random.key(seed))
+    u = model.dit(mx.concatenate([lat, z], axis=-1), mx.ones((1,)), feats)
+    out = model.vae.decode((z - u) * std + mean)
+    return np.array(out[0])
+
+
+def upscale(model: Model, rgb: np.ndarray, scale: int = 4, noise=None, seed: int = 42,
+            color_fix: str = "wavelet") -> np.ndarray:
+    """(h, w, 3) uint8 -> (h*scale, w*scale, 3) uint8, as inference_vosr_onestep.py (bicubic upscale with PIL first).
+    `noise`: (1, H/8, W/8, 16) float32 to fix the starting noise (H, W: the output size rounded up to 16)."""
+    from PIL import Image
+    h, w = rgb.shape[:2]
+    H, W = h * scale, w * scale
+    up = np.asarray(Image.fromarray(rgb).resize((W, H), Image.BICUBIC))
+    ph, pw = -H % 16, -W % 16
+    x = np.pad(up, [(0, ph), (0, pw), (0, 0)], mode="reflect") if ph or pw else up
+    sr = restore(model, x.astype(np.float32) / 255.0 * 2 - 1, noise, seed)[:H, :W]
+    sr = (np.clip(sr * 0.5 + 0.5, 0, 1) * 255).astype(np.uint8)        # ToPILImage truncates
+    return wavelet_color_fix(sr, up) if color_fix == "wavelet" else sr
+
+
+# ------------------------------------------------------------------------------------------------ conversion
+
+def _conv(w):
+    return w.transpose(0, 2, 3, 1)
+
+
+def convert(source: Path) -> dict:
+    """runs/source (CSWRY/VOSR: VOSR2/, Qwen-Image-vae-2d/, torch_cache/checkpoints/dinov2_vitl14_pretrain.pth)
+    -> fp32 numpy in this file's layout."""
+    import torch
+    import torch.nn.functional as F
+    from safetensors.numpy import load_file
+
+    source, cfg, out = Path(source), CONFIG, {}
+
+    dit = load_file(str(source / "VOSR2" / "checkpoints" / "ema_model.safetensors"))
+    ren = {"t_embedder.mlp.0.": "t_embedder.fc1.", "t_embedder.mlp.2.": "t_embedder.fc2.", "t_block.1.": "t_block.",
+           "final_layer.norm_final.": "norm_final.", "final_layer.linear.": "final_linear.",
+           "final_layer.adaLN_modulation.1.": "final_adaln.", "mlp_ca.fc1.": "mlp_ca_fc1.",
+           "mlp_ca.fc2.": "mlp_ca_fc2.", "x_embedder.proj.": "x_embedder."}
+    for k, v in dit.items():
+        if "freqs_" in k:
+            continue                                     # RoPE tables: recomputed per size
+        name = k
+        for a, b in ren.items():
+            if name.startswith(a):
+                name = b + name[len(a):]
+        v = v.astype(np.float32)
+        out["dit." + name] = _conv(v) if v.ndim == 4 else v
+
+    vae = load_file(str(source / "Qwen-Image-vae-2d" / "diffusion_pytorch_model.safetensors"))
+    for k, v in vae.items():
+        name = k.replace("resample.1.", "conv.")
+        v = v.astype(np.float32)
+        if name.endswith(".gamma"):
+            v = v.reshape(-1)
+        out["vae." + name] = _conv(v) if v.ndim == 4 else v
+
+    sd = torch.load(source / "torch_cache" / "checkpoints" / "dinov2_vitl14_pretrain.pth", map_location="cpu",
+                    weights_only=True)
+    g = cfg["dino_size"] // cfg["dino_patch"]
+    pos = sd["pos_embed"].float()
+    m = int(math.sqrt(pos.shape[1] - 1))
+    patch = F.interpolate(pos[:, 1:].reshape(1, m, m, -1).permute(0, 3, 1, 2), mode="bicubic", antialias=False,
+                          scale_factor=((g + 0.1) / m, (g + 0.1) / m))
+    assert patch.shape[-2:] == (g, g)
+    out["dino.pos_embed"] = torch.cat([pos[:, :1], patch.permute(0, 2, 3, 1).reshape(1, g * g, -1)], 1).numpy()
+    out["dino.cls_token"] = sd["cls_token"].float().numpy()
+    out["dino.patch_embed.weight"] = _conv(sd["patch_embed.proj.weight"].float().numpy())
+    out["dino.patch_embed.bias"] = sd["patch_embed.proj.bias"].float().numpy()
+    bren = {"attn.qkv.": "qkv.", "attn.proj.": "proj.", "mlp.fc1.": "fc1.", "mlp.fc2.": "fc2.",
+            "ls1.gamma": "ls1", "ls2.gamma": "ls2"}
+    for k, v in sd.items():
+        if not k.startswith("blocks."):
+            continue
+        i = int(k.split(".")[1])
+        if i >= cfg["dino_depth"]:
+            continue
+        rest = k.split(".", 2)[2]
+        for a, b in bren.items():
+            if rest.startswith(a):
+                rest = b + rest[len(a):]
+        out[f"dino.blocks.{i}.{rest}"] = v.float().numpy()
+    return out

4bit/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e2f7b023f9ce25fac5dd49d7d0e4f90a19c4b11e3919a4e7f79b544f4fdee56
+size 1006876833