refactor

app.py

@@ -1,4 +1,4 @@
-import spaces
+import spaces  # must imported when using spaces
 import os
 from stablepy import Model_Diffusers
 from stablepy.diffusers_vanilla.style_prompt_config import STYLE_NAMES
@@ -31,63 +31,13 @@ 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 models.controlnet import preprocessor_controlnet
+from models.embeds import download_embeds
 from examples.examples import example_prompts
 from utils.download_utils import download_things
 from utils.model_utils import get_model_list
 
-preprocessor_controlnet = {
-    "openpose": [
-        "Openpose",
-        "None",
-    ],
-    "scribble": [
-        "HED",
-        "Pidinet",
-        "None",
-    ],
-    "softedge": [
-        "Pidinet",
-        "HED",
-        "HED safe",
-        "Pidinet safe",
-        "None",
-    ],
-    "segmentation": [
-        "UPerNet",
-        "None",
-    ],
-    "depth": [
-        "DPT",
-        "Midas",
-        "None",
-    ],
-    "normalbae": [
-        "NormalBae",
-        "None",
-    ],
-    "lineart": [
-        "Lineart",
-        "Lineart coarse",
-        "Lineart (anime)",
-        "None",
-        "None (anime)",
-    ],
-    "shuffle": [
-        "ContentShuffle",
-        "None",
-    ],
-    "canny": [
-        "Canny"
-    ],
-    "mlsd": [
-        "MLSD"
-    ],
-    "ip2p": [
-        "ip2p"
-    ]
-}
-
-task_stablepy = {
+task_stablepy: dict = {
     'txt2img': 'txt2img',
     'img2img': 'img2img',
     'inpaint': 'inpaint',
@@ -111,20 +61,34 @@ task_stablepy = {
     'optical pattern ControlNet': 'pattern',
     'tile realistic': 'sdxl_tile_realistic',
 }
-
-task_model_list = list(task_stablepy.keys())
-
-directory_models = 'models'
-os.makedirs(directory_models, exist_ok=True)
-directory_loras = 'loras'
-os.makedirs(directory_loras, exist_ok=True)
-directory_vaes = 'vaes'
-os.makedirs(directory_vaes, exist_ok=True)
-
 # LOAD ALL ENV TOKEN
 CIVITAI_API_KEY: str = os.environ.get("CIVITAI_API_KEY")
 hf_token: str = os.environ.get("HF_TOKEN")
 
+
+task_model_list = list(task_stablepy.keys())
+
+directory_models: str = 'models'
+os.makedirs(
+    directory_models,
+    exist_ok=True
+)
+directory_loras: str = 'loras'
+os.makedirs(
+    directory_loras,
+    exist_ok=True
+)
+directory_vaes: str = 'vaes'
+os.makedirs(
+    directory_vaes,
+    exist_ok=True
+)
+directory_embeds: str = 'embedings'
+os.makedirs(
+    directory_embeds,
+    exist_ok=True
+)
+
 # Download stuffs
 for url in [url.strip() for url in download_model.split(',')]:
     if not os.path.exists(f"./models/{url.split('/')[-1]}"):
@@ -153,18 +117,6 @@ for url in [url.strip() for url in download_lora.split(',')]:
             CIVITAI_API_KEY
         )
 
-# Download Embeddings
-directory_embeds = 'embedings'
-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',
-    'https://huggingface.co/embed/negative/resolve/main/bad-hands-5.pt',
-]
-
 for url_embed in download_embeds:
     if not os.path.exists(f"./embedings/{url_embed.split('/')[-1]}"):
         download_things(
@@ -185,11 +137,11 @@ 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]}"):
+    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,
+                __url,
                 hf_token,
                 CIVITAI_API_KEY
             )
@@ -212,47 +164,6 @@ def get_my_lora(link_url):
 
 print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
 
-
-def extract_parameters(input_string):
-    parameters = {}
-    input_string = input_string.replace("\n", "")
-
-    if not "Negative prompt:" in input_string:
-        print("Negative prompt not detected")
-        parameters["prompt"] = input_string
-        return parameters
-
-    parm = input_string.split("Negative prompt:")
-    parameters["prompt"] = parm[0]
-    if not "Steps:" in parm[1]:
-        print("Steps not detected")
-        parameters["neg_prompt"] = parm[1]
-        return parameters
-    parm = parm[1].split("Steps:")
-    parameters["neg_prompt"] = parm[0]
-    input_string = "Steps:" + parm[1]
-
-    # Extracting Steps
-    steps_match = re.search(r'Steps: (\d+)', input_string)
-    if steps_match:
-        parameters['Steps'] = int(steps_match.group(1))
-
-    # Extracting Size
-    size_match = re.search(r'Size: (\d+x\d+)', input_string)
-    if size_match:
-        parameters['Size'] = size_match.group(1)
-        width, height = map(int, parameters['Size'].split('x'))
-        parameters['width'] = width
-        parameters['height'] = height
-
-    # Extracting other parameters
-    other_parameters = re.findall(r'(\w+): (.*?)(?=, \w+|$)', input_string)
-    for param in other_parameters:
-        parameters[param[0]] = param[1].strip('"')
-
-    return parameters
-
-
 #######################
 # GUI
 #######################
@@ -263,6 +174,8 @@ import IPython.display
 import time, json
 from IPython.utils import capture
 import logging
+from utils.string_utils import extract_parameters
+from stablepy import logger
 
 logging.getLogger("diffusers").setLevel(logging.ERROR)
 import diffusers
@@ -285,7 +198,6 @@ warnings.filterwarnings(
     category=FutureWarning,
     module="transformers"
 )
-from stablepy import logger
 
 logger.setLevel(logging.DEBUG)
 
@@ -448,7 +360,7 @@ class GuiSD:
         vae_model = vae_model if vae_model != "None" else None
         loras_list = [lora1, lora2, lora3, lora4, lora5]
         vae_msg = f"VAE: {vae_model}" if vae_model else ""
-        msg_lora = []
+        msg_lora: list = []
 
         if model_name in model_list:
             model_is_xl = "xl" in model_name.lower()
@@ -471,13 +383,15 @@ class GuiSD:
                 vae_model = None
 
             for la in loras_list:
-                if la is not None and la != "None" and la in lora_model_list:
-                    print(la)
-                    lora_type = ("animetarot" in la.lower() or "Hyper-SD15-8steps".lower() in la.lower())
-                    if (model_is_xl and lora_type) or (not model_is_xl and not lora_type):
-                        msg_inc_lora = f"The LoRA {la} is for {'SD 1.5' if model_is_xl else 'SDXL'}, but you are using {model_type}."
-                        gr.Info(msg_inc_lora)
-                        msg_lora.append(msg_inc_lora)
+                if la is None or la == "None" or la not in lora_model_list:
+                    continue
+
+                print(la)
+                lora_type = ("animetarot" in la.lower() or "Hyper-SD15-8steps".lower() in la.lower())
+                if (model_is_xl and lora_type) or (not model_is_xl and not lora_type):
+                    msg_inc_lora = f"The LoRA {la} is for {'SD 1.5' if model_is_xl else 'SDXL'}, but you are using {model_type}."
+                    gr.Info(msg_inc_lora)
+                    msg_lora.append(msg_inc_lora)
 
         task = task_stablepy[task]
 
@@ -718,6 +632,7 @@ def update_task_options(model_name, task_name):
         )
 
 
+# APP
 with gr.Blocks(css=CSS) as app:
     gr.Markdown("# 🧩 (Ivan) DiffuseCraft")
     with gr.Tab("Generation"):
@@ -846,7 +761,7 @@ with gr.Blocks(css=CSS) as app:
                         }
                         valid_keys = list(valid_receptors.keys())
 
-                        parameters = extract_parameters(base_prompt)
+                        parameters: dict = extract_parameters(base_prompt)
                         for key, val in parameters.items():
                             # print(val)
                             if key in valid_keys:
