refactor download models script

downloading_weights.py

@@ -26,8 +26,15 @@ def download_models(
       os.makedirs(dir, exist_ok=True)
 
     for url, path in tqdm(zip(urls, paths)):
-        local_file_path = os.path.join("pretrained_weights", path)
-        wget.download(url, local_file_path)
+        filename = os.path.basename(url)
+        if filename == "yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth":
+            filename = "yolox_l_8x8_300e_coco.pth"
+
+        full_file_path = os.path.join(model_dir, path, filename)
+
+        if not os.path.exists(full_file_path):
+            print(f"Model '{filename}' does not exists. Downloading to '{full_file_path}'..")
+            wget.download(url, full_file_path)
 
     config_urls = ['https://huggingface.co/lambdalabs/sd-image-variations-diffusers/resolve/main/unet/config.json',
                'https://huggingface.co/lambdalabs/sd-image-variations-diffusers/resolve/main/image_encoder/config.json',
@@ -37,10 +44,8 @@ def download_models(
 
     # saving config files
     for url, path in tqdm(zip(config_urls, config_paths)):
-        local_file_path = os.path.join("pretrained_weights", path)
-        wget.download(url, local_file_path)
-
-    # renaming model name as given in readme
-    wrong_file_path = os.path.join("pretrained_weights", "dwpose", "yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth")
-    correct_file_path = os.path.join("pretrained_weights", "dwpose", "yolox_l_8x8_300e_coco.pth")
-    os.rename(wrong_file_path, correct_file_path)
+        filename = os.path.basename(url)
+        full_file_path = os.path.join(model_dir, path, filename)
+        if not os.path.exists(full_file_path):
+            print(f"Model '{filename}' does not exists. Downloading to '{full_file_path}'..")
+            wget.download(url, full_file_path)

musepose_inference.py

@@ -18,6 +18,7 @@ from musepose.models.unet_2d_condition import UNet2DConditionModel
 from musepose.models.unet_3d import UNet3DConditionModel
 from musepose.pipelines.pipeline_pose2vid_long import Pose2VideoPipeline
 from musepose.utils.util import get_fps, read_frames, save_videos_grid
+from downloading_weights import download_models
 
 
 class MusePoseInference:
@@ -42,6 +43,7 @@ class MusePoseInference:
         self.pose_guider = None
         self.image_enc = None
         self.pipe = None
+        self.model_dir = model_dir
         self.output_dir = output_dir
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
@@ -62,6 +64,7 @@ class MusePoseInference:
         fps: int,
         skip: int
     ):
+        download_models(model_dir=self.model_dir)
         print(f"Model Paths: {self.musepose_model_paths}\n{self.image_gen_model_paths}\n{self.inference_config_path}")
         print(f"Input Image Path: {ref_image_path}")
         print(f"Pose Video Path: {pose_video_path}")

pose_align.py

@@ -1,5 +1,4 @@
 import numpy as np
-import argparse
 import torch
 import copy
 import cv2
@@ -7,10 +6,10 @@ import os
 import moviepy.video.io.ImageSequenceClip
 from datetime import datetime
 import gc
-from huggingface_hub import hf_hub_download
 
 from pose.script.dwpose import DWposeDetector, draw_pose
 from pose.script.util import size_calculate, warpAffine_kps
+from downloading_weights import download_models
 
 
 '''
@@ -32,6 +31,7 @@ class PoseAlignmentInference:
             "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.model_dir = model_dir
         self.output_dir = output_dir
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
@@ -45,7 +45,7 @@ class PoseAlignmentInference:
         align_frame: int,
         max_frame: int,
     ):
-        self.download_models()
+        download_models(model_dir=self.model_dir)
         dt_file_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
         outfn=os.path.abspath(os.path.join(self.output_dir, f'{dt_file_name}_demo.mp4'))
         outfn_align_pose_video=os.path.abspath(os.path.join(self.output_dir, f'{dt_file_name}.mp4'))