fix vae nan bug

app.py

@@ -199,7 +199,8 @@ if NEW_MODEL:
         model_path = "./DINO_EMA_11M_b50_lr1e-5_epoch10_step550k.ckpt"
     else:
         raise ValueError(f"new model epoch should be either 6 or 7, got {MODEL_EPOCH}")
-    vae_path = './vae-ft-mse-840000-ema-pruned.ckpt'
+    # vae_path = './vae-ft-mse-840000-ema-pruned.ckpt'
+    vae_path = hf_hub_download(repo_id="Chaerin5/FoundHand-weights", filename="vae-ft-mse-840000-ema-pruned.ckpt", token=token)
     # sd_path = './sd-v1-4.ckpt'
     print('Load diffusion model...')
     diffusion = create_diffusion(str(opts.test_sampling_steps))
@@ -216,7 +217,7 @@ if NEW_MODEL:
     model.eval()
     print(missing_keys, extra_keys)
     assert len(missing_keys) == 0
-    vae_state_dict = torch.load(vae_path,map_location='cpu')['state_dict']
+    vae_state_dict = torch.load(vae_path, map_location='cpu')['state_dict']
     autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False)
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
     autoencoder = autoencoder.to(device)

vae-ft-mse-840000-ema-pruned.ckpt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:c6a580b13a5bc05a5e16e4dbb80608ff2ec251a162311590c1f34c013d7f3dab
-size 334695179