From cb0e7f2a8059948742aae8a59a792740067b2cc4 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Fri, 10 Apr 2026 11:25:00 +0800
Subject: [PATCH] =?UTF-8?q?build:=20=E4=BF=AE=E6=94=B9docker=20=E6=9E=84?=
 =?UTF-8?q?=E5=BB=BA=E6=B5=81=E7=A8=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .dockerignore            |  3 ++-
 Dockerfile               | 28 ++++++++++------------------
 scripts/demo/download.py | 38 +++++++++++++++++++++++++++++++++-----
 3 files changed, 45 insertions(+), 24 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index 22aee8a..a5bbc7b 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -3,6 +3,7 @@
 
 # Allow necessary files
 !astrai/
-!scripts/tools/
+!scripts/
+!assets/
 !pyproject.toml
 !README.md
diff --git a/Dockerfile b/Dockerfile
index 3b2de34..231a3e6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,49 +1,41 @@
-# AstrAI Dockerfile
-# Multi-stage build for optimized image size
+# AstrAI Dockerfile - Minimal
 
 # Build stage
 FROM python:3.12-slim AS builder
 
 WORKDIR /app
 
-# Install build dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
     build-essential \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy project files first for version extraction
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
 
-# Install dependencies
 RUN pip install --no-cache-dir --upgrade pip \
-    && pip install --no-cache-dir .[dev]
+    && pip install --no-cache-dir .
 
 # Production stage
-FROM python:3.12-slim AS production
+FROM nvidia/cuda:12.6.0-runtime-ubuntu22.04 AS production
 
 WORKDIR /app
 
-# Install runtime dependencies
 RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 \
+    libpython3.12 \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy installed packages from builder
 COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
 COPY --from=builder /usr/local/bin /usr/local/bin
 
-# Copy application code
 COPY astrai/ ./astrai/
-COPY scripts/tools/ ./scripts/tools/
+COPY scripts/ ./scripts/
+COPY assets/ ./assets/
 COPY pyproject.toml .
+COPY README.md .
 
-# Create non-root user
 RUN useradd -m -u 1000 astrai && chown -R astrai:astrai /app
 USER astrai
 
-# Set environment variables
-ENV PYTHONUNBUFFERED=1
-ENV PYTHONDONTWRITEBYTECODE=1
-
-# Default command
-CMD ["python", "-m", "astrai.inference.server"]
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1
\ No newline at end of file
diff --git a/scripts/demo/download.py b/scripts/demo/download.py
index 790c5a1..43148f4 100644
--- a/scripts/demo/download.py
+++ b/scripts/demo/download.py
@@ -1,13 +1,41 @@
+import argparse
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
-PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+DEFAULT_LOCAL_DIR = Path(PROJECT_ROOT, "params")
+DEFAULT_REPO_ID = "ViperEk/KHAOSZ"
 
 if __name__ == "__main__":
-    snapshot_download(
-        repo_id="ViperEk/KHAOSZ",
-        local_dir=PARAMETER_ROOT,
-        force_download=True,
+    parser = argparse.ArgumentParser(
+        description="Download model parameters from HuggingFace"
     )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default=DEFAULT_REPO_ID,
+        help=f"HuggingFace repo ID (default: {DEFAULT_REPO_ID})",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=DEFAULT_LOCAL_DIR,
+        help=f"Local directory to save model (default: {DEFAULT_LOCAL_DIR})",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force download even if files exist",
+    )
+    args = parser.parse_args()
+
+    print(f"Downloading model from {args.repo_id} to {args.local_dir}")
+
+    snapshot_download(
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+        force_download=args.force,
+    )
+
+    print("Download complete!")