From 8cf049c08c41fa0703bf26191dbbd4861f7b002a Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sat, 10 Oct 2026 09:35:09 +0530 Subject: [PATCH 1/4] Fix Flux crash on latent sizes not divisible by patch size --- .../ComfyUI_SGLDiffusion/executors/flux.py | 32 +++++++++++++------ .../test/unit/test_comfyui_adapters.py | 15 +++++++++ 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py index 87ecdd02bb61..1ad77b7ad6b9 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py @@ -14,14 +14,23 @@ 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 ) -> PackedForward: - packed = self._pack_latents(x) + 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 return PackedForward( @@ -30,21 +39,22 @@ def pack( 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 +63,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): diff --git a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py index 42c1772baf1f..74fc33af7e24 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py +++ b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py @@ -56,3 +56,18 @@ 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_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) From 71e73f8e4bd3cba504aa83e07ba5f86fb9ec6d9c Mon Sep 17 00:00:00 2001 From: OnePunchMonk Date: Sat, 10 Oct 2026 09:38:50 +0530 Subject: [PATCH 2/4] Zero-fill missing pooled vector and reject Kontext/ControlNet in FluxAdapter Fixes #43115, fixes #43488. --- .../ComfyUI_SGLDiffusion/executors/flux.py | 34 +++++++++++- .../test/unit/test_comfyui_adapters.py | 54 +++++++++++++++++++ 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py index 1ad77b7ad6b9..301e69a78901 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py @@ -6,6 +6,10 @@ 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: return 3.5 @@ -28,11 +32,37 @@ class FluxAdapter(ComfyUIModelAdapter): 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: + 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, diff --git a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py index 74fc33af7e24..fade30bcdb38 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py +++ b/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 ( @@ -58,6 +59,59 @@ def test_flux_pack_and_unpack_roundtrip() -> None: 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 From 0451e1fb35c97f58b111250138afbce86f2ee071 Mon Sep 17 00:00:00 2001 From: onepunchmonk Date: Fri, 9 Oct 2026 21:40:59 +0530 Subject: [PATCH 3/4] [diffusion] fix: resolve ComfyUI Flux guidance_embeds from checkpoint keys instead of assuming True ComfyUI-mode loading always set guidance_embeds=True for Flux checkpoints, assuming every checkpoint carries guidance_in.* weights. FLUX.1-schnell checkpoints don't, so the model ended up with a guidance embedder the checkpoint could never populate, and strict loading failed with "New parameter 'time_text_embed.guidance_embedder.linear_1.weight' is not supported". Resolve guidance_embeds from whether the checkpoint's safetensors header actually has guidance_in.in_layer.weight, matching how ComfyUI itself decides. The dead "if not has_guidance_embeds" skip in _convert_weights now does something. Fixes #43395. --- .../loader/comfyui_checkpoints/flux.py | 14 ++- .../test/unit/test_comfyui_profile.py | 111 +++++++++++++++++- 2 files changed, 117 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/flux.py b/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/flux.py index c2f4e17a9c10..9d46f2facd9a 100644 --- a/python/sglang/multimodal_gen/runtime/loader/comfyui_checkpoints/flux.py +++ b/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 diff --git a/python/sglang/multimodal_gen/test/unit/test_comfyui_profile.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_profile.py index f6d6dbcf1507..b879e67a8b07 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_profile.py +++ b/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() + ) From 1ec558209de80dddc91f1f5fc74e0b790aff9a5d Mon Sep 17 00:00:00 2001 From: onepunchmonk Date: Sat, 10 Oct 2026 12:04:06 +0530 Subject: [PATCH 4/4] lint: isort/ruff-format fixes --- .../apps/ComfyUI_SGLDiffusion/executors/flux.py | 1 - .../multimodal_gen/test/unit/test_comfyui_adapters.py | 8 +++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py index 301e69a78901..7e7d7c4b7813 100644 --- a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py +++ b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py @@ -5,7 +5,6 @@ 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 diff --git a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py index fade30bcdb38..de635c109b01 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py +++ b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py @@ -98,7 +98,13 @@ def test_flux_pack_rejects_kontext_reference_latents() -> None: 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)]) + 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: