reset and revert

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):

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

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

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