fixing loading model

app.py

@@ -34,7 +34,7 @@ model_folder = 'models/'
 max_img_size = 800
 ################################
 
-# @st.cache(allow_output_mutation=True)
+@st.cache(allow_output_mutation=True)
 def load_model(model_dir, option):
     if option.lower() == 'artistic':
         model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
@@ -132,34 +132,36 @@ color_option = st.sidebar.selectbox('Select colorizer mode',
 try:
     colorizer = load_model(model_folder, color_option)
 except:
+    colorizer = None
     print('Error while loading the model. Please refresh the page')
 
-title_message.markdown("**To begin, please upload an image** 👇")
+if colorizer is not None:
+    title_message.markdown("**To begin, please upload an image** 👇")
 
-#Choose your own image
-uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
+    #Choose your own image
+    uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
 
-# show = st.image(use_column_width='auto')
-input_img_pos = st.empty()
-output_img_pos = st.empty()
+    # show = st.image(use_column_width='auto')
+    input_img_pos = st.empty()
+    output_img_pos = st.empty()
 
-if uploaded_file is not None:
-    img_name = uploaded_file.name
+    if uploaded_file is not None:
+        img_name = uploaded_file.name
 
-    pil_img = PIL.Image.open(uploaded_file)
-    img_rgb = np.array(pil_img)
+        pil_img = PIL.Image.open(uploaded_file)
+        img_rgb = np.array(pil_img)
 
-    resized_img_rgb = resize_img(img_rgb, max_img_size)
-    resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
+        resized_img_rgb = resize_img(img_rgb, max_img_size)
+        resized_pil_img = PIL.Image.fromarray(resized_img_rgb)
 
-    title_message.markdown("**Processing your image, please wait** ⌛")
+        title_message.markdown("**Processing your image, please wait** ⌛")
 
-    output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
-    
-    title_message.markdown("**To begin, please upload an image** 👇")
+        output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
+        
+        title_message.markdown("**To begin, please upload an image** 👇")
 
-    # Plot images
-    input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
-    output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)
+        # Plot images
+        input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
+        output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)
 
-    st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)
+        st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)