Function main

diffusers/scripts/convert_sd3_to_diffusers.py:198–244

(args)

def main(args):
    original_ckpt = load_original_checkpoint(args.checkpoint_path)
    num_layers = list(set(int(k.split(".", 2)[1]) for k in original_ckpt if "joint_blocks" in k))[-1] + 1  # noqa: C401
    caption_projection_dim = 1536

    converted_transformer_state_dict = convert_sd3_transformer_checkpoint_to_diffusers(
        original_ckpt, num_layers, caption_projection_dim
    )

    with CTX():
        transformer = SD3Transformer2DModel(
            sample_size=64,
            patch_size=2,
            in_channels=16,
            joint_attention_dim=4096,
            num_layers=num_layers,
            caption_projection_dim=caption_projection_dim,
            num_attention_heads=24,
            pos_embed_max_size=192,
        )
    if is_accelerate_available():
        load_model_dict_into_meta(transformer, converted_transformer_state_dict)
    else:
        transformer.load_state_dict(converted_transformer_state_dict, strict=True)

    print("Saving SD3 Transformer in Diffusers format.")
    transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer")

    if is_vae_in_checkpoint(original_ckpt):
        with CTX():
            vae = AutoencoderKL.from_config(
                "stabilityai/stable-diffusion-xl-base-1.0",
                subfolder="vae",
                latent_channels=16,
                use_post_quant_conv=False,
                use_quant_conv=False,
                scaling_factor=1.5305,
                shift_factor=0.0609,
            )
        converted_vae_state_dict = convert_ldm_vae_checkpoint(original_ckpt, vae.config)
        if is_accelerate_available():
            load_model_dict_into_meta(vae, converted_vae_state_dict)
        else:
            vae.load_state_dict(converted_vae_state_dict, strict=True)

        print("Saving SD3 Autoencoder in Diffusers format.")
        vae.to(dtype).save_pretrained(f"{args.output_path}/vae")


if __name__ == "__main__":