feature(framework): diffusers (#3534)

Co-authored-by: Aaron Pham <[email protected]>
bentoml · Feb 15, 2023 · ed1a7b9 · ed1a7b9
1 parent 8d091c8
commit ed1a7b9
Show file tree

Hide file tree

Showing 6 changed files with 592 additions and 0 deletions.
diff --git a/.github/workflows/frameworks.yml b/.github/workflows/frameworks.yml
@@ -23,6 +23,7 @@ jobs:
     runs-on: ubuntu-latest
     outputs:
       catboost: ${{ steps.filter.outputs.catboost }}
+      diffusers: ${{ steps.filter.outputs.diffusers }}
       fastai: ${{ steps.filter.outputs.fastai }}
       keras: ${{ steps.filter.outputs.keras }}
       lightgbm: ${{ steps.filter.outputs.lightgbm }}
@@ -56,6 +57,11 @@ jobs:
               - src/bentoml/catboost.py
               - src/bentoml/_internal/frameworks/catboost.py
               - tests/integration/frameworks/models/catboost.py
+            diffusers:
+              - *related
+              - src/bentoml/diffusers.py
+              - src/bentoml/_internal/frameworks/diffusers.py
+              - tests/integration/frameworks/models/diffusers.py
             lightgbm:
               - *related
               - src/bentoml/lightgbm.py
@@ -173,6 +179,51 @@ jobs:
           files: ./coverage.xml
           token: ${{ secrets.CODECOV_TOKEN }}
 
+  diffusers_integration_tests:
+    needs: diff
+    if: ${{ (github.event_name == 'pull_request' && needs.diff.outputs.diffusers == 'true') || github.event_name == 'push' }}
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0 # fetch all tags and branches
+      - name: Setup python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.8
+
+      - name: Get pip cache dir
+        id: cache-dir
+        run: |
+          echo ::set-output name=dir::$(pip cache dir)
+
+      - name: Cache pip dependencies
+        uses: actions/cache@v3
+        id: cache-pip
+        with:
+          path: ${{ steps.cache-dir.outputs.dir }}
+          key: ${{ runner.os }}-tests-${{ hashFiles('requirements/tests-requirements.txt') }}
+
+      - name: Install dependencies
+        run: |
+          pip install .
+          pip install diffusers torch transformers
+          pip install -r requirements/tests-requirements.txt
+
+      - name: Run tests and generate coverage report
+        run: |
+          OPTS=(--cov-config pyproject.toml --cov src/bentoml --cov-append --framework diffusers)
+          coverage run -m pytest tests/integration/frameworks/test_frameworks.py "${OPTS[@]}"
+
+      - name: Generate coverage
+        run: coverage xml
+
+      - name: Upload test coverage to Codecov
+        uses: codecov/codecov-action@v3
+        with:
+          files: ./coverage.xml
+          token: ${{ secrets.CODECOV_TOKEN }}
+
   fastai_integration_tests:
     needs: diff
     if: ${{ (github.event_name == 'pull_request' && needs.diff.outputs.fastai == 'true') || github.event_name == 'push' }}

diff --git a/requirements/frameworks-requirements.txt b/requirements/frameworks-requirements.txt
@@ -17,3 +17,4 @@ pytorch-lightning
 # huggingface dependencies
 transformers
 tokenizer
+diffusers
diff --git a/src/bentoml/__init__.py b/src/bentoml/__init__.py
@@ -68,6 +68,7 @@
     from . import lightgbm
     from . import onnxmlir
     from . import detectron
+    from . import diffusers
     from . import tensorflow
     from . import statsmodels
     from . import torchscript
@@ -91,6 +92,7 @@
 
     catboost = _LazyLoader("bentoml.catboost", globals(), "bentoml.catboost")
     detectron = _LazyLoader("bentoml.detectron", globals(), "bentoml.detectron")
+    diffusers = _LazyLoader("bentoml.diffusers", globals(), "bentoml.diffusers")
     easyocr = _LazyLoader("bentoml.easyocr", globals(), "bentoml.easyocr")
     flax = _LazyLoader("bentoml.flax", globals(), "bentoml.flax")
     fastai = _LazyLoader("bentoml.fastai", globals(), "bentoml.fastai")
@@ -166,6 +168,7 @@
     # Framework specific modules
     "catboost",
     "detectron",
+    "diffusers",
     "easyocr",
     "flax",
     "fastai",