diff --git a/fp8_optimization_v2.py b/fp8_optimization_v2.py new file mode 100644 index 0000000..03b3e7a --- /dev/null +++ b/fp8_optimization_v2.py @@ -0,0 +1,113 @@ +import torch +import torch.nn as nn +from accelerate import init_empty_weights + +#based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py +def _replace_linear(model, compute_dtype, state_dict, prefix="", patches=None, scale_weights=None): + + has_children = list(model.children()) + if not has_children: + return + for name, module in model.named_children(): + module_prefix = prefix + name + "." + _replace_linear(module, compute_dtype, state_dict, module_prefix, patches, scale_weights) + + if isinstance(module, nn.Linear): + in_features = state_dict[module_prefix + "weight"].shape[1] + out_features = state_dict[module_prefix + "weight"].shape[0] + if scale_weights is not None: + scale_key = f"{module_prefix}scale_weight" + + with init_empty_weights(): + model._modules[name] = Fp8Linear( + in_features, + out_features, + module.bias is not None, + compute_dtype=compute_dtype, + scale_weight=scale_weights.get(scale_key) if scale_weights else None + ) + #set_lora_params(model._modules[name], patches, module_prefix) + model._modules[name].source_cls = type(module) + # Force requires_grad to False to avoid unexpected errors + model._modules[name].requires_grad_(False) + + return model + +def set_lora_params(module, patches, module_prefix=""): + # Recursively set lora_diffs and lora_strengths for all Fp8Linear layers + for name, child in module.named_children(): + child_prefix = (f"{module_prefix}{name}.") + set_lora_params(child, patches, child_prefix) + if isinstance(module, Fp8Linear): + key = f"diffusion_model.{module_prefix}weight" + patch = patches.get(key, []) + #print(f"Processing LoRA patches for {key}: {len(patch)} patches found") + if len(patch) != 0: + lora_diffs = [] + for p in patch: + lora_obj = p[1] + if "head" in key: + continue # For now skip LoRA for head layers + elif hasattr(lora_obj, "weights"): + lora_diffs.append(lora_obj.weights) + elif isinstance(lora_obj, tuple) and lora_obj[0] == "diff": + lora_diffs.append(lora_obj[1]) + else: + continue + lora_strengths = [p[0] for p in patch] + module.lora = (lora_diffs, lora_strengths) + module.step = 0 # Initialize step for LoRA scheduling + + +class Fp8Linear(nn.Linear): + def __init__( + self, + in_features, + out_features, + bias=False, + compute_dtype=None, + device=None, + scale_weight=None + ) -> None: + super().__init__(in_features, out_features, bias, device) + self.compute_dtype = compute_dtype + self.lora = None + self.step = 0 + self.scale_weight = scale_weight + + def forward(self, input): + weight = self.weight.to(input.dtype) + bias = self.bias.to(input.dtype) if self.bias is not None else None + if self.scale_weight is not None: + scale_weight = self.scale_weight.to(input.device) + if weight.numel() < input.numel(): + weight = weight * scale_weight + else: + input = input * scale_weight + + if self.lora is not None: + weight = self.apply_lora(weight).to(input.dtype) + + return torch.nn.functional.linear(input, weight, bias) + + @torch.compiler.disable() + def apply_lora(self, weight): + for lora_diff, lora_strength in zip(self.lora[0], self.lora[1]): + if isinstance(lora_strength, list): + lora_strength = lora_strength[self.step] + if lora_strength == 0.0: + continue + elif lora_strength == 0.0: + continue + patch_diff = torch.mm( + lora_diff[0].flatten(start_dim=1).to(weight.device), + lora_diff[1].flatten(start_dim=1).to(weight.device) + ).reshape(weight.shape) + alpha = lora_diff[2] / lora_diff[1].shape[0] if lora_diff[2] is not None else 1.0 + scale = lora_strength * alpha + weight = weight.add(patch_diff, alpha=scale) + return weight + +def remove_lora_from_module(module): + for name, submodule in module.named_modules(): + submodule.lora = None \ No newline at end of file diff --git a/nodes.py b/nodes.py index b66f59d..d836856 100644 --- a/nodes.py +++ b/nodes.py @@ -8,7 +8,7 @@ import hashlib from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from .wanvideo.modules.model import rope_params -from .fp8_optimization import convert_linear_with_lora_and_scale, remove_lora_from_module +from .fp8_optimization_v2 import remove_lora_from_module, set_lora_params as set_lora_params_fp8 from .wanvideo.schedulers import get_scheduler, get_sampling_sigmas, retrieve_timesteps, scheduler_list from .gguf.gguf import set_lora_params from .multitalk.multitalk import timestep_transform, add_noise @@ -1513,9 +1513,7 @@ class WanVideoSampler: log.info(f"Using {len(patcher.patches)} LoRA weight patches for WanVideo model") if not merge_loras and fp8_matmul: raise NotImplementedError("FP8 matmul with unmerged LoRAs is not supported") - convert_linear_with_lora_and_scale(transformer, patches=patcher.patches, scale_weight_keys=scale_weights) - elif patch_linear: - convert_linear_with_lora_and_scale(transformer, scale_weight_keys=scale_weights) + set_lora_params_fp8(transformer, patcher.patches) else: remove_lora_from_module(transformer) diff --git a/nodes_model_loading.py b/nodes_model_loading.py index 421f63a..a2b064e 100644 --- a/nodes_model_loading.py +++ b/nodes_model_loading.py @@ -1036,7 +1036,10 @@ class WanVideoModelLoader: for k, v in sd.items(): if k.endswith(".scale_weight"): scale_weights[k] = v - + if not merge_loras: + from .fp8_optimization_v2 import _replace_linear + transformer = _replace_linear(transformer, base_dtype, sd, scale_weights=scale_weights) + if "fp8_e4m3fn" in quantization: dtype = torch.float8_e4m3fn elif "fp8_e5m2" in quantization: @@ -1182,7 +1185,7 @@ class WanVideoModelLoader: patcher.model.is_patched = True - patch_linear = (True if "scaled" in quantization or not merge_loras else False) + patch_linear = (True if "scaled" in quantization or (lora is not None and not merge_loras) else False) if "fast" in quantization: if lora is not None and not merge_loras: