+
+
+'''
+
+
+def make_progress_html(number, text):
+ return progress_html.replace('*number*', str(number)).replace('*text*', text)
diff --git a/modules/inpaint_worker.py b/modules/inpaint_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..88ec39d6df930fc689f591e85877d74e07db50d9
--- /dev/null
+++ b/modules/inpaint_worker.py
@@ -0,0 +1,257 @@
+import torch
+import numpy as np
+
+from PIL import Image, ImageFilter
+from modules.util import resample_image, set_image_shape_ceil, get_image_shape_ceil
+from modules.upscaler import perform_upscale
+
+
+inpaint_head_model = None
+
+
+class InpaintHead(torch.nn.Module):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.head = torch.nn.Parameter(torch.empty(size=(320, 5, 3, 3), device='cpu'))
+
+ def __call__(self, x):
+ x = torch.nn.functional.pad(x, (1, 1, 1, 1), "replicate")
+ return torch.nn.functional.conv2d(input=x, weight=self.head)
+
+
+current_task = None
+
+
+def box_blur(x, k):
+ x = Image.fromarray(x)
+ x = x.filter(ImageFilter.BoxBlur(k))
+ return np.array(x)
+
+
+def max33(x):
+ x = Image.fromarray(x)
+ x = x.filter(ImageFilter.MaxFilter(3))
+ return np.array(x)
+
+
+def morphological_open(x):
+ x_int32 = np.zeros_like(x).astype(np.int32)
+ x_int32[x > 127] = 256
+ for _ in range(32):
+ maxed = max33(x_int32) - 8
+ x_int32 = np.maximum(maxed, x_int32)
+ return x_int32.clip(0, 255).astype(np.uint8)
+
+
+def up255(x, t=0):
+ y = np.zeros_like(x).astype(np.uint8)
+ y[x > t] = 255
+ return y
+
+
+def imsave(x, path):
+ x = Image.fromarray(x)
+ x.save(path)
+
+
+def regulate_abcd(x, a, b, c, d):
+ H, W = x.shape[:2]
+ if a < 0:
+ a = 0
+ if a > H:
+ a = H
+ if b < 0:
+ b = 0
+ if b > H:
+ b = H
+ if c < 0:
+ c = 0
+ if c > W:
+ c = W
+ if d < 0:
+ d = 0
+ if d > W:
+ d = W
+ return int(a), int(b), int(c), int(d)
+
+
+def compute_initial_abcd(x):
+ indices = np.where(x)
+ a = np.min(indices[0])
+ b = np.max(indices[0])
+ c = np.min(indices[1])
+ d = np.max(indices[1])
+ abp = (b + a) // 2
+ abm = (b - a) // 2
+ cdp = (d + c) // 2
+ cdm = (d - c) // 2
+ l = int(max(abm, cdm) * 1.15)
+ a = abp - l
+ b = abp + l + 1
+ c = cdp - l
+ d = cdp + l + 1
+ a, b, c, d = regulate_abcd(x, a, b, c, d)
+ return a, b, c, d
+
+
+def solve_abcd(x, a, b, c, d, k):
+ k = float(k)
+ assert 0.0 <= k <= 1.0
+
+ H, W = x.shape[:2]
+ if k == 1.0:
+ return 0, H, 0, W
+ while True:
+ if b - a >= H * k and d - c >= W * k:
+ break
+
+ add_h = (b - a) < (d - c)
+ add_w = not add_h
+
+ if b - a == H:
+ add_w = True
+
+ if d - c == W:
+ add_h = True
+
+ if add_h:
+ a -= 1
+ b += 1
+
+ if add_w:
+ c -= 1
+ d += 1
+
+ a, b, c, d = regulate_abcd(x, a, b, c, d)
+ return a, b, c, d
+
+
+def fooocus_fill(image, mask):
+ current_image = image.copy()
+ raw_image = image.copy()
+ area = np.where(mask < 127)
+ store = raw_image[area]
+
+ for k, repeats in [(512, 2), (256, 2), (128, 4), (64, 4), (33, 8), (15, 8), (5, 16), (3, 16)]:
+ for _ in range(repeats):
+ current_image = box_blur(current_image, k)
+ current_image[area] = store
+
+ return current_image
+
+
+class InpaintWorker:
+ def __init__(self, image, mask, use_fill=True, k=0.618):
+ a, b, c, d = compute_initial_abcd(mask > 0)
+ a, b, c, d = solve_abcd(mask, a, b, c, d, k=k)
+
+ # interested area
+ self.interested_area = (a, b, c, d)
+ self.interested_mask = mask[a:b, c:d]
+ self.interested_image = image[a:b, c:d]
+
+ # super resolution
+ if get_image_shape_ceil(self.interested_image) < 1024:
+ self.interested_image = perform_upscale(self.interested_image)
+
+ # resize to make images ready for diffusion
+ self.interested_image = set_image_shape_ceil(self.interested_image, 1024)
+ self.interested_fill = self.interested_image.copy()
+ H, W, C = self.interested_image.shape
+
+ # process mask
+ self.interested_mask = up255(resample_image(self.interested_mask, W, H), t=127)
+
+ # compute filling
+ if use_fill:
+ self.interested_fill = fooocus_fill(self.interested_image, self.interested_mask)
+
+ # soft pixels
+ self.mask = morphological_open(mask)
+ self.image = image
+
+ # ending
+ self.latent = None
+ self.latent_after_swap = None
+ self.swapped = False
+ self.latent_mask = None
+ self.inpaint_head_feature = None
+ return
+
+ def load_latent(self, latent_fill, latent_mask, latent_swap=None):
+ self.latent = latent_fill
+ self.latent_mask = latent_mask
+ self.latent_after_swap = latent_swap
+ return
+
+ def patch(self, inpaint_head_model_path, inpaint_latent, inpaint_latent_mask, model):
+ global inpaint_head_model
+
+ if inpaint_head_model is None:
+ inpaint_head_model = InpaintHead()
+ sd = torch.load(inpaint_head_model_path, map_location='cpu')
+ inpaint_head_model.load_state_dict(sd)
+
+ feed = torch.cat([
+ inpaint_latent_mask,
+ model.model.process_latent_in(inpaint_latent)
+ ], dim=1)
+
+ inpaint_head_model.to(device=feed.device, dtype=feed.dtype)
+ inpaint_head_feature = inpaint_head_model(feed)
+
+ def input_block_patch(h, transformer_options):
+ if transformer_options["block"][1] == 0:
+ h = h + inpaint_head_feature.to(h)
+ return h
+
+ m = model.clone()
+ m.set_model_input_block_patch(input_block_patch)
+ return m
+
+ def swap(self):
+ if self.swapped:
+ return
+
+ if self.latent is None:
+ return
+
+ if self.latent_after_swap is None:
+ return
+
+ self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
+ self.swapped = True
+ return
+
+ def unswap(self):
+ if not self.swapped:
+ return
+
+ if self.latent is None:
+ return
+
+ if self.latent_after_swap is None:
+ return
+
+ self.latent, self.latent_after_swap = self.latent_after_swap, self.latent
+ self.swapped = False
+ return
+
+ def color_correction(self, img):
+ fg = img.astype(np.float32)
+ bg = self.image.copy().astype(np.float32)
+ w = self.mask[:, :, None].astype(np.float32) / 255.0
+ y = fg * w + bg * (1 - w)
+ return y.clip(0, 255).astype(np.uint8)
+
+ def post_process(self, img):
+ a, b, c, d = self.interested_area
+ content = resample_image(img, d - c, b - a)
+ result = self.image.copy()
+ result[a:b, c:d] = content
+ result = self.color_correction(result)
+ return result
+
+ def visualize_mask_processing(self):
+ return [self.interested_fill, self.interested_mask, self.interested_image]
+
diff --git a/modules/launch_util.py b/modules/launch_util.py
new file mode 100644
index 0000000000000000000000000000000000000000..00fff8aeaffb603bb7da09e9744cc7f2bd61703a
--- /dev/null
+++ b/modules/launch_util.py
@@ -0,0 +1,107 @@
+import os
+import importlib
+import importlib.util
+import subprocess
+import sys
+import re
+import logging
+
+
+logging.getLogger("torch.distributed.nn").setLevel(logging.ERROR) # sshh...
+logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
+
+re_requirement = re.compile(r"\s*([-_a-zA-Z0-9]+)\s*(?:==\s*([-+_.a-zA-Z0-9]+))?\s*")
+
+python = sys.executable
+default_command_live = (os.environ.get('LAUNCH_LIVE_OUTPUT') == "1")
+index_url = os.environ.get('INDEX_URL', "")
+
+modules_path = os.path.dirname(os.path.realpath(__file__))
+script_path = os.path.dirname(modules_path)
+
+
+def is_installed(package):
+ try:
+ spec = importlib.util.find_spec(package)
+ except ModuleNotFoundError:
+ return False
+
+ return spec is not None
+
+
+def run(command, desc=None, errdesc=None, custom_env=None, live: bool = default_command_live) -> str:
+ if desc is not None:
+ print(desc)
+
+ run_kwargs = {
+ "args": command,
+ "shell": True,
+ "env": os.environ if custom_env is None else custom_env,
+ "encoding": 'utf8',
+ "errors": 'ignore',
+ }
+
+ if not live:
+ run_kwargs["stdout"] = run_kwargs["stderr"] = subprocess.PIPE
+
+ result = subprocess.run(**run_kwargs)
+
+ if result.returncode != 0:
+ error_bits = [
+ f"{errdesc or 'Error running command'}.",
+ f"Command: {command}",
+ f"Error code: {result.returncode}",
+ ]
+ if result.stdout:
+ error_bits.append(f"stdout: {result.stdout}")
+ if result.stderr:
+ error_bits.append(f"stderr: {result.stderr}")
+ raise RuntimeError("\n".join(error_bits))
+
+ return (result.stdout or "")
+
+
+def run_pip(command, desc=None, live=default_command_live):
+ try:
+ index_url_line = f' --index-url {index_url}' if index_url != '' else ''
+ return run(f'"{python}" -m pip {command} --prefer-binary{index_url_line}', desc=f"Installing {desc}",
+ errdesc=f"Couldn't install {desc}", live=live)
+ except Exception as e:
+ print(e)
+ print(f'CMD Failed {desc}: {command}')
+ return None
+
+
+def requirements_met(requirements_file):
+ """
+ Does a simple parse of a requirements.txt file to determine if all rerqirements in it
+ are already installed. Returns True if so, False if not installed or parsing fails.
+ """
+
+ import importlib.metadata
+ import packaging.version
+
+ with open(requirements_file, "r", encoding="utf8") as file:
+ for line in file:
+ if line.strip() == "":
+ continue
+
+ m = re.match(re_requirement, line)
+ if m is None:
+ return False
+
+ package = m.group(1).strip()
+ version_required = (m.group(2) or "").strip()
+
+ if version_required == "":
+ continue
+
+ try:
+ version_installed = importlib.metadata.version(package)
+ except Exception:
+ return False
+
+ if packaging.version.parse(version_required) != packaging.version.parse(version_installed):
+ return False
+
+ return True
diff --git a/modules/localization.py b/modules/localization.py
new file mode 100644
index 0000000000000000000000000000000000000000..b21d4a564d134ac0be00d83c7005627d601d206e
--- /dev/null
+++ b/modules/localization.py
@@ -0,0 +1,60 @@
+import json
+import os
+
+
+current_translation = {}
+localization_root = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'language')
+
+
+def localization_js(filename):
+ global current_translation
+
+ if isinstance(filename, str):
+ full_name = os.path.abspath(os.path.join(localization_root, filename + '.json'))
+ if os.path.exists(full_name):
+ try:
+ with open(full_name, encoding='utf-8') as f:
+ current_translation = json.load(f)
+ assert isinstance(current_translation, dict)
+ for k, v in current_translation.items():
+ assert isinstance(k, str)
+ assert isinstance(v, str)
+ except Exception as e:
+ print(str(e))
+ print(f'Failed to load localization file {full_name}')
+
+ # current_translation = {k: 'XXX' for k in current_translation.keys()} # use this to see if all texts are covered
+
+ return f"window.localization = {json.dumps(current_translation)}"
+
+
+def dump_english_config(components):
+ all_texts = []
+ for c in components:
+ label = getattr(c, 'label', None)
+ value = getattr(c, 'value', None)
+ choices = getattr(c, 'choices', None)
+ info = getattr(c, 'info', None)
+
+ if isinstance(label, str):
+ all_texts.append(label)
+ if isinstance(value, str):
+ all_texts.append(value)
+ if isinstance(info, str):
+ all_texts.append(info)
+ if isinstance(choices, list):
+ for x in choices:
+ if isinstance(x, str):
+ all_texts.append(x)
+ if isinstance(x, tuple):
+ for y in x:
+ if isinstance(y, str):
+ all_texts.append(y)
+
+ config_dict = {k: k for k in all_texts if k != "" and 'progress-container' not in k}
+ full_name = os.path.abspath(os.path.join(localization_root, 'en.json'))
+
+ with open(full_name, "w", encoding="utf-8") as json_file:
+ json.dump(config_dict, json_file, indent=4)
+
+ return
diff --git a/modules/lora.py b/modules/lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..9eca0ed268c914ff79fcadf08106662c41affe15
--- /dev/null
+++ b/modules/lora.py
@@ -0,0 +1,140 @@
+def match_lora(lora, to_load):
+ patch_dict = {}
+ loaded_keys = set()
+ for x in to_load:
+ real_load_key = to_load[x]
+ if real_load_key in lora:
+ patch_dict[real_load_key] = lora[real_load_key]
+ loaded_keys.add(real_load_key)
+ continue
+
+ alpha_name = "{}.alpha".format(x)
+ alpha = None
+ if alpha_name in lora.keys():
+ alpha = lora[alpha_name].item()
+ loaded_keys.add(alpha_name)
+
+ regular_lora = "{}.lora_up.weight".format(x)
+ diffusers_lora = "{}_lora.up.weight".format(x)
+ transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
+ A_name = None
+
+ if regular_lora in lora.keys():
+ A_name = regular_lora
+ B_name = "{}.lora_down.weight".format(x)
+ mid_name = "{}.lora_mid.weight".format(x)
+ elif diffusers_lora in lora.keys():
+ A_name = diffusers_lora
+ B_name = "{}_lora.down.weight".format(x)
+ mid_name = None
+ elif transformers_lora in lora.keys():
+ A_name = transformers_lora
+ B_name ="{}.lora_linear_layer.down.weight".format(x)
+ mid_name = None
+
+ if A_name is not None:
+ mid = None
+ if mid_name is not None and mid_name in lora.keys():
+ mid = lora[mid_name]
+ loaded_keys.add(mid_name)
+ patch_dict[to_load[x]] = (lora[A_name], lora[B_name], alpha, mid)
+ loaded_keys.add(A_name)
+ loaded_keys.add(B_name)
+
+
+ ######## loha
+ hada_w1_a_name = "{}.hada_w1_a".format(x)
+ hada_w1_b_name = "{}.hada_w1_b".format(x)
+ hada_w2_a_name = "{}.hada_w2_a".format(x)
+ hada_w2_b_name = "{}.hada_w2_b".format(x)
+ hada_t1_name = "{}.hada_t1".format(x)
+ hada_t2_name = "{}.hada_t2".format(x)
+ if hada_w1_a_name in lora.keys():
+ hada_t1 = None
+ hada_t2 = None
+ if hada_t1_name in lora.keys():
+ hada_t1 = lora[hada_t1_name]
+ hada_t2 = lora[hada_t2_name]
+ loaded_keys.add(hada_t1_name)
+ loaded_keys.add(hada_t2_name)
+
+ patch_dict[to_load[x]] = (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2)
+ loaded_keys.add(hada_w1_a_name)
+ loaded_keys.add(hada_w1_b_name)
+ loaded_keys.add(hada_w2_a_name)
+ loaded_keys.add(hada_w2_b_name)
+
+
+ ######## lokr
+ lokr_w1_name = "{}.lokr_w1".format(x)
+ lokr_w2_name = "{}.lokr_w2".format(x)
+ lokr_w1_a_name = "{}.lokr_w1_a".format(x)
+ lokr_w1_b_name = "{}.lokr_w1_b".format(x)
+ lokr_t2_name = "{}.lokr_t2".format(x)
+ lokr_w2_a_name = "{}.lokr_w2_a".format(x)
+ lokr_w2_b_name = "{}.lokr_w2_b".format(x)
+
+ lokr_w1 = None
+ if lokr_w1_name in lora.keys():
+ lokr_w1 = lora[lokr_w1_name]
+ loaded_keys.add(lokr_w1_name)
+
+ lokr_w2 = None
+ if lokr_w2_name in lora.keys():
+ lokr_w2 = lora[lokr_w2_name]
+ loaded_keys.add(lokr_w2_name)
+
+ lokr_w1_a = None
+ if lokr_w1_a_name in lora.keys():
+ lokr_w1_a = lora[lokr_w1_a_name]
+ loaded_keys.add(lokr_w1_a_name)
+
+ lokr_w1_b = None
+ if lokr_w1_b_name in lora.keys():
+ lokr_w1_b = lora[lokr_w1_b_name]
+ loaded_keys.add(lokr_w1_b_name)
+
+ lokr_w2_a = None
+ if lokr_w2_a_name in lora.keys():
+ lokr_w2_a = lora[lokr_w2_a_name]
+ loaded_keys.add(lokr_w2_a_name)
+
+ lokr_w2_b = None
+ if lokr_w2_b_name in lora.keys():
+ lokr_w2_b = lora[lokr_w2_b_name]
+ loaded_keys.add(lokr_w2_b_name)
+
+ lokr_t2 = None
+ if lokr_t2_name in lora.keys():
+ lokr_t2 = lora[lokr_t2_name]
+ loaded_keys.add(lokr_t2_name)
+
+ if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
+ patch_dict[to_load[x]] = (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2)
+
+ w_norm_name = "{}.w_norm".format(x)
+ b_norm_name = "{}.b_norm".format(x)
+ w_norm = lora.get(w_norm_name, None)
+ b_norm = lora.get(b_norm_name, None)
+
+ if w_norm is not None:
+ loaded_keys.add(w_norm_name)
+ patch_dict[to_load[x]] = (w_norm,)
+ if b_norm is not None:
+ loaded_keys.add(b_norm_name)
+ patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = (b_norm,)
+
+ diff_name = "{}.diff".format(x)
+ diff_weight = lora.get(diff_name, None)
+ if diff_weight is not None:
+ patch_dict[to_load[x]] = (diff_weight,)
+ loaded_keys.add(diff_name)
+
+ diff_bias_name = "{}.diff_b".format(x)
+ diff_bias = lora.get(diff_bias_name, None)
+ if diff_bias is not None:
+ patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = (diff_bias,)
+ loaded_keys.add(diff_bias_name)
+
+ remaining_dict = {x: y for x, y in lora.items() if x not in loaded_keys}
+ return patch_dict, remaining_dict
diff --git a/modules/model_loader.py b/modules/model_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ba336a915ae234b7cd5f9f2576d4edf779738ba
--- /dev/null
+++ b/modules/model_loader.py
@@ -0,0 +1,26 @@
+import os
+from urllib.parse import urlparse
+from typing import Optional
+
+
+def load_file_from_url(
+ url: str,
+ *,
+ model_dir: str,
+ progress: bool = True,
+ file_name: Optional[str] = None,
+) -> str:
+ """Download a file from `url` into `model_dir`, using the file present if possible.
+
+ Returns the path to the downloaded file.
+ """
+ os.makedirs(model_dir, exist_ok=True)
+ if not file_name:
+ parts = urlparse(url)
+ file_name = os.path.basename(parts.path)
+ cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+ if not os.path.exists(cached_file):
+ print(f'Downloading: "{url}" to {cached_file}\n')
+ from torch.hub import download_url_to_file
+ download_url_to_file(url, cached_file, progress=progress)
+ return cached_file
diff --git a/modules/patch.py b/modules/patch.py
new file mode 100644
index 0000000000000000000000000000000000000000..a65e171fa066b2f4b5eed71f2fbd4b9dc5aceac2
--- /dev/null
+++ b/modules/patch.py
@@ -0,0 +1,555 @@
+import os
+import torch
+import math
+import time
+import numpy as np
+import fcbh.model_base
+import fcbh.ldm.modules.diffusionmodules.openaimodel
+import fcbh.samplers
+import fcbh.model_management
+import modules.anisotropic as anisotropic
+import fcbh.ldm.modules.attention
+import fcbh.k_diffusion.sampling
+import fcbh.sd1_clip
+import modules.inpaint_worker as inpaint_worker
+import fcbh.ldm.modules.diffusionmodules.openaimodel
+import fcbh.ldm.modules.diffusionmodules.model
+import fcbh.sd
+import fcbh.cldm.cldm
+import fcbh.model_patcher
+import fcbh.samplers
+import fcbh.cli_args
+import modules.advanced_parameters as advanced_parameters
+import warnings
+import safetensors.torch
+import modules.constants as constants
+
+from einops import repeat
+from fcbh.k_diffusion.sampling import BatchedBrownianTree
+from fcbh.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, apply_control
+from fcbh.ldm.modules.diffusionmodules.util import make_beta_schedule
+
+
+sharpness = 2.0
+
+adm_scaler_end = 0.3
+positive_adm_scale = 1.5
+negative_adm_scale = 0.8
+
+adaptive_cfg = 7.0
+global_diffusion_progress = 0
+eps_record = None
+
+
+def calculate_weight_patched(self, patches, weight, key):
+ for p in patches:
+ alpha = p[0]
+ v = p[1]
+ strength_model = p[2]
+
+ if strength_model != 1.0:
+ weight *= strength_model
+
+ if isinstance(v, list):
+ v = (self.calculate_weight(v[1:], v[0].clone(), key),)
+
+ if len(v) == 1:
+ w1 = v[0]
+ if alpha != 0.0:
+ if w1.shape != weight.shape:
+ print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+ else:
+ weight += alpha * fcbh.model_management.cast_to_device(w1, weight.device, weight.dtype)
+ elif len(v) == 3:
+ # fooocus
+ w1 = fcbh.model_management.cast_to_device(v[0], weight.device, torch.float32)
+ w_min = fcbh.model_management.cast_to_device(v[1], weight.device, torch.float32)
+ w_max = fcbh.model_management.cast_to_device(v[2], weight.device, torch.float32)
+ w1 = (w1 / 255.0) * (w_max - w_min) + w_min
+ if alpha != 0.0:
+ if w1.shape != weight.shape:
+ print("WARNING SHAPE MISMATCH {} FOOOCUS WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+ else:
+ weight += alpha * fcbh.model_management.cast_to_device(w1, weight.device, weight.dtype)
+ elif len(v) == 4: # lora/locon
+ mat1 = fcbh.model_management.cast_to_device(v[0], weight.device, torch.float32)
+ mat2 = fcbh.model_management.cast_to_device(v[1], weight.device, torch.float32)
+ if v[2] is not None:
+ alpha *= v[2] / mat2.shape[0]
+ if v[3] is not None:
+ # locon mid weights, hopefully the math is fine because I didn't properly test it
+ mat3 = fcbh.model_management.cast_to_device(v[3], weight.device, torch.float32)
+ 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
+
+ if w1 is None:
+ dim = w1_b.shape[0]
+ w1 = torch.mm(fcbh.model_management.cast_to_device(w1_a, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w1_b, weight.device, torch.float32))
+ else:
+ w1 = fcbh.model_management.cast_to_device(w1, weight.device, torch.float32)
+
+ if w2 is None:
+ dim = w2_b.shape[0]
+ if t2 is None:
+ w2 = torch.mm(fcbh.model_management.cast_to_device(w2_a, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2_b, weight.device, torch.float32))
+ else:
+ w2 = torch.einsum('i j k l, j r, i p -> p r k l',
+ fcbh.model_management.cast_to_device(t2, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2_b, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2_a, weight.device, torch.float32))
+ else:
+ w2 = fcbh.model_management.cast_to_device(w2, weight.device, torch.float32)
+
+ if len(w2.shape) == 4:
+ w1 = w1.unsqueeze(2).unsqueeze(2)
+ if v[2] is not None and dim is not None:
+ alpha *= v[2] / dim
+
+ try:
+ weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
+ except Exception as e:
+ print("ERROR", key, e)
+ 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',
+ fcbh.model_management.cast_to_device(t1, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w1b, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w1a, weight.device, torch.float32))
+
+ m2 = torch.einsum('i j k l, j r, i p -> p r k l',
+ fcbh.model_management.cast_to_device(t2, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2b, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2a, weight.device, torch.float32))
+ else:
+ m1 = torch.mm(fcbh.model_management.cast_to_device(w1a, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w1b, weight.device, torch.float32))
+ m2 = torch.mm(fcbh.model_management.cast_to_device(w2a, weight.device, torch.float32),
+ fcbh.model_management.cast_to_device(w2b, weight.device, torch.float32))
+
+ try:
+ weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
+ except Exception as e:
+ print("ERROR", key, e)
+
+ return weight
+
+
+class BrownianTreeNoiseSamplerPatched:
+ transform = None
+ tree = None
+ global_sigma_min = 1.0
+ global_sigma_max = 1.0
+
+ @staticmethod
+ def global_init(x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
+ t0, t1 = transform(torch.as_tensor(sigma_min)), transform(torch.as_tensor(sigma_max))
+
+ BrownianTreeNoiseSamplerPatched.transform = transform
+ BrownianTreeNoiseSamplerPatched.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
+
+ BrownianTreeNoiseSamplerPatched.global_sigma_min = sigma_min
+ BrownianTreeNoiseSamplerPatched.global_sigma_max = sigma_max
+
+ def __init__(self, *args, **kwargs):
+ pass
+
+ @staticmethod
+ def __call__(sigma, sigma_next):
+ transform = BrownianTreeNoiseSamplerPatched.transform
+ tree = BrownianTreeNoiseSamplerPatched.tree
+
+ t0, t1 = transform(torch.as_tensor(sigma)), transform(torch.as_tensor(sigma_next))
+ return tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+
+def compute_cfg(uncond, cond, cfg_scale, t):
+ global adaptive_cfg
+
+ mimic_cfg = float(adaptive_cfg)
+ real_cfg = float(cfg_scale)
+
+ real_eps = uncond + real_cfg * (cond - uncond)
+
+ if cfg_scale > adaptive_cfg:
+ mimicked_eps = uncond + mimic_cfg * (cond - uncond)
+ return real_eps * t + mimicked_eps * (1 - t)
+ else:
+ return real_eps
+
+
+def patched_sampler_cfg_function(args):
+ global eps_record
+
+ positive_eps = args['cond']
+ negative_eps = args['uncond']
+ cfg_scale = args['cond_scale']
+ positive_x0 = args['input'] - positive_eps
+ sigma = args['sigma']
+
+ alpha = 0.001 * sharpness * global_diffusion_progress
+ positive_eps_degraded = anisotropic.adaptive_anisotropic_filter(x=positive_eps, g=positive_x0)
+ positive_eps_degraded_weighted = positive_eps_degraded * alpha + positive_eps * (1.0 - alpha)
+
+ final_eps = compute_cfg(uncond=negative_eps, cond=positive_eps_degraded_weighted,
+ cfg_scale=cfg_scale, t=global_diffusion_progress)
+
+ if eps_record is not None:
+ eps_record = (final_eps / sigma).cpu()
+
+ return final_eps
+
+
+def sdxl_encode_adm_patched(self, **kwargs):
+ global positive_adm_scale, negative_adm_scale
+
+ clip_pooled = fcbh.model_base.sdxl_pooled(kwargs, self.noise_augmentor)
+ width = kwargs.get("width", 768)
+ height = kwargs.get("height", 768)
+ target_width = width
+ target_height = height
+
+ if kwargs.get("prompt_type", "") == "negative":
+ width = float(width) * negative_adm_scale
+ height = float(height) * negative_adm_scale
+ elif kwargs.get("prompt_type", "") == "positive":
+ width = float(width) * positive_adm_scale
+ height = float(height) * positive_adm_scale
+
+ # Avoid artifacts
+ width = int(width)
+ height = int(height)
+ crop_w = 0
+ crop_h = 0
+ target_width = int(target_width)
+ target_height = int(target_height)
+
+ out_a = [self.embedder(torch.Tensor([height])), self.embedder(torch.Tensor([width])),
+ self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
+ self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
+ flat_a = torch.flatten(torch.cat(out_a)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+
+ out_b = [self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width])),
+ self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
+ self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
+ flat_b = torch.flatten(torch.cat(out_b)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+
+ return torch.cat((clip_pooled.to(flat_a.device), flat_a, clip_pooled.to(flat_b.device), flat_b), dim=1)
+
+
+def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
+ to_encode = list()
+ max_token_len = 0
+ has_weights = False
+ for x in token_weight_pairs:
+ tokens = list(map(lambda a: a[0], x))
+ max_token_len = max(len(tokens), max_token_len)
+ has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
+ to_encode.append(tokens)
+
+ sections = len(to_encode)
+ if has_weights or sections == 0:
+ to_encode.append(fcbh.sd1_clip.gen_empty_tokens(self.special_tokens, max_token_len))
+
+ out, pooled = self.encode(to_encode)
+ if pooled is not None:
+ first_pooled = pooled[0:1].cpu()
+ else:
+ first_pooled = pooled
+
+ output = []
+ for k in range(0, sections):
+ z = out[k:k + 1]
+ if has_weights:
+ original_mean = z.mean()
+ z_empty = out[-1]
+ for i in range(len(z)):
+ for j in range(len(z[i])):
+ weight = token_weight_pairs[k][j][1]
+ if weight != 1.0:
+ z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
+ new_mean = z.mean()
+ z = z * (original_mean / new_mean)
+ output.append(z)
+
+ if len(output) == 0:
+ return out[-1:].cpu(), first_pooled
+
+ return torch.cat(output, dim=-2).cpu(), first_pooled
+
+
+def patched_KSamplerX0Inpaint_forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
+ if inpaint_worker.current_task is not None:
+ latent_processor = self.inner_model.inner_model.process_latent_in
+ inpaint_latent = latent_processor(inpaint_worker.current_task.latent).to(x)
+ inpaint_mask = inpaint_worker.current_task.latent_mask.to(x)
+
+ if getattr(self, 'energy_generator', None) is None:
+ # avoid bad results by using different seeds.
+ self.energy_generator = torch.Generator(device='cpu').manual_seed((seed + 1) % constants.MAX_SEED)
+
+ energy_sigma = sigma.reshape([sigma.shape[0]] + [1] * (len(x.shape) - 1))
+ current_energy = torch.randn(
+ x.size(), dtype=x.dtype, generator=self.energy_generator, device="cpu").to(x) * energy_sigma
+ x = x * inpaint_mask + (inpaint_latent + current_energy) * (1.0 - inpaint_mask)
+
+ out = self.inner_model(x, sigma,
+ cond=cond,
+ uncond=uncond,
+ cond_scale=cond_scale,
+ model_options=model_options,
+ seed=seed)
+
+ out = out * inpaint_mask + inpaint_latent * (1.0 - inpaint_mask)
+ else:
+ out = self.inner_model(x, sigma,
+ cond=cond,
+ uncond=uncond,
+ cond_scale=cond_scale,
+ model_options=model_options,
+ seed=seed)
+ return out
+
+
+def timed_adm(y, timesteps):
+ if isinstance(y, torch.Tensor) and int(y.dim()) == 2 and int(y.shape[1]) == 5632:
+ y_mask = (timesteps > 999.0 * (1.0 - float(adm_scaler_end))).to(y)[..., None]
+ y_with_adm = y[..., :2816].clone()
+ y_without_adm = y[..., 2816:].clone()
+ return y_with_adm * y_mask + y_without_adm * (1.0 - y_mask)
+ return y
+
+
+def patched_timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+ # Consistent with Kohya to reduce differences between model training and inference.
+
+ if not repeat_only:
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=timesteps.device)
+ args = timesteps[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ else:
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
+ return embedding
+
+
+def patched_cldm_forward(self, x, hint, timesteps, context, y=None, **kwargs):
+ t_emb = fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(
+ timesteps, self.model_channels, repeat_only=False).to(self.dtype)
+
+ emb = self.time_embed(t_emb)
+
+ guided_hint = self.input_hint_block(hint, emb, context)
+
+ y = timed_adm(y, timesteps)
+
+ outs = []
+
+ hs = []
+ if self.num_classes is not None:
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ h = x.type(self.dtype)
+ for module, zero_conv in zip(self.input_blocks, self.zero_convs):
+ if guided_hint is not None:
+ h = module(h, emb, context)
+ h += guided_hint
+ guided_hint = None
+ else:
+ h = module(h, emb, context)
+ outs.append(zero_conv(h, emb, context))
+
+ h = self.middle_block(h, emb, context)
+ outs.append(self.middle_block_out(h, emb, context))
+
+ if advanced_parameters.controlnet_softness > 0:
+ for i in range(10):
+ k = 1.0 - float(i) / 9.0
+ outs[i] = outs[i] * (1.0 - advanced_parameters.controlnet_softness * k)
+
+ return outs
+
+
+def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+ global global_diffusion_progress
+
+ self.current_step = 1.0 - timesteps.to(x) / 999.0
+ global_diffusion_progress = float(self.current_step.detach().cpu().numpy().tolist()[0])
+
+ transformer_options["original_shape"] = list(x.shape)
+ transformer_options["current_index"] = 0
+ transformer_patches = transformer_options.get("patches", {})
+
+ y = timed_adm(y, timesteps)
+
+ hs = []
+ t_emb = fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(
+ timesteps, self.model_channels, repeat_only=False).to(self.dtype)
+ emb = self.time_embed(t_emb)
+
+ if self.num_classes is not None:
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ h = x.type(self.dtype)
+ for id, module in enumerate(self.input_blocks):
+ transformer_options["block"] = ("input", id)
+ h = forward_timestep_embed(module, h, emb, context, transformer_options)
+ h = apply_control(h, control, 'input')
+ if "input_block_patch" in transformer_patches:
+ patch = transformer_patches["input_block_patch"]
+ for p in patch:
+ h = p(h, transformer_options)
+
+ hs.append(h)
+ if "input_block_patch_after_skip" in transformer_patches:
+ patch = transformer_patches["input_block_patch_after_skip"]
+ for p in patch:
+ h = p(h, transformer_options)
+
+ transformer_options["block"] = ("middle", 0)
+ h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options)
+ h = apply_control(h, control, 'middle')
+
+ for id, module in enumerate(self.output_blocks):
+ transformer_options["block"] = ("output", id)
+ hsp = hs.pop()
+ hsp = apply_control(hsp, control, 'output')
+
+ if "output_block_patch" in transformer_patches:
+ patch = transformer_patches["output_block_patch"]
+ for p in patch:
+ h, hsp = p(h, hsp, transformer_options)
+
+ h = torch.cat([h, hsp], dim=1)
+ del hsp
+ if len(hs) > 0:
+ output_shape = hs[-1].shape
+ else:
+ output_shape = None
+ h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape)
+ h = h.type(x.dtype)
+ if self.predict_codebook_ids:
+ return self.id_predictor(h)
+ else:
+ return self.out(h)
+
+
+def patched_register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+ # Consistent with Kohya to reduce differences between model training and inference.
+
+ if given_betas is not None:
+ betas = given_betas
+ else:
+ betas = make_beta_schedule(
+ beta_schedule,
+ timesteps,
+ linear_start=linear_start,
+ linear_end=linear_end,
+ cosine_s=cosine_s)
+
+ alphas = 1. - betas
+ alphas_cumprod = np.cumprod(alphas, axis=0)
+ timesteps, = betas.shape
+ self.num_timesteps = int(timesteps)
+ self.linear_start = linear_start
+ self.linear_end = linear_end
+ sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
+ self.set_sigmas(sigmas)
+ return
+
+
+def patched_load_models_gpu(*args, **kwargs):
+ execution_start_time = time.perf_counter()
+ y = fcbh.model_management.load_models_gpu_origin(*args, **kwargs)
+ moving_time = time.perf_counter() - execution_start_time
+ if moving_time > 0.1:
+ print(f'[Fooocus Model Management] Moving model(s) has taken {moving_time:.2f} seconds')
+ return y
+
+
+def build_loaded(module, loader_name):
+ original_loader_name = loader_name + '_origin'
+
+ if not hasattr(module, original_loader_name):
+ setattr(module, original_loader_name, getattr(module, loader_name))
+
+ original_loader = getattr(module, original_loader_name)
+
+ def loader(*args, **kwargs):
+ result = None
+ try:
+ result = original_loader(*args, **kwargs)
+ except Exception as e:
+ result = None
+ exp = str(e) + '\n'
+ for path in list(args) + list(kwargs.values()):
+ if isinstance(path, str):
+ if os.path.exists(path):
+ exp += f'File corrupted: {path} \n'
+ corrupted_backup_file = path + '.corrupted'
+ if os.path.exists(corrupted_backup_file):
+ os.remove(corrupted_backup_file)
+ os.replace(path, corrupted_backup_file)
+ if os.path.exists(path):
+ os.remove(path)
+ exp += f'Fooocus has tried to move the corrupted file to {corrupted_backup_file} \n'
+ exp += f'You may try again now and Fooocus will download models again. \n'
+ raise ValueError(exp)
+ return result
+
+ setattr(module, loader_name, loader)
+ return
+
+
+def patch_all():
+ if not hasattr(fcbh.model_management, 'load_models_gpu_origin'):
+ fcbh.model_management.load_models_gpu_origin = fcbh.model_management.load_models_gpu
+
+ fcbh.model_management.load_models_gpu = patched_load_models_gpu
+ fcbh.model_patcher.ModelPatcher.calculate_weight = calculate_weight_patched
+ fcbh.cldm.cldm.ControlNet.forward = patched_cldm_forward
+ fcbh.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
+ fcbh.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
+ fcbh.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights_patched_with_a1111_method
+ fcbh.samplers.KSamplerX0Inpaint.forward = patched_KSamplerX0Inpaint_forward
+ fcbh.k_diffusion.sampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSamplerPatched
+ fcbh.ldm.modules.diffusionmodules.openaimodel.timestep_embedding = patched_timestep_embedding
+ fcbh.model_base.ModelSamplingDiscrete._register_schedule = patched_register_schedule
+
+ warnings.filterwarnings(action='ignore', module='torchsde')
+
+ build_loaded(safetensors.torch, 'load_file')
+ build_loaded(torch, 'load')
+
+ return
diff --git a/modules/private_logger.py b/modules/private_logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..c2c0409482bcbd003d44c9f7bfd10ba96221962b
--- /dev/null
+++ b/modules/private_logger.py
@@ -0,0 +1,57 @@
+import os
+import args_manager
+import modules.config
+
+from PIL import Image
+from modules.util import generate_temp_filename
+
+
+log_cache = {}
+
+
+def get_current_html_path():
+ date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs,
+ extension='png')
+ html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+ return html_name
+
+
+def log(img, dic, single_line_number=3):
+ if args_manager.args.disable_image_log:
+ return
+
+ date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs, extension='png')
+ os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
+ Image.fromarray(img).save(local_temp_filename)
+ html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+
+ existing_log = log_cache.get(html_name, None)
+
+ if existing_log is None:
+ if os.path.exists(html_name):
+ existing_log = open(html_name, encoding='utf-8').read()
+ else:
+ existing_log = f'
+
+
+ *text*
+Fooocus Log {date_string} (private)
\nAll images do not contain any hidden data.
' + + div_name = only_name.replace('.', '_') + item = f'\n'
+ item += f"
\n"
+ existing_log = item + existing_log
+
+ with open(html_name, 'w', encoding='utf-8') as f:
+ f.write(existing_log)
+
+ print(f'Image generated with private log at: {html_name}')
+
+ log_cache[html_name] = existing_log
+
+ return
diff --git a/modules/sample_hijack.py b/modules/sample_hijack.py
new file mode 100644
index 0000000000000000000000000000000000000000..eafda2dda990ffcc21ad4872df14efd22454396a
--- /dev/null
+++ b/modules/sample_hijack.py
@@ -0,0 +1,155 @@
+import torch
+import fcbh.samplers
+import fcbh.model_management
+
+from fcbh.model_base import SDXLRefiner, SDXL
+from fcbh.conds import CONDRegular
+from fcbh.sample import get_additional_models, get_models_from_cond, cleanup_additional_models
+from fcbh.samplers import resolve_areas_and_cond_masks, wrap_model, calculate_start_end_timesteps, \
+ create_cond_with_same_area_if_none, pre_run_control, apply_empty_x_to_equal_area, encode_model_conds
+
+
+current_refiner = None
+refiner_switch_step = -1
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate_inner(c, p, target_model=None, target_clip=None):
+ if target_model is None or isinstance(target_model, SDXLRefiner):
+ c = c[..., -1280:].clone()
+ elif isinstance(target_model, SDXL):
+ c = c.clone()
+ else:
+ p = None
+ c = c[..., :768].clone()
+
+ final_layer_norm = target_clip.cond_stage_model.clip_l.transformer.text_model.final_layer_norm
+
+ final_layer_norm_origin_device = final_layer_norm.weight.device
+ final_layer_norm_origin_dtype = final_layer_norm.weight.dtype
+
+ c_origin_device = c.device
+ c_origin_dtype = c.dtype
+
+ final_layer_norm.to(device='cpu', dtype=torch.float32)
+ c = c.to(device='cpu', dtype=torch.float32)
+
+ c = torch.chunk(c, int(c.size(1)) // 77, 1)
+ c = [final_layer_norm(ci) for ci in c]
+ c = torch.cat(c, dim=1)
+
+ final_layer_norm.to(device=final_layer_norm_origin_device, dtype=final_layer_norm_origin_dtype)
+ c = c.to(device=c_origin_device, dtype=c_origin_dtype)
+ return c, p
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate(cond, target_model=None, target_clip=None):
+ results = []
+
+ for c, px in cond:
+ p = px.get('pooled_output', None)
+ c, p = clip_separate_inner(c, p, target_model=target_model, target_clip=target_clip)
+ p = {} if p is None else {'pooled_output': p.clone()}
+ results.append([c, p])
+
+ return results
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def clip_separate_after_preparation(cond, target_model=None, target_clip=None):
+ results = []
+
+ for x in cond:
+ p = x.get('pooled_output', None)
+ c = x['model_conds']['c_crossattn'].cond
+
+ c, p = clip_separate_inner(c, p, target_model=target_model, target_clip=target_clip)
+
+ result = {'model_conds': {'c_crossattn': CONDRegular(c)}}
+
+ if p is not None:
+ result['pooled_output'] = p.clone()
+
+ results.append(result)
+
+ return results
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def sample_hacked(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
+ global current_refiner
+
+ positive = positive[:]
+ negative = negative[:]
+
+ resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
+ resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
+
+ model_wrap = wrap_model(model)
+
+ calculate_start_end_timesteps(model, negative)
+ calculate_start_end_timesteps(model, positive)
+
+ #make sure each cond area has an opposite one with the same area
+ for c in positive:
+ create_cond_with_same_area_if_none(negative, c)
+ for c in negative:
+ create_cond_with_same_area_if_none(positive, c)
+
+ # pre_run_control(model, negative + positive)
+ pre_run_control(model, positive) # negative is not necessary in Fooocus, 0.5s faster.
+
+ apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
+ apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
+
+ if latent_image is not None:
+ latent_image = model.process_latent_in(latent_image)
+
+ if hasattr(model, 'extra_conds'):
+ positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
+ negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
+
+ extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
+
+ if current_refiner is not None and hasattr(current_refiner.model, 'extra_conds'):
+ positive_refiner = clip_separate_after_preparation(positive, target_model=current_refiner.model)
+ negative_refiner = clip_separate_after_preparation(negative, target_model=current_refiner.model)
+
+ positive_refiner = encode_model_conds(current_refiner.model.extra_conds, positive_refiner, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask)
+ negative_refiner = encode_model_conds(current_refiner.model.extra_conds, negative_refiner, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask)
+
+ def refiner_switch():
+ cleanup_additional_models(set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control")))
+
+ extra_args["cond"] = positive_refiner
+ extra_args["uncond"] = negative_refiner
+
+ # clear ip-adapter for refiner
+ extra_args['model_options'] = {k: {} if k == 'transformer_options' else v for k, v in extra_args['model_options'].items()}
+
+ models, inference_memory = get_additional_models(positive_refiner, negative_refiner, current_refiner.model_dtype())
+ fcbh.model_management.load_models_gpu([current_refiner] + models, current_refiner.memory_required(noise.shape) + inference_memory)
+
+ model_wrap.inner_model = current_refiner.model
+ print('Refiner Swapped')
+ return
+
+ def callback_wrap(step, x0, x, total_steps):
+ if step == refiner_switch_step and current_refiner is not None:
+ refiner_switch()
+ if callback is not None:
+ # residual_noise_preview = x - x0
+ # residual_noise_preview /= residual_noise_preview.std()
+ # residual_noise_preview *= x0.std()
+ callback(step, x0, x, total_steps)
+
+ samples = sampler.sample(model_wrap, sigmas, extra_args, callback_wrap, noise, latent_image, denoise_mask, disable_pbar)
+ return model.process_latent_out(samples.to(torch.float32))
+
+
+fcbh.samplers.sample = sample_hacked
diff --git a/modules/sdxl_styles.py b/modules/sdxl_styles.py
new file mode 100644
index 0000000000000000000000000000000000000000..d74894557eb786f0ecbaa09e3f9ba26dc22245dc
--- /dev/null
+++ b/modules/sdxl_styles.py
@@ -0,0 +1,81 @@
+import os
+import re
+import json
+
+from modules.util import get_files_from_folder
+
+
+# cannot use modules.config - validators causing circular imports
+styles_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../sdxl_styles/'))
+wildcards_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../wildcards/'))
+wildcards_max_bfs_depth = 64
+
+
+def normalize_key(k):
+ k = k.replace('-', ' ')
+ words = k.split(' ')
+ words = [w[:1].upper() + w[1:].lower() for w in words]
+ k = ' '.join(words)
+ k = k.replace('3d', '3D')
+ k = k.replace('Sai', 'SAI')
+ k = k.replace('Mre', 'MRE')
+ k = k.replace('(s', '(S')
+ return k
+
+
+styles = {}
+
+styles_files = get_files_from_folder(styles_path, ['.json'])
+
+for x in ['sdxl_styles_fooocus.json',
+ 'sdxl_styles_sai.json',
+ 'sdxl_styles_mre.json',
+ 'sdxl_styles_twri.json',
+ 'sdxl_styles_diva.json']:
+ if x in styles_files:
+ styles_files.remove(x)
+ styles_files.append(x)
+
+for styles_file in styles_files:
+ try:
+ with open(os.path.join(styles_path, styles_file), encoding='utf-8') as f:
+ for entry in json.load(f):
+ name = normalize_key(entry['name'])
+ prompt = entry['prompt'] if 'prompt' in entry else ''
+ negative_prompt = entry['negative_prompt'] if 'negative_prompt' in entry else ''
+ styles[name] = (prompt, negative_prompt)
+ except Exception as e:
+ print(str(e))
+ print(f'Failed to load style file {styles_file}')
+
+style_keys = list(styles.keys())
+fooocus_expansion = "Fooocus V2"
+legal_style_names = [fooocus_expansion] + style_keys
+
+
+def apply_style(style, positive):
+ p, n = styles[style]
+ return p.replace('{prompt}', positive).splitlines(), n.splitlines()
+
+
+def apply_wildcards(wildcard_text, rng, directory=wildcards_path):
+ for _ in range(wildcards_max_bfs_depth):
+ placeholders = re.findall(r'__([\w-]+)__', wildcard_text)
+ if len(placeholders) == 0:
+ return wildcard_text
+
+ print(f'[Wildcards] processing: {wildcard_text}')
+ for placeholder in placeholders:
+ try:
+ words = open(os.path.join(directory, f'{placeholder}.txt'), encoding='utf-8').read().splitlines()
+ words = [x for x in words if x != '']
+ assert len(words) > 0
+ wildcard_text = wildcard_text.replace(f'__{placeholder}__', rng.choice(words), 1)
+ except:
+ print(f'[Wildcards] Warning: {placeholder}.txt missing or empty. '
+ f'Using "{placeholder}" as a normal word.')
+ wildcard_text = wildcard_text.replace(f'__{placeholder}__', placeholder)
+ print(f'[Wildcards] {wildcard_text}')
+
+ print(f'[Wildcards] BFS stack overflow. Current text: {wildcard_text}')
+ return wildcard_text
diff --git a/modules/style_sorter.py b/modules/style_sorter.py
new file mode 100644
index 0000000000000000000000000000000000000000..393e441dad152f97bce8766113ed456118bb9506
--- /dev/null
+++ b/modules/style_sorter.py
@@ -0,0 +1,56 @@
+import os
+import gradio as gr
+import modules.localization as localization
+import json
+
+
+all_styles = []
+
+
+def try_load_sorted_styles(style_names, default_selected):
+ global all_styles
+
+ all_styles = style_names
+
+ try:
+ if os.path.exists('sorted_styles.json'):
+ with open('sorted_styles.json', 'rt', encoding='utf-8') as fp:
+ sorted_styles = json.load(fp)
+ if len(sorted_styles) == len(all_styles):
+ if all(x in all_styles for x in sorted_styles):
+ if all(x in sorted_styles for x in all_styles):
+ all_styles = sorted_styles
+ except Exception as e:
+ print('Load style sorting failed.')
+ print(e)
+
+ unselected = [y for y in all_styles if y not in default_selected]
+ all_styles = default_selected + unselected
+
+ return
+
+
+def sort_styles(selected):
+ global all_styles
+ unselected = [y for y in all_styles if y not in selected]
+ sorted_styles = selected + unselected
+ try:
+ with open('sorted_styles.json', 'wt', encoding='utf-8') as fp:
+ json.dump(sorted_styles, fp, indent=4)
+ except Exception as e:
+ print('Write style sorting failed.')
+ print(e)
+ all_styles = sorted_styles
+ return gr.CheckboxGroup.update(choices=sorted_styles)
+
+
+def localization_key(x):
+ return x + localization.current_translation.get(x, '')
+
+
+def search_styles(selected, query):
+ unselected = [y for y in all_styles if y not in selected]
+ matched = [y for y in unselected if query.lower() in localization_key(y).lower()] if len(query.replace(' ', '')) > 0 else []
+ unmatched = [y for y in unselected if y not in matched]
+ sorted_styles = matched + selected + unmatched
+ return gr.CheckboxGroup.update(choices=sorted_styles)
diff --git a/modules/ui_gradio_extensions.py b/modules/ui_gradio_extensions.py
new file mode 100644
index 0000000000000000000000000000000000000000..e59b151bf448ebd3ac036b74f16919a1d49dd200
--- /dev/null
+++ b/modules/ui_gradio_extensions.py
@@ -0,0 +1,65 @@
+# based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/v1.6.0/modules/ui_gradio_extensions.py
+
+import os
+import gradio as gr
+import args_manager
+
+from modules.localization import localization_js
+
+
+GradioTemplateResponseOriginal = gr.routes.templates.TemplateResponse
+
+modules_path = os.path.dirname(os.path.realpath(__file__))
+script_path = os.path.dirname(modules_path)
+
+
+def webpath(fn):
+ if fn.startswith(script_path):
+ web_path = os.path.relpath(fn, script_path).replace('\\', '/')
+ else:
+ web_path = os.path.abspath(fn)
+
+ return f'file={web_path}?{os.path.getmtime(fn)}'
+
+
+def javascript_html():
+ script_js_path = webpath('javascript/script.js')
+ context_menus_js_path = webpath('javascript/contextMenus.js')
+ localization_js_path = webpath('javascript/localization.js')
+ zoom_js_path = webpath('javascript/zoom.js')
+ edit_attention_js_path = webpath('javascript/edit-attention.js')
+ viewer_js_path = webpath('javascript/viewer.js')
+ image_viewer_js_path = webpath('javascript/imageviewer.js')
+ head = f'\n'
+ head += f'\n'
+ head += f'\n'
+ head += f'\n'
+ head += f'\n'
+ head += f'\n'
+ head += f'\n'
+ head += f'\n'
+
+ if args_manager.args.theme:
+ head += f'\n'
+
+ return head
+
+
+def css_html():
+ style_css_path = webpath('css/style.css')
+ head = f''
+ return head
+
+
+def reload_javascript():
+ js = javascript_html()
+ css = css_html()
+
+ def template_response(*args, **kwargs):
+ res = GradioTemplateResponseOriginal(*args, **kwargs)
+ res.body = res.body.replace(b'', f'{js}'.encode("utf8"))
+ res.body = res.body.replace(b'{only_name}
\n" + for i, (k, v) in enumerate(dic): + if i < single_line_number: + item += f"{k}: {v}
\n" + else: + if (i - single_line_number) % 2 == 0: + item += f"{k}: {v}, " + else: + item += f"{k}: {v}
\n" + item += f"