Update

app.py

@@ -18,7 +18,10 @@ from pytube import YouTube
 from youtube_transcript_api import YouTubeTranscriptApi
 from youtube_transcript_api.formatters import WebVTTFormatter
 
-device = 'cpu'
+if torch.cuda.is_available():
+ device = 'cuda'
+else:
+ device = 'cpu'
 model_name = 'BridgeTower/bridgetower-large-itm-mlm-itc'
 model = BridgeTowerForITC.from_pretrained(model_name).to(device)
 text_model = BridgeTowerTextFeatureExtractor.from_pretrained(model_name).to(device)
@@ -282,7 +285,7 @@ def process(video_url, text_query):
  expanded=False,
  batch_size=8,
  )
- frame_paths, transcripts = run_query(video_id, text_query, path=output_dir)
+ frame_paths, transcripts = run_query(video_file, text_query, path=output_dir)
  return video_file, [(image, caption) for image, caption in zip(frame_paths, transcripts)]