Spaces:
Sleeping
Sleeping
| # Device setup utilities for GPU/CPU detection | |
| import torch | |
| import gc | |
| from config import GPU_MEMORY_FRACTION | |
| def setup_device(): | |
| # GPU or CPU setup, Returns device object and bool indicating if GPU is available | |
| if torch.cuda.is_available(): | |
| device = torch.device("cuda") | |
| gpu_name = torch.cuda.get_device_name(0) | |
| gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9 | |
| print(f"Using GPU: {gpu_name} ({gpu_memory:.1f}GB)") | |
| # Seting some fraction of memory to avoid OOM errors | |
| torch.cuda.set_per_process_memory_fraction(GPU_MEMORY_FRACTION) | |
| return device, True | |
| else: | |
| print("GPU not available, using CPU") | |
| return torch.device("cpu"), False | |
| def clear_gpu_memory(): | |
| """Clear GPU memory cache to free up VRAM""" | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| print("GPU memory cache cleared") | |
| def get_gpu_memory_info(): | |
| """Get current GPU memory usage information""" | |
| if torch.cuda.is_available(): | |
| gpu_memory_used = torch.cuda.memory_allocated(0) / 1e9 | |
| gpu_memory_total = torch.cuda.get_device_properties(0).total_memory / 1e9 | |
| return gpu_memory_used, gpu_memory_total | |
| return 0, 0 |