feat: add Krea 2 W4A4 support

Felipe Sztutman · Felipe Sztutman · commit 777644e02712 · 2026-07-20T08:41:01.000-03:00
Adds SVDQuant W4A4 support for krea/Krea-2-Turbo through the existing fused
low-bit kernels. No new CUDA was required; the current kernels already serve
this shape, so the port is pure Python.

The attention processor expands the grouped-query key/value heads explicitly
rather than passing enable_gqa=True, because Krea 2 always supplies an
attention mask and SDPA will not serve that flag together with a mask on the
flash backend. It falls back to the math backend silently, at roughly 3x the
cost end to end.

Also adds the example script, usage and API docs, and an LPIPS regression test
following the Z-Image layout.
diff --git a/README.md b/README.md
@@ -17,6 +17,7 @@ Join our user groups on [**Discord**](https://discord.gg/Wk6PnwX9Sm) and [**WeCh
 
 ## News
 
+- **[2026-07-20]** 🔥 Release **4-bit [krea/Krea-2-Turbo](https://huggingface.co/krea/Krea-2-Turbo)**! W4A4 checkpoint on [Hugging Face](https://huggingface.co/felipesztutman/Krea-2-Turbo-W4A4-Nunchaku), and try it with our [example script](examples/v1/krea-2-turbo.py).
 - **[2026-01-12]** 🚀 **v1.2.0 Released!** Enjoy a **20–30% Z-Image performance boost**, seamless **LoRA support with native ComfyUI nodes**, **INT4 support for 20-series GPUs**, and **compatibility with ComfyUI 0.7**!
 - **[2025-12-20]** 🔥 Release **4-bit [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo)**! Download on [Hugging Face](https://huggingface.co/nunchaku-ai/nunchaku-z-image-turbo) or [ModelScope](https://modelscope.cn/models/nunchaku-tech/nunchaku-z-image-turbo), and try it with our [example script](examples/v1/z-image-turbo.py).
 - **[2025-09-25]** 🔥 Release **4-bit [4/8-step lightning Qwen-Image-Edit-2509](https://huggingface.co/lightx2v/Qwen-Image-Lightning)**! Download on [Hugging Face](https://huggingface.co/nunchaku-ai/nunchaku-qwen-image-edit-2509) or [ModelScope](https://modelscope.cn/models/nunchaku-tech/nunchaku-qwen-image-edit-2509), and try it with our [example script](examples/v1/qwen-image-edit-2509-lightning.py).
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -28,6 +28,7 @@ Check out `DeepCompressor <github_deepcompressor_>`_ for the quantization librar
     usage/cache.rst
     usage/pulid.rst
     usage/ip_adapter.rst
+    usage/krea2.rst
     usage/zimage.rst
 
 .. toctree::
diff --git a/docs/source/python_api/nunchaku.models.attention_processors.krea2.rst b/docs/source/python_api/nunchaku.models.attention_processors.krea2.rst
@@ -0,0 +1,7 @@
+nunchaku.models.attention\_processors.krea2
+===========================================
+
+.. automodule:: nunchaku.models.attention_processors.krea2
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/python_api/nunchaku.models.attention_processors.rst b/docs/source/python_api/nunchaku.models.attention_processors.rst
@@ -6,4 +6,5 @@ nunchaku.models.attention_processors
 
    nunchaku.models.attention_processors.flux
    nunchaku.models.attention_processors.qwenimage
+   nunchaku.models.attention_processors.krea2
    nunchaku.models.attention_processors.zimage
diff --git a/docs/source/python_api/nunchaku.models.transformers.rst b/docs/source/python_api/nunchaku.models.transformers.rst
@@ -7,6 +7,7 @@ nunchaku.models.transformers
    nunchaku.models.transformers.transformer_flux
    nunchaku.models.transformers.transformer_flux_v2
    nunchaku.models.transformers.transformer_qwenimage
+   nunchaku.models.transformers.transformer_krea2
    nunchaku.models.transformers.transformer_zimage
    nunchaku.models.transformers.transformer_sana
    nunchaku.models.transformers.utils
diff --git a/docs/source/python_api/nunchaku.models.transformers.transformer_krea2.rst b/docs/source/python_api/nunchaku.models.transformers.transformer_krea2.rst
@@ -0,0 +1,7 @@
+nunchaku.models.transformers.transformer\_krea2
+===============================================
+
+.. automodule:: nunchaku.models.transformers.transformer_krea2
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/usage/krea2.rst b/docs/source/usage/krea2.rst
@@ -0,0 +1,22 @@
+Krea 2
+======
+
+The following is the example of running Nunchaku version of Krea 2 text-to-image pipeline.
+
+.. tabs::
+
+   .. tab:: Krea-2-Turbo
+
+      .. literalinclude:: ../../../examples/v1/krea-2-turbo.py
+         :language: python
+         :caption: Running Krea-2-Turbo (`examples/v1/krea-2-turbo.py <https://github.com/nunchaku-tech/nunchaku/blob/main/examples/v1/krea-2-turbo.py>`__)
+         :linenos:
+
+.. note::
+
+   Krea 2 uses grouped-query attention and always passes an attention mask, because text and
+   image share one sequence. PyTorch SDPA will not serve ``enable_gqa=True`` together with a
+   mask on the flash backend and falls back to the math backend silently, which costs roughly
+   3x. The processor therefore expands the key/value heads explicitly.
+
+For more details, see :class:`~nunchaku.models.transformers.transformer_krea2.NunchakuKrea2Transformer2DModel`.
diff --git a/examples/v1/krea-2-turbo.py b/examples/v1/krea-2-turbo.py
@@ -0,0 +1,33 @@
+import torch
+from diffusers import Krea2Pipeline
+
+from nunchaku import NunchakuKrea2Transformer2DModel
+from nunchaku.utils import get_precision, is_turing
+
+if __name__ == "__main__":
+    precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
+    rank = 32  # Rank of the SVDQuant low-rank branch that absorbs the outlier activations
+    dtype = torch.float16 if is_turing() else torch.bfloat16  # Use float16 when Turing (20- series) GPU is used.
+    transformer = NunchakuKrea2Transformer2DModel.from_pretrained(
+        f"felipesztutman/Krea-2-Turbo-W4A4-Nunchaku/svdq-{precision}_r{rank}-krea-2-turbo.safetensors",
+        torch_dtype=dtype,
+    )
+
+    pipe = Krea2Pipeline.from_pretrained(
+        "krea/Krea-2-Turbo", transformer=transformer, torch_dtype=dtype, low_cpu_mem_usage=False
+    )
+    pipe.enable_sequential_cpu_offload()  # enable sequential CPU offload for low vram
+    # pipe = pipe.to("cuda") # or else comment the line above and uncomment this line to put all components to GPU
+
+    prompt = "a red fox in deep snow at dusk, backlit, shallow depth of field"
+
+    image = pipe(
+        prompt=prompt,
+        height=1024,
+        width=1024,
+        num_inference_steps=8,
+        guidance_scale=0.0,  # Guidance should be 0 for the Turbo models
+        generator=torch.Generator().manual_seed(1000),
+    ).images[0]
+
+    image.save(f"krea-2-turbo-{precision}_r{rank}_{str(dtype)}.png")
diff --git a/nunchaku/__init__.py b/nunchaku/__init__.py
@@ -1,6 +1,7 @@
 from .models import (
     NunchakuFluxTransformer2dModel,
     NunchakuFluxTransformer2DModelV2,
+    NunchakuKrea2Transformer2DModel,
     NunchakuQwenImageTransformer2DModel,
     NunchakuSanaTransformer2DModel,
     NunchakuT5EncoderModel,
@@ -13,5 +14,6 @@
     "NunchakuT5EncoderModel",
     "NunchakuFluxTransformer2DModelV2",
     "NunchakuQwenImageTransformer2DModel",
+    "NunchakuKrea2Transformer2DModel",
     "NunchakuZImageTransformer2DModel",
 ]
diff --git a/nunchaku/merge_safetensors.py b/nunchaku/merge_safetensors.py
@@ -48,7 +48,7 @@ def merge_safetensors(
     pretrained_model_name_or_path : str or os.PathLike
         Path to the model directory or HuggingFace repo.
     model_class : str
-        Specify model class. E.g. NunchakuFluxTransformer2dModel or NunchakuZImageTransformer2DModel
+        Specify model class. E.g. NunchakuFluxTransformer2dModel, NunchakuZImageTransformer2DModel or NunchakuKrea2Transformer2DModel
     **kwargs
         Additional keyword arguments for subfolder, comfy_config_path, and HuggingFace download options.
 
@@ -163,7 +163,7 @@ def merge_safetensors(
         "--model-class",
         type=str,
         required=True,
-        help="Specify model class. E.g. NunchakuFluxTransformer2dModel or NunchakuZImageTransformer2DModel",
+        help="Specify model class. E.g. NunchakuFluxTransformer2dModel, NunchakuZImageTransformer2DModel or NunchakuKrea2Transformer2DModel",
     )
     parser.add_argument("-o", "--output-path", type=Path, required=True, help="Path to output path")
     args = parser.parse_args()
diff --git a/nunchaku/models/__init__.py b/nunchaku/models/__init__.py
@@ -2,6 +2,7 @@
 from .transformers import (
     NunchakuFluxTransformer2dModel,
     NunchakuFluxTransformer2DModelV2,
+    NunchakuKrea2Transformer2DModel,
     NunchakuQwenImageTransformer2DModel,
     NunchakuSanaTransformer2DModel,
     NunchakuZImageTransformer2DModel,
@@ -13,5 +14,6 @@
     "NunchakuT5EncoderModel",
     "NunchakuFluxTransformer2DModelV2",
     "NunchakuQwenImageTransformer2DModel",
+    "NunchakuKrea2Transformer2DModel",
     "NunchakuZImageTransformer2DModel",
 ]
diff --git a/nunchaku/models/attention_processors/krea2.py b/nunchaku/models/attention_processors/krea2.py
@@ -0,0 +1,74 @@
+"""
+Attention processor for :class:`~nunchaku.models.transformers.transformer_krea2.NunchakuKrea2Attention`.
+
+Mirrors ``diffusers.models.transformers.transformer_krea2.Krea2AttnProcessor``, with the q/k/v/gate/out
+projections replaced by :class:`~nunchaku.models.linear.SVDQW4A4Linear`. All other operations run
+unchanged from the reference.
+
+Notes
+-----
+The key/value heads are expanded explicitly instead of using ``enable_gqa=True``, because SDPA will not
+serve that flag together with an attention mask on the flash backend and silently falls back to the math
+backend at roughly 3x the cost.
+"""
+
+from typing import Optional
+
+import torch
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+class NunchakuKrea2AttnProcessor:
+    """W4A4 attention processor for Krea 2 single-stream blocks."""
+
+    _attention_backend = None
+    _parallel_config = None
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        # One fused W4A4 GEMM for q/k/v/gate (they share the same input, one smooth + one low-rank branch,
+        # matching how the checkpoint was calibrated), then split. GQA: q has num_heads, k/v num_kv_heads.
+        qkv_gate = attn.to_qkv_gate(hidden_states)
+        q_dim = attn.num_heads * attn.head_dim
+        kv_dim = attn.num_kv_heads * attn.head_dim
+        query, key, value, gate = torch.split(qkv_gate, [q_dim, kv_dim, kv_dim, attn.hidden_size], dim=-1)
+        query = query.unflatten(-1, (attn.num_heads, attn.head_dim))
+        key = key.unflatten(-1, (attn.num_kv_heads, attn.head_dim))
+        value = value.unflatten(-1, (attn.num_kv_heads, attn.head_dim))
+
+        # Zero-centered RMSNorm on q/k (kept in fp32 by the source module).
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+        # Expand the grouped-query k/v heads rather than passing ``enable_gqa=True``. Krea 2 always
+        # supplies an attention mask, and the flash backend will not serve enable_gqa together with a
+        # mask, so it falls back to the math backend and costs ~3x end to end (22.3s vs 7.7s at
+        # 1024px/8 steps). The expansion is what keeps FlashAttention engaged.
+        if attn.num_heads != attn.num_kv_heads:
+            n_rep = attn.num_heads // attn.num_kv_heads
+            key = key.repeat_interleave(n_rep, dim=2)
+            value = value.repeat_interleave(n_rep, dim=2)
+
+        dtype = value.dtype
+        query, key = query.to(dtype), key.to(dtype)
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states * torch.sigmoid(gate)
+        return attn.to_out[0](hidden_states)
diff --git a/nunchaku/models/transformers/__init__.py b/nunchaku/models/transformers/__init__.py
@@ -1,5 +1,6 @@
 from .transformer_flux import NunchakuFluxTransformer2dModel
 from .transformer_flux_v2 import NunchakuFluxTransformer2DModelV2
+from .transformer_krea2 import NunchakuKrea2Transformer2DModel
 from .transformer_qwenimage import NunchakuQwenImageTransformer2DModel
 from .transformer_sana import NunchakuSanaTransformer2DModel
 from .transformer_zimage import NunchakuZImageTransformer2DModel
@@ -8,6 +9,7 @@
     "NunchakuFluxTransformer2dModel",
     "NunchakuSanaTransformer2DModel",
     "NunchakuFluxTransformer2DModelV2",
+    "NunchakuKrea2Transformer2DModel",
     "NunchakuQwenImageTransformer2DModel",
     "NunchakuZImageTransformer2DModel",
 ]
diff --git a/nunchaku/models/transformers/transformer_krea2.py b/nunchaku/models/transformers/transformer_krea2.py
@@ -0,0 +1,145 @@
+"""
+Nunchaku W4A4 runtime for the Krea 2 Turbo single-stream MMDiT.
+
+Strategy (correctness-first MVP): subclass the diffusers ``Krea2Transformer2DModel`` and reuse its whole
+forward pass (text fusion, [text, image] concat, 3D rotary, per-block AdaLN modulation, final layer)
+unchanged. Only the linear projections inside the 28 ``Krea2TransformerBlock`` heavy blocks are replaced
+with :class:`~nunchaku.models.linear.SVDQW4A4Linear`, which is where the SVDQuant W4A4 speedup lives. The
+light per-token modules (img_in, time_embed, text_fusion, txt_in, final_layer) stay in bf16, matching what
+deepcompressor skipped during calibration.
+
+Follow-up optimizations (not needed to run): fuse q/k/v(+gate) into one GEMM with norm+rope applied inside
+the kernel (see the Z-Image ``NunchakuZImageFusedModule``), and carry the gate as a 4th fused output.
+"""
+
+import json
+import os
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.transformers.transformer_krea2 import (
+    Krea2Attention,
+    Krea2SwiGLU,
+    Krea2Transformer2DModel,
+    Krea2TransformerBlock,
+)
+from huggingface_hub import utils
+
+from ...utils import get_precision
+from ..attention import NunchakuBaseAttention
+from ..attention_processors.krea2 import NunchakuKrea2AttnProcessor
+from ..linear import SVDQW4A4Linear
+from ..utils import fuse_linears
+from .utils import NunchakuModelLoaderMixin, convert_fp16, patch_scale_key
+
+
+class NunchakuKrea2Attention(NunchakuBaseAttention):
+    """W4A4 wrapper over :class:`~diffusers...Krea2Attention`.
+
+    Keeps grouped-query metadata and the fp32 q/k RMSNorm modules from the source attention, and swaps the
+    five projections (to_q, to_k, to_v, to_gate, to_out[0]) for :class:`SVDQW4A4Linear`.
+    """
+
+    def __init__(self, orig_attn: Krea2Attention, **kwargs):
+        super().__init__("krea2")
+        self.num_heads = orig_attn.num_heads
+        self.num_kv_heads = orig_attn.num_kv_heads
+        self.head_dim = orig_attn.head_dim
+        self.hidden_size = orig_attn.hidden_size
+
+        # Kept eager (fp32-normalized), unchanged from the reference.
+        self.norm_q = orig_attn.norm_q
+        self.norm_k = orig_attn.norm_k
+
+        # q/k/v/gate share one input, so the checkpoint calibrated them as a single group (one smooth,
+        # one low-rank branch). Fuse into a single quantized GEMM; the processor splits the output.
+        with torch.device("meta"):
+            to_qkv_gate = fuse_linears([orig_attn.to_q, orig_attn.to_k, orig_attn.to_v, orig_attn.to_gate])
+        self.to_qkv_gate = SVDQW4A4Linear.from_linear(to_qkv_gate, **kwargs)
+        self.to_out = orig_attn.to_out
+        self.to_out[0] = SVDQW4A4Linear.from_linear(self.to_out[0], **kwargs)
+
+    def set_processor(self, processor: str):
+        # Only one processor for Krea 2; the string is accepted for interface parity.
+        self.processor = NunchakuKrea2AttnProcessor()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        return self.processor(self, hidden_states, attention_mask, image_rotary_emb)
+
+
+class NunchakuKrea2SwiGLU(nn.Module):
+    """W4A4 SwiGLU. Same forward as :class:`~diffusers...Krea2SwiGLU` with quantized gate/up/down."""
+
+    def __init__(self, orig_ff: Krea2SwiGLU, **kwargs):
+        super().__init__()
+        self.gate = SVDQW4A4Linear.from_linear(orig_ff.gate, **kwargs)
+        self.up = SVDQW4A4Linear.from_linear(orig_ff.up, **kwargs)
+        self.down = SVDQW4A4Linear.from_linear(orig_ff.down, **kwargs)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states))
+
+
+class NunchakuKrea2Transformer2DModel(Krea2Transformer2DModel, NunchakuModelLoaderMixin):
+    """Nunchaku-optimized Krea2Transformer2DModel. Inherits the diffusers forward; only the 28 heavy
+    blocks' attention and SwiGLU are quantized."""
+
+    def _patch_model(self, **kwargs):
+        for block in self.transformer_blocks:
+            assert isinstance(block, Krea2TransformerBlock)
+            block.attn = NunchakuKrea2Attention(block.attn, **kwargs)
+            block.ff = NunchakuKrea2SwiGLU(block.ff, **kwargs)
+        return self
+
+    @classmethod
+    @utils.validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs):
+        """Load a quantized Krea 2 transformer from a single Nunchaku-format safetensors file.
+
+        The file carries the quantized weights for the 28 blocks plus a ``quantization_config`` metadata
+        entry (rank, precision). Non-quantized modules (img_in, time_embed, text_fusion, txt_in,
+        final_layer) are loaded in ``torch_dtype``.
+        """
+        device = kwargs.get("device", "cpu")
+        if kwargs.get("offload", False):
+            raise NotImplementedError("Offload is not supported for Krea2Transformer2DModel")
+
+        torch_dtype = kwargs.get("torch_dtype", torch.bfloat16)
+        if isinstance(pretrained_model_name_or_path, str):
+            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+        assert pretrained_model_name_or_path.is_file() or pretrained_model_name_or_path.name.endswith(
+            (".safetensors", ".sft")
+        ), "Only safetensors are supported"
+
+        transformer, model_state_dict, metadata = cls._build_model(pretrained_model_name_or_path, **kwargs)
+        quantization_config = json.loads(metadata.get("quantization_config", "{}"))
+        rank = quantization_config.get("rank", 32)
+        transformer = transformer.to(torch_dtype)
+
+        # Metadata is authoritative when the converter wrote it; otherwise fall back to the
+        # hardware default, as the other ports do. Silently defaulting to int4 would load an
+        # nvfp4 checkpoint with the wrong kernels on Blackwell.
+        precision = quantization_config.get("precision")
+        if precision is None:
+            precision = get_precision()
+        if precision == "fp4":
+            precision = "nvfp4"
+
+        print(f"quantization_config: {quantization_config}, rank={rank}, precision={precision}")
+
+        transformer._patch_model(precision=precision, rank=rank)
+        transformer = transformer.to_empty(device=device)
+
+        patch_scale_key(transformer, model_state_dict)
+        if torch_dtype == torch.float16:
+            convert_fp16(transformer, model_state_dict)
+        transformer.load_state_dict(model_state_dict)
+        return transformer
diff --git a/tests/v1/krea2/__init__.py b/tests/v1/krea2/__init__.py
diff --git a/tests/v1/krea2/test_krea2_turbo.py b/tests/v1/krea2/test_krea2_turbo.py
