fix generate_frames target poses

main.py

@@ -288,7 +288,7 @@ def prepare_inputs_train(images, bg_remove, dwpose, rembg_session):
     return in_img, in_pose, train_imgs, train_poses
         
 
-def prepare_inputs_inference(in_img, in_vid, frames, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app=False):
+def prepare_inputs_inference(in_img, in_vid, frames, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app=False, target_poses=None):
     progress=gr.Progress(track_tqdm=True)
 
     print("prepare_inputs_inference")
@@ -310,7 +310,8 @@ def prepare_inputs_inference(in_img, in_vid, frames, fps, dwpose, rembg_session,
     print("vid: ", in_vid, fps)
 
     progress_bar = tqdm(range(len(frames)), initial=0, desc="Frames")
-    target_poses = []
+    if not target_poses:
+        target_poses = []
     target_poses_coords = []
     max_left = max_top = 999999
     max_right = max_bottom = 0
@@ -333,6 +334,7 @@ def prepare_inputs_inference(in_img, in_vid, frames, fps, dwpose, rembg_session,
         target_poses_coords.append(json.dumps(coords))
         progress_bar.update(1)
         
+    for tpose in target_poses:
         bbox = tpose.getbbox()
         left, top, right, bottom = bbox
         max_left = min(max_left, left)
@@ -498,7 +500,6 @@ def train(modelId, in_image, in_pose, train_images, train_poses, train_steps, pc
     logging_dir = 'outputs/logging'
     print('start train')
     
-    
     progress=gr.Progress(track_tqdm=True)
 
     accelerator = Accelerator(
@@ -1159,7 +1160,8 @@ def run_inference_impl(images, video_path, frames, train_steps=100, inference_st
         frames = [img[0] for img in frames]
 
     in_img, target_poses, in_pose, target_poses_coords, orig_frames = prepare_inputs_inference(in_img, video_path, frames, fps, dwpose, rembg_session, bg_remove, resize_inputs, is_app)
-    
+    target_poses[0].save('inf_pose.png')
+
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
     #urls = save_temp_imgs(results)
 
@@ -1189,19 +1191,20 @@ def generate_frame(images, target_poses, train_steps=100, inference_steps=10, mo
     is_app=True
 
     dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
-    
+
     if not os.path.exists(modelId+".pt"):
         run_train(images, train_steps, modelId, bg_remove, resize_inputs)
-    
+
     images = [img[0] for img in images]
     in_img = images[0]
     in_pose, _ = get_pose(in_img, dwpose, "in_pose.png")
 
     print(target_poses)
     target_poses = json.loads(target_poses)
-    width, height = in_img.size
-    target_poses = [Image.fromarray(draw_openpose(pose, height=height, width=width, include_hands=True, include_face=False)) for pose in target_poses]
-    target_poses[0].save('pose.png')
+    target_poses = [Image.fromarray(draw_openpose(pose, height=img_height, width=img_width, include_hands=True, include_face=False)) for pose in target_poses]
+
+    in_img, target_poses, in_pose, target_poses_coords, orig_frames = prepare_inputs_inference(in_img, None, [], 12, dwpose, rembg_session, bg_remove, resize_inputs, is_app, target_poses)
+    target_poses[0].save('gen_pose.png')
 
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
     #urls = save_temp_imgs(results)