| """Load the NitroFusion2 DiT refiner stacks using the official diffusers Z-Image blocks.""" |
| from pathlib import Path |
| import json |
| import torch |
| from torch import nn |
| from accelerate import init_empty_weights |
| from safetensors.torch import load_file |
| from diffusers.models.transformers.transformer_z_image import ZImageTransformerBlock, RopeEmbedder |
|
|
|
|
| class NitroFusion2Refiners(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| dim = config['dim'] |
| args = (dim, config['n_heads'], config['n_kv_heads'], config['norm_eps'], config['qk_norm']) |
| n = config['n_refiner_layers'] |
| self.noise_refiner = nn.ModuleList([ |
| ZImageTransformerBlock(1000 + i, *args, modulation=True) for i in range(n)]) |
| self.context_refiner = nn.ModuleList([ |
| ZImageTransformerBlock(i, *args, modulation=False) for i in range(n)]) |
| self.rope_embedder = RopeEmbedder( |
| theta=config['rope_theta'], axes_dims=config['axes_dims'], axes_lens=config['axes_lens']) |
|
|
| def forward_noise(self, x, attn_mask, position_ids, adaln_input): |
| freqs_cis = self.rope_embedder(position_ids).unsqueeze(0) |
| for layer in self.noise_refiner: |
| x = layer(x, attn_mask, freqs_cis, adaln_input) |
| return x |
|
|
| def forward_context(self, cap_feats, attn_mask, position_ids): |
| freqs_cis = self.rope_embedder(position_ids).unsqueeze(0) |
| for layer in self.context_refiner: |
| cap_feats = layer(cap_feats, attn_mask, freqs_cis) |
| return cap_feats |
|
|
|
|
| def load_model(root, dtype=None): |
| root = Path(root) |
| config = json.loads((root / 'config.json').read_text()) |
| with init_empty_weights(include_buffers=False): |
| model = NitroFusion2Refiners(config) |
| state = load_file(str(root / 'model.safetensors'), device='cpu') |
| model.load_state_dict(state, strict=True, assign=True) |
| model = model.eval() |
| return model if dtype is None else model.to(dtype) |
