Aníbal Fuentes
9acbf63bce
Fixed unet saver with correct key names. Added on the fly bnb checkpoint and unet loaders, load any model in nf4 or fp4. Removed bloat from the easy resolution picker.
146 lines
5.3 KiB
Python
146 lines
5.3 KiB
Python
from enum import Enum
|
|
import logging
|
|
|
|
from comfy import model_management
|
|
from comfy.ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
|
|
from comfy.ldm.cascade.stage_a import StageA
|
|
from comfy.ldm.cascade.stage_c_coder import StageC_coder
|
|
from comfy.ldm.audio.autoencoder import AudioOobleckVAE
|
|
import yaml
|
|
|
|
from comfy import clip_vision
|
|
from comfy import gligen
|
|
from comfy import diffusers_convert
|
|
from comfy import model_detection
|
|
|
|
from comfy import sd1_clip
|
|
from comfy import sdxl_clip
|
|
import comfy.text_encoders.sd2_clip
|
|
import comfy.text_encoders.sd3_clip
|
|
import comfy.text_encoders.sa_t5
|
|
import comfy.text_encoders.aura_t5
|
|
import comfy.text_encoders.hydit
|
|
import comfy.text_encoders.flux
|
|
|
|
import comfy.model_patcher
|
|
import comfy.lora
|
|
import comfy.t2i_adapter.adapter
|
|
import comfy.supported_models_base
|
|
import comfy.taesd.taesd
|
|
|
|
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.model_base
|
|
import comfy.model_management
|
|
import comfy.model_sampling
|
|
|
|
import torch
|
|
import folder_paths
|
|
import json
|
|
import os
|
|
import safetensors
|
|
|
|
from comfy.cli_args import args
|
|
|
|
|
|
def save_diffusion_model(output_path, model, metadata=None, extra_keys={}):
|
|
# Ensure the model is loaded on the GPU for saving
|
|
model_management.load_models_gpu([model], force_patch_weights=True)
|
|
|
|
# Get the state dictionary for saving, excluding CLIP and VAE
|
|
sd = model.model.state_dict_for_saving()
|
|
|
|
# Clean up the state dictionary keys by removing the "model.diffusion_model." prefix
|
|
sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items()}
|
|
|
|
# Add any extra keys to the state dictionary
|
|
for k in extra_keys:
|
|
sd[k] = extra_keys[k]
|
|
|
|
# Ensure all tensors are contiguous before saving
|
|
for k in sd:
|
|
t = sd[k]
|
|
if not t.is_contiguous():
|
|
sd[k] = t.contiguous()
|
|
|
|
# Save the state dictionary to the specified output path
|
|
comfy.utils.save_torch_file(sd, output_path, metadata=metadata)
|
|
|
|
|
|
|
|
def save_model(model, filename_prefix=None, output_dir=None, prompt=None, extra_pnginfo=None):
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, output_dir)
|
|
prompt_info = ""
|
|
if prompt is not None:
|
|
prompt_info = json.dumps(prompt)
|
|
|
|
metadata = {}
|
|
|
|
enable_modelspec = True
|
|
if isinstance(model.model, comfy.model_base.SDXL):
|
|
if isinstance(model.model, comfy.model_base.SDXL_instructpix2pix):
|
|
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-edit"
|
|
else:
|
|
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
|
|
elif isinstance(model.model, comfy.model_base.SDXLRefiner):
|
|
metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
|
|
elif isinstance(model.model, comfy.model_base.SVD_img2vid):
|
|
metadata["modelspec.architecture"] = "stable-video-diffusion-img2vid-v1"
|
|
elif isinstance(model.model, comfy.model_base.SD3):
|
|
metadata["modelspec.architecture"] = "stable-diffusion-v3-medium" #TODO: other SD3 variants
|
|
else:
|
|
enable_modelspec = False
|
|
|
|
if enable_modelspec:
|
|
metadata["modelspec.sai_model_spec"] = "1.0.0"
|
|
metadata["modelspec.implementation"] = "sgm"
|
|
metadata["modelspec.title"] = "{} {}".format(filename, counter)
|
|
|
|
#TODO:
|
|
# "stable-diffusion-v1", "stable-diffusion-v1-inpainting", "stable-diffusion-v2-512",
|
|
# "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h",
|
|
# "v2-inpainting"
|
|
|
|
extra_keys = {}
|
|
model_sampling = model.get_model_object("model_sampling")
|
|
if isinstance(model_sampling, comfy.model_sampling.ModelSamplingContinuousEDM):
|
|
if isinstance(model_sampling, comfy.model_sampling.V_PREDICTION):
|
|
extra_keys["edm_vpred.sigma_max"] = torch.tensor(model_sampling.sigma_max).float()
|
|
extra_keys["edm_vpred.sigma_min"] = torch.tensor(model_sampling.sigma_min).float()
|
|
|
|
if model.model.model_type == comfy.model_base.ModelType.EPS:
|
|
metadata["modelspec.predict_key"] = "epsilon"
|
|
elif model.model.model_type == comfy.model_base.ModelType.V_PREDICTION:
|
|
metadata["modelspec.predict_key"] = "v"
|
|
|
|
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)
|
|
|
|
save_diffusion_model(output_checkpoint, model, metadata=metadata, extra_keys=extra_keys)
|
|
|
|
|
|
class ModelSave:
|
|
NAME = "Save Diffusion Model"
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "model": ("MODEL",),
|
|
"filename_prefix": ("STRING", {"default": "unet_output/ComfyUI"}),},
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save"
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = "advanced/model_merging"
|
|
|
|
def save(self, model, filename_prefix, prompt=None, extra_pnginfo=None):
|
|
save_model(model, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
|
|
return {} |