setup.py

#!/usr/bin/env python

from setuptools import find_packages, setup
import os


def check_torch_version():
    try:
        import torch

        if not torch.cuda.is_available():
            raise RuntimeError(
                "\033[91mCUDA is not available. Please install a CUDA 11- or 12.1-enabled PyTorch version.\033[0m"
            )

        cuda_version = torch.version.cuda
        print(cuda_version)
        if cuda_version not in ["11.6", "11.7", "11.8", "12.1"]:
            print(
                f" Warning: \033[91mUnsupported CUDA version: {cuda_version}. Please install PyTorch with 11.6<=CUDA<=12.1.\033[0m"
            )
    except:
        print(
            "\033[93mPlease install torch and torchvision compiled with cuda 11.8 before installing vbench\033[0m"
        )
        print("\033[93mFor CUDA 11.8, run:\033[0m")
        print(
            "\033[93m    pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118\033[0m"
        )
        print("\033[93mFor CUDA 12.1, run:\033[0m")
        print(
            "\033[93m    pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121\033[0m"
        )
        print("\033[93mFor CUDA 11 with PyTorch < 2.0 run:\033[0m")
        print("\033[93m    pip install torch==1.13.1 torchvision==0.14.1\033[0m")
        raise


def fetch_readme():
    with open("README-pypi.md", encoding="utf-8") as f:
        text = f.read()
    return text


def fetch_requirements():
    filename = os.path.join(
        os.path.dirname(os.path.realpath(__file__)), "requirements.txt"
    )
    with open(filename, "r") as f:
        envs = [line.rstrip("\n") for line in f.readlines() if "@" not in line]
    return envs


install_requires = fetch_requirements()
setup(
    name="vbench",
    version="0.1.5",
    description="Video generation benchmark",
    long_description=fetch_readme(),
    long_description_content_type="text/markdown",
    project_urls={
        "Source": "https://github.com/Vchitect/VBench",
    },
    entry_points={"console_scripts": ["vbench=vbench.cli.vbench:main"]},
    install_requires=install_requires,
    packages=find_packages(),
    include_package_data=True,
    license="Apache Software License 2.0",
)