From 8a12434b8792d5d58f5ee688439d05c82799436f Mon Sep 17 00:00:00 2001
From: zhangpeihao <zhangpeihao@kingsoft.com>
Date: Tue, 3 Dec 2024 11:51:02 +0800
Subject: [PATCH] feat: add support for docker deploy

---
 .dockerignore       | 27 +++++++++++++++++++++++++++
 Dockerfile          | 17 +++++++++++++++++
 docker-compose.yaml | 12 ++++++++++++
 requirements.txt    |  2 +-
 4 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 .dockerignore
 create mode 100644 Dockerfile
 create mode 100644 docker-compose.yaml

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 000000000..2d46c794c
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,27 @@
+__pycache__/
+*.py[cod]
+/output/
+/input/
+!/input/example.png
+/models/
+/temp/
+/custom_nodes/
+!custom_nodes/example_node.py.example
+extra_model_paths.yaml
+/.vs
+.vscode/
+.idea/
+venv/
+.venv/
+/web/extensions/*
+!/web/extensions/logging.js.example
+!/web/extensions/core/
+/tests-ui/data/object_info.json
+/user/
+*.log
+web_custom_versions/
+.DS_Store
+
+Dockerfile
+docker-compose.yml
+*.whl
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 000000000..07b204f8c
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,17 @@
+FROM python:3.12-slim
+WORKDIR /app
+ENV DEBIAN_FRONTEND=noninteractive \
+    PIP_NO_CACHE_DIR=on
+RUN --mount=type=cache,target=/var/cache/apt \
+    apt-get update && \
+    apt-get install git -y && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt requirements.txt
+RUN pip install torch==2.3.1 torchvision==0.18.1 torchaudio --extra-index-url https://download.pytorch.org/whl/cu124 \
+    && pip install -r requirements.txt
+
+COPY . .
+
+CMD ["python3", "main.py"]
\ No newline at end of file
diff --git a/docker-compose.yaml b/docker-compose.yaml
new file mode 100644
index 000000000..7ebbb812b
--- /dev/null
+++ b/docker-compose.yaml
@@ -0,0 +1,12 @@
+services:
+  comfyui:
+    build: .
+    ports:
+      - "8188:8188"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: "nvidia"
+              count: "all"
+              capabilities: ["gpu"]
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 4c2c0b2b2..f805a4fd1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,4 @@
-torch
+torch==2.3.1
 torchsde
 torchvision
 torchaudio