Dit unit tests #68
Merged
Dit unit tests #68
Conversation
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
Contributor
sajadn
commented
Nov 20, 2025
sajadn
commented
- add unit tests
- very minor cleanings
- add missing iter to the DiffusionDataModuleConfig
Signed-off-by: sajadn <[email protected]>
Signed-off-by: Sajad Norouzi <[email protected]>
Signed-off-by: Sajad Norouzi <[email protected]>
Signed-off-by: sajadn <[email protected]>
sajadn
force-pushed
the
dit_unit_tests
branch
from
4b75aa7 to
2f46b43
Compare
Signed-off-by: sajadn <[email protected]>
Signed-off-by: sajadn <[email protected]>
Contributor
Author
|
/ok to test fa1b884 |
Signed-off-by: sajadn <[email protected]>
Contributor
Author
|
/ok to test c376224 |
Signed-off-by: sajadn <[email protected]>
sajadn
force-pushed
the
dit_unit_tests
branch
from
c376224 to
644970d
Compare
Signed-off-by: sajadn <[email protected]>
Contributor
|
/ok to test 3e93c2a |
abhinavg4
previously approved these changes
Dec 1, 2025
Contributor
abhinavg4
left a comment
abhinavg4
left a comment
There was a problem hiding this comment.
Looks good. Let;s merge once tets pass
| "output_params": ["-f", "mp4"], | ||
| } | ||
|
|
||
| print("video_save_path", video_save_path) |
Contributor
There was a problem hiding this comment.
Good cleanup! Removing debug print statements keeps the output clean in production.
| def build_datasets(self, context: DatasetBuildContext): | ||
| return self.dataset.train_dataloader(), self.dataset.val_dataloader(), self.dataset.test_dataloader() | ||
| return ( | ||
| iter(self.dataset.train_dataloader()), |
Contributor
There was a problem hiding this comment.
Nice improvement! Wrapping dataloaders with iter() makes the interface more explicit and reduces potential confusion in downstream usage.
Contributor
Author
|
/ok to test f859fbb |
huvunvidia
reviewed
Dec 2, 2025
| DiTModelProvider: Configuration for the DiT-S model. | ||
| """ | ||
| return DiTModelProvider( | ||
| return DiTXLModelProvider( |
Contributor
There was a problem hiding this comment.
Nit: In docstring it notes "DiT-S" but we provide DiTXLModelProvider.
| "mediapy>=1.2.4", | ||
| "megatron-bridge", | ||
| "wandb[media]>=0.23.0", | ||
| ] |
Contributor
There was a problem hiding this comment.
Just to make sure, are mediapy and wandb included in OSRB?
Any dependency needs to be approved for open-source.
Contributor
Author
There was a problem hiding this comment.
oops, this is a great question. I'm not sure about that.
Contributor
Author
|
/ok to test 99ada83 |
pablo-garay
approved these changes
Dec 3, 2025
Collaborator
|
Approved from Automation perspective |
Contributor
|
/ok to test f007a7f |
lbliii
pushed a commit
that referenced
this pull request
Dec 3, 2025
* edm and data preprocess tests. Signed-off-by: sajadn <[email protected]> * Minor cleanings for DiT. Signed-off-by: Sajad Norouzi <[email protected]> * add dit unit test. Signed-off-by: Sajad Norouzi <[email protected]> * add iter to the DiffusionDataModule. Signed-off-by: sajadn <[email protected]> * add missing copyright. Signed-off-by: sajadn <[email protected]> * use 'no caption' if caption is not present. Signed-off-by: sajadn <[email protected]> * fix dit inference bug. Add wanbd to inference code. Signed-off-by: sajadn <[email protected]> * update the DiT configs to be aligned with the original paper. Signed-off-by: sajadn <[email protected]> * add wandb[video] and mediapy to uv. Signed-off-by: sajadn <[email protected]> * adjust pos_ids in mock_dataset to have batch dimension, fuse adaLN layers, use DiTSelfAttention. Signed-off-by: sajadn <[email protected]> * fix the diffusion sample size bug. Signed-off-by: sajadn <[email protected]> * fix broken tests. Signed-off-by: sajadn <[email protected]> --------- Signed-off-by: sajadn <[email protected]> Signed-off-by: Sajad Norouzi <[email protected]> Co-authored-by: Abhinav Garg <[email protected]> Signed-off-by: Lawrence Lane <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.