Upload transformer_chroma.py with huggingface_hub

transformer_chroma.py

@@ -1,30 +1,19 @@
 from __future__ import annotations
 
+import gc
 import json
-import math
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Optional
+from warnings import warn
 
 import torch
+import torch.nn.functional as F
 from diffusers.configuration_utils import ConfigMixin
 from diffusers.models.modeling_utils import ModelMixin
 from torch import nn
 
 
-# Keep the C++ additive-attention exception log one-shot to avoid
-# repeating the same fallback message for every transformer block.
-_CPP_ADDITIVE_ATTN_EXCEPTION_LOGGED = False
-
-
-def _log_cpp_additive_attn_exception(reason: str):
-    global _CPP_ADDITIVE_ATTN_EXCEPTION_LOGGED
-    if _CPP_ADDITIVE_ATTN_EXCEPTION_LOGGED:
-        return
-    _CPP_ADDITIVE_ATTN_EXCEPTION_LOGGED = True
-    print(f"[nunchaku.chroma] cpp_additive_attn fallback: {reason}")
-
-
 @dataclass(frozen=True)
 class LoadReport:
     config: dict[str, Any]
@@ -32,198 +21,131 @@ class LoadReport:
     rank: int
 
 
-def _maybe_log(verbose: bool, *args):
-    if verbose:
-        print(*args)
+# ---------------------------------------------------------------------------
+# Checkpoint key conversion
+# ---------------------------------------------------------------------------
 
+_KEY_RENAMES = [
+    (".lora_down", ".proj_down"),
+    (".lora_up", ".proj_up"),
+    (".smooth_orig", ".smooth_factor_orig"),
+    (".smooth", ".smooth_factor"),
+]
 
-def _load_safetensors_state_dict(
-    path: str | Path,
-    *,
-    device: str,
-) -> tuple[dict[str, Any], dict[str, str]]:
-    from nunchaku.utils import load_state_dict_in_safetensors
 
-    sd, md = load_state_dict_in_safetensors(path, device=device, return_metadata=True)
-    return sd, md
-
-
-def _convert_checkpoint_key_for_svdq_linear(k: str) -> str:
-    # This safetensors uses nunchaku converter naming:
-    # - lora_down/lora_up are actually SVD proj_down/proj_up
-    # - smooth/smooth_orig are smooth_factor/smooth_factor_orig
-    if ".lora_down" in k:
-        k = k.replace(".lora_down", ".proj_down")
-    if ".lora_up" in k:
-        k = k.replace(".lora_up", ".proj_up")
-    if ".smooth_orig" in k:
-        k = k.replace(".smooth_orig", ".smooth_factor_orig")
-    elif ".smooth" in k:
-        k = k.replace(".smooth", ".smooth_factor")
+def _convert_checkpoint_key(k: str) -> str:
+    """Rename nunchaku-converter keys to runtime parameter names."""
+    for old, new in _KEY_RENAMES:
+        if old in k:
+            return k.replace(old, new)
     return k
 
 
 def _convert_checkpoint_state_dict(sd: dict[str, Any]) -> dict[str, Any]:
-    return {_convert_checkpoint_key_for_svdq_linear(k): v for k, v in sd.items()}
+    return {_convert_checkpoint_key(k): v for k, v in sd.items()}
 
 
 def _infer_rank_from_converted_state_dict(sd: dict[str, Any]) -> int:
-    # Look for any SVDQW4A4Linear proj_down shaped [in_features, rank]
     for k, v in sd.items():
         if k.endswith(".proj_down") and getattr(v, "ndim", None) == 2:
             return int(v.shape[1])
     raise ValueError("Cannot infer SVD rank from checkpoint (missing any '*.proj_down' tensors).")
 
 
