Hongsheng Liu
81e5e8e3bb
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
(cherry picked from commit dd6f32cb47)
307 lines
12 KiB
Python
307 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import MethodType
|
|
from typing import Any
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKLLTX2Video
|
|
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoCausalConv3d
|
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
|
from torch import nn
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
|
DistributedOperator,
|
|
DistributedVaeExecutor,
|
|
DistributedVaeMixin,
|
|
GridSpec,
|
|
TileTask,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
# Adapted from https://github.com/sgl-project/sglang/pull/27431.
|
|
def _weight_is_channels_last_3d(module: LTX2VideoCausalConv3d) -> bool:
|
|
weight = module.conv.weight
|
|
return (
|
|
weight.dim() == 5
|
|
and hasattr(torch, "channels_last_3d")
|
|
and weight.is_contiguous(memory_format=torch.channels_last_3d)
|
|
)
|
|
|
|
|
|
def _temporal_pad_channels_last_3d(
|
|
hidden_states: torch.Tensor,
|
|
*,
|
|
left: int,
|
|
right: int,
|
|
) -> torch.Tensor:
|
|
"""Replicate temporal edge frames directly into CL3D storage."""
|
|
batch, channels, frames, height, width = hidden_states.shape
|
|
padded = torch.empty(
|
|
(batch, channels, frames + left + right, height, width),
|
|
dtype=hidden_states.dtype,
|
|
device=hidden_states.device,
|
|
memory_format=torch.channels_last_3d,
|
|
)
|
|
padded[:, :, left : left + frames].copy_(hidden_states)
|
|
if left:
|
|
padded[:, :, :left].copy_(hidden_states[:, :, :1])
|
|
if right:
|
|
padded[:, :, left + frames :].copy_(hidden_states[:, :, -1:])
|
|
return padded
|
|
|
|
|
|
def _layout_preserving_causal_conv_forward(
|
|
module: LTX2VideoCausalConv3d,
|
|
hidden_states: torch.Tensor,
|
|
causal: bool = True,
|
|
) -> torch.Tensor:
|
|
"""Diffusers LTX causal-conv math with a CL3D-preserving pad path."""
|
|
time_kernel_size = module.kernel_size[0]
|
|
use_channels_last_pad = hidden_states.dim() == 5 and _weight_is_channels_last_3d(module)
|
|
|
|
if use_channels_last_pad:
|
|
if causal:
|
|
left, right = time_kernel_size - 1, 0
|
|
else:
|
|
left = right = (time_kernel_size - 1) // 2
|
|
if left == 0 and right == 0:
|
|
hidden_states = hidden_states.contiguous(memory_format=torch.channels_last_3d)
|
|
else:
|
|
hidden_states = _temporal_pad_channels_last_3d(
|
|
hidden_states,
|
|
left=left,
|
|
right=right,
|
|
)
|
|
elif causal:
|
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, time_kernel_size - 1, 1, 1))
|
|
hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
|
|
else:
|
|
pad_size = (time_kernel_size - 1) // 2
|
|
pad_left = hidden_states[:, :, :1, :, :].repeat((1, 1, pad_size, 1, 1))
|
|
pad_right = hidden_states[:, :, -1:, :, :].repeat((1, 1, pad_size, 1, 1))
|
|
hidden_states = torch.concatenate([pad_left, hidden_states, pad_right], dim=2)
|
|
|
|
return module.conv(hidden_states)
|
|
|
|
|
|
def _enable_channels_last_3d(model: nn.Module) -> tuple[int, int]:
|
|
"""Convert LTX Conv3d weights and install layout-preserving causal pads."""
|
|
if not hasattr(torch, "channels_last_3d"):
|
|
return 0, 0
|
|
|
|
converted_convs = 0
|
|
patched_causal_convs = 0
|
|
for module in model.modules():
|
|
if isinstance(module, LTX2VideoCausalConv3d):
|
|
module.forward = MethodType(_layout_preserving_causal_conv_forward, module)
|
|
patched_causal_convs += 1
|
|
if isinstance(module, nn.Conv3d):
|
|
module.weight.data = module.weight.data.contiguous(memory_format=torch.channels_last_3d)
|
|
converted_convs += 1
|
|
return converted_convs, patched_causal_convs
|
|
|
|
|
|
class LTX2VaeExecutor(DistributedVaeExecutor):
|
|
def _pack_local_tiles_without_meta(self, local_results, global_padding_shape, device, dtype):
|
|
tile_tensor = torch.zeros(global_padding_shape, device=device, dtype=dtype)
|
|
for idx, (_, t_tensor) in enumerate(local_results):
|
|
slices = tuple(slice(0, s) for s in t_tensor.shape)
|
|
tile_tensor[idx][slices] = t_tensor
|
|
return tile_tensor
|
|
|
|
def _unpack_tiles_with_known_metadata(self, tile_gather, assigned, tile_output_shapes):
|
|
coord_tensor_map = {}
|
|
for r, rank_tasks in enumerate(assigned):
|
|
tiles_src = tile_gather[r]
|
|
for idx, task in enumerate(rank_tasks):
|
|
output_shape = tile_output_shapes[task.tile_id]
|
|
slices = tuple(slice(0, dim) for dim in output_shape)
|
|
coord_tensor_map[task.grid_coord] = tiles_src[idx][slices]
|
|
|
|
return coord_tensor_map
|
|
|
|
def execute(self, z: torch.Tensor, operator: DistributedOperator, broadcast_result: bool = True):
|
|
pp_size = min(self.parallel_size, self.world_size)
|
|
|
|
tiletask_list, grid_spec = operator.split(z)
|
|
max_tile_output_shape = grid_spec.tile_spec.get("max_tile_output_shape")
|
|
tile_output_shapes = grid_spec.tile_spec.get("tile_output_shapes")
|
|
if max_tile_output_shape is None or tile_output_shapes is None:
|
|
return super().execute(z, operator, broadcast_result=broadcast_result)
|
|
|
|
assigned = self._balance_tasks(tiletask_list, pp_size)
|
|
local_tasks = assigned[self.rank] if self.rank < pp_size else []
|
|
local_results = [(t.tile_id, operator.exec(t)) for t in local_tasks]
|
|
|
|
output_dtype = grid_spec.output_dtype if grid_spec.output_dtype is not None else z.dtype
|
|
max_local_tile_count = max(len(tasks) for tasks in assigned)
|
|
global_padding_shape = [max_local_tile_count, *max_tile_output_shape]
|
|
local_tile_tensor = self._pack_local_tiles_without_meta(
|
|
local_results,
|
|
global_padding_shape,
|
|
z.device,
|
|
output_dtype,
|
|
)
|
|
tile_gather = self.gather_tensors(local_tile_tensor)
|
|
|
|
if self.rank != 0:
|
|
result = torch.empty(0, device=z.device)
|
|
else:
|
|
coord_tensor_map = self._unpack_tiles_with_known_metadata(tile_gather, assigned, tile_output_shapes)
|
|
result = operator.merge(coord_tensor_map, grid_spec)
|
|
|
|
if broadcast_result:
|
|
result = self._sync_final_result(result, z.ndim, z.device, output_dtype)
|
|
return result
|
|
|
|
|
|
class DistributedAutoencoderKLLTX2Video(AutoencoderKLLTX2Video, DistributedVaeMixin):
|
|
@classmethod
|
|
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
|
model = super().from_pretrained(*args, **kwargs)
|
|
model.init_distributed()
|
|
return model
|
|
|
|
def enable_channels_last_3d(self) -> None:
|
|
"""Opt into the layout-preserving CL3D path for this video VAE."""
|
|
converted_convs, patched_causal_convs = _enable_channels_last_3d(self)
|
|
logger.info(
|
|
"LTX2 video VAE: converted %d Conv3d weights to channels_last_3d and "
|
|
"installed %d layout-preserving causal conv paths",
|
|
converted_convs,
|
|
patched_causal_convs,
|
|
)
|
|
|
|
def init_distributed(self):
|
|
self.distributed_executor = LTX2VaeExecutor()
|
|
|
|
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
|
_, _, num_frames, height, width = z.shape
|
|
sample_height = height * self.spatial_compression_ratio
|
|
sample_width = width * self.spatial_compression_ratio
|
|
|
|
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
|
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
|
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
|
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
|
sample_num_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
|
|
output_channels = int(getattr(getattr(self, "config", None), "out_channels", 3))
|
|
|
|
tiletask_list = []
|
|
tile_output_shapes = {}
|
|
for i in range(0, height, tile_latent_stride_height):
|
|
for j in range(0, width, tile_latent_stride_width):
|
|
tile = z[:, :, :num_frames, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
|
tile_id = len(tiletask_list)
|
|
tile_output_shapes[tile_id] = (
|
|
z.shape[0],
|
|
output_channels,
|
|
sample_num_frames,
|
|
tile.shape[3] * self.spatial_compression_ratio,
|
|
tile.shape[4] * self.spatial_compression_ratio,
|
|
)
|
|
tiletask_list.append(
|
|
TileTask(
|
|
tile_id,
|
|
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
|
tile,
|
|
workload=tile.shape[2] * tile.shape[3] * tile.shape[4],
|
|
)
|
|
)
|
|
|
|
tile_spec = {
|
|
"sample_height": sample_height,
|
|
"sample_width": sample_width,
|
|
"blend_height": self.tile_sample_min_height - self.tile_sample_stride_height,
|
|
"blend_width": self.tile_sample_min_width - self.tile_sample_stride_width,
|
|
"tile_sample_stride_height": self.tile_sample_stride_height,
|
|
"tile_sample_stride_width": self.tile_sample_stride_width,
|
|
"max_tile_output_shape": (
|
|
z.shape[0],
|
|
output_channels,
|
|
sample_num_frames,
|
|
self.tile_sample_min_height,
|
|
self.tile_sample_min_width,
|
|
),
|
|
"tile_output_shapes": tile_output_shapes,
|
|
}
|
|
grid_spec = GridSpec(
|
|
split_dims=(3, 4),
|
|
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
|
tile_spec=tile_spec,
|
|
output_dtype=self.dtype,
|
|
)
|
|
return tiletask_list, grid_spec
|
|
|
|
def tile_exec(
|
|
self,
|
|
task: TileTask,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
) -> torch.Tensor:
|
|
if hasattr(self, "clear_cache"):
|
|
self.clear_cache()
|
|
return self.decoder(task.tensor, temb, causal=causal)
|
|
|
|
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
|
grid_h, grid_w = grid_spec.grid_shape
|
|
result_rows = []
|
|
|
|
if hasattr(self, "clear_cache"):
|
|
self.clear_cache()
|
|
|
|
for i in range(grid_h):
|
|
result_row = []
|
|
for j in range(grid_w):
|
|
tile = coord_tensor_map[(i, j)]
|
|
if i > 0:
|
|
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
|
if j > 0:
|
|
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
|
result_row.append(
|
|
tile[
|
|
:,
|
|
:,
|
|
:,
|
|
: grid_spec.tile_spec["tile_sample_stride_height"],
|
|
: grid_spec.tile_spec["tile_sample_stride_width"],
|
|
]
|
|
)
|
|
result_rows.append(torch.cat(result_row, dim=-1))
|
|
|
|
dec = torch.cat(result_rows, dim=3)[
|
|
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
|
]
|
|
return dec
|
|
|
|
def tiled_decode(
|
|
self,
|
|
z: torch.Tensor,
|
|
temb: torch.Tensor | None = None,
|
|
causal: bool | None = None,
|
|
return_dict: bool = True,
|
|
):
|
|
if not self.is_distributed_enabled():
|
|
if causal is None:
|
|
return super().tiled_decode(z, temb, return_dict=return_dict)
|
|
return super().tiled_decode(z, temb, causal=causal, return_dict=return_dict)
|
|
|
|
logger.debug("LTX2 video VAE decode running with distributed tiled executor")
|
|
result = self.distributed_executor.execute(
|
|
z,
|
|
DistributedOperator(
|
|
split=self.tile_split,
|
|
exec=lambda task: self.tile_exec(task, temb=temb, causal=causal),
|
|
merge=self.tile_merge,
|
|
),
|
|
broadcast_result=False,
|
|
)
|
|
if not return_dict:
|
|
return (result,)
|
|
|
|
return DecoderOutput(sample=result)
|