Function require_big_accelerator

tests/testing_utils.py:619–645

Decorator marking a test that requires a bigger hardware accelerator (24GB) for execution. Some example pipelines: Flux, SD3, Cog, etc.

(test_case)

def require_big_accelerator(test_case):
    """
    Decorator marking a test that requires a bigger hardware accelerator (24GB) for execution. Some example pipelines:
    Flux, SD3, Cog, etc.
    """
    import pytest

    test_case = pytest.mark.big_accelerator(test_case)

    if not is_torch_available():
        return pytest.mark.skip(reason="test requires PyTorch")(test_case)

    import torch

    if not (torch.cuda.is_available() or torch.xpu.is_available()):
        return pytest.mark.skip(reason="test requires PyTorch CUDA")(test_case)

    if torch.xpu.is_available():
        device_properties = torch.xpu.get_device_properties(0)
    else:
        device_properties = torch.cuda.get_device_properties(0)

    total_memory = device_properties.total_memory / (1024**3)
    return pytest.mark.skipif(
        total_memory < BIG_GPU_MEMORY,
        reason=f"test requires a hardware accelerator with at least {BIG_GPU_MEMORY} GB memory",
    )(test_case)


def require_torch_accelerator_with_training(test_case):