add ttm pipelines

pipelines/cog_pipeline.py

@@ -0,0 +1,524 @@
+# Copyright 2025 Noam Rotstein
+#
+# 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.
+#
+# Adapted from Hugging Face Diffusers (Apache-2.0):
+#   https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py
+
+try:
+    from dataclasses import dataclass
+    import math
+    from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+    import torch
+    from transformers import T5EncoderModel, T5Tokenizer
+    from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+    from diffusers.image_processor import PipelineImageInput
+    from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
+    from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
+    from diffusers.utils import (
+        is_torch_xla_available,
+        logging,
+        replace_example_docstring,
+    )
+    from diffusers.utils.torch_utils import randn_tensor
+    from diffusers.video_processor import VideoProcessor
+    from diffusers.pipelines.cogvideo.pipeline_output import CogVideoXPipelineOutput
+    from diffusers.pipelines.cogvideo.pipeline_cogvideox_image2video import retrieve_timesteps
+    from diffusers import CogVideoXImageToVideoPipeline
+
+    import torch.nn.functional as F
+    from pipelines.utils import load_video_to_tensor
+
+except ImportError as e:
+    raise ImportError(f"Required module not found: {e}. Please install it before running this script. "
+                     f"For installation instructions, see: https://github.com/zai-org/CogVideo")
+
+
+try:
+    if is_torch_xla_available():
+        import torch_xla.core.xla_model as xm
+        XLA_AVAILABLE = True
+    else:
+        XLA_AVAILABLE = False
+except ImportError:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+EXAMPLE_DOC_STRING = """
+"""
+
+
+class CogVideoXImageToVideoTTMPipeline(CogVideoXImageToVideoPipeline):
+    r"""
+    Pipeline for image-to-video generation using CogVideoX combined with Time to Move (TTM).
+    This model inherits from [`CogVideoXImageToVideoPipeline`].
+    """
+    _optional_components = []
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    _callback_tensor_inputs = [
+        "latents",
+        "prompt_embeds",
+        "negative_prompt_embeds",
+    ]
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKLCogVideoX,
+        transformer: CogVideoXTransformer3DModel,
+        scheduler: Union[CogVideoXDDIMScheduler, CogVideoXDPMScheduler],
+    ):
+        super().__init__(
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            vae=vae,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+
+        self.register_modules(
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            vae=vae,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor_spatial = (
+            2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        )
+        self.vae_scale_factor_temporal = (
+            self.vae.config.temporal_compression_ratio if getattr(self, "vae", None) else 4
+        )
+        self.vae_scaling_factor_image = self.vae.config.scaling_factor if getattr(self, "vae", None) else 0.7
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+
+    @torch.no_grad()
+    def encode_frames(self, frames: torch.Tensor, vae_scale_factor: float = None) -> torch.Tensor:
+        """Encode video frames into latent space with shape  (B, F, C, H, W). Input shape (B, C, F, H, W), expected range [-1, 1]."""
+        latents = self.vae.encode(frames)[0].sample()
+        # latents = self.vae.encode(frames)[0].mode()
+        vae_scale_factor = vae_scale_factor or self.vae_scaling_factor_image
+        latents = latents * vae_scale_factor
+        return latents.permute(0, 2, 1, 3, 4).contiguous() # shape (B, C, F, H, W) -> (B, F, C, H, W)
+
+
+    def convert_rgb_mask_to_latent_mask(self, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Convert a per-frame mask [T, 1, H, W] to latent resolution [1, T_latent, 1, H', W'].
+        T_latent groups frames by the temporal VAE downsample factor k = vae_scale_factor_temporal:
+        [0], [1..k], [k+1..2k], ...
+        """
+        k = self.vae_scale_factor_temporal
+
+        mask0 = mask[0:1]  # [1,1,H,W]
+        mask1 = mask[1::k]  # [T'-1,1,H,W]
+        sampled = torch.cat([mask0, mask1], dim=0)  # [T',1,H,W]
+        pooled = sampled.permute(1, 0, 2, 3).unsqueeze(0)
+
+        # Up-sample spatially to match latent spatial resolution
+        s = self.vae_scale_factor_spatial
+        H_latent = pooled.shape[-2] // s
+        W_latent = pooled.shape[-1] // s
+        pooled = F.interpolate(pooled, size=(pooled.shape[2], H_latent, W_latent), mode="nearest")
+
+        # Back to [1, T_latent, 1, H, W]
+        latent_mask = pooled.permute(0, 2, 1, 3, 4)
+
+        return latent_mask
+        
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: PipelineImageInput,
+        prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_frames: int = 49,
+        num_inference_steps: int = 50,
+        timesteps: Optional[List[int]] = None,
+        guidance_scale: float = 6,
+        use_dynamic_cfg: bool = False,
+        num_videos_per_prompt: int = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 226,
+        motion_signal_video_path: Optional[str] = None,
+        motion_signal_mask_path: Optional[str] = None,
+        tweak_index: int = 0,
+        tstrong_index: int = 0
+    ) -> Union[CogVideoXPipelineOutput, Tuple]:
+        """
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`PipelineImageInput`):
+                The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            height (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
+                The height in pixels of the generated image. This is set to 480 by default for the best results.
+            width (`int`, *optional*, defaults to self.transformer.config.sample_height * self.vae_scale_factor_spatial):
+                The width in pixels of the generated image. This is set to 720 by default for the best results.
+            num_frames (`int`, defaults to `48`):
+                Number of frames to generate. Must be divisible by self.vae_scale_factor_temporal. Generated video will
+                contain 1 extra frame because CogVideoX is conditioned with (num_seconds * fps + 1) frames where
+                num_seconds is 6 and fps is 8. However, since videos can be saved at any fps, the only condition that
+                needs to be satisfied is that of divisibility mentioned above.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion_xl.StableDiffusionXLPipelineOutput`] instead
+                of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, defaults to `226`):
+                Maximum sequence length in encoded prompt. Must be consistent with
+                `self.transformer.config.max_text_seq_length` otherwise may lead to poor results.
+            motion_signal_video_path (`str`):
+                Path to the video file containing the motion signal to guide the motion of the generated video.
+                It should be a crude version of the reference video, with pixels with motion dragged to their target.
+            motion_signal_mask_path (`str`):
+                Path to the mask video file containing the motion mask of TTM.
+                The mask should be a binary with the conditioning motion pixels being 1 and the rest being 0.
+            tweak_index (`int`):
+                The index of the tweak, from which the denoising process starts.
+            tstrong_index (`int`):
+                The index of the tweak, from which the denoising process starts in the motion conditioned region.
+        Examples:
+
+        Returns:
+            [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] or `tuple`:
+            [`~pipelines.cogvideo.pipeline_output.CogVideoXPipelineOutput`] if `return_dict` is True, otherwise a
+            `tuple`. When returning a tuple, the first element is a list with the generated images.
+        """
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        height = height or self.transformer.config.sample_height * self.vae_scale_factor_spatial
+        width = width or self.transformer.config.sample_width * self.vae_scale_factor_spatial
+        num_frames = num_frames or self.transformer.config.sample_frames
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            image=image,
+            prompt=prompt,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        if motion_signal_mask_path is None:
+            raise ValueError("`motion_signal_mask_path` is required for TTM.")
+        if motion_signal_video_path is None:
+            raise ValueError("`motion_signal_video_path` is required for TTM.")
+
+        # 2. Default call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=do_classifier_free_guidance,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+        self._num_timesteps = len(timesteps)
+
+        # 5. Prepare latents
+        latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+
+        # For CogVideoX 1.5, the latent frames should be padded to make it divisible by patch_size_t
+        patch_size_t = self.transformer.config.patch_size_t
+        additional_frames = 0
+        if patch_size_t is not None and latent_frames % patch_size_t != 0:
+            additional_frames = patch_size_t - latent_frames % patch_size_t
+            num_frames += additional_frames * self.vae_scale_factor_temporal
+
+        image = self.video_processor.preprocess(image, height=height, width=width).to(
+            device, dtype=prompt_embeds.dtype
+        )
+
+        latent_channels = self.transformer.config.in_channels // 2
+        latents, image_latents = self.prepare_latents(
+            image,
+            batch_size * num_videos_per_prompt,
+            latent_channels,
+            num_frames,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Create rotary embeds if required
+        image_rotary_emb = (
+            self._prepare_rotary_positional_embeddings(height, width, latents.size(1), device)
+            if self.transformer.config.use_rotary_positional_embeddings
+            else None
+        )
+
+        # 8. Create ofs embeds if required
+        ofs_emb = None if self.transformer.config.ofs_embed_dim is None else latents.new_full((1,), fill_value=2.0)
+
+        # 9. Initialize for TTM
+        ref_vid = load_video_to_tensor(motion_signal_video_path).to(device=device) # shape [1, C, T, H, W]
+        refB, refC, refT, refH, refW = ref_vid.shape
+        ref_vid = F.interpolate(
+            ref_vid.permute(0, 2, 1, 3, 4).reshape(refB*refT, refC, refH, refW),
+            size=(height, width), mode="bicubic", align_corners=True, 
+        ).reshape(refB, refT, refC, height, width).permute(0, 2, 1, 3, 4)
+
+        ref_vid = self.video_processor.normalize(ref_vid.to(dtype=self.vae.dtype)) # Normalize and convert dtype for VAE encoding 
+        ref_latents = self.encode_frames(ref_vid).float().detach() # shape [1, T, C, H, W]
+
+        ref_mask = load_video_to_tensor(motion_signal_mask_path).to(device=device) # shape [1, C, T, H, W]
+        mB, mC, mT, mH, mW = ref_mask.shape
+
+        ref_mask = F.interpolate(
+                ref_mask.permute(0, 2, 1, 3, 4).reshape(mB*mT, mC, mH, mW),
+                size=(height, width), mode="nearest", 
+            ).reshape(mB, mT, mC, height, width).permute(0, 2, 1, 3, 4)
+        ref_mask = ref_mask[0].permute(1, 0, 2, 3).contiguous() # (1, C, T, H, W) -> (T, H, W, 1)
+
+        if len(ref_mask.shape) == 4:
+            ref_mask = ref_mask.unsqueeze(0)
+
+        ref_mask = ref_mask[0,:,:1].contiguous() # (1, T, C, H, W) -> (T, 1, H, W)
+        ref_mask = (ref_mask > 0.5).float().max(dim=1, keepdim=True)[0] # [T, 1, H, W]
+        motion_mask = self.convert_rgb_mask_to_latent_mask(ref_mask)  # [1, T, 1, H, W]
+        background_mask = 1.0 - motion_mask
+
+        if tweak_index >= 0:
+            tweak = self.scheduler.timesteps[tweak_index]
+            fixed_noise = randn_tensor(
+                ref_latents.shape,
+                generator=generator,
+                device=ref_latents.device,
+                dtype=ref_latents.dtype,
+            )
+            noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, tweak.long())
+            latents = noisy_latents.to(dtype=latents.dtype, device=latents.device)
+        else:
+            tweak = torch.tensor(-1)
+            fixed_noise = randn_tensor(
+                ref_latents.shape,
+                generator=generator,
+                device=ref_latents.device,
+                dtype=ref_latents.dtype,
+            )
+            tweak_index = 0
+
+        # 10. Denoising loop
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+        # logging
+        # ------------------------------------------------------------------
+        with self.progress_bar(total=len(timesteps) - tweak_index) as progress_bar:
+            # for DPM-solver++
+            old_pred_original_sample = None
+            for i, t in enumerate(timesteps[tweak_index:]):
+                if self.interrupt:
+                    continue
+                
+                self._current_timestep = t
+                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                latent_image_input = torch.cat([image_latents] * 2) if do_classifier_free_guidance else image_latents
+                latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=2)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latent_model_input.shape[0])
+
+                # predict noise model_output
+                noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep=timestep,
+                    ofs=ofs_emb,
+                    image_rotary_emb=image_rotary_emb,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+                noise_pred = noise_pred.float()
+
+                # perform guidance
+                if use_dynamic_cfg:
+                    self._guidance_scale = 1 + guidance_scale * (
+                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                    )
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                if not isinstance(self.scheduler, CogVideoXDPMScheduler):
+                    latents, old_pred_original_sample = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)
+                else:
+                    latents, old_pred_original_sample = self.scheduler.step(
+                        noise_pred,
+                        old_pred_original_sample,
+                        t,
+                        timesteps[i - 1] if i > 0 else None,
+                        latents,
+                        **extra_step_kwargs,
+                        return_dict=False,
+                    )
+
+                # In between tweak and tstrong, replace mask with noisy reference latents
+                in_between_tweak_tstrong = (i+tweak_index) < tstrong_index
+                
+                if in_between_tweak_tstrong:
+                    if i+tweak_index+1 < len(timesteps):
+                        prev_t = timesteps[i+tweak_index+1]
+                        noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, prev_t.long()).to(dtype=latents.dtype, device=latents.device)
+                        latents = latents * background_mask + noisy_latents * motion_mask
+                    else:
+                        latents = latents * background_mask + ref_latents * motion_mask
+                    
+                latents = latents.to(prompt_embeds.dtype)
+
+                # call the callback, if provided
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if not output_type == "latent":
+            # Discard any padding frames that were added for CogVideoX 1.5
+            latents = latents[:, additional_frames:]
+            frames = self.decode_latents(latents)
+            video = self.video_processor.postprocess_video(video=frames, output_type=output_type)
+        else:
+            video = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video,)
+
+        return CogVideoXPipelineOutput(
+            frames=video,
+            )

pipelines/svd_pipeline.py

@@ -0,0 +1,624 @@
+# Copyright 2025 Noam Rotstein
+#
+# 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.
+#
+# Adapted from Hugging Face Diffusers (Apache-2.0):
+#   https://github.com/huggingface/diffusers/blob/8abc7aeb715c0149ee0a9982b2d608ce97f55215/src/diffusers/pipelines/stable_video_diffusion/pipeline_stable_video_diffusion.py#L147
+
+try:
+    import inspect
+    from dataclasses import dataclass
+    from typing import Callable, Dict, List, Optional, Union
+    import numpy as np
+    import PIL.Image
+    import torch
+    from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
+    from diffusers.models import AutoencoderKLTemporalDecoder, UNetSpatioTemporalConditionModel
+    from diffusers.schedulers import EulerDiscreteScheduler
+    from diffusers.utils import BaseOutput, is_torch_xla_available, logging, replace_example_docstring
+    from diffusers.utils.torch_utils import randn_tensor
+    from diffusers.video_processor import VideoProcessor
+    import torch.nn.functional as F
+    from diffusers.pipelines.stable_video_diffusion import StableVideoDiffusionPipeline
+    from pipelines.utils import load_video_to_tensor
+
+except ImportError as e:
+    raise ImportError(f"Required module not found: {e}. Please install it before running this script. "
+                     f"For installation instructions, see:https://github.com/Stability-AI/generative-models")
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+EXAMPLE_DOC_STRING = """
+"""
+
+
+def _append_dims(x, target_dims):
+    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
+    dims_to_append = target_dims - x.ndim
+    if dims_to_append < 0:
+        raise ValueError(f"input has {x.ndim} dims but target_dims is {target_dims}, which is less")
+    return x[(...,) + (None,) * dims_to_append]
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+@dataclass
+class StableVideoDiffusionPipelineOutput(BaseOutput):
+    r"""
+    Output class for Stable Video Diffusion pipeline.
+
+    Args:
+        frames (`[List[List[PIL.Image.Image]]`, `np.ndarray`, `torch.Tensor`]):
+            List of denoised PIL images of length `batch_size` or numpy array or torch tensor of shape `(batch_size,
+            num_frames, height, width, num_channels)`.
+    """
+
+    frames: Union[List[List[PIL.Image.Image]], np.ndarray, torch.Tensor]
+
+
+class StableVideoDiffusionTTMPipeline(StableVideoDiffusionPipeline):
+    r"""
+    Pipeline to generate video from an input image using Stable Video Diffusion combined with Time to Move (TTM).
+    This model inherits from [`StableVideoDiffusionPipeline`].
+    """
+
+    model_cpu_offload_seq = "image_encoder->unet->vae"
+    _callback_tensor_inputs = ["latents"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKLTemporalDecoder,
+        image_encoder: CLIPVisionModelWithProjection,
+        unet: UNetSpatioTemporalConditionModel,
+        scheduler: EulerDiscreteScheduler,
+        feature_extractor: CLIPImageProcessor,
+    ):
+        super().__init__(
+            vae=vae,
+            image_encoder=image_encoder,
+            unet=unet,
+            scheduler=scheduler,
+            feature_extractor=feature_extractor,
+        )
+
+        self.register_modules(
+            vae=vae,
+            image_encoder=image_encoder,
+            unet=unet,
+            scheduler=scheduler,
+            feature_extractor=feature_extractor,
+        )
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+        self.video_processor = VideoProcessor(do_resize=True, vae_scale_factor=self.vae_scale_factor)
+
+
+    def encode_frames(self, frames: torch.Tensor, num_frames: int, encode_chunk_size: int = 14):
+        """
+        Args:
+            frames: [B, C, T, H, W] tensor, preprocessed to VAE's expected range (e.g., [-1, 1]).
+            num_frames: T (used for reshaping back).
+            encode_chunk_size: process at most this many frames at a time to avoid OOM.
+
+        Returns:
+            latents: [B, T, C_latent, h, w], multiplied by self.vae.config.scaling_factor.
+
+        Notes:
+            - Stochastic: samples from posterior (latent_dist.sample()).
+            - If the VAE's compiled module hides the signature, we inspect the original .forward
+            and pass num_frames only if it's accepted (same pattern as decode).
+        """
+        if frames.dim() != 5:
+            raise ValueError(f"Expected frames with shape [B, C, T, H, W], got {list(frames.shape)}")
+        B, C, T, H, W = frames.shape
+
+        # [B, C, T, H, W] -> [B, T, C, H, W] -> [B*T, C, H, W]
+        frames_bt = frames.permute(0, 2, 1, 3, 4).reshape(-1, C, H, W)
+
+        # Use the *encode* signature (decoder may accept num_frames, encoder usually doesn't)
+        encode_fn = self.vae._orig_mod.encode if hasattr(self.vae, "_orig_mod") else self.vae.encode
+        try:
+            accepts_num_frames = ("num_frames" in inspect.signature(encode_fn).parameters)
+        except (TypeError, ValueError):
+            # Signature might be obscured by wrappers/compilation; be conservative
+            accepts_num_frames = False
+
+        latents_chunks = []
+        for i in range(0, frames_bt.shape[0], encode_chunk_size):
+            chunk = frames_bt[i : i + encode_chunk_size]
+
+            # match VAE device/dtype to avoid implicit casts
+            chunk = chunk.to(device=self.vae.device, dtype=self.vae.dtype)
+
+            encode_kwargs = {}
+            if accepts_num_frames:
+                # This will normally be False for AutoencoderKLTemporalDecoder.encode()
+                encode_kwargs["num_frames"] = chunk.shape[0]
+
+            # Be robust to unexpected wrappers hiding the signature
+            try:
+                enc_out = self.vae.encode(chunk, **encode_kwargs)
+            except TypeError as e:
+                if "unexpected keyword argument 'num_frames'" in str(e):
+                    enc_out = self.vae.encode(chunk)
+                else:
+                    raise
+
+            posterior = enc_out.latent_dist  # DiagonalGaussianDistribution
+            latents_chunks.append(posterior.sample())
+
+        latents = torch.cat(latents_chunks, dim=0)  # [B*T, C_lat, h, w]
+        latents = latents * self.vae.config.scaling_factor
+
+        # [B*T, C_lat, h, w] -> [B, T, C_lat, h, w]
+        latents = latents.reshape(B, num_frames, *latents.shape[1:])
+
+        return latents
+
+    def convert_rgb_mask_to_latent_mask(self, mask: torch.Tensor, first_different=True) -> torch.Tensor:
+        """
+        Args:
+            mask: [T, 1, H, W] tensor (0/1 or any float in [0,1]).
+        Returns:
+            latent_mask: [1, T_latent, 1, H, W], where
+                T_latent = ceil(T / self.vae_scale_factor_temporal)
+                For CogVideoX-style VAE (k=4), groups are [0], [1-4], [5-8], ..., achieved by
+                pre-padding zeros at the start before max-pooling with stride=k.
+        """
+        T, _, H, W = mask.shape
+
+        k = self.vae_scale_factor_temporal
+        # Pre-pad zeros along time so that the first pooled window corresponds to frame 0 alone
+        if first_different:
+            num_pad = (k - (T % k)) % k
+            pad = torch.zeros((num_pad, 1, H, W), device=mask.device, dtype=mask.dtype)
+            mask = torch.cat([pad, mask], dim=0)
+        
+
+        # [T,1,H,W] -> [1,1,T,H,W]
+        x = mask.permute(1, 0, 2, 3).unsqueeze(0)
+        if k > 1:
+            # Max-pool over time with kernel=stride=k (no spatial pooling)
+            pooled = F.max_pool3d(x, kernel_size=(k, 1, 1), stride=(k, 1, 1))
+        else:
+            pooled = x
+
+        # Up-sample spatially to match latent spatial resolution
+        s = self.vae_scale_factor_spatial
+        H_latent = pooled.shape[-2] // s
+        W_latent = pooled.shape[-1] // s
+        pooled = F.interpolate(pooled, size=(pooled.shape[2], H_latent, W_latent), mode="nearest")
+
+        # Back to [1, T_latent, 1, H, W]
+        latent_mask = pooled.permute(0, 2, 1, 3, 4)
+
+        return latent_mask
+    
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: Union[PIL.Image.Image, List[PIL.Image.Image], torch.Tensor],
+        height: int = 576,
+        width: int = 1024,
+        num_frames: Optional[int] = None,
+        num_inference_steps: int = 25,
+        sigmas: Optional[List[float]] = None,
+        min_guidance_scale: float = 1.0,
+        max_guidance_scale: float = 3.0,
+        fps: int = 7,
+        motion_bucket_id: int = 127,
+        noise_aug_strength: float = 0.02,
+        decode_chunk_size: Optional[int] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        return_dict: bool = True,
+        motion_signal_video_path: Optional[str] = None,
+        motion_signal_mask_path: Optional[str] = None,
+        tweak_index: int = 0,
+        tstrong_index: int = 0
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            image (`PIL.Image.Image` or `List[PIL.Image.Image]` or `torch.Tensor`):
+                Image(s) to guide image generation. If you provide a tensor, the expected value range is between `[0,
+                1]`.
+            height (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to `self.unet.config.sample_size * self.vae_scale_factor`):
+                The width in pixels of the generated image.
+            num_frames (`int`, *optional*):
+                The number of video frames to generate. Defaults to `self.unet.config.num_frames` (14 for
+                `stable-video-diffusion-img2vid` and to 25 for `stable-video-diffusion-img2vid-xt`).
+            num_inference_steps (`int`, *optional*, defaults to 25):
+                The number of denoising steps. More denoising steps usually lead to a higher quality video at the
+                expense of slower inference. This parameter is modulated by `strength`.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            min_guidance_scale (`float`, *optional*, defaults to 1.0):
+                The minimum guidance scale. Used for the classifier free guidance with first frame.
+            max_guidance_scale (`float`, *optional*, defaults to 3.0):
+                The maximum guidance scale. Used for the classifier free guidance with last frame.
+            fps (`int`, *optional*, defaults to 7):
+                Frames per second. The rate at which the generated images shall be exported to a video after
+                generation. Note that Stable Diffusion Video's UNet was micro-conditioned on fps-1 during training.
+            motion_bucket_id (`int`, *optional*, defaults to 127):
+                Used for conditioning the amount of motion for the generation. The higher the number the more motion
+                will be in the video.
+            noise_aug_strength (`float`, *optional*, defaults to 0.02):
+                The amount of noise added to the init image, the higher it is the less the video will look like the
+                init image. Increase it for more motion.
+            decode_chunk_size (`int`, *optional*):
+                The number of frames to decode at a time. Higher chunk size leads to better temporal consistency at the
+                expense of more memory usage. By default, the decoder decodes all frames at once for maximal quality.
+                For lower memory usage, reduce `decode_chunk_size`.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of videos to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            motion_signal_video_path (`str`):
+                Path to the video file containing the motion signal to guide the motion of the generated video.
+                It should be a crude version of the reference video, with pixels with motion dragged to their target.
+            motion_signal_mask_path (`str`):
+                Path to the mask video file containing the motion mask of TTM.
+                The mask should be a binary with the conditioning motion pixels being 1 and the rest being 0.
+            tweak_index (`int`):
+                The index of the tweak, from which the denoising process starts.
+            tstrong_index (`int`):
+                The index of the tweak, from which the denoising process starts in the motion conditioned region.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for video
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generated image. Choose between `pil`, `np` or `pt`.
+            callback_on_step_end (`Callable`, *optional*):
+                A function that is called at the end of each denoising step during inference. The function is called
+                with the following arguments:
+                    `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`.
+                `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableVideoDiffusionPipelineOutput`] is
+                returned, otherwise a `tuple` of (`List[List[PIL.Image.Image]]` or `np.ndarray` or `torch.Tensor`) is
+                returned.
+        """
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        num_frames = num_frames if num_frames is not None else self.unet.config.num_frames
+        decode_chunk_size = decode_chunk_size if decode_chunk_size is not None else num_frames
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(image, height, width)
+
+        if motion_signal_mask_path is None:
+            raise ValueError("`motion_signal_mask_path` is required for TTM.")
+        if motion_signal_video_path is None:
+            raise ValueError("`motion_signal_video_path` is required for TTM.")
+        
+        # 2. Define call parameters
+        if isinstance(image, PIL.Image.Image):
+            batch_size = 1
+        elif isinstance(image, list):
+            batch_size = len(image)
+        else:
+            batch_size = image.shape[0]
+        device = self._execution_device
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        self._guidance_scale = max_guidance_scale
+
+        # 3. Encode input image
+        image_embeddings = self._encode_image(image, device, num_videos_per_prompt, self.do_classifier_free_guidance)
+
+        # NOTE: Stable Video Diffusion was conditioned on fps - 1, which is why it is reduced here.
+        # See: https://github.com/Stability-AI/generative-models/blob/ed0997173f98eaf8f4edf7ba5fe8f15c6b877fd3/scripts/sampling/simple_video_sample.py#L188
+        fps = fps - 1
+
+        # 4. Encode input image using VAE
+        image = self.video_processor.preprocess(image, height=height, width=width).to(device)
+        noise = randn_tensor(image.shape, generator=generator, device=device, dtype=image.dtype)
+        image = image + noise_aug_strength * noise
+
+        needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float32)
+
+        image_latents = self._encode_vae_image(
+            image,
+            device=device,
+            num_videos_per_prompt=num_videos_per_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+        )
+        image_latents = image_latents.to(image_embeddings.dtype)
+
+        # cast back to fp16 if needed
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+
+        # Repeat the image latents for each frame so we can concatenate them with the noise
+        # image_latents [batch, channels, height, width] ->[batch, num_frames, channels, height, width]
+        image_latents = image_latents.unsqueeze(1).repeat(1, num_frames, 1, 1, 1)
+
+        # 5. Get Added Time IDs
+        added_time_ids = self._get_add_time_ids(
+            fps,
+            motion_bucket_id,
+            noise_aug_strength,
+            image_embeddings.dtype,
+            batch_size,
+            num_videos_per_prompt,
+            self.do_classifier_free_guidance,
+        )
+        added_time_ids = added_time_ids.to(device)
+
+        # 6. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, None, sigmas)
+
+        # ---- Sanity checks for TTM indices (0 ≤ tstrong < tweak < num_steps) ----
+        if not (0 <= tstrong_index < num_inference_steps):
+            raise ValueError(f"tstrong_index must be in [0, {num_inference_steps-1}], got {tstrong_index}.")
+        if not (0 <= tweak_index < num_inference_steps):
+            raise ValueError(f"tweak_index must be in [0, {num_inference_steps-1}], got {tweak_index}.")
+        if not (tstrong_index > tweak_index):
+            raise ValueError(f"Require tweak_index < tstrong_index, got {tweak_index} >= {tstrong_index}.")
+
+        # 7. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_videos_per_prompt,
+            num_frames,
+            num_channels_latents,
+            height,
+            width,
+            image_embeddings.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 8. Prepare guidance scale
+        guidance_scale = torch.linspace(min_guidance_scale, max_guidance_scale, num_frames).unsqueeze(0)
+        guidance_scale = guidance_scale.to(device, latents.dtype)
+        guidance_scale = guidance_scale.repeat(batch_size * num_videos_per_prompt, 1)
+        guidance_scale = _append_dims(guidance_scale, latents.ndim)
+
+        self._guidance_scale = guidance_scale
+
+        # 9. Initialize for TTM
+        ref_vid = load_video_to_tensor(motion_signal_video_path).to(device=device) # shape [1, C, T, H, W]
+        refB, refC, refT, refH, refW = ref_vid.shape
+
+        ref_vid = F.interpolate(
+            ref_vid.permute(0, 2, 1, 3, 4).reshape(refB*refT, refC, refH, refW),
+            size=(height, width), mode="bicubic", align_corners=True, 
+        ).reshape(refB, refT, refC, height, width).permute(0, 2, 1, 3, 4)
+
+        ref_vid = self.video_processor.normalize(ref_vid.to(dtype=self.vae.dtype)) # Normalize and convert dtype for VAE encoding 
+        
+        if num_frames < refT:
+            logger.warning(f"num_frames ({num_frames}) < input frames ({refT}); trimming reference video.")
+            ref_vid = ref_vid[:, :, :num_frames]
+        elif num_frames > refT:
+            raise ValueError(f"num_frames ({num_frames}) is greater than input frames ({refT}). This is not supported.")
+        
+        ref_latents = self.encode_frames(ref_vid, num_frames, decode_chunk_size).detach()
+        ref_latents = ref_latents.to(dtype=latents.dtype, device=device)
+        
+        if not hasattr(self, "vae_scale_factor_temporal"): # encode ref video to latents
+            if hasattr(self.vae, "scale_factor_temporal"):
+                self.vae_scale_factor_temporal = self.vae.scale_factor_temporal
+            else:
+                if ref_latents.shape[1] == num_frames:
+                    self.vae_scale_factor_temporal = 1
+                else:
+                    raise ValueError("Please configure the temporal scale factor of the VAE.")
+        
+        self.vae_scale_factor_spatial = self.vae_scale_factor
+    
+        ref_mask = load_video_to_tensor(motion_signal_mask_path).to(device=device) # shape [1, C, T, H, W]
+
+        mB, mC, mT, mH, mW = ref_mask.shape # do resizing with nearest neighbor to avoid interpolation artifacts
+        ref_mask = F.interpolate(
+            ref_mask.permute(0, 2, 1, 3, 4).reshape(mB*mT, mC, mH, mW),
+            size=(height, width), mode="nearest", 
+        ).reshape(mB, mT, mC, height, width).permute(0, 2, 1, 3, 4)
+        ref_mask = ref_mask[0].permute(1, 0, 2, 3).contiguous() # (1, C, T, H, W) -> (T, H, W, 1)
+        if ref_mask.shape[0] > num_frames:
+            print(f"Warning: num_frames ({num_frames}) is less than input mask frames ({mT}). Trimming to {num_frames}.")
+            ref_mask = ref_mask[:num_frames]
+        elif ref_mask.shape[0] < num_frames:
+            raise ValueError(f"num_frames ({num_frames}) is greater than input mask frames ({mT}). This is not supported.")
+        ref_mask = (ref_mask > 0.5).float().max(dim=1, keepdim=True)[0] # [T, 1, H, W]
+        motion_mask = self.convert_rgb_mask_to_latent_mask(ref_mask, False)  # [1, T, 1, H, W]
+        motion_mask = motion_mask.to(dtype=latents.dtype)
+        background_mask = 1.0 - motion_mask
+
+        if tweak_index >= 0:
+            tweak = self.scheduler.timesteps[tweak_index]
+            tweak = torch.tensor([tweak], device=device)
+            fixed_noise = randn_tensor(ref_latents.shape,
+                                       generator=generator,
+                                       device=ref_latents.device,
+                                       dtype=ref_latents.dtype)
+            noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, tweak)
+            latents = noisy_latents.to(dtype=latents.dtype, device=latents.device)
+        else:
+            tweak = torch.tensor(-1)
+            fixed_noise = randn_tensor(ref_latents.shape,
+                                       generator=generator,
+                                       device=ref_latents.device,
+                                       dtype=ref_latents.dtype)
+            tweak_index = 0
+
+
+        # 10. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+        with self.progress_bar(total=len(timesteps) - tweak_index) as progress_bar:
+            for i, t in enumerate(timesteps[tweak_index:]):
+                # expand the latents if we are doing classifier free guidance
+                latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                # Concatenate image_latents over channels dimension
+                latent_model_input = torch.cat([latent_model_input, image_latents], dim=2)
+
+                # predict the noise residual
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=image_embeddings,
+                    added_time_ids=added_time_ids,
+                    return_dict=False,
+                )[0]
+
+                # perform guidance
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents).prev_sample
+
+                # In between tweak and tstrong, replace mask with noisy reference latents
+                in_between_tweak_tstrong = (i+tweak_index) < tstrong_index
+
+                if in_between_tweak_tstrong:
+                    if i+tweak_index+1 < len(timesteps):
+                        prev_t = torch.tensor([timesteps[i+tweak_index+1]], device=device)
+                        noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, prev_t).to(dtype=latents.dtype, device=latents.device)
+                        latents = latents * background_mask + noisy_latents * motion_mask
+                    elif i+tweak_index+1  == len(timesteps):
+                        latents = latents * background_mask + ref_latents * motion_mask
+                    else:
+                        raise ValueError(f"Unexpected timestep index {i+tweak_index+1} >= {len(timesteps)}")
+                    
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if not output_type == "latent":
+            # cast back to fp16 if needed
+            if needs_upcasting:
+                self.vae.to(dtype=torch.float16)
+            frames = self.decode_latents(latents, num_frames, decode_chunk_size)
+            frames = self.video_processor.postprocess_video(video=frames, output_type=output_type)
+        else:
+            frames = latents
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return frames
+
+        return StableVideoDiffusionPipelineOutput(
+            frames=frames)

pipelines/utils.py

@@ -0,0 +1,45 @@
+import logging
+import os
+from pathlib import Path
+from typing import Tuple
+import numpy as np
+import cv2
+import torch
+
+def validate_inputs(image_path: str, mask_path: str, motion_path: str) -> None:
+    for p in [image_path, mask_path, motion_path]:
+        if not Path(p).exists():
+            raise FileNotFoundError(f"Required file not found: {p}")
+
+def compute_hw_from_area(
+    image_height: int,
+    image_width: int,
+    max_area: int,
+    mod_value: int,
+) -> Tuple[int, int]:
+    """Compute (height, width) with same math and rounding as original."""
+    aspect_ratio = image_height / image_width
+    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+    return int(height), int(width)
+
+
+def load_video_to_tensor(video_path):
+    """Returns a video tensor from a video file. shape [1, T, C, H, W], [0, 1] range."""
+        # load video
+    cap = cv2.VideoCapture(video_path)
+    frames = []
+    while 1:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        frames.append(frame)
+    cap.release()
+    # Convert frames to tensor, shape [T, H, W, C], [0, 1] range
+    frames = np.array(frames)
+
+    video_tensor = torch.tensor(frames)
+    video_tensor = video_tensor.permute(0, 3, 1, 2).float() / 255.0
+    video_tensor = video_tensor.unsqueeze(0).permute(0, 2, 1, 3, 4)
+    return video_tensor

pipelines/wan_pipeline.py

@@ -0,0 +1,559 @@
+# Copyright 2025 Noam Rotstein
+#
+# 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.
+#
+# Adapted from Hugging Face Diffusers (Apache-2.0):
+#   https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/wan/pipeline_wan_i2v.py
+
+
+try:
+    import html
+    from typing import Any, Callable, Dict, List, Optional, Union
+    import torch
+    from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
+    from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+    from diffusers.image_processor import PipelineImageInput
+    from diffusers.models import AutoencoderKLWan, WanTransformer3DModel
+    from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+    from diffusers.utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring
+    from diffusers.utils.torch_utils import randn_tensor
+    from diffusers.video_processor import VideoProcessor
+    from diffusers.pipelines.wan.pipeline_output import WanPipelineOutput
+    from diffusers.pipelines.wan.pipeline_wan_i2v import retrieve_latents, WanImageToVideoPipeline
+
+    import torch.nn.functional as F
+    from pipelines.utils import load_video_to_tensor
+
+    if is_torch_xla_available():
+        import torch_xla.core.xla_model as xm
+
+        XLA_AVAILABLE = True
+    else:
+        XLA_AVAILABLE = False
+except ImportError as e:
+    raise ImportError(f"Required module not found: {e}. Please install it before running this script. "
+                     f"For installation instructions, see: https://github.com/Wan-Video/Wan2.2")
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# after logger/is_ftfy_available
+_ftfy = None
+if is_ftfy_available():
+    import ftfy as _ftfy
+
+
+EXAMPLE_DOC_STRING = """
+"""
+
+    
+class WanImageToVideoTTMPipeline(WanImageToVideoPipeline):
+    r"""
+    Pipeline for image-to-video generation using Wan with Time-To-Move (TTM) conditioning.
+    This model inherits from [`WanImageToVideoPipeline`].
+    """
+
+    model_cpu_offload_seq = "text_encoder->image_encoder->transformer->transformer_2->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+    _optional_components = ["transformer", "transformer_2", "image_encoder", "image_processor"]
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        text_encoder: UMT5EncoderModel,
+        vae: AutoencoderKLWan,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        image_processor: CLIPImageProcessor = None,
+        image_encoder: CLIPVisionModel = None,
+        transformer: WanTransformer3DModel = None,
+        transformer_2: WanTransformer3DModel = None,
+        boundary_ratio: Optional[float] = None,
+        expand_timesteps: bool = False,
+    ):
+        super().__init__(
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            vae=vae,
+            scheduler=scheduler,
+            image_processor=image_processor,
+            image_encoder=image_encoder,
+            transformer=transformer,
+            transformer_2=transformer_2,
+            boundary_ratio=boundary_ratio,
+            expand_timesteps=expand_timesteps,
+        )
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            image_encoder=image_encoder,
+            transformer=transformer,
+            scheduler=scheduler,
+            image_processor=image_processor,
+            transformer_2=transformer_2,
+        )
+        self.register_to_config(boundary_ratio=boundary_ratio, expand_timesteps=expand_timesteps)
+
+        self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if getattr(self, "vae", None) else 4
+        self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if getattr(self, "vae", None) else 8
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+        self.image_processor = image_processor
+
+    
+    def convert_rgb_mask_to_latent_mask(self, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Convert a per-frame mask [T, 1, H, W] to latent resolution [1, T_latent, 1, H', W'].
+        T_latent groups frames by the temporal VAE downsample factor k = vae_scale_factor_temporal:
+        [0], [1..k], [k+1..2k], ...
+        """
+
+        k = self.vae_scale_factor_temporal    
+        mask0 = mask[0:1]  # [1,1,H,W]
+        mask1 = mask[1::k]  # [T'-1,1,H,W]
+        sampled = torch.cat([mask0, mask1], dim=0)  # [T',1,H,W]
+        pooled = sampled.permute(1, 0, 2, 3).unsqueeze(0)
+
+        # Up-sample spatially to match latent spatial resolution
+        spatial_downsample = self.vae_scale_factor_spatial
+        H_latent = pooled.shape[-2] // spatial_downsample
+        W_latent = pooled.shape[-1] // spatial_downsample
+        pooled = F.interpolate(pooled, size=(pooled.shape[2], H_latent, W_latent), mode="nearest")
+
+        # Back to [1, T_latent, 1, H, W]
+        latent_mask = pooled.permute(0, 2, 1, 3, 4)
+
+        return latent_mask
+        
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: PipelineImageInput,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 81,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 5.0,
+        guidance_scale_2: Optional[float] = None,
+        num_videos_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        image_embeds: Optional[torch.Tensor] = None,
+        last_image: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "np",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        motion_signal_video_path: Optional[str] = None,
+        motion_signal_mask_path: Optional[str] = None,
+        tweak_index: int = 0,
+        tstrong_index: int = 0
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            image (`PipelineImageInput`):
+                The input image to condition the generation on. Must be an image, a list of images or a `torch.Tensor`.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            height (`int`, defaults to `480`):
+                The height of the generated video.
+            width (`int`, defaults to `832`):
+                The width of the generated video.
+            num_frames (`int`, defaults to `81`):
+                The number of frames in the generated video.
+            num_inference_steps (`int`, defaults to `50`):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, defaults to `5.0`):
+                Guidance scale as defined in [Classifier-Free Diffusion
+                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
+                the text `prompt`, usually at the expense of lower image quality.
+            guidance_scale_2 (`float`, *optional*, defaults to `None`):
+                Guidance scale for the low-noise stage transformer (`transformer_2`). If `None` and the pipeline's
+                `boundary_ratio` is not None, uses the same value as `guidance_scale`. Only used when `transformer_2`
+                and the pipeline's `boundary_ratio` are not None.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor is generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `negative_prompt` input argument.
+            image_embeds (`torch.Tensor`, *optional*):
+                Pre-generated image embeddings. Can be used to easily tweak image inputs (weighting). If not provided,
+                image embeddings are generated from the `image` input argument.
+            output_type (`str`, *optional*, defaults to `"np"`):
+                The output format of the generated image. Choose between `PIL.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`WanPipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, `PipelineCallback`, `MultiPipelineCallbacks`, *optional*):
+                A function or a subclass of `PipelineCallback` or `MultiPipelineCallbacks` that is called at the end of
+                each denoising step during the inference. with the following arguments: `callback_on_step_end(self:
+                DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a
+                list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int`, defaults to `512`):
+                The maximum sequence length of the text encoder. If the prompt is longer than this, it will be
+                truncated. If the prompt is shorter, it will be padded to this length.
+            motion_signal_video_path (`str`):
+                Path to the video file containing the motion signal to guide the motion of the generated video.
+                It should be a crude version of the reference video, with pixels with motion dragged to their target.
+            motion_signal_mask_path (`str`):
+                Path to the mask video file containing the motion mask of TTM.
+                The mask should be a binary with the conditioning motion pixels being 1 and the rest being 0.
+            tweak_index (`int`):
+                The index of the tweak, from which the denoising process starts.
+            tstrong_index (`int`):
+                The index of the tweak, from which the denoising process starts in the motion conditioned region.
+        Examples:
+
+        Returns:
+            [`~WanPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`WanPipelineOutput`] is returned, otherwise a `tuple` is returned where
+                the first element is a list with the generated images and the second element is a list of `bool`s
+                indicating whether the corresponding generated image contains "not-safe-for-work" (nsfw) content.
+        """
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            negative_prompt,
+            image,
+            height,
+            width,
+            prompt_embeds,
+            negative_prompt_embeds,
+            image_embeds,
+            callback_on_step_end_tensor_inputs,
+            guidance_scale_2,
+        )
+
+        if motion_signal_video_path is None:
+            raise ValueError("`motion_signal_video_path` must be provided for TTM.")
+        if motion_signal_mask_path is None:
+            raise ValueError("`motion_signal_mask_path` must be provided for TTM.")
+
+        if num_frames % self.vae_scale_factor_temporal != 1:
+            logger.warning(
+                f"`num_frames - 1` has to be divisible by {self.vae_scale_factor_temporal}. Rounding to the nearest number."
+            )
+            num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
+        num_frames = max(num_frames, 1)
+
+        if self.config.boundary_ratio is not None and guidance_scale_2 is None:
+            guidance_scale_2 = guidance_scale
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        device = self._execution_device
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # 3. Encode input prompt
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            num_videos_per_prompt=num_videos_per_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            max_sequence_length=max_sequence_length,
+            device=device,
+        )
+
+        # Encode image embedding
+        transformer_dtype = self.transformer.dtype if self.transformer is not None else self.transformer_2.dtype
+        prompt_embeds = prompt_embeds.to(transformer_dtype)
+        if negative_prompt_embeds is not None:
+            negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
+
+        # only wan 2.1 i2v transformer accepts image_embeds
+        if self.transformer is not None and self.transformer.config.image_dim is not None:
+            if image_embeds is None:
+                if last_image is None:
+                    image_embeds = self.encode_image(image, device)
+                else:
+                    image_embeds = self.encode_image([image, last_image], device)
+            image_embeds = image_embeds.repeat(batch_size, 1, 1)
+            image_embeds = image_embeds.to(transformer_dtype)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        tweak_index = int(tweak_index)
+        tstrong_index = int(tstrong_index)
+
+        if tweak_index < -1:
+            raise ValueError(f"`tweak_index` ({tweak_index}) must be >= -1.")
+        if tweak_index >= len(timesteps):
+            raise ValueError(f"`tweak_index` ({tweak_index}) must be < {len(timesteps)}.")
+
+        if tstrong_index < 0:
+            raise ValueError(f"`tstrong_index` ({tstrong_index}) must be >= 0.")
+        if tstrong_index >= len(timesteps):
+            raise ValueError(f"`tstrong_index` ({tstrong_index}) must be < {len(timesteps)}.")
+        if tstrong_index < max(0, tweak_index):
+            raise ValueError(f"`tstrong_index` ({tstrong_index}) must be >= `tweak_index` ({tweak_index}).")
+                            
+        # 5. Prepare latent variables
+        num_channels_latents = self.vae.config.z_dim
+        image = self.video_processor.preprocess(image, height=height, width=width).to(device, dtype=torch.float32)
+        if last_image is not None:
+            last_image = self.video_processor.preprocess(last_image, height=height, width=width).to(
+                device, dtype=torch.float32
+            )
+
+        latents_outputs = self.prepare_latents(
+            image,
+            batch_size * num_videos_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            num_frames,
+            torch.float32,
+            device,
+            generator,
+            latents,
+            last_image,
+        )
+        if self.config.expand_timesteps:
+            latents, condition, first_frame_mask = latents_outputs
+        else:
+            latents, condition = latents_outputs
+
+         # 6. Initialize for TTM
+        ref_vid = load_video_to_tensor(motion_signal_video_path).to(device=device) # shape [1, C, T, H, W]
+        refB, refC, refT, refH, refW = ref_vid.shape
+
+        ref_vid = F.interpolate(
+            ref_vid.permute(0, 2, 1, 3, 4).reshape(refB*refT, refC, refH, refW),
+            size=(height, width), mode="bicubic", align_corners=True, 
+        ).reshape(refB, refT, refC, height, width).permute(0, 2, 1, 3, 4)
+
+        ref_vid = self.video_processor.normalize(ref_vid.to(dtype=self.vae.dtype))         # [1, C, T, H, W]
+        ref_latents = retrieve_latents(self.vae.encode(ref_vid), sample_mode="argmax")     # [1, z, T', H', W']
+        latents_mean = torch.tensor(self.vae.config.latents_mean)\
+            .view(1, self.vae.config.z_dim, 1, 1, 1).to(ref_latents.device, ref_latents.dtype)
+        latents_std = 1.0 / torch.tensor(self.vae.config.latents_std)\
+            .view(1, self.vae.config.z_dim, 1, 1, 1).to(ref_latents.device, ref_latents.dtype)
+        ref_latents = (ref_latents - latents_mean) * latents_std      
+
+        
+        ref_mask = load_video_to_tensor(motion_signal_mask_path).to(device=device) # shape [1, C, T, H, W]
+        mB, mC, mT, mH, mW = ref_mask.shape
+        ref_mask = F.interpolate(
+            ref_mask.permute(0, 2, 1, 3, 4).reshape(mB*mT, mC, mH, mW),
+            size=(height, width), mode="nearest", 
+        ).reshape(mB, mT, mC, height, width).permute(0, 2, 1, 3, 4) # [1, C, T, H, W] -> [T, C, H, W]
+        mask_tc_hw = ref_mask[0].permute(1, 0, 2, 3).contiguous()
+
+        if mask_tc_hw.shape[0] > num_frames: # Align time dimension to num_frames
+            logger.warning("Mask has %d frames but num_frames=%d; trimming.", mask_tc_hw.shape[0], num_frames)
+            mask_tc_hw = mask_tc_hw[:num_frames]
+        elif mask_tc_hw.shape[0] < num_frames:
+            raise ValueError(f"num_frames ({num_frames}) is greater than mask frames ({mask_tc_hw.shape[0]}). "
+                            "Please pad/extend your mask or lower num_frames.")
+        
+        if mask_tc_hw.shape[1] > 1: # Reduce channels if needed -> [T,1,H,W], binarize once
+            mask_t1_hw = (mask_tc_hw > 0.5).any(dim=1, keepdim=True).float()
+        else:
+            mask_t1_hw = (mask_tc_hw > 0.5).float()
+
+        motion_mask = self.convert_rgb_mask_to_latent_mask(mask_t1_hw).permute(0, 2, 1, 3, 4).contiguous()
+        background_mask = 1.0 - motion_mask
+
+        if tweak_index >= 0:
+            tweak = timesteps[tweak_index]
+            fixed_noise = randn_tensor(
+                ref_latents.shape,
+                generator=generator,
+                device=ref_latents.device,
+                dtype=ref_latents.dtype,
+            )
+            tweak = torch.as_tensor(tweak, device=ref_latents.device, dtype=torch.long).view(1)
+            noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, tweak.long())
+            latents = noisy_latents.to(dtype=latents.dtype, device=latents.device)
+        else:
+            tweak = torch.tensor(-1)
+            fixed_noise = randn_tensor(
+                ref_latents.shape,
+                generator=generator,
+                device=ref_latents.device,
+                dtype=ref_latents.dtype,
+            )
+            tweak_index = 0
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        self._num_timesteps = len(timesteps)
+
+        if self.config.boundary_ratio is not None:
+            boundary_timestep = self.config.boundary_ratio * self.scheduler.config.num_train_timesteps
+        else:
+            boundary_timestep = None
+
+        with self.progress_bar(total=len(timesteps) - tweak_index) as progress_bar:
+            for i, t in enumerate(timesteps[tweak_index:]):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                if boundary_timestep is None or t >= boundary_timestep:
+                    # wan2.1 or high-noise stage in wan2.2
+                    current_model = self.transformer
+                    current_guidance_scale = guidance_scale
+                else:
+                    # low-noise stage in wan2.2
+                    current_model = self.transformer_2
+                    current_guidance_scale = guidance_scale_2
+
+                if self.config.expand_timesteps:
+                    latent_model_input = (1 - first_frame_mask) * condition + first_frame_mask * latents
+                    latent_model_input = latent_model_input.to(transformer_dtype)
+
+                    temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten()
+                    timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
+                else:
+                    latent_model_input = torch.cat([latents, condition], dim=1).to(transformer_dtype)
+                    timestep = t.expand(latents.shape[0])
+
+                with current_model.cache_context("cond"):
+                    noise_pred = current_model(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        encoder_hidden_states_image=image_embeds,
+                        attention_kwargs=attention_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                if self.do_classifier_free_guidance:
+                    with current_model.cache_context("uncond"):
+                        noise_uncond = current_model(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            encoder_hidden_states_image=image_embeds,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                        noise_pred = noise_uncond + current_guidance_scale * (noise_pred - noise_uncond)
+
+                # In between tweak and tstrong, replace mask with noisy reference latents
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+                in_between_tweak_tstrong = (i+tweak_index) < tstrong_index
+
+                if in_between_tweak_tstrong:
+                    if i+tweak_index+1 < len(timesteps):
+                        prev_t = timesteps[i+tweak_index+1]
+                        prev_t = torch.as_tensor(prev_t, device=ref_latents.device, dtype=torch.long).view(1)
+                        noisy_latents = self.scheduler.add_noise(ref_latents, fixed_noise, prev_t.long()).to(dtype=latents.dtype, device=latents.device)
+                        latents = latents * background_mask + noisy_latents * motion_mask
+                    else:
+                        latents = latents * background_mask + ref_latents.to(dtype=latents.dtype, device=latents.device) * motion_mask
+                    
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                    negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - tweak_index - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if self.config.expand_timesteps:
+            latents = (1 - first_frame_mask) * condition + first_frame_mask * latents
+
+        if not output_type == "latent":
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            video = self.vae.decode(latents, return_dict=False)[0]
+
+            video = self.video_processor.postprocess_video(video, output_type=output_type)
+        else:
+            video = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (video,)
+
+        return WanPipelineOutput(frames=video)