Function decode_latents

animation/utils/utils.py:14–45

(
    vae,
    latents, 
    num_frames, 
    decode_chunk_size=8)

logger = logging.getLogger(__name__)

def decode_latents(
    vae,
    latents, 
    num_frames, 
    decode_chunk_size=8):
    # [batch, frames, channels, height, width] -> [batch*frames, channels, height, width]
    latents = latents.flatten(0, 1)

    latents = 1 / vae.config.scaling_factor * latents

    forward_vae_fn = vae._orig_mod.forward if is_compiled_module(vae) else vae.forward
    accepts_num_frames = "num_frames" in set(inspect.signature(forward_vae_fn).parameters.keys())

    # decode decode_chunk_size frames at a time to avoid OOM
    frames = []
    for i in range(0, latents.shape[0], decode_chunk_size):
        num_frames_in = latents[i: i + decode_chunk_size].shape[0]
        decode_kwargs = {}
        if accepts_num_frames:
            # we only pass num_frames_in if it's expected
            decode_kwargs["num_frames"] = num_frames_in

        frame = vae.decode(latents[i: i + decode_chunk_size], **decode_kwargs).sample
        frames.append(frame.cpu())
    frames = torch.cat(frames, dim=0)

    # [batch*frames, channels, height, width] -> [batch, channels, frames, height, width]
    frames = frames.reshape(-1, num_frames, *frames.shape[1:]).permute(0, 2, 1, 3, 4)

    # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
    frames = frames.float()
    return frames

def tensor2vid(video, processor, output_type="np"):
    batch_size, channels, num_frames, height, width = video.shape