forked from s0md3v/sd-webui-roop
-
Notifications
You must be signed in to change notification settings - Fork 7
/
nodes.py
114 lines (98 loc) · 4.35 KB
/
nodes.py
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import os
from modules.processing import StableDiffusionProcessingImg2Img
from scripts.faceswap import FaceSwapScript, get_models
from utils import batch_tensor_to_pil, batched_pil_to_tensor, tensor_to_pil
from logging_patch import apply_logging_patch
def model_names():
models = get_models()
return {os.path.basename(x): x for x in models}
ORDERINGS = ["left to right", "up to down", "largest to smallest"]
DEFAULT_ORDERING = ORDERINGS[0]
class roop:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"reference_image": ("IMAGE",),
"swap_model": (list(model_names().keys()),),
# Comma separated face number(s)
"faces_index": ("STRING", {"default": "0"}),
"reference_faces_index": ("STRING", {"default": "0"}),
# Allow user to change the logging amount, going from minimal to verbose
"console_logging_level": ([0, 1, 2],),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "roop"
def execute(self, image, reference_image, swap_model, faces_index, reference_faces_index, console_logging_level):
apply_logging_patch(console_logging_level)
script = FaceSwapScript()
pil_images = batch_tensor_to_pil(image)
source = tensor_to_pil(reference_image)
p = StableDiffusionProcessingImg2Img(pil_images)
face_order = DEFAULT_ORDERING
reverse_order = False
reference_order = DEFAULT_ORDERING
reverse_reference_order = False
script.process(
p=p, img=source, enable=True, faces_index=faces_index,
reference_faces_index=reference_faces_index,
face_order=face_order, reverse_order=reverse_order,
reference_order=reference_order, reverse_reference_order=reverse_reference_order,
model=swap_model,
face_restorer_name=None, face_restorer_visibility=None,
upscaler_name=None, upscaler_scale=None, upscaler_visibility=None,
swap_in_source=True, swap_in_generated=True
)
result = batched_pil_to_tensor(p.init_images)
return (result,)
class RoopImproved:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image": ("IMAGE",),
"reference_image": ("IMAGE",),
"swap_model": (list(model_names().keys()),),
# Comma separated face number(s)
"faces_index": ("STRING", {"default": "0"}),
"reference_faces_index": ("STRING", {"default": "0"}),
"face_order": (ORDERINGS, {"default": DEFAULT_ORDERING}),
"reverse_order": ("BOOLEAN", {"default": False}),
"reference_order": (ORDERINGS, {"default": DEFAULT_ORDERING}),
"reverse_reference_order": ("BOOLEAN", {"default": False}),
# Allow user to change the logging amount, going from minimal to verbose
"console_logging_level": ([0, 1, 2],),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "execute"
CATEGORY = "roop"
def execute(self, image, reference_image, swap_model, faces_index, reference_faces_index, face_order, reverse_order, reference_order, reverse_reference_order, console_logging_level):
apply_logging_patch(console_logging_level)
script = FaceSwapScript()
pil_images = batch_tensor_to_pil(image)
source = tensor_to_pil(reference_image)
p = StableDiffusionProcessingImg2Img(pil_images)
script.process(
p=p, img=source, enable=True, faces_index=faces_index,
reference_faces_index=reference_faces_index,
face_order=face_order, reverse_order=reverse_order,
reference_order=reference_order, reverse_reference_order=reverse_reference_order,
model=swap_model,
face_restorer_name=None, face_restorer_visibility=None,
upscaler_name=None, upscaler_scale=None, upscaler_visibility=None,
swap_in_source=True, swap_in_generated=True
)
result = batched_pil_to_tensor(p.init_images)
return (result,)
NODE_CLASS_MAPPINGS = {
"roop": roop,
"RoopImproved": RoopImproved,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"roop": "roop",
"RoopImproved": "Roop (Improved)",
}