ComfyUI "Support"
#4
by onixxexxd5555LOAF - opened
So seeing that Nunchaku project isn't particularly alive right now (Last commit was 2 months ago, 3 months for the ComfyUI extension) I decided to add ComfyUI support myself since it's very unlikely that anyone would merge the PR and then update the extension.
To clarify I don't properly understand what I am doing, the code is unreviewed AI slop (I just checked if it works or not and kept feeding it error messages a lot of times until it worked), it's most likely buggy, performs suboptimal and ultimately provided as-is.
It requires you to install the custom node normally and then manually add PR contents to your venv/lib/python3.13/site-packages/nunchaku/ (add nunchaku/models/transformers/transformer_chroma.py and then merge the contents of 3 init.pys)
Then go to ComfyUI-nunchaku folder and do these:
- Add this before NODE_DISPLAY_NAME_MAPPINGS in init.py:
try:
from .nodes.models.chroma import NunchakuChromaDiTLoader
NODE_CLASS_MAPPINGS["NunchakuChromaDiTLoader"] = NunchakuChromaDiTLoader
except ImportError:
logger.exception("Node `NunchakuChromaDiTLoader` import failed:")
- Go to wrappers folder, create a file named chroma.py and paste this:
"""
ComfyChromaWrapper
==================
Bridges the ComfyUI diffusion-model calling convention:
forward(x, timestep, context, guidance, control, transformer_options, **kwargs)
to the NunchakuChromaTransformer2dModel forward signature:
forward(hidden_states, encoder_hidden_states, timestep,
img_ids, txt_ids, attention_mask, return_dict)
Key design decisions
--------------------
* **guidance is dropped** — Chroma has `guidance_embeds: false`; the distillation
guidance is handled entirely inside the nunchaku model via
`ChromaCombinedTimestepTextProjEmbeddings` + `ChromaApproximator`. Passing it
externally would cause a TypeError and is semantically wrong.
* **timestep is passed as-is** — The nunchaku forward already does `timestep * 1000`
internally, matching the ComfyUI flow-model convention of [0, 1] sigmas.
* **Patch preparation** mirrors `comfy/ldm/chroma/model.py` `_forward` exactly
(pad → rearrange → build img_ids / txt_ids → call → rearrange back).
"""
import torch
from torch import nn
from einops import rearrange, repeat
import comfy.ldm.common_dit
import comfy.patcher_extension
class ComfyChromaWrapper(nn.Module):
"""
Wraps a ``NunchakuChromaTransformer2dModel`` so it can be used as
``model.diffusion_model`` inside a ComfyUI ``ModelPatcher``.
Parameters
----------
transformer : NunchakuChromaTransformer2dModel
The already-loaded nunchaku Chroma model.
config : dict
Minimal config dict; only ``patch_size`` is consumed here.
ctx_for_copy : dict, optional
Arbitrary context stored for debugging / potential copy operations.
"""
def __init__(self, transformer, config: dict, ctx_for_copy: dict = None):
super().__init__()
self.transformer = transformer
self.patch_size = int(config.get("patch_size", 2))
self.ctx_for_copy = ctx_for_copy or {}
# ------------------------------------------------------------------
# Attributes required by ComfyUI's model_base / model_patcher
# ------------------------------------------------------------------
@property
def img_in(self):
"""
``comfy/model_base.py`` ``concat_cond`` reads
``self.diffusion_model.img_in.weight.shape[1]`` to infer input
channels. Chroma's nunchaku model calls this layer ``x_embedder``,
so we proxy it here.
"""
return self.transformer.x_embedder
@property
def patch_size(self) -> int: # type: ignore[override]
return self._patch_size
@patch_size.setter
def patch_size(self, value: int):
self._patch_size = value
@property
def dtype(self) -> torch.dtype:
"""
``comfy/model_base.py`` line 227 does::
return self.diffusion_model.dtype
PyTorch ``nn.Module`` does not expose ``.dtype`` by default, so we
derive it from the first parameter of the inner nunchaku model.
Falls back to ``bfloat16`` if the model has no parameters yet
(e.g., during meta-device initialisation).
"""
try:
return next(self.transformer.parameters()).dtype
except StopIteration:
return torch.bfloat16
# ------------------------------------------------------------------
# Public forward — routes through ComfyUI's WrapperExecutor so that
# any registered diffusion-model patches (ControlNet, etc.) still work.
# ------------------------------------------------------------------
def forward(
self,
x: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
guidance: torch.Tensor = None, # accepted but intentionally ignored
control=None,
transformer_options: dict = {},
**kwargs,
) -> torch.Tensor:
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(
comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL,
transformer_options,
),
).execute(x, timestep, context, guidance, control, transformer_options, **kwargs)
# ------------------------------------------------------------------
# Internal forward — actual computation
# ------------------------------------------------------------------
def _forward(
self,
x: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
guidance: torch.Tensor = None, # dropped — Chroma is distilled
control=None, # ControlNet not supported by nunchaku Chroma
transformer_options: dict = {},
**kwargs,
) -> torch.Tensor:
bs, c, h, w = x.shape
patch_size = self.patch_size
# ── 1. Pad spatial dims to patch-size multiples ──────────────
x_padded = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
# ── 2. Patchify: (B, C, H, W) → (B, H*W, C*ph*pw) ───────────
img = rearrange(
x_padded,
"b c (h ph) (w pw) -> b (h w) (c ph pw)",
ph=patch_size,
pw=patch_size,
)
# ── 3. Build positional ID grids (same as ComfyUI Chroma) ─────
h_len = (h + (patch_size // 2)) // patch_size
w_len = (w + (patch_size // 2)) // patch_size
img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(
0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype
).unsqueeze(1)
img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(
0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype
).unsqueeze(0)
# (h_len, w_len, 3) → (B, h_len*w_len, 3)
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
# txt_ids: all zeros (text has no spatial position)
txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
# ── 4. Optional attention mask from sampler kwargs ────────────
attention_mask = kwargs.get("attention_mask", None)
# ── 5. Call the nunchaku model ────────────────────────────────
#
# What the nunchaku model does internally with these inputs:
# hidden_states → x_embedder → patch tokens
# encoder_hidden_states → context_embedder → text tokens
# timestep * 1000 → time_text_embed + distilled_guidance_layer
# → pooled modulation vectors (shape B×344×inner_dim)
# img_ids, txt_ids → pos_embed (FluxPosEmbed / rotary)
#
# NOTE: guidance is NOT forwarded. The nunchaku Chroma forward()
# has no `guidance` parameter — distillation conditioning is
# computed solely from `timestep` inside `time_text_embed`.
result = self.transformer(
hidden_states=img,
encoder_hidden_states=context,
timestep=timestep,
img_ids=img_ids,
txt_ids=txt_ids,
attention_mask=attention_mask,
return_dict=False,
)
out = result[0] # Transformer2DModelOutput.sample when return_dict=True,
# or tuple[Tensor] when return_dict=False
# ── 6. Un-patchify: (B, H*W, C*ph*pw) → (B, C, H, W) ────────
out = rearrange(
out,
"b (h w) (c ph pw) -> b c (h ph) (w pw)",
h=h_len,
w=w_len,
ph=patch_size,
pw=patch_size,
)
# Crop back to the original (un-padded) spatial size
return out[:, :, :h, :w]
- Go to nodes/models, create a file named chroma.py and then paste this:
"""
NunchakuChromaDiTLoader
=======================
ComfyUI node for loading a nunchaku-quantized Chroma DiT model.
Usage
-----
Place the nunchaku Chroma ``.safetensors`` file under
``ComfyUI/models/diffusion_models/`` and select it from the node's
``model_path`` dropdown.
Architecture notes
------------------
* Chroma has ``guidance_embeds: false`` — no guidance stream.
* Chroma uses a ``ChromaApproximator`` (distilled guidance layer) driven by
``ChromaCombinedTimestepTextProjEmbeddings`` rather than Flux's
``CombinedTimestepGuidanceTextProjEmbeddings``.
* The nunchaku backend (``transformer_chroma.py``) owns ALL of this logic;
the wrapper here only handles patch preparation and argument translation.
What this node does
-------------------
1. Calls ``NunchakuChromaTransformer2dModel.from_pretrained(path)`` to load
the quantized weights (rank, precision, and config are auto-detected from
safetensors metadata).
2. Wraps the result in a ``ComfyChromaWrapper``.
3. Creates a ComfyUI ``Chroma`` model-config shell with
``disable_unet_model_creation=True`` (so the expensive full bf16 backbone
is never built), sets ``.diffusion_model`` to our wrapper, and returns
a ``ModelPatcher``.
"""
import gc
import logging
import os
import comfy.model_management
import comfy.model_patcher
import torch
from nunchaku.models.transformers.transformer_chroma import NunchakuChromaTransformer2dModel
try:
from nunchaku.utils import is_turing
except ImportError:
# Older nunchaku — assume not Turing
def is_turing(device: str) -> bool: # type: ignore[misc]
return False
try:
from comfy.supported_models import Chroma as ChromaModelConfig
except ImportError as exc:
raise ImportError(
"Could not import `Chroma` from `comfy.supported_models`. "
"Make sure your ComfyUI version includes Chroma support."
) from exc
# Reuse the ComfyChromaWrapper — adjust the import path to match your
# custom node's package structure (mirrors how flux.py imports ComfyFluxWrapper).
from ...wrappers.chroma import ComfyChromaWrapper
from ..utils import get_filename_list, get_full_path_or_raise
# ---------------------------------------------------------------------------
# Logging
# ---------------------------------------------------------------------------
log_level = os.getenv("LOG_LEVEL", "INFO").upper()
logging.basicConfig(
level=getattr(logging, log_level, logging.INFO),
format="%(asctime)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Hardcoded ComfyUI model config for Chroma
#
# Mirrors the structure used by the Flux node's comfy_config JSON.
# `disable_unet_model_creation` prevents the full bf16 Chroma backbone from
# being instantiated (we swap in the nunchaku wrapper instead).
# ---------------------------------------------------------------------------
_CHROMA_COMFY_CONFIG = {
"model_class": "Chroma",
"model_config": {
"image_model": "chroma",
"disable_unet_model_creation": True,
# ── Keys read by comfy/model_base.py (concat_cond and friends) ──
# Primary path: diffusion_model.img_in.weight (proxied via wrapper)
# Fallback path: unet_config[key] (these values)
# All taken from Chroma's config.json; out_channels null → 64.
"in_channels": 64,
"out_channels": 64,
"patch_size": 2,
# Kept for completeness — some BaseFlux helpers read these too.
"num_heads": 24,
"hidden_size": 3072, # num_attention_heads * attention_head_dim
"depth": 19, # num_layers
"depth_single_blocks": 38, # num_single_layers
"context_in_dim": 4096, # joint_attention_dim
"axes_dim": [16, 56, 56],
"theta": 10000,
"qkv_bias": True,
},
}
# ---------------------------------------------------------------------------
# Node class
# ---------------------------------------------------------------------------
class NunchakuChromaDiTLoader:
"""
ComfyUI node — Nunchaku Chroma DiT Loader.
Returns a ``MODEL`` patcher compatible with the standard KSampler.
"""
def __init__(self):
self.transformer = None
self.model_path = None
self.device = None
self.cpu_offload = None
self.data_type = None
# ------------------------------------------------------------------
# ComfyUI interface
# ------------------------------------------------------------------
@classmethod
def INPUT_TYPES(cls):
safetensor_files = get_filename_list("diffusion_models")
ngpus = max(torch.cuda.device_count(), 1)
# Turing GPUs (RTX 20xx) do not support bfloat16
all_turing = all(is_turing(f"cuda:{i}") for i in range(torch.cuda.device_count()))
dtype_options = ["float16"] if all_turing else ["bfloat16", "float16"]
return {
"required": {
"model_path": (
safetensor_files,
{
"tooltip": (
"Nunchaku Chroma safetensors file "
"(place under models/diffusion_models/)."
)
},
),
"cpu_offload": (
["auto", "enable", "disable"],
{
"default": "auto",
"tooltip": (
"'auto' enables CPU offload when GPU VRAM < 14 GB. "
"Required for 8 GB cards."
),
},
),
"device_id": (
"INT",
{
"default": 0,
"min": 0,
"max": ngpus - 1,
"step": 1,
"display": "number",
"lazy": True,
"tooltip": "CUDA device index.",
},
),
"data_type": (
dtype_options,
{
"default": dtype_options[0],
"tooltip": (
"Non-quantized layers run at this dtype. "
"Use float16 on RTX 20xx."
),
},
),
}
}
RETURN_TYPES = ("MODEL",)
FUNCTION = "load_model"
CATEGORY = "Nunchaku"
TITLE = "Nunchaku Chroma DiT Loader"
# ------------------------------------------------------------------
# Main loading logic
# ------------------------------------------------------------------
def load_model(
self,
model_path: str,
cpu_offload: str,
device_id: int,
data_type: str,
):
device = torch.device(f"cuda:{device_id}")
if device_id >= torch.cuda.device_count():
raise ValueError(
f"device_id={device_id} is invalid; "
f"only {torch.cuda.device_count()} GPU(s) available."
)
model_path = get_full_path_or_raise("diffusion_models", model_path)
# ── GPU memory check for auto-offload ──────────────────────────
props = torch.cuda.get_device_properties(device_id)
gpu_mem_mib = props.total_memory / (1024 ** 2)
logger.debug(f"GPU {device_id} ({props.name}) VRAM: {gpu_mem_mib:.0f} MiB")
if cpu_offload == "auto":
cpu_offload_enabled = gpu_mem_mib < 14336 # 14 GiB threshold
logger.debug(
f"cpu_offload=auto → {'enabled' if cpu_offload_enabled else 'disabled'}"
)
else:
cpu_offload_enabled = cpu_offload == "enable"
torch_dtype = torch.float16 if data_type == "float16" else torch.bfloat16
# ── Load the nunchaku model (cached if params unchanged) ───────
need_reload = (
self.model_path != model_path
or self.device != device
or self.cpu_offload != cpu_offload_enabled
or self.data_type != data_type
)
if need_reload:
self._unload_existing(device)
logger.info(f"Loading nunchaku Chroma from: {model_path}")
self.transformer = NunchakuChromaTransformer2dModel.from_pretrained(
model_path,
offload=cpu_offload_enabled,
device=str(device),
torch_dtype=torch_dtype,
verbose=True,
)
self.model_path = model_path
self.device = device
self.cpu_offload = cpu_offload_enabled
self.data_type = data_type
transformer = self.transformer
patch_size = 2 #just trust me
# ── Build the ComfyUI model shell ──────────────────────────────
#
# We instantiate `ChromaModelConfig` (comfy.supported_models.Chroma)
# with `disable_unet_model_creation=True` so that no bf16 backbone
# is allocated. Then we swap in our wrapper as `.diffusion_model`.
comfy_model_config_dict = _CHROMA_COMFY_CONFIG["model_config"].copy()
comfy_model_config_dict.setdefault("disable_unet_model_creation", True)
model_config = ChromaModelConfig(comfy_model_config_dict)
model_config.set_inference_dtype(torch_dtype, None)
model_config.custom_operations = None
# `get_model({})` instantiates model_base.Chroma; with
# disable_unet_model_creation=True the internal diffusion_model is
# either None or a stub.
comfy_model = model_config.get_model({})
# Replace the diffusion backbone with our nunchaku wrapper
comfy_model.diffusion_model = ComfyChromaWrapper(
transformer=transformer,
config={"patch_size": patch_size},
ctx_for_copy={
"comfy_config": _CHROMA_COMFY_CONFIG,
"model_config": model_config,
"device": device,
"device_id": device_id,
},
)
patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, device_id)
return (patcher,)
# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def _unload_existing(self, target_device: torch.device):
"""Free the previously loaded transformer to reclaim VRAM."""
if self.transformer is None:
return
try:
model_size = comfy.model_management.module_size(self.transformer)
except Exception:
model_size = 0
transformer = self.transformer
self.transformer = None
try:
transformer.to("cpu")
except Exception:
pass
del transformer
gc.collect()
comfy.model_management.cleanup_models_gc()
comfy.model_management.soft_empty_cache()
if model_size > 0:
comfy.model_management.free_memory(model_size, target_device)
Relaunch ComfyUI.
Example workflow here:
It runs with twice the speed of q8 on my system. I only tested int4, because it's the only version I can test.
Open to someone who knows what they are doing providing feedback. Thank you tonera for making the quant.
Great job!