Update Dockerfile

Dockerfile

@@ -10,37 +10,19 @@ WORKDIR /workspace
 RUN git clone https://github.com/comfyanonymous/ComfyUI.git
 WORKDIR /workspace/ComfyUI
 
-# 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
+# Copy download script
+COPY download_flux.py /workspace/ComfyUI/download_flux.py
+
+# Use secret to download model files via huggingface_hub
 RUN --mount=type=secret,id=HF_TOKEN,mode=0444,required=true \
     . .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
+    python download_flux.py
 
 EXPOSE 7860