api/ltx/ltx_utils.py

# 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.

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

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

# ==============================================================================
# --- 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
    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}")


# ==============================================================================
# --- FUNÇÕES DE CONSTRUÇÃO DE MODELO E PIPELINE ---
# ==============================================================================

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...")



    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}")

    logging.info(f"Building LTX pipeline ckpt:{ckpt_path_str}")
    
    
    with safe_open(ckpt_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)
    
    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,
        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


# ==============================================================================
# --- FUNÇÕES AUXILIARES (Seed, Preparação de Imagem) ---
# ==============================================================================

def seed_everything(seed: int):
    """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 = 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)