File size: 6,047 Bytes
05aac64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
import re
import os
import sys
import shutil

sys.path.append(os.getcwd())

from modules.utils import HF_download_file
from modules import gdown, meganz, mediafire, pixeldrain

def move_files_from_directory(src_dir, dest_models, model_name):
    for root, _, files in os.walk(src_dir):
        for file in files:
            file_path = os.path.join(root, file)
            if file.endswith(".index"):
                filepath = os.path.join(dest_models, file.replace(' ', '_').replace('(', '').replace(')', '').replace('[', '').replace(']', '').replace(",", "").replace('"', "").replace("'", "").replace("|", "").strip())

                shutil.move(file_path, filepath)
            elif file.endswith(".pth") and not file.startswith("D_") and not file.startswith("G_"):
                pth_path = os.path.join(dest_models, model_name + ".pth")

                shutil.move(file_path, pth_path)

def save_drop_model(dropbox):
    model_folders = "rvc_models" 
    save_model_temp = "save_model_temp"

    if not os.path.exists(model_folders): os.makedirs(model_folders, exist_ok=True)
    if not os.path.exists(save_model_temp): os.makedirs(save_model_temp, exist_ok=True)

    shutil.move(dropbox, save_model_temp)

    try:
        print("[INFO] Start uploading...")

        file_name = os.path.basename(dropbox)
        model_folders = os.path.join(model_folders, file_name.replace(".zip", "").replace(".pth", "").replace(".index", ""))

        if file_name.endswith(".zip"):
            shutil.unpack_archive(os.path.join(save_model_temp, file_name), save_model_temp)
            move_files_from_directory(save_model_temp, model_folders, file_name.replace(".zip", ""))
        elif file_name.endswith(".pth"): 
            output_file = os.path.join(model_folders, file_name)
            shutil.move(os.path.join(save_model_temp, file_name), output_file)
        elif file_name.endswith(".index"):
            def extract_name_model(filename):
                match = re.search(r"([A-Za-z]+)(?=_v|\.|$)", filename)
                return match.group(1) if match else None
            
            model_logs = os.path.join(model_folders, extract_name_model(file_name))
            if not os.path.exists(model_logs): os.makedirs(model_logs, exist_ok=True)
            shutil.move(os.path.join(save_model_temp, file_name), model_logs)
        else: 
            print("[WARNING] Format not supported. Supported formats ('.zip', '.pth', '.index')")
            return
        
        print("[INFO] Completed upload.")
    except Exception as e:
        print(f"[ERROR] An error occurred during unpack: {e}")
    finally:
        shutil.rmtree(save_model_temp, ignore_errors=True)

def download_model(url=None, model=None):
    if not url: 
        print("[WARNING] Please provide a valid url.")
        return

    if not model: 
        print("[WARNING] Please provide a valid model name.")
        return

    model = model.replace(".pth", "").replace(".index", "").replace(".zip", "").replace(" ", "_").replace("(", "").replace(")", "").replace("[", "").replace("]", "").replace(",", "").replace('"', "").replace("'", "").replace("|", "").strip()
    url = url.replace("/blob/", "/resolve/").replace("?download=true", "").strip()

    download_dir = "download_model"
    model_folders = "rvc_models" 

    if not os.path.exists(download_dir): os.makedirs(download_dir, exist_ok=True)
    if not os.path.exists(model_folders): os.makedirs(model_folders, exist_ok=True)

    model_folders = os.path.join(model_folders, model)
    os.makedirs(model_folders, exist_ok=True)
    
    try:
        print("[INFO] Start downloading...")

        if url.endswith(".pth"): HF_download_file(url, os.path.join(model_folders, f"{model}.pth"))
        elif url.endswith(".index"): HF_download_file(url, os.path.join(model_folders, f"{model}.index"))
        elif url.endswith(".zip"):
            output_path = HF_download_file(url, os.path.join(download_dir, model + ".zip"))
            shutil.unpack_archive(output_path, download_dir)

            move_files_from_directory(download_dir, model_folders, model)
        else:
            if "drive.google.com" in url or "drive.usercontent.google.com" in url:
                file_id = None

                if "/file/d/" in url: file_id = url.split("/d/")[1].split("/")[0]
                elif "open?id=" in url: file_id = url.split("open?id=")[1].split("/")[0]
                elif "/download?id=" in url: file_id = url.split("/download?id=")[1].split("&")[0]
                
                if file_id:
                    file = gdown.gdown_download(id=file_id, output=download_dir)
                    if file.endswith(".zip"): shutil.unpack_archive(file, download_dir)

                    move_files_from_directory(download_dir, model_folders, model)
            elif "mega.nz" in url:
                meganz.mega_download_url(url, download_dir)

                file_download = next((f for f in os.listdir(download_dir)), None)
                if file_download.endswith(".zip"): shutil.unpack_archive(os.path.join(download_dir, file_download), download_dir)

                move_files_from_directory(download_dir, model_folders, model)
            elif "mediafire.com" in url:
                file = mediafire.Mediafire_Download(url, download_dir)
                if file.endswith(".zip"): shutil.unpack_archive(file, download_dir)

                move_files_from_directory(download_dir, model_folders, model)
            elif "pixeldrain.com" in url:
                file = pixeldrain.pixeldrain(url, download_dir)
                if file.endswith(".zip"): shutil.unpack_archive(file, download_dir)

                move_files_from_directory(download_dir, model_folders, model)
            else:
                print("[WARNING] The url path is not supported.")
                return
        
        print("[INFO] Model download complete.")
    except Exception as e:
        print(f"[INFO] An error has occurred: {e}")
    finally:
        shutil.rmtree(download_dir, ignore_errors=True)