Update api/ltx/ltx_aduc_manager.py

api/ltx/ltx_aduc_manager.py

@@ -5,7 +5,6 @@
 import logging
 import torch
 import sys
-import os
 from pathlib import Path
 import threading
 import queue
@@ -14,10 +13,9 @@ import yaml
 from huggingface_hub import hf_hub_download
 from typing import List, Optional, Callable, Any, Tuple, Dict
 
-# --- Importa o gerenciador de GPUs e os builders de baixo nível ---
+# --- Importa o gerenciador de GPUs e o builder de baixo nível ---
 from managers.gpu_manager import gpu_manager
-# Importa os builders que foram movidos para ltx_utils
-from api.ltx.ltx_utils import _build_ltx_transformer_pipeline, _build_vae, _build_latent_upscaler
+from api.ltx.ltx_utils import build_components_on_cpu
 
 # --- Adiciona o path do LTX-Video para importação de tipos ---
 LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
@@ -29,76 +27,56 @@ add_deps_to_path()
 
 from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
 from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
-from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
-
-# ==============================================================================
-# --- FUNÇÕES DE ORQUESTRAÇÃO DA CONSTRUÇÃO (Internas ao Manager) ---
-# ==============================================================================
-
-def _load_config() -> Dict:
-    """Helper centralizado para carregar o arquivo de configuração YAML."""
-    config_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
-    with open(config_path, "r") as file:
-        return yaml.safe_load(file)
-
-def get_main_ltx_pipeline() -> LTXVideoPipeline:
-    """Orquestra a construção do Pipeline Transformer principal (sem o VAE)."""
-    config = _load_config()
-    precision = config.get("precision", "bfloat16")
-    ckpt_path_str = hf_hub_download(repo_id="Lightricks/LTX-Video", filename=config["checkpoint_path"], cache_dir=os.environ.get("HF_HOME"))
-    return _build_ltx_transformer_pipeline(ckpt_path_str, config, precision)
-
-def get_main_vae() -> CausalVideoAutoencoder:
-    """Orquestra a construção do VAE principal."""
-    config = _load_config()
-    precision = config.get("precision", "bfloat16")
-    ckpt_path_str = hf_hub_download(repo_id="Lightricks/LTX-Video", filename=config["checkpoint_path"], cache_dir=os.environ.get("HF_HOME"))
-    return _build_vae(ckpt_path_str, precision)
 
 # ==============================================================================
 # --- CLASSES DE WORKER (Especialistas em Tarefas) ---
 # ==============================================================================
 
 class BaseWorker(threading.Thread):
-    def __init__(self, worker_id: int, device: torch.device):
+    """Classe base para nossos workers com gerenciamento de estado e saúde."""
+    def __init__(self, worker_id: int, device: torch.device, model: torch.nn.Module):
         super().__init__()
         self.worker_id = worker_id
         self.device = device
+        self.model = model
         self.is_healthy = False
         self.is_busy = False
         self.daemon = True
 
     def run(self):
+        """O loop de vida do worker, responsável por mover o modelo para a GPU."""
         try:
-            self._load_models()
+            logging.info(f"Worker {self.worker_id} ({self.__class__.__name__}) moving model to {self.device}...")
+            self.model.to(self.device)
+            self._post_load_hook()
             self.is_healthy = True
             logging.info(f"✅ Worker {self.worker_id} ({self.__class__.__name__}) on {self.device} is healthy and ready.")
         except Exception:
             self.is_healthy = False
             logging.error(f"❌ Worker {self.worker_id} on {self.device} FAILED to initialize!", exc_info=True)
 
-    def _load_models(self):
-        raise NotImplementedError
+    def _post_load_hook(self):
+        """Gancho para ações pós-carregamento, como chamar .eval()."""
+        pass
 
     def get_status(self) -> Tuple[bool, bool]:
         return self.is_healthy, self.is_busy
 
 class LTXMainWorker(BaseWorker):
-    def __init__(self, worker_id: int, device: torch.device):
-        super().__init__(worker_id, device)
-        self.pipeline: Optional[LTXVideoPipeline] = None
+    """Worker especialista para o pipeline principal do LTX."""
+    def __init__(self, worker_id: int, device: torch.device, pipeline: LTXVideoPipeline):
+        super().__init__(worker_id, device, pipeline)
+        self.pipeline = self.model # Alias para clareza
         self.autocast_dtype: torch.dtype = torch.float32
 
-    def _load_models(self):
-        logging.info(f"[LTXWorker-{self.worker_id}] Loading models to CPU...")
-        self.pipeline = get_main_ltx_pipeline()
+    def _post_load_hook(self):
         self._set_precision_policy()
-        logging.info(f"[LTXWorker-{self.worker_id}] Moving pipeline to {self.device}...")
-        self.pipeline.to(self.device)
 
     def _set_precision_policy(self):
         try:
-            config = _load_config()
+            config_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
+            with open(config_path, "r") as file:
+                config = yaml.safe_load(file)
             precision = str(config.get("precision", "bfloat16")).lower()
             if precision in ["float8_e4m3fn", "bfloat16"]: self.autocast_dtype = torch.bfloat16
             elif precision == "mixed_precision": self.autocast_dtype = torch.float16
@@ -118,15 +96,12 @@ class LTXMainWorker(BaseWorker):
             self.is_busy = False
 
 class VAEWorker(BaseWorker):
-    def __init__(self, worker_id: int, device: torch.device):
-        super().__init__(worker_id, device)
-        self.vae: Optional[CausalVideoAutoencoder] = None
-    
-    def _load_models(self):
-        logging.info(f"[VAEWorker-{self.worker_id}] Loading VAE model to CPU...")
-        self.vae = get_main_vae()
-        logging.info(f"[VAEWorker-{self.worker_id}] Moving VAE to {self.device}...")
-        self.vae.to(self.device)
+    """Worker especialista para o modelo VAE."""
+    def __init__(self, worker_id: int, device: torch.device, vae: CausalVideoAutoencoder):
+        super().__init__(worker_id, device, vae)
+        self.vae = self.model # Alias
+
+    def _post_load_hook(self):
         self.vae.eval()
 
     def execute(self, job_func: Callable, args: tuple, kwargs: dict) -> Any:
@@ -163,8 +138,12 @@ class LTXAducManager:
         self.vae_job_queue = queue.Queue()
         self.pool_lock = threading.Lock()
 
+        # Carrega os modelos na CPU antes de criar os workers
+        self.main_pipeline, self.main_vae = self._load_components_once()
+        
         self._initialize_workers()
         
+        # Inicia threads consumidores para processar as filas
         self.ltx_dispatcher = threading.Thread(target=self._dispatch_jobs, args=(self.ltx_job_queue, self.ltx_workers), daemon=True)
         self.vae_dispatcher = threading.Thread(target=self._dispatch_jobs, args=(self.vae_job_queue, self.vae_workers), daemon=True)
         self.health_monitor = threading.Thread(target=self._health_check_loop, daemon=True)
@@ -176,18 +155,29 @@ class LTXAducManager:
         self._initialized = True
         logging.info("✅ Advanced Pool Manager is running with all threads started.")
 
+    def _load_components_once(self) -> Tuple[LTXVideoPipeline, CausalVideoAutoencoder]:
+        """Orquestra a construção de TODOS os componentes na CPU uma única vez."""
+        logging.info("Manager loading all components onto CPU...")
+        config_path = LTX_VIDEO_REPO_DIR / "configs" / "ltxv-13b-0.9.8-distilled-fp8.yaml"
+        with open(config_path, "r") as file:
+            config = yaml.safe_load(file)
+        
+        ckpt_path = hf_hub_download(repo_id="Lightricks/LTX-Video", filename=config["checkpoint_path"], cache_dir=os.environ.get("HF_HOME"))
+        pipeline, vae = build_components_on_cpu(ckpt_path, config)
+        logging.info("✅ All components loaded to CPU successfully.")
+        return pipeline, vae
+
     def _initialize_workers(self):
+        """Cria e inicia os workers, injetando os modelos já carregados."""
         ltx_device = gpu_manager.get_ltx_device()
         vae_device = gpu_manager.get_ltx_vae_device()
 
         with self.pool_lock:
-            # Para esta fase, criamos um worker de cada tipo.
-            # No futuro, pode-se iterar sobre listas de GPUs de gpu_manager.
-            ltx_worker = LTXMainWorker(worker_id=0, device=ltx_device)
+            ltx_worker = LTXMainWorker(worker_id=0, device=ltx_device, pipeline=self.main_pipeline)
             self.ltx_workers.append(ltx_worker)
             ltx_worker.start()
 
-            vae_worker = VAEWorker(worker_id=0, device=vae_device)
+            vae_worker = VAEWorker(worker_id=0, device=vae_device, vae=self.main_vae)
             self.vae_workers.append(vae_worker)
             vae_worker.start()
 
@@ -195,8 +185,7 @@ class LTXAducManager:
         with self.pool_lock:
             for worker in worker_pool:
                 healthy, busy = worker.get_status()
-                if healthy and not busy:
-                    return worker
+                if healthy and not busy: return worker
         return None
 
     def _dispatch_jobs(self, job_queue: queue.Queue, worker_pool: List[BaseWorker]):
@@ -206,7 +195,6 @@ class LTXAducManager:
             while worker is None:
                 worker = self._get_available_worker(worker_pool)
                 if worker is None: time.sleep(0.1) 
-
             try:
                 result = worker.execute(job_func, args, kwargs)
                 future.put(result)
@@ -220,14 +208,14 @@ class LTXAducManager:
                 for i, worker in enumerate(self.ltx_workers):
                     if not worker.is_alive() or not worker.is_healthy:
                         logging.warning(f"LTX Worker {worker.worker_id} on {worker.device} is UNHEALTHY. Restarting...")
-                        new_worker = LTXMainWorker(worker.worker_id, worker.device)
+                        new_worker = LTXMainWorker(worker.worker_id, worker.device, self.main_pipeline)
                         self.ltx_workers[i] = new_worker
                         new_worker.start()
                 
                 for i, worker in enumerate(self.vae_workers):
                     if not worker.is_alive() or not worker.is_healthy:
                         logging.warning(f"VAE Worker {worker.worker_id} on {worker.device} is UNHEALTHY. Restarting...")
-                        new_worker = VAEWorker(worker.worker_id, worker.device)
+                        new_worker = VAEWorker(worker.worker_id, worker.device, self.main_vae)
                         self.vae_workers[i] = new_worker
                         new_worker.start()
 
@@ -237,9 +225,7 @@ class LTXAducManager:
         
         job_queue = self.ltx_job_queue if job_type == 'ltx' else self.vae_job_queue
         future = queue.Queue(1)
-
         job_queue.put((job_func, args, kwargs, future))
-        
         result = future.get()
         
         if isinstance(result, Exception):
@@ -247,11 +233,9 @@ class LTXAducManager:
         
         return result
 
-# ==============================================================================
 # --- INSTANCIAÇÃO GLOBAL ---
-# ==============================================================================
 try:
     ltx_aduc_manager = LTXAducManager()
-except Exception as e:
+except Exception:
     logging.critical("CRITICAL ERROR: Failed to initialize the LTXAducManager pool.", exc_info=True)
     ltx_aduc_manager = None