ruslanmv commited on
Commit
6eb84b2
1 Parent(s): 6cb51f3

Update utils/default_models.py

Browse files
Files changed (1) hide show
  1. utils/default_models.py +57 -56
utils/default_models.py CHANGED
@@ -1,56 +1,57 @@
1
- import urllib.request
2
- from pathlib import Path
3
- from threading import Thread
4
- from urllib.error import HTTPError
5
-
6
- from tqdm import tqdm
7
-
8
-
9
- default_models = {
10
- "encoder": ("https://drive.google.com/uc?export=download&id=1q8mEGwCkFy23KZsinbuvdKAQLqNKbYf1", 17090379),
11
- "synthesizer": ("https://drive.google.com/u/0/uc?id=1EqFMIbvxffxtjiVrtykroF6_mUh-5Z3s&export=download&confirm=t", 370554559),
12
- "vocoder": ("https://drive.google.com/uc?export=download&id=1cf2NO6FtI0jDuy8AV3Xgn6leO6dHjIgu", 53845290),
13
- }
14
-
15
-
16
- class DownloadProgressBar(tqdm):
17
- def update_to(self, b=1, bsize=1, tsize=None):
18
- if tsize is not None:
19
- self.total = tsize
20
- self.update(b * bsize - self.n)
21
-
22
-
23
- def download(url: str, target: Path, bar_pos=0):
24
- # Ensure the directory exists
25
- target.parent.mkdir(exist_ok=True, parents=True)
26
-
27
- desc = f"Downloading {target.name}"
28
- with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=desc, position=bar_pos, leave=False) as t:
29
- try:
30
- urllib.request.urlretrieve(url, filename=target, reporthook=t.update_to)
31
- except HTTPError:
32
- return
33
-
34
-
35
- def ensure_default_models(models_dir: Path):
36
- # Define download tasks
37
- jobs = []
38
- for model_name, (url, size) in default_models.items():
39
- target_path = models_dir / "default" / f"{model_name}.pt"
40
- if target_path.exists():
41
- if target_path.stat().st_size != size:
42
- print(f"File {target_path} is not of expected size, redownloading...")
43
- else:
44
- continue
45
-
46
- thread = Thread(target=download, args=(url, target_path, len(jobs)))
47
- thread.start()
48
- jobs.append((thread, target_path, size))
49
-
50
- # Run and join threads
51
- for thread, target_path, size in jobs:
52
- thread.join()
53
-
54
- assert target_path.exists() and target_path.stat().st_size == size, \
55
- f"Download for {target_path.name} failed. You may download models manually instead.\n" \
56
- f"https://drive.google.com/drive/folders/1fU6umc5uQAVR2udZdHX-lDgXYzTyqG_j"
 
1
+ import urllib.request
2
+ from pathlib import Path
3
+ from threading import Thread
4
+ from urllib.error import HTTPError
5
+
6
+ from tqdm import tqdm
7
+
8
+
9
+ default_models = {
10
+ "encoder": ("https://drive.google.com/uc?export=download&id=1q8mEGwCkFy23KZsinbuvdKAQLqNKbYf1", 17090379),
11
+ #"synthesizer": ("https://drive.google.com/u/0/uc?id=1EqFMIbvxffxtjiVrtykroF6_mUh-5Z3s&export=download&confirm=t", 370554559),
12
+ "synthesizer": ("https://sourceforge.net/projects/encoders/files/synthesizer.pt/download", 370554559),
13
+ "vocoder": ("https://drive.google.com/uc?export=download&id=1cf2NO6FtI0jDuy8AV3Xgn6leO6dHjIgu", 53845290),
14
+ }
15
+
16
+
17
+ class DownloadProgressBar(tqdm):
18
+ def update_to(self, b=1, bsize=1, tsize=None):
19
+ if tsize is not None:
20
+ self.total = tsize
21
+ self.update(b * bsize - self.n)
22
+
23
+
24
+ def download(url: str, target: Path, bar_pos=0):
25
+ # Ensure the directory exists
26
+ target.parent.mkdir(exist_ok=True, parents=True)
27
+
28
+ desc = f"Downloading {target.name}"
29
+ with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=desc, position=bar_pos, leave=False) as t:
30
+ try:
31
+ urllib.request.urlretrieve(url, filename=target, reporthook=t.update_to)
32
+ except HTTPError:
33
+ return
34
+
35
+
36
+ def ensure_default_models(models_dir: Path):
37
+ # Define download tasks
38
+ jobs = []
39
+ for model_name, (url, size) in default_models.items():
40
+ target_path = models_dir / "default" / f"{model_name}.pt"
41
+ if target_path.exists():
42
+ if target_path.stat().st_size != size:
43
+ print(f"File {target_path} is not of expected size, redownloading...")
44
+ else:
45
+ continue
46
+
47
+ thread = Thread(target=download, args=(url, target_path, len(jobs)))
48
+ thread.start()
49
+ jobs.append((thread, target_path, size))
50
+
51
+ # Run and join threads
52
+ for thread, target_path, size in jobs:
53
+ thread.join()
54
+
55
+ assert target_path.exists() and target_path.stat().st_size == size, \
56
+ f"Download for {target_path.name} failed. You may download models manually instead.\n" \
57
+ f"https://drive.google.com/drive/folders/1fU6umc5uQAVR2udZdHX-lDgXYzTyqG_j"