Skip to content

Switch SD3 and Flux2 noise generation to TorchRandomSource - #253

Draft
stikves wants to merge 1 commit into
apple:mainfrom
stikves:sukru/torch-rng-parity-151
Draft

stikves wants to merge 1 commit into
apple:mainfrom
stikves:sukru/torch-rng-parity-151

Conversation

@stikves

@stikves stikves commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

Summary

SD3 and Flux2 generated their initial latent noise with the NumPy RNG instead of torch.randn, which the diffusers reference pipelines use. For a given seed the generated image therefore did not match the reference. This routes both pipelines through the torch noise source and corrects TorchRandomSource so its output matches torch.randn in float32.

Changes

  • SD3Pipeline / Flux2Pipeline: generate initial noise with sourceType: .torch.
  • TorchRandomSource.normalArray: implement PyTorch's normal_fill_16 fill order and fix the trailing-remainder refill (nextFloat, not nextDouble).
  • generateNoise(.torch): route through normalArray rather than the scalar nextNormal() loop.
  • Tests: eight TorchRandomSource parity cases against torch.randn (seeds 0/42; counts 16/17/32/4096) covering the scalar, batch, boundary, remainder, and realistic-shape paths.

Validation

  • RNG parity against torch.randn is float32-exact (max abs error <= 1e-6, cosine 1.000000) up to the largest pipeline latent count (SD3 262144, Flux2 524288).
  • End to end at seed 42: both pipelines reproduce the diffusers torch-generator reference composition; before the change each produced a different image.

Closes #151

SD3 and Flux2 generated their initial latent noise with the NumPy RNG, so
they did not match torch.randn, which the diffusers reference pipelines use.
Three fixes to the torch RNG path:

1. TorchRandomSource.normalArray re-filled the trailing partial block with
   nextDouble(); PyTorch uses nextFloat(). PyTorch also fills every element
   with uniforms before applying Box-Muller, then re-fills the last 16 for
   the remainder. Matching this corrects count=17 and every size that is not
   a multiple of 16.

2. generateNoise(.torch) used the scalar nextNormal() loop, which yields a
   different sequence than torch.randn's batch-16 fill. Route it through
   normalArray.

3. SD3Pipeline and Flux2Pipeline called generateNoise without a sourceType
   and defaulted to .numPy. Set both to .torch.

Adds TorchRandomSource parity tests against torch.randn reference values
(seeds 0/42; counts 16/17/32/4096) covering the scalar, batch, boundary,
remainder, and realistic-shape cases.

Closes apple#151
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Cleanup] Diffusion: TorchRandomSource Never Used

1 participant