Support fp8_scaled models and allow running LoRAs unmerged on other models as well

fp8_optimization.py

@@ -27,13 +27,113 @@ def fp8_linear_forward(cls, original_dtype, input):
             return cls.original_forward(input.to(original_dtype))
     else:
         return cls.original_forward(input)
+    
+def fp8_scaled_linear_forward(cls, original_dtype, input):
+    weight = cls.weight.to(original_dtype)
+    scale_weight = cls.scale_weight.to(input.device)
+    bias = cls.bias.to(original_dtype) if cls.bias is not None else None
+
+    if weight.numel() < input.numel():
+        weight = weight * scale_weight
+    else:
+        input = input * scale_weight
+
+    lora = getattr(cls, "lora", None)
+    if lora is not None:
+        for lora_diff, lora_strength in zip(lora[0], lora[1]):
+            patch_diff = torch.mm(
+                lora_diff[0].flatten(start_dim=1).to(weight.device),
+                lora_diff[1].flatten(start_dim=1).to(weight.device)
+            ).reshape(weight.shape)
+            alpha = lora_diff[2] / lora_diff[1].shape[0] if lora_diff[2] is not None else 1.0
+            scale = lora_strength * alpha
+            weight = weight.add(patch_diff, alpha=scale).to(original_dtype)
+
+    return torch.nn.functional.linear(input, weight, bias)
+
+def linear_with_lora_forward(cls, original_dtype, input):
+    weight = cls.weight.to(original_dtype)
+    bias = cls.bias.to(original_dtype) if cls.bias is not None else None
+
+    lora = getattr(cls, "lora", None)
+    if lora is not None:
+        for lora_diff, lora_strength in zip(lora[0], lora[1]):
+            patch_diff = torch.mm(
+                lora_diff[0].flatten(start_dim=1).to(weight.device),
+                lora_diff[1].flatten(start_dim=1).to(weight.device)
+            ).reshape(weight.shape)
+            alpha = lora_diff[2] / lora_diff[1].shape[0] if lora_diff[2] is not None else 1.0
+            scale = lora_strength * alpha
+            weight = weight.add(patch_diff, alpha=scale).to(original_dtype)
+
+    return torch.nn.functional.linear(input, weight, bias)
+ 
 
 def convert_fp8_linear(module, original_dtype, params_to_keep={}):
     setattr(module, "fp8_matmul_enabled", True)
    
-    for name, module in module.named_modules():
+    for name, submodule in module.named_modules():
         if not any(keyword in name for keyword in params_to_keep):
-            if isinstance(module, nn.Linear):
-                original_forward = module.forward
-                setattr(module, "original_forward", original_forward)
-                setattr(module, "forward", lambda input, m=module: fp8_linear_forward(m, original_dtype, input))
+            if isinstance(submodule, nn.Linear):
+                original_forward = submodule.forward
+                setattr(submodule, "original_forward", original_forward)
+                setattr(submodule, "forward", lambda input, m=submodule: fp8_linear_forward(m, original_dtype, input))
+
+def convert_fp8_scaled_linear(module, sd, original_dtype, params_to_keep={}, patches=None):
+    setattr(module, "fp8_scaled_enabled", True)
+   
+    for name, submodule in module.named_modules():
+        if not any(keyword in name for keyword in params_to_keep):
+            scale_key = f"{name}.scale_weight"
+            has_scale = scale_key in sd
+            weight = getattr(submodule, 'weight', None)
+            has_fp8_weight = weight is not None and weight.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+            if has_scale:
+                setattr(submodule, "scale_weight", sd[scale_key])
+            
+            if patches is not None:
+                patch_key = f"diffusion_model.{name}.weight"
+                patch = patches.get(patch_key, [])
+                #print("Patches for", patch_key, ":", patch)
+                if len(patch) != 0:
+                    lora_diffs = []
+                    for p in patch:
+                        lora_obj = p[1]
+                        if hasattr(lora_obj, "weights"):
+                            lora_diffs.append(lora_obj.weights)
+                        elif isinstance(lora_obj, tuple) and lora_obj[0] == "diff":
+                            lora_diffs.append(lora_obj[1])
+                        else:
+                            continue
+                        
+                    lora_strengths = [p[0] for p in patch]
+                    lora = (lora_diffs, lora_strengths)
+                    setattr(submodule, "lora", lora)
+
+            if isinstance(submodule, nn.Linear) and (has_scale or has_fp8_weight):
+                original_forward = submodule.forward
+                setattr(submodule, "original_forward", original_forward)
+                setattr(submodule, "forward", lambda input, m=submodule: fp8_scaled_linear_forward(m, original_dtype, input))
+
+def convert_linear_with_lora(module, original_dtype, patches=None):
+    for name, submodule in module.named_modules():
+        if isinstance(submodule, nn.Linear):
+            patch_key = f"diffusion_model.{name}.weight"
+            patch = patches.get(patch_key, [])
+            if len(patch) != 0:
+                lora_diffs = []
+                for p in patch:
+                    lora_obj = p[1]
+                    if hasattr(lora_obj, "weights"):
+                        lora_diffs.append(lora_obj.weights)
+                    elif isinstance(lora_obj, tuple) and lora_obj[0] == "diff":
+                        lora_diffs.append(lora_obj[1])
+                    else:
+                        continue
+                    
+                lora_strengths = [p[0] for p in patch]
+                lora = (lora_diffs, lora_strengths)
+                setattr(submodule, "lora", lora)
+                # original_forward = submodule.forward
+                # setattr(submodule, "original_forward", original_forward)
+                setattr(submodule, "forward", lambda input, m=submodule: linear_with_lora_forward(m, original_dtype, input))