Files
CavinHuang-comfyui-nodes-docs/docs/VAESave.md
2024-05-15 08:09:49 +08:00

2.9 KiB

Documentation

VAESave节点负责保存变分自编码器(VAE)模型的状态。它提供了将模型的状态字典序列化到文件的功能,该文件可以用于以后的检索或部署。此节点对于保存训练模型的参数并确保可以在不丢失信息的情况下重新加载它们至关重要。

Input types

Required

  • vae
    • 'vae'参数是必需的,因为它代表了要保存的变分自编码器模型。它是节点运行所必需的,也是决定节点操作的主要输入。
    • Comfy dtype: VAE
    • Python dtype: comfy.model_base.VAE
  • filename_prefix
    • 'filename_prefix'参数用于定义保存VAE模型的输出文件名的前缀。它在组织和识别保存的模型文件中起着重要作用。
    • Comfy dtype: STRING
    • Python dtype: str

Optional

  • prompt
    • 'prompt'参数虽然是可选的,但可以用来包含有关模型预期用途的附加上下文或描述,这对于将来的参考或元数据目的可能很有用。
    • Comfy dtype: PROMPT
    • Python dtype: str
  • extra_pnginfo
    • 'extra_pnginfo'参数允许包含可以与模型一起存储的额外信息。这对于添加注释或其他可能需要的相关数据可能很有用。
    • Comfy dtype: EXTRA_PNGINFO
    • Python dtype: dict

Output types

  • output_checkpoint
    • 'output_checkpoint'参数表示保存的模型文件的路径。它标志着保存过程的成功完成,并提供了可以找到模型的位置。
    • Comfy dtype: STRING
    • Python dtype: str

Usage tips

  • Infra type: CPU

Source code

class VAESave:

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()

    @classmethod
    def INPUT_TYPES(s):
        return {'required': {'vae': ('VAE',), 'filename_prefix': ('STRING', {'default': 'vae/ComfyUI_vae'})}, 'hidden': {'prompt': 'PROMPT', 'extra_pnginfo': 'EXTRA_PNGINFO'}}
    RETURN_TYPES = ()
    FUNCTION = 'save'
    OUTPUT_NODE = True
    CATEGORY = 'advanced/model_merging'

    def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None):
        (full_output_folder, filename, counter, subfolder, filename_prefix) = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
        prompt_info = ''
        if prompt is not None:
            prompt_info = json.dumps(prompt)
        metadata = {}
        if not args.disable_metadata:
            metadata['prompt'] = prompt_info
            if extra_pnginfo is not None:
                for x in extra_pnginfo:
                    metadata[x] = json.dumps(extra_pnginfo[x])
        output_checkpoint = f'{filename}_{counter:05}_.safetensors'
        output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
        comfy.utils.save_torch_file(vae.get_sd(), output_checkpoint, metadata=metadata)
        return {}