365 lines
12 KiB
Python
365 lines
12 KiB
Python
import sys
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
from typing import Any, NamedTuple
|
|
|
|
import comfy.samplers
|
|
import numpy as np
|
|
import server
|
|
import torch
|
|
from comfy.comfy_types.node_typing import IO
|
|
from comfy_api.latest import io
|
|
from PIL import Image
|
|
|
|
from .nodes import SendImageWebSocket
|
|
|
|
|
|
class Publisher(NamedTuple):
|
|
name: str
|
|
id: str
|
|
workflow: dict
|
|
|
|
|
|
class WorkflowExchange:
|
|
def __init__(self, server: server.PromptServer):
|
|
self._server = server
|
|
self._publishers: dict[str, Publisher] = {}
|
|
self._subscribers: list[str] = []
|
|
|
|
async def publish(self, publisher_name: str, publisher_id: str, workflow: dict):
|
|
publisher = Publisher(publisher_name, publisher_id, workflow)
|
|
for client_id in self._subscribers:
|
|
await self._notify(client_id, publisher)
|
|
self._publishers[publisher_id] = publisher
|
|
|
|
async def subscribe(self, client_id: str):
|
|
if client_id in self._subscribers:
|
|
raise KeyError("Already subscribed")
|
|
self._subscribers.append(client_id)
|
|
for publisher in self._publishers.values():
|
|
await self._notify(client_id, publisher)
|
|
|
|
async def unsubscribe(self, client_id: str):
|
|
if client_id in self._subscribers:
|
|
self._subscribers.remove(client_id)
|
|
else:
|
|
raise KeyError("No subscriber found with id " + client_id)
|
|
|
|
async def _notify(self, client_id: str, publisher: Publisher):
|
|
data = {
|
|
"publisher": {"name": publisher.name, "id": publisher.id},
|
|
"workflow": publisher.workflow,
|
|
}
|
|
await self._server.send_json("etn_workflow_published", data, client_id)
|
|
|
|
|
|
def _placeholder_image():
|
|
path = Path(__file__).parent / "data" / "external-image-placeholder.webp"
|
|
image = Image.open(path).convert("RGB")
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
return torch.from_numpy(image)[None,]
|
|
|
|
|
|
class _BasicTypes(str):
|
|
"""Matches IO.PRIMITIVE, but also any list of choices"""
|
|
|
|
basic_types = IO.PRIMITIVE.split(",") # STRING, FLOAT, INT, BOOLEAN
|
|
|
|
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 OutputBatchMode(Enum):
|
|
default = "default"
|
|
images = "images"
|
|
animation = "animation"
|
|
layers = "layers"
|
|
|
|
|
|
class KritaOutput(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaOutput",
|
|
display_name="Krita Output",
|
|
category="krita",
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.Int.Input("x", "offset x", default=0),
|
|
io.Int.Input("y", "offset y", default=0),
|
|
io.String.Input("name", default=""),
|
|
io.Combo.Input(
|
|
"batch_mode", OutputBatchMode, "batch mode", default=OutputBatchMode.default
|
|
),
|
|
io.Boolean.Input("resize_canvas", "resize canvas", default=False),
|
|
],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute( # type: ignore
|
|
cls,
|
|
images: torch.Tensor,
|
|
x: int = 0,
|
|
y: int = 0,
|
|
name="",
|
|
batch_mode: OutputBatchMode | str = OutputBatchMode.default,
|
|
resize_canvas=False,
|
|
):
|
|
batch_mode = batch_mode.value if isinstance(batch_mode, OutputBatchMode) else batch_mode
|
|
info = {
|
|
"name": name,
|
|
"offset_x": x,
|
|
"offset_y": y,
|
|
"batch_mode": batch_mode,
|
|
"resize_canvas": resize_canvas,
|
|
}
|
|
output = SendImageWebSocket.execute(images, "PNG")
|
|
assert isinstance(output.ui, dict)
|
|
output.ui["info"] = [info]
|
|
return output
|
|
|
|
|
|
class KritaSendText(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaSendText",
|
|
display_name="Send Text",
|
|
category="krita",
|
|
inputs=[
|
|
io.AnyType.Input("value"),
|
|
io.String.Input("name", default="Output"),
|
|
io.Combo.Input("type", options=["text", "markdown", "html"], default="text"),
|
|
],
|
|
is_output_node=True,
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, value: Any, name: str, type: str): # type: ignore
|
|
mime = {
|
|
"text": "text/plain",
|
|
"markdown": "text/markdown",
|
|
"html": "text/html",
|
|
}[type]
|
|
text = "None"
|
|
if value is not None:
|
|
try:
|
|
text = str(value)
|
|
except Exception as e:
|
|
text = f"Could not convert to text: {e}"
|
|
|
|
return io.NodeOutput(ui={"text": [{"name": name, "text": text, "content-type": mime}]})
|
|
|
|
|
|
class KritaCanvas(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaCanvas",
|
|
display_name="Krita Canvas",
|
|
category="krita",
|
|
outputs=[
|
|
io.Image.Output(display_name="image"),
|
|
io.Int.Output(display_name="width"),
|
|
io.Int.Output(display_name="height"),
|
|
io.Int.Output(display_name="seed"),
|
|
io.Mask.Output(display_name="mask"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs):
|
|
return io.NodeOutput(_placeholder_image(), 512, 512, 0, torch.ones(1, 512, 512))
|
|
|
|
|
|
class SelectionContext(Enum):
|
|
automatic = "automatic"
|
|
entire_image = "entire image"
|
|
mask_bounds = "mask bounds"
|
|
|
|
|
|
_selection_context_help = """
|
|
Determines the section (crop bounding box) of the image and mask to transmit:
|
|
- automatic: area around the selection determined by Krita settings
|
|
- entire image: always use the entire canvas area
|
|
- mask bounds: tight bounding box of the current selection
|
|
|
|
This affects the Selection and Canvas nodes. The offset x/y outputs indicate the top-left corner of the context area relative to the full canvas."""
|
|
|
|
|
|
class KritaSelection(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaSelection",
|
|
display_name="Krita Selection",
|
|
category="krita",
|
|
inputs=[
|
|
io.Combo.Input(
|
|
"context",
|
|
options=SelectionContext,
|
|
default=SelectionContext.entire_image,
|
|
tooltip=_selection_context_help,
|
|
),
|
|
io.Int.Input("padding", "padding", default=0, min=0),
|
|
],
|
|
outputs=[
|
|
io.Mask.Output("mask", "mask"),
|
|
io.Boolean.Output("active", "active"),
|
|
io.Int.Output("x", "offset x"),
|
|
io.Int.Output("y", "offset y"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs):
|
|
return io.NodeOutput(torch.ones(1, 512, 512), False, 0, 0)
|
|
|
|
|
|
class KritaImageLayer(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaImageLayer",
|
|
display_name="Krita Image Layer",
|
|
category="krita",
|
|
inputs=[io.String.Input("name", default="Image")],
|
|
outputs=[
|
|
io.Image.Output(display_name="image"),
|
|
io.Mask.Output(display_name="mask"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, name: str): # type: ignore
|
|
return io.NodeOutput(_placeholder_image(), torch.ones(1, 512, 512))
|
|
|
|
|
|
class KritaMaskLayer(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaMaskLayer",
|
|
display_name="Krita Mask Layer",
|
|
category="krita",
|
|
inputs=[io.String.Input("name", default="Mask")],
|
|
outputs=[
|
|
io.Mask.Output(display_name="mask"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, name: str): # type: ignore
|
|
return io.NodeOutput(torch.ones(1, 512, 512))
|
|
|
|
|
|
_param_types = [
|
|
"auto",
|
|
"number",
|
|
"number (integer)",
|
|
"toggle",
|
|
"choice",
|
|
"text",
|
|
"prompt (positive)",
|
|
"prompt (negative)",
|
|
]
|
|
_fmax = sys.float_info.max
|
|
|
|
|
|
class Parameter(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_Parameter",
|
|
display_name="Parameter",
|
|
category="krita",
|
|
inputs=[
|
|
io.String.Input("name", default="Parameter"),
|
|
io.Combo.Input("type", options=_param_types, default="auto"),
|
|
io.String.Input("default", default=""),
|
|
io.Float.Input("min", default=-1e10, min=-_fmax, max=_fmax, optional=True),
|
|
io.Float.Input("max", default=1e10, min=-_fmax, max=_fmax, optional=True),
|
|
],
|
|
outputs=[io.AnyType.Output(display_name="value")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, name: str, type: str, default, min=0.0, max=1.0): # type: ignore
|
|
if type == "number":
|
|
return io.NodeOutput(float(default))
|
|
elif type == "number (integer)":
|
|
return io.NodeOutput(int(default))
|
|
return io.NodeOutput(default)
|
|
|
|
|
|
class KritaStyle(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaStyle",
|
|
display_name="Krita Style",
|
|
category="krita",
|
|
inputs=[
|
|
io.String.Input("name", default="Style"),
|
|
io.Combo.Input("sampler_preset", options=["auto", "regular", "live"]),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model"),
|
|
io.Clip.Output(display_name="clip"),
|
|
io.Vae.Output(display_name="vae"),
|
|
io.String.Output(display_name="positive prompt"),
|
|
io.String.Output(display_name="negative prompt"),
|
|
io.Combo.Output(
|
|
display_name="sampler name", options=comfy.samplers.KSampler.SAMPLERS
|
|
),
|
|
io.Combo.Output(
|
|
display_name="scheduler", options=comfy.samplers.KSampler.SCHEDULERS
|
|
),
|
|
io.Int.Output(display_name="steps"),
|
|
io.Float.Output(display_name="guidance"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, name: str, sampler_preset: str): # type: ignore
|
|
raise NotImplementedError("This workflow must be started from Krita!")
|
|
|
|
|
|
class KritaStyleAndPrompt(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="ETN_KritaStyleAndPrompt",
|
|
display_name="Krita Style & Prompt",
|
|
category="krita",
|
|
inputs=[
|
|
io.Combo.Input("sampler_preset", options=["auto", "regular", "live"]),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model (with loras)"),
|
|
io.Clip.Output(display_name="clip"),
|
|
io.Vae.Output(display_name="vae"),
|
|
io.String.Output(display_name="positive prompt (evaluated)"),
|
|
io.String.Output(display_name="negative prompt (evaluated)"),
|
|
io.Combo.Output(
|
|
display_name="sampler name", options=comfy.samplers.KSampler.SAMPLERS
|
|
),
|
|
io.Combo.Output(
|
|
display_name="scheduler", options=comfy.samplers.KSampler.SCHEDULERS
|
|
),
|
|
io.Int.Output(display_name="steps"),
|
|
io.Float.Output(display_name="guidance"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, name: str, sampler_preset: str): # type: ignore
|
|
raise NotImplementedError("This workflow must be started from Krita!")
|