Nodes/kubakub nodes/kubakub onnx style transfer
ComfyUI Node

kubakub onnx style transfer

Bring your own trained pix2pix model, run it at full resolution

By kubakubkub·Created about a month ago·Updated 2 days ago· 10
kubakub onnx style transfer
  • image
  • styled_image
◄model_path►
◄modetiled►
◄tile_size512►
◄overlap64►
◄blend1.00►
◄passes1►
◄keep_loadedfalse►

This is the odd one out in the pack: a small ONNX runtime that runs your model, whatever it is. A pix2pix or similar image-to-image network that you trained yourself - sketch to render, clay to stylised elevation, some house style you developed for this project - exported to .onnx, dropped in a folder, and applied at full resolution by running it over overlapping tiles.

Why ONNX? Because it decouples the model from ComfyUI's Python. A torch checkpoint ties you to a torch version, a diffusers version, an architecture someone else has to maintain. An .onnx file is a frozen graph: onnxruntime loads it and runs it, on GPU if that build has CUDA, on CPU if it doesn't. For a model you trained once and want to run on a facade for the next two years, that's the right kind of boring.

Inputs

image is what you're restyling. A batch is fine - each image is done on its own.

model_path is the .onnx file. A full path, or just a name like final_generator.onnx if it lives in ComfyUI/models/onnx. There is no sample model for this node and the pack ships none - if the field is empty you get a clear error rather than a fallback. That's unusual in this pack, where almost every file-reading node has a built-in sample, and it's because the model is yours; nobody else's would be useful.

mode decides the inference strategy. tiled (default) runs the model on overlapping squares at full resolution. single squeezes the whole image into one square, runs it, and scales it back - fast and soft, and genuinely useful for checking the model does roughly the thing you remember before you spend minutes tiling. tile_size (default 512) is the square the model works at, and you should set it to the size the model was trained at, usually 512. overlap (default 64) is the pixels neighbouring tiles share and blend across - that blending is what keeps a tiled run seam-free, and 64 is enough for most piecewise networks.

blend (default 1) mixes the styled result with the original: 1 is fully styled, 0 is the untouched input. Setting it around 0.7 is the standard move for a sketch-to-render model that's still learning, so you keep some of your own drawing's line quality.

passes runs the model again on its own output, 1 to 5 times. It's a blunt instrument - the style gets stronger and the result gets further from your input with every pass - but it's cheap to try, and 2 is sometimes exactly right where 1 is timid.

keep_loaded is the VRAM switch. Off by default, and it matters: with it off, the ONNX session is torn down after the run so the GPU goes back to your diffusion models, which is what you want on a 12 GB card juggling both. Turn it on when you're batching hundreds of images through the same model and want to skip the reload every time.

One output: styled_image, same size as the input. On a 3200 or 4000 px facade render that's the whole point of the tiled mode - you get your full-resolution image back, not a 512 px square scaled up.

Install

cd ComfyUI/custom_nodes && git clone https://github.com/kubakubkub/ComfyUI_KubaNodes
python -m pip install -r ComfyUI_KubaNodes/requirements.txt
python -m pip install onnxruntime        # CPU
# or, for GPU:
python -m pip install onnxruntime-gpu

onnxruntime is not in the pack's requirements.txt on purpose - exactly one node needs it, so it's optional, and the startup line tells you when it's missing ([Kub] not installed (only these features are off): onnx style transfer -> pip install onnxruntime). Don't install both builds; pick one. With onnxruntime-gpu and a CUDA that matches your torch, inference is fast; with the CPU build it still works, just slowly, which for a batch job is often fine.

Two honest notes

ONNX is for portability, not for speed. The runtime can be quicker or slower than the torch model it came from, depending on the graph - the common experience in this ecosystem is that a straight ONNX export of a heavy model is dramatically slower, not faster, and people export for deployment reasons rather than performance ones. If your torch model is already running in ComfyUI happily, there's no free lunch waiting here.

A per-pixel model has no idea what a facade is. A pix2pix network trained on your own pairs will do exactly what you taught it and nothing else - it won't fix a bad input, won't keep a window where the window should be, and will happily restyle a lamp post into a tree if that's what the training data suggests. Tiled mode makes that worse in a specific way: each tile is processed independently, so a strong style can produce slight inconsistencies across tile boundaries. overlap blending hides the seams, but if you see repeated motifs in a grid pattern, that's tile independence, not a bug - and the fix is a larger overlap, or a model trained on larger crops.

Where this node earns its place in a facade pipeline is pre-visualising a house style. Run a clay render through it, look at the result, decide whether the commissioned look is worth commissioning. It's the cheapest test you'll get, and it doesn't touch your diffusion models at all.

Categorykubakub/lab

Inputs (8)

NameTypeDefaultDescription
imageIMAGEThe image (or batch) to restyle; each image is done on its own.
model_pathSTRINGAn .onnx file: a full path, or a name inside ComfyUI/models/onnx.
modeoptCOMBOtiledtiled runs the model on overlapping squares at full resolution. single squeezes the whole image into one square and scales it back (fast, softer).
tile_sizeoptINT512256–1024Square size in pixels the model works at. Use the size it was trained at (usually 512).
overlapoptINT640–256tiled only. Pixels neighbouring tiles share and blend across, so no seams show.
blendoptFLOAT1.000–1Mix with the original: 1 = fully styled, 0 = the original image.
passesoptINT11–5Run the model again on its own result this many times for a stronger style. 1 = once.
keep_loadedoptBOOLEANfalseKeep the ONNX session (and its VRAM) between runs; off frees the GPU for the diffusion models.

Outputs (1)

NameTypeDescription
styled_imageIMAGEThe restyled image, same size as the input.