114 lines
4.8 KiB
Python
114 lines
4.8 KiB
Python
# components/model.py
|
|
from comfy.model_patcher import ModelPatcher
|
|
import torch
|
|
from typing import Dict, Tuple, Optional
|
|
|
|
from ..ddare.const import UTIL_CATEGORY, MODEL_MASK
|
|
from ..ddare.mask import ModelMask
|
|
from ..ddare.merge import merge_tensors, METHODS
|
|
from ..ddare.model import collect_layers, layers_for_mask
|
|
from ..ddare.util import cuda_memory_profiler, get_device, get_patched_state
|
|
|
|
|
|
class ModelNoiseInjector:
|
|
"""
|
|
Inject gaussian noise into layers of a model.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, tuple]:
|
|
"""
|
|
Defines the input types for the noise injection process.
|
|
|
|
Returns:
|
|
Dict[str, tuple]: A dictionary specifying the required model types and parameters.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"model": ('MODEL',),
|
|
"operation": (["random", "gaussian"], {'default': "gaussian"}),
|
|
"ratio": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"mean": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"std": ("FLOAT", {"default": 0.01, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"seed": ("INT", {"default": 1, "min":0, "max": 99999999999}),
|
|
# It would be cool if we could autopupulate layer names here from a dropdown
|
|
"layers": ("STRING", {"multiline": True, "default": "*.to_v*"}),
|
|
"method": (["comfy",] + METHODS, {"default": "comfy"}),
|
|
},
|
|
"optional": {
|
|
"model_mask": (MODEL_MASK,),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "edit"
|
|
CATEGORY = UTIL_CATEGORY
|
|
|
|
def edit(self, model : ModelPatcher, operation : str, mean : float, std : float, ratio : float, seed : int, layers : str, method : str, model_mask : Optional[ModelMask] = None, **kwargs) -> Tuple[ModelPatcher]:
|
|
"""
|
|
Injects noise into the model.
|
|
|
|
Args:
|
|
model (ModelPatcher): The model to inject noise into.
|
|
operation (str): The type of noise to inject. Either "random" or "gaussian".
|
|
mean (float): The first argument for the noise injection. Only used for "gaussian" noise.
|
|
std (float): The second argument for the noise injection. Only used for "gaussian" noise.
|
|
ratio (float): The strength of the noise to inject.
|
|
seed (int): The seed for the noise injection.
|
|
layers (str): The layers to inject noise into.
|
|
model_mask (ModelMask): A ModelMask instance to use for masking the model. Default is None.
|
|
|
|
Returns:
|
|
Tuple[ModelPatcher]: A tuple containing the modified ModelPatcher instance.
|
|
"""
|
|
|
|
device = get_device()
|
|
|
|
m = model.clone() # Clone model_a to keep its structure
|
|
|
|
with cuda_memory_profiler():
|
|
model_sd = get_patched_state(m)
|
|
|
|
keys = list(model_sd.keys())
|
|
collected_targets = collect_layers(layers, keys)
|
|
if len(collected_targets) == 0:
|
|
raise ValueError("No layers specified")
|
|
|
|
for i, target in enumerate(collected_targets):
|
|
for j, k in enumerate(layers_for_mask(target, keys)):
|
|
if k not in model_sd:
|
|
print("could not patch. key doesn't exist in model:", k)
|
|
continue
|
|
|
|
torch.manual_seed(seed + i * 1000 + j)
|
|
|
|
# Get our tensor
|
|
mask : torch.Tensor = model_mask.get_layer_mask(k) if model_mask is not None else None
|
|
a : torch.Tensor = model_sd[k]
|
|
|
|
if operation == "random":
|
|
# Create a random mask of the same shape as the given layer.
|
|
random = torch.rand(a.shape) - 0.5
|
|
else:
|
|
# Create a gaussian noise mask of the same shape as the given layer.
|
|
random = torch.normal(mean, std, size=a.shape) - mean
|
|
|
|
if mask is None:
|
|
result_tensor = a + random
|
|
else:
|
|
result_tensor = torch.where(mask.to(device), a.to(device) + random.to(device), a.to(device))
|
|
del random
|
|
|
|
# Merge our tensors
|
|
if method == "comfy":
|
|
strength_patch = 1.0 - ratio
|
|
strength_model = ratio
|
|
else:
|
|
result_tensor = merge_tensors(method, a.to(device), result_tensor, 1 - ratio)
|
|
|
|
strength_model = 0
|
|
strength_patch = 1.0
|
|
|
|
m.add_patches({k: (result_tensor.to('cpu'),)}, strength_patch, strength_model)
|
|
|
|
return (m,) |