chore: import upstream snapshot with attribution

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

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

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

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

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

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

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

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

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

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

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

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