Spaces:
Running
on
Zero
Running
on
Zero
| #---------------------------------------------------------------------------------------------------------------------# | |
| # Comfyroll Studio custom nodes by RockOfFire and Akatsuzi https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes | |
| # for ComfyUI https://github.com/comfyanonymous/ComfyUI | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| import torch | |
| import numpy as np | |
| from PIL import Image, ImageEnhance | |
| import os | |
| import sys | |
| import folder_paths | |
| from nodes import MAX_RESOLUTION, ControlNetApply | |
| from ..categories import icons | |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # SDXL Nodes | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| class CR_PromptMixPresets: | |
| def __init__(self): | |
| pass | |
| def INPUT_TYPES(s): | |
| return { | |
| "required":{ | |
| }, | |
| "optional":{ | |
| "prompt_positive": ("STRING", {"multiline": True, "default": "prompt_pos"}), | |
| "prompt_negative": ("STRING", {"multiline": True, "default": "prompt_neg"}), | |
| "style_positive": ("STRING", {"multiline": True, "default": "style_pos"}), | |
| "style_negative": ("STRING", {"multiline": True, "default": "style_neg"}), | |
| "preset": (["default with no style text", "default with style text", | |
| "style boost 1", "style boost 2", "style text to refiner"],), | |
| }, | |
| } | |
| RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING", ) | |
| RETURN_NAMES = ("pos_g", "pos_l", "pos_r", "neg_g", "neg_l", "neg_r", "show_help", ) | |
| FUNCTION = "mixer" | |
| CATEGORY = icons.get("Comfyroll/SDXL") | |
| def mixer(self, prompt_positive, prompt_negative, style_positive, style_negative, preset): | |
| if preset == "default with no style text": | |
| pos_g = prompt_positive | |
| pos_l = prompt_positive | |
| pos_r = prompt_positive | |
| neg_g = prompt_negative | |
| neg_l = prompt_negative | |
| neg_r = prompt_negative | |
| elif preset == "default with style text": | |
| pos_g = prompt_positive + style_positive | |
| pos_l = prompt_positive + style_positive | |
| pos_r = prompt_positive + style_positive | |
| neg_g = prompt_negative + style_negative | |
| neg_l = prompt_negative + style_negative | |
| neg_r = prompt_negative + style_negative | |
| elif preset == "style boost 1": | |
| pos_g = prompt_positive | |
| pos_l = style_positive | |
| pos_r = prompt_positive | |
| neg_g = prompt_negative | |
| neg_l = style_negative | |
| neg_r = prompt_negative | |
| elif preset == "style boost 2": | |
| pos_g = style_positive | |
| pos_l = prompt_positive | |
| pos_r = style_positive | |
| neg_g = style_negative | |
| neg_l = prompt_negative | |
| neg_r = style_negative | |
| elif preset == "style text to refiner": | |
| pos_g = prompt_positive | |
| pos_l = prompt_positive | |
| pos_r = style_positive | |
| neg_g = prompt_negative | |
| neg_l = prompt_negative | |
| neg_r = style_negative | |
| show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/SDXL-Nodes#cr-sdxl-prompt-mix-presets" | |
| return (pos_g, pos_l, pos_r, neg_g, neg_l, neg_r, show_help, ) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| class CR_SDXLStyleText: | |
| def INPUT_TYPES(s): | |
| return {"required": { | |
| "positive_style": ("STRING", {"default": "POS_STYLE", "multiline": True}), | |
| "negative_style": ("STRING", {"default": "NEG_STYLE", "multiline": True}), | |
| }, | |
| } | |
| RETURN_TYPES = ("STRING", "STRING", "STRING", ) | |
| RETURN_NAMES = ("positive_prompt_text_l", "negative_prompt_text_l" , "show_help", ) | |
| FUNCTION = "get_value" | |
| CATEGORY = icons.get("Comfyroll/SDXL") | |
| def get_value(self, positive_style, negative_style): | |
| show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/SDXL-Nodes#cr-sdxl-style-text" | |
| return (positive_style, negative_style, show_help, ) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| class CR_SDXLBasePromptEncoder: | |
| def INPUT_TYPES(s): | |
| return {"required": { | |
| "base_clip": ("CLIP", ), | |
| "pos_g": ("STRING", {"multiline": True, "default": "POS_G"}), | |
| "pos_l": ("STRING", {"multiline": True, "default": "POS_L"}), | |
| "neg_g": ("STRING", {"multiline": True, "default": "NEG_G"}), | |
| "neg_l": ("STRING", {"multiline": True, "default": "NEG_L"}), | |
| "preset": (["preset A", "preset B", "preset C"],), | |
| "base_width": ("INT", {"default": 4096.0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| "base_height": ("INT", {"default": 4096.0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| "target_width": ("INT", {"default": 4096.0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| "target_height": ("INT", {"default": 4096.0, "min": 0, "max": MAX_RESOLUTION, "step": 64}), | |
| }, | |
| } | |
| RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "STRING", ) | |
| RETURN_NAMES = ("base_positive", "base_negative", "show_help", ) | |
| FUNCTION = "encode" | |
| CATEGORY = icons.get("Comfyroll/SDXL") | |
| def encode(self, base_clip, pos_g, pos_l, neg_g, neg_l, base_width, base_height, crop_w, crop_h, target_width, target_height, preset,): | |
| empty = base_clip.tokenize("") | |
| # positive prompt | |
| tokens1 = base_clip.tokenize(pos_g) | |
| tokens1["l"] = base_clip.tokenize(pos_l)["l"] | |
| if len(tokens1["l"]) != len(tokens1["g"]): | |
| while len(tokens1["l"]) < len(tokens1["g"]): | |
| tokens1["l"] += empty["l"] | |
| while len(tokens1["l"]) > len(tokens1["g"]): | |
| tokens1["g"] += empty["g"] | |
| cond1, pooled1 = base_clip.encode_from_tokens(tokens1, return_pooled=True) | |
| res1 = [[cond1, {"pooled_output": pooled1, "width": base_width, "height": base_height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]] | |
| # negative prompt | |
| tokens2 = base_clip.tokenize(neg_g) | |
| tokens2["l"] = base_clip.tokenize(neg_l)["l"] | |
| if len(tokens2["l"]) != len(tokens2["g"]): | |
| while len(tokens2["l"]) < len(tokens2["g"]): | |
| tokens2["l"] += empty["l"] | |
| while len(tokens2["l"]) > len(tokens2["g"]): | |
| tokens2["g"] += empty["g"] | |
| cond2, pooled2 = base_clip.encode_from_tokens(tokens2, return_pooled=True) | |
| res2 = [[cond2, {"pooled_output": pooled2, "width": base_width, "height": base_height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]] | |
| # positive style | |
| tokens2 = base_clip.tokenize(pos_l) | |
| tokens2["l"] = base_clip.tokenize(neg_l)["l"] | |
| if len(tokens2["l"]) != len(tokens2["g"]): | |
| while len(tokens2["l"]) < len(tokens2["g"]): | |
| tokens2["l"] += empty["l"] | |
| while len(tokens2["l"]) > len(tokens2["g"]): | |
| tokens2["g"] += empty["g"] | |
| cond2, pooled2 = base_clip.encode_from_tokens(tokens2, return_pooled=True) | |
| res3 = [[cond2, {"pooled_output": pooled2, "width": base_width, "height": base_height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]] | |
| # negative style | |
| tokens2 = base_clip.tokenize(neg_l) | |
| tokens2["l"] = base_clip.tokenize(neg_l)["l"] | |
| if len(tokens2["l"]) != len(tokens2["g"]): | |
| while len(tokens2["l"]) < len(tokens2["g"]): | |
| tokens2["l"] += empty["l"] | |
| while len(tokens2["l"]) > len(tokens2["g"]): | |
| tokens2["g"] += empty["g"] | |
| cond2, pooled2 = base_clip.encode_from_tokens(tokens2, return_pooled=True) | |
| res4 = [[cond2, {"pooled_output": pooled2, "width": base_width, "height": base_height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]] | |
| if preset == "preset A": | |
| base_positive = res1 | |
| base_negative = res2 | |
| elif preset == "preset B": | |
| base_positive = res3 | |
| base_negative = res4 | |
| elif preset == "preset C": | |
| base_positive = res1 + res3 | |
| base_negative = res2 + res4 | |
| show_help = "https://github.com/Suzie1/ComfyUI_Comfyroll_CustomNodes/wiki/SDXL-Nodes#cr-sdxl-base-prompt-encoder" | |
| return (base_positive, base_negative, show_help, ) | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # MAPPINGS | |
| #---------------------------------------------------------------------------------------------------------------------# | |
| # For reference only, actual mappings are in __init__.py | |
| ''' | |
| NODE_CLASS_MAPPINGS = { | |
| "CR SDXL Style Text":CR_SDXLStyleText, | |
| "CR SDXL Base Prompt Encoder":CR_SDXLBasePromptEncoder, | |
| "CR SDXL Prompt Mix Presets":CR_PromptMixPresets, | |
| } | |
| ''' | |