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