Merge remote-tracking branch 'upstream/main' into feature/add-nsfw-filter

# Conflicts:
#	modules/async_worker.py
#	modules/config.py

modules/advanced_parameters.py

@@ -2,9 +2,10 @@ disable_preview, adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adapt
     scheduler_name, generate_image_grid, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
     overwrite_vary_strength, overwrite_upscale_strength, \
     mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-    debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, inpaint_engine, \
+    debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
     refiner_swap_method, \
-    freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2 = [None] * 28
+    freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+    debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field = [None] * 32
 
 
 def set_all_advanced_parameters(*args):
@@ -12,16 +13,18 @@ def set_all_advanced_parameters(*args):
         scheduler_name, generate_image_grid, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
         overwrite_vary_strength, overwrite_upscale_strength, \
         mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-        debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, inpaint_engine, \
+        debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
         refiner_swap_method, \
-        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+        debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field
 
     disable_preview, adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg, sampler_name, \
         scheduler_name, generate_image_grid, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
         overwrite_vary_strength, overwrite_upscale_strength, \
         mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-        debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, inpaint_engine, \
+        debugging_cn_preprocessor, skipping_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
         refiner_swap_method, \
-        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2 = args
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+        debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field = args
 
     return

modules/async_worker.py

@@ -27,18 +27,18 @@ def worker():
     import modules.flags as flags
     import modules.config
     import modules.patch
-    import fcbh.model_management
-    import fooocus_extras.preprocessors as preprocessors
+    import ldm_patched.modules.model_management
+    import extras.preprocessors as preprocessors
     import modules.inpaint_worker as inpaint_worker
     import modules.constants as constants
     import modules.advanced_parameters as advanced_parameters
-    import fooocus_extras.ip_adapter as ip_adapter
-    import fooocus_extras.face_crop
+    import extras.ip_adapter as ip_adapter
+    import extras.face_crop
 
     from modules.censor import censor_batch
     from modules.sdxl_styles import apply_style, apply_wildcards, fooocus_expansion
     from modules.private_logger import log
-    from modules.expansion import safe_str
+    from extras.expansion import safe_str
     from modules.util import remove_empty_str, HWC3, resize_image, \
         get_image_shape_ceil, set_image_shape_ceil, get_shape_ceil, resample_image
     from modules.upscaler import perform_upscale
@@ -135,13 +135,14 @@ def worker():
         base_model_name = args.pop()
         refiner_model_name = args.pop()
         refiner_switch = args.pop()
-        loras = [(args.pop(), args.pop()) for _ in range(5)]
+        loras = [[str(args.pop()), float(args.pop())] for _ in range(5)]
         input_image_checkbox = args.pop()
         current_tab = args.pop()
         uov_method = args.pop()
         uov_input_image = args.pop()
         outpaint_selections = args.pop()
         inpaint_input_image = args.pop()
+        inpaint_additional_prompt = args.pop()
 
         cn_tasks = {x: [] for x in flags.ip_list}
         for _ in range(4):
@@ -182,7 +183,7 @@ def worker():
         if performance_selection == 'Extreme Speed':
             print('Enter LCM mode.')
             progressbar(async_task, 1, 'Downloading LCM components ...')
-            base_model_additional_loras += [(modules.config.downloading_sdxl_lcm_lora(), 1.0)]
+            loras += [(modules.config.downloading_sdxl_lcm_lora(), 1.0)]
 
             if refiner_model_name != 'None':
                 print(f'Refiner disabled in LCM mode.')
@@ -208,8 +209,10 @@ def worker():
         modules.patch.positive_adm_scale = advanced_parameters.adm_scaler_positive
         modules.patch.negative_adm_scale = advanced_parameters.adm_scaler_negative
         modules.patch.adm_scaler_end = advanced_parameters.adm_scaler_end
-        print(
-            f'[Parameters] ADM Scale = {modules.patch.positive_adm_scale} : {modules.patch.negative_adm_scale} : {modules.patch.adm_scaler_end}')
+        print(f'[Parameters] ADM Scale = '
+              f'{modules.patch.positive_adm_scale} : '
+              f'{modules.patch.negative_adm_scale} : '
+              f'{modules.patch.adm_scaler_end}')
 
         cfg_scale = float(guidance_scale)
         print(f'[Parameters] CFG = {cfg_scale}')
@@ -217,7 +220,6 @@ def worker():
         initial_latent = None
         denoising_strength = 1.0
         tiled = False
-        inpaint_worker.current_task = None
 
         width, height = aspect_ratios_selection.replace('×', ' ').split(' ')[:2]
         width, height = int(width), int(height)
@@ -225,9 +227,14 @@ def worker():
         skip_prompt_processing = False
         refiner_swap_method = advanced_parameters.refiner_swap_method
 
+        inpaint_worker.current_task = None
+        inpaint_parameterized = advanced_parameters.inpaint_engine != 'None'
         inpaint_image = None
         inpaint_mask = None
         inpaint_head_model_path = None
+
+        use_synthetic_refiner = False
+
         controlnet_canny_path = None
         controlnet_cpds_path = None
         clip_vision_path, ip_negative_path, ip_adapter_path, ip_adapter_face_path = None, None, None, None
@@ -274,11 +281,24 @@ def worker():
                 inpaint_image = HWC3(inpaint_image)
                 if isinstance(inpaint_image, np.ndarray) and isinstance(inpaint_mask, np.ndarray) \
                         and (np.any(inpaint_mask > 127) or len(outpaint_selections) > 0):
-                    progressbar(async_task, 1, 'Downloading inpainter ...')
-                    inpaint_head_model_path, inpaint_patch_model_path = modules.config.downloading_inpaint_models(
-                        advanced_parameters.inpaint_engine)
-                    base_model_additional_loras += [(inpaint_patch_model_path, 1.0)]
-                    print(f'[Inpaint] Current inpaint model is {inpaint_patch_model_path}')
+                    if inpaint_parameterized:
+                        progressbar(async_task, 1, 'Downloading inpainter ...')
+                        modules.config.downloading_upscale_model()
+                        inpaint_head_model_path, inpaint_patch_model_path = modules.config.downloading_inpaint_models(
+                            advanced_parameters.inpaint_engine)
+                        base_model_additional_loras += [(inpaint_patch_model_path, 1.0)]
+                        print(f'[Inpaint] Current inpaint model is {inpaint_patch_model_path}')
+                        if refiner_model_name == 'None':
+                            use_synthetic_refiner = True
+                            refiner_switch = 0.5
+                    else:
+                        inpaint_head_model_path, inpaint_patch_model_path = None, None
+                        print(f'[Inpaint] Parameterized inpaint is disabled.')
+                    if inpaint_additional_prompt != '':
+                        if prompt == '':
+                            prompt = inpaint_additional_prompt
+                        else:
+                            prompt = inpaint_additional_prompt + '\n' + prompt
                     goals.append('inpaint')
             if current_tab == 'ip' or \
                     advanced_parameters.mixing_image_prompt_and_inpaint or \
@@ -337,7 +357,8 @@ def worker():
 
             progressbar(async_task, 3, 'Loading models ...')
             pipeline.refresh_everything(refiner_model_name=refiner_model_name, base_model_name=base_model_name,
-                                        loras=loras, base_model_additional_loras=base_model_additional_loras)
+                                        loras=loras, base_model_additional_loras=base_model_additional_loras,
+                                        use_synthetic_refiner=use_synthetic_refiner)
 
             progressbar(async_task, 3, 'Processing prompts ...')
             tasks = []
@@ -380,8 +401,8 @@ def worker():
                     uc=None,
                     positive_top_k=len(positive_basic_workloads),
                     negative_top_k=len(negative_basic_workloads),
-                    log_positive_prompt='\n'.join([task_prompt] + task_extra_positive_prompts),
-                    log_negative_prompt='\n'.join([task_negative_prompt] + task_extra_negative_prompts),
+                    log_positive_prompt='; '.join([task_prompt] + task_extra_positive_prompts),
+                    log_negative_prompt='; '.join([task_negative_prompt] + task_extra_negative_prompts),
                 ))
 
             if use_expansion:
@@ -426,7 +447,15 @@ def worker():
 
             initial_pixels = core.numpy_to_pytorch(uov_input_image)
             progressbar(async_task, 13, 'VAE encoding ...')
-            initial_latent = core.encode_vae(vae=pipeline.final_vae, pixels=initial_pixels)
+
+            candidate_vae, _ = pipeline.get_candidate_vae(
+                steps=steps,
+                switch=switch,
+                denoise=denoising_strength,
+                refiner_swap_method=refiner_swap_method
+            )
+
+            initial_latent = core.encode_vae(vae=candidate_vae, pixels=initial_pixels)
             B, C, H, W = initial_latent['samples'].shape
             width = W * 8
             height = H * 8
@@ -435,10 +464,7 @@ def worker():
         if 'upscale' in goals:
             H, W, C = uov_input_image.shape
             progressbar(async_task, 13, f'Upscaling image from {str((H, W))} ...')
-
-            uov_input_image = core.numpy_to_pytorch(uov_input_image)
             uov_input_image = perform_upscale(uov_input_image)
-            uov_input_image = core.pytorch_to_numpy(uov_input_image)[0]
             print(f'Image upscaled.')
 
             if '1.5x' in uov_method:
@@ -484,14 +510,20 @@ def worker():
             initial_pixels = core.numpy_to_pytorch(uov_input_image)
             progressbar(async_task, 13, 'VAE encoding ...')
 
+            candidate_vae, _ = pipeline.get_candidate_vae(
+                steps=steps,
+                switch=switch,
+                denoise=denoising_strength,
+                refiner_swap_method=refiner_swap_method
+            )
+
             initial_latent = core.encode_vae(
-                vae=pipeline.final_vae if pipeline.final_refiner_vae is None else pipeline.final_refiner_vae,
+                vae=candidate_vae,
                 pixels=initial_pixels, tiled=True)
             B, C, H, W = initial_latent['samples'].shape
             width = W * 8
             height = H * 8
             print(f'Final resolution is {str((height, width))}.')
-            refiner_swap_method = 'upscale'
 
         if 'inpaint' in goals:
             if len(outpaint_selections) > 0:
@@ -517,13 +549,19 @@ def worker():
 
                 inpaint_image = np.ascontiguousarray(inpaint_image.copy())
                 inpaint_mask = np.ascontiguousarray(inpaint_mask.copy())
+                advanced_parameters.inpaint_strength = 1.0
+                advanced_parameters.inpaint_respective_field = 1.0
 
-            inpaint_worker.current_task = inpaint_worker.InpaintWorker(image=inpaint_image, mask=inpaint_mask,
-                                                                       is_outpaint=len(outpaint_selections) > 0)
+            denoising_strength = advanced_parameters.inpaint_strength
 
-            pipeline.final_unet.model.diffusion_model.in_inpaint = True
+            inpaint_worker.current_task = inpaint_worker.InpaintWorker(
+                image=inpaint_image,
+                mask=inpaint_mask,
+                use_fill=denoising_strength > 0.99,
+                k=advanced_parameters.inpaint_respective_field
+            )
 
-            if advanced_parameters.debugging_cn_preprocessor:
+            if advanced_parameters.debugging_inpaint_preprocessor:
                 yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(),
                              do_not_show_finished_images=True)
                 return
@@ -534,33 +572,47 @@ def worker():
             inpaint_pixel_image = core.numpy_to_pytorch(inpaint_worker.current_task.interested_image)
             inpaint_pixel_mask = core.numpy_to_pytorch(inpaint_worker.current_task.interested_mask)
 
+            candidate_vae, candidate_vae_swap = pipeline.get_candidate_vae(
+                steps=steps,
+                switch=switch,
+                denoise=denoising_strength,
+                refiner_swap_method=refiner_swap_method
+            )
+
             latent_inpaint, latent_mask = core.encode_vae_inpaint(
                 mask=inpaint_pixel_mask,
-                vae=pipeline.final_vae,
+                vae=candidate_vae,
                 pixels=inpaint_pixel_image)
 
             latent_swap = None
-            if pipeline.final_refiner_vae is not None:
-                progressbar(async_task, 13, 'VAE Inpaint SD15 encoding ...')
+            if candidate_vae_swap is not None:
+                progressbar(async_task, 13, 'VAE SD15 encoding ...')
                 latent_swap = core.encode_vae(
-                    vae=pipeline.final_refiner_vae,
+                    vae=candidate_vae_swap,
                     pixels=inpaint_pixel_fill)['samples']
 
             progressbar(async_task, 13, 'VAE encoding ...')
             latent_fill = core.encode_vae(
-                vae=pipeline.final_vae,
+                vae=candidate_vae,
                 pixels=inpaint_pixel_fill)['samples']
 
-            inpaint_worker.current_task.load_latent(latent_fill=latent_fill,
-                                                    latent_inpaint=latent_inpaint,
-                                                    latent_mask=latent_mask,
-                                                    latent_swap=latent_swap,
-                                                    inpaint_head_model_path=inpaint_head_model_path)
+            inpaint_worker.current_task.load_latent(
+                latent_fill=latent_fill, latent_mask=latent_mask, latent_swap=latent_swap)
+
+            if inpaint_parameterized:
+                pipeline.final_unet = inpaint_worker.current_task.patch(
+                    inpaint_head_model_path=inpaint_head_model_path,
+                    inpaint_latent=latent_inpaint,
+                    inpaint_latent_mask=latent_mask,
+                    model=pipeline.final_unet
+                )
+
+            if not advanced_parameters.inpaint_disable_initial_latent:
+                initial_latent = {'samples': latent_fill}
 
             B, C, H, W = latent_fill.shape
             height, width = H * 8, W * 8
             final_height, final_width = inpaint_worker.current_task.image.shape[:2]
-            initial_latent = {'samples': latent_fill}
             print(f'Final resolution is {str((final_height, final_width))}, latent is {str((height, width))}.')
 
         if 'cn' in goals:
@@ -604,7 +656,7 @@ def worker():
                 cn_img = HWC3(cn_img)
 
                 if not advanced_parameters.skipping_cn_preprocessor:
-                    cn_img = fooocus_extras.face_crop.crop_image(cn_img)
+                    cn_img = extras.face_crop.crop_image(cn_img)
 
                 # https://github.com/tencent-ailab/IP-Adapter/blob/d580c50a291566bbf9fc7ac0f760506607297e6d/README.md?plain=1#L75
                 cn_img = resize_image(cn_img, width=224, height=224, resize_mode=0)
@@ -631,6 +683,15 @@ def worker():
 
         all_steps = steps * image_number
 
+        print(f'[Parameters] Denoising Strength = {denoising_strength}')
+
+        if isinstance(initial_latent, dict) and 'samples' in initial_latent:
+            log_shape = initial_latent['samples'].shape
+        else:
+            log_shape = f'Image Space {(height, width)}'
+
+        print(f'[Parameters] Initial Latent shape: {log_shape}')
+
         preparation_time = time.perf_counter() - execution_start_time
         print(f'Preparation time: {preparation_time:.2f} seconds')
 
@@ -726,7 +787,7 @@ def worker():
                     log(x, d, single_line_number=3)
 
                 yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1, progressbar_index=int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)))
-            except fcbh.model_management.InterruptProcessingException as e:
+            except ldm_patched.modules.model_management.InterruptProcessingException as e:
                 if shared.last_stop == 'skip':
                     print('User skipped')
                     continue

modules/config.py

@@ -303,6 +303,16 @@ default_overwrite_switch = get_config_item_or_set_default(
     default_value=-1,
     validator=lambda x: isinstance(x, int)
 )
+example_inpaint_prompts = get_config_item_or_set_default(
+    key='example_inpaint_prompts',
+    default_value=[
+        'highly detailed face', 'detailed girl face', 'detailed man face', 'detailed hand', 'beautiful eyes'
+    ],
+    validator=lambda x: isinstance(x, list) and all(isinstance(v, str) for v in x)
+)
+
+example_inpaint_prompts = [[x] for x in example_inpaint_prompts]
+
 default_black_out_nsfw = get_config_item_or_set_default(
     key='default_black_out_nsfw',
     default_value=False,
@@ -430,7 +440,7 @@ def downloading_sdxl_lcm_lora():
         model_dir=path_loras,
         file_name='sdxl_lcm_lora.safetensors'
     )
-    return os.path.join(path_loras, 'sdxl_lcm_lora.safetensors')
+    return 'sdxl_lcm_lora.safetensors'
 
 
 def downloading_controlnet_canny():

modules/core.py

@@ -8,26 +8,25 @@ import einops
 import torch
 import numpy as np
 
-import fcbh.model_management
-import fcbh.model_detection
-import fcbh.model_patcher
-import fcbh.utils
-import fcbh.controlnet
+import ldm_patched.modules.model_management
+import ldm_patched.modules.model_detection
+import ldm_patched.modules.model_patcher
+import ldm_patched.modules.utils
+import ldm_patched.modules.controlnet
 import modules.sample_hijack
-import fcbh.samplers
-import fcbh.latent_formats
+import ldm_patched.modules.samplers
+import ldm_patched.modules.latent_formats
 import modules.advanced_parameters
 
-from fcbh.sd import load_checkpoint_guess_config
-from nodes import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, \
+from ldm_patched.modules.sd import load_checkpoint_guess_config
+from ldm_patched.contrib.external import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, \
     ControlNetApplyAdvanced
-from fcbh_extras.nodes_freelunch import FreeU_V2
-from fcbh.sample import prepare_mask
-from modules.patch import patched_sampler_cfg_function
-from fcbh.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
+from ldm_patched.contrib.external_freelunch import FreeU_V2
+from ldm_patched.modules.sample import prepare_mask
+from modules.lora import match_lora
+from ldm_patched.modules.lora import model_lora_keys_unet, model_lora_keys_clip
 from modules.config import path_embeddings
-from modules.lora import load_dangerous_lora
-from fcbh_extras.nodes_model_advanced import ModelSamplingDiscrete
+from ldm_patched.contrib.external_model_advanced import ModelSamplingDiscrete
 
 
 opEmptyLatentImage = EmptyLatentImage()
@@ -50,13 +49,17 @@ class StableDiffusionModel:
         self.unet_with_lora = unet
         self.clip_with_lora = clip
         self.visited_loras = ''
-        self.lora_key_map = {}
 
-        if self.unet is not None and self.clip is not None:
-            self.lora_key_map = model_lora_keys_unet(self.unet.model, self.lora_key_map)
-            self.lora_key_map = model_lora_keys_clip(self.clip.cond_stage_model, self.lora_key_map)
-            self.lora_key_map.update({x: x for x in self.unet.model.state_dict().keys()})
-            self.lora_key_map.update({x: x for x in self.clip.cond_stage_model.state_dict().keys()})
+        self.lora_key_map_unet = {}
+        self.lora_key_map_clip = {}
+
+        if self.unet is not None:
+            self.lora_key_map_unet = model_lora_keys_unet(self.unet.model, self.lora_key_map_unet)
+            self.lora_key_map_unet.update({x: x for x in self.unet.model.state_dict().keys()})
+
+        if self.clip is not None:
+            self.lora_key_map_clip = model_lora_keys_clip(self.clip.cond_stage_model, self.lora_key_map_clip)
+            self.lora_key_map_clip.update({x: x for x in self.clip.cond_stage_model.state_dict().keys()})
 
     @torch.no_grad()
     @torch.inference_mode()
@@ -67,13 +70,14 @@ class StableDiffusionModel:
             return
 
         self.visited_loras = str(loras)
-        loras_to_load = []
 
         if self.unet is None:
             return
 
         print(f'Request to load LoRAs {str(loras)} for model [{self.filename}].')
 
+        loras_to_load = []
+
         for name, weight in loras:
             if name == 'None':
                 continue
@@ -93,27 +97,33 @@ class StableDiffusionModel:
         self.clip_with_lora = self.clip.clone() if self.clip is not None else None
 
         for lora_filename, weight in loras_to_load:
-            lora = fcbh.utils.load_torch_file(lora_filename, safe_load=False)
-            lora_items = load_dangerous_lora(lora, self.lora_key_map)
+            lora_unmatch = ldm_patched.modules.utils.load_torch_file(lora_filename, safe_load=False)
+            lora_unet, lora_unmatch = match_lora(lora_unmatch, self.lora_key_map_unet)
+            lora_clip, lora_unmatch = match_lora(lora_unmatch, self.lora_key_map_clip)
 
-            if len(lora_items) == 0:
+            if len(lora_unmatch) > 12:
+                # model mismatch
                 continue
-            
-            print(f'Loaded LoRA [{lora_filename}] for model [{self.filename}] with {len(lora_items)} keys at weight {weight}.')
 
