sync with original main Fooocus repo

modules/advanced_parameters.py

@@ -0,0 +1,33 @@
+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, \
+    refiner_swap_method, \
+    freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+    debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field, \
+    inpaint_mask_upload_checkbox, invert_mask_checkbox, inpaint_erode_or_dilate = [None] * 35
+
+
+def set_all_advanced_parameters(*args):
+    global 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, \
+        refiner_swap_method, \
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+        debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field, \
+        inpaint_mask_upload_checkbox, invert_mask_checkbox, inpaint_erode_or_dilate
+
+    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, \
+        refiner_swap_method, \
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2, \
+        debugging_inpaint_preprocessor, inpaint_disable_initial_latent, inpaint_engine, inpaint_strength, inpaint_respective_field, \
+        inpaint_mask_upload_checkbox, invert_mask_checkbox, inpaint_erode_or_dilate = args
+
+    return

modules/anisotropic.py

@@ -0,0 +1,200 @@
+import torch
+
+
+Tensor = torch.Tensor
+Device = torch.DeviceObjType
+Dtype = torch.Type
+pad = torch.nn.functional.pad
+
+
+def _compute_zero_padding(kernel_size: tuple[int, int] | int) -> tuple[int, int]:
+    ky, kx = _unpack_2d_ks(kernel_size)
+    return (ky - 1) // 2, (kx - 1) // 2
+
+
+def _unpack_2d_ks(kernel_size: tuple[int, int] | int) -> tuple[int, int]:
+    if isinstance(kernel_size, int):
+        ky = kx = kernel_size
+    else:
+        assert len(kernel_size) == 2, '2D Kernel size should have a length of 2.'
+        ky, kx = kernel_size
+
+    ky = int(ky)
+    kx = int(kx)
+    return ky, kx
+
+
+def gaussian(
+    window_size: int, sigma: Tensor | float, *, device: Device | None = None, dtype: Dtype | None = None
+) -> Tensor:
+
+    batch_size = sigma.shape[0]
+
+    x = (torch.arange(window_size, device=sigma.device, dtype=sigma.dtype) - window_size // 2).expand(batch_size, -1)
+
+    if window_size % 2 == 0:
+        x = x + 0.5
+
+    gauss = torch.exp(-x.pow(2.0) / (2 * sigma.pow(2.0)))
+
+    return gauss / gauss.sum(-1, keepdim=True)
+
+
+def get_gaussian_kernel1d(
+    kernel_size: int,
+    sigma: float | Tensor,
+    force_even: bool = False,
+    *,
+    device: Device | None = None,
+    dtype: Dtype | None = None,
+) -> Tensor:
+
+    return gaussian(kernel_size, sigma, device=device, dtype=dtype)
+
+
+def get_gaussian_kernel2d(
+    kernel_size: tuple[int, int] | int,
+    sigma: tuple[float, float] | Tensor,
+    force_even: bool = False,
+    *,
+    device: Device | None = None,
+    dtype: Dtype | None = None,
+) -> Tensor:
+
+    sigma = torch.Tensor([[sigma, sigma]]).to(device=device, dtype=dtype)
+
+    ksize_y, ksize_x = _unpack_2d_ks(kernel_size)
+    sigma_y, sigma_x = sigma[:, 0, None], sigma[:, 1, None]
+
+    kernel_y = get_gaussian_kernel1d(ksize_y, sigma_y, force_even, device=device, dtype=dtype)[..., None]
+    kernel_x = get_gaussian_kernel1d(ksize_x, sigma_x, force_even, device=device, dtype=dtype)[..., None]
+
+    return kernel_y * kernel_x.view(-1, 1, ksize_x)
+
+
+def _bilateral_blur(
+    input: Tensor,
+    guidance: Tensor | None,
+    kernel_size: tuple[int, int] | int,
+    sigma_color: float | Tensor,
+    sigma_space: tuple[float, float] | Tensor,
+    border_type: str = 'reflect',
+    color_distance_type: str = 'l1',
+) -> Tensor:
+
+    if isinstance(sigma_color, Tensor):
+        sigma_color = sigma_color.to(device=input.device, dtype=input.dtype).view(-1, 1, 1, 1, 1)
+
+    ky, kx = _unpack_2d_ks(kernel_size)
+    pad_y, pad_x = _compute_zero_padding(kernel_size)
+
+    padded_input = pad(input, (pad_x, pad_x, pad_y, pad_y), mode=border_type)
+    unfolded_input = padded_input.unfold(2, ky, 1).unfold(3, kx, 1).flatten(-2)  # (B, C, H, W, Ky x Kx)
+
+    if guidance is None:
+        guidance = input
+        unfolded_guidance = unfolded_input
+    else:
+        padded_guidance = pad(guidance, (pad_x, pad_x, pad_y, pad_y), mode=border_type)
+        unfolded_guidance = padded_guidance.unfold(2, ky, 1).unfold(3, kx, 1).flatten(-2)  # (B, C, H, W, Ky x Kx)
+
+    diff = unfolded_guidance - guidance.unsqueeze(-1)
+    if color_distance_type == "l1":
+        color_distance_sq = diff.abs().sum(1, keepdim=True).square()
+    elif color_distance_type == "l2":
+        color_distance_sq = diff.square().sum(1, keepdim=True)
+    else:
+        raise ValueError("color_distance_type only acceps l1 or l2")
+    color_kernel = (-0.5 / sigma_color**2 * color_distance_sq).exp()  # (B, 1, H, W, Ky x Kx)
+
+    space_kernel = get_gaussian_kernel2d(kernel_size, sigma_space, device=input.device, dtype=input.dtype)
+    space_kernel = space_kernel.view(-1, 1, 1, 1, kx * ky)
+
+    kernel = space_kernel * color_kernel
+    out = (unfolded_input * kernel).sum(-1) / kernel.sum(-1)
+    return out
+
+
+def bilateral_blur(
+    input: Tensor,
+    kernel_size: tuple[int, int] | int = (13, 13),
+    sigma_color: float | Tensor = 3.0,
+    sigma_space: tuple[float, float] | Tensor = 3.0,
+    border_type: str = 'reflect',
+    color_distance_type: str = 'l1',
+) -> Tensor:
+    return _bilateral_blur(input, None, kernel_size, sigma_color, sigma_space, border_type, color_distance_type)
+
+
+def adaptive_anisotropic_filter(x, g=None):
+    if g is None:
+        g = x
+    s, m = torch.std_mean(g, dim=(1, 2, 3), keepdim=True)
+    s = s + 1e-5
+    guidance = (g - m) / s
+    y = _bilateral_blur(x, guidance,
+                        kernel_size=(13, 13),
+                        sigma_color=3.0,
+                        sigma_space=3.0,
+                        border_type='reflect',
+                        color_distance_type='l1')
+    return y
+
+
+def joint_bilateral_blur(
+    input: Tensor,
+    guidance: Tensor,
+    kernel_size: tuple[int, int] | int,
+    sigma_color: float | Tensor,
+    sigma_space: tuple[float, float] | Tensor,
+    border_type: str = 'reflect',
+    color_distance_type: str = 'l1',
+) -> Tensor:
+    return _bilateral_blur(input, guidance, kernel_size, sigma_color, sigma_space, border_type, color_distance_type)
+
+
+class _BilateralBlur(torch.nn.Module):
+    def __init__(
+        self,
+        kernel_size: tuple[int, int] | int,
+        sigma_color: float | Tensor,
+        sigma_space: tuple[float, float] | Tensor,
+        border_type: str = 'reflect',
+        color_distance_type: str = "l1",
+    ) -> None:
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.sigma_color = sigma_color
+        self.sigma_space = sigma_space
+        self.border_type = border_type
+        self.color_distance_type = color_distance_type
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}"
+            f"(kernel_size={self.kernel_size}, "
+            f"sigma_color={self.sigma_color}, "
+            f"sigma_space={self.sigma_space}, "
+            f"border_type={self.border_type}, "
+            f"color_distance_type={self.color_distance_type})"
+        )
+
+
+class BilateralBlur(_BilateralBlur):
+    def forward(self, input: Tensor) -> Tensor:
+        return bilateral_blur(
+            input, self.kernel_size, self.sigma_color, self.sigma_space, self.border_type, self.color_distance_type
+        )
+
+
+class JointBilateralBlur(_BilateralBlur):
+    def forward(self, input: Tensor, guidance: Tensor) -> Tensor:
+        return joint_bilateral_blur(
+            input,
+            guidance,
+            self.kernel_size,
+            self.sigma_color,
+            self.sigma_space,
+            self.border_type,
+            self.color_distance_type,
+        )

modules/async_worker.py

@@ -0,0 +1,833 @@
+import threading
+
+
+class AsyncTask:
+    def __init__(self, args):
+        self.args = args
+        self.yields = []
+        self.results = []
+
+
+async_tasks = []
+
+
+def worker():
+    global async_tasks
+
+    import traceback
+    import math
+    import numpy as np
+    import torch
+    import time
+    import shared
+    import random
+    import copy
+    import modules.default_pipeline as pipeline
+    import modules.core as core
+    import modules.flags as flags
+    import modules.config
+    import modules.patch
+    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 extras.ip_adapter as ip_adapter
+    import extras.face_crop
+    import fooocus_version
+
+    from modules.sdxl_styles import apply_style, apply_wildcards, fooocus_expansion
+    from modules.private_logger import log
+    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, erode_or_dilate
+    from modules.upscaler import perform_upscale
+
+    try:
+        async_gradio_app = shared.gradio_root
+        flag = f'''App started successful. Use the app with {str(async_gradio_app.local_url)} or {str(async_gradio_app.server_name)}:{str(async_gradio_app.server_port)}'''
+        if async_gradio_app.share:
+            flag += f''' or {async_gradio_app.share_url}'''
+        print(flag)
+    except Exception as e:
+        print(e)
+
+    def progressbar(async_task, number, text):
+        print(f'[Fooocus] {text}')
+        async_task.yields.append(['preview', (number, text, None)])
+
+    def yield_result(async_task, imgs, do_not_show_finished_images=False):
+        if not isinstance(imgs, list):
+            imgs = [imgs]
+
+        async_task.results = async_task.results + imgs
+
+        if do_not_show_finished_images:
+            return
+
+        async_task.yields.append(['results', async_task.results])
+        return
+
+    def build_image_wall(async_task):
+        if not advanced_parameters.generate_image_grid:
+            return
+
+        results = async_task.results
+
+        if len(results) < 2:
+            return
+
+        for img in results:
+            if not isinstance(img, np.ndarray):
+                return
+            if img.ndim != 3:
+                return
+
+        H, W, C = results[0].shape
+
+        for img in results:
+            Hn, Wn, Cn = img.shape
+            if H != Hn:
+                return
+            if W != Wn:
+                return
+            if C != Cn:
+                return
+
+        cols = float(len(results)) ** 0.5
+        cols = int(math.ceil(cols))
+        rows = float(len(results)) / float(cols)
+        rows = int(math.ceil(rows))
+
+        wall = np.zeros(shape=(H * rows, W * cols, C), dtype=np.uint8)
+
+        for y in range(rows):
+            for x in range(cols):
+                if y * cols + x < len(results):
+                    img = results[y * cols + x]
+                    wall[y * H:y * H + H, x * W:x * W + W, :] = img
+
+        # must use deep copy otherwise gradio is super laggy. Do not use list.append() .
+        async_task.results = async_task.results + [wall]
+        return
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def handler(async_task):
+        execution_start_time = time.perf_counter()
+
+        args = async_task.args
+        args.reverse()
+
+        prompt = args.pop()
+        negative_prompt = args.pop()
+        style_selections = args.pop()
+        performance_selection = args.pop()
+        aspect_ratios_selection = args.pop()
+        image_number = args.pop()
+        image_seed = args.pop()
+        sharpness = args.pop()
+        guidance_scale = args.pop()
+        base_model_name = args.pop()
+        refiner_model_name = args.pop()
+        refiner_switch = args.pop()
+        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()
+        inpaint_mask_image_upload = args.pop()
+
+        cn_tasks = {x: [] for x in flags.ip_list}
+        for _ in range(4):
+            cn_img = args.pop()
+            cn_stop = args.pop()
+            cn_weight = args.pop()
+            cn_type = args.pop()
+            if cn_img is not None:
+                cn_tasks[cn_type].append([cn_img, cn_stop, cn_weight])
+
+        outpaint_selections = [o.lower() for o in outpaint_selections]
+        base_model_additional_loras = []
+        raw_style_selections = copy.deepcopy(style_selections)
+        uov_method = uov_method.lower()
+
+        if fooocus_expansion in style_selections:
+            use_expansion = True
+            style_selections.remove(fooocus_expansion)
+        else:
+            use_expansion = False
+
+        use_style = len(style_selections) > 0
+
+        if base_model_name == refiner_model_name:
+            print(f'Refiner disabled because base model and refiner are same.')
+            refiner_model_name = 'None'
+
+        assert performance_selection in ['Speed', 'Quality', 'Extreme Speed']
+
+        steps = 30
+
+        if performance_selection == 'Speed':
+            steps = 30
+
+        if performance_selection == 'Quality':
+            steps = 60
+
+        if performance_selection == 'Extreme Speed':
+            print('Enter LCM mode.')
+            progressbar(async_task, 1, 'Downloading LCM components ...')
+            loras += [(modules.config.downloading_sdxl_lcm_lora(), 1.0)]
+
+            if refiner_model_name != 'None':
+                print(f'Refiner disabled in LCM mode.')
+
+            refiner_model_name = 'None'
+            sampler_name = advanced_parameters.sampler_name = 'lcm'
+            scheduler_name = advanced_parameters.scheduler_name = 'lcm'
+            modules.patch.sharpness = sharpness = 0.0
+            cfg_scale = guidance_scale = 1.0
+            modules.patch.adaptive_cfg = advanced_parameters.adaptive_cfg = 1.0
+            refiner_switch = 1.0
+            modules.patch.positive_adm_scale = advanced_parameters.adm_scaler_positive = 1.0
+            modules.patch.negative_adm_scale = advanced_parameters.adm_scaler_negative = 1.0
+            modules.patch.adm_scaler_end = advanced_parameters.adm_scaler_end = 0.0
+            steps = 8
+
+        modules.patch.adaptive_cfg = advanced_parameters.adaptive_cfg
+        print(f'[Parameters] Adaptive CFG = {modules.patch.adaptive_cfg}')
+
+        modules.patch.sharpness = sharpness
+        print(f'[Parameters] Sharpness = {modules.patch.sharpness}')
+
+        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 = '
+              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}')
+
+        initial_latent = None
+        denoising_strength = 1.0
+        tiled = False
+
+        width, height = aspect_ratios_selection.replace('×', ' ').split(' ')[:2]
+        width, height = int(width), int(height)
+
+        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
+
+        seed = int(image_seed)
+        print(f'[Parameters] Seed = {seed}')
+
+        sampler_name = advanced_parameters.sampler_name
+        scheduler_name = advanced_parameters.scheduler_name
+
+        goals = []
+        tasks = []
+
+        if input_image_checkbox:
+            if (current_tab == 'uov' or (
+                    current_tab == 'ip' and advanced_parameters.mixing_image_prompt_and_vary_upscale)) \
+                    and uov_method != flags.disabled and uov_input_image is not None:
+                uov_input_image = HWC3(uov_input_image)
+                if 'vary' in uov_method:
+                    goals.append('vary')
+                elif 'upscale' in uov_method:
+                    goals.append('upscale')
+                    if 'fast' in uov_method:
+                        skip_prompt_processing = True
+                    else:
+                        steps = 18
+
+                        if performance_selection == 'Speed':
+                            steps = 18
+
+                        if performance_selection == 'Quality':
+                            steps = 36
+
+                        if performance_selection == 'Extreme Speed':
+                            steps = 8
+
+                    progressbar(async_task, 1, 'Downloading upscale models ...')
+                    modules.config.downloading_upscale_model()
+            if (current_tab == 'inpaint' or (
+                    current_tab == 'ip' and advanced_parameters.mixing_image_prompt_and_inpaint)) \
+                    and isinstance(inpaint_input_image, dict):
+                inpaint_image = inpaint_input_image['image']
+                inpaint_mask = inpaint_input_image['mask'][:, :, 0]
+                
+                if advanced_parameters.inpaint_mask_upload_checkbox:
+                    if isinstance(inpaint_mask_image_upload, np.ndarray):
+                        if inpaint_mask_image_upload.ndim == 3:
+                            H, W, C = inpaint_image.shape
+                            inpaint_mask_image_upload = resample_image(inpaint_mask_image_upload, width=W, height=H)
+                            inpaint_mask_image_upload = np.mean(inpaint_mask_image_upload, axis=2)
+                            inpaint_mask_image_upload = (inpaint_mask_image_upload > 127).astype(np.uint8) * 255
+                            inpaint_mask = np.maximum(inpaint_mask, inpaint_mask_image_upload)
+
+                if int(advanced_parameters.inpaint_erode_or_dilate) != 0:
+                    inpaint_mask = erode_or_dilate(inpaint_mask, advanced_parameters.inpaint_erode_or_dilate)
+
+                if advanced_parameters.invert_mask_checkbox:
+                    inpaint_mask = 255 - inpaint_mask
+
+                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 upscale models ...')
+                    modules.config.downloading_upscale_model()
+                    if inpaint_parameterized:
+                        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 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 \
+                    advanced_parameters.mixing_image_prompt_and_vary_upscale:
+                goals.append('cn')
+                progressbar(async_task, 1, 'Downloading control models ...')
+                if len(cn_tasks[flags.cn_canny]) > 0:
+                    controlnet_canny_path = modules.config.downloading_controlnet_canny()
+                if len(cn_tasks[flags.cn_cpds]) > 0:
+                    controlnet_cpds_path = modules.config.downloading_controlnet_cpds()
+                if len(cn_tasks[flags.cn_ip]) > 0:
+                    clip_vision_path, ip_negative_path, ip_adapter_path = modules.config.downloading_ip_adapters('ip')
+                if len(cn_tasks[flags.cn_ip_face]) > 0:
+                    clip_vision_path, ip_negative_path, ip_adapter_face_path = modules.config.downloading_ip_adapters(
+                        'face')
+                progressbar(async_task, 1, 'Loading control models ...')
+
+        # Load or unload CNs
+        pipeline.refresh_controlnets([controlnet_canny_path, controlnet_cpds_path])
+        ip_adapter.load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_path)
+        ip_adapter.load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_face_path)
+
+        switch = int(round(steps * refiner_switch))
+
+        if advanced_parameters.overwrite_step > 0:
+            steps = advanced_parameters.overwrite_step
+
+        if advanced_parameters.overwrite_switch > 0:
+            switch = advanced_parameters.overwrite_switch
+
+        if advanced_parameters.overwrite_width > 0:
+            width = advanced_parameters.overwrite_width
+
+        if advanced_parameters.overwrite_height > 0:
+            height = advanced_parameters.overwrite_height
+
+        print(f'[Parameters] Sampler = {sampler_name} - {scheduler_name}')
+        print(f'[Parameters] Steps = {steps} - {switch}')
+
+        progressbar(async_task, 1, 'Initializing ...')
+
+        if not skip_prompt_processing:
+
+            prompts = remove_empty_str([safe_str(p) for p in prompt.splitlines()], default='')
+            negative_prompts = remove_empty_str([safe_str(p) for p in negative_prompt.splitlines()], default='')
+
+            prompt = prompts[0]
+            negative_prompt = negative_prompts[0]
+
+            if prompt == '':
+                # disable expansion when empty since it is not meaningful and influences image prompt
+                use_expansion = False
+
+            extra_positive_prompts = prompts[1:] if len(prompts) > 1 else []
+            extra_negative_prompts = negative_prompts[1:] if len(negative_prompts) > 1 else []
+
+            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,
+                                        use_synthetic_refiner=use_synthetic_refiner)
+
+            progressbar(async_task, 3, 'Processing prompts ...')
+            tasks = []
+            for i in range(image_number):
+                task_seed = (seed + i) % (constants.MAX_SEED + 1)  # randint is inclusive, % is not
+                task_rng = random.Random(task_seed)  # may bind to inpaint noise in the future
+
+                task_prompt = apply_wildcards(prompt, task_rng)
+                task_negative_prompt = apply_wildcards(negative_prompt, task_rng)
+                task_extra_positive_prompts = [apply_wildcards(pmt, task_rng) for pmt in extra_positive_prompts]
+                task_extra_negative_prompts = [apply_wildcards(pmt, task_rng) for pmt in extra_negative_prompts]
+
+                positive_basic_workloads = []
+                negative_basic_workloads = []
+
+                if use_style:
+                    for s in style_selections:
+                        p, n = apply_style(s, positive=task_prompt)
+                        positive_basic_workloads = positive_basic_workloads + p
+                        negative_basic_workloads = negative_basic_workloads + n
+                else:
+                    positive_basic_workloads.append(task_prompt)
+
+                negative_basic_workloads.append(task_negative_prompt)  # Always use independent workload for negative.
+
+                positive_basic_workloads = positive_basic_workloads + task_extra_positive_prompts
+                negative_basic_workloads = negative_basic_workloads + task_extra_negative_prompts
+
+                positive_basic_workloads = remove_empty_str(positive_basic_workloads, default=task_prompt)
+                negative_basic_workloads = remove_empty_str(negative_basic_workloads, default=task_negative_prompt)
+
+                tasks.append(dict(
+                    task_seed=task_seed,
+                    task_prompt=task_prompt,
+                    task_negative_prompt=task_negative_prompt,
+                    positive=positive_basic_workloads,
+                    negative=negative_basic_workloads,
+                    expansion='',
+                    c=None,
+                    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),
+                ))
+
+            if use_expansion:
+                for i, t in enumerate(tasks):
+                    progressbar(async_task, 5, f'Preparing Fooocus text #{i + 1} ...')
+                    expansion = pipeline.final_expansion(t['task_prompt'], t['task_seed'])
+                    print(f'[Prompt Expansion] {expansion}')
+                    t['expansion'] = expansion
+                    t['positive'] = copy.deepcopy(t['positive']) + [expansion]  # Deep copy.
+
+            for i, t in enumerate(tasks):
+                progressbar(async_task, 7, f'Encoding positive #{i + 1} ...')
+                t['c'] = pipeline.clip_encode(texts=t['positive'], pool_top_k=t['positive_top_k'])
+
+            for i, t in enumerate(tasks):
+                if abs(float(cfg_scale) - 1.0) < 1e-4:
+                    t['uc'] = pipeline.clone_cond(t['c'])
+                else:
+                    progressbar(async_task, 10, f'Encoding negative #{i + 1} ...')
+                    t['uc'] = pipeline.clip_encode(texts=t['negative'], pool_top_k=t['negative_top_k'])
+
+        if len(goals) > 0:
+            progressbar(async_task, 13, 'Image processing ...')
+
+        if 'vary' in goals:
+            if 'subtle' in uov_method:
+                denoising_strength = 0.5
+            if 'strong' in uov_method:
+                denoising_strength = 0.85
+            if advanced_parameters.overwrite_vary_strength > 0:
+                denoising_strength = advanced_parameters.overwrite_vary_strength
+
+            shape_ceil = get_image_shape_ceil(uov_input_image)
+            if shape_ceil < 1024:
+                print(f'[Vary] Image is resized because it is too small.')
+                shape_ceil = 1024
+            elif shape_ceil > 2048:
+                print(f'[Vary] Image is resized because it is too big.')
+                shape_ceil = 2048
+
+            uov_input_image = set_image_shape_ceil(uov_input_image, shape_ceil)
+
+            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=candidate_vae, pixels=initial_pixels)
+            B, C, H, W = initial_latent['samples'].shape
+            width = W * 8
+            height = H * 8
+            print(f'Final resolution is {str((height, width))}.')
+
+        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 = perform_upscale(uov_input_image)
+            print(f'Image upscaled.')
+
+            if '1.5x' in uov_method:
+                f = 1.5
+            elif '2x' in uov_method:
+                f = 2.0
+            else:
+                f = 1.0
+
+            shape_ceil = get_shape_ceil(H * f, W * f)
+
+            if shape_ceil < 1024:
+                print(f'[Upscale] Image is resized because it is too small.')
+                uov_input_image = set_image_shape_ceil(uov_input_image, 1024)
+                shape_ceil = 1024
+            else:
+                uov_input_image = resample_image(uov_input_image, width=W * f, height=H * f)
+
+            image_is_super_large = shape_ceil > 2800
+
+            if 'fast' in uov_method:
+                direct_return = True
+            elif image_is_super_large:
+                print('Image is too large. Directly returned the SR image. '
+                      'Usually directly return SR image at 4K resolution '
+                      'yields better results than SDXL diffusion.')
+                direct_return = True
+            else:
+                direct_return = False
+
+            if direct_return:
+                d = [('Upscale (Fast)', '2x')]
+                log(uov_input_image, d)
+                yield_result(async_task, uov_input_image, do_not_show_finished_images=True)
+                return
+
+            tiled = True
+            denoising_strength = 0.382
+
+            if advanced_parameters.overwrite_upscale_strength > 0:
+                denoising_strength = advanced_parameters.overwrite_upscale_strength
+
+            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=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))}.')
+
+        if 'inpaint' in goals:
+            if len(outpaint_selections) > 0:
+                H, W, C = inpaint_image.shape
+                if 'top' in outpaint_selections:
+                    inpaint_image = np.pad(inpaint_image, [[int(H * 0.3), 0], [0, 0], [0, 0]], mode='edge')
+                    inpaint_mask = np.pad(inpaint_mask, [[int(H * 0.3), 0], [0, 0]], mode='constant',
+                                          constant_values=255)
+                if 'bottom' in outpaint_selections:
+                    inpaint_image = np.pad(inpaint_image, [[0, int(H * 0.3)], [0, 0], [0, 0]], mode='edge')
+                    inpaint_mask = np.pad(inpaint_mask, [[0, int(H * 0.3)], [0, 0]], mode='constant',
+                                          constant_values=255)
+
+                H, W, C = inpaint_image.shape
+                if 'left' in outpaint_selections:
+                    inpaint_image = np.pad(inpaint_image, [[0, 0], [int(H * 0.3), 0], [0, 0]], mode='edge')
+                    inpaint_mask = np.pad(inpaint_mask, [[0, 0], [int(H * 0.3), 0]], mode='constant',
+                                          constant_values=255)
+                if 'right' in outpaint_selections:
+                    inpaint_image = np.pad(inpaint_image, [[0, 0], [0, int(H * 0.3)], [0, 0]], mode='edge')
+                    inpaint_mask = np.pad(inpaint_mask, [[0, 0], [0, int(H * 0.3)]], mode='constant',
+                                          constant_values=255)
+
+                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
+
+            denoising_strength = advanced_parameters.inpaint_strength
+
+            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_inpaint_preprocessor:
+                yield_result(async_task, inpaint_worker.current_task.visualize_mask_processing(),
+                             do_not_show_finished_images=True)
+                return
+
+            progressbar(async_task, 13, 'VAE Inpaint encoding ...')
+
+            inpaint_pixel_fill = core.numpy_to_pytorch(inpaint_worker.current_task.interested_fill)
+            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=candidate_vae,
+                pixels=inpaint_pixel_image)
+
+            latent_swap = None
+            if candidate_vae_swap is not None:
+                progressbar(async_task, 13, 'VAE SD15 encoding ...')
+                latent_swap = core.encode_vae(
+                    vae=candidate_vae_swap,
+                    pixels=inpaint_pixel_fill)['samples']
+
+            progressbar(async_task, 13, 'VAE encoding ...')
+            latent_fill = core.encode_vae(
+                vae=candidate_vae,
+                pixels=inpaint_pixel_fill)['samples']
+
+            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]
+            print(f'Final resolution is {str((final_height, final_width))}, latent is {str((height, width))}.')
+
+        if 'cn' in goals:
+            for task in cn_tasks[flags.cn_canny]:
+                cn_img, cn_stop, cn_weight = task
+                cn_img = resize_image(HWC3(cn_img), width=width, height=height)
+
+                if not advanced_parameters.skipping_cn_preprocessor:
+                    cn_img = preprocessors.canny_pyramid(cn_img)
+
+                cn_img = HWC3(cn_img)
+                task[0] = core.numpy_to_pytorch(cn_img)
+                if advanced_parameters.debugging_cn_preprocessor:
+                    yield_result(async_task, cn_img, do_not_show_finished_images=True)
+                    return
+            for task in cn_tasks[flags.cn_cpds]:
+                cn_img, cn_stop, cn_weight = task
+                cn_img = resize_image(HWC3(cn_img), width=width, height=height)
+
+                if not advanced_parameters.skipping_cn_preprocessor:
+                    cn_img = preprocessors.cpds(cn_img)
+
+                cn_img = HWC3(cn_img)
+                task[0] = core.numpy_to_pytorch(cn_img)
+                if advanced_parameters.debugging_cn_preprocessor:
+                    yield_result(async_task, cn_img, do_not_show_finished_images=True)
+                    return
+            for task in cn_tasks[flags.cn_ip]:
+                cn_img, cn_stop, cn_weight = task
+                cn_img = HWC3(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)
+
+                task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_path)
+                if advanced_parameters.debugging_cn_preprocessor:
+                    yield_result(async_task, cn_img, do_not_show_finished_images=True)
+                    return
+            for task in cn_tasks[flags.cn_ip_face]:
+                cn_img, cn_stop, cn_weight = task
+                cn_img = HWC3(cn_img)
+
+                if not advanced_parameters.skipping_cn_preprocessor:
+                    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)
+
+                task[0] = ip_adapter.preprocess(cn_img, ip_adapter_path=ip_adapter_face_path)
+                if advanced_parameters.debugging_cn_preprocessor:
+                    yield_result(async_task, cn_img, do_not_show_finished_images=True)
+                    return
+
+            all_ip_tasks = cn_tasks[flags.cn_ip] + cn_tasks[flags.cn_ip_face]
+
+            if len(all_ip_tasks) > 0:
+                pipeline.final_unet = ip_adapter.patch_model(pipeline.final_unet, all_ip_tasks)
+
+        if advanced_parameters.freeu_enabled:
+            print(f'FreeU is enabled!')
+            pipeline.final_unet = core.apply_freeu(
+                pipeline.final_unet,
+                advanced_parameters.freeu_b1,
+                advanced_parameters.freeu_b2,
+                advanced_parameters.freeu_s1,
+                advanced_parameters.freeu_s2
+            )
+
+        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')
+
+        final_sampler_name = sampler_name
+        final_scheduler_name = scheduler_name
+
+        if scheduler_name == 'lcm':
+            final_scheduler_name = 'sgm_uniform'
+            if pipeline.final_unet is not None:
+                pipeline.final_unet = core.opModelSamplingDiscrete.patch(
+                    pipeline.final_unet,
+                    sampling='lcm',
+                    zsnr=False)[0]
+            if pipeline.final_refiner_unet is not None:
+                pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch(
+                    pipeline.final_refiner_unet,
+                    sampling='lcm',
+                    zsnr=False)[0]
+            print('Using lcm scheduler.')
+
+        async_task.yields.append(['preview', (13, 'Moving model to GPU ...', None)])
+
+        def callback(step, x0, x, total_steps, y):
+            done_steps = current_task_id * steps + step
+            async_task.yields.append(['preview', (
+                int(15.0 + 85.0 * float(done_steps) / float(all_steps)),
+                f'Step {step}/{total_steps} in the {current_task_id + 1}-th Sampling',
+                y)])
+
+        for current_task_id, task in enumerate(tasks):
+            execution_start_time = time.perf_counter()
+
+            try:
+                positive_cond, negative_cond = task['c'], task['uc']
+
+                if 'cn' in goals:
+                    for cn_flag, cn_path in [
+                        (flags.cn_canny, controlnet_canny_path),
+                        (flags.cn_cpds, controlnet_cpds_path)
+                    ]:
+                        for cn_img, cn_stop, cn_weight in cn_tasks[cn_flag]:
+                            positive_cond, negative_cond = core.apply_controlnet(
+                                positive_cond, negative_cond,
+                                pipeline.loaded_ControlNets[cn_path], cn_img, cn_weight, 0, cn_stop)
+
+                imgs = pipeline.process_diffusion(
+                    positive_cond=positive_cond,
+                    negative_cond=negative_cond,
+                    steps=steps,
+                    switch=switch,
+                    width=width,
+                    height=height,
+                    image_seed=task['task_seed'],
+                    callback=callback,
+                    sampler_name=final_sampler_name,
+                    scheduler_name=final_scheduler_name,
+                    latent=initial_latent,
+                    denoise=denoising_strength,
+                    tiled=tiled,
+                    cfg_scale=cfg_scale,
+                    refiner_swap_method=refiner_swap_method
+                )
+
+                del task['c'], task['uc'], positive_cond, negative_cond  # Save memory
+
+                if inpaint_worker.current_task is not None:
+                    imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
+
+                for x in imgs:
+                    d = [
+                        ('Prompt', task['log_positive_prompt']),
+                        ('Negative Prompt', task['log_negative_prompt']),
+                        ('Fooocus V2 Expansion', task['expansion']),
+                        ('Styles', str(raw_style_selections)),
+                        ('Performance', performance_selection),
+                        ('Resolution', str((width, height))),
+                        ('Sharpness', sharpness),
+                        ('Guidance Scale', guidance_scale),
+                        ('ADM Guidance', str((
+                            modules.patch.positive_adm_scale,
+                            modules.patch.negative_adm_scale,
+                            modules.patch.adm_scaler_end))),
+                        ('Base Model', base_model_name),
+                        ('Refiner Model', refiner_model_name),
+                        ('Refiner Switch', refiner_switch),
+                        ('Sampler', sampler_name),
+                        ('Scheduler', scheduler_name),
+                        ('Seed', task['task_seed']),
+                    ]
+                    for li, (n, w) in enumerate(loras):
+                        if n != 'None':
+                            d.append((f'LoRA {li + 1}', f'{n} : {w}'))
+                    d.append(('Version', 'v' + fooocus_version.version))
+                    log(x, d)
+
+                yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1)
+            except ldm_patched.modules.model_management.InterruptProcessingException as e:
+                if shared.last_stop == 'skip':
+                    print('User skipped')
+                    continue
+                else:
+                    print('User stopped')
+                    break
+
+            execution_time = time.perf_counter() - execution_start_time
+            print(f'Generating and saving time: {execution_time:.2f} seconds')
+
+        return
+
+    while True:
+        time.sleep(0.01)
+        if len(async_tasks) > 0:
+            task = async_tasks.pop(0)
+            try:
+                handler(task)
+                build_image_wall(task)
+                task.yields.append(['finish', task.results])
+                pipeline.prepare_text_encoder(async_call=True)
+            except:
+                traceback.print_exc()
+                task.yields.append(['finish', task.results])
+    pass
+
+
+threading.Thread(target=worker, daemon=True).start()

modules/auth.py

@@ -0,0 +1,41 @@
+import json
+import hashlib
+import modules.constants as constants
+
+from os.path import exists
+
+
+def auth_list_to_dict(auth_list):
+    auth_dict = {}
+    for auth_data in auth_list:
+        if 'user' in auth_data:
+            if 'hash' in auth_data:
+                auth_dict |= {auth_data['user']: auth_data['hash']}
+            elif 'pass' in auth_data:
+                auth_dict |= {auth_data['user']: hashlib.sha256(bytes(auth_data['pass'], encoding='utf-8')).hexdigest()}
+    return auth_dict
+
+
+def load_auth_data(filename=None):
+    auth_dict = None
+    if filename != None and exists(filename):
+        with open(filename, encoding='utf-8') as auth_file:
+            try:
+                auth_obj = json.load(auth_file)
+                if isinstance(auth_obj, list) and len(auth_obj) > 0:
+                    auth_dict = auth_list_to_dict(auth_obj)
+            except Exception as e:
+                print('load_auth_data, e: ' + str(e))
+    return auth_dict
+
+
+auth_dict = load_auth_data(constants.AUTH_FILENAME)
+
+auth_enabled = auth_dict != None
+
+
+def check_auth(user, password):
+    if user not in auth_dict:
+        return False
+    else:   
+        return hashlib.sha256(bytes(password, encoding='utf-8')).hexdigest() == auth_dict[user]

modules/config.py

@@ -0,0 +1,510 @@
+import os
+import json
+import math
+import numbers
+import args_manager
+import modules.flags
+import modules.sdxl_styles
+
+from modules.model_loader import load_file_from_url
+from modules.util import get_files_from_folder
+
+
+config_path = os.path.abspath("./config.txt")
+config_example_path = os.path.abspath("config_modification_tutorial.txt")
+config_dict = {}
+always_save_keys = []
+visited_keys = []
+
+try:
+    if os.path.exists(config_path):
+        with open(config_path, "r", encoding="utf-8") as json_file:
+            config_dict = json.load(json_file)
+            always_save_keys = list(config_dict.keys())
+except Exception as e:
+    print(f'Failed to load config file "{config_path}" . The reason is: {str(e)}')
+    print('Please make sure that:')
+    print(f'1. The file "{config_path}" is a valid text file, and you have access to read it.')
+    print('2. Use "\\\\" instead of "\\" when describing paths.')
+    print('3. There is no "," before the last "}".')
+    print('4. All key/value formats are correct.')
+
+
+def try_load_deprecated_user_path_config():
+    global config_dict
+
+    if not os.path.exists('user_path_config.txt'):
+        return
+
+    try:
+        deprecated_config_dict = json.load(open('user_path_config.txt', "r", encoding="utf-8"))
+
+        def replace_config(old_key, new_key):
+            if old_key in deprecated_config_dict:
+                config_dict[new_key] = deprecated_config_dict[old_key]
+                del deprecated_config_dict[old_key]
+
+        replace_config('modelfile_path', 'path_checkpoints')
+        replace_config('lorafile_path', 'path_loras')
+        replace_config('embeddings_path', 'path_embeddings')
+        replace_config('vae_approx_path', 'path_vae_approx')
+        replace_config('upscale_models_path', 'path_upscale_models')
+        replace_config('inpaint_models_path', 'path_inpaint')
+        replace_config('controlnet_models_path', 'path_controlnet')
+        replace_config('clip_vision_models_path', 'path_clip_vision')
+        replace_config('fooocus_expansion_path', 'path_fooocus_expansion')
+        replace_config('temp_outputs_path', 'path_outputs')
+
+        if deprecated_config_dict.get("default_model", None) == 'juggernautXL_version6Rundiffusion.safetensors':
+            os.replace('user_path_config.txt', 'user_path_config-deprecated.txt')
+            print('Config updated successfully in silence. '
+                  'A backup of previous config is written to "user_path_config-deprecated.txt".')
+            return
+
+        if input("Newer models and configs are available. "
+                 "Download and update files? [Y/n]:") in ['n', 'N', 'No', 'no', 'NO']:
+            config_dict.update(deprecated_config_dict)
+            print('Loading using deprecated old models and deprecated old configs.')
+            return
+        else:
+            os.replace('user_path_config.txt', 'user_path_config-deprecated.txt')
+            print('Config updated successfully by user. '
+                  'A backup of previous config is written to "user_path_config-deprecated.txt".')
+            return
+    except Exception as e:
+        print('Processing deprecated config failed')
+        print(e)
+    return
+
+
+try_load_deprecated_user_path_config()
+
+preset = args_manager.args.preset
+
+if isinstance(preset, str):
+    preset_path = os.path.abspath(f'./presets/{preset}.json')
+    try:
+        if os.path.exists(preset_path):
+            with open(preset_path, "r", encoding="utf-8") as json_file:
+                config_dict.update(json.load(json_file))
+                print(f'Loaded preset: {preset_path}')
+        else:
+            raise FileNotFoundError
+    except Exception as e:
+        print(f'Load preset [{preset_path}] failed')
+        print(e)
+
+
+def get_dir_or_set_default(key, default_value):
+    global config_dict, visited_keys, always_save_keys
+
+    if key not in visited_keys:
+        visited_keys.append(key)
+
+    if key not in always_save_keys:
+        always_save_keys.append(key)
+
+    v = config_dict.get(key, None)
+    if isinstance(v, str) and os.path.exists(v) and os.path.isdir(v):
+        return v
+    else:
+        if v is not None:
+            print(f'Failed to load config key: {json.dumps({key:v})} is invalid or does not exist; will use {json.dumps({key:default_value})} instead.')
+        dp = os.path.abspath(os.path.join(os.path.dirname(__file__), default_value))
+        os.makedirs(dp, exist_ok=True)
+        config_dict[key] = dp
+        return dp
+
+
+path_checkpoints = get_dir_or_set_default('path_checkpoints', '../models/checkpoints/')
+path_loras = get_dir_or_set_default('path_loras', '../models/loras/')
+path_embeddings = get_dir_or_set_default('path_embeddings', '../models/embeddings/')
+path_vae_approx = get_dir_or_set_default('path_vae_approx', '../models/vae_approx/')
+path_upscale_models = get_dir_or_set_default('path_upscale_models', '../models/upscale_models/')
+path_inpaint = get_dir_or_set_default('path_inpaint', '../models/inpaint/')
+path_controlnet = get_dir_or_set_default('path_controlnet', '../models/controlnet/')
+path_clip_vision = get_dir_or_set_default('path_clip_vision', '../models/clip_vision/')
+path_fooocus_expansion = get_dir_or_set_default('path_fooocus_expansion', '../models/prompt_expansion/fooocus_expansion')
+path_outputs = get_dir_or_set_default('path_outputs', '../outputs/')
+
+
+def get_config_item_or_set_default(key, default_value, validator, disable_empty_as_none=False):
+    global config_dict, visited_keys
+
+    if key not in visited_keys:
+        visited_keys.append(key)
+    
+    if key not in config_dict:
+        config_dict[key] = default_value
+        return default_value
+
+    v = config_dict.get(key, None)
+    if not disable_empty_as_none:
+        if v is None or v == '':
+            v = 'None'
+    if validator(v):
+        return v
+    else:
+        if v is not None:
+            print(f'Failed to load config key: {json.dumps({key:v})} is invalid; will use {json.dumps({key:default_value})} instead.')
+        config_dict[key] = default_value
+        return default_value
+
+
+default_base_model_name = get_config_item_or_set_default(
+    key='default_model',
+    default_value='juggernautXL_version6Rundiffusion.safetensors',
+    validator=lambda x: isinstance(x, str)
+)
+default_refiner_model_name = get_config_item_or_set_default(
+    key='default_refiner',
+    default_value='None',
+    validator=lambda x: isinstance(x, str)
+)
+default_refiner_switch = get_config_item_or_set_default(
+    key='default_refiner_switch',
+    default_value=0.5,
+    validator=lambda x: isinstance(x, numbers.Number) and 0 <= x <= 1
+)
+default_loras = get_config_item_or_set_default(
+    key='default_loras',
+    default_value=[
+        [
+            "sd_xl_offset_example-lora_1.0.safetensors",
+            0.1
+        ],
+        [
+            "None",
+            1.0
+        ],
+        [
+            "None",
+            1.0
+        ],
+        [
+            "None",
+            1.0
+        ],
+        [
+            "None",
+            1.0
+        ]
+    ],
+    validator=lambda x: isinstance(x, list) and all(len(y) == 2 and isinstance(y[0], str) and isinstance(y[1], numbers.Number) for y in x)
+)
+default_cfg_scale = get_config_item_or_set_default(
+    key='default_cfg_scale',
+    default_value=4.0,
+    validator=lambda x: isinstance(x, numbers.Number)
+)
+default_sample_sharpness = get_config_item_or_set_default(
+    key='default_sample_sharpness',
+    default_value=2.0,
+    validator=lambda x: isinstance(x, numbers.Number)
+)
+default_sampler = get_config_item_or_set_default(
+    key='default_sampler',
+    default_value='dpmpp_2m_sde_gpu',
+    validator=lambda x: x in modules.flags.sampler_list
+)
+default_scheduler = get_config_item_or_set_default(
+    key='default_scheduler',
+    default_value='karras',
+    validator=lambda x: x in modules.flags.scheduler_list
+)
+default_styles = get_config_item_or_set_default(
+    key='default_styles',
+    default_value=[
+        "Fooocus V2",
+        "Fooocus Enhance",
+        "Fooocus Sharp"
+    ],
+    validator=lambda x: isinstance(x, list) and all(y in modules.sdxl_styles.legal_style_names for y in x)
+)
+default_prompt_negative = get_config_item_or_set_default(
+    key='default_prompt_negative',
+    default_value='',
+    validator=lambda x: isinstance(x, str),
+    disable_empty_as_none=True
+)
+default_prompt = get_config_item_or_set_default(
+    key='default_prompt',
+    default_value='',
+    validator=lambda x: isinstance(x, str),
+    disable_empty_as_none=True
+)
+default_performance = get_config_item_or_set_default(
+    key='default_performance',
+    default_value='Speed',
+    validator=lambda x: x in modules.flags.performance_selections
+)
+default_advanced_checkbox = get_config_item_or_set_default(
+    key='default_advanced_checkbox',
+    default_value=False,
+    validator=lambda x: isinstance(x, bool)
+)
+default_max_image_number = get_config_item_or_set_default(
+    key='default_max_image_number',
+    default_value=32,
+    validator=lambda x: isinstance(x, int) and x >= 1
+)
+default_image_number = get_config_item_or_set_default(
+    key='default_image_number',
+    default_value=2,
+    validator=lambda x: isinstance(x, int) and 1 <= x <= default_max_image_number
+)
+checkpoint_downloads = get_config_item_or_set_default(
+    key='checkpoint_downloads',
+    default_value={
+        "juggernautXL_version6Rundiffusion.safetensors": "https://huggingface.co/lllyasviel/fav_models/resolve/main/fav/juggernautXL_version6Rundiffusion.safetensors"
+    },
+    validator=lambda x: isinstance(x, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in x.items())
+)
+lora_downloads = get_config_item_or_set_default(
+    key='lora_downloads',
+    default_value={
+        "sd_xl_offset_example-lora_1.0.safetensors": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_offset_example-lora_1.0.safetensors"
+    },
+    validator=lambda x: isinstance(x, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in x.items())
+)
+embeddings_downloads = get_config_item_or_set_default(
+    key='embeddings_downloads',
+    default_value={},
+    validator=lambda x: isinstance(x, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in x.items())
+)
+available_aspect_ratios = get_config_item_or_set_default(
+    key='available_aspect_ratios',
+    default_value=[
+        '704*1408', '704*1344', '768*1344', '768*1280', '832*1216', '832*1152',
+        '896*1152', '896*1088', '960*1088', '960*1024', '1024*1024', '1024*960',
+        '1088*960', '1088*896', '1152*896', '1152*832', '1216*832', '1280*768',
+        '1344*768', '1344*704', '1408*704', '1472*704', '1536*640', '1600*640',
+        '1664*576', '1728*576'
+    ],
+    validator=lambda x: isinstance(x, list) and all('*' in v for v in x) and len(x) > 1
+)
+default_aspect_ratio = get_config_item_or_set_default(
+    key='default_aspect_ratio',
+    default_value='1152*896' if '1152*896' in available_aspect_ratios else available_aspect_ratios[0],
+    validator=lambda x: x in available_aspect_ratios
+)
+default_inpaint_engine_version = get_config_item_or_set_default(
+    key='default_inpaint_engine_version',
+    default_value='v2.6',
+    validator=lambda x: x in modules.flags.inpaint_engine_versions
+)
+default_cfg_tsnr = get_config_item_or_set_default(
+    key='default_cfg_tsnr',
+    default_value=7.0,
+    validator=lambda x: isinstance(x, numbers.Number)
+)
+default_overwrite_step = get_config_item_or_set_default(
+    key='default_overwrite_step',
+    default_value=-1,
+    validator=lambda x: isinstance(x, int)
+)
+default_overwrite_switch = get_config_item_or_set_default(
+    key='default_overwrite_switch',
+    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]
+
+config_dict["default_loras"] = default_loras = default_loras[:5] + [['None', 1.0] for _ in range(5 - len(default_loras))]
+
+possible_preset_keys = [
+    "default_model",
+    "default_refiner",
+    "default_refiner_switch",
+    "default_loras",
+    "default_cfg_scale",
+    "default_sample_sharpness",
+    "default_sampler",
+    "default_scheduler",
+    "default_performance",
+    "default_prompt",
+    "default_prompt_negative",
+    "default_styles",
+    "default_aspect_ratio",
+    "checkpoint_downloads",
+    "embeddings_downloads",
+    "lora_downloads",
+]
+
+
+REWRITE_PRESET = False
+
+if REWRITE_PRESET and isinstance(args_manager.args.preset, str):
+    save_path = 'presets/' + args_manager.args.preset + '.json'
+    with open(save_path, "w", encoding="utf-8") as json_file:
+        json.dump({k: config_dict[k] for k in possible_preset_keys}, json_file, indent=4)
+    print(f'Preset saved to {save_path}. Exiting ...')
+    exit(0)
+
+
+def add_ratio(x):
+    a, b = x.replace('*', ' ').split(' ')[:2]
+    a, b = int(a), int(b)
+    g = math.gcd(a, b)
+    return f'{a}×{b} <span style="color: grey;"> \U00002223 {a // g}:{b // g}</span>'
+
+
+default_aspect_ratio = add_ratio(default_aspect_ratio)
+available_aspect_ratios = [add_ratio(x) for x in available_aspect_ratios]
+
+
+# Only write config in the first launch.
+if not os.path.exists(config_path):
+    with open(config_path, "w", encoding="utf-8") as json_file:
+        json.dump({k: config_dict[k] for k in always_save_keys}, json_file, indent=4)
+
+
+# Always write tutorials.
+with open(config_example_path, "w", encoding="utf-8") as json_file:
+    cpa = config_path.replace("\\", "\\\\")
+    json_file.write(f'You can modify your "{cpa}" using the below keys, formats, and examples.\n'
+                    f'Do not modify this file. Modifications in this file will not take effect.\n'
+                    f'This file is a tutorial and example. Please edit "{cpa}" to really change any settings.\n'
+                    + 'Remember to split the paths with "\\\\" rather than "\\", '
+                      'and there is no "," before the last "}". \n\n\n')
+    json.dump({k: config_dict[k] for k in visited_keys}, json_file, indent=4)
+
+
+os.makedirs(path_outputs, exist_ok=True)
+
+model_filenames = []
+lora_filenames = []
+
+
+def get_model_filenames(folder_path, name_filter=None):
+    return get_files_from_folder(folder_path, ['.pth', '.ckpt', '.bin', '.safetensors', '.fooocus.patch'], name_filter)
+
+
+def update_all_model_names():
+    global model_filenames, lora_filenames
+    model_filenames = get_model_filenames(path_checkpoints)
+    lora_filenames = get_model_filenames(path_loras)
+    return
+
+
+def downloading_inpaint_models(v):
+    assert v in modules.flags.inpaint_engine_versions
+
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/fooocus_inpaint_head.pth',
+        model_dir=path_inpaint,
+        file_name='fooocus_inpaint_head.pth'
+    )
+    head_file = os.path.join(path_inpaint, 'fooocus_inpaint_head.pth')
+    patch_file = None
+
+    if v == 'v1':
+        load_file_from_url(
+            url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/inpaint.fooocus.patch',
+            model_dir=path_inpaint,
+            file_name='inpaint.fooocus.patch'
+        )
+        patch_file = os.path.join(path_inpaint, 'inpaint.fooocus.patch')
+
+    if v == 'v2.5':
+        load_file_from_url(
+            url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/inpaint_v25.fooocus.patch',
+            model_dir=path_inpaint,
+            file_name='inpaint_v25.fooocus.patch'
+        )
+        patch_file = os.path.join(path_inpaint, 'inpaint_v25.fooocus.patch')
+
+    if v == 'v2.6':
+        load_file_from_url(
+            url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/inpaint_v26.fooocus.patch',
+            model_dir=path_inpaint,
+            file_name='inpaint_v26.fooocus.patch'
+        )
+        patch_file = os.path.join(path_inpaint, 'inpaint_v26.fooocus.patch')
+
+    return head_file, patch_file
+
+
+def downloading_sdxl_lcm_lora():
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/sdxl_lcm_lora.safetensors',
+        model_dir=path_loras,
+        file_name='sdxl_lcm_lora.safetensors'
+    )
+    return 'sdxl_lcm_lora.safetensors'
+
+
+def downloading_controlnet_canny():
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/control-lora-canny-rank128.safetensors',
+        model_dir=path_controlnet,
+        file_name='control-lora-canny-rank128.safetensors'
+    )
+    return os.path.join(path_controlnet, 'control-lora-canny-rank128.safetensors')
+
+
+def downloading_controlnet_cpds():
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_xl_cpds_128.safetensors',
+        model_dir=path_controlnet,
+        file_name='fooocus_xl_cpds_128.safetensors'
+    )
+    return os.path.join(path_controlnet, 'fooocus_xl_cpds_128.safetensors')
+
+
+def downloading_ip_adapters(v):
+    assert v in ['ip', 'face']
+
+    results = []
+
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/clip_vision_vit_h.safetensors',
+        model_dir=path_clip_vision,
+        file_name='clip_vision_vit_h.safetensors'
+    )
+    results += [os.path.join(path_clip_vision, 'clip_vision_vit_h.safetensors')]
+
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_ip_negative.safetensors',
+        model_dir=path_controlnet,
+        file_name='fooocus_ip_negative.safetensors'
+    )
+    results += [os.path.join(path_controlnet, 'fooocus_ip_negative.safetensors')]
+
+    if v == 'ip':
+        load_file_from_url(
+            url='https://huggingface.co/lllyasviel/misc/resolve/main/ip-adapter-plus_sdxl_vit-h.bin',
+            model_dir=path_controlnet,
+            file_name='ip-adapter-plus_sdxl_vit-h.bin'
+        )
+        results += [os.path.join(path_controlnet, 'ip-adapter-plus_sdxl_vit-h.bin')]
+
+    if v == 'face':
+        load_file_from_url(
+            url='https://huggingface.co/lllyasviel/misc/resolve/main/ip-adapter-plus-face_sdxl_vit-h.bin',
+            model_dir=path_controlnet,
+            file_name='ip-adapter-plus-face_sdxl_vit-h.bin'
+        )
+        results += [os.path.join(path_controlnet, 'ip-adapter-plus-face_sdxl_vit-h.bin')]
+
+    return results
+
+
+def downloading_upscale_model():
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_upscaler_s409985e5.bin',
+        model_dir=path_upscale_models,
+        file_name='fooocus_upscaler_s409985e5.bin'
+    )
+    return os.path.join(path_upscale_models, 'fooocus_upscaler_s409985e5.bin')
+
+
+update_all_model_names()

modules/constants.py

@@ -0,0 +1,5 @@
+# as in k-diffusion (sampling.py)
+MIN_SEED = 0
+MAX_SEED = 2**63 - 1
+
+AUTH_FILENAME = 'auth.json'

modules/core.py

@@ -0,0 +1,344 @@
+from modules.patch import patch_all
+
+patch_all()
+
+
+import os
+import einops
+import torch
+import numpy as np
+
+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 ldm_patched.modules.samplers
+import ldm_patched.modules.latent_formats
+import modules.advanced_parameters
+
+from ldm_patched.modules.sd import load_checkpoint_guess_config
+from ldm_patched.contrib.external import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, \
+    ControlNetApplyAdvanced
+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 ldm_patched.contrib.external_model_advanced import ModelSamplingDiscrete
+
+
+opEmptyLatentImage = EmptyLatentImage()
+opVAEDecode = VAEDecode()
+opVAEEncode = VAEEncode()
+opVAEDecodeTiled = VAEDecodeTiled()
+opVAEEncodeTiled = VAEEncodeTiled()
+opControlNetApplyAdvanced = ControlNetApplyAdvanced()
+opFreeU = FreeU_V2()
+opModelSamplingDiscrete = ModelSamplingDiscrete()
+
+
+class StableDiffusionModel:
+    def __init__(self, unet=None, vae=None, clip=None, clip_vision=None, filename=None):
+        self.unet = unet
+        self.vae = vae
+        self.clip = clip
+        self.clip_vision = clip_vision
+        self.filename = filename
+        self.unet_with_lora = unet
+        self.clip_with_lora = clip
+        self.visited_loras = ''
+
+        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()
+    def refresh_loras(self, loras):
+        assert isinstance(loras, list)
+
+        if self.visited_loras == str(loras):
+            return
+
+        self.visited_loras = str(loras)
+
+        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
+
+            if os.path.exists(name):
+                lora_filename = name
+            else:
+                lora_filename = os.path.join(modules.config.path_loras, name)
+
+            if not os.path.exists(lora_filename):
+                print(f'Lora file not found: {lora_filename}')
+                continue
+
+            loras_to_load.append((lora_filename, weight))
+
+        self.unet_with_lora = self.unet.clone() if self.unet is not None else None
+        self.clip_with_lora = self.clip.clone() if self.clip is not None else None
+
+        for lora_filename, weight in loras_to_load:
+            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_unmatch) > 12:
+                # model mismatch
+                continue
+
+            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.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)
+
+            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()
+@torch.inference_mode()
+def apply_freeu(model, b1, b2, s1, s2):
+    return opFreeU.patch(model=model, b1=b1, b2=b2, s1=s1, s2=s2)[0]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def load_controlnet(ckpt_filename):
+    return ldm_patched.modules.controlnet.load_controlnet(ckpt_filename)
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def apply_controlnet(positive, negative, control_net, image, strength, start_percent, end_percent):
+    return opControlNetApplyAdvanced.apply_controlnet(positive=positive, negative=negative, control_net=control_net,
+        image=image, strength=strength, start_percent=start_percent, end_percent=end_percent)
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def load_model(ckpt_filename):
+    unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename, embedding_directory=path_embeddings)
+    return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, filename=ckpt_filename)
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def generate_empty_latent(width=1024, height=1024, batch_size=1):
+    return opEmptyLatentImage.generate(width=width, height=height, batch_size=batch_size)[0]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def decode_vae(vae, latent_image, tiled=False):
+    if tiled:
+        return opVAEDecodeTiled.decode(samples=latent_image, vae=vae, tile_size=512)[0]
+    else:
+        return opVAEDecode.decode(samples=latent_image, vae=vae)[0]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def encode_vae(vae, pixels, tiled=False):
+    if tiled:
+        return opVAEEncodeTiled.encode(pixels=pixels, vae=vae, tile_size=512)[0]
+    else:
+        return opVAEEncode.encode(pixels=pixels, vae=vae)[0]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def encode_vae_inpaint(vae, pixels, mask):
+    assert mask.ndim == 3 and pixels.ndim == 4
+    assert mask.shape[-1] == pixels.shape[-2]
+    assert mask.shape[-2] == pixels.shape[-3]
+
+    w = mask.round()[..., None]
+    pixels = pixels * (1 - w) + 0.5 * w
+
+    latent = vae.encode(pixels)
+    B, C, H, W = latent.shape
+
+    latent_mask = mask[:, None, :, :]
+    latent_mask = torch.nn.functional.interpolate(latent_mask, size=(H * 8, W * 8), mode="bilinear").round()
+    latent_mask = torch.nn.functional.max_pool2d(latent_mask, (8, 8)).round().to(latent)
+
+    return latent, latent_mask
+
+
+class VAEApprox(torch.nn.Module):
+    def __init__(self):
+        super(VAEApprox, self).__init__()
+        self.conv1 = torch.nn.Conv2d(4, 8, (7, 7))
+        self.conv2 = torch.nn.Conv2d(8, 16, (5, 5))
+        self.conv3 = torch.nn.Conv2d(16, 32, (3, 3))
+        self.conv4 = torch.nn.Conv2d(32, 64, (3, 3))
+        self.conv5 = torch.nn.Conv2d(64, 32, (3, 3))
+        self.conv6 = torch.nn.Conv2d(32, 16, (3, 3))
+        self.conv7 = torch.nn.Conv2d(16, 8, (3, 3))
+        self.conv8 = torch.nn.Conv2d(8, 3, (3, 3))
+        self.current_type = None
+
+    def forward(self, x):
+        extra = 11
+        x = torch.nn.functional.interpolate(x, (x.shape[2] * 2, x.shape[3] * 2))
+        x = torch.nn.functional.pad(x, (extra, extra, extra, extra))
+        for layer in [self.conv1, self.conv2, self.conv3, self.conv4, self.conv5, self.conv6, self.conv7, self.conv8]:
+            x = layer(x)
+            x = torch.nn.functional.leaky_relu(x, 0.1)
+        return x
+
+
+VAE_approx_models = {}
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def get_previewer(model):
+    global VAE_approx_models
+
+    from modules.config import path_vae_approx
+    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:
+        VAE_approx_model = VAE_approx_models[vae_approx_filename]
+    else:
+        sd = torch.load(vae_approx_filename, map_location='cpu')
+        VAE_approx_model = VAEApprox()
+        VAE_approx_model.load_state_dict(sd)
+        del sd
+        VAE_approx_model.eval()
+
+        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(ldm_patched.modules.model_management.get_torch_device())
+        VAE_approx_models[vae_approx_filename] = VAE_approx_model
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def preview_function(x0, step, total_steps):
+        with torch.no_grad():
+            x_sample = x0.to(VAE_approx_model.current_type)
+            x_sample = VAE_approx_model(x_sample) * 127.5 + 127.5
+            x_sample = einops.rearrange(x_sample, 'b c h w -> b h w c')[0]
+            x_sample = x_sample.cpu().numpy().clip(0, 255).astype(np.uint8)
+            return x_sample
+
+    return preview_function
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sampler_name='dpmpp_2m_sde_gpu',
+             scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
+             force_full_denoise=False, callback_function=None, refiner=None, refiner_switch=-1,
+             previewer_start=None, previewer_end=None, sigmas=None, noise_mean=None):
+
+    if sigmas is not None:
+        sigmas = sigmas.clone().to(ldm_patched.modules.model_management.get_torch_device())
+
+    latent_image = latent["samples"]
+
+    if disable_noise:
+        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 = 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)
+
+    noise_mask = None
+    if "noise_mask" in latent:
+        noise_mask = latent["noise_mask"]
+
+    previewer = get_previewer(model)
+
+    if previewer_start is None:
+        previewer_start = 0
+
+    if previewer_end is None:
+        previewer_end = steps
+
+    def callback(step, x0, x, total_steps):
+        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)
+        if callback_function is not None:
+            callback_function(previewer_start + step, x0, x, previewer_end, y)
+
+    disable_pbar = False
+    modules.sample_hijack.current_refiner = refiner
+    modules.sample_hijack.refiner_switch_step = refiner_switch
+    ldm_patched.modules.samplers.sample = modules.sample_hijack.sample_hacked
+
+    try:
+        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
+    finally:
+        modules.sample_hijack.current_refiner = None
+
+    return out
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def pytorch_to_numpy(x):
+    return [np.clip(255. * y.cpu().numpy(), 0, 255).astype(np.uint8) for y in x]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def numpy_to_pytorch(x):
+    y = x.astype(np.float32) / 255.0
+    y = y[None]
+    y = np.ascontiguousarray(y.copy())
+    y = torch.from_numpy(y).float()
+    return y

modules/default_pipeline.py

@@ -0,0 +1,492 @@
+import modules.core as core
+import os
+import torch
+import modules.patch
+import modules.config
+import ldm_patched.modules.model_management
+import ldm_patched.modules.latent_formats
+import modules.inpaint_worker
+import extras.vae_interpose as vae_interpose
+from extras.expansion import FooocusExpansion
+
+from ldm_patched.modules.model_base import SDXL, SDXLRefiner
+from modules.sample_hijack import clip_separate
+
+
+model_base = core.StableDiffusionModel()
+model_refiner = core.StableDiffusionModel()
+
+final_expansion = None
+final_unet = None
+final_clip = None
+final_vae = None
+final_refiner_unet = None
+final_refiner_vae = None
+
+loaded_ControlNets = {}
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def refresh_controlnets(model_paths):
+    global loaded_ControlNets
+    cache = {}
+    for p in model_paths:
+        if p is not None:
+            if p in loaded_ControlNets:
+                cache[p] = loaded_ControlNets[p]
+            else:
+                cache[p] = core.load_controlnet(p)
+    loaded_ControlNets = cache
+    return
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def assert_model_integrity():
+    error_message = None
+
+    if not isinstance(model_base.unet_with_lora.model, SDXL):
+        error_message = 'You have selected base model other than SDXL. This is not supported yet.'
+
+    if error_message is not None:
+        raise NotImplementedError(error_message)
+
+    return True
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def refresh_base_model(name):
+    global model_base
+
+    filename = os.path.abspath(os.path.realpath(os.path.join(modules.config.path_checkpoints, name)))
+
+    if model_base.filename == filename:
+        return
+
+    model_base = core.StableDiffusionModel()
+    model_base = core.load_model(filename)
+    print(f'Base model loaded: {model_base.filename}')
+    return
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def refresh_refiner_model(name):
+    global model_refiner
+
+    filename = os.path.abspath(os.path.realpath(os.path.join(modules.config.path_checkpoints, name)))
+
+    if model_refiner.filename == filename:
+        return
+
+    model_refiner = core.StableDiffusionModel()
+
+    if name == 'None':
+        print(f'Refiner unloaded.')
+        return
+
+    model_refiner = core.load_model(filename)
+    print(f'Refiner model loaded: {model_refiner.filename}')
+
+    if isinstance(model_refiner.unet.model, SDXL):
+        model_refiner.clip = None
+        model_refiner.vae = None
+    elif isinstance(model_refiner.unet.model, SDXLRefiner):
+        model_refiner.clip = None
+        model_refiner.vae = None
+    else:
+        model_refiner.clip = None
+
+    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):
+    global model_base, model_refiner
+
+    if not isinstance(base_model_additional_loras, list):
+        base_model_additional_loras = []
+
+    model_base.refresh_loras(loras + base_model_additional_loras)
+    model_refiner.refresh_loras(loras)
+
+    return
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_encode_single(clip, text, verbose=False):
+    cached = clip.fcs_cond_cache.get(text, None)
+    if cached is not None:
+        if verbose:
+            print(f'[CLIP Cached] {text}')
+        return cached
+    tokens = clip.tokenize(text)
+    result = clip.encode_from_tokens(tokens, return_pooled=True)
+    clip.fcs_cond_cache[text] = result
+    if verbose:
+        print(f'[CLIP Encoded] {text}')
+    return result
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clone_cond(conds):
+    results = []
+
+    for c, p in conds:
+        p = p["pooled_output"]
+
+        if isinstance(c, torch.Tensor):
+            c = c.clone()
+
+        if isinstance(p, torch.Tensor):
+            p = p.clone()
+
+        results.append([c, {"pooled_output": p}])
+
+    return results
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_encode(texts, pool_top_k=1):
+    global final_clip
+
+    if final_clip is None:
+        return None
+    if not isinstance(texts, list):
+        return None
+    if len(texts) == 0:
+        return None
+
+    cond_list = []
+    pooled_acc = 0
+
+    for i, text in enumerate(texts):
+        cond, pooled = clip_encode_single(final_clip, text)
+        cond_list.append(cond)
+        if i < pool_top_k:
+            pooled_acc += pooled
+
+    return [[torch.cat(cond_list, dim=1), {"pooled_output": pooled_acc}]]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clear_all_caches():
+    final_clip.fcs_cond_cache = {}
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def prepare_text_encoder(async_call=True):
+    if async_call:
+        # TODO: make sure that this is always called in an async way so that users cannot feel it.
+        pass
+    assert_model_integrity()
+    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, use_synthetic_refiner=False):
+    global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, final_expansion
+
+    final_unet = None
+    final_clip = None
+    final_vae = None
+    final_refiner_unet = None
+    final_refiner_vae = None
+
+    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()
+
+    final_unet = model_base.unet_with_lora
+    final_clip = model_base.clip_with_lora
+    final_vae = model_base.vae
+
+    final_refiner_unet = model_refiner.unet_with_lora
+    final_refiner_vae = model_refiner.vae
+
+    if final_expansion is None:
+        final_expansion = FooocusExpansion()
+
+    prepare_text_encoder(async_call=True)
+    clear_all_caches()
+    return
+
+
+refresh_everything(
+    refiner_model_name=modules.config.default_refiner_model_name,
+    base_model_name=modules.config.default_base_model_name,
+    loras=modules.config.default_loras
+)
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def vae_parse(latent):
+    if final_refiner_vae is None:
+        return latent
+
+    result = vae_interpose.parse(latent["samples"])
+    return {'samples': result}
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def calculate_sigmas_all(sampler, model, scheduler, steps):
+    from ldm_patched.modules.samplers import calculate_sigmas_scheduler
+
+    discard_penultimate_sigma = False
+    if sampler in ['dpm_2', 'dpm_2_ancestral']:
+        steps += 1
+        discard_penultimate_sigma = True
+
+    sigmas = calculate_sigmas_scheduler(model, scheduler, steps)
+
+    if discard_penultimate_sigma:
+        sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+    return sigmas
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def calculate_sigmas(sampler, model, scheduler, steps, denoise):
+    if denoise is None or denoise > 0.9999:
+        sigmas = calculate_sigmas_all(sampler, model, scheduler, steps)
+    else:
+        new_steps = int(steps / denoise)
+        sigmas = calculate_sigmas_all(sampler, model, scheduler, new_steps)
+        sigmas = sigmas[-(steps + 1):]
+    return sigmas
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def get_candidate_vae(steps, switch, denoise=1.0, refiner_swap_method='joint'):
+    assert refiner_swap_method in ['joint', 'separate', '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:
+                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:
+        initial_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
+    else:
+        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()
+    sigma_min = float(sigma_min.cpu().numpy())
+    sigma_max = float(sigma_max.cpu().numpy())
+    print(f'[Sampler] sigma_min = {sigma_min}, sigma_max = {sigma_max}')
+
+    modules.patch.BrownianTreeNoiseSamplerPatched.global_init(
+        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=target_unet,
+            refiner=target_refiner_unet,
+            positive=positive_cond,
+            negative=negative_cond,
+            latent=initial_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,
+            refiner_switch=switch,
+            previewer_start=0,
+            previewer_end=steps,
+        )
+        decoded_latent = core.decode_vae(vae=target_vae, latent_image=sampled_latent, tiled=tiled)
+
+    if refiner_swap_method == 'separate':
+        sampled_latent = core.ksampler(
+            model=target_unet,
+            positive=positive_cond,
+            negative=negative_cond,
+            latent=initial_latent,
+            steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=False,
+            seed=image_seed,
+            denoise=denoise,
+            callback_function=callback,
+            cfg=cfg_scale,
+            sampler_name=sampler_name,
+            scheduler=scheduler_name,
+            previewer_start=0,
+            previewer_end=steps,
+        )
+        print('Refiner swapped by changing ksampler. Noise preserved.')
+
+        target_model = target_refiner_unet
+        if target_model is None:
+            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=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,
+            denoise=denoise,
+            callback_function=callback,
+            cfg=cfg_scale,
+            sampler_name=sampler_name,
+            scheduler=scheduler_name,
+            previewer_start=switch,
+            previewer_end=steps,
+        )
+
+        target_model = target_refiner_vae
+        if target_model is None:
+            target_model = target_vae
+        decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
+
+    if refiner_swap_method == 'vae':
+        modules.patch.eps_record = 'vae'
+
+        if modules.inpaint_worker.current_task is not None:
+            modules.inpaint_worker.current_task.unswap()
+
+        sampled_latent = core.ksampler(
+            model=target_unet,
+            positive=positive_cond,
+            negative=negative_cond,
+            latent=initial_latent,
+            steps=steps, start_step=0, last_step=switch, 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
+        )
+        print('Fooocus VAE-based swap.')
+
+        target_model = target_refiner_unet
+        if target_model is None:
+            target_model = target_unet
+            print('Use base model to refine itself - this may because of developer mode.')
+
+        sampled_latent = vae_parse(sampled_latent)
+
+        k_sigmas = 1.4
+        sigmas = calculate_sigmas(sampler=sampler_name,
+                                  scheduler=scheduler_name,
+                                  model=target_model.model,
+                                  steps=steps,
+                                  denoise=denoise)[switch:] * k_sigmas
+        len_sigmas = len(sigmas) - 1
+
+        noise_mean = torch.mean(modules.patch.eps_record, dim=1, keepdim=True)
+
+        if modules.inpaint_worker.current_task is not None:
+            modules.inpaint_worker.current_task.swap()
+
+        sampled_latent = core.ksampler(
+            model=target_model,
+            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,
+            denoise=denoise,
+            callback_function=callback,
+            cfg=cfg_scale,
+            sampler_name=sampler_name,
+            scheduler=scheduler_name,
+            previewer_start=switch,
+            previewer_end=steps,
+            sigmas=sigmas,
+            noise_mean=noise_mean
+        )
+
+        target_model = target_refiner_vae
+        if target_model is None:
+            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.patch.eps_record = None
+    return images

modules/flags.py

