Update app.py

app.py

@@ -249,8 +249,10 @@ class CLIPDemo:
 def draw_text(
     key,
     plot=False,
+    device=None,
 ):
 
+    
     image = Image.open("data/logo.png")
     st.image(image, use_column_width="always")
 
@@ -259,7 +261,7 @@ def draw_text(
             text_encoder = AutoModel.from_pretrained(CLIP_TEXT_MODEL_PATH, local_files_only=True)
             vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True)
             tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
-            model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer)
+            model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer, device=device)
             model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH + "/*.jpeg")[:1000])
             st.session_state["model"] = model
 
@@ -302,13 +304,12 @@ def draw_text(
 def draw_audio(
     key,
     plot=False,
+    device=None,
 ):
 
     image = Image.open("data/logo.png")
     st.image(image, use_column_width="always")
     
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-
     if 'model' not in st.session_state:
         #with st.spinner('We are orginizing your traks...'):
             text_encoder = AutoModel.from_pretrained(CLIP_TEXT_MODEL_PATH, local_files_only=True)
@@ -371,6 +372,7 @@ def draw_audio(
 def draw_camera(
     key,
     plot=False,
+    device=None,
 ):
 
     image = Image.open("data/logo.png")
@@ -381,7 +383,7 @@ def draw_camera(
             text_encoder = AutoModel.from_pretrained(CLIP_TEXT_MODEL_PATH, local_files_only=True)
             vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True)
             tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
-            model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer)
+            model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer, device=device)
             model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH + "/*.jpeg")[:5000])
             st.session_state["model"] = model
             #st.session_state['model'] = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer)
@@ -429,15 +431,17 @@ def draw_camera(
 selected = streamlit_menu(example=3)
 df = pd.read_csv('full_metadata.csv', index_col=False)
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
 if selected == "Text":
     # st.title(f"You have selected {selected}")
-    draw_text("text", plot=True)
+    draw_text("text", plot=True, device=device)
 if selected == "Audio":
     # st.title(f"You have selected {selected}")
-    draw_audio("audio", plot=True)
+    draw_audio("audio", plot=True, device=device)
 if selected == "Camera":
     # st.title(f"You have selected {selected}")
-    #draw_camera("camera", plot=True)
+    #draw_camera("camera", plot=True, device=device)
     pass
 
 # with st.sidebar: