working v1

This commit is contained in:
n.marchenko
2025-06-21 07:04:51 +03:00
parent 1b0f763b0a
commit fa6c19c5e8
2 changed files with 441 additions and 16 deletions
+290
View File
@@ -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)
+151 -16
View File
@@ -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,