nx_diffusion_loop.nx
buildroot/runtime/nx_diffusion_loop.nx
about
nx_diffusion_loop.nx -- N-step diffusion sampler composer.
L4 brick. Closes the image-gen forward path: given noise + a
UNet block per layer + a denoising sigma schedule, produces a
denoised latent over N sampler steps.
Composes:
nx_sampler (DPM++ 2M Euler-first + 2nd-order-mid,
+ Karras schedule via nx_root)
nx_unet_block (per-layer ResBlock; caller stacks N layers)
NxTensor (L1 4D containers)
nx_loop.LoopVerdict (control)
===== High-level shape ==========================================
noise -> x_0
for i in 0..n_steps:
sigma_now = schedule[i]
sigma_next = schedule[i+1]
// Denoising prediction: caller's UNet (stack of nx_unet_block).
// v1 takes a callback function pointer so the composer is
// architecture-agnostic.
d_now = denoise(x, sigma_now)
// DPM++ 2M step.
if i == 0:
x = nx_sampler_step_euler(x, d_now, sigma_now, sigma_next)
else:
x = nx_sampler_step_2nd(x, d_now, d_prev, ratio)
d_prev = d_now
return x
===== v1 scope ==================================================
v1 makes the denoiser callback signature explicit:
func denoise(x_lanes: *i64, n_lanes: nx_int, sigma_q14: i64,
dependencies 4 imports · 0 importers
imports: nx_syscalls.nxnx_tier.nxnx_loop.nxnx_sampler.nx
imported by: nobody (leaf or entry point)
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 71 | const NX_DL_Q10: nx_int = 1024 |
| 72 | const NX_DL_Q14: nx_int = 16384 |
| 76 | const NX_DL_OK: nx_int = 0 |
| 77 | const NX_DL_ERR_BAD_DIMS: nx_int = 1 |
| 78 | const NX_DL_ERR_BAD_SCHEDULE: nx_int = 2 |
| 79 | const NX_DL_ERR_OOM: nx_int = 3 |
| 80 | const NX_DL_ERR_DENOISER: nx_int = 4 |
| 81 | const NX_DL_N_VERDICTS: nx_int = 5 |
functions
| 83 | func nx_dl_verdict_is_valid(v: nx_int) -> nx_int called by 1: main |
| 95 | func _dl_ratio_q10(sigma_now_q14: i64, sigma_next_q14: i64) -> nx_int called by 1: nx_diffusion_loop |
| 116 | func nx_diffusion_loop(x_io: *i64, n_lanes: nx_int, |
| 186 | func _dl_identity_denoiser(x_lanes: *i64, n_lanes: nx_int, sigma_q14: i64, |
| 206 | func main() -> i64 |