Update api/gpu_manager.py

api/gpu_manager.py

@@ -1,56 +1,94 @@
-# api/gpu_manager.py
+# api/gpu_manager.py (Versão com VAE Dedicado)
 
 import os
 import torch
+import math
 
 class GPUManager:
     """
-    Gerencia e aloca GPUs disponíveis para diferentes serviços (LTX, SeedVR).
+    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_gpus = []
+        self.ltx_main_gpus = []
+        self.ltx_vae_gpu = []
         self.seedvr_gpus = []
+        self.vincie_gpus = []
         self._allocate_gpus()
 
     def _allocate_gpus(self):
-        """
-        Divide as GPUs disponíveis entre os serviços LTX e SeedVR.
-        """
-        print("="*50)
-        print("🤖 Gerenciador de GPUs inicializado.")
+        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 detectada. Modo de compartilhamento de memória será usado.")
-            # Ambos usarão a GPU 0, mas precisarão gerenciar a memória
-            self.ltx_gpus = [0]
+            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]
-        else:
-            # Divide as GPUs entre os dois serviços
-            mid_point = self.total_gpus // 2
-            self.ltx_gpus = list(range(0, mid_point))
-            self.seedvr_gpus = list(range(mid_point, self.total_gpus))
-            print(f"   > Alocação: LTX usará GPUs {self.ltx_gpus}, SeedVR usará GPUs {self.seedvr_gpus}.")
+            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("="*50)
+        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_device(self):
-        """Retorna o dispositivo principal para o LTX."""
-        if not self.ltx_gpus:
+    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")
-        # Por padrão, o modelo principal do LTX roda na primeira GPU do seu grupo
-        return torch.device(f"cuda:{self.ltx_gpus[0]}")
+        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 é necessário mover modelos entre CPU e GPU."""
-        return self.total_gpus < 2
+        """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 para ser importada por outros módulos
+# Instância global
 gpu_manager = GPUManager()