Update api/ltx/ltx_utils.py

api/ltx/ltx_utils.py

@@ -10,6 +10,8 @@ import logging
 import sys
 from pathlib import Path
 from typing import Dict, Tuple
+import torchvision.transforms.functional as TVF
+from PIL import Image
 
 import torch
 from safetensors import safe_open
@@ -114,6 +116,181 @@ def build_complete_pipeline_on_cpu(checkpoint_path: str, config: Dict) -> LTXVid
 # --- 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 ---
+# ==============================================================================
+
+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.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 ---
+# ==============================================================================
+
+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 = Fals
+    
+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.
+    """
+    if isinstance(image_input, str):
+        image = Image.open(image_input).convert("RGB")
+    elif isinstance(image_input, Image.Image):
+        image = image_input.convert("RGB")
+    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 = (input_width - new_width) // 2
+        image = image.crop((x_start, 0, x_start + new_width, new_height))
+    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))
+    
+    image = image.resize((target_width, target_height), Image.Resampling.LANCZOS)
+    
+    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 = (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.