131 lines
5.4 KiB
Python
131 lines
5.4 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
from diffusers.quantizers.gguf.utils import GGUFParameter, dequantize_gguf_tensor
|
|
import gguf
|
|
from diffusers.utils import is_accelerate_available
|
|
from contextlib import nullcontext
|
|
from ..utils import log
|
|
if is_accelerate_available():
|
|
from accelerate import init_empty_weights
|
|
|
|
def load_gguf(model_path):
|
|
from gguf import GGUFReader
|
|
reader = GGUFReader(model_path)
|
|
parsed_parameters = {}
|
|
for tensor in reader.tensors:
|
|
# if the tensor is a torch supported dtype do not use GGUFParameter
|
|
is_gguf_quant = tensor.tensor_type not in [gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16]
|
|
meta_tensor = torch.empty(tensor.data.shape, dtype=torch.from_numpy(np.empty(0, dtype=tensor.data.dtype)).dtype, device='meta')
|
|
parsed_parameters[tensor.name] = GGUFParameter(meta_tensor, quant_type=tensor.tensor_type) if is_gguf_quant else meta_tensor
|
|
return parsed_parameters, reader
|
|
|
|
#based on https://github.com/huggingface/diffusers/blob/main/src/diffusers/quantizers/gguf/utils.py
|
|
def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modules_to_not_convert=[], patches=None):
|
|
def _should_convert_to_gguf(state_dict, prefix):
|
|
weight_key = prefix + "weight"
|
|
return weight_key in state_dict and isinstance(state_dict[weight_key], GGUFParameter)
|
|
|
|
has_children = list(model.children())
|
|
if not has_children:
|
|
return
|
|
|
|
for name, module in model.named_children():
|
|
module_prefix = prefix + name + "."
|
|
_replace_with_gguf_linear(module, compute_dtype, state_dict, module_prefix, modules_to_not_convert, patches)
|
|
|
|
if (
|
|
isinstance(module, nn.Linear)
|
|
and not isinstance(module, GGUFLinear)
|
|
and _should_convert_to_gguf(state_dict, module_prefix)
|
|
and name not in modules_to_not_convert
|
|
):
|
|
in_features = state_dict[module_prefix + "weight"].shape[1]
|
|
out_features = state_dict[module_prefix + "weight"].shape[0]
|
|
|
|
ctx = init_empty_weights if is_accelerate_available() else nullcontext
|
|
with ctx():
|
|
model._modules[name] = GGUFLinear(
|
|
in_features,
|
|
out_features,
|
|
module.bias is not None,
|
|
compute_dtype=compute_dtype
|
|
)
|
|
|
|
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_gguf(module, patches, module_prefix=""):
|
|
# Recursively set lora_diffs and lora_strengths for all GGUFLinear layers
|
|
for name, child in module.named_children():
|
|
child_prefix = (f"{module_prefix}{name}.")
|
|
set_lora_params_gguf(child, patches, child_prefix)
|
|
if isinstance(module, GGUFLinear):
|
|
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 GGUFLinear(nn.Linear):
|
|
def __init__(
|
|
self,
|
|
in_features,
|
|
out_features,
|
|
bias=False,
|
|
compute_dtype=None,
|
|
device=None,
|
|
) -> None:
|
|
super().__init__(in_features, out_features, bias, device)
|
|
self.compute_dtype = compute_dtype
|
|
self.lora = None
|
|
self.step = 0
|
|
|
|
def forward(self, inputs):
|
|
weight = self.dequantize_without_compile()
|
|
weight = weight.to(self.compute_dtype)
|
|
bias = self.bias.to(self.compute_dtype) if self.bias is not None else None
|
|
|
|
if hasattr(self, "lora") and self.lora is not None:
|
|
weight = self.apply_lora(weight, self.step).to(self.compute_dtype)
|
|
|
|
output = torch.nn.functional.linear(inputs, weight, bias)
|
|
return output
|
|
|
|
@torch.compiler.disable()
|
|
def dequantize_without_compile(self):
|
|
return dequantize_gguf_tensor(self.weight)
|
|
|
|
@torch.compiler.disable()
|
|
def apply_lora(self, weight, step=None):
|
|
for lora_diff, lora_strength in zip(self.lora[0], self.lora[1]):
|
|
if isinstance(lora_strength, list):
|
|
lora_strength = lora_strength[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 |