ghostv2/utils/inference/inference_arguments.py at main · dimitribarbot/ghostv2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import os
from dataclasses import dataclass
from typing import Optional

from simple_parsing import choice, flag

from utils.embedding_models_arguments import EmbeddingModelsArguments


def make_real_path(relative_path):
    return os.path.join(os.path.dirname(os.path.realpath(__file__)), "../..", relative_path)

@dataclass
class InferenceArguments(EmbeddingModelsArguments):

    """ Data arguments """
    source_file_path: str = make_real_path("./examples/images/base/source1.jpg")
    target_file_path: str = make_real_path("./examples/images/base/target1.jpg")
    output_file_path: str = make_real_path("./examples/results/inference/source1_target1.jpg")

    """ Model params """
    G_path: str = make_real_path("./weights/GhostV2/G_unet_2blocks.safetensors")
    gfpgan_model_path: str = make_real_path("./weights/GFPGAN/GFPGANv1.4.safetensors")
    face_alignment_model_path: str = make_real_path("./weights/FaceAlignment/landmark.safetensors")
    face_parser_model_path: str = make_real_path("./weights/BiSeNet/79999_iter.safetensors")
    retina_face_model_path: str = make_real_path("./weights/RetinaFace/Resnet50_Final.safetensors")
    cvlface_aligner_model_path: str = make_real_path("./weights/CVLFace/cvlface_DFA_mobilenet.safetensors")
    face_embeddings: str = choice("facenet", "arcface", "adaface", "cvl_arcface", "cvl_adaface", "cvl_vit", default="cvl_vit")
    backbone: str = choice("unet", "linknet", "resnet", default="unet")
    num_blocks: int = 2
    align_corners: bool = flag(default=True, negative_prefix="--no-")
    precision: Optional[str] = choice(
        None,
        "64",
        "32",
        "16",
        "bf16",
        "transformer-engine",
        "transformer-engine-float16",
        "16-true",
        "16-mixed",
        "bf16-true",
        "bf16-mixed",
        "32-true",
        "64-true",
        default=None
    )

    """ Run arguments """
    device_id: int = 0
    detection_threshold: float = 0.97
    inpaint_output: bool = flag(default=True, negative_prefix="--no-")
    enhance_output: bool = flag(default=True, negative_prefix="--no-")
    source_face_index: int = 0
    target_face_index: int = 0
    align_mode: str = choice("facexlib", "insightface_v1", "insightface_v2", "mtcnn", "cvlface", default="insightface_v2")
    paste_back_mode: str = choice("facexlib_with_parser", "facexlib_without_parser", "insightface", "ghost", "basic", "none", default="ghost")

    """ Debug arguments """
    debug: bool = flag(default=False, negative_prefix="--no-")
    debug_ghost_landmarks: bool = flag(default=False, negative_prefix="--no-")
    debug_source_face_path: str = make_real_path("./examples/results/inference/source_face.jpg")
    debug_target_face_path: str = make_real_path("./examples/results/inference/target_face.jpg")
    debug_swapped_face_path: str = make_real_path("./examples/results/inference/swapped_face.jpg")
    debug_enhanced_face_path: str = make_real_path("./examples/results/inference/enhanced_face.jpg")