conditioning.py

import torch
import torch.nn.functional as F
import math

from torch  import Tensor
from typing import Optional, Callable, Tuple, Dict, Any, Union, TYPE_CHECKING, TypeVar, List

from dataclasses import dataclass, field

import copy
import base64
import pickle # used strictly for serializing conditioning in the ConditioningToBase64 and Base64ToConditioning nodes for API use. (Offloading T5 processing to another machine to avoid model shuffling.)

import comfy.supported_models
import node_helpers
import gc


from .sigmas  import get_sigmas

from .helper  import initialize_or_scale, precision_tool, get_res4lyf_scheduler_list, pad_tensor_list_to_max_len
from .latents import get_orthogonal, get_collinear
from .res4lyf import RESplain
from .beta.constants import MAX_STEPS
from .attention_masks import FullAttentionMask, FullAttentionMaskHiDream, CrossAttentionMask, SplitAttentionMask, RegionalContext

from .flux.redux import ReReduxImageEncoder

def multiply_nested_tensors(structure, scalar):
    if isinstance(structure, torch.Tensor):
        return structure * scalar
    elif isinstance(structure, list):
        return [multiply_nested_tensors(item, scalar) for item in structure]
    elif isinstance(structure, dict):
        return {key: multiply_nested_tensors(value, scalar) for key, value in structure.items()}
    else:
        return structure



def pad_to_same_tokens(x1, x2, pad_value=0.0):
    T1, T2 = x1.shape[1], x2.shape[1]
    if T1 == T2:
        return x1, x2
    max_T = max(T1, T2)
    x1_padded = F.pad(x1, (0, 0, 0, max_T - T1), value=pad_value)
    x2_padded = F.pad(x2, (0, 0, 0, max_T - T2), value=pad_value)
    return x1_padded, x2_padded



class ConditioningOrthoCollin:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "conditioning_0": ("CONDITIONING", ), 
            "conditioning_1": ("CONDITIONING", ),
            "t5_strength"   : ("FLOAT", {"default": 1.0, "min": -10000, "max": 10000, "step":0.01}),
            "clip_strength" : ("FLOAT", {"default": 1.0, "min": -10000, "max": 10000, "step":0.01}),
            }}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "combine"
    CATEGORY     = "RES4LYF/conditioning"
    EXPERIMENTAL = True

    def combine(self, conditioning_0, conditioning_1, t5_strength, clip_strength):

        t5_0_1_collin         = get_collinear (conditioning_0[0][0], conditioning_1[0][0])
        t5_1_0_ortho          = get_orthogonal(conditioning_1[0][0], conditioning_0[0][0])

        t5_combined           = t5_0_1_collin + t5_1_0_ortho
        
        t5_1_0_collin         = get_collinear (conditioning_1[0][0], conditioning_0[0][0])
        t5_0_1_ortho          = get_orthogonal(conditioning_0[0][0], conditioning_1[0][0])

        t5_B_combined         = t5_1_0_collin + t5_0_1_ortho

        pooled_0_1_collin     = get_collinear (conditioning_0[0][1]['pooled_output'].unsqueeze(0), conditioning_1[0][1]['pooled_output'].unsqueeze(0)).squeeze(0)
        pooled_1_0_ortho      = get_orthogonal(conditioning_1[0][1]['pooled_output'].unsqueeze(0), conditioning_0[0][1]['pooled_output'].unsqueeze(0)).squeeze(0)

        pooled_combined       = pooled_0_1_collin + pooled_1_0_ortho
        
        #conditioning_0[0][0] = conditioning_0[0][0] + t5_strength * (t5_combined - conditioning_0[0][0])
        
        #conditioning_0[0][0] = t5_strength * t5_combined + (1-t5_strength) * t5_B_combined
        
        conditioning_0[0][0]  = t5_strength * t5_0_1_collin + (1-t5_strength) * t5_1_0_collin
        
        conditioning_0[0][1]['pooled_output'] = conditioning_0[0][1]['pooled_output'] + clip_strength * (pooled_combined - conditioning_0[0][1]['pooled_output'])

        return (conditioning_0, )



class CLIPTextEncodeFluxUnguided:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {
            "clip"  : ("CLIP", ),
            "clip_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl" : ("STRING", {"multiline": True, "dynamicPrompts": True}),
            }}
    RETURN_TYPES = ("CONDITIONING","INT","INT",)
    RETURN_NAMES = ("conditioning", "clip_l_end", "t5xxl_end",)
    FUNCTION     = "encode"
    CATEGORY     = "RES4LYF/conditioning"
    EXPERIMENTAL = True


    def encode(self, clip, clip_l, t5xxl):
        tokens = clip.tokenize(clip_l)
        tokens["t5xxl"] = clip.tokenize(t5xxl)["t5xxl"]

        clip_l_end=0
        for i in range(len(tokens['l'][0])):
            if tokens['l'][0][i][0] == 49407:
                clip_l_end=i
                break
        t5xxl_end=0
        for i in range(len(tokens['l'][0])):   # bug? should this be t5xxl?
            if tokens['t5xxl'][0][i][0] == 1:
                t5xxl_end=i
                break

        output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
        cond = output.pop("cond")
        conditioning = [[cond, output]]
        conditioning[0][1]['clip_l_end'] = clip_l_end
        conditioning[0][1]['t5xxl_end'] = t5xxl_end
        return (conditioning, clip_l_end, t5xxl_end,)


class StyleModelApplyStyle: 
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning":       ("CONDITIONING", ),
                "style_model":        ("STYLE_MODEL", ),
                "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
                "strength":           ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.001}),
            }
    }
        
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"
    DESCRIPTION  = "Use with Flux Redux."
    EXPERIMENTAL = True

    def main(self, clip_vision_output, style_model, conditioning, strength=1.0):
        c = style_model.model.feature_match(conditioning, clip_vision_output)
        #cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
        #cond = strength * cond
        #c = []
        #for t in conditioning:
        #    n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
        #    c.append(n)
        return (c, )


class ConditioningZeroAndTruncate: 
    # needs updating to ensure dims are correct for arbitrary models without hardcoding. 
    # vanilla ConditioningZeroOut node doesn't truncate and SD3.5M degrades badly with large embeddings, even if zeroed out, as the negative conditioning
    @classmethod
    def INPUT_TYPES(cls):
        return { "required": {"conditioning": ("CONDITIONING", )}}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "zero_out"
    CATEGORY     = "RES4LYF/conditioning"
    DESCRIPTION  = "Use for negative conditioning with SD3.5. ConditioningZeroOut does not truncate the embedding, \
                    which results in severe degradation of image quality with SD3.5 when the token limit is exceeded."

    def zero_out(self, conditioning):
        c = []
        for t in conditioning:
            d = t[1].copy()
            pooled_output = d.get("pooled_output", None)
            if pooled_output is not None:
                d["pooled_output"] = torch.zeros((1,2048), dtype=t[0].dtype, device=t[0].device)
                n = [torch.zeros((1,154,4096), dtype=t[0].dtype, device=t[0].device), d]
            c.append(n)
        return (c, )


class ConditioningTruncate: 
    # needs updating to ensure dims are correct for arbitrary models without hardcoding. 
    @classmethod
    def INPUT_TYPES(cls):
        return { "required": {"conditioning": ("CONDITIONING", )}}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "zero_out"
    CATEGORY     = "RES4LYF/conditioning"
    DESCRIPTION  = "Use for positive conditioning with SD3.5. Tokens beyond 77 result in degradation of image quality."
    EXPERIMENTAL = True

    def zero_out(self, conditioning):
        c = []
        for t in conditioning:
            d = t[1].copy()
            pooled_output = d.get("pooled_output", None)
            if pooled_output is not None:
                d["pooled_output"] = d["pooled_output"][:, :2048]
                n = [t[0][:, :154, :4096], d]
            c.append(n)
        return (c, )


class ConditioningMultiply:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"conditioning": ("CONDITIONING", ), 
                            "multiplier": ("FLOAT", {"default": 1.0, "min": -1000000000.0, "max": 1000000000.0, "step": 0.01})
                            }}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def main(self, conditioning, multiplier):
        c = multiply_nested_tensors(conditioning, multiplier)
        return (c,)



class ConditioningAdd:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"conditioning_1": ("CONDITIONING", ), 
                            "conditioning_2": ("CONDITIONING", ), 
                            "multiplier": ("FLOAT", {"default": 1.0, "min": -1000000000.0, "max": 1000000000.0, "step": 0.01})
                            }}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def main(self, conditioning_1, conditioning_2, multiplier):
        
        conditioning_1[0][0]                  += multiplier * conditioning_2[0][0]
        conditioning_1[0][1]['pooled_output'] += multiplier * conditioning_2[0][1]['pooled_output'] 
        
        return (conditioning_1,)




class ConditioningCombine:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"conditioning_1": ("CONDITIONING", ), "conditioning_2": ("CONDITIONING", )}}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "combine"
    CATEGORY     = "RES4LYF/conditioning"

    def combine(self, conditioning_1, conditioning_2):
        return (conditioning_1 + conditioning_2, )



class ConditioningAverage :
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning_to":          ("CONDITIONING", ), 
                "conditioning_from":        ("CONDITIONING", ),
                "conditioning_to_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
                }
            }
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    CATEGORY     = "RES4LYF/conditioning"
    FUNCTION     = "addWeighted"

    def addWeighted(self, conditioning_to, conditioning_from, conditioning_to_strength):
        out = []

        if len(conditioning_from) > 1:
            RESplain("Warning: ConditioningAverage conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")

        cond_from = conditioning_from[0][0]
        pooled_output_from = conditioning_from[0][1].get("pooled_output", None)

        for i in range(len(conditioning_to)):
            t1 = conditioning_to[i][0]
            pooled_output_to = conditioning_to[i][1].get("pooled_output", pooled_output_from)
            t0 = cond_from[:,:t1.shape[1]]
            if t0.shape[1] < t1.shape[1]:
                t0 = torch.cat([t0] + [torch.zeros((1, (t1.shape[1] - t0.shape[1]), t1.shape[2]))], dim=1)

            tw = torch.mul(t1, conditioning_to_strength) + torch.mul(t0, (1.0 - conditioning_to_strength))
            t_to = conditioning_to[i][1].copy()
            if pooled_output_from is not None and pooled_output_to is not None:
                t_to["pooled_output"] = torch.mul(pooled_output_to, conditioning_to_strength) + torch.mul(pooled_output_from, (1.0 - conditioning_to_strength))
            elif pooled_output_from is not None:
                t_to["pooled_output"] = pooled_output_from

            n = [tw, t_to]
            out.append(n)
        return (out, )
    
class ConditioningSetTimestepRange:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {"conditioning": ("CONDITIONING", ),
                            "start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                            "end":   ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
                            }}
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "set_range"
    CATEGORY     = "RES4LYF/conditioning"

    def set_range(self, conditioning, start, end):
        c = node_helpers.conditioning_set_values(conditioning, {"start_percent": start,
                                                                "end_percent": end})
        return (c, )

class ConditioningAverageScheduler: # don't think this is implemented correctly. needs to be reworked
    @classmethod
    def INPUT_TYPES(cls):
        return {
                "required": {
                    "conditioning_0": ("CONDITIONING", ), 
                    "conditioning_1": ("CONDITIONING", ),
                    "ratio": ("SIGMAS", ),
                    }
            }
    
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"
    EXPERIMENTAL = True

    @staticmethod
    def addWeighted(conditioning_to, conditioning_from, conditioning_to_strength): #this function borrowed from comfyui
        out = []

        if len(conditioning_from) > 1:
            RESplain("Warning: ConditioningAverage conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")

        cond_from = conditioning_from[0][0]
        pooled_output_from = conditioning_from[0][1].get("pooled_output", None)

        for i in range(len(conditioning_to)):
            t1 = conditioning_to[i][0]
            pooled_output_to = conditioning_to[i][1].get("pooled_output", pooled_output_from)
            t0 = cond_from[:,:t1.shape[1]]
            if t0.shape[1] < t1.shape[1]:
                t0 = torch.cat([t0] + [torch.zeros((1, (t1.shape[1] - t0.shape[1]), t1.shape[2]))], dim=1)

            tw = torch.mul(t1, conditioning_to_strength) + torch.mul(t0, (1.0 - conditioning_to_strength))
            t_to = conditioning_to[i][1].copy()
            if pooled_output_from is not None and pooled_output_to is not None:
                t_to["pooled_output"] = torch.mul(pooled_output_to, conditioning_to_strength) + torch.mul(pooled_output_from, (1.0 - conditioning_to_strength))
            elif pooled_output_from is not None:
                t_to["pooled_output"] = pooled_output_from

            n = [tw, t_to]
            out.append(n)
        return out

    @staticmethod
    def create_percent_array(steps):
        step_size = 1.0 / steps
        return [{"start_percent": i * step_size, "end_percent": (i + 1) * step_size} for i in range(steps)]

    def main(self, conditioning_0, conditioning_1, ratio):
        steps = len(ratio)

        percents = self.create_percent_array(steps)

        cond = []
        for i in range(steps):
            average = self.addWeighted(conditioning_0, conditioning_1, ratio[i].item())
            cond += node_helpers.conditioning_set_values(average, {"start_percent": percents[i]["start_percent"], "end_percent": percents[i]["end_percent"]})

        return (cond,)



class StableCascade_StageB_Conditioning64:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "conditioning": ("CONDITIONING",),
                "stage_c":      ("LATENT",),
                }
            }
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "set_prior"
    CATEGORY     = "RES4LYF/conditioning"

    @precision_tool.cast_tensor
    def set_prior(self, conditioning, stage_c):
        c = []
        for t in conditioning:
            d = t[1].copy()
            d['stable_cascade_prior'] = stage_c['samples']
            n = [t[0], d]
            c.append(n)
        return (c, )



class Conditioning_Recast64:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required": { "cond_0": ("CONDITIONING",),
                            },
                "optional": { "cond_1": ("CONDITIONING",),}
                }
    RETURN_TYPES = ("CONDITIONING","CONDITIONING",)
    RETURN_NAMES = ("cond_0_recast","cond_1_recast",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/precision"
    EXPERIMENTAL = True

    @precision_tool.cast_tensor
    def main(self, cond_0, cond_1 = None):
        cond_0[0][0] = cond_0[0][0].to(torch.float64)
        if 'pooled_output' in cond_0[0][1]:
            cond_0[0][1]["pooled_output"] = cond_0[0][1]["pooled_output"].to(torch.float64)
        
        if cond_1 is not None:
            cond_1[0][0] = cond_1[0][0].to(torch.float64)
            if 'pooled_output' in cond_0[0][1]:
                cond_1[0][1]["pooled_output"] = cond_1[0][1]["pooled_output"].to(torch.float64)

        return (cond_0, cond_1,)


class ConditioningToBase64:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning": ("CONDITIONING",),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES   = ("STRING",)
    RETURN_NAMES   = ("string",)
    FUNCTION       = "notify"
    OUTPUT_NODE    = True
    OUTPUT_IS_LIST = (True,)
    CATEGORY       = "RES4LYF/utilities"

    def notify(self, unique_id=None, extra_pnginfo=None, conditioning=None):
        
        conditioning_pickle = pickle.dumps(conditioning)
        conditioning_base64 = base64.b64encode(conditioning_pickle).decode('utf-8')
        text = [conditioning_base64]
        
        if unique_id is not None and extra_pnginfo is not None:
            if not isinstance(extra_pnginfo, list):
                RESplain("Error: extra_pnginfo is not a list")
            elif (
                not isinstance(extra_pnginfo[0], dict)
                or "workflow" not in extra_pnginfo[0]
            ):
                RESplain("Error: extra_pnginfo[0] is not a dict or missing 'workflow' key")
            else:
                workflow = extra_pnginfo[0]["workflow"]
                node = next(
                    (x for x in workflow["nodes"] if str(x["id"]) == str(unique_id[0])),
                    None,
                )
                if node:
                    node["widgets_values"] = [text]

        return {"ui": {"text": text}, "result": (text,)}

class Base64ToConditioning:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "data": ("STRING", {"default": ""}),
            }
        }
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/utilities"

    def main(self, data):
        conditioning_pickle = base64.b64decode(data)
        conditioning = pickle.loads(conditioning_pickle)
        return (conditioning,)




class ConditioningDownsampleT5:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "conditioning": ("CONDITIONING",),
                "token_limit" : ("INT", {'default': 128, 'min': 1, 'max': 16384}),
            },
            "optional": {
            }
        }
        
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"
    EXPERIMENTAL = True

    def main(self, conditioning, token_limit):
        
        conditioning[0][0] = downsample_tokens(conditioning[0][0], token_limit)
        return (conditioning, )




"""class ConditioningBatch4:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "conditioning_0": ("CONDITIONING",),
                },
            "optional": {
                "conditioning_1": ("CONDITIONING",),
                "conditioning_2": ("CONDITIONING",),
                "conditioning_3": ("CONDITIONING",),
            }
            }
        
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def main(self, conditioning_0, conditioning_1=None, conditioning_2=None, conditioning_3=None, ):
        c = copy.deepcopy(conditioning_0)
        
        if conditioning_1 is not None:
            c.append(conditioning_1[0])
            
        if conditioning_2 is not None:
            c.append(conditioning_2[0])
            
        if conditioning_3 is not None:
            c.append(conditioning_3[0])

        return (c, )"""


class ConditioningBatch4:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "conditioning_0": ("CONDITIONING",),
                },
            "optional": {
                "conditioning_1": ("CONDITIONING",),
                "conditioning_2": ("CONDITIONING",),
                "conditioning_3": ("CONDITIONING",),
            }
            }
        
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def main(self, conditioning_0, conditioning_1=None, conditioning_2=None, conditioning_3=None, ):
        c = []
        c.append(conditioning_0)
        
        if conditioning_1 is not None:
            c.append(conditioning_1)
            
        if conditioning_2 is not None:
            c.append(conditioning_2)
            
        if conditioning_3 is not None:
            c.append(conditioning_3)

        return (c, )



class ConditioningBatch8:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "conditioning_0": ("CONDITIONING",),
                },
            "optional": {
                "conditioning_1": ("CONDITIONING",),
                "conditioning_2": ("CONDITIONING",),
                "conditioning_3": ("CONDITIONING",),
                "conditioning_4": ("CONDITIONING",),
                "conditioning_5": ("CONDITIONING",),
                "conditioning_6": ("CONDITIONING",),
                "conditioning_7": ("CONDITIONING",),
            }
            }
        
    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def main(self, conditioning_0, conditioning_1=None, conditioning_2=None, conditioning_3=None, conditioning_4=None, conditioning_5=None, conditioning_6=None, conditioning_7=None, ):
        c = []
        c.append(conditioning_0)
        
        if conditioning_1 is not None:
            c.append(conditioning_1)
            
        if conditioning_2 is not None:
            c.append(conditioning_2)
            
        if conditioning_3 is not None:
            c.append(conditioning_3)
            
        if conditioning_4 is not None:
            c.append(conditioning_4)
            
        if conditioning_5 is not None:
            c.append(conditioning_5)
            
        if conditioning_6 is not None:
            c.append(conditioning_6)
            
        if conditioning_7 is not None:
            c.append(conditioning_7)
            
        return (c, )



class EmptyConditioningGenerator:
    def __init__(self, model=None, conditioning=None, device=None, dtype=None):
        """ device, dtype currently unused """
        if model is not None:
                    
            self.device = device
            self.dtype  = dtype
        
            import comfy.supported_models
            self.model_config = model.model.model_config

            self.llama3_shape = None
            self.pooled_len    = 0

            if isinstance(self.model_config, comfy.supported_models.SD3):
                self.text_len_base = 154
                self.text_channels = 4096
                self.pooled_len    = 2048
            elif isinstance(self.model_config, (comfy.supported_models.Flux, comfy.supported_models.FluxSchnell, comfy.supported_models.Chroma)):
                self.text_len_base = 256
                self.text_channels = 4096
                self.pooled_len    = 768
            elif isinstance(self.model_config, comfy.supported_models.AuraFlow):
                self.text_len_base = 256
                self.text_channels = 2048
                #self.pooled_len    = 1
            elif isinstance(self.model_config, comfy.supported_models.Stable_Cascade_C):
                self.text_len_base = 77
                self.text_channels = 1280
                self.pooled_len    = 1280
            elif isinstance(self.model_config, comfy.supported_models.WAN21_T2V) or isinstance(self.model_config, comfy.supported_models.WAN21_I2V):
                self.text_len_base = 512
                self.text_channels = 5120 # sometimes needs to be 4096, like when initializing in samplers_py in shark?
                #self.pooled_len    = 1
            elif isinstance(self.model_config, comfy.supported_models.HiDream):
                self.text_len_base = 128
                self.text_channels = 4096 # sometimes needs to be 4096, like when initializing in samplers_py in shark?
                self.pooled_len    = 2048
                self.llama3_shape  = torch.Size([1,32,128,4096])
            elif isinstance(self.model_config, comfy.supported_models.LTXV):
                self.text_len_base = 128
                self.text_channels = 4096
                #self.pooled_len    = 1
            elif isinstance(self.model_config, comfy.supported_models.SD15):
                self.text_len_base = 77
                self.text_channels = 768
                self.pooled_len    = 768
            elif isinstance(self.model_config, comfy.supported_models.SDXL):
                self.text_len_base = 77
                self.text_channels = 2048
                self.pooled_len    = 1280
            elif isinstance(self.model_config, comfy.supported_models.HunyuanVideo) or \
                isinstance (self.model_config, comfy.supported_models.HunyuanVideoI2V) or \
                isinstance (self.model_config, comfy.supported_models.HunyuanVideoSkyreelsI2V):
                self.text_len_base = 128
                self.text_channels = 4096
                #self.pooled_len    = 1
            else:
                raise ValueError(f"Unknown model config: {type(self.model_config)}")
        elif conditioning is not None:
            self.device        = conditioning[0][0].device
            self.dtype         = conditioning[0][0].dtype
            self.text_len_base = conditioning[0][0].shape[-2]
            if 'pooled_output' in conditioning[0][1]:
                self.pooled_len = conditioning[0][1]['pooled_output'].shape[-1]
            else:
                self.pooled_len = 0
            self.text_channels = conditioning[0][0].shape[-1]
            
    def get_empty_conditioning(self):
        if self.llama3_shape is not None and self.pooled_len > 0:
            return [[
                torch.zeros((1, self.text_len_base, self.text_channels)),
                {
                    'pooled_output'      : torch.zeros((1, self.pooled_len)),
                    'conditioning_llama3': torch.zeros(self.llama3_shape),
                }
            ]]
        elif self.pooled_len > 0:
            return [[
                torch.zeros((1, self.text_len_base, self.text_channels)),
                {
                    'pooled_output': torch.zeros((1, self.pooled_len)),
                }
            ]]
        else:
            return [[
                torch.zeros((1, self.text_len_base, self.text_channels)),
            ]]

    def get_empty_conditionings(self, count):
        return [self.get_empty_conditioning() for _ in range(count)]
    
    def zero_none_conditionings_(self, *conds):
        if len(conds) == 1 and isinstance(conds[0], (list, tuple)):
            conds = conds[0]
        for i, cond in enumerate(conds):
            conds[i] = self.get_empty_conditioning() if cond is None else cond
        return conds

"""def zero_conditioning_from_list(conds):
    for cond in conds:
        if cond is not None:
            for i in range(len(cond)):
                pooled     = cond[i][1].get('pooled_output')
                pooled_len = pooled.shape[-1] if pooled is not None else 1    # 1 default pooled_output len for those without it
                
                cond_zero  = [[
                    torch.zeros_like(cond[i][0]),
                    {"pooled_output": torch.zeros((1,pooled_len), dtype=cond[i][0].dtype, device=cond[i][0].device)},
                ]]
                
            return cond_zero"""

def zero_conditioning_from_list(conds):
    for cond in conds:
        if cond is not None:
            for i in range(len(cond)):
                pooled = cond[i][1].get('pooled_output')
                llama3 = cond[i][1].get('conditioning_llama3')

                pooled_len = pooled.shape[-1] if pooled is not None else 1
                llama3_shape = llama3.shape if llama3 is not None else (1, 32, 128, 4096)

                cond_zero = [[
                    torch.zeros_like(cond[i][0]),
                    {
                        "pooled_output": torch.zeros((1, pooled_len), dtype=cond[i][0].dtype, device=cond[i][0].device),
                        "conditioning_llama3": torch.zeros(llama3_shape, dtype=cond[i][0].dtype, device=cond[i][0].device),
                    },
                ]]

            return cond_zero

class TemporalMaskGenerator:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                    {
                    "switch_frame": ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "frames":       ("INT", {"default": 65, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "invert_mask":             ("BOOLEAN",                                  {"default": False}),

                    },
                "optional": 
                    {
                    }
                }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("temporal_mask",) 
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/masks"
    EXPERIMENTAL = True
    
    def main(self,
            switch_frame = 33,
            frames = 65,
            invert_mask = False,
            ):
        
        switch_frame = switch_frame // 4
        frames = frames // 4 + 1
        
        temporal_mask = torch.ones((frames, 2, 2))
        
        temporal_mask[switch_frame:,...] = 0.0
        
        if invert_mask:
            temporal_mask = 1 - temporal_mask
        
