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.
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
"""
|
|
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 = []
|
|
|
|
@staticmethod
|
|
def should_suppress_logs(timestep):
|
|
"""Determine if logs should be suppressed based on timestep value."""
|
|
if torch.is_tensor(timestep):
|
|
return bool((timestep < 1.0).item())
|
|
return bool(timestep < 1.0)
|
|
|
|
def set_lora(self, lora_nickname=None, lora_path=None, strength=None, target=None):
|
|
"""Set LoRA adapter using SGLang Diffusion API."""
|
|
if len(lora_nickname) > 0:
|
|
self.generator.set_lora(
|
|
lora_nickname=lora_nickname,
|
|
lora_path=lora_path,
|
|
strength=strength,
|
|
target=target,
|
|
)
|
|
|
|
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
|