-def _build_attn_norms(*, head_dim: int, eps: float, with_added: bool, device, dtype):
+# ---------------------------------------------------------------------------
+# SVDQ linear helper
+# ---------------------------------------------------------------------------
+
+
+def _make_svdq_linear(in_features: int, out_features: int, *, rank: int, precision: str, device, dtype):
+    from nunchaku.models.linear import SVDQW4A4Linear
+
+    return SVDQW4A4Linear(
+        in_features=in_features,
+        out_features=out_features,
+        rank=rank,
+        bias=True,
+        precision=precision,
+        torch_dtype=dtype,
+        device=device,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Attention norms helper
+# ---------------------------------------------------------------------------
+
+
+def _build_attn_norms(*, head_dim: int, eps: float, with_added: bool, device, dtype) -> nn.ModuleDict:
     from diffusers.models.normalization import RMSNorm
 
-    m = nn.Module()
-    m.norm_q = RMSNorm(head_dim, eps=eps, elementwise_affine=True).to(device=device, dtype=dtype)
-    m.norm_k = RMSNorm(head_dim, eps=eps, elementwise_affine=True).to(device=device, dtype=dtype)
+    def _rms():
+        return RMSNorm(head_dim, eps=eps, elementwise_affine=True).to(device=device, dtype=dtype)
+
+    d: dict[str, nn.Module] = {"norm_q": _rms(), "norm_k": _rms()}
     if with_added:
-        m.norm_added_q = RMSNorm(head_dim, eps=eps, elementwise_affine=True).to(device=device, dtype=dtype)
-        m.norm_added_k = RMSNorm(head_dim, eps=eps, elementwise_affine=True).to(device=device, dtype=dtype)
-    return m
-
-
-def _should_use_cpp_additive_attn(*, attention_mask_1d, hidden_states, head_dim: int) -> bool:
-    return attention_mask_1d is not None and hidden_states.is_cuda and int(head_dim) == 128
-
-
-def _pad_to_multiple(n: int, multiple: int) -> int:
-    return int(math.ceil(n / multiple) * multiple)
-
-
-def _get_or_create_cpp_workspace(
-    owner,
-    cache_attr: str,
-    *,
-    batch_size: int,
-    num_tokens_pad: int,
-    heads: int,
-    head_dim: int,
-    device,
-    out_dtype,
-):
-    key = (batch_size, num_tokens_pad, heads, head_dim, str(device), out_dtype)
-    ws = getattr(owner, cache_attr, None)
-    if ws is None or ws.get("key") != key:
-        ws = {
-            "key": key,
-            "q": torch.empty((batch_size, heads, num_tokens_pad, head_dim), device=device, dtype=torch.float16),
-            "k": torch.empty((batch_size, heads, num_tokens_pad, head_dim), device=device, dtype=torch.float16),
-            "v": torch.empty((batch_size, heads, num_tokens_pad, head_dim), device=device, dtype=torch.float16),
-            "m": torch.empty((batch_size, num_tokens_pad), device=device, dtype=torch.float16),
-            "out": torch.empty((batch_size, num_tokens_pad, heads * head_dim), device=device, dtype=out_dtype),
-        }
-        setattr(owner, cache_attr, ws)
-    return ws
-
-
-def _get_cpp_workspace_tensors(cpp_workspace: dict):
-    return (
-        cpp_workspace["q"],
-        cpp_workspace["k"],
-        cpp_workspace["v"],
-        cpp_workspace["m"],
-        cpp_workspace["out"],
-    )
+        d["norm_added_q"] = _rms()
+        d["norm_added_k"] = _rms()
+    return nn.ModuleDict(d)
+
+
+# ---------------------------------------------------------------------------
+# Pure-Python QKV + Norm + Rotary helper
+# ---------------------------------------------------------------------------
+
 
+def _apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply rotary positional embedding.
 
-def _run_cpp_additive_attention(q, k, v, mask, out, *, context: str) -> bool:
-    try:
-        from nunchaku._C.ops import chroma_additive_attention_packed_fp16
+    Uses the same convention as diffusers' ``apply_rotary_emb`` from
+    ``diffusers.models.embeddings``.
 
-        chroma_additive_attention_packed_fp16(q, k, v, mask, out, 0.0)
-        return True
-    except Exception as e:
-        _log_cpp_additive_attn_exception(f"exception in {context}: {type(e).__name__}: {e}")
-        return False
+    Args:
+        x: (B, S, H, D) query or key tensor.
+        freqs_cis: (S, D/2, 2) or broadcastable rotary embedding (sin, cos stacked).
+    """
+    from diffusers.models.embeddings import apply_rotary_emb
+    return apply_rotary_emb(x, freqs_cis, sequence_dim=1)
 
 
-def _fused_qkv_heads(hidden_states, qkv_proj, norm_q, norm_k, rotary_emb, heads: int):
-    from nunchaku.ops.fused import fused_qkv_norm_rottary
+def _qkv_norm_rotary(hidden_states, qkv_proj, norm_q, norm_k, rotary_emb, heads: int):
+    """Pure-Python replacement for fused_qkv_norm_rottary.
 
-    qkv = fused_qkv_norm_rottary(hidden_states, qkv_proj, norm_q, norm_k, rotary_emb)
+    Steps: QKV projection → chunk → RMSNorm(Q, K) → reshape to multi-head → rotary.
+    """
+    qkv = qkv_proj(hidden_states)
     query, key, value = qkv.chunk(3, dim=-1)
-    return tuple(x.unflatten(-1, (heads, -1)) for x in (query, key, value))
-
-
-def _expand_batch_dim(x: torch.Tensor, batch_size: int) -> torch.Tensor:
-    if batch_size != int(x.shape[0]):
-        x = x.expand(batch_size, -1, -1).contiguous()
-    return x
-
-
-def _prepare_cpp_context(owner, hidden_states, attention_mask, *, txt_tokens: int, img_tokens: int):
-    heads = int(owner.config.num_attention_heads)
-    head_dim = int(owner.config.attention_head_dim)
-    batch_size = int(hidden_states.shape[0])
-    device = hidden_states.device
-    out_dtype = hidden_states.dtype
-
-    pad_size = 256
-    txt_pad = _pad_to_multiple(txt_tokens, pad_size)
-    img_pad = _pad_to_multiple(img_tokens, pad_size)
-    s_total = int(txt_tokens + img_tokens)
-    s_pad = _pad_to_multiple(s_total, pad_size)
-
-    ws_dual = _get_or_create_cpp_workspace(
-        owner,
-        "_nunchaku_cpp_ws_dual_shared",
-        batch_size=batch_size,
-        num_tokens_pad=txt_pad + img_pad,
-        heads=heads,
-        head_dim=head_dim,
-        device=device,
-        out_dtype=out_dtype,
-    )
-    ws_single = _get_or_create_cpp_workspace(
-        owner,
-        "_nunchaku_cpp_ws_single_shared",
-        batch_size=batch_size,
-        num_tokens_pad=s_pad,
-        heads=heads,
-        head_dim=head_dim,
-        device=device,
-        out_dtype=out_dtype,
-    )
 
-    attn_mask_fp16 = attention_mask.to(dtype=torch.float16)
+    # (B, S, dim) → (B, S, H, D)
+    query = query.unflatten(-1, (heads, -1))
+    key = key.unflatten(-1, (heads, -1))
+    value = value.unflatten(-1, (heads, -1))
+
+    query = norm_q(query)
+    key = norm_k(key)
+
+    if rotary_emb is not None:
+        query = _apply_rotary_emb(query, rotary_emb)
+        key = _apply_rotary_emb(key, rotary_emb)
 
-    mask_single = ws_single["m"]
-    mask_single.zero_()
-    mask_single[:, :s_total] = attn_mask_fp16
+    return query, key, value
 
-    mask_dual = ws_dual["m"]
-    mask_dual.zero_()
-    mask_dual[:, :txt_tokens] = attn_mask_fp16[:, :txt_tokens]
-    mask_dual[:, txt_pad : txt_pad + img_tokens] = attn_mask_fp16[:, txt_tokens : txt_tokens + img_tokens]
 
-    return ws_dual, ws_single, mask_dual, mask_single
+# ---------------------------------------------------------------------------
+# Attention dispatch
+# ---------------------------------------------------------------------------
 
 
 def _dispatch_attention(query, key, value, attention_mask):
-    """
-    Chroma attention dispatch.
+    """Chroma attention dispatch (pure-Python path).
 
-    Performance note:
-    This function must NOT call `.item()` on CUDA tensors (it would introduce a device sync per block).
+    ``query``/``key``/``value`` are (B, S, H, D).
     """
-    from diffusers.models.transformers.transformer_flux import dispatch_attention_fn
+    from diffusers.models.attention_dispatch import dispatch_attention_fn
 
-    # No mask: allow fastest backend selection (FLASH where available).
     if attention_mask is None:
         return dispatch_attention_fn(query, key, value, attn_mask=None, backend=None)
 
-    # Speed + quality path (Chroma-specific):
-    # The Chroma pipeline provides a 2D mask `m` (values in {0,1}, dtype usually bf16/fp16), which diffusers expands
-    # to a rank-1 outer-product bias `m_i * m_j` and passes as an additive SDPA mask.
-    # This is *not* a boolean hard-mask, but an additive bias in SDPA.
-    #
-    # We can fold this outer-product bias into the QK dot-product by augmenting Q/K with extra dims, and then run
-    # fast attention with attn_mask=None while preserving semantics closely.
     if attention_mask.ndim == 2 and query.shape[0] == 1:
         b, s = attention_mask.shape
         if b != 1:
@@ -233,313 +155,117 @@ def _dispatch_attention(query, key, value, attention_mask):
                 f"Mask/sequence length mismatch: mask S={int(s)}, query S={int(query.shape[1])}, key S={int(key.shape[1])}"
             )
 
-        # Expand to (B,S,H,1) and keep dtype aligned with Q/K.
         m1 = attention_mask.to(dtype=query.dtype)[:, :, None, None].expand(
             query.shape[0], query.shape[1], query.shape[2], 1
         )
-
         d = int(query.shape[-1])
         scale = float(d) ** -0.5
-
-        # Keep extra dims minimal but aligned (multiple of 8) to reduce overhead.
-        extra = 8
         sqrt_d = float(d) ** 0.5
+        extra = 8
 
-        q_extra = torch.cat([m1 * sqrt_d, m1.new_zeros((*m1.shape[:-1], extra - 1))], dim=-1)
-        k_extra = torch.cat([m1, m1.new_zeros((*m1.shape[:-1], extra - 1))], dim=-1)
-        v_extra = value.new_zeros((*value.shape[:-1], extra))
-
-        q_ext = torch.cat([query, q_extra], dim=-1)
-        k_ext = torch.cat([key, k_extra], dim=-1)
-        v_ext = torch.cat([value, v_extra], dim=-1)
+        zeros_qk = m1.new_zeros((*m1.shape[:-1], extra - 1))
+        q_ext = torch.cat([query, m1 * sqrt_d, zeros_qk], dim=-1)
+        k_ext = torch.cat([key, m1, zeros_qk], dim=-1)
+        v_ext = torch.cat([value, value.new_zeros((*value.shape[:-1], extra))], dim=-1)
 
-        # Prefer native flash kernel when available; pass explicit scale to preserve original head_dim scaling.
         try:
             out_ext = dispatch_attention_fn(q_ext, k_ext, v_ext, attn_mask=None, backend="_native_flash", scale=scale)
         except TypeError:
-            # Older diffusers may not expose `scale` in dispatch_attention_fn; fallback to correctness baseline.
             out_ext = None
         if out_ext is not None:
             return out_ext[..., :d]
 
-    # Fallback: preserve diffusers Chroma mask semantics (outer-product additive bias) and use SDPA efficient.
-    attn_mask_4d = NunchakuChromaTransformerBlockMixin._mask_to_4d(attention_mask)
+    attn_mask_4d = _mask_to_4d(attention_mask)
     return dispatch_attention_fn(query, key, value, attn_mask=attn_mask_4d, backend="_native_efficient")
 
 
-class NunchakuChromaTransformerBlockMixin:
-    @staticmethod
-    def _mask_to_4d(attention_mask):
-        # Match diffusers `transformer_chroma` behavior:
-        # Expand a 2D mask to a full QK mask (outer product).
-        #
-        # IMPORTANT: do NOT cast to bool here. Chroma's pipeline may provide a non-bool mask (e.g. bf16 0/1),
-        # and changing dtype/value semantics affects output quality.
-        if attention_mask is None:
-            return None
-        if attention_mask.ndim == 4:
-            return attention_mask
-        if attention_mask.ndim != 2:
-            raise ValueError(f"Unsupported attention_mask shape: {tuple(attention_mask.shape)}")
-        return attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
-
-
-class NunchakuChromaSingleTransformerBlock(nn.Module, NunchakuChromaTransformerBlockMixin):
-    """
-    Matches the checkpoint key layout under:
-      single_transformer_blocks.<i>.{qkv_proj,out_proj,mlp_fc1,mlp_fc2,attn.norm_{q,k},norm,proj_out?}
-    """
+def _mask_to_4d(attention_mask):
+    """Expand a 2D mask to a full QK mask (outer product), matching diffusers Chroma semantics."""
+    if attention_mask is None:
+        return None
+    if attention_mask.ndim == 4:
+        return attention_mask
+    if attention_mask.ndim != 2:
+        raise ValueError(f"Unsupported attention_mask shape: {tuple(attention_mask.shape)}")
+    return attention_mask[:, None, None, :] * attention_mask[:, None, :, None]
 
-    def __init__(
-        self,
-        *,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        mlp_ratio: float,
-        rank: int,
-        precision: str,
-        device,
-        dtype,
-        eps: float = 1e-6,
-    ):
+
+# ===========================================================================
+# Transformer blocks
+# ===========================================================================
+
+
+class NunchakuChromaSingleTransformerBlock(nn.Module):
+    def __init__(self, *, dim: int, num_attention_heads: int, attention_head_dim: int, mlp_ratio: float,
+                 rank: int, precision: str, device, dtype, eps: float = 1e-6):
         super().__init__()
         from diffusers.models.transformers.transformer_chroma import ChromaAdaLayerNormZeroSinglePruned
-        from nunchaku.models.linear import SVDQW4A4Linear
 
-        self.heads = int(num_attention_heads)
-        self.head_dim = int(attention_head_dim)
-        self.inner_dim = int(dim)
-        self.mlp_hidden_dim = int(dim * mlp_ratio)
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
+        mlp_hidden_dim = int(dim * mlp_ratio)
 
         self.norm = ChromaAdaLayerNormZeroSinglePruned(dim).to(device=device, dtype=dtype)
         self.attn = _build_attn_norms(head_dim=self.head_dim, eps=eps, with_added=False, device=device, dtype=dtype)
 
