faisalhr1997/codeformer

1#!/usr/bin/env python2 3from setuptools import find_packages, setup4 5import os6import subprocess7import sys8import time9import torch10from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension11 12version_file = './basicsr/version.py'13 14 15def readme():16    with open('README.md', encoding='utf-8') as f:17        content = f.read()18    return content19 20 21def get_git_hash():22 23    def _minimal_ext_cmd(cmd):24        # construct minimal environment25        env = {}26        for k in ['SYSTEMROOT', 'PATH', 'HOME']:27            v = os.environ.get(k)28            if v is not None:29                env[k] = v30        # LANGUAGE is used on win3231        env['LANGUAGE'] = 'C'32        env['LANG'] = 'C'33        env['LC_ALL'] = 'C'34        out = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=env).communicate()[0]35        return out36 37    try:38        out = _minimal_ext_cmd(['git', 'rev-parse', 'HEAD'])39        sha = out.strip().decode('ascii')40    except OSError:41        sha = 'unknown'42 43    return sha44 45 46def get_hash():47    if os.path.exists('.git'):48        sha = get_git_hash()[:7]49    elif os.path.exists(version_file):50        try:51            from version import __version__52            sha = __version__.split('+')[-1]53        except ImportError:54            raise ImportError('Unable to get git version')55    else:56        sha = 'unknown'57 58    return sha59 60 61def write_version_py():62    content = """# GENERATED VERSION FILE63# TIME: {}64__version__ = '{}'65__gitsha__ = '{}'66version_info = ({})67"""68    sha = get_hash()69    with open('./basicsr/VERSION', 'r') as f:70        SHORT_VERSION = f.read().strip()71    VERSION_INFO = ', '.join([x if x.isdigit() else f'"{x}"' for x in SHORT_VERSION.split('.')])72 73    version_file_str = content.format(time.asctime(), SHORT_VERSION, sha, VERSION_INFO)74    with open(version_file, 'w') as f:75        f.write(version_file_str)76 77 78def get_version():79    with open(version_file, 'r') as f:80        exec(compile(f.read(), version_file, 'exec'))81    return locals()['__version__']82 83 84def make_cuda_ext(name, module, sources, sources_cuda=None):85    if sources_cuda is None:86        sources_cuda = []87    define_macros = []88    extra_compile_args = {'cxx': []}89 90    if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1':91        define_macros += [('WITH_CUDA', None)]92        extension = CUDAExtension93        extra_compile_args['nvcc'] = [94            '-D__CUDA_NO_HALF_OPERATORS__',95            '-D__CUDA_NO_HALF_CONVERSIONS__',96            '-D__CUDA_NO_HALF2_OPERATORS__',97        ]98        sources += sources_cuda99    else:100        print(f'Compiling {name} without CUDA')101        extension = CppExtension102 103    return extension(104        name=f'{module}.{name}',105        sources=[os.path.join(*module.split('.'), p) for p in sources],106        define_macros=define_macros,107        extra_compile_args=extra_compile_args)108 109 110def get_requirements(filename='requirements.txt'):111    with open(os.path.join('.', filename), 'r') as f:112        requires = [line.replace('\n', '') for line in f.readlines()]113    return requires114 115 116if __name__ == '__main__':117    if '--cuda_ext' in sys.argv:118        ext_modules = [119            make_cuda_ext(120                name='deform_conv_ext',121                module='ops.dcn',122                sources=['src/deform_conv_ext.cpp'],123                sources_cuda=['src/deform_conv_cuda.cpp', 'src/deform_conv_cuda_kernel.cu']),124            make_cuda_ext(125                name='fused_act_ext',126                module='ops.fused_act',127                sources=['src/fused_bias_act.cpp'],128                sources_cuda=['src/fused_bias_act_kernel.cu']),129            make_cuda_ext(130                name='upfirdn2d_ext',131                module='ops.upfirdn2d',132                sources=['src/upfirdn2d.cpp'],133                sources_cuda=['src/upfirdn2d_kernel.cu']),134        ]135        sys.argv.remove('--cuda_ext')136    else:137        ext_modules = []138 139    write_version_py()140    setup(141        name='basicsr',142        version=get_version(),143        description='Open Source Image and Video Super-Resolution Toolbox',144        long_description=readme(),145        long_description_content_type='text/markdown',146        author='Xintao Wang',147        author_email='xintao.wang@outlook.com',148        keywords='computer vision, restoration, super resolution',149        url='https://github.com/xinntao/BasicSR',150        include_package_data=True,151        packages=find_packages(exclude=('options', 'datasets', 'experiments', 'results', 'tb_logger', 'wandb')),152        classifiers=[153            'Development Status :: 4 - Beta',154            'License :: OSI Approved :: Apache Software License',155            'Operating System :: OS Independent',156            'Programming Language :: Python :: 3',157            'Programming Language :: Python :: 3.7',158            'Programming Language :: Python :: 3.8',159        ],160        license='Apache License 2.0',161        setup_requires=['cython', 'numpy'],162        install_requires=get_requirements(),163        ext_modules=ext_modules,164        cmdclass={'build_ext': BuildExtension},165        zip_safe=False)166