[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>

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

@@ -0,0 +1,49 @@
+# ComfyUI SGLDiffusion Plugin
+
+A ComfyUI plugin for integrating with SGLang Diffusion server, supporting image and video generation capabilities.
+
+## Installation
+
+1. Copy this entire directory (`ComfyUI_SGLDiffusion`) to your ComfyUI `custom_nodes/` folder
+2. Restart ComfyUI to load the plugin
+
+## Usage
+
+### Prerequisites
+
+Before using this plugin, you need to start the SGLang Diffusion server. The plugin connects to the server via HTTP API calls.
+
+### Basic Workflow
+
+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
+
+### 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
+
+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
+
+## 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.

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

@@ -0,0 +1,7 @@
+"""
+ComfyUI SGLang Diffusion nodes package.
+"""
+
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

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

@@ -0,0 +1,14 @@
+"""
+Core components for SGLang Diffusion ComfyUI integration.
+Provides generator, model patcher, and server API client.
+"""
+
+from .generator import SGLDiffusionGenerator
+from .model_patcher import SGLDModelPatcher
+from .server_api import SGLDiffusionServerAPI
+
+__all__ = [
+    "SGLDiffusionGenerator",
+    "SGLDModelPatcher",
+    "SGLDiffusionServerAPI",
+]

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py

@@ -0,0 +1,219 @@
+"""
+Generator for SGLang Diffusion ComfyUI integration.
+"""
+
+import logging
+import os
+
+import psutil
+from comfy import model_detection, model_management
+from comfy.utils import (
+    calculate_parameters,
+    load_torch_file,
+    state_dict_prefix_replace,
+    unet_to_diffusers,
+)
+
+logger = logging.getLogger(__name__)
+
+try:
+    from sglang.multimodal_gen import DiffGenerator
+except ImportError:
+    logger.error(
+        "Error: sglang.multimodal_gen is not installed. Please install it using 'pip install sglang[diffusion]'"
+    )
+
+from ..executors import FluxExecutor, ZImageExecutor
+from .model_patcher import SGLDModelPatcher
+
+
+class SGLDiffusionGenerator:
+    """Generator for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(self):
+        self.model_path = None
+        self.generator = None
+        self.executor = None
+        self.last_options = None
+
+        self.pipeline_class_dict = {
+            "flux": "ComfyUIFluxPipeline",
+            "lumina2": "ComfyUIZImagePipeline",  # zimage
+        }
+        self.executor_class_dict = {
+            "flux": FluxExecutor,
+            "lumina2": ZImageExecutor,
+        }
+
+    def __del__(self):
+        self.close_generator()
+
+    def init_generator(
+        self, model_path: str, pipeline_class_name: str, kwargs: dict = None
+    ):
+        """Initialize the diffusion generator."""
+        if self.generator is not None:
+            return self.generator
+        if kwargs is None:
+            kwargs = {}
+        # Set comfyui_mode for ComfyUI integration
+        kwargs["comfyui_mode"] = True
+        self.generator = DiffGenerator.from_pretrained(
+            model_path=model_path,
+            pipeline_class_name=pipeline_class_name,
+            **kwargs,
+        )
+        return self.generator
+
+    def kill_generator(self):
+        """Kill worker processes manually because generator shutdown cannot terminate them."""
+        current_pid = os.getpid()
+        worker_processes = []
+        for proc in psutil.process_iter(["pid", "name", "cmdline"]):
+            try:
+                # Look for sglang-diffusionWorker processes
+                if proc.info["cmdline"]:
+                    cmdline = " ".join(proc.info["cmdline"])
+                    if "sgl_diffusion::" in cmdline:
+                        if proc.info["pid"] != current_pid:
+                            worker_processes.append(proc)
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+
+        if worker_processes:
+            logger.info(
+                f"Found {len(worker_processes)} worker processes to terminate..."
+            )
+            for proc in worker_processes:
+                try:
+                    logger.info(
+                        f"Terminating worker process {proc.info['pid']}: {proc.info['name']}"
+                    )
+                    proc.terminate()
+                    proc.wait(timeout=5)
+                except psutil.TimeoutExpired:
+                    logger.warning(
+                        f"Process {proc.info['pid']} did not terminate, forcing kill..."
+                    )
+                    try:
+                        proc.kill()
+                        proc.wait(timeout=2)
+                    except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+                        pass
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+
+    def close_generator(self):
+        """Close and cleanup the generator and all associated resources."""
+        if self.generator is not None:
+            self.generator.shutdown()
+            self.kill_generator()
+            # Clear other references
+            self.last_options = None
+            self.model_path = None
+            self.generator = None
+            self.executor = None
+
+    def get_comfyui_model(self, model_path: str, model_options: dict = None):
+        """Get ComfyUI model from model path."""
+        if model_options is None:
+            model_options = {}
+        dtype = model_options.get("dtype", None)
+        # Allow loading unets from checkpoint files
+        sd = load_torch_file(model_path)
+        diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
+        temp_sd = state_dict_prefix_replace(
+            sd, {diffusion_model_prefix: ""}, filter_keys=True
+        )
+        if len(temp_sd) > 0:
+            sd = temp_sd
+
+        parameters = calculate_parameters(sd)
+        load_device = model_management.get_torch_device()
+
+        model_detect_config = model_detection.detect_unet_config(sd, "")
+        model_type = model_detect_config.get("image_model", None)
+        if model_type is None or model_type not in self.pipeline_class_dict:
+            raise ValueError(f"Unsupported model type: {model_type}")
+        model_config = model_detection.model_config_from_unet(sd, "")
+
+        if model_config is not None:
+            new_sd = sd
+        else:
+            new_sd = model_detection.convert_diffusers_mmdit(sd, "")
+            if new_sd is not None:  # diffusers mmdit
+                model_config = model_detection.model_config_from_unet(new_sd, "")
+                if model_config is None:
+                    return None
+            else:  # diffusers unet
+                model_config = model_detection.model_config_from_diffusers_unet(sd)
+                if model_config is None:
+                    return None
+
+                diffusers_keys = unet_to_diffusers(model_config.unet_config)
+                new_sd = {}
+                for k in diffusers_keys:
+                    if k in sd:
+                        new_sd[diffusers_keys[k]] = sd.pop(k)
+        offload_device = model_management.unet_offload_device()
+        if dtype is None:
+            unet_dtype = model_management.unet_dtype(
+                model_params=parameters,
+                supported_dtypes=model_config.supported_inference_dtypes,
+            )
+        else:
+            unet_dtype = dtype
+
+        manual_cast_dtype = model_management.unet_manual_cast(
+            unet_dtype, load_device, model_config.supported_inference_dtypes
+        )
+        model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+        model_config.custom_operations = model_options.get("custom_operations", None)
+        model_config.unet_config["disable_unet_model_creation"] = True
+        comfyui_model = model_config.get_model({})
+        return comfyui_model, model_config, model_type
+
+    def load_model(
+        self, model_path: str, model_options: dict = None, sgld_options: dict = None
+    ):
+        """Load model and return model patcher."""
+        gather_options = {
+            "model_path": model_path,
+            "model_options": model_options,
+            "sgld_options": sgld_options,
+        }
+        if (
+            self.last_options is not None
+            and self.last_options == gather_options
+            and self.generator is not None
+        ):
+            return self.generator
+        else:
+            self.close_generator()
+
+        self.last_options = gather_options
+        self.model_path = model_path
+
+        comfyui_model, model_config, model_type = self.get_comfyui_model(
+            model_path, model_options
+        )
+        if model_type is None or model_type not in self.pipeline_class_dict:
+            raise ValueError(f"Unsupported model type: {model_type}")
+
+        pipeline_class_name = self.pipeline_class_dict[model_type]
+        self.generator = self.init_generator(
+            model_path, pipeline_class_name, sgld_options
+        )
+
+        executor_class = self.executor_class_dict[model_type]
+        self.executor = executor_class(
+            self.generator, model_path, comfyui_model, model_config
+        )
+        comfyui_model.diffusion_model = self.executor
+
+        load_device = model_management.get_torch_device()
+        offload_device = model_management.unet_offload_device()
+
+        return SGLDModelPatcher(
+            comfyui_model, load_device, offload_device, model_type=model_type
+        )

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/model_patcher.py

@@ -0,0 +1,82 @@
+"""
+Model patcher for SGLang Diffusion ComfyUI integration.
+"""
+
+import copy
+
+from comfy.model_patcher import ModelPatcher
+
+
+class SGLDModelPatcher(ModelPatcher):
+    """Model patcher for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(
+        self,
+        model,
+        load_device,
+        offload_device,
+        size=0,
+        weight_inplace_update=False,
+        model_type=None,
+    ):
+        super().__init__(
+            model, load_device, offload_device, size, weight_inplace_update
+        )
+        self.lora_cache = {}
+        self.model_type = model_type
+        self.model_size_dict = {
+            "flux": 27 * 1024 * 1024 * 1024,
+            "lumina2": 8 * 1024 * 1024 * 1024,
+        }
+
+    def clone(self):
+        """Clone the model patcher."""
+        n = SGLDModelPatcher(
+            self.model,
+            self.load_device,
+            self.offload_device,
+            self.size,
+            weight_inplace_update=self.weight_inplace_update,
+        )
+        n.patches = {}
+        for k in self.patches:
+            n.patches[k] = self.patches[k][:]
+        n.patches_uuid = self.patches_uuid
+
+        n.object_patches = self.object_patches.copy()
+        n.model_options = copy.deepcopy(self.model_options)
+        n.backup = self.backup
+        n.object_patches_backup = self.object_patches_backup
+        n.lora_cache = copy.copy(self.lora_cache)
+        return n
+
+    def model_size(self):
+        """Get the model size in bytes."""
+        if self.model_type in self.model_size_dict:
+            return self.model_size_dict[self.model_type]
+        else:
+            return 0
+
+    def load(
+        self,
+        device_to=None,
+        lowvram_model_memory=0,
+        force_patch_weights=False,
+        full_load=False,
+    ):
+        """Load model (no-op for SGLang Diffusion)."""
+        pass
+
+    def patch_model(
+        self,
+        device_to=None,
+        lowvram_model_memory=0,
+        load_weights=True,
+        force_patch_weights=False,
+    ):
+        """Patch model (no-op for SGLang Diffusion)."""
+        pass
+
+    def unpatch_model(self, device_to=None, unpatch_weights=True):
+        """Unpatch model (no-op for SGLang Diffusion)."""
+        pass

python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/server_api.py

@@ -0,0 +1,539 @@
+"""
+SGLang Diffusion Server API client.
+Provides a low-level interface for interacting with SGLang Diffusion HTTP server.
+"""
+
+import base64
+import io
+import os
+import time
+from typing import Any, Dict, Optional
+
+import requests
+from PIL import Image
+
+
+class SGLDiffusionServerAPI:
+    """Client for SGLang Diffusion HTTP server API."""
+
+    def __init__(self, base_url: str, api_key: str = "sk-proj-1234567890"):
+        """
+        Initialize the API client.
+
+        Args:
+            base_url: Base URL of the SGLang Diffusion server (e.g., "http://localhost:30010/v1")
+            api_key: API key for authentication (default: "sk-proj-1234567890")
+        """
+        # Ensure base_url doesn't end with /v1 if it's already there
+        if base_url.endswith("/v1"):
+            self.base_url = base_url
+        elif base_url.endswith("/v1/"):
+            self.base_url = base_url.rstrip("/")
+        else:
+            self.base_url = f"{base_url.rstrip('/')}/v1"
+
+        self.api_key = api_key
+        self.headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {api_key}",
+        }
+
+    def get_model_info(self) -> Dict[str, Any]:
+        """
+        Get information about the model served by this server.
+
+        Returns:
+            Dictionary containing model information including:
+            - model_path: Path to the model
+            - task_type: Type of task (e.g., "T2V", "I2I")
+            - pipeline_name: Name of the pipeline
+            - num_gpus: Number of GPUs
+            - dit_precision: DiT model precision
+            - vae_precision: VAE model precision
+        """
+        try:
+            # Remove /v1 from base_url for /models endpoint
+            models_url = self.base_url.removesuffix("/v1") + "/models"
+            response = requests.get(models_url, headers=self.headers, timeout=30)
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to get model info: {str(e)}")
+
+    def generate_image(
+        self,
+        prompt: str,
+        image_path: Optional[str] = None,
+        mask_path: Optional[str] = None,
+        size: Optional[str] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        n: int = 1,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        response_format: str = "b64_json",
+        quality: Optional[str] = "auto",
+        style: Optional[str] = "vivid",
+        background: Optional[str] = "auto",
+        output_format: Optional[str] = None,
+        generator_device: Optional[str] = "cuda",
+    ) -> Dict[str, Any]:
+        """
+        Generate or edit an image using SGLang Diffusion API.
+        If image_path is provided, calls the edit endpoint; otherwise calls the generation endpoint.
+
+        Args:
+            prompt: Text prompt for image generation/editing
+            image_path: Optional path to input image file for editing. If provided, uses edit API.
+            mask_path: Optional path to mask image file (only used when image_path is provided)
+            size: Image size in format "WIDTHxHEIGHT" (e.g., "1024x1024")
+            width: Image width (used if size is not provided)
+            height: Image height (used if size is not provided)
+            n: Number of images to generate (1-10)
+            negative_prompt: Negative prompt to avoid certain elements
+            guidance_scale: Classifier-free guidance scale
+            num_inference_steps: Number of denoising steps
+            seed: Random seed for reproducible generation
+            enable_teacache: Enable TEA cache acceleration
+            response_format: Response format ("b64_json" or "url")
+            quality: Image quality ("auto", "standard", "hd") - only for generation
+            style: Image style ("vivid" or "natural") - only for generation
+            background: Background type ("auto", "transparent", "opaque")
+            output_format: Output format ("png", "jpeg", "webp")
+            generator_device: Device for random generator ("cuda" or "cpu")
+
+        Returns:
+            Dictionary containing the API response with generated/edited image data
+        """
+        if not prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        # Determine size
+        if size is None:
+            if width is not None and height is not None:
+                size = f"{width}x{height}"
+            else:
+                size = "1024x1024"
+
+        # Build common parameters
+        common_params = self._build_image_common_params(
+            prompt=prompt,
+            size=size,
+            n=n,
+            response_format=response_format,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            seed=seed,
+            enable_teacache=enable_teacache,
+            background=background,
+            output_format=output_format,
+            generator_device=generator_device,
+        )
+
+        # If image_path is provided, use edit endpoint
+        if image_path:
+            if not os.path.exists(image_path):
+                raise FileNotFoundError(f"Image file not found: {image_path}")
+
+            # Prepare multipart form data for edit
+            files: Dict[str, Any] = {}
+            data = common_params.copy()
+
+            # Add image file
+            files["image"] = (
+                os.path.basename(image_path),
+                open(image_path, "rb"),
+                self._get_content_type(image_path),
+            )
+
+            # Add mask file if provided
+            if mask_path:
+                if not os.path.exists(mask_path):
+                    raise FileNotFoundError(f"Mask file not found: {mask_path}")
+                files["mask"] = (
+                    os.path.basename(mask_path),
+                    open(mask_path, "rb"),
+                    self._get_content_type(mask_path),
+                )
+
+            # Prepare headers for multipart form data
+            headers = {
+                "Authorization": f"Bearer {self.api_key}",
+            }
+
+            try:
+                response = requests.post(
+                    f"{self.base_url}/images/edits",
+                    files=files,
+                    data=data,
+                    headers=headers,
+                    timeout=300,  # 5 minutes timeout for generation
+                )
+                response.raise_for_status()
+                return response.json()
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError(f"Failed to edit image: {str(e)}")
+            finally:
+                # Close file handles
+                for file_tuple in files.values():
+                    if isinstance(file_tuple, tuple) and len(file_tuple) > 1:
+                        file_tuple[1].close()
+        else:
+            # Use generation endpoint - add generation-specific parameters
+            payload = common_params.copy()
+            if quality:
+                payload["quality"] = quality
+            if style:
+                payload["style"] = style
+
+            try:
+                response = requests.post(
+                    f"{self.base_url}/images/generations",
+                    json=payload,
+                    headers=self.headers,
+                    timeout=300,  # 5 minutes timeout for generation
+                )
+                response.raise_for_status()
+                return response.json()
+            except requests.exceptions.RequestException as e:
+                raise RuntimeError(f"Failed to generate image: {str(e)}")
+
+    def generate_video(
+        self,
+        prompt: str,
+        size: Optional[str] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        seconds: Optional[int] = 4,
+        fps: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        generator_device: Optional[str] = "cuda",
+        input_reference: Optional[str] = None,
+        output_path: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Generate a video using SGLang Diffusion API and wait for completion.
+
+        Args:
+            prompt: Text prompt for video generation
+            size: Video size in format "WIDTHxHEIGHT" (e.g., "1280x720")
+            width: Video width (used if size is not provided)
+            height: Video height (used if size is not provided)
+            seconds: Duration of the video in seconds
+            fps: Frames per second
+            num_frames: Number of frames (overrides seconds * fps if provided)
+            negative_prompt: Negative prompt to avoid certain elements
+            guidance_scale: Classifier-free guidance scale
+            num_inference_steps: Number of denoising steps
+            seed: Random seed for reproducible generation
+            enable_teacache: Enable TEA cache acceleration
+            generator_device: Device for random generator ("cuda" or "cpu")
+            input_reference: Path to input reference image for image-to-video
+
+        Returns:
+            Dictionary containing completed video job information with file_path
+        """
+        if not prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        # Determine size
+        if size is None:
+            if width is not None and height is not None:
+                size = f"{width}x{height}"
+            else:
+                size = "720x1280"
+
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "prompt": prompt,
+            "size": size,
+        }
+
+        # Add optional parameters
+        if seconds is not None:
+            payload["seconds"] = seconds
+        if fps is not None:
+            payload["fps"] = fps
+        if num_frames is not None:
+            payload["num_frames"] = num_frames
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if guidance_scale is not None:
+            payload["guidance_scale"] = guidance_scale
+        if num_inference_steps is not None:
+            payload["num_inference_steps"] = num_inference_steps
+        if seed is not None and seed >= 0:
+            payload["seed"] = seed
+        if enable_teacache:
+            payload["enable_teacache"] = True
+        if generator_device:
+            payload["generator_device"] = generator_device
+        if input_reference:
+            payload["input_reference"] = input_reference
+        if output_path:
+            payload["output_path"] = output_path
+
+        try:
+            # Create video generation job
+            response = requests.post(
+                f"{self.base_url}/videos",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            video_job = response.json()
+            video_id = video_job.get("id")
+
+            # Wait for completion with fixed polling
+            poll_interval = 5  # 5 seconds
+            max_wait_time = 3600  # 1 hour
+            max_consecutive_errors = 5
+            consecutive_errors = 0
+            start_time = time.time()
+
+            while time.time() - start_time < max_wait_time:
+                try:
+                    status_response = requests.get(
+                        f"{self.base_url}/videos/{video_id}",
+                        headers=self.headers,
+                        timeout=30,
+                    )
+                    status_response.raise_for_status()
+                    status = status_response.json()
+
+                    # Reset error counter on successful request
+                    consecutive_errors = 0
+
+                    if status.get("status") == "completed":
+                        return status
+                    elif status.get("status") == "failed":
+                        error = status.get("error", {})
+                        error_msg = (
+                            error.get("message", "Unknown error")
+                            if error
+                            else "Unknown error"
+                        )
+                        raise RuntimeError(f"Video generation failed: {error_msg}")
+                except requests.exceptions.ConnectionError as e:
+                    # Connection errors - likely server is down
+                    consecutive_errors += 1
+                    if consecutive_errors >= max_consecutive_errors:
+                        raise RuntimeError(
+                            f"Lost connection to server after {consecutive_errors} consecutive errors. "
+                            f"Server may be unavailable: {str(e)}"
+                        )
+                except requests.exceptions.RequestException as e:
+                    # Other network errors - continue polling but track errors
+                    consecutive_errors += 1
+                    if consecutive_errors >= max_consecutive_errors:
+                        raise RuntimeError(
+                            f"Network error after {consecutive_errors} consecutive failures: {str(e)}"
+                        )
+
+                time.sleep(poll_interval)
+
+            raise TimeoutError(
+                f"Video generation timed out after {max_wait_time} seconds"
+            )
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to generate video: {str(e)}")
+
+    def _build_image_common_params(
+        self,
+        prompt: str,
+        size: str,
+        n: int,
+        response_format: str,
+        negative_prompt: Optional[str] = None,
+        guidance_scale: Optional[float] = None,
+        num_inference_steps: Optional[int] = None,
+        seed: Optional[int] = None,
+        enable_teacache: bool = False,
+        background: Optional[str] = None,
+        output_format: Optional[str] = None,
+        generator_device: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Build common parameters for both image generation and editing.
+
+        Returns:
+            Dictionary containing common parameters
+        """
+        params: Dict[str, Any] = {
+            "prompt": prompt,
+            "size": size,
+            "n": max(1, min(n, 10)),
+            "response_format": response_format,
+        }
+
+        # Add optional parameters
+        if negative_prompt:
+            params["negative_prompt"] = negative_prompt
+        if guidance_scale is not None:
+            params["guidance_scale"] = guidance_scale
+        if num_inference_steps is not None:
+            params["num_inference_steps"] = num_inference_steps
+        if seed is not None and seed >= 0:
+            params["seed"] = seed
+        if enable_teacache:
+            params["enable_teacache"] = True
+        if background:
+            params["background"] = background
+        if output_format:
+            params["output_format"] = output_format
+        if generator_device:
+            params["generator_device"] = generator_device
+
+        return params
+
+    def _get_content_type(self, file_path: str) -> str:
+        """Get content type based on file extension."""
+        ext = os.path.splitext(file_path)[1].lower()
+        content_types = {
+            ".png": "image/png",
+            ".jpg": "image/jpeg",
+            ".jpeg": "image/jpeg",
+            ".webp": "image/webp",
+        }
+        return content_types.get(ext, "image/png")
+
+    def decode_image_from_response(
+        self, response_data: Dict[str, Any], index: int = 0
+    ) -> Image.Image:
+        """
+        Decode base64 image from API response.
+
+        Args:
+            response_data: API response dictionary
+            index: Index of the image in the response (default: 0)
+
+        Returns:
+            PIL Image object
+        """
+        if "data" not in response_data or not response_data["data"]:
+            raise ValueError("No image data in response")
+
+        if index >= len(response_data["data"]):
+            raise IndexError(f"Image index {index} out of range")
+
+        image_data = response_data["data"][index]
+        if "b64_json" not in image_data or not image_data["b64_json"]:
+            raise ValueError("No base64 image data found")
+
+        image_bytes = base64.b64decode(image_data["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes))
+
+        # Convert to RGB if needed
+        if image.mode != "RGB":
+            image = image.convert("RGB")
+
+        return image
+
+    def set_lora(
+        self,
+        lora_nickname: str,
+        lora_path: Optional[str] = None,
+        target: str = "all",
+    ) -> Dict[str, Any]:
+        """
+        Set a LoRA adapter for the specified transformer(s).
+
+        Args:
+            lora_nickname: The nickname of the adapter (required).
+            lora_path: Path to the LoRA adapter (local path or HF repo id).
+                      Required for the first load; optional if re-activating a cached nickname.
+            target: Which transformer(s) to apply the LoRA to. One of:
+                - "all": Apply to all transformers (default)
+                - "transformer": Apply only to the primary transformer (high noise for Wan2.2)
+                - "transformer_2": Apply only to transformer_2 (low noise for Wan2.2)
+                - "critic": Apply only to the critic model
+
+        Returns:
+            Dictionary containing the API response with status and message
+        """
+        if not lora_nickname:
+            raise ValueError("lora_nickname cannot be empty")
+
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "lora_nickname": lora_nickname,
+            "target": target,
+        }
+
+        # Add optional lora_path if provided
+        if lora_path:
+            payload["lora_path"] = lora_path
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/set_lora",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}")
+
+    def unset_lora(
+        self,
+        target: str = "all",
+    ) -> Dict[str, Any]:
+        """
+        Unset (unmerge) LoRA weights from the base model.
+
+        Args:
+            target: same as set_lora
+
+        Returns:
+            Dictionary containing the API response with status and message
+        """
+        # Prepare request payload
+        payload: Dict[str, Any] = {
+            "target": target,
+        }
+
+        try:
+            response = requests.post(
+                f"{self.base_url}/unmerge_lora_weights",
+                json=payload,
+                headers=self.headers,
+                timeout=30,
+            )
+            response.raise_for_status()
+            return response.json()
+        except requests.exceptions.RequestException as e:
+            raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}")
+
+
+if __name__ == "__main__":
+    api = SGLDiffusionServerAPI(
+        base_url="http://localhost:30010/v1", api_key="sk-proj-1234567890"
+    )
+    model_info = api.get_model_info()
+    print(api.get_model_info())
+    if model_info.get("task_type") == "T2V" or model_info.get("task_type") == "I2V":
+        print(
+            api.generate_video(
+                prompt="A calico cat playing a piano on stage",
+                num_inference_steps=1,
+                size="480x480",
+            )
+        )
+    else:
+        print(
+            api.generate_image(
+                prompt="A calico cat playing a piano on stage", size="1024x1024"
+            )
+        )

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