        return (temporal_mask,)




class TemporalSplitAttnMask_Midframe:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                    {
                    "self_attn_midframe":  ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "cross_attn_midframe": ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "self_attn_invert":    ("BOOLEAN",                                  {"default": False}),
                    "cross_attn_invert":   ("BOOLEAN",                                  {"default": False}),
                    "frames":             ("INT", {"default": 65, "min": 1, "step": 4, "max": 0xffffffffffffffff}),

                    },
                "optional": 
                    {
                    }
                }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("temporal_mask",) 
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/masks"
    EXPERIMENTAL = True
    
    def main(self,
            self_attn_midframe = 33,
            cross_attn_midframe = 33,
            self_attn_invert = False,
            cross_attn_invert = False,
            frames = 65,
            ):

        frames = frames // 4 + 1
        
        temporal_self_mask  = torch.ones((frames, 2, 2))
        temporal_cross_mask = torch.ones((frames, 2, 2))

        
        self_attn_midframe  = self_attn_midframe  // 4
        cross_attn_midframe = cross_attn_midframe // 4
        
        temporal_self_mask[self_attn_midframe  :,...] = 0.0
        temporal_cross_mask[cross_attn_midframe:,...] = 0.0
        
        if self_attn_invert:
            temporal_self_mask  = 1 - temporal_self_mask
            
        if cross_attn_invert:
            temporal_cross_mask = 1 - temporal_cross_mask
        
        temporal_attn_masks = torch.stack([temporal_cross_mask, temporal_self_mask])
        
        return (temporal_attn_masks,)




class TemporalSplitAttnMask:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                    {
                    "self_attn_start":  ("INT", {"default": 1,  "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "self_attn_stop":   ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "cross_attn_start": ("INT", {"default": 1,  "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "cross_attn_stop":  ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),

                    #"frames":           ("INT", {"default": 65, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    },
                "optional": 
                    {
                    }
                }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("temporal_mask",) 
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/masks"
    
    def main(self,
            self_attn_start  = 0,
            self_attn_stop   = 33,
            cross_attn_start = 0,
            cross_attn_stop  = 33,
            #frames           = 65,
            ):

        #frames = frames // 4 + 1
        
        self_attn_start  = self_attn_start  // 4 #+ 1
        self_attn_stop   = self_attn_stop   // 4 + 1
        cross_attn_start = cross_attn_start // 4 #+ 1
        cross_attn_stop  = cross_attn_stop  // 4 + 1
        
        max_stop = max(self_attn_stop, cross_attn_stop)
        
        temporal_self_mask  = torch.zeros((max_stop, 1, 1))
        temporal_cross_mask = torch.zeros((max_stop, 1, 1))

        temporal_self_mask [ self_attn_start: self_attn_stop,...] = 1.0
        temporal_cross_mask[cross_attn_start:cross_attn_stop,...] = 1.0
        
        temporal_attn_masks = torch.stack([temporal_cross_mask, temporal_self_mask])
        
        return (temporal_attn_masks,)




class TemporalCrossAttnMask:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                    {
                    "cross_attn_start": ("INT", {"default": 1,  "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    "cross_attn_stop":  ("INT", {"default": 33, "min": 1, "step": 4, "max": 0xffffffffffffffff}),
                    },
                "optional": 
                    {
                    }
                }

    RETURN_TYPES = ("MASK",)
    RETURN_NAMES = ("temporal_mask",) 
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/masks"
    
    def main(self,
            cross_attn_start = 0,
            cross_attn_stop  = 33,
            ):
        
        cross_attn_start = cross_attn_start // 4 #+ 1
        cross_attn_stop  = cross_attn_stop  // 4 + 1
        
        temporal_self_mask  = torch.zeros((cross_attn_stop, 1, 1))  # dummy to satisfy stack
        temporal_cross_mask = torch.zeros((cross_attn_stop, 1, 1))

        temporal_cross_mask[cross_attn_start:cross_attn_stop,...] = 1.0
        
        temporal_attn_masks = torch.stack([temporal_cross_mask, temporal_self_mask])
        
