[diffusion] comfyui: consolidate FluxAdapter correctness fixes

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

@@ -5,6 +5,9 @@
 from .adapter import ComfyUIModelAdapter, PackedForward
 from .base import SGLDiffusionExecutor
 
+# Width of Flux.1's CLIP-L pooled vector (vec_in_dim / pooled_projection_dim).
+_FLUX_POOLED_DIM = 768
+
 
 def _flux_guidance_scale(guidance) -> float:
     if guidance is None:
@@ -14,37 +17,73 @@ def _flux_guidance_scale(guidance) -> float:
     return float(guidance)
 
 
+def _pad_to_patch_size(latents: torch.Tensor, patch_size: int) -> torch.Tensor:
+    # Equivalent to comfy.ldm.common_dit.pad_to_patch_size, inlined so this
+    # module keeps loading without ComfyUI installed (see executors/__init__.py).
+    pad_h = (patch_size - latents.shape[-2] % patch_size) % patch_size
+    pad_w = (patch_size - latents.shape[-1] % patch_size) % patch_size
+    return torch.nn.functional.pad(latents, (0, pad_w, 0, pad_h), mode="circular")
+
+
 class FluxAdapter(ComfyUIModelAdapter):
     model_types = ("flux",)
     pipeline_class_name = "FluxPipeline"
+    patch_size = 2
 
     def pack(
-        self, x, timestep, context, y=None, guidance=None, **kwargs
+        self,
+        x,
+        timestep,
+        context,
+        y=None,
+        guidance=None,
+        ref_latents=None,
+        control=None,
+        **kwargs,
     ) -> PackedForward:
-        packed = self._pack_latents(x)
+        if ref_latents:
+            raise ValueError(
+                "FluxAdapter does not support FLUX.1 Kontext reference latents "
+                "(ComfyUI's ReferenceLatent node). Running the workflow would "
+                "silently drop the reference images and produce a plain "
+                "text-to-image result instead of a Kontext edit."
+            )
+        if control is not None:
+            raise ValueError(
+                "FluxAdapter does not support ControlNet conditioning. Running "
+                "the workflow would silently drop the ControlNet and produce an "
+                "unconditioned text-to-image result."
+            )
+        if y is None:
+            # Native ComfyUI Flux zero-fills a missing pooled vector.
+            y = torch.zeros(
+                x.shape[0], _FLUX_POOLED_DIM, device=context.device, dtype=context.dtype
+            )
+        packed, padded_height, padded_width = self._pack_latents(x)
         t5_seq = int(context.shape[-2]) if context.ndim >= 2 else int(context.shape[0])
-        clip_batch = int(y.shape[0]) if y is not None else 1
+        clip_batch = int(y.shape[0])
         return PackedForward(
             latents=packed,
             timesteps=timestep * 1000.0,
             prompt_embeds=[y, context],
             prompt_seq_lens=[[clip_batch], [t5_seq]],
             pooled_embeds=[y],
-            height=x.shape[-2] * 8,
-            width=x.shape[-1] * 8,
+            height=padded_height * 8,
+            width=padded_width * 8,
             guidance_scale=_flux_guidance_scale(guidance),
             unpack_ctx={
-                "height": x.shape[-2],
-                "width": x.shape[-1],
+                "height": padded_height,
+                "width": padded_width,
                 "channels": x.shape[1],
             },
         )
 
     def unpack(self, noise_pred, packed, x):
         ctx = packed.unpack_ctx
-        return self._unpack_latents(
+        latents = self._unpack_latents(
             noise_pred, ctx["height"], ctx["width"], ctx["channels"]
-        ).to(x.device)
+        )
+        return latents[:, :, : x.shape[-2], : x.shape[-1]].to(x.device)
 
     @staticmethod
     def _unpack_latents(latents, height, width, channels):
@@ -53,16 +92,18 @@ def _unpack_latents(latents, height, width, channels):
         latents = latents.permute(0, 3, 1, 4, 2, 5)
         return latents.reshape(batch_size, channels, height, width)
 
