refiners/scripts/conversion/convert_ella_adapter.py

103 lines
4.4 KiB
Python
Raw Permalink Normal View History

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)