api/aduc_ltx_latent_patch.py

# aduc_ltx_latent_patch.py
#
# Este módulo fornece um monkey patch para a classe LTXVideoPipeline da biblioteca ltx_video.
# A principal funcionalidade deste patch é otimizar o processo de condicionamento, permitindo
# que a pipeline aceite tensores de latentes pré-calculados diretamente através de um
# `ConditioningItem` modificado. Isso evita a re-codificação desnecessária de mídias (imagens/vídeos)
# pela VAE, resultando em um ganho de performance significativo quando os latentes já estão disponíveis.

import torch
from torch import Tensor
from typing import Optional, List, Tuple
from pathlib import Path
import os
import sys
from dataclasses import dataclass, replace

# --- CONFIGURAÇÃO DE PATH (Assume que LTXV_DEBUG e _run_setup_script existem no escopo que carrega este módulo) ---
# DEPS_DIR = Path("/data")
# LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
# def add_deps_to_path(repo_path: Path):
#     """Adiciona o diretório do repositório ao sys.path para importações locais."""
#     resolved_path = str(repo_path.resolve())
#     if resolved_path not in sys.path:
#         sys.path.insert(0, resolved_path)
# add_deps_to_path(LTX_VIDEO_REPO_DIR)


# Tenta importar as dependências necessárias do módulo original que será modificado.
try:
    from ltx_video.pipelines.pipeline_ltx_video import (
        LTXVideoPipeline,
        ConditioningItem as OriginalConditioningItem
    )
    from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
    from ltx_video.models.autoencoders.vae_encode import vae_encode, latent_to_pixel_coords
    from diffusers.utils.torch_utils import randn_tensor
except ImportError as e:
    print(f"FATAL ERROR: Could not import dependencies from 'ltx_video'. "
          f"Please ensure the environment is correctly set up. Error: {e}")
    raise

print("[INFO] Patch module 'aduc_ltx_latent_patch' loaded successfully.")

# ==============================================================================
# 1. NOVA DEFINIÇÃO DA DATACLASS `PatchedConditioningItem`
# ==============================================================================

@dataclass
class PatchedConditioningItem:
    """
    Versão modificada do `ConditioningItem` que aceita tensores de pixel (`media_item`)
    ou tensores de latentes pré-codificados (`latents`).

    Attributes:
        media_frame_number (int): Quadro inicial do item de condicionamento no vídeo.
        conditioning_strength (float): Força do condicionamento (0.0 a 1.0).
        media_item (Optional[Tensor]): Tensor de mídia (pixels). Usado se `latents` for None.
        media_x (Optional[int]): Coordenada X (esquerda) para posicionamento espacial.
        media_y (Optional[int]): Coordenada Y (topo) para posicionamento espacial.
        latents (Optional[Tensor]): Tensor de latentes pré-codificado. Terá precedência sobre `media_item`.
    """
    media_frame_number: int
    conditioning_strength: float
    media_item: Optional[Tensor] = None
    media_x: Optional[int] = None
    media_y: Optional[int] = None
    latents: Optional[Tensor] = None

    def __post_init__(self):
        """Valida o estado do objeto após a inicialização."""
        if self.media_item is None and self.latents is None:
            raise ValueError("A `PatchedConditioningItem` must have either 'media_item' or 'latents' defined.")
        if self.media_item is not None and self.latents is not None:
            print("[WARNING] `PatchedConditioningItem` received both 'media_item' and 'latents'. "
                  "The 'latents' tensor will take precedence.")

# ==============================================================================
# 2. NOVA IMPLEMENTAÇÃO DA FUNÇÃO `prepare_conditioning`
# ==============================================================================

def prepare_conditioning_with_latents(
    self: LTXVideoPipeline,
    conditioning_items: Optional[List[PatchedConditioningItem]],
    init_latents: Tensor,
    num_frames: int,
    height: int,
    width: int,
    vae_per_channel_normalize: bool = False,
    generator: Optional[torch.Generator] = None,
) -> Tuple[Tensor, Tensor, Optional[Tensor], int]:
    """
    Versão modificada de `prepare_conditioning` que prioriza o uso de latentes pré-calculados
    dos `conditioning_items`, evitando a re-codificação desnecessária pela VAE.
    """
    assert isinstance(self, LTXVideoPipeline), "This function must be called as a method of LTXVideoPipeline."
    assert isinstance(self.vae, CausalVideoAutoencoder), "VAE must be of type CausalVideoAutoencoder."

    if not conditioning_items:
        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
        )
        return init_latents, init_pixel_coords, None, 0

    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

    for item in conditioning_items:
        item_latents: Tensor

        if item.latents is not None:
            item_latents = item.latents.to(dtype=init_latents.dtype, device=init_latents.device)
            if item_latents.ndim != 5:
                raise ValueError(f"Latents must have 5 dimensions (b, c, f, h, w), but got {item_latents.ndim}")
        elif item.media_item is not None:
            resized_item = self._resize_conditioning_item(item, height, width)
            media_item = resized_item.media_item
            assert media_item.ndim == 5, f"media_item must have 5 dims, but got {media_item.ndim}"
            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)
        else:
            raise ValueError("ConditioningItem is invalid: it has neither 'latents' nor 'media_item'.")

        media_frame_number = item.media_frame_number
        strength = item.conditioning_strength
        
        if media_frame_number == 0:
            # --- INÍCIO DA MODIFICAÇÃO ---
            # Se `item.media_item` for None (nosso caso de uso otimizado), a função original `_get_latent_spatial_position`
            # quebraria. Para evitar isso, criamos um item temporário com um tensor de placeholder que contém
            # as informações de dimensão corretas, inferidas a partir dos próprios latentes.
            
            item_for_spatial_position = item
            if item.media_item is None:
                # Infere as dimensões em pixels a partir da forma dos latentes
                latent_h, latent_w = item_latents.shape[-2:]
                pixel_h = latent_h * self.vae_scale_factor
                pixel_w = latent_w * self.vae_scale_factor
                
                # Cria um tensor de placeholder com o shape esperado (o conteúdo não importa)
                placeholder_media_item = torch.empty(
                    (1, 1, 1, pixel_h, pixel_w), device=item_latents.device, dtype=item_latents.dtype
                )
                
                # Usa `dataclasses.replace` para criar uma cópia temporária do item com o placeholder
                item_for_spatial_position = replace(item, media_item=placeholder_media_item)

            # Chama a função original com um item que ela pode processar sem erro
            item_latents, l_x, l_y = self._get_latent_spatial_position(
                item_latents, item_for_spatial_position, height, width, strip_latent_border=True
            )
            # --- FIM DA MODIFICAÇÃO ---
            
            _, _, f_l, h_l, w_l = item_latents.shape
            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], item_latents, strength
            )
            init_conditioning_mask[:, :f_l, l_y : l_y + h_l, l_x : l_x + w_l] = strength
        else:
            if item_latents.shape[2] > 1:
                (init_latents, init_conditioning_mask, item_latents) = self._handle_non_first_conditioning_sequence(
                    init_latents, init_conditioning_mask, item_latents, media_frame_number, strength
                )

            if item_latents is not None:
                noise = randn_tensor(
                    item_latents.shape, generator=generator,
                    device=item_latents.device, dtype=item_latents.dtype
                )
                item_latents = torch.lerp(noise, item_latents, strength)
                item_latents, latent_coords = self.patchifier.patchify(latents=item_latents)
                pixel_coords = latent_to_pixel_coords(
                    latent_coords, self.vae,
                    causal_fix=self.transformer.config.causal_temporal_positioning
                )
                pixel_coords[:, 0] += media_frame_number
                extra_conditioning_num_latents += item_latents.shape[1]
                conditioning_mask = torch.full(
                    item_latents.shape[:2], strength,
                    dtype=torch.float32, device=init_latents.device
                )
                extra_conditioning_latents.append(item_latents)
                extra_conditioning_pixel_coords.append(pixel_coords)
                extra_conditioning_mask.append(conditioning_mask)

    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
    )
    init_conditioning_mask, _ = self.patchifier.patchify(latents=init_conditioning_mask.unsqueeze(1))
    init_conditioning_mask = init_conditioning_mask.squeeze(-1)

    if extra_conditioning_latents:
        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:
            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

# ==============================================================================
# 3. CLASSE DO MONKEY PATCHER
# ==============================================================================

class LTXLatentConditioningPatch:
    """
    Classe estática para aplicar e reverter o monkey patch na pipeline LTX-Video.
    """
    _original_prepare_conditioning = None
    _is_patched = False

    @staticmethod
    def apply():
        """
        Aplica o monkey patch à classe `LTXVideoPipeline`.
        """
        if LTXLatentConditioningPatch._is_patched:
            print("[WARNING] LTXLatentConditioningPatch has already been applied. Ignoring.")
            return

        print("[INFO] Applying monkey patch for latent-based conditioning...")
        
        LTXLatentConditioningPatch._original_prepare_conditioning = LTXVideoPipeline.prepare_conditioning
        LTXVideoPipeline.prepare_conditioning = prepare_conditioning_with_latents

        LTXLatentConditioningPatch._is_patched = True
        print("[SUCCESS] Monkey patch applied successfully.")
        print("  - `LTXVideoPipeline.prepare_conditioning` has been updated.")
        print("  - NOTE: Remember to use `aduc_ltx_latent_patch.PatchedConditioningItem` when creating conditioning items.")

    @staticmethod
    def revert():
        """
        Reverte o monkey patch, restaurando a implementação original.
        """
        if not LTXLatentConditioningPatch._is_patched:
            print("[WARNING] Patch is not currently applied. No action taken.")
            return

        if LTXLatentConditioningPatch._original_prepare_conditioning:
            print("[INFO] Reverting LTXLatentConditioningPatch...")
            LTXVideoPipeline.prepare_conditioning = LTXLatentConditioningPatch._original_prepare_conditioning
            LTXLatentConditioningPatch._is_patched = False
            print("[SUCCESS] Patch reverted successfully. Original functionality restored.")
        else:
            print("[ERROR] Cannot revert: original implementation was not saved.")