diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..574d856
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,76 @@
+_c='RiceRoundOutputImageNode'
+_b='RiceRoundEncryptNode'
+_a='RiceRoundOutputTextNode'
+_Z='RiceRoundOutputBooleanNode'
+_Y='RiceRoundOutputFloatNode'
+_X='RiceRoundOutputIntNode'
+_W='RiceRoundOutputMaskBridgeNode'
+_V='RiceRoundUploadImageNode'
+_U='RiceRoundImageUrlNode'
+_T='RiceRoundOutputImageBridgeNode'
+_S='RiceRoundDecryptNode'
+_R='RiceRoundStrToBooleanNode'
+_Q='RiceRoundStrToFloatNode'
+_P='RiceRoundStrToIntNode'
+_O='RiceRoundBooleanNode'
+_N='RiceRoundFloatNode'
+_M='RiceRoundIntNode'
+_L='RiceRoundDownloadMaskNode'
+_K='RiceRoundMaskBridgeNode'
+_J='RiceRoundInputTextNode'
+_I='RiceRoundRandomSeedNode'
+_H='RiceRoundDownloadImageNode'
+_G='RiceRoundSimpleImageNode'
+_F='RiceRoundImageBridgeNode'
+_E='RiceRoundDebugNode'
+_D='RiceRoundAdvancedChoiceNode'
+_C='RiceRoundSimpleChoiceNode'
+_B='dynamic_class'
+_A='display_name'
+import platform,random
+from urllib.parse import unquote
+import subprocess
+from server import PromptServer
+from.input_node import*
+from.output_node import*
+from.encrypt_node import*
+from.auth_unit import AuthUnit
+from aiohttp import web
+from functools import partial
+from.rice_prompt_handler import RiceRoundPromptHandler
+from.rice_url_config import RiceUrlConfig
+from.rice_prompt_info import RicePromptInfo
+def create_dynamic_nodes(base_class):
+ rice_prompt_info=RicePromptInfo();dynamic_classes={}
+ for(node_name,info)in rice_prompt_info.choice_classname_map.items():class_name=node_name;category='RiceRound/Advanced/Choice';dynamic_class=type(class_name,(base_class,),{'__node_name__':node_name,'CATEGORY':category});dynamic_classes[class_name]={_B:dynamic_class,_A:info.get(_A,class_name)}
+ return dynamic_classes
+dynamic_choice_nodes=create_dynamic_nodes(RiceRoundBaseChoiceNode)
+NODE_CLASS_MAPPINGS={_C:RiceRoundSimpleChoiceNode,_D:RiceRoundAdvancedChoiceNode,_E:RiceRoundDebugNode,_F:RiceRoundImageBridgeNode,_G:RiceRoundSimpleImageNode,_H:RiceRoundDownloadImageNode,_I:RiceRoundRandomSeedNode,_J:RiceRoundInputTextNode,_K:RiceRoundMaskBridgeNode,_L:RiceRoundDownloadMaskNode,_M:RiceRoundIntNode,_N:RiceRoundFloatNode,_O:RiceRoundBooleanNode,_P:RiceRoundStrToIntNode,_Q:RiceRoundStrToFloatNode,_R:RiceRoundStrToBooleanNode,_S:RiceRoundDecryptNode,_T:RiceRoundOutputImageBridgeNode,_U:RiceRoundImageUrlNode,_V:RiceRoundUploadImageNode,_W:RiceRoundOutputMaskBridgeNode,_X:RiceRoundOutputIntNode,_Y:RiceRoundOutputFloatNode,_Z:RiceRoundOutputBooleanNode,_a:RiceRoundOutputTextNode,_b:RiceRoundEncryptNode,_c:RiceRoundOutputImageNode,**{name:cls[_B]for(name,cls)in dynamic_choice_nodes.items()}}
+NODE_DISPLAY_NAME_MAPPINGS={_C:'Simple Choice',_D:'Advanced Choice',_E:'Debug',_F:'Image Bridge',_G:'Simple Image',_H:'Download Image',_I:'Random Seed',_J:'Input Text',_K:'Mask Bridge',_L:'Download Mask',_M:'RiceRound Int',_N:'RiceRound Float',_O:'RiceRound Boolean',_P:'RiceRound Str To Int',_Q:'RiceRound Str To Float',_R:'RiceRound Str To Boolean',_S:'Decrypt',_T:'Output Image Bridge',_U:'Image URL',_V:'Upload Image',_W:'Output Mask Bridge',_X:'Output Int',_Y:'Output Float',_Z:'Output Boolean',_a:'Output Text',_b:'Encrypt',_c:'Output Image',**{name:cls[_A]for(name,cls)in dynamic_choice_nodes.items()}}
+WEB_DIRECTORY='./js'
+__all__=['NODE_CLASS_MAPPINGS','NODE_DISPLAY_NAMES_MAPPINGS','WEB_DIRECTORY']
+handler_instance=RiceRoundPromptHandler()
+onprompt_callback=partial(handler_instance.onprompt_handler)
+PromptServer.instance.add_on_prompt_handler(onprompt_callback)
+routes=PromptServer.instance.routes
+url_config=RiceUrlConfig()
+@routes.post('/riceround/set_comfyui_web_server_info')
+async def set_comfyui_web_server_info(request):server_info=await request.json();url_config.set_comfyui_local_web_server_info(server_info);return web.json_response({},status=200)
+@routes.get('/riceround/auth_callback')
+async def auth_callback(request):
+ auth_query=request.query;token=auth_query.get('token','')
+ if token:token=unquote(token);AuthUnit().save_user_token(token)
+ raise web.HTTPFound('/')
+@routes.post('/riceround/set_node_additional_info')
+async def set_node_additional_info(request):additional_info=await request.json();RicePromptInfo().set_node_additional_info(additional_info);return web.json_response({},status=200)
+@routes.get('/riceround/open_selector_list_folder')
+async def open_selector_list_folder(request):
+ A='error';choice_server_folder=RicePromptInfo().get_choice_server_folder()
+ if not choice_server_folder.exists():return web.json_response({A:'Folder does not exist'},status=404)
+ system=platform.system()
+ try:
+ if system=='Windows':os.startfile(choice_server_folder)
+ elif system=='Darwin':subprocess.run(['open',choice_server_folder])
+ else:subprocess.run(['xdg-open',choice_server_folder])
+ return web.json_response({'status':'success'},status=200)
+ except Exception as e:return web.json_response({A:str(e)},status=500)
\ No newline at end of file
diff --git a/auth_unit.py b/auth_unit.py
new file mode 100644
index 0000000..c74e734
--- /dev/null
+++ b/auth_unit.py
@@ -0,0 +1,65 @@
+_D='user_token'
+_C=None
+_B='Auth'
+_A='utf-8'
+import base64,json,os
+from pathlib import Path
+import time
+from urllib.parse import urljoin
+from.utils import get_local_app_path,get_machine_id
+from.rice_url_config import RiceUrlConfig
+from server import PromptServer
+import urllib.request,urllib.error,configparser
+class AuthUnit:
+ _instance=_C
+ def __new__(A,*B,**C):
+ if A._instance is _C:A._instance=super(AuthUnit,A).__new__(A)
+ return A._instance
+ def __init__(A):
+ B=True
+ if not hasattr(A,'initialized'):A.machine_id=get_machine_id();A.url_config=RiceUrlConfig();A.callback_path='/riceround/auth_callback';A.shift_key=sum(ord(A)for A in A.machine_id[:8])%20+1;C=get_local_app_path();C.mkdir(parents=B,exist_ok=B);A.config_path=C/'config.ini';A.auto_login=False;A.load_config();A.last_check_time=0;A.initialized=B
+ def load_config(A):B=configparser.ConfigParser();B.read(A.config_path,encoding=_A);A.auto_login=B.get(_B,'auto_login',fallback=False)
+ def get_user_token(B):
+ A=B.read_user_token()
+ if A and time.time()-B.last_check_time>120:
+ try:
+ D=urllib.request.Request(B.url_config.login_api_url,headers={'Content-Type':'application/json','Authorization':f"Bearer {A}"},method='GET')
+ with urllib.request.urlopen(D)as C:
+ if C.status!=200:A=_C;print(f"Auth failed, token: {A}")
+ else:E=json.loads(C.read().decode(_A));B.url_config.set_server_info(E);B.last_check_time=time.time();return A
+ except urllib.error.URLError:A=_C
+ if not A:B.login_in()
+ return A
+ def login_in(A):
+ PromptServer.instance.send_sync('rice_round_login',{'machine_id':A.machine_id,'url_prefix':A.url_config.web_url_prefix})
+ if A.auto_login:import webbrowser as C;B='';D=urljoin(A.url_config.comfyui_local_base_url,A.callback_path);E=base64.b64encode(D.encode(_A)).decode(_A);B='&callback_url='+E;F=A.url_config.auth_web_url+'?machine_id='+A.machine_id+B;C.open(F)
+ @staticmethod
+ def _encrypt(text,shift_key):
+ 'Encrypt text using a simple character shift and base64 encoding.'
+ try:A=''.join(chr((ord(A)+shift_key)%65536)for A in text);return base64.b64encode(A.encode(_A)).decode(_A)
+ except Exception as B:print(f"Encryption error: {B}");return text
+ @staticmethod
+ def _decrypt(encoded_text,shift_key):
+ 'Decrypt text that was encrypted with the _encrypt method.';A=encoded_text
+ try:B=base64.b64decode(A.encode(_A)).decode(_A);return''.join(chr((ord(A)-shift_key)%65536)for A in B)
+ except Exception as C:print(f"Decryption error: {C}");return A
+ def read_user_token(A):
+ 'Retrieve and decrypt the user token.'
+ if not os.path.exists(A.config_path):return''
+ try:
+ B=configparser.ConfigParser();B.read(A.config_path,encoding=_A);C=B.get(_B,'test_user_token',fallback='')
+ if C:return C
+ D=B.get(_B,_D,fallback='')
+ if not D:return''
+ return AuthUnit._decrypt(D,A.shift_key)
+ except Exception as E:print(f"Error reading token: {E}");return''
+ def save_user_token(B,user_token):
+ 'Encrypt and save the user token.';C=user_token
+ if not C:return
+ try:
+ A=configparser.ConfigParser()
+ if os.path.exists(B.config_path):A.read(B.config_path,encoding=_A)
+ if _B not in A:A.add_section(_B)
+ E=AuthUnit._encrypt(C,B.shift_key);A[_B][_D]=E
+ with open(B.config_path,'w',encoding=_A)as F:A.write(F)
+ except Exception as D:print(f"Error saving token: {D}");raise RuntimeError(f"Failed to save token: {D}")
\ No newline at end of file
diff --git a/encrypt_node.py b/encrypt_node.py
new file mode 100644
index 0000000..287cc72
--- /dev/null
+++ b/encrypt_node.py
@@ -0,0 +1,248 @@
+_A0='00000000-0000-0000-0000-000000000000'
+_z='task_id'
+_y='RiceRoundImageUrlNode'
+_x='RiceRoundUploadImageNode'
+_w='RiceRoundOutputImageBridgeNode'
+_v='Node name for S&R'
+_u='properties'
+_t='RiceRoundEncryptNode'
+_s='input_anything'
+_r='main_link_type'
+_q='last_node_id'
+_p='owner_id'
+_o='RiceRound'
+_n='EXTRA_PNGINFO'
+_m='PROMPT'
+_l='UNIQUE_ID'
+_k='template_id'
+_j='hidden'
+_i='optional'
+_h='required'
+_g='RiceRoundInputTextNode'
+_f='class_type'
+_e='main_link_id'
+_d='slot_index'
+_c='RiceRoundAdvancedChoiceNode'
+_b='RiceRoundSimpleChoiceNode'
+_a='RiceRoundDecryptNode'
+_Z='extra_pnginfo'
+_Y='prompt'
+_X='RiceRoundStrToBooleanNode'
+_W='RiceRoundBooleanNode'
+_V='RiceRoundStrToFloatNode'
+_U='RiceRoundStrToIntNode'
+_T='RiceRoundFloatNode'
+_S='RiceRoundIntNode'
+_R='RiceRoundDownloadMaskNode'
+_Q='RiceRoundMaskBridgeNode'
+_P='RiceRoundImageBridgeNode'
+_O='RiceRoundSimpleImageNode'
+_N='name'
+_M='unique_id'
+_L='RiceRoundDownloadImageNode'
+_K='IMAGE'
+_J='default'
+_I='inputs'
+_H='STRING'
+_G='label'
+_F='outputs'
+_E='nodes'
+_D='links'
+_C='id'
+_B=None
+_A='type'
+from collections import defaultdict
+import copy
+from io import BytesIO
+import json,os,uuid,numpy as np,comfy.utils,time
+from PIL import Image
+from.utils import combine_files
+from.rice_url_config import machine_upload_image
+if __name__=='__main__':output_project_folder='D:\\output'
+else:import folder_paths;from server import PromptServer;from.rice_url_config import RiceUrlConfig;from.rice_prompt_info import RicePromptInfo;output_project_folder=folder_paths.output_directory
+INPUT_NODE_TYPES=[_b,_c,_O,_L,_P,_g,_Q,_R,_S,_T,_U,_V,_W,_X]
+class RiceRoundEncryptNode:
+ def __init__(A):A.template_id=uuid.uuid4().hex
+ @classmethod
+ def INPUT_TYPES(A):return{_h:{'project_name':(_H,{_J:'my_project'}),_k:(_H,{_J:uuid.uuid4().hex})},_i:{'output':(_K,)},_j:{_M:_l,_Y:_m,_Z:_n}}
+ @classmethod
+ def IS_CHANGED(A,**B):return float('NaN')
+ RETURN_TYPES=();OUTPUT_NODE=True;FUNCTION='encrypt';CATEGORY=_o
+ def encrypt(E,project_name,template_id,output,**A):F=A.pop(_M,_B);B=A.pop(_Z,_B);C=A.pop(_Y,_B);D=Encrypt(B['workflow'],C,project_name,template_id);D.do_encrypt();return{}
+class Encrypt:
+ def __init__(A,workflow,prompt,project_name,template_id):
+ A.original_workflow=workflow;A.original_prompt=prompt;A.template_id=template_id;A.project_name=project_name;A.project_folder=os.path.join(output_project_folder,A.project_name)
+ if not os.path.exists(A.project_folder):os.makedirs(A.project_folder)
+ A.last_node_id=0;A.last_link_id=0;A.link_owner_map=defaultdict(dict);A.workflow_nodes_dict={};A.node_prompt_map={};A.input_node_map={};A.related_node_ids=set()
+ def do_encrypt(A):A.invalid_workflow(A.original_workflow);A.load_workflow();A.load_prompt();A.analyze_input_from_workflow();A.assemble_new_workflow();A.output_template_json_file();A.assemble_new_prompt();RicePromptInfo().save_choice_classname(A.project_folder);A.output_file(A.original_workflow,f"original_workflow");A.output_file(A.original_prompt,f"original_prompt");A.save_rice_zip();A.clear()
+ def clear(A):A.original_workflow=_B;A.original_prompt=_B;A.template_id=_B;A.project_name=_B;A.project_folder=_B;A.last_node_id=0;A.last_link_id=0;RicePromptInfo().clear()
+ def load_workflow(A):
+ C=copy.deepcopy(A.original_workflow);A.workflow_nodes_dict={int(A[_C]):A for A in C[_E]}
+ for F in C[_E]:
+ G=F.get(_F,[])
+ if not G:continue
+ for D in G:
+ E=D.get(_D,[])
+ if not E:continue
+ for B in E:B=int(B);A.link_owner_map[B][_D]=copy.deepcopy(E);A.link_owner_map[B][_d]=D.get(_d,0);A.link_owner_map[B][_p]=int(F[_C]);A.link_owner_map[B][_A]=D.get(_A,'')
+ A.last_node_id=int(C[_q]);A.last_link_id=int(C['last_link_id'])
+ def load_prompt(A):B=copy.deepcopy(A.original_prompt);A.node_prompt_map={int(A):B for(A,B)in B.items()}
+ def analyze_input_from_workflow(A):
+ for(id,B)in A.workflow_nodes_dict.items():
+ E=B.get(_A,'')
+ if E in INPUT_NODE_TYPES:
+ A.input_node_map[id]=copy.deepcopy(B);C=B.get(_F,[])
+ if not C:continue
+ D=C[0].get(_D,[])
+ if not D:continue
+ F=int(D[0]);A.input_node_map[id][_e]=F;A.input_node_map[id][_r]=C[0].get(_A,_H)
+ def assemble_new_workflow(A):C=list(A.input_node_map.keys());B=copy.deepcopy(A.original_workflow);A.related_node_ids=A.find_workflow_related_nodes(B[_D],C);B[_E]=[B for B in B[_E]if int(B[_C])in A.related_node_ids];A.invalid_new_workflow(B);D=A.add_decrypt_node(B);A.remove_redundant_links(B);A.remove_unrelated_nodes(B,A.related_node_ids,D);A.replace_choice_template(B);A.replace_workflow_node(B);A.output_file(B,f"workflow")
+ def output_template_json_file(G):
+ U='模型选择';T='提示词';S='请上传蒙版';R='image/*';Q='tip';P='max_size';O='accept';N='请上传图片';L='number';K='max';J='min';I='数值';F='settings';E='describe';V=RicePromptInfo();M=[]
+ for(H,W)in G.input_node_map.items():
+ X=W[_s];C=G.workflow_nodes_dict[H][_A];D=G.node_prompt_map[H].get(_I,{});A=str(D.get(_N,''));B={_C:str(X),_A:'',E:'输入组件','node_id':str(H),F:{}}
+ if C in[_O,_L,_P]:B[_A]='image_upload';B[E]=A if A else N,;B[F]={_G:A if A else N,O:R,P:500000,Q:'请上传不超过500KB的图片'}
+ elif C in[_Q,_R]:B[_A]='mask_upload';B[E]=A if A else S,;B[F]={_G:A if A else S,O:R,P:50000,Q:'请上传不超过50KB的图片'}
+ elif C==_g:B[_A]='text';B[E]=A if A else T,;B[F]={_G:A if A else T,'placeholder':'请描述图片内容','multiline':True}
+ elif C==_c or C==_b:B[_A]='choice';B[E]=A if A else U,;B[F]={_G:A if A else U,'options':V.get_choice_value(H),_J:D.get(_J,'')}
+ elif C==_S or C==_U:B[_A]='number_int';B[E]=A if A else I,;B[F]={_G:A if A else I,J:D.get(J,0),K:D.get(K,1000),L:D.get(L,0)}
+ elif C==_T or C==_V:B[_A]='number_float';B[E]=A if A else I,;B[F]={_G:A if A else I,J:D.get(J,.0),K:D.get(K,1e3),L:D.get(L,.0)}
+ elif C==_W or C==_X:B[_A]='switch';B[E]=A if A else'开关';B[F]={_G:A if A else'开关',_J:D.get('value',False)}
+ else:raise ValueError(f"Error: The node {H} is not a valid RiceRound node.")
+ M.append(B)
+ Y={_k:G.template_id,'title':G.project_name,'elements':M};G.output_file(Y,f"{G.template_id}_template")
+ def assemble_new_prompt(A):'\n 组装新的prompt配置。主要完成:\n 1. 移除不需要的节点\n 2. 转换特定节点的类型和输入\n 3. 保存处理后的prompt配置\n ';B=A._create_filtered_prompt();A._transform_node_types(B);A.output_file(B,f"{A.template_id}_job")
+ def _create_filtered_prompt(B):
+ '\n 创建经过过滤的prompt副本,移除不需要的节点\n ';A=copy.deepcopy(B.original_prompt);C=B._get_exclude_node_ids(A)
+ for D in C:A.pop(str(D),_B)
+ return A
+ def save_rice_zip(A):
+ '\n 将工作流相关文件打包成rice.zip\n 包括:\n - workflow.json\n - {template_id}_job.json\n - {template_id}_template.json \n - choice_node文件夹(如果存在)\n - original_workflow.json\n - original_prompt.json\n ';J='job.bin';import pyzipper as F
+ try:
+ C=[];B=[];G=[];D=os.path.join(A.project_folder,J)
+ for K in[f"{A.template_id}_job.json",'original_workflow.json','original_prompt.json']:H=os.path.join(A.project_folder,K);G.append(H);B.append(H)
+ combine_files(G,A.template_id,D);C.append((D,J));B.append(D);I=os.path.join(A.project_folder,f"{A.template_id}_template.json");C.append((I,'template.bin'));B.append(I);L=os.path.join(A.project_folder,'rice.zip')
+ with F.AESZipFile(L,'w',compression=F.ZIP_DEFLATED)as M:
+ for(E,N)in C:M.write(E,N)
+ print(f"save zip success")
+ for E in B:os.remove(E)
+ except Exception as O:print(f"Error creating zip: {str(O)}");raise
+ def _get_exclude_node_ids(A,prompt):
+ '\n 获取需要从prompt中排除的节点ID集合\n ';C={_t,_a,'RiceRoundDebugNode'};B=A.related_node_ids.difference(set(A.input_node_map.keys()))
+ for(D,E)in prompt.items():
+ if E.get(_f,'')in C:B.add(int(D))
+ return B
+ def _transform_node_types(K,prompt):
+ '\n 转换节点类型和更新节点输入配置\n ';J='image_url';D='str';B='new_inputs';A='new_type';E={_P:{A:_L,B:{J:''}},_O:{A:_L,B:{J:''}},_Q:{A:_R,B:{'mask_url':''}},_S:{A:_U,B:{D:''}},_T:{A:_V,B:{D:''}},_W:{A:_X,B:{D:''}}}
+ for(L,C)in prompt.items():
+ C.pop('is_changed',_B);F=C.get(_f,'');G=C.get(_I,{})
+ if not G:continue
+ H=G.get(_N,'')
+ if F in E:
+ I=E[F];C[_f]=I[A];C[_I]=I[B].copy()
+ if H:C[_I][_N]=H
+ def add_decrypt_node(A,workflow):
+ C=workflow;G=set();A.last_node_id+=1;H={_C:A.last_node_id,_A:_a,'pos':[420,0],'size':[500,150],'flags':{},'mode':0,'order':20,_I:[],_F:[{_N:_K,_A:_K,_D:[],_G:_K,_d:0}],_u:{_v:_a},'widgets_values':[str(A.template_id)]}
+ for(B,(D,E))in enumerate(A.input_node_map.items()):
+ I=E[_e];F=E[_r];E[_s]=B;J={_N:f"input_anything{B if B>0 else''} ({D})",_A:'*','link':I,_G:f"input_anything{B if B>0 else''} ({D})"}
+ if B==0:J['shape']=7
+ H[_I].append(J)
+ if F not in[_K,_H]:F=_H
+ K=[I,D,0,A.last_node_id,B,F];C[_D].append(K)
+ G.add(A.last_node_id);C[_E].append(H);C[_q]=A.last_node_id;return G
+ def output_file(A,workflow,prefix):
+ B=os.path.join(A.project_folder,f"{prefix}.json")
+ with open(B,'w',encoding='utf-8')as C:json.dump(workflow,C,ensure_ascii=False,indent=4)
+ def remove_redundant_links(D,workflow):
+ B=workflow;E=set()
+ for C in B[_E]:
+ F=int(C[_C])
+ if F in D.input_node_map:
+ G=D.input_node_map[F][_e];A=C.get(_F,[])
+ if not A:continue
+ if len(A)!=1:raise ValueError(f"Error: The node {C[_C]} has an invalid number of outputs.")
+ H=A[0].get(_D,[])
+ if not H:continue
+ for I in H:
+ if I!=G:E.add(I)
+ A[0][_D]=[G]
+ B[_D]=[A for A in B[_D]if isinstance(A,list)and len(A)==6 and A[0]not in E]
+ def replace_choice_template(H,workflow):
+ D='extra';C=workflow;E=RicePromptInfo()
+ for A in C[_E]:
+ B=int(A[_C])
+ if A.get(_A,'')==_c:
+ F=E.get_choice_classname(B)
+ if F:A[_A]=F
+ else:print(f"Warning: The node {B} is not a valid RiceRound Choice node.")
+ G={}
+ for(B,A)in H.input_node_map.items():
+ if A.get(_A,'')==_b:I=E.get_choice_value(B);G[B]=I
+ if D not in C:C[D]={}
+ C[D]['choice_node_map']=G
+ def replace_workflow_node(J,workflow):
+ I='unknown';D=workflow;C='RiceRoundOutputTextNode';E={_P:_w,_O:_x,_L:_y,_Q:'RiceRoundOutputMaskBridgeNode',_R:'RiceRoundMaskUrlNode',_S:'RiceRoundOutputIntNode',_T:'RiceRoundOutputFloatNode',_W:'RiceRoundOutputBooleanNode',_X:C,_U:C,_V:C};F=set()
+ for A in D[_E]:
+ G=A.get(_A,'')
+ if G in E:
+ if _F not in A:raise ValueError(f"Node {A.get(_C,I)} missing outputs")
+ if not A[_F]or not isinstance(A[_F],list):raise ValueError(f"Invalid outputs format in node {A.get(_C,I)}")
+ H=E[G];A.update({_A:H,_F:[{_A:_H,**A[_F][0]}],_u:{_v:H}});F.add(int(A[_C]))
+ for B in D[_D]:
+ if len(B)==6 and B[1]in F:B[5]=_H
+ def remove_unrelated_nodes(E,workflow,related_node_ids,new_node_ids):
+ B=workflow;C=[];D=related_node_ids.union(new_node_ids)
+ for A in B[_D]:
+ if len(A)==6:
+ if A[1]in D and A[3]in D:C.append(A)
+ B[_D]=C
+ @staticmethod
+ def invalid_workflow(workflow):
+ C=_B;D=0
+ for A in workflow[_E]:
+ B=A.get(_A,'')
+ if B=='RiceRoundOutputImageNode':C=A;continue
+ elif B==_t:D+=1;continue
+ if B in[_a,_w,_x,_y]:raise ValueError(f"Error: The node {A[_C]} is not a valid RiceRound node.")
+ if D!=1:raise ValueError('Error: Multiple RiceRoundEncryptNode nodes are not allowed.')
+ if C is _B:raise ValueError("Error: The node is not an 'RiceRoundOutputImageNode'.")
+ def invalid_new_workflow(A,workflow):
+ for B in workflow[_E]:
+ C=B.get(_I,[])
+ for input in C:
+ D=int(input.get('link',0));E=A.link_owner_map[D][_p];F=A.workflow_nodes_dict[E][_A]
+ if F in INPUT_NODE_TYPES:raise ValueError(f"Error: The node {B[_C]} may have circular references, generation failed.")
+ def find_workflow_related_nodes(F,links,input_ids):
+ E=input_ids;B=set(E);C=list(E)
+ while C:
+ G=C.pop()
+ for D in links:
+ if len(D)==6 and D[3]==G:
+ A=D[1]
+ if A not in B:
+ if A in F.workflow_nodes_dict:B.add(A)
+ C.append(A)
+ return B
+class RiceRoundOutputImageNode:
+ def __init__(A):A.url_config=RiceUrlConfig()
+ @classmethod
+ def INPUT_TYPES(A):return{_h:{'images':(_K,)},_i:{_z:(_H,{_J:_A0})},_j:{_M:_l,_Y:_m,_Z:_n}}
+ RETURN_TYPES=();OUTPUT_NODE=True;FUNCTION='load';CATEGORY=_o
+ def load(S,images,task_id,**C):
+ L='image_type';E='PNG';B=task_id;A=images;F=C.pop(_M,_B);M=C.pop(_Y,_B);T=C.pop(_Z,_B);G=PromptServer.instance.client_id;H=''
+ if hasattr(PromptServer.instance,'last_prompt_id')and PromptServer.instance.last_prompt_id:H=PromptServer.instance.last_prompt_id
+ if F is _B:raise Exception("Warning: 'unique_id' is missing.")
+ if M is _B:raise Exception("Warning: 'prompt' is missing.")
+ if B==_A0 or B=='':
+ N=comfy.utils.ProgressBar(A.shape[0]);I=0
+ for D in A:O=255.*D.cpu().numpy();P=Image.fromarray(np.clip(O,0,255).astype(np.uint8));N.update_absolute(I,A.shape[0],(E,P,_B));I+=1
+ return{}
+ else:
+ print(f"RiceRoundOutputImageNode task_id: {B}")
+ if A.shape[0]>5:raise ValueError('Error: Cannot upload more than 5 images.')
+ J=[]
+ for D in A:
+ K=machine_upload_image(D,B)
+ if not K:raise ValueError('Error: Failed to upload image.')
+ J.append(K)
+ Q={L:E,'image_results':J};R={_z:B,_M:F,'client_id':G,'prompt_id':H,'timestamp':int(time.time()*1000),L:E,'result_data':Q};PromptServer.instance.send_sync('rice_round_done',R,sid=G);return{}
\ No newline at end of file
diff --git a/input_node.py b/input_node.py
new file mode 100644
index 0000000..30d8ae1
--- /dev/null
+++ b/input_node.py
@@ -0,0 +1,150 @@
+_X='number'
+_W='bridge'
+_V='only image.'
+_U='load_image'
+_T='prompt'
+_S='unique_id'
+_R='BOOLEAN'
+_Q='str'
+_P='MASK'
+_O='tooltip'
+_N='hidden'
+_M='数值'
+_L='IMAGE'
+_K='optional'
+_J='FLOAT'
+_I='load'
+_H='INT'
+_G='name'
+_F='value'
+_E='required'
+_D=True
+_C='RiceRound/Input'
+_B='STRING'
+_A='default'
+import hashlib,json,os,re,time,random
+from PIL import Image,ImageOps
+from PIL.PngImagePlugin import PngInfo
+from.rice_prompt_info import RicePromptInfo
+import folder_paths
+from nodes import LoadImage
+import requests,sys
+from.utils import pil2tensor
+from comfy.utils import ProgressBar
+class _BasicTypes(str):
+ basic_types=[_B]
+ def __eq__(B,other):A=other;return A in B.basic_types or isinstance(A,(list,_BasicTypes))
+ def __ne__(A,other):return not A.__eq__(other)
+BasicTypes=_BasicTypes('BASIC')
+class RiceRoundSimpleChoiceNode:
+ def __init__(A):A.prompt_info=RicePromptInfo()
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:'Parameter'}),_A:(_B,{_A:''})},_K:{},_N:{_S:'UNIQUE_ID',_T:'PROMPT','extra_pnginfo':'EXTRA_PNGINFO'}}
+ RETURN_TYPES=BasicTypes,;RETURN_NAMES=_F,;FUNCTION='placeholder';CATEGORY=_C
+ def placeholder(B,name,default,**C):
+ A=int(C.pop(_S,0));D=C.pop(_T,None);E=_D
+ if D:
+ for(G,F)in D.items():
+ if F.get('class_type','')=='RiceRoundDecryptNode':E=False;break
+ if E:
+ for H in range(10):
+ if A in B.prompt_info.choice_node_map:break
+ time.sleep(1)
+ if A not in B.prompt_info.choice_node_map:print(f"Warning: RiceRoundSimpleChoiceNode {A} not found in prompt_info.choice_node_map")
+ return default,
+class RiceRoundAdvancedChoiceNode(RiceRoundSimpleChoiceNode):
+ def __init__(A):super().__init__()
+class RiceRoundSimpleImageNode(LoadImage):
+ def __init__(A):super().__init__()
+ RETURN_TYPES=_L,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;CATEGORY=_C;FUNCTION=_U
+ def load_image(B,image):A,C=super().load_image(image);return A,
+class RiceRoundDownloadImageNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'image_url':(_B,{_A:''})},_K:{},_N:{}}
+ RETURN_TYPES=_L,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_U;CATEGORY=_C
+ def load_image(B,image_url,**C):A=Image.open(requests.get(image_url,stream=_D).raw);A=ImageOps.exif_transpose(A);return pil2tensor(A),
+class RiceRoundImageBridgeNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'images':(_L,{_O:_V})},_K:{}}
+ RETURN_TYPES=_L,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_W;CATEGORY=_C
+ def bridge(A,images,**B):return images,
+class RiceRoundMaskBridgeNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'mask':(_P,{_O:_V})}}
+ RETURN_TYPES=_P,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_W;CATEGORY=_C
+ def bridge(A,mask,**B):return mask,
+class RiceRoundDownloadMaskNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'mask_url':(_B,{_A:''})}}
+ RETURN_TYPES=_P,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION='load_mask';CATEGORY=_C
+ def load_mask(E,mask_url,**F):
+ C=mask_url
+ try:
+ D=requests.get(C,stream=_D,timeout=10);D.raise_for_status();A=Image.open(D.raw)
+ if A.mode!='L':A=A.convert('L')
+ return pil2tensor(A),
+ except requests.exceptions.RequestException as B:print(f"Error downloading mask from {C}: {str(B)}");raise
+ except Exception as B:print(f"Error processing mask: {str(B)}");raise
+class RiceRoundIntNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:_M}),_X:(_H,{_A:0}),'min':(_H,{_A:0}),'max':(_H,{_A:100})}}
+ RETURN_TYPES=_H,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(A,name,number,min,max,**B):return number,
+class RiceRoundStrToIntNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:_M}),_Q:(_B,)}}
+ RETURN_TYPES=_H,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(A,name,str,**B):return int(str),
+class RiceRoundFloatNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:_M}),_X:(_J,{_A:.0}),'min':(_J,{_A:.0}),'max':(_J,{_A:1e2})}}
+ RETURN_TYPES=_J,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(A,name,number,min,max,**B):return number,
+class RiceRoundStrToFloatNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:_M}),_Q:(_B,)}}
+ RETURN_TYPES=_J,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(A,name,str,**B):return float(str),
+class RiceRoundBooleanNode:
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:'开关'}),_F:(_R,{_A:False})}}
+ RETURN_TYPES=_R,;RETURN_NAMES=_F,;FUNCTION='execute';CATEGORY=_C
+ def execute(A,name,value):return value,
+class RiceRoundStrToBooleanNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{_G:(_B,{_A:'开关'}),_Q:(_B,)}}
+ RETURN_TYPES=_R,;RETURN_NAMES=_F,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(A,name,str,**B):return str.lower()=='true',
+class RiceRoundInputTextNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'text_info':(_B,{'multiline':_D,_O:'The text to be encoded.'})}}
+ RETURN_TYPES=_B,;OUTPUT_NODE=_D;FUNCTION=_I;CATEGORY=_C
+ def load(D,text_info,**E):
+ B=text_info;A=''
+ try:C=json.loads(B);A=C.get('content','')
+ except json.JSONDecodeError:A=B
+ return A,
+class RiceRoundRandomSeedNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{},_K:{},_N:{}}
+ RETURN_TYPES=_H,;FUNCTION='random';CATEGORY=_C
+ @classmethod
+ def IS_CHANGED(A):return random.randint(0,999999)
+ def random(B):A=random.randint(0,999999);print('产生随机数 ',A);return A,
+class RiceRoundDebugNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_E:{'text':(_B,{_A:'Hello, World!'})}}
+ RETURN_TYPES=_B,;FUNCTION='debug';CATEGORY=_C
+ def debug(A,text):print(f"RiceDebugNode text: {text}");return text,
\ No newline at end of file
diff --git a/js/riceconfig.js b/js/riceconfig.js
new file mode 100644
index 0000000..3dcb0ea
--- /dev/null
+++ b/js/riceconfig.js
@@ -0,0 +1 @@
+import{api}from"../../../scripts/api.js";import{ComfyApp,app}from"../../../scripts/app.js";app.registerExtension({name:"riceround.config",async setup(){app.ui.settings.addSetting({id:"RiceRound.Setting.name",name:"模型列表存放位置,可以制造需要的ChoiceNode",type:()=>{const e=document.createElement("tr"),t=document.createElement("td"),o=document.createElement("input");return o.type="button",o.value="打开文件夹",o.style.borderRadius="8px",o.style.padding="8px 16px",o.style.fontSize="14px",o.style.cursor="pointer",o.style.border="1px solid #666",o.style.backgroundColor="#444",o.style.color="#fff",o.onmouseover=()=>{o.style.backgroundColor="#555"},o.onmouseout=()=>{o.style.backgroundColor="#444"},o.onclick=()=>{api.fetchApi("/riceround/open_selector_list_folder",{method:"GET",headers:{"Content-Type":"application/json"}})},t.appendChild(o),e.appendChild(t),e}})}});
\ No newline at end of file
diff --git a/js/riceround.js b/js/riceround.js
new file mode 100644
index 0000000..e751d69
--- /dev/null
+++ b/js/riceround.js
@@ -0,0 +1 @@
+import{api}from"../../../scripts/api.js";import{ComfyApp,app}from"../../../scripts/app.js";const GetHost=()=>{const{protocol:e,hostname:t,port:o}=window.location,n=parseInt(o,10);let i=t;return 80===n&&"http:"===e||443===n&&"https:"===e?i=t:isNaN(n)||0===n||(i=`${t}:${o}`),i};function loadToastScript(){return new Promise(((e,t)=>{const o=document.createElement("script");o.src="https://cdn.jsdelivr.net/npm/toastify-js",o.onload=e,o.onerror=t,document.head.appendChild(o);const n=document.createElement("link");n.href="https://cdn.jsdelivr.net/npm/toastify-js/src/toastify.min.css",n.type="text/css",n.rel="stylesheet",document.head.appendChild(n)}))}let toastHasLoaded=!1;async function loadToast(){toastHasLoaded||(await loadToastScript(),toastHasLoaded=!0)}async function showToast(e,t="info",o=3e3){await loadToast(),"info"==t?Toastify({text:e,duration:o,close:!1,gravity:"top",position:"center",stopOnFocus:!1}).showToast():"error"==t?Toastify({text:e,duration:o,close:!0,gravity:"top",position:"center",backgroundColor:"#FF4444",stopOnFocus:!0}).showToast():"warning"==t&&Toastify({text:e,duration:o,close:!0,gravity:"top",position:"center",backgroundColor:"#FFA500",stopOnFocus:!0}).showToast()}function loadMessageBoxScript(){return new Promise(((e,t)=>{const o=document.createElement("script");o.src="https://cdn.jsdelivr.net/npm/sweetalert2@11",o.onload=e,o.onerror=t,document.head.appendChild(o)}))}let messageBoxHasLoaded=!1;async function loadMssageBox(){messageBoxHasLoaded||(await loadMessageBoxScript(),messageBoxHasLoaded=!0)}async function showHtmlMessageBox(e,t,o,n,i){await loadMssageBox(),Swal.fire({title:e,html:t,icon:o,confirmButtonText:n,preConfirm:()=>{i&&(window.location.href=i)}})}async function setNodeAdditionalInfo(e){try{const t={method:"POST",headers:{"Content-Type":"application/json"},body:JSON.stringify(e)},o=await api.fetchApi("/riceround/set_node_additional_info",t);await o.json()}catch(e){console.error("riceround set_node_additional_info error",e)}}function changeWidget(e,t,o,n){e.type=t,e.value=o,e.options=n}function changeWidgets(e,t,o,n){"customtext"===t&&(t="text"),console.log("changeWidgets",e,t,o,n);const i=n.options;var a=e.widgets[1].value;i?.values?.includes(a)||(a=n.value),"RiceRoundAdvancedChoiceNode"==e.comfyClass||"RiceRoundSimpleChoiceNode"==e.comfyClass?changeWidget(e.widgets[1],t,a,i):"RiceRoundIntNode"!=e.comfyClass&&"RiceRoundFloatNode"!=e.comfyClass||4==e.widgets.length&&(changeWidget(e.widgets[1],"number",a,i),changeWidget(e.widgets[2],"number",i?.min??0,i),changeWidget(e.widgets[3],"number",i?.max??100,i))}function adaptWidgetsBasedOnConnection(e,t,o,n){if(e.outputs[0].type=o.type,"name"===e.widgets[0].label){(["数值","文本","列表","参数","Parameter"].includes(e.widgets[0].value)||""==e.widgets[0].value)&&(e.widgets[0].value=o.label?o.label:o.name)}const i=o.widget?.name?o.widget?.name:o.name,a=t.widgets.find((e=>e.name===i));if(!a)return;changeWidgets(e,a.origType??a.type,t,a)}api.addEventListener("rice_round_toast",(e=>{showToast(e.detail.content,e.detail.type,e.detail.duration)})),api.addEventListener("rice_round_login",(e=>{const t=e.detail.machine_id,o=e.detail.url_prefix,n=GetHost(),i=`${window.location.protocol}//${n}`,a=`${o}/auth/login?machine_id=${t}&callback_url=${btoa(`${i}/riceround/auth_callback`)}`;showHtmlMessageBox("登录",`RiceRound登录失效,请点击链接重新登录`,"info","确定",a)})),api.addEventListener("execution_start",(async({detail:e})=>{let t="";const o={};for(const e of app.graph.nodes){if("RiceRoundDecryptNode"===e.type)return;if("RiceRoundEncryptNode"===e.type){const o=e.widgets?.find((e=>"template_id"===e.name&&e.value));if(o){if(t)return;t=o.value}}else if("RiceRoundAdvancedChoiceNode"===e.type||"RiceRoundSimpleChoiceNode"===e.type){if(1===!e.outputs?.[0]?.links?.length)continue;const t=e.widgets[1].options?.values??[];if(!t.length)continue;const n=e.graph.links[e.outputs[0].links[0]];if(!n)continue;const i=e.graph.getNodeById(n.target_id);if(!i?.inputs||"RiceRoundDecryptNode"===i.comfyClass||"RiceRoundEncryptNode"===i.comfyClass)continue;const a=i.inputs[n.target_slot];if(!a||!i.widgets)continue;const s=a.widget?.name||a.name;if(!s)continue;const d=`${i.comfyClass}.${s}`;o[e.id]={class_name:d,options_value:t,node_type:e.type}}}t&&Object.keys(o).length>0&&await setNodeAdditionalInfo({choice_node_map:o,template_id:t})}));let applySimpleChoiceNodeExtraLogicTimer=null;function applySimpleChoiceNodeExtraLogic(e,t){applySimpleChoiceNodeExtraLogicTimer&&clearTimeout(applySimpleChoiceNodeExtraLogicTimer),applySimpleChoiceNodeExtraLogicTimer=setTimeout((()=>{applySimpleChoiceNodeExtraLogicTimer=null;const o=t.graph.extra?.choice_node_map;if(e&&e.widgets&&2===e.widgets.length&&o&&o[e.id]){const t={values:o[e.id]};changeWidget(e.widgets[1],"combo",e.widgets[1].value,t)}}),200)}function adaptWidgetsToConnection(e){if(!e.outputs||0===e.outputs.length)return;const t=e.outputs[0].links;if(t&&1===t.length){const o=e.graph.links[t[0]];if(!o)return;const n=e.graph.getNodeById(o.target_id);if(!n||!n.inputs)return;if("RiceRoundDecryptNode"==n.comfyClass&&"RiceRoundSimpleChoiceNode"==e.comfyClass)return console.log("applySimpleChoiceNodeExtraLogic",e),void applySimpleChoiceNodeExtraLogic(e,app);if("RiceRoundEncryptNode"==n.comfyClass||"RiceRoundDecryptNode"==n.comfyClass)return;const i=n.inputs[o.target_slot];if(!i||void 0===n.widgets)return;adaptWidgetsBasedOnConnection(e,n,i,app)}else t&&0!==t.length||("RiceRoundAdvancedChoiceNode"==e.comfyClass||"RiceRoundSimpleChoiceNode"==e.comfyClass?(e.widgets[0].value="Parameter",e.outputs[0].type="*"):"RiceRoundIntNode"==e.comfyClass?4==e.widgets.length&&(e.widgets[0].value="数值",e.widgets[1].value=0,e.widgets[2].value=0,e.widgets[3].value=100):"RiceRoundFloatNode"==e.comfyClass&&4==e.widgets.length&&(e.widgets[0].value="数值",e.widgets[1].value=0,e.widgets[2].value=0,e.widgets[3].value=100))}function setupParameterNode(e){const t=e.prototype.onAdded;e.prototype.onAdded=function(){t?.apply(this,arguments),adaptWidgetsToConnection(this)};const o=e.prototype.onAfterGraphConfigured;e.prototype.onAfterGraphConfigured=function(){o?.apply(this,arguments),adaptWidgetsToConnection(this)};const n=e.prototype.onConnectOutput;e.prototype.onConnectOutput=function(e,t,o,i,a){return!(!o.widget&&!(o.type in["STRING","COMBO","combo"]))&&(!n||(result=n.apply(this,arguments),result))};const i=e.prototype.onConnectionsChange;e.prototype.onConnectionsChange=function(e,t,o,n,a){return 2!=e||o||!this?.type||"RiceRoundAdvancedChoiceNode"!=this.type&&"RiceRoundSimpleChoiceNode"!=this.type||(this.widgets[0].value="Parameter"),app.configuringGraph||adaptWidgetsToConnection(this),i?.apply(this,arguments)}}function generateUUID(){let e="";for(let t=0;t<32;t++){e+=Math.floor(16*Math.random()).toString(16)}return e}app.registerExtension({name:"riceround.custom",async beforeRegisterNodeDef(e,t,o){["RiceRoundAdvancedChoiceNode","RiceRoundSimpleChoiceNode","RiceRoundIntNode","RiceRoundFloatNode"].includes(t.name)&&setupParameterNode(e)},nodeCreated(e,t){if("RiceRoundEncryptNode"==e.comfyClass&&e.widgets&&e.widgets.length>0){const t=window.document.title;if(t){let o=t.replace(/\s*-\s*ComfyUI$/,"").replace(/[<>:"/\\|?*]/g,"_").replace(/^[^a-zA-Z]+|[^a-zA-Z]+$/g,"");for(let t=0;t{const t=generateUUID();for(let o=0;oTaskStatus.FINISHED:raise ValueError(C.progress_text)
+ else:raise ValueError('websocket failed')
+ K=J.get('image_results',[])
+ if not K:raise ValueError('Failed to get image results')
+ M=K[0];H=Image.open(requests.get(M,stream=_C).raw);H=ImageOps.exif_transpose(H);A.pbar=None;return pil2tensor(H),
+ def create_task(E,input_data,template_id,user_token):
+ '\n Create a task and return the task UUID.\n \n Args:\n task_url (str): The URL to send the task request to\n request_data (dict): The data to send in the request\n headers (dict): The headers to send with the request\n \n Returns:\n str: The task UUID if successful\n \n Raises:\n ValueError: If the request fails or response is invalid\n ';D='data';F=E.url_config.prompt_task_url;G={'Authorization':f"Bearer {user_token}",'Content-Type':'application/json'};H={'taskData':json.dumps(input_data),'workData':json.dumps({'template_id':template_id})};A=requests.post(F,json=H,headers=G)
+ if A.status_code==200:
+ B=A.json()
+ if B.get('code')==0 and D in B:
+ C=TaskInfo(B.get(D,{}))
+ if C.task_uuid:return C
+ else:raise ValueError('No task UUID in response')
+ else:raise ValueError(f"API error: {B.get('message','Unknown error')}")
+ else:raise ValueError(f"HTTP error {A.status_code}: {A.text}")
+class RiceRoundBaseChoiceNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(B):A=getattr(B,'__node_name__',None);C=RicePromptInfo().get_choice_node_options(A)if A else[];return{_D:{_H:(_A,{_B:'Parameter'}),_B:(C,)},_I:{},_L:{}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;FUNCTION='placeholder';CATEGORY=_E
+ def placeholder(A,default,**B):return default,
+def upload_images(images):
+ C=AuthUnit().get_user_token()
+ if not C:raise ValueError(_O)
+ A=[]
+ for D in images:E=user_upload_image(D,C);A.append(E)
+ B=''
+ if len(A)==1:B=A[0]
+ elif len(A)>1:B=','.join(A)
+ return RICE_ROUND_IMAGE_TAG+B,
+class RiceRoundImageUrlNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{'image_url':(_A,)}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_P;CATEGORY=_E
+ def load_image(A,image_url,**B):return RICE_ROUND_IMAGE_TAG+image_url,
+class RiceRoundUploadImageNode(LoadImage):
+ def __init__(A):super().__init__()
+ @classmethod
+ def INPUT_TYPES(C):A=folder_paths.get_input_directory();B=[B for B in os.listdir(A)if os.path.isfile(os.path.join(A,B))];return{_D:{'image':(sorted(B),{'image_upload':_C})},_I:{},_L:{}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_P;CATEGORY=_E
+ def load_image(B,image,**C):A,D=super().load_image(image);return upload_images(A[:1])
+class RiceRoundOutputImageBridgeNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{'images':('IMAGE',{_Q:'only image.'})},_I:{}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(A,images,**B):return upload_images(images)
+class RiceRoundOutputMaskBridgeNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{'mask':('MASK',)}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(E,mask,**F):
+ A=mask.cpu().numpy();A=(A*255).astype(np.uint8);C=Image.fromarray(A);B=AuthUnit().get_user_token()
+ if not B:raise ValueError(_O)
+ D=user_upload_image(C,B);return RICE_ROUND_MASK_TAG+D,
+class RiceRoundMaskUrlNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{'mask_url':(_A,)}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_P;CATEGORY=_E
+ def load_image(A,mask_url,**B):return RICE_ROUND_MASK_TAG+mask_url,
+class RiceRoundOutputIntNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{_H:(_A,{_B:'数值'}),_S:(_J,),_M:(_J,{_B:0}),_N:(_J,{_B:1000000})}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(A,name,number,min,max,**B):return str(number),
+class RiceRoundOutputFloatNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(B):A='FLOAT';return{_D:{_H:(_A,{_B:'数值'}),_S:(A,),_M:(A,{_B:.0}),_N:(A,{_B:1e6})}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(A,name,number,min,max,**B):return str(number),
+class RiceRoundOutputBooleanNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{_H:(_A,{_B:'开关'}),_F:('BOOLEAN',{_B:False})}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(B,name,value,**C):A='true'if value else'false';return A,
+class RiceRoundOutputTextNode:
+ def __init__(A):0
+ @classmethod
+ def INPUT_TYPES(A):return{_D:{_H:(_A,{_B:'文本'}),'str':(_A,)}}
+ RETURN_TYPES=_A,;RETURN_NAMES=_F,;OUTPUT_NODE=_C;FUNCTION=_G;CATEGORY=_E
+ def bridge(A,name,str,**B):return str,
\ No newline at end of file
diff --git a/rice_prompt_handler.py b/rice_prompt_handler.py
new file mode 100644
index 0000000..7db3cf8
--- /dev/null
+++ b/rice_prompt_handler.py
@@ -0,0 +1,40 @@
+_A='inputs'
+import json,os,random,tempfile,time
+from.utils import get_local_app_path
+from.rice_prompt_info import RicePromptInfo
+class RiceRoundPromptHandler:
+ _instance=None;_initialized=False
+ def __new__(A,*B,**C):
+ if A._instance is None:A._instance=super(RiceRoundPromptHandler,A).__new__(A)
+ return A._instance
+ def __init__(A):
+ if not A._initialized:A.client_id='';A.task_uuid='';A._initialized=True
+ def onprompt_handler(B,json_data):
+ '\n 处理传入的 JSON 数据\n :param json_data: 输入的 JSON 数据,包含各种任务信息\n ';H='input';G='template';F='task_uuid';E='client_id';D='prompt';A=json_data;RicePromptInfo().clear()
+ if E not in A:return A
+ B.client_id=A[E]
+ if F not in A:return A
+ B.task_uuid=A[F]
+ if D not in A:raise Exception("Warning: 'prompt' is missing.")
+ if G not in A:raise Exception("Warning: 'template' is missing.")
+ print(f"RiceRoundPromptHandler self.client_id={B.client_id!r} self.task_uuid={B.task_uuid!r}");I=A[H]if H in A else{};C=A[D];C=B.replace_output_prompt(C);J,K=B.parse_template(A[G]);C=B.replace_input_prompt(C,I,J,K);print(f"RiceRoundPromptHandler prompt_data={C!r}");A[D]=C;return A
+ def parse_template(E,template_data):
+ B={};C={};D=template_data['elements']
+ for A in D:id=A['id'];C[id]=A['node_id'];B[id]=A['type']
+ return B,C
+ def replace_output_prompt(D,prompt_data):
+ '\n 替换输出节点中的 task_id\n :param prompt_data: 任务的 prompt 数据\n :return: 替换后的 prompt 数据\n ';C='class_type';B=prompt_data
+ for(E,A)in B.items():
+ if A.get(C)=='RiceRoundOutputImageNode':A[_A]['task_id']=D.task_uuid
+ elif A.get(C)=='RiceRoundRandomSeedNode':A[_A]['seed']=random.randint(0,999999)
+ return B
+ def replace_input_prompt(F,prompt_data,input_data,id_type_map,node_id_map):
+ B='str';A=prompt_data;G={'text':'text_info','image_upload':'image_url','mask_upload':'mask_url','number_int':B,'number_float':B,'choice':'default','switch':B}
+ for(C,H)in input_data.items():
+ D=id_type_map.get(C,'');E=G.get(D)
+ if not E:print(f"RiceRoundPromptHandler replace_input_prompt unknown input_type {D}");continue
+ I=A[node_id_map[C]];I[_A][E]=str(H)
+ if os.environ.get('DEBUG')=='true':
+ J=tempfile.gettempdir()
+ with open(f"{J}//{F.task_uuid}_prompt_data.json",'w')as K:json.dump(A,K,indent=4)
+ return A
\ No newline at end of file
diff --git a/rice_prompt_info.py b/rice_prompt_info.py
new file mode 100644
index 0000000..ea840eb
--- /dev/null
+++ b/rice_prompt_info.py
@@ -0,0 +1,54 @@
+_E='file_name'
+_D='options_value'
+_C='choice_node'
+_B='python_class_name'
+_A=True
+import hashlib,json
+from pathlib import Path
+from server import PromptServer
+import re
+from.utils import get_local_app_path
+class RicePromptInfo:
+ _instance=None;_initialized=False
+ def __new__(A):
+ if A._instance is None:A._instance=super(RicePromptInfo,A).__new__(A)
+ return A._instance
+ def __init__(A):
+ if RicePromptInfo._initialized:return
+ A.task_info_folder=get_local_app_path()/'tasks'
+ if not A.task_info_folder.exists():A.task_info_folder.mkdir(parents=_A)
+ A.choice_node_map={};A.new_choice_file_map={};A.choice_classname_map={};A.load_choice_node_map();RicePromptInfo._initialized=_A
+ def clear(A):A.choice_node_map.clear();A.new_choice_file_map.clear()
+ def get_choice_server_folder(B):
+ A=B.task_info_folder/_C
+ if not A.exists():A.mkdir(parents=_A)
+ return A
+ def load_choice_node_map(A):
+ E=A.get_choice_server_folder()
+ for B in E.glob('*.json'):
+ try:
+ with open(B,'r',encoding='utf-8')as F:
+ C=json.load(F);D=C.get(_B)
+ if D:A.choice_classname_map[D]=C
+ except json.JSONDecodeError:print(f"Error parsing JSON from server file: {B}");continue
+ def get_choice_node_options(A,node_class_name):return A.choice_classname_map.get(node_class_name,{}).get(_D,[])
+ def get_choice_classname(A,node_id):return A.choice_node_map.get(node_id,{}).get(_B,'')
+ def get_choice_value(A,node_id):return A.choice_node_map.get(node_id,{}).get(_D,[])
+ def save_choice_classname(D,save_folder):
+ A=Path(save_folder)/_C
+ if not A.exists():A.mkdir(parents=_A)
+ for(F,B)in D.new_choice_file_map.items():
+ C=B.get(_E,'')
+ if C:
+ with open(A/f"{C}.json",'w',encoding='utf-8')as E:json.dump(B,E,ensure_ascii=False,indent=4)
+ def set_node_additional_info(B,node_additional_info):
+ H='template_id';C=node_additional_info
+ if C and isinstance(C,dict):
+ B.template_id=C.get(H,'');I=C.get('choice_node_map',{})
+ for(D,A)in I.items():
+ D=int(D);G=A.get('class_name','');A[H]=B.template_id;A['display_name']=G;J=A.get('node_type','')
+ if J=='RiceRoundAdvancedChoiceNode':
+ E=re.sub('[<>:"/\\\\|?*]','',G);F='RiceRoundAdvancedChoiceNode_'+hashlib.md5(E.lower().encode()).hexdigest()
+ if F in B.choice_classname_map:print(f"Error RicePromptInfo set_node_additional_info python_class_name {F} already exists");PromptServer.instance.send_sync('rice_round_toast',{'content':'选择名称重复,如果本身是同一类型的节点,可以忽略提示','type':'info','duration':5000})
+ A[_B]=F;A[_E]=E;B.new_choice_file_map[E]=A
+ B.choice_node_map[D]=A
\ No newline at end of file
diff --git a/rice_url_config.py b/rice_url_config.py
new file mode 100644
index 0000000..af32d7f
--- /dev/null
+++ b/rice_url_config.py
@@ -0,0 +1,107 @@
+_I='Content-Type'
+_H='download_url'
+_G='upload_sign_url'
+_F='file_type'
+_E='upload_type'
+_D='server_info.json'
+_C='comfyui_local_web_server_info.json'
+_B='image/png'
+_A='data'
+from enum import IntEnum
+import json,os
+from PIL import Image
+from io import BytesIO
+import numpy as np,requests
+from urllib.parse import urljoin
+from.utils import get_local_app_path
+DEFAULT_SUBDOMAIN='api'if os.getenv('DEBUG')!='true'else'test'
+DEFAULT_URL_PREFIX=f"https://{DEFAULT_SUBDOMAIN}.riceround.online"
+DEFAULT_WS_PREFIX=f"wss://{DEFAULT_SUBDOMAIN}.riceround.online"
+DEFAULT_WEB_URL_PREFIX=f"https://www.riceround.online"if os.getenv('DEBUG')!='true'else'http://localhost:5173/'
+DEFAULT_LOCAL_BASE_URL='http://127.0.0.1:8188'
+class UploadType(IntEnum):TEMPLATE_PUBLISH_IMAGE=1;USER_UPLOAD_TASK_IMAGE=2;MACHINE_TASK_RESULT=1000
+class RiceUrlConfig:
+ _instance=None;_initialized=False
+ def __new__(A,*B,**C):
+ if A._instance is None:A._instance=super(RiceUrlConfig,A).__new__(A)
+ return A._instance
+ def __init__(A):
+ if not A._initialized:A.comfyui_local_web_server_info={};A.load_comfyui_local_web_server_info();A.server_info={};A.load_server_info();A._initialized=True
+ def set_comfyui_local_web_server_info(B,local_server_info):
+ A=local_server_info;B.comfyui_local_web_server_info=A;C=get_local_app_path();D=C/_C
+ with open(D,'w')as E:json.dump(A,E)
+ def load_comfyui_local_web_server_info(A):
+ C=get_local_app_path();B=C/_C
+ if B.exists():
+ with open(B,'r')as D:A.comfyui_local_web_server_info=json.load(D)
+ else:A.comfyui_local_web_server_info={}
+ @property
+ def comfyui_local_base_url(self):
+ if self.comfyui_local_web_server_info:return self.comfyui_local_web_server_info.get('local_base_url','')
+ return DEFAULT_LOCAL_BASE_URL
+ def set_server_info(C,server_info):
+ B=server_info
+ if B and isinstance(B,dict):
+ A=B.get(_A)
+ if A and isinstance(A,dict):
+ C.server_info=A;D=get_local_app_path()
+ with open(D/_D,'w')as E:json.dump(A,E)
+ def load_server_info(B):
+ C=get_local_app_path();A=C/_D
+ if A.exists():
+ with open(A,'r')as D:B.server_info=json.load(D)
+ def get_server_url(A,url_path):return urljoin(A.url_prefix,url_path)
+ def get_ws_url(A,url_path):return urljoin(A.ws_prefix,url_path)
+ def get_web_url(A,url_path):return urljoin(A.web_url_prefix,url_path)
+ @property
+ def ws_prefix(self):
+ A=DEFAULT_WS_PREFIX
+ if self.server_info:A=self.server_info.get('ws_prefix')
+ if A:return A
+ return DEFAULT_WS_PREFIX
+ @property
+ def url_prefix(self):
+ A=DEFAULT_URL_PREFIX
+ if self.server_info:A=self.server_info.get('url_prefix')
+ if A:return A
+ return DEFAULT_URL_PREFIX
+ @property
+ def web_url_prefix(self):
+ A=DEFAULT_WEB_URL_PREFIX
+ if self.server_info:A=self.server_info.get('web_url_prefix')
+ if A:return A
+ return DEFAULT_WEB_URL_PREFIX
+ @property
+ def machine_upload_sign_url(self):return self.get_server_url('/api/machine_client/upload_image_sign_url')
+ @property
+ def user_upload_sign_url(self):return self.get_server_url('/api/user/upload_sign_url')
+ @property
+ def prompt_task_url(self):return self.get_server_url('/api/workflow/add_task')
+ @property
+ def preview_refresh_url(self):return self.get_server_url('/api/workflow/refresh_preview')
+ @property
+ def task_ws_url(self):return self.get_ws_url('/api/workflow/task_websocket')
+ @property
+ def login_api_url(self):return self.get_server_url('/api/workflow/get_info')
+ @property
+ def auth_web_url(self):return self.get_web_url('/auth/login')
+def user_upload_image(image,user_token):
+ D=RiceUrlConfig().user_upload_sign_url;G={'Authorization':f"Bearer {user_token}"};H={_E:UploadType.USER_UPLOAD_TASK_IMAGE.value,_F:_B};A=requests.get(D,headers=G,params=H);E='';B=''
+ if A.status_code==200:
+ C=A.json()
+ if C.get('code')==0:E=C.get(_A,{}).get(_G,'');B=C.get(_A,{}).get(_H,'')
+ else:raise ValueError(f"failed to upload image. Status code: {A.status_code}")
+ if not E or not B:raise ValueError(f"failed to upload image. upload_sign_url is empty")
+ I=255.*image.cpu().numpy();J=Image.fromarray(np.clip(I,0,255).astype(np.uint8));F=BytesIO();J.save(F,format='PNG',quality=95,compress_level=1);K=F.getvalue();A=requests.put(D,data=K,headers={_I:_B})
+ if A.status_code==200:return B
+ else:print(f"failed to upload image. Status code: {A.status_code}");raise ValueError(f"failed to upload image. Status code: {A.status_code}")
+def machine_upload_image(image,task_id):
+ F=RiceUrlConfig().machine_upload_sign_url;G=255.*image.cpu().numpy();H=Image.fromarray(np.clip(G,0,255).astype(np.uint8));E=BytesIO();H.save(E,format='PNG',quality=95,compress_level=1);I=E.getvalue();B='';C='';J={_E:UploadType.MACHINE_TASK_RESULT.value,_F:_B,'task_id':task_id};A=requests.get(F,params=J)
+ if A.status_code==200:
+ D=A.json()
+ if D.get('code')==0:B=D.get(_A,{}).get(_G,'');C=D.get(_A,{}).get(_H,'')
+ else:print(f"failed to upload image. Status code: {A.status_code}, Response: {A.text}");raise ValueError(f"failed to upload image. Status code: {A.status_code}")
+ if not B or not C:raise ValueError(f"failed to upload image. upload_sign_url is empty")
+ A=requests.put(B,data=I,headers={_I:_B})
+ if A.status_code==200:return C
+ else:print(f"failed to upload image. Status code: {A.status_code}, Response: {A.text}");raise ValueError(f"failed to upload image. Status code: {A.status_code}")
\ No newline at end of file
diff --git a/rice_websocket.py b/rice_websocket.py
new file mode 100644
index 0000000..95a189e
--- /dev/null
+++ b/rice_websocket.py
@@ -0,0 +1,146 @@
+_Q='WebSocket connection opened'
+_P='Message'
+_O='CommandType'
+_N='resultData'
+_M='templateName'
+_L='templateDescription'
+_K='templateId'
+_J='updateTime'
+_I='createTime'
+_H='prompt'
+_G='thumbnail'
+_F='progressText'
+_E='progress'
+_D='state'
+_C='taskUuid'
+_B=False
+_A=None
+from enum import Enum
+import json,time,threading
+from typing import Any,Callable,Optional
+import asyncio,websockets
+from websockets.exceptions import ConnectionClosedError
+import comfy.model_management as model_management
+COMMAND_TYPE_USER_SERVER_TASK_PROGRESS=5004
+COMMAND_TYPE_USER_CLIENT_WEB_COMMAND_CANCEL_TASK=4002
+class TaskStatus(Enum):
+ CREATED=0;PENDING=1;IN_PROGRESS=2;FINISHED=3;FAILED=4;CANCELLED=5
+ def __lt__(A,other):
+ B=other
+ if A.__class__ is B.__class__:return A.valueB.value
+ return NotImplemented
+ def __ge__(A,other):
+ B=other
+ if A.__class__ is B.__class__:return A.value>=B.value
+ return NotImplemented
+class TaskInfo:
+ def __init__(A,json_data):'\n Initialize TaskInfo from JSON data\n \n Args:\n json_data (dict): JSON data containing task information\n ';B=json_data;A.task_uuid=B.get(_C,'');A.state=TaskStatus(B.get(_D,0));A.progress=B.get(_E,0);A.progress_text=B.get(_F,'');A.thumbnail=B.get(_G,'');A.prompt=B.get(_H,'');A.create_time=B.get(_I,'');A.update_time=B.get(_J,'');A.template_id=B.get(_K,'');A.template_description=B.get(_L,'');A.template_name=B.get(_M,'');A.result_data=B.get(_N,_A);A.lock=threading.Lock();A.last_update_time=time.time();A.preview_refreshed=_B
+ def to_dict(A):'\n Convert TaskInfo to a dictionary\n \n Returns:\n dict: The dictionary representation of the TaskInfo object\n ';return{_C:A.task_uuid,_D:A.state.value,_E:A.progress,_F:A.progress_text,_G:A.thumbnail,_H:A.prompt,_I:A.create_time,_J:A.update_time,_K:A.template_id,_L:A.template_description,_M:A.template_name}
+ def update_progress(A,json_data):
+ C=json_data
+ with A.lock:
+ if C.get(_C,'')!=A.task_uuid:return _B
+ D=TaskStatus(C.get(_D,0))
+ if DA.progress:A.progress=E;A.progress_text=B
+ elif D==TaskStatus.FAILED:A.progress_text=B if B else'task failed'
+ elif D==TaskStatus.CANCELLED:A.progress_text=B if B else'task cancelled'
+ else:return _B
+ F=C.get(_N,_A)
+ if F:A.result_data=F
+ elif A.state==TaskStatus.IN_PROGRESS:A.result_data=_A
+ return True
+ def is_task_done(A):
+ with A.lock:return A.state>=TaskStatus.FINISHED
+ def is_task_timeout(A,timeout):
+ with A.lock:return time.time()-A.last_update_time>timeout
+ def __str__(A):'\n Get a string representation of the task\n \n Returns:\n str: A human-readable string describing the task status\n ';return f"Task {A.task_uuid}: {A.state.name} ({A.progress}%) - {A.progress_text}"
+class PackageMessage:
+ def __init__(A,CommandType,Message):A.CommandType=CommandType;A.Message=Message
+ def to_json(A):return json.dumps({_O:A.CommandType,_P:A.Message})
+ @classmethod
+ def from_json(B,data):A=json.loads(data);return B(A[_O],A[_P])
+class TaskWebSocket:
+ def __init__(A,url,token,machine_id,task_info,progress_callback):A.url=f"{url}?machine_id={machine_id}";A.token=token;A.task_info=task_info;A.stop_event=asyncio.Event();A.timeout=3600;A.progress_callback=progress_callback;A.last_progress_time=_A;A.message_timeout=600;A.websocket=_A;A.task=_A
+ async def connect(A):
+ try:
+ D=f"{A.url}&token={A.token}"
+ async with websockets.connect(D)as E:
+ A.websocket=E;await A.on_connection_open();print(_Q);F=asyncio.create_task(A.monitor_progress_timeout());G=asyncio.create_task(A.on_receive());H=asyncio.create_task(A.stop_event.wait());J,I=await asyncio.wait([F,G,H],return_when=asyncio.FIRST_COMPLETED)
+ for C in I:
+ C.cancel()
+ try:await C
+ except asyncio.CancelledError:pass
+ except ConnectionClosedError as B:print(f"Connection closed: {B}")
+ except asyncio.CancelledError:print('WebSocket connection cancelled')
+ except Exception as B:print(f"Connection error: {B}")
+ finally:
+ if A.websocket:await A.websocket.close();A.websocket=_A
+ async def on_receive(A):
+ try:
+ if not A.websocket or A.stop_event.is_set():return
+ async for B in A.websocket:await A.on_message(B)
+ except Exception as C:print(f"Error while listening to messages: {C}")
+ async def monitor_progress_timeout(A):
+ while not A.stop_event.is_set():
+ await asyncio.sleep(5)
+ try:model_management.throw_exception_if_processing_interrupted()
+ except Exception as B:
+ print(f"Processing interrupted during progress monitoring: {B}");A.stop_event.set();C=PackageMessage(CommandType=COMMAND_TYPE_USER_CLIENT_WEB_COMMAND_CANCEL_TASK,Message={'task_uuid':A.task_info.task_uuid})
+ try:await A.send_message(C)
+ except Exception as D:print(f"Failed to send cancel notification: {D}")
+ break
+ E=asyncio.get_event_loop().time()
+ if A.last_progress_time and E-A.last_progress_time>A.message_timeout:print('No task progress received within 10 minute, disconnecting...');A.stop_event.set();break
+ async def on_message(B,message):
+ try:
+ A=PackageMessage.from_json(message)
+ if A.CommandType==COMMAND_TYPE_USER_SERVER_TASK_PROGRESS:await B.handle_task_progress(A)
+ else:print(f"Unknown message type: {A.CommandType}")
+ except Exception as C:print(f"Message unpacking error: {C}")
+ async def on_connection_open(A):print(_Q)
+ async def send_message(A,message):
+ try:
+ if A.websocket:await A.websocket.send(message.to_json())
+ except Exception as B:print(f"Error sending message: {B}")
+ async def handle_task_progress(A,package):
+ if not A.task_info or A.stop_event.is_set():return
+ A.last_progress_time=time.time()
+ if A.task_info.update_progress(package.Message):
+ print(f"Task progress updated: {A.task_info}")
+ if A.progress_callback:A.progress_callback(A.task_info.task_uuid,A.task_info.progress_text,A.task_info.progress,A.task_info.preview_refreshed)
+ if A.task_info.is_task_done():A.stop_event.set()
+ async def start(A):
+ try:A.task=asyncio.create_task(A.connect());await asyncio.wait_for(A.task,timeout=A.timeout)
+ except asyncio.TimeoutError:print('WebSocket timeout 1 hour');A.stop_event.set()
+ except Exception as B:print(f"WebSocket connection error: {B}");raise
+ async def shutdown(A):
+ A.stop_event.set();A.progress_callback=_A
+ if A.websocket:
+ try:await A.websocket.close()
+ except websockets.ConnectionClosed:pass
+ finally:A.websocket=_A
+ if A.task:
+ if not A.task.done():
+ A.task.cancel()
+ try:await A.task
+ except asyncio.CancelledError:pass
+def start_and_wait_task_done(task_ws_url,user_token,machine_id,task_info,progress_callback,timeout=7200):
+ async def A():
+ A=TaskWebSocket(task_ws_url,user_token,machine_id,task_info,progress_callback)
+ try:await A.start()
+ except asyncio.CancelledError:print('Task cancelled')
+ finally:await A.shutdown()
+ asyncio.run(A())
\ No newline at end of file
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..221c22c
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,68 @@
+_B='RiceRound'
+_A='utf-8'
+import configparser,hashlib
+from io import BytesIO
+import os
+from pathlib import Path
+import uuid,torch,numpy as np
+from PIL import Image
+import platform,subprocess
+def pil2tensor(images):
+ 'Converts a PIL Image or a list of PIL Images to a tensor.';A=images
+ def B(image):
+ A=np.array(image).astype(np.float32)/255.
+ if A.ndim==2:return torch.from_numpy(A).unsqueeze(0)
+ else:return torch.from_numpy(A).unsqueeze(0)
+ if isinstance(A,Image.Image):return B(A)
+ else:return torch.cat([B(A)for A in A],dim=0)
+def calculate_machine_id():
+ '\n 获取跨平台的机器唯一标识符,类似于 gopsutil 的 HostID\n ';A=platform.system()
+ if A=='Linux':
+ try:
+ with open('/etc/machine-id','r')as B:return B.read().strip()
+ except FileNotFoundError:
+ try:
+ with open('/var/lib/dbus/machine-id','r')as B:return B.read().strip()
+ except FileNotFoundError:pass
+ try:C=subprocess.check_output(['cat','/sys/class/dmi/id/product_uuid']);return C.decode().strip()
+ except Exception:pass
+ elif A=='Windows':
+ try:
+ import winreg as D;F='SOFTWARE\\Microsoft\\Cryptography'
+ with D.OpenKey(D.HKEY_LOCAL_MACHINE,F)as G:H,I=D.QueryValueEx(G,'MachineGuid');return H
+ except Exception:pass
+ elif A=='Darwin':
+ try:
+ C=subprocess.check_output(['ioreg','-rd1','-c','IOPlatformExpertDevice'])
+ for E in C.decode().splitlines():
+ if'IOPlatformUUID'in E:return E.split('=')[-1].strip().strip('"')
+ except Exception:pass
+ return str(uuid.getnode())
+def normalize_machine_id(machine_id):'\n 接受一个机器标识符,并返回经过 MD5 哈希处理的规范化标识符\n ';A=_B;B=machine_id.strip();C=B.lower();D=C+A;E=hashlib.md5(D.encode(_A));return E.hexdigest()
+def get_local_app_path():A=Path.home();B=A/_B;return B
+def get_machine_id():
+ '\n 返回机器ID,为了兼容各个平台,各个语言,需要统一读写这个值\n ';F='machine_id';B='Machine';D=get_local_app_path();C=D/'machine.ini'
+ try:D.mkdir(parents=True,exist_ok=True)
+ except Exception as E:print(f"Error creating directory '{D}': {E}");return''
+ A=configparser.ConfigParser()
+ try:
+ if C.exists():
+ A.read(C,encoding=_A)
+ if B in A and F in A[B]:return A[B][F]
+ H=calculate_machine_id();G=normalize_machine_id(H)
+ if B not in A:A.add_section(B)
+ A.set(B,F,G)
+ with open(C,'w',encoding=_A)as I:A.write(I)
+ return G
+ except Exception as E:print(f"Error handling machine ID in '{C}': {E}");return''
+def combine_files(files,password,zip_file_path):
+ D=files;A=password;import pyzipper as C
+ for B in D:
+ if not os.path.exists(B):raise FileNotFoundError(f"file not found: {B}")
+ try:
+ with C.AESZipFile(zip_file_path,'w',compression=C.ZIP_DEFLATED,encryption=C.WZ_AES)as E:
+ if isinstance(A,str):A=A.encode(_A)
+ E.setpassword(A)
+ for(F,B)in enumerate(D,start=1):G=f"{F}.bin";E.write(B,G)
+ return True
+ except Exception as H:print(f"Error creating zip: {str(H)}");return False
\ No newline at end of file