Merge pull request #94 from AIGODLIKE/develop

Develop

MultiLineText/trie.py

@@ -240,6 +240,7 @@ def insert_internal(trie: Trie):
                         if c == "," and left_bracket == 0:
                             split_words.append(split_word.strip())
                             split_word = ""
+                            continue
                         split_word += c
                         if c == "(":
                             left_bracket += 1
@@ -249,6 +250,8 @@ def insert_internal(trie: Trie):
                     split_words = content.split(",")
                 # 处理 to -> replace
                 for split_word in split_words:
+                    if not split_word:
+                        continue
                     rrow = (row[0], split_word, row[2], row[1], row[4])
                     words.append(rrow)
     for w in words:

MultiLineText/words_collection.py

@@ -27,6 +27,7 @@ class Words:
 
     def __init__(self):
         self.word_list: list[tuple] = []
+        self.word_map: dict[str, tuple] = {}
 
     def read_tags(self):
         if self.CACHE_PATH.exists():
@@ -89,10 +90,13 @@ def read_csv(self):
             # words.append((word, 5000, ""))
         return words
 
+    def load_map(self):
+        for word in self.word_list:
+            self.word_map[word[0]] = word
 
 words = Words()
 words.read_tags()
-
+words.load_map()
 
 @bpy.app.handlers.persistent
 def f(_):
@@ -109,7 +113,9 @@ def f(_):
             it.name = f"{word[0]} <== {word[2]}"
         it.freq = int(word[1])
         count += 1
-    logger.info(f"Load MLT Words: {time.time()-ts:.4f}s")
+        if count > 1000:
+            break
+    logger.info(f"Update Search Words: {time.time()-ts:.4f}s")
 
 
 if f not in bpy.app.handlers.load_post:

SDNode/blueprints.py

@@ -192,6 +192,7 @@ def show_model_preview(self: NodeBase, context: bpy.types.Context, layout: bpy.t
                 mgr = bpy.context.window_manager
                 col.prop(stat, "addtext", icon="ADD", text="")
                 col.template_list("MLTWords_UL_UIList", prop, mgr, "mlt_words", mgr, "mlt_words_index")
+                col.prop(bpy.context.scene.sdn, "search_tag", icon="VIEWZOOM", text="")
                 col.template_list("MLTText_UL_UIList", "", stat, "texts", stat, "tindex")
             return True
 

SDNode/manager.py

@@ -6,6 +6,7 @@
 import json
 import time
 import atexit
+import struct
 from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from shutil import rmtree
@@ -53,6 +54,7 @@ def __init__(self, task=None, pre=None, post=None, tree=None) -> None:
         self.executing_node: NodeBase = None
         self.is_finished = False
         self.process = {}
+        self.binary_message = b""
         # 记录node的类型 防止节点树变更
         self.node_ref_map = {}
         if not tree:
@@ -97,6 +99,7 @@ def f(self: Task):
 
     def set_executing_node_id(self, node_id):
         self.executing_node_id = node_id
+        self.binary_message = b""
 
         def f(self: Task):
             from .nodes import NodeBase
@@ -584,6 +587,7 @@ def wait_connect(self) -> bool:
             if requests.get(f"{self.get_url()}/object_info", proxies={"http": None, "https": None}, timeout=10).status_code == 200:
                 self.server_connected = True
                 update_screen()
+                get_pref().preview_method = get_pref().preview_method
                 return True
         except requests.exceptions.ConnectionError as e:
             TaskManager.put_error_msg(str(e))
@@ -704,6 +708,7 @@ def wait_connect(self) -> bool:
             import requests
             try:
                 if requests.get(f"{self.get_url()}/object_info", proxies={"http": None, "https": None}, timeout=1).status_code == 200:
+                    get_pref().preview_method = get_pref().preview_method
                     return True
             except requests.exceptions.ConnectionError:
                 ...
@@ -896,10 +901,10 @@ def refresh_node():
             t1 = time.time()
             rtnode_unreg()
             t2 = time.time()
-            logger.info(_T("UnregNode Time:") + f" {t2-t1:.2f}s")
+            logger.info(_T("UnregNode Time:") + f" {t2 - t1:.2f}s")
             rtnode_reg()
             t3 = time.time()
-            logger.info(_T("RegNode Time:") + f" {t3-t2:.2f}s")
+            logger.info(_T("RegNode Time:") + f" {t3 - t2:.2f}s")
         if TaskManager.is_launching():
             return
 
@@ -919,7 +924,7 @@ def update_screen_timer():
                 TaskManager.init_server(fake=True, callback=callback)
             TaskManager.is_server_launching = False
             t2 = time.time()
-            logger.info(_T("Launch Time:") + f" {t2-t1:.2f}s")
+            logger.info(_T("Launch Time:") + f" {t2 - t1:.2f}s")
             bpy.app.timers.unregister(update_screen_timer)
         if fake:
             job()
@@ -1165,6 +1170,11 @@ def poll_res():
         SessionId = TaskManager.SessionId
 
         def on_message(ws, message):
+            if isinstance(message, bytes):
+                if tm.cur_task:
+                    tm.cur_task.binary_message = message
+                TaskManager.handle_binary_message(message)
+                return
             msg = json.loads(message)
             try:
                 from .custom_support import crystools_monitor, cup_monitor
@@ -1217,7 +1227,7 @@ def on_message(ws, message):
                 TaskManager.progress_bar = v
                 cf = "\033[92m" + "█" * v + "\033[0m"
                 cp = "\033[32m" + "░" * (m - v) + "\033[0m"
-                content = f"{v*100/m:3.0f}% " + cf + cp + f" {v}/{m}"
+                content = f"{v * 100 / m:3.0f}% " + cf + cp + f" {v}/{m}"
                 logger.info(content + "\r", extra={"same_line": True})
                 # sys.stdout.write(content)
                 # sys.stdout.flush()
@@ -1270,15 +1280,45 @@ def on_message(ws, message):
                 ...  # pass
             else:
                 logger.error(message)
-
-        ws = WebSocketApp(f"ws://{get_ip()}:{get_port()}/ws?clientId={SessionId['SessionId']}", on_message=on_message)
+        listen_addr = f"ws://{get_ip()}:{get_port()}/ws?clientId={SessionId['SessionId']}"
+        ws = WebSocketApp(listen_addr, on_message=on_message)
         TaskManager.ws = ws
         ws.run_forever()
+        if True:
+            ...
+        else:
+            # 备选方案
+            from ..External.websockets.sync.client import connect
+            from ..External.websockets import ConnectionClosedError
+            ws = connect(listen_addr)
+            TaskManager.ws = ws
+            try:
+                for msg in ws:
+                    on_message(None, msg)
+            except ConnectionClosedError:
+                ...
         logger.debug(_T("Poll Result Thread Exit"))
         TaskManager.ws = None
         if TaskManager.server.is_launched():
             Timer.put((TaskManager.restart_server, True))
 
+    @staticmethod
+    def handle_binary_message(data):
+        # 解析二进制数据的前4个字节获取事件类型
+        event_type = struct.unpack(">I", data[:4])[0]
+        # 根据事件类型处理数据
+        if event_type != 1:
+            logger.debug("Unknown binary event type: %s", event_type)
+            return
+        # 处理图像类型
+        image_type = struct.unpack(">I", data[4:8])[0]
+        image_mime = "image/png" if image_type == 2 else "image/jpeg"
+        # 假设剩余的数据是图像数据，可以保存或进一步处理
+        image_data = data[8:]
+        return
+        with open(f"/Users/karrycharon/Desktop/000.{image_mime.split('/')[1]}", "wb") as f:
+            f.write(image_data)
+
 
 def removetemp():
     tempdir = Path(__file__).parent / "temp"

SDNode/node_process.py

@@ -1,5 +1,10 @@
 import bpy
 import blf
+import gpu
+import struct
+import tempfile
+from pathlib import Path
+from functools import lru_cache
 from mathutils import Vector
 from .manager import TaskManager
 from ..utils import _T
@@ -13,6 +18,41 @@
 #         if sp.type == "NODE_EDITOR":
 #             return sp
 #     return None
+@lru_cache
+def load_texture(data: bytes) -> gpu.types.GPUTexture:
+    if not data:
+        return
+    # 解析二进制数据的前4个字节获取事件类型
+    event_type = struct.unpack(">I", data[:4])[0]
+    # 根据事件类型处理数据
+    if event_type != 1:
+        return
+    # 处理图像类型
+    image_type = struct.unpack(">I", data[4:8])[0]
+    image_mime = ".png" if image_type == 2 else ".jpeg"
+    img_path = Path(tempfile.gettempdir()).joinpath("rviewport").with_suffix(image_mime)
+    img_path.write_bytes(data[8:])
+    bl_img = bpy.data.images.load(img_path.as_posix())
+    gpu_img = gpu.texture.from_image(bl_img)
+    bpy.data.images.remove(bl_img)
+    img_path.unlink()
+    return gpu_img
+
+
+def display_texture(texture: bytes, loc, size):
+    """
+    显示图片纹理
+    """
+    tex = load_texture(texture)
+    if not tex:
+        return
+    loc = loc.copy()
+    loc.y += 20
+    pos = VecWorldToRegScale(loc)
+    from gpu_extras.presets import draw_texture_2d
+    w = size
+    h = tex.height / tex.width * size
+    draw_texture_2d(tex, pos, w, h)
 
 
 def display_text(text, pos, size=50, color=(0, 0.7, 0.0, 1.0)):
@@ -60,6 +100,7 @@ def draw():
     loc = n.location.copy()
     loc.y += 10
     pos = VecWorldToRegScale(loc)
+    display_texture(task.binary_message, loc, calc_size(view2d, n.width))
     display_text(head, pos, size, (0, 1, 0.0, 1.0))
     if not task.process:
         return
@@ -68,7 +109,7 @@ def draw():
     blf.size(FONT_ID, size)
     loc.x += blf.dimensions(FONT_ID, head)[0] / size * vsize
     pos = VecWorldToRegScale(loc)
-    display_text(f" {v/m*100:3.0f}% ", pos, size * 1.5, (1, 1, 0.0, 1.0))
+    display_text(f" {v / m * 100:3.0f}% ", pos, size * 1.5, (1, 1, 0.0, 1.0))
     draw_node_process(n, v / m)
 
 

SDNode/nodes.py

@@ -1021,7 +1021,8 @@ def draw_color_simple(cls):
 
 class Ops_Switch_Socket_Disp(bpy.types.Operator):
     bl_idname = "sdn.switch_socket_disp"
-    bl_label = "切换Socket显示隐藏"
+    bl_label = "Toggle socket visibility"
+    bl_translation_context = ctxt
     socket_name: bpy.props.StringProperty()
     node_name: bpy.props.StringProperty()
     action: bpy.props.StringProperty(default="")
@@ -1074,6 +1075,7 @@ class Ops_Switch_Socket_Widget(bpy.types.Operator):
     socket_name: bpy.props.StringProperty()
     node_name: bpy.props.StringProperty()
     action: bpy.props.StringProperty(default="")
+    bl_translation_context = ctxt
 
     def set_active_node(self, tree):
         if not tree:
@@ -1142,6 +1144,7 @@ class Ops_Add_SaveImage(bpy.types.Operator):
     bl_label = "Add SaveImage node"
     bl_description = "Add a SaveImage node and connect it to the image"
     node_name: bpy.props.StringProperty()
+    bl_translation_context = ctxt
 
     def execute(self, context):
         tree = get_default_tree()
@@ -1163,6 +1166,7 @@ class Ops_Active_Tex(bpy.types.Operator):
     bl_label = "选择纹理"
     img_name: bpy.props.StringProperty()
     node_name: bpy.props.StringProperty()
+    bl_translation_context = ctxt
 
     def execute(self, context):
         if not (img := bpy.data.images.get(self.img_name)):
@@ -1178,6 +1182,7 @@ class Ops_Link_Mask(bpy.types.Operator):
     bl_idname = "sdn.link_mask"
     bl_label = "链接遮照"
     bl_options = {"REGISTER", "UNDO"}
+    bl_translation_context = ctxt
     action: bpy.props.StringProperty(default="")
     cam_name: bpy.props.StringProperty(default="")
     node_name: bpy.props.StringProperty(default="")
@@ -1733,6 +1738,9 @@ def _parse(name, desc, _desc):
                         if not isinstance(inp_desc[0], str):
                             logger.warning("socket type not str[IGNORE]: %s.%s -> %s", name, inp, inp_desc[0])
                             inp_desc[0] = str(inp_desc[0])
+                        # 如果到这里仍然是空 则使用默认字符串
+                        if inp_desc[0] == "":
+                            inp_desc[0] = f"{name}_{inp}"
                         _desc.add(inp_desc[0])
                         self.SOCKET_TYPE[name][inp] = inp_desc[0]
             for index, out_type in enumerate(desc["output"]):
@@ -1774,7 +1782,9 @@ def _parse_sockets_clss(self):
         for stype in sockets:
             if stype in {"ENUM", }:
                 continue
-
+            # 过滤不安全socket
+            if stype == "":
+                continue
             def draw(self, context, layout, node: NodeBase, text):
                 if not node.is_registered_node_type():
                     return

__init__.py

@@ -1,7 +1,7 @@
 bl_info = {
     'name': '无限圣杯-节点 (ComfyUI Node Editor)',
     'author': '幻之境开发小组-会飞的键盘侠、只剩一瓶辣椒酱',
-    'version': (1, 5, 0),
+    'version': (1, 5, 1),
     'blender': (3, 0, 0),
     'location': '3DView->Panel',
     'category': 'Node', # This is automatically translated by Blender
@@ -38,7 +38,6 @@ def clear_pyc(path=None, depth=2):
 from .SDNode import rtnode_unreg, TaskManager
 from .MultiLineText import EnableMLT
 
-from .translations import translations_dict
 from .utils import Icon, FSWatcher, ScopeTimer
 from .timer import timer_reg, timer_unreg
 from .preference import pref_register, pref_unregister
@@ -110,6 +109,8 @@ def register():
     if bpy.app.background:
         dump_info()
         return
+
+    from .translations import translations_dict
     bpy.app.translations.register(__name__, translations_dict)
     reg()
     Icon.set_hq_preview()

datas.py

@@ -3,6 +3,14 @@
 GROUPS_DIR = Path(__file__).parent / "groups"
 IMG_SUFFIX = {".png", ".jpg", ".jpeg"}
 
+VERSION = ""
+
+def get_bl_version():
+    if VERSION: return VERSION
+    from . import bl_info
+    return ".".join([str(i) for i in bl_info['version']])
+
+VERSION = get_bl_version()
 
 class MetaIn(type):
     CACHE: dict[str, dict] = {}