@@ -0,0 +1,44 @@
+disabled = 'Disabled'
+enabled = 'Enabled'
+subtle_variation = 'Vary (Subtle)'
+strong_variation = 'Vary (Strong)'
+upscale_15 = 'Upscale (1.5x)'
+upscale_2 = 'Upscale (2x)'
+upscale_fast = 'Upscale (Fast 2x)'
+
+uov_list = [
+    disabled, subtle_variation, strong_variation, upscale_15, upscale_2, upscale_fast
+]
+
+KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
+                  "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", "turbo"]
+SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
+
+sampler_list = SAMPLER_NAMES
+scheduler_list = SCHEDULER_NAMES
+
+cn_ip = "ImagePrompt"
+cn_ip_face = "FaceSwap"
+cn_canny = "PyraCanny"
+cn_cpds = "CPDS"
+
+ip_list = [cn_ip, cn_canny, cn_cpds, cn_ip_face]
+default_ip = cn_ip
+
+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 = ['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/gradio_hijack.py

@@ -0,0 +1,480 @@
+"""gr.Image() component."""
+
+from __future__ import annotations
+
+import warnings
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+import PIL
+import PIL.ImageOps
+import gradio.routes
+import importlib
+
+from gradio_client import utils as client_utils
+from gradio_client.documentation import document, set_documentation_group
+from gradio_client.serializing import ImgSerializable
+from PIL import Image as _Image  # using _ to minimize namespace pollution
+
+from gradio import processing_utils, utils
+from gradio.components.base import IOComponent, _Keywords, Block
+from gradio.deprecation import warn_style_method_deprecation
+from gradio.events import (
+    Changeable,
+    Clearable,
+    Editable,
+    EventListenerMethod,
+    Selectable,
+    Streamable,
+    Uploadable,
+)
+from gradio.interpretation import TokenInterpretable
+
+set_documentation_group("component")
+_Image.init()  # fixes https://github.com/gradio-app/gradio/issues/2843
+
+
+@document()
+class Image(
+    Editable,
+    Clearable,
+    Changeable,
+    Streamable,
+    Selectable,
+    Uploadable,
+    IOComponent,
+    ImgSerializable,
+    TokenInterpretable,
+):
+    """
+    Creates an image component that can be used to upload/draw images (as an input) or display images (as an output).
+    Preprocessing: passes the uploaded image as a {numpy.array}, {PIL.Image} or {str} filepath depending on `type` -- unless `tool` is `sketch` AND source is one of `upload` or `webcam`. In these cases, a {dict} with keys `image` and `mask` is passed, and the format of the corresponding values depends on `type`.
+    Postprocessing: expects a {numpy.array}, {PIL.Image} or {str} or {pathlib.Path} filepath to an image and displays the image.
+    Examples-format: a {str} filepath to a local file that contains the image.
+    Demos: image_mod, image_mod_default_image
+    Guides: image-classification-in-pytorch, image-classification-in-tensorflow, image-classification-with-vision-transformers, building-a-pictionary_app, create-your-own-friends-with-a-gan
+    """
+
+    def __init__(
+        self,
+        value: str | _Image.Image | np.ndarray | None = None,
+        *,
+        shape: tuple[int, int] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        image_mode: Literal[
+            "1", "L", "P", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "I", "F"
+        ] = "RGB",
+        invert_colors: bool = False,
+        source: Literal["upload", "webcam", "canvas"] = "upload",
+        tool: Literal["editor", "select", "sketch", "color-sketch"] | None = None,
+        type: Literal["numpy", "pil", "filepath"] = "numpy",
+        label: str | None = None,
+        every: float | None = None,
+        show_label: bool | None = None,
+        show_download_button: bool = True,
+        container: bool = True,
+        scale: int | None = None,
+        min_width: int = 160,
+        interactive: bool | None = None,
+        visible: bool = True,
+        streaming: bool = False,
+        elem_id: str | None = None,
+        elem_classes: list[str] | str | None = None,
+        mirror_webcam: bool = True,
+        brush_radius: float | None = None,
+        brush_color: str = "#000000",
+        mask_opacity: float = 0.7,
+        show_share_button: bool | None = None,
+        **kwargs,
+    ):
+        """
+        Parameters:
+            value: A PIL Image, numpy array, path or URL for the default value that Image component is going to take. If callable, the function will be called whenever the app loads to set the initial value of the component.
+            shape: (width, height) shape to crop and resize image when passed to function. If None, matches input image size. Pass None for either width or height to only crop and resize the other.
+            height: Height of the displayed image in pixels.
+            width: Width of the displayed image in pixels.
+            image_mode: "RGB" if color, or "L" if black and white. See https://pillow.readthedocs.io/en/stable/handbook/concepts.html for other supported image modes and their meaning.
+            invert_colors: whether to invert the image as a preprocessing step.
+            source: Source of image. "upload" creates a box where user can drop an image file, "webcam" allows user to take snapshot from their webcam, "canvas" defaults to a white image that can be edited and drawn upon with tools.
+            tool: Tools used for editing. "editor" allows a full screen editor (and is the default if source is "upload" or "webcam"), "select" provides a cropping and zoom tool, "sketch" allows you to create a binary sketch (and is the default if source="canvas"), and "color-sketch" allows you to created a sketch in different colors. "color-sketch" can be used with source="upload" or "webcam" to allow sketching on an image. "sketch" can also be used with "upload" or "webcam" to create a mask over an image and in that case both the image and mask are passed into the function as a dictionary with keys "image" and "mask" respectively.
+            type: The format the image is converted to before being passed into the prediction function. "numpy" converts the image to a numpy array with shape (height, width, 3) and values from 0 to 255, "pil" converts the image to a PIL image object, "filepath" passes a str path to a temporary file containing the image.
+            label: component name in interface.
+            every: If `value` is a callable, run the function 'every' number of seconds while the client connection is open. Has no effect otherwise. Queue must be enabled. The event can be accessed (e.g. to cancel it) via this component's .load_event attribute.
+            show_label: if True, will display label.
+            show_download_button: If True, will display button to download image.
+            container: If True, will place the component in a container - providing some extra padding around the border.
+            scale: relative width compared to adjacent Components in a Row. For example, if Component A has scale=2, and Component B has scale=1, A will be twice as wide as B. Should be an integer.
+            min_width: minimum pixel width, will wrap if not sufficient screen space to satisfy this value. If a certain scale value results in this Component being narrower than min_width, the min_width parameter will be respected first.
+            interactive: if True, will allow users to upload and edit an image; if False, can only be used to display images. If not provided, this is inferred based on whether the component is used as an input or output.
+            visible: If False, component will be hidden.
+            streaming: If True when used in a `live` interface, will automatically stream webcam feed. Only valid is source is 'webcam'.
+            elem_id: An optional string that is assigned as the id of this component in the HTML DOM. Can be used for targeting CSS styles.
+            elem_classes: An optional list of strings that are assigned as the classes of this component in the HTML DOM. Can be used for targeting CSS styles.
+            mirror_webcam: If True webcam will be mirrored. Default is True.
+            brush_radius: Size of the brush for Sketch. Default is None which chooses a sensible default
+            brush_color: Color of the brush for Sketch as hex string. Default is "#000000".
+            mask_opacity: Opacity of mask drawn on image, as a value between 0 and 1.
+            show_share_button: If True, will show a share icon in the corner of the component that allows user to share outputs to Hugging Face Spaces Discussions. If False, icon does not appear. If set to None (default behavior), then the icon appears if this Gradio app is launched on Spaces, but not otherwise.
+        """
+        self.brush_radius = brush_radius
+        self.brush_color = brush_color
+        self.mask_opacity = mask_opacity
+        self.mirror_webcam = mirror_webcam
+        valid_types = ["numpy", "pil", "filepath"]
+        if type not in valid_types:
+            raise ValueError(
+                f"Invalid value for parameter `type`: {type}. Please choose from one of: {valid_types}"
+            )
+        self.type = type
+        self.shape = shape
+        self.height = height
+        self.width = width
+        self.image_mode = image_mode
+        valid_sources = ["upload", "webcam", "canvas"]
+        if source not in valid_sources:
+            raise ValueError(
+                f"Invalid value for parameter `source`: {source}. Please choose from one of: {valid_sources}"
+            )
+        self.source = source
+        if tool is None:
+            self.tool = "sketch" if source == "canvas" else "editor"
+        else:
+            self.tool = tool
+        self.invert_colors = invert_colors
+        self.streaming = streaming
+        self.show_download_button = show_download_button
+        if streaming and source != "webcam":
+            raise ValueError("Image streaming only available if source is 'webcam'.")
+        self.select: EventListenerMethod
+        """
+        Event listener for when the user clicks on a pixel within the image.
+        Uses event data gradio.SelectData to carry `index` to refer to the [x, y] coordinates of the clicked pixel.
+        See EventData documentation on how to use this event data.
+        """
+        self.show_share_button = (
+            (utils.get_space() is not None)
+            if show_share_button is None
+            else show_share_button
+        )
+        IOComponent.__init__(
+            self,
+            label=label,
+            every=every,
+            show_label=show_label,
+            container=container,
+            scale=scale,
+            min_width=min_width,
+            interactive=interactive,
+            visible=visible,
+            elem_id=elem_id,
+            elem_classes=elem_classes,
+            value=value,
+            **kwargs,
+        )
+        TokenInterpretable.__init__(self)
+
+    def get_config(self):
+        return {
+            "image_mode": self.image_mode,
+            "shape": self.shape,
+            "height": self.height,
+            "width": self.width,
+            "source": self.source,
+            "tool": self.tool,
+            "value": self.value,
+            "streaming": self.streaming,
+            "mirror_webcam": self.mirror_webcam,
+            "brush_radius": self.brush_radius,
+            "brush_color": self.brush_color,
+            "mask_opacity": self.mask_opacity,
+            "selectable": self.selectable,
+            "show_share_button": self.show_share_button,
+            "show_download_button": self.show_download_button,
+            **IOComponent.get_config(self),
+        }
+
+    @staticmethod
+    def update(
+        value: Any | Literal[_Keywords.NO_VALUE] | None = _Keywords.NO_VALUE,
+        height: int | None = None,
+        width: int | None = None,
+        label: str | None = None,
+        show_label: bool | None = None,
+        show_download_button: bool | None = None,
+        container: bool | None = None,
+        scale: int | None = None,
+        min_width: int | None = None,
+        interactive: bool | None = None,
+        visible: bool | None = None,
+        brush_radius: float | None = None,
+        brush_color: str | None = None,
+        mask_opacity: float | None = None,
+        show_share_button: bool | None = None,
+    ):
+        return {
+            "height": height,
+            "width": width,
+            "label": label,
+            "show_label": show_label,
+            "show_download_button": show_download_button,
+            "container": container,
+            "scale": scale,
+            "min_width": min_width,
+            "interactive": interactive,
+            "visible": visible,
+            "value": value,
+            "brush_radius": brush_radius,
+            "brush_color": brush_color,
+            "mask_opacity": mask_opacity,
+            "show_share_button": show_share_button,
+            "__type__": "update",
+        }
+
+    def _format_image(
+        self, im: _Image.Image | None
+    ) -> np.ndarray | _Image.Image | str | None:
+        """Helper method to format an image based on self.type"""
+        if im is None:
+            return im
+        fmt = im.format
+        if self.type == "pil":
+            return im
+        elif self.type == "numpy":
+            return np.array(im)
+        elif self.type == "filepath":
+            path = self.pil_to_temp_file(
+                im, dir=self.DEFAULT_TEMP_DIR, format=fmt or "png"
+            )
+            self.temp_files.add(path)
+            return path
+        else:
+            raise ValueError(
+                "Unknown type: "
+                + str(self.type)
+                + ". Please choose from: 'numpy', 'pil', 'filepath'."
+            )
+
+    def preprocess(
+        self, x: str | dict[str, str]
+    ) -> np.ndarray | _Image.Image | str | dict | None:
+        """
+        Parameters:
+            x: base64 url data, or (if tool == "sketch") a dict of image and mask base64 url data
+        Returns:
+            image in requested format, or (if tool == "sketch") a dict of image and mask in requested format
+        """
+        if x is None:
+            return x
+
+        mask = None
+
+        if self.tool == "sketch" and self.source in ["upload", "webcam"]:
+            if isinstance(x, dict):
+                x, mask = x["image"], x["mask"]
+
+        assert isinstance(x, str)
+        im = processing_utils.decode_base64_to_image(x)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            im = im.convert(self.image_mode)
+        if self.shape is not None:
+            im = processing_utils.resize_and_crop(im, self.shape)
+        if self.invert_colors:
+            im = PIL.ImageOps.invert(im)
+        if (
+            self.source == "webcam"
+            and self.mirror_webcam is True
+            and self.tool != "color-sketch"
+        ):
+            im = PIL.ImageOps.mirror(im)
+
+        if self.tool == "sketch" and self.source in ["upload", "webcam"]:
+            if mask is not None:
+                mask_im = processing_utils.decode_base64_to_image(mask)
+                if mask_im.mode == "RGBA":  # whiten any opaque pixels in the mask
+                    alpha_data = mask_im.getchannel("A").convert("L")
+                    mask_im = _Image.merge("RGB", [alpha_data, alpha_data, alpha_data])
+                return {
+                    "image": self._format_image(im),
+                    "mask": self._format_image(mask_im),
+                }
+            else:
+                return {
+                    "image": self._format_image(im),
+                    "mask": None,
+                }
+
+        return self._format_image(im)
+
+    def postprocess(
+        self, y: np.ndarray | _Image.Image | str | Path | None
+    ) -> str | None:
+        """
+        Parameters:
+            y: image as a numpy array, PIL Image, string/Path filepath, or string URL
+        Returns:
+            base64 url data
+        """
+        if y is None:
+            return None
+        if isinstance(y, np.ndarray):
+            return processing_utils.encode_array_to_base64(y)
+        elif isinstance(y, _Image.Image):
+            return processing_utils.encode_pil_to_base64(y)
+        elif isinstance(y, (str, Path)):
+            return client_utils.encode_url_or_file_to_base64(y)
+        else:
+            raise ValueError("Cannot process this value as an Image")
+
+    def set_interpret_parameters(self, segments: int = 16):
+        """
+        Calculates interpretation score of image subsections by splitting the image into subsections, then using a "leave one out" method to calculate the score of each subsection by whiting out the subsection and measuring the delta of the output value.
+        Parameters:
+            segments: Number of interpretation segments to split image into.
+        """
+        self.interpretation_segments = segments
+        return self
+
+    def _segment_by_slic(self, x):
+        """
+        Helper method that segments an image into superpixels using slic.
+        Parameters:
+            x: base64 representation of an image
+        """
+        x = processing_utils.decode_base64_to_image(x)
+        if self.shape is not None:
+            x = processing_utils.resize_and_crop(x, self.shape)
+        resized_and_cropped_image = np.array(x)
+        try:
+            from skimage.segmentation import slic
+        except (ImportError, ModuleNotFoundError) as err:
+            raise ValueError(
+                "Error: running this interpretation for images requires scikit-image, please install it first."
+            ) from err
+        try:
+            segments_slic = slic(
+                resized_and_cropped_image,
+                self.interpretation_segments,
+                compactness=10,
+                sigma=1,
+                start_label=1,
+            )
+        except TypeError:  # For skimage 0.16 and older
+            segments_slic = slic(
+                resized_and_cropped_image,
+                self.interpretation_segments,
+                compactness=10,
+                sigma=1,
+            )
+        return segments_slic, resized_and_cropped_image
+
+    def tokenize(self, x):
+        """
+        Segments image into tokens, masks, and leave-one-out-tokens
+        Parameters:
+            x: base64 representation of an image
+        Returns:
+            tokens: list of tokens, used by the get_masked_input() method
+            leave_one_out_tokens: list of left-out tokens, used by the get_interpretation_neighbors() method
+            masks: list of masks, used by the get_interpretation_neighbors() method
+        """
+        segments_slic, resized_and_cropped_image = self._segment_by_slic(x)
+        tokens, masks, leave_one_out_tokens = [], [], []
+        replace_color = np.mean(resized_and_cropped_image, axis=(0, 1))
+        for segment_value in np.unique(segments_slic):
+            mask = segments_slic == segment_value
+            image_screen = np.copy(resized_and_cropped_image)
+            image_screen[segments_slic == segment_value] = replace_color
+            leave_one_out_tokens.append(
+                processing_utils.encode_array_to_base64(image_screen)
+            )
+            token = np.copy(resized_and_cropped_image)
+            token[segments_slic != segment_value] = 0
+            tokens.append(token)
+            masks.append(mask)
+        return tokens, leave_one_out_tokens, masks
+
+    def get_masked_inputs(self, tokens, binary_mask_matrix):
+        masked_inputs = []
+        for binary_mask_vector in binary_mask_matrix:
+            masked_input = np.zeros_like(tokens[0], dtype=int)
+            for token, b in zip(tokens, binary_mask_vector):
+                masked_input = masked_input + token * int(b)
+            masked_inputs.append(processing_utils.encode_array_to_base64(masked_input))
+        return masked_inputs
+
+    def get_interpretation_scores(
+        self, x, neighbors, scores, masks, tokens=None, **kwargs
+    ) -> list[list[float]]:
+        """
+        Returns:
+            A 2D array representing the interpretation score of each pixel of the image.
+        """
+        x = processing_utils.decode_base64_to_image(x)
+        if self.shape is not None:
+            x = processing_utils.resize_and_crop(x, self.shape)
+        x = np.array(x)
+        output_scores = np.zeros((x.shape[0], x.shape[1]))
+
+        for score, mask in zip(scores, masks):
+            output_scores += score * mask
+
+        max_val, min_val = np.max(output_scores), np.min(output_scores)
+        if max_val > 0:
+            output_scores = (output_scores - min_val) / (max_val - min_val)
+        return output_scores.tolist()
+
+    def style(self, *, height: int | None = None, width: int | None = None, **kwargs):
+        """
+        This method is deprecated. Please set these arguments in the constructor instead.
+        """
+        warn_style_method_deprecation()
+        if height is not None:
+            self.height = height
+        if width is not None:
+            self.width = width
+        return self
+
+    def check_streamable(self):
+        if self.source != "webcam":
+            raise ValueError("Image streaming only available if source is 'webcam'.")
+
+    def as_example(self, input_data: str | None) -> str:
+        if input_data is None:
+            return ""
+        elif (
+            self.root_url
+        ):  # If an externally hosted image, don't convert to absolute path
+            return input_data
+        return str(utils.abspath(input_data))
+
+
+all_components = []
+
+if not hasattr(Block, 'original__init__'):
+    Block.original_init = Block.__init__
+
+
+def blk_ini(self, *args, **kwargs):
+    all_components.append(self)
+    return Block.original_init(self, *args, **kwargs)
+
+
+Block.__init__ = blk_ini
+
+
+gradio.routes.asyncio = importlib.reload(gradio.routes.asyncio)
+
+if not hasattr(gradio.routes.asyncio, 'original_wait_for'):
+    gradio.routes.asyncio.original_wait_for = gradio.routes.asyncio.wait_for
+
+
+def patched_wait_for(fut, timeout):
+    del timeout
+    return gradio.routes.asyncio.original_wait_for(fut, timeout=65535)
+
+
+gradio.routes.asyncio.wait_for = patched_wait_for
+

modules/html.py

@@ -0,0 +1,128 @@
+css = '''
+.loader-container {
+  display: flex; /* Use flex to align items horizontally */
+  align-items: center; /* Center items vertically within the container */
+  white-space: nowrap; /* Prevent line breaks within the container */
+}
+
+.loader {
+  border: 8px solid #f3f3f3; /* Light grey */
+  border-top: 8px solid #3498db; /* Blue */
+  border-radius: 50%;
+  width: 30px;
+  height: 30px;
+  animation: spin 2s linear infinite;
+}
+
+@keyframes spin {
+  0% { transform: rotate(0deg); }
+  100% { transform: rotate(360deg); }
+}
+
+/* Style the progress bar */
+progress {
+  appearance: none; /* Remove default styling */
+  height: 20px; /* Set the height of the progress bar */
+  border-radius: 5px; /* Round the corners of the progress bar */
+  background-color: #f3f3f3; /* Light grey background */
+  width: 100%;
+}
+
+/* Style the progress bar container */
+.progress-container {
+  margin-left: 20px;
+  margin-right: 20px;
+  flex-grow: 1; /* Allow the progress container to take up remaining space */
+}
+
+/* Set the color of the progress bar fill */
+progress::-webkit-progress-value {
+  background-color: #3498db; /* Blue color for the fill */
+}
+
+progress::-moz-progress-bar {
+  background-color: #3498db; /* Blue color for the fill in Firefox */
+}
+
+/* Style the text on the progress bar */
+progress::after {
+  content: attr(value '%'); /* Display the progress value followed by '%' */
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  color: white; /* Set text color */
+  font-size: 14px; /* Set font size */
+}
+
+/* Style other texts */
+.loader-container > span {
+  margin-left: 5px; /* Add spacing between the progress bar and the text */
+}
+
+.progress-bar > .generating {
+  display: none !important;
+}
+
+.progress-bar{
+  height: 30px !important;
+}
+
+.type_row{
+  height: 80px !important;
+}
+
+.type_row_half{
+  height: 32px !important;
+}
+
+.scroll-hide{
+  resize: none !important;
+}
+
+.refresh_button{
+  border: none !important;
+  background: none !important;
+  font-size: none !important;
+  box-shadow: none !important;
+}
+
+.advanced_check_row{
+  width: 250px !important;
+}
+
+.min_check{
+  min-width: min(1px, 100%) !important;
+}
+
+.resizable_area {
+  resize: vertical;
+  overflow: auto !important;
+}
+
+.aspect_ratios label {
+    width: 140px !important;
+}
+
+.aspect_ratios label span {
+    white-space: nowrap !important;
+}
+
+.aspect_ratios label input {
+    margin-left: -5px !important;
+}
+
+'''
+progress_html = '''
+<div class="loader-container">
+  <div class="loader"></div>
+  <div class="progress-container">
+    <progress value="*number*" max="100"></progress>
+  </div>
+  <span>*text*</span>
+</div>
+'''
+
+
+def make_progress_html(number, text):
+    return progress_html.replace('*number*', str(number)).replace('*text*', text)

modules/inpaint_worker.py

@@ -0,0 +1,257 @@
+import torch
+import numpy as np
+
+from PIL import Image, ImageFilter
+from modules.util import resample_image, set_image_shape_ceil, get_image_shape_ceil
+from modules.upscaler import perform_upscale
+
+
+inpaint_head_model = None
+
+
+class InpaintHead(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.head = torch.nn.Parameter(torch.empty(size=(320, 5, 3, 3), device='cpu'))
+
+    def __call__(self, x):
+        x = torch.nn.functional.pad(x, (1, 1, 1, 1), "replicate")
+        return torch.nn.functional.conv2d(input=x, weight=self.head)
+
+
+current_task = None
+
+
+def box_blur(x, k):
+    x = Image.fromarray(x)
+    x = x.filter(ImageFilter.BoxBlur(k))
+    return np.array(x)
+
+
+def max33(x):
+    x = Image.fromarray(x)
+    x = x.filter(ImageFilter.MaxFilter(3))
+    return np.array(x)
+
+
+def morphological_open(x):
+    x_int32 = np.zeros_like(x).astype(np.int32)
+    x_int32[x > 127] = 256
+    for _ in range(32):
+        maxed = max33(x_int32) - 8
+        x_int32 = np.maximum(maxed, x_int32)
+    return x_int32.clip(0, 255).astype(np.uint8)
+
+
+def up255(x, t=0):
+    y = np.zeros_like(x).astype(np.uint8)
+    y[x > t] = 255
+    return y
+
+
+def imsave(x, path):
+    x = Image.fromarray(x)
+    x.save(path)
+
+
+def regulate_abcd(x, a, b, c, d):
+    H, W = x.shape[:2]
+    if a < 0:
+        a = 0
+    if a > H:
+        a = H
+    if b < 0:
+        b = 0
+    if b > H:
+        b = H
+    if c < 0:
+        c = 0
+    if c > W:
+        c = W
+    if d < 0:
+        d = 0
+    if d > W:
+        d = W
+    return int(a), int(b), int(c), int(d)
+
+
+def compute_initial_abcd(x):
+    indices = np.where(x)
+    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 = int(max(abm, cdm) * 1.15)
+    a = abp - l
+    b = abp + l + 1
+    c = 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, k):
+    k = float(k)
+    assert 0.0 <= k <= 1.0
+
+    H, W = x.shape[:2]
+    if k == 1.0:
+        return 0, H, 0, W
+    while True:
+        if b - a >= H * k and d - c >= W * k:
+            break
+
+        add_h = (b - a) < (d - c)
+        add_w = not add_h
+
+        if b - a == H:
+            add_w = True
+
+        if d - c == W:
+            add_h = True
+
+        if add_h:
+            a -= 1
+            b += 1
+
+        if add_w:
+            c -= 1
+            d += 1
+
+        a, b, c, d = regulate_abcd(x, a, b, c, d)
+    return a, b, c, d
+
+
+def fooocus_fill(image, mask):
+    current_image = image.copy()
+    raw_image = image.copy()
+    area = np.where(mask < 127)
+    store = raw_image[area]
+
+    for k, repeats in [(512, 2), (256, 2), (128, 4), (64, 4), (33, 8), (15, 8), (5, 16), (3, 16)]:
+        for _ in range(repeats):
+            current_image = box_blur(current_image, k)
+            current_image[area] = store
+
+    return current_image
+
+
+class InpaintWorker:
+    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, 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)
+
+        # compute filling
+        if use_fill:
+            self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask)
+
+        # soft pixels
+        self.mask = morphological_open(mask)
+        self.image = image
+
+        # ending
+        self.latent = None
+        self.latent_after_swap = None
+        self.swapped = False
+        self.latent_mask = None
+        self.inpaint_head_feature = None
+        return
+
+    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
+
+    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_model.load_state_dict(sd)
+
+        feed = torch.cat([
+            inpaint_latent_mask,
+            model.model.process_latent_in(inpaint_latent)
+        ], dim=1)
+
+        inpaint_head_model.to(device=feed.device, dtype=feed.dtype)
+        inpaint_head_feature = inpaint_head_model(feed)
+
+        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:
+            return
+
+        if self.latent is None:
+            return
+
+        if self.latent_after_swap is None:
+            return
+
+        self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
+        self.swapped = True
+        return
+
+    def unswap(self):
+        if not self.swapped:
+            return
+
+        if self.latent is None:
+            return
+
+        if self.latent_after_swap is None:
+            return
+
+        self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
+        self.swapped = False
+        return
+
+    def color_correction(self, img):
+        fg = img.astype(np.float32)
+        bg = self.image.copy().astype(np.float32)
+        w = self.mask[:, :, None].astype(np.float32) / 255.0
+        y = fg * w + bg * (1 - w)
+        return y.clip(0, 255).astype(np.uint8)
+
+    def post_process(self, img):
+        a, b, c, d = self.interested_area
+        content = resample_image(img, d - c, b - a)
+        result = self.image.copy()
+        result[a:b, c:d] = content
+        result = self.color_correction(result)
+        return result
+
+    def visualize_mask_processing(self):
+        return [self.interested_fill, self.interested_mask, self.interested_image]
+

