Upload 4 files

app.py

@@ -6,8 +6,10 @@ from stablepy import (
     SCHEDULE_PREDICTION_TYPE_OPTIONS,
     check_scheduler_compatibility,
     TASK_AND_PREPROCESSORS,
+    FACE_RESTORATION_MODELS,
 )
 from constants import (
+    DIRECTORY_UPSCALERS,
     TASK_STABLEPY,
     TASK_MODEL_LIST,
     UPSCALER_DICT_GUI,
@@ -77,6 +79,10 @@ download_private_repo(HF_LORA_ESSENTIAL_PRIVATE_REPO, DIRECTORY_LORAS, True)
 download_private_repo(HF_VAE_PRIVATE_REPO, DIRECTORY_VAES, False)
 ## END MOD
 
+directories = [DIRECTORY_MODELS, DIRECTORY_LORAS, DIRECTORY_VAES, DIRECTORY_EMBEDS, DIRECTORY_UPSCALERS]
+for directory in directories:
+    os.makedirs(directory, exist_ok=True)
+
 # - **Download Models**
 DOWNLOAD_MODEL = ", ".join(DOWNLOAD_MODEL_LIST)
 # - **Download VAEs**
@@ -129,6 +135,7 @@ flux_pipe = FluxPipeline.from_pretrained(
 components = flux_pipe.components
 components.pop("transformer", None)
 delete_model(flux_repo)
+# components = None
 
 #######################
 # GUI
@@ -331,8 +338,8 @@ class GuiSD:
         syntax_weights,
         upscaler_model_path,
         upscaler_increases_size,
-        esrgan_tile,
-        esrgan_tile_overlap,
+        upscaler_tile_size,
+        upscaler_tile_overlap,
         hires_steps,
         hires_denoising_strength,
         hires_sampler,
@@ -396,6 +403,9 @@ class GuiSD:
         mode_ip2,
         scale_ip2,
         pag_scale,
+        face_restoration_model,
+        face_restoration_visibility,
+        face_restoration_weight,
     ):
         info_state = html_template_message("Navigating latent space...")
         yield info_state, gr.update(), gr.update()
@@ -448,18 +458,15 @@ 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 "https://" not in str(UPSCALER_DICT_GUI[upscaler_model_path]):
             upscaler_model = upscaler_model_path
         else:
-            directory_upscalers = 'upscalers'
-            os.makedirs(directory_upscalers, exist_ok=True)
-
             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)
+            if not os.path.exists(f"./{DIRECTORY_UPSCALERS}/{url_upscaler.split('/')[-1]}"):
+                download_things(DIRECTORY_UPSCALERS, url_upscaler, HF_TOKEN)
 
-            upscaler_model = f"./upscalers/{url_upscaler.split('/')[-1]}"
+            upscaler_model = f"./{DIRECTORY_UPSCALERS}/{url_upscaler.split('/')[-1]}"
 
         logging.getLogger("ultralytics").setLevel(logging.INFO if adetailer_verbose else logging.ERROR)
 
@@ -563,8 +570,8 @@ class GuiSD:
             "t2i_adapter_conditioning_factor": float(t2i_adapter_conditioning_factor),
             "upscaler_model_path": upscaler_model,
             "upscaler_increases_size": upscaler_increases_size,
-            "esrgan_tile": esrgan_tile,
-            "esrgan_tile_overlap": esrgan_tile_overlap,
+            "upscaler_tile_size": upscaler_tile_size,
+            "upscaler_tile_overlap": upscaler_tile_overlap,
             "hires_steps": hires_steps,
             "hires_denoising_strength": hires_denoising_strength,
             "hires_prompt": hires_prompt,
@@ -579,6 +586,9 @@ class GuiSD:
             "ip_adapter_model": params_ip_model,
             "ip_adapter_mode": params_ip_mode,
             "ip_adapter_scale": params_ip_scale,
+            "face_restoration_model": face_restoration_model,
+            "face_restoration_visibility": face_restoration_visibility,
+            "face_restoration_weight": face_restoration_weight,
         }
 
         # kwargs for diffusers pipeline
@@ -729,22 +739,26 @@ def sd_gen_generate_pipeline(*args):
 
 
 @spaces.GPU(duration=15)
-def esrgan_upscale(image, upscaler_name, upscaler_size):
+def process_upscale(image, upscaler_name, upscaler_size):
     if image is None: return None
 
     from stablepy.diffusers_vanilla.utils import save_pil_image_with_metadata
-    from stablepy import UpscalerESRGAN
+    from stablepy import load_upscaler_model
 
+    image = image.convert("RGB")
     exif_image = extract_exif_data(image)
 
-    url_upscaler = UPSCALER_DICT_GUI[upscaler_name]
-    directory_upscalers = 'upscalers'
-    os.makedirs(directory_upscalers, exist_ok=True)
-    if not os.path.exists(f"./upscalers/{url_upscaler.split('/')[-1]}"):
-        download_things(directory_upscalers, url_upscaler, HF_TOKEN)
+    name_upscaler = UPSCALER_DICT_GUI[upscaler_name]
+
+    if "https://" in str(name_upscaler):
 
-    scaler_beta = UpscalerESRGAN(0, 0)
-    image_up = scaler_beta.upscale(image, upscaler_size, f"./upscalers/{url_upscaler.split('/')[-1]}")
+        if not os.path.exists(f"./{DIRECTORY_UPSCALERS}/{name_upscaler.split('/')[-1]}"):
+            download_things(DIRECTORY_UPSCALERS, name_upscaler, HF_TOKEN)
+
+        name_upscaler = f"./{DIRECTORY_UPSCALERS}/{name_upscaler.split('/')[-1]}"
+
+    scaler_beta = load_upscaler_model(model=name_upscaler, tile=0, tile_overlap=8, device="cuda", half=True)
+    image_up = scaler_beta.upscale(image, upscaler_size, True)
 
     image_path = save_pil_image_with_metadata(image_up, f'{os.getcwd()}/up_images', exif_image)
 
@@ -875,9 +889,8 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                     actual_task_info = gr.HTML()
 
                     with gr.Accordion("History", open=False):
-                        history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", format="png", interactive=False, show_share_button=False,
-                        show_download_button=True)
                         history_files = gr.Files(interactive=False, visible=False)
+                        history_gallery = gr.Gallery(label="History", columns=6, object_fit="contain", format="png", interactive=False, show_share_button=False, show_download_button=True)
                         history_clear_button = gr.Button(value="Clear History", variant="secondary")
                         history_clear_button.click(lambda: ([], []), None, [history_gallery, history_files], queue=False, show_api=False)
 
@@ -1106,8 +1119,8 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                             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")
-                                esrgan_tile_overlap_gui = gr.Slider(minimum=1, maximum=200, step=1, value=10, label="ESRGAN Tile Overlap")
+                                upscaler_tile_size_gui = gr.Slider(minimum=0, maximum=512, step=16, value=0, label="Upscaler Tile Size", info="0 = no tiling")
+                                upscaler_tile_overlap_gui = gr.Slider(minimum=0, maximum=48, step=1, value=8, label="Upscaler Tile Overlap")
                             with gr.Row():
                                 hires_steps_gui = gr.Slider(minimum=0, value=30, maximum=100, step=1, label="Hires Steps")
                                 hires_denoising_strength_gui = gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=0.55, label="Hires Denoising Strength")
@@ -1159,7 +1172,16 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                                     mask_dilation_b_gui = gr.Number(label="Mask dilation:", value=4, minimum=1)
                                     mask_blur_b_gui = gr.Number(label="Mask blur:", value=4, minimum=1)
                                     mask_padding_b_gui = gr.Number(label="Mask padding:", value=32, minimum=1)
-                                    
+                        
+                        with gr.Accordion("Face restoration", open=False, visible=True):
+
+                            face_rest_options = [None] + FACE_RESTORATION_MODELS
+
+                            face_restoration_model_gui = gr.Dropdown(label="Face restoration model", choices=face_rest_options, value=face_rest_options[0])
+                            with gr.Row():
+                                face_restoration_visibility_gui = gr.Slider(minimum=0., maximum=1., step=0.001, value=1., label="Visibility")
+                                face_restoration_weight_gui = gr.Slider(minimum=0., maximum=1., step=0.001, value=.5, label="Weight", info="(0 = maximum effect, 1 = minimum effect)")
+
                         with gr.Accordion("Textual inversion", open=False, visible=True) as menu_ti:
                             active_textual_inversion_gui = gr.Checkbox(value=False, label="Active Textual Inversion in prompt")
                             use_textual_inversion_gui = gr.CheckboxGroup(choices=get_embed_list(get_model_pipeline(model_name_gui.value)) if active_textual_inversion_gui.value else [], value=None, label="Use Textual Invertion in prompt")
@@ -1175,7 +1197,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                             with gr.Row():
                                 strength_gui = gr.Slider(
                                     minimum=0.01, maximum=1.0, step=0.01, value=0.55, label="Strength",
-                                    info="This option adjusts the level of changes for img2img and inpainting."
+                                    info="This option adjusts the level of changes for img2img, repaint and inpaint."
                                 )
                                 image_resolution_gui = gr.Slider(
                                     minimum=64, maximum=2048, step=64, value=1024, label="Image Resolution",
@@ -1395,8 +1417,11 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
         with gr.Tab("Upscaler"):
             with gr.Row():
                 with gr.Column():
+
+                    USCALER_TAB_KEYS = [name for name in UPSCALER_KEYS[9:]]
+
                     image_up_tab = gr.Image(label="Image", type="pil", sources=["upload"])
-                    upscaler_tab = gr.Dropdown(label="Upscaler", choices=UPSCALER_KEYS[9:], value=UPSCALER_KEYS[11])
+                    upscaler_tab = gr.Dropdown(label="Upscaler", choices=USCALER_TAB_KEYS, value=USCALER_TAB_KEYS[5])
                     upscaler_size_tab = gr.Slider(minimum=1., maximum=4., step=0.1, value=1.1, label="Upscale by")
                     generate_button_up_tab = gr.Button(value="START UPSCALE", variant="primary")
 
@@ -1404,7 +1429,7 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                     result_up_tab = gr.Image(label="Result", type="pil", interactive=False, format="png")
 
                     generate_button_up_tab.click(
-                        fn=esrgan_upscale,
+                        fn=process_upscale,
                         inputs=[image_up_tab, upscaler_tab, upscaler_size_tab],
                         outputs=[result_up_tab],
                     )
@@ -1586,8 +1611,8 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                 prompt_syntax_gui,
                 upscaler_model_path_gui,
                 upscaler_increases_size_gui,
-                esrgan_tile_gui,
-                esrgan_tile_overlap_gui,
+                upscaler_tile_size_gui,
+                upscaler_tile_overlap_gui,
                 hires_steps_gui,
                 hires_denoising_strength_gui,
                 hires_sampler_gui,
@@ -1651,6 +1676,9 @@ with gr.Blocks(theme='NoCrypt/miku@>=1.2.2', elem_id="main", fill_width=True, cs
                 mode_ip2,
                 scale_ip2,
                 pag_scale_gui,
+                face_restoration_model_gui,
+                face_restoration_visibility_gui,
+                face_restoration_weight_gui,
                 load_lora_cpu_gui,
                 verbose_info_gui,
                 gpu_duration_gui,

constants.py

@@ -4,6 +4,7 @@ from stablepy import (
     scheduler_names,
     SD15_TASKS,
     SDXL_TASKS,
+    ALL_BUILTIN_UPSCALERS,
 )
 
 # - **Download Models**
@@ -43,12 +44,19 @@ LOAD_DIFFUSERS_FORMAT_MODEL = [
     'John6666/ntr-mix-illustrious-xl-noob-xl-ntrmix35-sdxl',
     'John6666/ntr-mix-illustrious-xl-noob-xl-v777-sdxl',
     'John6666/ntr-mix-illustrious-xl-noob-xl-v777forlora-sdxl',
+    'John6666/ntr-mix-illustrious-xl-noob-xl-xi-sdxl',
+    'John6666/mistoon-anime-v10illustrious-sdxl',
+    'John6666/hassaku-xl-illustrious-v10-sdxl',
+    'John6666/hassaku-xl-illustrious-v10style-sdxl',
     'John6666/haruki-mix-illustrious-v10-sdxl',
     'John6666/noobreal-v10-sdxl',
     'John6666/complicated-noobai-merge-vprediction-sdxl',
+    'Laxhar/noobai-XL-Vpred-0.75s',
+    'Laxhar/noobai-XL-Vpred-0.75',
     'Laxhar/noobai-XL-Vpred-0.65s',
     'Laxhar/noobai-XL-Vpred-0.65',
     'Laxhar/noobai-XL-Vpred-0.6',
+    'John6666/cat-tower-noobai-xl-checkpoint-v14vpred-sdxl',
     'John6666/noobai-xl-nai-xl-vpred05version-sdxl',
     'John6666/noobai-fusion2-vpred-itercomp-v1-sdxl',
     'John6666/noobai-xl-nai-xl-vpredtestversion-sdxl',
@@ -58,6 +66,7 @@ LOAD_DIFFUSERS_FORMAT_MODEL = [
     'John6666/illustrious-pencil-xl-v200-sdxl',
     'John6666/obsession-illustriousxl-v21-sdxl',
     'John6666/obsession-illustriousxl-v30-sdxl',
+    'John6666/obsession-illustriousxl-v31-sdxl',
     'John6666/wai-nsfw-illustrious-v70-sdxl',
     'John6666/illustrious-pony-mix-v3-sdxl',
     'John6666/nova-anime-xl-illustriousv10-sdxl',
@@ -81,6 +90,7 @@ LOAD_DIFFUSERS_FORMAT_MODEL = [
     'John6666/ras-real-anime-screencap-v1-sdxl',
     'John6666/duchaiten-pony-xl-no-score-v60-sdxl',
     'John6666/mistoon-anime-ponyalpha-sdxl',
+    'John6666/mistoon-xl-copper-v20fast-sdxl',
     'John6666/ebara-mfcg-pony-mix-v12-sdxl',
     'John6666/t-ponynai3-v51-sdxl',
     'John6666/t-ponynai3-v65-sdxl',
@@ -156,6 +166,7 @@ DIRECTORY_MODELS = 'models'
 DIRECTORY_LORAS = 'loras'
 DIRECTORY_VAES = 'vaes'
 DIRECTORY_EMBEDS = 'embedings'
+DIRECTORY_UPSCALERS = 'upscalers'
 
 CACHE_HF = "/home/user/.cache/huggingface/hub/"
 STORAGE_ROOT = "/home/user/"
@@ -184,27 +195,21 @@ TASK_STABLEPY = {
     'optical pattern ControlNet': 'pattern',
     'recolor ControlNet': 'recolor',
     'tile ControlNet': 'tile',
+    'repaint ControlNet': 'repaint',
 }
 
 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",
+    **{bu: bu for bu in ALL_BUILTIN_UPSCALERS if bu not in ["HAT x4", "DAT x4", "DAT x3", "DAT x2", "SwinIR 4x"]},
+    # "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",
+    # "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",
@@ -219,6 +224,7 @@ UPSCALER_KEYS = list(UPSCALER_DICT_GUI.keys())
 DIFFUSERS_CONTROLNET_MODEL = [
     "Automatic",
 
+    "brad-twinkl/controlnet-union-sdxl-1.0-promax",
     "xinsir/controlnet-union-sdxl-1.0",
     "xinsir/anime-painter",
     "Eugeoter/noob-sdxl-controlnet-canny",
@@ -241,7 +247,6 @@ DIFFUSERS_CONTROLNET_MODEL = [
     "r3gm/controlnet-recolor-sdxl-fp16",
     "r3gm/controlnet-openpose-twins-sdxl-1.0-fp16",
     "r3gm/controlnet-qr-pattern-sdxl-fp16",
-    "brad-twinkl/controlnet-union-sdxl-1.0-promax",
     "Yakonrus/SDXL_Controlnet_Tile_Realistic_v2",
     "TheMistoAI/MistoLine",
     "briaai/BRIA-2.3-ControlNet-Recoloring",

requirements.txt

@@ -1,4 +1,4 @@
-git+https://github.com/R3gm/stablepy.git@a9fe2dc  # -b refactor_sampler_fix
+git+https://github.com/R3gm/stablepy.git@5e66972  # -b refactor_sampler_fix
 torch==2.2.0
 numpy<2
 gdown

utils.py

@@ -62,7 +62,7 @@ class ModelInformation:
         self.download_url = json_data.get("downloadUrl", "")
         self.model_url = f"https://civitai.com/models/{self.model_id}?modelVersionId={self.model_version_id}"
         self.filename_url = next(
-            (v.get("name", "") for v in json_data.get("files", []) if str(self.model_version_id) in v.get("downloadUrl", "")), ""
+            (v.get("name", "") for v in reversed(json_data.get("files", [])) if str(self.model_version_id) in v.get("downloadUrl", "")), ""
         )
         self.filename_url = self.filename_url if self.filename_url else ""
         self.description = json_data.get("description", "")
@@ -300,7 +300,7 @@ def get_model_type(repo_id: str):
     default = "SD 1.5"
     try:
         if os.path.exists(repo_id):
-            tag = checkpoint_model_type(repo_id)
+            tag, _, _ = checkpoint_model_type(repo_id)
             return DIFFUSECRAFT_CHECKPOINT_NAME[tag]
         else:
             model = api.model_info(repo_id=repo_id, timeout=5.0)