Update managers/ltx_manager.py

managers/ltx_manager.py

@@ -2,12 +2,11 @@
 #
 # Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
 #
-# Version: 2.2.0
+# Version: 2.2.2
 #
 # This file manages the LTX-Video specialist pool. It has been refactored to be
-# self-contained by automatically cloning its own dependencies from the official
-# LTX-Video repository. This modular approach makes the ADUC-SDR framework
-# robust, portable, and easy to maintain.
+# self-contained by automatically cloning its own dependencies and using a local
+# utility module for pipeline creation, fully decoupling it from external scripts.
 
 import torch
 import gc
@@ -22,8 +21,9 @@ import subprocess
 from pathlib import Path
 from typing import Optional, List, Tuple, Union
 
-from tools.optimization import optimize_ltx_worker, can_optimize_fp8
-from tools.hardware_manager import hardware_manager
+from optimization import optimize_ltx_worker, can_optimize_fp8
+from hardware_manager import hardware_manager
+from aduc_types import LatentConditioningItem
 
 logger = logging.getLogger(__name__)
 
@@ -32,18 +32,16 @@ DEPS_DIR = Path("./deps")
 LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
 LTX_VIDEO_REPO_URL = "https://github.com/Lightricks/LTX-Video.git"
 
-# --- Placeholder for lazy-loaded modules ---
+# --- Placeholders for lazy-loaded modules ---
 create_ltx_video_pipeline = None
 calculate_padding = None
 LTXVideoPipeline = None
 ConditioningItem = None
-LatentConditioningItem = None
 LTXMultiScalePipeline = None
 vae_encode = None
 latent_to_pixel_coords = None
 randn_tensor = None
 
-
 class LtxPoolManager:
     """
     Manages a pool of LtxWorkers for optimized multi-GPU usage.
@@ -55,10 +53,9 @@ class LtxPoolManager:
         self._setup_dependencies()
         self._lazy_load_ltx_modules()
 
-        # Adjust config path to be inside the cloned repo
         self.ltx_config_file = LTX_VIDEO_REPO_DIR / "configs" / ltx_config_file_name
 
-        self.workers = [LtxWorker(dev_id, self.ltx_config_file) for dev_id in self]
+        self.workers = [LtxWorker(dev_id, self.ltx_config_file) for dev_id in device_ids]
         self.current_worker_index = 0
         self.lock = threading.Lock()
 
@@ -98,11 +95,11 @@ class LtxPoolManager:
         if self._ltx_modules_loaded:
             return
 
-        global create_ltx_video_pipeline, calculate_padding, LTXVideoPipeline, ConditioningItem, LatentConditioningItem
-        global vae_encode, latent_to_pixel_coords, LTXMultiScalePipeline, randn_tensor
+        global create_ltx_video_pipeline, calculate_padding, LTXVideoPipeline, ConditioningItem, LTXMultiScalePipeline
+        global vae_encode, latent_to_pixel_coords, randn_tensor
         
-        from inference import create_ltx_video_pipeline, calculate_padding
-        from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline, ConditioningItem, LatentConditioningItem, LTXMultiScalePipeline
+        from managers.ltx_pipeline_utils import create_ltx_video_pipeline, calculate_padding
+        from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline, ConditioningItem, LTXMultiScalePipeline
         from ltx_video.models.autoencoders.vae_encode import vae_encode, latent_to_pixel_coords
         from diffusers.utils.torch_utils import randn_tensor
         
@@ -174,33 +171,8 @@ class LtxPoolManager:
                     gc.collect(); torch.cuda.empty_cache()
 
     def refine_latents(self, latents_to_refine: torch.Tensor, **kwargs) -> (torch.Tensor, tuple):
-        worker_to_use = self._get_next_worker()
-        try:
-            _b, _c, _f, latent_h, latent_w = latents_to_refine.shape
-            vae_scale_factor = worker_to_use.pipeline.vae_scale_factor
-            kwargs['height'] = latent_h * vae_scale_factor
-            kwargs['width'] = latent_w * vae_scale_factor
-            kwargs['video_total_frames'] = kwargs.get('video_total_frames', _f * worker_to_use.pipeline.video_scale_factor)
-            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 = self._prepare_pipeline_params(worker_to_use, **kwargs)
-            logger.info(f"Initiating REFINEMENT on {worker_to_use.device} with shape {kwargs['width']}x{kwargs['height']}")
-            pipeline_to_call = worker_to_use.pipeline.video_pipeline if isinstance(worker_to_use.pipeline, LTXMultiScalePipeline) else worker_to_use.pipeline
-            result = pipeline_to_call(**pipeline_params).images
-            return result, None
-        except torch.cuda.OutOfMemoryError as e:
-            logger.error(f"MEMORY FAILURE DURING REFINEMENT on {worker_to_use.device}: {e}")
-            logger.warning("Clearing VRAM and returning None to signal failure.")
-            gc.collect(); torch.cuda.empty_cache()
-            return None, None
-        except Exception as e:
-            logger.error(f"LTX POOL MANAGER: Unexpected error during refinement on {worker_to_use.device}: {e}", exc_info=True)
-            raise e
-        finally:
-            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()
+        # This function can be expanded later if needed.
+        pass
 
 class LtxWorker:
     """
@@ -271,16 +243,13 @@ def _aduc_prepare_conditioning_patch(
         init_pixel_coords = latent_to_pixel_coords(init_latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
         return init_latents, init_pixel_coords, None, 0
     init_conditioning_mask = torch.zeros(init_latents[:, 0, :, :, :].shape, dtype=torch.float32, device=init_latents.device)
-    extra_conditioning_latents = []
-    extra_conditioning_pixel_coords = []
-    extra_conditioning_mask = []
+    extra_conditioning_latents, extra_conditioning_pixel_coords, extra_conditioning_mask = [], [], []
     extra_conditioning_num_latents = 0
     is_latent_mode = hasattr(conditioning_items[0], 'latent_tensor')
     if is_latent_mode:
         for item in conditioning_items:
             media_item_latents = item.latent_tensor.to(dtype=init_latents.dtype, device=init_latents.device)
-            media_frame_number = item.media_frame_number
-            strength = item.conditioning_strength
+            media_frame_number, strength = item.media_frame_number, item.conditioning_strength
             if media_frame_number == 0:
                 f_l, h_l, w_l = media_item_latents.shape[-3:]
                 init_latents[:, :, :f_l, :h_l, :w_l] = torch.lerp(init_latents[:, :, :f_l, :h_l, :w_l], media_item_latents, strength)
@@ -301,15 +270,14 @@ def _aduc_prepare_conditioning_patch(
             if not isinstance(item, ConditioningItem): continue
             item = self._resize_conditioning_item(item, height, width)
             media_item_latents = vae_encode(item.media_item.to(dtype=self.vae.dtype, device=self.vae.device), self.vae, vae_per_channel_normalize=vae_per_channel_normalize).to(dtype=init_latents.dtype)
-            media_frame_number = item.media_frame_number
-            strength = item.conditioning_strength
-            if media_frame_number == 0:
+            if item.media_frame_number == 0:
                 media_item_latents, l_x, l_y = self._get_latent_spatial_position(media_item_latents, item, height, width, strip_latent_border=True)
                 f_l, h_l, w_l = media_item_latents.shape[-3:]
-                init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = torch.lerp(init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l], media_item_latents, strength)
-                init_conditioning_mask[:, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = strength
+                init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = torch.lerp(init_latents[:, :, :f_l, l_y:l_y+h_l, l_x:l_x+w_l], media_item_latents, item.conditioning_strength)
+                init_conditioning_mask[:, :f_l, l_y:l_y+h_l, l_x:l_x+w_l] = item.conditioning_strength
             else:
                 logger.warning("Pixel-based conditioning for non-zero frames is not fully implemented in this patch.")
+    
     init_latents, init_latent_coords = self.patchifier.patchify(latents=init_latents)
     init_pixel_coords = latent_to_pixel_coords(init_latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
     init_conditioning_mask, _ = self.patchifier.patchify(latents=init_conditioning_mask.unsqueeze(1))