354 lines
13 KiB
Python
354 lines
13 KiB
Python
import os
|
|
import sys
|
|
import wget
|
|
import zipfile
|
|
from bs4 import BeautifulSoup
|
|
import requests
|
|
from urllib.parse import unquote, urlencode, parse_qs, urlparse
|
|
import re
|
|
import shutil
|
|
import six
|
|
|
|
|
|
def find_folder_parent(search_dir, folder_name):
|
|
for dirpath, dirnames, _ in os.walk(search_dir):
|
|
if folder_name in dirnames:
|
|
return os.path.abspath(dirpath)
|
|
return None
|
|
|
|
|
|
now_dir = os.getcwd()
|
|
sys.path.append(now_dir)
|
|
|
|
from rvc.lib.utils import format_title
|
|
|
|
from rvc.lib.tools import gdown
|
|
|
|
file_path = find_folder_parent(now_dir, "logs")
|
|
|
|
zips_path = os.getcwd() + "/logs/zips"
|
|
|
|
|
|
def search_pth_index(folder):
|
|
pth_paths = [
|
|
os.path.join(folder, file)
|
|
for file in os.listdir(folder)
|
|
if os.path.isfile(os.path.join(folder, file)) and file.endswith(".pth")
|
|
]
|
|
index_paths = [
|
|
os.path.join(folder, file)
|
|
for file in os.listdir(folder)
|
|
if os.path.isfile(os.path.join(folder, file)) and file.endswith(".index")
|
|
]
|
|
|
|
return pth_paths, index_paths
|
|
|
|
|
|
def get_mediafire_download_link(url):
|
|
response = requests.get(url)
|
|
response.raise_for_status()
|
|
soup = BeautifulSoup(response.text, "html.parser")
|
|
download_button = soup.find(
|
|
"a", {"class": "input popsok", "aria-label": "Download file"}
|
|
)
|
|
if download_button:
|
|
download_link = download_button.get("href")
|
|
return download_link
|
|
else:
|
|
return None
|
|
|
|
|
|
def download_from_url(url):
|
|
os.makedirs(zips_path, exist_ok=True)
|
|
if url != "":
|
|
if "drive.google.com" in url:
|
|
if "file/d/" in url:
|
|
file_id = url.split("file/d/")[1].split("/")[0]
|
|
elif "id=" in url:
|
|
file_id = url.split("id=")[1].split("&")[0]
|
|
else:
|
|
return None
|
|
|
|
if file_id:
|
|
os.chdir(zips_path)
|
|
try:
|
|
gdown.download(
|
|
f"https://drive.google.com/uc?id={file_id}",
|
|
quiet=True,
|
|
fuzzy=True,
|
|
)
|
|
except Exception as error:
|
|
error_message = str(error)
|
|
if (
|
|
"Too many users have viewed or downloaded this file recently"
|
|
in error_message
|
|
):
|
|
os.chdir(now_dir)
|
|
return "too much use"
|
|
elif (
|
|
"Cannot retrieve the public link of the file." in error_message
|
|
):
|
|
os.chdir(now_dir)
|
|
return "private link"
|
|
else:
|
|
print(error_message)
|
|
os.chdir(now_dir)
|
|
return None
|
|
elif "disk.yandex.ru" in url:
|
|
base_url = "https://cloud-api.yandex.net/v1/disk/public/resources/download?"
|
|
public_key = url
|
|
final_url = base_url + urlencode(dict(public_key=public_key))
|
|
response = requests.get(final_url)
|
|
download_url = response.json()["href"]
|
|
download_response = requests.get(download_url)
|
|
|
|
if download_response.status_code == 200:
|
|
filename = parse_qs(urlparse(unquote(download_url)).query).get(
|
|
"filename", [""]
|
|
)[0]
|
|
if filename:
|
|
os.chdir(zips_path)
|
|
with open(filename, "wb") as f:
|
|
f.write(download_response.content)
|
|
else:
|
|
print("Failed to get filename from URL.")
|
|
return None
|
|
|
|
elif "pixeldrain.com" in url:
|
|
try:
|
|
file_id = url.split("pixeldrain.com/u/")[1]
|
|
os.chdir(zips_path)
|
|
print(file_id)
|
|
response = requests.get(f"https://pixeldrain.com/api/file/{file_id}")
|
|
if response.status_code == 200:
|
|
file_name = (
|
|
response.headers.get("Content-Disposition")
|
|
.split("filename=")[-1]
|
|
.strip('";')
|
|
)
|
|
os.makedirs(zips_path, exist_ok=True)
|
|
with open(os.path.join(zips_path, file_name), "wb") as newfile:
|
|
newfile.write(response.content)
|
|
os.chdir(file_path)
|
|
return "downloaded"
|
|
else:
|
|
os.chdir(file_path)
|
|
return None
|
|
except Exception as e:
|
|
print(e)
|
|
os.chdir(file_path)
|
|
return None
|
|
|
|
elif "cdn.discordapp.com" in url:
|
|
file = requests.get(url)
|
|
os.chdir(zips_path)
|
|
if file.status_code == 200:
|
|
name = url.split("/")
|
|
with open(os.path.join(name[-1]), "wb") as newfile:
|
|
newfile.write(file.content)
|
|
else:
|
|
return None
|
|
elif "/blob/" in url or "/resolve/" in url:
|
|
os.chdir(zips_path)
|
|
if "/blob/" in url:
|
|
url = url.replace("/blob/", "/resolve/")
|
|
|
|
response = requests.get(url, stream=True)
|
|
if response.status_code == 200:
|
|
content_disposition = six.moves.urllib_parse.unquote(
|
|
response.headers["Content-Disposition"]
|
|
)
|
|
m = re.search(r'filename="([^"]+)"', content_disposition)
|
|
file_name = m.groups()[0]
|
|
file_name = file_name.replace(os.path.sep, "_")
|
|
total_size_in_bytes = int(response.headers.get("content-length", 0))
|
|
block_size = 1024
|
|
progress_bar_length = 50
|
|
progress = 0
|
|
|
|
with open(os.path.join(zips_path, file_name), "wb") as file:
|
|
for data in response.iter_content(block_size):
|
|
file.write(data)
|
|
progress += len(data)
|
|
progress_percent = int((progress / total_size_in_bytes) * 100)
|
|
num_dots = int(
|
|
(progress / total_size_in_bytes) * progress_bar_length
|
|
)
|
|
progress_bar = (
|
|
"["
|
|
+ "." * num_dots
|
|
+ " " * (progress_bar_length - num_dots)
|
|
+ "]"
|
|
)
|
|
print(
|
|
f"{progress_percent}% {progress_bar} {progress}/{total_size_in_bytes} ",
|
|
end="\r",
|
|
)
|
|
if progress_percent == 100:
|
|
print("\n")
|
|
|
|
else:
|
|
os.chdir(now_dir)
|
|
return None
|
|
elif "/tree/main" in url:
|
|
os.chdir(zips_path)
|
|
response = requests.get(url)
|
|
soup = BeautifulSoup(response.content, "html.parser")
|
|
temp_url = ""
|
|
for link in soup.find_all("a", href=True):
|
|
if link["href"].endswith(".zip"):
|
|
temp_url = link["href"]
|
|
break
|
|
if temp_url:
|
|
url = temp_url
|
|
url = url.replace("blob", "resolve")
|
|
if "huggingface.co" not in url:
|
|
url = "https://huggingface.co" + url
|
|
|
|
wget.download(url)
|
|
else:
|
|
os.chdir(now_dir)
|
|
return None
|
|
|
|
params = {"id": f"eq.{id_number}"}
|
|
response = requests.get(url, headers=headers, params=params)
|
|
if response.status_code == 200:
|
|
json_response = response.json()
|
|
print(json_response)
|
|
if json_response:
|
|
link = json_response[0]["link"]
|
|
verify = download_from_url(link)
|
|
if verify == "downloaded":
|
|
return "downloaded"
|
|
else:
|
|
return None
|
|
else:
|
|
return None
|
|
else:
|
|
try:
|
|
os.chdir(zips_path)
|
|
wget.download(url)
|
|
except Exception as error:
|
|
os.chdir(now_dir)
|
|
print(error)
|
|
return None
|
|
|
|
for currentPath, _, zipFiles in os.walk(zips_path):
|
|
for Files in zipFiles:
|
|
filePart = Files.split(".")
|
|
extensionFile = filePart[len(filePart) - 1]
|
|
filePart.pop()
|
|
nameFile = "_".join(filePart)
|
|
realPath = os.path.join(currentPath, Files)
|
|
os.rename(realPath, nameFile + "." + extensionFile)
|
|
|
|
os.chdir(now_dir)
|
|
return "downloaded"
|
|
|
|
os.chdir(now_dir)
|
|
return None
|
|
|
|
|
|
def extract_and_show_progress(zipfile_path, unzips_path):
|
|
try:
|
|
with zipfile.ZipFile(zipfile_path, "r") as zip_ref:
|
|
for file_info in zip_ref.infolist():
|
|
zip_ref.extract(file_info, unzips_path)
|
|
os.remove(zipfile_path)
|
|
return True
|
|
except Exception as error:
|
|
print(error)
|
|
return False
|
|
|
|
|
|
def unzip_file(zip_path, zip_file_name):
|
|
zip_file_path = os.path.join(zip_path, zip_file_name + ".zip")
|
|
extract_path = os.path.join(file_path, zip_file_name)
|
|
with zipfile.ZipFile(zip_file_path, "r") as zip_ref:
|
|
zip_ref.extractall(extract_path)
|
|
os.remove(zip_file_path)
|
|
|
|
|
|
def model_download_pipeline(url):
|
|
verify = download_from_url(url)
|
|
if verify == "downloaded":
|
|
extract_folder_path = ""
|
|
for filename in os.listdir(zips_path):
|
|
if filename.endswith(".zip"):
|
|
zipfile_path = os.path.join(zips_path, filename)
|
|
print("Proceeding with the extraction...")
|
|
|
|
model_zip = os.path.basename(zipfile_path)
|
|
model_name = format_title(model_zip.split(".zip")[0])
|
|
extract_folder_path = os.path.join(
|
|
"logs",
|
|
os.path.normpath(model_name),
|
|
)
|
|
|
|
success = extract_and_show_progress(zipfile_path, extract_folder_path)
|
|
|
|
subfolders = [
|
|
f
|
|
for f in os.listdir(extract_folder_path)
|
|
if os.path.isdir(os.path.join(extract_folder_path, f))
|
|
]
|
|
if len(subfolders) == 1:
|
|
subfolder_path = os.path.join(extract_folder_path, subfolders[0])
|
|
for item in os.listdir(subfolder_path):
|
|
s = os.path.join(subfolder_path, item)
|
|
d = os.path.join(extract_folder_path, item)
|
|
shutil.move(s, d)
|
|
os.rmdir(subfolder_path)
|
|
|
|
for item in os.listdir(extract_folder_path):
|
|
if ".pth" in item:
|
|
file_name = item.split(".pth")[0]
|
|
if file_name != model_name:
|
|
os.rename(
|
|
os.path.join(extract_folder_path, item),
|
|
os.path.join(extract_folder_path, model_name + ".pth"),
|
|
)
|
|
else:
|
|
if "v2" not in item:
|
|
file_name = item.split("_nprobe_1_")[1].split("_v1")[0]
|
|
if file_name != model_name:
|
|
new_file_name = (
|
|
item.split("_nprobe_1_")[0]
|
|
+ "_nprobe_1_"
|
|
+ model_name
|
|
+ "_v1"
|
|
)
|
|
os.rename(
|
|
os.path.join(extract_folder_path, item),
|
|
os.path.join(
|
|
extract_folder_path, new_file_name + ".index"
|
|
),
|
|
)
|
|
else:
|
|
file_name = item.split("_nprobe_1_")[1].split("_v2")[0]
|
|
if file_name != model_name:
|
|
new_file_name = (
|
|
item.split("_nprobe_1_")[0]
|
|
+ "_nprobe_1_"
|
|
+ model_name
|
|
+ "_v2"
|
|
)
|
|
os.rename(
|
|
os.path.join(extract_folder_path, item),
|
|
os.path.join(
|
|
extract_folder_path, new_file_name + ".index"
|
|
),
|
|
)
|
|
|
|
if success:
|
|
print(f"Model {model_name} downloaded!")
|
|
else:
|
|
print(f"Error downloading {model_name}")
|
|
sys.exit()
|
|
if extract_folder_path == "":
|
|
print("Zip file was not found.")
|
|
sys.exit()
|
|
result = search_pth_index(extract_folder_path)
|
|
else:
|
|
message = "Error"
|