[WIP] Add Flax LoRA Support to Dreambooth

I saw @patrickvonplaten is working on LoRA support for the non-Flax Dreambooth.
We've been taking a stab at implementing LoRA support for TPUs, taking example from
the patching method used by @cloneofsimo in cloneofsimo/lora.

I've got it successfully patching and training, but the output is currently no good.
I'm reaching the end of the time I have allocated for this—might pick it up in the future,
but for now I'm putting this up in case anyone finds it useful!

Author: yasyf

examples/dreambooth/train_dreambooth_flax.py

@@ -22,6 +22,7 @@
     FlaxStableDiffusionPipeline,
     FlaxUNet2DConditionModel,
 )
+from diffusers.experimental.lora.linear_with_lora_flax import FlaxLora
 from diffusers.pipelines.stable_diffusion import FlaxStableDiffusionSafetyChecker
 from diffusers.utils import check_min_version
 from flax import jax_utils
@@ -97,6 +98,7 @@ def parse_args():
             " class_data_dir, additional images will be sampled with class_prompt."
         ),
     )
+    parser.add_argument("--lora", action="store_true", help="Use LoRA (https://arxiv.org/abs/2106.09685)")
     parser.add_argument(
         "--output_dir",
         type=str,
@@ -444,9 +446,6 @@ def collate_fn(examples):
     vae, vae_params = FlaxAutoencoderKL.from_pretrained(
         args.pretrained_model_name_or_path, subfolder="vae", dtype=weight_dtype
     )
-    unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
-        args.pretrained_model_name_or_path, subfolder="unet", dtype=weight_dtype
-    )
 
     # Optimization
     if args.scale_lr:
@@ -467,6 +466,22 @@ def collate_fn(examples):
         adamw,
     )
 
