diff --git a/comfy/model_management.py b/comfy/model_management.py index f7630b62b..91e692ba2 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -691,9 +691,7 @@ def text_encoder_offload_device(): return torch.device("cpu") def text_encoder_device(): - if "," in args.cuda_device: - return torch.device("cuda", 1) - elif args.gpu_only: + if args.gpu_only: return get_torch_device() elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM: if should_use_fp16(prioritize_performance=False): diff --git a/comfy/ops.py b/comfy/ops.py index 43ed55adb..5fee97600 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -307,3 +307,15 @@ def pick_operations(weight_dtype, compute_dtype, load_device=None): if comfy.model_management.supports_fp8_compute(load_device): return fp8_ops return manual_cast + +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) + # Return both float8 data and the inverse scale (as float), + # as both required as inputs to torch._scaled_mm + qweight = qweight.to(qdtype) + scale = scale.float().reciprocal() + return qweight, scale \ No newline at end of file diff --git a/comfy/sd.py b/comfy/sd.py index 1e1a6594b..fa0d685e7 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -650,6 +650,7 @@ def load_diffusion_model_state_dict(sd, model_options={}): #load unet in diffuse model_config.custom_operations = model_options.get("custom_operations", None) model = model_config.get_model(new_sd, "") model = model.to(offload_device) + quantize_model(model, new_sd) model.load_model_weights(new_sd, "") left_over = sd.keys() if len(left_over) > 0: @@ -695,3 +696,36 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m sd[k] = t.contiguous() comfy.utils.save_torch_file(sd, output_path, metadata=metadata) + +def quantize_model(model: torch.nn.Module, new_state_dict): + for name, module in model.diffusion_model.named_modules(): + if not isinstance(module, torch.nn.Linear): + continue + if not (name.startswith("double_blocks") or name.startswith("single_blocks")): + print(name) + continue + weight = new_state_dict[f"{name}.weight"] + # if module.bias is not None: + # module.bias.data = module.bias.data.to(torch.bfloat16) + 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) + new_state_dict.pop(f"{name}.weight") + + +def fp8_quantize(weight, qdtype=torch.float8_e4m3fn): + device = weight.device + # weight, scale = quant_weights(weight, torch.int8, False) + finfo = torch.finfo(qdtype) + # Calculate the scale as dtype max divided by absmax + scale = finfo.max / weight.abs().max().clamp(min=1e-12) + # scale and clamp the tensor to bring it to + # the representative range of float8 data type + # (as default cast is unsaturated) + qweight = (weight * scale).clamp(min=finfo.min, max=finfo.max) + # Return both float8 data and the inverse scale (as float), + # as both required as inputs to torch._scaled_mm + qweight = qweight.to(qdtype) + scale = scale.float().reciprocal() + return qweight, scale diff --git a/comfy/text_encoders/flux.py b/comfy/text_encoders/flux.py index 91d1f249d..7332513eb 100644 --- a/comfy/text_encoders/flux.py +++ b/comfy/text_encoders/flux.py @@ -13,7 +13,7 @@ class T5XXLModel(sd1_clip.SDClipModel): class T5XXLTokenizer(sd1_clip.SDTokenizer): def __init__(self, embedding_directory=None, tokenizer_data={}): tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_tokenizer") - super().__init__(tokenizer_path, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256) + super().__init__(tokenizer_path, pad_with_end=False, embedding_size=4096, embedding_key='t5xxl', tokenizer_class=T5TokenizerFast, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=768) class FluxTokenizer: @@ -62,6 +62,7 @@ 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) return self.t5xxl.load_sd(sd) def flux_clip(dtype_t5=None): @@ -69,3 +70,30 @@ 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