# 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 inspect

import torch
import torch.nn.functional as F
from tqdm import tqdm

from ...configuration_utils import FrozenDict
from ...guiders import ClassifierFreeGuidance
from ...models import AutoencoderKLWan, WanAnimate2Transformer3DModel
from ...models.transformers.transformer_wan_animate_2 import WanAnimate2KVCache
from ...schedulers.scheduling_utils import SchedulerMixin
from ...utils import logging
from ...utils.torch_utils import randn_tensor
from ..modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .encoders import encode_vae, get_i2v_mask


logger = logging.get_logger(__name__)  # pylint: disable=invalid-name


def decode_vae(vae: AutoencoderKLWan, latents: torch.Tensor) -> torch.Tensor:
    """De-standardize latents and VAE-decode them to `[B, 3, T, H, W]` pixels in `[-1, 1]`."""
    latents = latents.to(vae.dtype)
    latents_mean = (
        torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype)
    )
    latents_recip_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(
        latents.device, latents.dtype
    )
    latents = latents / latents_recip_std + latents_mean
    return vae.decode(latents, return_dict=False)[0]


# ========================================
# Segment Loop Leaf Blocks
# ========================================


class WanAnimate2SegmentVaeEncoderStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that VAE-encodes this segment's slice of the driving video and stacks "
            "the i2v conditioning mask on top. The Wan VAE is causal in time, so encoding the whole video once "
            "and slicing the latents would not be equivalent — each segment restarts the temporal convolution on "
            "its own slice. A streaming mode would replace this block with one fed segments incrementally. This "
            "block should be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("vae", AutoencoderKLWan),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            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",
            ),
            InputParam(
                "reference_image_latents",
                required=True,
                type_hint=torch.Tensor,
                description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step",
            ),
            InputParam(
                "effective_segment",
                required=True,
                type_hint=int,
                description="Frames each segment advances: `segment_frame_length - prev_segment_conditioning_frames`, from the video preprocess step",
            ),
            InputParam(
                "segment_frame_length",
                type_hint=int,
                default=81,
                description="The number of frames in each inference segment",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "driving_video_latents",
                type_hint=torch.Tensor,
                description="VAE latents of this segment's driving-video slice",
            ),
            OutputParam(
                "driving_video_condition",
                type_hint=torch.Tensor,
                description="i2v mask + driving-slice latents, conditioning the reference-extraction pass",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        device = components._execution_device

        latent_height, latent_width = block_state.reference_image_latents.shape[-2:]

        start = k * block_state.effective_segment
        block_state.driving_video_latents = encode_vae(
            components.vae, block_state.driving_video_pixels[:, :, start : start + block_state.segment_frame_length]
        )

        condition_mask = get_i2v_mask(
            block_state.driving_video_latents.shape[2],
            latent_height,
            latent_width,
            block_state.segment_frame_length,
            device=device,
        ).to(block_state.driving_video_latents.dtype)
        block_state.driving_video_condition = torch.cat([condition_mask, block_state.driving_video_latents[0]], dim=0)

        return components, block_state


class WanAnimate2SegmentPrevFramesStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that builds the generation-side conditioning tensor `reference_latents`: the previous "
            "segment's tail frames (zeros for the first segment) are VAE-encoded, masked, and stacked under the "
            "reference half `reference_image_latents`. This is how motion continuity crosses segment boundaries — in pixel space, "
            "not latent space. This block should be used to compose the `sub_blocks` attribute of "
            "`WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("vae", AutoencoderKLWan),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "reference_image_latents",
                required=True,
                type_hint=torch.Tensor,
                description="i2v mask + reference image latents `[20, 1, latent_height, latent_width]`, from the image VAE encoder step",
            ),
            InputParam(
                "segment_frame_length",
                type_hint=int,
                default=81,
                description="The number of frames in each inference segment",
            ),
            InputParam(
                "prev_segment_conditioning_frames",
                type_hint=int,
                default=1,
                description="The number of conditioning frames carried over from the previous segment",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "reference_latents",
                type_hint=torch.Tensor,
                description="The full conditioning tensor: reference half stacked over the segment half",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        # `block_state.out_frames` is seeded by the loop wrapper and written by the decode step of the
        # previous iteration.
        device = components._execution_device

        latent_height, latent_width = block_state.reference_image_latents.shape[-2:]
        height = latent_height * components.vae_scale_factor_spatial
        width = latent_width * components.vae_scale_factor_spatial

        num_frames = block_state.segment_frame_length + 1
        mask_len = block_state.prev_segment_conditioning_frames if k > 0 else 0
        if mask_len > 0:
            prev_frames = block_state.out_frames[0, :, -mask_len:].clone().detach()
            prev_frames = F.interpolate(prev_frames.permute(1, 0, 2, 3), size=(height, width), mode="bicubic").permute(
                1, 0, 2, 3
            )
            cond_pixels = torch.cat(
                [
                    prev_frames,
                    torch.zeros(3, num_frames - mask_len - 1, height, width, device=device),
                ],
                dim=1,
            )
        else:
            cond_pixels = torch.zeros(3, num_frames - 1, height, width, device=device)

        prev_segment_cond_latents = encode_vae(components.vae, cond_pixels.unsqueeze(0)).squeeze(0)
        prev_segment_cond_mask = get_i2v_mask(
            prev_segment_cond_latents.shape[1], latent_height, latent_width, mask_len, device=device
        ).to(prev_segment_cond_latents.dtype)
        prev_segment_cond_latents = torch.cat([prev_segment_cond_mask, prev_segment_cond_latents], dim=0)

        block_state.reference_latents = torch.cat(
            [block_state.reference_image_latents, prev_segment_cond_latents], dim=1
        )

        return components, block_state


class WanAnimate2SegmentPrepareStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that draws this segment's initial noise and allocates a fresh KV cache "
            "for the reference-extraction pass. This block should be used to compose the `sub_blocks` attribute "
            "of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", WanAnimate2Transformer3DModel),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("generator"),
            InputParam(
                "reference_latents",
                required=True,
                type_hint=torch.Tensor,
                description="The full conditioning tensor: reference half stacked over the segment half",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("latents", type_hint=torch.Tensor, description="This segment's initial noise"),
            OutputParam(
                "kv_cache",
                type_hint=WanAnimate2KVCache,
                description="Fresh per-segment cache for the reference K/V",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        device = components._execution_device

        block_state.latents = randn_tensor(
            (
                components.num_channels_latents,
                block_state.reference_latents.shape[1],
                block_state.reference_latents.shape[-2],
                block_state.reference_latents.shape[-1],
            ),
            generator=block_state.generator,
            device=device,
            dtype=torch.float32,
        )
        block_state.kv_cache = WanAnimate2KVCache(components.transformer.config.num_layers)

        return components, block_state


class WanAnimate2SegmentSchedulerResetStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that resets the scheduler: each segment is an independent denoising "
            "trajectory, so the solver state and timesteps are re-prepared per segment. This block should be used "
            "to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("scheduler", SchedulerMixin),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("num_inference_steps", default=40),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam("timesteps", type_hint=torch.Tensor, description="This segment's denoising timesteps"),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        device = components._execution_device

        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
        block_state.timesteps = components.scheduler.timesteps

        return components, block_state


class WanAnimate2RefExtractStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that runs the transformer's reference-extraction pass "
            '(`kv_cache_mode="extract"`): the driving-video segment is encoded once and every layer\'s reference '
            "K/V is stored in the KV cache, which the denoising forwards then attend over. This block should be "
            "used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", WanAnimate2Transformer3DModel),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "driving_video_latents",
                required=True,
                type_hint=torch.Tensor,
                description="VAE latents of this segment's driving-video slice",
            ),
            InputParam(
                "driving_video_condition",
                required=True,
                type_hint=torch.Tensor,
                description="i2v mask + driving-slice latents, conditioning the reference-extraction pass",
            ),
            InputParam(
                "condition_clip_context",
                required=True,
                type_hint=torch.Tensor,
                description="CLIP vision features of the driving video's first frame",
            ),
            InputParam(
                "prompt_ref_embeds",
                required=True,
                type_hint=torch.Tensor,
                description="Text embeddings of the reference prompt, guiding the reference-extraction pass",
            ),
            InputParam(
                "kv_cache",
                required=True,
                type_hint=WanAnimate2KVCache,
                description="Per-segment cache holding every layer's reference K/V",
            ),
            InputParam(
                "timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="This segment's denoising timesteps",
            ),
            InputParam(
                "max_seq_len_ref",
                required=True,
                type_hint=int,
                description="Packed sequence length of the reference tokens",
            ),
            InputParam(
                "grid_sizes_ref",
                required=True,
                type_hint=torch.Tensor,
                description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        device = components._execution_device
        transformer_dtype = components.transformer.dtype

        t_ref = torch.tensor([block_state.timesteps[0].item()], device=device, dtype=transformer_dtype)
        components.transformer(
            [block_state.driving_video_latents[0].to(transformer_dtype)],
            timestep=t_ref,
            encoder_hidden_states=[block_state.prompt_ref_embeds[0].to(transformer_dtype)],
            encoder_hidden_states_image=block_state.condition_clip_context.to(transformer_dtype),
            condition_latents=[block_state.driving_video_condition.to(transformer_dtype)],
            kv_cache=block_state.kv_cache,
            kv_cache_mode="extract",
            seq_len=block_state.max_seq_len_ref,
            offset_grid_sizes=block_state.grid_sizes_ref,
        )

        return components, block_state


# ========================================
# Inner Denoising Blocks
# ========================================


class WanAnimate2SegmentDenoiseInner(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Inner timestep loop that denoises one segment with guidance, attending over the segment's cached "
            "reference K/V. The unconditional branch passes `is_uncondtion=True` to the transformer (it skips a "
            "dedicated layer on that branch), routed through the guider as a per-branch input. This block should "
            "be used to compose the `sub_blocks` attribute of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", WanAnimate2Transformer3DModel),
            ComponentSpec("scheduler", SchedulerMixin),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 3.0}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "latents",
                required=True,
                type_hint=torch.Tensor,
                description="This segment's latents",
            ),
            InputParam(
                "reference_latents",
                required=True,
                type_hint=torch.Tensor,
                description="The full conditioning tensor: reference half stacked over the segment half",
            ),
            InputParam(
                "kv_cache",
                required=True,
                type_hint=WanAnimate2KVCache,
                description="Per-segment cache holding every layer's reference K/V",
            ),
            InputParam(
                "timesteps",
                required=True,
                type_hint=torch.Tensor,
                description="This segment's denoising timesteps",
            ),
            InputParam.template("num_inference_steps", default=40),
            InputParam(
                "num_segments",
                required=True,
                type_hint=int,
                description="Total number of segments in the driving video, from the video preprocess step",
            ),
            InputParam(
                "max_seq_len",
                required=True,
                type_hint=int,
                description="Packed sequence length of the generation tokens",
            ),
            InputParam(
                "grid_sizes_ref",
                required=True,
                type_hint=torch.Tensor,
                description="Post-patch latent grid `[[T, H/2, W/2]]` of a driving-video segment",
            ),
            InputParam(
                "segment_frame_length",
                type_hint=int,
                default=81,
                description="The number of frames in each inference segment",
            ),
            InputParam(
                "height",
                required=True,
                type_hint=int,
                description="The resolved frame height in pixels",
            ),
            InputParam(
                "width",
                required=True,
                type_hint=int,
                description="The resolved frame width in pixels",
            ),
            InputParam.template("generator"),
            InputParam.template("prompt_embeds"),
            InputParam.template("negative_prompt_embeds"),
            InputParam.template("denoiser_input_fields"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam.template("latents"),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        transformer_dtype = components.transformer.dtype

        guider_inputs = {
            "encoder_hidden_states": (block_state.prompt_embeds, block_state.negative_prompt_embeds),
            "is_uncondtion": (False, True),
        }

        # Everything the transformer accepts from the tagged conditioning fields, minus what the guider manages
        # per branch (currently the reference image's CLIP features, `encoder_hidden_states_image`).
        transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
        shared_kwargs = {
            name: value.to(transformer_dtype) if isinstance(value, torch.Tensor) else value
            for name, value in block_state.denoiser_input_fields.items()
            if name in transformer_args and name not in guider_inputs
        }

        with tqdm(
            total=len(block_state.timesteps), desc=f"Segment {k + 1}/{block_state.num_segments}"
        ) as progress_bar:
            for i, t in enumerate(block_state.timesteps):
                timestep = torch.stack([t])

                components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t)
                guider_state = components.guider.prepare_inputs(guider_inputs)

                for guider_state_batch in guider_state:
                    components.guider.prepare_models(components.transformer)

                    guider_state_batch.noise_pred = components.transformer(
                        [block_state.latents.to(transformer_dtype)],
                        timestep=timestep,
                        encoder_hidden_states=[guider_state_batch.encoder_hidden_states[0].to(transformer_dtype)],
                        condition_latents=[block_state.reference_latents.to(transformer_dtype)],
                        kv_cache=block_state.kv_cache,
                        kv_cache_mode="cached",
                        seq_len=block_state.max_seq_len,
                        reference_grid_sizes=block_state.grid_sizes_ref,
                        origin_len=block_state.segment_frame_length,
                        origin_area=[block_state.height, block_state.width],
                        is_uncondtion=guider_state_batch.is_uncondtion,
                        **shared_kwargs,
                    ).sample[0]

                    components.guider.cleanup_models(components.transformer)

                noise_pred = components.guider(guider_state)[0]

                latents = components.scheduler.step(
                    noise_pred.unsqueeze(0),
                    t,
                    block_state.latents.unsqueeze(0),
                    return_dict=False,
                    generator=block_state.generator,
                )[0]
                block_state.latents = latents.squeeze(0)

                progress_bar.update()

        return components, block_state


