Spaces:
Running on Zero
Running on Zero
| """MiniT2I: pixel-space MMDiT text-to-image backbone. | |
| Port of the released model (https://huggingface.co/MiniT2I/MiniT2I) that matches | |
| the published checkpoints key for key, so RoPE and block internals must not be | |
| "simplified". | |
| There is no AdaLN: conditioning runs entirely through the joint text+image | |
| attention stream, with frozen FLAN-T5-Large embeddings as the text stream. The | |
| time and pooled-text `vec` exist in the checkpoint but the blocks ignore them. | |
| Flow matching: x_t = t*img + (1-t)*noise with noise scaled by 2; the net predicts | |
| x0, and v = (x0 - x_t) / clamp(1-t). t=1 is clean, t=0 is noise. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| from dataclasses import dataclass, replace | |
| from typing import Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| MINIT2I_HF_REPO = "MiniT2I/MiniT2I" | |
| # --------------------------------------------------------------------------- | |
| # Primitives (verbatim from upstream mmdit.py) | |
| # --------------------------------------------------------------------------- | |
| def rotate_half(x: torch.Tensor) -> torch.Tensor: | |
| # NOTE: interleaved variant (reshape → 2 → -1), NOT the half-split variant. | |
| # The checkpoint's RoPE depends on this exact form. | |
| x1, x2 = x.reshape(*x.shape[:-1], 2, -1).unbind(dim=-2) | |
| return torch.cat((-x2, x1), dim=-1) | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float = 1e-6): | |
| super().__init__() | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| self.eps = eps | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| y = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) | |
| return y * self.weight | |
| class TimestepEmbedder(nn.Module): | |
| """Kept for state-dict compatibility; its output is unused in the forward.""" | |
| def __init__(self, hidden_size: int, frequency_embedding_size: int = 256): | |
| super().__init__() | |
| self.frequency_embedding_size = frequency_embedding_size | |
| self.mlp = nn.Sequential( | |
| nn.Linear(frequency_embedding_size, hidden_size), | |
| nn.SiLU(), | |
| nn.Linear(hidden_size, hidden_size), | |
| ) | |
| def forward(self, t: torch.Tensor) -> torch.Tensor: | |
| half = self.frequency_embedding_size // 2 | |
| freqs = torch.exp( | |
| -math.log(10000.0) | |
| * torch.arange(half, device=t.device, dtype=torch.float32) | |
| / half | |
| ) | |
| args = t.float()[:, None] * freqs[None] | |
| emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) | |
| return self.mlp(emb.to(dtype=self.mlp[0].weight.dtype)) | |
| class BottleneckPatchEmbed(nn.Module): | |
| """proj1: (in_ch -> pca) strided conv; proj2: (pca -> D) 1x1 conv.""" | |
| def __init__(self, img_size=512, patch_size=16, in_channels=3, pca_channels=128, hidden_size=768): | |
| super().__init__() | |
| self.img_size = img_size | |
| self.patch_size = patch_size | |
| self.proj1 = nn.Conv2d(in_channels, pca_channels, kernel_size=patch_size, stride=patch_size, bias=False) | |
| self.proj2 = nn.Conv2d(pca_channels, hidden_size, kernel_size=1, stride=1, bias=True) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.proj2(self.proj1(x)) | |
| return x.flatten(2).transpose(1, 2) | |
| class SwiGLUMlp(nn.Module): | |
| def __init__(self, in_features: int, hidden_features: int): | |
| super().__init__() | |
| hidden_dim = (hidden_features + 7) // 8 * 8 | |
| self.w1 = nn.Linear(in_features, hidden_dim, bias=False) | |
| self.w3 = nn.Linear(in_features, hidden_dim, bias=False) | |
| self.w2 = nn.Linear(hidden_dim, in_features, bias=False) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.w2(F.silu(self.w1(x)) * self.w3(x)) | |
| class TextRotaryEmbedding1D(nn.Module): | |
| def __init__(self, head_dim: int, theta: float = 10000.0): | |
| super().__init__() | |
| self.head_dim = head_dim | |
| self.theta = theta | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| b, length, h, d = x.shape | |
| inv = 1.0 / (self.theta ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)) | |
| pos = torch.arange(length, device=x.device, dtype=torch.float32) | |
| angles = torch.einsum("l,f->lf", pos, inv) | |
| angles = torch.cat([angles, angles], dim=-1) | |
| cos = angles.cos().to(dtype=x.dtype) | |
| sin = angles.sin().to(dtype=x.dtype) | |
| return x * cos[None, :, None, :] + rotate_half(x) * sin[None, :, None, :] | |
| class VisionRotaryEmbeddingFast(nn.Module): | |
| def __init__(self, head_dim: int, theta: float = 10000.0): | |
| super().__init__() | |
| self.dim = head_dim // 2 | |
| self.theta = theta | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| length = x.shape[1] | |
| side = int(math.sqrt(length)) | |
| if side * side != length: | |
| raise ValueError(f"image token length must be square, got {length}") | |
| freqs = 1.0 / ( | |
| self.theta | |
| ** (torch.arange(0, self.dim, 2, device=x.device, dtype=torch.float32)[: self.dim // 2] / self.dim) | |
| ) | |
| t = torch.arange(side, device=x.device, dtype=torch.float32) | |
| base = torch.einsum("l,f->lf", t, freqs) | |
| f_h, f_w = torch.broadcast_tensors(base[:, None, :], base[None, :, :]) | |
| angles = torch.cat([f_h, f_w], dim=-1) | |
| angles = torch.cat([angles, angles], dim=-1).reshape(length, -1) | |
| cos = angles.cos().to(dtype=x.dtype) | |
| sin = angles.sin().to(dtype=x.dtype) | |
| return x * cos[None, :, None, :] + rotate_half(x) * sin[None, :, None, :] | |
| class MultiModalRotaryEmbeddingFast(nn.Module): | |
| def __init__(self, head_dim: int): | |
| super().__init__() | |
| self.text_rope = TextRotaryEmbedding1D(head_dim) | |
| self.vision_rope = VisionRotaryEmbeddingFast(head_dim) | |
| def forward(self, x: torch.Tensor, txt_len: int) -> torch.Tensor: | |
| txt = self.text_rope(x[:, :txt_len]) | |
| img = self.vision_rope(x[:, txt_len:]) | |
| return torch.cat([txt, img], dim=1) | |
| # --------------------------------------------------------------------------- | |
| # Blocks | |
| # --------------------------------------------------------------------------- | |
| class PlainTextTransformerBlock(nn.Module): | |
| def __init__(self, hidden_size=768, num_heads=12, head_dim=64, mlp_ratio=2.6667): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = head_dim | |
| inner_dim = num_heads * head_dim | |
| self.norm1 = RMSNorm(hidden_size) | |
| self.norm2 = RMSNorm(hidden_size) | |
| self.qkv = nn.Linear(hidden_size, inner_dim * 3) | |
| self.attn_proj = nn.Linear(inner_dim, hidden_size) | |
| self.mlp = SwiGLUMlp(hidden_size, int(hidden_size * mlp_ratio)) | |
| self.q_norm = RMSNorm(head_dim) | |
| self.k_norm = RMSNorm(head_dim) | |
| self.rope = TextRotaryEmbedding1D(head_dim) | |
| def forward(self, txt: torch.Tensor) -> torch.Tensor: | |
| b, length, _ = txt.shape | |
| qkv = self.qkv(self.norm1(txt)).reshape(b, length, 3, self.num_heads, self.head_dim) | |
| q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] | |
| q = self.rope(self.q_norm(q)) | |
| k = self.rope(self.k_norm(k)) | |
| attn = torch.einsum("bqhd,bkhd->bhqk", q, k) * (self.head_dim ** -0.5) | |
| out = torch.einsum("bhqk,bkhd->bqhd", attn.softmax(dim=-1), v).reshape(b, length, -1) | |
| txt = txt + self.attn_proj(out) | |
| txt = txt + self.mlp(self.norm2(txt)) | |
| return txt | |
| class DoubleStreamDiTBlock(nn.Module): | |
| def __init__(self, hidden_size=768, txt_hidden_size=768, num_heads=12, head_dim=64, mlp_ratio=2.6667): | |
| super().__init__() | |
| self.hidden_size = hidden_size | |
| self.txt_hidden_size = txt_hidden_size | |
| self.num_heads = num_heads | |
| self.head_dim = head_dim | |
| inner_dim = num_heads * head_dim | |
| self.img_norm1 = RMSNorm(hidden_size) | |
| self.img_norm2 = RMSNorm(hidden_size) | |
| self.txt_norm1 = RMSNorm(txt_hidden_size) | |
| self.txt_norm2 = RMSNorm(txt_hidden_size) | |
| self.img_qkv = nn.Linear(hidden_size, inner_dim * 3) | |
| self.txt_qkv = nn.Linear(txt_hidden_size, inner_dim * 3) | |
| self.q_norm = RMSNorm(head_dim) | |
| self.k_norm = RMSNorm(head_dim) | |
| self.rope = MultiModalRotaryEmbeddingFast(head_dim) | |
| self.img_attn_proj = nn.Linear(inner_dim, hidden_size) | |
| self.txt_attn_proj = nn.Linear(inner_dim, txt_hidden_size) | |
| self.img_mlp = SwiGLUMlp(hidden_size, int(hidden_size * mlp_ratio)) | |
| self.txt_mlp = SwiGLUMlp(txt_hidden_size, int(txt_hidden_size * mlp_ratio)) | |
| def forward(self, x: torch.Tensor, txt: torch.Tensor, vec=None): | |
| b, li, _ = x.shape | |
| lt = txt.shape[1] | |
| x_norm = self.img_norm1(x) | |
| txt_norm = self.txt_norm1(txt) | |
| qkv_i = self.img_qkv(x_norm).reshape(b, li, 3, self.num_heads, self.head_dim) | |
| qkv_t = self.txt_qkv(txt_norm).reshape(b, lt, 3, self.num_heads, self.head_dim) | |
| q_i, k_i, v_i = qkv_i[:, :, 0], qkv_i[:, :, 1], qkv_i[:, :, 2] | |
| q_t, k_t, v_t = qkv_t[:, :, 0], qkv_t[:, :, 1], qkv_t[:, :, 2] | |
| q_i, k_i = self.q_norm(q_i), self.k_norm(k_i) | |
| q_t, k_t = self.q_norm(q_t), self.k_norm(k_t) | |
| q = self.rope(torch.cat([q_t, q_i], dim=1), txt_len=lt) | |
| k = self.rope(torch.cat([k_t, k_i], dim=1), txt_len=lt) | |
| v = torch.cat([v_t, v_i], dim=1) | |
| attn = torch.einsum("bqhd,bkhd->bhqk", q, k) * (self.head_dim ** -0.5) | |
| out = torch.einsum("bhqk,bkhd->bqhd", attn.softmax(dim=-1), v) | |
| x = x + self.img_attn_proj(out[:, lt:].reshape(b, li, -1)) | |
| txt = txt + self.txt_attn_proj(out[:, :lt].reshape(b, lt, -1)) | |
| x = x + self.img_mlp(self.img_norm2(x)) | |
| txt = txt + self.txt_mlp(self.txt_norm2(txt)) | |
| return x, txt | |
| class FinalLayer(nn.Module): | |
| def __init__(self, hidden_size=768, patch_size=16, out_channels=3): | |
| super().__init__() | |
| self.patch_size = patch_size | |
| self.out_channels = out_channels | |
| self.norm_final = RMSNorm(hidden_size) | |
| self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels) | |
| def forward(self, x: torch.Tensor, vec=None) -> torch.Tensor: | |
| return self.linear(self.norm_final(x)) | |
| def get_1d_sincos_pos_embed(embed_dim: int, pos: torch.Tensor) -> torch.Tensor: | |
| omega = torch.arange(embed_dim // 2, device=pos.device, dtype=torch.float32) | |
| omega = 1.0 / (10000 ** (omega / (embed_dim / 2.0))) | |
| out = torch.einsum("m,d->md", pos.reshape(-1), omega) | |
| return torch.cat([out.sin(), out.cos()], dim=1) | |
| def get_2d_sincos_pos_embed(embed_dim: int, grid_size: int, device, dtype) -> torch.Tensor: | |
| grid_h = torch.arange(grid_size, device=device, dtype=torch.float32) | |
| grid_w = torch.arange(grid_size, device=device, dtype=torch.float32) | |
| grid = torch.meshgrid(grid_w, grid_h, indexing="xy") | |
| grid = torch.stack(grid, dim=0).reshape(2, 1, grid_size, grid_size) | |
| emb_h = get_1d_sincos_pos_embed(embed_dim // 2, grid[0]) | |
| emb_w = get_1d_sincos_pos_embed(embed_dim // 2, grid[1]) | |
| return torch.cat([emb_h, emb_w], dim=1).to(dtype=dtype) | |
| # --------------------------------------------------------------------------- | |
| # Config + net | |
| # --------------------------------------------------------------------------- | |
| class MMJiTConfig: | |
| image_size: int = 512 | |
| patch_size: int = 16 | |
| in_channels: int = 3 | |
| txt_input_size: int = 1024 # FLAN-T5-Large hidden size | |
| hidden_size: int = 768 | |
| txt_hidden_size: int = 768 | |
| cond_vec_size: int = 768 | |
| depth_double: int = 17 | |
| txt_preamble_depth: int = 2 | |
| num_heads: int = 12 | |
| head_dim: int = 64 | |
| mlp_ratio: float = 2.6667 | |
| pca_channels: int = 128 | |
| prompt_length: int = 256 | |
| n_T: int = 100 | |
| prediction: str = "x" | |
| sampler: str = "euler" | |
| cfg_channels: int = 3 | |
| cfg_interval: tuple = (0.0, 1.0) | |
| llm: str = "google/flan-t5-large" | |
| class MMJiT(nn.Module): | |
| """The `net` sub-tree. State-dict matches the `model.net.<...>` checkpoint keys.""" | |
| def __init__(self, cfg: MMJiTConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.latent_img_size = cfg.image_size // cfg.patch_size | |
| self.img_embedder = BottleneckPatchEmbed( | |
| cfg.image_size, cfg.patch_size, cfg.in_channels, cfg.pca_channels, cfg.hidden_size | |
| ) | |
| self.txt_embedder = nn.Linear(cfg.txt_input_size, cfg.txt_hidden_size, bias=False) | |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, cfg.txt_input_size)) | |
| self.t_embedder = TimestepEmbedder(cfg.cond_vec_size) | |
| self.pooled_embedder = nn.Linear(cfg.txt_input_size, cfg.cond_vec_size, bias=False) | |
| self.txt_preamble_blocks = nn.ModuleList( | |
| [ | |
| PlainTextTransformerBlock(cfg.txt_hidden_size, cfg.num_heads, cfg.head_dim, cfg.mlp_ratio) | |
| for _ in range(cfg.txt_preamble_depth) | |
| ] | |
| ) | |
| self.double_blocks = nn.ModuleList( | |
| [ | |
| DoubleStreamDiTBlock( | |
| cfg.hidden_size, cfg.txt_hidden_size, cfg.num_heads, cfg.head_dim, cfg.mlp_ratio | |
| ) | |
| for _ in range(cfg.depth_double) | |
| ] | |
| ) | |
| self.final_layer = FinalLayer(cfg.hidden_size, cfg.patch_size, cfg.in_channels) | |
| def unpatchify(self, x: torch.Tensor) -> torch.Tensor: | |
| b = x.shape[0] | |
| p = self.cfg.patch_size | |
| c = self.cfg.in_channels | |
| h = w = int(math.sqrt(x.shape[1])) | |
| x = x.reshape(b, h, w, p, p, c) | |
| x = torch.einsum("nhwpqc->nchpwq", x) | |
| return x.reshape(b, c, h * p, w * p) | |
| def forward(self, img: torch.Tensor, t: torch.Tensor, context: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: | |
| if img.ndim == 4 and img.shape[1] != self.cfg.in_channels: | |
| img = img.permute(0, 3, 1, 2) | |
| attn_mask = attn_mask.to(device=context.device) | |
| context = torch.where(attn_mask[:, :, None] > 0.5, context, self.mask_token.to(dtype=context.dtype)) | |
| x = self.img_embedder(img) | |
| pos = get_2d_sincos_pos_embed(self.cfg.hidden_size, self.latent_img_size, x.device, x.dtype) | |
| x = x + pos[None] | |
| t_vec = self.t_embedder(t) | |
| txt = self.txt_embedder(context.to(dtype=self.txt_embedder.weight.dtype)) | |
| pooled_text = context.mean(dim=1) | |
| vec = t_vec + self.pooled_embedder(pooled_text.to(dtype=self.pooled_embedder.weight.dtype)) | |
| for block in self.txt_preamble_blocks: | |
| txt = block(txt) | |
| for block in self.double_blocks: | |
| x, txt = block(x, txt, vec) | |
| combined = torch.cat([txt, x], dim=1) | |
| out = self.final_layer(combined, vec) | |
| img_out = out[:, txt.shape[1]:, :] | |
| return self.unpatchify(img_out) | |
| # --------------------------------------------------------------------------- | |
| # Wrapper (mirrors upstream DiffusionModel): flow-matching + Euler/CFG sampler | |
| # --------------------------------------------------------------------------- | |
| class MiniT2I(nn.Module): | |
| """Holds `self.net = MMJiT(cfg)` under the same name the checkpoint uses. | |
| `forward` returns the x0 prediction (`predict='x'`), so the repo's existing | |
| x→velocity machinery applies. `sample()` runs the native Euler + CFG loop. | |
| """ | |
| NOISE_SCALE = 2.0 | |
| T_EPS = 0.001 | |
| def __init__(self, cfg: MMJiTConfig): | |
| super().__init__() | |
| self.cfg = cfg | |
| self.net = MMJiT(cfg) | |
| def image_size(self) -> int: | |
| return self.cfg.image_size | |
| def forward(self, img: torch.Tensor, t: torch.Tensor, context: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: | |
| return self.net(img, t, context, attn_mask) | |
| def pred_velocity(self, x, t, context, mask): | |
| x0 = self.net(x, t, context, mask) | |
| return (x0 - x) / torch.clamp(1 - t[:, None, None, None], min=self.T_EPS) | |
| def cfg_velocity(self, x, t, context, mask, cfg_scale: float): | |
| b = x.shape[0] | |
| xx = torch.cat([x, x], dim=0) | |
| tt = torch.cat([t, t], dim=0) | |
| cc = torch.cat([context, context], dim=0) | |
| mm = torch.cat([mask, torch.zeros_like(mask)], dim=0) # null branch = all-masked | |
| out = self.pred_velocity(xx, tt, cc, mm) | |
| cond, uncond = out[:b], out[b:] | |
| lo, hi = self.cfg.cfg_interval | |
| use_cfg = ((t >= lo) & (t <= hi)).to(out.dtype)[:, None, None, None] | |
| scale = torch.where(use_cfg > 0, out.new_tensor(cfg_scale), out.new_tensor(1.0)) | |
| return uncond + (cond - uncond) * scale | |
| def sample(self, context, mask, cfg_scale: float = 2.5, steps: Optional[int] = None, | |
| generator=None, progress: bool = False) -> torch.Tensor: | |
| """Euler flow-matching sampler. Returns pixels in [-1, 1].""" | |
| steps = int(steps or self.cfg.n_T) | |
| b = context.shape[0] | |
| device = context.device | |
| dtype = next(self.parameters()).dtype | |
| x = torch.randn( | |
| b, self.cfg.in_channels, self.cfg.image_size, self.cfg.image_size, | |
| generator=generator, device=device, dtype=dtype, | |
| ) * self.NOISE_SCALE | |
| timesteps = torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=dtype) | |
| iterator = range(steps) | |
| if progress: | |
| from tqdm.auto import tqdm | |
| iterator = tqdm(iterator) | |
| context = context.to(dtype) | |
| mask = mask.to(dtype) | |
| for i in iterator: | |
| t_cur = timesteps[i].expand(b) | |
| t_next = timesteps[i + 1].expand(b) | |
| if cfg_scale != 1.0: | |
| v = self.cfg_velocity(x, t_cur, context, mask, cfg_scale) | |
| else: | |
| v = self.pred_velocity(x, t_cur, context, mask) | |
| x = x + (t_next - t_cur)[:, None, None, None] * v | |
| return x.clamp(-1, 1) | |
| # --------------------------------------------------------------------------- | |
| # Registry + pretrained loader | |
| # --------------------------------------------------------------------------- | |
| # cfg_vec_size / txt_hidden_size follow hidden_size for both released variants. | |
| MiniT2I_configs = { | |
| "MiniT2I-B/16": MMJiTConfig( | |
| hidden_size=768, txt_hidden_size=768, cond_vec_size=768, | |
| depth_double=17, num_heads=12, head_dim=64, mlp_ratio=2.6667, | |
| ), | |
| "MiniT2I-L/16": MMJiTConfig( | |
| hidden_size=1248, txt_hidden_size=1248, cond_vec_size=1248, | |
| depth_double=23, num_heads=24, head_dim=52, mlp_ratio=2.7051282051282053, | |
| ), | |
| } | |
| MiniT2I_HF_SUBFOLDER = { | |
| "MiniT2I-B/16": "minit2i-b-16", | |
| "MiniT2I-L/16": "minit2i-l-16", | |
| } | |
| def _build(variant: str): | |
| base_cfg = MiniT2I_configs[variant] | |
| def ctor(**overrides) -> MiniT2I: | |
| # Accept common repo kwargs (image_size, in_channels, patch_size, ...) as | |
| # dataclass field overrides; ignore anything not in MMJiTConfig. | |
| fields = {f for f in base_cfg.__dataclass_fields__} | |
| cfg = replace(base_cfg, **{k: v for k, v in overrides.items() if k in fields}) | |
| return MiniT2I(cfg) | |
| return ctor | |
| MiniT2I_models = {name: _build(name) for name in MiniT2I_configs} | |
| def default_cfg_scale(variant: str) -> float: | |
| """Recommended classifier-free guidance scale per released variant.""" | |
| return 6.0 if "L/" in variant else 2.5 | |
| def _cache_dir(cache_dir: Optional[str]) -> Optional[str]: | |
| # L/16 is 3.65 GB — keep off the boot nvme by default (see CLAUDE.md). | |
| return cache_dir or os.environ.get("MINIT2I_CACHE") or None | |
| def load_pretrained( | |
| variant: str, | |
| model: Optional[MiniT2I] = None, | |
| *, | |
| cache_dir: Optional[str] = None, | |
| device: str | torch.device = "cpu", | |
| dtype: Optional[torch.dtype] = None, | |
| ) -> MiniT2I: | |
| """Download the released diffusers safetensors and load them into `model`. | |
| Checkpoint keys are prefixed `model.net.`; we strip that and load into | |
| `model.net` with strict=True (the unused t/pooled embedders exist in the | |
| checkpoint, so strict passes cleanly). | |
| """ | |
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_file | |
| if variant not in MiniT2I_HF_SUBFOLDER: | |
| raise KeyError(f"unknown MiniT2I variant {variant!r}; choices: {list(MiniT2I_HF_SUBFOLDER)}") | |
| if model is None: | |
| model = MiniT2I_models[variant]() | |
| sub = MiniT2I_HF_SUBFOLDER[variant] | |
| path = hf_hub_download( | |
| MINIT2I_HF_REPO, | |
| f"{sub}/transformer/diffusion_pytorch_model.safetensors", | |
| cache_dir=_cache_dir(cache_dir), | |
| ) | |
| raw = load_file(path) | |
| prefix = "model.net." | |
| net_sd = {} | |
| for k, v in raw.items(): | |
| if k.startswith(prefix): | |
| net_sd[k[len(prefix):]] = v | |
| elif k.startswith("net."): | |
| net_sd[k[len("net."):]] = v | |
| else: | |
| net_sd[k] = v | |
| missing, unexpected = model.net.load_state_dict(net_sd, strict=True) | |
| assert not missing and not unexpected, f"key mismatch: missing={missing} unexpected={unexpected}" | |
| model = model.to(device) | |
| if dtype is not None: | |
| model = model.to(dtype) | |
| return model | |
| def load_pretrained_repatched( | |
| variant: str, | |
| model: MiniT2I, | |
| *, | |
| cache_dir: Optional[str] = None, | |
| device: str | torch.device = "cpu", | |
| dtype: Optional[torch.dtype] = None, | |
| ) -> tuple[MiniT2I, list]: | |
| """Load the released $512$/p16 weights into a RE-PATCHED model (e.g. p=32 @1024). | |
| Both p16@512 and p32@1024 tokenize to a $32{\\times}32$ grid, so the entire | |
| transformer trunk (blocks, text path, rope, embedders) transfers by shape. | |
| Only the patch conv (`img_embedder.proj1`) and the output head | |
| (`final_layer.linear`) depend on patch size and are left at init --- these must | |
| be finetuned. Returns (model, list of re-initialised keys).""" | |
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_file | |
| sub = MiniT2I_HF_SUBFOLDER[variant] | |
| path = hf_hub_download( | |
| MINIT2I_HF_REPO, f"{sub}/transformer/diffusion_pytorch_model.safetensors", | |
| cache_dir=_cache_dir(cache_dir)) | |
| raw = load_file(path) | |
| net_sd = {} | |
| for k, v in raw.items(): | |
| for pre in ("model.net.", "net."): | |
| if k.startswith(pre): | |
| k = k[len(pre):]; break | |
| net_sd[k] = v | |
| tgt = model.net.state_dict() | |
| keep = {k: v for k, v in net_sd.items() if k in tgt and tgt[k].shape == v.shape} | |
| reinit = sorted(set(tgt) - set(keep)) | |
| model.net.load_state_dict(keep, strict=False) | |
| if int(os.environ.get("RANK", "0")) == 0: | |
| print(f"[repatch] transferred {len(keep)}/{len(tgt)} trunk tensors; " | |
| f"re-init {len(reinit)}: {reinit}", flush=True) | |
| model = model.to(device) | |
| if dtype is not None: | |
| model = model.to(dtype) | |
| return model, reinit | |