Update Dockerfile

Dockerfile

@@ -2,26 +2,45 @@
 FROM ubuntu:24.04
 ENV DEBIAN_FRONTEND=noninteractive
 
-RUN apt-get update && apt-get install -y python3-venv python3-pip git wget && \
-    rm -rf /var/lib/apt/lists/*
+RUN apt-get update && apt-get install -y \
+    python3 python3-venv python3-pip git wget \
+ && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
 RUN git clone https://github.com/comfyanonymous/ComfyUI.git
 WORKDIR /workspace/ComfyUI
 
-RUN python3 -m venv .venv && . .venv/bin/activate && \
-    pip install --upgrade pip && pip install xformers!=0.0.18 && \
-    pip install -r requirements.txt
+# Install dependencies including huggingface_hub
+RUN python3 -m venv .venv \
+ && . .venv/bin/activate \
+ && pip install --upgrade pip \
+ && pip install huggingface_hub xformers!=0.0.18 \
+ && pip install -r requirements.txt
 
+# Download model files securely using hf_hub_download
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true \
-    mkdir -p models/unet models/vae && \
-    hf=$(cat /run/secrets/HF_TOKEN) && \
-    wget --header="Authorization: Bearer $hf" \
-      https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors \
-      -O models/unet/flux1-schnell.safetensors && \
-    wget --header="Authorization: Bearer $hf" \
-      https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors \
-      -O models/vae/ae.safetensors
+    . .venv/bin/activate && \
+    python3 - << 'EOF'
+import os
+from huggingface_hub import hf_hub_download
+hf = open('/run/secrets/HF_TOKEN').read().strip()
+# UNET
+hf_hub_download(
+    repo_id="black-forest-labs/FLUX.1-schnell",
+    filename="flux1-schnell.safetensors",
+    token=hf,
+    local_dir="models/unet",
+    force_download=True
+)
+# VAE
+hf_hub_download(
+    repo_id="black-forest-labs/FLUX.1-schnell",
+    filename="ae.safetensors",
+    token=hf,
+    local_dir="models/vae",
+    force_download=True
+)
+EOF
 
 EXPOSE 7860