diff --git a/src/bizyair/image_utils.py b/src/bizyair/image_utils.py index 1e6ef54a..d8e34c18 100644 --- a/src/bizyair/image_utils.py +++ b/src/bizyair/image_utils.py @@ -398,18 +398,18 @@ def _(input: str): @singledispatch -def encode_data(output): +def encode_data(output, disable_image_marker=False): raise NotImplementedError(f"Unsupported type: {type(output)}") @encode_data.register(dict) -def _(output): - return {k: encode_data(v) for k, v in output.items()} +def _(output, **kwargs): + return {k: encode_data(v, **kwargs) for k, v in output.items()} @encode_data.register(list) -def _(output): - return [encode_data(x) for x in output] +def _(output, **kwargs): + return [encode_data(x, **kwargs) for x in output] def is_image_tensor(tensor) -> bool: @@ -440,14 +440,14 @@ def is_image_tensor(tensor) -> bool: @encode_data.register(torch.Tensor) -def _(output): - if is_image_tensor(output): +def _(output, **kwargs): + if is_image_tensor(output) and not kwargs.get("disable_image_marker", False): return IMAGE_MARKER + encode_comfy_image(output, image_format="WEBP") return TENSOR_MARKER + tensor_to_base64(output) @encode_data.register(BizyAirNodeIO) -def _(output: BizyAirNodeIO): +def _(output: BizyAirNodeIO, **kwargs): origin_id = output.node_id origin_slot = output.nodes[origin_id]["outputs"]["slot_index"] return [origin_id, origin_slot] @@ -458,5 +458,5 @@ def _(output: BizyAirNodeIO): @encode_data.register(str) @encode_data.register(bool) @encode_data.register(type(None)) -def _(output): +def _(output, **kwargs): return output diff --git a/supernode.py b/supernode.py index cff82b24..8859cc39 100644 --- a/supernode.py +++ b/supernode.py @@ -217,7 +217,7 @@ def super_resolution(self, image): "content-type": "application/json", "authorization": auth, } - input_image = encode_data(image) + input_image = encode_data(image, disable_image_marker=True) payload["image"] = input_image payload["is_compress"] = True