gatilin/nitrofusion2-dit-refiners

NitroFusion2 DiT Refiners

Community extraction of the two refiner stacks from NitroFusion2, revision 62d1a7000af9c46a0c4348d936bd7670ff33c1b9.

NitroFusion2 is a Z-Image-Turbo derived MMDiT. Before the 30-layer trunk runs, the model passes the noisy latent patches through a noise refiner and the caption features through a context refiner. This repository contains those two stacks and nothing else: 56 tensors, 715,746,304 parameters, in the original bfloat16.

Stack	Layers	Tensors	Parameters	Modulation
noise_refiner	2	30	361,820,672	adaLN from the timestep embedding, tanh gates
context_refiner	2	26	353,925,632	none

Each block is a single-stream transformer: RMSNorm-preceded attention with per-head RMSNorm on q/k (head dim 128), a SwiGLU feed-forward (3840 → 10240 → 3840), and post-norm on both residual branches. The noise refiner additionally carries adaLN_modulation, a Linear(256 → 15360) that produces shift and gate parameters for attention and MLP.

The 30-layer trunk (layers.*, 10.85 GB), the embedders, the final layer and the pad tokens are not included. This is a component, not a generation pipeline: it cannot produce images on its own.

Provenance

Item	Value
Source repository	ChenDY/NitroFusion2
Source revision	62d1a7000af9c46a0c4348d936bd7670ff33c1b9
Source file	transformer/diffusion_pytorch_model-00001-of-00003.safetensors
Extracted scope	noise_refiner.* and context_refiner.*
Tensor count	56
Tensor bytes	1,431,492,608
Output model.safetensors SHA-256	c871ae336b02e763a4be2018231a1b616082fd4edb04a51189b2798dcfb62284
Extraction method	HTTP range requests over two contiguous byte ranges; no full shard download

Both stacks occupy a single contiguous byte range each in the first shard, so only 1.43 GB of the 12.31 GB transformer was transferred. Every assembled range was compared byte-for-byte with the cached source ranges before the output file was hashed. Tensor values, shapes and dtypes are unchanged: no training, quantization, dtype conversion or reordering.

Original checkpoint tensor names are preserved, so load_state_dict(strict=True) works with no prefix stripping. tensor_mapping.json records the source shard, dtype, shape and byte offsets of every tensor.

Loading

The source configuration declares diffusers 0.36.0.dev0. ZImageTransformerBlock is not present in 0.35.x. Install:

pip install "diffusers==0.36.0" torch accelerate safetensors

import sys, torch
from huggingface_hub import snapshot_download

root = snapshot_download('gatilin/nitrofusion2-dit-refiners')
sys.path.insert(0, root)
from component import load_model

model = load_model(root)          # original bfloat16 weights

# x: [batch, seq, 3840] embedded noisy latent patches (seq must be a multiple of 32)
# attn_mask: [batch, seq] bool, True for real tokens
# position_ids: [seq, 3] int, from the upstream create_coordinate_grid layout
# adaln_input: [batch, 256] timestep embedding
with torch.inference_mode():
    refined = model.forward_noise(x, attn_mask, position_ids, adaln_input)

# cap_feats: [batch, cap_seq, 3840] embedded caption features
with torch.inference_mode():
    refined_cap = model.forward_context(cap_feats, attn_mask, position_ids)

load_model(root, dtype=torch.float32) runs the same weights in float32.

Input adapters are not included. forward_noise expects features already passed through the source model's all_x_embedder (Linear(64 → 3840) on 2×2 patches of 16-channel latents) and adaLN conditioning from t_embedder; forward_context expects features from cap_embedder (RMSNorm(2560) then Linear(2560 → 3840)). Take those from the source repository. The RoPE helper is included as model.rope_embedder, configured with the source axes_dims [32, 48, 48], axes_lens [1536, 512, 512] and rope_theta 256.

Position ids follow the upstream layout: caption tokens occupy (1..cap_len, 0, 0), and image tokens start at (cap_len + 1, 0, 0) — the image grid is offset by the padded caption length.

Validation

Measured on CPU with diffusers==0.36.0 and torch==2.14.1. Full results are in validation.json.

Check	Result
Strict load	56/56 tensors, no missing or unexpected keys, no prefix stripping
Parameters	715,746,304, all torch.bfloat16
noise_refiner input	three real COCO128 images → authentic Z-Image VAE → latent [1, 16, 32, 32] → 2×2 patches [256, 64]
noise_refiner output	[1, 256, 3840], all finite on all three images
context_refiner output	[1, 64, 3840], all finite
Determinism	repeated identical input, max absolute error exactly 0.0 on both stacks

The noise_refiner path uses real image data end to end. The VAE is the local Flux VAE, verified byte-identical to Tongyi-MAI/Z-Image-Turbo's vae/diffusion_pytorch_model.safetensors (SHA-256 f5b59a26851551b67ae1fe58d32e76486e1e812def4696a4bea97f16604d40a3). Latents are scaled with the source shift_factor 0.1159 and scaling_factor 0.3611, noised at flow-matching t = 0.5, then passed through the authentic all_x_embedder and t_embedder extracted from the same pinned revision. RoPE uses the upstream create_coordinate_grid layout.

Two documented assumptions. The official Qwen3 text encoder was not downloaded, so (1) context_refiner receives seeded synthetic caption features rather than real text-encoder output — they still pass through the authentic cap_embedder; and (2) the padded caption length is fixed at 64, which sets the offset of the image position ids. Neither stack is claimed to reproduce full-pipeline outputs.

Limitations

• This is a component-level extraction. Without the trunk, final layer, VAE and text encoder there is no image generation capability here.
• The refiners are designed to run before the trunk; their outputs are intermediate features, not images or final latents.
• bfloat16 CPU execution is not claimed to be bit-identical to the original GPU inference path.
• COCO128 images used for validation may overlap the model's training data; they are not an independent generalization benchmark.

License

Apache License 2.0, inherited from the source repository. LICENSE is the upstream Apache 2.0 text and NOTICE records attribution. No official source file was modified; component.py, component_config.json, extraction.json, tensor_mapping.json, validation.json and this README are added by the extractor.