diff --git a/.github/workflows/ci-cd.yaml b/.github/workflows/ci-cd.yaml
index dbaa2eb..cf85aba 100644
--- a/.github/workflows/ci-cd.yaml
+++ b/.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: |
diff --git a/integration-tests/test-model.py b/integration-tests/test-model.py
index 7fdb4bb..94f4552 100644
--- a/integration-tests/test-model.py
+++ b/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)
diff --git a/predict.py b/predict.py
index 049726f..ea2fbc8 100644
--- a/predict.py
+++ b/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
diff --git a/requirements_test.txt b/requirements_test.txt
new file mode 100644
index 0000000..3184d45
--- /dev/null
+++ b/requirements_test.txt
@@ -0,0 +1,5 @@
+numpy
+pytest
+replicate
+requests
+Pillow
\ No newline at end of file
diff --git a/torch_compile.py b/torch_compile.py
index d141e03..a4f5eb8 100644
--- a/torch_compile.py
+++ b/torch_compile.py
@@ -56,4 +56,3 @@
 # 1440:1440
 # torch.Size([1, 16, 180, 180])
 # torch.Size([1, 8100, 64])
-_
\ No newline at end of file