try multiprocessing

main.py

@@ -65,7 +65,8 @@ import json
 
 from huggingface_hub import hf_hub_download, HfApi
 from numba import cuda
-from multiprocessing import Pool
+from multiprocessing import Pool, Process, Queue
+import torch.multiprocessing as mp
 
 # Inputs ===================================================================================================
 
@@ -1111,7 +1112,7 @@ def run(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remo
     return results
 
 
-def run_train(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
     is_app=True
     images = [img[0] for img in images]
@@ -1127,8 +1128,18 @@ def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=Tru
     
     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)
 
+def run_train(images, train_steps=100, modelId="fine_tuned_pcdms", bg_remove=True, resize_inputs=True):
+    run_train_impl(images, train_steps, modelId, bg_remove, resize_inputs)
 
-def run_inference(images, video_path, frames, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
+    """
+    mp.set_start_method('spawn', force=True)
+    p = mp.Process(target=run_train_impl, args=(images, train_steps, modelId, bg_remove, resize_inputs))
+    p.start()
+    p.join()
+    """
+
+
+def run_inference_impl(images, video_path, frames, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
     finetune=True
     is_app=True
     
@@ -1160,16 +1171,17 @@ def run_inference(images, video_path, frames, train_steps=100, inference_steps=1
     #results = [img_pad(img, img_width, img_height, True) for img in results]
     
     print("Done!")
-    
+
     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):
+def run_inference(images, video_path, frames, train_steps=100, inference_steps=10, fps=12, modelId="fine_tuned_pcdms", img_width=1920, img_height=1080, bg_remove=True, resize_inputs=True):
+    return run_inference_impl(images, video_path, frames, train_steps, inference_steps, fps, modelId, img_width, img_height, bg_remove, resize_inputs)
+
+
+def 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
 
-    
-    
     dwpose, rembg_session, pcdms_model, noise_scheduler, image_encoder_p, image_encoder_g, vae, unet = load_models()
     
     if not os.path.exists(modelId+".pt"):
@@ -1198,6 +1210,9 @@ def run_generate_frame(images, target_poses, train_steps=100, inference_steps=10
 
     return results, getThumbnails(results)
 
+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):
+    return generate_frame(images, target_poses, train_steps, inference_steps, modelId, img_width, img_height, bg_remove, resize_inputs)
+
 
 def run_app(images, video_path, train_steps=100, inference_steps=10, fps=12, bg_remove=False, resize_inputs=True):
     
@@ -1246,33 +1261,9 @@ def interpolate_frames(frame1, frame2, times_to_interp):
 def run_interpolate_frames(frame1, frame2, times_to_interp):
     with Pool() as pool:
         results = pool.starmap(interpolate_frames, [(frame1, frame2, times_to_interp)])
-    print(results[0])
     return results[0]
 
 
-"""
-train_steps = 100
-inference_steps = 10
-fps = 12
-"""
-
-"""
-iface = gr.Interface(
-    fn=run,
-    inputs=[
-        gr.Gallery(type="pil", label="Images of the Character"),
-        gr.Video(label="Motion-Capture Video"),
-        gr.Number(label="Training steps", value=100),
-        gr.Number(label="Inference steps", value=10),
-        gr.Number(label="Output frame rate", value=12),
-        gr.Checkbox(label="Remove background", value=False),
-    ],
-    outputs=[gr.Video(label="Result"), gr.Gallery(type="pil", label="Frames")],
-    title="Keyframes AI",
-    description="Upload images of your character and a motion-capture video to generate an animation of the character.",
-)
-"""
-