Merge branch 'main' of https://huggingface.co/spaces/ALM/CALM into main

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():
             
@@ -249,17 +249,19 @@ class CLIPDemo:
 def draw_text(
     key,
     plot=False,
+    device=None,
 ):
 
+    
     image = Image.open("data/logo.png")
     st.image(image, use_column_width="always")
 
     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)
-            vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True)
+            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)
+            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,18 +304,19 @@ def draw_text(
 def draw_audio(
     key,
     plot=False,
+    device=None,
 ):
 
     image = Image.open("data/logo.png")
     st.image(image, use_column_width="always")
-
+    
     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)
-            vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True)
+            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)
-            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH+"/*.jpeg")[:5000])
+            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
             #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])
@@ -369,6 +372,7 @@ def draw_audio(
 def draw_camera(
     key,
     plot=False,
+    device=None,
 ):
 
     image = Image.open("data/logo.png")
@@ -377,10 +381,10 @@ def draw_camera(
     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)
-            vision_encoder = CLIPVisionModel.from_pretrained(CLIP_VISION_MODEL_PATH, local_files_only=True)
+            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)
-            model.compute_image_embeddings(glob.glob(SPECTROGRAMS_PATH + "/*.jpeg")[:5000])
+            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
             #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])
@@ -427,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:

requirements.txt

@@ -7,6 +7,7 @@ bokeh
 streamlit_bokeh_events
 streamlit-webcam-example
 torch
+torchvision
 numpy
 pandas
 tqdm