Function main

scripts/convert_flux_to_diffusers.py:275–304

(args)

def main(args):
    original_ckpt = load_original_checkpoint(args)
    has_guidance = any("guidance" in k for k in original_ckpt)

    if args.transformer:
        num_layers = 19
        num_single_layers = 38
        inner_dim = 3072
        mlp_ratio = 4.0

        converted_transformer_state_dict = convert_flux_transformer_checkpoint_to_diffusers(
            original_ckpt, num_layers, num_single_layers, inner_dim, mlp_ratio=mlp_ratio
        )
        transformer = FluxTransformer2DModel(
            in_channels=args.in_channels, out_channels=args.out_channels, guidance_embeds=has_guidance
        )
        transformer.load_state_dict(converted_transformer_state_dict, strict=True)

        print(
            f"Saving Flux Transformer in Diffusers format. Variant: {'guidance-distilled' if has_guidance else 'timestep-distilled'}"
        )
        transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer")

    if args.vae:
        config = AutoencoderKL.load_config("stabilityai/stable-diffusion-3-medium-diffusers", subfolder="vae")
        vae = AutoencoderKL.from_config(config, scaling_factor=0.3611, shift_factor=0.1159).to(torch.bfloat16)

        converted_vae_state_dict = convert_ldm_vae_checkpoint(original_ckpt, vae.config)
        vae.load_state_dict(converted_vae_state_dict, strict=True)
        vae.to(dtype).save_pretrained(f"{args.output_path}/vae")


if __name__ == "__main__":