mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-09-19 19:05:28 +00:00
103 lines
4.4 KiB
Python
103 lines
4.4 KiB
Python
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)
|