Spaces:
Paused
Paused
File size: 3,829 Bytes
0f26b1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 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() |