Update app.py

app.py

@@ -188,7 +188,7 @@ class CLIPDemo:
     def compute_image_embeddings(self, image_paths: list):
         self.image_paths = image_paths
         dataloader = DataLoader(VisionDataset(
-            image_paths=image_paths), batch_size=self.batch_size, num_workers=8)
+            image_paths=image_paths), batch_size=self.batch_size)
         embeddings = []
         with torch.no_grad():
             
@@ -316,7 +316,7 @@ def draw_audio(
             vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True).to(device)
             tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
             model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer, device=device)
-            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH+"/*.jpeg")[:5000])
+            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH+"/*.jpeg")[:1000])
             st.session_state["model"] = model
             #st.session_state['model'] = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer)
             #st.session_state.model.compute_image_embeddings(glob.glob("/data1/mlaquatra/TSOAI_hack/data/spectrograms/*.jpeg")[:100])
@@ -384,7 +384,7 @@ def draw_camera(
             vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True).to(device)
             tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
             model = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer, device=device)
-            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH + "/*.jpeg")[:5000])
+            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH + "/*.jpeg")[:1000])
             st.session_state["model"] = model
             #st.session_state['model'] = CLIPDemo(vision_encoder=vision_encoder, text_encoder=text_encoder, tokenizer=tokenizer)
             #st.session_state.model.compute_image_embeddings(glob.glob("/data1/mlaquatra/TSOAI_hack/data/spectrograms/*.jpeg")[:100])