linting

.github/workflows/ci-cd.yaml

@@ -22,6 +22,9 @@ jobs:
           - event: push
             model: 'dev', 'schnell'
             env: 'prod'
+          - event: workflow_dispatch
+            model: 'dev', 'schnell'
+            env: 'prod'
 
     steps:
     - name: Checkout
@@ -95,7 +98,7 @@ jobs:
 
     - name: Test model
       env:
-        REPLICATE_API_TOKEN: ${{secrets.REPLICATE_API_TOKEN}}
+        REPLICATE_API_TOKEN: ${{ secrets.REPLICATE_API_TOKEN }}
         MODEL: ${{ matrix.model }}
         TEST_ENV: ${{ matrix.env }}
       run: |

integration-tests/test-model.py

@@ -92,12 +92,12 @@ def wait_for_server_to_be_ready(url, timeout=400):
 def inference_func():
     if ENV == 'local':
         return partial(local_run, LOCAL_ENDPOINT)
-    elif ENV in {'staging', 'prod'}:
+    elif ENV in {'test', 'prod'}:
         model = replicate.models.get(MODEL)
         version = model.versions.list()[0]
         return partial(replicate_run, version)
     else:
-        raise Exception(f"env should be local, staging, or prod but was {ENV}")
+        raise Exception(f"env should be local, test, or prod but was {ENV}")
 
 
 @pytest.fixture(scope="session", autouse=True)

predict.py

@@ -1,5 +1,4 @@
 import os
-import pickle
 import time
 from typing import Optional
 
@@ -8,7 +7,6 @@
 
 import torch
 import numpy as np
-from einops import rearrange
 from PIL import Image
 from typing import List
 from einops import rearrange

requirements_test.txt

@@ -0,0 +1,5 @@
+numpy
+pytest
+replicate
+requests
+Pillow

torch_compile.py

@@ -56,4 +56,3 @@
 # 1440:1440
 # torch.Size([1, 16, 180, 180])
 # torch.Size([1, 8100, 64])
-_