utils.py

def install_packages():
    import subprocess
    import sys
    import importlib

    def _is_package_available(name) -> bool:
        try:
            importlib.import_module(name)
            return True
        except (ImportError, ModuleNotFoundError):
            return False

    # upgrade pip
    subprocess.run(
        f"{sys.executable} -m pip install --upgrade pip", shell=True)
    subprocess.run(
        f"{sys.executable} -m pip install --upgrade wheel setuptools packaging", shell=True)

    # install ninja
    if not _is_package_available("ninja"):
        subprocess.run(f"{sys.executable} -m pip install ninja", shell=True)

    # install flash attention
    if not _is_package_available("flash_attn"):
        subprocess.run(
            f"{sys.executable} -m pip install flash-attn --no-build-isolation",
            env={"MAX_JOBS": "4"},
            shell=True,
        )

    # install xformers
    if not _is_package_available("xformers"):
        subprocess.run(
            f"{sys.executable} -m pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers",
            shell=True,
        )