diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/README.md b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/README.md new file mode 100644 index 000000000..ec786a53d --- /dev/null +++ b/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. diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/__init__.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/__init__.py new file mode 100644 index 000000000..8282ee737 --- /dev/null +++ b/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"] diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/__init__.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/__init__.py new file mode 100644 index 000000000..a7d01106b --- /dev/null +++ b/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", +] diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/generator.py new file mode 100644 index 000000000..a20895d48 --- /dev/null +++ b/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 + ) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/model_patcher.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/model_patcher.py new file mode 100644 index 000000000..c88efd29f --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/server_api.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/core/server_api.py new file mode 100644 index 000000000..2f86efe79 --- /dev/null +++ b/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" + ) + ) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/__init__.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/__init__.py new file mode 100644 index 000000000..cd9e4a126 --- /dev/null +++ b/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", +] diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/base.py new file mode 100644 index 000000000..c61dcca4f --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/flux.py new file mode 100644 index 000000000..d7119da53 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/zimage.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/executors/zimage.py new file mode 100644 index 000000000..e85bacf74 --- /dev/null +++ b/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) diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/nodes.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/nodes.py new file mode 100644 index 000000000..4765e4e56 --- /dev/null +++ b/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", +} diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/utils.py b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/utils.py new file mode 100644 index 000000000..093b901e4 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_image2video.json b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_image2video.json new file mode 100644 index 000000000..e3d5456ff --- /dev/null +++ b/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" + } + } + } diff --git a/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_text2img.json b/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion/workflows/sgld_text2img.json new file mode 100644 index 000000000..65e74c7b3 --- /dev/null +++ b/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" + } + } + } diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index ab09b8112..361551a30 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/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 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 3c62b6952..ff62bcc39 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/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 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 113ad0e5f..b4cead4d5 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py index 8b4518b99..a4b0bd660 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/image_api.py +++ b/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( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index b6331cf6e..e56861a33 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py index a4e9a6e27..9fb36d6e1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/video_api.py +++ b/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()), } diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index 2f1d0996b..815286ae5 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_comfyui_passthrough.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_comfyui_passthrough.py new file mode 100644 index 000000000..e87f558b8 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py new file mode 100644 index 000000000..27f617752 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py new file mode 100644 index 000000000..e35ec90de --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py index f5df51d80..f2e42ec6c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index 28a150982..a8a6dd0fe 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index f395f5525..7cfb17854 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/comfyui_latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/comfyui_latent_preparation.py new file mode 100644 index 000000000..060aa2259 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index e6a9412a2..d4041aea4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 1683fd0a2..e055f5f3b 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/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