I tried a small Colab Free smoke run, and here is what I found:
I’m looking at this from a first-time reader / small-contributor perspective, not as a full training review or model-quality evaluation.
My main takeaway: the repo already has a useful educational component split — VAE, latent video, DiT, flow matching, conditioning, CFG — so most of my suggestions are about making the contracts between those components easier to test and document.
I’m framing these as cheap checks that might help future readers and contributors, not as criticism of model quality.
This also seems aligned with the general style of small examples used in libraries like Diffusers, where examples are meant to be self-contained, easy to tweak, beginner-friendly, and focused on one purpose.
High-level summary
If I were adding a small onboarding/debug layer, I would probably add these:
| Area |
Useful small addition |
| install |
minimal dependency note or requirements file |
| data |
tiny synthetic .npy video example |
| VAE |
valid num_frames / temporal shape-contract note |
| VAE → DiT boundary |
native VAE checkpoint interface smoke test |
| DiT + flow |
no-download fake-token smoke test |
| config validation |
patch-size divisibility check |
| scheduler helper |
z_img wrapper/lambda pattern note |
| external models |
keep SVD / other I2V pipelines as references, not drop-in checkpoints |
The main idea is to give readers a cheap path like:
compile/import
→ synthetic .npy preprocessing
→ tiny VAE shape roundtrip
→ native VAE checkpoint interface
→ fake-token DiT/flow loss
→ tiny overfit
→ real conditioning / real training later
That way, readers can catch shape/interface issues before they spend time on expensive runs.
Scope of my smoke check
I only ran tiny no-checkpoint checks.
I tested:
- Python compile check
- import probes
- synthetic
.npy preprocessing
- tiny VAE frame-count sweep
- tiny VAE backward / toy overfit
- native VAE checkpoint path simulation
- tiny
VideoDiT + FlowMatchingScheduler forward/loss/backward
- scheduler helper direct-call vs
z_img wrapper behavior
- patch-size grid sanity check
- tiny fixed-batch DiT/flow overfit
I did not test:
- full VAE training
- full DiT training
- realistic I2V inference
- output quality
- benchmark metrics
- external I2V pipelines such as SVD / CogVideoX / LTX / Wan
So I would not infer “the model works” or “the model does not work” from this. I would only treat the results as cheap onboarding / regression-test observations.
1. Install / dependency note
In my Colab/T4 run, the import probes passed. But I did not see a dependency manifest such as:
requirements.txtpyproject.tomlenvironment.yml
A small install block or requirements file would help first-time readers reproduce the initial smoke path.
It may be useful to separate dependency layers:
| Layer |
Example dependency category |
| Core VAE / DiT / flow smoke tests |
PyTorch and basic Python libs |
| Conditioning |
transformers, image/text processors |
| Training/logging |
TensorBoard or equivalent logging deps |
| Video data utilities |
video decoding / preprocessing deps |
| Optional metrics |
LPIPS or other evaluation helpers |
This would let someone run the core VAE/DiT/flow smoke path before they need T5/CLIP downloads, video datasets, or metric dependencies.
Why I think this helps
For an educational repo, the first failure mode for a new reader is often not “the architecture is wrong”; it is usually:
- missing dependency
- wrong runtime
- wrong dtype
- wrong tensor layout
- wrong file shape
- accidentally loading a heavy dependency too early
A minimal install section could make the reader path clearer:
Minimal smoke-test install:
- enough to import VAE / DiT / flow
- no pretrained encoders
- no external video datasets
Full training install:
- conditioning encoders
- logging
- video decoding
- metrics
That also makes it easier to report issues, because users can say which dependency layer failed.
2. A tiny synthetic .npy example would help onboarding
The dataset contract is easy to test with a tiny synthetic video.
A minimal docs/example path could be:
raw .npy video: [T, H, W, 3], uint8, [0, 255]
dataset output: [3, T, H, W], float32, [-1, 1]
A small moving-square .npy example would let readers test:
- axis order
- dtype
- normalization range
- short/exact/long sequence handling
- accidental double-normalization
This avoids gated datasets, large downloads, and unrelated video-decoding issues.
In my small check, float input was rejected with an assertion. I think that is useful, because it catches accidental double-normalization early.
Suggested synthetic dataset smoke path
A small example could be something like:
examples/
synthetic_moving_square.py
smoke_dataset_preprocessing.py
The test does not need real video quality. It only needs to verify:
input: [T, H, W, 3], uint8
output: [3, T, H, W], float32
range: [-1, 1]
This checks the raw video contract before introducing:
- real video datasets
- decoders/codecs
- dataset download failures
- caption files
- training loops
- GPU memory constraints
For a first reader, that is often more helpful than immediately starting from a real dataset.
3. Document valid debug num_frames values for the keep-first temporal VAE path
This was one of the more useful shape-contract checks.
In a tiny VAE frame sweep, I saw the following pattern:
Input T |
Round-tripped through VAE? |
| 5 |
yes |
| 9 |
yes |
| 13 |
yes |
| 17 |
yes |
| 21 |
yes |
| 25 |
yes |
| 33 |
yes |
| 8 |
no |
| 16 |
no |
| 18 |
no |
This looked consistent with a keep-first temporal down/up path where debug frame counts like 4k + 1 are safer, at least for the tiny temporal_ds=2 setup I tested.
The default 17 therefore looks well chosen.
I would frame this as a debug-config / shape-contract documentation item, not as a quality issue.
A short note like this would help:
For temporal_ds=2, use debug num_frames such as 5, 9, 13, 17, 21, ...
Avoid arbitrary even values such as 8 or 16 unless you intentionally handle temporal length mismatch.
Why this is useful
Video models are shape-sensitive. A reader may naturally try a smaller debug setting like num_frames=8 to save memory. If that decodes to a shorter temporal length, the failure may appear later as a loss-shape mismatch.
So documenting valid debug frame counts can save a lot of confusion.
This is also normal for video-model docs: some video pipelines have model-specific frame-count assumptions or recommendations. The Diffusers video generation guide explicitly points readers to model-specific pipeline docs/cards for those details.
For NanoI2V, I think a tiny table near the VAE config would be enough:
temporal_ds |
Safe debug num_frames examples |
| 1 |
document from implementation |
| 2 |
5, 9, 13, 17, 21, … |
| larger |
document if supported |
The exact rule should come from the intended implementation, but the docs should prevent readers from guessing arbitrary frame counts.
4. Native VAE checkpoint path: high-value interface smoke test
This is the check I would prioritize most.
NanoI2V has both a native VAE path and a pretrained/wrapped VAE path. That means the DiT training/inference code needs a clear VAE inference contract.
A useful test would be:
Can a native VAE3D checkpoint be loaded by the same train_dit.py / inference_dit.py path
and expose the same inference interface as the pretrained VAE wrapper?
In my tiny checkpoint simulation, I saw this interface shape:
| Object / path |
Observed |
VAE3D |
had encode, decode, forward |
VAE3D |
did not expose encode_mean |
VAE3D |
did not expose set_inference_mode |
native train_dit.py VAE path |
expected set_inference_mode |
| native inference path |
loaded VAE3D, but later expected encode_mean |
| small adapter-style probe |
worked |
I would not call this a model-quality issue. It is more of an interface-contract item between VAE training and DiT training/inference.
A small bridge may be enough, depending on the intended semantics:
native VAE inference contract:
- encode_mean(video) -> latent mean or deterministic latent
- decode(latent) -> reconstructed video
- eval / inference mode behavior
The important part is that the native VAE checkpoint path and the pretrained VAE wrapper path expose the same small inference API.
Why this seems important
This check connects two major stages:
VAE training
→ VAE checkpoint
→ DiT training
→ DiT inference
If the native VAE checkpoint can be trained but cannot be loaded into the DiT path with the expected interface, readers may get stuck before they can test I2V training/inference.
A tiny checkpoint simulation is cheap and can catch this before anyone runs expensive training.
Possible test:
1. Instantiate a tiny VAE3D.
2. Save a tiny checkpoint.
3. Load it through the same code path used by train_dit.py.
4. Check that the loaded object exposes:
- encode_mean
- decode
- expected eval/inference behavior
5. Probe one tiny video:
- input: [1, 3, T, H, W]
- latent: [1, C, T', H', W']
- recon: [1, 3, T, H, W]
This should stay as a tiny smoke test, not a full training test.
One possible design choice:
| Choice |
Meaning |
Add encode_mean to VAE3D |
native VAE directly matches wrapper contract |
Add a NativeVAEWrapper |
keep training VAE class separate from inference API |
Change DiT code to call encode(...) explicitly |
avoid extra method, but document expected return tuple |
| Keep both paths separate |
possible, but then docs should clearly say which VAE path supports which script |
I do not know which is best for your intended design. The useful part is probably the smoke test itself.
5. No-download DiT + flow smoke test
A no-download DiT/flow test seems very useful for this project.
Before loading T5/CLIP, before loading a VAE checkpoint, and before using a real dataset, it is possible to test the core DiT/flow interface with fake tensors:
random z_t
randomflow test seems very useful for this project.
Before loading T5/CLIP, before loading a VAE checkpoint, and before using a real dataset, z_img
fake text tokens
fake image tokens
FlowMatchingScheduler.add_noise
velocity target
VideoDiT forward
MSE loss
backward
In my toy run, the following worked:
VideoDiT output shape matched the velocity target shape.- flow math was easy to verify independently.
- loss was finite.
- backward passed.
- a tiny fixed-batch DiT/flow overfit decreased loss over 25 steps.
Again, I would not treat that as quality evidence. It is only an interface / optimizer sanity check.
This is especially useful because video generation mixes time, space, conditioning, and latent shapes. The Diffusers video generation guide describes video generation as extending image generation into video-specific parameters and memory constraints, so isolating the DiT/flow contract before adding all expensive components seems valuable.
Suggested no-download DiT/flow smoke test
A minimal test could check:
Inputs:
- z_t: [B, C, T, H, W]
- z_img: [B, C, 1, H, W]
- txt_tokens: [B, L_text, D_text]
- img_tokens: [B, L_img, D_img]
Expected:
- model output: [B, C, T, H, W]
- flow target: [B, C, T, H, W]
- finite MSE loss
- backward pass works
This test would not validate text/image semantics. It only validates that the DiT + flow path is wired correctly.
A possible ladder:
Test 1: flow math only
Test 2: DiT forward shape
Test 3: DiT + flow MSE loss
Test 4: backward pass
Test 5: tiny fixed-batch overfit
Test 6: real T5/CLIP tokens
Test 7: real VAE latents
That order makes it easier to localize failures.
6. FlowMatchingScheduler.training_loss and the z_img wrapper pattern
One small interface-clarity point:
A direct reusable helper call like this:
FlowMatchingScheduler.training_loss(dit, ...)
does not naturally pass the I2V-specific z_img argument expected by VideoDiT.
In my small check:
- direct helper call failed because
VideoDiT.forward() expected z_img
- a small wrapper/lambda that closes over
z_img worked
- Euler sampling with the same wrapper pattern also worked
This does not look like a blocker, because the current training/inference paths can handle z_img manually. But if training_loss is intended to be reused as a generic helper, documenting the wrapper pattern could avoid confusion.
Possible wording in docs
Something like:
# Pseudocode
def model_with_image_latent(x_t, t, txt_tokens, img_tokens):
return dit(x_t, t, txt_tokens, img_tokens, z_img)
loss, pred, target = scheduler.training_loss(
model_with_image_latent,
x0,
txt_tokens,
img_tokens,
null_txt_tokens,
null_img_tokens,
)
The exact API may differ, but documenting the idea would make the helper easier to reuse.
Possible options:
| Option |
When useful |
| Document wrapper/lambda pattern |
minimal change |
Add z_img to scheduler helper API |
if helper is meant to be I2V-aware |
| Keep helper generic and avoid it in I2V training |
also fine, but docs should say so |
7. Patch-size divisibility could be documented or asserted
The DiT patch-size grid is another small shape-contract item.
In a tiny patch-grid check:
| Patch size |
Latent shape divisible? |
Output shape matched input latent? |
(1, 1, 1) |
yes |
yes |
(1, 2, 2) |
yes |
yes |
(3, 2, 2) |
yes |
yes |
(2, 2, 2) with latent T=3 |
no |
no |
(1, 3, 2) with latent H=4 |
no |
no |
(1, 2, 3) with latent W=4 |
no |
no |
The non-divisible cases could still run, but produced shorter output shapes. That can be confusing later.
So it may be useful to document or assert early:
patch_t must divide latent_t
patch_h must divide latent_h
patch_w must divide latent_w
This is a good candidate for a small config validation check.
Why this is worth an early assert
Patchification is a boundary where a config error can become a later shape mismatch.
For readers, an early error like this is easier to understand:
Invalid patch_size=(2,2,2) for latent_shape=(3,4,4).
patch_t must divide latent_t.
than a later mismatch in loss, unpatchify, or decoded output shape.
This could be part of a general validate_config() or small smoke test.
8. External I2V models as references, not drop-in checkpoints
External I2V pipelines can still be useful, but I would keep them separate from NanoI2V-native smoke tests.
For example, Stable Video Diffusion in Diffusers is useful as a reference for image-to-video input/output conventions and memory tradeoffs. Its docs discuss settings such as CPU offload and decode_chunk_size, where lower memory can trade off with temporal consistency.
But external I2V models are not drop-in NanoI2V checkpoints. I would separate the docs into something like:
| Path |
Purpose |
| NanoI2V no-checkpoint smoke tests |
compile/import/shape/interface checks |
| NanoI2V native checkpoints |
exact VAE/DiT/config pairing |
| External I2V references |
compare pipeline conventions and memory tradeoffs |
This avoids mixing up “reference pipeline” with “compatible weights”.
Suggested separation
Possible documentation sections:
1. Learn the architecture
- read modules
- run fake-token tests
- use synthetic data
2. Train NanoI2V components
- VAE training
- native VAE checkpoint
- DiT training
- exact config pairing
3. Run NanoI2V inference
- native VAE checkpoint
- DiT checkpoint
- matching config
4. Compare with external I2V pipelines
- SVD
- CogVideoX-I2V style systems
- LTX / Wan / others
- clearly marked as references only
This keeps the educational repo path clean while still giving readers useful context.
Suggested small additions
If I were turning these observations into repo/docs tasks, I would suggest:
Documentation
- Minimal install block or dependency file.
- Tensor-shape table:
- raw video
- dataset output
- VAE input
- VAE latent
- image latent
- DiT input/output
- decoded video
- Valid debug
num_frames note for the keep-first temporal VAE path.
- DiT patch-size divisibility note.
- Clear native VAE vs pretrained VAE wrapper interface contract.
Examples / smoke tests
examples/synthetic_moving_square.pyexamples/smoke_dataset_preprocessing.pyexamples/smoke_vae_frame_contract.pyexamples/smoke_native_vae_checkpoint_interface.pyexamples/smoke_dit_flow_fake_tokens.py
Cheap test ladder
Level 0: py_compile + import probes
Level 1: synthetic .npy preprocessing
Level 2: tiny VAE frame roundtrip
Level 3: native VAE checkpoint interface
Level 4: fake-token DiT + flow loss/backward
Level 5: tiny fixed-batch overfit
Level 6: real conditioning encoders
Level 7: full training / inference
The main idea is: let readers validate component contracts before they spend time on expensive training.
Possible first-reader path
A first reader could follow:
1. Read repo overview.
2. Run compile/import smoke test.
3. Generate a tiny synthetic .npy video.
4. Check dataset preprocessing.
5. Run VAE frame-count sweep.
6. Run native VAE checkpoint interface test.
7. Run fake-token DiT/flow smoke test.
8. Only then try real conditioning encoders or real training.
That path would make the repo easier to learn from without reducing the project’s scope.
Closing note
Overall, I think the component separation is already useful for an educational I2V repo. My main suggestion is to make the contracts between components more visible and cheaply testable.
I would frame the smoke tests above as:
- not benchmarks
- not quality evaluation
- not a requirement that the full project taget Colab Free
- just cheap checks that help future readers and contributors avoid shape/interface confusion before expensive runs