jhj0517 commited on
Commit
85f37ae
·
1 Parent(s): 1a63918

Update pytube bot detection exception

Browse files
Files changed (2) hide show
  1. tests/test_config.py +19 -0
  2. tests/test_transcription.py +10 -9
tests/test_config.py CHANGED
@@ -1,4 +1,7 @@
 
 
1
  from modules.utils.paths import *
 
2
 
3
  import os
4
  import torch
@@ -13,5 +16,21 @@ TEST_SUBTITLE_SRT_PATH = os.path.join(WEBUI_DIR, "tests", "test_srt.srt")
13
  TEST_SUBTITLE_VTT_PATH = os.path.join(WEBUI_DIR, "tests", "test_vtt.vtt")
14
 
15
 
 
16
  def is_cuda_available():
17
  return torch.cuda.is_available()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import functools
2
+
3
  from modules.utils.paths import *
4
+ from modules.utils.youtube_manager import *
5
 
6
  import os
7
  import torch
 
16
  TEST_SUBTITLE_VTT_PATH = os.path.join(WEBUI_DIR, "tests", "test_vtt.vtt")
17
 
18
 
19
+ @functools.lru_cache
20
  def is_cuda_available():
21
  return torch.cuda.is_available()
22
+
23
+
24
+ @functools.lru_cache
25
+ def is_pytube_detected_bot(url: str = TEST_YOUTUBE_URL):
26
+ try:
27
+ yt_temp_path = os.path.join("modules", "yt_tmp.wav")
28
+ if os.path.exists(yt_temp_path):
29
+ return False
30
+ yt = get_ytdata(url)
31
+ audio = get_ytaudio(yt)
32
+ return False
33
+ except Exception as e:
34
+ print(f"Pytube has detected as a bot: {e}")
35
+ return True
36
+
tests/test_transcription.py CHANGED
@@ -66,15 +66,16 @@ def test_transcribe(
66
  assert isinstance(subtitle_str, str) and subtitle_str
67
  assert isinstance(file_path[0], str) and file_path
68
 
69
- whisper_inferencer.transcribe_youtube(
70
- TEST_YOUTUBE_URL,
71
- "SRT",
72
- False,
73
- gr.Progress(),
74
- *hparams,
75
- )
76
- assert isinstance(subtitle_str, str) and subtitle_str
77
- assert isinstance(file_path[0], str) and file_path
 
78
 
79
  whisper_inferencer.transcribe_mic(
80
  audio_path,
 
66
  assert isinstance(subtitle_str, str) and subtitle_str
67
  assert isinstance(file_path[0], str) and file_path
68
 
69
+ if not is_pytube_detected_bot():
70
+ whisper_inferencer.transcribe_youtube(
71
+ TEST_YOUTUBE_URL,
72
+ "SRT",
73
+ False,
74
+ gr.Progress(),
75
+ *hparams,
76
+ )
77
+ assert isinstance(subtitle_str, str) and subtitle_str
78
+ assert isinstance(file_path[0], str) and file_path
79
 
80
  whisper_inferencer.transcribe_mic(
81
  audio_path,