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", }