api for generating frames from poses

app.py

@@ -1,4 +1,4 @@
-from main import run_app, run_train, run_inference
+from main import run_app, run_train, run_inference, run_generate_frame
 
 import spaces
 from PIL import Image
@@ -11,6 +11,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             char_imgs = gr.Gallery(type="pil", label="Images of the Character")
             mocap = gr.Video(label="Motion-Capture Video")
+            poses = gr.JSON(label="Pose Coordinates")
             tr_steps = gr.Number(label="Training steps", value=10)
             inf_steps = gr.Number(label="Inference steps", value=10)
             fps = gr.Number(label="Output frame rate", value=12)
@@ -21,6 +22,7 @@ with gr.Blocks() as demo:
             img_height = gr.Number(label="Output height", value=1080)
             train_btn = gr.Button(value="Train")
             inference_btn = gr.Button(value="Inference")
+            generate_frame_btn = gr.Button(value="Generate Frame")
             submit_btn = gr.Button(value="Generate")
         with gr.Column():
             animation = gr.Video(label="Result")
@@ -41,6 +43,10 @@ with gr.Blocks() as demo:
         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]
     )
 
+    generate_frame_btn.click(
+        run_generate_frame, inputs=[char_imgs, poses, tr_steps, inf_steps, modelId, img_width, img_height, remove_bg, resize_inputs], outputs=[frames, frames_thumb]
+    )
+
 
 demo.launch(share=True)
 

main.py

@@ -1155,6 +1155,33 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     return out_vid+'.webm', results, getThumbnails(results), target_poses_coords, getThumbnails(orig_frames)
     
 
+def run_generate_frame(images, target_poses, train_steps=100, inference_steps=10, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
+    finetune=True
+    is_app=True
+
+    print(target_poses)
+    
+    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")
+
+    results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
+    #urls = save_temp_imgs(results)
+    
+    # postprocessing
+    results = [removebg(img, rembg_session, True) for img in results]
+    #results = [img_pad(img, img_width, img_height, True) for img in results]
+    
+    print("Done!")
+    
+    return results, getThumbnails(results)
+
+
 def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True):
     
     images = [img[0] for img in images]