[diffusion] app: add ComfyUI plugin support for SGLang-Diffusion (#15271)

Co-authored-by: niehen6174 <niehen.6174@gmail.com> Co-authored-by: Mick <mickjagger19@icloud.com> Co-authored-by: niehen6174 <nihen6174@gmail.com>
2026-01-12 21:58:16 +08:00
parent 9a628744fc
commit 76f69b7753
30 changed files with 3688 additions and 12 deletions
@@ -154,6 +154,7 @@ async def generations(
                ImageResponseData(
                    b64_json=b64,
                    revised_prompt=request.prompt,
+                    file_path=os.path.abspath(save_file_path),
                )
            ],
        }
@@ -264,14 +265,24 @@ async def edits(
            with open(save_file_path, "rb") as f:
                b64 = base64.b64encode(f.read()).decode("utf-8")
                response_kwargs["data"].append(
-                    ImageResponseData(b64_json=b64, revised_prompt=prompt)
+                    ImageResponseData(
+                        b64_json=b64,
+                        revised_prompt=prompt,
+                        file_path=os.path.abspath(save_file_path),
+                    )
                )
        if result.peak_memory_mb and result.peak_memory_mb > 0:
            response_kwargs["peak_memory_mb"] = result.peak_memory_mb
    else:
        url = f"/v1/images/{request_id}/content"
        response_kwargs = {
-            "data": [ImageResponseData(url=url, revised_prompt=prompt)],
+            "data": [
+                ImageResponseData(
+                    url=url,
+                    revised_prompt=prompt,
+                    file_path=os.path.abspath(save_file_path),
+                )
+            ],
        }

    response_kwargs = add_common_data_to_response(
@@ -9,6 +9,7 @@ class ImageResponseData(BaseModel):
    b64_json: Optional[str] = None
    url: Optional[str] = None
    revised_prompt: Optional[str] = None
+    file_path: Optional[str] = None


 class ImageResponse(BaseModel):
@@ -57,6 +58,7 @@ class VideoResponse(BaseModel):
    completed_at: Optional[int] = None
    expires_at: Optional[int] = None
    error: Optional[Dict[str, Any]] = None
+    file_path: Optional[str] = None
    peak_memory_mb: Optional[float] = None


@@ -79,6 +81,7 @@ class VideoGenerationsRequest(BaseModel):
    )
    negative_prompt: Optional[str] = None
    enable_teacache: Optional[bool] = False
+    output_path: Optional[str] = None
    diffusers_kwargs: Optional[Dict[str, Any]] = None  # kwargs for diffusers backend


@@ -85,6 +85,8 @@ def _build_sampling_params_from_request(
        sampling_kwargs["negative_prompt"] = request.negative_prompt
    if request.enable_teacache is not None:
        sampling_kwargs["enable_teacache"] = request.enable_teacache
+    if request.output_path is not None:
+        sampling_kwargs["output_path"] = request.output_path
    sampling_params = SamplingParams.from_user_sampling_params_args(
        model_path=server_args.model_path,
        server_args=server_args,
@@ -117,7 +119,7 @@ def _video_job_from_sampling(
        "size": size_str,
        "seconds": str(seconds),
        "quality": "standard",
-        "file_path": sampling.output_file_path(),
+        "file_path": os.path.abspath(sampling.output_file_path()),
    }


@@ -132,6 +132,7 @@ class GPUWorker:
                    timings=result.timings,
                    trajectory_timesteps=getattr(result, "trajectory_timesteps", None),
                    trajectory_latents=getattr(result, "trajectory_latents", None),
+                    noise_pred=getattr(result, "noise_pred", None),
                    trajectory_decoded=getattr(result, "trajectory_decoded", None),
                )
            else:
@@ -0,0 +1,193 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+"""
+Pass-through scheduler for ComfyUI integration.
+
+This scheduler does not modify latents - it simply returns the input sample unchanged.
+The actual denoising logic is handled by ComfyUI.
+"""
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+from diffusers.utils import BaseOutput
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class ComfyUIPassThroughSchedulerOutput(BaseOutput):
+    """
+    Output class for the scheduler's `step` function output.
+
+    Args:
+        prev_sample (`torch.FloatTensor`): The input sample unchanged (pass-through).
+    """
+
+    prev_sample: torch.FloatTensor
+
+
+class ComfyUIPassThroughScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
+    """
+    Pass-through scheduler for ComfyUI integration.
+
+    This scheduler does not modify latents. It is used when the denoising logic
+    is handled externally by ComfyUI. The scheduler simply returns the input
+    sample unchanged, allowing ComfyUI to manage the denoising process.
+
+    Usage:
+        - num_inference_steps is always 1 (each step is handled separately)
+        - timesteps are provided externally by ComfyUI
+        - step() returns the input sample unchanged
+    """
+
+    config_name = "scheduler_config.json"
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps=1000,
+        *args,
+        **kwargs,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        # Initialize timesteps as empty - will be set externally
+        self.timesteps = torch.tensor([], dtype=torch.long)
+        self.shift = 0.0
+        self._step_index = 0  # Track current step index
+        self._begin_index: int | None = None  # For compatibility with DenoisingStage
+
+    def set_timesteps(
+        self,
+        num_inference_steps=1,  # Always 1 for ComfyUI
+        timesteps=None,  # Can be provided externally
+        device=None,
+        **kwargs,
+    ):
+        """
+        Set timesteps. For ComfyUI, timesteps are provided externally.
+
+        Args:
+            num_inference_steps: Ignored (always 1 for ComfyUI)
+            timesteps: External timesteps provided by ComfyUI
+            device: Device to place timesteps on
+        """
+        if timesteps is not None:
+            # Use externally provided timesteps
+            if isinstance(timesteps, torch.Tensor):
+                self.timesteps = timesteps
+            else:
+                self.timesteps = torch.tensor(timesteps, dtype=torch.long)
+            if device is not None:
+                self.timesteps = self.timesteps.to(device)
+        else:
+            # Create a single timestep if none provided
+            if device is None:
+                device = torch.device("cpu")
+            self.timesteps = torch.tensor([0], dtype=torch.long, device=device)
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: torch.FloatTensor | int,
+        sample: torch.FloatTensor,
+        return_dict: bool = False,
+        **kwargs,
+    ) -> tuple | ComfyUIPassThroughSchedulerOutput:
+        """
+        Pass-through step: returns the input sample unchanged.
+
+        This scheduler does not modify latents. The actual denoising is handled
+        by ComfyUI, so we simply return the input sample as-is.
+
+        Args:
+            model_output: Predicted noise (ignored, but kept for API compatibility)
+            timestep: Current timestep (ignored, but kept for API compatibility)
+            sample: Input latents (returned unchanged)
+            return_dict: Whether to return a dict or tuple
+
+        Returns:
+            The input sample unchanged (prev_sample = sample)
+        """
+        # Increment step index for tracking
+        self._step_index += 1
+
+        # Simply return the input sample unchanged
+        prev_sample = sample
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return ComfyUIPassThroughSchedulerOutput(prev_sample=prev_sample)
+
+    def scale_model_input(
+        self, sample: torch.Tensor, timestep: int | None = None
+    ) -> torch.Tensor:
+        """
+        Scale model input. For pass-through scheduler, returns input unchanged.
+
+        Args:
+            sample: Input sample
+            timestep: Timestep (ignored)
+
+        Returns:
+            Input sample unchanged
+        """
+        return sample
+
+    def set_shift(self, shift: float) -> None:
+        """
+        Set shift parameter (no-op for pass-through scheduler).
+
+        Args:
+            shift: Shift value (ignored)
+        """
+        self.shift = shift
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index: The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    @property
+    def begin_index(self) -> int | None:
+        """
+        The index for the first timestep.
+        """
+        return self._begin_index
+
+    @property
+    def step_index(self) -> int:
+        """
+        The index counter for current timestep.
+        """
+        return self._step_index
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Add noise to samples. For pass-through scheduler, returns original samples.
+
+        Args:
+            original_samples: Original clean samples
+            noise: Noise to add (ignored)
+            timestep: Timestep (ignored)
+
+        Returns:
+            Original samples unchanged
+        """
+        return original_samples
+
+
+EntryClass = ComfyUIPassThroughScheduler
@@ -0,0 +1,645 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import re
+from typing import Any, Generator
+
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.flux import FluxConfig
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+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
+
+logger = init_logger(__name__)
+
+
+class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Simplified pipeline for ComfyUI integration with only denoising stage.
+
+    This pipeline requires pre-processed inputs:
+    - prompt_embeds: Pre-encoded text embeddings (list of tensors)
+    - negative_prompt_embeds: Pre-encoded negative prompt embeddings (if using CFG)
+    - latents: Optional initial noise latents (will be generated if not provided)
+
+    Usage:
+        generator = DiffGenerator.from_pretrained(
+            model_path="path/to/model",
+            pipeline_class_name="ComfyUIFluxPipeline",
+            device="cuda",
+        )
+    """
+
+    pipeline_name = "ComfyUIFluxPipeline"
+
+    # Configuration classes for safetensors files without model_index.json
+    from sglang.multimodal_gen.configs.pipeline_configs.flux import FluxPipelineConfig
+    from sglang.multimodal_gen.configs.sample.flux import FluxSamplingParams
+
+    pipeline_config_cls = FluxPipelineConfig
+    sampling_params_cls = FluxSamplingParams
+
+    _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
+        )
+
+        if hasattr(server_args.pipeline_config, "vae_config"):
+            vae_config = server_args.pipeline_config.vae_config
+            if hasattr(vae_config, "post_init") and not hasattr(
+                vae_config, "_post_init_called"
+            ):
+                vae_config.post_init()
+                logger.info(
+                    "Called vae_config.post_init() to set spatial_compression_ratio. "
+                    f"spatial_compression_ratio={vae_config.arch_config.spatial_compression_ratio}"
+                )
+
+    def load_modules(
+        self,
+        server_args: ServerArgs,
+        loaded_modules: dict[str, torch.nn.Module] | None = None,
+    ) -> dict[str, Any]:
+        """
+        Load modules for ComfyUIFluxPipeline.
+
+        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_and_convert_weights_from_safetensors(
+        self,
+        model_cls: type,
+        dit_config: FluxConfig,
+        hf_config: dict,
+        safetensors_list: list[str],
+        updated_mapping: dict,
+        qkv_size: int,
+        mlp_hidden_dim: int,
+        has_guidance_embeds: bool,
+        default_dtype: torch.dtype,
+    ) -> tuple[torch.nn.Module, dict]:
+        """
+        Load and convert weights from safetensors file, then load them into the model.
+        """
+        from sglang.multimodal_gen.runtime.loader.utils import (
+            get_param_names_mapping,
+            set_default_torch_dtype,
+        )
+        from sglang.multimodal_gen.runtime.loader.weight_utils import (
+            safetensors_weights_iterator,
+        )
+
+        logger.info(
+            "Converting ComfyUI Flux weights to SGLang format and loading model..."
+        )
+
+        # Create model on target device
+        device = get_local_torch_device()
+        with set_default_torch_dtype(default_dtype):
+            model = model_cls(**{"config": dit_config, "hf_config": hf_config})
+            model = model.to(device)
+
+        # Verify model has guidance_embedder if config says it should
+        has_guidance_embedder = hasattr(model.time_text_embed, "guidance_embedder")
+        if has_guidance_embeds and not has_guidance_embedder:
+            logger.warning(
+                "Config has guidance_embeds=True but model doesn't have guidance_embedder. "
+                "This may indicate a configuration mismatch."
+            )
+        elif not has_guidance_embeds and has_guidance_embedder:
+            logger.warning(
+                "Config has guidance_embeds=False but model has guidance_embedder. "
+                "This may indicate a configuration mismatch."
+            )
+
+        # Note: guidance_in mappings are already included in comfyui_flux_mappings above.
+        # If model doesn't support guidance embeddings, the weights will be filtered out
+        # in _convert_comfyui_weights() based on has_guidance_embeds flag.
+
+        param_names_mapping_fn = get_param_names_mapping(updated_mapping)
+
+        weight_iterator = safetensors_weights_iterator(safetensors_list)
+        converted_weights = self._convert_comfyui_weights(
+            weight_iterator=weight_iterator,
+            qkv_size=qkv_size,
+            mlp_hidden_dim=mlp_hidden_dim,
+            has_guidance_embeds=has_guidance_embeds,
+        )
+
+        model_state_dict = model.state_dict()
+        missing_keys = set(model_state_dict.keys())
+        unexpected_keys = []
+        loaded_count = 0
+        reverse_param_names_mapping = {}
+
+        # Handle merged parameters (collect all parts before merging)
+        from collections import defaultdict
+
+        to_merge_params = defaultdict(dict)
+
+        # Process weights incrementally: load immediately after conversion
+        for source_name, tensor in converted_weights:
+            target_name, merge_index, num_params_to_merge = param_names_mapping_fn(
+                source_name
+            )
+            reverse_param_names_mapping[target_name] = (
+                source_name,
+                merge_index,
+                num_params_to_merge,
+            )
+
+            if merge_index is not None:
+                # Collect parts for merging
+                to_merge_params[target_name][merge_index] = tensor
+                if len(to_merge_params[target_name]) == num_params_to_merge:
+                    # All parts collected, merge them
+                    sorted_tensors = [
+                        to_merge_params[target_name][i]
+                        for i in range(num_params_to_merge)
+                    ]
+                    merged_tensor = torch.cat(sorted_tensors, dim=0)
+                    # Load immediately after merging
+                    if target_name in model_state_dict:
+                        param = model_state_dict[target_name]
+                        loaded_tensor = merged_tensor.to(
+                            device=param.device, dtype=param.dtype
+                        )
+                        param.data.copy_(loaded_tensor)
+                        missing_keys.discard(target_name)
+                        loaded_count += 1
+                        del merged_tensor, loaded_tensor
+                    else:
+                        unexpected_keys.append(target_name)
+                    # Clear merged parts
+                    del to_merge_params[target_name]
+                    for t in sorted_tensors:
+                        del t
+            else:
+                # Direct mapping, load immediately
+                if target_name in model_state_dict:
+                    param = model_state_dict[target_name]
+                    # Check shape compatibility
+                    if tensor.shape != param.shape:
+                        logger.warning(
+                            f"Shape mismatch for {target_name}: "
+                            f"loaded {tensor.shape} vs model {param.shape}, skipping. "
+                            f"Source: {source_name}"
+                        )
+                        unexpected_keys.append(target_name)
+                        del tensor
+                        continue
+
+                    # Debug logging for norm_out.linear to verify mapping
+                    if (
+                        "norm_out.linear" in target_name
+                        or "final_layer.adaLN_modulation" in source_name
+                    ):
+                        logger.info(
+                            f"Loading norm_out.linear: {source_name} -> {target_name}, "
+                            f"shape: {tensor.shape}"
+                        )
+
+                    loaded_tensor = tensor.to(device=param.device, dtype=param.dtype)
+                    param.data.copy_(loaded_tensor)
+                    missing_keys.discard(target_name)
+                    loaded_count += 1
+                    del tensor, loaded_tensor
+                else:
+                    # Debug logging for unmapped parameters
+                    if "norm_out.linear" in target_name:
+                        logger.warning(
+                            f"norm_out.linear parameter {target_name} not found in model state_dict. "
+                            f"Source: {source_name}"
+                        )
+                    unexpected_keys.append(target_name)
+
+        optional_missing_keys = []
+        required_missing_keys = []
+        for key in missing_keys:
+            if key.endswith(".bias"):
+                # Check if corresponding weight exists (if weight exists but bias doesn't, it's optional)
+                weight_key = key.replace(".bias", ".weight")
+                if weight_key not in missing_keys:
+                    optional_missing_keys.append(key)
+                else:
+                    required_missing_keys.append(key)
+            else:
+                required_missing_keys.append(key)
+
+        if required_missing_keys:
+            logger.warning(
+                f"Required missing keys (first 10): {required_missing_keys[:10]}..."
+            )
+        if optional_missing_keys:
+            logger.info(
+                f"Optional missing keys (bias parameters, {len(optional_missing_keys)} total): "
+                f"These will use default values (zeros)"
+            )
+        if unexpected_keys:
+            logger.warning(f"Unexpected keys (first 10): {unexpected_keys[:10]}...")
+
+        logger.info(f"Successfully loaded {loaded_count} weight tensors")
+
+        return model, reverse_param_names_mapping
+
+    def _convert_comfyui_weights(
+        self,
+        weight_iterator: Generator[tuple[str, torch.Tensor], None, None],
+        qkv_size: int,
+        mlp_hidden_dim: int,
+        has_guidance_embeds: bool,
+    ) -> 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.
+        Filters out guidance_in weights if model doesn't support guidance embeddings.
+        Handles scale/shift order difference between ComfyUI and AdaLayerNormContinuous.
+        """
+        for name, tensor in weight_iterator:
+            if not has_guidance_embeds and name.startswith("guidance_in."):
+                logger.debug(
+                    f"Skipping {name} (model doesn't support guidance embeddings)"
+                )
+                continue
+
+            match = re.match(r"single_blocks\.(\d+)\.linear1\.(weight|bias)$", name)
+            if match:
+                block_idx, param_type = match.groups()
+                expected_size = qkv_size + mlp_hidden_dim
+
+                if tensor.shape[0] < expected_size:
+                    logger.warning(
+                        f"linear1.{param_type} shape {tensor.shape} doesn't match "
+                        f"expected size {expected_size}, skipping"
+                    )
+                    continue
+
+                # Split tensor
+                qkv_tensor = (
+                    tensor[:qkv_size] if param_type == "bias" else tensor[:qkv_size, :]
+                )
+                mlp_tensor = (
+                    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
+                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]
+                # AdaLayerNormContinuous: expects [scale, shift]
+                # Need to swap the first half and second half of the weight matrix
+                # Weight shape: (2 * hidden_size, hidden_size)
+                # Split into two halves and swap them
+                half_size = tensor.shape[0] // 2
+                shift_weights = tensor[:half_size, :]
+                scale_weights = tensor[half_size:, :]
+                # Swap: put scale first, then shift
+                swapped_tensor = torch.cat([scale_weights, shift_weights], dim=0)
+                logger.info(
+                    f"Swapped scale/shift order for {name}: "
+                    f"shape {tensor.shape} -> {swapped_tensor.shape}"
+                )
+                yield name, swapped_tensor
+            elif name == "final_layer.adaLN_modulation.1.bias":
+                # Same swap for bias: (2 * hidden_size,)
+                half_size = tensor.shape[0] // 2
+                shift_bias = tensor[:half_size]
+                scale_bias = tensor[half_size:]
+                swapped_tensor = torch.cat([scale_bias, shift_bias], dim=0)
+                logger.info(
+                    f"Swapped scale/shift order for {name}: "
+                    f"shape {tensor.shape} -> {swapped_tensor.shape}"
+                )
+                yield name, swapped_tensor
+            else:
+                # Other weights pass through (handled by param_names_mapping)
+                yield name, tensor
+
+    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 file without model_index.json.
+        """
+        if loaded_modules is not None and "transformer" in loaded_modules:
+            logger.info("Using provided transformer module")
+            components = {
+                "transformer": loaded_modules["transformer"],
+                "scheduler": self.modules.get("scheduler"),
+            }
+            return components
+
+        if hasattr(server_args.pipeline_config, "dit_config"):
+            dit_config = server_args.pipeline_config.dit_config
+            if not isinstance(dit_config, FluxConfig):
+                logger.warning("dit_config is not FluxConfig, creating new FluxConfig")
+                dit_config = FluxConfig()
+                server_args.pipeline_config.dit_config = dit_config
+        else:
+            logger.info("Creating default FluxConfig")
+            dit_config = FluxConfig()
+            server_args.pipeline_config.dit_config = dit_config
+
+        # Set guidance_embeds to True for ComfyUI Flux models
+        dit_config.arch_config.guidance_embeds = True
+        logger.info("Set guidance_embeds=True for ComfyUI Flux model")
+
+        if dit_config.arch_config.param_names_mapping is None:
+            dit_config.arch_config.param_names_mapping = {}
+
+        # ComfyUI Flux uses different parameter names than SGLang Flux
+        # Key differences:
+        # - ComfyUI: single_blocks.{i}.linear1 (fused QKV + MLP input)
+        # - SGLang: single_transformer_blocks.{i}.attn.to_qkv + proj_mlp (separate)
+        # - ComfyUI: single_blocks.{i}.linear2
+        # - SGLang: single_transformer_blocks.{i}.proj_out
+        # - ComfyUI: double_blocks.{i}.img_attn.qkv / txt_attn.qkv
+        # - SGLang: transformer_blocks.{i}.attn.to_qkv / attn.to_added_qkv
+
+        # Note: For fused layers like linear1, we need custom weight splitting logic
+        # which will be handled in the weight conversion function below
+        comfyui_flux_mappings = {
+            # Double stream blocks - attention layers
+            r"double_blocks\.(\d+)\.img_attn\.qkv\.(weight|bias)$": (
+                r"transformer_blocks.\1.attn.to_qkv.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_attn\.qkv\.(weight|bias)$": (
+                r"transformer_blocks.\1.attn.to_added_qkv.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.img_attn\.proj\.(weight|bias)$": (
+                r"transformer_blocks.\1.attn.to_out.0.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_attn\.proj\.(weight|bias)$": (
+                r"transformer_blocks.\1.attn.to_add_out.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.img_attn\.norm\.query_norm\.scale$": (
+                r"transformer_blocks.\1.attn.norm_q.weight",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.img_attn\.norm\.key_norm\.scale$": (
+                r"transformer_blocks.\1.attn.norm_k.weight",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_attn\.norm\.query_norm\.scale$": (
+                r"transformer_blocks.\1.attn.norm_added_q.weight",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_attn\.norm\.key_norm\.scale$": (
+                r"transformer_blocks.\1.attn.norm_added_k.weight",
+                None,
+                None,
+            ),
+            # Double stream blocks - MLP layers (map to net structure)
+            r"double_blocks\.(\d+)\.img_mlp\.0\.(weight|bias)$": (
+                r"transformer_blocks.\1.ff.net.0.proj.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.img_mlp\.2\.(weight|bias)$": (
+                r"transformer_blocks.\1.ff.net.2.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_mlp\.0\.(weight|bias)$": (
+                r"transformer_blocks.\1.ff_context.net.0.proj.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_mlp\.2\.(weight|bias)$": (
+                r"transformer_blocks.\1.ff_context.net.2.\2",
+                None,
+                None,
+            ),
+            # Double stream blocks - modulation layers
+            r"double_blocks\.(\d+)\.img_mod\.lin\.(weight|bias)$": (
+                r"transformer_blocks.\1.norm1.linear.\2",
+                None,
+                None,
+            ),
+            r"double_blocks\.(\d+)\.txt_mod\.lin\.(weight|bias)$": (
+                r"transformer_blocks.\1.norm1_context.linear.\2",
+                None,
+                None,
+            ),
+            # Single stream blocks - linear2 maps to proj_out
+            r"single_blocks\.(\d+)\.linear2\.(weight|bias)$": (
+                r"single_transformer_blocks.\1.proj_out.\2",
+                None,
+                None,
+            ),
+            # Single stream blocks - norm layers (scale -> weight)
+            r"single_blocks\.(\d+)\.norm\.query_norm\.scale$": (
+                r"single_transformer_blocks.\1.attn.norm_q.weight",
+                None,
+                None,
+            ),
+            r"single_blocks\.(\d+)\.norm\.key_norm\.scale$": (
+                r"single_transformer_blocks.\1.attn.norm_k.weight",
+                None,
+                None,
+            ),
+            # Single stream blocks - modulation (maps to norm.linear)
+            r"single_blocks\.(\d+)\.modulation\.lin\.(weight|bias)$": (
+                r"single_transformer_blocks.\1.norm.linear.\2",
+                None,
+                None,
+            ),
+            # Time and guidance embeddings
+            r"^time_in\.in_layer\.(weight|bias)$": (
+                r"time_text_embed.timestep_embedder.linear_1.\1",
+                None,
+                None,
+            ),
+            r"^time_in\.out_layer\.(weight|bias)$": (
+                r"time_text_embed.timestep_embedder.linear_2.\1",
+                None,
+                None,
+            ),
+            r"^txt_in\.(weight|bias)$": (r"context_embedder.\1", None, None),
+            r"^vector_in\.in_layer\.(weight|bias)$": (
+                r"time_text_embed.text_embedder.linear_1.\1",
+                None,
+                None,
+            ),
+            r"^vector_in\.out_layer\.(weight|bias)$": (
+                r"time_text_embed.text_embedder.linear_2.\1",
+                None,
+                None,
+            ),
+            # Final layer mappings
+            r"^final_layer\.linear\.(weight|bias)$": (r"proj_out.\1", None, None),
+            r"^final_layer\.norm_final\.(weight|bias)$": (r"norm_out.\1", None, None),
+            r"^final_layer\.adaLN_modulation\.1\.(weight|bias)$": (
+                r"norm_out.linear.\1",
+                None,
+                None,
+            ),
+            # Image input embedding
+            r"^img_in\.(weight|bias)$": (r"x_embedder.\1", None, None),
+            # Guidance embeddings (if model supports guidance)
+            r"^guidance_in\.in_layer\.(weight|bias)$": (
+                r"time_text_embed.guidance_embedder.linear_1.\1",
+                None,
+                None,
+            ),
+            r"^guidance_in\.out_layer\.(weight|bias)$": (
+                r"time_text_embed.guidance_embedder.linear_2.\1",
+                None,
+                None,
+            ),
+        }
+
+        # Merge ComfyUI mappings with existing mappings (ComfyUI mappings take precedence)
+        updated_mapping = {
+            **dit_config.arch_config.param_names_mapping,
+            **comfyui_flux_mappings,
+        }
+        dit_config.arch_config.param_names_mapping = updated_mapping
+        logger.info(
+            "Added ComfyUI weight name mappings for Flux model. "
+            f"Total mappings: {len(updated_mapping)}"
+        )
+
+        cls_name = "FluxTransformer2DModel"
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+        logger.info("Resolved transformer class: %s", cls_name)
+
+        original_mapping = None
+        if comfyui_flux_mappings:
+            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)}"
+            )
+
+        safetensors_list = [self.model_path]
+        logger.info("Loading weights from: %s", safetensors_list)
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+        server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "."
+        hf_config = {}
+
+        hidden_size = (
+            dit_config.arch_config.num_attention_heads
+            * dit_config.arch_config.attention_head_dim
+        )
+        mlp_ratio = getattr(dit_config.arch_config, "mlp_ratio", 4.0)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        qkv_size = 3 * hidden_size
+        has_guidance_embeds = True
+
+        # Load and convert weights from safetensors file
+        model, reverse_param_names_mapping = (
+            self._load_and_convert_weights_from_safetensors(
+                model_cls=model_cls,
+                dit_config=dit_config,
+                hf_config=hf_config,
+                safetensors_list=safetensors_list,
+                updated_mapping=updated_mapping,
+                qkv_size=qkv_size,
+                mlp_hidden_dim=mlp_hidden_dim,
+                has_guidance_embeds=has_guidance_embeds,
+                default_dtype=default_dtype,
+            )
+        )
+
+        model = model.eval()
+        for param in model.parameters():
+            param.requires_grad = False
+
+        model.reverse_param_names_mapping = reverse_param_names_mapping
+
+        if original_mapping is not None:
+            model_cls.param_names_mapping = original_mapping
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9)
+
+        components = {
+            "transformer": model,
+            "scheduler": self.modules.get("scheduler"),
+        }
+
+        logger.info("Successfully loaded modules: %s", list(components.keys()))
+        return components
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        logger.info(
+            "ComfyUIFluxPipeline.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"ComfyUIFluxPipeline stages created: {list(self._stage_name_mapping.keys())}"
+        )
+
+
+EntryClass = ComfyUIFluxPipeline
@@ -0,0 +1,390 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import re
+from collections.abc import Generator
+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.zimage import ZImageDitConfig
+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 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 ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Simplified pipeline for ComfyUI integration with only denoising stage.
+
+    This pipeline requires pre-processed inputs:
+    - prompt_embeds: Pre-encoded text embeddings (list of tensors)
+    - negative_prompt_embeds: Pre-encoded negative prompt embeddings (if using CFG)
+    - latents: Optional initial noise latents (will be generated if not provided)
+
+    Usage:
+        generator = DiffGenerator.from_pretrained(
+            model_path="path/to/model",
+            pipeline_class_name="ComfyUIZImagePipeline",
+            device="cuda",
+        )
+    """
+
+    pipeline_name = "ComfyUIZImagePipeline"
+    from sglang.multimodal_gen.configs.pipeline_configs.zimage import (
+        ZImagePipelineConfig,
+    )
+    from sglang.multimodal_gen.configs.sample.zimage import ZImageSamplingParams
+
+    pipeline_config_cls = ZImagePipelineConfig
+    sampling_params_cls = ZImageSamplingParams
+
+    _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
+        # This is necessary because get_freqs_cis uses spatial_compression_ratio
+        if hasattr(server_args.pipeline_config, "vae_config"):
+            vae_config = server_args.pipeline_config.vae_config
+            if hasattr(vae_config, "post_init") and not hasattr(
+                vae_config, "_post_init_called"
+            ):
+                vae_config.post_init()
+                logger.info(
+                    "Called vae_config.post_init() to set spatial_compression_ratio. "
+                    f"spatial_compression_ratio={vae_config.arch_config.spatial_compression_ratio}"
+                )
+
+    def load_modules(
+        self,
+        server_args: ServerArgs,
+        loaded_modules: dict[str, torch.nn.Module] | None = None,
+    ) -> dict[str, Any]:
+        """
+        Load modules for ComfyUIZImagePipeline.
+
+        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 _convert_comfyui_qkv_weights(
+        self,
+        weight_iterator: Generator[tuple[str, torch.Tensor], None, None],
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+    ) -> Generator[tuple[str, torch.Tensor], None, None]:
+        """
+        Convert ComfyUI zimage qkv weights to SGLang format.
+        Splits merged qkv.weight into separate to_q, to_k, to_v weights.
+
+        Args:
+            weight_iterator: Iterator yielding (name, tensor) pairs from safetensors
+            dim: Model dimension
+            num_heads: Number of attention heads
+            num_kv_heads: Number of key-value heads
+
+        Yields:
+            (name, tensor) pairs with qkv weights split into to_q, to_k, to_v
+        """
+        head_dim = dim // num_heads
+        q_size = dim
+        k_size = head_dim * num_kv_heads
+        v_size = head_dim * num_kv_heads
+
+        for name, tensor in weight_iterator:
+            # Match qkv weights in layers, noise_refiner, or context_refiner
+            # Pattern: (layers|noise_refiner|context_refiner).{i}.attention.qkv.(weight|bias)
+            match = re.match(
+                r"(layers|noise_refiner|context_refiner)\.(\d+)\.attention\.qkv\.(weight|bias)$",
+                name,
+            )
+            if match:
+                module_name, layer_idx, param_type = match.groups()
+                base_name = f"{module_name}.{layer_idx}.attention"
+
+                if param_type == "weight":
+                    # Weight shape: (q_size + k_size + v_size, dim)
+                    # Split into q, k, v
+                    q_weight = tensor[:q_size, :]
+                    k_weight = tensor[q_size : q_size + k_size, :]
+                    v_weight = tensor[q_size + k_size :, :]
+
+                    logger.debug(
+                        f"Splitting {name} (shape {tensor.shape}) into "
+                        f"to_q ({q_weight.shape}), to_k ({k_weight.shape}), to_v ({v_weight.shape})"
+                    )
+
+                    yield f"{base_name}.to_q.weight", q_weight
+                    yield f"{base_name}.to_k.weight", k_weight
+                    yield f"{base_name}.to_v.weight", v_weight
+                else:  # bias
+                    # Bias shape: (q_size + k_size + v_size,)
+                    # Split into q, k, v
+                    q_bias = tensor[:q_size]
+                    k_bias = tensor[q_size : q_size + k_size]
+                    v_bias = tensor[q_size + k_size :]
+
+                    logger.debug(
+                        f"Splitting {name} (shape {tensor.shape}) into "
+                        f"to_q ({q_bias.shape}), to_k ({k_bias.shape}), to_v ({v_bias.shape})"
+                    )
+
+                    yield f"{base_name}.to_q.bias", q_bias
+                    yield f"{base_name}.to_k.bias", k_bias
+                    yield f"{base_name}.to_v.bias", v_bias
+            else:
+                # Pass through other weights unchanged
+                yield name, tensor
+
+    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 file without model_index.json.
+
+        This method:
+        1. Uses hardcoded ZImageDitConfig for zimage model
+        2. Loads transformer from the safetensors file
+        3. Uses ComfyUIPassThroughScheduler (already created in initialize_pipeline)
+        """
+        # Check if transformer is already provided
+        if loaded_modules is not None and "transformer" in loaded_modules:
+            logger.info("Using provided transformer module")
+            components = {
+                "transformer": loaded_modules["transformer"],
+                "scheduler": self.modules.get("scheduler"),
+            }
+            return components
+
+        if hasattr(server_args.pipeline_config, "dit_config"):
+            dit_config = server_args.pipeline_config.dit_config
+            if not isinstance(dit_config, ZImageDitConfig):
+                logger.warning(
+                    "dit_config is not ZImageDitConfig, creating new ZImageDitConfig"
+                )
+                dit_config = ZImageDitConfig()
+                server_args.pipeline_config.dit_config = dit_config
+        else:
+            logger.info("Creating default ZImageDitConfig")
+            dit_config = ZImageDitConfig()
+            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 = {}
+
+        # Add mappings for norm layers: map from ComfyUI format (k_norm/q_norm) to SGLang format (norm_k/norm_q)
+        # The regex matches the source name from safetensors, and the tuple specifies the target name in the model
+        # Note: qkv weights are handled separately by _convert_comfyui_qkv_weights function
+        comfyui_norm_mappings = {
+            r"(.*)\.attention\.k_norm\.weight$": (
+                r"\1.attention.norm_k.weight",
+                None,
+                None,
+            ),
+            r"(.*)\.attention\.q_norm\.weight$": (
+                r"\1.attention.norm_q.weight",
+                None,
+                None,
+            ),
+            r"(.*)\.attention\.out\.weight$": (
+                r"\1.attention.to_out.0.weight",
+                None,
+                None,
+            ),
+            r"^final_layer\.(.*)$": (r"all_final_layer.2-1.\1", None, None),
+            r"^x_embedder\.(.*)$": (r"all_x_embedder.2-1.\1", None, None),
+        }
+
+        # Merge ComfyUI mappings with existing mappings (ComfyUI mappings take precedence)
+        updated_mapping = {
+            **dit_config.arch_config.param_names_mapping,
+            **comfyui_norm_mappings,
+        }
+        dit_config.arch_config.param_names_mapping = updated_mapping
+        logger.info(
+            "Added ComfyUI weight name mappings (k_norm/q_norm -> norm_k/norm_q) to param_names_mapping. "
+            f"Total mappings: {len(updated_mapping)}"
+        )
+
+        cls_name = "ZImageTransformer2DModel"
+        model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
+        logger.info("Resolved transformer class: %s", cls_name)
+        safetensors_list = [self.model_path]
+        logger.info("Loading weights from: %s", safetensors_list)
+
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+        server_args.model_paths["transformer"] = os.path.dirname(self.model_path) or "."
+        hf_config = {}
+
+        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,
+        )
+
+        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:
+            # Create model first (same as maybe_load_fsdp_model)
+            from sglang.multimodal_gen.runtime.platforms import current_platform
+
+            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})
+
+            # Check if we should use FSDP
+            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:
+                world_size = server_args.hsdp_replicate_dim * server_args.hsdp_shard_dim
+                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,
+                )
+
+            # Get model dimensions for qkv splitting
+            arch_config = dit_config.arch_config
+            dim = arch_config.dim
+            num_heads = arch_config.num_attention_heads
+            num_kv_heads = arch_config.n_kv_heads
+
+            # Create weight iterator with qkv conversion
+            base_weight_iterator = safetensors_weights_iterator(safetensors_list)
+            converted_weight_iterator = self._convert_comfyui_qkv_weights(
+                base_weight_iterator, dim, num_heads, num_kv_heads
+            )
+
+            # Load weights
+            param_names_mapping_fn = get_param_names_mapping(updated_mapping)
+            load_model_from_full_model_state_dict(
+                model,
+                converted_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
+        finally:
+            model_cls.param_names_mapping = original_mapping
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded transformer with %.2fB parameters", total_params / 1e9)
+
+        components = {
+            "transformer": model,
+            "scheduler": self.modules.get("scheduler"),
+        }
+
+        logger.info("Successfully loaded modules: %s", list(components.keys()))
+        return components
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        logger.info(
+            "ComfyUIZImagePipeline.create_pipeline_stages() called - creating only denoising_stage"
+        )
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+        logger.info(
+            f"ComfyUIZImagePipeline stages created: {list(self._stage_name_mapping.keys())}"
+        )
+
+
+EntryClass = ComfyUIZImagePipeline
@@ -42,11 +42,33 @@ def build_pipeline(
    3. based on the config, determine the pipeline class
    """
    model_path = server_args.model_path
-    model_info = get_model_info(model_path, backend=server_args.backend)
-    if model_info is None:
-        raise ValueError(f"Unsupported model: {model_path}")

-    pipeline_cls = model_info.pipeline_cls
+    # Check if pipeline class is explicitly specified
+    if server_args.pipeline_class_name:
+        from sglang.multimodal_gen.registry import (
+            _PIPELINE_REGISTRY,
+            _discover_and_register_pipelines,
+        )
+
+        _discover_and_register_pipelines()
+        logger.info(f"Requested pipeline_class_name: {server_args.pipeline_class_name}")
+        logger.info(
+            f"Available pipelines in registry: {list(_PIPELINE_REGISTRY.keys())}"
+        )
+        pipeline_cls = _PIPELINE_REGISTRY.get(server_args.pipeline_class_name)
+        if pipeline_cls is None:
+            raise ValueError(
+                f"Pipeline class '{server_args.pipeline_class_name}' not found in registry. "
+                f"Available pipelines: {list(_PIPELINE_REGISTRY.keys())}"
+            )
+        logger.info(
+            f"✓ Using explicitly specified pipeline: {server_args.pipeline_class_name} (class: {pipeline_cls.__name__})"
+        )
+    else:
+        logger.info("No pipeline_class_name specified, using model_index.json")
+        model_info = get_model_info(model_path)
+        pipeline_cls = model_info.pipeline_cls
+        logger.info(f"Using pipeline from model_index.json: {pipeline_cls.__name__}")

    # instantiate the pipelines
    pipeline = pipeline_cls(model_path, server_args)
@@ -300,4 +300,7 @@ class OutputBatch:

    # logged timings info, directly from Req.timings
    timings: Optional["RequestTimings"] = None
+
+    # For ComfyUI integration: noise prediction from denoising stage
+    noise_pred: torch.Tensor | None = None
    peak_memory_mb: float = 0.0
@@ -12,6 +12,9 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineSta
 from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import (
    CausalDMDDenoisingStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.comfyui_latent_preparation import (
+    ComfyUILatentPreparationStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.conditioning import (
    ConditioningStage,
 )
@@ -43,6 +46,7 @@ __all__ = [
    "InputValidationStage",
    "TimestepPreparationStage",
    "LatentPreparationStage",
+    "ComfyUILatentPreparationStage",
    "ConditioningStage",
    "DenoisingStage",
    "DmdDenoisingStage",
@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+ComfyUI latent preparation stage with device mismatch fix.
+This stage extends LatentPreparationStage to handle device mismatch issues
+that occur when tensors are pickled and unpickled via broadcast_pyobj in
+multi-GPU scenarios.
+"""
+import dataclasses
+
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_sp_group,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_sp_world_size
+from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
+    LatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class ComfyUILatentPreparationStage(LatentPreparationStage):
+    """
+    ComfyUI-specific latent preparation stage with device mismatch fix.
+
+    This stage extends LatentPreparationStage to automatically fix device
+    mismatches for tensor fields on non-source ranks in multi-GPU scenarios.
+    """
+
+    @staticmethod
+    def _fix_tensor_device(value, target_device):
+        """Recursively fix tensor device, handling single tensors, lists, and tuples."""
+        if isinstance(value, torch.Tensor):
+            if value.device != target_device:
+                return value.detach().clone().to(target_device)
+            return value
+        elif isinstance(value, list):
+            return [
+                ComfyUILatentPreparationStage._fix_tensor_device(v, target_device)
+                for v in value
+            ]
+        elif isinstance(value, tuple):
+            return tuple(
+                ComfyUILatentPreparationStage._fix_tensor_device(v, target_device)
+                for v in value
+            )
+        return value
+
+    @staticmethod
+    def _has_tensor(value):
+        """Check if value contains any tensor."""
+        if isinstance(value, torch.Tensor):
+            return True
+        elif isinstance(value, (list, tuple)):
+            return any(ComfyUILatentPreparationStage._has_tensor(v) for v in value)
+        return False
+
+    def forward(self, batch, server_args):
+        """
+        Prepare latents with device mismatch fix for ComfyUI pipelines.
+
+        This method first fixes device mismatches for all tensor fields,
+        then calls the parent class's forward method, and ensures raw_latent_shape
+        is set correctly (before packing, for proper unpadding later).
+        """
+        # Fix device mismatch for tensor fields on non-source ranks
+        if get_sp_world_size() > 1:
+            sp_group = get_sp_group()
+            target_device = get_local_torch_device()
+
+            if sp_group.rank != 0:
+                logger.debug(
+                    f"[ComfyUILatentPreparationStage] Fixing tensor device on rank={sp_group.rank} "
+                    f"target_device={target_device}"
+                )
+
+                if dataclasses.is_dataclass(batch):
+                    for field in dataclasses.fields(batch):
+                        value = getattr(batch, field.name, None)
+                        if value is not None and self._has_tensor(value):
+                            fixed_value = self._fix_tensor_device(value, target_device)
+                            setattr(batch, field.name, fixed_value)
+                else:
+                    for attr_name in dir(batch):
+                        if not attr_name.startswith("_") and not callable(
+                            getattr(batch, attr_name, None)
+                        ):
+                            try:
+                                value = getattr(batch, attr_name, None)
+                                if value is not None and self._has_tensor(value):
+                                    fixed_value = self._fix_tensor_device(
+                                        value, target_device
+                                    )
+                                    setattr(batch, attr_name, fixed_value)
+                            except (AttributeError, TypeError):
+                                continue
+
+        original_latents_shape = None
+        if batch.latents is not None:
+            original_latents_shape = batch.latents.shape
+
+        # Call parent class's forward method
+        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
+
+        return result
@@ -683,6 +683,23 @@ class DenoisingStage(PipelineStage):
            batch, latents, trajectory_tensor
        )

+        # Gather noise_pred if using sequence parallelism
+        # noise_pred has the same shape as latents (sharded along sequence dimension)
+        if (
+            get_sp_world_size() > 1
+            and getattr(batch, "did_sp_shard_latents", False)
+            and server_args.comfyui_mode
+            and hasattr(batch, "noise_pred")
+            and batch.noise_pred is not None
+        ):
+            batch.noise_pred = server_args.pipeline_config.gather_latents_for_sp(
+                batch.noise_pred
+            )
+            if hasattr(batch, "raw_latent_shape"):
+                orig_s = batch.raw_latent_shape[1]
+                if batch.noise_pred.shape[1] > orig_s:
+                    batch.noise_pred = batch.noise_pred[:, :orig_s, :]
+
        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
            batch.trajectory_latents = trajectory_tensor.cpu()
@@ -1029,6 +1046,10 @@ class DenoisingStage(PipelineStage):
                            latents=latents,
                        )

+                        # Save noise_pred to batch for external access (e.g., ComfyUI)
+                        if server_args.comfyui_mode:
+                            batch.noise_pred = noise_pred
+
                        # Compute the previous noisy sample
                        latents = self.scheduler.step(
                            model_output=noise_pred,
@@ -265,6 +265,11 @@ class ServerArgs:

    pipeline_config: PipelineConfig = field(default_factory=PipelineConfig, repr=False)

+    # Pipeline override
+    pipeline_class_name: str | None = (
+        None  # Override pipeline class from model_index.json
+    )
+
    # LoRA parameters
    # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
    lora_path: str | None = None
@@ -285,6 +290,9 @@ class ServerArgs:
    use_fsdp_inference: bool = False
    pin_cpu_memory: bool = True

+    # ComfyUI integration
+    comfyui_mode: bool = False
+
    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE