__init__.py

from typing import List
from unittest.mock import MagicMock, patch
import random
import time
import os
import json
import pathlib
from aiohttp import web
import hashlib
import traceback

from rich.console import Console

from execution import validate_prompt
from comfy_execution.graph_utils import GraphBuilder, is_link
import execution
import folder_paths
from server import PromptServer
import nodes

console = Console(color_system="truecolor", force_terminal=True)

idempt = 0
current_node_id = 0

orignal_exe = execution.execute
def new_exe(*args):
    global idempt
    global current_node_id
    try:
        node_id = args[3]
        current_node_id = node_id
        node = args[1].get_node(node_id)
        class_type = node["class_type"];
        node_title = f"{class_type} #{node_id}"
        PromptServer.instance.send_sync("/viv/subgraph/executing", {"title": node_title, "idempt": idempt})
        idempt += 1
    except Exception as e:
        console.print(f"[red]{e}[/red]")
        pass

    return orignal_exe(*args)
execution.execute = new_exe

SUBNODE_FOLDER = pathlib.Path(folder_paths.base_path) / "subnodes"
if not SUBNODE_FOLDER.exists():
    os.mkdir(SUBNODE_FOLDER)


def get_workflow_names() -> list[str]:
    files = list(SUBNODE_FOLDER.glob("*.json"))
    return [file.name for file in files]

def load_workflow(workflow, raw=False):
    with open(SUBNODE_FOLDER / workflow, encoding='utf-8') as f:
        data = json.load(f)

    if not raw:
        if "api_prompt" in data:
            data = data["api_prompt"]
        if "nodes" in data:
            raise ValueError((
                    f"Invalid subgraph file: {workflow}\n"
                    "reason: full workflow export without injected `api_prompt` key\n"
                    "hint: Try loading and re exporting the workflow."
                ))

    return data

def get_outputs(workflow):
    data = load_workflow(workflow, raw=True)
    if "api_prompt" in data:
        prompt = data["api_prompt"]
        extra_data = True
    else:
        prompt = data
        extra_data = False

    results = []
    for (node_id, node) in prompt.items():
        if node["class_type"] == "VIV_Subgraph_Outputs":
            for index, name in enumerate(node["inputs"].keys()):
                type_ = name.split(".")[0]

                if extra_data:
                    for graph_node in data["nodes"]:
                        if str(graph_node["id"]) == node_id:
                            graph_input = graph_node["inputs"][index]
                            if "label" in graph_input:
                                name = graph_input["label"]

                results.append({"name": name, "type": type_})
    return results

def get_inputs(workflow):
    data = load_workflow(workflow, raw=True)

    if "api_prompt" in data:
        prompt = data["api_prompt"]
        extra_data = True
    else:
        prompt = data
        extra_data = False

    for (node_id, node) in prompt.items():
        if node["class_type"] == "VIV_Subgraph_Inputs":
            break;
    else:
        return []

    inputs = {}
    for (id, node) in prompt.items():
        for (name, input) in node["inputs"].items():
            if isinstance(input, list):
                if input[0] == node_id:
                    type_ = "*"

                    prompt_name = name
                    node_name = node["class_type"]

                    cls = nodes.NODE_CLASS_MAPPINGS[node_name]
                    cls_input_config = cls.INPUT_TYPES()
                    input_data = cls_input_config.get("required", {}).get(name) or cls_input_config.get("optional", {}).get(name)

                    widget_data = {}
                    if input_data is not None:
                        if len(input_data) >= 2:
                            if isinstance(input_data[1], dict):
                                widget_data = input_data[1]
                                if widget_data.get("multiline", False):
                                    widget_data["multiline"] = False

                        if isinstance(input_data[0], tuple | list):
                            widget_data["values"] = input_data[0]

                    if extra_data:
                        for graph_node in data["nodes"]:
                            if str(graph_node["id"]) == id:
                                if "title" in graph_node:
                                    node_name = graph_node["title"]
                                for graph_input in graph_node.get("inputs", []):
                                    if graph_input["name"] == name:
                                        type_ = graph_input["type"]
                                        if "label" in graph_input:
                                            prompt_name = graph_input["label"]

                    display_name = f"{prompt_name}.{node_name}.{id}"
                    inputs[input[1]] = {
                            "name": display_name,
                            "type": type_,
                            "widget": widget_data
                        }

    if extra_data:
        for graph_node in data["nodes"]:
            if str(graph_node["id"]) == node_id:
                for index, graph_input in enumerate(graph_node.get("outputs", [])):
                    if "label" in graph_input:
                        input_name = graph_input["label"]
                        inputs[index]["name"] = input_name
                    if graph_input["name"] in graph_node["properties"]:
                        inputs[index]["widget"]["default"] = graph_node["properties"][graph_input["name"]]

    return [name for (_, name) in sorted(inputs.items(), key=lambda x: x[0])]

class Any(str):
    def __ne__(self, value: object, /) -> bool:
        return False

class AcceptAnyWidgetInput(dict):
    def __contains__(self, key: object, /) -> bool:
        return True
    def __getitem__(self, key):
        return self.get(key, Any("*"))

class ReturnAnyAmount(tuple):
    def __getitem__(self, index):
        return Any("*")

in_subgraph = False

