Update api/ltx/ltx_utils.py

api/ltx/ltx_utils.py

@@ -1,263 +1,144 @@
 # FILE: api/ltx/ltx_utils.py
-# DESCRIPTION: A pure utility library for the LTX ecosystem.
-# Contains the official low-level builder function for the complete pipeline
-# and other stateless helper functions.
+# DESCRIPTION: Comprehensive, self-contained utility module for the LTX pipeline.
+# Handles dependency path injection, model loading, pipeline creation, and tensor preparation.
 
 import os
 import random
 import json
 import logging
+import time
 import sys
 from pathlib import Path
-from typing import Dict, Tuple, Union
-import torchvision.transforms.functional as TVF
-from PIL import Image
+from typing import Dict, Optional, Tuple, Union
 
+import numpy as np
 import torch
+import torchvision.transforms.functional as TVF
+from PIL import Image
 from safetensors import safe_open
 from transformers import T5EncoderModel, T5Tokenizer
 
 # ==============================================================================
-# --- CONFIGURAÇÃO DE PATH E IMPORTS DA BIBLIOTECA LTX ---
+# --- CRITICAL: DEPENDENCY PATH INJECTION ---
 # ==============================================================================
 
+# Define o caminho para o repositório clonado
 LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
 
 def add_deps_to_path():
-    """Adiciona o diretório do repositório LTX ao sys.path para importação de suas bibliotecas."""
+    """
+    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()
 
-try:
-    from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
-    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
-except ImportError as e:
-    logging.critical("Failed to import a core LTX-Video library component.", exc_info=True)
-    raise ImportError(f"Could not import from LTX-Video library. Check repo integrity at '{LTX_VIDEO_REPO_DIR}'. Error: {e}")
-
-# ==============================================================================
-# --- FUNÇÃO HELPER 'create_transformer' (Essencial) ---
-# ==============================================================================
-
-def create_transformer(ckpt_path: str, precision: str) -> Transformer3DModel:
-    """
-    Cria e carrega o modelo Transformer3D com a lógica de precisão correta,
-    incluindo suporte para a otimização float8_e4m3fn.
-    """
-    if precision == "float8_e4m3fn":
-        try:
-            from q8_kernels.integration.patch_transformer import patch_diffusers_transformer as patch_transformer_for_q8_kernels
-            transformer = Transformer3DModel.from_pretrained(ckpt_path, dtype=torch.float8_e4m3fn)
-            patch_transformer_for_q8_kernels(transformer)
-            return transformer
-        except ImportError:
-            raise ValueError("Q8-Kernels not found. To use FP8 checkpoint, please install Q8 kernels from the project's wheels.")
-    elif precision == "bfloat16":
-        return Transformer3DModel.from_pretrained(ckpt_path).to(torch.bfloat16)
-    else:
-        return Transformer3DModel.from_pretrained(ckpt_path)
-
-# ==============================================================================
-# --- BUILDER DE BAIXO NÍVEL OFICIAL ---
-# ==============================================================================
-
-def build_complete_pipeline_on_cpu(checkpoint_path: str, config: Dict) -> LTXVideoPipeline:
-    """
-    Constrói o pipeline LTX COMPLETO, incluindo o VAE, e o mantém na CPU.
-    Esta é a função de construção fundamental usada pelo LTXAducManager.
-    """
-    logging.info(f"Building complete LTX pipeline from checkpoint: {Path(checkpoint_path).name}")
-    
-    with safe_open(checkpoint_path, framework="pt") as f:
-        metadata = f.metadata() or {}
-        config_str = metadata.get("config", "{}")
-        allowed_inference_steps = json.loads(config_str).get("allowed_inference_steps")
-
-    precision = config.get("precision", "bfloat16")
-    
-    # Usa a função helper correta para criar o transformer
-    transformer = create_transformer(checkpoint_path, precision).to("cpu")
-    
-    scheduler = RectifiedFlowScheduler.from_pretrained(checkpoint_path)
-    text_encoder = T5EncoderModel.from_pretrained(config["text_encoder_model_name_or_path"], subfolder="text_encoder").to("cpu")
-    tokenizer = T5Tokenizer.from_pretrained(config["text_encoder_model_name_or_path"], subfolder="tokenizer")
-    patchifier = SymmetricPatchifier(patch_size=1)
-    vae = CausalVideoAutoencoder.from_pretrained(checkpoint_path).to("cpu")
-
-    if precision == "bfloat16":
-        text_encoder.to(torch.bfloat16)
-        vae.to(torch.bfloat16)
-        # O transformer já foi convertido para bfloat16 dentro de create_transformer, se aplicável
-
-    pipeline = LTXVideoPipeline(
-        transformer=transformer,
-        patchifier=patchifier,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        scheduler=scheduler,
-        vae=vae, # VAE é incluído para que o pipeline possa ser auto-suficiente
-        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,
-    )
-    
-    return pipeline
-
-# ==============================================================================
-# --- FUNÇÕES AUXILIARES GENÉRICAS ---
-# ==============================================================================
-
-
-# # FILE: api/ltx/ltx_utils.py
-# DESCRIPTION: A pure utility library for the LTX ecosystem.
-# Contains the official low-level builder function for the complete pipeline
-# and other stateless helper functions.
-
-import os
-import random
-import json
-import logging
-import sys
-from pathlib import Path
-from typing import Dict, Tuple
-
-import torch
-from safetensors import safe_open
-from transformers import T5EncoderModel, T5Tokenizer
 
 # ==============================================================================
-# --- CONFIGURAÇÃO DE PATH E IMPORTS DA BIBLIOTECA LTX ---
+# --- IMPORTAÇÕES DA BIBLIOTECA LTX-VIDEO (Após configuração do path) ---
 # ==============================================================================
-
-LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
-
-def add_deps_to_path():
-    """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}")
-
-add_deps_to_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
+    import ltx_video.pipelines.crf_compressor as crf_compressor
 except ImportError as e:
-    logging.critical("Failed to import a core LTX-Video library component.", exc_info=True)
-    raise ImportError(f"Could not import from LTX-Video library. Check repo integrity at '{LTX_VIDEO_REPO_DIR}'. Error: {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}")
+
 
 # ==============================================================================
-# --- FUNÇÃO HELPER 'create_transformer' (Essencial) ---
+# --- FUNÇÕES DE CONSTRUÇÃO DE MODELO E PIPELINE ---
 # ==============================================================================
 
-def create_transformer(ckpt_path: str, precision: str) -> Transformer3DModel:
-    """
-    Cria e carrega o modelo Transformer3D com a lógica de precisão correta,
-    incluindo suporte para a otimização float8_e4m3fn.
-    """
-    if precision == "float8_e4m3fn":
-        try:
-            from q8_kernels.integration.patch_transformer import patch_diffusers_transformer as patch_transformer_for_q8_kernels
-            transformer = Transformer3DModel.from_pretrained(ckpt_path, dtype=torch.float8_e4m3fn)
-            patch_transformer_for_q8_kernels(transformer)
-            return transformer
-        except ImportError:
-            raise ValueError("Q8-Kernels not found. To use FP8 checkpoint, please install Q8 kernels from the project's wheels.")
-    elif precision == "bfloat16":
-        return Transformer3DModel.from_pretrained(ckpt_path).to(torch.bfloat16)
-    else:
-        return Transformer3DModel.from_pretrained(ckpt_path)
+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
 
-# ==============================================================================
-# --- BUILDER DE BAIXO NÍVEL OFICIAL ---
-# ==============================================================================
+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_complete_pipeline_on_cpu(checkpoint_path: str, config: Dict) -> LTXVideoPipeline:
-    """
-    Constrói o pipeline LTX COMPLETO, incluindo o VAE, e o mantém na CPU.
-    Esta é a função de construção fundamental usada pelo LTXAducManager.
-    """
-    logging.info(f"Building complete LTX pipeline from checkpoint: {Path(checkpoint_path).name}")
-    
-    with safe_open(checkpoint_path, framework="pt") as f:
+    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", "{}")
-        allowed_inference_steps = json.loads(config_str).get("allowed_inference_steps")
+        configs = json.loads(config_str)
+        allowed_inference_steps = configs.get("allowed_inference_steps")
 
-    precision = config.get("precision", "bfloat16")
+    vae = CausalVideoAutoencoder.from_pretrained(ckpt_path).to("cpu")
+    transformer = Transformer3DModel.from_pretrained(ckpt_path).to("cpu")
+    scheduler = RectifiedFlowScheduler.from_pretrained(ckpt_path)
     
-    # Usa a função helper correta para criar o transformer
-    transformer = create_transformer(checkpoint_path, precision).to("cpu")
-    
-    scheduler = RectifiedFlowScheduler.from_pretrained(checkpoint_path)
-    text_encoder = T5EncoderModel.from_pretrained(config["text_encoder_model_name_or_path"], subfolder="text_encoder").to("cpu")
-    tokenizer = T5Tokenizer.from_pretrained(config["text_encoder_model_name_or_path"], subfolder="tokenizer")
+    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)
-    vae = CausalVideoAutoencoder.from_pretrained(checkpoint_path).to("cpu")
 
+    precision = config.get("precision", "bfloat16")
     if precision == "bfloat16":
-        text_encoder.to(torch.bfloat16)
         vae.to(torch.bfloat16)
-        # O transformer já foi convertido para bfloat16 dentro de create_transformer, se aplicável
-
+        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, # VAE é incluído para que o pipeline possa ser auto-suficiente
+        transformer=transformer, patchifier=patchifier, text_encoder=text_encoder,
+        tokenizer=tokenizer, scheduler=scheduler, vae=vae,
         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,
+        prompt_enhancer_image_caption_model=None, prompt_enhancer_image_caption_processor=None,
+        prompt_enhancer_llm_model=None, prompt_enhancer_llm_tokenizer=None,
     )
-    
-    return pipeline
+
+    latent_upsampler = None
+    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)
+
+    logging.info(f"LTX pipeline built on CPU in {time.perf_counter() - t0:.2f}s")
+    return pipeline, latent_upsampler
+
 
 # ==============================================================================
-# --- FUNÇÕES AUXILIARES GENÉRICAS ---
+# --- FUNÇÕES AUXILIARES (Seed, Preparação de Imagem) ---
 # ==============================================================================
 
 def seed_everything(seed: int):
-    """
-    Define a semente para PyTorch, NumPy e Python para garantir reprodutibilidade.
-    """
+    """Sets the seed for reproducibility."""
     random.seed(seed)
     os.environ['PYTHONHASHSEED'] = str(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = Fals
-    
+    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:
-    """
-    Carrega, redimensiona, corta e processa uma imagem para um tensor de pixel 5D,
-    normalizado para [-1, 1], pronto para ser enviado ao VAE para encoding.
-    """
+    """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.convert("RGB")
+        image = image_input
     else:
         raise ValueError("image_input must be a file path or a PIL Image object")
 
@@ -267,38 +148,22 @@ def load_image_to_tensor_with_resize_and_crop(
 
     if aspect_ratio_frame > aspect_ratio_target:
         new_width, new_height = int(input_height * aspect_ratio_target), input_height
-        x_start = (input_width - new_width) // 2
-        image = image.crop((x_start, 0, x_start + new_width, new_height))
+        x_start, y_start = (input_width - new_width) // 2, 0
     else:
-        new_height = int(input_width / aspect_ratio_target)
-        y_start = (input_height - new_height) // 2
-        image = image.crop((0, y_start, input_width, y_start + new_height))
-    
+        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 = TVF.to_tensor(image)
-    
-    # Esta parte depende de 'crf_compressor', então precisamos importá-lo aqui também
-    try:
-        from ltx_video.pipelines import crf_compressor
-        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)
-    except ImportError:
-        logging.warning("CRF Compressor not found. Skipping compression step.")
-    
+    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
-    return frame_tensor.unsqueeze(0).unsqueeze(2)
-
-
-def seed_everything(seed: int):
-    """
-    Define a semente para PyTorch, NumPy e Python para garantir reprodutibilidade.
-    """
-    random.seed(seed)
-    os.environ['PYTHONHASHSEED'] = str(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
+    
+    # Create 5D tensor: (batch_size=1, channels=3, num_frames=1, height, width)
+    return frame_tensor.unsqueeze(0).unsqueeze(2)