refactor

app.py

@@ -1,11 +1,8 @@
-import spaces
 import os
 from stablepy import Model_Diffusers
-from stablepy.diffusers_vanilla.model import scheduler_names
 from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
 import torch
 import re
-import shutil
 import random
 from stablepy import (
     CONTROLNET_MODEL_IDS,
@@ -22,7 +19,6 @@ from stablepy import (
     SD15_TASKS,
     SDXL_TASKS,
 )
-import urllib.parse
 from config import (
     MINIMUM_IMAGE_NUMBER,
     MAXIMUM_IMAGE_NUMBER,
@@ -31,7 +27,15 @@ from config import (
 )
 from models.vae import VAE_LIST as download_vae
 from models.checkpoints import CHECKPOINT_LIST as download_model
+from models.loras import LORA_LIST as download_lora
+from models.format_models import FORMAT_MODELS as load_diffusers_format_model
+from models.upscaler import upscaler_dict_gui
 from examples.examples import example_prompts
+from utils.download_utils import download_things
+from utils.model_utils import get_model_list
+
+
+# from utils.string_utils import process_string
 
 preprocessor_controlnet = {
     "openpose": [
@@ -112,73 +116,6 @@ task_stablepy = {
 
 task_model_list = list(task_stablepy.keys())
 
-
-def download_things(directory, url, hf_token="", civitai_api_key=""):
-    url = url.strip()
-
-    if "drive.google.com" in url:
-        original_dir = os.getcwd()
-        os.chdir(directory)
-        os.system(f"gdown --fuzzy {url}")
-        os.chdir(original_dir)
-    elif "huggingface.co" in url:
-        url = url.replace("?download=true", "")
-        # url = urllib.parse.quote(url, safe=':/')  # fix encoding
-        if "/blob/" in url:
-            url = url.replace("/blob/", "/resolve/")
-        user_header = f'"Authorization: Bearer {hf_token}"'
-        if hf_token:
-            os.system(
-                f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
-        else:
-            os.system(
-                f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k "
-                f"1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
-    elif "civitai.com" in url:
-        if "?" in url:
-            url = url.split("?")[0]
-        if civitai_api_key:
-            url = url + f"?token={civitai_api_key}"
-            os.system(
-                f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
-        else:
-            print("\033[91mYou need an API key to download Civitai models.\033[0m")
-    else:
-        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
-
-
-def get_model_list(directory_path):
-    model_list: list = []
-    valid_extensions = {
-        '.ckpt',
-        '.pt',
-        '.pth',
-        '.safetensors',
-        '.bin'
-    }
-
-    for filename in os.listdir(directory_path):
-        if os.path.splitext(filename)[1] in valid_extensions:
-            name_without_extension = os.path.splitext(filename)[0]
-            file_path = os.path.join(directory_path, filename)
-            # model_list.append((name_without_extension, file_path))
-            model_list.append(file_path)
-            print('\033[34mFILE: ' + file_path + '\033[0m')
-    return model_list
-
-
-def process_string(input_string):
-    parts = input_string.split('/')
-
-    if len(parts) == 2:
-        first_element = parts[1]
-        complete_string = input_string
-        result = (first_element, complete_string)
-        return result
-    else:
-        return None
-
-
 directory_models = 'models'
 os.makedirs(directory_models, exist_ok=True)
 directory_loras = 'loras'
@@ -186,74 +123,42 @@ os.makedirs(directory_loras, exist_ok=True)
 directory_vaes = 'vaes'
 os.makedirs(directory_vaes, exist_ok=True)
 
-# - **Download LoRAs**
-download_lora = (
-    "https://civitai.com/api/download/models/423719, "
-    "https://civitai.com/api/download/models/50503, "
-    "https://civitai.com/api/download/models/133160, "
-    "https://civitai.com/api/download/models/29332, "
-    "https://huggingface.co/Leopain/color/resolve/main/Coloring_book_-_LineArt.safetensors, "
-    "https://civitai.com/api/download/models/135867, "
-    "https://civitai.com/api/download/models/145907, "
-    "https://huggingface.co/Linaqruf/anime-detailer-xl-lora/resolve/main/anime-detailer-xl.safetensors?download=true, "
-    "https://huggingface.co/Linaqruf/style-enhancer-xl-lora/resolve/main/style-enhancer-xl.safetensors?download=true, "
-    "https://civitai.com/api/download/models/28609, "
-    "https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SD15-8steps-CFG-lora.safetensors?download=true, "
-    "https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SDXL-8steps-CFG-lora.safetensors?download=true, "
-    "https://civitai.com/api/download/models/30666 "
-)
-
-load_diffusers_format_model = [
-    'stabilityai/stable-diffusion-xl-base-1.0',
-    'cagliostrolab/animagine-xl-3.1',
-    'misri/epicrealismXL_v7FinalDestination',
-    'misri/juggernautXL_juggernautX',
-    'misri/zavychromaxl_v80',
-    'SG161222/RealVisXL_V4.0',
-    'misri/newrealityxlAllInOne_Newreality40',
-    'eienmojiki/Anything-XL',
-    'eienmojiki/Starry-XL-v5.2',
-    'gsdf/CounterfeitXL',
-    'kitty7779/ponyDiffusionV6XL',
-    'John6666/ebara-mfcg-pony-mix-v12-sdxl',
-    'John6666/t-ponynai3-v51-sdxl',
-    'yodayo-ai/kivotos-xl-2.0',
-    'yodayo-ai/holodayo-xl-2.1',
-    'digiplay/majicMIX_sombre_v2',
-    'digiplay/majicMIX_realistic_v6',
-    'digiplay/majicMIX_realistic_v7',
-    'digiplay/DreamShaper_8',
-    'digiplay/BeautifulArt_v1',
-    'digiplay/DarkSushi2.5D_v1',
-    'digiplay/darkphoenix3D_v1.1',
-    'digiplay/BeenYouLiteL11_diffusers',
-    'rubbrband/revAnimated_v2Rebirth',
-    'youknownothing/cyberrealistic_v50',
-    'votepurchase/counterfeitV30_v30',
-    'Meina/MeinaMix_V11',
-    'Meina/MeinaUnreal_V5',
-    'Meina/MeinaPastel_V7',
-    'rubbrband/realcartoon3d_v16',
-    'rubbrband/realcartoonRealistic_v14',
-]
-
+# LOAD ALL ENV TOKEN
 CIVITAI_API_KEY: str = os.environ.get("CIVITAI_API_KEY")
 hf_token: str = os.environ.get("HF_TOKEN")
 
 # Download stuffs
 for url in [url.strip() for url in download_model.split(',')]:
     if not os.path.exists(f"./models/{url.split('/')[-1]}"):
-        download_things(directory_models, url, hf_token, CIVITAI_API_KEY)
+        download_things(
+            directory_models,
+            url,
+            hf_token,
+            CIVITAI_API_KEY
+        )
 for url in [url.strip() for url in download_vae.split(',')]:
     if not os.path.exists(f"./vaes/{url.split('/')[-1]}"):
-        download_things(directory_vaes, url, hf_token, CIVITAI_API_KEY)
+        download_things(
+            directory_vaes,
+            url,
+            hf_token,
+            CIVITAI_API_KEY
+        )
 for url in [url.strip() for url in download_lora.split(',')]:
     if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
-        download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+        download_things(
+            directory_loras,
+            url,
+            hf_token,
+            CIVITAI_API_KEY
+        )
 
 # Download Embeddings
 directory_embeds = 'embedings'
-os.makedirs(directory_embeds, exist_ok=True)
+os.makedirs(
+    directory_embeds,
+    exist_ok=True
+)
 download_embeds = [
     'https://huggingface.co/datasets/Nerfgun3/bad_prompt/blob/main/bad_prompt_version2.pt',
     'https://huggingface.co/embed/negative/resolve/main/EasyNegativeV2.safetensors',
@@ -262,7 +167,12 @@ download_embeds = [
 
 for url_embed in download_embeds:
     if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
-        download_things(directory_embeds, url_embed, hf_token, CIVITAI_API_KEY)
+        download_things(
+            directory_embeds,
+            url_embed,
+            hf_token,
+            CIVITAI_API_KEY
+        )
 
 # Build list models
 embed_list = get_model_list(directory_embeds)
@@ -277,7 +187,12 @@ vae_model_list.insert(0, "None")
 def get_my_lora(link_url):
     for url in [url.strip() for url in link_url.split(',')]:
         if not os.path.exists(f"./loras/{url.split('/')[-1]}"):
-            download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+            download_things(
+                directory_loras,
+                url,
+                hf_token,
+                CIVITAI_API_KEY
+            )
     new_lora_model_list = get_model_list(directory_loras)
     new_lora_model_list.insert(0, "None")
 
@@ -291,32 +206,11 @@ def get_my_lora(link_url):
         choices=new_lora_model_list
     ), gr.update(
         choices=new_lora_model_list
-    ),
+    )
 
 
 print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
 
-upscaler_dict_gui = {
-    None: None,
-    "Lanczos": "Lanczos",
-    "Nearest": "Nearest",
-    "RealESRGAN_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
-    "RealESRNet_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
-    "RealESRGAN_x4plus_anime_6B": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
-    "RealESRGAN_x2plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
-    "realesr-animevideov3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
-    "realesr-general-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
-    "realesr-general-wdn-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
-    "4x-UltraSharp": "https://huggingface.co/Shandypur/ESRGAN-4x-UltraSharp/resolve/main/4x-UltraSharp.pth",
-    "4x_foolhardy_Remacri": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
-    "Remacri4xExtraSmoother": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/Remacri%204x%20ExtraSmoother.pth",
-    "AnimeSharp4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/AnimeSharp%204x.pth",
-    "lollypop": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/lollypop.pth",
-    "RealisticRescaler4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/RealisticRescaler%204x.pth",
-    "NickelbackFS4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/NickelbackFS%204x.pth"
-}
-
-
 def extract_parameters(input_string):
     parameters = {}
     input_string = input_string.replace("\n", "")
@@ -764,9 +658,9 @@ class GuiSD:
         # Maybe fix lora issue: 'Cannot copy out of meta tensor; no data!''
         self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
 
-        info_state = f"PROCESSING "
+        info_state = f"PROCESSING"
         for img, seed, data in self.model(**pipe_params):
-            info_state += ">"
+            info_state += "."
             if data:
                 info_state = f"COMPLETED. Seeds: {str(seed)}"
                 if vae_msg:
@@ -1018,6 +912,7 @@ with gr.Blocks(css=CSS) as app:
                 vae_model_gui = gr.Dropdown(
                     label="VAE Model",
                     choices=vae_model_list,
+                    value=vae_model_list[0]
                 )
 
                 with gr.Accordion("Hires fix", open=False, visible=True):
@@ -1642,7 +1537,7 @@ with gr.Blocks(css=CSS) as app:
             )
 
 
-            def send_img(img_source, img_result):
+            def send_img(img_source, img_result) -> tuple:
                 return img_source, img_result
 
 
@@ -1767,7 +1662,10 @@ with gr.Blocks(css=CSS) as app:
             mode_ip2,
             scale_ip2,
         ],
-        outputs=[result_images, actual_task_info],
+        outputs=[
+            result_images,
+            actual_task_info
+        ],
         queue=True,
         show_progress="minimal",
     )

models/format_models.py

@@ -0,0 +1,33 @@
+FORMAT_MODELS: list = [
+    'stabilityai/stable-diffusion-xl-base-1.0',
+    'cagliostrolab/animagine-xl-3.1',
+    'misri/epicrealismXL_v7FinalDestination',
+    'misri/juggernautXL_juggernautX',
+    'misri/zavychromaxl_v80',
+    'SG161222/RealVisXL_V4.0',
+    'misri/newrealityxlAllInOne_Newreality40',
+    'eienmojiki/Anything-XL',
+    'eienmojiki/Starry-XL-v5.2',
+    'gsdf/CounterfeitXL',
+    'kitty7779/ponyDiffusionV6XL',
+    'John6666/ebara-mfcg-pony-mix-v12-sdxl',
+    'John6666/t-ponynai3-v51-sdxl',
+    'yodayo-ai/kivotos-xl-2.0',
+    'yodayo-ai/holodayo-xl-2.1',
+    'digiplay/majicMIX_sombre_v2',
+    'digiplay/majicMIX_realistic_v6',
+    'digiplay/majicMIX_realistic_v7',
+    'digiplay/DreamShaper_8',
+    'digiplay/BeautifulArt_v1',
+    'digiplay/DarkSushi2.5D_v1',
+    'digiplay/darkphoenix3D_v1.1',
+    'digiplay/BeenYouLiteL11_diffusers',
+    'rubbrband/revAnimated_v2Rebirth',
+    'youknownothing/cyberrealistic_v50',
+    'votepurchase/counterfeitV30_v30',
+    'Meina/MeinaMix_V11',
+    'Meina/MeinaUnreal_V5',
+    'Meina/MeinaPastel_V7',
+    'rubbrband/realcartoon3d_v16',
+    'rubbrband/realcartoonRealistic_v14',
+]

models/loras.py

@@ -0,0 +1,15 @@
+LORA_LIST = (
+    "https://civitai.com/api/download/models/423719, "
+    "https://civitai.com/api/download/models/50503, "
+    "https://civitai.com/api/download/models/133160, "
+    "https://civitai.com/api/download/models/29332, "
+    "https://huggingface.co/Leopain/color/resolve/main/Coloring_book_-_LineArt.safetensors, "
+    "https://civitai.com/api/download/models/135867, "
+    "https://civitai.com/api/download/models/145907, "
+    "https://huggingface.co/Linaqruf/anime-detailer-xl-lora/resolve/main/anime-detailer-xl.safetensors?download=true, "
+    "https://huggingface.co/Linaqruf/style-enhancer-xl-lora/resolve/main/style-enhancer-xl.safetensors?download=true, "
+    "https://civitai.com/api/download/models/28609, "
+    "https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SD15-8steps-CFG-lora.safetensors?download=true, "
+    "https://huggingface.co/ByteDance/Hyper-SD/resolve/main/Hyper-SDXL-8steps-CFG-lora.safetensors?download=true, "
+    "https://civitai.com/api/download/models/30666 "
+)

models/upscaler.py

@@ -0,0 +1,19 @@
+upscaler_dict_gui: dict = {
+    None: None,
+    "Lanczos": "Lanczos",
+    "Nearest": "Nearest",
+    "RealESRGAN_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+    "RealESRNet_x4plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth",
+    "RealESRGAN_x4plus_anime_6B": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth",
+    "RealESRGAN_x2plus": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth",
+    "realesr-animevideov3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth",
+    "realesr-general-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth",
+    "realesr-general-wdn-x4v3": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth",
+    "4x-UltraSharp": "https://huggingface.co/Shandypur/ESRGAN-4x-UltraSharp/resolve/main/4x-UltraSharp.pth",
+    "4x_foolhardy_Remacri": "https://huggingface.co/FacehugmanIII/4x_foolhardy_Remacri/resolve/main/4x_foolhardy_Remacri.pth",
+    "Remacri4xExtraSmoother": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/Remacri%204x%20ExtraSmoother.pth",
+    "AnimeSharp4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/AnimeSharp%204x.pth",
+    "lollypop": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/lollypop.pth",
+    "RealisticRescaler4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/RealisticRescaler%204x.pth",
+    "NickelbackFS4x": "https://huggingface.co/hollowstrawberry/upscalers-backup/resolve/main/ESRGAN/NickelbackFS%204x.pth"
+}

utils/download_utils.py

@@ -0,0 +1,43 @@
+def download_things(
+        directory: str,
+        url: str,
+        hf_token: str = "",
+        civitai_api_key: str = ""):
+    """
+    :param directory: String, directory to save the downloaded file
+    :param url: String, URL to download the file from
+    :param hf_token: String, Hugging Face API token
+    :param civitai_api_key: String, Civitai API key
+    """
+    import os
+    url = url.strip()
+
+    if "drive.google.com" in url:
+        original_dir = os.getcwd()
+        os.chdir(directory)
+        os.system(f"gdown --fuzzy {url}")
+        os.chdir(original_dir)
+    elif "huggingface.co" in url:
+        url = url.replace("?download=true", "")
+        # url = urllib.parse.quote(url, safe=':/')  # fix encoding
+        if "/blob/" in url:
+            url = url.replace("/blob/", "/resolve/")
+        user_header = f'"Authorization: Bearer {hf_token}"'
+        if hf_token:
+            os.system(
+                f"aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+        else:
+            os.system(
+                f"aria2c --optimize-concurrent-downloads --console-log-level=error --summary-interval=10 -c -x 16 -k "
+                f"1M -s 16 {url} -d {directory}  -o {url.split('/')[-1]}")
+    elif "civitai.com" in url:
+        if "?" in url:
+            url = url.split("?")[0]
+        if civitai_api_key:
+            url = url + f"?token={civitai_api_key}"
+            os.system(
+                f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")
+        else:
+            print("\033[91mYou need an API key to download Civitai models.\033[0m")
+    else:
+        os.system(f"aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {directory} {url}")

utils/model_utils.py

@@ -0,0 +1,19 @@
+def get_model_list(directory_path):
+    import os
+    model_list: list = []
+    valid_extensions = {
+        '.ckpt',
+        '.pt',
+        '.pth',
+        '.safetensors',
+        '.bin'
+    }
+
+    for filename in os.listdir(directory_path):
+        if os.path.splitext(filename)[1] in valid_extensions:
+            name_without_extension = os.path.splitext(filename)[0]
+            file_path = os.path.join(directory_path, filename)
+            # model_list.append((name_without_extension, file_path))
+            model_list.append(file_path)
+            print('\033[34mFILE: ' + file_path + '\033[0m')
+    return model_list

utils/string_utils.py

@@ -0,0 +1,14 @@
+def process_string(input_string: str):
+    """
+    :param input_string:
+    :return:
+    """
+    parts = input_string.split('/')
+
+    if len(parts) == 2:
+        first_element = parts[1]
+        complete_string = input_string
+        result = (first_element, complete_string)
+        return result
+    else:
+        return None