Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -14,6 +14,7 @@ import huggingface_hub
 import time
 import threading
 import json
+from typing import Optional, List
 
 from optimization import optimize_ltx_worker, can_optimize_fp8
 from hardware_manager import hardware_manager
@@ -103,33 +104,88 @@ class LtxPoolManager:
             worker = self.workers[self.current_worker_index]
             self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
             return worker
+    
+    # --- [NOVO] Função centralizada para preparar parâmetros da pipeline ---
+    def _prepare_pipeline_params(self, worker: LtxWorker, **kwargs) -> dict:
+        """Prepara o dicionário de parâmetros para a pipeline, tratando casos especiais como modelos destilados."""
+        
+        # Parâmetros obrigatórios
+        height = kwargs['height']
+        width = kwargs['width']
+        num_frames = kwargs['video_total_frames']
+        
+        # Parâmetros com valores padrão
+        motion_prompt = kwargs.get('motion_prompt', "")
+        negative_prompt = kwargs.get('negative_prompt', "blurry, distorted, static, bad quality")
+        guidance_scale = kwargs.get('guidance_scale', 1.0)
+        stg_scale = kwargs.get('stg_scale', 0.0)
+        rescaling_scale = kwargs.get('rescaling_scale', 0.15)
+        num_inference_steps = kwargs.get('num_inference_steps', 20)
+        
+        # Parâmetros opcionais (para geração ou refinamento)
+        latents_input = kwargs.get('latents')
+        strength = kwargs.get('strength')
+        conditioning_data = kwargs.get('conditioning_items_data')
+
+        # Prepara os itens de condicionamento se existirem
+        final_conditioning_items = []
+        if conditioning_data:
+            for item in conditioning_data:
+                item.latent_tensor = item.latent_tensor.to(worker.device)
+                final_conditioning_items.append(item)
 
+        # Constrói o dicionário base de parâmetros
+        pipeline_params = {
+            "height": height, "width": width, "num_frames": num_frames,
+            "frame_rate": kwargs.get('video_fps', 24),
+            "generator": torch.Generator(device=worker.device).manual_seed(int(time.time()) + kwargs.get('current_fragment_index', 0)),
+            "is_video": True, "vae_per_channel_normalize": True,
+            "prompt": motion_prompt, "negative_prompt": negative_prompt,
+            "guidance_scale": guidance_scale, "stg_scale": stg_scale,
+            "rescaling_scale": rescaling_scale, "num_inference_steps": num_inference_steps,
+            "output_type": "latent"
+        }
+
+        # Adiciona parâmetros opcionais se eles foram fornecidos
+        if latents_input is not None:
+            pipeline_params["latents"] = latents_input.to(worker.device, dtype=worker.pipeline.transformer.dtype)
+        if strength is not None:
+            pipeline_params["strength"] = strength
+        if final_conditioning_items:
+            pipeline_params["conditioning_items"] = final_conditioning_items
+
+        # --- LÓGICA CENTRALIZADA E À PROVA DE ERRO ---
+        # Se o modelo for destilado, sobrescreve os passos com os timesteps fixos obrigatórios.
+        if worker.is_distilled:
+            logger.info(f"Worker {worker.device} está usando um modelo destilado. Usando timesteps fixos.")
+            fixed_timesteps = worker.config.get("first_pass", {}).get("timesteps")
+            pipeline_params["timesteps"] = fixed_timesteps
+            if fixed_timesteps:
+                pipeline_params["num_inference_steps"] = len(fixed_timesteps)
+
+        # Log dos parâmetros para depuração
+        log_params = {k: v for k, v in pipeline_params.items() if k not in ['generator', 'latents', 'conditioning_items']}
+        logger.info(f"Parâmetros preparados para a pipeline em {worker.device}:\n{json.dumps(log_params, indent=2)}")
+
+        return pipeline_params
+
+    # --- [REATORADO] Função de Geração simplificada ---
     def generate_latent_fragment(self, **kwargs) -> (torch.Tensor, tuple):
         worker_to_use = self._get_next_worker()
         try:
