Create gpu_manager.py

api/gpu_manager.py

@@ -0,0 +1,56 @@
+# api/gpu_manager.py
+
+import os
+import torch
+
+class GPUManager:
+    """
+    Gerencia e aloca GPUs disponíveis para diferentes serviços (LTX, SeedVR).
+    """
+    def __init__(self):
+        self.total_gpus = torch.cuda.device_count()
+        self.ltx_gpus = []
+        self.seedvr_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(f"   > Total de GPUs detectadas: {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]
+            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}.")
+        
+        print("="*50)
+
+    def get_ltx_device(self):
+        """Retorna o dispositivo principal para o LTX."""
+        if not self.ltx_gpus:
+            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]}")
+
+    def get_seedvr_devices(self) -> list:
+        """Retorna a lista de IDs de GPU para o SeedVR."""
+        return self.seedvr_gpus
+
+    def requires_memory_swap(self) -> bool:
+        """Verifica se é necessário mover modelos entre CPU e GPU."""
+        return self.total_gpus < 2
+
+# Instância global para ser importada por outros módulos
+gpu_manager = GPUManager()