89 lines
2.8 KiB
Python
89 lines
2.8 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.sd2_clip import SD2Tokenizer
|
|
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 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["samples"][0]["samples"]
|
|
return round(samples.shape[3] * scaled_by * 8)
|
|
|
|
|
|
def get_scaled_height(scaled_by, input_data):
|
|
samples = input_data["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, tokenizer = _extract_embedding_names(text, input_data)
|
|
embedding_hashes = []
|
|
for embedding_name in embedding_names:
|
|
embedding_file_path = get_embedding_file_path(embedding_name, tokenizer)
|
|
embedding_hashes.append(calc_hash(embedding_file_path))
|
|
|
|
return embedding_hashes
|
|
|
|
|
|
def _extract_embedding_names(text, input_data):
|
|
clip = input_data["clip"][0]
|
|
tokenizer = clip.tokenizer
|
|
if isinstance(tokenizer, SD1Tokenizer):
|
|
tokenizer = tokenizer.clip_l
|
|
elif isinstance(tokenizer, SD2Tokenizer):
|
|
tokenizer = tokenizer.clip_h
|
|
elif isinstance(tokenizer, SDXLTokenizer):
|
|
tokenizer = tokenizer.clip_l
|
|
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:
|
|
# if we find an embedding, deal with the embedding
|
|
if (
|
|
word.startswith(tokenizer.embedding_identifier)
|
|
and tokenizer.embedding_directory is not None
|
|
):
|
|
embedding_name = word[len(tokenizer.embedding_identifier) :].strip("\n")
|
|
|
|
embedding_names.append(embedding_name)
|
|
|
|
return embedding_names, tokenizer
|