@@ -930,7 +845,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]
+                    value=vae_model_list[1]
                 )
 
                 with gr.Accordion("Hires fix", open=False, visible=True):
@@ -1061,11 +976,16 @@ with gr.Blocks(css=CSS) as app:
                         button_lora.click(
                             get_my_lora,
                             [text_lora],
-                            [lora1_gui, lora2_gui, lora3_gui, lora4_gui, lora5_gui]
+                            [
+                                lora1_gui,
+                                lora2_gui,
+                                lora3_gui,
+                                lora4_gui,
+                                lora5_gui
+                            ]
                         )
 
-                with gr.Accordion("IP-Adapter", open=False, visible=True):  ##############
-
+                with gr.Accordion("IP-Adapter", open=False, visible=True):  # IP-Adapter
                     IP_MODELS = sorted(list(set(IP_ADAPTERS_SD + IP_ADAPTERS_SDXL)))
                     MODE_IP_OPTIONS = [
                         "original",
@@ -1075,17 +995,58 @@ with gr.Blocks(css=CSS) as app:
                     ]
 
                     with gr.Accordion("IP-Adapter 1", open=False, visible=True):
-                        image_ip1 = gr.Image(label="IP Image", type="filepath")
-                        mask_ip1 = gr.Image(label="IP Mask", type="filepath")
-                        model_ip1 = gr.Dropdown(value="plus_face", label="Model", choices=IP_MODELS)
-                        mode_ip1 = gr.Dropdown(value="original", label="Mode", choices=MODE_IP_OPTIONS)
-                        scale_ip1 = gr.Slider(minimum=0., maximum=2., step=0.01, value=0.7, label="Scale")
+                        image_ip1 = gr.Image(
+                            label="IP Image",
+                            type="filepath"
+                        )
+                        mask_ip1 = gr.Image(
+                            label="IP Mask",
+                            type="filepath"
+                        )
+                        model_ip1 = gr.Dropdown(
+                            value="plus_face",
+                            label="Model",
+                            choices=IP_MODELS
+                        )
+                        mode_ip1 = gr.Dropdown(
+                            value="original",
+                            label="Mode",
+                            choices=MODE_IP_OPTIONS
+                        )
+                        scale_ip1 = gr.Slider(
+                            minimum=0.,
+                            maximum=2.,
+                            step=0.01,
+                            value=0.7,
+                            label="Scale"
+                        )
+
                     with gr.Accordion("IP-Adapter 2", open=False, visible=True):
