update code

app.py

@@ -473,12 +473,12 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
     exp_img_base_dir = os.path.join(target_path, 'images512x512')
     motion_base_dir = os.path.join(target_path, 'motions')
     label_file_test = os.path.join(target_path, 'images512x512/dataset_realcam.json')
-    render_model.to(device)
-    image_encoder.to(device)
-    vae_triplane.to(device)
-    dinov2.to(device)
-    ws_avg.to(device)
-    DiT_model.to(device)
+    # render_model.to(device)
+    # image_encoder.to(device)
+    # vae_triplane.to(device)
+    # dinov2.to(device)
+    # ws_avg.to(device)
+    # DiT_model.to(device)
     # Set up face verse for amimation
     base_coff = np.load(
         'pretrained_model/temp.npy').astype(
@@ -526,7 +526,7 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
         samples = rearrange(samples, "b c f h w -> b f c h w")
         samples = samples * std + mean
         torch.cuda.empty_cache()
-
+        torch.cuda.ipc_collect()
         save_frames_path_out = os.path.join(save_path_base, image_name, 'out')
         save_frames_path_outshow = os.path.join(save_path_base, image_name, 'out_show')
         save_frames_path_depth = os.path.join(save_path_base, image_name, 'depth')
@@ -580,22 +580,22 @@ def avatar_generation(items, save_path_base, video_path_input, source_type, is_s
             # Load motion data
             motion = torch.tensor(np.load(motion_each_dir)).float().unsqueeze(0).to(device)
 
-            img_ref_double = duplicate_batch(img_ref, batch_size=2)
-            motion_app_double = duplicate_batch(motion_app, batch_size=2)
-            motion_double = duplicate_batch(motion, batch_size=2)
-            pose_double = torch.cat([pose_show, pose], dim=0)
-            exp_target_double = duplicate_batch(exp_target, batch_size=2)
-            samples_double = duplicate_batch(samples, batch_size=2)
+            # img_ref_double = duplicate_batch(img_ref, batch_size=2)
+            # motion_app_double = duplicate_batch(motion_app, batch_size=2)
+            # motion_double = duplicate_batch(motion, batch_size=2)
+            # pose_double = torch.cat([pose_show, pose], dim=0)
+            # exp_target_double = duplicate_batch(exp_target, batch_size=2)
+            # samples_double = duplicate_batch(samples, batch_size=2)
             # Select refine_net processing method
             final_out = render_model(
-                img_ref_double, None, motion_app_double, motion_double, c=pose_double, mesh=exp_target_double,
-                triplane_recon=samples_double,
+                img_ref, None, motion_app, motion, c=pose, mesh=exp_target,
+                triplane_recon=samples,
                 ws_avg=ws_avg, motion_scale=1.
             )
 
             # Process output image
             final_out_show = trans(final_out['image_sr'][0].unsqueeze(0))
-            final_out_notshow = trans(final_out['image_sr'][1].unsqueeze(0))
+            final_out_notshow = trans(final_out['image_sr'][0].unsqueeze(0))
             depth = final_out['image_depth'][0].unsqueeze(0)
             depth = -depth
             depth = (depth - depth.min()) / (depth.max() - depth.min()) * 2 - 1