XavierXiao / Dreambooth-Stable-Diffusion

Implementation of Dreambooth (https://arxiv.org/abs/2208.12242) with Stable Diffusion
MIT License
7.57k stars 790 forks source link

Size of the trained checkpoint (ckpt) file #160

Open Sunset-Q opened 8 months ago

Sunset-Q commented 8 months ago

I obtained a ckpt file from training the model using the command 'python main.py,' and the individual file size is 12GB. However, I noticed in the Custom Diffusion paper that the ckpt file for Dreambooth is only 3GB. Could you please help me understand if I made a mistake somewhere? The sd-v1-4.ckpt file I downloaded from Hugging Face has a size of 4GB, and I used the command 'python main.py -n dog --data_root /root/autodl-tmp/data/dog --reg_data_root /root/autodl-tmp/real_reg/samples_dog/dog --class_word dog.' Is it because the sd-v1-4.ckpt file I downloaded itself is large?

victorchall commented 8 months ago

The ckpt contains optimizer state in addition to the Unet, CLIP, and VAE weights in FP32 precision.

The optimizer state can be removed to produce a 4GB file.