Load A2SB Audio Model

Load A2SB Audio Model is the download-heavy front door to the comfyui-nvidia-audio-diffusion pack, and it's where your whole restoration workflow's personality gets decided. It grabs NVIDIA's Audio-to-Audio Schrödinger Bridges checkpoints from HuggingFace on first use (~6.8GB if you take the default), builds the U-Net, and hands the other two nodes a ready-to-run A2SB_MODEL bundle. Every knob on this node is a speed/quality/VRAM trade, and most of them you can ignore on the first pass - but it helps to know which ones are the honest ones.

What it actually does

The pack wraps the research release of A2SB (arXiv 2403.07634), an NVIDIA model that bridges between corrupted and clean audio instead of diffusing from pure noise. The loader pulls checkpoints from the nvidia/audio_to_audio_schrodinger_bridge repo into ComfyUI/models/A2SB, loads the weights into an attention-U-Net, and wraps it in ComfyUI's model patcher so VRAM is managed like any other model. It also flips on TF32 and cuDNN autotune, and lays tensors out channels-last for tensor cores - small things, but they're why this runs respectably on modern NVIDIA cards.

The one choice that actually changes the model: model_type. 2-split (recommended) loads two U-Nets, each handling half of the noise schedule (t below and above 0.5). That's why the default is recommended - each network only has to learn half the job, so quality goes up. 1-split loads a single checkpoint, saves ~3.4GB of download and VRAM, and trades away that quality. For long archives on a mid-range card, 1-split is a perfectly defensible call.

The inputs that matter

• model_type - 2-split (recommended) or 1-split. The quality vs. footprint decision above.
• precision - bf16 (default) is right for any modern card. fp32 is for ancient GPUs that choke on bf16, fp16 is the middle ground; don't overthink it.
• attention_type - sdpa (default) or sage. Pick sage only if you've installed SageAttention yourself (it's not in requirements.txt). If you pick it without the package, the node just falls back to SDPA with a console warning - no crash.
• quantize - none, int8wo, or int4wo, applied via torchao (weight-only). Useful on tight VRAM; you give up a little quality for it.
• use_compile - torch.compile with max-autotune. A permanent speed boost after a long one-time compile that can look like a hang the first run. Enable once, then never think about it again.
• use_ot_ode - a deterministic sampler toggle. It changes the bridge dynamics; leave it off unless you're chasing consistency between runs.

The single output is a2sb_model (A2SB_MODEL), which plugs into the a2sb_model input on A2SB Bandwidth Extension and A2SB Audio Inpainting. You load once and can feed both.

Installing it

Same path as the rest of the pack - one repo, three nodes:

cd ComfyUI/custom_nodes
git clone https://github.com/mmoalem/comfyui-nvidia-audio-diffusion
cd comfyui-nvidia-audio-diffusion
pip install -r requirements.txt   # or run install.py

That pulls einops, rotary-embedding-torch, librosa, soundfile, and torchao. If you want sage attention on RTX 40/50-series, pip install sageattention manually. ComfyUI Manager works too - search "comfyui-nvidia-audio-diffusion".

Common issues

• First load downloads a lot. The console prints [A2SB] Downloading ... (~3.4GB) per checkpoint into ComfyUI/models/A2SB. It's not hung; it's a few gigabytes of weights.
• use_compile seems frozen. That's the initial compilation, which is slow and quiet. It pays off on every run after.
• Model files land in an unexpected spot. The downloader moves them from the HF ckpt/ subfolder into models/A2SB on first load, so a second run finds them already there.
• Watch the license. The underlying A2SB weights ship under NVIDIA's non-commercial research license even though the wrapper code is MIT. Fine for personal restoration; not for a paid product.

For most people the right settings are 2-split + bf16 + sdpa and stop. Revisit sage, quantize, and use_compile only when you're processing long recordings and the wait actually hurts. It's a boring node that does its one job well - and it's the piece every A2SB workflow stands on.