Support for multiple ComfyUIs

[Yahweasel]

I have two instances of Comfy running at the same time, one for each GPU. I'd like to implement something that involves doing one step of generation on one GPU, and one step of generation on the other. As far as I can tell, the way ComfyScript works, there's a global variable for the API to use, and no way of interacting with multiple servers simultaneously in the same script. Is that a correct assessment? Is there any workaround?

[Yahweasel]

(Or, alternatively, in real mode, I would want a way of specifying where to load things, which I guess could theoretically be possible without loading two per se instances)

[Chaoses-Ib]

Supporting multi ComfyUIs for virtual mode should be easy to do. Mostly just replacing client.client with self.client, then you can specify the client like Workflow(client=Client('http://127.0.0.1:8188/')).

I'd like to implement something that involves doing one step of generation on one GPU, and one step of generation on the other.

However, I'm not sure how would you do this. ComfyUI doesn't have API to pause the workflow. You mean saving and loading the image each time?

in real mode, I would want a way of specifying where to load things, which I guess could theoretically be possible without loading two per se instances

At least PyTorch can do it, you can try torch.cuda.set_device, torch.set_default_device, os.environ['CUDA_VISIBLE_DEVICES']. I'm not sure how well ComfyUI can work with this though.

[Yahweasel]

Let me make it clearer with a specific use case I do.

Wan2.2 has two models: the "low noise" model and the "high noise" model. I can load either one, but I can't load both on the same GPU. So, I run the low noise model on one GPU, and the high noise model on the other GPU, then I have the low-noise workflow save it latents, and the high-noise workflow load those latents. Similarly, Hunyuan has a refiner model, but in that case, it actually has a different latent format, so it literally would be pulling the PIL image from one and putting it into the other.

All of this is doable, but a huge PITA and not at all automated. I was hoping to use ComfyScript to pipeline these sorts of multi-GPU workflows. E.g., one thread doing high noise, one thread doing low noise, and a queue feeding the one to the other.

[Yahweasel]

I made a very simple example to demonstrate:

#!/usr/bin/env python3
import multiprocessing

def worker1(in_q, out_q):
    import random

    import comfy_script.runtime as runtime
    runtime.load(args=runtime.ComfyUIArgs("--cuda-device", "0", "--port", "8188"))
    import comfy_script.runtime.nodes as nodes

    with runtime.Workflow():
        model, clip, _ = nodes.CheckpointLoaderSimple("sd_xl_base_1.0.safetensors")
        model, clip = nodes.LoraLoader(model, clip, "Hyper-SDXL-2steps-lora.safetensors", 1, 1)
        latent = nodes.EmptyLatentImage(1024, 1024, 1)
        prompt = in_q.get()
        positive = nodes.CLIPTextEncode(prompt, clip)
        negative = nodes.CLIPTextEncode("", clip)
        latent = nodes.KSamplerAdvanced(
            model, "enable", random.randint(1, 2**64-1), 2, 1, "euler",
            "simple", positive, negative, latent, 0, 10000, "disable"
        )
        out_q.put(latent)

def worker2(in_q, out_q):
    import comfy_script.runtime as runtime
    runtime.load(args=runtime.ComfyUIArgs("--cuda-device", "1", "--port", "8189"))
    import comfy_script.runtime.nodes as nodes

    with runtime.Workflow():
        latent = in_q.get()
        image = nodes.DecodeTAESD(latent, "taesdxl")
        images = runtime.util.get_images(image)
        out_q.put(images[0])

q1 = multiprocessing.Queue()
q2 = multiprocessing.Queue()
q3 = multiprocessing.Queue()

p1 = multiprocessing.Process(target=worker1, args=(q1, q2))
p2 = multiprocessing.Process(target=worker2, args=(q2, q3))
p1.start()
p2.start()

q1.put("women laughing and eating salad")
image = q3.get()
image.save("out.png")

I believe that this actually would work if I were passing through a PIL image retrieved with util.get_images, and so this might be a fine technique when higher-level data is being passed through. If there were a way to extract latents like the SaveLatents node (but without actually saving them to disk), I'd be more-or-less done. Probably would be prettier if I used shared memory to transfer over the latents, but other than that, it'd work.

[Chaoses-Ib]

Looks like using multiprocessing is simpler than I thought. This can even work with real mode, feels much powerful than the multiple Client solution I mentioned, and probably also faster due to no GIL contention...

If there were a way to extract latents like the SaveLatents node (but without actually saving them to disk), I'd be more-or-less done. Probably would be prettier if I used shared memory to transfer over the latents, but other than that, it'd work.

I've added some utilities for latent. You can update by:

python -m pip install "comfy-script[default] @ git+https://github.com/Chaoses-Ib/ComfyScript.git"

(or git pull)

For example:

# For print
queue.watch_display(False)

with Workflow():
    latent = EmptyLatentImage(batch_size=4)
    latent_path = util.save_latent_and_get_path(latent)
print(latent_path)

with Workflow():
    latent = util.load_latent_from_path(latent_path)
    # Do something with the latent, for example:
    SaveLatent(latent, 'latents/loaded')

Unfortunately, this needs disk IO. ComfyUI doesn't even have an API to extract images without saving them to disk, not to mention latents. get_images() uses PreviewImage internally, which would save them to ComfyUI/temp. But for latents, there is not even a way to save them to temp instead of output...

You have to use real mode or count on custom nodes if you really need to avoid disk IO. Or, if you are on Linux/macOS, an easy workaround is to link ComfyUI/output/latents to the in-memory tmpfs. This is the cheap way to do shared memory.

It is also possible to hook Python disk IO and turn them into RPC (via shared memory / mapped file) instead. But it's a bit tricky, probably not worth it unless disk IO is truly the bottleneck.

[Yahweasel]

Thanks for the quick response! I think this should at least get me to whatever the next catastrophe is ;)

I've written many Comfy custom nodes (I've bridged three image models to Comfy), so I may look into writing custom nodes to avoid disk I/O, but for right now I don't know if that's enough of a bottleneck that I care. Honestly, my concern with images and disk I/O would be more about the pointless PNG cycle than the disk per se, and that doesn't apply to latents.

[Yahweasel]

I guess between multiprocessing and util.latents my needs are met, so I'll close this.

I'll mention that the way util.latents is written currently, there's no way to delete a latent when you're done with it. I suspect that [in virtual mode] there's really no way of doing that anyway, so it's a "we'll need a custom node to do this right" situation, but it's worth mentioning. If you're generating stuff like this in a loop, you're going to end up littering a lot of latents :)

[Chaoses-Ib]

Yeah, either a node to save latent to temp or to delete output is needed. Though if you using it locally, you can just os.remove('path/to/ComfyUI/output/' + path.removesuffix(' [output]')).

[Chaoses-Ib]

Think about it again. It may be the time to add an universal util.get() to extract arbitrary values, albeit it can only work with standalone (load() from path) runtimes.

[Yahweasel]

If using standalone runtime, why not just use real mode? Having a get for arbitrary values but restritected to standalone-non-real seems like a pretty niche case.

I'm currently fussing with two different ideas, and one of them, I'm just using real mode. In the other one, the installations are on multiple computers, so no go, but the save/load latents solves my particular need.