2.1.839

modules/async_worker.py

@@ -37,7 +37,7 @@ def worker():
 
     from modules.sdxl_styles import apply_style, apply_wildcards, fooocus_expansion
     from modules.private_logger import log
-    from modules.expansion import safe_str
+    from extras.expansion import safe_str
     from modules.util import remove_empty_str, HWC3, resize_image, \
         get_image_shape_ceil, set_image_shape_ceil, get_shape_ceil, resample_image
     from modules.upscaler import perform_upscale

modules/core.py

@@ -23,7 +23,6 @@ from ldm_patched.contrib.external import VAEDecode, EmptyLatentImage, VAEEncode,
     ControlNetApplyAdvanced
 from ldm_patched.contrib.external_freelunch import FreeU_V2
 from ldm_patched.modules.sample import prepare_mask
-from modules.patch import patched_sampler_cfg_function
 from modules.lora import match_lora
 from ldm_patched.modules.lora import model_lora_keys_unet, model_lora_keys_clip
 from modules.config import path_embeddings
@@ -150,7 +149,6 @@ def apply_controlnet(positive, negative, control_net, image, strength, start_per
 @torch.inference_mode()
 def load_model(ckpt_filename):
     unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename, embedding_directory=path_embeddings)
-    unet.model_options['sampler_cfg_function'] = patched_sampler_cfg_function
     return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, filename=ckpt_filename)
 
 

modules/default_pipeline.py

@@ -7,9 +7,9 @@ import ldm_patched.modules.model_management
 import ldm_patched.modules.latent_formats
 import modules.inpaint_worker
 import extras.vae_interpose as vae_interpose
+from extras.expansion import FooocusExpansion
 
 from ldm_patched.modules.model_base import SDXL, SDXLRefiner
-from modules.expansion import FooocusExpansion
 from modules.sample_hijack import clip_separate
 
 

modules/expansion.py

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

modules/patch.py

@@ -1,11 +1,9 @@
 import os
 import torch
 import time
-import numpy as np
 import math
 import ldm_patched.modules.model_base
 import ldm_patched.ldm.modules.diffusionmodules.openaimodel
-import ldm_patched.modules.samplers
 import ldm_patched.modules.model_management
 import modules.anisotropic as anisotropic
 import ldm_patched.ldm.modules.attention
@@ -24,10 +22,9 @@ import warnings
 import safetensors.torch
 import modules.constants as constants
 
-from einops import repeat
+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 ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
 
 
 sharpness = 2.0
@@ -178,8 +175,6 @@ def calculate_weight_patched(self, patches, weight, key):
 class BrownianTreeNoiseSamplerPatched:
     transform = None
     tree = None
-    global_sigma_min = 1.0
-    global_sigma_max = 1.0
 
     @staticmethod
     def global_init(x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
@@ -191,9 +186,6 @@ class BrownianTreeNoiseSamplerPatched:
         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
 
@@ -221,34 +213,47 @@ def compute_cfg(uncond, cond, cfg_scale, t):
         return real_eps
 
 
-def patched_sampler_cfg_function(args):
+def patched_sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options=None, seed=None):
+    if math.isclose(cond_scale, 1.0):
+        return calc_cond_uncond_batch(model, cond, None, x, timestep, model_options)[0]
+
     global eps_record
 
-    positive_eps = args['cond']
-    negative_eps = args['uncond']
-    cfg_scale = args['cond_scale']
-    positive_x0 = args['input'] - positive_eps
-    sigma = args['sigma']
+    positive_x0, negative_x0 = calc_cond_uncond_batch(model, cond, uncond, x, timestep, model_options)
+
+    positive_eps = x - positive_x0
+    negative_eps = x - negative_x0
+    sigma = timestep
 
     alpha = 0.001 * sharpness * global_diffusion_progress
