endman100
453c5b277f
Introduce a new ComfyUI plugin to integrate SGLang Diffusion. Adds core integration pieces (generator, model patcher, server API client), executor implementations for multiple model types (flux, qwen_image, zimage, native, anima), node mappings, utilities, example workflows, and tests. Includes README with installation/usage instructions and example workflow JSON files. This commit implements both server (HTTP) and integrated modes, model loading/patching logic, and basic LoRA/video/image API support.
123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
"""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(f"✓ Successfully retrieved noise_pred from OutputBatch!")
|
|
print(f" noise_pred shape: {noise_pred.shape}")
|
|
print(f" noise_pred dtype: {noise_pred.dtype}")
|
|
print(f" noise_pred device: {noise_pred.device}")
|
|
|
|
latents = output_batch.output if output_batch.output is not None else req.latents
|
|
assert latents is not None, "latents should not be None"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-v"]))
|