Publish PixelModel v6: MMDiT + REPA, 150k steps, FID 23.62 at cfg 3.0

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples.png filter=lfs diff=lfs merge=lfs -text
+model.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,192 @@
+---
+license: mit
+pipeline_tag: text-to-image
+language:
+- en
+tags:
+- text-to-image
+- diffusion
+- rectified-flow
+- mmdit
+- tiny
+- custom-code
+---
+
+# PixelModel v6 🖼️
+
+<img src="examples.png" alt="Two 256px images generated by PixelModel v6: a lighthouse on a cliff at sunset, and a cup of coffee on a wooden table" style="width: 100%; max-width: 600px;">
+
+*"a lighthouse on a cliff at sunset" / "a cup of coffee on a wooden table", cfg 4.0, single draw each, not
+cherry-picked from multiple seeds*
+
+PixelModel v6 is a 155,090,448 parameter text-to-image diffusion transformer. It generates 256x256
+images in 50 sampling steps and it trained in about 55 hours across two A100s.
+
+Where [v5](https://huggingface.co/bench-labs/PixelModel-v5) changed the data and kept v4's
+architecture byte for byte, v6 changes the architecture and keeps v5's data. v5's own README ended
+with a diagnosis: the model had learned appearance thoroughly and topology poorly, and topology is
+what capacity buys. This is that bet.
+
+## What actually changed
+
+Six real architecture changes, not one, which is a departure from this line's usual practice of
+isolating a single variable per release. Each change has independent literature support; the
+combination does not, and that is worth being upfront about rather than implying six papers'
+results simply add up.
+
+**Joint attention instead of cross-attention.** v4 and v5 read text through cross-attention, the way
+PixArt-alpha does it. v6 concatenates image and text tokens into one sequence and runs joint
+self-attention over both, the way SD3's MMDiT does it, so text tokens can attend to image tokens and
+not just the reverse.
+
+**T5 instead of CLIP-only conditioning.** The joint stream is built from `google/flan-t5-base`
+hidden states rather than CLIP token embeddings. CLIP's pooled vector is kept, but only for the
+global adaLN conditioning signal. T5 conditioning is the specific change PixArt-alpha, DeepFloyd IF,
+and SD3 all made for prompt adherence, and v5's README named exactly this failure mode: prompts
+asking for one object in an empty scene often returned the scene without the object.
+
+**2D RoPE instead of additive sincos position embeddings.** Image tokens get axial rotary position
+encoding, text tokens get standard 1D rotary encoding. This is also what makes a future higher
+resolution fine-tune cheap: RoPE carries no learned parameters, so a checkpoint trained at 256px can
+extend to a larger token grid without retraining or interpolating a position embedding table.
+
+**QK-norm and SwiGLU.** RMSNorm on queries and keys before attention, and a SwiGLU MLP in place of
+GELU. Both are standard modern stabilizers, included because a bigger model trains less forgivingly
+than a 40M one.
+
+**REPA.** An auxiliary loss projecting mid-stack image token hidden states to align with frozen
+DINOv2-small features of the same image, cosine similarity, weight scheduled to peak early and decay
+to exactly zero by step 105,000. This is the one item on this list without solid precedent at this
+scale: the original REPA paper is class-conditional ImageNet at 675M parameters, and v5's own
+planning notes flagged that transfer to text-conditional generation at a fraction of that size as
+genuinely unverified. The decay schedule follows the follow-up finding that the alignment loss helps
+early and can hurt if left on too long.
+
+**More capacity.** dim 512, depth 16, heads 8: 155.1M trainable, versus v5's 40M. Still well short of
+SD1.5's 860M UNet or SDXL's 2.6B, but the largest this line has been.
+
+<img src="pixelmodel-v6-params.png" alt="Trainable backbone size by version, v4 and v5 at 40M, v6 at 155.1M" style="width: 100%; max-width: 700px;">
+
+The frozen VAE also moved from `sd-vae-ft-mse` to `madebyollin/sdxl-vae-fp16-fix`, a drop-in upgrade
+at the same latent shape.
+
+## Training
+
+<img src="pixelmodel-v6-loss.png" alt="Training loss, diffusion and REPA components, over 150,000 steps" style="width: 100%; max-width: 860px;">
+
+150,000 steps at batch 192 on the same 2,986,571 image CC12M-recaptioned set v5 used, about 9.7
+epochs. AdamW at 2e-4, cosine decay after a 1,500 step warmup, EMA at 0.9999, 10% caption dropout,
+gradient checkpointing throughout. The REPA loss ran at peak weight for the first 40% of training,
+decayed linearly through the next 30%, and sat at exactly zero for the final 30%, both to follow the
+early-stopping finding above and because computing DINOv2 targets costs real throughput: the run held
+0.86 steps/s while REPA was active and 1.39 steps/s once it turned off.
+
+<img src="pixelmodel-v6-val.png" alt="Held-out validation loss falling monotonically from 0.938 to 0.798" style="width: 100%; max-width: 860px;">
+
+Held-out loss on a fixed 4,096-pair split, evaluated every 2,000 steps, fell without interruption
+from 0.938 to 0.798 across the entire run. No overfitting turn-up at any point, same clean shape v5
+reported.
+
+This run was not a clean single continuous job. It moved between two rented A100s partway through
+(a direct box-to-box transfer once the economics of a cheaper card made sense, verified byte-identical
+by checksum before the first box was released), and survived two host-level outages on the vast.ai
+side, both recovered from the last saved checkpoint. Total lost progress across both interruptions
+was under a thousand steps out of 150,000. Worth stating plainly rather than pretending the run was
+tidier than it was.
+
+## Results
+
+Measured on 5,000 MS-COCO val2014 images at 256 center crop, `torchmetrics` FID at 2048 features and
+`torchmetrics` CLIPScore with `openai/clip-vit-base-patch32`. Identical protocol to v4 and v5.
+
+| | FID ↓ | CLIP ↑ |
+|---|---|---|
+| PixelModel v4 | 39.54 | 28.04 |
+| PixelModel v5 | 34.06 | 29.27 |
+| **PixelModel v6, cfg 3.0** | **23.62** | 30.43 |
+| real photographs | n/a | 30.05 |
+
+v6 beats v5's best FID by over 10 points at every guidance scale tested, not just at one cherry-picked
+value. Full sweep, same 5,000 image set and 50 sampling steps throughout:
+
+| cfg | FID ↓ | CLIP ↑ |
+|---|---|---|
+| 2.0 | 25.65 | 29.65 |
+| **3.0** | **23.62** | 30.43 |
+| 4.0 | 23.77 | 30.81 |
+| 5.0 | 24.42 | 31.06 |
+| 6.0 | 25.10 | 31.12 |
+| 7.0 | 26.26 | 31.16 |
+| 8.0 | 26.87 | 31.17 |
+
+FID bottoms at cfg 3.0 and rises steadily past it, the usual guidance tradeoff: more weight on the
+text condition pulls generations away from the real-image manifold even as it pulls them closer to
+the prompt. cfg 3.0 is the reported default. Worth flagging rather than celebrating: CLIP score
+crosses the real-photograph baseline (30.05) starting around cfg 4.0 and keeps climbing to 31.17.
+That is not v6 out-photographing real photographs, it is high guidance making generations more
+literal about the prompt than an actual photo ever is, which CLIP score rewards and FID does not. Use
+FID, not CLIP, to pick a guidance scale.
+
+## What it does well and what it does not
+
+Judged across several hundred samples during training, not just the checkpoint above.
+
+It is good at material, light, and increasingly at structure that v5 could not hold: food, faces of
+common animals, latte foam texture, boat masts and sails, neon-lit streets. Where v5's README singled
+out faces and animal heads as a weak point, v6 renders recognizable dog and fox faces reliably by the
+back half of training, something no earlier version in this line managed.
+
+It is still uneven on rare subjects. Across many samples of the same fox prompt, output quality
+varied far more than for common subjects like coffee or strawberries, which is a direct read on
+training data frequency in CC12M rather than a model defect: the architecture change did not fix data
+scarcity, it just makes better use of what is there.
+
+A watermark and ghost-text artifact shows up intermittently, most often on the lighthouse and boat
+prompts: faint illegible text-like marks bled in from CC12M's watermarked stock photography, the same
+failure family v5 documented for actual text rendering. It did not fully train out by step 150,000.
+
+## Specification
+
+| | |
+|---|---|
+| trainable parameters | 155,090,448 |
+| frozen VAE, `sdxl-vae-fp16-fix` | 83,653,863 |
+| frozen text encoder, CLIP ViT-B/32 | 63,165,952 |
+| frozen text encoder, T5-base | 109,628,544 |
+| **combined total** | **411,538,807** |
+| architecture | MMDiT, dim 512, depth 16, heads 8, patch 2 |
+| tokens | 256 image (16x16 over a 32x32x4 latent) + 32 text |
+| objective | rectified flow, target `x1 - x0`, logit-normal timesteps |
+| conditioning | joint attention on T5 tokens, plus adaLN-zero from CLIP pooled + timestep |
+| resolution | 256x256 |
+| sampling | 50 steps, cfg 3.0 default (23.62 FID); see cfg sweep above for the tradeoff |
+
+All parameter counts above are the released backbone. Training also carried a small REPA projection
+head (459,648 params, 155,550,096 total while training) that only ever fed the auxiliary alignment
+loss; it has no role in sampling and is dropped from the published weights.
+
+## Usage
+
+```bash
+python main.py "a bowl of ramen with a soft boiled egg" --out ramen.png --cfg 3.0 --steps 50
+```
+
+As with every model in this line the weights are also stored as a PNG. `model.png` is not a picture
+of the network, it is the network, with each parameter packed across the red and green channels of
+one pixel.
+
+## Files
+
+| file | contents |
+|---|---|
+| `model.safetensors` | the trainable MMDiT backbone only, fp32. VAE, CLIP, and T5 load fresh from their own HF repos at inference time, same convention as v4 and v5 despite what their READMEs said |
+| `model.png` | the same backbone weights as an image, fp16, packed via `png_codec.py`. Verified bit-exact against `model.safetensors` (rounded to fp16) before upload |
+| `model_png.json` | manifest `png_codec.py` needs to decode `model.png` back into a state dict |
+| `png_codec.py` | the PNG weight codec, encode and decode |
+| `dit_v6.py` | the model |
+| `train_v6.py` | training loop, REPA schedule, held-out split |
+| `prep_v6.py` | streams CC12M-recaptioned, encodes latents, caches tokens |
+| `build_eval_set.py` | builds the 5,000 image COCO val2014 eval set |
+| `eval_v6.py` | FID / CLIPScore protocol and cfg sweep |
+| `preview.py` | quick sampling CLI for spot checks |
+| `main.py` | inference |

build_eval_set.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import argparse
+import io
+import json
+import os
+import random
+import zipfile
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import requests
+from PIL import Image
+
+ANN_URL = "http://images.cocodataset.org/annotations/annotations_trainval2014.zip"
+
+def csr(img, size):
+    img = img.convert("RGB")
+    w, h = img.size
+    s = min(w, h)
+    l, t = (w - s) // 2, (h - s) // 2
+    return np.asarray(img.crop((l, t, l + s, t + s)).resize((size, size), Image.BICUBIC), dtype=np.uint8)
+
+def fetch_one(item, size):
+    url, cap = item
+    for _ in range(3):
+        try:
+            r = requests.get(url, timeout=15)
+            if r.status_code == 200:
+                return csr(Image.open(io.BytesIO(r.content)), size), cap
+        except Exception:
+            pass
+    return None
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--work", default="/root/v6cache")
+    ap.add_argument("--n", type=int, default=5000)
+    ap.add_argument("--size", type=int, default=256)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--workers", type=int, default=48)
+    args = ap.parse_args()
+    os.makedirs(args.work, exist_ok=True)
+
+    out_path = os.path.join(args.work, "eval_256.npz")
+    if os.path.exists(out_path):
+        print(f"[eval-set] already exists at {out_path}", flush=True)
+        return
+
+    ann_path = os.path.join(args.work, "captions_val2014.json")
+    if not os.path.exists(ann_path):
+        print("[eval-set] downloading annotations", flush=True)
+        z = os.path.join(args.work, "ann.zip")
+        with requests.get(ANN_URL, stream=True, timeout=120) as r:
+            with open(z, "wb") as f:
+                for chunk in r.iter_content(1 << 20):
+                    f.write(chunk)
+        with zipfile.ZipFile(z) as zf:
+            with zf.open("annotations/captions_val2014.json") as src, open(ann_path, "wb") as dst:
+                dst.write(src.read())
+        os.remove(z)
+
+    ann = json.load(open(ann_path))
+    url_by_id = {im["id"]: im["coco_url"] for im in ann["images"]}
+    cap_by_id = {}
+    for a in ann["annotations"]:
+        cap_by_id.setdefault(a["image_id"], a["caption"])
+    items = [(url_by_id[i], cap_by_id[i]) for i in cap_by_id if i in url_by_id]
+    random.Random(args.seed).shuffle(items)
+    print(f"[eval-set] {len(items)} val2014 pairs available, target {args.n}", flush=True)
+
+    imgs, caps = [], []
+    pool = ThreadPoolExecutor(max_workers=args.workers)
+    idx, batch = 0, 64
+    while len(imgs) < args.n and idx < len(items):
+        chunk = items[idx:idx + batch]
+        idx += batch
+        results = [r for r in pool.map(lambda it: fetch_one(it, args.size), chunk) if r is not None]
+        for a, c in results:
+            imgs.append(a); caps.append(c)
+        if idx % (batch * 20) == 0:
+            print(f"[eval-set] {len(imgs)}/{args.n}", flush=True)
+
+    imgs = np.stack(imgs[:args.n])
+    caps = np.array(caps[:args.n], dtype=object)
+    np.savez(out_path, images=imgs, captions=caps)
+    print(f"[eval-set] DONE {imgs.shape} -> {out_path}", flush=True)
+
+if __name__ == "__main__":
+    main()

config.json

@@ -0,0 +1,51 @@
+{
+  "architecture": "latent-mmdit",
+  "objective": "rectified-flow",
+  "dit": {
+    "dim": 512,
+    "depth": 16,
+    "heads": 8,
+    "patch": 2,
+    "latent_ch": 4,
+    "latent_size": 32,
+    "mlp_hidden": 1408,
+    "t5_len": 32
+  },
+  "vae": "madebyollin/sdxl-vae-fp16-fix",
+  "text_encoders": {
+    "joint_stream": "google/flan-t5-base",
+    "global_conditioning": "openai/clip-vit-base-patch32"
+  },
+  "max_tokens": {
+    "t5": 32,
+    "clip": 40
+  },
+  "trainable_parameters": 155090448,
+  "total_parameters_incl_frozen": 411538807,
+  "sampling": {
+    "steps": 50,
+    "cfg": 3.0
+  },
+  "training": {
+    "steps": 150000,
+    "batch": 192,
+    "dataset": "undefined443/cc12m-wds-coco-recaptioned, 2986571 images",
+    "final_val_loss": 0.7982857534661889
+  },
+  "eval": {
+    "fid": 23.62,
+    "clip_score": 30.43,
+    "cfg_sweep": [
+      {"cfg": 2.0, "fid": 25.65, "clip_score": 29.65},
+      {"cfg": 3.0, "fid": 23.62, "clip_score": 30.43},
+      {"cfg": 4.0, "fid": 23.77, "clip_score": 30.81},
+      {"cfg": 5.0, "fid": 24.42, "clip_score": 31.06},
+      {"cfg": 6.0, "fid": 25.10, "clip_score": 31.12},
+      {"cfg": 7.0, "fid": 26.26, "clip_score": 31.16},
+      {"cfg": 8.0, "fid": 26.87, "clip_score": 31.17}
+    ],
+    "n": 5000,
+    "dataset": "MS-COCO val2014, 256 center crop",
+    "protocol": "torchmetrics FrechetInceptionDistance + CLIPScore (openai/clip-vit-base-patch32)"
+  }
+}

dit_v6.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+def timestep_embedding(t, dim, max_period=10000):
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
+    args = t[:, None].float() * freqs[None]
+    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+    return emb
+
+def rope_freqs(positions, dim, base=10000.0):
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+    return torch.outer(positions.float(), inv_freq)
+
+def rope_cos_sin(freqs):
+    emb = torch.cat([freqs, freqs], dim=-1)
+    return emb.cos(), emb.sin()
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+def apply_rope(x, cos, sin):
+    return x * cos + rotate_half(x) * sin
+
+def apply_rope_2d(x, row_cos, row_sin, col_cos, col_sin):
+    x1, x2 = x.chunk(2, dim=-1)
+    x1 = apply_rope(x1, row_cos, row_sin)
+    x2 = apply_rope(x2, col_cos, col_sin)
+    return torch.cat([x1, x2], dim=-1)
+
+class RMSNormHead(nn.Module):
+    def __init__(self, head_dim, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(head_dim))
+        self.eps = eps
+
+    def forward(self, x):
+        n = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return x * n * self.weight
+
+class SwiGLU(nn.Module):
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden)
+        self.up = nn.Linear(dim, hidden)
+        self.down = nn.Linear(hidden, dim)
+
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+class JointBlock(nn.Module):
+    def __init__(self, dim, heads, mlp_hidden):
+        super().__init__()
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.qkv_img = nn.Linear(dim, 3 * dim)
+        self.qkv_txt = nn.Linear(dim, 3 * dim)
+        self.qn_img = RMSNormHead(self.head_dim)
+        self.kn_img = RMSNormHead(self.head_dim)
+        self.qn_txt = RMSNormHead(self.head_dim)
+        self.kn_txt = RMSNormHead(self.head_dim)
+        self.proj_img = nn.Linear(dim, dim)
+        self.proj_txt = nn.Linear(dim, dim)
+        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.mlp_img = SwiGLU(dim, mlp_hidden)
+        self.mlp_txt = SwiGLU(dim, mlp_hidden)
+        self.ada_img = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.ada_txt = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+
+    def forward(self, img, txt, c, rope_img, rope_txt, key_valid):
+        s1i, sc1i, g1i, s2i, sc2i, g2i = self.ada_img(c).chunk(6, dim=-1)
+        s1t, sc1t, g1t, s2t, sc2t, g2t = self.ada_txt(c).chunk(6, dim=-1)
+
+        xi = modulate(self.norm1_img(img), s1i, sc1i)
+        xt = modulate(self.norm1_txt(txt), s1t, sc1t)
+
+        B, Ni, C = xi.shape
+        Nt = xt.shape[1]
+        H, D = self.heads, self.head_dim
+
+        qi, ki, vi = self.qkv_img(xi).reshape(B, Ni, 3, H, D).permute(2, 0, 3, 1, 4)
+        qt, kt, vt = self.qkv_txt(xt).reshape(B, Nt, 3, H, D).permute(2, 0, 3, 1, 4)
+
+        qi, ki = self.qn_img(qi), self.kn_img(ki)
+        qt, kt = self.qn_txt(qt), self.kn_txt(kt)
+
+        row_cos, row_sin, col_cos, col_sin = rope_img
+        qi = apply_rope_2d(qi, row_cos, row_sin, col_cos, col_sin)
+        ki = apply_rope_2d(ki, row_cos, row_sin, col_cos, col_sin)
+
+        t_cos, t_sin = rope_txt
+        qt = apply_rope(qt, t_cos, t_sin)
+        kt = apply_rope(kt, t_cos, t_sin)
+
+        q = torch.cat([qi, qt], dim=2)
+        k = torch.cat([ki, kt], dim=2)
+        v = torch.cat([vi, vt], dim=2)
+
+        mask = key_valid[:, None, None, :]
+        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+        o = o.transpose(1, 2).reshape(B, Ni + Nt, C)
+        oi, ot = o[:, :Ni], o[:, Ni:]
+
+        img = img + g1i.unsqueeze(1) * self.proj_img(oi)
+        txt = txt + g1t.unsqueeze(1) * self.proj_txt(ot)
+
+        img = img + g2i.unsqueeze(1) * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i))
+        txt = txt + g2t.unsqueeze(1) * self.mlp_txt(modulate(self.norm2_txt(txt), s2t, sc2t))
+        return img, txt
+
+class MMDiT(nn.Module):
+    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=512, depth=16, heads=8,
+                 t5_dim=768, clip_dim=512, t5_len=32, mlp_hidden=1408,
+                 repa_dim=384, repa_layer=8):
+        super().__init__()
+        self.latent_ch = latent_ch
+        self.latent_size = latent_size
+        self.patch = patch
+        self.grid = latent_size // patch
+        self.patch_dim = latent_ch * patch * patch
+        self.dim = dim
+        self.depth = depth
+        self.heads = heads
+        self.head_dim = dim // heads
+        self.t5_len = t5_len
+        self.repa_layer = repa_layer
+
+        self.x_embed = nn.Linear(self.patch_dim, dim)
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.clip_proj = nn.Linear(clip_dim, dim)
+        self.t5_proj = nn.Linear(t5_dim, dim)
+
+        self.blocks = nn.ModuleList([JointBlock(dim, heads, mlp_hidden) for _ in range(depth)])
+
+        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
+        self.head = nn.Linear(dim, self.patch_dim)
+
+        self.repa_head = nn.Sequential(nn.Linear(dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, repa_dim))
+
+        hd2 = self.head_dim // 2
+        rows = torch.arange(self.grid).repeat_interleave(self.grid)
+        cols = torch.arange(self.grid).repeat(self.grid)
+        row_cos, row_sin = rope_cos_sin(rope_freqs(rows, hd2))
+        col_cos, col_sin = rope_cos_sin(rope_freqs(cols, hd2))
+        self.register_buffer("row_cos", row_cos, persistent=False)
+        self.register_buffer("row_sin", row_sin, persistent=False)
+        self.register_buffer("col_cos", col_cos, persistent=False)
+        self.register_buffer("col_sin", col_sin, persistent=False)
+        t_cos, t_sin = rope_cos_sin(rope_freqs(torch.arange(t5_len), self.head_dim))
+        self.register_buffer("t_cos", t_cos, persistent=False)
+        self.register_buffer("t_sin", t_sin, persistent=False)
+
+        self._init()
+
+    def _init(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        for b in self.blocks:
+            nn.init.zeros_(b.ada_img[-1].weight); nn.init.zeros_(b.ada_img[-1].bias)
+            nn.init.zeros_(b.ada_txt[-1].weight); nn.init.zeros_(b.ada_txt[-1].bias)
+        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
+        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)
+
+    def patchify(self, x):
+        B, C, H, W = x.shape
+        p = self.patch
+        x = x.reshape(B, C, H // p, p, W // p, p)
+        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
+        return x
+
+    def unpatchify(self, x):
+        B, N, _ = x.shape
+        p = self.patch
+        g = self.grid
+        C = self.latent_ch
+        x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5)
+        return x.reshape(B, C, g * p, g * p)
+
+    def forward(self, x, t, t5_seq, t5_mask, clip_pool, return_repa=False, use_checkpoint=False):
+        B = x.shape[0]
+        img = self.x_embed(self.patchify(x))
+        txt = self.t5_proj(t5_seq)
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.clip_proj(clip_pool)
+
+        key_valid = torch.cat([
+            torch.ones(B, img.shape[1], dtype=torch.bool, device=x.device),
+            t5_mask.bool(),
+        ], dim=1)
+
+        rope_img = (self.row_cos, self.row_sin, self.col_cos, self.col_sin)
+        rope_txt = (self.t_cos, self.t_sin)
+
+        repa_hidden = None
+        for i, blk in enumerate(self.blocks):
+            if use_checkpoint and self.training:
+                img, txt = torch.utils.checkpoint.checkpoint(
+                    blk, img, txt, c, rope_img, rope_txt, key_valid, use_reentrant=False)
+            else:
+                img, txt = blk(img, txt, c, rope_img, rope_txt, key_valid)
+            if return_repa and i == self.repa_layer:
+                repa_hidden = img
+
+        shift, scale = self.ada_out(c).chunk(2, dim=-1)
+        img = modulate(self.norm_out(img), shift, scale)
+        out = self.unpatchify(self.head(img))
+
+        if return_repa:
+            return out, self.repa_head(repa_hidden)
+        return out
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def num_backbone_params(self):
+        return sum(p.numel() for n, p in self.named_parameters() if not n.startswith("repa_head"))

