[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/apps/ComfyUI_SGLDiffusion/README.md

@@ -4,46 +4,56 @@ A ComfyUI plugin for integrating with SGLang Diffusion server, supporting image
 
 ## Installation
 
-1. Copy this entire directory (`ComfyUI_SGLDiffusion`) to your ComfyUI `custom_nodes/` folder
-2. Restart ComfyUI to load the plugin
+1. **Install SGLang**: Follow the [Installation Guide](../../docs/install.md) to install `sglang[diffusion]`.
+2. **Install Plugin**: Copy this entire directory (`ComfyUI_SGLDiffusion`) to your ComfyUI `custom_nodes/` folder.
+3. **Restart ComfyUI**: Restart ComfyUI to load the plugin.
 
 ## Usage
 
-### Prerequisites
+The plugin supports two modes of operation: **Server Mode** (via HTTP API) and **Integrated Mode** (tight integration with ComfyUI).
 
-Before using this plugin, you need to start the SGLang Diffusion server. The plugin connects to the server via HTTP API calls.
+### Supported Models
+- **Z-Image**: High-speed image generation models (e.g., `Z-Image-Turbo`)
+- **FLUX**: State-of-the-art text-to-image models (e.g., `FLUX.1-dev`)
+- **Qwen-Image**: Multi-modal image generation models (e.g., `Qwen-Image`,`Qwen-Image-2512`). *Note: Image editing support is currently experimental and may have some issues.*
 
-### Basic Workflow
+### Mode 1: Server Mode (HTTP API)
+Connect to a standalone SGLang Diffusion server.
 
-1. **Start SGLang Diffusion Server**: Ensure the SGLang Diffusion server is running and accessible
-2. **Connect to Server**: Use the `SGLDiffusion Server Model` node to connect to your server (default: `http://localhost:3000/v1`)
-3. **Generate Content**: Use the generation nodes to create images or videos:
-   - `SGLDiffusion Generate Image`: For text-to-image and image editing
-   - `SGLDiffusion Generate Video`: For text-to-video and image-to-video
-4. **Optional**: Use `SGLDiffusion Set LoRA` to load LoRA adapters for style customization, use `SGLDiffusion Unset LoRA` to remove LoRA
+1. **Start SGLang Diffusion Server**: Ensure the server is running and accessible.
+2. **Connect to Server**: Use the `SGLDiffusion Server Model` node to connect (default: `http://localhost:3000/v1`).
+3. **Generate Content**:
+   - `SGLDiffusion Generate Image`: For text-to-image and image editing.
+   - `SGLDiffusion Generate Video`: For text-to-video and image-to-video.
+4. **LoRA Support**: Use `SGLDiffusion Server Set LoRA` and `SGLDiffusion Server Unset LoRA`.
 
-### Example Workflows
+### Mode 2: Integrated Mode (Tight Integration)
+Leverage SGLang's high-performance sampling directly within ComfyUI while using ComfyUI's front-end nodes (CLIP, VAE, etc.).
+
+1. **Load Model**: Use the `SGLDiffusion UNET Loader` node to load your diffusion model.
+2. **Configure Options**: Use the `SGLDiffusion Options` node to set runtime parameters like `num_gpus`, `tp_size`, `model_type`, or `enable_torch_compile`.
+3. **Sample**: Connect the loaded model to standard ComfyUI samplers. SGLang will handle the sampling process efficiently.
+4. **LoRA Support**: Use the `SGLDiffusion LoRA Loader` for native LoRA integration.
+
+## Example Workflows
 
 Reference workflow files are provided in the `workflows/` directory:
 
-- **`sgld_text2img.json`**: Text-to-image generation with LoRA support
-- **`sgld_image2video.json`**: Image-to-video generation
+- **`flux_sgld_sp.json`**: Multi-GPU (Sequence Parallelism) workflow for FLUX models. High-performance inference across multiple cards.
+- **`qwen_image_sgld.json`**: Qwen-Image generation with LoRA support. Optimized for multi-modal image tasks.
+- **`z-image_sgld.json`**: High-speed image generation using Z-Image.
+- **`sgld_text2img.json`**: Server-mode text-to-image generation with LoRA support.
+- **`sgld_image2video.json`**: Server-mode image-to-video generation.
+
+For other workflows supporting the models, you can easily use SGLang by replacing the official `UNET Loader` node with the `SGLDUNETLoader` node. Similarly, for LoRA support, replace the official LoRA loader with the `SGLDiffusion LoRA Loader`.
 
 To use these workflows:
-1. Open ComfyUI
-2. Load the workflow JSON file from the `workflows/` directory
-3. Adjust the server URL and API key in the `SGLDiffusion Server Model` node if needed
-4. Modify prompts and parameters as desired
-5. Run the workflow
+1. Open ComfyUI.
+2. Load the workflow JSON file from the `workflows/` directory.
+3. Adjust the parameters and model paths as needed.
+4. Run the workflow.
+
 
 ## Current Implementation
 
-Currently, this plugin uses a server-based approach where it makes HTTP API calls to a running SGLang Diffusion server. This requires:
-
-- The SGLang Diffusion server to be running separately
-- Network connectivity between ComfyUI and the server
-- Proper server configuration and API key setup
-
-## Future Improvements
-
-Future versions will integrate more tightly with ComfyUI, utilizing ComfyUI's models and most of its nodes, while leveraging SGLD specifically for the computationally intensive sampling process. This will provide a more seamless and efficient workflow.
+This plugin provides a high-performance backend for diffusion models in ComfyUI. By leveraging SGLang's optimized kernels and parallelization techniques (Tensor Parallelism, TeaCache, etc.), it significantly accelerates the sampling process, especially for large models like FLUX.

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

@@ -2,6 +2,12 @@
 ComfyUI SGLang Diffusion nodes package.
 """
 
-from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+try:
+    from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+except ImportError:
+    # ComfyUI dependencies not available (e.g., in test environment)
+    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

@@ -23,7 +23,12 @@ except ImportError:
         "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
     )
 
-from ..executors import FluxExecutor, ZImageExecutor
+from ..executors import (
+    FluxExecutor,
+    QwenImageEditExecutor,
+    QwenImageExecutor,
+    ZImageExecutor,
+)
 from .model_patcher import SGLDModelPatcher
 
 
@@ -39,10 +44,14 @@ class SGLDiffusionGenerator:
         self.pipeline_class_dict = {
             "flux": "ComfyUIFluxPipeline",
             "lumina2": "ComfyUIZImagePipeline",  # zimage
+            "qwen_image": "ComfyUIQwenImagePipeline",
+            "qwen_image_edit": "ComfyUIQwenImageEditPipeline",
         }
         self.executor_class_dict = {
             "flux": FluxExecutor,
             "lumina2": ZImageExecutor,
+            "qwen_image": QwenImageExecutor,
+            "qwen_image_edit": QwenImageEditExecutor,
         }
 
     def __del__(self):
@@ -200,6 +209,10 @@ class SGLDiffusionGenerator:
         if model_type is None or model_type not in self.pipeline_class_dict:
             raise ValueError(f"Unsupported model type: {model_type}")
 
+        set_model_type = sgld_options.pop("model_type", None) if sgld_options else None
+        if set_model_type is not None and set_model_type in self.pipeline_class_dict:
+            model_type = set_model_type
+
         pipeline_class_name = self.pipeline_class_dict[model_type]
         self.generator = self.init_generator(
             model_path, pipeline_class_name, sgld_options

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

@@ -5,10 +5,13 @@ Provides executor classes for different model types.
 
 from .base import SGLDiffusionExecutor
 from .flux import FluxExecutor
+from .qwen_image import QwenImageEditExecutor, QwenImageExecutor
 from .zimage import ZImageExecutor
 
 __all__ = [
     "SGLDiffusionExecutor",
     "FluxExecutor",
     "ZImageExecutor",
+    "QwenImageExecutor",
+    "QwenImageEditExecutor",
 ]

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

@@ -17,6 +17,23 @@ class SGLDiffusionExecutor(torch.nn.Module):
         self.config = config
         self.loras = []
 
+    @staticmethod
+    def should_suppress_logs(timestep):
+        """Determine if logs should be suppressed based on timestep value."""
+        if torch.is_tensor(timestep):
+            return bool((timestep < 1.0).item())
+        return bool(timestep < 1.0)
+
+    def set_lora(self, lora_nickname=None, lora_path=None, strength=None, target=None):
+        """Set LoRA adapter using SGLang Diffusion API."""
+        if len(lora_nickname) > 0:
+            self.generator.set_lora(
+                lora_nickname=lora_nickname,
+                lora_path=lora_path,
+                strength=strength,
+                target=target,
+            )
+
     def _unpack_latents(self, latents, height, width, channels):
         """Unpack latents from packed format to standard format."""
         batch_size = latents.shape[0]

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

@@ -43,6 +43,7 @@ class FluxExecutor(SGLDiffusionExecutor):
             num_frames=1,
             num_inference_steps=1,
             save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
         )
 
         # Prepare request (converts SamplingParams to Req)

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

@@ -0,0 +1,174 @@
+"""
+QwenImage executor for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+try:
+    from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+    from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+except ImportError:
+    print(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+import comfy.ldm.common_dit
+
+from .base import SGLDiffusionExecutor
+
+
+class QwenImageExecutor(SGLDiffusionExecutor):
+    """Executor for QwenImage models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+        self.patch_size = 2
+
+    def _pack_latents(self, x):
+        """Process hidden states for QwenImage model."""
+        bs, c, t, h, w = x.shape
+        patch_size = self.patch_size
+        latents = comfy.ldm.common_dit.pad_to_patch_size(
+            x, (1, self.patch_size, self.patch_size)
+        )
+        orig_shape = latents.shape
+        latents = latents.view(
+            orig_shape[0],
+            orig_shape[1],
+            orig_shape[-3],
+            orig_shape[-2] // 2,
+            2,
+            orig_shape[-1] // 2,
+            2,
+        )
+        latents = latents.permute(0, 2, 3, 5, 1, 4, 6)
+        latents = latents.reshape(
+            orig_shape[0],
+            orig_shape[-3] * (orig_shape[-2] // 2) * (orig_shape[-1] // 2),
+            orig_shape[1] * 4,
+        )
+        return latents, orig_shape
+
+    def _unpack_latents(self, latents, num_embeds, orig_shape, x):
+        """Unpack hidden states from packed format to standard format."""
+        latents = latents[:, :num_embeds].view(
+            orig_shape[0],
+            orig_shape[-3],
+            orig_shape[-2] // 2,
+            orig_shape[-1] // 2,
+            orig_shape[1],
+            2,
+            2,
+        )
+        latents = latents.permute(0, 4, 1, 2, 5, 3, 6)
+        latents = latents.reshape(orig_shape)[:, :, :, : x.shape[-2], : x.shape[-1]]
+        return latents
+
+    def forward(self, x, timestep, context, **kwargs):
+        """Forward pass for QwenImage model."""
+        latents, orig_shape = self._pack_latents(x)
+        num_embeds = latents.shape[1]
+        height = orig_shape[-2] * 8
+        width = orig_shape[-1] * 8
+
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=1.0,
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents
+        req.timesteps = timestep * 1000.0
+        req.prompt_embeds = [context]
+        req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+
+        return self._unpack_latents(noise_pred, num_embeds, orig_shape, x)
+
+
+class QwenImageEditExecutor(QwenImageExecutor):
+    """Executor for QwenImageEdit models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(
+        self,
+        x,
+        timestep,
+        context,
+        attention_mask=None,
+        ref_latents=None,
+        additional_t_cond=None,
+        transformer_options={},
+        **kwargs
+    ):
+        """Forward pass for QwenImageEdit model."""
+        latents, orig_shape = self._pack_latents(x)
+        num_embeds = latents.shape[1]
+        height = orig_shape[-2] * 8
+        width = orig_shape[-1] * 8
+
+        # Prepare vae_image_sizes for the condition image (ref_latents)
+        vae_image_sizes = []
+        pack_ref_latents = None
+
+        # TODO: sgld now don't support multiple condition images, so we only support one condition image for now.
+        if ref_latents is not None and len(ref_latents) > 0:
+            pack_ref_latents, orig_ref_shape = self._pack_latents(ref_latents[0])
+            vae_image_sizes = [(orig_ref_shape[-1], orig_ref_shape[-2])]
+
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=1.0,
+            image_path="",
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+            suppress_logs=self.should_suppress_logs(timestep),
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents
+        req.image_latent = pack_ref_latents
+        req.timesteps = timestep * 1000.0
+        req.vae_image_sizes = vae_image_sizes
+        req.prompt_embeds = [context]
+        req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+
+        return self._unpack_latents(noise_pred, num_embeds, orig_shape, x)

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

@@ -36,6 +36,7 @@ class ZImageExecutor(SGLDiffusionExecutor):
             num_frames=1,  # For images
             num_inference_steps=1,  # Single step for ComfyUI
             save_output=False,
+            suppress_logs=self.should_suppress_logs(timesteps),
         )
 
         # Prepare request (converts SamplingParams to Req)

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

@@ -4,6 +4,7 @@ Provides nodes for connecting to SGLang Diffusion server and generating images/v
 """
 
 import os
+import uuid
 
 import folder_paths
 import torch
@@ -23,6 +24,10 @@ class SGLDOptions:
         return {
             "required": {},
             "optional": {
+                "model_type": (
+                    ["auto-detect", "qwen_image", "qwen_image_edit", "flux", "lumina2"],
+                    {"default": "auto-detect"},
+                ),
                 "enable_torch_compile": (
                     "BOOLEAN",
                     {"default": False},
@@ -56,10 +61,6 @@ class SGLDOptions:
                     "STRING",
                     {"default": ""},
                 ),
-                "cache_strategy": (
-                    "STRING",
-                    {"default": "none"},
-                ),
             },
         }
 
@@ -70,6 +71,7 @@ class SGLDOptions:
 
     def create_options(
         self,
+        model_type: str = "auto-detect",
         enable_torch_compile: bool = False,
         num_gpus: int = 1,
         tp_size: int = -1,
@@ -80,7 +82,6 @@ class SGLDOptions:
         dp_degree: int = 1,
         enable_cfg_parallel: bool = False,
         attention_backend: str = "",
-        cache_strategy: str = "none",
     ):
         """
         Build a dictionary of SGLang Diffusion runtime options.
@@ -91,6 +92,7 @@ class SGLDOptions:
         attention_backend = None if attention_backend == "" else attention_backend
 
         options = {
+            "model_type": model_type,
             "enable_torch_compile": enable_torch_compile,
             "num_gpus": num_gpus,
             "tp_size": tp_size,
@@ -101,7 +103,6 @@ class SGLDOptions:
             "dp_degree": dp_degree,
             "enable_cfg_parallel": enable_cfg_parallel,
             "attention_backend": attention_backend,
-            "cache_strategy": cache_strategy,
         }
 
         # Strip None to keep payload clean
@@ -109,6 +110,59 @@ class SGLDOptions:
         return (options,)
 
 
+class SGLDLoraLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "lora_name": (folder_paths.get_filename_list("loras"),),
+                "strength_model": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0, "max": 10, "step": 0.01},
+                ),
+                "nickname": ("STRING", {"default": ""}),
+                "target": (
+                    ["all", "transformer", "transformer_2", "critic"],
+                    {"default": "all"},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_lora"
+
+    CATEGORY = "SGLDiffusion"
+
+    def load_lora(
+        self, model, lora_name, strength_model=1.0, nickname="", target="all"
+    ):
+        """Load LoRA adapter using SGLang Diffusion API."""
+        lora_path = folder_paths.get_full_path("loras", lora_name)
+        assert model is not None
+        bi = model.clone()
+        nickname = nickname if nickname != "" else str("lora" + str(uuid.uuid4()))
+        # set lora in the model
+        bi.patches[nickname] = (lora_path, strength_model, target)
+
+        # prepare input for the SGLang Diffusion API
+        lora_input = {
+            "lora_nickname": [],
+            "lora_path": [],
+            "strength": [],
+            "target": [],
+        }
+        for nickname, lora_info in bi.patches.items():
+            lora_input["lora_nickname"].append(nickname)
+            lora_input["lora_path"].append(lora_info[0])
+            lora_input["strength"].append(lora_info[1])
+            lora_input["target"].append(lora_info[2])
+
+        # call the SGLang Diffusion API
+        model.model.diffusion_model.set_lora(**lora_input)
+        return (model,)
+
+
 class SGLDUNETLoader:
     def __init__(self):
         self.generator = SGLDiffusionGenerator()
@@ -646,6 +700,7 @@ NODE_CLASS_MAPPINGS = {
     "SGLDiffusionServerUnsetLora": SGLDiffusionServerUnsetLora,
     "SGLDUNETLoader": SGLDUNETLoader,
     "SGLDOptions": SGLDOptions,
+    "SGLDLoraLoader": SGLDLoraLoader,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -656,4 +711,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "SGLDiffusionServerUnsetLora": "SGLDiffusion Server Unset LoRA",
     "SGLDUNETLoader": "SGLDiffusion UNET Loader",
     "SGLDOptions": "SGLDiffusion Options",
+    "SGLDLoraLoader": "SGLDiffusion LoRA Loader",
 }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/README.md

@@ -0,0 +1,66 @@
+# ComfyUI SGLDiffusion Pipeline Tests
+
+This directory contains tests for each ComfyUI pipeline integration.
+
+## Test Files
+
+- `test_zimage_pipeline.py` - Tests for ComfyUIZImagePipeline
+- `test_flux_pipeline.py` - Tests for ComfyUIFluxPipeline
+- `test_qwen_image_pipeline.py` - Tests for ComfyUIQwenImagePipeline
+- `test_qwen_image_edit_pipeline.py` - Tests for ComfyUIQwenImageEditPipeline (I2I/edit mode)
+
+## Running Tests
+
+### Run all tests
+
+```bash
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/ -v -s
+```
+
+### Run a specific test file
+
+```bash
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+```
+
+## Environment Variables
+
+You can configure model paths via environment variables. Model paths support two formats:
+- **Safetensors file**: Path to a single `.safetensors` file (e.g., `/path/to/model.safetensors`)
+- **Diffusers format**: HuggingFace model ID or local diffusers directory (e.g., `Tongyi-MAI/Z-Image-Turbo`)
+
+Environment variables:
+- `SGLANG_TEST_ZIMAGE_MODEL_PATH` - Path to ZImage model (default: `Tongyi-MAI/Z-Image-Turbo`)
+- `SGLANG_TEST_FLUX_MODEL_PATH` - Path to Flux model (default: `black-forest-labs/FLUX.1-dev`)
+- `SGLANG_TEST_QWEN_IMAGE_MODEL_PATH` - Path to QwenImage model (default: `Qwen/Qwen-Image`)
+- `SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH` - Path to QwenImageEdit model (default: `Qwen/Qwen-Image-Edit-2511`)
+
+Examples:
+
+```bash
+# Using HuggingFace model ID (diffusers format)
+export SGLANG_TEST_ZIMAGE_MODEL_PATH="Tongyi-MAI/Z-Image-Turbo"
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+
+# Using safetensors file
+export SGLANG_TEST_ZIMAGE_MODEL_PATH="/path/to/z_image_turbo_bf16.safetensors"
+pytest python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py -v -s
+```
+
+## Test Structure
+
+Each test file follows a similar structure:
+
+1. **Setup**: Creates a `DiffGenerator` with the appropriate pipeline class
+2. **Input Preparation**: Creates dummy tensors for latents, timesteps, and embeddings
+3. **Request Preparation**: Uses `prepare_request` to convert `SamplingParams` to `Req`
+4. **ComfyUI Inputs**: Sets ComfyUI-specific inputs directly on the `Req` object
+5. **Execution**: Sends request to scheduler and waits for response
+6. **Validation**: Checks that `noise_pred` is retrieved from `OutputBatch`
+
+## Notes
+
+- These tests use `comfyui_mode=True` to enable ComfyUI-specific behavior
+- Tests use pre-processed inputs (latents, timesteps, embeddings) as ComfyUI would provide
+- The tests verify that `noise_pred` can be retrieved from the `OutputBatch` after processing
+- All tests use dummy/ones tensors for simplicity - in production, these would be actual model outputs

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

@@ -0,0 +1,9 @@
+"""
+Test suite for ComfyUI SGLDiffusion pipelines.
+
+This package contains tests for each ComfyUI pipeline integration:
+- ZImagePipeline
+- FluxPipeline
+- QwenImagePipeline
+- QwenImageEditPipeline
+"""

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_flux_pipeline.py

@@ -0,0 +1,162 @@
+"""Test for ComfyUIFluxPipeline with pass-through scheduler."""
+
+import os
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_flux_pipeline_direct() -> None:
+    """Test ComfyUIFluxPipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_FLUX_MODEL_PATH",
+        "black-forest-labs/FLUX.1-dev",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIFluxPipeline",
+        num_gpus=2,
+        comfyui_mode=True,
+    )
+
+    batch_size = 1
+    hidden_states_seq_len = 3600
+    hidden_states_dim = 64
+    height = 1280
+    width = 720
+
+    encoder_seq_len = 512
+    encoder_dim = 4096
+    pooled_dim = 768
+
+    hidden_states = torch.ones(
+        batch_size,
+        hidden_states_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    pooled_projections = torch.ones(
+        batch_size,
+        pooled_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt="a beautiful girl",
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        save_output=True,
+        return_trajectory_latents=True,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = hidden_states
+    req.timesteps = timesteps
+    req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+    clip_dim = 768
+    dummy_clip_embedding = torch.zeros(
+        batch_size,
+        77,
+        clip_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    req.prompt_embeds = [pooled_projections, encoder_hidden_states]
+
+    if req.guidance_scale > 1.0:
+        dummy_neg_clip_embedding = torch.zeros(
+            batch_size,
+            77,
+            clip_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        negative_encoder_hidden_states = torch.ones(
+            batch_size,
+            encoder_seq_len,
+            encoder_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        req.negative_prompt_embeds = [
+            dummy_neg_clip_embedding,
+            negative_encoder_hidden_states,
+        ]
+    else:
+        req.negative_prompt_embeds = None
+
+    req.pooled_embeds = [pooled_projections]
+    req.neg_pooled_embeds = []
+
+    if (
+        req.guidance_scale > 1.0
+        and req.negative_prompt_embeds is not None
+        and len(req.negative_prompt_embeds) > 0
+    ):
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cuda" if generator_device == "cuda" else "cpu"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print(f"✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+    print(f"latents.shape: {latents.shape}")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_edit_pipeline.py

@@ -0,0 +1,135 @@
+"""Test for ComfyUIQwenImageEditPipeline with pass-through scheduler (I2I/edit mode)."""
+
+import os
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_qwen_image_edit_pipeline_direct() -> None:
+    """Test ComfyUIQwenImageEditPipeline with edit mode (I2I) and custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_QWEN_IMAGE_EDIT_MODEL_PATH",
+        "Qwen/Qwen-Image-Edit-2511",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIQwenImageEditPipeline",
+        num_gpus=1,
+        comfyui_mode=True,
+        dit_layerwise_offload=False,
+    )
+
+    batch_size = 1
+    noisy_image_seq_len = 3600
+    hidden_states_dim = 64
+    condition_image_seq_len = 6889
+    condition_image_dim = 64
+    encoder_seq_len = 45
+    encoder_dim = 3584
+    height = 720
+    width = 1280
+
+    vae_scale_factor = 8
+    condition_height_latent = 1328 // vae_scale_factor
+    condition_width_latent = 1328 // vae_scale_factor
+
+    noisy_image_latents = torch.ones(
+        batch_size,
+        noisy_image_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    condition_image_latents = torch.ones(
+        batch_size,
+        condition_image_seq_len,
+        condition_image_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt=" ",
+        guidance_scale=1.0,
+        height=height,
+        width=width,
+        image_path="",
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = noisy_image_latents
+    req.image_latent = condition_image_latents
+    req.timesteps = timesteps
+    req.prompt_embeds = [encoder_hidden_states]
+    req.negative_prompt_embeds = None
+    req.vae_image_sizes = [(condition_width_latent, condition_height_latent)]
+    req.raw_latent_shape = torch.tensor(noisy_image_latents.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print(f"✓ Successfully retrieved noise_pred from OutputBatch (Edit Mode)!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_qwen_image_pipeline.py

@@ -0,0 +1,119 @@
+"""Test for ComfyUIQwenImagePipeline with pass-through scheduler."""
+
+import os
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_qwen_image_pipeline_direct() -> None:
+    """Test ComfyUIQwenImagePipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_QWEN_IMAGE_MODEL_PATH",
+        "Qwen/Qwen-Image",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIQwenImagePipeline",
+        num_gpus=2,
+        comfyui_mode=True,
+        dit_layerwise_offload=False,
+    )
+
+    batch_size = 1
+    hidden_states_seq_len = 6889
+    hidden_states_dim = 64
+    encoder_seq_len = 45
+    encoder_dim = 3584
+    height = 1328
+    width = 1328
+    dtype = torch.bfloat16
+
+    hidden_states = torch.ones(
+        batch_size,
+        hidden_states_seq_len,
+        hidden_states_dim,
+        device="cuda",
+        dtype=dtype,
+    )
+
+    encoder_hidden_states = torch.ones(
+        batch_size,
+        encoder_seq_len,
+        encoder_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt=" ",
+        guidance_scale=3.0,
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = hidden_states
+    req.timesteps = timesteps
+    req.prompt_embeds = [encoder_hidden_states]
+    req.negative_prompt_embeds = [encoder_hidden_states]
+    req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print(f"✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/test/test_zimage_pipeline.py

@@ -0,0 +1,121 @@
+"""Test for ComfyUIZImagePipeline with pass-through scheduler."""
+
+import os
+
+import pytest
+import torch
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+from sglang.multimodal_gen.runtime.entrypoints.diffusion_generator import DiffGenerator
+from sglang.multimodal_gen.runtime.entrypoints.utils import prepare_request
+
+
+def test_comfyui_zimage_pipeline_direct() -> None:
+    """Test ComfyUIZImagePipeline with custom inputs."""
+    model_path = os.environ.get(
+        "SGLANG_TEST_ZIMAGE_MODEL_PATH",
+        "Tongyi-MAI/Z-Image-Turbo",  # Supports both safetensors file and diffusers format
+    )
+
+    generator = DiffGenerator.from_pretrained(
+        model_path=model_path,
+        pipeline_class_name="ComfyUIZImagePipeline",
+        num_gpus=1,
+        sp_degree=1,
+        comfyui_mode=True,
+    )
+
+    batch_size = 1
+    num_channels = 16
+    num_frames = 1
+    height = 720
+    width = 1280
+    latent_height = height // 8
+    latent_width = width // 8
+
+    latents = torch.ones(
+        batch_size,
+        num_channels,
+        num_frames,
+        latent_height,
+        latent_width,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    timesteps = torch.tensor([1000], dtype=torch.long, device="cuda")
+
+    context_seq_len = 19
+    context_dim = 2560
+    context = torch.ones(
+        context_seq_len,
+        context_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+
+    sampling_params = SamplingParams.from_user_sampling_params_args(
+        generator.server_args.model_path,
+        server_args=generator.server_args,
+        prompt="a beautiful girl",
+        guidance_scale=1.0,
+        height=height,
+        width=width,
+        num_frames=1,
+        num_inference_steps=1,
+        seed=42,
+        save_output=False,
+        return_frames=False,
+    )
+
+    req = prepare_request(
+        server_args=generator.server_args,
+        sampling_params=sampling_params,
+    )
+
+    req.latents = latents
+    req.timesteps = timesteps
+    req.prompt_embeds = [context]
+    req.negative_prompt_embeds = None
+    req.raw_latent_shape = torch.tensor(latents.shape, dtype=torch.long)
+
+    if req.guidance_scale > 1.0 and req.negative_prompt_embeds is not None:
+        req.do_classifier_free_guidance = True
+    else:
+        req.do_classifier_free_guidance = False
+
+    if req.seed is not None:
+        generator_device = req.generator_device
+        device_str = "cpu" if generator_device == "cpu" else "cuda"
+        req.generator = [
+            torch.Generator(device_str).manual_seed(req.seed + i)
+            for i in range(req.num_outputs_per_prompt)
+        ]
+    else:
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+    output_batch = generator._send_to_scheduler_and_wait_for_response([req])
+    noise_pred = output_batch.noise_pred
+
+    assert noise_pred is not None, "noise_pred should not be None in OutputBatch"
+    assert isinstance(noise_pred, torch.Tensor), "noise_pred should be a torch.Tensor"
+    assert (
+        noise_pred.device.type == "cuda"
+    ), f"noise_pred should be on cuda, got {noise_pred.device}"
+    assert (
+        noise_pred.dtype == torch.bfloat16
+    ), f"noise_pred should be bfloat16, got {noise_pred.dtype}"
+
+    print(f"✓ Successfully retrieved noise_pred from OutputBatch!")
+    print(f"  noise_pred shape: {noise_pred.shape}")
+    print(f"  noise_pred dtype: {noise_pred.dtype}")
+    print(f"  noise_pred device: {noise_pred.device}")
+
+    latents = output_batch.output if output_batch.output is not None else req.latents
+    assert latents is not None, "latents should not be None"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/flux_sgld_sp.json

@@ -0,0 +1,222 @@
+{
+    "8": {
+      "inputs": {
+        "samples": [
+          "40",
+          0
+        ],
+        "vae": [
+          "10",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "10": {
+      "inputs": {
+        "vae_name": "ae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "Load VAE"
+      }
+    },
+    "11": {
+      "inputs": {
+        "clip_name1": "t5xxl_fp16.safetensors",
+        "clip_name2": "clip_l.safetensors",
+        "type": "flux",
+        "device": "default"
+      },
+      "class_type": "DualCLIPLoader",
+      "_meta": {
+        "title": "DualCLIPLoader"
+      }
+    },
+    "17": {
+      "inputs": {
+        "scheduler": "normal",
+        "steps": 25,
+        "denoise": 1,
+        "model": [
+          "46",
+          0
+        ]
+      },
+      "class_type": "BasicScheduler",
+      "_meta": {
+        "title": "BasicScheduler"
+      }
+    },
+    "38": {
+      "inputs": {
+        "model": [
+          "46",
+          0
+        ],
+        "conditioning": [
+          "42",
+          0
+        ]
+      },
+      "class_type": "BasicGuider",
+      "_meta": {
+        "title": "BasicGuider"
+      }
+    },
+    "39": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "8",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "40": {
+      "inputs": {
+        "noise": [
+          "45",
+          0
+        ],
+        "guider": [
+          "38",
+          0
+        ],
+        "sampler": [
+          "47",
+          0
+        ],
+        "sigmas": [
+          "17",
+          0
+        ],
+        "latent_image": [
+          "44",
+          0
+        ]
+      },
+      "class_type": "SamplerCustomAdvanced",
+      "_meta": {
+        "title": "SamplerCustomAdvanced"
+      }
+    },
+    "42": {
+      "inputs": {
+        "guidance": 3.5,
+        "conditioning": [
+          "43",
+          0
+        ]
+      },
+      "class_type": "FluxGuidance",
+      "_meta": {
+        "title": "FluxGuidance"
+      }
+    },
+    "43": {
+      "inputs": {
+        "text": "beautiful photography of a gonger haired artist with Lots of Colorful coloursplashes in face and pn her hands, she is natural, having her hair in a casual bun, looking happily into camera, cinematic,",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "11",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Prompt)"
+      }
+    },
+    "44": {
+      "inputs": {
+        "width": 1024,
+        "height": 1024,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "45": {
+      "inputs": {
+        "noise_seed": 747172083610812
+      },
+      "class_type": "RandomNoise",
+      "_meta": {
+        "title": "RandomNoise"
+      }
+    },
+    "46": {
+      "inputs": {
+        "max_shift": 1.15,
+        "base_shift": 0.5,
+        "width": 1024,
+        "height": 1024,
+        "model": [
+          "51",
+          0
+        ]
+      },
+      "class_type": "ModelSamplingFlux",
+      "_meta": {
+        "title": "ModelSamplingFlux"
+      }
+    },
+    "47": {
+      "inputs": {
+        "sampler_name": "euler"
+      },
+      "class_type": "KSamplerSelect",
+      "_meta": {
+        "title": "KSamplerSelect"
+      }
+    },
+    "51": {
+      "inputs": {
+        "unet_name": "flux1-dev.safetensors",
+        "weight_dtype": "default",
+        "sgld_options": [
+          "52",
+          0
+        ]
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    },
+    "52": {
+      "inputs": {
+        "model_type": "auto-detect",
+        "enable_torch_compile": false,
+        "num_gpus": 2,
+        "tp_size": -1,
+        "sp_degree": -1,
+        "ulysses_degree": -1,
+        "ring_degree": -1,
+        "dp_size": 1,
+        "dp_degree": 1,
+        "enable_cfg_parallel": false,
+        "attention_backend": "",
+        "cache_strategy": "none"
+      },
+      "class_type": "SGLDOptions",
+      "_meta": {
+        "title": "SGLDiffusion Options"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/qwen_image_sgld.json

@@ -0,0 +1,165 @@
+{
+    "3": {
+      "inputs": {
+        "seed": 808633539418610,
+        "steps": 4,
+        "cfg": 1,
+        "sampler_name": "euler",
+        "scheduler": "simple",
+        "denoise": 1,
+        "model": [
+          "66",
+          0
+        ],
+        "positive": [
+          "6",
+          0
+        ],
+        "negative": [
+          "7",
+          0
+        ],
+        "latent_image": [
+          "58",
+          0
+        ]
+      },
+      "class_type": "KSampler",
+      "_meta": {
+        "title": "KSampler"
+      }
+    },
+    "6": {
+      "inputs": {
+        "text": "\"A vibrant, warm neon-lit street scene in Hong Kong at the afternoon, with a mix of colorful Chinese and English signs glowing brightly. The atmosphere is lively, cinematic, and rain-washed with reflections on the pavement. The colors are vivid, full of pink, blue, red, and green hues. Crowded buildings with overlapping neon signs. 1980s Hong Kong style. Signs include:\n\"龍鳳冰室\" \"金華燒臘\" \"HAPPY HAIR\" \"鴻運茶餐廳\" \"EASY BAR\" \"永發魚蛋粉\" \"添記粥麵\" \"SUNSHINE MOTEL\" \"美都餐室\" \"富記糖水\" \"太平館\" \"雅芳髮型屋\" \"STAR KTV\" \"銀河娛樂城\" \"百樂門舞廳\" \"BUBBLE CAFE\" \"萬豪麻雀館\" \"CITY LIGHTS BAR\" \"瑞祥香燭莊\" \"文記文具\" \"GOLDEN JADE HOTEL\" \"LOVELY BEAUTY\" \"合興百貨\" \"興旺電器\" And the background is warm yellow street and with all stores' lights on.",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "38",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Positive Prompt)"
+      }
+    },
+    "7": {
+      "inputs": {
+        "text": "",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "38",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Negative Prompt)"
+      }
+    },
+    "8": {
+      "inputs": {
+        "samples": [
+          "3",
+          0
+        ],
+        "vae": [
+          "39",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "38": {
+      "inputs": {
+        "clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors",
+        "type": "qwen_image",
+        "device": "default"
+      },
+      "class_type": "CLIPLoader",
+      "_meta": {
+        "title": "Load CLIP"
+      }
+    },
+    "39": {
+      "inputs": {
+        "vae_name": "qwen_image_vae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "Load VAE"
+      }
+    },
+    "58": {
+      "inputs": {
+        "width": 1328,
+        "height": 1328,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "60": {
+      "inputs": {
+        "filename_prefix": "ComfyUI"
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "66": {
+      "inputs": {
+        "shift": 3.1000000000000005,
+        "model": [
+          "78",
+          0
+        ]
+      },
+      "class_type": "ModelSamplingAuraFlow",
+      "_meta": {
+        "title": "ModelSamplingAuraFlow"
+      }
+    },
+    "77": {
+      "inputs": {
+        "unet_name": "qwen_image_2512_bf16.safetensors",
+        "weight_dtype": "default"
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    },
+    "78": {
+      "inputs": {
+        "lora_name": "Qwen-Image-2512-Lightning-4steps-V1.0-bf16.safetensors",
+        "strength_model": 1,
+        "nickname": "",
+        "target": "all",
+        "model": [
+          "77",
+          0
+        ]
+      },
+      "class_type": "SGLDLoraLoader",
+      "_meta": {
+        "title": "SGLDiffusion LoRA Loader"
+      }
+    }
+  }

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/z-image_sgld.json

@@ -0,0 +1,140 @@
+{
+    "3": {
+      "inputs": {
+        "seed": 3338398,
+        "steps": 9,
+        "cfg": 1,
+        "sampler_name": "euler",
+        "scheduler": "simple",
+        "denoise": 1,
+        "model": [
+          "28",
+          0
+        ],
+        "positive": [
+          "6",
+          0
+        ],
+        "negative": [
+          "7",
+          0
+        ],
+        "latent_image": [
+          "13",
+          0
+        ]
+      },
+      "class_type": "KSampler",
+      "_meta": {
+        "title": "KSampler"
+      }
+    },
+    "6": {
+      "inputs": {
+        "text": "cute anime style girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black gold leaf pattern dress and a white apron, it is a postcard held by a hand in front of a beautiful realistic city at sunset and there is cursive writing that says \"ZImage, Now in ComfyUI\"",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "18",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Positive Prompt)"
+      }
+    },
+    "7": {
+      "inputs": {
+        "text": "blurry ugly bad",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        },
+        "clip": [
+          "18",
+          0
+        ]
+      },
+      "class_type": "CLIPTextEncode",
+      "_meta": {
+        "title": "CLIP Text Encode (Negative Prompt)"
+      }
+    },
+    "8": {
+      "inputs": {
+        "samples": [
+          "3",
+          0
+        ],
+        "vae": [
+          "17",
+          0
+        ]
+      },
+      "class_type": "VAEDecode",
+      "_meta": {
+        "title": "VAE Decode"
+      }
+    },
+    "9": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "8",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "Save Image"
+      }
+    },
+    "13": {
+      "inputs": {
+        "width": 1024,
+        "height": 1024,
+        "batch_size": 1
+      },
+      "class_type": "EmptySD3LatentImage",
+      "_meta": {
+        "title": "EmptySD3LatentImage"
+      }
+    },
+    "17": {
+      "inputs": {
+        "vae_name": "ae.safetensors"
+      },
+      "class_type": "VAELoader",
+      "_meta": {
+        "title": "VAE Loader"
+      }
+    },
+    "18": {
+      "inputs": {
+        "clip_name": "qwen_3_4b.safetensors",
+        "type": "lumina2",
+        "device": "default"
+      },
+      "class_type": "CLIPLoader",
+      "_meta": {
+        "title": "CLIP Loader"
+      }
+    },
+    "28": {
+      "inputs": {
+        "unet_name": "z_image_turbo_bf16.safetensors",
+        "weight_dtype": "default"
+      },
+      "class_type": "SGLDUNETLoader",
+      "_meta": {
+        "title": "SGLDiffusion UNET Loader"
+      }
+    }
+}

python/sglang/multimodal_gen/configs/pipeline_configs/base.py

@@ -131,8 +131,14 @@ def shard_rotary_emb_for_sp(emb):
 
 def maybe_unpad_latents(latents, batch):
     # If SP padding was applied, remove extra tokens before reshaping
-    width, height = batch.raw_latent_shape[-1], batch.raw_latent_shape[-2]
-    target_tokens = width * height
+    raw_shape = batch.raw_latent_shape
+    if len(raw_shape) == 3:
+        # Sequence format [B, S, D]: use seq_len directly
+        target_tokens = raw_shape[1]
+    else:
+        # Spatial format [B, C, H, W] or [B, C, T, H, W]: use width * height
+        width, height = raw_shape[-1], raw_shape[-2]
+        target_tokens = width * height
     if latents.shape[1] > target_tokens:
         latents = latents[:, :target_tokens, :]
     return latents

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -200,6 +200,7 @@ class ZImagePipelineConfig(ImagePipelineConfig):
 
     def gather_latents_for_sp(self, latents):
         # Gather on effective-H dim=3 (matches shard_latents_for_sp); swap-back is handled in post_denoising_loop.
+        latents = latents.contiguous()
         if get_sp_world_size() <= 1 or latents.dim() != 5:
             return latents
         return sequence_model_parallel_all_gather(latents, dim=3)

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -149,6 +149,8 @@ class SamplingParams:
     no_override_protected_fields: bool = False
     # whether to adjust num_frames for multi-GPU friendly splitting (default: True)
     adjust_frames: bool = True
+    # if True, suppress verbose logging for this request
+    suppress_logs: bool = False
 
     def _set_output_file_ext(self):
         # add extension if needed

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