-            if self.unet_with_lora is not None:
-                loaded_unet_keys = self.unet_with_lora.add_patches(lora_items, weight)
-            else:
-                loaded_unet_keys = []
+            if len(lora_unmatch) > 0:
+                print(f'Loaded LoRA [{lora_filename}] for model [{self.filename}] '
+                      f'with unmatched keys {list(lora_unmatch.keys())}')
 
-            if self.clip_with_lora is not None:
-                loaded_clip_keys = self.clip_with_lora.add_patches(lora_items, weight)
-            else:
-                loaded_clip_keys = []
+            if self.unet_with_lora is not None and len(lora_unet) > 0:
+                loaded_keys = self.unet_with_lora.add_patches(lora_unet, weight)
+                print(f'Loaded LoRA [{lora_filename}] for UNet [{self.filename}] '
+                      f'with {len(loaded_keys)} keys at weight {weight}.')
+                for item in lora_unet:
+                    if item not in loaded_keys:
+                        print("UNet LoRA key skipped: ", item)
 
-            for item in lora_items:
-                if item not in set(list(loaded_unet_keys) + list(loaded_clip_keys)):
-                    print("LoRA key skipped: ", item)
+            if self.clip_with_lora is not None and len(lora_clip) > 0:
+                loaded_keys = self.clip_with_lora.add_patches(lora_clip, weight)
+                print(f'Loaded LoRA [{lora_filename}] for CLIP [{self.filename}] '
+                      f'with {len(loaded_keys)} keys at weight {weight}.')
+                for item in lora_clip:
+                    if item not in loaded_keys:
+                        print("CLIP LoRA key skipped: ", item)
 
 
 @torch.no_grad()
@@ -125,7 +135,7 @@ def apply_freeu(model, b1, b2, s1, s2):
 @torch.no_grad()
 @torch.inference_mode()
 def load_controlnet(ckpt_filename):
-    return fcbh.controlnet.load_controlnet(ckpt_filename)
+    return ldm_patched.modules.controlnet.load_controlnet(ckpt_filename)
 
 
 @torch.no_grad()
@@ -139,40 +149,9 @@ def apply_controlnet(positive, negative, control_net, image, strength, start_per
 @torch.inference_mode()
 def load_model(ckpt_filename):
     unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename, embedding_directory=path_embeddings)
-    unet.model_options['sampler_cfg_function'] = patched_sampler_cfg_function
     return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, filename=ckpt_filename)
 
 
-@torch.no_grad()
-@torch.inference_mode()
-def load_sd_lora(model, lora_filename, strength_model=1.0, strength_clip=1.0):
-    if strength_model == 0 and strength_clip == 0:
-        return model
-
-    lora = fcbh.utils.load_torch_file(lora_filename, safe_load=False)
-
-    if lora_filename.lower().endswith('.fooocus.patch'):
-        loaded = lora
-    else:
-        key_map = model_lora_keys_unet(model.unet.model)
-        key_map = model_lora_keys_clip(model.clip.cond_stage_model, key_map)
-        loaded = load_lora(lora, key_map)
-
-    new_unet = model.unet.clone()
-    loaded_unet_keys = new_unet.add_patches(loaded, strength_model)
-
-    new_clip = model.clip.clone()
-    loaded_clip_keys = new_clip.add_patches(loaded, strength_clip)
-
-    loaded_keys = set(list(loaded_unet_keys) + list(loaded_clip_keys))
-
-    for x in loaded:
-        if x not in loaded_keys:
-            print("Lora key not loaded: ", x)
-
-    return StableDiffusionModel(unet=new_unet, clip=new_clip, vae=model.vae, clip_vision=model.clip_vision)
-
-
 @torch.no_grad()
 @torch.inference_mode()
 def generate_empty_latent(width=1024, height=1024, batch_size=1):
@@ -249,7 +228,7 @@ def get_previewer(model):
     global VAE_approx_models
 
     from modules.config import path_vae_approx
-    is_sdxl = isinstance(model.model.latent_format, fcbh.latent_formats.SDXL)
+    is_sdxl = isinstance(model.model.latent_format, ldm_patched.modules.latent_formats.SDXL)
     vae_approx_filename = os.path.join(path_vae_approx, 'xlvaeapp.pth' if is_sdxl else 'vaeapp_sd15.pth')
 
     if vae_approx_filename in VAE_approx_models:
@@ -261,14 +240,14 @@ def get_previewer(model):
         del sd
         VAE_approx_model.eval()
 
-        if fcbh.model_management.should_use_fp16():
+        if ldm_patched.modules.model_management.should_use_fp16():
             VAE_approx_model.half()
             VAE_approx_model.current_type = torch.float16
         else:
             VAE_approx_model.float()
             VAE_approx_model.current_type = torch.float32
 
-        VAE_approx_model.to(fcbh.model_management.get_torch_device())
+        VAE_approx_model.to(ldm_patched.modules.model_management.get_torch_device())
         VAE_approx_models[vae_approx_filename] = VAE_approx_model
 
     @torch.no_grad()
@@ -292,7 +271,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
              previewer_start=None, previewer_end=None, sigmas=None, noise_mean=None):
 
     if sigmas is not None:
-        sigmas = sigmas.clone().to(fcbh.model_management.get_torch_device())
+        sigmas = sigmas.clone().to(ldm_patched.modules.model_management.get_torch_device())
 
     latent_image = latent["samples"]
 
@@ -300,7 +279,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
         noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
     else:
         batch_inds = latent["batch_index"] if "batch_index" in latent else None
-        noise = fcbh.sample.prepare_noise(latent_image, seed, batch_inds)
+        noise = ldm_patched.modules.sample.prepare_noise(latent_image, seed, batch_inds)
 
     if isinstance(noise_mean, torch.Tensor):
         noise = noise + noise_mean - torch.mean(noise, dim=1, keepdim=True)
@@ -318,7 +297,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
         previewer_end = steps
 
     def callback(step, x0, x, total_steps):
-        fcbh.model_management.throw_exception_if_processing_interrupted()
+        ldm_patched.modules.model_management.throw_exception_if_processing_interrupted()
         y = None
         if previewer is not None and not modules.advanced_parameters.disable_preview:
             y = previewer(x0, previewer_start + step, previewer_end)
@@ -328,14 +307,18 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
     disable_pbar = False
     modules.sample_hijack.current_refiner = refiner
     modules.sample_hijack.refiner_switch_step = refiner_switch
-    fcbh.samplers.sample = modules.sample_hijack.sample_hacked
+    ldm_patched.modules.samplers.sample = modules.sample_hijack.sample_hacked
 
     try:
-        samples = fcbh.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
-                                      denoise=denoise, disable_noise=disable_noise, start_step=start_step,
-                                      last_step=last_step,
-                                      force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback,
-                                      disable_pbar=disable_pbar, seed=seed, sigmas=sigmas)
+        samples = ldm_patched.modules.sample.sample(model,
+                                                    noise, steps, cfg, sampler_name, scheduler,
+                                                    positive, negative, latent_image,
+                                                    denoise=denoise, disable_noise=disable_noise,
+                                                    start_step=start_step,
+                                                    last_step=last_step,
+                                                    force_full_denoise=force_full_denoise, noise_mask=noise_mask,
+                                                    callback=callback,
+                                                    disable_pbar=disable_pbar, seed=seed, sigmas=sigmas)
 
         out = latent.copy()
         out["samples"] = samples

modules/default_pipeline.py

@@ -3,13 +3,13 @@ import os
 import torch
 import modules.patch
 import modules.config
-import fcbh.model_management
-import fcbh.latent_formats
+import ldm_patched.modules.model_management
+import ldm_patched.modules.latent_formats
 import modules.inpaint_worker
-import fooocus_extras.vae_interpose as vae_interpose
+import extras.vae_interpose as vae_interpose
+from extras.expansion import FooocusExpansion
 
-from fcbh.model_base import SDXL, SDXLRefiner
-from modules.expansion import FooocusExpansion
+from ldm_patched.modules.model_base import SDXL, SDXLRefiner
 from modules.sample_hijack import clip_separate
 
 
@@ -102,6 +102,26 @@ def refresh_refiner_model(name):
     return
 
 
+@torch.no_grad()
+@torch.inference_mode()
+def synthesize_refiner_model():
+    global model_base, model_refiner
+
+    print('Synthetic Refiner Activated')
+    model_refiner = core.StableDiffusionModel(
+        unet=model_base.unet,
+        vae=model_base.vae,
+        clip=model_base.clip,
+        clip_vision=model_base.clip_vision,
+        filename=model_base.filename
+    )
+    model_refiner.vae = None
+    model_refiner.clip = None
+    model_refiner.clip_vision = None
+
+    return
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def refresh_loras(loras, base_model_additional_loras=None):
@@ -188,13 +208,14 @@ def prepare_text_encoder(async_call=True):
         # TODO: make sure that this is always called in an async way so that users cannot feel it.
         pass
     assert_model_integrity()
-    fcbh.model_management.load_models_gpu([final_clip.patcher, final_expansion.patcher])
+    ldm_patched.modules.model_management.load_models_gpu([final_clip.patcher, final_expansion.patcher])
     return
 
 
 @torch.no_grad()
 @torch.inference_mode()
-def refresh_everything(refiner_model_name, base_model_name, loras, base_model_additional_loras=None):
+def refresh_everything(refiner_model_name, base_model_name, loras,
+                       base_model_additional_loras=None, use_synthetic_refiner=False):
     global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, final_expansion
 
     final_unet = None
