446 lines
13 KiB
Python
446 lines
13 KiB
Python
import copy
|
|
import glob
|
|
import os
|
|
import textwrap
|
|
from pathlib import Path
|
|
from typing import Any, TypedDict
|
|
|
|
import numpy as np
|
|
import torch
|
|
import yaml
|
|
from PIL import Image
|
|
|
|
NODE_CLASS_MAPPINGS = {}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
|
|
|
|
def register_node(identifier: str, display_name: str):
|
|
def decorator(cls):
|
|
NODE_CLASS_MAPPINGS[identifier] = cls
|
|
NODE_DISPLAY_NAME_MAPPINGS[identifier] = display_name
|
|
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
def load_image(path, convert="RGB"):
|
|
img = Image.open(path).convert(convert)
|
|
img = np.array(img).astype(np.float32) / 255.0
|
|
img = torch.from_numpy(img).unsqueeze(0)
|
|
return img
|
|
|
|
|
|
class RangedConfig:
|
|
def __init__(self, definition: dict[str, Any], range_key: str = "ranges") -> None:
|
|
self.definition = definition
|
|
self.range_key = range_key
|
|
self._validate()
|
|
|
|
def _validate(self):
|
|
for k, v in self.definition.items():
|
|
if k == self.range_key:
|
|
assert isinstance(v, dict), f"{type(v)!r}"
|
|
for kk, vv in v.items():
|
|
# sub prompts ranges
|
|
assert isinstance(kk, int), f"{type(kk)!r}"
|
|
assert isinstance(vv, dict), f"{type(vv)!r}"
|
|
for kkk, vvv in vv.items():
|
|
# actual sub prompts
|
|
assert isinstance(kkk, str), f"{type(kkk)!r}"
|
|
if vvv is None:
|
|
vvv = ""
|
|
assert isinstance(vvv, (int, float, str)), f"{type(vvv)!r}"
|
|
else:
|
|
# base prompts
|
|
assert isinstance(k, str), f"{type(k)!r}"
|
|
if v is None:
|
|
v = ""
|
|
assert isinstance(v, (int, float, str)), f"{type(v)!r}"
|
|
|
|
def get_ranges(self):
|
|
return sorted(self.definition[self.range_key].keys())
|
|
|
|
def _get_range_start(self, i: int) -> int | None:
|
|
if len(self.definition[self.range_key]) == 0:
|
|
return None
|
|
|
|
range_starts = sorted(self.definition[self.range_key].keys())
|
|
|
|
for range_start_idx, range_start in enumerate(range_starts):
|
|
if i < range_start:
|
|
if range_start_idx == 0:
|
|
return None
|
|
else:
|
|
return range_starts[range_start_idx - 1]
|
|
|
|
return range_starts[-1]
|
|
|
|
def _get_raw_sub_prompt(self, i: int):
|
|
range_start = self._get_range_start(i)
|
|
if range_start is None:
|
|
# not in range, just use base definition
|
|
return {**self.definition, "i": i}
|
|
else:
|
|
raw_sub_prompt = self.definition[self.range_key][self._get_range_start(i)]
|
|
return {**self.definition, **raw_sub_prompt, "i": i}
|
|
|
|
def get_sub_prompt(self, i: int):
|
|
raw_sub_prompt = self._get_raw_sub_prompt(i)
|
|
sub_prompt = {}
|
|
for k, v in raw_sub_prompt.items():
|
|
if k == self.range_key:
|
|
continue
|
|
|
|
if isinstance(v, str):
|
|
v = v.format(**raw_sub_prompt)
|
|
|
|
sub_prompt[k] = v
|
|
return sub_prompt
|
|
|
|
|
|
DEFAULT_CONFIG = """\
|
|
p: |
|
|
masterpiece, best quality,
|
|
{sp},
|
|
|
|
n: |
|
|
{sn},
|
|
embedding:EasyNegative, embedding:bad-artist, embedding:bad-hands-5, embedding:bad-image-v2-39000,
|
|
lowres, ((bad anatomy)), ((bad hands)), text, missing finger, extra digits, fewer digits, blurry, ((mutated hands and fingers)), (poorly drawn face), ((mutation)), ((deformed face)), (ugly), ((bad proportions)), ((extra limbs)), extra face, (double head), (extra head), ((extra feet)), monster, logo, cropped, worst quality, low quality, normal quality, jpeg, humpbacked, long body, long neck, ((jpeg artifacts)),
|
|
|
|
path: "{i:04d}.png"
|
|
|
|
example: 0
|
|
|
|
ranges:
|
|
1:
|
|
sp: positive subprompt for 1-4
|
|
sn: negative subprompt for 1-4
|
|
5:
|
|
sp: positive subprompt for 5-...
|
|
sn: negative subprompt for 5-...
|
|
example: 1
|
|
"""
|
|
|
|
|
|
@register_node("JWInfoHashFromRangedInfo", "Info Hash From Ranged Config")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"config": (
|
|
"STRING",
|
|
{"default": DEFAULT_CONFIG, "multiline": True, "dynamicPrompts": False},
|
|
),
|
|
"i": ("INT", {"default": 1, "min": 0, "step": 1, "max": 999999}),
|
|
"ranges_key": ("STRING", {"default": "ranges", "multiline": False}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("INFO_HASH",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, config: str, i: int, ranges_key: str):
|
|
config = yaml.safe_load(config)
|
|
|
|
info = RangedConfig(config, range_key=ranges_key)
|
|
|
|
return (info.get_sub_prompt(i),)
|
|
|
|
|
|
@register_node("JWInfoHashListFromRangedInfo", "Info Hash List From Ranged Config")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"config": (
|
|
"STRING",
|
|
{"default": DEFAULT_CONFIG, "multiline": True, "dynamicPrompts": False},
|
|
),
|
|
"i_start": ("INT", {"default": 0, "min": 0, "step": 1, "max": 999999}),
|
|
"i_stop": ("INT", {"default": 16, "min": 0, "step": 1, "max": 999999}),
|
|
"ranges_key": ("STRING", {"default": "ranges", "multiline": False}),
|
|
"inclusive": (("false", "true"), {"default": "false"}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("INFO_HASH_LIST",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(
|
|
self, config: str, i_start: int, i_stop: int, ranges_key: str, inclusive: str
|
|
):
|
|
inclusive: bool = inclusive == "true"
|
|
|
|
config = yaml.safe_load(config)
|
|
|
|
info = RangedConfig(config, range_key=ranges_key)
|
|
subinfos = [
|
|
info.get_sub_prompt(i)
|
|
for i in range(i_start, i_stop + 1 if inclusive else i_stop)
|
|
]
|
|
|
|
return (subinfos,)
|
|
|
|
|
|
def calculate_batches(
|
|
i_start: int, # start of i
|
|
i_stop: int, # end of i, excludes end
|
|
range_starts: int, # scene cuts, batch will be terminated before this
|
|
max_batch_size: int, # maximum length of batch
|
|
):
|
|
"""
|
|
:param int i_start: start of i
|
|
:param int i_stop: end of i, excludes end
|
|
:param int range_starts: scene cuts, batch will be terminated before this
|
|
:param int max_batch_size: maximum length of batch
|
|
:return: a list of 2-tuples, each represents (batch start frame, batch stop frame), where stop frame is exclusive
|
|
"""
|
|
batch_starts: list[int] = [] # also includes end frame
|
|
i = i_start - 1
|
|
counter = -1
|
|
while True:
|
|
i += 1
|
|
counter += 1
|
|
if i >= i_stop:
|
|
batch_starts.append(i)
|
|
break
|
|
|
|
if i in range_starts:
|
|
batch_starts.append(i)
|
|
counter = 0
|
|
continue
|
|
|
|
if counter >= max_batch_size:
|
|
batch_starts.append(i)
|
|
counter = 0
|
|
continue
|
|
|
|
if counter == 0:
|
|
batch_starts.append(i)
|
|
continue
|
|
|
|
batches = list(zip(batch_starts[:-1], batch_starts[1:]))
|
|
|
|
return batches
|
|
|
|
|
|
@register_node("JWRangedInfoCalculateSubBatch", "Calculate Sub Batch for Ranged Info")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"config": (
|
|
"STRING",
|
|
{"default": DEFAULT_CONFIG, "multiline": True, "dynamicPrompts": False},
|
|
),
|
|
"ranges_key": ("STRING", {"default": "ranges", "multiline": False}),
|
|
"batch_idx": ("INT", {"default": 0, "min": 0, "step": 1, "max": 999999}),
|
|
"i_start": ("INT", {"default": 1, "min": 0, "step": 1, "max": 999999}),
|
|
"i_stop": ("INT", {"default": 100, "min": 0, "step": 1, "max": 999999}),
|
|
"max_batch_size": (
|
|
"INT",
|
|
{"default": 16, "min": 1, "step": 1, "max": 999999},
|
|
),
|
|
"inclusive": (("false", "true"), {"default": "false"}),
|
|
}
|
|
}
|
|
RETURN_NAMES = ("BATCH_I_START", "BATCH_I_STOP")
|
|
RETURN_TYPES = ("INT", "INT")
|
|
FUNCTION = "execute"
|
|
|
|
def execute(
|
|
self,
|
|
config: str,
|
|
ranges_key: str,
|
|
batch_idx: int,
|
|
i_start: int,
|
|
i_stop: int,
|
|
max_batch_size: int,
|
|
inclusive: str,
|
|
):
|
|
inclusive: bool = inclusive == "true"
|
|
|
|
config = yaml.safe_load(config)
|
|
|
|
info = RangedConfig(config, range_key=ranges_key)
|
|
|
|
range_starts = set(info.get_ranges())
|
|
|
|
# get images in selected batch
|
|
batches = calculate_batches(
|
|
i_start, i_stop + 1 if inclusive else i_stop, range_starts, max_batch_size
|
|
)
|
|
batch = batches[batch_idx]
|
|
|
|
return (batch[0], batch[1])
|
|
|
|
|
|
@register_node(
|
|
"JWInfoHashFromRangedInfoAndLoadSubsequences",
|
|
"Info Hash From Ranged Config and Load Batch",
|
|
)
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"config": (
|
|
"STRING",
|
|
{"default": DEFAULT_CONFIG, "multiline": True, "dynamicPrompts": False},
|
|
),
|
|
"ranges_key": ("STRING", {"default": "ranges", "multiline": False}),
|
|
"path_key": ("STRING", {"default": "path", "multiline": False}),
|
|
"batch_idx": ("INT", {"default": 0, "min": 0, "step": 1, "max": 999999}),
|
|
"i_start": ("INT", {"default": 1, "min": 0, "step": 1, "max": 999999}),
|
|
"i_stop": ("INT", {"default": 100, "min": 0, "step": 1, "max": 999999}),
|
|
"max_batch_size": (
|
|
"INT",
|
|
{"default": 16, "min": 1, "step": 1, "max": 999999},
|
|
),
|
|
"inclusive": (("false", "true"), {"default": "false"}),
|
|
}
|
|
}
|
|
RETURN_NAMES = ("INFO_HASH", "IMAGE", "BATCH_I_START", "BATCH_I_STOP")
|
|
RETURN_TYPES = ("INFO_HASH", "IMAGE", "INT", "INT")
|
|
FUNCTION = "execute"
|
|
|
|
def execute(
|
|
self,
|
|
config: str,
|
|
ranges_key: str,
|
|
path_key: str,
|
|
batch_idx: int,
|
|
i_start: int,
|
|
i_stop: int,
|
|
max_batch_size: int,
|
|
inclusive: str,
|
|
):
|
|
inclusive: bool = inclusive == "true"
|
|
|
|
config = yaml.safe_load(config)
|
|
|
|
info = RangedConfig(config, range_key=ranges_key)
|
|
|
|
range_starts = set(info.get_ranges())
|
|
|
|
# get images in selected batch
|
|
batches = calculate_batches(
|
|
i_start, i_stop + 1 if inclusive else i_stop, range_starts, max_batch_size
|
|
)
|
|
batch = batches[batch_idx]
|
|
|
|
print(f"Getting images in batch: {batch}")
|
|
|
|
images = []
|
|
for i in range(batch[0], batch[1]):
|
|
subinfo = info.get_sub_prompt(i)
|
|
path = subinfo[path_key]
|
|
print(f" Loading: {path}")
|
|
img = load_image(path)
|
|
images.append(img)
|
|
images = torch.cat(images, dim=0)
|
|
|
|
return (info.get_sub_prompt(batch[0]), images, batch[0], batch[1])
|
|
|
|
|
|
@register_node("JWInfoHashExtractInteger", "Info Hash Extract Integer")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash": ("INFO_HASH",),
|
|
"key": ("STRING", {"default": "i", "multiline": False}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash: dict, key: str):
|
|
val = int(info_hash[key])
|
|
return (val,)
|
|
|
|
|
|
@register_node("JWInfoHashExtractFloat", "Info Hash Extract Float")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash": ("INFO_HASH",),
|
|
"key": ("STRING", {"default": "key", "multiline": False}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash: dict, key: str):
|
|
val = float(info_hash[key])
|
|
return (val,)
|
|
|
|
|
|
@register_node("JWInfoHashExtractString", "Info Hash Extract String")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash": ("INFO_HASH",),
|
|
"key": ("STRING", {"default": "p", "multiline": False}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("STRING",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash: dict, key: str):
|
|
val = str(info_hash[key])
|
|
return (val,)
|
|
|
|
|
|
@register_node("JWInfoHashListExtractStringList", "Info Hash List Extract String List")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash_list": ("INFO_HASH_LIST",),
|
|
"key": ("STRING", {"default": "p", "multiline": False}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("STRING_LIST",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash_list: list[dict], key: str):
|
|
val = [str(info_hash[key]) for info_hash in info_hash_list]
|
|
return (val,)
|
|
|
|
|
|
@register_node("JWInfoHashFromInfoHashList", "Extract Info Hash From Info Hash List")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash_list": ("INFO_HASH_LIST",),
|
|
"i": ("INT", {"default": 0, "step": 1, "min": -99999999, "max": 99999999}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("INFO_HASH",)
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash_list: list[dict], i: int):
|
|
return (info_hash_list[i],)
|
|
|
|
|
|
@register_node("JWInfoHashPrint", "Print Info Hash (Debug)")
|
|
class _:
|
|
CATEGORY = "jamesWalker55"
|
|
INPUT_TYPES = lambda: {
|
|
"required": {
|
|
"info_hash": ("INFO_HASH",),
|
|
}
|
|
}
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "execute"
|
|
|
|
def execute(self, info_hash: dict):
|
|
from pprint import pformat, pprint
|
|
|
|
pprint(info_hash)
|
|
raise ValueError(pformat(info_hash))
|