-                        image_ip2 = gr.Image(label="IP Image", type="filepath")
-                        mask_ip2 = gr.Image(label="IP Mask (optional)", type="filepath")
-                        model_ip2 = gr.Dropdown(value="base", label="Model", choices=IP_MODELS)
-                        mode_ip2 = gr.Dropdown(value="style", label="Mode", choices=MODE_IP_OPTIONS)
-                        scale_ip2 = gr.Slider(minimum=0., maximum=2., step=0.01, value=0.7, label="Scale")
+                        image_ip2 = gr.Image(
+                            label="IP Image",
+                            type="filepath"
+                        )
+                        mask_ip2 = gr.Image(
+                            label="IP Mask (optional)",
+                            type="filepath"
+                        )
+                        model_ip2 = gr.Dropdown(
+                            value="base",
+                            label="Model",
+                            choices=IP_MODELS
+                        )
+                        mode_ip2 = gr.Dropdown(
+                            value="style",
+                            label="Mode",
+                            choices=MODE_IP_OPTIONS
+                        )
+                        scale_ip2 = gr.Slider(
+                            minimum=0.,
+                            maximum=2.,
+                            step=0.01,
+                            value=0.7,
+                            label="Scale"
+                        )
 
                 with gr.Accordion("ControlNet / Img2img / Inpaint", open=False, visible=True):
                     image_control = gr.Image(
@@ -1122,7 +1083,10 @@ with gr.Blocks(css=CSS) as app:
                             choices_task = preprocessor_controlnet[task]
                         else:
                             choices_task = preprocessor_controlnet["canny"]
-                        return gr.update(choices=choices_task, value=choices_task[0])
+                        return gr.update(
+                            choices=choices_task,
+                            value=choices_task[0]
+                        )
 
 
                     task_gui.change(
@@ -1689,7 +1653,6 @@ with gr.Blocks(css=CSS) as app:
     )
 
 app.queue()
-
 app.launch(
     show_error=True,
     debug=True,

models/controlnet.py

@@ -0,0 +1,51 @@
+preprocessor_controlnet: dict = {
+    "openpose": [
+        "Openpose",
+        "None",
+    ],
+    "scribble": [
+        "HED",
+        "Pidinet",
+        "None",
+    ],
+    "softedge": [
+        "Pidinet",
+        "HED",
+        "HED safe",
+        "Pidinet safe",
+        "None",
+    ],
+    "segmentation": [
+        "UPerNet",
+        "None",
+    ],
+    "depth": [
+        "DPT",
+        "Midas",
+        "None",
+    ],
+    "normalbae": [
+        "NormalBae",
+        "None",
+    ],
+    "lineart": [
+        "Lineart",
+        "Lineart coarse",
+        "Lineart (anime)",
+        "None",
+        "None (anime)",
+    ],
+    "shuffle": [
+        "ContentShuffle",
+        "None",
+    ],
+    "canny": [
+        "Canny"
+    ],
+    "mlsd": [
+        "MLSD"
+    ],
+    "ip2p": [
+        "ip2p"
+    ]
+}

models/embeds.py

@@ -0,0 +1,5 @@
+download_embeds: list = [
+    'https://huggingface.co/datasets/Nerfgun3/bad_prompt/blob/main/bad_prompt_version2.pt',
+    'https://huggingface.co/embed/negative/resolve/main/EasyNegativeV2.safetensors',
+    'https://huggingface.co/embed/negative/resolve/main/bad-hands-5.pt',
+]

utils/model_utils.py

@@ -1,4 +1,8 @@
 def get_model_list(directory_path):
+    """
+    :param directory_path:
+    :return:
+    """
     import os
     model_list: list = []
     valid_extensions = {
@@ -16,4 +20,4 @@ def get_model_list(directory_path):
             # model_list.append((name_without_extension, file_path))
             model_list.append(file_path)
             print('\033[34mFILE: ' + file_path + '\033[0m')
-    return model_list
+    return model_list

utils/string_utils.py

@@ -12,3 +12,48 @@ def process_string(input_string: str):
         return result
     else:
         return None
+
+
+def extract_parameters(input_string: str) -> dict:
+    """
+    :param input_string:
+    :return:
+    """
+    import re
+    parameters: dict = {}
+    input_string: str = input_string.replace("\n", "")
+
+    if not "Negative prompt:" in input_string:
+        print("Negative prompt not detected")
+        parameters["prompt"] = input_string
+        return parameters
+
+    parm: list = input_string.split("Negative prompt:")
+    parameters["prompt"] = parm[0]
+    if not "Steps:" in parm[1]:
+        print("Steps not detected")
+        parameters["neg_prompt"] = parm[1]
+        return parameters
+    parm = parm[1].split("Steps:")
+    parameters["neg_prompt"] = parm[0]
+    input_string = "Steps:" + parm[1]
+
+    # Extracting Steps
+    steps_match = re.search(r'Steps: (\d+)', input_string)
+    if steps_match:
+        parameters['Steps'] = int(steps_match.group(1))
+
+    # Extracting Size
+    size_match = re.search(r'Size: (\d+x\d+)', input_string)
+    if size_match:
+        parameters['Size'] = size_match.group(1)
+        width, height = map(int, parameters['Size'].split('x'))
+        parameters['width'] = width
+        parameters['height'] = height
+
+    # Extracting other parameters
+    other_parameters = re.findall(r'(\w+): (.*?)(?=, \w+|$)', input_string)
+    for param in other_parameters:
+        parameters[param[0]] = param[1].strip('"')
+
+    return parameters