From c6168327e5ae9f4ff4d6d1d52519e7d5db27ac34 Mon Sep 17 00:00:00 2001
From: awwaawwa <8493196+awwaawwa@users.noreply.github.com>
Date: Wed, 18 Dec 2024 19:37:47 +0800
Subject: [PATCH] =?UTF-8?q?=E8=B0=83=E6=95=B4docker=E6=9E=84=E5=BB=BA?=
 =?UTF-8?q?=E9=A1=BA=E5=BA=8F=EF=BC=8C=E5=85=88=E4=B8=8B=E6=A8=A1=E5=9E=8B?=
 =?UTF-8?q?=E5=90=8E=E8=A3=85=E5=8C=85=EF=BC=8C=E5=B9=B6=E5=B0=86=E9=A2=84?=
 =?UTF-8?q?=E7=83=AD=E6=9B=B4=E6=94=B9=E4=B8=BAdockerfile=20RUN=E6=8C=87?=
 =?UTF-8?q?=E4=BB=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Dockerfile          |  8 +++++---
 pdf2zh/doclayout.py |  7 +------
 pdf2zh/warmup.py    | 10 ----------
 3 files changed, 6 insertions(+), 19 deletions(-)
 delete mode 100644 pdf2zh/warmup.py

diff --git a/Dockerfile b/Dockerfile
index 5a006be1..91b8920e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,7 +2,6 @@ FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
 
 WORKDIR /app
 
-COPY . .
 
 EXPOSE 7860
 
@@ -10,8 +9,11 @@ ENV PYTHONUNBUFFERED=1
 ADD "https://github.com/satbyy/go-noto-universal/releases/download/v7.0/GoNotoKurrent-Regular.ttf" /app
 RUN apt-get update && \
      apt-get install --no-install-recommends -y libgl1 && \
-     rm -rf /var/lib/apt/lists/*
+     rm -rf /var/lib/apt/lists/* && uv pip install --system --no-cache huggingface-hub && \
+     python3 -c "from huggingface_hub import hf_hub_download; hf_hub_download('wybxc/DocLayout-YOLO-DocStructBench-onnx','doclayout_yolo_docstructbench_imgsz1024.onnx');"
+
+COPY . .
 
-RUN uv pip install --system --no-cache . && uv run pdf2zh/warmup.py
+RUN uv pip install --system --no-cache .
 
 CMD ["pdf2zh", "-i"]
\ No newline at end of file
diff --git a/pdf2zh/doclayout.py b/pdf2zh/doclayout.py
index cd064d90..c2b4fa1f 100644
--- a/pdf2zh/doclayout.py
+++ b/pdf2zh/doclayout.py
@@ -83,12 +83,7 @@ def from_pretrained(repo_id: str, filename: str):
             model_dir = snapshot_download(repo_mapping[repo_id])
             pth = os.path.join(model_dir, filename)
         else:
-            try:
-                pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1, local_files_only=True)
-                print("Using local DocLayout-YOLO-DocStructBench-onnx file", pth)
-            except huggingface_hub.utils.LocalEntryNotFoundError:
-                print("Downloading DocLayout-YOLO-DocStructBench-onnx from Huggingface Hub...")
-                pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1)
+            pth = hf_hub_download(repo_id=repo_id, filename=filename, etag_timeout=1)
         return OnnxModel(pth)
 
     @property
diff --git a/pdf2zh/warmup.py b/pdf2zh/warmup.py
deleted file mode 100644
index 89c671a9..00000000
--- a/pdf2zh/warmup.py
+++ /dev/null
@@ -1,10 +0,0 @@
-from doclayout import DocLayoutModel
-
-def warmup():
-    print('Warming up the model...')
-    DocLayoutModel.load_available()
-    print('Warm-up completed!')
-
-
-if __name__ == "__main__":
-    warmup()
\ No newline at end of file