xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-13 23:27:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

refine flux.py

Browse Source
This commit is contained in:
Yao Chi 2024-08-26 00:02:14 +08:00
parent 9d2fcc7d06
commit 7b7c014142
1 changed files with 2 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
comfy/text_encoders/flux.py
View File

@ -62,7 +62,8 @@ class FluxClipModel(torch.nn.Module):
if "text_model.encoder.layers.1.mlp.fc1.weight" in sd: if "text_model.encoder.layers.1.mlp.fc1.weight" in sd:
return self.clip_l.load_sd(sd) return self.clip_l.load_sd(sd)
else: 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) return self.t5xxl.load_sd(sd)
def flux_clip(dtype_t5=None): 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={}): def __init__(self, device="cpu", dtype=None, model_options={}):
super().__init__(dtype_t5=dtype_t5, device=device, dtype=dtype, model_options=model_options) super().__init__(dtype_t5=dtype_t5, device=device, dtype=dtype, model_options=model_options)
return FluxClipModel_ 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