Delete api/aduc_ltx_latent_patch.py

api/aduc_ltx_latent_patch.py

@@ -1,268 +0,0 @@
-# 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
-
-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)
-        if LTXV_DEBUG:
-            print(f"[DEBUG] Adicionado ao sys.path: {resolved_path}")
-
-# --- Execução da configuração inicial ---
-if not LTX_VIDEO_REPO_DIR.exists():
-    _run_setup_script()
-add_deps_to_path(LTX_VIDEO_REPO_DIR)
-
-
-# Tenta importar as dependências necessárias do módulo original que será modificado.
-# Isso requer que o ambiente Python tenha o pacote `ltx_video` acessível em seu sys.path.
-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}")
-    # Interrompe a execução se as dependências essenciais não puderem ser encontradas.
-    raise
-
-print("[INFO] Patch module 'aduc_ltx_latent_patch' loaded successfully.")
-
-# ==============================================================================
-# 1. NOVA DEFINIÇÃO DA DATACLASS `ConditioningItem`
-# ==============================================================================
-
-from dataclasses import dataclass
-
-@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."
-
-    # Se não há itens de condicionamento, apenas patchifica os latentes e retorna.
-    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
-
-    # Inicializa tensores para acumular resultados
-    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
-
-        # --- LÓGICA CENTRAL DO PATCH ---
-        if item.latents is not None:
-            # 1. Se latentes pré-calculados existem, use-os diretamente.
-            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:
-            # 2. Caso contrário, volte para o fluxo original de codificação da VAE.
-            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:
-            # Este caso é prevenido pelo __post_init__ do dataclass, mas é bom ter uma checagem.
-            raise ValueError("ConditioningItem is invalid: it has neither 'latents' nor 'media_item'.")
-        # --- FIM DA LÓGICA DO PATCH ---
-
-        media_frame_number = item.media_frame_number
-        strength = item.conditioning_strength
-        
-        # O resto da lógica da função original é aplicado sobre `item_latents`.
-        if media_frame_number == 0:
-            item_latents, l_x, l_y = self._get_latent_spatial_position(
-                item_latents, item, height, width, strip_latent_border=True
-            )
-            _, _, 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)
-
-    # Patchifica os latentes principais e a máscara de condicionamento
-    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)
-
-    # Concatena os latentes extras (se houver)
-    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.
-
-    Esta classe substitui o método `prepare_conditioning` da `LTXVideoPipeline`
-    pela versão otimizada que suporta latentes pré-calculados, e implicitamente
-    requer o uso da `PatchedConditioningItem`.
-    """
-    _original_prepare_conditioning = None
-    _is_patched = False
-
-    @staticmethod
-    def apply():
-        """
-        Aplica o monkey patch à classe `LTXVideoPipeline`.
-
-        Guarda o método original e o substitui pela nova implementação.
-        É idempotente; aplicar múltiplas vezes não causa efeito adicional.
-        """
-        if LTXLatentConditioningPatch._is_patched:
-            print("[WARNING] LTXLatentConditioningPatch has already been applied. Ignoring.")
-            return
-
-        print("[INFO] Applying monkey patch for latent-based conditioning...")
-
-        # Guarda a implementação original para permitir a reversão.
-        LTXLatentConditioningPatch._original_prepare_conditioning = LTXVideoPipeline.prepare_conditioning
-
-        # Substitui o método na classe LTXVideoPipeline.
-        # Todas as instâncias futuras e existentes da classe usarão este novo método.
-        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.")