Rishi Desai commited on
Commit
431ce81
·
1 Parent(s): 9a054ff

raw workflow in python

Browse files
Files changed (1) hide show
  1. 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()