Spaces:
Paused
Paused
File size: 3,829 Bytes
99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd 99dd76a 8fd9bdd |
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() |