diff --git a/adv_encode.py b/adv_encode.py new file mode 100644 index 0000000..9e0cbb1 --- /dev/null +++ b/adv_encode.py @@ -0,0 +1,290 @@ +import torch +import numpy as np +import itertools +from math import gcd + +from comfy import model_management +from comfy.sdxl_clip import SDXLClipModel, SDXLRefinerClipModel, SDXLClipG + +def _grouper(n, iterable): + it = iter(iterable) + while True: + chunk = list(itertools.islice(it, n)) + if not chunk: + return + yield chunk + +def _norm_mag(w, n): + d = w - 1 + return 1 + np.sign(d) * np.sqrt(np.abs(d)**2 / n) + #return np.sign(w) * np.sqrt(np.abs(w)**2 / n) + +def divide_length(word_ids, weights): + sums = dict(zip(*np.unique(word_ids, return_counts=True))) + sums[0] = 1 + weights = [[_norm_mag(w, sums[id]) if id != 0 else 1.0 + for w, id in zip(x, y)] for x, y in zip(weights, word_ids)] + return weights + +def shift_mean_weight(word_ids, weights): + delta = 1 - np.mean([w for x, y in zip(weights, word_ids) for w, id in zip(x,y) if id != 0]) + weights = [[w if id == 0 else w+delta + for w, id in zip(x, y)] for x, y in zip(weights, word_ids)] + return weights + +def scale_to_norm(weights, word_ids, w_max): + top = np.max(weights) + w_max = min(top, w_max) + weights = [[w_max if id == 0 else (w/top) * w_max + for w, id in zip(x, y)] for x, y in zip(weights, word_ids)] + return weights + +def from_zero(weights, base_emb): + weight_tensor = torch.tensor(weights, dtype=base_emb.dtype, device=base_emb.device) + weight_tensor = weight_tensor.reshape(1,-1,1).expand(base_emb.shape) + return base_emb * weight_tensor + +def mask_word_id(tokens, word_ids, target_id, mask_token): + new_tokens = [[mask_token if wid == target_id else t + for t, wid in zip(x,y)] for x,y in zip(tokens, word_ids)] + mask = np.array(word_ids) == target_id + return (new_tokens, mask) + +def batched_clip_encode(tokens, length, encode_func, num_chunks): + embs = [] + for e in _grouper(32, tokens): + enc, pooled = encode_func(e) + enc = enc.reshape((len(e), length, -1)) + embs.append(enc) + embs = torch.cat(embs) + embs = embs.reshape((len(tokens) // num_chunks, length * num_chunks, -1)) + return embs + +def from_masked(tokens, weights, word_ids, base_emb, length, encode_func, m_token=266): + pooled_base = base_emb[0,length-1:length,:] + wids, inds = np.unique(np.array(word_ids).reshape(-1), return_index=True) + weight_dict = dict((id,w) + for id,w in zip(wids ,np.array(weights).reshape(-1)[inds]) + if w != 1.0) + + if len(weight_dict) == 0: + return torch.zeros_like(base_emb), base_emb[0,length-1:length,:] + + weight_tensor = torch.tensor(weights, dtype=base_emb.dtype, device=base_emb.device) + weight_tensor = weight_tensor.reshape(1,-1,1).expand(base_emb.shape) + + #m_token = (clip.tokenizer.end_token, 1.0) if clip.tokenizer.pad_with_end else (0,1.0) + #TODO: find most suitable masking token here + m_token = (m_token, 1.0) + + ws = [] + masked_tokens = [] + masks = [] + + #create prompts + for id, w in weight_dict.items(): + masked, m = mask_word_id(tokens, word_ids, id, m_token) + masked_tokens.extend(masked) + + m = torch.tensor(m, dtype=base_emb.dtype, device=base_emb.device) + m = m.reshape(1,-1,1).expand(base_emb.shape) + masks.append(m) + + ws.append(w) + + #batch process prompts + embs = batched_clip_encode(masked_tokens, length, encode_func, len(tokens)) + masks = torch.cat(masks) + + embs = (base_emb.expand(embs.shape) - embs) + pooled = embs[0,length-1:length,:] + + embs *= masks + embs = embs.sum(axis=0, keepdim=True) + + pooled_start = pooled_base.expand(len(ws), -1) + ws = torch.tensor(ws).reshape(-1,1).expand(pooled_start.shape) + pooled = (pooled - pooled_start) * (ws - 1) + pooled = pooled.mean(axis=0, keepdim=True) + + return ((weight_tensor - 1) * embs), pooled_base + pooled + +def mask_inds(tokens, inds, mask_token): + clip_len = len(tokens[0]) + inds_set = set(inds) + new_tokens = [[mask_token if i*clip_len + j in inds_set else t + for j, t in enumerate(x)] for i, x in enumerate(tokens)] + return new_tokens + +def down_weight(tokens, weights, word_ids, base_emb, length, encode_func, m_token=266): + w, w_inv = np.unique(weights,return_inverse=True) + + if np.sum(w < 1) == 0: + return base_emb, tokens, base_emb[0,length-1:length,:] + #m_token = (clip.tokenizer.end_token, 1.0) if clip.tokenizer.pad_with_end else (0,1.0) + #using the comma token as a masking token seems to work better than aos tokens for SD 1.x + m_token = (m_token, 1.0) + + masked_tokens = [] + + masked_current = tokens + for i in range(len(w)): + if w[i] >= 1: + continue + masked_current = mask_inds(masked_current, np.where(w_inv == i)[0], m_token) + masked_tokens.extend(masked_current) + + embs = batched_clip_encode(masked_tokens, length, encode_func, len(tokens)) + embs = torch.cat([base_emb, embs]) + w = w[w<=1.0] + w_mix = np.diff([0] + w.tolist()) + w_mix = torch.tensor(w_mix, dtype=embs.dtype, device=embs.device).reshape((-1,1,1)) + + weighted_emb = (w_mix * embs).sum(axis=0, keepdim=True) + return weighted_emb, masked_current, weighted_emb[0,length-1:length,:] + +def scale_emb_to_mag(base_emb, weighted_emb): + norm_base = torch.linalg.norm(base_emb) + norm_weighted = torch.linalg.norm(weighted_emb) + embeddings_final = (norm_base / norm_weighted) * weighted_emb + return embeddings_final + +def recover_dist(base_emb, weighted_emb): + fixed_std = (base_emb.std() / weighted_emb.std()) * (weighted_emb - weighted_emb.mean()) + embeddings_final = fixed_std + (base_emb.mean() - fixed_std.mean()) + return embeddings_final + +def A1111_renorm(base_emb, weighted_emb): + embeddings_final = (base_emb.mean() / weighted_emb.mean()) * weighted_emb + return embeddings_final + +def advanced_encode_from_tokens(tokenized, token_normalization, weight_interpretation, encode_func, m_token=266, length=77, w_max=1.0, return_pooled=False, apply_to_pooled=False): + tokens = [[t for t,_,_ in x] for x in tokenized] + weights = [[w for _,w,_ in x] for x in tokenized] + word_ids = [[wid for _,_,wid in x] for x in tokenized] + + #weight normalization + #==================== + + #distribute down/up weights over word lengths + if token_normalization.startswith("length"): + weights = divide_length(word_ids, weights) + + #make mean of word tokens 1 + if token_normalization.endswith("mean"): + weights = shift_mean_weight(word_ids, weights) + + #weight interpretation + #===================== + pooled = None + + if weight_interpretation == "comfy": + weighted_tokens = [[(t,w) for t, w in zip(x, y)] for x, y in zip(tokens, weights)] + weighted_emb, pooled_base = encode_func(weighted_tokens) + pooled = pooled_base + else: + unweighted_tokens = [[(t,1.0) for t, _,_ in x] for x in tokenized] + base_emb, pooled_base = encode_func(unweighted_tokens) + + if weight_interpretation == "A1111": + weighted_emb = from_zero(weights, base_emb) + weighted_emb = A1111_renorm(base_emb, weighted_emb) + pooled = pooled_base + + if weight_interpretation == "compel": + pos_tokens = [[(t,w) if w >= 1.0 else (t,1.0) for t, w in zip(x, y)] for x, y in zip(tokens, weights)] + weighted_emb, _ = encode_func(pos_tokens) + weighted_emb, _, pooled = down_weight(pos_tokens, weights, word_ids, weighted_emb, length, encode_func) + + if weight_interpretation == "comfy++": + weighted_emb, tokens_down, _ = down_weight(unweighted_tokens, weights, word_ids, base_emb, length, encode_func) + weights = [[w if w > 1.0 else 1.0 for w in x] for x in weights] + #unweighted_tokens = [[(t,1.0) for t, _,_ in x] for x in tokens_down] + embs, pooled = from_masked(unweighted_tokens, weights, word_ids, base_emb, length, encode_func) + weighted_emb += embs + + if weight_interpretation == "down_weight": + weights = scale_to_norm(weights, word_ids, w_max) + weighted_emb, _, pooled = down_weight(unweighted_tokens, weights, word_ids, base_emb, length, encode_func) + + if return_pooled: + if apply_to_pooled: + return weighted_emb, pooled + else: + return weighted_emb, pooled_base + return weighted_emb, None + +def encode_token_weights_g(model, token_weight_pairs): + return model.clip_g.encode_token_weights(token_weight_pairs) + +def encode_token_weights_l(model, token_weight_pairs): + l_out, _ = model.clip_l.encode_token_weights(token_weight_pairs) + return l_out, None + +def encode_token_weights(model, token_weight_pairs, encode_func): + if model.layer_idx is not None: + model.cond_stage_model.clip_layer(model.layer_idx) + + model_management.load_model_gpu(model.patcher) + return encode_func(model.cond_stage_model, token_weight_pairs) + +def prepareXL(embs_l, embs_g, pooled, clip_balance): + l_w = 1 - max(0, clip_balance - .5) * 2 + g_w = 1 - max(0, .5 - clip_balance) * 2 + if embs_l is not None: + return torch.cat([embs_l * l_w, embs_g * g_w], dim=-1), pooled + else: + return embs_g, pooled + +def advanced_encode(clip, text, token_normalization, weight_interpretation, w_max=1.0, clip_balance=.5, apply_to_pooled=True): + tokenized = clip.tokenize(text, return_word_ids=True) + if isinstance(clip.cond_stage_model, (SDXLClipModel, SDXLRefinerClipModel, SDXLClipG)): + embs_l = None + embs_g = None + pooled = None + if 'l' in tokenized and isinstance(clip.cond_stage_model, SDXLClipModel): + embs_l, _ = advanced_encode_from_tokens(tokenized['l'], + token_normalization, + weight_interpretation, + lambda x: encode_token_weights(clip, x, encode_token_weights_l), + w_max=w_max, + return_pooled=False) + if 'g' in tokenized: + embs_g, pooled = advanced_encode_from_tokens(tokenized['g'], + token_normalization, + weight_interpretation, + lambda x: encode_token_weights(clip, x, encode_token_weights_g), + w_max=w_max, + return_pooled=True, + apply_to_pooled=apply_to_pooled) + return prepareXL(embs_l, embs_g, pooled, clip_balance) + else: + return advanced_encode_from_tokens(tokenized['l'], + token_normalization, + weight_interpretation, + lambda x: (clip.encode_from_tokens({'l': x}), None), + w_max=w_max) +def advanced_encode_XL(clip, text1, text2, token_normalization, weight_interpretation, w_max=1.0, clip_balance=.5, apply_to_pooled=True): + tokenized1 = clip.tokenize(text1, return_word_ids=True) + tokenized2 = clip.tokenize(text2, return_word_ids=True) + + embs_l, _ = advanced_encode_from_tokens(tokenized1['l'], + token_normalization, + weight_interpretation, + lambda x: encode_token_weights(clip, x, encode_token_weights_l), + w_max=w_max, + return_pooled=False) + + embs_g, pooled = advanced_encode_from_tokens(tokenized2['g'], + token_normalization, + weight_interpretation, + lambda x: encode_token_weights(clip, x, encode_token_weights_g), + w_max=w_max, + return_pooled=True, + apply_to_pooled=apply_to_pooled) + + gcd_num = gcd(embs_l.shape[1], embs_g.shape[1]) + repeat_l = int((embs_g.shape[1] / gcd_num) * embs_l.shape[1]) + repeat_g = int((embs_l.shape[1] / gcd_num) * embs_g.shape[1]) + + return prepareXL(embs_l.expand((-1,repeat_l,-1)), embs_g.expand((-1,repeat_g,-1)), pooled, clip_balance) diff --git a/advancedksampler.py b/advancedksampler.py index 00b5fcb..4aadd8b 100644 --- a/advancedksampler.py +++ b/advancedksampler.py @@ -1,9 +1,12 @@ + from torch import Tensor from PIL import Image, ImageOps, ImageDraw, ImageFont from PIL.PngImagePlugin import PngInfo import numpy as np import torch +import folder_paths + import ast from pathlib import Path from importlib import import_module @@ -14,6 +17,10 @@ import subprocess import json import psutil +import re +import nodes as native +from .adv_encode import advanced_encode #, advanced_encode_XL + from comfy_extras.nodes_align_your_steps import AlignYourStepsScheduler from comfy_extras.nodes_gits import GITSScheduler import comfy_extras.nodes_model_advanced @@ -45,17 +52,136 @@ from comfy import samplers sys.path.append(custom_nodes_dir) + +class LoraTagLoader: + def __init__(self): + self.loaded_lora = None + self.tag_pattern = "\<[0-9a-zA-Z\:\_\-\.\s\/\(\)\\\\]+\>" + + @classmethod + def INPUT_TYPES(s): + return {"required": { "model": ("MODEL",), + "clip": ("CLIP", ), + "text": ("STRING", {"multiline": True}), + }} + RETURN_TYPES = ("MODEL", "CLIP", "STRING") + RETURN_NAMES = ("MODEL", "CLIP", "STRING") + FUNCTION = "load_lora" + + CATEGORY = "loaders" + + def load_lora(self, model, clip, text): + # print(f"\nLoraTagLoader input text: { text }") + + founds = re.findall(self.tag_pattern, text) + # print(f"\nfoound lora tags: { founds }") + + if len(founds) < 1: + return (model, clip, text) + + model_lora = model + clip_lora = clip + + lora_files = folder_paths.get_filename_list("loras") + for f in founds: + tag = f[1:-1] + pak = tag.split(":") + type = pak[0] + if type != 'lora': + continue + name = None + if len(pak) > 1 and len(pak[1]) > 0: + name = pak[1] + else: + continue + wModel = wClip = 0 + try: + if len(pak) > 2 and len(pak[2]) > 0: + wModel = float(pak[2]) + wClip = wModel + if len(pak) > 3 and len(pak[3]) > 0: + wClip = float(pak[3]) + except ValueError: + continue + if name == None: + continue + lora_name = None + for lora_file in lora_files: + if Path(lora_file).name.startswith(name) or lora_file.startswith(name): + lora_name = lora_file + break + if lora_name == None: + print(f"bypassed lora tag: { (type, name, wModel, wClip) } >> { lora_name }") + continue + print(f"detected lora tag: { (type, name, wModel, wClip) } >> { lora_name }") + + lora_path = folder_paths.get_full_path("loras", lora_name) + lora = None + if self.loaded_lora is not None: + if self.loaded_lora[0] == lora_path: + lora = self.loaded_lora[1] + else: + temp = self.loaded_lora + self.loaded_lora = None + del temp + + if lora is None: + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + self.loaded_lora = (lora_path, lora) + + model_lora, clip_lora = comfy.sd.load_lora_for_models(model_lora, clip_lora, lora, wModel, wClip) + + plain_prompt = re.sub(self.tag_pattern, "", text) + return (model_lora, clip_lora, plain_prompt) + + +class AdvancedCLIPTextEncodeWithBreak: + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "text": ("STRING", {"multiline": True}), + "clip": ("CLIP", ), + "token_normalization": (["none", "mean", "length", "length+mean"],), + "weight_interpretation": (["comfy", "A1111", "compel", "comfy++" ,"down_weight"],), + }} + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "encode" + + CATEGORY = "conditioning/advanced" + + def _encode(self, clip, text, token_normalization, weight_interpretation): + embeddings_final, pooled = advanced_encode(clip, text, token_normalization, weight_interpretation, w_max=1.0) + return ([[embeddings_final, {"pooled_output": pooled}]], ) + + def encode(self, clip, text, token_normalization, weight_interpretation): + prompts = re.split(r"\s*\bBREAK\b\s*", text) + # encode first prompt fragment + prompt = prompts.pop(0) + # print(f"prompt: {prompt}") + out = self._encode(clip, prompt, token_normalization, weight_interpretation) + # encode and concatenate the rest of the prompt + for prompt in prompts: + # print(f"prompt: {prompt}") + cond_to = self._encode(clip, prompt, token_normalization, weight_interpretation) + out = native.ConditioningConcat.concat(self, cond_to[0], out[0]) + return out + + + class TwoModelAdvancedKsampler: @classmethod def INPUT_TYPES(s): return { "required": { "model1": ("MODEL",), - "positive1": ("CONDITIONING",), - "negative1": ("CONDITIONING",), + "clip1": ("CLIP", ), "model2": ("MODEL",), - "positive2": ("CONDITIONING",), - "negative2": ("CONDITIONING",), + "clip2": ("CLIP", ), + "positive": ("STRING",), + "negative": ("STRING",), + "lora_name": ("STRING",), "rescaled_steps": ("INT",{"default": 8, "min": 0, "max": 100},), "rescale_multiplier": ( "FLOAT", @@ -98,6 +224,8 @@ class TwoModelAdvancedKsampler: "steps1_start": ("INT", {"default": 0, "min": 0, "max": 100}), "steps1_end": ("INT", {"default": 0, "min": 0, "max": 100}), "start_steps2_shift": ("INT", {"default": 0, "min": -20, "max": 100}), + "token_normalization": (["none", "mean", "length", "length+mean"],), + "weight_interpretation": (["comfy", "A1111", "compel", "comfy++" ,"down_weight"],), "latent_image": ("LATENT",), } } @@ -108,25 +236,32 @@ class TwoModelAdvancedKsampler: CATEGORY = "sampling" def sample(self, - model1, positive1, negative1, - model2, positive2, negative2, + model1, clip1, + model2, clip2, + positive, negative, lora_name, rescaled_steps, rescale_multiplier, total_steps_original, total_steps_shift, noise_seed, sampler_name, scheduler, starting_cfg, cfg_shift, steps1_start, steps1_end, start_steps2_shift, + token_normalization,weight_interpretation, latent_image, **kwargs): + + model1, clip1, text1 = LoraTagLoader().load_lora(model1, clip1, lora_name) + model2, clip2, text2 = LoraTagLoader().load_lora(model2, clip2, lora_name) + + positive_conditioning1 = AdvancedCLIPTextEncodeWithBreak().encode(clip1, positive, token_normalization, weight_interpretation) + positive_conditioning2 = AdvancedCLIPTextEncodeWithBreak().encode(clip2, positive, token_normalization, weight_interpretation) + negative_conditioning1 = AdvancedCLIPTextEncodeWithBreak().encode(clip1, negative, token_normalization, weight_interpretation) + negative_conditioning2 = AdvancedCLIPTextEncodeWithBreak().encode(clip2, negative, token_normalization, weight_interpretation) # Debug print all input parameters print("\n=== TwoModelAdvancedKsampler Debug Information ===") print(f"Model1: {type(model1)}") print(f"Model2: {type(model2)}") - print(f"Positive1: {len(positive1)} elements") - print(f"Negative1: {len(negative1)} elements") - print(f"Positive2: {len(positive2)} elements") - print(f"Negative2: {len(negative2)} elements") + print(f"Total Steps Original: {total_steps_original}") print(f"Total Steps Shift: {total_steps_shift} (New total: {total_steps_original + total_steps_shift})") print(f"Noise Seed: {noise_seed}") @@ -153,8 +288,8 @@ class TwoModelAdvancedKsampler: starting_cfg, sampler_name, scheduler, - positive1, - negative1, + positive_conditioning1[0], + negative_conditioning1[0], latent_image, steps1_start, rescaled_steps, @@ -179,8 +314,8 @@ class TwoModelAdvancedKsampler: starting_cfg, sampler_name, scheduler, - positive1, - negative1, + positive_conditioning1[0], + negative_conditioning1[0], samples[0], actual_steps, steps1_end, @@ -214,8 +349,8 @@ class TwoModelAdvancedKsampler: new_cfg, sampler_name, scheduler, - positive2, - negative2, + positive_conditioning2[0], + negative_conditioning2[0], samples[0], model2_start_step, model2_end_step,