73 lines
3.1 KiB
Python
73 lines
3.1 KiB
Python
# Please support 311_code on https://ko-fi.com/311_code
|
|
|
|
import torch
|
|
from nodes import MAX_RESOLUTION
|
|
|
|
class CLIPTextEncodeSDXL_311_code:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
|
|
"height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
|
|
"text_g": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
"clip": ("CLIP", ),
|
|
"text_l": ("STRING", {"multiline": True, "dynamicPrompts": True}),
|
|
"clip": ("CLIP", ),
|
|
"clip_g_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
|
|
"clip_l_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
|
|
"size_cond_factor": ("INT", {"default": 4, "min": 1, "max": 16}),
|
|
"layer_idx": ("INT", {"default": -2, "min": -33, "max": 33}), # Added layer_idx
|
|
}}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "essentials"
|
|
|
|
def execute(self, clip, width, height, size_cond_factor, text_g, text_l, clip_g_strength, clip_l_strength, layer_idx):
|
|
width = width * size_cond_factor
|
|
height = height * size_cond_factor
|
|
|
|
cond_combined = None
|
|
pooled_combined = None
|
|
|
|
# Modify CLIP layers based on `layer_idx`
|
|
clip = clip.clone() # Clone the clip model to ensure it's modifiable
|
|
clip.clip_layer(layer_idx) # Set the stopping point for the CLIP model based on layer_idx
|
|
|
|
# Tokenize both texts, or use empty tokens if both are empty
|
|
tokens_g = clip.tokenize(text_g) if text_g.strip() else clip.tokenize("")
|
|
tokens_l = clip.tokenize(text_l) if text_l.strip() else clip.tokenize("")
|
|
|
|
# Ensure equal length by padding with empty tokens if needed
|
|
if len(tokens_l["l"]) != len(tokens_g["g"]):
|
|
empty_g = clip.tokenize("")["g"]
|
|
empty_l = clip.tokenize("")["l"]
|
|
while len(tokens_l["l"]) < len(tokens_g["g"]):
|
|
tokens_l["l"] += empty_l
|
|
while len(tokens_g["g"]) < len(tokens_l["l"]):
|
|
tokens_g["g"] += empty_g
|
|
|
|
# Encode tokens
|
|
tokens = {"g": tokens_g["g"], "l": tokens_l["l"]}
|
|
cond_combined, pooled_combined = clip.encode_from_tokens(tokens, return_pooled=True)
|
|
|
|
# Apply strengths to g and l if needed
|
|
if cond_combined.dim() == 3:
|
|
if text_g.strip(): # Only apply strength to g if text_g is not empty
|
|
cond_combined[:, :, :] *= clip_g_strength
|
|
if text_l.strip(): # Only apply strength to l if text_l is not empty
|
|
cond_combined[:, :, :] *= clip_l_strength
|
|
|
|
# Return the conditioning even if both texts were empty
|
|
return ([[cond_combined, {"pooled_output": pooled_combined, "width": width, "height": height}]], )
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CLIPTextEncodeSDXL_311_code": CLIPTextEncodeSDXL_311_code
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CLIPTextEncodeSDXL_311_code": "311_code SDXL Clip Text Encode"
|
|
}
|
|
|