Proof of concept fix for Flax LCM PR #1
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Splitting the key makes the denoising loop work without the artifacts previously observed.
This is not meant to be the final solution, but a proof-of-concept demonstration. I'm not keen on introducing RNG Key management at the pipeline level; instead, I wonder if we could set the key as part of the scheduler state and have the scheduler manage it. In any case, I suggest to merge this in the in-progress PR and move the discussion there.
Note also that I believe the solution in
FlaxDDPMScheduler
is wrong, as it will always generate the same random variance every time the function is called. We do need to carry the RNG state.I have tested with
DEBUG
True
andFalse
, but without jitting the pipeline.(Commit 7e0a0cd is unrelated, sorry; it's already in main huggingface#6272)