-        self.qkv_proj = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=3 * dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.out_proj = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.mlp_fc1 = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=self.mlp_hidden_dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.mlp_fc2 = SVDQW4A4Linear(
-            in_features=self.mlp_hidden_dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-
-        self.norm_q = self.attn.norm_q
-        self.norm_k = self.attn.norm_k
-
-    def forward(
-        self,
-        hidden_states,
-        temb,
-        image_rotary_emb=None,
-        attention_mask_1d=None,
-        cpp_workspace: dict | None = None,
-        cpp_mask: torch.Tensor | None = None,
-    ):
-        from nunchaku.ops.fused import fused_gelu_mlp, fused_qkv_norm_rottary
+        lin = dict(rank=rank, precision=precision, device=device, dtype=dtype)
+        self.qkv_proj = _make_svdq_linear(dim, 3 * dim, **lin)
+        self.out_proj = _make_svdq_linear(dim, dim, **lin)
+        self.mlp_fc1 = _make_svdq_linear(dim, mlp_hidden_dim, **lin)
+        self.mlp_fc2 = _make_svdq_linear(mlp_hidden_dim, dim, **lin)
 
+    def forward(self, hidden_states, temb, image_rotary_emb=None, attention_mask_1d=None):
         residual = hidden_states
         norm_hidden_states, gate = self.norm(hidden_states, emb=temb)
 
-        mlp_out = fused_gelu_mlp(norm_hidden_states, self.mlp_fc1, self.mlp_fc2)
+        mlp_hidden = self.mlp_fc1(norm_hidden_states)
+        mlp_hidden = F.gelu(mlp_hidden, approximate="tanh")
+        mlp_out = self.mlp_fc2(mlp_hidden)
 
-        # Optional C++/CUDA additive attention backend (exact Chroma semantics, B=1 only).
-        use_cpp = _should_use_cpp_additive_attn(
-            attention_mask_1d=attention_mask_1d,
-            hidden_states=norm_hidden_states,
-            head_dim=self.head_dim,
+        query, key, value = _qkv_norm_rotary(
+            norm_hidden_states, self.qkv_proj, self.attn["norm_q"], self.attn["norm_k"],
+            image_rotary_emb, self.heads,
         )
-        if use_cpp:
-            assert cpp_workspace is not None and cpp_mask is not None
-            _, s, _ = norm_hidden_states.shape
-            q, k, v, _, out = _get_cpp_workspace_tensors(cpp_workspace)
-            fused_qkv_norm_rottary(
-                norm_hidden_states,
-                self.qkv_proj,
-                self.attn.norm_q,
-                self.attn.norm_k,
-                image_rotary_emb,
-                output=(q, k, v),
-                attn_tokens=int(s),
-            )
-            if _run_cpp_additive_attention(q, k, v, cpp_mask, out, context="single-block cpp path"):
-                attn_out = out[:, :s, :]
-            else:
-                use_cpp = False
-
-        if not use_cpp:
-            query, key, value = _fused_qkv_heads(
-                norm_hidden_states, self.qkv_proj, self.attn.norm_q, self.attn.norm_k, image_rotary_emb, self.heads
-            )
-            attn_out = _dispatch_attention(query, key, value, attention_mask_1d)
-            attn_out = attn_out.flatten(2, 3).to(query.dtype)
+        attn_out = _dispatch_attention(query, key, value, attention_mask_1d)
+        attn_out = attn_out.flatten(2, 3).to(query.dtype)
 
-        proj = self.out_proj(attn_out) + mlp_out
-        hidden_states = residual + gate.unsqueeze(1) * proj
+        hidden_states = residual + gate.unsqueeze(1) * (self.out_proj(attn_out) + mlp_out)
 
         if hidden_states.dtype == torch.float16:
             hidden_states = hidden_states.clip(-65504, 65504)
         return hidden_states
 
 
-class NunchakuChromaTransformerBlock(nn.Module, NunchakuChromaTransformerBlockMixin):
-    """
-    Matches the checkpoint key layout under:
-      transformer_blocks.<i>.{qkv_proj,qkv_proj_context,out_proj,out_proj_context,mlp_fc1,mlp_fc2,mlp_context_fc1,mlp_context_fc2,attn.*}
-    """
-
-    def __init__(
-        self,
-        *,
-        dim: int,
-        num_attention_heads: int,
-        attention_head_dim: int,
-        rank: int,
-        precision: str,
-        device,
-        dtype,
-        eps: float = 1e-6,
-    ):
+class NunchakuChromaTransformerBlock(nn.Module):
+    def __init__(self, *, dim: int, num_attention_heads: int, attention_head_dim: int,
+                 rank: int, precision: str, device, dtype, eps: float = 1e-6):
         super().__init__()
         from diffusers.models.transformers.transformer_chroma import ChromaAdaLayerNormZeroPruned
-        from nunchaku.models.linear import SVDQW4A4Linear
 
-        self.heads = int(num_attention_heads)
-        self.head_dim = int(attention_head_dim)
-        self.inner_dim = int(dim)
+        self.heads = num_attention_heads
+        self.head_dim = attention_head_dim
 
         self.norm1 = ChromaAdaLayerNormZeroPruned(dim).to(device=device, dtype=dtype)
         self.norm1_context = ChromaAdaLayerNormZeroPruned(dim).to(device=device, dtype=dtype)
-
         self.attn = _build_attn_norms(head_dim=self.head_dim, eps=eps, with_added=True, device=device, dtype=dtype)
 
-        self.qkv_proj = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=3 * dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.qkv_proj_context = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=3 * dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.out_proj = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.out_proj_context = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
+        lin = dict(rank=rank, precision=precision, device=device, dtype=dtype)
+        self.qkv_proj = _make_svdq_linear(dim, 3 * dim, **lin)
+        self.qkv_proj_context = _make_svdq_linear(dim, 3 * dim, **lin)
+        self.out_proj = _make_svdq_linear(dim, dim, **lin)
+        self.out_proj_context = _make_svdq_linear(dim, dim, **lin)
 
         self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6).to(device=device, dtype=dtype)
         self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6).to(device=device, dtype=dtype)
 
-        self.mlp_fc1 = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=4 * dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.mlp_fc2 = SVDQW4A4Linear(
-            in_features=4 * dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-
-        self.mlp_context_fc1 = SVDQW4A4Linear(
-            in_features=dim,
-            out_features=4 * dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        self.mlp_context_fc2 = SVDQW4A4Linear(
-            in_features=4 * dim,
-            out_features=dim,
-            rank=rank,
-            bias=True,
-            precision=precision,
-            torch_dtype=dtype,
-            device=device,
-        )
-        # Chroma int4 compatibility:
-        # the context-stream MLP down-projection also needs the signed
-        # activation path for stable parity and image quality.
+        self.mlp_fc1 = _make_svdq_linear(dim, 4 * dim, **lin)
+        self.mlp_fc2 = _make_svdq_linear(4 * dim, dim, **lin)
+        self.mlp_context_fc1 = _make_svdq_linear(dim, 4 * dim, **lin)
+        self.mlp_context_fc2 = _make_svdq_linear(4 * dim, dim, **lin)
         self.mlp_context_fc2.act_unsigned = False
 
-        self.norm_q = self.attn.norm_q
-        self.norm_k = self.attn.norm_k
-        self.norm_added_q = self.attn.norm_added_q
-        self.norm_added_k = self.attn.norm_added_k
-
-    def forward(
-        self,
-        hidden_states,
-        encoder_hidden_states,
-        temb,
-        image_rotary_emb=None,
-        attention_mask_1d=None,
-        cpp_workspace: dict | None = None,
-        cpp_mask: torch.Tensor | None = None,
-    ):
-        from nunchaku.ops.fused import fused_gelu_mlp, fused_qkv_norm_rottary
-
+    def forward(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=None,
+                attention_mask_1d=None):
         temb_img, temb_txt = temb[:, :6], temb[:, 6:]
         norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb_img)
         norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(
@@ -547,207 +273,228 @@ class NunchakuChromaTransformerBlock(nn.Module, NunchakuChromaTransformerBlockMi
         )
 
         rotary_img, rotary_txt = image_rotary_emb
-        use_cpp = _should_use_cpp_additive_attn(
-            attention_mask_1d=attention_mask_1d,
-            hidden_states=norm_hidden_states,
-            head_dim=self.head_dim,
-        )
-
         txt_len = int(norm_encoder_hidden_states.shape[1])
-        img_len = int(norm_hidden_states.shape[1])
-
-        if use_cpp:
-            assert cpp_workspace is not None and cpp_mask is not None
-            txt_pad = _pad_to_multiple(txt_len, 256)
-            q, k, v, _, out = _get_cpp_workspace_tensors(cpp_workspace)
-            fused_qkv_norm_rottary(
-                norm_hidden_states,
-                self.qkv_proj,
-                self.attn.norm_q,
-                self.attn.norm_k,
-                rotary_img,
-                output=(q[:, :, txt_pad:], k[:, :, txt_pad:], v[:, :, txt_pad:]),
-                attn_tokens=img_len,
-            )
-            fused_qkv_norm_rottary(
-                norm_encoder_hidden_states,
-                self.qkv_proj_context,
-                self.attn.norm_added_q,
-                self.attn.norm_added_k,
-                rotary_txt,
-                output=(q[:, :, :txt_pad], k[:, :, :txt_pad], v[:, :, :txt_pad]),
-                attn_tokens=txt_len,
-            )
-            if _run_cpp_additive_attention(q, k, v, cpp_mask, out, context="dual-block cpp path"):
-                context_attn_output = out[:, :txt_len, :]
-                attn_output = out[:, txt_pad : txt_pad + img_len, :]
-            else:
-                use_cpp = False
-
-        if not use_cpp:
-            query, key, value = _fused_qkv_heads(
-                norm_hidden_states, self.qkv_proj, self.attn.norm_q, self.attn.norm_k, rotary_img, self.heads
-            )
-            c_query, c_key, c_value = _fused_qkv_heads(
-                norm_encoder_hidden_states,
-                self.qkv_proj_context,
-                self.attn.norm_added_q,
-                self.attn.norm_added_k,
-                rotary_txt,
-                self.heads,
-            )
 
-            query = torch.cat([c_query, query], dim=1)
-            key = torch.cat([c_key, key], dim=1)
-            value = torch.cat([c_value, value], dim=1)
-
-            attn_out = _dispatch_attention(query, key, value, attention_mask_1d)
-            attn_out = attn_out.flatten(2, 3).to(query.dtype)
+        query, key, value = _qkv_norm_rotary(
+            norm_hidden_states, self.qkv_proj, self.attn["norm_q"], self.attn["norm_k"], rotary_img, self.heads
+        )
+        c_query, c_key, c_value = _qkv_norm_rotary(
+            norm_encoder_hidden_states, self.qkv_proj_context,
+            self.attn["norm_added_q"], self.attn["norm_added_k"], rotary_txt, self.heads,
+        )
+        query = torch.cat([c_query, query], dim=1)
+        key = torch.cat([c_key, key], dim=1)
+        value = torch.cat([c_value, value], dim=1)
 
-            context_attn_output, attn_output = attn_out.split_with_sizes([txt_len, attn_out.shape[1] - txt_len], dim=1)
+        attn_out = _dispatch_attention(query, key, value, attention_mask_1d)
+        attn_out = attn_out.flatten(2, 3).to(query.dtype)
+        context_attn_output, attn_output = attn_out.split_with_sizes([txt_len, attn_out.shape[1] - txt_len], dim=1)
 
         attn_output = self.out_proj(attn_output)
         context_attn_output = self.out_proj_context(context_attn_output)
 
         hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
-        nh = self.norm2(hidden_states)
-        nh = nh * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
-        ff = fused_gelu_mlp(nh, self.mlp_fc1, self.mlp_fc2)
-        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff
+        nh = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        mlp_h = F.gelu(self.mlp_fc1(nh), approximate="tanh")
+        hidden_states = hidden_states + gate_mlp.unsqueeze(1) * self.mlp_fc2(mlp_h)
 
         encoder_hidden_states = encoder_hidden_states + c_gate_msa.unsqueeze(1) * context_attn_output
-        ne = self.norm2_context(encoder_hidden_states)
-        ne = ne * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
-        c_ff = fused_gelu_mlp(ne, self.mlp_context_fc1, self.mlp_context_fc2)
-        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * c_ff
+        ne = self.norm2_context(encoder_hidden_states) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+        mlp_e = F.gelu(self.mlp_context_fc1(ne), approximate="tanh")
+        encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * self.mlp_context_fc2(mlp_e)
 
         if encoder_hidden_states.dtype == torch.float16:
             encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
         return encoder_hidden_states, hidden_states
 
 
+# ===========================================================================
+# Top-level model
+# ===========================================================================
+
+
 class NunchakuChromaTransformer2dModel(ModelMixin, ConfigMixin):
-    """
-    A Chroma-faithful transformer that loads the exact DeepCompressor/nunchaku-ext safetensors layout.
-    """
+    """Chroma transformer that loads the DeepCompressor/nunchaku-ext safetensors layout."""
 
-    def __init__(
-        self,
-        *,
-        config: dict[str, Any],
-        rank: int,
-        precision: str,
-        device,
-        dtype,
-    ):
+    def __init__(self, *, config: dict[str, Any], rank: int, precision: str, device, dtype):
         super().__init__()
         from diffusers.models.transformers.transformer_chroma import (
             ChromaAdaLayerNormContinuousPruned,
             ChromaApproximator,
             ChromaCombinedTimestepTextProjEmbeddings,
         )
-        from nunchaku.models.embeddings import NunchakuFluxPosEmbed
+        from diffusers.models.transformers.transformer_flux import FluxPosEmbed
+
+        patch_size = int(config["patch_size"])
+        in_channels = int(config["in_channels"])
+        out_channels = int(config.get("out_channels") or in_channels)
+        num_layers = int(config["num_layers"])
+        num_single_layers = int(config["num_single_layers"])
+        attention_head_dim = int(config["attention_head_dim"])
+        num_attention_heads = int(config["num_attention_heads"])
+        joint_attention_dim = int(config["joint_attention_dim"])
+        axes_dims_rope = tuple(config.get("axes_dims_rope", (16, 56, 56)))
+        approx_channels = int(config.get("approximator_num_channels", 64))
+        approx_hidden = int(config.get("approximator_hidden_dim", 5120))
+        approx_layers = int(config.get("approximator_layers", 5))
 
         self.register_to_config(
-            patch_size=int(config["patch_size"]),
-            in_channels=int(config["in_channels"]),
-            out_channels=config.get("out_channels", None),
-            num_layers=int(config["num_layers"]),
-            num_single_layers=int(config["num_single_layers"]),
-            attention_head_dim=int(config["attention_head_dim"]),
-            num_attention_heads=int(config["num_attention_heads"]),
-            joint_attention_dim=int(config["joint_attention_dim"]),
-            axes_dims_rope=tuple(config.get("axes_dims_rope", (16, 56, 56))),
-            approximator_num_channels=int(config.get("approximator_num_channels", 64)),
-            approximator_hidden_dim=int(config.get("approximator_hidden_dim", 5120)),
-            approximator_layers=int(config.get("approximator_layers", 5)),
+            patch_size=patch_size, in_channels=in_channels, out_channels=out_channels,
+            num_layers=num_layers, num_single_layers=num_single_layers,
+            attention_head_dim=attention_head_dim, num_attention_heads=num_attention_heads,
+            joint_attention_dim=joint_attention_dim, axes_dims_rope=axes_dims_rope,
+            approximator_num_channels=approx_channels, approximator_hidden_dim=approx_hidden,
+            approximator_layers=approx_layers,
         )
         self.nunchaku_precision = str(precision)
         self.nunchaku_rank = int(rank)
 
-        patch_size = int(self.config.patch_size)
-        in_channels = int(self.config.in_channels)
-        out_channels = int(getattr(self.config, "out_channels", None) or in_channels)
-        num_layers = int(self.config.num_layers)
-        num_single_layers = int(self.config.num_single_layers)
-        attention_head_dim = int(self.config.attention_head_dim)
-        num_attention_heads = int(self.config.num_attention_heads)
-        joint_attention_dim = int(self.config.joint_attention_dim)
-        axes_dims_rope = tuple(self.config.axes_dims_rope)
-
         self.out_channels = out_channels
-        self.inner_dim = num_attention_heads * attention_head_dim
+        inner_dim = num_attention_heads * attention_head_dim
+        self.inner_dim = inner_dim
 
-        self.pos_embed = NunchakuFluxPosEmbed(dim=self.inner_dim, theta=10000, axes_dim=list(axes_dims_rope)).to(
-            device=device, dtype=dtype
-        )
+        def _to(m):
+            return m.to(device=device, dtype=dtype)
 
-        self.time_text_embed = ChromaCombinedTimestepTextProjEmbeddings(
-            num_channels=int(self.config.approximator_num_channels) // 4,
+        self.pos_embed = _to(FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope))
+        self.time_text_embed = _to(ChromaCombinedTimestepTextProjEmbeddings(
+            num_channels=approx_channels // 4,
             out_dim=3 * num_single_layers + 2 * 6 * num_layers + 2,
-        ).to(device=device, dtype=dtype)
-
-        self.distilled_guidance_layer = ChromaApproximator(
-            in_dim=int(self.config.approximator_num_channels),
-            out_dim=self.inner_dim,
-            hidden_dim=int(self.config.approximator_hidden_dim),
-            n_layers=int(self.config.approximator_layers),
-        ).to(device=device, dtype=dtype)
-
-        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=True).to(device=device, dtype=dtype)
-        self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=True).to(device=device, dtype=dtype)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                NunchakuChromaTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                    rank=rank,
-                    precision=precision,
-                    device=device,
-                    dtype=dtype,
-                )
-                for _ in range(num_layers)
-            ]
-        )
-        self.single_transformer_blocks = nn.ModuleList(
-            [
-                NunchakuChromaSingleTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                    mlp_ratio=4.0,
-                    rank=rank,
-                    precision=precision,
-                    device=device,
-                    dtype=dtype,
-                )
-                for _ in range(num_single_layers)
-            ]
-        )
-
-        self.norm_out = ChromaAdaLayerNormContinuousPruned(
-            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6
-        ).to(device=device, dtype=dtype)
-        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * out_channels, bias=True).to(
-            device=device, dtype=dtype
-        )
+        ))
+        self.distilled_guidance_layer = _to(ChromaApproximator(
+            in_dim=approx_channels, out_dim=inner_dim, hidden_dim=approx_hidden, n_layers=approx_layers,
+        ))
+        self.context_embedder = _to(nn.Linear(joint_attention_dim, inner_dim, bias=True))
+        self.x_embedder = _to(nn.Linear(in_channels, inner_dim, bias=True))
+
+        block_kw = dict(num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim,
+                        rank=rank, precision=precision, device=device, dtype=dtype)
+        self.transformer_blocks = nn.ModuleList([
+            NunchakuChromaTransformerBlock(dim=inner_dim, **block_kw) for _ in range(num_layers)
+        ])
+        self.single_transformer_blocks = nn.ModuleList([
+            NunchakuChromaSingleTransformerBlock(dim=inner_dim, mlp_ratio=4.0, **block_kw) for _ in range(num_single_layers)
+        ])
+
+        self.norm_out = _to(ChromaAdaLayerNormContinuousPruned(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6))
+        self.proj_out = _to(nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True))
 
         self.encoder_hid_proj = None
+        self.offload = False
+        self.transformer_block_offload_manager = None
+        self.single_transformer_block_offload_manager = None
+
+    # ---- offload management ------------------------------------------------
+
+    def set_offload(self, offload: bool, **kwargs):
+        from nunchaku.models.utils import CPUOffloadManager
+
+        if offload == self.offload:
+            return
+        self.offload = offload
+        if offload:
+            use_pin_memory = kwargs.get("use_pin_memory", True)
+            num_blocks_on_gpu = kwargs.get("num_blocks_on_gpu", 1)
+            on_gpu = [self.pos_embed, self.time_text_embed, self.distilled_guidance_layer,
+                      self.context_embedder, self.x_embedder, self.norm_out, self.proj_out]
+            self.transformer_block_offload_manager = CPUOffloadManager(
+                self.transformer_blocks, use_pin_memory=use_pin_memory,
+                on_gpu_modules=on_gpu, num_blocks_on_gpu=num_blocks_on_gpu,
+            )
+            self.single_transformer_block_offload_manager = CPUOffloadManager(
+                self.single_transformer_blocks, use_pin_memory=use_pin_memory,
+                num_blocks_on_gpu=num_blocks_on_gpu,
+            )
+        else:
+            self.transformer_block_offload_manager = None
+            self.single_transformer_block_offload_manager = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def to(self, *args, **kwargs):
+        device_arg_or_kwarg_present = any(isinstance(arg, torch.device) for arg in args) or "device" in kwargs
+        dtype_present = "dtype" in kwargs or any(isinstance(arg, torch.dtype) for arg in args)
+        for arg in args:
+            if isinstance(arg, str):
+                try:
+                    torch.device(arg)
+                    device_arg_or_kwarg_present = True
+                except RuntimeError:
+                    pass
+        if dtype_present:
+            raise ValueError(
+                "Casting a quantized model to a new `dtype` is unsupported. Use the `torch_dtype` "
+                "argument in `from_pretrained` instead."
+            )
+        if getattr(self, "offload", False) and device_arg_or_kwarg_present:
+            warn("Skipping moving the model to GPU as offload is enabled", UserWarning)
+            return self
+        return super(type(self), self).to(*args, **kwargs)
+
+    def _iter_blocks(self, blocks, offload_manager):
+        """Yield (index, block) with transparent offload management."""
+        if offload_manager is not None:
+            compute_stream = torch.cuda.current_stream()
+            offload_manager.initialize(compute_stream)
+            for i in range(len(blocks)):
+                with torch.cuda.stream(compute_stream):
+                    yield i, offload_manager.get_block(i)
+                offload_manager.step(compute_stream)
+        else:
+            yield from enumerate(blocks)
+
+    def _strip_accelerate_hooks_if_needed(self):
+        """Neutralise accelerate per-module hooks that break quantized kernels."""
+        try:
+            from accelerate.hooks import remove_hook_from_module
+        except ImportError:
+            return
+
+        children = [m for m in self.modules() if m is not self]
+        has_hooks = any(getattr(m, "_hf_hook", None) is not None for m in children)
+        if not has_hooks:
+            return
+
+        if not self._sequential_offload_warned:
+            self._sequential_offload_warned = True
+            warn(
+                "Detected accelerate sequential hooks on sub-modules of "
+                "NunchakuChromaTransformer2dModel.  These are incompatible with "
+                "quantized CUDA kernels and will be removed.  "
+                "Use transformer.set_offload(True) and "
+                "pipe._exclude_from_cpu_offload.append('transformer') instead.",
+                UserWarning,
+            )
+
+        for m in children:
+            if getattr(m, "_hf_hook", None) is not None:
+                remove_hook_from_module(m)
+
+        top_hook = getattr(self, "_hf_hook", None)
+        if top_hook is not None:
+            top_hook.execution_device = None
+            top_hook.offload = False
+            if hasattr(top_hook, "io_same_device"):
+                top_hook.io_same_device = False
+            if hasattr(top_hook, "weights_map"):
+                top_hook.weights_map = None
+
+        if not self.offload:
+            self.set_offload(True)
+
+    # ---- loading -----------------------------------------------------------
 
     @classmethod
     def from_pretrained(
-        cls,
-        pretrained_model_name_or_path: str | Path,
-        *,
-        device: str = "cuda",
-        torch_dtype: Any = None,
-        precision: str | None = None,
-        rank: int | None = None,
-        verbose: bool = True,
-        return_report: bool = False,
+        cls, pretrained_model_name_or_path: str | Path, *,
+        device: str = "cuda", torch_dtype: Any = None, precision: str | None = None,
+        rank: int | None = None, pin_memory: bool | str = "auto",
+        offload: bool = False, verbose: bool = True, return_report: bool = False,
     ):
+        from nunchaku.models.transformers.utils import patch_scale_key
+        from nunchaku.utils import get_precision_from_quantization_config, load_state_dict_in_safetensors
+
         ckpt = Path(pretrained_model_name_or_path)
         if not ckpt.exists():
             raise FileNotFoundError(str(ckpt))
@@ -755,100 +502,75 @@ class NunchakuChromaTransformer2dModel(ModelMixin, ConfigMixin):
         if torch_dtype is None:
             torch_dtype = torch.bfloat16
 
-        sd_raw, md = _load_safetensors_state_dict(ckpt, device="cpu")
-        if "config" not in md:
-            raise ValueError("Missing required safetensors metadata: 'config'")
-        if "quantization_config" not in md:
-            raise ValueError("Missing required safetensors metadata: 'quantization_config'")
+        sd_raw, md = load_state_dict_in_safetensors(ckpt, device="cpu", return_metadata=True)
+        for required_key in ("config", "quantization_config"):
+            if required_key not in md:
+                raise ValueError(f"Missing required safetensors metadata: {required_key!r}")
 
         config = json.loads(md["config"])
-        if config.get("_class_name", None) != "ChromaTransformer2DModel":
-            raise ValueError(f"Unexpected config._class_name={config.get('_class_name')!r} (expected 'ChromaTransformer2DModel')")
-
-        quant_cfg = json.loads(md["quantization_config"])
-        from nunchaku.utils import get_precision_from_quantization_config
-
-        inferred_precision = get_precision_from_quantization_config(quant_cfg)
+        if config.get("_class_name") != "ChromaTransformer2DModel":
+            raise ValueError(f"Unexpected config._class_name={config.get('_class_name')!r}")
 
+        inferred_precision = get_precision_from_quantization_config(json.loads(md["quantization_config"]))
         sd = _convert_checkpoint_state_dict(sd_raw)
         inferred_rank = _infer_rank_from_converted_state_dict(sd)
 
         if precision is not None and str(precision) != str(inferred_precision):
-            raise ValueError(
-                f"precision mismatch: got precision={precision!r}, but checkpoint says {inferred_precision!r} "
-                f"(from safetensors metadata 'quantization_config')."
-            )
+            raise ValueError(f"precision mismatch: {precision!r} vs checkpoint {inferred_precision!r}")
         if rank is not None and int(rank) != int(inferred_rank):
-            raise ValueError(
-                f"rank mismatch: got rank={int(rank)}, but checkpoint implies rank={int(inferred_rank)} "
-                f"(from '*.proj_down' tensors)."
-            )
-
-        model = cls(
-            config=config,
-            rank=int(inferred_rank),
-            precision=str(inferred_precision),
-            device=torch.device(device),
-            dtype=torch_dtype,
-        )
+            raise ValueError(f"rank mismatch: {rank} vs checkpoint {inferred_rank}")
 
-        from nunchaku.models.transformers.utils import patch_scale_key
+        model = cls(config=config, rank=inferred_rank, precision=str(inferred_precision),
+                    device=torch.device(device), dtype=torch_dtype)
 
         patch_scale_key(model, sd)
-
         wanted = set(model.state_dict().keys())
         sd_filtered = {k: v for k, v in sd.items() if k in wanted}
         model.load_state_dict(sd_filtered, strict=True)
 
         if str(inferred_precision) == "int4":
-            # Chroma int4 compatibility:
-            # several dual-stream layers match the exported model much better
-            # when the runtime consumes `smooth_factor_orig` instead of
-            # `smooth_factor`. This is intentionally scoped to Chroma int4.
             for block in model.transformer_blocks:
-                block.qkv_proj.smooth_factor.data.copy_(block.qkv_proj.smooth_factor_orig.data)
-                block.qkv_proj_context.smooth_factor.data.copy_(block.qkv_proj_context.smooth_factor_orig.data)
-                block.mlp_context_fc2.smooth_factor.data.copy_(block.mlp_context_fc2.smooth_factor_orig.data)
-
-        _maybe_log(verbose, "[nunchaku.chroma] loaded:", str(ckpt))
-        # _maybe_log(verbose, "[nunchaku.chroma] precision:", inferred_precision, "rank:", inferred_rank, "dtype:", torch_dtype)
-        # _maybe_log(
-        #     verbose,
-        #     "[nunchaku.chroma] config.num_layers:",
-        #     int(config["num_layers"]),
-        #     "num_single_layers:",
-        #     int(config["num_single_layers"]),
-        # )
+                for attr in ("qkv_proj", "qkv_proj_context", "mlp_context_fc2"):
+                    layer = getattr(block, attr, None)
+                    if layer is not None and hasattr(layer, "smooth_factor_orig"):
+                        layer.smooth_factor.data.copy_(layer.smooth_factor_orig.data)
+
+        if verbose:
+            print("[nunchaku.chroma] loaded:", str(ckpt))
+
+        model.set_offload(offload)
 
         if return_report:
-            return model, LoadReport(config=config, precision=str(inferred_precision), rank=int(inferred_rank))
+            return model, LoadReport(config=config, precision=str(inferred_precision), rank=inferred_rank)
         return model
 
