feat: add download_models script

2024-05-29 16:43:38 +08:00
parent 4afb18ae21
commit 983128f0d3
4 changed files with 182 additions and 13 deletions
@@ -15,7 +15,7 @@ pip install -r requirements.txt

 ### Download model

-All models is same as [facefusion](https://github.com/facefusion/facefusion) which can be found in [facefusion assets](https://github.com/facefusion/facefusion-assets).
+All models is same as facefusion which can be found in [facefusion assets](https://github.com/facefusion/facefusion-assets).

 Now it will use the following model by default. In the future, dynamic model selection will be supported. Please download the model based on the following directory structure.

@@ -29,13 +29,24 @@ Put all into directory `CmofyUI/models/faceless`.
 |   `-- face_landmarker_68_5.onnx
 |-- face_recognizer
 |   `-- arcface_w600k_r50.onnx
+├── face_restoration
+│   └── gfpgan_1.4.onnx
 |-- face_swapper
-|   |-- inswapper_128.onnx
-|   `-- inswapper_128_fp16.onnx
+|   └── inswapper_128.onnx
 |-- gender_age.onnx
 `-- rmbg.pth
 ```

+You can download models by running `download_models.py` script under directory `CmofyUI/custom_nodes/comfyui-faceless-node`.
+
+```bash
+# Install required models
+python download_models.py
+
+# Install all models
+python download_models.py --all
+```
+
 ## Example workflows

 ### Image face swap
@@ -52,6 +63,6 @@ Put all into directory `CmofyUI/models/faceless`.

 It doesn't support preview video for now. All video will be save in directory `output/faceless`.

-## Thanks and Credits
+## Thanks

 Thanks to [Facefusion](https://github.com/facefusion/facefusion). This project is based on facefusion to implenment a special version for ComfyUI.
@@ -0,0 +1,157 @@
+import os
+import argparse
+from urllib.request import urlretrieve
+
+parser = argparse.ArgumentParser(description="Download faceless models")
+
+parser.add_argument('--all', action='store_true', help='Download all models, or just download required models', default=False)
+
+args = parser.parse_args()
+
+base_path = os.path.dirname(os.path.realpath(__file__))
+models_dir = os.path.normpath(os.path.join(base_path, "../../models/faceless"))
+
+# Define the directory structure
+model_sets = {
+    'face_detector': [
+        {
+            'name': 'yoloface_8n.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/yoloface_8n.onnx',
+            'required': True
+        }
+    ],
+    'face_landmarker': [
+        {
+            'name': '2dfan4.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/2dfan4.onnx',
+            'required': True
+        },
+        {
+            'name': 'face_landmarker_68_5.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/face_landmarker_68_5.onnx',
+            'required': True
+        }
+    ],
+    'face_recognizer': [
+        {
+            'name': 'arcface_w600k_r50.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/arcface_w600k_r50.onnx',
+            'required': True
+        }
+    ],
+    'face_swapper': [
+        {
+            'name': 'inswapper_128.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx',
+            'required': True
+        },
+        {
+            'name': 'inswapper_128_fp16.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128_fp16.onnx',
+            'required': False
+        },
+        {
+            'name': 'simswap_256.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_256.onnx',
+            'required': False
+        },
+        {
+            'name': 'simswap_512_unofficial.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/simswap_512_unofficial.onnx',
+            'required': False
+        },
+        {
+            'name': 'blendswap_256.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/blendswap_256.onnx',
+            'required': False
+        },
+        {
+            'name': 'uniface_256.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/uniface_256.onnx',
+            'required': False
+        }
+    ],
+    'face_restoration': [
+        {
+            'name': 'gfpgan_1.2.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.2.onnx',
+            'required': False
+        },
+        {
+            'name': 'gfpgan_1.3.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.3.onnx',
+            'required': False
+        },
+        {
+            'name': 'gfpgan_1.4.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.4.onnx',
+            'required': True
+        },
+        {
+            'name': 'codeformer.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/codeformer.onnx',
+            'required': False
+        },
+        {
+            'name': 'gpen_bfr_256.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_256.onnx',
+            'required': False
+        },
+        {
+            'name': 'gpen_bfr_512.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_512.onnx',
+            'required': False
+        },
+        {
+            'name': 'gpen_bfr_1024.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_1024.onnx',
+            'required': False
+        },
+        {
+            'name': 'gpen_bfr_2048.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_2048.onnx',
+            'required': False
+        },
+        {
+            'name': 'restoreformer_plus_plus.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/restoreformer_plus_plus.onnx',
+            'required': False
+        }
+    ],
+    '.': [
+        {
+            'name': 'gender_age.onnx',
+            'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gender_age.onnx',
+            'required': True
+        },
+        {
+            'name': 'rmbg.pth',
+            'url': 'https://huggingface.co/briaai/RMBG-1.4/resolve/main/model.pth?download=true',
+            'required': True
+        }
+    ]
+}
+
+# Create the directories they don't exist
+for model_type in model_sets:
+    full_dir = os.path.join(models_dir, model_type)
+    if not os.path.exists(full_dir):
+        os.makedirs(full_dir)
+
+# Download and save the models
+for model_type, models in model_sets.items():
+    for model in models:
+        name = model['name']
+        url = model['url']
+        required = model['required']
+
+        if not args.all and not required:
+            continue
+
+        file_path = os.path.join(models_dir, model_type, name)
+        if os.path.exists(file_path):
+            continue
+
+        print(f'Start downloading `{name}`')
+        urlretrieve(url, file_path)
+        print(f'Downloaded `{name}`')
@@ -32,6 +32,7 @@ class NodesRemoveBackground:
        return True

    def remove_images_background(self, images):
+        # Laod model first
        self.load_model()

        processed_images = []
@@ -36,63 +36,63 @@ MODELS : ModelSet =\
    'codeformer':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/codeformer.onnx',
-        'path': resolve_relative_path('../../../models/faceless/codeformer.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/codeformer.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    },
    'gfpgan_1.2':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.2.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gfpgan_1.2.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gfpgan_1.2.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    },
    'gfpgan_1.3':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.3.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gfpgan_1.3.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gfpgan_1.3.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    },
    'gfpgan_1.4':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gfpgan_1.4.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gfpgan_1.4.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gfpgan_1.4.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    },
    'gpen_bfr_256':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_256.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gpen_bfr_256.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gpen_bfr_256.onnx'),
        'template': 'arcface_128_v2',
        'size': (256, 256)
    },
    'gpen_bfr_512':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_512.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gpen_bfr_512.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gpen_bfr_512.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    },
    'gpen_bfr_1024':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_1024.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gpen_bfr_1024.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gpen_bfr_1024.onnx'),
        'template': 'ffhq_512',
        'size': (1024, 1024)
    },
    'gpen_bfr_2048':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/gpen_bfr_2048.onnx',
-        'path': resolve_relative_path('../../../models/faceless/gpen_bfr_2048.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/gpen_bfr_2048.onnx'),
        'template': 'ffhq_512',
        'size': (2048, 2048)
    },
    'restoreformer_plus_plus':
    {
        'url': 'https://github.com/facefusion/facefusion-assets/releases/download/models/restoreformer_plus_plus.onnx',
-        'path': resolve_relative_path('../../../models/faceless/restoreformer_plus_plus.onnx'),
+        'path': resolve_relative_path('../../../models/faceless/face_restoration/restoreformer_plus_plus.onnx'),
        'template': 'ffhq_512',
        'size': (512, 512)
    }