setup.py 3.98 KB
Newer Older
Soumith Chintala's avatar
Soumith Chintala committed
1
#!/usr/bin/env python
2
3
import os
import platform
4
5
import sys
import subprocess
Soumith Chintala's avatar
Soumith Chintala committed
6
7

from setuptools import setup, find_packages
8
from torch.utils.cpp_extension import BuildExtension, CppExtension
Soumith Chintala's avatar
Soumith Chintala committed
9

10
11
12
13
14

def check_env_flag(name, default=''):
    return os.getenv(name, default).upper() in set(['ON', '1', 'YES', 'TRUE', 'Y'])

DEBUG = check_env_flag('DEBUG')
15
16
17
18
19
IS_WHEEL = check_env_flag('IS_WHEEL')
IS_CONDA = check_env_flag('IS_CONDA')

print('DEBUG:', DEBUG, 'IS_WHEEL:', IS_WHEEL, 'IS_CONDA:', IS_CONDA)

20
21
22
23
24
25
26
27
28
eca = []
ela = []
if DEBUG:
    if platform.system() == 'Windows':
        ela += ['/DEBUG:FULL']
    else:
        eca += ['-O0', '-g']
        ela += ['-O0', '-g']

29
30
31
32
33
34
35
36
37
38
39

libraries = []
include_dirs = []
extra_objects = []

if IS_WHEEL:
    audio_path = os.path.dirname(os.path.abspath(__file__))

    include_dirs += [os.path.join(audio_path, 'third_party/flac/include')]
    include_dirs += [os.path.join(audio_path, 'third_party/lame/include')]
    include_dirs += [os.path.join(audio_path, 'third_party/sox/include')]
40
    include_dirs += [os.path.join(audio_path, 'third_party/mad/include')]
41

42
43
44
    # proper link order (sox, mad, flac, lame)
    # (the most important thing is that dependencies come after a libraryl
    # e.g., sox comes first)
45
    extra_objects += [os.path.join(audio_path, 'third_party/sox/lib/libsox.a')]
46
    extra_objects += [os.path.join(audio_path, 'third_party/mad/lib/libmad.a')]
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
    extra_objects += [os.path.join(audio_path, 'third_party/flac/lib/libFLAC.a')]
    extra_objects += [os.path.join(audio_path, 'third_party/lame/lib/libmp3lame.a')]
else:
    libraries += ['sox']

if IS_CONDA:
    # We want $PREFIX/include for conda (for sox.h)
    lib_path = os.path.dirname(sys.executable)
    include_dirs += [os.path.join(os.path.dirname(lib_path), 'include')]


# Creating the version file
cwd = os.path.dirname(os.path.abspath(__file__))
version = '0.2.0a0'
sha = 'Unknown'

try:
    sha = subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=cwd).decode('ascii').strip()
except Exception:
    pass

if os.getenv('TORCHAUDIO_BUILD_VERSION'):
    assert os.getenv('TORCHAUDIO_BUILD_NUMBER') is not None
    build_number = int(os.getenv('TORCHAUDIO_BUILD_NUMBER'))
    version = os.getenv('TORCHAUDIO_BUILD_VERSION')
    if build_number > 1:
        version += '.post' + str(build_number)
elif sha != 'Unknown':
    version += '+' + sha[:7]
print('-- Building version ' + version)

version_path = os.path.join(cwd, 'torchaudio', 'version.py')
with open(version_path, 'w') as f:
    f.write("__version__ = '{}'\n".format(version))
    f.write("git_version = {}\n".format(repr(sha)))

83
84
pytorch_package_name = os.getenv('TORCHAUDIO_PYTORCH_DEPENDENCY_NAME', 'torch')

Soumith Chintala's avatar
Soumith Chintala committed
85
86
setup(
    name="torchaudio",
David Pollack's avatar
David Pollack committed
87
    version="0.2",
Soumith Chintala's avatar
Soumith Chintala committed
88
89
    description="An audio package for PyTorch",
    url="https://github.com/pytorch/audio",
90
    author="Soumith Chintala, David Pollack, Sean Naren, Peter Goldsborough",
Soumith Chintala's avatar
Soumith Chintala committed
91
    author_email="soumith@pytorch.org",
Hong Xu's avatar
Hong Xu committed
92
93
94
95
96
97
98
99
100
    classifiers=[
        "Environment :: Plugins",
        "Intended Audience :: Developers",
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: BSD License",
        "Operating System :: MacOS :: MacOS X",
        "Operating System :: Microsoft :: Windows",
        "Operating System :: POSIX",
        "Programming Language :: C++",
101
102
        "Programming Language :: Python :: 2.7",
        "Programming Language :: Python :: 3",
Hong Xu's avatar
Hong Xu committed
103
104
105
106
        "Programming Language :: Python :: Implementation :: CPython",
        "Topic :: Multimedia :: Sound/Audio",
        "Topic :: Scientific/Engineering :: Artificial Intelligence"
    ],
Soumith Chintala's avatar
Soumith Chintala committed
107
108
    # Exclude the build files.
    packages=find_packages(exclude=["build"]),
109
110
    ext_modules=[
        CppExtension(
111
112
            '_torch_sox',
            ['torchaudio/torch_sox.cpp'],
113
114
            libraries=libraries,
            include_dirs=include_dirs,
115
            extra_compile_args=eca,
116
            extra_objects=extra_objects,
117
            extra_link_args=ela),
Soumith Chintala's avatar
Soumith Chintala committed
118
    ],
119
    cmdclass={'build_ext': BuildExtension},
120
    install_requires=[pytorch_package_name]
121
)