RiceRound-ComfyUI_CryptoCat/trim_workflow.py

from collections import defaultdict
import copy
import os
import random
from server import PromptServer
from .updown_workflow import DownloadWorkflow
import folder_paths
import json
from .file_compressor import FileCompressor
from typing import Dict, Any
from .auth_unit import AuthUnit


class CryptoWorkflow:
    def __init__(self, workflow, prompt, template_id):
        self.workflow = copy.deepcopy(workflow)
        self.prompt = copy.deepcopy(prompt)
        self.template_id = template_id
        self.workflow_nodes_dict = {}
        self.link_owner_map = defaultdict(dict)
        self.node_prompt_map = {}
        self.last_node_id = 0
        self.last_link_id = 0
        self.save_crypto_node_id = 0
        self.crypto_bridge_node_id = 0
        self.input_nodes_ids = set()
        self.output_nodes_ids = set()
        self.crypto_nodes_ids = set()
        self.crypto_result = {}

    def invalid_workflow(self):
        for node in self.workflow["nodes"]:
            node_type = node.get("type", "")
            if node_type == "SaveCryptoBridgeNode":
                if self.crypto_bridge_node_id == 0:
                    self.crypto_bridge_node_id = int(node["id"])
                else:
                    raise ValueError(
                        "Error: Multiple 'SaveCryptoBridgeNode' instances found."
                    )
                continue
            elif node_type == "SaveCryptoNode":
                if self.save_crypto_node_id == 0:
                    self.save_crypto_node_id = int(node["id"])
                else:
                    raise ValueError(
                        "Error: Multiple 'SaveCryptoNode' instances found."
                    )
                continue
        if self.save_crypto_node_id == 0:
            raise ValueError("Error: No 'SaveCryptoNode' instances found.")
        if self.crypto_bridge_node_id == 0:
            raise ValueError("Error: No 'SaveCryptoBridgeNode' instances found.")

    def load_workflow(self):
        simplify_workflow = self.workflow
        self.workflow_nodes_dict = {
            int(node["id"]): node for node in simplify_workflow["nodes"]
        }
        for node in simplify_workflow["nodes"]:
            output_nodes = node.get("outputs", [])
            if not output_nodes:
                continue
            for output in output_nodes:
                links = output.get("links", [])
                if not links:
                    continue
                for link in links:
                    link = int(link)
                    self.link_owner_map[link]["links"] = copy.deepcopy(links)
                    self.link_owner_map[link]["slot_index"] = output.get(
                        "slot_index", 0
                    )
                    self.link_owner_map[link]["owner_id"] = int(node["id"])
                    self.link_owner_map[link]["type"] = output.get("type", "")
        self.last_node_id = int(simplify_workflow.get("last_node_id", 0))
        self.last_link_id = int(simplify_workflow.get("last_link_id", 0))

    def load_prompt(self):
        simplify_prompt = self.prompt
        self.node_prompt_map = {
            int(node_id): node for (node_id, node) in simplify_prompt.items()
        }

    def analysis_node(self):
        self.input_nodes_ids.clear()
        self.output_nodes_ids.clear()
        self.crypto_nodes_ids.clear()

        def find_input_nodes(node_id, visited=None):
            if visited is None:
                visited = set()
            if node_id in visited:
                return
            visited.add(node_id)
            self.input_nodes_ids.add(node_id)
            node = self.workflow_nodes_dict.get(node_id)
            if not node:
                return
            for input_node in node.get("inputs", []):
                input_link = input_node.get("link")
                if input_link is not None and input_link in self.link_owner_map:
                    owner_id = self.link_owner_map[input_link]["owner_id"]
                    find_input_nodes(owner_id, visited)

        for input_node in self.workflow_nodes_dict[self.save_crypto_node_id].get(
            "inputs", []
        ):
            if input_node["name"] and input_node["name"].startswith("input_anything"):
                input_link = input_node["link"]
                if input_link is not None and input_link in self.link_owner_map:
                    owner_id = self.link_owner_map[input_link]["owner_id"]
                    find_input_nodes(owner_id)

        def find_output_nodes(node_id, visited=None):
            if visited is None:
                visited = set()
            if node_id in visited:
                return
            visited.add(node_id)
            self.output_nodes_ids.add(node_id)
            node = self.workflow_nodes_dict.get(node_id)
            if not node:
                return
            for output in node.get("outputs", []):
                if not output or not output.get("links"):
                    continue
                for link in output.get("links", []):
                    if not link:
                        continue
                    if link is not None:
                        for workflow_link in self.workflow.get("links", []):
                            if workflow_link[0] == link:
                                target_node_id = workflow_link[3]
                                find_output_nodes(target_node_id, visited)

        output_links = set()
        for output_node in self.workflow_nodes_dict[self.crypto_bridge_node_id].get(
            "outputs", []
        ):
            _links = output_node["links"]
            output_links.update(_links)
        for link in self.workflow.get("links", []):
            if len(link) > 3 and link[0] in output_links:
                find_output_nodes(link[3])
        self.crypto_nodes_ids = (
            self.workflow_nodes_dict.keys()
            - self.input_nodes_ids
            - self.output_nodes_ids
        )
        self.crypto_nodes_ids = self.crypto_nodes_ids - {
            self.save_crypto_node_id,
            self.crypto_bridge_node_id,
        }

    def calculate_crypto_result(self, crypto_file_name):
        self.crypto_result = {"prompt": {}, "workflow": {}, "outputs": []}
        for node_id in self.crypto_nodes_ids:
            if node_id in self.node_prompt_map:
                self.crypto_result["prompt"][node_id] = self.node_prompt_map[node_id]
            if node_id in self.workflow_nodes_dict:
                self.crypto_result["workflow"][node_id] = self.workflow_nodes_dict[
                    node_id
                ]
        crypto_bridge_node = self.node_prompt_map[self.crypto_bridge_node_id]
        for input_name, input_value in crypto_bridge_node.get("inputs", {}).items():
            if isinstance(input_value, list) and len(input_value) == 2:
                self.crypto_result["outputs"] = input_value[0], input_value[1]
        json_result = json.dumps(self.crypto_result, indent=4, ensure_ascii=False)
        save_dir = folder_paths.temp_directory
        with open(os.path.join(save_dir, crypto_file_name), "w", encoding="utf-8") as f:
            f.write(json_result)
        return save_dir

    def output_workflow_simple_shell(self, output_workflow_name):
        simplify_workflow = copy.deepcopy(self.workflow)
        save_crypto_node = None
        crypto_bridge_node = None
        for node in simplify_workflow["nodes"]:
            if node["id"] == self.save_crypto_node_id:
                save_crypto_node = node
            if node["id"] == self.crypto_bridge_node_id:
                crypto_bridge_node = node
        except_nodes_ids = self.crypto_nodes_ids
        except_nodes_ids.add(self.crypto_bridge_node_id)
        simplify_workflow["nodes"] = [
            node
            for node in simplify_workflow["nodes"]
            if int(node["id"]) not in except_nodes_ids
        ]
        if save_crypto_node is None:
            raise ValueError("SaveCryptoNode not found in workflow")
        if crypto_bridge_node is None:
            raise ValueError("CryptoBridgeNode not found in workflow")
        save_crypto_node["type"] = "DecodeCryptoNode"
        save_crypto_node["widgets_values"] = (
            [save_crypto_node.get("widgets_values", [None])[0]]
            if "widgets_values" in save_crypto_node
            else [None]
        )
        save_crypto_node["widgets_values"].append("")
        save_crypto_node["properties"] = {"Node name for S&R": "DecodeCryptoNode"}
        if "outputs" not in crypto_bridge_node:
            save_crypto_node["outputs"] = []
        else:
            save_crypto_node["outputs"] = copy.deepcopy(crypto_bridge_node["outputs"])
        output_nodes_ids = [int(node["id"]) for node in simplify_workflow["nodes"]]
        filtered_links = []
        for link in simplify_workflow["links"]:
            if len(link) < 5:
                continue
            if link[1] == self.crypto_bridge_node_id:
                link[1] = self.save_crypto_node_id
            if link[1] in output_nodes_ids and link[3] in output_nodes_ids:
                filtered_links.append(link)
            else:
                pass
        simplify_workflow["links"] = filtered_links
        simplify_workflow.pop("groups", None)
        save_dir = folder_paths.output_directory
        with open(
            os.path.join(folder_paths.output_directory, output_workflow_name),
            "w",
            encoding="utf-8",
        ) as f:
            json.dump(simplify_workflow, f, ensure_ascii=False, indent=4)
        return save_dir

    def save_original_workflow(self, output_workflow_name, save_dir):
        with open(
            os.path.join(save_dir, output_workflow_name), "w", encoding="utf-8"
        ) as f:
            json.dump(self.workflow, f, ensure_ascii=False, indent=4)
        return save_dir

    def save_original_prompt(self, output_prompt_name, save_dir):
        with open(
            os.path.join(save_dir, output_prompt_name), "w", encoding="utf-8"
        ) as f:
            json.dump(self.prompt, f, ensure_ascii=False, indent=4)
        return save_dir


class DecodeCryptoWorkflow:
    def __init__(self, prompt, workflow, template_id):
        self.prompt = prompt
        self.workflow = workflow
        self.template_id = template_id
        self.crypto_result = {}
        self.input_anything_map = {}

    def calculate_input_anything_map(self):
        self.input_anything_map.clear()
        for node_id, node in self.prompt.items():
            if node.get("class_type") == "DecodeCryptoNode":
                for input_name, input_value in node.get("inputs", {}).items():
                    if input_name.startswith("input_anything"):
                        if isinstance(input_value, list) and len(input_value) == 2:
                            input_link_key = f"{input_value[0]}_{input_value[1]}"
                            self.input_anything_map[input_link_key] = input_name
        return self.input_anything_map

    def load_crypto_prompt(self, serial_number_token, user_token=None):
        if not user_token:
            user_token = AuthUnit().read_user_token()
        content = DownloadWorkflow().download_workflow(
            self.template_id, serial_number_token, user_token
        )
        if not isinstance(content, str) or content == "":
            raise ValueError("failed to download workflow")
        self.crypto_result = json.loads(content)
        return self.crypto_result["prompt"]

    def get_hidden_input(self, input_value):
        if isinstance(input_value, list) and len(input_value) == 2:
            key = f"{input_value[0]}_{input_value[1]}"
            return self.input_anything_map.get(key, None)

    def get_outputs(self):
        return self.crypto_result["outputs"]


