Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -1,8 +1,125 @@
 # ltx_manager_helpers.py
-# ... (imports e outras classes) ...
+# Gerente de Pool de Workers LTX para revezamento assíncrono em múltiplas GPUs.
+# Este arquivo é parte do projeto Euia-AducSdr e está sob a licença AGPL v3.
+# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
+
+import torch
+import gc
+import os
+import yaml
+import numpy as np
+import imageio
+from pathlib import Path
+import huggingface_hub
+import threading
+from PIL import Image
+
+# Importa as funções e classes necessárias do inference.py
+from inference import (
+    create_ltx_video_pipeline,
+    create_latent_upsampler,
+    ConditioningItem,
+    calculate_padding,
+    prepare_conditioning
+)
+from ltx_video.pipelines.pipeline_ltx_video import LTXMultiScalePipeline
+
+class LtxWorker:
+    """
+    Representa uma única instância do pipeline LTX, associada a uma GPU específica.
+    O pipeline é carregado na CPU por padrão e movido para a GPU sob demanda.
+    """
+    def __init__(self, device_id='cuda:0'):
+        self.device = torch.device(device_id if torch.cuda.is_available() else 'cpu')
+        print(f"LTX Worker: Inicializando para o dispositivo {self.device} (carregando na CPU)...")
+        
+        config_file_path = "configs/ltxv-13b-0.9.8-distilled.yaml"
+        with open(config_file_path, "r") as file:
+            self.config = yaml.safe_load(file)
+
+        LTX_REPO = "Lightricks/LTX-Video"
+        models_dir = "downloaded_models_gradio"
+        
+        model_actual_path = huggingface_hub.hf_hub_download(
+            repo_id=LTX_REPO,
+            filename=self.config["checkpoint_path"],
+            local_dir=models_dir,
+            local_dir_use_symlinks=False
+        )
+        
+        self.pipeline = create_ltx_video_pipeline(
+            ckpt_path=model_actual_path,
+            precision=self.config["precision"],
+            text_encoder_model_name_or_path=self.config["text_encoder_model_name_or_path"],
+            sampler=self.config["sampler"],
+            device='cpu'
+        )
+
+        print(f"LTX Worker para {self.device}: Compilando o transformer (isso pode levar um momento)...")
+        self.pipeline.transformer.to(memory_format=torch.channels_last)
+        try:
+            self.pipeline.transformer = torch.compile(self.pipeline.transformer, mode="reduce-overhead", fullgraph=True)
+            print(f"LTX Worker para {self.device}: Transformer compilado com sucesso.")
+        except Exception as e:
+            print(f"AVISO: A compilação do Transformer falhou em {self.device}: {e}. Continuando sem compilação.")
+
+        self.latent_upsampler = None
+        if self.config.get("pipeline_type") == "multi-scale":
+            print(f"LTX Worker para {self.device}: Carregando Latent Upsampler (Multi-Scale)...")
+            upscaler_path = huggingface_hub.hf_hub_download(
+                repo_id=LTX_REPO,
+                filename=self.config["spatial_upscaler_model_path"],
+                local_dir=models_dir,
+                local_dir_use_symlinks=False
+            )
+            self.latent_upsampler = create_latent_upsampler(upscaler_path, 'cpu')
+
+        print(f"LTX Worker para {self.device} pronto na CPU.")
+
+    def to_gpu(self):
+        """Move o pipeline e o upsampler para a GPU designada."""
+        if self.device.type == 'cpu': return
+        print(f"LTX Worker: Movendo pipeline para {self.device}...")
+        self.pipeline.to(self.device)
+        if self.latent_upsampler:
+            print(f"LTX Worker: Movendo Latent Upsampler para {self.device}...")
+            self.latent_upsampler.to(self.device)
+        print(f"LTX Worker: Pipeline na GPU {self.device}.")
+
+    def to_cpu(self):
+        """Move o pipeline de volta para a CPU e limpa a memória da GPU."""
+        if self.device.type == 'cpu': return
+        print(f"LTX Worker: Descarregando pipeline da GPU {self.device}...")
+        self.pipeline.to('cpu')
+        if self.latent_upsampler:
+            self.latent_upsampler.to('cpu')
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print(f"LTX Worker: GPU {self.device} limpa.")
+    
+    def generate_video_fragment_internal(self, **kwargs):
+        """A lógica real da geração de vídeo, que espera estar na GPU."""
+        return self.pipeline(**kwargs)
 
 class LtxPoolManager:
