Add wheel building job

Author: NicolasHug

.github/workflows/linux_wheel.yaml

@@ -0,0 +1,54 @@
+name: Build Linux wheel
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+    tags:
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
+
+permissions:
+  id-token: write
+  contents: write
+
+defaults:
+  run:
+    shell: bash -l -eo pipefail {0}
+
+jobs:
+
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      with-xpu: disable
+      with-rocm: disable
+      with-cuda: disable
+      build-python-only: "disable"
+
+  build:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+    name: Build and Upload Linux wheel
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: pytorch-labs/torchvision-extra-decoders
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      pre-script: packaging/pre_build_script.sh
+      package-name: torchvision-extra-decoders
+      trigger-event: ${{ github.event_name }}

.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.0.1
+    hooks:
+      - id: check-docstring-first
+      - id: check-toml
+      - id: check-yaml
+        exclude: packaging/.*
+        args:
+          - --allow-multiple-documents
+      - id: mixed-line-ending
+        args: [--fix=lf]
+      - id: end-of-file-fixer
+
+  - repo: https://github.com/omnilib/ufmt
+    rev: v1.3.3
+    hooks:
+      - id: ufmt
+        additional_dependencies:
+          - black == 22.3.0
+          - usort == 1.0.2
+
+  - repo: https://github.com/PyCQA/flake8
+    rev: 5.0.4
+    hooks:
+      - id: flake8
+        args: [--config=setup.cfg]

packaging/pre_build_script.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+conda install libavif libheic -c conda-forge -yq

setup.cfg

@@ -0,0 +1,16 @@
+[bdist_wheel]
+universal=1
+
+[metadata]
+license_files = LICENSE
+
+[pep8]
+max-line-length = 120
+
+[flake8]
+# note: we ignore all 501s (line too long) anyway as they're taken care of by black
+max-line-length = 120
+ignore = E203, E402, W503, W504, F821, E501, B, C4, EXE
+per-file-ignores =
+    __init__.py: F401, F403, F405
+exclude = venv

setup.py

@@ -3,13 +3,15 @@
 # This software may be used and distributed according to the terms of the
 # GNU Lesser General Public License version 2.
 
-import sys
 import os
+import sys
 from pathlib import Path
+
 from setuptools import find_packages, setup
 
 from torch.utils.cpp_extension import BuildExtension, CppExtension
 
+
 def find_library(header):
     # returns (found, include dir, library dir)
     # if include dir or library dir is None, it means that the library is in
@@ -44,17 +46,19 @@ def find_library(header):
 
 def make_extension():
 
-    heic_found, heic_include_dir, heic_library_dir = find_library(header="libheif/heif_cxx.h")
+    heic_found, heic_include_dir, heic_library_dir = find_library(
+        header="libheif/heif_cxx.h"
+    )
     if not heic_found:
         raise RuntimeError("Couldn't find libheic!")
-    
+
     print(f"{heic_include_dir = }")
     print(f"{heic_library_dir = }")
 
     avif_found, avif_include_dir, avif_library_dir = find_library(header="avif/avif.h")
     if not avif_found:
         raise RuntimeError("Couldn't find libavif!")
-    
+
     print(f"{heic_include_dir = }")
     print(f"{heic_library_dir = }")
 
@@ -71,6 +75,14 @@ def make_extension():
     )
 
 
+def get_requirements():
+    pytorch_dep = os.getenv("TORCH_PACKAGE_NAME", "torch")
+    if os.getenv("PYTORCH_VERSION"):
+        pytorch_dep += "==" + os.getenv("PYTORCH_VERSION")
+
+    return [pytorch_dep]
+
+
 if __name__ == "__main__":
 
     with open("README.md") as f:
@@ -91,7 +103,7 @@ def make_extension():
         packages=find_packages(exclude=("test",)),
         package_data={PACKAGE_NAME: ["*.dll", "*.dylib", "*.so"]},
         zip_safe=False,
-        install_requires=[],
+        install_requires=get_requirements(),
         python_requires=">=3.9",
         ext_modules=[make_extension()],
         cmdclass={

torchvision_extra_decoders/__init__.py

@@ -4,6 +4,7 @@
 # GNU Lesser General Public License version 2.
 
 from pathlib import Path
+
 import torch
 
 

torchvision_extra_decoders/csrc/common.cpp

@@ -42,4 +42,4 @@ bool should_this_return_rgb_or_rgba(
   return !has_alpha;
 }
 
-} // namespace extra_decoders_ns
+} // namespace extra_decoders_ns

torchvision_extra_decoders/csrc/common.h

@@ -24,4 +24,4 @@ bool should_this_return_rgb_or_rgba(
     ImageReadMode mode,
     bool has_alpha);
 
-} // namespace extra_decoders_ns
+} // namespace extra_decoders_ns

torchvision_extra_decoders/csrc/decode_avif.cpp

@@ -86,4 +86,4 @@ torch::Tensor decode_avif(
   return out.permute({2, 0, 1}); // return CHW, channels-last
 }
 
-} // namespace extra_decoders_ns
+} // namespace extra_decoders_ns

torchvision_extra_decoders/csrc/decode_heic.cpp

@@ -125,4 +125,4 @@ torch::Tensor decode_heic(
   return out.permute({2, 0, 1});
 }
 
-} // namespace extra_decoders_ns
+} // namespace extra_decoders_ns