133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
import torch
|
|
import os
|
|
from diffusers.models.normalization import AdaGroupNorm
|
|
import safetensors.torch
|
|
|
|
import folder_paths
|
|
|
|
from comfy.utils import common_upscale
|
|
|
|
from .model import NoiseTransformer, SVDNoiseUnet
|
|
|
|
MODELS_DIR = os.path.join(folder_paths.models_dir, "npnets")
|
|
if "npnets" not in folder_paths.folder_names_and_paths:
|
|
current_paths = [MODELS_DIR]
|
|
else:
|
|
current_paths, _ = folder_paths.folder_names_and_paths["npnets"]
|
|
folder_paths.folder_names_and_paths["npnets"] = (current_paths, folder_paths.supported_pt_extensions)
|
|
|
|
class Noise_GoldenNoise:
|
|
def __init__(self, model, noise, cond):
|
|
self.model = model
|
|
self.noise = noise
|
|
self.cond = cond
|
|
self.seed = self.noise.seed
|
|
|
|
def generate_noise(self, input_latent):
|
|
og_shape = input_latent["samples"].shape
|
|
rescale = og_shape[-1] != 128 or og_shape[-2] != 128
|
|
|
|
if rescale:
|
|
input_latent = input_latent.copy()
|
|
input_latent["samples"] = common_upscale(input_latent["samples"], 128, 128, "nearest-exact", "disabled")
|
|
|
|
initial_noise = self.noise.generate_noise(input_latent).to(self.model["device"])
|
|
cond = self.cond[0].clone().to(self.model["device"])
|
|
|
|
try:
|
|
for _, cond in enumerate(torch.split(cond, 77, 1)):
|
|
cond = cond.float().view(cond.shape[0], -1)
|
|
text_emb = self.model["text_embedding"](initial_noise.float(), cond)
|
|
|
|
encoder_hidden_states_svd = initial_noise
|
|
encoder_hidden_states_embedding = initial_noise + text_emb
|
|
|
|
golden_embedding = self.model["unet_embedding"](encoder_hidden_states_embedding.float())
|
|
|
|
initial_noise = (
|
|
self.model["unet_svd"](encoder_hidden_states_svd.float())
|
|
+ (2 * torch.sigmoid(self.model["_alpha"]) - 1) * text_emb
|
|
+ self.model["_beta"] * golden_embedding
|
|
)
|
|
|
|
golden_noise = initial_noise
|
|
golden_noise.to("cpu")
|
|
|
|
if rescale:
|
|
golden_noise = common_upscale(golden_noise, og_shape[-1], og_shape[-2], "nearest-exact", "disabled")
|
|
|
|
except Exception as e:
|
|
print("Noise could not be turned golden:", e)
|
|
return initial_noise
|
|
|
|
return golden_noise
|
|
|
|
class GoldenNoise:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"noise": ("NOISE",),
|
|
"conditioning": ("CONDITIONING",),
|
|
"model_id": (["SDXL", "DreamShaper", "DiT"],),
|
|
"npnet_model": (folder_paths.get_filename_list("npnets"),),
|
|
"device": (["cuda", "cpu"],),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("NOISE",)
|
|
OUTPUT_TOOLTIPS = ("The golden noise",)
|
|
|
|
CATEGORY = "sampling/custom_sampling/noise"
|
|
FUNCTION = "get_noise"
|
|
DESCRIPTION = "Turn noise into golden noise."
|
|
|
|
def __init__(self):
|
|
self.models = {}
|
|
|
|
def load_model(self, model_id, npnet_model_path, device="cuda"):
|
|
if model_id not in self.models:
|
|
unet_embedding = NoiseTransformer(resolution=128).to(device, torch.float32)
|
|
unet_svd = SVDNoiseUnet(resolution=128).to(device, torch.float32)
|
|
|
|
if model_id == 'DiT':
|
|
text_embedding = AdaGroupNorm(1024 * 77, 4, 1, eps=1e-6).to(device, torch.float32)
|
|
else:
|
|
text_embedding = AdaGroupNorm(2048 * 77, 4, 1, eps=1e-6).to(device, torch.float32)
|
|
|
|
is_pth = ".pth" in npnet_model_path
|
|
|
|
if is_pth:
|
|
gloden_unet = torch.load(npnet_model_path, weights_only=True, map_location=device)
|
|
else:
|
|
gloden_unet = safetensors.torch.load_file(npnet_model_path, device=device)
|
|
|
|
unet_svd.load_state_dict(gloden_unet["unet_svd"] if is_pth else {k.replace("unet_svd.", ""): v for k, v in gloden_unet.items() if k.startswith("unet_svd.")})
|
|
unet_embedding.load_state_dict(gloden_unet["unet_embedding"] if is_pth else {k.replace("unet_embedding.", ""): v for k, v in gloden_unet.items() if k.startswith("unet_embedding.")})
|
|
text_embedding.load_state_dict(gloden_unet["embeeding"] if is_pth else {k.replace("embeeding.", ""): v for k, v in gloden_unet.items() if k.startswith("embeeding.")})
|
|
_alpha = gloden_unet["alpha"]
|
|
_beta = gloden_unet["beta"]
|
|
|
|
self.models[model_id] = {
|
|
"unet_svd": unet_svd,
|
|
"unet_embedding": unet_embedding,
|
|
"text_embedding": text_embedding,
|
|
"_alpha": _alpha,
|
|
"_beta": _beta,
|
|
"device": device
|
|
}
|
|
|
|
return self.models[model_id]
|
|
|
|
def get_noise(self, noise, conditioning, model_id, npnet_model, device):
|
|
npnet_model_path = folder_paths.get_full_path("npnets", npnet_model)
|
|
|
|
model = self.load_model(model_id, npnet_model_path, device)
|
|
cond = conditioning[0]
|
|
|
|
return (Noise_GoldenNoise(model, noise, cond),)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GoldenNoise": GoldenNoise,
|
|
}
|