Update app.py

app.py

@@ -1,4 +1,5 @@
-import os
+import os, shutil
+import subprocess, glob
 import sys
 from dotenv import load_dotenv
 from applio import *
@@ -7,11 +8,13 @@ sys.path.append(now_dir)
 load_dotenv()
 load_dotenv("sha256.env")
 
+
+
+
 if sys.platform == "darwin":
     os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
 
 from infer.modules.vc import VC, show_info, hash_similarity
-from infer.modules.uvr5.modules import uvr
 from infer.lib.train.process_ckpt import (
     change_info,
     extract_small_model,
@@ -799,6 +802,84 @@ def change_f0_method(f0method8):
     return {"visible": visible, "__type__": "update"}
 
 
+def show_available(filepath,format=None):
+    if format:
+        print(f"Format: {format}")
+        files = []
+        for file in os.listdir(filepath):
+            if file.endswith(format):
+                print(f"Matches format: {file}")
+                files.append(file)
+            else:
+                print(f"Does not match format: {file}")
+        print(f"Matches: {files}")
+        if len(files) < 1:
+            return ['']
+        return files
+    if len(os.listdir(filepath)) < 1:
+        return ['']
+    return os.listdir(filepath)
+
+
+def download_from_url(url, model):
+    if model =='':
+        try:
+            model = url.split('/')[-1].split('?')[0]
+        except:
+            return "You need to name your model. For example: My-Model", {"choices":show_available("assets/weights"),"__type__":"update"}
+    url=url.replace('/blob/main/','/resolve/main/')
+    model=model.replace('.pth','').replace('.index','').replace('.zip','')
+    print(f"Model name: {model}")
+    if url == '':
+        return "URL cannot be left empty.", {"choices":show_available("assets/weights"),"__type__":"update"}
+    url = url.strip()
+    zip_dirs = ["zips", "unzips"]
+    for directory in zip_dirs:
+        if os.path.exists(directory):
+            shutil.rmtree(directory)
+    os.makedirs("zips", exist_ok=True)
+    os.makedirs("unzips", exist_ok=True)
+    zipfile = model + '.zip'
+    zipfile_path = './zips/' + zipfile
+    try:
+        if url.endswith('.pth'):
+            subprocess.run(["wget", url, "-O", f'./assets/weights/{model}.pth'])
+            return f"Sucessfully downloaded as {model}.pth", {"choices":show_available("assets/weights"),"__type__":"update"}
+        elif url.endswith('.index'):
+            if not os.path.exists(f'./logs/{model}'): os.makedirs(f'./logs/{model}')
+            subprocess.run(["wget", url, "-O", f'./logs/{model}/added_{model}.index'])
+            return f"Successfully downloaded as added_{model}.index", {"choices":show_available("assets/weights"),"__type__":"update"}
+        if "drive.google.com" in url:
+            subprocess.run(["gdown", url, "--fuzzy", "-O", zipfile_path])
+        elif "mega.nz" in url:
+            m = Mega()
+            m.download_url(url, './zips')
+        else:
+            subprocess.run(["wget", url, "-O", zipfile_path])
+        for filename in os.listdir("./zips"):
+            if filename.endswith(".zip"):
+                zipfile_path = os.path.join("./zips/",filename)
+                shutil.unpack_archive(zipfile_path, "./unzips", 'zip')
+                for root, dirs, files in os.walk('./unzips'):
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        if file.endswith(".index"):
+                            os.mkdir(f'./logs/{model}')
+                            shutil.copy2(file_path,f'./logs/{model}')
+                        elif "G_" not in file and "D_" not in file and file.endswith(".pth"):
+                            shutil.copy(file_path,f'./assets/weights/{model}.pth')
+            elif filename.endswith(".pth"):
+                shutil.copy2(os.path.join("./zips/",filename),f'./assets/weights/{model}.pth')
+            elif filename.endswith(".index"):
+                os.mkdir(f'./logs/{model}')
+                shutil.copy2(os.path.join("./zips/",filename),f'./logs/{model}/')
+            else:
+                return "No zipfile found.", {"choices":show_available("assets/weights"),"__type__":"update"}
+        shutil.rmtree("zips")
+        shutil.rmtree("unzips")
+        return "Success.", {"choices":show_available("assets/weights"),"__type__":"update"}
+    except:
+        return "There's been an error.", {"choices":show_available("assets/weights"),"__type__":"update"}