| """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 |
|
|
|
|
| |
|
|
| 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) |
|
|
|
|
| |
|
|
| 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) |
|
|
|
|
| |
|
|
| 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:] |
|
|
|
|
| |
|
|
| 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) |
|
|
| 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 |
|
|
|
|
| |
|
|
| 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) |
| return wavelet_color_fix(sr, up) if color_fix == "wavelet" else sr |
|
|
|
|
| |
|
|
| 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 |
| 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 |
|
|
