Rename api/ltx/ltx_pool_manager.py to api/ltx/ltx_aduc_manager.py

api/ltx/{ltx_pool_manager.py → ltx_aduc_manager.py}

@@ -1,4 +1,4 @@
-# FILE: api/ltx/ltx_pool_manager.py
+# FILE: api/ltx/ltx_aduc_manager.py
 # DESCRIPTION: The "secret weapon". A pool manager for LTX that applies
 # runtime patches to the pipeline for full control and ADUC-SDR compatibility.
 
@@ -9,7 +9,7 @@ import torch
 from diffusers.utils.torch_utils import randn_tensor
 
 # --- Importações da nossa arquitetura ---
-from manager.gpu_manager import gpu_manager
+from managers.gpu_manager import gpu_manager
 from api.ltx.ltx_utils import build_ltx_pipeline_on_cpu
 
 def add_deps_to_path():
@@ -43,10 +43,9 @@ class LatentConditioningItem:
 # --- O MONKEY PATCH ---
 # Esta é a nossa versão customizada de `prepare_conditioning`
 # ==============================================================================
-
 def _aduc_prepare_conditioning_patch(
     self: "LTXVideoPipeline",
-    conditioning_items: Optional[List[Union[ConditioningItem, LatentConditioningItem]]],
+    conditioning_items: Optional[List[Union["ConditioningItem", "LatentConditioningItem"]]],
     init_latents: torch.Tensor,
     num_frames: int,
     height: int,
@@ -54,13 +53,52 @@ def _aduc_prepare_conditioning_patch(
     vae_per_channel_normalize: bool = False,
     generator=None,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    if not conditioning_items:
+        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)
+        return init_latents, init_pixel_coords, None, 0
     
-    # Esta função é uma cópia modificada da sua, com logging e pequenas melhorias.
+    init_conditioning_mask = torch.zeros_like(init_latents[:, 0, ...], dtype=torch.float32, device=init_latents.device)
+    extra_conditioning_latents, extra_conditioning_pixel_coords, extra_conditioning_mask = [], [], []
+    extra_conditioning_num_latents = 0
+    
+    for item in conditioning_items:
+        if not isinstance(item, LatentConditioningItem):
+            logger.warning("Patch ADUC: Item de condicionamento não é um LatentConditioningItem e será ignorado.")
+            continue
+        
+        media_item_latents = item.latent_tensor.to(dtype=init_latents.dtype, device=init_latents.device)
+        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)
+            init_conditioning_mask[..., :f_l, :h_l, :w_l] = strength
+        else:
+            noise = randn_tensor(media_item_latents.shape, generator=generator, device=media_item_latents.device, dtype=media_item_latents.dtype)
+            media_item_latents = torch.lerp(noise, media_item_latents, strength)
+            patched_latents, latent_coords = self.patchifier.patchify(latents=media_item_latents)
+            pixel_coords = latent_to_pixel_coords(latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
+            pixel_coords[:, 0] += media_frame_number
+            extra_conditioning_num_latents += patched_latents.shape[1]
+            new_mask = torch.full(patched_latents.shape[:2], strength, dtype=torch.float32, device=init_latents.device)
+            extra_conditioning_latents.append(patched_latents)
+            extra_conditioning_pixel_coords.append(pixel_coords)
+            extra_conditioning_mask.append(new_mask)
+
+    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))
+    init_conditioning_mask = init_conditioning_mask.squeeze(-1)
+    
+    if extra_conditioning_latents:
+        init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
+        init_pixel_coords = torch.cat([*extra_conditioning_pixel_coords, init_pixel_coords], dim=2)
+        init_conditioning_mask = torch.cat([*extra_conditioning_mask, init_conditioning_mask], dim=1)
     
-    # (O código do patch que você forneceu vai aqui, ligeiramente ajustado)
-    # ...
-
     return init_latents, init_pixel_coords, init_conditioning_mask, extra_conditioning_num_latents
+    
+
 
 
 # ==============================================================================