|
|
import base64 |
|
|
import io |
|
|
|
|
|
from PIL import Image |
|
|
|
|
|
import numpy as np |
|
|
import simpleeval |
|
|
import torch |
|
|
|
|
|
from ..utils import VyroParams |
|
|
|
|
|
class VyroParamUpdater: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return { |
|
|
"required": { |
|
|
"vyro_params": ("VYRO_PARAMS",), |
|
|
}, |
|
|
"optional": { |
|
|
"latents": ("LATENT",), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("VYRO_PARAMS",) |
|
|
RETURN_NAMES = ("vyro_params",) |
|
|
FUNCTION = "update" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def update(self, vyro_params:VyroParams, latents=None): |
|
|
if latents is not None: |
|
|
vyro_params.latents = latents |
|
|
return (vyro_params,) |
|
|
|
|
|
class VyroParamExtractor: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return { |
|
|
"required": { |
|
|
"vyro_params": ("VYRO_PARAMS",), |
|
|
} |
|
|
} |
|
|
RETURN_TYPES = ("LATENT","STRING","STRING","STRING","FLOAT","INT","INT","INT","INT","INT","FLOAT","FLOAT","FLOAT","FLOAT","STRING","STRING","STRING","BOOLEAN") |
|
|
RETURN_NAMES = ("latents","user_prompt","user_neg_prompt","mode","cfg","batch_size","steps","width","height","seed","denoise","stage1_strength","stage2_strength","efficiency_multiplier","style","final_positive_prompt","final_negative_prompt","is_raw") |
|
|
FUNCTION = "input" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def input(self, vyro_params:VyroParams): |
|
|
|
|
|
return (vyro_params.latents, vyro_params.user_prompt, vyro_params.user_neg_prompt, vyro_params.mode, vyro_params.cfg, vyro_params.batch_size, vyro_params.steps, vyro_params.width, vyro_params.height, vyro_params.seed, vyro_params.denoise, vyro_params.stage1_strength, vyro_params.stage2_strength, vyro_params.efficiency_multiplier, vyro_params.style, vyro_params.final_positive_prompt, vyro_params.final_negative_prompt, vyro_params.is_raw) |
|
|
|
|
|
class VyroModelStreamInput: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"model": ("MODEL",), |
|
|
"clip": ("CLIP",), |
|
|
}, |
|
|
"optional": { |
|
|
"refiner": ("MODEL",), |
|
|
"refiner_clip": ("CLIP",), |
|
|
}} |
|
|
|
|
|
RETURN_TYPES = ("*",) |
|
|
RETURN_NAMES = ("output",) |
|
|
FUNCTION = "input" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def input(self, model, clip, refiner=None, refiner_clip=None): |
|
|
if refiner is None: |
|
|
return {"model": model, "clip": clip, "refiner": None, "refiner_clip": None}, |
|
|
else: |
|
|
return {"model": model, "clip": clip, "refiner": refiner, "refiner_clip": refiner_clip}, |
|
|
|
|
|
class VyroModelStreamOutput: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return {"required": { |
|
|
"input": ("*",), |
|
|
}} |
|
|
RETURN_TYPES = ("MODEL","CLIP","MODEL","CLIP") |
|
|
RETURN_NAMES = ("model","clip","refiner_model","refiner_clip") |
|
|
FUNCTION = "output" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def output(self, input): |
|
|
if input["refiner"] is None: |
|
|
return input["model"], input["clip"], None, None |
|
|
else: |
|
|
return input["model"], input["clip"], input["refiner"], input["refiner_clip"] |
|
|
|
|
|
class VyroStyleSwitcher: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
input_types = {"required": { |
|
|
"style": ("STYLE", ), |
|
|
"default": ("*",) |
|
|
}, |
|
|
"optional": {} |
|
|
} |
|
|
for style in VyroParams.STYLES: |
|
|
input_types["optional"][style] = ("*",) |
|
|
return input_types |
|
|
|
|
|
RETURN_TYPES = ("*",) |
|
|
RETURN_NAMES = ("output",) |
|
|
FUNCTION = "switch" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def switch(self, style, default, **kwargs): |
|
|
return (kwargs.get(style, default), ) |
|
|
|
|
|
|
|
|
class VyroModeLatentMuxer: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
types = { |
|
|
"required": { |
|
|
"vyro_params": ("VYRO_PARAMS",), |
|
|
} |
|
|
} |
|
|
for mode in VyroParams.MODE: |
|
|
types["required"][mode] = ("LATENT",) |
|
|
return types |
|
|
|
|
|
RETURN_TYPES = ("LATENT",) |
|
|
RETURN_NAMES = ("latents",) |
|
|
FUNCTION = "mux_latents" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def mux_latents(self, vyro_params:VyroParams, **kwargs): |
|
|
mode = vyro_params.mode |
|
|
return (kwargs[mode],) |
|
|
|
|
|
def tensor2pil(image): |
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) |
|
|
|
|
|
class VyroImageToString: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
return { |
|
|
"required": { |
|
|
"image": ("IMAGE",), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("STRING",) |
|
|
RETURN_NAMES = ("string",) |
|
|
FUNCTION = "image_to_string" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def image_to_string(self, image): |
|
|
batch_size = image.shape[0] |
|
|
if batch_size > 1: |
|
|
split = torch.split(image, 1, dim=0) |
|
|
strs = [] |
|
|
for i in range(len(split)): |
|
|
image = split[i] |
|
|
image = image.squeeze(0) |
|
|
img = tensor2pil(image) |
|
|
|
|
|
buffered = io.BytesIO() |
|
|
img.save(buffered, format="PNG") |
|
|
img_str = base64.b64encode(buffered.getvalue()) |
|
|
strs.append(img_str.decode("utf-8")) |
|
|
return (strs,) |
|
|
|
|
|
image = image.squeeze(0) |
|
|
img = tensor2pil(image) |
|
|
|
|
|
buffered = io.BytesIO() |
|
|
img.save(buffered, format="PNG") |
|
|
img_str = base64.b64encode(buffered.getvalue()) |
|
|
|
|
|
return (img_str.decode("utf-8"),) |
|
|
|
|
|
class VyroModeFilter: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
|
|
|
req = { |
|
|
"required": { |
|
|
"vyro_params": ("VYRO_PARAMS",), |
|
|
} |
|
|
} |
|
|
for mode in VyroParams.MODE: |
|
|
req["required"][f'{mode}'] = (VyroParams.ALLOWED, {"default": VyroParams.ALLOWED[0]}) |
|
|
return req |
|
|
|
|
|
RETURN_TYPES = ("VYRO_PARAMS",) |
|
|
RETURN_NAMES = ("vyro_params",) |
|
|
FUNCTION = "filter" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def filter(self, vyro_params:VyroParams, **kwargs): |
|
|
mode = vyro_params.mode |
|
|
|
|
|
state = kwargs[mode] |
|
|
if state == VyroParams.ALLOWED[1]: |
|
|
raise ValueError("If you are receiving this error, it's because you're trying to execute the workflow in Comfy without detaching the preview nodes for the inactive modes. If you are reciving this error in the API, you are selecting the wrong output node.") |
|
|
|
|
|
return (vyro_params,) |
|
|
|
|
|
class VyroModelSwitcher: |
|
|
@classmethod |
|
|
def INPUT_TYPES(s): |
|
|
|
|
|
return { |
|
|
"required": { |
|
|
"vyro_params": ("VYRO_PARAMS",), |
|
|
"a": ("MODEL",), |
|
|
"b": ("MODEL",), |
|
|
"return_a_if_true": ("STRING", {"default": "face_swap_img is not None", "multiline": False}), |
|
|
} |
|
|
} |
|
|
|
|
|
RETURN_TYPES = ("MODEL",) |
|
|
FUNCTION = "switch" |
|
|
CATEGORY = "Vyro/Routing" |
|
|
|
|
|
def switch(self, vyro_params:VyroParams, a, b, return_a_if_true): |
|
|
names = {} |
|
|
for n in VyroParams.PARAMS: |
|
|
names[n] = getattr(vyro_params, n) |
|
|
|
|
|
result = simpleeval.simple_eval(return_a_if_true, names=names) |
|
|
if result: |
|
|
return (a,) |
|
|
else: |
|
|
return (b,) |
|
|
|
|
|
NODE_CLASS_MAPPINGS = { |
|
|
"Vyro Mode Latent Muxer": VyroModeLatentMuxer, |
|
|
"Vyro Style Switcher": VyroStyleSwitcher, |
|
|
"Vyro Model Stream Input": VyroModelStreamInput, |
|
|
"Vyro Model Stream Output": VyroModelStreamOutput, |
|
|
"Vyro Param Extractor": VyroParamExtractor, |
|
|
"Vyro Image to String": VyroImageToString, |
|
|
"Vyro Mode Filter": VyroModeFilter, |
|
|
"Vyro Model Switcher": VyroModelSwitcher, |
|
|
"Vyro Param Updater": VyroParamUpdater, |
|
|
} |
|
|
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = { |
|
|
"VyroModeLatentMuxer": "Vyro Mode Latent Muxer", |
|
|
"VyroStyleSwitcher": "Vyro Style Switcher", |
|
|
"VyroModelStreamInput": "Vyro Model Stream Input", |
|
|
"VyroModelStreamOutput": "Vyro Model Stream Output", |
|
|
"VyroParamExtractor": "Vyro Param Extractor", |
|
|
"VyroImageToString": "Vyro Image to String", |
|
|
"VyroModeFilter": "Vyro Mode Filter", |
|
|
"VyroModelSwitcher": "Vyro Model Switcher", |
|
|
"VyroParamUpdater": "Vyro Param Updater", |
|
|
} |
