414 lines
10 KiB
Python
414 lines
10 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
import random
|
|
from PIL import Image, ImageOps, ImageSequence
|
|
import numpy as np
|
|
import torch
|
|
import node_helpers
|
|
from .rice_prompt_info import RicePromptInfo
|
|
from nodes import LoadImage
|
|
import requests
|
|
from .utils import pil2tensor
|
|
|
|
|
|
class _BasicTypes(str):
|
|
basic_types = ["STRING"]
|
|
|
|
def __eq__(self, other):
|
|
return other in self.basic_types or isinstance(other, (list, _BasicTypes))
|
|
|
|
def __ne__(self, other):
|
|
return not self.__eq__(other)
|
|
|
|
|
|
BasicTypes = _BasicTypes("BASIC")
|
|
|
|
|
|
class RiceRoundSimpleChoiceNode:
|
|
def __init__(self):
|
|
self.prompt_info = RicePromptInfo()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {"default": "Parameter"}),
|
|
"default": ("STRING", {"default": ""}),
|
|
},
|
|
"optional": {},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (BasicTypes,)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "placeholder"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def placeholder(self, name, default, **kwargs):
|
|
unique_id = int(kwargs.pop("unique_id", 0))
|
|
prompt = kwargs.pop("prompt", None)
|
|
need_wait = True
|
|
if prompt:
|
|
for _, node in prompt.items():
|
|
if node.get("class_type", "") == "RiceRoundDecryptNode":
|
|
need_wait = False
|
|
break
|
|
if need_wait:
|
|
for i in range(10):
|
|
if unique_id in self.prompt_info.choice_node_map:
|
|
break
|
|
time.sleep(1)
|
|
if unique_id not in self.prompt_info.choice_node_map:
|
|
print(
|
|
f"Warning: RiceRoundSimpleChoiceNode {unique_id} not found in prompt_info.choice_node_map"
|
|
)
|
|
return (default,)
|
|
|
|
|
|
class RiceRoundAdvancedChoiceNode(RiceRoundSimpleChoiceNode):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
CATEGORY = "RiceRound/Advanced"
|
|
|
|
|
|
class RiceRoundSimpleImageNode(LoadImage):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("value",)
|
|
CATEGORY = "RiceRound/Input"
|
|
FUNCTION = "load_image"
|
|
|
|
def load_image(self, image):
|
|
output_image, _ = super().load_image(image)
|
|
return (output_image,)
|
|
|
|
|
|
class RiceRoundImageNode(LoadImage):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
RETURN_TYPES = "IMAGE", "MASK"
|
|
RETURN_NAMES = "image", "mask"
|
|
CATEGORY = "RiceRound/Input"
|
|
FUNCTION = "load_image"
|
|
|
|
def load_image(self, image):
|
|
return super().load_image(image)
|
|
|
|
|
|
class RiceRoundDownloadImageNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {"image_url": ("STRING", {"default": ""})},
|
|
"optional": {},
|
|
"hidden": {},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load_image(self, image_url, **kwargs):
|
|
image = Image.open(requests.get(image_url, stream=True).raw)
|
|
image = ImageOps.exif_transpose(image)
|
|
return (pil2tensor(image),)
|
|
|
|
|
|
class RiceRoundDownloadImageAndMaskNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"image_url": ("STRING", {"default": ""})}}
|
|
|
|
RETURN_TYPES = "IMAGE", "MASK"
|
|
RETURN_NAMES = "image", "mask"
|
|
FUNCTION = "load_image"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load_image(self, image_url, **kwargs):
|
|
img = Image.open(requests.get(image_url, stream=True).raw)
|
|
img = ImageOps.exif_transpose(img)
|
|
output_images = []
|
|
output_masks = []
|
|
w, h = None, None
|
|
excluded_formats = ["MPO"]
|
|
for i in ImageSequence.Iterator(img):
|
|
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
|
if i.mode == "I":
|
|
i = i.point(lambda i: i * (1 / 255))
|
|
image = i.convert("RGB")
|
|
if len(output_images) == 0:
|
|
w = image.size[0]
|
|
h = image.size[1]
|
|
if image.size[0] != w or image.size[1] != h:
|
|
continue
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
if "A" in i.getbands():
|
|
mask = np.array(i.getchannel("A")).astype(np.float32) / 255.0
|
|
mask = 1.0 - torch.from_numpy(mask)
|
|
else:
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
output_images.append(image)
|
|
output_masks.append(mask.unsqueeze(0))
|
|
if len(output_images) > 1 and img.format not in excluded_formats:
|
|
output_image = torch.cat(output_images, dim=0)
|
|
output_mask = torch.cat(output_masks, dim=0)
|
|
else:
|
|
output_image = output_images[0]
|
|
output_mask = output_masks[0]
|
|
return output_image, output_mask
|
|
|
|
|
|
class RiceRoundImageBridgeNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {"images": ("IMAGE", {"tooltip": "only image."})},
|
|
"optional": {},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "bridge"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def bridge(self, images, **kwargs):
|
|
return (images,)
|
|
|
|
|
|
class RiceRoundMaskBridgeNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"mask": ("MASK", {"tooltip": "only image."})}}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "bridge"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def bridge(self, mask, **kwargs):
|
|
return (mask,)
|
|
|
|
|
|
class RiceRoundDownloadMaskNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"mask_url": ("STRING", {"default": ""})}}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load_mask"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load_mask(self, mask_url, **kwargs):
|
|
try:
|
|
response = requests.get(mask_url, stream=True, timeout=10)
|
|
response.raise_for_status()
|
|
mask = Image.open(response.raw)
|
|
if mask.mode != "L":
|
|
mask = mask.convert("L")
|
|
return (pil2tensor(mask),)
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"Error downloading mask from {mask_url}: {str(e)}")
|
|
raise
|
|
except Exception as e:
|
|
print(f"Error processing mask: {str(e)}")
|
|
raise
|
|
|
|
|
|
class RiceRoundIntNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {"default": "数值"}),
|
|
"number": ("INT", {"default": 0}),
|
|
"min": ("INT", {"default": 0}),
|
|
"max": ("INT", {"default": 100}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, name, number, min, max, **kwargs):
|
|
return (number,)
|
|
|
|
|
|
class RiceRoundStrToIntNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"name": ("STRING", {"default": "数值"}), "str": ("STRING",)}}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("value",)
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, name, str, **kwargs):
|
|
return (int(str),)
|
|
|
|
|
|
class RiceRoundFloatNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {"default": "数值"}),
|
|
"number": ("FLOAT", {"default": 0.0}),
|
|
"min": ("FLOAT", {"default": 0.0}),
|
|
"max": ("FLOAT", {"default": 1e2}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, name, number, min, max, **kwargs):
|
|
return (number,)
|
|
|
|
|
|
class RiceRoundStrToFloatNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"name": ("STRING", {"default": "数值"}), "str": ("STRING",)}}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, name, str, **kwargs):
|
|
return (float(str),)
|
|
|
|
|
|
class RiceRoundBooleanNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {"default": "开关"}),
|
|
"value": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("BOOLEAN",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def execute(self, name, value):
|
|
return (value,)
|
|
|
|
|
|
class RiceRoundStrToBooleanNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"name": ("STRING", {"default": "开关"}), "str": ("STRING",)}}
|
|
|
|
RETURN_TYPES = ("BOOLEAN",)
|
|
RETURN_NAMES = ("value",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, name, str, **kwargs):
|
|
return (str.lower() == "true",)
|
|
|
|
|
|
class RiceRoundInputTextNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text_info": (
|
|
"STRING",
|
|
{"multiline": True, "tooltip": "The text to be encoded."},
|
|
)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
def load(self, text_info, **kwargs):
|
|
text = ""
|
|
try:
|
|
json_data = json.loads(text_info)
|
|
text = json_data.get("content", "")
|
|
except json.JSONDecodeError:
|
|
text = text_info
|
|
return (text,)
|
|
|
|
|
|
class RiceRoundRandomSeedNode:
|
|
def __init__(self):
|
|
0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {}, "optional": {}, "hidden": {}}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "random"
|
|
CATEGORY = "RiceRound/Input"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s):
|
|
return random.randint(0, 999999)
|
|
|
|
def random(self):
|
|
r = random.randint(0, 999999)
|
|
print("产生随机数 ", r)
|
|
return (r,)
|