import math import comfy.model_sampling import torch from dynamicprompts.generators import RandomPromptGenerator import hashlib import chardet import pandas import re from pathlib import Path import difflib import os import json import numpy as np from torchvision.transforms import InterpolationMode import torchvision.transforms.functional as F from urllib.parse import urlparse import requests import folder_paths import comfy_extras.nodes_model_advanced as nodes_model_advanced import nodes import collections import pytorch_lightning as pl import torch.nn as nn from PIL import Image, ImageOps, ImageSequence from comfy.k_diffusion.sampling import default_noise_sampler from comfy.ldm.modules.diffusionmodules.util import make_beta_schedule from comfy.model_sampling import EPS from comfy.samplers import KSAMPLER, calculate_sigmas from comfy_extras.nodes_model_advanced import ModelSamplingDiscreteDistilled from tqdm.auto import trange import comfy.model_detection as model_detection import comfy.model_management as model_management from transformers.dynamic_module_utils import get_imports from unittest.mock import patch from transformers import AutoProcessor, AutoModelForCausalLM, AutoTokenizer from comfy.utils import common_upscale import types import comfy.conds import node_helpers here = Path(__file__).parent.parent.absolute() comfy_dir = str(here.parent.parent) cache_dir = os.path.join(here, 'Nodes', '.cache') cache_file = os.path.join(cache_dir, '.cache.json') MAX_SEED = 999999999999999 SUPPORTED_FORMATS = [".png", ".jpg", ".jpeg", ".webp", ".preview.png", ".preview.jpg", ".preview.jpeg", ] STANDARD_SIDES = np.arange(64, 4096, 16).tolist() CASCADE_SIDES = np.arange(64, 4096, 16).tolist() MAX_RESOLUTION = 8192 VALID_SHAPES = np.arange(256, 4096, 256).tolist() # PREVIEW_ROOT = os.path.join(comfy_dir, "web", "extensions", "PrimerePreviews", "images") PREVIEW_ROOT = os.path.join(here, 'front_end', "images") KEYWORD_SELECTOR_VALUES = ["None", "Select in order", "Random select"] SUPPORTED_MODELS = ["SD1", "SD2", "SDXL", "Illustrious", "SD3", "StableCascade", "Chroma", "Z-Image", "Turbo", "Flux", "Nunchaku", "QwenGen", "QwenEdit", "WanImg", "KwaiKolors", "Hunyuan", "Playground", "Pony", "LCM", "Lightning", "Hyper", "PixartSigma", "SANA1024", "SANA512", "AuraFlow", "HiDream", "Mochi", "WanT2V", "WanI2V", "Cosmos", "Flux2", "SSD", "SegmindVega", "KOALA", "StableZero", "SV3D", "SD09", "StableAudio"] CONCEPT_RESOLUTIONS = { "512": ['SD09', 'SD1', "Turbo", "SANA512", "WanT2V", "WanI2V"], "768": ['SD2', "LCM"], "1024": ["SDXL", "Illustrious", "SD3", "StableCascade", "Chroma", "Z-Image", "Flux", "KwaiKolors", "Hunyuan", "Playground", "Pony", "Lightning", "Hyper", "PixartSigma", "SANA1024", "WanImg", "AuraFlow", "HiDream", "Mochi", "Cosmos", "Flux2"], "1328": ["QwenGen", "QwenEdit"], "1280": ["Nunchaku"] } PREVIEW_PATH_BY_TYPE = { "Checkpoint": os.path.join(PREVIEW_ROOT, "checkpoints"), "CSV Prompt": os.path.join(PREVIEW_ROOT, "styles"), "Lora": os.path.join(PREVIEW_ROOT, "loras"), "Lycoris": os.path.join(PREVIEW_ROOT, "lycoris"), "Hypernetwork": os.path.join(PREVIEW_ROOT, "hypernetworks"), "Embedding": os.path.join(PREVIEW_ROOT, "embeddings"), } WORKFLOW_SORT_LIST = ['exif_status', 'exif_data_count', 'meta_source', 'pic2story', 'positive', 'positive_l', 'positive_r', 'negative', 'negative_l', 'negative_r', 'prompt_state', 'decoded_positive', 'decoded_negative', 'pic2story_positive', 'model', 'model_concept', 'model_version', 'is_sdxl', 'model_hash', 'vae', 'vae_hash', 'vae_name_sd', 'vae_name_sdxl', 'sampler', 'scheduler', 'steps', 'cfg', 'seed', 'width', 'height', 'size_string', 'preferred', 'saved_image_width', 'saved_image_heigth', 'upscaler_ratio', 'vae_name_sd', 'vae_name_sdxl', 'concept_data'] vae_loader_class = nodes.VAELoader() def merge_str_to_tuple(item1, item2): if not isinstance(item1, tuple): item1 = (item1,) if not isinstance(item2, tuple): item2 = (item2,) return item1 + item2 def merge_dict(dict1, dict2): dict3 = dict1.copy() for k, v in dict2.items(): dict3[k] = merge_str_to_tuple(v, dict3[k]) if k in dict3 else v return dict3 def remove_quotes(string): return str(string).replace('"', "").replace("'", "") def add_quotes(string): return '"' + str(string) + '"' def get_square_shape(shape_a, shape_b): area = shape_a * shape_b square = math.sqrt(area) standard_square = min(VALID_SHAPES, key=lambda x: abs(square - x)) return standard_square def get_dimensions_by_shape(self, rationame: str, square: int, orientation: str = 'Vertical', round_to_standard: bool = False, calculate_by_custom: bool = False, custom_side_a: float = 1, custom_side_b: float = 1, standard: str = 'STANDARD'): def calculate_dim(ratio_1: float, ratio_2: float, square: int): FullPixels = square ** 2 ratio = ratio_2 / ratio_1 side_a = math.sqrt(FullPixels * ratio) side_b = side_a / ratio if round_to_standard == True: STANDARD_LIST = STANDARD_SIDES if standard == 'CASCADE': STANDARD_LIST = CASCADE_SIDES side_a = min(STANDARD_LIST, key=lambda x: abs(side_a - x)) side_b = round(FullPixels / side_a) side_b = min(STANDARD_LIST, key=lambda x: abs(x - side_b)) side_a = round(side_a) side_b = round(side_b) return sorted([side_a, side_b], reverse=True) if (calculate_by_custom == True and isinstance(custom_side_a, (int, float)) and isinstance(custom_side_b, (int, float)) and custom_side_a >= 1 and custom_side_b >= 1): ratio_x = custom_side_a ratio_y = custom_side_b else: RatioLabel = self.ratioNames[rationame] ratio_x = self.sd_ratios[RatioLabel]['side_x'] ratio_y = self.sd_ratios[RatioLabel]['side_y'] dimensions = calculate_dim(ratio_x, ratio_y, square) if (orientation == 'Vertical'): dimensions = sorted(dimensions) if (orientation == 'Vertical'): dimensions = dimensions + [ratio_y, ratio_x] else: dimensions = dimensions + [ratio_x, ratio_y] return dimensions def clear_prompt(NETWORK_START, NETWORK_END, promptstring, modelname=False): promptstring = promptstring + ' ' promptstring_temp = promptstring for LABEL in NETWORK_START: if LABEL in promptstring: LabelStartIndexes = [n for n in range(len(promptstring)) if promptstring.find(LABEL, n) == n] for LabelStartIndex in LabelStartIndexes: Matches = [] for endString in NETWORK_END: Match = promptstring.find(endString, (LabelStartIndex + 1)) if (Match > 0): Matches.append(Match) if len(Matches) > 0: LabelEndIndex = sorted(Matches)[0] MatchedString = promptstring[LabelStartIndex:(LabelEndIndex + 1)] if len(MatchedString) > 0 and (modelname == False or (modelname != False and modelname in MatchedString)): if '<' in MatchedString: endString = '>' Match = promptstring.find(endString, (LabelStartIndex + 1)) if (Match > 0): LabelEndIndex = Match MatchedString = promptstring[LabelStartIndex:(LabelEndIndex + 1)] promptstring_temp = promptstring_temp.replace(MatchedString, "") if '{' in MatchedString: endString = '}' Match = promptstring.find(endString, (LabelStartIndex + 1)) if (Match > 0): LabelEndIndex = Match MatchedString = promptstring[LabelStartIndex:(LabelEndIndex + 1)] promptstring_temp = promptstring_temp.replace(MatchedString, "") if ')' in MatchedString: MatchedString = promptstring[(LabelStartIndex - 1):(LabelEndIndex + 1)] promptstring_temp = promptstring_temp.replace(MatchedString, "") promptstring_temp = promptstring_temp.replace(MatchedString, "") return promptstring_temp.replace('()', '').replace(' , ,', ',').replace('||', '').replace('{,', '').replace(' ', ' ').replace(', ,', ',').strip(', ') def DiT_cleaner(prompt, length=0): cleanPrompt = re.sub("(:\d+\.\d+)|(:\d+)|[()]|BREAK|break", "", prompt).replace(' ', ' ') if length > 0: cleanPrompt = cleanPrompt[:length].rsplit(' ', 1)[0] return cleanPrompt def get_networks_prompt(NETWORK_START, NETWORK_END, promptstring): promptstring = promptstring + ' ' valid_networks = [] for LABEL in NETWORK_START: if LABEL in promptstring: LabelStartIndexes = [n for n in range(len(promptstring)) if promptstring.find(LABEL, n) == n] for LabelStartIndex in LabelStartIndexes: Matches = [] for endString in NETWORK_END: Match = promptstring.find(endString, (LabelStartIndex + 1)) if (Match > 0): Matches.append(Match) if len(Matches) > 0: LabelEndIndex = sorted(Matches)[0] MatchedString = promptstring[(LabelStartIndex + len(LABEL)):(LabelEndIndex)] if len(MatchedString) > 0: networkdata = MatchedString.split(":") if len(networkdata) == 1: networkdata.append('1') if LABEL == ' 0 and 'resolution' in header_json_dataset[0]: dataset_resolution = header_json_dataset[0]['resolution'] if len(dataset_resolution) == 2: model_version_res = (int(dataset_resolution[0]) * int(dataset_resolution[1])) if model_version_res < 1000 * 1000: model_version = 'SD1' else: model_version = 'SDXL' else: if '__metadata__' in header_json and 'modelspec.architecture' in header_json['__metadata__']: model_version = header_json['__metadata__']['modelspec.architecture'] else: return False except: try: sd = comfy.utils.load_torch_file(ckpt_path) diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix) model_version = type(model_config).__name__ except: return False else: try: sd = comfy.utils.load_torch_file(ckpt_path) diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix) model_version = type(model_config).__name__ except: return False else: return False if model_version is not None and model_version != 'NoneType' and (model_type == 'checkpoints' or model_type == 'loras' or model_type == 'lycoris'): model_version = model_version.replace("_", "").replace(".", "").replace("sdv1", "SD1") res = [ele for ele in SUPPORTED_MODELS if (ele.lower() in model_version.lower())] if len(res) > 0: model_version = res[0] return model_version def get_model_hash(filename): is_link = os.path.islink(str(filename)) if is_link == True: filename = Path(str(filename)).resolve() try: with open(filename, "rb") as file: m = hashlib.sha256() file.seek(0x100000) m.update(file.read(0x10000)) hash = m.hexdigest()[0:8] return hash except FileNotFoundError: return None def get_file_encoding(filename): try: # fileTest = open(filename, 'rb').readline() with open(filename, 'rb') as tempf: fileTest = tempf.readline() result = chardet.detect(fileTest) ENCODING = result['encoding'] if ENCODING == 'ascii': ENCODING = 'UTF-8' return ENCODING except FileNotFoundError: return None def open_file_by_chardet(filename): try: # fileTest = open(filename, 'rb').readline() with open(filename, 'rb') as tempf: fileTest = tempf.readline() result = chardet.detect(fileTest) ENCODING = result['encoding'] if ENCODING == 'ascii': ENCODING = 'UTF-8' with open(filename, "r", newline='', encoding=ENCODING) as the_file: return the_file.readlines() except FileNotFoundError: return None def load_external_csv(csv_full_path: str, header_cols: int): # fileTest = open(csv_full_path, 'rb').readline() with open(csv_full_path, 'rb') as tempf: fileTest = tempf.readline() result = chardet.detect(fileTest) ENCODING = result['encoding'] if ENCODING == 'ascii': ENCODING = 'UTF-8' with open(csv_full_path, "r", newline='', encoding=ENCODING) as csv_file: try: return pandas.read_csv(csv_file, header=header_cols, index_col=False, skipinitialspace=True) except pandas.errors.ParserError as e: errorstring = repr(e) matchre = re.compile('Expected (d+) fields in line (d+), saw (d+)') (expected, line, saw) = map(int, matchre.search(errorstring).groups()) print(f'Error at line {line}. Fields added : {saw - expected}.') return None def get_model_keywords(filename, modelhash, model_name): keywords = load_external_csv(filename, 3) if keywords is not None: selected_kv = keywords[keywords['#model_hash'] == modelhash]['keyword'].values if (len(selected_kv) > 1): selected_ckpt = keywords[keywords['#model_hash'] == modelhash]['filename.ckpt'].values basename = Path(model_name).stem cutoff_list = list(np.around(np.arange(0.1, 1.05, 0.05).tolist(), 2))[::-1] is_found = [] model_name_kw = None for trycut in cutoff_list: is_found = difflib.get_close_matches(basename, selected_ckpt, cutoff=trycut) if len(is_found) >= 1: model_name_kw = is_found[0] break if len(is_found) >= 0: if model_name_kw is not None: selected_kv = keywords[keywords['filename.ckpt'] == model_name_kw]['keyword'].values if (len(selected_kv) > 0): return selected_kv[0] else: return None else: return None def get_closest_element(value, netlist): cutoff_list = list(np.around(np.arange(0.1, 1.05, 0.05).tolist(), 2))[::-1] is_found = None for trycut in cutoff_list: is_found = difflib.get_close_matches(value, netlist, cutoff=trycut) if len(is_found) >= 1: return is_found[0] return is_found def get_category_from_cache(category): ifCacheExist = os.path.isfile(cache_file) if ifCacheExist == True: with open(cache_file, 'r') as openfile: try: saved_cache = json.load(openfile) try: return saved_cache[category] except Exception: return None except ValueError as e: return None else: return None def get_value_from_cache(category, key): ifCacheExist = os.path.isfile(cache_file) if ifCacheExist == True: with open(cache_file, 'r') as openfile: try: saved_cache = json.load(openfile) try: return saved_cache[category][key] except Exception: return None except ValueError as e: return None else: return None def update_value_in_cache(category, key, value): cacheData = {category: {key: value}} json_object = json.dumps(cacheData, indent=4) ifCacheExist = os.path.isfile(cache_file) if ifCacheExist == True: with open(cache_file, 'r') as openfile: try: saved_cache = json.load(openfile) if category in saved_cache and key in saved_cache[category]: saved_cache[category][key] = value else: saved_cache.update(cacheData) newJsonObject = json.dumps(saved_cache, indent=4) with open(cache_file, "w", encoding='utf-8') as outfile: outfile.write(newJsonObject) return True except ValueError as e: return None else: with open(cache_file, "w", encoding='utf-8') as outfile: outfile.write(json_object) return True def add_value_to_cache(category, key, value): cacheData = {category: {key: value}} json_object = json.dumps(cacheData, indent=4) ifCacheExist = os.path.isfile(cache_file) if ifCacheExist == True: with open(cache_file, 'r') as openfile: try: saved_cache = json.load(openfile) if category in saved_cache: saved_cache[category][key] = value else: saved_cache.update(cacheData) newJsonObject = json.dumps(saved_cache, indent=4) with open(cache_file, "w", encoding='utf-8') as outfile: outfile.write(newJsonObject) return True except ValueError as e: return False else: with open(cache_file, "w", encoding='utf-8') as outfile: outfile.write(json_object) return True def getLoraVersion(modelobject): VersionKeysBlock = [ 'lora_unet_down_blocks_0_attentions_0_proj_in.lora_up.weight', 'lora_unet_down_blocks_2_attentions_0_transformer_blocks_0_attn1_to_k.lora_up.weight', 'lora_te_text_model_encoder_layers_0_mlp_fc1.lora_up.weight', 'lora_unet_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_0_proj.lora_up.weight', 'lora_unet_input_blocks_4_1_transformer_blocks_0_ff_net_0_proj.lora_up.weight', 'lora_unet_input_blocks_4_1_transformer_blocks_0_ff_net_0_proj.hada_w1_a', 'lora_unet_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_0_proj.hada_w1_a', ] VersionHelper = 'Unknown' for index, value in modelobject.items(): try: type(value.item()).__name__ except Exception: if index in VersionKeysBlock: VersionHelper = len(value) break return VersionHelper def pil2numpy(image: Image.Image): return np.array(image).astype(np.float32) / 255.0 def numpy2pil(image: np.ndarray, mode=None): return Image.fromarray(np.clip(255.0 * image, 0, 255).astype(np.uint8), mode) def pil2tensor(image: Image.Image): return torch.from_numpy(pil2numpy(image)).unsqueeze(0) def tensor2pil(image: torch.Tensor, mode=None): return numpy2pil(image.cpu().numpy().squeeze(), mode=mode) def comfyimg2numpyarray(image): return np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) def image_scale_down(images, width, height, crop): if crop == "center": old_width = images.shape[2] old_height = images.shape[1] old_aspect = old_width / old_height new_aspect = width / height x = 0 y = 0 if old_aspect > new_aspect: x = round((old_width - old_width * (new_aspect / old_aspect)) / 2) elif old_aspect < new_aspect: y = round((old_height - old_height * (old_aspect / new_aspect)) / 2) s = images[:, y: old_height - y, x: old_width - x, :] else: s = images results = [] for image in s: img = tensor2pil(image).convert("RGB") img = img.resize((width, height), Image.LANCZOS) results.append(pil2tensor(img)) return (torch.cat(results, dim=0),) def image_scale_down_by_scale(images, scale_by): width = images.shape[2] height = images.shape[1] new_width = int(width * scale_by) new_height = int(height * scale_by) return image_scale_down(images, new_width, new_height, "center") def image_scale_down_by_dim(images, new_width, new_height): return image_scale_down(images, new_width, new_height, "center") def img_resizer(image: torch.Tensor, width: int, height: int, interpolation_mode: str): assert isinstance(image, torch.Tensor) assert isinstance(height, int) assert isinstance(width, int) assert isinstance(interpolation_mode, str) interpolation_mode = interpolation_mode.upper().replace(" ", "_") interpolation_mode = getattr(InterpolationMode, interpolation_mode) image = image.permute(0, 3, 1, 2) image = F.resize(image, (height, width), interpolation=interpolation_mode, antialias=True) image = image.permute(0, 2, 3, 1) return image def apply_variation_noise(latent_image, noise_device, variation_seed, variation_strength, mask=None): latent_size = latent_image.size() latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]] if noise_device == "cpu": variation_generator = torch.manual_seed(variation_seed) else: torch.cuda.manual_seed(variation_seed) variation_generator = None variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, generator=variation_generator, device=noise_device) variation_noise = variation_latent.expand(latent_image.size()[0], -1, -1, -1) if variation_strength == 0: return latent_image elif mask is None: result = (1 - variation_strength) * latent_image + variation_strength * variation_noise else: result = (mask == 1).float() * ((1 - variation_strength) * latent_image + variation_strength * variation_noise * mask) + (mask == 0).float() * latent_image return result def prepare_noise(latent_image, seed, noise_inds=None, noise_device="cpu", incremental_seed_mode="comfy", variation_seed=None, variation_strength=None): latent_size = latent_image.size() latent_size_1batch = [1, latent_size[1], latent_size[2], latent_size[3]] if variation_strength is not None and variation_strength > 0 or incremental_seed_mode.startswith("variation str inc"): if noise_device == "cpu": variation_generator = torch.manual_seed(variation_seed) else: torch.cuda.manual_seed(variation_seed) variation_generator = None variation_latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, generator=variation_generator, device=noise_device) else: variation_latent = None def apply_variation(input_latent, strength_up=None): if variation_latent is None: return input_latent else: strength = variation_strength if strength_up is not None: strength += strength_up variation_noise = variation_latent.expand(input_latent.size()[0], -1, -1, -1) result = (1 - strength) * input_latent + strength * variation_noise return result # method: incremental seed batch noise if noise_inds is None and incremental_seed_mode == "incremental": batch_cnt = latent_size[0] latents = None for i in range(batch_cnt): if noise_device == "cpu": generator = torch.manual_seed(seed + i) else: torch.cuda.manual_seed(seed + i) generator = None latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device=noise_device) latent = apply_variation(latent) if latents is None: latents = latent else: latents = torch.cat((latents, latent), dim=0) return latents # method: incremental variation batch noise elif noise_inds is None and incremental_seed_mode.startswith("variation str inc"): batch_cnt = latent_size[0] latents = None for i in range(batch_cnt): if noise_device == "cpu": generator = torch.manual_seed(seed) else: torch.cuda.manual_seed(seed) generator = None latent = torch.randn(latent_size_1batch, dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device=noise_device) step = float(incremental_seed_mode[18:]) latent = apply_variation(latent, step * i) if latents is None: latents = latent else: latents = torch.cat((latents, latent), dim=0) return latents # method: comfy batch noise if noise_device == "cpu": generator = torch.manual_seed(seed) else: torch.cuda.manual_seed(seed) generator = None if noise_inds is None: latents = torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device=noise_device) latents = apply_variation(latents) return latents unique_inds, inverse = np.unique(noise_inds, return_inverse=True) noises = [] for i in range(unique_inds[-1] + 1): noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device=noise_device) if i in unique_inds: noises.append(noise) noises = [noises[i] for i in inverse] noises = torch.cat(noises, axis=0) return noises def hf_downloader(repo_id, model_local_dir): from huggingface_hub import snapshot_download model_path = f"{model_local_dir}/{repo_id.split('/')[-1]}" snapshot_download(repo_id=repo_id, local_dir=model_path, local_dir_use_symlinks=True, max_workers=1) return model_path def ModelConceptNames(ckpt_name, model_concept, lightning_selector, lightning_model_step, hypersd_selector, hypersd_model_step, model_version='SDXL'): lora_name = None unet_name = None lightningModeValid = False hyperModeValid = False LoraList = getDownloadedFiles() if model_concept == 'Lightning': if lightning_selector == 'SAFETENSOR': allCheckpoints = folder_paths.get_filename_list("checkpoints") allLightning = list(filter(lambda a: 'sdxl_lightning_'.casefold() in a.casefold(), allCheckpoints)) if len(allLightning) > 0: finalLightning = list(filter(lambda a: str(lightning_model_step) + 'step'.casefold() in a.casefold(), allLightning)) if len(finalLightning) > 0: lightningModeValid = True ckpt_name = finalLightning[0] if lightning_selector == 'LORA': # LoraList = folder_paths.get_filename_list("loras") if len(LoraList) > 0: allLoraLightning = list(filter(lambda a: 'sdxl_lightning_'.casefold() in a.casefold(), LoraList)) if len(allLoraLightning) > 0: finalLightning = list(filter(lambda a: str(lightning_model_step) + 'step'.casefold() in a.casefold(), allLoraLightning)) if len(finalLightning) > 0: lightningModeValid = True lora_name = finalLightning[0] if lightning_selector == 'UNET': UnetList = folder_paths.get_filename_list("unet") if len(UnetList) > 0: allUnetLightning = list(filter(lambda a: 'sdxl_lightning_'.casefold() in a.casefold(), UnetList)) if len(allUnetLightning) > 0: finalLightning = list(filter(lambda a: str(lightning_model_step) + 'step'.casefold() in a.casefold(), allUnetLightning)) if len(finalLightning) > 0: lightningModeValid = True unet_name = finalLightning[0] if model_concept == 'Hyper': if hypersd_selector == 'LORA': if len(LoraList) > 0: if model_version == 'SDXL': allLoraHyper = list(filter(lambda a: 'Hyper-SDXL-'.casefold() in a.casefold(), LoraList)) else: allLoraHyper = list(filter(lambda a: 'Hyper-SD15-'.casefold() in a.casefold(), LoraList)) if len(allLoraHyper) > 0: pluralString = '' if hypersd_model_step > 1: pluralString = 's' finalHyper = list(filter(lambda a: str(hypersd_model_step) + 'step' + pluralString + '-lora'.casefold() in a.casefold() or str(hypersd_model_step) + 'step' + pluralString + '-CFG-lora'.casefold() in a.casefold(), allLoraHyper)) if len(finalHyper) > 0: hyperModeValid = True lora_name = finalHyper[0] if hypersd_selector == 'UNET': UnetList = folder_paths.get_filename_list("unet") if len(UnetList) > 0: allUnetHyper = list(filter(lambda a: 'Hyper-SDXL-1step-Unet-Comfyui'.casefold() in a.casefold(), UnetList)) if len(allUnetHyper) > 0: hyperModeValid = True unet_name = allUnetHyper[0] return {'ckpt_name': ckpt_name, 'lora_name': lora_name, 'unet_name': unet_name, 'lightningModeValid': lightningModeValid, 'hyperModeValid': hyperModeValid} def BDanceConceptHelper(self, model_concept, lightningModeValid, lightning_selector, lightning_model_step, OUTPUT_MODEL, lora_name, unet_name, ckpt_name, lora_model_strength=1): if model_concept == 'Lightning' and lightningModeValid == True and lightning_selector == 'LORA' and lora_name is not None: if lora_model_strength != 0: lora = None if self.loaded_lora is not None: if self.loaded_lora[0] == lora_name: lora = self.loaded_lora[1] else: temp = self.loaded_lora self.loaded_lora = None del temp if lora is None: lora = comfy.utils.load_torch_file(lora_name, safe_load=True) self.loaded_lora = (lora_name, lora) OUTPUT_MODEL = comfy.sd.load_lora_for_models(OUTPUT_MODEL, None, lora, lora_model_strength, 0)[0] if model_concept == 'Lightning' and lightningModeValid == True and lightning_selector == 'UNET' and unet_name is not None: OUTPUT_MODEL = nodes.UNETLoader.load_unet(self, unet_name, "default")[0] if model_concept == 'Lightning' and lightningModeValid == True and lightning_selector == 'SAFETENSOR' and ckpt_name is not None: OUTPUT_MODEL = nodes.CheckpointLoaderSimple.load_checkpoint(self, ckpt_name)[0] if model_concept == 'Lightning' and lightning_model_step == 1 and lightningModeValid == True: OUTPUT_MODEL = nodes_model_advanced.ModelSamplingDiscrete.patch(self, OUTPUT_MODEL, "x0", False)[0] if model_concept == 'Hyper' and lightningModeValid == True and lightning_selector == 'LORA' and lora_name is not None: if lora_model_strength != 0: lora = None if self.loaded_lora is not None: if self.loaded_lora[0] == lora_name: lora = self.loaded_lora[1] else: temp = self.loaded_lora self.loaded_lora = None del temp if lora is None: lora = comfy.utils.load_torch_file(lora_name, safe_load=True) self.loaded_lora = (lora_name, lora) OUTPUT_MODEL = comfy.sd.load_lora_for_models(OUTPUT_MODEL, None, lora, lora_model_strength, 0)[0] if model_concept == 'Hyper' and lightningModeValid == True and lightning_selector == 'UNET' and unet_name is not None: unet_path = folder_paths.get_full_path("unet", unet_name) if unet_path is None and os.path.isfile(unet_name): unet_path = unet_name OUTPUT_MODEL = comfy.sd.load_checkpoint_guess_config(unet_path) return OUTPUT_MODEL def get_hypersd_sigmas(model): timesteps = torch.tensor([800]) sigmas = model.model.model_sampling.sigma(timesteps) sigmas = torch.cat([sigmas, sigmas.new_zeros([1])]) return (sigmas,) class ModelSamplingDiscreteDistilledTCD(ModelSamplingDiscreteDistilled, EPS): def __init__(self, model_config=None): super().__init__(model_config) sampling_settings = model_config.sampling_settings if model_config is not None else {} beta_schedule = sampling_settings.get("beta_schedule", "linear") linear_start = sampling_settings.get("linear_start", 0.00085) linear_end = sampling_settings.get("linear_end", 0.012) betas = make_beta_schedule( beta_schedule, n_timestep=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3 ) alphas = 1.0 - betas alphas_cumprod = torch.cumprod(alphas, dim=0, dtype=torch.float32) self.register_buffer("alphas_cumprod", alphas_cumprod.clone().detach()) @torch.no_grad() def sample_tcd(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, eta=0.3, alpha_prod_s: torch.Tensor = None): extra_args = {} if extra_args is None else extra_args noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler s_in = x.new_ones([x.shape[0]]) beta_prod_s = 1 - alpha_prod_s for i in trange(len(sigmas) - 1, disable=disable): denoised = model(x, sigmas[i] * s_in, **extra_args) eps = (x - denoised) / sigmas[i] denoised = alpha_prod_s[i + 1].sqrt() * denoised + beta_prod_s[i + 1].sqrt() * eps if callback is not None: callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised}) x = denoised if eta > 0 and sigmas[i + 1] > 0: noise = noise_sampler(sigmas[i], sigmas[i + 1]) x = x / alpha_prod_s[i + 1].sqrt() + noise * (sigmas[i + 1] ** 2 + 1 - 1 / alpha_prod_s[i + 1]).sqrt() return x def TCDModelSamplingDiscrete(self, model, steps=4, scheduler="simple", denoise=1.0, eta=0.3): m = model.clone() ms = ModelSamplingDiscreteDistilledTCD(model.model.model_config) total_steps = steps if denoise <= 0.0: # raise error ? sigmas = torch.FloatTensor([]) elif denoise <= 1.0: total_steps = int(steps / denoise) sigmas = calculate_sigmas(ms, scheduler, total_steps).cpu() sigmas = sigmas[-(steps + 1):] m.add_object_patch("model_sampling", ms) timesteps_s = torch.floor((1 - eta) * ms.timestep(sigmas)).to(dtype=torch.long).detach() timesteps_s[-1] = 0 alpha_prod_s = ms.alphas_cumprod[timesteps_s] sampler = KSAMPLER(sample_tcd, extra_options={"eta": eta, "alpha_prod_s": alpha_prod_s}, inpaint_options={}) return (m, sampler, sigmas) def getLatentSize(samples): for tensor in samples['samples'][0]: if isinstance(tensor, torch.Tensor): shape = tensor.shape tensor_height = shape[-2] tensor_width = shape[-1] return (tensor_width, tensor_height) else: return (None, None) return (None, None) def MatchDimensions(width_1, height_1, width_2, height_2, axis_value): if axis_value == 1: rate = height_2 / height_1 new_heigth = height_1 new_width = width_2 / rate else: rate = width_2 / width_1 new_width = width_1 new_heigth = height_2 / rate return [round(new_width), round(new_heigth)] def ImageConcat(image1, image2, axis_value): image1_size = image1.size image2_size = image2.size new_image_dim = MatchDimensions(image1_size[0], image1_size[1], image2_size[0], image2_size[1], axis_value) img2_resized = image2.resize(new_image_dim) if axis_value == 1: new_image = Image.new('RGB', (new_image_dim[0] + image1_size[0], image1_size[1]), (250, 250, 250)) new_image.paste(image1, (0, 0)) new_image.paste(img2_resized, (image1_size[0], 0)) else: new_image = Image.new('RGB', (image1_size[0], new_image_dim[1] + image1_size[1]), (250, 250, 250)) new_image.paste(image1, (0, 0)) new_image.paste(img2_resized, (0, image1_size[1])) return new_image def getDataFromWorkflowById(workflow, nodeName, dataIndex): result = None for NODE_ITEMS in workflow: if 'type' in NODE_ITEMS: ITEM_TYPE = NODE_ITEMS['type'] if ITEM_TYPE == nodeName: if 'widgets_values' in NODE_ITEMS: ITEM_VALUES = NODE_ITEMS['widgets_values'] if len(ITEM_VALUES) >= dataIndex + 1: result = ITEM_VALUES[dataIndex] return result def getDataFromWorkflowByName(workflow, nodeName, inputName, prompt): results = None for node in workflow: node_id = None name = node["type"] if "properties" in node: if "Node name for S&R" in node["properties"]: name = node["properties"]["Node name for S&R"] if name == nodeName: node_id = node["id"] else: if "title" in node: name = node["title"] if name == nodeName: node_id = node["id"] if node_id is None: continue if str(node_id) in prompt: values = prompt[str(node_id)] if "inputs" in values and inputName in values["inputs"]: v = values["inputs"][inputName] if type(v).__name__ == 'list': values = prompt[v[0]] if "inputs" in values: v_list = list(values["inputs"].values()) return v_list[v[1]] return v return results def getInputsFromWorkflowByNode(workflow, nodeName, prompt): filtered = {} for node in workflow: node_id = None name = node["type"] if "properties" in node: if "Node name for S&R" in node["properties"]: name = node["properties"]["Node name for S&R"] if name == nodeName: node_id = node["id"] else: if "title" in node: name = node["title"] if name == nodeName: node_id = node["id"] if node_id is None: continue if str(node_id) in prompt: values = prompt[str(node_id)]["inputs"] filtered = {k: v for k, v in values.items() if not isinstance(v, list)} return filtered def collect_state(extra_pnginfo, prompt): workflow = extra_pnginfo["workflow"] results = {} if "links" in workflow: results["__links"] = workflow["links"] for node in workflow["nodes"]: node_id = str(node["id"]) name = node["type"] if "Debug" in name or "Show" in name or "Function" in name or "Evaluate" in name: continue if "widgets_values" in node and "inputs" not in node: results[node_id] = node["widgets_values"] elif node_id in prompt: values = prompt[node_id] if "inputs" in values: results[node_id] = {} for widget in values["inputs"].items(): (n, v) = widget if type(v) is not str and isinstance(v, collections.abc.Sequence): continue results[node_id][n] = v elif "widgets_values" in node: results[node_id] = node["widgets_values"] result = json.dumps(results, sort_keys=True) return hashlib.sha256(result.encode()).hexdigest() class MLP(pl.LightningModule): def __init__(self, input_size, xcol='emb', ycol='avg_rating'): super().__init__() self.input_size = input_size self.xcol = xcol self.ycol = ycol self.layers = nn.Sequential( nn.Linear(self.input_size, 1024), nn.Dropout(0.2), nn.Linear(1024, 128), nn.Dropout(0.2), nn.Linear(128, 64), nn.Dropout(0.1), nn.Linear(64, 16), nn.Linear(16, 1) ) def forward(self, x): return self.layers(x) def training_step(self, batch, batch_idx): x = batch[self.xcol] y = batch[self.ycol].reshape(-1, 1) x_hat = self.layers(x) loss = F.mse_loss(x_hat, y) return loss def validation_step(self, batch, batch_idx): x = batch[self.xcol] y = batch[self.ycol].reshape(-1, 1) x_hat = self.layers(x) loss = F.mse_loss(x_hat, y) return loss def configure_optimizers(self): optimizer = torch.optim.Adam(self.parameters(), lr=1e-3) return optimizer def normalized(a, axis=-1, order=2): import numpy as np # pylint: disable=import-outside-toplevel l2 = np.atleast_1d(np.linalg.norm(a, order, axis)) l2[l2 == 0] = 1 return a / np.expand_dims(l2, axis) def ImageLoaderFromPath(ImgPath, new_width=None, new_height=None): output_image = None if Path(ImgPath).is_file() == True: loaded_img = Image.open(ImgPath) output_images = [] for i in ImageSequence.Iterator(loaded_img): i = ImageOps.exif_transpose(i) if i.mode == 'I': i = i.point(lambda i: i * (1 / 255)) image = i.convert("RGB") if new_width is not None and new_height is not None: newsize = (new_width, new_height) image = image.resize(newsize) image = np.array(image).astype(np.float32) / 255.0 image = torch.from_numpy(image)[None,] output_images.append(image) if len(output_images) > 1: output_image = torch.cat(output_images, dim=0) else: output_image = output_images[0] return output_image def tensor_to_image(tensor): tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def florence_img2prompt(model, processor, image, max_new_tokens, num_beams, do_sample, text_input=None, llm_options=None): if text_input is None: text_input = 'detailed enhanced prompt for text2image models' task_prompt = 'detailed enhanced prompt for text2image models' inputs = processor(text=text_input, images=image, return_tensors="pt").to("cuda") default_settings = { "do_sample": do_sample, "temperature": 0.9, "top_k": 8, "max_new_tokens": max_new_tokens, "num_return_sequences": 1, "repetition_penalty": 1.2, "penalty_alpha": 0.6, "no_repeat_ngram_size": 1, "early_stopping": False, "top_p": 0.4, "num_beams": num_beams, } if llm_options is not None: if 'max_length' in llm_options: llm_options['max_new_tokens'] = llm_options['max_length'] del llm_options['max_length'] variant_params = llm_options settings = {**default_settings, **variant_params} else: settings = {**default_settings} print('--------------------------') print(settings) print('--------------------------') generated_ids = model.generate( input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], **settings ) generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height)) return parsed_answer def fixed_get_imports(filename) -> list[str]: """Workaround for FlashAttention""" if os.path.basename(filename) != "modeling_florence2.py": return get_imports(filename) imports = get_imports(filename) try: imports.remove("flash_attn") except: pass return imports def load_florence2_model(repo_id): model_path = Path(repo_id) attention = 'sdpa' device = torch.cuda.current_device() try: with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=torch.float32, trust_remote_code=True) processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) except Exception as e: try: model = AutoModelForCausalLM.from_pretrained(model_path, attn_implementation=attention, device_map=device, torch_dtype=torch.float32, trust_remote_code=True) processor = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) except Exception as e: sys.path.append(model_path) # Import the Florence modules if 'large-PromptGen-v1.5' in repo_id: from florence2_large.modeling_florence2 import Florence2ForConditionalGeneration from florence2_large.configuration_florence2 import Florence2Config elif 'base-PromptGen-v1.5' in repo_id: from florence2_base_ft.modeling_florence2 import Florence2ForConditionalGeneration from florence2_base_ft.configuration_florence2 import Florence2Config else: # log(f"Error loading model or tokenizer: {str(e)}", message_type='error') return (None, None) # Load the model configuration model_config = Florence2Config.from_pretrained(model_path) # Load the model with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports): model = Florence2ForConditionalGeneration.from_pretrained( model_path, config=model_config, attn_implementation=attention, device_map=device ).to(device) processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True) return (model.to(device), processor, device) def Pic2Story(repo_id, img, prompts, special_tokens_skip=True, clean_same_result=True, llm_options=None): story_out = None if "florence2" in repo_id: model = None processor = None version = None model, processor, device = load_florence2_model(repo_id) version = version florence2_model = {'model': model, 'processor': processor, 'version': version, 'device': device} model = florence2_model['model'] processor = florence2_model['processor'] pil_image = tensor2pil(img) results = florence_img2prompt(model, processor, pil_image, 512, 3, False, prompts[0], llm_options) story_out = next(iter(results.values()), None) story_out = story_out.removeprefix(prompts[0].lower()).strip().removeprefix("of").removeprefix("is") else: from transformers import BlipProcessor, BlipForConditionalGeneration import torch os.environ['TRANSFORMERS_OFFLINE'] = "1" processor = BlipProcessor.from_pretrained(repo_id) pil_image = tensor_to_image(img) try: model = BlipForConditionalGeneration.from_pretrained(repo_id, torch_dtype=torch.float16).to("cuda") if type(prompts) == str: inputs = processor(pil_image, prompts, return_tensors="pt").to("cuda", torch.float16) out = model.generate(**inputs) story_out = processor.decode(out[0], skip_special_tokens=special_tokens_skip) story_out = story_out.removeprefix(prompts.lower()).strip().removeprefix("of").removeprefix("is") elif type(prompts).__name__ == 'list': for prompt in prompts: inputs = processor(pil_image, prompt, return_tensors="pt").to("cuda", torch.float16) out = model.generate(**inputs) Processed = processor.decode(out[0], skip_special_tokens=special_tokens_skip) Processed = Processed.removeprefix(prompt.lower()).strip().removeprefix("of").removeprefix("is") + ', ' if story_out is not None: story_out = story_out + Processed else: story_out = Processed except Exception: print('Pic2Story Float 16 failed') if type(story_out) != str: try: model = BlipForConditionalGeneration.from_pretrained(repo_id).to("cuda") if type(prompts) == str: inputs = processor(pil_image, prompts, return_tensors="pt").to("cuda") out = model.generate(**inputs) story_out = processor.decode(out[0], skip_special_tokens=special_tokens_skip) story_out = story_out.removeprefix(prompts.lower()).strip().removeprefix("of").removeprefix("is") elif type(prompts).__name__ == 'list': for prompt in prompts: inputs = processor(pil_image, prompt, return_tensors="pt").to("cuda") out = model.generate(**inputs) Processed = processor.decode(out[0], skip_special_tokens=special_tokens_skip) Processed = Processed.removeprefix(prompt.lower()).strip().removeprefix("of").removeprefix("is") + ', ' if story_out is not None: story_out = story_out + Processed else: story_out = Processed except Exception: print('Pic2Story GPU failed') if type(story_out) != str: try: model = BlipForConditionalGeneration.from_pretrained(repo_id) if type(prompts) == str: inputs = processor(pil_image, prompts, return_tensors="pt") out = model.generate(**inputs) story_out = processor.decode(out[0], skip_special_tokens=special_tokens_skip) story_out = story_out.removeprefix(prompts.lower()).strip().removeprefix("of").removeprefix("is") elif type(prompts).__name__ == 'list': for prompt in prompts: inputs = processor(pil_image, prompt, return_tensors="pt") out = model.generate(**inputs) Processed = processor.decode(out[0], skip_special_tokens=special_tokens_skip) Processed = Processed.removeprefix(prompt.lower()).strip().removeprefix("of").removeprefix("is") + ', ' if story_out is not None: story_out = story_out + Processed else: story_out = Processed except Exception: print('Pic2Story CPU failed') if type(story_out) == str: if clean_same_result == True: story_out = ' '.join(dict.fromkeys(story_out.split())) return story_out.rstrip(', ').replace(' and ', ' ').replace(' an ', ' ').replace(' is ', ' ').replace(' are ', ' ').strip() else: return story_out.rstrip(', ').strip() def getDownloadedFiles(): DOWNLOAD_DIR = os.path.join(here, 'Nodes', 'Downloads') folder_paths.add_model_folder_path("primere_downloads", DOWNLOAD_DIR) downloaded_filelist = folder_paths.get_filename_list("primere_downloads") downloaded_filelist_filtered = folder_paths.filter_files_extensions(downloaded_filelist, ['.ckpt', '.safetensors']) return downloaded_filelist_filtered def downloader(from_url, to_path): if os.path.isfile(to_path) == False: pathparser = urlparse(from_url) TargetFilename = os.path.basename(pathparser.path) print('Downloading: ' + TargetFilename) Request = requests.get(from_url, allow_redirects=True) if Request.status_code == 200 and Request.ok == True: open(to_path, 'wb').write(Request.content) print('DOWNLOADED: ' + to_path) return True else: print('ERROR: Cannot download ' + TargetFilename) return False def fileDownloader(targetFILE, sourceURL): if os.path.exists(targetFILE) == False: print('Downloading from: ' + sourceURL + ' to: ' + str(targetFILE)) reqsdlcm = requests.get(sourceURL, allow_redirects=True) if reqsdlcm.status_code == 200 and reqsdlcm.ok == True: open(targetFILE, 'wb').write(reqsdlcm.content) return True else: print('ERROR: Cannot dowload required file to: ' + str(targetFILE)) return False return True def get_dtype_by_name(dtype): if dtype == 'Auto': try: if model_management.should_use_fp16(): dtype = torch.float16 elif model_management.should_use_bf16(): dtype = torch.bfloat16 else: dtype = torch.float32 except: raise AttributeError("ComfyUI version too old, can't autodetect properly. Set your dtypes manually.") elif dtype == "fp16": dtype = torch.float16 elif dtype == "bf16": dtype = torch.bfloat16 elif dtype == "fp32": dtype = torch.float32 elif dtype == "fp8_e4m3fn": dtype = torch.float8_e4m3fn elif dtype == "fp8_e4m3fnuz": dtype = torch.float8_e4m3fnuz elif dtype == "fp8_e5m2": dtype = torch.float8_e5m2 elif dtype == "fp8_e5m2fnuz": dtype = torch.float8_e5m2fnuz return dtype def getValidAscorerPaths(model_root): valid_ae_path = [] if os.path.exists(model_root): allsubdirs = list(os.listdir(Path(model_root))) for subdir in allsubdirs: path_config = os.path.join(model_root, subdir, 'config.json') path_model_bin = os.path.join(model_root, subdir, 'pytorch_model.bin') if os.path.exists(path_config) == True and os.path.exists(path_model_bin) == True: valid_ae_path.append(subdir) return valid_ae_path def edit_encoder(clip, prompt, vae, images=None): images_vl = [] ref_latents = [] llama_template = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, characters, style, mood, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n" image_prompt = "" for i, image in enumerate(images): if image is not None: samples = image.movedim(-1, 1) total = int(384 * 384) scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) width = round(samples.shape[3] * scale_by) height = round(samples.shape[2] * scale_by) s = comfy.utils.common_upscale(samples, width, height, "area", "disabled") images_vl.append(s.movedim(1, -1)) if vae is not None: total = int(1024 * 1024) scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) width = round(samples.shape[3] * scale_by / 8.0) * 8 height = round(samples.shape[2] * scale_by / 8.0) * 8 s = comfy.utils.common_upscale(samples, width, height, "area", "disabled") ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3])) image_prompt += "Picture {}: <|vision_start|><|image_pad|><|vision_end|>".format(i + 1) tokens = clip.tokenize(image_prompt + prompt, images=images_vl, llama_template=llama_template) conditioning = clip.encode_from_tokens_scheduled(tokens) if len(ref_latents) > 0: conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": ref_latents}, append=True) return conditioning def json2tuple(json_path): data_tuple = None ifConfigExist = os.path.isfile(json_path) if ifConfigExist: with open(json_path) as f: try: data_tuple = json.load(f) except ValueError as e: data_tuple = None return data_tuple