+
     positive_eps_degraded = anisotropic.adaptive_anisotropic_filter(x=positive_eps, g=positive_x0)
     positive_eps_degraded_weighted = positive_eps_degraded * alpha + positive_eps * (1.0 - alpha)
 
     final_eps = compute_cfg(uncond=negative_eps, cond=positive_eps_degraded_weighted,
-                            cfg_scale=cfg_scale, t=global_diffusion_progress)
+                            cfg_scale=cond_scale, t=global_diffusion_progress)
 
     if eps_record is not None:
         eps_record = (final_eps / sigma).cpu()
 
-    return final_eps
+    return x - final_eps
+
+
+def round_to_64(x):
+    h = float(x)
+    h = h / 64.0
+    h = round(h)
+    h = int(h)
+    h = h * 64
+    return h
 
 
 def sdxl_encode_adm_patched(self, **kwargs):
     global positive_adm_scale, negative_adm_scale
 
     clip_pooled = ldm_patched.modules.model_base.sdxl_pooled(kwargs, self.noise_augmentor)
-    width = kwargs.get("width", 768)
-    height = kwargs.get("height", 768)
+    width = kwargs.get("width", 1024)
+    height = kwargs.get("height", 1024)
     target_width = width
     target_height = height
 
@@ -259,25 +264,21 @@ def sdxl_encode_adm_patched(self, **kwargs):
         width = float(width) * positive_adm_scale
         height = float(height) * positive_adm_scale
 
-    # Avoid artifacts
-    width = int(width)
-    height = int(height)
-    crop_w = 0
-    crop_h = 0
-    target_width = int(target_width)
-    target_height = int(target_height)
+    def embedder(number_list):
+        h = [self.embedder(torch.Tensor([number])) for number in number_list]
+        y = torch.flatten(torch.cat(h)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+        return y
 
-    out_a = [self.embedder(torch.Tensor([height])), self.embedder(torch.Tensor([width])),
-             self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
-             self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
-    flat_a = torch.flatten(torch.cat(out_a)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+    width, height = round_to_64(width), round_to_64(height)
+    target_width, target_height = round_to_64(target_width), round_to_64(target_height)
 
-    out_b = [self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width])),
-             self.embedder(torch.Tensor([crop_h])), self.embedder(torch.Tensor([crop_w])),
-             self.embedder(torch.Tensor([target_height])), self.embedder(torch.Tensor([target_width]))]
-    flat_b = torch.flatten(torch.cat(out_b)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
+    adm_emphasized = embedder([height, width, 0, 0, target_height, target_width])
+    adm_consistent = embedder([target_height, target_width, 0, 0, target_height, target_width])
 
-    return torch.cat((clip_pooled.to(flat_a.device), flat_a, clip_pooled.to(flat_b.device), flat_b), dim=1)
+    clip_pooled = clip_pooled.to(adm_emphasized)
+    final_adm = torch.cat((clip_pooled, adm_emphasized, clip_pooled, adm_consistent), dim=1)
+
+    return final_adm
 
 
 def encode_token_weights_patched_with_a1111_method(self, token_weight_pairs):
@@ -512,48 +513,6 @@ def build_loaded(module, loader_name):
     return
 
 
-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_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():
     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
@@ -566,10 +525,7 @@ def patch_all():
     ldm_patched.modules.sd1_clip.ClipTokenWeightEncoder.encode_token_weights = encode_token_weights_patched_with_a1111_method
     ldm_patched.modules.samplers.KSamplerX0Inpaint.forward = patched_KSamplerX0Inpaint_forward
     ldm_patched.k_diffusion.sampling.BrownianTreeNoiseSampler = BrownianTreeNoiseSamplerPatched
-
-    # Precision fix
-    ldm_patched.ldm.modules.diffusionmodules.openaimodel.timestep_embedding = patched_timestep_embedding
-    ldm_patched.modules.model_base.ModelSamplingDiscrete._register_schedule = patched_register_schedule
+    ldm_patched.modules.samplers.sampling_function = patched_sampling_function
 
     warnings.filterwarnings(action='ignore', module='torchsde')