Update app.py

app.py

@@ -306,13 +306,15 @@ def draw_audio(
 
     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)
             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)