Upload 3 files

dc.py

@@ -1,12 +1,11 @@
 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
+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,
@@ -22,7 +21,7 @@ from stablepy import (
     SD15_TASKS,
     SDXL_TASKS,
 )
-import urllib.parse
+#import urllib.parse
 import gradio as gr
 from PIL import Image
 import IPython.display
@@ -40,7 +39,7 @@ from stablepy import logger
 logger.setLevel(logging.CRITICAL)
 
 from env import (
-    hf_token, hf_read_token, # to use only for private repos
+    HF_TOKEN, hf_read_token, # to use only for private repos
     CIVITAI_API_KEY, HF_LORA_PRIVATE_REPOS1, HF_LORA_PRIVATE_REPOS2,
     HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
     HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
@@ -49,7 +48,7 @@ from env import (
     load_diffusers_format_model, download_model_list, download_lora_list,
     download_vae_list, download_embeds)
 
-preprocessor_controlnet = {
+PREPROCESSOR_CONTROLNET = {
   "openpose": [
     "Openpose",
     "None",
@@ -86,6 +85,13 @@ preprocessor_controlnet = {
     "None",
     "None (anime)",
   ],
+  "lineart_anime": [
+    "Lineart",
+    "Lineart coarse",
+    "Lineart (anime)",
+    "None",
+    "None (anime)",
+  ],
   "shuffle": [
     "ContentShuffle",
     "None",
@@ -114,7 +120,7 @@ preprocessor_controlnet = {
   ],
 }
 
-task_stablepy = {
+TASK_STABLEPY = {
     'txt2img': 'txt2img',
     'img2img': 'img2img',
     'inpaint': 'inpaint',
@@ -140,7 +146,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()
@@ -171,21 +205,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 (to_list, list_uniq, list_sub, get_model_id_list, get_tupled_embed_list,
                       get_tupled_model_list, get_lora_model_list, download_private_repo)
@@ -203,24 +235,21 @@ download_private_repo(HF_VAE_PRIVATE_REPO, directory_vaes, False)
 load_diffusers_format_model = list_uniq(load_diffusers_format_model + get_model_id_list())
 ## END MOD
 
-CIVITAI_API_KEY = os.environ.get("CIVITAI_API_KEY")
-hf_token = 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
 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)
@@ -237,53 +266,45 @@ embed_sdxl_list = get_model_list(directory_embeds_sdxl) + get_model_list(directo
 
 def get_embed_list(pipeline_name):
     return get_tupled_embed_list(embed_sdxl_list if pipeline_name == "StableDiffusionXLPipeline" else embed_list)
-
-
 ## END MOD
 
 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"
+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]
 
 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
@@ -311,6 +332,17 @@ def extract_parameters(input_string):
 
     return parameters
 
+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
 
 ## BEGIN MOD
 class GuiSD:
@@ -319,12 +351,12 @@ class GuiSD:
     
         print("Loading model...")
         self.model = Model_Diffusers(
-            base_model_id="cagliostrolab/animagine-xl-3.1",
+            base_model_id="Lykon/dreamshaper-8",
             task_name="txt2img",
             vae_model=None,
             type_model_precision=torch.float16,
             retain_task_model_in_cache=False,
-            #device="cpu",
+            device="cpu",
         )
         self.model.device = torch.device("cpu") #
 
@@ -344,28 +376,27 @@ 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,
+            type_model_precision=dtype_model,
             retain_task_model_in_cache=False,
         )
         yield f"Model loaded: {model_name}"
 
     @spaces.GPU
+    @torch.inference_mode()
     def generate_pipeline(
         self,
         prompt,
@@ -479,38 +510,15 @@ 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
         prompt, neg_prompt = insert_model_recom_prompt(prompt, neg_prompt, model_name)
         global lora_model_list
         lora_model_list = get_lora_model_list()
 ## 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 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 = []
@@ -532,83 +540,53 @@ class GuiSD:
                 params_ip_mode.append(modeip)
                 params_ip_scale.append(scaleip)
 
-        # First load
-        self.model.device = torch.device("cuda:0")
-        model_precision = torch.float16
-        if not self.model:
-            from stablepy import Model_Diffusers
-
-            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,
-            )
-
         if task != "txt2img" and not image_control:
             raise ValueError("No control image found: To use this function, you have to upload an image in 'Image ControlNet/Inpaint/Img2img'")
 
         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,
@@ -690,8 +668,10 @@ class GuiSD:
             "ip_adapter_scale": params_ip_scale,
         }
 
-        # 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")
+        self.model.device = torch.device("cuda:0")
+        if hasattr(self.model.pipe, "transformer") and loras_list != ["None"] * 5:
+            self.model.pipe.transformer.to(self.model.device)
+            print("transformer to cuda")
 
         progress(1, desc="Inference preparation completed. Starting inference...")
 
@@ -699,8 +679,35 @@ class GuiSD:
         return self.infer_short(self.model, pipe_params, progress), info_state
 ## END MOD
 
+# 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
 
 from pathlib import Path
+import random
 from modutils import (safe_float, escape_lora_basename, to_lora_key, to_lora_path,
     get_local_model_list, get_private_lora_model_lists, get_valid_lora_name,
     get_valid_lora_path, get_valid_lora_wt, get_lora_info,

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**
@@ -50,8 +50,9 @@ load_diffusers_format_model = [
     'eienmojiki/Anything-XL',
     'eienmojiki/Starry-XL-v5.2',
     'gsdf/CounterfeitXL',
-    'kitty7779/ponyDiffusionV6XL',
+    'John6666/silvermoon-mix-01xl-v11-sdxl',
     'WhiteAiZ/autismmixSDXL_autismmixConfetti_diffusers',
+    'kitty7779/ponyDiffusionV6XL',
     'GraydientPlatformAPI/aniverse-pony',
     'John6666/mistoon-anime-ponyalpha-sdxl',
     'John6666/ebara-mfcg-pony-mix-v12-sdxl',
@@ -61,10 +62,10 @@ load_diffusers_format_model = [
     'John6666/cyberrealistic-pony-v63-sdxl',
     'GraydientPlatformAPI/realcartoon-pony-diffusion',
     'John6666/nova-anime-xl-pony-v5-sdxl',
+    'John6666/autismmix-sdxl-autismmix-pony-sdxl',
     'yodayo-ai/kivotos-xl-2.0',
     'yodayo-ai/holodayo-xl-2.1',
     'yodayo-ai/clandestine-xl-1.0',
-    'John6666/silvermoon-mix-01xl-v11-sdxl',
     'digiplay/majicMIX_sombre_v2',
     'digiplay/majicMIX_realistic_v6',
     'digiplay/majicMIX_realistic_v7',
@@ -93,6 +94,12 @@ load_diffusers_format_model = [
     'Eugeoter/artiwaifu-diffusion-2.0',
     'Raelina/Rae-Diffusion-XL-V2',
     'Raelina/Raemu-XL-V4',
+    "camenduru/FLUX.1-dev-diffusers",
+    "black-forest-labs/FLUX.1-schnell",
+    "sayakpaul/FLUX.1-merged",
+    "ostris/OpenFLUX.1",
+    "multimodalart/FLUX.1-dev2pro-full",
+    "Raelina/Raemu-Flux",
 ]
 
 # List all Models for specified user

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():
@@ -100,6 +100,24 @@ def safe_float(input):
     return output
 
 
+from PIL import Image
+def save_images(images: list[Image.Image], metadatas: list[str]):
+    from PIL import PngImagePlugin
+    import uuid
+    try:
+        output_images = []
+        for image, metadata in zip(images, metadatas):
+            info = PngImagePlugin.PngInfo()
+            info.add_text("metadata", metadata)
+            savefile = f"{str(uuid.uuid4())}.png"
+            image.save(savefile, "PNG", pnginfo=info)
+            output_images.append(str(Path(savefile).resolve()))
+        return output_images
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e
+
+
 def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
     from datetime import datetime, timezone, timedelta
     progress(0, desc="Updating gallery...")
@@ -209,11 +227,16 @@ def get_model_id_list():
         model_ids.append(model.id) if not model.private else ""
     anime_models = []
     real_models = []
+    anime_models_flux = []
+    real_models_flux = []
     for model in models_ex:
-        if not model.private and not model.gated and "diffusers:FluxPipeline" not in model.tags:
-            anime_models.append(model.id) if "anime" in model.tags else real_models.append(model.id)
+        if not model.private and not model.gated:
+            if "diffusers:FluxPipeline" in model.tags: anime_models_flux.append(model.id) if "anime" in model.tags else real_models_flux.append(model.id)
+            else: anime_models.append(model.id) if "anime" in model.tags else real_models.append(model.id)
     model_ids.extend(anime_models)
     model_ids.extend(real_models)
+    model_ids.extend(anime_models_flux)
+    model_ids.extend(real_models_flux)
     model_id_list = model_ids.copy()
     return model_ids
 
@@ -408,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)
@@ -670,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:
@@ -727,8 +750,7 @@ def move_file_lora(filepaths):
 
 
 def get_civitai_info(path):
-    global civitai_not_exists_list
-    global loras_url_to_path_dict
+    global civitai_not_exists_list, loras_url_to_path_dict
     import requests
     from requests.adapters import HTTPAdapter
     from urllib3.util import Retry
@@ -1224,7 +1246,7 @@ def get_model_pipeline(repo_id: str):
     try:
         if " " in repo_id or not api.repo_exists(repo_id): return default
         model = api.model_info(repo_id=repo_id)
-    except Exception as e:
+    except Exception:
         return default
     if model.private or model.gated: return default
     tags = model.tags