ella adapter implementation. tested with sd1.5 model

Author: ily-R

pyproject.toml

@@ -55,6 +55,7 @@ test = [
     # HQ-SAM missing dependency:
     # https://github.com/SysCV/sam-hq/pull/59
     "timm>=0.5.0",
+    "sentencepiece>=0.2.0",
 ]
 conversion = [
     "diffusers>=0.26.1",

scripts/conversion/convert_ella_adapter.py

@@ -0,0 +1,102 @@
+import argparse
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download  # type: ignore
+
+from refiners.fluxion.utils import load_from_safetensors, save_to_safetensors
+
+
+class Args(argparse.Namespace):
+    source_path: str
+    output_path: str | None
+    use_half: bool
+
+
+def convert(args: Args) -> dict[str, torch.Tensor]:
+    if Path(args.source_path).suffix != ".safetensors":
+        args.source_path = hf_hub_download(
+            repo_id=args.source_path, filename="ella-sd1.5-tsc-t5xl.safetensors", local_dir="tests/weights/ELLA-Adapter"
+        )
+    weights = load_from_safetensors(args.source_path)
+
+    for key in list(weights.keys()):
+        if "latents" in key:
+            new_key = "PerceiverResampler.Latents.ParameterInitialized.weight"
+            weights[new_key] = weights.pop(key)
+        elif "time_embedding" in key:
+            new_key = key.replace("time_embedding", "TimestepEncoder.RangeEncoder").replace("linear", "Linear")
+            weights[new_key] = weights.pop(key)
+        elif "proj_in" in key:
+            new_key = f"PerceiverResampler.Linear.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+        elif "time_aware" in key:
+            new_key = f"PerceiverResampler.Residual.Linear.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+        elif "attn.in_proj" in key:
+            layer_num = int(key.split(".")[2])
+            query_param, key_param, value_param = weights.pop(key).chunk(3, dim=0)
+            param_type = "weight" if "weight" in key else "bias"
+            for i, param in enumerate([query_param, key_param, value_param]):
+                new_key = f"PerceiverResampler.Transformer.TransformerLayer_{layer_num+1}.Residual_1.PerceiverAttention.Attention.Distribute.Linear_{i+1}.{param_type}"
+                weights[new_key] = param
+        elif "attn.out_proj" in key:
+            layer_num = int(key.split(".")[2])
+            new_key = f"PerceiverResampler.Transformer.TransformerLayer_{layer_num+1}.Residual_1.PerceiverAttention.Attention.Linear.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+        elif "ln_ff" in key:
+            layer_num = int(key.split(".")[2])
+            new_key = f"PerceiverResampler.Transformer.TransformerLayer_{layer_num+1}.Residual_2.AdaLayerNorm.Parallel.Chain.Linear.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+        elif "ln_1" in key or "ln_2" in key:
+            layer_num = int(key.split(".")[2])
+            n = 1 if int(key.split(".")[3].split("_")[-1]) == 2 else 2
+            new_key = f"PerceiverResampler.Transformer.TransformerLayer_{layer_num+1}.Residual_1.PerceiverAttention.Distribute.AdaLayerNorm_{n}.Parallel.Chain.Linear.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+        elif "mlp" in key:
+            layer_num = int(key.split(".")[2])
+            n = 1 if "c_fc" in key else 2
+            new_key = f"PerceiverResampler.Transformer.TransformerLayer_{layer_num+1}.Residual_2.FeedForward.Linear_{n}.{key.split('.')[-1]}"
+            weights[new_key] = weights.pop(key)
+
+    if args.use_half:
+        weights = {key: value.half() for key, value in weights.items()}
+
+    return weights
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert a pretrained Ella Adapter to refiners implementation")
+    parser.add_argument(
+        "--from",
+        type=str,
+        dest="source_path",
+        default="QQGYLab/ELLA",
+        help=(
+            "A path to a local .safetensors weights. If not provided, a repo from Hugging Face Hub will be used"
+            "Default to QQGYLab/ELLA"
+        ),
+    )
+
+    parser.add_argument(
+        "--to",
+        type=str,
+        dest="output_path",
+        default=None,
+        help=(
+            "Path to save the converted model (extension will be .safetensors). If not specified, the output path will"
+            " be the source path with the prefix set to refiners"
+        ),
+    )
+    parser.add_argument(
+        "--half",
+        action="store_true",
+        dest="use_half",
+        default=True,
+        help="Use this flag to save the output file as half precision (default: full precision).",
+    )
+    args = parser.parse_args(namespace=Args())
+    weights = convert(args)
+    if args.output_path is None:
+        args.output_path = f"{Path(args.source_path).stem}-refiners.safetensors"
+    save_to_safetensors(path=args.output_path, tensors=weights)

scripts/prepare_test_weights.py

@@ -353,6 +353,29 @@ def download_ip_adapter():
     download_files(urls, sdxl_models_folder)
 
 
+def download_t5xl_fp16():
+    base_folder = os.path.join(test_weights_dir, "QQGYLab", "T5XLFP16")
+    urls = [
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/config.json",
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/model.safetensors",
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/special_tokens_map.json",
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/spiece.model",
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/tokenizer.json",
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/models--google--flan-t5-xl--text_encoder/tokenizer_config.json",
+    ]
+    download_files(urls, base_folder)
+
+
+def download_ella_adapter():
+    download_t5xl_fp16()
+    base_folder = os.path.join(test_weights_dir, "QQGYLab", "ELLA")
+    download_file(
+        "https://huggingface.co/QQGYLab/ELLA/resolve/main/ella-sd1.5-tsc-t5xl.safetensors",
+        base_folder,
+        expected_hash="5af7b200",
+    )
+
+
 def download_t2i_adapter():
     base_folder = os.path.join(test_weights_dir, "TencentARC", "t2iadapter_depth_sd15v2")
     urls = [
@@ -689,6 +712,17 @@ def convert_ip_adapter():
     )
 
 
+def convert_ella_adapter():
+    os.makedirs("tests/weights/ELLA-Adapter", exist_ok=True)
+    run_conversion_script(
+        "convert_ella_adapter.py",
+        "tests/weights/QQGYLab/ELLA/ella-sd1.5-tsc-t5xl.safetensors",
+        "tests/weights/ELLA-Adapter/ella-sd1.5-tsc-t5xl.safetensors",
+        half=True,
+        expected_hash="b8244cb6",
+    )
+
+
 def convert_t2i_adapter():
     os.makedirs("tests/weights/T2I-Adapter", exist_ok=True)
     run_conversion_script(
@@ -860,6 +894,7 @@ def download_all():
     download_unclip()
     download_ip_adapter()
     download_t2i_adapter()
+    download_ella_adapter()
     download_sam()
     download_hq_sam()
     download_dinov2()
@@ -884,6 +919,7 @@ def convert_all():
     convert_unclip()
     convert_ip_adapter()
     convert_t2i_adapter()
+    convert_ella_adapter()
     convert_sam()
     convert_hq_sam()
     convert_dinov2()

src/refiners/foundationals/latent_diffusion/__init__.py

@@ -8,6 +8,7 @@
 from refiners.foundationals.latent_diffusion.solvers import DPMSolver, LCMSolver, Solver
 from refiners.foundationals.latent_diffusion.stable_diffusion_1 import (
     SD1ControlnetAdapter,
+    SD1ELLAAdapter,
     SD1IPAdapter,
     SD1T2IAdapter,
     SD1UNet,
@@ -32,6 +33,7 @@
     "SD1ControlnetAdapter",
     "SD1IPAdapter",
     "SD1T2IAdapter",
+    "SD1ELLAAdapter",
     "SDXLUNet",
     "DoubleTextEncoder",
     "SDXLIPAdapter",