Create ltx_utils.py

api/ltx/ltx_utils.py

@@ -0,0 +1,203 @@
+# FILE: api/ltx/ltx_utils.py
+# DESCRIPTION: Comprehensive, self-contained utility module for the LTX pipeline.
+# 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 dataclasses import dataclass
+from enum import Enum, auto
+
+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")
+
+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 (Agora devem funcionar) ---
+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}")
+
+
+# ==============================================================================
+# --- ESTRUTURAS DE DADOS E ENUMS (Centralizadas aqui) ---
+# ==============================================================================
+
+@dataclass
+class ConditioningItem:
+    """Defines 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."""
+    AttentionSkip = auto()
+    AttentionValues = auto()
+    Residual = auto()
+    TransformerBlock = auto()
+
+
+# ==============================================================================
+# --- 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 = 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", "{}")
+        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,
+    )
+
+    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 (Latent Processing, Seed, Image Prep) ---
+# ==============================================================================
+
+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)
+    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 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)
+    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)
+    frame_tensor = (frame_tensor * 2.0) - 1.0
+    
+    return frame_tensor.unsqueeze(0).unsqueeze(2)