Modified sampling to be closer to genies inference + made it faster#11
Open
maharajamihir wants to merge 2 commits intoFLAIROx:mainfrom
Open
Modified sampling to be closer to genies inference + made it faster#11maharajamihir wants to merge 2 commits intoFLAIROx:mainfrom
maharajamihir wants to merge 2 commits intoFLAIROx:mainfrom
Conversation
…avoid recompile compile, 3) jitted sampling, 4) scan instead of loop
maharajamihir
changed the title
- Overhauled Genie.sample() to support autoregressive generation of up to `seq_len` future frames, following the approach in Figure 8 of the paper.
- Input frames are tokenized once, and future frames are generated autoregressively in token space.
- All frames are detokenized in a single pass at the end.
- Added detailed docstring explaining the sampling process and tensor dimensions.
- Replaced the old MaskGIT loop with a step-wise scan over timesteps, masking and updating tokens for each future frame.
- Updated MaskGITStep mask update logic:
- Flattened and sorted token probabilities for mask updates using einops.
- Changed mask update to operate on the flattened token dimension, then reshaped back.
- Fixed mask indexing to ensure correct number of tokens are unmasked at each step.
- Added einops import for tensor reshaping.
- Minor type and variable name improvements for clarity and correctness.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
This PR is a mirror of p-doom#60. I have tested it on our repo on our Minecraft world-modelling dataset and checkpoints. Sampling of 16 frames went from 6min30s to 1min20s, with compilation being the largest fraction of the time.
Autoregressively samples up to
seq_lenfuture frames, following Figure 8 of the paperNote: