diffusers 0.33.1__py3-none-any.whl → 0.35.0__py3-none-any.whl

diffusers/pipelines/skyreels_v2/pipeline_skyreels_v2_diffusion_forcing.py

@@ -0,0 +1,978 @@
+# Copyright 2025 The SkyReels-V2 Team, The Wan Team and 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 html
+import math
+import re
+from copy import deepcopy
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import ftfy
+import torch
+from transformers import AutoTokenizer, UMT5EncoderModel
+
+from ...callbacks import MultiPipelineCallbacks, PipelineCallback
+from ...loaders import SkyReelsV2LoraLoaderMixin
+from ...models import AutoencoderKLWan, SkyReelsV2Transformer3DModel
+from ...schedulers import UniPCMultistepScheduler
+from ...utils import is_ftfy_available, is_torch_xla_available, logging, replace_example_docstring
+from ...utils.torch_utils import randn_tensor
+from ...video_processor import VideoProcessor
+from ..pipeline_utils import DiffusionPipeline
+from .pipeline_output import SkyReelsV2PipelineOutput
+
+
+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
+
+if is_ftfy_available():
+    import ftfy
+
+
+EXAMPLE_DOC_STRING = """\
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import (
+        ...     SkyReelsV2DiffusionForcingPipeline,
+        ...     UniPCMultistepScheduler,
+        ...     AutoencoderKLWan,
+        ... )
+        >>> from diffusers.utils import export_to_video
+
+        >>> # Load the pipeline
+        >>> # Available models:
+        >>> # - Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers
+        >>> # - Skywork/SkyReels-V2-DF-14B-540P-Diffusers
+        >>> # - Skywork/SkyReels-V2-DF-14B-720P-Diffusers
+        >>> vae = AutoencoderKLWan.from_pretrained(
+        ...     "Skywork/SkyReels-V2-DF-14B-720P-Diffusers",
+        ...     subfolder="vae",
+        ...     torch_dtype=torch.float32,
+        ... )
+        >>> pipe = SkyReelsV2DiffusionForcingPipeline.from_pretrained(
+        ...     "Skywork/SkyReels-V2-DF-14B-720P-Diffusers",
+        ...     vae=vae,
+        ...     torch_dtype=torch.bfloat16,
+        ... )
+        >>> flow_shift = 8.0  # 8.0 for T2V, 5.0 for I2V
+        >>> pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
+        >>> pipe = pipe.to("cuda")
+
+        >>> prompt = "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window."
+
+        >>> output = pipe(
+        ...     prompt=prompt,
+        ...     num_inference_steps=30,
+        ...     height=544,
+        ...     width=960,
+        ...     guidance_scale=6.0,  # 6.0 for T2V, 5.0 for I2V
+        ...     num_frames=97,
+        ...     ar_step=5,  # Controls asynchronous inference (0 for synchronous mode)
+        ...     causal_block_size=5,  # Number of frames processed together in a causal block
+        ...     overlap_history=None,  # Number of frames to overlap for smooth transitions in long videos
+        ...     addnoise_condition=20,  # Improves consistency in long video generation
+        ... ).frames[0]
+        >>> export_to_video(output, "video.mp4", fps=24, quality=8)
+        ```
+"""
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r"\s+", " ", text)
+    text = text.strip()
+    return text
+
+
+def prompt_clean(text):
+    text = whitespace_clean(basic_clean(text))
+    return text
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+class SkyReelsV2DiffusionForcingPipeline(DiffusionPipeline, SkyReelsV2LoraLoaderMixin):
+    """
+    Pipeline for Text-to-Video (t2v) generation using SkyReels-V2 with diffusion forcing.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods
+    implemented for all pipelines (downloading, saving, running on a specific device, etc.).
+
+    Args:
+        tokenizer ([`AutoTokenizer`]):
+            Tokenizer from [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5Tokenizer),
+            specifically the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
+        text_encoder ([`UMT5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/umt5-xxl](https://huggingface.co/google/umt5-xxl) variant.
+        transformer ([`SkyReelsV2Transformer3DModel`]):
+            Conditional Transformer to denoise the encoded image latents.
+        scheduler ([`UniPCMultistepScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKLWan`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
+
+    def __init__(
+        self,
+        tokenizer: AutoTokenizer,
+        text_encoder: UMT5EncoderModel,
+        transformer: SkyReelsV2Transformer3DModel,
+        vae: AutoencoderKLWan,
+        scheduler: UniPCMultistepScheduler,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+
+        self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
+        self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
+
+    # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline._get_t5_prompt_embeds
+    def _get_t5_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        num_videos_per_prompt: int = 1,
+        max_sequence_length: int = 226,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        prompt = [prompt_clean(u) for u in prompt]
+        batch_size = len(prompt)
+
+        text_inputs = self.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            add_special_tokens=True,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        text_input_ids, mask = text_inputs.input_ids, text_inputs.attention_mask
+        seq_lens = mask.gt(0).sum(dim=1).long()
+
+        prompt_embeds = self.text_encoder(text_input_ids.to(device), mask.to(device)).last_hidden_state
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+        prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
+        )
+
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.wan.pipeline_wan.WanPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        do_classifier_free_guidance: bool = True,
+        num_videos_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 226,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            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`).
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                Whether to use classifier free guidance or not.
+            num_videos_per_prompt (`int`, *optional*, defaults to 1):
+                Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+            prompt_embeds (`torch.Tensor`, *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.Tensor`, *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.
+            device: (`torch.device`, *optional*):
+                torch device
+            dtype: (`torch.dtype`, *optional*):
+                torch dtype
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        if prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            elif batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+                    f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+                    " the batch size of `prompt`."
+                )
+
+            negative_prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=negative_prompt,
+                num_videos_per_prompt=num_videos_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+                dtype=dtype,
+            )
+
+        return prompt_embeds, negative_prompt_embeds
+
+    def check_inputs(
+        self,
+        prompt,
+        negative_prompt,
+        height,
+        width,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        callback_on_step_end_tensor_inputs=None,
+        overlap_history=None,
+        num_frames=None,
+        base_num_frames=None,
+    ):
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`: {negative_prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+        elif negative_prompt is not None and (
+            not isinstance(negative_prompt, str) and not isinstance(negative_prompt, list)
+        ):
+            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+        if num_frames > base_num_frames and overlap_history is None:
+            raise ValueError(
+                "`overlap_history` is required when `num_frames` exceeds `base_num_frames` to ensure smooth transitions in long video generation. "
+                "Please specify a value for `overlap_history`. Recommended values are 17 or 37."
+            )
+
+    def prepare_latents(
+        self,
+        batch_size: int,
+        num_channels_latents: int = 16,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 97,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        base_latent_num_frames: Optional[int] = None,
+        video_latents: Optional[torch.Tensor] = None,
+        causal_block_size: Optional[int] = None,
+        overlap_history_latent_frames: Optional[int] = None,
+        long_video_iter: Optional[int] = None,
+    ) -> torch.Tensor:
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+
+        num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        latent_height = height // self.vae_scale_factor_spatial
+        latent_width = width // self.vae_scale_factor_spatial
+
+        prefix_video_latents = None
+        prefix_video_latents_frames = 0
+
+        if video_latents is not None:  # long video generation at the iterations other than the first one
+            prefix_video_latents = video_latents[:, :, -overlap_history_latent_frames:]
+
+            if prefix_video_latents.shape[2] % causal_block_size != 0:
+                truncate_len_latents = prefix_video_latents.shape[2] % causal_block_size
+                logger.warning(
+                    f"The length of prefix video latents is truncated by {truncate_len_latents} frames for the causal block size alignment. "
+                    f"This truncation ensures compatibility with the causal block size, which is required for proper processing. "
+                    f"However, it may slightly affect the continuity of the generated video at the truncation boundary."
+                )
+                prefix_video_latents = prefix_video_latents[:, :, :-truncate_len_latents]
+            prefix_video_latents_frames = prefix_video_latents.shape[2]
+
+            finished_frame_num = (
+                long_video_iter * (base_latent_num_frames - overlap_history_latent_frames)
+                + overlap_history_latent_frames
+            )
+            left_frame_num = num_latent_frames - finished_frame_num
+            num_latent_frames = min(left_frame_num + overlap_history_latent_frames, base_latent_num_frames)
+        elif base_latent_num_frames is not None:  # long video generation at the first iteration
+            num_latent_frames = base_latent_num_frames
+        else:  # short video generation
+            num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+
+        shape = (
+            batch_size,
+            num_channels_latents,
+            num_latent_frames,
+            latent_height,
+            latent_width,
+        )
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        return latents, num_latent_frames, prefix_video_latents, prefix_video_latents_frames
+
+    def generate_timestep_matrix(
+        self,
+        num_latent_frames: int,
+        step_template: torch.Tensor,
+        base_num_latent_frames: int,
+        ar_step: int = 5,
+        num_pre_ready: int = 0,
+        causal_block_size: int = 1,
+        shrink_interval_with_mask: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[tuple]]:
+        """
+        This function implements the core diffusion forcing algorithm that creates a coordinated denoising schedule
+        across temporal frames. It supports both synchronous and asynchronous generation modes:
+
+        **Synchronous Mode** (ar_step=0, causal_block_size=1):
+        - All frames are denoised simultaneously at each timestep
+        - Each frame follows the same denoising trajectory: [1000, 800, 600, ..., 0]
+        - Simpler but may have less temporal consistency for long videos
+
+        **Asynchronous Mode** (ar_step>0, causal_block_size>1):
+        - Frames are grouped into causal blocks and processed block/chunk-wise
+        - Each block is denoised in a staggered pattern creating a "denoising wave"
+        - Earlier blocks are more denoised, later blocks lag behind by ar_step timesteps
+        - Creates stronger temporal dependencies and better consistency
+
+        Args:
+            num_latent_frames (int): Total number of latent frames to generate
+            step_template (torch.Tensor): Base timestep schedule (e.g., [1000, 800, 600, ..., 0])
+            base_num_latent_frames (int): Maximum frames the model can process in one forward pass
+            ar_step (int, optional): Autoregressive step size for temporal lag.
+                                   0 = synchronous, >0 = asynchronous. Defaults to 5.
+            num_pre_ready (int, optional):
+                                         Number of frames already denoised (e.g., from prefix in a video2video task).
+                                         Defaults to 0.
+            causal_block_size (int, optional): Number of frames processed as a causal block.
+                                             Defaults to 1.
+            shrink_interval_with_mask (bool, optional): Whether to optimize processing intervals.
+                                                      Defaults to False.
+
+        Returns:
+            tuple containing:
+                - step_matrix (torch.Tensor): Matrix of timesteps for each frame at each iteration Shape:
+                  [num_iterations, num_latent_frames]
+                - step_index (torch.Tensor): Index matrix for timestep lookup Shape: [num_iterations,
+                  num_latent_frames]
+                - step_update_mask (torch.Tensor): Boolean mask indicating which frames to update Shape:
+                  [num_iterations, num_latent_frames]
+                - valid_interval (list[tuple]): List of (start, end) intervals for each iteration
+
+        Raises:
+            ValueError: If ar_step is too small for the given configuration
+        """
+        # Initialize lists to store the scheduling matrices and metadata
+        step_matrix, step_index = [], []  # Will store timestep values and indices for each iteration
+        update_mask, valid_interval = [], []  # Will store update masks and processing intervals
+
+        # Calculate total number of denoising iterations (add 1 for initial noise state)
+        num_iterations = len(step_template) + 1
+
+        # Convert frame counts to block counts for causal processing
+        # Each block contains causal_block_size frames that are processed together
+        # E.g.: 25 frames ÷ 5 = 5 blocks total
+        num_blocks = num_latent_frames // causal_block_size
+        base_num_blocks = base_num_latent_frames // causal_block_size
+
+        # Validate ar_step is sufficient for the given configuration
+        # In asynchronous mode, we need enough timesteps to create the staggered pattern
+        if base_num_blocks < num_blocks:
+            min_ar_step = len(step_template) / base_num_blocks
+            if ar_step < min_ar_step:
+                raise ValueError(f"`ar_step` should be at least {math.ceil(min_ar_step)} in your setting")
+
+        # Extend step_template with boundary values for easier indexing
+        # 999: dummy value for counter starting from 1
+        # 0: final timestep (completely denoised)
+        step_template = torch.cat(
+            [
+                torch.tensor([999], dtype=torch.int64, device=step_template.device),
+                step_template.long(),
+                torch.tensor([0], dtype=torch.int64, device=step_template.device),
+            ]
+        )
+
+        # Initialize the previous row state (tracks denoising progress for each block)
+        # 0 means not started, num_iterations means fully denoised
+        pre_row = torch.zeros(num_blocks, dtype=torch.long)
+
+        # Mark pre-ready frames (e.g., from prefix video for a video2video task) as already at final denoising state
+        if num_pre_ready > 0:
+            pre_row[: num_pre_ready // causal_block_size] = num_iterations
+
+        # Main loop: Generate denoising schedule until all frames are fully denoised
+        while not torch.all(pre_row >= (num_iterations - 1)):
+            # Create new row representing the next denoising step
+            new_row = torch.zeros(num_blocks, dtype=torch.long)
+
+            # Apply diffusion forcing logic for each block
+            for i in range(num_blocks):
+                if i == 0 or pre_row[i - 1] >= (
+                    num_iterations - 1
+                ):  # the first frame or the last frame is completely denoised
+                    new_row[i] = pre_row[i] + 1
+                else:
+                    # Asynchronous mode: lag behind previous block by ar_step timesteps
+                    # This creates the "diffusion forcing" staggered pattern
+                    new_row[i] = new_row[i - 1] - ar_step
+
+            # Clamp values to valid range [0, num_iterations]
+            new_row = new_row.clamp(0, num_iterations)
+
+            # Create update mask: True for blocks that need denoising update at this iteration
+            # Exclude blocks that haven't started (new_row != pre_row) or are finished (new_row != num_iterations)
+            # Final state example: [False, ..., False, True, True, True, True, True]
+            # where first 20 frames are done (False) and last 5 frames still need updates (True)
+            update_mask.append((new_row != pre_row) & (new_row != num_iterations))
+
+            # Store the iteration state
+            step_index.append(new_row)  # Index into step_template
+            step_matrix.append(step_template[new_row])  # Actual timestep values
+            pre_row = new_row  # Update for next iteration
+
+        # For videos longer than model capacity, we process in sliding windows
+        terminal_flag = base_num_blocks
+
+        # Optional optimization: shrink interval based on first update mask
+        if shrink_interval_with_mask:
+            idx_sequence = torch.arange(num_blocks, dtype=torch.int64)
+            update_mask = update_mask[0]
+            update_mask_idx = idx_sequence[update_mask]
+            last_update_idx = update_mask_idx[-1].item()
+            terminal_flag = last_update_idx + 1
+
+        # Each interval defines which frames to process in the current forward pass
+        for curr_mask in update_mask:
+            # Extend terminal flag if current mask has updates beyond current terminal
+            if terminal_flag < num_blocks and curr_mask[terminal_flag]:
+                terminal_flag += 1
+            # Create interval: [start, end) where start ensures we don't exceed model capacity
+            valid_interval.append((max(terminal_flag - base_num_blocks, 0), terminal_flag))
+
+        # Convert lists to tensors for efficient processing
+        step_update_mask = torch.stack(update_mask, dim=0)
+        step_index = torch.stack(step_index, dim=0)
+        step_matrix = torch.stack(step_matrix, dim=0)
+
+        # Each block's schedule is replicated to all frames within that block
+        if causal_block_size > 1:
+            # Expand each block to causal_block_size frames
+            step_update_mask = step_update_mask.unsqueeze(-1).repeat(1, 1, causal_block_size).flatten(1).contiguous()
+            step_index = step_index.unsqueeze(-1).repeat(1, 1, causal_block_size).flatten(1).contiguous()
+            step_matrix = step_matrix.unsqueeze(-1).repeat(1, 1, causal_block_size).flatten(1).contiguous()
+            # Scale intervals from block-level to frame-level
+            valid_interval = [(s * causal_block_size, e * causal_block_size) for s, e in valid_interval]
+
+        return step_matrix, step_index, step_update_mask, valid_interval
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 1.0
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        negative_prompt: Union[str, List[str]] = None,
+        height: int = 544,
+        width: int = 960,
+        num_frames: int = 97,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 6.0,
+        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,
+        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,
+        overlap_history: Optional[int] = None,
+        addnoise_condition: float = 0,
+        base_num_frames: int = 97,
+        ar_step: int = 0,
+        causal_block_size: Optional[int] = None,
+        fps: int = 24,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            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 `544`):
+                The height of the generated video.
+            width (`int`, defaults to `960`):
+                The width of the generated video.
+            num_frames (`int`, defaults to `97`):
+                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 `6.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. (**6.0 for T2V**, **5.0 for I2V**)
+            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.
+            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 [`SkyReelsV2PipelineOutput`] 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`, *optional*, defaults to `512`):
+                The maximum sequence length of the prompt.
+            overlap_history (`int`, *optional*, defaults to `None`):
+                Number of frames to overlap for smooth transitions in long videos. If `None`, the pipeline assumes
+                short video generation mode, and no overlap is applied. 17 and 37 are recommended to set.
+            addnoise_condition (`float`, *optional*, defaults to `0`):
+                This is used to help smooth the long video generation by adding some noise to the clean condition. Too
+                large noise can cause the inconsistency as well. 20 is a recommended value, and you may try larger
+                ones, but it is recommended to not exceed 50.
+            base_num_frames (`int`, *optional*, defaults to `97`):
+                97 or 121 | Base frame count (**97 for 540P**, **121 for 720P**)
+            ar_step (`int`, *optional*, defaults to `0`):
+                Controls asynchronous inference (0 for synchronous mode) You can set `ar_step=5` to enable asynchronous
+                inference. When asynchronous inference, `causal_block_size=5` is recommended while it is not supposed
+                to be set for synchronous generation. Asynchronous inference will take more steps to diffuse the whole
+                sequence which means it will be SLOWER than synchronous mode. In our experiments, asynchronous
+                inference may improve the instruction following and visual consistent performance.
+            causal_block_size (`int`, *optional*, defaults to `None`):
+                The number of frames in each block/chunk. Recommended when using asynchronous inference (when ar_step >
+                0)
+            fps (`int`, *optional*, defaults to `24`):
+                Frame rate of the generated video
+
+        Examples:
+
+        Returns:
+            [`~SkyReelsV2PipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`SkyReelsV2PipelineOutput`] 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,
+            height,
+            width,
+            prompt_embeds,
+            negative_prompt_embeds,
+            callback_on_step_end_tensor_inputs,
+            overlap_history,
+            num_frames,
+            base_num_frames,
+        )
+
+        if addnoise_condition > 60:
+            logger.warning(
+                f"The value of 'addnoise_condition' is too large ({addnoise_condition}) and may cause inconsistencies in long video generation. A value of 20 is recommended."
+            )
+
+        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)
+
+        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,
+        )
+
+        transformer_dtype = self.transformer.dtype
+        prompt_embeds = prompt_embeds.to(transformer_dtype)
+        if negative_prompt_embeds is not None:
+            negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        if causal_block_size is None:
+            causal_block_size = self.transformer.config.num_frame_per_block
+        else:
+            self.transformer._set_ar_attention(causal_block_size)
+
+        fps_embeds = [fps] * prompt_embeds.shape[0]
+        fps_embeds = [0 if i == 16 else 1 for i in fps_embeds]
+
+        # Determine if we're doing long video generation
+        is_long_video = overlap_history is not None and base_num_frames is not None and num_frames > base_num_frames
+        # Initialize accumulated_latents to store all latents in one tensor
+        accumulated_latents = None
+        if is_long_video:
+            # Long video generation setup
+            overlap_history_latent_frames = (overlap_history - 1) // self.vae_scale_factor_temporal + 1
+            num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+            base_latent_num_frames = (
+                (base_num_frames - 1) // self.vae_scale_factor_temporal + 1
+                if base_num_frames is not None
+                else num_latent_frames
+            )
+            n_iter = (
+                1
+                + (num_latent_frames - base_latent_num_frames - 1)
+                // (base_latent_num_frames - overlap_history_latent_frames)
+                + 1
+            )
+        else:
+            # Short video generation setup
+            n_iter = 1
+            base_latent_num_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+
+        # Loop through iterations (multiple iterations only for long videos)
+        for iter_idx in range(n_iter):
+            if is_long_video:
+                logger.debug(f"Processing iteration {iter_idx + 1}/{n_iter} for long video generation...")
+
+            # 5. Prepare latent variables
+            num_channels_latents = self.transformer.config.in_channels
+            latents, current_num_latent_frames, prefix_video_latents, prefix_video_latents_frames = (
+                self.prepare_latents(
+                    batch_size * num_videos_per_prompt,
+                    num_channels_latents,
+                    height,
+                    width,
+                    num_frames,
+                    torch.float32,
+                    device,
+                    generator,
+                    latents if iter_idx == 0 else None,
+                    video_latents=accumulated_latents,  # Pass latents directly instead of decoded video
+                    base_latent_num_frames=base_latent_num_frames if is_long_video else None,
+                    causal_block_size=causal_block_size,
+                    overlap_history_latent_frames=overlap_history_latent_frames if is_long_video else None,
+                    long_video_iter=iter_idx if is_long_video else None,
+                )
+            )
+
+            if prefix_video_latents_frames > 0:
+                latents[:, :, :prefix_video_latents_frames, :, :] = prefix_video_latents.to(transformer_dtype)
+
+            # 6. Prepare sample schedulers and timestep matrix
+            sample_schedulers = []
+            for _ in range(current_num_latent_frames):
+                sample_scheduler = deepcopy(self.scheduler)
+                sample_scheduler.set_timesteps(num_inference_steps, device=device)
+                sample_schedulers.append(sample_scheduler)
+
+            # Different matrix generation for short vs long video
+            step_matrix, _, step_update_mask, valid_interval = self.generate_timestep_matrix(
+                current_num_latent_frames,
+                timesteps,
+                current_num_latent_frames if is_long_video else base_latent_num_frames,
+                ar_step,
+                prefix_video_latents_frames,
+                causal_block_size,
+            )
+
+            # 7. Denoising loop
+            num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+            self._num_timesteps = len(step_matrix)
+
+            with self.progress_bar(total=len(step_matrix)) as progress_bar:
+                for i, t in enumerate(step_matrix):
+                    if self.interrupt:
+                        continue
+
+                    self._current_timestep = t
+                    valid_interval_start, valid_interval_end = valid_interval[i]
+                    latent_model_input = (
+                        latents[:, :, valid_interval_start:valid_interval_end, :, :].to(transformer_dtype).clone()
+                    )
+                    timestep = t.expand(latents.shape[0], -1)[:, valid_interval_start:valid_interval_end].clone()
+
+                    if addnoise_condition > 0 and valid_interval_start < prefix_video_latents_frames:
+                        noise_factor = 0.001 * addnoise_condition
+                        latent_model_input[:, :, valid_interval_start:prefix_video_latents_frames, :, :] = (
+                            latent_model_input[:, :, valid_interval_start:prefix_video_latents_frames, :, :]
+                            * (1.0 - noise_factor)
+                            + torch.randn_like(
+                                latent_model_input[:, :, valid_interval_start:prefix_video_latents_frames, :, :]
+                            )
+                            * noise_factor
+                        )
+                        timestep[:, valid_interval_start:prefix_video_latents_frames] = addnoise_condition
+
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        enable_diffusion_forcing=True,
+                        fps=fps_embeds,
+                        attention_kwargs=attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    if self.do_classifier_free_guidance:
+                        noise_uncond = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            enable_diffusion_forcing=True,
+                            fps=fps_embeds,
+                            attention_kwargs=attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                        noise_pred = noise_uncond + guidance_scale * (noise_pred - noise_uncond)
+
+                    update_mask_i = step_update_mask[i]
+                    for idx in range(valid_interval_start, valid_interval_end):
+                        if update_mask_i[idx].item():
+                            latents[:, :, idx, :, :] = sample_schedulers[idx].step(
+                                noise_pred[:, :, idx - valid_interval_start, :, :],
+                                t[idx],
+                                latents[:, :, idx, :, :],
+                                return_dict=False,
+                            )[0]
+
+                    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(step_matrix) - 1 or (
+                        (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                    ):
+                        progress_bar.update()
+
+                    if XLA_AVAILABLE:
+                        xm.mark_step()
+
+            # Handle latent accumulation for long videos or use the current latents for short videos
+            if is_long_video:
+                if accumulated_latents is None:
+                    accumulated_latents = latents
+                else:
+                    # Keep overlap frames for conditioning but don't include them in final output
+                    accumulated_latents = torch.cat(
+                        [accumulated_latents, latents[:, :, overlap_history_latent_frames:]], dim=2
+                    )
+
+        if is_long_video:
+            latents = accumulated_latents
+
+        self._current_timestep = None
+
+        # Final decoding step - convert latents to pixels
+        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 SkyReelsV2PipelineOutput(frames=video)