Spaces:
Running on Zero
Running on Zero
File size: 4,612 Bytes
3db1cbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | # Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import numpy as np
import torch
from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import InputParam, OutputParam
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class WanAnimate2PrepareSegmentsStep(ModularPipelineBlocks):
model_name = "wan-animate-2"
@property
def description(self) -> str:
return (
"Step that computes the segment-invariant geometry for the segment loop. The zigzag padding makes "
"every segment exactly `segment_frame_length` frames, so the latent grid, the packed sequence lengths, and the "
"noise shape are the same for all segments and are computed once here."
)
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"segment_frame_length",
type_hint=int,
default=81,
description="The number of frames in each inference segment",
),
InputParam(
"reference_image_latents",
required=True,
type_hint=torch.Tensor,
description="The reference conditioning tensor `[20, 1, latent_height, latent_width]`; "
"provides the latent grid",
),
InputParam(
"driving_video_pixels",
required=True,
type_hint=torch.Tensor,
description="The preprocessed driving video `[1, 3, T, H, W]`, from the video preprocess step",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"grid_sizes_ref",
type_hint=torch.Tensor,
description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment, used as the "
"offset grid of the reference-extraction pass and the reference grid of the denoising passes",
),
OutputParam(
"max_seq_len",
type_hint=int,
description="Packed sequence length of the generation tokens",
),
OutputParam(
"max_seq_len_ref",
type_hint=int,
description="Packed sequence length of the reference tokens",
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
latent_height, latent_width = block_state.reference_image_latents.shape[-2:]
expected = (
latent_height * components.vae_scale_factor_spatial,
latent_width * components.vae_scale_factor_spatial,
)
if tuple(block_state.driving_video_pixels.shape[-2:]) != expected:
raise ValueError(
f"`driving_video_pixels` is letterboxed to {tuple(block_state.driving_video_pixels.shape[-2:])} but "
f"the reference image conditioning is {expected} — the video and image preprocess steps must use the "
"same `height`/`width`."
)
latent_segment_frames = (block_state.segment_frame_length - 1) // components.vae_scale_factor_temporal + 1
ref_shape = [latent_segment_frames, latent_height, latent_width]
ref_shape_post = [ref_shape[0], ref_shape[1] // 2, ref_shape[2] // 2]
block_state.grid_sizes_ref = torch.tensor([ref_shape_post], dtype=torch.long)
# The noise tensor carries one extra latent frame: the reference image's slot.
latent_noise_frames = latent_segment_frames + 1
block_state.max_seq_len = int(math.ceil(np.prod([latent_noise_frames, latent_height // 2, latent_width // 2])))
block_state.max_seq_len_ref = int(math.ceil(np.prod(ref_shape) // 4))
self.set_block_state(state, block_state)
return components, state
|