Spaces:
Paused
Paused
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()
|