From 7b7c014142bcd588bd0c4a19f62dc171030ff607 Mon Sep 17 00:00:00 2001 From: Yao Chi Date: Mon, 26 Aug 2024 00:02:14 +0800 Subject: [PATCH] refine flux.py --- comfy/text_encoders/flux.py | 30 ++---------------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py index 7332513eb..0ddb313b7 100644 --- a/comfy/text_encoders/flux.py +++ b/comfy/text_encoders/flux.py @@ -62,7 +62,8 @@ class FluxClipModel(torch.nn.Module): if "text_model.encoder.layers.1.mlp.fc1.weight" in sd: return self.clip_l.load_sd(sd) else: - quantize_model(self.t5xxl.transformer, sd) + from bizyairenhancer import clip_quantize_model + clip_quantize_model(self.t5xxl.transformer, sd) return self.t5xxl.load_sd(sd) def flux_clip(dtype_t5=None): @@ -70,30 +71,3 @@ def flux_clip(dtype_t5=None): def __init__(self, device="cpu", dtype=None, model_options={}): super().__init__(dtype_t5=dtype_t5, device=device, dtype=dtype, model_options=model_options) return FluxClipModel_ - -def quantize_model(model: torch.nn.Module, new_state_dict): - for name, module in model.named_modules(): - device = next(module.parameters()).device - if not isinstance(module, torch.nn.Linear): - continue - weight = new_state_dict[f"{name}.weight"].to(device) - if module.bias is not None: - module.bias.data = module.bias.data.to(torch.float16).to(device) - qweight, scale = fp8_quantize(weight) - module.weight.data = module.weight.data.to(torch.float8_e4m3fn) - module.weight.data.copy_(qweight.data) - module.register_buffer("scale", scale.to(device)) - new_state_dict.pop(f"{name}.weight") - - -def fp8_quantize(weight, qdtype=torch.float8_e4m3fn): - device = weight.device - finfo = torch.finfo(qdtype) - - scale = finfo.max / weight.abs().max().clamp(min=1e-12) - - qweight = (weight * scale).clamp(min=finfo.min, max=finfo.max) - - qweight = qweight.to(qdtype) - scale = scale.float().reciprocal() - return qweight, scale \ No newline at end of file