Upload 2 files

api/debug_utils.py

@@ -0,0 +1,70 @@
+# FILE: api/utils/debug_utils.py
+# DESCRIPTION: A utility for detailed function logging and debugging.
+
+import os
+import functools
+import logging
+import torch
+
+# Define o nível de log. Mude para "INFO" para desativar os logs detalhados.
+# Você pode controlar isso com uma variável de ambiente.
+LOG_LEVEL = os.environ.get("ADUC_LOG_LEVEL", "DEBUG").upper()
+logging.basicConfig(level=LOG_LEVEL, format='[%(levelname)s] [%(name)s] %(message)s')
+logger = logging.getLogger("AducDebug")
+
+
+def _format_value(value):
+    """Formata os valores dos argumentos para uma exibição concisa e informativa."""
+    if isinstance(value, torch.Tensor):
+        return f"Tensor(shape={list(value.shape)}, device='{value.device}', dtype={value.dtype})"
+    if isinstance(value, str) and len(value) > 70:
+        return f"'{value[:70]}...'"
+    if isinstance(value, list) and len(value) > 5:
+        return f"List(len={len(value)})"
+    if isinstance(value, dict) and len(value.keys()) > 5:
+        return f"Dict(keys={list(value.keys())[:5]}...)"
+    return repr(value)
+
+def log_function_io(func):
+    """
+    Um decorador que registra as entradas, saídas e exceções de uma função.
+    Ele é ativado apenas se o nível de log estiver definido como DEBUG.
+    """
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # Só executa a lógica de log se o nível for DEBUG
+        if logger.isEnabledFor(logging.DEBUG):
+            # Obtém o nome do módulo e da função
+            func_name = f"{func.__module__}.{func.__name__}"
+            
+            # Formata os argumentos de entrada
+            args_repr = [_format_value(a) for a in args]
+            kwargs_repr = {k: _format_value(v) for k, v in kwargs.items()}
+            signature = ", ".join(args_repr + [f"{k}={v}" for k, v in kwargs_repr.items()])
+
+            # Log de Entrada
+            logger.debug(f"================ INÍCIO: {func_name} ================")
+            logger.debug(f"  -> ENTRADA: ({signature})")
+
+            try:
+                # Executa a função original
+                result = func(*args, **kwargs)
+                
+                # Formata e registra o resultado
+                result_repr = _format_value(result)
+                logger.debug(f"  <- SAÍDA: {result_repr}")
+                
+            except Exception as e:
+                # Registra qualquer exceção que ocorra
+                logger.error(f"  <-- ERRO em {func_name}: {e}", exc_info=True)
+                raise  # Re-lança a exceção para não alterar o comportamento do programa
+            finally:
+                # Log de Fim
+                logger.debug(f"================ FIM: {func_name} ================\n")
+            
+            return result
+        else:
+            # Se o log não estiver em modo DEBUG, executa a função sem nenhum overhead.
+            return func(*args, **kwargs)
+            
+    return wrapper

api/ltx_utils.py

@@ -0,0 +1,207 @@
+# 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 (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}")
+
+
+# ==============================================================================
+# --- 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 ---
+# ==============================================================================
+
+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)
+    # 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)