Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -1,5 +1,9 @@
-# ltx_manager_helpers.py (Com Lógica de Refinamento Especializada)
+# ltx_manager_helpers.py
 # Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
+#
+# Este programa é software livre: você pode redistribuí-lo e/ou modificá-lo
+# sob os termos da Licença Pública Geral Affero GNU...
+# AVISO DE PATENTE PENDENTE: Consulte NOTICE.md.
 
 import torch
 import gc
@@ -19,6 +23,10 @@ from ltx_video.pipelines.pipeline_ltx_video import LatentConditioningItem
 logger = logging.getLogger(__name__)
 
 class LtxWorker:
+    """
+    Representa uma única instância da pipeline LTX-Video em um dispositivo específico.
+    Gerencia o carregamento do modelo para a CPU e a movimentação de/para a GPU.
+    """
     def __init__(self, device_id, ltx_config_file):
         self.cpu_device = torch.device('cpu')
         self.device = torch.device(device_id if torch.cuda.is_available() else 'cpu')
@@ -28,6 +36,7 @@ class LtxWorker:
             self.config = yaml.safe_load(file)
         
         self.is_distilled = "distilled" in self.config.get("checkpoint_path", "")
+
         models_dir = "downloaded_models_gradio"
         
         logger.info(f"LTX Worker ({self.device}): Carregando modelo para a CPU...")
@@ -46,9 +55,11 @@ class LtxWorker:
         logger.info(f"LTX Worker ({self.device}): Modelo pronto na CPU. É um modelo destilado? {self.is_distilled}")
 
     def to_gpu(self):
+        """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)
+        
         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)
@@ -57,6 +68,7 @@ class LtxWorker:
             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."""
         if self.device.type == 'cpu': return
         logger.info(f"LTX Worker: Descarregando pipeline da GPU {self.device}...")
         self.pipeline.to('cpu')
@@ -64,14 +76,20 @@ class LtxWorker:
         if torch.cuda.is_available(): torch.cuda.empty_cache()
 
     def generate_video_fragment_internal(self, **kwargs):
+        """Invoca a pipeline de geração."""
         return self.pipeline(**kwargs).images
 
 class LtxPoolManager:
+    """
+    Gerencia um pool de LtxWorkers para otimizar o uso de múltiplas GPUs.
+    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()
+
         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:
@@ -93,12 +111,19 @@ class LtxPoolManager:
             padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
             padding_vals = calculate_padding(height, width, padded_h, padded_w)
             
-            conditioning_items = [item.to(worker_to_use.device) for item in kwargs.get('conditioning_items_data', [])]
+            # --- [INÍCIO DA CORREÇÃO] ---
+            # Move o tensor DENTRO de cada item de condicionamento para o dispositivo do worker.
+            conditioning_data = kwargs.get('conditioning_items_data', [])
+            final_conditioning_items = []
+            for item in conditioning_data:
+                item.latent_tensor = item.latent_tensor.to(worker_to_use.device)
+                final_conditioning_items.append(item)
+            # --- [FIM DA CORREÇÃO] ---
 
             pipeline_params = {
                 "height": padded_h, "width": padded_w, "num_frames": kwargs['video_total_frames'], 
                 "frame_rate": kwargs['video_fps'], "generator": torch.Generator(device=worker_to_use.device).manual_seed(int(time.time()) + kwargs['current_fragment_index']),
-                "conditioning_items": conditioning_items, "is_video": True, "vae_per_channel_normalize": True,
+                "conditioning_items": final_conditioning_items, "is_video": True, "vae_per_channel_normalize": True,
                 "prompt": kwargs['motion_prompt'], "negative_prompt": "blurry, distorted, static, bad quality",
                 "guidance_scale": kwargs['guidance_scale'], "stg_scale": kwargs['stg_scale'], 
                 "rescaling_scale": kwargs['rescaling_scale'], "num_inference_steps": kwargs['num_inference_steps']
@@ -120,9 +145,6 @@ class LtxPoolManager:
     def refine_latents(self, latents_to_refine: torch.Tensor, **kwargs) -> (torch.Tensor, tuple):
         worker_to_use = self._get_next_worker()
         try:
-            # --- [INÍCIO DA CORREÇÃO] ---
-            # Para refinamento, as dimensões são derivadas DIRETAMENTE do tensor latente.
-            # Não há padding. A resolução em pixels é passada, mas a forma latente é a fonte da verdade.
             height, width, num_frames = kwargs['height'], kwargs['width'], kwargs['video_total_frames']
             
             pipeline_params = {
@@ -131,16 +153,15 @@ class LtxPoolManager:
                 "generator": torch.Generator(device=worker_to_use.device).manual_seed(int(time.time()) + kwargs['current_fragment_index']),
                 "is_video": True, "vae_per_channel_normalize": True,
                 "prompt": kwargs['motion_prompt'], "negative_prompt": "blurry, distorted, static, bad quality",
-                "guidance_scale": kwargs.get('guidance_scale', 1.0), # Força 1.0 para refinamento incondicional se não especificado
+                "guidance_scale": kwargs.get('guidance_scale', 1.0),
                 "num_inference_steps": int(kwargs.get('refine_steps', 10)),
                 "strength": kwargs.get('denoise_strength', 0.4),
                 "output_type": "latent"
             }
-            # --- [FIM DA CORREÇÃO] ---
 
             logger.info("LTX POOL MANAGER: Iniciando passe de refinamento (denoise)...")
             result = worker_to_use.generate_video_fragment_internal(**pipeline_params)
-            return result, None # Nenhum padding é aplicado no refinamento
+            return result, None
         except Exception as e:
             logger.error(f"LTX POOL MANAGER: Erro durante o refinamento em {worker_to_use.device}: {e}", exc_info=True)
             raise e