class VIV_Subgraph:
    @classmethod
    def INPUT_TYPES(cls):
        return {
                "required": AcceptAnyWidgetInput({
                    "workflow": (tuple(get_workflow_names()), {"forceInput": True})
                    }),
                }

    RETURN_TYPES = ReturnAnyAmount()
    RETURN_NAMES = ReturnAnyAmount()
    FUNCTION = "run"
    CATEGORY = "sub_graph"

    @classmethod
    def IS_CHANGED(cls, workflow, **kwargs):
        workflows = {workflow}
        stack = [workflow]

        while stack:
            workflow = stack.pop()
            data = load_workflow(workflow)
            for node in data.values():
                if node["class_type"] == "VIV_Subgraph":
                    file = node["inputs"]["workflow"]
                    if file not in workflows:
                        workflows.add(file)
                        stack.append(file)

        m = hashlib.sha256()
        for workflow in sorted(workflows):
            with open(SUBNODE_FOLDER / workflow, "rb") as f:
                m.update(f.read())

        return m.digest().hex()

    def run(self, workflow: str, **kwargs):
        prompt = load_workflow(workflow)
        _, error, outputs, _ = validate_prompt(prompt) 
        if error is not None:
           console.print(f"[red]INVALID SUBGRAPH: {workflow}[/red]")
           console.print(error)
           raise ValueError(f"Invalid subgraph: {error['message']}")

        inputs = get_inputs(workflow)
        input_order = [inp["name"] for inp in inputs]
        kwargs = {key: value for (key, value) in kwargs.items() if key in input_order}
        for inp in inputs:
           if inp["name"] not in kwargs:
               kwargs[inp["name"]] = None
        kwargs = sorted(kwargs.items(), key=lambda x: input_order.index(x[0]))
        kwargs = (value for (_, value) in kwargs)

        graph = GraphBuilder()
        output_node = None
        for (node_id, orig_node) in prompt.items():
            node = graph.node(orig_node["class_type"], node_id)
            if orig_node["class_type"] == "VIV_Subgraph_Outputs":
                output_node = node

        for (node_id, orig_node) in prompt.items():
            node = graph.lookup_node(node_id)
            for k, v in orig_node["inputs"].items():
                if is_link(v):
                    parent = graph.lookup_node(v[0])
                    node.set_input(k, parent.out(v[1]))
                else:
                    node.set_input(k, v)
        for (node_id, orig_node) in prompt.items():
            if orig_node["class_type"] == "VIV_Subgraph_Inputs":
                for index, value in enumerate(kwargs):
                    graph.replace_node_output(node_id, index, value)

        if output_node:
            result = tuple(output_node.inputs.values())
        else:
            result = ()
        return {
                "result": result,
                "expand": graph.finalize(),
        }


OUTPUT_RESULTS = []

class VIV_Subgraph_outputs:
    @classmethod
    def INPUT_TYPES(cls):
        return {
                "required": {},
                }

    RETURN_TYPES = ()
    RETURN_NAMES = ()
    FUNCTION = "run"
    CATEGORY = "sub_graph"
    OUTPUT_NODE = True

    def run(self, **kwargs):
        return ()

    # @classmethod
    # def IS_CHANGED(cls, **kwargs):
    #     return time.time()

# INPUTS = []

class VIV_Subgraph_inputs:
    @classmethod
    def INPUT_TYPES(cls):
        return {
                "required": {},
                }

    RETURN_TYPES = ReturnAnyAmount()
    RETURN_NAMES = ReturnAnyAmount()
    FUNCTION = "run"
    CATEGORY = "sub_graph"
    OUTPUT_NODE = True

    def run(self, **kwargs):
        # return tuple(INPUTS.pop())
        return ()

    # @classmethod
    # def IS_CHANGED(cls, **args):
    #     return time.time()

class VIV_Default:
    @classmethod
    def INPUT_TYPES(cls):
        return {
                "required": {
                    "default": Any("*"),
                    },
                "optional": {
                    "inp": Any("*"),
                    }
                }

    DEPRECATED = True
    RETURN_TYPES = (Any("*"),)
    RETURN_NAMES = ("result",)
    FUNCTION = "run"
    CATEGORY = "sub_graph"

    def run(self, default, inp=None):
        if inp is None:
            return (default,)
        return (inp,)

@PromptServer.instance.routes.get("/viv/subgraph")
async def get_workflow(request):
    workflow = request.rel_url.query["workflow"] + ".json"
    data = load_workflow(workflow, raw=True)
    return web.json_response(data)

@PromptServer.instance.routes.get("/viv/input_outputs")
async def get_input_outputs(request):
    workflow = request.rel_url.query["workflow"] + ".json"

    try:
        load_workflow(workflow) # in order to catch errors
        response = {
                "outputs": get_outputs(workflow),
                "inputs": get_inputs(workflow)
                }
        return web.json_response({"data": response})
    except Exception as err:
        traceback.print_exception(err)
        return web.json_response({"error": str(err)})


console.print("[green]Sub nodes, yay![/green]")
console.print("[red]WARNING: This is a very experimental setup :P![/red]")

NODE_CLASS_MAPPINGS = {
        "VIV_Subgraph": VIV_Subgraph,
        "VIV_Subgraph_Outputs": VIV_Subgraph_outputs,
        "VIV_Subgraph_Inputs": VIV_Subgraph_inputs,
        "VIV_Default": VIV_Default,
        }

MANIFEST = {
        "name": "VIVAX",
        }
WEB_DIRECTORY = "web"