diff --git a/python/sglang/multimodal_gen/configs/models/dits/__init__.py b/python/sglang/multimodal_gen/configs/models/dits/__init__.py index 2c6e5b0c6..ba114c18f 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/__init__.py +++ b/python/sglang/multimodal_gen/configs/models/dits/__init__.py @@ -1,5 +1,6 @@ # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo +from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.dits.hunyuan3d import Hunyuan3DDiTConfig from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig @@ -7,6 +8,7 @@ from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig __all__ = [ + "HeliosConfig", "HunyuanVideoConfig", "WanVideoConfig", "Hunyuan3DDiTConfig", diff --git a/python/sglang/multimodal_gen/configs/models/dits/helios.py b/python/sglang/multimodal_gen/configs/models/dits/helios.py new file mode 100644 index 000000000..15f73fb02 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/helios.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig + + +def is_blocks(n: str, m) -> bool: + return "blocks" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class HeliosArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks]) + + param_names_mapping: dict = field( + default_factory=lambda: { + # Patch embeddings + r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1", + # Condition embedder: text + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + # Condition embedder: time + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + # Blocks: self-attention (keep attn1. prefix, drop .0. from to_out) + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": r"blocks.\1.attn1.to_out.\2", + # Blocks: cross-attention output (drop .0. from to_out) + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": r"blocks.\1.attn2.to_out.\2", + # Blocks: feed-forward + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + # Blocks: cross-attn residual norm + r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.self_attn_residual_norm.\2", + } + ) + + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + lora_param_names_mapping: dict = field(default_factory=lambda: {}) + + patch_size: tuple[int, int, int] = (1, 2, 2) + text_len: int = 226 + num_attention_heads: int = 40 + attention_head_dim: int = 128 + in_channels: int = 16 + out_channels: int = 16 + text_dim: int = 4096 + freq_dim: int = 256 + ffn_dim: int = 13824 + num_layers: int = 40 + cross_attn_norm: bool = True + qk_norm: str = "rms_norm_across_heads" + eps: float = 1e-6 + added_kv_proj_dim: int | None = None + rope_max_seq_len: int = 1024 + pos_embed_seq_len: int | None = None + exclude_lora_layers: list[str] = field(default_factory=lambda: ["embedder"]) + + # Helios-specific + rope_dim: tuple[int, int, int] = (44, 42, 42) + rope_theta: float = 10000.0 + guidance_cross_attn: bool = True + zero_history_timestep: bool = True + has_multi_term_memory_patch: bool = True + is_amplify_history: bool = False + history_scale_mode: str = "per_head" + + def __post_init__(self): + super().__post_init__() + self.out_channels = self.out_channels or self.in_channels + self.hidden_size = self.num_attention_heads * self.attention_head_dim + self.num_channels_latents = self.out_channels + + +@dataclass +class HeliosConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=HeliosArchConfig) + + prefix: str = "Helios" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py index c7005a33f..f54b0ae9a 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py @@ -15,6 +15,11 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import ( from sglang.multimodal_gen.configs.pipeline_configs.flux_finetuned import ( Flux2FinetunedPipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.helios import ( + HeliosDistilledConfig, + HeliosMidConfig, + HeliosT2VConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import ( FastHunyuanConfig, HunyuanConfig, @@ -35,6 +40,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.zimage import ZImagePipeline __all__ = [ "DiffusersGenericPipelineConfig", + "HeliosDistilledConfig", + "HeliosMidConfig", + "HeliosT2VConfig", "HunyuanConfig", "FastHunyuanConfig", "Hunyuan3D2PipelineConfig", diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py b/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py new file mode 100644 index 000000000..d14e92769 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/helios.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +from collections.abc import Callable +from dataclasses import dataclass, field + +import torch + +from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAEConfig +from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig +from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput, T5Config +from sglang.multimodal_gen.configs.models.encoders.t5 import T5ArchConfig +from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig +from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ModelTaskType, + PipelineConfig, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# Helios UMT5 max sequence length (used for both tokenizer and post-processing padding) +HELIOS_MAX_SEQUENCE_LENGTH = 226 + + +def umt5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: + """Post-process UMT5 text encoder outputs, padding to HELIOS_MAX_SEQUENCE_LENGTH tokens.""" + max_seq_len = HELIOS_MAX_SEQUENCE_LENGTH + mask: torch.Tensor = outputs.attention_mask + hidden_state: torch.Tensor = outputs.last_hidden_state + seq_lens = mask.gt(0).sum(dim=1).long() + assert torch.isnan(hidden_state).sum() == 0 + prompt_embeds = [u[:v] for u, v in zip(hidden_state, seq_lens, strict=True)] + prompt_embeds_tensor: torch.Tensor = torch.stack( + [ + torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0, + ) + return prompt_embeds_tensor + + +@dataclass +class HeliosT2VConfig(PipelineConfig): + """Configuration for the Helios T2V pipeline.""" + + task_type: ModelTaskType = ModelTaskType.T2V + + # DiT + dit_config: DiTConfig = field(default_factory=HeliosConfig) + + # VAE (same as Wan) + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + vae_tiling: bool = False + vae_sp: bool = False + + # Denoising stage + flow_shift: float | None = 1.0 + + # Text encoding stage (UMT5 is T5-compatible) + text_encoder_configs: tuple[EncoderConfig, ...] = field( + default_factory=lambda: ( + T5Config(arch_config=T5ArchConfig(text_len=HELIOS_MAX_SEQUENCE_LENGTH)), + ) + ) + postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = ( + field(default_factory=lambda: (umt5_postprocess_text,)) + ) + + # Precision for each component + precision: str = "bf16" + vae_precision: str = "fp32" + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("fp32",)) + + # Helios-specific chunked denoising params + num_latent_frames_per_chunk: int = 9 + history_sizes: list[int] = field(default_factory=lambda: [16, 2, 1]) + is_cfg_zero_star: bool = False + zero_steps: int = 1 + keep_first_frame: bool = True + + # Stage 2 (Pyramid SR) & Stage 3 (DMD) params + is_enable_stage2: bool = False + pyramid_num_stages: int = 3 + pyramid_num_inference_steps_list: list[int] = field( + default_factory=lambda: [10, 10, 10] + ) + is_distilled: bool = False + is_amplify_first_chunk: bool = False + scheduler_type: str = "unipc" + gamma: float = 1 / 3 + + def __post_init__(self): + self.vae_config.load_encoder = False + self.vae_config.load_decoder = True + + +@dataclass +class HeliosMidConfig(HeliosT2VConfig): + """Configuration for Helios-Mid (Stage 1 + Stage 2 pyramid SR).""" + + is_enable_stage2: bool = True + is_cfg_zero_star: bool = True + pyramid_num_inference_steps_list: list[int] = field( + default_factory=lambda: [20, 20, 20] + ) + + +@dataclass +class HeliosDistilledConfig(HeliosT2VConfig): + """Configuration for Helios-Distilled (Stage 1 + Stage 2 + Stage 3 DMD).""" + + is_enable_stage2: bool = True + is_distilled: bool = True + is_amplify_first_chunk: bool = True + scheduler_type: str = "dmd" + pyramid_num_inference_steps_list: list[int] = field( + default_factory=lambda: [10, 10, 10] + ) diff --git a/python/sglang/multimodal_gen/configs/sample/helios.py b/python/sglang/multimodal_gen/configs/sample/helios.py new file mode 100644 index 000000000..28ccd2093 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/helios.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams + + +@dataclass +class HeliosT2VSamplingParams(SamplingParams): + # Video parameters + height: int = 384 + width: int = 640 + num_frames: int = 99 + fps: int = 24 + + # Denoising stage + guidance_scale: float = 5.0 + negative_prompt: str = ( + "Bright tones, overexposed, static, blurred details, subtitles, style, " + "works, paintings, images, static, overall gray, worst quality, low quality, " + "JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, " + "poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, " + "still picture, messy background, three legs, many people in the background, " + "walking backwards" + ) + num_inference_steps: int = 50 + + # Helios T2V supported resolutions + supported_resolutions: list[tuple[int, int]] | None = field( + default_factory=lambda: [ + (640, 384), # ~5:3 + (384, 640), # ~3:5 + (832, 480), # ~16:9-ish + (480, 832), # ~9:16-ish + ] + ) + + +@dataclass +class HeliosMidSamplingParams(HeliosT2VSamplingParams): + """Sampling params for Helios-Mid (Stage 2 pyramid SR).""" + + num_inference_steps: int = 20 + + +@dataclass +class HeliosDistilledSamplingParams(HeliosT2VSamplingParams): + """Sampling params for Helios-Distilled (DMD, no CFG needed).""" + + guidance_scale: float = 1.0 + num_inference_steps: int = 10 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index 90a8f217a..690327218 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -30,6 +30,9 @@ if TYPE_CHECKING: from sglang.multimodal_gen.configs.pipeline_configs import ( FastHunyuanConfig, FluxPipelineConfig, + HeliosDistilledConfig, + HeliosMidConfig, + HeliosT2VConfig, HunyuanConfig, WanI2V480PConfig, WanI2V720PConfig, @@ -74,6 +77,11 @@ from sglang.multimodal_gen.configs.sample.flux import ( FluxSamplingParams, ) from sglang.multimodal_gen.configs.sample.glmimage import GlmImageSamplingParams +from sglang.multimodal_gen.configs.sample.helios import ( + HeliosDistilledSamplingParams, + HeliosMidSamplingParams, + HeliosT2VSamplingParams, +) from sglang.multimodal_gen.configs.sample.hunyuan import ( FastHunyuanSamplingParam, HunyuanSamplingParams, @@ -726,6 +734,34 @@ def _register_configs(): model_detectors=[lambda hf_id: "hunyuan3d" in hf_id.lower()], ) + # Helios + register_configs( + sampling_param_cls=HeliosT2VSamplingParams, + pipeline_config_cls=HeliosT2VConfig, + hf_model_paths=[ + "BestWishYsh/Helios-Base", + ], + model_detectors=[ + lambda hf_id: "helios" in hf_id.lower() + and "mid" not in hf_id.lower() + and "distill" not in hf_id.lower() + ], + ) + register_configs( + sampling_param_cls=HeliosMidSamplingParams, + pipeline_config_cls=HeliosMidConfig, + hf_model_paths=[ + "BestWishYsh/Helios-Mid", + ], + ) + register_configs( + sampling_param_cls=HeliosDistilledSamplingParams, + pipeline_config_cls=HeliosDistilledConfig, + hf_model_paths=[ + "BestWishYsh/Helios-Distilled", + ], + ) + _register_configs() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/helios.py b/python/sglang/multimodal_gen/runtime/models/dits/helios.py new file mode 100644 index 000000000..f4eecd42e --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/dits/helios.py @@ -0,0 +1,825 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from Helios diffusers transformer: +# https://github.com/BestWishYsh/Helios +""" +Helios Transformer 3D model for video generation. + +Implements the HeliosTransformer3DModel with multi-term memory patches, +3D rotary position embeddings, and per-block scale-shift modulation. +""" + +import math +from functools import lru_cache +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig +from sglang.multimodal_gen.runtime.distributed import ( + divide, + get_tp_world_size, +) +from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.layernorm import ( + FP32LayerNorm, + RMSNorm, + tensor_parallel_rms_norm, +) +from sglang.multimodal_gen.runtime.layers.linear import ( + ColumnParallelLinear, + RowParallelLinear, +) +from sglang.multimodal_gen.runtime.layers.mlp import MLP +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, +) +from sglang.multimodal_gen.runtime.layers.visual_embedding import ( + ModulateProjection, + PatchEmbed, + TimestepEmbedder, +) +from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT +from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Utility functions +# --------------------------------------------------------------------------- + + +def pad_for_3d_conv(x, kernel_size): + """Pad input to make it divisible by kernel_size using replicate mode.""" + b, c, t, h, w = x.shape + pt, ph, pw = kernel_size + pad_t = (pt - (t % pt)) % pt + pad_h = (ph - (h % ph)) % ph + pad_w = (pw - (w % pw)) % pw + return F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t), mode="replicate") + + +def center_down_sample_3d(x, kernel_size): + """Average pooling for 3D downsampling.""" + return F.avg_pool3d(x, kernel_size, stride=kernel_size) + + +def apply_rotary_emb_transposed(hidden_states, freqs_cis): + """Apply rotary positional embeddings with transposed cos/sin format.""" + x_1, x_2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs_cis.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(hidden_states) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(hidden_states) + + +# --------------------------------------------------------------------------- +# Output norm +# --------------------------------------------------------------------------- + + +class HeliosOutputNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + self.norm = FP32LayerNorm(dim, eps, elementwise_affine=False) + + def forward(self, hidden_states, temb, original_context_length): + temb = temb[:, -original_context_length:, :] + shift, scale = ( + self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2) + ).chunk(2, dim=2) + shift = shift.squeeze(2).to(hidden_states.device) + scale = scale.squeeze(2).to(hidden_states.device) + hidden_states = hidden_states[:, -original_context_length:, :] + hidden_states = ( + self.norm(hidden_states.float()) * (1 + scale) + shift + ).type_as(hidden_states) + return hidden_states + + +# --------------------------------------------------------------------------- +# Rotary Positional Embedding (3D) +# --------------------------------------------------------------------------- + + +class HeliosRotaryPosEmbed(nn.Module): + """3D rotary position embeddings for (time, height, width).""" + + def __init__(self, rope_dim, theta): + super().__init__() + self.DT, self.DY, self.DX = rope_dim + self.theta = theta + # Store as plain attributes (not buffers) to avoid meta-device issues + # during FSDP loading. They'll be re-created on the correct device in forward. + self._freqs_base_t = None + self._freqs_base_y = None + self._freqs_base_x = None + + def _get_freqs_base(self, dim): + return 1.0 / ( + self.theta + ** (torch.arange(0, dim, 2, dtype=torch.float32)[: (dim // 2)] / dim) + ) + + def _ensure_freqs_base(self, device): + """Lazily create frequency bases on the correct device.""" + if self._freqs_base_t is None or self._freqs_base_t.device != device: + self._freqs_base_t = self._get_freqs_base(self.DT).to(device) + self._freqs_base_y = self._get_freqs_base(self.DY).to(device) + self._freqs_base_x = self._get_freqs_base(self.DX).to(device) + + @torch.no_grad() + def get_frequency_batched(self, freqs_base, pos): + freqs = torch.einsum("d,bthw->dbthw", freqs_base, pos) + freqs = freqs.repeat_interleave(2, dim=0) + return freqs.cos(), freqs.sin() + + @torch.no_grad() + @lru_cache(maxsize=32) + def _get_spatial_meshgrid(self, height, width, device_str): + device = torch.device(device_str) + grid_y_coords = torch.arange(height, device=device, dtype=torch.float32) + grid_x_coords = torch.arange(width, device=device, dtype=torch.float32) + grid_y, grid_x = torch.meshgrid(grid_y_coords, grid_x_coords, indexing="ij") + return grid_y, grid_x + + @torch.no_grad() + def forward(self, frame_indices, height, width, device): + self._ensure_freqs_base(device) + batch_size = frame_indices.shape[0] + num_frames = frame_indices.shape[1] + + frame_indices = frame_indices.to(device=device, dtype=torch.float32) + grid_y, grid_x = self._get_spatial_meshgrid(height, width, str(device)) + + grid_t = frame_indices[:, :, None, None].expand( + batch_size, num_frames, height, width + ) + grid_y_batch = grid_y[None, None, :, :].expand(batch_size, num_frames, -1, -1) + grid_x_batch = grid_x[None, None, :, :].expand(batch_size, num_frames, -1, -1) + + freqs_cos_t, freqs_sin_t = self.get_frequency_batched( + self._freqs_base_t, grid_t + ) + freqs_cos_y, freqs_sin_y = self.get_frequency_batched( + self._freqs_base_y, grid_y_batch + ) + freqs_cos_x, freqs_sin_x = self.get_frequency_batched( + self._freqs_base_x, grid_x_batch + ) + + result = torch.cat( + [ + freqs_cos_t, + freqs_cos_y, + freqs_cos_x, + freqs_sin_t, + freqs_sin_y, + freqs_sin_x, + ], + dim=0, + ) + return result.permute(1, 0, 2, 3, 4) + + +# --------------------------------------------------------------------------- +# Condition Embedder +# --------------------------------------------------------------------------- + + +class HeliosTimeTextEmbedding(nn.Module): + """Condition embedder combining timestep and text embeddings.""" + + def __init__(self, dim, time_freq_dim, time_proj_dim, text_embed_dim): + super().__init__() + self.time_embedder = TimestepEmbedder( + dim, frequency_embedding_size=time_freq_dim, act_layer="silu" + ) + self.time_modulation = ModulateProjection(dim, factor=6, act_layer="silu") + self.text_embedder = MLP( + text_embed_dim, dim, dim, bias=True, act_type="gelu_pytorch_tanh" + ) + + def forward( + self, timestep, encoder_hidden_states, is_return_encoder_hidden_states=True + ): + temb = self.time_embedder(timestep) + timestep_proj = self.time_modulation(temb) + + if encoder_hidden_states is not None and is_return_encoder_hidden_states: + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + + return temb, timestep_proj, encoder_hidden_states + + +# --------------------------------------------------------------------------- +# Self-Attention for Helios +# --------------------------------------------------------------------------- + + +class HeliosSelfAttention(nn.Module): + """Self-attention with RMSNorm Q/K, optional history key amplification.""" + + def __init__( + self, + dim: int, + num_heads: int, + eps: float = 1e-6, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + tp_size = get_tp_world_size() + self.local_num_heads = divide(num_heads, tp_size) + + self.to_q = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_k = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_v = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_out = RowParallelLinear( + dim, dim, bias=True, reduce_results=True, quant_config=quant_config + ) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + self.tp_rmsnorm = tp_size > 1 + + self.attn = USPAttention( + num_heads=self.local_num_heads, + head_size=self.head_dim, + causal=False, + is_cross_attention=False, + ) + + self.is_amplify_history = is_amplify_history + if is_amplify_history: + if history_scale_mode == "scalar": + self.history_key_scale = nn.Parameter(torch.ones(1)) + elif history_scale_mode == "per_head": + self.history_key_scale = nn.Parameter(torch.ones(num_heads)) + else: + raise ValueError(f"Unknown history_scale_mode: {history_scale_mode}") + self.history_scale_mode = history_scale_mode + self.max_scale = 10.0 + + def forward(self, hidden_states, rotary_emb=None, original_context_length=None): + q, _ = self.to_q(hidden_states) + k, _ = self.to_k(hidden_states) + v, _ = self.to_v(hidden_states) + + if self.tp_rmsnorm: + q = tensor_parallel_rms_norm(q, self.norm_q) + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + q = self.norm_q(q) + k = self.norm_k(k) + + q = q.unflatten(2, (self.local_num_heads, self.head_dim)) + k = k.unflatten(2, (self.local_num_heads, self.head_dim)) + v = v.unflatten(2, (self.local_num_heads, self.head_dim)) + + if rotary_emb is not None: + q = apply_rotary_emb_transposed(q, rotary_emb) + k = apply_rotary_emb_transposed(k, rotary_emb) + + if self.is_amplify_history and original_context_length is not None: + history_seq_len = hidden_states.shape[1] - original_context_length + if history_seq_len > 0: + scale_key = 1.0 + torch.sigmoid(self.history_key_scale) * ( + self.max_scale - 1.0 + ) + if self.history_scale_mode == "per_head": + scale_key = scale_key.view(1, 1, -1, 1) + k = torch.cat( + [k[:, :history_seq_len] * scale_key, k[:, history_seq_len:]], + dim=1, + ) + + x = self.attn(q, k, v) + x = x.flatten(2) + x, _ = self.to_out(x) + return x + + +# --------------------------------------------------------------------------- +# Cross-Attention for Helios +# --------------------------------------------------------------------------- + + +class HeliosCrossAttention(nn.Module): + """Cross-attention with RMSNorm Q/K normalization.""" + + def __init__( + self, + dim: int, + num_heads: int, + eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + tp_size = get_tp_world_size() + self.local_num_heads = divide(num_heads, tp_size) + + self.to_q = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_k = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_v = ColumnParallelLinear( + dim, dim, bias=True, gather_output=False, quant_config=quant_config + ) + self.to_out = RowParallelLinear( + dim, dim, bias=True, reduce_results=True, quant_config=quant_config + ) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + self.tp_rmsnorm = tp_size > 1 + + self.attn = USPAttention( + num_heads=self.local_num_heads, + head_size=self.head_dim, + causal=False, + is_cross_attention=True, + ) + + def forward(self, hidden_states, encoder_hidden_states): + q, _ = self.to_q(hidden_states) + k, _ = self.to_k(encoder_hidden_states) + v, _ = self.to_v(encoder_hidden_states) + + if self.tp_rmsnorm: + q = tensor_parallel_rms_norm(q, self.norm_q) + k = tensor_parallel_rms_norm(k, self.norm_k) + else: + q = self.norm_q(q) + k = self.norm_k(k) + + q = q.unflatten(2, (self.local_num_heads, self.head_dim)) + k = k.unflatten(2, (self.local_num_heads, self.head_dim)) + v = v.unflatten(2, (self.local_num_heads, self.head_dim)) + + x = self.attn(q, k, v) + x = x.flatten(2) + x, _ = self.to_out(x) + return x + + +# --------------------------------------------------------------------------- +# Transformer Block +# --------------------------------------------------------------------------- + + +class HeliosTransformerBlock(nn.Module): + """ + Single transformer block with self-attention, cross-attention, FFN, + and scale-shift modulation from timestep embeddings. + """ + + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + cross_attn_norm: bool = True, + eps: float = 1e-6, + guidance_cross_attn: bool = True, + is_amplify_history: bool = False, + history_scale_mode: str = "per_head", + quant_config: QuantizationConfig | None = None, + ): + super().__init__() + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = HeliosSelfAttention( + dim=dim, + num_heads=num_heads, + eps=eps, + is_amplify_history=is_amplify_history, + history_scale_mode=history_scale_mode, + quant_config=quant_config, + ) + + # 2. Cross-attention + self.attn2 = HeliosCrossAttention( + dim=dim, + num_heads=num_heads, + eps=eps, + quant_config=quant_config, + ) + self.self_attn_residual_norm = ( + FP32LayerNorm(dim, eps, elementwise_affine=True) + if cross_attn_norm + else nn.Identity() + ) + + # 3. Feed-forward + self.ffn = MLP( + dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config + ) + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + # 4. Guidance cross-attention flag + self.guidance_cross_attn = guidance_cross_attn + + def forward( + self, + hidden_states, + encoder_hidden_states, + temb, + rotary_emb, + original_context_length=None, + ): + if temb.ndim == 4: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table.unsqueeze(0) + temb.float() + ).chunk(6, dim=2) + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + else: + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( + self.scale_shift_table + temb.float() + ).chunk(6, dim=1) + + # 1. Self-attention + norm_hidden_states = ( + self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa + ).type_as(hidden_states) + attn_output = self.attn1( + norm_hidden_states, rotary_emb, original_context_length + ) + hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as( + hidden_states + ) + + # 2. Cross-attention + if self.guidance_cross_attn: + history_seq_len = hidden_states.shape[1] - original_context_length + history_hidden_states, current_hidden_states = torch.split( + hidden_states, [history_seq_len, original_context_length], dim=1 + ) + norm_hidden_states = self.self_attn_residual_norm( + current_hidden_states.float() + ).type_as(current_hidden_states) + attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) + current_hidden_states = current_hidden_states + attn_output + hidden_states = torch.cat( + [history_hidden_states, current_hidden_states], dim=1 + ) + else: + norm_hidden_states = self.self_attn_residual_norm( + hidden_states.float() + ).type_as(hidden_states) + attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) + hidden_states = hidden_states + attn_output + + # 3. Feed-forward + norm_hidden_states = ( + self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa + ).type_as(hidden_states) + ff_output = self.ffn(norm_hidden_states) + hidden_states = ( + hidden_states.float() + ff_output.float() * c_gate_msa + ).type_as(hidden_states) + + return hidden_states + + +# --------------------------------------------------------------------------- +# Main model +# --------------------------------------------------------------------------- + + +class HeliosTransformer3DModel(CachableDiT, OffloadableDiTMixin): + """ + Helios Transformer 3D model for video generation. + + Implements multi-scale history patches, 3D RoPE, and chunked denoising + with zero_history_timestep and guidance_cross_attn. + """ + + _fsdp_shard_conditions = HeliosConfig()._fsdp_shard_conditions + _compile_conditions = HeliosConfig()._compile_conditions + _supported_attention_backends = HeliosConfig()._supported_attention_backends + param_names_mapping = HeliosConfig().param_names_mapping + reverse_param_names_mapping = HeliosConfig().reverse_param_names_mapping + lora_param_names_mapping = HeliosConfig().lora_param_names_mapping + + def __init__( + self, + config: HeliosConfig, + hf_config: dict[str, Any], + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + self.text_len = config.text_len + self.inner_dim = inner_dim + + # Helios-specific config + self.zero_history_timestep = config.zero_history_timestep + self.has_multi_term_memory_patch = config.has_multi_term_memory_patch + self.guidance_cross_attn = config.guidance_cross_attn + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed( + in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False, + ) + + # 2. Rotary position embeddings + self.rope = HeliosRotaryPosEmbed( + rope_dim=config.rope_dim, theta=config.rope_theta + ) + + # 3. Multi-term memory patches + if self.has_multi_term_memory_patch: + self.patch_short = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=config.patch_size, + stride=config.patch_size, + ) + self.patch_mid = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=tuple(2 * p for p in config.patch_size), + stride=tuple(2 * p for p in config.patch_size), + ) + self.patch_long = nn.Conv3d( + config.in_channels, + inner_dim, + kernel_size=tuple(4 * p for p in config.patch_size), + stride=tuple(4 * p for p in config.patch_size), + ) + + # 4. Condition embeddings + self.condition_embedder = HeliosTimeTextEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=config.text_dim, + ) + + # 5. Transformer blocks + self.blocks = nn.ModuleList( + [ + HeliosTransformerBlock( + dim=inner_dim, + ffn_dim=config.ffn_dim, + num_heads=config.num_attention_heads, + cross_attn_norm=config.cross_attn_norm, + eps=config.eps, + guidance_cross_attn=config.guidance_cross_attn, + is_amplify_history=config.is_amplify_history, + history_scale_mode=config.history_scale_mode, + quant_config=quant_config, + ) + for _ in range(config.num_layers) + ] + ) + + # 6. Output norm & projection + self.norm_out = HeliosOutputNorm(inner_dim, config.eps) + self.proj_out = ColumnParallelLinear( + inner_dim, + config.out_channels * math.prod(config.patch_size), + bias=True, + gather_output=True, + quant_config=quant_config, + ) + + self.cnt = 0 + self.__post_init__() + self.layer_names = ["blocks"] + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + # Stage 1 history inputs + indices_hidden_states=None, + indices_latents_history_short=None, + indices_latents_history_mid=None, + indices_latents_history_long=None, + latents_history_short=None, + latents_history_mid=None, + latents_history_long=None, + **kwargs, + ) -> torch.Tensor: + orig_dtype = hidden_states.dtype + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + + batch_size = hidden_states.shape[0] + p_t, p_h, p_w = self.patch_size + + # 1. Patch embed the noisy latents + hidden_states = self.patch_embedding(hidden_states) + _, _, post_patch_num_frames, post_patch_height, post_patch_width = ( + hidden_states.shape + ) + + if indices_hidden_states is None: + indices_hidden_states = ( + torch.arange(0, post_patch_num_frames) + .unsqueeze(0) + .expand(batch_size, -1) + ) + + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + # 2. Compute rotary embeddings + rotary_emb = self.rope( + frame_indices=indices_hidden_states, + height=post_patch_height, + width=post_patch_width, + device=hidden_states.device, + ) + rotary_emb = rotary_emb.flatten(2).transpose(1, 2) + original_context_length = hidden_states.shape[1] + + # 3. Process short history + if ( + latents_history_short is not None + and indices_latents_history_short is not None + ): + latents_history_short = latents_history_short.to(hidden_states) + latents_history_short = self.patch_short(latents_history_short) + _, _, _, H1, W1 = latents_history_short.shape + latents_history_short = latents_history_short.flatten(2).transpose(1, 2) + + rotary_emb_history_short = self.rope( + frame_indices=indices_latents_history_short, + height=H1, + width=W1, + device=latents_history_short.device, + ) + rotary_emb_history_short = rotary_emb_history_short.flatten(2).transpose( + 1, 2 + ) + hidden_states = torch.cat([latents_history_short, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_short, rotary_emb], dim=1) + + # 4. Process mid history + if latents_history_mid is not None and indices_latents_history_mid is not None: + latents_history_mid = latents_history_mid.to(hidden_states) + latents_history_mid = pad_for_3d_conv(latents_history_mid, (2, 4, 4)) + latents_history_mid = self.patch_mid(latents_history_mid) + latents_history_mid = latents_history_mid.flatten(2).transpose(1, 2) + + rotary_emb_history_mid = self.rope( + frame_indices=indices_latents_history_mid, + height=H1, + width=W1, + device=latents_history_mid.device, + ) + rotary_emb_history_mid = pad_for_3d_conv(rotary_emb_history_mid, (2, 2, 2)) + rotary_emb_history_mid = center_down_sample_3d( + rotary_emb_history_mid, (2, 2, 2) + ) + rotary_emb_history_mid = rotary_emb_history_mid.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_mid, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_mid, rotary_emb], dim=1) + + # 5. Process long history + if ( + latents_history_long is not None + and indices_latents_history_long is not None + ): + latents_history_long = latents_history_long.to(hidden_states) + latents_history_long = pad_for_3d_conv(latents_history_long, (4, 8, 8)) + latents_history_long = self.patch_long(latents_history_long) + latents_history_long = latents_history_long.flatten(2).transpose(1, 2) + + rotary_emb_history_long = self.rope( + frame_indices=indices_latents_history_long, + height=H1, + width=W1, + device=latents_history_long.device, + ) + rotary_emb_history_long = pad_for_3d_conv( + rotary_emb_history_long, (4, 4, 4) + ) + rotary_emb_history_long = center_down_sample_3d( + rotary_emb_history_long, (4, 4, 4) + ) + rotary_emb_history_long = rotary_emb_history_long.flatten(2).transpose(1, 2) + + hidden_states = torch.cat([latents_history_long, hidden_states], dim=1) + rotary_emb = torch.cat([rotary_emb_history_long, rotary_emb], dim=1) + + history_context_length = hidden_states.shape[1] - original_context_length + + # 6. Compute condition embeddings + if indices_hidden_states is not None and self.zero_history_timestep: + timestep_t0 = torch.zeros( + (1,), dtype=timestep.dtype, device=timestep.device + ) + temb_t0, timestep_proj_t0, _ = self.condition_embedder( + timestep_t0, + encoder_hidden_states, + is_return_encoder_hidden_states=False, + ) + temb_t0 = temb_t0.unsqueeze(1).expand( + batch_size, history_context_length, -1 + ) + timestep_proj_t0 = ( + timestep_proj_t0.unflatten(-1, (6, -1)) + .view(1, 6, 1, -1) + .expand(batch_size, -1, history_context_length, -1) + ) + + temb, timestep_proj, encoder_hidden_states = self.condition_embedder( + timestep, encoder_hidden_states + ) + timestep_proj = timestep_proj.unflatten(-1, (6, -1)) + + if indices_hidden_states is not None and not self.zero_history_timestep: + main_repeat_size = hidden_states.shape[1] + else: + main_repeat_size = original_context_length + temb = temb.view(batch_size, 1, -1).expand(batch_size, main_repeat_size, -1) + timestep_proj = timestep_proj.view(batch_size, 6, 1, -1).expand( + batch_size, 6, main_repeat_size, -1 + ) + + if indices_hidden_states is not None and self.zero_history_timestep: + temb = torch.cat([temb_t0, temb], dim=1) + timestep_proj = torch.cat([timestep_proj_t0, timestep_proj], dim=2) + + if timestep_proj.ndim == 4: + timestep_proj = timestep_proj.permute(0, 2, 1, 3) + + # 7. Transformer blocks + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + rotary_emb = rotary_emb.contiguous() + + for block in self.blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + rotary_emb, + original_context_length, + ) + + self.cnt += 1 + + # 8. Output norm & projection + hidden_states = self.norm_out(hidden_states, temb, original_context_length) + hidden_states, _ = self.proj_out(hidden_states) + + # 9. Unpatchify + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + return output + + +EntryClass = HeliosTransformer3DModel diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py new file mode 100644 index 000000000..3b3f64e92 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_helios.py @@ -0,0 +1,734 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from Helios diffusers scheduler: +# https://github.com/BestWishYsh/Helios +""" +Helios scheduler implementing flow-matching with UniPC/Euler solvers. + +For Phase 1 T2V (stages=1), this simplifies to standard flow-matching +with dynamic shifting and UniPC multistep solver. +""" + +import math +from dataclasses import dataclass + +import numpy as np +import torch + + +@dataclass +class HeliosSchedulerOutput: + prev_sample: torch.FloatTensor + model_outputs: torch.FloatTensor | None = None + last_sample: torch.FloatTensor | None = None + this_order: int | None = None + + +class HeliosSchedulerConfig: + """Mimics diffusers config interface for scheduler parameters.""" + + def __init__(self, **kwargs): + for k, v in kwargs.items(): + setattr(self, k, v) + + def get(self, key, default=None): + return getattr(self, key, default) + + +class HeliosScheduler: + """ + Helios multi-stage scheduler supporting Euler, UniPC, and DMD solvers. + + For Phase 1 T2V with stages=1, this is a standard flow-matching scheduler + with optional time shifting and UniPC multistep updates. + """ + + order = 1 + + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + stages: int = 1, + stage_range: list | None = None, + gamma: float = 1 / 3, + thresholding: bool = False, + prediction_type: str = "flow_prediction", + solver_order: int = 2, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: list[int] | None = None, + use_flow_sigmas: bool = True, + scheduler_type: str = "unipc", + use_dynamic_shifting: bool = False, + time_shift_type: str = "linear", + **kwargs, + ): + if stage_range is None: + # Evenly divide [0, 1] into 3 stages for pyramid SR + stage_range = [0, 1 / 3, 2 / 3, 1] + if disable_corrector is None: + disable_corrector = [] + + self.config = HeliosSchedulerConfig( + num_train_timesteps=num_train_timesteps, + shift=shift, + stages=stages, + stage_range=stage_range, + gamma=gamma, + thresholding=thresholding, + prediction_type=prediction_type, + solver_order=solver_order, + predict_x0=predict_x0, + solver_type=solver_type, + lower_order_final=lower_order_final, + disable_corrector=disable_corrector, + use_flow_sigmas=use_flow_sigmas, + scheduler_type=scheduler_type, + use_dynamic_shifting=use_dynamic_shifting, + time_shift_type=time_shift_type, + ) + + self.timestep_ratios = {} + self.timesteps_per_stage = {} + self.sigmas_per_stage = {} + self.start_sigmas = {} + self.end_sigmas = {} + self.ori_start_sigmas = {} + + self.init_sigmas_for_each_stage() + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + self.gamma = gamma + + if solver_type not in ["bh1", "bh2"]: + raise NotImplementedError(f"{solver_type} is not implemented") + + self.predict_x0 = predict_x0 + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = None + self.last_sample = None + self._step_index = None + self._begin_index = None + self.num_inference_steps = None + + def init_sigmas(self): + num_train_timesteps = self.config.num_train_timesteps + shift = self.config.shift + + alphas = np.linspace(1, 1 / num_train_timesteps, num_train_timesteps + 1) + sigmas = 1.0 - alphas + sigmas = np.flip(shift * sigmas / (1 + (shift - 1) * sigmas))[:-1].copy() + sigmas = torch.from_numpy(sigmas) + timesteps = (sigmas * num_train_timesteps).clone() + + self._step_index = None + self._begin_index = None + self.timesteps = timesteps + self.sigmas = sigmas.to("cpu") + + def init_sigmas_for_each_stage(self): + self.init_sigmas() + + stage_distance = [] + stages = self.config.stages + training_steps = self.config.num_train_timesteps + stage_range = self.config.stage_range + + for i_s in range(stages): + start_indice = int(stage_range[i_s] * training_steps) + start_indice = max(start_indice, 0) + end_indice = int(stage_range[i_s + 1] * training_steps) + end_indice = min(end_indice, training_steps) + start_sigma = self.sigmas[start_indice].item() + end_sigma = ( + self.sigmas[end_indice].item() if end_indice < training_steps else 0.0 + ) + self.ori_start_sigmas[i_s] = start_sigma + + if i_s != 0: + ori_sigma = 1 - start_sigma + gamma = self.config.gamma + corrected_sigma = ( + 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + ) * ori_sigma + start_sigma = 1 - corrected_sigma + + stage_distance.append(start_sigma - end_sigma) + self.start_sigmas[i_s] = start_sigma + self.end_sigmas[i_s] = end_sigma + + tot_distance = sum(stage_distance) + for i_s in range(stages): + if i_s == 0: + start_ratio = 0.0 + else: + start_ratio = sum(stage_distance[:i_s]) / tot_distance + if i_s == stages - 1: + # Use value just below 1.0 to avoid out-of-bounds indexing + end_ratio = 1.0 - 1e-16 + else: + end_ratio = sum(stage_distance[: i_s + 1]) / tot_distance + self.timestep_ratios[i_s] = (start_ratio, end_ratio) + + for i_s in range(stages): + timestep_ratio = self.timestep_ratios[i_s] + # Clamp to max valid timestep (num_train_timesteps - 1) + timestep_max = min( + self.timesteps[int(timestep_ratio[0] * training_steps)], 999 + ) + timestep_min = self.timesteps[ + min(int(timestep_ratio[1] * training_steps), training_steps - 1) + ] + timesteps = np.linspace(timestep_max, timestep_min, training_steps + 1) + self.timesteps_per_stage[i_s] = ( + timesteps[:-1] + if isinstance(timesteps, torch.Tensor) + else torch.from_numpy(timesteps[:-1]) + ) + # Sigma range [0.999, 0]: start just below 1.0 to avoid singularity + stage_sigmas = np.linspace(0.999, 0, training_steps + 1) + self.sigmas_per_stage[i_s] = torch.from_numpy(stage_sigmas[:-1]) + + @property + def step_index(self): + return self._step_index + + @property + def begin_index(self): + return self._begin_index + + def set_begin_index(self, begin_index: int = 0): + self._begin_index = begin_index + + def time_shift(self, mu, sigma, t): + if self.config.time_shift_type == "exponential": + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + elif self.config.time_shift_type == "linear": + return mu / (mu + (1 / t - 1) ** sigma) + + def set_timesteps( + self, + num_inference_steps: int, + stage_index: int | None = None, + device: str | torch.device = None, + sigmas=None, + mu=None, + is_amplify_first_chunk: bool = False, + ): + if self.config.scheduler_type == "dmd": + if is_amplify_first_chunk: + num_inference_steps = num_inference_steps * 2 + 1 + else: + num_inference_steps = num_inference_steps + 1 + + self.num_inference_steps = num_inference_steps + self.init_sigmas() + + if self.config.stages == 1: + if sigmas is None: + sigmas = np.linspace( + 1, + 1 / self.config.num_train_timesteps, + num_inference_steps + 1, + )[:-1].astype(np.float32) + if self.config.shift != 1.0: + assert not self.config.use_dynamic_shifting + sigmas = self.time_shift(self.config.shift, 1.0, sigmas) + timesteps = (sigmas * self.config.num_train_timesteps).copy() + sigmas = torch.from_numpy(sigmas) + else: + stage_timesteps = self.timesteps_per_stage[stage_index] + timesteps = np.linspace( + stage_timesteps[0].item(), + stage_timesteps[-1].item(), + num_inference_steps, + ) + stage_sigmas = self.sigmas_per_stage[stage_index] + ratios = np.linspace( + stage_sigmas[0].item(), stage_sigmas[-1].item(), num_inference_steps + ) + sigmas = torch.from_numpy(ratios) + + self.timesteps = torch.from_numpy(timesteps).to(device=device) + self.sigmas = torch.cat([sigmas, torch.zeros(1)]).to(device=device) + + self._step_index = None + self.reset_scheduler_history() + + if self.config.scheduler_type == "dmd": + self.timesteps = self.timesteps[:-1] + self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]]) + + if self.config.use_dynamic_shifting: + assert self.config.shift == 1.0 + self.sigmas = self.time_shift(mu, 1.0, self.sigmas) + if self.config.stages == 1: + self.timesteps = self.sigmas[:-1] * self.config.num_train_timesteps + else: + self.timesteps = self.timesteps_per_stage[ + stage_index + ].min() + self.sigmas[:-1] * ( + self.timesteps_per_stage[stage_index].max() + - self.timesteps_per_stage[stage_index].min() + ) + + # ---------------------------------- Euler ---------------------------------- + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + indices = (schedule_timesteps == timestep).nonzero() + pos = 1 if len(indices) > 1 else 0 + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step_euler( + self, + model_output: torch.FloatTensor, + timestep=None, + sample: torch.FloatTensor = None, + return_dict: bool = True, + **kwargs, + ) -> HeliosSchedulerOutput | tuple: + if self.step_index is None: + self._step_index = 0 + + sample = sample.to(torch.float32) + sigma = self.sigmas[self.step_index] + sigma_next = self.sigmas[self.step_index + 1] + + prev_sample = sample + (sigma_next - sigma) * model_output + prev_sample = prev_sample.to(model_output.dtype) + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- UniPC ---------------------------------- + def _sigma_to_alpha_sigma_t(self, sigma): + if self.config.use_flow_sigmas: + alpha_t = 1 - sigma + sigma_t = torch.clamp(sigma, min=1e-8) + else: + alpha_t = 1 / ((sigma**2 + 1) ** 0.5) + sigma_t = sigma * alpha_t + return alpha_t, sigma_t + + def convert_model_output(self, model_output, sample=None, sigma=None, **kwargs): + flag = False + if sigma is None: + flag = True + sigma = self.sigmas[self.step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + if flag: + sigma_t = self.sigmas[self.step_index] + else: + sigma_t = sigma + x0_pred = sample - sigma_t * model_output + elif self.config.prediction_type == "epsilon": + x0_pred = (sample - sigma_t * model_output) / alpha_t + elif self.config.prediction_type == "sample": + x0_pred = model_output + elif self.config.prediction_type == "v_prediction": + x0_pred = alpha_t * sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type {self.config.prediction_type} not supported" + ) + return x0_pred + else: + if self.config.prediction_type == "epsilon": + return model_output + elif self.config.prediction_type == "sample": + return (sample - alpha_t * model_output) / sigma_t + elif self.config.prediction_type == "v_prediction": + return alpha_t * model_output + sigma_t * sample + else: + raise ValueError( + f"prediction_type {self.config.prediction_type} not supported" + ) + + def multistep_uni_p_bh_update( + self, model_output, sample=None, order=None, sigma=None, sigma_next=None + ): + model_output_list = self.model_outputs + m0 = model_output_list[-1] + x = sample + + if sigma_next is None and sigma is None: + sigma_t, sigma_s0 = ( + self.sigmas[self.step_index + 1], + self.sigmas[self.step_index], + ) + else: + sigma_t, sigma_s0 = sigma_next, sigma + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + device = sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + if order == 2: + rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1]).to(device).to(x.dtype) + else: + D1s = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + pred_res = ( + torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 + ) + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + pred_res = ( + torch.einsum("k,bkc...->bc...", rhos_p, D1s) if D1s is not None else 0 + ) + x_t = x_t_ - sigma_t * B_h * pred_res + + return x_t.to(x.dtype) + + def multistep_uni_c_bh_update( + self, + this_model_output, + last_sample=None, + this_sample=None, + order=None, + sigma_before=None, + sigma=None, + ): + model_output_list = self.model_outputs + m0 = model_output_list[-1] + x = last_sample + model_t = this_model_output + + if sigma_before is None and sigma is None: + sigma_t, sigma_s0 = ( + self.sigmas[self.step_index], + self.sigmas[self.step_index - 1], + ) + else: + sigma_t, sigma_s0 = sigma, sigma_before + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + h = lambda_t - lambda_s0 + device = this_sample.device + + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(1.0) + rks = torch.tensor(rks, device=device) + + R = [] + b = [] + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R) + b = torch.tensor(b, device=device) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + if order == 1: + rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device) + else: + rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + corr_res = ( + torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + if D1s is not None + else 0 + ) + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + corr_res = ( + torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + if D1s is not None + else 0 + ) + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + + return x_t.to(x.dtype) + + def step_unipc( + self, + model_output, + timestep=None, + sample=None, + return_dict: bool = True, + **kwargs, + ) -> HeliosSchedulerOutput | tuple: + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', run 'set_timesteps' first" + ) + + if self.step_index is None: + self._step_index = 0 + + use_corrector = ( + self.step_index > 0 + and self.step_index - 1 not in self.disable_corrector + and self.last_sample is not None + ) + + model_output_convert = self.convert_model_output(model_output, sample=sample) + + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min( + self.config.solver_order, len(self.timesteps) - self.step_index + ) + else: + this_order = self.config.solver_order + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + ) + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + self._step_index += 1 + + if not return_dict: + return (prev_sample,) + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- DMD ---------------------------------- + def add_noise(self, original_samples, noise, timestep, sigmas, timesteps): + sigmas = sigmas.to(noise.device) + timesteps = timesteps.to(noise.device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 + ) + sigma = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def convert_flow_pred_to_x0(self, flow_pred, xt, timestep, sigmas, timesteps): + original_dtype = flow_pred.dtype + device = flow_pred.device + flow_pred, xt, sigmas, timesteps = ( + x.double().to(device) for x in (flow_pred, xt, sigmas, timesteps) + ) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1 + ) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1, 1) + x0_pred = xt - sigma_t * flow_pred + return x0_pred.to(original_dtype) + + def step_dmd( + self, + model_output: torch.FloatTensor, + timestep=None, + sample: torch.FloatTensor = None, + return_dict: bool = True, + cur_sampling_step: int = 0, + dmd_noisy_tensor: torch.FloatTensor | None = None, + dmd_sigmas: torch.FloatTensor | None = None, + dmd_timesteps: torch.FloatTensor | None = None, + all_timesteps: torch.FloatTensor | None = None, + **kwargs, + ) -> HeliosSchedulerOutput | tuple: + pred_image_or_video = self.convert_flow_pred_to_x0( + flow_pred=model_output, + xt=sample, + timestep=torch.full( + (model_output.shape[0],), + timestep, + dtype=torch.long, + device=model_output.device, + ), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + if cur_sampling_step < len(all_timesteps) - 1: + prev_sample = self.add_noise( + pred_image_or_video, + dmd_noisy_tensor, + torch.full( + (model_output.shape[0],), + all_timesteps[cur_sampling_step + 1], + dtype=torch.long, + device=model_output.device, + ), + sigmas=dmd_sigmas, + timesteps=dmd_timesteps, + ) + else: + prev_sample = pred_image_or_video + + if not return_dict: + return (prev_sample,) + return HeliosSchedulerOutput(prev_sample=prev_sample) + + # ---------------------------------- Main step ---------------------------------- + def step( + self, + model_output, + timestep=None, + sample=None, + return_dict: bool = True, + **kwargs, + ) -> HeliosSchedulerOutput | tuple: + if self.config.scheduler_type == "euler": + return self.step_euler( + model_output=model_output, + timestep=timestep, + sample=sample, + return_dict=return_dict, + ) + elif self.config.scheduler_type == "unipc": + return self.step_unipc( + model_output=model_output, + timestep=timestep, + sample=sample, + return_dict=return_dict, + ) + elif self.config.scheduler_type == "dmd": + return self.step_dmd( + model_output=model_output, + timestep=timestep, + sample=sample, + return_dict=return_dict, + **kwargs, + ) + else: + raise NotImplementedError( + f"Scheduler type '{self.config.scheduler_type}' not implemented" + ) + + def reset_scheduler_history(self): + self.model_outputs = [None] * self.config.solver_order + self.timestep_list = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.disable_corrector = self.config.disable_corrector + self.solver_p = None + self.last_sample = None + self._step_index = None + self._begin_index = None + + def set_shift(self, shift: float): + """Update the shift parameter (called by SchedulerLoader after loading).""" + self.config.shift = shift + self.shift = shift + + def __len__(self): + return self.config.num_train_timesteps + + +EntryClass = HeliosScheduler diff --git a/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py new file mode 100644 index 000000000..eb6a9e398 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/helios_pipeline.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Helios video diffusion pipeline implementation. + +This module contains an implementation of the Helios video diffusion pipeline +using the modular pipeline architecture. Phase 1: T2V only. +""" + +from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( + ComposedPipelineBase, +) +from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline +from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + InputValidationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import ( + HeliosChunkedDenoisingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +class HeliosPipeline(LoRAPipeline, ComposedPipelineBase): + """ + Helios video diffusion pipeline with LoRA support. + + Implements the Helios T2V pipeline with chunked denoising, + multi-term memory history, and CFG Zero Star guidance. + """ + + pipeline_name = "HeliosPipeline" + + _required_config_modules = [ + "text_encoder", + "tokenizer", + "vae", + "transformer", + "scheduler", + ] + + def initialize_pipeline(self, server_args: ServerArgs): + # Use the scheduler loaded from model's scheduler_config.json as-is. + # It contains critical config: use_dynamic_shifting=true, + # time_shift_type="exponential", etc. + scheduler = self.modules.get("scheduler") + if scheduler is not None and server_args.pipeline_config.flow_shift is not None: + scheduler.set_shift(server_args.pipeline_config.flow_shift) + + # Configure scheduler for Stage 2/3 if enabled + pipeline_config = server_args.pipeline_config + if scheduler is not None and pipeline_config.is_enable_stage2: + scheduler.config.stages = pipeline_config.pyramid_num_stages + scheduler.config.scheduler_type = pipeline_config.scheduler_type + scheduler.config.gamma = pipeline_config.gamma + scheduler.init_sigmas_for_each_stage() + + def create_pipeline_stages(self, server_args: ServerArgs) -> None: + self.add_stage(InputValidationStage()) + self.add_standard_text_encoding_stage() + self.add_standard_latent_preparation_stage() + # Skip standard timestep preparation — the Helios denoising stage + # handles scheduler.set_timesteps internally per-chunk with mu. + self.add_stage( + HeliosChunkedDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.modules["scheduler"], + ), + "helios_chunked_denoising_stage", + ) + # Standard DecodingStage handles VAE decode of the denoised latents + self.add_standard_decoding_stage() + + +class HeliosPyramidPipeline(HeliosPipeline): + """Helios pyramid SR pipeline (used by Helios-Mid and Helios-Distilled).""" + + pipeline_name = "HeliosPyramidPipeline" + + +EntryClass = [HeliosPipeline, HeliosPyramidPipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py new file mode 100644 index 000000000..ed7c48c08 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/helios_denoising.py @@ -0,0 +1,666 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Helios-specific chunked denoising stage. + +Implements Stage 1 chunked denoising with multi-term memory history +and CFG Zero Star guidance. VAE decoding is handled by the standard +DecodingStage downstream. +""" + +import math + +import numpy as np +import torch +import torch.nn.functional as F + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( + PipelineStage, + StageParallelismType, +) +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__) + + +def optimized_scale(positive_flat, negative_flat): + """CFG Zero Star: compute optimal guidance scale.""" + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + squared_norm = torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8 + return dot_product / squared_norm + + +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +def sample_block_noise( + batch_size, channel, num_frames, height, width, gamma, patch_size=(1, 2, 2) +): + """Generate spatially-correlated block noise for pyramid SR.""" + _, ph, pw = patch_size + block_size = ph * pw + + cov = ( + torch.eye(block_size) * (1 + gamma) - torch.ones(block_size, block_size) * gamma + ) + dist = torch.distributions.MultivariateNormal( + torch.zeros(block_size, device=cov.device), covariance_matrix=cov + ) + block_number = batch_size * channel * num_frames * (height // ph) * (width // pw) + + noise = dist.sample((block_number,)) + noise = noise.view( + batch_size, channel, num_frames, height // ph, width // pw, ph, pw + ) + noise = noise.permute(0, 1, 2, 3, 5, 4, 6).reshape( + batch_size, channel, num_frames, height, width + ) + return noise + + +class HeliosChunkedDenoisingStage(PipelineStage): + """ + Helios chunked denoising stage implementing Stage 1 loop. + + Iterates over video chunks, manages history buffers (short/mid/long), + runs transformer per chunk with CFG guidance, scheduler step, + and accumulates denoised latents. VAE decoding is left to DecodingStage. + """ + + def __init__(self, transformer, scheduler): + super().__init__() + self.transformer = transformer + self.scheduler = scheduler + + @property + def parallelism_type(self): + return StageParallelismType.REPLICATED + + def _denoise_one_chunk( + self, + latents, + prompt_embeds, + negative_prompt_embeds, + timesteps, + guidance_scale, + indices_hidden_states, + indices_latents_history_short, + indices_latents_history_mid, + indices_latents_history_long, + latents_history_short, + latents_history_mid, + latents_history_long, + target_dtype, + device, + is_cfg_zero_star=True, + use_zero_init=True, + zero_steps=1, + batch=None, + server_args=None, + ): + """Denoise a single chunk with full timestep loop.""" + batch_size = latents.shape[0] + do_cfg = guidance_scale > 1.0 + + for i, t in enumerate(timesteps): + timestep = t.expand(batch_size) + latent_model_input = latents.to(target_dtype) + + with set_forward_context( + current_timestep=t, + forward_batch=None, + attn_metadata=None, + ): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), + ) + + if do_cfg: + with set_forward_context( + current_timestep=t, + forward_batch=None, + attn_metadata=None, + ): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), + ) + + if is_cfg_zero_star: + noise_pred_text = noise_pred + positive_flat = noise_pred_text.reshape(batch_size, -1) + negative_flat = noise_uncond.reshape(batch_size, -1) + + alpha = optimized_scale(positive_flat, negative_flat) + alpha = alpha.view( + batch_size, *([1] * (len(noise_pred_text.shape) - 1)) + ) + alpha = alpha.to(noise_pred_text.dtype) + + if (i <= zero_steps) and use_zero_init: + noise_pred = noise_pred_text * 0.0 + else: + noise_pred = noise_uncond * alpha + guidance_scale * ( + noise_pred_text - noise_uncond * alpha + ) + else: + noise_pred = noise_uncond + guidance_scale * ( + noise_pred - noise_uncond + ) + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + return latents + + def _denoise_one_chunk_stage2( + self, + latents, + prompt_embeds, + negative_prompt_embeds, + guidance_scale, + indices_hidden_states, + indices_latents_history_short, + indices_latents_history_mid, + indices_latents_history_long, + latents_history_short, + latents_history_mid, + latents_history_long, + target_dtype, + device, + pyramid_num_stages, + pyramid_num_inference_steps_list, + is_distilled, + is_amplify_first_chunk, + gamma, + is_cfg_zero_star=True, + use_zero_init=True, + zero_steps=1, + batch=None, + server_args=None, + ): + """Denoise a single chunk using pyramid super-resolution (Stage 2).""" + batch_size, num_channel, num_frames, height, width = latents.shape + patch_size = self.transformer.patch_size + + # Downsample to lowest pyramid level + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames, num_channel, height, width + ) + for _ in range(pyramid_num_stages - 1): + height //= 2 + width //= 2 + latents = F.interpolate(latents, size=(height, width), mode="bilinear") * 2 + latents = latents.reshape( + batch_size, num_frames, num_channel, height, width + ).permute(0, 2, 1, 3, 4) + + start_point_list = None + if is_distilled: + start_point_list = [latents] + + do_cfg = guidance_scale > 1.0 + + for i_s in range(pyramid_num_stages): + # Compute mu for current resolution + image_seq_len = ( + latents.shape[-1] + * latents.shape[-2] + * latents.shape[-3] + // (patch_size[0] * patch_size[1] * patch_size[2]) + ) + mu = calculate_shift(image_seq_len) + + self.scheduler.set_timesteps( + pyramid_num_inference_steps_list[i_s], + i_s, + device=device, + mu=mu, + is_amplify_first_chunk=is_amplify_first_chunk, + ) + timesteps = self.scheduler.timesteps + + if i_s > 0: + # Upsample 2x nearest-neighbor + height *= 2 + width *= 2 + latents = latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames, + num_channel, + height // 2, + width // 2, + ) + latents = F.interpolate(latents, size=(height, width), mode="nearest") + latents = latents.reshape( + batch_size, num_frames, num_channel, height, width + ).permute(0, 2, 1, 3, 4) + + # Renoise with correlated block noise + ori_sigma = 1 - self.scheduler.ori_start_sigmas[i_s] + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + + bs, ch, nf, h, w = latents.shape + noise = sample_block_noise(bs, ch, nf, h, w, gamma, patch_size) + noise = noise.to(device=device, dtype=target_dtype) + latents = alpha * latents + beta * noise + + if is_distilled: + start_point_list.append(latents) + + # Denoising loop for this pyramid stage + for idx, t in enumerate(timesteps): + timestep = t.expand(batch_size) + latent_model_input = latents.to(target_dtype) + + with set_forward_context( + current_timestep=t, + forward_batch=None, + attn_metadata=None, + ): + noise_pred = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), + ) + + if do_cfg: + with set_forward_context( + current_timestep=t, + forward_batch=None, + attn_metadata=None, + ): + noise_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=( + latents_history_short.to(target_dtype) + if latents_history_short is not None + else None + ), + latents_history_mid=( + latents_history_mid.to(target_dtype) + if latents_history_mid is not None + else None + ), + latents_history_long=( + latents_history_long.to(target_dtype) + if latents_history_long is not None + else None + ), + ) + + if is_cfg_zero_star: + noise_pred_text = noise_pred + positive_flat = noise_pred_text.reshape(batch_size, -1) + negative_flat = noise_uncond.reshape(batch_size, -1) + + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view( + batch_size, + *([1] * (len(noise_pred_text.shape) - 1)), + ) + alpha_cfg = alpha_cfg.to(noise_pred_text.dtype) + + if (i_s == 0 and idx <= zero_steps) and use_zero_init: + noise_pred = noise_pred_text * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + guidance_scale * ( + noise_pred_text - noise_uncond * alpha_cfg + ) + else: + noise_pred = noise_uncond + guidance_scale * ( + noise_pred - noise_uncond + ) + + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + cur_sampling_step=idx, + dmd_noisy_tensor=( + start_point_list[i_s] if start_point_list is not None else None + ), + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=timesteps, + )[0] + + return latents + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + """Run the Helios chunked denoising loop.""" + pipeline_config = server_args.pipeline_config + device = ( + batch.latents.device + if hasattr(batch, "latents") and batch.latents is not None + else torch.device("cuda") + ) + target_dtype = PRECISION_TO_TYPE.get( + server_args.pipeline_config.precision, torch.bfloat16 + ) + + # Get config params + num_latent_frames_per_chunk = pipeline_config.num_latent_frames_per_chunk + history_sizes = sorted(list(pipeline_config.history_sizes), reverse=True) + is_cfg_zero_star = pipeline_config.is_cfg_zero_star + zero_steps = pipeline_config.zero_steps + keep_first_frame = pipeline_config.keep_first_frame + guidance_scale = batch.guidance_scale + num_inference_steps = batch.num_inference_steps + + # Stage 2 params + is_enable_stage2 = pipeline_config.is_enable_stage2 + pyramid_num_stages = pipeline_config.pyramid_num_stages + pyramid_num_inference_steps_list = ( + pipeline_config.pyramid_num_inference_steps_list + ) + is_distilled = pipeline_config.is_distilled + is_amplify_first_chunk = pipeline_config.is_amplify_first_chunk + gamma = pipeline_config.gamma + + # Move transformer to GPU if CPU-offloaded + if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: + if next(self.transformer.parameters()).device.type == "cpu": + self.transformer.to(get_local_torch_device()) + + # Get encoder outputs (prompt_embeds is a list of tensors, one per encoder) + prompt_embeds = batch.prompt_embeds + if isinstance(prompt_embeds, list): + prompt_embeds = prompt_embeds[0] + prompt_embeds = prompt_embeds.to(target_dtype) + negative_prompt_embeds = batch.negative_prompt_embeds + if isinstance(negative_prompt_embeds, list): + negative_prompt_embeds = ( + negative_prompt_embeds[0] if negative_prompt_embeds else None + ) + if negative_prompt_embeds is not None: + negative_prompt_embeds = negative_prompt_embeds.to(target_dtype) + + # Scale factors inherited from the Wan VAE used by Helios + # (AutoencoderKLWan: temporal_compression_ratio=4, spatial_compression_ratio=8) + vae_scale_factor_temporal = 4 + vae_scale_factor_spatial = 8 + + # Compute chunking + height = batch.height + width = batch.width + num_frames = batch.num_frames + num_channels_latents = self.transformer.in_channels + + window_num_frames = ( + num_latent_frames_per_chunk - 1 + ) * vae_scale_factor_temporal + 1 + num_latent_chunk = max( + 1, (num_frames + window_num_frames - 1) // window_num_frames + ) + num_history_latent_frames = sum(history_sizes) + batch_size = 1 # Helios processes one video at a time + + # Prepare history latents + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // vae_scale_factor_spatial, + width // vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + + # Build frame indices + if keep_first_frame: + indices = torch.arange( + 0, sum([1, *history_sizes, num_latent_frames_per_chunk]) + ) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat( + [indices_prefix, indices_latents_history_1x], dim=0 + ) + else: + indices = torch.arange( + 0, sum([*history_sizes, num_latent_frames_per_chunk]) + ) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + + indices_hidden_states = indices_hidden_states.unsqueeze(0) + indices_latents_history_short = indices_latents_history_short.unsqueeze(0) + indices_latents_history_mid = indices_latents_history_mid.unsqueeze(0) + indices_latents_history_long = indices_latents_history_long.unsqueeze(0) + + # Set up scheduler + patch_size = self.transformer.patch_size + image_seq_len = ( + num_latent_frames_per_chunk + * (height // vae_scale_factor_spatial) + * (width // vae_scale_factor_spatial) + // (patch_size[0] * patch_size[1] * patch_size[2]) + ) + # Sigma schedule from near-1.0 (pure noise) to 0.0 (clean); 0.999 avoids singularity + sigmas = np.linspace(0.999, 0.0, num_inference_steps + 1)[:-1] + mu = calculate_shift(image_seq_len) + + # Chunk loop + image_latents = None + total_generated_latent_frames = 0 + + self.log_info( + f"Starting chunked denoising: {num_latent_chunk} chunks, " + f"{num_inference_steps} steps each" + ) + + for k in range(num_latent_chunk): + is_first_chunk = k == 0 + + # Extract history + if keep_first_frame: + ( + latents_history_long, + latents_history_mid, + latents_history_1x, + ) = history_latents[:, :, -num_history_latent_frames:].split( + history_sizes, dim=2 + ) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + batch_size, + num_channels_latents, + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat( + [latents_prefix, latents_history_1x], dim=2 + ) + else: + ( + latents_history_long, + latents_history_mid, + latents_history_short, + ) = history_latents[:, :, -num_history_latent_frames:].split( + history_sizes, dim=2 + ) + + # Generate noise latents for this chunk + latents = torch.randn( + batch_size, + num_channels_latents, + (window_num_frames - 1) // vae_scale_factor_temporal + 1, + height // vae_scale_factor_spatial, + width // vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + + if is_enable_stage2: + # Stage 2: Pyramid SR denoising (handles scheduler internally) + latents = self._denoise_one_chunk_stage2( + latents=latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + guidance_scale=guidance_scale, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + target_dtype=target_dtype, + device=device, + pyramid_num_stages=pyramid_num_stages, + pyramid_num_inference_steps_list=pyramid_num_inference_steps_list, + is_distilled=is_distilled, + is_amplify_first_chunk=(is_amplify_first_chunk and is_first_chunk), + gamma=gamma, + is_cfg_zero_star=is_cfg_zero_star, + use_zero_init=True, + zero_steps=zero_steps, + batch=batch, + server_args=server_args, + ) + else: + # Stage 1: Standard flat denoising + self.scheduler.set_timesteps( + num_inference_steps, device=device, sigmas=sigmas, mu=mu + ) + timesteps = self.scheduler.timesteps + + latents = self._denoise_one_chunk( + latents=latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + timesteps=timesteps, + guidance_scale=guidance_scale, + indices_hidden_states=indices_hidden_states, + indices_latents_history_short=indices_latents_history_short, + indices_latents_history_mid=indices_latents_history_mid, + indices_latents_history_long=indices_latents_history_long, + latents_history_short=latents_history_short, + latents_history_mid=latents_history_mid, + latents_history_long=latents_history_long, + target_dtype=target_dtype, + device=device, + is_cfg_zero_star=is_cfg_zero_star, + use_zero_init=True, + zero_steps=zero_steps, + batch=batch, + server_args=server_args, + ) + + # Extract first frame as image_latents for subsequent chunks + if keep_first_frame and is_first_chunk and image_latents is None: + image_latents = latents[:, :, 0:1, :, :] + + # Update history + total_generated_latent_frames += latents.shape[2] + history_latents = torch.cat([history_latents, latents], dim=2) + + # Move transformer back to CPU after denoising + if server_args.dit_cpu_offload and not server_args.use_fsdp_inference: + if next(self.transformer.parameters()).device.type != "cpu": + self.transformer.to("cpu") + torch.cuda.empty_cache() + + # Store denoised latents for the standard DecodingStage to decode + batch.latents = history_latents[:, :, -total_generated_latent_frames:] + + return batch