Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from super_image import EdsrModel, ImageLoader
 from PIL import Image
+import cv2
 
 st.title('Take your picture to the next level')
 
@@ -24,11 +25,12 @@ with col1:
             inputs = ImageLoader.load_image(image)
             preds = model(inputs)
 
-            for i in range(101):
-                st.progress(i, text='Magic is happening')
+            ImageLoader.save_image(preds, './scaled_2x.png')
+            img = cv2.imread('./scaled_2x.png')
+            
             st.subheader('New image')
-            st.image(preds, caption=f'Upscaled {choice} times image')
-            st.download_button('Download', preds)
+            st.image(img, caption=f'Upscaled {choice} times image')
+            st.download_button('Download', img)
             
         else:
             st.text('Not supported yet. Comeback soon for fun releases')