modules/launch_util.py

@@ -0,0 +1,107 @@
+import os
+import importlib
+import importlib.util
+import subprocess
+import sys
+import re
+import logging
+
+
+logging.getLogger("torch.distributed.nn").setLevel(logging.ERROR)  # sshh...
+logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
+
+re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
+
+python = sys.executable
+default_command_live = (os.environ.get('LAUNCH_LIVE_OUTPUT') == "1")
+index_url = os.environ.get('INDEX_URL', "")
+
+modules_path = os.path.dirname(os.path.realpath(__file__))
+script_path = os.path.dirname(modules_path)
+
+
+def is_installed(package):
+    try:
+        spec = importlib.util.find_spec(package)
+    except ModuleNotFoundError:
+        return False
+
+    return spec is not None
+
+
+def run(command, desc=None, errdesc=None, custom_env=None, live: bool = default_command_live) -> str:
+    if desc is not None:
+        print(desc)
+
+    run_kwargs = {
+        "args": command,
+        "shell": True,
+        "env": os.environ if custom_env is None else custom_env,
+        "encoding": 'utf8',
+        "errors": 'ignore',
+    }
+
+    if not live:
+        run_kwargs["stdout"] = run_kwargs["stderr"] = subprocess.PIPE
+
+    result = subprocess.run(**run_kwargs)
+
+    if result.returncode != 0:
+        error_bits = [
+            f"{errdesc or 'Error running command'}.",
+            f"Command: {command}",
+            f"Error code: {result.returncode}",
+        ]
+        if result.stdout:
+            error_bits.append(f"stdout: {result.stdout}")
+        if result.stderr:
+            error_bits.append(f"stderr: {result.stderr}")
+        raise RuntimeError("\n".join(error_bits))
+
+    return (result.stdout or "")
+
+
+def run_pip(command, desc=None, live=default_command_live):
+    try:
+        index_url_line = f' --index-url {index_url}' if index_url != '' else ''
+        return run(f'"{python}" -m pip {command} --prefer-binary{index_url_line}', desc=f"Installing {desc}",
+                   errdesc=f"Couldn't install {desc}", live=live)
+    except Exception as e:
+        print(e)
+        print(f'CMD Failed {desc}: {command}')
+        return None
+
+
+def requirements_met(requirements_file):
+    """
+    Does a simple parse of a requirements.txt file to determine if all rerqirements in it
+    are already installed. Returns True if so, False if not installed or parsing fails.
+    """
+
+    import importlib.metadata
+    import packaging.version
+
+    with open(requirements_file, "r", encoding="utf8") as file:
+        for line in file:
+            if line.strip() == "":
+                continue
+
+            m = re.match(re_requirement, line)
+            if m is None:
+                return False
+
+            package = m.group(1).strip()
+            version_required = (m.group(2) or "").strip()
+
+            if version_required == "":
+                continue
+
+            try:
+                version_installed = importlib.metadata.version(package)
+            except Exception:
+                return False
+
+            if packaging.version.parse(version_required) != packaging.version.parse(version_installed):
+                return False
+
+    return True

modules/localization.py

@@ -0,0 +1,60 @@
+import json
+import os
+
+
+current_translation = {}
+localization_root = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'language')
+
+
+def localization_js(filename):
+    global current_translation
+
+    if isinstance(filename, str):
+        full_name = os.path.abspath(os.path.join(localization_root, filename + '.json'))
+        if os.path.exists(full_name):
+            try:
+                with open(full_name, encoding='utf-8') as f:
+                    current_translation = json.load(f)
+                    assert isinstance(current_translation, dict)
+                    for k, v in current_translation.items():
+                        assert isinstance(k, str)
+                        assert isinstance(v, str)
+            except Exception as e:
+                print(str(e))
+                print(f'Failed to load localization file {full_name}')
+
+    # current_translation = {k: 'XXX' for k in current_translation.keys()}  # use this to see if all texts are covered
+
+    return f"window.localization = {json.dumps(current_translation)}"
+
+
+def dump_english_config(components):
+    all_texts = []
+    for c in components:
+        label = getattr(c, 'label', None)
+        value = getattr(c, 'value', None)
+        choices = getattr(c, 'choices', None)
+        info = getattr(c, 'info', None)
+
+        if isinstance(label, str):
+            all_texts.append(label)
+        if isinstance(value, str):
+            all_texts.append(value)
+        if isinstance(info, str):
+            all_texts.append(info)
+        if isinstance(choices, list):
+            for x in choices:
+                if isinstance(x, str):
+                    all_texts.append(x)
+                if isinstance(x, tuple):
+                    for y in x:
+                        if isinstance(y, str):
+                            all_texts.append(y)
+
+    config_dict = {k: k for k in all_texts if k != "" and 'progress-container' not in k}
+    full_name = os.path.abspath(os.path.join(localization_root, 'en.json'))
+
+    with open(full_name, "w", encoding="utf-8") as json_file:
+        json.dump(config_dict, json_file, indent=4)
+
+    return

modules/lora.py

@@ -0,0 +1,152 @@
+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] = ('fooocus', lora[real_load_key])
+            loaded_keys.add(real_load_key)
+            continue
+
+        alpha_name = "{}.alpha".format(x)
+        alpha = None
+        if alpha_name in lora.keys():
+            alpha = lora[alpha_name].item()
+            loaded_keys.add(alpha_name)
+
+        regular_lora = "{}.lora_up.weight".format(x)
+        diffusers_lora = "{}_lora.up.weight".format(x)
+        transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
+        A_name = None
+
+        if regular_lora in lora.keys():
+            A_name = regular_lora
+            B_name = "{}.lora_down.weight".format(x)
+            mid_name = "{}.lora_mid.weight".format(x)
+        elif diffusers_lora in lora.keys():
+            A_name = diffusers_lora
+            B_name = "{}_lora.down.weight".format(x)
+            mid_name = None
+        elif transformers_lora in lora.keys():
+            A_name = transformers_lora
+            B_name ="{}.lora_linear_layer.down.weight".format(x)
+            mid_name = None
+
+        if A_name is not None:
+            mid = None
+            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", (lora[A_name], lora[B_name], alpha, mid))
+            loaded_keys.add(A_name)
+            loaded_keys.add(B_name)
+
+
+        ######## loha
+        hada_w1_a_name = "{}.hada_w1_a".format(x)
+        hada_w1_b_name = "{}.hada_w1_b".format(x)
+        hada_w2_a_name = "{}.hada_w2_a".format(x)
+        hada_w2_b_name = "{}.hada_w2_b".format(x)
+        hada_t1_name = "{}.hada_t1".format(x)
+        hada_t2_name = "{}.hada_t2".format(x)
+        if hada_w1_a_name in lora.keys():
+            hada_t1 = None
+            hada_t2 = None
+            if hada_t1_name in lora.keys():
+                hada_t1 = lora[hada_t1_name]
+                hada_t2 = lora[hada_t2_name]
+                loaded_keys.add(hada_t1_name)
+                loaded_keys.add(hada_t2_name)
+
+            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)
+            loaded_keys.add(hada_w2_b_name)
+
+
+        ######## lokr
+        lokr_w1_name = "{}.lokr_w1".format(x)
+        lokr_w2_name = "{}.lokr_w2".format(x)
+        lokr_w1_a_name = "{}.lokr_w1_a".format(x)
+        lokr_w1_b_name = "{}.lokr_w1_b".format(x)
+        lokr_t2_name = "{}.lokr_t2".format(x)
+        lokr_w2_a_name = "{}.lokr_w2_a".format(x)
+        lokr_w2_b_name = "{}.lokr_w2_b".format(x)
+
+        lokr_w1 = None
+        if lokr_w1_name in lora.keys():
+            lokr_w1 = lora[lokr_w1_name]
+            loaded_keys.add(lokr_w1_name)
+
+        lokr_w2 = None
+        if lokr_w2_name in lora.keys():
+            lokr_w2 = lora[lokr_w2_name]
+            loaded_keys.add(lokr_w2_name)
+
+        lokr_w1_a = None
+        if lokr_w1_a_name in lora.keys():
+            lokr_w1_a = lora[lokr_w1_a_name]
+            loaded_keys.add(lokr_w1_a_name)
+
+        lokr_w1_b = None
+        if lokr_w1_b_name in lora.keys():
+            lokr_w1_b = lora[lokr_w1_b_name]
+            loaded_keys.add(lokr_w1_b_name)
+
+        lokr_w2_a = None
+        if lokr_w2_a_name in lora.keys():
+            lokr_w2_a = lora[lokr_w2_a_name]
+            loaded_keys.add(lokr_w2_a_name)
+
+        lokr_w2_b = None
+        if lokr_w2_b_name in lora.keys():
+            lokr_w2_b = lora[lokr_w2_b_name]
+            loaded_keys.add(lokr_w2_b_name)
+
+        lokr_t2 = None
+        if lokr_t2_name in lora.keys():
+            lokr_t2 = lora[lokr_t2_name]
+            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", (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)
+        w_norm = lora.get(w_norm_name, None)
+        b_norm = lora.get(b_norm_name, None)
+
+        if w_norm is not None:
+            loaded_keys.add(w_norm_name)
+            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")])] = ("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", (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", (diff_bias,))
+            loaded_keys.add(diff_bias_name)
+
+    remaining_dict = {x: y for x, y in lora.items() if x not in loaded_keys}
+    return patch_dict, remaining_dict

modules/meta_parser.py

@@ -0,0 +1,148 @@
+import json
+import gradio as gr
+import modules.config
+
+
+def load_parameter_button_click(raw_prompt_txt, is_generating):
+    loaded_parameter_dict = json.loads(raw_prompt_txt)
+    assert isinstance(loaded_parameter_dict, dict)
+
+    results = [True, 1]
+
+    try:
+        h = loaded_parameter_dict.get('Prompt', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Negative Prompt', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Styles', None)
+        h = eval(h)
+        assert isinstance(h, list)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Performance', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Resolution', None)
+        width, height = eval(h)
+        formatted = modules.config.add_ratio(f'{width}*{height}')
+        if formatted in modules.config.available_aspect_ratios:
+            results.append(formatted)
+            results.append(-1)
+            results.append(-1)
+        else:
+            results.append(gr.update())
+            results.append(width)
+            results.append(height)
+    except:
+        results.append(gr.update())
+        results.append(gr.update())
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Sharpness', None)
+        assert h is not None
+        h = float(h)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Guidance Scale', None)
+        assert h is not None
+        h = float(h)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('ADM Guidance', None)
+        p, n, e = eval(h)
+        results.append(float(p))
+        results.append(float(n))
+        results.append(float(e))
+    except:
+        results.append(gr.update())
+        results.append(gr.update())
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Base Model', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Refiner Model', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Refiner Switch', None)
+        assert h is not None
+        h = float(h)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Sampler', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Scheduler', None)
+        assert isinstance(h, str)
+        results.append(h)
+    except:
+        results.append(gr.update())
+
+    try:
+        h = loaded_parameter_dict.get('Seed', None)
+        assert h is not None
+        h = int(h)
+        results.append(False)
+        results.append(h)
+    except:
+        results.append(gr.update())
+        results.append(gr.update())
+
+    if is_generating:
+        results.append(gr.update())
+    else:
+        results.append(gr.update(visible=True))
+    
+    results.append(gr.update(visible=False))
+
+    for i in range(1, 6):
+        try:
+            n, w = loaded_parameter_dict.get(f'LoRA {i}').split(' : ')
+            w = float(w)
+            results.append(n)
+            results.append(w)
+        except:
+            results.append(gr.update())
+            results.append(gr.update())
+
+    return results

modules/model_loader.py

@@ -0,0 +1,26 @@
+import os
+from urllib.parse import urlparse
+from typing import Optional
+
+
+def load_file_from_url(
+        url: str,
+        *,
+        model_dir: str,
+        progress: bool = True,
+        file_name: Optional[str] = None,
+) -> str:
+    """Download a file from `url` into `model_dir`, using the file present if possible.
+
+    Returns the path to the downloaded file.
+    """
+    os.makedirs(model_dir, exist_ok=True)
+    if not file_name:
+        parts = urlparse(url)
+        file_name = os.path.basename(parts.path)
+    cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+    if not os.path.exists(cached_file):
+        print(f'Downloading: "{url}" to {cached_file}\n')
+        from torch.hub import download_url_to_file
+        download_url_to_file(url, cached_file, progress=progress)
+    return cached_file

modules/ops.py

@@ -0,0 +1,19 @@
+import torch
+import contextlib
+
+
+@contextlib.contextmanager
+def use_patched_ops(operations):
+    op_names = ['Linear', 'Conv2d', 'Conv3d', 'GroupNorm', 'LayerNorm']
+    backups = {op_name: getattr(torch.nn, op_name) for op_name in op_names}
+
+    try:
+        for op_name in op_names:
+            setattr(torch.nn, op_name, getattr(operations, op_name))
+
+        yield
+
+    finally:
+        for op_name in op_names:
+            setattr(torch.nn, op_name, backups[op_name])
+    return

modules/patch.py

@@ -0,0 +1,507 @@
+import os
+import torch
+import time
+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 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 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 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
+from modules.patch_precision import patch_all_precision
+from modules.patch_clip import patch_all_clip
+
+
+sharpness = 2.0
+
+adm_scaler_end = 0.3
+positive_adm_scale = 1.5
+negative_adm_scale = 0.8
+
+adaptive_cfg = 7.0
+global_diffusion_progress = 0
+eps_record = None
+
+
+def calculate_weight_patched(self, patches, weight, key):
+    for p in patches:
+        alpha = p[0]
+        v = p[1]
+        strength_model = p[2]
+
+        if strength_model != 1.0:
+            weight *= strength_model
+
+        if isinstance(v, list):
+            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 * 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:
+                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)
+            try:
+                weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(
+                    weight.shape).type(weight.dtype)
+            except Exception as e:
+                print("ERROR", key, e)
+        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]
+            w1_b = v[4]
+            w2_a = v[5]
+            w2_b = v[6]
+            t2 = v[7]
+            dim = None
+
+            if w1 is None:
+                dim = w1_b.shape[0]
+                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 = 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(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',
+                                      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 = ldm_patched.modules.model_management.cast_to_device(w2, weight.device, torch.float32)
+
+            if len(w2.shape) == 4:
+                w1 = w1.unsqueeze(2).unsqueeze(2)
+            if v[2] is not None and dim is not None:
+                alpha *= v[2] / dim
+
+            try:
+                weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
+            except Exception as e:
+                print("ERROR", key, e)
+        elif patch_type == "loha":
+            w1a = v[0]
+            w1b = v[1]
+            if v[2] is not None:
+                alpha *= v[2] / w1b.shape[0]
+            w2a = v[3]
+            w2b = v[4]
+            if v[5] is not None:  # cp decomposition
+                t1 = v[5]
+                t2 = v[6]
+                m1 = torch.einsum('i j k l, j r, i p -> p r k l',
+                                  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',
+                                  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(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
+
+
+class BrownianTreeNoiseSamplerPatched:
+    transform = None
+    tree = None
+
+    @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)
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    @staticmethod
+    def __call__(sigma, sigma_next):
+        transform = BrownianTreeNoiseSamplerPatched.transform
+        tree = BrownianTreeNoiseSamplerPatched.tree
+
+        t0, t1 = transform(torch.as_tensor(sigma)), transform(torch.as_tensor(sigma_next))
+        return tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+
+def compute_cfg(uncond, cond, cfg_scale, t):
+    global adaptive_cfg
+
+    mimic_cfg = float(adaptive_cfg)
+    real_cfg = float(cfg_scale)
+
+    real_eps = uncond + real_cfg * (cond - uncond)
+
+    if cfg_scale > adaptive_cfg:
+        mimicked_eps = uncond + mimic_cfg * (cond - uncond)
+        return real_eps * t + mimicked_eps * (1 - t)
+    else:
+        return real_eps
+
+
+def patched_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options=None, seed=None):
+    global eps_record
+
+    if math.isclose(cond_scale, 1.0) and not model_options.get("disable_cfg1_optimization", False):
+        final_x0 = calc_cond_uncond_batch(model, cond, None, x, timestep, model_options)[0]
+
+        if eps_record is not None:
+            eps_record = ((x - final_x0) / timestep).cpu()
+
+        return final_x0
+
+    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
+
+    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=cond_scale, t=global_diffusion_progress)
+
+    if eps_record is not None:
+        eps_record = (final_eps / timestep).cpu()
+
+    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 = 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
+
+    if kwargs.get("prompt_type", "") == "negative":
+        width = float(width) * negative_adm_scale
+        height = float(height) * negative_adm_scale
+    elif kwargs.get("prompt_type", "") == "positive":
+        width = float(width) * positive_adm_scale
+        height = float(height) * positive_adm_scale
+
+    def embedder(number_list):
+        h = self.embedder(torch.tensor(number_list, dtype=torch.float32))
+        h = torch.flatten(h).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+        return h
+
+    width, height = int(width), int(height)
+    target_width, target_height = round_to_64(target_width), round_to_64(target_height)
+
+    adm_emphasized = embedder([height, width, 0, 0, target_height, target_width])
+    adm_consistent = embedder([target_height, target_width, 0, 0, target_height, target_width])
+
+    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 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)
+
+        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
+        x = x * inpaint_mask + (inpaint_latent + current_energy) * (1.0 - inpaint_mask)
+
+        out = self.inner_model(x, sigma,
+                               cond=cond,
+                               uncond=uncond,
+                               cond_scale=cond_scale,
+                               model_options=model_options,
+                               seed=seed)
+
+        out = out * inpaint_mask + inpaint_latent * (1.0 - inpaint_mask)
+    else:
+        out = self.inner_model(x, sigma,
+                               cond=cond,
+                               uncond=uncond,
+                               cond_scale=cond_scale,
+                               model_options=model_options,
+                               seed=seed)
+    return out
+
+
+def timed_adm(y, timesteps):
+    if isinstance(y, torch.Tensor) and int(y.dim()) == 2 and int(y.shape[1]) == 5632:
+        y_mask = (timesteps > 999.0 * (1.0 - float(adm_scaler_end))).to(y)[..., None]
+        y_with_adm = y[..., :2816].clone()
+        y_without_adm = y[..., 2816:].clone()
+        return y_with_adm * y_mask + y_without_adm * (1.0 - y_mask)
+    return y
+
+
+def patched_cldm_forward(self, x, hint, timesteps, context, y=None, **kwargs):
+    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)
+
+    y = timed_adm(y, timesteps)
+
+    outs = []
+
+    hs = []
+    if self.num_classes is not None:
+        assert y.shape[0] == x.shape[0]
+        emb = emb + self.label_emb(y)
+
+    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)
+            h += guided_hint
+            guided_hint = None
+        else:
+            h = module(h, emb, context)
+        outs.append(zero_conv(h, emb, context))
+
+    h = self.middle_block(h, emb, context)
+    outs.append(self.middle_block_out(h, emb, context))
+
+    if advanced_parameters.controlnet_softness > 0:
+        for i in range(10):
+            k = 1.0 - float(i) / 9.0
+            outs[i] = outs[i] * (1.0 - advanced_parameters.controlnet_softness * k)
+
+    return outs
+
+
+def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+    global global_diffusion_progress
+
+    self.current_step = 1.0 - timesteps.to(x) / 999.0
+    global_diffusion_progress = float(self.current_step.detach().cpu().numpy().tolist()[0])
+
+    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 = 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
+    for id, module in enumerate(self.input_blocks):
+        transformer_options["block"] = ("input", id)
+        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"]
+            for p in patch:
+                h = p(h, transformer_options)
+
+        hs.append(h)
+        if "input_block_patch_after_skip" in transformer_patches:
+            patch = transformer_patches["input_block_patch_after_skip"]
+            for p in patch:
+                h = p(h, transformer_options)
+
+    transformer_options["block"] = ("middle", 0)
+    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):
+        transformer_options["block"] = ("output", id)
+        hsp = hs.pop()
+        hsp = apply_control(hsp, control, 'output')
+
+        if "output_block_patch" in transformer_patches:
+            patch = transformer_patches["output_block_patch"]
+            for p in patch:
+                h, hsp = p(h, hsp, transformer_options)
+
+        h = torch.cat([h, hsp], dim=1)
+        del hsp
+        if len(hs) > 0:
+            output_shape = hs[-1].shape
+        else:
+            output_shape = None
+        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)
+    else:
+        return self.out(h)
+
+
+def patched_load_models_gpu(*args, **kwargs):
+    execution_start_time = time.perf_counter()
+    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')
+    return y
+
+
+def build_loaded(module, loader_name):
+    original_loader_name = loader_name + '_origin'
+
+    if not hasattr(module, original_loader_name):
+        setattr(module, original_loader_name, getattr(module, loader_name))
+
+    original_loader = getattr(module, original_loader_name)
+
+    def loader(*args, **kwargs):
+        result = None
+        try:
+            result = original_loader(*args, **kwargs)
+        except Exception as e:
+            result = None
+            exp = str(e) + '\n'
+            for path in list(args) + list(kwargs.values()):
+                if isinstance(path, str):
+                    if os.path.exists(path):
+                        exp += f'File corrupted: {path} \n'
+                        corrupted_backup_file = path + '.corrupted'
+                        if os.path.exists(corrupted_backup_file):
+                            os.remove(corrupted_backup_file)
+                        os.replace(path, corrupted_backup_file)
+                        if os.path.exists(path):
+                            os.remove(path)
+                        exp += f'Fooocus has tried to move the corrupted file to {corrupted_backup_file} \n'
+                        exp += f'You may try again now and Fooocus will download models again. \n'
+            raise ValueError(exp)
+        return result
+
+    setattr(module, loader_name, loader)
+    return
+
+
+def patch_all():
+    if ldm_patched.modules.model_management.directml_enabled:
+        ldm_patched.modules.model_management.lowvram_available = True
+        ldm_patched.modules.model_management.OOM_EXCEPTION = Exception
+    
+    patch_all_precision()
+    patch_all_clip()
+
+    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
+
+    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.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')
+
+    build_loaded(safetensors.torch, 'load_file')
+    build_loaded(torch, 'load')
+
+    return

modules/patch_clip.py

@@ -0,0 +1,195 @@
+# Consistent with Kohya/A1111 to reduce differences between model training and inference.
+
+import os
+import torch
+import ldm_patched.controlnet.cldm
+import ldm_patched.k_diffusion.sampling
+import ldm_patched.ldm.modules.attention
+import ldm_patched.ldm.modules.diffusionmodules.model
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.args_parser
+import ldm_patched.modules.model_base
+import ldm_patched.modules.model_management
+import ldm_patched.modules.model_patcher
+import ldm_patched.modules.samplers
+import ldm_patched.modules.sd
+import ldm_patched.modules.sd1_clip
+import ldm_patched.modules.clip_vision
+import ldm_patched.modules.ops as ops
+
+from modules.ops import use_patched_ops
+from transformers import CLIPTextModel, CLIPTextConfig, modeling_utils, CLIPVisionConfig, CLIPVisionModelWithProjection
+
+
+def patched_encode_token_weights(self, token_weight_pairs):
+    to_encode = list()
+    max_token_len = 0
+    has_weights = False
+    for x in token_weight_pairs:
+        tokens = list(map(lambda a: a[0], x))
+        max_token_len = max(len(tokens), max_token_len)
+        has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
+        to_encode.append(tokens)
+
+    sections = len(to_encode)
+    if has_weights or sections == 0:
+        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].to(ldm_patched.modules.model_management.intermediate_device())
+    else:
+        first_pooled = pooled
+
+    output = []
+    for k in range(0, sections):
+        z = out[k:k + 1]
+        if has_weights:
+            original_mean = z.mean()
+            z_empty = out[-1]
+            for i in range(len(z)):
+                for j in range(len(z[i])):
+                    weight = token_weight_pairs[k][j][1]
+                    if weight != 1.0:
+                        z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
+            new_mean = z.mean()
+            z = z * (original_mean / new_mean)
+        output.append(z)
+
+    if len(output) == 0:
+        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_SDClipModel__init__(self, max_length=77, freeze=True, layer="last", layer_idx=None,
+                                textmodel_json_config=None, dtype=None, special_tokens=None,
+                                layer_norm_hidden_state=True, **kwargs):
+    torch.nn.Module.__init__(self)
+    assert layer in self.LAYERS
+
+    if special_tokens is None:
+        special_tokens = {"start": 49406, "end": 49407, "pad": 49407}
+
+    if textmodel_json_config is None:
+        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)),
+                                             "sd1_clip_config.json")
+
+    config = CLIPTextConfig.from_json_file(textmodel_json_config)
+    self.num_layers = config.num_hidden_layers
+
+    with use_patched_ops(ops.manual_cast):
+        with modeling_utils.no_init_weights():
+            self.transformer = CLIPTextModel(config)
+
+    if dtype is not None:
+        self.transformer.to(dtype)
+
+    self.transformer.text_model.embeddings.to(torch.float32)
+
+    if freeze:
+        self.freeze()
+
+    self.max_length = max_length
+    self.layer = layer
+    self.layer_idx = None
+    self.special_tokens = special_tokens
+    self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
+    self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
+    self.enable_attention_masks = False
+
+    self.layer_norm_hidden_state = layer_norm_hidden_state
+    if layer == "hidden":
+        assert layer_idx is not None
+        assert abs(layer_idx) < self.num_layers
+        self.clip_layer(layer_idx)
+    self.layer_default = (self.layer, self.layer_idx)
+
+
+def patched_SDClipModel_forward(self, tokens):
+    backup_embeds = self.transformer.get_input_embeddings()
+    device = backup_embeds.weight.device
+    tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
+    tokens = torch.LongTensor(tokens).to(device)
+
+    attention_mask = None
+    if self.enable_attention_masks:
+        attention_mask = torch.zeros_like(tokens)
+        max_token = self.transformer.get_input_embeddings().weight.shape[0] - 1
+        for x in range(attention_mask.shape[0]):
+            for y in range(attention_mask.shape[1]):
+                attention_mask[x, y] = 1
+                if tokens[x, y] == max_token:
+                    break
+
+    outputs = self.transformer(input_ids=tokens, attention_mask=attention_mask,
+                               output_hidden_states=self.layer == "hidden")
+    self.transformer.set_input_embeddings(backup_embeds)
+
+    if self.layer == "last":
+        z = outputs.last_hidden_state
+    elif self.layer == "pooled":
+        z = outputs.pooler_output[:, None, :]
+    else:
+        z = outputs.hidden_states[self.layer_idx]
+        if self.layer_norm_hidden_state:
+            z = self.transformer.text_model.final_layer_norm(z)
+
+    if hasattr(outputs, "pooler_output"):
+        pooled_output = outputs.pooler_output.float()
+    else:
+        pooled_output = None
+
+    if self.text_projection is not None and pooled_output is not None:
+        pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
+
+    return z.float(), pooled_output
+
+
+def patched_ClipVisionModel__init__(self, json_config):
+    config = CLIPVisionConfig.from_json_file(json_config)
+
+    self.load_device = ldm_patched.modules.model_management.text_encoder_device()
+    self.offload_device = ldm_patched.modules.model_management.text_encoder_offload_device()
+
+    if ldm_patched.modules.model_management.should_use_fp16(self.load_device, prioritize_performance=False):
+        self.dtype = torch.float16
+    else:
+        self.dtype = torch.float32
+
+    with use_patched_ops(ops.manual_cast):
+        with modeling_utils.no_init_weights():
+            self.model = CLIPVisionModelWithProjection(config)
+
+    self.model.to(self.dtype)
+    self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(
+        self.model,
+        load_device=self.load_device,
+        offload_device=self.offload_device
+    )
+
+
+def patched_ClipVisionModel_encode_image(self, image):
+    ldm_patched.modules.model_management.load_model_gpu(self.patcher)
+    pixel_values = ldm_patched.modules.clip_vision.clip_preprocess(image.to(self.load_device))
+    outputs = self.model(pixel_values=pixel_values, output_hidden_states=True)
+
+    for k in outputs:
+        t = outputs[k]
+        if t is not None:
+            if k == 'hidden_states':
+                outputs["penultimate_hidden_states"] = t[-2].to(ldm_patched.modules.model_management.intermediate_device())
+                outputs["hidden_states"] = None
+            else:
+                outputs[k] = t.to(ldm_patched.modules.model_management.intermediate_device())
+
+    return outputs
+
+
+def patch_all_clip():
+    ldm_patched.modules.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = patched_encode_token_weights
+    ldm_patched.modules.sd1_clip.SDClipModel.__init__ = patched_SDClipModel__init__
+    ldm_patched.modules.sd1_clip.SDClipModel.forward = patched_SDClipModel_forward
+    ldm_patched.modules.clip_vision.ClipVisionModel.__init__ = patched_ClipVisionModel__init__
+    ldm_patched.modules.clip_vision.ClipVisionModel.encode_image = patched_ClipVisionModel_encode_image
+    return

modules/patch_precision.py

@@ -0,0 +1,60 @@
+# Consistent with Kohya to reduce differences between model training and inference.
+
+import torch
+import math
+import einops
+import numpy as np
+
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.model_sampling
+import ldm_patched.modules.sd1_clip
+
+from ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
+
+
+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 = einops.repeat(timesteps, 'b -> b d', d=dim)
+    return embedding
+
+
+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 patch_all_precision():
+    ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding = patched_timestep_embedding
+    ldm_patched.modules.model_sampling.ModelSamplingDiscrete._register_schedule = patched_register_schedule
+    return

modules/private_logger.py

@@ -0,0 +1,108 @@
+import os
+import args_manager
+import modules.config
+import json
+import urllib.parse
+
+from PIL import Image
+from modules.util import generate_temp_filename
+
+
+log_cache = {}
+
+
+def get_current_html_path():
+    date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs,
+                                                                         extension='png')
+    html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+    return html_name
+
+
+def log(img, dic):
+    if args_manager.args.disable_image_log:
+        return
+
+    date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs, extension='png')
+    os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
+    Image.fromarray(img).save(local_temp_filename)
+    html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+
+    css_styles = (
+        "<style>"
+        "body { background-color: #121212; color: #E0E0E0; } "
+        "a { color: #BB86FC; } "
+        ".metadata { border-collapse: collapse; width: 100%; } "
+        ".metadata .key { width: 15%; } "
+        ".metadata .value { width: 85%; font-weight: bold; } "
+        ".metadata th, .metadata td { border: 1px solid #4d4d4d; padding: 4px; } "
+        ".image-container img { height: auto; max-width: 512px; display: block; padding-right:10px; } "
+        ".image-container div { text-align: center; padding: 4px; } "
+        "hr { border-color: gray; } "
+        "button { background-color: black; color: white; border: 1px solid grey; border-radius: 5px; padding: 5px 10px; text-align: center; display: inline-block; font-size: 16px; cursor: pointer; }"
+        "button:hover {background-color: grey; color: black;}"
+        "</style>"
+    )
+
+    js = (
+        """<script>
+        function to_clipboard(txt) { 
+        txt = decodeURIComponent(txt);
+        if (navigator.clipboard && navigator.permissions) {
+            navigator.clipboard.writeText(txt)
+        } else {
+            const textArea = document.createElement('textArea')
+            textArea.value = txt
+            textArea.style.width = 0
+            textArea.style.position = 'fixed'
+            textArea.style.left = '-999px'
+            textArea.style.top = '10px'
+            textArea.setAttribute('readonly', 'readonly')
+            document.body.appendChild(textArea)
+
+            textArea.select()
+            document.execCommand('copy')
+            document.body.removeChild(textArea)
+        }
+        alert('Copied to Clipboard!\\nPaste to prompt area to load parameters.\\nCurrent clipboard content is:\\n\\n' + txt);
+        }
+        </script>"""
+    )
+
+    begin_part = f"<html><head><title>Fooocus Log {date_string}</title>{css_styles}</head><body>{js}<p>Fooocus Log {date_string} (private)</p>\n<p>All images are clean, without any hidden data/meta, and safe to share with others.</p><!--fooocus-log-split-->\n\n"
+    end_part = f'\n<!--fooocus-log-split--></body></html>'
+
+    middle_part = log_cache.get(html_name, "")
+
+    if middle_part == "":
+        if os.path.exists(html_name):
+            existing_split = open(html_name, 'r', encoding='utf-8').read().split('<!--fooocus-log-split-->')
+            if len(existing_split) == 3:
+                middle_part = existing_split[1]
+            else:
+                middle_part = existing_split[0]
+
+    div_name = only_name.replace('.', '_')
+    item = f"<div id=\"{div_name}\" class=\"image-container\"><hr><table><tr>\n"
+    item += f"<td><a href=\"{only_name}\" target=\"_blank\"><img src='{only_name}' onerror=\"this.closest('.image-container').style.display='none';\" loading='lazy'></img></a><div>{only_name}</div></td>"
+    item += "<td><table class='metadata'>"
+    for key, value in dic:
+        value_txt = str(value).replace('\n', ' </br> ')
+        item += f"<tr><td class='key'>{key}</td><td class='value'>{value_txt}</td></tr>\n"
+    item += "</table>"
+
+    js_txt = urllib.parse.quote(json.dumps({k: v for k, v in dic}, indent=0), safe='')
+    item += f"</br><button onclick=\"to_clipboard('{js_txt}')\">Copy to Clipboard</button>"
+
+    item += "</td>"
+    item += "</tr></table></div>\n\n"
+
+    middle_part = item + middle_part
+
+    with open(html_name, 'w', encoding='utf-8') as f:
+        f.write(begin_part + middle_part + end_part)
+
+    print(f'Image generated with private log at: {html_name}')
+
+    log_cache[html_name] = middle_part
+
+    return

