fix vae nan bug

app.py

@@ -276,7 +276,7 @@ def get_ref_anno(ref):
             None,
         )
     missing_keys, extra_keys = autoencoder.load_state_dict(vae_state_dict, strict=False)
-    
+
     img = ref["composite"][..., :3]
     img = cv2.resize(img, opts.image_size, interpolation=cv2.INTER_AREA)
     keypts = np.zeros((42, 2))
@@ -566,7 +566,7 @@ def sample_diff(ref_cond, target_cond, target_keypts, num_gen, seed, cfg):
     print(f"results[0].max(): {results[0].max()}")
     return results, results_pose
 
-@spaces.GPU(duration=120)
+# @spaces.GPU(duration=120)
 def ready_sample(img_ori, inpaint_mask, keypts):
     img = cv2.resize(img_ori[..., :3], opts.image_size, interpolation=cv2.INTER_AREA)
     sam_predictor.set_image(img)
@@ -638,7 +638,7 @@ def ready_sample(img_ori, inpaint_mask, keypts):
             inpaint_mask, dsize=opts.latent_size, interpolation=cv2.INTER_NEAREST
         ),
         dtype=torch.float,
-        device=device,
+        # device=device,
     ).unsqueeze(0)[None, ...]
 
     def make_ref_cond(
@@ -656,7 +656,7 @@ def ready_sample(img_ori, inpaint_mask, keypts):
                 Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
             ]
         )
-        image = image_transform(img).to(device)
+        image = image_transform(img)
         kpts_valid = check_keypoints_validity(keypts, target_size)
         heatmaps = torch.tensor(
             keypoint_heatmap(
@@ -664,7 +664,7 @@ def ready_sample(img_ori, inpaint_mask, keypts):
             )
             * kpts_valid[:, None, None],
             dtype=torch.float,
-            device=device,
+            # device=device,
         )[None, ...]
         mask = torch.tensor(
             cv2.resize(
@@ -673,7 +673,7 @@ def ready_sample(img_ori, inpaint_mask, keypts):
                 interpolation=cv2.INTER_NEAREST,
             ),
             dtype=torch.float,
-            device=device,
+            # device=device,
         ).unsqueeze(0)[None, ...]
         return image[None, ...], heatmaps, mask
 
@@ -744,7 +744,7 @@ def sample_inpaint(
     target_cond_N = target_cond.repeat(N, 1, 1, 1)
     ref_cond_N = ref_cond.repeat(N, 1, 1, 1)
     # novel view synthesis mode = off
-    nvs = torch.zeros(N, dtype=torch.int, device=device)
+    nvs = torch.zeros(N, dtype=torch.int, device=z.device)
     z = torch.cat([z, z], 0)
     model_kwargs = dict(
         target_cond=torch.cat([target_cond_N, torch.zeros_like(target_cond_N)]),
@@ -762,7 +762,7 @@ def sample_inpaint(
         clip_denoised=False,
         model_kwargs=model_kwargs,
         progress=True,
-        device=device,
+        device=z.device,
         jump_length=jump_length,
         jump_n_sample=jump_n_sample,
     ).chunk(2)
@@ -1078,7 +1078,7 @@ with gr.Blocks(css=custom_css) as demo:
                 )
                 run = gr.Button(value="Run", interactive=False)
                 gr.Markdown(
-                    """<p style="text-align: center;">~20s per generation. <br>(For example, if you set Number of generations as 2, it would take around 40s)</p>"""
+                    """<p style="text-align: center;">~20s per generation with RTX3090. ~50s with A100. <br>(For example, if you set Number of generations as 2, it would take around 40s)</p>"""
                 )
                 results = gr.Gallery(
                     type="numpy",