File size: 575 Bytes
2e1a430 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 | from diffsynth import load_state_dict
import torch
from safetensors.torch import save_file
from tqdm import tqdm
def dequantize(source_path, target_path, device="cuda", torch_dtype=torch.bfloat16):
sd = load_state_dict(source_path, device=device)
for k in tqdm([k for k in sd if k.endswith(".weight_scale")]):
weight_key = k[:-13] + ".weight"
weight = sd.pop(weight_key).to(torch_dtype)
scale = sd.pop(k).to(torch_dtype).unsqueeze(1)
sd[weight_key] = weight * scale
if target_path is not None:
save_file(sd, target_path)
|