Create easyfuncs.py

easyfuncs.py

@@ -0,0 +1,112 @@
+import os, subprocess
+import gradio as gr
+import shutil
+from mega import Mega
+
+import pandas as pd
+import os
+
+# Class to handle caching model urls from a spreadsheet
+class CachedModels:
+    def __init__(self):
+        csv_url = "https://docs.google.com/spreadsheets/d/1tAUaQrEHYgRsm1Lvrnj14HFHDwJWl0Bd9x0QePewNco/export?format=csv&gid=1977693859"
+        if os.path.exists("spreadsheet.csv"):
+            self.cached_data = pd.read_csv("spreadsheet.csv") 
+        else:
+            self.cached_data = pd.read_csv(csv_url)
+            self.cached_data.to_csv("spreadsheet.csv", index=False)
+        # Cache model urls        
+        self.models = {}
+        for _, row in self.cached_data.iterrows():
+            filename = row['Filename']
+            url = None
+            for value in row.values:
+                if isinstance(value, str) and "huggingface" in value:
+                    url = value
+                    break
+            if url:
+                self.models[filename] = url
+    # Get cached model urls    
+    def get_models(self):
+        return self.models
+        
+def show(path,ext,on_error=None):
+    try:
+        return list(filter(lambda x: x.endswith(ext), os.listdir(path)))
+    except:
+        return on_error
+    
+def run_subprocess(command):
+    try:
+        subprocess.run(command, check=True)
+        return True, None
+    except Exception as e:
+        return False, e
+        
+def download_from_url(url=None, model=None):
+    if not url:
+        try:
+            url = model[f'{model}']
+        except:
+            gr.Warning("Failed")
+            return ''
+    if model == '':
+        try:
+            model = url.split('/')[-1].split('?')[0]
+        except:
+            gr.Warning('Please name the model')
+            return
+    model = model.replace('.pth', '').replace('.index', '').replace('.zip', '')
+    url = url.replace('/blob/main/', '/resolve/main/').strip()
+
+    for directory in ["downloads", "unzips","zip"]:
+        #shutil.rmtree(directory, ignore_errors=True)
+        os.makedirs(directory, exist_ok=True)
+
+    try:
+        if url.endswith('.pth'):
+            subprocess.run(["wget", url, "-O", f'assets/weights/{model}.pth'])
+        elif url.endswith('.index'):
+            os.makedirs(f'logs/{model}', exist_ok=True)
+            subprocess.run(["wget", url, "-O", f'logs/{model}/added_{model}.index'])
+        elif url.endswith('.zip'):
+            subprocess.run(["wget", url, "-O", f'downloads/{model}.zip'])
+        else:
+            if "drive.google.com" in url:
+                url = url.split('/')[0]
+                subprocess.run(["gdown", url, "--fuzzy", "-O", f'downloads/{model}'])
+            elif "mega.nz" in url:
+                Mega().download_url(url, 'downloads')
+            else:
+                subprocess.run(["wget", url, "-O", f'downloads/{model}'])
+
+        downloaded_file = next((f for f in os.listdir("downloads")), None)
+        if downloaded_file:
+            if downloaded_file.endswith(".zip"):
+                shutil.unpack_archive(f'downloads/{downloaded_file}', "unzips", 'zip')
+                for root, _, files in os.walk('unzips'):
+                    for file in files:
+                        file_path = os.path.join(root, file)
+                        if file.endswith(".index"):
+                            os.makedirs(f'logs/{model}', exist_ok=True)
+                            shutil.copy2(file_path, f'logs/{model}')
+                        elif file.endswith(".pth") and "G_" not in file and "D_" not in file:
+                            shutil.copy(file_path, f'assets/weights/{model}.pth')
+            elif downloaded_file.endswith(".pth"):
+                shutil.copy(f'downloads/{downloaded_file}', f'assets/weights/{model}.pth')
+            elif downloaded_file.endswith(".index"):
+                os.makedirs(f'logs/{model}', exist_ok=True)
+                shutil.copy(f'downloads/{downloaded_file}', f'logs/{model}/added_{model}.index')
+            else:
+                gr.Warning("Failed to download file")
+                return 'Failed'
+
+        gr.Info("Done")
+    except Exception as e:
+        gr.Warning(f"There's been an error: {str(e)}")
+    finally:
+        shutil.rmtree("downloads", ignore_errors=True)
+        shutil.rmtree("unzips", ignore_errors=True)
+        shutil.rmtree("zip", ignore_errors=True)
+        return 'Done'
+