Update Dockerfile

Dockerfile

@@ -1,36 +1,32 @@
-# Dockerfile
 FROM ubuntu:24.04
 
+ARG HF_TOKEN
+ENV HF_TOKEN=${HF_TOKEN}
 ENV DEBIAN_FRONTEND=noninteractive
 
-# Install essentials
 RUN apt-get update && apt-get install -y \
     python3 python3-venv python3-pip git wget \
   && rm -rf /var/lib/apt/lists/*
 
 WORKDIR /workspace
 
-# Clone ComfyUI
 RUN git clone https://github.com/comfyanonymous/ComfyUI.git
-
 WORKDIR /workspace/ComfyUI
 
-# Set up Python venv & install
 RUN python3 -m venv .venv \
   && . .venv/bin/activate \
   && pip install --upgrade pip \
   && pip install xformers!=0.0.18 \
   && pip install -r requirements.txt
 
-# Download Flux.1 Schnell UNET & VAE
 RUN mkdir -p models/unet models/vae \
-  && wget -c https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors \
-     -P models/unet/ \
-  && wget -c https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors \
-     -P models/vae/
+  && wget --header="Authorization: Bearer ${HF_TOKEN}" \
+     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_TOKEN}" \
+     https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors \
+     -O models/vae/ae.safetensors
 
-# Expose ComfyUI port
 EXPOSE 7860
 
-# Launch script
 CMD ["/bin/bash", "-lc", "source .venv/bin/activate && python3 main.py --host 0.0.0.0 --port 7860 --cpu"]