modules/sample_hijack.py

@@ -0,0 +1,184 @@
+import torch
+import ldm_patched.modules.samplers
+import ldm_patched.modules.model_management
+
+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
+
+
+current_refiner = None
+refiner_switch_step = -1
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate_inner(c, p, target_model=None, target_clip=None):
+    if target_model is None or isinstance(target_model, SDXLRefiner):
+        c = c[..., -1280:].clone()
+    elif isinstance(target_model, SDXL):
+        c = c.clone()
+    else:
+        p = None
+        c = c[..., :768].clone()
+
+        final_layer_norm = target_clip.cond_stage_model.clip_l.transformer.text_model.final_layer_norm
+
+        final_layer_norm_origin_device = final_layer_norm.weight.device
+        final_layer_norm_origin_dtype = final_layer_norm.weight.dtype
+
+        c_origin_device = c.device
+        c_origin_dtype = c.dtype
+
+        final_layer_norm.to(device='cpu', dtype=torch.float32)
+        c = c.to(device='cpu', dtype=torch.float32)
+
+        c = torch.chunk(c, int(c.size(1)) // 77, 1)
+        c = [final_layer_norm(ci) for ci in c]
+        c = torch.cat(c, dim=1)
+
+        final_layer_norm.to(device=final_layer_norm_origin_device, dtype=final_layer_norm_origin_dtype)
+        c = c.to(device=c_origin_device, dtype=c_origin_dtype)
+    return c, p
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate(cond, target_model=None, target_clip=None):
+    results = []
+
+    for c, px in cond:
+        p = px.get('pooled_output', None)
+        c, p = clip_separate_inner(c, p, target_model=target_model, target_clip=target_clip)
+        p = {} if p is None else {'pooled_output': p.clone()}
+        results.append([c, p])
+
+    return results
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate_after_preparation(cond, target_model=None, target_clip=None):
+    results = []
+
+    for x in cond:
+        p = x.get('pooled_output', None)
+        c = x['model_conds']['c_crossattn'].cond
+
+        c, p = clip_separate_inner(c, p, target_model=target_model, target_clip=target_clip)
+
+        result = {'model_conds': {'c_crossattn': CONDRegular(c)}}
+
+        if p is not None:
+            result['pooled_output'] = p.clone()
+
+        results.append(result)
+
+    return results
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
+    global current_refiner
+
+    positive = positive[:]
+    negative = negative[:]
+
+    resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
+    resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
+
+    model_wrap = wrap_model(model)
+
+    calculate_start_end_timesteps(model, negative)
+    calculate_start_end_timesteps(model, positive)
+
+    if latent_image is not None:
+        latent_image = model.process_latent_in(latent_image)
+
+    if hasattr(model, 'extra_conds'):
+        positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
+        negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
+
+    #make sure each cond area has an opposite one with the same area
+    for c in positive:
+        create_cond_with_same_area_if_none(negative, c)
+    for c in negative:
+        create_cond_with_same_area_if_none(positive, c)
+
+    # pre_run_control(model, negative + positive)
+    pre_run_control(model, positive)  # negative is not necessary in Fooocus, 0.5s faster.
+
+    apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
+    apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
+
+    extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
+
+    if current_refiner is not None and hasattr(current_refiner.model, 'extra_conds'):
+        positive_refiner = clip_separate_after_preparation(positive, target_model=current_refiner.model)
+        negative_refiner = clip_separate_after_preparation(negative, target_model=current_refiner.model)
+
+        positive_refiner = encode_model_conds(current_refiner.model.extra_conds, positive_refiner, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
+        negative_refiner = encode_model_conds(current_refiner.model.extra_conds, negative_refiner, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
+
+    def refiner_switch():
+        cleanup_additional_models(set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")))
+
+        extra_args["cond"] = positive_refiner
+        extra_args["uncond"] = negative_refiner
+
+        # clear ip-adapter for refiner
+        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())
+        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')
+        return
+
+    def callback_wrap(step, x0, x, total_steps):
+        if step == refiner_switch_step and current_refiner is not None:
+            refiner_switch()
+        if callback is not None:
+            # residual_noise_preview = x - x0
+            # residual_noise_preview /= residual_noise_preview.std()
+            # residual_noise_preview *= x0.std()
+            callback(step, x0, x, total_steps)
+
+    samples = sampler.sample(model_wrap, sigmas, extra_args, callback_wrap, noise, latent_image, denoise_mask, disable_pbar)
+    return model.process_latent_out(samples.to(torch.float32))
+
+
+@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=steps, denoise=1.0)[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/sdxl_styles.py

@@ -0,0 +1,82 @@
+import os
+import re
+import json
+
+from modules.util import get_files_from_folder
+
+
+# cannot use modules.config - validators causing circular imports
+styles_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../sdxl_styles/'))
+wildcards_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../wildcards/'))
+wildcards_max_bfs_depth = 64
+
+
+def normalize_key(k):
+    k = k.replace('-', ' ')
+    words = k.split(' ')
+    words = [w[:1].upper() + w[1:].lower() for w in words]
+    k = ' '.join(words)
+    k = k.replace('3d', '3D')
+    k = k.replace('Sai', 'SAI')
+    k = k.replace('Mre', 'MRE')
+    k = k.replace('(s', '(S')
+    return k
+
+
+styles = {}
+
+styles_files = get_files_from_folder(styles_path, ['.json'])
+
+for x in ['sdxl_styles_fooocus.json',
+          'sdxl_styles_sai.json',
+          'sdxl_styles_mre.json',
+          'sdxl_styles_twri.json',
+          'sdxl_styles_diva.json',
+          'sdxl_styles_marc_k3nt3l.json']:
+    if x in styles_files:
+        styles_files.remove(x)
+        styles_files.append(x)
+
+for styles_file in styles_files:
+    try:
+        with open(os.path.join(styles_path, styles_file), encoding='utf-8') as f:
+            for entry in json.load(f):
+                name = normalize_key(entry['name'])
+                prompt = entry['prompt'] if 'prompt' in entry else ''
+                negative_prompt = entry['negative_prompt'] if 'negative_prompt' in entry else ''
+                styles[name] = (prompt, negative_prompt)
+    except Exception as e:
+        print(str(e))
+        print(f'Failed to load style file {styles_file}')
+
+style_keys = list(styles.keys())
+fooocus_expansion = "Fooocus V2"
+legal_style_names = [fooocus_expansion] + style_keys
+
+
+def apply_style(style, positive):
+    p, n = styles[style]
+    return p.replace('{prompt}', positive).splitlines(), n.splitlines()
+
+
+def apply_wildcards(wildcard_text, rng, directory=wildcards_path):
+    for _ in range(wildcards_max_bfs_depth):
+        placeholders = re.findall(r'__([\w-]+)__', wildcard_text)
+        if len(placeholders) == 0:
+            return wildcard_text
+
+        print(f'[Wildcards] processing: {wildcard_text}')
+        for placeholder in placeholders:
+            try:
+                words = open(os.path.join(directory, f'{placeholder}.txt'), encoding='utf-8').read().splitlines()
+                words = [x for x in words if x != '']
+                assert len(words) > 0
+                wildcard_text = wildcard_text.replace(f'__{placeholder}__', rng.choice(words), 1)
+            except:
+                print(f'[Wildcards] Warning: {placeholder}.txt missing or empty. '
+                      f'Using "{placeholder}" as a normal word.')
+                wildcard_text = wildcard_text.replace(f'__{placeholder}__', placeholder)
+            print(f'[Wildcards] {wildcard_text}')
+
+    print(f'[Wildcards] BFS stack overflow. Current text: {wildcard_text}')
+    return wildcard_text

modules/style_sorter.py

@@ -0,0 +1,59 @@
+import os
+import gradio as gr
+import modules.localization as localization
+import json
+
+
+all_styles = []
+
+
+def try_load_sorted_styles(style_names, default_selected):
+    global all_styles
+
+    all_styles = style_names
+
+    try:
+        if os.path.exists('sorted_styles.json'):
+            with open('sorted_styles.json', 'rt', encoding='utf-8') as fp:
+                sorted_styles = []
+                for x in json.load(fp):
+                    if x in all_styles:
+                        sorted_styles.append(x)
+                for x in all_styles:
+                    if x not in sorted_styles:
+                        sorted_styles.append(x)
+                all_styles = sorted_styles
+    except Exception as e:
+        print('Load style sorting failed.')
+        print(e)
+
+    unselected = [y for y in all_styles if y not in default_selected]
+    all_styles = default_selected + unselected
+
+    return
+
+
+def sort_styles(selected):
+    global all_styles
+    unselected = [y for y in all_styles if y not in selected]
+    sorted_styles = selected + unselected
+    try:
+        with open('sorted_styles.json', 'wt', encoding='utf-8') as fp:
+            json.dump(sorted_styles, fp, indent=4)
+    except Exception as e:
+        print('Write style sorting failed.')
+        print(e)
+    all_styles = sorted_styles
+    return gr.CheckboxGroup.update(choices=sorted_styles)
+
+
+def localization_key(x):
+    return x + localization.current_translation.get(x, '')
+
+
+def search_styles(selected, query):
+    unselected = [y for y in all_styles if y not in selected]
+    matched = [y for y in unselected if query.lower() in localization_key(y).lower()] if len(query.replace(' ', '')) > 0 else []
+    unmatched = [y for y in unselected if y not in matched]
+    sorted_styles = matched + selected + unmatched
+    return gr.CheckboxGroup.update(choices=sorted_styles)

modules/ui_gradio_extensions.py

@@ -0,0 +1,67 @@
+# based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/v1.6.0/modules/ui_gradio_extensions.py
+
+import os
+import gradio as gr
+import args_manager
+
+from modules.localization import localization_js
+
+
+GradioTemplateResponseOriginal = gr.routes.templates.TemplateResponse
+
+modules_path = os.path.dirname(os.path.realpath(__file__))
+script_path = os.path.dirname(modules_path)
+
+
+def webpath(fn):
+    if fn.startswith(script_path):
+        web_path = os.path.relpath(fn, script_path).replace('\\', '/')
+    else:
+        web_path = os.path.abspath(fn)
+
+    return f'file={web_path}?{os.path.getmtime(fn)}'
+
+
+def javascript_html():
+    script_js_path = webpath('javascript/script.js')
+    context_menus_js_path = webpath('javascript/contextMenus.js')
+    localization_js_path = webpath('javascript/localization.js')
+    zoom_js_path = webpath('javascript/zoom.js')
+    edit_attention_js_path = webpath('javascript/edit-attention.js')
+    viewer_js_path = webpath('javascript/viewer.js')
+    image_viewer_js_path = webpath('javascript/imageviewer.js')
+    samples_path = webpath(os.path.abspath('./sdxl_styles/samples/fooocus_v2.jpg'))
+    head = f'<script type="text/javascript">{localization_js(args_manager.args.language)}</script>\n'
+    head += f'<script type="text/javascript" src="{script_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{context_menus_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{localization_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{zoom_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{edit_attention_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{viewer_js_path}"></script>\n'
+    head += f'<script type="text/javascript" src="{image_viewer_js_path}"></script>\n'
+    head += f'<meta name="samples-path" content="{samples_path}"></meta>\n'
+
+    if args_manager.args.theme:
+        head += f'<script type="text/javascript">set_theme(\"{args_manager.args.theme}\");</script>\n'
+
+    return head
+
+
+def css_html():
+    style_css_path = webpath('css/style.css')
+    head = f'<link rel="stylesheet" property="stylesheet" href="{style_css_path}">'
+    return head
+
+
+def reload_javascript():
+    js = javascript_html()
+    css = css_html()
+
+    def template_response(*args, **kwargs):
+        res = GradioTemplateResponseOriginal(*args, **kwargs)
+        res.body = res.body.replace(b'</head>', f'{js}</head>'.encode("utf8"))
+        res.body = res.body.replace(b'</body>', f'{css}</body>'.encode("utf8"))
+        res.init_headers()
+        return res
+
+    gr.routes.templates.TemplateResponse = template_response

modules/upscaler.py

@@ -0,0 +1,34 @@
+import os
+import torch
+import modules.core as core
+
+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
+
+model_filename = os.path.join(path_upscale_models, 'fooocus_upscaler_s409985e5.bin')
+opImageUpscaleWithModel = ImageUpscaleWithModel()
+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()
+        for k, v in sd.items():
+            sdo[k.replace('residual_block_', 'RDB')] = v
+        del sd
+        model = ESRGAN(sdo)
+        model.cpu()
+        model.eval()
+
+    img = core.numpy_to_pytorch(img)
+    img = opImageUpscaleWithModel.upscale(model, img)[0]
+    img = core.pytorch_to_numpy(img)[0]
+
+    return img

modules/util.py

@@ -0,0 +1,177 @@
+import numpy as np
+import datetime
+import random
+import math
+import os
+import cv2
+
+from PIL import Image
+
+
+LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
+
+
+def erode_or_dilate(x, k):
+    k = int(k)
+    if k > 0:
+        return cv2.dilate(x, kernel=np.ones(shape=(3, 3), dtype=np.uint8), iterations=k)
+    if k < 0:
+        return cv2.erode(x, kernel=np.ones(shape=(3, 3), dtype=np.uint8), iterations=-k)
+    return x
+
+
+def resample_image(im, width, height):
+    im = Image.fromarray(im)
+    im = im.resize((int(width), int(height)), resample=LANCZOS)
+    return np.array(im)
+
+
+def resize_image(im, width, height, resize_mode=1):
+    """
+    Resizes an image with the specified resize_mode, width, and height.
+
+    Args:
+        resize_mode: The mode to use when resizing the image.
+            0: Resize the image to the specified width and height.
+            1: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess.
+            2: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image.
+        im: The image to resize.
+        width: The width to resize the image to.
+        height: The height to resize the image to.
+    """
+
+    im = Image.fromarray(im)
+
+    def resize(im, w, h):
+        return im.resize((w, h), resample=LANCZOS)
+
+    if resize_mode == 0:
+        res = resize(im, width, height)
+
+    elif resize_mode == 1:
+        ratio = width / height
+        src_ratio = im.width / im.height
+
+        src_w = width if ratio > src_ratio else im.width * height // im.height
+        src_h = height if ratio <= src_ratio else im.height * width // im.width
+
+        resized = resize(im, src_w, src_h)
+        res = Image.new("RGB", (width, height))
+        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
+
+    else:
+        ratio = width / height
+        src_ratio = im.width / im.height
+
+        src_w = width if ratio < src_ratio else im.width * height // im.height
+        src_h = height if ratio >= src_ratio else im.height * width // im.width
+
+        resized = resize(im, src_w, src_h)
+        res = Image.new("RGB", (width, height))
+        res.paste(resized, box=(width // 2 - src_w // 2, height // 2 - src_h // 2))
+
+        if ratio < src_ratio:
+            fill_height = height // 2 - src_h // 2
+            if fill_height > 0:
+                res.paste(resized.resize((width, fill_height), box=(0, 0, width, 0)), box=(0, 0))
+                res.paste(resized.resize((width, fill_height), box=(0, resized.height, width, resized.height)), box=(0, fill_height + src_h))
+        elif ratio > src_ratio:
+            fill_width = width // 2 - src_w // 2
+            if fill_width > 0:
+                res.paste(resized.resize((fill_width, height), box=(0, 0, 0, height)), box=(0, 0))
+                res.paste(resized.resize((fill_width, height), box=(resized.width, 0, resized.width, height)), box=(fill_width + src_w, 0))
+
+    return np.array(res)
+
+
+def get_shape_ceil(h, w):
+    return math.ceil(((h * w) ** 0.5) / 64.0) * 64.0
+
+
+def get_image_shape_ceil(im):
+    H, W = im.shape[:2]
+    return get_shape_ceil(H, W)
+
+
+def set_image_shape_ceil(im, shape_ceil):
+    shape_ceil = float(shape_ceil)
+
+    H_origin, W_origin, _ = im.shape
+    H, W = H_origin, W_origin
+    
+    for _ in range(256):
+        current_shape_ceil = get_shape_ceil(H, W)
+        if abs(current_shape_ceil - shape_ceil) < 0.1:
+            break
+        k = shape_ceil / current_shape_ceil
+        H = int(round(float(H) * k / 64.0) * 64)
+        W = int(round(float(W) * k / 64.0) * 64)
+
+    if H == H_origin and W == W_origin:
+        return im
+
+    return resample_image(im, width=W, height=H)
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def remove_empty_str(items, default=None):
+    items = [x for x in items if x != ""]
+    if len(items) == 0 and default is not None:
+        return [default]
+    return items
+
+
+def join_prompts(*args, **kwargs):
+    prompts = [str(x) for x in args if str(x) != ""]
+    if len(prompts) == 0:
+        return ""
+    if len(prompts) == 1:
+        return prompts[0]
+    return ', '.join(prompts)
+
+
+def generate_temp_filename(folder='./outputs/', extension='png'):
+    current_time = datetime.datetime.now()
+    date_string = current_time.strftime("%Y-%m-%d")
+    time_string = current_time.strftime("%Y-%m-%d_%H-%M-%S")
+    random_number = random.randint(1000, 9999)
+    filename = f"{time_string}_{random_number}.{extension}"
+    result = os.path.join(folder, date_string, filename)
+    return date_string, os.path.abspath(os.path.realpath(result)), filename
+
+
+def get_files_from_folder(folder_path, exensions=None, name_filter=None):
+    if not os.path.isdir(folder_path):
+        raise ValueError("Folder path is not a valid directory.")
+
+    filenames = []
+
+    for root, dirs, files in os.walk(folder_path):
+        relative_path = os.path.relpath(root, folder_path)
+        if relative_path == ".":
+            relative_path = ""
+        for filename in files:
+            _, file_extension = os.path.splitext(filename)
+            if (exensions == None or file_extension.lower() in exensions) and (name_filter == None or name_filter in _):
+                path = os.path.join(relative_path, filename)
+                filenames.append(path)
+
+    return sorted(filenames, key=lambda x: -1 if os.sep in x else 1)