Update api/ltx/ltx_utils.py

api/ltx/ltx_utils.py

@@ -1,16 +1,18 @@
 # FILE: api/ltx/ltx_utils.py
 # DESCRIPTION: Comprehensive, self-contained utility module for the LTX pipeline.
-# Handles dependency path injection, model loading, pipeline creation, and tensor preparation.
+# Handles dependency path injection, model loading, data structures, and helper functions.
 
 import os
 import random
 import json
-
+import logging
 import time
 import sys
 from pathlib import Path
 from typing import Dict, Optional, Tuple, Union
-from huggingface_hub import hf_hub_download
+from dataclasses import dataclass
+from enum import Enum, auto
+
 import numpy as np
 import torch
 import torchvision.transforms.functional as TVF
@@ -18,50 +20,69 @@ from PIL import Image
 from safetensors import safe_open
 from transformers import T5EncoderModel, T5Tokenizer
 
-import logging
-import warnings
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=FutureWarning)
-warnings.filterwarnings("ignore", message=".*")
-from huggingface_hub import logging as ll
-ll.set_verbosity_error()
-ll.set_verbosity_warning()
-ll.set_verbosity_info()
-from utils.debug_utils import log_function_io
-
-ll.set_verbosity_debug()
-
-
 # ==============================================================================
 # --- CRITICAL: DEPENDENCY PATH INJECTION ---
 # ==============================================================================
 
 # Define o caminho para o repositório clonado
 LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
-LTX_REPO_ID = "Lightricks/LTX-Video"
-CACHE_DIR = os.environ.get("HF_HOME")
+
+def add_deps_to_path():
+    """
+    Adiciona o diretório do repositório LTX ao sys.path para garantir que suas
+    bibliotecas possam ser importadas.
+    """
+    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
+    if repo_path not in sys.path:
+        sys.path.insert(0, repo_path)
+        logging.info(f"[ltx_utils] LTX-Video repository added to sys.path: {repo_path}")
+
+# Executa a função imediatamente para configurar o ambiente antes de qualquer importação.
+add_deps_to_path()
+
 
 # ==============================================================================
 # --- IMPORTAÇÕES DA BIBLIOTECA LTX-VIDEO (Após configuração do path) ---
 # ==============================================================================
+try:
+    from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
+    from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
+    from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
+    from ltx_video.models.transformers.transformer3d import Transformer3DModel
+    from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
+    from ltx_video.schedulers.rf import RectifiedFlowScheduler
+    from ltx_video.models.autoencoders.vae_encode import un_normalize_latents, normalize_latents
+    import ltx_video.pipelines.crf_compressor as crf_compressor
+except ImportError as e:
+    raise ImportError(f"Could not import from LTX-Video library even after setting sys.path. Check repo integrity at '{LTX_VIDEO_REPO_DIR}'. Error: {e}")
 
-repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
-if repo_path not in sys.path:
-    sys.path.insert(0, repo_path)
 
-from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
-from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
-from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
-from ltx_video.models.transformers.transformer3d import Transformer3DModel
-from ltx_video.models.transformers.symmetric_patchifier import SymmetricPatchifier
-from ltx_video.schedulers.rf import RectifiedFlowScheduler
-import ltx_video.pipelines.crf_compressor as crf_compressor
+# ==============================================================================
+# --- ESTRUTURAS DE DADOS E ENUMS (Centralizadas aqui) ---
+# ==============================================================================
+
+@dataclass
+class ConditioningItem:
+    """Define a single frame-conditioning item, used to guide the generation pipeline."""
+    media_item: torch.Tensor
+    media_frame_number: int
+    conditioning_strength: float
+    media_x: Optional[int] = None
+    media_y: Optional[int] = None
+
+
+class SkipLayerStrategy(Enum):
+    """Defines the strategy for how spatio-temporal guidance is applied across transformer blocks."""
+    AttentionSkip = auto()
+    AttentionValues = auto()
+    Residual = auto()
+    TransformerBlock = auto()
+
 
 # ==============================================================================
 # --- FUNÇÕES DE CONSTRUÇÃO DE MODELO E PIPELINE ---
 # ==============================================================================
 
-@log_function_io
 def create_latent_upsampler(latent_upsampler_model_path: str, device: str) -> LatentUpsampler:
     """Loads the Latent Upsampler model from a checkpoint path."""
     logging.info(f"Loading Latent Upsampler from: {latent_upsampler_model_path} to device: {device}")
@@ -70,17 +91,15 @@ def create_latent_upsampler(latent_upsampler_model_path: str, device: str) -> La
     latent_upsampler.eval()
     return latent_upsampler
 