-    @staticmethod
-    def _pack_latents(latents):
+    @classmethod
+    def _pack_latents(cls, latents):
+        latents = _pad_to_patch_size(latents, cls.patch_size)
         batch_size, num_channels_latents, height, width = latents.shape
         latents = latents.view(
             batch_size, num_channels_latents, height // 2, 2, width // 2, 2
         )
         latents = latents.permute(0, 2, 4, 1, 3, 5)
-        return latents.reshape(
+        latents = latents.reshape(
             batch_size, (height // 2) * (width // 2), num_channels_latents * 4
         )
+        return latents, height, width
 
 
 class FluxExecutor(SGLDiffusionExecutor):

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

@@ -5,6 +5,7 @@
 from typing import Any
 
 import torch
+from safetensors import safe_open
 
 from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
 from sglang.multimodal_gen.runtime.loader.comfyui_checkpoints.spec import (
@@ -24,11 +25,20 @@ def _build_dit_config(server_args: ServerArgs) -> FluxConfig:
         dit_config = FluxConfig()
         server_args.pipeline_config.dit_config = dit_config
 
-    # ComfyUI Flux checkpoints always carry guidance_in weights.
-    dit_config.arch_config.guidance_embeds = True
+    # ComfyUI decides guidance_embeds the same way: whether the checkpoint
+    # actually carries guidance_in weights. FLUX.1-dev has them; FLUX.1-schnell
+    # does not.
+    dit_config.arch_config.guidance_embeds = _checkpoint_has_guidance_embeds(
+        server_args.model_path
+    )
     return dit_config
 
 
+def _checkpoint_has_guidance_embeds(model_path: str) -> bool:
+    with safe_open(model_path, framework="pt", device="cpu") as checkpoint:
+        return "guidance_in.in_layer.weight" in checkpoint.keys()
+
+
 def _split_sizes(dit_config: FluxConfig) -> tuple[int, int]:
     arch_config = dit_config.arch_config
     hidden_size = arch_config.num_attention_heads * arch_config.attention_head_dim

python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 """Pack / unpack contract for ComfyUI model adapters."""
 
+import pytest
 import torch
 
 from sglang.multimodal_gen.apps.ComfyUI_SGLDiffusion.executors.adapter import (
@@ -56,3 +57,77 @@ def test_flux_pack_and_unpack_roundtrip() -> None:
 
     default = adapter.pack(x, timestep, context, y=y)
     assert default.guidance_scale == 3.5
+
+
+def test_flux_pack_zero_fills_missing_pooled() -> None:
+    from sglang.multimodal_gen.runtime.layers.visual_embedding import (
+        CombinedTimestepGuidanceTextProjEmbeddings,
+    )
+
+    # ComfyUI calls Flux with y=None when the cond has no pooled_output; its
+    # native model zero-fills y, so the adapter must too.
+    adapter = FluxAdapter()
+    x = torch.zeros(2, 16, 8, 8)
+    context = torch.ones(2, 8, 4096, dtype=torch.bfloat16)
+    packed = adapter.pack(x, torch.tensor([0.5, 0.5]), context, y=None)
+    y = packed.pooled_embeds[0]
+    assert y.shape == (2, 768)
+    assert y.dtype == context.dtype
+    assert not y.any()
+    assert packed.prompt_embeds[0] is y
+
+    # The JIT timestep_embedding kernel only accepts CUDA tensors, so run the
+    # embedding where the kernel can run.
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    embed = CombinedTimestepGuidanceTextProjEmbeddings(
+        embedding_dim=32, pooled_projection_dim=768
+    ).to(device)
+    out = embed(
+        torch.tensor([500.0, 500.0], device=device),
+        torch.tensor([3.5, 3.5], device=device),
+        y.float().to(device),
+    )
+    assert out.shape == (2, 32)
+
+
+def test_flux_pack_rejects_kontext_reference_latents() -> None:
+    """ComfyUI's Flux forward passes ref_latents from the ReferenceLatent node;
+    FluxAdapter must reject it instead of silently running plain T2I."""
+    adapter = FluxAdapter()
+    x = torch.ones(1, 16, 8, 8)
+    timestep = torch.tensor([0.5])
+    context = torch.ones(1, 8, 4096)
+    with pytest.raises(ValueError, match="Kontext"):
+        adapter.pack(
+            x,
+            timestep,
+            context,
+            y=torch.ones(1, 768),
+            ref_latents=[torch.ones(1, 16, 8, 8)],
+        )
+
+
+def test_flux_pack_rejects_controlnet_control() -> None:
+    """ComfyUI's Flux forward passes control from a ControlNet node; FluxAdapter
+    must reject it instead of silently running unconditioned."""
+    adapter = FluxAdapter()
+    x = torch.ones(1, 16, 8, 8)
+    timestep = torch.tensor([0.5])
+    context = torch.ones(1, 8, 4096)
+    with pytest.raises(ValueError, match="ControlNet"):
+        adapter.pack(x, timestep, context, y=torch.ones(1, 768), control={"input": []})
+
+
+def test_flux_pack_and_unpack_roundtrip_odd_latent_size() -> None:
+    """Regression: a width/height not divisible by patch_size=2 (e.g. a
+    1032px-wide ComfyUI latent, 1032 // 8 = 129) raised a view() RuntimeError
+    in _pack_latents before it was padded like QwenImageAdapter."""
+    adapter = FluxAdapter()
+    x = torch.arange(1 * 16 * 8 * 9, dtype=torch.float32).reshape(1, 16, 8, 9)
+    timestep = torch.tensor([0.5])
+    context = torch.ones(1, 8, 4096)
+    y = torch.ones(1, 768)
+    packed = adapter.pack(x, timestep, context, y=y, guidance=torch.tensor([1.0]))
+    out = adapter.unpack(packed.latents, packed, x)
+    assert out.shape == x.shape
+    assert torch.equal(out, x)

python/sglang/multimodal_gen/test/unit/test_comfyui_profile.py

@@ -47,7 +47,9 @@ def _shrink(arch) -> None:
     arch.guidance_embeds = True
 
 
-def _write_comfyui_flux_checkpoint(path: str, config: FluxConfig) -> dict:
+def _write_comfyui_flux_checkpoint(
+    path: str, config: FluxConfig, with_guidance: bool = True
+) -> dict:
     """Write a checkpoint using ComfyUI's parameter names and fused layouts."""
     arch = config.arch_config
     hidden = arch.num_attention_heads * arch.attention_head_dim
@@ -93,11 +95,10 @@ def put(name: str, *shape: int) -> None:
         put(f"single_blocks.{b}.modulation.lin.weight", 3 * hidden, hidden)
         put(f"single_blocks.{b}.modulation.lin.bias", 3 * hidden)
 
-    for stem, in_dim in (
-        ("time_in", 256),
-        ("vector_in", arch.pooled_projection_dim),
-        ("guidance_in", 256),
-    ):
+    stems = [("time_in", 256), ("vector_in", arch.pooled_projection_dim)]
+    if with_guidance:
+        stems.append(("guidance_in", 256))
+    for stem, in_dim in stems:
         put(f"{stem}.in_layer.weight", hidden, in_dim)
         put(f"{stem}.in_layer.bias", hidden)
         put(f"{stem}.out_layer.weight", hidden, hidden)
@@ -203,3 +204,101 @@ def test_adaln_modulation_scale_and_shift_are_swapped(comfyui_flux_pipeline):
     half = source.shape[0] // 2
     # ComfyUI emits [shift, scale]; AdaLayerNormContinuous expects [scale, shift].
     assert torch.equal(actual, torch.cat([source[half:], source[:half]], dim=0))
+
+
+def test_comfyui_flux_guidance_embeds_follows_checkpoint_keys():
+    """FLUX.1-schnell checkpoints have no guidance_in.* tensors; the loader must
+    not force guidance_embeds=True for them, or the model ends up with a
+    guidance embedder that the checkpoint can never populate."""
+    spec = get_comfyui_checkpoint_spec("FluxPipeline")
+
+    for with_guidance in (True, False):
+        tmpdir = tempfile.mkdtemp(prefix="comfyui_flux_guidance_")
+        checkpoint = os.path.join(tmpdir, "flux_comfyui.safetensors")
+
+        seed_config = FluxConfig()
+        _shrink(seed_config.arch_config)
+        _write_comfyui_flux_checkpoint(checkpoint, seed_config, with_guidance)
+
+        server_args = ServerArgs.from_kwargs(
+            model_path=checkpoint,
+            pipeline_class_name="FluxPipeline",
+            comfyui_mode=True,
+            num_gpus=1,
+        )
+        _shrink(server_args.pipeline_config.dit_config.arch_config)
+
+        dit_config = spec.build_dit_config(server_args)
+        assert dit_config.arch_config.guidance_embeds is with_guidance
+
+
+def test_comfyui_flux_schnell_checkpoint_loads_without_guidance_embedder():
+    """Regression test for the ComfyUI-mode loader always assuming
+    guidance_embeds=True: a FLUX.1-schnell-shaped checkpoint (no guidance_in.*
+    tensors) must load on its own, not just dev-shaped ones."""
+    _ensure_single_process_parallel_runtime()
+
+    from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+    from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
+    from sglang.multimodal_gen.runtime.loader.weight_utils import (
+        safetensors_weights_iterator,
+    )
+    from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+
+    spec = get_comfyui_checkpoint_spec("FluxPipeline")
+
+    tmpdir = tempfile.mkdtemp(prefix="comfyui_flux_schnell_")
+    checkpoint = os.path.join(tmpdir, "flux_comfyui.safetensors")
+
+    seed_config = FluxConfig()
+    _shrink(seed_config.arch_config)
+    _write_comfyui_flux_checkpoint(checkpoint, seed_config, with_guidance=False)
+
+    server_args = ServerArgs.from_kwargs(
+        model_path=checkpoint,
+        pipeline_class_name="FluxPipeline",
+        comfyui_mode=True,
+        num_gpus=1,
+    )
+    _shrink(server_args.pipeline_config.dit_config.arch_config)
+    set_global_server_args(server_args)
+
+    dit_config = spec.build_dit_config(server_args)
+    mapping = dict(spec.param_names_mapping)
+    if spec.inherit_config_mapping:
+        mapping = {**(dit_config.arch_config.param_names_mapping or {}), **mapping}
+    dit_config.arch_config.param_names_mapping = mapping
+
+    model_cls, _ = ModelRegistry.resolve_model_cls(spec.dit_cls_name)
+    weights_iterator = spec.convert_weights(
+        safetensors_weights_iterator([checkpoint]), dit_config
+    )
+
+    original_mapping = model_cls.param_names_mapping
+    model_cls.param_names_mapping = mapping
+    try:
+        model = maybe_load_fsdp_model(
+            model_cls=model_cls,
+            init_params={"config": dit_config, "hf_config": {}},
+            weight_dir_list=[checkpoint],
+            device=get_local_torch_device(),
+            hsdp_replicate_dim=1,
+            hsdp_shard_dim=1,
+            component_starts_on_cpu=True,
+            pin_cpu_memory=False,
+            fsdp_inference=False,
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.float32,
+            output_dtype=None,
+            strict=spec.strict,
+            weights_iterator=weights_iterator,
+        )
+    finally:
+        model_cls.param_names_mapping = original_mapping
+
+    unloaded = [name for name, p in model.named_parameters() if p.is_meta]
+    assert not unloaded, f"parameters never received checkpoint weights: {unloaded[:5]}"
+    assert not any(
+        name.startswith("time_text_embed.guidance_embedder")
+        for name, _ in model.named_parameters()
+    )