fix eval script

evaluate.py

@@ -60,12 +60,11 @@ with open('metrics.json', 'r') as file:
     metrics = json.load(file)
 
 def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False):
-    print(item)
-    
+
     images = []
     for path in image_paths:
         img = Image.open(path)
-        images.append(img)
+        images.append([img])
     
     gt_frames = extract_frames(video_path, fps)
     
@@ -76,13 +75,13 @@ def get_score(item, image_paths, video_path, train_steps=100, inference_steps=10
         frame.save("out/"+item+"/frame_"+str(i)+".png")
     
     #results = run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, finetune=True)
-    results = run_eval(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False)
+    results, results_base = run_eval(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False)
     
+    print(results)
+
     for i, result in enumerate(results):
         result.save("out/"+item+"/result_"+str(i)+".png")
         
-    results_base = run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, finetune=False)
-        
     for i, result in enumerate(results_base):
         result.save("out/"+item+"/base_"+str(i)+".png")
     
@@ -156,23 +155,29 @@ def run_evaluate():
     snapshot_download(repo_id="acmyu/KeyframesAI-eval", local_dir="test", repo_type="dataset")
 
     items = os.listdir('test') 
-    items = ['test/woody'] #['sidewalk', 'aaa', 'azri', 'dead', 'frankgirl', 'kobold', 'ramona', 'renee', 'walk', 'woody']
+    items = [it for it in items if not it[0]=='.' and not os.path.isfile('test/'+it)]
+    print(items)
+    items = ['sidewalk'] #['sidewalk', 'aaa', 'azri', 'dead', 'frankgirl', 'kobold', 'ramona', 'renee', 'walk', 'woody']
     
     for item in items: 
         if item in metrics:
             continue
-
-        name = os.path.basename(os.path.dirname(item))
-        print(name)
+        print(item)
         
-        files = get_files(item)
-        videos = (x for x in files if x.endswith('.mp4'))
-        if len(videos) == 1:
-            get_score(name, list(filter(lambda x: not x.endswith('.mp4'), files)), videos[0])
-            #get_score(item, ['test/'+item+'/1.jpg', 'test/'+item+'/2.jpg', 'test/'+item+'/3.jpg'], 'test/'+item+'/v.mp4')
-        else:
-            print('Error: mp4 not found')
-
+        try:
+            files = get_files('test/'+item)
+            images = list(filter(lambda x: not x.endswith('.mp4'), files))
+            images = ['test/'+item+'/'+img for img in images]
+            videos = [x for x in files if x.endswith('.mp4')]
+            print(images, videos)
+
+            if len(videos) == 1:
+                get_score(item, images, 'test/'+item+'/'+videos[0])
+                #get_score(item, ['test/'+item+'/1.jpg', 'test/'+item+'/2.jpg', 'test/'+item+'/3.jpg'], 'test/'+item+'/v.mp4')
+            else:
+                print('Error: mp4 not found')
+        except:
+            print("Error", item)
 
 
     ssim = []

main.py

@@ -708,7 +708,18 @@ def train(modelId, in_image, in_pose, train_images, train_poses, train_steps, pc
     if not finetune or train_steps == 0:
         accelerator.wait_for_everyone()
         accelerator.end_training()
-        return {k: v.cpu() for k, v in sd_model.state_dict().items()}
+
+        checkpoint_state_dict = {
+            "epoch": 0,
+            "module": {k: v.cpu() for k, v in sd_model.state_dict().items()}, #sd_model.state_dict(),
+        }
+        torch.save(checkpoint_state_dict, modelId+".pt")
+        
+        del sd_model
+        gc.collect()
+        torch.cuda.empty_cache()
+        return
+        #return {k: v.cpu() for k, v in sd_model.state_dict().items()}
 
 
     it = range(starting_epoch, args.num_train_epochs)
@@ -1117,7 +1128,7 @@ def run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remo
     return results
 
 
-def run_train_impl(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
+def run_train_impl(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True, finetune=True):
     finetune=True
     is_app=True
     images = [img[0] for img in images]
@@ -1250,25 +1261,28 @@ def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_
     return out_vid+'.webm', results
 
 
-def run_eval(images, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False):
-    finetune=True
+def run_eval(images_orig, video_path, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=False, resize_inputs=False):
     is_app=False
 
     dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
 
-    run_train_impl(images, train_steps, modelId, bg_remove, resize_inputs)
-    
-    images = [img[0] for img in images]
-    in_img = images[0]
+    images = [img[0] for img in images_orig]
 
-    in_img, target_poses, in_pose, target_poses_coords, orig_frames = prepare_inputs_inference(in_img, video_path, [], fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
+    in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
+    in_img, target_poses, in_pose, _, _ = prepare_inputs_inference(in_img, video_path, [], fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
+    
+    finetune = True
+    train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
+    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
 
-    _, results, _, _, _ = run_inference_impl(images, video_path, frames, train_steps, inference_steps, fps, modelId, img_width, img_height, bg_remove, resize_inputs)
+    finetune = False
+    train(modelId, in_img, in_pose, train_imgs, train_poses, train_steps, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet, finetune, is_app)
+    results_base = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
 
     gc.collect()
     torch.cuda.empty_cache()
 
-    return results
+    return results, results_base