xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-15 19:23:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Set scale device

Browse Source
This commit is contained in:
lspindler 2025-08-25 10:09:49 +02:00
parent 70cb59eed6
commit 1404916962
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ops.py
View File

@ -183,7 +183,6 @@ class disable_weight_init:
self.register_buffer('weight', weight) self.register_buffer('weight', weight)
return return
device = state_dict[f"{prefix}weight"].device
weight_dtype = state_dict[f"{prefix}weight"].dtype weight_dtype = state_dict[f"{prefix}weight"].dtype
weight = torch.nn.Parameter( weight = torch.nn.Parameter(
torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype) torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype)
@ -196,11 +195,11 @@ class disable_weight_init:
scale_weight = state_dict.get(f"{prefix}scale_weight", None) scale_weight = state_dict.get(f"{prefix}scale_weight", None)
if scale_weight is None: if scale_weight is None:
raise Exception("Using quantized Weights requires a scale to be present!") raise Exception("Using quantized Weights requires a scale to be present!")
self.register_buffer('scale_weight', scale_weight.float()) self.register_buffer('scale_weight', scale_weight.to(device=self.device, dtype=torch.float32))
scale_input = state_dict.get(f"{prefix}scale_input", None) scale_input = state_dict.get(f"{prefix}scale_input", None)
if scale_input is not None: if scale_input is not None:
self.register_buffer('scale_input', scale_input.float()) self.register_buffer('scale_input', scale_input.to(device=self.device, dtype=torch.float32))
self.forward = types.MethodType(get_quantized_forward(scale_weight, scale_input), self) self.forward = types.MethodType(get_quantized_forward(scale_weight, scale_input), self)
setattr(self, "quantizer", get_quantizer_fn(scale_weight, scale_input)) setattr(self, "quantizer", get_quantizer_fn(scale_weight, scale_input))