+            # Padding
             height, width = kwargs['height'], kwargs['width']
             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_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)
-
-            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": 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'],
-                "output_type": "latent"
-            }
-            if worker_to_use.is_distilled:
-                pipeline_params["timesteps"] = worker_to_use.config.get("first_pass", {}).get("timesteps")
-                pipeline_params["num_inference_steps"] = len(pipeline_params["timesteps"]) if pipeline_params["timesteps"] else 20
+            # Atualiza kwargs com as dimensões com padding
+            kwargs['height'] = padded_h
+            kwargs['width'] = padded_w
 
+            # Prepara os parâmetros usando a função centralizada
+            pipeline_params = self._prepare_pipeline_params(worker_to_use, **kwargs)
+            
+            # Executa a geração
             if isinstance(worker_to_use.pipeline, LTXMultiScalePipeline):
                 result = worker_to_use.pipeline.video_pipeline(**pipeline_params).images
             else:
@@ -144,35 +200,30 @@ class LtxPoolManager:
                 with torch.cuda.device(worker_to_use.device):
                     gc.collect(); torch.cuda.empty_cache()
 
+    # --- [REATORADO] Função de Refinamento simplificada ---
     def refine_latents(self, latents_to_refine: torch.Tensor, **kwargs) -> (torch.Tensor, tuple):
         worker_to_use = self._get_next_worker()
         try:
-            height, width, num_frames = kwargs['height'], kwargs['width'], kwargs['video_total_frames']
+            # Adiciona os tensores e a força de denoise aos kwargs para a função auxiliar
+            kwargs['latents'] = latents_to_refine
+            kwargs['strength'] = kwargs.get('denoise_strength', 0.4)
+            kwargs['num_inference_steps'] = int(kwargs.get('refine_steps', 10))
             
-            pipeline_params = {
-                "latents": latents_to_refine.to(worker_to_use.device, dtype=worker_to_use.pipeline.transformer.dtype),
-                "height": height, "width": width, "num_frames": num_frames, "frame_rate": kwargs['video_fps'],
-                "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),
-                "num_inference_steps": int(kwargs.get('refine_steps', 10)),
-                "strength": kwargs.get('denoise_strength', 0.4),
-                "output_type": "latent"
-            }
+            # Prepara os parâmetros usando a mesma função centralizada
+            pipeline_params = self._prepare_pipeline_params(worker_to_use, **kwargs)
 
             logger.info("LTX POOL MANAGER: Iniciando passe de refinamento (denoise)...")
             
             pipeline_to_call = worker_to_use.pipeline.video_pipeline if isinstance(worker_to_use.pipeline, LTXMultiScalePipeline) else worker_to_use.pipeline
 
-            try:
-                result = pipeline_to_call(**pipeline_params).images
-                return result, None
-            except torch.cuda.OutOfMemoryError as e:
-                logger.error(f"FALHA DE MEMÓRIA DURANTE O REFINAMENTO em {worker_to_use.device}: {e}")
-                logger.warning("Limpando VRAM e retornando None para sinalizar a falha.")
-                gc.collect(); torch.cuda.empty_cache()
-                return None, None
+            result = pipeline_to_call(**pipeline_params).images
+            return result, None
+        
+        except torch.cuda.OutOfMemoryError as e:
+            logger.error(f"FALHA DE MEMÓRIA DURANTE O REFINAMENTO em {worker_to_use.device}: {e}")
+            logger.warning("Limpando VRAM e retornando None para sinalizar a falha.")
+            gc.collect(); torch.cuda.empty_cache()
+            return None, None
         except Exception as e:
             logger.error(f"LTX POOL MANAGER: Erro inesperado durante o refinamento em {worker_to_use.device}: {e}", exc_info=True)
             raise e