Upload 4 files

app.py

@@ -2,12 +2,11 @@ import spaces
 import gradio as gr
 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
+from stablepy.diffusers_vanilla.constants import FLUX_CN_UNION_MODES
 import torch
 import re
-import shutil
-import random
+from huggingface_hub import HfApi
 from stablepy import (
     CONTROLNET_MODEL_IDS,
     VALID_TASKS,
@@ -23,9 +22,9 @@ from stablepy import (
     SD15_TASKS,
     SDXL_TASKS,
 )
-import urllib.parse
+#import urllib.parse
 
-preprocessor_controlnet = {
+PREPROCESSOR_CONTROLNET = {
   "openpose": [
     "Openpose",
     "None",
@@ -97,7 +96,7 @@ preprocessor_controlnet = {
   ],
 }
 
-task_stablepy = {
+TASK_STABLEPY = {
     'txt2img': 'txt2img',
     'img2img': 'img2img',
     'inpaint': 'inpaint',
@@ -123,8 +122,35 @@ task_stablepy = {
     'tile ControlNet': 'tile',
 }
 
-task_model_list = list(task_stablepy.keys())
+TASK_MODEL_LIST = list(TASK_STABLEPY.keys())
 
+UPSCALER_DICT_GUI = {
+    None: None,
+    "Lanczos": "Lanczos",
+    "Nearest": "Nearest",
+    'Latent': 'Latent',
+    'Latent (antialiased)': 'Latent (antialiased)',
+    'Latent (bicubic)': 'Latent (bicubic)',
+    'Latent (bicubic antialiased)': 'Latent (bicubic antialiased)',
+    'Latent (nearest)': 'Latent (nearest)',
+    'Latent (nearest-exact)': 'Latent (nearest-exact)',
+    "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"
+}
+
+UPSCALER_KEYS = list(UPSCALER_DICT_GUI.keys())
 
 def download_things(directory, url, hf_token="", civitai_api_key=""):
     url = url.strip()
@@ -155,21 +181,19 @@ def download_things(directory, url, hf_token="", civitai_api_key=""):
     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 = []
     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]
+            # 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
 
-
 ## BEGIN MOD
 from modutils import (list_uniq, download_private_repo, get_model_id_list, get_tupled_embed_list,
     get_lora_model_list, get_all_lora_tupled_list, update_loras, apply_lora_prompt, set_prompt_loras,
@@ -178,7 +202,7 @@ from modutils import (list_uniq, download_private_repo, get_model_id_list, get_t
     get_tupled_model_list, save_gallery_images, set_optimization, set_sampler_settings,
     set_quick_presets, process_style_prompt, optimization_list, save_images,
     preset_styles, preset_quality, preset_sampler_setting, translate_to_en)
-from env import (hf_token, CIVITAI_API_KEY, HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
+from env import (HF_TOKEN, CIVITAI_API_KEY, HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
     HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
     directory_models, directory_loras, directory_vaes, directory_embeds, directory_embeds_sdxl,
     directory_embeds_positive_sdxl, load_diffusers_format_model,
@@ -200,18 +224,18 @@ load_diffusers_format_model = list_uniq(get_model_id_list() + load_diffusers_for
 # 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
 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)
@@ -232,46 +256,80 @@ def get_embed_list(pipeline_name):
 
 print('\033[33m🏁 Download and listing of valid models completed.\033[0m')
 
-upscaler_dict_gui = {
-    None: None,
-    "Lanczos": "Lanczos",
-    "Nearest": "Nearest",
-    'Latent': 'Latent',
-    'Latent (antialiased)': 'Latent (antialiased)',
-    'Latent (bicubic)': 'Latent (bicubic)',
-    'Latent (bicubic antialiased)': 'Latent (bicubic antialiased)',
-    'Latent (nearest)': 'Latent (nearest)',
-    'Latent (nearest-exact)': 'Latent (nearest-exact)',
-    "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"
+
+#######################
+# GUI
+#######################
+import gradio as gr
+import logging
+logging.getLogger("diffusers").setLevel(logging.ERROR)
+import diffusers
+diffusers.utils.logging.set_verbosity(40)
+import warnings
+warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffusers")
+warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
+warnings.filterwarnings(action="ignore", category=FutureWarning, module="transformers")
+## BEGIN MOD
+from stablepy import logger
+logger.setLevel(logging.CRITICAL)
+
+from v2 import V2_ALL_MODELS, v2_random_prompt, v2_upsampling_prompt
+from utils import (gradio_copy_text, COPY_ACTION_JS, gradio_copy_prompt,
+    V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS, V2_LENGTH_OPTIONS, V2_IDENTITY_OPTIONS)
+from tagger import (predict_tags_wd, convert_danbooru_to_e621_prompt,
+    remove_specific_prompt, insert_recom_prompt, insert_model_recom_prompt,
+    compose_prompt_to_copy, translate_prompt, select_random_character)
+def description_ui():
+    gr.Markdown(
+        """
+## Danbooru Tags Transformer V2 Demo with WD Tagger
+(Image =>) Prompt => Upsampled longer prompt
+- Mod of p1atdev's [Danbooru Tags Transformer V2 Demo](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) and [WD Tagger with 🤗 transformers](https://huggingface.co/spaces/p1atdev/wd-tagger-transformers).
+- Models: p1atdev's [wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf), [dart-v2-moe-sft](https://huggingface.co/p1atdev/dart-v2-moe-sft)
+"""
+    )
+## END MOD
+
+msg_inc_vae = (
+    "Use the right VAE for your model to maintain image quality. The wrong"
+    " VAE can lead to poor results, like blurriness in the generated images."
+)
+
+SDXL_TASK = [k for k, v in TASK_STABLEPY.items() if v in SDXL_TASKS]
+SD_TASK = [k for k, v in TASK_STABLEPY.items() if v in SD15_TASKS]
+FLUX_TASK = list(TASK_STABLEPY.keys())[:3] + [k for k, v in TASK_STABLEPY.items() if v in FLUX_CN_UNION_MODES.keys()]
+
+MODEL_TYPE_TASK = {
+    "SD 1.5": SD_TASK,
+    "SDXL": SDXL_TASK,
+    "FLUX": FLUX_TASK,
 }
 
-upscaler_keys = list(upscaler_dict_gui.keys())
+MODEL_TYPE_CLASS = {
+    "diffusers:StableDiffusionPipeline": "SD 1.5",
+    "diffusers:StableDiffusionXLPipeline": "SDXL",
+    "diffusers:FluxPipeline": "FLUX",
+}
+
+POST_PROCESSING_SAMPLER = ["Use same sampler"] + scheduler_names[:-2]
+
+SUBTITLE_GUI = (
+    "### This demo uses [diffusers](https://github.com/huggingface/diffusers)"
+    " to perform different tasks in image generation."
+)
 
 def extract_parameters(input_string):
     parameters = {}
     input_string = input_string.replace("\n", "")
 
-    if not "Negative prompt:" in input_string:
+    if "Negative prompt:" not 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]:
+    if "Steps:" not in parm[1]:
         print("Steps not detected")
         parameters["neg_prompt"] = parm[1]
         return parameters
@@ -299,45 +357,6 @@ def extract_parameters(input_string):
 
     return parameters
 
-
-#######################
-# GUI
-#######################
-import gradio as gr
-from PIL import Image
-import IPython.display
-import time, json
-from IPython.utils import capture
-import logging
-logging.getLogger("diffusers").setLevel(logging.ERROR)
-import diffusers
-diffusers.utils.logging.set_verbosity(40)
-import warnings
-warnings.filterwarnings(action="ignore", category=FutureWarning, module="diffusers")
-warnings.filterwarnings(action="ignore", category=UserWarning, module="diffusers")
-warnings.filterwarnings(action="ignore", category=FutureWarning, module="transformers")
-## BEGIN MOD
-from stablepy import logger
-logger.setLevel(logging.CRITICAL)
-
-from v2 import V2_ALL_MODELS, v2_random_prompt, v2_upsampling_prompt
-from utils import (gradio_copy_text, COPY_ACTION_JS, gradio_copy_prompt,
-    V2_ASPECT_RATIO_OPTIONS, V2_RATING_OPTIONS, V2_LENGTH_OPTIONS, V2_IDENTITY_OPTIONS)
-from tagger import (predict_tags_wd, convert_danbooru_to_e621_prompt,
-    remove_specific_prompt, insert_recom_prompt, insert_model_recom_prompt,
-    compose_prompt_to_copy, translate_prompt, select_random_character)
-def description_ui():
-    gr.Markdown(
-        """
-## Danbooru Tags Transformer V2 Demo with WD Tagger
-(Image =>) Prompt => Upsampled longer prompt
-- Mod of p1atdev's [Danbooru Tags Transformer V2 Demo](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) and [WD Tagger with 🤗 transformers](https://huggingface.co/spaces/p1atdev/wd-tagger-transformers).
-- Models: p1atdev's [wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf), [dart-v2-moe-sft](https://huggingface.co/p1atdev/dart-v2-moe-sft)
-"""
-    )
-## END MOD
-
-
 def info_html(json_data, title, subtitle):
     return f"""
         <div style='padding: 0; border-radius: 10px;'>
@@ -349,6 +368,18 @@ def info_html(json_data, title, subtitle):
         </div>
         """
 
+def get_model_type(repo_id: str):
+    api = HfApi(token=os.environ.get("HF_TOKEN"))  # if use private or gated model
+    default = "SD 1.5"
+    try:
+        model = api.model_info(repo_id=repo_id, timeout=5.0)
+        tags = model.tags
+        for tag in tags:
+            if tag in MODEL_TYPE_CLASS.keys(): return MODEL_TYPE_CLASS.get(tag, default)
+    except Exception:
+        return default
+    return default
+
 class GuiSD:
     def __init__(self):
         self.model = None
@@ -368,23 +399,21 @@ class GuiSD:
         yield f"Loading model: {model_name}"
         
         vae_model = vae_model if vae_model != "None" else None
+        model_type = get_model_type(model_name)
 
-        if model_name in model_list:
-            model_is_xl = "xl" in model_name.lower()
-            sdxl_in_vae = vae_model and "sdxl" in vae_model.lower()
-            model_type = "SDXL" if model_is_xl else "SD 1.5"
-            incompatible_vae = (model_is_xl and vae_model and not sdxl_in_vae) or (not model_is_xl and sdxl_in_vae)
-
-            if incompatible_vae:
-                vae_model = None
+        if vae_model:
+            vae_type = "SDXL" if "sdxl" in vae_model.lower() else "SD 1.5"
+            if model_type != vae_type:
+                gr.Info(msg_inc_vae)
 
         self.model.device = torch.device("cpu")
+        dtype_model = torch.bfloat16 if model_type == "FLUX" else torch.float16
 
         self.model.load_pipe(
             model_name,
-            task_name=task_stablepy[task],
+            task_name=TASK_STABLEPY[task],
             vae_model=vae_model if vae_model != "None" else None,
-            type_model_precision=torch.float16 if "flux" not in model_name.lower() else torch.bfloat16,
+            type_model_precision=dtype_model,
             retain_task_model_in_cache=False,
         )
         yield f"Model loaded: {model_name}"
@@ -504,6 +533,8 @@ class GuiSD:
         vae_msg = f"VAE: {vae_model}" if vae_model else ""
         msg_lora = []
 
+        print("Config model:", model_name, vae_model, loras_list)
+
 ## BEGIN MOD
         global lora_model_list
         lora_model_list = get_lora_model_list()
@@ -512,33 +543,8 @@ class GuiSD:
                                   lora_scale3, lora4, lora_scale4, lora5, lora_scale5)
         prompt, neg_prompt = insert_model_recom_prompt(prompt, neg_prompt, model_name)
 ## END MOD
-        
-        if model_name in model_list:
-            model_is_xl = "xl" in model_name.lower()
-            sdxl_in_vae = vae_model and "sdxl" in vae_model.lower()
-            model_type = "SDXL" if model_is_xl else "SD 1.5"
-            incompatible_vae = (model_is_xl and vae_model and not sdxl_in_vae) or (not model_is_xl and sdxl_in_vae)
-
-            if incompatible_vae:
-                msg_inc_vae = (
-                    f"The selected VAE is for a { 'SD 1.5' if model_is_xl else 'SDXL' } model, but you"
-                    f" are using a { model_type } model. The default VAE "
-                    "will be used."
-                )
-                gr.Info(msg_inc_vae)
-                vae_msg = msg_inc_vae
-                vae_model = None
 
-            for la in loras_list:
-                if la is not None and la != "None" and la != "" 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)
-
-        task = task_stablepy[task]
+        task = TASK_STABLEPY[task]
 
         params_ip_img = []
         params_ip_msk = []
@@ -560,18 +566,6 @@ class GuiSD:
                 params_ip_mode.append(modeip)
                 params_ip_scale.append(scaleip)
 
-        model_precision = torch.float16 if "flux" not in model_name.lower() else torch.bfloat16
-
-        # First load
-        if not self.model:
-            print("Loading model...")
-            self.model = Model_Diffusers(
-                base_model_id=model_name,
-                task_name=task,
-                vae_model=vae_model if vae_model != "None" else None,
-                type_model_precision=model_precision,
-                retain_task_model_in_cache=retain_task_cache_gui,
-            )
         self.model.stream_config(concurrency=5, latent_resize_by=1, vae_decoding=False)
 
         if task != "txt2img" and not image_control:
@@ -580,62 +574,47 @@ class GuiSD:
         if task == "inpaint" and not image_mask:
             raise ValueError("No mask image found: Specify one in 'Image Mask'")
 
-        if upscaler_model_path in upscaler_keys[:9]:
+        if upscaler_model_path in UPSCALER_KEYS[:9]:
             upscaler_model = upscaler_model_path
         else:
             directory_upscalers = 'upscalers'
             os.makedirs(directory_upscalers, exist_ok=True)
 
-            url_upscaler = upscaler_dict_gui[upscaler_model_path]
+            url_upscaler = UPSCALER_DICT_GUI[upscaler_model_path]
 
             if not os.path.exists(f"./upscalers/{url_upscaler.split('/')[-1]}"):
-                download_things(directory_upscalers, url_upscaler, hf_token)
+                download_things(directory_upscalers, url_upscaler, HF_TOKEN)
 
             upscaler_model = f"./upscalers/{url_upscaler.split('/')[-1]}"
 
         logging.getLogger("ultralytics").setLevel(logging.INFO if adetailer_verbose else logging.ERROR)
 
-        print("Config model:", model_name, vae_model, loras_list)
-
-        self.model.load_pipe(
-            model_name,
-            task_name=task,
-            vae_model=vae_model if vae_model != "None" else None,
-            type_model_precision=model_precision,
-            retain_task_model_in_cache=retain_task_cache_gui,
-        )
-
-## BEGIN MOD
-#        if textual_inversion and self.model.class_name == "StableDiffusionXLPipeline":
-#            print("No Textual inversion for SDXL")
-## END MOD
-
         adetailer_params_A = {
-            "face_detector_ad" : face_detector_ad_a,
-            "person_detector_ad" : person_detector_ad_a,
-            "hand_detector_ad" : hand_detector_ad_a,
+            "face_detector_ad": face_detector_ad_a,
+            "person_detector_ad": person_detector_ad_a,
+            "hand_detector_ad": hand_detector_ad_a,
             "prompt": prompt_ad_a,
-            "negative_prompt" : negative_prompt_ad_a,
-            "strength" : strength_ad_a,
+            "negative_prompt": negative_prompt_ad_a,
+            "strength": strength_ad_a,
             # "image_list_task" : None,
-            "mask_dilation" : mask_dilation_a,
-            "mask_blur" : mask_blur_a,
-            "mask_padding" : mask_padding_a,
-            "inpaint_only" : adetailer_inpaint_only,
-            "sampler" : adetailer_sampler,
+            "mask_dilation": mask_dilation_a,
+            "mask_blur": mask_blur_a,
+            "mask_padding": mask_padding_a,
+            "inpaint_only": adetailer_inpaint_only,
+            "sampler": adetailer_sampler,
         }
 
         adetailer_params_B = {
-            "face_detector_ad" : face_detector_ad_b,
-            "person_detector_ad" : person_detector_ad_b,
-            "hand_detector_ad" : hand_detector_ad_b,
+            "face_detector_ad": face_detector_ad_b,
+            "person_detector_ad": person_detector_ad_b,
+            "hand_detector_ad": hand_detector_ad_b,
             "prompt": prompt_ad_b,
-            "negative_prompt" : negative_prompt_ad_b,
-            "strength" : strength_ad_b,
+            "negative_prompt": negative_prompt_ad_b,
+            "strength": strength_ad_b,
             # "image_list_task" : None,
-            "mask_dilation" : mask_dilation_b,
-            "mask_blur" : mask_blur_b,
-            "mask_padding" : mask_padding_b,
+            "mask_dilation": mask_dilation_b,
+            "mask_blur": mask_blur_b,
+            "mask_padding": mask_padding_b,
         }
         pipe_params = {
             "prompt": prompt,
@@ -748,6 +727,40 @@ class GuiSD:
                 
             yield img, info_state
 
+def update_task_options(model_name, task_name):
+    new_choices = MODEL_TYPE_TASK[get_model_type(model_name)]
+
+    if task_name not in new_choices:
+        task_name = "txt2img"
+
+    return gr.update(value=task_name, choices=new_choices)
+
+# def sd_gen_generate_pipeline(*args):
+
+#     # Load lora in CPU
+#     status_lora = sd_gen.model.lora_merge(
+#         lora_A=args[7] if args[7] != "None" else None, lora_scale_A=args[8],
+#         lora_B=args[9] if args[9] != "None" else None, lora_scale_B=args[10],
+#         lora_C=args[11] if args[11] != "None" else None, lora_scale_C=args[12],
+#         lora_D=args[13] if args[13] != "None" else None, lora_scale_D=args[14],
+#         lora_E=args[15] if args[15] != "None" else None, lora_scale_E=args[16],
+#     )
+
+#     lora_list = [args[7], args[9], args[11], args[13], args[15]]
+#     print(status_lora)
+#     for status, lora in zip(status_lora, lora_list):
+#         if status:
+#             gr.Info(f"LoRA loaded: {lora}")
+#         elif status is not None:
+#             gr.Warning(f"Failed to load LoRA: {lora}")
+
+#     # if status_lora == [None] * 5 and self.model.lora_memory != [None] * 5:
+#     #     gr.Info(f"LoRAs in cache: {", ".join(str(x) for x in self.model.lora_memory if x is not None)}")
+
+#     yield from sd_gen.generate_pipeline(*args)
+
+
+# sd_gen_generate_pipeline.zerogpu = True
 sd_gen = GuiSD()
 
 ## BEGIN MOD
@@ -759,27 +772,7 @@ CSS ="""
 .lora { min-width:480px; !important; }
 #model-info { text-align:center; }
 """
-## END MOD
 
-sdxl_task = [k for k, v in task_stablepy.items() if v in SDXL_TASKS ]
-sd_task = [k for k, v in task_stablepy.items() if v in SD15_TASKS ]
-def update_task_options(model_name, task_name):
-    if model_name in model_list:
-        if "xl" in model_name.lower() or "pony" in model_name.lower():
-            new_choices = sdxl_task
-        else:
-            new_choices = sd_task
-
-        if task_name not in new_choices:
-            task_name = "txt2img"
-
-        return gr.update(value=task_name, choices=new_choices)
-    else:
-        return gr.update(value=task_name, choices=task_model_list)
-
-POST_PROCESSING_SAMPLER = ["Use same sampler"] + scheduler_names[:-2]
-
-## BEGIN MOD
 with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, css=CSS, delete_cache=(60, 3600)) as app:
     gr.Markdown("# 🧩 DiffuseCraft Mod")
     gr.Markdown(
@@ -793,7 +786,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                 with gr.Column(scale=2):
                     interface_mode_gui = gr.Radio(label="Quick settings", choices=["Simple", "Standard", "Fast", "LoRA"], value="Standard")
                     with gr.Accordion("Model and Task", open=False) as menu_model:
-                        task_gui = gr.Dropdown(label="Task", choices=sdxl_task, value=task_model_list[0])
+                        task_gui = gr.Dropdown(label="Task", choices=SDXL_TASK, value=TASK_MODEL_LIST[0])
                         with gr.Group():
                             model_name_gui = gr.Dropdown(label="Model", info="You can enter a huggingface model repo_id to want to use.", choices=get_tupled_model_list(model_list), value="votepurchase/animagine-xl-3.1", allow_custom_value=True)
                             model_info_gui = gr.Markdown(elem_id="model-info")
@@ -1036,7 +1029,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
 
                     with gr.Column() as menu_advanced:
                         with gr.Accordion("Hires fix", open=False, visible=True) as menu_hires:
-                            upscaler_model_path_gui = gr.Dropdown(label="Upscaler", choices=upscaler_keys, value=upscaler_keys[0])
+                            upscaler_model_path_gui = gr.Dropdown(label="Upscaler", choices=UPSCALER_KEYS, value=UPSCALER_KEYS[0])
                             with gr.Row():
                                 upscaler_increases_size_gui = gr.Slider(minimum=1.1, maximum=6., step=0.1, value=1.0, label="Upscale by")
                                 esrgan_tile_gui = gr.Slider(minimum=0, value=100, maximum=500, step=1, label="ESRGAN Tile")
@@ -1108,21 +1101,21 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                                     info="This option adjusts the level of changes for img2img and inpainting."
                                 )
                                 image_resolution_gui = gr.Slider(minimum=64, maximum=2048, step=64, value=1024, label="Image Resolution")
-                                preprocessor_name_gui = gr.Dropdown(label="Preprocessor Name", choices=preprocessor_controlnet["canny"], value=preprocessor_controlnet["canny"][0])
-
-                            def change_preprocessor_choices(task):
-                                task = task_stablepy[task]
-                                if task in preprocessor_controlnet.keys():
-                                    choices_task = preprocessor_controlnet[task]
-                                else:
-                                    choices_task = preprocessor_controlnet["canny"]
-                                return gr.update(choices=choices_task, value=choices_task[0])
-
-                            task_gui.change(
-                                change_preprocessor_choices,
-                                [task_gui],
-                                [preprocessor_name_gui],
-                            )
+                                preprocessor_name_gui = gr.Dropdown(label="Preprocessor Name", choices=PREPROCESSOR_CONTROLNET["canny"])
+
+                                def change_preprocessor_choices(task):
+                                    task = TASK_STABLEPY[task]
+                                    if task in PREPROCESSOR_CONTROLNET.keys():
+                                        choices_task = PREPROCESSOR_CONTROLNET[task]
+                                    else:
+                                        choices_task = PREPROCESSOR_CONTROLNET["canny"]
+                                    return gr.update(choices=choices_task, value=choices_task[0])
+
+                                task_gui.change(
+                                    change_preprocessor_choices,
+                                    [task_gui],
+                                    [preprocessor_name_gui],
+                                )
                             with gr.Row():
                                 preprocess_resolution_gui = gr.Slider(minimum=64, maximum=2048, step=64, value=512, label="Preprocess Resolution")
                                 low_threshold_gui = gr.Slider(minimum=1, maximum=255, step=1, value=100, label="Canny low threshold")
@@ -1166,7 +1159,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
 
                             try:
                                 style_names_found = sd_gen.model.STYLE_NAMES
-                            except:
+                            except Exception:
                                 style_names_found = STYLE_NAMES
 
                             style_prompt_gui = gr.Dropdown(
@@ -1293,7 +1286,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                         img_width_gui,
                         model_name_gui,
                     ],
-                    outputs=[result_images],
+                    outputs=[result_images, actual_task_info],
                     cache_examples=False,
                     #elem_id="examples",
                 )

env.py

@@ -1,7 +1,7 @@
 import os
 
 CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")
-hf_token = os.environ.get("HF_TOKEN")
+HF_TOKEN = os.environ.get("HF_TOKEN")
 hf_read_token = os.environ.get('HF_READ_TOKEN') # only use for private repo
 
 # - **List Models**

modutils.py

@@ -8,7 +8,7 @@ from pathlib import Path
 
 from env import (HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
     HF_MODEL_USER_EX, HF_MODEL_USER_LIKES,
-    directory_loras, hf_read_token, hf_token, CIVITAI_API_KEY)
+    directory_loras, hf_read_token, HF_TOKEN, CIVITAI_API_KEY)
 
 
 def get_user_agent():
@@ -431,7 +431,7 @@ def download_lora(dl_urls: str):
     for url in [url.strip() for url in dl_urls.split(',')]:
         local_path = f"{directory_loras}/{url.split('/')[-1]}"
         if not Path(local_path).exists():
-            download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+            download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
             urls.append(url)
     after = get_local_model_list(directory_loras)
     new_files = list_sub(after, before)
@@ -693,7 +693,7 @@ def get_my_lora(link_url):
     before = get_local_model_list(directory_loras)
     for url in [url.strip() for url in link_url.split(',')]:
         if not Path(f"{directory_loras}/{url.split('/')[-1]}").exists():
-            download_things(directory_loras, url, hf_token, CIVITAI_API_KEY)
+            download_things(directory_loras, url, HF_TOKEN, CIVITAI_API_KEY)
     after = get_local_model_list(directory_loras)
     new_files = list_sub(after, before)
     for file in new_files: