PixelModel v5: same architecture, 36x the data

.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,179 @@
+---
+license: mit
+pipeline_tag: text-to-image
+language:
+- en
+tags:
+- text-to-image
+- diffusion
+- rectified-flow
+- dit
+- tiny
+- custom-code
+---
+
+# PixelModel v5 🖼️
+
+<img src="examples.png" alt="Nine 256px images generated by PixelModel v5" style="width: 100%; max-width: 860px;">
+
+PixelModel v5 is a 40,013,980 parameter text-to-image diffusion transformer. It generates 256x256
+images in 50 sampling steps and it trained in about 7 hours on one A100.
+
+The architecture is identical to [v4](https://huggingface.co/bench-labs/PixelModel-v4). Same DiT,
+same width, same depth, same VAE, same text encoder, same objective. Nothing in the model changed.
+The only difference is that v5 saw 36 times more data.
+
+That was the whole experiment, and it is worth being blunt about why.
+
+## v4 was starving, not undersized
+
+v4 trained on MS-COCO, which is about 83,000 images. At batch 256 for 80,000 steps that works out to
+roughly 247 passes over the same photographs. A model that sees the same 83,000 images 247 times is
+not learning what the world looks like, it is memorising a small album.
+
+We spent a while looking at architecture changes first. Linear attention, deeper stacks, a
+higher compression autoencoder, representation alignment. All of them are real techniques and none
+of them address the actual problem, because the bottleneck was never the shape of the network.
+
+<img src="pixelmodel-v5-data.png" alt="Training set size, v4 at 83k images versus v5 at 2.99M" style="width: 100%; max-width: 780px;">
+
+v5 trains on 2,986,571 image and caption pairs, which is about 10 epochs instead of 247.
+
+## The captions matter as much as the count
+
+The images come from CC12M. Raw CC12M alt text is close to unusable for this, because it is scraped
+SEO copy. A typical caption reads like "This is the best recipe I have ever tried for Cuban bread.
+Cuban Recipes, Bread Recipes, Cooking Recipes, Pan Cubano".
+
+v5 uses a recaptioned version where a vision model rewrote every caption in COCO style:
+
+> A purple-lit subway train sits at a station platform.
+>
+> An old red tractor with gold fenders and red wheels stands in a snowy field.
+>
+> A compact kitchenette in a studio apartment features a stainless steel refrigerator, a microwave,
+> and a sink with a glass faucet.
+
+Those describe what is in the frame, which is what a text-to-image model needs. They also match the
+caption distribution that FID and CLIP score are computed against, so the training signal and the
+evaluation are speaking the same language.
+
+## Training
+
+<img src="pixelmodel-v5-loss.png" alt="Training and held-out loss over 120,000 steps" style="width: 100%; max-width: 860px;">
+
+120,000 steps at batch 256. AdamW at 2e-4 with betas (0.9, 0.99), no weight decay, gradient clipping
+at 1.0, cosine decay to a tenth of peak after a 1,000 step warmup, EMA at 0.9999 starting after
+warmup, and 10% caption dropout for classifier free guidance. It held 4.6 steps per second at
+44.2 GiB the whole way.
+
+The held-out curve climbs for the first 10,000 steps or so, which looks alarming and is not. EMA at
+0.9999 has a time constant of roughly 10,000 steps, so early on those weights are still mostly
+initialisation. The number only becomes meaningful once EMA has caught up, and from there it falls
+without interruption.
+
+<img src="pixelmodel-v5-val.png" alt="Held-out loss falling monotonically with no turn-up" style="width: 100%; max-width: 860px;">
+
+**v4 had no validation split at all**, so it shipped whichever weights the last step happened to
+produce. v5 holds out 4,096 pairs, evaluates every 2,000 steps, and writes the best checkpoint to a
+separate file. Across the whole run the held-out loss improved at every single check and never once
+turned upward, which is the clearest evidence that the data change worked. Our sibling model
+[VoxelModel v1](https://huggingface.co/bench-labs/VoxelModel-v1), trained on 28,415 meshes for 811
+epochs, peaked at step 64,000 and got 3.08% worse by the end. Same recipe, same team, different data
+regime, opposite outcome.
+
+## Results
+
+Measured on 5,000 MS-COCO val2014 images at 256 center crop, using `torchmetrics` FID at 2048
+features and `torchmetrics` CLIPScore with `openai/clip-vit-base-patch32`. This is v4's protocol.
+
+| | FID ↓ | CLIP ↑ |
+|---|---|---|
+| PixelModel v4 | 39.54 | 28.04 |
+| **PixelModel v5, cfg 5.0** | **34.06** | **29.27** |
+| PixelModel v5, cfg 3.0 | 35.03 | 28.59 |
+| PixelModel v5, cfg 6.0 | 35.18 | 29.28 |
+| real photographs | n/a | 30.05 |
+
+FID is 13.9% lower and CLIP score is 1.23 higher, at an identical parameter count. Nothing in the
+network changed between v4 and v5, so the entire difference is the training data.
+
+The last row is there because a metric without a reference point is hard to read. Feeding the real
+COCO photographs and their own captions through the same CLIPScore gives 30.05, so v5 at 29.27 is
+within 0.78 of the ceiling this metric can award. That number also served as a sanity check on the
+harness, since `torchmetrics` CLIPScore is incompatible with transformers 5.x and needed patching
+before it would run at all.
+
+Guidance scale was swept rather than assumed. FID bottoms out at cfg 5.0 and CLIP keeps climbing
+slightly to cfg 6.0, which is the usual tension between the two. Anyone reproducing this should
+also know that FID is heavily biased at small sample sizes. The same checkpoint at cfg 5.0 scores
+74.35 on 1,000 images and 34.06 on 5,000, so sample count has to match before two FID numbers can be
+compared at all.
+
+One caveat worth stating plainly. v4's cached evaluation set is not published in its repository, so
+the reference images here were rebuilt from COCO val2014 with a fixed seed. The protocol is the same
+but the exact 5,000 photographs may not be. CLIP score is unaffected by this since it compares each
+generated image against its own caption, but the FID comparison carries that asterisk.
+
+## Specification
+
+| | |
+|---|---|
+| trainable parameters | 40,013,980 |
+| frozen VAE, `sd-vae-ft-mse` | 83,653,863 |
+| frozen text encoder, CLIP ViT-B/32 | 63,165,952 |
+| **combined total** | **186,833,795** |
+| architecture | DiT, dim 384, depth 12, heads 6, patch 2 |
+| tokens | 256, over a 32x32x4 latent |
+| objective | rectified flow, target `x1 - x0`, logit-normal timesteps |
+| conditioning | cross attention on 40 CLIP text tokens, plus adaLN-zero |
+| resolution | 256x256 |
+| sampling | 50 steps, cfg 5.0 for best FID, cfg 6.0 for best CLIP |
+
+A note on that combined number, because it is higher than what v4 published. v4's card reports
+161,500,000 total. That figure counts the CLIP text transformer but leaves out its 49,408 by 512
+token embedding table, which is another 25.3M parameters. The encoder does not run without it, so it
+counts. The correct total for both v4 and v5 is 186,833,795.
+
+## What it does well and what it does not
+
+After a few hundred samples the boundary is consistent.
+
+It is good at material and light. Food, landscapes, skies, water, snow, foliage, stone, interiors,
+depth of field. A bowl of ramen, a wet cobblestone street, a wheat field under storm light, a library
+of wooden shelves. These come back looking like photographs.
+
+It is bad at structure. Faces and animal heads deform. Limbs merge. Wings on aircraft are stubby or
+fused to the fuselage. Any text it renders, on a shop sign or an airliner livery, is convincing
+looking gibberish. Prompts asking for one isolated object in an empty scene often return the scene
+without the object, so "a lighthouse on a cliff at sunset" gives a very nice sunset over a cliff and
+no lighthouse.
+
+Frequency in the training data matters more than complexity. Aircraft are structurally complicated
+and come out fine because CC12M is full of them. Hot air balloons are simpler and fail, because they
+are rare.
+
+That failure profile is what a 40M model looks like once data is no longer the constraint. It has
+learned appearance thoroughly and topology poorly, and topology is what capacity buys. v6 is where
+that gets spent.
+
+## Usage
+
+```bash
+python main.py "a bowl of ramen with a soft boiled egg" --out ramen.png --cfg 6.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` | combined weights, backbone and VAE and text encoder |
+| `model.png` | the same backbone weights as an image |
+| `dit.py` | the model |
+| `train_v5.py` | training loop with the held-out split |
+| `prep_v5.py` | streams the dataset, encodes latents, caches captions |
+| `main.py` | inference |

config.json

@@ -0,0 +1,52 @@
+{
+  "architecture": "latent-diffusion-dit",
+  "objective": "rectified-flow",
+  "dit": {
+    "dim": 384,
+    "depth": 12,
+    "heads": 6,
+    "patch": 2,
+    "latent_ch": 4,
+    "latent_size": 32,
+    "text_dim": 512,
+    "mlp_ratio": 4.0
+  },
+  "vae": "stabilityai/sd-vae-ft-mse",
+  "text_encoder": "openai/clip-vit-base-patch32",
+  "max_tokens": 40,
+  "resolution": 256,
+  "trainable_parameters": 40013980,
+  "vae_parameters": 83653863,
+  "text_encoder_parameters": 63165952,
+  "total_parameters_incl_frozen": 186833795,
+  "training": {
+    "steps": 120000,
+    "batch_size": 256,
+    "pairs": 2986571,
+    "val_pairs": 4096,
+    "optimizer": "adamw",
+    "lr": 0.0002,
+    "betas": [
+      0.9,
+      0.99
+    ],
+    "warmup": 1000,
+    "ema": 0.9999,
+    "cfg_dropout": 0.1,
+    "best_val_loss": 0.7852968573570251,
+    "dataset": "undefined443/cc12m-wds-coco-recaptioned"
+  },
+  "sampling": {
+    "steps": 50,
+    "cfg": 5.0
+  },
+  "eval": {
+    "fid": 34.06,
+    "clip_score": 29.27,
+    "n": 5000,
+    "cfg": 5.0,
+    "dataset": "MS-COCO val2014, 256 center crop",
+    "protocol": "torchmetrics FrechetInceptionDistance + CLIPScore (openai/clip-vit-base-patch32)",
+    "real_image_clip_score": 30.05
+  }
+}

dit.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+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 sincos_2d(dim, grid_size):
+    g = np.arange(grid_size, dtype=np.float32)
+    gx, gy = np.meshgrid(g, g, indexing="xy")
+    assert dim % 4 == 0
+    d4 = dim // 4
+    omega = 1.0 / (10000 ** (np.arange(d4, dtype=np.float32) / d4))
+    def emb1(p):
+        out = p.reshape(-1)[:, None] * omega[None]
+        return np.concatenate([np.sin(out), np.cos(out)], axis=1)
+    pe = np.concatenate([emb1(gx), emb1(gy)], axis=1)
+    return torch.from_numpy(pe).float()
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads):
+        super().__init__()
+        self.heads = heads
+        self.q = nn.Linear(dim, dim)
+        self.kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, ctx=None):
+        ctx = x if ctx is None else ctx
+        B, N, C = x.shape
+        M = ctx.shape[1]
+        h = self.heads
+        q = self.q(x).reshape(B, N, h, C // h).transpose(1, 2)
+        kv = self.kv(ctx).reshape(B, M, 2, h, C // h).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+        o = F.scaled_dot_product_attention(q, k, v)
+        o = o.transpose(1, 2).reshape(B, N, C)
+        return self.proj(o)
+
+class Block(nn.Module):
+    def __init__(self, dim, heads, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.attn = Attention(dim, heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross = Attention(dim, heads)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"),
+                                 nn.Linear(hidden, dim))
+        self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
+        self.cross_gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, c, text):
+        shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(c).chunk(6, dim=1)
+        x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1))
+        x = x + self.cross_gate * self.cross(self.norm_ca(x), text)
+        x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift2, scale2))
+        return x
+
+class DiT(nn.Module):
+    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=384, depth=12,
+                 heads=6, text_dim=512, mlp_ratio=4.0):
+        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.x_embed = nn.Linear(self.patch_dim, dim)
+        self.register_buffer("pos", sincos_2d(dim, self.grid).unsqueeze(0))
+        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.text_proj = nn.Linear(text_dim, dim)
+        self.text_pool = nn.Linear(text_dim, dim)
+        self.blocks = nn.ModuleList([Block(dim, heads, mlp_ratio) 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.dim = dim
+        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[-1].weight); nn.init.zeros_(b.ada[-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, text_seq, text_pool):
+        x = self.x_embed(self.patchify(x)) + self.pos
+        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.text_pool(text_pool)
+        text = self.text_proj(text_seq)
+        for blk in self.blocks:
+            x = blk(x, c, text)
+        shift, scale = self.ada_out(c).chunk(2, dim=1)
+        x = modulate(self.norm_out(x), shift, scale)
+        x = self.head(x)
+        return self.unpatchify(x)
+
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())

eval_v5.py

@@ -0,0 +1,108 @@
+import argparse, json, os, sys, numpy as np, torch
+from diffusers import AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPModel
+sys.path.insert(0, '/root')
+from dit import DiT
+
+SCALE = 0.18215
+
+def _img_feats(self, pixel_values=None, **kw):
+    return self.visual_projection(self.vision_model(pixel_values=pixel_values).pooler_output)
+
+def _txt_feats(self, input_ids=None, attention_mask=None, **kw):
+    return self.text_projection(self.text_model(input_ids=input_ids,
+                                                attention_mask=attention_mask).pooler_output)
+
+CLIPModel.get_image_features = _img_feats
+CLIPModel.get_text_features = _txt_feats
+
+@torch.no_grad()
+def sample(model, seq, pool, ns, npool, steps, cfg, dev):
+    B = seq.shape[0]
+    x = torch.randn(B, 4, 32, 32, device=dev)
+    a, b = ns.expand(B, -1, -1), npool.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, pool)
+            vu = model(x, t, a, b)
+        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/pm5/best.pt')
+    ap.add_argument('--data', default='/root/pm5eval/eval_256.npz')
+    ap.add_argument('--n', type=int, default=5000)
+    ap.add_argument('--batch', type=int, default=100)
+    ap.add_argument('--steps', type=int, default=50)
+    ap.add_argument('--cfgs', default='2,3,4,5,6')
+    ap.add_argument('--out', default='/root/pm5_eval.json')
+    args = ap.parse_args()
+    dev = 'cuda'
+
+    ck = torch.load(args.ckpt, map_location=dev)
+    model = DiT(dim=384, depth=12, heads=6).to(dev).eval()
+    model.load_state_dict(ck['ema'])
+    print(f"[eval] ckpt step {ck.get('step')} val {ck.get('val')}", flush=True)
+
+    vae = AutoencoderKL.from_pretrained('stabilityai/sd-vae-ft-mse').to(dev).half().eval()
+    tok = CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32')
+    txt = CLIPTextModel.from_pretrained('openai/clip-vit-base-patch32').to(dev).half().eval()
+    e = tok([''], padding='max_length', max_length=40, truncation=True, return_tensors='pt').to(dev)
+    o = txt(**e)
+    ns, npool = o.last_hidden_state.float(), o.pooler_output.float()
+
+    d = np.load(args.data, allow_pickle=True)
+    real = d['images'][:args.n]
+    caps = [str(x) for x in d['captions'][:args.n]]
+    n = len(caps)
+    print(f'[eval] {n} real images and captions', flush=True)
+
+    from torchmetrics.image.fid import FrechetInceptionDistance
+    from torchmetrics.multimodal.clip_score import CLIPScore
+
+    ref = CLIPScore(model_name_or_path='openai/clip-vit-base-patch32').to(dev)
+    for i in range(0, n, args.batch):
+        rb = torch.from_numpy(real[i:i + args.batch]).permute(0, 3, 1, 2).to(dev)
+        ref.update(rb, caps[i:i + args.batch])
+    real_clip = float(ref.compute().item())
+    print(f'[eval] REAL images CLIP score {real_clip:.2f} (metric sanity check / ceiling)', flush=True)
+    del ref
+    torch.cuda.empty_cache()
+
+    results = []
+    for cfg in [float(x) for x in args.cfgs.split(',')]:
+        fid = FrechetInceptionDistance(feature=2048, normalize=True).to(dev)
+        clip = CLIPScore(model_name_or_path='openai/clip-vit-base-patch32').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)
+        for i in range(0, n, args.batch):
+            cb = caps[i:i + args.batch]
+            t = tok(cb, padding='max_length', max_length=40, truncation=True, return_tensors='pt').to(dev)
+            oo = txt(**t)
+            z = sample(model, oo.last_hidden_state.float(), oo.pooler_output.float(),
+                       ns, npool, args.steps, cfg, dev)
+            img = vae.decode((z / SCALE).half()).sample.float()
+            img = (img.clamp(-1, 1) + 1) / 2
+            fid.update(img, real=False)
+            clip.update((img * 255).to(torch.uint8), cb)
+        f = float(fid.compute().item()); c = float(clip.compute().item())
+        results.append({'cfg': cfg, 'fid': round(f, 2), 'clip_score': round(c, 2),
+                        'n': n, 'steps': args.steps})
+        print(f'[eval] cfg {cfg}: FID {f:.2f}  CLIP {c:.2f}', flush=True)
+        del fid, clip
+        torch.cuda.empty_cache()
+
+    best = min(results, key=lambda r: r['fid'])
+    json.dump({'ckpt_step': ck.get('step'), 'val_loss': ck.get('val'),
+               'real_image_clip_score': round(real_clip, 2),
+               'results': results, 'best_fid': best}, open(args.out, 'w'), indent=2)
+    print(f"[eval] BEST FID {best['fid']} at cfg {best['cfg']} (CLIP {best['clip_score']})", flush=True)
+    print('EVALDONE', flush=True)
+
+if __name__ == '__main__':
+    main()

