Update api/ltx/ltx_utils.py

api/ltx/ltx_utils.py

@@ -1,7 +1,6 @@
 # FILE: api/ltx/ltx_utils.py
 # DESCRIPTION: A pure utility library for the LTX ecosystem.
-# REFACTORED to contain only the official, low-level builder function for core components
-# and other stateless helper functions.
+# Contains the official low-level builder function for core components and other stateless helpers.
 
 import os
 import random
@@ -13,7 +12,7 @@ from typing import Dict, Tuple
 
 import torch
 from safetensors import safe_open
-from transformers import T5EncoderModel, T5Tokenizer, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
+from transformers import T5EncoderModel, T5Tokenizer
 
 # ==============================================================================
 # --- CONFIGURAÇÃO DE PATH E IMPORTS DA BIBLIOTECA LTX ---
@@ -33,64 +32,73 @@ 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 create_transformer
+    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}")
 
 # ==============================================================================
-# --- BUILDER DE BAIXO NÍVEL OFICIAL ---
-# (Esta é a única função de construção, usada pelo LTXAducManager)
+# --- FUNÇÃO HELPER 'create_transformer' (Essencial) ---
 # ==============================================================================
 
-def build_components_on_cpu(checkpoint_path: str, config: Dict) -> Tuple[LTXVideoPipeline, CausalVideoAutoencoder]:
+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.
     """
-    Constrói o pipeline LTX principal (sem VAE) e o modelo VAE separadamente,
-    mantendo ambos os componentes na CPU. Esta é a função de construção fundamental
-    usada pelo Manager antes de distribuir os modelos para as GPUs.
+    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)
 
-    Args:
-        checkpoint_path (str): Caminho absoluto para o arquivo de checkpoint principal.
-        config (Dict): O dicionário de configuração carregado do arquivo YAML.
+# ==============================================================================
+# --- BUILDER DE BAIXO NÍVEL OFICIAL ---
+# ==============================================================================
 
-    Returns:
-        Tuple[LTXVideoPipeline, CausalVideoAutoencoder]: Uma tupla contendo o pipeline principal
-                                                         e o modelo VAE, ambos na CPU.
+def build_components_on_cpu(checkpoint_path: str, config: Dict) -> Tuple[LTXVideoPipeline, CausalVideoAutoencoder]:
+    """
+    Constrói o pipeline LTX principal (sem VAE) e o modelo VAE separadamente, na CPU.
+    Esta é a função de construção fundamental usada pelo LTXAducManager.
     """
     logging.info(f"Building LTX components from checkpoint: {Path(checkpoint_path).name}")
     
     with safe_open(checkpoint_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")
+        allowed_inference_steps = json.loads(config_str).get("allowed_inference_steps")
 
-    # --- Construir componentes na CPU ---
     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)
-    
-    # Construir o VAE separadamente, também na CPU
     vae = CausalVideoAutoencoder.from_pretrained(checkpoint_path).to("cpu")
 
-    # Aplicar precisão bfloat16 se configurado
     if precision == "bfloat16":
-        transformer.to(torch.bfloat16)
         text_encoder.to(torch.bfloat16)
         vae.to(torch.bfloat16)
 
-    # Montar o pipeline principal, passando 'vae=None' para garantir o desacoplamento
     pipeline = LTXVideoPipeline(
         transformer=transformer,
         patchifier=patchifier,
         text_encoder=text_encoder,
         tokenizer=tokenizer,
         scheduler=scheduler,
-        vae=None,  # VAE é explicitamente desacoplado do pipeline principal
+        vae=None, # VAE é desacoplado para ser gerenciado por um worker separado
         allowed_inference_steps=allowed_inference_steps,
         prompt_enhancer_image_caption_model=None,
         prompt_enhancer_image_caption_processor=None,
@@ -106,8 +114,7 @@ def build_components_on_cpu(checkpoint_path: str, config: Dict) -> Tuple[LTXVide
 
 def seed_everything(seed: int):
     """
-    Define a semente para PyTorch, NumPy e Python para garantir reprodutibilidade
-    em experimentos ou gerações.
+    Define a semente para PyTorch, NumPy e Python para garantir reprodutibilidade.
     """
     random.seed(seed)
     os.environ['PYTHONHASHSEED'] = str(seed)