Upload gpu_manager.py

api/gpu_manager.py

@@ -0,0 +1,94 @@
+# api/gpu_manager.py (Versão com VAE Dedicado)
+
+import os
+import torch
+import math
+
+class GPUManager:
+    """
+    Gerencia e aloca GPUs de forma inteligente entre LTX (com VAE dedicado),
+    SeedVR e VINCIE.
+    """
+    def __init__(self):
+        self.total_gpus = torch.cuda.device_count()
+        self.ltx_main_gpus = []
+        self.ltx_vae_gpu = []
+        self.seedvr_gpus = []
+        self.vincie_gpus = []
+        self._allocate_gpus()
+
+    def _allocate_gpus(self):
+        print("="*60)
+        print("🤖 Gerenciador de GPUs (com VAE dedicado para LTX)")
+        print(f"   > Total de GPUs detectadas: {self.total_gpus}")
+
+        all_indices = list(range(self.total_gpus))
+
+        if self.total_gpus == 0:
+            print("   > Nenhuma GPU detectada. Operando em modo CPU.")
+        elif self.total_gpus == 1:
+            print("   > 1 GPU: Modo de compartilhamento total.")
+            self.ltx_main_gpus = [0]
+            self.ltx_vae_gpu = [0] # Compartilha com o principal
+            self.seedvr_gpus = [0]
+            self.vincie_gpus = [0]
+        elif self.total_gpus == 2:
+            print("   > 2 GPUs: LTX com VAE dedicado, outros compartilham a GPU principal.")
+            self.ltx_main_gpus = [0]
+            self.ltx_vae_gpu = [1] # VAE fica com a segunda GPU
+            self.seedvr_gpus = [0] # Compartilha com LTX principal
+            self.vincie_gpus = [0] # Compartilha com LTX principal
+        else: # 3 ou mais GPUs
+            print(f"   > {self.total_gpus} GPUs: Alocação distribuída.")
+            # LTX sempre fica com as duas primeiras GPUs se disponíveis
+            self.ltx_main_gpus = [0]
+            self.ltx_vae_gpu = [1]
+            
+            remaining_gpus = all_indices[2:]
+            
+            if not remaining_gpus: # Caso de exatamente 2 GPUs, já coberto, mas por segurança
+                self.seedvr_gpus = [0]
+                self.vincie_gpus = [0]
+            else:
+                # O resto é dividido entre SeedVR e VINCIE
+                vincie_count = max(1, math.ceil(len(remaining_gpus) / 2))
+                seedvr_count = len(remaining_gpus) - vincie_count
+                
+                self.vincie_gpus = remaining_gpus[:vincie_count]
+                self.seedvr_gpus = remaining_gpus[vincie_count:]
+        
+        print(f"   > Alocação Final:")
+        print(f"     - LTX (Transformer): GPUs {self.ltx_main_gpus}")
+        print(f"     - LTX (VAE):         GPUs {self.ltx_vae_gpu}")
+        print(f"     - SeedVR:            GPUs {self.seedvr_gpus}")
+        print(f"     - VINCIE:            GPUs {self.vincie_gpus}")
+        print("="*60)
+
+    def get_ltx_device(self) -> torch.device:
+        """Retorna o dispositivo principal para o Transformer do LTX."""
+        if not self.ltx_main_gpus:
+            return torch.device("cpu")
+        return torch.device(f"cuda:{self.ltx_main_gpus[0]}")
+
+    def get_ltx_vae_device(self) -> torch.device:
+        """Retorna o dispositivo dedicado para o VAE do LTX."""
+        if not self.ltx_vae_gpu:
+            return torch.device("cpu")
+        return torch.device(f"cuda:{self.ltx_vae_gpu[0]}")
+
+    def get_seedvr_devices(self) -> list:
+        """Retorna a lista de IDs de GPU para o SeedVR."""
+        return self.seedvr_gpus
+    
+    def get_vincie_devices(self) -> list:
+        """Retorna a lista de IDs de GPU para o VINCIE."""
+        return self.vincie_gpus
+
+    def requires_memory_swap(self) -> bool:
+        """Verifica se múltiplos serviços estão compartilhando a mesma GPU."""
+        all_allocations = self.ltx_main_gpus + self.seedvr_gpus + self.vincie_gpus
+        # O swap é necessário se o número de alocações for maior que o número de GPUs únicas
+        return len(all_allocations) > len(set(all_allocations))
+
+# Instância global
+gpu_manager = GPUManager()