JulianPhillips commited on
Commit
1cae688
·
verified ·
1 Parent(s): 6a6e71d

Update Dockerfile

Browse files
Files changed (1) hide show
  1. Dockerfile +11 -3
Dockerfile CHANGED
@@ -1,7 +1,7 @@
1
  # Base image with PyTorch and CUDA for GPU support
2
  FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-runtime
3
 
4
- # Install system dependencies and Hugging Face CLI
5
  RUN apt-get update && apt-get install -y \
6
  git \
7
  wget \
@@ -24,8 +24,16 @@ RUN mkdir -p /tmp/cache
24
  RUN mkdir -p /models/sapiens_pose /models/motionbert
25
 
26
  # Python script to download models
27
- RUN echo "from transformers import AutoModel, AutoTokenizer\n\
28
- AutoModel.from_pretrained('facebook/sapiens-pose-1b-torchscript').save_pretrained('/models/sapiens_pose')\n\
 
 
 
 
 
 
 
 
29
  AutoModel.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')\n\
30
  AutoTokenizer.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')" > download_models.py
31
 
 
1
  # Base image with PyTorch and CUDA for GPU support
2
  FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-runtime
3
 
4
+ # Install system dependencies
5
  RUN apt-get update && apt-get install -y \
6
  git \
7
  wget \
 
24
  RUN mkdir -p /models/sapiens_pose /models/motionbert
25
 
26
  # Python script to download models
27
+ RUN echo "import requests\n\
28
+ url = 'https://huggingface.co/facebook/sapiens-pose-1b-torchscript/resolve/main/model.pt'\n\
29
+ response = requests.get(url)\n\
30
+ if response.status_code == 200:\n\
31
+ with open('/models/sapiens_pose/model.pt', 'wb') as f:\n\
32
+ f.write(response.content)\n\
33
+ else:\n\
34
+ raise Exception(f'Failed to download model: {response.status_code}')\n\
35
+ \n\
36
+ from transformers import AutoModel, AutoTokenizer\n\
37
  AutoModel.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')\n\
38
  AutoTokenizer.from_pretrained('walterzhu/MotionBERT').save_pretrained('/models/motionbert')" > download_models.py
39