+    # ---- forward -----------------------------------------------------------
+
+    _sequential_offload_warned: bool = False
+
     def forward(
-        self,
-        hidden_states,
-        encoder_hidden_states=None,
-        timestep=None,
-        img_ids=None,
-        txt_ids=None,
-        attention_mask=None,
+        self, hidden_states, encoder_hidden_states=None, timestep=None,
+        img_ids=None, txt_ids=None, attention_mask=None,
         joint_attention_kwargs: Optional[dict[str, Any]] = None,
-        controlnet_block_samples=None,
-        controlnet_single_block_samples=None,
-        return_dict: bool = True,
-        controlnet_blocks_repeat: bool = False,
+        controlnet_block_samples=None, controlnet_single_block_samples=None,
+        return_dict: bool = True, controlnet_blocks_repeat: bool = False,
     ):
         del controlnet_blocks_repeat
 
         from diffusers.models.modeling_outputs import Transformer2DModelOutput
-        from nunchaku.models.embeddings import pack_rotemb
-        from nunchaku.utils import pad_tensor
 
         if controlnet_block_samples is not None or controlnet_single_block_samples is not None:
             raise NotImplementedError("ControlNet is not supported in NunchakuChromaTransformer2dModel")
         if joint_attention_kwargs:
             raise NotImplementedError("joint_attention_kwargs is not supported in NunchakuChromaTransformer2dModel")
 
+        self._strip_accelerate_hooks_if_needed()
+
+        if self.offload:
+            device = hidden_states.device
+            self.transformer_block_offload_manager.set_device(device)
+            self.single_transformer_block_offload_manager.set_device(device)
+
         if txt_ids.ndim == 3:
             txt_ids = txt_ids[0]
         if img_ids.ndim == 3:
@@ -858,83 +580,51 @@ class NunchakuChromaTransformer2dModel(ModelMixin, ConfigMixin):
         timestep = timestep.to(hidden_states.dtype) * 1000
         batch_size = int(hidden_states.shape[0])
 
-        input_vec = self.time_text_embed(timestep)
-        pooled_temb = self.distilled_guidance_layer(input_vec)
-
+        pooled_temb = self.distilled_guidance_layer(self.time_text_embed(timestep))
         encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
         ids = torch.cat((txt_ids, img_ids), dim=0)
         image_rotary_emb = self.pos_embed(ids)
 
         txt_tokens = int(encoder_hidden_states.shape[1])
         img_tokens = int(hidden_states.shape[1])
         attn_mask_1d = attention_mask
-        assert image_rotary_emb.ndim == 6
-        assert image_rotary_emb.shape[0] == 1
-        assert image_rotary_emb.shape[1] == 1
-        assert image_rotary_emb.shape[2] == 1 * (txt_tokens + img_tokens)
-        image_rotary_emb = image_rotary_emb.reshape([1, txt_tokens + img_tokens, *image_rotary_emb.shape[3:]])
-        rotary_emb_txt = pack_rotemb(pad_tensor(image_rotary_emb[:, :txt_tokens, ...], 256, 1))
-        rotary_emb_img = pack_rotemb(pad_tensor(image_rotary_emb[:, txt_tokens:, ...], 256, 1))
-        rotary_emb_single = pack_rotemb(pad_tensor(image_rotary_emb, 256, 1))
-
-        rotary_emb_txt = _expand_batch_dim(rotary_emb_txt, batch_size)
-        rotary_emb_img = _expand_batch_dim(rotary_emb_img, batch_size)
-        rotary_emb_single = _expand_batch_dim(rotary_emb_single, batch_size)
-
-        use_cpp_ws = _should_use_cpp_additive_attn(
-            attention_mask_1d=attn_mask_1d,
-            hidden_states=hidden_states,
-            head_dim=int(self.config.attention_head_dim),
-        )
-        ws_dual: dict | None = None
-        ws_single: dict | None = None
-        mask_dual: torch.Tensor | None = None
-        mask_single: torch.Tensor | None = None
-        if use_cpp_ws:
-            ws_dual, ws_single, mask_dual, mask_single = _prepare_cpp_context(
-                self, hidden_states, attention_mask, txt_tokens=txt_tokens, img_tokens=img_tokens
-            )
+
+        freqs_cos, freqs_sin = image_rotary_emb
+        rotary_emb_txt = (freqs_cos[:txt_tokens], freqs_sin[:txt_tokens])
+        rotary_emb_img = (freqs_cos[txt_tokens:], freqs_sin[txt_tokens:])
+        rotary_emb_single = image_rotary_emb
 
         num_layers = len(self.transformer_blocks)
         num_single = len(self.single_transformer_blocks)
         img_offset = 3 * num_single
         txt_offset = img_offset + 6 * num_layers
 
-        for i, block in enumerate(self.transformer_blocks):
+        dual_mgr = self.transformer_block_offload_manager if self.offload else None
+        single_mgr = self.single_transformer_block_offload_manager if self.offload else None
+
+        for i, block in self._iter_blocks(self.transformer_blocks, dual_mgr):
             img_mod = img_offset + 6 * i
             txt_mod = txt_offset + 6 * i
             temb = torch.cat(
-                (pooled_temb[:, img_mod : img_mod + 6], pooled_temb[:, txt_mod : txt_mod + 6]),
-                dim=1,
+                (pooled_temb[:, img_mod : img_mod + 6], pooled_temb[:, txt_mod : txt_mod + 6]), dim=1,
             )
             encoder_hidden_states, hidden_states = block(
-                hidden_states=hidden_states,
-                encoder_hidden_states=encoder_hidden_states,
-                temb=temb,
-                image_rotary_emb=(rotary_emb_img, rotary_emb_txt),
-                attention_mask_1d=attn_mask_1d,
-                cpp_workspace=ws_dual,
-                cpp_mask=mask_dual,
+                hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb,
+                image_rotary_emb=(rotary_emb_img, rotary_emb_txt), attention_mask_1d=attn_mask_1d,
             )
 
         hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
-        for i, block in enumerate(self.single_transformer_blocks):
-            start = 3 * i
-            temb = pooled_temb[:, start : start + 3]
+        for i, block in self._iter_blocks(self.single_transformer_blocks, single_mgr):
+            temb = pooled_temb[:, 3 * i : 3 * i + 3]
             hidden_states = block(
-                hidden_states=hidden_states,
-                temb=temb,
-                image_rotary_emb=rotary_emb_single,
+                hidden_states=hidden_states, temb=temb, image_rotary_emb=rotary_emb_single,
                 attention_mask_1d=attn_mask_1d,
-                cpp_workspace=ws_single,
-                cpp_mask=mask_single,
             )
 
         hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
-
-        temb = pooled_temb[:, -2:]
-        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.norm_out(hidden_states, pooled_temb[:, -2:])
         output = self.proj_out(hidden_states)
 
         if not return_dict: