nodes.py

import torch
import requests
import os
import folder_paths
import io
import time
from comfy.utils import ProgressBar

TEMP_DIR = folder_paths.get_temp_directory()
BRIDGE_DIR = os.path.join(TEMP_DIR, "tensor_bridge")


def move_to_cpu_recursive(obj):
    if isinstance(obj, torch.Tensor):
        return obj.cpu()
    elif isinstance(obj, dict):
        return {k: move_to_cpu_recursive(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [move_to_cpu_recursive(v) for v in obj]
    elif isinstance(obj, tuple):
        return tuple(move_to_cpu_recursive(v) for v in obj)
    return obj


def send_tensor_http(url, channel, data):
    cpu_data = move_to_cpu_recursive(data)

    buffer = io.BytesIO()
    torch.save(cpu_data, buffer)
    buffer.seek(0)

    base_url = url.rstrip('/')
    target_url = f"{base_url}/tensor_bridge/upload"
    files = {'file': buffer}
    post_data = {'channel': channel}
    headers = {"ngrok-skip-browser-warning": "true", "User-Agent": "ComfyTensorBridge"}

    file_size = len(buffer.getbuffer()) / (1024 * 1024)
    print(f"[TensorBridge] Sending {file_size:.2f} MB to {target_url} ({channel})...")

    try:
        response = requests.post(target_url, files=files, data=post_data, headers=headers, timeout=120)
        response.raise_for_status()
    except Exception as e:
        print(f"[TensorBridge] Upload Failed: {e}")
        raise e


def load_tensor_disk(channel, timeout=300):
    filename = f"{channel}.pt"
    filepath = os.path.join(BRIDGE_DIR, filename)
    pbar = ProgressBar(timeout)
    start_time = time.time()

    print(f"[TensorBridge] Waiting for new data on channel: '{channel}'...")

    while time.time() - start_time < timeout:
        if os.path.exists(filepath):
            time.sleep(0.5)
            try:
                data = torch.load(filepath, map_location="cpu")
                os.remove(filepath)
                print(f"[TensorBridge] Data received and consumed.")
                return data
            except Exception as e:
                print(f"[TensorBridge] Read failed (retrying): {e}")

        time.sleep(0.5)
        pbar.update(1)

    raise TimeoutError(f"[TensorBridge] Timed out waiting for data on channel {channel}")


# --- NODE CLASSES ---

class BridgeSenderLatent:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "latent": ("LATENT",),
                "url": ("STRING", {"default": "http://127.0.0.1:8188"}),
                "channel": ("STRING", {"default": "latent_pipe"}),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "send"
    CATEGORY = "TensorBridge"
    OUTPUT_NODE = True

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def send(self, latent, url, channel):
        send_tensor_http(url, channel, {"payload": latent})
        return (latent,)


class BridgeReceiverLatent:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "channel": ("STRING", {"default": "latent_pipe"}),
                "timeout_sec": ("INT", {"default": 60, "min": 5, "max": 604800}),
            },
            "optional": {
                "trigger": ("*",),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "receive"
    CATEGORY = "TensorBridge"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def receive(self, channel, timeout_sec, trigger=None):
        data = load_tensor_disk(channel, timeout=timeout_sec)
        return (data["payload"],)


class BridgeSenderImage:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "url": ("STRING", {"default": "http://127.0.0.1:8188"}),
                "channel": ("STRING", {"default": "image_pipe"}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "send"
    CATEGORY = "TensorBridge"
    OUTPUT_NODE = True

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def send(self, image, url, channel):
        send_tensor_http(url, channel, {"payload": image})
        return (image,)


class BridgeReceiverImage:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "channel": ("STRING", {"default": "image_pipe"}),
                "timeout_sec": ("INT", {"default": 60, "min": 5, "max": 604800}),
            },
            "optional": {
                "trigger": ("*",),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "receive"
    CATEGORY = "TensorBridge"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def receive(self, channel, timeout_sec, trigger=None):
        data = load_tensor_disk(channel, timeout=timeout_sec)
        return (data["payload"],)


class BridgeSenderConditioning:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "conditioning": ("CONDITIONING",),
                "url": ("STRING", {"default": "http://127.0.0.1:8188"}),
                "channel": ("STRING", {"default": "cond_pipe"}),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "send"
    CATEGORY = "TensorBridge"
    OUTPUT_NODE = True

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def send(self, conditioning, url, channel):
        send_tensor_http(url, channel, {"payload": conditioning})
        return (conditioning,)


class BridgeReceiverConditioning:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "channel": ("STRING", {"default": "cond_pipe"}),
                "timeout_sec": ("INT", {"default": 60, "min": 5, "max": 604800}),
            },
            "optional": {"trigger": ("*",), }
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "receive"
    CATEGORY = "TensorBridge"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def receive(self, channel, timeout_sec, trigger=None):
        data = load_tensor_disk(channel, timeout=timeout_sec)
        payload = data["payload"]

        if not isinstance(payload, list):
            raise ValueError(f"[TensorBridge] Type Error: Expected CONDITIONING (List), got {type(payload)}.")

        return (payload,)


class BridgeExecutionSync:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "optional": {
                "wait_for_1": ("*",),
                "wait_for_2": ("*",),
                "wait_for_3": ("*",),
                "wait_for_4": ("*",),
            }
        }

    RETURN_TYPES = ("*",)
    FUNCTION = "sync"
    CATEGORY = "TensorBridge"

    @classmethod
    def IS_CHANGED(s, **kwargs):
        return float("NaN")

    def sync(self, wait_for_1=None, wait_for_2=None, wait_for_3=None, wait_for_4=None):
        return (True,)