Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -1,15 +1,6 @@
 # ltx_manager_helpers.py
 # Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
-#
-# ORIGINAL SOURCE: LTX-Video by Lightricks Ltd. & other open-source projects.
-# Licensed under the Apache License, Version 2.0
-# https://github.com/Lightricks/LTX-Video
-#
-# MODIFICATIONS FOR ADUC-SDR_Video:
-# This file is part of ADUC-SDR_Video, a derivative work based on LTX-Video.
-# It has been modified to manage pools of LTX workers, handle GPU memory,
-# and prepare parameters for the ADUC-SDR orchestration framework.
-# All modifications are also licensed under the Apache License, Version 2.0.
+# (Licenciamento e cabeçalhos permanecem os mesmos)
 
 import torch
 import gc
@@ -34,6 +25,7 @@ class LtxWorker:
     Gerencia o carregamento do modelo para a CPU e a movimentação de/para a GPU.
     """
     def __init__(self, device_id, ltx_config_file):
+        # ... (código do LtxWorker __init__ permanece o mesmo) ...
         self.cpu_device = torch.device('cpu')
         self.device = torch.device(device_id if torch.cuda.is_available() else 'cpu')
         logger.info(f"LTX Worker ({self.device}): Inicializando com config '{ltx_config_file}'...")
@@ -60,20 +52,19 @@ class LtxWorker:
         )
         logger.info(f"LTX Worker ({self.device}): Modelo pronto na CPU. É um modelo destilado? {self.is_distilled}")
 
-        if self.device.type == 'cuda' and can_optimize_fp8():
-            logger.info(f"LTX Worker ({self.device}): GPU com suporte a FP8 detectada. Iniciando otimização...")
-            self.pipeline.to(self.device)
-            optimize_ltx_worker(self)
-            self.pipeline.to(self.cpu_device)
-            logger.info(f"LTX Worker ({self.device}): Otimização concluída. Modelo pronto.")
-        elif self.device.type == 'cuda':
-            logger.info(f"LTX Worker ({self.device}): Otimização FP8 não suportada ou desativada. Usando modelo padrão.")
-
     def to_gpu(self):
-        """Move o pipeline para a GPU designada."""
+        """Move o pipeline para a GPU designada E OTIMIZA SE POSSÍVEL."""
         if self.device.type == 'cpu': return
         logger.info(f"LTX Worker: Movendo pipeline para a GPU {self.device}...")
         self.pipeline.to(self.device)
+        
+        # A otimização agora ocorre aqui, uma única vez, quando o modelo vai para a GPU.
+        if self.device.type == 'cuda' and can_optimize_fp8():
+            logger.info(f"LTX Worker ({self.device}): GPU com suporte a FP8 detectada. Iniciando otimização...")
+            optimize_ltx_worker(self)
+            logger.info(f"LTX Worker ({self.device}): Otimização concluída.")
+        elif self.device.type == 'cuda':
+            logger.info(f"LTX Worker ({self.device}): Otimização FP8 não suportada ou desativada.")
 
     def to_cpu(self):
         """Move o pipeline de volta para a CPU e libera a memória da GPU."""
@@ -89,23 +80,29 @@ class LtxWorker:
 
 class LtxPoolManager:
     """
-    Gerencia um pool de LtxWorkers para otimizar o uso de múltiplas GPUs,
-    alternando o worker ativo para permitir que o anterior descarregue da VRAM em segundo plano.
+    Gerencia um pool de LtxWorkers para otimizar o uso de múltiplas GPUs.
+    NOVO MODO "HOT START": Mantém todos os modelos carregados na VRAM para latência mínima.
     """
     def __init__(self, device_ids, ltx_config_file):
         logger.info(f"LTX POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
         self.workers = [LtxWorker(dev_id, ltx_config_file) for dev_id in device_ids]
         self.current_worker_index = 0
         self.lock = threading.Lock()
-        self.last_cleanup_thread = None
 
-    def _cleanup_worker_thread(self, worker):
-        """Thread para descarregar um worker da GPU em segundo plano."""
-        logger.info(f"LTX CLEANUP THREAD: Iniciando limpeza de {worker.device} em background...")
-        worker.to_cpu()
+        # ######################################################################
+        # ##               MUDANÇA 1: PRÉ-AQUECIMENTO DAS GPUs                ##
+        # ######################################################################
+        if all(w.device.type == 'cuda' for w in self.workers):
+            logger.info("LTX POOL MANAGER: MODO HOT START ATIVADO. Pré-aquecendo todas as GPUs...")
+            for worker in self.workers:
+                worker.to_gpu()
+            logger.info("LTX POOL MANAGER: Todas as GPUs estão quentes e prontas.")
+        else:
+            logger.info("LTX POOL MANAGER: Operando em modo CPU ou misto. O pré-aquecimento de GPU foi ignorado.")
+        # ######################################################################
 
     def _prepare_and_log_params(self, worker_to_use, **kwargs):
-        """Prepara e registra os parâmetros para a chamada da pipeline LTX."""
+        # ... (Esta função permanece exatamente a mesma) ...
         target_device = worker_to_use.device
         height, width = kwargs['height'], kwargs['width']
         
@@ -113,7 +110,6 @@ class LtxPoolManager:
         final_conditioning_items = []
         conditioning_log_details = []
         for i, item in enumerate(conditioning_data):
-            # Lida tanto com LatentConditioningItem quanto ConditioningItem (se usado no futuro)
             if hasattr(item, 'latent_tensor'):
                 item.latent_tensor = item.latent_tensor.to(target_device)
                 final_conditioning_items.append(item)
@@ -146,7 +142,6 @@ class LtxPoolManager:
         else:
             pipeline_params["num_inference_steps"] = int(kwargs.get('num_inference_steps', 20))
         
-        # Log detalhado dos parâmetros para depuração.
         log_friendly_params = pipeline_params.copy()
         log_friendly_params.pop('generator', None)
         log_friendly_params.pop('conditioning_items', None)
@@ -162,80 +157,60 @@ class LtxPoolManager:
         
         return pipeline_params, padding_vals
     
-    def generate_latent_fragment(self, **kwargs) -> (torch.Tensor, tuple):
+    def _execute_on_worker(self, execution_fn, **kwargs):
         """
-        Orquestra a geração de um novo fragmento de vídeo a partir do zero (ruído).
+        Função unificada para selecionar um worker e executar uma tarefa,
+        sem a lógica de carregar/descarregar.
         """
         worker_to_use = None
         try:
             with self.lock:
-                if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
-                    self.last_cleanup_thread.join()
                 worker_to_use = self.workers[self.current_worker_index]
-                previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
-                worker_to_cleanup = self.workers[previous_worker_index]
-                cleanup_thread = threading.Thread(target=self._cleanup_worker_thread, args=(worker_to_cleanup,))
-                cleanup_thread.start()
-                self.last_cleanup_thread = cleanup_thread
-                worker_to_use.to_gpu()
                 self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
             
             pipeline_params, padding_vals = self._prepare_and_log_params(worker_to_use, **kwargs)
-            pipeline_params['output_type'] = "latent"
-
-            with torch.no_grad():
-                result_tensor = worker_to_use.generate_video_fragment_internal(**pipeline_params)
             
-            return result_tensor, padding_vals
+            result = execution_fn(worker_to_use, pipeline_params, **kwargs)
+            
+            return result, padding_vals
+            
         except Exception as e:
-            logger.error(f"LTX POOL MANAGER: Erro durante a geração de latentes: {e}", exc_info=True)
+            logger.error(f"LTX POOL MANAGER: Erro durante a execução em {worker_to_use.device if worker_to_use else 'N/A'}: {e}", exc_info=True)
             raise e
         finally:
-            if worker_to_use:
-                logger.info(f"LTX POOL MANAGER: Executando limpeza final para {worker_to_use.device}...")
-                worker_to_use.to_cpu()
+            # Apenas limpa o cache da GPU, não descarrega o modelo.
+            if worker_to_use and worker_to_use.device.type == 'cuda':
+                with torch.cuda.device(worker_to_use.device):
+                    gc.collect()
+                    torch.cuda.empty_cache()
+    
+    def generate_latent_fragment(self, **kwargs) -> (torch.Tensor, tuple):
+        """
+        Orquestra a geração de um novo fragmento de vídeo a partir do ruído.
+        """
+        def execution_logic(worker, params, **inner_kwargs):
+            params['output_type'] = "latent"
+            with torch.no_grad():
+                return worker.generate_video_fragment_internal(**params)
+
+        return self._execute_on_worker(execution_logic, **kwargs)
 
     def refine_latents(self, upscaled_latents: torch.Tensor, **kwargs) -> (torch.Tensor, tuple):
         """
-        Orquestra um passe de difusão curto em latentes já existentes para refinar texturas.
-        Usado na etapa de pós-produção de upscale.
+        Orquestra um passe de difusão curto em latentes já existentes para refinamento.
         """
-        worker_to_use = None
-        try:
-            with self.lock:
-                if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
-                    self.last_cleanup_thread.join()
-                worker_to_use = self.workers[self.current_worker_index]
-                previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
-                worker_to_cleanup = self.workers[previous_worker_index]
-                cleanup_thread = threading.Thread(target=self._cleanup_worker_thread, args=(worker_to_cleanup,))
-                cleanup_thread.start()
-                self.last_cleanup_thread = cleanup_thread
-                worker_to_use.to_gpu()
-                self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
-
-            pipeline_params, padding_vals = self._prepare_and_log_params(worker_to_use, **kwargs)
-            
-            # Parâmetros específicos para o passe de refinamento (denoise)
-            pipeline_params['latents'] = upscaled_latents.to(worker_to_use.device, dtype=worker_to_use.pipeline.transformer.dtype)
-            pipeline_params['strength'] = kwargs.get('denoise_strength', 0.4)
-            pipeline_params['num_inference_steps'] = int(kwargs.get('refine_steps', 10))
-            pipeline_params['output_type'] = "latent"
+        def execution_logic(worker, params, **inner_kwargs):
+            params['latents'] = upscaled_latents.to(worker.device, dtype=worker.pipeline.transformer.dtype)
+            params['strength'] = inner_kwargs.get('denoise_strength', 0.4)
+            params['num_inference_steps'] = int(inner_kwargs.get('refine_steps', 10))
+            params['output_type'] = "latent"
             
             logger.info("LTX POOL MANAGER: Iniciando passe de refinamento (denoise) em latentes de alta resolução.")
             
             with torch.no_grad():
-                refined_tensor = worker_to_use.generate_video_fragment_internal(**pipeline_params)
-            
-            return refined_tensor, padding_vals
+                return worker.generate_video_fragment_internal(**params)
 
-        except Exception as e:
-            logger.error(f"LTX POOL MANAGER: Erro durante o refinamento de latentes: {e}", exc_info=True)
-            raise e
-        finally:
-            if worker_to_use:
-                logger.info(f"LTX POOL MANAGER: Executando limpeza final para {worker_to_use.device}...")
-                worker_to_use.to_cpu()
+        return self._execute_on_worker(execution_logic, upscaled_latents=upscaled_latents, **kwargs)
 
 # --- Instanciação Singleton ---
 logger.info("Lendo config.yaml para inicializar o LTX Pool Manager...")