config/gpu_allocation.py

import torch

class GPUManager:
    def __init__(self, allocation_map):
        self.map = allocation_map
        self.device_count = torch.cuda.device_count()

    def get_device(self, agent_name):
        gpu_id = self.map.get(agent_name, 0)
        if gpu_id >= self.device_count:
            gpu_id = 0
        return f"cuda:{gpu_id}"

    def get_vram_info(self, gpu_id):
        if not torch.cuda.is_available():
            return {"total": 0, "free": 0, "used": 0}
        total = torch.cuda.get_device_properties(gpu_id).total_memory
        reserved = torch.cuda.memory_reserved(gpu_id)
        allocated = torch.cuda.memory_allocated(gpu_id)
        return {
            "total": total / 1024**3,
            "used": allocated / 1024**3,
            "free": (total - reserved) / 1024**3,
        }

    def all_devices(self):
        return [f"cuda:{i}" for i in range(self.device_count)]

gpu_manager = GPUManager({
    "orchestrator": 7,
    "video_primary": 0,
    "video_backup": 1,
    "face_restore": 2,
    "frame_interpolate": 3,
    "video_upscale": 4,
    "image_generator": 5,
    "audio_generator": 6,
    "quality_evaluator": 6,
})