@@ -203,8 +224,14 @@ def refresh_everything(refiner_model_name, base_model_name, loras, base_model_ad
     final_refiner_unet = None
     final_refiner_vae = None
 
-    refresh_refiner_model(refiner_model_name)
-    refresh_base_model(base_model_name)
+    if use_synthetic_refiner and refiner_model_name == 'None':
+        print('Synthetic Refiner Activated')
+        refresh_base_model(base_model_name)
+        synthesize_refiner_model()
+    else:
+        refresh_refiner_model(refiner_model_name)
+        refresh_base_model(base_model_name)
+
     refresh_loras(loras, base_model_additional_loras=base_model_additional_loras)
     assert_model_integrity()
 
@@ -212,14 +239,9 @@ def refresh_everything(refiner_model_name, base_model_name, loras, base_model_ad
     final_clip = model_base.clip_with_lora
     final_vae = model_base.vae
 
-    final_unet.model.diffusion_model.in_inpaint = False
-
     final_refiner_unet = model_refiner.unet_with_lora
     final_refiner_vae = model_refiner.vae
 
-    if final_refiner_unet is not None:
-        final_refiner_unet.model.diffusion_model.in_inpaint = False
-
     if final_expansion is None:
         final_expansion = FooocusExpansion()
 
@@ -248,7 +270,7 @@ def vae_parse(latent):
 @torch.no_grad()
 @torch.inference_mode()
 def calculate_sigmas_all(sampler, model, scheduler, steps):
-    from fcbh.samplers import calculate_sigmas_scheduler
+    from ldm_patched.modules.samplers import calculate_sigmas_scheduler
 
     discard_penultimate_sigma = False
     if sampler in ['dpm_2', 'dpm_2_ancestral']:
@@ -276,32 +298,52 @@ def calculate_sigmas(sampler, model, scheduler, steps, denoise):
 
 @torch.no_grad()
 @torch.inference_mode()
-def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0, refiner_swap_method='joint'):
-    global final_unet, final_refiner_unet, final_vae, final_refiner_vae
+def get_candidate_vae(steps, switch, denoise=1.0, refiner_swap_method='joint'):
+    assert refiner_swap_method in ['joint', 'separate', 'vae']
 
-    assert refiner_swap_method in ['joint', 'separate', 'vae', 'upscale']
-
-    refiner_use_different_vae = final_refiner_vae is not None and final_refiner_unet is not None
-
-    if refiner_swap_method == 'upscale':
-        if not refiner_use_different_vae:
-            refiner_swap_method = 'joint'
-    else:
-        if refiner_use_different_vae:
-            if denoise > 0.95:
-                refiner_swap_method = 'vae'
+    if final_refiner_vae is not None and final_refiner_unet is not None:
+        if denoise > 0.9:
+            return final_vae, final_refiner_vae
+        else:
+            if denoise > (float(steps - switch) / float(steps)) ** 0.834:  # karras 0.834
+                return final_vae, None
             else:
-                # VAE swap only support full denoise
-                # Disable refiner to avoid SD15 in joint/separate swap
-                final_refiner_unet = None
-                final_refiner_vae = None
+                return final_refiner_vae, None
+
+    return final_vae, final_refiner_vae
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0, refiner_swap_method='joint'):
+    target_unet, target_vae, target_refiner_unet, target_refiner_vae, target_clip \
+        = final_unet, final_vae, final_refiner_unet, final_refiner_vae, final_clip
+
+    assert refiner_swap_method in ['joint', 'separate', 'vae']
+
+    if final_refiner_vae is not None and final_refiner_unet is not None:
+        # Refiner Use Different VAE (then it is SD15)
+        if denoise > 0.9:
+            refiner_swap_method = 'vae'
+        else:
+            refiner_swap_method = 'joint'
+            if denoise > (float(steps - switch) / float(steps)) ** 0.834:  # karras 0.834
+                target_unet, target_vae, target_refiner_unet, target_refiner_vae \
+                    = final_unet, final_vae, None, None
+                print(f'[Sampler] only use Base because of partial denoise.')
+            else:
+                positive_cond = clip_separate(positive_cond, target_model=final_refiner_unet.model, target_clip=final_clip)
+                negative_cond = clip_separate(negative_cond, target_model=final_refiner_unet.model, target_clip=final_clip)
+                target_unet, target_vae, target_refiner_unet, target_refiner_vae \
+                    = final_refiner_unet, final_refiner_vae, None, None
+                print(f'[Sampler] only use Refiner because of partial denoise.')
 
     print(f'[Sampler] refiner_swap_method = {refiner_swap_method}')
 
     if latent is None:
-        empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
+        initial_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
     else:
-        empty_latent = latent
+        initial_latent = latent
 
     minmax_sigmas = calculate_sigmas(sampler=sampler_name, scheduler=scheduler_name, model=final_unet.model, steps=steps, denoise=denoise)
     sigma_min, sigma_max = minmax_sigmas[minmax_sigmas > 0].min(), minmax_sigmas.max()
@@ -310,18 +352,18 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
     print(f'[Sampler] sigma_min = {sigma_min}, sigma_max = {sigma_max}')
 
     modules.patch.BrownianTreeNoiseSamplerPatched.global_init(
-        empty_latent['samples'].to(fcbh.model_management.get_torch_device()),
+        initial_latent['samples'].to(ldm_patched.modules.model_management.get_torch_device()),
         sigma_min, sigma_max, seed=image_seed, cpu=False)
 
     decoded_latent = None
 
     if refiner_swap_method == 'joint':
         sampled_latent = core.ksampler(
-            model=final_unet,
-            refiner=final_refiner_unet,
+            model=target_unet,
+            refiner=target_refiner_unet,
             positive=positive_cond,
             negative=negative_cond,
-            latent=empty_latent,
+            latent=initial_latent,
             steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
             denoise=denoise,
@@ -333,32 +375,14 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             previewer_start=0,
             previewer_end=steps,
         )
-        decoded_latent = core.decode_vae(vae=final_vae, latent_image=sampled_latent, tiled=tiled)
-
-    if refiner_swap_method == 'upscale':
-        sampled_latent = core.ksampler(
-            model=final_refiner_unet,
-            positive=clip_separate(positive_cond, target_model=final_refiner_unet.model, target_clip=final_clip),
-            negative=clip_separate(negative_cond, target_model=final_refiner_unet.model, target_clip=final_clip),
-            latent=empty_latent,
-            steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
-            seed=image_seed,
-            denoise=denoise,
-            callback_function=callback,
-            cfg=cfg_scale,
-            sampler_name=sampler_name,
-            scheduler=scheduler_name,
-            previewer_start=0,
-            previewer_end=steps,
-        )
-        decoded_latent = core.decode_vae(vae=final_refiner_vae, latent_image=sampled_latent, tiled=tiled)
+        decoded_latent = core.decode_vae(vae=target_vae, latent_image=sampled_latent, tiled=tiled)
 
     if refiner_swap_method == 'separate':
         sampled_latent = core.ksampler(
-            model=final_unet,
+            model=target_unet,
             positive=positive_cond,
             negative=negative_cond,
-            latent=empty_latent,
+            latent=initial_latent,
             steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=False,
             seed=image_seed,
             denoise=denoise,
@@ -371,15 +395,15 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
         )
         print('Refiner swapped by changing ksampler. Noise preserved.')
 
-        target_model = final_refiner_unet
+        target_model = target_refiner_unet
         if target_model is None:
-            target_model = final_unet
+            target_model = target_unet
             print('Use base model to refine itself - this may because of developer mode.')
 
         sampled_latent = core.ksampler(
             model=target_model,
-            positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=final_clip),
-            negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=final_clip),
+            positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=target_clip),
+            negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=target_clip),
             latent=sampled_latent,
             steps=steps, start_step=switch, last_step=steps, disable_noise=True, force_full_denoise=True,
             seed=image_seed,
@@ -392,9 +416,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             previewer_end=steps,
         )
 
-        target_model = final_refiner_vae
+        target_model = target_refiner_vae
         if target_model is None:
-            target_model = final_vae
+            target_model = target_vae
         decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
 
     if refiner_swap_method == 'vae':
@@ -404,10 +428,10 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             modules.inpaint_worker.current_task.unswap()
 
         sampled_latent = core.ksampler(
-            model=final_unet,
+            model=target_unet,
             positive=positive_cond,
             negative=negative_cond,
-            latent=empty_latent,
+            latent=initial_latent,
             steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
             denoise=denoise,
@@ -420,9 +444,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
         )
         print('Fooocus VAE-based swap.')
 
-        target_model = final_refiner_unet
+        target_model = target_refiner_unet
         if target_model is None:
-            target_model = final_unet
+            target_model = target_unet
             print('Use base model to refine itself - this may because of developer mode.')
 
         sampled_latent = vae_parse(sampled_latent)
@@ -442,8 +466,8 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
 
         sampled_latent = core.ksampler(
             model=target_model,
-            positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=final_clip),
-            negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=final_clip),
+            positive=clip_separate(positive_cond, target_model=target_model.model, target_clip=target_clip),
+            negative=clip_separate(negative_cond, target_model=target_model.model, target_clip=target_clip),
             latent=sampled_latent,
             steps=len_sigmas, start_step=0, last_step=len_sigmas, disable_noise=False, force_full_denoise=True,
             seed=image_seed+1,
@@ -458,9 +482,9 @@ def process_diffusion(positive_cond, negative_cond, steps, switch, width, height
             noise_mean=noise_mean
         )
 
-        target_model = final_refiner_vae
+        target_model = target_refiner_vae
         if target_model is None:
-            target_model = final_vae
+            target_model = target_vae
         decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
 
     images = core.pytorch_to_numpy(decoded_latent)

modules/expansion.py

