Update update_log.md (#468)

Added information about diffusion and total execution time (#463)

.github/CODEOWNERS

@@ -0,0 +1 @@
+*       @lllyasviel

.gitignore

@@ -2,7 +2,13 @@ __pycache__
 *.ckpt
 *.safetensors
 *.pth
-!taesdxl_decoder.pth
+*.bin
+*.patch
+lena.png
+lena_result.png
+lena_test.py
+user_path_config.txt
+/modules/*.png
 /repositories
 /venv
 /tmp

colab.ipynb

@@ -10,22 +10,9 @@
    "source": [
     "%cd /content\n",
     "!git clone https://github.com/lllyasviel/Fooocus\n",
-    "\n",
-    "!apt -y update -qq\n",
-    "!wget https://github.com/camenduru/gperftools/releases/download/v1.0/libtcmalloc_minimal.so.4 -O /content/libtcmalloc_minimal.so.4\n",
-    "%env LD_PRELOAD=/content/libtcmalloc_minimal.so.4\n",
-    "\n",
-    "!pip install torchsde==0.2.5 einops==0.4.1 transformers==4.30.2 safetensors==0.3.1 accelerate==0.21.0\n",
-    "!pip install pytorch_lightning==1.9.4 omegaconf==2.2.3 gradio==3.39.0 xformers==0.0.20 triton==2.0.0 pygit2==1.12.2\n",
-    "\n",
-    "!apt -y install -qq aria2\n",
-    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/ckpt/sd_xl_base_1.0/resolve/main/sd_xl_base_1.0_0.9vae.safetensors -d /content/Fooocus/models/checkpoints -o sd_xl_base_1.0_0.9vae.safetensors\n",
-    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/ckpt/sd_xl_refiner_1.0/resolve/main/sd_xl_refiner_1.0_0.9vae.safetensors -d /content/Fooocus/models/checkpoints -o sd_xl_refiner_1.0_0.9vae.safetensors\n",
-    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_offset_example-lora_1.0.safetensors -d /content/Fooocus/models/loras -o sd_xl_offset_example-lora_1.0.safetensors\n",
-    "\n",
     "%cd /content/Fooocus\n",
-    "!git pull\n",
-    "!python launch.py --share\n"
+    "!pip install pygit2==1.12.2\n",
+    "!python entry_with_update.py --share\n"
    ]
   }
  ],

fooocus_version.py

@@ -1 +1 @@
-version = '1.0.35'
+version = '2.0.78'

launch.py

@@ -1,4 +1,6 @@
 import os
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
+
 import sys
 import platform
 import fooocus_version
@@ -6,11 +8,10 @@ import fooocus_version
 from modules.launch_util import is_installed, run, python, \
     run_pip, repo_dir, git_clone, requirements_met, script_path, dir_repos
 from modules.model_loader import load_file_from_url
-from modules.path import modelfile_path, lorafile_path
+from modules.path import modelfile_path, lorafile_path, vae_approx_path, fooocus_expansion_path, upscale_models_path
 
 REINSTALL_ALL = False
 
-
 def prepare_environment():
     torch_index_url = os.environ.get('TORCH_INDEX_URL', "https://download.pytorch.org/whl/cu118")
     torch_command = os.environ.get('TORCH_COMMAND',
@@ -19,7 +20,7 @@ def prepare_environment():
 
     xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.20')
 
-    comfy_repo = os.environ.get('COMFY_REPO', "https://github.com/comfyanonymous/ComfyUI")
+    comfy_repo = os.environ.get('COMFY_REPO', "https://github.com/lllyasviel/ComfyUI_2bc12d.git")
     comfy_commit_hash = os.environ.get('COMFY_COMMIT_HASH', "2bc12d3d22efb5c63ae3a7fc342bb2dd16b31735")
 
     print(f"Python {sys.version}")
@@ -63,25 +64,53 @@ lora_filenames = [
      'https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_offset_example-lora_1.0.safetensors')
 ]
 
+vae_approx_filenames = [
+    ('xlvaeapp.pth',
+     'https://huggingface.co/lllyasviel/misc/resolve/main/xlvaeapp.pth')
+]
+
+
+upscaler_filenames = [
+    ('fooocus_upscaler_s409985e5.bin',
+     'https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_upscaler_s409985e5.bin')
+]
+
 
 def download_models():
     for file_name, url in model_filenames:
         load_file_from_url(url=url, model_dir=modelfile_path, file_name=file_name)
     for file_name, url in lora_filenames:
         load_file_from_url(url=url, model_dir=lorafile_path, file_name=file_name)
+    for file_name, url in vae_approx_filenames:
+        load_file_from_url(url=url, model_dir=vae_approx_path, file_name=file_name)
+    for file_name, url in upscaler_filenames:
+        load_file_from_url(url=url, model_dir=upscale_models_path, file_name=file_name)
+
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/misc/resolve/main/fooocus_expansion.bin',
+        model_dir=fooocus_expansion_path,
+        file_name='pytorch_model.bin'
+    )
+
     return
 
 
-def cuda_malloc():
+def clear_comfy_args():
     argv = sys.argv
     sys.argv = [sys.argv[0]]
-    import cuda_malloc
+    from comfy.cli_args import args as comfy_args
+    comfy_args.disable_cuda_malloc = True
     sys.argv = argv
 
 
+def cuda_malloc():
+    import cuda_malloc
+
+
 prepare_environment()
 
-cuda_malloc()
+clear_comfy_args()
+# cuda_malloc()
 
 download_models()
 

models/prompt_expansion/fooocus_expansion/config.json

@@ -0,0 +1,40 @@
+{
+  "_name_or_path": "gpt2",
+  "activation_function": "gelu_new",
+  "architectures": [
+    "GPT2LMHeadModel"
+  ],
+  "attn_pdrop": 0.1,
+  "bos_token_id": 50256,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 50256,
+  "pad_token_id": 50256,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "gpt2",
+  "n_ctx": 1024,
+  "n_embd": 768,
+  "n_head": 12,
+  "n_inner": null,
+  "n_layer": 12,
+  "n_positions": 1024,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "task_specific_params": {
+    "text-generation": {
+      "do_sample": true,
+      "max_length": 50
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.23.0.dev0",
+  "use_cache": true,
+  "vocab_size": 50257
+}

models/prompt_expansion/fooocus_expansion/special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

models/prompt_expansion/fooocus_expansion/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1024,
+  "name_or_path": "gpt2",
+  "special_tokens_map_file": null,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

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

@@ -8,15 +8,26 @@ outputs = []
 def worker():
     global buffer, outputs
 
+    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.path
     import modules.patch
+    import modules.virtual_memory as virtual_memory
+    import comfy.model_management
+    import modules.inpaint_worker as inpaint_worker
 
-    from modules.sdxl_styles import apply_style, aspect_ratios
+    from modules.sdxl_styles import apply_style, aspect_ratios, fooocus_expansion
     from modules.private_logger import log
+    from modules.expansion import safe_str
+    from modules.util import join_prompts, remove_empty_str, HWC3, resize_image, image_is_generated_in_current_ui
+    from modules.upscaler import perform_upscale
 
     try:
         async_gradio_app = shared.gradio_root
@@ -27,21 +38,42 @@ def worker():
     except Exception as e:
         print(e)
 
+    def progressbar(number, text):
+        print(f'[Fooocus] {text}')
+        outputs.append(['preview', (number, text, None)])
+
+    @torch.no_grad()
+    @torch.inference_mode()
     def handler(task):
-        prompt, negative_prompt, style_selction, performance_selction, \
-        aspect_ratios_selction, image_number, image_seed, sharpness, base_model_name, refiner_model_name, \
-        l1, w1, l2, w2, l3, w3, l4, w4, l5, w5 = task
+        prompt, negative_prompt, style_selections, performance_selction, \
+            aspect_ratios_selction, image_number, image_seed, sharpness, \
+            base_model_name, refiner_model_name, \
+            l1, w1, l2, w2, l3, w3, l4, w4, l5, w5, \
+            input_image_checkbox, current_tab, \
+            uov_method, uov_input_image, outpaint_selections, inpaint_input_image = task
+
+        outpaint_selections = [o.lower() for o in outpaint_selections]
 
         loras = [(l1, w1), (l2, w2), (l3, w3), (l4, w4), (l5, w5)]
+        loras_user_raw_input = copy.deepcopy(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
         modules.patch.sharpness = sharpness
-
-        pipeline.refresh_base_model(base_model_name)
-        pipeline.refresh_refiner_model(refiner_model_name)
-        pipeline.refresh_loras(loras)
-        pipeline.clean_prompt_cond_caches()
-
-        p_txt, n_txt = apply_style(style_selction, prompt, negative_prompt)
+        modules.patch.negative_adm = True
+        initial_latent = None
+        denoising_strength = 1.0
+        tiled = False
+        inpaint_worker.current_task = None
 
         if performance_selction == 'Speed':
             steps = 30
@@ -50,49 +82,304 @@ def worker():
             steps = 60
             switch = 40
 
+        pipeline.clear_all_caches()  # save memory
+
         width, height = aspect_ratios[aspect_ratios_selction]
 
-        results = []
+        if input_image_checkbox:
+            progressbar(0, 'Image processing ...')
+            if current_tab == 'uov' and uov_method != flags.disabled and uov_input_image is not None:
+                uov_input_image = HWC3(uov_input_image)
+                if 'vary' in uov_method:
+                    if not image_is_generated_in_current_ui(uov_input_image, ui_width=width, ui_height=height):
+                        uov_input_image = resize_image(uov_input_image, width=width, height=height)
+                        print(f'Resolution corrected - users are uploading their own images.')
+                    else:
+                        print(f'Processing images generated by Fooocus.')
+                    if 'subtle' in uov_method:
+                        denoising_strength = 0.5
+                    if 'strong' in uov_method:
+                        denoising_strength = 0.85
+                    initial_pixels = core.numpy_to_pytorch(uov_input_image)
+                    progressbar(0, 'VAE encoding ...')
+                    initial_latent = core.encode_vae(vae=pipeline.xl_base_patched.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))}.')
+                elif 'upscale' in uov_method:
+                    H, W, C = uov_input_image.shape
+                    progressbar(0, f'Upscaling image from {str((H, W))} ...')
+
+                    uov_input_image = core.numpy_to_pytorch(uov_input_image)
+                    uov_input_image = perform_upscale(uov_input_image)
+                    uov_input_image = core.pytorch_to_numpy(uov_input_image)[0]
+                    print(f'Image upscaled.')
+
+                    if '1.5x' in uov_method:
+                        f = 1.5
+                    elif '2x' in uov_method:
+                        f = 2.0
+                    else:
+                        f = 1.0
+
+                    width_f = int(width * f)
+                    height_f = int(height * f)
+
+                    if image_is_generated_in_current_ui(uov_input_image, ui_width=width_f, ui_height=height_f):
+                        uov_input_image = resize_image(uov_input_image, width=int(W * f), height=int(H * f))
+                        print(f'Processing images generated by Fooocus.')
+                    else:
+                        uov_input_image = resize_image(uov_input_image, width=width_f, height=height_f)
+                        print(f'Resolution corrected - users are uploading their own images.')
+
+                    H, W, C = uov_input_image.shape
+                    image_is_super_large = H * W > 2800 * 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, single_line_number=1)
+                        outputs.append(['results', [uov_input_image]])
+                        return
+
+                    tiled = True
+                    denoising_strength = 1.0 - 0.618
+                    steps = int(steps * 0.618)
+                    switch = int(steps * 0.67)
+                    initial_pixels = core.numpy_to_pytorch(uov_input_image)
+                    progressbar(0, 'VAE encoding ...')
+
+                    initial_latent = core.encode_vae(vae=pipeline.xl_base_patched.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 current_tab == 'inpaint' and isinstance(inpaint_input_image, dict):
+                inpaint_image = inpaint_input_image['image']
+                inpaint_mask = inpaint_input_image['mask'][:, :, 0]
+                if isinstance(inpaint_image, np.ndarray) and isinstance(inpaint_mask, np.ndarray) \
+                        and (np.any(inpaint_mask > 127) or len(outpaint_selections) > 0):
+                    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())
+
+                    inpaint_worker.current_task = inpaint_worker.InpaintWorker(image=inpaint_image, mask=inpaint_mask,
+                                                                               is_outpaint=len(outpaint_selections) > 0)
+
+                    # print(f'Inpaint task: {str((height, width))}')
+                    # outputs.append(['results', inpaint_worker.current_task.visualize_mask_processing()])
+                    # return
+
+                    progressbar(0, 'Downloading inpainter ...')
+                    inpaint_head_model_path, inpaint_patch_model_path = modules.path.downloading_inpaint_models()
+                    loras += [(inpaint_patch_model_path, 1.0)]
+
+                    inpaint_pixels = core.numpy_to_pytorch(inpaint_worker.current_task.image_ready)
+                    progressbar(0, 'VAE encoding ...')
+                    initial_latent = core.encode_vae(vae=pipeline.xl_base_patched.vae, pixels=inpaint_pixels)
+                    inpaint_latent = initial_latent['samples']
+                    B, C, H, W = inpaint_latent.shape
+                    inpaint_mask = core.numpy_to_pytorch(inpaint_worker.current_task.mask_ready[None])
+                    inpaint_mask = torch.nn.functional.avg_pool2d(inpaint_mask, (8, 8))
+                    inpaint_mask = torch.nn.functional.interpolate(inpaint_mask, (H, W), mode='bilinear')
+                    inpaint_worker.current_task.load_latent(latent=inpaint_latent, mask=inpaint_mask)
+
+                    progressbar(0, 'VAE inpaint encoding ...')
+
+                    inpaint_mask = (inpaint_worker.current_task.mask_ready > 0).astype(np.float32)
+                    inpaint_mask = torch.tensor(inpaint_mask).float()
+
+                    vae_dict = core.encode_vae_inpaint(
+                        mask=inpaint_mask, vae=pipeline.xl_base_patched.vae, pixels=inpaint_pixels)
+
+                    inpaint_latent = vae_dict['samples']
+                    inpaint_mask = vae_dict['noise_mask']
+                    inpaint_worker.current_task.load_inpaint_guidance(latent=inpaint_latent, mask=inpaint_mask, model_path=inpaint_head_model_path)
+
+                    B, C, H, W = inpaint_latent.shape
+                    height, width = inpaint_worker.current_task.image_raw.shape[:2]
+                    print(f'Final resolution is {str((height, width))}, latent is {str((H * 8, W * 8))}.')
+
+        progressbar(1, 'Initializing ...')
+
+        raw_prompt = prompt
+        raw_negative_prompt = negative_prompt
+
+        prompts = remove_empty_str([safe_str(p) for p in prompt.split('\n')], default='')
+        negative_prompts = remove_empty_str([safe_str(p) for p in negative_prompt.split('\n')], default='')
+
+        prompt = prompts[0]
+        negative_prompt = negative_prompts[0]
+
+        extra_positive_prompts = prompts[1:] if len(prompts) > 1 else []
+        extra_negative_prompts = negative_prompts[1:] if len(negative_prompts) > 1 else []
+
         seed = image_seed
         max_seed = int(1024 * 1024 * 1024)
-
         if not isinstance(seed, int):
             seed = random.randint(1, max_seed)
         if seed < 0:
             seed = - seed
         seed = seed % max_seed
 
+        progressbar(3, 'Loading models ...')
+
+        pipeline.refresh_everything(
+            refiner_model_name=refiner_model_name,
+            base_model_name=base_model_name,
+            loras=loras)
+
+        progressbar(3, 'Processing prompts ...')
+
+        positive_basic_workloads = []
+        negative_basic_workloads = []
+
+        if use_style:
+            for s in style_selections:
+                p, n = apply_style(s, positive=prompt)
+                positive_basic_workloads.append(p)
+                negative_basic_workloads.append(n)
+        else:
+            positive_basic_workloads.append(prompt)
+
+        negative_basic_workloads.append(negative_prompt)  # Always use independent workload for negative.
+
+        positive_basic_workloads = positive_basic_workloads + extra_positive_prompts
+        negative_basic_workloads = negative_basic_workloads + extra_negative_prompts
+
+        positive_basic_workloads = remove_empty_str(positive_basic_workloads, default=prompt)
+        negative_basic_workloads = remove_empty_str(negative_basic_workloads, default=negative_prompt)
+
+        positive_top_k = len(positive_basic_workloads)
+        negative_top_k = len(negative_basic_workloads)
+
+        tasks = [dict(
+            task_seed=seed + i,
+            positive=positive_basic_workloads,
+            negative=negative_basic_workloads,
+            expansion='',
+            c=[None, None],
+            uc=[None, None],
+        ) for i in range(image_number)]
+
+        if use_expansion:
+            for i, t in enumerate(tasks):
+                progressbar(5, f'Preparing Fooocus text #{i + 1} ...')
+                expansion = pipeline.expansion(prompt, t['task_seed'])
+                print(f'[Prompt Expansion] New suffix: {expansion}')
+                t['expansion'] = expansion
+                t['positive'] = copy.deepcopy(t['positive']) + [join_prompts(prompt, expansion)]  # Deep copy.
+
+        for i, t in enumerate(tasks):
+            progressbar(7, f'Encoding base positive #{i + 1} ...')
+            t['c'][0] = pipeline.clip_encode(sd=pipeline.xl_base_patched, texts=t['positive'],
+                                             pool_top_k=positive_top_k)
+
+        for i, t in enumerate(tasks):
+            progressbar(9, f'Encoding base negative #{i + 1} ...')
+            t['uc'][0] = pipeline.clip_encode(sd=pipeline.xl_base_patched, texts=t['negative'],
+                                              pool_top_k=negative_top_k)
+
+        if pipeline.xl_refiner is not None:
+            virtual_memory.load_from_virtual_memory(pipeline.xl_refiner.clip.cond_stage_model)
+
+            for i, t in enumerate(tasks):
+                progressbar(11, f'Encoding refiner positive #{i + 1} ...')
+                t['c'][1] = pipeline.clip_encode(sd=pipeline.xl_refiner, texts=t['positive'],
+                                                 pool_top_k=positive_top_k)
+
+            for i, t in enumerate(tasks):
+                progressbar(13, f'Encoding refiner negative #{i + 1} ...')
+                t['uc'][1] = pipeline.clip_encode(sd=pipeline.xl_refiner, texts=t['negative'],
+                                                  pool_top_k=negative_top_k)
+
+            virtual_memory.try_move_to_virtual_memory(pipeline.xl_refiner.clip.cond_stage_model)
+
+        results = []
         all_steps = steps * image_number
 
         def callback(step, x0, x, total_steps, y):
-            done_steps = i * steps + step
+            done_steps = current_task_id * steps + step
             outputs.append(['preview', (
-                int(100.0 * float(done_steps) / float(all_steps)),
-                f'Step {step}/{total_steps} in the {i}-th Sampling',
+                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 i in range(image_number):
-            imgs = pipeline.process(p_txt, n_txt, steps, switch, width, height, seed, callback=callback)
+        print(f'[ADM] Negative ADM = {modules.patch.negative_adm}')
+
+        outputs.append(['preview', (13, 'Starting tasks ...', None)])
+        for current_task_id, task in enumerate(tasks):
+            try:
+                execution_start_time = time.perf_counter()
+
+                imgs = pipeline.process_diffusion(
+                    positive_cond=task['c'],
+                    negative_cond=task['uc'],
+                    steps=steps,
+                    switch=switch,
+                    width=width,
+                    height=height,
+                    image_seed=task['task_seed'],
+                    callback=callback,
+                    latent=initial_latent,
+                    denoise=denoising_strength,
+                    tiled=tiled
+                )
+
+                if inpaint_worker.current_task is not None:
+                    imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
+
+                execution_time = time.perf_counter() - execution_start_time
+                print(f'Diffusion time: {execution_time:.2f} seconds')
 
                 for x in imgs:
                     d = [
-                    ('Prompt', prompt),
-                    ('Negative Prompt', negative_prompt),
-                    ('Style', style_selction),
+                        ('Prompt', raw_prompt),
+                        ('Negative Prompt', raw_negative_prompt),
+                        ('Fooocus V2 Expansion', task['expansion']),
+                        ('Styles', str(raw_style_selections)),
                         ('Performance', performance_selction),
                         ('Resolution', str((width, height))),
                         ('Sharpness', sharpness),
                         ('Base Model', base_model_name),
                         ('Refiner Model', refiner_model_name),
-                    ('Seed', seed)
+                        ('Seed', task['task_seed'])
                     ]
-                for n, w in loras:
+                    for n, w in loras_user_raw_input:
                         if n != 'None':
                             d.append((f'LoRA [{n}] weight', w))
-                log(x, d)
+                    log(x, d, single_line_number=3)
 
-            seed += 1
                 results += imgs
+            except comfy.model_management.InterruptProcessingException as e:
+                print('User stopped')
+                break
 
         outputs.append(['results', results])
         return

modules/core.py

@@ -1,3 +1,8 @@
+from modules.patch import patch_all
+
+patch_all()
+
+
 import os
 import random
 import einops
@@ -8,20 +13,31 @@ import comfy.model_management
 import comfy.utils
 
 from comfy.sd import load_checkpoint_guess_config
-from nodes import VAEDecode, EmptyLatentImage, CLIPTextEncode
+from nodes import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, VAEEncodeForInpaint
 from comfy.sample import prepare_mask, broadcast_cond, load_additional_models, cleanup_additional_models
-from modules.samplers_advanced import KSampler, KSamplerWithRefiner
-from modules.patch import patch_all
+from comfy.model_base import SDXLRefiner
+from comfy.sd import model_lora_keys_unet, model_lora_keys_clip, load_lora
+from modules.samplers_advanced import KSamplerBasic, KSamplerWithRefiner
 
 
-patch_all()
-opCLIPTextEncode = CLIPTextEncode()
 opEmptyLatentImage = EmptyLatentImage()
 opVAEDecode = VAEDecode()
+opVAEEncode = VAEEncode()
+opVAEDecodeTiled = VAEDecodeTiled()
+opVAEEncodeTiled = VAEEncodeTiled()
+opVAEEncodeForInpaint = VAEEncodeForInpaint()
 
 
 class StableDiffusionModel:
-    def __init__(self, unet, vae, clip, clip_vision):
+    def __init__(self, unet, vae, clip, clip_vision, model_filename=None):
+        if isinstance(model_filename, str):
+            is_refiner = isinstance(unet.model, SDXLRefiner)
+            if unet is not None:
+                unet.model.model_file = dict(filename=model_filename, prefix='model')
+            if clip is not None:
+                clip.cond_stage_model.model_file = dict(filename=model_filename, prefix='refiner_clip' if is_refiner else 'base_clip')
+            if vae is not None:
+                vae.first_stage_model.model_file = dict(filename=model_filename, prefix='first_stage_model')
         self.unet = unet
         self.vae = vae
         self.clip = clip
@@ -37,62 +53,132 @@ class StableDiffusionModel:
 
 
 @torch.no_grad()
+@torch.inference_mode()
 def load_model(ckpt_filename):
     unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename)
-    return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision)
+    return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, model_filename=ckpt_filename)
 
 
 @torch.no_grad()
-def load_lora(model, lora_filename, strength_model=1.0, strength_clip=1.0):
+@torch.inference_mode()
+def load_sd_lora(model, lora_filename, strength_model=1.0, strength_clip=1.0):
     if strength_model == 0 and strength_clip == 0:
         return model
 
-    lora = comfy.utils.load_torch_file(lora_filename, safe_load=True)
-    unet, clip = comfy.sd.load_lora_for_models(model.unet, model.clip, lora, strength_model, strength_clip)
+    lora = comfy.utils.load_torch_file(lora_filename, safe_load=False)
+
+    if lora_filename.lower().endswith('.fooocus.patch'):
+        loaded = lora
+    else:
+        key_map = model_lora_keys_unet(model.unet.model)
+        key_map = model_lora_keys_clip(model.clip.cond_stage_model, key_map)
+        loaded = load_lora(lora, key_map)
+
+    new_modelpatcher = model.unet.clone()
+    k = new_modelpatcher.add_patches(loaded, strength_model)
+
+    new_clip = model.clip.clone()
+    k1 = new_clip.add_patches(loaded, strength_clip)
+
+    k = set(k)
+    k1 = set(k1)
+    for x in loaded:
+        if (x not in k) and (x not in k1):
+            print("Lora missed: ", x)
+
+    unet, clip = new_modelpatcher, new_clip
     return StableDiffusionModel(unet=unet, clip=clip, vae=model.vae, clip_vision=model.clip_vision)
 
 
 @torch.no_grad()
-def encode_prompt_condition(clip, prompt):
-    return opCLIPTextEncode.encode(clip=clip, text=prompt)[0]
-
-
-@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()
-def decode_vae(vae, latent_image):
-    return opVAEDecode.decode(samples=latent_image, vae=vae)[0]
-
-
-def get_previewer(device, latent_format):
-    from latent_preview import TAESD, TAESDPreviewerImpl
-    taesd_decoder_path = os.path.abspath(os.path.realpath(os.path.join("models", "vae_approx",
-                                                                       latent_format.taesd_decoder_name)))
-
-    if not os.path.exists(taesd_decoder_path):
-        print(f"Warning: TAESD previews enabled, but could not find {taesd_decoder_path}")
-        return None
-
-    taesd = TAESD(None, taesd_decoder_path).to(device)
-
-    def preview_function(x0, step, total_steps):
-        global cv2_is_top
-        with torch.no_grad():
-            x_sample = taesd.decoder(torch.nn.functional.avg_pool2d(x0, kernel_size=(2, 2))).detach() * 255.0
-            x_sample = einops.rearrange(x_sample, 'b c h w -> b h w c')
-            x_sample = x_sample.cpu().numpy().clip(0, 255).astype(np.uint8)
-            return x_sample[0]
-
-    taesd.preview = preview_function
-
-    return taesd
+@torch.inference_mode()
+def decode_vae(vae, latent_image, tiled=False):
+    return (opVAEDecodeTiled if tiled else opVAEDecode).decode(samples=latent_image, vae=vae)[0]
 
 
 @torch.no_grad()
-def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sampler_name='dpmpp_2m_sde_gpu',
+@torch.inference_mode()
+def encode_vae(vae, pixels, tiled=False):
+    return (opVAEEncodeTiled if tiled else opVAEEncode).encode(pixels=pixels, vae=vae)[0]
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def encode_vae_inpaint(vae, pixels, mask):
+    return opVAEEncodeForInpaint.encode(pixels=pixels, vae=vae, mask=mask)[0]
+
+
+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_model = None
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def get_previewer(device, latent_format):
+    global VAE_approx_model
+
+    if VAE_approx_model is None:
+        from modules.path import vae_approx_path
+        vae_approx_filename = os.path.join(vae_approx_path, 'xlvaeapp.pth')
+        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 comfy.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(comfy.model_management.get_torch_device())
+
+    @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_fooocus_2m_sde_inpaint_seamless',
              scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
              force_full_denoise=False, callback_function=None):
     # SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
@@ -121,8 +207,8 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
 
     def callback(step, x0, x, total_steps):
         y = None
-        if previewer and step % 3 == 0:
-            y = previewer.preview(x0, step, total_steps)
+        if previewer is not None:
+            y = previewer(x0, step, total_steps)
         if callback_function is not None:
             callback_function(step, x0, x, total_steps, y)
         pbar.update_absolute(step + 1, total_steps, None)
@@ -144,7 +230,7 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
 
     models = load_additional_models(positive, negative, model.model_dtype())
 
-    sampler = KSampler(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler,
+    sampler = KSamplerBasic(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler,
                        denoise=denoise, model_options=model.model_options)
 
     samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image,
@@ -163,8 +249,9 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
 
 
 @torch.no_grad()
+@torch.inference_mode()
 def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive, refiner_negative, latent,
-                          seed=None, steps=30, refiner_switch_step=20, cfg=7.0, sampler_name='dpmpp_2m_sde_gpu',
+                          seed=None, steps=30, refiner_switch_step=20, cfg=7.0, sampler_name='dpmpp_fooocus_2m_sde_inpaint_seamless',
                           scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
                           force_full_denoise=False, callback_function=None):
     # SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
@@ -193,8 +280,8 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
 
     def callback(step, x0, x, total_steps):
         y = None
-        if previewer and step % 3 == 0:
-            y = previewer.preview(x0, step, total_steps)
+        if previewer is not None:
+            y = previewer(x0, step, total_steps)
         if callback_function is not None:
             callback_function(step, x0, x, total_steps, y)
         pbar.update_absolute(step + 1, total_steps, None)
@@ -240,5 +327,16 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
 
 
 @torch.no_grad()
-def image_to_numpy(x):
+@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

@@ -2,8 +2,12 @@ import modules.core as core
 import os
 import torch
 import modules.path
+import modules.virtual_memory as virtual_memory
+import comfy.model_management
 
 from comfy.model_base import SDXL, SDXLRefiner
+from modules.patch import cfg_patched, patched_model_function
+from modules.expansion import FooocusExpansion
 
 
 xl_base: core.StableDiffusionModel = None
@@ -16,12 +20,16 @@ xl_base_patched: core.StableDiffusionModel = None
 xl_base_patched_hash = ''
 
 
+@torch.no_grad()
+@torch.inference_mode()
 def refresh_base_model(name):
     global xl_base, xl_base_hash, xl_base_patched, xl_base_patched_hash
-    if xl_base_hash == str(name):
-        return
 
-    filename = os.path.join(modules.path.modelfile_path, name)
+    filename = os.path.abspath(os.path.realpath(os.path.join(modules.path.modelfile_path, name)))
+    model_hash = filename
+
+    if xl_base_hash == model_hash:
+        return
 
     if xl_base is not None:
         xl_base.to_meta()
@@ -33,22 +41,27 @@ def refresh_base_model(name):
         xl_base = None
         xl_base_hash = ''
         refresh_base_model(modules.path.default_base_model_name)
-        xl_base_hash = name
+        xl_base_hash = model_hash
         xl_base_patched = xl_base
         xl_base_patched_hash = ''
         return
 
-    xl_base_hash = name
+    xl_base_hash = model_hash
     xl_base_patched = xl_base
     xl_base_patched_hash = ''
-    print(f'Base model loaded: {xl_base_hash}')
-
+    print(f'Base model loaded: {model_hash}')
     return
 
 
+@torch.no_grad()
+@torch.inference_mode()
 def refresh_refiner_model(name):
     global xl_refiner, xl_refiner_hash
-    if xl_refiner_hash == str(name):
+
+    filename = os.path.abspath(os.path.realpath(os.path.join(modules.path.modelfile_path, name)))
+    model_hash = filename
+
+    if xl_refiner_hash == model_hash:
         return
 
     if name == 'None':
@@ -57,8 +70,6 @@ def refresh_refiner_model(name):
         print(f'Refiner unloaded.')
         return
 
-    filename = os.path.join(modules.path.modelfile_path, name)
-
     if xl_refiner is not None:
         xl_refiner.to_meta()
         xl_refiner = None
@@ -71,14 +82,16 @@ def refresh_refiner_model(name):
         print(f'Refiner unloaded.')
         return
 
-    xl_refiner_hash = name
-    print(f'Refiner model loaded: {xl_refiner_hash}')
+    xl_refiner_hash = model_hash
+    print(f'Refiner model loaded: {model_hash}')
 
     xl_refiner.vae.first_stage_model.to('meta')
     xl_refiner.vae = None
     return
 
 
+@torch.no_grad()
+@torch.inference_mode()
 def refresh_loras(loras):
     global xl_base, xl_base_patched, xl_base_patched_hash
     if xl_base_patched_hash == str(loras):
@@ -89,8 +102,14 @@ def refresh_loras(loras):
         if name == 'None':
             continue
 
+        if os.path.exists(name):
+            filename = name
+        else:
             filename = os.path.join(modules.path.lorafile_path, name)
-        model = core.load_lora(model, filename, strength_model=weight, strength_clip=weight)
+
+        assert os.path.exists(filename), 'Lora file not found!'
+
+        model = core.load_sd_lora(model, filename, strength_model=weight, strength_clip=weight)
     xl_base_patched = model
     xl_base_patched_hash = str(loras)
     print(f'LoRAs loaded: {xl_base_patched_hash}')
@@ -98,73 +117,152 @@ def refresh_loras(loras):
     return
 
 
-refresh_base_model(modules.path.default_base_model_name)
-refresh_refiner_model(modules.path.default_refiner_model_name)
-refresh_loras([(modules.path.default_lora_name, 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5)])
-
-positive_conditions_cache = None
-negative_conditions_cache = None
-positive_conditions_refiner_cache = None
-negative_conditions_refiner_cache = None
+@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
 
 
-def clean_prompt_cond_caches():
-    global positive_conditions_cache, negative_conditions_cache, \
-        positive_conditions_refiner_cache, negative_conditions_refiner_cache
-    positive_conditions_cache = None
-    negative_conditions_cache = None
-    positive_conditions_refiner_cache = None
-    negative_conditions_refiner_cache = None
+@torch.no_grad()
+@torch.inference_mode()
+def clip_encode(sd, texts, pool_top_k=1):
+    if sd is None:
+        return None
+    if sd.clip is None:
+        return None
+    if not isinstance(texts, list):
+        return None
+    if len(texts) == 0:
+        return None
+
+    clip = sd.clip
+    cond_list = []
+    pooled_acc = 0
+
+    for i, text in enumerate(texts):
+        cond, pooled = clip_encode_single(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_sd_cond_cache(sd):
+    if sd is None:
+        return None
+    if sd.clip is None:
+        return None
+    sd.clip.fcs_cond_cache = {}
     return
 
 
 @torch.no_grad()
-def process(positive_prompt, negative_prompt, steps, switch, width, height, image_seed, callback):
-    global positive_conditions_cache, negative_conditions_cache, \
-        positive_conditions_refiner_cache, negative_conditions_refiner_cache
+@torch.inference_mode()
+def clear_all_caches():
+    clear_sd_cond_cache(xl_base_patched)
+    clear_sd_cond_cache(xl_refiner)
 
-    positive_conditions = core.encode_prompt_condition(clip=xl_base_patched.clip, prompt=positive_prompt) if positive_conditions_cache is None else positive_conditions_cache
-    negative_conditions = core.encode_prompt_condition(clip=xl_base_patched.clip, prompt=negative_prompt) if negative_conditions_cache is None else negative_conditions_cache
 
-    positive_conditions_cache = positive_conditions
-    negative_conditions_cache = negative_conditions
+@torch.no_grad()
+@torch.inference_mode()
+def refresh_everything(refiner_model_name, base_model_name, loras):
+    refresh_refiner_model(refiner_model_name)
+    if xl_refiner is not None:
+        virtual_memory.try_move_to_virtual_memory(xl_refiner.unet.model)
+        virtual_memory.try_move_to_virtual_memory(xl_refiner.clip.cond_stage_model)
 
-    empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
+    refresh_base_model(base_model_name)
+    virtual_memory.load_from_virtual_memory(xl_base.unet.model)
+
+    refresh_loras(loras)
+    clear_all_caches()
+    return
+
+
+refresh_everything(
+    refiner_model_name=modules.path.default_refiner_model_name,
+    base_model_name=modules.path.default_base_model_name,
+    loras=[(modules.path.default_lora_name, 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5), ('None', 0.5)]
+)
+
+expansion = FooocusExpansion()
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def patch_all_models():
+    assert xl_base is not None
+    assert xl_base_patched is not None
+
+    xl_base.unet.model_options['sampler_cfg_function'] = cfg_patched
+    xl_base.unet.model_options['model_function_wrapper'] = patched_model_function
+
+    xl_base_patched.unet.model_options['sampler_cfg_function'] = cfg_patched
+    xl_base_patched.unet.model_options['model_function_wrapper'] = patched_model_function
 
     if xl_refiner is not None:
-        positive_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=positive_prompt) if positive_conditions_refiner_cache is None else positive_conditions_refiner_cache
-        negative_conditions_refiner = core.encode_prompt_condition(clip=xl_refiner.clip, prompt=negative_prompt) if negative_conditions_refiner_cache is None else negative_conditions_refiner_cache
+        xl_refiner.unet.model_options['sampler_cfg_function'] = cfg_patched
+        xl_refiner.unet.model_options['model_function_wrapper'] = patched_model_function
 
-        positive_conditions_refiner_cache = positive_conditions_refiner
-        negative_conditions_refiner_cache = negative_conditions_refiner
+    return
 
+
+@torch.no_grad()
+@torch.inference_mode()
+def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, latent=None, denoise=1.0, tiled=False):
+    patch_all_models()
+
+    if xl_refiner is not None:
+        virtual_memory.try_move_to_virtual_memory(xl_refiner.unet.model)
+    virtual_memory.load_from_virtual_memory(xl_base.unet.model)
+
+    if latent is None:
+        empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
+    else:
+        empty_latent = latent
+
+    if xl_refiner is not None:
         sampled_latent = core.ksampler_with_refiner(
             model=xl_base_patched.unet,
-            positive=positive_conditions,
-            negative=negative_conditions,
+            positive=positive_cond[0],
+            negative=negative_cond[0],
             refiner=xl_refiner.unet,
-            refiner_positive=positive_conditions_refiner,
-            refiner_negative=negative_conditions_refiner,
+            refiner_positive=positive_cond[1],
+            refiner_negative=negative_cond[1],
             refiner_switch_step=switch,
             latent=empty_latent,
             steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
+            denoise=denoise,
             callback_function=callback
         )
-
     else:
         sampled_latent = core.ksampler(
             model=xl_base_patched.unet,
-            positive=positive_conditions,
-            negative=negative_conditions,
+            positive=positive_cond[0],
+            negative=negative_cond[0],
             latent=empty_latent,
             steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
             seed=image_seed,
+            denoise=denoise,
             callback_function=callback
         )
 
-    decoded_latent = core.decode_vae(vae=xl_base_patched.vae, latent_image=sampled_latent)
-
-    images = core.image_to_numpy(decoded_latent)
+    decoded_latent = core.decode_vae(vae=xl_base_patched.vae, latent_image=sampled_latent, tiled=tiled)
+    images = core.pytorch_to_numpy(decoded_latent)
 
+    comfy.model_management.soft_empty_cache()
     return images

modules/expansion.py

@@ -0,0 +1,71 @@
+import torch
+
+import comfy.model_management as model_management
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed
+from modules.path import fooocus_expansion_path
+from comfy.sd import ModelPatcher
+
+
+fooocus_magic_split = [
+    ', extremely',
+    ', intricate,',
+]
+dangrous_patterns = '[]【】()（）|:：'
+
+
+def safe_str(x):
+    x = str(x)
+    for _ in range(16):
+        x = x.replace('  ', ' ')
+    return x.strip(",. \r\n")
+
+
+def remove_pattern(x, pattern):
+    for p in pattern:
+        x = x.replace(p, '')
+    return x
+
+
+class FooocusExpansion:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(fooocus_expansion_path)
+        self.model = AutoModelForCausalLM.from_pretrained(fooocus_expansion_path)
+        self.model.eval()
+
+        load_device = model_management.text_encoder_device()
+
+        if 'mps' in load_device.type:
+            load_device = torch.device('cpu')
+
+        if 'cpu' not in load_device.type and model_management.should_use_fp16():
+            self.model.half()
+
+        offload_device = model_management.text_encoder_offload_device()
+        self.patcher = ModelPatcher(self.model, load_device=load_device, offload_device=offload_device)
+
+        print(f'Fooocus Expansion engine loaded for {load_device}.')
+
+    def __call__(self, prompt, seed):
+        model_management.load_model_gpu(self.patcher)
+        seed = int(seed)
+        set_seed(seed)
+        origin = safe_str(prompt)
+        prompt = origin + fooocus_magic_split[seed % len(fooocus_magic_split)]
+
+        tokenized_kwargs = self.tokenizer(prompt, return_tensors="pt")
+        tokenized_kwargs.data['input_ids'] = tokenized_kwargs.data['input_ids'].to(self.patcher.load_device)
+        tokenized_kwargs.data['attention_mask'] = tokenized_kwargs.data['attention_mask'].to(self.patcher.load_device)
+
+        # https://huggingface.co/blog/introducing-csearch
+        # https://huggingface.co/docs/transformers/generation_strategies
+        features = self.model.generate(**tokenized_kwargs,
+                                       num_beams=1,
+                                       max_new_tokens=256,
+                                       do_sample=True)
+
+        response = self.tokenizer.batch_decode(features, skip_special_tokens=True)
+        result = response[0][len(origin):]
+        result = safe_str(result)
+        result = remove_pattern(result, dangrous_patterns)
+        return result

modules/filters.py

@@ -1,32 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-
-
-def gaussian_kernel(kernel_size, sigma):
-    kernel = np.fromfunction(
-        lambda x, y: (1 / (2 * np.pi * sigma ** 2)) *
-                     np.exp(-((x - (kernel_size - 1) / 2) ** 2 + (y - (kernel_size - 1) / 2) ** 2) / (2 * sigma ** 2)),
-        (kernel_size, kernel_size)
-    )
-    return kernel / np.sum(kernel)
-
-
-class GaussianBlur(nn.Module):
-    def __init__(self, channels, kernel_size, sigma):
-        super(GaussianBlur, self).__init__()
-        self.channels = channels
-        self.kernel_size = kernel_size
-        self.sigma = sigma
-        self.padding = kernel_size // 2  # Ensure output size matches input size
-        self.register_buffer('kernel', torch.tensor(gaussian_kernel(kernel_size, sigma), dtype=torch.float32))
-        self.kernel = self.kernel.view(1, 1, kernel_size, kernel_size)
-        self.kernel = self.kernel.expand(self.channels, -1, -1, -1)  # Repeat the kernel for each input channel
-
-    def forward(self, x):
-        x = F.conv2d(x, self.kernel.to(x), padding=self.padding, groups=self.channels)
-        return x
-
-
-gaussian_filter_2d = GaussianBlur(4, 7, 0.8)

modules/flags.py

@@ -0,0 +1,11 @@
+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
+]

modules/gradio_hijack.py

@@ -0,0 +1,448 @@
+"""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
+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
+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))

modules/html.py

@@ -83,6 +83,14 @@ progress::after {
   box-shadow: none !important;
 }
 
+.advanced_check_row{
+  width: 250px !important;
+}
+
+.min_check{
+  min-width: min(1px, 100%) !important;
+}
+
 '''
 progress_html = '''
 <div class="loader-container">

modules/inpaint_worker.py

@@ -0,0 +1,226 @@
+import os.path
+
+import torch
+import numpy as np
+import modules.default_pipeline as pipeline
+
+from PIL import Image, ImageFilter
+from modules.util import resample_image
+from modules.path import inpaint_models_path
+
+
+inpaint_head = 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 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]) - 64
+    b = np.max(indices[0]) + 65
+    c = np.min(indices[1]) - 64
+    d = np.max(indices[1]) + 65
+    a, b, c, d = regulate_abcd(x, a, b, c, d)
+    return a, b, c, d
+
+
+def area_abcd(a, b, c, d):
+    return (b - a) * (d - c)
+
+
+def solve_abcd(x, a, b, c, d, k, outpaint):
+    H, W = x.shape[:2]
+    if outpaint:
+        return 0, H, 0, W
+    min_area = H * W * k
+    max_area = H * W
+    while True:
+        if area_abcd(a, b, c, d) > min_area and abs((b - a) - (d - c)) < 16:
+            break
+        if area_abcd(a, b, c, d) >= max_area:
+            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 [(64, 4), (32, 4), (16, 4), (4, 4), (2, 4)]:
+        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, is_outpaint):
+        # mask processing
+        self.mask_raw_soft = morphological_open(mask)
+        self.mask_raw_fg = (self.mask_raw_soft == 255).astype(np.uint8) * 255
+        self.mask_raw_bg = (self.mask_raw_soft == 0).astype(np.uint8) * 255
+        self.mask_raw_trim = 255 - np.maximum(self.mask_raw_fg, self.mask_raw_bg)
+
+        # image processing
+        self.image_raw = fooocus_fill(image, self.mask_raw_fg)
+
+        # log all images
+        # imsave(self.image_raw, 'image_raw.png')
+        # imsave(self.mask_raw_soft, 'mask_raw_soft.png')
+        # imsave(self.mask_raw_fg, 'mask_raw_fg.png')
+        # imsave(self.mask_raw_bg, 'mask_raw_bg.png')
+        # imsave(self.mask_raw_trim, 'mask_raw_trim.png')
+
+        # compute abcd
+        a, b, c, d = compute_initial_abcd(self.mask_raw_bg < 127)
+        a, b, c, d = solve_abcd(self.mask_raw_bg, a, b, c, d, k=0.618, outpaint=is_outpaint)
+
+        # interested area
+        self.interested_area = (a, b, c, d)
+        self.mask_interested_soft = self.mask_raw_soft[a:b, c:d]
+        self.mask_interested_fg = self.mask_raw_fg[a:b, c:d]
+        self.mask_interested_bg = self.mask_raw_bg[a:b, c:d]
+        self.mask_interested_trim = self.mask_raw_trim[a:b, c:d]
+        self.image_interested = self.image_raw[a:b, c:d]
+
+        # resize to make images ready for diffusion
+        H, W, C = self.image_interested.shape
+        k = (1024.0 ** 2.0 / float(H * W)) ** 0.5
+        H = int(np.ceil(float(H) * k / 16.0)) * 16
+        W = int(np.ceil(float(W) * k / 16.0)) * 16
+        self.image_ready = resample_image(self.image_interested, W, H)
+        self.mask_ready = resample_image(self.mask_interested_soft, W, H)
+
+        # ending
+        self.latent = None
+        self.latent_mask = None
+        self.inpaint_head_feature = None
+        return
+
+    def load_inpaint_guidance(self, latent, mask, model_path):
+        global inpaint_head
+        if inpaint_head is None:
+            inpaint_head = InpaintHead()
+            sd = torch.load(model_path, map_location='cpu')
+            inpaint_head.load_state_dict(sd)
+        process_latent_in = pipeline.xl_base_patched.unet.model.process_latent_in
+
+        latent = process_latent_in(latent)
+        B, C, H, W = latent.shape
+
+        mask = torch.nn.functional.interpolate(mask, size=(H, W), mode="bilinear")
+        mask = mask.round()
+
+        feed = torch.cat([mask, latent], dim=1)
+
+        inpaint_head.to(device=feed.device, dtype=feed.dtype)
+        self.inpaint_head_feature = inpaint_head(feed)
+        return
+
+    def load_latent(self, latent, mask):
+        self.latent = latent
+        self.latent_mask = mask
+
+    def color_correction(self, img):
+        fg = img.astype(np.float32)
+        bg = self.image_raw.copy().astype(np.float32)
+        w = self.mask_raw_soft[:, :, 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_raw.copy()
+        result[a:b, c:d] = content
+        result = self.color_correction(result)
+        return result
+
+    def visualize_mask_processing(self):
+        result = self.image_raw // 4
+        a, b, c, d = self.interested_area
+        result[a:b, c:d] += 64
+        result[self.mask_raw_trim > 127] += 64
+        result[self.mask_raw_fg > 127] += 128
+        return [result, self.mask_raw_soft, self.image_ready, self.mask_ready]
+

modules/launch_util.py

@@ -91,9 +91,14 @@ def run(command, desc=None, errdesc=None, custom_env=None, live: bool = default_
 
 
 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
 
 
 re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")

modules/model_loader.py

@@ -1,13 +1,13 @@
 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: str | None = None,
+        file_name: Optional[str] = None,
 ) -> str:
     """Download a file from `url` into `model_dir`, using the file present if possible.
 

modules/patch.py

@@ -1,312 +1,317 @@
 import torch
+import contextlib
 import comfy.model_base
 import comfy.ldm.modules.diffusionmodules.openaimodel
 import comfy.samplers
+import comfy.k_diffusion.external
+import comfy.model_management
+import modules.anisotropic as anisotropic
+import comfy.ldm.modules.attention
+import comfy.k_diffusion.sampling
+import comfy.sd1_clip
+import modules.inpaint_worker as inpaint_worker
+import comfy.ldm.modules.diffusionmodules.openaimodel
+import comfy.ldm.modules.diffusionmodules.model
+import comfy.sd
 
-from comfy.samplers import model_management, lcm, math
+from comfy.k_diffusion import utils
+from comfy.k_diffusion.sampling import BrownianTreeNoiseSampler, trange
 from comfy.ldm.modules.diffusionmodules.openaimodel import timestep_embedding, forward_timestep_embed
-from modules.filters import gaussian_filter_2d
+
 
 sharpness = 2.0
+negative_adm = True
+
+cfg_x0 = 0.0
+cfg_s = 1.0
+cfg_cin = 1.0
 
 
-def sampling_function_patched(model_function, x, timestep, uncond, cond, cond_scale, cond_concat=None, model_options={},
-                      seed=None):
-    def get_area_and_mult(cond, x_in, cond_concat_in, timestep_in):
-        area = (x_in.shape[2], x_in.shape[3], 0, 0)
-        strength = 1.0
-        if 'timestep_start' in cond[1]:
-            timestep_start = cond[1]['timestep_start']
-            if timestep_in[0] > timestep_start:
-                return None
-        if 'timestep_end' in cond[1]:
-            timestep_end = cond[1]['timestep_end']
-            if timestep_in[0] < timestep_end:
-                return None
-        if 'area' in cond[1]:
-            area = cond[1]['area']
-        if 'strength' in cond[1]:
-            strength = cond[1]['strength']
-
-        adm_cond = None
-        if 'adm_encoded' in cond[1]:
-            adm_cond = cond[1]['adm_encoded']
-
-        input_x = x_in[:, :, area[2]:area[0] + area[2], area[3]:area[1] + area[3]]
-        if 'mask' in cond[1]:
-            # Scale the mask to the size of the input
-            # The mask should have been resized as we began the sampling process
-            mask_strength = 1.0
-            if "mask_strength" in cond[1]:
-                mask_strength = cond[1]["mask_strength"]
-            mask = cond[1]['mask']
-            assert (mask.shape[1] == x_in.shape[2])
-            assert (mask.shape[2] == x_in.shape[3])
-            mask = mask[:, area[2]:area[0] + area[2], area[3]:area[1] + area[3]] * mask_strength
-            mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
-        else:
-            mask = torch.ones_like(input_x)
-        mult = mask * strength
-
-        if 'mask' not in cond[1]:
-            rr = 8
-            if area[2] != 0:
-                for t in range(rr):
-                    mult[:, :, t:1 + t, :] *= ((1.0 / rr) * (t + 1))
-            if (area[0] + area[2]) < x_in.shape[2]:
-                for t in range(rr):
-                    mult[:, :, area[0] - 1 - t:area[0] - t, :] *= ((1.0 / rr) * (t + 1))
-            if area[3] != 0:
-                for t in range(rr):
-                    mult[:, :, :, t:1 + t] *= ((1.0 / rr) * (t + 1))
-            if (area[1] + area[3]) < x_in.shape[3]:
-                for t in range(rr):
-                    mult[:, :, :, area[1] - 1 - t:area[1] - t] *= ((1.0 / rr) * (t + 1))
-
-        conditionning = {}
-        conditionning['c_crossattn'] = cond[0]
-        if cond_concat_in is not None and len(cond_concat_in) > 0:
-            cropped = []
-            for x in cond_concat_in:
-                cr = x[:, :, area[2]:area[0] + area[2], area[3]:area[1] + area[3]]
-                cropped.append(cr)
-            conditionning['c_concat'] = torch.cat(cropped, dim=1)
-
-        if adm_cond is not None:
-            conditionning['c_adm'] = adm_cond
-
-        control = None
-        if 'control' in cond[1]:
-            control = cond[1]['control']
-
-        patches = None
-        if 'gligen' in cond[1]:
-            gligen = cond[1]['gligen']
-            patches = {}
-            gligen_type = gligen[0]
-            gligen_model = gligen[1]
-            if gligen_type == "position":
-                gligen_patch = gligen_model.set_position(input_x.shape, gligen[2], input_x.device)
-            else:
-                gligen_patch = gligen_model.set_empty(input_x.shape, input_x.device)
-
-            patches['middle_patch'] = [gligen_patch]
-
-        return (input_x, mult, conditionning, area, control, patches)
-
-    def cond_equal_size(c1, c2):
-        if c1 is c2:
-            return True
-        if c1.keys() != c2.keys():
-            return False
-        if 'c_crossattn' in c1:
-            s1 = c1['c_crossattn'].shape
-            s2 = c2['c_crossattn'].shape
-            if s1 != s2:
-                if s1[0] != s2[0] or s1[2] != s2[2]:  # these 2 cases should not happen
-                    return False
-
-                mult_min = lcm(s1[1], s2[1])
-                diff = mult_min // min(s1[1], s2[1])
-                if diff > 4:  # arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
-                    return False
-        if 'c_concat' in c1:
-            if c1['c_concat'].shape != c2['c_concat'].shape:
-                return False
-        if 'c_adm' in c1:
-            if c1['c_adm'].shape != c2['c_adm'].shape:
-                return False
-        return True
-
-    def can_concat_cond(c1, c2):
-        if c1[0].shape != c2[0].shape:
-            return False
-
-        # control
-        if (c1[4] is None) != (c2[4] is None):
-            return False
-        if c1[4] is not None:
-            if c1[4] is not c2[4]:
-                return False
-
-        # patches
-        if (c1[5] is None) != (c2[5] is None):
-            return False
-        if (c1[5] is not None):
-            if c1[5] is not c2[5]:
-                return False
-
-        return cond_equal_size(c1[2], c2[2])
-
-    def cond_cat(c_list):
-        c_crossattn = []
-        c_concat = []
-        c_adm = []
-        crossattn_max_len = 0
-        for x in c_list:
-            if 'c_crossattn' in x:
-                c = x['c_crossattn']
-                if crossattn_max_len == 0:
-                    crossattn_max_len = c.shape[1]
-                else:
-                    crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
-                c_crossattn.append(c)
-            if 'c_concat' in x:
-                c_concat.append(x['c_concat'])
-            if 'c_adm' in x:
-                c_adm.append(x['c_adm'])
-        out = {}
-        c_crossattn_out = []
-        for c in c_crossattn:
-            if c.shape[1] < crossattn_max_len:
-                c = c.repeat(1, crossattn_max_len // c.shape[1], 1)  # padding with repeat doesn't change result
-            c_crossattn_out.append(c)
-
-        if len(c_crossattn_out) > 0:
-            out['c_crossattn'] = [torch.cat(c_crossattn_out)]
-        if len(c_concat) > 0:
-            out['c_concat'] = [torch.cat(c_concat)]
-        if len(c_adm) > 0:
-            out['c_adm'] = torch.cat(c_adm)
-        return out
-
-    def calc_cond_uncond_batch(model_function, cond, uncond, x_in, timestep, max_total_area, cond_concat_in,
-                               model_options):
-        out_cond = torch.zeros_like(x_in)
-        out_count = torch.ones_like(x_in) / 100000.0
-
-        out_uncond = torch.zeros_like(x_in)
-        out_uncond_count = torch.ones_like(x_in) / 100000.0
-
-        COND = 0
-        UNCOND = 1
-
-        to_run = []
-        for x in cond:
-            p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
-            if p is None:
-                continue
-
-            to_run += [(p, COND)]
-        if uncond is not None:
-            for x in uncond:
-                p = get_area_and_mult(x, x_in, cond_concat_in, timestep)
-                if p is None:
-                    continue
-
-                to_run += [(p, UNCOND)]
-
-        while len(to_run) > 0:
-            first = to_run[0]
-            first_shape = first[0][0].shape
-            to_batch_temp = []
-            for x in range(len(to_run)):
-                if can_concat_cond(to_run[x][0], first[0]):
-                    to_batch_temp += [x]
-
-            to_batch_temp.reverse()
-            to_batch = to_batch_temp[:1]
-
-            for i in range(1, len(to_batch_temp) + 1):
-                batch_amount = to_batch_temp[:len(to_batch_temp) // i]
-                if (len(batch_amount) * first_shape[0] * first_shape[2] * first_shape[3] < max_total_area):
-                    to_batch = batch_amount
-                    break
-
-            input_x = []
-            mult = []
-            c = []
-            cond_or_uncond = []
-            area = []
-            control = None
-            patches = None
-            for x in to_batch:
-                o = to_run.pop(x)
-                p = o[0]
-                input_x += [p[0]]
-                mult += [p[1]]
-                c += [p[2]]
-                area += [p[3]]
-                cond_or_uncond += [o[1]]
-                control = p[4]
-                patches = p[5]
-
-            batch_chunks = len(cond_or_uncond)
-            input_x = torch.cat(input_x)
-            c = cond_cat(c)
-            timestep_ = torch.cat([timestep] * batch_chunks)
-
-            if control is not None:
-                c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
-
-            transformer_options = {}
-            if 'transformer_options' in model_options:
-                transformer_options = model_options['transformer_options'].copy()
-
-            if patches is not None:
-                if "patches" in transformer_options:
-                    cur_patches = transformer_options["patches"].copy()
+def calculate_weight_patched(self, patches, weight, key):
     for p in patches:
-                        if p in cur_patches:
-                            cur_patches[p] = cur_patches[p] + patches[p]
+        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:
+            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:
-                            cur_patches[p] = patches[p]
+                    weight += alpha * w1.type(weight.dtype).to(weight.device)
+        elif len(v) == 3:
+            # fooocus
+            w1 = v[0].float()
+            w_min = v[1].float()
+            w_max = v[2].float()
+            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:
-                    transformer_options["patches"] = patches
+                    weight += alpha * w1.type(weight.dtype).to(weight.device)
+        elif len(v) == 4:  # lora/locon
+            mat1 = v[0].float().to(weight.device)
+            mat2 = v[1].float().to(weight.device)
+            if v[2] is not None:
+                alpha *= v[2] / mat2.shape[0]
+            if v[3] is not None:
+                mat3 = v[3].float().to(weight.device)
+                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 len(v) == 8:  # 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
 
-            c['transformer_options'] = transformer_options
-
-            transformer_options['uc_mask'] = torch.Tensor(cond_or_uncond).to(input_x).float()[:, None, None, None]
-
-            if 'model_function_wrapper' in model_options:
-                output = model_options['model_function_wrapper'](model_function,
-                                                                 {"input": input_x, "timestep": timestep_, "c": c,
-                                                                  "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
+            if w1 is None:
+                dim = w1_b.shape[0]
+                w1 = torch.mm(w1_a.float(), w1_b.float())
             else:
-                output = model_function(input_x, timestep_, **c).chunk(batch_chunks)
-            del input_x
+                w1 = w1.float().to(weight.device)
 
-            model_management.throw_exception_if_processing_interrupted()
-
-            for o in range(batch_chunks):
-                if cond_or_uncond[o] == COND:
-                    out_cond[:, :, area[o][2]:area[o][0] + area[o][2], area[o][3]:area[o][1] + area[o][3]] += output[
-                                                                                                                  o] * \
-                                                                                                              mult[o]
-                    out_count[:, :, area[o][2]:area[o][0] + area[o][2], area[o][3]:area[o][1] + area[o][3]] += mult[o]
+            if w2 is None:
+                dim = w2_b.shape[0]
+                if t2 is None:
+                    w2 = torch.mm(w2_a.float().to(weight.device), w2_b.float().to(weight.device))
                 else:
-                    out_uncond[:, :, area[o][2]:area[o][0] + area[o][2], area[o][3]:area[o][1] + area[o][3]] += output[
-                                                                                                                    o] * \
-                                                                                                                mult[o]
-                    out_uncond_count[:, :, area[o][2]:area[o][0] + area[o][2], area[o][3]:area[o][1] + area[o][3]] += \
-                    mult[o]
-            del mult
-
-        out_cond /= out_count
-        del out_count
-        out_uncond /= out_uncond_count
-        del out_uncond_count
-
-        return out_cond, out_uncond
-
-    max_total_area = model_management.maximum_batch_area()
-    if math.isclose(cond_scale, 1.0):
-        uncond = None
-
-    cond, uncond = calc_cond_uncond_batch(model_function, cond, uncond, x, timestep, max_total_area, cond_concat,
-                                          model_options)
-    if "sampler_cfg_function" in model_options:
-        args = {"cond": cond, "uncond": uncond, "cond_scale": cond_scale, "timestep": timestep}
-        return model_options["sampler_cfg_function"](args)
+                    w2 = torch.einsum('i j k l, j r, i p -> p r k l', t2.float().to(weight.device),
+                                      w2_b.float().to(weight.device), w2_a.float().to(weight.device))
             else:
+                w2 = w2.float().to(weight.device)
+
+            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)
+        else:  # 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', t1.float().to(weight.device),
+                                  w1b.float().to(weight.device), w1a.float().to(weight.device))
+                m2 = torch.einsum('i j k l, j r, i p -> p r k l', t2.float().to(weight.device),
+                                  w2b.float().to(weight.device), w2a.float().to(weight.device))
+            else:
+                m1 = torch.mm(w1a.float().to(weight.device), w1b.float().to(weight.device))
+                m2 = torch.mm(w2a.float().to(weight.device), w2b.float().to(weight.device))
+
+            try:
+                weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
+            except Exception as e:
+                print("ERROR", key, e)
+
+    return weight
+
+
+def cfg_patched(args):
+    global cfg_x0, cfg_s
+    positive_eps = args['cond'].clone()
+    positive_x0 = args['cond'] * cfg_s + cfg_x0
+    uncond = args['uncond'] * cfg_s + cfg_x0
+    cond_scale = args['cond_scale']
+    t = args['timestep']
+
+    alpha = 1.0 - (t / 999.0)[:, None, None, None].clone()
+    alpha *= 0.001 * sharpness
+
+    eps_degraded = anisotropic.adaptive_anisotropic_filter(x=positive_eps, g=positive_x0)
+    eps_degraded_weighted = eps_degraded * alpha + positive_eps * (1.0 - alpha)
+
+    cond = eps_degraded_weighted * cfg_s + cfg_x0
+
     return uncond + (cond - uncond) * cond_scale
 
 
-def unet_forward_patched(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
-    uc_mask = transformer_options['uc_mask']
+def patched_discrete_eps_ddpm_denoiser_forward(self, input, sigma, **kwargs):
+    global cfg_x0, cfg_s, cfg_cin
+    c_out, c_in = [utils.append_dims(x, input.ndim) for x in self.get_scalings(sigma)]
+    cfg_x0 = input
+    cfg_s = c_out
+    cfg_cin = c_in
+    return self.get_eps(input * c_in, self.sigma_to_t(sigma), **kwargs)
+
+
+def patched_model_function(func, args):
+    global cfg_cin
+    x = args['input']
+    t = args['timestep']
+    c = args['c']
+    # is_uncond = torch.tensor(args['cond_or_uncond'])[:, None, None, None].to(x) * 5e-3
+    return func(x, t, **c)
+
+
+def sdxl_encode_adm_patched(self, **kwargs):
+    global negative_adm
+
+    clip_pooled = kwargs["pooled_output"]
+    width = kwargs.get("width", 768)
+    height = kwargs.get("height", 768)
+    crop_w = kwargs.get("crop_w", 0)
+    crop_h = kwargs.get("crop_h", 0)
+    target_width = kwargs.get("target_width", width)
+    target_height = kwargs.get("target_height", height)
+
+    if negative_adm:
+        if kwargs.get("prompt_type", "") == "negative":
+            width *= 0.8
+            height *= 0.8
+        elif kwargs.get("prompt_type", "") == "positive":
+            width *= 1.5
+            height *= 1.5
+
+    out = []
+    out.append(self.embedder(torch.Tensor([height])))
+    out.append(self.embedder(torch.Tensor([width])))
+    out.append(self.embedder(torch.Tensor([crop_h])))
+    out.append(self.embedder(torch.Tensor([crop_w])))
+    out.append(self.embedder(torch.Tensor([target_height])))
+    out.append(self.embedder(torch.Tensor([target_width])))
+    flat = torch.flatten(torch.cat(out))[None, ]
+    return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
+
+
+def text_encoder_device_patched():
+    # Fooocus's style system uses text encoder much more times than comfy so this makes things much faster.
+    return comfy.model_management.get_torch_device()
+
+
+def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
+    to_encode = list(self.empty_tokens)
+    for x in token_weight_pairs:
+        tokens = list(map(lambda a: a[0], x))
+        to_encode.append(tokens)
+
+    out, pooled = self.encode(to_encode)
+
+    z_empty = out[0:1]
+    if pooled.shape[0] > 1:
+        first_pooled = pooled[1:2]
+    else:
+        first_pooled = pooled[0:1]
+
+    output = []
+    for k in range(1, out.shape[0]):
+        z = out[k:k + 1]
+        original_mean = z.mean()
+
+        for i in range(len(z)):
+            for j in range(len(z[i])):
+                weight = token_weight_pairs[k - 1][j][1]
+                z[i][j] = (z[i][j] - z_empty[0][j]) * weight + z_empty[0][j]
+
+        new_mean = z.mean()
+        z = z * (original_mean / new_mean)
+        output.append(z)
+
+    if len(output) == 0:
+        return z_empty.cpu(), first_pooled.cpu()
+    return torch.cat(output, dim=-2).cpu(), first_pooled.cpu()
+
+
+@torch.no_grad()
+def sample_dpmpp_fooocus_2m_sde_inpaint_seamless(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, **kwargs):
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
+
+    seed = extra_args.get("seed", None)
+    assert isinstance(seed, int)
+
+    energy_generator = torch.Generator(device='cpu')
+    energy_generator.manual_seed(seed + 1)  # avoid bad results by using different seeds.
+
+    def get_energy():
+        return torch.randn(x.size(), dtype=x.dtype, generator=energy_generator, device="cpu").to(x)
+
+    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    old_denoised, h_last, h = None, None, None
+
+    latent_processor = model.inner_model.inner_model.inner_model.process_latent_in
+    inpaint_latent = None
+    inpaint_mask = None
+
+    if inpaint_worker.current_task is not None:
+        inpaint_latent = latent_processor(inpaint_worker.current_task.latent).to(x)
+        inpaint_mask = inpaint_worker.current_task.latent_mask.to(x)
+
+    def blend_latent(a, b, w):
+        return a * w + b * (1 - w)
+
+    for i in trange(len(sigmas) - 1, disable=disable):
+        if inpaint_latent is None:
+            denoised = model(x, sigmas[i] * s_in, **extra_args)
+        else:
+            energy = get_energy() * sigmas[i] + inpaint_latent
+            x_prime = blend_latent(x, energy, inpaint_mask)
+            denoised = model(x_prime, sigmas[i] * s_in, **extra_args)
+            denoised = blend_latent(denoised, inpaint_latent, inpaint_mask)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+        if sigmas[i + 1] == 0:
+            x = denoised
+        else:
+            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
+            h = s - t
+            eta_h = eta * h
+
+            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
+            if old_denoised is not None:
+                r = h_last / h
+                x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
+
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (
+                        -2 * eta_h).expm1().neg().sqrt() * s_noise
+
+        old_denoised = denoised
+        h_last = h
+
+    return x
+
+
+def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+    inpaint_fix = None
+    if inpaint_worker.current_task is not None:
+        inpaint_fix = inpaint_worker.current_task.inpaint_head_feature
+
     transformer_options["original_shape"] = list(x.shape)
     transformer_options["current_index"] = 0
 
+    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 = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(self.dtype)
     emb = self.time_embed(t_emb)
@@ -319,6 +324,12 @@ def unet_forward_patched(self, x, timesteps=None, context=None, y=None, control=
     for id, module in enumerate(self.input_blocks):
         transformer_options["block"] = ("input", id)
         h = forward_timestep_embed(module, h, emb, context, transformer_options)
+
+        if inpaint_fix is not None:
+            if int(h.shape[1]) == int(inpaint_fix.shape[1]):
+                h = h + inpaint_fix.to(h)
+                inpaint_fix = None
+
         if control is not None and 'input' in control and len(control['input']) > 0:
             ctrl = control['input'].pop()
             if ctrl is not None:
@@ -345,45 +356,94 @@ def unet_forward_patched(self, x, timesteps=None, context=None, y=None, control=
             output_shape = None
         h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape)
     h = h.type(x.dtype)
-    x0 = self.out(h)
-
-    alpha = 1.0 - (timesteps / 999.0)[:, None, None, None].clone()
-    alpha *= 0.001 * sharpness
-    degraded_x0 = gaussian_filter_2d(x0) * alpha + x0 * (1.0 - alpha)
-
-    x0 = x0 * uc_mask + degraded_x0 * (1.0 - uc_mask)
-
-    return x0
+    if self.predict_codebook_ids:
+        return self.id_predictor(h)
+    else:
+        return self.out(h)
 
 
-def sdxl_encode_adm_patched(self, **kwargs):
-    clip_pooled = kwargs["pooled_output"]
-    width = kwargs.get("width", 768)
-    height = kwargs.get("height", 768)
-    crop_w = kwargs.get("crop_w", 0)
-    crop_h = kwargs.get("crop_h", 0)
-    target_width = kwargs.get("target_width", width)
-    target_height = kwargs.get("target_height", height)
+def patched_SD1ClipModel_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)
 
-    if kwargs.get("prompt_type", "") == "negative":
-        width *= 0.8
-        height *= 0.8
-    elif kwargs.get("prompt_type", "") == "positive":
-        width *= 1.5
-        height *= 1.5
+        if backup_embeds.weight.dtype != torch.float32:
+            precision_scope = torch.autocast
+        else:
+            precision_scope = contextlib.nullcontext
 
-    out = []
-    out.append(self.embedder(torch.Tensor([height])))
-    out.append(self.embedder(torch.Tensor([width])))
-    out.append(self.embedder(torch.Tensor([crop_h])))
-    out.append(self.embedder(torch.Tensor([crop_w])))
-    out.append(self.embedder(torch.Tensor([target_height])))
-    out.append(self.embedder(torch.Tensor([target_width])))
-    flat = torch.flatten(torch.cat(out))[None, ]
-    return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
+        with precision_scope(comfy.model_management.get_autocast_device(device)):
+            outputs = self.transformer(input_ids=tokens, 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)
+
+            pooled_output = outputs.pooler_output
+            if self.text_projection is not None:
+                pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
+        return z.float(), pooled_output.float()
+
+
+VAE_DTYPE = None
+
+
+def vae_dtype_patched():
+    global VAE_DTYPE
+    if VAE_DTYPE is None:
+        VAE_DTYPE = torch.float32
+        if comfy.model_management.is_nvidia():
+            torch_version = torch.version.__version__
+            if int(torch_version[0]) >= 2:
+                if torch.cuda.is_bf16_supported():
+                    VAE_DTYPE = torch.bfloat16
+                    print('BFloat16 VAE: Enabled')
+    return VAE_DTYPE
+
+
+def vae_bf16_upsample_forward(self, x):
+    try:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+    except:  # operation not implemented for bf16
+        b, c, h, w = x.shape
+        out = torch.empty((b, c, h * 2, w * 2), dtype=x.dtype, layout=x.layout, device=x.device)
+        split = 8
+        l = out.shape[1] // split
+        for i in range(0, out.shape[1], l):
+            out[:, i:i + l] = torch.nn.functional.interpolate(x[:, i:i + l].to(torch.float32), scale_factor=2.0,
+                                                              mode="nearest").to(x.dtype)
+        del x
+        x = out
+
+    if self.with_conv:
+        x = self.conv(x)
+    return x
 
 
 def patch_all():
-    comfy.samplers.sampling_function = sampling_function_patched
+    comfy.model_management.vae_dtype = vae_dtype_patched
+    comfy.ldm.modules.diffusionmodules.model.Upsample.forward = vae_bf16_upsample_forward
+
+    comfy.sd1_clip.SD1ClipModel.forward = patched_SD1ClipModel_forward
+
+    comfy.sd.ModelPatcher.calculate_weight = calculate_weight_patched
+    comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
+
+    comfy.ldm.modules.attention.print = lambda x: None
+    comfy.k_diffusion.sampling.sample_dpmpp_fooocus_2m_sde_inpaint_seamless = sample_dpmpp_fooocus_2m_sde_inpaint_seamless
+
+    comfy.model_management.text_encoder_device = text_encoder_device_patched
+    print(f'Fooocus Text Processing Pipelines are retargeted to {str(comfy.model_management.text_encoder_device())}')
+
+    comfy.k_diffusion.external.DiscreteEpsDDPMDenoiser.forward = patched_discrete_eps_ddpm_denoiser_forward
     comfy.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
-    comfy.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = unet_forward_patched
+
+    comfy.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights_patched_with_a1111_method
+    return

modules/path.py

@@ -1,8 +1,46 @@
 import os
+import json
+from modules.model_loader import load_file_from_url
+
+
+config_path = "user_path_config.txt"
+config_dict = {}
+
+
+try:
+    if os.path.exists(config_path):
+        with open(config_path, "r", encoding="utf-8") as json_file:
+            config_dict = json.load(json_file)
+except Exception as e:
+    print('Load path config failed')
+    print(e)
+
+
+def get_config_or_set_default(key, default):
+    global config_dict
+    v = config_dict.get(key, None)
+    if isinstance(v, str) and os.path.exists(v) and os.path.isdir(v):
+        return v
+    else:
+        dp = os.path.abspath(os.path.join(os.path.dirname(__file__), default))
+        os.makedirs(dp, exist_ok=True)
+        config_dict[key] = dp
+        return dp
+
+
+modelfile_path = get_config_or_set_default('modelfile_path', '../models/checkpoints/')
+lorafile_path = get_config_or_set_default('lorafile_path', '../models/loras/')
+vae_approx_path = get_config_or_set_default('vae_approx_path', '../models/vae_approx/')
+upscale_models_path = get_config_or_set_default('upscale_models_path', '../models/upscale_models/')
+inpaint_models_path = get_config_or_set_default('inpaint_models_path', '../models/inpaint/')
+fooocus_expansion_path = get_config_or_set_default('fooocus_expansion_path',
+                                                   '../models/prompt_expansion/fooocus_expansion')
+
+temp_outputs_path = get_config_or_set_default('temp_outputs_path', '../outputs/')
+
+with open(config_path, "w", encoding="utf-8") as json_file:
+    json.dump(config_dict, json_file, indent=4)
 
-modelfile_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../models/checkpoints/'))
-lorafile_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../models/loras/'))
-temp_outputs_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../outputs/'))
 
 os.makedirs(temp_outputs_path, exist_ok=True)
 
@@ -22,9 +60,10 @@ def get_model_filenames(folder_path):
     filenames = []
     for filename in os.listdir(folder_path):
         if os.path.isfile(os.path.join(folder_path, filename)):
-            _, file_extension = os.path.splitext(filename)
-            if file_extension.lower() in ['.pth', '.ckpt', '.bin', '.safetensors']:
+            for ends in ['.pth', '.ckpt', '.bin', '.safetensors', '.fooocus.patch']:
+                if filename.lower().endswith(ends):
                     filenames.append(filename)
+                    break
 
     return filenames
 
@@ -36,4 +75,18 @@ def update_all_model_names():
     return
 
 
+def downloading_inpaint_models():
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/fooocus_inpaint_head.pth',
+        model_dir=inpaint_models_path,
+        file_name='fooocus_inpaint_head.pth'
+    )
+    load_file_from_url(
+        url='https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/inpaint.fooocus.patch',
+        model_dir=inpaint_models_path,
+        file_name='inpaint.fooocus.patch'
+    )
+    return os.path.join(inpaint_models_path, 'fooocus_inpaint_head.pth'), os.path.join(inpaint_models_path, 'inpaint.fooocus.patch')
+
+
 update_all_model_names()

modules/private_logger.py

@@ -5,27 +5,27 @@ from PIL import Image
 from modules.util import generate_temp_filename
 
 
-def log(img, dic):
+def log(img, dic, single_line_number=3):
     date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.path.temp_outputs_path, 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')
 
     if not os.path.exists(html_name):
-        with open(html_name, 'a+') as f:
+        with open(html_name, 'a+', encoding='utf-8') as f:
             f.write(f"<p>Fooocus Log {date_string} (private)</p>\n")
             f.write(f"<p>All images do not contain any hidden data.</p>")
 
-    with open(html_name, 'a+') as f:
+    with open(html_name, 'a+', encoding='utf-8') as f:
         div_name = only_name.replace('.', '_')
         f.write(f'<div id="{div_name}"><hr>\n')
         f.write(f"<p>{only_name}</p>\n")
         i = 0
         for k, v in dic:
-            if i < 2:
+            if i < single_line_number:
                 f.write(f"<p>{k}: <b>{v}</b> </p>\n")
             else:
-                if i % 2 == 0:
+                if (i - single_line_number) % 2 == 0:
                     f.write(f"<p>{k}: <b>{v}</b>, ")
                 else:
                     f.write(f"{k}: <b>{v}</b></p>\n")

modules/samplers_advanced.py

@@ -1,13 +1,212 @@
 from comfy.samplers import *
 
 import comfy.model_management
+import modules.virtual_memory
+
+
+class KSamplerBasic:
+    SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
+    SAMPLERS = ["euler", "euler_ancestral", "heun", "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", "ddim", "uni_pc", "uni_pc_bh2", "dpmpp_fooocus_2m_sde_inpaint_seamless"]
+
+    def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
+        self.model = model
+        self.model_denoise = CFGNoisePredictor(self.model)
+        if self.model.model_type == model_base.ModelType.V_PREDICTION:
+            self.model_wrap = CompVisVDenoiser(self.model_denoise, quantize=True)
+        else:
+            self.model_wrap = k_diffusion_external.CompVisDenoiser(self.model_denoise, quantize=True)
+
+        self.model_k = KSamplerX0Inpaint(self.model_wrap)
+        self.device = device
+        if scheduler not in self.SCHEDULERS:
+            scheduler = self.SCHEDULERS[0]
+        if sampler not in self.SAMPLERS:
+            sampler = self.SAMPLERS[0]
+        self.scheduler = scheduler
+        self.sampler = sampler
+        self.sigma_min=float(self.model_wrap.sigma_min)
+        self.sigma_max=float(self.model_wrap.sigma_max)
+        self.set_steps(steps, denoise)
+        self.denoise = denoise
+        self.model_options = model_options
+
+    def calculate_sigmas(self, steps):
+        sigmas = None
+
+        discard_penultimate_sigma = False
+        if self.sampler in ['dpm_2', 'dpm_2_ancestral']:
+            steps += 1
+            discard_penultimate_sigma = True
+
+        if self.scheduler == "karras":
+            sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
+        elif self.scheduler == "exponential":
+            sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=self.sigma_min, sigma_max=self.sigma_max)
+        elif self.scheduler == "normal":
+            sigmas = self.model_wrap.get_sigmas(steps)
+        elif self.scheduler == "simple":
+            sigmas = simple_scheduler(self.model_wrap, steps)
+        elif self.scheduler == "ddim_uniform":
+            sigmas = ddim_scheduler(self.model_wrap, steps)
+        else:
+            print("error invalid scheduler", self.scheduler)
+
+        if discard_penultimate_sigma:
+            sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
+        return sigmas
+
+    def set_steps(self, steps, denoise=None):
+        self.steps = steps
+        if denoise is None or denoise > 0.9999:
+            self.sigmas = self.calculate_sigmas(steps).to(self.device)
+        else:
+            new_steps = int(steps/denoise)
+            sigmas = self.calculate_sigmas(new_steps).to(self.device)
+            self.sigmas = sigmas[-(steps + 1):]
+
+    def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
+        if sigmas is None:
+            sigmas = self.sigmas
+        sigma_min = self.sigma_min
+
+        if last_step is not None and last_step < (len(sigmas) - 1):
+            sigma_min = sigmas[last_step]
+            sigmas = sigmas[:last_step + 1]
+            if force_full_denoise:
+                sigmas[-1] = 0
+
+        if start_step is not None:
+            if start_step < (len(sigmas) - 1):
+                sigmas = sigmas[start_step:]
+            else:
+                if latent_image is not None:
+                    return latent_image
+                else:
+                    return torch.zeros_like(noise)
+
+        positive = positive[:]
+        negative = negative[:]
+
+        resolve_cond_masks(positive, noise.shape[2], noise.shape[3], self.device)
+        resolve_cond_masks(negative, noise.shape[2], noise.shape[3], self.device)
+
+        calculate_start_end_timesteps(self.model_wrap, negative)
+        calculate_start_end_timesteps(self.model_wrap, positive)
+
+        #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(self.model_wrap, negative + positive)
+
+        apply_empty_x_to_equal_area(list(filter(lambda c: c[1].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])
+
+        if self.model.is_adm():
+            positive = encode_adm(self.model, positive, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "positive")
+            negative = encode_adm(self.model, negative, noise.shape[0], noise.shape[3], noise.shape[2], self.device, "negative")
+
+        if latent_image is not None:
+            latent_image = self.model.process_latent_in(latent_image)
+
+        extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": self.model_options, "seed":seed}
+
+        cond_concat = None
+        if hasattr(self.model, 'concat_keys'): #inpaint
+            cond_concat = []
+            for ck in self.model.concat_keys:
+                if denoise_mask is not None:
+                    if ck == "mask":
+                        cond_concat.append(denoise_mask[:,:1])
+                    elif ck == "masked_image":
+                        cond_concat.append(latent_image) #NOTE: the latent_image should be masked by the mask in pixel space
+                else:
+                    if ck == "mask":
+                        cond_concat.append(torch.ones_like(noise)[:,:1])
+                    elif ck == "masked_image":
+                        cond_concat.append(blank_inpaint_image_like(noise))
+            extra_args["cond_concat"] = cond_concat
+
+        if sigmas[0] != self.sigmas[0] or (self.denoise is not None and self.denoise < 1.0):
+            max_denoise = False
+        else:
+            max_denoise = True
+
+
+        if self.sampler == "uni_pc":
+            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)
+        elif self.sampler == "uni_pc_bh2":
+            samples = uni_pc.sample_unipc(self.model_wrap, noise, latent_image, sigmas, sampling_function=sampling_function, max_denoise=max_denoise, extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
+        elif self.sampler == "ddim":
+            timesteps = []
+            for s in range(sigmas.shape[0]):
+                timesteps.insert(0, self.model_wrap.sigma_to_discrete_timestep(sigmas[s]))
+            noise_mask = None
+            if denoise_mask is not None:
+                noise_mask = 1.0 - denoise_mask
+
+            ddim_callback = None
+            if callback is not None:
+                total_steps = len(timesteps) - 1
+                ddim_callback = lambda pred_x0, i: callback(i, pred_x0, None, total_steps)
+
+            sampler = DDIMSampler(self.model, device=self.device)
+            sampler.make_schedule_timesteps(ddim_timesteps=timesteps, verbose=False)
+            z_enc = sampler.stochastic_encode(latent_image, torch.tensor([len(timesteps) - 1] * noise.shape[0]).to(self.device), noise=noise, max_denoise=max_denoise)
+            samples, _ = sampler.sample_custom(ddim_timesteps=timesteps,
+                                                    conditioning=positive,
+                                                    batch_size=noise.shape[0],
+                                                    shape=noise.shape[1:],
+                                                    verbose=False,
+                                                    unconditional_guidance_scale=cfg,
+                                                    unconditional_conditioning=negative,
+                                                    eta=0.0,
+                                                    x_T=z_enc,
+                                                    x0=latent_image,
+                                                    img_callback=ddim_callback,
+                                                    denoise_function=self.model_wrap.predict_eps_discrete_timestep,
+                                                    extra_args=extra_args,
+                                                    mask=noise_mask,
+                                                    to_zero=sigmas[-1]==0,
+                                                    end_step=sigmas.shape[0] - 1,
+                                                    disable_pbar=disable_pbar)
+
+        else:
+            extra_args["denoise_mask"] = denoise_mask
+            self.model_k.latent_image = latent_image
+            self.model_k.noise = noise
+
+            if max_denoise:
+                noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
+            else:
+                noise = noise * sigmas[0]
+
+            k_callback = None
+            total_steps = len(sigmas) - 1
+            if callback is not None:
+                k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
+
+            if latent_image is not None:
+                noise += latent_image
+            if self.sampler == "dpm_fast":
+                samples = k_diffusion_sampling.sample_dpm_fast(self.model_k, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+            elif self.sampler == "dpm_adaptive":
+                samples = k_diffusion_sampling.sample_dpm_adaptive(self.model_k, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+            else:
+                samples = getattr(k_diffusion_sampling, "sample_{}".format(self.sampler))(self.model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar)
+
+        return self.model.process_latent_out(samples.to(torch.float32))
 
 
 class KSamplerWithRefiner:
     SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
     SAMPLERS = ["euler", "euler_ancestral", "heun", "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", "ddim", "uni_pc", "uni_pc_bh2"]
+                "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2", "dpmpp_fooocus_2m_sde_inpaint_seamless"]
 
     def __init__(self, model, refiner_model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
         self.model_patcher = model
@@ -152,6 +351,8 @@ class KSamplerWithRefiner:
                                           noise.shape[3], noise.shape[2], self.device, "negative")
 
         def refiner_switch():
+            modules.virtual_memory.try_move_to_virtual_memory(self.model_denoise.inner_model)
+            modules.virtual_memory.load_from_virtual_memory(self.refiner_model_denoise.inner_model)
             comfy.model_management.load_model_gpu(self.refiner_model_patcher)
             self.model_denoise.inner_model = self.refiner_model_denoise.inner_model
             for i in range(len(positive)):

modules/sdxl_styles.py

@@ -1,13 +1,14 @@
+from modules.util import join_prompts
+
+
+fooocus_expansion = "Fooocus V2"
+default_styles = ["Default (Slightly Cinematic)"]
+
 # https://github.com/twri/sdxl_prompt_styler/blob/main/sdxl_styles.json
 
 styles = [
     {
-        "name": "None",
-        "prompt": "{prompt}",
-        "negative_prompt": ""
-    },
-    {
-        "name": "cinematic-default",
+        "name": "Default (Slightly Cinematic)",
         "prompt": "cinematic still {prompt} . emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
         "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured"
     },
@@ -525,13 +526,419 @@ styles = [
         "name": "photo-tilt-shift",
         "prompt": "tilt-shift photo of {prompt} . Selective focus, miniature effect, blurred background, highly detailed, vibrant, perspective control",
         "negative_prompt": "blurry, noisy, deformed, flat, low contrast, unrealistic, oversaturated, underexposed"
+    },
+    {
+        "name": "cinematic-diva",
+        "prompt": "UHD, 8K, ultra detailed, a cinematic photograph of {prompt}, beautiful lighting, great composition",
+        "negative_prompt": "ugly, deformed, noisy, blurry, NSFW"
+    },
+    {
+        "name": "Abstract Expressionism",
+        "prompt": "Abstract Expressionism Art, {prompt}, High contrast, minimalistic, colorful, stark, dramatic, expressionism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realism, photorealistic"
+    },
+    {
+        "name": "Academia",
+        "prompt": "Academia, {prompt}, preppy Ivy League style, stark, dramatic, chic boarding school, academia",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, grunge, sloppy, unkempt"
+    },
+    {
+        "name": "Action Figure",
+        "prompt": "Action Figure, {prompt}, plastic collectable action figure, collectable toy action figure",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Adorable 3D Character",
+        "prompt": "Adorable 3D Character, {prompt}, 3D render, adorable character, 3D art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, grunge, sloppy, unkempt, photograph, photo, realistic"
+    },
+    {
+        "name": "Adorable Kawaii",
+        "prompt": "Adorable Kawaii, {prompt}, pretty, cute, adorable, kawaii",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, gothic, dark, moody, monochromatic"
+    },
+    {
+        "name": "Art Deco",
+        "prompt": "Art Deco, {prompt}, sleek, geometric forms, art deco style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Art Nouveau",
+        "prompt": "Art Nouveau, beautiful art, {prompt}, sleek, organic forms, long, sinuous, art nouveau style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, industrial, mechanical"
+    },
+    {
+        "name": "Astral Aura",
+        "prompt": "Astral Aura, {prompt}, astral, colorful aura, vibrant energy",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Avant-garde",
+        "prompt": "Avant-garde, {prompt}, unusual, experimental, avant-garde art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Baroque",
+        "prompt": "Baroque, {prompt}, dramatic, exuberant, grandeur, baroque art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Bauhaus-Style Poster",
+        "prompt": "Bauhaus-Style Poster, {prompt}, simple geometric shapes, clean lines, primary colors, Bauhaus-Style Poster",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Blueprint Schematic Drawing",
+        "prompt": "Blueprint Schematic Drawing, {prompt}, technical drawing, blueprint, schematic",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Caricature",
+        "prompt": "Caricature, {prompt}, exaggerated, comical, caricature",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realistic"
+    },
+    {
+        "name": "Cel Shaded Art",
+        "prompt": "Cel Shaded Art, {prompt}, 2D, flat color, toon shading, cel shaded style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Character Design Sheet",
+        "prompt": "Character Design Sheet, {prompt}, character reference sheet, character turn around",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Classicism Art",
+        "prompt": "Classicism Art, {prompt}, inspired by Roman and Greek culture, clarity, harmonious, classicism art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Color Field Painting",
+        "prompt": "Color Field Painting, {prompt}, abstract, simple, geometic, color field painting style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Colored Pencil Art",
+        "prompt": "Colored Pencil Art, {prompt}, colored pencil strokes, light color, visible paper texture, colored pencil art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Conceptual Art",
+        "prompt": "Conceptual Art, {prompt}, concept art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Constructivism",
+        "prompt": "Constructivism Art, {prompt}, minimalistic, geometric forms, constructivism art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Cubism",
+        "prompt": "Cubism Art, {prompt}, flat geometric forms, cubism art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Dadaism",
+        "prompt": "Dadaism Art, {prompt}, satirical, nonsensical, dadaism art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Dark Fantasy",
+        "prompt": "Dark Fantasy Art, {prompt}, dark, moody, dark fantasy style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, bright, sunny"
+    },
+    {
+        "name": "Dark Moody Atmosphere",
+        "prompt": "Dark Moody Atmosphere, {prompt}, dramatic, mysterious, dark moody atmosphere",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, vibrant, colorful, bright"
+    },
+    {
+        "name": "DMT Art Style",
+        "prompt": "DMT Art Style, {prompt}, bright colors, surreal visuals, swirling patterns, DMT art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Doodle Art",
+        "prompt": "Doodle Art Style, {prompt}, drawing, freeform, swirling patterns, doodle art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Double Exposure",
+        "prompt": "Double Exposure Style, {prompt}, double image ghost effect, image combination, double exposure style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Dripping Paint Splatter Art",
+        "prompt": "Dripping Paint Splatter Art, {prompt}, dramatic, paint drips, splatters, dripping paint",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Expressionism",
+        "prompt": "Expressionism Art Style, {prompt}, movement, contrast, emotional, exaggerated forms, expressionism art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Faded Polaroid Photo",
+        "prompt": "Faded Polaroid Photo, {prompt}, analog, old faded photo, old polaroid",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, vibrant, colorful"
+    },
+    {
+        "name": "Fauvism",
+        "prompt": "Fauvism Art, {prompt}, painterly, bold colors, textured brushwork, fauvism art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Flat 2D Art",
+        "prompt": "Flat 2D Art, {prompt}, simple flat color, 2-dimensional, Flat 2D Art Style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, 3D, photo, realistic"
+    },
+    {
+        "name": "Fortnite Art Style",
+        "prompt": "Fortnite Art Style, {prompt}, 3D cartoon, colorful, Fortnite Art Style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photo, realistic"
+    },
+    {
+        "name": "Futurism",
+        "prompt": "Futurism Art Style, {prompt}, dynamic, dramatic, Futurism Art Style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Glitchcore",
+        "prompt": "Glitchcore Art Style, {prompt}, dynamic, dramatic, distorted, vibrant colors, glitchcore art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Glo-fi",
+        "prompt": "Glo-fi Art Style, {prompt}, dynamic, dramatic, vibrant colors, glo-fi art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Googie Art Style",
+        "prompt": "Googie Art Style, {prompt}, dynamic, dramatic, 1950's futurism, bold boomerang angles, Googie art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Graffiti Art",
+        "prompt": "Graffiti Art Style, {prompt}, dynamic, dramatic, vibrant colors, graffiti art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Harlem Renaissance Art",
+        "prompt": "Harlem Renaissance Art Style, {prompt}, dynamic, dramatic, 1920s African American culture, Harlem Renaissance art style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "High Fashion",
+        "prompt": "High Fashion, {prompt}, dynamic, dramatic, haute couture, elegant, ornate clothing, High Fashion",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Idyllic",
+        "prompt": "Idyllic, {prompt}, peaceful, happy, pleasant, happy, harmonious, picturesque, charming",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Impressionism",
+        "prompt": "Impressionism, {prompt}, painterly, small brushstrokes, visible brushstrokes, impressionistic style",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Infographic Drawing",
+        "prompt": "Infographic Drawing, {prompt}, diagram, infographic",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Ink Dripping Drawing",
+        "prompt": "Ink Dripping Drawing, {prompt}, ink drawing, dripping ink",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, colorful, vibrant"
+    },
+    {
+        "name": "Japanese Ink Drawing",
+        "prompt": "Japanese Ink Drawing, {prompt}, ink drawing, inkwash, Japanese Ink Drawing",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, colorful, vibrant"
+    },
+    {
+        "name": "Knolling Photography",
+        "prompt": "Knolling Photography, {prompt}, flat lay photography, object arrangment, knolling photography",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Light Cheery Atmosphere",
+        "prompt": "Light Cheery Atmosphere, {prompt}, happy, joyful, cheerful, carefree, gleeful, lighthearted, pleasant atmosphere",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, monochromatic, dark, moody"
+    },
+    {
+        "name": "Logo Design",
+        "prompt": "Logo Design, {prompt}, dynamic graphic art, vector art, minimalist, professional logo design",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Luxurious Elegance",
+        "prompt": "Luxurious Elegance, {prompt}, extravagant, ornate, designer, opulent, picturesque, lavish",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Macro Photography",
+        "prompt": "Macro Photography, {prompt}, close-up, macro 100mm, macro photography",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Mandola Art",
+        "prompt": "Mandola art style, {prompt}, complex, circular design, mandola",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Marker Drawing",
+        "prompt": "Marker Drawing, {prompt}, bold marker lines, visibile paper texture, marker drawing",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photograph, realistic"
+    },
+    {
+        "name": "Medievalism",
+        "prompt": "Medievalism, {prompt}, inspired by The Middle Ages, medieval art, elaborate patterns and decoration, Medievalism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Minimalism",
+        "prompt": "Minimalism, {prompt}, abstract, simple geometic shapes, hard edges, sleek contours, Minimalism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Neo-Baroque",
+        "prompt": "Neo-Baroque, {prompt}, ornate and elaborate, dynaimc, Neo-Baroque",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Neo-Byzantine",
+        "prompt": "Neo-Byzantine, {prompt}, grand decorative religious style, Orthodox Christian inspired, Neo-Byzantine",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Neo-Futurism",
+        "prompt": "Neo-Futurism, {prompt}, high-tech, curves, spirals, flowing lines, idealistic future, Neo-Futurism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Neo-Impressionism",
+        "prompt": "Neo-Impressionism, {prompt}, tiny dabs of color, Pointillism, painterly, Neo-Impressionism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photograph, realistic"
+    },
+    {
+        "name": "Neo-Rococo",
+        "prompt": "Neo-Rococo, {prompt}, curved forms, naturalistic ornamentation, elaborate, decorative, gaudy, Neo-Rococo",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Neoclassicism",
+        "prompt": "Neoclassicism, {prompt}, ancient Rome and Greece inspired, idealic, sober colors, Neoclassicism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Op Art",
+        "prompt": "Op Art, {prompt}, optical illusion, abstract, geometric pattern, impression of movement, Op Art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Ornate and Intricate",
+        "prompt": "Ornate and Intricate, {prompt}, decorative, highly detailed, elaborate, ornate, intricate",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Pencil Sketch Drawing",
+        "prompt": "Pencil Sketch Drawing, {prompt}, black and white drawing, graphite drawing",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Pop Art 2",
+        "prompt": "Pop Art, {prompt}, vivid colors, flat color, 2D, strong lines, Pop Art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photo, realistic"
+    },
+    {
+        "name": "Rococo",
+        "prompt": "Rococo, {prompt}, flamboyant, pastel colors, curved lines, elaborate detail, Rococo",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Silhouette Art",
+        "prompt": "Silhouette Art, {prompt}, high contrast, well defined, Silhouette Art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Simple Vector Art",
+        "prompt": "Simple Vector Art, {prompt}, 2D flat, simple shapes, minimalistic, professional graphic, flat color, high contrast, Simple Vector Art",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, 3D, photo, realistic"
+    },
+    {
+        "name": "Sketchup",
+        "prompt": "Sketchup, {prompt}, CAD, professional design, Sketchup",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photo, photograph"
+    },
+    {
+        "name": "Steampunk 2",
+        "prompt": "Steampunk, {prompt}, retrofuturistic science fantasy, steam-powered tech, vintage industry, gears, neo-victorian, steampunk",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Surrealism",
+        "prompt": "Surrealism, {prompt}, expressive, dramatic, organic lines and forms, dreamlike and mysterious, Surrealism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realistic"
+    },
+    {
+        "name": "Suprematism",
+        "prompt": "Suprematism, {prompt}, abstract, limited color palette, geometric forms, Suprematism",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, realistic"
+    },
+    {
+        "name": "Terragen",
+        "prompt": "Terragen, {prompt}, beautiful massive landscape, epic scenery, Terragen",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Tranquil Relaxing Atmosphere",
+        "prompt": "Tranquil Relaxing Atmosphere, {prompt}, calming style, soothing colors, peaceful, idealic, Tranquil Relaxing Atmosphere",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, oversaturated"
+    },
+    {
+        "name": "Sticker Designs",
+        "prompt": "Vector Art Stickers, {prompt}, professional vector design, sticker designs, Sticker Sheet",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Vibrant Rim Light",
+        "prompt": "Vibrant Rim Light, {prompt}, bright rim light, high contrast, bold edge light",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Volumetric Lighting",
+        "prompt": "Volumetric Lighting, {prompt}, light depth, dramatic atmospheric lighting, Volumetric Lighting",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
+    },
+    {
+        "name": "Watercolor 2",
+        "prompt": "Watercolor style painting, {prompt}, visible paper texture, colorwash, watercolor",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, photo, realistic"
+    },
+    {
+        "name": "Whimsical and Playful",
+        "prompt": "Whimsical and Playful, {prompt}, imaginative, fantastical, bight colors, stylized, happy, Whimsical and Playful",
+        "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, drab, boring, moody"
     }
 ]
 
-styles = {k['name']: (k['prompt'], k['negative_prompt']) for k in styles}
-default_style = styles['None']
-style_keys = list(styles.keys())
 
+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('(s', '(S')
+    return k
+
+
+default_styles = [normalize_key(x) for x in default_styles]
+styles = {normalize_key(k['name']): (k['prompt'], k['negative_prompt']) for k in styles}
+style_keys = list(styles.keys())
 
 SD_XL_BASE_RATIOS = {
     "0.5": (704, 1408),
@@ -562,10 +969,9 @@ SD_XL_BASE_RATIOS = {
     "3.0": (1728, 576),
 }
 
-
 aspect_ratios = {str(v[0]) + '×' + str(v[1]): v for k, v in SD_XL_BASE_RATIOS.items()}
 
 
-def apply_style(style, positive, negative):
-    p, n = styles.get(style, default_style)
-    return p.replace('{prompt}', positive), n + ', ' + negative
+def apply_style(style, positive):
+    p, n = styles[style]
+    return p.replace('{prompt}', positive), n

modules/upscaler.py

@@ -0,0 +1,25 @@
+import os
+import torch
+
+from comfy_extras.chainner_models.architecture.RRDB import RRDBNet as ESRGAN
+from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel
+from collections import OrderedDict
+from modules.path import upscale_models_path
+
+model_filename = os.path.join(upscale_models_path, 'fooocus_upscaler_s409985e5.bin')
+opImageUpscaleWithModel = ImageUpscaleWithModel()
+model = None
+
+
+def perform_upscale(img):
+    global model
+    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()
+    return opImageUpscaleWithModel.upscale(model, img)[0]

modules/util.py

@@ -1,7 +1,131 @@
+import numpy as np
 import datetime
 import random
 import os
 
+from PIL import Image
+
+
+def image_is_generated_in_current_ui(image, ui_width, ui_height):
+    H, W, C = image.shape
+    
+    if H < ui_height:
+        return False
+
+    if W < ui_width:
+        return False
+
+    # k1 = float(H) / float(W)
+    # k2 = float(ui_height) / float(ui_width)
+    # d = abs(k1 - k2)
+    #
+    # if d > 0.01:
+    #     return False
+
+    return True
+
+
+LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)
+
+
+def resample_image(im, width, height):
+    im = Image.fromarray(im)
+    im = im.resize((width, 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 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()

modules/virtual_memory.py

@@ -0,0 +1,185 @@
+import torch
+import time
+import gc
+
+from safetensors import safe_open
+from comfy import model_management
+from comfy.diffusers_convert import textenc_conversion_lst
+
+
+ALWAYS_USE_VM = None
+
+if isinstance(ALWAYS_USE_VM, bool):
+    print(f'[Virtual Memory System] Forced = {ALWAYS_USE_VM}')
+
+if 'cpu' in model_management.unet_offload_device().type.lower():
+    logic_memory = model_management.total_ram
+    global_virtual_memory_activated = ALWAYS_USE_VM if isinstance(ALWAYS_USE_VM, bool) else logic_memory < 30000
+    print(f'[Virtual Memory System] Logic target is CPU, memory = {logic_memory}')
+else:
+    logic_memory = model_management.total_vram
+    global_virtual_memory_activated = ALWAYS_USE_VM if isinstance(ALWAYS_USE_VM, bool) else logic_memory < 22000
+    print(f'[Virtual Memory System] Logic target is GPU, memory = {logic_memory}')
+
+
+print(f'[Virtual Memory System] Activated = {global_virtual_memory_activated}')
+
+
+@torch.no_grad()
+def recursive_set(obj, key, value):
+    if obj is None:
+        return
+    if '.' in key:
+        k1, k2 = key.split('.', 1)
+        recursive_set(getattr(obj, k1, None), k2, value)
+    else:
+        setattr(obj, key, value)
+
+
+@torch.no_grad()
+def recursive_del(obj, key):
+    if obj is None:
+        return
+    if '.' in key:
+        k1, k2 = key.split('.', 1)
+        recursive_del(getattr(obj, k1, None), k2)
+    else:
+        delattr(obj, key)
+
+
+@torch.no_grad()
+def force_load_state_dict(model, state_dict):
+    for k in list(state_dict.keys()):
+        p = torch.nn.Parameter(state_dict[k], requires_grad=False)
+        recursive_set(model, k, p)
+        del state_dict[k]
+    return
+
+
+@torch.no_grad()
+def only_load_safetensors_keys(filename):
+    try:
+        with safe_open(filename, framework="pt", device='cpu') as f:
+            result = list(f.keys())
+        assert len(result) > 0
+        return result
+    except:
+        return None
+
+
+@torch.no_grad()
+def move_to_virtual_memory(model, comfy_unload=True):
+    timer = time.time()
+
+    if comfy_unload:
+        model_management.unload_model()
+
+    virtual_memory_dict = getattr(model, 'virtual_memory_dict', None)
+    if isinstance(virtual_memory_dict, dict):
+        # Already in virtual memory.
+        return
+
+    model_file = getattr(model, 'model_file', None)
+    assert isinstance(model_file, dict)
+
+    filename = model_file['filename']
+    prefix = model_file['prefix']
+
+    if hasattr(model, 'virtual_memory_safetensors_keys'):
+        safetensors_keys = model.virtual_memory_safetensors_keys
+    else:
+        safetensors_keys = only_load_safetensors_keys(filename)
+        model.virtual_memory_safetensors_keys = safetensors_keys
+
+    if safetensors_keys is None:
+        print(f'[Virtual Memory System] Error: The Virtual Memory System currently only support safetensors models!')
+        return
+
+    sd = model.state_dict()
+    original_device = list(sd.values())[0].device.type
+    model_file['original_device'] = original_device
+
+    virtual_memory_dict = {}
+
+    for k, v in sd.items():
+        current_key = k
+        current_flag = None
+        if prefix == 'refiner_clip':
+            current_key_in_safetensors = k
+
+            for a, b in textenc_conversion_lst:
+                current_key_in_safetensors = current_key_in_safetensors.replace(b, a)
+
+            current_key_in_safetensors = current_key_in_safetensors.replace('clip_g.transformer.text_model.encoder.layers.', 'conditioner.embedders.0.model.transformer.resblocks.')
+            current_key_in_safetensors = current_key_in_safetensors.replace('clip_g.text_projection', 'conditioner.embedders.0.model.text_projection')
+            current_key_in_safetensors = current_key_in_safetensors.replace('clip_g.logit_scale', 'conditioner.embedders.0.model.logit_scale')
+            current_key_in_safetensors = current_key_in_safetensors.replace('clip_g.', 'conditioner.embedders.0.model.')
+
+            for e in ["weight", "bias"]:
+                for i, k in enumerate(['q', 'k', 'v']):
+                    e_flag = f'.{k}_proj.{e}'
+                    if current_key_in_safetensors.endswith(e_flag):
+                        current_key_in_safetensors = current_key_in_safetensors[:-len(e_flag)] + f'.in_proj_{e}'
+                        current_flag = (1280 * i, 1280 * (i + 1))
+        else:
+            current_key_in_safetensors = prefix + '.' + k
+        current_device = torch.device(index=v.device.index, type=v.device.type)
+        if current_key_in_safetensors in safetensors_keys:
+            virtual_memory_dict[current_key] = (current_key_in_safetensors, current_device, current_flag)
+            recursive_del(model, current_key)
+        else:
+            # print(f'[Virtual Memory System] Missed key: {current_key}')
+            pass
+
+    del sd
+    gc.collect()
+    model_management.soft_empty_cache()
+
+    model.virtual_memory_dict = virtual_memory_dict
+
+    print(f'[Virtual Memory System] time = {str("%.5f" % (time.time() - timer))}s: {prefix} released from {original_device}: {filename}')
+    return
+
+
+@torch.no_grad()
+def load_from_virtual_memory(model):
+    timer = time.time()
+
+    virtual_memory_dict = getattr(model, 'virtual_memory_dict', None)
+    if not isinstance(virtual_memory_dict, dict):
+        # Not in virtual memory.
+        return
+
+    model_file = getattr(model, 'model_file', None)
+    assert isinstance(model_file, dict)
+
+    filename = model_file['filename']
+    prefix = model_file['prefix']
+    original_device = model_file['original_device']
+
+    with safe_open(filename, framework="pt", device=original_device) as f:
+        for current_key, (current_key_in_safetensors, current_device, current_flag) in virtual_memory_dict.items():
+            tensor = f.get_tensor(current_key_in_safetensors).to(current_device)
+            if isinstance(current_flag, tuple) and len(current_flag) == 2:
+                a, b = current_flag
+                tensor = tensor[a:b]
+            parameter = torch.nn.Parameter(tensor, requires_grad=False)
+            recursive_set(model, current_key, parameter)
+
+    print(f'[Virtual Memory System] time = {str("%.5f" % (time.time() - timer))}s: {prefix} loaded to {original_device}: {filename}')
+    del model.virtual_memory_dict
+    return
+
+
+@torch.no_grad()
+def try_move_to_virtual_memory(model, comfy_unload=True):
+    if not global_virtual_memory_activated:
+        return
+
+    import modules.default_pipeline as pipeline
+
+    if pipeline.xl_refiner is None:
+        # If users do not use refiner, no need to use this.
+        return
+
+    move_to_virtual_memory(model, comfy_unload)

readme.md

@@ -1,8 +1,8 @@
 # Fooocus
 
-<img src="https://github.com/lllyasviel/Fooocus/assets/19834515/bcb0336b-5c79-4de2-b0cb-f7f68c753a88" width=100%>
+<img src="https://github.com/lllyasviel/Fooocus/assets/19834515/f79c5981-cf80-4ee3-b06b-3fef3f8bfbc7" width=100%>
 
-Fooocus is an image generating software.
+Fooocus is an image generating software (based on [Gradio](https://www.gradio.app/)).
 
 Fooocus is a rethinking of Stable Diffusion and Midjourney’s designs:
 
@@ -24,7 +24,7 @@ Fooocus also developed many "fooocus-only" features for advanced users to get pe
 
 You can directly download Fooocus with:
 
-**[>>> Click here to download <<<](https://github.com/lllyasviel/Fooocus/releases/download/release/Fooocus_win64_1-1-10.7z)**
+**[>>> Click here to download <<<](https://github.com/lllyasviel/Fooocus/releases/download/release/Fooocus_win64_2-0-50.7z)**
 
 After you download the file, please uncompress it, and then run the "run.bat".
 
@@ -34,55 +34,102 @@ In the first time you launch the software, it will automatically download models
 
 1. It will download [sd_xl_base_1.0_0.9vae.safetensors from here](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0_0.9vae.safetensors) as the file "Fooocus\models\checkpoints\sd_xl_base_1.0_0.9vae.safetensors".
 2. It will download [sd_xl_refiner_1.0_0.9vae.safetensors from here](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0_0.9vae.safetensors) as the file "Fooocus\models\checkpoints\sd_xl_refiner_1.0_0.9vae.safetensors".
+3. Note that if you use inpaint, at the first time you inpaint an image, it will download [Fooocus's own inpaint control model from here](https://huggingface.co/lllyasviel/fooocus_inpaint/resolve/main/inpaint.fooocus.patch) as the file "Fooocus\models\inpaint\inpaint.fooocus.patch" (the size of this file is 1.28GB).
 
 ![image](https://github.com/lllyasviel/Fooocus/assets/19834515/d386f817-4bd7-490c-ad89-c1e228c23447)
 
 If you already have these files, you can copy them to the above locations to speed up installation.
 
+Note that if you see **"MetadataIncompleteBuffer"**, then your model files are corrupted. Please download models again.
+
 Below is a test on a relatively low-end laptop with **16GB System RAM** and **6GB VRAM** (Nvidia 3060 laptop). The speed on this machine is about 1.35 seconds per iteration. Pretty impressive – nowadays laptops with 3060 are usually at very acceptable price.
 
 ![image](https://github.com/lllyasviel/Fooocus/assets/19834515/938737a5-b105-4f19-b051-81356cb7c495)
 
-Note that the minimal requirement is **4GB Nvidia GPU memory (4GB VRAM)** and **8GB system memory (8GB RAM)**. This requires using Microsoft’s Virtual Swap technique, which is automatically enabled by your Windows installation in most cases, so you often do not need to do anything about it. However, if you are not sure, or if you manually turned it off (would anyone really do that?), you can enable it here:
+Besides, recently many other software report that Nvidia driver above 532 is sometimes 10x slower than Nvidia driver 531. If your generation time is very long, consider download [Nvidia Driver 531 Laptop](https://www.nvidia.com/download/driverResults.aspx/199991/en-us/) or [Nvidia Driver 531 Desktop](https://www.nvidia.com/download/driverResults.aspx/199990/en-us/).
+
+Note that the minimal requirement is **4GB Nvidia GPU memory (4GB VRAM)** and **8GB system memory (8GB RAM)**. This requires using Microsoft’s Virtual Swap technique, which is automatically enabled by your Windows installation in most cases, so you often do not need to do anything about it. However, if you are not sure, or if you manually turned it off (would anyone really do that?), or **if you see any "RuntimeError: CPUAllocator"**, you can enable it here:
 
 <details>
 <summary>Click here to the see the image instruction. </summary>
 
 ![image](https://github.com/lllyasviel/Fooocus/assets/19834515/2a06b130-fe9b-4504-94f1-2763be4476e9)
 
+**And make sure that you have at least 40GB free space on each drive if you still see "RuntimeError: CPUAllocator" !**
+
 </details>
 
 Please open an issue if you use similar devices but still cannot achieve acceptable performances.
 
 ### Colab
 
-(Last tested - 2023 Aug 14)
+(Last tested - 2023 Sep 13)
 
 | Colab | Info
 | --- | --- |
 [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lllyasviel/Fooocus/blob/main/colab.ipynb) | Fooocus Colab (Official Version)
 
-Note that sometimes this Colab will say like "you must restart the runtime in order to use newly installed XX". This can be safely ignored.
+Thanks to [camenduru](https://github.com/camenduru)!
 
-Thanks to [camenduru](https://github.com/camenduru)'s codes!
+### Linux (Using Anaconda)
 
-### Linux
-
-The command lines are
+If you want to use Anaconda/Miniconda, you can
 
     git clone https://github.com/lllyasviel/Fooocus.git
     cd Fooocus
     conda env create -f environment.yaml
     conda activate fooocus
-    pip install -r requirements_versions.txt
+    pip install pygit2==1.12.2
 
 Then download the models: download [sd_xl_base_1.0_0.9vae.safetensors from here](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0_0.9vae.safetensors) as the file "Fooocus\models\checkpoints\sd_xl_base_1.0_0.9vae.safetensors", and download [sd_xl_refiner_1.0_0.9vae.safetensors from here](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/resolve/main/sd_xl_refiner_1.0_0.9vae.safetensors) as the file "Fooocus\models\checkpoints\sd_xl_refiner_1.0_0.9vae.safetensors". **Or let Fooocus automatically download the models** using the launcher:
 
-    python launch.py
+    conda activate fooocus
+    python entry_with_update.py
 
 Or if you want to open a remote port, use
 
-    python launch.py --listen
+    conda activate fooocus
+    python entry_with_update.py --listen
+
+### Linux (Using Python Venv)
+
+Your Linux needs to have **Python 3.10** installed, and lets say your Python can be called with command **python3** with your venv system working, you can
+
+    git clone https://github.com/lllyasviel/Fooocus.git
+    cd Fooocus
+    python3 -m venv fooocus_env
+    source fooocus_env/bin/activate
+    pip install pygit2==1.12.2
+
+See the above sections for model downloads. You can launch the software with:
+
+    source fooocus_env/bin/activate
+    python entry_with_update.py
+
+Or if you want to open a remote port, use
+
+    source fooocus_env/bin/activate
+    python entry_with_update.py --listen
+
+### Linux (Using native system Python)
+
+If you know what you are doing, and your Linux already has **Python 3.10** installed, and your Python can be called with command **python3** (and Pip with **pip3**), you can
+
+    git clone https://github.com/lllyasviel/Fooocus.git
+    cd Fooocus
+    pip3 install pygit2==1.12.2
+
+See the above sections for model downloads. You can launch the software with:
+
+    python3 entry_with_update.py
+
+Or if you want to open a remote port, use
+
+    python3 entry_with_update.py --listen
+
+### Linux (AMD GPUs)
+
+Installation is the same as Linux part. It has been tested for 6700XT. Works for both Pytorch 1.13 and Pytorch 2. 
 
 ### Mac/Windows(AMD GPUs)
 
@@ -93,17 +140,19 @@ Coming soon ...
 
 Below things are already inside the software, and **users do not need to do anything about these**.
 
-Note that some of these tricks are currently (2023 Aug 11) impossible to reproduce in Automatic1111's interface or ComfyUI's node system. You may expect better results from Fooocus than other software even when they use similar models/pipelines.
-
-1. Native refiner swap inside one single k-sampler. The advantage is that now the refiner model can reuse the base model's momentum (or ODE's history parameters) collected from k-sampling to achieve more coherent sampling. In Automatic1111's high-res fix and ComfyUI's node system, the base model and refiner use two independent k-samplers, which means the momentum is largely wasted, and the sampling continuity is broken. Fooocus uses its own advanced k-diffusion sampling that ensures seamless, native, and continuous swap in a refiner setup. (Update Aug 13: Actually I discussed this with Automatic1111 several days ago and it seems that the “native refiner swap inside one single k-sampler” is [merged]( https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12371) into the dev branch of webui. Great!)
-2. Negative ADM guidance. Because the highest resolution level of XL Base does not have cross attentions, the positive and negative signals for XL's highest resolution level cannot receive enough contrasts during the CFG sampling, causing the results look a bit plastic or overly smooth in certain cases. Fortunately, since the XL's highest resolution level is still conditioned on image aspect ratios (ADM), we can modify the adm on the positive/negative side to compensate for the lack of CFG contrast in the highest resolution level. (Update Aug 16, the IOS App [Drawing Things](https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820) will support Negative ADM Guidance. Great!)
-3. We implemented a carefully tuned variation of the Section 5.1 of ["Improving Sample Quality of Diffusion Models Using Self-Attention Guidance"](https://arxiv.org/pdf/2210.00939.pdf). The weight is set to very low, but this is Fooocus's final guarantee to make sure that the XL will never yield overly smooth or plastic appearance. This can almostly eliminate all cases that XL still occasionally produce overly smooth results even with negative ADM guidance.
-4. We modified the style templates a bit and added the "cinematic-default".
-5. We tested the "sd_xl_offset_example-lora_1.0.safetensors" and it seems that when the lora weight is below 0.5, the results are always better than XL without lora.
-6. The parameters of samplers are carefully tuned.
-7. Because XL uses positional encoding for generation resolution, images generated by several fixed resolutions look a bit better than that from arbitrary resolutions (because the positional encoding is not very good at handling int numbers that are unseen during training). This suggests that the resolutions in UI may be hard coded for best results.
-8. Separated prompts for two different text encoders seem unnecessary. Separated prompts for base model and refiner may work but the effects are random, and we refrain from implement this.
-9. DPM family seems well-suited for XL, since XL sometimes generates overly smooth texture but DPM family sometimes generate overly dense detail in texture. Their joint effect looks neutral and appealing to human perception.
+1. GPT2-based [prompt expansion as a dynamic style "Fooocus V2".](https://github.com/lllyasviel/Fooocus/discussions/117#raw) (similar to Midjourney's hidden pre-processsing and "raw" mode, or the LeonardoAI's Prompt Magic).
+2. Native refiner swap inside one single k-sampler. The advantage is that now the refiner model can reuse the base model's momentum (or ODE's history parameters) collected from k-sampling to achieve more coherent sampling. In Automatic1111's high-res fix and ComfyUI's node system, the base model and refiner use two independent k-samplers, which means the momentum is largely wasted, and the sampling continuity is broken. Fooocus uses its own advanced k-diffusion sampling that ensures seamless, native, and continuous swap in a refiner setup. (Update Aug 13: Actually I discussed this with Automatic1111 several days ago and it seems that the “native refiner swap inside one single k-sampler” is [merged]( https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12371) into the dev branch of webui. Great!)
+3. Negative ADM guidance. Because the highest resolution level of XL Base does not have cross attentions, the positive and negative signals for XL's highest resolution level cannot receive enough contrasts during the CFG sampling, causing the results look a bit plastic or overly smooth in certain cases. Fortunately, since the XL's highest resolution level is still conditioned on image aspect ratios (ADM), we can modify the adm on the positive/negative side to compensate for the lack of CFG contrast in the highest resolution level. (Update Aug 16, the IOS App [Drawing Things](https://apps.apple.com/us/app/draw-things-ai-generation/id6444050820) will support Negative ADM Guidance. Great!)
+4. We implemented a carefully tuned variation of the Section 5.1 of ["Improving Sample Quality of Diffusion Models Using Self-Attention Guidance"](https://arxiv.org/pdf/2210.00939.pdf). The weight is set to very low, but this is Fooocus's final guarantee to make sure that the XL will never yield overly smooth or plastic appearance (examples [here](https://github.com/lllyasviel/Fooocus/discussions/117#sharpness)). This can almostly eliminate all cases that XL still occasionally produce overly smooth results even with negative ADM guidance. (Update 2023 Aug 18, the Gaussian kernel of SAG is changed to an anisotropic kernel for better structure preservation and fewer artifacts.)
+5. We modified the style templates a bit and added the "cinematic-default".
+6. We tested the "sd_xl_offset_example-lora_1.0.safetensors" and it seems that when the lora weight is below 0.5, the results are always better than XL without lora.
+7. The parameters of samplers are carefully tuned.
+8. Because XL uses positional encoding for generation resolution, images generated by several fixed resolutions look a bit better than that from arbitrary resolutions (because the positional encoding is not very good at handling int numbers that are unseen during training). This suggests that the resolutions in UI may be hard coded for best results.
+9. Separated prompts for two different text encoders seem unnecessary. Separated prompts for base model and refiner may work but the effects are random, and we refrain from implement this.
+10. DPM family seems well-suited for XL, since XL sometimes generates overly smooth texture but DPM family sometimes generate overly dense detail in texture. Their joint effect looks neutral and appealing to human perception.
+11. A carefully designed system for balancing multiple styles as well as prompt expansion.
+12. Using automatic1111's method to normalize prompt emphasizing. This significantly improve results when users directly copy prompts from civitai.
+13. The joint swap system of refiner now also support img2img and upscale in a seamless way.
 
 ## Advanced Features
 

shared.py

@@ -1,2 +1 @@
 gradio_root = None
-

update_log.md

@@ -1,3 +1,92 @@
+(2023 sep 21) The feature updating of Fooocus will be paused for about two or three weeks because we have some events and travelling - we will come back in early or mid October. 
+
+### 2.0.72
+
+* Allow users to choose path of models.
+
+### 2.0.65
+
+* Inpaint model released.
+
+### 2.0.50
+
+* Variation/Upscale (Midjourney Toolbar) implemented.
+
+### 2.0.16
+
+* Virtual memory system implemented. Now Colab can run both base model and refiner model with 7.8GB RAM + 5.3GB VRAM, and it never crashes.
+* If you are lucky enough to read this line, keep in mind that ComfyUI cannot do this. This is very reasonable that Fooocus is more optimized because it only need to handle a fixed pipeline, but ComfyUI need to consider arbitrary pipelines. 
+* But if we just consider the optimization of this fixed workload, after 2.0.16, Fooocus has become the most optimized SDXL app, outperforming ComfyUI.
+
+### 2.0.0
+
+* V2 released.
+* completely rewrite text processing pipeline (higher image quality and prompt understanding).
+* support multi-style.
+* In 100 tests (prompts written by ChatGPT), V2 default results outperform V1 default results in 87 cases, evaluated by two human.
+* In 100 tests (prompts written by ChatGPT), V2 prompt understanding outperform V1 prompt understanding in 81 cases, evaluated by two human, in both default setting and multi/single style mode.
+* Because the above number is above 80%, we view this as a major update and directly jump to 2.0.0.
+* Some other things are renamed.
+
+### 1.0.67
+
+* Use dynamic weighting and lower weights for prompt expansion.
+
+### 1.0.64
+
+* Fixed a small OOM problem.
+
+### 1.0.62
+
+* Change prompt expansion to suffix mode for better balance of semantic and style (and debugging).
+
+### 1.0.60
+
+* Tune the balance between style and Prompt Expansion.
+
+### 1.0.56
+
+* Begin to use magic split.
+
+### 1.0.55
+
+* Minor changes of Prompt Expansion.
+
+### 1.0.52
+
+* Reduce the semantic corruption of Prompt Expansion.
+
+### 1.0.51
+
+* Speed up Prompt Expansion a bit.
+
+### 1.0.50
+
+* Prompt expansion and a "Raw mode" to turn it off (similar to Midjourney's "raw").
+
+### 1.0.45
+
+* Reworked SAG, removed unnecessary patch
+* Reworked anisotropic filters for faster compute.
+* Replaced with guided anisotropic filter for less distortion.
+
+### 1.0.41
+
+(The update of Fooocus will be paused for a period of time for AUTOMATIC1111 sd-webui 1.6.X, and some features will also be implemented as webui extensions)
+
+### 1.0.40
+
+* Behaviors reverted to 1.0.36 again (refiner steps). The 1.0.36 is too perfect and too typical; beating 1.0.36 is just impossible.
+
+### 1.0.39
+
+* Reverted unstable changes between 1.0.37 and 1.0.38 .
+* Increased refiner steps to half of sampling steps.
+
+### 1.0.36
+
+* Change gaussian kernel to anisotropic kernel.
+
 ### 1.0.34
 
 * Random seed restoring.

webui.py

@@ -7,13 +7,17 @@ import modules.path
 import fooocus_version
 import modules.html
 import modules.async_worker as worker
+import modules.flags as flags
+import modules.gradio_hijack as grh
+import comfy.model_management as model_management
 
-from modules.sdxl_styles import style_keys, aspect_ratios
+from modules.sdxl_styles import style_keys, aspect_ratios, fooocus_expansion, default_styles
 
 
 def generate_clicked(*args):
-    yield gr.update(interactive=False), \
-        gr.update(visible=True, value=modules.html.make_progress_html(1, 'Processing text encoding ...')), \
+    execution_start_time = time.perf_counter()
+
+    yield gr.update(visible=True, value=modules.html.make_progress_html(1, 'Initializing ...')), \
         gr.update(visible=True, value=None), \
         gr.update(visible=False)
 
@@ -26,16 +30,17 @@ def generate_clicked(*args):
             flag, product = worker.outputs.pop(0)
             if flag == 'preview':
                 percentage, title, image = product
-                yield gr.update(interactive=False), \
-                    gr.update(visible=True, value=modules.html.make_progress_html(percentage, title)), \
+                yield gr.update(visible=True, value=modules.html.make_progress_html(percentage, title)), \
                     gr.update(visible=True, value=image) if image is not None else gr.update(), \
                     gr.update(visible=False)
             if flag == 'results':
-                yield gr.update(interactive=True), \
-                    gr.update(visible=False), \
+                yield gr.update(visible=False), \
                     gr.update(visible=False), \
                     gr.update(visible=True, value=product)
                 finished = True
+
+    execution_time = time.perf_counter() - execution_start_time
+    print(f'Total time: {execution_time:.2f} seconds')
     return
 
 
@@ -43,23 +48,76 @@ shared.gradio_root = gr.Blocks(title='Fooocus ' + fooocus_version.version, css=m
 with shared.gradio_root:
     with gr.Row():
         with gr.Column():
-            progress_window = gr.Image(label='Preview', show_label=True, height=640, visible=False)
+            progress_window = grh.Image(label='Preview', show_label=True, height=640, visible=False)
             progress_html = gr.HTML(value=modules.html.make_progress_html(32, 'Progress 32%'), visible=False, elem_id='progress-bar', elem_classes='progress-bar')
             gallery = gr.Gallery(label='Gallery', show_label=False, object_fit='contain', height=720, visible=True)
             with gr.Row(elem_classes='type_row'):
                 with gr.Column(scale=0.85):
                     prompt = gr.Textbox(show_label=False, placeholder="Type prompt here.", container=False, autofocus=True, elem_classes='type_row', lines=1024)
                 with gr.Column(scale=0.15, min_width=0):
-                    run_button = gr.Button(label="Generate", value="Generate", elem_classes='type_row')
+                    run_button = gr.Button(label="Generate", value="Generate", elem_classes='type_row', visible=True)
+                    stop_button = gr.Button(label="Stop", value="Stop", elem_classes='type_row', visible=False)
+
+                    def stop_clicked():
+                        model_management.interrupt_current_processing()
+                        return gr.update(interactive=False)
+
+                    stop_button.click(stop_clicked, outputs=stop_button, queue=False)
+            with gr.Row(elem_classes='advanced_check_row'):
+                input_image_checkbox = gr.Checkbox(label='Input Image', value=False, container=False, elem_classes='min_check')
+                advanced_checkbox = gr.Checkbox(label='Advanced', value=False, container=False, elem_classes='min_check')
+            with gr.Row(visible=False) as image_input_panel:
+                with gr.Tabs():
+                    with gr.TabItem(label='Upscale or Variation') as uov_tab:
                         with gr.Row():
-                advanced_checkbox = gr.Checkbox(label='Advanced', value=False, container=False)
+                            with gr.Column():
+                                uov_input_image = grh.Image(label='Drag above image to here', source='upload', type='numpy')
+                            with gr.Column():
+                                uov_method = gr.Radio(label='Upscale or Variation:', choices=flags.uov_list, value=flags.disabled)
+                                gr.HTML('<a href="https://github.com/lllyasviel/Fooocus/discussions/390">\U0001F4D4 Document</a>')
+                    with gr.TabItem(label='Inpaint or Outpaint (beta)') as inpaint_tab:
+                        inpaint_input_image = grh.Image(label='Drag above image to here', source='upload', type='numpy', tool='sketch', height=500, brush_color="#FFFFFF")
+                        gr.HTML('Outpaint Expansion (<a href="https://github.com/lllyasviel/Fooocus/discussions/414">\U0001F4D4 Document</a>):')
+                        outpaint_selections = gr.CheckboxGroup(choices=['Left', 'Right', 'Top', 'Bottom'], value=[], label='Outpaint', show_label=False, container=False)
+                        gr.HTML('* \"Inpaint or Outpaint\" is powered by the sampler \"DPMPP Fooocus Seamless 2M SDE Karras Inpaint Sampler\" (beta)')
+
+            input_image_checkbox.change(lambda x: gr.update(visible=x), inputs=input_image_checkbox, outputs=image_input_panel, queue=False,
+                                        _js="(x) => {if(x){setTimeout(() => window.scrollTo({ top: window.scrollY + 500, behavior: 'smooth' }), 50);}else{setTimeout(() => window.scrollTo({ top: 0, behavior: 'smooth' }), 50);} return x}")
+
+            current_tab = gr.Textbox(value='uov', visible=False)
+
+            default_image = None
+
+            def update_default_image(x):
+                global default_image
+                if isinstance(x, dict):
+                    default_image = x['image']
+                else:
+                    default_image = x
+                return
+
+            def clear_default_image():
+                global default_image
+                default_image = None
+                return
+
+            uov_input_image.upload(update_default_image, inputs=uov_input_image, queue=False)
+            inpaint_input_image.upload(update_default_image, inputs=inpaint_input_image, queue=False)
+
+            uov_input_image.clear(clear_default_image, queue=False)
+            inpaint_input_image.clear(clear_default_image, queue=False)
+
+            uov_tab.select(lambda: ['uov', default_image], outputs=[current_tab, uov_input_image], queue=False)
+            inpaint_tab.select(lambda: ['inpaint', default_image], outputs=[current_tab, inpaint_input_image], queue=False)
+
         with gr.Column(scale=0.5, visible=False) as right_col:
             with gr.Tab(label='Setting'):
                 performance_selction = gr.Radio(label='Performance', choices=['Speed', 'Quality'], value='Speed')
-                aspect_ratios_selction = gr.Radio(label='Aspect Ratios (width × height)', choices=list(aspect_ratios.keys()),
-                                                  value='1152×896')
+                aspect_ratios_selction = gr.Radio(label='Aspect Ratios', choices=list(aspect_ratios.keys()),
+                                                  value='1152×896', info='width × height')
                 image_number = gr.Slider(label='Image Number', minimum=1, maximum=32, step=1, value=2)
-                negative_prompt = gr.Textbox(label='Negative Prompt', show_label=True, placeholder="Type prompt here.")
+                negative_prompt = gr.Textbox(label='Negative Prompt', show_label=True, placeholder="Type prompt here.",
+                                             info='Describing objects that you do not want to see.')
                 seed_random = gr.Checkbox(label='Random', value=True)
                 image_seed = gr.Number(label='Seed', value=0, precision=0, visible=False)
 
@@ -72,11 +130,13 @@ with shared.gradio_root:
                     else:
                         return s
 
-                seed_random.change(random_checked, inputs=[seed_random], outputs=[image_seed])
+                seed_random.change(random_checked, inputs=[seed_random], outputs=[image_seed], queue=False)
 
             with gr.Tab(label='Style'):
-                style_selction = gr.Radio(show_label=False, container=True,
-                                          choices=style_keys, value='cinematic-default')
+                style_selections = gr.CheckboxGroup(show_label=False, container=False,
+                                                    choices=[fooocus_expansion] + style_keys,
+                                                    value=[fooocus_expansion] + default_styles,
+                                                    label='Image Style')
             with gr.Tab(label='Advanced'):
                 with gr.Row():
                     base_model = gr.Dropdown(label='SDXL Base Model', choices=modules.path.model_filenames, value=modules.path.default_base_model_name, show_label=True)
@@ -91,7 +151,7 @@ with shared.gradio_root:
                 with gr.Row():
                     model_refresh = gr.Button(label='Refresh', value='\U0001f504 Refresh All Files', variant='secondary', elem_classes='refresh_button')
                 with gr.Accordion(label='Advanced', open=False):
-                    sharpness = gr.Slider(label='Sampling Sharpness', minimum=0.0, maximum=40.0, step=0.01, value=2.0)
+                    sharpness = gr.Slider(label='Sampling Sharpness', minimum=0.0, maximum=30.0, step=0.01, value=2.0)
                     gr.HTML('<a href="https://github.com/lllyasviel/Fooocus/discussions/117">\U0001F4D4 Document</a>')
 
                 def model_refresh_clicked():
@@ -102,16 +162,22 @@ with shared.gradio_root:
                         results += [gr.update(choices=['None'] + modules.path.lora_filenames), gr.update()]
                     return results
 
-                model_refresh.click(model_refresh_clicked, [], [base_model, refiner_model] + lora_ctrls)
+                model_refresh.click(model_refresh_clicked, [], [base_model, refiner_model] + lora_ctrls, queue=False)
 
-        advanced_checkbox.change(lambda x: gr.update(visible=x), advanced_checkbox, right_col)
+        advanced_checkbox.change(lambda x: gr.update(visible=x), advanced_checkbox, right_col, queue=False)
         ctrls = [
-            prompt, negative_prompt, style_selction,
+            prompt, negative_prompt, style_selections,
             performance_selction, aspect_ratios_selction, image_number, image_seed, sharpness
         ]
         ctrls += [base_model, refiner_model] + lora_ctrls
-        run_button.click(fn=refresh_seed, inputs=[seed_random, image_seed], outputs=image_seed)\
-            .then(fn=generate_clicked, inputs=ctrls, outputs=[run_button, progress_html, progress_window, gallery])
+        ctrls += [input_image_checkbox, current_tab]
+        ctrls += [uov_method, uov_input_image]
+        ctrls += [outpaint_selections, inpaint_input_image]
+
+        run_button.click(lambda: (gr.update(visible=True, interactive=True), gr.update(visible=False), []), outputs=[stop_button, run_button, gallery])\
+            .then(fn=refresh_seed, inputs=[seed_random, image_seed], outputs=image_seed)\
+            .then(fn=generate_clicked, inputs=ctrls, outputs=[progress_html, progress_window, gallery])\
+            .then(lambda: (gr.update(visible=True), gr.update(visible=False)), outputs=[run_button, stop_button])
 
 
 parser = argparse.ArgumentParser()