nazunaex/tsubuDiT-beta

tsubuDiT beta

tsubuDiT is a hobby text-to-image project designed for training on a single consumer NVIDIA GPU. This release contains Small and Base DiT weights, their trainable text adapters, and a shared frozen text encoder.

Architecture

Component	Small	Base
Trainable parameters, including text adapter	86,437,248	164,224,128
Hidden width / attention heads	512 / 8	640 / 10
DiT blocks: input + middle + output	6 + 4 + 6	8 + 4 + 8
SwiGLU intermediate width	1,536	1,920
Trainable bidirectional text adapter	2 Transformer blocks	2 Transformer blocks

The DiT uses self-attention with 2D RoPE and QK normalization, text cross-attention in every block, RMSNorm, shared AdaLN modulation, and U-shaped long skips fused by concatenation and a linear projection. Conditioning combines timestep, global text, image size, aspect ratio, and crop information.

• Frozen text encoder: the first 12 layers of Qwen3.5-0.8B-Base, producing 1,024-dimensional hidden states with a maximum length of 96 tokens. The bundled BF16 extraction omits the original final norm and language-model output head.
• Frozen VAE: REPA-E/e2e-vavae-hf, with 16× spatial downsampling and 32 latent channels. Latents use the VAE's channel-wise mean/std normalization and 2×2 DiT patches: 512×512 pixels become 256 image tokens.
• Objective: Boundary Rectified Flow. The network predicts a residual used to assemble the velocity; training minimizes velocity MSE. Default sampling uses 28 Heun steps and CFG 4.0.

Training and checkpoints

Both models used the same one-million-image subset of the linked LAION bucketed dataset: 998,000 training images and 2,000 validation images, across 19 aspect-ratio buckets around 512-pixel resolution. Captions are model-generated English descriptions. The text encoder and VAE stayed frozen.

Model	Final epoch / optimizer step	Retained BF16 EMA snapshots
Small	30 / 467,813	5, 10, 15, 20, 25, 30 epochs
Base	15 / 233,907	5, 10, 15 epochs

Training used AdamW, an effective batch size of 64, learning rate 1e-4 with 2,344 warmup steps, and EMA decay 0.9999. Full settings, pinned dependency revisions, and the required Boundary RF mean are preserved in each checkpoint's metadata.json.

small/ and base/ each contain the latest full training checkpoint and snapshots/ with the EMA inference checkpoints above. Full checkpoints include FP32 model/EMA weights and optimizer, sampler, and RNG state. Resuming training also requires compatible training data and caches, which are not included. artifacts/qwen-truncated/ contains the shared frozen encoder and tokenizer. The VAE is downloaded separately at the revision pinned in checkpoint metadata.

Usage

The original training/inference source is not distributed with this release. The snippets below describe the loading and inference contract for developers or coding agents implementing a compatible DiT. A compatible model must be supplied using the reconstruction notes below; this is not an installed tsubuDiT package or a ready-made Diffusers pipeline. The release contains no project Python source files.

Download a snapshot and inspect its specification and state-dict keys:

import json
from pathlib import Path
import torch
from safetensors.torch import load_file
from huggingface_hub import snapshot_download

root = Path(snapshot_download(
    "YOUR_NAMESPACE/tsubuDiT-beta",
    allow_patterns=["small/snapshots/snapshot-00467813/*",
                    "artifacts/qwen-truncated/*"],
))
checkpoint = root / "small/snapshots/snapshot-00467813"
meta = json.loads((checkpoint / "metadata.json").read_text())
weights = load_file(checkpoint / "ema.safetensors")
print(meta["model_spec"])
print({key: tuple(value.shape) for key, value in weights.items()})
# After implementing the nn.Module below, match its parameter names exactly:
# model.load_state_dict(weights, strict=True)
# model = model.cuda().eval()  # FP32 parameters, BF16 autocast for the DiT

