[diffusion] comfyui: support Qwen-Image, Multi-GPU Z-Image, and Enhanced ComfyUI Integration (#17678)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -149,11 +149,12 @@ class GPUWorker:
                 can_stay_resident = self.get_can_stay_resident_components(
                     remaining_gpu_mem_gb
                 )
-                logger.info(
-                    f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
-                    f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
-                    f"Components that can stay resident: {can_stay_resident}"
-                )
+                if not req.suppress_logs:
+                    logger.info(
+                        f"Peak GPU memory: {peak_memory_gb:.2f} GB, "
+                        f"Remaining GPU memory at peak: {remaining_gpu_mem_gb:.2f} GB. "
+                        f"Components that can stay resident: {can_stay_resident}"
+                    )
 
             duration_ms = (time.monotonic() - start_time) * 1000
             output_batch.timings.total_duration_ms = duration_ms

python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py

@@ -288,7 +288,7 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase):
     ) -> Generator[tuple[str, torch.Tensor], None, None]:
         """
         Convert ComfyUI Flux weights to SGLang format.
-        Splits fused linear1 into separate to_qkv and proj_mlp weights.
+        Splits fused qkv weights into to_q/to_k/to_v plus proj_mlp.
         Filters out guidance_in weights if model doesn't support guidance embeddings.
         Handles scale/shift order difference between ComfyUI and AdaLayerNormContinuous.
         """
@@ -299,6 +299,41 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase):
                 )
                 continue
 
+            # Split fused qkv in double blocks into separate q/k/v projections
+            match = re.match(
+                r"double_blocks\.(\d+)\.(img_attn|txt_attn)\.qkv\.(weight|bias)$", name
+            )
+            if match:
+                block_idx, attn_type, param_type = match.groups()
+                hidden_size = qkv_size // 3
+
+                if tensor.shape[0] < 3 * hidden_size:
+                    logger.warning(
+                        f"{name} shape {tensor.shape} smaller than expected qkv size {3 * hidden_size}, skipping"
+                    )
+                    continue
+
+                if param_type == "bias":
+                    q_tensor = tensor[:hidden_size]
+                    k_tensor = tensor[hidden_size : 2 * hidden_size]
+                    v_tensor = tensor[2 * hidden_size : 3 * hidden_size]
+                else:
+                    q_tensor = tensor[:hidden_size, :]
+                    k_tensor = tensor[hidden_size : 2 * hidden_size, :]
+                    v_tensor = tensor[2 * hidden_size : 3 * hidden_size, :]
+
+                target_prefix = f"transformer_blocks.{block_idx}.attn"
+                if attn_type == "img_attn":
+                    yield f"{target_prefix}.to_q.{param_type}", q_tensor
+                    yield f"{target_prefix}.to_k.{param_type}", k_tensor
+                    yield f"{target_prefix}.to_v.{param_type}", v_tensor
+                else:
+                    # txt_attn corresponds to encoder projections
+                    yield f"{target_prefix}.add_q_proj.{param_type}", q_tensor
+                    yield f"{target_prefix}.add_k_proj.{param_type}", k_tensor
+                    yield f"{target_prefix}.add_v_proj.{param_type}", v_tensor
+                continue
+
             match = re.match(r"single_blocks\.(\d+)\.linear1\.(weight|bias)$", name)
             if match:
                 block_idx, param_type = match.groups()
@@ -319,8 +354,20 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase):
                     tensor[qkv_size:] if param_type == "bias" else tensor[qkv_size:, :]
                 )
 
-                # Yield split weights
-                yield f"single_transformer_blocks.{block_idx}.attn.to_qkv.{param_type}", qkv_tensor
+                # Split qkv into q/k/v for single blocks
+                hidden_size = qkv_size // 3
+                if param_type == "bias":
+                    q_tensor = qkv_tensor[:hidden_size]
+                    k_tensor = qkv_tensor[hidden_size : 2 * hidden_size]
+                    v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size]
+                else:
+                    q_tensor = qkv_tensor[:hidden_size, :]
+                    k_tensor = qkv_tensor[hidden_size : 2 * hidden_size, :]
+                    v_tensor = qkv_tensor[2 * hidden_size : 3 * hidden_size, :]
+
+                yield f"single_transformer_blocks.{block_idx}.attn.to_q.{param_type}", q_tensor
+                yield f"single_transformer_blocks.{block_idx}.attn.to_k.{param_type}", k_tensor
+                yield f"single_transformer_blocks.{block_idx}.attn.to_v.{param_type}", v_tensor
                 yield f"single_transformer_blocks.{block_idx}.proj_mlp.{param_type}", mlp_tensor
             elif name == "final_layer.adaLN_modulation.1.weight":
                 # ComfyUI: output order is [shift, scale]

python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py

@@ -0,0 +1,349 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from itertools import chain
+from typing import Any
+
+import torch
+from torch.distributed import init_device_mesh
+from torch.distributed.fsdp import MixedPrecisionPolicy
+
+from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.loader.fsdp_load import (
+    load_model_from_full_model_state_dict,
+    set_default_dtype,
+    shard_model,
+)
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+from sglang.multimodal_gen.runtime.loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_comfyui_passthrough import (
+    ComfyUIPassThroughScheduler,
+)
+from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    ComfyUILatentPreparationStage,
+    DenoisingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, set_mixed_precision_policy
+
+logger = init_logger(__name__)
+
+
+class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
+    """
+    Base pipeline for ComfyUI QwenImage integration with only denoising stage.
+
+    This pipeline requires pre-processed inputs:
+    - prompt_embeds: Pre-encoded text embeddings (list of tensors)
+    - latents: Pre-processed image latents in sequence format [B, S, D]
+
+    Usage:
+        generator = DiffGenerator.from_pretrained(
+            model_path="path/to/model",
+            pipeline_class_name="ComfyUIQwenImagePipeline",
+            device="cuda",
+        )
+    """
+
+    # Subclasses should override this
+    zero_cond_t: bool = False
+
+    pipeline_name = "ComfyUIQwenImagePipeline"
+
+    _required_config_modules = [
+        "transformer",
+        "scheduler",
+    ]
+
+    def initialize_pipeline(self, server_args: ServerArgs):
+        """
+        Initialize the pipeline with ComfyUI pass-through scheduler.
+        This scheduler does not modify latents, allowing ComfyUI to handle denoising.
+        """
+        self.modules["scheduler"] = ComfyUIPassThroughScheduler(
+            num_train_timesteps=1000
+        )
+
+        # Ensure VAE config is properly initialized even though we don't load the VAE model
+        vae_config = server_args.pipeline_config.vae_config
+        vae_config.post_init()
+        logger.info(
+            "Called vae_config.post_init() to set vae_scale_factor. "
+            f"vae_scale_factor={vae_config.arch_config.vae_scale_factor}"
+        )
+
+    def load_modules(
+        self,
+        server_args: ServerArgs,
+        loaded_modules: dict[str, torch.nn.Module] | None = None,
+    ) -> dict[str, Any]:
+        """
+        Load modules for ComfyUIQwenImagePipeline.
+
+        If model_path is a safetensors file, load transformer directly from it
+        without requiring model_index.json. Otherwise, fall back to default loading.
+        """
+        if os.path.isfile(self.model_path) and self.model_path.endswith(".safetensors"):
+            logger.info(
+                "Detected safetensors file, loading transformer directly from: %s",
+                self.model_path,
+            )
+            return self._load_transformer_from_safetensors(server_args, loaded_modules)
+        else:
+            logger.info(
+                "Model path is a directory, using default loading method: %s",
+                self.model_path,
+            )
+            return super().load_modules(server_args, loaded_modules)
+
+    def _load_transformer_from_safetensors(
+        self,
+        server_args: ServerArgs,
+        loaded_modules: dict[str, torch.nn.Module] | None = None,
+    ) -> dict[str, Any]:
+        """Load transformer directly from safetensors without model_index.json."""
+
+        # 1) Fast path: use provided module
+        if loaded_modules is not None and "transformer" in loaded_modules:
+            logger.info("Using provided transformer module")
+            return {
+                "transformer": loaded_modules["transformer"],
+                "scheduler": self.modules.get("scheduler"),
+            }
+
+        # 2) Build config and mappings
+        dit_config, updated_mapping, model_cls, default_dtype = (
+            self._prepare_dit_config_and_mapping(server_args)
+        )
+        safetensors_list = [self.model_path]
+        logger.info("Loading weights from: %s", safetensors_list)
+
+        # 3) Instantiate model (meta) and optionally shard
+        model = self._instantiate_model(
+            model_cls, dit_config, default_dtype, updated_mapping, server_args
+        )
+
+        # 4) Load weights
+        self._load_weights_into_model(
+            model, safetensors_list, default_dtype, updated_mapping, server_args
+        )
+
+        components = {
+            "transformer": model,
+            "scheduler": self.modules.get("scheduler"),
+        }
+        logger.info("Successfully loaded modules: %s", list(components.keys()))
+        return components
+
+    def _prepare_dit_config_and_mapping(self, server_args: ServerArgs):
+        from sglang.multimodal_gen.configs.models.dits.qwenimage import (
+            QwenImageArchConfig,
+        )
+
+        comfyui_arch_config = QwenImageArchConfig(
+            patch_size=2,
+            in_channels=64,
+            out_channels=16,
+            num_layers=60,
+            attention_head_dim=128,
+            num_attention_heads=24,
+            joint_attention_dim=3584,
+            pooled_projection_dim=768,
+            guidance_embeds=False,
+            axes_dims_rope=(16, 56, 56),
+            zero_cond_t=self.zero_cond_t,
+        )
+        dit_config = QwenImageDitConfig(arch_config=comfyui_arch_config)
+        server_args.pipeline_config.dit_config = dit_config
+
+        if dit_config.arch_config.param_names_mapping is None:
+            dit_config.arch_config.param_names_mapping = {}
+
+        comfyui_qwen_mappings = {r"^model\.diffusion_model\.(.*)$": r"\1"}
+        updated_mapping = {
+            **dit_config.arch_config.param_names_mapping,
+            **comfyui_qwen_mappings,
+        }
+        dit_config.arch_config.param_names_mapping = updated_mapping
+        logger.info(
+            "Added ComfyUI weight name mappings to param_names_mapping. "
+            f"Total mappings: {len(updated_mapping)}"
+        )
+
+        cls_name = "QwenImageTransformer2DModel"
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+        logger.info("Resolved transformer class: %s", cls_name)
+
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+        server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "."
+        assert server_args.hsdp_shard_dim is not None, "hsdp_shard_dim must be set"
+        logger.info(
+            "Loading %s from safetensors file, default_dtype: %s",
+            cls_name,
+            default_dtype,
+        )
+        return dit_config, updated_mapping, model_cls, default_dtype
+
+    def _instantiate_model(
+        self,
+        model_cls,
+        dit_config,
+        default_dtype,
+        updated_mapping,
+        server_args: ServerArgs,
+    ):
+        from sglang.multimodal_gen.runtime.platforms import current_platform
+
+        hf_config = {}
+        original_mapping = model_cls.param_names_mapping
+        model_cls.param_names_mapping = updated_mapping
+        logger.info(
+            "Temporarily updated model class param_names_mapping with ComfyUI mappings. "
+            f"Total mappings: {len(updated_mapping)}"
+        )
+
+        try:
+            mp_policy = MixedPrecisionPolicy(
+                torch.bfloat16, torch.float32, None, cast_forward_inputs=False
+            )
+            set_mixed_precision_policy(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                output_dtype=None,
+                mp_policy=mp_policy,
+            )
+
+            with set_default_dtype(default_dtype), torch.device("meta"):
+                model = model_cls(**{"config": dit_config, "hf_config": hf_config})
+
+            use_fsdp = server_args.use_fsdp_inference
+            if current_platform.is_mps():
+                use_fsdp = False
+                logger.info("Disabling FSDP for MPS platform as it's not compatible")
+
+            if use_fsdp:
+                device_mesh = init_device_mesh(
+                    current_platform.device_type,
+                    mesh_shape=(
+                        server_args.hsdp_replicate_dim,
+                        server_args.hsdp_shard_dim,
+                    ),
+                    mesh_dim_names=("replicate", "shard"),
+                )
+                shard_model(
+                    model,
+                    cpu_offload=server_args.dit_cpu_offload,
+                    reshard_after_forward=True,
+                    mp_policy=mp_policy,
+                    mesh=device_mesh,
+                    fsdp_shard_conditions=model._fsdp_shard_conditions,
+                    pin_cpu_memory=server_args.pin_cpu_memory,
+                )
+        finally:
+            model_cls.param_names_mapping = original_mapping
+
+        return model
+
+    def _load_weights_into_model(
+        self,
+        model,
+        safetensors_list,
+        default_dtype,
+        updated_mapping,
+        server_args: ServerArgs,
+    ):
+        # Create weight iterator for loading
+        weight_iterator = safetensors_weights_iterator(safetensors_list)
+
+        # Load weights
+        param_names_mapping_fn = get_param_names_mapping(updated_mapping)
+        load_model_from_full_model_state_dict(
+            model,
+            weight_iterator,
+            get_local_torch_device(),
+            default_dtype,
+            strict=True,
+            cpu_offload=server_args.dit_cpu_offload,
+            param_names_mapping=param_names_mapping_fn,
+        )
+
+        # Check for meta parameters
+        for n, p in chain(model.named_parameters(), model.named_buffers()):
+            if p.is_meta:
+                raise RuntimeError(f"Unexpected param or buffer {n} on meta device.")
+            if isinstance(p, torch.nn.Parameter):
+                p.requires_grad = False
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9)
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        logger.info(
+            f"{self.__class__.__name__}.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage"
+        )
+
+        # Add ComfyUILatentPreparationStage to handle latents properly for SP
+        # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=ComfyUILatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        # Add DenoisingStage for the actual denoising process
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+        logger.info(
+            f"{self.__class__.__name__} stages created: {list(self._stage_name_mapping.keys())}"
+        )
+
+
+class ComfyUIQwenImagePipeline(ComfyUIQwenImagePipelineBase):
+    """ComfyUI QwenImage pipeline for text-to-image generation."""
+
+    pipeline_name = "ComfyUIQwenImagePipeline"
+    zero_cond_t = False
+
+    from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+        QwenImagePipelineConfig,
+    )
+    from sglang.multimodal_gen.configs.sample.qwenimage import QwenImageSamplingParams
+
+    pipeline_config_cls = QwenImagePipelineConfig
+    sampling_params_cls = QwenImageSamplingParams
+
+
+class ComfyUIQwenImageEditPipeline(ComfyUIQwenImagePipelineBase):
+    """ComfyUI QwenImage pipeline for image-to-image editing."""
+
+    pipeline_name = "ComfyUIQwenImageEditPipeline"
+    zero_cond_t = True
+
+    from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
+        QwenImageEditPlusPipelineConfig,
+    )
+    from sglang.multimodal_gen.configs.sample.qwenimage import (
+        QwenImageEditPlusSamplingParams,
+    )
+
+    pipeline_config_cls = QwenImageEditPlusPipelineConfig
+    sampling_params_cls = QwenImageEditPlusSamplingParams
+
+
+EntryClass = [ComfyUIQwenImagePipeline, ComfyUIQwenImageEditPipeline]

