211 lines
8.5 KiB
Python
211 lines
8.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKLLTX2Video
|
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
|
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__)
|
|
|
|
|
|
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 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)
|