refactor output dir and models dir

app.py

@@ -1,15 +1,23 @@
 import gradio as gr
+import argparse
 import os
-from huggingface_hub import hf_hub_download
 
 from musepose_inference import MusePoseInference
 from pose_align import PoseAlignmentInference
+from downloading_weights import download_models
 
 
 class App:
-    def __init__(self):
-        self.pose_alignment_infer = PoseAlignmentInference()
-        self.musepose_infer = MusePoseInference()
+    def __init__(self, args):
+        self.pose_alignment_infer = PoseAlignmentInference(
+            model_dir=args.model_dir,
+            output_dir=args.output_dir
+        )
+        self.musepose_infer = MusePoseInference(
+            model_dir=args.model_dir,
+            output_dir=args.output_dir
+        )
+        download_models(args.model_dir)
 
     def musepose_demo(self):
         with gr.Blocks() as demo:
@@ -82,5 +90,10 @@ class App:
 
 
 if __name__ == "__main__":
-    app = App()
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_dir', type=str, default=os.path.join("pretrained_weights"), help='Pretrained models directory for MusePose')
+    parser.add_argument('--output_dir', type=str, default=os.path.join("assets", "videos"), help='Output directory for the result')
+    args = parser.parse_args()
+
+    app = App(args=args)
     app.launch()

downloading_weights.py

@@ -4,7 +4,7 @@ from tqdm import tqdm
 
 
 def download_models(
-    models_dir:str = os.makedirs('pretrained_weights', exist_ok=True)
+    models_dir: str = os.makedirs('pretrained_weights', exist_ok=True)
 ):
     os.makedirs(models_dir, exist_ok=True)
 

musepose_inference.py

@@ -21,17 +21,19 @@ from musepose.utils.util import get_fps, read_frames, save_videos_grid
 
 
 class MusePoseInference:
-    def __init__(self):
+    def __init__(self,
+                 model_dir,
+                 output_dir):
         self.image_gen_model_paths = {
-            "pretrained_base_model": os.path.join("pretrained_weights", "sd-image-variations-diffusers"),
-            "pretrained_vae": os.path.join("pretrained_weights", "sd-vae-ft-mse"),
-            "image_encoder": os.path.join("pretrained_weights", "image_encoder"),
+            "pretrained_base_model": os.path.join(model_dir, "sd-image-variations-diffusers"),
+            "pretrained_vae": os.path.join(model_dir, "sd-vae-ft-mse"),
+            "image_encoder": os.path.join(model_dir, "image_encoder"),
         }
         self.musepose_model_paths = {
-            "denoising_unet": os.path.join("pretrained_weights", "MusePose", "denoising_unet.pth"),
-            "reference_unet": os.path.join("pretrained_weights", "MusePose", "reference_unet.pth"),
-            "pose_guider": os.path.join("pretrained_weights", "MusePose", "pose_guider.pth"),
-            "motion_module": os.path.join("pretrained_weights", "MusePose", "motion_module.pth"),
+            "denoising_unet": os.path.join(model_dir, "MusePose", "denoising_unet.pth"),
+            "reference_unet": os.path.join(model_dir, "MusePose", "reference_unet.pth"),
+            "pose_guider": os.path.join(model_dir, "MusePose", "pose_guider.pth"),
+            "motion_module": os.path.join(model_dir, "MusePose", "motion_module.pth"),
         }
         self.inference_config_path = os.path.join("configs", "inference_v2.yaml")
         self.vae = None
@@ -40,7 +42,7 @@ class MusePoseInference:
         self.pose_guider = None
         self.image_enc = None
         self.pipe = None
-        self.output_dir = os.path.join("assets", "videos")
+        self.output_dir = output_dir
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
 

pose_align.py

@@ -20,17 +20,19 @@ from pose.script.util import size_calculate, warpAffine_kps
     scales: scale parameters
 '''
 class PoseAlignmentInference:
-    def __init__(self):
+    def __init__(self,
+                 model_dir,
+                 output_dir):
         self.detector = None
         self.model_paths = {
-            "det_ckpt": os.path.join("pretrained_weights", "dwpose", "yolox_l_8x8_300e_coco.pth"),
-            "pose_ckpt": os.path.join("pretrained_weights", "dwpose", "dw-ll_ucoco_384.pth")
+            "det_ckpt": os.path.join(model_dir, "dwpose", "yolox_l_8x8_300e_coco.pth"),
+            "pose_ckpt": os.path.join(model_dir, "dwpose", "dw-ll_ucoco_384.pth")
         }
         self.config_paths = {
             "pose_config": os.path.join("pose", "config", "dwpose-l_384x288.py"),
             "det_config": os.path.join("pose", "config", "yolox_l_8xb8-300e_coco.py"),
         }
-        self.output_dir = os.path.join("assets", "videos")
+        self.output_dir = output_dir
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)