Update preprocess/humanparsing/run_parsing.py

preprocess/humanparsing/run_parsing.py

@@ -10,6 +10,8 @@ import torch
 
 from huggingface_hub import hf_hub_download
 
+
+
 class Parsing:
     def __init__(self, gpu_id: int):
         self.gpu_id = gpu_id
@@ -19,14 +21,13 @@ class Parsing:
         session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
         session_options.add_session_config_entry('gpu_id', str(gpu_id))
 
-        print('start download')
-        hf_hub_download(repo_id="levihsu/OOTDiffusion", 
-                        filename="checkpoints/humanparsing/parsing_atr.onnx", 
-                        local_dir=os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing'))
-        hf_hub_download(repo_id="levihsu/OOTDiffusion", 
-                        filename="checkpoints/humanparsing/parsing_lip.onnx", 
-                        local_dir=os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing'))
-        print('finish download')
+        parsing_ckpt_path = os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing')
+        atr_model_path = 'https://huggingface.co/levihsu/OOTDiffusion/blob/main/checkpoints/humanparsing/parsing_atr.onnx'
+        lip_model_path = 'https://huggingface.co/levihsu/OOTDiffusion/blob/main/checkpoints/humanparsing/parsing_lip.onnx'
+        
+        from basicsr.utils.download_util import load_file_from_url
+        load_file_from_url(atr_model_path, model_dir=parsing_ckpt_path)
+        load_file_from_url(lip_model_path, model_dir=parsing_ckpt_path)
         
         self.session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing/parsing_atr.onnx'),
                                             sess_options=session_options, providers=['CPUExecutionProvider'])