Upload app.py

app.py

@@ -16,10 +16,21 @@ from sentence_transformers import SentenceTransformer
 from diffusers import DiffusionPipeline
 
 
-# Environment variables
+###############################
+## --- GLOBAL VARIABLES ---- ##
+###############################
+
+
+IS_MODEL_LOADED = False
+PATH_JSON = '/home/user/app/.kaggle/kaggle.json'
+
+
+
+# Environment variables to authenticate Kaggle account
 os.environ['KAGGLE_USERNAME'] = st.secrets['username']
 os.environ['KAGGLE_KEY'] = st.secrets['key']
-#os.environ['KAGGLE_CONFIG_DIR'] = '/home/appuser/.kaggle/kaggle.json'
+os.environ['KAGGLE_CONFIG_DIR'] = PATH_JSON
+
 from kaggle.api.kaggle_api_extended import KaggleApi
 
 
@@ -48,10 +59,13 @@ def load_dataset():
 
     return data
 
-@st.cache
+@st.cache(persist=True, show_spinner=False, allow_output_mutation=True, suppress_st_warning=True)
 def load_model():
-    #return DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    return DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+    model = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    IS_MODEL_LOADED = True
+    return model
+    #return DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+    #return DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-2")
 
 
 def query_summarization(text):
@@ -79,12 +93,14 @@ def generate_poster(movie_data):
         synopsis_sum = query_summarization(movie_data.overview.values[0])
     st.text(synopsis_sum)
 
+    # Load text-to-image model
+    if not IS_MODEL_LOADED:
+        with st.spinner("Loading Text to Image model..."):
+            pipeline = load_model()
+    
     # Get image based on synopsis
-    pipeline = load_model()
-    #pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-2")
-
-    #image = pipe(prompt).images[0]
-    #st.image(image, caption=movie_data.title)
+    image = pipeline(prompt).images[0]
+    st.image(image, caption=movie_data.title)
 
     return None #poster_image
 
@@ -93,30 +109,25 @@ def generate_poster(movie_data):
 ## --- CONNECT TO KAGGLE --- ##
 ###############################
 
-'''
-# Authenticate Kaggle account
-os.environ['KAGGLE_USERNAME'] = st.secrets['username']
-os.environ['KAGGLE_KEY'] = st.secrets['key']
-os.environ['KAGGLE_CONFIG_DIR'] = '/home/appuser/.kaggle/kaggle.json'
 
+if not os.path.exists(PATH_JSON):
+    api_token = {"username":st.secrets['username'],"key":st.secrets['key']}
+    with open(PATH_JSON, 'w') as file:
+        json.dump(api_token, file)
 
-api_token = {"username":st.secrets['username'],"key":st.secrets['key']}
-with open('/home/appuser/.kaggle/kaggle.json', 'w') as file:
-    json.dump(api_token, file)
-'''
 
 # Activate Kaggle API
 api = KaggleApi()
 api.authenticate()
 
-'''
+
 try:
     api.authenticate()
 except:
     with open('/home/appuser/.kaggle/kaggle.json', 'w') as file:
         json.dump(api_token, file)
     api.authenticate() 
-'''
+
 
 
 ###############################