xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-15 07:44:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update nodes.py

Browse Source
This commit is contained in:
kijai 2024-11-22 17:15:50 +02:00
parent 44620cb566
commit 8f057eb563
1 changed files with 17 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
nodes/nodes.py
View File

@ -2172,11 +2172,6 @@ class CheckpointLoaderKJ:
def patch(self, ckpt_name, patch_cublaslinear, sage_attention): def patch(self, ckpt_name, patch_cublaslinear, sage_attention):
from comfy.ops import disable_weight_init, CastWeightBiasOp, cast_bias_weight from comfy.ops import disable_weight_init, CastWeightBiasOp, cast_bias_weight
from nodes import CheckpointLoaderSimple from nodes import CheckpointLoaderSimple
if patch_cublaslinear:
try:
from cublas_ops import CublasLinear
except ImportError:
raise Exception("Can't import 'torch-cublas-hgemm', install it from here https://github.com/aredden/torch-cublas-hgemm")
if sage_attention: if sage_attention:
from sageattention import sageattn from sageattention import sageattn
@ -2215,6 +2210,11 @@ class CheckpointLoaderKJ:
else: else:
return super().forward(*args, **kwargs) return super().forward(*args, **kwargs)
if patch_cublaslinear:
try:
from cublas_ops import CublasLinear
except ImportError:
raise Exception("Can't import 'torch-cublas-hgemm', install it from here https://github.com/aredden/torch-cublas-hgemm")
class PatchedLinear(CublasLinear, CastWeightBiasOp): class PatchedLinear(CublasLinear, CastWeightBiasOp):
def reset_parameters(self): def reset_parameters(self):
return None return None
@ -2228,8 +2228,6 @@ class CheckpointLoaderKJ:
return self.forward_comfy_cast_weights(*args, **kwargs) return self.forward_comfy_cast_weights(*args, **kwargs)
else: else:
return super().forward(*args, **kwargs) return super().forward(*args, **kwargs)
if patch_cublaslinear:
disable_weight_init.Linear = PatchedLinear disable_weight_init.Linear = PatchedLinear
else: else:
disable_weight_init.Linear = OriginalLinear disable_weight_init.Linear = OriginalLinear