lib/tools/server.py

from flask import Flask, request, Response
import os
import sys
import requests
import shutil
import subprocess
import wget
import signal
from bs4 import BeautifulSoup
import logging
import click


app = Flask(__name__)

# Disable flask starting message
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

def secho(text, file=None, nl=None, err=None, color=None, **styles):
    pass

def echo(text, file=None, nl=None, err=None, color=None, **styles):
    pass

click.echo = echo
click.secho = secho

# Get the current directory path
now_dir = os.path.dirname(os.path.abspath(__file__))

# Go up two levels in the directory hierarchy
for _ in range(2):
    now_dir = os.path.dirname(now_dir)

# Add now_dir to sys.path so Python can find modules in that location
sys.path.append(now_dir)

from assets.i18n.i18n import I18nAuto
i18n = I18nAuto()

# Use the code from the resources module but with some changes
def find_folder_parent(search_dir, folder_name):
    for dirpath, dirnames, filenames in os.walk(search_dir):
        if folder_name in dirnames:
            return os.path.abspath(dirpath)
    return None

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):
    file_path = find_folder_parent(now_dir, "assets")
    print(file_path)
    zips_path = os.path.join(file_path, "assets", "zips")
    print(zips_path)
    os.makedirs(zips_path, exist_ok=True)
    if url != "":
        print(i18n("Downloading the file: ") + f"{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)
                result = subprocess.run(
                    ["gdown", f"https://drive.google.com/uc?id={file_id}", "--fuzzy"],
                    capture_output=True,
                    text=True,
                    encoding="utf-8",
                )
                if (
                    "Too many users have viewed or downloaded this file recently"
                    in str(result.stderr)
                ):
                    return "too much use"
                if "Cannot retrieve the public link of the file." in str(result.stderr):
                    return "private link"
                print(result.stderr)

        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:
                file_name = url.split("/")[-1]
                file_name = file_name.replace("%20", "_")
                total_size_in_bytes = int(response.headers.get('content-length', 0))
                block_size = 1024  # 1 Kibibyte
                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(file_path)
                return None
        elif "mega.nz" in url:
            if "#!" in url:
                file_id = url.split("#!")[1].split("!")[0]
            elif "file/" in url:
                file_id = url.split("file/")[1].split("/")[0]
            else:
                return None
            if file_id:
                print("Mega.nz is unsupported due mega.py deprecation")
        elif "/tree/main" in url:
            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:
                print("No .zip file found on the page.")
        elif "cdn.discordapp.com" in url:
            file = requests.get(url)
            os.chdir("./assets/zips")
            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 "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 "mediafire.com" in url:
            download_link = get_mediafire_download_link(url)
            if download_link:
                os.chdir(zips_path)
                wget.download(download_link)
            else:
                return None
        elif "www.weights.gg" in url:
            #Pls weights creator dont fix this because yes. c:
            url_parts = url.split("/")
            weights_gg_index = url_parts.index("www.weights.gg")
            if weights_gg_index != -1 and weights_gg_index < len(url_parts) - 1:
                model_part = "/".join(url_parts[weights_gg_index + 1:])
                if "models" in model_part:
                    model_part = model_part.split("models/")[-1]
                    print(model_part)
                    if model_part:
                        download_url = f"https://www.weights.gg/es/models/{model_part}"
                        response = requests.get(download_url)
                        if response.status_code == 200:
                            soup = BeautifulSoup(response.text, "html.parser")
                            button_link = soup.find("a", class_="bg-black text-white px-3 py-2 rounded-lg flex items-center gap-1")
                            if button_link:
                                download_link = button_link["href"]
                                result = download_from_url(download_link)
                                if result == "downloaded":
                                    return "downloaded"
                                else:
                                    return None
                            else:
                                return None
                        else:
                            return None
                    else:
                        return None
                else:
                    return None
            else:
                return None
        else:
            os.chdir(zips_path)
            wget.download(url)

        # Fix points in the zips
        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(file_path)
        print(i18n("Full download"))
        return "downloaded"
    else:
        return None

def extract_and_show_progress(zipfile_path, unzips_path):
    try:
        # Use shutil because zipfile not working
        shutil.unpack_archive(zipfile_path, unzips_path)
        return True
    except Exception as e:
        print(f"Error al descomprimir {zipfile_path}: {e}")
        return False


@app.route('/download/<path:url>', methods=['GET'])

def load_downloaded_model(url):
    parent_path = find_folder_parent(now_dir, "assets")
    response = requests.get(url)
    response.raise_for_status()
    try:
        zips_path = os.path.join(parent_path, "assets", "zips")
        unzips_path = os.path.join(parent_path, "assets", "unzips")
        weights_path = os.path.join(parent_path, "logs", "weights")
        logs_dir = ""

        if os.path.exists(zips_path):
            shutil.rmtree(zips_path)
        if os.path.exists(unzips_path):
            shutil.rmtree(unzips_path)

        os.mkdir(zips_path)
        os.mkdir(unzips_path)

        download_file = download_from_url(url)
        if not download_file:
            print(i18n("The file could not be downloaded."))
        elif download_file == "downloaded":
            print(i18n("It has been downloaded successfully."))
        elif download_file == "too much use":
            raise Exception(
                i18n("Too many users have recently viewed or downloaded this file")
            )
        elif download_file == "private link":
            raise Exception(i18n("Cannot get file from this private link"))

        for filename in os.listdir(zips_path):
            if filename.endswith(".zip"):
                zipfile_path = os.path.join(zips_path, filename)
                print(i18n("Proceeding with the extraction..."))
                model_name = os.path.basename(zipfile_path)
                logs_dir = os.path.join(
                    parent_path,
                    "logs",
                    os.path.normpath(str(model_name).replace(".zip", "")),
                )
                
                success = extract_and_show_progress(zipfile_path, unzips_path)
                if success:
                    print(f"Extracción exitosa: {model_name}")
                else:
                    print(f"Fallo en la extracción: {model_name}")
            else:
                print(i18n("Unzip error."))
                return ""

        index_file = False
        model_file = False

        for path, subdirs, files in os.walk(unzips_path):
            for item in files:
                item_path = os.path.join(path, item)
                if not "G_" in item and not "D_" in item and item.endswith(".pth"):
                    model_file = True
                    model_name = item.replace(".pth", "")
                    logs_dir = os.path.join(parent_path, "logs", model_name)
                    if os.path.exists(logs_dir):
                        shutil.rmtree(logs_dir)
                    os.mkdir(logs_dir)
                    if not os.path.exists(weights_path):
                        os.mkdir(weights_path)
                    if os.path.exists(os.path.join(weights_path, item)):
                        os.remove(os.path.join(weights_path, item))
                    if os.path.exists(item_path):
                        shutil.move(item_path, weights_path)

        if not model_file and not os.path.exists(logs_dir):
            os.mkdir(logs_dir)
        for path, subdirs, files in os.walk(unzips_path):
            for item in files:
                item_path = os.path.join(path, item)
                if item.startswith("added_") and item.endswith(".index"):
                    index_file = True
                    if os.path.exists(item_path):
                        if os.path.exists(os.path.join(logs_dir, item)):
                            os.remove(os.path.join(logs_dir, item))
                        shutil.move(item_path, logs_dir)
                if item.startswith("total_fea.npy") or item.startswith("events."):
                    if os.path.exists(item_path):
                        if os.path.exists(os.path.join(logs_dir, item)):
                            os.remove(os.path.join(logs_dir, item))
                        shutil.move(item_path, logs_dir)

        result = ""
        if model_file:
            if index_file:
                print(i18n("The model works for inference, and has the .index file."))
            else:
                print(
                    i18n(
                        "The model works for inference, but it doesn't have the .index file."
                    )
                )

        if not index_file and not model_file:
            print(i18n("No relevant file was found to upload."))
        
        os.chdir(parent_path)
        return result
    except Exception as e:
        os.chdir(parent_path)
        if "too much use" in str(e):
            print(i18n("Too many users have recently viewed or downloaded this file"))
        elif "private link" in str(e):
            print(i18n("Cannot get file from this private link"))
        else:
            print(i18n("An error occurred downloading"))
            print(e)
    finally:
        os.chdir(parent_path)

@app.route('/shutdown', methods=['POST'])
def shoutdown():
    print("This flask server is shutting down please close the window")   
    pid = os.getpid()
    os.kill(pid, signal.SIGTERM)

if __name__ == '__main__':
    app.run(host='localhost', port=8000)