minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: support LoRA for LTX2.3 (#23649)

Browse Source
This commit is contained in:
Mick
2026-04-25 01:52:41 +08:00
committed by GitHub
parent baa0aa670f
commit 05696527ea
15 changed files with 810 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/lora/linear.py
+70 -22
View File
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
# Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
import os
import torch
from torch import nn
@@ -37,10 +37,17 @@ torch._dynamo.config.recompile_limit = 64
LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024
LoRAWeightEntry = tuple[
torch.nn.Parameter,
torch.nn.Parameter,
str | None,
float,
int | None,
int | None,
]
class BaseLayerWithLoRA(nn.Module):
def __init__(
self,
base_layer: nn.Module,
@@ -60,9 +67,7 @@ class BaseLayerWithLoRA(nn.Module):
self.disable_lora: bool = True
self.lora_rank = lora_rank
self.lora_alpha = lora_alpha
self.lora_weights_list: list[
tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
] = []
self.lora_weights_list: list[LoRAWeightEntry] = []
self.lora_path: str | None = None
self.strength: float = 1.0
@@ -147,7 +152,16 @@ class BaseLayerWithLoRA(nn.Module):
self.strength = 1.0
# Add to list for multi-LoRA support
self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength))
self.lora_weights_list.append(
(
lora_A_param,
lora_B_param,
lora_path,
strength,
self.lora_rank,
self.lora_alpha,
)
)
# Set backward compatibility attributes to point to the last LoRA (for single LoRA case)
# This ensures backward compatibility while supporting multiple LoRA
@@ -166,29 +180,27 @@ class BaseLayerWithLoRA(nn.Module):
def _merge_lora_into_data(
self,
data: torch.Tensor,
lora_list: list[
tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float]
],
lora_list: list[LoRAWeightEntry],
) -> None:
"""
Merge all LoRA adapters into the data tensor in-place.
Args:
data: The base weight tensor to merge LoRA into (modified in-place)
lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples
lora_list: List of (lora_A, lora_B, lora_path, lora_strength, rank, alpha) tuples
"""
# Merge all LoRA adapters in order
for lora_A, lora_B, _, lora_strength in lora_list:
for lora_A, lora_B, _, lora_strength, lora_rank, lora_alpha in lora_list:
lora_A_sliced = self.slice_lora_a_weights(lora_A.to(data))
lora_B_sliced = self.slice_lora_b_weights(lora_B.to(data))
scale = lora_strength
if (
self.lora_alpha is not None
and self.lora_rank is not None
and self.lora_alpha != self.lora_rank
lora_alpha is not None
and lora_rank is not None
and lora_alpha != lora_rank
):
scale *= self.lora_alpha / self.lora_rank
scale *= lora_alpha / lora_rank
if not isinstance(lora_B_sliced, torch.Tensor):
lora_delta = lora_B_sliced @ lora_A_sliced
@@ -222,6 +234,17 @@ class BaseLayerWithLoRA(nn.Module):
chunk_delta = lora_B_2d[start:end] @ lora_A_sliced
data_2d[start:end].add_(chunk_delta, alpha=scale)
def _should_merge_in_fp32(
self,
lora_list: list[LoRAWeightEntry],
) -> bool:
if os.getenv("SGLANG_DIFFUSION_LORA_MERGE_FP32", "0") != "1":
return False
for _, _, lora_path, _, _, _ in lora_list:
if lora_path and "distilled-lora" in lora_path.lower():
return False
return True
@torch.no_grad()
def merge_lora_weights(self, strength: float | None = None) -> None:
if strength is not None:
@@ -236,11 +259,22 @@ class BaseLayerWithLoRA(nn.Module):
# Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility
lora_list = self.lora_weights_list if self.lora_weights_list else []
if not lora_list and self.lora_A is not None and self.lora_B is not None:
lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)]
lora_list = [
(
self.lora_A,
self.lora_B,
self.lora_path,
self.strength,
self.lora_rank,
self.lora_alpha,
)
]
if not lora_list:
raise ValueError("LoRA weights not set. Please set them first.")
merge_in_fp32 = self._should_merge_in_fp32(lora_list)
if isinstance(self.base_layer.weight, DTensor):
mesh = self.base_layer.weight.data.device_mesh
unsharded_base_layer = ReplicatedLinear(
@@ -257,10 +291,19 @@ class BaseLayerWithLoRA(nn.Module):
data = self.base_layer.weight.data.to(
get_local_torch_device()
).full_tensor()
target_dtype = data.dtype
if (
merge_in_fp32
and data.is_floating_point()
and data.dtype != torch.float32
):
data = data.to(torch.float32)
self._merge_lora_into_data(data, lora_list)
unsharded_base_layer.weight = nn.Parameter(data.to(current_device))
unsharded_base_layer.weight = nn.Parameter(
data.to(current_device, dtype=target_dtype)
)
if isinstance(getattr(self.base_layer, "bias", None), DTensor):
unsharded_base_layer.bias = nn.Parameter(
self.base_layer.bias.to(get_local_torch_device(), non_blocking=True)
@@ -282,10 +325,19 @@ class BaseLayerWithLoRA(nn.Module):
else:
current_device = self.base_layer.weight.data.device
data = self.base_layer.weight.data.to(get_local_torch_device())
target_dtype = data.dtype
if (
merge_in_fp32
and data.is_floating_point()
and data.dtype != torch.float32
):
data = data.to(torch.float32)
self._merge_lora_into_data(data, lora_list)
self.base_layer.weight.data = data.to(current_device, non_blocking=True)
self.base_layer.weight.data = data.to(
current_device, dtype=target_dtype, non_blocking=True
)
self.merged = True
@@ -342,7 +394,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
def __init__(
self,
base_layer: ColumnParallelLinear,
@@ -400,7 +451,6 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
def __init__(
self,
base_layer: MergedColumnParallelLinear,
@@ -422,7 +472,6 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
def __init__(
self,
base_layer: QKVParallelLinear,
@@ -455,7 +504,6 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
def __init__(
self,
base_layer: RowParallelLinear,
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+27 -8
View File
@@ -494,19 +494,25 @@ class LTX2TwoStageDeviceManager:
if "stage2" in self._phase_ready_events:
return
if self._snapshot_low_vram_mode:
stage1_module = self.pipeline.get_module("transformer")
stage1_param = (
next(stage1_module.parameters(), None)
if stage1_module is not None
else None
)
if stage1_param is not None and stage1_param.device.type == "cuda":
self._release_module_to_cpu_snapshot("transformer")
self._release_stage1_for_low_vram()
self._schedule_phase_prefetch(
"stage2", self.pipeline.get_module("transformer_2")
)
def prepare_upsample_after_stage1(self) -> bool:
if (
not self.should_use_premerged
or self.mode != "snapshot"
or not self.server_args.dit_cpu_offload
or not self._snapshot_low_vram_mode
):
return False
if "stage2" in self._phase_ready_events:
return False
self._release_stage1_for_low_vram()
return True
def ensure_phase_ready(self, phase: str | None) -> None:
if not self.should_use_premerged or phase not in ("stage1", "stage2"):
return
@@ -616,6 +622,16 @@ class LTX2TwoStageDeviceManager:
phase = "stage2" if module_name == "transformer_2" else "stage1"
self._phase_ready_events.pop(phase, None)
def _release_stage1_for_low_vram(self) -> None:
stage1_module = self.pipeline.get_module("transformer")
stage1_param = (
next(stage1_module.parameters(), None)
if stage1_module is not None
else None
)
if stage1_param is not None and stage1_param.device.type == "cuda":
self._release_module_to_cpu_snapshot("transformer")
def _ensure_on_gpu(self, module_name: str) -> None:
module = self.pipeline.get_module(module_name)
if module is None:
@@ -807,6 +823,9 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
def prefetch_ltx2_stage2_after_stage1(self) -> None:
self._device_manager.prefetch_stage2_after_stage1()
def prepare_ltx2_upsample_after_stage1(self) -> bool:
return self._device_manager.prepare_upsample_after_stage1()
def should_skip_ltx2_lora_switch_stage(self) -> bool:
return self._use_premerged_stage2_transformer and self._device_manager.mode in (
"snapshot",
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
+38 -7
View File
@@ -99,7 +99,9 @@ class LoRAPipeline(ComposedPipelineBase):
if self.lora_path is not None:
self.convert_to_lora_layers()
self.set_lora(
self.lora_nickname, self.lora_path, strength=self.server_args.lora_scale # type: ignore
self.lora_nickname,
self.lora_path,
strength=self.server_args.lora_scale, # type: ignore
) # type: ignore
def is_target_layer(self, module_name: str) -> bool:
@@ -426,6 +428,8 @@ class LoRAPipeline(ComposedPipelineBase):
)
adapted_count = 0
missing_layers_by_adapter = [[] for _ in lora_nicknames]
applied_count_by_adapter = [0 for _ in lora_nicknames]
for name, layer in lora_layers.items():
# Apply all LoRA adapters in order
for idx, (nickname, path, lora_strength) in enumerate(
@@ -465,13 +469,9 @@ class LoRAPipeline(ComposedPipelineBase):
), # Only clear on first LoRA
)
adapted_count += 1
applied_count_by_adapter[idx] += 1
else:
if rank == 0 and idx == 0: # Only warn for first missing LoRA
logger.warning(
"LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.",
path,
name,
)
missing_layers_by_adapter[idx].append(name)
# Only disable if no LoRA was applied at all
if idx == len(lora_nicknames) - 1:
has_any_lora = any(
@@ -481,6 +481,37 @@ class LoRAPipeline(ComposedPipelineBase):
)
if not has_any_lora:
layer.disable_lora = True
if rank == 0:
total_layers = len(lora_layers)
example_limit = 8
for idx, path in enumerate(lora_paths):
missing_layers = missing_layers_by_adapter[idx]
if not missing_layers:
continue
missing_count = len(missing_layers)
applied_count = applied_count_by_adapter[idx]
examples = ", ".join(missing_layers[:example_limit])
if missing_count > example_limit:
examples += ", ..."
if applied_count == 0:
logger.warning(
"LoRA adapter %s did not match any LoRA layer. "
"Checked %d layers; examples: %s",
path,
total_layers,
examples,
)
else:
logger.info(
"LoRA adapter %s covers %d/%d LoRA layers; "
"%d layers use base weights. Examples: %s",
path,
applied_count,
total_layers,
missing_count,
examples,
)
return adapted_count
def is_lora_effective(self, target: str = "all") -> bool:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+9
View File
@@ -217,6 +217,14 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
"""Run the distilled refinement schedule on top of the shared AV denoiser."""
batch.extra["ltx2_phase"] = "stage2"
pipeline = self.pipeline() if self.pipeline else None
ensure_phase_ready = (
getattr(pipeline, "ensure_ltx2_phase_ready", None)
if pipeline is not None
else None
)
if callable(ensure_phase_ready):
ensure_phase_ready("stage2")
original_clean_latent_background = getattr(
batch, "ltx2_ti2v_clean_latent_background", None
)
@@ -250,6 +258,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
renoise_generator = None
if is_native_ti2v:
prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state(
batch=batch,
latents=batch.latents,
image_latent=batch.image_latent,
num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py
+79 -56
View File
@@ -417,25 +417,48 @@ class LTX2ImageEncodingStage(PipelineStage):
):
return self._condition_image_encoder(video_condition)
@staticmethod
def _normalize_ltx2_image_paths(image_path: str | list[str]) -> list[str]:
image_paths = image_path if isinstance(image_path, list) else [image_path]
if len(image_paths) > 2:
raise ValueError(
"LTX-2 TI2V currently supports at most two conditioning images "
"([first_frame, last_frame])."
)
return image_paths
@staticmethod
def _normalize_ltx2_image_latents(
image_latent: torch.Tensor | list[torch.Tensor] | None,
) -> list[torch.Tensor]:
if image_latent is None:
return []
return image_latent if isinstance(image_latent, list) else [image_latent]
# -- forward ---------------------------------------------------------
@torch.no_grad()
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
if batch.image_path is None:
return batch
image_paths = self._normalize_ltx2_image_paths(batch.image_path)
vae_sf = int(server_args.pipeline_config.vae_scale_factor)
patch = int(server_args.pipeline_config.patch_size)
expected_tokens = (int(batch.height) // vae_sf // patch) * (
int(batch.width) // vae_sf // patch
)
if (
batch.image_latent is not None
and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0
):
# Re-encode if resolution changed (e.g. two-stage upsample between stages)
vae_sf = int(server_args.pipeline_config.vae_scale_factor)
patch = int(server_args.pipeline_config.patch_size)
expected = (int(batch.height) // vae_sf // patch) * (
int(batch.width) // vae_sf // patch
)
if int(batch.image_latent.shape[1]) == expected:
existing_latents = self._normalize_ltx2_image_latents(batch.image_latent)
if len(existing_latents) == len(image_paths) and all(
int(latent.shape[1]) == expected_tokens for latent in existing_latents
):
return batch
# Resolution mismatch — clear and re-encode below
# Resolution or reference-count mismatch — clear and re-encode below
batch.image_latent = None
batch.ltx2_num_image_tokens = 0
@@ -447,20 +470,23 @@ class LTX2ImageEncodingStage(PipelineStage):
from sglang.multimodal_gen.runtime.models.vision_utils import load_image
# 1. Load image, apply codec compression, resize for condition_image
image_path = (
batch.image_path[0]
if isinstance(batch.image_path, list)
else batch.image_path
)
img = load_image(image_path)
arr = np.array(img).astype(np.uint8)[..., :3]
arr = self._apply_video_codec_compression(arr, crf=33)
conditioned_img = PIL.Image.fromarray(arr)
batch.condition_image = conditioned_img.resize(
(int(batch.width), int(batch.height)),
resample=PIL.Image.Resampling.BILINEAR,
)
# 1. Load images, apply codec compression, resize for condition_image
conditioned_imgs = []
for image_path in image_paths:
img = load_image(image_path)
arr = np.array(img).astype(np.uint8)[..., :3]
arr = self._apply_video_codec_compression(arr, crf=33)
conditioned_img = PIL.Image.fromarray(arr)
conditioned_imgs.append(conditioned_img)
batch.condition_image = [
img.resize(
(int(batch.width), int(batch.height)),
resample=PIL.Image.Resampling.BILINEAR,
)
for img in conditioned_imgs
]
if len(batch.condition_image) == 1:
batch.condition_image = batch.condition_image[0]
# 2. Load encoder(s) to device, cast to encode_dtype
use_condition_encoder = self._ensure_condition_image_encoder(server_args)
@@ -478,21 +504,35 @@ class LTX2ImageEncodingStage(PipelineStage):
else:
self.vae = self.vae.to(dtype=encode_dtype)
video_condition = self._pil_to_video_tensor(
conditioned_img,
width=int(batch.width),
height=int(batch.height),
device=device,
dtype=encode_dtype,
)
# 3. Encode
if use_condition_encoder:
latent = self._condition_encode(video_condition, server_args).to(
dtype=encode_dtype
packed_latents = []
for conditioned_img in conditioned_imgs:
video_condition = self._pil_to_video_tensor(
conditioned_img,
width=int(batch.width),
height=int(batch.height),
device=device,
dtype=encode_dtype,
)
else:
latent = self._vae_encode(video_condition, server_args, batch.generator)
# 3. Encode
if use_condition_encoder:
latent = self._condition_encode(video_condition, server_args).to(
dtype=encode_dtype
)
else:
latent = self._vae_encode(video_condition, server_args, batch.generator)
packed = server_args.pipeline_config.maybe_pack_latents(
latent, latent.shape[0], batch
)
if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
raise ValueError("Expected packed image latents [B, S0, D].")
if int(packed.shape[1]) != expected_tokens:
raise ValueError(
f"LTX-2 conditioning token count mismatch: "
f"{packed.shape[1]=} {expected_tokens=}."
)
packed_latents.append(packed)
# Restore VAE to its config dtype (shared with decoding stage)
if not use_condition_encoder:
@@ -501,32 +541,16 @@ class LTX2ImageEncodingStage(PipelineStage):
]
self.vae = self.vae.to(dtype=original_dtype)
# 4. Pack into token latents and validate
packed = server_args.pipeline_config.maybe_pack_latents(
latent, latent.shape[0], batch
batch.image_latent = (
packed_latents[0] if len(packed_latents) == 1 else packed_latents
)
if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
raise ValueError("Expected packed image latents [B, S0, D].")
vae_sf = int(server_args.pipeline_config.vae_scale_factor)
patch = int(server_args.pipeline_config.patch_size)
expected_tokens = (int(batch.height) // vae_sf // patch) * (
int(batch.width) // vae_sf // patch
)
if int(packed.shape[1]) != expected_tokens:
raise ValueError(
f"LTX-2 conditioning token count mismatch: "
f"{packed.shape[1]=} {expected_tokens=}."
)
batch.image_latent = packed
batch.ltx2_num_image_tokens = int(packed.shape[1])
batch.ltx2_num_image_tokens = int(packed_latents[0].shape[1])
if batch.debug:
logger.info(
"LTX2 TI2V: %d tokens (shape=%s) for %sx%s",
batch.ltx2_num_image_tokens,
tuple(batch.image_latent.shape),
tuple(packed_latents[0].shape),
batch.width,
batch.height,
)
@@ -703,7 +727,6 @@ class ImageVAEEncodingStage(PipelineStage):
self,
image: torch.Tensor | PIL.Image.Image,
) -> torch.Tensor:
if isinstance(image, PIL.Image.Image):
image = pil_to_numpy(image) # to np
image = numpy_to_pt(image) # to pt
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+120 -27
View File
@@ -516,24 +516,111 @@ class LTX2DenoisingStage(DenoisingStage):
return next_video, next_audio
@staticmethod
def _prepare_ltx2_ti2v_clean_state(
def _normalize_ltx2_condition_latents(
image_latent: torch.Tensor | list[torch.Tensor] | None,
) -> list[torch.Tensor]:
if image_latent is None:
return []
return image_latent if isinstance(image_latent, list) else [image_latent]
@classmethod
def _get_ltx2_condition_spans(
cls,
batch: Req,
latents: torch.Tensor,
image_latent: torch.Tensor,
image_latent: torch.Tensor | list[torch.Tensor] | None,
num_img_tokens: int,
) -> list[tuple[int, torch.Tensor]]:
if num_img_tokens <= 0:
return []
if not (isinstance(latents, torch.Tensor) and latents.ndim == 3):
raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
condition_latents = cls._normalize_ltx2_condition_latents(image_latent)
if not condition_latents:
return []
if len(condition_latents) > 2:
raise ValueError(
"LTX-2 TI2V currently supports at most two conditioning images."
)
for cond in condition_latents:
if not (isinstance(cond, torch.Tensor) and cond.ndim == 3):
raise ValueError(
"Expected LTX-2 conditioning latents to be packed tensors [B, S, D]."
)
if int(cond.shape[1]) < int(num_img_tokens):
raise ValueError(
"LTX-2 conditioning latent is shorter than one frame token span."
)
did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
if not did_sp_shard:
if int(latents.shape[1]) < int(num_img_tokens):
raise ValueError(
"LTX-2 latent sequence is shorter than one conditioning frame."
)
if len(condition_latents) == 1:
return [(0, condition_latents[0])]
return [
(0, condition_latents[0]),
(int(latents.shape[1]) - int(num_img_tokens), condition_latents[1]),
]
tokens_per_frame = int(getattr(batch, "sp_video_tokens_per_frame", 0))
if tokens_per_frame <= 0:
raise ValueError(
"SP-sharded LTX-2 TI2V requires batch.sp_video_tokens_per_frame."
)
if int(num_img_tokens) != int(tokens_per_frame):
raise ValueError(
"LTX-2 conditioning token count must match one latent frame when using SP."
)
raw_shape = getattr(batch, "raw_latent_shape", None)
if raw_shape is None:
raise ValueError("SP-sharded LTX-2 TI2V requires batch.raw_latent_shape.")
global_seq_len = int(raw_shape[1])
if global_seq_len % tokens_per_frame != 0:
raise ValueError(
"SP-sharded LTX-2 TI2V expected raw seq_len divisible by tokens_per_frame."
)
global_num_frames = global_seq_len // tokens_per_frame
local_start_frame = int(getattr(batch, "sp_video_start_frame", 0))
local_num_frames = int(getattr(batch, "sp_video_latent_num_frames", 0))
local_end_frame = local_start_frame + local_num_frames
spans: list[tuple[int, torch.Tensor]] = []
if local_start_frame == 0:
spans.append((0, condition_latents[0]))
if len(condition_latents) == 2:
last_global_frame = global_num_frames - 1
if local_start_frame <= last_global_frame < local_end_frame:
local_last_frame = last_global_frame - local_start_frame
spans.append(
(local_last_frame * tokens_per_frame, condition_latents[1])
)
return spans
@classmethod
def _prepare_ltx2_ti2v_clean_state(
cls,
batch: Req,
latents: torch.Tensor,
image_latent: torch.Tensor | list[torch.Tensor] | None,
num_img_tokens: int,
zero_clean_latent: bool,
clean_latent_background: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
latents = latents.clone()
conditioned = image_latent[:, :num_img_tokens, :].to(
device=latents.device, dtype=latents.dtype
)
latents[:, :num_img_tokens, :] = conditioned
denoise_mask = torch.ones(
(latents.shape[0], latents.shape[1], 1),
device=latents.device,
dtype=torch.float32,
)
denoise_mask[:, :num_img_tokens, :] = 0.0
if clean_latent_background is not None:
clean_latent = (
clean_latent_background.detach()
@@ -544,7 +631,24 @@ class LTX2DenoisingStage(DenoisingStage):
clean_latent = torch.zeros_like(latents)
else:
clean_latent = latents.detach().clone()
clean_latent[:, :num_img_tokens, :] = conditioned
spans = cls._get_ltx2_condition_spans(
batch=batch,
latents=latents,
image_latent=image_latent,
num_img_tokens=num_img_tokens,
)
for start, cond in spans:
stop = int(start) + int(num_img_tokens)
conditioned = cls._repeat_batch_dim(
cond[:, :num_img_tokens, :].to(
device=latents.device, dtype=latents.dtype
),
int(latents.shape[0]),
)
latents[:, start:stop, :] = conditioned
denoise_mask[:, start:stop, :] = 0.0
clean_latent[:, start:stop, :] = conditioned
return latents, denoise_mask, clean_latent
@staticmethod
@@ -915,23 +1019,6 @@ class LTX2DenoisingStage(DenoisingStage):
super()._preprocess_sp_latents(batch, server_args)
batch.image_latent = saved
@staticmethod
def _should_apply_ltx2_ti2v(batch: Req) -> bool:
"""True if we have an image-latent token prefix to condition with.
SP note: when token latents are time-sharded, only the rank that owns the
*global* first latent frame should apply TI2V conditioning (rank with start_frame==0).
"""
if (
batch.image_latent is None
or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0
):
return False
did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
if not did_sp_shard:
return True
return int(getattr(batch, "sp_video_start_frame", 0)) == 0
@staticmethod
def _should_use_native_hq_res2s_sde_noise(server_args: ServerArgs) -> bool:
return server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
@@ -998,8 +1085,6 @@ class LTX2DenoisingStage(DenoisingStage):
# Video and audio keep separate scheduler state throughout the denoising loop.
ctx.audio_scheduler = copy.deepcopy(self.scheduler)
do_ti2v = self._should_apply_ltx2_ti2v(batch)
if ctx.use_ltx23_legacy_one_stage:
batch.ltx23_audio_replicated_for_sp = False
batch.did_sp_shard_audio_latents = False
@@ -1038,6 +1123,13 @@ class LTX2DenoisingStage(DenoisingStage):
batch.width
// server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
)
ti2v_spans = self._get_ltx2_condition_spans(
batch=batch,
latents=ctx.latents,
image_latent=batch.image_latent,
num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
)
do_ti2v = bool(ti2v_spans)
if do_ti2v:
if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3):
raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
@@ -1052,6 +1144,7 @@ class LTX2DenoisingStage(DenoisingStage):
# Keep conditioned tokens clean and reuse the mask during every step update.
ctx.latents, ctx.denoise_mask, ctx.clean_latent = (
self._prepare_ltx2_ti2v_clean_state(
batch=batch,
latents=ctx.latents,
image_latent=batch.image_latent,
num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)),
python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py
+16 -3
View File
@@ -84,8 +84,8 @@ class LTX2UpsampleStage(PipelineStage):
device=device, dtype=latents.dtype
)
latents = self.spatial_upsampler(latents)
if server_args.vae_cpu_offload:
self.spatial_upsampler = self.spatial_upsampler.to("cpu")
# Keep the small spatial upsampler resident after warmup; moving it
# every request dominates the measured two-stage upsample latency.
latents = (latents - vae_mean) / vae_std
return latents
@@ -118,18 +118,31 @@ class LTX2UpsampleStage(PipelineStage):
)
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
delay_stage2_prefetch = False
if self.pipeline is not None:
prepare_upsample = getattr(
self.pipeline, "prepare_ltx2_upsample_after_stage1", None
)
if callable(prepare_upsample):
delay_stage2_prefetch = prepare_upsample()
prefetch_stage2 = (
getattr(self.pipeline, "prefetch_ltx2_stage2_after_stage1", None)
if self.pipeline is not None
else None
)
if callable(prefetch_stage2):
if callable(prefetch_stage2) and not delay_stage2_prefetch:
prefetch_stage2()
device = get_local_torch_device()
latents = self._upsample_video_latents(batch.latents, server_args, device)
if callable(prefetch_stage2) and delay_stage2_prefetch:
prefetch_stage2()
logger.info("Upsampled video latents: %s", list(latents.shape))
self._restore_full_resolution(batch)
batch.image_latent = None
batch.ltx2_num_image_tokens = 0
batch.did_sp_shard_latents = False
batch.did_sp_shard_audio_latents = False
self._pack_video_latents(batch, latents, server_args)
logger.info(
"Packed video latents for Stage 2: %s (resolution %dx%d)",
python/sglang/multimodal_gen/test/server/accuracy_config.py
-11
View File
@@ -142,17 +142,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = {
"Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TRANSFORMER: ComponentSkip(
"Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
ComponentType.TEXT_ENCODER: ComponentSkip(
"Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
),
},
"wan2_1_t2v_1.3b_teacache_enabled": {
ComponentType.VAE: ComponentSkip(
"Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology"
python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py
-1
View File
@@ -31,7 +31,6 @@ ACCURACY_ONE_GPU_CASE_IDS = (
"flux_2_image_t2i_upscaling_4x",
"mova_360p_1gpu",
"wan2_1_t2v_1.3b",
"wan2_1_t2v_1.3b_text_encoder_cpu_offload",
"wan2_1_t2v_1.3b_teacache_enabled",
"wan2_1_t2v_1.3b_frame_interp_2x",
"wan2_1_t2v_1.3b_upscaling_4x",
python/sglang/multimodal_gen/test/server/gpu_cases.py
-8
View File
@@ -175,14 +175,6 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
),
T2V_sampling_params,
),
DiffusionTestCase(
"wan2_1_t2v_1.3b_text_encoder_cpu_offload",
DiffusionServerArgs(
model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST,
text_encoder_cpu_offload=True,
),
T2V_sampling_params,
),
# TeaCache acceleration test for Wan video model
DiffusionTestCase(
"wan2_1_t2v_1.3b_teacache_enabled",
python/sglang/multimodal_gen/test/server/perf_baselines.json
-66
View File
@@ -951,72 +951,6 @@
"expected_median_denoise_ms": 145.57,
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_text_encoder_cpu_offload": {
"stages_ms": {
"DecodingStage": 675.91,
"TextEncodingStage": 1072.93,
"TimestepPreparationStage": 2.43,
"LatentPreparationStage": 0.14,
"InputValidationStage": 0.07,
"DenoisingStage": 7221.86
},
"denoise_step_ms": {
"0": 128.64,
"1": 108.46,
"2": 140.93,
"3": 140.8,
"4": 141.21,
"5": 141.84,
"6": 141.56,
"7": 142.35,
"8": 142.08,
"9": 141.68,
"10": 141.69,
"11": 141.46,
"12": 141.39,
"13": 141.36,
"14": 141.65,
"15": 141.55,
"16": 142.02,
"17": 141.53,
"18": 140.98,
"19": 142.4,
"20": 141.84,
"21": 141.3,
"22": 141.41,
"23": 141.47,
"24": 141.78,
"25": 141.6,
"26": 142.19,
"27": 141.09,
"28": 141.2,
"29": 141.22,
"30": 141.2,
"31": 141.23,
"32": 141.41,
"33": 141.5,
"34": 141.56,
"35": 141.51,
"36": 141.25,
"37": 141.49,
"38": 141.56,
"39": 141.52,
"40": 141.17,
"41": 141.83,
"42": 141.72,
"43": 142.31,
"44": 141.56,
"45": 141.91,
"46": 141.93,
"47": 141.9,
"48": 141.52,
"49": 141.34
},
"expected_e2e_ms": 9296.17,
"expected_avg_denoise_ms": 144.33,
"expected_median_denoise_ms": 144.84,
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_cfg_parallel": {
"stages_ms": {
"LatentPreparationStage": 0.29,