python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py

@@ -30,7 +30,10 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    ComfyUILatentPreparationStage,
+    DenoisingStage,
+)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE, set_mixed_precision_policy
@@ -373,8 +376,20 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
 
     def create_pipeline_stages(self, server_args: ServerArgs):
         logger.info(
-            "ComfyUIZImagePipeline.create_pipeline_stages() called - creating only denoising_stage"
+            "ComfyUIZImagePipeline.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage"
         )
+
+        # Add ComfyUILatentPreparationStage to handle latents properly for SP
+        # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=ComfyUILatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        # Add DenoisingStage for the actual denoising process
         self.add_stage(
             stage_name="denoising_stage",
             stage=DenoisingStage(

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -330,7 +330,7 @@ class ComposedPipelineBase(ABC):
         batch.log(server_args=server_args)
 
         # Execute each stage
-        if not batch.is_warmup:
+        if not batch.is_warmup and not batch.suppress_logs:
             logger.info(
                 "Running pipeline stages: %s",
                 list(self._stage_name_mapping.keys()),

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -106,6 +106,7 @@ class Req:
     # vae-encoded condition image
     image_latent: torch.Tensor | list[torch.Tensor] | None = None
     condition_image_latent_ids: torch.Tensor | list[torch.Tensor] | None = None
+    vae_image_sizes: list[tuple[int, int]] | None = None
 
     # Latent dimensions
     height_latents: list[int] | int | None = None
@@ -302,7 +303,8 @@ class Req:
                  save_output: {self.save_output}
             output_file_path: {self.output_file_path()}
         """  # type: ignore[attr-defined]
-        logger.info(debug_str)
+        if not self.suppress_logs:
+            logger.info(debug_str)
 
 
 @dataclass

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -54,6 +54,8 @@ class PipelineStage(ABC):
 
     def log_info(self, msg, *args):
         """Logs an informational message with the stage name as a prefix."""
+        if self.server_args.comfyui_mode:
+            return
         logger.info(f"[{self.__class__.__name__}] {msg}", *args)
 
     def log_warning(self, msg, *args):
@@ -195,7 +197,8 @@ class PipelineStage(ABC):
             logger=logger,
             timings=batch.timings,
             perf_dump_path_provided=batch.perf_dump_path is not None,
-            log_stage_start_end=not batch.is_warmup,
+            log_stage_start_end=not batch.is_warmup
+            and not (self.server_args and self.server_args.comfyui_mode),
         ):
             result = self.forward(batch, server_args)
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/comfyui_latent_preparation.py

@@ -106,17 +106,8 @@ class ComfyUILatentPreparationStage(LatentPreparationStage):
         result = super().forward(batch, server_args)
 
         if original_latents_shape is not None:
-            current_shape = result.latents.shape if result.latents is not None else None
-            if (
-                current_shape is not None
-                and len(current_shape) == 3
-                and len(original_latents_shape) == 4
-            ):
-                # Keep original shape for raw_latent_shape
-                result.raw_latent_shape = original_latents_shape
-            elif current_shape is not None and current_shape == original_latents_shape:
-                result.raw_latent_shape = current_shape
-            else:
-                result.raw_latent_shape = original_latents_shape
+            # Preserve the original shape before any potential packing/conversion
+            # (e.g., 4D spatial -> 3D sequence) to ensure proper unpadding later.
+            result.raw_latent_shape = original_latents_shape
 
         return result