File size: 10,242 Bytes
943b415 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | import torch
import torch.nn.functional as F
import comfy.conds
import comfy.latent_formats
import comfy.model_base
import comfy.model_detection
import comfy.model_management
import comfy.model_patcher
import comfy.sd
import comfy.storage
import comfy.supported_models_base
import comfy.utils
import folder_paths
from .model import Nanosaur2Transformer2DModel
from .text_encoder import Nanosaur2Tokenizer, te
from .vae import Nanosaur2VAE
from .inpaint import Nanosaur2InpaintControl
GUIDANCE_MODES = ["alternate", "cfg", "path_drop"]
class Nanosaur2LatentFormat(comfy.latent_formats.LatentFormat):
latent_channels = 64
spacial_downscale_ratio = 16
class CONDTextPadded(comfy.conds.CONDRegular):
"""Batches text of different lengths by zero padding the token dimension; zero token weights mask the padding."""
def can_concat(self, other):
s1, s2 = self.cond.shape, other.cond.shape
return s1[0] == s2[0] and s1[2:] == s2[2:] and self.cond.device == other.cond.device
def concat(self, others):
conds = [self.cond] + [x.cond for x in others]
length = max(c.shape[1] for c in conds)
return torch.cat([F.pad(c, (0, 0) * (c.ndim - 2) + (0, length - c.shape[1])) for c in conds])
def sparse_skip_rows(batch, transformer_options):
"""Rows of the negative (uncond) chunks when this step uses path-drop guidance, else None.
"alternate" runs full CFG on even sampling steps and path-drop on odd ones.
"""
guidance = transformer_options.get("nanosaur2_guidance", "alternate")
cond_or_uncond = transformer_options.get("cond_or_uncond", [])
if guidance == "cfg" or 1 not in cond_or_uncond:
return None
if guidance == "alternate":
sigma = float(transformer_options["sigmas"][0])
sample_sigmas = transformer_options["sample_sigmas"].tolist()
step = min(range(len(sample_sigmas)), key=lambda i: abs(sample_sigmas[i] - sigma))
if step % 2 == 0:
return None
chunk = batch // len(cond_or_uncond)
return [c == 1 for c in cond_or_uncond for _ in range(chunk)]
class Nanosaur2Model(comfy.model_base.BaseModel):
def __init__(self, model_config, device=None):
super().__init__(model_config, comfy.model_base.ModelType.FLOW, device=device, unet_model=Nanosaur2Transformer2DModel)
def _apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
return super()._apply_model(x, t, c_concat, c_crossattn, control, transformer_options, sparse_skip=sparse_skip_rows(x.shape[0], transformer_options), **kwargs)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
cross_attn = kwargs.get("cross_attn", None)
if cross_attn is not None:
out["c_crossattn"] = CONDTextPadded(cross_attn)
token_weights = kwargs.get("token_weights", None)
if token_weights is not None:
out["token_weights"] = CONDTextPadded(token_weights)
return out
class Nanosaur2ModelConfig(comfy.supported_models_base.BASE):
unet_extra_config = {}
sampling_settings = {"multiplier": 1.0, "shift": 3.0}
latent_format = Nanosaur2LatentFormat
memory_usage_factor = 4.0
supported_inference_dtypes = [torch.bfloat16, torch.float32]
def get_model(self, state_dict, prefix="", device=None):
return Nanosaur2Model(self, device=device)
def detect_unet_config(sd):
hidden_size = sd["s_embedder.proj.weight"].shape[0]
return {
"in_channels": sd["final_layer.linear.weight"].shape[0],
"hidden_size": hidden_size,
"num_heads": hidden_size // sd["blocks.0.attn.q_norm.weight"].shape[0],
"num_blocks": comfy.model_detection.count_blocks(sd, "blocks.{}."),
"num_text_blocks": comfy.model_detection.count_blocks(sd, "text_refine_blocks.{}."),
"mlp_hidden": sd["blocks.0.mlp.w12.weight"].shape[0] // 2,
}
def load_diffusion_model(unet_path, disable_dynamic=False):
sd = comfy.utils.load_torch_file(unet_path)
model_config = Nanosaur2ModelConfig(detect_unet_config(sd))
load_device = comfy.model_management.get_torch_device()
offload_device = comfy.model_management.unet_offload_device()
unet_dtype = comfy.model_management.unet_dtype(model_params=comfy.utils.calculate_parameters(sd), supported_dtypes=model_config.supported_inference_dtypes, weight_dtype=comfy.utils.weight_dtype(sd))
manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
model = model_config.get_model(sd)
ModelPatcher = comfy.model_patcher.ModelPatcher if disable_dynamic else comfy.model_patcher.CoreModelPatcher
model_patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(sd))
if not comfy.model_management.is_device_cpu(offload_device):
model.to(offload_device)
model.load_model_weights(sd, "", assign=model_patcher.is_dynamic())
model_patcher.cached_patcher_init = (load_diffusion_model, (unet_path,))
return model_patcher
def load_clip(text_encoder_path, disable_dynamic=False):
sd = comfy.utils.load_torch_file(text_encoder_path, safe_load=True)
class ClipTarget:
params = {}
clip = te(dtype_llama=sd["model.embed_tokens.weight"].dtype, llama_quantization_metadata=comfy.utils.detect_layer_quantization(sd, ""))
tokenizer = Nanosaur2Tokenizer
clip = comfy.sd.CLIP(ClipTarget, tokenizer_data={"spiece_model": sd.pop("spiece_model")}, parameters=comfy.utils.calculate_parameters(sd), state_dict=[sd], disable_dynamic=disable_dynamic)
clip.patcher.cached_patcher_init = (load_clip_patcher, (text_encoder_path,))
return clip
def load_clip_patcher(text_encoder_path, disable_dynamic=False):
return load_clip(text_encoder_path, disable_dynamic=disable_dynamic).patcher
class Nanosaur2VAEWrapper(comfy.sd.VAE):
def __init__(self, sd, disable_dynamic=False):
self.latent_channels = sd["latent_mean"].shape[1]
self.first_stage_model = Nanosaur2VAE(latent_channels=self.latent_channels).eval()
self.memory_used_encode = lambda shape, dtype: 200 * shape[2] * shape[3] * comfy.model_management.dtype_size(dtype)
self.memory_used_decode = lambda shape, dtype: 2178 * shape[2] * shape[3] * 256 * comfy.model_management.dtype_size(dtype)
self.downscale_ratio = 16
self.upscale_ratio = 16
self.latent_dim = 2
self.output_channels = 3
self.pad_channel_value = None
self.process_input = lambda image: image * 2.0 - 1.0
self.process_output = lambda image: image.add_(1.0).div_(2.0).clamp_(0.0, 1.0)
self.working_dtypes = [torch.bfloat16, torch.float32]
self.disable_offload = False
self.not_video = False
self.size = None
self.downscale_index_formula = None
self.upscale_index_formula = None
self.extra_1d_channel = None
self.crop_input = True
self.handles_tiling = False
self.format_encoded = None
self.device = comfy.model_management.vae_device()
offload_device = comfy.model_management.vae_offload_device()
self.vae_dtype = comfy.model_management.vae_dtype(self.device, self.working_dtypes)
self.first_stage_model.to(self.vae_dtype)
comfy.model_management.archive_model_dtypes(self.first_stage_model)
self.output_device = comfy.model_management.intermediate_device()
ModelPatcher = comfy.model_patcher.ModelPatcher if disable_dynamic else comfy.model_patcher.CoreModelPatcher
self.patcher = ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device, fast_disk=comfy.storage.state_dict_fast_disk(sd))
self.first_stage_model.load_state_dict(sd, assign=self.patcher.is_dynamic())
if not self.patcher.is_dynamic():
self.first_stage_model.to(self.vae_dtype)
self.model_size()
def load_vae(vae_path, disable_dynamic=False):
vae = Nanosaur2VAEWrapper(comfy.utils.load_torch_file(vae_path), disable_dynamic=disable_dynamic)
vae.patcher.cached_patcher_init = (load_vae_patcher, (vae_path,))
return vae
def load_vae_patcher(vae_path, disable_dynamic=False):
return load_vae(vae_path, disable_dynamic=disable_dynamic).patcher
class Nanosaur2Loader:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"unet_name": (folder_paths.get_filename_list("diffusion_models"), {"default": "nanosaur2_diffusion_model.safetensors"}),
"text_encoder_name": (folder_paths.get_filename_list("text_encoders"), {"default": "nanosaur2_text_encoder.safetensors"}),
"vae_name": (folder_paths.get_filename_list("vae"), {"default": "nanosaur2_vae.safetensors"}),
"guidance": (GUIDANCE_MODES, {"default": "alternate", "tooltip": "Unconditional pass for CFG. cfg: full model on the negative prompt. path_drop: the negative prompt without the sparse middle blocks. alternate: cfg on even steps, path_drop on odd steps (recommended)."}),
},
}
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
FUNCTION = "load"
CATEGORY = "loaders"
DESCRIPTION = "Loads the Nanosaur2 diffusion model, its Gemma3-270M text encoder and its Semantic VAE."
def load(self, unet_name, text_encoder_name, vae_name, guidance):
model = load_diffusion_model(folder_paths.get_full_path_or_raise("diffusion_models", unet_name))
model.model_options["transformer_options"]["nanosaur2_guidance"] = guidance
clip = load_clip(folder_paths.get_full_path_or_raise("text_encoders", text_encoder_name))
vae = load_vae(folder_paths.get_full_path_or_raise("vae", vae_name))
return (model, clip, vae)
NODE_CLASS_MAPPINGS = {
"Nanosaur2Loader": Nanosaur2Loader,
"Nanosaur2InpaintControl": Nanosaur2InpaintControl,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Nanosaur2Loader": "Load Nanosaur2",
"Nanosaur2InpaintControl": "Apply Nanosaur2 Inpaint Control",
}
|