[diffusion] model: support LTX2.3 high quality pipeline (#23366)
This commit is contained in:
Mick
@@ -33,23 +33,29 @@ 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 | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ |
|
||||
| Wan2.1 Fun 1.3B InP | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ |
|
||||
| Helios Base | `BestWishYsh/Helios-Base` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| Helios Mid | `BestWishYsh/Helios-Mid` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| Helios Distilled | `BestWishYsh/Helios-Distilled` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| LTX-2 (one and two stages) | `Lightricks/LTX-2` | 768×512<br>1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| LTX-2.3 (one and two stages) | `Lightricks/LTX-2.3` | 768×512<br>1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| Wan2.1 Fun 1.3B InP | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ |
|
||||
| Helios Base | `BestWishYsh/Helios-Base` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| Helios Mid | `BestWishYsh/Helios-Mid` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| Helios Distilled | `BestWishYsh/Helios-Distilled` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| LTX-2 (one/two-stage/TI2V) | `Lightricks/LTX-2` | 768×512<br>1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| LTX-2.3 (one/two-stage/TI2V/HQ) | `Lightricks/LTX-2.3` | 768×512<br>1536×1024<br>1920×1088 (HQ default) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
|
||||
**Note**:
|
||||
|
||||
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 and LTX-2.3 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`.
|
||||
3. LTX pipeline selection:
|
||||
- One-stage: `--pipeline-class-name LTX2Pipeline`
|
||||
- Two-stage: `--pipeline-class-name LTX2TwoStagePipeline`
|
||||
- Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`)
|
||||
- LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ).
|
||||
- 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`.
|
||||
- For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
|
||||
4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {legacy,snapshot,resident}`:
|
||||
4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`:
|
||||
- `snapshot` is the default and recommended mode.
|
||||
- `resident` usually provides the best latency/throughput but uses much more VRAM.
|
||||
- `legacy` preserves the historical switching path for fallback/debug.
|
||||
- `original` keeps official two-stage semantics without the premerged stage-2 transformer path.
|
||||
- Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`.
|
||||
|
||||
### Image Generation Models
|
||||
|
||||
|
||||
@@ -285,7 +285,7 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2 (one and two stages)</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2 (one/two-stage/TI2V)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2</code></td>
|
||||
<td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
@@ -297,9 +297,9 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one and two stages)</td>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>LTX-2.3 (one/two-stage/TI2V/HQ)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>Lightricks/LTX-2.3</code></td>
|
||||
<td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024</td>
|
||||
<td style={{padding: "9px 8px", backgroundColor: "rgba(255,255,255,0.02)"}}>768×512<br />1536×1024<br />1920×1088 (HQ default)</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>❌</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>❌</td>
|
||||
@@ -315,12 +315,18 @@ Optimization columns are abbreviated to keep the matrix readable:
|
||||
|
||||
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 and LTX-2.3 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`.
|
||||
- For LTX models, the `Resolution` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
|
||||
4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {legacy,snapshot,resident}`:
|
||||
3. LTX pipeline selection:
|
||||
- One-stage: `--pipeline-class-name LTX2Pipeline`
|
||||
- Two-stage: `--pipeline-class-name LTX2TwoStagePipeline`
|
||||
- Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`)
|
||||
- LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ).
|
||||
- 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`.
|
||||
- For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`.
|
||||
4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`:
|
||||
- `snapshot` is the default and recommended mode.
|
||||
- `resident` usually provides the best latency/throughput but uses much more VRAM.
|
||||
- `legacy` preserves the historical switching path for fallback/debug.
|
||||
- `original` keeps official two-stage semantics without the premerged stage-2 transformer path.
|
||||
- Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`.
|
||||
|
||||
### Image Generation Models
|
||||
|
||||
|
||||
@@ -244,6 +244,9 @@ class LTX2PipelineConfig(PipelineConfig):
|
||||
def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict:
|
||||
# Adapted from diffusers_pipeline.py _get_gemma_prompt_embeds
|
||||
# But we only need tokenization here, the embedding happens in TextEncodingStage
|
||||
# Official LTX Gemma tokenizer trims surrounding whitespace before
|
||||
# tokenization.
|
||||
prompt = [text.strip() for text in prompt]
|
||||
|
||||
# Gemma expects left padding for chat-style prompts
|
||||
tokenizer.padding_side = "left"
|
||||
|
||||
@@ -82,3 +82,38 @@ class LTX23SamplingParams(LTX2SamplingParams):
|
||||
"audio_stg_blocks": self.audio_stg_blocks,
|
||||
}
|
||||
return extra
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LTX23HQSamplingParams(LTX23SamplingParams):
|
||||
"""Sampling parameters matching official LTX-2.3 HQ two-stage defaults."""
|
||||
|
||||
height: int = 1088
|
||||
width: int = 1920
|
||||
num_inference_steps: int = 15
|
||||
distilled_lora_strength_stage_1: float = 0.25
|
||||
distilled_lora_strength_stage_2: float = 0.5
|
||||
|
||||
video_cfg_scale: float = 3.0
|
||||
video_stg_scale: float = 0.0
|
||||
video_rescale_scale: float = 0.45
|
||||
video_modality_scale: float = 3.0
|
||||
video_skip_step: int = 0
|
||||
video_stg_blocks: list[int] = field(default_factory=list)
|
||||
|
||||
audio_cfg_scale: float = 7.0
|
||||
audio_stg_scale: float = 0.0
|
||||
audio_rescale_scale: float = 1.0
|
||||
audio_modality_scale: float = 3.0
|
||||
audio_skip_step: int = 0
|
||||
audio_stg_blocks: list[int] = field(default_factory=list)
|
||||
|
||||
def build_request_extra(self) -> dict[str, Any]:
|
||||
extra = super().build_request_extra()
|
||||
extra["ltx2_distilled_lora_strength_stage_1"] = float(
|
||||
self.distilled_lora_strength_stage_1
|
||||
)
|
||||
extra["ltx2_distilled_lora_strength_stage_2"] = float(
|
||||
self.distilled_lora_strength_stage_2
|
||||
)
|
||||
return extra
|
||||
|
||||
@@ -571,18 +571,28 @@ class SamplingParams:
|
||||
def from_user_sampling_params_args(
|
||||
model_path: str, server_args: "ServerArgs", *args, **kwargs
|
||||
):
|
||||
pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
|
||||
try:
|
||||
sampling_params = SamplingParams.from_pretrained(
|
||||
model_path, backend=server_args.backend, model_id=server_args.model_id
|
||||
)
|
||||
except (AttributeError, ValueError) as e:
|
||||
sampling_params = None
|
||||
if pipeline_class_name:
|
||||
from sglang.multimodal_gen.registry import get_pipeline_config_classes
|
||||
|
||||
config_classes = get_pipeline_config_classes(pipeline_class_name)
|
||||
if config_classes is not None:
|
||||
_, sampling_params_cls = config_classes
|
||||
sampling_params = sampling_params_cls()
|
||||
|
||||
if sampling_params is None:
|
||||
sampling_params = SamplingParams.from_pretrained(
|
||||
model_path,
|
||||
backend=server_args.backend,
|
||||
model_id=server_args.model_id,
|
||||
)
|
||||
except (AttributeError, ValueError):
|
||||
# Handle safetensors files or other cases where model_index.json is not available
|
||||
# Use appropriate SamplingParams based on pipeline_class_name from registry
|
||||
if os.path.isfile(model_path) and model_path.endswith(".safetensors"):
|
||||
# Determine which sampling params to use based on pipeline_class_name
|
||||
pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
|
||||
|
||||
# Try to get SamplingParams from registry
|
||||
from sglang.multimodal_gen.registry import get_pipeline_config_classes
|
||||
|
||||
config_classes = (
|
||||
@@ -616,7 +626,7 @@ class SamplingParams:
|
||||
user_kwargs = dict(kwargs)
|
||||
user_kwargs.pop("diffusers_kwargs", None)
|
||||
|
||||
user_sampling_params = SamplingParams(*args, **user_kwargs)
|
||||
user_sampling_params = type(sampling_params)(*args, **user_kwargs)
|
||||
# TODO: refactor
|
||||
sampling_params._merge_with_user_params(
|
||||
user_sampling_params, explicit_fields=set(user_kwargs.keys())
|
||||
@@ -982,7 +992,14 @@ class SamplingParams:
|
||||
for field in dataclasses.fields(user_params):
|
||||
field_name = field.name
|
||||
user_value = getattr(user_params, field_name)
|
||||
default_class_value = getattr(SamplingParams, field_name)
|
||||
if hasattr(SamplingParams, field_name):
|
||||
default_class_value = getattr(SamplingParams, field_name)
|
||||
elif field.default is not dataclasses.MISSING:
|
||||
default_class_value = field.default
|
||||
elif field.default_factory is not dataclasses.MISSING:
|
||||
default_class_value = field.default_factory()
|
||||
else:
|
||||
default_class_value = dataclasses.MISSING
|
||||
|
||||
is_user_modified = user_value != default_class_value or (
|
||||
explicit_fields is not None and field_name in explicit_fields
|
||||
|
||||
@@ -99,6 +99,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
|
||||
from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.ltx_2 import (
|
||||
LTX2SamplingParams,
|
||||
LTX23HQSamplingParams,
|
||||
LTX23SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.mova import (
|
||||
@@ -632,6 +633,11 @@ def _register_configs():
|
||||
lambda path: "ltx-2.3" in path.lower(),
|
||||
],
|
||||
)
|
||||
# register dedicated sampling params for LTX2TwoStageHQPipeline
|
||||
_PIPELINE_CONFIG_REGISTRY.setdefault(
|
||||
"LTX2TwoStageHQPipeline",
|
||||
(LTX2PipelineConfig, LTX23HQSamplingParams),
|
||||
)
|
||||
|
||||
# Hunyuan
|
||||
register_configs(
|
||||
|
||||
@@ -31,6 +31,32 @@ from sglang.multimodal_gen.utils import FlexibleArgumentParser
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _resolve_cli_sampling_params_cls(server_args: ServerArgs) -> type[SamplingParams]:
|
||||
pipeline_class_name = getattr(server_args, "pipeline_class_name", None)
|
||||
if pipeline_class_name:
|
||||
from sglang.multimodal_gen.registry import get_pipeline_config_classes
|
||||
|
||||
config_classes = get_pipeline_config_classes(pipeline_class_name)
|
||||
if config_classes is not None:
|
||||
_, sampling_params_cls = config_classes
|
||||
return sampling_params_cls
|
||||
|
||||
try:
|
||||
from sglang.multimodal_gen.registry import get_model_info
|
||||
|
||||
model_info = get_model_info(
|
||||
server_args.model_path,
|
||||
backend=server_args.backend,
|
||||
model_id=server_args.model_id,
|
||||
)
|
||||
if model_info is not None:
|
||||
return model_info.sampling_param_cls
|
||||
except Exception as exc:
|
||||
logger.debug("Falling back to base SamplingParams for CLI parsing: %s", exc)
|
||||
|
||||
return SamplingParams
|
||||
|
||||
|
||||
def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser):
|
||||
"""Add the arguments for the generate command."""
|
||||
parser.add_argument(
|
||||
@@ -130,6 +156,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
|
||||
args.request_id = "mocked_fake_id_for_offline_generate"
|
||||
|
||||
server_args = ServerArgs.from_cli_args(args, unknown_args)
|
||||
sampling_params_cls = _resolve_cli_sampling_params_cls(server_args)
|
||||
|
||||
sampling_params_kwargs = {}
|
||||
config_file = getattr(args, "config", None)
|
||||
@@ -137,7 +164,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
|
||||
if config_file:
|
||||
config_args = ServerArgs.load_config_file(config_file) or {}
|
||||
sampling_param_fields = {
|
||||
field.name for field in dataclasses.fields(SamplingParams)
|
||||
field.name for field in dataclasses.fields(sampling_params_cls)
|
||||
}
|
||||
sampling_params_kwargs.update(
|
||||
{
|
||||
@@ -147,7 +174,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
|
||||
}
|
||||
)
|
||||
|
||||
sampling_params_kwargs.update(SamplingParams.get_cli_args(args))
|
||||
sampling_params_kwargs.update(sampling_params_cls.get_cli_args(args))
|
||||
_apply_output_file_path_override(args, sampling_params_kwargs)
|
||||
sampling_params_kwargs["request_id"] = generate_request_id()
|
||||
|
||||
|
||||
@@ -36,6 +36,9 @@ from sglang.multimodal_gen.utils import get_mixed_precision_state
|
||||
torch._dynamo.config.recompile_limit = 64
|
||||
|
||||
|
||||
LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024
|
||||
|
||||
|
||||
class BaseLayerWithLoRA(nn.Module):
|
||||
|
||||
def __init__(
|
||||
@@ -176,16 +179,48 @@ class BaseLayerWithLoRA(nn.Module):
|
||||
"""
|
||||
# Merge all LoRA adapters in order
|
||||
for lora_A, lora_B, _, lora_strength in lora_list:
|
||||
lora_delta = self.slice_lora_b_weights(
|
||||
lora_B.to(data)
|
||||
) @ self.slice_lora_a_weights(lora_A.to(data))
|
||||
# Apply lora_alpha / lora_rank scaling for consistency with forward()
|
||||
if self.lora_alpha is not None and self.lora_rank is not None:
|
||||
if self.lora_alpha != self.lora_rank:
|
||||
lora_delta = lora_delta * (self.lora_alpha / self.lora_rank)
|
||||
if lora_delta.dim() > 2:
|
||||
lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
|
||||
data += lora_strength * lora_delta
|
||||
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
|
||||
):
|
||||
scale *= self.lora_alpha / self.lora_rank
|
||||
|
||||
if not isinstance(lora_B_sliced, torch.Tensor):
|
||||
lora_delta = lora_B_sliced @ lora_A_sliced
|
||||
if isinstance(lora_delta, torch.Tensor) and lora_delta.dim() > 2:
|
||||
lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
|
||||
data.add_(lora_delta, alpha=scale)
|
||||
continue
|
||||
|
||||
if lora_A_sliced.dim() > 2 or lora_B_sliced.dim() > 2:
|
||||
lora_delta = lora_B_sliced @ lora_A_sliced
|
||||
if lora_delta.dim() > 2:
|
||||
lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1])
|
||||
data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
|
||||
data_2d.add_(lora_delta, alpha=scale)
|
||||
continue
|
||||
|
||||
data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data
|
||||
lora_B_2d = (
|
||||
lora_B_sliced.reshape(-1, lora_B_sliced.shape[-1])
|
||||
if lora_B_sliced.dim() > 2
|
||||
else lora_B_sliced
|
||||
)
|
||||
|
||||
chunk_rows = max(
|
||||
1,
|
||||
LORA_MERGE_CHUNK_BYTES
|
||||
// (data_2d.shape[-1] * max(1, data_2d.element_size())),
|
||||
)
|
||||
for start in range(0, lora_B_2d.shape[0], chunk_rows):
|
||||
end = min(start + chunk_rows, lora_B_2d.shape[0])
|
||||
chunk_delta = lora_B_2d[start:end] @ lora_A_sliced
|
||||
data_2d[start:end].add_(chunk_delta, alpha=scale)
|
||||
|
||||
@torch.no_grad()
|
||||
def merge_lora_weights(self, strength: float | None = None) -> None:
|
||||
|
||||
@@ -1219,6 +1219,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
return timestep
|
||||
return timestep.amax(dim=tuple(range(1, timestep.ndim)))
|
||||
|
||||
def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor:
|
||||
ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2"))
|
||||
if ltx_variant == "ltx_2_3" and bool(
|
||||
getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False)
|
||||
):
|
||||
return timestep * float(self.timestep_scale_multiplier)
|
||||
return timestep
|
||||
|
||||
def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None:
|
||||
"""Validate TP-related dimension constraints (fail-fast)."""
|
||||
if tp_size < 1:
|
||||
@@ -1663,8 +1671,10 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
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
|
||||
timestep_for_adaln = self._scale_timestep_for_adaln(timestep)
|
||||
audio_timestep_for_adaln = self._scale_timestep_for_adaln(audio_timestep)
|
||||
temb, embedded_timestep = self.adaln_single(
|
||||
timestep.flatten(),
|
||||
timestep_for_adaln.flatten(),
|
||||
hidden_dtype=hidden_states.dtype,
|
||||
)
|
||||
temb = temb.view(batch_size, -1, temb.size(-1))
|
||||
@@ -1673,7 +1683,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
)
|
||||
|
||||
temb_audio, audio_embedded_timestep = self.audio_adaln_single(
|
||||
audio_timestep.flatten(),
|
||||
audio_timestep_for_adaln.flatten(),
|
||||
hidden_dtype=audio_hidden_states.dtype,
|
||||
)
|
||||
temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
|
||||
@@ -1688,8 +1698,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
if prompt_timestep is None
|
||||
else prompt_timestep
|
||||
)
|
||||
prompt_timestep_for_adaln = self._scale_timestep_for_adaln(prompt_timestep)
|
||||
temb_prompt, _ = self.prompt_adaln_single(
|
||||
prompt_timestep.flatten(), hidden_dtype=hidden_states.dtype
|
||||
prompt_timestep_for_adaln.flatten(), hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1))
|
||||
if self.audio_prompt_adaln_single is not None:
|
||||
@@ -1698,8 +1709,11 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
if audio_prompt_timestep is None
|
||||
else audio_prompt_timestep
|
||||
)
|
||||
audio_prompt_timestep_for_adaln = self._scale_timestep_for_adaln(
|
||||
audio_prompt_timestep
|
||||
)
|
||||
temb_audio_prompt, _ = self.audio_prompt_adaln_single(
|
||||
audio_prompt_timestep.flatten(),
|
||||
audio_prompt_timestep_for_adaln.flatten(),
|
||||
hidden_dtype=audio_hidden_states.dtype,
|
||||
)
|
||||
temb_audio_prompt = temb_audio_prompt.view(
|
||||
@@ -1714,8 +1728,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
prompt_timestep,
|
||||
audio_prompt_timestep,
|
||||
)
|
||||
av_ca_video_timestep_for_adaln = self._scale_timestep_for_adaln(
|
||||
av_ca_video_timestep
|
||||
)
|
||||
av_ca_audio_timestep_for_adaln = self._scale_timestep_for_adaln(
|
||||
av_ca_audio_timestep
|
||||
)
|
||||
temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single(
|
||||
av_ca_video_timestep.flatten(), hidden_dtype=hidden_dtype
|
||||
av_ca_video_timestep_for_adaln.flatten(), hidden_dtype=hidden_dtype
|
||||
)
|
||||
temb_ca_scale_shift = temb_ca_scale_shift.view(
|
||||
batch_size, -1, temb_ca_scale_shift.shape[-1]
|
||||
@@ -1723,20 +1743,21 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
|
||||
|
||||
av_ca_gate_factor = self._get_av_ca_gate_timestep_factor()
|
||||
temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single(
|
||||
av_ca_video_timestep.flatten() * av_ca_gate_factor,
|
||||
av_ca_video_timestep_for_adaln.flatten() * av_ca_gate_factor,
|
||||
hidden_dtype=hidden_dtype,
|
||||
)
|
||||
temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1])
|
||||
|
||||
temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single(
|
||||
av_ca_audio_timestep.flatten(), hidden_dtype=audio_hidden_states.dtype
|
||||
av_ca_audio_timestep_for_adaln.flatten(),
|
||||
hidden_dtype=audio_hidden_states.dtype,
|
||||
)
|
||||
temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view(
|
||||
batch_size, -1, temb_ca_audio_scale_shift.shape[-1]
|
||||
)
|
||||
|
||||
temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single(
|
||||
av_ca_audio_timestep.flatten() * av_ca_gate_factor,
|
||||
av_ca_audio_timestep_for_adaln.flatten() * av_ca_gate_factor,
|
||||
hidden_dtype=audio_hidden_states.dtype,
|
||||
)
|
||||
temb_ca_audio_gate = temb_ca_audio_gate.view(
|
||||
|
||||
@@ -6,9 +6,11 @@ import torch
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
|
||||
LTX2PipelineConfig,
|
||||
is_ltx23_native_variant,
|
||||
sync_ltx23_runtime_vae_markers,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.ltx_2 import LTX23HQSamplingParams
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
|
||||
PipelineComponentLoader,
|
||||
@@ -120,12 +122,18 @@ def build_official_ltx2_sigmas(
|
||||
stretch: bool = True,
|
||||
terminal: float = 0.1,
|
||||
default_number_of_tokens: int = MAX_SHIFT_ANCHOR,
|
||||
number_of_tokens: int | None = None,
|
||||
) -> list[float]:
|
||||
sigmas = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32)
|
||||
|
||||
mm = (max_shift - base_shift) / (MAX_SHIFT_ANCHOR - BASE_SHIFT_ANCHOR)
|
||||
b = base_shift - mm * BASE_SHIFT_ANCHOR
|
||||
sigma_shift = float(default_number_of_tokens) * mm + b
|
||||
tokens = (
|
||||
int(number_of_tokens)
|
||||
if number_of_tokens is not None
|
||||
else int(default_number_of_tokens)
|
||||
)
|
||||
sigma_shift = float(tokens) * mm + b
|
||||
|
||||
non_zero_mask = sigmas != 0
|
||||
shifted = torch.where(
|
||||
@@ -136,8 +144,9 @@ def build_official_ltx2_sigmas(
|
||||
|
||||
if stretch:
|
||||
one_minus_z = 1.0 - shifted[non_zero_mask]
|
||||
scale_factor = one_minus_z[-1] / (1.0 - terminal)
|
||||
shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor)
|
||||
if bool(torch.any(one_minus_z != 0)):
|
||||
scale_factor = one_minus_z[-1] / (1.0 - terminal)
|
||||
shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor)
|
||||
|
||||
return shifted[:-1].tolist()
|
||||
|
||||
@@ -148,7 +157,28 @@ class LTX2SigmaPreparationStage(PipelineStage):
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
|
||||
batch.extra["ltx2_phase"] = "stage1"
|
||||
if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config):
|
||||
batch.sigmas = build_official_ltx2_sigmas(int(batch.num_inference_steps))
|
||||
# Resolution-aware sigma shift is only required for the HQ pipeline
|
||||
# (which targets 1080p+ resolutions and was aligned against official
|
||||
# LTX-2.3 HQ sigmas). Legacy one-stage and two-stage LTX-2.3 paths
|
||||
# were baselined against the constant-anchor schedule.
|
||||
if server_args.pipeline_class_name == "LTX2TwoStageHQPipeline":
|
||||
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
|
||||
)
|
||||
latent_width = int(batch.width) // int(
|
||||
server_args.pipeline_config.vae_scale_factor
|
||||
)
|
||||
batch.sigmas = build_official_ltx2_sigmas(
|
||||
int(batch.num_inference_steps),
|
||||
number_of_tokens=latent_num_frames * latent_height * latent_width,
|
||||
)
|
||||
else:
|
||||
batch.sigmas = build_official_ltx2_sigmas(
|
||||
int(batch.num_inference_steps)
|
||||
)
|
||||
else:
|
||||
batch.sigmas = np.linspace(
|
||||
1.0,
|
||||
@@ -171,7 +201,11 @@ def _add_ltx2_front_stages(pipeline: ComposedPipelineBase):
|
||||
)
|
||||
|
||||
|
||||
def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase):
|
||||
def _add_ltx2_stage1_generation_stages(
|
||||
pipeline: ComposedPipelineBase,
|
||||
*,
|
||||
denoising_sampler_name: str = "euler",
|
||||
):
|
||||
pipeline.add_stage(LTX2SigmaPreparationStage())
|
||||
pipeline.add_standard_timestep_preparation_stage(
|
||||
prepare_extra_kwargs=[prepare_ltx2_mu]
|
||||
@@ -191,6 +225,7 @@ def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase):
|
||||
scheduler=pipeline.get_module("scheduler"),
|
||||
vae=pipeline.get_module("vae"),
|
||||
audio_vae=pipeline.get_module("audio_vae"),
|
||||
sampler_name=denoising_sampler_name,
|
||||
pipeline=pipeline,
|
||||
),
|
||||
]
|
||||
@@ -542,18 +577,25 @@ class LTX2TwoStageDeviceManager:
|
||||
module.to("cpu")
|
||||
return
|
||||
|
||||
pin_memory = bool(
|
||||
self.server_args.pin_cpu_memory and torch.get_device_module().is_available()
|
||||
)
|
||||
for name, param in module.named_parameters():
|
||||
snapshot = param_snapshots.get(name)
|
||||
if snapshot is None:
|
||||
raise KeyError(
|
||||
f"Missing CPU parameter snapshot for {module_name}.{name}"
|
||||
snapshot = self._clone_cpu_tensor_snapshot(
|
||||
param.data, pin_memory=pin_memory
|
||||
)
|
||||
param_snapshots[name] = snapshot
|
||||
param.data = snapshot
|
||||
|
||||
for name, buffer in module.named_buffers():
|
||||
snapshot = buffer_snapshots.get(name)
|
||||
if snapshot is None:
|
||||
raise KeyError(f"Missing CPU buffer snapshot for {module_name}.{name}")
|
||||
snapshot = self._clone_cpu_tensor_snapshot(
|
||||
buffer.data, pin_memory=pin_memory
|
||||
)
|
||||
buffer_snapshots[name] = snapshot
|
||||
# Preserve runtime-updated buffers (e.g., lazily built caches) when
|
||||
# releasing back to CPU snapshots.
|
||||
if buffer.device.type == "cuda":
|
||||
@@ -666,6 +708,10 @@ class LTX2TwoStageDeviceManager:
|
||||
class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
pipeline_name = "LTX2TwoStagePipeline"
|
||||
STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
|
||||
STAGE_1_DISTILLED_LORA_STRENGTH = 0.0
|
||||
STAGE_2_DISTILLED_LORA_STRENGTH = 1.0
|
||||
STAGE_1_DENOISING_SAMPLER_NAME = "euler"
|
||||
STAGE_2_DENOISING_SAMPLER_NAME = "euler"
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
@@ -712,6 +758,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
self._stage1_lora_path = server_args.lora_path
|
||||
self._stage1_lora_scale = float(server_args.lora_scale)
|
||||
self._active_lora_phase = None
|
||||
self._active_lora_signature = None
|
||||
self._use_premerged_stage2_transformer = False
|
||||
|
||||
def _initialize_premerged_stage2_transformer(self, server_args: ServerArgs) -> None:
|
||||
@@ -733,7 +780,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
lora_nickname="ltx2_stage2_distilled",
|
||||
lora_path=self._distilled_lora_path,
|
||||
target="transformer_2",
|
||||
strength=1.0,
|
||||
strength=self.STAGE_2_DISTILLED_LORA_STRENGTH,
|
||||
merge_weights=True,
|
||||
)
|
||||
|
||||
@@ -757,16 +804,48 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
"resident",
|
||||
)
|
||||
|
||||
def _can_short_circuit_lora_switch(self, phase: str) -> bool:
|
||||
return (
|
||||
phase in ("stage1", "stage2")
|
||||
and self._use_premerged_stage2_transformer
|
||||
and self._stage1_lora_path is None
|
||||
)
|
||||
def _get_stage_distilled_lora_strength(
|
||||
self, phase: str, batch: Req | None
|
||||
) -> float:
|
||||
if phase == "stage1":
|
||||
default_strength = self.STAGE_1_DISTILLED_LORA_STRENGTH
|
||||
extra_key = "ltx2_distilled_lora_strength_stage_1"
|
||||
elif phase == "stage2":
|
||||
default_strength = self.STAGE_2_DISTILLED_LORA_STRENGTH
|
||||
extra_key = "ltx2_distilled_lora_strength_stage_2"
|
||||
else:
|
||||
raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
|
||||
|
||||
if batch is None:
|
||||
return float(default_strength)
|
||||
|
||||
request_strength = batch.extra.get(extra_key)
|
||||
if request_strength is None:
|
||||
return float(default_strength)
|
||||
return float(request_strength)
|
||||
|
||||
def _can_short_circuit_lora_switch(
|
||||
self, phase: str, batch: Req | None = None
|
||||
) -> bool:
|
||||
distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
|
||||
if phase == "stage1":
|
||||
return (
|
||||
self._use_premerged_stage2_transformer
|
||||
and self._stage1_lora_path is None
|
||||
and distilled_lora_strength == 0.0
|
||||
)
|
||||
if phase == "stage2":
|
||||
return (
|
||||
self._use_premerged_stage2_transformer
|
||||
and self._stage1_lora_path is None
|
||||
and distilled_lora_strength == self.STAGE_2_DISTILLED_LORA_STRENGTH
|
||||
)
|
||||
return False
|
||||
|
||||
def _build_lora_switch_spec(
|
||||
self, phase: str
|
||||
self, phase: str, batch: Req | None = None
|
||||
) -> tuple[list[str], list[str], list[float], list[str]]:
|
||||
distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
|
||||
lora_nicknames: list[str] = []
|
||||
lora_paths: list[str] = []
|
||||
lora_strengths: list[float] = []
|
||||
@@ -778,33 +857,42 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
lora_paths.append(self._stage1_lora_path)
|
||||
lora_strengths.append(self._stage1_lora_scale)
|
||||
lora_targets.append("transformer")
|
||||
if distilled_lora_strength != 0.0:
|
||||
lora_nicknames.append("ltx2_stage1_distilled")
|
||||
lora_paths.append(self._distilled_lora_path)
|
||||
lora_strengths.append(distilled_lora_strength)
|
||||
lora_targets.append("transformer")
|
||||
elif phase == "stage2":
|
||||
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")
|
||||
if distilled_lora_strength != 0.0:
|
||||
lora_nicknames.append("ltx2_stage2_distilled")
|
||||
lora_paths.append(self._distilled_lora_path)
|
||||
lora_strengths.append(distilled_lora_strength)
|
||||
lora_targets.append("transformer")
|
||||
else:
|
||||
raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
|
||||
|
||||
return lora_nicknames, lora_paths, lora_strengths, lora_targets
|
||||
|
||||
def switch_lora_phase(self, phase: str) -> None:
|
||||
if phase == self._active_lora_phase:
|
||||
def switch_lora_phase(self, phase: str, batch: Req | None = None) -> None:
|
||||
distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch)
|
||||
phase_signature = (phase, distilled_lora_strength)
|
||||
if phase_signature == self._active_lora_signature:
|
||||
return
|
||||
|
||||
if self._device_manager.switch_phase(
|
||||
phase
|
||||
) and self._can_short_circuit_lora_switch(phase):
|
||||
) and self._can_short_circuit_lora_switch(phase, batch):
|
||||
self._active_lora_phase = phase
|
||||
self._active_lora_signature = phase_signature
|
||||
return
|
||||
|
||||
lora_nicknames, lora_paths, lora_strengths, lora_targets = (
|
||||
self._build_lora_switch_spec(phase)
|
||||
self._build_lora_switch_spec(phase, batch)
|
||||
)
|
||||
if lora_nicknames:
|
||||
set_lora_kwargs = dict(
|
||||
@@ -830,6 +918,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
self.deactivate_lora_weights(target="transformer")
|
||||
|
||||
self._active_lora_phase = phase
|
||||
self._active_lora_signature = phase_signature
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs):
|
||||
_add_ltx2_front_stages(self)
|
||||
@@ -837,7 +926,10 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
self.add_stage(
|
||||
LTX2LoRASwitchStage(pipeline=self, phase="stage1"),
|
||||
)
|
||||
_add_ltx2_stage1_generation_stages(self)
|
||||
_add_ltx2_stage1_generation_stages(
|
||||
self,
|
||||
denoising_sampler_name=self.STAGE_1_DENOISING_SAMPLER_NAME,
|
||||
)
|
||||
self.add_stages(
|
||||
[
|
||||
LTX2UpsampleStage(
|
||||
@@ -863,10 +955,21 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
|
||||
vae=self.get_module("vae"),
|
||||
audio_vae=self.get_module("audio_vae"),
|
||||
pipeline=self,
|
||||
sampler_name=self.STAGE_2_DENOISING_SAMPLER_NAME,
|
||||
),
|
||||
]
|
||||
)
|
||||
_add_ltx2_decoding_stage(self)
|
||||
|
||||
|
||||
EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline]
|
||||
class LTX2TwoStageHQPipeline(LTX2TwoStagePipeline):
|
||||
pipeline_name = "LTX2TwoStageHQPipeline"
|
||||
pipeline_config_cls = LTX2PipelineConfig
|
||||
sampling_params_cls = LTX23HQSamplingParams
|
||||
STAGE_1_DISTILLED_LORA_STRENGTH = 0.25
|
||||
STAGE_2_DISTILLED_LORA_STRENGTH = 0.5
|
||||
STAGE_1_DENOISING_SAMPLER_NAME = "res2s"
|
||||
STAGE_2_DENOISING_SAMPLER_NAME = "res2s"
|
||||
|
||||
|
||||
EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline, LTX2TwoStageHQPipeline]
|
||||
|
||||
@@ -112,8 +112,16 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
|
||||
vae=None,
|
||||
audio_vae=None,
|
||||
pipeline=None,
|
||||
sampler_name: str = "euler",
|
||||
):
|
||||
super().__init__(transformer, scheduler, vae, audio_vae, pipeline=pipeline)
|
||||
super().__init__(
|
||||
transformer,
|
||||
scheduler,
|
||||
vae,
|
||||
audio_vae,
|
||||
pipeline=pipeline,
|
||||
sampler_name=sampler_name,
|
||||
)
|
||||
self.distilled_sigmas = torch.tensor(distilled_sigmas)
|
||||
|
||||
@staticmethod
|
||||
@@ -171,6 +179,41 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
|
||||
return False
|
||||
return "LTX-2.3" not in str(getattr(server_args, "model_path", ""))
|
||||
|
||||
@staticmethod
|
||||
def _build_stage2_renoise_generator(
|
||||
batch: Req, reference_tensor: torch.Tensor
|
||||
) -> torch.Generator:
|
||||
seeds = getattr(batch, "seeds", None)
|
||||
if seeds:
|
||||
seed = int(seeds[0])
|
||||
else:
|
||||
seed = int(getattr(batch, "seed", 10))
|
||||
device = reference_tensor.device
|
||||
dtype = reference_tensor.dtype
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
video_shape = batch.extra.get("ltx2_stage1_packed_video_shape")
|
||||
audio_shape = batch.extra.get("ltx2_stage1_packed_audio_shape")
|
||||
if video_shape is not None:
|
||||
_ = torch.randn(
|
||||
tuple(video_shape), device=device, dtype=dtype, generator=generator
|
||||
)
|
||||
if audio_shape is not None:
|
||||
_ = torch.randn(
|
||||
tuple(audio_shape), device=device, dtype=dtype, generator=generator
|
||||
)
|
||||
return generator
|
||||
|
||||
@staticmethod
|
||||
def _ltx2_renoise_like(
|
||||
reference_tensor: torch.Tensor, generator: torch.Generator
|
||||
) -> torch.Tensor:
|
||||
return torch.randn(
|
||||
reference_tensor.shape,
|
||||
device=reference_tensor.device,
|
||||
dtype=reference_tensor.dtype,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
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"
|
||||
@@ -191,6 +234,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
|
||||
if self._should_reset_stage2_generators(server_args):
|
||||
self._reset_stage2_generators(batch)
|
||||
noise_scale = float(self.distilled_sigmas[0].item())
|
||||
# HQ pipeline uses a dedicated, deterministic renoise generator seeded
|
||||
# from the request seed and advanced by stage-1 packed shapes to match
|
||||
# official LTX-2.3 HQ output. Legacy two-stage paths were baselined
|
||||
# against `batch.generator`'s natural advance through stage-1, so keep
|
||||
# them on the original `_randn_like_with_batch_generators` sampling.
|
||||
is_hq_pipeline = server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
|
||||
if is_hq_pipeline:
|
||||
video_reference_for_gen = (
|
||||
batch.latents if isinstance(batch.latents, torch.Tensor) else None
|
||||
)
|
||||
if video_reference_for_gen is None:
|
||||
video_reference_for_gen = batch.audio_latents
|
||||
renoise_generator = self._build_stage2_renoise_generator(
|
||||
batch, video_reference_for_gen
|
||||
)
|
||||
else:
|
||||
renoise_generator = None
|
||||
if is_native_ti2v:
|
||||
prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state(
|
||||
latents=batch.latents,
|
||||
@@ -199,34 +259,63 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
|
||||
zero_clean_latent=True,
|
||||
clean_latent_background=batch.ltx2_ti2v_clean_latent_background,
|
||||
)
|
||||
video_noise = self._randn_like_with_batch_generators(
|
||||
prepared_latents, batch
|
||||
)
|
||||
if is_hq_pipeline:
|
||||
video_noise = self._ltx2_renoise_like(
|
||||
prepared_latents, renoise_generator
|
||||
)
|
||||
else:
|
||||
video_noise = self._randn_like_with_batch_generators(
|
||||
prepared_latents, batch
|
||||
)
|
||||
scaled_mask = (
|
||||
denoise_mask.to(device=prepared_latents.device, dtype=torch.float32)
|
||||
* noise_scale
|
||||
)
|
||||
batch.latents = (
|
||||
video_noise * scaled_mask + prepared_latents * (1 - scaled_mask)
|
||||
).to(prepared_latents.dtype)
|
||||
if is_hq_pipeline:
|
||||
batch.latents = (
|
||||
video_noise.float() * scaled_mask
|
||||
+ prepared_latents.float() * (1.0 - scaled_mask)
|
||||
).to(prepared_latents.dtype)
|
||||
else:
|
||||
batch.latents = (
|
||||
video_noise * scaled_mask + prepared_latents * (1 - scaled_mask)
|
||||
).to(prepared_latents.dtype)
|
||||
else:
|
||||
video_noise = self._randn_like_with_batch_generators(batch.latents, batch)
|
||||
batch.latents = (
|
||||
video_noise * noise_scale + batch.latents * (1 - noise_scale)
|
||||
).to(batch.latents.dtype)
|
||||
if is_hq_pipeline:
|
||||
video_noise = self._ltx2_renoise_like(batch.latents, renoise_generator)
|
||||
batch.latents = (
|
||||
video_noise.float() * noise_scale
|
||||
+ batch.latents.float() * (1.0 - noise_scale)
|
||||
).to(batch.latents.dtype)
|
||||
else:
|
||||
video_noise = self._randn_like_with_batch_generators(
|
||||
batch.latents, batch
|
||||
)
|
||||
batch.latents = (
|
||||
video_noise * noise_scale + batch.latents * (1 - noise_scale)
|
||||
).to(batch.latents.dtype)
|
||||
|
||||
if isinstance(batch.audio_latents, torch.Tensor):
|
||||
audio_noise = self._randn_like_with_batch_generators(
|
||||
batch.audio_latents, batch
|
||||
)
|
||||
audio_scaled_mask = (
|
||||
torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32)
|
||||
* noise_scale
|
||||
)
|
||||
batch.audio_latents = (
|
||||
audio_noise * audio_scaled_mask
|
||||
+ batch.audio_latents * (1 - audio_scaled_mask)
|
||||
).to(batch.audio_latents.dtype)
|
||||
if is_hq_pipeline:
|
||||
audio_noise = self._ltx2_renoise_like(
|
||||
batch.audio_latents, renoise_generator
|
||||
)
|
||||
batch.audio_latents = (
|
||||
audio_noise.float() * noise_scale
|
||||
+ batch.audio_latents.float() * (1.0 - noise_scale)
|
||||
).to(batch.audio_latents.dtype)
|
||||
else:
|
||||
audio_noise = self._randn_like_with_batch_generators(
|
||||
batch.audio_latents, batch
|
||||
)
|
||||
audio_scaled_mask = (
|
||||
torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32)
|
||||
* noise_scale
|
||||
)
|
||||
batch.audio_latents = (
|
||||
audio_noise * audio_scaled_mask
|
||||
+ batch.audio_latents * (1 - audio_scaled_mask)
|
||||
).to(batch.audio_latents.dtype)
|
||||
if not is_ltx23_native_variant(
|
||||
server_args.pipeline_config.vae_config.arch_config
|
||||
):
|
||||
@@ -244,8 +333,25 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
|
||||
|
||||
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
|
||||
# HQ pipeline extends the sigma schedule so the final step targets a
|
||||
# small non-zero sigma (0.0011) instead of 0.0, matching official
|
||||
# LTX-2.3 HQ's last-step behavior. Legacy two-stage baselines used the
|
||||
# un-extended schedule (final step goes to 0.0).
|
||||
if (
|
||||
server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
|
||||
and self.distilled_sigmas[-1].item() == 0.0
|
||||
):
|
||||
scheduler_sigmas = torch.cat(
|
||||
[
|
||||
self.distilled_sigmas[:-1],
|
||||
torch.tensor([0.0011, 0.0], dtype=self.distilled_sigmas.dtype),
|
||||
],
|
||||
dim=0,
|
||||
)
|
||||
else:
|
||||
scheduler_sigmas = self.distilled_sigmas
|
||||
self.scheduler.sigmas = scheduler_sigmas.to(distilled_device)
|
||||
self.scheduler.num_inference_steps = num_steps
|
||||
self.scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to(
|
||||
distilled_device
|
||||
|
||||
@@ -155,14 +155,16 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage):
|
||||
latent_shape = server_args.pipeline_config.prepare_latent_shape(
|
||||
batch, batch_size, num_frames
|
||||
)
|
||||
packed_video_shape = self._packed_video_latent_shape(
|
||||
latent_shape, server_args.pipeline_config
|
||||
)
|
||||
latents = randn_tensor(
|
||||
self._packed_video_latent_shape(
|
||||
latent_shape, server_args.pipeline_config
|
||||
),
|
||||
packed_video_shape,
|
||||
generator=generator,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
batch.extra["ltx2_stage1_packed_video_shape"] = tuple(packed_video_shape)
|
||||
|
||||
latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents)
|
||||
if latent_ids is not None:
|
||||
@@ -196,13 +198,14 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage):
|
||||
latent_shape = server_args.pipeline_config.prepare_audio_latent_shape(
|
||||
batch, batch_size, batch.num_frames
|
||||
)
|
||||
|
||||
packed_audio_shape = self._packed_audio_latent_shape(latent_shape)
|
||||
audio_latents = randn_tensor(
|
||||
self._packed_audio_latent_shape(latent_shape),
|
||||
packed_audio_shape,
|
||||
generator=generator,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
batch.extra["ltx2_stage1_packed_audio_shape"] = tuple(packed_audio_shape)
|
||||
else:
|
||||
audio_latents = audio_latents.to(device)
|
||||
audio_latents = server_args.pipeline_config.maybe_pack_audio_latents(
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -47,7 +47,6 @@ class LTX2TextConnectorStage(PipelineStage):
|
||||
|
||||
# Handle CFG: Concatenate negative and positive inputs
|
||||
if batch.do_classifier_free_guidance:
|
||||
|
||||
# Concatenate: [Negative, Positive]
|
||||
prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0)
|
||||
prompt_attention_mask = torch.cat(
|
||||
@@ -57,7 +56,9 @@ class LTX2TextConnectorStage(PipelineStage):
|
||||
# Prepare additive mask for connectors (as per Diffusers implementation)
|
||||
dtype = prompt_embeds.dtype
|
||||
|
||||
additive_attention_mask = (1 - prompt_attention_mask.to(dtype)) * -1000000.0
|
||||
additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to(
|
||||
dtype
|
||||
) * torch.finfo(dtype).max
|
||||
|
||||
# Call connectors
|
||||
# Expects: prompt_embeds, attention_mask, additive_mask=True
|
||||
|
||||
@@ -55,7 +55,7 @@ class LTX2LoRASwitchStage(PipelineStage):
|
||||
raise ValueError(
|
||||
"LTX2LoRASwitchStage requires pipeline.switch_lora_phase()"
|
||||
)
|
||||
switch_fn(self.phase)
|
||||
switch_fn(self.phase, batch=batch)
|
||||
batch.extra["ltx2_phase"] = self.phase
|
||||
return batch
|
||||
|
||||
|
||||
@@ -73,7 +73,7 @@ logger = init_logger(__name__)
|
||||
# GPUs on the faster no-offload default while preserving some headroom.
|
||||
WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
|
||||
LTX2_TWO_STAGE_DEVICE_MODES = ("original", "snapshot", "resident")
|
||||
LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline",)
|
||||
LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline")
|
||||
# H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident.
|
||||
LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130
|
||||
|
||||
|
||||
@@ -26,10 +26,8 @@ logger = init_logger(__name__)
|
||||
# Built-in diffusion model overlay registry.
|
||||
BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
|
||||
"Lightricks/LTX-2.3": {
|
||||
# TODO: consider move to lmsys hf repo
|
||||
"overlay_repo_id": "MickJ/LTX-2.3-overlay",
|
||||
"overlay_revision": "main",
|
||||
"bundled_overlay_subdir": "ltx_2_3",
|
||||
"overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -332,6 +332,16 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [
|
||||
# num_frames=33,
|
||||
# ),
|
||||
# ),
|
||||
DiffusionTestCase(
|
||||
"ltx_2_3_hq_pipeline",
|
||||
DiffusionServerArgs(
|
||||
model_path="Lightricks/LTX-2.3",
|
||||
extras=[
|
||||
"--pipeline-class-name LTX2TwoStageHQPipeline --ltx2-two-stage-device-mode snapshot"
|
||||
],
|
||||
),
|
||||
T2I_sampling_params,
|
||||
),
|
||||
]
|
||||
|
||||
# Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm.
|
||||
@@ -519,6 +529,7 @@ TWO_GPU_CASES = [
|
||||
"ltx_2_3_two_stage_ti2v_2gpus",
|
||||
DiffusionServerArgs(
|
||||
model_path="Lightricks/LTX-2.3",
|
||||
ulysses_degree=2,
|
||||
extras=[
|
||||
"--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original"
|
||||
],
|
||||
@@ -537,6 +548,7 @@ TWO_GPU_CASES = [
|
||||
"ltx_2.3_two_stage_t2v_2gpus",
|
||||
DiffusionServerArgs(
|
||||
model_path="Lightricks/LTX-2.3",
|
||||
ulysses_degree=2,
|
||||
extras=[
|
||||
"--pipeline-class-name LTX2TwoStagePipeline",
|
||||
"--ltx2-two-stage-device-mode original",
|
||||
@@ -619,6 +631,7 @@ TWO_GPU_CASES = [
|
||||
"ltx_2.3_one_stage_ti2v",
|
||||
DiffusionServerArgs(
|
||||
model_path="Lightricks/LTX-2.3",
|
||||
ulysses_degree=2,
|
||||
),
|
||||
TI2V_sampling_params,
|
||||
),
|
||||
|
||||
@@ -2629,6 +2629,47 @@
|
||||
"expected_avg_denoise_ms": 890.17,
|
||||
"expected_median_denoise_ms": 896.09,
|
||||
"estimated_full_test_time_s": 155.3
|
||||
},
|
||||
"ltx_2_3_hq_pipeline": {
|
||||
"stages_ms": {
|
||||
"InputValidationStage": 0.11,
|
||||
"TextEncodingStage": 984.78,
|
||||
"LTX2TextConnectorStage": 30.42,
|
||||
"LTX2HalveResolutionStage": 0.1,
|
||||
"LTX2LoRASwitchStage": 0.01,
|
||||
"LTX2SigmaPreparationStage": 0.36,
|
||||
"TimestepPreparationStage": 21.28,
|
||||
"LTX2AVLatentPreparationStage": 0.13,
|
||||
"LTX2ImageEncodingStage": 0.03,
|
||||
"LTX2AVDenoisingStage": 20227.05,
|
||||
"LTX2UpsampleStage": 157.73,
|
||||
"LTX2RefinementStage": 1676.07,
|
||||
"LTX2AVDecodingStage": 521.04,
|
||||
"per_frame_generation": null
|
||||
},
|
||||
"denoise_step_ms": {
|
||||
"0": 1406.0,
|
||||
"1": 1362.75,
|
||||
"2": 1306.18,
|
||||
"3": 1299.65,
|
||||
"4": 1282.16,
|
||||
"5": 1290.32,
|
||||
"6": 1284.64,
|
||||
"7": 1265.01,
|
||||
"8": 1304.06,
|
||||
"9": 1246.21,
|
||||
"10": 1102.19,
|
||||
"11": 1379.4,
|
||||
"12": 1467.28,
|
||||
"13": 1469.49,
|
||||
"14": 734.96,
|
||||
"15": 547.19,
|
||||
"16": 543.27,
|
||||
"17": 539.0
|
||||
},
|
||||
"expected_e2e_ms": 24150.97,
|
||||
"expected_avg_denoise_ms": 1157.21,
|
||||
"expected_median_denoise_ms": 1287.48
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user