50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import torch
|
|
|
|
from .zeval import LatentMap, parse
|
|
|
|
|
|
class ZEvalNode:
|
|
"""
|
|
ComfyUI Custom Node with dynamic latent input count.
|
|
"""
|
|
|
|
CATEGORY = "latent"
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "evaluate"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"expression": ("STRING", {"multiline": True, "default": "l1"}),
|
|
"latent_1": ("LATENT",),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, expression, **kwargs):
|
|
"""Force re-evaluation when expression changes."""
|
|
return expression or kwargs
|
|
|
|
def evaluate(self, expression: str, **kwargs) -> tuple[dict[str, torch.Tensor]]:
|
|
"""
|
|
Evaluate the expression using the provided latents.
|
|
"""
|
|
|
|
latents = [
|
|
(name[7:], latent['samples'])
|
|
for name, latent in kwargs.items()
|
|
if name.startswith("latent_")
|
|
]
|
|
latents.sort()
|
|
latent_map = LatentMap.create(latent for _, latent in latents)
|
|
try:
|
|
ast = parse(expression)
|
|
print(f"ZEvalNode: evaluating expression {ast!r}")
|
|
result = ast(latent_map)
|
|
|
|
except Exception as e:
|
|
raise RuntimeError(f"Latent expression parsing failed: {e}") from e
|
|
|
|
return ({"samples": result},)
|