eval_v6.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import AutoencoderKL
+from transformers import CLIPModel, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from dit_v6 import MMDiT
+
+def patch_clip_score():
+    def get_image_features(self, pixel_values, **kw):
+        pooled = self.vision_model(pixel_values=pixel_values).pooler_output
+        return self.visual_projection(pooled)
+
+    def get_text_features(self, input_ids=None, attention_mask=None, **kw):
+        pooled = self.text_model(input_ids=input_ids, attention_mask=attention_mask).pooler_output
+        return self.text_projection(pooled)
+
+    CLIPModel.get_image_features = get_image_features
+    CLIPModel.get_text_features = get_text_features
+
+@torch.no_grad()
+def sample(model, seq, mask, pool, null_seq, null_mask, null_pool, steps, cfg, dev):
+    B = seq.shape[0]
+    x = torch.randn(B, 4, 32, 32, device=dev)
+    ns, nm, npo = null_seq.expand(B, -1, -1), null_mask.expand(B, -1), null_pool.expand(B, -1)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=dev)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            vc = model(x, t, seq, mask, pool)
+            vu = model(x, t, ns, nm, npo)
+        x = x + (vu + cfg * (vc - vu)).float() * dt
+    return x
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--work", default="/root/v6cache")
+    ap.add_argument("--ckpt", default="/root/runs/pm6/best.pt")
+    ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--t5", default="google/flan-t5-base")
+    ap.add_argument("--n", type=int, default=5000)
+    ap.add_argument("--batch", type=int, default=50)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--cfg", type=float, nargs="+", default=[5.0])
+    ap.add_argument("--t5-len", type=int, default=32)
+    ap.add_argument("--clip-len", type=int, default=40)
+    ap.add_argument("--out", default="/root/runs/pm6/eval_results.jsonl")
+    ap.add_argument("--preview", default="")
+    args = ap.parse_args()
+    dev = "cuda"
+
+    ck = torch.load(args.ckpt, map_location=dev)
+    c = ck["cfg"]
+    model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"],
+                  t5_len=c["t5_len"]).to(dev).eval()
+    model.load_state_dict(ck["ema"])
+    print(f"[eval] loaded {args.ckpt} step {ck['step']} params {model.num_params():,}", flush=True)
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    vae_scale = vae.config.scaling_factor
+    t5_tok = T5TokenizerFast.from_pretrained(args.t5)
+    t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval()
+    clip_tok = CLIPTokenizer.from_pretrained(args.clip)
+    clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+
+    def enc(strings):
+        te = t5_tok(strings, padding="max_length", max_length=args.t5_len, truncation=True, return_tensors="pt").to(dev)
+        seq = t5(input_ids=te["input_ids"], attention_mask=te["attention_mask"]).last_hidden_state.float()
+        ce = clip_tok(strings, padding="max_length", max_length=args.clip_len, truncation=True,
+                       return_tensors="pt").to(dev)
+        pool = clip_txt(input_ids=ce["input_ids"]).pooler_output.float()
+        return seq, te["attention_mask"].float(), pool
+
+    null_seq, null_mask, null_pool = enc([""])
+
+    d = np.load(os.path.join(args.work, "eval_256.npz"), allow_pickle=True)
+    real = d["images"][:args.n]
+    caps = [str(x) for x in d["captions"][:args.n]]
+    n = len(caps)
+
+    patch_clip_score()
+    from torchmetrics.image.fid import FrechetInceptionDistance
+    from torchmetrics.multimodal.clip_score import CLIPScore
+
+    results = []
+    for cfg_val in args.cfg:
+        fid = FrechetInceptionDistance(feature=2048, normalize=True).to(dev)
+        clip_metric = CLIPScore(model_name_or_path=args.clip).to(dev)
+        for i in range(0, n, args.batch):
+            rb = torch.from_numpy(real[i:i + args.batch].astype(np.float32) / 255.0).permute(0, 3, 1, 2).to(dev)
+            fid.update(rb, real=True)
+
+        preview_imgs = []
+        for i in range(0, n, args.batch):
+            cb = caps[i:i + args.batch]
+            seq, mask, pool = enc(cb)
+            z = sample(model, seq, mask, pool, null_seq, null_mask, null_pool, args.steps, cfg_val, dev)
+            img = vae.decode((z / vae_scale).half()).sample.float()
+            img = (img.clamp(-1, 1) + 1) / 2
+            fid.update(img, real=False)
+            clip_metric.update((img * 255).to(torch.uint8), cb)
+            if args.preview and cfg_val == args.cfg[0] and len(preview_imgs) < 12:
+                for j in range(min(len(cb), 12 - len(preview_imgs))):
+                    a = (img[j].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
+                    preview_imgs.append((a, cb[j]))
+            if i % (args.batch * 10) == 0:
+                print(f"[eval] cfg={cfg_val} generated {i}/{n}", flush=True)
+
+        fid_v = float(fid.compute().item())
+        clip_v = float(clip_metric.compute().item())
+        res = {"n": n, "fid": round(fid_v, 2), "clip_score": round(clip_v, 2), "steps": args.steps,
+               "cfg": cfg_val, "render_res": 256, "fid_size": 256, "clip_model": args.clip, "step": ck["step"]}
+        results.append(res)
+        print(f"[eval] cfg={cfg_val} FID={fid_v:.2f} CLIP={clip_v:.2f} (n={n}, steps={args.steps})", flush=True)
+
+        with open(args.out, "a") as f:
+            f.write(json.dumps(res) + "\n")
+
+        if args.preview and cfg_val == args.cfg[0] and preview_imgs:
+            cell, pad = 256, 8
+            cols = 4
+            rows = (len(preview_imgs) + cols - 1) // cols
+            sheet = Image.new("RGB", (cols * cell + (cols + 1) * pad, rows * cell + (rows + 1) * pad), (245, 246, 248))
+            for k, (a, cap) in enumerate(preview_imgs):
+                r, cc = divmod(k, cols)
+                sheet.paste(Image.fromarray(a), (pad + cc * (cell + pad), pad + r * (cell + pad)))
+            sheet.save(args.preview)
+            print(f"[eval] wrote preview {args.preview}", flush=True)
+
+    print(json.dumps(results, indent=2))
+
+if __name__ == "__main__":
+    main()

main.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import argparse
+import json
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+from safetensors.torch import load_file
+from diffusers import AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from dit_v6 import MMDiT
+
+@torch.no_grad()
+def sample(model, seq, mask, pool, null_seq, null_mask, null_pool, steps, cfg, dev):
+    B = seq.shape[0]
+    x = torch.randn(B, 4, 32, 32, device=dev)
+    ns, nm, npo = null_seq.expand(B, -1, -1), null_mask.expand(B, -1), null_pool.expand(B, -1)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=dev)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            vc = model(x, t, seq, mask, pool)
+            vu = model(x, t, ns, nm, npo)
+        x = x + (vu + cfg * (vc - vu)).float() * dt
+    return x
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompt")
+    ap.add_argument("--out", default="out.png")
+    ap.add_argument("--cfg", type=float, default=5.0)
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--safetensors", default="model.safetensors")
+    ap.add_argument("--config", default="config.json")
+    ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--t5", default="google/flan-t5-base")
+    ap.add_argument("--t5-len", type=int, default=32)
+    ap.add_argument("--clip-len", type=int, default=40)
+    args = ap.parse_args()
+    dev = args.device
+
+    d = json.load(open(args.config))["dit"] if os.path.exists(args.config) else \
+        {"dim": 512, "depth": 16, "heads": 8, "mlp_hidden": 1408, "t5_len": 32}
+    model = MMDiT(dim=d["dim"], depth=d["depth"], heads=d["heads"], mlp_hidden=d["mlp_hidden"],
+                  t5_len=d["t5_len"]).to(dev).eval()
+    model.load_state_dict(load_file(args.safetensors))
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    vae_scale = vae.config.scaling_factor
+    t5_tok = T5TokenizerFast.from_pretrained(args.t5)
+    t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval()
+    clip_tok = CLIPTokenizer.from_pretrained(args.clip)
+    clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+
+    def enc(strings):
+        te = t5_tok(strings, padding="max_length", max_length=args.t5_len, truncation=True,
+                     return_tensors="pt").to(dev)
+        seq = t5(input_ids=te["input_ids"], attention_mask=te["attention_mask"]).last_hidden_state.float()
+        ce = clip_tok(strings, padding="max_length", max_length=args.clip_len, truncation=True,
+                       return_tensors="pt").to(dev)
+        pool = clip_txt(input_ids=ce["input_ids"]).pooler_output.float()
+        return seq, te["attention_mask"].float(), pool
+
+    seq, mask, pool = enc([args.prompt])
+    null_seq, null_mask, null_pool = enc([""])
+    z = sample(model, seq, mask, pool, null_seq, null_mask, null_pool, args.steps, args.cfg, dev)
+    img = vae.decode((z / vae_scale).half()).sample.float()
+    img = ((img.clamp(-1, 1) + 1) / 2)[0].permute(1, 2, 0).cpu().numpy()
+    Image.fromarray((img * 255).round().astype(np.uint8)).save(args.out)
+    print(f'[main] "{args.prompt}" -> {args.out} (cfg {args.cfg}, {args.steps} steps)')
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