@@ -0,0 +1,14 @@
+"""
+ComfyUI SGLang Diffusion executors package.
+Provides executor classes for different model types.
+"""
+
+from .base import SGLDiffusionExecutor
+from .flux import FluxExecutor
+from .zimage import ZImageExecutor
+
+__all__ = [
+    "SGLDiffusionExecutor",
+    "FluxExecutor",
+    "ZImageExecutor",
+]

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

@@ -0,0 +1,39 @@
+"""
+Base executor class for SGLang Diffusion ComfyUI integration.
+"""
+
+import torch
+
+
+class SGLDiffusionExecutor(torch.nn.Module):
+    """Base executor class for SGLang Diffusion models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super(SGLDiffusionExecutor, self).__init__()
+        self.generator = generator
+        self.model_path = model_path
+        self.model = model
+        self.dtype = config.unet_config["dtype"]
+        self.config = config
+        self.loras = []
+
+    def _unpack_latents(self, latents, height, width, channels):
+        """Unpack latents from packed format to standard format."""
+        batch_size = latents.shape[0]
+        latents = latents.view(batch_size, height // 2, width // 2, channels, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        latents = latents.reshape(batch_size, channels, height, width)
+
+        return latents
+
+    def _pack_latents(self, latents):
+        """Pack latents from standard format to packed format."""
+        batch_size, num_channels_latents, height, width = latents.shape
+        latents = latents.view(
+            batch_size, num_channels_latents, height // 2, 2, width // 2, 2
+        )
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size, (height // 2) * (width // 2), num_channels_latents * 4
+        )
+        return latents

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

@@ -0,0 +1,68 @@
+"""
+Flux 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]'"
+    )
+
+from .base import SGLDiffusionExecutor
+
+
+class FluxExecutor(SGLDiffusionExecutor):
+    """Executor for Flux models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(self, x, timestep, context, y=None, guidance=None, **kwargs):
+        """Forward pass for Flux model."""
+        hidden_states = self._pack_latents(x)
+        timesteps = timestep * 1000.0
+        encoder_hidden_states = context
+        pooled_projections = y
+        guidance = guidance * 1000.0
+
+        B, C, H, W = x.shape
+        height = H * 8
+        width = W * 8
+        # Create SamplingParams
+        sampling_params = SamplingParams.from_user_sampling_params_args(
+            self.model_path,
+            server_args=self.generator.server_args,
+            prompt=" ",
+            guidance_scale=3.5,  # Flux typically uses embedded_cfg_scale=3.5
+            height=height,
+            width=width,
+            num_frames=1,
+            num_inference_steps=1,
+            save_output=False,
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        req.latents = hidden_states  # Set as [B, S, D] format directly
+        req.timesteps = timesteps  # ComfyUI's timesteps parameter
+        req.prompt_embeds = [pooled_projections, encoder_hidden_states]  # [CLIP, T5]
+        req.raw_latent_shape = torch.tensor(hidden_states.shape, dtype=torch.long)
+
+        # Set pooled_projections (required by Flux)
+        req.pooled_embeds = [pooled_projections]  # List format as per Req definition
+        req.do_classifier_free_guidance = False
+        req.generator = [
+            torch.Generator("cuda") for _ in range(req.num_outputs_per_prompt)
+        ]
+
+        # Send request to scheduler
+        output_batch = self.generator._send_to_scheduler_and_wait_for_response([req])
+        noise_pred = output_batch.noise_pred
+        return self._unpack_latents(noise_pred, H, W, C).to(x.device)

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

@@ -0,0 +1,63 @@
+"""
+ZImage 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]'"
+    )
+
+from .base import SGLDiffusionExecutor
+
+
+class ZImageExecutor(SGLDiffusionExecutor):
+    """Executor for ZImage models in ComfyUI."""
+
+    def __init__(self, generator, model_path, model, config):
+        super().__init__(generator, model_path, model, config)
+
+    def forward(self, x, timesteps, context, **kwargs):
+        """Forward pass for ZImage model."""
+        B, C, H, W = x.shape
+        height = H * 8
+        width = W * 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,  # For images
+            num_inference_steps=1,  # Single step for ComfyUI
+            save_output=False,
+        )
+
+        # Prepare request (converts SamplingParams to Req)
+        req = prepare_request(
+            server_args=self.generator.server_args,
+            sampling_params=sampling_params,
+        )
+        latents = x.unsqueeze(2)
+        context = context.squeeze(0)
+        # Set ComfyUI-specific inputs directly on the Req object
+        req.latents = latents  # ComfyUI's x parameter
+        req.timesteps = timesteps * 1000.0  # ComfyUI's timesteps parameter
+        req.prompt_embeds = [
+            context
+        ]  # ComfyUI's context parameter (must be List[Tensor])
+        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 noise_pred.permute(1, 0, 2, 3).to(x.device)

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

