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