Upload pipeline_ltx_video.py

pipeline_ltx_video.py

@@ -0,0 +1,1902 @@
+# Adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py
+import copy
+import inspect
+import math
+import re
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models import AutoencoderKL
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from diffusers.utils import deprecate, logging
+from diffusers.utils.torch_utils import randn_tensor
+from einops import rearrange
+from transformers import (
+    T5EncoderModel,
+    T5Tokenizer,
+    AutoModelForCausalLM,
+    AutoProcessor,
+    AutoTokenizer,
+)
+
+from ltx_video.models.autoencoders.causal_video_autoencoder import (
+    CausalVideoAutoencoder,
+)
+from ltx_video.models.autoencoders.vae_encode import (
+    get_vae_size_scale_factor,
+    latent_to_pixel_coords,
+    vae_decode,
+    vae_encode,
+)
+from ltx_video.models.transformers.symmetric_patchifier import Patchifier
+from ltx_video.models.transformers.transformer3d import Transformer3DModel
+from ltx_video.schedulers.rf import TimestepShifter
+from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy
+from ltx_video.utils.prompt_enhance_utils import generate_cinematic_prompt
+from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
+from ltx_video.models.autoencoders.vae_encode import (
+    un_normalize_latents,
+    normalize_latents,
+)
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+ASPECT_RATIO_1024_BIN = {
+    "0.25": [512.0, 2048.0],
+    "0.28": [512.0, 1856.0],
+    "0.32": [576.0, 1792.0],
+    "0.33": [576.0, 1728.0],
+    "0.35": [576.0, 1664.0],
+    "0.4": [640.0, 1600.0],
+    "0.42": [640.0, 1536.0],
+    "0.48": [704.0, 1472.0],
+    "0.5": [704.0, 1408.0],
+    "0.52": [704.0, 1344.0],
+    "0.57": [768.0, 1344.0],
+    "0.6": [768.0, 1280.0],
+    "0.68": [832.0, 1216.0],
+    "0.72": [832.0, 1152.0],
+    "0.78": [896.0, 1152.0],
+    "0.82": [896.0, 1088.0],
+    "0.88": [960.0, 1088.0],
+    "0.94": [960.0, 1024.0],
+    "1.0": [1024.0, 1024.0],
+    "1.07": [1024.0, 960.0],
+    "1.13": [1088.0, 960.0],
+    "1.21": [1088.0, 896.0],
+    "1.29": [1152.0, 896.0],
+    "1.38": [1152.0, 832.0],
+    "1.46": [1216.0, 832.0],
+    "1.67": [1280.0, 768.0],
+    "1.75": [1344.0, 768.0],
+    "2.0": [1408.0, 704.0],
+    "2.09": [1472.0, 704.0],
+    "2.4": [1536.0, 640.0],
+    "2.5": [1600.0, 640.0],
+    "3.0": [1728.0, 576.0],
+    "4.0": [2048.0, 512.0],
+}
+
+ASPECT_RATIO_512_BIN = {
+    "0.25": [256.0, 1024.0],
+    "0.28": [256.0, 928.0],
+    "0.32": [288.0, 896.0],
+    "0.33": [288.0, 864.0],
+    "0.35": [288.0, 832.0],
+    "0.4": [320.0, 800.0],
+    "0.42": [320.0, 768.0],
+    "0.48": [352.0, 736.0],
+    "0.5": [352.0, 704.0],
+    "0.52": [352.0, 672.0],
+    "0.57": [384.0, 672.0],
+    "0.6": [384.0, 640.0],
+    "0.68": [416.0, 608.0],
+    "0.72": [416.0, 576.0],
+    "0.78": [448.0, 576.0],
+    "0.82": [448.0, 544.0],
+    "0.88": [480.0, 544.0],
+    "0.94": [480.0, 512.0],
+    "1.0": [512.0, 512.0],
+    "1.07": [512.0, 480.0],
+    "1.13": [544.0, 480.0],
+    "1.21": [544.0, 448.0],
+    "1.29": [576.0, 448.0],
+    "1.38": [576.0, 416.0],
+    "1.46": [608.0, 416.0],
+    "1.67": [640.0, 384.0],
+    "1.75": [672.0, 384.0],
+    "2.0": [704.0, 352.0],
+    "2.09": [736.0, 352.0],
+    "2.4": [768.0, 320.0],
+    "2.5": [800.0, 320.0],
+    "3.0": [864.0, 288.0],
+    "4.0": [1024.0, 256.0],
+}
+
+
+# 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,
+    skip_initial_inference_steps: int = 0,
+    skip_final_inference_steps: int = 0,
+    **kwargs,
+):
+    """
+    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 support arbitrary spacing between timesteps. If `None`, then the default
+            timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+            must be `None`.
+        max_timestep ('float', *optional*, defaults to 1.0):
+            The initial noising level for image-to-image/video-to-video. The list if timestamps will be
+            truncated to start with a timestamp greater or equal to this.
+
+    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:
+        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)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+
+        if (
+            skip_initial_inference_steps < 0
+            or skip_final_inference_steps < 0
+            or skip_initial_inference_steps + skip_final_inference_steps
+            >= num_inference_steps
+        ):
+            raise ValueError(
+                "invalid skip inference step values: must be non-negative and the sum of skip_initial_inference_steps and skip_final_inference_steps must be less than the number of inference steps"
+            )
+
+        timesteps = timesteps[
+            skip_initial_inference_steps : len(timesteps) - skip_final_inference_steps
+        ]
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        num_inference_steps = len(timesteps)
+
+    return timesteps, num_inference_steps
+
+
+@dataclass
+class ConditioningItem:
+    """
+    Defines a single frame-conditioning item - a single frame or a sequence of frames.
+
+    Attributes:
+        media_item (torch.Tensor): shape=(b, 3, f, h, w). The media item to condition on.
+        media_frame_number (int): The start-frame number of the media item in the generated video.
+        conditioning_strength (float): The strength of the conditioning (1.0 = full conditioning).
+        media_x (Optional[int]): Optional left x coordinate of the media item in the generated frame.
+        media_y (Optional[int]): Optional top y coordinate of the media item in the generated frame.
+    """
+
+    media_item: torch.Tensor
+    media_frame_number: int
+    conditioning_strength: float
+    media_x: Optional[int] = None
+    media_y: Optional[int] = None
+
+
+class LTXVideoPipeline(DiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using LTX-Video.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`T5EncoderModel`]):
+            Frozen text-encoder. This uses
+            [T5](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5EncoderModel), specifically the
+            [t5-v1_1-xxl](https://huggingface.co/PixArt-alpha/PixArt-alpha/tree/main/t5-v1_1-xxl) variant.
+        tokenizer (`T5Tokenizer`):
+            Tokenizer of class
+            [T5Tokenizer](https://huggingface.co/docs/transformers/model_doc/t5#transformers.T5Tokenizer).
+        transformer ([`Transformer2DModel`]):
+            A text conditioned `Transformer2DModel` to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+    """
+
+    bad_punct_regex = re.compile(
+        r"["
+        + "#®•©™&@·º½¾¿¡§~"
+        + r"\)"
+        + r"\("
+        + r"\]"
+        + r"\["
+        + r"\}"
+        + r"\{"
+        + r"\|"
+        + "\\"
+        + r"\/"
+        + r"\*"
+        + r"]{1,}"
+    )  # noqa
+
+    _optional_components = [
+        "tokenizer",
+        "text_encoder",
+        "prompt_enhancer_image_caption_model",
+        "prompt_enhancer_image_caption_processor",
+        "prompt_enhancer_llm_model",
+        "prompt_enhancer_llm_tokenizer",
+    ]
+    model_cpu_offload_seq = "prompt_enhancer_image_caption_model->prompt_enhancer_llm_model->text_encoder->transformer->vae"
+
+    def __init__(
+        self,
+        tokenizer: T5Tokenizer,
+        text_encoder: T5EncoderModel,
+        vae: AutoencoderKL,
+        transformer: Transformer3DModel,
+        scheduler: DPMSolverMultistepScheduler,
+        patchifier: Patchifier,
+        prompt_enhancer_image_caption_model: AutoModelForCausalLM,
+        prompt_enhancer_image_caption_processor: AutoProcessor,
+        prompt_enhancer_llm_model: AutoModelForCausalLM,
+        prompt_enhancer_llm_tokenizer: AutoTokenizer,
+        allowed_inference_steps: Optional[List[float]] = None,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            tokenizer=tokenizer,
+            text_encoder=text_encoder,
+            vae=vae,
+            transformer=transformer,
+            scheduler=scheduler,
+            patchifier=patchifier,
+            prompt_enhancer_image_caption_model=prompt_enhancer_image_caption_model,
+            prompt_enhancer_image_caption_processor=prompt_enhancer_image_caption_processor,
+            prompt_enhancer_llm_model=prompt_enhancer_llm_model,
+            prompt_enhancer_llm_tokenizer=prompt_enhancer_llm_tokenizer,
+        )
+
+        self.video_scale_factor, self.vae_scale_factor, _ = get_vae_size_scale_factor(
+            self.vae
+        )
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+
+        self.allowed_inference_steps = allowed_inference_steps
+
+    def mask_text_embeddings(self, emb, mask):
+        if emb.shape[0] == 1:
+            keep_index = mask.sum().item()
+            return emb[:, :, :keep_index, :], keep_index
+        else:
+            masked_feature = emb * mask[:, None, :, None]
+            return masked_feature, emb.shape[2]
+
+    # Adapted from diffusers.pipelines.deepfloyd_if.pipeline_if.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        do_classifier_free_guidance: bool = True,
+        negative_prompt: str = "",
+        num_images_per_prompt: int = 1,
+        device: Optional[torch.device] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        text_encoder_max_tokens: int = 256,
+        **kwargs,
+    ):
+        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 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`). For
+                This should be "".
+            do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+                whether to use classifier free guidance or not
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                number of images that should be generated per prompt
+            device: (`torch.device`, *optional*):
+                torch device to place the resulting embeddings on
+            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.
+        """
+
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        if device is None:
+            device = self._execution_device
+
+        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]
+
+        # See Section 3.1. of the paper.
+        max_length = (
+            text_encoder_max_tokens  # TPU supports only lengths multiple of 128
+        )
+        if prompt_embeds is None:
+            assert (
+                self.text_encoder is not None
+            ), "You should provide either prompt_embeds or self.text_encoder should not be None,"
+            text_enc_device = next(self.text_encoder.parameters()).device
+            prompt = self._text_preprocessing(prompt)
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(
+                prompt, padding="longest", return_tensors="pt"
+            ).input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[
+                -1
+            ] and not torch.equal(text_input_ids, untruncated_ids):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, max_length - 1 : -1]
+                )
+                logger.warning(
+                    "The following part of your input was truncated because CLIP can only handle sequences up to"
+                    f" {max_length} tokens: {removed_text}"
+                )
+
+            prompt_attention_mask = text_inputs.attention_mask
+            prompt_attention_mask = prompt_attention_mask.to(text_enc_device)
+            prompt_attention_mask = prompt_attention_mask.to(device)
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(text_enc_device), attention_mask=prompt_attention_mask
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            dtype = self.text_encoder.dtype
+        elif self.transformer is not None:
+            dtype = self.transformer.dtype
+        else:
+            dtype = None
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(
+            bs_embed * num_images_per_prompt, seq_len, -1
+        )
+        prompt_attention_mask = prompt_attention_mask.repeat(1, num_images_per_prompt)
+        prompt_attention_mask = prompt_attention_mask.view(
+            bs_embed * num_images_per_prompt, -1
+        )
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens = self._text_preprocessing(negative_prompt)
+            uncond_tokens = uncond_tokens * batch_size
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_attention_mask=True,
+                add_special_tokens=True,
+                return_tensors="pt",
+            )
+            negative_prompt_attention_mask = uncond_input.attention_mask
+            negative_prompt_attention_mask = negative_prompt_attention_mask.to(
+                text_enc_device
+            )
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(text_enc_device),
+                attention_mask=negative_prompt_attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(
+                dtype=dtype, device=device
+            )
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(
+                1, num_images_per_prompt, 1
+            )
+            negative_prompt_embeds = negative_prompt_embeds.view(
+                batch_size * num_images_per_prompt, seq_len, -1
+            )
+
+            negative_prompt_attention_mask = negative_prompt_attention_mask.repeat(
+                1, num_images_per_prompt
+            )
+            negative_prompt_attention_mask = negative_prompt_attention_mask.view(
+                bs_embed * num_images_per_prompt, -1
+            )
+        else:
+            negative_prompt_embeds = None
+            negative_prompt_attention_mask = None
+
+        return (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        )
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(
+            inspect.signature(self.scheduler.step).parameters.keys()
+        )
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_attention_mask=None,
+        negative_prompt_attention_mask=None,
+        enhance_prompt=False,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(
+                f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
+            )
+
+        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 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)}"
+            )
+
+        if prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if 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`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_attention_mask is None:
+            raise ValueError(
+                "Must provide `prompt_attention_mask` when specifying `prompt_embeds`."
+            )
+
+        if (
+            negative_prompt_embeds is not None
+            and negative_prompt_attention_mask is None
+        ):
+            raise ValueError(
+                "Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+            if prompt_attention_mask.shape != negative_prompt_attention_mask.shape:
+                raise ValueError(
+                    "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but"
+                    f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`"
+                    f" {negative_prompt_attention_mask.shape}."
+                )
+
+        if enhance_prompt:
+            assert (
+                self.prompt_enhancer_image_caption_model is not None
+            ), "Image caption model must be initialized if enhance_prompt is True"
+            assert (
+                self.prompt_enhancer_image_caption_processor is not None
+            ), "Image caption processor must be initialized if enhance_prompt is True"
+            assert (
+                self.prompt_enhancer_llm_model is not None
+            ), "Text prompt enhancer model must be initialized if enhance_prompt is True"
+            assert (
+                self.prompt_enhancer_llm_tokenizer is not None
+            ), "Text prompt enhancer tokenizer must be initialized if enhance_prompt is True"
+
+    def _text_preprocessing(self, text):
+        if not isinstance(text, (tuple, list)):
+            text = [text]
+
+        def process(text: str):
+            text = text.strip()
+            return text
+
+        return [process(t) for t in text]
+
+    @staticmethod
+    def add_noise_to_image_conditioning_latents(
+        t: float,
+        init_latents: torch.Tensor,
+        latents: torch.Tensor,
+        noise_scale: float,
+        conditioning_mask: torch.Tensor,
+        generator,
+        eps=1e-6,
+    ):
+        """
+        Add timestep-dependent noise to the hard-conditioning latents.
+        This helps with motion continuity, especially when conditioned on a single frame.
+        """
+        noise = randn_tensor(
+            latents.shape,
+            generator=generator,
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+        # Add noise only to hard-conditioning latents (conditioning_mask = 1.0)
+        need_to_noise = (conditioning_mask > 1.0 - eps).unsqueeze(-1)
+        noised_latents = init_latents + noise_scale * noise * (t**2)
+        latents = torch.where(need_to_noise, noised_latents, latents)
+        return latents
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_latents
+    def prepare_latents(
+        self,
+        latents: torch.Tensor | None,
+        media_items: torch.Tensor | None,
+        timestep: float,
+        latent_shape: torch.Size | Tuple[Any, ...],
+        dtype: torch.dtype,
+        device: torch.device,
+        generator: torch.Generator | List[torch.Generator],
+        vae_per_channel_normalize: bool = True,
+    ):
+        """
+        Prepare the initial latent tensor to be denoised.
+        The latents are either pure noise or a noised version of the encoded media items.
+        Args:
+            latents (`torch.FloatTensor` or `None`):
+                The latents to use (provided by the user) or `None` to create new latents.
+            media_items (`torch.FloatTensor` or `None`):
+                An image or video to be updated using img2img or vid2vid. The media item is encoded and noised.
+            timestep (`float`):
+                The timestep to noise the encoded media_items to.
+            latent_shape (`torch.Size`):
+                The target latent shape.
+            dtype (`torch.dtype`):
+                The target dtype.
+            device (`torch.device`):
+                The target device.
+            generator (`torch.Generator` or `List[torch.Generator]`):
+                Generator(s) to be used for the noising process.
+            vae_per_channel_normalize ('bool'):
+                When encoding the media_items, whether to normalize the latents per-channel.
+        Returns:
+            `torch.FloatTensor`: The latents to be used for the denoising process. This is a tensor of shape
+            (batch_size, num_channels, height, width).
+        """
+        if isinstance(generator, list) and len(generator) != latent_shape[0]:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {latent_shape[0]}. Make sure the batch size matches the length of the generators."
+            )
+
+        # Initialize the latents with the given latents or encoded media item, if provided
+        assert (
+            latents is None or media_items is None
+        ), "Cannot provide both latents and media_items. Please provide only one of the two."
+
+        assert (
+            latents is None and media_items is None or timestep < 1.0
+        ), "Input media_item or latents are provided, but they will be replaced with noise."
+
+        if media_items is not None:
+            latents = vae_encode(
+                media_items.to(dtype=self.vae.dtype, device=self.vae.device),
+                self.vae,
+                vae_per_channel_normalize=vae_per_channel_normalize,
+            )
+        if latents is not None:
+            assert (
+                latents.shape == latent_shape
+            ), f"Latents have to be of shape {latent_shape} but are {latents.shape}."
+            latents = latents.to(device=device, dtype=dtype)
+
+        # For backward compatibility, generate in the "patchified" shape and rearrange
+        b, c, f, h, w = latent_shape
+        noise = randn_tensor(
+            (b, f * h * w, c), generator=generator, device=device, dtype=dtype
+        )
+        noise = rearrange(noise, "b (f h w) c -> b c f h w", f=f, h=h, w=w)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        noise = noise * self.scheduler.init_noise_sigma
+
+        if latents is None:
+            latents = noise
+        else:
+            # Noise the latents to the required (first) timestep
+            latents = timestep * noise + (1 - timestep) * latents
+
+        return latents
+
+    @staticmethod
+    def classify_height_width_bin(
+        height: int, width: int, ratios: dict
+    ) -> Tuple[int, int]:
+        """Returns binned height and width."""
+        ar = float(height / width)
+        closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - ar))
+        default_hw = ratios[closest_ratio]
+        return int(default_hw[0]), int(default_hw[1])
+
+    @staticmethod
+    def resize_and_crop_tensor(
+        samples: torch.Tensor, new_width: int, new_height: int
+    ) -> torch.Tensor:
+        n_frames, orig_height, orig_width = samples.shape[-3:]
+
+        # Check if resizing is needed
+        if orig_height != new_height or orig_width != new_width:
+            ratio = max(new_height / orig_height, new_width / orig_width)
+            resized_width = int(orig_width * ratio)
+            resized_height = int(orig_height * ratio)
+
+            # Resize
+            samples = LTXVideoPipeline.resize_tensor(
+                samples, resized_height, resized_width
+            )
+
+            # Center Crop
+            start_x = (resized_width - new_width) // 2
+            end_x = start_x + new_width
+            start_y = (resized_height - new_height) // 2
+            end_y = start_y + new_height
+            samples = samples[..., start_y:end_y, start_x:end_x]
+
+        return samples
+
+    @staticmethod
+    def resize_tensor(media_items, height, width):
+        n_frames = media_items.shape[2]
+        if media_items.shape[-2:] != (height, width):
+            media_items = rearrange(media_items, "b c n h w -> (b n) c h w")
+            media_items = F.interpolate(
+                media_items,
+                size=(height, width),
+                mode="bilinear",
+                align_corners=False,
+            )
+            media_items = rearrange(media_items, "(b n) c h w -> b c n h w", n=n_frames)
+        return media_items
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        skip_initial_inference_steps: int = 0,
+        skip_final_inference_steps: int = 0,
+        timesteps: List[int] = None,
+        guidance_scale: Union[float, List[float]] = 4.5,
+        cfg_star_rescale: bool = False,
+        skip_layer_strategy: Optional[SkipLayerStrategy] = None,
+        skip_block_list: Optional[Union[List[List[int]], List[int]]] = None,
+        stg_scale: Union[float, List[float]] = 1.0,
+        rescaling_scale: Union[float, List[float]] = 0.7,
+        guidance_timesteps: Optional[List[int]] = None,
+        num_images_per_prompt: Optional[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,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        conditioning_items: Optional[List[ConditioningItem]] = None,
+        decode_timestep: Union[List[float], float] = 0.0,
+        decode_noise_scale: Optional[List[float]] = None,
+        mixed_precision: bool = False,
+        offload_to_cpu: bool = False,
+        enhance_prompt: bool = False,
+        text_encoder_max_tokens: int = 256,
+        stochastic_sampling: bool = False,
+        media_items: Optional[torch.Tensor] = None,
+        tone_map_compression_ratio: float = 0.0,
+        **kwargs,
+    ) -> Union[ImagePipelineOutput, Tuple]:
+        """
+        Function invoked when calling 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`).
+            num_inference_steps (`int`, *optional*, defaults to 100):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference. If `timesteps` is provided, this parameter is ignored.
+            skip_initial_inference_steps (`int`, *optional*, defaults to 0):
+                The number of initial timesteps to skip. After calculating the timesteps, this number of timesteps will
+                be removed from the beginning of the timesteps list. Meaning the highest-timesteps values will not run.
+            skip_final_inference_steps (`int`, *optional*, defaults to 0):
+                The number of final timesteps to skip. After calculating the timesteps, this number of timesteps will
+                be removed from the end of the timesteps list. Meaning the lowest-timesteps values will not run.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process. If not defined, equal spaced `num_inference_steps`
+                timesteps are used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 4.5):
+                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.
+            cfg_star_rescale (`bool`, *optional*, defaults to `False`):
+                If set to `True`, applies the CFG star rescale. Scales the negative prediction according to dot
+                product between positive and negative.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size):
+                The width in pixels of the generated image.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            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 ge 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.
+            prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for text embeddings.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. This negative prompt should be "". If not
+                provided, negative_prompt_embeds will be generated from `negative_prompt` input argument.
+            negative_prompt_attention_mask (`torch.FloatTensor`, *optional*):
+                Pre-generated attention mask for negative text embeddings.
+            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 to return a [`~pipelines.stable_diffusion.IFPipelineOutput`] instead of a plain tuple.
+            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`.
+            use_resolution_binning (`bool` defaults to `True`):
+                If set to `True`, the requested height and width are first mapped to the closest resolutions using
+                `ASPECT_RATIO_1024_BIN`. After the produced latents are decoded into images, they are resized back to
+                the requested resolution. Useful for generating non-square images.
+            enhance_prompt (`bool`, *optional*, defaults to `False`):
+                If set to `True`, the prompt is enhanced using a LLM model.
+            text_encoder_max_tokens (`int`, *optional*, defaults to `256`):
+                The maximum number of tokens to use for the text encoder.
+            stochastic_sampling (`bool`, *optional*, defaults to `False`):
+                If set to `True`, the sampling is stochastic. If set to `False`, the sampling is deterministic.
+            media_items ('torch.Tensor', *optional*):
+                The input media item used for image-to-image / video-to-video.
+            tone_map_compression_ratio: compression ratio for tone mapping, defaults to 0.0.
+                        If set to 0.0, no tone mapping is applied. If set to 1.0 - full compression is applied.
+        Examples:
+
+        Returns:
+            [`~pipelines.ImagePipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.ImagePipelineOutput`] is returned, otherwise a `tuple` is
+                returned where the first element is a list with the generated images
+        """
+        if "mask_feature" in kwargs:
+            deprecation_message = "The use of `mask_feature` is deprecated. It is no longer used in any computation and that doesn't affect the end results. It will be removed in a future version."
+            deprecate("mask_feature", "1.0.0", deprecation_message, standard_warn=False)
+
+        is_video = kwargs.get("is_video", False)
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        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
+
+        self.video_scale_factor = self.video_scale_factor if is_video else 1
+        vae_per_channel_normalize = kwargs.get("vae_per_channel_normalize", True)
+        image_cond_noise_scale = kwargs.get("image_cond_noise_scale", 0.0)
+
+        latent_height = height // self.vae_scale_factor
+        latent_width = width // self.vae_scale_factor
+        latent_num_frames = num_frames // self.video_scale_factor
+        if isinstance(self.vae, CausalVideoAutoencoder) and is_video:
+            latent_num_frames += 1
+        latent_shape = (
+            batch_size * num_images_per_prompt,
+            self.transformer.config.in_channels,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+        )
+
+        # Prepare the list of denoising time-steps
+
+        retrieve_timesteps_kwargs = {}
+        if isinstance(self.scheduler, TimestepShifter):
+            retrieve_timesteps_kwargs["samples_shape"] = latent_shape
+
+        assert (
+            skip_initial_inference_steps == 0
+            or latents is not None
+            or media_items is not None
+        ), (
+            f"skip_initial_inference_steps ({skip_initial_inference_steps}) is used for image-to-image/video-to-video - "
+            "media_item or latents should be provided."
+        )
+
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            skip_initial_inference_steps=skip_initial_inference_steps,
+            skip_final_inference_steps=skip_final_inference_steps,
+            **retrieve_timesteps_kwargs,
+        )
+
+        if self.allowed_inference_steps is not None:
+            for timestep in [round(x, 4) for x in timesteps.tolist()]:
+                assert (
+                    timestep in self.allowed_inference_steps
+                ), f"Invalid inference timestep {timestep}. Allowed timesteps are {self.allowed_inference_steps}."
+
+
+        # Versão Corrigida em pipeline_ltx_video.py
+        if guidance_timesteps:
+            guidance_mapping = []
+            for timestep in timesteps:
+                indices = [
+                    i for i, val in enumerate(guidance_timesteps) if val <= timestep
+                ]
+                guidance_mapping.append(
+                    indices[0] if len(indices) > 0 else (len(guidance_timesteps) - 1)
+                )
+        else:
+            # CORREÇÃO: Garante que guidance_mapping exista mesmo se guidance_timesteps não for fornecido.
+            # Cria um mapeamento de identidade seguro.
+            guidance_mapping = list(range(len(timesteps)))
+            
+            
+            
+        # 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.
+        if not isinstance(guidance_scale, List):
+            guidance_scale = [guidance_scale] * len(timesteps)
+        else:
+            guidance_scale = [
+                guidance_scale[guidance_mapping[i]] for i in range(len(timesteps))
+            ]
+
+        if not isinstance(stg_scale, List):
+            stg_scale = [stg_scale] * len(timesteps)
+        else:
+            stg_scale = [stg_scale[guidance_mapping[i]] for i in range(len(timesteps))]
+
+        if not isinstance(rescaling_scale, List):
+            rescaling_scale = [rescaling_scale] * len(timesteps)
+        else:
+            rescaling_scale = [
+                rescaling_scale[guidance_mapping[i]] for i in range(len(timesteps))
+            ]
+
+        # Normalize skip_block_list to always be None or a list of lists matching timesteps
+        if skip_block_list is not None:
+            # Convert single list to list of lists if needed
+            if len(skip_block_list) == 0 or not isinstance(skip_block_list[0], list):
+                skip_block_list = [skip_block_list] * len(timesteps)
+            else:
+                new_skip_block_list = []
+                for i, timestep in enumerate(timesteps):
+                    new_skip_block_list.append(skip_block_list[guidance_mapping[i]])
+                skip_block_list = new_skip_block_list
+
+        if enhance_prompt:
+            self.prompt_enhancer_image_caption_model = (
+                self.prompt_enhancer_image_caption_model.to(self._execution_device)
+            )
+            self.prompt_enhancer_llm_model = self.prompt_enhancer_llm_model.to(
+                self._execution_device
+            )
+
+            prompt = generate_cinematic_prompt(
+                self.prompt_enhancer_image_caption_model,
+                self.prompt_enhancer_image_caption_processor,
+                self.prompt_enhancer_llm_model,
+                self.prompt_enhancer_llm_tokenizer,
+                prompt,
+                conditioning_items,
+                max_new_tokens=text_encoder_max_tokens,
+            )
+
+        # 3. Encode input prompt
+        if self.text_encoder is not None:
+            self.text_encoder = self.text_encoder.to(self._execution_device)
+
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            True,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            text_encoder_max_tokens=text_encoder_max_tokens,
+        )
+
+        if offload_to_cpu and self.text_encoder is not None:
+            self.text_encoder = self.text_encoder.cpu()
+
+        self.transformer = self.transformer.to(self._execution_device)
+
+        prompt_embeds_batch = prompt_embeds
+        prompt_attention_mask_batch = prompt_attention_mask
+        negative_prompt_embeds = (
+            torch.zeros_like(prompt_embeds)
+            if negative_prompt_embeds is None
+            else negative_prompt_embeds
+        )
+        negative_prompt_attention_mask = (
+            torch.zeros_like(prompt_attention_mask)
+            if negative_prompt_attention_mask is None
+            else negative_prompt_attention_mask
+        )
+
+        prompt_embeds_batch = torch.cat(
+            [negative_prompt_embeds, prompt_embeds, prompt_embeds], dim=0
+        )
+        prompt_attention_mask_batch = torch.cat(
+            [
+                negative_prompt_attention_mask,
+                prompt_attention_mask,
+                prompt_attention_mask,
+            ],
+            dim=0,
+        )
+        # 4. Prepare the initial latents using the provided media and conditioning items
+
+        # Prepare the initial latents tensor, shape = (b, c, f, h, w)
+        latents = self.prepare_latents(
+            latents=latents,
+            media_items=media_items,
+            timestep=timesteps[0],
+            latent_shape=latent_shape,
+            dtype=prompt_embeds.dtype,
+            device=device,
+            generator=generator,
+            vae_per_channel_normalize=vae_per_channel_normalize,
+        )
+
+        # Update the latents with the conditioning items and patchify them into (b, n, c)
+        latents, pixel_coords, conditioning_mask, num_cond_latents = (
+            self.prepare_conditioning(
+                conditioning_items=conditioning_items,
+                init_latents=latents,
+                num_frames=num_frames,
+                height=height,
+                width=width,
+                vae_per_channel_normalize=vae_per_channel_normalize,
+                generator=generator,
+            )
+        )
+        init_latents = latents.clone()  # Used for image_cond_noise_update
+
+        # 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. Denoising loop
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+
+        orig_conditioning_mask = conditioning_mask
+
+        # Befor compiling this code please be aware:
+        # This code might generate different input shapes if some timesteps have no STG or CFG.
+        # This means that the codes might need to be compiled mutliple times.
+        # To avoid that, use the same STG and CFG values for all timesteps.
+
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                do_classifier_free_guidance = guidance_scale[i] > 1.0
+                do_spatio_temporal_guidance = stg_scale[i] > 0
+                do_rescaling = rescaling_scale[i] != 1.0
+
+                num_conds = 1
+                if do_classifier_free_guidance:
+                    num_conds += 1
+                if do_spatio_temporal_guidance:
+                    num_conds += 1
+
+                if do_classifier_free_guidance and do_spatio_temporal_guidance:
+                    indices = slice(batch_size * 0, batch_size * 3)
+                elif do_classifier_free_guidance:
+                    indices = slice(batch_size * 0, batch_size * 2)
+                elif do_spatio_temporal_guidance:
+                    indices = slice(batch_size * 1, batch_size * 3)
+                else:
+                    indices = slice(batch_size * 1, batch_size * 2)
+
+                # Prepare skip layer masks
+                skip_layer_mask: Optional[torch.Tensor] = None
+                if do_spatio_temporal_guidance:
+                    if skip_block_list is not None:
+                        skip_layer_mask = self.transformer.create_skip_layer_mask(
+                            batch_size, num_conds, num_conds - 1, skip_block_list[i]
+                        )
+
+                batch_pixel_coords = torch.cat([pixel_coords] * num_conds)
+                conditioning_mask = orig_conditioning_mask
+                if conditioning_mask is not None and is_video:
+                    assert num_images_per_prompt == 1
+                    conditioning_mask = torch.cat([conditioning_mask] * num_conds)
+                fractional_coords = batch_pixel_coords.to(torch.float32)
+                fractional_coords[:, 0] = fractional_coords[:, 0] * (1.0 / frame_rate)
+
+                if conditioning_mask is not None and image_cond_noise_scale > 0.0:
+                    latents = self.add_noise_to_image_conditioning_latents(
+                        t,
+                        init_latents,
+                        latents,
+                        image_cond_noise_scale,
+                        orig_conditioning_mask,
+                        generator,
+                    )
+
+                latent_model_input = (
+                    torch.cat([latents] * num_conds) if num_conds > 1 else latents
+                )
+                latent_model_input = self.scheduler.scale_model_input(
+                    latent_model_input, t
+                )
+
+                current_timestep = t
+                if not torch.is_tensor(current_timestep):
+                    # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                    # This would be a good case for the `match` statement (Python 3.10+)
+                    is_mps = latent_model_input.device.type == "mps"
+                    if isinstance(current_timestep, float):
+                        dtype = torch.float32 if is_mps else torch.float64
+                    else:
+                        dtype = torch.int32 if is_mps else torch.int64
+                    current_timestep = torch.tensor(
+                        [current_timestep],
+                        dtype=dtype,
+                        device=latent_model_input.device,
+                    )
+                elif len(current_timestep.shape) == 0:
+                    current_timestep = current_timestep[None].to(
+                        latent_model_input.device
+                    )
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                current_timestep = current_timestep.expand(
+                    latent_model_input.shape[0]
+                ).unsqueeze(-1)
+
+                if conditioning_mask is not None:
+                    # Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
+                    # and will start to be denoised when the current timestep is lower than their conditioning timestep.
+                    current_timestep = torch.min(
+                        current_timestep, 1.0 - conditioning_mask
+                    )
+
+                # Choose the appropriate context manager based on `mixed_precision`
+                if mixed_precision:
+                    context_manager = torch.autocast(device.type, dtype=torch.bfloat16)
+                else:
+                    context_manager = nullcontext()  # Dummy context manager
+
+                # predict noise model_output
+                with context_manager:
+                    noise_pred = self.transformer(
+                        latent_model_input.to(self.transformer.dtype),
+                        indices_grid=fractional_coords,
+                        encoder_hidden_states=prompt_embeds_batch[indices].to(
+                            self.transformer.dtype
+                        ),
+                        encoder_attention_mask=prompt_attention_mask_batch[indices],
+                        timestep=current_timestep,
+                        skip_layer_mask=skip_layer_mask,
+                        skip_layer_strategy=skip_layer_strategy,
+                        return_dict=False,
+                    )[0]
+
+                # perform guidance
+                if do_spatio_temporal_guidance:
+                    noise_pred_text, noise_pred_text_perturb = noise_pred.chunk(
+                        num_conds
+                    )[-2:]
+                if do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(num_conds)[:2]
+
+                    if cfg_star_rescale:
+                        # Rescales the unconditional noise prediction using the projection of the conditional prediction onto it:
+                        # α = (⟨ε_text, ε_uncond⟩ / ||ε_uncond||²), then ε_uncond ← α * ε_uncond
+                        # where ε_text is the conditional noise prediction and ε_uncond is the unconditional one.
+                        positive_flat = noise_pred_text.view(batch_size, -1)
+                        negative_flat = noise_pred_uncond.view(batch_size, -1)
+                        dot_product = torch.sum(
+                            positive_flat * negative_flat, dim=1, keepdim=True
+                        )
+                        squared_norm = (
+                            torch.sum(negative_flat**2, dim=1, keepdim=True) + 1e-8
+                        )
+                        alpha = dot_product / squared_norm
+                        noise_pred_uncond = alpha * noise_pred_uncond
+
+                    noise_pred = noise_pred_uncond + guidance_scale[i] * (
+                        noise_pred_text - noise_pred_uncond
+                    )
+                elif do_spatio_temporal_guidance:
+                    noise_pred = noise_pred_text
+                if do_spatio_temporal_guidance:
+                    noise_pred = noise_pred + stg_scale[i] * (
+                        noise_pred_text - noise_pred_text_perturb
+                    )
+                    if do_rescaling and stg_scale[i] > 0.0:
+                        noise_pred_text_std = noise_pred_text.view(batch_size, -1).std(
+                            dim=1, keepdim=True
+                        )
+                        noise_pred_std = noise_pred.view(batch_size, -1).std(
+                            dim=1, keepdim=True
+                        )
+
+                        factor = noise_pred_text_std / noise_pred_std
+                        factor = rescaling_scale[i] * factor + (1 - rescaling_scale[i])
+
+                        noise_pred = noise_pred * factor.view(batch_size, 1, 1)
+
+                current_timestep = current_timestep[:1]
+                # learned sigma
+                if (
+                    self.transformer.config.out_channels // 2
+                    == self.transformer.config.in_channels
+                ):
+                    noise_pred = noise_pred.chunk(2, dim=1)[0]
+
+                # compute previous image: x_t -> x_t-1
+                latents = self.denoising_step(
+                    latents,
+                    noise_pred,
+                    current_timestep,
+                    orig_conditioning_mask,
+                    t,
+                    extra_step_kwargs,
+                    stochastic_sampling=stochastic_sampling,
+                )
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or (
+                    (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+                ):
+                    progress_bar.update()
+
+                if callback_on_step_end is not None:
+                    callback_on_step_end(self, i, t, {})
+
+        if offload_to_cpu:
+            self.transformer = self.transformer.cpu()
+            if self._execution_device == "cuda":
+                torch.cuda.empty_cache()
+
+        # Remove the added conditioning latents
+        latents = latents[:, num_cond_latents:]
+
+        latents = self.patchifier.unpatchify(
+            latents=latents,
+            output_height=latent_height,
+            output_width=latent_width,
+            out_channels=self.transformer.in_channels
+            // math.prod(self.patchifier.patch_size),
+        )
+        if output_type != "latent":
+            if self.vae.decoder.timestep_conditioning:
+                noise = torch.randn_like(latents)
+                if not isinstance(decode_timestep, list):
+                    decode_timestep = [decode_timestep] * latents.shape[0]
+                if decode_noise_scale is None:
+                    decode_noise_scale = decode_timestep
+                elif not isinstance(decode_noise_scale, list):
+                    decode_noise_scale = [decode_noise_scale] * latents.shape[0]
+
+                decode_timestep = torch.tensor(decode_timestep).to(latents.device)
+                decode_noise_scale = torch.tensor(decode_noise_scale).to(
+                    latents.device
+                )[:, None, None, None, None]
+                latents = (
+                    latents * (1 - decode_noise_scale) + noise * decode_noise_scale
+                )
+            else:
+                decode_timestep = None
+            latents = self.tone_map_latents(latents, tone_map_compression_ratio)
+            image = vae_decode(
+                latents,
+                self.vae,
+                is_video,
+                vae_per_channel_normalize=kwargs["vae_per_channel_normalize"],
+                timestep=decode_timestep,
+            )
+
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        else:
+            image = latents
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return ImagePipelineOutput(images=image)
+
+    def denoising_step(
+        self,
+        latents: torch.Tensor,
+        noise_pred: torch.Tensor,
+        current_timestep: torch.Tensor,
+        conditioning_mask: torch.Tensor,
+        t: float,
+        extra_step_kwargs,
+        t_eps=1e-6,
+        stochastic_sampling=False,
+    ):
+        """
+        Perform the denoising step for the required tokens, based on the current timestep and
+        conditioning mask:
+        Conditioning latents have an initial timestep and noising level of (1.0 - conditioning_mask)
+        and will start to be denoised when the current timestep is equal or lower than their
+        conditioning timestep.
+        (hard-conditioning latents with conditioning_mask = 1.0 are never denoised)
+        """
+        # Denoise the latents using the scheduler
+        denoised_latents = self.scheduler.step(
+            noise_pred,
+            t if current_timestep is None else current_timestep,
+            latents,
+            **extra_step_kwargs,
+            return_dict=False,
+            stochastic_sampling=stochastic_sampling,
+        )[0]
+
+        if conditioning_mask is None:
+            return denoised_latents
+
+        tokens_to_denoise_mask = (t - t_eps < (1.0 - conditioning_mask)).unsqueeze(-1)
+        return torch.where(tokens_to_denoise_mask, denoised_latents, latents)
+
+    def prepare_conditioning(
+        self,
+        conditioning_items: Optional[List[ConditioningItem]],
+        init_latents: torch.Tensor,
+        num_frames: int,
+        height: int,
+        width: int,
+        vae_per_channel_normalize: bool = False,
+        generator=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+        """
+        Prepare conditioning tokens based on the provided conditioning items.
+
+        This method encodes provided conditioning items (video frames or single frames) into latents
+        and integrates them with the initial latent tensor. It also calculates corresponding pixel
+        coordinates, a mask indicating the influence of conditioning latents, and the total number of
+        conditioning latents.
+
+        Args:
+            conditioning_items (Optional[List[ConditioningItem]]): A list of ConditioningItem objects.
+            init_latents (torch.Tensor): The initial latent tensor of shape (b, c, f_l, h_l, w_l), where
+                `f_l` is the number of latent frames, and `h_l` and `w_l` are latent spatial dimensions.
+            num_frames, height, width: The dimensions of the generated video.
+            vae_per_channel_normalize (bool, optional): Whether to normalize channels during VAE encoding.
+                Defaults to `False`.
+            generator: The random generator
+
+        Returns:
+            Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+                - `init_latents` (torch.Tensor): The updated latent tensor including conditioning latents,
+                  patchified into (b, n, c) shape.
+                - `init_pixel_coords` (torch.Tensor): The pixel coordinates corresponding to the updated
+                  latent tensor.
+                - `conditioning_mask` (torch.Tensor): A mask indicating the conditioning-strength of each
+                  latent token.
+                - `num_cond_latents` (int): The total number of latent tokens added from conditioning items.
+
+        Raises:
+            AssertionError: If input shapes, dimensions, or conditions for applying conditioning are invalid.
+        """
+        assert isinstance(self.vae, CausalVideoAutoencoder)
+
+        if conditioning_items:
+            batch_size, _, num_latent_frames = init_latents.shape[:3]
+
+            init_conditioning_mask = torch.zeros(
+                init_latents[:, 0, :, :, :].shape,
+                dtype=torch.float32,
+                device=init_latents.device,
+            )
+
+            extra_conditioning_latents = []
+            extra_conditioning_pixel_coords = []
+            extra_conditioning_mask = []
+            extra_conditioning_num_latents = 0  # Number of extra conditioning latents added (should be removed before decoding)
+
+            # Process each conditioning item
+            for conditioning_item in conditioning_items:
+                conditioning_item = self._resize_conditioning_item(
+                    conditioning_item, height, width
+                )
+                media_item = conditioning_item.media_item
+                media_frame_number = conditioning_item.media_frame_number
+                strength = conditioning_item.conditioning_strength
+                assert media_item.ndim == 5  # (b, c, f, h, w)
+                b, c, n_frames, h, w = media_item.shape
+                assert (
+                    height == h and width == w
+                ) or media_frame_number == 0, f"Dimensions do not match: {height}x{width} != {h}x{w} - allowed only when media_frame_number == 0"
+                assert n_frames % 8 == 1
+                assert (
+                    media_frame_number >= 0
+                    and media_frame_number + n_frames <= num_frames
+                )
+
+                # Encode the provided conditioning media item
+                media_item_latents = vae_encode(
+                    media_item.to(dtype=self.vae.dtype, device=self.vae.device),
+                    self.vae,
+                    vae_per_channel_normalize=vae_per_channel_normalize,
+                ).to(dtype=init_latents.dtype)
+
+                # Handle the different conditioning cases
+                if media_frame_number == 0:
+                    # Get the target spatial position of the latent conditioning item
+                    media_item_latents, l_x, l_y = self._get_latent_spatial_position(
+                        media_item_latents,
+                        conditioning_item,
+                        height,
+                        width,
+                        strip_latent_border=True,
+                    )
+                    b, c_l, f_l, h_l, w_l = media_item_latents.shape
+
+                    # First frame or sequence - just update the initial noise latents and the mask
+                    init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l] = (
+                        torch.lerp(
+                            init_latents[:, :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l],
+                            media_item_latents,
+                            strength,
+                        )
+                    )
+                    init_conditioning_mask[
+                        :, :f_l, l_y : l_y + h_l, l_x : l_x + w_l
+                    ] = strength
+                else:
+                    # Non-first frame or sequence
+                    if n_frames > 1:
+                        # Handle non-first sequence.
+                        # Encoded latents are either fully consumed, or the prefix is handled separately below.
+                        (
+                            init_latents,
+                            init_conditioning_mask,
+                            media_item_latents,
+                        ) = self._handle_non_first_conditioning_sequence(
+                            init_latents,
+                            init_conditioning_mask,
+                            media_item_latents,
+                            media_frame_number,
+                            strength,
+                        )
+
+                    # Single frame or sequence-prefix latents
+                    if media_item_latents is not None:
+                        noise = randn_tensor(
+                            media_item_latents.shape,
+                            generator=generator,
+                            device=media_item_latents.device,
+                            dtype=media_item_latents.dtype,
+                        )
+
+                        media_item_latents = torch.lerp(
+                            noise, media_item_latents, strength
+                        )
+
+                        # Patchify the extra conditioning latents and calculate their pixel coordinates
+                        media_item_latents, latent_coords = self.patchifier.patchify(
+                            latents=media_item_latents
+                        )
+                        pixel_coords = latent_to_pixel_coords(
+                            latent_coords,
+                            self.vae,
+                            causal_fix=self.transformer.config.causal_temporal_positioning,
+                        )
+
+                        # Update the frame numbers to match the target frame number
+                        pixel_coords[:, 0] += media_frame_number
+                        extra_conditioning_num_latents += media_item_latents.shape[1]
+
+                        conditioning_mask = torch.full(
+                            media_item_latents.shape[:2],
+                            strength,
+                            dtype=torch.float32,
+                            device=init_latents.device,
+                        )
+
+                        extra_conditioning_latents.append(media_item_latents)
+                        extra_conditioning_pixel_coords.append(pixel_coords)
+                        extra_conditioning_mask.append(conditioning_mask)
+
+        # Patchify the updated latents and calculate their pixel coordinates
+        init_latents, init_latent_coords = self.patchifier.patchify(
+            latents=init_latents
+        )
+        init_pixel_coords = latent_to_pixel_coords(
+            init_latent_coords,
+            self.vae,
+            causal_fix=self.transformer.config.causal_temporal_positioning,
+        )
+
+        if not conditioning_items:
+            return init_latents, init_pixel_coords, None, 0
+
+        init_conditioning_mask, _ = self.patchifier.patchify(
+            latents=init_conditioning_mask.unsqueeze(1)
+        )
+        init_conditioning_mask = init_conditioning_mask.squeeze(-1)
+
+        if extra_conditioning_latents:
+            # Stack the extra conditioning latents, pixel coordinates and mask
+            init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
+            init_pixel_coords = torch.cat(
+                [*extra_conditioning_pixel_coords, init_pixel_coords], dim=2
+            )
+            init_conditioning_mask = torch.cat(
+                [*extra_conditioning_mask, init_conditioning_mask], dim=1
+            )
+
+            if self.transformer.use_tpu_flash_attention:
+                # When flash attention is used, keep the original number of tokens by removing
+                #   tokens from the end.
+                init_latents = init_latents[:, :-extra_conditioning_num_latents]
+                init_pixel_coords = init_pixel_coords[
+                    :, :, :-extra_conditioning_num_latents
+                ]
+                init_conditioning_mask = init_conditioning_mask[
+                    :, :-extra_conditioning_num_latents
+                ]
+
+        return (
+            init_latents,
+            init_pixel_coords,
+            init_conditioning_mask,
+            extra_conditioning_num_latents,
+        )
+
+    @staticmethod
+    def _resize_conditioning_item(
+        conditioning_item: ConditioningItem,
+        height: int,
+        width: int,
+    ):
+        if conditioning_item.media_x or conditioning_item.media_y:
+            raise ValueError(
+                "Provide media_item in the target size for spatial conditioning."
+            )
+        new_conditioning_item = copy.copy(conditioning_item)
+        new_conditioning_item.media_item = LTXVideoPipeline.resize_tensor(
+            conditioning_item.media_item, height, width
+        )
+        return new_conditioning_item
+
+    def _get_latent_spatial_position(
+        self,
+        latents: torch.Tensor,
+        conditioning_item: ConditioningItem,
+        height: int,
+        width: int,
+        strip_latent_border,
+    ):
+        """
+        Get the spatial position of the conditioning item in the latent space.
+        If requested, strip the conditioning latent borders that do not align with target borders.
+        (border latents look different then other latents and might confuse the model)
+        """
+        scale = self.vae_scale_factor
+        h, w = conditioning_item.media_item.shape[-2:]
+        assert (
+            h <= height and w <= width
+        ), f"Conditioning item size {h}x{w} is larger than target size {height}x{width}"
+        assert h % scale == 0 and w % scale == 0
+
+        # Compute the start and end spatial positions of the media item
+        x_start, y_start = conditioning_item.media_x, conditioning_item.media_y
+        x_start = (width - w) // 2 if x_start is None else x_start
+        y_start = (height - h) // 2 if y_start is None else y_start
+        x_end, y_end = x_start + w, y_start + h
+        assert (
+            x_end <= width and y_end <= height
+        ), f"Conditioning item {x_start}:{x_end}x{y_start}:{y_end} is out of bounds for target size {width}x{height}"
+
+        if strip_latent_border:
+            # Strip one latent from left/right and/or top/bottom, update x, y accordingly
+            if x_start > 0:
+                x_start += scale
+                latents = latents[:, :, :, :, 1:]
+
+            if y_start > 0:
+                y_start += scale
+                latents = latents[:, :, :, 1:, :]
+
+            if x_end < width:
+                latents = latents[:, :, :, :, :-1]
+
+            if y_end < height:
+                latents = latents[:, :, :, :-1, :]
+
+        return latents, x_start // scale, y_start // scale
+
+    @staticmethod
+    def _handle_non_first_conditioning_sequence(
+        init_latents: torch.Tensor,
+        init_conditioning_mask: torch.Tensor,
+        latents: torch.Tensor,
+        media_frame_number: int,
+        strength: float,
+        num_prefix_latent_frames: int = 2,
+        prefix_latents_mode: str = "concat",
+        prefix_soft_conditioning_strength: float = 0.15,
+    ):
+        """
+        Special handling for a conditioning sequence that does not start on the first frame.
+        The special handling is required to allow a short encoded video to be used as middle
+        (or last) sequence in a longer video.
+        Args:
+            init_latents (torch.Tensor): The initial noise latents to be updated.
+            init_conditioning_mask (torch.Tensor): The initial conditioning mask to be updated.
+            latents (torch.Tensor): The encoded conditioning item.
+            media_frame_number (int): The target frame number of the first frame in the conditioning sequence.
+            strength (float): The conditioning strength for the conditioning latents.
+            num_prefix_latent_frames (int, optional): The length of the sequence prefix, to be handled
+                separately. Defaults to 2.
+            prefix_latents_mode (str, optional): Special treatment for prefix (boundary) latents.
+                - "drop": Drop the prefix latents.
+                - "soft": Use the prefix latents, but with soft-conditioning
+                - "concat": Add the prefix latents as extra tokens (like single frames)
+            prefix_soft_conditioning_strength (float, optional): The strength of the soft-conditioning for
+                the prefix latents, relevant if `prefix_latents_mode` is "soft". Defaults to 0.1.
+
+        """
+        f_l = latents.shape[2]
+        f_l_p = num_prefix_latent_frames
+        assert f_l >= f_l_p
+        assert media_frame_number % 8 == 0
+        if f_l > f_l_p:
+            # Insert the conditioning latents **excluding the prefix** into the sequence
+            f_l_start = media_frame_number // 8 + f_l_p
+            f_l_end = f_l_start + f_l - f_l_p
+            init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
+                init_latents[:, :, f_l_start:f_l_end],
+                latents[:, :, f_l_p:],
+                strength,
+            )
+            # Mark these latent frames as conditioning latents
+            init_conditioning_mask[:, f_l_start:f_l_end] = strength
+
+        # Handle the prefix-latents
+        if prefix_latents_mode == "soft":
+            if f_l_p > 1:
+                # Drop the first (single-frame) latent and soft-condition the remaining prefix
+                f_l_start = media_frame_number // 8 + 1
+                f_l_end = f_l_start + f_l_p - 1
+                strength = min(prefix_soft_conditioning_strength, strength)
+                init_latents[:, :, f_l_start:f_l_end] = torch.lerp(
+                    init_latents[:, :, f_l_start:f_l_end],
+                    latents[:, :, 1:f_l_p],
+                    strength,
+                )
+                # Mark these latent frames as conditioning latents
+                init_conditioning_mask[:, f_l_start:f_l_end] = strength
+            latents = None  # No more latents to handle
+        elif prefix_latents_mode == "drop":
+            # Drop the prefix latents
+            latents = None
+        elif prefix_latents_mode == "concat":
+            # Pass-on the prefix latents to be handled as extra conditioning frames
+            latents = latents[:, :, :f_l_p]
+        else:
+            raise ValueError(f"Invalid prefix_latents_mode: {prefix_latents_mode}")
+        return (
+            init_latents,
+            init_conditioning_mask,
+            latents,
+        )
+
+    def trim_conditioning_sequence(
+        self, start_frame: int, sequence_num_frames: int, target_num_frames: int
+    ):
+        """
+        Trim a conditioning sequence to the allowed number of frames.
+
+        Args:
+            start_frame (int): The target frame number of the first frame in the sequence.
+            sequence_num_frames (int): The number of frames in the sequence.
+            target_num_frames (int): The target number of frames in the generated video.
+
+        Returns:
+            int: updated sequence length
+        """
+        scale_factor = self.video_scale_factor
+        num_frames = min(sequence_num_frames, target_num_frames - start_frame)
+        # Trim down to a multiple of temporal_scale_factor frames plus 1
+        num_frames = (num_frames - 1) // scale_factor * scale_factor + 1
+        return num_frames
+
+    @staticmethod
+    def tone_map_latents(
+        latents: torch.Tensor,
+        compression: float,
+    ) -> torch.Tensor:
+        """
+        Applies a non-linear tone-mapping function to latent values to reduce their dynamic range
+        in a perceptually smooth way using a sigmoid-based compression.
+
+        This is useful for regularizing high-variance latents or for conditioning outputs
+        during generation, especially when controlling dynamic behavior with a `compression` factor.
+
+        Parameters:
+        ----------
+        latents : torch.Tensor
+            Input latent tensor with arbitrary shape. Expected to be roughly in [-1, 1] or [0, 1] range.
+        compression : float
+            Compression strength in the range [0, 1].
+            - 0.0: No tone-mapping (identity transform)
+            - 1.0: Full compression effect
+
+        Returns:
+        -------
+        torch.Tensor
+            The tone-mapped latent tensor of the same shape as input.
+        """
+        if not (0 <= compression <= 1):
+            raise ValueError("Compression must be in the range [0, 1]")
+
+        # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
+        scale_factor = compression * 0.75
+        abs_latents = torch.abs(latents)
+
+        # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
+        # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
+        sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
+        scales = 1.0 - 0.8 * scale_factor * sigmoid_term
+
+        filtered = latents * scales
+        return filtered
+
+
+def adain_filter_latent(
+    latents: torch.Tensor, reference_latents: torch.Tensor, factor=1.0
+):
+    """
+    Applies Adaptive Instance Normalization (AdaIN) to a latent tensor based on
+    statistics from a reference latent tensor.
+
+    Args:
+        latent (torch.Tensor): Input latents to normalize
+        reference_latent (torch.Tensor): The reference latents providing style statistics.
+        factor (float): Blending factor between original and transformed latent.
+                       Range: -10.0 to 10.0, Default: 1.0
+
+    Returns:
+        torch.Tensor: The transformed latent tensor
+    """
+    result = latents.clone()
+
+    for i in range(latents.size(0)):
+        for c in range(latents.size(1)):
+            r_sd, r_mean = torch.std_mean(
+                reference_latents[i, c], dim=None
+            )  # index by original dim order
+            i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
+
+            result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
+
+    result = torch.lerp(latents, result, factor)
+    return result
+
+
+class LTXMultiScalePipeline:
+    def _upsample_latents(
+        self, latest_upsampler: LatentUpsampler, latents: torch.Tensor
+    ):
+        assert latents.device == latest_upsampler.device
+
+        latents = un_normalize_latents(
+            latents, self.vae, vae_per_channel_normalize=True
+        )
+        upsampled_latents = latest_upsampler(latents)
+        upsampled_latents = normalize_latents(
+            upsampled_latents, self.vae, vae_per_channel_normalize=True
+        )
+        return upsampled_latents
+
+    def __init__(
+        self, video_pipeline: LTXVideoPipeline, latent_upsampler: LatentUpsampler
+    ):
+        self.video_pipeline = video_pipeline
+        self.vae = video_pipeline.vae
+        self.latent_upsampler = latent_upsampler
+
+    def __call__(
+        self,
+        downscale_factor: float,
+        first_pass: dict,
+        second_pass: dict,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        original_kwargs = kwargs.copy()
+        original_output_type = kwargs["output_type"]
+        original_width = kwargs["width"]
+        original_height = kwargs["height"]
+
+        x_width = int(kwargs["width"] * downscale_factor)
+        downscaled_width = x_width - (x_width % self.video_pipeline.vae_scale_factor)
+        x_height = int(kwargs["height"] * downscale_factor)
+        downscaled_height = x_height - (x_height % self.video_pipeline.vae_scale_factor)
+
+        kwargs["output_type"] = "latent"
+        kwargs["width"] = downscaled_width
+        kwargs["height"] = downscaled_height
+        kwargs.update(**first_pass)
+        result = self.video_pipeline(*args, **kwargs)
+        latents = result.images
+
+        upsampled_latents = self._upsample_latents(self.latent_upsampler, latents)
+        upsampled_latents = adain_filter_latent(
+            latents=upsampled_latents, reference_latents=latents
+        )
+
+        kwargs = original_kwargs
+
+        kwargs["latents"] = upsampled_latents
+        kwargs["output_type"] = original_output_type
+        kwargs["width"] = downscaled_width * 2
+        kwargs["height"] = downscaled_height * 2
+        kwargs.update(**second_pass)
+
+        result = self.video_pipeline(*args, **kwargs)
+        if original_output_type != "latent":
+            num_frames = result.images.shape[2]
+            videos = rearrange(result.images, "b c f h w -> (b f) c h w")
+
+            videos = F.interpolate(
+                videos,
+                size=(original_height, original_width),
+                mode="bilinear",
+                align_corners=False,
+            )
+            videos = rearrange(videos, "(b f) c h w -> b c f h w", f=num_frames)
+            result.images = videos
+
+        return result