"""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)