extract_vae_latent.sh CUDA out of memory #167
Description
我运行以下脚本来extract vae latent，分辨率是768p（即WIDTH=1280，HEIGHT=768），NUM_FRAMES降到17了还是CUDA out of memory，这要怎么解决呢？
用的是8卡H800，80G显存。
#!/bin/bash
# This script is used for batch extract the vae latents for video generation training
# Since the video latent extract is very slow, pre-extract the video vae latents will save the training time
GPUS=8 # The gpu number
MODEL_NAME=pyramid_flux # The model name, `pyramid_flux` or `pyramid_mmdit`
VAE_MODEL_PATH=pretrained_weights/pyramid-flow-miniflux/causal_video_vae # The VAE CKPT dir.
ANNO_FILE=path # The video annotation file path
WIDTH=1280
HEIGHT=768
NUM_FRAMES=17
torchrun --nproc_per_node $GPUS \
tools/extract_video_vae_latents.py \
--batch_size 1 \
--model_dtype bf16 \
--model_path $VAE_MODEL_PATH \
--anno_file $ANNO_FILE \
--width $WIDTH \
--height $HEIGHT \
--num_frames $NUM_FRAMES