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..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,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): 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_adapters.py b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py index f1963f6a920e..a7c4b9349ea7 100644 --- a/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py +++ b/python/sglang/multimodal_gen/test/unit/test_comfyui_adapters.py @@ -69,6 +69,80 @@ 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 + 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) + + def test_worker_error_is_raised_not_unpacked() -> None: """Unpacking a failed reply replaced the worker's message with a misleading adapter TypeError about noise_pred being None.""" 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() + )