-    # ... (__init__, _cleanup_worker) ...
+    """
+    Gerencia um pool de LtxWorkers, orquestrando um revezamento entre GPUs
+    para permitir que a limpeza de uma GPU ocorra em paralelo com a computação em outra.
+    """
+    def __init__(self, device_ids=['cuda:2', 'cuda:3']):
+        print(f"LTX POOL MANAGER: Criando workers para os dispositivos: {device_ids}")
+        self.workers = [LtxWorker(device_id) for device_id in device_ids]
+        self.current_worker_index = 0
+        self.lock = threading.Lock()
+        self.last_cleanup_thread = None
+
+    def _cleanup_worker(self, worker):
+        """Função alvo para a thread de limpeza."""
+        print(f"CLEANUP THREAD: Iniciando limpeza da GPU {worker.device} em background...")
+        worker.to_cpu()
+        print(f"CLEANUP THREAD: Limpeza da GPU {worker.device} concluída.")
+
     def generate_video_fragment(
         self,
         motion_prompt: str, conditioning_items_data: list,
@@ -13,10 +130,44 @@ class LtxPoolManager:
     ):
         worker_to_use = None
         try:
-            # ... (lógica de seleção e limpeza de worker) ...
+            with self.lock:
+                if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
+                    print("LTX POOL MANAGER: Aguardando limpeza da GPU anterior...")
+                    self.last_cleanup_thread.join()
+
+                worker_to_use = self.workers[self.current_worker_index]
+                previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
+                worker_to_cleanup = self.workers[previous_worker_index]
+
+                cleanup_thread = threading.Thread(target=self._cleanup_worker, args=(worker_to_cleanup,))
+                cleanup_thread.start()
+                self.last_cleanup_thread = cleanup_thread
+                
+                worker_to_use.to_gpu()
+                
+                self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
+            
+            target_device = worker_to_use.device
             
-            # ... (lógica de preparação de padding e conditioning_items) ...
+            if use_attention_slicing:
+                worker_to_use.pipeline.enable_attention_slicing()
+
+            media_paths = [item[0] for item in conditioning_items_data]
+            start_frames = [item[1] for item in conditioning_items_data]
+            strengths = [item[2] for item in conditioning_items_data]
+
+            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 = prepare_conditioning(
+                conditioning_media_paths=media_paths, conditioning_strengths=strengths,
+                conditioning_start_frames=start_frames, height=height, width=width,
+                num_frames=video_total_frames, padding=padding_vals, pipeline=worker_to_use.pipeline,
+            )
             
+            for item in conditioning_items:
+                item.media_item = item.media_item.to(target_device)
+
             kwargs = {
                 "prompt": motion_prompt,
                 "negative_prompt": "blurry, distorted, bad quality, artifacts",
@@ -38,28 +189,36 @@ class LtxPoolManager:
                 "enhance_prompt": False,
             }
 
-            # --- CORREÇÃO AQUI ---
             # Verifica se o config do modelo especifica uma lista de timesteps.
             # Se sim, usa essa lista. Se não, usa o num_inference_steps da UI.
             first_pass_config = worker_to_use.config.get("first_pass", {})
             if "timesteps" in first_pass_config:
                 print("Usando timesteps customizados do arquivo de configuração para o modelo distilled.")
                 kwargs["timesteps"] = first_pass_config["timesteps"]
-                # Quando usamos timesteps customizados, o num_inference_steps é inferido
-                # a partir do tamanho da lista de timesteps.
                 kwargs["num_inference_steps"] = len(first_pass_config["timesteps"])
+                # Para modelos distilled, a UI de steps é ignorada, mas outros params do config são usados
+                kwargs.update({k: v for k, v in first_pass_config.items() if k != "timesteps"})
             else:
-                # Comportamento antigo para modelos não-destilados
                 print(f"Usando num_inference_steps da UI: {num_inference_steps}")
                 kwargs["num_inference_steps"] = int(num_inference_steps)
-            # --- FIM DA CORREÇÃO ---
-
+            
             progress(0.1, desc=f"[Câmera LTX em {worker_to_use.device}] Filmando Cena {current_fragment_index}...")
             result_tensor = worker_to_use.generate_video_fragment_internal(**kwargs).images
             
-            # ... (lógica de cropping e salvamento) ...
+            pad_l, pad_r, pad_t, pad_b = map(int, padding_vals)
+            slice_h = -pad_b if pad_b > 0 else None
+            slice_w = -pad_r if pad_r > 0 else None
+            cropped_tensor = result_tensor[:, :, :video_total_frames, pad_t:slice_h, pad_l:slice_w]
+            video_np = (cropped_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
+            
+            with imageio.get_writer(output_path, fps=video_fps, codec='libx264', quality=8) as writer:
+                for frame in video_np:
+                    writer.append_data(frame)
             
             return output_path, video_total_frames
 
         finally:
-            # ...
+            if use_attention_slicing and worker_to_use and worker_to_use.pipeline:
+                worker_to_use.pipeline.disable_attention_slicing()
+
+ltx_manager_singleton = LtxPoolManager(device_ids=['cuda:2', 'cuda:3'])