擴散模型生成流程中不可或缺的 NaN 檢查

I generate 512x768 images on a 4GB GTX 1650. About 150 seconds per image in fp32 with model offload. That is the number I plan around, and it does not leave room for rerunning a batch because one file was quietly broken.

How fp16 on this card produced NaN without raising

I switched the pipeline to fp16 to see if I could buy time back. Nothing crashed. No exception, no warning, no non-zero exit code. The script finished, PIL wrote the files, and the batch looked like a success. Some of the PNGs were pure noise, or flat grey, or black.

I did not trust my eyes, so I diffed the two runs. The only change was the precision flag:

# before
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
# after
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)

Same seed, same prompt, same scheduler. fp32 gave me an image every time. fp16 gave me NaN tensors that survived all the way to image.save(). The cause was the dtype, not the prompt and not the offload. On a 4GB card, offload plus fp16 is exactly where the numeric headroom goes missing, and nothing in the library reports it.

A guard that runs before the file is written

The fix is cheap and it belongs inside the generation loop, not after it. Before decoding and before writing:

latents = latents.float()
if not torch.isfinite(latents).all():
    continue  # drop, log, retry — do not save

isfinite catches the NaN and the inf. It does not catch the case where the tensor is finite but the decode collapses to a flat image, so I added a pixel-stddev floor on the decoded array. If stddev is below my floor, treat it as a failure. That is two checks: one on the tensor, one on the pixels. Both run before save(). A bad sample costs me a retry inside the loop instead of a file on disk.

The ordering matters. The check has to sit between sampling and writing, because once the PNG exists, every downstream step believes it. I log the drop with the seed and the prompt so a retry is reproducible.

What a NaN costs downstream

A NaN that reaches disk is not one bad file. I generated and pinned 77 pieces. Of those, 7 are minted and 70 are still to mint, 5 per wallet, at a mint price of 0.00063 ETH on Base. The on-chain preload is about 0.0000008 ETH per piece, so the gas side is trivial — the cost is the manual work.

A pinned file with a mint page pointing at it means I have to find it, unpin it, regenerate it, re-pin it, and fix the page. The pipeline had already moved on to the next piece by then. That cleanup is the whole reason the guard exists.

I have about 271 tests in the suite now, and the isfinite-plus-stddev check is one of the cheapest ones to run.

Takeaway: before you call save(), run torch.isfinite(tensor).all() and a pixel-stddev floor. Two lines, run every image, and you never debug a silent NaN through a pinned file.

The collection is at https://thunderous-sunflower-a910b4.netlify.app