class WanAnimate2DistilledSegmentDenoiseInner(WanAnimate2SegmentDenoiseInner):
    model_name = "wan-animate-2-distilled"

    @property
    def description(self) -> str:
        return (
            "Inner timestep loop that denoises one segment for the distilled model, which is trained for few-step "
            "sampling without classifier-free guidance — the guider defaults to `guidance_scale=1.0`, so only the "
            "conditional branch runs. This block should be used to compose the `sub_blocks` attribute of "
            "`WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("transformer", WanAnimate2Transformer3DModel),
            ComponentSpec("scheduler", SchedulerMixin),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 1.0}),
                default_creation_method="from_config",
            ),
        ]


# ========================================
# Post-Denoise
# ========================================


class WanAnimate2SegmentDecodeStep(ModularPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Step within the segment loop that VAE-decodes the denoised segment. Decoding happens inside the loop "
            "because the next segment conditions on this segment's decoded pixels. Finished frames move to CPU and "
            "the per-segment KV cache and latents are freed — holding them across segments fragments the "
            "allocator enough to OOM at high resolution. This block should be used to compose the `sub_blocks` "
            "attribute of `WanAnimate2SegmentLoopWrapper`."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("vae", AutoencoderKLWan),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "latents",
                required=True,
                type_hint=torch.Tensor,
                description="This segment's latents",
            ),
            InputParam(
                "kv_cache",
                required=True,
                type_hint=WanAnimate2KVCache,
                description="Per-segment cache holding every layer's reference K/V",
            ),
            InputParam(
                "prev_segment_conditioning_frames",
                type_hint=int,
                default=1,
                description="The number of conditioning frames carried over from the previous segment",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "out_frames",
                type_hint=torch.Tensor,
                description="This segment's decoded frames on device; the next segment conditions on its tail",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, block_state: BlockState, k: int):
        latents = block_state.latents.to(torch.float32)
        # The first latent frame is the reference image's slot, not video content.
        out_frames = decode_vae(components.vae, latents[:, 1:])

        if k > 0:
            out_frames = out_frames[:, :, block_state.prev_segment_conditioning_frames :]

        block_state.segment_frames.append(out_frames.cpu())
        block_state.out_frames = out_frames

        block_state.kv_cache.clear()
        block_state.kv_cache = None
        block_state.latents = None
        torch.cuda.empty_cache()

        return components, block_state


# ========================================
# Segment Loop Wrapper
# ========================================


class WanAnimate2SegmentLoopWrapper(LoopSequentialPipelineBlocks):
    model_name = "wan-animate-2"

    @property
    def description(self) -> str:
        return (
            "Pipeline block that iterates over the driving video's segments. At each segment it runs sub-blocks "
            "for per-segment encoding, preparation, reference extraction, denoising, and decoding; each segment "
            "conditions on the previous one's decoded tail frames."
        )

    @property
    def loop_inputs(self) -> list[InputParam]:
        return [
            InputParam(
                "num_segments",
                required=True,
                type_hint=int,
                description="Total number of segments in the driving video, from the video preprocess step",
            ),
        ]

    @property
    def loop_intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                "segment_frames",
                type_hint=list[torch.Tensor],
                description="Per-segment decoded frames on CPU, each `[1, 3, T, H, W]`",
            ),
        ]

    @torch.no_grad()
    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        # Seed the loop-carried state: `segment_frames` collects each segment's decoded frames (the decode step
        # appends to it); `out_frames` is the previous segment's decoded frames — written by the decode step, read
        # by the prev-frames step of the next iteration. `None` marks "no previous segment" for the first iteration.
        block_state.segment_frames = []
        block_state.out_frames = None

        for k in range(block_state.num_segments):
            components, block_state = self.loop_step(components, block_state, k=k)

        self.set_block_state(state, block_state)
        return components, state


# ========================================
# Composed Segment Denoise Steps
# ========================================


class WanAnimate2DenoiseStep(WanAnimate2SegmentLoopWrapper):
    block_classes = [
        WanAnimate2SegmentVaeEncoderStep,
        WanAnimate2SegmentPrevFramesStep,
        WanAnimate2SegmentPrepareStep,
        WanAnimate2SegmentSchedulerResetStep,
        WanAnimate2RefExtractStep,
        WanAnimate2SegmentDenoiseInner,
        WanAnimate2SegmentDecodeStep,
    ]
    block_names = [
        "vae_encoder",
        "prev_frames",
        "prepare",
        "scheduler_reset",
        "ref_extract",
        "denoise_inner",
        "decode",
    ]

    @property
    def description(self) -> str:
        return (
            "Segment denoise step that iterates over the driving video's segments.\n"
            "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> "
            "denoise_inner -> decode."
        )


class WanAnimate2DistilledDenoiseStep(WanAnimate2SegmentLoopWrapper):
    model_name = "wan-animate-2-distilled"

    block_classes = [
        WanAnimate2SegmentVaeEncoderStep,
        WanAnimate2SegmentPrevFramesStep,
        WanAnimate2SegmentPrepareStep,
        WanAnimate2SegmentSchedulerResetStep,
        WanAnimate2RefExtractStep,
        WanAnimate2DistilledSegmentDenoiseInner,
        WanAnimate2SegmentDecodeStep,
    ]
    block_names = [
        "vae_encoder",
        "prev_frames",
        "prepare",
        "scheduler_reset",
        "ref_extract",
        "denoise_inner",
        "decode",
    ]

    @property
    def description(self) -> str:
        return (
            "Segment denoise step for the distilled model that iterates over the driving video's segments.\n"
            "At each segment: vae_encoder -> prev_frames -> prepare -> scheduler_reset -> ref_extract -> "
            "denoise_inner (no classifier-free guidance) -> decode."
        )