MaciejBalaNV
c482626450
Signed-off-by: Maciej Bala <mbala@nvidia.com> Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
63 lines
2.6 KiB
Python
63 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Chunked Wan decode with a preallocated output buffer."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify
|
|
|
|
|
|
def decode_frames(vae, z: torch.Tensor) -> torch.Tensor:
|
|
"""``AutoencoderKLWan._decode``'s frame loop, writing every chunk straight into the result.
|
|
|
|
Upstream grows the result with ``torch.cat([out, out_], 2)`` on every chunk
|
|
(quadratic copying) and then materializes ``unpatchify`` and ``clamp`` as two
|
|
more full-size copies, so three or four copies of the decoded video are live
|
|
at once. Here the final ``[B, C, T, H, W]`` buffer is allocated after the
|
|
first chunk and each chunk is unpatchified and clamped directly into its
|
|
slot. Every step is a permutation or an elementwise clamp, so the values are
|
|
identical to upstream. Tiling dispatch is the caller's responsibility.
|
|
"""
|
|
vae.clear_cache()
|
|
x = vae.post_quant_conv(z)
|
|
patch_size = vae.config.patch_size
|
|
num_frames = x.shape[2]
|
|
# Each temporal upsampler doubles the frames of every chunk after the first
|
|
# (the first chunk drops the extra leading frames), so steady-state chunks
|
|
# produce ``2 ** num_temporal_upsamplers`` frames each.
|
|
frames_per_chunk = 2 ** sum(bool(flag) for flag in vae.decoder.temperal_upsample)
|
|
|
|
output: torch.Tensor | None = None
|
|
offset = 0
|
|
for index in range(num_frames):
|
|
vae._conv_idx = [0]
|
|
chunk = vae.decoder(
|
|
x[:, :, index : index + 1],
|
|
feat_cache=vae._feat_map,
|
|
feat_idx=vae._conv_idx,
|
|
first_chunk=index == 0,
|
|
)
|
|
if patch_size is not None:
|
|
chunk = unpatchify(chunk, patch_size=patch_size)
|
|
if output is None:
|
|
total_frames = chunk.shape[2] + (num_frames - 1) * frames_per_chunk
|
|
output = chunk.new_empty((chunk.shape[0], chunk.shape[1], total_frames, chunk.shape[3], chunk.shape[4]))
|
|
end = offset + chunk.shape[2]
|
|
if end > output.shape[2]:
|
|
raise RuntimeError(
|
|
f"Wan decoder produced more frames than expected: chunk {index} ends at {end}, "
|
|
f"allocated {output.shape[2]}"
|
|
)
|
|
torch.clamp(chunk, min=-1.0, max=1.0, out=output[:, :, offset:end])
|
|
offset = end
|
|
|
|
if output is None or offset != output.shape[2]:
|
|
expected = None if output is None else output.shape[2]
|
|
raise RuntimeError(f"Wan decoder produced {offset} frames, expected {expected}")
|
|
vae.clear_cache()
|
|
return output
|
|
|
|
|
|
__all__ = ["decode_frames"]
|