diffusers/scripts/convert_tiny_autoencoder_to_diffusers.py at main · koji/diffusers · GitHub

History

71 lines (58 loc) · 2.33 KB

Raw

import argparse

import safetensors.torch

from diffusers import AutoencoderTiny

"""

Example - From the diffusers root directory:

Download the weights:

```sh

$ wget -q https://huggingface.co/madebyollin/taesd/resolve/main/taesd_encoder.safetensors

$ wget -q https://huggingface.co/madebyollin/taesd/resolve/main/taesd_decoder.safetensors

```

Convert the model:

```sh

$ python scripts/convert_tiny_autoencoder_to_diffusers.py \

--encoder_ckpt_path taesd_encoder.safetensors \

--decoder_ckpt_path taesd_decoder.safetensors \

--dump_path taesd-diffusers

```

"""

if __name__ == "__main__":

parser = argparse.ArgumentParser()

parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")

parser.add_argument(

"--encoder_ckpt_path",

default=None,

type=str,

required=True,

help="Path to the encoder ckpt.",

)

parser.add_argument(

"--decoder_ckpt_path",

default=None,

type=str,

required=True,

help="Path to the decoder ckpt.",

)

parser.add_argument(

"--use_safetensors", action="store_true", help="Whether to serialize in the safetensors format."

)

args = parser.parse_args()

print("Loading the original state_dicts of the encoder and the decoder...")

encoder_state_dict = safetensors.torch.load_file(args.encoder_ckpt_path)

decoder_state_dict = safetensors.torch.load_file(args.decoder_ckpt_path)

print("Populating the state_dicts in the diffusers format...")

tiny_autoencoder = AutoencoderTiny()

new_state_dict = {}

# Modify the encoder state dict.

for k in encoder_state_dict:

new_state_dict.update({f"encoder.layers.{k}": encoder_state_dict[k]})

# Modify the decoder state dict.

for k in decoder_state_dict:

layer_id = int(k.split(".")[0]) - 1

new_k = str(layer_id) + "." + ".".join(k.split(".")[1:])

new_state_dict.update({f"decoder.layers.{new_k}": decoder_state_dict[k]})

# Assertion tests with the original implementation can be found here:

# https://gist.github.com/sayakpaul/337b0988f08bd2cf2b248206f760e28f

tiny_autoencoder.load_state_dict(new_state_dict)

print("Population successful, serializing...")

tiny_autoencoder.save_pretrained(args.dump_path, safe_serialization=args.use_safetensors)

diffusers/scripts/convert_tiny_autoencoder_to_diffusers.py at main · koji/diffusers · GitHub

Expand file tree

Breadcrumbs

convert_tiny_autoencoder_to_diffusers.py

Latest commit

History

Breadcrumbs

convert_tiny_autoencoder_to_diffusers.py

File metadata and controls