Spaces:
Running
on
Zero
Running
on
Zero
Rishi Desai
commited on
Commit
·
431ce81
1
Parent(s):
9a054ff
raw workflow in python
Browse files- FaceEnhancementProd.py +299 -0
FaceEnhancementProd.py
ADDED
@@ -0,0 +1,299 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import random
|
3 |
+
import sys
|
4 |
+
from typing import Sequence, Mapping, Any, Union
|
5 |
+
import torch
|
6 |
+
|
7 |
+
|
8 |
+
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
9 |
+
"""Returns the value at the given index of a sequence or mapping.
|
10 |
+
|
11 |
+
If the object is a sequence (like list or string), returns the value at the given index.
|
12 |
+
If the object is a mapping (like a dictionary), returns the value at the index-th key.
|
13 |
+
|
14 |
+
Some return a dictionary, in these cases, we look for the "results" key
|
15 |
+
|
16 |
+
Args:
|
17 |
+
obj (Union[Sequence, Mapping]): The object to retrieve the value from.
|
18 |
+
index (int): The index of the value to retrieve.
|
19 |
+
|
20 |
+
Returns:
|
21 |
+
Any: The value at the given index.
|
22 |
+
|
23 |
+
Raises:
|
24 |
+
IndexError: If the index is out of bounds for the object and the object is not a mapping.
|
25 |
+
"""
|
26 |
+
try:
|
27 |
+
return obj[index]
|
28 |
+
except KeyError:
|
29 |
+
return obj["result"][index]
|
30 |
+
|
31 |
+
|
32 |
+
def find_path(name: str, path: str = None) -> str:
|
33 |
+
"""
|
34 |
+
Recursively looks at parent folders starting from the given path until it finds the given name.
|
35 |
+
Returns the path as a Path object if found, or None otherwise.
|
36 |
+
"""
|
37 |
+
# If no path is given, use the current working directory
|
38 |
+
if path is None:
|
39 |
+
path = os.getcwd()
|
40 |
+
|
41 |
+
# Check if the current directory contains the name
|
42 |
+
if name in os.listdir(path):
|
43 |
+
path_name = os.path.join(path, name)
|
44 |
+
print(f"{name} found: {path_name}")
|
45 |
+
return path_name
|
46 |
+
|
47 |
+
# Get the parent directory
|
48 |
+
parent_directory = os.path.dirname(path)
|
49 |
+
|
50 |
+
# If the parent directory is the same as the current directory, we've reached the root and stop the search
|
51 |
+
if parent_directory == path:
|
52 |
+
return None
|
53 |
+
|
54 |
+
# Recursively call the function with the parent directory
|
55 |
+
return find_path(name, parent_directory)
|
56 |
+
|
57 |
+
|
58 |
+
def add_comfyui_directory_to_sys_path() -> None:
|
59 |
+
"""
|
60 |
+
Add 'ComfyUI' to the sys.path
|
61 |
+
"""
|
62 |
+
comfyui_path = find_path("ComfyUI")
|
63 |
+
if comfyui_path is not None and os.path.isdir(comfyui_path):
|
64 |
+
sys.path.append(comfyui_path)
|
65 |
+
print(f"'{comfyui_path}' added to sys.path")
|
66 |
+
|
67 |
+
|
68 |
+
def add_extra_model_paths() -> None:
|
69 |
+
"""
|
70 |
+
Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
|
71 |
+
"""
|
72 |
+
try:
|
73 |
+
from main import load_extra_path_config
|
74 |
+
except ImportError:
|
75 |
+
print(
|
76 |
+
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
|
77 |
+
)
|
78 |
+
from utils.extra_config import load_extra_path_config
|
79 |
+
|
80 |
+
extra_model_paths = find_path("extra_model_paths.yaml")
|
81 |
+
|
82 |
+
if extra_model_paths is not None:
|
83 |
+
load_extra_path_config(extra_model_paths)
|
84 |
+
else:
|
85 |
+
print("Could not find the extra_model_paths config file.")
|
86 |
+
|
87 |
+
|
88 |
+
add_comfyui_directory_to_sys_path()
|
89 |
+
add_extra_model_paths()
|
90 |
+
|
91 |
+
|
92 |
+
def import_custom_nodes() -> None:
|
93 |
+
"""Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
|
94 |
+
|
95 |
+
This function sets up a new asyncio event loop, initializes the PromptServer,
|
96 |
+
creates a PromptQueue, and initializes the custom nodes.
|
97 |
+
"""
|
98 |
+
import asyncio
|
99 |
+
import execution
|
100 |
+
from nodes import init_extra_nodes
|
101 |
+
import server
|
102 |
+
|
103 |
+
# Creating a new event loop and setting it as the default loop
|
104 |
+
loop = asyncio.new_event_loop()
|
105 |
+
asyncio.set_event_loop(loop)
|
106 |
+
|
107 |
+
# Creating an instance of PromptServer with the loop
|
108 |
+
server_instance = server.PromptServer(loop)
|
109 |
+
execution.PromptQueue(server_instance)
|
110 |
+
|
111 |
+
# Initializing custom nodes
|
112 |
+
init_extra_nodes()
|
113 |
+
|
114 |
+
|
115 |
+
from nodes import (
|
116 |
+
LoadImage,
|
117 |
+
SaveImage,
|
118 |
+
NODE_CLASS_MAPPINGS,
|
119 |
+
CLIPTextEncode,
|
120 |
+
VAELoader,
|
121 |
+
VAEEncode,
|
122 |
+
DualCLIPLoader,
|
123 |
+
VAEDecode,
|
124 |
+
UNETLoader,
|
125 |
+
ControlNetLoader,
|
126 |
+
ControlNetApplyAdvanced,
|
127 |
+
)
|
128 |
+
|
129 |
+
|
130 |
+
def main():
|
131 |
+
import_custom_nodes()
|
132 |
+
with torch.inference_mode():
|
133 |
+
dualcliploader = DualCLIPLoader()
|
134 |
+
dualcliploader_94 = dualcliploader.load_clip(
|
135 |
+
clip_name1="t5xxl_fp16.safetensors",
|
136 |
+
clip_name2="clip_l.safetensors",
|
137 |
+
type="flux",
|
138 |
+
device="default",
|
139 |
+
)
|
140 |
+
|
141 |
+
cliptextencode = CLIPTextEncode()
|
142 |
+
cliptextencode_23 = cliptextencode.encode(
|
143 |
+
text="", clip=get_value_at_index(dualcliploader_94, 0)
|
144 |
+
)
|
145 |
+
|
146 |
+
loadimage = LoadImage()
|
147 |
+
loadimage_24 = loadimage.load_image(image="woman_face.jpg")
|
148 |
+
|
149 |
+
loadimage_40 = loadimage.load_image(image="woman1_gpt_2.png")
|
150 |
+
|
151 |
+
vaeloader = VAELoader()
|
152 |
+
vaeloader_95 = vaeloader.load_vae(vae_name="ae.safetensors")
|
153 |
+
|
154 |
+
vaeencode = VAEEncode()
|
155 |
+
vaeencode_35 = vaeencode.encode(
|
156 |
+
pixels=get_value_at_index(loadimage_40, 0),
|
157 |
+
vae=get_value_at_index(vaeloader_95, 0),
|
158 |
+
)
|
159 |
+
|
160 |
+
randomnoise = NODE_CLASS_MAPPINGS["RandomNoise"]()
|
161 |
+
randomnoise_39 = randomnoise.get_noise(noise_seed=random.randint(1, 2**64))
|
162 |
+
|
163 |
+
cliptextencode_42 = cliptextencode.encode(
|
164 |
+
text="", clip=get_value_at_index(dualcliploader_94, 0)
|
165 |
+
)
|
166 |
+
|
167 |
+
pulidfluxmodelloader = NODE_CLASS_MAPPINGS["PulidFluxModelLoader"]()
|
168 |
+
pulidfluxmodelloader_44 = pulidfluxmodelloader.load_model(
|
169 |
+
pulid_file="pulid_flux_v0.9.1.safetensors"
|
170 |
+
)
|
171 |
+
|
172 |
+
pulidfluxevacliploader = NODE_CLASS_MAPPINGS["PulidFluxEvaClipLoader"]()
|
173 |
+
pulidfluxevacliploader_45 = pulidfluxevacliploader.load_eva_clip()
|
174 |
+
|
175 |
+
pulidfluxinsightfaceloader = NODE_CLASS_MAPPINGS["PulidFluxInsightFaceLoader"]()
|
176 |
+
pulidfluxinsightfaceloader_46 = pulidfluxinsightfaceloader.load_insightface(
|
177 |
+
provider="CUDA"
|
178 |
+
)
|
179 |
+
|
180 |
+
controlnetloader = ControlNetLoader()
|
181 |
+
controlnetloader_49 = controlnetloader.load_controlnet(
|
182 |
+
control_net_name="Flux_Dev_ControlNet_Union_Pro_ShakkerLabs.safetensors"
|
183 |
+
)
|
184 |
+
|
185 |
+
ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
|
186 |
+
ksamplerselect_50 = ksamplerselect.get_sampler(sampler_name="euler")
|
187 |
+
|
188 |
+
unetloader = UNETLoader()
|
189 |
+
unetloader_93 = unetloader.load_unet(
|
190 |
+
unet_name="flux1-dev.safetensors", weight_dtype="default"
|
191 |
+
)
|
192 |
+
|
193 |
+
faceanalysismodels = NODE_CLASS_MAPPINGS["FaceAnalysisModels"]()
|
194 |
+
faceanalysismodels_118 = faceanalysismodels.load_models(
|
195 |
+
library="insightface", provider="CUDA"
|
196 |
+
)
|
197 |
+
|
198 |
+
applypulidflux = NODE_CLASS_MAPPINGS["ApplyPulidFlux"]()
|
199 |
+
setunioncontrolnettype = NODE_CLASS_MAPPINGS["SetUnionControlNetType"]()
|
200 |
+
controlnetapplyadvanced = ControlNetApplyAdvanced()
|
201 |
+
basicguider = NODE_CLASS_MAPPINGS["BasicGuider"]()
|
202 |
+
basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
|
203 |
+
samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
|
204 |
+
vaedecode = VAEDecode()
|
205 |
+
faceembeddistance = NODE_CLASS_MAPPINGS["FaceEmbedDistance"]()
|
206 |
+
display_any_rgthree = NODE_CLASS_MAPPINGS["Display Any (rgthree)"]()
|
207 |
+
image_comparer_rgthree = NODE_CLASS_MAPPINGS["Image Comparer (rgthree)"]()
|
208 |
+
saveimage = SaveImage()
|
209 |
+
|
210 |
+
for q in range(10):
|
211 |
+
applypulidflux_133 = applypulidflux.apply_pulid_flux(
|
212 |
+
weight=0.7500000000000001,
|
213 |
+
start_at=0.10000000000000002,
|
214 |
+
end_at=1,
|
215 |
+
fusion="mean",
|
216 |
+
fusion_weight_max=1,
|
217 |
+
fusion_weight_min=0,
|
218 |
+
train_step=1000,
|
219 |
+
use_gray=True,
|
220 |
+
model=get_value_at_index(unetloader_93, 0),
|
221 |
+
pulid_flux=get_value_at_index(pulidfluxmodelloader_44, 0),
|
222 |
+
eva_clip=get_value_at_index(pulidfluxevacliploader_45, 0),
|
223 |
+
face_analysis=get_value_at_index(pulidfluxinsightfaceloader_46, 0),
|
224 |
+
image=get_value_at_index(loadimage_24, 0),
|
225 |
+
unique_id=1674270197144619516,
|
226 |
+
)
|
227 |
+
|
228 |
+
setunioncontrolnettype_41 = setunioncontrolnettype.set_controlnet_type(
|
229 |
+
type="tile", control_net=get_value_at_index(controlnetloader_49, 0)
|
230 |
+
)
|
231 |
+
|
232 |
+
controlnetapplyadvanced_37 = controlnetapplyadvanced.apply_controlnet(
|
233 |
+
strength=1,
|
234 |
+
start_percent=0.1,
|
235 |
+
end_percent=0.8,
|
236 |
+
positive=get_value_at_index(cliptextencode_42, 0),
|
237 |
+
negative=get_value_at_index(cliptextencode_23, 0),
|
238 |
+
control_net=get_value_at_index(setunioncontrolnettype_41, 0),
|
239 |
+
image=get_value_at_index(loadimage_40, 0),
|
240 |
+
vae=get_value_at_index(vaeloader_95, 0),
|
241 |
+
)
|
242 |
+
|
243 |
+
basicguider_122 = basicguider.get_guider(
|
244 |
+
model=get_value_at_index(applypulidflux_133, 0),
|
245 |
+
conditioning=get_value_at_index(controlnetapplyadvanced_37, 0),
|
246 |
+
)
|
247 |
+
|
248 |
+
basicscheduler_131 = basicscheduler.get_sigmas(
|
249 |
+
scheduler="beta",
|
250 |
+
steps=28,
|
251 |
+
denoise=0.75,
|
252 |
+
model=get_value_at_index(applypulidflux_133, 0),
|
253 |
+
)
|
254 |
+
|
255 |
+
samplercustomadvanced_1 = samplercustomadvanced.sample(
|
256 |
+
noise=get_value_at_index(randomnoise_39, 0),
|
257 |
+
guider=get_value_at_index(basicguider_122, 0),
|
258 |
+
sampler=get_value_at_index(ksamplerselect_50, 0),
|
259 |
+
sigmas=get_value_at_index(basicscheduler_131, 0),
|
260 |
+
latent_image=get_value_at_index(vaeencode_35, 0),
|
261 |
+
)
|
262 |
+
|
263 |
+
vaedecode_114 = vaedecode.decode(
|
264 |
+
samples=get_value_at_index(samplercustomadvanced_1, 0),
|
265 |
+
vae=get_value_at_index(vaeloader_95, 0),
|
266 |
+
)
|
267 |
+
|
268 |
+
faceembeddistance_117 = faceembeddistance.analize(
|
269 |
+
similarity_metric="cosine",
|
270 |
+
filter_thresh=100,
|
271 |
+
filter_best=0,
|
272 |
+
generate_image_overlay=True,
|
273 |
+
analysis_models=get_value_at_index(faceanalysismodels_118, 0),
|
274 |
+
reference=get_value_at_index(loadimage_24, 0),
|
275 |
+
image=get_value_at_index(vaedecode_114, 0),
|
276 |
+
)
|
277 |
+
|
278 |
+
display_any_rgthree_121 = display_any_rgthree.main(
|
279 |
+
source=get_value_at_index(faceembeddistance_117, 1)
|
280 |
+
)
|
281 |
+
|
282 |
+
image_comparer_rgthree_123 = image_comparer_rgthree.compare_images(
|
283 |
+
image_a=get_value_at_index(loadimage_40, 0),
|
284 |
+
image_b=get_value_at_index(vaedecode_114, 0),
|
285 |
+
)
|
286 |
+
|
287 |
+
saveimage_128 = saveimage.save_images(
|
288 |
+
filename_prefix="FaceEnhanced",
|
289 |
+
images=get_value_at_index(vaedecode_114, 0),
|
290 |
+
)
|
291 |
+
|
292 |
+
saveimage_129 = saveimage.save_images(
|
293 |
+
filename_prefix="FaceEmbedDist",
|
294 |
+
images=get_value_at_index(faceembeddistance_117, 0),
|
295 |
+
)
|
296 |
+
|
297 |
+
|
298 |
+
if __name__ == "__main__":
|
299 |
+
main()
|