resolve failed to read downloaded model from gdrive

app.py

@@ -1,19 +1,27 @@
 import os
+import builtins
 import math
 import streamlit as st
-from google_drive_downloader import GoogleDriveDownloader as gdd
+import gdown
+#from google_drive_downloader import GoogleDriveDownloader as gdd
 
 from demo.src.models import load_trained_model
 from demo.src.utils import render_predict_from_pose, predict_to_image
 from demo.src.config import MODEL_DIR, MODEL_NAME, FILE_ID
 
+st.set_page_config(page_title="DietNeRF Demo")
+
 
-if not os.path.isfile('models'):
-    model_path = os.path.join(MODEL_DIR, MODEL_NAME)
-    gdd.download_file_from_google_drive(file_id=FILE_ID,
-                                        dest_path=model_path,
-                                        unzip=False)
-    print(f'model downloaded from google drive: {model_path}')
+@st.cache
+def download_model():
+    os.makedirs(MODEL_DIR, exist_ok=True)
+    _model_path = os.path.join(MODEL_DIR, MODEL_NAME)
+    # gdd.download_file_from_google_drive(file_id=FILE_ID,
+    #                                     dest_path=_model_path,
+    #                                     unzip=True)
+    url = f'https://drive.google.com/uc?id={FILE_ID}'
+    gdown.download(url, _model_path, quiet=False)
+    print(f'model downloaded from google drive: {_model_path}')
 
 
 @st.cache(show_spinner=False, allow_output_mutation=True)
@@ -22,9 +30,12 @@ def fetch_model():
     return model, state
 
 
+model_path = os.path.join(MODEL_DIR, MODEL_NAME)
+if not os.path.isfile(model_path):
+    download_model()
+
 model, state = fetch_model()
 pi = math.pi
-st.set_page_config(page_title="DietNeRF Demo")
 st.sidebar.header('SELECT YOUR VIEW DIRECTION')
 theta = st.sidebar.slider("Theta", min_value=0., max_value=2.*pi,
                           step=0.5, value=0.)

demo/src/config.py

@@ -1,6 +1,6 @@
 # for downloading model from google drive
-FILE_ID = "1iytA1n2z4go3uVCwE__vIKouTKyIDjEq"
-MODEL_DIR = './models'
+FILE_ID = "1msgvx_jiI-Fr5BrB-otyCwRMSBcUyu-h"
+MODEL_DIR = 'models'
 MODEL_NAME = 'trained_model'
 
 

demo/src/models.py

@@ -5,7 +5,7 @@ from flax.training import checkpoints
 
 from jaxnerf.nerf import models
 from jaxnerf.nerf import utils
-from demo.src.config import NerfConfig
+from demo.src.config import NerfConfig, MODEL_DIR, MODEL_NAME, FILE_ID
 
 rng = random.PRNGKey(0)
 # TODO @Alex: make image size flexible if needed

demo/src/utils.py

@@ -1,3 +1,4 @@
+import os
 from functools import partial
 import jax
 from jax import random

requirements.txt

@@ -11,4 +11,4 @@ tensorflow-hub>=0.11.0
 transformers==4.8.2
 tqdm==4.61.2
 streamlit==0.84.1
-googledrivedownloader==0.4
+pygdrive3==0.8.1