[Fix] comfyui: Fix GGUF detection, import guard and VRAM accounting in SGLDiffusion plugin

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/__init__.py

@@ -2,14 +2,24 @@
 ComfyUI SGLang Diffusion nodes package.
 """
 
+import logging
+
+logger = logging.getLogger(__name__)
+
 WEB_DIRECTORY = "./web"
 
 try:
     from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
     __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
-except ImportError:
-    # ComfyUI dependencies not available (e.g., in test environment)
+except ImportError as exc:
+    # Missing ComfyUI or SGLang deps; log why, or every node vanishes silently.
+    logger.error(
+        "ComfyUI_SGLDiffusion failed to register nodes: %r. If this is a "
+        "real ComfyUI install, install the diffusion extras with "
+        "'pip install sglang[diffusion]'.",
+        exc,
+    )
     NODE_CLASS_MAPPINGS = {}
     NODE_DISPLAY_NAME_MAPPINGS = {}
     __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py

@@ -104,48 +104,6 @@ def _looks_like_gguf(path: str) -> bool:
     )
 
 
-def _h3_detect_companion(gguf_ref: str) -> str:
-    """ComfyUI detect needs H3 safetensors keys; GGUF only overrides the DiT."""
-    name = os.path.basename(gguf_ref).lower()
-    prefer = []
-    if "ref2va" in name:
-        prefer.append("minimax_h3_ref2va_bf16.safetensors")
-    prefer.extend(
-        [
-            "minimax_h3_fl2va_bf16.safetensors",
-            "minimax_h3_ref2va_bf16.safetensors",
-        ]
-    )
-    folders = (
-        [os.path.dirname(os.path.abspath(gguf_ref))]
-        if os.path.dirname(gguf_ref)
-        else []
-    )
-    try:
-        import folder_paths
-
-        folders.extend(folder_paths.get_folder_paths("diffusion_models"))
-    except Exception:
-        pass
-    seen: set[str] = set()
-    for folder in folders:
-        if not folder:
-            continue
-        for candidate in prefer:
-            path = os.path.join(folder, candidate)
-            if path in seen:
-                continue
-            seen.add(path)
-            if os.path.isfile(path):
-                return path
-    raise ValueError(
-        "GGUF transformer needs a MiniMax-H3 BF16 safetensors in "
-        "models/diffusion_models for ComfyUI architecture detect "
-        f"(looked for {prefer}). Keep --model-path / unet_name on the "
-        "BF16 file and pass the GGUF via transformer_weights_path."
-    )
-
-
 _NATIVE_MODEL_TYPES = {"FastH3Pipeline": "fast_h3", "VDNH3Pipeline": "vdn_h3"}
 NATIVE_MODEL_TYPES = tuple(_NATIVE_MODEL_TYPES.values())
 
@@ -210,10 +168,16 @@ def _h3_unet_config_from_native(model_dir: str) -> dict:
 
 def _load_state_dict_for_detection(model_path: str):
     if _looks_like_gguf(model_path):
-        raise ValueError(
-            f"Cannot detect architecture from GGUF {model_path!r}. "
-            "Use a MiniMax-H3 BF16 safetensors for detect."
+        # Same header-only trick as safetensors: tensor names and logical
+        # (unpacked) shapes are all detect_unet_config looks at.
+        from sglang.multimodal_gen.runtime.loader.gguf_weights import (
+            read_gguf_tensor_meta,
         )
+
+        return {
+            name: _HeaderTensor(meta.logical_shape)
+            for name, meta in read_gguf_tensor_meta(model_path).items()
+        }
     if model_path.endswith(".safetensors"):
         from safetensors import safe_open
 
@@ -527,11 +491,11 @@ def get_comfyui_model(self, model_path: str, model_options: dict = None):
             if new_sd is not None:  # diffusers mmdit
                 model_config = model_detection.model_config_from_unet(new_sd, "")
                 if model_config is None:
-                    return None
+                    raise ValueError(f"No ComfyUI model config for {model_path!r}")
             else:  # diffusers unet
                 model_config = model_detection.model_config_from_diffusers_unet(sd)
                 if model_config is None:
-                    return None
+                    raise ValueError(f"No ComfyUI model config for {model_path!r}")
 
                 diffusers_keys = unet_to_diffusers(model_config.unet_config)
                 new_sd = {}
@@ -596,8 +560,6 @@ def load_model(
         else:
             sgld_options.pop("transformer_weights_path", None)
         detect_path = model_path
-        if _looks_like_gguf(model_path):
-            detect_path = _h3_detect_companion(sgld_options["transformer_weights_path"])
         gather_options = {
             "model_path": detect_path,
             "model_options": model_options,

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/minimax_h3.py

@@ -10,28 +10,38 @@
 
 import torch
 
-from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
-from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams
-from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
-    MiniMaxH3ReleaseMetadata,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.comfyui_step import (
-    DEFAULT_SIGMA_SHIFT_AUDIO,
-    DEFAULT_SIGMA_SHIFT_VIDEO,
-    serialize_comfyui_layout,
-    time_shift_sigma,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
-    minimax_h3_time_shift_sigmas,
-)
-
 from .adapter import ComfyUIModelAdapter, PackedForward
 from .base import SGLDiffusionExecutor
 
-# H3 pins loop/CFG fields with init=False; ComfyUI owns those.
-_H3_PINNED_SAMPLING_FIELDS = frozenset(
-    f.name for f in fields(MiniMaxH3SamplingParams) if not f.init
-)
+# Same guard as base.py: generator.py imports this module unconditionally, so
+# a missing runtime must surface on first use instead of hiding every node.
+try:
+    from sglang.multimodal_gen.configs.sample.minimax_h3 import MiniMaxH3SamplingParams
+    from sglang.multimodal_gen.configs.sample.minimax_h3_vdn import VDNH3SamplingParams
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
+        MiniMaxH3ReleaseMetadata,
+    )
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.stages.comfyui_step import (
+        DEFAULT_SIGMA_SHIFT_AUDIO,
+        DEFAULT_SIGMA_SHIFT_VIDEO,
+        serialize_comfyui_layout,
+        time_shift_sigma,
+    )
+    from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import (
+        minimax_h3_time_shift_sigmas,
+    )
+except ImportError as exc:
+    _H3_RUNTIME_IMPORT_ERROR: ImportError | None = exc
+    MiniMaxH3SamplingParams = VDNH3SamplingParams = MiniMaxH3ReleaseMetadata = None
+    DEFAULT_SIGMA_SHIFT_AUDIO = DEFAULT_SIGMA_SHIFT_VIDEO = 0.0
+    serialize_comfyui_layout = time_shift_sigma = minimax_h3_time_shift_sigmas = None
+    _H3_PINNED_SAMPLING_FIELDS: frozenset[str] = frozenset()
+else:
+    _H3_RUNTIME_IMPORT_ERROR = None
+    # H3 pins loop/CFG fields with init=False; ComfyUI owns those.
+    _H3_PINNED_SAMPLING_FIELDS = frozenset(
+        f.name for f in fields(MiniMaxH3SamplingParams) if not f.init
+    )
 
 
 def drop_h3_pinned_sampling_fields(kwargs: dict[str, Any]) -> dict[str, Any]:
@@ -223,6 +233,13 @@ class MiniMaxH3Executor(SGLDiffusionExecutor):
     # ComfyUI MiniMaxH3.extra_conds reads this when a denoise_mask is present.
     patch_size = (1, 2, 2)
 
+    def __init__(self, *args, **kwargs):
+        if _H3_RUNTIME_IMPORT_ERROR is not None:
+            raise RuntimeError(
+                "SGLang diffusion runtime failed to import"
+            ) from _H3_RUNTIME_IMPORT_ERROR
+        super().__init__(*args, **kwargs)
+
     def preprocess_text_embeds(self, text_states):
         """ComfyUI extra_conds runs this once per sample.
 
@@ -368,4 +385,7 @@ def forward(self, x, timestep, context, **kwargs):
 
 class VDNH3Executor(_DistilledH3Executor):
     adapter_cls = VDNH3Adapter
-    expected_sigmas = distilled_h3_sigmas("vdn_h3")
+    # Without the runtime __init__ raises anyway; keep the import working.
+    expected_sigmas = (
+        distilled_h3_sigmas("vdn_h3") if _H3_RUNTIME_IMPORT_ERROR is None else []
+    )

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/nodes.py

@@ -13,14 +13,19 @@
 from .options import SGLDOptions
 
 
-def _enable_gguf_in_diffusion_models() -> None:
-    """Let SGLDUNETLoader list ``.gguf`` DiTs. Quantization lives in the file."""
-    entry = folder_paths.folder_names_and_paths.get("diffusion_models")
-    if entry is not None and len(entry) >= 2 and isinstance(entry[1], set):
-        entry[1].add(".gguf")
+def _list_unet_names_including_gguf() -> list[str]:
+    """SGLDUNETLoader's file list plus ``.gguf`` DiTs.
+
+    Kept local: registering ``.gguf`` in ``folder_paths`` would make ComfyUI's
+    stock loader offer files it cannot load.
+    """
+    names = set(folder_paths.get_filename_list("diffusion_models"))
+    for folder in folder_paths.get_folder_paths("diffusion_models"):
+        files, _ = folder_paths.recursive_search(folder)
+        names.update(folder_paths.filter_files_extensions(files, {".gguf"}))
+    return sorted(names)
 
 
-_enable_gguf_in_diffusion_models()
 from .utils import (
     convert_b64_to_tensor_image,
     convert_video_to_comfy_video,
@@ -138,7 +143,7 @@ def __init__(self):
     def INPUT_TYPES(s):
         return {
             "required": {
-                "unet_name": (folder_paths.get_filename_list("diffusion_models"),),
+                "unet_name": (_list_unet_names_including_gguf(),),
                 "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
             },
             "optional": {

python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/flux.py

@@ -35,6 +35,12 @@ def _build_dit_config(server_args: ServerArgs) -> FluxConfig:
 
 
 def _checkpoint_has_guidance_embeds(model_path: str) -> bool:
+    if model_path.endswith(".gguf"):
+        from sglang.multimodal_gen.runtime.loader.gguf_weights import (
+            read_gguf_tensor_meta,
+        )
+
+        return "guidance_in.in_layer.weight" in read_gguf_tensor_meta(model_path)
     with safe_open(model_path, framework="pt", device="cpu") as checkpoint:
         return "guidance_in.in_layer.weight" in checkpoint.keys()
 

python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/spec.py

@@ -14,8 +14,9 @@
 from __future__ import annotations
 
 import os
+from collections import defaultdict
 from collections.abc import Callable, Iterator
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -27,6 +28,7 @@
     inspect_minimax_h3_safetensors,
     resolve_minimax_h3_checkpoint_quantization,
 )
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
 from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
@@ -115,8 +117,78 @@ def get_registered_comfyui_pipeline_names() -> list[str]:
 
 
 def is_comfyui_single_file(model_path: str) -> bool:
-    """ComfyUI ships DiTs as one safetensors file with no model_index.json."""
-    return os.path.isfile(model_path) and model_path.endswith(".safetensors")
+    """ComfyUI ships DiTs as one safetensors or GGUF file with no model_index.json."""
+    return os.path.isfile(model_path) and model_path.endswith((".safetensors", ".gguf"))
+
+
+def _gguf_runtime_layout(tensor_meta: dict, mapping: dict) -> tuple[dict, dict]:
+    """Key GGUF metadata by runtime names; return it with the fused-shard plan.
+
+    Shards of one fused runtime layer (Z-Image w1/w3 -> w13) keep their own
+    entries and are concatenated on load, under the fused name quant lookup uses.
+    """
+    from sglang.multimodal_gen.runtime.loader.gguf_weights import (
+        remap_gguf_tensor_meta,
+    )
+
+    map_name = get_param_names_mapping(mapping)
+    plain, fused = {}, defaultdict(dict)
+    for name, meta in tensor_meta.items():
+        target, index, _ = map_name(name)
+        if index is None:
+            plain[name] = meta
+        else:
+            fused[target][index] = (name, meta)
+
+    layout = remap_gguf_tensor_meta(plain, lambda name: map_name(name)[0])
+    shard_plan = {}
+    for target, parts in fused.items():
+        metas = [parts[i][1] for i in sorted(parts)]
+        first = metas[0]
+        if len({(m.ggml_type, m.is_packed, m.logical_shape[1:]) for m in metas}) > 1:
+            raise ValueError(f"GGUF shards of {target} differ in type or width")
+        param_name = (
+            f"{target.removesuffix('.weight')}.qweight" if first.is_packed else target
+        )
+        layout[target] = replace(
+            first,
+            logical_shape=(
+                sum(m.logical_shape[0] for m in metas),
+                *first.logical_shape[1:],
+            ),
+            stored_shape=(
+                sum(m.stored_shape[0] for m in metas),
+                *first.stored_shape[1:],
+            ),
+            param_name=param_name,
+        )
+        for index, (name, meta) in parts.items():
+            layout[name] = meta
+            shard_plan[meta.param_name] = (param_name, index, len(parts))
+    return layout, shard_plan
+
+
+def _runtime_gguf_weights(weights, layout: dict, shard_plan: dict, quant_config):
+    """Fuse planned shards; dequantize packed tensors whose layer built without
+    GGUF (e.g. Z-Image adaLN ReplicatedLinear). Runs after model construction,
+    so quant_config.selected is final."""
+    import gguf
+
+    packed_types = {m.param_name: m.weight_type for m in layout.values() if m.is_packed}
+    pending = defaultdict(dict)
+    for name, tensor in weights:
+        if name in shard_plan:
+            name, index, total = shard_plan[name]
+            pending[name][index] = tensor
+            if len(pending[name]) < total:
+                continue
+            parts = pending.pop(name)
+            tensor = torch.cat([parts[i] for i in range(total)])
+        prefix = name.removesuffix(".qweight")
+        if name in packed_types and prefix not in quant_config.selected:
+            value = gguf.dequantize(tensor.numpy(), packed_types[name])
+            name, tensor = f"{prefix}.weight", torch.from_numpy(value).bfloat16()
+        yield name, tensor
 
 
 def _load_comfyui_gguf_transformer(
@@ -146,7 +218,10 @@ def _load_comfyui_gguf_transformer(
             [gguf_file], str(server_args.model_variant or "fl2va")
         )
 
-    tensor_meta = read_gguf_tensor_meta(gguf_file)
+    # Quant lookup and weight loading must both see runtime names.
+    tensor_meta, shard_plan = _gguf_runtime_layout(
+        read_gguf_tensor_meta(gguf_file), mapping
+    )
     curve = tensor_meta.get("adaln_t_table")
     if curve is not None:
         if curve.is_quantized or len(curve.logical_shape) != 2:
@@ -188,7 +263,12 @@ def _load_comfyui_gguf_transformer(
             reduce_dtype=torch.float32,
             output_dtype=None,
             strict=False,
-            weights_iterator=gguf_weights_iterator(gguf_file, tensor_meta),
+            weights_iterator=_runtime_gguf_weights(
+                gguf_weights_iterator(gguf_file, tensor_meta),
+                tensor_meta,
+                shard_plan,
+                quant_config,
+            ),
         )
     finally:
         model_cls.param_names_mapping = original_mapping
@@ -283,6 +363,9 @@ def load_comfyui_transformer(
     )
 
     gguf_file = resolve_transformer_gguf_to_load(server_args, "transformer")
+    if gguf_file is None and model_path.endswith(".gguf"):
+        # A GGUF unet picked in ComfyUI is the model path itself.
+        gguf_file = model_path
     if gguf_file is not None:
         model = _load_comfyui_gguf_transformer(
             model_cls=model_cls,

python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/zimage.py

@@ -46,6 +46,8 @@ def _convert_weights(weights: WeightIterator, dit_config: Any) -> WeightIterator
 
 
 _PARAM_NAMES_MAPPING = {
+    # GGUF skips _convert_weights; its fused qkv already matches to_qkv.
+    r"(.*)\.attention\.qkv\.(weight|bias)$": (r"\1.attention.to_qkv.\2", None, None),
     r"(.*)\.attention\.k_norm\.weight$": (r"\1.attention.norm_k.weight", None, None),
     r"(.*)\.attention\.q_norm\.weight$": (r"\1.attention.norm_q.weight", None, None),
     r"(.*)\.attention\.out\.weight$": (r"\1.attention.to_out.0.weight", None, None),