Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -1,6 +1,15 @@
 # ltx_manager_helpers.py
 # Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
-# (Licenciamento e cabeçalhos permanecem os mesmos)
+#
+# 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.
 
 import torch
 import gc
@@ -111,7 +120,6 @@ class LtxPoolManager:
 
         first_pass_config = worker_to_use.config.get("first_pass", {})
         
-        # 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)
@@ -146,7 +154,6 @@ class LtxPoolManager:
         
         logger.info("="*60)
         logger.info(f"CHAMADA AO PIPELINE LTX NO DISPOSITIVO: {worker_to_use.device}")
-        # ... (resto do logging)
         
         return pipeline_params, padding_vals
     
@@ -184,29 +191,36 @@ class LtxPoolManager:
         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))
+            num_refine_steps_requested = int(inner_kwargs.get('refine_steps', 10))
             
-            scheduler = worker.pipeline.scheduler
-            scheduler.set_timesteps(num_refine_steps, device=worker.device)
-            timesteps = scheduler.timesteps
+            allowed_timesteps = worker.config.get("first_pass", {}).get("timesteps")
+            
+            if allowed_timesteps is None:
+                scheduler = worker.pipeline.scheduler
+                scheduler.set_timesteps(num_refine_steps_requested, device=worker.device)
+                timesteps = scheduler.timesteps
+            else:
+                timesteps = torch.tensor(allowed_timesteps, device=worker.device)
+
+            num_total_timesteps = len(timesteps)
+            start_timestep_idx = int(num_total_timesteps * strength)
+            if start_timestep_idx >= num_total_timesteps:
+                start_timestep_idx = num_total_timesteps - 1
             
-            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)
+            noisy_latents = worker.pipeline.scheduler.add_noise(upscaled_latents.to(worker.device), noise, start_timestep)
             
+            final_timesteps = timesteps[start_timestep_idx:]
             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['timesteps'] = final_timesteps
+            pipeline_params['num_inference_steps'] = len(final_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.")
+            logger.info(f"LTX POOL MANAGER: Iniciando refinamento com {len(final_timesteps)} passos a partir do timestep {start_timestep.item():.4f}.")
             
             with torch.no_grad():
                 refined_tensor = worker.generate_video_fragment_internal(**pipeline_params)