@@ -0,0 +1,659 @@
+"""
+ComfyUI nodes for SGLang Diffusion integration.
+Provides nodes for connecting to SGLang Diffusion server and generating images/videos.
+"""
+
+import os
+
+import folder_paths
+import torch
+
+from .core import SGLDiffusionGenerator, SGLDiffusionServerAPI
+from .utils import (
+    convert_b64_to_tensor_image,
+    convert_video_to_comfy_video,
+    get_image_path,
+    is_empty_image,
+)
+
+
+class SGLDOptions:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {},
+            "optional": {
+                "enable_torch_compile": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "num_gpus": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "tp_size": ("INT", {"default": -1, "min": -1, "step": 1}),
+                "sp_degree": ("INT", {"default": -1, "min": -1, "step": 1}),
+                "ulysses_degree": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                    },
+                ),
+                "ring_degree": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                    },
+                ),
+                "dp_size": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "dp_degree": ("INT", {"default": 1, "min": 1, "step": 1}),
+                "enable_cfg_parallel": (
+                    "BOOLEAN",
+                    {"default": False},
+                ),
+                "attention_backend": (
+                    "STRING",
+                    {"default": ""},
+                ),
+                "cache_strategy": (
+                    "STRING",
+                    {"default": "none"},
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_OPTIONS",)
+    RETURN_NAMES = ("sgld_options",)
+    FUNCTION = "create_options"
+    CATEGORY = "SGLDiffusion"
+
+    def create_options(
+        self,
+        enable_torch_compile: bool = False,
+        num_gpus: int = 1,
+        tp_size: int = -1,
+        sp_degree: int = -1,
+        ulysses_degree: int = -1,
+        ring_degree: int = -1,
+        dp_size: int = 1,
+        dp_degree: int = 1,
+        enable_cfg_parallel: bool = False,
+        attention_backend: str = "",
+        cache_strategy: str = "none",
+    ):
+        """
+        Build a dictionary of SGLang Diffusion runtime options.
+        """
+        # Convert -1 to None for optional parameters (matching ServerArgs defaults)
+        ulysses_degree = None if ulysses_degree == -1 else ulysses_degree
+        ring_degree = None if ring_degree == -1 else ring_degree
+        attention_backend = None if attention_backend == "" else attention_backend
+
+        options = {
+            "enable_torch_compile": enable_torch_compile,
+            "num_gpus": num_gpus,
+            "tp_size": tp_size,
+            "sp_degree": sp_degree,
+            "ulysses_degree": ulysses_degree,
+            "ring_degree": ring_degree,
+            "dp_size": dp_size,
+            "dp_degree": dp_degree,
+            "enable_cfg_parallel": enable_cfg_parallel,
+            "attention_backend": attention_backend,
+            "cache_strategy": cache_strategy,
+        }
+
+        # Strip None to keep payload clean
+        options = {k: v for k, v in options.items() if v is not None}
+        return (options,)
+
+
+class SGLDUNETLoader:
+    def __init__(self):
+        self.generator = SGLDiffusionGenerator()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "unet_name": (folder_paths.get_filename_list("diffusion_models"),),
+                "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
+            },
+            "optional": {
+                "sgld_options": ("SGLD_OPTIONS",),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "load_unet"
+
+    CATEGORY = "SGLDiffusion"
+
+    def load_unet(self, unet_name, weight_dtype, sgld_options: dict = None):
+        model_options = {}
+        if weight_dtype == "fp8_e4m3fn":
+            model_options["dtype"] = torch.float8_e4m3fn
+        elif weight_dtype == "fp8_e5m2":
+            model_options["dtype"] = torch.float8_e5m2
+
+        unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
+
+        model = self.generator.load_model(
+            unet_path, model_options=model_options, sgld_options=sgld_options
+        )
+        return (model,)
+
+
+class SGLDiffusionServerModel:
+    """Node to load and manage SGLang Diffusion server connection."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "base_url": (
+                    "STRING",
+                    {
+                        "default": "http://localhost:3000/v1",
+                        "multiline": False,
+                    },
+                ),
+                "api_key": (
+                    "STRING",
+                    {
+                        "default": "sk-proj-1234567890",
+                        "multiline": False,
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT", "STRING")
+    RETURN_NAMES = ("sgld_client", "model_info")
+    FUNCTION = "load_server"
+    CATEGORY = "SGLDiffusion"
+
+    def load_server(self, base_url: str, api_key: str):
+        """Initialize OpenAI client for SGLang Diffusion server."""
+        client = SGLDiffusionServerAPI(base_url=base_url, api_key=api_key)
+        try:
+            model_info = client.get_model_info()
+            # Format model_info as a readable string
+            info_lines = ["=== SGLDiffusion Model Info ==="]
+            for key, value in model_info.items():
+                info_lines.append(f"{key}: {value}")
+            model_info_str = "\n".join(info_lines)
+        except Exception as e:
+            model_info_str = f"Failed to get model info: {str(e)}"
+        return (client, model_info_str)
+
+
+class SGLDiffusionGenerateImage:
+    """Node to generate images using SGLang Diffusion."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "positive_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Text prompt for image generation",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Negative prompt to avoid certain elements",
+                    },
+                ),
+                "image": (
+                    "IMAGE",
+                    {
+                        "default": None,
+                        "tooltip": "input image to use for editing",
+                    },
+                ),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": -1,
+                        "max": 2**32 - 1,
+                    },
+                ),
+                "steps": (
+                    "INT",
+                    {
+                        "default": 6,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                    },
+                ),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 7.0,
+                        "min": 1.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                    },
+                ),
+                "width": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 64,
+                    },
+                ),
+                "height": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 64,
+                    },
+                ),
+                "enable_teacache": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate_image"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def generate_image(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        positive_prompt: str,
+        negative_prompt: str = "",
+        image: torch.Tensor = None,
+        seed: int = 1024,
+        steps: int = 6,
+        cfg: float = 7.0,
+        width: int = 1024,
+        height: int = 1024,
+        enable_teacache: bool = False,
+    ):
+        """Generate image using SGLang Diffusion API."""
+        if not positive_prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        size = f"{width}x{height}"
+
+        # Prepare request parameters
+        request_params = {
+            "prompt": positive_prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+
+        # Add optional parameters if provided
+        if negative_prompt:
+            request_params["negative_prompt"] = negative_prompt
+        if cfg is not None:
+            request_params["guidance_scale"] = cfg
+        if steps is not None:
+            request_params["num_inference_steps"] = steps
+        if seed is not None and seed >= 0:
+            request_params["seed"] = seed
+        if enable_teacache:
+            request_params["enable_teacache"] = True
+        if image is not None:
+            # If the image is empty, use the size of the image to generate the image
+            if is_empty_image(image):
+                width, height = image.shape[2], image.shape[1]
+                size = f"{width}x{height}"
+                request_params["size"] = size
+            else:
+                request_params["image_path"] = get_image_path(image)
+
+        # Call API
+        try:
+            response = sgld_client.generate_image(**request_params)
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate image: {str(e)}")
+
+        # Decode base64 image
+        if not response["data"] or not response["data"][0]["b64_json"]:
+            raise RuntimeError("No image data in response")
+        image_data = response["data"][0]["b64_json"]
+        image = convert_b64_to_tensor_image(image_data)
+
+        return (image,)
+
+
+class SGLDiffusionGenerateVideo:
+    """Node to generate videos using SGLang Diffusion."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "positive_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Text prompt for video generation",
+                    },
+                ),
+            },
+            "optional": {
+                "negative_prompt": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "Negative prompt to avoid certain elements",
+                    },
+                ),
+                "image": (
+                    "IMAGE",
+                    {
+                        "default": None,
+                        "tooltip": "input image to use for image-to-video",
+                    },
+                ),
+                "seed": (
+                    "INT",
+                    {
+                        "default": 1024,
+                        "min": -1,
+                        "max": 2**32 - 1,
+                    },
+                ),
+                "steps": (
+                    "INT",
+                    {
+                        "default": 6,
+                        "min": 1,
+                        "max": 100,
+                        "step": 1,
+                    },
+                ),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 7.0,
+                        "min": 1.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                    },
+                ),
+                "width": (
+                    "INT",
+                    {
+                        "default": 1280,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 1,
+                    },
+                ),
+                "height": (
+                    "INT",
+                    {
+                        "default": 720,
+                        "min": 256,
+                        "max": 4096,
+                        "step": 1,
+                    },
+                ),
+                "num_frames": (
+                    "INT",
+                    {
+                        "default": 120,
+                        "min": 1,
+                        "max": 1000,
+                        "step": 1,
+                    },
+                ),
+                "fps": (
+                    "INT",
+                    {
+                        "default": 24,
+                        "min": 1,
+                        "max": 60,
+                        "step": 1,
+                    },
+                ),
+                "seconds": (
+                    "INT",
+                    {
+                        "default": 5,
+                        "min": 1,
+                        "max": 60,
+                        "step": 1,
+                    },
+                ),
+                "enable_teacache": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO", "STRING")
+    RETURN_NAMES = ("video", "video_path")
+    FUNCTION = "generate_video"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def generate_video(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        positive_prompt: str,
+        negative_prompt: str = "",
+        image: torch.Tensor = None,
+        seed: int = 1024,
+        steps: int = 6,
+        cfg: float = 7.0,
+        width: int = 1280,
+        height: int = 720,
+        num_frames: int = 120,
+        fps: int = 24,
+        seconds: int = 5,
+        enable_teacache: bool = False,
+    ):
+        """Generate video using SGLang Diffusion API."""
+        if not positive_prompt:
+            raise ValueError("Prompt cannot be empty")
+
+        size = f"{width}x{height}"
+        output_dir = folder_paths.get_temp_directory()
+
+        # Prepare request parameters
+        request_params = {
+            "prompt": positive_prompt,
+            "size": size,
+            "seconds": seconds,
+            "fps": fps,
+            "output_path": output_dir,
+        }
+
+        # Add optional parameters if provided
+        if negative_prompt:
+            request_params["negative_prompt"] = negative_prompt
+        if cfg is not None:
+            request_params["guidance_scale"] = cfg
+        if steps is not None:
+            request_params["num_inference_steps"] = steps
+        if seed is not None and seed >= 0:
+            request_params["seed"] = seed
+        if enable_teacache:
+            request_params["enable_teacache"] = True
+        if num_frames is not None:
+            request_params["num_frames"] = num_frames
+        if image is not None:
+            # If the image is empty, use the size of the image to generate the video
+            if is_empty_image(image):
+                width, height = image.shape[2], image.shape[1]
+                size = f"{width}x{height}"
+                request_params["size"] = size
+            else:
+                request_params["input_reference"] = get_image_path(image)
+
+        # Call API
+        try:
+            response = sgld_client.generate_video(**request_params)
+            video_path = response.get("file_path", "")
+            video = convert_video_to_comfy_video(video_path, height, width)
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate video: {str(e)}")
+
+        return (video, video_path)
+
+
+class SGLDiffusionServerSetLora:
+    """Node to set LoRA adapter for SGLang Diffusion server."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+                "lora_name": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "The name of the LoRA adapter",
+                    },
+                ),
+            },
+            "optional": {
+                "lora_nickname": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "tooltip": "The nickname of the LoRA adapter",
+                    },
+                ),
+                "target": (
+                    [
+                        "all",
+                        "transformer",
+                        "transformer_2",
+                        "critic",
+                    ],
+                    {
+                        "default": "all",
+                        "tooltip": "Which transformer(s) to apply the LoRA to",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT",)
+    RETURN_NAMES = ("sgld_client",)
+    FUNCTION = "set_lora"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def set_lora(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        lora_name: str = "",
+        lora_nickname: str = "",
+        target: str = "all",
+    ):
+        """Set LoRA adapter using SGLang Diffusion API."""
+        if lora_nickname == "":
+            lora_nickname = os.path.splitext(lora_name)[0]
+
+        # Prepare request parameters
+        request_params = {
+            "lora_nickname": lora_nickname,
+            "lora_path": lora_name,
+            "target": target,
+        }
+
+        # Call API
+        try:
+            response = sgld_client.set_lora(**request_params)
+            return (sgld_client,)
+        except Exception as e:
+            raise RuntimeError(f"Failed to set LoRA adapter: {str(e)}")
+
+
+class SGLDiffusionServerUnsetLora:
+    """Node to unset LoRA adapter for SGLang Diffusion server."""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "sgld_client": ("SGLD_CLIENT",),
+            },
+            "optional": {
+                "target": (
+                    [
+                        "all",
+                        "transformer",
+                        "transformer_2",
+                        "critic",
+                    ],
+                    {
+                        "default": "all",
+                        "tooltip": "Which transformer(s) to unset the LoRA from",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("SGLD_CLIENT",)
+    RETURN_NAMES = ("sgld_client",)
+    FUNCTION = "unset_lora"
+    CATEGORY = "SGLDiffusion"
+    OUTPUT_NODE = False
+
+    def unset_lora(
+        self,
+        sgld_client: SGLDiffusionServerAPI,
+        target: str = "all",
+    ):
+        """Unset LoRA adapter using SGLang Diffusion API."""
+        try:
+            response = sgld_client.unset_lora(target=target)
+            return (sgld_client,)
+        except Exception as e:
+            raise RuntimeError(f"Failed to unset LoRA adapter: {str(e)}")
+
+
+# Register nodes
+NODE_CLASS_MAPPINGS = {
+    "SGLDiffusionServerModel": SGLDiffusionServerModel,
+    "SGLDiffusionGenerateImage": SGLDiffusionGenerateImage,
+    "SGLDiffusionGenerateVideo": SGLDiffusionGenerateVideo,
+    "SGLDiffusionServerSetLora": SGLDiffusionServerSetLora,
+    "SGLDiffusionServerUnsetLora": SGLDiffusionServerUnsetLora,
+    "SGLDUNETLoader": SGLDUNETLoader,
+    "SGLDOptions": SGLDOptions,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SGLDiffusionServerModel": "SGLDiffusion Server Model",
+    "SGLDiffusionGenerateImage": "SGLDiffusion Generate Image",
+    "SGLDiffusionGenerateVideo": "SGLDiffusion Generate Video",
+    "SGLDiffusionServerSetLora": "SGLDiffusion Server Set LoRA",
+    "SGLDiffusionServerUnsetLora": "SGLDiffusion Server Unset LoRA",
+    "SGLDUNETLoader": "SGLDiffusion UNET Loader",
+    "SGLDOptions": "SGLDiffusion Options",
+}

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

@@ -0,0 +1,176 @@
+import base64
+import io
+import os
+import shutil
+import time
+import uuid
+
+import folder_paths
+import numpy as np
+import torch
+from comfy_api.input import VideoInput
+from PIL import Image
+
+
+def _ensure_dir(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+
+
+def _to_numpy_image(image: torch.Tensor) -> np.ndarray:
+    """Convert ComfyUI image tensor to uint8 numpy array (H, W, C)."""
+    if image.dim() == 4:
+        image = image[0]
+    if image.dim() == 3 and image.shape[0] in (1, 3, 4):
+        image = image.permute(1, 2, 0)
+    elif image.dim() == 2:
+        image = image.unsqueeze(-1)
+    np_img = image.detach().cpu().numpy()
+    np_img = np.clip(np_img, 0.0, 1.0)
+    np_img = (np_img * 255).astype(np.uint8)
+    if np_img.shape[-1] == 1:
+        np_img = np.repeat(np_img, 3, axis=-1)
+    return np_img
+
+
+def _to_hwc_tensor(image: torch.Tensor) -> torch.Tensor:
+    """Convert ComfyUI image tensor to HWC format (normalized [0, 1])."""
+    img = image.clone()
+    if img.dim() == 4:
+        img = img[0]
+    if img.dim() == 3 and img.shape[0] in (1, 3, 4):
+        img = img.permute(1, 2, 0)
+    elif img.dim() == 2:
+        img = img.unsqueeze(-1)
+
+    img = torch.clamp(img, 0.0, 1.0)
+    if img.shape[-1] == 1:
+        img = img.repeat(1, 1, 3)
+
+    return img
+
+
+def is_empty_image(image: torch.Tensor, tolerance: float = 1e-6) -> bool:
+    """
+    Check if the input image is an empty/solid color image (like ComfyUI's empty image).
+    Args:
+        image: Input tensor image in ComfyUI format (BCHW, CHW, HWC, etc.)
+        tolerance: Tolerance for floating point comparison (default: 1e-6)
+
+    Returns:
+        True if the image is empty (all pixels have same color), False otherwise
+    """
+    if image is None:
+        return True
+
+    # Convert to HWC format
+    img_hwc = _to_hwc_tensor(image)
+
+    # Get the first pixel's RGB values
+    first_pixel = img_hwc[0, 0, :]
+
+    h, w, c = img_hwc.shape
+    pixels = img_hwc.reshape(-1, c)
+
+    diff = torch.abs(pixels - first_pixel)
+    max_diff = torch.max(diff)
+
+    return max_diff.item() <= tolerance
+
+
+def get_image_path(image: torch.Tensor) -> str:
+    """
+    Save tensor image to ComfyUI temp directory as PNG and return the path.
+    """
+    temp_dir = folder_paths.get_temp_directory()
+
+    # Build file name
+    ts = time.strftime("%Y%m%d-%H%M%S")
+    unique = uuid.uuid4().hex[:8]
+    file_name = f"sgl_output_{ts}_{unique}.png"
+    file_path = os.path.join(temp_dir, file_name)
+
+    # Save image
+    np_img = _to_numpy_image(image)
+    img = Image.fromarray(np_img)
+    img.save(file_path, format="PNG")
+
+    return file_path
+
+
+def convert_b64_to_tensor_image(b64_image: str) -> torch.Tensor:
+    """
+    Convert base64 encoded image to ComfyUI IMAGE format (torch.Tensor).
+
+    Args:
+        b64_image: Base64 encoded image string
+
+    Returns:
+        torch.Tensor with shape [batch_size, height, width, channels] (BHWC format),
+        values normalized to [0, 1] range, RGB format (3 channels)
+    """
+    # Decode base64
+    image_bytes = base64.b64decode(b64_image)
+
+    # Open image and convert to RGB
+    pil_image = Image.open(io.BytesIO(image_bytes))
+    if pil_image.mode != "RGB":
+        pil_image = pil_image.convert("RGB")
+
+    # Convert to numpy array and normalize to [0, 1]
+    image_array = np.array(pil_image).astype(np.float32) / 255.0
+
+    # Add batch dimension: [height, width, channels] -> [1, height, width, channels]
+    image_array = image_array[np.newaxis, ...]
+
+    # Convert to torch.Tensor
+    tensor_image = torch.from_numpy(image_array)
+
+    return tensor_image
+
+
+class SGLDVideoInput(VideoInput):
+    def __init__(self, video_path: str, height: int, width: int):
+        super().__init__()
+
+        self.video_path = video_path
+        self.height = height
+        self.width = width
+
+    def get_dimensions(self) -> tuple[int, int]:
+        """
+        Returns the dimensions of the video input.
+
+        Returns:
+            Tuple of (width, height)
+        """
+        return self.width, self.height
+
+    def get_components(self):
+        """
+        Returns the components of the video input.
+        This is required by the VideoInput abstract base class.
+        """
+        return [self.video_path]
+
+    def save_to(self, path: str, format=None, codec=None, metadata=None):
+        """
+        Abstract method to save the video input to a file.
+        """
+        save_path = path
+        # Copy video file from video_path to save_path
+        if os.path.exists(self.video_path):
+            # Ensure destination directory exists
+            save_dir = os.path.dirname(save_path)
+            if save_dir:
+                os.makedirs(save_dir, exist_ok=True)
+            shutil.copy2(self.video_path, save_path)
+
+
+def convert_video_to_comfy_video(
+    video_path: str, height: int, width: int
+) -> VideoInput:
+    """
+    Convert video to ComfyUI VIDEO format (VideoInput).
+    """
+    video_input = SGLDVideoInput(video_path, height, width)
+    return video_input

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

@@ -0,0 +1,97 @@
+{
+    "1": {
+      "inputs": {
+        "base_url": "http://localhost:3000/v1",
+        "api_key": "sk-proj-1234567890"
+      },
+      "class_type": "SGLDiffusionServerModel",
+      "_meta": {
+        "title": "SGLDiffusion Server Model"
+      }
+    },
+    "3": {
+      "inputs": {
+        "prompt": "The girl turn the body and spin around in place.",
+        "main": "none",
+        "lighting": "none",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        }
+      },
+      "class_type": "easy prompt",
+      "_meta": {
+        "title": "Prompt"
+      }
+    },
+    "4": {
+      "inputs": {
+        "text": "",
+        "anything": [
+          "1",
+          1
+        ]
+      },
+      "class_type": "easy showAnything",
+      "_meta": {
+        "title": "Show Any"
+      }
+    },
+    "15": {
+      "inputs": {
+        "positive_prompt": [
+          "3",
+          0
+        ],
+        "negative_prompt": "",
+        "seed": 2435791308,
+        "steps": 50,
+        "cfg": 4,
+        "width": 704,
+        "height": 1280,
+        "num_frames": 16,
+        "fps": 16,
+        "seconds": 1,
+        "enable_teacache": false,
+        "sgld_client": [
+          "1",
+          0
+        ],
+        "image": [
+          "17",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionGenerateVideo",
+      "_meta": {
+        "title": "SGLDiffusion Generate Video"
+      }
+    },
+    "16": {
+      "inputs": {
+        "filename_prefix": "video/ComfyUI",
+        "format": "auto",
+        "codec": "auto",
+        "video-preview": "",
+        "video": [
+          "15",
+          0
+        ]
+      },
+      "class_type": "SaveVideo",
+      "_meta": {
+        "title": "save video"
+      }
+    },
+    "17": {
+      "inputs": {
+        "image": "tmpe_w0bd_0.jpg"
+      },
+      "class_type": "LoadImage",
+      "_meta": {
+        "title": "load image"
+      }
+    }
+  }

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

@@ -0,0 +1,109 @@
+{
+    "1": {
+      "inputs": {
+        "base_url": "http://localhost:3000/v1",
+        "api_key": "sk-proj-1234567890"
+      },
+      "class_type": "SGLDiffusionServerModel",
+      "_meta": {
+        "title": "SGLDiffusion Server Model"
+      }
+    },
+    "3": {
+      "inputs": {
+        "prompt": "a bicycle, illustration in the style of SMPL, thick black lines on a white background",
+        "main": "none",
+        "lighting": "none",
+        "speak_and_recognation": {
+          "__value__": [
+            false,
+            true
+          ]
+        }
+      },
+      "class_type": "easy prompt",
+      "_meta": {
+        "title": "Prompt"
+      }
+    },
+    "4": {
+      "inputs": {
+        "text": "",
+        "anything": [
+          "1",
+          1
+        ]
+      },
+      "class_type": "easy showAnything",
+      "_meta": {
+        "title": "Show Any"
+      }
+    },
+    "5": {
+      "inputs": {
+        "filename_prefix": "ComfyUI",
+        "images": [
+          "6",
+          0
+        ]
+      },
+      "class_type": "SaveImage",
+      "_meta": {
+        "title": "save image"
+      }
+    },
+    "6": {
+      "inputs": {
+        "positive_prompt": [
+          "3",
+          0
+        ],
+        "negative_prompt": "",
+        "seed": 4215918563,
+        "steps": 50,
+        "cfg": 4,
+        "width": 512,
+        "height": 512,
+        "enable_teacache": false,
+        "sgld_client": [
+          "11",
+          0
+        ],
+        "image": [
+          "14",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionGenerateImage",
+      "_meta": {
+        "title": "SGLDiffusion Generate Image"
+      }
+    },
+    "11": {
+      "inputs": {
+        "lora_name": "dvyio/flux-lora-simple-illustration",
+        "lora_nickname": "",
+        "target": "all",
+        "sgld_client": [
+          "1",
+          0
+        ]
+      },
+      "class_type": "SGLDiffusionSetLora",
+      "_meta": {
+        "title": "SGLDiffusion Set LoRA"
+      }
+    },
+    "14": {
+      "inputs": {
+        "width": 512,
+        "height": 512,
+        "batch_size": 1,
+        "color": 0
+      },
+      "class_type": "EmptyImage",
+      "_meta": {
+        "title": "empty image"
+      }
+    }
+  }

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

@@ -2,6 +2,7 @@
 
 # SPDX-License-Identifier: Apache-2.0
 import json
+import os
 from collections.abc import Callable
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum, auto
@@ -518,15 +519,55 @@ class PipelineConfig:
         if model_path is None:
             raise ValueError("model_path is required in kwargs")
 
+        # Check if model_path is a safetensors file and pipeline_class_name is specified
+        pipeline_class_name = kwargs.get(
+            prefix_with_dot + "pipeline_class_name"
+        ) or kwargs.get("pipeline_class_name")
+        is_safetensors_file = os.path.isfile(model_path) and model_path.endswith(
+            ".safetensors"
+        )
+
         # 1. Get the pipeline config class from the registry
         from sglang.multimodal_gen.configs.pipeline_configs.flux import (
             Flux2PipelineConfig,
         )
+        from sglang.multimodal_gen.registry import get_pipeline_config_classes
 
-        model_info = get_model_info(model_path)
+        # If model_path is a safetensors file and pipeline_class_name is specified,
+        # try to get PipelineConfig from the registry first
+        if is_safetensors_file and pipeline_class_name:
+            config_classes = get_pipeline_config_classes(pipeline_class_name)
+            if config_classes is not None:
+                pipeline_config_cls, _ = config_classes
+                logger.info(
+                    f"Detected safetensors file with {pipeline_class_name}, "
+                    f"using {pipeline_config_cls.__name__} directly without model_index.json"
+                )
+            else:
+                model_info = get_model_info(model_path)
+                if model_info is None:
+                    from sglang.multimodal_gen.registry import (
+                        _PIPELINE_CONFIG_REGISTRY,
+                        _discover_and_register_pipelines,
+                    )
 
-        # 1.5. Adjust pipeline config for fine-tuned VAE if needed
-        pipeline_config_cls = model_info.pipeline_config_cls
+                    _discover_and_register_pipelines()
+                    available_pipelines = list(_PIPELINE_CONFIG_REGISTRY.keys())
+                    raise ValueError(
+                        f"Could not get model info for '{model_path}'. "
+                        f"If using a safetensors file, please specify a valid pipeline_class_name. "
+                        f"Available pipelines with config classes: {available_pipelines}"
+                    )
+                pipeline_config_cls = model_info.pipeline_config_cls
+        else:
+            model_info = get_model_info(model_path)
+            if model_info is None:
+                raise ValueError(
+                    f"Could not get model info for '{model_path}'. "
+                    f"If using a safetensors file, please specify pipeline_class_name"
+                )
+            # 1.5. Adjust pipeline config for fine-tuned VAE if needed
+            pipeline_config_cls = model_info.pipeline_config_cls
         vae_path = kwargs.get(prefix_with_dot + "vae_path") or kwargs.get("vae_path")
 
         # Check if this is a Flux2 model with fal/FLUX.2-Tiny-AutoEncoder

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

@@ -405,7 +405,8 @@ class SamplingParams:
                 )
                 self.num_frames = new_num_frames
 
-        self._set_output_file_name()
+        if not server_args.comfyui_mode:
+            self._set_output_file_name()
 
     @classmethod
     def from_pretrained(cls, model_path: str, **kwargs) -> "SamplingParams":
@@ -417,7 +418,45 @@ class SamplingParams:
 
     @staticmethod
     def from_user_sampling_params_args(model_path: str, server_args, *args, **kwargs):
-        sampling_params = SamplingParams.from_pretrained(model_path)
+        try:
+            sampling_params = SamplingParams.from_pretrained(model_path)
+        except (AttributeError, ValueError) as e:
+            # Handle safetensors files or other cases where model_index.json is not available
+            # Use appropriate SamplingParams based on pipeline_class_name from registry
+            if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
+                # Determine which sampling params to use based on pipeline_class_name
+                pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
+
+                # Try to get SamplingParams from registry
+                from sglang.multimodal_gen.registry import get_pipeline_config_classes
+
+                config_classes = (
+                    get_pipeline_config_classes(pipeline_class_name)
+                    if pipeline_class_name
+                    else None
+                )
+
+                if config_classes is not None:
+                    _, sampling_params_cls = config_classes
+                    try:
+                        sampling_params = sampling_params_cls()
+                        logger.info(
+                            f"Using {sampling_params_cls.__name__} for {pipeline_class_name} safetensors file (no model_index.json): %s",
+                            model_path,
+                        )
+                    except Exception as import_error:
+                        logger.warning(
+                            f"Failed to instantiate {sampling_params_cls.__name__}: {import_error}. "
+                            "Using default SamplingParams"
+                        )
+                        sampling_params = SamplingParams()
+                else:
+                    raise ValueError(
+                        f"Could not get pipeline config classes for {pipeline_class_name}"
+                    )
+            else:
+                # Re-raise if it's not a safetensors file issue
+                raise
 
         user_sampling_params = SamplingParams(*args, **kwargs)
         # TODO: refactor

python/sglang/multimodal_gen/registry.py

@@ -94,6 +94,10 @@ logger = init_logger(__name__)
 
 _PIPELINE_REGISTRY: Dict[str, Type[ComposedPipelineBase]] = {}
 
+# Registry for pipeline configuration classes (for safetensors files without model_index.json)
+# Maps pipeline_class_name -> (PipelineConfig class, SamplingParams class)
+_PIPELINE_CONFIG_REGISTRY: Dict[str, Tuple[Type[PipelineConfig], Type[Any]]] = {}
+
 
 def _discover_and_register_pipelines():
     """
@@ -126,11 +130,43 @@ def _discover_and_register_pipelines():
                                 f"Duplicate pipeline name '{cls.pipeline_name}' found. Overwriting."
                             )
                         _PIPELINE_REGISTRY[cls.pipeline_name] = cls
+
+                        # Auto-register config classes if Pipeline class has them defined
+                        # because comfyui get model from a single weight file, so we need to register the config classes here
+                        if hasattr(cls, "pipeline_config_cls") and hasattr(
+                            cls, "sampling_params_cls"
+                        ):
+                            _PIPELINE_CONFIG_REGISTRY[cls.pipeline_name] = (
+                                cls.pipeline_config_cls,
+                                cls.sampling_params_cls,
+                            )
+                            logger.debug(
+                                f"Auto-registered config classes for pipeline '{cls.pipeline_name}': "
+                                f"PipelineConfig={cls.pipeline_config_cls.__name__}, "
+                                f"SamplingParams={cls.sampling_params_cls.__name__}"
+                            )
     logger.debug(
         f"Registering pipelines complete, {len(_PIPELINE_REGISTRY)} pipelines registered"
     )
 
 
+def get_pipeline_config_classes(
+    pipeline_class_name: str,
+) -> Tuple[Type[PipelineConfig], Type[Any]] | None:
+    """
+    Get the configuration classes for a pipeline.
+
+    Args:
+        pipeline_class_name: The name of the pipeline class
+
+    Returns:
+        A tuple of (PipelineConfig class, SamplingParams class) if found, None otherwise
+    """
+    # Ensure pipelines are discovered first
+    _discover_and_register_pipelines()
+    return _PIPELINE_CONFIG_REGISTRY.get(pipeline_class_name)
+
+
 # --- Part 2: Config Registration ---
 @dataclasses.dataclass
 class ConfigInfo:

python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py

@@ -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(

python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py

@@ -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
 
 

python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py

@@ -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()),
     }
 
 

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

@@ -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:

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_comfyui_passthrough.py

@@ -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

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

@@ -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

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

@@ -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

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

@@ -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)

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

@@ -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

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

@@ -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",

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

@@ -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

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

@@ -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,

python/sglang/multimodal_gen/runtime/server_args.py

@@ -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