xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-01-06 22:36:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update nodes.py

Browse Source
This commit is contained in:
kijai 2024-06-15 01:19:00 +03:00
parent 3ba3ddf0b1
commit 36a8633aff
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
nodes/nodes.py
View File

@ -1666,6 +1666,7 @@ class CheckpointPerturbWeights:
return {"required": {
"model": ("MODEL",),
"joint_blocks": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}),
"final_layer": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}),
"rest_of_the_blocks": ("FLOAT", {"default": 0.02, "min": 0.001, "max": 10.0, "step": 0.001}),
}
}
@ -1675,7 +1676,7 @@ class CheckpointPerturbWeights:
CATEGORY = "KJNodes/experimental"
def mod(self, model, joint_blocks, rest_of_the_blocks):
def mod(self, model, joint_blocks, final_layer, rest_of_the_blocks):
import copy
device = model_management.get_torch_device()
model_copy = copy.deepcopy(model)
@ -1688,10 +1689,12 @@ class CheckpointPerturbWeights:
pbar = ProgressBar(len(keys))
for k in keys:
v = dict[k] # This is common to both conditions, so we can take it out
v = dict[k]
print(f'{k}: {v.std()}')
if k.startswith('joint_blocks'):
multiplier = joint_blocks
elif k.startswith('final_layer'):
multiplier = final_layer
else:
multiplier = rest_of_the_blocks
dict[k] += torch.normal(torch.zeros_like(v) * v.mean(), torch.ones_like(v) * v.std() * multiplier).to(device)