diffusers/optimized-diffusers-code

1import subprocess2import psutil3import functools4from torch._inductor.runtime.hints import DeviceProperties5from torch._inductor.utils import get_gpu_type6from typing import Union7import torch8 9 10@functools.cache11def get_system_ram_gb():12    """13    Gets the total physical system RAM in Gigabytes.14 15    Returns:16        float: Total system RAM in GB, or None if it cannot be determined.17    """18    try:19        # Get virtual memory details20        virtual_memory = psutil.virtual_memory()21        # Total physical memory in bytes22        total_ram_bytes = virtual_memory.total23        # Convert bytes to gigabytes (1 GB = 1024^3 bytes)24        total_ram_gb = total_ram_bytes / (1024**3)25        return total_ram_gb26    except Exception as e:27        print(f"Error getting system RAM: {e}")28        return None29 30 31@functools.cache32def get_gpu_vram_gb():33    """34    Gets the total GPU VRAM in Gigabytes using the nvidia-smi command.35    This function is intended for NVIDIA GPUs.36 37    Returns:38        float: Total GPU VRAM in GB, or None if it cannot be determined.39    """40    try:41        # Execute the nvidia-smi command to get GPU memory info42        # The command queries for the total memory and outputs it in MiB43        result = subprocess.run(44            ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"],45            capture_output=True,46            text=True,47            check=True,48        )49        # The output will be a string like "12288\n" for the first GPU50        # We take the first line in case there are multiple GPUs51        vram_mib = int(result.stdout.strip().split("\n")[0])52        # Convert MiB to Gigabytes (1 GB = 1024 MiB)53        vram_gb = vram_mib / 102454        return vram_gb55    except FileNotFoundError:56        # This error occurs if nvidia-smi is not installed or not in the PATH57        print("INFO: 'nvidia-smi' command not found. Cannot determine GPU VRAM.")58        print("      This is expected if you don't have an NVIDIA GPU or drivers installed.")59        return None60    except (subprocess.CalledProcessError, ValueError, IndexError) as e:61        # Handles other potential errors like command failure or parsing issues62        print(f"Error getting GPU VRAM: {e}")63        return None64 65 66def categorize_ram(ram_gb):67    """68    Categorizes RAM into 'small', 'medium', or 'large'.69 70    Args:71        ram_gb (float): The amount of RAM in GB.72 73    Returns:74        str: The category ('small', 'medium', 'large') or 'unknown'.75    """76    if ram_gb is None:77        return "unknown"78    if ram_gb <= 20:79        return "small"80    elif 20 < ram_gb <= 40:81        return "medium"82    else:  # ram_gb > 4083        return "large"84 85 86def categorize_vram(vram_gb):87    """88    Categorizes VRAM into 'small', 'medium', or 'large'.89 90    Args:91        vram_gb (float): The amount of VRAM in GB.92 93    Returns:94        str: The category ('small', 'medium', 'large') or 'not applicable/unknown'.95    """96    if vram_gb is None:97        return "not applicable/unknown"98    if vram_gb <= 8:99        return "small"100    elif 8 < vram_gb <= 24:101        return "medium"102    else:  # vram_gb > 24103        return "large"104 105 106@functools.cache107def is_compile_friendly_gpu(index_or_device: Union[int, str, torch.device] = 0) -> bool:108    """Hand-coded rules from experiments. Don't take seriously."""109    if isinstance(index_or_device, torch.device):110        device = index_or_device111    elif isinstance(index_or_device, str):112        device = torch.device(index_or_device)113    else:114        device = torch.device(get_gpu_type(), index_or_device)115 116    prop = DeviceProperties.create(device)117    return prop.major >= 8118 119 120@functools.lru_cache()121def is_sm_version(major: int, minor: int) -> bool:122    """Check if the CUDA version is exactly major.minor"""123    is_cuda = torch.cuda.is_available() and torch.version.cuda124    return torch.cuda.get_device_capability() == (major, minor) if is_cuda else False125 126 127def is_fp8_friendly():128    return is_sm_version(8, 9)129