[question] any plan for higher compression ratios?

[eisneim]

LTX-video is the fastest open source video model, its has compression ration of 32x32x8 but with 128 channels, reduce spatial dimension and temporal dimension but add more channels will make the diffusion process much faster;

So I wonder if LeanVAE will experiment this direction?

[hyc9]

Thanks for the suggestion — we'll consider this! Feel free to train such a model using the training code we have provided. You can adjust the patch_size and latent_dim to achieve higher compression ratios and increased latent channels!

[eisneim]

@hyc9 I use this config to match the LTX's compression ratio

arg_dc = {
    'embedding_dim': 1024, 
    'latent_dim': 128, 
    'ista_iter_num': 2, 
    'ista_layer_num': 2, 
    'l_dim': 128 * 2,  # i think it should be multiply 2x4x4=32 to compensate bigger patch size losing pixel space context 
    'h_dim': 384 * 2,  # but if l_dim and h_dim are too big that will make the whole vae too big;
    'sep_num_layer': 2, 
    'fusion_num_layer': 4,
    'sample_rate': 1, 
    'dynamic_sample': False, 
    'use_tile_inference': False,
    "patch_size": (4, 16, 16), # add this option to Encoder_Arch and Decoder_Arch
}
args = argparse.Namespace(**arg_dc)

vae = LeanVAE(args)

test_video = torch.randn(1, 3, 9, 256, 256)
z = vae.encode(test_video)

print(z.shape)
# torch.Size([1, 128, 2, 8, 8])

I will start doing some test training then train LTX as diffusion model using this new VAE

[hyc9]

@eisneim It looks great. Looking forward to seeing the results!

[eisneim]

Reporting my training progress, using 140K video clips each 2 seconds long(48~60 frames) size: 512x288, 384x384, sample 17 frames for this first test run, Hardware: 2 RTX 4090 48GB;

CUDA_VISIBLE_DEVICES=0,1  python leanvae_train.py \
    --default_root_dir "../LeanVAE_ex/run1" \
    --bf16 \
    --num_nodes 1 \
    --gpus 2 \
    --grad_clip_val 1.0 \
    --lr 5e-5 \
    --lr_min 1e-5 \
    --warmup_steps 5000 \
    --discriminator_iter_start 600000 \
    --progress_bar_refresh_rate 50 \
    --max_steps 700000 \
    --data_path '' \
    --train_datalist "/home/eisneim/dataset/leanVAE_data/file_lists" \
    --val_datalist "/home/eisneim/dataset/leanVAE_data/val_list" \
    --batch_size 10 \
    --num_workers 10 \
    --sample_rate 1 \
    --sequence_length 17 \
    --latent_dim 128 \
    --ista_iter_num 2 \
    --ista_layer_num 2 \
    --embedding_dim 1024 \
    --l_dim 256 \
    --h_dim 768 \
    --sep_num_layer 2 \
    --fusion_num_layer 4 \

Currently Progress:

gs-154428_e-000016_b-007500_reconstructions.mp4

gs-200343_e-000021_b-007500_reconstructions.mp4

Gan loss has not kicked in yet, the results are still pretty bad, but i think I'm on the right track. Also I'm still cropping and triming video clips from film and TV shows, 140k video is not be enough.

@hyc9 Any segguestions, tips or insights would be appreciated! 🤝

[hyc9]

@eisneim Thanks for sharing! Here are some training insights from my own experience that might help:

1. Avoid letting the KL loss dominate the total loss. Ideally, its magnitude should be at least an order of magnitude lower than the other loss terms. Otherwise, you might observe KL going down while reconstruction-related losses increase, which is not what you want.

2. Keep training! The recon and perceptual losses tend to decrease more slowly later in training, but you'll still see steady improvements across epochs. In my experience, with enough training steps, the model can eventually reconstruct samples without visible artifacts (like in your example), although the results are over-smoothed and blurred. Achieving a low recon/perceptual loss baseline is really helpful for stabilizing GAN training later. After introducing the GAN loss, the recon quality will become much sharper.

3. It looks like increasing the temporal compression ratio significantly affects reconstruction quality. I'm not sure if that's just due to undertraining, but if the issue persists, it might be worth trying other compression setups like 64×64×4.

[hyc9]

@eisneim I just realized an important detail I forgot to mention earlier. In my training setup, the temporal compression ratio was set to 4, so k was also set to 4 accordingly (

LeanVAE/LeanVAE/models/autoencoder_pl.py

Line 85 in 77186e3

k = 4

). In your case, using k = 8 will lead to more balanced training and much better performance.

[eisneim]

@hyc9 Great thanks! i'll try and let you know the results after the training is done 🤝

[eisneim]

@hyc9 after training 17 frame for few days, the results still has artifacts, i guess when use large temporal compression like 8, 17 frames is not enough, 17 -> 3 in latent, which is not enough for cnv3d to learn; so i'm start from this model weight to train 25 frames;

here is the results training 17 frames

gs-407520_e-000015_b-018000_inputs.mp4

gs-407520_e-000015_b-018000_reconstructions.mp4

charts:

Also, previously i made a mistake, i simply just cut movies into 4 seconds clip, this will cause a clip have hard cut in it; now i'm using optical flow to detect scene change, this method is more accurate than FFMPEG's scene split or python lib "pySceneDetect"; I also use optical flow to calculate motion scores to filter out dialog scenes which has no camera movement; my new dataset should be better for the model to learn temporal change when using clean clips and longer frame sequence(17 -> 25)

Is there anything i need to change for VAE encode / decode when i use 25 as default training length?

[hyc9]

@eisneim Thanks for the update! I think just changing args.sequence_length to 25 should be ok.

I suggest further training in Stage 1 for the unnatural motion. In my experience, Stage 2 mainly improves LPIPS, but if Stage 1 isn't sufficiently trained, GAN training becomes unstable and hard to train. Training Video VAE is very time-consuming — I needed several weeks on 8 GPUs to get good results.

[eisneim]

@hyc9 after many days of training, the model is still struggles to handle large temporal changes, I guess i need to increase the layer size that handles high frequency output of wavelet transform to work harder on temporal changes? should i increase ista_iter_num ? or h_dim or fusion_num_layer ?

    --ista_iter_num 2 \
    --ista_layer_num 2 \
    --embedding_dim 1024 \
    --l_dim 256 \
    --h_dim 768 \
    --sep_num_layer 2 \
    --fusion_num_layer 4 \

https://github.com/user-attachments/assets/7f1047fb-626a-4475-a450-de489ec4cc8c
https://github.com/user-attachments/assets/bc6d2292-a127-43a7-bec6-c4fddf725e91

[hyc9]

@eisneim There isn't a definitive answer to this. I think it's worth trying to increase both ista_iter_num and fusion_num_layer. Let us know how it goes!

[NilanEkanayake]

I've been working on a similar project for a while. Using 4x8x8 patching but discrete (FSQ) with a ~4k codebook, I had to train around 300k steps with BS8 to get good temporal motion.

If you're looking for a good video dataset, I use https://huggingface.co/datasets/Vchitect/Vchitect_T2V_DataVerse
I also re-encoded it to 256p with the fastdecode h264 flag for fast processing: https://huggingface.co/datasets/NilanE/Vchitect_T2V_DataVerse_256p_8fps_wds (ignore the 8fps in the name, I didn't end up changing the original FPS)

You can find my dataloader code here: https://github.com/NilanEkanayake/TiTok-Video/blob/main/dataset/video_dataset.py

I used a non-causal Encoder/Decoder setup with ResNAF blocks as a platform to test my GAN code as it converges faster than the 1D transformer tokenizer I intend to do final training with.

Here are some reconstructions from the first stage (128p, 16 frames, 8fps, L1+LPIPS):

The model I test with is very small, however (~15m params).

[NilanEkanayake]

@eisneim I just realized an important detail I forgot to mention earlier. In my training setup, the temporal compression ratio was set to 4, so k was also set to 4 accordingly (

LeanVAE/LeanVAE/models/autoencoder_pl.py

Line 85 in 77186e3
k = 4
). In your case, using k = 8 will lead to more balanced training and much better performance.

In the code, it seems k refers to the number of frames used to sample the perceptual loss from (+1 for the first frame), if I'm reading it correctly? I do the same in my codebase, but I use 3 (no first frame as well) for 16-frame videos without a noticeable quality sacrifice. It does speed up training a lot and reduce memory usage.

@eisneim you could try a transformer-based 1D tokenizer like I'm using. It supports arbitrary output token counts irrespective of the patch size. Essentially, you concatenate some empty/noised 'latent tokens' to the video patches after the initial linear layer, pass the lot through the encoder, then discard the patches and send the latent tokens to the decoder, which concatenates them with empty video patches, then projects those patches back to pixel space after the transformer layers, discarding the latent tokens. In this way, you can use a 4x8x8 patch size but still have a small token count, which avoids the blur. You're forced to use transformers, however, for the global attention. Also, I'm not sure how compatible those tokens would be with a typical video diffusion model. You could set a latent token count the same as what a higher patch size would return? What might be the largest draw is that 1D tokenizers have a fixed resolution and token count after training. See https://github.com/hywang66/LARP.

[eisneim]

@NilanEkanayake thanks for the info! I'v tried FSQ + Transformer autoregressive method before, but the quality is not on part with VAE + Flow matching models for video generation;

4x8x8 is not suitable for 720p video generations, Wan-video 2.1 1.3B VAE is what I'm currently using, 4x8x8 ratio make the transformer token sequence way way to long, causing scale_dot_product_attention to OOM on 48G GPU I had to to using many other tricks to reduce VRAM usage;

LTX-Video has good ratio: 8x32x32 so the token sequence length is short enough, the reconstruction quality is good if training DiT + VAE all together, but I only have 48GB vram i can only train one model at once;

LeanVAE is very promising, I'm still testing bigger LeanVAE, if 8x32x32 won't be able to get good quality, i'll go with 4x16x16, which is still better than Wan-video VAE or HunyuanVideo VAE;

Diffusion forcing,, causVid and FramePack is more promising autoregressive video model

Thanks for the dataset links!

[NilanEkanayake]

@NilanEkanayake thanks for the info! I'v tried FSQ + Transformer autoregressive method before, but the quality is not on part with VAE + Flow matching models for video generation;

4x8x8 is not suitable for 720p video generations, Wan-video 2.1 1.3B VAE is what I'm currently using, 4x8x8 ratio make the transformer token sequence way way to long, causing scale_dot_product_attention to OOM on 48G GPU I had to to using many other tricks to reduce VRAM usage;

LTX-Video has good ratio: 8x32x32 so the token sequence length is short enough, the reconstruction quality is good if training DiT + VAE all together, but I only have 48GB vram i can only train one model at once;

LeanVAE is very promising, I'm still testing bigger LeanVAE, if 8x32x32 won't be able to get good quality, i'll go with 4x16x16, which is still better than Wan-video VAE or HunyuanVideo VAE;

Diffusion forcing,, causVid and FramePack is more promising autoregressive video model

Thanks for the dataset links!

You can use continuous tokens with a 1D transformer encoder (change the bottleneck to VAE). Then you're free to use normal diffusion models and losses. You can then set a large channel count in the bottleneck (eg. 128). I've been toying with the idea of setting a single latent token out with a massive channel count (~512) to see if it's possible to get good reconstruction and generation results. Would certainly be very fast.