🚀 ComfyUI Any-Device Offload

Force any Model, VAE, or CLIP to any GPU or CPU—and keep it there (or don't).

This custom node gives you total control over where your models run and where they live in memory. It solves common "Out of Memory" (OOM) errors, enables multi-GPU workflows, and fixes persistent crashes when trying to run modern workflows (like Z_Image/Flux) on CPUs or secondary GPUs.

✨ Key Features

👇 Manual Device Selection: Force specific models to run on cuda:0, cuda:1, cpu, or Mac mps.
🧠 Intelligent VRAM Management:
- Keep in Memory (True): Pin models to VRAM for instant switching.
- Keep in Memory (False): The "Kill Switch." Automatically unloads the Model, CLIP, and VAE from VRAM immediately after generation finishes. It aggressively triggers garbage collection (gc.collect) and clears the CUDA cache (empty_cache) to free up space for other apps or workflows.
🛠️ VAE Patcher (The Crash Fixer):
- Fixes RuntimeError: Input type (c10::BFloat16) and bias type (float) should be the same.
- Fixes CuDNN error: GET was unable to find an engine to execute this computation.
- Automatically converts inputs and weights to matching types (Float32/BFloat16) on the fly without breaking the workflow.
🛡️ CPU "Safe Mode":
- Automatically intercepts xFormers calls on CPU (which usually crash with "No operator found") and redirects them to standard PyTorch attention.
- Auto-casts Float16 inputs to Float32 when running on CPU to prevent "mat1 and mat2" dtype errors.

📦 Installation

Navigate to your ComfyUI custom nodes directory:
```
cd ComfyUI/custom_nodes/
```

Clone this repository:

git clone https://github.com/FearL0rd/ComfyUI-AnyDeviceOffload.git

Restart ComfyUI.

📖 Usage

Add the node "Offload Anything (GPU/CPU)" to your workflow (found under utils/hardware).

Inputs

target_device:
- Select the hardware you want to use (e.g., cuda:0, cuda:1, cpu).
- Note: If you select CPU, the node automatically enables "Safe Mode" to prevent xFormers crashes.
vae_mode:
- Original (Default): Uses the standard VAE behavior.
- Vae Patched: ⚠️ IMPORTANT: Select this if you get black images, NaN errors, or PyTorch crashes (especially with Flux, SD3, or secondary GPUs). It forces Float32 precision and auto-casts incoming tensors to ensure compatibility.
keep_in_memory:
- True: Models stay loaded on the device (faster for repeated runs).
- False: Models are aggressively purged from VRAM immediately after the image is saved. Great for low-VRAM cards.

Connections

Connect your MODEL, CLIP, and/or VAE into the inputs.
Connect the outputs to your KSampler or VAE Decode.

💡 Troubleshooting / FAQ

Q: My generation crashes with "Input type and bias type should be the same".

A: Switch vae_mode to Vae Patched. This forces the VAE to handle the data type mismatch (usually caused by Flux/SD3 generating BFloat16 latents while the VAE expects Float32).

Q: I get an xFormers error when using CPU.

A: This node automatically patches xFormers to work on CPU. If you still see it, ensure you are running the latest version of this node and have restarted ComfyUI.

Q: Why is my GPU memory not clearing?

A: Uncheck keep_in_memory. This activates the "Kill Switch" which forces a torch.cuda.empty_cache() immediately after the VAE finishes decoding.

👨‍💻 Credits

Created to solve complex memory management and precision mismatches in ComfyUI multi-device workflows.

License: MIT

4.2 KiB Raw Permalink Blame History