| import comfy, folder_paths, execution
|
| from server import PromptServer
|
| from comfy import samplers
|
| import functools
|
|
|
|
|
|
|
| class SwarmEmbedLoaderListProvider:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {
|
| "required": {
|
| "embed_name": (folder_paths.get_filename_list("embeddings"), )
|
| }
|
| }
|
|
|
| CATEGORY = "SwarmUI/internal"
|
| RETURN_TYPES = ("EMBEDDING",)
|
| FUNCTION = "load_embed"
|
| DESCRIPTION = "Internal node just intended to provide a list of currently known embeddings to Swarm. You can also use it to blindly load an embedding file if you need to."
|
|
|
| def load_embed(self, embed_name):
|
| embed_path = folder_paths.get_full_path("embedding", embed_name)
|
| sd = comfy.utils.load_torch_file(embed_path)
|
| return (sd,)
|
|
|
|
|
| class SwarmJustLoadTheModelPlease:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {
|
| "required": {
|
| "model": ("MODEL",),
|
| "clip": ("CLIP,GEMMA",),
|
| "vae": ("VAE",),
|
| }
|
| }
|
|
|
| CATEGORY = "SwarmUI/internal"
|
| RETURN_TYPES = ()
|
| FUNCTION = "just_load"
|
| OUTPUT_NODE = True
|
| DESCRIPTION = "Internal node that acts as a final output for a model/clip/vae. This allows swarm to load models when needed without generating anything."
|
|
|
| def just_load(self, model, clip, vae):
|
| if model is None:
|
| raise ValueError("The model failed to load")
|
| if clip is None:
|
| raise ValueError("The text encoders (CLIP) failed to load")
|
| if vae is None:
|
| raise ValueError("The VAE failed to load")
|
| return {}
|
|
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "SwarmEmbedLoaderListProvider": SwarmEmbedLoaderListProvider,
|
| "SwarmJustLoadTheModelPlease": SwarmJustLoadTheModelPlease
|
| }
|
|
|
|
|
|
|
| ORIG_EXECUTION_VALIDATE = execution.validate_inputs
|
| async def validate_inputs(prompt_id, prompt, item, validated):
|
| raw_result = await ORIG_EXECUTION_VALIDATE(prompt_id, prompt, item, validated)
|
| if raw_result is None:
|
| return None
|
| (did_succeed, errors, unique_id) = raw_result
|
| if did_succeed:
|
| return raw_result
|
| for error in errors:
|
| if error['type'] == "return_type_mismatch":
|
| o_id = error['extra_info']['linked_node'][0]
|
| o_class_type = prompt[o_id]['class_type']
|
| if o_class_type == "SwarmInputModelName" or o_class_type == "SwarmInputDropdown":
|
| errors.remove(error)
|
| did_succeed = len(errors) == 0
|
| return (did_succeed, errors, unique_id)
|
|
|
| execution.validate_inputs = validate_inputs
|
|
|
|
|
| try:
|
| from app import logger
|
| def patch_interceptor(interceptor):
|
| if interceptor:
|
| orig = interceptor.write
|
| def write(self, data):
|
| orig(data)
|
| self.flush()
|
| interceptor.write = functools.partial(write, interceptor)
|
|
|
|
|
| interceptor.reconfigure(encoding='utf-8')
|
| patch_interceptor(logger.stdout_interceptor)
|
| patch_interceptor(logger.stderr_interceptor)
|
| except Exception as e:
|
| import traceback
|
| traceback.print_exc()
|
|
|
|
|
| try:
|
| server = PromptServer.instance
|
| orig = server.send_bytes
|
| async def send_bytes(self, event, data, sid=None):
|
| if event == 9999123:
|
| event = 4
|
| await orig(event, data, sid=sid)
|
| server.send_bytes = functools.partial(send_bytes, server)
|
| except Exception as e:
|
| import traceback
|
| traceback.print_exc()
|
|
|
