vace/models/ltx/pipelines/pipeline_ltx_video.py

# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import html
import inspect
from dataclasses import dataclass

import math
import re
import urllib.parse as ul
import numpy as np
import PIL
from typing import Callable, Dict, List, Optional, Tuple, Union


import torch
from torch import Tensor
import torch.nn.functional as F
from contextlib import nullcontext
from diffusers.image_processor import VaeImageProcessor
from diffusers.models import AutoencoderKL
from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput, BaseOutput
from diffusers.schedulers import DPMSolverMultistepScheduler
from diffusers.utils import (
    BACKENDS_MAPPING,
    deprecate,
    is_bs4_available,
    is_ftfy_available,
    logging,
)
from diffusers.utils.torch_utils import randn_tensor
from einops import rearrange
from transformers import T5EncoderModel, T5Tokenizer

from ltx_video.models.transformers.transformer3d import Transformer3DModel
from ltx_video.models.transformers.symmetric_patchifier import Patchifier
from ltx_video.models.autoencoders.vae_encode import (
    get_vae_size_scale_factor,
    vae_decode,
    vae_encode,
)
from ltx_video.models.autoencoders.causal_video_autoencoder import (
    CausalVideoAutoencoder,
)
from ltx_video.schedulers.rf import TimestepShifter
from ltx_video.utils.conditioning_method import ConditioningMethod
from ltx_video.utils.skip_layer_strategy import SkipLayerStrategy

from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline, retrieve_timesteps

from ...utils.preprocessor import prepare_source



@dataclass
class ImagePipelineOutput(BaseOutput):
    images: Union[List[PIL.Image.Image], np.ndarray]
    info: Optional[Dict] = None


class VaceLTXVideoPipeline(LTXVideoPipeline):

    @torch.no_grad()
    def __call__(
            self,
            src_video: torch.FloatTensor = None,
            src_mask: torch.FloatTensor = None,
            src_ref_images: List[torch.FloatTensor] = None,
            height: int = 512,
            width: int = 768,
            num_frames: int = 97,
            frame_rate: float = 25,
            prompt: Union[str, List[str]] = None,
            negative_prompt: str = "",
            num_inference_steps: int = 20,
            timesteps: List[int] = None,
            guidance_scale: float = 4.5,
            context_scale: float = 1.0,
            skip_layer_strategy: Optional[SkipLayerStrategy] = None,
            skip_block_list: List[int] = None,
            stg_scale: float = 1.0,
            do_rescaling: bool = True,
            rescaling_scale: float = 0.7,
            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,
            clean_caption: bool = True,
            media_items: Optional[torch.FloatTensor] = 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,
            decouple_with_mask: bool = True,
            use_mask: bool = True,
            decode_all_frames: bool = False,
            mask_downsample: [list] = [2, 8, 8],
            **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.
            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.
            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 or not 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`.
            clean_caption (`bool`, *optional*, defaults to `True`):
                Whether or not to clean the caption before creating embeddings. Requires `beautifulsoup4` and `ftfy` to
                be installed. If the dependencies are not installed, the embeddings will be created from the raw
                prompt.
            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.

        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

        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
        # corresponds to doing no classifier free guidance.
        do_classifier_free_guidance = guidance_scale > 1.0
        do_spatio_temporal_guidance = stg_scale > 0.0

        num_conds = 1
        if do_classifier_free_guidance:
            num_conds += 1
        if do_spatio_temporal_guidance:
            num_conds += 1

        skip_layer_mask = None
        if do_spatio_temporal_guidance:
            skip_layer_mask = self.transformer.create_skip_layer_mask(
                skip_block_list, batch_size, num_conds, 2
            )

        # 3. Encode input prompt
        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,
            do_classifier_free_guidance,
            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,
            clean_caption=clean_caption,
        )

        if offload_to_cpu:
            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
        if do_classifier_free_guidance:
            prompt_embeds_batch = torch.cat(
                [negative_prompt_embeds, prompt_embeds], dim=0
            )
            prompt_attention_mask_batch = torch.cat(
                [negative_prompt_attention_mask, prompt_attention_mask], dim=0
            )
        if do_spatio_temporal_guidance:
            prompt_embeds_batch = torch.cat([prompt_embeds_batch, prompt_embeds], dim=0)
            prompt_attention_mask_batch = torch.cat(
                [
                    prompt_attention_mask_batch,
                    prompt_attention_mask,
                ],
                dim=0,
            )

        # 3b. Encode and prepare conditioning data
        self.video_scale_factor = self.video_scale_factor if is_video else 1
        conditioning_method = kwargs.get("conditioning_method", None)
        vae_per_channel_normalize = kwargs.get("vae_per_channel_normalize", False)
        image_cond_noise_scale = kwargs.get("image_cond_noise_scale", 0.0)
        init_latents, conditioning_mask = self.prepare_conditioning(
            media_items,
            num_frames,
            height,
            width,
            conditioning_method,
            vae_per_channel_normalize,
        )

        #------------------------ VACE Part ------------------------#
        # 4. Prepare latents.
        image_size = (height, width)
        src_ref_images = [None] * batch_size if src_ref_images is None else src_ref_images
        source_ref_len = max([len(ref_imgs) if ref_imgs is not None else 0 for ref_imgs in src_ref_images])
        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_frame_rate = frame_rate / self.video_scale_factor
        num_latent_patches = latent_height * latent_width * (latent_num_frames + source_ref_len)
        latents = self.prepare_latents(
            batch_size=batch_size * num_images_per_prompt,
            num_latent_channels=self.transformer.config.in_channels,
            num_patches=num_latent_patches,
            dtype=prompt_embeds_batch.dtype,
            device=device,
            generator=generator,
            latents=init_latents,
            latents_mask=conditioning_mask,
        )
        src_video, src_mask, src_ref_images = prepare_source(src_video, src_mask, src_ref_images, num_frames, image_size, latents.device)

        # Prepare source_latents
        if decouple_with_mask:
            unchanged = [i * (1 - m) + 0 * m for i, m in zip(src_video, src_mask)]
            changed = [i * m + 0 * (1 - m) for i, m in zip(src_video, src_mask)]
            unchanged = torch.stack(unchanged, dim=0).to(self.vae.dtype).to(device) if isinstance(unchanged, list) else unchanged.to(self.vae.dtype).to(device)  # [B, C, F, H, W]
            changed = torch.stack(changed, dim=0).to(self.vae.dtype).to(device) if isinstance(changed, list) else changed.to(self.vae.dtype).to(device)  # [B, C, F, H, W]
            unchanged_latents = vae_encode(unchanged, vae=self.vae, vae_per_channel_normalize=vae_per_channel_normalize).float()
            changed_latents = vae_encode(changed, vae=self.vae, vae_per_channel_normalize=vae_per_channel_normalize).float()
            source_latents = torch.stack([torch.cat((u, c), dim=0) for u, c in zip(unchanged_latents, changed_latents)], dim=0)
        else:
            src_video = torch.stack(src_video, dim=0).to(self.vae.dtype).to(device) if isinstance(src_video, list) else src_video.to(self.vae.dtype).to(device)  # [B, C, F, H, W]
            source_latents = vae_encode(src_video, vae=self.vae, vae_per_channel_normalize=vae_per_channel_normalize).float()

        # Prepare source_ref_latents
        use_ref = all(ref_imgs is not None and len(ref_imgs) > 0 for ref_imgs in src_ref_images)
        if use_ref:
            source_ref_latents = []
            for i, ref_imgs in enumerate(src_ref_images):
                # [(C=3, F=1, H, W), ...]  ->  (N_REF, C'=128, F=1, H', W')
                ref_imgs = torch.stack(ref_imgs, dim=0).to(self.vae.dtype).to(device) if isinstance(ref_imgs, list) else ref_imgs.to(self.vae.dtype).to(device)  # [B, C, F, H, W]
                ref_latents = vae_encode(ref_imgs, vae=self.vae, vae_per_channel_normalize=vae_per_channel_normalize).float()
                # (N_REF, C'=128, F=1, H', W') -> (1, C'=128, N_REF, H', W')
                ref_latents = ref_latents.permute(2, 1, 0, 3, 4)
                if decouple_with_mask:
                    ref_latents = torch.cat([ref_latents, torch.zeros_like(ref_latents)], dim=1)  # [unchanged, changed]
                source_ref_latents.append(ref_latents)
            # (B, C'=128, N_REF, H', W')
            source_ref_latents = torch.cat(source_ref_latents, dim=0)
        else:
            source_ref_latents = None

        # Prepare source_latents
        if source_ref_latents is not None:
            source_latents = torch.cat([source_ref_latents, source_latents], dim=2)
        source_latents = self.patchifier.patchify(latents=source_latents).to(self.transformer.dtype).to(device)

        # Prepare source_mask_latents
        if use_mask and src_mask is not None:
            source_mask_latents = []
            for submask in src_mask:
                submask = F.interpolate(submask.unsqueeze(0),
                                        size=(latent_num_frames * mask_downsample[0],
                                              latent_height * mask_downsample[1],
                                              latent_width * mask_downsample[2]),
                                        mode='trilinear', align_corners=True)
                submask = rearrange(submask, "b c (f p1) (h p2) (w p3) -> b (c p1 p2 p3) f h w", p1=mask_downsample[0], p2=mask_downsample[1], p3=mask_downsample[2]).to(device)
                if source_ref_latents is not None:
                    if decouple_with_mask:
                        submask = torch.cat([torch.zeros_like(source_ref_latents[:, :latents.shape[-1], :]), submask], dim=2)
                    else:
                        submask = torch.cat([torch.zeros_like(source_ref_latents), submask], dim=2)
                submask = self.patchifier.patchify(submask)
                source_mask_latents.append(submask)
            source_mask_latents = torch.cat(source_mask_latents, dim=0).to(self.transformer.dtype).to(device)
        else:
            source_mask_latents = None
        #------------------------ VACE Part ------------------------#

        orig_conditiong_mask = 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)
                if num_conds > 1
                else conditioning_mask
            )

        # 5. Prepare timesteps
        retrieve_timesteps_kwargs = {}
        if isinstance(self.scheduler, TimestepShifter):
            retrieve_timesteps_kwargs["samples"] = latents
        timesteps, num_inference_steps = retrieve_timesteps(
            self.scheduler,
            num_inference_steps,
            device,
            timesteps,
            **retrieve_timesteps_kwargs,
        )

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

        with self.progress_bar(total=num_inference_steps) as progress_bar:
            for i, t in enumerate(timesteps):
                if conditioning_method == ConditioningMethod.FIRST_FRAME:
                    latents = self.image_cond_noise_update(
                        t,
                        init_latents,
                        latents,
                        image_cond_noise_scale,
                        orig_conditiong_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
                )

                latent_frame_rates = (
                        torch.ones(
                            latent_model_input.shape[0], 1, device=latent_model_input.device
                        )
                        * latent_frame_rate
                )

                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)
                scale_grid = (
                    (
                        1 / latent_frame_rates,
                        self.vae_scale_factor,
                        self.vae_scale_factor,
                    )
                    if self.transformer.use_rope
                    else None
                )
                indices_grid = self.patchifier.get_grid(
                    orig_num_frames=latent_num_frames + source_ref_len,
                    orig_height=latent_height,
                    orig_width=latent_width,
                    batch_size=latent_model_input.shape[0],
                    scale_grid=scale_grid,
                    device=latents.device,
                )

                if conditioning_mask is not None:
                    current_timestep = current_timestep * (1 - conditioning_mask)
                # Choose the appropriate context manager based on `mixed_precision`
                if mixed_precision:
                    if "xla" in device.type:
                        raise NotImplementedError(
                            "Mixed precision is not supported yet on XLA devices."
                        )

                    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,
                        source_latents=source_latents,
                        source_mask_latents=source_mask_latents if use_mask else None,
                        context_scale=context_scale,
                        encoder_hidden_states=prompt_embeds_batch.to(
                            self.transformer.dtype
                        ),
                        encoder_attention_mask=prompt_attention_mask_batch,
                        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_perturb = noise_pred[-1:]
                if do_classifier_free_guidance:
                    noise_pred_uncond, noise_pred_text = noise_pred[:2].chunk(2)
                    noise_pred = noise_pred_uncond + guidance_scale * (
                            noise_pred_text - noise_pred_uncond
                    )
                if do_spatio_temporal_guidance:
                    noise_pred = noise_pred + stg_scale * (
                            noise_pred_text - noise_pred_text_perturb
                    )
                    if do_rescaling:
                        factor = noise_pred_text.std() / noise_pred.std()
                        factor = rescaling_scale * factor + (1 - rescaling_scale)
                        noise_pred = noise_pred * factor

                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.scheduler.step(
                    noise_pred,
                    t if current_timestep is None else current_timestep,
                    latents,
                    **extra_step_kwargs,
                    return_dict=False,
                )[0]

                # 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()

        latents = self.patchifier.unpatchify(
            latents=latents,
            output_height=latent_height,
            output_width=latent_width,
            output_num_frames=latent_num_frames + source_ref_len,
            out_channels=self.transformer.config.in_channels
                         // math.prod(self.patchifier.patch_size),
        )

        if not decode_all_frames:
            latents = latents[:, :, source_ref_len:]

        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
            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()

        info = {
            "height": height,
            "width": width,
            "num_frames": num_frames,
            "frame_rate": frame_rate
        }

        if not return_dict:
            return (image,)

        return ImagePipelineOutput(images=image, info=info)