| | 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: |
| | |
| | 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_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') and not 'G_' in file and not 'D_' in file: |
| | resources['pth'] = os.path.relpath(os.path.join(root, file), start=weight_path) |
| | return resources |
