fix vae nan bug

app.py

@@ -217,7 +217,21 @@ 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']
+    print(f"vae_state_dict encoder dtype: {vae_state_dict['encoder.conv_in.weight'].dtype}")
+    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False)
+    print(f"autoencoder encoder dtype: {next(autoencoder.encoder.parameters()).dtype}")
+    print(f"encoder before load_state_dict parameters min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
+    print(f"encoder before load_state_dict parameters max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
+    missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
+    print(f"encoder after load_state_dict parameters min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
+    print(f"encoder after load_state_dict parameters max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
+    autoencoder = autoencoder.to(device)
+    autoencoder.eval()
+    print(f"encoder after eval() min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
+    print(f"encoder after eval() max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
+    print(f"autoencoder encoder after eval() dtype: {next(autoencoder.encoder.parameters()).dtype}")
+    assert len(missing_keys) == 0
 # else:
 #     opts = HandDiffOpts()
 #     model_path = './finetune_epoch=5-step=130000.ckpt'
@@ -261,24 +275,8 @@ def get_ref_anno(ref):
             None,
             None,
         )
-
-    vae_state_dict = torch.load(vae_path, map_location='cpu')['state_dict']
-    print(f"vae_state_dict encoder dtype: {vae_state_dict['encoder.conv_in.weight'].dtype}")
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False)
-    print(f"autoencoder encoder dtype: {next(autoencoder.encoder.parameters()).dtype}")
-    print(f"encoder before load_state_dict parameters min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
-    print(f"encoder before load_state_dict parameters max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
-    print(f"encoder after load_state_dict parameters min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
-    print(f"encoder after load_state_dict parameters max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
-    autoencoder = autoencoder.to(device)
-    autoencoder.eval()
-    print(f"encoder after eval() min: {min([p.min() for p in autoencoder.encoder.parameters()])}")
-    print(f"encoder after eval() max: {max([p.max() for p in autoencoder.encoder.parameters()])}")
-    print(f"autoencoder encoder after eval() dtype: {next(autoencoder.encoder.parameters()).dtype}")
-    assert len(missing_keys) == 0
     
-
     img = ref["composite"][..., :3]
     img = cv2.resize(img, opts.image_size, interpolation=cv2.INTER_AREA)
     keypts = np.zeros((42, 2))