KSampler: use the same noise samples across a batch

Using an ancestral sampler on a batch currently yields results that can
only be reproduced at the same index in a batch of the same size.  For
example, if you want to reproduce a single result from a batch
generation, you have to re-run the entire batch.

This change ensures that the same noise samples are used for every
latent in the batch, leading to consistent and reproducible results
regardless of batch size/order.  In addition, the noise is now a view
of a C×H×W tensor, which uses less memory than a full B×C×H×W.
This commit is contained in:
Devin J. Pohly 2025-04-18 13:20:01 -05:00
parent f3b09b9f2d
commit 80378f5b86

View File

@ -77,7 +77,9 @@ def default_noise_sampler(x, seed=None):
else:
generator = None
return lambda sigma, sigma_next: torch.randn(x.size(), dtype=x.dtype, layout=x.layout, device=x.device, generator=generator)
# To ensure reproducible generations regardless of batch order, use the same
# noise sample for every latent in the batch.
return lambda sigma, sigma_next: torch.randn(x.shape[1:], dtype=x.dtype, layout=x.layout, device=x.device, generator=generator).expand(x.shape)
class BatchedBrownianTree: