+
+
+'''
+
+
+def make_progress_html(number, text):
+ return progress_html.replace('*number*', str(number)).replace('*text*', text)
diff --git a/Fooocus/modules/inpaint_worker.py b/Fooocus/modules/inpaint_worker.py
new file mode 100644
index 0000000000000000000000000000000000000000..88ec39d6df930fc689f591e85877d74e07db50d9
--- /dev/null
+++ b/Fooocus/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/Fooocus/modules/launch_util.py b/Fooocus/modules/launch_util.py
new file mode 100644
index 0000000000000000000000000000000000000000..00fff8aeaffb603bb7da09e9744cc7f2bd61703a
--- /dev/null
+++ b/Fooocus/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/Fooocus/modules/localization.py b/Fooocus/modules/localization.py
new file mode 100644
index 0000000000000000000000000000000000000000..b21d4a564d134ac0be00d83c7005627d601d206e
--- /dev/null
+++ b/Fooocus/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/Fooocus/modules/lora.py b/Fooocus/modules/lora.py
new file mode 100644
index 0000000000000000000000000000000000000000..088545c708943aa8e51e8b2bfd32b2a9968b639f
--- /dev/null
+++ b/Fooocus/modules/lora.py
@@ -0,0 +1,152 @@
+def match_lora(lora, to_load):
+ patch_dict = {}
+ loaded_keys = set()
+ for x in to_load:
+ real_load_key = to_load[x]
+ if real_load_key in lora:
+ patch_dict[real_load_key] = ('fooocus', lora[real_load_key])
+ loaded_keys.add(real_load_key)
+ continue
+
+ alpha_name = "{}.alpha".format(x)
+ alpha = None
+ if alpha_name in lora.keys():
+ alpha = lora[alpha_name].item()
+ loaded_keys.add(alpha_name)
+
+ regular_lora = "{}.lora_up.weight".format(x)
+ diffusers_lora = "{}_lora.up.weight".format(x)
+ transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
+ A_name = None
+
+ if regular_lora in lora.keys():
+ A_name = regular_lora
+ B_name = "{}.lora_down.weight".format(x)
+ mid_name = "{}.lora_mid.weight".format(x)
+ elif diffusers_lora in lora.keys():
+ A_name = diffusers_lora
+ B_name = "{}_lora.down.weight".format(x)
+ mid_name = None
+ elif transformers_lora in lora.keys():
+ A_name = transformers_lora
+ B_name ="{}.lora_linear_layer.down.weight".format(x)
+ mid_name = None
+
+ if A_name is not None:
+ mid = None
+ if mid_name is not None and mid_name in lora.keys():
+ mid = lora[mid_name]
+ loaded_keys.add(mid_name)
+ patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid))
+ loaded_keys.add(A_name)
+ loaded_keys.add(B_name)
+
+
+ ######## loha
+ hada_w1_a_name = "{}.hada_w1_a".format(x)
+ hada_w1_b_name = "{}.hada_w1_b".format(x)
+ hada_w2_a_name = "{}.hada_w2_a".format(x)
+ hada_w2_b_name = "{}.hada_w2_b".format(x)
+ hada_t1_name = "{}.hada_t1".format(x)
+ hada_t2_name = "{}.hada_t2".format(x)
+ if hada_w1_a_name in lora.keys():
+ hada_t1 = None
+ hada_t2 = None
+ if hada_t1_name in lora.keys():
+ hada_t1 = lora[hada_t1_name]
+ hada_t2 = lora[hada_t2_name]
+ loaded_keys.add(hada_t1_name)
+ loaded_keys.add(hada_t2_name)
+
+ patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2))
+ loaded_keys.add(hada_w1_a_name)
+ loaded_keys.add(hada_w1_b_name)
+ loaded_keys.add(hada_w2_a_name)
+ loaded_keys.add(hada_w2_b_name)
+
+
+ ######## lokr
+ lokr_w1_name = "{}.lokr_w1".format(x)
+ lokr_w2_name = "{}.lokr_w2".format(x)
+ lokr_w1_a_name = "{}.lokr_w1_a".format(x)
+ lokr_w1_b_name = "{}.lokr_w1_b".format(x)
+ lokr_t2_name = "{}.lokr_t2".format(x)
+ lokr_w2_a_name = "{}.lokr_w2_a".format(x)
+ lokr_w2_b_name = "{}.lokr_w2_b".format(x)
+
+ lokr_w1 = None
+ if lokr_w1_name in lora.keys():
+ lokr_w1 = lora[lokr_w1_name]
+ loaded_keys.add(lokr_w1_name)
+
+ lokr_w2 = None
+ if lokr_w2_name in lora.keys():
+ lokr_w2 = lora[lokr_w2_name]
+ loaded_keys.add(lokr_w2_name)
+
+ lokr_w1_a = None
+ if lokr_w1_a_name in lora.keys():
+ lokr_w1_a = lora[lokr_w1_a_name]
+ loaded_keys.add(lokr_w1_a_name)
+
+ lokr_w1_b = None
+ if lokr_w1_b_name in lora.keys():
+ lokr_w1_b = lora[lokr_w1_b_name]
+ loaded_keys.add(lokr_w1_b_name)
+
+ lokr_w2_a = None
+ if lokr_w2_a_name in lora.keys():
+ lokr_w2_a = lora[lokr_w2_a_name]
+ loaded_keys.add(lokr_w2_a_name)
+
+ lokr_w2_b = None
+ if lokr_w2_b_name in lora.keys():
+ lokr_w2_b = lora[lokr_w2_b_name]
+ loaded_keys.add(lokr_w2_b_name)
+
+ lokr_t2 = None
+ if lokr_t2_name in lora.keys():
+ lokr_t2 = lora[lokr_t2_name]
+ loaded_keys.add(lokr_t2_name)
+
+ if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
+ patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2))
+
+ #glora
+ a1_name = "{}.a1.weight".format(x)
+ a2_name = "{}.a2.weight".format(x)
+ b1_name = "{}.b1.weight".format(x)
+ b2_name = "{}.b2.weight".format(x)
+ if a1_name in lora:
+ patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha))
+ loaded_keys.add(a1_name)
+ loaded_keys.add(a2_name)
+ loaded_keys.add(b1_name)
+ loaded_keys.add(b2_name)
+
+ w_norm_name = "{}.w_norm".format(x)
+ b_norm_name = "{}.b_norm".format(x)
+ w_norm = lora.get(w_norm_name, None)
+ b_norm = lora.get(b_norm_name, None)
+
+ if w_norm is not None:
+ loaded_keys.add(w_norm_name)
+ patch_dict[to_load[x]] = ("diff", (w_norm,))
+ if b_norm is not None:
+ loaded_keys.add(b_norm_name)
+ patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (b_norm,))
+
+ diff_name = "{}.diff".format(x)
+ diff_weight = lora.get(diff_name, None)
+ if diff_weight is not None:
+ patch_dict[to_load[x]] = ("diff", (diff_weight,))
+ loaded_keys.add(diff_name)
+
+ diff_bias_name = "{}.diff_b".format(x)
+ diff_bias = lora.get(diff_bias_name, None)
+ if diff_bias is not None:
+ patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,))
+ loaded_keys.add(diff_bias_name)
+
+ remaining_dict = {x: y for x, y in lora.items() if x not in loaded_keys}
+ return patch_dict, remaining_dict
diff --git a/Fooocus/modules/meta_parser.py b/Fooocus/modules/meta_parser.py
new file mode 100644
index 0000000000000000000000000000000000000000..07b42a160fc0c6ad5ccc5f3945b225fd11dde9e2
--- /dev/null
+++ b/Fooocus/modules/meta_parser.py
@@ -0,0 +1,148 @@
+import json
+import gradio as gr
+import modules.config
+
+
+def load_parameter_button_click(raw_prompt_txt, is_generating):
+ loaded_parameter_dict = json.loads(raw_prompt_txt)
+ assert isinstance(loaded_parameter_dict, dict)
+
+ results = [True, 1]
+
+ try:
+ h = loaded_parameter_dict.get('Prompt', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Negative Prompt', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Styles', None)
+ h = eval(h)
+ assert isinstance(h, list)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Performance', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Resolution', None)
+ width, height = eval(h)
+ formatted = modules.config.add_ratio(f'{width}*{height}')
+ if formatted in modules.config.available_aspect_ratios:
+ results.append(formatted)
+ results.append(-1)
+ results.append(-1)
+ else:
+ results.append(gr.update())
+ results.append(width)
+ results.append(height)
+ except:
+ results.append(gr.update())
+ results.append(gr.update())
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Sharpness', None)
+ assert h is not None
+ h = float(h)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Guidance Scale', None)
+ assert h is not None
+ h = float(h)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('ADM Guidance', None)
+ p, n, e = eval(h)
+ results.append(float(p))
+ results.append(float(n))
+ results.append(float(e))
+ except:
+ results.append(gr.update())
+ results.append(gr.update())
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Base Model', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Refiner Model', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Refiner Switch', None)
+ assert h is not None
+ h = float(h)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Sampler', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Scheduler', None)
+ assert isinstance(h, str)
+ results.append(h)
+ except:
+ results.append(gr.update())
+
+ try:
+ h = loaded_parameter_dict.get('Seed', None)
+ assert h is not None
+ h = int(h)
+ results.append(False)
+ results.append(h)
+ except:
+ results.append(gr.update())
+ results.append(gr.update())
+
+ if is_generating:
+ results.append(gr.update())
+ else:
+ results.append(gr.update(visible=True))
+
+ results.append(gr.update(visible=False))
+
+ for i in range(1, 6):
+ try:
+ n, w = loaded_parameter_dict.get(f'LoRA {i}').split(' : ')
+ w = float(w)
+ results.append(n)
+ results.append(w)
+ except:
+ results.append(gr.update())
+ results.append(gr.update())
+
+ return results
diff --git a/Fooocus/modules/model_loader.py b/Fooocus/modules/model_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ba336a915ae234b7cd5f9f2576d4edf779738ba
--- /dev/null
+++ b/Fooocus/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/Fooocus/modules/ops.py b/Fooocus/modules/ops.py
new file mode 100644
index 0000000000000000000000000000000000000000..ee0e775634314d1b71811258cff87b2178e1c740
--- /dev/null
+++ b/Fooocus/modules/ops.py
@@ -0,0 +1,19 @@
+import torch
+import contextlib
+
+
+@contextlib.contextmanager
+def use_patched_ops(operations):
+ op_names = ['Linear', 'Conv2d', 'Conv3d', 'GroupNorm', 'LayerNorm']
+ backups = {op_name: getattr(torch.nn, op_name) for op_name in op_names}
+
+ try:
+ for op_name in op_names:
+ setattr(torch.nn, op_name, getattr(operations, op_name))
+
+ yield
+
+ finally:
+ for op_name in op_names:
+ setattr(torch.nn, op_name, backups[op_name])
+ return
diff --git a/Fooocus/modules/patch.py b/Fooocus/modules/patch.py
new file mode 100644
index 0000000000000000000000000000000000000000..2e2409c54dacd3da749afd6e33fb4441c73871c5
--- /dev/null
+++ b/Fooocus/modules/patch.py
@@ -0,0 +1,507 @@
+import os
+import torch
+import time
+import math
+import ldm_patched.modules.model_base
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.model_management
+import modules.anisotropic as anisotropic
+import ldm_patched.ldm.modules.attention
+import ldm_patched.k_diffusion.sampling
+import ldm_patched.modules.sd1_clip
+import modules.inpaint_worker as inpaint_worker
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.ldm.modules.diffusionmodules.model
+import ldm_patched.modules.sd
+import ldm_patched.controlnet.cldm
+import ldm_patched.modules.model_patcher
+import ldm_patched.modules.samplers
+import ldm_patched.modules.args_parser
+import modules.advanced_parameters as advanced_parameters
+import warnings
+import safetensors.torch
+import modules.constants as constants
+
+from ldm_patched.modules.samplers import calc_cond_uncond_batch
+from ldm_patched.k_diffusion.sampling import BatchedBrownianTree
+from ldm_patched.ldm.modules.diffusionmodules.openaimodel import forward_timestep_embed, apply_control
+from modules.patch_precision import patch_all_precision
+from modules.patch_clip import patch_all_clip
+
+
+sharpness = 2.0
+
+adm_scaler_end = 0.3
+positive_adm_scale = 1.5
+negative_adm_scale = 0.8
+
+adaptive_cfg = 7.0
+global_diffusion_progress = 0
+eps_record = None
+
+
+def calculate_weight_patched(self, patches, weight, key):
+ for p in patches:
+ alpha = p[0]
+ v = p[1]
+ strength_model = p[2]
+
+ if strength_model != 1.0:
+ weight *= strength_model
+
+ if isinstance(v, list):
+ v = (self.calculate_weight(v[1:], v[0].clone(), key),)
+
+ if len(v) == 1:
+ patch_type = "diff"
+ elif len(v) == 2:
+ patch_type = v[0]
+ v = v[1]
+
+ if patch_type == "diff":
+ w1 = v[0]
+ if alpha != 0.0:
+ if w1.shape != weight.shape:
+ print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+ else:
+ weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
+ elif patch_type == "lora":
+ mat1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
+ mat2 = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
+ if v[2] is not None:
+ alpha *= v[2] / mat2.shape[0]
+ if v[3] is not None:
+ mat3 = ldm_patched.modules.model_management.cast_to_device(v[3], weight.device, torch.float32)
+ final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
+ mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1),
+ mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
+ try:
+ weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(
+ weight.shape).type(weight.dtype)
+ except Exception as e:
+ print("ERROR", key, e)
+ elif patch_type == "fooocus":
+ w1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
+ w_min = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
+ w_max = ldm_patched.modules.model_management.cast_to_device(v[2], weight.device, torch.float32)
+ w1 = (w1 / 255.0) * (w_max - w_min) + w_min
+ if alpha != 0.0:
+ if w1.shape != weight.shape:
+ print("WARNING SHAPE MISMATCH {} FOOOCUS WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
+ else:
+ weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
+ elif patch_type == "lokr":
+ w1 = v[0]
+ w2 = v[1]
+ w1_a = v[3]
+ w1_b = v[4]
+ w2_a = v[5]
+ w2_b = v[6]
+ t2 = v[7]
+ dim = None
+
+ if w1 is None:
+ dim = w1_b.shape[0]
+ w1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1_a, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w1_b, weight.device, torch.float32))
+ else:
+ w1 = ldm_patched.modules.model_management.cast_to_device(w1, weight.device, torch.float32)
+
+ if w2 is None:
+ dim = w2_b.shape[0]
+ if t2 is None:
+ w2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32))
+ else:
+ w2 = torch.einsum('i j k l, j r, i p -> p r k l',
+ ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32))
+ else:
+ w2 = ldm_patched.modules.model_management.cast_to_device(w2, weight.device, torch.float32)
+
+ if len(w2.shape) == 4:
+ w1 = w1.unsqueeze(2).unsqueeze(2)
+ if v[2] is not None and dim is not None:
+ alpha *= v[2] / dim
+
+ try:
+ weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
+ except Exception as e:
+ print("ERROR", key, e)
+ elif patch_type == "loha":
+ w1a = v[0]
+ w1b = v[1]
+ if v[2] is not None:
+ alpha *= v[2] / w1b.shape[0]
+ w2a = v[3]
+ w2b = v[4]
+ if v[5] is not None: # cp decomposition
+ t1 = v[5]
+ t2 = v[6]
+ m1 = torch.einsum('i j k l, j r, i p -> p r k l',
+ ldm_patched.modules.model_management.cast_to_device(t1, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32))
+
+ m2 = torch.einsum('i j k l, j r, i p -> p r k l',
+ ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32))
+ else:
+ m1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32))
+ m2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32),
+ ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32))
+
+ try:
+ weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
+ except Exception as e:
+ print("ERROR", key, e)
+ elif patch_type == "glora":
+ if v[4] is not None:
+ alpha *= v[4] / v[0].shape[0]
+
+ a1 = ldm_patched.modules.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
+ a2 = ldm_patched.modules.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
+ b1 = ldm_patched.modules.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
+ b2 = ldm_patched.modules.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)
+
+ weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype)
+ else:
+ print("patch type not recognized", patch_type, key)
+
+ return weight
+
+
+class BrownianTreeNoiseSamplerPatched:
+ transform = None
+ tree = None
+
+ @staticmethod
+ def global_init(x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
+ if ldm_patched.modules.model_management.directml_enabled:
+ cpu = True
+
+ t0, t1 = transform(torch.as_tensor(sigma_min)), transform(torch.as_tensor(sigma_max))
+
+ BrownianTreeNoiseSamplerPatched.transform = transform
+ BrownianTreeNoiseSamplerPatched.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
+
+ def __init__(self, *args, **kwargs):
+ pass
+
+ @staticmethod
+ def __call__(sigma, sigma_next):
+ transform = BrownianTreeNoiseSamplerPatched.transform
+ tree = BrownianTreeNoiseSamplerPatched.tree
+
+ t0, t1 = transform(torch.as_tensor(sigma)), transform(torch.as_tensor(sigma_next))
+ return tree(t0, t1) / (t1 - t0).abs().sqrt()
+
+
+def compute_cfg(uncond, cond, cfg_scale, t):
+ global adaptive_cfg
+
+ mimic_cfg = float(adaptive_cfg)
+ real_cfg = float(cfg_scale)
+
+ real_eps = uncond + real_cfg * (cond - uncond)
+
+ if cfg_scale > adaptive_cfg:
+ mimicked_eps = uncond + mimic_cfg * (cond - uncond)
+ return real_eps * t + mimicked_eps * (1 - t)
+ else:
+ return real_eps
+
+
+def patched_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options=None, seed=None):
+ global eps_record
+
+ if math.isclose(cond_scale, 1.0) and not model_options.get("disable_cfg1_optimization", False):
+ final_x0 = calc_cond_uncond_batch(model, cond, None, x, timestep, model_options)[0]
+
+ if eps_record is not None:
+ eps_record = ((x - final_x0) / timestep).cpu()
+
+ return final_x0
+
+ positive_x0, negative_x0 = calc_cond_uncond_batch(model, cond, uncond, x, timestep, model_options)
+
+ positive_eps = x - positive_x0
+ negative_eps = x - negative_x0
+
+ alpha = 0.001 * sharpness * global_diffusion_progress
+
+ positive_eps_degraded = anisotropic.adaptive_anisotropic_filter(x=positive_eps, g=positive_x0)
+ positive_eps_degraded_weighted = positive_eps_degraded * alpha + positive_eps * (1.0 - alpha)
+
+ final_eps = compute_cfg(uncond=negative_eps, cond=positive_eps_degraded_weighted,
+ cfg_scale=cond_scale, t=global_diffusion_progress)
+
+ if eps_record is not None:
+ eps_record = (final_eps / timestep).cpu()
+
+ return x - final_eps
+
+
+def round_to_64(x):
+ h = float(x)
+ h = h / 64.0
+ h = round(h)
+ h = int(h)
+ h = h * 64
+ return h
+
+
+def sdxl_encode_adm_patched(self, **kwargs):
+ global positive_adm_scale, negative_adm_scale
+
+ clip_pooled = ldm_patched.modules.model_base.sdxl_pooled(kwargs, self.noise_augmentor)
+ width = kwargs.get("width", 1024)
+ height = kwargs.get("height", 1024)
+ target_width = width
+ target_height = height
+
+ if kwargs.get("prompt_type", "") == "negative":
+ width = float(width) * negative_adm_scale
+ height = float(height) * negative_adm_scale
+ elif kwargs.get("prompt_type", "") == "positive":
+ width = float(width) * positive_adm_scale
+ height = float(height) * positive_adm_scale
+
+ def embedder(number_list):
+ h = self.embedder(torch.tensor(number_list, dtype=torch.float32))
+ h = torch.flatten(h).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+ return h
+
+ width, height = int(width), int(height)
+ target_width, target_height = round_to_64(target_width), round_to_64(target_height)
+
+ adm_emphasized = embedder([height, width, 0, 0, target_height, target_width])
+ adm_consistent = embedder([target_height, target_width, 0, 0, target_height, target_width])
+
+ clip_pooled = clip_pooled.to(adm_emphasized)
+ final_adm = torch.cat((clip_pooled, adm_emphasized, clip_pooled, adm_consistent), dim=1)
+
+ return final_adm
+
+
+def patched_KSamplerX0Inpaint_forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
+ if inpaint_worker.current_task is not None:
+ latent_processor = self.inner_model.inner_model.process_latent_in
+ inpaint_latent = latent_processor(inpaint_worker.current_task.latent).to(x)
+ inpaint_mask = inpaint_worker.current_task.latent_mask.to(x)
+
+ if getattr(self, 'energy_generator', None) is None:
+ # avoid bad results by using different seeds.
+ self.energy_generator = torch.Generator(device='cpu').manual_seed((seed + 1) % constants.MAX_SEED)
+
+ energy_sigma = sigma.reshape([sigma.shape[0]] + [1] * (len(x.shape) - 1))
+ current_energy = torch.randn(
+ x.size(), dtype=x.dtype, generator=self.energy_generator, device="cpu").to(x) * energy_sigma
+ x = x * inpaint_mask + (inpaint_latent + current_energy) * (1.0 - inpaint_mask)
+
+ out = self.inner_model(x, sigma,
+ cond=cond,
+ uncond=uncond,
+ cond_scale=cond_scale,
+ model_options=model_options,
+ seed=seed)
+
+ out = out * inpaint_mask + inpaint_latent * (1.0 - inpaint_mask)
+ else:
+ out = self.inner_model(x, sigma,
+ cond=cond,
+ uncond=uncond,
+ cond_scale=cond_scale,
+ model_options=model_options,
+ seed=seed)
+ return out
+
+
+def timed_adm(y, timesteps):
+ if isinstance(y, torch.Tensor) and int(y.dim()) == 2 and int(y.shape[1]) == 5632:
+ y_mask = (timesteps > 999.0 * (1.0 - float(adm_scaler_end))).to(y)[..., None]
+ y_with_adm = y[..., :2816].clone()
+ y_without_adm = y[..., 2816:].clone()
+ return y_with_adm * y_mask + y_without_adm * (1.0 - y_mask)
+ return y
+
+
+def patched_cldm_forward(self, x, hint, timesteps, context, y=None, **kwargs):
+ t_emb = ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
+ emb = self.time_embed(t_emb)
+
+ guided_hint = self.input_hint_block(hint, emb, context)
+
+ y = timed_adm(y, timesteps)
+
+ outs = []
+
+ hs = []
+ if self.num_classes is not None:
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ h = x
+ for module, zero_conv in zip(self.input_blocks, self.zero_convs):
+ if guided_hint is not None:
+ h = module(h, emb, context)
+ h += guided_hint
+ guided_hint = None
+ else:
+ h = module(h, emb, context)
+ outs.append(zero_conv(h, emb, context))
+
+ h = self.middle_block(h, emb, context)
+ outs.append(self.middle_block_out(h, emb, context))
+
+ if advanced_parameters.controlnet_softness > 0:
+ for i in range(10):
+ k = 1.0 - float(i) / 9.0
+ outs[i] = outs[i] * (1.0 - advanced_parameters.controlnet_softness * k)
+
+ return outs
+
+
+def patched_unet_forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
+ global global_diffusion_progress
+
+ self.current_step = 1.0 - timesteps.to(x) / 999.0
+ global_diffusion_progress = float(self.current_step.detach().cpu().numpy().tolist()[0])
+
+ y = timed_adm(y, timesteps)
+
+ transformer_options["original_shape"] = list(x.shape)
+ transformer_options["transformer_index"] = 0
+ transformer_patches = transformer_options.get("patches", {})
+
+ num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
+ image_only_indicator = kwargs.get("image_only_indicator", self.default_image_only_indicator)
+ time_context = kwargs.get("time_context", None)
+
+ assert (y is not None) == (
+ self.num_classes is not None
+ ), "must specify y if and only if the model is class-conditional"
+ hs = []
+ t_emb = ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
+ emb = self.time_embed(t_emb)
+
+ if self.num_classes is not None:
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ h = x
+ for id, module in enumerate(self.input_blocks):
+ transformer_options["block"] = ("input", id)
+ h = forward_timestep_embed(module, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
+ h = apply_control(h, control, 'input')
+ if "input_block_patch" in transformer_patches:
+ patch = transformer_patches["input_block_patch"]
+ for p in patch:
+ h = p(h, transformer_options)
+
+ hs.append(h)
+ if "input_block_patch_after_skip" in transformer_patches:
+ patch = transformer_patches["input_block_patch_after_skip"]
+ for p in patch:
+ h = p(h, transformer_options)
+
+ transformer_options["block"] = ("middle", 0)
+ h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
+ h = apply_control(h, control, 'middle')
+
+ for id, module in enumerate(self.output_blocks):
+ transformer_options["block"] = ("output", id)
+ hsp = hs.pop()
+ hsp = apply_control(hsp, control, 'output')
+
+ if "output_block_patch" in transformer_patches:
+ patch = transformer_patches["output_block_patch"]
+ for p in patch:
+ h, hsp = p(h, hsp, transformer_options)
+
+ h = torch.cat([h, hsp], dim=1)
+ del hsp
+ if len(hs) > 0:
+ output_shape = hs[-1].shape
+ else:
+ output_shape = None
+ h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
+ h = h.type(x.dtype)
+ if self.predict_codebook_ids:
+ return self.id_predictor(h)
+ else:
+ return self.out(h)
+
+
+def patched_load_models_gpu(*args, **kwargs):
+ execution_start_time = time.perf_counter()
+ y = ldm_patched.modules.model_management.load_models_gpu_origin(*args, **kwargs)
+ moving_time = time.perf_counter() - execution_start_time
+ if moving_time > 0.1:
+ print(f'[Fooocus Model Management] Moving model(s) has taken {moving_time:.2f} seconds')
+ return y
+
+
+def build_loaded(module, loader_name):
+ original_loader_name = loader_name + '_origin'
+
+ if not hasattr(module, original_loader_name):
+ setattr(module, original_loader_name, getattr(module, loader_name))
+
+ original_loader = getattr(module, original_loader_name)
+
+ def loader(*args, **kwargs):
+ result = None
+ try:
+ result = original_loader(*args, **kwargs)
+ except Exception as e:
+ result = None
+ exp = str(e) + '\n'
+ for path in list(args) + list(kwargs.values()):
+ if isinstance(path, str):
+ if os.path.exists(path):
+ exp += f'File corrupted: {path} \n'
+ corrupted_backup_file = path + '.corrupted'
+ if os.path.exists(corrupted_backup_file):
+ os.remove(corrupted_backup_file)
+ os.replace(path, corrupted_backup_file)
+ if os.path.exists(path):
+ os.remove(path)
+ exp += f'Fooocus has tried to move the corrupted file to {corrupted_backup_file} \n'
+ exp += f'You may try again now and Fooocus will download models again. \n'
+ raise ValueError(exp)
+ return result
+
+ setattr(module, loader_name, loader)
+ return
+
+
+def patch_all():
+ if ldm_patched.modules.model_management.directml_enabled:
+ ldm_patched.modules.model_management.lowvram_available = True
+ ldm_patched.modules.model_management.OOM_EXCEPTION = Exception
+
+ patch_all_precision()
+ patch_all_clip()
+
+ if not hasattr(ldm_patched.modules.model_management, 'load_models_gpu_origin'):
+ ldm_patched.modules.model_management.load_models_gpu_origin = ldm_patched.modules.model_management.load_models_gpu
+
+ ldm_patched.modules.model_management.load_models_gpu = patched_load_models_gpu
+ ldm_patched.modules.model_patcher.ModelPatcher.calculate_weight = calculate_weight_patched
+ ldm_patched.controlnet.cldm.ControlNet.forward = patched_cldm_forward
+ ldm_patched.ldm.modules.diffusionmodules.openaimodel.UNetModel.forward = patched_unet_forward
+ ldm_patched.modules.model_base.SDXL.encode_adm = sdxl_encode_adm_patched
+ ldm_patched.modules.samplers.KSamplerX0Inpaint.forward = patched_KSamplerX0Inpaint_forward
+ ldm_patched.k_diffusion.sampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSamplerPatched
+ ldm_patched.modules.samplers.sampling_function = patched_sampling_function
+
+ warnings.filterwarnings(action='ignore', module='torchsde')
+
+ build_loaded(safetensors.torch, 'load_file')
+ build_loaded(torch, 'load')
+
+ return
diff --git a/Fooocus/modules/patch_clip.py b/Fooocus/modules/patch_clip.py
new file mode 100644
index 0000000000000000000000000000000000000000..06b7f01bb857b01995ff7b0326813f98f92ea76d
--- /dev/null
+++ b/Fooocus/modules/patch_clip.py
@@ -0,0 +1,195 @@
+# Consistent with Kohya/A1111 to reduce differences between model training and inference.
+
+import os
+import torch
+import ldm_patched.controlnet.cldm
+import ldm_patched.k_diffusion.sampling
+import ldm_patched.ldm.modules.attention
+import ldm_patched.ldm.modules.diffusionmodules.model
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.args_parser
+import ldm_patched.modules.model_base
+import ldm_patched.modules.model_management
+import ldm_patched.modules.model_patcher
+import ldm_patched.modules.samplers
+import ldm_patched.modules.sd
+import ldm_patched.modules.sd1_clip
+import ldm_patched.modules.clip_vision
+import ldm_patched.modules.ops as ops
+
+from modules.ops import use_patched_ops
+from transformers import CLIPTextModel, CLIPTextConfig, modeling_utils, CLIPVisionConfig, CLIPVisionModelWithProjection
+
+
+def patched_encode_token_weights(self, token_weight_pairs):
+ to_encode = list()
+ max_token_len = 0
+ has_weights = False
+ for x in token_weight_pairs:
+ tokens = list(map(lambda a: a[0], x))
+ max_token_len = max(len(tokens), max_token_len)
+ has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
+ to_encode.append(tokens)
+
+ sections = len(to_encode)
+ if has_weights or sections == 0:
+ to_encode.append(ldm_patched.modules.sd1_clip.gen_empty_tokens(self.special_tokens, max_token_len))
+
+ out, pooled = self.encode(to_encode)
+ if pooled is not None:
+ first_pooled = pooled[0:1].to(ldm_patched.modules.model_management.intermediate_device())
+ else:
+ first_pooled = pooled
+
+ output = []
+ for k in range(0, sections):
+ z = out[k:k + 1]
+ if has_weights:
+ original_mean = z.mean()
+ z_empty = out[-1]
+ for i in range(len(z)):
+ for j in range(len(z[i])):
+ weight = token_weight_pairs[k][j][1]
+ if weight != 1.0:
+ z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
+ new_mean = z.mean()
+ z = z * (original_mean / new_mean)
+ output.append(z)
+
+ if len(output) == 0:
+ return out[-1:].to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
+ return torch.cat(output, dim=-2).to(ldm_patched.modules.model_management.intermediate_device()), first_pooled
+
+
+def patched_SDClipModel__init__(self, max_length=77, freeze=True, layer="last", layer_idx=None,
+ textmodel_json_config=None, dtype=None, special_tokens=None,
+ layer_norm_hidden_state=True, **kwargs):
+ torch.nn.Module.__init__(self)
+ assert layer in self.LAYERS
+
+ if special_tokens is None:
+ special_tokens = {"start": 49406, "end": 49407, "pad": 49407}
+
+ if textmodel_json_config is None:
+ textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(ldm_patched.modules.sd1_clip.__file__)),
+ "sd1_clip_config.json")
+
+ config = CLIPTextConfig.from_json_file(textmodel_json_config)
+ self.num_layers = config.num_hidden_layers
+
+ with use_patched_ops(ops.manual_cast):
+ with modeling_utils.no_init_weights():
+ self.transformer = CLIPTextModel(config)
+
+ if dtype is not None:
+ self.transformer.to(dtype)
+
+ self.transformer.text_model.embeddings.to(torch.float32)
+
+ if freeze:
+ self.freeze()
+
+ self.max_length = max_length
+ self.layer = layer
+ self.layer_idx = None
+ self.special_tokens = special_tokens
+ self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
+ self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
+ self.enable_attention_masks = False
+
+ self.layer_norm_hidden_state = layer_norm_hidden_state
+ if layer == "hidden":
+ assert layer_idx is not None
+ assert abs(layer_idx) < self.num_layers
+ self.clip_layer(layer_idx)
+ self.layer_default = (self.layer, self.layer_idx)
+
+
+def patched_SDClipModel_forward(self, tokens):
+ backup_embeds = self.transformer.get_input_embeddings()
+ device = backup_embeds.weight.device
+ tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
+ tokens = torch.LongTensor(tokens).to(device)
+
+ attention_mask = None
+ if self.enable_attention_masks:
+ attention_mask = torch.zeros_like(tokens)
+ max_token = self.transformer.get_input_embeddings().weight.shape[0] - 1
+ for x in range(attention_mask.shape[0]):
+ for y in range(attention_mask.shape[1]):
+ attention_mask[x, y] = 1
+ if tokens[x, y] == max_token:
+ break
+
+ outputs = self.transformer(input_ids=tokens, attention_mask=attention_mask,
+ output_hidden_states=self.layer == "hidden")
+ self.transformer.set_input_embeddings(backup_embeds)
+
+ if self.layer == "last":
+ z = outputs.last_hidden_state
+ elif self.layer == "pooled":
+ z = outputs.pooler_output[:, None, :]
+ else:
+ z = outputs.hidden_states[self.layer_idx]
+ if self.layer_norm_hidden_state:
+ z = self.transformer.text_model.final_layer_norm(z)
+
+ if hasattr(outputs, "pooler_output"):
+ pooled_output = outputs.pooler_output.float()
+ else:
+ pooled_output = None
+
+ if self.text_projection is not None and pooled_output is not None:
+ pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
+
+ return z.float(), pooled_output
+
+
+def patched_ClipVisionModel__init__(self, json_config):
+ config = CLIPVisionConfig.from_json_file(json_config)
+
+ self.load_device = ldm_patched.modules.model_management.text_encoder_device()
+ self.offload_device = ldm_patched.modules.model_management.text_encoder_offload_device()
+
+ if ldm_patched.modules.model_management.should_use_fp16(self.load_device, prioritize_performance=False):
+ self.dtype = torch.float16
+ else:
+ self.dtype = torch.float32
+
+ with use_patched_ops(ops.manual_cast):
+ with modeling_utils.no_init_weights():
+ self.model = CLIPVisionModelWithProjection(config)
+
+ self.model.to(self.dtype)
+ self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(
+ self.model,
+ load_device=self.load_device,
+ offload_device=self.offload_device
+ )
+
+
+def patched_ClipVisionModel_encode_image(self, image):
+ ldm_patched.modules.model_management.load_model_gpu(self.patcher)
+ pixel_values = ldm_patched.modules.clip_vision.clip_preprocess(image.to(self.load_device))
+ outputs = self.model(pixel_values=pixel_values, output_hidden_states=True)
+
+ for k in outputs:
+ t = outputs[k]
+ if t is not None:
+ if k == 'hidden_states':
+ outputs["penultimate_hidden_states"] = t[-2].to(ldm_patched.modules.model_management.intermediate_device())
+ outputs["hidden_states"] = None
+ else:
+ outputs[k] = t.to(ldm_patched.modules.model_management.intermediate_device())
+
+ return outputs
+
+
+def patch_all_clip():
+ ldm_patched.modules.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = patched_encode_token_weights
+ ldm_patched.modules.sd1_clip.SDClipModel.__init__ = patched_SDClipModel__init__
+ ldm_patched.modules.sd1_clip.SDClipModel.forward = patched_SDClipModel_forward
+ ldm_patched.modules.clip_vision.ClipVisionModel.__init__ = patched_ClipVisionModel__init__
+ ldm_patched.modules.clip_vision.ClipVisionModel.encode_image = patched_ClipVisionModel_encode_image
+ return
diff --git a/Fooocus/modules/patch_precision.py b/Fooocus/modules/patch_precision.py
new file mode 100644
index 0000000000000000000000000000000000000000..83569bdd15f5ab0cac2c57353626c4e843bd264d
--- /dev/null
+++ b/Fooocus/modules/patch_precision.py
@@ -0,0 +1,60 @@
+# Consistent with Kohya to reduce differences between model training and inference.
+
+import torch
+import math
+import einops
+import numpy as np
+
+import ldm_patched.ldm.modules.diffusionmodules.openaimodel
+import ldm_patched.modules.model_sampling
+import ldm_patched.modules.sd1_clip
+
+from ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
+
+
+def patched_timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+ # Consistent with Kohya to reduce differences between model training and inference.
+
+ if not repeat_only:
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=timesteps.device)
+ args = timesteps[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ else:
+ embedding = einops.repeat(timesteps, 'b -> b d', d=dim)
+ return embedding
+
+
+def patched_register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+ # Consistent with Kohya to reduce differences between model training and inference.
+
+ if given_betas is not None:
+ betas = given_betas
+ else:
+ betas = make_beta_schedule(
+ beta_schedule,
+ timesteps,
+ linear_start=linear_start,
+ linear_end=linear_end,
+ cosine_s=cosine_s)
+
+ alphas = 1. - betas
+ alphas_cumprod = np.cumprod(alphas, axis=0)
+ timesteps, = betas.shape
+ self.num_timesteps = int(timesteps)
+ self.linear_start = linear_start
+ self.linear_end = linear_end
+ sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
+ self.set_sigmas(sigmas)
+ return
+
+
+def patch_all_precision():
+ ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding = patched_timestep_embedding
+ ldm_patched.modules.model_sampling.ModelSamplingDiscrete._register_schedule = patched_register_schedule
+ return
diff --git a/Fooocus/modules/private_logger.py b/Fooocus/modules/private_logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..968bd4f5d0d91697a4fd1f6e5b55e99eaae871de
--- /dev/null
+++ b/Fooocus/modules/private_logger.py
@@ -0,0 +1,108 @@
+import os
+import args_manager
+import modules.config
+import json
+import urllib.parse
+
+from PIL import Image
+from modules.util import generate_temp_filename
+
+
+log_cache = {}
+
+
+def get_current_html_path():
+ date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs,
+ extension='png')
+ html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+ return html_name
+
+
+def log(img, dic):
+ if args_manager.args.disable_image_log:
+ return
+
+ date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs, extension='png')
+ os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
+ Image.fromarray(img).save(local_temp_filename)
+ html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
+
+ css_styles = (
+ ""
+ )
+
+ js = (
+ """"""
+ )
+
+ begin_part = f"
+
+
+ *text*
+Fooocus Log {date_string} (private)
\nAll images are clean, without any hidden data/meta, and safe to share with others.
\n\n" + end_part = f'\n' + + middle_part = log_cache.get(html_name, "") + + if middle_part == "": + if os.path.exists(html_name): + existing_split = open(html_name, 'r', encoding='utf-8').read().split('') + if len(existing_split) == 3: + middle_part = existing_split[1] + else: + middle_part = existing_split[0] + + div_name = only_name.replace('.', '_') + item = f"{only_name} | "
+ item += "
| "
+ item += "