Canlin GuoandHongsheng Liu
e771842781
Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
from diffusers.models.autoencoders import AutoencoderKLQwenImage
|
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
|
DistributedOperator,
|
|
DistributedVaeMixin,
|
|
GridSpec,
|
|
TileTask,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
|
|
@classmethod
|
|
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
|
model = super().from_pretrained(*args, **kwargs)
|
|
model.init_distributed()
|
|
return model
|
|
|
|
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
|
# mostly copy from AutoencoderKL
|
|
_, _, 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
|
|
|
|
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
|
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
|
|
|
# Split z into overlapping tiles and decode them separately.
|
|
# The tiles have an overlap to avoid seams between tiles.
|
|
tiletask_list = []
|
|
for i in range(0, height, tile_latent_stride_height):
|
|
for j in range(0, width, tile_latent_stride_width):
|
|
time_list = []
|
|
for k in range(num_frames):
|
|
self._conv_idx = [0]
|
|
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
|
time_list.append(tile)
|
|
tiletask_list.append(
|
|
TileTask(
|
|
len(tiletask_list),
|
|
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
|
time_list,
|
|
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
|
)
|
|
)
|
|
tile_spec = {
|
|
"sample_height": sample_height,
|
|
"sample_width": sample_width,
|
|
"blend_height": blend_height,
|
|
"blend_width": blend_width,
|
|
}
|
|
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) -> torch.Tensor:
|
|
"""Decode a single latent tile into RGB space."""
|
|
self.clear_cache()
|
|
time = []
|
|
for k in range(len(task.tensor)):
|
|
self._conv_idx = [0]
|
|
tile = self.post_quant_conv(task.tensor[k])
|
|
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
|
|
time.append(decoded)
|
|
result = torch.cat(time, dim=2)
|
|
return result
|
|
|
|
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
|
"""Merge decoded tiles into a full image."""
|
|
grid_h, grid_w = grid_spec.grid_shape
|
|
self.clear_cache()
|
|
|
|
result_rows = []
|
|
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[:, :, :, : self.tile_sample_stride_height, : self.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, return_dict: bool = True):
|
|
if not self.is_distributed_enabled():
|
|
return super().tiled_decode(z, return_dict=return_dict)
|
|
|
|
logger.debug("Decode running with distributed executor")
|
|
result = self.distributed_executor.execute(
|
|
z,
|
|
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
|
broadcast_result=True,
|
|
)
|
|
if not return_dict:
|
|
return (result,)
|
|
|
|
return DecoderOutput(sample=result)
|