Update Dockerfile

Dockerfile

@@ -1,5 +1,9 @@
 FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime
 
+# Set ARG and ENV for HF token (used in snapshot_download)
+ARG HF_TOKEN
+ENV HF_TOKEN=${HF_TOKEN}
+
 # Install OS-level dependencies
 RUN apt-get update && apt-get install -y \
     git ffmpeg curl \
@@ -11,28 +15,41 @@ RUN apt-get update && apt-get install -y \
 WORKDIR /app
 COPY . /app
 
-# Create directory for weights
+# Create weights directory
 RUN mkdir -p /weights/FVSM /weights/OMSM /weights/output /weights/assets/manual_poses
 
-# Download FVSM model weights
-RUN curl -L -o /weights/FVSM/FloVD_FVSM_Controlnet.pt https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/FloVD_FVSM_Controlnet.pt
+# Install required Python packages
+RUN pip install --upgrade pip
+RUN pip install av huggingface_hub
+RUN pip install -r requirements.txt
+RUN pip install --no-binary :all: av  # optional av re-install
 
-# Clone OMSM weights from Hugging Face
-RUN git clone https://huggingface.co/datasets/mutqa/FloVD-HF-OMSM /weights/OMSM
+# --------------------------
+# Download FloVD FVSM model
+# --------------------------
+RUN curl -L -o /weights/FVSM/FloVD_FVSM_Controlnet.pt https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/FloVD_FVSM_Controlnet.pt
 
+# --------------------------
+# Download OMSM weights using huggingface_hub (requires HF_TOKEN)
+# --------------------------
+RUN python3 -c "\
+from huggingface_hub import snapshot_download; \
+snapshot_download(\
+    repo_id='mutqa/FloVD-HF-OMSM', \
+    repo_type='dataset', \
+    local_dir='/weights/OMSM', \
+    token='${HF_TOKEN}'\
+)"
+
+# --------------------------
 # Download example camera pose
+# --------------------------
 RUN curl -L -o /weights/assets/manual_poses/example.txt https://huggingface.co/datasets/mutqa/FloVD-HF-Assets/resolve/main/example.txt
 
-# Ensure permissions
-RUN chmod -R 777 /app
-
-# Install Python dependencies
-RUN pip install --upgrade pip
-RUN pip install av  # workaround for some ffmpeg+av issues
-RUN pip install -r requirements.txt
-RUN pip install --no-binary :all: av  # optional re-install to ensure compatibility
+# Set permissions
+RUN chmod -R 777 /weights /app
 
 EXPOSE 7860
 
-# Run app
+# Start the app
 CMD ["python", "app.py"]