utils.py

"""
Utility functions và PyTorch optimizations
"""
import torch
import os

def setup_pytorch_optimizations():
    """Cấu hình PyTorch cho tốc độ tối đa"""
    if torch.cuda.is_available():
        # Tối ưu CUDA settings
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        # Tối ưu memory allocation - tận dụng hết GPU RAM
        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True,roundup_power2_divisions:16'
        # Tối ưu cuDNN
        os.environ['CUDNN_BENCHMARK'] = '1'
        # Tối ưu cuBLAS - tăng workspace để nhanh hơn
        os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':8192:8'
        # Tối ưu thêm
        os.environ['TORCH_CUDNN_V8_API_ENABLED'] = '1'  # CUDNN v8 API
        print("⚡ FULL PERFORMANCE MODE: Đã bật tất cả tối ưu CUDA")
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        print(f"⚡ GPU Memory: {gpu_memory:.2f} GB - Sẽ tận dụng tối đa!")
        return True
    return False

def get_gpu_decorator():
    """Lấy GPU decorator từ spaces module nếu có"""
    try:
        from spaces import GPU
        return GPU
    except ImportError:
        # Fallback if spaces module is not available
        def GPU(func):
            return func
        return GPU