For Base, select base/snapshots/snapshot-00233907/* and its directory instead. Verify the checkpoint's integrity.json hashes before loading weights. The original environment used PyTorch 2.13, Transformers 5.15, Diffusers 0.39, and Safetensors 0.8; CUDA with BF16 is the reference execution environment.

DiT reconstruction notes

Use metadata.json["model_spec"] and the tensor keys/shapes as the module schema. All Linear and Conv2d projections have bias. RMSNorm uses FP32 statistics and epsilon 1e-6; its affine scale is enabled only in the text adapter, including its Q/K norms. DiT block and DiT Q/K norms have no affine parameters.

• Image tokens: patch_embed is Conv2d(32, d, kernel_size=2, stride=2). Flatten its grid in row-major order. Self-attention projects fused QKV, normalizes Q/K per head, then applies 2D RoPE: base 10000, unscaled integer grid coordinates, x then y halves of each head, rotating adjacent pairs. Cross-attention uses separate query and fused key/value projections, per-head Q/K normalization, and no RoPE. SDPA is noncausal; boolean mask True means a valid text token.
• Text adapter: text_conditioner.norm is affine RMSNorm(1024), followed by projection to d, two pre-norm residual self-attention/SwiGLU blocks (bidirectional, no positional embedding), and affine output_norm. Mask padding to zero. global_mlp is Linear(d,d) → SiLU → Linear(d,d), applied to the last valid adapted token. SwiGLU is output(silu(gate) * value), where fused input projects to 2×FFN width and splits as [gate, value].
• Time: concatenate cos then sin of 1000*t*exp(-log(10000)*arange(128)/128), giving 256 values, then apply time_conditioner.mlp: Linear(256,d) → SiLU → Linear(d,d).
• Size/crop: seven scalars are [clip(log2(original_w/512),-3,3), clip(log2(original_h/512),-3,3), log2(target_w/512), log2(target_h/512), crop_left/max(resized_w-target_w,1), crop_top/max(resized_h-target_h,1), log(target_w/target_h)/log(4)]. Each scalar becomes [sin(pi*s*2**j), cos(pi*s*2**j)] for j=0..31, sin values first. Flatten the resulting 7×64 features and apply micro_conditioner.mlp: Linear(448,d) → SiLU → Linear(d,d). For uncropped 512×512 generation the seven scalars are all zero, but the Fourier features are not.
• Modulation: c = time + global_text + micro. shared_modulation(silu(c)) produces [B,9,d]; add each block's modulation_offset. The order is [shift_sa, scale_sa, gate_sa, shift_ca, scale_ca, gate_ca, shift_ff, scale_ff, gate_ff]. Each branch receives RMSNorm(x)*(1+scale)+shift, then its output is added to x multiplied by its gate. Branch order: self-attention, cross-attention, SwiGLU.
• Long skips: save the output of each input block i. Immediately before output block depth-1-i, use its skip_proj(cat([current, saved_i], dim=-1)).
• Output: final_modulation(silu(c)) splits as [shift, scale]; apply final_norm(x)*(1+scale)+shift, then final_head to 128 values per token. Reshape patches as [B,grid_h,grid_w,32,2,2] and permute to [B,32,grid_h,2,grid_w,2] before flattening to the latent image. The result is a Boundary RF residual, not the velocity directly.

The required forward signature is model(latent, timestep, text_hidden, text_mask, micro_scalars) -> residual, with shapes [B,32,H/16,W/16], [B], [B,L,1024], [B,L], [B,7] respectively. Residual and latent shapes match.

Frozen models and text preparation

The bundled encoder can be reconstructed with standard Transformers classes:

from transformers import AutoTokenizer, Qwen3_5TextConfig, Qwen3_5TextModel
from diffusers import AutoencoderKL

qdir = root / "artifacts/qwen-truncated"
qmeta = json.loads((qdir / "tsubudit-qwen.json").read_text())
encoder = Qwen3_5TextModel(Qwen3_5TextConfig.from_dict(qmeta["text_config"]))
encoder.norm = torch.nn.Identity()  # Do not apply the original final RMSNorm.
encoder = encoder.to(dtype=torch.bfloat16)
encoder.load_state_dict(load_file(qdir / "model.safetensors"), strict=True)
encoder = encoder.cuda().eval().requires_grad_(False)
tokenizer = AutoTokenizer.from_pretrained(qdir / "tokenizer")
vae_lock = meta["model_lock"]["models"]["vae"]
vae = AutoencoderKL.from_pretrained(
    vae_lock["id"], revision=vae_lock["revision"], torch_dtype=torch.float32,
).cuda().eval().requires_grad_(False)

Normalize captions to NFC, replace CR/LF/TAB with spaces, remove other Unicode control characters, and collapse whitespace. Tokenize plain text with add_special_tokens=False and no chat template. Keep the first 95 tokens, then append tokenizer.eos_token_id (also for the empty unconditional prompt). Right-pad each batch to its longest sequence, using the pad ID or EOS if no pad ID exists. Pass the boolean validity mask to the encoder with use_cache=False; take last_hidden_state without an extra norm. These BF16 hidden states and masks are cond, cond_mask and uncond, uncond_mask below. The frozen encoder retains its original causal attention behavior; only the trainable text adapter is bidirectional.

Boundary RF sampling contract

Given a reconstructed, loaded CUDA model, the frozen vae, and encoded prompts, this is the reference 512×512 sampling procedure. Keep integration state and velocity assembly in FP32. Unconditional conditioning uses the empty prompt and the same size/crop scalars.

import math

@torch.inference_mode()
def sample_512(model, vae, meta, cond, cond_mask, uncond, uncond_mask,
               steps=28, cfg=4.0, seed=12345):
    batch = cond.shape[0]
    micro = torch.zeros(batch, 7, device="cuda")
    mean = torch.tensor(meta["boundary_mean"], device="cuda").view(1, 32, 1, 1)
    generator = torch.Generator(device="cuda").manual_seed(seed)
    x = torch.randn(batch, 32, 32, 32, generator=generator, device="cuda")

    def velocity(state, time):
        t = time.expand(batch)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            ru = model(state, t, uncond, uncond_mask, micro)
            rc = model(state, t, cond, cond_mask, micro)
        residual = ru.float() + cfg * (rc.float() - ru.float())
        return (torch.cos(math.pi * time / 2) * (mean - state)
                + torch.sin(math.pi * time / 2) * state
                + torch.sin(math.pi * time) * residual)

    times = torch.linspace(0, 1, steps + 1, device="cuda")
    for i in range(steps):
        t, next_t = times[i], times[i + 1]
        dt = next_t - t
        v1 = velocity(x, t)
        predicted = x + dt * v1
        # Analytic endpoint: v(x, 1) = x.
        v2 = predicted if i == steps - 1 else velocity(predicted, next_t)
        x = x + dt * (v1 + v2) / 2

    std = x.new_tensor(vae.config.latents_std).view(1, 32, 1, 1)
    vae_mean = x.new_tensor(vae.config.latents_mean).view(1, 32, 1, 1)
    pixels = vae.decode(x * std + vae_mean).sample  # No scaling_factor/shift_factor.
    return pixels.float().clamp(-1, 1).add(1).div(2)  # [B,3,512,512], RGB in [0,1]

Other supported sizes are listed in checkpoint metadata; width and height must match a trained bucket. Change the latent shape to [B,32,H/16,W/16] and compute the corresponding seven micro-conditioning scalars.

License and limitations

The custom DiT and trainable adapter weights are released under Apache-2.0. The bundled truncated Qwen encoder and tokenizer retain their upstream Apache-2.0 license and attribution. The separately downloaded VAE is MIT-licensed. See third-party licenses.

This is an experimental hobby release without standardized quality benchmarks. Outputs may contain artifacts and reflect biases in the training data. The model license does not relicense the source images, captions, or other third-party materials.