-@log_function_io
 def build_ltx_pipeline_on_cpu(config: Dict) -> Tuple[LTXVideoPipeline, Optional[torch.nn.Module]]:
     """Builds the complete LTX pipeline and upsampler on the CPU."""
     t0 = time.perf_counter()
     logging.info("Building LTX pipeline on CPU...")
 
-    ckpt_path_str = hf_hub_download(repo_id=LTX_REPO_ID, filename=config["checkpoint_path"], cache_dir=CACHE_DIR)
-    ckpt_path = Path(ckpt_path_str)
+    ckpt_path = Path(config["checkpoint_path"])
     if not ckpt_path.is_file():
         raise FileNotFoundError(f"Main checkpoint file not found: {ckpt_path}")
-  
+
     with safe_open(ckpt_path, framework="pt") as f:
         metadata = f.metadata() or {}
         config_str = metadata.get("config", "{}")
@@ -111,11 +130,8 @@ def build_ltx_pipeline_on_cpu(config: Dict) -> Tuple[LTXVideoPipeline, Optional[
     )
 
     latent_upsampler = None
-    if config.get("spatial_upscaler_model_path"):  
-        spatial_path_str = hf_hub_download(repo_id=LTX_REPO_ID, filename=config["spatial_upscaler_model_path"], cache_dir=CACHE_DIR)
-        spatial_path = Path(spatial_path_str)
-        if not ckpt_path.is_file():
-            raise FileNotFoundError(f"Main checkpoint file not found: {spatial_path}")
+    if config.get("spatial_upscaler_model_path"):
+        spatial_path = config["spatial_upscaler_model_path"]
         latent_upsampler = create_latent_upsampler(spatial_path, device="cpu")
         if precision == "bfloat16":
             latent_upsampler.to(torch.bfloat16)
@@ -123,11 +139,24 @@ def build_ltx_pipeline_on_cpu(config: Dict) -> Tuple[LTXVideoPipeline, Optional[
     logging.info(f"LTX pipeline built on CPU in {time.perf_counter() - t0:.2f}s")
     return pipeline, latent_upsampler
 
+
 # ==============================================================================
-# --- FUNÇÕES AUXILIARES (Seed, Preparação de Imagem) ---
+# --- FUNÇÕES AUXILIARES (Latent Processing, Seed, Image Prep) ---
 # ==============================================================================
 
-@log_function_io
+def adain_filter_latent(
+    latents: torch.Tensor, reference_latents: torch.Tensor, factor=1.0
+) -> torch.Tensor:
+    """Applies AdaIN to transfer the style from a reference latent to another."""
+    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)
+            i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
+            if i_sd > 1e-6:
+                result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
+    return torch.lerp(latents, result, factor)
+
 def seed_everything(seed: int):
     """Sets the seed for reproducibility."""
     random.seed(seed)
@@ -138,13 +167,12 @@ def seed_everything(seed: int):
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
-@log_function_io
 def load_image_to_tensor_with_resize_and_crop(
     image_input: Union[str, Image.Image],
     target_height: int,
     target_width: int,
 ) -> torch.Tensor:
-    """Loads and processes an image into a 5D pixel tensor compatible with the LTX pipeline."""
+    """Loads and processes an image into a 5D tensor compatible with the LTX pipeline."""
     if isinstance(image_input, str):
         image = Image.open(image_input).convert("RGB")
     elif isinstance(image_input, Image.Image):
@@ -166,15 +194,14 @@ def load_image_to_tensor_with_resize_and_crop(
     image = image.crop((x_start, y_start, x_start + new_width, y_start + new_height))
     image = image.resize((target_width, target_height), Image.Resampling.LANCZOS)
 
-    frame_tensor = TVF.to_tensor(image)  # PIL -> tensor (C, H, W) in [0, 1] range
+    frame_tensor = TVF.to_tensor(image)
     frame_tensor = TVF.gaussian_blur(frame_tensor, kernel_size=(3, 3))
     
     frame_tensor_hwc = frame_tensor.permute(1, 2, 0)
     frame_tensor_hwc = crf_compressor.compress(frame_tensor_hwc)
     frame_tensor = frame_tensor_hwc.permute(2, 0, 1)
-    # Normalize to [-1, 1] range, which the VAE expects for encoding
+    # Normalize to [-1, 1] range
     frame_tensor = (frame_tensor * 2.0) - 1.0
     
     # Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
-    return frame_tensor.unsqueeze(0).unsqueeze(2)
-    
+    return frame_tensor.unsqueeze(0).unsqueeze(2)