From bd9401a6964d11f1abed02211667f7daf7c1ebdd Mon Sep 17 00:00:00 2001
From: mbukeRepo <mbuke404@gmail.com>
Date: Sat, 6 Jul 2024 13:13:39 +0200
Subject: [PATCH] chore: setup replicate deployment

---
 .dockerignore | 17 +++++++++++++++++
 cog.yaml      | 27 +++++++++++++++++++++++++++
 predict.py    | 29 +++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 .dockerignore
 create mode 100644 cog.yaml
 create mode 100644 predict.py

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..4522d57
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,17 @@
+# The .dockerignore file excludes files from the container build process.
+#
+# https://docs.docker.com/engine/reference/builder/#dockerignore-file
+
+# Exclude Git files
+.git
+.github
+.gitignore
+
+# Exclude Python cache files
+__pycache__
+.mypy_cache
+.pytest_cache
+.ruff_cache
+
+# Exclude Python virtual environment
+/venv
diff --git a/cog.yaml b/cog.yaml
new file mode 100644
index 0000000..a02f5f1
--- /dev/null
+++ b/cog.yaml
@@ -0,0 +1,27 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  gpu: true
+  python_version: "3.11"
+  python_packages:
+    - "torch==2.3.0"
+    - "torchvision==0.18.0"
+    - "torchaudio==2.3.0"
+    - "numpy==1.26.4"
+    - "pyyaml==6.0.1"
+    - "opencv-python==4.10.0.84"
+    - "scipy==1.13.1"
+    - "imageio==2.34.2"
+    - "lmdb==1.4.1"
+    - "tqdm==4.66.4"
+    - "rich==13.7.1"
+    - "ffmpeg==1.4"
+    - "onnxruntime-gpu==1.18.0"
+    - "onnx==1.16.1"
+    - "scikit-image==0.24.0"
+    - "albumentations==1.4.10"
+    - "matplotlib==3.9.0"
+    - "imageio-ffmpeg==0.5.1"
+    - "tyro==0.8.5"
+predict: "predict.py:Predictor"
diff --git a/predict.py b/predict.py
new file mode 100644
index 0000000..f0957b2
--- /dev/null
+++ b/predict.py
@@ -0,0 +1,29 @@
+from cog import BasePredictor, Input, Path
+
+from src.config.argument_config import ArgumentConfig
+from src.config.inference_config import InferenceConfig
+from src.config.crop_config import CropConfig
+from src.live_portrait_pipeline import LivePortraitPipeline
+
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+        self.live_portrait_pipeline = LivePortraitPipeline(
+            inference_cfg=InferenceConfig(),
+            crop_cfg=CropConfig()
+        )
+
+    def predict(
+        self,
+        image: Path = Input(description="Portrait image"),
+        driving_info: Path = Input(
+            description="driving video or template (.pkl format)"
+        ),
+    ) -> Path:
+        """Run a single prediction on the model"""
+        video_path, _ = self.live_portrait_pipeline.execute(
+            ArgumentConfig(source_image=image, driving_info=driving_info, output_dir="/tmp/")
+        )
+
+        return Path(video_path)