106 lines
3.5 KiB
Python
106 lines
3.5 KiB
Python
import os
|
|
|
|
import folder_paths
|
|
|
|
from ..utils.hash import calc_hash
|
|
from ..utils.embedding import get_embedding_file_path
|
|
|
|
from comfy.sd1_clip import escape_important, token_weights, unescape_important
|
|
from comfy.sd1_clip import SD1Tokenizer
|
|
from comfy.text_encoders.sd2_clip import SD2Tokenizer
|
|
from comfy.text_encoders.sd3_clip import SD3Tokenizer
|
|
from comfy.text_encoders.flux import FluxTokenizer
|
|
from comfy.sdxl_clip import SDXLTokenizer
|
|
|
|
cache_model_hash = {}
|
|
|
|
|
|
def calc_model_hash(model_name, input_data):
|
|
filename = folder_paths.get_full_path("checkpoints", model_name)
|
|
return calc_hash(filename)
|
|
|
|
|
|
def calc_vae_hash(model_name, input_data):
|
|
filename = folder_paths.get_full_path("vae", model_name)
|
|
return calc_hash(filename)
|
|
|
|
|
|
def calc_lora_hash(model_name, input_data):
|
|
filename = folder_paths.get_full_path("loras", model_name)
|
|
return calc_hash(filename)
|
|
|
|
|
|
def calc_unet_hash(model_name, input_data):
|
|
filename = folder_paths.get_full_path("unet", model_name)
|
|
return calc_hash(filename)
|
|
|
|
|
|
def convert_skip_clip(stop_at_clip_layer, input_data):
|
|
return stop_at_clip_layer * -1
|
|
|
|
|
|
def get_scaled_width(scaled_by, input_data):
|
|
samples = input_data[0]["samples"][0]["samples"]
|
|
return round(samples.shape[3] * scaled_by * 8)
|
|
|
|
|
|
def get_scaled_height(scaled_by, input_data):
|
|
samples = input_data[0]["samples"][0]["samples"]
|
|
return round(samples.shape[2] * scaled_by * 8)
|
|
|
|
|
|
def extract_embedding_names(text, input_data):
|
|
embedding_names, _ = _extract_embedding_names(text, input_data)
|
|
|
|
return [os.path.basename(embedding_name) for embedding_name in embedding_names]
|
|
|
|
|
|
def extract_embedding_hashes(text, input_data):
|
|
embedding_names, clip = _extract_embedding_names(text, input_data)
|
|
embedding_hashes = []
|
|
for embedding_name in embedding_names:
|
|
embedding_file_path = get_embedding_file_path(embedding_name, clip)
|
|
embedding_hashes.append(calc_hash(embedding_file_path))
|
|
|
|
return embedding_hashes
|
|
|
|
|
|
def _extract_embedding_names(text, input_data):
|
|
embedding_identifier = "embedding:"
|
|
clip_ = input_data[0]["clip"][0]
|
|
clip = None
|
|
if clip_ is not None:
|
|
tokenizer = clip_.tokenizer
|
|
if isinstance(tokenizer, SD1Tokenizer):
|
|
clip = tokenizer.clip_l
|
|
elif isinstance(tokenizer, SD2Tokenizer):
|
|
clip = tokenizer.clip_h
|
|
elif isinstance(tokenizer, SDXLTokenizer):
|
|
clip = tokenizer.clip_l
|
|
elif isinstance(tokenizer, SD3Tokenizer):
|
|
clip = tokenizer.clip_l
|
|
elif isinstance(tokenizer, FluxTokenizer):
|
|
clip = tokenizer.clip_l
|
|
if clip is not None and hasattr(clip, "embedding_identifier"):
|
|
embedding_identifier = clip.embedding_identifier
|
|
if not isinstance(text, str):
|
|
text = "".join(str(item) if item is not None else "" for item in text)
|
|
text = escape_important(text)
|
|
parsed_weights = token_weights(text, 1.0)
|
|
|
|
# tokenize words
|
|
embedding_names = []
|
|
for weighted_segment, weight in parsed_weights:
|
|
to_tokenize = unescape_important(weighted_segment).replace("\n", " ").split(" ")
|
|
to_tokenize = [x for x in to_tokenize if x != ""]
|
|
for word in to_tokenize:
|
|
# find an embedding, deal with the embedding
|
|
if (
|
|
word.startswith(embedding_identifier)
|
|
and clip.embedding_directory is not None
|
|
):
|
|
embedding_name = word[len(embedding_identifier) :].strip("\n")
|
|
embedding_names.append(embedding_name)
|
|
|
|
return embedding_names, clip
|