Automated PR - 2026-07-07
This commit is contained in:
github-actions[bot]
@@ -11,17 +11,38 @@ This package provides ready-to-use pipelines for video generation:
|
||||
- RetakePipeline: Regenerate a time region (retake) of an existing video
|
||||
For more detailed components and utilities, import from specific submodules
|
||||
like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
|
||||
Pipeline classes are imported lazily (PEP 562). Importing this package therefore
|
||||
does not eagerly pull in every pipeline module, which keeps `import ltx_pipelines`
|
||||
light and avoids the runpy double-import warning when a pipeline is run as a module
|
||||
(e.g. `python -m ltx_pipelines.distilled`).
|
||||
"""
|
||||
|
||||
from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
|
||||
from ltx_pipelines.distilled import DistilledPipeline
|
||||
from ltx_pipelines.ic_lora import ICLoraPipeline
|
||||
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
|
||||
from ltx_pipelines.lipdub import LipDubPipeline
|
||||
from ltx_pipelines.retake import RetakePipeline
|
||||
from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
|
||||
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
|
||||
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
|
||||
from ltx_pipelines.distilled import DistilledPipeline
|
||||
from ltx_pipelines.ic_lora import ICLoraPipeline
|
||||
from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
|
||||
from ltx_pipelines.lipdub import LipDubPipeline
|
||||
from ltx_pipelines.retake import RetakePipeline
|
||||
from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
|
||||
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
|
||||
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
|
||||
|
||||
# Public name -> module that defines it. Used for lazy resolution in __getattr__.
|
||||
_EXPORTS = {
|
||||
"A2VidPipelineTwoStage": "ltx_pipelines.a2vid_two_stage",
|
||||
"DistilledPipeline": "ltx_pipelines.distilled",
|
||||
"ICLoraPipeline": "ltx_pipelines.ic_lora",
|
||||
"KeyframeInterpolationPipeline": "ltx_pipelines.keyframe_interpolation",
|
||||
"LipDubPipeline": "ltx_pipelines.lipdub",
|
||||
"RetakePipeline": "ltx_pipelines.retake",
|
||||
"T2AOneStagePipeline": "ltx_pipelines.t2a_one_stage",
|
||||
"TI2VidOneStagePipeline": "ltx_pipelines.ti2vid_one_stage",
|
||||
"TI2VidTwoStagesPipeline": "ltx_pipelines.ti2vid_two_stages",
|
||||
}
|
||||
|
||||
__all__ = [
|
||||
"A2VidPipelineTwoStage",
|
||||
@@ -34,3 +55,16 @@ __all__ = [
|
||||
"TI2VidOneStagePipeline",
|
||||
"TI2VidTwoStagesPipeline",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str) -> object:
|
||||
module_path = _EXPORTS.get(name)
|
||||
if module_path is None:
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
value = getattr(importlib.import_module(module_path), name)
|
||||
globals()[name] = value # cache so later lookups skip __getattr__
|
||||
return value
|
||||
|
||||
|
||||
def __dir__() -> list[str]:
|
||||
return sorted({*globals(), *_EXPORTS})
|
||||
|
||||
@@ -13,6 +13,7 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import default_2_stage_arg_parser
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioConditioner,
|
||||
@@ -43,7 +44,7 @@ class A2VidPipelineTwoStage:
|
||||
both video and audio using a distilled LoRA for higher quality output.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
checkpoint_path: str,
|
||||
distilled_lora: list[LoraPathStrengthAndSDOps],
|
||||
@@ -55,17 +56,28 @@ class A2VidPipelineTwoStage:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
self._scheduler = LTX2Scheduler()
|
||||
|
||||
self.prompt_encoder = PromptEncoder(
|
||||
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
|
||||
checkpoint_path,
|
||||
gemma_root,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_conditioner = AudioConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.audio_conditioner = AudioConditioner(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -74,9 +86,10 @@ class A2VidPipelineTwoStage:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -85,11 +98,19 @@ class A2VidPipelineTwoStage:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
def __call__( # noqa: PLR0913
|
||||
self,
|
||||
|
||||
@@ -10,10 +10,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
default_2_stage_distilled_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -26,7 +27,6 @@ from ltx_pipelines.utils.blocks import (
|
||||
from ltx_pipelines.utils.constants import (
|
||||
DISTILLED_SIGMAS,
|
||||
STAGE_2_DISTILLED_SIGMAS,
|
||||
detect_params,
|
||||
)
|
||||
from ltx_pipelines.utils.denoisers import SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import (
|
||||
@@ -56,6 +56,7 @@ class DistilledPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
@@ -67,9 +68,16 @@ class DistilledPipeline:
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.stage = DiffusionStage(
|
||||
self.image_conditioner = ImageConditioner(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -78,12 +86,30 @@ class DistilledPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
distilled_checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
def __call__( # noqa: PLR0913
|
||||
self,
|
||||
@@ -182,8 +208,7 @@ class DistilledPipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path(distilled=True)
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params(distilled=True)
|
||||
parser = default_2_stage_distilled_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
pipeline = DistilledPipeline(
|
||||
|
||||
@@ -0,0 +1,193 @@
|
||||
"""Multi-GPU distilled video runner.
|
||||
Runs :class:`DistilledPipeline` across multiple GPUs with:
|
||||
- **Shared stage** -- sequence parallelism (SP); the same DiffusionStage is
|
||||
reused for both stage 1 (half-res) and stage 2 (full-res), so a single
|
||||
SP wrapping covers both invocations.
|
||||
- **Gemma** -- Accelerate-based parallelization
|
||||
- **VAE** -- distributed decoding
|
||||
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Iterator
|
||||
from multiprocessing import SimpleQueue
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.loader.registry import StateDictRegistry
|
||||
from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import get_video_chunks_number
|
||||
from ltx_core.model.video_vae.tiling import TilingConfig
|
||||
from ltx_core.multigpu.transformer.attention import AttentionManager
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
|
||||
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
|
||||
from ltx_pipelines.distilled import DistilledPipeline
|
||||
from ltx_pipelines.multigpu.controller import MGPUController
|
||||
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner
|
||||
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
|
||||
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
|
||||
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.media_io import encode_video
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Stage 2 (full-res) dominates: 1024x1536, 121 frames ~= 24576 video tokens + audio tokens.
|
||||
_DEFAULT_SP_MAX_TOKENS = 32768
|
||||
# Rank that collects distributed-VAE tiles and encodes the assembled video.
|
||||
_DRIVER_RANK = 0
|
||||
|
||||
|
||||
class DistilledRunner(MGPURunner):
|
||||
"""Distributed :class:`DistilledPipeline`: SP shared stage + Accelerate Gemma + distributed VAE."""
|
||||
|
||||
@torch.inference_mode()
|
||||
def setup(
|
||||
self,
|
||||
*,
|
||||
distilled_checkpoint_path: str,
|
||||
gemma_root: str,
|
||||
spatial_upsampler_path: str,
|
||||
vae_queue: SimpleQueue,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
|
||||
quantization: Callable[[], QuantizationPolicy] | None = None,
|
||||
) -> None:
|
||||
# quantization is a picklable zero-arg builder (built per worker, post-spawn); default fp8-cast.
|
||||
quantization_policy = (
|
||||
quantization() if quantization is not None else _build_fp8_cast_policy(distilled_checkpoint_path)
|
||||
)
|
||||
registry = StateDictRegistry()
|
||||
pipeline = DistilledPipeline(
|
||||
distilled_checkpoint_path=distilled_checkpoint_path,
|
||||
gemma_root=gemma_root,
|
||||
spatial_upsampler_path=spatial_upsampler_path,
|
||||
loras=[],
|
||||
registry=registry,
|
||||
quantization=quantization_policy,
|
||||
compilation_config=compilation_config,
|
||||
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
|
||||
)
|
||||
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
|
||||
|
||||
# Shared stage: sequence parallelism (covers both stage 1 and stage 2 invocations).
|
||||
model_cfg = pipeline.stage._transformer_builder.model_config().get("transformer", {})
|
||||
attn_mgr = AttentionManager(
|
||||
max_tokens=sp_max_tokens,
|
||||
num_heads=model_cfg["num_attention_heads"],
|
||||
head_dim=model_cfg["attention_head_dim"],
|
||||
tensor_dtype=pipeline.dtype,
|
||||
group=self.groups.transformer_group,
|
||||
)
|
||||
pipeline.stage._transformer_builder = SequenceParallelBuilder(
|
||||
inner=pipeline.stage._transformer_builder,
|
||||
attn_mgr=attn_mgr,
|
||||
registry=registry,
|
||||
tracker=tracker,
|
||||
)
|
||||
|
||||
# Accelerate Gemma parallelization.
|
||||
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
|
||||
gemma_root_path=gemma_root,
|
||||
gemma_group=self.groups.gemma_group,
|
||||
broadcast_group=self.groups.transformer_group,
|
||||
registry=registry,
|
||||
src_rank=_DRIVER_RANK,
|
||||
dtype=pipeline.dtype,
|
||||
)
|
||||
|
||||
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
|
||||
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
|
||||
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
|
||||
vae_tiling = TileCountConfig(
|
||||
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
|
||||
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
|
||||
)
|
||||
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder( # type: ignore[assignment]
|
||||
inner=pipeline.video_decoder._decoder_builder,
|
||||
queue=vae_queue,
|
||||
vae_group=self.groups.vae_group,
|
||||
vae_tiling=vae_tiling,
|
||||
driver_rank=_DRIVER_RANK,
|
||||
registry=registry,
|
||||
)
|
||||
|
||||
self._pipeline = pipeline
|
||||
|
||||
@torch.inference_mode()
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
output_path: str,
|
||||
prompt: str,
|
||||
seed: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
frame_rate: float,
|
||||
images: list[Any] | None = None,
|
||||
) -> Iterator[str | None]:
|
||||
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
|
||||
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
|
||||
video, audio = self._pipeline(
|
||||
prompt=prompt,
|
||||
seed=seed,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
images=images or [],
|
||||
tiling_config=None,
|
||||
)
|
||||
if dist.get_rank() != _DRIVER_RANK:
|
||||
yield None # workers: nothing to encode
|
||||
return
|
||||
encode_video(
|
||||
video=video,
|
||||
fps=frame_rate,
|
||||
audio=audio,
|
||||
output_path=output_path,
|
||||
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
|
||||
)
|
||||
yield output_path
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from ltx_pipelines.utils.args import (
|
||||
default_2_stage_distilled_arg_parser,
|
||||
resolve_cli_params,
|
||||
)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
|
||||
params = resolve_cli_params(distilled=True)
|
||||
args = default_2_stage_distilled_arg_parser(params=params).parse_args()
|
||||
|
||||
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
|
||||
controller = MGPUController(DistilledRunner)
|
||||
controller.start(
|
||||
distilled_checkpoint_path=args.distilled_checkpoint_path,
|
||||
gemma_root=args.gemma_root,
|
||||
spatial_upsampler_path=args.spatial_upsampler_path,
|
||||
vae_queue=vae_queue,
|
||||
compilation_config=args.compile,
|
||||
)
|
||||
try:
|
||||
for _ in controller.stream(
|
||||
output_path=args.output_path,
|
||||
prompt=args.prompt,
|
||||
seed=args.seed,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
frame_rate=args.frame_rate,
|
||||
images=args.images,
|
||||
):
|
||||
pass # drive the job to completion; the runner writes the file as a side effect
|
||||
finally:
|
||||
controller.shutdown()
|
||||
@@ -39,6 +39,7 @@ from ltx_core.conditioning import (
|
||||
ConditioningItem,
|
||||
VideoConditionByReferenceLatent,
|
||||
)
|
||||
from ltx_core.devices import empty_device_cache
|
||||
from ltx_core.hdr import apply_hdr_decode_postprocess
|
||||
from ltx_core.loader import LoraPathStrengthAndSDOps
|
||||
from ltx_core.loader.registry import Registry
|
||||
@@ -49,6 +50,7 @@ from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
from ltx_core.types import VideoLatentShape, VideoPixelShape
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
DiffusionStage,
|
||||
ImageConditioner,
|
||||
@@ -199,7 +201,7 @@ class HDRICLoraPipeline:
|
||||
Tonemapping and EXR saving are the caller's responsibility.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
distilled_checkpoint_path: str,
|
||||
spatial_upsampler_path: str,
|
||||
@@ -211,6 +213,7 @@ class HDRICLoraPipeline:
|
||||
hdr_lora_config: HdrLoraConfig | None = None,
|
||||
tiled_vae_encode_pixel_threshold: int = TILED_VAE_ENCODE_PIXEL_THRESHOLD,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -252,17 +255,14 @@ class HDRICLoraPipeline:
|
||||
f.get_tensor("audio_context"),
|
||||
)
|
||||
|
||||
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.image_conditioner = ImageConditioner(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
loras=loras,
|
||||
quantization=quantization,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -270,11 +270,33 @@ class HDRICLoraPipeline:
|
||||
quantization=quantization,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
loras=loras,
|
||||
quantization=quantization,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
distilled_checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
# HDR config: explicit override, or auto-detect from LoRA metadata.
|
||||
if hdr_lora_config is not None:
|
||||
@@ -721,7 +743,7 @@ def _process_single_video( # noqa: PLR0913
|
||||
|
||||
del hdr_video
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
empty_device_cache()
|
||||
|
||||
if not skip_mp4:
|
||||
# Wait for EXR saves to finish before encoding.
|
||||
|
||||
@@ -16,12 +16,13 @@ from ltx_pipelines.iclora_utils import (
|
||||
read_lora_reference_downscale_factor,
|
||||
read_lora_reference_temporal_scale_factor,
|
||||
)
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
VideoConditioningAction,
|
||||
VideoMaskConditioningAction,
|
||||
default_2_stage_distilled_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -34,7 +35,6 @@ from ltx_pipelines.utils.blocks import (
|
||||
from ltx_pipelines.utils.constants import (
|
||||
DISTILLED_SIGMAS,
|
||||
STAGE_2_DISTILLED_SIGMAS,
|
||||
detect_params,
|
||||
)
|
||||
from ltx_pipelines.utils.denoisers import SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
|
||||
@@ -64,6 +64,7 @@ class ICLoraPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
@@ -75,9 +76,16 @@ class ICLoraPipeline:
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.image_conditioner = ImageConditioner(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -86,8 +94,9 @@ class ICLoraPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -96,12 +105,30 @@ class ICLoraPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
distilled_checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
# Read reference scale factors from LoRA metadata.
|
||||
# IC-LoRAs trained with scaled reference videos store these factors
|
||||
@@ -344,8 +371,7 @@ class ICLoraPipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path(distilled=True)
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params(distilled=True)
|
||||
parser = default_2_stage_distilled_arg_parser(params=params)
|
||||
parser.add_argument(
|
||||
"--video-conditioning",
|
||||
|
||||
@@ -16,10 +16,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio, VideoPixelShape
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
default_2_stage_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -31,7 +32,6 @@ from ltx_pipelines.utils.blocks import (
|
||||
)
|
||||
from ltx_pipelines.utils.constants import (
|
||||
STAGE_2_DISTILLED_SIGMAS,
|
||||
detect_params,
|
||||
)
|
||||
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import (
|
||||
@@ -53,7 +53,7 @@ class KeyframeInterpolationPipeline:
|
||||
as the upsampled video already has good quality and just needs refinement.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
checkpoint_path: str,
|
||||
distilled_lora: list[LoraPathStrengthAndSDOps],
|
||||
@@ -65,16 +65,25 @@ class KeyframeInterpolationPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
self._scheduler = LTX2Scheduler()
|
||||
|
||||
self.prompt_encoder = PromptEncoder(
|
||||
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
|
||||
checkpoint_path,
|
||||
gemma_root,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -83,9 +92,10 @@ class KeyframeInterpolationPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
stage_2_loras = (*tuple(loras), *tuple(distilled_lora))
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -94,12 +104,22 @@ class KeyframeInterpolationPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
def __call__( # noqa: PLR0913
|
||||
self,
|
||||
@@ -235,8 +255,7 @@ class KeyframeInterpolationPipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path()
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params()
|
||||
parser = default_2_stage_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
pipeline = KeyframeInterpolationPipeline(
|
||||
|
||||
@@ -21,10 +21,11 @@ from ltx_pipelines.iclora_utils import (
|
||||
append_ic_lora_reference_video_conditionings,
|
||||
read_lora_reference_downscale_factor,
|
||||
)
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
detect_checkpoint_path,
|
||||
lipdub_arg_parser,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioConditioner,
|
||||
@@ -35,7 +36,7 @@ from ltx_pipelines.utils.blocks import (
|
||||
VideoDecoder,
|
||||
VideoUpsampler,
|
||||
)
|
||||
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS, detect_params
|
||||
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, STAGE_2_DISTILLED_SIGMAS
|
||||
from ltx_pipelines.utils.denoisers import SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import assert_resolution, combined_image_conditionings, get_device
|
||||
from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video, get_videostream_metadata
|
||||
@@ -62,6 +63,7 @@ class LipDubPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
@@ -75,15 +77,23 @@ class LipDubPipeline:
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_conditioner = AudioConditioner(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage = DiffusionStage(
|
||||
self.stage = DiffusionStage.from_checkpoint(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -92,12 +102,30 @@ class LipDubPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.upsampler = VideoUpsampler(
|
||||
distilled_checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
distilled_checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
distilled_checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.reference_downscale_factor = read_lora_reference_downscale_factor(ic_lora.path)
|
||||
|
||||
def _create_stage_conditionings(
|
||||
@@ -291,8 +319,7 @@ def patchify_lipdub_audio_reference_latent(
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path(distilled=True)
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params(distilled=True)
|
||||
parser = lipdub_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
"""Multi-GPU inference controller.
|
||||
Public API:
|
||||
- ``MGPUController``: controller-driven persistent multi-GPU fleet (start / stream / drain / shutdown)
|
||||
- ``MGPURunner``: base class implemented per pipeline (setup + __call__; __call__ is a generator)
|
||||
- ``Stream``: handle returned by ``controller.stream(...)``; iterate it for each rank's yields as they arrive
|
||||
- ``RunnerError``: raised by a runner for a recoverable, symmetric failure
|
||||
- ``SymmetricRunnerError`` / ``AsymmetricRunnerError``: caller-side exceptions raised from a job's result
|
||||
- ``ControllerBusyError``: ``stream()`` called while a previous job is still uncollected
|
||||
- ``NCCLGroups``: per-component NCCL process groups passed to a runner's ``setup``
|
||||
"""
|
||||
|
||||
from ltx_pipelines.multigpu.controller import (
|
||||
AsymmetricRunnerError,
|
||||
ControllerBusyError,
|
||||
MGPUController,
|
||||
Stream,
|
||||
SymmetricRunnerError,
|
||||
)
|
||||
from ltx_pipelines.multigpu.nccl_groups import NCCLGroups
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
|
||||
|
||||
__all__ = [
|
||||
"AsymmetricRunnerError",
|
||||
"ControllerBusyError",
|
||||
"MGPUController",
|
||||
"MGPURunner",
|
||||
"NCCLGroups",
|
||||
"RunnerError",
|
||||
"Stream",
|
||||
"SymmetricRunnerError",
|
||||
]
|
||||
@@ -0,0 +1,68 @@
|
||||
"""Store-based broadcast signaling between the relay rank and workers.
|
||||
Workers poll an incrementing counter in the process group's store instead of
|
||||
blocking directly on a collective, so an idle fleet doesn't trip the NCCL
|
||||
watchdog timeout while waiting for the next job.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
class BroadcastCoordinator:
|
||||
"""Coordinates broadcast signaling between the relay and workers via a distributed store.
|
||||
Uses an incrementing counter to signal when a broadcast is ready. Workers
|
||||
poll for counter changes to avoid NCCL timeout issues during idle periods.
|
||||
The counter wraps at a large value to prevent overflow.
|
||||
"""
|
||||
|
||||
def __init__(self, store: dist.Store, is_driver: bool) -> None:
|
||||
"""Initialize the broadcast coordinator.
|
||||
Args:
|
||||
store: Distributed store (TCPStore, PrefixStore, etc.) for coordination across ranks.
|
||||
is_driver: Whether this rank is the relay (signals broadcasts).
|
||||
"""
|
||||
self.store = store
|
||||
self.is_driver = is_driver
|
||||
self.broadcast_key = "dist_pipeline_broadcast_ready"
|
||||
self.signal = 0
|
||||
self.signal_max = 2**31 - 1 # ~2.1 billion calls before wraparound
|
||||
|
||||
if is_driver:
|
||||
self.store.set(self.broadcast_key, str(self.signal))
|
||||
|
||||
self.last_seen_signal = self.get_current_signal() if not is_driver else 0
|
||||
|
||||
@contextmanager
|
||||
def broadcast_context(self) -> Iterator[None]:
|
||||
"""Signal a broadcast (driver only).
|
||||
Increments the signal counter on entry to notify workers that a
|
||||
broadcast is ready. The counter stays incremented for change detection.
|
||||
Raises:
|
||||
RuntimeError: If called on a non-driver rank.
|
||||
"""
|
||||
if not self.is_driver:
|
||||
raise RuntimeError("broadcast_context can only be called on driver rank")
|
||||
self.signal = (self.signal + 1) % self.signal_max
|
||||
self.store.set(self.broadcast_key, str(self.signal))
|
||||
yield
|
||||
|
||||
def wait_for_signal_change(self, poll_interval: float = 0.01) -> None:
|
||||
"""Wait until the signal changes from the last seen value (worker only).
|
||||
Args:
|
||||
poll_interval: Seconds to sleep between polls (default: 0.01).
|
||||
"""
|
||||
while True:
|
||||
current = self.get_current_signal()
|
||||
if current != self.last_seen_signal:
|
||||
self.last_seen_signal = current
|
||||
return
|
||||
time.sleep(poll_interval)
|
||||
|
||||
def get_current_signal(self) -> int:
|
||||
"""Get the current signal value from the store."""
|
||||
return int(self.store.get(self.broadcast_key).decode("utf-8"))
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Batch-parallel Gemma text encoder builder.
|
||||
Each rank materialises a full :class:`GemmaTextEncoder` on its own local
|
||||
CUDA device via the standard :class:`SingleGPUModelBuilder` pipeline (the
|
||||
same code path used by the non-MGPU pipelines). No Accelerate, no
|
||||
``device_map``, no per-layer dispatch hooks.
|
||||
Every ``build()`` reconstructs the encoder through
|
||||
:class:`SingleGPUModelBuilder`: a fresh meta module is created, the
|
||||
``GEMMA_MODEL_OPS`` chain re-runs (recomputing the rotary / position
|
||||
buffers that live outside the safetensors file), and the trained weights
|
||||
are bound from the provided :class:`Registry`. The registry caches the
|
||||
loaded state dict so subsequent calls skip disk I/O while still rebuilding
|
||||
the module tree -- mirroring the rebuild logic of
|
||||
:class:`AccelerateGemmaBuilder` on this branch. Encoder-instance caching is
|
||||
intentionally left out; it will arrive later as a global builder refactor.
|
||||
The result is wrapped in :class:`BatchParallelGemmaWrapper`, which
|
||||
partitions prompt lists across ranks in ``encode`` and routes
|
||||
non-deterministic sampling (``enhance_t2v`` / ``enhance_i2v``) through a
|
||||
single ``src_rank``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.loader.primitives import BuilderProtocol
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
|
||||
from ltx_core.multigpu.gemma.batch_parallel_wrapper import BatchParallelGemmaWrapper
|
||||
from ltx_core.text_encoders.gemma import (
|
||||
GEMMA_LLM_KEY_OPS,
|
||||
GEMMA_MODEL_OPS,
|
||||
GemmaTextEncoderConfigurator,
|
||||
module_ops_from_gemma_root,
|
||||
)
|
||||
from ltx_core.utils import find_matching_file
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BatchParallelGemmaBuilder(BuilderProtocol):
|
||||
"""Per-rank Gemma replica builder for the batch-parallel encode path.
|
||||
Mirrors the inline Gemma builder construction inside
|
||||
:class:`PromptEncoder` (single-GPU path) and adds the MGPU wiring --
|
||||
broadcast group + source rank for non-deterministic methods. Each
|
||||
``build()`` reconstructs the encoder via :class:`SingleGPUModelBuilder`;
|
||||
the registry caches the state dict so only disk I/O is skipped across
|
||||
calls.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gemma_root_path: str,
|
||||
broadcast_group: dist.ProcessGroup | None,
|
||||
registry: Registry,
|
||||
*,
|
||||
src_rank: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
model_folder = find_matching_file(gemma_root_path, "model*.safetensors").parent
|
||||
weight_paths = tuple(str(p) for p in model_folder.rglob("*.safetensors"))
|
||||
self._inner = Builder(
|
||||
model_path=weight_paths,
|
||||
model_class_configurator=GemmaTextEncoderConfigurator,
|
||||
model_sd_ops=GEMMA_LLM_KEY_OPS,
|
||||
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(gemma_root_path)),
|
||||
registry=registry,
|
||||
)
|
||||
self._broadcast_group = broadcast_group
|
||||
self._src_rank = src_rank
|
||||
self._dtype = dtype
|
||||
|
||||
def model_config(self) -> dict:
|
||||
return {}
|
||||
|
||||
def build(
|
||||
self,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
) -> BatchParallelGemmaWrapper:
|
||||
dtype = dtype or self._dtype
|
||||
encoder = self._inner.build(device=device, dtype=dtype).eval()
|
||||
return BatchParallelGemmaWrapper(
|
||||
encoder=encoder,
|
||||
broadcast_group=self._broadcast_group,
|
||||
src_rank=dist.get_group_rank(self._broadcast_group, self._src_rank),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
@@ -0,0 +1,382 @@
|
||||
"""Synchronous multi-GPU controller.
|
||||
Spawns one worker process per GPU. Rank 0 is the *relay*: the controller hands it a job over a
|
||||
queue and it NCCL-broadcasts the job to every rank. All ranks run the user's runner -- a
|
||||
*generator* -- in SPMD lockstep. `stream(**kwargs)` dispatches and returns a `Stream` you iterate;
|
||||
each yielded value is forwarded straight to you, one element per yield, as it comes off the result
|
||||
queue (no gathering across ranks). One job at a time -- no job queue, no pipelining.
|
||||
Constraints and contract:
|
||||
- Single machine only: MASTER_ADDR is localhost, RANK == LOCAL_RANK, one rank per GPU. By default
|
||||
rank r runs on cuda:r; pass `devices=[...]` to place the fleet on a specific physical GPU subset
|
||||
(rank r -> cuda:devices[r]), e.g. to run two controllers side by side on disjoint GPUs.
|
||||
- The runner is a generator and runs in SPMD LOCKSTEP. Yields are forwarded individually, not
|
||||
gathered, so each rank's yields appear as their own stream elements (in result-queue order). Only
|
||||
the ranks' terminals are collected: once all have ended, the controller classifies and ends iteration.
|
||||
- Job kwargs (non-tensor parts) cross the queue and are pickled again through the NCCL broadcast;
|
||||
yielded values ride only the result queue (pickled once). All must be picklable and small.
|
||||
Tensors are the exception: pass them as top-level kwargs and the relay broadcasts them over NCCL
|
||||
instead of pickling (see `stream`); tensors inside a yielded value ride the result queue by
|
||||
shared memory / CUDA IPC.
|
||||
- Each job belongs to the thread that dispatched it: only that thread may iterate it, enforced in
|
||||
`Stream`.
|
||||
- One job in flight at a time: consume the `Stream` to the end before dispatching the next.
|
||||
Abandoning it is NOT cleaned up -- the next `stream()` raises `ControllerBusyError` until you
|
||||
`stream.drain()` or `shutdown()`.
|
||||
- A runner that *raises* an unexpected exception kills the controller (a desynced NCCL collective
|
||||
cannot be unwound; make a new one). For a recoverable failure the runner raises a `RunnerError`
|
||||
(or a `ValueError`, which the controller turns into one) identically on every rank: the worker loop
|
||||
catches it, the fleet survives, and iterating the `Stream` re-raises it as `SymmetricRunnerError`
|
||||
-- or `AsymmetricRunnerError` if only some ranks raised.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator, Sequence
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.multiprocessing as torch_mp
|
||||
from torch.distributed.elastic.multiprocessing.api import DefaultLogsSpecs, LogsSpecs, Std
|
||||
|
||||
from ltx_pipelines.multigpu.fleet import _POLL_S, _Channels, _Job, _RunnersFleet
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEFAULT_INIT_TIMEOUT = timedelta(minutes=30)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Errors raised to the caller: a busy dispatch, or a runner failure classified across ranks.
|
||||
# =============================================================================
|
||||
class SymmetricRunnerError(Exception):
|
||||
"""Every rank raised a RunnerError -- the contract was honored. Recoverable: fix the
|
||||
input and retry; the controller is still alive."""
|
||||
|
||||
def __init__(self, errors: list[RunnerError]) -> None:
|
||||
self.errors = errors
|
||||
super().__init__(errors[0].message)
|
||||
|
||||
|
||||
class AsymmetricRunnerError(Exception):
|
||||
"""Some ranks raised a RunnerError and some finished cleanly. The fleet is provably healthy
|
||||
(a mix of terminals means every rank ran to completion and reported), but the runner's error
|
||||
path is non-deterministic across ranks -- a latent hang risk. Loud by design; does NOT kill
|
||||
the fleet."""
|
||||
|
||||
def __init__(self, terminals: list[Any]) -> None:
|
||||
self.terminals = terminals # the per-rank terminals (RunnerErrors mixed with clean StopIterations)
|
||||
super().__init__("runner raised RunnerError on some ranks but not all")
|
||||
|
||||
|
||||
class ControllerBusyError(RuntimeError):
|
||||
"""`stream()` was called while a job is still in flight (uncollected).
|
||||
`stream` returns IMMEDIATELY, so it raises rather than silently draining the in-flight job
|
||||
(which would block for the full job, possibly one owned by another thread). The in-flight job belongs
|
||||
to the thread that dispatched it: consume it (`try: ... finally: stream.drain()`) so you never
|
||||
wedge yourself, and let a concurrent caller that loses the dispatch race simply bounce.
|
||||
"""
|
||||
|
||||
def __init__(self, job_id: int) -> None:
|
||||
self.job_id = job_id # the in-flight job's id, for catch-site logging
|
||||
super().__init__(f"MGPU job {job_id} is still in flight; consume it (stream.drain()) first.")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The caller's streaming handle: a thin wrapper over the collector generator.
|
||||
# Iterating it drives the job and forwards each rank's yields as they arrive.
|
||||
# =============================================================================
|
||||
class Stream:
|
||||
"""The iterable returned by `stream()` and the handle to one in-flight job. Each element is one
|
||||
rank's yielded value, forwarded as it comes off the result queue -- yields are NOT gathered
|
||||
across ranks. When every rank's generator has finished, iteration ends: all clean -> the
|
||||
`StopIteration` value is the per-rank list of returns (normally Nones, since results are yielded
|
||||
-- `result = yield from stream` to read it); all `RunnerError` -> `SymmetricRunnerError`; a mix
|
||||
-> `AsymmetricRunnerError`.
|
||||
The job belongs to the thread that dispatched it: only that thread may iterate (or `drain`) this
|
||||
Stream -- a second thread advancing the same ordered collection is the one hazard forbidden here.
|
||||
Consume the Stream to completion before the next `stream()`; abandoning it leaves the job in
|
||||
flight (the next `stream()` raises `ControllerBusyError` until it is drained or shut down). The
|
||||
controller does not clean up after you -- the recommended pattern is `try: ... finally:
|
||||
stream.drain()`.
|
||||
"""
|
||||
|
||||
def __init__(self, pump: Iterator[Any], job_thread: int, job_id: int) -> None:
|
||||
self._pump = pump # generator: yields each rank's values, returns the per-rank returns
|
||||
self._job_thread = job_thread # the thread that dispatched this job; only it may iterate/drain
|
||||
self.job_id = job_id # this job's id (for messages/debugging)
|
||||
|
||||
def __iter__(self) -> Stream:
|
||||
return self
|
||||
|
||||
def __next__(self) -> Any: # noqa: ANN401
|
||||
if threading.get_ident() != self._job_thread:
|
||||
raise RuntimeError(
|
||||
f"Stream for job {self.job_id} is single-threaded: dispatched on thread "
|
||||
f"{self._job_thread}, iterated from thread {threading.get_ident()} -- a job belongs "
|
||||
f"to its dispatching thread."
|
||||
)
|
||||
# The pump yields each rank's value directly; once all ranks have ended it raises
|
||||
# StopIteration(per-rank returns), or Symmetric/Asymmetric if any rank raised.
|
||||
return next(self._pump)
|
||||
|
||||
def drain(self) -> None:
|
||||
"""Exhaust the Stream and free the controller, discarding any unconsumed yields. drain() must
|
||||
be called from the same thread that called stream() (it iterates the Stream, so the owner
|
||||
check applies); the recommended pattern is `try: ... finally: stream.drain()`. A recoverable
|
||||
Symmetric/AsymmetricRunnerError is swallowed (cleanup); a dead/timed-out/desynced fleet still
|
||||
surfaces -- you need to know.
|
||||
"""
|
||||
try:
|
||||
for _ in self: # iterate via __next__, so the dispatching-thread check applies
|
||||
pass
|
||||
except (SymmetricRunnerError, AsymmetricRunnerError):
|
||||
pass # cleanup: a recoverable runner error isn't worth surfacing when draining
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The controller.
|
||||
# =============================================================================
|
||||
class MGPUController:
|
||||
"""Persistent one-job-at-a-time GPU fleet.
|
||||
controller = MGPUController(MyRunner, num_gpus=8)
|
||||
controller.start(**setup_kwargs)
|
||||
stream = controller.stream(prompt="...")
|
||||
try:
|
||||
for item in stream: # one element per yield, as it arrives (not gathered)
|
||||
show(item)
|
||||
finally:
|
||||
stream.drain() # free the controller even on early exit -- abandoning wedges it
|
||||
controller.shutdown()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
runner_cls: type[MGPURunner],
|
||||
*,
|
||||
num_gpus: int | None = None,
|
||||
devices: Sequence[int] | None = None,
|
||||
logs_specs: LogsSpecs | None = None,
|
||||
) -> None:
|
||||
"""`num_gpus` uses GPUs 0..num_gpus-1 (default: all). `devices` places the fleet on specific
|
||||
physical GPUs (e.g. `[2, 3]` -> rank r on cuda:devices[r]), so several controllers can share a
|
||||
box on disjoint GPU sets; the two are mutually exclusive. Indices are as the controller sees
|
||||
them; every GPU stays visible to each worker, which simply binds to its assigned one.
|
||||
"""
|
||||
if devices is not None:
|
||||
if num_gpus is not None:
|
||||
raise ValueError("Pass either num_gpus or devices, not both.")
|
||||
if len(devices) == 0 or len(set(devices)) != len(devices):
|
||||
raise ValueError(f"devices must be non-empty and unique: {list(devices)}.")
|
||||
self._runner_cls = runner_cls
|
||||
self._devices = list(devices) if devices is not None else None
|
||||
self._num_gpus = num_gpus
|
||||
self._logs_specs = logs_specs
|
||||
self._spawn_ctx = torch_mp.get_context("spawn")
|
||||
|
||||
self._fleet: _RunnersFleet | None = None
|
||||
self._channels: _Channels | None = None
|
||||
self._next_job_id = 0 # monotonic job-id generator; persists across jobs (powers desync detection)
|
||||
self._inflight: Stream | None = None # the one in-flight job, as its handle; None between jobs
|
||||
|
||||
# baton-lock: guards ONLY the _inflight check-and-set; never held across dispatch / iteration / _collect.
|
||||
self._lock = threading.Lock()
|
||||
self._fatal_error: BaseException | None = None # one-way: set once, then every call raises
|
||||
self._started = False
|
||||
|
||||
@property
|
||||
def is_alive(self) -> bool:
|
||||
"""True while the fleet is up and unpoisoned -- a health check that needs no try/except."""
|
||||
return self._started and self._fatal_error is None
|
||||
|
||||
# ---------------------------------------------------------------- lifecycle
|
||||
def start(
|
||||
self,
|
||||
*,
|
||||
timeout: timedelta = _DEFAULT_INIT_TIMEOUT,
|
||||
**setup_kwargs: Any, # noqa: ANN401
|
||||
) -> None:
|
||||
"""Spawn the fleet, run setup() on every rank, block until all report ready.
|
||||
`timeout` bounds both the NCCL `init_process_group` and the controller's wait for
|
||||
every rank to finish CUDA init, `create_local_nccl_groups`, and `setup()`, so a
|
||||
worker wedged ALIVE turns into a clear error instead of an infinite hang (poll()
|
||||
only sees a process *exit*, never a wedge). It must comfortably exceed your slowest
|
||||
model load; pass a large value to effectively wait forever.
|
||||
"""
|
||||
if self._started:
|
||||
raise RuntimeError("MGPUController.start() called twice")
|
||||
if self._devices is not None:
|
||||
num_gpus = len(self._devices)
|
||||
device_ids: list[int] | None = list(self._devices)
|
||||
else:
|
||||
num_gpus = torch.cuda.device_count() if self._num_gpus is None else self._num_gpus
|
||||
device_ids = None
|
||||
if num_gpus <= 0:
|
||||
raise ValueError(f"No GPUs available: num_gpus={num_gpus}.")
|
||||
self._num_gpus = num_gpus # resolve "all GPUs" to the actual count == world size (one rank per GPU)
|
||||
|
||||
self._channels = _Channels(
|
||||
jobs=self._spawn_ctx.Queue(),
|
||||
results=self._spawn_ctx.Queue(),
|
||||
ready=self._spawn_ctx.Queue(),
|
||||
)
|
||||
self._fleet = _RunnersFleet.spawn(
|
||||
runner_cls=self._runner_cls,
|
||||
setup_kwargs=setup_kwargs,
|
||||
init_timeout=timeout,
|
||||
channels=self._channels,
|
||||
num_gpus=num_gpus,
|
||||
logs_specs=self._logs_specs or DefaultLogsSpecs(tee=Std.ALL),
|
||||
device_ids=device_ids,
|
||||
)
|
||||
try:
|
||||
self._await_ready(num_gpus, timeout.total_seconds())
|
||||
except BaseException:
|
||||
self.shutdown() # tear the half-up fleet down so the caller need not
|
||||
raise
|
||||
self._started = True
|
||||
logger.info("MGPU fleet ready (%d workers).", num_gpus)
|
||||
|
||||
def _await_ready(self, num_gpus: int, timeout: float) -> None:
|
||||
assert self._channels is not None
|
||||
assert self._fleet is not None
|
||||
deadline = time.monotonic() + timeout
|
||||
seen: set[int] = set() # which ranks have checked in (ready.put sends the rank)
|
||||
while len(seen) < num_gpus:
|
||||
died = self._fleet.poll()
|
||||
if died is not None: # a worker exited -- catches crashes, not wedges
|
||||
raise died
|
||||
if time.monotonic() > deadline: # catches the wedges poll() can't see
|
||||
missing = sorted(set(range(num_gpus)) - seen)
|
||||
raise TimeoutError(
|
||||
f"MGPU startup: {len(seen)}/{num_gpus} workers ready after {timeout:.0f}s; "
|
||||
f"ranks {missing} never checked in -- stuck in init_process_group / "
|
||||
f"create_local_nccl_groups / setup()? Check the worker logs."
|
||||
)
|
||||
try:
|
||||
seen.add(self._channels.ready.get(timeout=_POLL_S))
|
||||
except Exception:
|
||||
continue # queue.Empty: re-check death + deadline, then retry
|
||||
|
||||
def shutdown(self, *, graceful_timeout: float = 60.0) -> None:
|
||||
"""Tell the fleet to exit, give it a moment, then make sure it is gone.
|
||||
Both teardown and kill switch -- safe to call from another thread while a job is in flight: it
|
||||
force-kills the fleet, so a thread wedged on the Stream surfaces an error and unwedges (this is
|
||||
how you recover a job you can't drain). A mid-job shutdown waits out `graceful_timeout` before
|
||||
forcing; pass 0 to skip it.
|
||||
"""
|
||||
if self._fleet is None:
|
||||
return
|
||||
if self._fatal_error is None and self._channels is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
self._channels.jobs.put(None) # relay broadcasts the sentinel to all ranks
|
||||
self._fleet.drain(graceful_timeout)
|
||||
self._fleet.terminate()
|
||||
self._fleet = None
|
||||
self._started = False
|
||||
|
||||
# ---------------------------------------------------------------- dispatch
|
||||
def stream(self, *, timeout: float | None = None, **kwargs: Any) -> Stream: # noqa: ANN401
|
||||
"""Dispatch one job and return IMMEDIATELY; collect later by iterating the returned Stream.
|
||||
The workers run the job on their own while the controller does nothing; iterate the Stream
|
||||
for the runner's yields -- each rank's yield is forwarded as its own element, as it arrives
|
||||
(not gathered across ranks). Worker death or a blown timeout (measured from dispatch)
|
||||
surfaces when you come back to iterate, not before.
|
||||
TENSORS ARE TRANSPARENT. Pass them as top-level kwargs (`stream(latent=t, steps=30)`): the
|
||||
tensor rides the queue to the relay by shared memory / CUDA IPC and the relay broadcasts it
|
||||
to every rank over NCCL, so `__call__` receives it on the local GPU. (Only top-level kwargs
|
||||
are broadcast this way; tensors nested in a list/dict ride the pickle path.) To send one
|
||||
back, `yield` it from every rank; the yield rides the result queue, so its tensors come back
|
||||
without pickling.
|
||||
ONE JOB AT A TIME. Consume the Stream to the end before dispatching the next. Abandoning a
|
||||
Stream is NOT cleaned up: the job stays in flight and the next `stream()` raises
|
||||
`ControllerBusyError` until you `stream.drain()` or `shutdown()`. Use `try: ... finally:
|
||||
stream.drain()`.
|
||||
"""
|
||||
if not self._started:
|
||||
raise RuntimeError("MGPUController not started; call start() first")
|
||||
if self._fatal_error is not None:
|
||||
raise RuntimeError("MGPUController is dead; create a new one.") from self._fatal_error
|
||||
assert self._channels is not None
|
||||
job_thread = threading.get_ident() # this job belongs to the dispatching thread (checked in Stream)
|
||||
|
||||
# The baton-lock guards exactly this check-and-set -- nothing else (see the lock's comment).
|
||||
with self._lock:
|
||||
if self._inflight is not None:
|
||||
raise ControllerBusyError(self._inflight.job_id)
|
||||
job_id = self._next_job_id
|
||||
self._next_job_id += 1
|
||||
deadline = None if timeout is None else time.monotonic() + timeout
|
||||
stream = Stream(self._collect(job_id, deadline, timeout), job_thread, job_id)
|
||||
self._inflight = stream # the single source of "a job is in flight"
|
||||
|
||||
self._channels.jobs.put(_Job(job_id=job_id, kwargs=kwargs)) # dispatch OUTSIDE the lock; fleet starts NOW
|
||||
return stream
|
||||
|
||||
@staticmethod
|
||||
def _is_terminal(value: object) -> bool:
|
||||
"""True if `value` is a rank's end marker on the wire: a clean StopIteration or a RunnerError."""
|
||||
return isinstance(value, (StopIteration, RunnerError))
|
||||
|
||||
@staticmethod
|
||||
def _classify_terminal(values: list[Any]) -> list[Any]:
|
||||
"""Classify every rank's terminal (a StopIteration or a RunnerError), ordered by rank. All
|
||||
RunnerError -> SymmetricRunnerError; a mix of RunnerError and clean StopIteration ->
|
||||
AsymmetricRunnerError; all clean -> the per-rank return values (the caller's StopIteration.value).
|
||||
Neither typed error kills the fleet."""
|
||||
errs = [v for v in values if isinstance(v, RunnerError)]
|
||||
if errs and len(errs) == len(values):
|
||||
raise SymmetricRunnerError(errs)
|
||||
if errs:
|
||||
logger.error("asymmetric RunnerError across ranks: %r", values)
|
||||
raise AsymmetricRunnerError(values)
|
||||
return [v.value for v in values] # all StopIteration -> their return values
|
||||
|
||||
def _collect(self, job_id: int, deadline: float | None, timeout: float | None) -> Iterator[Any]:
|
||||
"""Drain the result queue for the in-flight job: forward each yielded value to the caller as
|
||||
soon as it arrives (no gathering), collecting each rank's terminal as it ends. Once every
|
||||
rank has ended, classify them -- all clean -> return the per-rank returns (the caller's
|
||||
StopIteration.value); all/some RunnerError -> Symmetric/Asymmetric. Watches for a dead worker
|
||||
/ blown timeout meanwhile. Runs on the caller's thread when they come back to the Stream.
|
||||
"""
|
||||
assert self._channels is not None
|
||||
assert self._fleet is not None
|
||||
assert self._num_gpus is not None
|
||||
channels, fleet = self._channels, self._fleet
|
||||
n = self._num_gpus
|
||||
terminals: dict[int, Any] = {} # rank -> its end marker (StopIteration | RunnerError)
|
||||
while True:
|
||||
try:
|
||||
msg = channels.results.get(timeout=_POLL_S)
|
||||
except Exception:
|
||||
# Queue empty: nothing ready, so NOW (and only now) check for death / timeout.
|
||||
died = fleet.poll()
|
||||
if died is not None:
|
||||
self._fatal_error = died
|
||||
fleet.terminate()
|
||||
raise died from None
|
||||
if deadline is not None and time.monotonic() > deadline:
|
||||
self._fatal_error = TimeoutError(f"MGPU job {job_id} exceeded {timeout}s")
|
||||
fleet.terminate()
|
||||
raise self._fatal_error from None
|
||||
continue
|
||||
|
||||
if msg.job_id != job_id: # a rank ran a different job than we dispatched -> SPMD desync
|
||||
self._fatal_error = RuntimeError(
|
||||
f"MGPU fleet desync: rank {msg.rank} sent job {msg.job_id}, expected {job_id}."
|
||||
)
|
||||
fleet.terminate()
|
||||
raise self._fatal_error from None
|
||||
|
||||
if not self._is_terminal(msg.value):
|
||||
yield msg.value # forward this rank's yield straight to the caller -- no gathering
|
||||
continue
|
||||
|
||||
terminals[msg.rank] = msg.value # a rank ended; hold its terminal for classification
|
||||
if len(terminals) == n: # every rank has ended -> fleet free; classify and finish
|
||||
self._inflight = None
|
||||
return self._classify_terminal([terminals[r] for r in range(n)])
|
||||
@@ -0,0 +1,87 @@
|
||||
"""Base class for multigpu transformer builders that delegate to an inner builder.
|
||||
Shared boilerplate for SP and TDP builders — both wrap a
|
||||
:class:`SingleGPUModelBuilder` and forward the ``ModelBuilderProtocol`` surface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, Generic, TypeVar
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.fuse_loras import FuseRule
|
||||
from ltx_core.loader.module_ops import ModuleOps
|
||||
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.loader.sd_ops import SDOps
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
InnerModelT = TypeVar("InnerModelT", bound=torch.nn.Module)
|
||||
|
||||
|
||||
class DelegatingBuilder(Generic[InnerModelT]):
|
||||
"""Thin wrapper that delegates all ``ModelBuilderProtocol`` accessors to *inner*.
|
||||
``InnerModelT`` is the type produced by the inner builder. Subclasses only
|
||||
need to implement ``__init__`` and ``build`` (whose return type may differ).
|
||||
"""
|
||||
|
||||
_inner: Builder[InnerModelT]
|
||||
|
||||
# -- delegated properties / with_* methods --------------------------------
|
||||
|
||||
@property
|
||||
def checkpoint(self) -> str | tuple[str, ...]:
|
||||
return self._inner.checkpoint
|
||||
|
||||
@property
|
||||
def model_sd_ops(self) -> SDOps | None:
|
||||
return self._inner.model_sd_ops
|
||||
|
||||
def with_sd_ops(self, sd_ops: SDOps | None) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_sd_ops(sd_ops)
|
||||
return clone
|
||||
|
||||
@property
|
||||
def module_ops(self) -> tuple[ModuleOps, ...]:
|
||||
return self._inner.module_ops
|
||||
|
||||
def with_module_ops(self, module_ops: tuple[ModuleOps, ...]) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_module_ops(module_ops)
|
||||
return clone
|
||||
|
||||
@property
|
||||
def loras(self) -> tuple[LoraPathStrengthAndSDOps, ...]:
|
||||
return self._inner.loras
|
||||
|
||||
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_loras(loras)
|
||||
return clone
|
||||
|
||||
@property
|
||||
def registry(self) -> Registry:
|
||||
return self._inner.registry
|
||||
|
||||
def with_registry(self, registry: Registry) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_registry(registry)
|
||||
return clone
|
||||
|
||||
def with_lora_load_device(self, device: torch.device) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_lora_load_device(device)
|
||||
return clone
|
||||
|
||||
def with_fuse_rule(self, fuse_rule: FuseRule) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_fuse_rule(fuse_rule)
|
||||
return clone
|
||||
|
||||
def model_config(self) -> dict:
|
||||
return self._inner.model_config()
|
||||
@@ -0,0 +1,375 @@
|
||||
"""Worker fleet, wire protocol, and SPMD job execution driven by the MGPU controller.
|
||||
Everything the controller (``MGPUController`` in ``controller.py``) drives lives here: the on-the-wire
|
||||
payloads, the per-job NCCL input broadcast (``_RankLink`` / ``_run_job``), the worker entrypoint +
|
||||
loops, the persistent worker fleet (``_RunnersFleet``), and ``_RunnerShipper`` (ships the runner
|
||||
class to workers by value). The runner contract it executes (``MGPURunner`` / ``RunnerError``)
|
||||
lives in ``runner.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from multiprocessing import Queue
|
||||
from typing import Any
|
||||
|
||||
import cloudpickle
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed.elastic.multiprocessing import start_processes
|
||||
from torch.distributed.elastic.multiprocessing.api import LogsSpecs
|
||||
from torch.distributed.elastic.multiprocessing.errors import ProcessFailure, record
|
||||
|
||||
from ltx_pipelines.multigpu._broadcast import BroadcastCoordinator
|
||||
from ltx_pipelines.multigpu.nccl_groups import create_local_nccl_groups
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner, RunnerError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_RELAY_RANK = 0
|
||||
_POLL_S = 0.2 # how often the controller re-checks "did a worker die?" while waiting for a result
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The wire: what crosses the queues between controller and workers.
|
||||
# =============================================================================
|
||||
@dataclass
|
||||
class _TensorPlaceholder:
|
||||
"""Marker left in a kwargs/return dict where a tensor was lifted out for separate
|
||||
transport. Carries the shape/dtype so receivers can preallocate before NCCL fills it.
|
||||
Picklable and tiny -- this is what rides the object collective in the tensor's place.
|
||||
"""
|
||||
|
||||
idx: int # position in the lifted-out tensor list
|
||||
shape: tuple[int, ...]
|
||||
dtype: Any # torch.dtype
|
||||
|
||||
|
||||
def _replace_tensors_by_placeholders(d: dict[str, Any]) -> tuple[dict[str, Any], list[torch.Tensor]]:
|
||||
"""Split a dict into (skeleton, tensors): top-level Tensor values become _TensorPlaceholders.
|
||||
Top level only -- a tensor buried inside a list or nested dict is left alone and
|
||||
will ride the pickle path. Keep tensors as direct kwargs/return values.
|
||||
"""
|
||||
skeleton: dict[str, Any] = {}
|
||||
tensors: list[torch.Tensor] = []
|
||||
for k, v in d.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
skeleton[k] = _TensorPlaceholder(len(tensors), tuple(v.shape), v.dtype)
|
||||
tensors.append(v)
|
||||
else:
|
||||
skeleton[k] = v
|
||||
return skeleton, tensors
|
||||
|
||||
|
||||
def _fill_tensors_into_placeholders(skeleton: dict[str, Any], tensors: list[torch.Tensor]) -> dict[str, Any]:
|
||||
"""Inverse of _replace_tensors_by_placeholders: put the tensors back where the _TensorPlaceholders are."""
|
||||
return {k: (tensors[v.idx] if isinstance(v, _TensorPlaceholder) else v) for k, v in skeleton.items()}
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Job:
|
||||
"""One dispatched job: the controller queues it, the relay (rank 0) broadcasts it to every
|
||||
rank, and all ranks run it. `kwargs` may hold top-level tensors -- broadcast over NCCL rather
|
||||
than pickled (see `_RankLink`). A `None` on the job queue is the shutdown sentinel.
|
||||
"""
|
||||
|
||||
job_id: int # every rank echoes this back; a mismatch means the fleet desynced
|
||||
kwargs: dict[str, Any]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _JobResult:
|
||||
"""worker -> controller: one item from one rank -- a yielded value, a RunnerError, or the
|
||||
terminating StopIteration itself (carrying `.value`, the generator's `return`). A yielded value
|
||||
is forwarded to the caller as soon as it arrives. Terminals are collected per rank: once all
|
||||
`world_size` ranks have ended, the controller classifies them -- all StopIteration -> end
|
||||
iteration with `StopIteration([returns...])`; all/some RunnerError -> Symmetric/Asymmetric.
|
||||
Tensors in `value` ride the queue by shared memory / CUDA IPC, the same for every rank.
|
||||
"""
|
||||
|
||||
job_id: int # the job this is for; the controller rejects a mismatch as a desync
|
||||
rank: int # which rank produced this -- used to collect one terminal per rank (and order returns)
|
||||
value: Any # a yielded value, a RunnerError, or the StopIteration that ended this rank
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Channels:
|
||||
"""The queues bridging the controller and the workers."""
|
||||
|
||||
jobs: Queue # type: ignore[type-arg] # controller -> relay: _Job or None
|
||||
results: Queue # type: ignore[type-arg] # all ranks -> controller: chunk items + each rank's _JobResult
|
||||
ready: Queue # type: ignore[type-arg] # workers -> controller: rank, on setup-complete
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The NCCL core: the per-job input broadcast (relay -> all ranks).
|
||||
# =============================================================================
|
||||
class _RankLink:
|
||||
"""One rank's handle to a job's *input* broadcast (relay -> all ranks, over NCCL). A job's data
|
||||
crosses in two directions and this covers only the inbound one: inputs arrive here, while each
|
||||
rank sends its results back out-of-band on the result queue -- never through a collective -- so
|
||||
the controller assembles them without a second NCCL gather.
|
||||
Between jobs the coordinator parks idle non-relay ranks on a cheap store signal instead of
|
||||
leaving them blocked inside a pending broadcast (which would trip the NCCL watchdog); a rank
|
||||
enters the collective only once the relay signals it has a job to send.
|
||||
"""
|
||||
|
||||
def __init__(self, coordinator: BroadcastCoordinator, device: torch.device) -> None:
|
||||
self._coordinator = coordinator
|
||||
self._device = device # this rank's GPU; where received tensors land
|
||||
|
||||
# ---- relay side
|
||||
def announce_job(self, job: _Job | None) -> _Job | None:
|
||||
"""Broadcast the job to every rank, sending kwargs tensors over NCCL instead of
|
||||
pickling them. Returns the job with its kwargs tensors now on THIS rank's GPU,
|
||||
ready for the relay to run; None for the shutdown sentinel.
|
||||
"""
|
||||
with self._coordinator.broadcast_context():
|
||||
if job is None:
|
||||
dist.broadcast_object_list([None], src=_RELAY_RANK)
|
||||
return None
|
||||
skeleton, tensors = _replace_tensors_by_placeholders(job.kwargs)
|
||||
gpu = [t.to(self._device).contiguous() for t in tensors] # NCCL needs CUDA + contiguous
|
||||
dist.broadcast_object_list([_Job(job.job_id, skeleton)], src=_RELAY_RANK)
|
||||
for t in gpu: # same order every rank, driven by the skeleton broadcast above
|
||||
dist.broadcast(t, src=_RELAY_RANK)
|
||||
job.kwargs = _fill_tensors_into_placeholders(skeleton, gpu)
|
||||
return job
|
||||
|
||||
# ---- non-relay side
|
||||
def await_job(self) -> _Job | None:
|
||||
self._coordinator.wait_for_signal_change()
|
||||
payload: list[Any] = [None]
|
||||
dist.broadcast_object_list(payload, src=_RELAY_RANK)
|
||||
shell: _Job | None = payload[0]
|
||||
if shell is None:
|
||||
return None
|
||||
holes = sorted((v for v in shell.kwargs.values() if isinstance(v, _TensorPlaceholder)), key=lambda h: h.idx)
|
||||
tensors: list[Any] = []
|
||||
for h in holes: # receive in idx order -- matches the relay's send order
|
||||
buf = torch.empty(h.shape, dtype=h.dtype, device=self._device)
|
||||
dist.broadcast(buf, src=_RELAY_RANK)
|
||||
tensors.append(buf)
|
||||
shell.kwargs = _fill_tensors_into_placeholders(shell.kwargs, tensors)
|
||||
return shell
|
||||
|
||||
|
||||
def _run_job(runner: MGPURunner, kwargs: dict[str, Any], job_id: int, rank: int, result: Queue[_JobResult]) -> None:
|
||||
"""Run one job on this rank and put each item on the result queue: every yield is a `_JobResult`
|
||||
(tagged by rank), forwarded straight to the caller; the rank's end -- a clean StopIteration or a
|
||||
raised RunnerError (recoverable) -- is itself a `_JobResult`, which the controller collects and
|
||||
classifies. Runners must be generators; a non-generator return is not iterable, so `next` raises
|
||||
(fatal).
|
||||
"""
|
||||
out = runner(**kwargs)
|
||||
while True:
|
||||
try:
|
||||
value = next(out)
|
||||
except StopIteration as stop: # clean finish: carries the (normally None) return value
|
||||
result.put(_JobResult(job_id, rank, stop))
|
||||
return
|
||||
except RunnerError as err: # runner raised it explicitly -> recoverable; collect and classify
|
||||
result.put(_JobResult(job_id, rank, err))
|
||||
return
|
||||
except ValueError as err: # input validation (symmetric) -> synthesize a recoverable RunnerError
|
||||
result.put(_JobResult(job_id, rank, RunnerError(str(err))))
|
||||
return
|
||||
result.put(_JobResult(job_id, rank, value))
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The worker process: entrypoint + the two loops.
|
||||
# =============================================================================
|
||||
def _relay_loop(runner: MGPURunner, link: _RankLink, channels: _Channels) -> None:
|
||||
"""Rank 0: take a job, broadcast it (tensors over NCCL), run it, stream every item on the queue."""
|
||||
while True:
|
||||
job: _Job | None = channels.jobs.get()
|
||||
job = link.announce_job(job) # returns the job with kwargs tensors on this GPU
|
||||
if job is None: # shutdown sentinel (already broadcast to the others)
|
||||
return
|
||||
_run_job(runner, job.kwargs, job.job_id, _RELAY_RANK, channels.results)
|
||||
|
||||
|
||||
def _worker_loop(runner: MGPURunner, link: _RankLink, channels: _Channels, rank: int) -> None:
|
||||
"""Non-relay ranks: wait for the broadcast, run in SPMD, stream every item on the queue."""
|
||||
while True:
|
||||
job = link.await_job()
|
||||
if job is None: # shutdown sentinel
|
||||
return
|
||||
_run_job(runner, job.kwargs, job.job_id, rank, channels.results)
|
||||
|
||||
|
||||
def _shutdown_distributed() -> None:
|
||||
# Drop torch.compile state before tearing down NCCL: compiled artifacts hold CUDA
|
||||
# pool/stream refs that ncclCommDestroy waits on, and destroy_process_group can
|
||||
# deadlock without this.
|
||||
torch._dynamo.reset()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@record
|
||||
def _worker_entrypoint(
|
||||
runner_cls: type[MGPURunner],
|
||||
setup_kwargs: dict[str, Any],
|
||||
init_timeout: timedelta,
|
||||
channels: _Channels,
|
||||
device_ids: list[int] | None = None,
|
||||
) -> None:
|
||||
"""Per-worker entry. @record turns a crash into a readable failure on the controller."""
|
||||
local_rank = int(os.environ["LOCAL_RANK"])
|
||||
# All GPUs stay visible; we bind this rank to its physical GPU by index. (Setting
|
||||
# CUDA_VISIBLE_DEVICES here is unreliable -- the spawn bootstrap can touch CUDA, freezing
|
||||
# the device list, before a late env override would apply.)
|
||||
device_index = local_rank if device_ids is None else device_ids[local_rank]
|
||||
torch.cuda.set_device(device_index)
|
||||
# Build the runner before NCCL init so its __init__ runs as the per-worker pre-init hook, for setup
|
||||
# that must precede init_process_group (e.g. tests set torch.use_deterministic_algorithms there).
|
||||
runner = runner_cls()
|
||||
if not dist.is_initialized():
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
device_id=torch.device("cuda", device_index),
|
||||
timeout=init_timeout,
|
||||
)
|
||||
|
||||
is_relay = dist.get_rank() == _RELAY_RANK
|
||||
device = torch.device("cuda", device_index)
|
||||
groups = create_local_nccl_groups()
|
||||
store = dist.distributed_c10d.PrefixStore("ltx_pipeline_broadcast/", dist.distributed_c10d._get_default_store())
|
||||
link = _RankLink(BroadcastCoordinator(store=store, is_driver=is_relay), device)
|
||||
|
||||
runner._groups = groups
|
||||
runner.setup(**setup_kwargs)
|
||||
channels.ready.put(dist.get_rank())
|
||||
|
||||
try:
|
||||
if is_relay:
|
||||
_relay_loop(runner, link, channels)
|
||||
else:
|
||||
_worker_loop(runner, link, channels, dist.get_rank())
|
||||
finally:
|
||||
_shutdown_distributed()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# The fleet: spawn, poll for death, terminate. No job knowledge.
|
||||
# =============================================================================
|
||||
class _RunnerShipper:
|
||||
"""Ships a runner CLASS to spawned workers by value (cloudpickle), not by reference.
|
||||
A runner defined in ``__main__`` (a plain script or ``python -m``) or in a test module is
|
||||
not importable under that name in a freshly spawned worker, so the stock by-reference pickle
|
||||
the elastic launcher uses would raise ModuleNotFoundError on every rank. This proxy serializes
|
||||
the class by value; its ``__reduce__`` targets ``cloudpickle.loads`` (importable everywhere),
|
||||
so only the runner crosses by value -- the controller's own channel payloads stay on the stock
|
||||
pickler. The worker unpickles it straight back to the runner class.
|
||||
"""
|
||||
|
||||
def __init__(self, runner_cls: type[MGPURunner]) -> None:
|
||||
module = sys.modules.get(runner_cls.__module__)
|
||||
if module is None:
|
||||
self._payload = cloudpickle.dumps(runner_cls)
|
||||
return
|
||||
# cloudpickle pickles a class by reference when its module looks importable; force
|
||||
# by-value so a __main__/test-module runner survives the spawn.
|
||||
cloudpickle.register_pickle_by_value(module)
|
||||
try:
|
||||
self._payload = cloudpickle.dumps(runner_cls)
|
||||
finally:
|
||||
cloudpickle.unregister_pickle_by_value(module)
|
||||
|
||||
def __reduce__(self) -> tuple[object, tuple[bytes]]:
|
||||
return (cloudpickle.loads, (self._payload,))
|
||||
|
||||
|
||||
def _format_failures(failures: dict[int, ProcessFailure]) -> RuntimeError:
|
||||
lines = [f" rank {r} (pid {f.pid}) exit {f.exitcode}:\n{f.message}" for r, f in failures.items()]
|
||||
return RuntimeError("MGPU worker(s) failed:\n" + "\n".join(lines))
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
class _RunnersFleet:
|
||||
"""The worker processes. All three methods call torch.elastic from the controller
|
||||
thread only -- that is the whole reason there is no lock anywhere.
|
||||
"""
|
||||
|
||||
def __init__(self, pcontext: Any) -> None: # noqa: ANN401
|
||||
self._pcontext = pcontext
|
||||
|
||||
@classmethod
|
||||
def spawn(
|
||||
cls,
|
||||
*,
|
||||
runner_cls: type[MGPURunner],
|
||||
setup_kwargs: dict[str, Any],
|
||||
init_timeout: timedelta,
|
||||
channels: _Channels,
|
||||
num_gpus: int,
|
||||
logs_specs: LogsSpecs,
|
||||
device_ids: list[int] | None = None,
|
||||
) -> _RunnersFleet:
|
||||
port = _find_free_port()
|
||||
envs = {
|
||||
r: {
|
||||
"RANK": str(r),
|
||||
"LOCAL_RANK": str(r),
|
||||
"WORLD_SIZE": str(num_gpus),
|
||||
"MASTER_ADDR": "127.0.0.1",
|
||||
"MASTER_PORT": str(port),
|
||||
}
|
||||
for r in range(num_gpus)
|
||||
}
|
||||
# device_ids maps rank -> physical GPU (None = identity); the worker binds to it. The runner
|
||||
# class ships by value so a __main__/test-module runner survives the spawn (see _RunnerShipper).
|
||||
packed = (_RunnerShipper(runner_cls), setup_kwargs, init_timeout, channels, device_ids)
|
||||
args = dict.fromkeys(range(num_gpus), packed)
|
||||
logger.info("Spawning %d MGPU workers...", num_gpus)
|
||||
return cls(
|
||||
start_processes(
|
||||
name="ltx_mgpu_worker",
|
||||
entrypoint=_worker_entrypoint,
|
||||
args=args,
|
||||
envs=envs,
|
||||
logs_specs=logs_specs,
|
||||
start_method="spawn",
|
||||
)
|
||||
)
|
||||
|
||||
def poll(self) -> RuntimeError | None:
|
||||
"""None while all workers are alive. Once any has exited, an error describing it.
|
||||
Used only mid-job, where ANY exit is unexpected (workers only exit on the
|
||||
shutdown sentinel), so a clean exit is reported as an error too.
|
||||
"""
|
||||
result = self._pcontext.wait(timeout=0)
|
||||
if result is None:
|
||||
return None
|
||||
if result.failures:
|
||||
return _format_failures(result.failures)
|
||||
return RuntimeError("MGPU workers exited unexpectedly")
|
||||
|
||||
def drain(self, timeout: float) -> bool:
|
||||
"""Wait up to `timeout` for all workers to exit on their own. True if they did."""
|
||||
end = time.monotonic() + timeout
|
||||
while time.monotonic() < end:
|
||||
if self._pcontext.wait(timeout=0) is not None:
|
||||
return True
|
||||
time.sleep(_POLL_S)
|
||||
return False
|
||||
|
||||
def terminate(self) -> None:
|
||||
"""SIGTERM -> SIGKILL. Never raises; safe to call more than once."""
|
||||
with contextlib.suppress(Exception):
|
||||
self._pcontext.close()
|
||||
@@ -0,0 +1,150 @@
|
||||
"""Multi-GPU Gemma text encoder builder.
|
||||
Replaces the text encoder builder on the ``PromptEncoder`` block with an
|
||||
:class:`AccelerateGemmaBuilder` that uses ``device_map="auto"`` on the
|
||||
source rank and a broadcast stub elsewhere.
|
||||
On the source rank the first ``build()`` loads via HuggingFace
|
||||
``from_pretrained`` and caches the full state dict (including non-persistent
|
||||
buffers) in the provided :class:`Registry`. Subsequent calls recreate the
|
||||
model from cache and reinstall accelerate dispatch hooks — no disk I/O.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from accelerate import dispatch_model
|
||||
from transformers import Gemma3ForConditionalGeneration
|
||||
|
||||
from ltx_core.loader.primitives import BuilderProtocol, StateDict
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.multigpu.gemma.accelerate_wrapper import AccelerateGemmaWrapper
|
||||
from ltx_core.multigpu.gemma.loader import load_gemma_with_device_map
|
||||
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AccelerateGemmaBuilder(BuilderProtocol):
|
||||
"""Builder that loads Gemma with ``device_map="auto"`` on the source rank.
|
||||
Conforms to the builder interface expected by ``PromptEncoder``:
|
||||
``build(device, dtype) -> model``. Non-source ranks get a lightweight
|
||||
:class:`AccelerateGemmaWrapper` that receives embeddings via broadcast.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gemma_root_path: str,
|
||||
gemma_group: dist.ProcessGroup | None,
|
||||
broadcast_group: dist.ProcessGroup | None,
|
||||
registry: Registry,
|
||||
*,
|
||||
src_rank: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
) -> None:
|
||||
self._gemma_root_path = gemma_root_path
|
||||
self._gemma_group = gemma_group
|
||||
self._broadcast_group = broadcast_group
|
||||
self._registry = registry
|
||||
self._src_rank = src_rank
|
||||
self._is_src = dist.get_rank() == src_rank
|
||||
self._dtype = dtype
|
||||
# Cached on the src rank after first build (non-tensor objects).
|
||||
self._config: object | None = None
|
||||
self._hf_device_map: dict[str, int | str] | None = None
|
||||
self._tokenizer: object | None = None
|
||||
self._processor: object | None = None
|
||||
|
||||
@property
|
||||
def registry(self) -> Registry:
|
||||
return self._registry
|
||||
|
||||
def with_registry(self, registry: Registry) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._registry = registry
|
||||
return clone
|
||||
|
||||
def model_config(self) -> dict:
|
||||
return {}
|
||||
|
||||
def build(
|
||||
self,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
**_kwargs: Any, # noqa: ANN401
|
||||
) -> AccelerateGemmaWrapper:
|
||||
dtype = dtype or self._dtype
|
||||
|
||||
encoder = self._build_encoder(dtype) if self._is_src else None
|
||||
|
||||
return AccelerateGemmaWrapper(
|
||||
encoder=encoder,
|
||||
broadcast_group=self._broadcast_group,
|
||||
src_rank=dist.get_group_rank(self._broadcast_group, self._src_rank),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# -- src-rank helpers ---------------------------------------------------
|
||||
|
||||
def _build_encoder(self, dtype: torch.dtype) -> GemmaTextEncoder:
|
||||
cached = self._registry.get([self._gemma_root_path], None)
|
||||
if cached is not None:
|
||||
logger.info("Rebuilding Gemma from cached state dict (no disk I/O).")
|
||||
return self._rebuild_from_cache(cached, dtype)
|
||||
|
||||
encoder = load_gemma_with_device_map(self._gemma_root_path, dtype)
|
||||
|
||||
# Cache non-tensor objects on the builder instance.
|
||||
self._config = encoder.model.config
|
||||
self._hf_device_map = encoder.model.hf_device_map
|
||||
self._tokenizer = encoder.tokenizer
|
||||
self._processor = encoder.processor
|
||||
|
||||
# Cache full state dict including non-persistent buffers.
|
||||
sd = encoder.model.state_dict()
|
||||
for name, buf in encoder.model.named_buffers():
|
||||
if name not in sd:
|
||||
sd[name] = buf
|
||||
total_size = sum(t.nelement() * t.element_size() for t in sd.values())
|
||||
dtypes = {t.dtype for t in sd.values()}
|
||||
self._registry.add(
|
||||
[self._gemma_root_path],
|
||||
None,
|
||||
StateDict(sd=sd, device=torch.device("meta"), size=total_size, dtype=dtypes),
|
||||
)
|
||||
logger.info("Cached Gemma state dict in registry (%d entries).", len(sd))
|
||||
|
||||
return encoder
|
||||
|
||||
def _rebuild_from_cache(self, cached: StateDict, dtype: torch.dtype) -> GemmaTextEncoder:
|
||||
with torch.device("meta"):
|
||||
model = Gemma3ForConditionalGeneration(self._config)
|
||||
|
||||
# Split into persistent (load_state_dict) and non-persistent (manual assign).
|
||||
expected_keys = set(model.state_dict().keys())
|
||||
persistent_sd = {k: v for k, v in cached.sd.items() if k in expected_keys}
|
||||
non_persistent_sd = {k: v for k, v in cached.sd.items() if k not in expected_keys}
|
||||
|
||||
model.load_state_dict(persistent_sd, strict=True, assign=True)
|
||||
for name, tensor in non_persistent_sd.items():
|
||||
parent_path, attr = name.rsplit(".", 1)
|
||||
module = model
|
||||
for part in parent_path.split("."):
|
||||
module = getattr(module, part)
|
||||
setattr(module, attr, tensor)
|
||||
|
||||
dispatch_model(model, self._hf_device_map)
|
||||
|
||||
return GemmaTextEncoder(
|
||||
model=model,
|
||||
tokenizer=self._tokenizer,
|
||||
processor=self._processor,
|
||||
dtype=dtype,
|
||||
)
|
||||
@@ -0,0 +1,31 @@
|
||||
"""NCCL process group management."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
@dataclass
|
||||
class NCCLGroups:
|
||||
"""Container for the NCCL process groups used by each pipeline component."""
|
||||
|
||||
gemma_group: dist.ProcessGroup
|
||||
transformer_group: dist.ProcessGroup
|
||||
vae_group: dist.ProcessGroup
|
||||
|
||||
|
||||
def create_local_nccl_groups() -> NCCLGroups:
|
||||
"""Create NCCL process groups for each pipeline component.
|
||||
All ranks must call this collectively because ``dist.new_group`` is a
|
||||
collective operation. All ranks participate in every group.
|
||||
Returns:
|
||||
NCCLGroups with one process group per component.
|
||||
"""
|
||||
all_ranks = list(range(dist.get_world_size()))
|
||||
return NCCLGroups(
|
||||
gemma_group=dist.new_group(ranks=all_ranks),
|
||||
transformer_group=dist.new_group(ranks=all_ranks),
|
||||
vae_group=dist.new_group(ranks=all_ranks),
|
||||
)
|
||||
@@ -0,0 +1,62 @@
|
||||
"""The runner contract: the base class a pipeline subclasses and the error it raises.
|
||||
``MGPURunner`` is what a pipeline subclasses (``setup`` + a generator ``__call__``); ``RunnerError``
|
||||
is the recoverable, symmetric failure a runner raises. The fleet (``fleet.py``) runs runners and
|
||||
ships them to the spawned workers; the controller (``controller.py``) classifies their results.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from ltx_pipelines.multigpu.nccl_groups import NCCLGroups
|
||||
|
||||
|
||||
class RunnerError(Exception):
|
||||
"""Raised by a runner (or synthesized by the controller from a ValueError) for a recoverable,
|
||||
SYMMETRIC failure. The worker loop catches it and puts it on the result queue as that rank's
|
||||
end; the controller collects it with the other terminals and classifies them -- every rank ->
|
||||
SymmetricRunnerError, a mix with clean finishes -> AsymmetricRunnerError. Raise it IDENTICALLY on
|
||||
every rank, outside any collective (e.g. validating the broadcast kwargs before the first one);
|
||||
a RunnerError on only some ranks is the contract violation AsymmetricRunnerError flags.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
self.message = message
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class MGPURunner(ABC):
|
||||
"""Subclass this. The controller builds one per worker, injects groups,
|
||||
calls setup() once, then calls the instance per job. Define it anywhere -- the controller
|
||||
ships the class to workers by value, so a runner in ``__main__`` or a test module works.
|
||||
setup() and __call__() run on EVERY rank. __call__ MUST be a generator (use `yield`, even
|
||||
once): each yield is forwarded to the Stream on its own, as it arrives -- yields are NOT gathered
|
||||
across ranks. Results are *yielded*; a `return` value, if any, rides the terminal
|
||||
StopIteration.value (the rare path), collected per rank once every rank has ended. The framework
|
||||
does not wrap them in inference mode -- annotate your setup()/__call__ with @torch.inference_mode()
|
||||
if you want it.
|
||||
Tensors are transparent: pass them as kwargs (e.g. `stream(latent=t)`) and the relay
|
||||
broadcasts them to every rank over NCCL, so `__call__` receives them already on the local GPU.
|
||||
Yield tensors back the same way -- as values in a yielded dict -- and they come back to the
|
||||
controller without being pickled. See the module docstring.
|
||||
Raising an unexpected exception from __call__ is FATAL: it kills the worker, poisons the
|
||||
controller, and needs a new one. For a RECOVERABLE failure raise a `RunnerError` (or a `ValueError`,
|
||||
which the controller converts) -- identically on every rank, outside any collective. The worker
|
||||
loop catches it and the fleet stays alive; iterating the Stream re-raises it as
|
||||
SymmetricRunnerError.
|
||||
"""
|
||||
|
||||
_groups: NCCLGroups
|
||||
|
||||
@property
|
||||
def groups(self) -> NCCLGroups:
|
||||
return self._groups
|
||||
|
||||
@abstractmethod
|
||||
def setup(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, *args: Any, **kwargs: Any) -> Any: # noqa: ANN401
|
||||
...
|
||||
@@ -0,0 +1,57 @@
|
||||
"""Sequence parallel transformer builder.
|
||||
Wrapping builder that produces a transformer model with sequence parallelism applied.
|
||||
Requires ``ltx-kernels`` to be installed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Generic
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.loader.primitives import ModelBuilderProtocol
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
|
||||
from ltx_core.model.model_protocol import LTXModelProtocol
|
||||
from ltx_core.multigpu.transformer.attention import AttentionManager
|
||||
from ltx_core.multigpu.transformer.sequence_parallel import (
|
||||
SequenceParallelModelWrapper,
|
||||
create_video_self_attention_module_ops,
|
||||
)
|
||||
from ltx_pipelines.multigpu.delegating_builder import DelegatingBuilder, InnerModelT
|
||||
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
|
||||
|
||||
|
||||
class SequenceParallelBuilder(DelegatingBuilder[InnerModelT], Generic[InnerModelT]):
|
||||
"""Builder that injects SP module ops and wraps with :class:`SequenceParallelModelWrapper`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: ModelBuilderProtocol[LTXModelProtocol],
|
||||
attn_mgr: AttentionManager,
|
||||
registry: Registry,
|
||||
tracker: TransformerWeightTracker,
|
||||
) -> None:
|
||||
if not isinstance(inner, Builder):
|
||||
raise TypeError(f"SequenceParallelBuilder wraps a SingleGPUModelBuilder, got {type(inner).__name__}")
|
||||
cuda_device = torch.device(f"cuda:{torch.cuda.current_device()}")
|
||||
inner = inner.with_registry(registry).with_lora_load_device(cuda_device)
|
||||
sp_ops = create_video_self_attention_module_ops(attn_mgr)
|
||||
self._inner = inner.with_module_ops((*inner.module_ops, sp_ops))
|
||||
self._tracker = tracker
|
||||
self._attn_mgr = attn_mgr
|
||||
|
||||
@property
|
||||
def all2all_timeout_seconds(self) -> float:
|
||||
"""The SP all2all barrier timeout (seconds); forwards to the AttentionManager that owns the refs."""
|
||||
return self._attn_mgr.all2all_timeout_seconds
|
||||
|
||||
@all2all_timeout_seconds.setter
|
||||
def all2all_timeout_seconds(self, seconds: float) -> None:
|
||||
self._attn_mgr.all2all_timeout_seconds = seconds
|
||||
|
||||
def build(
|
||||
self, device: torch.device | None = None, dtype: torch.dtype | None = None, **kwargs: object
|
||||
) -> SequenceParallelModelWrapper:
|
||||
model = self._tracker.build(self._inner, device=device, dtype=dtype, **kwargs)
|
||||
return SequenceParallelModelWrapper(model, self._attn_mgr)
|
||||
@@ -0,0 +1,70 @@
|
||||
"""Tiled data parallel transformer builder.
|
||||
Wrapping builder that produces a transformer model with tiled data parallelism applied.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Generic
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.loader.primitives import ModelBuilderProtocol
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
|
||||
from ltx_core.model.model_protocol import LTXModelProtocol
|
||||
from ltx_core.multigpu.transformer.tiled_data_parallel import (
|
||||
TiledDataParallelModelWrapper,
|
||||
)
|
||||
from ltx_core.tiling import TileCountConfig
|
||||
from ltx_core.tools import VideoLatentTools
|
||||
from ltx_pipelines.multigpu.delegating_builder import DelegatingBuilder, InnerModelT
|
||||
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
|
||||
|
||||
|
||||
class TiledDataParallelBuilder(DelegatingBuilder[InnerModelT], Generic[InnerModelT]):
|
||||
"""Builder conforming to :class:`ModelBuilderProtocol` that wraps with
|
||||
:class:`TiledDataParallelModelWrapper`.
|
||||
Requires ``video_tools`` as a keyword argument to :meth:`build` so the
|
||||
wrapper can compute the tile for this rank.
|
||||
The underlying model must accept ``(video, audio, perturbations)`` and return
|
||||
``(denoised_video, denoised_audio)`` — i.e. conform to the ``X0Model`` forward
|
||||
signature used by the LTX transformer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: ModelBuilderProtocol[LTXModelProtocol],
|
||||
group: dist.ProcessGroup,
|
||||
tiling: TileCountConfig,
|
||||
registry: Registry,
|
||||
tracker: TransformerWeightTracker,
|
||||
normalize_positions: bool = True,
|
||||
) -> None:
|
||||
if not isinstance(inner, Builder):
|
||||
raise TypeError(f"TiledDataParallelBuilder wraps a SingleGPUModelBuilder, got {type(inner).__name__}")
|
||||
cuda_device = torch.device(f"cuda:{torch.cuda.current_device()}")
|
||||
self._inner = inner.with_registry(registry).with_lora_load_device(cuda_device)
|
||||
self._tracker = tracker
|
||||
self._group = group
|
||||
self._tiling = tiling
|
||||
self._normalize_positions = normalize_positions
|
||||
|
||||
def build(
|
||||
self,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
*,
|
||||
video_tools: VideoLatentTools | None = None,
|
||||
**_kwargs: object,
|
||||
) -> TiledDataParallelModelWrapper:
|
||||
if video_tools is None:
|
||||
raise ValueError("TiledDataParallelBuilder.build() requires video_tools")
|
||||
model = self._tracker.build(self._inner, device=device, dtype=dtype, **_kwargs)
|
||||
return TiledDataParallelModelWrapper(
|
||||
model,
|
||||
video_tools=video_tools,
|
||||
tiling=self._tiling,
|
||||
group=self._group,
|
||||
normalize_positions=self._normalize_positions,
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Multi-GPU VAE decoder builder.
|
||||
Wrapping builder that produces a :class:`DistributedVideoDecoder`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.multiprocessing import Queue
|
||||
|
||||
from ltx_core.loader.primitives import BuilderProtocol
|
||||
from ltx_core.loader.registry import Registry
|
||||
from ltx_core.multigpu.vae.distributed_decoder import DistributedVideoDecoder
|
||||
from ltx_core.tiling import TileCountConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
class DistributedDecoderBuilder(BuilderProtocol):
|
||||
"""Builder that wraps a base decoder builder with distributed logic."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner: BuilderProtocol,
|
||||
queue: Queue, # type: ignore[type-arg]
|
||||
vae_group: dist.ProcessGroup,
|
||||
vae_tiling: TileCountConfig,
|
||||
driver_rank: int,
|
||||
registry: Registry,
|
||||
) -> None:
|
||||
self._inner = inner.with_registry(registry)
|
||||
self._queue = queue
|
||||
self._vae_group = vae_group
|
||||
self._vae_tiling = vae_tiling
|
||||
self._driver_rank = driver_rank
|
||||
|
||||
@property
|
||||
def registry(self) -> Registry:
|
||||
return self._inner.registry
|
||||
|
||||
def with_registry(self, registry: Registry) -> Self:
|
||||
clone = copy.copy(self)
|
||||
clone._inner = self._inner.with_registry(registry)
|
||||
return clone
|
||||
|
||||
def build(
|
||||
self,
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
**kwargs: Any, # noqa: ANN401
|
||||
) -> DistributedVideoDecoder:
|
||||
base_decoder = self._inner.build(device=device, dtype=dtype, **kwargs)
|
||||
return DistributedVideoDecoder(
|
||||
base_decoder,
|
||||
queue=self._queue,
|
||||
vae_group=self._vae_group,
|
||||
vae_tiling=self._vae_tiling,
|
||||
driver_rank=dist.get_group_rank(self._vae_group, self._driver_rank),
|
||||
)
|
||||
@@ -0,0 +1,184 @@
|
||||
"""Distributed transformer weight tracker with LoRA hot-swap.
|
||||
Shared infrastructure used by both TDP and SP builders.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.loader.fuse_loras import fuse_lora_weights
|
||||
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps, LoraStateDictWithStrength, StateDict
|
||||
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
|
||||
from ltx_core.model.model_protocol import ModelType
|
||||
from ltx_core.multigpu.sharded_sd import ShardedSD
|
||||
|
||||
|
||||
def _apply_loras_inplace(
|
||||
source: dict[str, torch.Tensor],
|
||||
target: dict[str, torch.Tensor],
|
||||
loras: tuple[LoraPathStrengthAndSDOps, ...],
|
||||
builder: Builder, # type: ignore[type-arg]
|
||||
lora_keys: frozenset[str],
|
||||
) -> None:
|
||||
"""Reset *target* to clean weights from *source*, then fuse all LoRAs in one pass."""
|
||||
for key, clean_weight in source.items():
|
||||
target[key].copy_(clean_weight)
|
||||
|
||||
lora_sds = [
|
||||
LoraStateDictWithStrength(
|
||||
builder.load_sd(
|
||||
[lora.path],
|
||||
sd_ops=lora.sd_ops.with_additional_allowed_keys(lora_keys),
|
||||
registry=builder.registry,
|
||||
device=builder.lora_load_device,
|
||||
),
|
||||
lora.strength,
|
||||
)
|
||||
for lora in loras
|
||||
if lora.strength != 0
|
||||
]
|
||||
target_sd = StateDict(
|
||||
sd=target, device=next(iter(target.values())).device, size=0, dtype={next(iter(target.values())).dtype}
|
||||
)
|
||||
for key, fused in fuse_lora_weights(target_sd, lora_sds, fuse_rule=builder.fuse_rule):
|
||||
target[key].copy_(fused)
|
||||
|
||||
|
||||
class TransformerWeightTracker:
|
||||
"""Tracks cached transformer weights with distributed LoRA hot-swap.
|
||||
Shared across stage builders that operate on the same checkpoint.
|
||||
Does **not** own the model weights — it references tensors stored in a
|
||||
:class:`Registry` and receives a builder at :meth:`build` time.
|
||||
Uses two :class:`ShardedSD` instances (created on first :meth:`build` call):
|
||||
- ``stored_sd`` — cloned backup of the original (pre-LoRA) weights.
|
||||
Used to restore registry tensors before applying a different LoRA set.
|
||||
- ``broadcast_sd`` — zero-copy view into the registry tensors.
|
||||
After in-place LoRA fusion on the owning rank, this broadcasts the
|
||||
fused results to all other ranks so every rank sees the same weights.
|
||||
Both are created together and are always either both ``None`` or both set.
|
||||
With ``no_lora_swap``, the LoRA set is assumed fixed (none, or one set):
|
||||
the backup clone is skipped and any swap or reset raises.
|
||||
"""
|
||||
|
||||
def __init__(self, group: dist.ProcessGroup, bucket_mb: int = 256, no_lora_swap: bool = False) -> None:
|
||||
if bucket_mb <= 0:
|
||||
raise ValueError("bucket_mb must be > 0")
|
||||
self._group = group
|
||||
self._bucket_mb = bucket_mb
|
||||
self._no_lora_swap = no_lora_swap
|
||||
self._staging: torch.Tensor | None = None
|
||||
self.stored_sd: ShardedSD | None = None
|
||||
self.broadcast_sd: ShardedSD | None = None
|
||||
self.loras: list[tuple[str, float, str]] = []
|
||||
|
||||
@property
|
||||
def staging(self) -> torch.Tensor:
|
||||
"""The single broadcast scratch buffer, shared by both SDs, allocated on first use."""
|
||||
if self._staging is None:
|
||||
device = torch.device(f"cuda:{torch.cuda.current_device()}")
|
||||
self._staging = torch.empty(self._bucket_mb * 1024 * 1024, dtype=torch.uint8, device=device)
|
||||
return self._staging
|
||||
|
||||
def loras_match(self, lora_list: list[tuple[str, float, str]]) -> bool:
|
||||
if len(lora_list) != len(self.loras):
|
||||
return False
|
||||
return sorted(lora_list) == sorted(self.loras)
|
||||
|
||||
def reset_loras(self, target_sd: dict[str, torch.Tensor]) -> None:
|
||||
"""Restore *target_sd* to original (pre-LoRA) weights.
|
||||
No-op if no LoRAs are currently applied. This is a cooperative
|
||||
operation — all ranks must call it simultaneously.
|
||||
"""
|
||||
if not self.loras:
|
||||
return
|
||||
if self._no_lora_swap:
|
||||
raise RuntimeError("no_lora_swap tracker has no backup to reset from")
|
||||
if self.stored_sd is None:
|
||||
raise RuntimeError("stored_sd must be initialised before reset_loras")
|
||||
self.loras = []
|
||||
self.stored_sd.broadcast_shards_into(target_sd, self.staging)
|
||||
|
||||
def _local_lora_keys(self) -> frozenset[str]:
|
||||
"""Derive LoRA key names from the locally owned model keys."""
|
||||
if self.stored_sd is None:
|
||||
return frozenset()
|
||||
|
||||
keys: set[str] = set()
|
||||
for k in self.stored_sd.local_shard:
|
||||
if k.endswith(".weight"):
|
||||
prefix = k[: -len(".weight")]
|
||||
keys.add(f"{prefix}.lora_A.weight")
|
||||
keys.add(f"{prefix}.lora_B.weight")
|
||||
return frozenset(keys)
|
||||
|
||||
def apply_loras_(
|
||||
self,
|
||||
target_sd: dict[str, torch.Tensor],
|
||||
loras: tuple[LoraPathStrengthAndSDOps, ...],
|
||||
builder: Builder, # type: ignore[type-arg]
|
||||
) -> None:
|
||||
"""Fuse *loras* into *target_sd* in-place (trailing ``_`` denotes in-place).
|
||||
Skips work when the requested LoRAs already match. Restores stored
|
||||
weights before applying new LoRAs. This is a cooperative operation —
|
||||
all ranks must call it simultaneously.
|
||||
"""
|
||||
new_loras = [(lora.path, lora.strength, lora.sd_ops.name) for lora in loras]
|
||||
|
||||
if self.loras_match(new_loras):
|
||||
return
|
||||
|
||||
if self._no_lora_swap and self.loras:
|
||||
raise RuntimeError(f"no_lora_swap tracker cannot change LoRAs: have {self.loras}, requested {new_loras}")
|
||||
|
||||
if all(lora.strength == 0 for lora in loras):
|
||||
self.reset_loras(target_sd)
|
||||
return
|
||||
|
||||
if self.stored_sd is None or self.broadcast_sd is None:
|
||||
raise RuntimeError("ShardedSDs must be initialised before apply_loras_ (call build first)")
|
||||
|
||||
source = self.stored_sd.local_shard
|
||||
target = {k: v for k, v in target_sd.items() if k in source}
|
||||
lora_keys = self._local_lora_keys()
|
||||
_apply_loras_inplace(source, target, loras, builder, lora_keys)
|
||||
|
||||
self.broadcast_sd.broadcast_shards_into(target_sd, self.staging)
|
||||
self.loras = new_loras
|
||||
|
||||
def build(
|
||||
self,
|
||||
builder: Builder[ModelType],
|
||||
device: torch.device | None = None,
|
||||
dtype: torch.dtype | None = None,
|
||||
**kwargs: object,
|
||||
) -> ModelType:
|
||||
"""Build the transformer model with distributed LoRA hot-swap.
|
||||
Populates the registry with clean weights on first call, then applies
|
||||
LoRAs in-place and broadcasts to all ranks. Assumes the builder carries
|
||||
a non-dummy :class:`Registry` so that weights can be cached and reused
|
||||
across calls.
|
||||
"""
|
||||
loras = builder.loras
|
||||
clean_builder = builder.with_loras(())
|
||||
|
||||
model_paths = list(builder.model_path) if isinstance(builder.model_path, tuple) else [builder.model_path]
|
||||
|
||||
# First call: populate the registry with clean weights.
|
||||
if clean_builder.registry.get(model_paths, clean_builder.model_sd_ops) is None:
|
||||
clean_builder.build(device=device, dtype=dtype, **kwargs)
|
||||
|
||||
cached_sd = clean_builder.registry.get(model_paths, clean_builder.model_sd_ops)
|
||||
if cached_sd is None:
|
||||
raise RuntimeError("Expected model state dict in registry but found None")
|
||||
|
||||
if self.stored_sd is None:
|
||||
self.stored_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=not self._no_lora_swap)
|
||||
self.broadcast_sd = ShardedSD.from_state_dict(cached_sd.sd, self._group, clone=False)
|
||||
|
||||
if loras:
|
||||
self.apply_loras_(cached_sd.sd, loras, builder)
|
||||
else:
|
||||
self.reset_loras(cached_sd.sd)
|
||||
|
||||
return clean_builder.build(device=device, dtype=dtype, **kwargs)
|
||||
@@ -17,6 +17,7 @@ from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import (
|
||||
SpatioTemporalScaleFactors,
|
||||
)
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import video_editing_arg_parser
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioConditioner,
|
||||
@@ -76,6 +77,7 @@ class RetakePipeline:
|
||||
distilled: bool = True,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
@@ -89,20 +91,23 @@ class RetakePipeline:
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_conditioner = AudioConditioner(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage = DiffusionStage(
|
||||
self.stage = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
@@ -111,18 +116,21 @@ class RetakePipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------- #
|
||||
|
||||
@@ -16,16 +16,16 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio
|
||||
from ltx_pipelines.utils import get_device
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
default_1_stage_t2a_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
DiffusionStage,
|
||||
PromptEncoder,
|
||||
)
|
||||
from ltx_pipelines.utils.constants import detect_params
|
||||
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
|
||||
from ltx_pipelines.utils.media_io import encode_audio
|
||||
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
|
||||
@@ -56,6 +56,7 @@ class T2AOneStagePipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.dtype = torch.bfloat16
|
||||
self.device = device or get_device()
|
||||
@@ -67,12 +68,13 @@ class T2AOneStagePipeline:
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
# Audio-only: build an audio-only transformer (model_configurator) so the video
|
||||
# weights are never instantiated, plus a use-case-specific SDOps that restricts
|
||||
# checkpoint reads to the audio model's keys, so the video weights are never even
|
||||
# read from disk (the loader skips any key the SDOps maps to None).
|
||||
self.stage = DiffusionStage(
|
||||
self.stage = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
@@ -83,12 +85,14 @@ class T2AOneStagePipeline:
|
||||
offload_mode=offload_mode,
|
||||
model_configurator=LTXAudioOnlyModelConfigurator,
|
||||
model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
def __call__(
|
||||
@@ -153,8 +157,7 @@ class T2AOneStagePipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path()
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params()
|
||||
parser = default_1_stage_t2a_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
pipeline = T2AOneStagePipeline(
|
||||
|
||||
@@ -21,10 +21,11 @@ from ltx_pipelines.utils import (
|
||||
combined_image_conditionings,
|
||||
get_device,
|
||||
)
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
default_1_stage_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -33,7 +34,6 @@ from ltx_pipelines.utils.blocks import (
|
||||
PromptEncoder,
|
||||
VideoDecoder,
|
||||
)
|
||||
from ltx_pipelines.utils.constants import detect_params
|
||||
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
|
||||
from ltx_pipelines.utils.media_io import encode_video
|
||||
from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
|
||||
@@ -58,6 +58,7 @@ class TI2VidOneStagePipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.dtype = torch.bfloat16
|
||||
self.device = device or get_device()
|
||||
@@ -69,14 +70,16 @@ class TI2VidOneStagePipeline:
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage = DiffusionStage(
|
||||
self.stage = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
@@ -85,18 +88,21 @@ class TI2VidOneStagePipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
def __call__( # noqa: PLR0913
|
||||
@@ -187,8 +193,7 @@ class TI2VidOneStagePipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path()
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params()
|
||||
parser = default_1_stage_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
pipeline = TI2VidOneStagePipeline(
|
||||
|
||||
@@ -16,10 +16,11 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio, VideoPixelShape
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import (
|
||||
ImageConditioningInput,
|
||||
default_2_stage_arg_parser,
|
||||
detect_checkpoint_path,
|
||||
resolve_cli_params,
|
||||
)
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -31,7 +32,6 @@ from ltx_pipelines.utils.blocks import (
|
||||
)
|
||||
from ltx_pipelines.utils.constants import (
|
||||
STAGE_2_DISTILLED_SIGMAS,
|
||||
detect_params,
|
||||
)
|
||||
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
|
||||
from ltx_pipelines.utils.helpers import (
|
||||
@@ -52,7 +52,7 @@ class TI2VidTwoStagesPipeline:
|
||||
images parameter.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
def __init__( # noqa: PLR0913
|
||||
self,
|
||||
checkpoint_path: str,
|
||||
distilled_lora: list[LoraPathStrengthAndSDOps],
|
||||
@@ -64,22 +64,40 @@ class TI2VidTwoStagesPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
self._scheduler = LTX2Scheduler()
|
||||
|
||||
self.prompt_encoder = PromptEncoder(
|
||||
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
|
||||
checkpoint_path,
|
||||
gemma_root,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.upsampler = VideoUpsampler(
|
||||
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -88,8 +106,9 @@ class TI2VidTwoStagesPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -98,6 +117,7 @@ class TI2VidTwoStagesPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
def __call__( # noqa: PLR0913
|
||||
@@ -196,7 +216,9 @@ class TI2VidTwoStagesPipeline:
|
||||
)
|
||||
)
|
||||
|
||||
video_state, audio_state = self.stage_2(
|
||||
# Stage 2 refines video only; discard its audio. On the multi-GPU path stage-2 audio
|
||||
# runs under partial tiled/TDP video context, so the full-context stage-1 audio is kept.
|
||||
video_state, _ = self.stage_2(
|
||||
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
|
||||
sigmas=stage_2_sigmas,
|
||||
noiser=noiser,
|
||||
@@ -225,8 +247,7 @@ class TI2VidTwoStagesPipeline:
|
||||
@torch.inference_mode()
|
||||
def main() -> None:
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
checkpoint_path = detect_checkpoint_path()
|
||||
params = detect_params(checkpoint_path)
|
||||
params = resolve_cli_params()
|
||||
parser = default_2_stage_arg_parser(params=params)
|
||||
args = parser.parse_args()
|
||||
pipeline = TI2VidTwoStagesPipeline(
|
||||
|
||||
@@ -13,6 +13,7 @@ from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
|
||||
from ltx_pipelines.utils.blocks import (
|
||||
AudioDecoder,
|
||||
@@ -63,6 +64,7 @@ class TI2VidTwoStagesHQPipeline:
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
):
|
||||
self.device = device or get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
@@ -80,16 +82,33 @@ class TI2VidTwoStagesHQPipeline:
|
||||
)
|
||||
|
||||
self.prompt_encoder = PromptEncoder(
|
||||
checkpoint_path, gemma_root, self.dtype, self.device, registry=registry, offload_mode=offload_mode
|
||||
checkpoint_path,
|
||||
gemma_root,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.upsampler = VideoUpsampler(
|
||||
checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
|
||||
checkpoint_path,
|
||||
spatial_upsampler_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
registry=registry,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.video_decoder = VideoDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.audio_decoder = AudioDecoder(
|
||||
checkpoint_path, self.dtype, self.device, registry=registry, alloc_trim_strategy=alloc_trim_strategy
|
||||
)
|
||||
self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
|
||||
|
||||
self.stage_1 = DiffusionStage(
|
||||
self.stage_1 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -98,8 +117,9 @@ class TI2VidTwoStagesHQPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
self.stage_2 = DiffusionStage(
|
||||
self.stage_2 = DiffusionStage.from_checkpoint(
|
||||
checkpoint_path,
|
||||
self.dtype,
|
||||
self.device,
|
||||
@@ -108,6 +128,7 @@ class TI2VidTwoStagesHQPipeline:
|
||||
registry=registry,
|
||||
compilation_config=compilation_config,
|
||||
offload_mode=offload_mode,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
@torch.inference_mode()
|
||||
@@ -213,7 +234,9 @@ class TI2VidTwoStagesHQPipeline:
|
||||
)
|
||||
)
|
||||
|
||||
video_state, audio_state = self.stage_2(
|
||||
# Stage 2 refines video only; discard its audio. On the multi-GPU path stage-2 audio
|
||||
# runs under partial tiled/TDP video context, so the full-context stage-1 audio is kept.
|
||||
video_state, _ = self.stage_2(
|
||||
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
|
||||
sigmas=stage_2_sigmas,
|
||||
noiser=noiser,
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
"""Multi-GPU two-stage HQ text/image-to-video runner.
|
||||
Runs :class:`TI2VidTwoStagesHQPipeline` across multiple GPUs with:
|
||||
- **Stage 1** -- sequence parallelism (SP) at half resolution
|
||||
- **Stage 2** -- tiled data parallelism (TDP) on height + width with overlap,
|
||||
at full resolution
|
||||
- **Gemma** -- Accelerate-based parallelization
|
||||
- **VAE** -- distributed decoding
|
||||
The HQ pipeline applies the distilled LoRA in both stages with separate
|
||||
strengths and uses the res_2s second-order sampler.
|
||||
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Iterator
|
||||
from multiprocessing import SimpleQueue
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.components.guiders import MultiModalGuiderParams
|
||||
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
|
||||
from ltx_core.loader.registry import StateDictRegistry
|
||||
from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import get_video_chunks_number
|
||||
from ltx_core.model.video_vae.tiling import TilingConfig
|
||||
from ltx_core.multigpu.transformer.attention import AttentionManager
|
||||
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
|
||||
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
|
||||
from ltx_pipelines.multigpu.controller import MGPUController
|
||||
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner
|
||||
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
|
||||
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
|
||||
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
|
||||
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
|
||||
from ltx_pipelines.ti2vid_two_stages_hq import TI2VidTwoStagesHQPipeline
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.constants import TDP_DISTILLED_SIGMAS
|
||||
from ltx_pipelines.utils.media_io import encode_video
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Stage 1 at half-res, 121 frames = 6144 video tokens + audio tokens.
|
||||
_DEFAULT_SP_MAX_TOKENS = 32768
|
||||
# Rank that collects distributed-VAE tiles and encodes the assembled video.
|
||||
_DRIVER_RANK = 0
|
||||
|
||||
|
||||
class TI2VidTwoStagesHQRunner(MGPURunner):
|
||||
"""Distributed HQ pipeline: SP stage 1 (half-res) + TDP stage 2 (full-res) + Gemma + distributed VAE."""
|
||||
|
||||
@torch.inference_mode()
|
||||
def setup(
|
||||
self,
|
||||
*,
|
||||
checkpoint_path: str,
|
||||
gemma_root: str,
|
||||
spatial_upsampler_path: str,
|
||||
vae_queue: SimpleQueue,
|
||||
distilled_lora_path: str,
|
||||
distilled_lora_strength_stage_1: float,
|
||||
distilled_lora_strength_stage_2: float,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
|
||||
) -> None:
|
||||
distilled_lora = [LoraPathStrengthAndSDOps(distilled_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP)]
|
||||
registry = StateDictRegistry()
|
||||
pipeline = TI2VidTwoStagesHQPipeline(
|
||||
checkpoint_path=checkpoint_path,
|
||||
distilled_lora=distilled_lora,
|
||||
distilled_lora_strength_stage_1=distilled_lora_strength_stage_1,
|
||||
distilled_lora_strength_stage_2=distilled_lora_strength_stage_2,
|
||||
spatial_upsampler_path=spatial_upsampler_path,
|
||||
gemma_root=gemma_root,
|
||||
loras=(),
|
||||
registry=registry,
|
||||
quantization=_build_fp8_cast_policy(checkpoint_path),
|
||||
compilation_config=compilation_config,
|
||||
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
|
||||
)
|
||||
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
|
||||
|
||||
# Stage 1: sequence parallelism.
|
||||
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
|
||||
attn_mgr = AttentionManager(
|
||||
max_tokens=sp_max_tokens,
|
||||
num_heads=model_cfg["num_attention_heads"],
|
||||
head_dim=model_cfg["attention_head_dim"],
|
||||
tensor_dtype=pipeline.dtype,
|
||||
group=self.groups.transformer_group,
|
||||
)
|
||||
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
|
||||
inner=pipeline.stage_1._transformer_builder,
|
||||
attn_mgr=attn_mgr,
|
||||
registry=registry,
|
||||
tracker=tracker,
|
||||
)
|
||||
|
||||
# Stage 2: tiled data parallelism -- balanced 2D spatial grid over the group (one tile/rank).
|
||||
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
|
||||
tdp_height_tiles, tdp_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
|
||||
tdp_tiling = TileCountConfig(
|
||||
height=DimensionTilingConfig(num_tiles=tdp_height_tiles, overlap=5),
|
||||
width=DimensionTilingConfig(num_tiles=tdp_width_tiles, overlap=5),
|
||||
)
|
||||
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
|
||||
inner=pipeline.stage_2._transformer_builder,
|
||||
group=self.groups.transformer_group,
|
||||
tiling=tdp_tiling,
|
||||
registry=registry,
|
||||
tracker=tracker,
|
||||
)
|
||||
|
||||
# Accelerate Gemma parallelization.
|
||||
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
|
||||
gemma_root_path=gemma_root,
|
||||
gemma_group=self.groups.gemma_group,
|
||||
broadcast_group=self.groups.transformer_group,
|
||||
registry=registry,
|
||||
src_rank=_DRIVER_RANK,
|
||||
dtype=pipeline.dtype,
|
||||
)
|
||||
|
||||
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
|
||||
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
|
||||
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
|
||||
vae_tiling = TileCountConfig(
|
||||
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
|
||||
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
|
||||
)
|
||||
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder( # type: ignore[assignment]
|
||||
inner=pipeline.video_decoder._decoder_builder,
|
||||
queue=vae_queue,
|
||||
vae_group=self.groups.vae_group,
|
||||
vae_tiling=vae_tiling,
|
||||
driver_rank=_DRIVER_RANK,
|
||||
registry=registry,
|
||||
)
|
||||
|
||||
self._pipeline = pipeline
|
||||
|
||||
@torch.inference_mode()
|
||||
def __call__( # noqa: PLR0913
|
||||
self,
|
||||
*,
|
||||
output_path: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
seed: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
frame_rate: float,
|
||||
num_inference_steps: int,
|
||||
video_guider_params: MultiModalGuiderParams,
|
||||
audio_guider_params: MultiModalGuiderParams,
|
||||
images: list | None = None,
|
||||
) -> Iterator[str | None]:
|
||||
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
|
||||
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
|
||||
video, audio = self._pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
num_inference_steps=num_inference_steps,
|
||||
video_guider_params=video_guider_params,
|
||||
audio_guider_params=audio_guider_params,
|
||||
images=images or [],
|
||||
tiling_config=None,
|
||||
stage_2_sigmas=TDP_DISTILLED_SIGMAS,
|
||||
)
|
||||
if dist.get_rank() != _DRIVER_RANK:
|
||||
yield None # workers: nothing to encode
|
||||
return
|
||||
encode_video(
|
||||
video=video,
|
||||
fps=frame_rate,
|
||||
audio=audio,
|
||||
output_path=output_path,
|
||||
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
|
||||
)
|
||||
yield output_path
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from ltx_pipelines.utils.args import hq_2_stage_arg_parser
|
||||
from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
|
||||
args = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS).parse_args()
|
||||
|
||||
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
|
||||
controller = MGPUController(TI2VidTwoStagesHQRunner)
|
||||
controller.start(
|
||||
checkpoint_path=args.checkpoint_path,
|
||||
gemma_root=args.gemma_root,
|
||||
spatial_upsampler_path=args.spatial_upsampler_path,
|
||||
vae_queue=vae_queue,
|
||||
distilled_lora_path=args.distilled_lora[0].path,
|
||||
distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1,
|
||||
distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2,
|
||||
compilation_config=args.compile,
|
||||
)
|
||||
try:
|
||||
for _ in controller.stream(
|
||||
output_path=args.output_path,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
seed=args.seed,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
frame_rate=args.frame_rate,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
video_guider_params=MultiModalGuiderParams(
|
||||
cfg_scale=args.video_cfg_guidance_scale,
|
||||
stg_scale=args.video_stg_guidance_scale,
|
||||
rescale_scale=args.video_rescale_scale,
|
||||
modality_scale=args.a2v_guidance_scale,
|
||||
skip_step=args.video_skip_step,
|
||||
stg_blocks=args.video_stg_blocks,
|
||||
),
|
||||
audio_guider_params=MultiModalGuiderParams(
|
||||
cfg_scale=args.audio_cfg_guidance_scale,
|
||||
stg_scale=args.audio_stg_guidance_scale,
|
||||
rescale_scale=args.audio_rescale_scale,
|
||||
modality_scale=args.v2a_guidance_scale,
|
||||
skip_step=args.audio_skip_step,
|
||||
stg_blocks=args.audio_stg_blocks,
|
||||
),
|
||||
images=args.images,
|
||||
):
|
||||
pass # drive the job to completion; the runner writes the file as a side effect
|
||||
finally:
|
||||
controller.shutdown()
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Multi-GPU two-stage text/image-to-video runner.
|
||||
Runs :class:`TI2VidTwoStagesPipeline` across multiple GPUs with:
|
||||
- **Stage 1** -- sequence parallelism (SP)
|
||||
- **Stage 2** -- tiled data parallelism (TDP) on height + width with overlap
|
||||
- **Gemma** -- Accelerate-based parallelization
|
||||
- **VAE** -- distributed decoding
|
||||
Requires ``ltx-kernels`` to be installed (transitive via SP builder).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable, Iterator
|
||||
from multiprocessing import SimpleQueue
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from ltx_core.components.guiders import MultiModalGuiderParams
|
||||
from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
|
||||
from ltx_core.loader.registry import StateDictRegistry
|
||||
from ltx_core.model.transformer.compiling import CompilationConfig
|
||||
from ltx_core.model.video_vae import get_video_chunks_number
|
||||
from ltx_core.model.video_vae.tiling import TilingConfig
|
||||
from ltx_core.multigpu.transformer.attention import AttentionManager
|
||||
from ltx_core.quantization import QuantizationPolicy
|
||||
from ltx_core.quantization.fp8_cast import build_policy as _build_fp8_cast_policy
|
||||
from ltx_core.tiling import DimensionTilingConfig, TileCountConfig, balanced_tile_split
|
||||
from ltx_pipelines.multigpu.controller import MGPUController
|
||||
from ltx_pipelines.multigpu.gemma_builders import AccelerateGemmaBuilder
|
||||
from ltx_pipelines.multigpu.runner import MGPURunner
|
||||
from ltx_pipelines.multigpu.sp_builder import SequenceParallelBuilder
|
||||
from ltx_pipelines.multigpu.tdp_builder import TiledDataParallelBuilder
|
||||
from ltx_pipelines.multigpu.vae_builders import DistributedDecoderBuilder
|
||||
from ltx_pipelines.multigpu.weight_tracker import TransformerWeightTracker
|
||||
from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.constants import TDP_DISTILLED_SIGMAS
|
||||
from ltx_pipelines.utils.media_io import encode_video
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Stage 1 at 512x768, 121 frames = 6144 video tokens + audio tokens.
|
||||
_DEFAULT_SP_MAX_TOKENS = 32768
|
||||
# Rank that collects distributed-VAE tiles and encodes the assembled video.
|
||||
_DRIVER_RANK = 0
|
||||
|
||||
|
||||
class TI2VidTwoStagesRunner(MGPURunner):
|
||||
"""Distributed :class:`TI2VidTwoStagesPipeline`: SP stage 1 + TDP stage 2 + Gemma + distributed VAE."""
|
||||
|
||||
@torch.inference_mode()
|
||||
def setup(
|
||||
self,
|
||||
*,
|
||||
checkpoint_path: str,
|
||||
gemma_root: str,
|
||||
spatial_upsampler_path: str,
|
||||
vae_queue: SimpleQueue,
|
||||
distilled_lora_path: str,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
sp_max_tokens: int = _DEFAULT_SP_MAX_TOKENS,
|
||||
quantization: Callable[[], QuantizationPolicy] | None = None,
|
||||
) -> None:
|
||||
# quantization is a picklable zero-arg builder (built per worker, post-spawn); default fp8-cast.
|
||||
quantization_policy = quantization() if quantization is not None else _build_fp8_cast_policy(checkpoint_path)
|
||||
distilled_lora = [LoraPathStrengthAndSDOps(distilled_lora_path, 1.0, LTXV_LORA_COMFY_RENAMING_MAP)]
|
||||
registry = StateDictRegistry()
|
||||
pipeline = TI2VidTwoStagesPipeline(
|
||||
checkpoint_path=checkpoint_path,
|
||||
distilled_lora=distilled_lora,
|
||||
spatial_upsampler_path=spatial_upsampler_path,
|
||||
gemma_root=gemma_root,
|
||||
loras=[],
|
||||
registry=registry,
|
||||
quantization=quantization_policy,
|
||||
compilation_config=compilation_config,
|
||||
alloc_trim_strategy=AllocatorTrimStrategy.DEFER,
|
||||
)
|
||||
tracker = TransformerWeightTracker(group=self.groups.transformer_group)
|
||||
|
||||
# Stage 1: sequence parallelism.
|
||||
model_cfg = pipeline.stage_1._transformer_builder.model_config().get("transformer", {})
|
||||
attn_mgr = AttentionManager(
|
||||
max_tokens=sp_max_tokens,
|
||||
num_heads=model_cfg["num_attention_heads"],
|
||||
head_dim=model_cfg["attention_head_dim"],
|
||||
tensor_dtype=pipeline.dtype,
|
||||
group=self.groups.transformer_group,
|
||||
)
|
||||
pipeline.stage_1._transformer_builder = SequenceParallelBuilder(
|
||||
inner=pipeline.stage_1._transformer_builder,
|
||||
attn_mgr=attn_mgr,
|
||||
registry=registry,
|
||||
tracker=tracker,
|
||||
)
|
||||
|
||||
# Stage 2: tiled data parallelism -- balanced 2D spatial grid over the group (one tile/rank).
|
||||
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
|
||||
tdp_height_tiles, tdp_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.transformer_group))
|
||||
tdp_tiling = TileCountConfig(
|
||||
height=DimensionTilingConfig(num_tiles=tdp_height_tiles, overlap=5),
|
||||
width=DimensionTilingConfig(num_tiles=tdp_width_tiles, overlap=5),
|
||||
)
|
||||
pipeline.stage_2._transformer_builder = TiledDataParallelBuilder(
|
||||
inner=pipeline.stage_2._transformer_builder,
|
||||
group=self.groups.transformer_group,
|
||||
tiling=tdp_tiling,
|
||||
registry=registry,
|
||||
tracker=tracker,
|
||||
)
|
||||
|
||||
# Accelerate Gemma parallelization.
|
||||
pipeline.prompt_encoder._text_encoder_builder = AccelerateGemmaBuilder(
|
||||
gemma_root_path=gemma_root,
|
||||
gemma_group=self.groups.gemma_group,
|
||||
broadcast_group=self.groups.transformer_group,
|
||||
registry=registry,
|
||||
src_rank=_DRIVER_RANK,
|
||||
dtype=pipeline.dtype,
|
||||
)
|
||||
|
||||
# Distributed VAE decoding: balanced 2D spatial grid over the group (one tile/rank).
|
||||
# height takes the smaller factor of world_size, width the larger; size-aware split is a follow-up.
|
||||
vae_height_tiles, vae_width_tiles = balanced_tile_split(dist.get_world_size(self.groups.vae_group))
|
||||
vae_tiling = TileCountConfig(
|
||||
height=DimensionTilingConfig(num_tiles=vae_height_tiles, overlap=4),
|
||||
width=DimensionTilingConfig(num_tiles=vae_width_tiles, overlap=4),
|
||||
)
|
||||
pipeline.video_decoder._decoder_builder = DistributedDecoderBuilder(
|
||||
inner=pipeline.video_decoder._decoder_builder,
|
||||
queue=vae_queue,
|
||||
vae_group=self.groups.vae_group,
|
||||
vae_tiling=vae_tiling,
|
||||
driver_rank=_DRIVER_RANK,
|
||||
registry=registry,
|
||||
)
|
||||
|
||||
self._pipeline = pipeline
|
||||
|
||||
@torch.inference_mode()
|
||||
def __call__( # noqa: PLR0913
|
||||
self,
|
||||
*,
|
||||
output_path: str,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
seed: int,
|
||||
height: int,
|
||||
width: int,
|
||||
num_frames: int,
|
||||
frame_rate: float,
|
||||
num_inference_steps: int,
|
||||
video_guider_params: MultiModalGuiderParams,
|
||||
audio_guider_params: MultiModalGuiderParams,
|
||||
images: list | None = None,
|
||||
) -> Iterator[str | None]:
|
||||
# The pipeline raises ValueError on invalid input (symmetric across ranks); the controller
|
||||
# catches that and turns it into a recoverable RunnerError. Anything else is fatal.
|
||||
video, audio = self._pipeline(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
seed=seed,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
num_inference_steps=num_inference_steps,
|
||||
video_guider_params=video_guider_params,
|
||||
audio_guider_params=audio_guider_params,
|
||||
images=images or [],
|
||||
tiling_config=None,
|
||||
stage_2_sigmas=TDP_DISTILLED_SIGMAS,
|
||||
)
|
||||
if dist.get_rank() != _DRIVER_RANK:
|
||||
yield None # workers: nothing to encode
|
||||
return
|
||||
encode_video(
|
||||
video=video,
|
||||
fps=frame_rate,
|
||||
audio=audio,
|
||||
output_path=output_path,
|
||||
video_chunks_number=get_video_chunks_number(num_frames, TilingConfig.default()),
|
||||
)
|
||||
yield output_path
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from ltx_pipelines.utils.args import (
|
||||
default_2_stage_arg_parser,
|
||||
resolve_cli_params,
|
||||
)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s")
|
||||
params = resolve_cli_params()
|
||||
args = default_2_stage_arg_parser(params=params).parse_args()
|
||||
|
||||
vae_queue = torch.multiprocessing.get_context("spawn").SimpleQueue()
|
||||
controller = MGPUController(TI2VidTwoStagesRunner)
|
||||
controller.start(
|
||||
checkpoint_path=args.checkpoint_path,
|
||||
gemma_root=args.gemma_root,
|
||||
spatial_upsampler_path=args.spatial_upsampler_path,
|
||||
vae_queue=vae_queue,
|
||||
distilled_lora_path=args.distilled_lora[0].path,
|
||||
compilation_config=args.compile,
|
||||
)
|
||||
try:
|
||||
for _ in controller.stream(
|
||||
output_path=args.output_path,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
seed=args.seed,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
frame_rate=args.frame_rate,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
video_guider_params=MultiModalGuiderParams(
|
||||
cfg_scale=args.video_cfg_guidance_scale,
|
||||
stg_scale=args.video_stg_guidance_scale,
|
||||
rescale_scale=args.video_rescale_scale,
|
||||
modality_scale=args.a2v_guidance_scale,
|
||||
skip_step=args.video_skip_step,
|
||||
stg_blocks=args.video_stg_blocks,
|
||||
),
|
||||
audio_guider_params=MultiModalGuiderParams(
|
||||
cfg_scale=args.audio_cfg_guidance_scale,
|
||||
stg_scale=args.audio_stg_guidance_scale,
|
||||
rescale_scale=args.audio_rescale_scale,
|
||||
modality_scale=args.v2a_guidance_scale,
|
||||
skip_step=args.audio_skip_step,
|
||||
stg_blocks=args.audio_stg_blocks,
|
||||
),
|
||||
images=args.images,
|
||||
):
|
||||
pass # drive the job to completion; the runner writes the file as a side effect
|
||||
finally:
|
||||
controller.shutdown()
|
||||
@@ -0,0 +1,8 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class AllocatorTrimStrategy(Enum):
|
||||
"""How a block releases its model's memory when its scope exits."""
|
||||
|
||||
TRIM = "trim" # sync, release storage (to meta), and empty_cache() back to the OS
|
||||
DEFER = "defer" # skip teardown; let GC reclaim it and keep the CUDA cache warm
|
||||
@@ -1,5 +1,6 @@
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
@@ -14,6 +15,7 @@ from ltx_pipelines.utils.constants import (
|
||||
LTX_2_3_HQ_PARAMS,
|
||||
LTX_2_3_PARAMS,
|
||||
PipelineParams,
|
||||
detect_params,
|
||||
)
|
||||
from ltx_pipelines.utils.quantization_factory import QuantizationKind
|
||||
from ltx_pipelines.utils.types import OffloadMode
|
||||
@@ -263,6 +265,24 @@ def detect_checkpoint_path(distilled: bool = False) -> str:
|
||||
return known.distilled_checkpoint_path if distilled else known.checkpoint_path
|
||||
|
||||
|
||||
def help_requested() -> bool:
|
||||
"""Whether ``-h``/``--help`` appears on the command line."""
|
||||
return "-h" in sys.argv or "--help" in sys.argv
|
||||
|
||||
|
||||
def resolve_cli_params(distilled: bool = False) -> PipelineParams:
|
||||
"""Return the model params a pipeline CLI uses to build its argument parser.
|
||||
Reads the model version from the checkpoint named on the command line so the
|
||||
parser's defaults match the target model.
|
||||
Args:
|
||||
distilled: Whether the pipeline takes a distilled checkpoint
|
||||
(``--distilled-checkpoint-path``) rather than a full one (``--checkpoint-path``).
|
||||
"""
|
||||
if help_requested():
|
||||
return LTX_2_3_PARAMS
|
||||
return detect_params(detect_checkpoint_path(distilled=distilled))
|
||||
|
||||
|
||||
def basic_arg_parser(
|
||||
params: PipelineParams = LTX_2_3_PARAMS,
|
||||
distilled: bool = False,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
"""Pipeline blocks — each block owns its model lifecycle.
|
||||
Blocks build a model on each ``__call__``, use it, then free GPU memory.
|
||||
This eliminates manual ``del model; cleanup_memory()`` in pipelines and
|
||||
removes the need for :class:`ModelLedger`.
|
||||
This eliminates manual ``del model; cleanup_memory()`` in pipelines: each
|
||||
block is self-contained, so no central model-coordinator object is needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -77,6 +77,7 @@ from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcesso
|
||||
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
|
||||
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
|
||||
from ltx_core.utils import find_matching_file
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.gpu_model import gpu_model
|
||||
from ltx_pipelines.utils.helpers import (
|
||||
cleanup_memory,
|
||||
@@ -136,23 +137,25 @@ def _apply_compile_ops(
|
||||
@contextmanager
|
||||
def _streaming_model(
|
||||
builder: StreamingModelBuilder,
|
||||
offload_mode: OffloadMode,
|
||||
target_device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> Iterator:
|
||||
"""Build a streaming wrapper, yield it, then tear down and free memory."""
|
||||
cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
|
||||
wrapped = builder.build(
|
||||
device=target_device,
|
||||
dtype=dtype,
|
||||
cpu_slots_count=cpu_slots_count,
|
||||
)
|
||||
"""Build a streaming wrapper, yield it, then tear down and free memory.
|
||||
The builder's own ``cpu_slots_count`` selects RAM vs disk streaming.
|
||||
``teardown()`` always runs -- it releases non-memory resources (forward
|
||||
hooks, the disk I/O worker thread, open file handles) that GC would not
|
||||
reclaim promptly. ``alloc_trim_strategy=DEFER`` only skips the eager allocator
|
||||
reclaim (``to("meta")`` + ``cleanup_memory()``), leaving param storage for GC.
|
||||
"""
|
||||
wrapped = builder.build(device=target_device, dtype=dtype)
|
||||
try:
|
||||
yield wrapped
|
||||
finally:
|
||||
wrapped.teardown()
|
||||
wrapped.to("meta")
|
||||
cleanup_memory()
|
||||
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
|
||||
wrapped.to("meta")
|
||||
cleanup_memory()
|
||||
|
||||
|
||||
def _build_state(
|
||||
@@ -177,9 +180,13 @@ def _build_state(
|
||||
return state
|
||||
|
||||
|
||||
def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterator[torch.Tensor]:
|
||||
"""Wrap an iterator to clean up *model* memory once it is exhausted or abandoned."""
|
||||
with gpu_model(model):
|
||||
def _cleanup_iter(
|
||||
it: Iterator[torch.Tensor],
|
||||
model: torch.nn.Module,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> Iterator[torch.Tensor]:
|
||||
"""Wrap an iterator to release *model* memory (per ``alloc_trim_strategy``) once exhausted or abandoned."""
|
||||
with gpu_model(model, alloc_trim_strategy=alloc_trim_strategy):
|
||||
yield from it
|
||||
|
||||
|
||||
@@ -190,12 +197,41 @@ def _cleanup_iter(it: Iterator[torch.Tensor], model: torch.nn.Module) -> Iterato
|
||||
|
||||
class DiffusionStage:
|
||||
"""Owns transformer lifecycle. Builds on each call, frees on exit.
|
||||
Replaces the manual ``model_ledger.transformer()`` / ``del transformer``
|
||||
pattern in every pipeline.
|
||||
Replaces the manual build-transformer / ``del transformer`` pattern that
|
||||
every pipeline previously repeated.
|
||||
"""
|
||||
|
||||
def __init__( # noqa: PLR0913
|
||||
def __init__(
|
||||
self,
|
||||
transformer_builder: ModelBuilderProtocol[LTXModelProtocol],
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
*,
|
||||
quantization: QuantizationPolicy | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
"""Construct a stage from a single pre-built transformer ``builder``.
|
||||
Holds only that builder plus build-time configuration (dtype, device,
|
||||
quantization, compilation). Turning a checkpoint path + LoRA set into a
|
||||
builder -- and choosing a :class:`StreamingModelBuilder` when offloading --
|
||||
lives in :meth:`from_checkpoint`, which is how pipelines normally create a
|
||||
stage. A :class:`StreamingModelBuilder` selects the block-streaming build
|
||||
path; any other builder uses the standard (all-on-GPU) path.
|
||||
``quantization`` and ``compilation_config`` are applied lazily on the
|
||||
standard path; on the streaming path they are already baked into the
|
||||
streaming builder by :meth:`from_checkpoint` and these fields are unused.
|
||||
"""
|
||||
self._transformer_builder = transformer_builder
|
||||
self._dtype = dtype
|
||||
self._device = device
|
||||
self._quantization = quantization
|
||||
self._compilation_config = compilation_config
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint( # noqa: PLR0913
|
||||
cls,
|
||||
checkpoint_path: str,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
@@ -203,17 +239,24 @@ class DiffusionStage:
|
||||
quantization: QuantizationPolicy | None = None,
|
||||
registry: Registry | None = None,
|
||||
compilation_config: CompilationConfig | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
transformer_builder: ModelBuilderProtocol[LTXModelProtocol] | None = None,
|
||||
model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
|
||||
model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
|
||||
) -> None:
|
||||
self._checkpoint_path = checkpoint_path
|
||||
self._dtype = dtype
|
||||
self._device = device
|
||||
self._quantization = quantization
|
||||
self._compilation_config = compilation_config
|
||||
self._offload_mode = offload_mode
|
||||
) -> "DiffusionStage":
|
||||
"""Build a stage from a checkpoint path and LoRA set.
|
||||
Constructs a single transformer builder from ``checkpoint_path`` +
|
||||
``loras`` + ``quantization`` and delegates to ``__init__``. When
|
||||
``offload_mode != OffloadMode.NONE`` that builder is a
|
||||
:class:`StreamingModelBuilder` (with quantization/compilation baked in and
|
||||
its ``cpu_slots_count`` set for the requested mode); otherwise it is the
|
||||
standard single-GPU builder. This is the high-level entry point used by
|
||||
pipelines; ``__init__`` itself takes an already-built builder.
|
||||
``model_configurator`` / ``model_sd_ops`` let callers (e.g. the audio-only
|
||||
T2A pipeline) override the model class configurator and the state-dict key
|
||||
mapping. A quantization policy that pins its own configurator takes
|
||||
precedence over ``model_configurator``.
|
||||
"""
|
||||
# A quantization policy may pin its own configurator; otherwise use the one
|
||||
# provided by the caller (defaults to the audio-video LTXModelConfigurator).
|
||||
configurator = (
|
||||
@@ -221,51 +264,75 @@ class DiffusionStage:
|
||||
if quantization is not None and quantization.model_configurator is not None
|
||||
else model_configurator
|
||||
)
|
||||
if transformer_builder is not None:
|
||||
self._transformer_builder = transformer_builder
|
||||
else:
|
||||
self._transformer_builder = Builder(
|
||||
|
||||
transformer_builder: ModelBuilderProtocol[LTXModelProtocol]
|
||||
if offload_mode == OffloadMode.NONE:
|
||||
transformer_builder = Builder(
|
||||
model_path=checkpoint_path,
|
||||
model_class_configurator=configurator,
|
||||
model_sd_ops=model_sd_ops,
|
||||
loras=tuple(loras),
|
||||
registry=registry or DummyRegistry(),
|
||||
)
|
||||
|
||||
if offload_mode != OffloadMode.NONE:
|
||||
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
|
||||
# (no companion-key emission). Quantization policies that emit
|
||||
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
|
||||
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
|
||||
raise ValueError(
|
||||
"Block streaming is not supported with this quantization policy "
|
||||
"(only bf16 and fp8_cast are currently supported)."
|
||||
)
|
||||
streaming_sd_ops: SDOps = model_sd_ops
|
||||
streaming_module_ops: tuple[ModuleOps, ...] = ()
|
||||
streaming_loras = tuple(loras)
|
||||
|
||||
if compilation_config:
|
||||
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
|
||||
streaming_sd_ops, streaming_module_ops, streaming_loras = _apply_compile_ops(
|
||||
streaming_sd_ops, streaming_module_ops, streaming_loras, number_of_layers
|
||||
)
|
||||
if quantization is not None:
|
||||
streaming_sd_ops, streaming_module_ops = _chain_quantization(
|
||||
streaming_sd_ops, streaming_module_ops, quantization
|
||||
)
|
||||
self._streaming_builder = StreamingModelBuilder(
|
||||
model_class_configurator=configurator,
|
||||
model_path=checkpoint_path,
|
||||
model_sd_ops=streaming_sd_ops,
|
||||
module_ops=streaming_module_ops,
|
||||
loras=streaming_loras,
|
||||
else:
|
||||
transformer_builder = cls._build_streaming_builder(
|
||||
checkpoint_path=checkpoint_path,
|
||||
configurator=configurator,
|
||||
model_sd_ops=model_sd_ops,
|
||||
loras=tuple(loras),
|
||||
quantization=quantization,
|
||||
registry=registry or DummyRegistry(),
|
||||
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
|
||||
blocks_attr="transformer_blocks",
|
||||
blocks_prefix="transformer_blocks",
|
||||
offload_mode=offload_mode,
|
||||
)
|
||||
|
||||
return cls(
|
||||
transformer_builder,
|
||||
dtype,
|
||||
device,
|
||||
quantization=quantization,
|
||||
compilation_config=compilation_config,
|
||||
alloc_trim_strategy=alloc_trim_strategy,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_streaming_builder(
|
||||
*,
|
||||
checkpoint_path: str,
|
||||
configurator: type[ModelConfigurator],
|
||||
model_sd_ops: SDOps,
|
||||
loras: tuple[LoraPathStrengthAndSDOps, ...],
|
||||
quantization: QuantizationPolicy | None,
|
||||
registry: Registry,
|
||||
offload_mode: OffloadMode,
|
||||
) -> StreamingModelBuilder:
|
||||
"""Construct the streaming transformer builder for an offloading stage.
|
||||
Holds only raw config (``model_sd_ops`` / ``loras``); compilation and
|
||||
quantization are applied at build time by :meth:`_prepared_builder`, exactly
|
||||
as on the standard path -- so the builder's LoRA set stays raw and
|
||||
:meth:`with_loras` swaps it consistently. ``cpu_slots_count`` is pinned for
|
||||
the requested ``offload_mode`` (disk streaming uses a small slot count;
|
||||
CPU/RAM streaming pins every block).
|
||||
"""
|
||||
# WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
|
||||
# (no companion-key emission). Quantization policies that emit
|
||||
# companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
|
||||
if quantization is not None and quantization.fuse_rule is not fp8_cast_fuse_rule:
|
||||
raise ValueError(
|
||||
"Block streaming is not supported with this quantization policy "
|
||||
"(only bf16 and fp8_cast are currently supported)."
|
||||
)
|
||||
return StreamingModelBuilder(
|
||||
model_class_configurator=configurator,
|
||||
model_path=checkpoint_path,
|
||||
model_sd_ops=model_sd_ops,
|
||||
loras=loras,
|
||||
registry=registry,
|
||||
fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
|
||||
blocks_attr="transformer_blocks",
|
||||
blocks_prefix="transformer_blocks",
|
||||
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
|
||||
)
|
||||
|
||||
def with_attention(self, attention: AttentionFunction | AttentionCallable | None) -> "DiffusionStage":
|
||||
"""Return a new ``DiffusionStage`` that pins the transformer build to ``attention``.
|
||||
Functional: never mutates ``self``. The returned stage shares all other
|
||||
@@ -280,41 +347,64 @@ class DiffusionStage:
|
||||
new._transformer_builder = self._transformer_builder.with_module_ops(
|
||||
(*self._transformer_builder.module_ops, op),
|
||||
)
|
||||
if self._offload_mode != OffloadMode.NONE:
|
||||
new._streaming_builder = self._streaming_builder.with_module_ops(
|
||||
(*self._streaming_builder.module_ops, op),
|
||||
)
|
||||
return new
|
||||
|
||||
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
|
||||
target = device or self._device
|
||||
sd_ops = self._transformer_builder.model_sd_ops
|
||||
module_ops = self._transformer_builder.module_ops
|
||||
loras = self._transformer_builder.loras
|
||||
def with_builder(self, builder: ModelBuilderProtocol[LTXModelProtocol]) -> "DiffusionStage":
|
||||
"""Return a new ``DiffusionStage`` that builds its transformer from ``builder``.
|
||||
Functional: never mutates ``self``; shares all other configuration (dtype, device,
|
||||
quantization, compilation). Affects the standard (non-offload) build path.
|
||||
"""
|
||||
new = copy.copy(self)
|
||||
new._transformer_builder = builder
|
||||
return new
|
||||
|
||||
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "DiffusionStage":
|
||||
"""Return a new ``DiffusionStage`` built with exactly ``loras`` (replacing the current set)."""
|
||||
return self.with_builder(self._transformer_builder.with_loras(loras))
|
||||
|
||||
def _prepared_builder(self) -> ModelBuilderProtocol[LTXModelProtocol]:
|
||||
"""Return the configured builder with the stage's build-time ops applied.
|
||||
Compilation and quantization live on the stage (not on the builder) and are
|
||||
applied here, lazily, for both the standard and streaming paths. This keeps
|
||||
the builder holding only raw sd_ops/module_ops/LoRAs, so ``with_loras`` /
|
||||
``with_builder`` swap them consistently regardless of the build path. The
|
||||
returned copy preserves the builder's concrete type (e.g. a
|
||||
``StreamingModelBuilder`` stays one).
|
||||
"""
|
||||
builder = self._transformer_builder
|
||||
sd_ops = builder.model_sd_ops
|
||||
module_ops = builder.module_ops
|
||||
loras = builder.loras
|
||||
if self._compilation_config is not None:
|
||||
number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
|
||||
number_of_layers = builder.model_config()["transformer"]["num_layers"]
|
||||
sd_ops, module_ops, loras = _apply_compile_ops(
|
||||
sd_ops, module_ops, loras, number_of_layers, self._compilation_config
|
||||
)
|
||||
if self._quantization is not None:
|
||||
sd_ops, module_ops = _chain_quantization(sd_ops, module_ops, self._quantization)
|
||||
|
||||
builder = self._transformer_builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
|
||||
if self._quantization is not None:
|
||||
builder = builder.with_fuse_rule(self._quantization.fuse_rule)
|
||||
return X0Model(builder.build(device=target, **kwargs)).to(target).eval()
|
||||
return builder.with_module_ops(module_ops).with_sd_ops(sd_ops).with_loras(loras)
|
||||
|
||||
def _build_transformer(self, *, device: torch.device | None = None, **kwargs: object) -> X0Model:
|
||||
target = device or self._device
|
||||
return X0Model(self._prepared_builder().build(device=target, **kwargs)).to(target).eval()
|
||||
|
||||
@property
|
||||
def _is_streaming(self) -> bool:
|
||||
"""Whether the configured builder uses the block-streaming build path."""
|
||||
return isinstance(self._transformer_builder, StreamingModelBuilder)
|
||||
|
||||
@contextmanager
|
||||
def _streaming_transformer_ctx(self) -> Iterator[X0Model]:
|
||||
with _streaming_model(
|
||||
self._streaming_builder, self._offload_mode, self._device, self._dtype
|
||||
) as streaming_wrapper:
|
||||
builder = self._prepared_builder()
|
||||
assert isinstance(builder, StreamingModelBuilder)
|
||||
with _streaming_model(builder, self._device, self._dtype, self._alloc_trim_strategy) as streaming_wrapper:
|
||||
yield X0Model(streaming_wrapper).eval()
|
||||
|
||||
def _transformer_ctx(self, **kwargs: object) -> AbstractContextManager:
|
||||
if self._offload_mode != OffloadMode.NONE:
|
||||
if self._is_streaming:
|
||||
return self._streaming_transformer_ctx()
|
||||
return gpu_model(self._build_transformer(**kwargs))
|
||||
return gpu_model(self._build_transformer(**kwargs), alloc_trim_strategy=self._alloc_trim_strategy)
|
||||
|
||||
def model_context(self, **kwargs: object) -> AbstractContextManager:
|
||||
"""Build the transformer, yield it, then free its memory on exit.
|
||||
@@ -419,8 +509,8 @@ class DiffusionStage:
|
||||
v_shape = VideoLatentShape.from_pixel_shape(pixel_shape)
|
||||
video_tools = VideoLatentTools(VideoLatentPatchifier(patch_size=1), v_shape, fps)
|
||||
|
||||
mode = "streaming" if self._offload_mode != OffloadMode.NONE else "standard"
|
||||
logger.info("Building transformer (%s) from %s", mode, self._checkpoint_path)
|
||||
mode = "streaming" if self._is_streaming else "standard"
|
||||
logger.info("Building transformer (%s) from %s", mode, self._transformer_builder.checkpoint)
|
||||
with self._transformer_ctx(video_tools=video_tools) as transformer:
|
||||
logger.info(
|
||||
"Running denoising loop (%d steps, %dx%d %d frames @ %.1f fps)",
|
||||
@@ -467,12 +557,14 @@ class PromptEncoder:
|
||||
registry: Registry | None = None,
|
||||
offload_mode: OffloadMode = OffloadMode.NONE,
|
||||
text_encoder_builder: BuilderProtocol | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._gemma_root = gemma_root
|
||||
self._checkpoint_path = checkpoint_path
|
||||
self._dtype = dtype
|
||||
self._device = device
|
||||
self._offload_mode = offload_mode
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
if text_encoder_builder is not None:
|
||||
if offload_mode != OffloadMode.NONE:
|
||||
@@ -501,6 +593,7 @@ class PromptEncoder:
|
||||
registry=registry or DummyRegistry(),
|
||||
blocks_attr="model.model.language_model.layers",
|
||||
blocks_prefix="model.model.language_model.layers",
|
||||
cpu_slots_count=DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None,
|
||||
)
|
||||
self._embeddings_processor_builder = Builder(
|
||||
model_path=checkpoint_path,
|
||||
@@ -519,8 +612,10 @@ class PromptEncoder:
|
||||
|
||||
def _text_encoder_ctx(self) -> AbstractContextManager:
|
||||
if self._offload_mode != OffloadMode.NONE:
|
||||
return _streaming_model(self._streaming_text_encoder_builder, self._offload_mode, self._device, self._dtype)
|
||||
return gpu_model(self._build_text_encoder())
|
||||
return _streaming_model(
|
||||
self._streaming_text_encoder_builder, self._device, self._dtype, self._alloc_trim_strategy
|
||||
)
|
||||
return gpu_model(self._build_text_encoder(), alloc_trim_strategy=self._alloc_trim_strategy)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
@@ -541,7 +636,9 @@ class PromptEncoder:
|
||||
raw_outputs = text_encoder.encode(prompts)
|
||||
logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
|
||||
|
||||
with gpu_model(self._build_embeddings_processor()) as embeddings_processor:
|
||||
with gpu_model(
|
||||
self._build_embeddings_processor(), alloc_trim_strategy=self._alloc_trim_strategy
|
||||
) as embeddings_processor:
|
||||
result = [embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs]
|
||||
logger.info("Prompt encoding complete")
|
||||
return result
|
||||
@@ -563,6 +660,7 @@ class ImageConditioner:
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
registry: Registry | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._dtype = dtype
|
||||
self._device = device
|
||||
@@ -572,13 +670,14 @@ class ImageConditioner:
|
||||
model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
|
||||
registry=registry or DummyRegistry(),
|
||||
)
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
def _build_encoder(self) -> VideoEncoder:
|
||||
return self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()
|
||||
|
||||
def __call__(self, fn: Callable[[VideoEncoder], T]) -> T:
|
||||
"""Build video encoder → call *fn(encoder)* → free encoder."""
|
||||
with gpu_model(self._build_encoder()) as encoder:
|
||||
with gpu_model(self._build_encoder(), alloc_trim_strategy=self._alloc_trim_strategy) as encoder:
|
||||
return fn(encoder)
|
||||
|
||||
|
||||
@@ -597,6 +696,7 @@ class VideoUpsampler:
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
registry: Registry | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._upsampler_path = upsampler_path
|
||||
self._dtype = dtype
|
||||
@@ -612,13 +712,20 @@ class VideoUpsampler:
|
||||
model_class_configurator=LatentUpsamplerConfigurator,
|
||||
registry=registry or DummyRegistry(),
|
||||
)
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
def __call__(self, latent: torch.Tensor) -> torch.Tensor:
|
||||
"""Upsample *latent* using video encoder + spatial upsampler, then free both."""
|
||||
logger.info("Building video encoder + spatial upsampler from %s", self._upsampler_path)
|
||||
with (
|
||||
gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder,
|
||||
gpu_model(self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval()) as upsampler,
|
||||
gpu_model(
|
||||
self._encoder_builder.build(device=self._device, dtype=self._dtype).eval(),
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
) as encoder,
|
||||
gpu_model(
|
||||
self._upsampler_builder.build(device=self._device, dtype=self._dtype).eval(),
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
) as upsampler,
|
||||
):
|
||||
return upsample_video(latent=latent, video_encoder=encoder, upsampler=upsampler)
|
||||
|
||||
@@ -641,6 +748,7 @@ class VideoDecoder:
|
||||
registry: Registry | None = None,
|
||||
memory_efficient: bool = True,
|
||||
decoder_builder: BuilderProtocol | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._checkpoint_path = checkpoint_path
|
||||
self._dtype = dtype
|
||||
@@ -655,6 +763,7 @@ class VideoDecoder:
|
||||
registry=registry or DummyRegistry(),
|
||||
module_ops=(MEMORY_EFFICIENT_DECODE,) if memory_efficient else (),
|
||||
)
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
@@ -665,7 +774,11 @@ class VideoDecoder:
|
||||
"""Decode *latent* to pixel-space video chunks. Decoder freed after exhaustion."""
|
||||
logger.info("Building video decoder from %s", self._checkpoint_path)
|
||||
decoder = self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()
|
||||
return _cleanup_iter(decoder.decode_video(latent, tiling_config, generator), decoder)
|
||||
return _cleanup_iter(
|
||||
decoder.decode_video(latent, tiling_config, generator),
|
||||
decoder,
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -682,6 +795,7 @@ class AudioDecoder:
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
registry: Registry | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._checkpoint_path = checkpoint_path
|
||||
self._dtype = dtype
|
||||
@@ -698,13 +812,25 @@ class AudioDecoder:
|
||||
model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
|
||||
registry=registry or DummyRegistry(),
|
||||
)
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
|
||||
def __call__(self, latent: torch.Tensor) -> Audio:
|
||||
"""Decode audio *latent* through VAE decoder + vocoder, then free both."""
|
||||
logger.info("Building audio decoder + vocoder from %s", self._checkpoint_path)
|
||||
# The vocoder always runs in fp32 (bf16 accumulation degrades spectral
|
||||
# metrics). On CUDA/CPU it is stored in bf16 and autocast upcasts per-op to
|
||||
# save memory; MPS has no fp32 autocast, so store it in fp32 directly and
|
||||
# avoid the per-call cast. Negligible footprint for this small model.
|
||||
vocoder_dtype = torch.float32 if self._device.type == "mps" else self._dtype
|
||||
with (
|
||||
gpu_model(self._decoder_builder.build(device=self._device, dtype=self._dtype).eval()) as decoder,
|
||||
gpu_model(self._vocoder_builder.build(device=self._device, dtype=self._dtype).eval()) as vocoder,
|
||||
gpu_model(
|
||||
self._decoder_builder.build(device=self._device, dtype=self._dtype).eval(),
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
) as decoder,
|
||||
gpu_model(
|
||||
self._vocoder_builder.build(device=self._device, dtype=vocoder_dtype).eval(),
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
) as vocoder,
|
||||
):
|
||||
return vae_decode_audio(latent, decoder, vocoder)
|
||||
|
||||
@@ -726,9 +852,11 @@ class AudioConditioner:
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
registry: Registry | None = None,
|
||||
alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM,
|
||||
) -> None:
|
||||
self._dtype = dtype
|
||||
self._device = device
|
||||
self._alloc_trim_strategy = alloc_trim_strategy
|
||||
self._encoder_builder = Builder(
|
||||
model_path=checkpoint_path,
|
||||
model_class_configurator=AudioEncoderConfigurator,
|
||||
@@ -738,5 +866,8 @@ class AudioConditioner:
|
||||
|
||||
def __call__(self, fn: Callable[[torch.nn.Module], T]) -> T:
|
||||
"""Build audio encoder → call *fn(encoder)* → free encoder."""
|
||||
with gpu_model(self._encoder_builder.build(device=self._device, dtype=self._dtype).eval()) as encoder:
|
||||
with gpu_model(
|
||||
self._encoder_builder.build(device=self._device, dtype=self._dtype).eval(),
|
||||
alloc_trim_strategy=self._alloc_trim_strategy,
|
||||
) as encoder:
|
||||
return fn(encoder)
|
||||
|
||||
@@ -20,6 +20,8 @@ STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
|
||||
|
||||
DISTILLED_SIGMAS = torch.tensor(DISTILLED_SIGMA_VALUES)
|
||||
STAGE_2_DISTILLED_SIGMAS = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES)
|
||||
# Stage 2 schedule for the tiled-data-parallel multi-GPU runner.
|
||||
TDP_DISTILLED_SIGMAS = torch.tensor([0.625, 0.4, 0.0])
|
||||
|
||||
|
||||
# =============================================================================
|
||||
|
||||
@@ -164,17 +164,24 @@ def _guided_denoise( # noqa: PLR0913,PLR0915
|
||||
enabled=not a_skip,
|
||||
)
|
||||
|
||||
# Replicate each pass's PerturbationConfig to all `orig_b` samples it
|
||||
# carries, so `BatchedPerturbationConfig.mask_like` returns a per-sample
|
||||
# mask (length n*orig_b) instead of a per-pass mask (length n). Without
|
||||
# this expansion the mask is broadcast against a (n*orig_b, T, D) tensor
|
||||
# and the multiplication fails with a batch-dim mismatch whenever
|
||||
# `orig_b > 1` (e.g. multi-prompt benchmark panels).
|
||||
# Replicate each pass's PerturbationConfig to all `orig_b` samples it carries, so the keep-mask
|
||||
# has one row per sample (length n*orig_b) instead of per-pass (length n). Without this
|
||||
# expansion the mask broadcasts against a (n*orig_b, T, D) tensor and the multiplication fails
|
||||
# with a batch-dim mismatch whenever `orig_b > 1` (e.g. multi-prompt benchmark panels).
|
||||
batched_ptb_configs = [ptb for ptb in ptb_configs for _ in range(orig_b)]
|
||||
|
||||
all_v, all_a = transformer(
|
||||
video=batched_video, audio=batched_audio, perturbations=BatchedPerturbationConfig(batched_ptb_configs)
|
||||
# Build the config with num_blocks/device/dtype so it precomputes its per-block mask tensor
|
||||
# on init (transformer.num_blocks delegates through the SP/TDP/BatchSplit/X0 wrappers). The
|
||||
# compiled forward then reads perturbation as a runtime tensor, not by querying the config
|
||||
# in-graph -- so it doesn't recompile per perturbation config.
|
||||
ref_modality = batched_video if batched_video is not None else batched_audio
|
||||
perturbations = BatchedPerturbationConfig(
|
||||
batched_ptb_configs,
|
||||
num_blocks=transformer.num_blocks,
|
||||
device=ref_modality.latent.device,
|
||||
dtype=ref_modality.latent.dtype,
|
||||
)
|
||||
all_v, all_a = transformer(video=batched_video, audio=batched_audio, perturbations=perturbations)
|
||||
|
||||
# Split results back and combine via guiders.
|
||||
splits_v = list(all_v.chunk(n)) if all_v is not None else [0.0] * n
|
||||
|
||||
@@ -4,27 +4,32 @@ from typing import TypeVar
|
||||
|
||||
import torch
|
||||
|
||||
from ltx_core.devices import synchronize_device
|
||||
from ltx_pipelines.utils.allocator_trim_strategy import AllocatorTrimStrategy
|
||||
from ltx_pipelines.utils.helpers import cleanup_memory
|
||||
|
||||
_M = TypeVar("_M", bound=torch.nn.Module)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def gpu_model(model: _M) -> Iterator[_M]:
|
||||
def gpu_model(model: _M, alloc_trim_strategy: AllocatorTrimStrategy = AllocatorTrimStrategy.TRIM) -> Iterator[_M]:
|
||||
"""Context manager that yields a model and releases its memory on exit.
|
||||
Moves all parameters and buffers to ``meta`` device on exit, which
|
||||
immediately releases the underlying storage on **both** GPU and CPU,
|
||||
then runs ``cleanup_memory()`` to reclaim fragmented CUDA memory.
|
||||
On ``TRIM`` (default): synchronize, move parameters/buffers to the ``meta``
|
||||
device (releasing GPU+CPU storage), then ``cleanup_memory()`` to return
|
||||
cached blocks to the OS. ``DEFER`` skips this -- the model's storage is
|
||||
reclaimed by normal GC and the CUDA caching allocator stays warm for the
|
||||
next build (cheaper for back-to-back runs).
|
||||
Usage::
|
||||
with gpu_model(build_encoder()) as encoder:
|
||||
... # use encoder — typed as the concrete class
|
||||
... # use encoder -- typed as the concrete class
|
||||
# GPU + CPU memory freed automatically
|
||||
"""
|
||||
try:
|
||||
yield model
|
||||
finally:
|
||||
torch.cuda.synchronize()
|
||||
# .to("meta") releases storage for all parameters/buffers regardless
|
||||
# of their original device (CUDA or CPU).
|
||||
model.to("meta")
|
||||
cleanup_memory()
|
||||
if alloc_trim_strategy == AllocatorTrimStrategy.TRIM:
|
||||
synchronize_device()
|
||||
# .to("meta") releases storage for all parameters/buffers regardless
|
||||
# of their original device (CUDA or CPU).
|
||||
model.to("meta")
|
||||
cleanup_memory()
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import gc
|
||||
import logging
|
||||
|
||||
import torch
|
||||
@@ -9,6 +8,7 @@ from ltx_core.conditioning import (
|
||||
VideoConditionByKeyframeIndex,
|
||||
VideoConditionByLatentIndex,
|
||||
)
|
||||
from ltx_core.devices import cleanup_accelerator_memory, get_preferred_device
|
||||
from ltx_core.model.audio_vae import encode_audio
|
||||
from ltx_core.model.transformer import Modality
|
||||
from ltx_core.model.video_vae import TilingConfig, VideoEncoder
|
||||
@@ -28,20 +28,11 @@ from ltx_pipelines.utils.media_io import (
|
||||
|
||||
|
||||
def get_device() -> torch.device:
|
||||
if torch.cuda.is_available():
|
||||
return torch.device("cuda", torch.cuda.current_device())
|
||||
return torch.device("cpu")
|
||||
return get_preferred_device()
|
||||
|
||||
|
||||
def cleanup_memory() -> None:
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
try:
|
||||
if hasattr(torch._C, "_host_emptyCache"):
|
||||
torch._C._host_emptyCache()
|
||||
except Exception:
|
||||
logging.warning("Host empty cache cleanup failed; ignoring.", exc_info=True)
|
||||
cleanup_accelerator_memory()
|
||||
|
||||
|
||||
def _conform_latent_length(latent: torch.Tensor, expected_frames_count: int) -> torch.Tensor:
|
||||
|
||||
@@ -258,16 +258,22 @@ def decode_image(image_path: str) -> np.ndarray:
|
||||
return np_array
|
||||
|
||||
|
||||
def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
|
||||
def _validate_audio_waveform(audio: Audio) -> None:
|
||||
"""Raise ValueError if the waveform is empty or not stereo ``(2, N)`` / ``(N, 2)``."""
|
||||
samples = audio.waveform
|
||||
if samples.ndim == 1:
|
||||
samples = samples[:, None]
|
||||
if samples.numel() == 0:
|
||||
raise ValueError("audio.waveform is empty; pass audio=None for no audio.")
|
||||
if samples.ndim != 2 or 2 not in samples.shape:
|
||||
raise ValueError(f"audio.waveform must be stereo (2, N) or (N, 2); got shape {tuple(samples.shape)}.")
|
||||
|
||||
if samples.shape[1] != 2 and samples.shape[0] == 2:
|
||||
samples = samples.T
|
||||
|
||||
if samples.shape[1] != 2:
|
||||
raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
|
||||
def _normalize_audio_waveform(samples: torch.Tensor) -> torch.Tensor:
|
||||
"""Transpose a validated stereo waveform to channel-last ``(N, 2)``."""
|
||||
return samples.T if samples.shape[1] != 2 else samples
|
||||
|
||||
|
||||
def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
|
||||
samples = _normalize_audio_waveform(audio.waveform)
|
||||
|
||||
# Convert to int16 packed for ingestion; resampler converts to encoder fmt.
|
||||
if samples.dtype != torch.int16:
|
||||
@@ -335,13 +341,35 @@ def encode_video(
|
||||
preset: str = "veryfast",
|
||||
thread_count: int = 0,
|
||||
) -> None:
|
||||
"""Encode RGB frames to an H.264 file, optionally muxing an audio track.
|
||||
Args:
|
||||
video: RGB frames as a ``(F, H, W, C)`` float ``[0, 1]`` tensor, or an iterator of
|
||||
such per-chunk tensors (e.g. the VAE decoder output). An empty iterator raises.
|
||||
fps: Output frame rate.
|
||||
audio: Audio track to mux, or None for a video-only file. Waveform must be stereo
|
||||
``(2, N)`` or ``(N, 2)``.
|
||||
output_path: Destination path. Partial output is removed if encoding fails.
|
||||
video_chunks_number: Number of chunks yielded by ``video``, for the progress bar.
|
||||
frame_converter: Float-to-pixel converter (default YUV420p BT.709).
|
||||
crf: libx264 constant rate factor; lower is higher quality (0-51).
|
||||
preset: libx264 speed/compression preset.
|
||||
thread_count: libx264 thread count (0 = auto).
|
||||
Raises:
|
||||
ValueError: On an empty ``video`` or a non-stereo ``audio.waveform``.
|
||||
"""
|
||||
if audio is not None:
|
||||
_validate_audio_waveform(audio)
|
||||
|
||||
if isinstance(video, torch.Tensor):
|
||||
video = iter([video])
|
||||
|
||||
def convert(chunk: torch.Tensor) -> torch.Tensor:
|
||||
return frame_converter(chunk.movedim(-1, -3))
|
||||
|
||||
first_chunk = convert(next(video))
|
||||
first_raw_chunk = next(video, None)
|
||||
if first_raw_chunk is None:
|
||||
raise ValueError("video is empty; expected at least one frame chunk.")
|
||||
first_chunk = convert(first_raw_chunk)
|
||||
|
||||
if frame_converter.pixel_format == PixelFormat.RGB24:
|
||||
height, width = first_chunk.shape[-3], first_chunk.shape[-2]
|
||||
@@ -397,6 +425,7 @@ def encode_audio(audio: Audio, output_path: str) -> None:
|
||||
the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
|
||||
the AAC stream used for muxed video.
|
||||
"""
|
||||
_validate_audio_waveform(audio)
|
||||
container = av.open(output_path, mode="w")
|
||||
audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
|
||||
audio_stream.codec_context.sample_rate = audio.sampling_rate
|
||||
|
||||
@@ -8,6 +8,7 @@ from tqdm import tqdm
|
||||
|
||||
from ltx_core.components.diffusion_steps import EulerCfgPpDiffusionStep, Res2sDiffusionStep
|
||||
from ltx_core.components.protocols import DiffusionStepProtocol
|
||||
from ltx_core.devices import highest_precision_float
|
||||
from ltx_core.model.transformer import X0Model
|
||||
from ltx_core.utils import to_denoised, to_velocity
|
||||
from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
|
||||
@@ -157,7 +158,10 @@ def _channelwise_normalize(x: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
|
||||
def _get_new_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
|
||||
noise = torch.randn(x.shape, generator=generator, dtype=torch.float64, device=generator.device)
|
||||
# float64 on CUDA/CPU for numerical stability; MPS has no float64, so degrade to float32.
|
||||
noise = torch.randn(
|
||||
x.shape, generator=generator, dtype=highest_precision_float(generator.device), device=generator.device
|
||||
)
|
||||
noise = (noise - noise.mean()) / noise.std()
|
||||
return _channelwise_normalize(noise)
|
||||
|
||||
@@ -175,10 +179,11 @@ def _inject_sde_noise(
|
||||
eta: float = 0.5,
|
||||
) -> torch.Tensor:
|
||||
sigmas_copy = sigmas.clone()
|
||||
hp = highest_precision_float(state.denoise_mask.device)
|
||||
new_noise = new_noise_fn(state.latent, step_noise_generator)
|
||||
if not legacy_mode:
|
||||
timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx].double())
|
||||
next_timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx + 1].double())
|
||||
timesteps = timesteps_from_mask(state.denoise_mask.to(hp), sigmas_copy[step_idx].to(hp))
|
||||
next_timesteps = timesteps_from_mask(state.denoise_mask.to(hp), sigmas_copy[step_idx + 1].to(hp))
|
||||
sigmas = torch.stack([timesteps, next_timesteps])
|
||||
step_idx = 0
|
||||
x_next = stepper.step(
|
||||
@@ -249,6 +254,8 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
if present_state is None:
|
||||
raise ValueError("At least one of video_state or audio_state must be provided")
|
||||
state_device = present_state.latent.device
|
||||
# float64 on CUDA/CPU for ODE numerical stability; MPS has no float64, so degrade to float32.
|
||||
hp = highest_precision_float(state_device)
|
||||
|
||||
# Initialize noise generators with different seeds
|
||||
if noise_seed_substep is None:
|
||||
@@ -270,19 +277,19 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
if sigmas[-1] == 0:
|
||||
sigmas = torch.cat([sigmas[:-1], torch.tensor([0.0011, 0.0], device=sigmas.device)], dim=0)
|
||||
# Compute step sizes in hyperbolic space
|
||||
hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
|
||||
hs = -torch.log(sigmas[1:].to(hp).cpu() / (sigmas[:-1].to(hp).cpu()))
|
||||
|
||||
# Initialize phi cache for reuse across loop iterations
|
||||
phi_cache = {}
|
||||
c2 = 0.5 # Midpoint for res_2s
|
||||
|
||||
for step_idx in tqdm(range(n_full_steps)):
|
||||
sigma = sigmas[step_idx].double()
|
||||
sigma_next = sigmas[step_idx + 1].double()
|
||||
sigma = sigmas[step_idx].to(hp)
|
||||
sigma_next = sigmas[step_idx + 1].to(hp)
|
||||
|
||||
# Initialize anchor point
|
||||
x_anchor_video = video_state.latent.clone().double() if video_state is not None else None
|
||||
x_anchor_audio = audio_state.latent.clone().double() if audio_state is not None else None
|
||||
x_anchor_video = video_state.latent.clone().to(hp) if video_state is not None else None
|
||||
x_anchor_audio = audio_state.latent.clone().to(hp) if audio_state is not None else None
|
||||
|
||||
# ====================================================================
|
||||
# STAGE 1: Evaluate at current point
|
||||
@@ -307,15 +314,15 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
# Compute substep x using RK coefficient a21
|
||||
# ====================================================================
|
||||
if x_anchor_video is not None and denoised_video_1 is not None:
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
|
||||
eps_1_video = denoised_video_1.to(hp) - x_anchor_video
|
||||
x_mid_video = x_anchor_video.to(hp) + h * a21 * eps_1_video
|
||||
else:
|
||||
eps_1_video = None
|
||||
x_mid_video = None
|
||||
|
||||
if x_anchor_audio is not None and denoised_audio_1 is not None:
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
|
||||
eps_1_audio = denoised_audio_1.to(hp) - x_anchor_audio
|
||||
x_mid_audio = x_anchor_audio.to(hp) + h * a21 * eps_1_audio
|
||||
else:
|
||||
eps_1_audio = None
|
||||
x_mid_audio = None
|
||||
@@ -347,10 +354,10 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
for _ in range(bongmath_max_iter):
|
||||
if x_mid_video is not None and eps_1_video is not None:
|
||||
x_anchor_video = x_mid_video - h * a21 * eps_1_video
|
||||
eps_1_video = denoised_video_1.double() - x_anchor_video
|
||||
eps_1_video = denoised_video_1.to(hp) - x_anchor_video
|
||||
if x_mid_audio is not None and eps_1_audio is not None:
|
||||
x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
|
||||
eps_1_audio = denoised_audio_1.double() - x_anchor_audio
|
||||
eps_1_audio = denoised_audio_1.to(hp) - x_anchor_audio
|
||||
|
||||
# ====================================================================
|
||||
# STAGE 2: Evaluate at substep point (WITH NOISE)
|
||||
@@ -384,13 +391,13 @@ def res2s_audio_video_denoising_loop( # noqa: PLR0913,PLR0915,PLR0912
|
||||
# FINAL COMBINATION: Compute x_next using RK coefficients
|
||||
# ====================================================================
|
||||
if x_anchor_video is not None and eps_1_video is not None and denoised_video_2 is not None:
|
||||
eps_2_video = denoised_video_2.double() - x_anchor_video
|
||||
eps_2_video = denoised_video_2.to(hp) - x_anchor_video
|
||||
x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
|
||||
else:
|
||||
x_next_video = None
|
||||
|
||||
if x_anchor_audio is not None and eps_1_audio is not None and denoised_audio_2 is not None:
|
||||
eps_2_audio = denoised_audio_2.double() - x_anchor_audio
|
||||
eps_2_audio = denoised_audio_2.to(hp) - x_anchor_audio
|
||||
x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
|
||||
else:
|
||||
x_next_audio = None
|
||||
|
||||
Reference in New Issue
Block a user