explainingai-code / StableDiffusion-PyTorch

This repo implements a Stable Diffusion model in PyTorch with all the essential components.
139 stars 30 forks source link

Error loading vgg weights #23

Open Recoil77 opened 2 months ago

Recoil77 commented 2 months ago

Hello. Did all as described in manual. Python 3.8 pip install -r requirements.txt all fine Loadaded vgg.pth Prepared mnist dataset.

But while training VQVAE got this error

(.conda) (base) root@gpu:~/code/StableDiffusion-PyTorch# python -m tools.train_vqvae --config config/mnist.yaml

{'dataset_params': {'im_path': 'data/mnist/train/images', 'im_channels': 1, 'im_size': 28, 'name': 'mnist'}, 'diffusion_params': {'num_timesteps': 1000, 'beta_start': 0.0015, 'beta_end': 0.0195}, 'ldm_params': {'down_channels': [128, 256, 256, 256], 'mid_channels': [256, 256], 'down_sample': [False, False, False], 'attn_down': [True, True, True], 'time_emb_dim': 256, 'norm_channels': 32, 'num_heads': 16, 'conv_out_channels': 128, 'num_down_layers': 2, 'num_mid_layers': 2, 'num_up_layers': 2}, 'autoencoder_params': {'z_channels': 3, 'codebook_size': 20, 'down_channels': [32, 64, 128], 'mid_channels': [128, 128], 'down_sample': [True, True], 'attn_down': [False, False], 'norm_channels': 32, 'num_heads': 16, 'num_down_layers': 1, 'num_mid_layers': 1, 'num_up_layers': 1}, 'train_params': {'seed': 1111, 'task_name': 'mnist', 'ldm_batch_size': 64, 'autoencoder_batch_size': 64, 'disc_start': 1000, 'disc_weight': 0.5, 'codebook_weight': 1, 'commitment_beta': 0.2, 'perceptual_weight': 1, 'kl_weight': 5e-06, 'ldm_epochs': 100, 'autoencoder_epochs': 10, 'num_samples': 25, 'num_grid_rows': 5, 'ldm_lr': 1e-05, 'autoencoder_lr': 0.0001, 'autoencoder_acc_steps': 1, 'autoencoder_img_save_steps': 8, 'save_latents': False, 'vae_latent_dir_name': 'vae_latents', 'vqvae_latent_dir_name': 'vqvae_latents', 'ldm_ckpt_name': 'ddpm_ckpt.pth', 'vqvae_autoencoder_ckpt_name': 'vqvae_autoencoder_ckpt.pth', 'vae_autoencoder_ckpt_name': 'vae_autoencoder_ckpt.pth', 'vqvae_discriminator_ckpt_name': 'vqvae_discriminator_ckpt.pth', 'vae_discriminator_ckpt_name': 'vae_discriminator_ckpt.pth'}} 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 11/11 [00:00<00:00, 36.81it/s] Found 60000 images for split train

Loading model from: /root/code/StableDiffusion-PyTorch/models/weights/v0.1/vgg.pth Traceback (most recent call last): File "/root/code/StableDiffusion-PyTorch/.conda/lib/python3.8/runpy.py", line 194, in _run_module_as_main return _run_code(code, main_globals, None, File "/root/code/StableDiffusion-PyTorch/.conda/lib/python3.8/runpy.py", line 87, in _run_code exec(code, run_globals) File "/root/code/StableDiffusion-PyTorch/tools/train_vqvae.py", line 204, in train(args) File "/root/code/StableDiffusion-PyTorch/tools/train_vqvae.py", line 73, in train lpips_model = LPIPS().eval().to(device) File "/root/code/StableDiffusion-PyTorch/models/lpips.py", line 94, in init self.load_state_dict(torch.load(model_path, map_location=device), strict=False) File "/root/code/StableDiffusion-PyTorch/.conda/lib/python3.8/site-packages/torch/serialization.py", line 713, in load return _legacy_load(opened_file, map_location, pickle_module, pickle_load_args) File "/root/code/StableDiffusion-PyTorch/.conda/lib/python3.8/site-packages/torch/serialization.py", line 920, in _legacy_load magic_number = pickle_module.load(f, pickle_load_args) _pickle.UnpicklingError: invalid load key, '\x0a'.

Please help.

soulslicer commented 1 month ago

You have to download it manually and drop it in the folder, wget does not work

explainingai-code commented 1 month ago

Thank you @soulslicer (have now updated the readme to indicate the same) @Recoil77 So sorry for completely missing to respond on this. I know its been long but were you able to get it working ?