from typing import Literal, Optional

from setuptools import find_packages, setup


deps = [
    "diffusers==0.25.0",
    "transformers",
    "accelerate",
    "fire",
    "einops",
    "omegaconf",
    "imageio",
    "timm==0.6.7",
    "lightning",
    "peft",
    "av",
    "decord",
    "pillow",
    "pywin32;sys_platform == 'win32'",
]


def get_cuda_version_from_torch() -> Optional[Literal["11", "12"]]:
    try:
        import torch

        return torch.version.cuda.split(".")[0]
    except ImportError:
        raise ImportError("Please install PyTorch first. See https://pytorch.org/get-started/locally/.")


cu = get_cuda_version_from_torch()

assert cu in ["11", "12"], f"Unsupported CUDA version: {cu}"

deps_tensorrt = [
    "onnx==1.16.0",
    "onnxruntime==1.16.3",
    "protobuf==5.27.0",
    "polygraphy",
    "onnx-graphsurgeon",
    "cuda-python",
    f"tensorrt_cu{cu}_libs==10.0.1",
    f"tensorrt_cu{cu}_bindings==10.0.1",
    "tensorrt==10.0.1",
    "colored",
]
extras = {"tensorrt": deps_tensorrt}


if __name__ == "__main__":
    setup(
        name="Live2Diff",
        version="0.1",
        description="real-time interactive video translation pipeline",
        long_description=open("README.md", "r", encoding="utf-8").read(),
        long_description_content_type="text/markdown",
        keywords="deep learning diffusion pytorch stable diffusion streamdiffusion real-time next-frame prediction",
        license="Apache 2.0 License",
        author="leo",
        author_email="xingzhening@pjlab.org.cn",
        url="https://github.com/LeoXing1996/NextFramePredictionPreview",
        package_dir={"": "live2diff"},
        packages=find_packages("live2diff"),
        python_requires=">=3.10.0",
        install_requires=deps,
        extras_require=extras,
    )