RiceRound
d62754850e
Added serial number validity period and days feature. Fixed the bug where the serial number could only be used for 7 days. Code optimization.
267 lines
8.2 KiB
Python
267 lines
8.2 KiB
Python
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),)
|