class WorkflowTrimHandler:
    @staticmethod
    def onprompt_handler(json_data: Dict[str, Any]) -> Dict[str, Any]:
        prompt = json_data["prompt"]
        has_new_component = False
        has_old_component = False
        for node in prompt.values():
            class_type = node.get("class_type")
            if class_type == "SaveCryptoNode":
                has_new_component = True
                break
            if class_type == "ExcuteCryptoNode":
                has_old_component = True
                break
        if has_new_component:
            user_token, error_msg, error_code = AuthUnit().get_user_token()
            print(
                f"user_token: {user_token}, error_msg: {error_msg}, error_code: {error_code}"
            )
            if not user_token:
                if error_code == 401 or error_code == -3:
                    AuthUnit().login_dialog("输出CryptoCat加密工作流，请先完成登录")
                    json_data["prompt"] = {}
                else:
                    PromptServer.instance.send_sync(
                        "cryptocat_toast",
                        {"content": f"无法完成鉴权登录，{error_msg}", "type": "error"},
                    )
        elif has_old_component:
            prompt = WorkflowTrimHandler.replace_prompt(prompt)
            json_data["prompt"] = prompt
        return json_data

    @staticmethod
    def replace_prompt(prompt: Dict[str, Any]) -> Dict[str, Any]:
        if not prompt:
            raise ValueError("Invalid JSON format.")
        crypto_file_path = ""
        excute_crypto_id = None
        for node_id, node in list(prompt.items()):
            if node.get("class_type") == "ExcuteCryptoNode":
                crypto_file_path = node["inputs"].get("crypto_file_path")
                excute_crypto_id = node_id
                del prompt[node_id]
                break
        if not crypto_file_path:
            print("No 'ExcuteCryptoNode' found in prompt.")
            return prompt
        inject_json = FileCompressor.decompress_from_json(crypto_file_path, "19040822")
        output_images_ids = inject_json.pop("output_images_ids")
        for node in prompt.values():
            if node.get("class_type") == "CryptoCatImage":
                if excute_crypto_id in node["inputs"]["images"]:
                    node["inputs"]["images"] = output_images_ids
        random_seed_node = next(
            (
                node
                for node in inject_json.values()
                if node.get("class_type") == "RandomSeedNode"
            ),
            None,
        )
        if random_seed_node:
            random_seed_node["inputs"]["is_changed"] = random.randint(0, 999999)
        prompt.update(inject_json)
        return prompt