crypto_node.py

import json
import os
import random
import uuid
import torch
from comfy_execution.graph import ExecutionBlocker
from comfy_execution.graph_utils import GraphBuilder
from .updown_workflow import UploadWorkflow
from server import PromptServer
from .trim_workflow import CryptoWorkflow, DecodeCryptoWorkflow
from .auth_unit import AuthUnit


class SaveCryptoNode:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {"template_id": ("STRING", {"default": uuid.uuid4().hex})},
            "optional": {"input_anything": ("*", {})},
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "crypto"
    CATEGORY = "advanced/CryptoCat"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    @classmethod
    def VALIDATE_INPUTS(s, input_types):
        return True

    def crypto(self, template_id, **kwargs):
        unique_id = kwargs.pop("unique_id", None)
        prompt = kwargs.pop("prompt", None)
        extra_pnginfo = kwargs.pop("extra_pnginfo", None)
        if unique_id is None:
            raise Exception("Warning: 'unique_id' is missing.")
        if prompt is None:
            raise Exception("Warning: 'prompt' is missing.")
        if len(template_id) != 32:
            raise Exception("Warning: 'template_id' length is not 32.")
        crypto_workflow = CryptoWorkflow(extra_pnginfo["workflow"], prompt, template_id)
        crypto_workflow.invalid_workflow()
        crypto_workflow.load_workflow()
        crypto_workflow.load_prompt()
        crypto_workflow.analysis_node()
        crypto_dir = crypto_workflow.calculate_crypto_result(
            f"crypto_{template_id}.json"
        )
        output_dir = crypto_workflow.output_workflow_simple_shell(
            f"cryptocat_{template_id}.json"
        )
        crypto_workflow.save_original_workflow(
            f"original_workflow_{template_id}.json", crypto_dir
        )
        crypto_workflow.save_original_prompt(
            f"original_prompt_{template_id}.json", crypto_dir
        )
        user_token, error_msg, error_code = AuthUnit().get_user_token()
        if not user_token:
            print(f"crypto cat upload failed, {error_msg}")
            if error_code == 401 or error_code == -3:
                AuthUnit().login_dialog("输出CryptoCat加密工作流，请先完成登录")
            else:
                PromptServer.instance.send_sync(
                    "cryptocat_toast",
                    {"content": "无法完成鉴权登录，请检查网络或完成登录步骤", "type": "error"},
                )
            return (ExecutionBlocker(None),)
        upload = UploadWorkflow(user_token)
        ret = upload.upload_workflow(template_id, crypto_dir)
        if not ret:
            print("crypto cat upload failed")
            return (ExecutionBlocker(None),)
        serial_numbers = upload.generate_serial_number(template_id, count=10)
        if serial_numbers:
            sn_file_path = os.path.join(output_dir, f"serial_numbers_{template_id}.txt")
            with open(sn_file_path, "w", encoding="utf-8") as f:
                f.write("\n".join(serial_numbers))
            PromptServer.instance.send_sync(
                "cryptocat_toast",
                {
                    "content": f"加密完成，在{sn_file_path}中查看序列号",
                    "type": "info",
                    "duration": 5000,
                },
            )
        return (template_id,)


class AnyType(str):
    def __ne__(self, __value: object) -> bool:
        return False


any = AnyType("*")


class SaveCryptoBridgeNode:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {"value": (any,)}}

    @classmethod
    def VALIDATE_INPUTS(s, input_types):
        return True

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    RETURN_TYPES = (any,)
    FUNCTION = "doit"
    CATEGORY = "advanced/CryptoCat"

    def doit(self, value):
        return (value,)


def is_link(obj):
    if not isinstance(obj, list):
        return False
    if len(obj) != 2:
        return False
    if not isinstance(obj[0], str):
        return False
    if (
        not isinstance(obj[1], int)
        and not isinstance(obj[1], float)
        and not isinstance(obj[1], str)
    ):
        return False
    return True


class AlwaysEqualProxy(str):
    def __eq__(self, _):
        return True

    def __ne__(self, _):
        return False


class AlwaysTupleZero(tuple):
    def __getitem__(self, _):
        return AlwaysEqualProxy(super().__getitem__(0))


class DecodeCryptoNode:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "template_id": ("STRING",),
                "serial_number": (
                    "STRING",
                    {"multiline": True, "placeholder": "请输入序列号"},
                ),
            },
            "optional": {"input_anything": (any,)},
            "hidden": {
                "unique_id": "UNIQUE_ID",
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    @classmethod
    def VALIDATE_INPUTS(s, input_types):
        return True

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    RETURN_TYPES = AlwaysTupleZero(AlwaysEqualProxy("*"))
    FUNCTION = "decode"
    CATEGORY = "advanced/CryptoCat"
    OUTPUT_NODE = False

    def decode(self, template_id, serial_number, **kwargs):
        unique_id = kwargs.pop("unique_id", None)
        prompt = kwargs.pop("prompt", None)
        extra_pnginfo = kwargs.pop("extra_pnginfo", None)
        decode_crypto_workflow = DecodeCryptoWorkflow(
            prompt, extra_pnginfo["workflow"], template_id
        )
        crypto_prompt = decode_crypto_workflow.load_crypto_prompt(serial_number)
        decode_crypto_workflow.calculate_input_anything_map()
        processed_nodes = {}
        graph = GraphBuilder()

        def get_node_result(nodeData, id):
            inputKeys = []
            for ikey in nodeData["inputs"].keys():
                input_value = nodeData["inputs"][ikey]
                if (
                    is_link(input_value)
                    and decode_crypto_workflow.get_hidden_input(input_value) is None
                    and input_value[0] not in processed_nodes
                ):
                    inputKeys.append(input_value[0])
            for ikey in inputKeys:
                if ikey not in crypto_prompt:
                    continue
                node = get_node_result(crypto_prompt[ikey], ikey)
                processed_nodes[ikey] = node
            inputs = nodeData["inputs"]
            newInputs = {}
            for ikey in inputs.keys():
                if is_link(inputs[ikey]):
                    hidden_input_name = decode_crypto_workflow.get_hidden_input(
                        inputs[ikey]
                    )
                    if hidden_input_name:
                        if hidden_input_name in kwargs:
                            newInputs[ikey] = kwargs[hidden_input_name]
                    elif inputs[ikey][0] in processed_nodes:
                        newInputs[ikey] = processed_nodes[inputs[ikey][0]].out(
                            inputs[ikey][1]
                        )
                else:
                    newInputs[ikey] = inputs[ikey]
            return graph.node(nodeData["class_type"], id, **newInputs)

        node_id, link_idx = decode_crypto_workflow.get_outputs()
        nodeData = crypto_prompt[node_id]
        node = get_node_result(nodeData, node_id)
        value = node.out(link_idx)
        return {"result": tuple([value]), "expand": graph.finalize()}


class RandomSeedNode:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {"required": {}, "optional": {}, "hidden": {}}

    RETURN_TYPES = ("INT",)
    FUNCTION = "random"
    CATEGORY = "advanced/CryptoCat"

    def IS_CHANGED():
        return float("NaN")

    def random(self):
        return (random.randint(0, 999999),)