File size: 1,963 Bytes
8a653bd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """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) |