fix vae nan bug

app.py

@@ -228,29 +228,32 @@ if NEW_MODEL:
     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'
-    sd_path = './sd-v1-4.ckpt'
-    print('Load diffusion model...')
-    diffusion = create_diffusion(str(opts.test_sampling_steps))
-    model = vit.DiT_XL_2(
-        input_size=opts.latent_size[0],
-        latent_dim=opts.latent_dim,
-        in_channels=opts.latent_dim+opts.n_keypoints+opts.n_mask,
-        learn_sigma=True,
-    ).to(device)
-    ckpt_state_dict = torch.load(model_path)['state_dict']
-    dit_state_dict = {remove_prefix(k, 'diffusion_backbone.'): v for k, v in ckpt_state_dict.items() if k.startswith('diffusion_backbone')}
-    vae_state_dict = {remove_prefix(k, 'autoencoder.'): v for k, v in ckpt_state_dict.items() if k.startswith('autoencoder')}
-    missing_keys, extra_keys = model.load_state_dict(dit_state_dict, strict=False)
-    model.eval()
-    assert len(missing_keys) == 0 and len(extra_keys) == 0
-    autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).to(device)
-    missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
-    autoencoder.eval()
-    assert len(missing_keys) == 0 and len(extra_keys) == 0
+# else:
+#     opts = HandDiffOpts()
+#     model_path = './finetune_epoch=5-step=130000.ckpt'
+#     sd_path = './sd-v1-4.ckpt'
+#     print('Load diffusion model...')
+#     diffusion = create_diffusion(str(opts.test_sampling_steps))
+#     model = vit.DiT_XL_2(
+#         input_size=opts.latent_size[0],
+#         latent_dim=opts.latent_dim,
+#         in_channels=opts.latent_dim+opts.n_keypoints+opts.n_mask,
+#         learn_sigma=True,
+#     ).to(device)
+#     ckpt_state_dict = torch.load(model_path)['state_dict']
+#     dit_state_dict = {remove_prefix(k, 'diffusion_backbone.'): v for k, v in ckpt_state_dict.items() if k.startswith('diffusion_backbone')}
+#     vae_state_dict = {remove_prefix(k, 'autoencoder.'): v for k, v in ckpt_state_dict.items() if k.startswith('autoencoder')}
+#     missing_keys, extra_keys = model.load_state_dict(dit_state_dict, strict=False)
+#     model.eval()
+#     assert len(missing_keys) == 0 and len(extra_keys) == 0
+#     autoencoder = vqvae.create_model(3, 3, opts.latent_dim).eval().requires_grad_(False).to(device)
+#     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
+#     autoencoder.eval()
+#     assert len(missing_keys) == 0 and len(extra_keys) == 0
 sam_path = hf_hub_download(repo_id="Chaerin5/FoundHand-weights", filename="sam_vit_h_4b8939.pth", token=token)
 sam_predictor = init_sam(ckpt_path=sam_path, device='cpu')
 
@@ -492,6 +495,7 @@ def get_ref_anno(ref):
     print(f"opts.latent_scaling_factor: {opts.latent_scaling_factor}")
     print(f"autoencoder encoder before operating max: {min([p.min() for p in autoencoder.encoder.parameters()])}")
     print(f"autoencoder encoder before operating min: {max([p.max() for p in autoencoder.encoder.parameters()])}")
+    print(f"autoencoder encoder before operating dtype: {next(autoencoder.encoder.parameters()).dtype}")
     latent = opts.latent_scaling_factor * autoencoder.encode(image).sample()
     print(f"latent.max(): {latent.max()}, latent.min(): {latent.min()}")
     if not REF_POSE_MASK: