Update Dockerfile

Dockerfile

@@ -7,32 +7,36 @@ RUN apt-get update && apt-get install -y \
     wget \
     && rm -rf /var/lib/apt/lists/*
 
-# Install Python packages including Hugging Face Transformers, TorchScript, and Flask
+# Install Python packages including Hugging Face Transformers, TorchScript, Flask, and TensorFlow
 RUN pip install --no-cache-dir \
     torch \
     torchvision \
     transformers \
     requests \
     Flask \
-    Pillow
+    Pillow \
+    huggingface_hub \
+    tensorflow
 
 # Set Hugging Face cache to a guaranteed writable directory
 ENV TRANSFORMERS_CACHE=/tmp/cache
 RUN mkdir -p /tmp/cache
 
 # Create directories for the models
-RUN mkdir -p /models/sapiens_pose /models/motionbert
-
-# Python script to download models
-RUN echo "import requests\n\
-url = 'https://huggingface.co/facebook/sapiens-pose-1b-torchscript/resolve/main/model.pt'\n\
-response = requests.get(url)\n\
-if response.status_code == 200:\n\
-    with open('/models/sapiens_pose/model.pt', 'wb') as f:\n\
-        f.write(response.content)\n\
-else:\n\
-    raise Exception(f'Failed to download model: {response.status_code}')\n\
+RUN mkdir -p /models/movenet /models/motionbert
+
+# Python script to download models using tensorflow_hub and huggingface_hub
+RUN echo "import os\n\
+import tensorflow_hub as hub\n\
+\n\
+# Download MoveNet model from TensorFlow Hub\n\
+movenet_model = hub.load('https://tfhub.dev/google/movenet/singlepose/lightning/4')\n\
+movenet_model_path = '/models/movenet/movenet_lightning'\n\
+os.makedirs(movenet_model_path, exist_ok=True)\n\
+movenet_model.save(movenet_model_path)\n\
 \n\
+# Download MotionBERT model and tokenizer using huggingface_hub\n\
+from huggingface_hub import hf_hub_download\n\
 from transformers import AutoModel, AutoTokenizer\n\
 AutoModel.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')\n\
 AutoTokenizer.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')" > download_models.py
@@ -48,3 +52,4 @@ EXPOSE 7860
 
 # Run the Flask app
 CMD ["python", "/app/app.py"]
+