| 52 | DEFAULT_VIDEO_END_TOKEN = "</vid>" |
| 53 | |
| 54 | def get_gpu_memory(max_gpus=None): |
| 55 | gpu_memory = [] |
| 56 | num_gpus = ( |
| 57 | torch.cuda.device_count() |
| 58 | if max_gpus is None |
| 59 | else min(max_gpus, torch.cuda.device_count()) |
| 60 | ) |
| 61 | |
| 62 | for gpu_id in range(num_gpus): |
| 63 | with torch.cuda.device(gpu_id): |
| 64 | device = torch.cuda.current_device() |
| 65 | gpu_properties = torch.cuda.get_device_properties(device) |
| 66 | total_memory = gpu_properties.total_memory / (1024 ** 3) |
| 67 | allocated_memory = torch.cuda.memory_allocated() / (1024 ** 3) |
| 68 | available_memory = total_memory - allocated_memory |
| 69 | gpu_memory.append(available_memory) |
| 70 | return gpu_memory |
| 71 | |
| 72 | def load_model( |
| 73 | model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, lora_weights=None |