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.
71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
"""
|
|
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,
|
|
suppress_logs=self.should_suppress_logs(timesteps),
|
|
)
|
|
|
|
# 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
|
|
|
|
# SGLang returns 5D tensor in (F, C, B, H, W) format for video/image.
|
|
# For image generation F=1, so squeeze dim 0 to get (C, B, H, W),
|
|
# then permute(1,0,2,3) converts to (B, C, H, W) as ComfyUI expects.
|
|
if noise_pred.dim() == 5:
|
|
noise_pred = noise_pred.squeeze(0) # (F=1, C, B, H, W) -> (C, B, H, W)
|
|
|
|
return noise_pred.permute(1, 0, 2, 3).to(x.device)
|