File size: 4,114 Bytes
20d05ae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91bc9e7
20d05ae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import shutil
from mega import Mega
import gdown
import re
import wget
import sys
import uuid
import zipfile


class InvalidDriveId(Exception):
    def __init__(self, message="Error de la url"):
        self.message = message
        super().__init__(self.message)


def model_downloader(url, zip_path, dest_path):
    """Download and unzip a file from Google Drive or Mega."""
    
    def drive_download(url, dest_folder):
        print(f"Descargando desde drive...")
        try:
            filename = gdown.download(url, os.path.join(dest_folder, f"{uuid.uuid4()}.zip"), fuzzy=True)
            return os.path.basename(filename)
        except:
            print("El intento de descargar con drive no funcionó")
            return None

    def mega_download(url, dest_folder):
        try:
            file_id = None
            if "#!" in url:
                file_id = url.split("#!")[1].split("!")[0]
            elif "file/" in url:
                file_id = url.split("file/")[1].split("/")[0]
            else:
                file_id = None

            print(f"Descargando desde mega...")
            if file_id:
                mega = Mega()
                m = mega.login()
                filename = m.download_url(url, dest_path=dest_folder, dest_filename=f"{uuid.uuid4()}.zip")

                return os.path.basename(filename)
            else:
                return None

        except Exception as e:
            print("Ocurrio un error**")
            print(e)
            return None

    def download(url, dest_folder):
        try:
            print(f"Descargando desde url generica...")
            dest_path = wget.download(url=url, out=os.path.join(dest_folder, f"{uuid.uuid4()}.zip"))

            return os.path.basename(dest_path)
        except Exception as e:
            print(f"Error al descargar el archivo: {str(e)}")

    filename = ""

    if not os.path.exists(zip_path):
        os.mkdir(zip_path)

    if url and 'drive.google.com' in url:
        # Descargar el elemento si la URL es de Google Drive
        filename = drive_download(url, zip_path)
    elif url and 'mega.nz' in url:
        filename = mega_download(url, zip_path)
    elif url and 'pixeldrain' in url:
        print("No se puede descargar de pixeldrain")
        sys.exit()
    else:
        filename = download(url, zip_path)

    if filename:
        print(f"Descomprimiendo {filename}...")
        modelname = str(filename).replace(".zip", "")
        zip_file_path = os.path.join(zip_path, filename)
        
        try:
            shutil.unpack_archive(zip_file_path, os.path.join(dest_path, modelname))
        except Exception as e:
            try:
                with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                    zip_ref.extractall(dest_path)
            except zipfile.BadZipFile as e:
                print(f"Error: El archivo ZIP no es válido - {e}")
            except Exception as e:
                print(f"Error inesperado: {e}")
        
        if os.path.exists(zip_file_path):
            os.remove(zip_file_path)

        return modelname
    else:
        return None


def get_models(weight_path):
    # Obtener todos los elementos en la ruta
    files = os.listdir(weight_path)
    # Filtrar solo los directorios
    return [file for file in files if os.path.isdir(os.path.join(weight_path, file))]


def get_model(weight_path, modelname):
    resources = {}
    for root, dirs, files in os.walk(os.path.join(weight_path, modelname)):
        for file in files:
            if file.endswith('.index'):
                resources['index'] =  os.path.relpath(os.path.join(root, file))
            if file.endswith('.pth'):
                resources['pth'] =  os.path.relpath(os.path.join(root, file), start=weight_path)
    return resources


def get_audios(audios_path):
    # Obtener todos los elementos en la ruta
    files = os.listdir(audios_path)
    # Filtrar solo los directorios
    return [file for file in files if not os.path.isdir(os.path.join(audios_path, file)) and os.path.join(audios_path, file).endswith(('.mp3', '.wav'))]