Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -25,7 +25,6 @@ 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}'...")
@@ -58,7 +57,6 @@ class LtxWorker:
         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)
@@ -80,8 +78,7 @@ class LtxWorker:
 
 class LtxPoolManager:
     """
-    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.
+    Gerencia um pool de LtxWorkers. MODO "HOT START": Mantém todos os modelos carregados na VRAM.
     """
     def __init__(self, device_ids, ltx_config_file):
         logger.info(f"LTX POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
@@ -89,9 +86,6 @@ class LtxPoolManager:
         self.current_worker_index = 0
         self.lock = threading.Lock()
 
-        # ######################################################################
-        # ##               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:
@@ -99,10 +93,8 @@ class LtxPoolManager:
             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):
-        # ... (Esta função permanece exatamente a mesma) ...
         target_device = worker_to_use.device
         height, width = kwargs['height'], kwargs['width']
         
@@ -118,8 +110,14 @@ class LtxPoolManager:
                 )
 
         first_pass_config = worker_to_use.config.get("first_pass", {})
-        padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
-        padding_vals = calculate_padding(height, width, padded_h, padded_w)
+        
+        # Correção para o modo de refinamento: não recalcular padding
+        if 'latents' in kwargs and kwargs['latents'] is not None:
+            padded_h, padded_w = height, width
+            padding_vals = (0, 0, 0, 0)
+        else:
+            padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
+            padding_vals = calculate_padding(height, width, padded_h, padded_w)
 
         pipeline_params = {
             "height": padded_h, "width": padded_w, 
@@ -148,29 +146,18 @@ class LtxPoolManager:
         
         logger.info("="*60)
         logger.info(f"CHAMADA AO PIPELINE LTX NO DISPOSITIVO: {worker_to_use.device}")
-        logger.info(f"Modelo: {'Distilled' if worker_to_use.is_distilled else 'Base'}")
-        logger.info("-" * 20 + " PARÂMETROS DA PIPELINE " + "-" * 20)
-        logger.info(json.dumps(log_friendly_params, indent=2))
-        logger.info("-" * 20 + " ITENS DE CONDICIONAMENTO " + "-" * 19)
-        logger.info("\n".join(conditioning_log_details) if conditioning_log_details else "  - Nenhum")
-        logger.info("="*60)
+        # ... (resto do logging)
         
         return pipeline_params, padding_vals
     
     def _execute_on_worker(self, execution_fn, **kwargs):
-        """
-        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:
                 worker_to_use = self.workers[self.current_worker_index]
                 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)
-            
-            result = execution_fn(worker_to_use, pipeline_params, **kwargs)
+            result, padding_vals = execution_fn(worker_to_use, **kwargs)
             
             return result, padding_vals
             
@@ -178,37 +165,53 @@ class LtxPoolManager:
             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:
-            # 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"
+        def execution_logic(worker, **inner_kwargs):
+            pipeline_params, padding_vals = self._prepare_and_log_params(worker, **inner_kwargs)
+            pipeline_params['output_type'] = "latent"
             with torch.no_grad():
-                return worker.generate_video_fragment_internal(**params)
+                result_tensor = worker.generate_video_fragment_internal(**pipeline_params)
+            return result_tensor, padding_vals
 
         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 refinamento.
-        """
-        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"
+        def execution_logic(worker, **inner_kwargs):
+            pipeline_params, padding_vals = self._prepare_and_log_params(worker, **inner_kwargs)
+            
+            # --- LÓGICA DE REFINAMENTO EXPLÍCITA E CORRIGIDA ---
+            strength = inner_kwargs.get('denoise_strength', 0.4)
+            num_refine_steps = int(inner_kwargs.get('refine_steps', 10))
             
-            logger.info("LTX POOL MANAGER: Iniciando passe de refinamento (denoise) em latentes de alta resolução.")
+            scheduler = worker.pipeline.scheduler
+            scheduler.set_timesteps(num_refine_steps, device=worker.device)
+            timesteps = scheduler.timesteps
+            
+            start_timestep_idx = int(num_refine_steps * strength)
+            if start_timestep_idx >= len(timesteps):
+                start_timestep_idx = len(timesteps) - 1
+            start_timestep = timesteps[start_timestep_idx]
+            
+            noise = torch.randn_like(upscaled_latents, device=worker.device)
+            noisy_latents = scheduler.add_noise(upscaled_latents.to(worker.device), noise, start_timestep)
+            
+            pipeline_params['latents'] = noisy_latents.to(worker.device, dtype=worker.pipeline.transformer.dtype)
+            pipeline_params['timesteps'] = timesteps[start_timestep_idx:]
+            pipeline_params['num_inference_steps'] = len(pipeline_params['timesteps'])
+            pipeline_params.pop('strength', None)
+            pipeline_params['output_type'] = "latent"
+            
+            logger.info("LTX POOL MANAGER: Iniciando passe de refinamento (denoise) com controle manual de ruído.")
             
             with torch.no_grad():
-                return worker.generate_video_fragment_internal(**params)
+                refined_tensor = worker.generate_video_fragment_internal(**pipeline_params)
+            
+            return refined_tensor, padding_vals
 
         return self._execute_on_worker(execution_logic, upscaled_latents=upscaled_latents, **kwargs)