comfyanonymous/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-27 18:34:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactoring: use the same code for "scan task" and realtime DB population

Browse Source
This commit is contained in:
bigcat88
2025-08-25 13:31:56 +03:00
parent d7464e9e73
commit 09dabf95bc
7 changed files with 178 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
nodes.py
View File

@@ -31,7 +31,6 @@ from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator
from comfy_api.internal import register_versions, ComfyAPIWithVersion
from comfy_api.version_list import supported_versions
from comfy_api.latest import io, ComfyExtension
from app.assets_manager import populate_db_with_asset
import comfy.clip_vision
@@ -555,9 +554,7 @@ class CheckpointLoader:
def load_checkpoint(self, config_name, ckpt_name):
config_path = folder_paths.get_full_path("configs", config_name)
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
out = comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
populate_db_with_asset(["models", "checkpoint"], ckpt_name, ckpt_path)
return out
return comfy.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
class CheckpointLoaderSimple:
@classmethod
@@ -579,7 +576,6 @@ class CheckpointLoaderSimple:
def load_checkpoint(self, ckpt_name):
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
populate_db_with_asset(["models", "checkpoint"], ckpt_name, ckpt_path)
return out[:3]
class DiffusersLoader:
@@ -622,7 +618,6 @@ class unCLIPCheckpointLoader:
def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)
out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
populate_db_with_asset(["models", "checkpoint"], ckpt_name, ckpt_path)
return out
class CLIPSetLastLayer:
@@ -681,7 +676,6 @@ class LoraLoader:
self.loaded_lora = (lora_path, lora)
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
populate_db_with_asset(["models", "lora"], lora_name, lora_path)
return (model_lora, clip_lora)
class LoraLoaderModelOnly(LoraLoader):
@@ -746,15 +740,11 @@ class VAELoader:
encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes))
decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes))
encoder_path = folder_paths.get_full_path_or_raise("vae_approx", encoder)
populate_db_with_asset(["models", "vae-approx", "encoder"], name, encoder_path)
enc = comfy.utils.load_torch_file(encoder_path)
enc = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", encoder))
for k in enc:
sd["taesd_encoder.{}".format(k)] = enc[k]
decoder_path = folder_paths.get_full_path_or_raise("vae_approx", decoder)
populate_db_with_asset(["models", "vae-approx", "decoder"], name, decoder_path)
dec = comfy.utils.load_torch_file(decoder_path)
dec = comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("vae_approx", decoder))
for k in dec:
sd["taesd_decoder.{}".format(k)] = dec[k]
@@ -787,7 +777,6 @@ class VAELoader:
else:
vae_path = folder_paths.get_full_path_or_raise("vae", vae_name)
sd = comfy.utils.load_torch_file(vae_path)
populate_db_with_asset(["models", "vae"], vae_name, vae_path)
vae = comfy.sd.VAE(sd=sd)
vae.throw_exception_if_invalid()
return (vae,)
@@ -807,7 +796,6 @@ class ControlNetLoader:
controlnet = comfy.controlnet.load_controlnet(controlnet_path)
if controlnet is None:
raise RuntimeError("ERROR: controlnet file is invalid and does not contain a valid controlnet model.")
populate_db_with_asset(["models", "controlnet"], control_net_name, controlnet_path)
return (controlnet,)
class DiffControlNetLoader:
@@ -824,7 +812,6 @@ class DiffControlNetLoader:
def load_controlnet(self, model, control_net_name):
controlnet_path = folder_paths.get_full_path_or_raise("controlnet", control_net_name)
controlnet = comfy.controlnet.load_controlnet(controlnet_path, model)
populate_db_with_asset(["models", "controlnet"], control_net_name, controlnet_path)
return (controlnet,)
@@ -932,7 +919,6 @@ class UNETLoader:
unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name)
model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options)
populate_db_with_asset(["models", "diffusion-model"], unet_name, unet_path)
return (model,)
class CLIPLoader:
@@ -960,7 +946,6 @@ class CLIPLoader:
clip_path = folder_paths.get_full_path_or_raise("text_encoders", clip_name)
clip = comfy.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
populate_db_with_asset(["models", "text-encoder"], clip_name, clip_path)
return (clip,)
class DualCLIPLoader:
@@ -991,8 +976,6 @@ class DualCLIPLoader:
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
clip = comfy.sd.load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=folder_paths.get_folder_paths("embeddings"), clip_type=clip_type, model_options=model_options)
populate_db_with_asset(["models", "text-encoder"], clip_name1, clip_path1)
populate_db_with_asset(["models", "text-encoder"], clip_name2, clip_path2)
return (clip,)
class CLIPVisionLoader:
@@ -1010,7 +993,6 @@ class CLIPVisionLoader:
clip_vision = comfy.clip_vision.load(clip_path)
if clip_vision is None:
raise RuntimeError("ERROR: clip vision file is invalid and does not contain a valid vision model.")
populate_db_with_asset(["models", "clip-vision"], clip_name, clip_path)
return (clip_vision,)
class CLIPVisionEncode:
@@ -1045,7 +1027,6 @@ class StyleModelLoader:
def load_style_model(self, style_model_name):
style_model_path = folder_paths.get_full_path_or_raise("style_models", style_model_name)
style_model = comfy.sd.load_style_model(style_model_path)
populate_db_with_asset(["models", "style-model"], style_model_name, style_model_path)
return (style_model,)
@@ -1143,7 +1124,6 @@ class GLIGENLoader:
def load_gligen(self, gligen_name):
gligen_path = folder_paths.get_full_path_or_raise("gligen", gligen_name)
gligen = comfy.sd.load_gligen(gligen_path)
populate_db_with_asset(["models", "gligen"], gligen_name, gligen_path)
return (gligen,)
class GLIGENTextBoxApply: