103 lines
3.8 KiB
Python
103 lines
3.8 KiB
Python
import torch
|
|
from comfy_api.v0_0_2 import io
|
|
from ... import ROOT_NAME
|
|
|
|
CATEGORY_NAME = ROOT_NAME + "cd-tuner"
|
|
|
|
class CDTuner(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="CD_Tuner|cgem156",
|
|
display_name="CD Tuner 🍌",
|
|
category=CATEGORY_NAME,
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Float.Input("detail_1", default=0, min=-10, max=10, step=0.1),
|
|
io.Float.Input("detail_2", default=0, min=-10, max=10, step=0.1),
|
|
io.Float.Input("contrast_1", default=0, min=-20, max=20, step=0.1),
|
|
io.Int.Input("start", default=0, min=0, max=1000, step=1, display_mode=io.NumberDisplay.number),
|
|
io.Int.Input("end", default=1000, min=0, max=1000, step=1, display_mode=io.NumberDisplay.number),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, detail_1, detail_2, contrast_1, start, end) -> io.NodeOutput:
|
|
'''
|
|
detail_1: 最初のConv層のweightを減らしbiasを増やすことで、detailを増やす・・？
|
|
detail_2: 最後のConv層前のGroupNormの以下略
|
|
contrast_1: 最後のConv層のbiasの0チャンネル目を増やすことでコントラストを増やす・・・？
|
|
'''
|
|
new_model = model.clone()
|
|
ratios = fineman([detail_1, detail_2, contrast_1])
|
|
storedweights = {}
|
|
|
|
# unet計算前後のパッチ
|
|
def apply_cdtuner(model_function, kwargs):
|
|
t = new_model.model.model_sampling.timestep(kwargs["timestep"])
|
|
if t[0] < (1000 - end) or t[0] > (1000 - start):
|
|
return model_function(kwargs["input"], kwargs["timestep"], **kwargs["c"])
|
|
for i, name in enumerate(ADJUSTS):
|
|
# 元の重みをロード
|
|
storedweights[name] = getset_nested_module_tensor(True, new_model, name).clone()
|
|
if 4 > i:
|
|
new_weight = storedweights[name] * ratios[i]
|
|
else:
|
|
device = storedweights[name].device
|
|
dtype = storedweights[name].dtype
|
|
new_weight = storedweights[name] + torch.tensor(ratios[i], device=device, dtype=dtype)
|
|
# 重みを書き換え
|
|
getset_nested_module_tensor(False, new_model, name, new_tensor=new_weight)
|
|
retval = model_function(kwargs["input"], kwargs["timestep"], **kwargs["c"])
|
|
|
|
# 重みを元に戻す
|
|
for name in ADJUSTS:
|
|
getset_nested_module_tensor(False, new_model, name, new_tensor=storedweights[name])
|
|
|
|
return retval
|
|
|
|
new_model.set_model_unet_function_wrapper(apply_cdtuner)
|
|
|
|
return io.NodeOutput(new_model)
|
|
|
|
|
|
def getset_nested_module_tensor(clone, model, tensor_path, new_tensor=None):
|
|
sdmodules = tensor_path.split('.')
|
|
target_module = model
|
|
last_attr = None
|
|
|
|
for module_name in sdmodules if clone else sdmodules[:-1]:
|
|
if module_name.isdigit():
|
|
target_module = target_module[int(module_name)]
|
|
else:
|
|
target_module = getattr(target_module, module_name)
|
|
|
|
if clone:
|
|
return target_module
|
|
|
|
last_attr = sdmodules[-1]
|
|
setattr(target_module, last_attr, torch.nn.Parameter(new_tensor))
|
|
|
|
# なんでfineman?
|
|
def fineman(fine):
|
|
fine = [
|
|
1 - fine[0] * 0.01,
|
|
1 + fine[0] * 0.02,
|
|
1 - fine[1] * 0.01,
|
|
1 + fine[1] * 0.02,
|
|
[fine[2] * 0.02, 0, 0, 0]
|
|
]
|
|
return fine
|
|
|
|
|
|
ADJUSTS = [
|
|
"model.diffusion_model.input_blocks.0.0.weight",
|
|
"model.diffusion_model.input_blocks.0.0.bias",
|
|
"model.diffusion_model.out.0.weight",
|
|
"model.diffusion_model.out.0.bias",
|
|
"model.diffusion_model.out.2.bias",
|
|
]
|