lenML/ChatTTS-Forge

1import logging2import sys3from functools import lru_cache4 5import torch6 7from modules import config8 9logger = logging.getLogger(__name__)10 11if sys.platform == "darwin":12    from modules.devices import mac_devices13 14 15def has_mps() -> bool:16    if sys.platform != "darwin":17        return False18    else:19        return mac_devices.has_mps20 21 22def get_cuda_device_id():23    return (24        int(config.runtime_env_vars.device_id)25        if config.runtime_env_vars.device_id is not None26        and config.runtime_env_vars.device_id.isdigit()27        else 028    ) or torch.cuda.current_device()29 30 31def get_cuda_device_string():32    if config.runtime_env_vars.device_id is not None:33        return f"cuda:{config.runtime_env_vars.device_id}"34 35    return "cuda"36 37 38def get_available_gpus() -> list[tuple[int, int]]:39    """40    Get the list of available GPUs and their free memory.41 42    :return: A list of tuples where each tuple contains (GPU index, free memory in bytes).43    """44    available_gpus = []45    for i in range(torch.cuda.device_count()):46        props = torch.cuda.get_device_properties(i)47        free_memory = props.total_memory - torch.cuda.memory_reserved(i)48        available_gpus.append((i, free_memory))49    return available_gpus50 51 52def get_memory_available_gpus(min_memory=2048):53    available_gpus = get_available_gpus()54    memory_available_gpus = [55        gpu for gpu, free_memory in available_gpus if free_memory > min_memory56    ]57    return memory_available_gpus58 59 60def get_target_device_id_or_memory_available_gpu():61    memory_available_gpus = get_memory_available_gpus()62    device_id = get_cuda_device_id()63    if device_id not in memory_available_gpus:64        if len(memory_available_gpus) != 0:65            logger.warning(66                f"Device {device_id} is not available or does not have enough memory. will try to use {memory_available_gpus}"67            )68            config.runtime_env_vars.device_id = str(memory_available_gpus[0])69        else:70            logger.warning(71                f"Device {device_id} is not available or does not have enough memory. Using CPU instead."72            )73            return "cpu"74    return get_cuda_device_string()75 76 77def get_optimal_device_name():78    if config.runtime_env_vars.use_cpu == "all":79        return "cpu"80 81    if torch.cuda.is_available():82        return get_target_device_id_or_memory_available_gpu()83 84    if has_mps():85        return "mps"86 87    return "cpu"88 89 90def get_optimal_device():91    return torch.device(get_optimal_device_name())92 93 94def get_device_for(task):95    if (96        task in config.runtime_env_vars.use_cpu97        or "all" in config.runtime_env_vars.use_cpu98    ):99        return cpu100 101    return get_optimal_device()102 103 104def torch_gc():105    try:106        if torch.cuda.is_available():107            with torch.cuda.device(get_cuda_device_string()):108                torch.cuda.empty_cache()109                torch.cuda.ipc_collect()110 111        if has_mps():112            mac_devices.torch_mps_gc()113    except Exception as e:114        logger.error(f"Error in torch_gc", exc_info=True)115 116 117cpu: torch.device = torch.device("cpu")118device: torch.device = None119dtype: torch.dtype = torch.float32120dtype_dvae: torch.dtype = torch.float32121dtype_vocos: torch.dtype = torch.float32122dtype_gpt: torch.dtype = torch.float32123dtype_decoder: torch.dtype = torch.float32124 125 126def reset_device():127    global device128    global dtype129    global dtype_dvae130    global dtype_vocos131    global dtype_gpt132    global dtype_decoder133 134    if config.runtime_env_vars.use_cpu is None:135        config.runtime_env_vars.use_cpu = []136 137    if "all" in config.runtime_env_vars.use_cpu and not config.runtime_env_vars.no_half:138        logger.warning(139            "Cannot use half precision with CPU, using full precision instead"140        )141        config.runtime_env_vars.no_half = True142 143    if not config.runtime_env_vars.no_half:144        dtype = torch.float16145        dtype_dvae = torch.float16146        dtype_vocos = torch.float16147        dtype_gpt = torch.float16148        dtype_decoder = torch.float16149 150        logger.info("Using half precision: torch.float16")151    else:152        dtype = torch.float32153        dtype_dvae = torch.float32154        dtype_vocos = torch.float32155        dtype_gpt = torch.float32156        dtype_decoder = torch.float32157 158        logger.info("Using full precision: torch.float32")159 160    if "all" in config.runtime_env_vars.use_cpu:161        device = cpu162    else:163        device = get_optimal_device()164 165    logger.info(f"Using device: {device}")166 167 168@lru_cache169def first_time_calculation():170    """171    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and172    spends about 2.7 seconds doing that, at least wih NVidia.173    """174 175    x = torch.zeros((1, 1)).to(device, dtype)176    linear = torch.nn.Linear(1, 1).to(device, dtype)177    linear(x)178 179    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)180    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)181    conv2d(x)182