@@ -1,126 +0,0 @@
-# Fooocus GPT2 Expansion
-# Algorithm created by Lvmin Zhang at 2023, Stanford
-# If used inside Fooocus, any use is permitted.
-# If used outside Fooocus, only non-commercial use is permitted (CC-By NC 4.0).
-# This applies to the word list, vocab, model, and algorithm.
-
-
-import os
-import torch
-import math
-import fcbh.model_management as model_management
-
-from transformers.generation.logits_process import LogitsProcessorList
-from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
-from modules.config import path_fooocus_expansion
-from fcbh.model_patcher import ModelPatcher
-
-
-# limitation of np.random.seed(), called from transformers.set_seed()
-SEED_LIMIT_NUMPY = 2**32
-neg_inf = - 8192.0
-
-
-def safe_str(x):
-    x = str(x)
-    for _ in range(16):
-        x = x.replace('  ', ' ')
-    return x.strip(",. \r\n")
-
-
-def remove_pattern(x, pattern):
-    for p in pattern:
-        x = x.replace(p, '')
-    return x
-
-
-class FooocusExpansion:
-    def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained(path_fooocus_expansion)
-
-        positive_words = open(os.path.join(path_fooocus_expansion, 'positive.txt'),
-                              encoding='utf-8').read().splitlines()
-        positive_words = ['Ġ' + x.lower() for x in positive_words if x != '']
-
-        self.logits_bias = torch.zeros((1, len(self.tokenizer.vocab)), dtype=torch.float32) + neg_inf
-
-        debug_list = []
-        for k, v in self.tokenizer.vocab.items():
-            if k in positive_words:
-                self.logits_bias[0, v] = 0
-                debug_list.append(k[1:])
-
-        print(f'Fooocus V2 Expansion: Vocab with {len(debug_list)} words.')
-
-        # debug_list = '\n'.join(sorted(debug_list))
-        # print(debug_list)
-
-        # t11 = self.tokenizer(',', return_tensors="np")
-        # t198 = self.tokenizer('\n', return_tensors="np")
-        # eos = self.tokenizer.eos_token_id
-
-        self.model = AutoModelForCausalLM.from_pretrained(path_fooocus_expansion)
-        self.model.eval()
-
-        load_device = model_management.text_encoder_device()
-        offload_device = model_management.text_encoder_offload_device()
-
-        # MPS hack
-        if model_management.is_device_mps(load_device):
-            load_device = torch.device('cpu')
-            offload_device = torch.device('cpu')
-
-        use_fp16 = model_management.should_use_fp16(device=load_device)
-
-        if use_fp16:
-            self.model.half()
-
-        self.patcher = ModelPatcher(self.model, load_device=load_device, offload_device=offload_device)
-        print(f'Fooocus Expansion engine loaded for {load_device}, use_fp16 = {use_fp16}.')
-
-    @torch.no_grad()
-    @torch.inference_mode()
-    def logits_processor(self, input_ids, scores):
-        assert scores.ndim == 2 and scores.shape[0] == 1
-        self.logits_bias = self.logits_bias.to(scores)
-
-        bias = self.logits_bias.clone()
-        bias[0, input_ids[0].to(bias.device).long()] = neg_inf
-        bias[0, 11] = 0
-
-        return scores + bias
-
-    @torch.no_grad()
-    @torch.inference_mode()
-    def __call__(self, prompt, seed):
-        if prompt == '':
-            return ''
-
-        if self.patcher.current_device != self.patcher.load_device:
-            print('Fooocus Expansion loaded by itself.')
-            model_management.load_model_gpu(self.patcher)
-
-        seed = int(seed) % SEED_LIMIT_NUMPY
-        set_seed(seed)
-        prompt = safe_str(prompt) + ','
-
-        tokenized_kwargs = self.tokenizer(prompt, return_tensors="pt")
-        tokenized_kwargs.data['input_ids'] = tokenized_kwargs.data['input_ids'].to(self.patcher.load_device)
-        tokenized_kwargs.data['attention_mask'] = tokenized_kwargs.data['attention_mask'].to(self.patcher.load_device)
-
-        current_token_length = int(tokenized_kwargs.data['input_ids'].shape[1])
-        max_token_length = 75 * int(math.ceil(float(current_token_length) / 75.0))
-        max_new_tokens = max_token_length - current_token_length
-
-        # https://huggingface.co/blog/introducing-csearch
-        # https://huggingface.co/docs/transformers/generation_strategies
-        features = self.model.generate(**tokenized_kwargs,
-                                       top_k=100,
-                                       max_new_tokens=max_new_tokens,
-                                       do_sample=True,
-                                       logits_processor=LogitsProcessorList([self.logits_processor]))
-
-        response = self.tokenizer.batch_decode(features, skip_special_tokens=True)
-        result = safe_str(response[0])
-
-        return result

modules/flags.py

@@ -14,7 +14,7 @@ KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_
                   "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
                   "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
 
-SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo"]
 SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
 
 sampler_list = SAMPLER_NAMES
@@ -32,5 +32,13 @@ default_parameters = {
     cn_ip: (0.5, 0.6), cn_ip_face: (0.9, 0.75), cn_canny: (0.5, 1.0), cn_cpds: (0.5, 1.0)
 }  # stop, weight
 
-inpaint_engine_versions = ['v1', 'v2.5', 'v2.6']
+inpaint_engine_versions = ['None', 'v1', 'v2.5', 'v2.6']
 performance_selections = ['Speed', 'Quality', 'Extreme Speed']
+
+inpaint_option_default = 'Inpaint or Outpaint (default)'
+inpaint_option_detail = 'Improve Detail (face, hand, eyes, etc.)'
+inpaint_option_modify = 'Modify Content (add objects, change background, etc.)'
+inpaint_options = [inpaint_option_default, inpaint_option_detail, inpaint_option_modify]
+
+desc_type_photo = 'Photograph'
+desc_type_anime = 'Art/Anime'

modules/inpaint_worker.py

@@ -1,12 +1,12 @@
 import torch
 import numpy as np
-import modules.default_pipeline as pipeline
 
 from PIL import Image, ImageFilter
-from modules.util import resample_image, set_image_shape_ceil
+from modules.util import resample_image, set_image_shape_ceil, get_image_shape_ceil
+from modules.upscaler import perform_upscale
 
 
-inpaint_head = None
+inpaint_head_model = None
 
 
 class InpaintHead(torch.nn.Module):
@@ -77,29 +77,32 @@ def regulate_abcd(x, a, b, c, d):
 
 def compute_initial_abcd(x):
     indices = np.where(x)
-    a = np.min(indices[0]) - 64
-    b = np.max(indices[0]) + 65
-    c = np.min(indices[1]) - 64
-    d = np.max(indices[1]) + 65
+    a = np.min(indices[0])
+    b = np.max(indices[0])
+    c = np.min(indices[1])
+    d = np.max(indices[1])
     abp = (b + a) // 2
     abm = (b - a) // 2
     cdp = (d + c) // 2
     cdm = (d - c) // 2
-    l = max(abm, cdm)
+    l = int(max(abm, cdm) * 1.15)
     a = abp - l
-    b = abp + l
+    b = abp + l + 1
     c = cdp - l
-    d = cdp + l
+    d = cdp + l + 1
     a, b, c, d = regulate_abcd(x, a, b, c, d)
     return a, b, c, d
 
 
-def solve_abcd(x, a, b, c, d, outpaint):
+def solve_abcd(x, a, b, c, d, k):
+    k = float(k)
+    assert 0.0 <= k <= 1.0
+
     H, W = x.shape[:2]
-    if outpaint:
+    if k == 1.0:
         return 0, H, 0, W
     while True:
-        if b - a > H * 0.618 and d - c > W * 0.618:
+        if b - a >= H * k and d - c >= W * k:
             break
 
         add_h = (b - a) < (d - c)
@@ -138,21 +141,30 @@ def fooocus_fill(image, mask):
 
 
 class InpaintWorker:
-    def __init__(self, image, mask, is_outpaint):
+    def __init__(self, image, mask, use_fill=True, k=0.618):
         a, b, c, d = compute_initial_abcd(mask > 0)
-        a, b, c, d = solve_abcd(mask, a, b, c, d, outpaint=is_outpaint)
+        a, b, c, d = solve_abcd(mask, a, b, c, d, k=k)
 
         # interested area
         self.interested_area = (a, b, c, d)
         self.interested_mask = mask[a:b, c:d]
         self.interested_image = image[a:b, c:d]
 
+        # super resolution
+        if get_image_shape_ceil(self.interested_image) < 1024:
+            self.interested_image = perform_upscale(self.interested_image)
+
         # resize to make images ready for diffusion
         self.interested_image = set_image_shape_ceil(self.interested_image, 1024)
+        self.interested_fill = self.interested_image.copy()
         H, W, C = self.interested_image.shape
 
+        # process mask
         self.interested_mask = up255(resample_image(self.interested_mask, W, H), t=127)
-        self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask)
+
+        # compute filling
+        if use_fill:
+            self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask)
 
         # soft pixels
         self.mask = morphological_open(mask)
@@ -166,34 +178,36 @@ class InpaintWorker:
         self.inpaint_head_feature = None
         return
 
-    def load_latent(self,
-                    latent_fill,
-                    latent_inpaint,
-                    latent_mask,
-                    latent_swap=None,
-                    inpaint_head_model_path=None):
-
-        global inpaint_head
-        assert inpaint_head_model_path is not None
-
+    def load_latent(self, latent_fill, latent_mask, latent_swap=None):
         self.latent = latent_fill
         self.latent_mask = latent_mask
         self.latent_after_swap = latent_swap
+        return
 
-        if inpaint_head is None:
-            inpaint_head = InpaintHead()
+    def patch(self, inpaint_head_model_path, inpaint_latent, inpaint_latent_mask, model):
+        global inpaint_head_model
+
+        if inpaint_head_model is None:
+            inpaint_head_model = InpaintHead()
             sd = torch.load(inpaint_head_model_path, map_location='cpu')
-            inpaint_head.load_state_dict(sd)
+            inpaint_head_model.load_state_dict(sd)
 
         feed = torch.cat([
-            latent_mask,
-            pipeline.final_unet.model.process_latent_in(latent_inpaint)
+            inpaint_latent_mask,
+            model.model.process_latent_in(inpaint_latent)
         ], dim=1)
 
-        inpaint_head.to(device=feed.device, dtype=feed.dtype)
-        self.inpaint_head_feature = inpaint_head(feed)
+        inpaint_head_model.to(device=feed.device, dtype=feed.dtype)
+        inpaint_head_feature = inpaint_head_model(feed)
 
-        return
+        def input_block_patch(h, transformer_options):
+            if transformer_options["block"][1] == 0:
+                h = h + inpaint_head_feature.to(h)
+            return h
+
+        m = model.clone()
+        m.set_model_input_block_patch(input_block_patch)
+        return m
 
     def swap(self):
         if self.swapped:
@@ -239,5 +253,5 @@ class InpaintWorker:
         return result
 
     def visualize_mask_processing(self):
-        return [self.interested_fill, self.interested_mask, self.image, self.mask]
+        return [self.interested_fill, self.interested_mask, self.interested_image]
 

modules/lora.py

@@ -1,10 +1,10 @@
-def load_dangerous_lora(lora, to_load):
+def match_lora(lora, to_load):
     patch_dict = {}
     loaded_keys = set()
     for x in to_load:
         real_load_key = to_load[x]
         if real_load_key in lora:
-            patch_dict[real_load_key] = lora[real_load_key]
+            patch_dict[real_load_key] = ('fooocus', lora[real_load_key])
             loaded_keys.add(real_load_key)
             continue
 
@@ -37,7 +37,7 @@ def load_dangerous_lora(lora, to_load):
             if mid_name is not None and mid_name in lora.keys():
                 mid = lora[mid_name]
                 loaded_keys.add(mid_name)
-            patch_dict[to_load[x]] = (lora[A_name], lora[B_name], alpha, mid)
+            patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid))
             loaded_keys.add(A_name)
             loaded_keys.add(B_name)
 
@@ -58,7 +58,7 @@ def load_dangerous_lora(lora, to_load):
                 loaded_keys.add(hada_t1_name)
                 loaded_keys.add(hada_t2_name)
 
-            patch_dict[to_load[x]] = (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2)
+            patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2))
             loaded_keys.add(hada_w1_a_name)
             loaded_keys.add(hada_w1_b_name)
             loaded_keys.add(hada_w2_a_name)
@@ -110,7 +110,19 @@ def load_dangerous_lora(lora, to_load):
             loaded_keys.add(lokr_t2_name)
 
         if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
-            patch_dict[to_load[x]] = (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2)
+            patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2))
+
+        #glora
+        a1_name = "{}.a1.weight".format(x)
+        a2_name = "{}.a2.weight".format(x)
+        b1_name = "{}.b1.weight".format(x)
+        b2_name = "{}.b2.weight".format(x)
+        if a1_name in lora:
+            patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha))
+            loaded_keys.add(a1_name)
+            loaded_keys.add(a2_name)
+            loaded_keys.add(b1_name)
+            loaded_keys.add(b2_name)
 
         w_norm_name = "{}.w_norm".format(x)
         b_norm_name = "{}.b_norm".format(x)
@@ -119,30 +131,22 @@ def load_dangerous_lora(lora, to_load):
 
         if w_norm is not None:
             loaded_keys.add(w_norm_name)
-            patch_dict[to_load[x]] = (w_norm,)
+            patch_dict[to_load[x]] = ("diff", (w_norm,))
             if b_norm is not None:
                 loaded_keys.add(b_norm_name)
-                patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = (b_norm,)
+                patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (b_norm,))
 
         diff_name = "{}.diff".format(x)
         diff_weight = lora.get(diff_name, None)
         if diff_weight is not None:
-            patch_dict[to_load[x]] = (diff_weight,)
+            patch_dict[to_load[x]] = ("diff", (diff_weight,))
             loaded_keys.add(diff_name)
 
         diff_bias_name = "{}.diff_b".format(x)
         diff_bias = lora.get(diff_bias_name, None)
         if diff_bias is not None:
-            patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = (diff_bias,)
+            patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,))
             loaded_keys.add(diff_bias_name)
 
-    remaining_keys = [x for x in lora.keys() if x not in loaded_keys]
-
-    if len(remaining_keys) == 0:
-        return patch_dict
-
-    if len(remaining_keys) > 12:
-        return {}
-
-    print(f'LoRA loaded with extra keys: {remaining_keys}')
-    return patch_dict
+    remaining_dict = {x: y for x, y in lora.items() if x not in loaded_keys}
+    return patch_dict, remaining_dict

modules/patch.py

@@ -1,33 +1,30 @@
 import os
 import torch
-import math
 import time
-import numpy as np
-import fcbh.model_base
-import fcbh.ldm.modules.diffusionmodules.openaimodel
-import fcbh.samplers
-import fcbh.model_management
+import math
+import ldm_patched.modules.model_base
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.model_management
 import modules.anisotropic as anisotropic
-import fcbh.ldm.modules.attention
-import fcbh.k_diffusion.sampling
-import fcbh.sd1_clip
+import ldm_patched.ldm.modules.attention
+import ldm_patched.k_diffusion.sampling
+import ldm_patched.modules.sd1_clip
 import modules.inpaint_worker as inpaint_worker
-import fcbh.ldm.modules.diffusionmodules.openaimodel
-import fcbh.ldm.modules.diffusionmodules.model
-import fcbh.sd
-import fcbh.cldm.cldm
-import fcbh.model_patcher
-import fcbh.samplers
-import fcbh.cli_args
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.ldm.modules.diffusionmodules.model
+import ldm_patched.modules.sd
+import ldm_patched.controlnet.cldm
+import ldm_patched.modules.model_patcher
+import ldm_patched.modules.samplers
+import ldm_patched.modules.args_parser
 import modules.advanced_parameters as advanced_parameters
 import warnings
 import safetensors.torch
 import modules.constants as constants
 
-from einops import repeat
-from fcbh.k_diffusion.sampling import BatchedBrownianTree
-from fcbh.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, apply_control
-from fcbh.ldm.modules.diffusionmodules.util import make_beta_schedule
+from ldm_patched.modules.samplers import calc_cond_uncond_batch
+from ldm_patched.k_diffusion.sampling import BatchedBrownianTree
+from ldm_patched.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, apply_control
 
 
 sharpness = 2.0
@@ -54,31 +51,25 @@ def calculate_weight_patched(self, patches, weight, key):
             v = (self.calculate_weight(v[1:], v[0].clone(), key),)
 
         if len(v) == 1:
+            patch_type = "diff"
+        elif len(v) == 2:
+            patch_type = v[0]
+            v = v[1]
+
+        if patch_type == "diff":
             w1 = v[0]
             if alpha != 0.0:
                 if w1.shape != weight.shape:
                     print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
                 else:
-                    weight += alpha * fcbh.model_management.cast_to_device(w1, weight.device, weight.dtype)
-        elif len(v) == 3:
-            # fooocus
-            w1 = fcbh.model_management.cast_to_device(v[0], weight.device, torch.float32)
-            w_min = fcbh.model_management.cast_to_device(v[1], weight.device, torch.float32)
-            w_max = fcbh.model_management.cast_to_device(v[2], weight.device, torch.float32)
-            w1 = (w1 / 255.0) * (w_max - w_min) + w_min
-            if alpha != 0.0:
-                if w1.shape != weight.shape:
-                    print("WARNING SHAPE MISMATCH {} FOOOCUS WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
-                else:
-                    weight += alpha * fcbh.model_management.cast_to_device(w1, weight.device, weight.dtype)
-        elif len(v) == 4:  # lora/locon
-            mat1 = fcbh.model_management.cast_to_device(v[0], weight.device, torch.float32)
-            mat2 = fcbh.model_management.cast_to_device(v[1], weight.device, torch.float32)
+                    weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
+        elif patch_type == "lora":
+            mat1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
+            mat2 = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
             if v[2] is not None:
                 alpha *= v[2] / mat2.shape[0]
             if v[3] is not None:
-                # locon mid weights, hopefully the math is fine because I didn't properly test it
-                mat3 = fcbh.model_management.cast_to_device(v[3], weight.device, torch.float32)
+                mat3 = ldm_patched.modules.model_management.cast_to_device(v[3], weight.device, torch.float32)
                 final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
                 mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1),
                                 mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
@@ -87,7 +78,17 @@ def calculate_weight_patched(self, patches, weight, key):
                     weight.shape).type(weight.dtype)
             except Exception as e:
                 print("ERROR", key, e)
-        elif len(v) == 8:  # lokr
+        elif patch_type == "fooocus":
+            w1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
+            w_min = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
+            w_max = ldm_patched.modules.model_management.cast_to_device(v[2], weight.device, torch.float32)
+            w1 = (w1 / 255.0) * (w_max - w_min) + w_min
+            if alpha != 0.0:
+                if w1.shape != weight.shape:
+                    print("WARNING SHAPE MISMATCH {} FOOOCUS WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+                else:
+                    weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
+        elif patch_type == "lokr":
             w1 = v[0]
             w2 = v[1]
             w1_a = v[3]
@@ -99,23 +100,23 @@ def calculate_weight_patched(self, patches, weight, key):
 
             if w1 is None:
                 dim = w1_b.shape[0]
-                w1 = torch.mm(fcbh.model_management.cast_to_device(w1_a, weight.device, torch.float32),
-                              fcbh.model_management.cast_to_device(w1_b, weight.device, torch.float32))
+                w1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1_a, weight.device, torch.float32),
+                              ldm_patched.modules.model_management.cast_to_device(w1_b, weight.device, torch.float32))
             else:
-                w1 = fcbh.model_management.cast_to_device(w1, weight.device, torch.float32)
+                w1 = ldm_patched.modules.model_management.cast_to_device(w1, weight.device, torch.float32)
 
             if w2 is None:
                 dim = w2_b.shape[0]
                 if t2 is None:
-                    w2 = torch.mm(fcbh.model_management.cast_to_device(w2_a, weight.device, torch.float32),
-                                  fcbh.model_management.cast_to_device(w2_b, weight.device, torch.float32))
+                    w2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32),
+                                  ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32))
                 else:
                     w2 = torch.einsum('i j k l, j r, i p -> p r k l',
-                                      fcbh.model_management.cast_to_device(t2, weight.device, torch.float32),
-                                      fcbh.model_management.cast_to_device(w2_b, weight.device, torch.float32),
-                                      fcbh.model_management.cast_to_device(w2_a, weight.device, torch.float32))
+                                      ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
+                                      ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32),
+                                      ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32))
             else:
-                w2 = fcbh.model_management.cast_to_device(w2, weight.device, torch.float32)
+                w2 = ldm_patched.modules.model_management.cast_to_device(w2, weight.device, torch.float32)
 
             if len(w2.shape) == 4:
                 w1 = w1.unsqueeze(2).unsqueeze(2)
@@ -126,7 +127,7 @@ def calculate_weight_patched(self, patches, weight, key):
                 weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
             except Exception as e:
                 print("ERROR", key, e)
-        else:  # loha
+        elif patch_type == "loha":
             w1a = v[0]
             w1b = v[1]
             if v[2] is not None:
@@ -137,24 +138,36 @@ def calculate_weight_patched(self, patches, weight, key):
                 t1 = v[5]
                 t2 = v[6]
                 m1 = torch.einsum('i j k l, j r, i p -> p r k l',
-                                  fcbh.model_management.cast_to_device(t1, weight.device, torch.float32),
-                                  fcbh.model_management.cast_to_device(w1b, weight.device, torch.float32),
-                                  fcbh.model_management.cast_to_device(w1a, weight.device, torch.float32))
+                                  ldm_patched.modules.model_management.cast_to_device(t1, weight.device, torch.float32),
+                                  ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32),
+                                  ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32))
 
                 m2 = torch.einsum('i j k l, j r, i p -> p r k l',
-                                  fcbh.model_management.cast_to_device(t2, weight.device, torch.float32),
-                                  fcbh.model_management.cast_to_device(w2b, weight.device, torch.float32),
-                                  fcbh.model_management.cast_to_device(w2a, weight.device, torch.float32))
+                                  ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
+                                  ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32),
+                                  ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32))
             else:
-                m1 = torch.mm(fcbh.model_management.cast_to_device(w1a, weight.device, torch.float32),
-                              fcbh.model_management.cast_to_device(w1b, weight.device, torch.float32))
-                m2 = torch.mm(fcbh.model_management.cast_to_device(w2a, weight.device, torch.float32),
-                              fcbh.model_management.cast_to_device(w2b, weight.device, torch.float32))
+                m1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32),
+                              ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32))
+                m2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32),
+                              ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32))
 
             try:
                 weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
             except Exception as e:
                 print("ERROR", key, e)
+        elif patch_type == "glora":
+            if v[4] is not None:
+                alpha *= v[4] / v[0].shape[0]
+
+            a1 = ldm_patched.modules.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
+            a2 = ldm_patched.modules.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
+            b1 = ldm_patched.modules.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
+            b2 = ldm_patched.modules.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)
+
+            weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype)
+        else:
+            print("patch type not recognized", patch_type, key)
 
     return weight
 
@@ -162,19 +175,17 @@ def calculate_weight_patched(self, patches, weight, key):
 class BrownianTreeNoiseSamplerPatched:
     transform = None
     tree = None
-    global_sigma_min = 1.0
-    global_sigma_max = 1.0
 
     @staticmethod
     def global_init(x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
+        if ldm_patched.modules.model_management.directml_enabled:
+            cpu = True
+
         t0, t1 = transform(torch.as_tensor(sigma_min)), transform(torch.as_tensor(sigma_max))
 
         BrownianTreeNoiseSamplerPatched.transform = transform
         BrownianTreeNoiseSamplerPatched.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
 
-        BrownianTreeNoiseSamplerPatched.global_sigma_min = sigma_min
-        BrownianTreeNoiseSamplerPatched.global_sigma_max = sigma_max
-
     def __init__(self, *args, **kwargs):
         pass
 
@@ -202,34 +213,47 @@ def compute_cfg(uncond, cond, cfg_scale, t):
         return real_eps
 
 
-def patched_sampler_cfg_function(args):
+def patched_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options=None, seed=None):
+    if math.isclose(cond_scale, 1.0):
+        return calc_cond_uncond_batch(model, cond, None, x, timestep, model_options)[0]
+
     global eps_record
 
-    positive_eps = args['cond']
-    negative_eps = args['uncond']
-    cfg_scale = args['cond_scale']
-    positive_x0 = args['input'] - positive_eps
-    sigma = args['sigma']
+    positive_x0, negative_x0 = calc_cond_uncond_batch(model, cond, uncond, x, timestep, model_options)
+
+    positive_eps = x - positive_x0
+    negative_eps = x - negative_x0
+    sigma = timestep
 
     alpha = 0.001 * sharpness * global_diffusion_progress
+
     positive_eps_degraded = anisotropic.adaptive_anisotropic_filter(x=positive_eps, g=positive_x0)
     positive_eps_degraded_weighted = positive_eps_degraded * alpha + positive_eps * (1.0 - alpha)
 
     final_eps = compute_cfg(uncond=negative_eps, cond=positive_eps_degraded_weighted,
-                            cfg_scale=cfg_scale, t=global_diffusion_progress)
+                            cfg_scale=cond_scale, t=global_diffusion_progress)
 
     if eps_record is not None:
         eps_record = (final_eps / sigma).cpu()
 
-    return final_eps
+    return x - final_eps
+
+
+def round_to_64(x):
+    h = float(x)
+    h = h / 64.0
+    h = round(h)
+    h = int(h)
+    h = h * 64
+    return h
 
 
 def sdxl_encode_adm_patched(self, **kwargs):
     global positive_adm_scale, negative_adm_scale
 
-    clip_pooled = fcbh.model_base.sdxl_pooled(kwargs, self.noise_augmentor)
-    width = kwargs.get("width", 768)
-    height = kwargs.get("height", 768)
+    clip_pooled = ldm_patched.modules.model_base.sdxl_pooled(kwargs, self.noise_augmentor)
+    width = kwargs.get("width", 1024)
+    height = kwargs.get("height", 1024)
     target_width = width
     target_height = height
 
@@ -240,25 +264,21 @@ def sdxl_encode_adm_patched(self, **kwargs):
         width = float(width) * positive_adm_scale
         height = float(height) * positive_adm_scale
 
-    # Avoid artifacts
-    width = int(width)
-    height = int(height)
-    crop_w = 0
-    crop_h = 0
-    target_width = int(target_width)
-    target_height = int(target_height)
+    def embedder(number_list):
+        h = [self.embedder(torch.Tensor([number])) for number in number_list]
+        y = torch.flatten(torch.cat(h)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+        return y
 
-    out_a = [self.embedder(torch.Tensor([height])), self.embedder(torch.Tensor([width])),
-             self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
-             self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
-    flat_a = torch.flatten(torch.cat(out_a)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+    width, height = round_to_64(width), round_to_64(height)
+    target_width, target_height = round_to_64(target_width), round_to_64(target_height)
 
-    out_b = [self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width])),
-             self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
-             self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
-    flat_b = torch.flatten(torch.cat(out_b)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+    adm_emphasized = embedder([height, width, 0, 0, target_height, target_width])
+    adm_consistent = embedder([target_height, target_width, 0, 0, target_height, target_width])
 
-    return torch.cat((clip_pooled.to(flat_a.device), flat_a, clip_pooled.to(flat_b.device), flat_b), dim=1)
+    clip_pooled = clip_pooled.to(adm_emphasized)
+    final_adm = torch.cat((clip_pooled, adm_emphasized, clip_pooled, adm_consistent), dim=1)
+
+    return final_adm
 
 
 def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
@@ -273,11 +293,11 @@ def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
 
     sections = len(to_encode)
     if has_weights or sections == 0:
-        to_encode.append(fcbh.sd1_clip.gen_empty_tokens(self.special_tokens, max_token_len))
+        to_encode.append(ldm_patched.modules.sd1_clip.gen_empty_tokens(self.special_tokens, max_token_len))
 
     out, pooled = self.encode(to_encode)
     if pooled is not None:
-        first_pooled = pooled[0:1].cpu()
+        first_pooled = pooled[0:1].to(ldm_patched.modules.model_management.intermediate_device())
     else:
         first_pooled = pooled
 
@@ -297,22 +317,23 @@ def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
         output.append(z)
 
     if len(output) == 0:
-        return out[-1:].cpu(), first_pooled
-
-    return torch.cat(output, dim=-2).cpu(), first_pooled
+        return out[-1:].to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
+    return torch.cat(output, dim=-2).to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
 
 
 def patched_KSamplerX0Inpaint_forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
     if inpaint_worker.current_task is not None:
+        latent_processor = self.inner_model.inner_model.process_latent_in
+        inpaint_latent = latent_processor(inpaint_worker.current_task.latent).to(x)
+        inpaint_mask = inpaint_worker.current_task.latent_mask.to(x)
+
         if getattr(self, 'energy_generator', None) is None:
             # avoid bad results by using different seeds.
             self.energy_generator = torch.Generator(device='cpu').manual_seed((seed + 1) % constants.MAX_SEED)
 
-        latent_processor = self.inner_model.inner_model.process_latent_in
-        inpaint_latent = latent_processor(inpaint_worker.current_task.latent).to(x)
-        inpaint_mask = inpaint_worker.current_task.latent_mask.to(x)
         energy_sigma = sigma.reshape([sigma.shape[0]] + [1] * (len(x.shape) - 1))
-        current_energy = torch.randn(x.size(), dtype=x.dtype, generator=self.energy_generator, device="cpu").to(x) * energy_sigma
+        current_energy = torch.randn(
+            x.size(), dtype=x.dtype, generator=self.energy_generator, device="cpu").to(x) * energy_sigma
         x = x * inpaint_mask + (inpaint_latent + current_energy) * (1.0 - inpaint_mask)
 
         out = self.inner_model(x, sigma,
@@ -342,27 +363,8 @@ def timed_adm(y, timesteps):
     return y
 
 
-def patched_timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
-    # Consistent with Kohya to reduce differences between model training and inference.
-
-    if not repeat_only:
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=timesteps.device)
-        args = timesteps[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-    else:
-        embedding = repeat(timesteps, 'b -> b d', d=dim)
-    return embedding
-
-
 def patched_cldm_forward(self, x, hint, timesteps, context, y=None, **kwargs):
-    t_emb = fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(
-        timesteps, self.model_channels, repeat_only=False).to(self.dtype)
-
+    t_emb = ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
     emb = self.time_embed(t_emb)
 
     guided_hint = self.input_hint_block(hint, emb, context)
@@ -376,7 +378,7 @@ def patched_cldm_forward(self, x, hint, timesteps, context, y=None, **kwargs):
         assert y.shape[0] == x.shape[0]
         emb = emb + self.label_emb(y)
 
-    h = x.type(self.dtype)
+    h = x
     for module, zero_conv in zip(self.input_blocks, self.zero_convs):
         if guided_hint is not None:
             h = module(h, emb, context)
@@ -403,35 +405,31 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
     self.current_step = 1.0 - timesteps.to(x) / 999.0
     global_diffusion_progress = float(self.current_step.detach().cpu().numpy().tolist()[0])
 
-    inpaint_fix = None
-    if getattr(self, 'in_inpaint', False) and inpaint_worker.current_task is not None:
-        inpaint_fix = inpaint_worker.current_task.inpaint_head_feature
-
-    transformer_options["original_shape"] = list(x.shape)
-    transformer_options["current_index"] = 0
-    transformer_patches = transformer_options.get("patches", {})
-
     y = timed_adm(y, timesteps)
 
+    transformer_options["original_shape"] = list(x.shape)
+    transformer_options["transformer_index"] = 0
+    transformer_patches = transformer_options.get("patches", {})
+
+    num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
+    image_only_indicator = kwargs.get("image_only_indicator", self.default_image_only_indicator)
+    time_context = kwargs.get("time_context", None)
+
+    assert (y is not None) == (
+            self.num_classes is not None
+    ), "must specify y if and only if the model is class-conditional"
     hs = []
-    t_emb = fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(
-        timesteps, self.model_channels, repeat_only=False).to(self.dtype)
+    t_emb = ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
     emb = self.time_embed(t_emb)
 
     if self.num_classes is not None:
         assert y.shape[0] == x.shape[0]
         emb = emb + self.label_emb(y)
 
-    h = x.type(self.dtype)
+    h = x
     for id, module in enumerate(self.input_blocks):
         transformer_options["block"] = ("input", id)
-        h = forward_timestep_embed(module, h, emb, context, transformer_options)
-
-        if inpaint_fix is not None:
-            if int(h.shape[1]) == int(inpaint_fix.shape[1]):
-                h = h + inpaint_fix.to(h)
-                inpaint_fix = None
-
+        h = forward_timestep_embed(module, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
         h = apply_control(h, control, 'input')
         if "input_block_patch" in transformer_patches:
             patch = transformer_patches["input_block_patch"]
@@ -445,7 +443,7 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
                 h = p(h, transformer_options)
 
     transformer_options["block"] = ("middle", 0)
-    h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options)
+    h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
     h = apply_control(h, control, 'middle')
 
     for id, module in enumerate(self.output_blocks):
@@ -464,7 +462,7 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
             output_shape = hs[-1].shape
         else:
             output_shape = None
-        h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape)
+        h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
     h = h.type(x.dtype)
     if self.predict_codebook_ids:
         return self.id_predictor(h)
@@ -472,34 +470,9 @@ def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=
         return self.out(h)
 
 
-def patched_register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
-                          linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
-    # Consistent with Kohya to reduce differences between model training and inference.
-
-    if given_betas is not None:
-        betas = given_betas
-    else:
-        betas = make_beta_schedule(
-            beta_schedule,
-            timesteps,
-            linear_start=linear_start,
-            linear_end=linear_end,
-            cosine_s=cosine_s)
-
-    alphas = 1. - betas
-    alphas_cumprod = np.cumprod(alphas, axis=0)
-    timesteps, = betas.shape
-    self.num_timesteps = int(timesteps)
-    self.linear_start = linear_start
-    self.linear_end = linear_end
-    sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
-    self.set_sigmas(sigmas)
-    return
-
-
 def patched_load_models_gpu(*args, **kwargs):
     execution_start_time = time.perf_counter()
-    y = fcbh.model_management.load_models_gpu_origin(*args, **kwargs)
+    y = ldm_patched.modules.model_management.load_models_gpu_origin(*args, **kwargs)
     moving_time = time.perf_counter() - execution_start_time
     if moving_time > 0.1:
         print(f'[Fooocus Model Management] Moving model(s) has taken {moving_time:.2f} seconds')
@@ -541,19 +514,18 @@ def build_loaded(module, loader_name):
 
 
 def patch_all():
-    if not hasattr(fcbh.model_management, 'load_models_gpu_origin'):
-        fcbh.model_management.load_models_gpu_origin = fcbh.model_management.load_models_gpu
+    if not hasattr(ldm_patched.modules.model_management, 'load_models_gpu_origin'):
+        ldm_patched.modules.model_management.load_models_gpu_origin = ldm_patched.modules.model_management.load_models_gpu
 
-    fcbh.model_management.load_models_gpu = patched_load_models_gpu
-    fcbh.model_patcher.ModelPatcher.calculate_weight = calculate_weight_patched
-    fcbh.cldm.cldm.ControlNet.forward = patched_cldm_forward
-    fcbh.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
-    fcbh.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
-    fcbh.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights_patched_with_a1111_method
-    fcbh.samplers.KSamplerX0Inpaint.forward = patched_KSamplerX0Inpaint_forward
-    fcbh.k_diffusion.sampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSamplerPatched
-    fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding = patched_timestep_embedding
-    fcbh.model_base.ModelSamplingDiscrete._register_schedule = patched_register_schedule
+    ldm_patched.modules.model_management.load_models_gpu = patched_load_models_gpu
+    ldm_patched.modules.model_patcher.ModelPatcher.calculate_weight = calculate_weight_patched
+    ldm_patched.controlnet.cldm.ControlNet.forward = patched_cldm_forward
+    ldm_patched.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
+    ldm_patched.modules.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
+    ldm_patched.modules.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights_patched_with_a1111_method
+    ldm_patched.modules.samplers.KSamplerX0Inpaint.forward = patched_KSamplerX0Inpaint_forward
+    ldm_patched.k_diffusion.sampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSamplerPatched
+    ldm_patched.modules.samplers.sampling_function = patched_sampling_function
 
     warnings.filterwarnings(action='ignore', module='torchsde')
 

modules/private_logger.py

@@ -35,16 +35,19 @@ def log(img, dic, single_line_number=3):
 
     div_name = only_name.replace('.', '_')
     item = f'<div id="{div_name}">\n'
-    item += f"<p>{only_name}</p>\n"
+    item += "<table><tr>"
+    item += f"<td><img src=\"{only_name}\" width=auto height=100% loading=lazy style=\"height:auto;max-width:512px\" onerror=\"document.getElementById('{div_name}').style.display = 'none';\"></img></p></td>"
+    item += f"<td style=\"padding-left:10px;\"><p>{only_name}</p>\n"
     for i, (k, v) in enumerate(dic):
         if i < single_line_number:
-            item += f"<p>{k}: <b>{v}</b> </p>\n"
+            item += f"<p>{k}: <b>{v}</b></p>\n"
         else:
             if (i - single_line_number) % 2 == 0:
                 item += f"<p>{k}: <b>{v}</b>, "
             else:
                 item += f"{k}: <b>{v}</b></p>\n"
-    item += f"<p><img src=\"{only_name}\" width=512 onerror=\"document.getElementById('{div_name}').style.display = 'none';\"></img></p><hr></div>\n"
+    item += "</td>"
+    item += "</tr></table><hr></div>\n"
     existing_log = item + existing_log
 
     with open(html_name, 'w', encoding='utf-8') as f:

modules/sample_hijack.py

@@ -1,11 +1,15 @@
 import torch
-import fcbh.samplers
-import fcbh.model_management
+import ldm_patched.modules.samplers
+import ldm_patched.modules.model_management
 
-from fcbh.model_base import SDXLRefiner, SDXL
-from fcbh.conds import CONDRegular
-from fcbh.sample import get_additional_models, get_models_from_cond, cleanup_additional_models
-from fcbh.samplers import resolve_areas_and_cond_masks, wrap_model, calculate_start_end_timesteps, \
+from collections import namedtuple
+from ldm_patched.contrib.external_custom_sampler import SDTurboScheduler
+from ldm_patched.k_diffusion import sampling as k_diffusion_sampling
+from ldm_patched.modules.samplers import normal_scheduler, simple_scheduler, ddim_scheduler
+from ldm_patched.modules.model_base import SDXLRefiner, SDXL
+from ldm_patched.modules.conds import CONDRegular
+from ldm_patched.modules.sample import get_additional_models, get_models_from_cond, cleanup_additional_models
+from ldm_patched.modules.samplers import resolve_areas_and_cond_masks, wrap_model, calculate_start_end_timesteps, \
     create_cond_with_same_area_if_none, pre_run_control, apply_empty_x_to_equal_area, encode_model_conds
 
 
@@ -133,7 +137,9 @@ def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas
         extra_args['model_options'] = {k: {} if k == 'transformer_options' else v for k, v in extra_args['model_options'].items()}
 
         models, inference_memory = get_additional_models(positive_refiner, negative_refiner, current_refiner.model_dtype())
-        fcbh.model_management.load_models_gpu([current_refiner] + models, current_refiner.memory_required(noise.shape) + inference_memory)
+        ldm_patched.modules.model_management.load_models_gpu(
+            [current_refiner] + models,
+            model.memory_required([noise.shape[0] * 2] + list(noise.shape[1:])) + inference_memory)
 
         model_wrap.inner_model = current_refiner.model
         print('Refiner Swapped')
@@ -152,4 +158,27 @@ def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas
     return model.process_latent_out(samples.to(torch.float32))
 
 
-fcbh.samplers.sample = sample_hacked
+@torch.no_grad()
+@torch.inference_mode()
+def calculate_sigmas_scheduler_hacked(model, scheduler_name, steps):
+    if scheduler_name == "karras":
+        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
+    elif scheduler_name == "exponential":
+        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
+    elif scheduler_name == "normal":
+        sigmas = normal_scheduler(model, steps)
+    elif scheduler_name == "simple":
+        sigmas = simple_scheduler(model, steps)
+    elif scheduler_name == "ddim_uniform":
+        sigmas = ddim_scheduler(model, steps)
+    elif scheduler_name == "sgm_uniform":
+        sigmas = normal_scheduler(model, steps, sgm=True)
+    elif scheduler_name == "turbo":
+        sigmas = SDTurboScheduler().get_sigmas(namedtuple('Patcher', ['model'])(model=model), steps)[0]
+    else:
+        raise TypeError("error invalid scheduler")
+    return sigmas
+
+
+ldm_patched.modules.samplers.calculate_sigmas_scheduler = calculate_sigmas_scheduler_hacked
+ldm_patched.modules.samplers.sample = sample_hacked

modules/upscaler.py

@@ -1,8 +1,9 @@
 import os
 import torch
+import modules.core as core
 
-from fcbh_extras.chainner_models.architecture.RRDB import RRDBNet as ESRGAN
-from fcbh_extras.nodes_upscale_model import ImageUpscaleWithModel
+from ldm_patched.pfn.architecture.RRDB import RRDBNet as ESRGAN
+from ldm_patched.contrib.external_upscale_model import ImageUpscaleWithModel
 from collections import OrderedDict
 from modules.config import path_upscale_models
 
@@ -13,6 +14,9 @@ model = None
 
 def perform_upscale(img):
     global model
+
+    print(f'Upscaling image with shape {str(img.shape)} ...')
+
     if model is None:
         sd = torch.load(model_filename)
         sdo = OrderedDict()
@@ -22,4 +26,9 @@ def perform_upscale(img):
         model = ESRGAN(sdo)
         model.cpu()
         model.eval()
-    return opImageUpscaleWithModel.upscale(model, img)[0]
+
+    img = core.numpy_to_pytorch(img)
+    img = opImageUpscaleWithModel.upscale(model, img)[0]
+    img = core.pytorch_to_numpy(img)[0]
+
+    return img

modules/util.py

@@ -79,7 +79,7 @@ def get_shape_ceil(h, w):
 
 
 def get_image_shape_ceil(im):
-    H, W, _ = im.shape
+    H, W = im.shape[:2]
     return get_shape_ceil(H, W)