Update api/ltx/ltx_utils.py

api/ltx/ltx_utils.py

@@ -1,6 +1,6 @@
 # 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.
+# DESCRIPTION: A pure utility library for the LTX ecosystem.
+# Contains low-level, stateless builders for core components and general-purpose helper functions.
 
 import os
 import random
@@ -10,42 +10,30 @@ import time
 import sys
 from pathlib import Path
 from typing import Dict, Optional, Tuple, Union
-from huggingface_hub import hf_hub_download
 
+from huggingface_hub import hf_hub_download
 import numpy as np
 import torch
 import torchvision.transforms.functional as TVF
 from PIL import Image
-from safetensors import safe_open
+from safensors import safe_open
 from transformers import T5EncoderModel, T5Tokenizer
 
 # ==============================================================================
-# --- CRITICAL: DEPENDENCY PATH INJECTION ---
+# --- CONFIGURAÇÃO DE PATH E IMPORTS DA BIBLIOTECA LTX ---
 # ==============================================================================
 
-# 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.
-    """
+    """Adiciona o diretório do repositório LTX ao sys.path para importação de suas bibliotecas."""
     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
@@ -55,93 +43,65 @@ try:
     from ltx_video.schedulers.rf import RectifiedFlowScheduler
     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}")
-
+    raise ImportError(f"Could not import from LTX-Video library. Check repo integrity at '{LTX_VIDEO_REPO_DIR}'. Error: {e}")
 
 # ==============================================================================
-# --- FUNÇÕES DE CONSTRUÇÃO DE MODELO E PIPELINE ---
+# --- CAMADA DE BUILDERS DE BAIXO NÍVEL ---
+# (Responsáveis por construir um único componente na CPU)
 # ==============================================================================
 
-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}")
-    latent_upsampler = LatentUpsampler.from_pretrained(latent_upsampler_model_path)
-    latent_upsampler.to(device)
-    latent_upsampler.eval()
-    return latent_upsampler
-
-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...")
-
-
+def _build_vae(checkpoint_path: str, precision: str) -> CausalVideoAutoencoder:
+    """Constrói o CausalVideoAutoencoder a partir de um checkpoint, sempre na CPU."""
+    logging.info(f"Building VAE from checkpoint: {Path(checkpoint_path).name}")
+    vae = CausalVideoAutoencoder.from_pretrained(checkpoint_path).to("cpu")
+    if precision == "bfloat16":
+        vae.to(torch.bfloat16)
+    return vae
 
-    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)
-    if not ckpt_path.is_file():
-        raise FileNotFoundError(f"Main checkpoint file not found: {ckpt_path}")
+def _build_latent_upscaler(upscaler_path: str, precision: str) -> LatentUpsampler:
+    """Constrói o LatentUpsampler a partir de um checkpoint, sempre na CPU."""
+    logging.info(f"Building Latent Upscaler from: {Path(upscaler_path).name}")
+    upscaler = LatentUpsampler.from_pretrained(upscaler_path).to("cpu")
+    if precision == "bfloat16":
+        upscaler.to(torch.bfloat16)
+    return upscaler
 
-    logging.info(f"Building LTX pipeline ckpt:{ckpt_path_str}")
-    
+def _build_ltx_transformer_pipeline(checkpoint_path: str, config: Dict, precision: str) -> LTXVideoPipeline:
+    """Constrói o LTXVideoPipeline principal (sem VAE), sempre na CPU."""
+    logging.info(f"Building core LTX Transformer Pipeline from checkpoint: {Path(checkpoint_path).name}")
     
-    with safe_open(ckpt_path, framework="pt") as f:
+    with safe_open(checkpoint_path, framework="pt") as f:
         metadata = f.metadata() or {}
         config_str = metadata.get("config", "{}")
         configs = json.loads(config_str)
         allowed_inference_steps = configs.get("allowed_inference_steps")
-
-        
-    vae = CausalVideoAutoencoder.from_pretrained(ckpt_path).to("cpu")
-    transformer = Transformer3DModel.from_pretrained(ckpt_path).to("cpu")
-    scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)
     
+    transformer = Transformer3DModel.from_pretrained(checkpoint_path).to("cpu")
+    scheduler = RectifiedFlowScheduler.from_pretrained(checkpoint_path)
     text_encoder_path = config["text_encoder_model_name_or_path"]
     text_encoder = T5EncoderModel.from_pretrained(text_encoder_path, subfolder="text_encoder").to("cpu")
     tokenizer = T5Tokenizer.from_pretrained(text_encoder_path, subfolder="tokenizer")
     patchifier = SymmetricPatchifier(patch_size=1)
 
-    precision = config.get("precision", "bfloat16")
     if precision == "bfloat16":
-        vae.to(torch.bfloat16)
         transformer.to(torch.bfloat16)
         text_encoder.to(torch.bfloat16)
     
     pipeline = LTXVideoPipeline(
         transformer=transformer, patchifier=patchifier, text_encoder=text_encoder,
-        tokenizer=tokenizer, scheduler=scheduler, vae=vae,
+        tokenizer=tokenizer, scheduler=scheduler, vae=None, # VAE é explicitamente None
         allowed_inference_steps=allowed_inference_steps,
         prompt_enhancer_image_caption_model=None, prompt_enhancer_image_caption_processor=None,
         prompt_enhancer_llm_model=None, prompt_enhancer_llm_tokenizer=None,
     )
-
-
-    vae = CausalVideoAutoencoder.from_pretrained(ckpt_path).to("cpu")
-    if precision == "bfloat16":
-        vae.to(torch.bfloat16)
-
-    latent_upsampler = None
-    if config.get("spatial_upscaler_model_path"):
-        spatial_path = config["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 spatial_path.is_file():
-            raise FileNotFoundError(f"Main checkpoint upscaler file not found: {spatial_path_str}")
-        logging.info(f"Building UPSCALER pipeline ckpt:{spatial_path_str}")
-        latent_upsampler = create_latent_upsampler(spatial_path, device="cpu")
-        if precision == "bfloat16":
-            latent_upsampler.to(torch.bfloat16)
-
-    logging.info(f"LTX pipeline built on CPU in {time.perf_counter() - t0:.2f}s")
-    return pipeline, latent_upsampler, vae
-
+    return pipeline
 
 # ==============================================================================
-# --- FUNÇÕES AUXILIARES (Seed, Preparação de Imagem) ---
+# --- FUNÇÕES AUXILIARES GENÉRICAS ---
 # ==============================================================================
 
 def seed_everything(seed: int):
-    """Sets the seed for reproducibility."""
+    """Define a semente para PyTorch, NumPy e Python para garantir reprodutibilidade."""
     random.seed(seed)
     os.environ['PYTHONHASHSEED'] = str(seed)
     np.random.seed(seed)
@@ -150,41 +110,7 @@ def seed_everything(seed: int):
     torch.backends.cudnn.deterministic = True
     torch.backends.cudnn.benchmark = False
 
-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."""
-    if isinstance(image_input, str):
-        image = Image.open(image_input).convert("RGB")
-    elif isinstance(image_input, Image.Image):
-        image = image_input
-    else:
-        raise ValueError("image_input must be a file path or a PIL Image object")
-
-    input_width, input_height = image.size
-    aspect_ratio_target = target_width / target_height
-    aspect_ratio_frame = input_width / input_height
-
-    if aspect_ratio_frame > aspect_ratio_target:
-        new_width, new_height = int(input_height * aspect_ratio_target), input_height
-        x_start, y_start = (input_width - new_width) // 2, 0
-    else:
-        new_width, new_height = input_width, int(input_width / aspect_ratio_target)
-        x_start, y_start = 0, (input_height - new_height) // 2
-
-    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.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
-    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)
+# NOTA: A função load_image_to_tensor_with_resize_and_crop foi movida para o
+# cliente vae_aduc_pipeline.py, pois é uma dependência direta de pré-processamento
+# para as tarefas de VAE, tornando aquele módulo mais autocontido. Se for
+# necessária em outros locais, este seria o lugar para centralizá-la novamente.