Support Diffusers Flux LoRA (#95)

* support diffusers lora

* support flux diffusers lora

Author: tenderness-git

diffsynth_engine/pipelines/base.py

@@ -196,8 +196,14 @@ def prepare_latents(
             # if you have any questions about this, please ask @dizhipeng.dzp for more details
             latents = latents * sigmas[0] / ((sigmas[0] ** 2 + 1) ** 0.5)
             init_latents = latents.clone()
-        sigmas, timesteps = sigmas.to(device=self.device, dtype=self.dtype), timesteps.to(device=self.device, dtype=self.dtype)
-        init_latents, latents = init_latents.to(device=self.device, dtype=self.dtype), latents.to(device=self.device, dtype=self.dtype)
+        sigmas, timesteps = (
+            sigmas.to(device=self.device, dtype=self.dtype),
+            timesteps.to(device=self.device, dtype=self.dtype),
+        )
+        init_latents, latents = (
+            init_latents.to(device=self.device, dtype=self.dtype),
+            latents.to(device=self.device, dtype=self.dtype),
+        )
         return init_latents, latents, sigmas, timesteps
 
     def eval(self):

diffsynth_engine/pipelines/flux_image.py

@@ -163,9 +163,8 @@ def _from_diffusers(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str,
         dit_dict = {}
         for key, param in lora_state_dict.items():
             origin_key = key
-            if ".alpha" not in key:
+            if "lora_A.weight" not in key or "lora_up.weight" not in key:
                 continue
-            key = key.replace(".alpha", ".weight")
             key = key.replace("transformer.", "")
             if "single_transformer_blocks" in key:  # transformer.single_transformer_blocks.0.attn.to_k.weight
                 key = key.replace(
@@ -208,10 +207,17 @@ def _from_diffusers(self, lora_state_dict: Dict[str, torch.Tensor]) -> Dict[str,
             else:
                 raise ValueError(f"Unsupported key: {key}")
             lora_args = {}
-            lora_args["alpha"] = param
-            lora_args["up"] = lora_state_dict[origin_key.replace(".alpha", ".lora_up.weight")]
-            lora_args["down"] = lora_state_dict[origin_key.replace(".alpha", ".lora_down.weight")]
+            lora_args["up"] = param
+            lora_args["down"] = lora_state_dict[
+                origin_key.replace("lora_A.weight", "lora_B.weight").replace("lora_up.weight", "lora_down.weight")
+            ]
             lora_args["rank"] = lora_args["up"].shape[1]
+            alpha_key = origin_key.replace("lora_A.weight", "alpha").replace("lora_up.weight", "alpha")
+            if alpha_key in lora_state_dict:
+                alpha = lora_state_dict[alpha_key]
+            else:
+                alpha = lora_args["rank"]  # 如果alpha不存在，则取alpha/rank = 1
+            lora_args["alpha"] = alpha
             key = key.replace(".weight", "")
             dit_dict[key] = lora_args
         return {"dit": dit_dict}

diffsynth_engine/tools/flux_inpainting_tool.py

@@ -1,5 +1,5 @@
 from diffsynth_engine import fetch_model, FluxControlNet, ControlNetParams, FluxImagePipeline
-from typing import List, Tuple, Optional
+from typing import List, Tuple, Optional, Callable
 from PIL import Image
 import torch
 
@@ -34,6 +34,7 @@ def __call__(
         inpainting_scale: float = 0.9,
         seed: int = 42,
         num_inference_steps: int = 20,
+        progress_callback: Optional[Callable] = None,  # def progress_callback(current, total, status)
     ):
         assert image.size == mask.size
         return self.pipe(
@@ -49,4 +50,5 @@ def __call__(
                 mask=mask,
                 scale=inpainting_scale,
             ),
+            progress_callback=progress_callback,
         )

tests/common/test_case.py

@@ -1,6 +1,5 @@
 import unittest
 import os
-import time
 import numpy as np
 import torch
 from pathlib import Path
@@ -80,7 +79,7 @@ def assertImageEqualAndSaveFailed(self, input_image: Image.Image, expect_image_p
             self.assertImageEqual(input_image, expect_image, threshold=threshold)
         except Exception as e:
             name = expect_image_path.split("/")[-1]
-            input_image.save(f"save_{time.time()}_{name}")
+            input_image.save(f"{name}")
             raise e
 
 

tests/data/expect/flux/flux_diffusers_lora.png

@@ -1,5 +1,6 @@
 import unittest
 
+
 from tests.common.test_case import ImageTestCase
 from diffsynth_engine.pipelines import FluxImagePipeline, FluxModelConfig
 from diffsynth_engine import fetch_model
@@ -47,6 +48,21 @@ def test_unfused_lora(self):
         self.pipe.unload_loras()
         self.assertImageEqualAndSaveFailed(image, "flux/flux_lora.png", threshold=0.98)
 
+    def test_diffusers_lora_patch(self):
+        lora_model_path = fetch_model(
+            "InstantX/FLUX.1-dev-LoRA-Ghibli", revision="master", path="ghibli_style.safetensors"
+        )
+        self.pipe.load_loras([(lora_model_path, 0.8)], fused=True, save_original_weight=True)
+        image = self.pipe(
+            prompt="ghibli style, a shepherd boy floating on a wooly cloud-whale, holding a glowing dandelion staff to guide sheep-shaped cumulus, miniature storm clouds grazing nearby, his patched jacket flapping in high-altitude winds, aurora-like ribbons in peach and lavender stretching across the sky",
+            width=960,
+            height=1280,
+            num_inference_steps=24,
+            seed=42,
+        )
+        self.pipe.unload_loras()
+        self.assertImageEqualAndSaveFailed(image, "flux/flux_diffusers_lora.png", threshold=0.99)
+
 
 class TestFLUXGGUF(ImageTestCase):
     @classmethod