diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml
index b995253..f1fba58 100644
--- a/.github/workflows/push.yml
+++ b/.github/workflows/push.yml
@@ -13,52 +13,33 @@ jobs:
   build_test_push:
     name: Build, Test & Push Model
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        model: [yolo11n]
     steps:
-      #- name: Cleanup disk space
-      #  uses: ultralytics/actions/cleanup-disk@main
-
       - name: Checkout
         uses: actions/checkout@v4
-
       - name: Setup Python
         uses: actions/setup-python@v5
         with:
           python-version: 3.12
-
       - name: Setup uv
         uses: astral-sh/setup-uv@v6
-
       - name: Install dependencies
         run: uv pip install --system ultralytics --extra-index-url https://download.pytorch.org/whl/cpu
-
       - name: Setup Cog
         uses: replicate/setup-cog@v2
         with:
           token: ${{ secrets.REPLICATE_API_TOKEN }}
-
       - name: Download YOLO11n weights
-        run: |
-          python << 'EOF'
-          from ultralytics import YOLO
-          model = YOLO('yolo11n.pt')
-          EOF
-          ls -la yolo11n*
-          mv yolo11n.pt yolo11n/
-          echo "Files in yolo11n directory:"
-          ls -la yolo11n/
-
+        run: python ${{ matrix.model }}/download.py
       - name: Build model image
-        run: |
-          cd yolo11n
-          cog build
-
+        working-directory: ${{ matrix.model }}
+        run: cog build
       - name: Test model
-        run: |
-          cd yolo11n
-          cog predict -i image=@../assets/bus.jpg -i conf=0.25 -i iou=0.45
-
+        working-directory: ${{ matrix.model }}
+        run: cog predict -i image=@../assets/bus.jpg -i conf=0.25 -i iou=0.45
       - name: Push to Replicate
         if: github.ref == 'refs/heads/main'
-        run: |
-          cd yolo11n
-          cog push
+        working-directory: ${{ matrix.model }}
+        run: cog push
diff --git a/yolo11n/cog.yaml b/yolo11n/cog.yaml
index 3889754..8b5fd54 100644
--- a/yolo11n/cog.yaml
+++ b/yolo11n/cog.yaml
@@ -9,11 +9,6 @@ build:
     - "ffmpeg"
   python_packages:
     - "ultralytics>=8.3.0"
-    - "torch==2.3.1"
-    - "torchvision"
-    - "numpy>=1.24.0"
-    - "opencv-python"
-    - "pillow"
 
 predict: predict.py:Predictor
 image: r8.im/ultralytics/yolo11n
diff --git a/yolo11n/download.py b/yolo11n/download.py
new file mode 100644
index 0000000..8940470
--- /dev/null
+++ b/yolo11n/download.py
@@ -0,0 +1,20 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from pathlib import Path
+
+from ultralytics import YOLO
+
+
+def main():
+    """Download YOLO11n weights and move to model directory."""
+    current_dir = Path(__file__).parent
+    YOLO(current_dir / "yolo11n.pt")
+
+    # List files in model directory
+    print(f"Files in {current_dir.name} directory:")
+    for file in sorted(current_dir.iterdir()):
+        print(f"  {file.stat().st_size:>10} {file.name}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/yolo11n/predict.py b/yolo11n/predict.py
index 4d80c42..3f23445 100644
--- a/yolo11n/predict.py
+++ b/yolo11n/predict.py
@@ -1,9 +1,9 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-import json
 from typing import Optional
+
+from cog import BaseModel, BasePredictor, Input, Path
 from ultralytics import YOLO
-from cog import BasePredictor, Input, Path, BaseModel
 
 
 class Output(BaseModel):
@@ -32,9 +32,6 @@ def predict(
         result.save(image_path)
 
         if return_json:
-            return Output(
-                image=Path(image_path), 
-                json_str=result.to_json()
-            )
+            return Output(image=Path(image_path), json_str=result.to_json())
         else:
             return Output(image=Path(image_path))