        return (temporal_attn_masks,)




@dataclass
class RegionalParameters:
    weights    : List[float] = field(default_factory=list)
    floors     : List[float] = field(default_factory=list)



REG_MASK_TYPE_2 = [
    "gradient",
    "gradient_masked",
    "gradient_unmasked",
    "boolean",
    "boolean_masked",
    "boolean_unmasked",
]

REG_MASK_TYPE_3 = [
    "gradient",
    "gradient_A",
    "gradient_B",
    "gradient_unmasked",
    "gradient_AB",
    "gradient_A,unmasked",
    "gradient_B,unmasked",

    "boolean",
    "boolean_A",
    "boolean_B",
    "boolean_unmasked",
    "boolean_AB",
    "boolean_A,unmasked",
    "boolean_B,unmasked",
]

REG_MASK_TYPE_AB = [
    "gradient",
    "gradient_A",
    "gradient_B",
    "boolean",
    "boolean_A",
    "boolean_B",
]

REG_MASK_TYPE_ABC = [
    "gradient",
    "gradient_A",
    "gradient_B",
    "gradient_C",
    "gradient_AB",
    "gradient_AC",
    "gradient_BC",

    "boolean",
    "boolean_A",
    "boolean_B",
    "boolean_C",
    "boolean_AB",
    "boolean_AC",
    "boolean_BC",

]




class ClownRegionalConditioning_AB:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "weight":                  ("FLOAT",                                     {"default": 1.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed":            ("FLOAT",                                     {"default": 0.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed_start_step": ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "weight_scheduler":        (["constant"] + get_res4lyf_scheduler_list(), {"default": "constant"},),
                "start_step":              ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "end_step":                ("INT",                                       {"default": -1,  "min": -1,        "max": 10000}),
                "mask_type":               (REG_MASK_TYPE_AB,                            {"default": "boolean"}),
                "edge_width":              ("INT",                                       {"default": 0,  "min": 0,          "max": 10000}),
                "invert_mask":             ("BOOLEAN",                                   {"default": False}),
            }, 
            "optional": {
                "conditioning_A":          ("CONDITIONING", ),
                "conditioning_B":          ("CONDITIONING", ),
                "mask_A":                  ("MASK", ),
                "mask_B":                  ("MASK", ),
                "weights":                 ("SIGMAS", ),
                "region_bleeds":           ("SIGMAS", ),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def create_callback(self, **kwargs):
        def callback(model):
            kwargs["model"] = model  
            pos_cond, = self.prepare_regional_cond(**kwargs)
            return pos_cond
        return callback

    def main(self,
            weight                   : float  = 1.0,
            start_sigma              : float  = 0.0,
            end_sigma                : float  = 1.0,
            weight_scheduler                  = None,
            start_step               : int    = 0,
            end_step                 : int    = -1,
            conditioning_A                    = None,
            conditioning_B                    = None,
            weights                  : Tensor = None,
            region_bleeds            : Tensor = None,
            region_bleed             : float  = 0.0,
            region_bleed_start_step  : int    = 0,
            mask_type                : str    = "boolean",
            edge_width               : int    = 0,
            mask_A                            = None,
            mask_B                            = None,
            invert_mask              : bool   = False
            ) -> Tuple[Tensor]:
        
        mask   = mask_A
        unmask = mask_B
        
        if end_step == -1:
            end_step = MAX_STEPS

        callback = self.create_callback(weight                   = weight,
                                        start_sigma              = start_sigma,
                                        end_sigma                = end_sigma,
                                        weight_scheduler         = weight_scheduler,
                                        start_step               = start_step,
                                        end_step                 = end_step,
                                        weights                  = weights,
                                        region_bleeds            = region_bleeds,
                                        region_bleed             = region_bleed,
                                        region_bleed_start_step  = region_bleed_start_step,
                                        mask_type                = mask_type,
                                        edge_width               = edge_width,
                                        mask                     = mask,
                                        unmask                   = unmask,
                                        invert_mask              = invert_mask,
                                        conditioning_A           = conditioning_A,
                                        conditioning_B           = conditioning_B,
                                        )

        cond = zero_conditioning_from_list([conditioning_A, conditioning_B])
        
        cond[0][1]['callback_regional'] = callback
        
        return (cond,)



    def prepare_regional_cond(self,
                                model,
                                weight                   : float  = 1.0,
                                start_sigma              : float  = 0.0,
                                end_sigma                : float  = 1.0,
                                weight_scheduler                  = None,
                                start_step               : int    = 0,
                                end_step                 : int    = -1,
                                conditioning_A                    = None,
                                conditioning_B                    = None,
                                weights                  : Tensor = None,
                                region_bleeds            : Tensor = None,
                                region_bleed             : float  = 0.0,
                                region_bleed_start_step  : int    = 0,
                                mask_type                : str    = "gradient",
                                edge_width               : int    = 0,
                                mask                              = None,
                                unmask                            = None,
                                invert_mask              : bool   = False,
                                ) -> Tuple[Tensor]:

        default_dtype  = torch.float64
        default_device = torch.device("cuda") 
        
        if end_step == -1:
            end_step = MAX_STEPS
        
        if weights is None and weight_scheduler != "constant":
            total_steps = end_step - start_step
            weights     = get_sigmas(model, weight_scheduler, total_steps, 1.0).to(dtype=default_dtype, device=default_device) #/ model.inner_model.inner_model.model_sampling.sigma_max  #scaling doesn't matter as this is a flux-only node
            prepend     = torch.zeros(start_step,                                  dtype=default_dtype, device=default_device)
            weights     = torch.cat((prepend, weights), dim=0)
        
        if invert_mask and mask is not None:
            mask   = 1-mask
            unmask = 1-unmask

        floor, floors = region_bleed, region_bleeds
        
        weights = initialize_or_scale(weights, weight, end_step).to(default_dtype).to(default_device)
        weights = F.pad(weights, (0, MAX_STEPS), value=0.0)
        
        prepend = torch.full((region_bleed_start_step,),  0.0, dtype=default_dtype, device=default_device)
        floors  = initialize_or_scale(floors,  floor,  end_step).to(default_dtype).to(default_device)
        floors  = F.pad(floors,  (0, MAX_STEPS), value=0.0)
        floors  = torch.cat((prepend, floors), dim=0)

        if (conditioning_A is None) and (conditioning_B is None):
            cond = None

        elif mask is not None:
            EmptyCondGen = EmptyConditioningGenerator(model)
            conditioning_A, conditioning_B = EmptyCondGen.zero_none_conditionings_([conditioning_A, conditioning_B])
            
            cond = copy.deepcopy(conditioning_A)
            
            if isinstance(model.model.model_config, (comfy.supported_models.WAN21_T2V, comfy.supported_models.WAN21_I2V)):
                if model.model.diffusion_model.blocks[0].self_attn.winderz_type != "false":
                    AttnMask = CrossAttentionMask(mask_type, edge_width)
                else:
                    AttnMask = SplitAttentionMask(mask_type, edge_width)
            elif isinstance(model.model.model_config, comfy.supported_models.HiDream):
                AttnMask = FullAttentionMaskHiDream(mask_type, edge_width)
            elif isinstance(model.model.model_config, (comfy.supported_models.SDXL, comfy.supported_models.SD15, comfy.supported_models.Stable_Cascade_C)):
                AttnMask = SplitAttentionMask(mask_type, edge_width)
            else:
                AttnMask = FullAttentionMask(mask_type, edge_width)

            RegContext = RegionalContext()
            
            if isinstance(model.model.model_config, comfy.supported_models.HiDream):

                AttnMask.add_region_sizes(
                    [
                        conditioning_A[0][0].shape[-2],
                        conditioning_A[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        conditioning_A[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    mask)
                AttnMask.add_region_sizes(
                    [
                        conditioning_B[0][0].shape[-2],
                        conditioning_B[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        conditioning_B[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    unmask)

                RegContext.add_region_llama3(conditioning_A[0][1]['conditioning_llama3'])
                RegContext.add_region_llama3(conditioning_B[0][1]['conditioning_llama3'])
            else:
                AttnMask.add_region(conditioning_A[0][0],   mask)
                AttnMask.add_region(conditioning_B[0][0], unmask)
            
            RegContext.add_region(conditioning_A[0][0], conditioning_A[0][1].get('pooled_output'))
            RegContext.add_region(conditioning_B[0][0], conditioning_B[0][1].get('pooled_output'))
            
            if 'clip_vision_output' in conditioning_A[0][1]: # For WAN... dicey results
                RegContext.add_region_clip_fea(conditioning_A[0][1]['clip_vision_output'].penultimate_hidden_states)
                RegContext.add_region_clip_fea(conditioning_B[0][1]['clip_vision_output'].penultimate_hidden_states)
            if 'unclip_conditioning' in conditioning_A[0][1]:
                RegContext.add_region_clip_fea(conditioning_A[0][1]['unclip_conditioning'][0]['clip_vision_output'].image_embeds) #['penultimate_hidden_states'])
            if 'unclip_conditioning' in conditioning_B[0][1]:
                RegContext.add_region_clip_fea(conditioning_B[0][1]['unclip_conditioning'][0]['clip_vision_output'].image_embeds) #['penultimate_hidden_states'])
                
            cond[0][1]['AttnMask'] = AttnMask
            cond[0][1]['RegContext'] = RegContext
            
            cond = merge_with_base(base=cond, others=[conditioning_A, conditioning_B])
            
            if 'pooled_output' in cond[0][1] and cond[0][1]['pooled_output'] is not None:
                cond[0][1]['pooled_output'] = (conditioning_A[0][1]['pooled_output'] + conditioning_B[0][1]['pooled_output']) / 2
            
            #if 'conditioning_llama3' in cond[0][1] and cond[0][1]['conditioning_llama3'] is not None:
            #    cond[0][1]['conditioning_llama3'] = (conditioning_A[0][1]['conditioning_llama3'] + conditioning_B[0][1]['conditioning_llama3']) / 2
            #cond[0] = list(cond[0])
            #cond[0][0] = (conditioning_A[0][0] + conditioning_B[0][0]) / 2
            #cond[0] = tuple(cond[0])
            
        else:
            cond = conditioning_A
            
        cond[0][1]['RegParam'] = RegionalParameters(weights, floors)
        
        return (cond,)





class ClownRegionalConditioning_ABC:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "weight":                  ("FLOAT",                                     {"default": 1.0,  "min": -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed":            ("FLOAT",                                     {"default": 0.0,  "min": -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed_start_step": ("INT",                                       {"default": 0,    "min":  0,       "max": 10000}),
                "weight_scheduler":        (["constant"] + get_res4lyf_scheduler_list(), {"default": "constant"},),
                "start_step":              ("INT",                                       {"default": 0,    "min":  0,       "max": 10000}),
                "end_step":                ("INT",                                       {"default": 100,  "min": -1,       "max": 10000}),
                "mask_type":               (REG_MASK_TYPE_ABC,                           {"default": "boolean"}),
                "edge_width":              ("INT",                                       {"default": 0,    "min": 0,        "max": 10000}),
                "invert_mask":             ("BOOLEAN",                                   {"default": False}),
            }, 
            "optional": {
                "conditioning_A":          ("CONDITIONING", ),
                "conditioning_B":          ("CONDITIONING", ),
                "conditioning_C":          ("CONDITIONING", ),
                "mask_A":                  ("MASK", ),
                "mask_B":                  ("MASK", ),
                "mask_C":                  ("MASK", ),
                "weights":                 ("SIGMAS", ),
                "region_bleeds":           ("SIGMAS", ),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def create_callback(self, **kwargs):
        def callback(model):
            kwargs["model"] = model  
            pos_cond, = self.prepare_regional_cond(**kwargs)
            return pos_cond
        return callback

    def main(self,
            weight                   : float  = 1.0,
            start_sigma              : float  = 0.0,
            end_sigma                : float  = 1.0,
            weight_scheduler                  = None,
            start_step               : int    = 0,
            end_step                 : int    = -1,
            conditioning_A                    = None,
            conditioning_B                    = None,
            conditioning_C                    = None,
            weights                  : Tensor = None,
            region_bleeds            : Tensor = None,
            region_bleed             : float  = 0.0,
            region_bleed_start_step  : int    = 0,

            mask_type                : str    = "boolean",
            edge_width               : int    = 0,
            mask_A                              = None,
            mask_B                              = None,
            mask_C                              = None,
            invert_mask              : bool   = False
            ) -> Tuple[Tensor]:

        if end_step == -1:
            end_step = MAX_STEPS
        
        callback = self.create_callback(weight                   = weight,
                                        start_sigma              = start_sigma,
                                        end_sigma                = end_sigma,
                                        weight_scheduler         = weight_scheduler,
                                        start_step               = start_step,
                                        end_step                 = end_step,
                                        weights                  = weights,
                                        region_bleeds            = region_bleeds,
                                        region_bleed             = region_bleed,
                                        region_bleed_start_step  = region_bleed_start_step,
                                        mask_type                = mask_type,
                                        edge_width               = edge_width,
                                        mask_A                   = mask_A,
                                        mask_B                   = mask_B,
                                        mask_C                   = mask_C,
                                        invert_mask              = invert_mask,
                                        conditioning_A           = conditioning_A,
                                        conditioning_B           = conditioning_B,
                                        conditioning_C           = conditioning_C,
                                        )

        cond = zero_conditioning_from_list([conditioning_A, conditioning_B, conditioning_C])
        
        cond[0][1]['callback_regional'] = callback
        
        return (cond,)



    def prepare_regional_cond(self,
                                model,
                                weight                   : float  = 1.0,
                                start_sigma              : float  = 0.0,
                                end_sigma                : float  = 1.0,
                                weight_scheduler                  = None,
                                start_step               : int    =  0,
                                end_step                 : int    = -1,
                                conditioning_A                    = None,
                                conditioning_B                    = None,

                                conditioning_C                    = None,
                                weights                  : Tensor = None,
                                region_bleeds            : Tensor = None,
                                region_bleed             : float  = 0.0,
                                region_bleed_start_step  : int    = 0,

                                mask_type                : str    = "boolean",
                                edge_width               : int    = 0,
                                mask_A                            = None,
                                mask_B                            = None,
                                mask_C                            = None,
                                invert_mask              : bool   = False,
                                ) -> Tuple[Tensor]:

        default_dtype  = torch.float64
        default_device = torch.device("cuda") 
        
        if end_step == -1:
            end_step = MAX_STEPS
        
        if weights is None and weight_scheduler != "constant":
            total_steps = end_step - start_step
            weights     = get_sigmas(model, weight_scheduler, total_steps, 1.0).to(dtype=default_dtype, device=default_device) #/ model.inner_model.inner_model.model_sampling.sigma_max  #scaling doesn't matter as this is a flux-only node
            prepend     = torch.zeros(start_step,                                  dtype=default_dtype, device=default_device)
            weights     = torch.cat((prepend, weights), dim=0)
        
        if invert_mask and mask_A is not None:
            mask_A = 1-mask_A
            
        if invert_mask and mask_B is not None:
            mask_B = 1-mask_B
        
        mask_AB_inv = mask_C
        if invert_mask and mask_AB_inv is not None:
            mask_AB_inv = 1-mask_AB_inv
        
        floor, floors = region_bleed, region_bleeds
        
        weights = initialize_or_scale(weights, weight, end_step).to(default_dtype)
        weights = F.pad(weights, (0, MAX_STEPS), value=0.0)
        
        prepend    = torch.full((region_bleed_start_step,),  0.0, dtype=default_dtype, device=default_device)
        floors  = initialize_or_scale(floors,  floor,  end_step).to(default_dtype).to(default_device)
        floors  = F.pad(floors,  (0, MAX_STEPS), value=0.0)
        floors  = torch.cat((prepend, floors), dim=0)

        if (conditioning_A is None) and (conditioning_B is None) and (conditioning_C is None):
            conditioning = None

        elif mask_A is not None:
            
            EmptyCondGen = EmptyConditioningGenerator(model)
            conditioning_A, conditioning_B, conditioning_C = EmptyCondGen.zero_none_conditionings_([conditioning_A, conditioning_B, conditioning_C])

            conditioning = copy.deepcopy(conditioning_A)
            
            if isinstance(model.model.model_config, (comfy.supported_models.WAN21_T2V, comfy.supported_models.WAN21_I2V)):
                if model.model.diffusion_model.blocks[0].self_attn.winderz_type != "false":
                    AttnMask = CrossAttentionMask(mask_type, edge_width)
                else:
                    AttnMask = SplitAttentionMask(mask_type, edge_width)
            elif isinstance(model.model.model_config, comfy.supported_models.HiDream):
                AttnMask = FullAttentionMaskHiDream(mask_type, edge_width)
            elif isinstance(model.model.model_config, (comfy.supported_models.SDXL, comfy.supported_models.SD15, comfy.supported_models.Stable_Cascade_C)):
                AttnMask = SplitAttentionMask(mask_type, edge_width)
            else:
                AttnMask = FullAttentionMask(mask_type, edge_width)
                
            RegContext = RegionalContext()
            
            if isinstance(model.model.model_config, comfy.supported_models.HiDream):
                AttnMask.add_region_sizes(
                    [
                        conditioning_A[0][0].shape[-2],
                        conditioning_A[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        conditioning_A[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    mask_A)
                AttnMask.add_region_sizes(
                    [
                        conditioning_B[0][0].shape[-2],
                        conditioning_B[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        conditioning_B[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    mask_B)
                AttnMask.add_region_sizes(
                    [
                        conditioning_C[0][0].shape[-2],
                        conditioning_C[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        conditioning_C[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    mask_AB_inv)
                
                RegContext.add_region_llama3(conditioning_A[0][1]['conditioning_llama3'])
                RegContext.add_region_llama3(conditioning_B[0][1]['conditioning_llama3'])
                RegContext.add_region_llama3(conditioning_C[0][1]['conditioning_llama3'])
            else:
                AttnMask.add_region(conditioning_A[0][0], mask_A)
                AttnMask.add_region(conditioning_B[0][0], mask_B)
                AttnMask.add_region(conditioning_C[0][0], mask_AB_inv)
            
            RegContext.add_region(conditioning_A[0][0], conditioning_A[0][1].get('pooled_output'))
            RegContext.add_region(conditioning_B[0][0], conditioning_B[0][1].get('pooled_output'))
            RegContext.add_region(conditioning_C[0][0], conditioning_C[0][1].get('pooled_output'))
            
            #if 'pooled_output' in conditioning_A[0][1]:
            #    RegContext.pooled_output = conditioning_A[0][1]['pooled_output'] + conditioning_B[0][1]['pooled_output'] + conditioning_C[0][1]['pooled_output']
            
            conditioning[0][1]['AttnMask']   = AttnMask
            conditioning[0][1]['RegContext'] = RegContext
            
            conditioning = merge_with_base(base=conditioning, others=[conditioning_A, conditioning_B, conditioning_C])
            
            if 'pooled_output' in conditioning[0][1] and conditioning[0][1]['pooled_output'] is not None:
                conditioning[0][1]['pooled_output'] = (conditioning_A[0][1]['pooled_output'] + conditioning_B[0][1]['pooled_output'] + conditioning_C[0][1]['pooled_output']) / 3
            
        else:
            conditioning = conditioning_A

        conditioning[0][1]['RegParam'] = RegionalParameters(weights, floors)
        
        return (conditioning,)



class ClownRegionalConditioning2(ClownRegionalConditioning_AB):
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "weight":                  ("FLOAT",                                     {"default": 1.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed":            ("FLOAT",                                     {"default": 0.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed_start_step": ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "weight_scheduler":        (["constant"] + get_res4lyf_scheduler_list(), {"default": "constant"},),
                "start_step":              ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "end_step":                ("INT",                                       {"default": -1,  "min": -1,        "max": 10000}),
                "mask_type":               (REG_MASK_TYPE_2,                             {"default": "boolean"}),
                "edge_width":              ("INT",                                       {"default": 0,  "min": -10000,          "max": 10000}),
                "invert_mask":             ("BOOLEAN",                                   {"default": False}),
            }, 
            "optional": {
                "conditioning_masked":     ("CONDITIONING", ),
                "conditioning_unmasked":   ("CONDITIONING", ),
                "mask":                    ("MASK", ),
                "weights":                 ("SIGMAS", ),
                "region_bleeds":           ("SIGMAS", ),
            }
        }

    def main(self, conditioning_masked, conditioning_unmasked, mask, **kwargs):
        return super().main(
            conditioning_A = conditioning_masked,
            conditioning_B = conditioning_unmasked,
            mask_A         =   mask,
            mask_B         = 1-mask,
            **kwargs
        )    



class ClownRegionalConditioning3(ClownRegionalConditioning_ABC):
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "weight":                  ("FLOAT",                                     {"default": 1.0,  "min": -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed":            ("FLOAT",                                     {"default": 0.0,  "min": -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed_start_step": ("INT",                                       {"default": 0,    "min":  0,       "max": 10000}),
                "weight_scheduler":        (["constant"] + get_res4lyf_scheduler_list(), {"default": "constant"},),
                "start_step":              ("INT",                                       {"default": 0,    "min":  0,       "max": 10000}),
                "end_step":                ("INT",                                       {"default": 100,  "min": -1,       "max": 10000}),
                "mask_type":               (REG_MASK_TYPE_3,                             {"default": "boolean"}),
                "edge_width":              ("INT",                                       {"default": 0,    "min": 0,        "max": 10000}),
                "invert_mask":             ("BOOLEAN",                                   {"default": False}),
            },
            "optional": {
                "conditioning_A":          ("CONDITIONING", ),
                "conditioning_B":          ("CONDITIONING", ),
                "conditioning_unmasked":   ("CONDITIONING", ),
                "mask_A":                  ("MASK", ),
                "mask_B":                  ("MASK", ),
                "weights":                 ("SIGMAS", ),
                "region_bleeds":           ("SIGMAS", ),
            }
        }

    def main(self, conditioning_unmasked, mask_A, mask_B, **kwargs):
        
        mask_AB_inv = torch.ones_like(mask_A) - mask_A - mask_B
        mask_AB_inv[mask_AB_inv < 0] = 0
        
        return super().main(
            conditioning_C = conditioning_unmasked,
            mask_A         = mask_A,
            mask_B         = mask_B,
            mask_C         = mask_AB_inv,
            **kwargs
        )    




class ClownRegionalConditioning:
    @classmethod
    def INPUT_TYPES(cls):
        return {"required":
                    {
                    "spineless":    ("BOOLEAN", {"default": False}),
                    "edge_width":   ("INT",     {"default": 0,  "min": -10000,  "max": 10000}),
                    },
                "optional": 
                    {
                    "cond_regions": ("COND_REGIONS", ),
                    "conditioning": ("CONDITIONING", ),
                    "mask":         ("MASK", ),
                    }
                }

    RETURN_TYPES = ("COND_REGIONS",)
    RETURN_NAMES = ("cond_regions",) 
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"
    
    def main(self,
            spineless    = False,
            edge_width   = 0,
            cond_regions = None,
            conditioning = None,
            mask         = None,
            ):
        
        cond_reg = [] if cond_regions is None else copy.deepcopy(cond_regions)
        
        if mask is None:
            mask = torch.ones_like(cond_reg[0]['mask'])
            for i in range(len(cond_reg)):
                if mask.dtype == torch.bool:
                    mask &= cond_reg[i]['mask'].to(cond_reg[0]['mask'].dtype)
                else:
                    mask = mask - cond_reg[i]['mask'].to(cond_reg[0]['mask'].dtype)
                    mask[mask < 0] = 0.0
                    
        
        cond_reg.append(
            {
                'use_self_attn_mask': not spineless,
                'edge_width'        : edge_width,
                'conditioning'      : conditioning,
                'mask'              : mask,
            }
        )

        return (cond_reg,)





class ClownRegionalConditionings:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": { 
                "weight":                  ("FLOAT",                                     {"default": 1.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed":            ("FLOAT",                                     {"default": 0.0, "min":  -10000.0, "max": 10000.0, "step": 0.01}),
                "region_bleed_start_step": ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "weight_scheduler":        (["constant"] + get_res4lyf_scheduler_list(), {"default": "constant"},),
                "start_step":              ("INT",                                       {"default": 0,   "min":  0,        "max": 10000}),
                "end_step":                ("INT",                                       {"default": -1,  "min": -1,        "max": 10000}),
                "mask_type":               (["gradient", "boolean"],                     {"default": "boolean"}),
                "invert_masks":            ("BOOLEAN",                                   {"default": False}),
            },
            "optional": {
                "cond_regions":            ("COND_REGIONS", ),
                "weights":                 ("SIGMAS", ),
                "region_bleeds":           ("SIGMAS", ),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION     = "main"
    CATEGORY     = "RES4LYF/conditioning"

    def create_callback(self, **kwargs):
        def callback(model):
            kwargs["model"] = model  
            pos_cond, = self.prepare_regional_cond(**kwargs)
            return pos_cond
        return callback

    def main(self,
            weight                   : float  = 1.0,
            start_sigma              : float  = 0.0,
            end_sigma                : float  = 1.0,
            weight_scheduler                  = None,
            start_step               : int    = 0,
            end_step                 : int    = -1,
            cond_regions                      = None,
            weights                  : Tensor = None,
            region_bleeds            : Tensor = None,
            region_bleed             : float  = 0.0,
            region_bleed_start_step  : int    = 0,
            mask_type                : str    = "boolean",
            invert_masks             : bool   = False
            ) -> Tuple[Tensor]:
                
        if end_step == -1:
            end_step = MAX_STEPS

        callback = self.create_callback(weight                   = weight,
                                        start_sigma              = start_sigma,
                                        end_sigma                = end_sigma,
                                        weight_scheduler         = weight_scheduler,
                                        start_step               = start_step,
                                        end_step                 = end_step,
                                        weights                  = weights,
                                        region_bleeds            = region_bleeds,
                                        region_bleed             = region_bleed,
                                        region_bleed_start_step  = region_bleed_start_step,
                                        mask_type                = mask_type,
                                        invert_masks             = invert_masks,
                                        cond_regions             = cond_regions,
                                        )

        cond_list = [region['conditioning'] for region in cond_regions]
        conditioning = zero_conditioning_from_list(cond_list)
        
        conditioning[0][1]['callback_regional'] = callback
        
        return (conditioning,)



    def prepare_regional_cond(self,
                                model,
                                weight                   : float  = 1.0,
                                start_sigma              : float  = 0.0,
                                end_sigma                : float  = 1.0,
                                weight_scheduler                  = None,
                                start_step               : int    = 0,
                                end_step                 : int    = -1,
                                weights                  : Tensor = None,
                                region_bleeds            : Tensor = None,
                                region_bleed             : float  = 0.0,
                                region_bleed_start_step  : int    = 0,
                                mask_type                : str    = "gradient",
                                cond_regions                      = None,
                                invert_masks             : bool   = False,
                                ) -> Tuple[Tensor]:

        default_dtype  = torch.float64
        default_device = torch.device("cuda") 
        
        cond_list               = [region['conditioning']       for region in cond_regions]
        mask_list               = [region['mask']               for region in cond_regions]
        edge_width_list         = [region['edge_width']         for region in cond_regions]
        use_self_attn_mask_list = [region['use_self_attn_mask'] for region in cond_regions]
        if end_step == -1:
            end_step = MAX_STEPS
        if weights is None and weight_scheduler != "constant":
            total_steps = end_step - start_step
            weights     = get_sigmas(model, weight_scheduler, total_steps, 1.0).to(dtype=default_dtype, device=default_device) #/ model.inner_model.inner_model.model_sampling.sigma_max  #scaling doesn't matter as this is a flux-only node
            prepend     = torch.zeros(start_step,                                  dtype=default_dtype, device=default_device)
            weights     = torch.cat((prepend, weights), dim=0)
        
        if invert_masks:
            for i in range(len(mask_list)):
                if mask_list[i].dtype == torch.bool:
                    mask_list[i] = ~mask_list[i]
                else:
                    mask_list[i] = 1 - mask_list[i]
                    
        floor, floors = region_bleed, region_bleeds
        
        weights = initialize_or_scale(weights, weight, end_step).to(default_dtype).to(default_device)
        weights = F.pad(weights, (0, MAX_STEPS), value=0.0)
        
        prepend = torch.full((region_bleed_start_step,),  0.0, dtype=default_dtype, device=default_device)
        floors  = initialize_or_scale(floors,  floor,  end_step).to(default_dtype).to(default_device)
        floors  = F.pad(floors,  (0, MAX_STEPS), value=0.0)
        floors  = torch.cat((prepend, floors), dim=0)

        EmptyCondGen = EmptyConditioningGenerator(model)
        cond_list = EmptyCondGen.zero_none_conditionings_(cond_list)
        
        conditioning = copy.deepcopy(cond_list[0])
        
        if isinstance(model.model.model_config, comfy.supported_models.WAN21_T2V) or isinstance(model.model.model_config, comfy.supported_models.WAN21_I2V):
            if model.model.diffusion_model.blocks[0].self_attn.winderz_type != "false":
                AttnMask = CrossAttentionMask  (mask_type, edge_width_list=edge_width_list, use_self_attn_mask_list=use_self_attn_mask_list)
            else:
                AttnMask = SplitAttentionMask  (mask_type, edge_width_list=edge_width_list, use_self_attn_mask_list=use_self_attn_mask_list)
        elif isinstance(model.model.model_config, comfy.supported_models.HiDream):
            AttnMask = FullAttentionMaskHiDream(mask_type, edge_width_list=edge_width_list, use_self_attn_mask_list=use_self_attn_mask_list)
        elif isinstance(model.model.model_config, comfy.supported_models.SDXL) or isinstance(model.model.model_config, comfy.supported_models.SD15):
            AttnMask = SplitAttentionMask(mask_type, edge_width_list=edge_width_list, use_self_attn_mask_list=use_self_attn_mask_list)
        else:
            AttnMask = FullAttentionMask       (mask_type, edge_width_list=edge_width_list, use_self_attn_mask_list=use_self_attn_mask_list)

        RegContext = RegionalContext()
        
        for cond, mask in zip(cond_list, mask_list):
            if isinstance(model.model.model_config, comfy.supported_models.HiDream):
                
                AttnMask.add_region_sizes(
                    [
                        cond[0][0].shape[-2],
                        cond[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                        cond[0][1]['conditioning_llama3'][0,0,...].shape[-2],
                    ],
                    mask)

                RegContext.add_region_llama3(cond[0][1]['conditioning_llama3'])
            else:
                AttnMask.add_region(cond[0][0],   mask)
            
            RegContext.add_region(cond[0][0])
            
            if 'clip_vision_output' in cond[0][1]: # For WAN... dicey results
                RegContext.add_region_clip_fea(cond[0][1]['clip_vision_output'].penultimate_hidden_states)
            
        conditioning[0][1]['AttnMask']   = AttnMask
        conditioning[0][1]['RegContext'] = RegContext
        conditioning[0][1]['RegParam']   = RegionalParameters(weights, floors)
        
        conditioning = merge_with_base(base=conditioning, others=cond_list)
        
        if 'pooled_output' in conditioning[0][1] and conditioning[0][1]['pooled_output'] is not None:
            conditioning[0][1]['pooled_output'] = torch.stack([cond_tmp[0][1]['pooled_output'] for cond_tmp in cond_list]).mean(dim=0)

            #conditioning[0][1]['pooled_output'] = cond_list[0][0][1]['pooled_output']

        return (conditioning,)







def merge_with_base(
    base   : List[     Tuple[torch.Tensor, Dict[str, Any]]],
    others : List[List[Tuple[torch.Tensor, Dict[str, Any]]]],
    dim    : int = -2
) ->         List[     Tuple[torch.Tensor, Dict[str, Any]]]:
    """
    Merge `base` plus an arbitrary list of other conditioning objects:
        - base: zero out its tensors, for use as an accumulator
        - For each level ℓ:
            • Collect the base’s zeroed tensor + all others’ ℓ-tensors.
            • Pad them along `dim` to the same length and sum.
            • Replace merged[ℓ][0] with that sum.
        - For each tensor-valued key in the base’s info-dict at level ℓ:
            • Gather a zeroed tensor + that key from all others.
            • Pad & sum, and store back under that key.
        - Any non-tensor entries in the base’s info are preserved untouched.
    """
    max_levels = max(len(base), *(len(p) for p in others))

    for lvl in range(max_levels):
        if lvl >= len(base): # if base lacks this level, skip entirely
            continue

        # --- tokens merge ---
        base_tokens, base_info = base[lvl]
        zero_tokens = torch.zeros_like(base_tokens)
        toks = [zero_tokens]

        # zero-out any tensor fields in base_info
        for key, val in base_info.items():
            if isinstance(val, torch.Tensor):
                base_info[key] = torch.zeros_like(val)

        # collect same-level tokens from each other
        for pos in others:
            if lvl < len(pos):
                toks.append(pos[lvl][0])

        toks = pad_tensor_list_to_max_len(toks, dim=dim)
        base_tokens = sum(toks)
        base[lvl] = (base_tokens, base_info)

        # --- info-dict tensor merge ---
        for key, val in list(base_info.items()):
            if not isinstance(val, torch.Tensor):
                continue
            pieces = [val]  # zeroed base tensor
            for pos in others:
                if lvl < len(pos):
                    info_i = pos[lvl][1]
                    if key in info_i and isinstance(info_i[key], torch.Tensor):
                        pieces.append(info_i[key])
            pieces = pad_tensor_list_to_max_len(pieces, dim=dim)
            base[lvl][1][key] = sum(pieces)

    return base





def best_hw(n): # get factor pair closesst to a true square
    best = (1, n)
    min_diff = n
    for i in range(1, int(n**0.5) + 1):
        if n % i == 0:
            j = n // i
            if abs(i - j) < min_diff:
                best = (i, j)
                min_diff = abs(i - j)
    return best

def downsample_tokens(cond: torch.Tensor, target_tokens: int, mode="bicubic") -> torch.Tensor:
    B, T, D = cond.shape

    def next_square(n: int):
        root = math.ceil(n**0.5)
        return root * root

    padded_len = next_square(T)
    pad_amount = padded_len - T
    if pad_amount > 0:
        pad_tensor = torch.zeros(B, pad_amount, D, dtype=cond.dtype, device=cond.device)
        cond = torch.cat([cond, pad_tensor], dim=1)

    side_len = int(math.sqrt(padded_len))
    cond_reshaped = cond.view(B, side_len, side_len, D).permute(0, 3, 1, 2)  # [B, D, H, W]

    H_target, W_target = best_hw(target_tokens)
    cond_interp = F.interpolate(cond_reshaped, size=(H_target, W_target), mode=mode)

    cond_final = cond_interp.permute(0, 2, 3, 1).reshape(B, -1, D)
    cond_final = cond_final[:, :target_tokens, :]

    return cond_final





class CrossAttn_EraseReplace_HiDream:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "clip": ("CLIP", ),
            "t5xxl_erase":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "llama_erase":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_replace": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "llama_replace": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_erase_token":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "llama_erase_token":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_replace_token": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "llama_replace_token": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            }}
    RETURN_TYPES = ("CONDITIONING","CONDITIONING",)
    RETURN_NAMES = ("positive",    "negative",)
    FUNCTION = "encode"

    CATEGORY = "advanced/conditioning"
    EXPERIMENTAL = True

    def encode(self, clip, t5xxl_erase, llama_erase, t5xxl_replace, llama_replace, t5xxl_erase_token, llama_erase_token, t5xxl_replace_token, llama_replace_token):

        tokens_erase      = clip.tokenize("")
        tokens_erase["l"] = clip.tokenize("")["l"]
        tokens_replace      = clip.tokenize("")
        tokens_replace["l"] = clip.tokenize("")["l"]
        
        tokens_erase  ["t5xxl"] = clip.tokenize(t5xxl_erase)  ["t5xxl"]
        tokens_erase  ["llama"] = clip.tokenize(llama_erase)  ["llama"]
        tokens_replace["t5xxl"] = clip.tokenize(t5xxl_replace)["t5xxl"]
        tokens_replace["llama"] = clip.tokenize(llama_replace)["llama"]
        
        
        tokens_erase_token      = clip.tokenize("")
        tokens_erase_token["l"] = clip.tokenize("")["l"]
        tokens_replace_token      = clip.tokenize("")
        tokens_replace_token["l"] = clip.tokenize("")["l"]
        
        tokens_erase_token  ["t5xxl"] = clip.tokenize(t5xxl_erase_token)  ["t5xxl"]
        tokens_erase_token  ["llama"] = clip.tokenize(llama_erase_token)  ["llama"]
        tokens_replace_token["t5xxl"] = clip.tokenize(t5xxl_replace_token)["t5xxl"]
        tokens_replace_token["llama"] = clip.tokenize(llama_replace_token)["llama"]
        
        
        encoded_erase   = clip.encode_from_tokens_scheduled(tokens_erase)
        encoded_replace = clip.encode_from_tokens_scheduled(tokens_replace)
        
        return (encoded_replace, encoded_erase, )








class CrossAttn_EraseReplace_Flux:
    @classmethod
    def INPUT_TYPES(s):
        return {"required": {
            "clip": ("CLIP", ),
            "t5xxl_erase":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_replace": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_erase_token":   ("STRING", {"multiline": True, "dynamicPrompts": True}),
            "t5xxl_replace_token": ("STRING", {"multiline": True, "dynamicPrompts": True}),
            }}
    RETURN_TYPES = ("CONDITIONING","CONDITIONING",)
    RETURN_NAMES = ("positive",    "negative",)
    FUNCTION = "encode"

    CATEGORY = "advanced/conditioning"
    EXPERIMENTAL = True

    def encode(self, clip, t5xxl_erase, llama_erase, t5xxl_replace, llama_replace, t5xxl_erase_token, llama_erase_token, t5xxl_replace_token, llama_replace_token):

        tokens_erase      = clip.tokenize("")
        tokens_erase["l"] = clip.tokenize("")["l"]
        tokens_replace      = clip.tokenize("")
        tokens_replace["l"] = clip.tokenize("")["l"]
        
        tokens_erase  ["t5xxl"] = clip.tokenize(t5xxl_erase)  ["t5xxl"]
        tokens_erase  ["llama"] = clip.tokenize(llama_erase)  ["llama"]
        tokens_replace["t5xxl"] = clip.tokenize(t5xxl_replace)["t5xxl"]
        tokens_replace["llama"] = clip.tokenize(llama_replace)["llama"]
        
        
        tokens_erase_token      = clip.tokenize("")
        tokens_erase_token["l"] = clip.tokenize("")["l"]
        tokens_replace_token      = clip.tokenize("")
        tokens_replace_token["l"] = clip.tokenize("")["l"]
        
        tokens_erase_token  ["t5xxl"] = clip.tokenize(t5xxl_erase_token)  ["t5xxl"]
        tokens_erase_token  ["llama"] = clip.tokenize(llama_erase_token)  ["llama"]
        tokens_replace_token["t5xxl"] = clip.tokenize(t5xxl_replace_token)["t5xxl"]
        tokens_replace_token["llama"] = clip.tokenize(llama_replace_token)["llama"]
        
        
        encoded_erase   = clip.encode_from_tokens_scheduled(tokens_erase)
        encoded_replace = clip.encode_from_tokens_scheduled(tokens_replace)
        
        return (encoded_replace, encoded_erase, )