diff --git a/modules/comfy_ops.py b/modules/comfy_ops.py index d11d4d9..8fdf9e1 100644 --- a/modules/comfy_ops.py +++ b/modules/comfy_ops.py @@ -24,16 +24,20 @@ def dequantize_block_weight(weight, scale_weight): dequantized_weight = dequantized_weight.view_as(weight) return dequantized_weight -def block_scaled_fp8_ops(override_dtype=None, block_size=16): +def block_scaled_fp8_ops(override_dtype=None, block_size=64): logging.info(f"Using block scaled fp8: fp8 dtype: {override_dtype}, block size: {block_size}") + class block_scaled_fp8_op(manual_cast): class Linear(manual_cast.Linear): def __init__(self, *args, **kwargs): + # override Linear dtype if specified if override_dtype is not None: kwargs['dtype'] = override_dtype super().__init__(*args, **kwargs) + # override Linear's reset_parameters def reset_parameters(self): + # add scale_weight parameter to load scale factor for each block if not hasattr(self, 'scale_weight'): out_features, in_features = self.weight.shape if in_features % block_size != 0: @@ -46,7 +50,8 @@ def block_scaled_fp8_ops(override_dtype=None, block_size=16): ) return None - + + # override manual_cast.Linear's forward_comfy_cast_weights def forward_comfy_cast_weights(self, input): weight, bias = cast_bias_weight(self, input) device, dtype = weight.device, weight.dtype @@ -54,14 +59,16 @@ def block_scaled_fp8_ops(override_dtype=None, block_size=16): return torch.nn.functional.linear(input, dequantize_block_weight(weight, self.scale_weight), bias) + # comfyui load lora will call convert_weight def convert_weight(self, weight, inplace=False, **kwargs): scale_weight = self.scale_weight.to(device=weight.device, dtype=weight.dtype) if inplace: weight = dequantize_block_weight(weight, scale_weight) return weight else: - return dequantize_block_weight(weight, self.scale_weight) + return dequantize_block_weight(weight, scale_weight) + # comfyui load lora will call set_weight def set_weight(self, weight, inplace_update=False, seed=None, **kwargs): scale_weight = self.scale_weight.to(device=weight.device, dtype=weight.dtype) weight = comfy.float.stochastic_rounding(quantize_block_weight(weight, scale_weight), self.weight.dtype, seed=seed) diff --git a/modules/nodes.py b/modules/nodes.py index 6cd89ba..caaf5e6 100644 --- a/modules/nodes.py +++ b/modules/nodes.py @@ -11,8 +11,6 @@ from .comfy_ops import block_scaled_fp8_ops class MusubiUNETLoader: - """ - """ @classmethod def INPUT_TYPES(cls): return {