model_png.json

@@ -0,0 +1 @@
+{"cfg": {"dim": 512, "depth": 16, "heads": 8, "mlp_hidden": 1408, "t5_len": 32}, "params": [{"name": "x_embed.weight", "shape": [512, 16], "numel": 8192}, {"name": "x_embed.bias", "shape": [512], "numel": 512}, {"name": "t_mlp.0.weight", "shape": [512, 512], "numel": 262144}, {"name": "t_mlp.0.bias", "shape": [512], "numel": 512}, {"name": "t_mlp.2.weight", "shape": [512, 512], "numel": 262144}, {"name": "t_mlp.2.bias", "shape": [512], "numel": 512}, {"name": "clip_proj.weight", "shape": [512, 512], "numel": 262144}, {"name": "clip_proj.bias", "shape": [512], "numel": 512}, {"name": "t5_proj.weight", "shape": [512, 768], "numel": 393216}, {"name": "t5_proj.bias", "shape": [512], "numel": 512}, {"name": "blocks.0.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.0.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.0.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.0.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.0.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.0.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.0.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.0.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.0.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.0.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.0.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.0.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.0.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.0.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.0.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.0.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.0.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.0.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.0.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.0.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.0.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.0.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.0.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.0.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.0.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.0.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.0.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.0.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.1.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.1.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.1.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.1.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.1.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.1.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.1.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.1.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.1.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.1.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.1.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.1.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.1.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.1.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.1.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.1.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.1.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.1.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.1.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.1.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.1.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.1.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.1.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.1.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.1.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.1.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.1.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.1.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.2.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.2.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.2.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.2.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.2.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.2.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.2.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.2.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.2.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.2.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.2.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.2.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.2.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.2.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.2.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.2.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.2.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.2.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.2.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.2.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.2.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.2.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.2.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.2.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.2.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.2.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.2.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.2.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.3.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.3.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.3.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.3.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.3.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.3.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.3.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.3.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.3.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.3.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.3.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.3.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.3.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.3.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.3.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.3.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.3.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.3.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.3.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.3.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.3.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.3.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.3.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.3.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.3.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.3.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.3.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.3.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.4.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.4.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.4.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.4.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.4.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.4.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.4.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.4.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.4.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.4.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.4.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.4.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.4.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.4.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.4.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.4.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.4.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.4.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.4.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.4.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.4.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.4.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.4.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.4.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.4.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.4.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.4.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.4.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.5.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.5.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.5.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.5.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.5.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.5.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.5.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.5.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.5.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.5.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.5.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.5.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.5.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.5.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.5.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.5.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.5.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.5.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.5.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.5.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.5.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.5.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.5.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.5.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.5.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.5.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.5.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.5.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.6.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.6.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.6.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.6.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.6.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.6.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.6.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.6.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.6.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.6.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.6.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.6.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.6.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.6.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.6.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.6.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.6.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.6.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.6.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.6.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.6.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.6.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.6.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.6.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.6.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.6.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.6.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.6.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.7.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.7.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.7.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.7.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.7.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.7.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.7.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.7.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.7.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.7.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.7.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.7.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.7.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.7.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.7.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.7.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.7.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.7.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.7.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.7.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.7.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.7.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.7.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.7.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.7.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.7.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.7.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.7.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.8.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.8.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.8.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.8.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.8.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.8.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.8.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.8.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.8.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.8.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.8.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.8.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.8.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.8.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.8.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.8.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.8.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.8.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.8.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.8.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.8.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.8.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.8.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.8.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.8.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.8.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.8.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.8.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.9.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.9.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.9.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.9.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.9.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.9.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.9.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.9.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.9.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.9.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.9.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.9.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.9.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.9.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.9.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.9.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.9.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.9.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.9.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.9.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.9.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.9.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.9.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.9.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.9.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.9.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.9.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.9.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.10.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.10.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.10.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.10.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.10.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.10.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.10.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.10.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.10.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.10.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.10.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.10.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.10.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.10.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.10.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.10.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.10.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.10.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.10.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.10.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.10.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.10.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.10.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.10.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.10.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.10.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.10.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.10.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.11.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.11.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.11.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.11.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.11.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.11.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.11.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.11.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.11.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.11.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.11.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.11.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.11.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.11.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.11.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.11.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.11.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.11.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.11.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.11.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.11.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.11.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.11.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.11.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.11.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.11.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.11.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.11.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.12.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.12.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.12.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.12.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.12.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.12.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.12.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.12.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.12.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.12.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.12.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.12.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.12.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.12.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.12.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.12.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.12.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.12.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.12.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.12.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.12.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.12.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.12.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.12.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.12.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.12.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.12.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.12.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.13.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.13.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.13.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.13.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.13.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.13.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.13.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.13.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.13.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.13.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.13.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.13.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.13.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.13.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.13.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.13.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.13.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.13.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.13.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.13.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.13.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.13.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.13.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.13.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.13.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.13.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.13.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.13.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.14.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.14.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.14.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.14.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.14.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.14.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.14.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.14.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.14.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.14.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.14.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.14.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.14.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.14.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.14.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.14.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.14.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.14.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.14.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.14.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.14.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.14.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.14.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.14.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.14.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.14.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.14.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.14.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.15.qkv_img.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.15.qkv_img.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.15.qkv_txt.weight", "shape": [1536, 512], "numel": 786432}, {"name": "blocks.15.qkv_txt.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.15.qn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.15.kn_img.weight", "shape": [64], "numel": 64}, {"name": "blocks.15.qn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.15.kn_txt.weight", "shape": [64], "numel": 64}, {"name": "blocks.15.proj_img.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.15.proj_img.bias", "shape": [512], "numel": 512}, {"name": "blocks.15.proj_txt.weight", "shape": [512, 512], "numel": 262144}, {"name": "blocks.15.proj_txt.bias", "shape": [512], "numel": 512}, {"name": "blocks.15.mlp_img.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.15.mlp_img.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.15.mlp_img.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.15.mlp_img.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.15.mlp_img.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.15.mlp_img.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.15.mlp_txt.gate.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.15.mlp_txt.gate.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.15.mlp_txt.up.weight", "shape": [1408, 512], "numel": 720896}, {"name": "blocks.15.mlp_txt.up.bias", "shape": [1408], "numel": 1408}, {"name": "blocks.15.mlp_txt.down.weight", "shape": [512, 1408], "numel": 720896}, {"name": "blocks.15.mlp_txt.down.bias", "shape": [512], "numel": 512}, {"name": "blocks.15.ada_img.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.15.ada_img.1.bias", "shape": [3072], "numel": 3072}, {"name": "blocks.15.ada_txt.1.weight", "shape": [3072, 512], "numel": 1572864}, {"name": "blocks.15.ada_txt.1.bias", "shape": [3072], "numel": 3072}, {"name": "ada_out.1.weight", "shape": [1024, 512], "numel": 524288}, {"name": "ada_out.1.bias", "shape": [1024], "numel": 1024}, {"name": "head.weight", "shape": [16, 512], "numel": 8192}, {"name": "head.bias", "shape": [16], "numel": 16}], "total_parameters": 155090448, "side": 12454, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"}

png_codec.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+import argparse, json, math
+import numpy as np
+import torch
+from PIL import Image
+from dit_v6 import MMDiT
+
+def encode(ckpt, png_path, cfg_path, which="ema"):
+    ck = torch.load(ckpt, map_location="cpu")
+    c = ck["cfg"]
+    model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"])
+    model.load_state_dict(ck[which])
+    parts, manifest = [], []
+    for name, p in model.named_parameters():
+        a = p.detach().to(torch.float16).contiguous().view(-1).numpy()
+        parts.append(a)
+        manifest.append({"name": name, "shape": list(p.shape), "numel": int(a.size)})
+    flat = np.concatenate(parts)
+    N = flat.size
+    side = math.ceil(math.sqrt(N))
+    u16 = flat.view(np.uint16)
+    img = np.zeros((side * side, 3), dtype=np.uint8)
+    img[:N, 0] = (u16 >> 8).astype(np.uint8)
+    img[:N, 1] = (u16 & 0xFF).astype(np.uint8)
+    Image.fromarray(img.reshape(side, side, 3), "RGB").save(png_path)
+    total = sum(m["numel"] for m in manifest)
+    with open(cfg_path, "w") as f:
+        json.dump({"cfg": c, "params": manifest, "total_parameters": total,
+                   "side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused"}, f)
+    import os
+    mb = os.path.getsize(png_path) / 1e6
+    print(f"[png] encoded {total:,} params -> {side}x{side} PNG ({mb:.1f} MB)", flush=True)
+    return total, side
+
+def load_model_png(png_path, cfg_path, device="cpu"):
+    meta = json.load(open(cfg_path))
+    c = meta["cfg"]
+    model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"])
+    arr = np.asarray(Image.open(png_path).convert("RGB")).reshape(-1, 3)
+    total = meta["total_parameters"]
+    hi = arr[:total, 0].astype(np.uint16)
+    lo = arr[:total, 1].astype(np.uint16)
+    flat = ((hi << 8) | lo).astype(np.uint16).view(np.float16)
+    sd = dict(model.named_parameters())
+    off = 0
+    with torch.no_grad():
+        for m in meta["params"]:
+            n = m["numel"]
+            chunk = flat[off:off + n].astype(np.float16)
+            t = torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32)
+            sd[m["name"]].copy_(t)
+            off += n
+    return model.to(device).eval()
+
+def decode_and_verify(png_path, cfg_path, ckpt=None, which="ema"):
+    model = load_model_png(png_path, cfg_path)
+    print(f"[png] decoded -> MMDiT with {sum(p.numel() for p in model.parameters()):,} params", flush=True)
+    if ckpt:
+        ck = torch.load(ckpt, map_location="cpu")
+        c = ck["cfg"]
+        ref = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"], t5_len=c["t5_len"])
+        ref.load_state_dict(ck[which])
+        maxdiff = 0.0
+        for (n1, p1), (n2, p2) in zip(model.named_parameters(), ref.named_parameters()):
+            maxdiff = max(maxdiff, (p1.float() - p2.half().float()).abs().max().item())
+        print(f"[png] max |decoded - original(fp16)| = {maxdiff:.2e}  (0 == lossless)", flush=True)
+    return model
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("mode", choices=["encode", "decode"])
+    ap.add_argument("--ckpt", default="ckpt/final.pt")
+    ap.add_argument("--png", default="model.png")
+    ap.add_argument("--config", default="model_png.json")
+    ap.add_argument("--which", default="ema")
+    args = ap.parse_args()
+    if args.mode == "encode":
+        encode(args.ckpt, args.png, args.config, args.which)
+    else:
+        decode_and_verify(args.png, args.config, args.ckpt, args.which)

prep_v6.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import argparse
+import io
+import os
+import tarfile
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import AutoencoderKL
+from transformers import CLIPTokenizer, T5TokenizerFast
+from huggingface_hub import hf_hub_download
+
+REPO = "undefined443/cc12m-wds-coco-recaptioned"
+
+def csr(img, size):
+    img = img.convert("RGB")
+    w, h = img.size
+    s = min(w, h)
+    l, t = (w - s) // 2, (h - s) // 2
+    return np.asarray(img.crop((l, t, l + s, t + s)).resize((size, size), Image.BICUBIC), dtype=np.uint8)
+
+def load_shard_items(tar_path, size):
+    t = tarfile.open(tar_path)
+    raw = {}
+    for m in t.getmembers():
+        if not m.isfile():
+            continue
+        key, ext = m.name.rsplit(".", 1)
+        raw.setdefault(key, {})[ext] = t.extractfile(m).read()
+    t.close()
+
+    def proc(kv):
+        _, d = kv
+        if "jpg" not in d or "txt" not in d:
+            return None
+        try:
+            arr = csr(Image.open(io.BytesIO(d["jpg"])), size)
+            cap = d["txt"].decode("utf-8", "ignore").strip()
+            if not cap:
+                return None
+            return arr, cap
+        except Exception:
+            return None
+
+    results = []
+    with ThreadPoolExecutor(max_workers=32) as pool:
+        for r in pool.map(proc, raw.items()):
+            if r is not None:
+                results.append(r)
+    return results
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--out", default="/root/v6cache/shards")
+    ap.add_argument("--tmp", default="/root/v6cache/tars")
+    ap.add_argument("--size", type=int, default=256)
+    ap.add_argument("--t5-len", type=int, default=32)
+    ap.add_argument("--clip-len", type=int, default=40)
+    ap.add_argument("--batch", type=int, default=128)
+    ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--t5", default="google/flan-t5-base")
+    ap.add_argument("--start", type=int, default=0)
+    ap.add_argument("--end", type=int, default=598)
+    ap.add_argument("--prefetch", type=int, default=2)
+    args = ap.parse_args()
+    os.makedirs(args.out, exist_ok=True)
+    os.makedirs(args.tmp, exist_ok=True)
+    dev = "cuda"
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    scale = vae.config.scaling_factor
+    print(f"[prep] vae={args.vae} scaling_factor={scale}", flush=True)
+    clip_tok = CLIPTokenizer.from_pretrained(args.clip)
+    t5_tok = T5TokenizerFast.from_pretrained(args.t5)
+
+    shard_names = [f"cc12m-coco-{i:04d}.tar" for i in range(args.start, args.end)]
+
+    def fetch(name):
+        return hf_hub_download(REPO, name, repo_type="dataset", local_dir=args.tmp)
+
+    fpool = ThreadPoolExecutor(max_workers=args.prefetch)
+    futures = {}
+    def ensure_fetch(idx):
+        if idx < len(shard_names) and idx not in futures:
+            futures[idx] = fpool.submit(fetch, shard_names[idx])
+
+    for k in range(args.prefetch):
+        ensure_fetch(k)
+
+    t0 = time.time()
+    total = 0
+    for i, name in enumerate(shard_names):
+        out_path = f"{args.out}/shard_{args.start+i:04d}.npz"
+        if os.path.exists(out_path):
+            total += np.load(out_path)["latents"].shape[0]
+            futures.pop(i, None)
+            ensure_fetch(i + args.prefetch)
+            continue
+
+        tar_path = futures.pop(i).result()
+        ensure_fetch(i + args.prefetch)
+
+        items = load_shard_items(tar_path, args.size)
+        os.remove(tar_path)
+        if not items:
+            print(f"[prep] shard {args.start+i:04d} EMPTY, skipping", flush=True)
+            continue
+        imgs = [a for a, c in items]
+        caps = [c for a, c in items]
+
+        lat_chunks = []
+        for j in range(0, len(imgs), args.batch):
+            chunk = np.stack(imgs[j:j + args.batch]).astype(np.float32) / 127.5 - 1.0
+            x = torch.from_numpy(chunk).permute(0, 3, 1, 2).to(dev).half()
+            z = vae.encode(x).latent_dist.mean * scale
+            lat_chunks.append(z.cpu().numpy().astype(np.float16))
+        latents = np.concatenate(lat_chunks)
+
+        t5o = t5_tok(caps, padding="max_length", max_length=args.t5_len, truncation=True, return_tensors="np")
+        clip_ids = clip_tok(caps, padding="max_length", max_length=args.clip_len, truncation=True,
+                             return_tensors="np")["input_ids"]
+
+        np.savez(out_path, latents=latents,
+                 t5_ids=t5o["input_ids"].astype(np.int32),
+                 t5_mask=t5o["attention_mask"].astype(np.int8),
+                 clip_ids=clip_ids.astype(np.int64))
+        total += len(imgs)
+        el = time.time() - t0
+        print(f"[prep] shard {args.start+i:04d} +{len(imgs)} total={total} "
+              f"({total/el:.1f} img/s, {el/3600:.2f}h elapsed)", flush=True)
+
+    print(f"[prep] DONE total={total}", flush=True)
+
+if __name__ == "__main__":
+    main()

preview.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import argparse
+import numpy as np
+import torch
+from PIL import Image
+from diffusers import AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from dit_v6 import MMDiT
+
+PROMPTS = [
+    "a bowl of ramen with a soft boiled egg",
+    "a red fox sitting in a snowy forest",
+    "a lighthouse on a cliff at sunset",
+    "a wooden cabin in the mountains",
+    "a cup of coffee on a wooden table",
+    "a golden retriever running on a beach",
+    "a city street at night with neon signs",
+    "a bowl of fresh strawberries",
+    "a sailboat on a calm lake",
+]
+
+@torch.no_grad()
+def sample(model, seq, mask, pool, null_seq, null_mask, null_pool, steps, cfg, dev):
+    B = seq.shape[0]
+    x = torch.randn(B, 4, 32, 32, device=dev)
+    ns, nm, npo = null_seq.expand(B, -1, -1), null_mask.expand(B, -1), null_pool.expand(B, -1)
+    dt = 1.0 / steps
+    for i in range(steps):
+        t = torch.full((B,), i * dt, device=dev)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            vc = model(x, t, seq, mask, pool)
+            vu = model(x, t, ns, nm, npo)
+        x = x + (vu + cfg * (vc - vu)).float() * dt
+    return x
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="/root/runs/pm6/best.pt")
+    ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--t5", default="google/flan-t5-base")
+    ap.add_argument("--steps", type=int, default=50)
+    ap.add_argument("--cfg", type=float, default=5.0)
+    ap.add_argument("--t5-len", type=int, default=32)
+    ap.add_argument("--clip-len", type=int, default=40)
+    ap.add_argument("--out", default="/root/preview.png")
+    args = ap.parse_args()
+    dev = "cuda"
+
+    ck = torch.load(args.ckpt, map_location=dev)
+    c = ck["cfg"]
+    model = MMDiT(dim=c["dim"], depth=c["depth"], heads=c["heads"], mlp_hidden=c["mlp_hidden"],
+                  t5_len=c["t5_len"]).to(dev).eval()
+    model.load_state_dict(ck["ema"])
+    print(f"[preview] loaded {args.ckpt} step {ck['step']}", flush=True)
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    vae_scale = vae.config.scaling_factor
+    t5_tok = T5TokenizerFast.from_pretrained(args.t5)
+    t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval()
+    clip_tok = CLIPTokenizer.from_pretrained(args.clip)
+    clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+
+    def enc(strings):
+        te = t5_tok(strings, padding="max_length", max_length=args.t5_len, truncation=True,
+                     return_tensors="pt").to(dev)
+        seq = t5(input_ids=te["input_ids"], attention_mask=te["attention_mask"]).last_hidden_state.float()
+        ce = clip_tok(strings, padding="max_length", max_length=args.clip_len, truncation=True,
+                       return_tensors="pt").to(dev)
+        pool = clip_txt(input_ids=ce["input_ids"]).pooler_output.float()
+        return seq, te["attention_mask"].float(), pool
+
+    null_seq, null_mask, null_pool = enc([""])
+
+    cell, pad, cols = 256, 8, 3
+    rows = (len(PROMPTS) + cols - 1) // cols
+    sheet = Image.new("RGB", (cols * cell + (cols + 1) * pad, rows * cell + (rows + 1) * pad), (245, 246, 248))
+
+    for i, prompt in enumerate(PROMPTS):
+        seq, mask, pool = enc([prompt])
+        z = sample(model, seq, mask, pool, null_seq, null_mask, null_pool, args.steps, args.cfg, dev)
+        img = vae.decode((z / vae_scale).half()).sample.float()
+        img = ((img.clamp(-1, 1) + 1) / 2)[0].permute(1, 2, 0).cpu().numpy()
+        a = (img * 255).round().astype(np.uint8)
+        r, cc = divmod(i, cols)
+        sheet.paste(Image.fromarray(a), (pad + cc * (cell + pad), pad + r * (cell + pad)))
+        print(f"[preview] {i+1}/{len(PROMPTS)}: {prompt}", flush=True)
+
+    sheet.save(args.out)
+    print(f"[preview] wrote {args.out}", flush=True)
+
+if __name__ == "__main__":
+    main()

