tmp save poc

Author: ccssu

bizyair_extras/nodes_flux_trainer.py

@@ -1,8 +1,4 @@
 import io
-import json
-import os
-import re
-import time
 
 import matplotlib.pyplot as plt
 

nodes.py

@@ -125,16 +125,16 @@ def INPUT_TYPES(s):
         }
 
     RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
+    # FUNCTION = "sample"
 
     CATEGORY = "sampling"
 
-    def sample(self, model, **kwargs):
-        new_model: BizyAirNodeIO = model.copy(self.assigned_id)
-        kwargs["model"] = model
-        new_model.add_node_data(class_type="KSamplerAdvanced", inputs=kwargs)
-        progress_callback = ProgressCallback()
-        return new_model.send_request(progress_callback=progress_callback)
+    # def sample(self, model, **kwargs):
+    #     new_model: BizyAirNodeIO = model.copy(self.assigned_id)
+    #     kwargs["model"] = model
+    #     new_model.add_node_data(class_type="KSamplerAdvanced", inputs=kwargs)
+    #     progress_callback = ProgressCallback()
+    #     return new_model.send_request(progress_callback=progress_callback)
 
 
 class BizyAir_CheckpointLoaderSimple(BizyAirBaseNode):
@@ -222,21 +222,21 @@ def INPUT_TYPES(s):
 
     RETURN_TYPES = ("IMAGE",)
     RETURN_NAMES = (f"IMAGE",)
-    FUNCTION = "decode"
+    # FUNCTION = "decode"
 
     CATEGORY = f"{PREFIX}/latent"
 
-    def decode(self, vae, samples):
-        new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
-        new_vae.add_node_data(
-            class_type="VAEDecode",
-            inputs={
-                "samples": samples,
-                "vae": vae,
-            },
-            outputs={"slot_index": 0},
-        )
-        return new_vae.send_request()
+    # def decode(self, vae, samples):
+    #     new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
+    #     new_vae.add_node_data(
+    #         class_type="VAEDecode",
+    #         inputs={
+    #             "samples": samples,
+    #             "vae": vae,
+    #         },
+    #         outputs={"slot_index": 0},
+    #     )
+    #     return new_vae.send_request()
 
 
 class BizyAir_LoraLoader(BizyAirBaseNode):
@@ -295,20 +295,20 @@ def INPUT_TYPES(s):
 
     RETURN_TYPES = ("LATENT",)
     RETURN_NAMES = (f"LATENT",)
-    FUNCTION = "encode"
+    # FUNCTION = "encode"
     CATEGORY = f"{PREFIX}/latent"
 
-    def encode(self, vae, pixels):
-        new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
-        new_vae.add_node_data(
-            class_type="VAEEncode",
-            inputs={
-                "vae": vae,
-                "pixels": pixels,
-            },
-            outputs={"slot_index": 0},
-        )
-        return new_vae.send_request()
+    # def encode(self, vae, pixels):
+    #     new_vae: BizyAirNodeIO = vae.copy(self.assigned_id)
+    #     new_vae.add_node_data(
+    #         class_type="VAEEncode",
+    #         inputs={
+    #             "vae": vae,
+    #             "pixels": pixels,
+    #         },
+    #         outputs={"slot_index": 0},
+    #     )
+    #     return new_vae.send_request()
 
 
 class BizyAir_VAEEncodeForInpaint(BizyAirBaseNode):

src/bizyair/commands/servers/pub_sub_sse.py

@@ -118,6 +118,7 @@ def __init__(self, name, mediator):
         self.mediator = mediator
         self.messages = queue.Queue()
         self.is_subscribed = True  # 订阅状态监测
+        self.tmp_result = {}
 
     def receive(self, message):
         self.messages.put(message)
@@ -141,6 +142,11 @@ def unsubscribe(self):
 
     def get_result(self, node_id, timeout=12):
         while True:
+            if node_id in self.tmp_result:
+                out = self.tmp_result[node_id]
+                del self.tmp_result[node_id]
+                return out
+
             result = self.pop(timeout=timeout)
             if result is None:
                 return None
@@ -153,6 +159,8 @@ def get_result(self, node_id, timeout=12):
                     event_node_id = result["message"]["data"]["node"]
                     if event_node_id == node_id:
                         return result["data"]["payload"]
+                    else:
+                        self.tmp_result[event_node_id] = result["data"]["payload"]
             except Exception as e:
                 print(f"Error processing message for {self.name}: {e}")
                 return None

src/bizyair/image_utils.py

@@ -176,14 +176,18 @@ def tensor_to_base64(tensor: torch.Tensor, compress=True) -> str:
     return tensor_b64
 
 
-def base64_to_tensor(tensor_b64: str, compress=True) -> torch.Tensor:
+def base64_to_numpy(tensor_b64: str, compress=True) -> np.ndarray:
     tensor_bytes = base64.b64decode(tensor_b64)
 
     if compress:
         tensor_bytes = zlib.decompress(tensor_bytes)
 
     tensor_np = pickle.loads(tensor_bytes)
+    return tensor_np
+
 
+def base64_to_tensor(tensor_b64: str, compress=True) -> torch.Tensor:
+    tensor_np = base64_to_numpy(tensor_b64, compress=compress)
     tensor = torch.from_numpy(tensor_np)
     return tensor
 
@@ -222,7 +226,7 @@ def _(input: str, **kwargs):
         return decode_comfy_image(tensor_b64, old_version=old_version)
     elif input.startswith(NUMPY_MARKER):
         tensor_b64 = input[len(NUMPY_MARKER) :]
-        return base64_to_tensor(tensor_b64)
+        return base64_to_numpy(tensor_b64)
     return input
 
 

src/bizyair/nodes_base.py

@@ -166,6 +166,8 @@ def default_function(self, **kwargs):
             result = self._pre_run()
             if result:
                 return self._merge_results(result, node_ios)
+            else:
+                warnings.warn("Pre-run result is None")
 
         if len(send_request_datatype_list) == len(self.RETURN_TYPES):
             return self._process_all_send_request_types(node_ios)