output original video frame

app.py

@@ -27,6 +27,7 @@ with gr.Blocks() as demo:
             frames = gr.Gallery(type="pil", label="Frames", format="png")
             frames_thumb = gr.Gallery(type="pil", label="Thumbnails", format="png")
             pose_coords = gr.JSON(label="Pose Coordinates")
+            reference = gr.Gallery(type="pil", label="Reference Images", format="png")
 
     submit_btn.click(
         run_app, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, remove_bg, resize_inputs], outputs=[animation, frames]
@@ -37,7 +38,7 @@ with gr.Blocks() as demo:
     )
     
     inference_btn.click(
-        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, img_width, img_height, remove_bg, resize_inputs], outputs=[animation, frames, frames_thumb, pose_coords]
+        run_inference, inputs=[char_imgs, mocap, tr_steps, inf_steps, fps, modelId, img_width, img_height, remove_bg, resize_inputs], outputs=[animation, frames, frames_thumb, pose_coords, reference]
     )
 
 

main.py

@@ -345,14 +345,14 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remo
             tpose.save("out/"+"tar_pose"+str(len(target_poses_cropped))+".png")
         target_poses_cropped.append(tpose)
         
-    return in_img, target_poses_cropped, in_pose, target_poses_coords
+    return in_img, target_poses_cropped, in_pose, target_poses_coords, frames
 
 
 def prepare_inputs(images, in_vid, fps, bg_remove, dwpose, rembg_session, resize='target', is_app=False):
     
     in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
     
-    in_img, target_poses_cropped, _, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize, is_app)
+    in_img, target_poses_cropped, _, _, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize, is_app)
     
     
     return in_img, in_pose, train_imgs, train_poses, target_poses_cropped
@@ -1138,7 +1138,7 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     images = [img[0] for img in images]
     in_img = images[0]
 
-    in_img, target_poses, in_pose, target_poses_coords = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', is_app)
+    in_img, target_poses, in_pose, target_poses_coords, orig_frames = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', is_app)
     
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
     #urls = save_temp_imgs(results)
@@ -1156,7 +1156,7 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     
     print("Done!")
     
-    return out_vid+'.webm', results, getThumbnails(results), target_poses_coords
+    return out_vid+'.webm', results, getThumbnails(results), target_poses_coords, orig_frames
     
 
 def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True):