diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index 392f3d9b9..03ebd5098 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/docs/diffusion/compatibility_matrix.md @@ -33,23 +33,26 @@ default parameters when initializing and generating videos. | TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | +| LTX-2 | `Lightricks/LTX-2` | 1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | **Note**: -1.Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. -2.SageSLA Based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` +1. Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. +2. SageSLA is based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` +3. LTX-2 two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`. ### Image Generation Models -| Model Name | HuggingFace Model ID | Resolutions | -|:-----------------|:----------------------------------------|:---------------| -| FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | Any resolution | -| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Any resolution | -| FLUX.2-Klein | `black-forest-labs/FLUX.2-klein-4B` | Any resolution | -| Z-Image-Turbo | `Tongyi-MAI/Z-Image-Turbo` | Any resolution | -| GLM-Image | `zai-org/GLM-Image` | Any resolution | -| Qwen Image | `Qwen/Qwen-Image` | Any resolution | -| Qwen Image 2512 | `Qwen/Qwen-Image-2512` | Any resolution | -| Qwen Image Edit | `Qwen/Qwen-Image-Edit` | Any resolution | +| Model Name | HuggingFace Model ID | +|:---------------------|:------------------------------------| +| FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | +| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | +| FLUX.2-Klein | `black-forest-labs/FLUX.2-klein-4B` | +| Z-Image-Turbo | `Tongyi-MAI/Z-Image-Turbo` | +| GLM-Image | `zai-org/GLM-Image` | +| Qwen Image | `Qwen/Qwen-Image` | +| Qwen Image 2512 | `Qwen/Qwen-Image-2512` | +| Qwen Image Edit | `Qwen/Qwen-Image-Edit` | +| Qwen Image Edit 2511 | `Qwen/Qwen-Image-Edit-2511` | ## Verified LoRA Examples diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index 5523eb9a5..be8981b13 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -27,12 +27,12 @@ def _is_overlay_diffusion_model(model_path: str) -> bool: def _is_registered_diffusion_model(model_path: str) -> bool: try: - # if diffusion dependencies are not installed - from sglang.multimodal_gen.registry import get_model_info + from sglang.multimodal_gen.registry import has_registered_diffusion_model_path except ImportError: + # if diffusion dependencies are not installed return False - return get_model_info(model_path, backend="sglang") is not None + return has_registered_diffusion_model_path(model_path) def _is_diffusers_model_dir(model_dir: str) -> bool: @@ -93,9 +93,7 @@ def get_is_diffusion_model(model_path: str) -> bool: return _is_diffusers_model_dir(os.path.dirname(file_path)) except Exception as e: logger.debug("Failed to auto-detect diffusion model for %s: %s", model_path, e) - # For gated repos, file download fails but model card is still accessible. - # Check library_name from HF metadata as a fallback. - return _is_gated_diffusion_repo(model_path) + return False def get_model_path(extra_argv): diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 8dceb1934..471b91a1a 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -9,7 +9,7 @@ SGLang diffusion features an end-to-end unified pipeline for accelerating diffus ## Key Features SGLang Diffusion has the following features: - - Broad model support: Wan series, FastWan series, Hunyuan, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image + - Broad model support: Wan series, FastWan series, Hunyuan, LTX-2, Qwen-Image, Qwen-Image-Edit, Flux, Z-Image, GLM-Image - Fast inference speed: enpowered by highly optimized kernel from sgl-kernel and efficient scheduler loop - Ease of use: OpenAI-compatible api, CLI, and python sdk support - Multi-platform support: @@ -76,6 +76,11 @@ sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --save-output ``` +For LTX-2 two-stage generation, use `--pipeline-class-name LTX2TwoStagePipeline`. The +spatial upsampler and distilled LoRA are auto-resolved from the same model snapshot by +default, and can still be overridden with `--spatial-upsampler-path` and +`--distilled-lora-path` when needed. + ### LoRA support Apply LoRA adapters via `--lora-path`: diff --git a/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py b/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py index dfdb940eb..54a2b5a0e 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py @@ -8,6 +8,7 @@ from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAECon @dataclass class LTXAudioVAEArchConfig(VAEArchConfig): # Architecture params + temporal_compression_ratio: int = 4 causality_axis: str = "height" attn_resolutions: Optional[Tuple[int, ...]] = None base_channels: int = 128 @@ -20,6 +21,7 @@ class LTXAudioVAEArchConfig(VAEArchConfig): mid_block_add_attention: bool = False sample_rate: int = 16000 mel_hop_length: int = 160 + mel_compression_ratio: int = 4 is_causal: bool = True mel_bins: Optional[int] = 64 double_z: bool = True diff --git a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py index 3757cce63..02964d48a 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py @@ -11,6 +11,8 @@ class LTXVideoVAEArchConfig(VAEArchConfig): in_channels: int = 3 latent_channels: int = 128 out_channels: int = 3 + temporal_compression_ratio: int = 8 + spatial_compression_ratio: int = 32 block_out_channels: List[int] = field( default_factory=lambda: [256, 512, 1024, 2048] ) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py index d15506899..936a14ac9 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/base.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/base.py @@ -169,6 +169,7 @@ class PipelineConfig: # controls the timestep embedding generation should_use_guidance: bool = True embedded_cfg_scale: float = 6.0 + generator_device: str | None = None flow_shift: float | None = None disable_autocast: bool = False @@ -420,6 +421,9 @@ class PipelineConfig: def prepare_neg_cond_kwargs(self, batch, device, rotary_emb, dtype): return {} + def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_vae): + raise NotImplementedError("not yet implemented") + @staticmethod def add_cli_args( parser: FlexibleArgumentParser, prefix: str = "" diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 54fafce0d..827633c3d 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -1,6 +1,6 @@ import dataclasses from dataclasses import field -from typing import Callable +from typing import Callable, Optional import torch @@ -11,6 +11,7 @@ from sglang.multimodal_gen.configs.models.encoders import ( ) from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig +from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig from sglang.multimodal_gen.configs.pipeline_configs.base import ( ModelTaskType, PipelineConfig, @@ -93,8 +94,11 @@ def pack_text_embeds( def _gemma_postprocess_func( - outputs: BaseEncoderOutput, text_inputs: dict + outputs: BaseEncoderOutput, + text_inputs: dict, + pipeline_config: Optional["LTX2PipelineConfig"] = None, ) -> torch.Tensor: + _ = pipeline_config # LTX-2 requires all hidden states concatenated for the connector if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: # outputs.hidden_states is a tuple of tensors @@ -116,6 +120,7 @@ class LTX2PipelineConfig(PipelineConfig): task_type: ModelTaskType = ModelTaskType.TI2V skip_input_image_preprocess: bool = True + generator_device: str = "cpu" dit_config: LTX2Config = field(default_factory=LTX2Config) # Model architecture @@ -125,35 +130,23 @@ class LTX2PipelineConfig(PipelineConfig): patch_size_t: int = 1 # Audio VAE configuration + vae_config: LTXVideoVAEConfig = field(default_factory=LTXVideoVAEConfig) audio_vae_config: LTXAudioVAEConfig = field(default_factory=LTXAudioVAEConfig) audio_vae_precision: str = "fp32" - audio_vae_temporal_compression_ratio: int = 4 - audio_vae_mel_compression_ratio: int = 4 @property def vae_scale_factor(self): - return getattr(self.vae_config.arch_config, "spatial_compression_ratio", 32) + return self.vae_config.arch_config.spatial_compression_ratio @property def vae_temporal_compression(self): - return getattr(self.vae_config.arch_config, "temporal_compression_ratio", 8) + return self.vae_config.arch_config.temporal_compression_ratio def prepare_latent_shape(self, batch, batch_size, num_frames): - """Return packed latent shape [B, seq, C] directly.""" + """Return unpacked latent shape [B, C, F, H, W].""" height = batch.height // self.vae_scale_factor width = batch.width // self.vae_scale_factor - - post_patch_num_frames = num_frames // self.patch_size_t - post_patch_height = height // self.patch_size - post_patch_width = width // self.patch_size - seq_len = post_patch_num_frames * post_patch_height * post_patch_width - - num_channels = ( - self.in_channels * self.patch_size_t * self.patch_size * self.patch_size - ) - - shape = (batch_size, seq_len, num_channels) - return shape + return (batch_size, self.in_channels, num_frames, height, width) def prepare_audio_latent_shape(self, batch, batch_size, num_frames): # Adapted from diffusers pipeline prepare_audio_latents @@ -161,7 +154,9 @@ class LTX2PipelineConfig(PipelineConfig): sample_rate = self.audio_vae_config.arch_config.sample_rate hop_length = self.audio_vae_config.arch_config.mel_hop_length - temporal_compression = self.audio_vae_temporal_compression_ratio + temporal_compression = ( + self.audio_vae_config.arch_config.temporal_compression_ratio + ) latents_per_second = ( float(sample_rate) / float(hop_length) / float(temporal_compression) @@ -169,15 +164,13 @@ class LTX2PipelineConfig(PipelineConfig): latent_length = round(duration_s * latents_per_second) num_mel_bins = self.audio_vae_config.arch_config.mel_bins - mel_compression_ratio = self.audio_vae_mel_compression_ratio + mel_compression_ratio = self.audio_vae_config.arch_config.mel_compression_ratio latent_mel_bins = num_mel_bins // mel_compression_ratio # Default to 8 num_channels_latents = self.audio_vae_config.arch_config.latent_channels - shape = (batch_size, latent_length, num_channels_latents * latent_mel_bins) - - return shape + return (batch_size, num_channels_latents, latent_length, latent_mel_bins) # Text encoding stage (Gemma) # LTX-2 needs separate contexts for video/audio streams. We model this as @@ -221,6 +214,7 @@ class LTX2PipelineConfig(PipelineConfig): padding="max_length", max_length=max_sequence_length, truncation=True, + add_special_tokens=True, return_tensors="pt", ) return text_inputs @@ -524,7 +518,9 @@ class LTX2PipelineConfig(PipelineConfig): sample_rate = self.audio_vae_config.arch_config.sample_rate hop_length = self.audio_vae_config.arch_config.mel_hop_length - temporal_compression = self.audio_vae_temporal_compression_ratio + temporal_compression = ( + self.audio_vae_config.arch_config.temporal_compression_ratio + ) duration_s = num_frames / batch.fps latents_per_second = ( @@ -533,43 +529,9 @@ class LTX2PipelineConfig(PipelineConfig): audio_num_frames = round(duration_s * latents_per_second) num_mel_bins = self.audio_vae_config.arch_config.mel_bins - mel_compression_ratio = self.audio_vae_mel_compression_ratio + mel_compression_ratio = self.audio_vae_config.arch_config.mel_compression_ratio latent_mel_bins = num_mel_bins // mel_compression_ratio - audio_latents_mean = getattr(audio_vae, "latents_mean", None) - audio_latents_std = getattr(audio_vae, "latents_std", None) - if ( - isinstance(audio_latents_mean, torch.Tensor) - and isinstance(audio_latents_std, torch.Tensor) - and audio_latents_mean.numel() == audio_latents_std.numel() - ): - audio_latents_mean = audio_latents_mean.to( - device=audio_latents.device, dtype=audio_latents.dtype - ) - audio_latents_std = audio_latents_std.to( - device=audio_latents.device, dtype=audio_latents.dtype - ) - if audio_latents.ndim == 3: - if audio_latents.shape[-1] != audio_latents_mean.numel(): - raise ValueError( - f"audio_latents last dim {audio_latents.shape[-1]} " - f"does not match audio_vae stats {audio_latents_mean.numel()}" - ) - audio_latents = audio_latents * audio_latents_std.view( - 1, 1, -1 - ) + audio_latents_mean.view(1, 1, -1) - elif audio_latents.ndim == 2: - if audio_latents.shape[-1] != audio_latents_mean.numel(): - raise ValueError( - f"audio_latents last dim {audio_latents.shape[-1]} " - f"does not match audio_vae stats {audio_latents_mean.numel()}" - ) - audio_latents = audio_latents * audio_latents_std.view( - 1, -1 - ) + audio_latents_mean.view(1, -1) - else: - audio_latents = audio_latents * audio_latents_std + audio_latents_mean - audio_latents = self._unpack_audio_latents( audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins ) diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index 5d2e92b58..ef0b35981 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -10,6 +10,7 @@ class LTX2SamplingParams(SamplingParams): # Match the reference defaults used by ltx-pipelines (one-stage). # See: LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py seed: int = 10 + generator_device: str = "cpu" # Video parameters height: int = 512 diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index f4d353874..617e2b5e2 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -128,7 +128,7 @@ class SamplingParams: # Batch info num_outputs_per_prompt: int = 1 seed: int = 42 - generator_device: str = "cuda" # Device for random generator: "cuda" or "cpu" + generator_device: str | None = None # None means use the pipeline/model default # Original dimensions (before VAE scaling) num_frames: int = 1 # Default for image models @@ -685,7 +685,7 @@ class SamplingParams: "--generator-device", type=str, choices=["cuda", "musa", "cpu"], - help="Device for random generator (cuda, musa or cpu). Default: cuda", + help="Device for random generator (cuda, musa or cpu). Default: use the model-specific setting.", ) add_argument( "--num-frames", diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index b9e03c3d9..ca5db43ac 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -264,6 +264,29 @@ def _normalize_hf_cache_path(path: str) -> str: return os.path.normpath(path).lower().replace("\\", "/") +def has_registered_diffusion_model_path(model_path: str) -> bool: + all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True) + + if model_path in _MODEL_HF_PATH_TO_NAME: + return True + + model_short_name = get_model_short_name(model_path.lower()) + for registered_model_hf_id in all_model_hf_paths: + registered_model_name = get_model_short_name(registered_model_hf_id.lower()) + if registered_model_name in model_short_name: + return True + + normalized_model_path = _normalize_hf_cache_path(model_path) + for registered_model_hf_id in all_model_hf_paths: + cache_repo_fragment = ( + f"models--{registered_model_hf_id.lower().replace('/', '--')}" + ) + if cache_repo_fragment in normalized_model_path: + return True + + return False + + @lru_cache(maxsize=1) def _get_config_info( model_path: str, model_id: Optional[str] = None @@ -570,6 +593,9 @@ def _register_configs(): register_configs( sampling_param_cls=LTX2SamplingParams, pipeline_config_cls=LTX2PipelineConfig, + hf_model_paths=[ + "Lightricks/LTX-2", + ], model_detectors=[ lambda path: "ltx" in path.lower() and "video" in path.lower(), lambda path: "ltx-2" in path.lower(), diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 21b6e1a1f..5450cc409 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -270,7 +270,6 @@ def _maybe_mux_audio_into_mp4( sample_rate=selected_sr, ffmpeg_exe=ffmpeg_exe, ) - logger.info(f"Merged video saved to {CYAN}{save_file_path}{RESET}") except Exception as e: logger.warning( "Failed to mux audio into mp4 (saved silent video): %s", diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index b60dfd9dc..fac06a8d6 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py @@ -266,6 +266,7 @@ class LocalAttention(nn.Module): q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: """ Apply local attention between query, key and value tensors. @@ -284,6 +285,35 @@ class LocalAttention(nn.Module): forward_context: ForwardContext = get_forward_context() ctx_attn_metadata = forward_context.attn_metadata + if attn_mask is not None: + q_ = q.transpose(1, 2) + k_ = k.transpose(1, 2) + v_ = v.transpose(1, 2) + + if torch.is_floating_point(attn_mask): + mask = attn_mask.to(dtype=q_.dtype, device=q_.device) + if mask.dim() == 2: + mask = mask[:, None, None, :] + elif mask.dim() == 3: + mask = mask[:, None, :, :] + else: + mask = attn_mask.to(dtype=q_.dtype, device=q_.device) + if mask.dim() == 2: + mask = mask[:, None, None, :] + elif mask.dim() == 3: + mask = mask[:, None, :, :] + mask = (mask - 1.0) * torch.finfo(q_.dtype).max + + return torch.nn.functional.scaled_dot_product_attention( + q_, + k_, + v_, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + scale=self.softmax_scale, + ).transpose(1, 2) + output = self.attn_impl.forward(q, k, v, attn_metadata=ctx_attn_metadata) return output diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 6b426ac64..ed318ee69 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -84,18 +84,22 @@ class BaseLayerWithLoRA(nn.Module): # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: - lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) - lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) - delta = x @ lora_A_sliced.T @ lora_B_sliced.T + lora_dtype = lora_A.dtype + x_lora = x.to(dtype=lora_dtype) + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(device=x.device, non_blocking=True) + ) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(device=x.device, non_blocking=True) + ) + delta = x_lora @ lora_A_sliced.T @ lora_B_sliced.T if self.lora_alpha != self.lora_rank: delta = delta * ( self.lora_alpha / self.lora_rank # type: ignore ) # type: ignore delta = delta * self.strength - if delta.dim() > 2: - delta = delta.reshape(-1, delta.shape[-1]) out, output_bias = self.base_layer(x) - return out + delta, output_bias + return out + delta.to(dtype=out.dtype), output_bias else: out, output_bias = self.base_layer(x) return out, output_bias @@ -113,6 +117,7 @@ class BaseLayerWithLoRA(nn.Module): lora_path: str | None = None, strength: float = 1.0, clear_existing: bool = False, + merge_weights: bool = True, ) -> None: """ Set LoRA weights. Supports multiple LoRA adapters. @@ -149,7 +154,10 @@ class BaseLayerWithLoRA(nn.Module): self.strength = strength self.disable_lora = False - self.merge_lora_weights() + if merge_weights: + self.merge_lora_weights() + elif self.merged: + self.unmerge_lora_weights() @torch.no_grad() def _merge_lora_into_data( @@ -309,11 +317,34 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): super().__init__(base_layer, lora_rank, lora_alpha) def forward(self, input_: torch.Tensor) -> torch.Tensor: - # duplicate the logic in ColumnParallelLinear + lora_A = self.lora_A + lora_B = self.lora_B + if isinstance(self.lora_B, DTensor): + lora_B = self.lora_B.to_local() + lora_A = self.lora_A.to_local() + bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_, bias ) + if not self.merged and not self.disable_lora: + lora_dtype = lora_A.dtype + input_lora = input_.to(dtype=lora_dtype) + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(device=input_.device, non_blocking=True) + ) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(device=input_.device, non_blocking=True) + ) + delta_parallel = input_lora @ lora_A_sliced.T @ lora_B_sliced.T + if self.lora_alpha != self.lora_rank: + delta_parallel = delta_parallel * ( + self.lora_alpha / self.lora_rank # type: ignore + ) # type: ignore + delta_parallel = delta_parallel * self.strength + output_parallel = output_parallel + delta_parallel.to( + dtype=output_parallel.dtype + ) if self.base_layer.gather_output: output = tensor_model_parallel_all_gather(output_parallel) else: @@ -399,7 +430,12 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): super().__init__(base_layer, lora_rank, lora_alpha) def forward(self, input_: torch.Tensor): - # duplicate the logic in RowParallelLinear + lora_A = self.lora_A + lora_B = self.lora_B + if isinstance(self.lora_B, DTensor): + lora_B = self.lora_B.to_local() + lora_A = self.lora_A.to_local() + if self.base_layer.input_is_parallel: input_parallel = input_ else: @@ -411,6 +447,24 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA): output_parallel = self.base_layer.quant_method.apply( self.base_layer, input_parallel ) + if not self.merged and not self.disable_lora: + lora_dtype = lora_A.dtype + input_parallel_lora = input_parallel.to(dtype=lora_dtype) + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(device=input_parallel.device, non_blocking=True) + ) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(device=input_parallel.device, non_blocking=True) + ) + delta_parallel = input_parallel_lora @ lora_A_sliced.T @ lora_B_sliced.T + if self.lora_alpha != self.lora_rank: + delta_parallel = delta_parallel * ( + self.lora_alpha / self.lora_rank # type: ignore + ) # type: ignore + delta_parallel = delta_parallel * self.strength + output_parallel = output_parallel + delta_parallel.to( + dtype=output_parallel.dtype + ) if self.base_layer.reduce_results and self.base_layer.tp_size > 1: output_ = tensor_model_parallel_all_reduce(output_parallel) @@ -466,19 +520,23 @@ class LinearWithLoRA(BaseLayerWithLoRA): # TODO: Support multiple LoRA adapters when use not merged mode if not self.merged and not self.disable_lora: - lora_A_sliced = self.slice_lora_a_weights(lora_A.to(x, non_blocking=True)) - lora_B_sliced = self.slice_lora_b_weights(lora_B.to(x, non_blocking=True)) - delta = x @ lora_A_sliced.T @ lora_B_sliced.T + lora_dtype = lora_A.dtype + x_lora = x.to(dtype=lora_dtype) + lora_A_sliced = self.slice_lora_a_weights( + lora_A.to(device=x.device, non_blocking=True) + ) + lora_B_sliced = self.slice_lora_b_weights( + lora_B.to(device=x.device, non_blocking=True) + ) + delta = x_lora @ lora_A_sliced.T @ lora_B_sliced.T if self.lora_alpha != self.lora_rank: delta = delta * ( self.lora_alpha / self.lora_rank # type: ignore ) # type: ignore delta = delta * self.strength - if delta.dim() > 2: - delta = delta.reshape(-1, delta.shape[-1]) # nn.Linear.forward() returns a single tensor, not a tuple out = self.base_layer(x) - return out + delta + return out + delta.to(dtype=out.dtype) else: # nn.Linear.forward() returns a single tensor out = self.base_layer(x) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py index b416e6bbd..47d88c5ed 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/component_loader.py @@ -15,6 +15,7 @@ from torch import nn from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer from sglang.multimodal_gen.configs.models import ModelConfig +from sglang.multimodal_gen.configs.pipeline_configs.flux import Flux2PipelineConfig from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.utils import ( _normalize_component_type, @@ -294,9 +295,15 @@ class TokenizerLoader(ComponentLoader): def load_customized( self, component_model_path: str, server_args: ServerArgs, component_name: str ) -> Any: + # Flux.2 aligns to the tokenizer defaults from the original baseline. + # TODO: abstract this + if isinstance(server_args.pipeline_config, Flux2PipelineConfig): + return AutoTokenizer.from_pretrained(component_model_path) + return AutoTokenizer.from_pretrained( component_model_path, - padding_size="right", + padding_side="right", + use_fast=True, ) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py new file mode 100644 index 000000000..6ad78e066 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/upsampler_loader.py @@ -0,0 +1,223 @@ +import glob +import json +import os +import re + +import safetensors +import torch +from safetensors.torch import load_file as safetensors_load_file + +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + ComponentLoader, +) +from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( + LatentUpsampler, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + +UPSAMPLER_CONSTRUCTOR_KEYS = { + "in_channels", + "mid_channels", + "num_blocks_per_stage", + "dims", + "spatial_upsample", + "temporal_upsample", + "spatial_scale", + "rational_resampler", +} + +_HF_BLOB_URL_RE = re.compile( + r"https?://huggingface\.co/([^/]+/[^/]+)/blob/([^/]+)/(.*)" +) +_HF_RESOLVE_URL_RE = re.compile( + r"https?://huggingface\.co/([^/]+/[^/]+)/resolve/([^/]+)/(.*)" +) + + +def _parse_hf_url(path: str): + m = _HF_BLOB_URL_RE.match(path) or _HF_RESOLVE_URL_RE.match(path) + if m: + return m.group(1), m.group(2), m.group(3) + return None + + +def _download_hf_file(repo_id: str, filename: str, revision: str = "main") -> str: + from huggingface_hub import hf_hub_download + + logger.info("Downloading %s from %s (revision=%s)", filename, repo_id, revision) + return hf_hub_download(repo_id=repo_id, filename=filename, revision=revision) + + +def _find_safetensors_file(path: str) -> str: + """Resolve path to a single safetensors file (local path, directory, HF URL, or HF repo id).""" + if os.path.isfile(path) and path.endswith(".safetensors"): + return path + + if os.path.isdir(path): + files = sorted(glob.glob(os.path.join(path, "*.safetensors"))) + if len(files) == 1: + return files[0] + elif len(files) > 1: + raise ValueError( + f"Found {len(files)} safetensors files in {path}, expected 1" + ) + + hf = _parse_hf_url(path) + if hf: + repo_id, revision, filename = hf + return _download_hf_file(repo_id, filename, revision) + + try: + maybe_downloaded = maybe_download_model(path) + if os.path.isdir(maybe_downloaded): + files = sorted(glob.glob(os.path.join(maybe_downloaded, "*.safetensors"))) + if len(files) == 1: + return files[0] + elif len(files) > 1: + raise ValueError( + f"Found {len(files)} safetensors files in {maybe_downloaded}, expected 1" + ) + except Exception: + pass + + raise FileNotFoundError( + f"No safetensors file found at {path}. " + "Provide a local .safetensors file, a directory containing one, " + "a HuggingFace URL (https://huggingface.co//blob/main/), " + "or a HuggingFace repo id." + ) + + +def _normalize_config(raw: dict) -> dict: + """Map diffusers / original-repo config fields to LatentUpsampler kwargs.""" + config = {k: v for k, v in raw.items() if k in UPSAMPLER_CONSTRUCTOR_KEYS} + + # diffusers uses rational_spatial_scale instead of rational_resampler + spatial_scale + if "rational_spatial_scale" in raw and "rational_resampler" not in config: + config["rational_resampler"] = True + config.setdefault("spatial_scale", raw["rational_spatial_scale"]) + + return config + + +def _infer_config_from_state_dict(state_dict: dict[str, torch.Tensor]) -> dict: + """Infer LatentUpsampler kwargs from weight shapes and key names. + + Works even when no config.json or safetensors metadata is available. + """ + config: dict = {} + + w = state_dict.get("initial_conv.weight") + if w is not None: + config["mid_channels"] = w.shape[0] + config["in_channels"] = w.shape[1] + config["dims"] = 3 if w.ndim == 5 else 2 + + num_blocks = sum( + 1 + for k in state_dict + if k.startswith("res_blocks.") and k.endswith(".conv1.weight") + ) + if num_blocks > 0: + config["num_blocks_per_stage"] = num_blocks + + # Detect upsampler type from key patterns + has_rational = any(k.startswith("upsampler.blur_down.") for k in state_dict) + if has_rational: + config["rational_resampler"] = True + config["spatial_upsample"] = True + config["temporal_upsample"] = False + config["spatial_scale"] = 2.0 + else: + up_w = state_dict.get("upsampler.0.weight") + if up_w is not None and up_w.ndim == 5: + ratio = up_w.shape[0] // up_w.shape[1] + if ratio == 8: + config["spatial_upsample"] = True + config["temporal_upsample"] = True + elif ratio == 2: + config["spatial_upsample"] = False + config["temporal_upsample"] = True + else: + config["spatial_upsample"] = True + config["temporal_upsample"] = False + else: + config["spatial_upsample"] = True + config["temporal_upsample"] = False + + return config + + +def _load_config( + safetensors_path: str, + original_path: str, + state_dict: dict[str, torch.Tensor], +) -> dict: + """Load upsampler config with fallback chain: + 1. safetensors metadata ("config" key) - original LTX-2 repo format + 2. sibling config.json - diffusers format + 3. config.json from HF (if original_path was a URL) + 4. infer from state dict shapes (always works) + """ + with safetensors.safe_open(safetensors_path, framework="pt") as f: + meta = f.metadata() + if meta and "config" in meta: + logger.info("Using config from safetensors metadata") + return _normalize_config(json.loads(meta["config"])) + + config_json_path = os.path.join(os.path.dirname(safetensors_path), "config.json") + if os.path.isfile(config_json_path): + with open(config_json_path) as fp: + logger.info("Using config from sibling config.json") + return _normalize_config(json.load(fp)) + + hf = _parse_hf_url(original_path) + if hf: + repo_id, revision, filename = hf + config_filename = os.path.dirname(filename) + "/config.json" + try: + local = _download_hf_file(repo_id, config_filename, revision) + with open(local) as fp: + logger.info("Using config from HF config.json") + return _normalize_config(json.load(fp)) + except Exception: + pass + + logger.info("No explicit config found, inferring from state dict") + return _infer_config_from_state_dict(state_dict) + + +class UpsamplerLoader(ComponentLoader): + component_names = ["spatial_upsampler"] + expected_library = "diffusers" + + def should_offload(self, server_args: ServerArgs, model_config=None): + return server_args.vae_cpu_offload + + def load_customized( + self, + component_model_path: str, + server_args: ServerArgs, + component_name: str, + ): + safetensors_path = _find_safetensors_file(component_model_path) + state_dict = safetensors_load_file(safetensors_path) + config = _load_config(safetensors_path, component_model_path, state_dict) + + logger.info("Loading LatentUpsampler with config: %s", config) + + should_offload = self.should_offload(server_args) + target_device = self.target_device(should_offload) + + with torch.device("meta"): + model = LatentUpsampler(**config) + + model.load_state_dict(state_dict, assign=True) + model = model.to(device=target_device, dtype=torch.bfloat16).eval() + + logger.info("Loaded LatentUpsampler to %s", target_device) + return model diff --git a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py index 1ec861673..7bc5cd4db 100644 --- a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py @@ -46,7 +46,7 @@ def apply_split_rotary_emb( r = last // 2 # (..., 2, r) - split_x = x.reshape(*x.shape[:-1], 2, r) + split_x = x.reshape(*x.shape[:-1], 2, r).float() first_x = split_x[..., :1, :] # (..., 1, r) second_x = split_x[..., 1:, :] # (..., 1, r) @@ -153,12 +153,6 @@ class LTX2Attention(torch.nn.Module): query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: - batch_size, sequence_length, _ = ( - hidden_states.shape - if encoder_hidden_states is None - else encoder_hidden_states.shape - ) - if encoder_hidden_states is None: encoder_hidden_states = hidden_states @@ -183,16 +177,26 @@ class LTX2Attention(torch.nn.Module): key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, ) - query = query.unflatten(2, (self.heads, -1)) - key = key.unflatten(2, (self.heads, -1)) - value = value.unflatten(2, (self.heads, -1)) + query = query.unflatten(2, (self.heads, -1)).transpose(1, 2) + key = key.unflatten(2, (self.heads, -1)).transpose(1, 2) + value = value.unflatten(2, (self.heads, -1)).transpose(1, 2) - hidden_states = self.attn( + if attention_mask is not None: + if attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + elif attention_mask.ndim == 3: + attention_mask = attention_mask[:, None, :, :] + attention_mask = attention_mask.to(dtype=query.dtype) + + hidden_states = F.scaled_dot_product_attention( query, key, value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, ) - hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) hidden_states = hidden_states.to(query.dtype) hidden_states = self.to_out[0](hidden_states) @@ -464,9 +468,7 @@ class LTX2ConnectorTransformer1d(nn.Module): attention_mask = torch.zeros_like(attention_mask) # 2. Calculate 1D RoPE positional embeddings - rotary_emb = self.rope( - batch_size, seq_len, device=hidden_states.device, dtype=hidden_states.dtype - ) + rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device) # 3. Run 1D transformer blocks for block in self.transformer_blocks: diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index f8c493578..eae637529 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -21,7 +21,7 @@ from sglang.multimodal_gen.runtime.distributed import ( from sglang.multimodal_gen.runtime.distributed.communication_op import ( tensor_model_parallel_all_reduce, ) -from sglang.multimodal_gen.runtime.layers.attention import USPAttention +from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention from sglang.multimodal_gen.runtime.layers.linear import ( ColumnParallelLinear, RowParallelLinear, @@ -66,7 +66,7 @@ def apply_split_rotary_emb( ) r = last // 2 - split_x = x.reshape(*x.shape[:-1], 2, r) + split_x = x.reshape(*x.shape[:-1], 2, r).float() first_x = split_x[..., :1, :] second_x = split_x[..., 1:, :] @@ -137,7 +137,6 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): self.causal_offset = int(causal_offset) self.modality = modality - self.coords_dtype = torch.bfloat16 if modality == "video" else torch.float32 if self.modality not in ["video", "audio"]: raise ValueError( f"Modality {modality} is not supported. Supported modalities are `video` and `audio`." @@ -244,7 +243,6 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): device = device or coords.device num_pos_dims = coords.shape[1] - coords = coords.to(self.coords_dtype) if coords.ndim == 4: coords_start, coords_end = coords.chunk(2, dim=-1) coords = (coords_start + coords_end) / 2.0 @@ -309,9 +307,7 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module): cos_freqs = torch.swapaxes(cos_freq, 1, 2) sin_freqs = torch.swapaxes(sin_freq, 1, 2) - # Cast to bf16 to match model weights dtype. coords_dtype controls - # intermediate coordinate precision (fp32 for audio) and differs. - return cos_freqs.to(torch.bfloat16), sin_freqs.to(torch.bfloat16) + return cos_freqs, sin_freqs def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor: @@ -450,6 +446,7 @@ class LTX2Attention(nn.Module): dim_head: int = 64, norm_eps: float = 1e-6, qk_norm: bool = True, + use_local_attention: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, @@ -463,6 +460,7 @@ class LTX2Attention(nn.Module): self.inner_dim = self.heads * self.dim_head self.norm_eps = float(norm_eps) self.qk_norm = bool(qk_norm) + self.use_local_attention = bool(use_local_attention) tp_size = get_tp_world_size() if tp_size <= 0: @@ -531,16 +529,27 @@ class LTX2Attention(nn.Module): nn.Identity(), ) - self.attn = USPAttention( - num_heads=self.local_heads, - head_size=self.dim_head, - num_kv_heads=self.local_heads, - dropout_rate=0, - softmax_scale=None, - causal=False, - supported_attention_backends=supported_attention_backends, - prefix=f"{prefix}.attn", - ) + if self.use_local_attention: + self.attn = LocalAttention( + num_heads=self.local_heads, + head_size=self.dim_head, + num_kv_heads=self.local_heads, + softmax_scale=None, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn", + ) + else: + self.attn = USPAttention( + num_heads=self.local_heads, + head_size=self.dim_head, + num_kv_heads=self.local_heads, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn", + ) def forward( self, @@ -549,66 +558,56 @@ class LTX2Attention(nn.Module): mask: torch.Tensor | None = None, pe: tuple[torch.Tensor, torch.Tensor] | None = None, k_pe: tuple[torch.Tensor, torch.Tensor] | None = None, + perturbation_mask: torch.Tensor | None = None, + all_perturbed: bool = False, ) -> torch.Tensor: - q, _ = self.to_q(x) context_ = x if context is None else context - k, _ = self.to_k(context_) v, _ = self.to_v(context_) + use_attention = not all_perturbed - if self.qk_norm: - assert self.q_norm is not None and self.k_norm is not None - q = self.q_norm(q) - k = self.k_norm(k) + if use_attention: + q, _ = self.to_q(x) + k, _ = self.to_k(context_) - if pe is not None: - cos, sin = pe - k_cos, k_sin = pe if k_pe is None else k_pe - tp_size = get_tp_world_size() - if tp_size > 1: - tp_rank = get_tp_rank() - cos, sin = self._slice_rope_for_tp( - cos, sin, tp_rank=tp_rank, tp_size=tp_size - ) - k_cos, k_sin = self._slice_rope_for_tp( - k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size - ) - if cos.dim() == 3: - q = apply_interleaved_rotary_emb(q, (cos, sin)) - k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) - else: - q = apply_split_rotary_emb(q, (cos, sin)) - k = apply_split_rotary_emb(k, (k_cos, k_sin)) + if self.qk_norm: + assert self.q_norm is not None and self.k_norm is not None + q = self.q_norm(q) + k = self.k_norm(k) + + if pe is not None: + cos, sin = pe + k_cos, k_sin = pe if k_pe is None else k_pe + tp_size = get_tp_world_size() + if tp_size > 1: + tp_rank = get_tp_rank() + cos, sin = self._slice_rope_for_tp( + cos, sin, tp_rank=tp_rank, tp_size=tp_size + ) + k_cos, k_sin = self._slice_rope_for_tp( + k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size + ) + if cos.dim() == 3: + q = apply_interleaved_rotary_emb(q, (cos, sin)) + k = apply_interleaved_rotary_emb(k, (k_cos, k_sin)) + else: + q = apply_split_rotary_emb(q, (cos, sin)) + k = apply_split_rotary_emb(k, (k_cos, k_sin)) - q = q.view(*q.shape[:-1], self.local_heads, self.dim_head) - k = k.view(*k.shape[:-1], self.local_heads, self.dim_head) v = v.view(*v.shape[:-1], self.local_heads, self.dim_head) + if use_attention: + q = q.view(*q.shape[:-1], self.local_heads, self.dim_head) + k = k.view(*k.shape[:-1], self.local_heads, self.dim_head) - if mask is not None: - # Fallback to SDPA for masked attention - q_ = q.transpose(1, 2) - k_ = k.transpose(1, 2) - v_ = v.transpose(1, 2) - - if torch.is_floating_point(mask): - m = mask - if m.dim() == 2: - m = m[:, None, None, :] - elif m.dim() == 3: - m = m[:, None, :, :] - sdpa_mask = m.to(dtype=q_.dtype, device=q_.device) + if self.use_local_attention: + out = self.attn(q, k, v, attn_mask=mask) else: - m = mask.to(dtype=q_.dtype, device=q_.device) - if m.dim() == 2: - m = m[:, None, None, :] - elif m.dim() == 3: - m = m[:, None, :, :] - sdpa_mask = (m - 1.0) * torch.finfo(q_.dtype).max + out = self.attn(q, k, v) - out = torch.nn.functional.scaled_dot_product_attention( - q_, k_, v_, attn_mask=sdpa_mask, dropout_p=0.0, is_causal=False - ).transpose(1, 2) - else: - out = self.attn(q, k, v) + if perturbation_mask is not None: + out = out * perturbation_mask + v * (1 - perturbation_mask) + + if not use_attention: + out = v out = out.flatten(2) out, _ = self.to_out[0](out) @@ -720,6 +719,8 @@ class LTX2TransformerBlock(nn.Module): ) # 2. Prompt Cross-Attention + # Prompt KV is replicated across SP ranks, so prompt cross-attn should + # stay local and preserve the explicit KV mask semantics from official. self.attn2 = LTX2Attention( query_dim=dim, context_dim=cross_attention_dim, @@ -727,6 +728,7 @@ class LTX2TransformerBlock(nn.Module): dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, + use_local_attention=True, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn2", quant_config=quant_config, @@ -738,6 +740,7 @@ class LTX2TransformerBlock(nn.Module): dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, + use_local_attention=True, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn2", quant_config=quant_config, @@ -822,6 +825,10 @@ class LTX2TransformerBlock(nn.Module): audio_encoder_attention_mask: Optional[torch.Tensor] = None, a2v_cross_attention_mask: Optional[torch.Tensor] = None, v2a_cross_attention_mask: Optional[torch.Tensor] = None, + skip_video_self_attn: bool = False, + skip_audio_self_attn: bool = False, + skip_a2v_cross_attn: bool = False, + skip_v2a_cross_attn: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.size(0) @@ -833,7 +840,11 @@ class LTX2TransformerBlock(nn.Module): norm_hidden_states = ( rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa ) - attn_hidden_states = self.attn1(norm_hidden_states, pe=video_rotary_emb) + attn_hidden_states = self.attn1( + norm_hidden_states, + pe=video_rotary_emb, + all_perturbed=skip_video_self_attn, + ) hidden_states = hidden_states + attn_hidden_states * vgate_msa ashift_msa, ascale_msa, agate_msa = self.get_ada_values( @@ -843,10 +854,11 @@ class LTX2TransformerBlock(nn.Module): rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + ashift_msa ) attn_audio_hidden_states = self.audio_attn1( - norm_audio_hidden_states, pe=audio_rotary_emb + norm_audio_hidden_states, + pe=audio_rotary_emb, + all_perturbed=skip_audio_self_attn, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * agate_msa - # 2. Prompt Cross-Attention norm_hidden_states = rms_norm(hidden_states, self.norm_eps) attn_hidden_states = self.attn2( @@ -863,7 +875,6 @@ class LTX2TransformerBlock(nn.Module): mask=audio_encoder_attention_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states - # 3. Audio-to-Video and Video-to-Audio Cross-Attention norm_hidden_states = rms_norm(hidden_states, self.norm_eps) norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) @@ -934,14 +945,15 @@ class LTX2TransformerBlock(nn.Module): norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift ) - a2v_attn_hidden_states = self.audio_to_video_attn( - mod_norm_hidden_states, - context=mod_norm_audio_hidden_states, - pe=ca_video_rotary_emb, - k_pe=ca_audio_rotary_emb, - mask=a2v_cross_attention_mask, - ) - hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states + if not skip_a2v_cross_attn: + a2v_attn_hidden_states = self.audio_to_video_attn( + mod_norm_hidden_states, + context=mod_norm_audio_hidden_states, + pe=ca_video_rotary_emb, + k_pe=ca_audio_rotary_emb, + mask=a2v_cross_attention_mask, + ) + hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states # V2A mod_norm_hidden_states = ( @@ -951,15 +963,17 @@ class LTX2TransformerBlock(nn.Module): norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift ) - v2a_attn_hidden_states = self.video_to_audio_attn( - mod_norm_audio_hidden_states, - context=mod_norm_hidden_states, - pe=ca_audio_rotary_emb, - k_pe=ca_video_rotary_emb, - mask=v2a_cross_attention_mask, - ) - audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states - + if not skip_v2a_cross_attn: + v2a_attn_hidden_states = self.video_to_audio_attn( + mod_norm_audio_hidden_states, + context=mod_norm_hidden_states, + pe=ca_audio_rotary_emb, + k_pe=ca_video_rotary_emb, + mask=v2a_cross_attention_mask, + ) + audio_hidden_states = ( + audio_hidden_states + v2a_gate * v2a_attn_hidden_states + ) # 4. Feedforward vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( self.scale_shift_table, batch_size, temb, slice(3, None) @@ -978,7 +992,6 @@ class LTX2TransformerBlock(nn.Module): ) audio_ff_output = self.audio_ff(norm_audio_hidden_states) audio_hidden_states = audio_hidden_states + audio_ff_output * agate_mlp - return hidden_states, audio_hidden_states @@ -1275,6 +1288,10 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_num_frames: Optional[int] = None, video_coords: Optional[torch.Tensor] = None, audio_coords: Optional[torch.Tensor] = None, + skip_video_self_attn_blocks: Optional[tuple[int, ...]] = None, + skip_audio_self_attn_blocks: Optional[tuple[int, ...]] = None, + disable_a2v_cross_attn: bool = False, + disable_v2a_cross_attn: bool = False, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: @@ -1333,7 +1350,6 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): # 2. Patchify input projections hidden_states, _ = self.patchify_proj(hidden_states) audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states) - # 3. Prepare timestep embeddings # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters temb, embedded_timestep = self.adaln_single( @@ -1391,8 +1407,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_encoder_hidden_states = self.audio_caption_projection( audio_encoder_hidden_states ) - # 5. Run blocks + skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) + skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) for block in self.transformer_blocks: hidden_states, audio_hidden_states = block( hidden_states, @@ -1414,6 +1431,10 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): ca_audio_rotary_emb=ca_audio_rotary_emb, encoder_attention_mask=encoder_attention_mask, audio_encoder_attention_mask=audio_encoder_attention_mask, + skip_video_self_attn=block.idx in skip_video_self_attn_blocks, + skip_audio_self_attn=block.idx in skip_audio_self_attn_blocks, + skip_a2v_cross_attn=disable_a2v_cross_attn, + skip_v2a_cross_attn=disable_v2a_cross_attn, ) # 6. Output layers @@ -1439,7 +1460,6 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_hidden_states = self.audio_norm_out(audio_hidden_states) audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states) - # Unpatchify if requested (default True for pipeline compatibility) return_latents = kwargs.get("return_latents", True) diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py index 0927645fb..2db60639e 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py @@ -146,23 +146,61 @@ class Gemma3Attention(nn.Module): prefix=f"{prefix}.o_proj", ) + self.layer_type = ( + config.text_config.layer_types[layer_id] + if hasattr(config.text_config, "layer_types") + else None + ) self.is_sliding = ( config.text_config.layer_types[layer_id] == "sliding_attention" ) + rope_parameters = getattr(config.text_config, "rope_parameters", None) or {} + layer_rope_params = {} + if self.layer_type is not None and isinstance(rope_parameters, dict): + layer_rope_params = dict(rope_parameters.get(self.layer_type) or {}) + # Initialize the rotary embedding. if self.is_sliding: # Local attention. - self.rope_theta = config.text_config.rope_local_base_freq - rope_scaling = None # Default + self.rope_theta = float( + layer_rope_params.get( + "rope_theta", + getattr( + config.text_config, + "rope_local_base_freq", + getattr( + getattr(config.text_config, "default_theta", {}), + "get", + lambda *_: 10_000.0, + )("local", 10_000.0), + ), + ) + ) + rope_scaling = layer_rope_params or None # sliding window self.sliding_window = get_attention_sliding_window_size(config.text_config) # (left, right) = (window, 0) effectively for causal self.window_size = (self.sliding_window, 0) else: # Global attention. - self.rope_theta = config.text_config.rope_theta - rope_scaling = config.text_config.rope_scaling + self.rope_theta = float( + layer_rope_params.get( + "rope_theta", + getattr( + config.text_config, + "rope_theta", + getattr( + getattr(config.text_config, "default_theta", {}), + "get", + lambda *_: 1_000_000.0, + )("global", 1_000_000.0), + ), + ) + ) + rope_scaling = layer_rope_params or getattr( + config.text_config, "rope_scaling", None + ) self.sliding_window = None self.window_size = (-1, -1) @@ -734,7 +772,9 @@ class Gemma3TextModel(nn.Module): layer_id=i, config=config, quant_config=self.quant_config, - prefix=f"{config.text_config.prefix}.layers.{i}", + prefix=add_prefix( + f"layers.{i}", getattr(config.text_config, "prefix", "") + ), ) for i in range(config.text_config.num_hidden_layers) ] diff --git a/python/sglang/multimodal_gen/runtime/models/upsampler/__init__.py b/python/sglang/multimodal_gen/runtime/models/upsampler/__init__.py new file mode 100644 index 000000000..fe8edc84a --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/upsampler/__init__.py @@ -0,0 +1,5 @@ +from sglang.multimodal_gen.runtime.models.upsampler.latent_upsampler import ( + LatentUpsampler, +) + +__all__ = ["LatentUpsampler"] diff --git a/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py b/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py new file mode 100644 index 000000000..c62ec85f3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/upsampler/latent_upsampler.py @@ -0,0 +1,268 @@ +# Ported from https://github.com/Lightricks/LTX-2 +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from einops import rearrange + + +class BlurDownsample(torch.nn.Module): + """Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.""" + + def __init__(self, dims: int, stride: int, kernel_size: int = 5) -> None: + super().__init__() + assert dims in (2, 3) + assert isinstance(stride, int) and stride >= 1 + assert kernel_size >= 3 and kernel_size % 2 == 1 + self.dims = dims + self.stride = stride + self.kernel_size = kernel_size + + k = torch.tensor([math.comb(kernel_size - 1, i) for i in range(kernel_size)]) + k2d = k[:, None] @ k[None, :] + k2d = (k2d / k2d.sum()).float() + self.register_buffer("kernel", k2d[None, None, :, :]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.stride == 1: + return x + if self.dims == 2: + return self._apply_2d(x) + b, _, f, _, _ = x.shape + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self._apply_2d(x) + h2, w2 = x.shape[-2:] + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2) + return x + + def _apply_2d(self, x2d: torch.Tensor) -> torch.Tensor: + c = x2d.shape[1] + weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size) + x2d = F.conv2d( + x2d, + weight=weight, + bias=None, + stride=self.stride, + padding=self.kernel_size // 2, + groups=c, + ) + return x2d + + +class PixelShuffleND(torch.nn.Module): + """N-dimensional pixel shuffle for upsampling tensors.""" + + def __init__(self, dims: int, upscale_factors: Tuple[int, int, int] = (2, 2, 2)): + super().__init__() + assert dims in [1, 2, 3] + self.dims = dims + self.upscale_factors = upscale_factors + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.dims == 3: + return rearrange( + x, + "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)", + p1=self.upscale_factors[0], + p2=self.upscale_factors[1], + p3=self.upscale_factors[2], + ) + elif self.dims == 2: + return rearrange( + x, + "b (c p1 p2) h w -> b c (h p1) (w p2)", + p1=self.upscale_factors[0], + p2=self.upscale_factors[1], + ) + elif self.dims == 1: + return rearrange( + x, + "b (c p1) f h w -> b c (f p1) h w", + p1=self.upscale_factors[0], + ) + else: + raise ValueError(f"Unsupported dims: {self.dims}") + + +class ResBlock(torch.nn.Module): + """Residual block with two conv layers, group norm, and SiLU activation.""" + + def __init__( + self, channels: int, mid_channels: Optional[int] = None, dims: int = 3 + ): + super().__init__() + if mid_channels is None: + mid_channels = channels + conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1) + self.norm1 = torch.nn.GroupNorm(32, mid_channels) + self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1) + self.norm2 = torch.nn.GroupNorm(32, channels) + self.activation = torch.nn.SiLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.conv1(x) + x = self.norm1(x) + x = self.activation(x) + x = self.conv2(x) + x = self.norm2(x) + x = self.activation(x + residual) + return x + + +def _rational_for_scale(scale: float) -> Tuple[int, int]: + mapping = {0.75: (3, 4), 1.5: (3, 2), 2.0: (2, 1), 4.0: (4, 1)} + if float(scale) not in mapping: + raise ValueError( + f"Unsupported scale {scale}. Choose from {list(mapping.keys())}" + ) + return mapping[float(scale)] + + +class SpatialRationalResampler(torch.nn.Module): + """Fully-learned rational spatial scaling via PixelShuffle + anti-aliased downsample.""" + + def __init__(self, mid_channels: int, scale: float): + super().__init__() + self.scale = float(scale) + self.num, self.den = _rational_for_scale(self.scale) + self.conv = torch.nn.Conv2d( + mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1 + ) + self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num)) + self.blur_down = BlurDownsample(dims=2, stride=self.den) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + b, _, f, _, _ = x.shape + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self.conv(x) + x = self.pixel_shuffle(x) + x = self.blur_down(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + return x + + +class LatentUpsampler(torch.nn.Module): + """ + Upsample VAE latents spatially and/or temporally. + + Args: + in_channels: Number of channels in the input latent. + mid_channels: Number of channels in the middle layers. + num_blocks_per_stage: Number of ResBlocks per stage (pre/post upsampling). + dims: Dimensionality of convolutions (2 or 3). + spatial_upsample: Whether to spatially upsample. + temporal_upsample: Whether to temporally upsample. + spatial_scale: Scale factor for spatial upsampling. + rational_resampler: Whether to use rational resampler for spatial upsampling. + """ + + def __init__( + self, + in_channels: int = 128, + mid_channels: int = 512, + num_blocks_per_stage: int = 4, + dims: int = 3, + spatial_upsample: bool = True, + temporal_upsample: bool = False, + spatial_scale: float = 2.0, + rational_resampler: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.mid_channels = mid_channels + self.num_blocks_per_stage = num_blocks_per_stage + self.dims = dims + self.spatial_upsample = spatial_upsample + self.temporal_upsample = temporal_upsample + self.spatial_scale = float(spatial_scale) + self.rational_resampler = rational_resampler + + conv = torch.nn.Conv2d if dims == 2 else torch.nn.Conv3d + + self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1) + self.initial_norm = torch.nn.GroupNorm(32, mid_channels) + self.initial_activation = torch.nn.SiLU() + + self.res_blocks = torch.nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + if spatial_upsample and temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d( + mid_channels, 8 * mid_channels, kernel_size=3, padding=1 + ), + PixelShuffleND(3), + ) + elif spatial_upsample: + if rational_resampler: + self.upsampler = SpatialRationalResampler( + mid_channels=mid_channels, scale=self.spatial_scale + ) + else: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv2d( + mid_channels, 4 * mid_channels, kernel_size=3, padding=1 + ), + PixelShuffleND(2), + ) + elif temporal_upsample: + self.upsampler = torch.nn.Sequential( + torch.nn.Conv3d( + mid_channels, 2 * mid_channels, kernel_size=3, padding=1 + ), + PixelShuffleND(1), + ) + else: + raise ValueError( + "Either spatial_upsample or temporal_upsample must be True" + ) + + self.post_upsample_res_blocks = torch.nn.ModuleList( + [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)] + ) + + self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + b, _, f, _, _ = latent.shape + + if self.dims == 2: + x = rearrange(latent, "b c f h w -> (b f) c h w") + x = self.initial_conv(x) + x = self.initial_norm(x) + x = self.initial_activation(x) + for block in self.res_blocks: + x = block(x) + x = self.upsampler(x) + for block in self.post_upsample_res_blocks: + x = block(x) + x = self.final_conv(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + else: + x = self.initial_conv(latent) + x = self.initial_norm(x) + x = self.initial_activation(x) + for block in self.res_blocks: + x = block(x) + + if self.temporal_upsample: + x = self.upsampler(x) + x = x[:, :, 1:, :, :] + elif isinstance(self.upsampler, SpatialRationalResampler): + x = self.upsampler(x) + else: + x = rearrange(x, "b c f h w -> (b f) c h w") + x = self.upsampler(x) + x = rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f) + + for block in self.post_upsample_res_blocks: + x = block(x) + x = self.final_conv(x) + + return x diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py index 0341771a0..bf1613e5a 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py @@ -843,8 +843,8 @@ class AutoencoderKLLTX2Audio(ParallelTiledVAE): # Per-channel statistics for normalizing and denormalizing the latent representation. This statistics is computed over # the entire dataset and stored in model's checkpoint under AudioVAE state_dict - latents_std = torch.zeros((base_channels,)) - latents_mean = torch.ones((base_channels,)) + latents_std = torch.ones((base_channels,)) + latents_mean = torch.zeros((base_channels,)) self.register_buffer("latents_mean", latents_mean, persistent=True) self.register_buffer("latents_std", latents_std, persistent=True) diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py index 77c4e512b..53176ed98 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py @@ -1135,6 +1135,28 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE): config.arch_config.decoder_spatio_temporal_scaling ) decoder_layers_per_block = config.arch_config.decoder_layers_per_block + decoder_inject_noise = getattr( + config.arch_config, "decoder_inject_noise", (False, False, False, False) + ) + if isinstance(decoder_inject_noise, bool): + decoder_inject_noise = (decoder_inject_noise,) * 4 + else: + decoder_inject_noise = tuple(decoder_inject_noise) + upsample_residual = getattr( + config.arch_config, "upsample_residual", (True, True, True) + ) + if isinstance(upsample_residual, bool): + upsample_residual = (upsample_residual,) * 3 + else: + upsample_residual = tuple(upsample_residual) + upsample_factor = getattr(config.arch_config, "upsample_factor", (2, 2, 2)) + if isinstance(upsample_factor, int): + upsample_factor = (upsample_factor,) * 3 + else: + upsample_factor = tuple(upsample_factor) + timestep_conditioning = getattr( + config.arch_config, "timestep_conditioning", False + ) decoder_causal = config.arch_config.decoder_causal decoder_spatial_padding_mode = config.arch_config.decoder_spatial_padding_mode @@ -1154,16 +1176,20 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE): ) self.decoder = LTX2VideoDecoder3d( - latent_channels, - out_channels, - decoder_block_out_channels, - decoder_spatio_temporal_scaling, - decoder_layers_per_block, - patch_size, - patch_size_t, - resnet_norm_eps, - decoder_causal, - decoder_spatial_padding_mode, + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=decoder_block_out_channels, + spatio_temporal_scaling=decoder_spatio_temporal_scaling, + layers_per_block=decoder_layers_per_block, + patch_size=patch_size, + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=decoder_causal, + inject_noise=decoder_inject_noise, + timestep_conditioning=timestep_conditioning, + upsample_residual=upsample_residual, + upsample_factor=upsample_factor, + spatial_padding_mode=decoder_spatial_padding_mode, ) latents_mean = torch.zeros((latent_channels,), requires_grad=False) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index d8798ecf2..311a18c79 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -1,5 +1,3 @@ -import inspect -import json import math import os @@ -7,103 +5,183 @@ import numpy as np import torch from diffusers import FlowMatchEulerDiscreteScheduler +from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( + PipelineComponentLoader, +) 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.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, LTX2AVDecodingStage, LTX2AVDenoisingStage, LTX2AVLatentPreparationStage, + LTX2HalveResolutionStage, + LTX2LoRASwitchStage, + LTX2RefinementStage, LTX2TextConnectorStage, + LTX2UpsampleStage, TextEncodingStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) - -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 +BASE_SHIFT_ANCHOR = 1024 +MAX_SHIFT_ANCHOR = 4096 -def prepare_mu(batch: Req, server_args: ServerArgs): - height = batch.height - width = batch.width - num_frames = batch.num_frames +def _resolve_ltx2_two_stage_component_paths( + model_path: str, component_paths: dict[str, str] +) -> dict[str, str]: + resolved = dict(component_paths) + auto_resolved = [] - vae_arch = getattr( - getattr(server_args.pipeline_config, "vae_config", None), "arch_config", None + if "spatial_upsampler" not in resolved: + spatial_candidates = [ + os.path.join(model_path, "latent_upsampler"), + os.path.join(model_path, "ltx-2-spatial-upscaler-x2-1.0.safetensors"), + ] + for candidate in spatial_candidates: + if os.path.exists(candidate): + resolved["spatial_upsampler"] = candidate + auto_resolved.append(f"spatial_upsampler={candidate}") + break + + if "distilled_lora" not in resolved: + distilled_lora = os.path.join( + model_path, "ltx-2-19b-distilled-lora-384.safetensors" + ) + if os.path.exists(distilled_lora): + resolved["distilled_lora"] = distilled_lora + auto_resolved.append(f"distilled_lora={distilled_lora}") + + if auto_resolved: + logger.info( + "Auto-resolved LTX2 two-stage components: %s", ", ".join(auto_resolved) + ) + + return resolved + + +def calculate_ltx2_shift( + image_seq_len: int, + base_seq_len: int = BASE_SHIFT_ANCHOR, + max_seq_len: int = MAX_SHIFT_ANCHOR, + base_shift: float = 0.95, + max_shift: float = 2.05, +) -> float: + mm = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - mm * base_seq_len + return image_seq_len * mm + b + + +def prepare_ltx2_mu(batch: Req, server_args: ServerArgs): + latent_num_frames = (int(batch.num_frames) - 1) // int( + server_args.pipeline_config.vae_temporal_compression + ) + 1 + latent_height = int(batch.height) // int( + server_args.pipeline_config.vae_scale_factor ) - vae_scale_factor = ( - getattr(vae_arch, "spatial_compression_ratio", None) - or getattr(vae_arch, "vae_scale_factor", None) - or getattr(server_args.pipeline_config, "vae_scale_factor", None) + latent_width = int(batch.width) // int(server_args.pipeline_config.vae_scale_factor) + video_sequence_length = latent_num_frames * latent_height * latent_width + return "mu", calculate_ltx2_shift(video_sequence_length) + + +class LTX2SigmaPreparationStage(PipelineStage): + """Prepare native LTX-2 sigma schedule before timestep setup.""" + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + batch.extra["ltx2_phase"] = "stage1" + batch.sigmas = np.linspace( + 1.0, + 1.0 / int(batch.num_inference_steps), + int(batch.num_inference_steps), + ).tolist() + return batch + + +def _add_ltx2_front_stages(pipeline: ComposedPipelineBase): + pipeline.add_stages( + [ + InputValidationStage(), + TextEncodingStage( + text_encoders=[pipeline.get_module("text_encoder")], + tokenizers=[pipeline.get_module("tokenizer")], + ), + LTX2TextConnectorStage(connectors=pipeline.get_module("connectors")), + ] ) - vae_temporal_compression = getattr( - vae_arch, "temporal_compression_ratio", None - ) or getattr(server_args.pipeline_config, "vae_temporal_compression", None) - # Values from LTX2Pipeline in diffusers - mu = calculate_shift( - 4096, - base_seq_len=1024, - max_seq_len=4096, - base_shift=0.95, - max_shift=2.05, + +def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase): + pipeline.add_stage(LTX2SigmaPreparationStage()) + pipeline.add_standard_timestep_preparation_stage( + prepare_extra_kwargs=[prepare_ltx2_mu] + ) + pipeline.add_stages( + [ + LTX2AVLatentPreparationStage( + scheduler=pipeline.get_module("scheduler"), + transformer=pipeline.get_module("transformer"), + audio_vae=pipeline.get_module("audio_vae"), + ), + LTX2AVDenoisingStage( + transformer=pipeline.get_module("transformer"), + scheduler=pipeline.get_module("scheduler"), + vae=pipeline.get_module("vae"), + audio_vae=pipeline.get_module("audio_vae"), + pipeline=pipeline, + ), + ] ) - return "mu", mu -def _load_component_config(model_path: str, component_name: str): - """Helper to load component config from model_index.json or config.json""" - try: - # Try loading model_index.json first - index_path = os.path.join(model_path, "model_index.json") - if os.path.exists(index_path): - with open(index_path, "r") as f: - index = json.load(f) - - if component_name in index: - # It's a subfolder - subfolder = index[component_name][1] - config_path = os.path.join(model_path, subfolder, "config.json") - if os.path.exists(config_path): - with open(config_path, "r") as f: - return json.load(f) - - # Fallback to direct config.json in subfolder if standard structure - config_path = os.path.join(model_path, component_name, "config.json") - if os.path.exists(config_path): - with open(config_path, "r") as f: - return json.load(f) - - except Exception as e: - logger.warning(f"Failed to load config for {component_name}: {e}") - - return {} - - -def _filter_kwargs_for_cls(cls, kwargs): - """Filter kwargs to only include those accepted by cls.__init__""" - sig = inspect.signature(cls.__init__) - return {k: v for k, v in kwargs.items() if k in sig.parameters} +def _add_ltx2_decoding_stage(pipeline: ComposedPipelineBase): + pipeline.add_stage( + LTX2AVDecodingStage( + vae=pipeline.get_module("vae"), + audio_vae=pipeline.get_module("audio_vae"), + vocoder=pipeline.get_module("vocoder"), + pipeline=pipeline, + ) + ) class LTX2FlowMatchScheduler(FlowMatchEulerDiscreteScheduler): """Override ``_time_shift_exponential`` to use torch f32 instead of numpy f64.""" + def set_timesteps( + self, + num_inference_steps=None, + device=None, + sigmas=None, + mu=None, + timesteps=None, + ): + if sigmas is not None and timesteps is None and mu is None: + sigmas = torch.tensor(sigmas, dtype=torch.float32, device=device) + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + self.num_inference_steps = len(timesteps) + self.timesteps = timesteps + self.sigmas = sigmas + self._step_index = None + self._begin_index = None + return + + return super().set_timesteps( + num_inference_steps=num_inference_steps, + device=device, + sigmas=sigmas, + mu=mu, + timesteps=timesteps, + ) + def _time_shift_exponential(self, mu, sigma, t): if isinstance(t, np.ndarray): t_torch = torch.from_numpy(t).to(torch.float32) @@ -112,10 +190,7 @@ class LTX2FlowMatchScheduler(FlowMatchEulerDiscreteScheduler): return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) -class LTX2Pipeline(ComposedPipelineBase): - # NOTE: must match `model_index.json`'s `_class_name` for native dispatch. - pipeline_name = "LTX2Pipeline" - +class _BaseLTX2Pipeline(LoRAPipeline): _required_config_modules = [ "transformer", "text_encoder", @@ -131,41 +206,127 @@ class LTX2Pipeline(ComposedPipelineBase): orig = self.get_module("scheduler") self.modules["scheduler"] = LTX2FlowMatchScheduler.from_config(orig.config) + +class LTX2Pipeline(_BaseLTX2Pipeline): + # Must match model_index.json `_class_name`. + pipeline_name = "LTX2Pipeline" + def create_pipeline_stages(self, server_args: ServerArgs): - self.add_stages( - [ - InputValidationStage(), - TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - ), - LTX2TextConnectorStage(connectors=self.get_module("connectors")), - ] + _add_ltx2_front_stages(self) + _add_ltx2_stage1_generation_stages(self) + _add_ltx2_decoding_stage(self) + + +class LTX2TwoStagePipeline(_BaseLTX2Pipeline): + pipeline_name = "LTX2TwoStagePipeline" + STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0] + + def initialize_pipeline(self, server_args: ServerArgs): + super().initialize_pipeline(server_args) + server_args.component_paths = _resolve_ltx2_two_stage_component_paths( + self.model_path, server_args.component_paths ) - self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) + upsampler_path = server_args.component_paths.get("spatial_upsampler") + if not upsampler_path: + raise ValueError( + "LTX2TwoStagePipeline requires --spatial-upsampler-path " + "(component_paths['spatial_upsampler'])." + ) + module, memory_usage = PipelineComponentLoader.load_component( + component_name="spatial_upsampler", + component_model_path=upsampler_path, + transformers_or_diffusers="diffusers", + server_args=server_args, + ) + self.modules["spatial_upsampler"] = module + self.memory_usages["spatial_upsampler"] = memory_usage + distilled_lora_path = server_args.component_paths.get("distilled_lora") + if not distilled_lora_path: + raise ValueError( + "LTX2TwoStagePipeline requires --distilled-lora-path " + "(component_paths['distilled_lora'])." + ) + self._distilled_lora_path = distilled_lora_path + self._stage1_lora_path = server_args.lora_path + self._stage1_lora_scale = float(server_args.lora_scale) + self._active_lora_phase = None + + def switch_lora_phase(self, phase: str) -> None: + if phase == self._active_lora_phase: + return + + if phase == "stage1": + if self._stage1_lora_path: + self.set_lora( + lora_nickname="ltx2_stage1_base", + lora_path=self._stage1_lora_path, + target="transformer", + strength=self._stage1_lora_scale, + ) + else: + # Stage 1 must run on the base transformer weights. If stage 2 left the + # distilled adapter active, stage 1 quality drifts away from the official + # two-stage pipeline immediately. + self.deactivate_lora_weights(target="transformer") + elif phase == "stage2": + lora_nicknames = [] + lora_paths = [] + lora_strengths = [] + lora_targets = [] + if self._stage1_lora_path: + lora_nicknames.append("ltx2_stage1_base") + lora_paths.append(self._stage1_lora_path) + lora_strengths.append(self._stage1_lora_scale) + lora_targets.append("transformer") + lora_nicknames.append("ltx2_stage2_distilled") + lora_paths.append(self._distilled_lora_path) + lora_strengths.append(1.0) + lora_targets.append("transformer") + self.set_lora( + lora_nickname=lora_nicknames, + lora_path=lora_paths, + target=lora_targets, + strength=lora_strengths, + # Keep the distilled adapter unmerged when it is the only active LoRA. + # Merging it into the base weights makes the subsequent switch back to + # stage 1 depend on unmerge bookkeeping instead of the original base. + merge_weights=self._stage1_lora_path is not None, + ) + else: + raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}") + + self._active_lora_phase = phase + + def create_pipeline_stages(self, server_args: ServerArgs): + _add_ltx2_front_stages(self) + self.add_stage(LTX2HalveResolutionStage()) + self.add_stage( + LTX2LoRASwitchStage(pipeline=self, phase="stage1"), + ) + _add_ltx2_stage1_generation_stages(self) self.add_stages( [ - LTX2AVLatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"), - audio_vae=self.get_module("audio_vae"), - ), - LTX2AVDenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"), + LTX2UpsampleStage( + spatial_upsampler=self.get_module("spatial_upsampler"), vae=self.get_module("vae"), audio_vae=self.get_module("audio_vae"), ), - LTX2AVDecodingStage( + ( + LTX2LoRASwitchStage(pipeline=self, phase="stage2"), + "ltx2_lora_switch_stage2", + ), + LTX2RefinementStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + distilled_sigmas=self.STAGE_2_DISTILLED_SIGMA_VALUES, vae=self.get_module("vae"), audio_vae=self.get_module("audio_vae"), - vocoder=self.get_module("vocoder"), - pipeline=self, ), ] ) + _add_ltx2_decoding_stage(self) -EntryClass = LTX2Pipeline +EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 45ee048cd..2c2e7fe9f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py @@ -8,7 +8,6 @@ This module defines the base class for pipelines that are composed of multiple s """ import os -import re from abc import ABC, abstractmethod from typing import Any, Callable, Literal, cast @@ -333,12 +332,7 @@ class ComposedPipelineBase(ABC): @staticmethod def _infer_stage_name(stage: PipelineStage) -> str: - class_name = stage.__class__.__name__ - # snake_case - name = re.sub(r"(? int: """ Apply LoRA weights to the given lora_layers. Supports multiple LoRA adapters. @@ -435,38 +436,29 @@ class LoRAPipeline(ComposedPipelineBase): lora_A_name in self.lora_adapters[nickname] and lora_B_name in self.lora_adapters[nickname] ): - # Some LoRA checkpoints (e.g. Lightning distill) store per-layer alpha as ".alpha". - # If present, we must apply the standard LoRA scaling: scale = alpha / rank. - try: - inferred_rank = int( - self.lora_adapters[nickname][lora_A_name].shape[0] - ) - except Exception: - inferred_rank = None - # Default to None for some checkpoints without ".alpha" - inferred_alpha: int | None = None + inferred_rank = int( + self.lora_adapters[nickname][lora_A_name].shape[0] + ) alpha_key = name + ".alpha" if alpha_key in self.lora_adapters[nickname]: - try: - inferred_alpha = int( - self.lora_adapters[nickname][alpha_key].item() - ) - except Exception: - inferred_alpha = None - - if inferred_rank is not None: - layer.lora_rank = inferred_rank - layer.lora_alpha = ( - inferred_alpha - if inferred_alpha is not None - else inferred_rank + inferred_alpha = int( + self.lora_adapters[nickname][alpha_key].item() ) + else: + # Some distilled LoRAs omit per-layer alpha and rely on the + # default LoRA scale of alpha == rank. Falling back to rank + # keeps the effective delta consistent with the official path. + inferred_alpha = inferred_rank + + layer.lora_rank = inferred_rank + layer.lora_alpha = inferred_alpha layer.set_lora_weights( self.lora_adapters[nickname][lora_A_name], self.lora_adapters[nickname][lora_B_name], lora_path=path, strength=lora_strength, + merge_weights=merge_weights, clear_existing=( clear_existing and idx == 0 ), # Only clear on first LoRA @@ -589,6 +581,7 @@ class LoRAPipeline(ComposedPipelineBase): lora_path: str | None | list[str | None] = None, target: str | list[str] = "all", strength: float | list[float] = 1.0, + merge_weights: bool = True, ): # type: ignore """ Load LoRA adapter(s) into the pipeline and apply them to the specified transformer(s). @@ -682,6 +675,7 @@ class LoRAPipeline(ComposedPipelineBase): rank, tgt_strengths, clear_existing=True, + merge_weights=merge_weights, ) adapted_count += count self.cur_adapter_name[module_name] = merged_name @@ -689,7 +683,7 @@ class LoRAPipeline(ComposedPipelineBase): str(p or self.loaded_adapter_paths.get(n, "")) for n, p in zip(tgt_nicknames, tgt_paths) ) - self.is_lora_merged[module_name] = True + self.is_lora_merged[module_name] = merge_weights self.cur_adapter_strength[module_name] = tgt_strengths[0] # Store full configuration for multi-LoRA support (preserves order and all strengths) self.cur_adapter_config[module_name] = ( @@ -698,7 +692,7 @@ class LoRAPipeline(ComposedPipelineBase): ) logger.info( - "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s)", + "Rank %d: LoRA adapter(s) %s applied to %d layers (targets: %s, strengths: %s, merge_weights=%s)", rank, ", ".join(map(str, lora_paths)) if lora_paths else None, adapted_count, @@ -708,8 +702,42 @@ class LoRAPipeline(ComposedPipelineBase): if len(strengths) > 1 else f"{strengths[0]:.2f}" ), + merge_weights, ) + def deactivate_lora_weights(self, target: str = "all") -> None: + """ + Disable LoRA for the specified target, regardless of whether weights were + merged into the base model or are still active in the wrapped LoRA path. + """ + target_modules, error = self._get_target_lora_layers(target) + if error: + logger.warning("deactivate_lora_weights: %s", error) + if not target_modules: + return + + modules_requiring_unmerge = [] + for module_name, lora_layers_dict in target_modules: + if self.is_lora_merged.get(module_name, False) or any( + layer.merged for layer in lora_layers_dict.values() + ): + modules_requiring_unmerge.append((module_name, lora_layers_dict)) + + offload_context = self._temporarily_disable_offload( + target_modules=modules_requiring_unmerge + ) + with offload_context: + for module_name, lora_layers_dict in target_modules: + for layer in lora_layers_dict.values(): + if layer.merged: + layer.unmerge_lora_weights() + if not layer.disable_lora: + layer.disable_lora = True + self.is_lora_merged[module_name] = False + self.cur_adapter_strength.pop(module_name, None) + self.cur_adapter_config.pop(module_name, None) + logger.info("LoRA weights deactivated for %s", module_name) + def merge_lora_weights(self, target: str = "all", strength: float = 1.0) -> None: """ Merge LoRA weights into the base model for the specified target. diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index 6f8e14221..ba6fc5f46 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -22,6 +22,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import ( from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import ( LTX2AVDenoisingStage, + LTX2RefinementStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import ( DmdDenoisingStage, @@ -64,6 +65,11 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation import ( TimestepPreparationStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.upsampling import ( + LTX2HalveResolutionStage, + LTX2LoRASwitchStage, + LTX2UpsampleStage, +) __all__ = [ "PipelineStage", @@ -92,4 +98,9 @@ __all__ = [ "Hunyuan3DPaintPreprocessStage", "Hunyuan3DPaintTexGenStage", "Hunyuan3DPaintPostprocessStage", + # LTX-2 two-stage + "LTX2RefinementStage", + "LTX2HalveResolutionStage", + "LTX2LoRASwitchStage", + "LTX2UpsampleStage", ] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py index 28df749bd..e56622189 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py @@ -106,6 +106,23 @@ class LTX2AVDecodingStage(DecodingStage): logger.warning( "audio_vae.latents_std is all zeros; audio denorm may be incorrect." ) + try: + latents_mean = self.audio_vae.latents_mean + except AttributeError: + latents_mean = None + if isinstance(latents_mean, torch.Tensor) and isinstance( + latents_std, torch.Tensor + ): + latents_mean = latents_mean.to(device=device, dtype=dtype) + latents_std = latents_std.to(device=device, dtype=dtype) + if audio_latents.ndim == 4: + latents_mean = latents_mean.view( + 1, audio_latents.shape[1], 1, audio_latents.shape[3] + ) + latents_std = latents_std.view( + 1, audio_latents.shape[1], 1, audio_latents.shape[3] + ) + audio_latents = audio_latents * latents_std + latents_mean with torch.no_grad(): # Decode latents to spectrogram diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 2f348d68c..fcc4a429a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -710,6 +710,8 @@ class DenoisingStage(PipelineStage): trajectory_timesteps: list, server_args: ServerArgs, is_warmup: bool = False, + *args, + **kwargs, ): # Gather results if using sequence parallelism if trajectory_latents: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 7bd5a7e95..50118f6d0 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -9,6 +9,7 @@ import PIL.Image import torch from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.models.vision_utils import ( @@ -70,7 +71,11 @@ class LTX2AVDenoisingStage(DenoisingStage): return int(batch.sp_video_latent_num_frames) pc = server_args.pipeline_config - return int((batch.num_frames - 1) // int(pc.vae_temporal_compression) + 1) + return int( + (batch.num_frames - 1) + // int(pc.vae_config.arch_config.temporal_compression_ratio) + + 1 + ) @staticmethod def _truncate_sp_padded_token_latents( @@ -106,6 +111,57 @@ class LTX2AVDenoisingStage(DenoisingStage): return return super()._maybe_enable_cache_dit(num_inference_steps, batch) + def _get_ltx2_stage1_guider_params( + self, batch: Req, server_args: ServerArgs, stage: str + ) -> dict[str, object] | None: + if stage != "stage1": + return None + + pipeline_ref = getattr(self, "pipeline", None) + pipeline = pipeline_ref() if callable(pipeline_ref) else pipeline_ref + pipeline_name = getattr(pipeline, "pipeline_name", None) + if pipeline_name != "LTX2TwoStagePipeline": + return None + + return batch.extra.get("ltx2_stage1_guider_params") + + @staticmethod + def _ltx2_should_skip_step(step_index: int, skip_step: int) -> bool: + if skip_step == 0: + return False + return step_index % (skip_step + 1) != 0 + + @staticmethod + def _ltx2_apply_rescale( + cond: torch.Tensor, pred: torch.Tensor, rescale_scale: float + ) -> torch.Tensor: + if rescale_scale == 0.0: + return pred + factor = cond.std() / pred.std() + factor = rescale_scale * factor + (1.0 - rescale_scale) + return pred * factor + + @classmethod + def _ltx2_calculate_guided_x0( + cls, + *, + cond: torch.Tensor, + uncond_text: torch.Tensor | float, + uncond_perturbed: torch.Tensor | float, + uncond_modality: torch.Tensor | float, + cfg_scale: float, + stg_scale: float, + rescale_scale: float, + modality_scale: float, + ) -> torch.Tensor: + pred = ( + cond + + (cfg_scale - 1.0) * (cond - uncond_text) + + stg_scale * (cond - uncond_perturbed) + + (modality_scale - 1.0) * (cond - uncond_modality) + ) + return cls._ltx2_apply_rescale(cond, pred, rescale_scale) + @staticmethod def _resize_center_crop( img: PIL.Image.Image, *, width: int, height: int @@ -328,22 +384,16 @@ class LTX2AVDenoisingStage(DenoisingStage): # Prepare variables for the denoising loop prepared_vars = self._prepare_denoising_loop(batch, server_args) - extra_step_kwargs = prepared_vars["extra_step_kwargs"] target_dtype = prepared_vars["target_dtype"] autocast_enabled = prepared_vars["autocast_enabled"] timesteps = prepared_vars["timesteps"] num_inference_steps = prepared_vars["num_inference_steps"] num_warmup_steps = prepared_vars["num_warmup_steps"] - image_kwargs = prepared_vars["image_kwargs"] - pos_cond_kwargs = prepared_vars["pos_cond_kwargs"] - neg_cond_kwargs = prepared_vars["neg_cond_kwargs"] latents = prepared_vars["latents"] boundary_timestep = prepared_vars["boundary_timestep"] z = prepared_vars["z"] reserved_frames_mask = prepared_vars["reserved_frames_mask"] - seq_len = prepared_vars["seq_len"] - guidance = prepared_vars["guidance"] - + stage = batch.extra.get("ltx2_phase", "stage1") audio_latents = batch.audio_latents audio_scheduler = copy.deepcopy(self.scheduler) @@ -356,8 +406,14 @@ class LTX2AVDenoisingStage(DenoisingStage): latent_num_frames_for_model = self._get_video_latent_num_frames_for_model( batch=batch, server_args=server_args, latents=latents ) - latent_height = batch.height // server_args.pipeline_config.vae_scale_factor - latent_width = batch.width // server_args.pipeline_config.vae_scale_factor + latent_height = ( + batch.height + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) + latent_width = ( + batch.width + // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + ) # Initialize lists for ODE trajectory trajectory_timesteps: list[torch.Tensor] = [] @@ -394,7 +450,6 @@ class LTX2AVDenoisingStage(DenoisingStage): clean_latent[:, :num_img_tokens, :] = batch.image_latent[ :, :num_img_tokens, : ].to(device=latents.device, dtype=latents.dtype) - with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, @@ -438,7 +493,9 @@ class LTX2AVDenoisingStage(DenoisingStage): latent_model_input = latents.to(target_dtype) audio_latent_model_input = audio_latents.to(target_dtype) - + stage1_guider_params = self._get_ltx2_stage1_guider_params( + batch, server_args, stage + ) latent_num_frames = latent_num_frames_for_model # Audio latent dims @@ -468,56 +525,58 @@ class LTX2AVDenoisingStage(DenoisingStage): timestep_video = timestep timestep_audio = timestep - # Conditions - encoder_hidden_states = batch.prompt_embeds[0] - audio_encoder_hidden_states = batch.audio_prompt_embeds[0] - encoder_attention_mask = batch.prompt_attention_mask - - # Follow ltx-pipelines structure: separate pos/neg forward passes, - # then apply CFG on denoised (x0) predictions. - with set_forward_context( - current_timestep=i, attn_metadata=attn_metadata - ): - v_pos, a_v_pos = current_model( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - timestep=timestep_video, - audio_timestep=timestep_audio, - encoder_attention_mask=encoder_attention_mask, - audio_encoder_attention_mask=encoder_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=batch.fps, - audio_num_frames=audio_num_frames_latent, - video_coords=video_coords, - audio_coords=audio_coords, - return_latents=False, - return_dict=False, - ) - + use_official_cfg_path = stage1_guider_params is None + if use_official_cfg_path: + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = batch.prompt_attention_mask if batch.do_classifier_free_guidance: - neg_encoder_hidden_states = ( - batch.negative_prompt_embeds[0] + latent_model_input = torch.cat( + [latent_model_input] * 2, dim=0 ) - neg_audio_encoder_hidden_states = ( - batch.negative_audio_prompt_embeds[0] + audio_latent_model_input = torch.cat( + [audio_latent_model_input] * 2, dim=0 ) - neg_encoder_attention_mask = ( - batch.negative_attention_mask + encoder_hidden_states = torch.cat( + [ + batch.negative_prompt_embeds[0], + encoder_hidden_states, + ], + dim=0, + ) + audio_encoder_hidden_states = torch.cat( + [ + batch.negative_audio_prompt_embeds[0], + audio_encoder_hidden_states, + ], + dim=0, + ) + encoder_attention_mask = torch.cat( + [ + batch.negative_attention_mask, + encoder_attention_mask, + ], + dim=0, + ) + timestep_video = timestep_video.expand( + int(latent_model_input.shape[0]) + ) + timestep_audio = timestep_audio.expand( + int(latent_model_input.shape[0]) ) - v_neg, a_v_neg = current_model( + with set_forward_context( + current_timestep=i, attn_metadata=attn_metadata + ): + model_video, model_audio = current_model( hidden_states=latent_model_input, audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=neg_encoder_hidden_states, - audio_encoder_hidden_states=neg_audio_encoder_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, timestep=timestep_video, audio_timestep=timestep_audio, - encoder_attention_mask=neg_encoder_attention_mask, - audio_encoder_attention_mask=neg_encoder_attention_mask, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=encoder_attention_mask, num_frames=latent_num_frames, height=latent_height, width=latent_width, @@ -528,16 +587,132 @@ class LTX2AVDenoisingStage(DenoisingStage): return_latents=False, return_dict=False, ) - else: - v_neg = None - a_v_neg = None - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - if v_neg is not None: - v_neg = v_neg.float() - if a_v_neg is not None: - a_v_neg = a_v_neg.float() + model_video = model_video.float() + model_audio = model_audio.float() + if batch.do_classifier_free_guidance: + model_video_uncond, model_video_text = ( + model_video.chunk(2) + ) + model_audio_uncond, model_audio_text = ( + model_audio.chunk(2) + ) + model_video = model_video_uncond + ( + batch.guidance_scale + * (model_video_text - model_video_uncond) + ) + model_audio = model_audio_uncond + ( + batch.guidance_scale + * (model_audio_text - model_audio_uncond) + ) + v_pos = model_video + a_v_pos = model_audio + v_neg = None + a_v_neg = None + + latents = self.scheduler.step( + v_pos, t_device, latents, return_dict=False + )[0] + audio_latents = audio_scheduler.step( + a_v_pos, t_device, audio_latents, return_dict=False + )[0] + if do_ti2v: + latents[:, :num_img_tokens, :] = batch.image_latent[ + :, :num_img_tokens, : + ].to(device=latents.device, dtype=latents.dtype) + + latents = self.post_forward_for_ti2v_task( + batch, server_args, reserved_frames_mask, latents, z + ) + + if batch.return_trajectory_latents: + trajectory_timesteps.append(t_host) + trajectory_latents.append(latents) + if audio_latents is not None: + trajectory_audio_latents.append(audio_latents) + + if i == num_timesteps - 1 or ( + (i + 1) > num_warmup_steps + and (i + 1) % self.scheduler.order == 0 + and progress_bar is not None + ): + progress_bar.update() + + if not is_warmup: + self.step_profile() + continue + else: + # Follow ltx-pipelines structure: separate pos/neg forward passes, + # then apply CFG on denoised (x0) predictions. + encoder_hidden_states = batch.prompt_embeds[0] + audio_encoder_hidden_states = batch.audio_prompt_embeds[0] + encoder_attention_mask = batch.prompt_attention_mask + with set_forward_context( + current_timestep=i, attn_metadata=attn_metadata + ): + v_pos, a_v_pos = current_model( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep_video, + audio_timestep=timestep_audio, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=encoder_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=batch.fps, + audio_num_frames=audio_num_frames_latent, + video_coords=video_coords, + audio_coords=audio_coords, + return_latents=False, + return_dict=False, + ) + + if ( + stage1_guider_params is not None + or batch.do_classifier_free_guidance + ): + neg_encoder_hidden_states = ( + batch.negative_prompt_embeds[0] + ) + neg_audio_encoder_hidden_states = ( + batch.negative_audio_prompt_embeds[0] + ) + neg_encoder_attention_mask = ( + batch.negative_attention_mask + ) + + v_neg, a_v_neg = current_model( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=neg_encoder_hidden_states, + audio_encoder_hidden_states=neg_audio_encoder_hidden_states, + timestep=timestep_video, + audio_timestep=timestep_audio, + encoder_attention_mask=neg_encoder_attention_mask, + audio_encoder_attention_mask=neg_encoder_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=batch.fps, + audio_num_frames=audio_num_frames_latent, + video_coords=video_coords, + audio_coords=audio_coords, + return_latents=False, + return_dict=False, + ) + else: + v_neg = None + a_v_neg = None + + v_pos = v_pos.float() + a_v_pos = a_v_pos.float() + if v_neg is not None: + v_neg = v_neg.float() + if a_v_neg is not None: + a_v_neg = a_v_neg.float() # Velocity -> denoised (x0): x0 = x - sigma * v sigma_val = float(sigma.item()) @@ -547,9 +722,18 @@ class LTX2AVDenoisingStage(DenoisingStage): denoised_audio = ( audio_latents.float() - sigma_val * a_v_pos ).to(audio_latents.dtype) + denoised_video_neg = None + denoised_audio_neg = None + denoised_video_perturbed = None + denoised_audio_perturbed = None + denoised_video_modality = None + denoised_audio_modality = None if ( - batch.do_classifier_free_guidance + ( + stage1_guider_params is not None + or batch.do_classifier_free_guidance + ) and v_neg is not None and a_v_neg is not None ): @@ -559,6 +743,159 @@ class LTX2AVDenoisingStage(DenoisingStage): denoised_audio_neg = ( audio_latents.float() - sigma_val * a_v_neg ).to(audio_latents.dtype) + if stage1_guider_params is not None: + video_skip = self._ltx2_should_skip_step( + i, int(stage1_guider_params["video_skip_step"]) + ) + audio_skip = self._ltx2_should_skip_step( + i, int(stage1_guider_params["audio_skip_step"]) + ) + + need_perturbed = ( + float(stage1_guider_params["video_stg_scale"]) != 0.0 + or float(stage1_guider_params["audio_stg_scale"]) != 0.0 + ) + if need_perturbed: + with set_forward_context( + current_timestep=i, attn_metadata=attn_metadata + ): + v_ptb, a_v_ptb = current_model( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep_video, + audio_timestep=timestep_audio, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=encoder_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=batch.fps, + audio_num_frames=audio_num_frames_latent, + video_coords=video_coords, + audio_coords=audio_coords, + return_latents=False, + return_dict=False, + skip_video_self_attn_blocks=tuple( + stage1_guider_params["video_stg_blocks"] + ), + skip_audio_self_attn_blocks=tuple( + stage1_guider_params["audio_stg_blocks"] + ), + ) + denoised_video_perturbed = ( + latents.float() - sigma_val * v_ptb.float() + ).to(latents.dtype) + denoised_audio_perturbed = ( + audio_latents.float() - sigma_val * a_v_ptb.float() + ).to(audio_latents.dtype) + + need_modality = ( + float(stage1_guider_params["video_modality_scale"]) + != 1.0 + or float(stage1_guider_params["audio_modality_scale"]) + != 1.0 + ) + if need_modality: + with set_forward_context( + current_timestep=i, attn_metadata=attn_metadata + ): + v_mod, a_v_mod = current_model( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + timestep=timestep_video, + audio_timestep=timestep_audio, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=encoder_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=batch.fps, + audio_num_frames=audio_num_frames_latent, + video_coords=video_coords, + audio_coords=audio_coords, + return_latents=False, + return_dict=False, + disable_a2v_cross_attn=True, + disable_v2a_cross_attn=True, + ) + denoised_video_modality = ( + latents.float() - sigma_val * v_mod.float() + ).to(latents.dtype) + denoised_audio_modality = ( + audio_latents.float() - sigma_val * a_v_mod.float() + ).to(audio_latents.dtype) + + if not video_skip: + denoised_video = self._ltx2_calculate_guided_x0( + cond=denoised_video, + uncond_text=( + denoised_video_neg + if denoised_video_neg is not None + else denoised_video + ), + uncond_perturbed=( + denoised_video_perturbed + if denoised_video_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_video_modality + if denoised_video_modality is not None + else 0.0 + ), + cfg_scale=float( + stage1_guider_params["video_cfg_scale"] + ), + stg_scale=float( + stage1_guider_params["video_stg_scale"] + ), + rescale_scale=float( + stage1_guider_params["video_rescale_scale"] + ), + modality_scale=float( + stage1_guider_params["video_modality_scale"] + ), + ) + if not audio_skip: + denoised_audio = self._ltx2_calculate_guided_x0( + cond=denoised_audio, + uncond_text=( + denoised_audio_neg + if denoised_audio_neg is not None + else denoised_audio + ), + uncond_perturbed=( + denoised_audio_perturbed + if denoised_audio_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_audio_modality + if denoised_audio_modality is not None + else 0.0 + ), + cfg_scale=float( + stage1_guider_params["audio_cfg_scale"] + ), + stg_scale=float( + stage1_guider_params["audio_stg_scale"] + ), + rescale_scale=float( + stage1_guider_params["audio_rescale_scale"] + ), + modality_scale=float( + stage1_guider_params["audio_modality_scale"] + ), + ) + elif ( + batch.do_classifier_free_guidance + and denoised_video_neg is not None + and denoised_audio_neg is not None + ): denoised_video = denoised_video + ( batch.guidance_scale - 1.0 ) * (denoised_video - denoised_video_neg) @@ -576,7 +913,6 @@ class LTX2AVDenoisingStage(DenoisingStage): denoised_video * denoise_mask + clean_latent.float() * (1.0 - denoise_mask) ) - # Euler step in sigma space: x_next = x + (sigma_next - sigma) * v, # where v = (x - x0) / sigma. if sigma_val == 0.0: @@ -655,6 +991,8 @@ class LTX2AVDenoisingStage(DenoisingStage): trajectory_audio_latents: list, server_args: ServerArgs, is_warmup: bool = False, + *args, + **kwargs, ): # 1. Handle Trajectory (Video) - Copy from base if trajectory_latents: @@ -705,7 +1043,8 @@ class LTX2AVDenoisingStage(DenoisingStage): batch.latents = latents batch.audio_latents = audio_latents - + # 4. Cleanup + # TODO: make this a general denoising-stage hook if isinstance(self.transformer, OffloadableDiTMixin): for manager in self.transformer.layerwise_offload_managers: manager.release_all() @@ -753,9 +1092,6 @@ class LTX2AVDenoisingStage(DenoisingStage): ) return result - def do_classifier_free_guidance(self, batch: Req) -> bool: - return batch.guidance_scale > 1.0 - class LTX2RefinementStage(LTX2AVDenoisingStage): def __init__( @@ -764,34 +1100,109 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): super().__init__(transformer, scheduler, vae, audio_vae) self.distilled_sigmas = torch.tensor(distilled_sigmas) + @staticmethod + def _randn_like_with_batch_generators( + reference_tensor: torch.Tensor, batch: Req + ) -> torch.Tensor: + generator = getattr(batch, "generator", None) + if isinstance(generator, list): + bsz = int(reference_tensor.shape[0]) + valid_generators = [g for g in generator if isinstance(g, torch.Generator)] + if len(valid_generators) == 1: + generator = valid_generators[0] + elif len(valid_generators) >= bsz: + generator = valid_generators[:bsz] + else: + generator = None + elif not isinstance(generator, torch.Generator): + generator = None + + return randn_tensor( + reference_tensor.shape, + generator=generator, + device=reference_tensor.device, + dtype=reference_tensor.dtype, + ) + + @staticmethod + def _reset_stage2_generators(batch: Req) -> None: + generator = getattr(batch, "generator", None) + if isinstance(generator, list) and generator: + generator_device = str(generator[0].device) + elif isinstance(generator, torch.Generator): + generator_device = str(generator.device) + else: + generator_device = "cpu" + + seeds = getattr(batch, "seeds", None) + if not seeds: + seed = getattr(batch, "seed", None) + if seed is None: + return + seeds = [int(seed)] + + batch.generator = [ + torch.Generator(device=generator_device).manual_seed(int(seed)) + for seed in seeds + ] + def forward(self, batch: Req, server_args: ServerArgs) -> Req: - # 1. Add noise to latents + batch.extra["ltx2_phase"] = "stage2" + self._reset_stage2_generators(batch) noise_scale = self.distilled_sigmas[0].to(batch.latents.device) - noise = torch.randn_like(batch.latents) - batch.latents = batch.latents + noise * noise_scale + video_noise = self._randn_like_with_batch_generators(batch.latents, batch) + batch.latents = video_noise * noise_scale + batch.latents * (1 - noise_scale) - # 2. Run denoising loop with distilled_sigmas - # Save original sigmas - original_sigmas = self.scheduler.sigmas - original_timesteps = self.scheduler.timesteps - original_num_inference_steps = self.scheduler.num_inference_steps + if isinstance(batch.audio_latents, torch.Tensor): + audio_noise = self._randn_like_with_batch_generators( + batch.audio_latents, batch + ) + audio_noise_scale = noise_scale.to( + batch.audio_latents.device, batch.audio_latents.dtype + ) + batch.audio_latents = ( + audio_noise * audio_noise_scale + + batch.audio_latents * (1 - audio_noise_scale) + ) + batch.latents = batch.latents.to( + device=batch.latents.device, dtype=torch.float32 + ) + if isinstance(batch.audio_latents, torch.Tensor): + batch.audio_latents = batch.audio_latents.to( + device=batch.audio_latents.device, dtype=torch.float32 + ) - # Set distilled sigmas - self.scheduler.sigmas = self.distilled_sigmas.to(self.scheduler.sigmas.device) - # Approximation for timesteps - self.scheduler.timesteps = self.scheduler.sigmas * 1000 - self.scheduler.num_inference_steps = len(self.distilled_sigmas) - 1 + # Stage 2 runs at full resolution, so Stage 1 TI2V conditioning is invalid. + batch.image_latent = None + batch.ltx2_num_image_tokens = 0 + + # Use a private scheduler copy to avoid mutating shared state. + original_scheduler = self.scheduler + original_batch_timesteps = batch.timesteps + original_batch_num_inference_steps = batch.num_inference_steps + + self.scheduler = copy.deepcopy(original_scheduler) + distilled_device = self.scheduler.sigmas.device + self.scheduler.sigmas = self.distilled_sigmas.to(distilled_device) + num_steps = len(self.distilled_sigmas) - 1 + self.scheduler.num_inference_steps = num_steps + self.scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to( + distilled_device + ) + self.scheduler._step_index = None + self.scheduler._begin_index = None + + batch.timesteps = self.scheduler.timesteps + batch.num_inference_steps = num_steps + original_do_cfg = batch.do_classifier_free_guidance + batch.do_classifier_free_guidance = False - # Call parent forward try: batch = super().forward(batch, server_args) finally: - # Restore original sigmas - self.scheduler.sigmas = original_sigmas - self.scheduler.timesteps = original_timesteps - self.scheduler.num_inference_steps = original_num_inference_steps + self.scheduler = original_scheduler + batch.timesteps = original_batch_timesteps + batch.num_inference_steps = original_batch_num_inference_steps + batch.do_classifier_free_guidance = original_do_cfg return batch - - def do_classifier_free_guidance(self, batch: Req) -> bool: - return False # Stage 2 uses simple denoising (no CFG) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py index fc27f657f..99cc92d7d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py @@ -79,6 +79,11 @@ class InputValidationStage(PipelineStage): # Create generators based on generator_device parameter # Note: This will overwrite any existing batch.generator generator_device = batch.generator_device + if generator_device is None: + generator_device = ( + getattr(server_args.pipeline_config, "generator_device", None) + or current_platform.device_type + ) if generator_device == "cpu": device_str = "cpu" diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index fb9879a5f..5051b85b1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -37,6 +37,15 @@ class LatentPreparationStage(PipelineStage): self.scheduler = scheduler self.transformer = transformer + def _get_latent_dtype( + self, + batch: Req, + server_args: ServerArgs, + ): + return server_args.pipeline_config.get_latent_dtype( + batch.prompt_embeds[0].dtype + ) + def forward( self, batch: Req, @@ -57,9 +66,7 @@ class LatentPreparationStage(PipelineStage): batch_size = batch.batch_size # Get required parameters - dtype = server_args.pipeline_config.get_latent_dtype( - batch.prompt_embeds[0].dtype - ) + dtype = self._get_latent_dtype(batch, server_args) device = get_local_torch_device() generator = batch.generator latents = batch.latents diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py index 4440a3060..130953346 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py @@ -53,6 +53,13 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): result.add_check("latents", batch.latents, V.none_or_tensor) return result + def _get_latent_dtype( + self, + batch: Req, + server_args: ServerArgs, + ): + return torch.float32 + def forward(self, batch: Req, server_args: ServerArgs) -> Req: # 1. Prepare Video Latents using base class logic # This sets batch.latents and batch.raw_latent_shape @@ -70,12 +77,7 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): return batch device = get_local_torch_device() - if isinstance(batch.prompt_embeds, list) and batch.prompt_embeds: - dtype = batch.prompt_embeds[0].dtype - elif isinstance(batch.prompt_embeds, torch.Tensor): - dtype = batch.prompt_embeds.dtype - else: - dtype = torch.float16 + dtype = self._get_latent_dtype(batch, server_args) generator = batch.generator audio_latents = batch.audio_latents diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index e1184cbc0..61293b306 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -7,6 +7,8 @@ Prompt encoding stages for diffusion pipelines. This module contains implementations of prompt encoding stages for diffusion pipelines. """ +import inspect + import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput @@ -263,14 +265,22 @@ class TextEncodingStage(PipelineStage): attention_mask = torch.ones(input_ids.shape[:2], device=target_device) else: attention_mask = text_inputs["attention_mask"] + encoder_forward_kwargs = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "output_hidden_states": True, + } + if "use_cache" in inspect.signature(text_encoder.forward).parameters: + encoder_forward_kwargs["use_cache"] = False with set_forward_context(current_timestep=0, attn_metadata=None): - outputs: BaseEncoderOutput = text_encoder( - input_ids=input_ids, - attention_mask=attention_mask, - output_hidden_states=True, - use_cache=False, - ) - prompt_embeds = postprocess_func(outputs, text_inputs) + outputs: BaseEncoderOutput = text_encoder(**encoder_forward_kwargs) + postprocess_sig = inspect.signature(postprocess_func) + + postprocess_kwargs = {} + if "pipeline_config" in postprocess_sig.parameters: + # required by models like LTX + postprocess_kwargs["pipeline_config"] = server_args.pipeline_config + prompt_embeds = postprocess_func(outputs, text_inputs, **postprocess_kwargs) if dtype is not None: prompt_embeds = prompt_embeds.to(dtype=dtype) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py new file mode 100644 index 000000000..d4e6e0e56 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py @@ -0,0 +1,126 @@ +import torch + +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +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 LTX2HalveResolutionStage(PipelineStage): + """Halve batch height/width for two-stage Stage 1 (low-res generation).""" + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + original_h, original_w = batch.height, batch.width + + vae_scale_factor = getattr(server_args.pipeline_config, "vae_scale_factor", 32) + required_alignment = max(64, int(vae_scale_factor) * 2) + if original_h % required_alignment != 0 or original_w % required_alignment != 0: + raise ValueError( + "LTX-2 two-stage requires resolution divisible by " + f"{required_alignment}, got ({original_h}x{original_w})." + ) + + batch.height = batch.height // 2 + batch.width = batch.width // 2 + logger.info( + "Halved resolution: %dx%d -> %dx%d", + original_h, + original_w, + batch.height, + batch.width, + ) + return batch + + +class LTX2LoRASwitchStage(PipelineStage): + """Switch LoRA configuration for the requested two-stage phase.""" + + def __init__(self, pipeline, phase: str): + super().__init__() + self.pipeline = pipeline + self.phase = phase + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + switch_fn = getattr(self.pipeline, "switch_lora_phase", None) + if not callable(switch_fn): + raise ValueError( + "LTX2LoRASwitchStage requires pipeline.switch_lora_phase()" + ) + switch_fn(self.phase) + batch.extra["ltx2_phase"] = self.phase + return batch + + +class LTX2UpsampleStage(PipelineStage): + """Upsample Stage-1 video latents and prepare Stage-2 inputs.""" + + def __init__(self, spatial_upsampler, vae, audio_vae=None): + super().__init__() + self.spatial_upsampler = spatial_upsampler + self.vae = vae + self.audio_vae = audio_vae + + def _upsample_video_latents( + self, latents: torch.Tensor, server_args: ServerArgs, device: torch.device + ) -> torch.Tensor: + vae_mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to( + device=device, dtype=latents.dtype + ) + vae_std = self.vae.latents_std.view(1, -1, 1, 1, 1).to( + device=device, dtype=latents.dtype + ) + latents = latents * vae_std + vae_mean + self.spatial_upsampler = self.spatial_upsampler.to( + device=device, dtype=latents.dtype + ) + latents = self.spatial_upsampler(latents) + if server_args.vae_cpu_offload: + self.spatial_upsampler = self.spatial_upsampler.to("cpu") + latents = (latents - vae_mean) / vae_std + return latents + + @staticmethod + def _restore_full_resolution(batch: Req) -> None: + batch.height *= 2 + batch.width *= 2 + + @staticmethod + def _pack_video_latents( + batch: Req, latents: torch.Tensor, server_args: ServerArgs + ) -> None: + batch_size = latents.shape[0] + latents = server_args.pipeline_config.maybe_pack_latents( + latents, batch_size, batch + ) + batch.latents = latents + batch.raw_latent_shape = latents.shape + + def _repack_audio_latents(self, batch: Req, server_args: ServerArgs) -> None: + if batch.audio_latents is None or self.audio_vae is None: + return + audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( + batch.audio_latents, batch.audio_latents.shape[0], batch + ) + batch.audio_latents = audio_latents + batch.raw_audio_latent_shape = audio_latents.shape + logger.info( + "Re-packed audio latents for Stage 2: %s", list(audio_latents.shape) + ) + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + device = get_local_torch_device() + latents = self._upsample_video_latents(batch.latents, server_args, device) + logger.info("Upsampled video latents: %s", list(latents.shape)) + self._restore_full_resolution(batch) + self._pack_video_latents(batch, latents, server_args) + logger.info( + "Packed video latents for Stage 2: %s (resolution %dx%d)", + list(batch.latents.shape), + batch.height, + batch.width, + ) + self._repack_audio_latents(batch, server_args) + return batch diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 62edb3dca..f8578a495 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -48,7 +48,6 @@ from sglang.multimodal_gen.utils import ( FlexibleArgumentParser, StoreBoolean, expand_path_fields, - expand_path_kwargs, ) logger = init_logger(__name__) @@ -573,6 +572,15 @@ class ServerArgs: "(e.g. 'Qwen-Image' for 'Qwen/Qwen-Image')." ), ) + parser.add_argument( + "--pipeline-class-name", + type=str, + default=ServerArgs.pipeline_class_name, + help=( + "Override pipeline class selection from model_index.json. " + "Must match a registered pipeline_name." + ), + ) # attention parser.add_argument( "--attention-backend", @@ -936,7 +944,11 @@ class ServerArgs: unknown_args: list[str], ) -> tuple[dict[str, str], list[str]]: """ - Extract dynamic ``---path`` args from unrecognised CLI args. + Extract dynamic component path args from unrecognised CLI args. + + Supported forms: + - ``---path /path/to/component`` + - ``--component-paths. /path/to/component`` (expanded from config) """ component_paths: dict[str, str] = {} remaining: list[str] = [] @@ -944,8 +956,15 @@ class ServerArgs: while i < len(unknown_args): arg = unknown_args[i] key_part = arg.split("=", 1)[0] if "=" in arg else arg - if key_part.startswith("--") and key_part.endswith("-path"): + component = None + if key_part.startswith("--component-paths."): + component = key_part[len("--component-paths.") :].replace("-", "_") + elif key_part.startswith("--component_paths."): + component = key_part[len("--component_paths.") :].replace("-", "_") + elif key_part.startswith("--") and key_part.endswith("-path"): component = key_part[2:-5].replace("-", "_") + + if component is not None: if "=" in arg: component_paths[component] = arg.split("=", 1)[1] elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith( @@ -997,7 +1016,6 @@ class ServerArgs: @classmethod def from_dict(cls, kwargs: dict[str, Any]) -> "ServerArgs": """Create a ServerArgs object from a dictionary.""" - kwargs = expand_path_kwargs(dict(kwargs)) attrs = [attr.name for attr in dataclasses.fields(cls)] server_args_kwargs: dict[str, Any] = {} diff --git a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py index f12f30411..25af21d61 100644 --- a/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py @@ -545,7 +545,7 @@ def maybe_download_model_index(model_name_or_path: str) -> dict[str, Any]: ) return config except EntryNotFoundError: - logger.warning( + logger.debug( "model_index.json not found for %s. Assuming it is a single model and downloading it.", model_name_or_path, ) diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 0be696864..99926bd78 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -1089,6 +1089,70 @@ "expected_avg_denoise_ms": 319.61, "expected_median_denoise_ms": 127.39 }, + "ltx_2_two_stage_t2v": { + "stages_ms": { + "InputValidationStage": 0.05, + "TextEncodingStage": 1827.03, + "LTX2TextConnectorStage": 11.54, + "LTX2HalveResolutionStage": 0.1, + "LTX2LoRASwitchStage": 13014.75, + "LTX2SigmaPreparationStage": 0.25, + "TimestepPreparationStage": 19.26, + "LTX2AVLatentPreparationStage": 0.37, + "LTX2AVDenoisingStage": 53324.98, + "LTX2UpsampleStage": 1894.17, + "LTX2RefinementStage": 4330.57, + "LTX2AVDecodingStage": 337.68 + }, + "denoise_step_ms": { + "0": 1206.17, + "1": 1335.08, + "2": 1336.0, + "3": 1337.77, + "4": 1335.76, + "5": 1334.02, + "6": 1332.78, + "7": 1333.1, + "8": 1334.86, + "9": 1335.26, + "10": 1335.62, + "11": 1334.21, + "12": 1336.52, + "13": 1335.84, + "14": 1338.81, + "15": 1330.74, + "16": 1334.75, + "17": 1337.27, + "18": 1334.05, + "19": 1335.62, + "20": 1349.22, + "21": 1338.2, + "22": 1341.3, + "23": 1358.03, + "24": 1341.14, + "25": 1339.98, + "26": 1332.75, + "27": 1333.31, + "28": 1333.9, + "29": 1333.36, + "30": 1333.96, + "31": 1336.93, + "32": 1335.04, + "33": 1334.48, + "34": 1334.99, + "35": 1333.58, + "36": 1334.49, + "37": 1333.83, + "38": 1332.2, + "39": 1333.38, + "40": 1652.67, + "41": 1330.67, + "42": 1340.97 + }, + "expected_e2e_ms": 75229.22, + "expected_avg_denoise_ms": 1340.53, + "expected_median_denoise_ms": 1334.99 + }, "wan2_2_ti2v_5b": { "stages_ms": { "InputValidationStage": 96.27, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 638fc78f6..151653ee5 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -384,6 +384,10 @@ MULTI_FRAME_I2I_sampling_params = DiffusionSamplingParams( T2V_PROMPT = "A curious raccoon" +T2V_sampling_params = DiffusionSamplingParams( + prompt=T2V_PROMPT, +) + TI2V_sampling_params = DiffusionSamplingParams( prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.", image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg", @@ -399,7 +403,6 @@ TURBOWAN_I2V_sampling_params = DiffusionSamplingParams( fps=4, ) - # All test cases with clean default values # To test different models, simply add more DiffusionCase entries ONE_GPU_CASES_A: list[DiffusionTestCase] = [ @@ -577,9 +580,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ modality="video", custom_validator="video", ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ), DiffusionTestCase( "wan2_1_t2v_1.3b_text_encoder_cpu_offload", @@ -589,9 +590,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ custom_validator="video", text_encoder_cpu_offload=True, ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ), # TeaCache acceleration test for Wan video model DiffusionTestCase( @@ -706,9 +705,7 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ modality="video", custom_validator="video", ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ), # === Text and Image to Video (TI2V) === DiffusionTestCase( @@ -792,9 +789,7 @@ if not current_platform.is_hip(): modality="video", custom_validator="video", ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ) ) @@ -828,9 +823,7 @@ TWO_GPU_CASES_A = [ custom_validator="video", num_gpus=2, ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ), # TeaCache smoke test for Wan2.2 T2V A14B — verifies enable_teacache=True # doesn't crash. Perf check disabled because Wan2.2-specific TeaCache @@ -886,9 +879,7 @@ TWO_GPU_CASES_A = [ num_gpus=2, cfg_parallel=True, ), - DiffusionSamplingParams( - prompt=T2V_PROMPT, - ), + T2V_sampling_params, ), DiffusionTestCase( "fsdp-inference", @@ -938,6 +929,17 @@ TWO_GPU_CASES_A = [ TI2V_sampling_params, run_perf_check=False, ), + DiffusionTestCase( + "ltx_2_two_stage_t2v", + DiffusionServerArgs( + model_path="Lightricks/LTX-2", + modality="video", + num_gpus=2, + dit_layerwise_offload=True, + extras=["--pipeline-class-name LTX2TwoStagePipeline"], + ), + T2V_sampling_params, + ), ] TWO_GPU_CASES_B = [ @@ -1054,7 +1056,6 @@ if not current_platform.is_hip(): ) ) - # Load global configuration BASELINE_CONFIG = BaselineConfig.load( Path(__file__).with_name("perf_baselines.json") diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index b2dbf701d..5c6e519e6 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -102,5 +102,25 @@ class TestPipelineResolutionCliOverride(unittest.TestCase): self.assertEqual(server_args.pipeline_config.resolution, 768) +class TestComponentPathParsing(unittest.TestCase): + def test_extract_component_paths_accepts_config_expanded_keys(self): + component_paths, remaining = ServerArgs._extract_component_paths( + [ + "--component-paths.spatial-upsampler", + "/tmp/latent_upsampler", + "--component_paths.distilled-lora=/tmp/distilled.safetensors", + ] + ) + + self.assertEqual( + component_paths, + { + "spatial_upsampler": "/tmp/latent_upsampler", + "distilled_lora": "/tmp/distilled.safetensors", + }, + ) + self.assertEqual(remaining, []) + + if __name__ == "__main__": unittest.main() diff --git a/python/sglang/multimodal_gen/utils.py b/python/sglang/multimodal_gen/utils.py index 2de9132fe..359fd35f3 100644 --- a/python/sglang/multimodal_gen/utils.py +++ b/python/sglang/multimodal_gen/utils.py @@ -35,25 +35,21 @@ logger = init_logger(__name__) T = TypeVar("T") -def _expand_path_value(field_name: str, value: Any) -> Any: - eu = os.path.expanduser - if field_name.endswith("_path") and isinstance(value, str): - return eu(value) - if field_name.endswith("_path") and isinstance(value, list): - return [eu(x) if isinstance(x, str) else x for x in value] - if field_name.endswith("_paths") and isinstance(value, dict): - return {k: eu(p) if isinstance(p, str) else p for k, p in value.items()} - return value - - -def expand_path_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - return {key: _expand_path_value(key, value) for key, value in kwargs.items()} - - def expand_path_fields(obj) -> None: """In-place expanduser on all dataclass fields whose name ends with '_path' or '_paths'.""" + eu = os.path.expanduser for f in fields(obj): - setattr(obj, f.name, _expand_path_value(f.name, getattr(obj, f.name))) + v = getattr(obj, f.name) + if f.name.endswith("_path") and isinstance(v, str): + setattr(obj, f.name, eu(v)) + elif f.name.endswith("_path") and isinstance(v, list): + setattr(obj, f.name, [eu(x) if isinstance(x, str) else x for x in v]) + elif f.name.endswith("_paths") and isinstance(v, dict): + setattr( + obj, + f.name, + {k: eu(p) if isinstance(p, str) else p for k, p in v.items()}, + ) # TODO(will): used to convert server_args.precision to torch.dtype. Find a diff --git a/scripts/ci/utils/diffusion/comparison_configs.json b/scripts/ci/utils/diffusion/comparison_configs.json index 7f655a1ab..a6f1be874 100644 --- a/scripts/ci/utils/diffusion/comparison_configs.json +++ b/scripts/ci/utils/diffusion/comparison_configs.json @@ -1,156 +1,156 @@ { - "_comment": "Per-model comparison config. Only frameworks listed under each case are tested. vLLM-Omni disabled until dep install issues resolved.", - "test_image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", - "cases": [ - { - "id": "flux1_dev_t2i_1024", - "model": "black-forest-labs/FLUX.1-dev", - "task": "text-to-image", - "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", - "width": 1024, - "height": 1024, - "num_inference_steps": 50, - "guidance_scale": 4.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup --dit-layerwise-offload false", - "extra_env": {} + "_comment": "Per-model comparison config. Only frameworks listed under each case are tested. vLLM-Omni disabled until dep install issues resolved.", + "test_image_url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png", + "cases": [ + { + "id": "flux1_dev_t2i_1024", + "model": "black-forest-labs/FLUX.1-dev", + "task": "text-to-image", + "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", + "width": 1024, + "height": 1024, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --dit-layerwise-offload false", + "extra_env": {} + } + } + }, + { + "id": "flux2_dev_t2i_1024", + "model": "black-forest-labs/FLUX.2-dev", + "task": "text-to-image", + "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", + "width": 1024, + "height": 1024, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --dit-layerwise-offload false", + "extra_env": {} + } + } + }, + { + "id": "qwen_image_2512_t2i_1024", + "model": "Qwen/Qwen-Image-2512", + "task": "text-to-image", + "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", + "width": 1024, + "height": 1024, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup", + "extra_env": {} + } + } + }, + { + "id": "qwen_image_edit_2511", + "model": "Qwen/Qwen-Image-Edit-2511", + "task": "image-edit", + "prompt": "Make the cat wear a red hat", + "reference_image": true, + "width": 1024, + "height": 1024, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup", + "extra_env": {} + } + } + }, + { + "id": "zimage_turbo_t2i_1024", + "model": "Tongyi-MAI/Z-Image-Turbo", + "task": "text-to-image", + "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", + "width": 1024, + "height": 1024, + "num_inference_steps": 9, + "guidance_scale": 4.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup", + "extra_env": {} + } + } + }, + { + "id": "wan22_t2v_a14b_720p", + "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", + "task": "text-to-video", + "prompt": "A cat and a dog baking a cake together in a kitchen.", + "width": 1280, + "height": 720, + "num_frames": 81, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "seed": 42, + "num_gpus": 4, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory", + "extra_env": {} + } + } + }, + { + "id": "wan22_ti2v_5b_720p", + "model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", + "task": "text-image-to-video", + "prompt": "The cat starts walking slowly towards the camera.", + "reference_image": true, + "width": 1280, + "height": 720, + "num_frames": 81, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "seed": 42, + "num_gpus": 1, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup", + "extra_env": {} + } + } + }, + { + "id": "wan22_i2v_a14b_720p", + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "task": "image-to-video", + "prompt": "The cat starts walking slowly towards the camera.", + "reference_image": true, + "width": 1280, + "height": 720, + "num_frames": 81, + "num_inference_steps": 2, + "guidance_scale": 5.0, + "seed": 42, + "num_gpus": 4, + "frameworks": { + "sglang": { + "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory", + "extra_env": {} + } + } } - } - }, - { - "id": "flux2_dev_t2i_1024", - "model": "black-forest-labs/FLUX.2-dev", - "task": "text-to-image", - "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", - "width": 1024, - "height": 1024, - "num_inference_steps": 50, - "guidance_scale": 4.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup --dit-layerwise-offload false", - "extra_env": {} - } - } - }, - { - "id": "qwen_image_2512_t2i_1024", - "model": "Qwen/Qwen-Image-2512", - "task": "text-to-image", - "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", - "width": 1024, - "height": 1024, - "num_inference_steps": 50, - "guidance_scale": 4.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup", - "extra_env": {} - } - } - }, - { - "id": "qwen_image_edit_2511", - "model": "Qwen/Qwen-Image-Edit-2511", - "task": "image-edit", - "prompt": "Make the cat wear a red hat", - "reference_image": true, - "width": 1024, - "height": 1024, - "num_inference_steps": 50, - "guidance_scale": 4.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup", - "extra_env": {} - } - } - }, - { - "id": "zimage_turbo_t2i_1024", - "model": "Tongyi-MAI/Z-Image-Turbo", - "task": "text-to-image", - "prompt": "A futuristic cyberpunk city at night, neon lights reflecting on wet streets", - "width": 1024, - "height": 1024, - "num_inference_steps": 9, - "guidance_scale": 4.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup", - "extra_env": {} - } - } - }, - { - "id": "wan22_t2v_a14b_720p", - "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers", - "task": "text-to-video", - "prompt": "A cat and a dog baking a cake together in a kitchen.", - "width": 1280, - "height": 720, - "num_frames": 81, - "num_inference_steps": 2, - "guidance_scale": 5.0, - "seed": 42, - "num_gpus": 4, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory", - "extra_env": {} - } - } - }, - { - "id": "wan22_ti2v_5b_720p", - "model": "Wan-AI/Wan2.2-TI2V-5B-Diffusers", - "task": "text-image-to-video", - "prompt": "The cat starts walking slowly towards the camera.", - "reference_image": true, - "width": 1280, - "height": 720, - "num_frames": 81, - "num_inference_steps": 50, - "guidance_scale": 5.0, - "seed": 42, - "num_gpus": 1, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup", - "extra_env": {} - } - } - }, - { - "id": "wan22_i2v_a14b_720p", - "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - "task": "image-to-video", - "prompt": "The cat starts walking slowly towards the camera.", - "reference_image": true, - "width": 1280, - "height": 720, - "num_frames": 81, - "num_inference_steps": 2, - "guidance_scale": 5.0, - "seed": 42, - "num_gpus": 4, - "frameworks": { - "sglang": { - "serve_args": "--enable-torch-compile --warmup --enable-cfg-parallel --ulysses-degree 2 --text-encoder-cpu-offload --pin-cpu-memory", - "extra_env": {} - } - } - } - ] + ] }