publish_v6.py

@@ -0,0 +1,13 @@
+import os
+from huggingface_hub import HfApi
+
+tok = os.environ["HF_TOKEN"]
+api = HfApi(token=tok)
+url = api.upload_folder(
+    repo_id="bench-labs/PixelModel-v6",
+    repo_type="model",
+    folder_path="/root/pm6release",
+    ignore_patterns=["__pycache__/*", "*.pyc"],
+    commit_message="Publish PixelModel v6: MMDiT + REPA, 150k steps, FID 23.62 at cfg 3.0",
+)
+print("done:", url)

train_v6.py

@@ -0,0 +1,286 @@
+from __future__ import annotations
+
+import argparse
+import copy
+import glob
+import json
+import math
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from diffusers import AutoencoderKL
+from transformers import AutoModel, CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
+
+from dit_v6 import MMDiT
+
+IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
+IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
+
+def build_cache(shards_dir, cache):
+    shards = sorted(glob.glob(f"{shards_dir}/shard_*.npz"))
+    if not shards:
+        raise SystemExit(f"no shards in {shards_dir}")
+    lat, t5i, t5m, ci = [], [], [], []
+    for i, s in enumerate(shards):
+        z = np.load(s)
+        lat.append(z["latents"]); t5i.append(z["t5_ids"]); t5m.append(z["t5_mask"]); ci.append(z["clip_ids"])
+        if (i + 1) % 50 == 0:
+            print(f"  loaded {i+1}/{len(shards)} shards", flush=True)
+    lat = np.concatenate(lat); t5i = np.concatenate(t5i); t5m = np.concatenate(t5m); ci = np.concatenate(ci)
+    np.save(f"{cache}_lat.npy", lat)
+    np.save(f"{cache}_t5ids.npy", t5i)
+    np.save(f"{cache}_t5mask.npy", t5m)
+    np.save(f"{cache}_clipids.npy", ci)
+    return lat, t5i, t5m, ci
+
+def repa_weight_at(step, total, warm_frac=0.40, zero_frac=0.70, peak=0.5):
+    p = step / total
+    if p < warm_frac:
+        return peak
+    if p < zero_frac:
+        return peak * (1 - (p - warm_frac) / (zero_frac - warm_frac))
+    return 0.0
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--shards", default="/root/v6cache/shards")
+    ap.add_argument("--cache", default="/root/v6cache/cache")
+    ap.add_argument("--out", default="/root/runs/pm6")
+    ap.add_argument("--steps", type=int, default=150000)
+    ap.add_argument("--batch", type=int, default=192)
+    ap.add_argument("--lr", type=float, default=2e-4)
+    ap.add_argument("--warmup", type=int, default=1500)
+    ap.add_argument("--dim", type=int, default=512)
+    ap.add_argument("--depth", type=int, default=16)
+    ap.add_argument("--heads", type=int, default=8)
+    ap.add_argument("--mlp-hidden", type=int, default=1408)
+    ap.add_argument("--t5-len", type=int, default=32)
+    ap.add_argument("--repa-layer", type=int, default=8)
+    ap.add_argument("--repa-peak", type=float, default=0.5)
+    ap.add_argument("--repa-batch", type=int, default=48)
+    ap.add_argument("--cfg-dropout", type=float, default=0.1)
+    ap.add_argument("--ema", type=float, default=0.9999)
+    ap.add_argument("--val-size", type=int, default=4096)
+    ap.add_argument("--val-every", type=int, default=2000)
+    ap.add_argument("--log-every", type=int, default=200)
+    ap.add_argument("--save-every", type=int, default=5000)
+    ap.add_argument("--vae", default="madebyollin/sdxl-vae-fp16-fix")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--t5", default="google/flan-t5-base")
+    ap.add_argument("--dinov2", default="facebook/dinov2-small")
+    ap.add_argument("--resume", default="")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--grad-ckpt", action="store_true")
+    args = ap.parse_args()
+
+    dev = "cuda"
+    os.makedirs(args.out, exist_ok=True)
+    torch.manual_seed(args.seed)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    if os.path.exists(f"{args.cache}_lat.npy"):
+        lat = np.load(f"{args.cache}_lat.npy", mmap_mode="r")
+        t5i = np.load(f"{args.cache}_t5ids.npy")
+        t5m = np.load(f"{args.cache}_t5mask.npy")
+        ci = np.load(f"{args.cache}_clipids.npy")
+    else:
+        lat, t5i, t5m, ci = build_cache(args.shards, args.cache)
+
+    N = len(lat)
+    perm = np.random.RandomState(args.seed).permutation(N)
+    val_i = np.sort(perm[:args.val_size])
+    tr_i = perm[args.val_size:]
+    print(f"[data] {N} pairs, {len(tr_i)} train, {len(val_i)} val", flush=True)
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    vae_scale = vae.config.scaling_factor
+    for p in vae.parameters():
+        p.requires_grad_(False)
+
+    clip_txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+    for p in clip_txt.parameters():
+        p.requires_grad_(False)
+
+    t5 = T5EncoderModel.from_pretrained(args.t5).to(dev).eval()
+    for p in t5.parameters():
+        p.requires_grad_(False)
+
+    dinov2 = AutoModel.from_pretrained(args.dinov2).to(dev).eval()
+    for p in dinov2.parameters():
+        p.requires_grad_(False)
+    imagenet_mean = IMAGENET_MEAN.to(dev)
+    imagenet_std = IMAGENET_STD.to(dev)
+
+    @torch.no_grad()
+    def encode_text(t5_ids, t5_mask, clip_ids):
+        t5_out = t5(input_ids=t5_ids, attention_mask=t5_mask).last_hidden_state.float()
+        clip_pool = clip_txt(input_ids=clip_ids).pooler_output.float()
+        return t5_out, clip_pool
+
+    @torch.no_grad()
+    def dino_features(x1_latent):
+        px = vae.decode((x1_latent / vae_scale).to(vae.dtype)).sample.float()
+        px = (px.clamp(-1, 1) + 1) / 2
+        px = F.interpolate(px, size=(224, 224), mode="bilinear", align_corners=False)
+        px = (px - imagenet_mean) / imagenet_std
+        out = dinov2(pixel_values=px.to(dinov2.dtype)).last_hidden_state
+        return out[:, 1:, :].float()
+
+    t5_tok = T5TokenizerFast.from_pretrained(args.t5)
+    clip_tok = CLIPTokenizer.from_pretrained(args.clip)
+    null_t5_enc = t5_tok([""], padding="max_length", max_length=args.t5_len, truncation=True, return_tensors="pt")
+    null_t5_ids = null_t5_enc["input_ids"].to(dev)
+    null_t5_mask = null_t5_enc["attention_mask"].to(dev)
+    null_clip_ids = clip_tok([""], padding="max_length", max_length=ci.shape[1], truncation=True,
+                              return_tensors="pt")["input_ids"].to(dev)
+    null_t5_seq, null_clip_pool = encode_text(null_t5_ids, null_t5_mask, null_clip_ids)
+
+    t5i_t = torch.from_numpy(t5i.astype(np.int64))
+    t5m_t = torch.from_numpy(t5m.astype(np.int64))
+    ci_t = torch.from_numpy(ci.astype(np.int64))
+
+    vlat = torch.from_numpy(np.asarray(lat[val_i])).float()
+    vg = torch.Generator(device=dev).manual_seed(1234)
+    vx1 = vlat.to(dev)
+    vx0 = torch.randn(vx1.shape, device=dev, generator=vg)
+    vt = torch.sigmoid(torch.randn(vx1.shape[0], device=dev, generator=vg))
+    with torch.no_grad():
+        vseq, vmask, vpool = [], [], []
+        for i in range(0, len(val_i), 256):
+            s, p = encode_text(t5i_t[val_i[i:i+256]].to(dev), t5m_t[val_i[i:i+256]].to(dev), ci_t[val_i[i:i+256]].to(dev))
+            vseq.append(s); vmask.append(t5m_t[val_i[i:i+256]].to(dev)); vpool.append(p)
+        vseq = torch.cat(vseq); vmask = torch.cat(vmask); vpool = torch.cat(vpool)
+
+    model = MMDiT(dim=args.dim, depth=args.depth, heads=args.heads, mlp_hidden=args.mlp_hidden,
+                  t5_len=args.t5_len).to(dev)
+    ema = copy.deepcopy(model).eval()
+    for p in ema.parameters():
+        p.requires_grad_(False)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=0.0)
+    print(f"[model] {model.num_params()/1e6:.2f}M trainable, {model.num_backbone_params()/1e6:.2f}M backbone", flush=True)
+
+    start, best = 0, float("inf")
+    if args.resume and os.path.exists(args.resume):
+        ck = torch.load(args.resume, map_location=dev)
+        model.load_state_dict(ck["model"]); ema.load_state_dict(ck["ema"])
+        opt.load_state_dict(ck["opt"]); start = ck["step"] + 1; best = ck.get("best", best)
+        print(f"[resume] from step {start}", flush=True)
+
+    def lr_at(s):
+        if s < args.warmup:
+            return args.lr * (s + 1) / args.warmup
+        p = (s - args.warmup) / max(1, args.steps - args.warmup)
+        return args.lr * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, p))))
+
+    @torch.no_grad()
+    def val_loss():
+        tot, n = 0.0, 0
+        for j in range(0, vx1.shape[0], args.batch):
+            sl = slice(j, j + args.batch)
+            m = vx1[sl].shape[0]
+            tb = vt[sl].view(-1, 1, 1, 1)
+            xt = (1 - tb) * vx0[sl] + tb * vx1[sl]
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                v = ema(xt, vt[sl], vseq[sl], vmask[sl], vpool[sl])
+            tot += F.mse_loss(v.float(), vx1[sl] - vx0[sl]).item() * m
+            n += m
+        return tot / n
+
+    logf = open(f"{args.out}/log.jsonl", "a")
+    gen = torch.Generator(device=dev).manual_seed(args.seed)
+    run, run_diff, run_repa = 0.0, 0.0, 0.0
+    t0 = time.time()
+    for step in range(start, args.steps):
+        i = tr_i[np.random.randint(0, len(tr_i), args.batch)]
+        i_sorted = np.sort(i)
+        x1 = torch.from_numpy(np.asarray(lat[i_sorted])).to(dev).float()
+        t5_ids_b = t5i_t[i_sorted].to(dev)
+        t5_mask_b = t5m_t[i_sorted].to(dev)
+        clip_ids_b = ci_t[i_sorted].to(dev)
+        seq, pool = encode_text(t5_ids_b, t5_mask_b, clip_ids_b)
+        mask = t5_mask_b
+
+        drop = torch.rand(x1.shape[0], device=dev, generator=gen) < args.cfg_dropout
+        seq = torch.where(drop[:, None, None], null_t5_seq, seq)
+        mask = torch.where(drop[:, None], null_t5_mask, mask)
+        pool = torch.where(drop[:, None], null_clip_pool, pool)
+
+        x0 = torch.randn(x1.shape, device=dev, generator=gen)
+        t = torch.sigmoid(torch.randn(x1.shape[0], device=dev, generator=gen))
+        tb = t.view(-1, 1, 1, 1)
+        xt = (1 - tb) * x0 + tb * x1
+        target = x1 - x0
+
+        rw = repa_weight_at(step, args.steps, peak=args.repa_peak)
+
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            if rw > 0:
+                v, repa_pred = model(xt, t, seq, mask, pool, return_repa=True, use_checkpoint=args.grad_ckpt)
+            else:
+                v = model(xt, t, seq, mask, pool, use_checkpoint=args.grad_ckpt)
+            loss_diff = F.mse_loss(v.float(), target)
+            if rw > 0:
+                rb = min(args.repa_batch, x1.shape[0])
+                with torch.no_grad():
+                    dino_tgt = dino_features(x1[:rb])
+                loss_repa = 1.0 - F.cosine_similarity(repa_pred[:rb].float(), dino_tgt, dim=-1).mean()
+                loss = loss_diff + rw * loss_repa
+            else:
+                loss_repa = torch.zeros((), device=dev)
+                loss = loss_diff
+
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        d = args.ema if step > args.warmup else 0.0
+        with torch.no_grad():
+            for pe, pm in zip(ema.parameters(), model.parameters()):
+                pe.mul_(d).add_(pm.detach(), alpha=1 - d)
+            for be, bm in zip(ema.buffers(), model.buffers()):
+                be.copy_(bm)
+
+        run += loss.item(); run_diff += loss_diff.item(); run_repa += loss_repa.item()
+        if (step + 1) % args.log_every == 0:
+            el = time.time() - t0
+            sps = args.log_every / el
+            print(f"[s{step+1:06d}] loss={run/args.log_every:.4f} diff={run_diff/args.log_every:.4f} "
+                  f"repa={run_repa/args.log_every:.4f} rw={rw:.3f} lr={lr_at(step):.2e} gnorm={gn:.2f} "
+                  f"{sps:.2f} steps/s eta={(args.steps-step-1)/sps/3600:.1f}h", flush=True)
+            logf.write(json.dumps({"step": step + 1, "loss": run/args.log_every, "loss_diff": run_diff/args.log_every,
+                                    "loss_repa": run_repa/args.log_every, "repa_weight": rw,
+                                    "steps_per_s": sps}) + "\n"); logf.flush()
+            run, run_diff, run_repa, t0 = 0.0, 0.0, 0.0, time.time()
+
+        if (step + 1) % args.val_every == 0 or step + 1 == args.steps:
+            vl = val_loss()
+            tag = ""
+            if vl < best:
+                best = vl
+                torch.save({"ema": ema.state_dict(), "step": step, "val": vl,
+                            "cfg": {"dim": args.dim, "depth": args.depth, "heads": args.heads,
+                                    "mlp_hidden": args.mlp_hidden, "t5_len": args.t5_len}},
+                           f"{args.out}/best.pt")
+                tag = "  *best*"
+            print(f"[s{step+1:06d}] val_loss={vl:.5f}{tag}", flush=True)
+            logf.write(json.dumps({"step": step + 1, "val_loss": vl}) + "\n"); logf.flush()
+            t0 = time.time()
+
+        if (step + 1) % args.save_every == 0 or step + 1 == args.steps:
+            torch.save({"model": model.state_dict(), "ema": ema.state_dict(), "opt": opt.state_dict(),
+                        "step": step, "best": best,
+                        "cfg": {"dim": args.dim, "depth": args.depth, "heads": args.heads,
+                                "mlp_hidden": args.mlp_hidden, "t5_len": args.t5_len}},
+                       f"{args.out}/latest.pt")
+
+    print("TRAINDONE best_val", best, flush=True)
+
+if __name__ == "__main__":
+    main()