Create linfusion.py

src/linfusion/linfusion.py

@@ -0,0 +1,116 @@
+from diffusers.models.attention_processor import Attention
+from diffusers import ModelMixin, ConfigMixin
+import functools
+
+from .attention import GeneralizedLinearAttention
+
+
+model_dict = {
+    "runwayml/stable-diffusion-v1-5": "Yuanshi/LinFusion-1-5",
+    "stablediffusionapi/realistic-vision-v51": "Yuanshi/LinFusion-1-5",
+    "Lykon/dreamshaper-8": "Yuanshi/LinFusion-1-5",
+}
+
+
+def replace_submodule(model, module_name, new_submodule):
+    path, attr = module_name.rsplit(".", 1)
+    parent_module = functools.reduce(getattr, path.split("."), model)
+    setattr(parent_module, attr, new_submodule)
+
+
+class LinFusion(ModelMixin, ConfigMixin):
+    def __init__(self, modules_list, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+        self.modules_dict = {}
+        self.register_to_config(modules_list=modules_list)
+
+        for i, attention_config in enumerate(modules_list):
+            dim_n = attention_config["dim_n"]
+            heads = attention_config["heads"]
+            projection_mid_dim = attention_config["projection_mid_dim"]
+            linear_attention = GeneralizedLinearAttention(
+                query_dim=dim_n,
+                out_dim=dim_n,
+                dim_head=dim_n // heads,
+                projection_mid_dim=projection_mid_dim,
+            )
+            self.add_module(f"{i}", linear_attention)
+            self.modules_dict[attention_config["module_name"]] = linear_attention
+
+    @classmethod
+    def get_default_config(
+        cls,
+        pipeline=None,
+        unet=None,
+    ):
+        """
+        Get the default configuration for the LinFusion model.
+        (The `projection_mid_dim` is same as the `query_dim` by default.)
+        """
+        assert unet is not None or pipeline.unet is not None
+        unet = unet or pipeline.unet
+        modules_list = []
+        for module_name, module in unet.named_modules():
+            if not isinstance(module, Attention):
+                continue
+            if "attn1" not in module_name:
+                continue
+            dim_n = module.to_q.weight.shape[0]
+            # modules_list.append((module_name, dim_n, module.heads))
+            modules_list.append(
+                {
+                    "module_name": module_name,
+                    "dim_n": dim_n,
+                    "heads": module.heads,
+                    "projection_mid_dim": None,
+                }
+            )
+        return {"modules_list": modules_list}
+
+    @classmethod
+    def construct_for(
+        cls,
+        pipeline=None,
+        unet=None,
+        load_pretrained=True,
+        pretrained_model_name_or_path=None,
+    ) -> "LinFusion":
+        """
+        Construct a LinFusion object for the given pipeline.
+        """
+        assert unet is not None or pipeline.unet is not None
+        unet = unet or pipeline.unet
+        if load_pretrained:
+            # Load from pretrained
+            pipe_name_path = pipeline._internal_dict._name_or_path
+            if not pretrained_model_name_or_path:
+                pretrained_model_name_or_path = model_dict.get(pipe_name_path, None)
+                if pretrained_model_name_or_path:
+                    print(
+                        f"Matching LinFusion '{pretrained_model_name_or_path}' for pipeline '{pipe_name_path}'."
+                    )
+                else:
+                    raise Exception(
+                        f"LinFusion not found for pipeline [{pipe_name_path}], please provide the path."
+                    )
+            linfusion = (
+                LinFusion.from_pretrained(pretrained_model_name_or_path)
+                .to(pipeline.device)
+                .to(pipeline.dtype)
+            )
+        else:
+            # Create from scratch without pretrained parameters
+            default_config = LinFusion.get_default_config(pipeline)
+            linfusion = (
+                LinFusion(**default_config).to(pipeline.device).to(pipeline.dtype)
+            )
+        linfusion.mount_to(unet)
+        return linfusion
+
+    def mount_to(self, unet) -> None:
+        """
+        Mounts the modules in the `modules_dict` to the given `pipeline`.
+        """
+        for module_name, module in self.modules_dict.items():
+            replace_submodule(unet, module_name, module)