main.py

@@ -0,0 +1,89 @@
+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
+
+from dit import DiT
+
+SCALE = 0.18215
+
+@torch.no_grad()
+def sample(model, seq, pool, null_seq, null_pool, steps, cfg, dev, seed=None):
+    B = seq.shape[0]
+    g = None
+    if seed is not None:
+        g = torch.Generator(device=dev).manual_seed(seed)
+    x = torch.randn(B, 4, 32, 32, device=dev, generator=g)
+    ns, npool = null_seq.expand(B, -1, -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, pool)
+            vu = model(x, t, ns, npool)
+        x = x + (vu + cfg * (vc - vu)).float() * dt
+    return x
+
+@torch.no_grad()
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("prompts", nargs="+")
+    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("--seed", type=int, default=None)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--weights", default="model.safetensors")
+    ap.add_argument("--config", default="config.json")
+    ap.add_argument("--vae", default="stabilityai/sd-vae-ft-mse")
+    ap.add_argument("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--max-tokens", 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": 384, "depth": 12, "heads": 6}
+    model = DiT(dim=d["dim"], depth=d["depth"], heads=d["heads"]).to(dev).eval()
+    sd = load_file(args.weights)
+    model.load_state_dict({k[len("dit."):]: v for k, v in sd.items() if k.startswith("dit.")})
+
+    vae = AutoencoderKL.from_pretrained(args.vae).to(dev).half().eval()
+    tok = CLIPTokenizer.from_pretrained(args.clip)
+    txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+
+    def enc(strings):
+        t = tok(strings, padding="max_length", max_length=args.max_tokens,
+                truncation=True, return_tensors="pt").to(dev)
+        o = txt(**t)
+        return o.last_hidden_state.float(), o.pooler_output.float()
+
+    seq, pool = enc(args.prompts)
+    null_seq, null_pool = enc([""])
+    z = sample(model, seq, pool, null_seq, null_pool, args.steps, args.cfg, dev, args.seed)
+    img = vae.decode((z / SCALE).half()).sample.float()
+    img = ((img.clamp(-1, 1) + 1) / 2).permute(0, 2, 3, 1).cpu().numpy()
+
+    n = len(args.prompts)
+    if n == 1:
+        Image.fromarray((img[0] * 255).round().astype(np.uint8)).save(args.out)
+        paths = [args.out]
+    else:
+        root, ext = os.path.splitext(args.out)
+        paths = []
+        for i in range(n):
+            p = f"{root}_{i}{ext}"
+            Image.fromarray((img[i] * 255).round().astype(np.uint8)).save(p)
+            paths.append(p)
+    for p, s in zip(paths, args.prompts):
+        print(f'[pixelmodel] "{s}" -> {p} (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:927864ec43cc1534849ed27c9be86d6f5a59b4fe13fd451806df4b7a66939178
+size 747804492

model_png.json

@@ -0,0 +1 @@
+{"cfg": {"dim": 384, "depth": 12, "heads": 6}, "params": [{"name": "x_embed.weight", "shape": [384, 16], "numel": 6144}, {"name": "x_embed.bias", "shape": [384], "numel": 384}, {"name": "t_mlp.0.weight", "shape": [384, 384], "numel": 147456}, {"name": "t_mlp.0.bias", "shape": [384], "numel": 384}, {"name": "t_mlp.2.weight", "shape": [384, 384], "numel": 147456}, {"name": "t_mlp.2.bias", "shape": [384], "numel": 384}, {"name": "text_proj.weight", "shape": [384, 512], "numel": 196608}, {"name": "text_proj.bias", "shape": [384], "numel": 384}, {"name": "text_pool.weight", "shape": [384, 512], "numel": 196608}, {"name": "text_pool.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.0.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.0.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.0.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.0.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.0.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.0.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.0.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.0.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.0.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.0.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.0.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.1.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.1.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.1.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.1.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.1.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.1.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.1.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.1.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.1.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.1.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.1.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.1.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.2.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.2.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.2.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.2.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.2.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.2.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.2.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.2.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.2.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.2.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.2.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.2.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.3.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.3.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.3.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.3.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.3.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.3.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.3.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.3.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.3.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.3.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.3.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.3.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.4.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.4.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.4.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.4.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.4.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.4.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.4.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.4.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.4.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.4.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.4.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.4.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.5.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.5.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.5.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.5.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.5.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.5.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.5.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.5.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.5.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.5.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.5.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.5.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.6.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.6.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.6.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.6.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.6.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.6.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.6.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.6.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.6.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.6.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.6.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.6.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.7.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.7.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.7.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.7.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.7.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.7.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.7.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.7.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.7.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.7.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.7.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.7.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.8.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.8.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.8.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.8.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.8.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.8.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.8.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.8.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.8.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.8.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.8.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.8.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.9.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.9.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.9.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.9.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.9.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.9.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.9.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.9.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.9.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.9.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.9.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.9.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.10.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.10.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.10.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.10.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.10.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.10.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.10.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.10.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.10.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.10.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.10.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.10.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "blocks.11.cross_gate", "shape": [1], "numel": 1}, {"name": "blocks.11.attn.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.attn.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.attn.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.11.attn.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.11.attn.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.attn.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.cross.q.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.cross.q.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.cross.kv.weight", "shape": [768, 384], "numel": 294912}, {"name": "blocks.11.cross.kv.bias", "shape": [768], "numel": 768}, {"name": "blocks.11.cross.proj.weight", "shape": [384, 384], "numel": 147456}, {"name": "blocks.11.cross.proj.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.mlp.0.weight", "shape": [1536, 384], "numel": 589824}, {"name": "blocks.11.mlp.0.bias", "shape": [1536], "numel": 1536}, {"name": "blocks.11.mlp.2.weight", "shape": [384, 1536], "numel": 589824}, {"name": "blocks.11.mlp.2.bias", "shape": [384], "numel": 384}, {"name": "blocks.11.ada.1.weight", "shape": [2304, 384], "numel": 884736}, {"name": "blocks.11.ada.1.bias", "shape": [2304], "numel": 2304}, {"name": "ada_out.1.weight", "shape": [768, 384], "numel": 294912}, {"name": "ada_out.1.bias", "shape": [768], "numel": 768}, {"name": "head.weight", "shape": [16, 384], "numel": 6144}, {"name": "head.bias", "shape": [16], "numel": 16}], "total_parameters": 40013980, "side": 6326, "dtype": "float16", "channels": "R=hi,G=lo,B=unused", "step": 119999, "val_loss": 0.7852968573570251}

png_codec.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+
+import numpy as np
+import torch
+from PIL import Image
+
+from dit import DiT
+
+Image.MAX_IMAGE_PIXELS = None
+
+def build(cfg):
+    return DiT(dim=cfg["dim"], depth=cfg["depth"], heads=cfg["heads"])
+
+def encode(ckpt, png_path, cfg_path, dim=384, depth=12, heads=6, which="ema"):
+    ck = torch.load(ckpt, map_location="cpu")
+    cfg = {"dim": dim, "depth": depth, "heads": heads}
+    model = build(cfg)
+    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)
+    json.dump({"cfg": cfg, "params": manifest, "total_parameters": int(N),
+               "side": side, "dtype": "float16", "channels": "R=hi,G=lo,B=unused",
+               "step": ck.get("step"), "val_loss": ck.get("val")},
+              open(cfg_path, "w"))
+    print(f"[png] encoded {N:,} params into a {side}x{side} PNG "
+          f"({os.path.getsize(png_path)/1e6:.1f} MB)")
+    return N, side
+
+def load_model_png(png_path, cfg_path, device="cpu"):
+    meta = json.load(open(cfg_path))
+    model = build(meta["cfg"])
+    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)
+            sd[m["name"]].copy_(torch.from_numpy(chunk.copy()).view(*m["shape"]).to(torch.float32))
+            off += n
+    return model.to(device).eval()
+
+def verify(ckpt, png_path, cfg_path, which="ema"):
+    model = load_model_png(png_path, cfg_path)
+    ref = torch.load(ckpt, map_location="cpu")[which]
+    worst, name = 0.0, ""
+    for k, p in model.named_parameters():
+        d = (p.detach() - ref[k].float()).abs().max().item()
+        if d > worst:
+            worst, name = d, k
+    scale = max(1e-12, ref[name].float().abs().max().item())
+    print(f"[png] round-trip max abs error {worst:.3e} on {name} (relative {worst/scale:.2e})")
+    return worst
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="/root/runs/pm5/best.pt")
+    ap.add_argument("--png", default="model.png")
+    ap.add_argument("--manifest", default="model_png.json")
+    args = ap.parse_args()
+    encode(args.ckpt, args.png, args.manifest)
+    verify(args.ckpt, args.png, args.manifest)
+
+if __name__ == "__main__":
+    main()

prep_v5.py

@@ -0,0 +1,111 @@
+import os, sys, io, gc, time, tarfile, argparse, numpy as np, torch
+from concurrent.futures import ThreadPoolExecutor
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from diffusers import AutoencoderKL
+from transformers import CLIPTokenizer
+
+REPO = 'undefined443/cc12m-wds-coco-recaptioned'
+OUT = '/root/v5data'
+TMP = '/root/v5tmp'
+SCALE = 0.18215
+RES = 256
+MAXTOK = 40
+
+ap = argparse.ArgumentParser()
+ap.add_argument('--start', type=int, default=0)
+ap.add_argument('--end', type=int, default=598)
+ap.add_argument('--batch', type=int, default=64)
+args = ap.parse_args()
+
+os.makedirs(OUT, exist_ok=True)
+os.makedirs(TMP, exist_ok=True)
+dev = 'cuda'
+vae = AutoencoderKL.from_pretrained('stabilityai/sd-vae-ft-mse').to(dev).half().eval()
+for p in vae.parameters():
+    p.requires_grad_(False)
+tok = CLIPTokenizer.from_pretrained('openai/clip-vit-base-patch32')
+
+def prep_img(raw):
+    try:
+        im = Image.open(io.BytesIO(raw)).convert('RGB')
+    except Exception:
+        return None
+    w, h = im.size
+    s = RES / min(w, h)
+    im = im.resize((max(RES, int(w * s + 0.5)), max(RES, int(h * s + 0.5))), Image.BICUBIC)
+    w, h = im.size
+    l, t = (w - RES) // 2, (h - RES) // 2
+    return np.asarray(im.crop((l, t, l + RES, t + RES)), dtype=np.uint8)
+
+pool = ThreadPoolExecutor(16)
+t_start = time.time()
+done_total = 0
+
+def fetch(si):
+    return hf_hub_download(REPO, f'cc12m-coco-{si:04d}.tar', repo_type='dataset', local_dir=TMP)
+
+todo = [i for i in range(args.start, args.end) if not os.path.exists(f'{OUT}/shard_{i:04d}.npz')]
+print(f'{len(todo)} shards to process', flush=True)
+dl = ThreadPoolExecutor(2)
+nxt = dl.submit(fetch, todo[0]) if todo else None
+
+for k, si in enumerate(todo):
+    dst = f'{OUT}/shard_{si:04d}.npz'
+    try:
+        p = nxt.result()
+    except Exception as e:
+        print(f'  shard {si} download failed: {type(e).__name__}', flush=True)
+        nxt = dl.submit(fetch, todo[k + 1]) if k + 1 < len(todo) else None
+        continue
+    nxt = dl.submit(fetch, todo[k + 1]) if k + 1 < len(todo) else None
+
+    pairs = {}
+    try:
+        with tarfile.open(p) as tf:
+            for m in tf:
+                if not m.isfile():
+                    continue
+                key, _, ext = m.name.partition('.')
+                if ext == 'jpg':
+                    pairs.setdefault(key, {})['img'] = tf.extractfile(m).read()
+                elif ext == 'txt':
+                    pairs.setdefault(key, {})['cap'] = tf.extractfile(m).read().decode('utf-8', 'ignore').strip()
+    except Exception as e:
+        print(f'  shard {si} tar failed: {type(e).__name__}', flush=True)
+        os.remove(p); continue
+    os.remove(p)
+
+    items = [(v['img'], v['cap']) for v in pairs.values() if 'img' in v and 'cap' in v and v['cap']]
+    del pairs
+    if not items:
+        continue
+
+    lat_all, tok_all = [], []
+    for i in range(0, len(items), args.batch):
+        chunk = items[i:i + args.batch]
+        arrs = list(pool.map(prep_img, [c[0] for c in chunk]))
+        keep = [j for j, a in enumerate(arrs) if a is not None]
+        if not keep:
+            continue
+        x = np.stack([arrs[j] for j in keep])
+        t = torch.from_numpy(x).to(dev).half().permute(0, 3, 1, 2).div_(127.5).sub_(1.0)
+        with torch.no_grad():
+            lat = vae.encode(t).latent_dist.sample() * SCALE
+        lat_all.append(lat.cpu().numpy().astype(np.float16))
+        ids = tok([chunk[j][1] for j in keep], padding='max_length', max_length=MAXTOK,
+                  truncation=True, return_tensors='np')['input_ids']
+        tok_all.append(ids.astype(np.int32))
+
+    if not lat_all:
+        continue
+    lat = np.concatenate(lat_all); ids = np.concatenate(tok_all)
+    np.savez(dst, latents=lat, tokens=ids)
+    done_total += len(lat)
+    el = time.time() - t_start
+    print(f'  shard {si:04d}: {len(lat)} imgs  total {done_total}  {done_total/el:.0f} img/s  '
+          f'{el/60:.1f} min elapsed', flush=True)
+    del lat_all, tok_all, lat, ids, items
+    gc.collect()
+
+print('PREPDONE', done_total, flush=True)

train_v5.py

@@ -0,0 +1,193 @@
+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 transformers import CLIPTextModel
+
+from dit import DiT
+
+def build_cache(data, cache):
+    shards = sorted(glob.glob(f"{data}/shard_*.npz"))
+    if not shards:
+        raise SystemExit(f"no shards in {data}")
+    lat, tok = [], []
+    for i, s in enumerate(shards):
+        z = np.load(s)
+        lat.append(z["latents"]); tok.append(z["tokens"])
+        if (i + 1) % 50 == 0:
+            print(f"  loaded {i+1}/{len(shards)} shards", flush=True)
+    lat = np.concatenate(lat); tok = np.concatenate(tok)
+    np.save(f"{cache}_lat.npy", lat); np.save(f"{cache}_tok.npy", tok)
+    return lat, tok
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default="/root/v5data")
+    ap.add_argument("--cache", default="/root/v5cache")
+    ap.add_argument("--out", default="/root/runs/pm5")
+    ap.add_argument("--steps", type=int, default=120000)
+    ap.add_argument("--batch", type=int, default=256)
+    ap.add_argument("--lr", type=float, default=2e-4)
+    ap.add_argument("--warmup", type=int, default=1000)
+    ap.add_argument("--dim", type=int, default=384)
+    ap.add_argument("--depth", type=int, default=12)
+    ap.add_argument("--heads", type=int, default=6)
+    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("--clip", default="openai/clip-vit-base-patch32")
+    ap.add_argument("--resume", default="")
+    ap.add_argument("--seed", type=int, default=0)
+    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")
+        tok = np.load(f"{args.cache}_tok.npy")
+    else:
+        lat, tok = build_cache(args.data, 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)
+
+    txt = CLIPTextModel.from_pretrained(args.clip).to(dev).eval()
+    for p in txt.parameters():
+        p.requires_grad_(False)
+
+    @torch.no_grad()
+    def encode(ids):
+        o = txt(input_ids=ids)
+        return o.last_hidden_state.float(), o.pooler_output.float()
+
+    null_ids = torch.full((1, tok.shape[1]), 0, dtype=torch.long, device=dev)
+    null_ids[0, 0] = 49406; null_ids[0, 1:] = 49407
+    null_seq, null_pool = encode(null_ids)
+
+    tok_t = torch.from_numpy(tok.astype(np.int64))
+    vlat = torch.from_numpy(np.asarray(lat[val_i])).float()
+    vseq, vpool = [], []
+    with torch.no_grad():
+        for i in range(0, len(val_i), 512):
+            s, p = encode(tok_t[val_i[i:i+512]].to(dev))
+            vseq.append(s); vpool.append(p)
+    vseq = torch.cat(vseq); vpool = torch.cat(vpool)
+    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))
+
+    model = DiT(dim=args.dim, depth=args.depth, heads=args.heads).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] {sum(p.numel() for p in model.parameters())/1e6:.2f}M trainable", 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)
+
+    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], 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, t0 = 0.0, time.time()
+    for step in range(start, args.steps):
+        i = tr_i[np.random.randint(0, len(tr_i), args.batch)]
+        x1 = torch.from_numpy(np.asarray(lat[np.sort(i)])).to(dev).float()
+        seq, pool = encode(tok_t[np.sort(i)].to(dev))
+        drop = torch.rand(x1.shape[0], device=dev, generator=gen) < args.cfg_dropout
+        seq = torch.where(drop[:, None, None], null_seq, seq)
+        pool = torch.where(drop[:, None], null_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
+
+        for g in opt.param_groups:
+            g["lr"] = lr_at(step)
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            v = model(xt, t, seq, pool)
+            loss = F.mse_loss(v.float(), target)
+        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()
+        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} lr={lr_at(step):.2e} "
+                  f"gnorm={gn:.2f} {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,
+                                   "steps_per_s": sps}) + "\n"); logf.flush()
+            run, t0 = 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}, 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}, f"{args.out}/latest.pt")
+
+    print("TRAINDONE best_val", best, flush=True)
+
+if __name__ == "__main__":
+    main()