+    if args.lora:
+        unet, unet_params = FlaxLora(FlaxUNet2DConditionModel).from_pretrained(
+            args.pretrained_model_name_or_path,
+            subfolder="unet",
+            dtype=weight_dtype,
+            revision=args.revision,
+        )
+        optimizer = optax.masked(optimizer, mask=unet.get_mask)
+    else:
+        unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
+            args.pretrained_model_name_or_path,
+            subfolder="unet",
+            dtype=weight_dtype,
+            revision=args.revision,
+        )
+
     unet_state = train_state.TrainState.create(apply_fn=unet.__call__, params=unet_params, tx=optimizer)
     text_encoder_state = train_state.TrainState.create(
         apply_fn=text_encoder.__call__, params=text_encoder.params, tx=optimizer

src/diffusers/experimental/lora/linear_with_lora_flax.py

@@ -0,0 +1,111 @@
+import copy
+from collections import defaultdict
+from typing import Dict, List, Type, Union, cast
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+from diffusers.modeling_flax_utils import FlaxModelMixin
+from flax.core.frozen_dict import FrozenDict
+from flax.traverse_util import flatten_dict, unflatten_dict
+
+
+class FlaxLinearWithLora(nn.Module):
+    out_features: int
+    rank: int = 5
+    in_features: int = 1
+    scale: float = 1.0
+    use_bias: bool = True
+
+    def setup(self):
+        self.linear = nn.Dense(features=self.out_features, use_bias=self.use_bias)
+        self.lora_up = nn.Dense(features=self.out_features, use_bias=False)
+        self.lora_down = nn.Dense(features=4, use_bias=False)
+
+    def init_weights(self, rng: jax.random.PRNGKey) -> FrozenDict:
+        return self.init(rng, jnp.zeros((self.in_features, self.out_features)))
+
+    def __call__(self, input):
+        return self.linear(input) + self.lora_up(self.lora_down(input)) * self.scale
+
+
+class FlaxLoraBase(nn.Module):
+    @staticmethod
+    def _get_children(model: nn.Module) -> Dict[str, nn.Module]:
+        model._try_setup(shallow=True)
+        return {k: v for k, v in model._state.children.items() if isinstance(v, nn.Module)}
+
+    @staticmethod
+    def _wrap_dense(params: dict, parent: nn.Module, model: Union[nn.Dense, nn.Module], name: str):
+        if not isinstance(model, nn.Dense):
+            return params, {}
+
+        params_to_optimize = defaultdict(dict)
+
+        parent._in_setup = True
+        lora = FlaxLinearWithLora(
+            out_features=model.features,
+            use_bias=model.use_bias,
+            name=name,
+            parent=parent,
+        )
+
+        lora_params = lora.init_weights(jax.random.PRNGKey(0)).unfreeze()["params"]
+        lora_params["linear"] = params
+        lora = lora.bind({"params": lora_params})
+
+        for k, v in parent.__dict__.items():
+            if isinstance(v, nn.Module) and v.name == name:
+                setattr(model.parent, k, lora)
+
+        parent._in_setup = False
+
+        for n in ["lora_up", "lora_down"]:
+            params_to_optimize[n] = {k: True for k in lora_params[n].keys()}
+        params_to_optimize["linear"] = {k: False for k in lora_params["linear"].keys()}
+
+        return lora_params, dict(params_to_optimize)
+
+    @staticmethod
+    def inject(
+        params: Union[dict, FrozenDict],
+        model: nn.Module,
+        targets: List[str],
+        is_target: bool = False,
+    ):
+        params = params.unfreeze() if isinstance(params, FrozenDict) else copy.copy(params)
+        params_to_optimize = {}
+
+        for name, child in FlaxLoraBase._get_children(model).items():
+            if is_target:
+                results = FlaxLoraBase._wrap_dense(params.get(name, {}), model, child, name)
+            elif child.__class__.__name__ in targets:
+                results = FlaxLoraBase.inject(params.get(name, {}), child, targets=targets, is_target=True)
+            else:
+                results = FlaxLoraBase.inject(params.get(name, {}), child, targets=targets)
+
+            params[name], params_to_optimize[name] = results
+
+        return params, params_to_optimize
+
+
+def FlaxLora(model: Type[nn.Module], targets=["FlaxAttentionBlock"]):
+    class _FlaxLora(model):
+        def setup(self):
+            super().setup()
+            params = cast(FlaxModelMixin, self).init_weights(jax.random.PRNGKey(0))
+            FlaxLoraBase.inject(params, self, targets=targets)
+
+        @classmethod
+        def from_pretrained(cls, *args, **kwargs):
+            instance, params = cast(Type[FlaxModelMixin], model).from_pretrained(*args, **kwargs)
+            params, mask = FlaxLoraBase.inject(params, instance, targets=targets)
+            mask_values = flatten_dict(mask)
+            instance.get_mask = lambda params: unflatten_dict(
+                {k: mask_values.get(k, False) for k in flatten_dict(params, keep_empty_nodes=True).keys()}
+            )
+            return instance, params
+
+    _FlaxLora.__name__ = f"{model.__name__}Lora"
+
+    return _FlaxLora

src/diffusers/experimental/lora/test_lora.py

@@ -0,0 +1,26 @@
+import os
+import pdb
+
+import optax
+from diffusers import FlaxUNet2DConditionModel
+from diffusers.experimental.lora.linear_with_lora_flax import FlaxLora
+from flax.training import train_state
+from jax.config import config
+from jax.experimental.compilation_cache import compilation_cache as cc
+
+
+config.update("jax_traceback_filtering", "off")
+cc.initialize_cache(os.path.expanduser("~/.cache/jax/compilation_cache"))
+
+if __name__ == "__main__":
+    unet, unet_params = FlaxLora(FlaxUNet2DConditionModel).from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        subfolder="unet",
+        revision="flax",
+    )
+    get_mask = unet.get_mask
+
+    optimizer = optax.masked(optax.adamw(1e-6), mask=get_mask)
+    unet_state = train_state.TrainState.create(apply_fn=unet.__call__, params=unet_params, tx=optimizer)
+
+    pdb.set_trace()