From 691072b864c94deb73ab9fb9692e19020ac8ce92 Mon Sep 17 00:00:00 2001 From: hlky Date: Tue, 8 Sep 2026 10:07:55 +0100 Subject: [PATCH] Unify checkpoint conversions with bidirectional component rules --- MANIFEST.in | 2 + docs/source/en/api/loaders/single_file.md | 27 +- .../source/en/api/pipelines/stable_audio_3.md | 2 +- .../en/using-diffusers/other-formats.md | 112 +- .../ko/using-diffusers/other-formats.md | 2 +- scripts/README.md | 188 ++ scripts/build_pipeline.py | 33 + scripts/convert_amused.py | 523 --- scripts/convert_anima_to_diffusers.py | 316 -- ...rt_animatediff_motion_lora_to_diffusers.py | 69 - ..._animatediff_motion_module_to_diffusers.py | 62 - ...ert_animatediff_sparsectrl_to_diffusers.py | 83 - .../convert_asymmetric_vqgan_to_diffusers.py | 184 -- scripts/convert_aura_flow_to_diffusers.py | 131 - scripts/convert_blipdiffusion_to_diffusers.py | 344 -- scripts/convert_checkpoint.py | 103 + scripts/convert_consistency_decoder.py | 1131 ------- scripts/convert_consistency_to_diffusers.py | 315 -- scripts/convert_cosmos_to_diffusers.py | 976 ------ .../convert_dance_diffusion_to_diffusers.py | 346 -- scripts/convert_dcae_to_diffusers.py | 323 -- ...t_ddpm_original_checkpoint_to_diffusers.py | 431 --- .../convert_diffusers_sdxl_lora_to_webui.py | 56 - scripts/convert_diffusers_to_original_sdxl.py | 350 -- ..._diffusers_to_original_stable_diffusion.py | 353 -- scripts/convert_dit_to_diffusers.py | 162 - scripts/convert_flux2_to_diffusers.py | 536 --- scripts/convert_flux_to_diffusers.py | 308 -- scripts/convert_gligen_to_diffusers.py | 581 ---- scripts/convert_hunyuan_image_to_diffusers.py | 1044 ------ .../convert_hunyuan_video1_5_to_diffusers.py | 875 ----- scripts/convert_hunyuan_video_to_diffusers.py | 353 -- ...vert_hunyuandit_controlnet_to_diffusers.py | 241 -- scripts/convert_hunyuandit_to_diffusers.py | 266 -- scripts/convert_i2vgen_to_diffusers.py | 510 --- scripts/convert_if.py | 1250 ------- scripts/convert_joyimage_edit_to_diffusers.py | 355 -- scripts/convert_k_upscaler_to_diffusers.py | 297 -- ...convert_kakao_brain_unclip_to_diffusers.py | 1159 ------- scripts/convert_kandinsky3_unet.py | 98 - scripts/convert_kandinsky_to_diffusers.py | 1411 -------- ...rt_ldm_original_checkpoint_to_diffusers.py | 359 -- .../convert_lora_safetensor_to_diffusers.py | 128 - scripts/convert_ltx2_to_diffusers.py | 1703 ---------- scripts/convert_ltx_to_diffusers.py | 525 --- scripts/convert_lumina_to_diffusers.py | 142 - scripts/convert_minimax_h3_to_diffusers.py | 947 ------ .../convert_minimax_music3_to_diffusers.py | 269 -- scripts/convert_mochi_to_diffusers.py | 463 --- .../convert_models_diffuser_to_diffusers.py | 100 - .../convert_ms_text_to_video_to_diffusers.py | 428 --- .../convert_music_spectrogram_to_diffusers.py | 203 -- ...ncsnpp_original_checkpoint_to_diffusers.py | 185 -- ...convert_original_audioldm2_to_diffusers.py | 1135 ------- .../convert_original_audioldm_to_diffusers.py | 1042 ------ .../convert_original_musicldm_to_diffusers.py | 1056 ------ scripts/convert_original_t2i_adapter.py | 250 -- scripts/convert_ovis_image_to_diffusers.py | 263 -- scripts/convert_pixart_alpha_to_diffusers.py | 198 -- scripts/convert_pixart_sigma_to_diffusers.py | 223 -- .../convert_sana_controlnet_to_diffusers.py | 216 -- .../convert_sd3_controlnet_to_diffusers.py | 185 -- scripts/convert_sd3_to_diffusers.py | 351 -- scripts/convert_shap_e_to_diffusers.py | 1080 ------ scripts/convert_skyreelsv2_to_diffusers.py | 637 ---- scripts/convert_stable_audio.py | 279 -- scripts/convert_svd_to_diffusers.py | 730 ---- .../convert_tiny_autoencoder_to_diffusers.py | 71 - scripts/convert_unidiffuser_to_diffusers.py | 786 ----- scripts/convert_vae_pt_to_diffusers.py | 177 - ...onvert_versatile_diffusion_to_diffusers.py | 791 ----- scripts/convert_vq_diffusion_to_diffusers.py | 916 ----- scripts/convert_wan_to_diffusers.py | 1287 -------- scripts/convert_wuerstchen.py | 115 - scripts/convert_zero123_to_diffusers.py | 807 ----- scripts/export_pipeline_checkpoint.py | 30 + ...rt_stable_diffusion_checkpoint_to_onnx.py} | 0 ...rt_stable_diffusion_controlnet_to_onnx.py} | 4 +- ...table_diffusion_controlnet_to_tensorrt.py} | 4 +- ..._to_onnx.py => export_vae_diff_to_onnx.py} | 0 scripts/merge_lora.py | 26 + scripts/recipes/README.md | 39 + .../ace_step.py} | 81 +- scripts/recipes/amused.py | 231 ++ scripts/recipes/anima.py | 159 + .../anyflow.py} | 2 +- scripts/recipes/audioldm.py | 361 ++ scripts/recipes/audioldm2.py | 420 +++ scripts/recipes/blip_diffusion.py | 106 + .../cogvideox.py} | 160 +- .../cogview3.py} | 106 +- .../cogview4.py} | 115 +- .../cogview4_megatron.py} | 135 +- scripts/recipes/consistency.py | 87 + scripts/recipes/consistency_decoder.py | 36 + .../controlnet.py} | 0 scripts/recipes/cosmos.py | 456 +++ scripts/recipes/dance_diffusion.py | 159 + scripts/recipes/ddpm.py | 66 + scripts/recipes/deepfloyd_if.py | 313 ++ scripts/recipes/dit.py | 89 + scripts/recipes/flux2.py | 133 + .../flux_ip_adapter.py} | 42 +- scripts/recipes/gligen.py | 190 ++ scripts/recipes/hunyuan_dit.py | 81 + scripts/recipes/hunyuan_image.py | 312 ++ scripts/recipes/hunyuan_video.py | 132 + scripts/recipes/hunyuan_video15.py | 303 ++ scripts/recipes/i2vgen.py | 95 + scripts/recipes/joyimage.py | 152 + scripts/recipes/k_upscaler.py | 98 + scripts/recipes/kandinsky.py | 310 ++ scripts/recipes/ldm.py | 69 + .../longcat_audio.py} | 16 +- scripts/recipes/ltx.py | 193 ++ scripts/recipes/ltx2.py | 555 ++++ scripts/recipes/lumina.py | 72 + scripts/recipes/minimax_h3.py | 363 ++ scripts/recipes/minimax_music3.py | 182 + scripts/recipes/mochi.py | 95 + scripts/recipes/music_spectrogram.py | 131 + scripts/recipes/musicldm.py | 365 ++ scripts/recipes/ncsnpp.py | 81 + .../omnigen.py} | 40 +- scripts/recipes/pixart_alpha.py | 79 + scripts/recipes/pixart_sigma.py | 87 + .../prx.py} | 130 +- .../rae.py} | 102 +- .../sana.py} | 175 +- scripts/recipes/sana_controlnet.py | 114 + .../sana_video.py} | 125 +- scripts/recipes/shap_e.py | 252 ++ scripts/recipes/skyreels_v2.py | 124 + scripts/recipes/stable_audio.py | 183 + .../stable_audio3.py} | 477 +-- .../stable_cascade.py} | 15 +- .../stable_cascade_lite.py} | 15 +- .../stable_diffusion.py} | 0 scripts/recipes/unclip.py | 431 +++ .../unclip_image_variation.py} | 0 scripts/recipes/unidiffuser.py | 220 ++ scripts/recipes/versatile_diffusion.py | 183 + scripts/recipes/vq_diffusion.py | 487 +++ scripts/recipes/wan.py | 310 ++ scripts/recipes/wuerstchen.py | 87 + scripts/recipes/zero123.py | 173 + setup.py | 2 +- src/diffusers/loaders/conversion/README.md | 230 ++ src/diffusers/loaders/conversion/__init__.py | 30 + src/diffusers/loaders/conversion/ace_step.py | 73 + .../conversion/ace_step_conditioner.py | 37 + .../conversion/ace_step_detokenizer.py | 24 + .../loaders/conversion/ace_step_tokenizer.py | 33 + .../loaders/conversion/anima_conditioner.py | 42 + .../loaders/conversion/animatediff.py | 74 + src/diffusers/loaders/conversion/anyflow.py | 26 + .../loaders/conversion/anyflow_far.py | 23 + .../loaders/conversion/asymmetric_vae.py | 51 + .../conversion/audioldm2_projection.py | 31 + .../loaders/conversion/audioldm2_unet.py | 58 + src/diffusers/loaders/conversion/auraflow.py | 65 + .../loaders/conversion/autoencoder_dc.py | 108 + .../loaders/conversion/blip_qformer.py | 91 + .../loaders/conversion/checkpoint.py | 238 ++ src/diffusers/loaders/conversion/chroma.py | 72 + .../loaders/conversion/chronoedit.py | 21 + .../loaders/conversion/clap_audio.py | 87 + src/diffusers/loaders/conversion/clap_text.py | 51 + src/diffusers/loaders/conversion/clip.py | 104 + .../loaders/conversion/clip_vision.py | 67 + src/diffusers/loaders/conversion/cogvideox.py | 239 ++ .../loaders/conversion/cogview3plus.py | 55 + src/diffusers/loaders/conversion/cogview4.py | 27 + .../loaders/conversion/cogview4_megatron.py | 63 + .../loaders/conversion/configs/__init__.py | 61 + .../loaders/conversion/configs/anima.py | 64 + .../loaders/conversion/configs/audioldm.py | 178 + .../loaders/conversion/configs/audioldm2.py | 180 + .../conversion/configs/blip_diffusion.py | 35 + .../loaders/conversion/configs/catalog.json | 145 + .../loaders/conversion/configs/consistency.py | 119 + .../loaders/conversion/configs/cosmos.py | 266 ++ .../conversion/configs/deepfloyd_if.py | 200 ++ .../loaders/conversion/configs/flux2.py | 85 + .../loaders/conversion/configs/gligen.py | 85 + .../conversion/configs/hunyuan_video.py | 77 + .../conversion/configs/hunyuan_video15.py | 103 + .../loaders/conversion/configs/joyimage.py | 30 + .../loaders/conversion/configs/kandinsky.py | 142 + .../loaders/conversion/configs/ltx.py | 208 ++ .../loaders/conversion/configs/ltx2.py | 737 +++++ .../loaders/conversion/configs/minimax_h3.py | 156 + .../loaders/conversion/configs/musicldm.py | 178 + .../loaders/conversion/configs/prx.py | 80 + .../loaders/conversion/configs/rae.py | 38 + .../loaders/conversion/configs/shap_e.py | 52 + .../loaders/conversion/configs/skyreels_v2.py | 253 ++ .../conversion/configs/stable_audio3.py | 219 ++ .../loaders/conversion/configs/unclip.py | 86 + .../loaders/conversion/configs/unidiffuser.py | 214 ++ .../conversion/configs/versatile_diffusion.py | 182 + .../loaders/conversion/configs/wan.py | 273 ++ .../loaders/conversion/configs/zero123.py | 147 + .../loaders/conversion/consistency_decoder.py | 100 + .../loaders/conversion/controlnet.py | 21 + src/diffusers/loaders/conversion/core.py | 119 + src/diffusers/loaders/conversion/cosmos.py | 83 + .../loaders/conversion/cosmos_controlnet.py | 41 + .../loaders/conversion/cosmos_vae.py | 78 + .../loaders/conversion/dance_diffusion.py | 104 + .../loaders/conversion/diffuser_rl.py | 80 + src/diffusers/loaders/conversion/dit.py | 71 + .../loaders/conversion/ernie_image.py | 48 + src/diffusers/loaders/conversion/flux.py | 87 + src/diffusers/loaders/conversion/flux2.py | 74 + src/diffusers/loaders/conversion/flux2_vae.py | 40 + .../loaders/conversion/flux_ip_adapter.py | 30 + src/diffusers/loaders/conversion/gligen.py | 44 + src/diffusers/loaders/conversion/hidream.py | 45 + src/diffusers/loaders/conversion/hifigan.py | 47 + .../loaders/conversion/hunyuan_dit.py | 94 + .../conversion/hunyuan_dit_controlnet.py | 22 + .../loaders/conversion/hunyuan_image.py | 83 + .../conversion/hunyuan_image_refiner_vae.py | 22 + .../loaders/conversion/hunyuan_image_vae.py | 74 + .../loaders/conversion/hunyuan_video.py | 115 + .../loaders/conversion/hunyuan_video15.py | 51 + .../loaders/conversion/hunyuan_video15_vae.py | 59 + .../loaders/conversion/hunyuan_video_vae.py | 46 + src/diffusers/loaders/conversion/i2vgen_xl.py | 49 + src/diffusers/loaders/conversion/if_unet.py | 77 + src/diffusers/loaders/conversion/io.py | 277 ++ src/diffusers/loaders/conversion/joy_image.py | 42 + .../loaders/conversion/k_upscaler.py | 134 + .../loaders/conversion/kandinsky3.py | 115 + src/diffusers/loaders/conversion/ldm_unet.py | 149 + src/diffusers/loaders/conversion/ldm_vae.py | 89 + .../loaders/conversion/longcat_audio.py | 54 + .../loaders/conversion/longcat_audio_vae.py | 46 + src/diffusers/loaders/conversion/lora.py | 99 + src/diffusers/loaders/conversion/ltx.py | 44 + src/diffusers/loaders/conversion/ltx2.py | 93 + .../loaders/conversion/ltx2_audio_vae.py | 63 + .../loaders/conversion/ltx2_connectors.py | 42 + .../conversion/ltx2_diffusion_decoder.py | 75 + .../loaders/conversion/ltx2_duration.py | 34 + .../loaders/conversion/ltx2_upsampler.py | 31 + src/diffusers/loaders/conversion/ltx2_vae.py | 21 + .../loaders/conversion/ltx2_vocoder.py | 54 + .../loaders/conversion/ltx_upsampler.py | 21 + src/diffusers/loaders/conversion/ltx_vae.py | 103 + src/diffusers/loaders/conversion/lumina.py | 76 + src/diffusers/loaders/conversion/lumina2.py | 62 + .../loaders/conversion/minimax_h3.py | 81 + .../conversion/minimax_h3_audio_vae.py | 76 + .../loaders/conversion/minimax_h3_vae.py | 75 + .../loaders/conversion/minimax_music3.py | 46 + .../conversion/minimax_music3_conditioner.py | 27 + .../loaders/conversion/minimax_music3_rvq.py | 42 + .../conversion/minimax_music3_vocoder.py | 37 + src/diffusers/loaders/conversion/mochi.py | 70 + src/diffusers/loaders/conversion/mochi_vae.py | 82 + .../loaders/conversion/motif_video.py | 69 + src/diffusers/loaders/conversion/movq.py | 32 + src/diffusers/loaders/conversion/ncsnpp.py | 90 + src/diffusers/loaders/conversion/omnigen.py | 57 + src/diffusers/loaders/conversion/oobleck.py | 62 + .../loaders/conversion/ovis_image.py | 55 + src/diffusers/loaders/conversion/paella.py | 50 + src/diffusers/loaders/conversion/pipeline.py | 97 + .../loaders/conversion/pipeline_components.py | 93 + src/diffusers/loaders/conversion/pixart.py | 70 + .../presets/asymmetric-vae-1.5.json | 36 + .../conversion/presets/asymmetric-vae-2.json | 36 + .../conversion/presets/diffuser-rl-128.json | 32 + .../conversion/presets/diffuser-rl-32.json | 35 + .../conversion/presets/diffuser-rl-value.json | 31 + .../presets/modelscope-text-to-video.json | 3 + .../loaders/conversion/presets/sd-vae.json | 28 + .../loaders/conversion/presets/tiny-vae.json | 3 + src/diffusers/loaders/conversion/prior.py | 105 + src/diffusers/loaders/conversion/prx.py | 66 + src/diffusers/loaders/conversion/qwen3.py | 45 + .../loaders/conversion/qwen_image.py | 44 + .../loaders/conversion/qwen_image_vae.py | 23 + src/diffusers/loaders/conversion/rae.py | 128 + src/diffusers/loaders/conversion/registry.py | 433 +++ src/diffusers/loaders/conversion/same.py | 56 + src/diffusers/loaders/conversion/sana.py | 90 + .../loaders/conversion/sana_controlnet.py | 37 + .../loaders/conversion/sana_video.py | 28 + src/diffusers/loaders/conversion/sd3.py | 89 + .../loaders/conversion/sd3_controlnet.py | 63 + .../loaders/conversion/shap_e_renderer.py | 40 + .../loaders/conversion/shap_e_tables.py | 494 +++ .../loaders/conversion/skyreels_v2.py | 31 + src/diffusers/loaders/conversion/source.py | 132 + .../loaders/conversion/sparse_controlnet.py | 64 + .../conversion/spectrogram_continuous.py | 25 + .../loaders/conversion/spectrogram_notes.py | 40 + .../loaders/conversion/stable_audio.py | 65 + .../loaders/conversion/stable_audio3.py | 74 + .../conversion/stable_audio3_duration.py | 33 + .../conversion/stable_audio_projection.py | 28 + .../loaders/conversion/stable_cascade.py | 78 + src/diffusers/loaders/conversion/svd.py | 76 + src/diffusers/loaders/conversion/svd_vae.py | 78 + .../loaders/conversion/t2i_adapter.py | 39 + src/diffusers/loaders/conversion/t5.py | 29 + src/diffusers/loaders/conversion/t5_film.py | 49 + src/diffusers/loaders/conversion/tiny_vae.py | 48 + .../loaders/conversion/transforms.py | 327 ++ src/diffusers/loaders/conversion/umt5.py | 25 + .../conversion/unclip_text_projection.py | 29 + .../loaders/conversion/unclip_unet.py | 74 + src/diffusers/loaders/conversion/unet_1d.py | 24 + src/diffusers/loaders/conversion/unet_2d.py | 174 + src/diffusers/loaders/conversion/unet_3d.py | 97 + .../loaders/conversion/unidiffuser.py | 78 + .../loaders/conversion/unidiffuser_text.py | 33 + src/diffusers/loaders/conversion/uvit.py | 117 + .../conversion/versatile_image_unet.py | 25 + .../loaders/conversion/versatile_text_unet.py | 37 + .../loaders/conversion/vq_diffusion.py | 45 + src/diffusers/loaders/conversion/vq_model.py | 28 + src/diffusers/loaders/conversion/wan.py | 59 + .../loaders/conversion/wan_animate.py | 77 + .../loaders/conversion/wan_animate2.py | 53 + src/diffusers/loaders/conversion/wan_vace.py | 37 + src/diffusers/loaders/conversion/wan_vae.py | 93 + .../loaders/conversion/wuerstchen_decoder.py | 42 + .../loaders/conversion/wuerstchen_prior.py | 63 + src/diffusers/loaders/conversion/z_image.py | 80 + .../loaders/conversion/z_image_controlnet.py | 54 + .../loaders/conversion/zero123_projection.py | 19 + src/diffusers/loaders/single_file_model.py | 166 +- src/diffusers/loaders/single_file_utils.py | 2933 +---------------- src/diffusers/models/adapter.py | 3 +- .../autoencoders/autoencoder_asym_kl.py | 3 +- .../autoencoders/autoencoder_kl_cosmos.py | 3 +- .../autoencoder_kl_hunyuan_video.py | 3 +- .../autoencoder_kl_hunyuanimage_refiner.py | 3 +- .../autoencoder_kl_hunyuanvideo15.py | 3 +- .../autoencoders/autoencoder_kl_ltx2_audio.py | 3 +- .../autoencoders/autoencoder_kl_minimax_h3.py | 3 +- .../autoencoder_kl_minimax_h3_audio.py | 3 +- .../autoencoders/autoencoder_kl_mochi.py | 3 +- .../autoencoder_kl_temporal_decoder.py | 3 +- .../autoencoder_longcat_audio_dit.py | 3 +- .../autoencoders/autoencoder_oobleck.py | 3 +- .../models/autoencoders/autoencoder_rae.py | 9 +- .../models/autoencoders/autoencoder_same.py | 3 +- .../models/autoencoders/autoencoder_tiny.py | 3 +- .../autoencoders/consistency_decoder_vae.py | 3 +- .../autoencoders/ltx2_diffusion_decoder.py | 3 +- .../autoencoders/minimax_music3_vocoder.py | 3 +- src/diffusers/models/autoencoders/vq_model.py | 3 +- .../condition_embedder_anima.py | 4 +- .../condition_embedder_minimax_music3.py | 3 +- .../models/controlnets/controlnet_hunyuan.py | 3 +- .../models/controlnets/controlnet_sana.py | 4 +- .../transformers/ace_step_transformer.py | 6 +- .../transformers/cogvideox_transformer_3d.py | 6 +- .../models/transformers/dit_transformer_2d.py | 3 +- .../transformers/hunyuan_transformer_2d.py | 3 +- .../models/transformers/lumina_nextdit2d.py | 3 +- .../minimax_music3_rvq_depth_decoder.py | 3 +- .../transformers/pixart_transformer_2d.py | 3 +- .../models/transformers/prior_transformer.py | 6 +- .../transformers/stable_audio_transformer.py | 3 +- .../transformers/t5_film_transformer.py | 3 +- .../models/transformers/transformer_2d.py | 3 +- .../transformers/transformer_cogview3plus.py | 3 +- .../transformers/transformer_cogview4.py | 4 +- .../transformers/transformer_joyimage.py | 3 +- .../transformer_joyimage_edit_plus.py | 3 +- .../transformer_longcat_audio_dit.py | 3 +- .../transformers/transformer_minimax_h3.py | 6 +- .../transformer_minimax_music3.py | 3 +- .../transformers/transformer_omnigen.py | 3 +- .../models/transformers/transformer_prx.py | 3 +- .../transformers/transformer_stable_audio3.py | 3 +- src/diffusers/models/unets/unet_1d.py | 3 +- src/diffusers/models/unets/unet_2d.py | 3 +- .../models/unets/unet_3d_condition.py | 6 +- src/diffusers/models/unets/unet_i2vgen_xl.py | 4 +- src/diffusers/models/unets/unet_kandinsky3.py | 3 +- .../unets/unet_spatio_temporal_condition.py | 6 +- src/diffusers/models/unets/uvit_2d.py | 4 +- .../pipelines/ace_step/modeling_ace_step.py | 7 +- .../pipelines/audioldm2/modeling_audioldm2.py | 8 +- .../continuous_encoder.py | 3 +- .../spectrogram_diffusion/notes_encoder.py | 3 +- .../pipelines/deprecated/unclip/text_proj.py | 3 +- .../unidiffuser/modeling_text_decoder.py | 3 +- .../deprecated/unidiffuser/modeling_uvit.py | 3 +- .../versatile_diffusion/modeling_text_unet.py | 3 +- .../vq_diffusion/pipeline_vq_diffusion.py | 3 +- .../wuerstchen/modeling_paella_vq_model.py | 3 +- .../modeling_wuerstchen_diffnext.py | 3 +- .../wuerstchen/modeling_wuerstchen_prior.py | 6 +- .../ltx/modeling_latent_upsampler.py | 3 +- src/diffusers/pipelines/ltx2/connectors.py | 4 +- src/diffusers/pipelines/ltx2/duration_head.py | 3 +- .../pipelines/ltx2/latent_upsampler.py | 3 +- src/diffusers/pipelines/ltx2/vocoder.py | 5 +- src/diffusers/pipelines/shap_e/renderer.py | 3 +- .../stable_audio/modeling_stable_audio.py | 3 +- .../stable_audio_3/modeling_stable_audio_3.py | 3 +- .../stable_diffusion/convert_from_ckpt.py | 724 +--- tests/others/test_checkpoint_conversion.py | 245 ++ tests/others/test_conversion_io.py | 418 +++ tests/others/test_pipeline_conversion.py | 199 ++ .../single_file/test_auxiliary_conversions.py | 1044 ++++++ .../single_file/test_cogvideox_conversion.py | 308 ++ tests/single_file/test_model_conversions.py | 198 ++ .../test_text_encoder_conversions.py | 224 ++ 417 files changed, 28649 insertions(+), 38867 deletions(-) create mode 100644 scripts/README.md create mode 100644 scripts/build_pipeline.py delete mode 100644 scripts/convert_amused.py delete mode 100644 scripts/convert_anima_to_diffusers.py delete mode 100644 scripts/convert_animatediff_motion_lora_to_diffusers.py delete mode 100644 scripts/convert_animatediff_motion_module_to_diffusers.py delete mode 100644 scripts/convert_animatediff_sparsectrl_to_diffusers.py delete mode 100644 scripts/convert_asymmetric_vqgan_to_diffusers.py delete mode 100644 scripts/convert_aura_flow_to_diffusers.py delete mode 100644 scripts/convert_blipdiffusion_to_diffusers.py create mode 100644 scripts/convert_checkpoint.py delete mode 100644 scripts/convert_consistency_decoder.py delete mode 100644 scripts/convert_consistency_to_diffusers.py delete mode 100644 scripts/convert_cosmos_to_diffusers.py delete mode 100755 scripts/convert_dance_diffusion_to_diffusers.py delete mode 100644 scripts/convert_dcae_to_diffusers.py delete mode 100644 scripts/convert_ddpm_original_checkpoint_to_diffusers.py delete mode 100644 scripts/convert_diffusers_sdxl_lora_to_webui.py delete mode 100644 scripts/convert_diffusers_to_original_sdxl.py delete mode 100644 scripts/convert_diffusers_to_original_stable_diffusion.py delete mode 100644 scripts/convert_dit_to_diffusers.py delete mode 100644 scripts/convert_flux2_to_diffusers.py delete mode 100644 scripts/convert_flux_to_diffusers.py delete mode 100644 scripts/convert_gligen_to_diffusers.py delete mode 100644 scripts/convert_hunyuan_image_to_diffusers.py delete mode 100644 scripts/convert_hunyuan_video1_5_to_diffusers.py delete mode 100644 scripts/convert_hunyuan_video_to_diffusers.py delete mode 100644 scripts/convert_hunyuandit_controlnet_to_diffusers.py delete mode 100644 scripts/convert_hunyuandit_to_diffusers.py delete mode 100644 scripts/convert_i2vgen_to_diffusers.py delete mode 100644 scripts/convert_if.py delete mode 100644 scripts/convert_joyimage_edit_to_diffusers.py delete mode 100644 scripts/convert_k_upscaler_to_diffusers.py delete mode 100644 scripts/convert_kakao_brain_unclip_to_diffusers.py delete mode 100644 scripts/convert_kandinsky3_unet.py delete mode 100644 scripts/convert_kandinsky_to_diffusers.py delete mode 100644 scripts/convert_ldm_original_checkpoint_to_diffusers.py delete mode 100644 scripts/convert_lora_safetensor_to_diffusers.py delete mode 100644 scripts/convert_ltx2_to_diffusers.py delete mode 100644 scripts/convert_ltx_to_diffusers.py delete mode 100644 scripts/convert_lumina_to_diffusers.py delete mode 100644 scripts/convert_minimax_h3_to_diffusers.py delete mode 100644 scripts/convert_minimax_music3_to_diffusers.py delete mode 100644 scripts/convert_mochi_to_diffusers.py delete mode 100644 scripts/convert_models_diffuser_to_diffusers.py delete mode 100644 scripts/convert_ms_text_to_video_to_diffusers.py delete mode 100644 scripts/convert_music_spectrogram_to_diffusers.py delete mode 100644 scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py delete mode 100644 scripts/convert_original_audioldm2_to_diffusers.py delete mode 100644 scripts/convert_original_audioldm_to_diffusers.py delete mode 100644 scripts/convert_original_musicldm_to_diffusers.py delete mode 100644 scripts/convert_original_t2i_adapter.py delete mode 100644 scripts/convert_ovis_image_to_diffusers.py delete mode 100644 scripts/convert_pixart_alpha_to_diffusers.py delete mode 100644 scripts/convert_pixart_sigma_to_diffusers.py delete mode 100644 scripts/convert_sana_controlnet_to_diffusers.py delete mode 100644 scripts/convert_sd3_controlnet_to_diffusers.py delete mode 100644 scripts/convert_sd3_to_diffusers.py delete mode 100644 scripts/convert_shap_e_to_diffusers.py delete mode 100644 scripts/convert_skyreelsv2_to_diffusers.py delete mode 100644 scripts/convert_stable_audio.py delete mode 100644 scripts/convert_svd_to_diffusers.py delete mode 100644 scripts/convert_tiny_autoencoder_to_diffusers.py delete mode 100644 scripts/convert_unidiffuser_to_diffusers.py delete mode 100644 scripts/convert_vae_pt_to_diffusers.py delete mode 100644 scripts/convert_versatile_diffusion_to_diffusers.py delete mode 100644 scripts/convert_vq_diffusion_to_diffusers.py delete mode 100644 scripts/convert_wan_to_diffusers.py delete mode 100644 scripts/convert_wuerstchen.py delete mode 100644 scripts/convert_zero123_to_diffusers.py create mode 100644 scripts/export_pipeline_checkpoint.py rename scripts/{convert_stable_diffusion_checkpoint_to_onnx.py => export_stable_diffusion_checkpoint_to_onnx.py} (100%) rename scripts/{convert_stable_diffusion_controlnet_to_onnx.py => export_stable_diffusion_controlnet_to_onnx.py} (99%) rename scripts/{convert_stable_diffusion_controlnet_to_tensorrt.py => export_stable_diffusion_controlnet_to_tensorrt.py} (97%) rename scripts/{convert_vae_diff_to_onnx.py => export_vae_diff_to_onnx.py} (100%) create mode 100644 scripts/merge_lora.py create mode 100644 scripts/recipes/README.md rename scripts/{convert_ace_step_to_diffusers.py => recipes/ace_step.py} (85%) create mode 100644 scripts/recipes/amused.py create mode 100644 scripts/recipes/anima.py rename scripts/{convert_anyflow_to_diffusers.py => recipes/anyflow.py} (99%) create mode 100644 scripts/recipes/audioldm.py create mode 100644 scripts/recipes/audioldm2.py create mode 100644 scripts/recipes/blip_diffusion.py rename scripts/{convert_cogvideox_to_diffusers.py => recipes/cogvideox.py} (56%) rename scripts/{convert_cogview3_to_diffusers.py => recipes/cogview3.py} (52%) rename scripts/{convert_cogview4_to_diffusers.py => recipes/cogview4.py} (51%) rename scripts/{convert_cogview4_to_diffusers_megatron.py => recipes/cogview4_megatron.py} (58%) create mode 100644 scripts/recipes/consistency.py create mode 100644 scripts/recipes/consistency_decoder.py rename scripts/{convert_original_controlnet_to_diffusers.py => recipes/controlnet.py} (100%) create mode 100644 scripts/recipes/cosmos.py create mode 100644 scripts/recipes/dance_diffusion.py create mode 100644 scripts/recipes/ddpm.py create mode 100644 scripts/recipes/deepfloyd_if.py create mode 100644 scripts/recipes/dit.py create mode 100644 scripts/recipes/flux2.py rename scripts/{convert_flux_xlabs_ipadapter_to_diffusers.py => recipes/flux_ip_adapter.py} (53%) create mode 100644 scripts/recipes/gligen.py create mode 100644 scripts/recipes/hunyuan_dit.py create mode 100644 scripts/recipes/hunyuan_image.py create mode 100644 scripts/recipes/hunyuan_video.py create mode 100644 scripts/recipes/hunyuan_video15.py create mode 100644 scripts/recipes/i2vgen.py create mode 100644 scripts/recipes/joyimage.py create mode 100644 scripts/recipes/k_upscaler.py create mode 100644 scripts/recipes/kandinsky.py create mode 100644 scripts/recipes/ldm.py rename scripts/{convert_longcat_audio_dit_to_diffusers.py => recipes/longcat_audio.py} (91%) create mode 100644 scripts/recipes/ltx.py create mode 100644 scripts/recipes/ltx2.py create mode 100644 scripts/recipes/lumina.py create mode 100644 scripts/recipes/minimax_h3.py create mode 100644 scripts/recipes/minimax_music3.py create mode 100644 scripts/recipes/mochi.py create mode 100644 scripts/recipes/music_spectrogram.py create mode 100644 scripts/recipes/musicldm.py create mode 100644 scripts/recipes/ncsnpp.py rename scripts/{convert_omnigen_to_diffusers.py => recipes/omnigen.py} (74%) create mode 100644 scripts/recipes/pixart_alpha.py create mode 100644 scripts/recipes/pixart_sigma.py rename scripts/{convert_prx_to_diffusers.py => recipes/prx.py} (73%) rename scripts/{convert_rae_to_diffusers.py => recipes/rae.py} (76%) rename scripts/{convert_sana_to_diffusers.py => recipes/sana.py} (55%) create mode 100644 scripts/recipes/sana_controlnet.py rename scripts/{convert_sana_video_to_diffusers.py => recipes/sana_video.py} (56%) create mode 100644 scripts/recipes/shap_e.py create mode 100644 scripts/recipes/skyreels_v2.py create mode 100644 scripts/recipes/stable_audio.py rename scripts/{convert_stable_audio_3_to_diffusers.py => recipes/stable_audio3.py} (54%) rename scripts/{convert_stable_cascade.py => recipes/stable_cascade.py} (95%) rename scripts/{convert_stable_cascade_lite.py => recipes/stable_cascade_lite.py} (95%) rename scripts/{convert_original_stable_diffusion_to_diffusers.py => recipes/stable_diffusion.py} (100%) create mode 100644 scripts/recipes/unclip.py rename scripts/{convert_unclip_txt2img_to_image_variation.py => recipes/unclip_image_variation.py} (100%) create mode 100644 scripts/recipes/unidiffuser.py create mode 100644 scripts/recipes/versatile_diffusion.py create mode 100644 scripts/recipes/vq_diffusion.py create mode 100644 scripts/recipes/wan.py create mode 100644 scripts/recipes/wuerstchen.py create mode 100644 scripts/recipes/zero123.py create mode 100644 src/diffusers/loaders/conversion/README.md create mode 100644 src/diffusers/loaders/conversion/__init__.py create mode 100644 src/diffusers/loaders/conversion/ace_step.py create mode 100644 src/diffusers/loaders/conversion/ace_step_conditioner.py create mode 100644 src/diffusers/loaders/conversion/ace_step_detokenizer.py create mode 100644 src/diffusers/loaders/conversion/ace_step_tokenizer.py create mode 100644 src/diffusers/loaders/conversion/anima_conditioner.py create mode 100644 src/diffusers/loaders/conversion/animatediff.py create mode 100644 src/diffusers/loaders/conversion/anyflow.py create mode 100644 src/diffusers/loaders/conversion/anyflow_far.py create mode 100644 src/diffusers/loaders/conversion/asymmetric_vae.py create mode 100644 src/diffusers/loaders/conversion/audioldm2_projection.py create mode 100644 src/diffusers/loaders/conversion/audioldm2_unet.py create mode 100644 src/diffusers/loaders/conversion/auraflow.py create mode 100644 src/diffusers/loaders/conversion/autoencoder_dc.py create mode 100644 src/diffusers/loaders/conversion/blip_qformer.py create mode 100644 src/diffusers/loaders/conversion/checkpoint.py create mode 100644 src/diffusers/loaders/conversion/chroma.py create mode 100644 src/diffusers/loaders/conversion/chronoedit.py create mode 100644 src/diffusers/loaders/conversion/clap_audio.py create mode 100644 src/diffusers/loaders/conversion/clap_text.py create mode 100644 src/diffusers/loaders/conversion/clip.py create mode 100644 src/diffusers/loaders/conversion/clip_vision.py create mode 100644 src/diffusers/loaders/conversion/cogvideox.py create mode 100644 src/diffusers/loaders/conversion/cogview3plus.py create mode 100644 src/diffusers/loaders/conversion/cogview4.py create mode 100644 src/diffusers/loaders/conversion/cogview4_megatron.py create mode 100644 src/diffusers/loaders/conversion/configs/__init__.py create mode 100644 src/diffusers/loaders/conversion/configs/anima.py create mode 100644 src/diffusers/loaders/conversion/configs/audioldm.py create mode 100644 src/diffusers/loaders/conversion/configs/audioldm2.py create mode 100644 src/diffusers/loaders/conversion/configs/blip_diffusion.py create mode 100644 src/diffusers/loaders/conversion/configs/catalog.json create mode 100644 src/diffusers/loaders/conversion/configs/consistency.py create mode 100644 src/diffusers/loaders/conversion/configs/cosmos.py create mode 100644 src/diffusers/loaders/conversion/configs/deepfloyd_if.py create mode 100644 src/diffusers/loaders/conversion/configs/flux2.py create mode 100644 src/diffusers/loaders/conversion/configs/gligen.py create mode 100644 src/diffusers/loaders/conversion/configs/hunyuan_video.py create mode 100644 src/diffusers/loaders/conversion/configs/hunyuan_video15.py create mode 100644 src/diffusers/loaders/conversion/configs/joyimage.py create mode 100644 src/diffusers/loaders/conversion/configs/kandinsky.py create mode 100644 src/diffusers/loaders/conversion/configs/ltx.py create mode 100644 src/diffusers/loaders/conversion/configs/ltx2.py create mode 100644 src/diffusers/loaders/conversion/configs/minimax_h3.py create mode 100644 src/diffusers/loaders/conversion/configs/musicldm.py create mode 100644 src/diffusers/loaders/conversion/configs/prx.py create mode 100644 src/diffusers/loaders/conversion/configs/rae.py create mode 100644 src/diffusers/loaders/conversion/configs/shap_e.py create mode 100644 src/diffusers/loaders/conversion/configs/skyreels_v2.py create mode 100644 src/diffusers/loaders/conversion/configs/stable_audio3.py create mode 100644 src/diffusers/loaders/conversion/configs/unclip.py create mode 100644 src/diffusers/loaders/conversion/configs/unidiffuser.py create mode 100644 src/diffusers/loaders/conversion/configs/versatile_diffusion.py create mode 100644 src/diffusers/loaders/conversion/configs/wan.py create mode 100644 src/diffusers/loaders/conversion/configs/zero123.py create mode 100644 src/diffusers/loaders/conversion/consistency_decoder.py create mode 100644 src/diffusers/loaders/conversion/controlnet.py create mode 100644 src/diffusers/loaders/conversion/core.py create mode 100644 src/diffusers/loaders/conversion/cosmos.py create mode 100644 src/diffusers/loaders/conversion/cosmos_controlnet.py create mode 100644 src/diffusers/loaders/conversion/cosmos_vae.py create mode 100644 src/diffusers/loaders/conversion/dance_diffusion.py create mode 100644 src/diffusers/loaders/conversion/diffuser_rl.py create mode 100644 src/diffusers/loaders/conversion/dit.py create mode 100644 src/diffusers/loaders/conversion/ernie_image.py create mode 100644 src/diffusers/loaders/conversion/flux.py create mode 100644 src/diffusers/loaders/conversion/flux2.py create mode 100644 src/diffusers/loaders/conversion/flux2_vae.py create mode 100644 src/diffusers/loaders/conversion/flux_ip_adapter.py create mode 100644 src/diffusers/loaders/conversion/gligen.py create mode 100644 src/diffusers/loaders/conversion/hidream.py create mode 100644 src/diffusers/loaders/conversion/hifigan.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_dit.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_dit_controlnet.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_image.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_image_refiner_vae.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_image_vae.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_video.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_video15.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_video15_vae.py create mode 100644 src/diffusers/loaders/conversion/hunyuan_video_vae.py create mode 100644 src/diffusers/loaders/conversion/i2vgen_xl.py create mode 100644 src/diffusers/loaders/conversion/if_unet.py create mode 100644 src/diffusers/loaders/conversion/io.py create mode 100644 src/diffusers/loaders/conversion/joy_image.py create mode 100644 src/diffusers/loaders/conversion/k_upscaler.py create mode 100644 src/diffusers/loaders/conversion/kandinsky3.py create mode 100644 src/diffusers/loaders/conversion/ldm_unet.py create mode 100644 src/diffusers/loaders/conversion/ldm_vae.py create mode 100644 src/diffusers/loaders/conversion/longcat_audio.py create mode 100644 src/diffusers/loaders/conversion/longcat_audio_vae.py create mode 100644 src/diffusers/loaders/conversion/lora.py create mode 100644 src/diffusers/loaders/conversion/ltx.py create mode 100644 src/diffusers/loaders/conversion/ltx2.py create mode 100644 src/diffusers/loaders/conversion/ltx2_audio_vae.py create mode 100644 src/diffusers/loaders/conversion/ltx2_connectors.py create mode 100644 src/diffusers/loaders/conversion/ltx2_diffusion_decoder.py create mode 100644 src/diffusers/loaders/conversion/ltx2_duration.py create mode 100644 src/diffusers/loaders/conversion/ltx2_upsampler.py create mode 100644 src/diffusers/loaders/conversion/ltx2_vae.py create mode 100644 src/diffusers/loaders/conversion/ltx2_vocoder.py create mode 100644 src/diffusers/loaders/conversion/ltx_upsampler.py create mode 100644 src/diffusers/loaders/conversion/ltx_vae.py create mode 100644 src/diffusers/loaders/conversion/lumina.py create mode 100644 src/diffusers/loaders/conversion/lumina2.py create mode 100644 src/diffusers/loaders/conversion/minimax_h3.py create mode 100644 src/diffusers/loaders/conversion/minimax_h3_audio_vae.py create mode 100644 src/diffusers/loaders/conversion/minimax_h3_vae.py create mode 100644 src/diffusers/loaders/conversion/minimax_music3.py create mode 100644 src/diffusers/loaders/conversion/minimax_music3_conditioner.py create mode 100644 src/diffusers/loaders/conversion/minimax_music3_rvq.py create mode 100644 src/diffusers/loaders/conversion/minimax_music3_vocoder.py create mode 100644 src/diffusers/loaders/conversion/mochi.py create mode 100644 src/diffusers/loaders/conversion/mochi_vae.py create mode 100644 src/diffusers/loaders/conversion/motif_video.py create mode 100644 src/diffusers/loaders/conversion/movq.py create mode 100644 src/diffusers/loaders/conversion/ncsnpp.py create mode 100644 src/diffusers/loaders/conversion/omnigen.py create mode 100644 src/diffusers/loaders/conversion/oobleck.py create mode 100644 src/diffusers/loaders/conversion/ovis_image.py create mode 100644 src/diffusers/loaders/conversion/paella.py create mode 100644 src/diffusers/loaders/conversion/pipeline.py create mode 100644 src/diffusers/loaders/conversion/pipeline_components.py create mode 100644 src/diffusers/loaders/conversion/pixart.py create mode 100644 src/diffusers/loaders/conversion/presets/asymmetric-vae-1.5.json create mode 100644 src/diffusers/loaders/conversion/presets/asymmetric-vae-2.json create mode 100644 src/diffusers/loaders/conversion/presets/diffuser-rl-128.json create mode 100644 src/diffusers/loaders/conversion/presets/diffuser-rl-32.json create mode 100644 src/diffusers/loaders/conversion/presets/diffuser-rl-value.json create mode 100644 src/diffusers/loaders/conversion/presets/modelscope-text-to-video.json create mode 100644 src/diffusers/loaders/conversion/presets/sd-vae.json create mode 100644 src/diffusers/loaders/conversion/presets/tiny-vae.json create mode 100644 src/diffusers/loaders/conversion/prior.py create mode 100644 src/diffusers/loaders/conversion/prx.py create mode 100644 src/diffusers/loaders/conversion/qwen3.py create mode 100644 src/diffusers/loaders/conversion/qwen_image.py create mode 100644 src/diffusers/loaders/conversion/qwen_image_vae.py create mode 100644 src/diffusers/loaders/conversion/rae.py create mode 100644 src/diffusers/loaders/conversion/registry.py create mode 100644 src/diffusers/loaders/conversion/same.py create mode 100644 src/diffusers/loaders/conversion/sana.py create mode 100644 src/diffusers/loaders/conversion/sana_controlnet.py create mode 100644 src/diffusers/loaders/conversion/sana_video.py create mode 100644 src/diffusers/loaders/conversion/sd3.py create mode 100644 src/diffusers/loaders/conversion/sd3_controlnet.py create mode 100644 src/diffusers/loaders/conversion/shap_e_renderer.py create mode 100644 src/diffusers/loaders/conversion/shap_e_tables.py create mode 100644 src/diffusers/loaders/conversion/skyreels_v2.py create mode 100644 src/diffusers/loaders/conversion/source.py create mode 100644 src/diffusers/loaders/conversion/sparse_controlnet.py create mode 100644 src/diffusers/loaders/conversion/spectrogram_continuous.py create mode 100644 src/diffusers/loaders/conversion/spectrogram_notes.py create mode 100644 src/diffusers/loaders/conversion/stable_audio.py create mode 100644 src/diffusers/loaders/conversion/stable_audio3.py create mode 100644 src/diffusers/loaders/conversion/stable_audio3_duration.py create mode 100644 src/diffusers/loaders/conversion/stable_audio_projection.py create mode 100644 src/diffusers/loaders/conversion/stable_cascade.py create mode 100644 src/diffusers/loaders/conversion/svd.py create mode 100644 src/diffusers/loaders/conversion/svd_vae.py create mode 100644 src/diffusers/loaders/conversion/t2i_adapter.py create mode 100644 src/diffusers/loaders/conversion/t5.py create mode 100644 src/diffusers/loaders/conversion/t5_film.py create mode 100644 src/diffusers/loaders/conversion/tiny_vae.py create mode 100644 src/diffusers/loaders/conversion/transforms.py create mode 100644 src/diffusers/loaders/conversion/umt5.py create mode 100644 src/diffusers/loaders/conversion/unclip_text_projection.py create mode 100644 src/diffusers/loaders/conversion/unclip_unet.py create mode 100644 src/diffusers/loaders/conversion/unet_1d.py create mode 100644 src/diffusers/loaders/conversion/unet_2d.py create mode 100644 src/diffusers/loaders/conversion/unet_3d.py create mode 100644 src/diffusers/loaders/conversion/unidiffuser.py create mode 100644 src/diffusers/loaders/conversion/unidiffuser_text.py create mode 100644 src/diffusers/loaders/conversion/uvit.py create mode 100644 src/diffusers/loaders/conversion/versatile_image_unet.py create mode 100644 src/diffusers/loaders/conversion/versatile_text_unet.py create mode 100644 src/diffusers/loaders/conversion/vq_diffusion.py create mode 100644 src/diffusers/loaders/conversion/vq_model.py create mode 100644 src/diffusers/loaders/conversion/wan.py create mode 100644 src/diffusers/loaders/conversion/wan_animate.py create mode 100644 src/diffusers/loaders/conversion/wan_animate2.py create mode 100644 src/diffusers/loaders/conversion/wan_vace.py create mode 100644 src/diffusers/loaders/conversion/wan_vae.py create mode 100644 src/diffusers/loaders/conversion/wuerstchen_decoder.py create mode 100644 src/diffusers/loaders/conversion/wuerstchen_prior.py create mode 100644 src/diffusers/loaders/conversion/z_image.py create mode 100644 src/diffusers/loaders/conversion/z_image_controlnet.py create mode 100644 src/diffusers/loaders/conversion/zero123_projection.py create mode 100644 tests/others/test_checkpoint_conversion.py create mode 100644 tests/others/test_conversion_io.py create mode 100644 tests/others/test_pipeline_conversion.py create mode 100644 tests/single_file/test_auxiliary_conversions.py create mode 100644 tests/single_file/test_cogvideox_conversion.py create mode 100644 tests/single_file/test_model_conversions.py create mode 100644 tests/single_file/test_text_encoder_conversions.py diff --git a/MANIFEST.in b/MANIFEST.in index b22fe1a28a1e..d9c4907fc34a 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,4 @@ include LICENSE include src/diffusers/utils/model_card_template.md +recursive-include src/diffusers/loaders/conversion/configs *.json +recursive-include src/diffusers/loaders/conversion/presets *.json diff --git a/docs/source/en/api/loaders/single_file.md b/docs/source/en/api/loaders/single_file.md index 94fcda9f1765..2655252eed6e 100644 --- a/docs/source/en/api/loaders/single_file.md +++ b/docs/source/en/api/loaders/single_file.md @@ -46,12 +46,27 @@ The [`~loaders.FromSingleFileMixin.from_single_file`] method allows you to load: ## Supported models -- [`UNet2DConditionModel`] -- [`StableCascadeUNet`] -- [`AutoencoderKL`] -- [`ControlNetModel`] -- [`SD3Transformer2DModel`] -- [`FluxTransformer2DModel`] +Diffusers model components with a registered bidirectional conversion expose +[`~loaders.FromOriginalModelMixin.from_single_file`], including transformers, UNets, autoencoders, ControlNets, +adapters, and audio components. See the [component registry](https://github.com/huggingface/diffusers/tree/main/src/diffusers/loaders/conversion) +for the supported classes and original format variants. + +Pass the matching Diffusers component configuration explicitly when loading a component whose architecture cannot +be inferred from its checkpoint. Existing Stable Diffusion and other established checkpoint detection paths remain +available. + +```python +from diffusers import DiTTransformer2DModel + +transformer = DiTTransformer2DModel.from_single_file( + "./original-transformer.safetensors", + config="./dit/transformer", + local_files_only=True, +) +``` + +Component support does not add single-file loading to its full pipeline. Transformers-owned text/vision encoders and +LoRA layouts use the conversion API or their existing pipeline loaders. ## FromSingleFileMixin diff --git a/docs/source/en/api/pipelines/stable_audio_3.md b/docs/source/en/api/pipelines/stable_audio_3.md index 358de2df5a25..6ed51bee9373 100644 --- a/docs/source/en/api/pipelines/stable_audio_3.md +++ b/docs/source/en/api/pipelines/stable_audio_3.md @@ -40,7 +40,7 @@ The Stability AI checkpoints are not published in diffusers format, so convert t checkpoint's `model_config.json` and selects the scheduler from its `diffusion_objective`: ```bash -python scripts/convert_stable_audio_3_to_diffusers.py \ +python scripts/build_pipeline.py stable_audio3 \ --checkpoint_path stabilityai/stable-audio-3-medium-base \ --text_encoder_repo google/t5gemma-b-b-ul2 \ --output_dir /tmp/sa3-diffusers-euler \ diff --git a/docs/source/en/using-diffusers/other-formats.md b/docs/source/en/using-diffusers/other-formats.md index a4a7620d587b..69bda6c77e27 100644 --- a/docs/source/en/using-diffusers/other-formats.md +++ b/docs/source/en/using-diffusers/other-formats.md @@ -222,12 +222,12 @@ pipeline = DiffusionPipeline.from_single_file( Diffusers provides scripts and methods to convert format and files to enable broader support across the diffusion ecosystem. -Take a look at the [diffusers/scripts](https://github.com/huggingface/diffusers/tree/main/scripts) folder to find a conversion script. Scripts with `"to_diffusers` appended at the end converts a model to the Diffusers format. Each script has a specific set of arguments for configuring the conversion. Make sure you check what arguments are available. +Use the shared component converter below for tensor layout conversion. The [scripts guide](https://github.com/huggingface/diffusers/blob/main/scripts/README.md) describes the remaining scripts for complete pipeline assembly, source preparation, and graph export. The example below converts a model stored in Diffusers format to a single-file format. Provide the path to the model to convert and where to save the converted model. You can optionally specify what file type and data type to save the model as. ```bash -python convert_diffusers_to_original_sdxl.py --model_path path/to/model/to/convert --checkpoint_path path/to/save/model/to --use_safetensors +python scripts/export_pipeline_checkpoint.py --input ./sdxl --output ./sdxl.safetensors ``` The [`~DiffusionPipeline.save_pretrained`] method also saves a model in Diffusers format and takes care of creating subfolders for each model. It saves the files as safetensor files by default. @@ -243,6 +243,114 @@ pipeline.save_pretrained() Finally, you can use a Space like [SD To Diffusers](https://hf.co/spaces/diffusers/sd-to-diffusers) or [SD-XL To Diffusers](https://hf.co/spaces/diffusers/sdxl-to-diffusers) to convert models to the Diffusers format. It'll open a PR on your model repository with the converted files. This is the easiest way to convert a model, but it may fail for more complicated models. Using a conversion script is more reliable. +### Bidirectional component conversion + +The `diffusers.loaders.conversion` package defines each component's tensor layout once and uses the same definition in +both directions. It covers the model component classes used by the conversion scripts and single-file model loader, +including transformers, UNets, autoencoders, ControlNets, adapters, text/audio encoders, and LoRA layouts. List the +registered classes from a repository checkout with: + +```bash +python scripts/convert_checkpoint.py --list-models +``` + +Use the matching **Diffusers component configuration** in either direction. The conversion does not instantiate a +model, change tensor dtypes, or need a record of an earlier import: + +```python +from diffusers.loaders.conversion import get_conversion + +conversion = get_conversion("FluxTransformer2DModel", transformer_config) +original_weights = conversion.to_original(diffusers_weights) +diffusers_weights = conversion.to_diffusers(original_weights) +``` + +The shared file converter accepts a local component directory, safetensors/PyTorch file, or shard index. It writes a +new component directory with safetensors weights, optionally sharded. For example: + +```bash +python scripts/convert_checkpoint.py --direction to-original \ + --input ./flux/transformer --output ./flux-original + +python scripts/convert_checkpoint.py --direction to-diffusers \ + --input ./flux-original --config ./flux-original/conversion_config.json \ + --output ./flux-restored +``` + +For bundled original inputs, `--input-prefix model.diffusion_model.` selects and strips a component prefix. +`--input-wrapper state_dict` selects a nested PyTorch dictionary; repeat the option for nested wrappers. Set +`--output-prefix` to prepend an original component prefix, and `--max-shard-size` to set a shard limit in bytes. +One tensor can exceed that limit. Existing outputs are not overwritten. + +Use `--output-format pytorch` to write a single checkpoint file instead of a directory, and repeat +`--output-wrapper` to nest its state dict. For example, the original SAT transformer container uses: + +```bash +python scripts/convert_checkpoint.py --direction to-original --input ./cogvideox/transformer \ + --output ./mp_rank_00_model_states.pt --output-format pytorch \ + --output-wrapper module --output-prefix model.diffusion_model. +``` + +The PyTorch exporter regenerates CogVideoX's fixed positional embedding from config. A SAT VAE uses +`--output-wrapper state_dict` with no prefix. See the [scripts migration guide](https://github.com/huggingface/diffusers/blob/main/scripts/README.md) +for the component commands replaced by this shared converter and the pipeline preparation scripts that remain. + +Use `--list-presets` to inspect reusable configuration presets and original-config helpers. Select one with `--preset`, +pass helper keyword arguments as JSON with `--preset-args`, and specify `--model-class` when the config has no class name. +For components stored across several files, `--input-manifest` accepts a JSON list of namespaced sources. See the +[scripts guide](https://github.com/huggingface/diffusers/blob/main/scripts/README.md) for manifest examples, original +runtime requirements, and the complete command migration table. Complete pipeline assembly uses `build_pipeline.py`; +SD/SDXL checkpoint packaging uses `export_pipeline_checkpoint.py`, and adapter fusion uses `merge_lora.py`. + +Some architectures have several source layouts. Set `original_format` in the config, or pass `--original-format` to +the script. Examples include CLIP/OpenCLIP, Cosmos 1/2, CogView4 Megatron, and MiniMax H3's tensor-parallel shard layout. +See the [component definitions and format notes](https://github.com/huggingface/diffusers/tree/main/src/diffusers/loaders/conversion) +for the supported choices. + +These APIs convert component tensor layouts. The generic script expects the canonical keys declared by the selected +definition after wrapper/prefix selection; it does not download external encoder weights, assemble tensor-parallel +ranks, or reconstruct original runtime configuration, tokenizers, training state, or pipeline packaging. Exports carry +`conversion_config.json` for repeat conversion, rather than an original runtime config. Shared source preparation handles known auxiliary state such as training counters; such state is not recoverable from Diffusers weights. The generic command can also write original PyTorch containers, including SAT checkpoints. + +Most operations preserve tensors exactly. A conversion with `lossless=False` performs a documented normalization, +such as folding LTX2 decoder gates into linear weights; its inverse emits a canonical factorization. Shared original +parameters can only be exported when their separate Diffusers copies agree. LoRA conversion preserves the factors and +alpha tensors; it does not merge adapters into base weights. Quantized packed weights and graph exports such as ONNX +or TensorRT are outside this tensor conversion API. + +### Defining a reversible conversion + +Single-file model loading and the shared file converter use the registry in `diffusers.loaders.conversion`. +To define another component conversion, use `Conversion` for exact key renames and `Rule` for grouped tensor operations: + +```python +from diffusers.loaders.conversion import Conversion, Rule, Split + +conversion = Conversion( + mapping={"time_embed.0.weight": "time_embedding.linear_1.weight"}, + rules=( + Rule( + original=("attention.qkv.weight",), + diffusers=("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight"), + transform=Split((64, 64, 64)), + ), + ), +) + +# Supply a component state dict containing exactly the keys declared above. +diffusers_weights = conversion.to_diffusers(original_weights) +original_weights = conversion.to_original(diffusers_weights) +``` + +Build exact keys and split sizes from the component config with ordinary Python loops. A custom tensor transform +implements `forward(tensors)` and `inverse(tensors)`, each returning an ordered tuple of tensors. Transform inputs must +not be modified, and dtype/device must be preserved. Outputs may share storage with inputs. + +Each original key and each Diffusers key must appear once in the definition. Missing keys, unknown keys, duplicate +destinations, and incompatible split shapes raise errors. Handle component prefixes and known auxiliary keys before +calling the conversion. Conversion definitions describe tensor layouts; configuration conversion, checkpoint wrappers, +and file I/O remain the responsibility of the model's format-specific code. + ## Resources - Learn more about the design decisions and why safetensor files are preferred for saving and loading model weights in the [Safetensors audited as really safe and becoming the default](https://blog.eleuther.ai/safetensors-security-audit/) blog post. diff --git a/docs/source/ko/using-diffusers/other-formats.md b/docs/source/ko/using-diffusers/other-formats.md index 41f07818f5e6..994211fbba1a 100644 --- a/docs/source/ko/using-diffusers/other-formats.md +++ b/docs/source/ko/using-diffusers/other-formats.md @@ -71,7 +71,7 @@ git checkout pr/13 4. 이제 스크립트를 실행하여 .ckpt 파일을 변환할 수 있습니다: ```bash -python ../diffusers/scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path temporalnetv3.ckpt --original_config_file cldm_v15.yaml --dump_path ./ --controlnet +python ../diffusers/scripts/build_pipeline.py stable_diffusion --checkpoint_path temporalnetv3.ckpt --original_config_file cldm_v15.yaml --dump_path ./ --controlnet ``` 5. 변환이 완료되면 변환된 모델을 업로드하고 결과물을 pull request [pull request](https://huggingface.co/CiaraRowles/TemporalNet/discussions/13)를 테스트하세요! diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 000000000000..af3050c922b3 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,188 @@ +# Checkpoint conversion and pipeline assembly + +[`convert_checkpoint.py`](convert_checkpoint.py) is the single component tensor-conversion command. Every model uses +the same reversible `Conversion(mapping=..., rules=...)` definition in the library, single-file loader, and recipes. +Use the matching Diffusers component configuration; tensor conversion preserves dtypes and does not construct a model. + +```bash +python scripts/convert_checkpoint.py --list-models +python scripts/convert_checkpoint.py --direction to-diffusers \ + --input ./original.safetensors --config ./transformer/config.json --output ./converted +python scripts/convert_checkpoint.py --direction to-original --input ./converted --output ./original-export +``` + +The command accepts local tensor files, component directories, and shard indexes. Use `--input-prefix` to select a +component (for example `model.diffusion_model.`), and repeat `--input-wrapper` for nested PyTorch dictionaries. +Known component-specific auxiliary state is handled by shared source preparation. Unknown component tensors fail +coverage validation. Download Hub inputs before invoking this local file command. + +## Configuration and source preparation + +Configuration helpers and presets live in [`diffusers.loaders.conversion.configs`](../src/diffusers/loaders/conversion/configs). +List them with `--list-presets`. Select a JSON preset, or call a named helper with JSON keyword arguments. Use +`--preset-component` to select one entry from a variant collection. Supply `--model-class` if the config has no `_class_name`. + +```bash +python scripts/convert_checkpoint.py --direction to-diffusers \ + --input ./wan.safetensors --output ./wan-transformer --model-class WanTransformer3DModel \ + --preset wan.get_transformer_config --preset-args '{"model_type":"Wan-T2V-1.3B"}' +``` + +Use `--input-manifest sources.json` when a component spans several files. Paths are relative to the manifest. +Each source can select an `input_prefix`, add an `output_prefix`, and specify a list of dictionary `wrapper` keys. +Duplicate resulting keys are rejected. For example, a tiny VAE's separate encoder and decoder files use: + +```json +{ + "sources": [ + {"path": "encoder.pth", "output_prefix": "encoder."}, + {"path": "decoder.pth", "output_prefix": "decoder."} + ] +} +``` + +```bash +python scripts/convert_checkpoint.py --direction to-diffusers \ + --input-manifest sources.json --preset tiny-vae --output ./tiny-vae +``` + +The default source `format` is `tensors`. `torchscript` extracts an archive's state dict. The explicit `python-model` +format extracts a pickled module's state dict (including wrapped modules, such as legacy Diffuser RL checkpoints). +Use `python-model` only with trusted files: it executes pickle code and requires the original Python model package. +Tensor-parallel rank concatenation and external JAX/TensorFlow/reference-model preparation remain in the relevant +pipeline recipe; the manifest combines disjoint keys, not tensor-parallel slices. + +## Containers, pipelines, adapters, and graph export + +Safetensors output is a new component directory, optionally sharded with `--max-shard-size`. PyTorch output is a +single file: use `--output-format pytorch`, repeated `--output-wrapper`, and `--output-prefix` for its container. +The original SAT transformer container, including regenerated fixed CogVideoX positional embeddings, uses: + +```bash +python scripts/convert_checkpoint.py --direction to-original --input ./cogvideox/transformer \ + --output ./mp_rank_00_model_states.pt --output-format pytorch \ + --output-wrapper module --output-prefix model.diffusion_model. +``` + +A SAT VAE uses its VAE component directory, `--output-wrapper state_dict`, and no prefix. Original runtime settings +and separately loaded assets still come from the original implementation. `conversion_config.json` records the +Diffusers component configuration for repeat conversion. + +[`build_pipeline.py`](build_pipeline.py) runs the preserved assembly recipes. Recipes prepare original containers, +resolve model configurations, gather encoders/tokenizers, and construct pipelines; component tensor layouts come from +the shared conversion package. Existing recipe-specific arguments are retained. Some recipes require their original +runtime dependencies or GPU hardware. See the [recipe dependency guide](recipes/README.md) and the recipe's help +for requirements and the scope of startup checks. + +```bash +python scripts/build_pipeline.py --list +python scripts/build_pipeline.py cogvideox --help +``` + +[`export_pipeline_checkpoint.py`](export_pipeline_checkpoint.py) packages a local SD/SDXL pipeline's shared component +exports into one original checkpoint. [`merge_lora.py`](merge_lora.py) loads and fuses adapters through the pipeline's +public LoRA API. LoRA factor-format conversion itself uses `convert_checkpoint.py` with an explicit module list; +merging changes base weights and is not reversible. Deployment graph commands are named `export_*`. + +```bash +python scripts/export_pipeline_checkpoint.py --input ./sdxl --output ./sdxl.safetensors +python scripts/merge_lora.py --model ./base --adapter ./adapter.safetensors --output ./merged +``` + +## Migration table + +For component rows, provide `--config` or the listed preset and `--model-class`. Convert each component separately +when several classes are listed. Pipeline rows retain the original assembly behavior through the named recipe. + +| Former command | Replacement | +| --- | --- | +| `convert_ace_step_to_diffusers.py` | `build_pipeline.py ace_step` ([recipe](recipes/ace_step.py)) | +| `convert_amused.py` | `build_pipeline.py amused` ([recipe](recipes/amused.py)) | +| `convert_anima_to_diffusers.py` | `build_pipeline.py anima` ([recipe](recipes/anima.py)) | +| `convert_animatediff_motion_lora_to_diffusers.py` | `convert_checkpoint.py`: LoRA (`original_format=animatediff`) | +| `convert_animatediff_motion_module_to_diffusers.py` | `convert_checkpoint.py`: MotionAdapter | +| `convert_animatediff_sparsectrl_to_diffusers.py` | `convert_checkpoint.py`: SparseControlNetModel | +| `convert_anyflow_to_diffusers.py` | `build_pipeline.py anyflow` ([recipe](recipes/anyflow.py)) | +| `convert_asymmetric_vqgan_to_diffusers.py` | `convert_checkpoint.py`: AsymmetricAutoencoderKL; preset `asymmetric-vae-1.5` or `asymmetric-vae-2` | +| `convert_aura_flow_to_diffusers.py` | `convert_checkpoint.py`: AuraFlowTransformer2DModel | +| `convert_blipdiffusion_to_diffusers.py` | `build_pipeline.py blip_diffusion` ([recipe](recipes/blip_diffusion.py)) | +| `convert_cogvideox_to_diffusers.py` | `build_pipeline.py cogvideox` ([recipe](recipes/cogvideox.py)) | +| `convert_cogview3_to_diffusers.py` | `build_pipeline.py cogview3` ([recipe](recipes/cogview3.py)) | +| `convert_cogview4_to_diffusers.py` | `build_pipeline.py cogview4` ([recipe](recipes/cogview4.py)) | +| `convert_cogview4_to_diffusers_megatron.py` | `build_pipeline.py cogview4_megatron` ([recipe](recipes/cogview4_megatron.py)) | +| `convert_consistency_decoder.py` | `build_pipeline.py consistency_decoder` ([recipe](recipes/consistency_decoder.py)) | +| `convert_consistency_to_diffusers.py` | `build_pipeline.py consistency` ([recipe](recipes/consistency.py)) | +| `convert_cosmos_to_diffusers.py` | `build_pipeline.py cosmos` ([recipe](recipes/cosmos.py)) | +| `convert_dance_diffusion_to_diffusers.py` | `build_pipeline.py dance_diffusion` ([recipe](recipes/dance_diffusion.py)) | +| `convert_dcae_to_diffusers.py` | `convert_checkpoint.py`: AutoencoderDC | +| `convert_ddpm_original_checkpoint_to_diffusers.py` | `build_pipeline.py ddpm` ([recipe](recipes/ddpm.py)) | +| `convert_diffusers_sdxl_lora_to_webui.py` | `convert_checkpoint.py`: LoRA (`original_format=kohya`) | +| `convert_diffusers_to_original_sdxl.py` | `export_pipeline_checkpoint.py` | +| `convert_diffusers_to_original_stable_diffusion.py` | `export_pipeline_checkpoint.py` | +| `convert_dit_to_diffusers.py` | `build_pipeline.py dit` ([recipe](recipes/dit.py)) | +| `convert_flux2_to_diffusers.py` | `build_pipeline.py flux2` ([recipe](recipes/flux2.py)) | +| `convert_flux_to_diffusers.py` | `convert_checkpoint.py`: FluxTransformer2DModel / AutoencoderKL | +| `convert_flux_xlabs_ipadapter_to_diffusers.py` | `build_pipeline.py flux_ip_adapter` ([recipe](recipes/flux_ip_adapter.py)) | +| `convert_gligen_to_diffusers.py` | `build_pipeline.py gligen` ([recipe](recipes/gligen.py)) | +| `convert_hunyuan_image_to_diffusers.py` | `build_pipeline.py hunyuan_image` ([recipe](recipes/hunyuan_image.py)) | +| `convert_hunyuan_video1_5_to_diffusers.py` | `build_pipeline.py hunyuan_video15` ([recipe](recipes/hunyuan_video15.py)) | +| `convert_hunyuan_video_to_diffusers.py` | `build_pipeline.py hunyuan_video` ([recipe](recipes/hunyuan_video.py)) | +| `convert_hunyuandit_controlnet_to_diffusers.py` | `convert_checkpoint.py`: HunyuanDiT2DControlNetModel | +| `convert_hunyuandit_to_diffusers.py` | `build_pipeline.py hunyuan_dit` ([recipe](recipes/hunyuan_dit.py)) | +| `convert_i2vgen_to_diffusers.py` | `build_pipeline.py i2vgen` ([recipe](recipes/i2vgen.py)) | +| `convert_if.py` | `build_pipeline.py deepfloyd_if` ([recipe](recipes/deepfloyd_if.py)) | +| `convert_joyimage_edit_to_diffusers.py` | `build_pipeline.py joyimage` ([recipe](recipes/joyimage.py)) | +| `convert_k_upscaler_to_diffusers.py` | `build_pipeline.py k_upscaler` ([recipe](recipes/k_upscaler.py)) | +| `convert_kakao_brain_unclip_to_diffusers.py` | `build_pipeline.py unclip` ([recipe](recipes/unclip.py)) | +| `convert_kandinsky3_unet.py` | `convert_checkpoint.py`: Kandinsky3UNet | +| `convert_kandinsky_to_diffusers.py` | `build_pipeline.py kandinsky` ([recipe](recipes/kandinsky.py)) | +| `convert_ldm_original_checkpoint_to_diffusers.py` | `build_pipeline.py ldm` ([recipe](recipes/ldm.py)) | +| `convert_longcat_audio_dit_to_diffusers.py` | `build_pipeline.py longcat_audio` ([recipe](recipes/longcat_audio.py)) | +| `convert_lora_safetensor_to_diffusers.py` | `merge_lora.py` | +| `convert_ltx2_to_diffusers.py` | `build_pipeline.py ltx2` ([recipe](recipes/ltx2.py)) | +| `convert_ltx_to_diffusers.py` | `build_pipeline.py ltx` ([recipe](recipes/ltx.py)) | +| `convert_lumina_to_diffusers.py` | `build_pipeline.py lumina` ([recipe](recipes/lumina.py)) | +| `convert_minimax_h3_to_diffusers.py` | `build_pipeline.py minimax_h3` ([recipe](recipes/minimax_h3.py)) | +| `convert_minimax_music3_to_diffusers.py` | `build_pipeline.py minimax_music3` ([recipe](recipes/minimax_music3.py)) | +| `convert_mochi_to_diffusers.py` | `build_pipeline.py mochi` ([recipe](recipes/mochi.py)) | +| `convert_models_diffuser_to_diffusers.py` | `convert_checkpoint.py`: UNet1DModel; preset `diffuser-rl-32`, `diffuser-rl-128`, or `diffuser-rl-value` | +| `convert_ms_text_to_video_to_diffusers.py` | `convert_checkpoint.py`: UNet3DConditionModel; preset `modelscope-text-to-video` | +| `convert_music_spectrogram_to_diffusers.py` | `build_pipeline.py music_spectrogram` ([recipe](recipes/music_spectrogram.py)) | +| `convert_ncsnpp_original_checkpoint_to_diffusers.py` | `build_pipeline.py ncsnpp` ([recipe](recipes/ncsnpp.py)) | +| `convert_omnigen_to_diffusers.py` | `build_pipeline.py omnigen` ([recipe](recipes/omnigen.py)) | +| `convert_original_audioldm2_to_diffusers.py` | `build_pipeline.py audioldm2` ([recipe](recipes/audioldm2.py)) | +| `convert_original_audioldm_to_diffusers.py` | `build_pipeline.py audioldm` ([recipe](recipes/audioldm.py)) | +| `convert_original_controlnet_to_diffusers.py` | `build_pipeline.py controlnet` ([recipe](recipes/controlnet.py)) | +| `convert_original_musicldm_to_diffusers.py` | `build_pipeline.py musicldm` ([recipe](recipes/musicldm.py)) | +| `convert_original_stable_diffusion_to_diffusers.py` | `build_pipeline.py stable_diffusion` ([recipe](recipes/stable_diffusion.py)) | +| `convert_original_t2i_adapter.py` | `convert_checkpoint.py`: T2IAdapter | +| `convert_ovis_image_to_diffusers.py` | `convert_checkpoint.py`: OvisImageTransformer2DModel | +| `convert_pixart_alpha_to_diffusers.py` | `build_pipeline.py pixart_alpha` ([recipe](recipes/pixart_alpha.py)) | +| `convert_pixart_sigma_to_diffusers.py` | `build_pipeline.py pixart_sigma` ([recipe](recipes/pixart_sigma.py)) | +| `convert_prx_to_diffusers.py` | `build_pipeline.py prx` ([recipe](recipes/prx.py)) | +| `convert_rae_to_diffusers.py` | `build_pipeline.py rae` ([recipe](recipes/rae.py)) | +| `convert_sana_controlnet_to_diffusers.py` | `build_pipeline.py sana_controlnet` ([recipe](recipes/sana_controlnet.py)) | +| `convert_sana_to_diffusers.py` | `build_pipeline.py sana` ([recipe](recipes/sana.py)) | +| `convert_sana_video_to_diffusers.py` | `build_pipeline.py sana_video` ([recipe](recipes/sana_video.py)) | +| `convert_sd3_controlnet_to_diffusers.py` | `convert_checkpoint.py`: SD3ControlNetModel | +| `convert_sd3_to_diffusers.py` | `convert_checkpoint.py`: SD3Transformer2DModel / AutoencoderKL | +| `convert_shap_e_to_diffusers.py` | `build_pipeline.py shap_e` ([recipe](recipes/shap_e.py)) | +| `convert_skyreelsv2_to_diffusers.py` | `build_pipeline.py skyreels_v2` ([recipe](recipes/skyreels_v2.py)) | +| `convert_stable_audio.py` | `build_pipeline.py stable_audio` ([recipe](recipes/stable_audio.py)) | +| `convert_stable_audio_3_to_diffusers.py` | `build_pipeline.py stable_audio3` ([recipe](recipes/stable_audio3.py)) | +| `convert_stable_cascade.py` | `build_pipeline.py stable_cascade` ([recipe](recipes/stable_cascade.py)) | +| `convert_stable_cascade_lite.py` | `build_pipeline.py stable_cascade_lite` ([recipe](recipes/stable_cascade_lite.py)) | +| `convert_stable_diffusion_checkpoint_to_onnx.py` | `export_stable_diffusion_checkpoint_to_onnx.py` | +| `convert_stable_diffusion_controlnet_to_onnx.py` | `export_stable_diffusion_controlnet_to_onnx.py` | +| `convert_stable_diffusion_controlnet_to_tensorrt.py` | `export_stable_diffusion_controlnet_to_tensorrt.py` | +| `convert_svd_to_diffusers.py` | `convert_checkpoint.py`: UNetSpatioTemporalConditionModel / AutoencoderKLTemporalDecoder | +| `convert_tiny_autoencoder_to_diffusers.py` | `convert_checkpoint.py`: AutoencoderTiny; preset `tiny-vae` and source manifest | +| `convert_unclip_txt2img_to_image_variation.py` | `build_pipeline.py unclip_image_variation` ([recipe](recipes/unclip_image_variation.py)) | +| `convert_unidiffuser_to_diffusers.py` | `build_pipeline.py unidiffuser` ([recipe](recipes/unidiffuser.py)) | +| `convert_vae_diff_to_onnx.py` | `export_vae_diff_to_onnx.py` | +| `convert_vae_pt_to_diffusers.py` | `convert_checkpoint.py`: AutoencoderKL; preset `sd-vae` | +| `convert_versatile_diffusion_to_diffusers.py` | `build_pipeline.py versatile_diffusion` ([recipe](recipes/versatile_diffusion.py)) | +| `convert_vq_diffusion_to_diffusers.py` | `build_pipeline.py vq_diffusion` ([recipe](recipes/vq_diffusion.py)) | +| `convert_wan_to_diffusers.py` | `build_pipeline.py wan` ([recipe](recipes/wan.py)) | +| `convert_wuerstchen.py` | `build_pipeline.py wuerstchen` ([recipe](recipes/wuerstchen.py)) | +| `convert_zero123_to_diffusers.py` | `build_pipeline.py zero123` ([recipe](recipes/zero123.py)) | diff --git a/scripts/build_pipeline.py b/scripts/build_pipeline.py new file mode 100644 index 000000000000..64663dd9b737 --- /dev/null +++ b/scripts/build_pipeline.py @@ -0,0 +1,33 @@ +"""Run a pipeline assembly recipe; component tensor mappings live in diffusers.loaders.conversion.""" + +import argparse +import runpy +import sys +from pathlib import Path + + +def main(): + directory = Path(__file__).resolve().parent / "recipes" + recipes = {path.stem: path for path in directory.glob("*.py") if not path.name.startswith("_")} + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("recipe", nargs="?", choices=sorted(recipes)) + parser.add_argument("--list", action="store_true", help="List available recipes") + parser.add_argument("recipe_args", nargs=argparse.REMAINDER) + args = parser.parse_args() + if args.list: + print("\n".join(sorted(recipes))) + return + if not args.recipe: + parser.error("Choose a recipe, or use --list") + previous_argv, previous_path = sys.argv[:], sys.path[:] + try: + sys.argv = [str(recipes[args.recipe]), *args.recipe_args] + sys.path.insert(0, str(directory)) + runpy.run_path(str(recipes[args.recipe]), run_name="__main__") + finally: + sys.argv[:] = previous_argv + sys.path[:] = previous_path + + +if __name__ == "__main__": + main() diff --git a/scripts/convert_amused.py b/scripts/convert_amused.py deleted file mode 100644 index ddd1bf508b6d..000000000000 --- a/scripts/convert_amused.py +++ /dev/null @@ -1,523 +0,0 @@ -import inspect -import os -from argparse import ArgumentParser - -import numpy as np -import torch -from muse import MaskGiTUViT, VQGANModel -from muse import PipelineMuse as OldPipelineMuse -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from diffusers import VQModel -from diffusers.models.attention_processor import AttnProcessor -from diffusers.models.unets.uvit_2d import UVit2DModel -from diffusers.pipelines.amused.pipeline_amused import AmusedPipeline -from diffusers.schedulers import AmusedScheduler - - -torch.backends.cuda.enable_flash_sdp(False) -torch.backends.cuda.enable_mem_efficient_sdp(False) -torch.backends.cuda.enable_math_sdp(True) - -os.environ["CUDA_LAUNCH_BLOCKING"] = "1" -os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" -torch.use_deterministic_algorithms(True) - -# Enable CUDNN deterministic mode -torch.backends.cudnn.deterministic = True -torch.backends.cudnn.benchmark = False -torch.backends.cuda.matmul.allow_tf32 = False - -device = "cuda" - - -def main(): - args = ArgumentParser() - args.add_argument("--model_256", action="store_true") - args.add_argument("--write_to", type=str, required=False, default=None) - args.add_argument("--transformer_path", type=str, required=False, default=None) - args = args.parse_args() - - transformer_path = args.transformer_path - subfolder = "transformer" - - if transformer_path is None: - if args.model_256: - transformer_path = "openMUSE/muse-256" - else: - transformer_path = ( - "../research-run-512-checkpoints/research-run-512-with-downsample-checkpoint-554000/unwrapped_model/" - ) - subfolder = None - - old_transformer = MaskGiTUViT.from_pretrained(transformer_path, subfolder=subfolder) - - old_transformer.to(device) - - old_vae = VQGANModel.from_pretrained("openMUSE/muse-512", subfolder="vae") - old_vae.to(device) - - vqvae = make_vqvae(old_vae) - - tokenizer = CLIPTokenizer.from_pretrained("openMUSE/muse-512", subfolder="text_encoder") - - text_encoder = CLIPTextModelWithProjection.from_pretrained("openMUSE/muse-512", subfolder="text_encoder") - text_encoder.to(device) - - transformer = make_transformer(old_transformer, args.model_256) - - scheduler = AmusedScheduler(mask_token_id=old_transformer.config.mask_token_id) - - new_pipe = AmusedPipeline( - vqvae=vqvae, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, scheduler=scheduler - ) - - old_pipe = OldPipelineMuse( - vae=old_vae, transformer=old_transformer, text_encoder=text_encoder, tokenizer=tokenizer - ) - old_pipe.to(device) - - if args.model_256: - transformer_seq_len = 256 - orig_size = (256, 256) - else: - transformer_seq_len = 1024 - orig_size = (512, 512) - - old_out = old_pipe( - "dog", - generator=torch.Generator(device).manual_seed(0), - transformer_seq_len=transformer_seq_len, - orig_size=orig_size, - timesteps=12, - )[0] - - new_out = new_pipe("dog", generator=torch.Generator(device).manual_seed(0)).images[0] - - old_out = np.array(old_out) - new_out = np.array(new_out) - - diff = np.abs(old_out.astype(np.float64) - new_out.astype(np.float64)) - - # assert diff diff.sum() == 0 - print("skipping pipeline full equivalence check") - - print(f"max diff: {diff.max()}, diff.sum() / diff.size {diff.sum() / diff.size}") - - if args.model_256: - assert diff.max() <= 3 - assert diff.sum() / diff.size < 0.7 - else: - assert diff.max() <= 1 - assert diff.sum() / diff.size < 0.4 - - if args.write_to is not None: - new_pipe.save_pretrained(args.write_to) - - -def make_transformer(old_transformer, model_256): - args = dict(old_transformer.config) - force_down_up_sample = args["force_down_up_sample"] - - signature = inspect.signature(UVit2DModel.__init__) - - args_ = { - "downsample": force_down_up_sample, - "upsample": force_down_up_sample, - "block_out_channels": args["block_out_channels"][0], - "sample_size": 16 if model_256 else 32, - } - - for s in list(signature.parameters.keys()): - if s in ["self", "downsample", "upsample", "sample_size", "block_out_channels"]: - continue - - args_[s] = args[s] - - new_transformer = UVit2DModel(**args_) - new_transformer.to(device) - - new_transformer.set_attn_processor(AttnProcessor()) - - state_dict = old_transformer.state_dict() - - state_dict["cond_embed.linear_1.weight"] = state_dict.pop("cond_embed.0.weight") - state_dict["cond_embed.linear_2.weight"] = state_dict.pop("cond_embed.2.weight") - - for i in range(22): - state_dict[f"transformer_layers.{i}.norm1.norm.weight"] = state_dict.pop( - f"transformer_layers.{i}.attn_layer_norm.weight" - ) - state_dict[f"transformer_layers.{i}.norm1.linear.weight"] = state_dict.pop( - f"transformer_layers.{i}.self_attn_adaLN_modulation.mapper.weight" - ) - - state_dict[f"transformer_layers.{i}.attn1.to_q.weight"] = state_dict.pop( - f"transformer_layers.{i}.attention.query.weight" - ) - state_dict[f"transformer_layers.{i}.attn1.to_k.weight"] = state_dict.pop( - f"transformer_layers.{i}.attention.key.weight" - ) - state_dict[f"transformer_layers.{i}.attn1.to_v.weight"] = state_dict.pop( - f"transformer_layers.{i}.attention.value.weight" - ) - state_dict[f"transformer_layers.{i}.attn1.to_out.0.weight"] = state_dict.pop( - f"transformer_layers.{i}.attention.out.weight" - ) - - state_dict[f"transformer_layers.{i}.norm2.norm.weight"] = state_dict.pop( - f"transformer_layers.{i}.crossattn_layer_norm.weight" - ) - state_dict[f"transformer_layers.{i}.norm2.linear.weight"] = state_dict.pop( - f"transformer_layers.{i}.cross_attn_adaLN_modulation.mapper.weight" - ) - - state_dict[f"transformer_layers.{i}.attn2.to_q.weight"] = state_dict.pop( - f"transformer_layers.{i}.crossattention.query.weight" - ) - state_dict[f"transformer_layers.{i}.attn2.to_k.weight"] = state_dict.pop( - f"transformer_layers.{i}.crossattention.key.weight" - ) - state_dict[f"transformer_layers.{i}.attn2.to_v.weight"] = state_dict.pop( - f"transformer_layers.{i}.crossattention.value.weight" - ) - state_dict[f"transformer_layers.{i}.attn2.to_out.0.weight"] = state_dict.pop( - f"transformer_layers.{i}.crossattention.out.weight" - ) - - state_dict[f"transformer_layers.{i}.norm3.norm.weight"] = state_dict.pop( - f"transformer_layers.{i}.ffn.pre_mlp_layer_norm.weight" - ) - state_dict[f"transformer_layers.{i}.norm3.linear.weight"] = state_dict.pop( - f"transformer_layers.{i}.ffn.adaLN_modulation.mapper.weight" - ) - - wi_0_weight = state_dict.pop(f"transformer_layers.{i}.ffn.wi_0.weight") - wi_1_weight = state_dict.pop(f"transformer_layers.{i}.ffn.wi_1.weight") - proj_weight = torch.concat([wi_1_weight, wi_0_weight], dim=0) - state_dict[f"transformer_layers.{i}.ff.net.0.proj.weight"] = proj_weight - - state_dict[f"transformer_layers.{i}.ff.net.2.weight"] = state_dict.pop(f"transformer_layers.{i}.ffn.wo.weight") - - if force_down_up_sample: - state_dict["down_block.downsample.norm.weight"] = state_dict.pop("down_blocks.0.downsample.0.norm.weight") - state_dict["down_block.downsample.conv.weight"] = state_dict.pop("down_blocks.0.downsample.1.weight") - - state_dict["up_block.upsample.norm.weight"] = state_dict.pop("up_blocks.0.upsample.0.norm.weight") - state_dict["up_block.upsample.conv.weight"] = state_dict.pop("up_blocks.0.upsample.1.weight") - - state_dict["mlm_layer.layer_norm.weight"] = state_dict.pop("mlm_layer.layer_norm.norm.weight") - - for i in range(3): - state_dict[f"down_block.res_blocks.{i}.norm.weight"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.norm.norm.weight" - ) - state_dict[f"down_block.res_blocks.{i}.channelwise_linear_1.weight"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.channelwise.0.weight" - ) - state_dict[f"down_block.res_blocks.{i}.channelwise_norm.gamma"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.channelwise.2.gamma" - ) - state_dict[f"down_block.res_blocks.{i}.channelwise_norm.beta"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.channelwise.2.beta" - ) - state_dict[f"down_block.res_blocks.{i}.channelwise_linear_2.weight"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.channelwise.4.weight" - ) - state_dict[f"down_block.res_blocks.{i}.cond_embeds_mapper.weight"] = state_dict.pop( - f"down_blocks.0.res_blocks.{i}.adaLN_modulation.mapper.weight" - ) - - state_dict[f"down_block.attention_blocks.{i}.norm1.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.attn_layer_norm.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn1.to_q.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.attention.query.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn1.to_k.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.attention.key.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn1.to_v.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.attention.value.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn1.to_out.0.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.attention.out.weight" - ) - - state_dict[f"down_block.attention_blocks.{i}.norm2.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.crossattn_layer_norm.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn2.to_q.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.crossattention.query.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn2.to_k.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.crossattention.key.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn2.to_v.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.crossattention.value.weight" - ) - state_dict[f"down_block.attention_blocks.{i}.attn2.to_out.0.weight"] = state_dict.pop( - f"down_blocks.0.attention_blocks.{i}.crossattention.out.weight" - ) - - state_dict[f"up_block.res_blocks.{i}.norm.weight"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.norm.norm.weight" - ) - state_dict[f"up_block.res_blocks.{i}.channelwise_linear_1.weight"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.channelwise.0.weight" - ) - state_dict[f"up_block.res_blocks.{i}.channelwise_norm.gamma"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.channelwise.2.gamma" - ) - state_dict[f"up_block.res_blocks.{i}.channelwise_norm.beta"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.channelwise.2.beta" - ) - state_dict[f"up_block.res_blocks.{i}.channelwise_linear_2.weight"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.channelwise.4.weight" - ) - state_dict[f"up_block.res_blocks.{i}.cond_embeds_mapper.weight"] = state_dict.pop( - f"up_blocks.0.res_blocks.{i}.adaLN_modulation.mapper.weight" - ) - - state_dict[f"up_block.attention_blocks.{i}.norm1.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.attn_layer_norm.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn1.to_q.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.attention.query.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn1.to_k.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.attention.key.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn1.to_v.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.attention.value.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn1.to_out.0.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.attention.out.weight" - ) - - state_dict[f"up_block.attention_blocks.{i}.norm2.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.crossattn_layer_norm.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn2.to_q.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.crossattention.query.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn2.to_k.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.crossattention.key.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn2.to_v.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.crossattention.value.weight" - ) - state_dict[f"up_block.attention_blocks.{i}.attn2.to_out.0.weight"] = state_dict.pop( - f"up_blocks.0.attention_blocks.{i}.crossattention.out.weight" - ) - - for key in list(state_dict.keys()): - if key.startswith("up_blocks.0"): - key_ = "up_block." + ".".join(key.split(".")[2:]) - state_dict[key_] = state_dict.pop(key) - - if key.startswith("down_blocks.0"): - key_ = "down_block." + ".".join(key.split(".")[2:]) - state_dict[key_] = state_dict.pop(key) - - new_transformer.load_state_dict(state_dict) - - input_ids = torch.randint(0, 10, (1, 32, 32), device=old_transformer.device) - encoder_hidden_states = torch.randn((1, 77, 768), device=old_transformer.device) - cond_embeds = torch.randn((1, 768), device=old_transformer.device) - micro_conds = torch.tensor([[512, 512, 0, 0, 6]], dtype=torch.float32, device=old_transformer.device) - - old_out = old_transformer(input_ids.reshape(1, -1), encoder_hidden_states, cond_embeds, micro_conds) - old_out = old_out.reshape(1, 32, 32, 8192).permute(0, 3, 1, 2) - - new_out = new_transformer(input_ids, encoder_hidden_states, cond_embeds, micro_conds) - - # NOTE: these differences are solely due to using the geglu block that has a single linear layer of - # double output dimension instead of two different linear layers - max_diff = (old_out - new_out).abs().max() - total_diff = (old_out - new_out).abs().sum() - print(f"Transformer max_diff: {max_diff} total_diff: {total_diff}") - assert max_diff < 0.01 - assert total_diff < 1500 - - return new_transformer - - -def make_vqvae(old_vae): - new_vae = VQModel( - act_fn="silu", - block_out_channels=[128, 256, 256, 512, 768], - down_block_types=[ - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - ], - in_channels=3, - latent_channels=64, - layers_per_block=2, - norm_num_groups=32, - num_vq_embeddings=8192, - out_channels=3, - sample_size=32, - up_block_types=[ - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - ], - mid_block_add_attention=False, - lookup_from_codebook=True, - ) - new_vae.to(device) - - # fmt: off - - new_state_dict = {} - - old_state_dict = old_vae.state_dict() - - new_state_dict["encoder.conv_in.weight"] = old_state_dict.pop("encoder.conv_in.weight") - new_state_dict["encoder.conv_in.bias"] = old_state_dict.pop("encoder.conv_in.bias") - - convert_vae_block_state_dict(old_state_dict, "encoder.down.0", new_state_dict, "encoder.down_blocks.0") - convert_vae_block_state_dict(old_state_dict, "encoder.down.1", new_state_dict, "encoder.down_blocks.1") - convert_vae_block_state_dict(old_state_dict, "encoder.down.2", new_state_dict, "encoder.down_blocks.2") - convert_vae_block_state_dict(old_state_dict, "encoder.down.3", new_state_dict, "encoder.down_blocks.3") - convert_vae_block_state_dict(old_state_dict, "encoder.down.4", new_state_dict, "encoder.down_blocks.4") - - new_state_dict["encoder.mid_block.resnets.0.norm1.weight"] = old_state_dict.pop("encoder.mid.block_1.norm1.weight") - new_state_dict["encoder.mid_block.resnets.0.norm1.bias"] = old_state_dict.pop("encoder.mid.block_1.norm1.bias") - new_state_dict["encoder.mid_block.resnets.0.conv1.weight"] = old_state_dict.pop("encoder.mid.block_1.conv1.weight") - new_state_dict["encoder.mid_block.resnets.0.conv1.bias"] = old_state_dict.pop("encoder.mid.block_1.conv1.bias") - new_state_dict["encoder.mid_block.resnets.0.norm2.weight"] = old_state_dict.pop("encoder.mid.block_1.norm2.weight") - new_state_dict["encoder.mid_block.resnets.0.norm2.bias"] = old_state_dict.pop("encoder.mid.block_1.norm2.bias") - new_state_dict["encoder.mid_block.resnets.0.conv2.weight"] = old_state_dict.pop("encoder.mid.block_1.conv2.weight") - new_state_dict["encoder.mid_block.resnets.0.conv2.bias"] = old_state_dict.pop("encoder.mid.block_1.conv2.bias") - new_state_dict["encoder.mid_block.resnets.1.norm1.weight"] = old_state_dict.pop("encoder.mid.block_2.norm1.weight") - new_state_dict["encoder.mid_block.resnets.1.norm1.bias"] = old_state_dict.pop("encoder.mid.block_2.norm1.bias") - new_state_dict["encoder.mid_block.resnets.1.conv1.weight"] = old_state_dict.pop("encoder.mid.block_2.conv1.weight") - new_state_dict["encoder.mid_block.resnets.1.conv1.bias"] = old_state_dict.pop("encoder.mid.block_2.conv1.bias") - new_state_dict["encoder.mid_block.resnets.1.norm2.weight"] = old_state_dict.pop("encoder.mid.block_2.norm2.weight") - new_state_dict["encoder.mid_block.resnets.1.norm2.bias"] = old_state_dict.pop("encoder.mid.block_2.norm2.bias") - new_state_dict["encoder.mid_block.resnets.1.conv2.weight"] = old_state_dict.pop("encoder.mid.block_2.conv2.weight") - new_state_dict["encoder.mid_block.resnets.1.conv2.bias"] = old_state_dict.pop("encoder.mid.block_2.conv2.bias") - new_state_dict["encoder.conv_norm_out.weight"] = old_state_dict.pop("encoder.norm_out.weight") - new_state_dict["encoder.conv_norm_out.bias"] = old_state_dict.pop("encoder.norm_out.bias") - new_state_dict["encoder.conv_out.weight"] = old_state_dict.pop("encoder.conv_out.weight") - new_state_dict["encoder.conv_out.bias"] = old_state_dict.pop("encoder.conv_out.bias") - new_state_dict["quant_conv.weight"] = old_state_dict.pop("quant_conv.weight") - new_state_dict["quant_conv.bias"] = old_state_dict.pop("quant_conv.bias") - new_state_dict["quantize.embedding.weight"] = old_state_dict.pop("quantize.embedding.weight") - new_state_dict["post_quant_conv.weight"] = old_state_dict.pop("post_quant_conv.weight") - new_state_dict["post_quant_conv.bias"] = old_state_dict.pop("post_quant_conv.bias") - new_state_dict["decoder.conv_in.weight"] = old_state_dict.pop("decoder.conv_in.weight") - new_state_dict["decoder.conv_in.bias"] = old_state_dict.pop("decoder.conv_in.bias") - new_state_dict["decoder.mid_block.resnets.0.norm1.weight"] = old_state_dict.pop("decoder.mid.block_1.norm1.weight") - new_state_dict["decoder.mid_block.resnets.0.norm1.bias"] = old_state_dict.pop("decoder.mid.block_1.norm1.bias") - new_state_dict["decoder.mid_block.resnets.0.conv1.weight"] = old_state_dict.pop("decoder.mid.block_1.conv1.weight") - new_state_dict["decoder.mid_block.resnets.0.conv1.bias"] = old_state_dict.pop("decoder.mid.block_1.conv1.bias") - new_state_dict["decoder.mid_block.resnets.0.norm2.weight"] = old_state_dict.pop("decoder.mid.block_1.norm2.weight") - new_state_dict["decoder.mid_block.resnets.0.norm2.bias"] = old_state_dict.pop("decoder.mid.block_1.norm2.bias") - new_state_dict["decoder.mid_block.resnets.0.conv2.weight"] = old_state_dict.pop("decoder.mid.block_1.conv2.weight") - new_state_dict["decoder.mid_block.resnets.0.conv2.bias"] = old_state_dict.pop("decoder.mid.block_1.conv2.bias") - new_state_dict["decoder.mid_block.resnets.1.norm1.weight"] = old_state_dict.pop("decoder.mid.block_2.norm1.weight") - new_state_dict["decoder.mid_block.resnets.1.norm1.bias"] = old_state_dict.pop("decoder.mid.block_2.norm1.bias") - new_state_dict["decoder.mid_block.resnets.1.conv1.weight"] = old_state_dict.pop("decoder.mid.block_2.conv1.weight") - new_state_dict["decoder.mid_block.resnets.1.conv1.bias"] = old_state_dict.pop("decoder.mid.block_2.conv1.bias") - new_state_dict["decoder.mid_block.resnets.1.norm2.weight"] = old_state_dict.pop("decoder.mid.block_2.norm2.weight") - new_state_dict["decoder.mid_block.resnets.1.norm2.bias"] = old_state_dict.pop("decoder.mid.block_2.norm2.bias") - new_state_dict["decoder.mid_block.resnets.1.conv2.weight"] = old_state_dict.pop("decoder.mid.block_2.conv2.weight") - new_state_dict["decoder.mid_block.resnets.1.conv2.bias"] = old_state_dict.pop("decoder.mid.block_2.conv2.bias") - - convert_vae_block_state_dict(old_state_dict, "decoder.up.0", new_state_dict, "decoder.up_blocks.4") - convert_vae_block_state_dict(old_state_dict, "decoder.up.1", new_state_dict, "decoder.up_blocks.3") - convert_vae_block_state_dict(old_state_dict, "decoder.up.2", new_state_dict, "decoder.up_blocks.2") - convert_vae_block_state_dict(old_state_dict, "decoder.up.3", new_state_dict, "decoder.up_blocks.1") - convert_vae_block_state_dict(old_state_dict, "decoder.up.4", new_state_dict, "decoder.up_blocks.0") - - new_state_dict["decoder.conv_norm_out.weight"] = old_state_dict.pop("decoder.norm_out.weight") - new_state_dict["decoder.conv_norm_out.bias"] = old_state_dict.pop("decoder.norm_out.bias") - new_state_dict["decoder.conv_out.weight"] = old_state_dict.pop("decoder.conv_out.weight") - new_state_dict["decoder.conv_out.bias"] = old_state_dict.pop("decoder.conv_out.bias") - - # fmt: on - - assert len(old_state_dict.keys()) == 0 - - new_vae.load_state_dict(new_state_dict) - - input = torch.randn((1, 3, 512, 512), device=device) - input = input.clamp(-1, 1) - - old_encoder_output = old_vae.quant_conv(old_vae.encoder(input)) - new_encoder_output = new_vae.quant_conv(new_vae.encoder(input)) - assert (old_encoder_output == new_encoder_output).all() - - old_decoder_output = old_vae.decoder(old_vae.post_quant_conv(old_encoder_output)) - new_decoder_output = new_vae.decoder(new_vae.post_quant_conv(new_encoder_output)) - - # assert (old_decoder_output == new_decoder_output).all() - print("kipping vae decoder equivalence check") - print(f"vae decoder diff {(old_decoder_output - new_decoder_output).float().abs().sum()}") - - old_output = old_vae(input)[0] - new_output = new_vae(input)[0] - - # assert (old_output == new_output).all() - print("skipping full vae equivalence check") - print(f"vae full diff {(old_output - new_output).float().abs().sum()}") - - return new_vae - - -def convert_vae_block_state_dict(old_state_dict, prefix_from, new_state_dict, prefix_to): - # fmt: off - - new_state_dict[f"{prefix_to}.resnets.0.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.norm1.weight") - new_state_dict[f"{prefix_to}.resnets.0.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.norm1.bias") - new_state_dict[f"{prefix_to}.resnets.0.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.conv1.weight") - new_state_dict[f"{prefix_to}.resnets.0.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.conv1.bias") - new_state_dict[f"{prefix_to}.resnets.0.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.norm2.weight") - new_state_dict[f"{prefix_to}.resnets.0.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.norm2.bias") - new_state_dict[f"{prefix_to}.resnets.0.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.conv2.weight") - new_state_dict[f"{prefix_to}.resnets.0.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.conv2.bias") - - if f"{prefix_from}.block.0.nin_shortcut.weight" in old_state_dict: - new_state_dict[f"{prefix_to}.resnets.0.conv_shortcut.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.nin_shortcut.weight") - new_state_dict[f"{prefix_to}.resnets.0.conv_shortcut.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.nin_shortcut.bias") - - new_state_dict[f"{prefix_to}.resnets.1.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.norm1.weight") - new_state_dict[f"{prefix_to}.resnets.1.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.norm1.bias") - new_state_dict[f"{prefix_to}.resnets.1.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.conv1.weight") - new_state_dict[f"{prefix_to}.resnets.1.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.conv1.bias") - new_state_dict[f"{prefix_to}.resnets.1.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.norm2.weight") - new_state_dict[f"{prefix_to}.resnets.1.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.norm2.bias") - new_state_dict[f"{prefix_to}.resnets.1.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.conv2.weight") - new_state_dict[f"{prefix_to}.resnets.1.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.conv2.bias") - - if f"{prefix_from}.downsample.conv.weight" in old_state_dict: - new_state_dict[f"{prefix_to}.downsamplers.0.conv.weight"] = old_state_dict.pop(f"{prefix_from}.downsample.conv.weight") - new_state_dict[f"{prefix_to}.downsamplers.0.conv.bias"] = old_state_dict.pop(f"{prefix_from}.downsample.conv.bias") - - if f"{prefix_from}.upsample.conv.weight" in old_state_dict: - new_state_dict[f"{prefix_to}.upsamplers.0.conv.weight"] = old_state_dict.pop(f"{prefix_from}.upsample.conv.weight") - new_state_dict[f"{prefix_to}.upsamplers.0.conv.bias"] = old_state_dict.pop(f"{prefix_from}.upsample.conv.bias") - - if f"{prefix_from}.block.2.norm1.weight" in old_state_dict: - new_state_dict[f"{prefix_to}.resnets.2.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.norm1.weight") - new_state_dict[f"{prefix_to}.resnets.2.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.norm1.bias") - new_state_dict[f"{prefix_to}.resnets.2.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.conv1.weight") - new_state_dict[f"{prefix_to}.resnets.2.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.conv1.bias") - new_state_dict[f"{prefix_to}.resnets.2.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.norm2.weight") - new_state_dict[f"{prefix_to}.resnets.2.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.norm2.bias") - new_state_dict[f"{prefix_to}.resnets.2.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.conv2.weight") - new_state_dict[f"{prefix_to}.resnets.2.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.conv2.bias") - - # fmt: on - - -if __name__ == "__main__": - main() diff --git a/scripts/convert_anima_to_diffusers.py b/scripts/convert_anima_to_diffusers.py deleted file mode 100644 index bf8e006ad266..000000000000 --- a/scripts/convert_anima_to_diffusers.py +++ /dev/null @@ -1,316 +0,0 @@ -""" -Convert Anima checkpoints to Diffusers format. - -Example: -```bash -python scripts/convert_anima_to_diffusers.py \ - --transformer_ckpt_path anima_model/anima-preview3-base.safetensors \ - --text_encoder_ckpt_path anima_model/qwen_3_06b_base.safetensors \ - --vae_ckpt_path anima_model/qwen_image_vae.safetensors \ - --qwen_tokenizer_path path/to/qwen25_tokenizer \ - --t5_tokenizer_path path/to/t5_tokenizer \ - --output_path anima_model/anima-preview3-diffusers \ - --save_pipeline -``` -""" - -import argparse -import pathlib -import sys -from typing import Any - -import torch -from accelerate import init_empty_weights -from convert_cosmos_to_diffusers import convert_transformer -from safetensors.torch import load_file -from transformers import AutoTokenizer, Qwen3Config, Qwen3Model, T5TokenizerFast - -from diffusers import ( - AnimaAutoBlocks, - AnimaTextConditioner, - AutoencoderKLQwenImage, - FlowMatchEulerDiscreteScheduler, -) - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -def rename_residual_key(key: str) -> str: - replacements = { - ".residual.0.": ".norm1.", - ".residual.2.": ".conv1.", - ".residual.3.": ".norm2.", - ".residual.6.": ".conv2.", - ".shortcut.": ".conv_shortcut.", - } - for old, new in replacements.items(): - key = key.replace(old, new) - return key - - -def rename_mid_key(key: str) -> str: - replacements = { - ".middle.0.": ".mid_block.resnets.0.", - ".middle.1.": ".mid_block.attentions.0.", - ".middle.2.": ".mid_block.resnets.1.", - } - for old, new in replacements.items(): - key = key.replace(old, new) - return rename_residual_key(key) - - -def rename_decoder_upsample_key(key: str) -> str: - prefix = "decoder.upsamples." - suffix = key.removeprefix(prefix) - index_str, rest = suffix.split(".", 1) - index = int(index_str) - - if index in (3, 7, 11): - block_index = (index - 3) // 4 - new_key = f"decoder.up_blocks.{block_index}.upsamplers.0.{rest}" - else: - block_index = index // 4 - resnet_index = index % 4 - new_key = f"decoder.up_blocks.{block_index}.resnets.{resnet_index}.{rest}" - - return rename_residual_key(new_key) - - -def convert_qwen_image_vae_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - converted_state_dict = {} - for key, value in state_dict.items(): - if key.startswith("conv1."): - new_key = key.replace("conv1.", "quant_conv.", 1) - elif key.startswith("conv2."): - new_key = key.replace("conv2.", "post_quant_conv.", 1) - elif key.startswith("encoder.conv1."): - new_key = key.replace("encoder.conv1.", "encoder.conv_in.", 1) - elif key.startswith("decoder.conv1."): - new_key = key.replace("decoder.conv1.", "decoder.conv_in.", 1) - elif key.startswith("encoder.downsamples."): - new_key = rename_residual_key(key.replace("encoder.downsamples.", "encoder.down_blocks.", 1)) - elif key.startswith("decoder.upsamples."): - new_key = rename_decoder_upsample_key(key) - elif key.startswith("encoder.middle.") or key.startswith("decoder.middle."): - new_key = rename_mid_key(key) - elif key.startswith("encoder.head.0."): - new_key = key.replace("encoder.head.0.", "encoder.norm_out.", 1) - elif key.startswith("encoder.head.2."): - new_key = key.replace("encoder.head.2.", "encoder.conv_out.", 1) - elif key.startswith("decoder.head.0."): - new_key = key.replace("decoder.head.0.", "decoder.norm_out.", 1) - elif key.startswith("decoder.head.2."): - new_key = key.replace("decoder.head.2.", "decoder.conv_out.", 1) - else: - new_key = rename_residual_key(key) - - if new_key in converted_state_dict: - raise ValueError(f"Duplicate converted VAE key: {new_key}") - converted_state_dict[new_key] = value - - return converted_state_dict - - -def convert_qwen_image_vae(state_dict: dict[str, torch.Tensor]) -> AutoencoderKLQwenImage: - converted_state_dict = convert_qwen_image_vae_state_dict(state_dict) - with init_empty_weights(): - vae = AutoencoderKLQwenImage() - - expected_keys = set(vae.state_dict().keys()) - converted_keys = set(converted_state_dict.keys()) - missing_keys = expected_keys - converted_keys - unexpected_keys = converted_keys - expected_keys - if missing_keys or unexpected_keys: - if missing_keys: - print(f"ERROR: missing VAE keys ({len(missing_keys)}):", file=sys.stderr) - for key in sorted(missing_keys): - print(key, file=sys.stderr) - if unexpected_keys: - print(f"ERROR: unexpected VAE keys ({len(unexpected_keys)}):", file=sys.stderr) - for key in sorted(unexpected_keys): - print(key, file=sys.stderr) - sys.exit(1) - - vae.load_state_dict(converted_state_dict, strict=True, assign=True) - return vae - - -def infer_text_conditioner_config(state_dict: dict[str, torch.Tensor]) -> dict[str, Any]: - model_dim = state_dict["blocks.0.self_attn.q_proj.weight"].shape[0] - source_dim = state_dict["blocks.0.cross_attn.k_proj.weight"].shape[1] - target_vocab_size, target_dim = state_dict["embed.weight"].shape - attention_head_dim = state_dict["blocks.0.self_attn.q_norm.weight"].shape[0] - num_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("blocks.")) - - return { - "source_dim": source_dim, - "target_dim": target_dim, - "model_dim": model_dim, - "num_layers": num_layers, - "num_attention_heads": model_dim // attention_head_dim, - "target_vocab_size": target_vocab_size, - } - - -def convert_text_conditioner(state_dict: dict[str, torch.Tensor]) -> AnimaTextConditioner: - config = infer_text_conditioner_config(state_dict) - with init_empty_weights(): - text_conditioner = AnimaTextConditioner(**config) - - expected_keys = set(text_conditioner.state_dict().keys()) - converted_keys = set(state_dict.keys()) - missing_keys = expected_keys - converted_keys - unexpected_keys = converted_keys - expected_keys - if missing_keys or unexpected_keys: - if missing_keys: - print(f"ERROR: missing text conditioner keys ({len(missing_keys)}):", file=sys.stderr) - for key in sorted(missing_keys): - print(key, file=sys.stderr) - if unexpected_keys: - print(f"ERROR: unexpected text conditioner keys ({len(unexpected_keys)}):", file=sys.stderr) - for key in sorted(unexpected_keys): - print(key, file=sys.stderr) - sys.exit(1) - - text_conditioner.load_state_dict(state_dict, strict=True, assign=True) - return text_conditioner - - -def infer_qwen3_config(state_dict: dict[str, torch.Tensor]) -> Qwen3Config: - vocab_size, hidden_size = state_dict["embed_tokens.weight"].shape - intermediate_size = state_dict["layers.0.mlp.gate_proj.weight"].shape[0] - num_hidden_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("layers.")) - head_dim = state_dict["layers.0.self_attn.q_norm.weight"].shape[0] - num_attention_heads = state_dict["layers.0.self_attn.q_proj.weight"].shape[0] // head_dim - num_key_value_heads = state_dict["layers.0.self_attn.k_proj.weight"].shape[0] // head_dim - - return Qwen3Config( - vocab_size=vocab_size, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - num_hidden_layers=num_hidden_layers, - num_attention_heads=num_attention_heads, - num_key_value_heads=num_key_value_heads, - max_position_embeddings=32768, - rms_norm_eps=1e-6, - rope_theta=1000000.0, - head_dim=head_dim, - attention_bias=False, - tie_word_embeddings=False, - ) - - -def convert_text_encoder(state_dict: dict[str, torch.Tensor]) -> Qwen3Model: - state_dict = {key.removeprefix("model."): value for key, value in state_dict.items()} - config = infer_qwen3_config(state_dict) - with init_empty_weights(): - text_encoder = Qwen3Model(config) - - expected_keys = set(text_encoder.state_dict().keys()) - converted_keys = set(state_dict.keys()) - missing_keys = expected_keys - converted_keys - unexpected_keys = converted_keys - expected_keys - if missing_keys or unexpected_keys: - if missing_keys: - print(f"ERROR: missing Qwen3 keys ({len(missing_keys)}):", file=sys.stderr) - for key in sorted(missing_keys): - print(key, file=sys.stderr) - if unexpected_keys: - print(f"ERROR: unexpected Qwen3 keys ({len(unexpected_keys)}):", file=sys.stderr) - for key in sorted(unexpected_keys): - print(key, file=sys.stderr) - sys.exit(1) - - text_encoder.load_state_dict(state_dict, strict=True, assign=True) - return text_encoder - - -def split_anima_transformer_checkpoint( - state_dict: dict[str, torch.Tensor], -) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - transformer_state_dict = {} - text_conditioner_state_dict = {} - adapter_prefix = "net.llm_adapter." - - for key, value in state_dict.items(): - if key.startswith(adapter_prefix): - text_conditioner_state_dict[key.removeprefix(adapter_prefix)] = value - else: - transformer_state_dict[key] = value - - return transformer_state_dict, text_conditioner_state_dict - - -def save_pipeline(args, transformer, text_conditioner, text_encoder, vae): - tokenizer = AutoTokenizer.from_pretrained(args.qwen_tokenizer_path) - t5_tokenizer = T5TokenizerFast.from_pretrained(args.t5_tokenizer_path) - scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) - - pipe = AnimaAutoBlocks().init_pipeline() - pipe.update_components( - text_encoder=text_encoder, - tokenizer=tokenizer, - t5_tokenizer=t5_tokenizer, - text_conditioner=text_conditioner, - transformer=transformer, - vae=vae, - scheduler=scheduler, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size=args.max_shard_size) - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--transformer_ckpt_path", type=str, required=True, help="Path to Anima DiT safetensors") - parser.add_argument("--text_encoder_ckpt_path", type=str, required=True, help="Path to Qwen3 text encoder") - parser.add_argument("--vae_ckpt_path", type=str, required=True, help="Path to Qwen-Image VAE safetensors") - parser.add_argument("--qwen_tokenizer_path", type=str, default=None) - parser.add_argument("--t5_tokenizer_path", type=str, default=None) - parser.add_argument("--output_path", type=str, required=True) - parser.add_argument("--save_pipeline", action="store_true") - parser.add_argument("--dtype", default="bf16", choices=list(DTYPE_MAPPING.keys())) - parser.add_argument("--max_shard_size", default="5GB") - return parser.parse_args() - - -if __name__ == "__main__": - args = get_args() - output_path = pathlib.Path(args.output_path) - dtype = DTYPE_MAPPING[args.dtype] - - raw_transformer_state_dict = load_file(args.transformer_ckpt_path, device="cpu") - transformer_state_dict, text_conditioner_state_dict = split_anima_transformer_checkpoint( - raw_transformer_state_dict - ) - transformer = convert_transformer( - "Cosmos-2.0-Diffusion-2B-Text2Image", state_dict=transformer_state_dict, weights_only=True - ).to(dtype=dtype) - text_conditioner = convert_text_conditioner(text_conditioner_state_dict).to(dtype=dtype) - - text_encoder_state_dict = load_file(args.text_encoder_ckpt_path, device="cpu") - text_encoder = convert_text_encoder(text_encoder_state_dict).to(dtype=dtype) - - vae_state_dict = load_file(args.vae_ckpt_path, device="cpu") - vae = convert_qwen_image_vae(vae_state_dict).to(dtype=dtype) - - if args.save_pipeline: - if args.qwen_tokenizer_path is None or args.t5_tokenizer_path is None: - raise ValueError("`--qwen_tokenizer_path` and `--t5_tokenizer_path` are required with `--save_pipeline`.") - save_pipeline(args, transformer, text_conditioner, text_encoder, vae) - else: - output_path.mkdir(parents=True, exist_ok=True) - transformer.save_pretrained( - output_path / "transformer", safe_serialization=True, max_shard_size=args.max_shard_size - ) - text_conditioner.save_pretrained( - output_path / "text_conditioner", safe_serialization=True, max_shard_size=args.max_shard_size - ) - text_encoder.save_pretrained( - output_path / "text_encoder", safe_serialization=True, max_shard_size=args.max_shard_size - ) - vae.save_pretrained(output_path / "vae", safe_serialization=True, max_shard_size=args.max_shard_size) diff --git a/scripts/convert_animatediff_motion_lora_to_diffusers.py b/scripts/convert_animatediff_motion_lora_to_diffusers.py deleted file mode 100644 index 21567ffa9e7a..000000000000 --- a/scripts/convert_animatediff_motion_lora_to_diffusers.py +++ /dev/null @@ -1,69 +0,0 @@ -import argparse -import os - -import torch -from huggingface_hub import create_repo, upload_folder -from safetensors.torch import load_file, save_file - - -def convert_motion_module(original_state_dict): - converted_state_dict = {} - for k, v in original_state_dict.items(): - if "pos_encoder" in k: - continue - - else: - converted_state_dict[ - k.replace(".norms.0", ".norm1") - .replace(".norms.1", ".norm2") - .replace(".ff_norm", ".norm3") - .replace(".attention_blocks.0", ".attn1") - .replace(".attention_blocks.1", ".attn2") - .replace(".temporal_transformer", "") - ] = v - - return converted_state_dict - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--ckpt_path", type=str, required=True, help="Path to checkpoint") - parser.add_argument("--output_path", type=str, required=True, help="Path to output directory") - parser.add_argument( - "--push_to_hub", - action="store_true", - default=False, - help="Whether to push the converted model to the HF or not", - ) - - return parser.parse_args() - - -if __name__ == "__main__": - args = get_args() - - if args.ckpt_path.endswith(".safetensors"): - state_dict = load_file(args.ckpt_path) - else: - state_dict = torch.load(args.ckpt_path, map_location="cpu") - - if "state_dict" in state_dict.keys(): - state_dict = state_dict["state_dict"] - - conv_state_dict = convert_motion_module(state_dict) - - # convert to new format - output_dict = {} - for module_name, params in conv_state_dict.items(): - if type(params) is not torch.Tensor: - continue - output_dict.update({f"unet.{module_name}": params}) - - os.makedirs(args.output_path, exist_ok=True) - - filepath = os.path.join(args.output_path, "diffusion_pytorch_model.safetensors") - save_file(output_dict, filepath) - - if args.push_to_hub: - repo_id = create_repo(args.output_path, exist_ok=True).repo_id - upload_folder(repo_id=repo_id, folder_path=args.output_path, repo_type="model") diff --git a/scripts/convert_animatediff_motion_module_to_diffusers.py b/scripts/convert_animatediff_motion_module_to_diffusers.py deleted file mode 100644 index e188a6a533e8..000000000000 --- a/scripts/convert_animatediff_motion_module_to_diffusers.py +++ /dev/null @@ -1,62 +0,0 @@ -import argparse - -import torch -from safetensors.torch import load_file - -from diffusers import MotionAdapter - - -def convert_motion_module(original_state_dict): - converted_state_dict = {} - for k, v in original_state_dict.items(): - if "pos_encoder" in k: - continue - - else: - converted_state_dict[ - k.replace(".norms.0", ".norm1") - .replace(".norms.1", ".norm2") - .replace(".ff_norm", ".norm3") - .replace(".attention_blocks.0", ".attn1") - .replace(".attention_blocks.1", ".attn2") - .replace(".temporal_transformer", "") - ] = v - - return converted_state_dict - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--ckpt_path", type=str, required=True) - parser.add_argument("--output_path", type=str, required=True) - parser.add_argument("--use_motion_mid_block", action="store_true") - parser.add_argument("--motion_max_seq_length", type=int, default=32) - parser.add_argument("--block_out_channels", nargs="+", default=[320, 640, 1280, 1280], type=int) - parser.add_argument("--save_fp16", action="store_true") - - return parser.parse_args() - - -if __name__ == "__main__": - args = get_args() - - if args.ckpt_path.endswith(".safetensors"): - state_dict = load_file(args.ckpt_path) - else: - state_dict = torch.load(args.ckpt_path, map_location="cpu") - - if "state_dict" in state_dict.keys(): - state_dict = state_dict["state_dict"] - - conv_state_dict = convert_motion_module(state_dict) - adapter = MotionAdapter( - block_out_channels=args.block_out_channels, - use_motion_mid_block=args.use_motion_mid_block, - motion_max_seq_length=args.motion_max_seq_length, - ) - # skip loading position embeddings - adapter.load_state_dict(conv_state_dict, strict=False) - adapter.save_pretrained(args.output_path) - - if args.save_fp16: - adapter.to(dtype=torch.float16).save_pretrained(args.output_path, variant="fp16") diff --git a/scripts/convert_animatediff_sparsectrl_to_diffusers.py b/scripts/convert_animatediff_sparsectrl_to_diffusers.py deleted file mode 100644 index b96b42d60f44..000000000000 --- a/scripts/convert_animatediff_sparsectrl_to_diffusers.py +++ /dev/null @@ -1,83 +0,0 @@ -import argparse -from typing import Dict - -import torch -import torch.nn as nn - -from diffusers import SparseControlNetModel - - -KEYS_RENAME_MAPPING = { - ".attention_blocks.0": ".attn1", - ".attention_blocks.1": ".attn2", - ".attn1.pos_encoder": ".pos_embed", - ".ff_norm": ".norm3", - ".norms.0": ".norm1", - ".norms.1": ".norm2", - ".temporal_transformer": "", -} - - -def convert(original_state_dict: Dict[str, nn.Module]) -> dict[str, nn.Module]: - converted_state_dict = {} - - for key in list(original_state_dict.keys()): - renamed_key = key - for new_name, old_name in KEYS_RENAME_MAPPING.items(): - renamed_key = renamed_key.replace(new_name, old_name) - converted_state_dict[renamed_key] = original_state_dict.pop(key) - - return converted_state_dict - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--ckpt_path", type=str, required=True, help="Path to checkpoint") - parser.add_argument("--output_path", type=str, required=True, help="Path to output directory") - parser.add_argument( - "--max_motion_seq_length", - type=int, - default=32, - help="Max motion sequence length supported by the motion adapter", - ) - parser.add_argument( - "--conditioning_channels", type=int, default=4, help="Number of channels in conditioning input to controlnet" - ) - parser.add_argument( - "--use_simplified_condition_embedding", - action="store_true", - default=False, - help="Whether or not to use simplified condition embedding. When `conditioning_channels==4` i.e. latent inputs, set this to `True`. When `conditioning_channels==3` i.e. image inputs, set this to `False`", - ) - parser.add_argument( - "--save_fp16", - action="store_true", - default=False, - help="Whether or not to save model in fp16 precision along with fp32", - ) - parser.add_argument( - "--push_to_hub", action="store_true", default=False, help="Whether or not to push saved model to the HF hub" - ) - return parser.parse_args() - - -if __name__ == "__main__": - args = get_args() - - state_dict = torch.load(args.ckpt_path, map_location="cpu") - if "state_dict" in state_dict.keys(): - state_dict: dict = state_dict["state_dict"] - - controlnet = SparseControlNetModel( - conditioning_channels=args.conditioning_channels, - motion_max_seq_length=args.max_motion_seq_length, - use_simplified_condition_embedding=args.use_simplified_condition_embedding, - ) - - state_dict = convert(state_dict) - controlnet.load_state_dict(state_dict, strict=True) - - controlnet.save_pretrained(args.output_path, push_to_hub=args.push_to_hub) - if args.save_fp16: - controlnet = controlnet.to(dtype=torch.float16) - controlnet.save_pretrained(args.output_path, variant="fp16", push_to_hub=args.push_to_hub) diff --git a/scripts/convert_asymmetric_vqgan_to_diffusers.py b/scripts/convert_asymmetric_vqgan_to_diffusers.py deleted file mode 100644 index a9f83a46c2d0..000000000000 --- a/scripts/convert_asymmetric_vqgan_to_diffusers.py +++ /dev/null @@ -1,184 +0,0 @@ -import argparse -import time -from pathlib import Path -from typing import Any, Dict, Literal - -import torch - -from diffusers import AsymmetricAutoencoderKL - - -ASYMMETRIC_AUTOENCODER_KL_x_1_5_CONFIG = { - "in_channels": 3, - "out_channels": 3, - "down_block_types": [ - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - ], - "down_block_out_channels": [128, 256, 512, 512], - "layers_per_down_block": 2, - "up_block_types": [ - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - ], - "up_block_out_channels": [192, 384, 768, 768], - "layers_per_up_block": 3, - "act_fn": "silu", - "latent_channels": 4, - "norm_num_groups": 32, - "sample_size": 256, - "scaling_factor": 0.18215, -} - -ASYMMETRIC_AUTOENCODER_KL_x_2_CONFIG = { - "in_channels": 3, - "out_channels": 3, - "down_block_types": [ - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - "DownEncoderBlock2D", - ], - "down_block_out_channels": [128, 256, 512, 512], - "layers_per_down_block": 2, - "up_block_types": [ - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - "UpDecoderBlock2D", - ], - "up_block_out_channels": [256, 512, 1024, 1024], - "layers_per_up_block": 5, - "act_fn": "silu", - "latent_channels": 4, - "norm_num_groups": 32, - "sample_size": 256, - "scaling_factor": 0.18215, -} - - -def convert_asymmetric_autoencoder_kl_state_dict(original_state_dict: Dict[str, Any]) -> dict[str, Any]: - converted_state_dict = {} - for k, v in original_state_dict.items(): - if k.startswith("encoder."): - converted_state_dict[ - k.replace("encoder.down.", "encoder.down_blocks.") - .replace("encoder.mid.", "encoder.mid_block.") - .replace("encoder.norm_out.", "encoder.conv_norm_out.") - .replace(".downsample.", ".downsamplers.0.") - .replace(".nin_shortcut.", ".conv_shortcut.") - .replace(".block.", ".resnets.") - .replace(".block_1.", ".resnets.0.") - .replace(".block_2.", ".resnets.1.") - .replace(".attn_1.k.", ".attentions.0.to_k.") - .replace(".attn_1.q.", ".attentions.0.to_q.") - .replace(".attn_1.v.", ".attentions.0.to_v.") - .replace(".attn_1.proj_out.", ".attentions.0.to_out.0.") - .replace(".attn_1.norm.", ".attentions.0.group_norm.") - ] = v - elif k.startswith("decoder.") and "up_layers" not in k: - converted_state_dict[ - k.replace("decoder.encoder.", "decoder.condition_encoder.") - .replace(".norm_out.", ".conv_norm_out.") - .replace(".up.0.", ".up_blocks.3.") - .replace(".up.1.", ".up_blocks.2.") - .replace(".up.2.", ".up_blocks.1.") - .replace(".up.3.", ".up_blocks.0.") - .replace(".block.", ".resnets.") - .replace("mid", "mid_block") - .replace(".0.upsample.", ".0.upsamplers.0.") - .replace(".1.upsample.", ".1.upsamplers.0.") - .replace(".2.upsample.", ".2.upsamplers.0.") - .replace(".nin_shortcut.", ".conv_shortcut.") - .replace(".block_1.", ".resnets.0.") - .replace(".block_2.", ".resnets.1.") - .replace(".attn_1.k.", ".attentions.0.to_k.") - .replace(".attn_1.q.", ".attentions.0.to_q.") - .replace(".attn_1.v.", ".attentions.0.to_v.") - .replace(".attn_1.proj_out.", ".attentions.0.to_out.0.") - .replace(".attn_1.norm.", ".attentions.0.group_norm.") - ] = v - elif k.startswith("quant_conv."): - converted_state_dict[k] = v - elif k.startswith("post_quant_conv."): - converted_state_dict[k] = v - else: - print(f" skipping key `{k}`") - # fix weights shape - for k, v in converted_state_dict.items(): - if ( - (k.startswith("encoder.mid_block.attentions.0") or k.startswith("decoder.mid_block.attentions.0")) - and k.endswith("weight") - and ("to_q" in k or "to_k" in k or "to_v" in k or "to_out" in k) - ): - converted_state_dict[k] = converted_state_dict[k][:, :, 0, 0] - - return converted_state_dict - - -def get_asymmetric_autoencoder_kl_from_original_checkpoint( - scale: Literal["1.5", "2"], original_checkpoint_path: str, map_location: torch.device -) -> AsymmetricAutoencoderKL: - print("Loading original state_dict") - original_state_dict = torch.load(original_checkpoint_path, map_location=map_location) - original_state_dict = original_state_dict["state_dict"] - print("Converting state_dict") - converted_state_dict = convert_asymmetric_autoencoder_kl_state_dict(original_state_dict) - kwargs = ASYMMETRIC_AUTOENCODER_KL_x_1_5_CONFIG if scale == "1.5" else ASYMMETRIC_AUTOENCODER_KL_x_2_CONFIG - print("Initializing AsymmetricAutoencoderKL model") - asymmetric_autoencoder_kl = AsymmetricAutoencoderKL(**kwargs) - print("Loading weight from converted state_dict") - asymmetric_autoencoder_kl.load_state_dict(converted_state_dict) - asymmetric_autoencoder_kl.eval() - print("AsymmetricAutoencoderKL successfully initialized") - return asymmetric_autoencoder_kl - - -if __name__ == "__main__": - start = time.time() - parser = argparse.ArgumentParser() - parser.add_argument( - "--scale", - default=None, - type=str, - required=True, - help="Asymmetric VQGAN scale: `1.5` or `2`", - ) - parser.add_argument( - "--original_checkpoint_path", - default=None, - type=str, - required=True, - help="Path to the original Asymmetric VQGAN checkpoint", - ) - parser.add_argument( - "--output_path", - default=None, - type=str, - required=True, - help="Path to save pretrained AsymmetricAutoencoderKL model", - ) - parser.add_argument( - "--map_location", - default="cpu", - type=str, - required=False, - help="The device passed to `map_location` when loading the checkpoint", - ) - args = parser.parse_args() - - assert args.scale in ["1.5", "2"], f"{args.scale} should be `1.5` of `2`" - assert Path(args.original_checkpoint_path).is_file() - - asymmetric_autoencoder_kl = get_asymmetric_autoencoder_kl_from_original_checkpoint( - scale=args.scale, - original_checkpoint_path=args.original_checkpoint_path, - map_location=torch.device(args.map_location), - ) - print("Saving pretrained AsymmetricAutoencoderKL") - asymmetric_autoencoder_kl.save_pretrained(args.output_path) - print(f"Done in {time.time() - start:.2f} seconds") diff --git a/scripts/convert_aura_flow_to_diffusers.py b/scripts/convert_aura_flow_to_diffusers.py deleted file mode 100644 index 74c34f4851ff..000000000000 --- a/scripts/convert_aura_flow_to_diffusers.py +++ /dev/null @@ -1,131 +0,0 @@ -import argparse - -import torch -from huggingface_hub import hf_hub_download - -from diffusers.models.transformers.auraflow_transformer_2d import AuraFlowTransformer2DModel - - -def load_original_state_dict(args): - model_pt = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename="aura_diffusion_pytorch_model.bin") - state_dict = torch.load(model_pt, map_location="cpu") - return state_dict - - -def calculate_layers(state_dict_keys, key_prefix): - dit_layers = set() - for k in state_dict_keys: - if key_prefix in k: - dit_layers.add(int(k.split(".")[2])) - print(f"{key_prefix}: {len(dit_layers)}") - return len(dit_layers) - - -# similar to SD3 but only for the last norm layer -def swap_scale_shift(weight, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_transformer(state_dict): - converted_state_dict = {} - state_dict_keys = list(state_dict.keys()) - - converted_state_dict["register_tokens"] = state_dict.pop("model.register_tokens") - converted_state_dict["pos_embed.pos_embed"] = state_dict.pop("model.positional_encoding") - converted_state_dict["pos_embed.proj.weight"] = state_dict.pop("model.init_x_linear.weight") - converted_state_dict["pos_embed.proj.bias"] = state_dict.pop("model.init_x_linear.bias") - - converted_state_dict["time_step_proj.linear_1.weight"] = state_dict.pop("model.t_embedder.mlp.0.weight") - converted_state_dict["time_step_proj.linear_1.bias"] = state_dict.pop("model.t_embedder.mlp.0.bias") - converted_state_dict["time_step_proj.linear_2.weight"] = state_dict.pop("model.t_embedder.mlp.2.weight") - converted_state_dict["time_step_proj.linear_2.bias"] = state_dict.pop("model.t_embedder.mlp.2.bias") - - converted_state_dict["context_embedder.weight"] = state_dict.pop("model.cond_seq_linear.weight") - - mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers") - single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers") - - # MMDiT blocks 🎸. - for i in range(mmdit_layers): - # feed-forward - path_mapping = {"mlpX": "ff", "mlpC": "ff_context"} - weight_mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for orig_k, diffuser_k in path_mapping.items(): - for k, v in weight_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.{v}.weight"] = state_dict.pop( - f"model.double_layers.{i}.{orig_k}.{k}.weight" - ) - - # norms - path_mapping = {"modX": "norm1", "modC": "norm1_context"} - for orig_k, diffuser_k in path_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.linear.weight"] = state_dict.pop( - f"model.double_layers.{i}.{orig_k}.1.weight" - ) - - # attns - x_attn_mapping = {"w2q": "to_q", "w2k": "to_k", "w2v": "to_v", "w2o": "to_out.0"} - context_attn_mapping = {"w1q": "add_q_proj", "w1k": "add_k_proj", "w1v": "add_v_proj", "w1o": "to_add_out"} - for attn_mapping in [x_attn_mapping, context_attn_mapping]: - for k, v in attn_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.attn.{v}.weight"] = state_dict.pop( - f"model.double_layers.{i}.attn.{k}.weight" - ) - - # Single-DiT blocks. - for i in range(single_dit_layers): - # feed-forward - mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for k, v in mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.ff.{v}.weight"] = state_dict.pop( - f"model.single_layers.{i}.mlp.{k}.weight" - ) - - # norms - converted_state_dict[f"single_transformer_blocks.{i}.norm1.linear.weight"] = state_dict.pop( - f"model.single_layers.{i}.modCX.1.weight" - ) - - # attns - x_attn_mapping = {"w1q": "to_q", "w1k": "to_k", "w1v": "to_v", "w1o": "to_out.0"} - for k, v in x_attn_mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.attn.{v}.weight"] = state_dict.pop( - f"model.single_layers.{i}.attn.{k}.weight" - ) - - # Final blocks. - converted_state_dict["proj_out.weight"] = state_dict.pop("model.final_linear.weight") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(state_dict.pop("model.modF.1.weight"), dim=None) - - return converted_state_dict - - -@torch.no_grad() -def populate_state_dict(args): - original_state_dict = load_original_state_dict(args) - state_dict_keys = list(original_state_dict.keys()) - mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers") - single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers") - - converted_state_dict = convert_transformer(original_state_dict) - model_diffusers = AuraFlowTransformer2DModel( - num_mmdit_layers=mmdit_layers, num_single_dit_layers=single_dit_layers - ) - model_diffusers.load_state_dict(converted_state_dict, strict=True) - - return model_diffusers - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--original_state_dict_repo_id", default="AuraDiffusion/auradiffusion-v0.1a0", type=str) - parser.add_argument("--dump_path", default="aura-flow", type=str) - parser.add_argument("--hub_id", default=None, type=str) - args = parser.parse_args() - - model_diffusers = populate_state_dict(args) - model_diffusers.save_pretrained(args.dump_path) - if args.hub_id is not None: - model_diffusers.push_to_hub(args.hub_id) diff --git a/scripts/convert_blipdiffusion_to_diffusers.py b/scripts/convert_blipdiffusion_to_diffusers.py deleted file mode 100644 index 2c286ea0fdc7..000000000000 --- a/scripts/convert_blipdiffusion_to_diffusers.py +++ /dev/null @@ -1,344 +0,0 @@ -""" -This script requires you to build `LAVIS` from source, since the pip version doesn't have BLIP Diffusion. Follow instructions here: https://github.com/salesforce/LAVIS/tree/main. -""" - -import argparse -import os -import tempfile - -import torch -from lavis.models import load_model_and_preprocess -from transformers import CLIPTokenizer -from transformers.models.blip_2.configuration_blip_2 import Blip2Config - -from diffusers import ( - AutoencoderKL, - PNDMScheduler, - UNet2DConditionModel, -) -from diffusers.pipelines import BlipDiffusionPipeline -from diffusers.pipelines.blip_diffusion.blip_image_processing import BlipImageProcessor -from diffusers.pipelines.blip_diffusion.modeling_blip2 import Blip2QFormerModel -from diffusers.pipelines.blip_diffusion.modeling_ctx_clip import ContextCLIPTextModel - - -BLIP2_CONFIG = { - "vision_config": { - "hidden_size": 1024, - "num_hidden_layers": 23, - "num_attention_heads": 16, - "image_size": 224, - "patch_size": 14, - "intermediate_size": 4096, - "hidden_act": "quick_gelu", - }, - "qformer_config": { - "cross_attention_frequency": 1, - "encoder_hidden_size": 1024, - "vocab_size": 30523, - }, - "num_query_tokens": 16, -} -blip2config = Blip2Config(**BLIP2_CONFIG) - - -def qformer_model_from_original_config(): - qformer = Blip2QFormerModel(blip2config) - return qformer - - -def embeddings_from_original_checkpoint(model, diffuser_embeddings_prefix, original_embeddings_prefix): - embeddings = {} - embeddings.update( - { - f"{diffuser_embeddings_prefix}.word_embeddings.weight": model[ - f"{original_embeddings_prefix}.word_embeddings.weight" - ] - } - ) - embeddings.update( - { - f"{diffuser_embeddings_prefix}.position_embeddings.weight": model[ - f"{original_embeddings_prefix}.position_embeddings.weight" - ] - } - ) - embeddings.update( - {f"{diffuser_embeddings_prefix}.LayerNorm.weight": model[f"{original_embeddings_prefix}.LayerNorm.weight"]} - ) - embeddings.update( - {f"{diffuser_embeddings_prefix}.LayerNorm.bias": model[f"{original_embeddings_prefix}.LayerNorm.bias"]} - ) - return embeddings - - -def proj_layer_from_original_checkpoint(model, diffuser_proj_prefix, original_proj_prefix): - proj_layer = {} - proj_layer.update({f"{diffuser_proj_prefix}.dense1.weight": model[f"{original_proj_prefix}.dense1.weight"]}) - proj_layer.update({f"{diffuser_proj_prefix}.dense1.bias": model[f"{original_proj_prefix}.dense1.bias"]}) - proj_layer.update({f"{diffuser_proj_prefix}.dense2.weight": model[f"{original_proj_prefix}.dense2.weight"]}) - proj_layer.update({f"{diffuser_proj_prefix}.dense2.bias": model[f"{original_proj_prefix}.dense2.bias"]}) - proj_layer.update({f"{diffuser_proj_prefix}.LayerNorm.weight": model[f"{original_proj_prefix}.LayerNorm.weight"]}) - proj_layer.update({f"{diffuser_proj_prefix}.LayerNorm.bias": model[f"{original_proj_prefix}.LayerNorm.bias"]}) - return proj_layer - - -def attention_from_original_checkpoint(model, diffuser_attention_prefix, original_attention_prefix): - attention = {} - attention.update( - { - f"{diffuser_attention_prefix}.attention.query.weight": model[ - f"{original_attention_prefix}.self.query.weight" - ] - } - ) - attention.update( - {f"{diffuser_attention_prefix}.attention.query.bias": model[f"{original_attention_prefix}.self.query.bias"]} - ) - attention.update( - {f"{diffuser_attention_prefix}.attention.key.weight": model[f"{original_attention_prefix}.self.key.weight"]} - ) - attention.update( - {f"{diffuser_attention_prefix}.attention.key.bias": model[f"{original_attention_prefix}.self.key.bias"]} - ) - attention.update( - { - f"{diffuser_attention_prefix}.attention.value.weight": model[ - f"{original_attention_prefix}.self.value.weight" - ] - } - ) - attention.update( - {f"{diffuser_attention_prefix}.attention.value.bias": model[f"{original_attention_prefix}.self.value.bias"]} - ) - attention.update( - {f"{diffuser_attention_prefix}.output.dense.weight": model[f"{original_attention_prefix}.output.dense.weight"]} - ) - attention.update( - {f"{diffuser_attention_prefix}.output.dense.bias": model[f"{original_attention_prefix}.output.dense.bias"]} - ) - attention.update( - { - f"{diffuser_attention_prefix}.output.LayerNorm.weight": model[ - f"{original_attention_prefix}.output.LayerNorm.weight" - ] - } - ) - attention.update( - { - f"{diffuser_attention_prefix}.output.LayerNorm.bias": model[ - f"{original_attention_prefix}.output.LayerNorm.bias" - ] - } - ) - return attention - - -def output_layers_from_original_checkpoint(model, diffuser_output_prefix, original_output_prefix): - output_layers = {} - output_layers.update({f"{diffuser_output_prefix}.dense.weight": model[f"{original_output_prefix}.dense.weight"]}) - output_layers.update({f"{diffuser_output_prefix}.dense.bias": model[f"{original_output_prefix}.dense.bias"]}) - output_layers.update( - {f"{diffuser_output_prefix}.LayerNorm.weight": model[f"{original_output_prefix}.LayerNorm.weight"]} - ) - output_layers.update( - {f"{diffuser_output_prefix}.LayerNorm.bias": model[f"{original_output_prefix}.LayerNorm.bias"]} - ) - return output_layers - - -def encoder_from_original_checkpoint(model, diffuser_encoder_prefix, original_encoder_prefix): - encoder = {} - for i in range(blip2config.qformer_config.num_hidden_layers): - encoder.update( - attention_from_original_checkpoint( - model, f"{diffuser_encoder_prefix}.{i}.attention", f"{original_encoder_prefix}.{i}.attention" - ) - ) - encoder.update( - attention_from_original_checkpoint( - model, f"{diffuser_encoder_prefix}.{i}.crossattention", f"{original_encoder_prefix}.{i}.crossattention" - ) - ) - - encoder.update( - { - f"{diffuser_encoder_prefix}.{i}.intermediate.dense.weight": model[ - f"{original_encoder_prefix}.{i}.intermediate.dense.weight" - ] - } - ) - encoder.update( - { - f"{diffuser_encoder_prefix}.{i}.intermediate.dense.bias": model[ - f"{original_encoder_prefix}.{i}.intermediate.dense.bias" - ] - } - ) - encoder.update( - { - f"{diffuser_encoder_prefix}.{i}.intermediate_query.dense.weight": model[ - f"{original_encoder_prefix}.{i}.intermediate_query.dense.weight" - ] - } - ) - encoder.update( - { - f"{diffuser_encoder_prefix}.{i}.intermediate_query.dense.bias": model[ - f"{original_encoder_prefix}.{i}.intermediate_query.dense.bias" - ] - } - ) - - encoder.update( - output_layers_from_original_checkpoint( - model, f"{diffuser_encoder_prefix}.{i}.output", f"{original_encoder_prefix}.{i}.output" - ) - ) - encoder.update( - output_layers_from_original_checkpoint( - model, f"{diffuser_encoder_prefix}.{i}.output_query", f"{original_encoder_prefix}.{i}.output_query" - ) - ) - return encoder - - -def visual_encoder_layer_from_original_checkpoint(model, diffuser_prefix, original_prefix): - visual_encoder_layer = {} - - visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm1.weight": model[f"{original_prefix}.ln_1.weight"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm1.bias": model[f"{original_prefix}.ln_1.bias"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm2.weight": model[f"{original_prefix}.ln_2.weight"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.layer_norm2.bias": model[f"{original_prefix}.ln_2.bias"]}) - visual_encoder_layer.update( - {f"{diffuser_prefix}.self_attn.qkv.weight": model[f"{original_prefix}.attn.in_proj_weight"]} - ) - visual_encoder_layer.update( - {f"{diffuser_prefix}.self_attn.qkv.bias": model[f"{original_prefix}.attn.in_proj_bias"]} - ) - visual_encoder_layer.update( - {f"{diffuser_prefix}.self_attn.projection.weight": model[f"{original_prefix}.attn.out_proj.weight"]} - ) - visual_encoder_layer.update( - {f"{diffuser_prefix}.self_attn.projection.bias": model[f"{original_prefix}.attn.out_proj.bias"]} - ) - visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc1.weight": model[f"{original_prefix}.mlp.c_fc.weight"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc1.bias": model[f"{original_prefix}.mlp.c_fc.bias"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc2.weight": model[f"{original_prefix}.mlp.c_proj.weight"]}) - visual_encoder_layer.update({f"{diffuser_prefix}.mlp.fc2.bias": model[f"{original_prefix}.mlp.c_proj.bias"]}) - - return visual_encoder_layer - - -def visual_encoder_from_original_checkpoint(model, diffuser_prefix, original_prefix): - visual_encoder = {} - - visual_encoder.update( - { - f"{diffuser_prefix}.embeddings.class_embedding": model[f"{original_prefix}.class_embedding"] - .unsqueeze(0) - .unsqueeze(0) - } - ) - visual_encoder.update( - { - f"{diffuser_prefix}.embeddings.position_embedding": model[ - f"{original_prefix}.positional_embedding" - ].unsqueeze(0) - } - ) - visual_encoder.update( - {f"{diffuser_prefix}.embeddings.patch_embedding.weight": model[f"{original_prefix}.conv1.weight"]} - ) - visual_encoder.update({f"{diffuser_prefix}.pre_layernorm.weight": model[f"{original_prefix}.ln_pre.weight"]}) - visual_encoder.update({f"{diffuser_prefix}.pre_layernorm.bias": model[f"{original_prefix}.ln_pre.bias"]}) - - for i in range(blip2config.vision_config.num_hidden_layers): - visual_encoder.update( - visual_encoder_layer_from_original_checkpoint( - model, f"{diffuser_prefix}.encoder.layers.{i}", f"{original_prefix}.transformer.resblocks.{i}" - ) - ) - - visual_encoder.update({f"{diffuser_prefix}.post_layernorm.weight": model["blip.ln_vision.weight"]}) - visual_encoder.update({f"{diffuser_prefix}.post_layernorm.bias": model["blip.ln_vision.bias"]}) - - return visual_encoder - - -def qformer_original_checkpoint_to_diffusers_checkpoint(model): - qformer_checkpoint = {} - qformer_checkpoint.update(embeddings_from_original_checkpoint(model, "embeddings", "blip.Qformer.bert.embeddings")) - qformer_checkpoint.update({"query_tokens": model["blip.query_tokens"]}) - qformer_checkpoint.update(proj_layer_from_original_checkpoint(model, "proj_layer", "proj_layer")) - qformer_checkpoint.update( - encoder_from_original_checkpoint(model, "encoder.layer", "blip.Qformer.bert.encoder.layer") - ) - qformer_checkpoint.update(visual_encoder_from_original_checkpoint(model, "visual_encoder", "blip.visual_encoder")) - return qformer_checkpoint - - -def get_qformer(model): - print("loading qformer") - - qformer = qformer_model_from_original_config() - qformer_diffusers_checkpoint = qformer_original_checkpoint_to_diffusers_checkpoint(model) - - load_checkpoint_to_model(qformer_diffusers_checkpoint, qformer) - - print("done loading qformer") - return qformer - - -def load_checkpoint_to_model(checkpoint, model): - with tempfile.NamedTemporaryFile(delete=False) as file: - torch.save(checkpoint, file.name) - del checkpoint - model.load_state_dict(torch.load(file.name), strict=False) - - os.remove(file.name) - - -def save_blip_diffusion_model(model, args): - qformer = get_qformer(model) - qformer.eval() - - text_encoder = ContextCLIPTextModel.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="text_encoder" - ) - vae = AutoencoderKL.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="vae") - unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet") - vae.eval() - text_encoder.eval() - scheduler = PNDMScheduler( - beta_start=0.00085, - beta_end=0.012, - beta_schedule="scaled_linear", - set_alpha_to_one=False, - skip_prk_steps=True, - ) - tokenizer = CLIPTokenizer.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="tokenizer") - image_processor = BlipImageProcessor() - blip_diffusion = BlipDiffusionPipeline( - tokenizer=tokenizer, - text_encoder=text_encoder, - vae=vae, - unet=unet, - scheduler=scheduler, - qformer=qformer, - image_processor=image_processor, - ) - blip_diffusion.save_pretrained(args.checkpoint_path) - - -def main(args): - model, _, _ = load_model_and_preprocess("blip_diffusion", "base", device="cpu", is_eval=True) - save_blip_diffusion_model(model.state_dict(), args) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") - args = parser.parse_args() - - main(args) diff --git a/scripts/convert_checkpoint.py b/scripts/convert_checkpoint.py new file mode 100644 index 000000000000..f378601efe5e --- /dev/null +++ b/scripts/convert_checkpoint.py @@ -0,0 +1,103 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Convert any registered component's tensor checkpoint in either direction, without constructing a model.""" + +import argparse +import json +from pathlib import Path + +from diffusers.loaders.conversion.configs import get_config_preset, list_config_presets +from diffusers.loaders.conversion.io import convert_checkpoint +from diffusers.loaders.conversion.registry import CONVERSION_BUILDERS +from diffusers.loaders.conversion.source import load_source_manifest + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--list-models", action="store_true", help="List registered component classes and exit") + parser.add_argument("--list-presets", action="store_true", help="List configuration presets and helpers") + inputs = parser.add_mutually_exclusive_group() + inputs.add_argument("--input", help="Local component directory, tensor file, or shard index") + inputs.add_argument("--input-manifest", help="JSON sources manifest for a component stored across several files") + parser.add_argument("--output", help="New safetensors component directory or PyTorch checkpoint file") + configs = parser.add_mutually_exclusive_group() + configs.add_argument("--config", help="Matching Diffusers component config JSON (required for original inputs)") + configs.add_argument("--preset", help="Built-in configuration preset or qualified original-config helper") + parser.add_argument("--preset-args", default="{}", help="JSON keyword arguments for a configuration helper") + parser.add_argument("--preset-component", help="Variant key when a preset holds several configurations") + parser.add_argument("--model-class", help="Override the config's _class_name") + parser.add_argument("--direction", choices=("to-original", "to-diffusers"), default="to-original") + parser.add_argument("--original-format", help="Select a model definition's original_format variant") + parser.add_argument("--input-prefix", default="", help="Select and strip this component key prefix") + parser.add_argument( + "--input-wrapper", action="append", default=[], help="Nested PyTorch wrapper key; repeat for nesting" + ) + parser.add_argument("--output-prefix", default="", help="Prepend this prefix to output keys") + parser.add_argument( + "--output-format", + choices=("safetensors", "pytorch"), + default="safetensors", + help="Write a safetensors directory or a single PyTorch file", + ) + parser.add_argument( + "--output-wrapper", action="append", default=[], help="PyTorch output wrapper key; repeat for nesting" + ) + parser.add_argument( + "--max-shard-size", + type=int, + default=5_000_000_000, + help="Maximum shard size in bytes (one tensor may exceed it)", + ) + args = parser.parse_args() + if args.list_models: + print("\n".join(sorted(CONVERSION_BUILDERS))) + return + if args.list_presets: + print("\n".join(list_config_presets())) + return + if not (args.input or args.input_manifest) or not args.output: + parser.error("--input or --input-manifest, and --output are required") + config_path = Path(args.config) if args.config else Path(args.input or args.input_manifest).parent / "config.json" + if args.input and not args.config and Path(args.input).is_dir(): + config_path = Path(args.input) / "config.json" + if not args.preset and not config_path.is_file(): + parser.error("Provide --config with the matching Diffusers component configuration") + config = ( + get_config_preset(args.preset, arguments=json.loads(args.preset_args), component=args.preset_component) + if args.preset + else json.loads(config_path.read_text(encoding="utf-8")) + ) + if args.original_format: + config["original_format"] = args.original_format + if not args.model_class and not config.get("_class_name"): + parser.error("Provide --model-class when the configuration does not contain _class_name") + output = convert_checkpoint( + load_source_manifest(args.input_manifest) if args.input_manifest else args.input, + args.output, + config=config, + model_class=args.model_class, + reverse=args.direction == "to-original", + input_prefix=args.input_prefix, + input_wrapper=args.input_wrapper, + output_prefix=args.output_prefix, + output_format=args.output_format, + output_wrapper=args.output_wrapper, + max_shard_size=args.max_shard_size, + ) + print(output) + + +if __name__ == "__main__": + main() diff --git a/scripts/convert_consistency_decoder.py b/scripts/convert_consistency_decoder.py deleted file mode 100644 index 9e289457752b..000000000000 --- a/scripts/convert_consistency_decoder.py +++ /dev/null @@ -1,1131 +0,0 @@ -import math -import os -import urllib -import warnings -from argparse import ArgumentParser - -import torch -import torch.nn as nn -import torch.nn.functional as F -from huggingface_hub.utils import insecure_hashlib -from safetensors.torch import load_file as stl -from tqdm import tqdm - -from diffusers import AutoencoderKL, ConsistencyDecoderVAE, DiffusionPipeline, StableDiffusionPipeline, UNet2DModel -from diffusers.models.autoencoders.vae import Encoder -from diffusers.models.embeddings import TimestepEmbedding -from diffusers.models.unets.unet_2d_blocks import ResnetDownsampleBlock2D, ResnetUpsampleBlock2D, UNetMidBlock2D - - -args = ArgumentParser() -args.add_argument("--save_pretrained", required=False, default=None, type=str) -args.add_argument("--test_image", required=True, type=str) -args = args.parse_args() - - -def _extract_into_tensor(arr, timesteps, broadcast_shape): - # from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 - # """ - res = arr[timesteps].float() - dims_to_append = len(broadcast_shape) - len(res.shape) - return res[(...,) + (None,) * dims_to_append] - - -def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999): - # from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L45 - betas = [] - for i in range(num_diffusion_timesteps): - t1 = i / num_diffusion_timesteps - t2 = (i + 1) / num_diffusion_timesteps - betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta)) - return torch.tensor(betas) - - -def _download(url: str, root: str): - os.makedirs(root, exist_ok=True) - filename = os.path.basename(url) - - expected_sha256 = url.split("/")[-2] - download_target = os.path.join(root, filename) - - if os.path.exists(download_target) and not os.path.isfile(download_target): - raise RuntimeError(f"{download_target} exists and is not a regular file") - - if os.path.isfile(download_target): - if insecure_hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256: - return download_target - else: - warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file") - - with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: - with tqdm( - total=int(source.info().get("Content-Length")), - ncols=80, - unit="iB", - unit_scale=True, - unit_divisor=1024, - ) as loop: - while True: - buffer = source.read(8192) - if not buffer: - break - - output.write(buffer) - loop.update(len(buffer)) - - if insecure_hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256: - raise RuntimeError("Model has been downloaded but the SHA256 checksum does not match") - - return download_target - - -class ConsistencyDecoder: - def __init__(self, device="cuda:0", download_root=os.path.expanduser("~/.cache/clip")): - self.n_distilled_steps = 64 - download_target = _download( - "https://openaipublic.azureedge.net/diff-vae/c9cebd3132dd9c42936d803e33424145a748843c8f716c0814838bdc8a2fe7cb/decoder.pt", - download_root, - ) - self.ckpt = torch.jit.load(download_target).to(device) - self.device = device - sigma_data = 0.5 - betas = betas_for_alpha_bar(1024, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2).to(device) - alphas = 1.0 - betas - alphas_cumprod = torch.cumprod(alphas, dim=0) - self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) - self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) - sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod) - sigmas = torch.sqrt(1.0 / alphas_cumprod - 1) - self.c_skip = sqrt_recip_alphas_cumprod * sigma_data**2 / (sigmas**2 + sigma_data**2) - self.c_out = sigmas * sigma_data / (sigmas**2 + sigma_data**2) ** 0.5 - self.c_in = sqrt_recip_alphas_cumprod / (sigmas**2 + sigma_data**2) ** 0.5 - - @staticmethod - def round_timesteps(timesteps, total_timesteps, n_distilled_steps, truncate_start=True): - with torch.no_grad(): - space = torch.div(total_timesteps, n_distilled_steps, rounding_mode="floor") - rounded_timesteps = (torch.div(timesteps, space, rounding_mode="floor") + 1) * space - if truncate_start: - rounded_timesteps[rounded_timesteps == total_timesteps] -= space - else: - rounded_timesteps[rounded_timesteps == total_timesteps] -= space - rounded_timesteps[rounded_timesteps == 0] += space - return rounded_timesteps - - @staticmethod - def ldm_transform_latent(z, extra_scale_factor=1): - channel_means = [0.38862467, 0.02253063, 0.07381133, -0.0171294] - channel_stds = [0.9654121, 1.0440036, 0.76147926, 0.77022034] - - if len(z.shape) != 4: - raise ValueError() - - z = z * 0.18215 - channels = [z[:, i] for i in range(z.shape[1])] - - channels = [extra_scale_factor * (c - channel_means[i]) / channel_stds[i] for i, c in enumerate(channels)] - return torch.stack(channels, dim=1) - - @torch.no_grad() - def __call__( - self, - features: torch.Tensor, - schedule=[1.0, 0.5], - generator=None, - ): - features = self.ldm_transform_latent(features) - ts = self.round_timesteps( - torch.arange(0, 1024), - 1024, - self.n_distilled_steps, - truncate_start=False, - ) - shape = ( - features.size(0), - 3, - 8 * features.size(2), - 8 * features.size(3), - ) - x_start = torch.zeros(shape, device=features.device, dtype=features.dtype) - schedule_timesteps = [int((1024 - 1) * s) for s in schedule] - for i in schedule_timesteps: - t = ts[i].item() - t_ = torch.tensor([t] * features.shape[0]).to(self.device) - # noise = torch.randn_like(x_start) - noise = torch.randn(x_start.shape, dtype=x_start.dtype, generator=generator).to(device=x_start.device) - x_start = ( - _extract_into_tensor(self.sqrt_alphas_cumprod, t_, x_start.shape) * x_start - + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t_, x_start.shape) * noise - ) - c_in = _extract_into_tensor(self.c_in, t_, x_start.shape) - - import torch.nn.functional as F - - from diffusers import UNet2DModel - - if isinstance(self.ckpt, UNet2DModel): - input = torch.concat([c_in * x_start, F.upsample_nearest(features, scale_factor=8)], dim=1) - model_output = self.ckpt(input, t_).sample - else: - model_output = self.ckpt(c_in * x_start, t_, features=features) - - B, C = x_start.shape[:2] - model_output, _ = torch.split(model_output, C, dim=1) - pred_xstart = ( - _extract_into_tensor(self.c_out, t_, x_start.shape) * model_output - + _extract_into_tensor(self.c_skip, t_, x_start.shape) * x_start - ).clamp(-1, 1) - x_start = pred_xstart - return x_start - - -def save_image(image, name): - import numpy as np - from PIL import Image - - image = image[0].cpu().numpy() - image = (image + 1.0) * 127.5 - image = image.clip(0, 255).astype(np.uint8) - image = Image.fromarray(image.transpose(1, 2, 0)) - image.save(name) - - -def load_image(uri, size=None, center_crop=False): - import numpy as np - from PIL import Image - - image = Image.open(uri) - if center_crop: - image = image.crop( - ( - (image.width - min(image.width, image.height)) // 2, - (image.height - min(image.width, image.height)) // 2, - (image.width + min(image.width, image.height)) // 2, - (image.height + min(image.width, image.height)) // 2, - ) - ) - if size is not None: - image = image.resize(size) - image = torch.tensor(np.array(image).transpose(2, 0, 1)).unsqueeze(0).float() - image = image / 127.5 - 1.0 - return image - - -class TimestepEmbedding_(nn.Module): - def __init__(self, n_time=1024, n_emb=320, n_out=1280) -> None: - super().__init__() - self.emb = nn.Embedding(n_time, n_emb) - self.f_1 = nn.Linear(n_emb, n_out) - self.f_2 = nn.Linear(n_out, n_out) - - def forward(self, x) -> torch.Tensor: - x = self.emb(x) - x = self.f_1(x) - x = F.silu(x) - return self.f_2(x) - - -class ImageEmbedding(nn.Module): - def __init__(self, in_channels=7, out_channels=320) -> None: - super().__init__() - self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) - - def forward(self, x) -> torch.Tensor: - return self.f(x) - - -class ImageUnembedding(nn.Module): - def __init__(self, in_channels=320, out_channels=6) -> None: - super().__init__() - self.gn = nn.GroupNorm(32, in_channels) - self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) - - def forward(self, x) -> torch.Tensor: - return self.f(F.silu(self.gn(x))) - - -class ConvResblock(nn.Module): - def __init__(self, in_features=320, out_features=320) -> None: - super().__init__() - self.f_t = nn.Linear(1280, out_features * 2) - - self.gn_1 = nn.GroupNorm(32, in_features) - self.f_1 = nn.Conv2d(in_features, out_features, kernel_size=3, padding=1) - - self.gn_2 = nn.GroupNorm(32, out_features) - self.f_2 = nn.Conv2d(out_features, out_features, kernel_size=3, padding=1) - - skip_conv = in_features != out_features - self.f_s = nn.Conv2d(in_features, out_features, kernel_size=1, padding=0) if skip_conv else nn.Identity() - - def forward(self, x, t): - x_skip = x - t = self.f_t(F.silu(t)) - t = t.chunk(2, dim=1) - t_1 = t[0].unsqueeze(dim=2).unsqueeze(dim=3) + 1 - t_2 = t[1].unsqueeze(dim=2).unsqueeze(dim=3) - - gn_1 = F.silu(self.gn_1(x)) - f_1 = self.f_1(gn_1) - - gn_2 = self.gn_2(f_1) - - return self.f_s(x_skip) + self.f_2(F.silu(gn_2 * t_1 + t_2)) - - -# Also ConvResblock -class Downsample(nn.Module): - def __init__(self, in_channels=320) -> None: - super().__init__() - self.f_t = nn.Linear(1280, in_channels * 2) - - self.gn_1 = nn.GroupNorm(32, in_channels) - self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) - self.gn_2 = nn.GroupNorm(32, in_channels) - - self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) - - def forward(self, x, t) -> torch.Tensor: - x_skip = x - - t = self.f_t(F.silu(t)) - t_1, t_2 = t.chunk(2, dim=1) - t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1 - t_2 = t_2.unsqueeze(2).unsqueeze(3) - - gn_1 = F.silu(self.gn_1(x)) - avg_pool2d = F.avg_pool2d(gn_1, kernel_size=(2, 2), stride=None) - - f_1 = self.f_1(avg_pool2d) - gn_2 = self.gn_2(f_1) - - f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2))) - - return f_2 + F.avg_pool2d(x_skip, kernel_size=(2, 2), stride=None) - - -# Also ConvResblock -class Upsample(nn.Module): - def __init__(self, in_channels=1024) -> None: - super().__init__() - self.f_t = nn.Linear(1280, in_channels * 2) - - self.gn_1 = nn.GroupNorm(32, in_channels) - self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) - self.gn_2 = nn.GroupNorm(32, in_channels) - - self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) - - def forward(self, x, t) -> torch.Tensor: - x_skip = x - - t = self.f_t(F.silu(t)) - t_1, t_2 = t.chunk(2, dim=1) - t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1 - t_2 = t_2.unsqueeze(2).unsqueeze(3) - - gn_1 = F.silu(self.gn_1(x)) - upsample = F.upsample_nearest(gn_1, scale_factor=2) - f_1 = self.f_1(upsample) - gn_2 = self.gn_2(f_1) - - f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2))) - - return f_2 + F.upsample_nearest(x_skip, scale_factor=2) - - -class ConvUNetVAE(nn.Module): - def __init__(self) -> None: - super().__init__() - self.embed_image = ImageEmbedding() - self.embed_time = TimestepEmbedding_() - - down_0 = nn.ModuleList( - [ - ConvResblock(320, 320), - ConvResblock(320, 320), - ConvResblock(320, 320), - Downsample(320), - ] - ) - down_1 = nn.ModuleList( - [ - ConvResblock(320, 640), - ConvResblock(640, 640), - ConvResblock(640, 640), - Downsample(640), - ] - ) - down_2 = nn.ModuleList( - [ - ConvResblock(640, 1024), - ConvResblock(1024, 1024), - ConvResblock(1024, 1024), - Downsample(1024), - ] - ) - down_3 = nn.ModuleList( - [ - ConvResblock(1024, 1024), - ConvResblock(1024, 1024), - ConvResblock(1024, 1024), - ] - ) - self.down = nn.ModuleList( - [ - down_0, - down_1, - down_2, - down_3, - ] - ) - - self.mid = nn.ModuleList( - [ - ConvResblock(1024, 1024), - ConvResblock(1024, 1024), - ] - ) - - up_3 = nn.ModuleList( - [ - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 * 2, 1024), - Upsample(1024), - ] - ) - up_2 = nn.ModuleList( - [ - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 * 2, 1024), - ConvResblock(1024 + 640, 1024), - Upsample(1024), - ] - ) - up_1 = nn.ModuleList( - [ - ConvResblock(1024 + 640, 640), - ConvResblock(640 * 2, 640), - ConvResblock(640 * 2, 640), - ConvResblock(320 + 640, 640), - Upsample(640), - ] - ) - up_0 = nn.ModuleList( - [ - ConvResblock(320 + 640, 320), - ConvResblock(320 * 2, 320), - ConvResblock(320 * 2, 320), - ConvResblock(320 * 2, 320), - ] - ) - self.up = nn.ModuleList( - [ - up_0, - up_1, - up_2, - up_3, - ] - ) - - self.output = ImageUnembedding() - - def forward(self, x, t, features) -> torch.Tensor: - converted = hasattr(self, "converted") and self.converted - - x = torch.cat([x, F.upsample_nearest(features, scale_factor=8)], dim=1) - - if converted: - t = self.time_embedding(self.time_proj(t)) - else: - t = self.embed_time(t) - - x = self.embed_image(x) - - skips = [x] - for i, down in enumerate(self.down): - if converted and i in [0, 1, 2, 3]: - x, skips_ = down(x, t) - for skip in skips_: - skips.append(skip) - else: - for block in down: - x = block(x, t) - skips.append(x) - print(x.float().abs().sum()) - - if converted: - x = self.mid(x, t) - else: - for i in range(2): - x = self.mid[i](x, t) - print(x.float().abs().sum()) - - for i, up in enumerate(self.up[::-1]): - if converted and i in [0, 1, 2, 3]: - skip_4 = skips.pop() - skip_3 = skips.pop() - skip_2 = skips.pop() - skip_1 = skips.pop() - skips_ = (skip_1, skip_2, skip_3, skip_4) - x = up(x, skips_, t) - else: - for block in up: - if isinstance(block, ConvResblock): - x = torch.concat([x, skips.pop()], dim=1) - x = block(x, t) - - return self.output(x) - - -def rename_state_dict_key(k): - k = k.replace("blocks.", "") - for i in range(5): - k = k.replace(f"down_{i}_", f"down.{i}.") - k = k.replace(f"conv_{i}.", f"{i}.") - k = k.replace(f"up_{i}_", f"up.{i}.") - k = k.replace(f"mid_{i}", f"mid.{i}") - k = k.replace("upsamp.", "4.") - k = k.replace("downsamp.", "3.") - k = k.replace("f_t.w", "f_t.weight").replace("f_t.b", "f_t.bias") - k = k.replace("f_1.w", "f_1.weight").replace("f_1.b", "f_1.bias") - k = k.replace("f_2.w", "f_2.weight").replace("f_2.b", "f_2.bias") - k = k.replace("f_s.w", "f_s.weight").replace("f_s.b", "f_s.bias") - k = k.replace("f.w", "f.weight").replace("f.b", "f.bias") - k = k.replace("gn_1.g", "gn_1.weight").replace("gn_1.b", "gn_1.bias") - k = k.replace("gn_2.g", "gn_2.weight").replace("gn_2.b", "gn_2.bias") - k = k.replace("gn.g", "gn.weight").replace("gn.b", "gn.bias") - return k - - -def rename_state_dict(sd, embedding): - sd = {rename_state_dict_key(k): v for k, v in sd.items()} - sd["embed_time.emb.weight"] = embedding["weight"] - return sd - - -# encode with stable diffusion vae -pipe = StableDiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16 -) -pipe.vae.cuda() - -# construct original decoder with jitted model -decoder_consistency = ConsistencyDecoder(device="cuda:0") - -# construct UNet code, overwrite the decoder with conv_unet_vae -model = ConvUNetVAE() -model.load_state_dict( - rename_state_dict( - stl("consistency_decoder.safetensors"), - stl("embedding.safetensors"), - ) -) -model = model.cuda() - -decoder_consistency.ckpt = model - -image = load_image(args.test_image, size=(256, 256), center_crop=True) -latent = pipe.vae.encode(image.half().cuda()).latent_dist.sample() - -# decode with gan -sample_gan = pipe.vae.decode(latent).sample.detach() -save_image(sample_gan, "gan.png") - -# decode with conv_unet_vae -sample_consistency_orig = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) -save_image(sample_consistency_orig, "con_orig.png") - - -########### conversion - -print("CONVERSION") - -print("DOWN BLOCK ONE") - -block_one_sd_orig = model.down[0].state_dict() -block_one_sd_new = {} - -for i in range(3): - block_one_sd_new[f"resnets.{i}.norm1.weight"] = block_one_sd_orig.pop(f"{i}.gn_1.weight") - block_one_sd_new[f"resnets.{i}.norm1.bias"] = block_one_sd_orig.pop(f"{i}.gn_1.bias") - block_one_sd_new[f"resnets.{i}.conv1.weight"] = block_one_sd_orig.pop(f"{i}.f_1.weight") - block_one_sd_new[f"resnets.{i}.conv1.bias"] = block_one_sd_orig.pop(f"{i}.f_1.bias") - block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_one_sd_orig.pop(f"{i}.f_t.weight") - block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_one_sd_orig.pop(f"{i}.f_t.bias") - block_one_sd_new[f"resnets.{i}.norm2.weight"] = block_one_sd_orig.pop(f"{i}.gn_2.weight") - block_one_sd_new[f"resnets.{i}.norm2.bias"] = block_one_sd_orig.pop(f"{i}.gn_2.bias") - block_one_sd_new[f"resnets.{i}.conv2.weight"] = block_one_sd_orig.pop(f"{i}.f_2.weight") - block_one_sd_new[f"resnets.{i}.conv2.bias"] = block_one_sd_orig.pop(f"{i}.f_2.bias") - -block_one_sd_new["downsamplers.0.norm1.weight"] = block_one_sd_orig.pop("3.gn_1.weight") -block_one_sd_new["downsamplers.0.norm1.bias"] = block_one_sd_orig.pop("3.gn_1.bias") -block_one_sd_new["downsamplers.0.conv1.weight"] = block_one_sd_orig.pop("3.f_1.weight") -block_one_sd_new["downsamplers.0.conv1.bias"] = block_one_sd_orig.pop("3.f_1.bias") -block_one_sd_new["downsamplers.0.time_emb_proj.weight"] = block_one_sd_orig.pop("3.f_t.weight") -block_one_sd_new["downsamplers.0.time_emb_proj.bias"] = block_one_sd_orig.pop("3.f_t.bias") -block_one_sd_new["downsamplers.0.norm2.weight"] = block_one_sd_orig.pop("3.gn_2.weight") -block_one_sd_new["downsamplers.0.norm2.bias"] = block_one_sd_orig.pop("3.gn_2.bias") -block_one_sd_new["downsamplers.0.conv2.weight"] = block_one_sd_orig.pop("3.f_2.weight") -block_one_sd_new["downsamplers.0.conv2.bias"] = block_one_sd_orig.pop("3.f_2.bias") - -assert len(block_one_sd_orig) == 0 - -block_one = ResnetDownsampleBlock2D( - in_channels=320, - out_channels=320, - temb_channels=1280, - num_layers=3, - add_downsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -block_one.load_state_dict(block_one_sd_new) - -print("DOWN BLOCK TWO") - -block_two_sd_orig = model.down[1].state_dict() -block_two_sd_new = {} - -for i in range(3): - block_two_sd_new[f"resnets.{i}.norm1.weight"] = block_two_sd_orig.pop(f"{i}.gn_1.weight") - block_two_sd_new[f"resnets.{i}.norm1.bias"] = block_two_sd_orig.pop(f"{i}.gn_1.bias") - block_two_sd_new[f"resnets.{i}.conv1.weight"] = block_two_sd_orig.pop(f"{i}.f_1.weight") - block_two_sd_new[f"resnets.{i}.conv1.bias"] = block_two_sd_orig.pop(f"{i}.f_1.bias") - block_two_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_two_sd_orig.pop(f"{i}.f_t.weight") - block_two_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_two_sd_orig.pop(f"{i}.f_t.bias") - block_two_sd_new[f"resnets.{i}.norm2.weight"] = block_two_sd_orig.pop(f"{i}.gn_2.weight") - block_two_sd_new[f"resnets.{i}.norm2.bias"] = block_two_sd_orig.pop(f"{i}.gn_2.bias") - block_two_sd_new[f"resnets.{i}.conv2.weight"] = block_two_sd_orig.pop(f"{i}.f_2.weight") - block_two_sd_new[f"resnets.{i}.conv2.bias"] = block_two_sd_orig.pop(f"{i}.f_2.bias") - - if i == 0: - block_two_sd_new[f"resnets.{i}.conv_shortcut.weight"] = block_two_sd_orig.pop(f"{i}.f_s.weight") - block_two_sd_new[f"resnets.{i}.conv_shortcut.bias"] = block_two_sd_orig.pop(f"{i}.f_s.bias") - -block_two_sd_new["downsamplers.0.norm1.weight"] = block_two_sd_orig.pop("3.gn_1.weight") -block_two_sd_new["downsamplers.0.norm1.bias"] = block_two_sd_orig.pop("3.gn_1.bias") -block_two_sd_new["downsamplers.0.conv1.weight"] = block_two_sd_orig.pop("3.f_1.weight") -block_two_sd_new["downsamplers.0.conv1.bias"] = block_two_sd_orig.pop("3.f_1.bias") -block_two_sd_new["downsamplers.0.time_emb_proj.weight"] = block_two_sd_orig.pop("3.f_t.weight") -block_two_sd_new["downsamplers.0.time_emb_proj.bias"] = block_two_sd_orig.pop("3.f_t.bias") -block_two_sd_new["downsamplers.0.norm2.weight"] = block_two_sd_orig.pop("3.gn_2.weight") -block_two_sd_new["downsamplers.0.norm2.bias"] = block_two_sd_orig.pop("3.gn_2.bias") -block_two_sd_new["downsamplers.0.conv2.weight"] = block_two_sd_orig.pop("3.f_2.weight") -block_two_sd_new["downsamplers.0.conv2.bias"] = block_two_sd_orig.pop("3.f_2.bias") - -assert len(block_two_sd_orig) == 0 - -block_two = ResnetDownsampleBlock2D( - in_channels=320, - out_channels=640, - temb_channels=1280, - num_layers=3, - add_downsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -block_two.load_state_dict(block_two_sd_new) - -print("DOWN BLOCK THREE") - -block_three_sd_orig = model.down[2].state_dict() -block_three_sd_new = {} - -for i in range(3): - block_three_sd_new[f"resnets.{i}.norm1.weight"] = block_three_sd_orig.pop(f"{i}.gn_1.weight") - block_three_sd_new[f"resnets.{i}.norm1.bias"] = block_three_sd_orig.pop(f"{i}.gn_1.bias") - block_three_sd_new[f"resnets.{i}.conv1.weight"] = block_three_sd_orig.pop(f"{i}.f_1.weight") - block_three_sd_new[f"resnets.{i}.conv1.bias"] = block_three_sd_orig.pop(f"{i}.f_1.bias") - block_three_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_three_sd_orig.pop(f"{i}.f_t.weight") - block_three_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_three_sd_orig.pop(f"{i}.f_t.bias") - block_three_sd_new[f"resnets.{i}.norm2.weight"] = block_three_sd_orig.pop(f"{i}.gn_2.weight") - block_three_sd_new[f"resnets.{i}.norm2.bias"] = block_three_sd_orig.pop(f"{i}.gn_2.bias") - block_three_sd_new[f"resnets.{i}.conv2.weight"] = block_three_sd_orig.pop(f"{i}.f_2.weight") - block_three_sd_new[f"resnets.{i}.conv2.bias"] = block_three_sd_orig.pop(f"{i}.f_2.bias") - - if i == 0: - block_three_sd_new[f"resnets.{i}.conv_shortcut.weight"] = block_three_sd_orig.pop(f"{i}.f_s.weight") - block_three_sd_new[f"resnets.{i}.conv_shortcut.bias"] = block_three_sd_orig.pop(f"{i}.f_s.bias") - -block_three_sd_new["downsamplers.0.norm1.weight"] = block_three_sd_orig.pop("3.gn_1.weight") -block_three_sd_new["downsamplers.0.norm1.bias"] = block_three_sd_orig.pop("3.gn_1.bias") -block_three_sd_new["downsamplers.0.conv1.weight"] = block_three_sd_orig.pop("3.f_1.weight") -block_three_sd_new["downsamplers.0.conv1.bias"] = block_three_sd_orig.pop("3.f_1.bias") -block_three_sd_new["downsamplers.0.time_emb_proj.weight"] = block_three_sd_orig.pop("3.f_t.weight") -block_three_sd_new["downsamplers.0.time_emb_proj.bias"] = block_three_sd_orig.pop("3.f_t.bias") -block_three_sd_new["downsamplers.0.norm2.weight"] = block_three_sd_orig.pop("3.gn_2.weight") -block_three_sd_new["downsamplers.0.norm2.bias"] = block_three_sd_orig.pop("3.gn_2.bias") -block_three_sd_new["downsamplers.0.conv2.weight"] = block_three_sd_orig.pop("3.f_2.weight") -block_three_sd_new["downsamplers.0.conv2.bias"] = block_three_sd_orig.pop("3.f_2.bias") - -assert len(block_three_sd_orig) == 0 - -block_three = ResnetDownsampleBlock2D( - in_channels=640, - out_channels=1024, - temb_channels=1280, - num_layers=3, - add_downsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -block_three.load_state_dict(block_three_sd_new) - -print("DOWN BLOCK FOUR") - -block_four_sd_orig = model.down[3].state_dict() -block_four_sd_new = {} - -for i in range(3): - block_four_sd_new[f"resnets.{i}.norm1.weight"] = block_four_sd_orig.pop(f"{i}.gn_1.weight") - block_four_sd_new[f"resnets.{i}.norm1.bias"] = block_four_sd_orig.pop(f"{i}.gn_1.bias") - block_four_sd_new[f"resnets.{i}.conv1.weight"] = block_four_sd_orig.pop(f"{i}.f_1.weight") - block_four_sd_new[f"resnets.{i}.conv1.bias"] = block_four_sd_orig.pop(f"{i}.f_1.bias") - block_four_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_four_sd_orig.pop(f"{i}.f_t.weight") - block_four_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_four_sd_orig.pop(f"{i}.f_t.bias") - block_four_sd_new[f"resnets.{i}.norm2.weight"] = block_four_sd_orig.pop(f"{i}.gn_2.weight") - block_four_sd_new[f"resnets.{i}.norm2.bias"] = block_four_sd_orig.pop(f"{i}.gn_2.bias") - block_four_sd_new[f"resnets.{i}.conv2.weight"] = block_four_sd_orig.pop(f"{i}.f_2.weight") - block_four_sd_new[f"resnets.{i}.conv2.bias"] = block_four_sd_orig.pop(f"{i}.f_2.bias") - -assert len(block_four_sd_orig) == 0 - -block_four = ResnetDownsampleBlock2D( - in_channels=1024, - out_channels=1024, - temb_channels=1280, - num_layers=3, - add_downsample=False, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -block_four.load_state_dict(block_four_sd_new) - - -print("MID BLOCK 1") - -mid_block_one_sd_orig = model.mid.state_dict() -mid_block_one_sd_new = {} - -for i in range(2): - mid_block_one_sd_new[f"resnets.{i}.norm1.weight"] = mid_block_one_sd_orig.pop(f"{i}.gn_1.weight") - mid_block_one_sd_new[f"resnets.{i}.norm1.bias"] = mid_block_one_sd_orig.pop(f"{i}.gn_1.bias") - mid_block_one_sd_new[f"resnets.{i}.conv1.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_1.weight") - mid_block_one_sd_new[f"resnets.{i}.conv1.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_1.bias") - mid_block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_t.weight") - mid_block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_t.bias") - mid_block_one_sd_new[f"resnets.{i}.norm2.weight"] = mid_block_one_sd_orig.pop(f"{i}.gn_2.weight") - mid_block_one_sd_new[f"resnets.{i}.norm2.bias"] = mid_block_one_sd_orig.pop(f"{i}.gn_2.bias") - mid_block_one_sd_new[f"resnets.{i}.conv2.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_2.weight") - mid_block_one_sd_new[f"resnets.{i}.conv2.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_2.bias") - -assert len(mid_block_one_sd_orig) == 0 - -mid_block_one = UNetMidBlock2D( - in_channels=1024, - temb_channels=1280, - num_layers=1, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, - add_attention=False, -) - -mid_block_one.load_state_dict(mid_block_one_sd_new) - -print("UP BLOCK ONE") - -up_block_one_sd_orig = model.up[-1].state_dict() -up_block_one_sd_new = {} - -for i in range(4): - up_block_one_sd_new[f"resnets.{i}.norm1.weight"] = up_block_one_sd_orig.pop(f"{i}.gn_1.weight") - up_block_one_sd_new[f"resnets.{i}.norm1.bias"] = up_block_one_sd_orig.pop(f"{i}.gn_1.bias") - up_block_one_sd_new[f"resnets.{i}.conv1.weight"] = up_block_one_sd_orig.pop(f"{i}.f_1.weight") - up_block_one_sd_new[f"resnets.{i}.conv1.bias"] = up_block_one_sd_orig.pop(f"{i}.f_1.bias") - up_block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_one_sd_orig.pop(f"{i}.f_t.weight") - up_block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_one_sd_orig.pop(f"{i}.f_t.bias") - up_block_one_sd_new[f"resnets.{i}.norm2.weight"] = up_block_one_sd_orig.pop(f"{i}.gn_2.weight") - up_block_one_sd_new[f"resnets.{i}.norm2.bias"] = up_block_one_sd_orig.pop(f"{i}.gn_2.bias") - up_block_one_sd_new[f"resnets.{i}.conv2.weight"] = up_block_one_sd_orig.pop(f"{i}.f_2.weight") - up_block_one_sd_new[f"resnets.{i}.conv2.bias"] = up_block_one_sd_orig.pop(f"{i}.f_2.bias") - up_block_one_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_one_sd_orig.pop(f"{i}.f_s.weight") - up_block_one_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_one_sd_orig.pop(f"{i}.f_s.bias") - -up_block_one_sd_new["upsamplers.0.norm1.weight"] = up_block_one_sd_orig.pop("4.gn_1.weight") -up_block_one_sd_new["upsamplers.0.norm1.bias"] = up_block_one_sd_orig.pop("4.gn_1.bias") -up_block_one_sd_new["upsamplers.0.conv1.weight"] = up_block_one_sd_orig.pop("4.f_1.weight") -up_block_one_sd_new["upsamplers.0.conv1.bias"] = up_block_one_sd_orig.pop("4.f_1.bias") -up_block_one_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_one_sd_orig.pop("4.f_t.weight") -up_block_one_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_one_sd_orig.pop("4.f_t.bias") -up_block_one_sd_new["upsamplers.0.norm2.weight"] = up_block_one_sd_orig.pop("4.gn_2.weight") -up_block_one_sd_new["upsamplers.0.norm2.bias"] = up_block_one_sd_orig.pop("4.gn_2.bias") -up_block_one_sd_new["upsamplers.0.conv2.weight"] = up_block_one_sd_orig.pop("4.f_2.weight") -up_block_one_sd_new["upsamplers.0.conv2.bias"] = up_block_one_sd_orig.pop("4.f_2.bias") - -assert len(up_block_one_sd_orig) == 0 - -up_block_one = ResnetUpsampleBlock2D( - in_channels=1024, - prev_output_channel=1024, - out_channels=1024, - temb_channels=1280, - num_layers=4, - add_upsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -up_block_one.load_state_dict(up_block_one_sd_new) - -print("UP BLOCK TWO") - -up_block_two_sd_orig = model.up[-2].state_dict() -up_block_two_sd_new = {} - -for i in range(4): - up_block_two_sd_new[f"resnets.{i}.norm1.weight"] = up_block_two_sd_orig.pop(f"{i}.gn_1.weight") - up_block_two_sd_new[f"resnets.{i}.norm1.bias"] = up_block_two_sd_orig.pop(f"{i}.gn_1.bias") - up_block_two_sd_new[f"resnets.{i}.conv1.weight"] = up_block_two_sd_orig.pop(f"{i}.f_1.weight") - up_block_two_sd_new[f"resnets.{i}.conv1.bias"] = up_block_two_sd_orig.pop(f"{i}.f_1.bias") - up_block_two_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_two_sd_orig.pop(f"{i}.f_t.weight") - up_block_two_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_two_sd_orig.pop(f"{i}.f_t.bias") - up_block_two_sd_new[f"resnets.{i}.norm2.weight"] = up_block_two_sd_orig.pop(f"{i}.gn_2.weight") - up_block_two_sd_new[f"resnets.{i}.norm2.bias"] = up_block_two_sd_orig.pop(f"{i}.gn_2.bias") - up_block_two_sd_new[f"resnets.{i}.conv2.weight"] = up_block_two_sd_orig.pop(f"{i}.f_2.weight") - up_block_two_sd_new[f"resnets.{i}.conv2.bias"] = up_block_two_sd_orig.pop(f"{i}.f_2.bias") - up_block_two_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_two_sd_orig.pop(f"{i}.f_s.weight") - up_block_two_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_two_sd_orig.pop(f"{i}.f_s.bias") - -up_block_two_sd_new["upsamplers.0.norm1.weight"] = up_block_two_sd_orig.pop("4.gn_1.weight") -up_block_two_sd_new["upsamplers.0.norm1.bias"] = up_block_two_sd_orig.pop("4.gn_1.bias") -up_block_two_sd_new["upsamplers.0.conv1.weight"] = up_block_two_sd_orig.pop("4.f_1.weight") -up_block_two_sd_new["upsamplers.0.conv1.bias"] = up_block_two_sd_orig.pop("4.f_1.bias") -up_block_two_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_two_sd_orig.pop("4.f_t.weight") -up_block_two_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_two_sd_orig.pop("4.f_t.bias") -up_block_two_sd_new["upsamplers.0.norm2.weight"] = up_block_two_sd_orig.pop("4.gn_2.weight") -up_block_two_sd_new["upsamplers.0.norm2.bias"] = up_block_two_sd_orig.pop("4.gn_2.bias") -up_block_two_sd_new["upsamplers.0.conv2.weight"] = up_block_two_sd_orig.pop("4.f_2.weight") -up_block_two_sd_new["upsamplers.0.conv2.bias"] = up_block_two_sd_orig.pop("4.f_2.bias") - -assert len(up_block_two_sd_orig) == 0 - -up_block_two = ResnetUpsampleBlock2D( - in_channels=640, - prev_output_channel=1024, - out_channels=1024, - temb_channels=1280, - num_layers=4, - add_upsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -up_block_two.load_state_dict(up_block_two_sd_new) - -print("UP BLOCK THREE") - -up_block_three_sd_orig = model.up[-3].state_dict() -up_block_three_sd_new = {} - -for i in range(4): - up_block_three_sd_new[f"resnets.{i}.norm1.weight"] = up_block_three_sd_orig.pop(f"{i}.gn_1.weight") - up_block_three_sd_new[f"resnets.{i}.norm1.bias"] = up_block_three_sd_orig.pop(f"{i}.gn_1.bias") - up_block_three_sd_new[f"resnets.{i}.conv1.weight"] = up_block_three_sd_orig.pop(f"{i}.f_1.weight") - up_block_three_sd_new[f"resnets.{i}.conv1.bias"] = up_block_three_sd_orig.pop(f"{i}.f_1.bias") - up_block_three_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_three_sd_orig.pop(f"{i}.f_t.weight") - up_block_three_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_three_sd_orig.pop(f"{i}.f_t.bias") - up_block_three_sd_new[f"resnets.{i}.norm2.weight"] = up_block_three_sd_orig.pop(f"{i}.gn_2.weight") - up_block_three_sd_new[f"resnets.{i}.norm2.bias"] = up_block_three_sd_orig.pop(f"{i}.gn_2.bias") - up_block_three_sd_new[f"resnets.{i}.conv2.weight"] = up_block_three_sd_orig.pop(f"{i}.f_2.weight") - up_block_three_sd_new[f"resnets.{i}.conv2.bias"] = up_block_three_sd_orig.pop(f"{i}.f_2.bias") - up_block_three_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_three_sd_orig.pop(f"{i}.f_s.weight") - up_block_three_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_three_sd_orig.pop(f"{i}.f_s.bias") - -up_block_three_sd_new["upsamplers.0.norm1.weight"] = up_block_three_sd_orig.pop("4.gn_1.weight") -up_block_three_sd_new["upsamplers.0.norm1.bias"] = up_block_three_sd_orig.pop("4.gn_1.bias") -up_block_three_sd_new["upsamplers.0.conv1.weight"] = up_block_three_sd_orig.pop("4.f_1.weight") -up_block_three_sd_new["upsamplers.0.conv1.bias"] = up_block_three_sd_orig.pop("4.f_1.bias") -up_block_three_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_three_sd_orig.pop("4.f_t.weight") -up_block_three_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_three_sd_orig.pop("4.f_t.bias") -up_block_three_sd_new["upsamplers.0.norm2.weight"] = up_block_three_sd_orig.pop("4.gn_2.weight") -up_block_three_sd_new["upsamplers.0.norm2.bias"] = up_block_three_sd_orig.pop("4.gn_2.bias") -up_block_three_sd_new["upsamplers.0.conv2.weight"] = up_block_three_sd_orig.pop("4.f_2.weight") -up_block_three_sd_new["upsamplers.0.conv2.bias"] = up_block_three_sd_orig.pop("4.f_2.bias") - -assert len(up_block_three_sd_orig) == 0 - -up_block_three = ResnetUpsampleBlock2D( - in_channels=320, - prev_output_channel=1024, - out_channels=640, - temb_channels=1280, - num_layers=4, - add_upsample=True, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -up_block_three.load_state_dict(up_block_three_sd_new) - -print("UP BLOCK FOUR") - -up_block_four_sd_orig = model.up[-4].state_dict() -up_block_four_sd_new = {} - -for i in range(4): - up_block_four_sd_new[f"resnets.{i}.norm1.weight"] = up_block_four_sd_orig.pop(f"{i}.gn_1.weight") - up_block_four_sd_new[f"resnets.{i}.norm1.bias"] = up_block_four_sd_orig.pop(f"{i}.gn_1.bias") - up_block_four_sd_new[f"resnets.{i}.conv1.weight"] = up_block_four_sd_orig.pop(f"{i}.f_1.weight") - up_block_four_sd_new[f"resnets.{i}.conv1.bias"] = up_block_four_sd_orig.pop(f"{i}.f_1.bias") - up_block_four_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_four_sd_orig.pop(f"{i}.f_t.weight") - up_block_four_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_four_sd_orig.pop(f"{i}.f_t.bias") - up_block_four_sd_new[f"resnets.{i}.norm2.weight"] = up_block_four_sd_orig.pop(f"{i}.gn_2.weight") - up_block_four_sd_new[f"resnets.{i}.norm2.bias"] = up_block_four_sd_orig.pop(f"{i}.gn_2.bias") - up_block_four_sd_new[f"resnets.{i}.conv2.weight"] = up_block_four_sd_orig.pop(f"{i}.f_2.weight") - up_block_four_sd_new[f"resnets.{i}.conv2.bias"] = up_block_four_sd_orig.pop(f"{i}.f_2.bias") - up_block_four_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_four_sd_orig.pop(f"{i}.f_s.weight") - up_block_four_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_four_sd_orig.pop(f"{i}.f_s.bias") - -assert len(up_block_four_sd_orig) == 0 - -up_block_four = ResnetUpsampleBlock2D( - in_channels=320, - prev_output_channel=640, - out_channels=320, - temb_channels=1280, - num_layers=4, - add_upsample=False, - resnet_time_scale_shift="scale_shift", - resnet_eps=1e-5, -) - -up_block_four.load_state_dict(up_block_four_sd_new) - -print("initial projection (conv_in)") - -conv_in_sd_orig = model.embed_image.state_dict() -conv_in_sd_new = {} - -conv_in_sd_new["weight"] = conv_in_sd_orig.pop("f.weight") -conv_in_sd_new["bias"] = conv_in_sd_orig.pop("f.bias") - -assert len(conv_in_sd_orig) == 0 - -block_out_channels = [320, 640, 1024, 1024] - -in_channels = 7 -conv_in_kernel = 3 -conv_in_padding = (conv_in_kernel - 1) // 2 -conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding) - -conv_in.load_state_dict(conv_in_sd_new) - -print("out projection (conv_out) (conv_norm_out)") -out_channels = 6 -norm_num_groups = 32 -norm_eps = 1e-5 -act_fn = "silu" -conv_out_kernel = 3 -conv_out_padding = (conv_out_kernel - 1) // 2 -conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps) -# uses torch.functional in orig -# conv_act = get_activation(act_fn) -conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding) - -conv_norm_out.load_state_dict(model.output.gn.state_dict()) -conv_out.load_state_dict(model.output.f.state_dict()) - -print("timestep projection (time_proj) (time_embedding)") - -f1_sd = model.embed_time.f_1.state_dict() -f2_sd = model.embed_time.f_2.state_dict() - -time_embedding_sd = { - "linear_1.weight": f1_sd.pop("weight"), - "linear_1.bias": f1_sd.pop("bias"), - "linear_2.weight": f2_sd.pop("weight"), - "linear_2.bias": f2_sd.pop("bias"), -} - -assert len(f1_sd) == 0 -assert len(f2_sd) == 0 - -time_embedding_type = "learned" -num_train_timesteps = 1024 -time_embedding_dim = 1280 - -time_proj = nn.Embedding(num_train_timesteps, block_out_channels[0]) -timestep_input_dim = block_out_channels[0] - -time_embedding = TimestepEmbedding(timestep_input_dim, time_embedding_dim) - -time_proj.load_state_dict(model.embed_time.emb.state_dict()) -time_embedding.load_state_dict(time_embedding_sd) - -print("CONVERT") - -time_embedding.to("cuda") -time_proj.to("cuda") -conv_in.to("cuda") - -block_one.to("cuda") -block_two.to("cuda") -block_three.to("cuda") -block_four.to("cuda") - -mid_block_one.to("cuda") - -up_block_one.to("cuda") -up_block_two.to("cuda") -up_block_three.to("cuda") -up_block_four.to("cuda") - -conv_norm_out.to("cuda") -conv_out.to("cuda") - -model.time_proj = time_proj -model.time_embedding = time_embedding -model.embed_image = conv_in - -model.down[0] = block_one -model.down[1] = block_two -model.down[2] = block_three -model.down[3] = block_four - -model.mid = mid_block_one - -model.up[-1] = up_block_one -model.up[-2] = up_block_two -model.up[-3] = up_block_three -model.up[-4] = up_block_four - -model.output.gn = conv_norm_out -model.output.f = conv_out - -model.converted = True - -sample_consistency_new = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) -save_image(sample_consistency_new, "con_new.png") - -assert (sample_consistency_orig == sample_consistency_new).all() - -print("making unet") - -unet = UNet2DModel( - in_channels=in_channels, - out_channels=out_channels, - down_block_types=( - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - ), - up_block_types=( - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - ), - block_out_channels=block_out_channels, - layers_per_block=3, - norm_num_groups=norm_num_groups, - norm_eps=norm_eps, - resnet_time_scale_shift="scale_shift", - time_embedding_type="learned", - num_train_timesteps=num_train_timesteps, - add_attention=False, -) - -unet_state_dict = {} - - -def add_state_dict(prefix, mod): - for k, v in mod.state_dict().items(): - unet_state_dict[f"{prefix}.{k}"] = v - - -add_state_dict("conv_in", conv_in) -add_state_dict("time_proj", time_proj) -add_state_dict("time_embedding", time_embedding) -add_state_dict("down_blocks.0", block_one) -add_state_dict("down_blocks.1", block_two) -add_state_dict("down_blocks.2", block_three) -add_state_dict("down_blocks.3", block_four) -add_state_dict("mid_block", mid_block_one) -add_state_dict("up_blocks.0", up_block_one) -add_state_dict("up_blocks.1", up_block_two) -add_state_dict("up_blocks.2", up_block_three) -add_state_dict("up_blocks.3", up_block_four) -add_state_dict("conv_norm_out", conv_norm_out) -add_state_dict("conv_out", conv_out) - -unet.load_state_dict(unet_state_dict) - -print("running with diffusers unet") - -unet.to("cuda") - -decoder_consistency.ckpt = unet - -sample_consistency_new_2 = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) -save_image(sample_consistency_new_2, "con_new_2.png") - -assert (sample_consistency_orig == sample_consistency_new_2).all() - -print("running with diffusers model") - -Encoder.old_constructor = Encoder.__init__ - - -def new_constructor(self, **kwargs): - self.old_constructor(**kwargs) - self.constructor_arguments = kwargs - - -Encoder.__init__ = new_constructor - - -vae = AutoencoderKL.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="vae") -consistency_vae = ConsistencyDecoderVAE( - encoder_args=vae.encoder.constructor_arguments, - decoder_args=unet.config, - scaling_factor=vae.config.scaling_factor, - block_out_channels=vae.config.block_out_channels, - latent_channels=vae.config.latent_channels, -) -consistency_vae.encoder.load_state_dict(vae.encoder.state_dict()) -consistency_vae.quant_conv.load_state_dict(vae.quant_conv.state_dict()) -consistency_vae.decoder_unet.load_state_dict(unet.state_dict()) - -consistency_vae.to(dtype=torch.float16, device="cuda") - -sample_consistency_new_3 = consistency_vae.decode( - 0.18215 * latent, generator=torch.Generator("cpu").manual_seed(0) -).sample - -print("max difference") -print((sample_consistency_orig - sample_consistency_new_3).abs().max()) -print("total difference") -print((sample_consistency_orig - sample_consistency_new_3).abs().sum()) -# assert (sample_consistency_orig == sample_consistency_new_3).all() - -print("running with diffusers pipeline") - -pipe = DiffusionPipeline.from_pretrained( - "stable-diffusion-v1-5/stable-diffusion-v1-5", vae=consistency_vae, torch_dtype=torch.float16 -) -pipe.to("cuda") - -pipe("horse", generator=torch.Generator("cpu").manual_seed(0)).images[0].save("horse.png") - - -if args.save_pretrained is not None: - consistency_vae.save_pretrained(args.save_pretrained) diff --git a/scripts/convert_consistency_to_diffusers.py b/scripts/convert_consistency_to_diffusers.py deleted file mode 100644 index 2b918280ca05..000000000000 --- a/scripts/convert_consistency_to_diffusers.py +++ /dev/null @@ -1,315 +0,0 @@ -import argparse -import os - -import torch - -from diffusers import ( - CMStochasticIterativeScheduler, - ConsistencyModelPipeline, - UNet2DModel, -) - - -TEST_UNET_CONFIG = { - "sample_size": 32, - "in_channels": 3, - "out_channels": 3, - "layers_per_block": 2, - "num_class_embeds": 1000, - "block_out_channels": [32, 64], - "attention_head_dim": 8, - "down_block_types": [ - "ResnetDownsampleBlock2D", - "AttnDownBlock2D", - ], - "up_block_types": [ - "AttnUpBlock2D", - "ResnetUpsampleBlock2D", - ], - "resnet_time_scale_shift": "scale_shift", - "attn_norm_num_groups": 32, - "upsample_type": "resnet", - "downsample_type": "resnet", -} - -IMAGENET_64_UNET_CONFIG = { - "sample_size": 64, - "in_channels": 3, - "out_channels": 3, - "layers_per_block": 3, - "num_class_embeds": 1000, - "block_out_channels": [192, 192 * 2, 192 * 3, 192 * 4], - "attention_head_dim": 64, - "down_block_types": [ - "ResnetDownsampleBlock2D", - "AttnDownBlock2D", - "AttnDownBlock2D", - "AttnDownBlock2D", - ], - "up_block_types": [ - "AttnUpBlock2D", - "AttnUpBlock2D", - "AttnUpBlock2D", - "ResnetUpsampleBlock2D", - ], - "resnet_time_scale_shift": "scale_shift", - "attn_norm_num_groups": 32, - "upsample_type": "resnet", - "downsample_type": "resnet", -} - -LSUN_256_UNET_CONFIG = { - "sample_size": 256, - "in_channels": 3, - "out_channels": 3, - "layers_per_block": 2, - "num_class_embeds": None, - "block_out_channels": [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4], - "attention_head_dim": 64, - "down_block_types": [ - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "AttnDownBlock2D", - "AttnDownBlock2D", - "AttnDownBlock2D", - ], - "up_block_types": [ - "AttnUpBlock2D", - "AttnUpBlock2D", - "AttnUpBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - ], - "resnet_time_scale_shift": "default", - "upsample_type": "resnet", - "downsample_type": "resnet", -} - -CD_SCHEDULER_CONFIG = { - "num_train_timesteps": 40, - "sigma_min": 0.002, - "sigma_max": 80.0, -} - -CT_IMAGENET_64_SCHEDULER_CONFIG = { - "num_train_timesteps": 201, - "sigma_min": 0.002, - "sigma_max": 80.0, -} - -CT_LSUN_256_SCHEDULER_CONFIG = { - "num_train_timesteps": 151, - "sigma_min": 0.002, - "sigma_max": 80.0, -} - - -def str2bool(v): - """ - https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse - """ - if isinstance(v, bool): - return v - if v.lower() in ("yes", "true", "t", "y", "1"): - return True - elif v.lower() in ("no", "false", "f", "n", "0"): - return False - else: - raise argparse.ArgumentTypeError("boolean value expected") - - -def convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix, has_skip=False): - new_checkpoint[f"{new_prefix}.norm1.weight"] = checkpoint[f"{old_prefix}.in_layers.0.weight"] - new_checkpoint[f"{new_prefix}.norm1.bias"] = checkpoint[f"{old_prefix}.in_layers.0.bias"] - new_checkpoint[f"{new_prefix}.conv1.weight"] = checkpoint[f"{old_prefix}.in_layers.2.weight"] - new_checkpoint[f"{new_prefix}.conv1.bias"] = checkpoint[f"{old_prefix}.in_layers.2.bias"] - new_checkpoint[f"{new_prefix}.time_emb_proj.weight"] = checkpoint[f"{old_prefix}.emb_layers.1.weight"] - new_checkpoint[f"{new_prefix}.time_emb_proj.bias"] = checkpoint[f"{old_prefix}.emb_layers.1.bias"] - new_checkpoint[f"{new_prefix}.norm2.weight"] = checkpoint[f"{old_prefix}.out_layers.0.weight"] - new_checkpoint[f"{new_prefix}.norm2.bias"] = checkpoint[f"{old_prefix}.out_layers.0.bias"] - new_checkpoint[f"{new_prefix}.conv2.weight"] = checkpoint[f"{old_prefix}.out_layers.3.weight"] - new_checkpoint[f"{new_prefix}.conv2.bias"] = checkpoint[f"{old_prefix}.out_layers.3.bias"] - - if has_skip: - new_checkpoint[f"{new_prefix}.conv_shortcut.weight"] = checkpoint[f"{old_prefix}.skip_connection.weight"] - new_checkpoint[f"{new_prefix}.conv_shortcut.bias"] = checkpoint[f"{old_prefix}.skip_connection.bias"] - - return new_checkpoint - - -def convert_attention(checkpoint, new_checkpoint, old_prefix, new_prefix, attention_dim=None): - weight_q, weight_k, weight_v = checkpoint[f"{old_prefix}.qkv.weight"].chunk(3, dim=0) - bias_q, bias_k, bias_v = checkpoint[f"{old_prefix}.qkv.bias"].chunk(3, dim=0) - - new_checkpoint[f"{new_prefix}.group_norm.weight"] = checkpoint[f"{old_prefix}.norm.weight"] - new_checkpoint[f"{new_prefix}.group_norm.bias"] = checkpoint[f"{old_prefix}.norm.bias"] - - new_checkpoint[f"{new_prefix}.to_q.weight"] = weight_q.squeeze(-1).squeeze(-1) - new_checkpoint[f"{new_prefix}.to_q.bias"] = bias_q.squeeze(-1).squeeze(-1) - new_checkpoint[f"{new_prefix}.to_k.weight"] = weight_k.squeeze(-1).squeeze(-1) - new_checkpoint[f"{new_prefix}.to_k.bias"] = bias_k.squeeze(-1).squeeze(-1) - new_checkpoint[f"{new_prefix}.to_v.weight"] = weight_v.squeeze(-1).squeeze(-1) - new_checkpoint[f"{new_prefix}.to_v.bias"] = bias_v.squeeze(-1).squeeze(-1) - - new_checkpoint[f"{new_prefix}.to_out.0.weight"] = ( - checkpoint[f"{old_prefix}.proj_out.weight"].squeeze(-1).squeeze(-1) - ) - new_checkpoint[f"{new_prefix}.to_out.0.bias"] = checkpoint[f"{old_prefix}.proj_out.bias"].squeeze(-1).squeeze(-1) - - return new_checkpoint - - -def con_pt_to_diffuser(checkpoint_path: str, unet_config): - checkpoint = torch.load(checkpoint_path, map_location="cpu") - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = checkpoint["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = checkpoint["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = checkpoint["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = checkpoint["time_embed.2.bias"] - - if unet_config["num_class_embeds"] is not None: - new_checkpoint["class_embedding.weight"] = checkpoint["label_emb.weight"] - - new_checkpoint["conv_in.weight"] = checkpoint["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = checkpoint["input_blocks.0.0.bias"] - - down_block_types = unet_config["down_block_types"] - layers_per_block = unet_config["layers_per_block"] - attention_head_dim = unet_config["attention_head_dim"] - channels_list = unet_config["block_out_channels"] - current_layer = 1 - prev_channels = channels_list[0] - - for i, layer_type in enumerate(down_block_types): - current_channels = channels_list[i] - downsample_block_has_skip = current_channels != prev_channels - if layer_type == "ResnetDownsampleBlock2D": - for j in range(layers_per_block): - new_prefix = f"down_blocks.{i}.resnets.{j}" - old_prefix = f"input_blocks.{current_layer}.0" - has_skip = True if j == 0 and downsample_block_has_skip else False - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix, has_skip=has_skip) - current_layer += 1 - - elif layer_type == "AttnDownBlock2D": - for j in range(layers_per_block): - new_prefix = f"down_blocks.{i}.resnets.{j}" - old_prefix = f"input_blocks.{current_layer}.0" - has_skip = True if j == 0 and downsample_block_has_skip else False - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix, has_skip=has_skip) - new_prefix = f"down_blocks.{i}.attentions.{j}" - old_prefix = f"input_blocks.{current_layer}.1" - new_checkpoint = convert_attention( - checkpoint, new_checkpoint, old_prefix, new_prefix, attention_head_dim - ) - current_layer += 1 - - if i != len(down_block_types) - 1: - new_prefix = f"down_blocks.{i}.downsamplers.0" - old_prefix = f"input_blocks.{current_layer}.0" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix) - current_layer += 1 - - prev_channels = current_channels - - # hardcoded the mid-block for now - new_prefix = "mid_block.resnets.0" - old_prefix = "middle_block.0" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix) - new_prefix = "mid_block.attentions.0" - old_prefix = "middle_block.1" - new_checkpoint = convert_attention(checkpoint, new_checkpoint, old_prefix, new_prefix, attention_head_dim) - new_prefix = "mid_block.resnets.1" - old_prefix = "middle_block.2" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix) - - current_layer = 0 - up_block_types = unet_config["up_block_types"] - - for i, layer_type in enumerate(up_block_types): - if layer_type == "ResnetUpsampleBlock2D": - for j in range(layers_per_block + 1): - new_prefix = f"up_blocks.{i}.resnets.{j}" - old_prefix = f"output_blocks.{current_layer}.0" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix, has_skip=True) - current_layer += 1 - - if i != len(up_block_types) - 1: - new_prefix = f"up_blocks.{i}.upsamplers.0" - old_prefix = f"output_blocks.{current_layer - 1}.1" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix) - elif layer_type == "AttnUpBlock2D": - for j in range(layers_per_block + 1): - new_prefix = f"up_blocks.{i}.resnets.{j}" - old_prefix = f"output_blocks.{current_layer}.0" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix, has_skip=True) - new_prefix = f"up_blocks.{i}.attentions.{j}" - old_prefix = f"output_blocks.{current_layer}.1" - new_checkpoint = convert_attention( - checkpoint, new_checkpoint, old_prefix, new_prefix, attention_head_dim - ) - current_layer += 1 - - if i != len(up_block_types) - 1: - new_prefix = f"up_blocks.{i}.upsamplers.0" - old_prefix = f"output_blocks.{current_layer - 1}.2" - new_checkpoint = convert_resnet(checkpoint, new_checkpoint, old_prefix, new_prefix) - - new_checkpoint["conv_norm_out.weight"] = checkpoint["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = checkpoint["out.0.bias"] - new_checkpoint["conv_out.weight"] = checkpoint["out.2.weight"] - new_checkpoint["conv_out.bias"] = checkpoint["out.2.bias"] - - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--unet_path", default=None, type=str, required=True, help="Path to the unet.pt to convert.") - parser.add_argument( - "--dump_path", default=None, type=str, required=True, help="Path to output the converted UNet model." - ) - parser.add_argument("--class_cond", default=True, type=str, help="Whether the model is class-conditional.") - - args = parser.parse_args() - args.class_cond = str2bool(args.class_cond) - - ckpt_name = os.path.basename(args.unet_path) - print(f"Checkpoint: {ckpt_name}") - - # Get U-Net config - if "imagenet64" in ckpt_name: - unet_config = IMAGENET_64_UNET_CONFIG - elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)): - unet_config = LSUN_256_UNET_CONFIG - elif "test" in ckpt_name: - unet_config = TEST_UNET_CONFIG - else: - raise ValueError(f"Checkpoint type {ckpt_name} is not currently supported.") - - if not args.class_cond: - unet_config["num_class_embeds"] = None - - converted_unet_ckpt = con_pt_to_diffuser(args.unet_path, unet_config) - - image_unet = UNet2DModel(**unet_config) - image_unet.load_state_dict(converted_unet_ckpt) - - # Get scheduler config - if "cd" in ckpt_name or "test" in ckpt_name: - scheduler_config = CD_SCHEDULER_CONFIG - elif "ct" in ckpt_name and "imagenet64" in ckpt_name: - scheduler_config = CT_IMAGENET_64_SCHEDULER_CONFIG - elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)): - scheduler_config = CT_LSUN_256_SCHEDULER_CONFIG - else: - raise ValueError(f"Checkpoint type {ckpt_name} is not currently supported.") - - cm_scheduler = CMStochasticIterativeScheduler(**scheduler_config) - - consistency_model = ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler) - consistency_model.save_pretrained(args.dump_path) diff --git a/scripts/convert_cosmos_to_diffusers.py b/scripts/convert_cosmos_to_diffusers.py deleted file mode 100644 index b6f321b5b2fb..000000000000 --- a/scripts/convert_cosmos_to_diffusers.py +++ /dev/null @@ -1,976 +0,0 @@ -""" -# Cosmos 2 Predict - -Download checkpoint -```bash -hf download nvidia/Cosmos-Predict2-2B-Text2Image -``` - -convert checkpoint -```bash -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2-2B-Text2Image/snapshots/acdb5fde992a73ef0355f287977d002cbfd127e0/model.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_ckpt_path $transformer_ckpt_path \ - --transformer_type Cosmos-2.0-Diffusion-2B-Text2Image \ - --text_encoder_path google-t5/t5-11b \ - --tokenizer_path google-t5/t5-11b \ - --vae_type wan2.1 \ - --output_path converted/cosmos-p2-t2i-2b \ - --save_pipeline -``` - -# Cosmos 2.5 Predict - -Download checkpoint -```bash -hf download nvidia/Cosmos-Predict2.5-2B -``` - -Convert checkpoint -```bash -# pre-trained -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/pre-trained/d20b7120-df3e-4911-919d-db6e08bad31c_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Predict-Base-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/2b/d20b7120-df3e-4911-919d-db6e08bad31c \ - --save_pipeline - -# post-trained -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Predict-Base-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/2b/81edfebe-bd6a-4039-8c1d-737df1a790bf \ - --save_pipeline -``` - -## 14B - -```bash -hf download nvidia/Cosmos-Predict2.5-14B -``` - -```bash -# pre-trained -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/pre-trained/54937b8c-29de-4f04-862c-e67b04ec41e8_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Predict-Base-14B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/14b/54937b8c-29de-4f04-862c-e67b04ec41e8/ \ - --save_pipeline - -# post-trained -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Predict-Base-14B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/14b/e21d2a49-4747-44c8-ba44-9f6f9243715f/ \ - --save_pipeline -``` - -# Cosmos 2.5 Transfer - -Download checkpoint -```bash -hf download nvidia/Cosmos-Transfer2.5-2B -``` - -Convert checkpoint -```bash -# depth -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/depth/pipeline \ - --save_pipeline - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/depth/models - -# edge -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/edge/pipeline \ - --save_pipeline - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/edge/models - -# blur -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/blur/pipeline \ - --save_pipeline - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/blur/models - -# seg -transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.pt - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/seg/pipeline \ - --save_pipeline - -python scripts/convert_cosmos_to_diffusers.py \ - --transformer_type Cosmos-2.5-Transfer-General-2B \ - --transformer_ckpt_path $transformer_ckpt_path \ - --vae_type wan2.1 \ - --output_path converted/transfer/2b/general/seg/models -``` -""" - -import argparse -import pathlib -import sys -from typing import Any, Dict, Optional - -import torch -from accelerate import init_empty_weights -from huggingface_hub import snapshot_download -from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration, T5EncoderModel, T5TokenizerFast - -from diffusers import ( - AutoencoderKLCosmos, - AutoencoderKLWan, - Cosmos2TextToImagePipeline, - Cosmos2VideoToWorldPipeline, - CosmosControlNetModel, - CosmosTextToWorldPipeline, - CosmosTransformer3DModel, - CosmosVideoToWorldPipeline, - EDMEulerScheduler, - FlowMatchEulerDiscreteScheduler, - UniPCMultistepScheduler, -) -from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBasePipeline -from diffusers.pipelines.cosmos.pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline - - -def remove_keys_(key: str, state_dict: Dict[str, Any]): - state_dict.pop(key) - - -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def rename_transformer_blocks_(key: str, state_dict: Dict[str, Any]): - block_index = int(key.split(".")[1].removeprefix("block")) - new_key = key - - old_prefix = f"blocks.block{block_index}" - new_prefix = f"transformer_blocks.{block_index}" - new_key = new_prefix + new_key.removeprefix(old_prefix) - - state_dict[new_key] = state_dict.pop(key) - - -TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 = { - "t_embedder.1": "time_embed.t_embedder", - "affline_norm": "time_embed.norm", - ".blocks.0.block.attn": ".attn1", - ".blocks.1.block.attn": ".attn2", - ".blocks.2.block": ".ff", - ".blocks.0.adaLN_modulation.1": ".norm1.linear_1", - ".blocks.0.adaLN_modulation.2": ".norm1.linear_2", - ".blocks.1.adaLN_modulation.1": ".norm2.linear_1", - ".blocks.1.adaLN_modulation.2": ".norm2.linear_2", - ".blocks.2.adaLN_modulation.1": ".norm3.linear_1", - ".blocks.2.adaLN_modulation.2": ".norm3.linear_2", - "to_q.0": "to_q", - "to_q.1": "norm_q", - "to_k.0": "to_k", - "to_k.1": "norm_k", - "to_v.0": "to_v", - "layer1": "net.0.proj", - "layer2": "net.2", - "proj.1": "proj", - "x_embedder": "patch_embed", - "extra_pos_embedder": "learnable_pos_embed", - "final_layer.adaLN_modulation.1": "norm_out.linear_1", - "final_layer.adaLN_modulation.2": "norm_out.linear_2", - "final_layer.linear": "proj_out", -} - -TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 = { - "blocks.block": rename_transformer_blocks_, - "logvar.0.freqs": remove_keys_, - "logvar.0.phases": remove_keys_, - "logvar.1.weight": remove_keys_, - "pos_embedder.seq": remove_keys_, -} - -TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 = { - "t_embedder.1": "time_embed.t_embedder", - "t_embedding_norm": "time_embed.norm", - "blocks": "transformer_blocks", - "adaln_modulation_self_attn.1": "norm1.linear_1", - "adaln_modulation_self_attn.2": "norm1.linear_2", - "adaln_modulation_cross_attn.1": "norm2.linear_1", - "adaln_modulation_cross_attn.2": "norm2.linear_2", - "adaln_modulation_mlp.1": "norm3.linear_1", - "adaln_modulation_mlp.2": "norm3.linear_2", - "self_attn": "attn1", - "cross_attn": "attn2", - "q_proj": "to_q", - "k_proj": "to_k", - "v_proj": "to_v", - "output_proj": "to_out.0", - "q_norm": "norm_q", - "k_norm": "norm_k", - "mlp.layer1": "ff.net.0.proj", - "mlp.layer2": "ff.net.2", - "x_embedder.proj.1": "patch_embed.proj", - "final_layer.adaln_modulation.1": "norm_out.linear_1", - "final_layer.adaln_modulation.2": "norm_out.linear_2", - "final_layer.linear": "proj_out", -} - -TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 = { - "accum_video_sample_counter": remove_keys_, - "accum_image_sample_counter": remove_keys_, - "accum_iteration": remove_keys_, - "accum_train_in_hours": remove_keys_, - "pos_embedder.seq": remove_keys_, - "pos_embedder.dim_spatial_range": remove_keys_, - "pos_embedder.dim_temporal_range": remove_keys_, - "_extra_state": remove_keys_, -} - - -TRANSFORMER_CONFIGS = { - "Cosmos-1.0-Diffusion-7B-Text2World": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 32, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (2.0, 1.0, 1.0), - "concat_padding_mask": True, - "extra_pos_embed_type": "learnable", - }, - "Cosmos-1.0-Diffusion-7B-Video2World": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 32, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (2.0, 1.0, 1.0), - "concat_padding_mask": True, - "extra_pos_embed_type": "learnable", - }, - "Cosmos-1.0-Diffusion-14B-Text2World": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 40, - "attention_head_dim": 128, - "num_layers": 36, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (2.0, 2.0, 2.0), - "concat_padding_mask": True, - "extra_pos_embed_type": "learnable", - }, - "Cosmos-1.0-Diffusion-14B-Video2World": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 40, - "attention_head_dim": 128, - "num_layers": 36, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (2.0, 2.0, 2.0), - "concat_padding_mask": True, - "extra_pos_embed_type": "learnable", - }, - "Cosmos-2.0-Diffusion-2B-Text2Image": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 16, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 4.0, 4.0), - "concat_padding_mask": True, - "extra_pos_embed_type": None, - }, - "Cosmos-2.0-Diffusion-14B-Text2Image": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 40, - "attention_head_dim": 128, - "num_layers": 36, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 4.0, 4.0), - "concat_padding_mask": True, - "extra_pos_embed_type": None, - }, - "Cosmos-2.0-Diffusion-2B-Video2World": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 16, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 3.0, 3.0), - "concat_padding_mask": True, - "extra_pos_embed_type": None, - }, - "Cosmos-2.0-Diffusion-14B-Video2World": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 40, - "attention_head_dim": 128, - "num_layers": 36, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (20 / 24, 2.0, 2.0), - "concat_padding_mask": True, - "extra_pos_embed_type": None, - }, - "Cosmos-2.5-Predict-Base-2B": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 16, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 3.0, 3.0), - "concat_padding_mask": True, - # NOTE: source config has pos_emb_learnable: 'True' - but params are missing - "extra_pos_embed_type": None, - "use_crossattn_projection": True, - "crossattn_proj_in_channels": 100352, - "encoder_hidden_states_channels": 1024, - }, - "Cosmos-2.5-Predict-Base-14B": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 40, - "attention_head_dim": 128, - "num_layers": 36, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 3.0, 3.0), - "concat_padding_mask": True, - # NOTE: source config has pos_emb_learnable: 'True' - but params are missing - "extra_pos_embed_type": None, - "use_crossattn_projection": True, - "crossattn_proj_in_channels": 100352, - "encoder_hidden_states_channels": 1024, - }, - "Cosmos-2.5-Transfer-General-2B": { - "in_channels": 16 + 1, - "out_channels": 16, - "num_attention_heads": 16, - "attention_head_dim": 128, - "num_layers": 28, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "max_size": (128, 240, 240), - "patch_size": (1, 2, 2), - "rope_scale": (1.0, 3.0, 3.0), - "concat_padding_mask": True, - "extra_pos_embed_type": None, - "use_crossattn_projection": True, - "crossattn_proj_in_channels": 100352, - "encoder_hidden_states_channels": 1024, - "controlnet_block_every_n": 7, - "img_context_dim_in": 1152, - "img_context_dim_out": 2048, - "img_context_num_tokens": 256, - }, -} - -CONTROLNET_CONFIGS = { - "Cosmos-2.5-Transfer-General-2B": { - "n_controlnet_blocks": 4, - "model_channels": 2048, - "in_channels": 130, - "latent_channels": 18, # (16 latent + 1 condition_mask) + 1 padding_mask = 18 - "num_attention_heads": 16, - "attention_head_dim": 128, - "mlp_ratio": 4.0, - "text_embed_dim": 1024, - "adaln_lora_dim": 256, - "patch_size": (1, 2, 2), - "max_size": (128, 240, 240), - "rope_scale": (1.0, 3.0, 3.0), - "extra_pos_embed_type": None, - "img_context_dim_in": 1152, - "img_context_dim_out": 2048, - "use_crossattn_projection": True, - "crossattn_proj_in_channels": 100352, - "encoder_hidden_states_channels": 1024, - }, -} - -CONTROLNET_KEYS_RENAME_DICT = { - **TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0, - "blocks": "blocks", - "control_embedder.proj.1": "patch_embed.proj", -} - - -CONTROLNET_SPECIAL_KEYS_REMAP = {**TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0} - -VAE_KEYS_RENAME_DICT = { - "down.0": "down_blocks.0", - "down.1": "down_blocks.1", - "down.2": "down_blocks.2", - "up.0": "up_blocks.2", - "up.1": "up_blocks.1", - "up.2": "up_blocks.0", - ".block.": ".resnets.", - "downsample": "downsamplers.0", - "upsample": "upsamplers.0", - "mid.block_1": "mid_block.resnets.0", - "mid.attn_1.0": "mid_block.attentions.0", - "mid.attn_1.1": "mid_block.temp_attentions.0", - "mid.block_2": "mid_block.resnets.1", - ".q.conv3d": ".to_q", - ".k.conv3d": ".to_k", - ".v.conv3d": ".to_v", - ".proj_out.conv3d": ".to_out.0", - ".0.conv3d": ".conv_s", - ".1.conv3d": ".conv_t", - "conv1.conv3d": "conv1", - "conv2.conv3d": "conv2", - "conv3.conv3d": "conv3", - "nin_shortcut.conv3d": "conv_shortcut", - "quant_conv.conv3d": "quant_conv", - "post_quant_conv.conv3d": "post_quant_conv", -} - -VAE_SPECIAL_KEYS_REMAP = { - "wavelets": remove_keys_, - "_arange": remove_keys_, - "patch_size_buffer": remove_keys_, -} - -VAE_CONFIGS = { - "CV8x8x8-0.1": { - "name": "nvidia/Cosmos-0.1-Tokenizer-CV8x8x8", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 16, - "encoder_block_out_channels": (128, 256, 512, 512), - "decode_block_out_channels": (256, 512, 512, 512), - "attention_resolutions": (32,), - "resolution": 1024, - "num_layers": 2, - "patch_size": 4, - "patch_type": "haar", - "scaling_factor": 1.0, - "spatial_compression_ratio": 8, - "temporal_compression_ratio": 8, - "latents_mean": None, - "latents_std": None, - }, - }, - "CV8x8x8-1.0": { - "name": "nvidia/Cosmos-1.0-Tokenizer-CV8x8x8", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 16, - "encoder_block_out_channels": (128, 256, 512, 512), - "decode_block_out_channels": (256, 512, 512, 512), - "attention_resolutions": (32,), - "resolution": 1024, - "num_layers": 2, - "patch_size": 4, - "patch_type": "haar", - "scaling_factor": 1.0, - "spatial_compression_ratio": 8, - "temporal_compression_ratio": 8, - "latents_mean": None, - "latents_std": None, - }, - }, -} - - -def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: - state_dict = saved_dict - if "model" in saved_dict.keys(): - state_dict = state_dict["model"] - if "module" in saved_dict.keys(): - state_dict = state_dict["module"] - if "state_dict" in saved_dict.keys(): - state_dict = state_dict["state_dict"] - return state_dict - - -def convert_transformer( - transformer_type: str, - state_dict: Optional[Dict[str, Any]] = None, - weights_only: bool = True, -): - PREFIX_KEY = "net." - - if "Cosmos-1.0" in transformer_type: - TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 - TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 - elif "Cosmos-2.0" in transformer_type: - TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 - TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 - elif "Cosmos-2.5" in transformer_type: - TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 - TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 - else: - assert False - - with init_empty_weights(): - config = TRANSFORMER_CONFIGS[transformer_type] - transformer = CosmosTransformer3DModel(**config) - - old2new = {} - new2old = {} - for key in list(state_dict.keys()): - new_key = key[:] - if new_key.startswith(PREFIX_KEY): - new_key = new_key.removeprefix(PREFIX_KEY) - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - print(key, "->", new_key, flush=True) - assert new_key not in new2old, f"new key {new_key} already mapped" - assert key not in old2new, f"old key {key} already mapped" - old2new[key] = new_key - new2old[new_key] = key - update_state_dict_(state_dict, key, new_key) - - for key in list(state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, state_dict) - - expected_keys = set(transformer.state_dict().keys()) - mapped_keys = set(state_dict.keys()) - missing_keys = expected_keys - mapped_keys - unexpected_keys = mapped_keys - expected_keys - if missing_keys: - print(f"ERROR: missing keys ({len(missing_keys)} from state_dict:", flush=True, file=sys.stderr) - for k in missing_keys: - print(k) - sys.exit(1) - if unexpected_keys: - print(f"ERROR: unexpected keys ({len(unexpected_keys)}) from state_dict:", flush=True, file=sys.stderr) - for k in unexpected_keys: - print(k) - sys.exit(2) - - transformer.load_state_dict(state_dict, strict=True, assign=True) - return transformer - - -def convert_controlnet( - transformer_type: str, - control_state_dict: Dict[str, Any], - base_state_dict: Dict[str, Any], - weights_only: bool = True, -): - """ - Convert controlnet weights. - - Args: - transformer_type: The type of transformer/controlnet - control_state_dict: State dict containing controlnet-specific weights - base_state_dict: State dict containing base transformer weights (for shared modules) - weights_only: Whether to use weights_only loading - """ - if transformer_type not in CONTROLNET_CONFIGS: - raise AssertionError(f"{transformer_type} does not define a ControlNet config") - - PREFIX_KEY = "net." - - # Process control-specific keys - for key in list(control_state_dict.keys()): - new_key = key[:] - if new_key.startswith(PREFIX_KEY): - new_key = new_key.removeprefix(PREFIX_KEY) - for replace_key, rename_key in CONTROLNET_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(control_state_dict, key, new_key) - - for key in list(control_state_dict.keys()): - for special_key, handler_fn_inplace in CONTROLNET_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, control_state_dict) - - # Copy shared weights from base transformer to controlnet - # These are the duplicated modules: patch_embed_base, time_embed, learnable_pos_embed, img_context_proj, crossattn_proj - shared_module_mappings = { - # transformer key prefix -> controlnet key prefix - "patch_embed.": "patch_embed_base.", - "time_embed.": "time_embed.", - "learnable_pos_embed.": "learnable_pos_embed.", - "img_context_proj.": "img_context_proj.", - "crossattn_proj.": "crossattn_proj.", - } - - for key in list(base_state_dict.keys()): - for transformer_prefix, controlnet_prefix in shared_module_mappings.items(): - if key.startswith(transformer_prefix): - controlnet_key = controlnet_prefix + key[len(transformer_prefix) :] - control_state_dict[controlnet_key] = base_state_dict[key].clone() - print(f"Copied shared weight: {key} -> {controlnet_key}", flush=True) - break - - cfg = CONTROLNET_CONFIGS[transformer_type] - controlnet = CosmosControlNetModel(**cfg) - - expected_keys = set(controlnet.state_dict().keys()) - mapped_keys = set(control_state_dict.keys()) - missing_keys = expected_keys - mapped_keys - unexpected_keys = mapped_keys - expected_keys - if missing_keys: - print(f"WARNING: missing controlnet keys ({len(missing_keys)}):", file=sys.stderr, flush=True) - for k in sorted(missing_keys): - print(k, file=sys.stderr) - sys.exit(3) - if unexpected_keys: - print(f"WARNING: unexpected controlnet keys ({len(unexpected_keys)}):", file=sys.stderr, flush=True) - for k in sorted(unexpected_keys): - print(k, file=sys.stderr) - sys.exit(4) - - controlnet.load_state_dict(control_state_dict, strict=True, assign=True) - return controlnet - - -def convert_vae(vae_type: str): - model_name = VAE_CONFIGS[vae_type]["name"] - snapshot_directory = snapshot_download(model_name, repo_type="model") - directory = pathlib.Path(snapshot_directory) - - autoencoder_file = directory / "autoencoder.jit" - mean_std_file = directory / "mean_std.pt" - - original_state_dict = torch.jit.load(autoencoder_file.as_posix()).state_dict() - if mean_std_file.exists(): - mean_std = torch.load(mean_std_file, map_location="cpu", weights_only=True) - else: - mean_std = (None, None) - - config = VAE_CONFIGS[vae_type]["diffusers_config"] - config.update( - { - "latents_mean": mean_std[0].detach().cpu().numpy().tolist(), - "latents_std": mean_std[1].detach().cpu().numpy().tolist(), - } - ) - vae = AutoencoderKLCosmos(**config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def save_pipeline_cosmos_1_0(args, transformer, vae): - text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16) - tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path) - # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly. - # So, the sigma_min values that is used is the default value of 0.002. - scheduler = EDMEulerScheduler( - sigma_min=0.002, - sigma_max=80, - sigma_data=0.5, - sigma_schedule="karras", - num_train_timesteps=1000, - prediction_type="epsilon", - rho=7.0, - final_sigmas_type="sigma_min", - ) - - pipe_cls = CosmosTextToWorldPipeline if "Text2World" in args.transformer_type else CosmosVideoToWorldPipeline - pipe = pipe_cls( - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - vae=vae, - scheduler=scheduler, - safety_checker=lambda *args, **kwargs: None, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - -def save_pipeline_cosmos_2_0(args, transformer, vae): - text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16) - tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path) - - scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) - - pipe_cls = Cosmos2TextToImagePipeline if "Text2Image" in args.transformer_type else Cosmos2VideoToWorldPipeline - pipe = pipe_cls( - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - vae=vae, - scheduler=scheduler, - safety_checker=lambda *args, **kwargs: None, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - -def save_pipeline_cosmos2_5_predict(args, transformer, vae): - text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" - tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" - - text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - text_encoder_path, torch_dtype="auto", device_map="cpu" - ) - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - scheduler = UniPCMultistepScheduler( - use_karras_sigmas=True, - use_flow_sigmas=True, - prediction_type="flow_prediction", - sigma_max=200.0, - sigma_min=0.01, - ) - - pipe = Cosmos2_5_PredictBasePipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - vae=vae, - scheduler=scheduler, - safety_checker=lambda *args, **kwargs: None, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - -def save_pipeline_cosmos2_5_transfer(args, transformer, controlnet, vae): - text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" - tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" - - text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( - text_encoder_path, torch_dtype="auto", device_map="cpu" - ) - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - - scheduler = UniPCMultistepScheduler( - use_karras_sigmas=True, - use_flow_sigmas=True, - prediction_type="flow_prediction", - sigma_max=200.0, - sigma_min=0.01, - ) - - pipe = Cosmos2_5_TransferPipeline( - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - controlnet=controlnet, - vae=vae, - scheduler=scheduler, - safety_checker=lambda *args, **kwargs: None, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys())) - parser.add_argument( - "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" - ) - parser.add_argument( - "--vae_type", type=str, default="wan2.1", choices=["wan2.1", *list(VAE_CONFIGS.keys())], help="Type of VAE" - ) - parser.add_argument("--text_encoder_path", type=str, default=None) - parser.add_argument("--tokenizer_path", type=str, default=None) - parser.add_argument("--save_pipeline", action="store_true") - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") - parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.") - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -if __name__ == "__main__": - args = get_args() - - transformer = None - controlnet = None - dtype = DTYPE_MAPPING[args.dtype] - - if args.save_pipeline: - assert args.transformer_ckpt_path is not None - assert args.vae_type is not None - - raw_state_dict = None - if args.transformer_ckpt_path is not None: - weights_only = "Cosmos-1.0" in args.transformer_type - raw_state_dict = get_state_dict( - torch.load(args.transformer_ckpt_path, map_location="cpu", weights_only=weights_only) - ) - - if raw_state_dict is not None: - if "Transfer" in args.transformer_type: - base_state_dict = {} - control_state_dict = {} - for k, v in raw_state_dict.items(): - plain_key = k.removeprefix("net.") if k.startswith("net.") else k - if "control" in plain_key.lower(): - control_state_dict[k] = v - else: - base_state_dict[k] = v - assert len(base_state_dict.keys() & control_state_dict.keys()) == 0 - - # Convert transformer first to get the processed base state dict - transformer = convert_transformer( - args.transformer_type, state_dict=base_state_dict, weights_only=weights_only - ) - transformer = transformer.to(dtype=dtype) - - # Get converted transformer state dict to copy shared weights to controlnet - converted_base_state_dict = transformer.state_dict() - - # Convert controlnet with both control-specific and shared weights from transformer - controlnet = convert_controlnet( - args.transformer_type, control_state_dict, converted_base_state_dict, weights_only=weights_only - ) - controlnet = controlnet.to(dtype=dtype) - - if not args.save_pipeline: - transformer.save_pretrained( - pathlib.Path(args.output_path) / "transformer", safe_serialization=True, max_shard_size="5GB" - ) - controlnet.save_pretrained( - pathlib.Path(args.output_path) / "controlnet", safe_serialization=True, max_shard_size="5GB" - ) - else: - transformer = convert_transformer( - args.transformer_type, state_dict=raw_state_dict, weights_only=weights_only - ) - transformer = transformer.to(dtype=dtype) - if not args.save_pipeline: - transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.vae_type is not None: - if "Cosmos-1.0" in args.transformer_type: - vae = convert_vae(args.vae_type) - elif "Cosmos-2.0" in args.transformer_type or "Cosmos-2.5" in args.transformer_type: - vae = AutoencoderKLWan.from_pretrained( - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32 - ) - else: - raise AssertionError(f"{args.transformer_type} not supported") - - if not args.save_pipeline: - vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - else: - vae = None - - if args.save_pipeline: - if "Cosmos-1.0" in args.transformer_type: - assert args.text_encoder_path is not None - assert args.tokenizer_path is not None - save_pipeline_cosmos_1_0(args, transformer, vae) - elif "Cosmos-2.0" in args.transformer_type: - assert args.text_encoder_path is not None - assert args.tokenizer_path is not None - save_pipeline_cosmos_2_0(args, transformer, vae) - elif "Cosmos-2.5" in args.transformer_type: - if "Predict" in args.transformer_type: - save_pipeline_cosmos2_5_predict(args, transformer, vae) - elif "Transfer" in args.transformer_type: - save_pipeline_cosmos2_5_transfer(args, transformer, None, vae) - else: - raise AssertionError(f"{args.transformer_type} not supported") - else: - raise AssertionError(f"{args.transformer_type} not supported") diff --git a/scripts/convert_dance_diffusion_to_diffusers.py b/scripts/convert_dance_diffusion_to_diffusers.py deleted file mode 100755 index e269a49070cc..000000000000 --- a/scripts/convert_dance_diffusion_to_diffusers.py +++ /dev/null @@ -1,346 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import math -import os -from copy import deepcopy - -import requests -import torch -from audio_diffusion.models import DiffusionAttnUnet1D -from diffusion import sampling -from torch import nn - -from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNet1DModel -from diffusers.utils.constants import DIFFUSERS_REQUEST_TIMEOUT - - -MODELS_MAP = { - "gwf-440k": { - "url": "https://model-server.zqevans2.workers.dev/gwf-440k.ckpt", - "sample_rate": 48000, - "sample_size": 65536, - }, - "jmann-small-190k": { - "url": "https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt", - "sample_rate": 48000, - "sample_size": 65536, - }, - "jmann-large-580k": { - "url": "https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt", - "sample_rate": 48000, - "sample_size": 131072, - }, - "maestro-uncond-150k": { - "url": "https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt", - "sample_rate": 16000, - "sample_size": 65536, - }, - "unlocked-uncond-250k": { - "url": "https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt", - "sample_rate": 16000, - "sample_size": 65536, - }, - "honk-140k": { - "url": "https://model-server.zqevans2.workers.dev/honk-140k.ckpt", - "sample_rate": 16000, - "sample_size": 65536, - }, -} - - -def alpha_sigma_to_t(alpha, sigma): - """Returns a timestep, given the scaling factors for the clean image and for - the noise.""" - return torch.atan2(sigma, alpha) / math.pi * 2 - - -def get_crash_schedule(t): - sigma = torch.sin(t * math.pi / 2) ** 2 - alpha = (1 - sigma**2) ** 0.5 - return alpha_sigma_to_t(alpha, sigma) - - -class Object(object): - pass - - -class DiffusionUncond(nn.Module): - def __init__(self, global_args): - super().__init__() - - self.diffusion = DiffusionAttnUnet1D(global_args, n_attn_layers=4) - self.diffusion_ema = deepcopy(self.diffusion) - self.rng = torch.quasirandom.SobolEngine(1, scramble=True) - - -def download(model_name): - url = MODELS_MAP[model_name]["url"] - r = requests.get(url, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT) - - local_filename = f"./{model_name}.ckpt" - with open(local_filename, "wb") as fp: - for chunk in r.iter_content(chunk_size=8192): - fp.write(chunk) - - return local_filename - - -DOWN_NUM_TO_LAYER = { - "1": "resnets.0", - "2": "attentions.0", - "3": "resnets.1", - "4": "attentions.1", - "5": "resnets.2", - "6": "attentions.2", -} -UP_NUM_TO_LAYER = { - "8": "resnets.0", - "9": "attentions.0", - "10": "resnets.1", - "11": "attentions.1", - "12": "resnets.2", - "13": "attentions.2", -} -MID_NUM_TO_LAYER = { - "1": "resnets.0", - "2": "attentions.0", - "3": "resnets.1", - "4": "attentions.1", - "5": "resnets.2", - "6": "attentions.2", - "8": "resnets.3", - "9": "attentions.3", - "10": "resnets.4", - "11": "attentions.4", - "12": "resnets.5", - "13": "attentions.5", -} -DEPTH_0_TO_LAYER = { - "0": "resnets.0", - "1": "resnets.1", - "2": "resnets.2", - "4": "resnets.0", - "5": "resnets.1", - "6": "resnets.2", -} - -RES_CONV_MAP = { - "skip": "conv_skip", - "main.0": "conv_1", - "main.1": "group_norm_1", - "main.3": "conv_2", - "main.4": "group_norm_2", -} - -ATTN_MAP = { - "norm": "group_norm", - "qkv_proj": ["query", "key", "value"], - "out_proj": ["proj_attn"], -} - - -def convert_resconv_naming(name): - if name.startswith("skip"): - return name.replace("skip", RES_CONV_MAP["skip"]) - - # name has to be of format main.{digit} - if not name.startswith("main."): - raise ValueError(f"ResConvBlock error with {name}") - - return name.replace(name[:6], RES_CONV_MAP[name[:6]]) - - -def convert_attn_naming(name): - for key, value in ATTN_MAP.items(): - if name.startswith(key) and not isinstance(value, list): - return name.replace(key, value) - elif name.startswith(key): - return [name.replace(key, v) for v in value] - raise ValueError(f"Attn error with {name}") - - -def rename(input_string, max_depth=13): - string = input_string - - if string.split(".")[0] == "timestep_embed": - return string.replace("timestep_embed", "time_proj") - - depth = 0 - if string.startswith("net.3."): - depth += 1 - string = string[6:] - elif string.startswith("net."): - string = string[4:] - - while string.startswith("main.7."): - depth += 1 - string = string[7:] - - if string.startswith("main."): - string = string[5:] - - # mid block - if string[:2].isdigit(): - layer_num = string[:2] - string_left = string[2:] - else: - layer_num = string[0] - string_left = string[1:] - - if depth == max_depth: - new_layer = MID_NUM_TO_LAYER[layer_num] - prefix = "mid_block" - elif depth > 0 and int(layer_num) < 7: - new_layer = DOWN_NUM_TO_LAYER[layer_num] - prefix = f"down_blocks.{depth}" - elif depth > 0 and int(layer_num) > 7: - new_layer = UP_NUM_TO_LAYER[layer_num] - prefix = f"up_blocks.{max_depth - depth - 1}" - elif depth == 0: - new_layer = DEPTH_0_TO_LAYER[layer_num] - prefix = f"up_blocks.{max_depth - 1}" if int(layer_num) > 3 else "down_blocks.0" - - if not string_left.startswith("."): - raise ValueError(f"Naming error with {input_string} and string_left: {string_left}.") - - string_left = string_left[1:] - - if "resnets" in new_layer: - string_left = convert_resconv_naming(string_left) - elif "attentions" in new_layer: - new_string_left = convert_attn_naming(string_left) - string_left = new_string_left - - if not isinstance(string_left, list): - new_string = prefix + "." + new_layer + "." + string_left - else: - new_string = [prefix + "." + new_layer + "." + s for s in string_left] - return new_string - - -def rename_orig_weights(state_dict): - new_state_dict = {} - for k, v in state_dict.items(): - if k.endswith("kernel"): - # up- and downsample layers, don't have trainable weights - continue - - new_k = rename(k) - - # check if we need to transform from Conv => Linear for attention - if isinstance(new_k, list): - new_state_dict = transform_conv_attns(new_state_dict, new_k, v) - else: - new_state_dict[new_k] = v - - return new_state_dict - - -def transform_conv_attns(new_state_dict, new_k, v): - if len(new_k) == 1: - if len(v.shape) == 3: - # weight - new_state_dict[new_k[0]] = v[:, :, 0] - else: - # bias - new_state_dict[new_k[0]] = v - else: - # qkv matrices - trippled_shape = v.shape[0] - single_shape = trippled_shape // 3 - for i in range(3): - if len(v.shape) == 3: - new_state_dict[new_k[i]] = v[i * single_shape : (i + 1) * single_shape, :, 0] - else: - new_state_dict[new_k[i]] = v[i * single_shape : (i + 1) * single_shape] - return new_state_dict - - -def main(args): - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - model_name = args.model_path.split("/")[-1].split(".")[0] - if not os.path.isfile(args.model_path): - assert model_name == args.model_path, ( - f"Make sure to provide one of the official model names {MODELS_MAP.keys()}" - ) - args.model_path = download(model_name) - - sample_rate = MODELS_MAP[model_name]["sample_rate"] - sample_size = MODELS_MAP[model_name]["sample_size"] - - config = Object() - config.sample_size = sample_size - config.sample_rate = sample_rate - config.latent_dim = 0 - - diffusers_model = UNet1DModel(sample_size=sample_size, sample_rate=sample_rate) - diffusers_state_dict = diffusers_model.state_dict() - - orig_model = DiffusionUncond(config) - orig_model.load_state_dict(torch.load(args.model_path, map_location=device)["state_dict"]) - orig_model = orig_model.diffusion_ema.eval() - orig_model_state_dict = orig_model.state_dict() - renamed_state_dict = rename_orig_weights(orig_model_state_dict) - - renamed_minus_diffusers = set(renamed_state_dict.keys()) - set(diffusers_state_dict.keys()) - diffusers_minus_renamed = set(diffusers_state_dict.keys()) - set(renamed_state_dict.keys()) - - assert len(renamed_minus_diffusers) == 0, f"Problem with {renamed_minus_diffusers}" - assert all(k.endswith("kernel") for k in list(diffusers_minus_renamed)), f"Problem with {diffusers_minus_renamed}" - - for key, value in renamed_state_dict.items(): - assert diffusers_state_dict[key].squeeze().shape == value.squeeze().shape, ( - f"Shape for {key} doesn't match. Diffusers: {diffusers_state_dict[key].shape} vs. {value.shape}" - ) - if key == "time_proj.weight": - value = value.squeeze() - - diffusers_state_dict[key] = value - - diffusers_model.load_state_dict(diffusers_state_dict) - - steps = 100 - seed = 33 - - diffusers_scheduler = IPNDMScheduler(num_train_timesteps=steps) - - generator = torch.manual_seed(seed) - noise = torch.randn([1, 2, config.sample_size], generator=generator).to(device) - - t = torch.linspace(1, 0, steps + 1, device=device)[:-1] - step_list = get_crash_schedule(t) - - pipe = DanceDiffusionPipeline(unet=diffusers_model, scheduler=diffusers_scheduler) - - generator = torch.manual_seed(33) - audio = pipe(num_inference_steps=steps, generator=generator).audios - - generated = sampling.iplms_sample(orig_model, noise, step_list, {}) - generated = generated.clamp(-1, 1) - - diff_sum = (generated - audio).abs().sum() - diff_max = (generated - audio).abs().max() - - if args.save: - pipe.save_pretrained(args.checkpoint_path) - - print("Diff sum", diff_sum) - print("Diff max", diff_max) - - assert diff_max < 1e-3, f"Diff max: {diff_max} is too much :-/" - - print(f"Conversion for {model_name} successful!") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.") - parser.add_argument( - "--save", default=True, type=bool, required=False, help="Whether to save the converted model or not." - ) - parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") - args = parser.parse_args() - - main(args) diff --git a/scripts/convert_dcae_to_diffusers.py b/scripts/convert_dcae_to_diffusers.py deleted file mode 100644 index 1a94c4b06640..000000000000 --- a/scripts/convert_dcae_to_diffusers.py +++ /dev/null @@ -1,323 +0,0 @@ -import argparse -from typing import Any, Dict - -import torch -from huggingface_hub import hf_hub_download -from safetensors.torch import load_file - -from diffusers import AutoencoderDC - - -def remap_qkv_(key: str, state_dict: Dict[str, Any]): - qkv = state_dict.pop(key) - q, k, v = torch.chunk(qkv, 3, dim=0) - parent_module, _, _ = key.rpartition(".qkv.conv.weight") - state_dict[f"{parent_module}.to_q.weight"] = q.squeeze() - state_dict[f"{parent_module}.to_k.weight"] = k.squeeze() - state_dict[f"{parent_module}.to_v.weight"] = v.squeeze() - - -def remap_proj_conv_(key: str, state_dict: Dict[str, Any]): - parent_module, _, _ = key.rpartition(".proj.conv.weight") - state_dict[f"{parent_module}.to_out.weight"] = state_dict.pop(key).squeeze() - - -AE_KEYS_RENAME_DICT = { - # common - "main.": "", - "op_list.": "", - "context_module": "attn", - "local_module": "conv_out", - # NOTE: The below two lines work because scales in the available configs only have a tuple length of 1 - # If there were more scales, there would be more layers, so a loop would be better to handle this - "aggreg.0.0": "to_qkv_multiscale.0.proj_in", - "aggreg.0.1": "to_qkv_multiscale.0.proj_out", - "depth_conv.conv": "conv_depth", - "inverted_conv.conv": "conv_inverted", - "point_conv.conv": "conv_point", - "point_conv.norm": "norm", - "conv.conv.": "conv.", - "conv1.conv": "conv1", - "conv2.conv": "conv2", - "conv2.norm": "norm", - "proj.norm": "norm_out", - # encoder - "encoder.project_in.conv": "encoder.conv_in", - "encoder.project_out.0.conv": "encoder.conv_out", - "encoder.stages": "encoder.down_blocks", - # decoder - "decoder.project_in.conv": "decoder.conv_in", - "decoder.project_out.0": "decoder.norm_out", - "decoder.project_out.2.conv": "decoder.conv_out", - "decoder.stages": "decoder.up_blocks", -} - -AE_F32C32_KEYS = { - # encoder - "encoder.project_in.conv": "encoder.conv_in.conv", - # decoder - "decoder.project_out.2.conv": "decoder.conv_out.conv", -} - -AE_F64C128_KEYS = { - # encoder - "encoder.project_in.conv": "encoder.conv_in.conv", - # decoder - "decoder.project_out.2.conv": "decoder.conv_out.conv", -} - -AE_F128C512_KEYS = { - # encoder - "encoder.project_in.conv": "encoder.conv_in.conv", - # decoder - "decoder.project_out.2.conv": "decoder.conv_out.conv", -} - -AE_SPECIAL_KEYS_REMAP = { - "qkv.conv.weight": remap_qkv_, - "proj.conv.weight": remap_proj_conv_, -} - - -def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: - state_dict = saved_dict - if "model" in saved_dict.keys(): - state_dict = state_dict["model"] - if "module" in saved_dict.keys(): - state_dict = state_dict["module"] - if "state_dict" in saved_dict.keys(): - state_dict = state_dict["state_dict"] - return state_dict - - -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def convert_ae(config_name: str, dtype: torch.dtype): - config = get_ae_config(config_name) - hub_id = f"mit-han-lab/{config_name}" - ckpt_path = hf_hub_download(hub_id, "model.safetensors") - original_state_dict = get_state_dict(load_file(ckpt_path)) - - ae = AutoencoderDC(**config).to(dtype=dtype) - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in AE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in AE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - ae.load_state_dict(original_state_dict, strict=True) - return ae - - -def get_ae_config(name: str): - if name in ["dc-ae-f32c32-sana-1.0"]: - config = { - "latent_channels": 32, - "encoder_block_types": ( - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ), - "decoder_block_types": ( - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ), - "encoder_block_out_channels": (128, 256, 512, 512, 1024, 1024), - "decoder_block_out_channels": (128, 256, 512, 512, 1024, 1024), - "encoder_qkv_multiscales": ((), (), (), (5,), (5,), (5,)), - "decoder_qkv_multiscales": ((), (), (), (5,), (5,), (5,)), - "encoder_layers_per_block": (2, 2, 2, 3, 3, 3), - "decoder_layers_per_block": [3, 3, 3, 3, 3, 3], - "downsample_block_type": "conv", - "upsample_block_type": "interpolate", - "decoder_norm_types": "rms_norm", - "decoder_act_fns": "silu", - "scaling_factor": 0.41407, - } - elif name in ["dc-ae-f32c32-in-1.0", "dc-ae-f32c32-mix-1.0"]: - AE_KEYS_RENAME_DICT.update(AE_F32C32_KEYS) - config = { - "latent_channels": 32, - "encoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "decoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "encoder_block_out_channels": [128, 256, 512, 512, 1024, 1024], - "decoder_block_out_channels": [128, 256, 512, 512, 1024, 1024], - "encoder_layers_per_block": [0, 4, 8, 2, 2, 2], - "decoder_layers_per_block": [0, 5, 10, 2, 2, 2], - "encoder_qkv_multiscales": ((), (), (), (), (), ()), - "decoder_qkv_multiscales": ((), (), (), (), (), ()), - "decoder_norm_types": ["batch_norm", "batch_norm", "batch_norm", "rms_norm", "rms_norm", "rms_norm"], - "decoder_act_fns": ["relu", "relu", "relu", "silu", "silu", "silu"], - } - if name == "dc-ae-f32c32-in-1.0": - config["scaling_factor"] = 0.3189 - elif name == "dc-ae-f32c32-mix-1.0": - config["scaling_factor"] = 0.4552 - elif name in ["dc-ae-f64c128-in-1.0", "dc-ae-f64c128-mix-1.0"]: - AE_KEYS_RENAME_DICT.update(AE_F64C128_KEYS) - config = { - "latent_channels": 128, - "encoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "decoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "encoder_block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048], - "decoder_block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048], - "encoder_layers_per_block": [0, 4, 8, 2, 2, 2, 2], - "decoder_layers_per_block": [0, 5, 10, 2, 2, 2, 2], - "encoder_qkv_multiscales": ((), (), (), (), (), (), ()), - "decoder_qkv_multiscales": ((), (), (), (), (), (), ()), - "decoder_norm_types": [ - "batch_norm", - "batch_norm", - "batch_norm", - "rms_norm", - "rms_norm", - "rms_norm", - "rms_norm", - ], - "decoder_act_fns": ["relu", "relu", "relu", "silu", "silu", "silu", "silu"], - } - if name == "dc-ae-f64c128-in-1.0": - config["scaling_factor"] = 0.2889 - elif name == "dc-ae-f64c128-mix-1.0": - config["scaling_factor"] = 0.4538 - elif name in ["dc-ae-f128c512-in-1.0", "dc-ae-f128c512-mix-1.0"]: - AE_KEYS_RENAME_DICT.update(AE_F128C512_KEYS) - config = { - "latent_channels": 512, - "encoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "decoder_block_types": [ - "ResBlock", - "ResBlock", - "ResBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - "EfficientViTBlock", - ], - "encoder_block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048, 2048], - "decoder_block_out_channels": [128, 256, 512, 512, 1024, 1024, 2048, 2048], - "encoder_layers_per_block": [0, 4, 8, 2, 2, 2, 2, 2], - "decoder_layers_per_block": [0, 5, 10, 2, 2, 2, 2, 2], - "encoder_qkv_multiscales": ((), (), (), (), (), (), (), ()), - "decoder_qkv_multiscales": ((), (), (), (), (), (), (), ()), - "decoder_norm_types": [ - "batch_norm", - "batch_norm", - "batch_norm", - "rms_norm", - "rms_norm", - "rms_norm", - "rms_norm", - "rms_norm", - ], - "decoder_act_fns": ["relu", "relu", "relu", "silu", "silu", "silu", "silu", "silu"], - } - if name == "dc-ae-f128c512-in-1.0": - config["scaling_factor"] = 0.4883 - elif name == "dc-ae-f128c512-mix-1.0": - config["scaling_factor"] = 0.3620 - else: - raise ValueError("Invalid config name provided.") - - return config - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--config_name", - type=str, - default="dc-ae-f32c32-sana-1.0", - choices=[ - "dc-ae-f32c32-sana-1.0", - "dc-ae-f32c32-in-1.0", - "dc-ae-f32c32-mix-1.0", - "dc-ae-f64c128-in-1.0", - "dc-ae-f64c128-mix-1.0", - "dc-ae-f128c512-in-1.0", - "dc-ae-f128c512-mix-1.0", - ], - help="The DCAE checkpoint to convert", - ) - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") - parser.add_argument("--dtype", default="fp32", help="Torch dtype to save the model in.") - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - -VARIANT_MAPPING = { - "fp32": None, - "fp16": "fp16", - "bf16": "bf16", -} - - -if __name__ == "__main__": - args = get_args() - - dtype = DTYPE_MAPPING[args.dtype] - variant = VARIANT_MAPPING[args.dtype] - - ae = convert_ae(args.config_name, dtype) - ae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", variant=variant) diff --git a/scripts/convert_ddpm_original_checkpoint_to_diffusers.py b/scripts/convert_ddpm_original_checkpoint_to_diffusers.py deleted file mode 100644 index 46595784b0ba..000000000000 --- a/scripts/convert_ddpm_original_checkpoint_to_diffusers.py +++ /dev/null @@ -1,431 +0,0 @@ -import argparse -import json - -import torch - -from diffusers import AutoencoderKL, DDPMPipeline, DDPMScheduler, UNet2DModel, VQModel - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - mapping = [] - for old_item in old_list: - new_item = old_item - new_item = new_item.replace("block.", "resnets.") - new_item = new_item.replace("conv_shorcut", "conv1") - new_item = new_item.replace("in_shortcut", "conv_shortcut") - new_item = new_item.replace("temb_proj", "time_emb_proj") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0, in_mid=False): - mapping = [] - for old_item in old_list: - new_item = old_item - - # In `model.mid`, the layer is called `attn`. - if not in_mid: - new_item = new_item.replace("attn", "attentions") - new_item = new_item.replace(".k.", ".key.") - new_item = new_item.replace(".v.", ".value.") - new_item = new_item.replace(".q.", ".query.") - - new_item = new_item.replace("proj_out", "proj_attn") - new_item = new_item.replace("norm", "group_norm") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - if attention_paths_to_split is not None: - if config is None: - raise ValueError("Please specify the config if setting 'attention_paths_to_split' to 'True'.") - - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config.get("num_head_channels", 1) // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape).squeeze() - checkpoint[path_map["key"]] = key.reshape(target_shape).squeeze() - checkpoint[path_map["value"]] = value.reshape(target_shape).squeeze() - - for path in paths: - new_path = path["new"] - - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - new_path = new_path.replace("down.", "down_blocks.") - new_path = new_path.replace("up.", "up_blocks.") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - if "attentions" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]].squeeze() - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def convert_ddpm_checkpoint(checkpoint, config): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = checkpoint["temb.dense.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = checkpoint["temb.dense.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = checkpoint["temb.dense.1.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = checkpoint["temb.dense.1.bias"] - - new_checkpoint["conv_norm_out.weight"] = checkpoint["norm_out.weight"] - new_checkpoint["conv_norm_out.bias"] = checkpoint["norm_out.bias"] - - new_checkpoint["conv_in.weight"] = checkpoint["conv_in.weight"] - new_checkpoint["conv_in.bias"] = checkpoint["conv_in.bias"] - new_checkpoint["conv_out.weight"] = checkpoint["conv_out.weight"] - new_checkpoint["conv_out.bias"] = checkpoint["conv_out.bias"] - - num_down_blocks = len({".".join(layer.split(".")[:2]) for layer in checkpoint if "down" in layer}) - down_blocks = { - layer_id: [key for key in checkpoint if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - num_up_blocks = len({".".join(layer.split(".")[:2]) for layer in checkpoint if "up" in layer}) - up_blocks = {layer_id: [key for key in checkpoint if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)} - - for i in range(num_down_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - - if any("downsample" in layer for layer in down_blocks[i]): - new_checkpoint[f"down_blocks.{i}.downsamplers.0.conv.weight"] = checkpoint[ - f"down.{i}.downsample.op.weight" - ] - new_checkpoint[f"down_blocks.{i}.downsamplers.0.conv.bias"] = checkpoint[f"down.{i}.downsample.op.bias"] - # new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.weight'] = checkpoint[f'down.{i}.downsample.conv.weight'] - # new_checkpoint[f'down_blocks.{i}.downsamplers.0.op.bias'] = checkpoint[f'down.{i}.downsample.conv.bias'] - - if any("block" in layer for layer in down_blocks[i]): - num_blocks = len( - {".".join(shave_segments(layer, 2).split(".")[:2]) for layer in down_blocks[i] if "block" in layer} - ) - blocks = { - layer_id: [key for key in down_blocks[i] if f"block.{layer_id}" in key] - for layer_id in range(num_blocks) - } - - if num_blocks > 0: - for j in range(config["layers_per_block"]): - paths = renew_resnet_paths(blocks[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint) - - if any("attn" in layer for layer in down_blocks[i]): - num_attn = len( - {".".join(shave_segments(layer, 2).split(".")[:2]) for layer in down_blocks[i] if "attn" in layer} - ) - attns = { - layer_id: [key for key in down_blocks[i] if f"attn.{layer_id}" in key] - for layer_id in range(num_blocks) - } - - if num_attn > 0: - for j in range(config["layers_per_block"]): - paths = renew_attention_paths(attns[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, config=config) - - mid_block_1_layers = [key for key in checkpoint if "mid.block_1" in key] - mid_block_2_layers = [key for key in checkpoint if "mid.block_2" in key] - mid_attn_1_layers = [key for key in checkpoint if "mid.attn_1" in key] - - # Mid new 2 - paths = renew_resnet_paths(mid_block_1_layers) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "block_1", "new": "resnets.0"}], - ) - - paths = renew_resnet_paths(mid_block_2_layers) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "block_2", "new": "resnets.1"}], - ) - - paths = renew_attention_paths(mid_attn_1_layers, in_mid=True) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "attn_1", "new": "attentions.0"}], - ) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - - if any("upsample" in layer for layer in up_blocks[i]): - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = checkpoint[ - f"up.{i}.upsample.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = checkpoint[f"up.{i}.upsample.conv.bias"] - - if any("block" in layer for layer in up_blocks[i]): - num_blocks = len( - {".".join(shave_segments(layer, 2).split(".")[:2]) for layer in up_blocks[i] if "block" in layer} - ) - blocks = { - layer_id: [key for key in up_blocks[i] if f"block.{layer_id}" in key] for layer_id in range(num_blocks) - } - - if num_blocks > 0: - for j in range(config["layers_per_block"] + 1): - replace_indices = {"old": f"up_blocks.{i}", "new": f"up_blocks.{block_id}"} - paths = renew_resnet_paths(blocks[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, additional_replacements=[replace_indices]) - - if any("attn" in layer for layer in up_blocks[i]): - num_attn = len( - {".".join(shave_segments(layer, 2).split(".")[:2]) for layer in up_blocks[i] if "attn" in layer} - ) - attns = { - layer_id: [key for key in up_blocks[i] if f"attn.{layer_id}" in key] for layer_id in range(num_blocks) - } - - if num_attn > 0: - for j in range(config["layers_per_block"] + 1): - replace_indices = {"old": f"up_blocks.{i}", "new": f"up_blocks.{block_id}"} - paths = renew_attention_paths(attns[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, additional_replacements=[replace_indices]) - - new_checkpoint = {k.replace("mid_new_2", "mid_block"): v for k, v in new_checkpoint.items()} - return new_checkpoint - - -def convert_vq_autoenc_checkpoint(checkpoint, config): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - new_checkpoint = {} - - new_checkpoint["encoder.conv_norm_out.weight"] = checkpoint["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = checkpoint["encoder.norm_out.bias"] - - new_checkpoint["encoder.conv_in.weight"] = checkpoint["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = checkpoint["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = checkpoint["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = checkpoint["encoder.conv_out.bias"] - - new_checkpoint["decoder.conv_norm_out.weight"] = checkpoint["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = checkpoint["decoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = checkpoint["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = checkpoint["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = checkpoint["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = checkpoint["decoder.conv_out.bias"] - - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in checkpoint if "down" in layer}) - down_blocks = { - layer_id: [key for key in checkpoint if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in checkpoint if "up" in layer}) - up_blocks = {layer_id: [key for key in checkpoint if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)} - - for i in range(num_down_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - - if any("downsample" in layer for layer in down_blocks[i]): - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = checkpoint[ - f"encoder.down.{i}.downsample.conv.weight" - ] - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = checkpoint[ - f"encoder.down.{i}.downsample.conv.bias" - ] - - if any("block" in layer for layer in down_blocks[i]): - num_blocks = len( - {".".join(shave_segments(layer, 3).split(".")[:3]) for layer in down_blocks[i] if "block" in layer} - ) - blocks = { - layer_id: [key for key in down_blocks[i] if f"block.{layer_id}" in key] - for layer_id in range(num_blocks) - } - - if num_blocks > 0: - for j in range(config["layers_per_block"]): - paths = renew_resnet_paths(blocks[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint) - - if any("attn" in layer for layer in down_blocks[i]): - num_attn = len( - {".".join(shave_segments(layer, 3).split(".")[:3]) for layer in down_blocks[i] if "attn" in layer} - ) - attns = { - layer_id: [key for key in down_blocks[i] if f"attn.{layer_id}" in key] - for layer_id in range(num_blocks) - } - - if num_attn > 0: - for j in range(config["layers_per_block"]): - paths = renew_attention_paths(attns[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, config=config) - - mid_block_1_layers = [key for key in checkpoint if "mid.block_1" in key] - mid_block_2_layers = [key for key in checkpoint if "mid.block_2" in key] - mid_attn_1_layers = [key for key in checkpoint if "mid.attn_1" in key] - - # Mid new 2 - paths = renew_resnet_paths(mid_block_1_layers) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "block_1", "new": "resnets.0"}], - ) - - paths = renew_resnet_paths(mid_block_2_layers) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "block_2", "new": "resnets.1"}], - ) - - paths = renew_attention_paths(mid_attn_1_layers, in_mid=True) - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[{"old": "mid.", "new": "mid_new_2."}, {"old": "attn_1", "new": "attentions.0"}], - ) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - - if any("upsample" in layer for layer in up_blocks[i]): - new_checkpoint[f"decoder.up_blocks.{block_id}.upsamplers.0.conv.weight"] = checkpoint[ - f"decoder.up.{i}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{block_id}.upsamplers.0.conv.bias"] = checkpoint[ - f"decoder.up.{i}.upsample.conv.bias" - ] - - if any("block" in layer for layer in up_blocks[i]): - num_blocks = len( - {".".join(shave_segments(layer, 3).split(".")[:3]) for layer in up_blocks[i] if "block" in layer} - ) - blocks = { - layer_id: [key for key in up_blocks[i] if f"block.{layer_id}" in key] for layer_id in range(num_blocks) - } - - if num_blocks > 0: - for j in range(config["layers_per_block"] + 1): - replace_indices = {"old": f"up_blocks.{i}", "new": f"up_blocks.{block_id}"} - paths = renew_resnet_paths(blocks[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, additional_replacements=[replace_indices]) - - if any("attn" in layer for layer in up_blocks[i]): - num_attn = len( - {".".join(shave_segments(layer, 3).split(".")[:3]) for layer in up_blocks[i] if "attn" in layer} - ) - attns = { - layer_id: [key for key in up_blocks[i] if f"attn.{layer_id}" in key] for layer_id in range(num_blocks) - } - - if num_attn > 0: - for j in range(config["layers_per_block"] + 1): - replace_indices = {"old": f"up_blocks.{i}", "new": f"up_blocks.{block_id}"} - paths = renew_attention_paths(attns[j]) - assign_to_checkpoint(paths, new_checkpoint, checkpoint, additional_replacements=[replace_indices]) - - new_checkpoint = {k.replace("mid_new_2", "mid_block"): v for k, v in new_checkpoint.items()} - new_checkpoint["quant_conv.weight"] = checkpoint["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = checkpoint["quant_conv.bias"] - if "quantize.embedding.weight" in checkpoint: - new_checkpoint["quantize.embedding.weight"] = checkpoint["quantize.embedding.weight"] - new_checkpoint["post_quant_conv.weight"] = checkpoint["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = checkpoint["post_quant_conv.bias"] - - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - - parser.add_argument( - "--config_file", - default=None, - type=str, - required=True, - help="The config json file corresponding to the architecture.", - ) - - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - - args = parser.parse_args() - checkpoint = torch.load(args.checkpoint_path) - - with open(args.config_file) as f: - config = json.loads(f.read()) - - # unet case - key_prefix_set = {key.split(".")[0] for key in checkpoint.keys()} - if "encoder" in key_prefix_set and "decoder" in key_prefix_set: - converted_checkpoint = convert_vq_autoenc_checkpoint(checkpoint, config) - else: - converted_checkpoint = convert_ddpm_checkpoint(checkpoint, config) - - if "ddpm" in config: - del config["ddpm"] - - if config["_class_name"] == "VQModel": - model = VQModel(**config) - model.load_state_dict(converted_checkpoint) - model.save_pretrained(args.dump_path) - elif config["_class_name"] == "AutoencoderKL": - model = AutoencoderKL(**config) - model.load_state_dict(converted_checkpoint) - model.save_pretrained(args.dump_path) - else: - model = UNet2DModel(**config) - model.load_state_dict(converted_checkpoint) - - scheduler = DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) - - pipe = DDPMPipeline(unet=model, scheduler=scheduler) - pipe.save_pretrained(args.dump_path) diff --git a/scripts/convert_diffusers_sdxl_lora_to_webui.py b/scripts/convert_diffusers_sdxl_lora_to_webui.py deleted file mode 100644 index dfb3871275cb..000000000000 --- a/scripts/convert_diffusers_sdxl_lora_to_webui.py +++ /dev/null @@ -1,56 +0,0 @@ -# Script for converting a Hugging Face Diffusers trained SDXL LoRAs to Kohya format -# This means that you can input your diffusers-trained LoRAs and -# Get the output to work with WebUIs such as AUTOMATIC1111, ComfyUI, SD.Next and others. - -# To get started you can find some cool `diffusers` trained LoRAs such as this cute Corgy -# https://huggingface.co/ignasbud/corgy_dog_LoRA/, download its `pytorch_lora_weights.safetensors` file -# and run the script: -# python convert_diffusers_sdxl_lora_to_webui.py --input_lora pytorch_lora_weights.safetensors --output_lora corgy.safetensors -# now you can use corgy.safetensors in your WebUI of choice! - -# To train your own, here are some diffusers training scripts and utils that you can use and then convert: -# LoRA Ease - no code SDXL Dreambooth LoRA trainer: https://huggingface.co/spaces/multimodalart/lora-ease -# Dreambooth Advanced Training Script - state of the art techniques such as pivotal tuning and prodigy optimizer: -# - Script: https://github.com/huggingface/diffusers/blob/main/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py -# - Colab (only on Pro): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_Dreambooth_LoRA_advanced_example.ipynb -# Canonical diffusers training scripts: -# - Script: https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py -# - Colab (runs on free tier): https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/SDXL_DreamBooth_LoRA_.ipynb - -import argparse -import os - -from safetensors.torch import load_file, save_file - -from diffusers.utils import convert_all_state_dict_to_peft, convert_state_dict_to_kohya - - -def convert_and_save(input_lora, output_lora=None): - if output_lora is None: - base_name = os.path.splitext(input_lora)[0] - output_lora = f"{base_name}_webui.safetensors" - - diffusers_state_dict = load_file(input_lora) - peft_state_dict = convert_all_state_dict_to_peft(diffusers_state_dict) - kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict) - save_file(kohya_state_dict, output_lora) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert LoRA model to PEFT and then to Kohya format.") - parser.add_argument( - "--input_lora", - type=str, - required=True, - help="Path to the input LoRA model file in the diffusers format.", - ) - parser.add_argument( - "--output_lora", - type=str, - required=False, - help="Path for the converted LoRA (safetensors format for AUTOMATIC1111, ComfyUI, etc.). Optional, defaults to input name with a _webui suffix.", - ) - - args = parser.parse_args() - - convert_and_save(args.input_lora, args.output_lora) diff --git a/scripts/convert_diffusers_to_original_sdxl.py b/scripts/convert_diffusers_to_original_sdxl.py deleted file mode 100644 index 1aa792b3f06a..000000000000 --- a/scripts/convert_diffusers_to_original_sdxl.py +++ /dev/null @@ -1,350 +0,0 @@ -# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint. -# *Only* converts the UNet, VAE, and Text Encoder. -# Does not convert optimizer state or any other thing. - -import argparse -import os.path as osp -import re - -import torch -from safetensors.torch import load_file, save_file - - -# =================# -# UNet Conversion # -# =================# - -unet_conversion_map = [ - # (stable-diffusion, HF Diffusers) - ("time_embed.0.weight", "time_embedding.linear_1.weight"), - ("time_embed.0.bias", "time_embedding.linear_1.bias"), - ("time_embed.2.weight", "time_embedding.linear_2.weight"), - ("time_embed.2.bias", "time_embedding.linear_2.bias"), - ("input_blocks.0.0.weight", "conv_in.weight"), - ("input_blocks.0.0.bias", "conv_in.bias"), - ("out.0.weight", "conv_norm_out.weight"), - ("out.0.bias", "conv_norm_out.bias"), - ("out.2.weight", "conv_out.weight"), - ("out.2.bias", "conv_out.bias"), - # the following are for sdxl - ("label_emb.0.0.weight", "add_embedding.linear_1.weight"), - ("label_emb.0.0.bias", "add_embedding.linear_1.bias"), - ("label_emb.0.2.weight", "add_embedding.linear_2.weight"), - ("label_emb.0.2.bias", "add_embedding.linear_2.bias"), -] - -unet_conversion_map_resnet = [ - # (stable-diffusion, HF Diffusers) - ("in_layers.0", "norm1"), - ("in_layers.2", "conv1"), - ("out_layers.0", "norm2"), - ("out_layers.3", "conv2"), - ("emb_layers.1", "time_emb_proj"), - ("skip_connection", "conv_shortcut"), -] - -unet_conversion_map_layer = [] -# hardcoded number of downblocks and resnets/attentions... -# would need smarter logic for other networks. -for i in range(3): - # loop over downblocks/upblocks - - for j in range(2): - # loop over resnets/attentions for downblocks - hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}." - sd_down_res_prefix = f"input_blocks.{3 * i + j + 1}.0." - unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix)) - - if i > 0: - hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}." - sd_down_atn_prefix = f"input_blocks.{3 * i + j + 1}.1." - unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix)) - - for j in range(4): - # loop over resnets/attentions for upblocks - hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}." - sd_up_res_prefix = f"output_blocks.{3 * i + j}.0." - unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix)) - - if i < 2: - # no attention layers in up_blocks.0 - hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}." - sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1." - unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix)) - - if i < 3: - # no downsample in down_blocks.3 - hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv." - sd_downsample_prefix = f"input_blocks.{3 * (i + 1)}.0.op." - unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix)) - - # no upsample in up_blocks.3 - hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." - sd_upsample_prefix = f"output_blocks.{3 * i + 2}.{1 if i == 0 else 2}." - unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix)) -unet_conversion_map_layer.append(("output_blocks.2.2.conv.", "output_blocks.2.1.conv.")) - -hf_mid_atn_prefix = "mid_block.attentions.0." -sd_mid_atn_prefix = "middle_block.1." -unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix)) -for j in range(2): - hf_mid_res_prefix = f"mid_block.resnets.{j}." - sd_mid_res_prefix = f"middle_block.{2 * j}." - unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix)) - - -def convert_unet_state_dict(unet_state_dict): - # buyer beware: this is a *brittle* function, - # and correct output requires that all of these pieces interact in - # the exact order in which I have arranged them. - mapping = {k: k for k in unet_state_dict.keys()} - for sd_name, hf_name in unet_conversion_map: - mapping[hf_name] = sd_name - for k, v in mapping.items(): - if "resnets" in k: - for sd_part, hf_part in unet_conversion_map_resnet: - v = v.replace(hf_part, sd_part) - mapping[k] = v - for k, v in mapping.items(): - for sd_part, hf_part in unet_conversion_map_layer: - v = v.replace(hf_part, sd_part) - mapping[k] = v - new_state_dict = {sd_name: unet_state_dict[hf_name] for hf_name, sd_name in mapping.items()} - return new_state_dict - - -# ================# -# VAE Conversion # -# ================# - -vae_conversion_map = [ - # (stable-diffusion, HF Diffusers) - ("nin_shortcut", "conv_shortcut"), - ("norm_out", "conv_norm_out"), - ("mid.attn_1.", "mid_block.attentions.0."), -] - -for i in range(4): - # down_blocks have two resnets - for j in range(2): - hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}." - sd_down_prefix = f"encoder.down.{i}.block.{j}." - vae_conversion_map.append((sd_down_prefix, hf_down_prefix)) - - if i < 3: - hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0." - sd_downsample_prefix = f"down.{i}.downsample." - vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix)) - - hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." - sd_upsample_prefix = f"up.{3 - i}.upsample." - vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix)) - - # up_blocks have three resnets - # also, up blocks in hf are numbered in reverse from sd - for j in range(3): - hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}." - sd_up_prefix = f"decoder.up.{3 - i}.block.{j}." - vae_conversion_map.append((sd_up_prefix, hf_up_prefix)) - -# this part accounts for mid blocks in both the encoder and the decoder -for i in range(2): - hf_mid_res_prefix = f"mid_block.resnets.{i}." - sd_mid_res_prefix = f"mid.block_{i + 1}." - vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix)) - - -vae_conversion_map_attn = [ - # (stable-diffusion, HF Diffusers) - ("norm.", "group_norm."), - # the following are for SDXL - ("q.", "to_q."), - ("k.", "to_k."), - ("v.", "to_v."), - ("proj_out.", "to_out.0."), -] - - -def reshape_weight_for_sd(w): - # convert HF linear weights to SD conv2d weights - if not w.ndim == 1: - return w.reshape(*w.shape, 1, 1) - else: - return w - - -def convert_vae_state_dict(vae_state_dict): - mapping = {k: k for k in vae_state_dict.keys()} - for k, v in mapping.items(): - for sd_part, hf_part in vae_conversion_map: - v = v.replace(hf_part, sd_part) - mapping[k] = v - for k, v in mapping.items(): - if "attentions" in k: - for sd_part, hf_part in vae_conversion_map_attn: - v = v.replace(hf_part, sd_part) - mapping[k] = v - new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()} - weights_to_convert = ["q", "k", "v", "proj_out"] - for k, v in new_state_dict.items(): - for weight_name in weights_to_convert: - if f"mid.attn_1.{weight_name}.weight" in k: - print(f"Reshaping {k} for SD format") - new_state_dict[k] = reshape_weight_for_sd(v) - return new_state_dict - - -# =========================# -# Text Encoder Conversion # -# =========================# - - -textenc_conversion_lst = [ - # (stable-diffusion, HF Diffusers) - ("transformer.resblocks.", "text_model.encoder.layers."), - ("ln_1", "layer_norm1"), - ("ln_2", "layer_norm2"), - (".c_fc.", ".fc1."), - (".c_proj.", ".fc2."), - (".attn", ".self_attn"), - ("ln_final.", "text_model.final_layer_norm."), - ("token_embedding.weight", "text_model.embeddings.token_embedding.weight"), - ("positional_embedding", "text_model.embeddings.position_embedding.weight"), -] -protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst} -textenc_pattern = re.compile("|".join(protected.keys())) - -# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp -code2idx = {"q": 0, "k": 1, "v": 2} - - -def convert_openclip_text_enc_state_dict(text_enc_dict): - new_state_dict = {} - capture_qkv_weight = {} - capture_qkv_bias = {} - for k, v in text_enc_dict.items(): - if ( - k.endswith(".self_attn.q_proj.weight") - or k.endswith(".self_attn.k_proj.weight") - or k.endswith(".self_attn.v_proj.weight") - ): - k_pre = k[: -len(".q_proj.weight")] - k_code = k[-len("q_proj.weight")] - if k_pre not in capture_qkv_weight: - capture_qkv_weight[k_pre] = [None, None, None] - capture_qkv_weight[k_pre][code2idx[k_code]] = v - continue - - if ( - k.endswith(".self_attn.q_proj.bias") - or k.endswith(".self_attn.k_proj.bias") - or k.endswith(".self_attn.v_proj.bias") - ): - k_pre = k[: -len(".q_proj.bias")] - k_code = k[-len("q_proj.bias")] - if k_pre not in capture_qkv_bias: - capture_qkv_bias[k_pre] = [None, None, None] - capture_qkv_bias[k_pre][code2idx[k_code]] = v - continue - - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k) - new_state_dict[relabelled_key] = v - - for k_pre, tensors in capture_qkv_weight.items(): - if None in tensors: - raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors) - - for k_pre, tensors in capture_qkv_bias.items(): - if None in tensors: - raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors) - - return new_state_dict - - -def convert_openai_text_enc_state_dict(text_enc_dict): - return text_enc_dict - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.") - parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--half", action="store_true", help="Save weights in half precision.") - parser.add_argument( - "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt." - ) - - args = parser.parse_args() - - assert args.model_path is not None, "Must provide a model path!" - - assert args.checkpoint_path is not None, "Must provide a checkpoint path!" - - # Path for safetensors - unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors") - vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors") - text_enc_path = osp.join(args.model_path, "text_encoder", "model.safetensors") - text_enc_2_path = osp.join(args.model_path, "text_encoder_2", "model.safetensors") - - # Load models from safetensors if it exists, if it doesn't pytorch - if osp.exists(unet_path): - unet_state_dict = load_file(unet_path, device="cpu") - else: - unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin") - unet_state_dict = torch.load(unet_path, map_location="cpu") - - if osp.exists(vae_path): - vae_state_dict = load_file(vae_path, device="cpu") - else: - vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin") - vae_state_dict = torch.load(vae_path, map_location="cpu") - - if osp.exists(text_enc_path): - text_enc_dict = load_file(text_enc_path, device="cpu") - else: - text_enc_path = osp.join(args.model_path, "text_encoder", "pytorch_model.bin") - text_enc_dict = torch.load(text_enc_path, map_location="cpu") - - if osp.exists(text_enc_2_path): - text_enc_2_dict = load_file(text_enc_2_path, device="cpu") - else: - text_enc_2_path = osp.join(args.model_path, "text_encoder_2", "pytorch_model.bin") - text_enc_2_dict = torch.load(text_enc_2_path, map_location="cpu") - - # Convert the UNet model - unet_state_dict = convert_unet_state_dict(unet_state_dict) - unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()} - - # Convert the VAE model - vae_state_dict = convert_vae_state_dict(vae_state_dict) - vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()} - - # Convert text encoder 1 - text_enc_dict = convert_openai_text_enc_state_dict(text_enc_dict) - text_enc_dict = {"conditioner.embedders.0.transformer." + k: v for k, v in text_enc_dict.items()} - - # Convert text encoder 2 - text_enc_2_dict = convert_openclip_text_enc_state_dict(text_enc_2_dict) - text_enc_2_dict = {"conditioner.embedders.1.model." + k: v for k, v in text_enc_2_dict.items()} - # We call the `.T.contiguous()` to match what's done in - # https://github.com/huggingface/diffusers/blob/84905ca7287876b925b6bf8e9bb92fec21c78764/src/diffusers/loaders/single_file_utils.py#L1085 - text_enc_2_dict["conditioner.embedders.1.model.text_projection"] = text_enc_2_dict.pop( - "conditioner.embedders.1.model.text_projection.weight" - ).T.contiguous() - - # Put together new checkpoint - state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict, **text_enc_2_dict} - - if args.half: - state_dict = {k: v.half() for k, v in state_dict.items()} - - if args.use_safetensors: - save_file(state_dict, args.checkpoint_path) - else: - state_dict = {"state_dict": state_dict} - torch.save(state_dict, args.checkpoint_path) diff --git a/scripts/convert_diffusers_to_original_stable_diffusion.py b/scripts/convert_diffusers_to_original_stable_diffusion.py deleted file mode 100644 index 049dda7d42a7..000000000000 --- a/scripts/convert_diffusers_to_original_stable_diffusion.py +++ /dev/null @@ -1,353 +0,0 @@ -# Script for converting a HF Diffusers saved pipeline to a Stable Diffusion checkpoint. -# *Only* converts the UNet, VAE, and Text Encoder. -# Does not convert optimizer state or any other thing. - -import argparse -import os.path as osp -import re - -import torch -from safetensors.torch import load_file, save_file - - -# =================# -# UNet Conversion # -# =================# - -unet_conversion_map = [ - # (stable-diffusion, HF Diffusers) - ("time_embed.0.weight", "time_embedding.linear_1.weight"), - ("time_embed.0.bias", "time_embedding.linear_1.bias"), - ("time_embed.2.weight", "time_embedding.linear_2.weight"), - ("time_embed.2.bias", "time_embedding.linear_2.bias"), - ("input_blocks.0.0.weight", "conv_in.weight"), - ("input_blocks.0.0.bias", "conv_in.bias"), - ("out.0.weight", "conv_norm_out.weight"), - ("out.0.bias", "conv_norm_out.bias"), - ("out.2.weight", "conv_out.weight"), - ("out.2.bias", "conv_out.bias"), -] - -unet_conversion_map_resnet = [ - # (stable-diffusion, HF Diffusers) - ("in_layers.0", "norm1"), - ("in_layers.2", "conv1"), - ("out_layers.0", "norm2"), - ("out_layers.3", "conv2"), - ("emb_layers.1", "time_emb_proj"), - ("skip_connection", "conv_shortcut"), -] - -unet_conversion_map_layer = [] -# hardcoded number of downblocks and resnets/attentions... -# would need smarter logic for other networks. -for i in range(4): - # loop over downblocks/upblocks - - for j in range(2): - # loop over resnets/attentions for downblocks - hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}." - sd_down_res_prefix = f"input_blocks.{3 * i + j + 1}.0." - unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix)) - - if i < 3: - # no attention layers in down_blocks.3 - hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}." - sd_down_atn_prefix = f"input_blocks.{3 * i + j + 1}.1." - unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix)) - - for j in range(3): - # loop over resnets/attentions for upblocks - hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}." - sd_up_res_prefix = f"output_blocks.{3 * i + j}.0." - unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix)) - - if i > 0: - # no attention layers in up_blocks.0 - hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}." - sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1." - unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix)) - - if i < 3: - # no downsample in down_blocks.3 - hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv." - sd_downsample_prefix = f"input_blocks.{3 * (i + 1)}.0.op." - unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix)) - - # no upsample in up_blocks.3 - hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." - sd_upsample_prefix = f"output_blocks.{3 * i + 2}.{1 if i == 0 else 2}." - unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix)) - -hf_mid_atn_prefix = "mid_block.attentions.0." -sd_mid_atn_prefix = "middle_block.1." -unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix)) - -for j in range(2): - hf_mid_res_prefix = f"mid_block.resnets.{j}." - sd_mid_res_prefix = f"middle_block.{2 * j}." - unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix)) - - -def convert_unet_state_dict(unet_state_dict): - # buyer beware: this is a *brittle* function, - # and correct output requires that all of these pieces interact in - # the exact order in which I have arranged them. - mapping = {k: k for k in unet_state_dict.keys()} - for sd_name, hf_name in unet_conversion_map: - mapping[hf_name] = sd_name - for k, v in mapping.items(): - if "resnets" in k: - for sd_part, hf_part in unet_conversion_map_resnet: - v = v.replace(hf_part, sd_part) - mapping[k] = v - for k, v in mapping.items(): - for sd_part, hf_part in unet_conversion_map_layer: - v = v.replace(hf_part, sd_part) - mapping[k] = v - new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()} - return new_state_dict - - -# ================# -# VAE Conversion # -# ================# - -vae_conversion_map = [ - # (stable-diffusion, HF Diffusers) - ("nin_shortcut", "conv_shortcut"), - ("norm_out", "conv_norm_out"), - ("mid.attn_1.", "mid_block.attentions.0."), -] - -for i in range(4): - # down_blocks have two resnets - for j in range(2): - hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}." - sd_down_prefix = f"encoder.down.{i}.block.{j}." - vae_conversion_map.append((sd_down_prefix, hf_down_prefix)) - - if i < 3: - hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0." - sd_downsample_prefix = f"down.{i}.downsample." - vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix)) - - hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0." - sd_upsample_prefix = f"up.{3 - i}.upsample." - vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix)) - - # up_blocks have three resnets - # also, up blocks in hf are numbered in reverse from sd - for j in range(3): - hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}." - sd_up_prefix = f"decoder.up.{3 - i}.block.{j}." - vae_conversion_map.append((sd_up_prefix, hf_up_prefix)) - -# this part accounts for mid blocks in both the encoder and the decoder -for i in range(2): - hf_mid_res_prefix = f"mid_block.resnets.{i}." - sd_mid_res_prefix = f"mid.block_{i + 1}." - vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix)) - - -vae_conversion_map_attn = [ - # (stable-diffusion, HF Diffusers) - ("norm.", "group_norm."), - ("q.", "query."), - ("k.", "key."), - ("v.", "value."), - ("proj_out.", "proj_attn."), -] - -# This is probably not the most ideal solution, but it does work. -vae_extra_conversion_map = [ - ("to_q", "q"), - ("to_k", "k"), - ("to_v", "v"), - ("to_out.0", "proj_out"), -] - - -def reshape_weight_for_sd(w): - # convert HF linear weights to SD conv2d weights - if not w.ndim == 1: - return w.reshape(*w.shape, 1, 1) - else: - return w - - -def convert_vae_state_dict(vae_state_dict): - mapping = {k: k for k in vae_state_dict.keys()} - for k, v in mapping.items(): - for sd_part, hf_part in vae_conversion_map: - v = v.replace(hf_part, sd_part) - mapping[k] = v - for k, v in mapping.items(): - if "attentions" in k: - for sd_part, hf_part in vae_conversion_map_attn: - v = v.replace(hf_part, sd_part) - mapping[k] = v - new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()} - weights_to_convert = ["q", "k", "v", "proj_out"] - keys_to_rename = {} - for k, v in new_state_dict.items(): - for weight_name in weights_to_convert: - if f"mid.attn_1.{weight_name}.weight" in k: - print(f"Reshaping {k} for SD format") - new_state_dict[k] = reshape_weight_for_sd(v) - for weight_name, real_weight_name in vae_extra_conversion_map: - if f"mid.attn_1.{weight_name}.weight" in k or f"mid.attn_1.{weight_name}.bias" in k: - keys_to_rename[k] = k.replace(weight_name, real_weight_name) - for k, v in keys_to_rename.items(): - if k in new_state_dict: - print(f"Renaming {k} to {v}") - new_state_dict[v] = reshape_weight_for_sd(new_state_dict[k]) - del new_state_dict[k] - return new_state_dict - - -# =========================# -# Text Encoder Conversion # -# =========================# - - -textenc_conversion_lst = [ - # (stable-diffusion, HF Diffusers) - ("resblocks.", "text_model.encoder.layers."), - ("ln_1", "layer_norm1"), - ("ln_2", "layer_norm2"), - (".c_fc.", ".fc1."), - (".c_proj.", ".fc2."), - (".attn", ".self_attn"), - ("ln_final.", "transformer.text_model.final_layer_norm."), - ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"), - ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"), -] -protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst} -textenc_pattern = re.compile("|".join(protected.keys())) - -# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp -code2idx = {"q": 0, "k": 1, "v": 2} - - -def convert_text_enc_state_dict_v20(text_enc_dict): - new_state_dict = {} - capture_qkv_weight = {} - capture_qkv_bias = {} - for k, v in text_enc_dict.items(): - if ( - k.endswith(".self_attn.q_proj.weight") - or k.endswith(".self_attn.k_proj.weight") - or k.endswith(".self_attn.v_proj.weight") - ): - k_pre = k[: -len(".q_proj.weight")] - k_code = k[-len("q_proj.weight")] - if k_pre not in capture_qkv_weight: - capture_qkv_weight[k_pre] = [None, None, None] - capture_qkv_weight[k_pre][code2idx[k_code]] = v - continue - - if ( - k.endswith(".self_attn.q_proj.bias") - or k.endswith(".self_attn.k_proj.bias") - or k.endswith(".self_attn.v_proj.bias") - ): - k_pre = k[: -len(".q_proj.bias")] - k_code = k[-len("q_proj.bias")] - if k_pre not in capture_qkv_bias: - capture_qkv_bias[k_pre] = [None, None, None] - capture_qkv_bias[k_pre][code2idx[k_code]] = v - continue - - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k) - new_state_dict[relabelled_key] = v - - for k_pre, tensors in capture_qkv_weight.items(): - if None in tensors: - raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors) - - for k_pre, tensors in capture_qkv_bias.items(): - if None in tensors: - raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing") - relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre) - new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors) - - return new_state_dict - - -def convert_text_enc_state_dict(text_enc_dict): - return text_enc_dict - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.") - parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--half", action="store_true", help="Save weights in half precision.") - parser.add_argument( - "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt." - ) - - args = parser.parse_args() - - assert args.model_path is not None, "Must provide a model path!" - - assert args.checkpoint_path is not None, "Must provide a checkpoint path!" - - # Path for safetensors - unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors") - vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors") - text_enc_path = osp.join(args.model_path, "text_encoder", "model.safetensors") - - # Load models from safetensors if it exists, if it doesn't pytorch - if osp.exists(unet_path): - unet_state_dict = load_file(unet_path, device="cpu") - else: - unet_path = osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin") - unet_state_dict = torch.load(unet_path, map_location="cpu") - - if osp.exists(vae_path): - vae_state_dict = load_file(vae_path, device="cpu") - else: - vae_path = osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin") - vae_state_dict = torch.load(vae_path, map_location="cpu") - - if osp.exists(text_enc_path): - text_enc_dict = load_file(text_enc_path, device="cpu") - else: - text_enc_path = osp.join(args.model_path, "text_encoder", "pytorch_model.bin") - text_enc_dict = torch.load(text_enc_path, map_location="cpu") - - # Convert the UNet model - unet_state_dict = convert_unet_state_dict(unet_state_dict) - unet_state_dict = {"model.diffusion_model." + k: v for k, v in unet_state_dict.items()} - - # Convert the VAE model - vae_state_dict = convert_vae_state_dict(vae_state_dict) - vae_state_dict = {"first_stage_model." + k: v for k, v in vae_state_dict.items()} - - # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper - is_v20_model = "text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict - - if is_v20_model: - # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm - text_enc_dict = {"transformer." + k: v for k, v in text_enc_dict.items()} - text_enc_dict = convert_text_enc_state_dict_v20(text_enc_dict) - text_enc_dict = {"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()} - else: - text_enc_dict = convert_text_enc_state_dict(text_enc_dict) - text_enc_dict = {"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()} - - # Put together new checkpoint - state_dict = {**unet_state_dict, **vae_state_dict, **text_enc_dict} - if args.half: - state_dict = {k: v.half() for k, v in state_dict.items()} - - if args.use_safetensors: - save_file(state_dict, args.checkpoint_path) - else: - state_dict = {"state_dict": state_dict} - torch.save(state_dict, args.checkpoint_path) diff --git a/scripts/convert_dit_to_diffusers.py b/scripts/convert_dit_to_diffusers.py deleted file mode 100644 index dc127f69555c..000000000000 --- a/scripts/convert_dit_to_diffusers.py +++ /dev/null @@ -1,162 +0,0 @@ -import argparse -import os - -import torch -from torchvision.datasets.utils import download_url - -from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, Transformer2DModel - - -pretrained_models = {512: "DiT-XL-2-512x512.pt", 256: "DiT-XL-2-256x256.pt"} - - -def download_model(model_name): - """ - Downloads a pre-trained DiT model from the web. - """ - local_path = f"pretrained_models/{model_name}" - if not os.path.isfile(local_path): - os.makedirs("pretrained_models", exist_ok=True) - web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}" - download_url(web_path, "pretrained_models") - model = torch.load(local_path, map_location=lambda storage, loc: storage) - return model - - -def main(args): - state_dict = download_model(pretrained_models[args.image_size]) - - state_dict["pos_embed.proj.weight"] = state_dict["x_embedder.proj.weight"] - state_dict["pos_embed.proj.bias"] = state_dict["x_embedder.proj.bias"] - state_dict.pop("x_embedder.proj.weight") - state_dict.pop("x_embedder.proj.bias") - - for depth in range(28): - state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_1.weight"] = state_dict[ - "t_embedder.mlp.0.weight" - ] - state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_1.bias"] = state_dict[ - "t_embedder.mlp.0.bias" - ] - state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_2.weight"] = state_dict[ - "t_embedder.mlp.2.weight" - ] - state_dict[f"transformer_blocks.{depth}.norm1.emb.timestep_embedder.linear_2.bias"] = state_dict[ - "t_embedder.mlp.2.bias" - ] - state_dict[f"transformer_blocks.{depth}.norm1.emb.class_embedder.embedding_table.weight"] = state_dict[ - "y_embedder.embedding_table.weight" - ] - - state_dict[f"transformer_blocks.{depth}.norm1.linear.weight"] = state_dict[ - f"blocks.{depth}.adaLN_modulation.1.weight" - ] - state_dict[f"transformer_blocks.{depth}.norm1.linear.bias"] = state_dict[ - f"blocks.{depth}.adaLN_modulation.1.bias" - ] - - q, k, v = torch.chunk(state_dict[f"blocks.{depth}.attn.qkv.weight"], 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(state_dict[f"blocks.{depth}.attn.qkv.bias"], 3, dim=0) - - state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - state_dict[f"transformer_blocks.{depth}.attn1.to_q.bias"] = q_bias - state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - state_dict[f"transformer_blocks.{depth}.attn1.to_k.bias"] = k_bias - state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - state_dict[f"transformer_blocks.{depth}.attn1.to_v.bias"] = v_bias - - state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict[ - f"blocks.{depth}.attn.proj.weight" - ] - state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict[f"blocks.{depth}.attn.proj.bias"] - - state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.weight"] = state_dict[f"blocks.{depth}.mlp.fc1.weight"] - state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.bias"] = state_dict[f"blocks.{depth}.mlp.fc1.bias"] - state_dict[f"transformer_blocks.{depth}.ff.net.2.weight"] = state_dict[f"blocks.{depth}.mlp.fc2.weight"] - state_dict[f"transformer_blocks.{depth}.ff.net.2.bias"] = state_dict[f"blocks.{depth}.mlp.fc2.bias"] - - state_dict.pop(f"blocks.{depth}.attn.qkv.weight") - state_dict.pop(f"blocks.{depth}.attn.qkv.bias") - state_dict.pop(f"blocks.{depth}.attn.proj.weight") - state_dict.pop(f"blocks.{depth}.attn.proj.bias") - state_dict.pop(f"blocks.{depth}.mlp.fc1.weight") - state_dict.pop(f"blocks.{depth}.mlp.fc1.bias") - state_dict.pop(f"blocks.{depth}.mlp.fc2.weight") - state_dict.pop(f"blocks.{depth}.mlp.fc2.bias") - state_dict.pop(f"blocks.{depth}.adaLN_modulation.1.weight") - state_dict.pop(f"blocks.{depth}.adaLN_modulation.1.bias") - - state_dict.pop("t_embedder.mlp.0.weight") - state_dict.pop("t_embedder.mlp.0.bias") - state_dict.pop("t_embedder.mlp.2.weight") - state_dict.pop("t_embedder.mlp.2.bias") - state_dict.pop("y_embedder.embedding_table.weight") - - state_dict["proj_out_1.weight"] = state_dict["final_layer.adaLN_modulation.1.weight"] - state_dict["proj_out_1.bias"] = state_dict["final_layer.adaLN_modulation.1.bias"] - state_dict["proj_out_2.weight"] = state_dict["final_layer.linear.weight"] - state_dict["proj_out_2.bias"] = state_dict["final_layer.linear.bias"] - - state_dict.pop("final_layer.linear.weight") - state_dict.pop("final_layer.linear.bias") - state_dict.pop("final_layer.adaLN_modulation.1.weight") - state_dict.pop("final_layer.adaLN_modulation.1.bias") - - # DiT XL/2 - transformer = Transformer2DModel( - sample_size=args.image_size // 8, - num_layers=28, - attention_head_dim=72, - in_channels=4, - out_channels=8, - patch_size=2, - attention_bias=True, - num_attention_heads=16, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - norm_type="ada_norm_zero", - norm_elementwise_affine=False, - ) - transformer.load_state_dict(state_dict, strict=True) - - scheduler = DDIMScheduler( - num_train_timesteps=1000, - beta_schedule="linear", - prediction_type="epsilon", - clip_sample=False, - ) - - vae = AutoencoderKL.from_pretrained(args.vae_model) - - pipeline = DiTPipeline(transformer=transformer, vae=vae, scheduler=scheduler) - - if args.save: - pipeline.save_pretrained(args.checkpoint_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--image_size", - default=256, - type=int, - required=False, - help="Image size of pretrained model, either 256 or 512.", - ) - parser.add_argument( - "--vae_model", - default="stabilityai/sd-vae-ft-ema", - type=str, - required=False, - help="Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.", - ) - parser.add_argument( - "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." - ) - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the output pipeline." - ) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_flux2_to_diffusers.py b/scripts/convert_flux2_to_diffusers.py deleted file mode 100644 index a8fa6f87eee1..000000000000 --- a/scripts/convert_flux2_to_diffusers.py +++ /dev/null @@ -1,536 +0,0 @@ -import argparse -from contextlib import nullcontext -from typing import Any, Dict, Tuple - -import safetensors.torch -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download -from transformers import AutoProcessor, GenerationConfig, Mistral3ForConditionalGeneration - -from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2Pipeline, Flux2Transformer2DModel -from diffusers.utils.import_utils import is_accelerate_available - - -""" -# VAE - -python scripts/convert_flux2_to_diffusers.py \ ---original_state_dict_repo_id "diffusers-internal-dev/new-model-image" \ ---vae_filename "flux2-vae.sft" \ ---output_path "/raid/yiyi/dummy-flux2-diffusers" \ ---vae - -# DiT - -python scripts/convert_flux2_to_diffusers.py \ - --original_state_dict_repo_id diffusers-internal-dev/new-model-image \ - --dit_filename flux-dev-dummy.sft \ - --dit \ - --output_path . - -# Full pipe - -python scripts/convert_flux2_to_diffusers.py \ - --original_state_dict_repo_id diffusers-internal-dev/new-model-image \ - --dit_filename flux-dev-dummy.sft \ - --vae_filename "flux2-vae.sft" \ - --dit --vae --full_pipe \ - --output_path . -""" - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - -parser = argparse.ArgumentParser() -parser.add_argument("--original_state_dict_repo_id", default=None, type=str) -parser.add_argument("--vae_filename", default="flux2-vae.sft", type=str) -parser.add_argument("--dit_filename", default="flux2-dev.safetensors", type=str) -parser.add_argument("--vae", action="store_true") -parser.add_argument("--dit", action="store_true") -parser.add_argument("--vae_dtype", type=str, default="fp32") -parser.add_argument("--dit_dtype", type=str, default="bf16") -parser.add_argument("--checkpoint_path", default=None, type=str) -parser.add_argument("--full_pipe", action="store_true") -parser.add_argument("--output_path", type=str) - -args = parser.parse_args() - - -def load_original_checkpoint(args, filename): - if args.original_state_dict_repo_id is not None: - ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename) - elif args.checkpoint_path is not None: - ckpt_path = args.checkpoint_path - else: - raise ValueError(" please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") - - original_state_dict = safetensors.torch.load_file(ckpt_path) - return original_state_dict - - -DIFFUSERS_VAE_TO_FLUX2_MAPPING = { - "encoder.conv_in.weight": "encoder.conv_in.weight", - "encoder.conv_in.bias": "encoder.conv_in.bias", - "encoder.conv_out.weight": "encoder.conv_out.weight", - "encoder.conv_out.bias": "encoder.conv_out.bias", - "encoder.conv_norm_out.weight": "encoder.norm_out.weight", - "encoder.conv_norm_out.bias": "encoder.norm_out.bias", - "decoder.conv_in.weight": "decoder.conv_in.weight", - "decoder.conv_in.bias": "decoder.conv_in.bias", - "decoder.conv_out.weight": "decoder.conv_out.weight", - "decoder.conv_out.bias": "decoder.conv_out.bias", - "decoder.conv_norm_out.weight": "decoder.norm_out.weight", - "decoder.conv_norm_out.bias": "decoder.norm_out.bias", - "quant_conv.weight": "encoder.quant_conv.weight", - "quant_conv.bias": "encoder.quant_conv.bias", - "post_quant_conv.weight": "decoder.post_quant_conv.weight", - "post_quant_conv.bias": "decoder.post_quant_conv.bias", - "bn.running_mean": "bn.running_mean", - "bn.running_var": "bn.running_var", -} - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.conv_attn_to_linear -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def update_vae_resnet_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ldm_key.replace(mapping["old"], mapping["new"]).replace("nin_shortcut", "conv_shortcut") - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_vae_attentions_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ( - ldm_key.replace(mapping["old"], mapping["new"]) - .replace("norm.weight", "group_norm.weight") - .replace("norm.bias", "group_norm.bias") - .replace("q.weight", "to_q.weight") - .replace("q.bias", "to_q.bias") - .replace("k.weight", "to_k.weight") - .replace("k.bias", "to_k.bias") - .replace("v.weight", "to_v.weight") - .replace("v.bias", "to_v.bias") - .replace("proj_out.weight", "to_out.0.weight") - .replace("proj_out.bias", "to_out.0.bias") - ) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - # proj_attn.weight has to be converted from conv 1D to linear - shape = new_checkpoint[diffusers_key].shape - - if len(shape) == 3: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0] - elif len(shape) == 4: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0, 0] - - -def convert_flux2_vae_checkpoint_to_diffusers(vae_state_dict, config): - new_checkpoint = {} - for diffusers_key, ldm_key in DIFFUSERS_VAE_TO_FLUX2_MAPPING.items(): - if ldm_key not in vae_state_dict: - continue - new_checkpoint[diffusers_key] = vae_state_dict[ldm_key] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len(config["down_block_types"]) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}, - ) - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.bias" - ) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len(config["up_block_types"]) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}, - ) - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - conv_attn_to_linear(new_checkpoint) - - return new_checkpoint - - -FLUX2_TRANSFORMER_KEYS_RENAME_DICT = { - # Image and text input projections - "img_in": "x_embedder", - "txt_in": "context_embedder", - # Timestep and guidance embeddings - "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", - "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", - "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", - "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", - # Modulation parameters - "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", - "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", - "single_stream_modulation.lin": "single_stream_modulation.linear", - # Final output layer - # "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params - "final_layer.linear": "proj_out", -} - - -FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear", -} - - -FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = { - # Handle fused QKV projections separately as we need to break into Q, K, V projections - "img_attn.norm.query_norm": "attn.norm_q", - "img_attn.norm.key_norm": "attn.norm_k", - "img_attn.proj": "attn.to_out.0", - "img_mlp.0": "ff.linear_in", - "img_mlp.2": "ff.linear_out", - "txt_attn.norm.query_norm": "attn.norm_added_q", - "txt_attn.norm.key_norm": "attn.norm_added_k", - "txt_attn.proj": "attn.to_add_out", - "txt_mlp.0": "ff_context.linear_in", - "txt_mlp.2": "ff_context.linear_out", -} - - -FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = { - "linear1": "attn.to_qkv_mlp_proj", - "linear2": "attn.to_out", - "norm.query_norm": "attn.norm_q", - "norm.key_norm": "attn.norm_k", -} - - -# in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; -# while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use -# diffusers implementation -def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_ada_layer_norm_weights(key: str, state_dict: Dict[str, Any]) -> None: - # Skip if not a weight - if ".weight" not in key: - return - - # If adaLN_modulation is in the key, swap scale and shift parameters - # Original implementation is (shift, scale); diffusers implementation is (scale, shift) - if "adaLN_modulation" in key: - key_without_param_type, param_type = key.rsplit(".", maxsplit=1) - # Assume all such keys are in the AdaLayerNorm key map - new_key_without_param_type = FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP[key_without_param_type] - new_key = ".".join([new_key_without_param_type, param_type]) - - swapped_weight = swap_scale_shift(state_dict.pop(key)) - state_dict[new_key] = swapped_weight - return - - -def convert_flux2_double_stream_blocks(key: str, state_dict: Dict[str, Any]) -> None: - # Skip if not a weight, bias, or scale - if ".weight" not in key and ".bias" not in key and ".scale" not in key: - return - - new_prefix = "transformer_blocks" - if "double_blocks." in key: - parts = key.split(".") - block_idx = parts[1] - modality_block_name = parts[2] # img_attn, img_mlp, txt_attn, txt_mlp - within_block_name = ".".join(parts[2:-1]) - param_type = parts[-1] - - if param_type == "scale": - param_type = "weight" - - if "qkv" in within_block_name: - fused_qkv_weight = state_dict.pop(key) - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - if "img" in modality_block_name: - # double_blocks.{N}.img_attn.qkv --> transformer_blocks.{N}.attn.{to_q|to_k|to_v} - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - new_q_name = "attn.to_q" - new_k_name = "attn.to_k" - new_v_name = "attn.to_v" - elif "txt" in modality_block_name: - # double_blocks.{N}.txt_attn.qkv --> transformer_blocks.{N}.attn.{add_q_proj|add_k_proj|add_v_proj} - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - new_q_name = "attn.add_q_proj" - new_k_name = "attn.add_k_proj" - new_v_name = "attn.add_v_proj" - new_q_key = ".".join([new_prefix, block_idx, new_q_name, param_type]) - new_k_key = ".".join([new_prefix, block_idx, new_k_name, param_type]) - new_v_key = ".".join([new_prefix, block_idx, new_v_name, param_type]) - state_dict[new_q_key] = to_q_weight - state_dict[new_k_key] = to_k_weight - state_dict[new_v_key] = to_v_weight - else: - new_within_block_name = FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP[within_block_name] - new_key = ".".join([new_prefix, block_idx, new_within_block_name, param_type]) - - param = state_dict.pop(key) - state_dict[new_key] = param - return - - -def convert_flux2_single_stream_blocks(key: str, state_dict: Dict[str, Any]) -> None: - # Skip if not a weight, bias, or scale - if ".weight" not in key and ".bias" not in key and ".scale" not in key: - return - - # Mapping: - # - single_blocks.{N}.linear1 --> single_transformer_blocks.{N}.attn.to_qkv_mlp_proj - # - single_blocks.{N}.linear2 --> single_transformer_blocks.{N}.attn.to_out - # - single_blocks.{N}.norm.query_norm.scale --> single_transformer_blocks.{N}.attn.norm_q.weight - # - single_blocks.{N}.norm.key_norm.scale --> single_transformer_blocks.{N}.attn.norm_k.weight - new_prefix = "single_transformer_blocks" - if "single_blocks." in key: - parts = key.split(".") - block_idx = parts[1] - within_block_name = ".".join(parts[2:-1]) - param_type = parts[-1] - - if param_type == "scale": - param_type = "weight" - - new_within_block_name = FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP[within_block_name] - new_key = ".".join([new_prefix, block_idx, new_within_block_name, param_type]) - - param = state_dict.pop(key) - state_dict[new_key] = param - return - - -TRANSFORMER_SPECIAL_KEYS_REMAP = { - "adaLN_modulation": convert_ada_layer_norm_weights, - "double_blocks": convert_flux2_double_stream_blocks, - "single_blocks": convert_flux2_single_stream_blocks, -} - - -def update_state_dict(state_dict: Dict[str, Any], old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - -def get_flux2_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]: - if model_type == "flux2-dev": - config = { - "model_id": "black-forest-labs/FLUX.2-dev", - "diffusers_config": { - "patch_size": 1, - "in_channels": 128, - "num_layers": 8, - "num_single_layers": 48, - "attention_head_dim": 128, - "num_attention_heads": 48, - "joint_attention_dim": 15360, - "timestep_guidance_channels": 256, - "mlp_ratio": 3.0, - "axes_dims_rope": (32, 32, 32, 32), - "rope_theta": 2000, - "eps": 1e-6, - }, - } - rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "klein-4b": - config = { - "model_id": "diffusers-internal-dev/dummy0115", - "diffusers_config": { - "patch_size": 1, - "in_channels": 128, - "num_layers": 5, - "num_single_layers": 20, - "attention_head_dim": 128, - "num_attention_heads": 24, - "joint_attention_dim": 7680, - "timestep_guidance_channels": 256, - "mlp_ratio": 3.0, - "axes_dims_rope": (32, 32, 32, 32), - "rope_theta": 2000, - "eps": 1e-6, - "guidance_embeds": False, - }, - } - rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP - - elif model_type == "klein-9b": - config = { - "model_id": "diffusers-internal-dev/dummy0115", - "diffusers_config": { - "patch_size": 1, - "in_channels": 128, - "num_layers": 8, - "num_single_layers": 24, - "attention_head_dim": 128, - "num_attention_heads": 32, - "joint_attention_dim": 12288, - "timestep_guidance_channels": 256, - "mlp_ratio": 3.0, - "axes_dims_rope": (32, 32, 32, 32), - "rope_theta": 2000, - "eps": 1e-6, - "guidance_embeds": False, - }, - } - rename_dict = FLUX2_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = TRANSFORMER_SPECIAL_KEYS_REMAP - - else: - raise ValueError(f"Unknown model_type: {model_type}. Choose from: flux2-dev, klein-4b, klein-9b") - - return config, rename_dict, special_keys_remap - - -def convert_flux2_transformer_to_diffusers(original_state_dict: Dict[str, torch.Tensor], model_type: str): - config, rename_dict, special_keys_remap = get_flux2_transformer_config(model_type) - - diffusers_config = config["diffusers_config"] - - with init_empty_weights(): - transformer = Flux2Transformer2DModel.from_config(diffusers_config) - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict(original_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - transformer.load_state_dict(original_state_dict, strict=True, assign=True) - return transformer - - -def main(args): - if args.vae: - original_vae_ckpt = load_original_checkpoint(args, filename=args.vae_filename) - vae = AutoencoderKLFlux2() - converted_vae_state_dict = convert_flux2_vae_checkpoint_to_diffusers(original_vae_ckpt, vae.config) - vae.load_state_dict(converted_vae_state_dict, strict=True) - if not args.full_pipe: - vae_dtype = torch.bfloat16 if args.vae_dtype == "bf16" else torch.float32 - vae.to(vae_dtype).save_pretrained(f"{args.output_path}/vae") - - if args.dit: - original_dit_ckpt = load_original_checkpoint(args, filename=args.dit_filename) - - if "klein-4b" in args.dit_filename: - model_type = "klein-4b" - elif "klein-9b" in args.dit_filename: - model_type = "klein-9b" - else: - model_type = "flux2-dev" - transformer = convert_flux2_transformer_to_diffusers(original_dit_ckpt, model_type) - if not args.full_pipe: - dit_dtype = torch.bfloat16 if args.dit_dtype == "bf16" else torch.float32 - transformer.to(dit_dtype).save_pretrained(f"{args.output_path}/transformer") - - if args.full_pipe: - tokenizer_id = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" - text_encoder_id = "mistralai/Mistral-Small-3.2-24B-Instruct-2506" - generate_config = GenerationConfig.from_pretrained(text_encoder_id) - generate_config.do_sample = True - text_encoder = Mistral3ForConditionalGeneration.from_pretrained( - text_encoder_id, generation_config=generate_config, torch_dtype=torch.bfloat16 - ) - tokenizer = AutoProcessor.from_pretrained(tokenizer_id) - scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - "black-forest-labs/FLUX.1-dev", subfolder="scheduler" - ) - - if_distilled = "base" not in args.dit_filename - - pipe = Flux2Pipeline( - vae=vae, - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - if_distilled=if_distilled, - ) - pipe.save_pretrained(args.output_path) - - -if __name__ == "__main__": - main(args) diff --git a/scripts/convert_flux_to_diffusers.py b/scripts/convert_flux_to_diffusers.py deleted file mode 100644 index ec31d842d4db..000000000000 --- a/scripts/convert_flux_to_diffusers.py +++ /dev/null @@ -1,308 +0,0 @@ -import argparse -from contextlib import nullcontext - -import safetensors.torch -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download - -from diffusers import AutoencoderKL, FluxTransformer2DModel -from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint -from diffusers.utils.import_utils import is_accelerate_available - - -""" -# Transformer - -python scripts/convert_flux_to_diffusers.py \ ---original_state_dict_repo_id "black-forest-labs/FLUX.1-schnell" \ ---filename "flux1-schnell.sft" ---output_path "flux-schnell" \ ---transformer -""" - -""" -# VAE - -python scripts/convert_flux_to_diffusers.py \ ---original_state_dict_repo_id "black-forest-labs/FLUX.1-schnell" \ ---filename "ae.sft" ---output_path "flux-schnell" \ ---vae -""" - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - -parser = argparse.ArgumentParser() -parser.add_argument("--original_state_dict_repo_id", default=None, type=str) -parser.add_argument("--filename", default="flux.safetensors", type=str) -parser.add_argument("--checkpoint_path", default=None, type=str) -parser.add_argument("--in_channels", type=int, default=64) -parser.add_argument("--out_channels", type=int, default=None) -parser.add_argument("--vae", action="store_true") -parser.add_argument("--transformer", action="store_true") -parser.add_argument("--output_path", type=str) -parser.add_argument("--dtype", type=str, default="bf16") - -args = parser.parse_args() -dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - - -def load_original_checkpoint(args): - if args.original_state_dict_repo_id is not None: - ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=args.filename) - elif args.checkpoint_path is not None: - ckpt_path = args.checkpoint_path - else: - raise ValueError(" please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") - - original_state_dict = safetensors.torch.load_file(ckpt_path) - return original_state_dict - - -# in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; -# while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation -def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_flux_transformer_checkpoint_to_diffusers( - original_state_dict, num_layers, num_single_layers, inner_dim, mlp_ratio=4.0 -): - converted_state_dict = {} - - ## time_text_embed.timestep_embedder <- time_in - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "time_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "time_in.in_layer.bias" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "time_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "time_in.out_layer.bias" - ) - - ## time_text_embed.text_embedder <- vector_in - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop( - "vector_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop( - "vector_in.in_layer.bias" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop( - "vector_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop( - "vector_in.out_layer.bias" - ) - - # guidance - has_guidance = any("guidance" in k for k in original_state_dict) - if has_guidance: - converted_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = original_state_dict.pop( - "guidance_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = original_state_dict.pop( - "guidance_in.in_layer.bias" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = original_state_dict.pop( - "guidance_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = original_state_dict.pop( - "guidance_in.out_layer.bias" - ) - - # context_embedder - converted_state_dict["context_embedder.weight"] = original_state_dict.pop("txt_in.weight") - converted_state_dict["context_embedder.bias"] = original_state_dict.pop("txt_in.bias") - - # x_embedder - converted_state_dict["x_embedder.weight"] = original_state_dict.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = original_state_dict.pop("img_in.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - # norms. - ## norm1 - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.bias" - ) - ## norm1_context - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.bias" - ) - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0 - ) - context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - # qk_norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.key_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.key_norm.scale" - ) - # ff img_mlp - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.0.bias" - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.2.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.0.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.2.bias" - ) - # output projections. - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.bias" - ) - - # single transformer blocks - for i in range(num_single_layers): - block_prefix = f"single_transformer_blocks.{i}." - # norm.linear <- single_blocks.0.modulation.lin - converted_state_dict[f"{block_prefix}norm.linear.weight"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm.linear.bias"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.bias" - ) - # Q, K, V, mlp - mlp_hidden_dim = int(inner_dim * mlp_ratio) - split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) - q, k, v, mlp = torch.split(original_state_dict.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0) - q_bias, k_bias, v_bias, mlp_bias = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp]) - converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias]) - # qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.key_norm.scale" - ) - # output projections. - converted_state_dict[f"{block_prefix}proj_out.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.weight" - ) - converted_state_dict[f"{block_prefix}proj_out.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.bias" - ) - - converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.weight") - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.bias") - ) - - return converted_state_dict - - -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__": - main(args) diff --git a/scripts/convert_gligen_to_diffusers.py b/scripts/convert_gligen_to_diffusers.py deleted file mode 100644 index 83c1f928e407..000000000000 --- a/scripts/convert_gligen_to_diffusers.py +++ /dev/null @@ -1,581 +0,0 @@ -import argparse -import re - -import torch -import yaml -from transformers import ( - CLIPProcessor, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers import ( - AutoencoderKL, - DDIMScheduler, - StableDiffusionGLIGENPipeline, - StableDiffusionGLIGENTextImagePipeline, - UNet2DConditionModel, -) -from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( - assign_to_checkpoint, - conv_attn_to_linear, - protected, - renew_attention_paths, - renew_resnet_paths, - renew_vae_attention_paths, - renew_vae_resnet_paths, - shave_segments, - textenc_conversion_map, - textenc_pattern, -) - - -def convert_open_clip_checkpoint(checkpoint): - checkpoint = checkpoint["text_encoder"] - text_model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") - - keys = list(checkpoint.keys()) - - text_model_dict = {} - - if "cond_stage_model.model.text_projection" in checkpoint: - d_model = int(checkpoint["cond_stage_model.model.text_projection"].shape[0]) - else: - d_model = 1024 - - for key in keys: - if "resblocks.23" in key: # Diffusers drops the final layer and only uses the penultimate layer - continue - if key in textenc_conversion_map: - text_model_dict[textenc_conversion_map[key]] = checkpoint[key] - # if key.startswith("cond_stage_model.model.transformer."): - new_key = key[len("transformer.") :] - if new_key.endswith(".in_proj_weight"): - new_key = new_key[: -len(".in_proj_weight")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.weight"] = checkpoint[key][:d_model, :] - text_model_dict[new_key + ".k_proj.weight"] = checkpoint[key][d_model : d_model * 2, :] - text_model_dict[new_key + ".v_proj.weight"] = checkpoint[key][d_model * 2 :, :] - elif new_key.endswith(".in_proj_bias"): - new_key = new_key[: -len(".in_proj_bias")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.bias"] = checkpoint[key][:d_model] - text_model_dict[new_key + ".k_proj.bias"] = checkpoint[key][d_model : d_model * 2] - text_model_dict[new_key + ".v_proj.bias"] = checkpoint[key][d_model * 2 :] - else: - if key != "transformer.text_model.embeddings.position_ids": - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - - text_model_dict[new_key] = checkpoint[key] - - if key == "transformer.text_model.embeddings.token_embedding.weight": - text_model_dict["text_model.embeddings.token_embedding.weight"] = checkpoint[key] - - text_model_dict.pop("text_model.embeddings.transformer.text_model.embeddings.token_embedding.weight") - - text_model.load_state_dict(text_model_dict) - - return text_model - - -def convert_gligen_vae_checkpoint(checkpoint, config): - checkpoint = checkpoint["autoencoder"] - vae_state_dict = {} - vae_key = "first_stage_model." - keys = list(checkpoint.keys()) - for key in keys: - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for key in new_checkpoint.keys(): - if "encoder.mid_block.attentions.0" in key or "decoder.mid_block.attentions.0" in key: - if "query" in key: - new_checkpoint[key.replace("query", "to_q")] = new_checkpoint.pop(key) - if "value" in key: - new_checkpoint[key.replace("value", "to_v")] = new_checkpoint.pop(key) - if "key" in key: - new_checkpoint[key.replace("key", "to_k")] = new_checkpoint.pop(key) - if "proj_attn" in key: - new_checkpoint[key.replace("proj_attn", "to_out.0")] = new_checkpoint.pop(key) - - return new_checkpoint - - -def convert_gligen_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - unet_state_dict = {} - checkpoint = checkpoint["model"] - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has bot EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - for key in keys: - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - for key in keys: - if "position_net" in key: - new_checkpoint[key] = unet_state_dict[key] - - return new_checkpoint - - -def create_vae_config(original_config, image_size: int): - vae_params = original_config["autoencoder"]["params"]["ddconfig"] - _ = original_config["autoencoder"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - } - - return config - - -def create_unet_config(original_config, image_size: int, attention_type): - unet_params = original_config["model"]["params"] - vae_params = original_config["autoencoder"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None - use_linear_projection = ( - unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False - ) - if use_linear_projection: - if head_dim is None: - head_dim = [5, 10, 20, 20] - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": unet_params["context_dim"], - "attention_head_dim": head_dim, - "use_linear_projection": use_linear_projection, - "attention_type": attention_type, - } - - return config - - -def convert_gligen_to_diffusers( - checkpoint_path: str, - original_config_file: str, - attention_type: str, - image_size: int = 512, - extract_ema: bool = False, - num_in_channels: int = None, - device: str = None, -): - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - checkpoint = torch.load(checkpoint_path, map_location=device) - else: - checkpoint = torch.load(checkpoint_path, map_location=device) - - if "global_step" in checkpoint: - checkpoint["global_step"] - else: - print("global_step key not found in model") - - original_config = yaml.safe_load(original_config_file) - - if num_in_channels is not None: - original_config["model"]["params"]["in_channels"] = num_in_channels - - num_train_timesteps = original_config["diffusion"]["params"]["timesteps"] - beta_start = original_config["diffusion"]["params"]["linear_start"] - beta_end = original_config["diffusion"]["params"]["linear_end"] - - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type="epsilon", - ) - - # Convert the UNet2DConditionalModel model - unet_config = create_unet_config(original_config, image_size, attention_type) - unet = UNet2DConditionModel(**unet_config) - - converted_unet_checkpoint = convert_gligen_unet_checkpoint( - checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema - ) - - unet.load_state_dict(converted_unet_checkpoint) - - # Convert the VAE model - vae_config = create_vae_config(original_config, image_size) - converted_vae_checkpoint = convert_gligen_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - - # Convert the text model - text_encoder = convert_open_clip_checkpoint(checkpoint) - tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") - - if attention_type == "gated-text-image": - image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") - processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") - - pipe = StableDiffusionGLIGENTextImagePipeline( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - image_encoder=image_encoder, - processor=processor, - unet=unet, - scheduler=scheduler, - safety_checker=None, - feature_extractor=None, - ) - elif attention_type == "gated": - pipe = StableDiffusionGLIGENPipeline( - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - safety_checker=None, - feature_extractor=None, - ) - - return pipe - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--original_config_file", - default=None, - type=str, - required=True, - help="The YAML config file corresponding to the gligen architecture.", - ) - parser.add_argument( - "--num_in_channels", - default=None, - type=int, - help="The number of input channels. If `None` number of input channels will be automatically inferred.", - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument( - "--attention_type", - default=None, - type=str, - required=True, - help="Type of attention ex: gated or gated-text-image", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--device", type=str, help="Device to use.") - parser.add_argument("--half", action="store_true", help="Save weights in half precision.") - - args = parser.parse_args() - - pipe = convert_gligen_to_diffusers( - checkpoint_path=args.checkpoint_path, - original_config_file=args.original_config_file, - attention_type=args.attention_type, - extract_ema=args.extract_ema, - num_in_channels=args.num_in_channels, - device=args.device, - ) - - if args.half: - pipe.to(dtype=torch.float16) - - pipe.save_pretrained(args.dump_path) diff --git a/scripts/convert_hunyuan_image_to_diffusers.py b/scripts/convert_hunyuan_image_to_diffusers.py deleted file mode 100644 index c41e934cc3d4..000000000000 --- a/scripts/convert_hunyuan_image_to_diffusers.py +++ /dev/null @@ -1,1044 +0,0 @@ -import argparse -import logging - -import torch -from safetensors import safe_open - -from diffusers import AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, HunyuanImageTransformer2DModel - - -logger = logging.getLogger(__name__) # pylint: disable=invalid-name - - -""" -Usage examples -============== - -python scripts/convert_hunyuan_image_to_diffusers.py \ - --model_type hunyuanimage2.1 \ - --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1.safetensors" \ - --vae_checkpoint_path "HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ - --output_path "/raid/yiyi/test-hy21-diffusers" \ - --dtype fp32 - -python scripts/convert_hunyuan_image_to_diffusers.py \ - --model_type hunyuanimage2.1-distilled \ - --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1-distilled.safetensors" \ - --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ - --output_path "/raid/yiyi/test-hy21-distilled-diffusers" \ - --dtype fp32 - - -python scripts/convert_hunyuan_image_to_diffusers.py \ - --model_type hunyuanimage-refiner \ - --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage-refiner.safetensors" \ - --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_refiner/pytorch_model.pt" \ - --output_path "/raid/yiyi/test-hy2-refiner-diffusers" \ - --dtype fp32 -""" - -parser = argparse.ArgumentParser() -parser.add_argument( - "--model_type", type=str, default=None -) # hunyuanimage2.1, hunyuanimage2.1-distilled, hunyuanimage-refiner -parser.add_argument("--transformer_checkpoint_path", default=None, type=str) # ckpts/dit/hunyuanimage2.1.safetensors -parser.add_argument("--vae_checkpoint_path", default=None, type=str) # ckpts/vae/vae_2_1/pytorch_model.ckpt -parser.add_argument("--output_path", type=str) -parser.add_argument("--dtype", type=str, default="fp32") - -args = parser.parse_args() -dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - - -# copied from https://github.com/Tencent-Hunyuan/HunyuanImage-2.1/hyimage/models/hunyuan/modules/hunyuanimage_dit.py#L21 -def convert_hunyuan_dict_for_tensor_parallel(state_dict): - """ - Convert a Hunyuan model state dict to be compatible with tensor parallel architectures. - - Args: - state_dict: Original state dict - - Returns: - new_dict: Converted state dict - """ - new_dict = {} - for k, w in state_dict.items(): - if k.startswith("double_blocks") and "attn_qkv.weight" in k: - hidden_size = w.shape[1] - k1 = k.replace("attn_qkv.weight", "attn_q.weight") - w1 = w[:hidden_size, :] - new_dict[k1] = w1 - k2 = k.replace("attn_qkv.weight", "attn_k.weight") - w2 = w[hidden_size : 2 * hidden_size, :] - new_dict[k2] = w2 - k3 = k.replace("attn_qkv.weight", "attn_v.weight") - w3 = w[-hidden_size:, :] - new_dict[k3] = w3 - elif k.startswith("double_blocks") and "attn_qkv.bias" in k: - hidden_size = w.shape[0] // 3 - k1 = k.replace("attn_qkv.bias", "attn_q.bias") - w1 = w[:hidden_size] - new_dict[k1] = w1 - k2 = k.replace("attn_qkv.bias", "attn_k.bias") - w2 = w[hidden_size : 2 * hidden_size] - new_dict[k2] = w2 - k3 = k.replace("attn_qkv.bias", "attn_v.bias") - w3 = w[-hidden_size:] - new_dict[k3] = w3 - elif k.startswith("single_blocks") and "linear1" in k: - hidden_size = state_dict[k.replace("linear1", "linear2")].shape[0] - k1 = k.replace("linear1", "linear1_q") - w1 = w[:hidden_size] - new_dict[k1] = w1 - k2 = k.replace("linear1", "linear1_k") - w2 = w[hidden_size : 2 * hidden_size] - new_dict[k2] = w2 - k3 = k.replace("linear1", "linear1_v") - w3 = w[2 * hidden_size : 3 * hidden_size] - new_dict[k3] = w3 - k4 = k.replace("linear1", "linear1_mlp") - w4 = w[3 * hidden_size :] - new_dict[k4] = w4 - elif k.startswith("single_blocks") and "linear2" in k: - k1 = k.replace("linear2", "linear2.fc") - new_dict[k1] = w - else: - new_dict[k] = w - return new_dict - - -def load_original_vae_checkpoint(args): - # "ckpts/vae/vae_2_1/pytorch_model.ckpt" - state_dict = torch.load(args.vae_checkpoint_path) - - if "state_dict" in state_dict: - state_dict = state_dict["state_dict"] - vae_state_dict = {} - for k, v in state_dict.items(): - if k.startswith("vae."): - vae_state_dict[k.replace("vae.", "")] = v - - for k, v in vae_state_dict.items(): - if "weight" in k: - if len(v.shape) == 5 and v.shape[2] == 1: - vae_state_dict[k] = v.squeeze(2) - else: - vae_state_dict[k] = v - else: - vae_state_dict[k] = v - return vae_state_dict - - -def load_original_refiner_vae_checkpoint(args): - # "ckpts/vae/vae_refiner/pytorch_model.pt" - state_dict = torch.load(args.vae_checkpoint_path) - - if "state_dict" in state_dict: - state_dict = state_dict["state_dict"] - vae_state_dict = {} - for k, v in state_dict.items(): - if k.startswith("vae."): - vae_state_dict[k.replace("vae.", "")] = v - return vae_state_dict - - -def load_original_transformer_checkpoint(args): - # ckpts/dit/hunyuanimage-refiner.safetensors" - # ckpts/dit/hunyuanimage2.1.safetensors" - state_dict = {} - with safe_open(args.transformer_checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - state_dict[key] = f.get_tensor(key) - if args.model_type == "hunyuanimage-2.1": - state_dict = convert_hunyuan_dict_for_tensor_parallel(state_dict) - return state_dict - - -def convert_hunyuan_image_transformer_checkpoint_to_diffusers( - original_state_dict, use_byt5=True, guidance_distilled=False, use_meanflow=False -): - converted_state_dict = {} - - # 1. byt5_in -> context_embedder_2 - if use_byt5: - converted_state_dict["context_embedder_2.norm.weight"] = original_state_dict.pop("byt5_in.layernorm.weight") - converted_state_dict["context_embedder_2.norm.bias"] = original_state_dict.pop("byt5_in.layernorm.bias") - converted_state_dict["context_embedder_2.linear_1.weight"] = original_state_dict.pop("byt5_in.fc1.weight") - converted_state_dict["context_embedder_2.linear_1.bias"] = original_state_dict.pop("byt5_in.fc1.bias") - converted_state_dict["context_embedder_2.linear_2.weight"] = original_state_dict.pop("byt5_in.fc2.weight") - converted_state_dict["context_embedder_2.linear_2.bias"] = original_state_dict.pop("byt5_in.fc2.bias") - converted_state_dict["context_embedder_2.linear_3.weight"] = original_state_dict.pop("byt5_in.fc3.weight") - converted_state_dict["context_embedder_2.linear_3.bias"] = original_state_dict.pop("byt5_in.fc3.bias") - - # 2. img_in -> x_embedder - converted_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("img_in.proj.weight") - converted_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("img_in.proj.bias") - - # 3. txt_in -> context_embedder (complex mapping) - # txt_in.input_embedder -> context_embedder.proj_in - converted_state_dict["context_embedder.proj_in.weight"] = original_state_dict.pop("txt_in.input_embedder.weight") - converted_state_dict["context_embedder.proj_in.bias"] = original_state_dict.pop("txt_in.input_embedder.bias") - - # txt_in.t_embedder -> context_embedder.time_text_embed.timestep_embedder - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.weight"] = ( - original_state_dict.pop("txt_in.t_embedder.mlp.0.weight") - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "txt_in.t_embedder.mlp.0.bias" - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.weight"] = ( - original_state_dict.pop("txt_in.t_embedder.mlp.2.weight") - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "txt_in.t_embedder.mlp.2.bias" - ) - - # txt_in.c_embedder -> context_embedder.time_text_embed.text_embedder - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop( - "txt_in.c_embedder.linear_1.weight" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop( - "txt_in.c_embedder.linear_1.bias" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop( - "txt_in.c_embedder.linear_2.weight" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop( - "txt_in.c_embedder.linear_2.bias" - ) - - # txt_in.individual_token_refiner -> context_embedder.token_refiner - for i in range(2): # 2 refiner blocks - block_prefix = f"context_embedder.token_refiner.refiner_blocks.{i}." - # norm1 - converted_state_dict[f"{block_prefix}norm1.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.norm1.weight" - ) - converted_state_dict[f"{block_prefix}norm1.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.norm1.bias" - ) - # norm2 - converted_state_dict[f"{block_prefix}norm2.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.norm2.weight" - ) - converted_state_dict[f"{block_prefix}norm2.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.norm2.bias" - ) - - # Split QKV - qkv_weight = original_state_dict.pop(f"txt_in.individual_token_refiner.blocks.{i}.self_attn_qkv.weight") - qkv_bias = original_state_dict.pop(f"txt_in.individual_token_refiner.blocks.{i}.self_attn_qkv.bias") - q_weight, k_weight, v_weight = torch.chunk(qkv_weight, 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) - - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = q_weight - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = q_bias - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = k_weight - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = k_bias - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = v_weight - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = v_bias - - # attn projection - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.self_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.self_attn_proj.bias" - ) - - # MLP - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.mlp.fc2.bias" - ) - - # norm_out - converted_state_dict[f"{block_prefix}norm_out.linear.weight"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.adaLN_modulation.1.weight" - ) - converted_state_dict[f"{block_prefix}norm_out.linear.bias"] = original_state_dict.pop( - f"txt_in.individual_token_refiner.blocks.{i}.adaLN_modulation.1.bias" - ) - - # 4. time_in -> time_text_embed.timestep_embedder - converted_state_dict["time_guidance_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "time_in.mlp.0.weight" - ) - converted_state_dict["time_guidance_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "time_in.mlp.0.bias" - ) - converted_state_dict["time_guidance_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "time_in.mlp.2.weight" - ) - converted_state_dict["time_guidance_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "time_in.mlp.2.bias" - ) - - # time_r_in -> time_guidance_embed.timestep_r_embedder - if use_meanflow: - converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_1.weight"] = original_state_dict.pop( - "time_r_in.mlp.0.weight" - ) - converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_1.bias"] = original_state_dict.pop( - "time_r_in.mlp.0.bias" - ) - converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_2.weight"] = original_state_dict.pop( - "time_r_in.mlp.2.weight" - ) - converted_state_dict["time_guidance_embed.timestep_embedder_r.linear_2.bias"] = original_state_dict.pop( - "time_r_in.mlp.2.bias" - ) - - # guidance_in -> time_guidance_embed.guidance_embedder - if guidance_distilled: - converted_state_dict["time_guidance_embed.guidance_embedder.linear_1.weight"] = original_state_dict.pop( - "guidance_in.mlp.0.weight" - ) - converted_state_dict["time_guidance_embed.guidance_embedder.linear_1.bias"] = original_state_dict.pop( - "guidance_in.mlp.0.bias" - ) - converted_state_dict["time_guidance_embed.guidance_embedder.linear_2.weight"] = original_state_dict.pop( - "guidance_in.mlp.2.weight" - ) - converted_state_dict["time_guidance_embed.guidance_embedder.linear_2.bias"] = original_state_dict.pop( - "guidance_in.mlp.2.bias" - ) - - # 5. double_blocks -> transformer_blocks - for i in range(20): # 20 double blocks - block_prefix = f"transformer_blocks.{i}." - - # norm1 (img_mod) - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.linear.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.linear.bias" - ) - - # norm1_context (txt_mod) - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.linear.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.linear.bias" - ) - - # img attention - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_q.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_q.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_k.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_k.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_v.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_v.bias" - ) - - # img attention norms - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_q_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_k_norm.weight" - ) - - # img attention projection - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn_proj.bias" - ) - - # img MLP - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.fc2.bias" - ) - - # txt attention (additional projections) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_q.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_q.bias" - ) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_k.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_k.bias" - ) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_v.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_v.bias" - ) - - # txt attention norms - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_q_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_k_norm.weight" - ) - - # txt attention projection - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn_proj.bias" - ) - - # txt MLP (ff_context) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.fc2.bias" - ) - - # 6. single_blocks -> single_transformer_blocks - for i in range(40): # 40 single blocks - block_prefix = f"single_transformer_blocks.{i}." - - # norm - converted_state_dict[f"{block_prefix}norm.linear.weight"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.linear.weight" - ) - converted_state_dict[f"{block_prefix}norm.linear.bias"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.linear.bias" - ) - - # attention Q, K, V - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_q.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_q.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_k.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_k.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_v.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_v.bias" - ) - - # attention norms - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"single_blocks.{i}.q_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"single_blocks.{i}.k_norm.weight" - ) - - # MLP projection - converted_state_dict[f"{block_prefix}proj_mlp.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_mlp.weight" - ) - converted_state_dict[f"{block_prefix}proj_mlp.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear1_mlp.bias" - ) - - # output projection - converted_state_dict[f"{block_prefix}proj_out.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.fc.weight" - ) - converted_state_dict[f"{block_prefix}proj_out.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.fc.bias" - ) - - # 7. final_layer -> norm_out + proj_out - converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - shift_w, scale_w = original_state_dict.pop("final_layer.adaLN_modulation.1.weight").chunk(2, dim=0) - shift_b, scale_b = original_state_dict.pop("final_layer.adaLN_modulation.1.bias").chunk(2, dim=0) - converted_state_dict["norm_out.linear.weight"] = torch.cat([scale_w, shift_w], dim=0) - converted_state_dict["norm_out.linear.bias"] = torch.cat([scale_b, shift_b], dim=0) - - return converted_state_dict, original_state_dict - - -def convert_hunyuan_image_vae_checkpoint_to_diffusers( - original_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 -): - """Convert original VAE state dict to Diffusers format.""" - converted = {} - - # 1. Encoder - # 1.1 conv_in - converted["encoder.conv_in.weight"] = original_state_dict.pop("encoder.conv_in.weight") - converted["encoder.conv_in.bias"] = original_state_dict.pop("encoder.conv_in.bias") - - # 1.2 down blocks - diffusers_block_idx = 0 - - for block_index in range(len(block_out_channels)): - for resnet_block_index in range(layers_per_block): - orig_prefix = f"encoder.down.{block_index}.block.{resnet_block_index}" - diff_prefix = f"encoder.down_blocks.{diffusers_block_idx}" - - # resnet blocks - converted[f"{diff_prefix}.norm1.weight"] = original_state_dict.pop(f"{orig_prefix}.norm1.weight") - converted[f"{diff_prefix}.norm1.bias"] = original_state_dict.pop(f"{orig_prefix}.norm1.bias") - converted[f"{diff_prefix}.conv1.weight"] = original_state_dict.pop(f"{orig_prefix}.conv1.weight") - converted[f"{diff_prefix}.conv1.bias"] = original_state_dict.pop(f"{orig_prefix}.conv1.bias") - converted[f"{diff_prefix}.norm2.weight"] = original_state_dict.pop(f"{orig_prefix}.norm2.weight") - converted[f"{diff_prefix}.norm2.bias"] = original_state_dict.pop(f"{orig_prefix}.norm2.bias") - converted[f"{diff_prefix}.conv2.weight"] = original_state_dict.pop(f"{orig_prefix}.conv2.weight") - converted[f"{diff_prefix}.conv2.bias"] = original_state_dict.pop(f"{orig_prefix}.conv2.bias") - - diffusers_block_idx += 1 - - # downsample blocks - if f"encoder.down.{block_index}.downsample.conv.weight" in original_state_dict: - converted[f"encoder.down_blocks.{diffusers_block_idx}.conv.weight"] = original_state_dict.pop( - f"encoder.down.{block_index}.downsample.conv.weight" - ) - converted[f"encoder.down_blocks.{diffusers_block_idx}.conv.bias"] = original_state_dict.pop( - f"encoder.down.{block_index}.downsample.conv.bias" - ) - diffusers_block_idx += 1 - - # 1.3 mid block - converted["encoder.mid_block.resnets.0.norm1.weight"] = original_state_dict.pop("encoder.mid.block_1.norm1.weight") - converted["encoder.mid_block.resnets.0.norm1.bias"] = original_state_dict.pop("encoder.mid.block_1.norm1.bias") - converted["encoder.mid_block.resnets.0.conv1.weight"] = original_state_dict.pop("encoder.mid.block_1.conv1.weight") - converted["encoder.mid_block.resnets.0.conv1.bias"] = original_state_dict.pop("encoder.mid.block_1.conv1.bias") - converted["encoder.mid_block.resnets.0.norm2.weight"] = original_state_dict.pop("encoder.mid.block_1.norm2.weight") - converted["encoder.mid_block.resnets.0.norm2.bias"] = original_state_dict.pop("encoder.mid.block_1.norm2.bias") - converted["encoder.mid_block.resnets.0.conv2.weight"] = original_state_dict.pop("encoder.mid.block_1.conv2.weight") - converted["encoder.mid_block.resnets.0.conv2.bias"] = original_state_dict.pop("encoder.mid.block_1.conv2.bias") - - converted["encoder.mid_block.resnets.1.norm1.weight"] = original_state_dict.pop("encoder.mid.block_2.norm1.weight") - converted["encoder.mid_block.resnets.1.norm1.bias"] = original_state_dict.pop("encoder.mid.block_2.norm1.bias") - converted["encoder.mid_block.resnets.1.conv1.weight"] = original_state_dict.pop("encoder.mid.block_2.conv1.weight") - converted["encoder.mid_block.resnets.1.conv1.bias"] = original_state_dict.pop("encoder.mid.block_2.conv1.bias") - converted["encoder.mid_block.resnets.1.norm2.weight"] = original_state_dict.pop("encoder.mid.block_2.norm2.weight") - converted["encoder.mid_block.resnets.1.norm2.bias"] = original_state_dict.pop("encoder.mid.block_2.norm2.bias") - converted["encoder.mid_block.resnets.1.conv2.weight"] = original_state_dict.pop("encoder.mid.block_2.conv2.weight") - converted["encoder.mid_block.resnets.1.conv2.bias"] = original_state_dict.pop("encoder.mid.block_2.conv2.bias") - - converted["encoder.mid_block.attentions.0.norm.weight"] = original_state_dict.pop("encoder.mid.attn_1.norm.weight") - converted["encoder.mid_block.attentions.0.norm.bias"] = original_state_dict.pop("encoder.mid.attn_1.norm.bias") - converted["encoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("encoder.mid.attn_1.q.weight") - converted["encoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("encoder.mid.attn_1.q.bias") - converted["encoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("encoder.mid.attn_1.k.weight") - converted["encoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("encoder.mid.attn_1.k.bias") - converted["encoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("encoder.mid.attn_1.v.weight") - converted["encoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("encoder.mid.attn_1.v.bias") - converted["encoder.mid_block.attentions.0.proj.weight"] = original_state_dict.pop( - "encoder.mid.attn_1.proj_out.weight" - ) - converted["encoder.mid_block.attentions.0.proj.bias"] = original_state_dict.pop("encoder.mid.attn_1.proj_out.bias") - - # 1.4 encoder output - converted["encoder.norm_out.weight"] = original_state_dict.pop("encoder.norm_out.weight") - converted["encoder.norm_out.bias"] = original_state_dict.pop("encoder.norm_out.bias") - converted["encoder.conv_out.weight"] = original_state_dict.pop("encoder.conv_out.weight") - converted["encoder.conv_out.bias"] = original_state_dict.pop("encoder.conv_out.bias") - - # 2. Decoder - # 2.1 conv_in - converted["decoder.conv_in.weight"] = original_state_dict.pop("decoder.conv_in.weight") - converted["decoder.conv_in.bias"] = original_state_dict.pop("decoder.conv_in.bias") - - # 2.2 mid block - converted["decoder.mid_block.resnets.0.norm1.weight"] = original_state_dict.pop("decoder.mid.block_1.norm1.weight") - converted["decoder.mid_block.resnets.0.norm1.bias"] = original_state_dict.pop("decoder.mid.block_1.norm1.bias") - converted["decoder.mid_block.resnets.0.conv1.weight"] = original_state_dict.pop("decoder.mid.block_1.conv1.weight") - converted["decoder.mid_block.resnets.0.conv1.bias"] = original_state_dict.pop("decoder.mid.block_1.conv1.bias") - converted["decoder.mid_block.resnets.0.norm2.weight"] = original_state_dict.pop("decoder.mid.block_1.norm2.weight") - converted["decoder.mid_block.resnets.0.norm2.bias"] = original_state_dict.pop("decoder.mid.block_1.norm2.bias") - converted["decoder.mid_block.resnets.0.conv2.weight"] = original_state_dict.pop("decoder.mid.block_1.conv2.weight") - converted["decoder.mid_block.resnets.0.conv2.bias"] = original_state_dict.pop("decoder.mid.block_1.conv2.bias") - - converted["decoder.mid_block.resnets.1.norm1.weight"] = original_state_dict.pop("decoder.mid.block_2.norm1.weight") - converted["decoder.mid_block.resnets.1.norm1.bias"] = original_state_dict.pop("decoder.mid.block_2.norm1.bias") - converted["decoder.mid_block.resnets.1.conv1.weight"] = original_state_dict.pop("decoder.mid.block_2.conv1.weight") - converted["decoder.mid_block.resnets.1.conv1.bias"] = original_state_dict.pop("decoder.mid.block_2.conv1.bias") - converted["decoder.mid_block.resnets.1.norm2.weight"] = original_state_dict.pop("decoder.mid.block_2.norm2.weight") - converted["decoder.mid_block.resnets.1.norm2.bias"] = original_state_dict.pop("decoder.mid.block_2.norm2.bias") - converted["decoder.mid_block.resnets.1.conv2.weight"] = original_state_dict.pop("decoder.mid.block_2.conv2.weight") - converted["decoder.mid_block.resnets.1.conv2.bias"] = original_state_dict.pop("decoder.mid.block_2.conv2.bias") - - converted["decoder.mid_block.attentions.0.norm.weight"] = original_state_dict.pop("decoder.mid.attn_1.norm.weight") - converted["decoder.mid_block.attentions.0.norm.bias"] = original_state_dict.pop("decoder.mid.attn_1.norm.bias") - converted["decoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("decoder.mid.attn_1.q.weight") - converted["decoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("decoder.mid.attn_1.q.bias") - converted["decoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("decoder.mid.attn_1.k.weight") - converted["decoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("decoder.mid.attn_1.k.bias") - converted["decoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("decoder.mid.attn_1.v.weight") - converted["decoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("decoder.mid.attn_1.v.bias") - converted["decoder.mid_block.attentions.0.proj.weight"] = original_state_dict.pop( - "decoder.mid.attn_1.proj_out.weight" - ) - converted["decoder.mid_block.attentions.0.proj.bias"] = original_state_dict.pop("decoder.mid.attn_1.proj_out.bias") - - # 2.3 up blocks - diffusers_block_idx = 0 - for up_block_index in range(len(block_out_channels)): - # resnet blocks - for resnet_block_index in range(layers_per_block + 1): - orig_prefix = f"decoder.up.{up_block_index}.block.{resnet_block_index}" - diff_prefix = f"decoder.up_blocks.{diffusers_block_idx}" - - converted[f"{diff_prefix}.norm1.weight"] = original_state_dict.pop(f"{orig_prefix}.norm1.weight") - converted[f"{diff_prefix}.norm1.bias"] = original_state_dict.pop(f"{orig_prefix}.norm1.bias") - converted[f"{diff_prefix}.conv1.weight"] = original_state_dict.pop(f"{orig_prefix}.conv1.weight") - converted[f"{diff_prefix}.conv1.bias"] = original_state_dict.pop(f"{orig_prefix}.conv1.bias") - converted[f"{diff_prefix}.norm2.weight"] = original_state_dict.pop(f"{orig_prefix}.norm2.weight") - converted[f"{diff_prefix}.norm2.bias"] = original_state_dict.pop(f"{orig_prefix}.norm2.bias") - converted[f"{diff_prefix}.conv2.weight"] = original_state_dict.pop(f"{orig_prefix}.conv2.weight") - converted[f"{diff_prefix}.conv2.bias"] = original_state_dict.pop(f"{orig_prefix}.conv2.bias") - - diffusers_block_idx += 1 - - # upsample blocks - if f"decoder.up.{up_block_index}.upsample.conv.weight" in original_state_dict: - converted[f"decoder.up_blocks.{diffusers_block_idx}.conv.weight"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.weight" - ) - converted[f"decoder.up_blocks.{diffusers_block_idx}.conv.bias"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.bias" - ) - diffusers_block_idx += 1 - - # 2.4 decoder output - converted["decoder.norm_out.weight"] = original_state_dict.pop("decoder.norm_out.weight") - converted["decoder.norm_out.bias"] = original_state_dict.pop("decoder.norm_out.bias") - converted["decoder.conv_out.weight"] = original_state_dict.pop("decoder.conv_out.weight") - converted["decoder.conv_out.bias"] = original_state_dict.pop("decoder.conv_out.bias") - - return converted, original_state_dict - - -def convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( - original_state_dict, block_out_channels=[128, 256, 512, 1024, 1024], layers_per_block=2 -): - converted = {} - - # 1. Encoder - # 1.1 conv_in - converted["encoder.conv_in.conv.weight"] = original_state_dict.pop("encoder.conv_in.conv.weight") - converted["encoder.conv_in.conv.bias"] = original_state_dict.pop("encoder.conv_in.conv.bias") - - # 1.2 Down blocks - for down_block_index in range(len(block_out_channels)): # 0 to 4 - # ResNet blocks - for resnet_block_index in range(layers_per_block): # 0 to 1 - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm1.gamma") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( - original_state_dict.pop( - f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.weight" - ) - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.bias") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm2.gamma") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( - original_state_dict.pop( - f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.weight" - ) - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.bias") - ) - - # Downsample (if exists) - if f"encoder.down.{down_block_index}.downsample.conv.conv.weight" in original_state_dict: - converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.weight"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.weight") - ) - converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.bias") - ) - - # 1.3 Mid block - converted["encoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm1.gamma") - converted["encoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_1.conv1.conv.weight" - ) - converted["encoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_1.conv1.conv.bias" - ) - converted["encoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm2.gamma") - converted["encoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_1.conv2.conv.weight" - ) - converted["encoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_1.conv2.conv.bias" - ) - - converted["encoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm1.gamma") - converted["encoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_2.conv1.conv.weight" - ) - converted["encoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_2.conv1.conv.bias" - ) - converted["encoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm2.gamma") - converted["encoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_2.conv2.conv.weight" - ) - converted["encoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_2.conv2.conv.bias" - ) - - # Attention block - converted["encoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("encoder.mid.attn_1.norm.gamma") - converted["encoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("encoder.mid.attn_1.q.weight") - converted["encoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("encoder.mid.attn_1.q.bias") - converted["encoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("encoder.mid.attn_1.k.weight") - converted["encoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("encoder.mid.attn_1.k.bias") - converted["encoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("encoder.mid.attn_1.v.weight") - converted["encoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("encoder.mid.attn_1.v.bias") - converted["encoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( - "encoder.mid.attn_1.proj_out.weight" - ) - converted["encoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( - "encoder.mid.attn_1.proj_out.bias" - ) - - # 1.4 Encoder output - converted["encoder.norm_out.gamma"] = original_state_dict.pop("encoder.norm_out.gamma") - converted["encoder.conv_out.conv.weight"] = original_state_dict.pop("encoder.conv_out.conv.weight") - converted["encoder.conv_out.conv.bias"] = original_state_dict.pop("encoder.conv_out.conv.bias") - - # 2. Decoder - # 2.1 conv_in - converted["decoder.conv_in.conv.weight"] = original_state_dict.pop("decoder.conv_in.conv.weight") - converted["decoder.conv_in.conv.bias"] = original_state_dict.pop("decoder.conv_in.conv.bias") - - # 2.2 Mid block - converted["decoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm1.gamma") - converted["decoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_1.conv1.conv.weight" - ) - converted["decoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_1.conv1.conv.bias" - ) - converted["decoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm2.gamma") - converted["decoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_1.conv2.conv.weight" - ) - converted["decoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_1.conv2.conv.bias" - ) - - converted["decoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm1.gamma") - converted["decoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_2.conv1.conv.weight" - ) - converted["decoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_2.conv1.conv.bias" - ) - converted["decoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm2.gamma") - converted["decoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_2.conv2.conv.weight" - ) - converted["decoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_2.conv2.conv.bias" - ) - - # Decoder attention block - converted["decoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("decoder.mid.attn_1.norm.gamma") - converted["decoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("decoder.mid.attn_1.q.weight") - converted["decoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("decoder.mid.attn_1.q.bias") - converted["decoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("decoder.mid.attn_1.k.weight") - converted["decoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("decoder.mid.attn_1.k.bias") - converted["decoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("decoder.mid.attn_1.v.weight") - converted["decoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("decoder.mid.attn_1.v.bias") - converted["decoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( - "decoder.mid.attn_1.proj_out.weight" - ) - converted["decoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( - "decoder.mid.attn_1.proj_out.bias" - ) - - # 2.3 Up blocks - for up_block_index in range(len(block_out_channels)): # 0 to 5 - # ResNet blocks - for resnet_block_index in range(layers_per_block + 1): # 0 to 2 (decoder has 3 resnets per level) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm1.gamma") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.weight") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.bias") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm2.gamma") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.weight") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.bias") - ) - - # Upsample (if exists) - if f"decoder.up.{up_block_index}.upsample.conv.conv.weight" in original_state_dict: - converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.weight"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.conv.weight" - ) - converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.bias"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.conv.bias" - ) - - # 2.4 Decoder output - converted["decoder.norm_out.gamma"] = original_state_dict.pop("decoder.norm_out.gamma") - converted["decoder.conv_out.conv.weight"] = original_state_dict.pop("decoder.conv_out.conv.weight") - converted["decoder.conv_out.conv.bias"] = original_state_dict.pop("decoder.conv_out.conv.bias") - - return converted, original_state_dict - - -def main(args): - if args.model_type == "hunyuanimage2.1": - original_transformer_state_dict = load_original_transformer_checkpoint(args) - original_vae_state_dict = load_original_vae_checkpoint(args) - - transformer_config = { - "in_channels": 64, - "out_channels": 64, - "num_attention_heads": 28, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "patch_size": (1, 1), - "qk_norm": "rms_norm", - "guidance_embeds": False, - "text_embed_dim": 3584, - "text_embed_2_dim": 1472, - "rope_theta": 256.0, - "rope_axes_dim": (64, 64), - } - - converted_transformer_state_dict, original_transformer_state_dict = ( - convert_hunyuan_image_transformer_checkpoint_to_diffusers( - original_transformer_state_dict, use_byt5=True, guidance_distilled=False - ) - ) - - if original_transformer_state_dict: - logger.warning( - f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" - ) - - transformer = HunyuanImageTransformer2DModel(**transformer_config) - missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) - - if missing_keys: - logger.warning(f"Missing keys for transformer: {missing_keys}") - if unexpected_key: - logger.warning(f"Unexpected keys for transformer: {unexpected_key}") - - transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") - - vae_config_diffusers = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 64, - "block_out_channels": [128, 256, 512, 512, 1024, 1024], - "layers_per_block": 2, - "spatial_compression_ratio": 32, - "sample_size": 384, - "scaling_factor": 0.75289, - "downsample_match_channel": True, - "upsample_match_channel": True, - } - converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( - original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 - ) - if original_vae_state_dict: - logger.warning( - f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" - ) - - vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) - missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) - - if missing_keys: - logger.warning(f"Missing keys for vae: {missing_keys}") - if unexpected_key: - logger.warning(f"Unexpected keys for vae: {unexpected_key}") - - vae.to(dtype).save_pretrained(f"{args.output_path}/vae") - - elif args.model_type == "hunyuanimage2.1-distilled": - original_transformer_state_dict = load_original_transformer_checkpoint(args) - original_vae_state_dict = load_original_vae_checkpoint(args) - - transformer_config = { - "in_channels": 64, - "out_channels": 64, - "num_attention_heads": 28, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "patch_size": (1, 1), - "qk_norm": "rms_norm", - "guidance_embeds": True, - "text_embed_dim": 3584, - "text_embed_2_dim": 1472, - "rope_theta": 256.0, - "rope_axes_dim": (64, 64), - "use_meanflow": True, - } - - converted_transformer_state_dict, original_transformer_state_dict = ( - convert_hunyuan_image_transformer_checkpoint_to_diffusers( - original_transformer_state_dict, use_byt5=True, guidance_distilled=True, use_meanflow=True - ) - ) - - if original_transformer_state_dict: - logger.warning( - f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" - ) - - transformer = HunyuanImageTransformer2DModel(**transformer_config) - missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) - - if missing_keys: - logger.warning(f"Missing keys for transformer: {missing_keys}") - if unexpected_key: - logger.warning(f"Unexpected keys for transformer: {unexpected_key}") - - transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") - - vae_config_diffusers = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 64, - "block_out_channels": [128, 256, 512, 512, 1024, 1024], - "layers_per_block": 2, - "spatial_compression_ratio": 32, - "sample_size": 384, - "scaling_factor": 0.75289, - "downsample_match_channel": True, - "upsample_match_channel": True, - } - converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( - original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 - ) - if original_vae_state_dict: - logger.warning( - f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" - ) - - vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) - missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) - - if missing_keys: - logger.warning(f"Missing keys for vae: {missing_keys}") - if unexpected_key: - logger.warning(f"Unexpected keys for vae: {unexpected_key}") - - vae.to(dtype).save_pretrained(f"{args.output_path}/vae") - - elif args.model_type == "hunyuanimage-refiner": - original_transformer_state_dict = load_original_transformer_checkpoint(args) - original_vae_state_dict = load_original_refiner_vae_checkpoint(args) - - transformer_config = { - "in_channels": 128, - "out_channels": 64, - "num_layers": 20, - "num_single_layers": 40, - "rope_axes_dim": [16, 56, 56], - "num_attention_heads": 26, - "attention_head_dim": 128, - "mlp_ratio": 4, - "patch_size": (1, 1, 1), - "text_embed_dim": 3584, - "guidance_embeds": True, - } - converted_transformer_state_dict, original_transformer_state_dict = ( - convert_hunyuan_image_transformer_checkpoint_to_diffusers( - original_transformer_state_dict, use_byt5=False, guidance_distilled=True - ) - ) - if original_transformer_state_dict: - logger.warning( - f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" - ) - - transformer = HunyuanImageTransformer2DModel(**transformer_config) - missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) - if missing_keys: - logger.warning(f"Missing keys for transformer: {missing_keys}") - if unexpected_key: - logger.warning(f"Unexpected keys for transformer: {unexpected_key}") - - transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") - - vae = AutoencoderKLHunyuanImageRefiner() - - converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( - original_vae_state_dict - ) - if original_vae_state_dict: - logger.warning( - f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" - ) - - missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) - logger.warning(f"Missing keys for vae: {missing_keys}") - logger.warning(f"Unexpected keys for vae: {unexpected_key}") - - vae.to(dtype).save_pretrained(f"{args.output_path}/vae") - - -if __name__ == "__main__": - main(args) diff --git a/scripts/convert_hunyuan_video1_5_to_diffusers.py b/scripts/convert_hunyuan_video1_5_to_diffusers.py deleted file mode 100644 index 89e5cdb16956..000000000000 --- a/scripts/convert_hunyuan_video1_5_to_diffusers.py +++ /dev/null @@ -1,875 +0,0 @@ -import argparse -import json -import os -import pathlib - -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download, snapshot_download -from safetensors.torch import load_file -from transformers import ( - AutoModel, - AutoTokenizer, - SiglipImageProcessor, - SiglipVisionModel, - T5EncoderModel, -) - -from diffusers import ( - AutoencoderKLHunyuanVideo15, - ClassifierFreeGuidance, - FlowMatchEulerDiscreteScheduler, - HunyuanVideo15ImageToVideoPipeline, - HunyuanVideo15Pipeline, - HunyuanVideo15Transformer3DModel, -) - - -# to convert only transformer -""" -python scripts/convert_hunyuan_video1_5_to_diffusers.py \ - --original_state_dict_repo_id tencent/HunyuanVideo-1.5\ - --output_path /fsx/yiyi/HunyuanVideo-1.5-Diffusers/transformer\ - --transformer_type 480p_t2v -""" - -# to convert full pipeline -""" -python scripts/convert_hunyuan_video1_5_to_diffusers.py \ - --original_state_dict_repo_id tencent/HunyuanVideo-1.5\ - --output_path /fsx/yiyi/HunyuanVideo-1.5-Diffusers \ - --save_pipeline \ - --byt5_path /fsx/yiyi/hy15/text_encoder/Glyph-SDXL-v2\ - --transformer_type 480p_t2v -""" - - -TRANSFORMER_CONFIGS = { - "480p_t2v": { - "target_size": 640, - "task_type": "i2v", - }, - "720p_t2v": { - "target_size": 960, - "task_type": "t2v", - }, - "720p_i2v": { - "target_size": 960, - "task_type": "i2v", - }, - "480p_t2v_distilled": { - "target_size": 640, - "task_type": "t2v", - }, - "480p_i2v_distilled": { - "target_size": 640, - "task_type": "i2v", - }, - "720p_i2v_distilled": { - "target_size": 960, - "task_type": "i2v", - }, - "480p_i2v_step_distilled": { - "target_size": 640, - "task_type": "i2v", - "use_meanflow": True, - }, -} - -SCHEDULER_CONFIGS = { - "480p_t2v": { - "shift": 5.0, - }, - "480p_i2v": { - "shift": 5.0, - }, - "720p_t2v": { - "shift": 9.0, - }, - "720p_i2v": { - "shift": 7.0, - }, - "480p_t2v_distilled": { - "shift": 5.0, - }, - "480p_i2v_distilled": { - "shift": 5.0, - }, - "720p_i2v_distilled": { - "shift": 7.0, - }, - "480p_i2v_step_distilled": { - "shift": 7.0, - }, -} - -GUIDANCE_CONFIGS = { - "480p_t2v": { - "guidance_scale": 6.0, - }, - "480p_i2v": { - "guidance_scale": 6.0, - }, - "720p_t2v": { - "guidance_scale": 6.0, - }, - "720p_i2v": { - "guidance_scale": 6.0, - }, - "480p_t2v_distilled": { - "guidance_scale": 1.0, - }, - "480p_i2v_distilled": { - "guidance_scale": 1.0, - }, - "720p_i2v_distilled": { - "guidance_scale": 1.0, - }, - "480p_i2v_step_distilled": { - "guidance_scale": 1.0, - }, -} - - -def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_hyvideo15_transformer_to_diffusers(original_state_dict, config=None): - """ - Convert HunyuanVideo 1.5 original checkpoint to Diffusers format. - """ - converted_state_dict = {} - - # 1. time_embed.timestep_embedder <- time_in - converted_state_dict["time_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "time_in.mlp.0.weight" - ) - converted_state_dict["time_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop("time_in.mlp.0.bias") - converted_state_dict["time_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "time_in.mlp.2.weight" - ) - converted_state_dict["time_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop("time_in.mlp.2.bias") - - if config.use_meanflow: - converted_state_dict["time_embed.timestep_embedder_r.linear_1.weight"] = original_state_dict.pop( - "time_r_in.mlp.0.weight" - ) - converted_state_dict["time_embed.timestep_embedder_r.linear_1.bias"] = original_state_dict.pop( - "time_r_in.mlp.0.bias" - ) - converted_state_dict["time_embed.timestep_embedder_r.linear_2.weight"] = original_state_dict.pop( - "time_r_in.mlp.2.weight" - ) - converted_state_dict["time_embed.timestep_embedder_r.linear_2.bias"] = original_state_dict.pop( - "time_r_in.mlp.2.bias" - ) - - # 2. context_embedder.time_text_embed.timestep_embedder <- txt_in.t_embedder - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.weight"] = ( - original_state_dict.pop("txt_in.t_embedder.mlp.0.weight") - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "txt_in.t_embedder.mlp.0.bias" - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.weight"] = ( - original_state_dict.pop("txt_in.t_embedder.mlp.2.weight") - ) - converted_state_dict["context_embedder.time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "txt_in.t_embedder.mlp.2.bias" - ) - - # 3. context_embedder.time_text_embed.text_embedder <- txt_in.c_embedder - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop( - "txt_in.c_embedder.linear_1.weight" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop( - "txt_in.c_embedder.linear_1.bias" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop( - "txt_in.c_embedder.linear_2.weight" - ) - converted_state_dict["context_embedder.time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop( - "txt_in.c_embedder.linear_2.bias" - ) - - # 4. context_embedder.proj_in <- txt_in.input_embedder - converted_state_dict["context_embedder.proj_in.weight"] = original_state_dict.pop("txt_in.input_embedder.weight") - converted_state_dict["context_embedder.proj_in.bias"] = original_state_dict.pop("txt_in.input_embedder.bias") - - # 5. context_embedder.token_refiner <- txt_in.individual_token_refiner - num_refiner_blocks = 2 - for i in range(num_refiner_blocks): - block_prefix = f"context_embedder.token_refiner.refiner_blocks.{i}." - orig_prefix = f"txt_in.individual_token_refiner.blocks.{i}." - - # norm1 - converted_state_dict[f"{block_prefix}norm1.weight"] = original_state_dict.pop(f"{orig_prefix}norm1.weight") - converted_state_dict[f"{block_prefix}norm1.bias"] = original_state_dict.pop(f"{orig_prefix}norm1.bias") - - # Split self_attn_qkv into to_q, to_k, to_v - qkv_weight = original_state_dict.pop(f"{orig_prefix}self_attn_qkv.weight") - qkv_bias = original_state_dict.pop(f"{orig_prefix}self_attn_qkv.bias") - q, k, v = torch.chunk(qkv_weight, 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) - - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = q - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = q_bias - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = k - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = k_bias - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = v - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = v_bias - - # self_attn_proj -> attn.to_out.0 - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"{orig_prefix}self_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"{orig_prefix}self_attn_proj.bias" - ) - - # norm2 - converted_state_dict[f"{block_prefix}norm2.weight"] = original_state_dict.pop(f"{orig_prefix}norm2.weight") - converted_state_dict[f"{block_prefix}norm2.bias"] = original_state_dict.pop(f"{orig_prefix}norm2.bias") - - # mlp -> ff - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( - f"{orig_prefix}mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( - f"{orig_prefix}mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"{orig_prefix}mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop(f"{orig_prefix}mlp.fc2.bias") - - # adaLN_modulation -> norm_out - converted_state_dict[f"{block_prefix}norm_out.linear.weight"] = original_state_dict.pop( - f"{orig_prefix}adaLN_modulation.1.weight" - ) - converted_state_dict[f"{block_prefix}norm_out.linear.bias"] = original_state_dict.pop( - f"{orig_prefix}adaLN_modulation.1.bias" - ) - - # 6. context_embedder_2 <- byt5_in - converted_state_dict["context_embedder_2.norm.weight"] = original_state_dict.pop("byt5_in.layernorm.weight") - converted_state_dict["context_embedder_2.norm.bias"] = original_state_dict.pop("byt5_in.layernorm.bias") - converted_state_dict["context_embedder_2.linear_1.weight"] = original_state_dict.pop("byt5_in.fc1.weight") - converted_state_dict["context_embedder_2.linear_1.bias"] = original_state_dict.pop("byt5_in.fc1.bias") - converted_state_dict["context_embedder_2.linear_2.weight"] = original_state_dict.pop("byt5_in.fc2.weight") - converted_state_dict["context_embedder_2.linear_2.bias"] = original_state_dict.pop("byt5_in.fc2.bias") - converted_state_dict["context_embedder_2.linear_3.weight"] = original_state_dict.pop("byt5_in.fc3.weight") - converted_state_dict["context_embedder_2.linear_3.bias"] = original_state_dict.pop("byt5_in.fc3.bias") - - # 7. image_embedder <- vision_in - converted_state_dict["image_embedder.norm_in.weight"] = original_state_dict.pop("vision_in.proj.0.weight") - converted_state_dict["image_embedder.norm_in.bias"] = original_state_dict.pop("vision_in.proj.0.bias") - converted_state_dict["image_embedder.linear_1.weight"] = original_state_dict.pop("vision_in.proj.1.weight") - converted_state_dict["image_embedder.linear_1.bias"] = original_state_dict.pop("vision_in.proj.1.bias") - converted_state_dict["image_embedder.linear_2.weight"] = original_state_dict.pop("vision_in.proj.3.weight") - converted_state_dict["image_embedder.linear_2.bias"] = original_state_dict.pop("vision_in.proj.3.bias") - converted_state_dict["image_embedder.norm_out.weight"] = original_state_dict.pop("vision_in.proj.4.weight") - converted_state_dict["image_embedder.norm_out.bias"] = original_state_dict.pop("vision_in.proj.4.bias") - - # 8. x_embedder <- img_in - converted_state_dict["x_embedder.proj.weight"] = original_state_dict.pop("img_in.proj.weight") - converted_state_dict["x_embedder.proj.bias"] = original_state_dict.pop("img_in.proj.bias") - - # 9. cond_type_embed <- cond_type_embedding - converted_state_dict["cond_type_embed.weight"] = original_state_dict.pop("cond_type_embedding.weight") - - # 10. transformer_blocks <- double_blocks - num_layers = 54 - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - orig_prefix = f"double_blocks.{i}." - - # norm1 (img_mod) - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop( - f"{orig_prefix}img_mod.linear.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop( - f"{orig_prefix}img_mod.linear.bias" - ) - - # norm1_context (txt_mod) - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_mod.linear.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_mod.linear.bias" - ) - - # img attention (to_q, to_k, to_v) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_q.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = original_state_dict.pop( - f"{orig_prefix}img_attn_q.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_k.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = original_state_dict.pop( - f"{orig_prefix}img_attn_k.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_v.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = original_state_dict.pop( - f"{orig_prefix}img_attn_v.bias" - ) - - # img attention qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_q_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_k_norm.weight" - ) - - # img attention output projection - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"{orig_prefix}img_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"{orig_prefix}img_attn_proj.bias" - ) - - # txt attention (add_q_proj, add_k_proj, add_v_proj) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_q.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_q.bias" - ) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_k.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_k.bias" - ) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_v.weight" - ) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_v.bias" - ) - - # txt attention qk norm - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_q_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_k_norm.weight" - ) - - # txt attention output projection - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_attn_proj.bias" - ) - - # norm2 and norm2_context (these don't have weights in the original, they're LayerNorm with elementwise_affine=False) - # So we skip them - - # img_mlp -> ff - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = original_state_dict.pop( - f"{orig_prefix}img_mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = original_state_dict.pop( - f"{orig_prefix}img_mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"{orig_prefix}img_mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( - f"{orig_prefix}img_mlp.fc2.bias" - ) - - # txt_mlp -> ff_context - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_mlp.fc1.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_mlp.fc1.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop( - f"{orig_prefix}txt_mlp.fc2.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop( - f"{orig_prefix}txt_mlp.fc2.bias" - ) - - # 11. norm_out and proj_out <- final_layer - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.weight") - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.bias") - ) - converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - - return converted_state_dict - - -def convert_hunyuan_video_15_vae_checkpoint_to_diffusers( - original_state_dict, block_out_channels=[128, 256, 512, 1024, 1024], layers_per_block=2 -): - converted = {} - - # 1. Encoder - # 1.1 conv_in - converted["encoder.conv_in.conv.weight"] = original_state_dict.pop("encoder.conv_in.conv.weight") - converted["encoder.conv_in.conv.bias"] = original_state_dict.pop("encoder.conv_in.conv.bias") - - # 1.2 Down blocks - for down_block_index in range(len(block_out_channels)): # 0 to 4 - # ResNet blocks - for resnet_block_index in range(layers_per_block): # 0 to 1 - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm1.gamma") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( - original_state_dict.pop( - f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.weight" - ) - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv1.conv.bias") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.norm2.gamma") - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( - original_state_dict.pop( - f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.weight" - ) - ) - converted[f"encoder.down_blocks.{down_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.block.{resnet_block_index}.conv2.conv.bias") - ) - - # Downsample (if exists) - if f"encoder.down.{down_block_index}.downsample.conv.conv.weight" in original_state_dict: - converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.weight"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.weight") - ) - converted[f"encoder.down_blocks.{down_block_index}.downsamplers.0.conv.conv.bias"] = ( - original_state_dict.pop(f"encoder.down.{down_block_index}.downsample.conv.conv.bias") - ) - - # 1.3 Mid block - converted["encoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm1.gamma") - converted["encoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_1.conv1.conv.weight" - ) - converted["encoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_1.conv1.conv.bias" - ) - converted["encoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_1.norm2.gamma") - converted["encoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_1.conv2.conv.weight" - ) - converted["encoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_1.conv2.conv.bias" - ) - - converted["encoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm1.gamma") - converted["encoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_2.conv1.conv.weight" - ) - converted["encoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_2.conv1.conv.bias" - ) - converted["encoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("encoder.mid.block_2.norm2.gamma") - converted["encoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( - "encoder.mid.block_2.conv2.conv.weight" - ) - converted["encoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( - "encoder.mid.block_2.conv2.conv.bias" - ) - - # Attention block - converted["encoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("encoder.mid.attn_1.norm.gamma") - converted["encoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("encoder.mid.attn_1.q.weight") - converted["encoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("encoder.mid.attn_1.q.bias") - converted["encoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("encoder.mid.attn_1.k.weight") - converted["encoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("encoder.mid.attn_1.k.bias") - converted["encoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("encoder.mid.attn_1.v.weight") - converted["encoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("encoder.mid.attn_1.v.bias") - converted["encoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( - "encoder.mid.attn_1.proj_out.weight" - ) - converted["encoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( - "encoder.mid.attn_1.proj_out.bias" - ) - - # 1.4 Encoder output - converted["encoder.norm_out.gamma"] = original_state_dict.pop("encoder.norm_out.gamma") - converted["encoder.conv_out.conv.weight"] = original_state_dict.pop("encoder.conv_out.conv.weight") - converted["encoder.conv_out.conv.bias"] = original_state_dict.pop("encoder.conv_out.conv.bias") - - # 2. Decoder - # 2.1 conv_in - converted["decoder.conv_in.conv.weight"] = original_state_dict.pop("decoder.conv_in.conv.weight") - converted["decoder.conv_in.conv.bias"] = original_state_dict.pop("decoder.conv_in.conv.bias") - - # 2.2 Mid block - converted["decoder.mid_block.resnets.0.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm1.gamma") - converted["decoder.mid_block.resnets.0.conv1.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_1.conv1.conv.weight" - ) - converted["decoder.mid_block.resnets.0.conv1.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_1.conv1.conv.bias" - ) - converted["decoder.mid_block.resnets.0.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_1.norm2.gamma") - converted["decoder.mid_block.resnets.0.conv2.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_1.conv2.conv.weight" - ) - converted["decoder.mid_block.resnets.0.conv2.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_1.conv2.conv.bias" - ) - - converted["decoder.mid_block.resnets.1.norm1.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm1.gamma") - converted["decoder.mid_block.resnets.1.conv1.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_2.conv1.conv.weight" - ) - converted["decoder.mid_block.resnets.1.conv1.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_2.conv1.conv.bias" - ) - converted["decoder.mid_block.resnets.1.norm2.gamma"] = original_state_dict.pop("decoder.mid.block_2.norm2.gamma") - converted["decoder.mid_block.resnets.1.conv2.conv.weight"] = original_state_dict.pop( - "decoder.mid.block_2.conv2.conv.weight" - ) - converted["decoder.mid_block.resnets.1.conv2.conv.bias"] = original_state_dict.pop( - "decoder.mid.block_2.conv2.conv.bias" - ) - - # Decoder attention block - converted["decoder.mid_block.attentions.0.norm.gamma"] = original_state_dict.pop("decoder.mid.attn_1.norm.gamma") - converted["decoder.mid_block.attentions.0.to_q.weight"] = original_state_dict.pop("decoder.mid.attn_1.q.weight") - converted["decoder.mid_block.attentions.0.to_q.bias"] = original_state_dict.pop("decoder.mid.attn_1.q.bias") - converted["decoder.mid_block.attentions.0.to_k.weight"] = original_state_dict.pop("decoder.mid.attn_1.k.weight") - converted["decoder.mid_block.attentions.0.to_k.bias"] = original_state_dict.pop("decoder.mid.attn_1.k.bias") - converted["decoder.mid_block.attentions.0.to_v.weight"] = original_state_dict.pop("decoder.mid.attn_1.v.weight") - converted["decoder.mid_block.attentions.0.to_v.bias"] = original_state_dict.pop("decoder.mid.attn_1.v.bias") - converted["decoder.mid_block.attentions.0.proj_out.weight"] = original_state_dict.pop( - "decoder.mid.attn_1.proj_out.weight" - ) - converted["decoder.mid_block.attentions.0.proj_out.bias"] = original_state_dict.pop( - "decoder.mid.attn_1.proj_out.bias" - ) - - # 2.3 Up blocks - for up_block_index in range(len(block_out_channels)): # 0 to 5 - # ResNet blocks - for resnet_block_index in range(layers_per_block + 1): # 0 to 2 (decoder has 3 resnets per level) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm1.gamma"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm1.gamma") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.weight"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.weight") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv1.conv.bias"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv1.conv.bias") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.norm2.gamma"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.norm2.gamma") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.weight"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.weight") - ) - converted[f"decoder.up_blocks.{up_block_index}.resnets.{resnet_block_index}.conv2.conv.bias"] = ( - original_state_dict.pop(f"decoder.up.{up_block_index}.block.{resnet_block_index}.conv2.conv.bias") - ) - - # Upsample (if exists) - if f"decoder.up.{up_block_index}.upsample.conv.conv.weight" in original_state_dict: - converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.weight"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.conv.weight" - ) - converted[f"decoder.up_blocks.{up_block_index}.upsamplers.0.conv.conv.bias"] = original_state_dict.pop( - f"decoder.up.{up_block_index}.upsample.conv.conv.bias" - ) - - # 2.4 Decoder output - converted["decoder.norm_out.gamma"] = original_state_dict.pop("decoder.norm_out.gamma") - converted["decoder.conv_out.conv.weight"] = original_state_dict.pop("decoder.conv_out.conv.weight") - converted["decoder.conv_out.conv.bias"] = original_state_dict.pop("decoder.conv_out.conv.bias") - - return converted - - -def load_sharded_safetensors(dir: pathlib.Path): - file_paths = list(dir.glob("diffusion_pytorch_model*.safetensors")) - state_dict = {} - for path in file_paths: - state_dict.update(load_file(path)) - return state_dict - - -def load_original_transformer_state_dict(args): - if args.original_state_dict_repo_id is not None: - model_dir = snapshot_download( - args.original_state_dict_repo_id, - repo_type="model", - allow_patterns="transformer/" + args.transformer_type + "/*", - ) - elif args.original_state_dict_folder is not None: - model_dir = pathlib.Path(args.original_state_dict_folder) - else: - raise ValueError("Please provide either `original_state_dict_repo_id` or `original_state_dict_folder`") - model_dir = pathlib.Path(model_dir) - model_dir = model_dir / "transformer" / args.transformer_type - return load_sharded_safetensors(model_dir) - - -def load_original_vae_state_dict(args): - if args.original_state_dict_repo_id is not None: - ckpt_path = hf_hub_download( - repo_id=args.original_state_dict_repo_id, filename="vae/diffusion_pytorch_model.safetensors" - ) - elif args.original_state_dict_folder is not None: - model_dir = pathlib.Path(args.original_state_dict_folder) - ckpt_path = model_dir / "vae/diffusion_pytorch_model.safetensors" - else: - raise ValueError("Please provide either `original_state_dict_repo_id` or `original_state_dict_folder`") - - original_state_dict = load_file(ckpt_path) - return original_state_dict - - -def convert_transformer(args): - original_state_dict = load_original_transformer_state_dict(args) - - config = TRANSFORMER_CONFIGS[args.transformer_type] - with init_empty_weights(): - transformer = HunyuanVideo15Transformer3DModel(**config) - state_dict = convert_hyvideo15_transformer_to_diffusers(original_state_dict, config=transformer.config) - transformer.load_state_dict(state_dict, strict=True, assign=True) - - return transformer - - -def convert_vae(args): - original_state_dict = load_original_vae_state_dict(args) - with init_empty_weights(): - vae = AutoencoderKLHunyuanVideo15() - state_dict = convert_hunyuan_video_15_vae_checkpoint_to_diffusers(original_state_dict) - vae.load_state_dict(state_dict, strict=True, assign=True) - return vae - - -def load_mllm(): - print(" loading from Qwen/Qwen2.5-VL-7B-Instruct") - text_encoder = AutoModel.from_pretrained( - "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True - ) - if hasattr(text_encoder, "language_model"): - text_encoder = text_encoder.language_model - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", padding_side="right") - return text_encoder, tokenizer - - -# copied from https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/910da2a829c484ea28982e8cff3bbc2cacdf1681/hyvideo/models/text_encoders/byT5/__init__.py#L89 -def add_special_token( - tokenizer, - text_encoder, - add_color=True, - add_font=True, - multilingual=True, - color_ann_path="assets/color_idx.json", - font_ann_path="assets/multilingual_10-lang_idx.json", -): - """ - Add special tokens for color and font to tokenizer and text encoder. - - Args: - tokenizer: Huggingface tokenizer. - text_encoder: Huggingface T5 encoder. - add_color (bool): Whether to add color tokens. - add_font (bool): Whether to add font tokens. - color_ann_path (str): Path to color annotation JSON. - font_ann_path (str): Path to font annotation JSON. - multilingual (bool): Whether to use multilingual font tokens. - """ - with open(font_ann_path, "r") as f: - idx_font_dict = json.load(f) - with open(color_ann_path, "r") as f: - idx_color_dict = json.load(f) - - if multilingual: - font_token = [f"<{font_code[:2]}-font-{idx_font_dict[font_code]}>" for font_code in idx_font_dict] - else: - font_token = [f"" for i in range(len(idx_font_dict))] - color_token = [f"" for i in range(len(idx_color_dict))] - additional_special_tokens = [] - if add_color: - additional_special_tokens += color_token - if add_font: - additional_special_tokens += font_token - - tokenizer.add_tokens(additional_special_tokens, special_tokens=True) - # Set mean_resizing=False to avoid PyTorch LAPACK dependency - text_encoder.resize_token_embeddings(len(tokenizer), mean_resizing=False) - - -def load_byt5(args): - """ - Load ByT5 encoder with Glyph-SDXL-v2 weights and save in HuggingFace format. - """ - - # 1. Load base tokenizer and encoder - tokenizer = AutoTokenizer.from_pretrained("google/byt5-small") - - # Load as T5EncoderModel - encoder = T5EncoderModel.from_pretrained("google/byt5-small") - - byt5_checkpoint_path = os.path.join(args.byt5_path, "checkpoints/byt5_model.pt") - color_ann_path = os.path.join(args.byt5_path, "assets/color_idx.json") - font_ann_path = os.path.join(args.byt5_path, "assets/multilingual_10-lang_idx.json") - - # 2. Add special tokens - add_special_token( - tokenizer=tokenizer, - text_encoder=encoder, - add_color=True, - add_font=True, - color_ann_path=color_ann_path, - font_ann_path=font_ann_path, - multilingual=True, - ) - - # 3. Load Glyph-SDXL-v2 checkpoint - print(f"\n3. Loading Glyph-SDXL-v2 checkpoint: {byt5_checkpoint_path}") - checkpoint = torch.load(byt5_checkpoint_path, map_location="cpu") - - # Handle different checkpoint formats - if "state_dict" in checkpoint: - state_dict = checkpoint["state_dict"] - else: - state_dict = checkpoint - - # add 'encoder.' prefix to the keys - # Remove 'module.text_tower.encoder.' prefix if present - cleaned_state_dict = {} - for key, value in state_dict.items(): - if key.startswith("module.text_tower.encoder."): - new_key = "encoder." + key[len("module.text_tower.encoder.") :] - cleaned_state_dict[new_key] = value - else: - new_key = "encoder." + key - cleaned_state_dict[new_key] = value - - # 4. Load weights - missing_keys, unexpected_keys = encoder.load_state_dict(cleaned_state_dict, strict=False) - if unexpected_keys: - raise ValueError(f"Unexpected keys: {unexpected_keys}") - if "shared.weight" in missing_keys: - print(" Missing shared.weight as expected") - missing_keys.remove("shared.weight") - if missing_keys: - raise ValueError(f"Missing keys: {missing_keys}") - - return encoder, tokenizer - - -def load_siglip(): - image_encoder = SiglipVisionModel.from_pretrained( - "black-forest-labs/FLUX.1-Redux-dev", subfolder="image_encoder", torch_dtype=torch.bfloat16 - ) - feature_extractor = SiglipImageProcessor.from_pretrained( - "black-forest-labs/FLUX.1-Redux-dev", subfolder="feature_extractor" - ) - return image_encoder, feature_extractor - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--original_state_dict_repo_id", type=str, default=None, help="Path to original hub_id for the model" - ) - parser.add_argument( - "--original_state_dict_folder", type=str, default=None, help="Local folder name of the original state dict" - ) - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model(s) should be saved") - parser.add_argument("--transformer_type", type=str, default="480p_i2v", choices=list(TRANSFORMER_CONFIGS.keys())) - parser.add_argument( - "--byt5_path", - type=str, - default=None, - help=( - "path to the downloaded byt5 checkpoint & assets. " - "Note: They use Glyph-SDXL-v2 as byt5 encoder. You can download from modelscope like: " - "`modelscope download --model AI-ModelScope/Glyph-SDXL-v2 --local_dir ./ckpts/text_encoder/Glyph-SDXL-v2` " - "or manually download following the instructions on " - "https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/910da2a829c484ea28982e8cff3bbc2cacdf1681/checkpoints-download.md. " - "The path should point to the Glyph-SDXL-v2 folder which should contain an `assets` folder and a `checkpoints` folder, " - "like: Glyph-SDXL-v2/assets/... and Glyph-SDXL-v2/checkpoints/byt5_model.pt" - ), - ) - parser.add_argument("--save_pipeline", action="store_true") - return parser.parse_args() - - -if __name__ == "__main__": - args = get_args() - - if args.save_pipeline and args.byt5_path is None: - raise ValueError("Please provide --byt5_path when saving pipeline") - - transformer = None - - transformer = convert_transformer(args) - if not args.save_pipeline: - transformer.save_pretrained(args.output_path, safe_serialization=True) - else: - task_type = transformer.config.task_type - - vae = convert_vae(args) - - text_encoder, tokenizer = load_mllm() - text_encoder_2, tokenizer_2 = load_byt5(args) - - flow_shift = SCHEDULER_CONFIGS[args.transformer_type]["shift"] - scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) - - guidance_scale = GUIDANCE_CONFIGS[args.transformer_type]["guidance_scale"] - guider = ClassifierFreeGuidance(guidance_scale=guidance_scale) - - if task_type == "i2v": - image_encoder, feature_extractor = load_siglip() - pipeline = HunyuanVideo15ImageToVideoPipeline( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - guider=guider, - scheduler=scheduler, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - ) - elif task_type == "t2v": - pipeline = HunyuanVideo15Pipeline( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - transformer=transformer, - guider=guider, - scheduler=scheduler, - ) - else: - raise ValueError(f"Task type {task_type} is not supported") - - pipeline.save_pretrained(args.output_path, safe_serialization=True) diff --git a/scripts/convert_hunyuan_video_to_diffusers.py b/scripts/convert_hunyuan_video_to_diffusers.py deleted file mode 100644 index fb604c98b65b..000000000000 --- a/scripts/convert_hunyuan_video_to_diffusers.py +++ /dev/null @@ -1,353 +0,0 @@ -import argparse -from typing import Any, Dict - -import torch -from accelerate import init_empty_weights -from transformers import ( - AutoModel, - AutoTokenizer, - CLIPImageProcessor, - CLIPTextModel, - CLIPTokenizer, - LlavaForConditionalGeneration, -) - -from diffusers import ( - AutoencoderKLHunyuanVideo, - FlowMatchEulerDiscreteScheduler, - HunyuanVideoImageToVideoPipeline, - HunyuanVideoPipeline, - HunyuanVideoTransformer3DModel, -) - - -def remap_norm_scale_shift_(key, state_dict): - weight = state_dict.pop(key) - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - state_dict[key.replace("final_layer.adaLN_modulation.1", "norm_out.linear")] = new_weight - - -def remap_txt_in_(key, state_dict): - def rename_key(key): - new_key = key.replace("individual_token_refiner.blocks", "token_refiner.refiner_blocks") - new_key = new_key.replace("adaLN_modulation.1", "norm_out.linear") - new_key = new_key.replace("txt_in", "context_embedder") - new_key = new_key.replace("t_embedder.mlp.0", "time_text_embed.timestep_embedder.linear_1") - new_key = new_key.replace("t_embedder.mlp.2", "time_text_embed.timestep_embedder.linear_2") - new_key = new_key.replace("c_embedder", "time_text_embed.text_embedder") - new_key = new_key.replace("mlp", "ff") - return new_key - - if "self_attn_qkv" in key: - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_q"))] = to_q - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_k"))] = to_k - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_v"))] = to_v - else: - state_dict[rename_key(key)] = state_dict.pop(key) - - -def remap_img_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = to_q - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = to_k - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = to_v - - -def remap_txt_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = to_q - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = to_k - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = to_v - - -def remap_single_transformer_blocks_(key, state_dict): - hidden_size = 3072 - - if "linear1.weight" in key: - linear1_weight = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_weight.size(0) - 3 * hidden_size) - q, k, v, mlp = torch.split(linear1_weight, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.weight") - state_dict[f"{new_key}.attn.to_q.weight"] = q - state_dict[f"{new_key}.attn.to_k.weight"] = k - state_dict[f"{new_key}.attn.to_v.weight"] = v - state_dict[f"{new_key}.proj_mlp.weight"] = mlp - - elif "linear1.bias" in key: - linear1_bias = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_bias.size(0) - 3 * hidden_size) - q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.bias") - state_dict[f"{new_key}.attn.to_q.bias"] = q_bias - state_dict[f"{new_key}.attn.to_k.bias"] = k_bias - state_dict[f"{new_key}.attn.to_v.bias"] = v_bias - state_dict[f"{new_key}.proj_mlp.bias"] = mlp_bias - - else: - new_key = key.replace("single_blocks", "single_transformer_blocks") - new_key = new_key.replace("linear2", "proj_out") - new_key = new_key.replace("q_norm", "attn.norm_q") - new_key = new_key.replace("k_norm", "attn.norm_k") - state_dict[new_key] = state_dict.pop(key) - - -TRANSFORMER_KEYS_RENAME_DICT = { - "img_in": "x_embedder", - "time_in.mlp.0": "time_text_embed.timestep_embedder.linear_1", - "time_in.mlp.2": "time_text_embed.timestep_embedder.linear_2", - "guidance_in.mlp.0": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "double_blocks": "transformer_blocks", - "img_attn_q_norm": "attn.norm_q", - "img_attn_k_norm": "attn.norm_k", - "img_attn_proj": "attn.to_out.0", - "txt_attn_q_norm": "attn.norm_added_q", - "txt_attn_k_norm": "attn.norm_added_k", - "txt_attn_proj": "attn.to_add_out", - "img_mod.linear": "norm1.linear", - "img_norm1": "norm1.norm", - "img_norm2": "norm2", - "img_mlp": "ff", - "txt_mod.linear": "norm1_context.linear", - "txt_norm1": "norm1.norm", - "txt_norm2": "norm2_context", - "txt_mlp": "ff_context", - "self_attn_proj": "attn.to_out.0", - "modulation.linear": "norm.linear", - "pre_norm": "norm.norm", - "final_layer.norm_final": "norm_out.norm", - "final_layer.linear": "proj_out", - "fc1": "net.0.proj", - "fc2": "net.2", - "input_embedder": "proj_in", -} - -TRANSFORMER_SPECIAL_KEYS_REMAP = { - "txt_in": remap_txt_in_, - "img_attn_qkv": remap_img_attn_qkv_, - "txt_attn_qkv": remap_txt_attn_qkv_, - "single_blocks": remap_single_transformer_blocks_, - "final_layer.adaLN_modulation.1": remap_norm_scale_shift_, -} - -VAE_KEYS_RENAME_DICT = {} - -VAE_SPECIAL_KEYS_REMAP = {} - - -TRANSFORMER_CONFIGS = { - "HYVideo-T/2-cfgdistill": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 24, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "mlp_ratio": 4.0, - "patch_size": 2, - "patch_size_t": 1, - "qk_norm": "rms_norm", - "guidance_embeds": True, - "text_embed_dim": 4096, - "pooled_projection_dim": 768, - "rope_theta": 256.0, - "rope_axes_dim": (16, 56, 56), - "image_condition_type": None, - }, - "HYVideo-T/2-I2V-33ch": { - "in_channels": 16 * 2 + 1, - "out_channels": 16, - "num_attention_heads": 24, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "mlp_ratio": 4.0, - "patch_size": 2, - "patch_size_t": 1, - "qk_norm": "rms_norm", - "guidance_embeds": False, - "text_embed_dim": 4096, - "pooled_projection_dim": 768, - "rope_theta": 256.0, - "rope_axes_dim": (16, 56, 56), - "image_condition_type": "latent_concat", - }, - "HYVideo-T/2-I2V-16ch": { - "in_channels": 16, - "out_channels": 16, - "num_attention_heads": 24, - "attention_head_dim": 128, - "num_layers": 20, - "num_single_layers": 40, - "num_refiner_layers": 2, - "mlp_ratio": 4.0, - "patch_size": 2, - "patch_size_t": 1, - "qk_norm": "rms_norm", - "guidance_embeds": True, - "text_embed_dim": 4096, - "pooled_projection_dim": 768, - "rope_theta": 256.0, - "rope_axes_dim": (16, 56, 56), - "image_condition_type": "token_replace", - }, -} - - -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: - state_dict = saved_dict - if "model" in saved_dict.keys(): - state_dict = state_dict["model"] - if "module" in saved_dict.keys(): - state_dict = state_dict["module"] - if "state_dict" in saved_dict.keys(): - state_dict = state_dict["state_dict"] - return state_dict - - -def convert_transformer(ckpt_path: str, transformer_type: str): - original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True)) - config = TRANSFORMER_CONFIGS[transformer_type] - - with init_empty_weights(): - transformer = HunyuanVideoTransformer3DModel(**config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - transformer.load_state_dict(original_state_dict, strict=True, assign=True) - return transformer - - -def convert_vae(ckpt_path: str): - original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True)) - - with init_empty_weights(): - vae = AutoencoderKLHunyuanVideo() - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" - ) - parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original VAE checkpoint") - parser.add_argument("--text_encoder_path", type=str, default=None, help="Path to original llama checkpoint") - parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to original llama tokenizer") - parser.add_argument("--text_encoder_2_path", type=str, default=None, help="Path to original clip checkpoint") - parser.add_argument("--save_pipeline", action="store_true") - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") - parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.") - parser.add_argument( - "--transformer_type", type=str, default="HYVideo-T/2-cfgdistill", choices=list(TRANSFORMER_CONFIGS.keys()) - ) - parser.add_argument("--flow_shift", type=float, default=7.0) - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -if __name__ == "__main__": - args = get_args() - - transformer = None - dtype = DTYPE_MAPPING[args.dtype] - - if args.save_pipeline: - assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None - assert args.text_encoder_path is not None - assert args.tokenizer_path is not None - assert args.text_encoder_2_path is not None - - if args.transformer_ckpt_path is not None: - transformer = convert_transformer(args.transformer_ckpt_path, args.transformer_type) - transformer = transformer.to(dtype=dtype) - if not args.save_pipeline: - transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.vae_ckpt_path is not None: - vae = convert_vae(args.vae_ckpt_path) - if not args.save_pipeline: - vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.save_pipeline: - if args.transformer_type == "HYVideo-T/2-cfgdistill": - text_encoder = AutoModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.float16) - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, padding_side="right") - text_encoder_2 = CLIPTextModel.from_pretrained(args.text_encoder_2_path, torch_dtype=torch.float16) - tokenizer_2 = CLIPTokenizer.from_pretrained(args.text_encoder_2_path) - scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) - - pipe = HunyuanVideoPipeline( - transformer=transformer, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - scheduler=scheduler, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - else: - text_encoder = LlavaForConditionalGeneration.from_pretrained( - args.text_encoder_path, torch_dtype=torch.float16 - ) - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, padding_side="right") - text_encoder_2 = CLIPTextModel.from_pretrained(args.text_encoder_2_path, torch_dtype=torch.float16) - tokenizer_2 = CLIPTokenizer.from_pretrained(args.text_encoder_2_path) - scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) - image_processor = CLIPImageProcessor.from_pretrained(args.text_encoder_path) - - pipe = HunyuanVideoImageToVideoPipeline( - transformer=transformer, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - text_encoder_2=text_encoder_2, - tokenizer_2=tokenizer_2, - scheduler=scheduler, - image_processor=image_processor, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") diff --git a/scripts/convert_hunyuandit_controlnet_to_diffusers.py b/scripts/convert_hunyuandit_controlnet_to_diffusers.py deleted file mode 100644 index 5cef46c98983..000000000000 --- a/scripts/convert_hunyuandit_controlnet_to_diffusers.py +++ /dev/null @@ -1,241 +0,0 @@ -import argparse - -import torch - -from diffusers import HunyuanDiT2DControlNetModel - - -def main(args): - state_dict = torch.load(args.pt_checkpoint_path, map_location="cpu") - - if args.load_key != "none": - try: - state_dict = state_dict[args.load_key] - except KeyError: - raise KeyError( - f"{args.load_key} not found in the checkpoint." - "Please load from the following keys:{state_dict.keys()}" - ) - device = "cuda" - - model_config = HunyuanDiT2DControlNetModel.load_config( - "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", subfolder="transformer" - ) - model_config["use_style_cond_and_image_meta_size"] = ( - args.use_style_cond_and_image_meta_size - ) ### version <= v1.1: True; version >= v1.2: False - print(model_config) - - for key in state_dict: - print("local:", key) - - model = HunyuanDiT2DControlNetModel.from_config(model_config).to(device) - - for key in model.state_dict(): - print("diffusers:", key) - - num_layers = 19 - for i in range(num_layers): - # attn1 - # Wkqv -> to_q, to_k, to_v - q, k, v = torch.chunk(state_dict[f"blocks.{i}.attn1.Wqkv.weight"], 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(state_dict[f"blocks.{i}.attn1.Wqkv.bias"], 3, dim=0) - state_dict[f"blocks.{i}.attn1.to_q.weight"] = q - state_dict[f"blocks.{i}.attn1.to_q.bias"] = q_bias - state_dict[f"blocks.{i}.attn1.to_k.weight"] = k - state_dict[f"blocks.{i}.attn1.to_k.bias"] = k_bias - state_dict[f"blocks.{i}.attn1.to_v.weight"] = v - state_dict[f"blocks.{i}.attn1.to_v.bias"] = v_bias - state_dict.pop(f"blocks.{i}.attn1.Wqkv.weight") - state_dict.pop(f"blocks.{i}.attn1.Wqkv.bias") - - # q_norm, k_norm -> norm_q, norm_k - state_dict[f"blocks.{i}.attn1.norm_q.weight"] = state_dict[f"blocks.{i}.attn1.q_norm.weight"] - state_dict[f"blocks.{i}.attn1.norm_q.bias"] = state_dict[f"blocks.{i}.attn1.q_norm.bias"] - state_dict[f"blocks.{i}.attn1.norm_k.weight"] = state_dict[f"blocks.{i}.attn1.k_norm.weight"] - state_dict[f"blocks.{i}.attn1.norm_k.bias"] = state_dict[f"blocks.{i}.attn1.k_norm.bias"] - - state_dict.pop(f"blocks.{i}.attn1.q_norm.weight") - state_dict.pop(f"blocks.{i}.attn1.q_norm.bias") - state_dict.pop(f"blocks.{i}.attn1.k_norm.weight") - state_dict.pop(f"blocks.{i}.attn1.k_norm.bias") - - # out_proj -> to_out - state_dict[f"blocks.{i}.attn1.to_out.0.weight"] = state_dict[f"blocks.{i}.attn1.out_proj.weight"] - state_dict[f"blocks.{i}.attn1.to_out.0.bias"] = state_dict[f"blocks.{i}.attn1.out_proj.bias"] - state_dict.pop(f"blocks.{i}.attn1.out_proj.weight") - state_dict.pop(f"blocks.{i}.attn1.out_proj.bias") - - # attn2 - # kq_proj -> to_k, to_v - k, v = torch.chunk(state_dict[f"blocks.{i}.attn2.kv_proj.weight"], 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict[f"blocks.{i}.attn2.kv_proj.bias"], 2, dim=0) - state_dict[f"blocks.{i}.attn2.to_k.weight"] = k - state_dict[f"blocks.{i}.attn2.to_k.bias"] = k_bias - state_dict[f"blocks.{i}.attn2.to_v.weight"] = v - state_dict[f"blocks.{i}.attn2.to_v.bias"] = v_bias - state_dict.pop(f"blocks.{i}.attn2.kv_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.kv_proj.bias") - - # q_proj -> to_q - state_dict[f"blocks.{i}.attn2.to_q.weight"] = state_dict[f"blocks.{i}.attn2.q_proj.weight"] - state_dict[f"blocks.{i}.attn2.to_q.bias"] = state_dict[f"blocks.{i}.attn2.q_proj.bias"] - state_dict.pop(f"blocks.{i}.attn2.q_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.q_proj.bias") - - # q_norm, k_norm -> norm_q, norm_k - state_dict[f"blocks.{i}.attn2.norm_q.weight"] = state_dict[f"blocks.{i}.attn2.q_norm.weight"] - state_dict[f"blocks.{i}.attn2.norm_q.bias"] = state_dict[f"blocks.{i}.attn2.q_norm.bias"] - state_dict[f"blocks.{i}.attn2.norm_k.weight"] = state_dict[f"blocks.{i}.attn2.k_norm.weight"] - state_dict[f"blocks.{i}.attn2.norm_k.bias"] = state_dict[f"blocks.{i}.attn2.k_norm.bias"] - - state_dict.pop(f"blocks.{i}.attn2.q_norm.weight") - state_dict.pop(f"blocks.{i}.attn2.q_norm.bias") - state_dict.pop(f"blocks.{i}.attn2.k_norm.weight") - state_dict.pop(f"blocks.{i}.attn2.k_norm.bias") - - # out_proj -> to_out - state_dict[f"blocks.{i}.attn2.to_out.0.weight"] = state_dict[f"blocks.{i}.attn2.out_proj.weight"] - state_dict[f"blocks.{i}.attn2.to_out.0.bias"] = state_dict[f"blocks.{i}.attn2.out_proj.bias"] - state_dict.pop(f"blocks.{i}.attn2.out_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.out_proj.bias") - - # switch norm 2 and norm 3 - norm2_weight = state_dict[f"blocks.{i}.norm2.weight"] - norm2_bias = state_dict[f"blocks.{i}.norm2.bias"] - state_dict[f"blocks.{i}.norm2.weight"] = state_dict[f"blocks.{i}.norm3.weight"] - state_dict[f"blocks.{i}.norm2.bias"] = state_dict[f"blocks.{i}.norm3.bias"] - state_dict[f"blocks.{i}.norm3.weight"] = norm2_weight - state_dict[f"blocks.{i}.norm3.bias"] = norm2_bias - - # norm1 -> norm1.norm - # default_modulation.1 -> norm1.linear - state_dict[f"blocks.{i}.norm1.norm.weight"] = state_dict[f"blocks.{i}.norm1.weight"] - state_dict[f"blocks.{i}.norm1.norm.bias"] = state_dict[f"blocks.{i}.norm1.bias"] - state_dict[f"blocks.{i}.norm1.linear.weight"] = state_dict[f"blocks.{i}.default_modulation.1.weight"] - state_dict[f"blocks.{i}.norm1.linear.bias"] = state_dict[f"blocks.{i}.default_modulation.1.bias"] - state_dict.pop(f"blocks.{i}.norm1.weight") - state_dict.pop(f"blocks.{i}.norm1.bias") - state_dict.pop(f"blocks.{i}.default_modulation.1.weight") - state_dict.pop(f"blocks.{i}.default_modulation.1.bias") - - # mlp.fc1 -> ff.net.0, mlp.fc2 -> ff.net.2 - state_dict[f"blocks.{i}.ff.net.0.proj.weight"] = state_dict[f"blocks.{i}.mlp.fc1.weight"] - state_dict[f"blocks.{i}.ff.net.0.proj.bias"] = state_dict[f"blocks.{i}.mlp.fc1.bias"] - state_dict[f"blocks.{i}.ff.net.2.weight"] = state_dict[f"blocks.{i}.mlp.fc2.weight"] - state_dict[f"blocks.{i}.ff.net.2.bias"] = state_dict[f"blocks.{i}.mlp.fc2.bias"] - state_dict.pop(f"blocks.{i}.mlp.fc1.weight") - state_dict.pop(f"blocks.{i}.mlp.fc1.bias") - state_dict.pop(f"blocks.{i}.mlp.fc2.weight") - state_dict.pop(f"blocks.{i}.mlp.fc2.bias") - - # after_proj_list -> controlnet_blocks - state_dict[f"controlnet_blocks.{i}.weight"] = state_dict[f"after_proj_list.{i}.weight"] - state_dict[f"controlnet_blocks.{i}.bias"] = state_dict[f"after_proj_list.{i}.bias"] - state_dict.pop(f"after_proj_list.{i}.weight") - state_dict.pop(f"after_proj_list.{i}.bias") - - # before_proj -> input_block - state_dict["input_block.weight"] = state_dict["before_proj.weight"] - state_dict["input_block.bias"] = state_dict["before_proj.bias"] - state_dict.pop("before_proj.weight") - state_dict.pop("before_proj.bias") - - # pooler -> time_extra_emb - state_dict["time_extra_emb.pooler.positional_embedding"] = state_dict["pooler.positional_embedding"] - state_dict["time_extra_emb.pooler.k_proj.weight"] = state_dict["pooler.k_proj.weight"] - state_dict["time_extra_emb.pooler.k_proj.bias"] = state_dict["pooler.k_proj.bias"] - state_dict["time_extra_emb.pooler.q_proj.weight"] = state_dict["pooler.q_proj.weight"] - state_dict["time_extra_emb.pooler.q_proj.bias"] = state_dict["pooler.q_proj.bias"] - state_dict["time_extra_emb.pooler.v_proj.weight"] = state_dict["pooler.v_proj.weight"] - state_dict["time_extra_emb.pooler.v_proj.bias"] = state_dict["pooler.v_proj.bias"] - state_dict["time_extra_emb.pooler.c_proj.weight"] = state_dict["pooler.c_proj.weight"] - state_dict["time_extra_emb.pooler.c_proj.bias"] = state_dict["pooler.c_proj.bias"] - state_dict.pop("pooler.k_proj.weight") - state_dict.pop("pooler.k_proj.bias") - state_dict.pop("pooler.q_proj.weight") - state_dict.pop("pooler.q_proj.bias") - state_dict.pop("pooler.v_proj.weight") - state_dict.pop("pooler.v_proj.bias") - state_dict.pop("pooler.c_proj.weight") - state_dict.pop("pooler.c_proj.bias") - state_dict.pop("pooler.positional_embedding") - - # t_embedder -> time_embedding (`TimestepEmbedding`) - state_dict["time_extra_emb.timestep_embedder.linear_1.bias"] = state_dict["t_embedder.mlp.0.bias"] - state_dict["time_extra_emb.timestep_embedder.linear_1.weight"] = state_dict["t_embedder.mlp.0.weight"] - state_dict["time_extra_emb.timestep_embedder.linear_2.bias"] = state_dict["t_embedder.mlp.2.bias"] - state_dict["time_extra_emb.timestep_embedder.linear_2.weight"] = state_dict["t_embedder.mlp.2.weight"] - - state_dict.pop("t_embedder.mlp.0.bias") - state_dict.pop("t_embedder.mlp.0.weight") - state_dict.pop("t_embedder.mlp.2.bias") - state_dict.pop("t_embedder.mlp.2.weight") - - # x_embedder -> pos_embd (`PatchEmbed`) - state_dict["pos_embed.proj.weight"] = state_dict["x_embedder.proj.weight"] - state_dict["pos_embed.proj.bias"] = state_dict["x_embedder.proj.bias"] - state_dict.pop("x_embedder.proj.weight") - state_dict.pop("x_embedder.proj.bias") - - # mlp_t5 -> text_embedder - state_dict["text_embedder.linear_1.bias"] = state_dict["mlp_t5.0.bias"] - state_dict["text_embedder.linear_1.weight"] = state_dict["mlp_t5.0.weight"] - state_dict["text_embedder.linear_2.bias"] = state_dict["mlp_t5.2.bias"] - state_dict["text_embedder.linear_2.weight"] = state_dict["mlp_t5.2.weight"] - state_dict.pop("mlp_t5.0.bias") - state_dict.pop("mlp_t5.0.weight") - state_dict.pop("mlp_t5.2.bias") - state_dict.pop("mlp_t5.2.weight") - - # extra_embedder -> extra_embedder - state_dict["time_extra_emb.extra_embedder.linear_1.bias"] = state_dict["extra_embedder.0.bias"] - state_dict["time_extra_emb.extra_embedder.linear_1.weight"] = state_dict["extra_embedder.0.weight"] - state_dict["time_extra_emb.extra_embedder.linear_2.bias"] = state_dict["extra_embedder.2.bias"] - state_dict["time_extra_emb.extra_embedder.linear_2.weight"] = state_dict["extra_embedder.2.weight"] - state_dict.pop("extra_embedder.0.bias") - state_dict.pop("extra_embedder.0.weight") - state_dict.pop("extra_embedder.2.bias") - state_dict.pop("extra_embedder.2.weight") - - # style_embedder - if model_config["use_style_cond_and_image_meta_size"]: - print(state_dict["style_embedder.weight"]) - print(state_dict["style_embedder.weight"].shape) - state_dict["time_extra_emb.style_embedder.weight"] = state_dict["style_embedder.weight"][0:1] - state_dict.pop("style_embedder.weight") - - model.load_state_dict(state_dict) - - if args.save: - model.save_pretrained(args.output_checkpoint_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." - ) - parser.add_argument( - "--pt_checkpoint_path", default=None, type=str, required=True, help="Path to the .pt pretrained model." - ) - parser.add_argument( - "--output_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the output converted diffusers pipeline.", - ) - parser.add_argument( - "--load_key", default="none", type=str, required=False, help="The key to load from the pretrained .pt file" - ) - parser.add_argument( - "--use_style_cond_and_image_meta_size", - type=bool, - default=False, - help="version <= v1.1: True; version >= v1.2: False", - ) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_hunyuandit_to_diffusers.py b/scripts/convert_hunyuandit_to_diffusers.py deleted file mode 100644 index 65fcccb22a1a..000000000000 --- a/scripts/convert_hunyuandit_to_diffusers.py +++ /dev/null @@ -1,266 +0,0 @@ -import argparse - -import torch - -from diffusers import HunyuanDiT2DModel - - -def main(args): - state_dict = torch.load(args.pt_checkpoint_path, map_location="cpu") - - if args.load_key != "none": - try: - state_dict = state_dict[args.load_key] - except KeyError: - raise KeyError( - f"{args.load_key} not found in the checkpoint.Please load from the following keys:{state_dict.keys()}" - ) - - device = "cuda" - model_config = HunyuanDiT2DModel.load_config("Tencent-Hunyuan/HunyuanDiT-Diffusers", subfolder="transformer") - model_config["use_style_cond_and_image_meta_size"] = ( - args.use_style_cond_and_image_meta_size - ) ### version <= v1.1: True; version >= v1.2: False - - # input_size -> sample_size, text_dim -> cross_attention_dim - for key in state_dict: - print("local:", key) - - model = HunyuanDiT2DModel.from_config(model_config).to(device) - - for key in model.state_dict(): - print("diffusers:", key) - - num_layers = 40 - for i in range(num_layers): - # attn1 - # Wkqv -> to_q, to_k, to_v - q, k, v = torch.chunk(state_dict[f"blocks.{i}.attn1.Wqkv.weight"], 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(state_dict[f"blocks.{i}.attn1.Wqkv.bias"], 3, dim=0) - state_dict[f"blocks.{i}.attn1.to_q.weight"] = q - state_dict[f"blocks.{i}.attn1.to_q.bias"] = q_bias - state_dict[f"blocks.{i}.attn1.to_k.weight"] = k - state_dict[f"blocks.{i}.attn1.to_k.bias"] = k_bias - state_dict[f"blocks.{i}.attn1.to_v.weight"] = v - state_dict[f"blocks.{i}.attn1.to_v.bias"] = v_bias - state_dict.pop(f"blocks.{i}.attn1.Wqkv.weight") - state_dict.pop(f"blocks.{i}.attn1.Wqkv.bias") - - # q_norm, k_norm -> norm_q, norm_k - state_dict[f"blocks.{i}.attn1.norm_q.weight"] = state_dict[f"blocks.{i}.attn1.q_norm.weight"] - state_dict[f"blocks.{i}.attn1.norm_q.bias"] = state_dict[f"blocks.{i}.attn1.q_norm.bias"] - state_dict[f"blocks.{i}.attn1.norm_k.weight"] = state_dict[f"blocks.{i}.attn1.k_norm.weight"] - state_dict[f"blocks.{i}.attn1.norm_k.bias"] = state_dict[f"blocks.{i}.attn1.k_norm.bias"] - - state_dict.pop(f"blocks.{i}.attn1.q_norm.weight") - state_dict.pop(f"blocks.{i}.attn1.q_norm.bias") - state_dict.pop(f"blocks.{i}.attn1.k_norm.weight") - state_dict.pop(f"blocks.{i}.attn1.k_norm.bias") - - # out_proj -> to_out - state_dict[f"blocks.{i}.attn1.to_out.0.weight"] = state_dict[f"blocks.{i}.attn1.out_proj.weight"] - state_dict[f"blocks.{i}.attn1.to_out.0.bias"] = state_dict[f"blocks.{i}.attn1.out_proj.bias"] - state_dict.pop(f"blocks.{i}.attn1.out_proj.weight") - state_dict.pop(f"blocks.{i}.attn1.out_proj.bias") - - # attn2 - # kq_proj -> to_k, to_v - k, v = torch.chunk(state_dict[f"blocks.{i}.attn2.kv_proj.weight"], 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict[f"blocks.{i}.attn2.kv_proj.bias"], 2, dim=0) - state_dict[f"blocks.{i}.attn2.to_k.weight"] = k - state_dict[f"blocks.{i}.attn2.to_k.bias"] = k_bias - state_dict[f"blocks.{i}.attn2.to_v.weight"] = v - state_dict[f"blocks.{i}.attn2.to_v.bias"] = v_bias - state_dict.pop(f"blocks.{i}.attn2.kv_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.kv_proj.bias") - - # q_proj -> to_q - state_dict[f"blocks.{i}.attn2.to_q.weight"] = state_dict[f"blocks.{i}.attn2.q_proj.weight"] - state_dict[f"blocks.{i}.attn2.to_q.bias"] = state_dict[f"blocks.{i}.attn2.q_proj.bias"] - state_dict.pop(f"blocks.{i}.attn2.q_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.q_proj.bias") - - # q_norm, k_norm -> norm_q, norm_k - state_dict[f"blocks.{i}.attn2.norm_q.weight"] = state_dict[f"blocks.{i}.attn2.q_norm.weight"] - state_dict[f"blocks.{i}.attn2.norm_q.bias"] = state_dict[f"blocks.{i}.attn2.q_norm.bias"] - state_dict[f"blocks.{i}.attn2.norm_k.weight"] = state_dict[f"blocks.{i}.attn2.k_norm.weight"] - state_dict[f"blocks.{i}.attn2.norm_k.bias"] = state_dict[f"blocks.{i}.attn2.k_norm.bias"] - - state_dict.pop(f"blocks.{i}.attn2.q_norm.weight") - state_dict.pop(f"blocks.{i}.attn2.q_norm.bias") - state_dict.pop(f"blocks.{i}.attn2.k_norm.weight") - state_dict.pop(f"blocks.{i}.attn2.k_norm.bias") - - # out_proj -> to_out - state_dict[f"blocks.{i}.attn2.to_out.0.weight"] = state_dict[f"blocks.{i}.attn2.out_proj.weight"] - state_dict[f"blocks.{i}.attn2.to_out.0.bias"] = state_dict[f"blocks.{i}.attn2.out_proj.bias"] - state_dict.pop(f"blocks.{i}.attn2.out_proj.weight") - state_dict.pop(f"blocks.{i}.attn2.out_proj.bias") - - # switch norm 2 and norm 3 - norm2_weight = state_dict[f"blocks.{i}.norm2.weight"] - norm2_bias = state_dict[f"blocks.{i}.norm2.bias"] - state_dict[f"blocks.{i}.norm2.weight"] = state_dict[f"blocks.{i}.norm3.weight"] - state_dict[f"blocks.{i}.norm2.bias"] = state_dict[f"blocks.{i}.norm3.bias"] - state_dict[f"blocks.{i}.norm3.weight"] = norm2_weight - state_dict[f"blocks.{i}.norm3.bias"] = norm2_bias - - # norm1 -> norm1.norm - # default_modulation.1 -> norm1.linear - state_dict[f"blocks.{i}.norm1.norm.weight"] = state_dict[f"blocks.{i}.norm1.weight"] - state_dict[f"blocks.{i}.norm1.norm.bias"] = state_dict[f"blocks.{i}.norm1.bias"] - state_dict[f"blocks.{i}.norm1.linear.weight"] = state_dict[f"blocks.{i}.default_modulation.1.weight"] - state_dict[f"blocks.{i}.norm1.linear.bias"] = state_dict[f"blocks.{i}.default_modulation.1.bias"] - state_dict.pop(f"blocks.{i}.norm1.weight") - state_dict.pop(f"blocks.{i}.norm1.bias") - state_dict.pop(f"blocks.{i}.default_modulation.1.weight") - state_dict.pop(f"blocks.{i}.default_modulation.1.bias") - - # mlp.fc1 -> ff.net.0, mlp.fc2 -> ff.net.2 - state_dict[f"blocks.{i}.ff.net.0.proj.weight"] = state_dict[f"blocks.{i}.mlp.fc1.weight"] - state_dict[f"blocks.{i}.ff.net.0.proj.bias"] = state_dict[f"blocks.{i}.mlp.fc1.bias"] - state_dict[f"blocks.{i}.ff.net.2.weight"] = state_dict[f"blocks.{i}.mlp.fc2.weight"] - state_dict[f"blocks.{i}.ff.net.2.bias"] = state_dict[f"blocks.{i}.mlp.fc2.bias"] - state_dict.pop(f"blocks.{i}.mlp.fc1.weight") - state_dict.pop(f"blocks.{i}.mlp.fc1.bias") - state_dict.pop(f"blocks.{i}.mlp.fc2.weight") - state_dict.pop(f"blocks.{i}.mlp.fc2.bias") - - # pooler -> time_extra_emb - state_dict["time_extra_emb.pooler.positional_embedding"] = state_dict["pooler.positional_embedding"] - state_dict["time_extra_emb.pooler.k_proj.weight"] = state_dict["pooler.k_proj.weight"] - state_dict["time_extra_emb.pooler.k_proj.bias"] = state_dict["pooler.k_proj.bias"] - state_dict["time_extra_emb.pooler.q_proj.weight"] = state_dict["pooler.q_proj.weight"] - state_dict["time_extra_emb.pooler.q_proj.bias"] = state_dict["pooler.q_proj.bias"] - state_dict["time_extra_emb.pooler.v_proj.weight"] = state_dict["pooler.v_proj.weight"] - state_dict["time_extra_emb.pooler.v_proj.bias"] = state_dict["pooler.v_proj.bias"] - state_dict["time_extra_emb.pooler.c_proj.weight"] = state_dict["pooler.c_proj.weight"] - state_dict["time_extra_emb.pooler.c_proj.bias"] = state_dict["pooler.c_proj.bias"] - state_dict.pop("pooler.k_proj.weight") - state_dict.pop("pooler.k_proj.bias") - state_dict.pop("pooler.q_proj.weight") - state_dict.pop("pooler.q_proj.bias") - state_dict.pop("pooler.v_proj.weight") - state_dict.pop("pooler.v_proj.bias") - state_dict.pop("pooler.c_proj.weight") - state_dict.pop("pooler.c_proj.bias") - state_dict.pop("pooler.positional_embedding") - - # t_embedder -> time_embedding (`TimestepEmbedding`) - state_dict["time_extra_emb.timestep_embedder.linear_1.bias"] = state_dict["t_embedder.mlp.0.bias"] - state_dict["time_extra_emb.timestep_embedder.linear_1.weight"] = state_dict["t_embedder.mlp.0.weight"] - state_dict["time_extra_emb.timestep_embedder.linear_2.bias"] = state_dict["t_embedder.mlp.2.bias"] - state_dict["time_extra_emb.timestep_embedder.linear_2.weight"] = state_dict["t_embedder.mlp.2.weight"] - - state_dict.pop("t_embedder.mlp.0.bias") - state_dict.pop("t_embedder.mlp.0.weight") - state_dict.pop("t_embedder.mlp.2.bias") - state_dict.pop("t_embedder.mlp.2.weight") - - # x_embedder -> pos_embd (`PatchEmbed`) - state_dict["pos_embed.proj.weight"] = state_dict["x_embedder.proj.weight"] - state_dict["pos_embed.proj.bias"] = state_dict["x_embedder.proj.bias"] - state_dict.pop("x_embedder.proj.weight") - state_dict.pop("x_embedder.proj.bias") - - # mlp_t5 -> text_embedder - state_dict["text_embedder.linear_1.bias"] = state_dict["mlp_t5.0.bias"] - state_dict["text_embedder.linear_1.weight"] = state_dict["mlp_t5.0.weight"] - state_dict["text_embedder.linear_2.bias"] = state_dict["mlp_t5.2.bias"] - state_dict["text_embedder.linear_2.weight"] = state_dict["mlp_t5.2.weight"] - state_dict.pop("mlp_t5.0.bias") - state_dict.pop("mlp_t5.0.weight") - state_dict.pop("mlp_t5.2.bias") - state_dict.pop("mlp_t5.2.weight") - - # extra_embedder -> extra_embedder - state_dict["time_extra_emb.extra_embedder.linear_1.bias"] = state_dict["extra_embedder.0.bias"] - state_dict["time_extra_emb.extra_embedder.linear_1.weight"] = state_dict["extra_embedder.0.weight"] - state_dict["time_extra_emb.extra_embedder.linear_2.bias"] = state_dict["extra_embedder.2.bias"] - state_dict["time_extra_emb.extra_embedder.linear_2.weight"] = state_dict["extra_embedder.2.weight"] - state_dict.pop("extra_embedder.0.bias") - state_dict.pop("extra_embedder.0.weight") - state_dict.pop("extra_embedder.2.bias") - state_dict.pop("extra_embedder.2.weight") - - # model.final_adaLN_modulation.1 -> norm_out.linear - def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - state_dict["norm_out.linear.weight"] = swap_scale_shift(state_dict["final_layer.adaLN_modulation.1.weight"]) - state_dict["norm_out.linear.bias"] = swap_scale_shift(state_dict["final_layer.adaLN_modulation.1.bias"]) - state_dict.pop("final_layer.adaLN_modulation.1.weight") - state_dict.pop("final_layer.adaLN_modulation.1.bias") - - # final_linear -> proj_out - state_dict["proj_out.weight"] = state_dict["final_layer.linear.weight"] - state_dict["proj_out.bias"] = state_dict["final_layer.linear.bias"] - state_dict.pop("final_layer.linear.weight") - state_dict.pop("final_layer.linear.bias") - - # style_embedder - if model_config["use_style_cond_and_image_meta_size"]: - print(state_dict["style_embedder.weight"]) - print(state_dict["style_embedder.weight"].shape) - state_dict["time_extra_emb.style_embedder.weight"] = state_dict["style_embedder.weight"][0:1] - state_dict.pop("style_embedder.weight") - - model.load_state_dict(state_dict) - - from diffusers import HunyuanDiTPipeline - - if args.use_style_cond_and_image_meta_size: - pipe = HunyuanDiTPipeline.from_pretrained( - "Tencent-Hunyuan/HunyuanDiT-Diffusers", transformer=model, torch_dtype=torch.float32 - ) - else: - pipe = HunyuanDiTPipeline.from_pretrained( - "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", transformer=model, torch_dtype=torch.float32 - ) - pipe.to("cuda") - pipe.to(dtype=torch.float32) - - if args.save: - pipe.save_pretrained(args.output_checkpoint_path) - - # ### NOTE: HunyuanDiT supports both Chinese and English inputs - prompt = "一个宇航员在骑马" - # prompt = "An astronaut riding a horse" - generator = torch.Generator(device="cuda").manual_seed(0) - image = pipe( - height=1024, width=1024, prompt=prompt, generator=generator, num_inference_steps=25, guidance_scale=5.0 - ).images[0] - - image.save("img.png") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." - ) - parser.add_argument( - "--pt_checkpoint_path", default=None, type=str, required=True, help="Path to the .pt pretrained model." - ) - parser.add_argument( - "--output_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the output converted diffusers pipeline.", - ) - parser.add_argument( - "--load_key", default="none", type=str, required=False, help="The key to load from the pretrained .pt file" - ) - parser.add_argument( - "--use_style_cond_and_image_meta_size", - type=bool, - default=False, - help="version <= v1.1: True; version >= v1.2: False", - ) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_i2vgen_to_diffusers.py b/scripts/convert_i2vgen_to_diffusers.py deleted file mode 100644 index 643780caac2d..000000000000 --- a/scripts/convert_i2vgen_to_diffusers.py +++ /dev/null @@ -1,510 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the LDM checkpoints.""" - -import argparse - -import torch -from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection - -from diffusers import DDIMScheduler, I2VGenXLPipeline, I2VGenXLUNet, StableDiffusionPipeline - - -CLIP_ID = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - weight = old_checkpoint[path["old"]] - names = ["proj_attn.weight"] - names_2 = ["proj_out.weight", "proj_in.weight"] - if any(k in new_path for k in names): - checkpoint[new_path] = weight[:, :, 0] - elif any(k in new_path for k in names_2) and len(weight.shape) > 2 and ".attentions." not in new_path: - checkpoint[new_path] = weight[:, :, 0] - else: - checkpoint[new_path] = weight - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_temp_conv_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - mapping.append({"old": old_item, "new": old_item}) - - return mapping - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - if "temopral_conv" not in old_item: - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has both EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - additional_embedding_substrings = [ - "local_image_concat", - "context_embedding", - "local_image_embedding", - "fps_embedding", - ] - for k in unet_state_dict: - if any(substring in k for substring in additional_embedding_substrings): - diffusers_key = k.replace("local_image_concat", "image_latents_proj_in").replace( - "local_image_embedding", "image_latents_context_embedding" - ) - new_checkpoint[diffusers_key] = unet_state_dict[k] - - # temporal encoder. - new_checkpoint["image_latents_temporal_encoder.norm1.weight"] = unet_state_dict[ - "local_temporal_encoder.layers.0.0.norm.weight" - ] - new_checkpoint["image_latents_temporal_encoder.norm1.bias"] = unet_state_dict[ - "local_temporal_encoder.layers.0.0.norm.bias" - ] - - # attention - qkv = unet_state_dict["local_temporal_encoder.layers.0.0.fn.to_qkv.weight"] - q, k, v = torch.chunk(qkv, 3, dim=0) - new_checkpoint["image_latents_temporal_encoder.attn1.to_q.weight"] = q - new_checkpoint["image_latents_temporal_encoder.attn1.to_k.weight"] = k - new_checkpoint["image_latents_temporal_encoder.attn1.to_v.weight"] = v - new_checkpoint["image_latents_temporal_encoder.attn1.to_out.0.weight"] = unet_state_dict[ - "local_temporal_encoder.layers.0.0.fn.to_out.0.weight" - ] - new_checkpoint["image_latents_temporal_encoder.attn1.to_out.0.bias"] = unet_state_dict[ - "local_temporal_encoder.layers.0.0.fn.to_out.0.bias" - ] - - # feedforward - new_checkpoint["image_latents_temporal_encoder.ff.net.0.proj.weight"] = unet_state_dict[ - "local_temporal_encoder.layers.0.1.net.0.0.weight" - ] - new_checkpoint["image_latents_temporal_encoder.ff.net.0.proj.bias"] = unet_state_dict[ - "local_temporal_encoder.layers.0.1.net.0.0.bias" - ] - new_checkpoint["image_latents_temporal_encoder.ff.net.2.weight"] = unet_state_dict[ - "local_temporal_encoder.layers.0.1.net.2.weight" - ] - new_checkpoint["image_latents_temporal_encoder.ff.net.2.bias"] = unet_state_dict[ - "local_temporal_encoder.layers.0.1.net.2.bias" - ] - - if "class_embed_type" in config: - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - first_temp_attention = [v for v in unet_state_dict if v.startswith("input_blocks.0.1")] - paths = renew_attention_paths(first_temp_attention) - meta_path = {"old": "input_blocks.0.1", "new": "transformer_in"} - assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config) - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - temp_attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.2" in key] - - if f"input_blocks.{i}.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temporal_convs = [key for key in resnets if "temopral_conv" in key] - paths = renew_temp_conv_paths(temporal_convs) - meta_path = { - "old": f"input_blocks.{i}.0.temopral_conv", - "new": f"down_blocks.{block_id}.temp_convs.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(temp_attentions): - paths = renew_attention_paths(temp_attentions) - meta_path = { - "old": f"input_blocks.{i}.2", - "new": f"down_blocks.{block_id}.temp_attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - temporal_convs_0 = [key for key in resnet_0 if "temopral_conv" in key] - attentions = middle_blocks[1] - temp_attentions = middle_blocks[2] - resnet_1 = middle_blocks[3] - temporal_convs_1 = [key for key in resnet_1 if "temopral_conv" in key] - - resnet_0_paths = renew_resnet_paths(resnet_0) - meta_path = {"old": "middle_block.0", "new": "mid_block.resnets.0"} - assign_to_checkpoint( - resnet_0_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - temp_conv_0_paths = renew_temp_conv_paths(temporal_convs_0) - meta_path = {"old": "middle_block.0.temopral_conv", "new": "mid_block.temp_convs.0"} - assign_to_checkpoint( - temp_conv_0_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - resnet_1_paths = renew_resnet_paths(resnet_1) - meta_path = {"old": "middle_block.3", "new": "mid_block.resnets.1"} - assign_to_checkpoint( - resnet_1_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - temp_conv_1_paths = renew_temp_conv_paths(temporal_convs_1) - meta_path = {"old": "middle_block.3.temopral_conv", "new": "mid_block.temp_convs.1"} - assign_to_checkpoint( - temp_conv_1_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temp_attentions_paths = renew_attention_paths(temp_attentions) - meta_path = {"old": "middle_block.2", "new": "mid_block.temp_attentions.0"} - assign_to_checkpoint( - temp_attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - temp_attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.2" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temporal_convs = [key for key in resnets if "temopral_conv" in key] - paths = renew_temp_conv_paths(temporal_convs) - meta_path = { - "old": f"output_blocks.{i}.0.temopral_conv", - "new": f"up_blocks.{block_id}.temp_convs.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(temp_attentions): - paths = renew_attention_paths(temp_attentions) - meta_path = { - "old": f"output_blocks.{i}.2", - "new": f"up_blocks.{block_id}.temp_attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - new_checkpoint[new_path] = unet_state_dict[old_path] - - temopral_conv_paths = [l for l in output_block_layers if "temopral_conv" in l] - for path in temopral_conv_paths: - pruned_path = path.split("temopral_conv.")[-1] - old_path = ".".join(["output_blocks", str(i), str(block_id), "temopral_conv", pruned_path]) - new_path = ".".join(["up_blocks", str(block_id), "temp_convs", str(layer_in_block_id), pruned_path]) - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--unet_checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--push_to_hub", action="store_true") - args = parser.parse_args() - - # UNet - unet_checkpoint = torch.load(args.unet_checkpoint_path, map_location="cpu") - unet_checkpoint = unet_checkpoint["state_dict"] - unet = I2VGenXLUNet(sample_size=32) - - converted_ckpt = convert_ldm_unet_checkpoint(unet_checkpoint, unet.config) - - diff_0 = set(unet.state_dict().keys()) - set(converted_ckpt.keys()) - diff_1 = set(converted_ckpt.keys()) - set(unet.state_dict().keys()) - - assert len(diff_0) == len(diff_1) == 0, "Converted weights don't match" - - unet.load_state_dict(converted_ckpt, strict=True) - - # vae - temp_pipe = StableDiffusionPipeline.from_single_file( - "https://huggingface.co/ali-vilab/i2vgen-xl/blob/main/models/v2-1_512-ema-pruned.ckpt" - ) - vae = temp_pipe.vae - del temp_pipe - - # text encoder and tokenizer - text_encoder = CLIPTextModel.from_pretrained(CLIP_ID) - tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID) - - # image encoder and feature extractor - image_encoder = CLIPVisionModelWithProjection.from_pretrained(CLIP_ID) - feature_extractor = CLIPImageProcessor.from_pretrained(CLIP_ID) - - # scheduler - # https://github.com/ali-vilab/i2vgen-xl/blob/main/configs/i2vgen_xl_train.yaml - scheduler = DDIMScheduler( - beta_schedule="squaredcos_cap_v2", - rescale_betas_zero_snr=True, - set_alpha_to_one=True, - clip_sample=False, - steps_offset=1, - timestep_spacing="leading", - prediction_type="v_prediction", - ) - - # final - pipeline = I2VGenXLPipeline( - unet=unet, - vae=vae, - image_encoder=image_encoder, - feature_extractor=feature_extractor, - text_encoder=text_encoder, - tokenizer=tokenizer, - scheduler=scheduler, - ) - - pipeline.save_pretrained(args.dump_path, push_to_hub=args.push_to_hub) diff --git a/scripts/convert_if.py b/scripts/convert_if.py deleted file mode 100644 index 85c739ca92f0..000000000000 --- a/scripts/convert_if.py +++ /dev/null @@ -1,1250 +0,0 @@ -import argparse -import inspect -import os - -import numpy as np -import torch -import yaml -from torch.nn import functional as F -from transformers import CLIPConfig, CLIPImageProcessor, CLIPVisionModelWithProjection, T5EncoderModel, T5Tokenizer - -from diffusers import DDPMScheduler, IFPipeline, IFSuperResolutionPipeline, UNet2DConditionModel -from diffusers.pipelines.deepfloyd_if.safety_checker import IFSafetyChecker - - -def parse_args(): - parser = argparse.ArgumentParser() - - parser.add_argument("--dump_path", required=False, default=None, type=str) - - parser.add_argument("--dump_path_stage_2", required=False, default=None, type=str) - - parser.add_argument("--dump_path_stage_3", required=False, default=None, type=str) - - parser.add_argument("--unet_config", required=False, default=None, type=str, help="Path to unet config file") - - parser.add_argument( - "--unet_checkpoint_path", required=False, default=None, type=str, help="Path to unet checkpoint file" - ) - - parser.add_argument( - "--unet_checkpoint_path_stage_2", - required=False, - default=None, - type=str, - help="Path to stage 2 unet checkpoint file", - ) - - parser.add_argument( - "--unet_checkpoint_path_stage_3", - required=False, - default=None, - type=str, - help="Path to stage 3 unet checkpoint file", - ) - - parser.add_argument("--p_head_path", type=str, required=True) - - parser.add_argument("--w_head_path", type=str, required=True) - - args = parser.parse_args() - - return args - - -def main(args): - tokenizer = T5Tokenizer.from_pretrained("google/t5-v1_1-xxl") - text_encoder = T5EncoderModel.from_pretrained("google/t5-v1_1-xxl") - - feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14") - safety_checker = convert_safety_checker(p_head_path=args.p_head_path, w_head_path=args.w_head_path) - - if args.unet_config is not None and args.unet_checkpoint_path is not None and args.dump_path is not None: - convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args) - - if args.unet_checkpoint_path_stage_2 is not None and args.dump_path_stage_2 is not None: - convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=2) - - if args.unet_checkpoint_path_stage_3 is not None and args.dump_path_stage_3 is not None: - convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=3) - - -def convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args): - unet = get_stage_1_unet(args.unet_config, args.unet_checkpoint_path) - - scheduler = DDPMScheduler( - variance_type="learned_range", - beta_schedule="squaredcos_cap_v2", - prediction_type="epsilon", - thresholding=True, - dynamic_thresholding_ratio=0.95, - sample_max_value=1.5, - ) - - pipe = IFPipeline( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - requires_safety_checker=True, - ) - - pipe.save_pretrained(args.dump_path) - - -def convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage): - if stage == 2: - unet_checkpoint_path = args.unet_checkpoint_path_stage_2 - sample_size = None - dump_path = args.dump_path_stage_2 - elif stage == 3: - unet_checkpoint_path = args.unet_checkpoint_path_stage_3 - sample_size = 1024 - dump_path = args.dump_path_stage_3 - else: - assert False - - unet = get_super_res_unet(unet_checkpoint_path, verify_param_count=False, sample_size=sample_size) - - image_noising_scheduler = DDPMScheduler( - beta_schedule="squaredcos_cap_v2", - ) - - scheduler = DDPMScheduler( - variance_type="learned_range", - beta_schedule="squaredcos_cap_v2", - prediction_type="epsilon", - thresholding=True, - dynamic_thresholding_ratio=0.95, - sample_max_value=1.0, - ) - - pipe = IFSuperResolutionPipeline( - tokenizer=tokenizer, - text_encoder=text_encoder, - unet=unet, - scheduler=scheduler, - image_noising_scheduler=image_noising_scheduler, - safety_checker=safety_checker, - feature_extractor=feature_extractor, - requires_safety_checker=True, - ) - - pipe.save_pretrained(dump_path) - - -def get_stage_1_unet(unet_config, unet_checkpoint_path): - original_unet_config = yaml.safe_load(unet_config) - original_unet_config = original_unet_config["params"] - - unet_diffusers_config = create_unet_diffusers_config(original_unet_config) - - unet = UNet2DConditionModel(**unet_diffusers_config) - - device = "cuda" if torch.cuda.is_available() else "cpu" - unet_checkpoint = torch.load(unet_checkpoint_path, map_location=device) - - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - unet_checkpoint, unet_diffusers_config, path=unet_checkpoint_path - ) - - unet.load_state_dict(converted_unet_checkpoint) - - return unet - - -def convert_safety_checker(p_head_path, w_head_path): - state_dict = {} - - # p head - - p_head = np.load(p_head_path) - - p_head_weights = p_head["weights"] - p_head_weights = torch.from_numpy(p_head_weights) - p_head_weights = p_head_weights.unsqueeze(0) - - p_head_biases = p_head["biases"] - p_head_biases = torch.from_numpy(p_head_biases) - p_head_biases = p_head_biases.unsqueeze(0) - - state_dict["p_head.weight"] = p_head_weights - state_dict["p_head.bias"] = p_head_biases - - # w head - - w_head = np.load(w_head_path) - - w_head_weights = w_head["weights"] - w_head_weights = torch.from_numpy(w_head_weights) - w_head_weights = w_head_weights.unsqueeze(0) - - w_head_biases = w_head["biases"] - w_head_biases = torch.from_numpy(w_head_biases) - w_head_biases = w_head_biases.unsqueeze(0) - - state_dict["w_head.weight"] = w_head_weights - state_dict["w_head.bias"] = w_head_biases - - # vision model - - vision_model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") - vision_model_state_dict = vision_model.state_dict() - - for key, value in vision_model_state_dict.items(): - key = f"vision_model.{key}" - state_dict[key] = value - - # full model - - config = CLIPConfig.from_pretrained("openai/clip-vit-large-patch14") - safety_checker = IFSafetyChecker(config) - - safety_checker.load_state_dict(state_dict) - - return safety_checker - - -def create_unet_diffusers_config(original_unet_config, class_embed_type=None): - attention_resolutions = parse_list(original_unet_config["attention_resolutions"]) - attention_resolutions = [original_unet_config["image_size"] // int(res) for res in attention_resolutions] - - channel_mult = parse_list(original_unet_config["channel_mult"]) - block_out_channels = [original_unet_config["model_channels"] * mult for mult in channel_mult] - - down_block_types = [] - resolution = 1 - - for i in range(len(block_out_channels)): - if resolution in attention_resolutions: - block_type = "SimpleCrossAttnDownBlock2D" - elif original_unet_config["resblock_updown"]: - block_type = "ResnetDownsampleBlock2D" - else: - block_type = "DownBlock2D" - - down_block_types.append(block_type) - - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - if resolution in attention_resolutions: - block_type = "SimpleCrossAttnUpBlock2D" - elif original_unet_config["resblock_updown"]: - block_type = "ResnetUpsampleBlock2D" - else: - block_type = "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - head_dim = original_unet_config["num_head_channels"] - - use_linear_projection = ( - original_unet_config["use_linear_in_transformer"] - if "use_linear_in_transformer" in original_unet_config - else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim = [5, 10, 20, 20] - - projection_class_embeddings_input_dim = None - - if class_embed_type is None: - if "num_classes" in original_unet_config: - if original_unet_config["num_classes"] == "sequential": - class_embed_type = "projection" - assert "adm_in_channels" in original_unet_config - projection_class_embeddings_input_dim = original_unet_config["adm_in_channels"] - else: - raise NotImplementedError( - f"Unknown conditional unet num_classes config: {original_unet_config['num_classes']}" - ) - - config = { - "sample_size": original_unet_config["image_size"], - "in_channels": original_unet_config["in_channels"], - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": original_unet_config["num_res_blocks"], - "cross_attention_dim": original_unet_config["encoder_channels"], - "attention_head_dim": head_dim, - "use_linear_projection": use_linear_projection, - "class_embed_type": class_embed_type, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "out_channels": original_unet_config["out_channels"], - "up_block_types": tuple(up_block_types), - "upcast_attention": False, # TODO: guessing - "cross_attention_norm": "group_norm", - "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", - "addition_embed_type": "text", - "act_fn": "gelu", - } - - if original_unet_config["use_scale_shift_norm"]: - config["resnet_time_scale_shift"] = "scale_shift" - - if "encoder_dim" in original_unet_config: - config["encoder_hid_dim"] = original_unet_config["encoder_dim"] - - return config - - -def convert_ldm_unet_checkpoint(unet_state_dict, config, path=None): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] in [None, "identity"]: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - - # TODO need better check than i in [4, 8, 12, 16] - block_type = config["down_block_types"][block_id] - if (block_type == "ResnetDownsampleBlock2D" or block_type == "SimpleCrossAttnDownBlock2D") and i in [ - 4, - 8, - 12, - 16, - ]: - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.downsamplers.0"} - else: - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - old_path = f"input_blocks.{i}.1" - new_path = f"down_blocks.{block_id}.attentions.{layer_in_block_id}" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - paths = renew_attention_paths(attentions) - meta_path = {"old": old_path, "new": new_path} - assign_to_checkpoint( - paths, - new_checkpoint, - unet_state_dict, - additional_replacements=[meta_path], - config=config, - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - old_path = "middle_block.1" - new_path = "mid_block.attentions.0" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - # len(output_block_list) == 1 -> resnet - # len(output_block_list) == 2 -> resnet, attention - # len(output_block_list) == 3 -> resnet, attention, upscale resnet - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - old_path = f"output_blocks.{i}.1" - new_path = f"up_blocks.{block_id}.attentions.{layer_in_block_id}" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - paths = renew_attention_paths(attentions) - meta_path = { - "old": old_path, - "new": new_path, - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(output_block_list) == 3: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.2" in key] - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"output_blocks.{i}.2", "new": f"up_blocks.{block_id}.upsamplers.0"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - if "encoder_proj.weight" in unet_state_dict: - new_checkpoint["encoder_hid_proj.weight"] = unet_state_dict.pop("encoder_proj.weight") - new_checkpoint["encoder_hid_proj.bias"] = unet_state_dict.pop("encoder_proj.bias") - - if "encoder_pooling.0.weight" in unet_state_dict: - new_checkpoint["add_embedding.norm1.weight"] = unet_state_dict.pop("encoder_pooling.0.weight") - new_checkpoint["add_embedding.norm1.bias"] = unet_state_dict.pop("encoder_pooling.0.bias") - - new_checkpoint["add_embedding.pool.positional_embedding"] = unet_state_dict.pop( - "encoder_pooling.1.positional_embedding" - ) - new_checkpoint["add_embedding.pool.k_proj.weight"] = unet_state_dict.pop("encoder_pooling.1.k_proj.weight") - new_checkpoint["add_embedding.pool.k_proj.bias"] = unet_state_dict.pop("encoder_pooling.1.k_proj.bias") - new_checkpoint["add_embedding.pool.q_proj.weight"] = unet_state_dict.pop("encoder_pooling.1.q_proj.weight") - new_checkpoint["add_embedding.pool.q_proj.bias"] = unet_state_dict.pop("encoder_pooling.1.q_proj.bias") - new_checkpoint["add_embedding.pool.v_proj.weight"] = unet_state_dict.pop("encoder_pooling.1.v_proj.weight") - new_checkpoint["add_embedding.pool.v_proj.bias"] = unet_state_dict.pop("encoder_pooling.1.v_proj.bias") - - new_checkpoint["add_embedding.proj.weight"] = unet_state_dict.pop("encoder_pooling.2.weight") - new_checkpoint["add_embedding.proj.bias"] = unet_state_dict.pop("encoder_pooling.2.bias") - - new_checkpoint["add_embedding.norm2.weight"] = unet_state_dict.pop("encoder_pooling.3.weight") - new_checkpoint["add_embedding.norm2.bias"] = unet_state_dict.pop("encoder_pooling.3.bias") - - return new_checkpoint - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - if "qkv" in new_item: - continue - - if "encoder_kv" in new_item: - continue - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = new_item.replace("norm_encoder.weight", "norm_cross.weight") - new_item = new_item.replace("norm_encoder.bias", "norm_cross.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_attention_to_checkpoint(new_checkpoint, unet_state_dict, old_path, new_path, config): - qkv_weight = unet_state_dict.pop(f"{old_path}.qkv.weight") - qkv_weight = qkv_weight[:, :, 0] - - qkv_bias = unet_state_dict.pop(f"{old_path}.qkv.bias") - - is_cross_attn_only = "only_cross_attention" in config and config["only_cross_attention"] - - split = 1 if is_cross_attn_only else 3 - - weights, bias = split_attentions( - weight=qkv_weight, - bias=qkv_bias, - split=split, - chunk_size=config["attention_head_dim"], - ) - - if is_cross_attn_only: - query_weight, q_bias = weights, bias - new_checkpoint[f"{new_path}.to_q.weight"] = query_weight[0] - new_checkpoint[f"{new_path}.to_q.bias"] = q_bias[0] - else: - [query_weight, key_weight, value_weight], [q_bias, k_bias, v_bias] = weights, bias - new_checkpoint[f"{new_path}.to_q.weight"] = query_weight - new_checkpoint[f"{new_path}.to_q.bias"] = q_bias - new_checkpoint[f"{new_path}.to_k.weight"] = key_weight - new_checkpoint[f"{new_path}.to_k.bias"] = k_bias - new_checkpoint[f"{new_path}.to_v.weight"] = value_weight - new_checkpoint[f"{new_path}.to_v.bias"] = v_bias - - encoder_kv_weight = unet_state_dict.pop(f"{old_path}.encoder_kv.weight") - encoder_kv_weight = encoder_kv_weight[:, :, 0] - - encoder_kv_bias = unet_state_dict.pop(f"{old_path}.encoder_kv.bias") - - [encoder_k_weight, encoder_v_weight], [encoder_k_bias, encoder_v_bias] = split_attentions( - weight=encoder_kv_weight, - bias=encoder_kv_bias, - split=2, - chunk_size=config["attention_head_dim"], - ) - - new_checkpoint[f"{new_path}.add_k_proj.weight"] = encoder_k_weight - new_checkpoint[f"{new_path}.add_k_proj.bias"] = encoder_k_bias - new_checkpoint[f"{new_path}.add_v_proj.weight"] = encoder_v_weight - new_checkpoint[f"{new_path}.add_v_proj.bias"] = encoder_v_bias - - -def assign_to_checkpoint(paths, checkpoint, old_checkpoint, additional_replacements=None, config=None): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - for path in paths: - new_path = path["new"] - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path or "to_out.0.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) -def split_attentions(*, weight, bias, split, chunk_size): - weights = [None] * split - biases = [None] * split - - weights_biases_idx = 0 - - for starting_row_index in range(0, weight.shape[0], chunk_size): - row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size) - - weight_rows = weight[row_indices, :] - bias_rows = bias[row_indices] - - if weights[weights_biases_idx] is None: - weights[weights_biases_idx] = weight_rows - biases[weights_biases_idx] = bias_rows - else: - assert weights[weights_biases_idx] is not None - weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows]) - biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows]) - - weights_biases_idx = (weights_biases_idx + 1) % split - - return weights, biases - - -def parse_list(value): - if isinstance(value, str): - value = value.split(",") - value = [int(v) for v in value] - elif isinstance(value, list): - pass - else: - raise ValueError(f"Can't parse list for type: {type(value)}") - - return value - - -# below is copy and pasted from original convert_if_stage_2.py script - - -def get_super_res_unet(unet_checkpoint_path, verify_param_count=True, sample_size=None): - orig_path = unet_checkpoint_path - - original_unet_config = yaml.safe_load(os.path.join(orig_path, "config.yml")) - original_unet_config = original_unet_config["params"] - - unet_diffusers_config = superres_create_unet_diffusers_config(original_unet_config) - unet_diffusers_config["time_embedding_dim"] = original_unet_config["model_channels"] * int( - original_unet_config["channel_mult"].split(",")[-1] - ) - if original_unet_config["encoder_dim"] != original_unet_config["encoder_channels"]: - unet_diffusers_config["encoder_hid_dim"] = original_unet_config["encoder_dim"] - unet_diffusers_config["class_embed_type"] = "timestep" - unet_diffusers_config["addition_embed_type"] = "text" - - unet_diffusers_config["time_embedding_act_fn"] = "gelu" - unet_diffusers_config["resnet_skip_time_act"] = True - unet_diffusers_config["resnet_out_scale_factor"] = 1 / 0.7071 - unet_diffusers_config["mid_block_scale_factor"] = 1 / 0.7071 - unet_diffusers_config["only_cross_attention"] = ( - bool(original_unet_config["disable_self_attentions"]) - if ( - "disable_self_attentions" in original_unet_config - and isinstance(original_unet_config["disable_self_attentions"], int) - ) - else True - ) - - if sample_size is None: - unet_diffusers_config["sample_size"] = original_unet_config["image_size"] - else: - # The second upscaler unet's sample size is incorrectly specified - # in the config and is instead hardcoded in source - unet_diffusers_config["sample_size"] = sample_size - - unet_checkpoint = torch.load(os.path.join(unet_checkpoint_path, "pytorch_model.bin"), map_location="cpu") - - if verify_param_count: - # check that architecture matches - is a bit slow - verify_param_count(orig_path, unet_diffusers_config) - - converted_unet_checkpoint = superres_convert_ldm_unet_checkpoint( - unet_checkpoint, unet_diffusers_config, path=unet_checkpoint_path - ) - converted_keys = converted_unet_checkpoint.keys() - - model = UNet2DConditionModel(**unet_diffusers_config) - expected_weights = model.state_dict().keys() - - diff_c_e = set(converted_keys) - set(expected_weights) - diff_e_c = set(expected_weights) - set(converted_keys) - - assert len(diff_e_c) == 0, f"Expected, but not converted: {diff_e_c}" - assert len(diff_c_e) == 0, f"Converted, but not expected: {diff_c_e}" - - model.load_state_dict(converted_unet_checkpoint) - - return model - - -def superres_create_unet_diffusers_config(original_unet_config): - attention_resolutions = parse_list(original_unet_config["attention_resolutions"]) - attention_resolutions = [original_unet_config["image_size"] // int(res) for res in attention_resolutions] - - channel_mult = parse_list(original_unet_config["channel_mult"]) - block_out_channels = [original_unet_config["model_channels"] * mult for mult in channel_mult] - - down_block_types = [] - resolution = 1 - - for i in range(len(block_out_channels)): - if resolution in attention_resolutions: - block_type = "SimpleCrossAttnDownBlock2D" - elif original_unet_config["resblock_updown"]: - block_type = "ResnetDownsampleBlock2D" - else: - block_type = "DownBlock2D" - - down_block_types.append(block_type) - - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - if resolution in attention_resolutions: - block_type = "SimpleCrossAttnUpBlock2D" - elif original_unet_config["resblock_updown"]: - block_type = "ResnetUpsampleBlock2D" - else: - block_type = "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - head_dim = original_unet_config["num_head_channels"] - use_linear_projection = ( - original_unet_config["use_linear_in_transformer"] - if "use_linear_in_transformer" in original_unet_config - else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim = [5, 10, 20, 20] - - class_embed_type = None - projection_class_embeddings_input_dim = None - - if "num_classes" in original_unet_config: - if original_unet_config["num_classes"] == "sequential": - class_embed_type = "projection" - assert "adm_in_channels" in original_unet_config - projection_class_embeddings_input_dim = original_unet_config["adm_in_channels"] - else: - raise NotImplementedError( - f"Unknown conditional unet num_classes config: {original_unet_config['num_classes']}" - ) - - config = { - "in_channels": original_unet_config["in_channels"], - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": tuple(original_unet_config["num_res_blocks"]), - "cross_attention_dim": original_unet_config["encoder_channels"], - "attention_head_dim": head_dim, - "use_linear_projection": use_linear_projection, - "class_embed_type": class_embed_type, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "out_channels": original_unet_config["out_channels"], - "up_block_types": tuple(up_block_types), - "upcast_attention": False, # TODO: guessing - "cross_attention_norm": "group_norm", - "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", - "act_fn": "gelu", - } - - if original_unet_config["use_scale_shift_norm"]: - config["resnet_time_scale_shift"] = "scale_shift" - - return config - - -def superres_convert_ldm_unet_checkpoint(unet_state_dict, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["aug_proj.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["aug_proj.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["aug_proj.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["aug_proj.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - if "encoder_proj.weight" in unet_state_dict: - new_checkpoint["encoder_hid_proj.weight"] = unet_state_dict["encoder_proj.weight"] - new_checkpoint["encoder_hid_proj.bias"] = unet_state_dict["encoder_proj.bias"] - - if "encoder_pooling.0.weight" in unet_state_dict: - mapping = { - "encoder_pooling.0": "add_embedding.norm1", - "encoder_pooling.1": "add_embedding.pool", - "encoder_pooling.2": "add_embedding.proj", - "encoder_pooling.3": "add_embedding.norm2", - } - for key in unet_state_dict.keys(): - if key.startswith("encoder_pooling"): - prefix = key[: len("encoder_pooling.0")] - new_key = key.replace(prefix, mapping[prefix]) - new_checkpoint[new_key] = unet_state_dict[key] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key] - for layer_id in range(num_output_blocks) - } - if not isinstance(config["layers_per_block"], int): - layers_per_block_list = [e + 1 for e in config["layers_per_block"]] - layers_per_block_cumsum = list(np.cumsum(layers_per_block_list)) - downsampler_ids = layers_per_block_cumsum - else: - # TODO need better check than i in [4, 8, 12, 16] - downsampler_ids = [4, 8, 12, 16] - - for i in range(1, num_input_blocks): - if isinstance(config["layers_per_block"], int): - layers_per_block = config["layers_per_block"] - block_id = (i - 1) // (layers_per_block + 1) - layer_in_block_id = (i - 1) % (layers_per_block + 1) - else: - block_id = next(k for k, n in enumerate(layers_per_block_cumsum) if (i - 1) < n) - passed_blocks = layers_per_block_cumsum[block_id - 1] if block_id > 0 else 0 - layer_in_block_id = (i - 1) - passed_blocks - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - - block_type = config["down_block_types"][block_id] - if ( - block_type == "ResnetDownsampleBlock2D" or block_type == "SimpleCrossAttnDownBlock2D" - ) and i in downsampler_ids: - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.downsamplers.0"} - else: - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - old_path = f"input_blocks.{i}.1" - new_path = f"down_blocks.{block_id}.attentions.{layer_in_block_id}" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - paths = renew_attention_paths(attentions) - meta_path = {"old": old_path, "new": new_path} - assign_to_checkpoint( - paths, - new_checkpoint, - unet_state_dict, - additional_replacements=[meta_path], - config=config, - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - old_path = "middle_block.1" - new_path = "mid_block.attentions.0" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - if not isinstance(config["layers_per_block"], int): - layers_per_block_list = list(reversed([e + 1 for e in config["layers_per_block"]])) - layers_per_block_cumsum = list(np.cumsum(layers_per_block_list)) - - for i in range(num_output_blocks): - if isinstance(config["layers_per_block"], int): - layers_per_block = config["layers_per_block"] - block_id = i // (layers_per_block + 1) - layer_in_block_id = i % (layers_per_block + 1) - else: - block_id = next(k for k, n in enumerate(layers_per_block_cumsum) if i < n) - passed_blocks = layers_per_block_cumsum[block_id - 1] if block_id > 0 else 0 - layer_in_block_id = i - passed_blocks - - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - # len(output_block_list) == 1 -> resnet - # len(output_block_list) == 2 -> resnet, attention or resnet, upscale resnet - # len(output_block_list) == 3 -> resnet, attention, upscale resnet - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - - has_attention = True - if len(output_block_list) == 2 and any("in_layers" in k for k in output_block_list["1"]): - has_attention = False - - maybe_attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # this layer was no attention - has_attention = False - maybe_attentions = [] - - if has_attention: - old_path = f"output_blocks.{i}.1" - new_path = f"up_blocks.{block_id}.attentions.{layer_in_block_id}" - - assign_attention_to_checkpoint( - new_checkpoint=new_checkpoint, - unet_state_dict=unet_state_dict, - old_path=old_path, - new_path=new_path, - config=config, - ) - - paths = renew_attention_paths(maybe_attentions) - meta_path = { - "old": old_path, - "new": new_path, - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(output_block_list) == 3 or (not has_attention and len(maybe_attentions) > 0): - layer_id = len(output_block_list) - 1 - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.{layer_id}" in key] - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"output_blocks.{i}.{layer_id}", "new": f"up_blocks.{block_id}.upsamplers.0"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -def verify_param_count(orig_path, unet_diffusers_config): - if "-II-" in orig_path: - from deepfloyd_if.modules import IFStageII - - if_II = IFStageII(device="cpu", dir_or_name=orig_path) - elif "-III-" in orig_path: - from deepfloyd_if.modules import IFStageIII - - if_II = IFStageIII(device="cpu", dir_or_name=orig_path) - else: - assert f"Weird name. Should have -II- or -III- in path: {orig_path}" - - unet = UNet2DConditionModel(**unet_diffusers_config) - - # in params - assert_param_count(unet.time_embedding, if_II.model.time_embed) - assert_param_count(unet.conv_in, if_II.model.input_blocks[:1]) - - # downblocks - assert_param_count(unet.down_blocks[0], if_II.model.input_blocks[1:4]) - assert_param_count(unet.down_blocks[1], if_II.model.input_blocks[4:7]) - assert_param_count(unet.down_blocks[2], if_II.model.input_blocks[7:11]) - - if "-II-" in orig_path: - assert_param_count(unet.down_blocks[3], if_II.model.input_blocks[11:17]) - assert_param_count(unet.down_blocks[4], if_II.model.input_blocks[17:]) - if "-III-" in orig_path: - assert_param_count(unet.down_blocks[3], if_II.model.input_blocks[11:15]) - assert_param_count(unet.down_blocks[4], if_II.model.input_blocks[15:20]) - assert_param_count(unet.down_blocks[5], if_II.model.input_blocks[20:]) - - # mid block - assert_param_count(unet.mid_block, if_II.model.middle_block) - - # up block - if "-II-" in orig_path: - assert_param_count(unet.up_blocks[0], if_II.model.output_blocks[:6]) - assert_param_count(unet.up_blocks[1], if_II.model.output_blocks[6:12]) - assert_param_count(unet.up_blocks[2], if_II.model.output_blocks[12:16]) - assert_param_count(unet.up_blocks[3], if_II.model.output_blocks[16:19]) - assert_param_count(unet.up_blocks[4], if_II.model.output_blocks[19:]) - if "-III-" in orig_path: - assert_param_count(unet.up_blocks[0], if_II.model.output_blocks[:5]) - assert_param_count(unet.up_blocks[1], if_II.model.output_blocks[5:10]) - assert_param_count(unet.up_blocks[2], if_II.model.output_blocks[10:14]) - assert_param_count(unet.up_blocks[3], if_II.model.output_blocks[14:18]) - assert_param_count(unet.up_blocks[4], if_II.model.output_blocks[18:21]) - assert_param_count(unet.up_blocks[5], if_II.model.output_blocks[21:24]) - - # out params - assert_param_count(unet.conv_norm_out, if_II.model.out[0]) - assert_param_count(unet.conv_out, if_II.model.out[2]) - - # make sure all model architecture has same param count - assert_param_count(unet, if_II.model) - - -def assert_param_count(model_1, model_2): - count_1 = sum(p.numel() for p in model_1.parameters()) - count_2 = sum(p.numel() for p in model_2.parameters()) - assert count_1 == count_2, f"{model_1.__class__}: {count_1} != {model_2.__class__}: {count_2}" - - -def superres_check_against_original(dump_path, unet_checkpoint_path): - model_path = dump_path - model = UNet2DConditionModel.from_pretrained(model_path) - model.to("cuda") - orig_path = unet_checkpoint_path - - if "-II-" in orig_path: - from deepfloyd_if.modules import IFStageII - - if_II_model = IFStageII(device="cuda", dir_or_name=orig_path, model_kwargs={"precision": "fp32"}).model - elif "-III-" in orig_path: - from deepfloyd_if.modules import IFStageIII - - if_II_model = IFStageIII(device="cuda", dir_or_name=orig_path, model_kwargs={"precision": "fp32"}).model - - batch_size = 1 - channels = model.config.in_channels // 2 - height = model.config.sample_size - width = model.config.sample_size - height = 1024 - width = 1024 - - torch.manual_seed(0) - - latents = torch.randn((batch_size, channels, height, width), device=model.device) - image_small = torch.randn((batch_size, channels, height // 4, width // 4), device=model.device) - - interpolate_antialias = {} - if "antialias" in inspect.signature(F.interpolate).parameters: - interpolate_antialias["antialias"] = True - image_upscaled = F.interpolate( - image_small, size=[height, width], mode="bicubic", align_corners=False, **interpolate_antialias - ) - - latent_model_input = torch.cat([latents, image_upscaled], dim=1).to(model.dtype) - t = torch.tensor([5], device=model.device).to(model.dtype) - - seq_len = 64 - encoder_hidden_states = torch.randn((batch_size, seq_len, model.config.encoder_hid_dim), device=model.device).to( - model.dtype - ) - - fake_class_labels = torch.tensor([t], device=model.device).to(model.dtype) - - with torch.no_grad(): - out = if_II_model(latent_model_input, t, aug_steps=fake_class_labels, text_emb=encoder_hidden_states) - - if_II_model.to("cpu") - del if_II_model - import gc - - torch.cuda.empty_cache() - gc.collect() - print(50 * "=") - - with torch.no_grad(): - noise_pred = model( - sample=latent_model_input, - encoder_hidden_states=encoder_hidden_states, - class_labels=fake_class_labels, - timestep=t, - ).sample - - print("Out shape", noise_pred.shape) - print("Diff", (out - noise_pred).abs().sum()) - - -if __name__ == "__main__": - main(parse_args()) diff --git a/scripts/convert_joyimage_edit_to_diffusers.py b/scripts/convert_joyimage_edit_to_diffusers.py deleted file mode 100644 index 3fe4a7b12cd3..000000000000 --- a/scripts/convert_joyimage_edit_to_diffusers.py +++ /dev/null @@ -1,355 +0,0 @@ -"""Convert JoyImage Edit / Edit Plus checkpoints to diffusers format. - -Supports both JoyImage-Edit (single-image editing) and JoyImage-Edit-Plus -(multi-image editing). The transformer weight layout is identical; only the -target model class and pipeline differ. - -Usage: - # Convert JoyImage Edit (default) - python convert_joyimage_edit_to_diffusers.py \ - --transformer_ckpt_path /path/to/transformer.pt \ - --vae_ckpt_path /path/to/vae.pt \ - --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ - --output_path /path/to/output \ - --save_pipeline - - # Convert JoyImage Edit Plus - python convert_joyimage_edit_to_diffusers.py \ - --model_type edit_plus \ - --transformer_ckpt_path /path/to/transformer.pt \ - --vae_ckpt_path /path/to/vae.pt \ - --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ - --output_path /path/to/output \ - --save_pipeline -""" - -import argparse - -import torch -from accelerate import init_empty_weights -from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration - -from diffusers import ( - AutoencoderKLWan, - JoyImageEditPipeline, - JoyImageEditTransformer3DModel, -) -from diffusers.models.transformers.transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel -from diffusers.pipelines.joyimage.pipeline_joyimage_edit_plus import JoyImageEditPlusPipeline -from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler, -) - - -def convert_vae(vae_ckpt_path): - old_state_dict = torch.load(vae_ckpt_path, weights_only=True) - new_state_dict = {} - - middle_key_mapping = { - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - attention_mapping = { - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - head_mapping = { - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - for key, value in old_state_dict.items(): - if key in middle_key_mapping: - new_state_dict[middle_key_mapping[key]] = value - elif key in attention_mapping: - new_state_dict[attention_mapping[key]] = value - elif key in head_mapping: - new_state_dict[head_mapping[key]] = value - elif key in quant_mapping: - new_state_dict[quant_mapping[key]] = value - elif key == "encoder.conv1.weight": - new_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - new_state_dict["encoder.conv_in.bias"] = value - elif key == "decoder.conv1.weight": - new_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - new_state_dict["decoder.conv_in.bias"] = value - elif key.startswith("encoder.downsamples."): - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - new_state_dict[new_key] = value - elif key.startswith("decoder.upsamples."): - parts = key.split(".") - block_idx = int(parts[2]) - - if "residual" in key: - if block_idx in [0, 1, 2]: - new_block_idx = 0 - resnet_idx = block_idx - elif block_idx in [4, 5, 6]: - new_block_idx = 1 - resnet_idx = block_idx - 4 - elif block_idx in [8, 9, 10]: - new_block_idx = 2 - resnet_idx = block_idx - 8 - elif block_idx in [12, 13, 14]: - new_block_idx = 3 - resnet_idx = block_idx - 12 - else: - new_state_dict[key] = value - continue - - if ".residual.0.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" - elif ".residual.2.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias" - elif ".residual.2.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight" - elif ".residual.3.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma" - elif ".residual.6.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias" - elif ".residual.6.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" - else: - new_key = key - new_state_dict[new_key] = value - - elif ".shortcut." in key: - if block_idx == 4: - new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") - new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - new_state_dict[new_key] = value - - elif ".resample." in key or ".time_conv." in key: - if block_idx == 3: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") - elif block_idx == 7: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") - elif block_idx == 11: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_state_dict[new_key] = value - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_state_dict[new_key] = value - else: - new_state_dict[key] = value - - with init_empty_weights(): - vae = AutoencoderKLWan() - vae.load_state_dict(new_state_dict, strict=True, assign=True) - return vae - - -TRANSFORMER_CONFIG = { - "hidden_size": 4096, - "in_channels": 16, - "num_attention_heads": 32, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "rope_dim_list": [16, 56, 56], - "text_dim": 4096, - "rope_type": "rope", - "theta": 10000, -} - - -def convert_transformer(ckpt_path: str, model_type: str = "edit"): - checkpoint = torch.load(ckpt_path, weights_only=True) - if "model" in checkpoint: - original_state_dict = checkpoint["model"] - else: - original_state_dict = checkpoint - - attn_suffixes = ( - "img_attn_qkv.", - "img_attn_q_norm.", - "img_attn_k_norm.", - "img_attn_proj.", - "txt_attn_qkv.", - "txt_attn_q_norm.", - "txt_attn_k_norm.", - "txt_attn_proj.", - ) - remapped = {} - for key, value in original_state_dict.items(): - new_key = key - if key.startswith("double_blocks."): - for suffix in attn_suffixes: - if "." + suffix in key and ".attn." + suffix not in key: - new_key = key.replace("." + suffix, ".attn." + suffix) - break - remapped[new_key] = value - - transformer_cls = ( - JoyImageEditPlusTransformer3DModel if model_type == "edit_plus" else JoyImageEditTransformer3DModel - ) - with init_empty_weights(): - transformer = transformer_cls(**TRANSFORMER_CONFIG) - transformer.load_state_dict(remapped, strict=True, assign=True) - return transformer - - -def get_args(): - parser = argparse.ArgumentParser(description="Convert JoyImage Edit / Edit Plus checkpoints to diffusers format") - parser.add_argument( - "--model_type", - type=str, - choices=["edit", "edit_plus"], - default="edit", - help="Model type: 'edit' for JoyImage-Edit, 'edit_plus' for JoyImage-Edit-Plus", - ) - parser.add_argument( - "--transformer_ckpt_path", - type=str, - default=None, - help="Path to original transformer checkpoint", - ) - parser.add_argument( - "--vae_ckpt_path", - type=str, - default=None, - help="Path to original VAE checkpoint", - ) - parser.add_argument( - "--text_encoder_path", - type=str, - default=None, - help="Path to Qwen3-VL text encoder (e.g. Qwen/Qwen3-VL-8B-Instruct)", - ) - parser.add_argument("--save_pipeline", action="store_true") - parser.add_argument( - "--output_path", - type=str, - required=True, - help="Path where converted model should be saved", - ) - parser.add_argument("--dtype", default="bf16", help="Torch dtype (fp32, fp16, bf16)") - parser.add_argument("--flow_shift", type=float, default=1.5) - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - -if __name__ == "__main__": - args = get_args() - transformer = None - vae = None - dtype = DTYPE_MAPPING[args.dtype] - - if args.save_pipeline: - assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None - assert args.text_encoder_path is not None - - if args.transformer_ckpt_path is not None: - transformer = convert_transformer(args.transformer_ckpt_path, model_type=args.model_type) - transformer = transformer.to(dtype=dtype) - if not args.save_pipeline: - transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.vae_ckpt_path is not None: - vae = convert_vae(args.vae_ckpt_path) - vae = vae.to(dtype=dtype) - if not args.save_pipeline: - vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.save_pipeline: - processor = AutoProcessor.from_pretrained(args.text_encoder_path) - text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( - args.text_encoder_path, torch_dtype=torch.bfloat16 - ).to("cuda") - tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path) - scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.flow_shift) - transformer = transformer.to("cuda") - vae = vae.to("cuda") - - if args.model_type == "edit_plus": - pipe = JoyImageEditPlusPipeline( - processor=processor, - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ).to("cuda") - else: - pipe = JoyImageEditPipeline( - processor=processor, - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ).to("cuda") - - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - processor.save_pretrained(f"{args.output_path}/processor") diff --git a/scripts/convert_k_upscaler_to_diffusers.py b/scripts/convert_k_upscaler_to_diffusers.py deleted file mode 100644 index cff845ef8099..000000000000 --- a/scripts/convert_k_upscaler_to_diffusers.py +++ /dev/null @@ -1,297 +0,0 @@ -import argparse - -import huggingface_hub -import k_diffusion as K -import torch - -from diffusers import UNet2DConditionModel - - -UPSCALER_REPO = "pcuenq/k-upscaler" - - -def resnet_to_diffusers_checkpoint(resnet, checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - rv = { - # norm1 - f"{diffusers_resnet_prefix}.norm1.linear.weight": checkpoint[f"{resnet_prefix}.main.0.mapper.weight"], - f"{diffusers_resnet_prefix}.norm1.linear.bias": checkpoint[f"{resnet_prefix}.main.0.mapper.bias"], - # conv1 - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.main.2.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.main.2.bias"], - # norm2 - f"{diffusers_resnet_prefix}.norm2.linear.weight": checkpoint[f"{resnet_prefix}.main.4.mapper.weight"], - f"{diffusers_resnet_prefix}.norm2.linear.bias": checkpoint[f"{resnet_prefix}.main.4.mapper.bias"], - # conv2 - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.main.6.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.main.6.bias"], - } - - if resnet.conv_shortcut is not None: - rv.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{resnet_prefix}.skip.weight"], - } - ) - - return rv - - -def self_attn_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix): - weight_q, weight_k, weight_v = checkpoint[f"{attention_prefix}.qkv_proj.weight"].chunk(3, dim=0) - bias_q, bias_k, bias_v = checkpoint[f"{attention_prefix}.qkv_proj.bias"].chunk(3, dim=0) - rv = { - # norm - f"{diffusers_attention_prefix}.norm1.linear.weight": checkpoint[f"{attention_prefix}.norm_in.mapper.weight"], - f"{diffusers_attention_prefix}.norm1.linear.bias": checkpoint[f"{attention_prefix}.norm_in.mapper.bias"], - # to_q - f"{diffusers_attention_prefix}.attn1.to_q.weight": weight_q.squeeze(-1).squeeze(-1), - f"{diffusers_attention_prefix}.attn1.to_q.bias": bias_q, - # to_k - f"{diffusers_attention_prefix}.attn1.to_k.weight": weight_k.squeeze(-1).squeeze(-1), - f"{diffusers_attention_prefix}.attn1.to_k.bias": bias_k, - # to_v - f"{diffusers_attention_prefix}.attn1.to_v.weight": weight_v.squeeze(-1).squeeze(-1), - f"{diffusers_attention_prefix}.attn1.to_v.bias": bias_v, - # to_out - f"{diffusers_attention_prefix}.attn1.to_out.0.weight": checkpoint[f"{attention_prefix}.out_proj.weight"] - .squeeze(-1) - .squeeze(-1), - f"{diffusers_attention_prefix}.attn1.to_out.0.bias": checkpoint[f"{attention_prefix}.out_proj.bias"], - } - - return rv - - -def cross_attn_to_diffusers_checkpoint( - checkpoint, *, diffusers_attention_prefix, diffusers_attention_index, attention_prefix -): - weight_k, weight_v = checkpoint[f"{attention_prefix}.kv_proj.weight"].chunk(2, dim=0) - bias_k, bias_v = checkpoint[f"{attention_prefix}.kv_proj.bias"].chunk(2, dim=0) - - rv = { - # norm2 (ada groupnorm) - f"{diffusers_attention_prefix}.norm{diffusers_attention_index}.linear.weight": checkpoint[ - f"{attention_prefix}.norm_dec.mapper.weight" - ], - f"{diffusers_attention_prefix}.norm{diffusers_attention_index}.linear.bias": checkpoint[ - f"{attention_prefix}.norm_dec.mapper.bias" - ], - # layernorm on encoder_hidden_state - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.norm_cross.weight": checkpoint[ - f"{attention_prefix}.norm_enc.weight" - ], - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.norm_cross.bias": checkpoint[ - f"{attention_prefix}.norm_enc.bias" - ], - # to_q - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_q.weight": checkpoint[ - f"{attention_prefix}.q_proj.weight" - ] - .squeeze(-1) - .squeeze(-1), - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_q.bias": checkpoint[ - f"{attention_prefix}.q_proj.bias" - ], - # to_k - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_k.weight": weight_k.squeeze(-1).squeeze(-1), - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_k.bias": bias_k, - # to_v - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_v.weight": weight_v.squeeze(-1).squeeze(-1), - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_v.bias": bias_v, - # to_out - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_out.0.weight": checkpoint[ - f"{attention_prefix}.out_proj.weight" - ] - .squeeze(-1) - .squeeze(-1), - f"{diffusers_attention_prefix}.attn{diffusers_attention_index}.to_out.0.bias": checkpoint[ - f"{attention_prefix}.out_proj.bias" - ], - } - - return rv - - -def block_to_diffusers_checkpoint(block, checkpoint, block_idx, block_type): - block_prefix = "inner_model.u_net.u_blocks" if block_type == "up" else "inner_model.u_net.d_blocks" - block_prefix = f"{block_prefix}.{block_idx}" - - diffusers_checkpoint = {} - - if not hasattr(block, "attentions"): - n = 1 # resnet only - elif not block.attentions[0].add_self_attention: - n = 2 # resnet -> cross-attention - else: - n = 3 # resnet -> self-attention -> cross-attention) - - for resnet_idx, resnet in enumerate(block.resnets): - # diffusers_resnet_prefix = f"{diffusers_up_block_prefix}.resnets.{resnet_idx}" - diffusers_resnet_prefix = f"{block_type}_blocks.{block_idx}.resnets.{resnet_idx}" - idx = n * resnet_idx if block_type == "up" else n * resnet_idx + 1 - resnet_prefix = f"{block_prefix}.{idx}" if block_type == "up" else f"{block_prefix}.{idx}" - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - if hasattr(block, "attentions"): - for attention_idx, attention in enumerate(block.attentions): - diffusers_attention_prefix = f"{block_type}_blocks.{block_idx}.attentions.{attention_idx}" - idx = n * attention_idx + 1 if block_type == "up" else n * attention_idx + 2 - self_attention_prefix = f"{block_prefix}.{idx}" - cross_attention_prefix = f"{block_prefix}.{idx}" - cross_attention_index = 1 if not attention.add_self_attention else 2 - idx = ( - n * attention_idx + cross_attention_index - if block_type == "up" - else n * attention_idx + cross_attention_index + 1 - ) - cross_attention_prefix = f"{block_prefix}.{idx}" - - diffusers_checkpoint.update( - cross_attn_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - diffusers_attention_index=2, - attention_prefix=cross_attention_prefix, - ) - ) - - if attention.add_self_attention is True: - diffusers_checkpoint.update( - self_attn_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - attention_prefix=self_attention_prefix, - ) - ) - - return diffusers_checkpoint - - -def unet_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # pre-processing - diffusers_checkpoint.update( - { - "conv_in.weight": checkpoint["inner_model.proj_in.weight"], - "conv_in.bias": checkpoint["inner_model.proj_in.bias"], - } - ) - - # timestep and class embedding - diffusers_checkpoint.update( - { - "time_proj.weight": checkpoint["inner_model.timestep_embed.weight"].squeeze(-1), - "time_embedding.linear_1.weight": checkpoint["inner_model.mapping.0.weight"], - "time_embedding.linear_1.bias": checkpoint["inner_model.mapping.0.bias"], - "time_embedding.linear_2.weight": checkpoint["inner_model.mapping.2.weight"], - "time_embedding.linear_2.bias": checkpoint["inner_model.mapping.2.bias"], - "time_embedding.cond_proj.weight": checkpoint["inner_model.mapping_cond.weight"], - } - ) - - # down_blocks - for down_block_idx, down_block in enumerate(model.down_blocks): - diffusers_checkpoint.update(block_to_diffusers_checkpoint(down_block, checkpoint, down_block_idx, "down")) - - # up_blocks - for up_block_idx, up_block in enumerate(model.up_blocks): - diffusers_checkpoint.update(block_to_diffusers_checkpoint(up_block, checkpoint, up_block_idx, "up")) - - # post-processing - diffusers_checkpoint.update( - { - "conv_out.weight": checkpoint["inner_model.proj_out.weight"], - "conv_out.bias": checkpoint["inner_model.proj_out.bias"], - } - ) - - return diffusers_checkpoint - - -def unet_model_from_original_config(original_config): - in_channels = original_config["input_channels"] + original_config["unet_cond_dim"] - out_channels = original_config["input_channels"] + (1 if original_config["has_variance"] else 0) - - block_out_channels = original_config["channels"] - - assert len(set(original_config["depths"])) == 1, ( - "UNet2DConditionModel currently do not support blocks with different number of layers" - ) - layers_per_block = original_config["depths"][0] - - class_labels_dim = original_config["mapping_cond_dim"] - cross_attention_dim = original_config["cross_cond_dim"] - - attn1_types = [] - attn2_types = [] - for s, c in zip(original_config["self_attn_depths"], original_config["cross_attn_depths"]): - if s: - a1 = "self" - a2 = "cross" if c else None - elif c: - a1 = "cross" - a2 = None - else: - a1 = None - a2 = None - attn1_types.append(a1) - attn2_types.append(a2) - - unet = UNet2DConditionModel( - in_channels=in_channels, - out_channels=out_channels, - down_block_types=("KDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D"), - mid_block_type=None, - up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D"), - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - act_fn="gelu", - norm_num_groups=None, - cross_attention_dim=cross_attention_dim, - attention_head_dim=64, - time_cond_proj_dim=class_labels_dim, - resnet_time_scale_shift="scale_shift", - time_embedding_type="fourier", - timestep_post_act="gelu", - conv_in_kernel=1, - conv_out_kernel=1, - ) - - return unet - - -def main(args): - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - orig_config_path = huggingface_hub.hf_hub_download(UPSCALER_REPO, "config_laion_text_cond_latent_upscaler_2.json") - orig_weights_path = huggingface_hub.hf_hub_download( - UPSCALER_REPO, "laion_text_cond_latent_upscaler_2_1_00470000_slim.pth" - ) - print(f"loading original model configuration from {orig_config_path}") - print(f"loading original model checkpoint from {orig_weights_path}") - - print("converting to diffusers unet") - orig_config = K.config.load_config(open(orig_config_path))["model"] - model = unet_model_from_original_config(orig_config) - - orig_checkpoint = torch.load(orig_weights_path, map_location=device)["model_ema"] - converted_checkpoint = unet_to_diffusers_checkpoint(model, orig_checkpoint) - - model.load_state_dict(converted_checkpoint, strict=True) - model.save_pretrained(args.dump_path) - print(f"saving converted unet model in {args.dump_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - args = parser.parse_args() - - main(args) diff --git a/scripts/convert_kakao_brain_unclip_to_diffusers.py b/scripts/convert_kakao_brain_unclip_to_diffusers.py deleted file mode 100644 index 5135eaed5b98..000000000000 --- a/scripts/convert_kakao_brain_unclip_to_diffusers.py +++ /dev/null @@ -1,1159 +0,0 @@ -import argparse -import tempfile - -import torch -from accelerate import load_checkpoint_and_dispatch -from transformers import CLIPTextModelWithProjection, CLIPTokenizer - -from diffusers import UnCLIPPipeline, UNet2DConditionModel, UNet2DModel -from diffusers.models.transformers.prior_transformer import PriorTransformer -from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel -from diffusers.schedulers.scheduling_unclip import UnCLIPScheduler - - -r""" -Example - From the diffusers root directory: - -Download weights: -```sh -$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/efdf6206d8ed593961593dc029a8affa/decoder-ckpt-step%3D01000000-of-01000000.ckpt -$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/4226b831ae0279020d134281f3c31590/improved-sr-ckpt-step%3D1.2M.ckpt -$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/85626483eaca9f581e2a78d31ff905ca/prior-ckpt-step%3D01000000-of-01000000.ckpt -$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/0b62380a75e56f073e2844ab5199153d/ViT-L-14_stats.th -``` - -Convert the model: -```sh -$ python scripts/convert_kakao_brain_unclip_to_diffusers.py \ - --decoder_checkpoint_path ./decoder-ckpt-step\=01000000-of-01000000.ckpt \ - --super_res_unet_checkpoint_path ./improved-sr-ckpt-step\=1.2M.ckpt \ - --prior_checkpoint_path ./prior-ckpt-step\=01000000-of-01000000.ckpt \ - --clip_stat_path ./ViT-L-14_stats.th \ - --dump_path -``` -""" - - -# prior - -PRIOR_ORIGINAL_PREFIX = "model" - -# Uses default arguments -PRIOR_CONFIG = {} - - -def prior_model_from_original_config(): - model = PriorTransformer(**PRIOR_CONFIG) - - return model - - -def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint, clip_stats_checkpoint): - diffusers_checkpoint = {} - - # .time_embed.0 -> .time_embedding.linear_1 - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.weight"], - "time_embedding.linear_1.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.bias"], - } - ) - - # .clip_img_proj -> .proj_in - diffusers_checkpoint.update( - { - "proj_in.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.weight"], - "proj_in.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.bias"], - } - ) - - # .text_emb_proj -> .embedding_proj - diffusers_checkpoint.update( - { - "embedding_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.weight"], - "embedding_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.bias"], - } - ) - - # .text_enc_proj -> .encoder_hidden_states_proj - diffusers_checkpoint.update( - { - "encoder_hidden_states_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.weight"], - "encoder_hidden_states_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.bias"], - } - ) - - # .positional_embedding -> .positional_embedding - diffusers_checkpoint.update({"positional_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.positional_embedding"]}) - - # .prd_emb -> .prd_embedding - diffusers_checkpoint.update({"prd_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.prd_emb"]}) - - # .time_embed.2 -> .time_embedding.linear_2 - diffusers_checkpoint.update( - { - "time_embedding.linear_2.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.weight"], - "time_embedding.linear_2.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.bias"], - } - ) - - # .resblocks. -> .transformer_blocks. - for idx in range(len(model.transformer_blocks)): - diffusers_transformer_prefix = f"transformer_blocks.{idx}" - original_transformer_prefix = f"{PRIOR_ORIGINAL_PREFIX}.transformer.resblocks.{idx}" - - # .attn -> .attn1 - diffusers_attention_prefix = f"{diffusers_transformer_prefix}.attn1" - original_attention_prefix = f"{original_transformer_prefix}.attn" - diffusers_checkpoint.update( - prior_attention_to_diffusers( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - original_attention_prefix=original_attention_prefix, - attention_head_dim=model.attention_head_dim, - ) - ) - - # .mlp -> .ff - diffusers_ff_prefix = f"{diffusers_transformer_prefix}.ff" - original_ff_prefix = f"{original_transformer_prefix}.mlp" - diffusers_checkpoint.update( - prior_ff_to_diffusers( - checkpoint, diffusers_ff_prefix=diffusers_ff_prefix, original_ff_prefix=original_ff_prefix - ) - ) - - # .ln_1 -> .norm1 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm1.weight": checkpoint[ - f"{original_transformer_prefix}.ln_1.weight" - ], - f"{diffusers_transformer_prefix}.norm1.bias": checkpoint[f"{original_transformer_prefix}.ln_1.bias"], - } - ) - - # .ln_2 -> .norm3 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm3.weight": checkpoint[ - f"{original_transformer_prefix}.ln_2.weight" - ], - f"{diffusers_transformer_prefix}.norm3.bias": checkpoint[f"{original_transformer_prefix}.ln_2.bias"], - } - ) - - # .final_ln -> .norm_out - diffusers_checkpoint.update( - { - "norm_out.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.weight"], - "norm_out.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.bias"], - } - ) - - # .out_proj -> .proj_to_clip_embeddings - diffusers_checkpoint.update( - { - "proj_to_clip_embeddings.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.weight"], - "proj_to_clip_embeddings.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.bias"], - } - ) - - # clip stats - clip_mean, clip_std = clip_stats_checkpoint - clip_mean = clip_mean[None, :] - clip_std = clip_std[None, :] - - diffusers_checkpoint.update({"clip_mean": clip_mean, "clip_std": clip_std}) - - return diffusers_checkpoint - - -def prior_attention_to_diffusers( - checkpoint, *, diffusers_attention_prefix, original_attention_prefix, attention_head_dim -): - diffusers_checkpoint = {} - - # .c_qkv -> .{to_q, to_k, to_v} - [q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions( - weight=checkpoint[f"{original_attention_prefix}.c_qkv.weight"], - bias=checkpoint[f"{original_attention_prefix}.c_qkv.bias"], - split=3, - chunk_size=attention_head_dim, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_q.weight": q_weight, - f"{diffusers_attention_prefix}.to_q.bias": q_bias, - f"{diffusers_attention_prefix}.to_k.weight": k_weight, - f"{diffusers_attention_prefix}.to_k.bias": k_bias, - f"{diffusers_attention_prefix}.to_v.weight": v_weight, - f"{diffusers_attention_prefix}.to_v.bias": v_bias, - } - ) - - # .c_proj -> .to_out.0 - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{original_attention_prefix}.c_proj.weight"], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{original_attention_prefix}.c_proj.bias"], - } - ) - - return diffusers_checkpoint - - -def prior_ff_to_diffusers(checkpoint, *, diffusers_ff_prefix, original_ff_prefix): - diffusers_checkpoint = { - # .c_fc -> .net.0.proj - f"{diffusers_ff_prefix}.net.{0}.proj.weight": checkpoint[f"{original_ff_prefix}.c_fc.weight"], - f"{diffusers_ff_prefix}.net.{0}.proj.bias": checkpoint[f"{original_ff_prefix}.c_fc.bias"], - # .c_proj -> .net.2 - f"{diffusers_ff_prefix}.net.{2}.weight": checkpoint[f"{original_ff_prefix}.c_proj.weight"], - f"{diffusers_ff_prefix}.net.{2}.bias": checkpoint[f"{original_ff_prefix}.c_proj.bias"], - } - - return diffusers_checkpoint - - -# done prior - - -# decoder - -DECODER_ORIGINAL_PREFIX = "model" - -# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can -# update then. -DECODER_CONFIG = { - "sample_size": 64, - "layers_per_block": 3, - "down_block_types": ( - "ResnetDownsampleBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - ), - "up_block_types": ( - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "ResnetUpsampleBlock2D", - ), - "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", - "block_out_channels": (320, 640, 960, 1280), - "in_channels": 3, - "out_channels": 6, - "cross_attention_dim": 1536, - "class_embed_type": "identity", - "attention_head_dim": 64, - "resnet_time_scale_shift": "scale_shift", -} - - -def decoder_model_from_original_config(): - model = UNet2DConditionModel(**DECODER_CONFIG) - - return model - - -def decoder_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - original_unet_prefix = DECODER_ORIGINAL_PREFIX - num_head_channels = DECODER_CONFIG["attention_head_dim"] - - diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix)) - - # .input_blocks -> .down_blocks - - original_down_block_idx = 1 - - for diffusers_down_block_idx in range(len(model.down_blocks)): - checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_down_block_idx=diffusers_down_block_idx, - original_down_block_idx=original_down_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=num_head_channels, - ) - - original_down_block_idx += num_original_down_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .input_blocks -> .down_blocks - - diffusers_checkpoint.update( - unet_midblock_to_diffusers_checkpoint( - model, - checkpoint, - original_unet_prefix=original_unet_prefix, - num_head_channels=num_head_channels, - ) - ) - - # .output_blocks -> .up_blocks - - original_up_block_idx = 0 - - for diffusers_up_block_idx in range(len(model.up_blocks)): - checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_up_block_idx=diffusers_up_block_idx, - original_up_block_idx=original_up_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=num_head_channels, - ) - - original_up_block_idx += num_original_up_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .output_blocks -> .up_blocks - - diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix)) - - return diffusers_checkpoint - - -# done decoder - -# text proj - - -def text_proj_from_original_config(): - # From the conditional unet constructor where the dimension of the projected time embeddings is - # constructed - time_embed_dim = DECODER_CONFIG["block_out_channels"][0] * 4 - - cross_attention_dim = DECODER_CONFIG["cross_attention_dim"] - - model = UnCLIPTextProjModel(time_embed_dim=time_embed_dim, cross_attention_dim=cross_attention_dim) - - return model - - -# Note that the input checkpoint is the original decoder checkpoint -def text_proj_original_checkpoint_to_diffusers_checkpoint(checkpoint): - diffusers_checkpoint = { - # .text_seq_proj.0 -> .encoder_hidden_states_proj - "encoder_hidden_states_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.weight"], - "encoder_hidden_states_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.0.bias"], - # .text_seq_proj.1 -> .text_encoder_hidden_states_norm - "text_encoder_hidden_states_norm.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.weight"], - "text_encoder_hidden_states_norm.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_seq_proj.1.bias"], - # .clip_tok_proj -> .clip_extra_context_tokens_proj - "clip_extra_context_tokens_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.weight"], - "clip_extra_context_tokens_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.clip_tok_proj.bias"], - # .text_feat_proj -> .embedding_proj - "embedding_proj.weight": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.weight"], - "embedding_proj.bias": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.text_feat_proj.bias"], - # .cf_param -> .learned_classifier_free_guidance_embeddings - "learned_classifier_free_guidance_embeddings": checkpoint[f"{DECODER_ORIGINAL_PREFIX}.cf_param"], - # .clip_emb -> .clip_image_embeddings_project_to_time_embeddings - "clip_image_embeddings_project_to_time_embeddings.weight": checkpoint[ - f"{DECODER_ORIGINAL_PREFIX}.clip_emb.weight" - ], - "clip_image_embeddings_project_to_time_embeddings.bias": checkpoint[ - f"{DECODER_ORIGINAL_PREFIX}.clip_emb.bias" - ], - } - - return diffusers_checkpoint - - -# done text proj - -# super res unet first steps - -SUPER_RES_UNET_FIRST_STEPS_PREFIX = "model_first_steps" - -SUPER_RES_UNET_FIRST_STEPS_CONFIG = { - "sample_size": 256, - "layers_per_block": 3, - "down_block_types": ( - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - ), - "up_block_types": ( - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - ), - "block_out_channels": (320, 640, 960, 1280), - "in_channels": 6, - "out_channels": 3, - "add_attention": False, -} - - -def super_res_unet_first_steps_model_from_original_config(): - model = UNet2DModel(**SUPER_RES_UNET_FIRST_STEPS_CONFIG) - - return model - - -def super_res_unet_first_steps_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - original_unet_prefix = SUPER_RES_UNET_FIRST_STEPS_PREFIX - - diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix)) - - # .input_blocks -> .down_blocks - - original_down_block_idx = 1 - - for diffusers_down_block_idx in range(len(model.down_blocks)): - checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_down_block_idx=diffusers_down_block_idx, - original_down_block_idx=original_down_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - - original_down_block_idx += num_original_down_blocks - - diffusers_checkpoint.update(checkpoint_update) - - diffusers_checkpoint.update( - unet_midblock_to_diffusers_checkpoint( - model, - checkpoint, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - ) - - # .output_blocks -> .up_blocks - - original_up_block_idx = 0 - - for diffusers_up_block_idx in range(len(model.up_blocks)): - checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_up_block_idx=diffusers_up_block_idx, - original_up_block_idx=original_up_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - - original_up_block_idx += num_original_up_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .output_blocks -> .up_blocks - - diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix)) - - return diffusers_checkpoint - - -# done super res unet first steps - -# super res unet last step - -SUPER_RES_UNET_LAST_STEP_PREFIX = "model_last_step" - -SUPER_RES_UNET_LAST_STEP_CONFIG = { - "sample_size": 256, - "layers_per_block": 3, - "down_block_types": ( - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - "ResnetDownsampleBlock2D", - ), - "up_block_types": ( - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - "ResnetUpsampleBlock2D", - ), - "block_out_channels": (320, 640, 960, 1280), - "in_channels": 6, - "out_channels": 3, - "add_attention": False, -} - - -def super_res_unet_last_step_model_from_original_config(): - model = UNet2DModel(**SUPER_RES_UNET_LAST_STEP_CONFIG) - - return model - - -def super_res_unet_last_step_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - original_unet_prefix = SUPER_RES_UNET_LAST_STEP_PREFIX - - diffusers_checkpoint.update(unet_time_embeddings(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_in(checkpoint, original_unet_prefix)) - - # .input_blocks -> .down_blocks - - original_down_block_idx = 1 - - for diffusers_down_block_idx in range(len(model.down_blocks)): - checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_down_block_idx=diffusers_down_block_idx, - original_down_block_idx=original_down_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - - original_down_block_idx += num_original_down_blocks - - diffusers_checkpoint.update(checkpoint_update) - - diffusers_checkpoint.update( - unet_midblock_to_diffusers_checkpoint( - model, - checkpoint, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - ) - - # .output_blocks -> .up_blocks - - original_up_block_idx = 0 - - for diffusers_up_block_idx in range(len(model.up_blocks)): - checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_up_block_idx=diffusers_up_block_idx, - original_up_block_idx=original_up_block_idx, - original_unet_prefix=original_unet_prefix, - num_head_channels=None, - ) - - original_up_block_idx += num_original_up_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .output_blocks -> .up_blocks - - diffusers_checkpoint.update(unet_conv_norm_out(checkpoint, original_unet_prefix)) - diffusers_checkpoint.update(unet_conv_out(checkpoint, original_unet_prefix)) - - return diffusers_checkpoint - - -# done super res unet last step - - -# unet utils - - -# .time_embed -> .time_embedding -def unet_time_embeddings(checkpoint, original_unet_prefix): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint[f"{original_unet_prefix}.time_embed.0.weight"], - "time_embedding.linear_1.bias": checkpoint[f"{original_unet_prefix}.time_embed.0.bias"], - "time_embedding.linear_2.weight": checkpoint[f"{original_unet_prefix}.time_embed.2.weight"], - "time_embedding.linear_2.bias": checkpoint[f"{original_unet_prefix}.time_embed.2.bias"], - } - ) - - return diffusers_checkpoint - - -# .input_blocks.0 -> .conv_in -def unet_conv_in(checkpoint, original_unet_prefix): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_in.weight": checkpoint[f"{original_unet_prefix}.input_blocks.0.0.weight"], - "conv_in.bias": checkpoint[f"{original_unet_prefix}.input_blocks.0.0.bias"], - } - ) - - return diffusers_checkpoint - - -# .out.0 -> .conv_norm_out -def unet_conv_norm_out(checkpoint, original_unet_prefix): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_norm_out.weight": checkpoint[f"{original_unet_prefix}.out.0.weight"], - "conv_norm_out.bias": checkpoint[f"{original_unet_prefix}.out.0.bias"], - } - ) - - return diffusers_checkpoint - - -# .out.2 -> .conv_out -def unet_conv_out(checkpoint, original_unet_prefix): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_out.weight": checkpoint[f"{original_unet_prefix}.out.2.weight"], - "conv_out.bias": checkpoint[f"{original_unet_prefix}.out.2.bias"], - } - ) - - return diffusers_checkpoint - - -# .input_blocks -> .down_blocks -def unet_downblock_to_diffusers_checkpoint( - model, checkpoint, *, diffusers_down_block_idx, original_down_block_idx, original_unet_prefix, num_head_channels -): - diffusers_checkpoint = {} - - diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.resnets" - original_down_block_prefix = f"{original_unet_prefix}.input_blocks" - - down_block = model.down_blocks[diffusers_down_block_idx] - - num_resnets = len(down_block.resnets) - - if down_block.downsamplers is None: - downsampler = False - else: - assert len(down_block.downsamplers) == 1 - downsampler = True - # The downsample block is also a resnet - num_resnets += 1 - - for resnet_idx_inc in range(num_resnets): - full_resnet_prefix = f"{original_down_block_prefix}.{original_down_block_idx + resnet_idx_inc}.0" - - if downsampler and resnet_idx_inc == num_resnets - 1: - # this is a downsample block - full_diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.downsamplers.0" - else: - # this is a regular resnet block - full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}" - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix - ) - ) - - if hasattr(down_block, "attentions"): - num_attentions = len(down_block.attentions) - diffusers_attention_prefix = f"down_blocks.{diffusers_down_block_idx}.attentions" - - for attention_idx_inc in range(num_attentions): - full_attention_prefix = f"{original_down_block_prefix}.{original_down_block_idx + attention_idx_inc}.1" - full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}" - - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - attention_prefix=full_attention_prefix, - diffusers_attention_prefix=full_diffusers_attention_prefix, - num_head_channels=num_head_channels, - ) - ) - - num_original_down_blocks = num_resnets - - return diffusers_checkpoint, num_original_down_blocks - - -# .middle_block -> .mid_block -def unet_midblock_to_diffusers_checkpoint(model, checkpoint, *, original_unet_prefix, num_head_channels): - diffusers_checkpoint = {} - - # block 0 - - original_block_idx = 0 - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, - diffusers_resnet_prefix="mid_block.resnets.0", - resnet_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}", - ) - ) - - original_block_idx += 1 - - # optional block 1 - - if hasattr(model.mid_block, "attentions") and model.mid_block.attentions[0] is not None: - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix="mid_block.attentions.0", - attention_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}", - num_head_channels=num_head_channels, - ) - ) - original_block_idx += 1 - - # block 1 or block 2 - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, - diffusers_resnet_prefix="mid_block.resnets.1", - resnet_prefix=f"{original_unet_prefix}.middle_block.{original_block_idx}", - ) - ) - - return diffusers_checkpoint - - -# .output_blocks -> .up_blocks -def unet_upblock_to_diffusers_checkpoint( - model, checkpoint, *, diffusers_up_block_idx, original_up_block_idx, original_unet_prefix, num_head_channels -): - diffusers_checkpoint = {} - - diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.resnets" - original_up_block_prefix = f"{original_unet_prefix}.output_blocks" - - up_block = model.up_blocks[diffusers_up_block_idx] - - num_resnets = len(up_block.resnets) - - if up_block.upsamplers is None: - upsampler = False - else: - assert len(up_block.upsamplers) == 1 - upsampler = True - # The upsample block is also a resnet - num_resnets += 1 - - has_attentions = hasattr(up_block, "attentions") - - for resnet_idx_inc in range(num_resnets): - if upsampler and resnet_idx_inc == num_resnets - 1: - # this is an upsample block - if has_attentions: - # There is a middle attention block that we skip - original_resnet_block_idx = 2 - else: - original_resnet_block_idx = 1 - - # we add the `minus 1` because the last two resnets are stuck together in the same output block - full_resnet_prefix = ( - f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc - 1}.{original_resnet_block_idx}" - ) - - full_diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.upsamplers.0" - else: - # this is a regular resnet block - full_resnet_prefix = f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc}.0" - full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}" - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix - ) - ) - - if has_attentions: - num_attentions = len(up_block.attentions) - diffusers_attention_prefix = f"up_blocks.{diffusers_up_block_idx}.attentions" - - for attention_idx_inc in range(num_attentions): - full_attention_prefix = f"{original_up_block_prefix}.{original_up_block_idx + attention_idx_inc}.1" - full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}" - - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - attention_prefix=full_attention_prefix, - diffusers_attention_prefix=full_diffusers_attention_prefix, - num_head_channels=num_head_channels, - ) - ) - - num_original_down_blocks = num_resnets - 1 if upsampler else num_resnets - - return diffusers_checkpoint, num_original_down_blocks - - -def resnet_to_diffusers_checkpoint(checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - diffusers_checkpoint = { - f"{diffusers_resnet_prefix}.norm1.weight": checkpoint[f"{resnet_prefix}.in_layers.0.weight"], - f"{diffusers_resnet_prefix}.norm1.bias": checkpoint[f"{resnet_prefix}.in_layers.0.bias"], - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.in_layers.2.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.in_layers.2.bias"], - f"{diffusers_resnet_prefix}.time_emb_proj.weight": checkpoint[f"{resnet_prefix}.emb_layers.1.weight"], - f"{diffusers_resnet_prefix}.time_emb_proj.bias": checkpoint[f"{resnet_prefix}.emb_layers.1.bias"], - f"{diffusers_resnet_prefix}.norm2.weight": checkpoint[f"{resnet_prefix}.out_layers.0.weight"], - f"{diffusers_resnet_prefix}.norm2.bias": checkpoint[f"{resnet_prefix}.out_layers.0.bias"], - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.out_layers.3.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.out_layers.3.bias"], - } - - skip_connection_prefix = f"{resnet_prefix}.skip_connection" - - if f"{skip_connection_prefix}.weight" in checkpoint: - diffusers_checkpoint.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{skip_connection_prefix}.weight"], - f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{skip_connection_prefix}.bias"], - } - ) - - return diffusers_checkpoint - - -def attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix, num_head_channels): - diffusers_checkpoint = {} - - # .norm -> .group_norm - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.group_norm.weight": checkpoint[f"{attention_prefix}.norm.weight"], - f"{diffusers_attention_prefix}.group_norm.bias": checkpoint[f"{attention_prefix}.norm.bias"], - } - ) - - # .qkv -> .{query, key, value} - [q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions( - weight=checkpoint[f"{attention_prefix}.qkv.weight"][:, :, 0], - bias=checkpoint[f"{attention_prefix}.qkv.bias"], - split=3, - chunk_size=num_head_channels, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_q.weight": q_weight, - f"{diffusers_attention_prefix}.to_q.bias": q_bias, - f"{diffusers_attention_prefix}.to_k.weight": k_weight, - f"{diffusers_attention_prefix}.to_k.bias": k_bias, - f"{diffusers_attention_prefix}.to_v.weight": v_weight, - f"{diffusers_attention_prefix}.to_v.bias": v_bias, - } - ) - - # .encoder_kv -> .{context_key, context_value} - [encoder_k_weight, encoder_v_weight], [encoder_k_bias, encoder_v_bias] = split_attentions( - weight=checkpoint[f"{attention_prefix}.encoder_kv.weight"][:, :, 0], - bias=checkpoint[f"{attention_prefix}.encoder_kv.bias"], - split=2, - chunk_size=num_head_channels, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.add_k_proj.weight": encoder_k_weight, - f"{diffusers_attention_prefix}.add_k_proj.bias": encoder_k_bias, - f"{diffusers_attention_prefix}.add_v_proj.weight": encoder_v_weight, - f"{diffusers_attention_prefix}.add_v_proj.bias": encoder_v_bias, - } - ) - - # .proj_out (1d conv) -> .proj_attn (linear) - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][ - :, :, 0 - ], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj_out.bias"], - } - ) - - return diffusers_checkpoint - - -# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) -def split_attentions(*, weight, bias, split, chunk_size): - weights = [None] * split - biases = [None] * split - - weights_biases_idx = 0 - - for starting_row_index in range(0, weight.shape[0], chunk_size): - row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size) - - weight_rows = weight[row_indices, :] - bias_rows = bias[row_indices] - - if weights[weights_biases_idx] is None: - assert weights[weights_biases_idx] is None - weights[weights_biases_idx] = weight_rows - biases[weights_biases_idx] = bias_rows - else: - assert weights[weights_biases_idx] is not None - weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows]) - biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows]) - - weights_biases_idx = (weights_biases_idx + 1) % split - - return weights, biases - - -# done unet utils - - -# Driver functions - - -def text_encoder(): - print("loading CLIP text encoder") - - clip_name = "openai/clip-vit-large-patch14" - - # sets pad_value to 0 - pad_token = "!" - - tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto") - - assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0 - - text_encoder_model = CLIPTextModelWithProjection.from_pretrained( - clip_name, - # `CLIPTextModel` does not support device_map="auto" - # device_map="auto" - ) - - print("done loading CLIP text encoder") - - return text_encoder_model, tokenizer_model - - -def prior(*, args, checkpoint_map_location): - print("loading prior") - - prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) - prior_checkpoint = prior_checkpoint["state_dict"] - - clip_stats_checkpoint = torch.load(args.clip_stat_path, map_location=checkpoint_map_location) - - prior_model = prior_model_from_original_config() - - prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint( - prior_model, prior_checkpoint, clip_stats_checkpoint - ) - - del prior_checkpoint - del clip_stats_checkpoint - - load_checkpoint_to_model(prior_diffusers_checkpoint, prior_model, strict=True) - - print("done loading prior") - - return prior_model - - -def decoder(*, args, checkpoint_map_location): - print("loading decoder") - - decoder_checkpoint = torch.load(args.decoder_checkpoint_path, map_location=checkpoint_map_location) - decoder_checkpoint = decoder_checkpoint["state_dict"] - - decoder_model = decoder_model_from_original_config() - - decoder_diffusers_checkpoint = decoder_original_checkpoint_to_diffusers_checkpoint( - decoder_model, decoder_checkpoint - ) - - # text proj interlude - - # The original decoder implementation includes a set of parameters that are used - # for creating the `encoder_hidden_states` which are what the U-net is conditioned - # on. The diffusers conditional unet directly takes the encoder_hidden_states. We pull - # the parameters into the UnCLIPTextProjModel class - text_proj_model = text_proj_from_original_config() - - text_proj_checkpoint = text_proj_original_checkpoint_to_diffusers_checkpoint(decoder_checkpoint) - - load_checkpoint_to_model(text_proj_checkpoint, text_proj_model, strict=True) - - # done text proj interlude - - del decoder_checkpoint - - load_checkpoint_to_model(decoder_diffusers_checkpoint, decoder_model, strict=True) - - print("done loading decoder") - - return decoder_model, text_proj_model - - -def super_res_unet(*, args, checkpoint_map_location): - print("loading super resolution unet") - - super_res_checkpoint = torch.load(args.super_res_unet_checkpoint_path, map_location=checkpoint_map_location) - super_res_checkpoint = super_res_checkpoint["state_dict"] - - # model_first_steps - - super_res_first_model = super_res_unet_first_steps_model_from_original_config() - - super_res_first_steps_checkpoint = super_res_unet_first_steps_original_checkpoint_to_diffusers_checkpoint( - super_res_first_model, super_res_checkpoint - ) - - # model_last_step - super_res_last_model = super_res_unet_last_step_model_from_original_config() - - super_res_last_step_checkpoint = super_res_unet_last_step_original_checkpoint_to_diffusers_checkpoint( - super_res_last_model, super_res_checkpoint - ) - - del super_res_checkpoint - - load_checkpoint_to_model(super_res_first_steps_checkpoint, super_res_first_model, strict=True) - - load_checkpoint_to_model(super_res_last_step_checkpoint, super_res_last_model, strict=True) - - print("done loading super resolution unet") - - return super_res_first_model, super_res_last_model - - -def load_checkpoint_to_model(checkpoint, model, strict=False): - with tempfile.NamedTemporaryFile() as file: - torch.save(checkpoint, file.name) - del checkpoint - if strict: - model.load_state_dict(torch.load(file.name), strict=True) - else: - load_checkpoint_and_dispatch(model, file.name, device_map="auto") - - -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( - "--prior_checkpoint_path", - default=None, - type=str, - required=True, - help="Path to the prior checkpoint to convert.", - ) - - parser.add_argument( - "--decoder_checkpoint_path", - default=None, - type=str, - required=True, - help="Path to the decoder checkpoint to convert.", - ) - - parser.add_argument( - "--super_res_unet_checkpoint_path", - default=None, - type=str, - required=True, - help="Path to the super resolution checkpoint to convert.", - ) - - parser.add_argument( - "--clip_stat_path", default=None, type=str, required=True, help="Path to the clip stats checkpoint to convert." - ) - - parser.add_argument( - "--checkpoint_load_device", - default="cpu", - type=str, - required=False, - help="The device passed to `map_location` when loading checkpoints.", - ) - - parser.add_argument( - "--debug", - default=None, - type=str, - required=False, - help="Only run a specific stage of the convert script. Used for debugging", - ) - - args = parser.parse_args() - - print(f"loading checkpoints to {args.checkpoint_load_device}") - - checkpoint_map_location = torch.device(args.checkpoint_load_device) - - if args.debug is not None: - print(f"debug: only executing {args.debug}") - - if args.debug is None: - text_encoder_model, tokenizer_model = text_encoder() - - prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) - - decoder_model, text_proj_model = decoder(args=args, checkpoint_map_location=checkpoint_map_location) - - super_res_first_model, super_res_last_model = super_res_unet( - args=args, checkpoint_map_location=checkpoint_map_location - ) - - prior_scheduler = UnCLIPScheduler( - variance_type="fixed_small_log", - prediction_type="sample", - num_train_timesteps=1000, - clip_sample_range=5.0, - ) - - decoder_scheduler = UnCLIPScheduler( - variance_type="learned_range", - prediction_type="epsilon", - num_train_timesteps=1000, - ) - - super_res_scheduler = UnCLIPScheduler( - variance_type="fixed_small_log", - prediction_type="epsilon", - num_train_timesteps=1000, - ) - - print(f"saving Kakao Brain unCLIP to {args.dump_path}") - - pipe = UnCLIPPipeline( - prior=prior_model, - decoder=decoder_model, - text_proj=text_proj_model, - tokenizer=tokenizer_model, - text_encoder=text_encoder_model, - super_res_first=super_res_first_model, - super_res_last=super_res_last_model, - prior_scheduler=prior_scheduler, - decoder_scheduler=decoder_scheduler, - super_res_scheduler=super_res_scheduler, - ) - pipe.save_pretrained(args.dump_path) - - print("done writing Kakao Brain unCLIP") - elif args.debug == "text_encoder": - text_encoder_model, tokenizer_model = text_encoder() - elif args.debug == "prior": - prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) - elif args.debug == "decoder": - decoder_model, text_proj_model = decoder(args=args, checkpoint_map_location=checkpoint_map_location) - elif args.debug == "super_res_unet": - super_res_first_model, super_res_last_model = super_res_unet( - args=args, checkpoint_map_location=checkpoint_map_location - ) - else: - raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/convert_kandinsky3_unet.py b/scripts/convert_kandinsky3_unet.py deleted file mode 100644 index 4fe8c54eb7fc..000000000000 --- a/scripts/convert_kandinsky3_unet.py +++ /dev/null @@ -1,98 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import fnmatch - -from safetensors.torch import load_file - -from diffusers import Kandinsky3UNet - - -MAPPING = { - "to_time_embed.1": "time_embedding.linear_1", - "to_time_embed.3": "time_embedding.linear_2", - "in_layer": "conv_in", - "out_layer.0": "conv_norm_out", - "out_layer.2": "conv_out", - "down_samples": "down_blocks", - "up_samples": "up_blocks", - "projection_lin": "encoder_hid_proj.projection_linear", - "projection_ln": "encoder_hid_proj.projection_norm", - "feature_pooling": "add_time_condition", - "to_query": "to_q", - "to_key": "to_k", - "to_value": "to_v", - "output_layer": "to_out.0", - "self_attention_block": "attentions.0", -} - -DYNAMIC_MAP = { - "resnet_attn_blocks.*.0": "resnets_in.*", - "resnet_attn_blocks.*.1": ("attentions.*", 1), - "resnet_attn_blocks.*.2": "resnets_out.*", -} -# MAPPING = {} - - -def convert_state_dict(unet_state_dict): - """ - Convert the state dict of a U-Net model to match the key format expected by Kandinsky3UNet model. - Args: - unet_model (torch.nn.Module): The original U-Net model. - unet_kandi3_model (torch.nn.Module): The Kandinsky3UNet model to match keys with. - - Returns: - OrderedDict: The converted state dictionary. - """ - # Example of renaming logic (this will vary based on your model's architecture) - converted_state_dict = {} - for key in unet_state_dict: - new_key = key - for pattern, new_pattern in MAPPING.items(): - new_key = new_key.replace(pattern, new_pattern) - - for dyn_pattern, dyn_new_pattern in DYNAMIC_MAP.items(): - has_matched = False - if fnmatch.fnmatch(new_key, f"*.{dyn_pattern}.*") and not has_matched: - star = int(new_key.split(dyn_pattern.split(".")[0])[-1].split(".")[1]) - - if isinstance(dyn_new_pattern, tuple): - new_star = star + dyn_new_pattern[-1] - dyn_new_pattern = dyn_new_pattern[0] - else: - new_star = star - - pattern = dyn_pattern.replace("*", str(star)) - new_pattern = dyn_new_pattern.replace("*", str(new_star)) - - new_key = new_key.replace(pattern, new_pattern) - has_matched = True - - converted_state_dict[new_key] = unet_state_dict[key] - - return converted_state_dict - - -def main(model_path, output_path): - # Load your original U-Net model - unet_state_dict = load_file(model_path) - - # Initialize your Kandinsky3UNet model - config = {} - - # Convert the state dict - converted_state_dict = convert_state_dict(unet_state_dict) - - unet = Kandinsky3UNet(config) - unet.load_state_dict(converted_state_dict) - - unet.save_pretrained(output_path) - print(f"Converted model saved to {output_path}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Convert U-Net PyTorch model to Kandinsky3UNet format") - parser.add_argument("--model_path", type=str, required=True, help="Path to the original U-Net PyTorch model") - parser.add_argument("--output_path", type=str, required=True, help="Path to save the converted model") - - args = parser.parse_args() - main(args.model_path, args.output_path) diff --git a/scripts/convert_kandinsky_to_diffusers.py b/scripts/convert_kandinsky_to_diffusers.py deleted file mode 100644 index 8d3f7b63d0c1..000000000000 --- a/scripts/convert_kandinsky_to_diffusers.py +++ /dev/null @@ -1,1411 +0,0 @@ -import argparse -import os -import tempfile - -import torch -from accelerate import load_checkpoint_and_dispatch - -from diffusers import UNet2DConditionModel -from diffusers.models.transformers.prior_transformer import PriorTransformer -from diffusers.models.vq_model import VQModel - - -""" -Example - From the diffusers root directory: - -Download weights: -```sh -$ wget https://huggingface.co/ai-forever/Kandinsky_2.1/blob/main/prior_fp16.ckpt -``` - -Convert the model: -```sh -python scripts/convert_kandinsky_to_diffusers.py \ - --prior_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/prior_fp16.ckpt \ - --clip_stat_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/ViT-L-14_stats.th \ - --text2img_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/decoder_fp16.ckpt \ - --inpaint_text2img_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/inpainting_fp16.ckpt \ - --movq_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/movq_final.ckpt \ - --dump_path /home/yiyi_huggingface_co/dump \ - --debug decoder -``` -""" - - -# prior - -PRIOR_ORIGINAL_PREFIX = "model" - -# Uses default arguments -PRIOR_CONFIG = {} - - -def prior_model_from_original_config(): - model = PriorTransformer(**PRIOR_CONFIG) - - return model - - -def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint, clip_stats_checkpoint): - diffusers_checkpoint = {} - - # .time_embed.0 -> .time_embedding.linear_1 - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.weight"], - "time_embedding.linear_1.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.0.bias"], - } - ) - - # .clip_img_proj -> .proj_in - diffusers_checkpoint.update( - { - "proj_in.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.weight"], - "proj_in.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_img_proj.bias"], - } - ) - - # .text_emb_proj -> .embedding_proj - diffusers_checkpoint.update( - { - "embedding_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.weight"], - "embedding_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_emb_proj.bias"], - } - ) - - # .text_enc_proj -> .encoder_hidden_states_proj - diffusers_checkpoint.update( - { - "encoder_hidden_states_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.weight"], - "encoder_hidden_states_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.text_enc_proj.bias"], - } - ) - - # .positional_embedding -> .positional_embedding - diffusers_checkpoint.update({"positional_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.positional_embedding"]}) - - # .prd_emb -> .prd_embedding - diffusers_checkpoint.update({"prd_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.prd_emb"]}) - - # .time_embed.2 -> .time_embedding.linear_2 - diffusers_checkpoint.update( - { - "time_embedding.linear_2.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.weight"], - "time_embedding.linear_2.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.2.bias"], - } - ) - - # .resblocks. -> .transformer_blocks. - for idx in range(len(model.transformer_blocks)): - diffusers_transformer_prefix = f"transformer_blocks.{idx}" - original_transformer_prefix = f"{PRIOR_ORIGINAL_PREFIX}.transformer.resblocks.{idx}" - - # .attn -> .attn1 - diffusers_attention_prefix = f"{diffusers_transformer_prefix}.attn1" - original_attention_prefix = f"{original_transformer_prefix}.attn" - diffusers_checkpoint.update( - prior_attention_to_diffusers( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - original_attention_prefix=original_attention_prefix, - attention_head_dim=model.attention_head_dim, - ) - ) - - # .mlp -> .ff - diffusers_ff_prefix = f"{diffusers_transformer_prefix}.ff" - original_ff_prefix = f"{original_transformer_prefix}.mlp" - diffusers_checkpoint.update( - prior_ff_to_diffusers( - checkpoint, diffusers_ff_prefix=diffusers_ff_prefix, original_ff_prefix=original_ff_prefix - ) - ) - - # .ln_1 -> .norm1 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm1.weight": checkpoint[ - f"{original_transformer_prefix}.ln_1.weight" - ], - f"{diffusers_transformer_prefix}.norm1.bias": checkpoint[f"{original_transformer_prefix}.ln_1.bias"], - } - ) - - # .ln_2 -> .norm3 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm3.weight": checkpoint[ - f"{original_transformer_prefix}.ln_2.weight" - ], - f"{diffusers_transformer_prefix}.norm3.bias": checkpoint[f"{original_transformer_prefix}.ln_2.bias"], - } - ) - - # .final_ln -> .norm_out - diffusers_checkpoint.update( - { - "norm_out.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.weight"], - "norm_out.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.final_ln.bias"], - } - ) - - # .out_proj -> .proj_to_clip_embeddings - diffusers_checkpoint.update( - { - "proj_to_clip_embeddings.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.weight"], - "proj_to_clip_embeddings.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.out_proj.bias"], - } - ) - - # clip stats - clip_mean, clip_std = clip_stats_checkpoint - clip_mean = clip_mean[None, :] - clip_std = clip_std[None, :] - - diffusers_checkpoint.update({"clip_mean": clip_mean, "clip_std": clip_std}) - - return diffusers_checkpoint - - -def prior_attention_to_diffusers( - checkpoint, *, diffusers_attention_prefix, original_attention_prefix, attention_head_dim -): - diffusers_checkpoint = {} - - # .c_qkv -> .{to_q, to_k, to_v} - [q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions( - weight=checkpoint[f"{original_attention_prefix}.c_qkv.weight"], - bias=checkpoint[f"{original_attention_prefix}.c_qkv.bias"], - split=3, - chunk_size=attention_head_dim, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_q.weight": q_weight, - f"{diffusers_attention_prefix}.to_q.bias": q_bias, - f"{diffusers_attention_prefix}.to_k.weight": k_weight, - f"{diffusers_attention_prefix}.to_k.bias": k_bias, - f"{diffusers_attention_prefix}.to_v.weight": v_weight, - f"{diffusers_attention_prefix}.to_v.bias": v_bias, - } - ) - - # .c_proj -> .to_out.0 - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{original_attention_prefix}.c_proj.weight"], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{original_attention_prefix}.c_proj.bias"], - } - ) - - return diffusers_checkpoint - - -def prior_ff_to_diffusers(checkpoint, *, diffusers_ff_prefix, original_ff_prefix): - diffusers_checkpoint = { - # .c_fc -> .net.0.proj - f"{diffusers_ff_prefix}.net.{0}.proj.weight": checkpoint[f"{original_ff_prefix}.c_fc.weight"], - f"{diffusers_ff_prefix}.net.{0}.proj.bias": checkpoint[f"{original_ff_prefix}.c_fc.bias"], - # .c_proj -> .net.2 - f"{diffusers_ff_prefix}.net.{2}.weight": checkpoint[f"{original_ff_prefix}.c_proj.weight"], - f"{diffusers_ff_prefix}.net.{2}.bias": checkpoint[f"{original_ff_prefix}.c_proj.bias"], - } - - return diffusers_checkpoint - - -# done prior - -# unet - -# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can -# update then. - -UNET_CONFIG = { - "act_fn": "silu", - "addition_embed_type": "text_image", - "addition_embed_type_num_heads": 64, - "attention_head_dim": 64, - "block_out_channels": [384, 768, 1152, 1536], - "center_input_sample": False, - "class_embed_type": None, - "class_embeddings_concat": False, - "conv_in_kernel": 3, - "conv_out_kernel": 3, - "cross_attention_dim": 768, - "cross_attention_norm": None, - "down_block_types": [ - "ResnetDownsampleBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - ], - "downsample_padding": 1, - "dual_cross_attention": False, - "encoder_hid_dim": 1024, - "encoder_hid_dim_type": "text_image_proj", - "flip_sin_to_cos": True, - "freq_shift": 0, - "in_channels": 4, - "layers_per_block": 3, - "mid_block_only_cross_attention": None, - "mid_block_scale_factor": 1, - "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", - "norm_eps": 1e-05, - "norm_num_groups": 32, - "num_class_embeds": None, - "only_cross_attention": False, - "out_channels": 8, - "projection_class_embeddings_input_dim": None, - "resnet_out_scale_factor": 1.0, - "resnet_skip_time_act": False, - "resnet_time_scale_shift": "scale_shift", - "sample_size": 64, - "time_cond_proj_dim": None, - "time_embedding_act_fn": None, - "time_embedding_dim": None, - "time_embedding_type": "positional", - "timestep_post_act": None, - "up_block_types": [ - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "ResnetUpsampleBlock2D", - ], - "upcast_attention": False, - "use_linear_projection": False, -} - - -def unet_model_from_original_config(): - model = UNet2DConditionModel(**UNET_CONFIG) - - return model - - -def unet_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - num_head_channels = UNET_CONFIG["attention_head_dim"] - - diffusers_checkpoint.update(unet_time_embeddings(checkpoint)) - diffusers_checkpoint.update(unet_conv_in(checkpoint)) - diffusers_checkpoint.update(unet_add_embedding(checkpoint)) - diffusers_checkpoint.update(unet_encoder_hid_proj(checkpoint)) - - # .input_blocks -> .down_blocks - - original_down_block_idx = 1 - - for diffusers_down_block_idx in range(len(model.down_blocks)): - checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_down_block_idx=diffusers_down_block_idx, - original_down_block_idx=original_down_block_idx, - num_head_channels=num_head_channels, - ) - - original_down_block_idx += num_original_down_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .input_blocks -> .down_blocks - - diffusers_checkpoint.update( - unet_midblock_to_diffusers_checkpoint( - model, - checkpoint, - num_head_channels=num_head_channels, - ) - ) - - # .output_blocks -> .up_blocks - - original_up_block_idx = 0 - - for diffusers_up_block_idx in range(len(model.up_blocks)): - checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_up_block_idx=diffusers_up_block_idx, - original_up_block_idx=original_up_block_idx, - num_head_channels=num_head_channels, - ) - - original_up_block_idx += num_original_up_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .output_blocks -> .up_blocks - - diffusers_checkpoint.update(unet_conv_norm_out(checkpoint)) - diffusers_checkpoint.update(unet_conv_out(checkpoint)) - - return diffusers_checkpoint - - -# done unet - -# inpaint unet - -# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can -# update then. - -INPAINT_UNET_CONFIG = { - "act_fn": "silu", - "addition_embed_type": "text_image", - "addition_embed_type_num_heads": 64, - "attention_head_dim": 64, - "block_out_channels": [384, 768, 1152, 1536], - "center_input_sample": False, - "class_embed_type": None, - "class_embeddings_concat": None, - "conv_in_kernel": 3, - "conv_out_kernel": 3, - "cross_attention_dim": 768, - "cross_attention_norm": None, - "down_block_types": [ - "ResnetDownsampleBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - "SimpleCrossAttnDownBlock2D", - ], - "downsample_padding": 1, - "dual_cross_attention": False, - "encoder_hid_dim": 1024, - "encoder_hid_dim_type": "text_image_proj", - "flip_sin_to_cos": True, - "freq_shift": 0, - "in_channels": 9, - "layers_per_block": 3, - "mid_block_only_cross_attention": None, - "mid_block_scale_factor": 1, - "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", - "norm_eps": 1e-05, - "norm_num_groups": 32, - "num_class_embeds": None, - "only_cross_attention": False, - "out_channels": 8, - "projection_class_embeddings_input_dim": None, - "resnet_out_scale_factor": 1.0, - "resnet_skip_time_act": False, - "resnet_time_scale_shift": "scale_shift", - "sample_size": 64, - "time_cond_proj_dim": None, - "time_embedding_act_fn": None, - "time_embedding_dim": None, - "time_embedding_type": "positional", - "timestep_post_act": None, - "up_block_types": [ - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "SimpleCrossAttnUpBlock2D", - "ResnetUpsampleBlock2D", - ], - "upcast_attention": False, - "use_linear_projection": False, -} - - -def inpaint_unet_model_from_original_config(): - model = UNet2DConditionModel(**INPAINT_UNET_CONFIG) - - return model - - -def inpaint_unet_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - num_head_channels = INPAINT_UNET_CONFIG["attention_head_dim"] - - diffusers_checkpoint.update(unet_time_embeddings(checkpoint)) - diffusers_checkpoint.update(unet_conv_in(checkpoint)) - diffusers_checkpoint.update(unet_add_embedding(checkpoint)) - diffusers_checkpoint.update(unet_encoder_hid_proj(checkpoint)) - - # .input_blocks -> .down_blocks - - original_down_block_idx = 1 - - for diffusers_down_block_idx in range(len(model.down_blocks)): - checkpoint_update, num_original_down_blocks = unet_downblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_down_block_idx=diffusers_down_block_idx, - original_down_block_idx=original_down_block_idx, - num_head_channels=num_head_channels, - ) - - original_down_block_idx += num_original_down_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .input_blocks -> .down_blocks - - diffusers_checkpoint.update( - unet_midblock_to_diffusers_checkpoint( - model, - checkpoint, - num_head_channels=num_head_channels, - ) - ) - - # .output_blocks -> .up_blocks - - original_up_block_idx = 0 - - for diffusers_up_block_idx in range(len(model.up_blocks)): - checkpoint_update, num_original_up_blocks = unet_upblock_to_diffusers_checkpoint( - model, - checkpoint, - diffusers_up_block_idx=diffusers_up_block_idx, - original_up_block_idx=original_up_block_idx, - num_head_channels=num_head_channels, - ) - - original_up_block_idx += num_original_up_blocks - - diffusers_checkpoint.update(checkpoint_update) - - # done .output_blocks -> .up_blocks - - diffusers_checkpoint.update(unet_conv_norm_out(checkpoint)) - diffusers_checkpoint.update(unet_conv_out(checkpoint)) - - return diffusers_checkpoint - - -# done inpaint unet - - -# unet utils - - -# .time_embed -> .time_embedding -def unet_time_embeddings(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint["time_embed.0.weight"], - "time_embedding.linear_1.bias": checkpoint["time_embed.0.bias"], - "time_embedding.linear_2.weight": checkpoint["time_embed.2.weight"], - "time_embedding.linear_2.bias": checkpoint["time_embed.2.bias"], - } - ) - - return diffusers_checkpoint - - -# .input_blocks.0 -> .conv_in -def unet_conv_in(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_in.weight": checkpoint["input_blocks.0.0.weight"], - "conv_in.bias": checkpoint["input_blocks.0.0.bias"], - } - ) - - return diffusers_checkpoint - - -def unet_add_embedding(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "add_embedding.text_norm.weight": checkpoint["ln_model_n.weight"], - "add_embedding.text_norm.bias": checkpoint["ln_model_n.bias"], - "add_embedding.text_proj.weight": checkpoint["proj_n.weight"], - "add_embedding.text_proj.bias": checkpoint["proj_n.bias"], - "add_embedding.image_proj.weight": checkpoint["img_layer.weight"], - "add_embedding.image_proj.bias": checkpoint["img_layer.bias"], - } - ) - - return diffusers_checkpoint - - -def unet_encoder_hid_proj(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "encoder_hid_proj.image_embeds.weight": checkpoint["clip_to_seq.weight"], - "encoder_hid_proj.image_embeds.bias": checkpoint["clip_to_seq.bias"], - "encoder_hid_proj.text_proj.weight": checkpoint["to_model_dim_n.weight"], - "encoder_hid_proj.text_proj.bias": checkpoint["to_model_dim_n.bias"], - } - ) - - return diffusers_checkpoint - - -# .out.0 -> .conv_norm_out -def unet_conv_norm_out(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_norm_out.weight": checkpoint["out.0.weight"], - "conv_norm_out.bias": checkpoint["out.0.bias"], - } - ) - - return diffusers_checkpoint - - -# .out.2 -> .conv_out -def unet_conv_out(checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update( - { - "conv_out.weight": checkpoint["out.2.weight"], - "conv_out.bias": checkpoint["out.2.bias"], - } - ) - - return diffusers_checkpoint - - -# .input_blocks -> .down_blocks -def unet_downblock_to_diffusers_checkpoint( - model, checkpoint, *, diffusers_down_block_idx, original_down_block_idx, num_head_channels -): - diffusers_checkpoint = {} - - diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.resnets" - original_down_block_prefix = "input_blocks" - - down_block = model.down_blocks[diffusers_down_block_idx] - - num_resnets = len(down_block.resnets) - - if down_block.downsamplers is None: - downsampler = False - else: - assert len(down_block.downsamplers) == 1 - downsampler = True - # The downsample block is also a resnet - num_resnets += 1 - - for resnet_idx_inc in range(num_resnets): - full_resnet_prefix = f"{original_down_block_prefix}.{original_down_block_idx + resnet_idx_inc}.0" - - if downsampler and resnet_idx_inc == num_resnets - 1: - # this is a downsample block - full_diffusers_resnet_prefix = f"down_blocks.{diffusers_down_block_idx}.downsamplers.0" - else: - # this is a regular resnet block - full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}" - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix - ) - ) - - if hasattr(down_block, "attentions"): - num_attentions = len(down_block.attentions) - diffusers_attention_prefix = f"down_blocks.{diffusers_down_block_idx}.attentions" - - for attention_idx_inc in range(num_attentions): - full_attention_prefix = f"{original_down_block_prefix}.{original_down_block_idx + attention_idx_inc}.1" - full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}" - - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - attention_prefix=full_attention_prefix, - diffusers_attention_prefix=full_diffusers_attention_prefix, - num_head_channels=num_head_channels, - ) - ) - - num_original_down_blocks = num_resnets - - return diffusers_checkpoint, num_original_down_blocks - - -# .middle_block -> .mid_block -def unet_midblock_to_diffusers_checkpoint(model, checkpoint, *, num_head_channels): - diffusers_checkpoint = {} - - # block 0 - - original_block_idx = 0 - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, - diffusers_resnet_prefix="mid_block.resnets.0", - resnet_prefix=f"middle_block.{original_block_idx}", - ) - ) - - original_block_idx += 1 - - # optional block 1 - - if hasattr(model.mid_block, "attentions") and model.mid_block.attentions[0] is not None: - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix="mid_block.attentions.0", - attention_prefix=f"middle_block.{original_block_idx}", - num_head_channels=num_head_channels, - ) - ) - original_block_idx += 1 - - # block 1 or block 2 - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, - diffusers_resnet_prefix="mid_block.resnets.1", - resnet_prefix=f"middle_block.{original_block_idx}", - ) - ) - - return diffusers_checkpoint - - -# .output_blocks -> .up_blocks -def unet_upblock_to_diffusers_checkpoint( - model, checkpoint, *, diffusers_up_block_idx, original_up_block_idx, num_head_channels -): - diffusers_checkpoint = {} - - diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.resnets" - original_up_block_prefix = "output_blocks" - - up_block = model.up_blocks[diffusers_up_block_idx] - - num_resnets = len(up_block.resnets) - - if up_block.upsamplers is None: - upsampler = False - else: - assert len(up_block.upsamplers) == 1 - upsampler = True - # The upsample block is also a resnet - num_resnets += 1 - - has_attentions = hasattr(up_block, "attentions") - - for resnet_idx_inc in range(num_resnets): - if upsampler and resnet_idx_inc == num_resnets - 1: - # this is an upsample block - if has_attentions: - # There is a middle attention block that we skip - original_resnet_block_idx = 2 - else: - original_resnet_block_idx = 1 - - # we add the `minus 1` because the last two resnets are stuck together in the same output block - full_resnet_prefix = ( - f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc - 1}.{original_resnet_block_idx}" - ) - - full_diffusers_resnet_prefix = f"up_blocks.{diffusers_up_block_idx}.upsamplers.0" - else: - # this is a regular resnet block - full_resnet_prefix = f"{original_up_block_prefix}.{original_up_block_idx + resnet_idx_inc}.0" - full_diffusers_resnet_prefix = f"{diffusers_resnet_prefix}.{resnet_idx_inc}" - - diffusers_checkpoint.update( - resnet_to_diffusers_checkpoint( - checkpoint, resnet_prefix=full_resnet_prefix, diffusers_resnet_prefix=full_diffusers_resnet_prefix - ) - ) - - if has_attentions: - num_attentions = len(up_block.attentions) - diffusers_attention_prefix = f"up_blocks.{diffusers_up_block_idx}.attentions" - - for attention_idx_inc in range(num_attentions): - full_attention_prefix = f"{original_up_block_prefix}.{original_up_block_idx + attention_idx_inc}.1" - full_diffusers_attention_prefix = f"{diffusers_attention_prefix}.{attention_idx_inc}" - - diffusers_checkpoint.update( - attention_to_diffusers_checkpoint( - checkpoint, - attention_prefix=full_attention_prefix, - diffusers_attention_prefix=full_diffusers_attention_prefix, - num_head_channels=num_head_channels, - ) - ) - - num_original_down_blocks = num_resnets - 1 if upsampler else num_resnets - - return diffusers_checkpoint, num_original_down_blocks - - -def resnet_to_diffusers_checkpoint(checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - diffusers_checkpoint = { - f"{diffusers_resnet_prefix}.norm1.weight": checkpoint[f"{resnet_prefix}.in_layers.0.weight"], - f"{diffusers_resnet_prefix}.norm1.bias": checkpoint[f"{resnet_prefix}.in_layers.0.bias"], - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.in_layers.2.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.in_layers.2.bias"], - f"{diffusers_resnet_prefix}.time_emb_proj.weight": checkpoint[f"{resnet_prefix}.emb_layers.1.weight"], - f"{diffusers_resnet_prefix}.time_emb_proj.bias": checkpoint[f"{resnet_prefix}.emb_layers.1.bias"], - f"{diffusers_resnet_prefix}.norm2.weight": checkpoint[f"{resnet_prefix}.out_layers.0.weight"], - f"{diffusers_resnet_prefix}.norm2.bias": checkpoint[f"{resnet_prefix}.out_layers.0.bias"], - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.out_layers.3.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.out_layers.3.bias"], - } - - skip_connection_prefix = f"{resnet_prefix}.skip_connection" - - if f"{skip_connection_prefix}.weight" in checkpoint: - diffusers_checkpoint.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{skip_connection_prefix}.weight"], - f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{skip_connection_prefix}.bias"], - } - ) - - return diffusers_checkpoint - - -def attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix, num_head_channels): - diffusers_checkpoint = {} - - # .norm -> .group_norm - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.group_norm.weight": checkpoint[f"{attention_prefix}.norm.weight"], - f"{diffusers_attention_prefix}.group_norm.bias": checkpoint[f"{attention_prefix}.norm.bias"], - } - ) - - # .qkv -> .{query, key, value} - [q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions( - weight=checkpoint[f"{attention_prefix}.qkv.weight"][:, :, 0], - bias=checkpoint[f"{attention_prefix}.qkv.bias"], - split=3, - chunk_size=num_head_channels, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_q.weight": q_weight, - f"{diffusers_attention_prefix}.to_q.bias": q_bias, - f"{diffusers_attention_prefix}.to_k.weight": k_weight, - f"{diffusers_attention_prefix}.to_k.bias": k_bias, - f"{diffusers_attention_prefix}.to_v.weight": v_weight, - f"{diffusers_attention_prefix}.to_v.bias": v_bias, - } - ) - - # .encoder_kv -> .{context_key, context_value} - [encoder_k_weight, encoder_v_weight], [encoder_k_bias, encoder_v_bias] = split_attentions( - weight=checkpoint[f"{attention_prefix}.encoder_kv.weight"][:, :, 0], - bias=checkpoint[f"{attention_prefix}.encoder_kv.bias"], - split=2, - chunk_size=num_head_channels, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.add_k_proj.weight": encoder_k_weight, - f"{diffusers_attention_prefix}.add_k_proj.bias": encoder_k_bias, - f"{diffusers_attention_prefix}.add_v_proj.weight": encoder_v_weight, - f"{diffusers_attention_prefix}.add_v_proj.bias": encoder_v_bias, - } - ) - - # .proj_out (1d conv) -> .proj_attn (linear) - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][ - :, :, 0 - ], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj_out.bias"], - } - ) - - return diffusers_checkpoint - - -# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) -def split_attentions(*, weight, bias, split, chunk_size): - weights = [None] * split - biases = [None] * split - - weights_biases_idx = 0 - - for starting_row_index in range(0, weight.shape[0], chunk_size): - row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size) - - weight_rows = weight[row_indices, :] - bias_rows = bias[row_indices] - - if weights[weights_biases_idx] is None: - assert weights[weights_biases_idx] is None - weights[weights_biases_idx] = weight_rows - biases[weights_biases_idx] = bias_rows - else: - assert weights[weights_biases_idx] is not None - weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows]) - biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows]) - - weights_biases_idx = (weights_biases_idx + 1) % split - - return weights, biases - - -# done unet utils - - -def prior(*, args, checkpoint_map_location): - print("loading prior") - - prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) - - clip_stats_checkpoint = torch.load(args.clip_stat_path, map_location=checkpoint_map_location) - - prior_model = prior_model_from_original_config() - - prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint( - prior_model, prior_checkpoint, clip_stats_checkpoint - ) - - del prior_checkpoint - del clip_stats_checkpoint - - load_checkpoint_to_model(prior_diffusers_checkpoint, prior_model, strict=True) - - print("done loading prior") - - return prior_model - - -def text2img(*, args, checkpoint_map_location): - print("loading text2img") - - text2img_checkpoint = torch.load(args.text2img_checkpoint_path, map_location=checkpoint_map_location) - - unet_model = unet_model_from_original_config() - - unet_diffusers_checkpoint = unet_original_checkpoint_to_diffusers_checkpoint(unet_model, text2img_checkpoint) - - del text2img_checkpoint - - load_checkpoint_to_model(unet_diffusers_checkpoint, unet_model, strict=True) - - print("done loading text2img") - - return unet_model - - -def inpaint_text2img(*, args, checkpoint_map_location): - print("loading inpaint text2img") - - inpaint_text2img_checkpoint = torch.load( - args.inpaint_text2img_checkpoint_path, map_location=checkpoint_map_location - ) - - inpaint_unet_model = inpaint_unet_model_from_original_config() - - inpaint_unet_diffusers_checkpoint = inpaint_unet_original_checkpoint_to_diffusers_checkpoint( - inpaint_unet_model, inpaint_text2img_checkpoint - ) - - del inpaint_text2img_checkpoint - - load_checkpoint_to_model(inpaint_unet_diffusers_checkpoint, inpaint_unet_model, strict=True) - - print("done loading inpaint text2img") - - return inpaint_unet_model - - -# movq - -MOVQ_CONFIG = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 4, - "down_block_types": ("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "AttnDownEncoderBlock2D"), - "up_block_types": ("AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"), - "num_vq_embeddings": 16384, - "block_out_channels": (128, 256, 256, 512), - "vq_embed_dim": 4, - "layers_per_block": 2, - "norm_type": "spatial", -} - - -def movq_model_from_original_config(): - movq = VQModel(**MOVQ_CONFIG) - return movq - - -def movq_encoder_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # conv_in - diffusers_checkpoint.update( - { - "encoder.conv_in.weight": checkpoint["encoder.conv_in.weight"], - "encoder.conv_in.bias": checkpoint["encoder.conv_in.bias"], - } - ) - - # down_blocks - for down_block_idx, down_block in enumerate(model.encoder.down_blocks): - diffusers_down_block_prefix = f"encoder.down_blocks.{down_block_idx}" - down_block_prefix = f"encoder.down.{down_block_idx}" - - # resnets - for resnet_idx, resnet in enumerate(down_block.resnets): - diffusers_resnet_prefix = f"{diffusers_down_block_prefix}.resnets.{resnet_idx}" - resnet_prefix = f"{down_block_prefix}.block.{resnet_idx}" - - diffusers_checkpoint.update( - movq_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - # downsample - - # do not include the downsample when on the last down block - # There is no downsample on the last down block - if down_block_idx != len(model.encoder.down_blocks) - 1: - # There's a single downsample in the original checkpoint but a list of downsamples - # in the diffusers model. - diffusers_downsample_prefix = f"{diffusers_down_block_prefix}.downsamplers.0.conv" - downsample_prefix = f"{down_block_prefix}.downsample.conv" - diffusers_checkpoint.update( - { - f"{diffusers_downsample_prefix}.weight": checkpoint[f"{downsample_prefix}.weight"], - f"{diffusers_downsample_prefix}.bias": checkpoint[f"{downsample_prefix}.bias"], - } - ) - - # attentions - - if hasattr(down_block, "attentions"): - for attention_idx, _ in enumerate(down_block.attentions): - diffusers_attention_prefix = f"{diffusers_down_block_prefix}.attentions.{attention_idx}" - attention_prefix = f"{down_block_prefix}.attn.{attention_idx}" - diffusers_checkpoint.update( - movq_attention_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - attention_prefix=attention_prefix, - ) - ) - - # mid block - - # mid block attentions - - # There is a single hardcoded attention block in the middle of the VQ-diffusion encoder - diffusers_attention_prefix = "encoder.mid_block.attentions.0" - attention_prefix = "encoder.mid.attn_1" - diffusers_checkpoint.update( - movq_attention_to_diffusers_checkpoint( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # mid block resnets - - for diffusers_resnet_idx, resnet in enumerate(model.encoder.mid_block.resnets): - diffusers_resnet_prefix = f"encoder.mid_block.resnets.{diffusers_resnet_idx}" - - # the hardcoded prefixes to `block_` are 1 and 2 - orig_resnet_idx = diffusers_resnet_idx + 1 - # There are two hardcoded resnets in the middle of the VQ-diffusion encoder - resnet_prefix = f"encoder.mid.block_{orig_resnet_idx}" - - diffusers_checkpoint.update( - movq_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - diffusers_checkpoint.update( - { - # conv_norm_out - "encoder.conv_norm_out.weight": checkpoint["encoder.norm_out.weight"], - "encoder.conv_norm_out.bias": checkpoint["encoder.norm_out.bias"], - # conv_out - "encoder.conv_out.weight": checkpoint["encoder.conv_out.weight"], - "encoder.conv_out.bias": checkpoint["encoder.conv_out.bias"], - } - ) - - return diffusers_checkpoint - - -def movq_decoder_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # conv in - diffusers_checkpoint.update( - { - "decoder.conv_in.weight": checkpoint["decoder.conv_in.weight"], - "decoder.conv_in.bias": checkpoint["decoder.conv_in.bias"], - } - ) - - # up_blocks - - for diffusers_up_block_idx, up_block in enumerate(model.decoder.up_blocks): - # up_blocks are stored in reverse order in the VQ-diffusion checkpoint - orig_up_block_idx = len(model.decoder.up_blocks) - 1 - diffusers_up_block_idx - - diffusers_up_block_prefix = f"decoder.up_blocks.{diffusers_up_block_idx}" - up_block_prefix = f"decoder.up.{orig_up_block_idx}" - - # resnets - for resnet_idx, resnet in enumerate(up_block.resnets): - diffusers_resnet_prefix = f"{diffusers_up_block_prefix}.resnets.{resnet_idx}" - resnet_prefix = f"{up_block_prefix}.block.{resnet_idx}" - - diffusers_checkpoint.update( - movq_resnet_to_diffusers_checkpoint_spatial_norm( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - # upsample - - # there is no up sample on the last up block - if diffusers_up_block_idx != len(model.decoder.up_blocks) - 1: - # There's a single upsample in the VQ-diffusion checkpoint but a list of downsamples - # in the diffusers model. - diffusers_downsample_prefix = f"{diffusers_up_block_prefix}.upsamplers.0.conv" - downsample_prefix = f"{up_block_prefix}.upsample.conv" - diffusers_checkpoint.update( - { - f"{diffusers_downsample_prefix}.weight": checkpoint[f"{downsample_prefix}.weight"], - f"{diffusers_downsample_prefix}.bias": checkpoint[f"{downsample_prefix}.bias"], - } - ) - - # attentions - - if hasattr(up_block, "attentions"): - for attention_idx, _ in enumerate(up_block.attentions): - diffusers_attention_prefix = f"{diffusers_up_block_prefix}.attentions.{attention_idx}" - attention_prefix = f"{up_block_prefix}.attn.{attention_idx}" - diffusers_checkpoint.update( - movq_attention_to_diffusers_checkpoint_spatial_norm( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - attention_prefix=attention_prefix, - ) - ) - - # mid block - - # mid block attentions - - # There is a single hardcoded attention block in the middle of the VQ-diffusion decoder - diffusers_attention_prefix = "decoder.mid_block.attentions.0" - attention_prefix = "decoder.mid.attn_1" - diffusers_checkpoint.update( - movq_attention_to_diffusers_checkpoint_spatial_norm( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # mid block resnets - - for diffusers_resnet_idx, resnet in enumerate(model.encoder.mid_block.resnets): - diffusers_resnet_prefix = f"decoder.mid_block.resnets.{diffusers_resnet_idx}" - - # the hardcoded prefixes to `block_` are 1 and 2 - orig_resnet_idx = diffusers_resnet_idx + 1 - # There are two hardcoded resnets in the middle of the VQ-diffusion decoder - resnet_prefix = f"decoder.mid.block_{orig_resnet_idx}" - - diffusers_checkpoint.update( - movq_resnet_to_diffusers_checkpoint_spatial_norm( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - diffusers_checkpoint.update( - { - # conv_norm_out - "decoder.conv_norm_out.norm_layer.weight": checkpoint["decoder.norm_out.norm_layer.weight"], - "decoder.conv_norm_out.norm_layer.bias": checkpoint["decoder.norm_out.norm_layer.bias"], - "decoder.conv_norm_out.conv_y.weight": checkpoint["decoder.norm_out.conv_y.weight"], - "decoder.conv_norm_out.conv_y.bias": checkpoint["decoder.norm_out.conv_y.bias"], - "decoder.conv_norm_out.conv_b.weight": checkpoint["decoder.norm_out.conv_b.weight"], - "decoder.conv_norm_out.conv_b.bias": checkpoint["decoder.norm_out.conv_b.bias"], - # conv_out - "decoder.conv_out.weight": checkpoint["decoder.conv_out.weight"], - "decoder.conv_out.bias": checkpoint["decoder.conv_out.bias"], - } - ) - - return diffusers_checkpoint - - -def movq_resnet_to_diffusers_checkpoint(resnet, checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - rv = { - # norm1 - f"{diffusers_resnet_prefix}.norm1.weight": checkpoint[f"{resnet_prefix}.norm1.weight"], - f"{diffusers_resnet_prefix}.norm1.bias": checkpoint[f"{resnet_prefix}.norm1.bias"], - # conv1 - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.conv1.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.conv1.bias"], - # norm2 - f"{diffusers_resnet_prefix}.norm2.weight": checkpoint[f"{resnet_prefix}.norm2.weight"], - f"{diffusers_resnet_prefix}.norm2.bias": checkpoint[f"{resnet_prefix}.norm2.bias"], - # conv2 - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.conv2.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.conv2.bias"], - } - - if resnet.conv_shortcut is not None: - rv.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{resnet_prefix}.nin_shortcut.weight"], - f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{resnet_prefix}.nin_shortcut.bias"], - } - ) - - return rv - - -def movq_resnet_to_diffusers_checkpoint_spatial_norm(resnet, checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - rv = { - # norm1 - f"{diffusers_resnet_prefix}.norm1.norm_layer.weight": checkpoint[f"{resnet_prefix}.norm1.norm_layer.weight"], - f"{diffusers_resnet_prefix}.norm1.norm_layer.bias": checkpoint[f"{resnet_prefix}.norm1.norm_layer.bias"], - f"{diffusers_resnet_prefix}.norm1.conv_y.weight": checkpoint[f"{resnet_prefix}.norm1.conv_y.weight"], - f"{diffusers_resnet_prefix}.norm1.conv_y.bias": checkpoint[f"{resnet_prefix}.norm1.conv_y.bias"], - f"{diffusers_resnet_prefix}.norm1.conv_b.weight": checkpoint[f"{resnet_prefix}.norm1.conv_b.weight"], - f"{diffusers_resnet_prefix}.norm1.conv_b.bias": checkpoint[f"{resnet_prefix}.norm1.conv_b.bias"], - # conv1 - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.conv1.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.conv1.bias"], - # norm2 - f"{diffusers_resnet_prefix}.norm2.norm_layer.weight": checkpoint[f"{resnet_prefix}.norm2.norm_layer.weight"], - f"{diffusers_resnet_prefix}.norm2.norm_layer.bias": checkpoint[f"{resnet_prefix}.norm2.norm_layer.bias"], - f"{diffusers_resnet_prefix}.norm2.conv_y.weight": checkpoint[f"{resnet_prefix}.norm2.conv_y.weight"], - f"{diffusers_resnet_prefix}.norm2.conv_y.bias": checkpoint[f"{resnet_prefix}.norm2.conv_y.bias"], - f"{diffusers_resnet_prefix}.norm2.conv_b.weight": checkpoint[f"{resnet_prefix}.norm2.conv_b.weight"], - f"{diffusers_resnet_prefix}.norm2.conv_b.bias": checkpoint[f"{resnet_prefix}.norm2.conv_b.bias"], - # conv2 - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.conv2.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.conv2.bias"], - } - - if resnet.conv_shortcut is not None: - rv.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{resnet_prefix}.nin_shortcut.weight"], - f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{resnet_prefix}.nin_shortcut.bias"], - } - ) - - return rv - - -def movq_attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix): - return { - # norm - f"{diffusers_attention_prefix}.group_norm.weight": checkpoint[f"{attention_prefix}.norm.weight"], - f"{diffusers_attention_prefix}.group_norm.bias": checkpoint[f"{attention_prefix}.norm.bias"], - # query - f"{diffusers_attention_prefix}.to_q.weight": checkpoint[f"{attention_prefix}.q.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_q.bias": checkpoint[f"{attention_prefix}.q.bias"], - # key - f"{diffusers_attention_prefix}.to_k.weight": checkpoint[f"{attention_prefix}.k.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_k.bias": checkpoint[f"{attention_prefix}.k.bias"], - # value - f"{diffusers_attention_prefix}.to_v.weight": checkpoint[f"{attention_prefix}.v.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_v.bias": checkpoint[f"{attention_prefix}.v.bias"], - # proj_attn - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj_out.bias"], - } - - -def movq_attention_to_diffusers_checkpoint_spatial_norm(checkpoint, *, diffusers_attention_prefix, attention_prefix): - return { - # norm - f"{diffusers_attention_prefix}.spatial_norm.norm_layer.weight": checkpoint[ - f"{attention_prefix}.norm.norm_layer.weight" - ], - f"{diffusers_attention_prefix}.spatial_norm.norm_layer.bias": checkpoint[ - f"{attention_prefix}.norm.norm_layer.bias" - ], - f"{diffusers_attention_prefix}.spatial_norm.conv_y.weight": checkpoint[ - f"{attention_prefix}.norm.conv_y.weight" - ], - f"{diffusers_attention_prefix}.spatial_norm.conv_y.bias": checkpoint[f"{attention_prefix}.norm.conv_y.bias"], - f"{diffusers_attention_prefix}.spatial_norm.conv_b.weight": checkpoint[ - f"{attention_prefix}.norm.conv_b.weight" - ], - f"{diffusers_attention_prefix}.spatial_norm.conv_b.bias": checkpoint[f"{attention_prefix}.norm.conv_b.bias"], - # query - f"{diffusers_attention_prefix}.to_q.weight": checkpoint[f"{attention_prefix}.q.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_q.bias": checkpoint[f"{attention_prefix}.q.bias"], - # key - f"{diffusers_attention_prefix}.to_k.weight": checkpoint[f"{attention_prefix}.k.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_k.bias": checkpoint[f"{attention_prefix}.k.bias"], - # value - f"{diffusers_attention_prefix}.to_v.weight": checkpoint[f"{attention_prefix}.v.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_v.bias": checkpoint[f"{attention_prefix}.v.bias"], - # proj_attn - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj_out.bias"], - } - - -def movq_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - diffusers_checkpoint.update(movq_encoder_to_diffusers_checkpoint(model, checkpoint)) - - # quant_conv - - diffusers_checkpoint.update( - { - "quant_conv.weight": checkpoint["quant_conv.weight"], - "quant_conv.bias": checkpoint["quant_conv.bias"], - } - ) - - # quantize - diffusers_checkpoint.update({"quantize.embedding.weight": checkpoint["quantize.embedding.weight"]}) - - # post_quant_conv - diffusers_checkpoint.update( - { - "post_quant_conv.weight": checkpoint["post_quant_conv.weight"], - "post_quant_conv.bias": checkpoint["post_quant_conv.bias"], - } - ) - - # decoder - diffusers_checkpoint.update(movq_decoder_to_diffusers_checkpoint(model, checkpoint)) - - return diffusers_checkpoint - - -def movq(*, args, checkpoint_map_location): - print("loading movq") - - movq_checkpoint = torch.load(args.movq_checkpoint_path, map_location=checkpoint_map_location) - - movq_model = movq_model_from_original_config() - - movq_diffusers_checkpoint = movq_original_checkpoint_to_diffusers_checkpoint(movq_model, movq_checkpoint) - - del movq_checkpoint - - load_checkpoint_to_model(movq_diffusers_checkpoint, movq_model, strict=True) - - print("done loading movq") - - return movq_model - - -def load_checkpoint_to_model(checkpoint, model, strict=False): - with tempfile.NamedTemporaryFile(delete=False) as file: - torch.save(checkpoint, file.name) - del checkpoint - if strict: - model.load_state_dict(torch.load(file.name), strict=True) - else: - load_checkpoint_and_dispatch(model, file.name, device_map="auto") - os.remove(file.name) - - -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( - "--prior_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the prior checkpoint to convert.", - ) - parser.add_argument( - "--clip_stat_path", - default=None, - type=str, - required=False, - help="Path to the clip stats checkpoint to convert.", - ) - parser.add_argument( - "--text2img_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the text2img checkpoint to convert.", - ) - parser.add_argument( - "--movq_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the text2img checkpoint to convert.", - ) - parser.add_argument( - "--inpaint_text2img_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the inpaint text2img checkpoint to convert.", - ) - parser.add_argument( - "--checkpoint_load_device", - default="cpu", - type=str, - required=False, - help="The device passed to `map_location` when loading checkpoints.", - ) - - parser.add_argument( - "--debug", - default=None, - type=str, - required=False, - help="Only run a specific stage of the convert script. Used for debugging", - ) - - args = parser.parse_args() - - print(f"loading checkpoints to {args.checkpoint_load_device}") - - checkpoint_map_location = torch.device(args.checkpoint_load_device) - - if args.debug is not None: - print(f"debug: only executing {args.debug}") - - if args.debug is None: - print("to-do") - elif args.debug == "prior": - prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) - prior_model.save_pretrained(args.dump_path) - elif args.debug == "text2img": - unet_model = text2img(args=args, checkpoint_map_location=checkpoint_map_location) - unet_model.save_pretrained(f"{args.dump_path}/unet") - elif args.debug == "inpaint_text2img": - inpaint_unet_model = inpaint_text2img(args=args, checkpoint_map_location=checkpoint_map_location) - inpaint_unet_model.save_pretrained(f"{args.dump_path}/inpaint_unet") - elif args.debug == "decoder": - decoder = movq(args=args, checkpoint_map_location=checkpoint_map_location) - decoder.save_pretrained(f"{args.dump_path}/decoder") - else: - raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/convert_ldm_original_checkpoint_to_diffusers.py b/scripts/convert_ldm_original_checkpoint_to_diffusers.py deleted file mode 100644 index cdaf317af752..000000000000 --- a/scripts/convert_ldm_original_checkpoint_to_diffusers.py +++ /dev/null @@ -1,359 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the LDM checkpoints.""" - -import argparse -import json - -import torch - -from diffusers import DDPMScheduler, LDMPipeline, UNet2DModel, VQModel - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("proj_out.weight", "proj_attn.weight") - new_item = new_item.replace("proj_out.bias", "proj_attn.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming - to them. It splits attention layers, and takes into account additional replacements - that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def convert_ldm_checkpoint(checkpoint, config): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = checkpoint["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = checkpoint["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = checkpoint["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = checkpoint["time_embed.2.bias"] - - new_checkpoint["conv_in.weight"] = checkpoint["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = checkpoint["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = checkpoint["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = checkpoint["out.0.bias"] - new_checkpoint["conv_out.weight"] = checkpoint["out.2.weight"] - new_checkpoint["conv_out.bias"] = checkpoint["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in checkpoint if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in checkpoint if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in checkpoint if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in checkpoint if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in checkpoint if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in checkpoint if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["num_res_blocks"] + 1) - layer_in_block_id = (i - 1) % (config["num_res_blocks"] + 1) - - resnets = [key for key in input_blocks[i] if f"input_blocks.{i}.0" in key] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in checkpoint: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = checkpoint[ - f"input_blocks.{i}.0.op.weight" - ] - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = checkpoint[ - f"input_blocks.{i}.0.op.bias" - ] - continue - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - resnet_op = {"old": "resnets.2.op", "new": "downsamplers.0.op"} - assign_to_checkpoint( - paths, new_checkpoint, checkpoint, additional_replacements=[meta_path, resnet_op], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"input_blocks.{i}.1", - "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}", - } - to_split = { - f"input_blocks.{i}.1.qkv.bias": { - "key": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias", - "query": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias", - "value": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias", - }, - f"input_blocks.{i}.1.qkv.weight": { - "key": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight", - "query": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight", - "value": f"down_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight", - }, - } - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[meta_path], - attention_paths_to_split=to_split, - config=config, - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, checkpoint, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, checkpoint, config=config) - - attentions_paths = renew_attention_paths(attentions) - to_split = { - "middle_block.1.qkv.bias": { - "key": "mid_block.attentions.0.key.bias", - "query": "mid_block.attentions.0.query.bias", - "value": "mid_block.attentions.0.value.bias", - }, - "middle_block.1.qkv.weight": { - "key": "mid_block.attentions.0.key.weight", - "query": "mid_block.attentions.0.query.weight", - "value": "mid_block.attentions.0.value.weight", - }, - } - assign_to_checkpoint( - attentions_paths, new_checkpoint, checkpoint, attention_paths_to_split=to_split, config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["num_res_blocks"] + 1) - layer_in_block_id = i % (config["num_res_blocks"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint(paths, new_checkpoint, checkpoint, additional_replacements=[meta_path], config=config) - - if ["conv.weight", "conv.bias"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.weight", "conv.bias"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = checkpoint[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = checkpoint[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - to_split = { - f"output_blocks.{i}.1.qkv.bias": { - "key": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.key.bias", - "query": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.query.bias", - "value": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.value.bias", - }, - f"output_blocks.{i}.1.qkv.weight": { - "key": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.key.weight", - "query": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.query.weight", - "value": f"up_blocks.{block_id}.attentions.{layer_in_block_id}.value.weight", - }, - } - assign_to_checkpoint( - paths, - new_checkpoint, - checkpoint, - additional_replacements=[meta_path], - attention_paths_to_split=to_split if any("qkv" in key for key in attentions) else None, - config=config, - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = checkpoint[old_path] - - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - - parser.add_argument( - "--config_file", - default=None, - type=str, - required=True, - help="The config json file corresponding to the architecture.", - ) - - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - - args = parser.parse_args() - - checkpoint = torch.load(args.checkpoint_path) - - with open(args.config_file) as f: - config = json.loads(f.read()) - - converted_checkpoint = convert_ldm_checkpoint(checkpoint, config) - - if "ldm" in config: - del config["ldm"] - - model = UNet2DModel(**config) - model.load_state_dict(converted_checkpoint) - - try: - scheduler = DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) - vqvae = VQModel.from_pretrained("/".join(args.checkpoint_path.split("/")[:-1])) - - pipe = LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae) - pipe.save_pretrained(args.dump_path) - except: # noqa: E722 - model.save_pretrained(args.dump_path) diff --git a/scripts/convert_lora_safetensor_to_diffusers.py b/scripts/convert_lora_safetensor_to_diffusers.py deleted file mode 100644 index 4237452e2ed1..000000000000 --- a/scripts/convert_lora_safetensor_to_diffusers.py +++ /dev/null @@ -1,128 +0,0 @@ -# coding=utf-8 -# Copyright 2024, Haofan Wang, Qixun Wang, All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Conversion script for the LoRA's safetensors checkpoints.""" - -import argparse - -import torch -from safetensors.torch import load_file - -from diffusers import StableDiffusionPipeline - - -def convert(base_model_path, checkpoint_path, LORA_PREFIX_UNET, LORA_PREFIX_TEXT_ENCODER, alpha): - # load base model - pipeline = StableDiffusionPipeline.from_pretrained(base_model_path, torch_dtype=torch.float32) - - # load LoRA weight from .safetensors - state_dict = load_file(checkpoint_path) - - visited = [] - - # directly update weight in diffusers model - for key in state_dict: - # it is suggested to print out the key, it usually will be something like below - # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight" - - # as we have set the alpha beforehand, so just skip - if ".alpha" in key or key in visited: - continue - - if "text" in key: - layer_infos = key.split(".")[0].split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_") - curr_layer = pipeline.text_encoder - else: - layer_infos = key.split(".")[0].split(LORA_PREFIX_UNET + "_")[-1].split("_") - curr_layer = pipeline.unet - - # find the target layer - temp_name = layer_infos.pop(0) - while len(layer_infos) > -1: - try: - curr_layer = curr_layer.__getattr__(temp_name) - if len(layer_infos) > 0: - temp_name = layer_infos.pop(0) - elif len(layer_infos) == 0: - break - except Exception: - if len(temp_name) > 0: - temp_name += "_" + layer_infos.pop(0) - else: - temp_name = layer_infos.pop(0) - - pair_keys = [] - if "lora_down" in key: - pair_keys.append(key.replace("lora_down", "lora_up")) - pair_keys.append(key) - else: - pair_keys.append(key) - pair_keys.append(key.replace("lora_up", "lora_down")) - - # update weight - if len(state_dict[pair_keys[0]].shape) == 4: - weight_up = state_dict[pair_keys[0]].squeeze(3).squeeze(2).to(torch.float32) - weight_down = state_dict[pair_keys[1]].squeeze(3).squeeze(2).to(torch.float32) - curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down).unsqueeze(2).unsqueeze(3) - else: - weight_up = state_dict[pair_keys[0]].to(torch.float32) - weight_down = state_dict[pair_keys[1]].to(torch.float32) - curr_layer.weight.data += alpha * torch.mm(weight_up, weight_down) - - # update visited list - for item in pair_keys: - visited.append(item) - - return pipeline - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--base_model_path", default=None, type=str, required=True, help="Path to the base model in diffusers format." - ) - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument( - "--lora_prefix_unet", default="lora_unet", type=str, help="The prefix of UNet weight in safetensors" - ) - parser.add_argument( - "--lora_prefix_text_encoder", - default="lora_te", - type=str, - help="The prefix of text encoder weight in safetensors", - ) - parser.add_argument("--alpha", default=0.75, type=float, help="The merging ratio in W = W0 + alpha * deltaW") - parser.add_argument( - "--to_safetensors", action="store_true", help="Whether to store pipeline in safetensors format or not." - ) - parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") - - args = parser.parse_args() - - base_model_path = args.base_model_path - checkpoint_path = args.checkpoint_path - dump_path = args.dump_path - lora_prefix_unet = args.lora_prefix_unet - lora_prefix_text_encoder = args.lora_prefix_text_encoder - alpha = args.alpha - - pipe = convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha) - - pipe = pipe.to(args.device) - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py deleted file mode 100644 index 33b91790ef1c..000000000000 --- a/scripts/convert_ltx2_to_diffusers.py +++ /dev/null @@ -1,1703 +0,0 @@ -import argparse -import os -from contextlib import nullcontext -from typing import Any - -import safetensors.torch -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download -from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer - -from diffusers import ( - AutoencoderKLLTX2Audio, - AutoencoderKLLTX2Video, - FlowMatchEulerDiscreteScheduler, - LTX2LatentUpsamplePipeline, - LTX2Pipeline, - LTX2VideoDiffusionDecoderModel, - LTX2VideoTransformer3DModel, -) -from diffusers.pipelines.ltx2 import ( - LTX2DurationHead, - LTX2LatentUpsamplerModel, - LTX2TextConnectors, - LTX2Vocoder, - LTX2VocoderWithBWE, -) -from diffusers.utils.import_utils import is_accelerate_available - - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - - -LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = { - # Input Patchify Projections - "patchify_proj": "proj_in", - "audio_patchify_proj": "audio_proj_in", - # Modulation Parameters - # Handle adaln_single --> time_embed, audioln_single --> audio_time_embed separately as the original keys are - # substrings of the other modulation parameters below - "av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift", - "av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate", - "av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift", - "av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate", - # Transformer Blocks - # Per-Block Cross Attention Modulatin Parameters - "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", - "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", - # Attention QK Norms - "q_norm": "norm_q", - "k_norm": "norm_k", -} - -LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT = { - **LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT, - "audio_prompt_adaln_single": "audio_prompt_adaln", - "prompt_adaln_single": "prompt_adaln", -} - -LTX_2_0_VIDEO_VAE_RENAME_DICT = { - # Encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.1", - "down_blocks.3": "down_blocks.1.downsamplers.0", - "down_blocks.4": "down_blocks.2", - "down_blocks.5": "down_blocks.2.downsamplers.0", - "down_blocks.6": "down_blocks.3", - "down_blocks.7": "down_blocks.3.downsamplers.0", - "down_blocks.8": "mid_block", - # Decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", - # Common - # For all 3D ResNets - "res_blocks": "resnets", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", -} - -LTX_2_3_VIDEO_VAE_RENAME_DICT = { - **LTX_2_0_VIDEO_VAE_RENAME_DICT, - # Decoder extra blocks - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", -} - -# LTX-2.5's diffusion decoder replaces the conv decoder while keeping the same encoder, so only the -# `decoder.*` half of the VAE checkpoint is renamed with these rules. -LTX_2_5_DIFFUSION_DECODER_RENAME_DICT = { - # The original `t_embedder` *is* diffusers' `PixArtAlphaCombinedTimestepSizeEmbeddings`, saved under - # shorter names, so only its two Linears need renaming. - "t_embedder.mlp.0.": "t_embedder.timestep_embedder.linear_1.", - "t_embedder.mlp.2.": "t_embedder.timestep_embedder.linear_2.", - ".attn.proj.": ".attn.to_out.0.", - ".attn.q_norm.": ".attn.norm_q.", - ".attn.k_norm.": ".attn.norm_k.", -} - -# Where a checkpoint carries static AdaLN gates, each is folded into the Linear it gates (W <- g * W) and -# dropped, because the decoder's residuals are ungated. Maps a renamed parameter to its gate's suffix. -LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS = { - ".attn.to_out.0.weight": ".gate_msa", - ".attn.to_out.0.bias": ".gate_msa", - ".mlp.w_down.weight": ".gate_mlp", - ".context_proj.weight": ".gate_ctx", - ".context_proj.bias": ".gate_ctx", -} - -LTX_2_0_AUDIO_VAE_RENAME_DICT = { - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", -} - -LTX_2_0_VOCODER_RENAME_DICT = { - "ups": "upsamplers", - "resblocks": "resnets", - "conv_pre": "conv_in", - "conv_post": "conv_out", -} - -LTX_2_3_VOCODER_RENAME_DICT = { - # Handle upsamplers ("ups" --> "upsamplers") due to name clash - "resblocks": "resnets", - "conv_pre": "conv_in", - "conv_post": "conv_out", - "act_post": "act_out", - "downsample.lowpass": "downsample", -} - -LTX_2_0_CONNECTORS_KEYS_RENAME_DICT = { - "connectors.": "", - "video_embeddings_connector": "video_connector", - "audio_embeddings_connector": "audio_connector", - "transformer_1d_blocks": "transformer_blocks", - "text_embedding_projection.aggregate_embed": "text_proj_in", - # Attention QK Norms - "q_norm": "norm_q", - "k_norm": "norm_k", -} - -LTX_2_3_CONNECTORS_KEYS_RENAME_DICT = { - "connectors.": "", - "video_embeddings_connector": "video_connector", - "audio_embeddings_connector": "audio_connector", - "transformer_1d_blocks": "transformer_blocks", - # LTX-2.3 uses per-modality embedding projections - "text_embedding_projection.audio_aggregate_embed": "audio_text_proj_in", - "text_embedding_projection.video_aggregate_embed": "video_text_proj_in", - # Attention QK Norms - "q_norm": "norm_q", - "k_norm": "norm_k", -} - - -def update_state_dict_inplace(state_dict: dict[str, Any], old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - -def remove_keys_inplace(key: str, state_dict: dict[str, Any]) -> None: - state_dict.pop(key) - - -def convert_ltx2_transformer_adaln_single(key: str, state_dict: dict[str, Any]) -> None: - # Skip if not a weight, bias - if ".weight" not in key and ".bias" not in key: - return - - if key.startswith("adaln_single."): - new_key = key.replace("adaln_single.", "time_embed.") - param = state_dict.pop(key) - state_dict[new_key] = param - - if key.startswith("audio_adaln_single."): - new_key = key.replace("audio_adaln_single.", "audio_time_embed.") - param = state_dict.pop(key) - state_dict[new_key] = param - - return - - -def convert_ltx2_audio_vae_per_channel_statistics(key: str, state_dict: dict[str, Any]) -> None: - if key.startswith("per_channel_statistics"): - new_key = ".".join(["decoder", key]) - param = state_dict.pop(key) - state_dict[new_key] = param - - return - - -def convert_ltx2_3_vocoder_upsamplers(key: str, state_dict: dict[str, Any]) -> None: - # Skip if not a weight, bias - if ".weight" not in key and ".bias" not in key: - return - - if ".ups." in key: - new_key = key.replace(".ups.", ".upsamplers.") - param = state_dict.pop(key) - state_dict[new_key] = param - return - - -LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = { - "video_embeddings_connector": remove_keys_inplace, - "audio_embeddings_connector": remove_keys_inplace, - "adaln_single": convert_ltx2_transformer_adaln_single, -} - -LTX_2_0_VAE_SPECIAL_KEYS_REMAP = { - "per_channel_statistics.channel": remove_keys_inplace, - "per_channel_statistics.mean-of-stds": remove_keys_inplace, -} - -LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP = {} - -LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP = {} - -LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP = { - ".ups.": convert_ltx2_3_vocoder_upsamplers, -} - -LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP = {} - - -def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: - connector_prefixes = ( - "video_embeddings_connector", - "audio_embeddings_connector", - "transformer_1d_blocks", - "text_embedding_projection", - "connectors.", - "video_connector", - "audio_connector", - "text_proj_in", - ) - - transformer_state_dict, connector_state_dict = {}, {} - for key, value in state_dict.items(): - if key.startswith(connector_prefixes): - connector_state_dict[key] = value - else: - transformer_state_dict[key] = value - - return transformer_state_dict, connector_state_dict - - -def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version == "test": - # Produces a transformer of the same size as used in test_models_transformer_ltx2.py - config = { - "model_id": "diffusers-internal-dev/dummy-ltx2", - "diffusers_config": { - "in_channels": 4, - "out_channels": 4, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 2, - "attention_head_dim": 8, - "cross_attention_dim": 16, - "vae_scale_factors": (8, 32, 32), - "pos_embed_max_pos": 20, - "base_height": 2048, - "base_width": 2048, - "audio_in_channels": 4, - "audio_out_channels": 4, - "audio_patch_size": 1, - "audio_patch_size_t": 1, - "audio_num_attention_heads": 2, - "audio_attention_head_dim": 4, - "audio_cross_attention_dim": 8, - "audio_scale_factor": 4, - "audio_pos_embed_max_pos": 20, - "audio_sampling_rate": 16000, - "audio_hop_length": 160, - "num_layers": 2, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 16, - "attention_bias": True, - "attention_out_bias": True, - "rope_theta": 10000.0, - "rope_double_precision": False, - "causal_offset": 1, - "timestep_scale_multiplier": 1000, - "cross_attn_timestep_scale_multiplier": 1, - }, - } - rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP - elif version == "2.0": - config = { - "model_id": "Lightricks/LTX-2", - "diffusers_config": { - "in_channels": 128, - "out_channels": 128, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 32, - "attention_head_dim": 128, - "cross_attention_dim": 4096, - "vae_scale_factors": (8, 32, 32), - "pos_embed_max_pos": 20, - "base_height": 2048, - "base_width": 2048, - "gated_attn": False, - "cross_attn_mod": False, - "audio_in_channels": 128, - "audio_out_channels": 128, - "audio_patch_size": 1, - "audio_patch_size_t": 1, - "audio_num_attention_heads": 32, - "audio_attention_head_dim": 64, - "audio_cross_attention_dim": 2048, - "audio_scale_factor": 4, - "audio_pos_embed_max_pos": 20, - "audio_sampling_rate": 16000, - "audio_hop_length": 160, - "audio_gated_attn": False, - "audio_cross_attn_mod": False, - "num_layers": 48, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 3840, - "attention_bias": True, - "attention_out_bias": True, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_offset": 1, - "timestep_scale_multiplier": 1000, - "cross_attn_timestep_scale_multiplier": 1000, - "rope_type": "split", - "use_prompt_embeddings": True, - "perturbed_attn": False, - }, - } - rename_dict = LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP - elif version == "2.3": - config = { - "model_id": "Lightricks/LTX-2.3", - "diffusers_config": { - "in_channels": 128, - "out_channels": 128, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 32, - "attention_head_dim": 128, - "cross_attention_dim": 4096, - "vae_scale_factors": (8, 32, 32), - "pos_embed_max_pos": 20, - "base_height": 2048, - "base_width": 2048, - "gated_attn": True, - "cross_attn_mod": True, - "audio_in_channels": 128, - "audio_out_channels": 128, - "audio_patch_size": 1, - "audio_patch_size_t": 1, - "audio_num_attention_heads": 32, - "audio_attention_head_dim": 64, - "audio_cross_attention_dim": 2048, - "audio_scale_factor": 4, - "audio_pos_embed_max_pos": 20, - "audio_sampling_rate": 16000, - "audio_hop_length": 160, - "audio_gated_attn": True, - "audio_cross_attn_mod": True, - "num_layers": 48, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 3840, - "attention_bias": True, - "attention_out_bias": True, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_offset": 1, - "timestep_scale_multiplier": 1000, - "cross_attn_timestep_scale_multiplier": 1000, - "rope_type": "split", - "use_prompt_embeddings": False, - "perturbed_attn": True, - }, - } - rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP - elif version == "2.5": - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - "in_channels": 128, - "out_channels": 128, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 32, - "attention_head_dim": 128, - "cross_attention_dim": 4096, - "vae_scale_factors": (8, 32, 32), - "pos_embed_max_pos": 20, - "base_height": 2048, - "base_width": 2048, - "gated_attn": True, - "cross_attn_mod": True, - "audio_in_channels": 128, - "audio_out_channels": 128, - "audio_patch_size": 1, - "audio_patch_size_t": 1, - "audio_num_attention_heads": 32, - "audio_attention_head_dim": 64, - "audio_cross_attention_dim": 2048, - "audio_scale_factor": 4, - "audio_pos_embed_max_pos": 20, - "audio_sampling_rate": 16000, - "audio_hop_length": 160, - "audio_gated_attn": True, - "audio_cross_attn_mod": True, - "num_layers": 48, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 3840, - "attention_bias": True, - "attention_out_bias": True, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_offset": 1, - "timestep_scale_multiplier": 1000, - "cross_attn_timestep_scale_multiplier": 1000, - "rope_type": "split", - "use_prompt_embeddings": False, - "perturbed_attn": True, - # The only transformer-level deltas from 2.3: the video FFN drops its bias (audio_ff_bias and - # use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5 carries a - # learned keyframe absolute-position embedding. - "ff_bias": False, - "use_keyframes_abs_pos_embedding": True, - }, - } - rename_dict = LTX_2_3_TRANSFORMER_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP - return config, rename_dict, special_keys_remap - - -def get_ltx2_connectors_config( - version: str, gemma_text_config: Any | None = None -) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version == "test": - config = { - "model_id": "diffusers-internal-dev/dummy-ltx2", - "diffusers_config": { - "caption_channels": 16, - "text_proj_in_factor": 3, - "video_connector_num_attention_heads": 4, - "video_connector_attention_head_dim": 8, - "video_connector_num_layers": 1, - "video_connector_num_learnable_registers": None, - "audio_connector_num_attention_heads": 4, - "audio_connector_attention_head_dim": 8, - "audio_connector_num_layers": 1, - "audio_connector_num_learnable_registers": None, - "connector_rope_base_seq_len": 32, - "rope_theta": 10000.0, - "rope_double_precision": False, - "causal_temporal_positioning": False, - }, - } - elif version == "2.0": - config = { - "model_id": "Lightricks/LTX-2", - "diffusers_config": { - "caption_channels": 3840, - "text_proj_in_factor": 49, - "video_connector_num_attention_heads": 30, - "video_connector_attention_head_dim": 128, - "video_connector_num_layers": 2, - "video_connector_num_learnable_registers": 128, - "video_gated_attn": False, - "audio_connector_num_attention_heads": 30, - "audio_connector_attention_head_dim": 128, - "audio_connector_num_layers": 2, - "audio_connector_num_learnable_registers": 128, - "audio_gated_attn": False, - "connector_rope_base_seq_len": 4096, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_temporal_positioning": False, - "rope_type": "split", - "per_modality_projections": False, - "proj_bias": False, - }, - } - rename_dict = LTX_2_0_CONNECTORS_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP - elif version == "2.3": - config = { - "model_id": "Lightricks/LTX-2.3", - "diffusers_config": { - "caption_channels": 3840, - "text_proj_in_factor": 49, - "video_connector_num_attention_heads": 32, - "video_connector_attention_head_dim": 128, - "video_connector_num_layers": 8, - "video_connector_num_learnable_registers": 128, - "video_gated_attn": True, - "audio_connector_num_attention_heads": 32, - "audio_connector_attention_head_dim": 64, - "audio_connector_num_layers": 8, - "audio_connector_num_learnable_registers": 128, - "audio_gated_attn": True, - "connector_rope_base_seq_len": 4096, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_temporal_positioning": False, - "rope_type": "split", - "per_modality_projections": True, - "video_hidden_dim": 4096, - "audio_hidden_dim": 2048, - "proj_bias": True, - }, - } - rename_dict = LTX_2_3_CONNECTORS_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP - elif version == "2.5": - if gemma_text_config is None: - raise ValueError("gemma_text_config is required to derive connector dims for LTX-2.5.") - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - # Derived from the Gemma 4 text config rather than hardcoded, since (unlike Gemma-3-12B) the - # 2.5 text encoder isn't a single fixed checkpoint. Formula matches the reference - # (`encoder_configurator._create_feature_extractor`): hidden_size, and num_hidden_layers + 1 - # for the embedding layer. - "caption_channels": gemma_text_config.hidden_size, - "text_proj_in_factor": gemma_text_config.num_hidden_layers + 1, - "video_connector_num_attention_heads": 32, - "video_connector_attention_head_dim": 128, - "video_connector_num_layers": 8, - "video_connector_num_learnable_registers": 128, - "video_gated_attn": True, - "audio_connector_num_attention_heads": 32, - "audio_connector_attention_head_dim": 64, - "audio_connector_num_layers": 8, - "audio_connector_num_learnable_registers": 128, - "audio_gated_attn": True, - "connector_rope_base_seq_len": 4096, - "rope_theta": 10000.0, - "rope_double_precision": True, - "causal_temporal_positioning": False, - "rope_type": "split", - "per_modality_projections": True, - "video_hidden_dim": 4096, - "audio_hidden_dim": 2048, - "proj_bias": True, - }, - } - rename_dict = LTX_2_3_CONNECTORS_KEYS_RENAME_DICT - special_keys_remap = LTX_2_0_CONNECTORS_SPECIAL_KEYS_REMAP - - return config, rename_dict, special_keys_remap - - -def convert_ltx2_transformer(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: - config, rename_dict, special_keys_remap = get_ltx2_transformer_config(version) - diffusers_config = config["diffusers_config"] - - transformer_state_dict, _ = split_transformer_and_connector_state_dict(original_state_dict) - - with init_empty_weights(): - transformer = LTX2VideoTransformer3DModel.from_config(diffusers_config) - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(transformer_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(transformer_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(transformer_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, transformer_state_dict) - - transformer.load_state_dict(transformer_state_dict, strict=True, assign=True) - return transformer - - -def convert_ltx2_connectors( - original_state_dict: dict[str, Any], version: str, gemma_text_config: Any | None = None -) -> LTX2TextConnectors: - config, rename_dict, special_keys_remap = get_ltx2_connectors_config(version, gemma_text_config=gemma_text_config) - diffusers_config = config["diffusers_config"] - - _, connector_state_dict = split_transformer_and_connector_state_dict(original_state_dict) - if len(connector_state_dict) == 0: - raise ValueError("No connector weights found in the provided state dict.") - - with init_empty_weights(): - connectors = LTX2TextConnectors.from_config(diffusers_config) - - for key in list(connector_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(connector_state_dict, key, new_key) - - for key in list(connector_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, connector_state_dict) - - connectors.load_state_dict(connector_state_dict, strict=True, assign=True) - return connectors - - -def convert_ltx2_duration_head(original_state_dict: dict[str, Any]) -> LTX2DurationHead | None: - """Builds an `LTX2DurationHead` from a duration-head state dict, or `None` if the checkpoint has none. - - The duration head ships from LTX-2.5 onward. Its hyperparameters are absent from checkpoint metadata, so the - dimensions are read back from the weight shapes; `num_pooler_heads` cannot be recovered that way and is fixed at - the value the head was trained with. - - The original checkpoint stores the pooler's projections fused, in `torch.nn.MultiheadAttention` layout. They are - split here into the separate q/k/v projections `LTX2DurationAttentionPooler` uses. - """ - state_dict = dict(original_state_dict) - if len(state_dict) == 0: - print("No duration_head weights found in the checkpoint; skipping (expected for pre-2.5 checkpoints).") - return None - - in_proj_weight = state_dict.pop("attention_pooler.cross_attn.in_proj_weight") - in_proj_bias = state_dict.pop("attention_pooler.cross_attn.in_proj_bias") - query_weight, key_weight, value_weight = in_proj_weight.chunk(3, dim=0) - query_bias, key_bias, value_bias = in_proj_bias.chunk(3, dim=0) - state_dict["attention_pooler.to_q.weight"] = query_weight - state_dict["attention_pooler.to_k.weight"] = key_weight - state_dict["attention_pooler.to_v.weight"] = value_weight - state_dict["attention_pooler.to_q.bias"] = query_bias - state_dict["attention_pooler.to_k.bias"] = key_bias - state_dict["attention_pooler.to_v.bias"] = value_bias - state_dict["attention_pooler.to_out.weight"] = state_dict.pop("attention_pooler.cross_attn.out_proj.weight") - state_dict["attention_pooler.to_out.bias"] = state_dict.pop("attention_pooler.cross_attn.out_proj.bias") - - diffusers_config = { - "video_cross_attention_dim": state_dict["video_input_proj.weight"].shape[1], - "audio_cross_attention_dim": state_dict["audio_input_proj.weight"].shape[1], - "pooler_hidden_dim": state_dict["attention_pooler.to_q.weight"].shape[1], - "num_queries": state_dict["attention_pooler.query_tokens"].shape[0], - "mlp_hidden_dim": state_dict["mlp_hidden.weight"].shape[0], - "num_pooler_heads": 4, - } - - with init_empty_weights(): - duration_head = LTX2DurationHead.from_config(diffusers_config) - - duration_head.load_state_dict(state_dict, strict=True, assign=True) - return duration_head - - -def get_ltx2_video_vae_config( - version: str, timestep_conditioning: bool = False -) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version == "test": - config = { - "model_id": "diffusers-internal-dev/dummy-ltx2", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (256, 512, 1024, 2048), - "down_block_types": ( - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 1024), - "layers_per_block": (4, 6, 6, 2, 2), - "decoder_layers_per_block": (5, 5, 5, 5), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True), - "decoder_inject_noise": (False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_residual": (True, True, True), - "upsample_factor": (2, 2, 2), - "timestep_conditioning": timestep_conditioning, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "encoder_causal": True, - "decoder_causal": False, - "encoder_spatial_padding_mode": "zeros", - "decoder_spatial_padding_mode": "reflect", - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - }, - } - rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP - elif version == "2.0": - config = { - "model_id": "Lightricks/LTX-2", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (256, 512, 1024, 2048), - "down_block_types": ( - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 1024), - "layers_per_block": (4, 6, 6, 2, 2), - "decoder_layers_per_block": (5, 5, 5, 5), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True), - "decoder_inject_noise": (False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_type": ("spatiotemporal", "spatiotemporal", "spatiotemporal"), - "upsample_residual": (True, True, True), - "upsample_factor": (2, 2, 2), - "timestep_conditioning": timestep_conditioning, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "encoder_causal": True, - "decoder_causal": False, - "encoder_spatial_padding_mode": "zeros", - "decoder_spatial_padding_mode": "reflect", - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - }, - } - rename_dict = LTX_2_0_VIDEO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP - elif version == "2.3": - config = { - "model_id": "Lightricks/LTX-2.3", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (256, 512, 1024, 1024), - "down_block_types": ( - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 512, 1024), - "layers_per_block": (4, 6, 4, 2, 2), - "decoder_layers_per_block": (4, 6, 4, 2, 2), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True, True), - "decoder_inject_noise": (False, False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"), - "upsample_residual": (False, False, False, False), - "upsample_factor": (2, 2, 1, 2), - "timestep_conditioning": timestep_conditioning, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "encoder_causal": True, - "decoder_causal": False, - "encoder_spatial_padding_mode": "zeros", - "decoder_spatial_padding_mode": "zeros", - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - }, - } - rename_dict = LTX_2_3_VIDEO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP - elif version == "2.5": - # Same block structure as 2.3 (32x32x8 compression); confirmed against the checkpoint's - # config["vae"]["encoder_blocks"]/["decoder_blocks"] metadata, which is byte-identical to 2.3's. - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (256, 512, 1024, 1024), - "down_block_types": ( - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - "LTX2VideoDownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 512, 1024), - "layers_per_block": (4, 6, 4, 2, 2), - "decoder_layers_per_block": (4, 6, 4, 2, 2), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True, True), - "decoder_inject_noise": (False, False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"), - "upsample_residual": (False, False, False, False), - "upsample_factor": (2, 2, 1, 2), - "timestep_conditioning": timestep_conditioning, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "encoder_causal": True, - "decoder_causal": False, - "encoder_spatial_padding_mode": "zeros", - "decoder_spatial_padding_mode": "zeros", - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - }, - } - rename_dict = LTX_2_3_VIDEO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_VAE_SPECIAL_KEYS_REMAP - return config, rename_dict, special_keys_remap - - -def convert_ltx2_video_vae( - original_state_dict: dict[str, Any], version: str, timestep_conditioning: bool -) -> dict[str, Any]: - config, rename_dict, special_keys_remap = get_ltx2_video_vae_config(version, timestep_conditioning) - diffusers_config = config["diffusers_config"] - - with init_empty_weights(): - vae = AutoencoderKLLTX2Video.from_config(diffusers_config) - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def get_ltx2_diffusion_video_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version != "2.5": - raise ValueError( - f"The diffusion decoder was introduced in LTX-2.5, which the converter handles under " - f"`--version 2.5`; got version {version!r}." - ) - # The encoder half is 2.5's conv VAE encoder, unchanged (its weights are byte identical to 2.3's), so - # those entries must stay in sync with `get_ltx2_video_vae_config("2.5")`. - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - "out_channels": 3, - "latent_channels": 128, - "patch_size": 4, - "decoder_head_dim": 64, - "decoder_stage_channels": (2048, 1024, 512, 512, 256), - "decoder_stage_depths": (4, 6, 4, 2, 8), - "decoder_stage_kernels": ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)), - "decoder_upsample_strides": ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)), - "decoder_upsample_channel_reductions": (2, 2, 1, 2), - "decoder_stage5_kernel": (11, 11, 11), - "decoder_t_emb_dim": 384, - "decoder_timestep_scale_multiplier": 1000.0, - "decoder_model_output_type": "x0", - "decoder_num_inference_steps": 1, - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - }, - } - return config, LTX_2_3_VIDEO_VAE_RENAME_DICT, LTX_2_0_VAE_SPECIAL_KEYS_REMAP - - -def convert_ltx2_diffusion_video_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: - config, rename_dict, special_keys_remap = get_ltx2_diffusion_video_vae_config(version) - diffusers_config = config["diffusers_config"] - - with init_empty_weights(): - vae = LTX2VideoDiffusionDecoderModel.from_config(diffusers_config) - - # The checkpoint is a whole VAE, but this model is decoder-only: encoding stays with - # `AutoencoderKLLTX2Video`. So keep the `decoder.` half and the per-channel statistics (which become - # `latents_mean` / `latents_std`), and drop the encoder outright rather than remapping weights that - # have nowhere to land. `load_state_dict` below is strict, so anything left over would raise. - decoder_state_dict = { - key.removeprefix("decoder."): value for key, value in original_state_dict.items() if key.startswith("decoder.") - } - for key in list(original_state_dict.keys()): - if key.startswith("decoder.") or not key.startswith("per_channel_statistics"): - del original_state_dict[key] - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - gates = { - key: value - for key, value in decoder_state_dict.items() - if key.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")) - } - for key, value in decoder_state_dict.items(): - # Bundled preview heads are not part of the decoder. The sft (dev) checkpoint carries neither these - # nor any gate, so both this and the fold below no-op on it; gated distilled checkpoints need them. - if key.startswith("coarse_") or ".coarse_" in key or key in gates: - continue - - new_key = key[:] - for replace_key, rename_key in LTX_2_5_DIFFUSION_DECODER_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - gated_leaf = next( - (leaf for leaf in LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS if new_key.endswith(leaf)), None - ) - if gated_leaf is not None: - # The gate is a sibling of the Linear it gates, so it shares the block prefix. - gate = gates.get(new_key[: -len(gated_leaf)] + LTX_2_5_DIFFUSION_DECODER_GATE_FOLD_TARGETS[gated_leaf]) - if gate is not None: - gate = gate.to(torch.float32) - folded = (gate.unsqueeze(1) if value.ndim == 2 else gate) * value.to(torch.float32) - value = folded.to(value.dtype) - - # The checkpoint stores one fused `Linear(dim, 3 * dim)`; the model owns three separate projections. - if new_key.endswith((".qkv.weight", ".qkv.bias")): - leaf = "weight" if new_key.endswith(".weight") else "bias" - prefix = new_key[: -len(f"qkv.{leaf}")] - if value.shape[0] % 3 != 0: - raise ValueError(f"Fused QKV param {key!r} leading dim {value.shape[0]} is not divisible by 3.") - chunk = value.shape[0] // 3 - original_state_dict[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone() - original_state_dict[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone() - original_state_dict[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone() - continue - - original_state_dict[f"decoder.{new_key}"] = value - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version == "2.0": - config = { - "model_id": "Lightricks/LTX-2", - "diffusers_config": { - "base_channels": 128, - "output_channels": 2, - "ch_mult": (1, 2, 4), - "num_res_blocks": 2, - "attn_resolutions": None, - "in_channels": 2, - "resolution": 256, - "latent_channels": 8, - "norm_type": "pixel", - "causality_axis": "height", - "dropout": 0.0, - "mid_block_add_attention": False, - "sample_rate": 16000, - "mel_hop_length": 160, - "is_causal": True, - "mel_bins": 64, - "double_z": True, - }, - } - rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP - elif version == "2.3": - config = { - "model_id": "Lightricks/LTX-2.3", - "diffusers_config": { - "base_channels": 128, - "output_channels": 2, - "ch_mult": (1, 2, 4), - "num_res_blocks": 2, - "attn_resolutions": None, - "in_channels": 2, - "resolution": 256, - "latent_channels": 8, - "norm_type": "pixel", - "causality_axis": "height", - "dropout": 0.0, - "mid_block_add_attention": False, - "sample_rate": 16000, - "mel_hop_length": 160, - "is_causal": True, - "mel_bins": 64, - "double_z": True, - }, # Same config as LTX-2.0 - } - rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP - elif version == "2.5": - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - "base_channels": 128, - "output_channels": 2, - "ch_mult": (1, 2, 4), - "num_res_blocks": 2, - "attn_resolutions": None, - "in_channels": 2, - "resolution": 256, - "latent_channels": 8, - "norm_type": "pixel", - "causality_axis": "height", - "dropout": 0.0, - "mid_block_add_attention": False, - "sample_rate": 16000, - "mel_hop_length": 160, - "is_causal": True, - "mel_bins": 64, - "double_z": True, - }, # Same config as LTX-2.0 / 2.3 - } - rename_dict = LTX_2_0_AUDIO_VAE_RENAME_DICT - special_keys_remap = LTX_2_0_AUDIO_VAE_SPECIAL_KEYS_REMAP - return config, rename_dict, special_keys_remap - - -def convert_ltx2_audio_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: - config, rename_dict, special_keys_remap = get_ltx2_audio_vae_config(version) - diffusers_config = config["diffusers_config"] - - with init_empty_weights(): - vae = AutoencoderKLLTX2Audio.from_config(diffusers_config) - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: - if version == "2.0": - config = { - "model_id": "Lightricks/LTX-2", - "diffusers_config": { - "in_channels": 128, - "hidden_channels": 1024, - "out_channels": 2, - "upsample_kernel_sizes": [16, 15, 8, 4, 4], - "upsample_factors": [6, 5, 2, 2, 2], - "resnet_kernel_sizes": [3, 7, 11], - "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "act_fn": "leaky_relu", - "leaky_relu_negative_slope": 0.1, - "antialias": False, - "final_act_fn": "tanh", - "final_bias": True, - "output_sampling_rate": 24000, - }, - } - rename_dict = LTX_2_0_VOCODER_RENAME_DICT - special_keys_remap = LTX_2_0_VOCODER_SPECIAL_KEYS_REMAP - elif version == "2.3": - config = { - "model_id": "Lightricks/LTX-2.3", - "diffusers_config": { - "in_channels": 128, - "hidden_channels": 1536, - "out_channels": 2, - "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4], - "upsample_factors": [5, 2, 2, 2, 2, 2], - "resnet_kernel_sizes": [3, 7, 11], - "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "act_fn": "snakebeta", - "leaky_relu_negative_slope": 0.1, - "antialias": True, - "antialias_ratio": 2, - "antialias_kernel_size": 12, - "final_act_fn": None, - "final_bias": False, - "bwe_in_channels": 128, - "bwe_hidden_channels": 512, - "bwe_out_channels": 2, - "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4], - "bwe_upsample_factors": [6, 5, 2, 2, 2], - "bwe_resnet_kernel_sizes": [3, 7, 11], - "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "bwe_act_fn": "snakebeta", - "bwe_leaky_relu_negative_slope": 0.1, - "bwe_antialias": True, - "bwe_antialias_ratio": 2, - "bwe_antialias_kernel_size": 12, - "bwe_final_act_fn": None, - "bwe_final_bias": False, - "filter_length": 512, - "hop_length": 80, - "window_length": 512, - "num_mel_channels": 64, - "input_sampling_rate": 16000, - "output_sampling_rate": 48000, - }, - } - rename_dict = LTX_2_3_VOCODER_RENAME_DICT - special_keys_remap = LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP - elif version == "2.5": - config = { - "model_id": "Lightricks/LTX-2.5", - "diffusers_config": { - "in_channels": 128, - "hidden_channels": 1536, - "out_channels": 2, - "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4], - "upsample_factors": [5, 2, 2, 2, 2, 2], - "resnet_kernel_sizes": [3, 7, 11], - "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "act_fn": "snakebeta", - "leaky_relu_negative_slope": 0.1, - "antialias": True, - "antialias_ratio": 2, - "antialias_kernel_size": 12, - "final_act_fn": None, - "final_bias": False, - "bwe_in_channels": 128, - "bwe_hidden_channels": 512, - "bwe_out_channels": 2, - "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4], - "bwe_upsample_factors": [6, 5, 2, 2, 2], - "bwe_resnet_kernel_sizes": [3, 7, 11], - "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "bwe_act_fn": "snakebeta", - "bwe_leaky_relu_negative_slope": 0.1, - "bwe_antialias": True, - "bwe_antialias_ratio": 2, - "bwe_antialias_kernel_size": 12, - "bwe_final_act_fn": None, - "bwe_final_bias": False, - "filter_length": 512, - "hop_length": 80, - "window_length": 512, - "num_mel_channels": 64, - "input_sampling_rate": 16000, - "output_sampling_rate": 48000, - }, # Same config as LTX-2.3 - } - rename_dict = LTX_2_3_VOCODER_RENAME_DICT - special_keys_remap = LTX_2_3_VOCODER_SPECIAL_KEYS_REMAP - return config, rename_dict, special_keys_remap - - -def convert_ltx2_vocoder(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]: - config, rename_dict, special_keys_remap = get_ltx2_vocoder_config(version) - diffusers_config = config["diffusers_config"] - if version in ("2.3", "2.5"): - vocoder_cls = LTX2VocoderWithBWE - else: - vocoder_cls = LTX2Vocoder - - with init_empty_weights(): - vocoder = vocoder_cls.from_config(diffusers_config) - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in rename_dict.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in special_keys_remap.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vocoder.load_state_dict(original_state_dict, strict=True, assign=True) - return vocoder - - -def get_ltx2_spatial_latent_upsampler_config(version: str): - if version == "2.0": - config = { - "in_channels": 128, - "mid_channels": 1024, - "num_blocks_per_stage": 4, - "dims": 3, - "spatial_upsample": True, - "temporal_upsample": False, - "rational_spatial_scale": 2.0, - "use_rational_resampler": True, - } - elif version in ("2.3", "2.5"): - config = { - "in_channels": 128, - "mid_channels": 1024, - "num_blocks_per_stage": 4, - "dims": 3, - "spatial_upsample": True, - "temporal_upsample": False, - "rational_spatial_scale": 2.0, - "use_rational_resampler": False, - } - else: - raise ValueError(f"Unsupported version: {version}") - return config - - -def get_ltx2_temporal_latent_upsampler_config(version: str): - if version != "2.5": - raise ValueError(f"Unsupported version: {version}") - # The temporal x2 upsampler is narrower than its spatial sibling and pixel-shuffles along time only. - return { - "in_channels": 128, - "mid_channels": 512, - "num_blocks_per_stage": 4, - "dims": 3, - "spatial_upsample": False, - "temporal_upsample": True, - } - - -def convert_ltx2_latent_upsampler(original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype): - with init_empty_weights(): - latent_upsampler = LTX2LatentUpsamplerModel(**config) - - latent_upsampler.load_state_dict(original_state_dict, strict=True, assign=True) - latent_upsampler.to(dtype) - return latent_upsampler - - -def load_original_checkpoint(args, filename: str | None) -> dict[str, Any]: - if args.original_state_dict_repo_id is not None: - ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename) - elif args.checkpoint_path is not None: - ckpt_path = args.checkpoint_path - else: - raise ValueError("Please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") - - original_state_dict = safetensors.torch.load_file(ckpt_path) - return original_state_dict - - -def load_hub_or_local_checkpoint(repo_id: str | None = None, filename: str | None = None) -> dict[str, Any]: - if repo_id is None and filename is None: - raise ValueError("Please supply at least one of `repo_id` or `filename`") - - if repo_id is not None: - if filename is None: - raise ValueError("If repo_id is specified, filename must also be specified.") - ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename) - else: - ckpt_path = filename - - _, ext = os.path.splitext(ckpt_path) - if ext in [".safetensors", ".sft"]: - state_dict = safetensors.torch.load_file(ckpt_path) - else: - state_dict = torch.load(ckpt_path, map_location="cpu") - - return state_dict - - -def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefix: str) -> dict[str, Any]: - # Ensure that the key prefix ends with a dot (.) - if not prefix.endswith("."): - prefix = prefix + "." - - model_state_dict = {} - for param_name, param in combined_ckpt.items(): - if param_name.startswith(prefix): - model_state_dict[param_name.removeprefix(prefix)] = param - - if prefix == "model.diffusion_model.": - # Some checkpoints store the text connector projection outside the diffusion model prefix. - connector_prefixes = ["text_embedding_projection"] - for param_name, param in combined_ckpt.items(): - for prefix in connector_prefixes: - if param_name.startswith(prefix): - # Check to make sure we're not overwriting an existing key - if param_name not in model_state_dict: - model_state_dict[param_name] = combined_ckpt[param_name] - - return model_state_dict - - -def get_args(): - parser = argparse.ArgumentParser() - - def none_or_str(value: str): - if isinstance(value, str) and value.lower() == "none": - return None - return value - - parser.add_argument( - "--original_state_dict_repo_id", - default="Lightricks/LTX-2", - type=none_or_str, - help="HF Hub repo id with LTX 2.0 checkpoint", - ) - parser.add_argument( - "--checkpoint_path", - default=None, - type=str, - help="Local checkpoint path for LTX 2.0. Will be used if `original_state_dict_repo_id` is not specified.", - ) - parser.add_argument( - "--version", - type=str, - default="2.0", - choices=["test", "2.0", "2.3", "2.5"], - help="Version of the LTX 2.0 model", - ) - - parser.add_argument( - "--combined_filename", - default="ltx-2-19b-dev.safetensors", - type=none_or_str, - help="Filename for combined checkpoint with all LTX 2.0 models (VAE, DiT, etc.)", - ) - parser.add_argument("--vae_prefix", default="vae.", type=str) - parser.add_argument("--audio_vae_prefix", default="audio_vae.", type=str) - parser.add_argument("--dit_prefix", default="model.diffusion_model.", type=str) - parser.add_argument("--vocoder_prefix", default="vocoder.", type=str) - parser.add_argument("--duration_head_prefix", default="duration_head.", type=str) - - parser.add_argument("--vae_filename", default=None, type=str, help="VAE filename; overrides combined ckpt if set") - parser.add_argument( - "--audio_vae_filename", default=None, type=str, help="Audio VAE filename; overrides combined ckpt if set" - ) - parser.add_argument("--dit_filename", default=None, type=str, help="DiT filename; overrides combined ckpt if set") - parser.add_argument( - "--vocoder_filename", default=None, type=str, help="Vocoder filename; overrides combined ckpt if set" - ) - parser.add_argument( - "--text_encoder_model_id", - default="google/gemma-3-12b-it-qat-q4_0-unquantized", - type=none_or_str, - help=( - "HF Hub id for the text encoder model. Default is Gemma 3, used by LTX 2.0/2.3. LTX-2.5 requires a " - "Gemma 4 (`gemma4_unified`) checkpoint here instead -- passing the Gemma 3 default with `--version 2.5` " - "raises an error." - ), - ) - parser.add_argument( - "--tokenizer_id", - default="google/gemma-3-12b-it-qat-q4_0-unquantized", - type=none_or_str, - help="HF Hub id for the text tokenizer. Should match --text_encoder_model_id's family (Gemma 3 vs Gemma 4).", - ) - parser.add_argument( - "--prompt_enhancer_model_id", - default=None, - type=none_or_str, - help=( - "HF Hub id for the prompt-enhancer model (used with --add_processor). For LTX-2.0/2.3, defaults to " - "--text_encoder_model_id (the same Gemma 3 checkpoint serves both roles). LTX-2.5's fine-tuned text " - "encoder is not trained for enhancement, so this must be set explicitly for --version 2.5 -- e.g. to " - "google/gemma-4-E2B-it or google/gemma-4-E4B-it." - ), - ) - parser.add_argument( - "--temporal_latent_upsampler_filename", - default="ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors", - type=none_or_str, - help="Temporal x2 latent upsampler filename (LTX-2.5, used by the DFR pipeline's temporal refine rounds)", - ) - parser.add_argument( - "--latent_upsampler_filename", - default="ltx-2-spatial-upscaler-x2-1.0.safetensors", - type=none_or_str, - help="Latent upsampler filename", - ) - - parser.add_argument( - "--timestep_conditioning", action="store_true", help="Whether to add timestep condition to the video VAE model" - ) - parser.add_argument("--vae", action="store_true", help="Whether to convert the video VAE model") - parser.add_argument( - "--diffusion_vae", - action="store_true", - help=( - "Whether to convert the LTX-2.5 diffusion decoder, saved to a `diffusion_decoder` subfolder — the " - "component name `LTX2VideoDiffusionDecodePipeline` and the modular blocks resolve it by — so " - "`from_pretrained` keeps returning the conv decoder in `vae` by default" - ), - ) - parser.add_argument("--audio_vae", action="store_true", help="Whether to convert the audio VAE model") - parser.add_argument("--dit", action="store_true", help="Whether to convert the DiT model") - parser.add_argument("--connectors", action="store_true", help="Whether to convert the connector model") - parser.add_argument( - "--duration_head", - action="store_true", - help="Whether to convert the duration head (present in LTX-2.5 and later checkpoints only)", - ) - parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model") - parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder") - parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler") - parser.add_argument( - "--temporal_latent_upsampler", - action="store_true", - help="Whether to convert the temporal x2 latent upsampler (LTX-2.5)", - ) - parser.add_argument( - "--full_pipeline", - action="store_true", - help="Whether to save the pipeline. This will attempt to convert all models (e.g. vae, dit, etc.)", - ) - parser.add_argument( - "--upsample_pipeline", - action="store_true", - help="Whether to save a latent upsampling pipeline", - ) - parser.add_argument( - "--add_processor", - action="store_true", - help="Whether to add a text-encoder processor to the pipeline for prompt enhancement.", - ) - - parser.add_argument("--vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--audio_vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--dit_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--vocoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) - parser.add_argument("--text_encoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) - - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") - parser.add_argument( - "--upsample_output_path", - type=str, - default=None, - help="Path where converted upsampling pipeline should be saved", - ) - - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - -VARIANT_MAPPING = { - "fp32": None, - "fp16": "fp16", - "bf16": "bf16", -} - - -def main(args): - vae_dtype = DTYPE_MAPPING[args.vae_dtype] - audio_vae_dtype = DTYPE_MAPPING[args.audio_vae_dtype] - dit_dtype = DTYPE_MAPPING[args.dit_dtype] - vocoder_dtype = DTYPE_MAPPING[args.vocoder_dtype] - text_encoder_dtype = DTYPE_MAPPING[args.text_encoder_dtype] - - combined_ckpt = None - load_combined_models = any( - [ - args.vae, - args.diffusion_vae, - args.audio_vae, - args.dit, - args.vocoder, - args.connectors, - args.full_pipeline, - args.upsample_pipeline, - ] - ) - if args.combined_filename is not None and load_combined_models: - combined_ckpt = load_original_checkpoint(args, filename=args.combined_filename) - - # LTX-2.5 only works with a Gemma 4 (`gemma4_unified`) text encoder; --text_encoder_model_id defaults to - # Gemma 3 (for 2.0/2.3), so silently proceeding would pair a 2.5 checkpoint with the wrong text encoder. - gemma_text_config = None - if args.version == "2.5" and (args.text_encoder or args.connectors or args.full_pipeline): - gemma_config = AutoConfig.from_pretrained(args.text_encoder_model_id) - if gemma_config.model_type != "gemma4_unified": - raise ValueError( - f"LTX-2.5 requires a Gemma 4 (`gemma4_unified`) text encoder, but --text_encoder_model_id=" - f"{args.text_encoder_model_id!r} has model_type={gemma_config.model_type!r}. Pass " - "--text_encoder_model_id pointing at a Gemma 4 checkpoint (the default is Gemma 3, for 2.0/2.3)." - ) - gemma_text_config = gemma_config.text_config - - # LTX-2.5's fine-tuned text encoder is never a valid prompt enhancer (unlike LTX-2.0/2.3, where the same Gemma 3 - # checkpoint serves both roles) -- require an explicit, separate --prompt_enhancer_model_id instead of silently - # falling back to --text_encoder_model_id. - if ( - args.version == "2.5" - and args.add_processor - and (args.text_encoder or args.full_pipeline) - and args.prompt_enhancer_model_id is None - ): - raise ValueError( - "LTX-2.5's text encoder is not trained for prompt enhancement, so --prompt_enhancer_model_id must be " - "set explicitly when --add_processor is used with --version 2.5 -- e.g. to google/gemma-4-E2B-it or " - "google/gemma-4-E4B-it." - ) - - if args.vae or args.full_pipeline or args.upsample_pipeline: - if args.vae_filename is not None: - original_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename) - elif combined_ckpt is not None: - original_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix) - vae = convert_ltx2_video_vae( - original_vae_ckpt, version=args.version, timestep_conditioning=args.timestep_conditioning - ) - if not args.full_pipeline and not args.upsample_pipeline: - vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "vae")) - - if args.diffusion_vae: - if args.vae_filename is not None: - original_diffusion_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename) - elif combined_ckpt is not None: - original_diffusion_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix) - diffusion_vae = convert_ltx2_diffusion_video_vae(original_diffusion_vae_ckpt, version=args.version) - # "diffusion_decoder", not "vae_diffusion": pipeline-level `from_pretrained` resolves each component - # from the subfolder named after it, so this folder name must match the `diffusion_decoder` component - # of `LTX2VideoDiffusionDecodePipeline` (and the modular `ComponentSpec`) for those loads to work. - diffusion_vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "diffusion_decoder")) - - if args.audio_vae or args.full_pipeline: - if args.audio_vae_filename is not None: - original_audio_vae_ckpt = load_hub_or_local_checkpoint(filename=args.audio_vae_filename) - elif combined_ckpt is not None: - original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.audio_vae_prefix) - audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=args.version) - if not args.full_pipeline: - audio_vae.to(audio_vae_dtype).save_pretrained(os.path.join(args.output_path, "audio_vae")) - - if args.dit or args.full_pipeline: - if args.dit_filename is not None: - original_dit_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) - elif combined_ckpt is not None: - original_dit_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) - transformer = convert_ltx2_transformer(original_dit_ckpt, version=args.version) - if not args.full_pipeline: - transformer.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "transformer")) - - if args.connectors or args.full_pipeline: - if args.dit_filename is not None: - original_connectors_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) - elif combined_ckpt is not None: - original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) - connectors = convert_ltx2_connectors( - original_connectors_ckpt, version=args.version, gemma_text_config=gemma_text_config - ) - if not args.full_pipeline: - connectors.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "connectors")) - - duration_head = None - if args.duration_head or args.full_pipeline: - if combined_ckpt is not None: - original_duration_head_ckpt = get_model_state_dict_from_combined_ckpt( - combined_ckpt, args.duration_head_prefix - ) - duration_head = convert_ltx2_duration_head(original_duration_head_ckpt) - if duration_head is not None and not args.full_pipeline: - duration_head.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "duration_head")) - - if args.vocoder or args.full_pipeline: - if args.vocoder_filename is not None: - original_vocoder_ckpt = load_hub_or_local_checkpoint(filename=args.vocoder_filename) - elif combined_ckpt is not None: - original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vocoder_prefix) - vocoder = convert_ltx2_vocoder(original_vocoder_ckpt, version=args.version) - if not args.full_pipeline: - vocoder.to(vocoder_dtype).save_pretrained(os.path.join(args.output_path, "vocoder")) - - if args.text_encoder or args.full_pipeline: - text_encoder = AutoModelForImageTextToText.from_pretrained(args.text_encoder_model_id) - if not args.full_pipeline: - text_encoder.to(text_encoder_dtype).save_pretrained(os.path.join(args.output_path, "text_encoder")) - - tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) - if not args.full_pipeline: - tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer")) - - processor = None - prompt_enhancer = None - if args.add_processor: - enhancer_model_id = args.prompt_enhancer_model_id or args.text_encoder_model_id - processor = AutoProcessor.from_pretrained(enhancer_model_id) - if not args.full_pipeline: - processor.save_pretrained(os.path.join(args.output_path, "processor")) - - if args.prompt_enhancer_model_id is not None: - # Separate, dedicated enhancer model (required for LTX-2.5); for LTX-2.0/2.3, the same - # `text_encoder` checkpoint already serves as its own enhancer, so nothing extra is saved. - prompt_enhancer = AutoModelForImageTextToText.from_pretrained(enhancer_model_id) - if not args.full_pipeline: - prompt_enhancer.to(text_encoder_dtype).save_pretrained( - os.path.join(args.output_path, "prompt_enhancer") - ) - - if args.latent_upsampler or args.upsample_pipeline: - original_latent_upsampler_ckpt = load_hub_or_local_checkpoint( - repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename - ) - latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version) - latent_upsampler = convert_ltx2_latent_upsampler( - original_latent_upsampler_ckpt, - latent_upsampler_config, - dtype=vae_dtype, - ) - if not args.full_pipeline and not args.upsample_pipeline: - latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler")) - - if args.temporal_latent_upsampler: - original_temporal_upsampler_ckpt = load_hub_or_local_checkpoint( - repo_id=args.original_state_dict_repo_id, filename=args.temporal_latent_upsampler_filename - ) - temporal_latent_upsampler = convert_ltx2_latent_upsampler( - original_temporal_upsampler_ckpt, - get_ltx2_temporal_latent_upsampler_config(args.version), - dtype=vae_dtype, - ) - temporal_latent_upsampler.save_pretrained(os.path.join(args.output_path, "temporal_latent_upsampler")) - - if args.full_pipeline: - is_distilled_ckpt = "distilled" in args.combined_filename - if is_distilled_ckpt: - # Disable dynamic shifting and terminal shift so that distilled sigmas are used as-is - scheduler = FlowMatchEulerDiscreteScheduler( - use_dynamic_shifting=False, - base_shift=0.95, - max_shift=2.05, - base_image_seq_len=1024, - max_image_seq_len=4096, - shift_terminal=None, - ) - else: - scheduler = FlowMatchEulerDiscreteScheduler( - use_dynamic_shifting=True, - base_shift=0.95, - max_shift=2.05, - base_image_seq_len=1024, - max_image_seq_len=4096, - shift_terminal=0.1, - ) - - pipe = LTX2Pipeline( - scheduler=scheduler, - vae=vae, - audio_vae=audio_vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - connectors=connectors, - transformer=transformer, - vocoder=vocoder, - processor=processor, - prompt_enhancer=prompt_enhancer, - ) - - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - if args.upsample_pipeline: - pipe = LTX2LatentUpsamplePipeline(vae=vae, latent_upsampler=latent_upsampler) - - # As two diffusers pipelines cannot be in the same directory, save the upsampling pipeline to its own directory - if args.upsample_output_path: - upsample_output_path = args.upsample_output_path - else: - upsample_output_path = args.output_path - pipe.save_pretrained(upsample_output_path, safe_serialization=True, max_shard_size="5GB") - - -if __name__ == "__main__": - args = get_args() - main(args) diff --git a/scripts/convert_ltx_to_diffusers.py b/scripts/convert_ltx_to_diffusers.py deleted file mode 100644 index b4ae57f9c8f0..000000000000 --- a/scripts/convert_ltx_to_diffusers.py +++ /dev/null @@ -1,525 +0,0 @@ -import argparse -from pathlib import Path -from typing import Any, Dict - -import torch -from accelerate import init_empty_weights -from safetensors.torch import load_file -from transformers import T5EncoderModel, T5Tokenizer - -from diffusers import ( - AutoencoderKLLTXVideo, - FlowMatchEulerDiscreteScheduler, - LTXConditionPipeline, - LTXLatentUpsamplePipeline, - LTXPipeline, - LTXVideoTransformer3DModel, -) -from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel - - -def remove_keys_(key: str, state_dict: Dict[str, Any]): - state_dict.pop(key) - - -TOKENIZER_MAX_LENGTH = 128 - -TRANSFORMER_KEYS_RENAME_DICT = { - "patchify_proj": "proj_in", - "adaln_single": "time_embed", - "q_norm": "norm_q", - "k_norm": "norm_k", -} - -TRANSFORMER_SPECIAL_KEYS_REMAP = { - "vae": remove_keys_, -} - -VAE_KEYS_RENAME_DICT = { - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0", - "up_blocks.2": "up_blocks.1.upsamplers.0", - "up_blocks.3": "up_blocks.1", - "up_blocks.4": "up_blocks.2.conv_in", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.conv_in", - "up_blocks.8": "up_blocks.3.upsamplers.0", - "up_blocks.9": "up_blocks.3", - # encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.0.conv_out", - "down_blocks.3": "down_blocks.1", - "down_blocks.4": "down_blocks.1.downsamplers.0", - "down_blocks.5": "down_blocks.1.conv_out", - "down_blocks.6": "down_blocks.2", - "down_blocks.7": "down_blocks.2.downsamplers.0", - "down_blocks.8": "down_blocks.3", - "down_blocks.9": "mid_block", - # common - "conv_shortcut": "conv_shortcut.conv", - "res_blocks": "resnets", - "norm3.norm": "norm3", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", -} - -VAE_091_RENAME_DICT = { - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", -} - -VAE_095_RENAME_DICT = { - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.1", - "down_blocks.3": "down_blocks.1.downsamplers.0", - "down_blocks.4": "down_blocks.2", - "down_blocks.5": "down_blocks.2.downsamplers.0", - "down_blocks.6": "down_blocks.3", - "down_blocks.7": "down_blocks.3.downsamplers.0", - "down_blocks.8": "mid_block", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", -} - -VAE_SPECIAL_KEYS_REMAP = { - "per_channel_statistics.channel": remove_keys_, - "per_channel_statistics.mean-of-means": remove_keys_, - "per_channel_statistics.mean-of-stds": remove_keys_, - "model.diffusion_model": remove_keys_, -} - - -def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: - state_dict = saved_dict - if "model" in saved_dict.keys(): - state_dict = state_dict["model"] - if "module" in saved_dict.keys(): - state_dict = state_dict["module"] - if "state_dict" in saved_dict.keys(): - state_dict = state_dict["state_dict"] - return state_dict - - -def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def convert_transformer(ckpt_path: str, config, dtype: torch.dtype): - PREFIX_KEY = "model.diffusion_model." - - original_state_dict = get_state_dict(load_file(ckpt_path)) - with init_empty_weights(): - transformer = LTXVideoTransformer3DModel(**config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - if new_key.startswith(PREFIX_KEY): - new_key = key[len(PREFIX_KEY) :] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - transformer.load_state_dict(original_state_dict, strict=True, assign=True) - return transformer - - -def convert_vae(ckpt_path: str, config, dtype: torch.dtype): - PREFIX_KEY = "vae." - - original_state_dict = get_state_dict(load_file(ckpt_path)) - with init_empty_weights(): - vae = AutoencoderKLLTXVideo(**config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - if new_key.startswith(PREFIX_KEY): - new_key = key[len(PREFIX_KEY) :] - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True, assign=True) - return vae - - -def convert_spatial_latent_upsampler(ckpt_path: str, config, dtype: torch.dtype): - original_state_dict = get_state_dict(load_file(ckpt_path)) - - with init_empty_weights(): - latent_upsampler = LTXLatentUpsamplerModel(**config) - - latent_upsampler.load_state_dict(original_state_dict, strict=True, assign=True) - latent_upsampler.to(dtype) - return latent_upsampler - - -def get_transformer_config(version: str) -> dict[str, Any]: - if version == "0.9.7": - config = { - "in_channels": 128, - "out_channels": 128, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 32, - "attention_head_dim": 128, - "cross_attention_dim": 4096, - "num_layers": 48, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 4096, - "attention_bias": True, - "attention_out_bias": True, - } - else: - config = { - "in_channels": 128, - "out_channels": 128, - "patch_size": 1, - "patch_size_t": 1, - "num_attention_heads": 32, - "attention_head_dim": 64, - "cross_attention_dim": 2048, - "num_layers": 28, - "activation_fn": "gelu-approximate", - "qk_norm": "rms_norm_across_heads", - "norm_elementwise_affine": False, - "norm_eps": 1e-6, - "caption_channels": 4096, - "attention_bias": True, - "attention_out_bias": True, - } - return config - - -def get_vae_config(version: str) -> dict[str, Any]: - if version in ["0.9.0"]: - config = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (128, 256, 512, 512), - "down_block_types": ( - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - ), - "decoder_block_out_channels": (128, 256, 512, 512), - "layers_per_block": (4, 3, 3, 3, 4), - "decoder_layers_per_block": (4, 3, 3, 3, 4), - "spatio_temporal_scaling": (True, True, True, False), - "decoder_spatio_temporal_scaling": (True, True, True, False), - "decoder_inject_noise": (False, False, False, False, False), - "downsample_type": ("conv", "conv", "conv", "conv"), - "upsample_residual": (False, False, False, False), - "upsample_factor": (1, 1, 1, 1), - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "scaling_factor": 1.0, - "encoder_causal": True, - "decoder_causal": False, - "timestep_conditioning": False, - } - elif version in ["0.9.1"]: - config = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (128, 256, 512, 512), - "down_block_types": ( - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - "LTXVideoDownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 1024), - "layers_per_block": (4, 3, 3, 3, 4), - "decoder_layers_per_block": (5, 6, 7, 8), - "spatio_temporal_scaling": (True, True, True, False), - "decoder_spatio_temporal_scaling": (True, True, True), - "decoder_inject_noise": (True, True, True, False), - "downsample_type": ("conv", "conv", "conv", "conv"), - "upsample_residual": (True, True, True), - "upsample_factor": (2, 2, 2), - "timestep_conditioning": True, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "scaling_factor": 1.0, - "encoder_causal": True, - "decoder_causal": False, - } - VAE_KEYS_RENAME_DICT.update(VAE_091_RENAME_DICT) - elif version in ["0.9.5"]: - config = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (128, 256, 512, 1024, 2048), - "down_block_types": ( - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 1024), - "layers_per_block": (4, 6, 6, 2, 2), - "decoder_layers_per_block": (5, 5, 5, 5), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True), - "decoder_inject_noise": (False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_residual": (True, True, True), - "upsample_factor": (2, 2, 2), - "timestep_conditioning": True, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "scaling_factor": 1.0, - "encoder_causal": True, - "decoder_causal": False, - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - } - VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT) - elif version in ["0.9.7"]: - config = { - "in_channels": 3, - "out_channels": 3, - "latent_channels": 128, - "block_out_channels": (128, 256, 512, 1024, 2048), - "down_block_types": ( - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - "LTXVideo095DownBlock3D", - ), - "decoder_block_out_channels": (256, 512, 1024), - "layers_per_block": (4, 6, 6, 2, 2), - "decoder_layers_per_block": (5, 5, 5, 5), - "spatio_temporal_scaling": (True, True, True, True), - "decoder_spatio_temporal_scaling": (True, True, True), - "decoder_inject_noise": (False, False, False, False), - "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), - "upsample_residual": (True, True, True), - "upsample_factor": (2, 2, 2), - "timestep_conditioning": True, - "patch_size": 4, - "patch_size_t": 1, - "resnet_norm_eps": 1e-6, - "scaling_factor": 1.0, - "encoder_causal": True, - "decoder_causal": False, - "spatial_compression_ratio": 32, - "temporal_compression_ratio": 8, - } - VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT) - return config - - -def get_spatial_latent_upsampler_config(version: str) -> dict[str, Any]: - if version == "0.9.7": - config = { - "in_channels": 128, - "mid_channels": 512, - "num_blocks_per_stage": 4, - "dims": 3, - "spatial_upsample": True, - "temporal_upsample": False, - } - elif version == "0.9.8": - config = { - "in_channels": 128, - "mid_channels": 512, - "num_blocks_per_stage": 4, - "dims": 3, - "spatial_upsample": True, - "temporal_upsample": False, - } - else: - raise ValueError(f"Unsupported version: {version}") - return config - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" - ) - parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint") - parser.add_argument( - "--spatial_latent_upsampler_path", - type=str, - default=None, - help="Path to original spatial latent upsampler checkpoint", - ) - parser.add_argument( - "--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory" - ) - parser.add_argument( - "--typecast_text_encoder", - action="store_true", - default=False, - help="Whether or not to apply fp16/bf16 precision to text_encoder", - ) - parser.add_argument("--save_pipeline", action="store_true") - parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") - parser.add_argument("--dtype", default="fp32", help="Torch dtype to save the model in.") - parser.add_argument( - "--version", - type=str, - default="0.9.0", - choices=["0.9.0", "0.9.1", "0.9.5", "0.9.7", "0.9.8"], - help="Version of the LTX model", - ) - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - -VARIANT_MAPPING = { - "fp32": None, - "fp16": "fp16", - "bf16": "bf16", -} - - -if __name__ == "__main__": - args = get_args() - - transformer = None - dtype = DTYPE_MAPPING[args.dtype] - variant = VARIANT_MAPPING[args.dtype] - output_path = Path(args.output_path) - - if args.transformer_ckpt_path is not None: - config = get_transformer_config(args.version) - transformer: LTXVideoTransformer3DModel = convert_transformer(args.transformer_ckpt_path, config, dtype) - if not args.save_pipeline: - transformer.save_pretrained( - output_path / "transformer", safe_serialization=True, max_shard_size="5GB", variant=variant - ) - - if args.vae_ckpt_path is not None: - config = get_vae_config(args.version) - vae: AutoencoderKLLTXVideo = convert_vae(args.vae_ckpt_path, config, dtype) - if not args.save_pipeline: - vae.save_pretrained(output_path / "vae", safe_serialization=True, max_shard_size="5GB", variant=variant) - - if args.spatial_latent_upsampler_path is not None: - config = get_spatial_latent_upsampler_config(args.version) - latent_upsampler: LTXLatentUpsamplerModel = convert_spatial_latent_upsampler( - args.spatial_latent_upsampler_path, config, dtype - ) - if not args.save_pipeline: - latent_upsampler.save_pretrained( - output_path / "latent_upsampler", safe_serialization=True, max_shard_size="5GB", variant=variant - ) - - if args.save_pipeline: - text_encoder_id = "google/t5-v1_1-xxl" - tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH) - text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir) - - if args.typecast_text_encoder: - text_encoder = text_encoder.to(dtype=dtype) - - # Apparently, the conversion does not work anymore without this :shrug: - for param in text_encoder.parameters(): - param.data = param.data.contiguous() - - if args.version in ["0.9.5", "0.9.7"]: - scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=False) - else: - scheduler = FlowMatchEulerDiscreteScheduler( - use_dynamic_shifting=True, - base_shift=0.95, - max_shift=2.05, - base_image_seq_len=1024, - max_image_seq_len=4096, - shift_terminal=0.1, - ) - - if args.version in ["0.9.0", "0.9.1", "0.9.5"]: - pipe = LTXPipeline( - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - ) - pipe.save_pretrained( - output_path.as_posix(), safe_serialization=True, variant=variant, max_shard_size="5GB" - ) - elif args.version in ["0.9.7"]: - pipe = LTXConditionPipeline( - scheduler=scheduler, - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - ) - pipe_upsample = LTXLatentUpsamplePipeline( - vae=vae, - latent_upsampler=latent_upsampler, - ) - pipe.save_pretrained( - (output_path / "ltx_pipeline").as_posix(), - safe_serialization=True, - variant=variant, - max_shard_size="5GB", - ) - pipe_upsample.save_pretrained( - (output_path / "ltx_upsample_pipeline").as_posix(), - safe_serialization=True, - variant=variant, - max_shard_size="5GB", - ) - else: - raise ValueError(f"Unsupported version: {args.version}") diff --git a/scripts/convert_lumina_to_diffusers.py b/scripts/convert_lumina_to_diffusers.py deleted file mode 100644 index c14aad3c6bf2..000000000000 --- a/scripts/convert_lumina_to_diffusers.py +++ /dev/null @@ -1,142 +0,0 @@ -import argparse -import os - -import torch -from safetensors.torch import load_file -from transformers import AutoModel, AutoTokenizer - -from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, LuminaNextDiT2DModel, LuminaPipeline - - -def main(args): - # checkpoint from https://huggingface.co/Alpha-VLLM/Lumina-Next-SFT or https://huggingface.co/Alpha-VLLM/Lumina-Next-T2I - all_sd = load_file(args.origin_ckpt_path, device="cpu") - converted_state_dict = {} - # pad token - converted_state_dict["pad_token"] = all_sd["pad_token"] - - # patch embed - converted_state_dict["patch_embedder.weight"] = all_sd["x_embedder.weight"] - converted_state_dict["patch_embedder.bias"] = all_sd["x_embedder.bias"] - - # time and caption embed - converted_state_dict["time_caption_embed.timestep_embedder.linear_1.weight"] = all_sd["t_embedder.mlp.0.weight"] - converted_state_dict["time_caption_embed.timestep_embedder.linear_1.bias"] = all_sd["t_embedder.mlp.0.bias"] - converted_state_dict["time_caption_embed.timestep_embedder.linear_2.weight"] = all_sd["t_embedder.mlp.2.weight"] - converted_state_dict["time_caption_embed.timestep_embedder.linear_2.bias"] = all_sd["t_embedder.mlp.2.bias"] - converted_state_dict["time_caption_embed.caption_embedder.0.weight"] = all_sd["cap_embedder.0.weight"] - converted_state_dict["time_caption_embed.caption_embedder.0.bias"] = all_sd["cap_embedder.0.bias"] - converted_state_dict["time_caption_embed.caption_embedder.1.weight"] = all_sd["cap_embedder.1.weight"] - converted_state_dict["time_caption_embed.caption_embedder.1.bias"] = all_sd["cap_embedder.1.bias"] - - for i in range(24): - # adaln - converted_state_dict[f"layers.{i}.gate"] = all_sd[f"layers.{i}.attention.gate"] - converted_state_dict[f"layers.{i}.adaLN_modulation.1.weight"] = all_sd[f"layers.{i}.adaLN_modulation.1.weight"] - converted_state_dict[f"layers.{i}.adaLN_modulation.1.bias"] = all_sd[f"layers.{i}.adaLN_modulation.1.bias"] - - # qkv - converted_state_dict[f"layers.{i}.attn1.to_q.weight"] = all_sd[f"layers.{i}.attention.wq.weight"] - converted_state_dict[f"layers.{i}.attn1.to_k.weight"] = all_sd[f"layers.{i}.attention.wk.weight"] - converted_state_dict[f"layers.{i}.attn1.to_v.weight"] = all_sd[f"layers.{i}.attention.wv.weight"] - - # cap - converted_state_dict[f"layers.{i}.attn2.to_q.weight"] = all_sd[f"layers.{i}.attention.wq.weight"] - converted_state_dict[f"layers.{i}.attn2.to_k.weight"] = all_sd[f"layers.{i}.attention.wk_y.weight"] - converted_state_dict[f"layers.{i}.attn2.to_v.weight"] = all_sd[f"layers.{i}.attention.wv_y.weight"] - - # output - converted_state_dict[f"layers.{i}.attn2.to_out.0.weight"] = all_sd[f"layers.{i}.attention.wo.weight"] - - # attention - # qk norm - converted_state_dict[f"layers.{i}.attn1.norm_q.weight"] = all_sd[f"layers.{i}.attention.q_norm.weight"] - converted_state_dict[f"layers.{i}.attn1.norm_q.bias"] = all_sd[f"layers.{i}.attention.q_norm.bias"] - - converted_state_dict[f"layers.{i}.attn1.norm_k.weight"] = all_sd[f"layers.{i}.attention.k_norm.weight"] - converted_state_dict[f"layers.{i}.attn1.norm_k.bias"] = all_sd[f"layers.{i}.attention.k_norm.bias"] - - converted_state_dict[f"layers.{i}.attn2.norm_q.weight"] = all_sd[f"layers.{i}.attention.q_norm.weight"] - converted_state_dict[f"layers.{i}.attn2.norm_q.bias"] = all_sd[f"layers.{i}.attention.q_norm.bias"] - - converted_state_dict[f"layers.{i}.attn2.norm_k.weight"] = all_sd[f"layers.{i}.attention.ky_norm.weight"] - converted_state_dict[f"layers.{i}.attn2.norm_k.bias"] = all_sd[f"layers.{i}.attention.ky_norm.bias"] - - # attention norm - converted_state_dict[f"layers.{i}.attn_norm1.weight"] = all_sd[f"layers.{i}.attention_norm1.weight"] - converted_state_dict[f"layers.{i}.attn_norm2.weight"] = all_sd[f"layers.{i}.attention_norm2.weight"] - converted_state_dict[f"layers.{i}.norm1_context.weight"] = all_sd[f"layers.{i}.attention_y_norm.weight"] - - # feed forward - converted_state_dict[f"layers.{i}.feed_forward.linear_1.weight"] = all_sd[f"layers.{i}.feed_forward.w1.weight"] - converted_state_dict[f"layers.{i}.feed_forward.linear_2.weight"] = all_sd[f"layers.{i}.feed_forward.w2.weight"] - converted_state_dict[f"layers.{i}.feed_forward.linear_3.weight"] = all_sd[f"layers.{i}.feed_forward.w3.weight"] - - # feed forward norm - converted_state_dict[f"layers.{i}.ffn_norm1.weight"] = all_sd[f"layers.{i}.ffn_norm1.weight"] - converted_state_dict[f"layers.{i}.ffn_norm2.weight"] = all_sd[f"layers.{i}.ffn_norm2.weight"] - - # final layer - converted_state_dict["final_layer.linear.weight"] = all_sd["final_layer.linear.weight"] - converted_state_dict["final_layer.linear.bias"] = all_sd["final_layer.linear.bias"] - - converted_state_dict["final_layer.adaLN_modulation.1.weight"] = all_sd["final_layer.adaLN_modulation.1.weight"] - converted_state_dict["final_layer.adaLN_modulation.1.bias"] = all_sd["final_layer.adaLN_modulation.1.bias"] - - # Lumina-Next-SFT 2B - transformer = LuminaNextDiT2DModel( - sample_size=128, - patch_size=2, - in_channels=4, - hidden_size=2304, - num_layers=24, - num_attention_heads=32, - num_kv_heads=8, - multiple_of=256, - ffn_dim_multiplier=None, - norm_eps=1e-5, - learn_sigma=True, - qk_norm=True, - cross_attention_dim=2048, - scaling_factor=1.0, - ) - transformer.load_state_dict(converted_state_dict, strict=True) - - num_model_params = sum(p.numel() for p in transformer.parameters()) - print(f"Total number of transformer parameters: {num_model_params}") - - if args.only_transformer: - transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) - else: - scheduler = FlowMatchEulerDiscreteScheduler() - - vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", torch_dtype=torch.float32) - - tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b") - text_encoder = AutoModel.from_pretrained("google/gemma-2b") - - pipeline = LuminaPipeline( - tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler - ) - pipeline.save_pretrained(args.dump_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--origin_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--image_size", - default=1024, - type=int, - choices=[256, 512, 1024], - required=False, - help="Image size of pretrained model, either 512 or 1024.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") - parser.add_argument("--only_transformer", default=True, type=bool, required=True) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_minimax_h3_to_diffusers.py b/scripts/convert_minimax_h3_to_diffusers.py deleted file mode 100644 index 337d6c0e6816..000000000000 --- a/scripts/convert_minimax_h3_to_diffusers.py +++ /dev/null @@ -1,947 +0,0 @@ -"""Convert an original MiniMax-H3 checkpoint into the diffusers layout. - -The transformer checkpoint is streamed shard by shard, so peak memory stays close to a single shard (~4.9 GiB) and -never approaches the 62 GiB of the full 33B DiT. - -Every source key maps onto a diffusers module by renaming alone, except for three things: - -* the fused QKV projection is split into `to_q` / `to_k` / `to_v`. Raw checkpoint shards store QKV per-head - interleaved; the shard streamer first reorders them into the reference's in-memory `[q_all; k_all; v_all]` layout - (`reorder_interleaved_qkv`, mirroring the reference's load-time transform), then `convert_transformer_key` splits - contiguous thirds (`split_fused_qkv`), -* the gated FFN's fused `mlp.fc1` becomes `ff.net.0.proj` and its two halves are swapped, because diffusers' - [`SwiGLU`] reads `[value; gate]` where the reference stores `[gate; value]` (the same transform the video VAE needs, - see `convert_video_vae_key`); `mlp.fc2` becomes `ff.net.2`, -* `rope.inv_freq` is dropped: it is a pure function of `rope_theta` and `rope_freq_dim`, the port recomputes it in a - non-persistent buffer, and the recomputed value is bitwise equal to the shipped one for both released variants. - -There are no transposes anywhere. - -The FL2VA and Ref2VA variants differ only in the transformer weights, so the variant is selected by pointing -`--checkpoint_path` at the corresponding folder. Both land in one repository, which carries a single -`modular_model_index.json`: MiniMax-H3 is integrated as Modular Diffusers blocks only, so no `model_index.json` is -written. - -Usage: - -```bash -# Validate the key mapping without any weights present. -python scripts/convert_minimax_h3_to_diffusers.py \ - --checkpoint_path /path/to/MiniMax-H3/FL2VA --output_path /tmp/h3-diffusers --dry_run - -# Convert, and point the component loading specs at the Hub id the result is published under. -python scripts/convert_minimax_h3_to_diffusers.py \ - --checkpoint_path /path/to/MiniMax-H3/FL2VA --output_path /tmp/h3-diffusers \ - --modular_repo_id MiniMaxAI/MiniMax-H3 -``` -""" - -import argparse -import glob -import json -import math -import os -import struct -from typing import Any - -import torch -from safetensors import safe_open -from safetensors.torch import save_file - -from diffusers.models.autoencoders.autoencoder_kl_minimax_h3_audio import AutoencoderKLMiniMaxH3Audio -from diffusers.utils.constants import SAFE_WEIGHTS_INDEX_NAME - - -# `MiniMaxH3Transformer3DModel` argument names. The original config uses the sglang-native names listed in the -# comments; everything else in the original config (`adaln_out_features`, `final_adaln_out_features`) is derived. -MINIMAX_H3_TRANSFORMER_CONFIG = { - "num_attention_heads": 56, - "attention_head_dim": 128, - "hidden_size": 5376, - "num_layers": 50, - "num_refiner_layers": 2, # token_refiner_num_layers - "ffn_dim": 14336, # ffn_hidden_size - "in_channels": 24, # latents_dim - "audio_in_channels": 32, # audio_latents_dim - "patch_size": [1, 2, 2], - "text_dim": 5120, - "freq_dim": 256, # timestep_input_dim - "time_embed_hidden_dim": 5376, # time_embed_hidden_size - "time_embed_dim": 2688, - "rope_freq_dim": 16, # rope_inv_freq_len - "rope_theta": 10000.0, - "norm_eps": 1e-05, - "qk_norm_eps": 1e-05, - "final_norm_eps": 1e-05, -} - -# A tiny configuration with the checkpoint-tied dimensions left intact, for building fixtures. -MINIMAX_H3_TEST_TRANSFORMER_CONFIG = { - **MINIMAX_H3_TRANSFORMER_CONFIG, - "num_attention_heads": 2, - "attention_head_dim": 32, - "hidden_size": 64, - "num_layers": 2, - "num_refiner_layers": 2, - "ffn_dim": 128, - "text_dim": 48, - "freq_dim": 16, - "time_embed_hidden_dim": 64, - "time_embed_dim": 32, - "rope_freq_dim": 4, -} - -# MiniMax-H3 ships a mixed-precision checkpoint. These *original* keys are float32; everything else is bfloat16 — -# including the AdaLN projections. -MINIMAX_H3_FP32_SOURCE_PREFIXES = ( - "video_patch_proj.", - "audio_patch_proj.", - "time_embedder.", - "final_layer.video_out.", - "final_layer.audio_out.", -) - -# `rope.inv_freq` is `1 / rope_theta ** (arange(0, 2 * rope_freq_dim, 2) / (2 * rope_freq_dim))`, which -# `MiniMaxH3RotaryPosEmbed` recomputes into a non-persistent buffer. The recomputed tensor is bitwise equal to the -# shipped one in both released variants, so the key is not carried into the diffusers checkpoint. -MINIMAX_H3_TRANSFORMER_DROPPED_KEYS = ("rope.inv_freq",) - - -def reorder_interleaved_qkv(weight: torch.Tensor, num_attention_heads: int, attention_head_dim: int) -> torch.Tensor: - """Reorder a *raw-checkpoint* per-head-interleaved fused QKV weight into `[q_all; k_all; v_all]`. - - The original checkpoint shards store rows as `[head0: q(head_dim) k(head_dim) v(head_dim), head1: q, k, v, ...]`. - The reference implementation applies exactly this reorder at load time (`_reorder_grouped_qkv_to_qkv` with one head - per query group), so `[q_all; k_all; v_all]` is the reference's in-memory / state-dict layout. There is no - transpose. - """ - expected_rows = num_attention_heads * 3 * attention_head_dim - if weight.shape[0] != expected_rows: - raise ValueError( - f"fused qkv weight has {weight.shape[0]} rows, expected " - f"{expected_rows} = {num_attention_heads} heads * 3 * {attention_head_dim}." - ) - grouped = weight.reshape(num_attention_heads, 3 * attention_head_dim, *weight.shape[1:]) - query, key, value = grouped.split(attention_head_dim, dim=1) - return torch.cat( - [ - tensor.reshape(num_attention_heads * attention_head_dim, *weight.shape[1:]) - for tensor in (query, key, value) - ], - dim=0, - ) - - -def split_fused_qkv( - weight: torch.Tensor, num_attention_heads: int, attention_head_dim: int -) -> tuple[torch.Tensor, ...]: - """Split a fused `[q_all; k_all; v_all]` QKV weight into separate `to_q` / `to_k` / `to_v` weights. - - The input is the *reference model* layout — what `MiniMaxH3DiTModel.state_dict()` holds after the reference's - load-time reorder — i.e. the three logical projection matrices stacked contiguously, NOT the raw checkpoint's - per-head interleave (see `reorder_interleaved_qkv`, which the shard streamer applies first). - """ - inner_dim = num_attention_heads * attention_head_dim - if weight.shape[0] != 3 * inner_dim: - raise ValueError( - f"fused qkv weight has {weight.shape[0]} rows, expected " - f"{3 * inner_dim} = 3 * {num_attention_heads} heads * {attention_head_dim}." - ) - query, key, value = weight.split(inner_dim, dim=0) - return tuple(tensor.contiguous() for tensor in (query, key, value)) - - -def get_transformer_key_plan(config: dict[str, Any]) -> dict[str, list[tuple[str, list[int]]]]: - """Map every original transformer key to the diffusers key(s) it produces, with the resulting shapes. - - The plan is derived from the config alone, so it can be printed and checked without any weights present. - """ - hidden_size = config["hidden_size"] - heads = config["num_attention_heads"] - head_dim = config["attention_head_dim"] - inner_dim = heads * head_dim - ffn_dim = config["ffn_dim"] - time_embed_dim = config["time_embed_dim"] - video_patch_dim = ( - config["in_channels"] * config["patch_size"][0] * config["patch_size"][1] * config["patch_size"][2] - ) - - plan: dict[str, list[tuple[str, list[int]]]] = { - "video_patch_proj.weight": [("proj_in.weight", [hidden_size, video_patch_dim])], - "video_patch_proj.bias": [("proj_in.bias", [hidden_size])], - "audio_patch_proj.weight": [("audio_proj_in.weight", [hidden_size, config["audio_in_channels"]])], - "audio_patch_proj.bias": [("audio_proj_in.bias", [hidden_size])], - "condition_proj.weight": [("context_embedder.weight", [hidden_size, config["text_dim"]])], - "condition_proj.bias": [("context_embedder.bias", [hidden_size])], - # `Timesteps` + `TimestepEmbedding` reproduce the reference sinusoid and MLP exactly, so the timestep MLP is - # renamed onto `TimestepEmbedding`'s `linear_1` / `linear_2`. - "time_embedder.proj_in.weight": [ - ("time_embedder.linear_1.weight", [config["time_embed_hidden_dim"], config["freq_dim"]]) - ], - "time_embedder.proj_in.bias": [("time_embedder.linear_1.bias", [config["time_embed_hidden_dim"]])], - "time_embedder.proj_out.weight": [ - ("time_embedder.linear_2.weight", [time_embed_dim, config["time_embed_hidden_dim"]]) - ], - "time_embedder.proj_out.bias": [("time_embedder.linear_2.bias", [time_embed_dim])], - "token_refiner.final_norm.weight": [("token_refiner.final_norm.weight", [hidden_size])], - "final_layer.norm.weight": [("norm_out.norm.weight", [hidden_size])], - "final_layer.adaln_proj.linear.weight": [("norm_out.linear.weight", [2 * hidden_size, time_embed_dim])], - "final_layer.adaln_proj.linear.bias": [("norm_out.linear.bias", [2 * hidden_size])], - "final_layer.video_out.weight": [("proj_out.weight", [video_patch_dim, hidden_size])], - "final_layer.video_out.bias": [("proj_out.bias", [video_patch_dim])], - "final_layer.audio_out.weight": [("audio_proj_out.weight", [config["audio_in_channels"], hidden_size])], - "final_layer.audio_out.bias": [("audio_proj_out.bias", [config["audio_in_channels"]])], - } - for key in MINIMAX_H3_TRANSFORMER_DROPPED_KEYS: - plan[key] = [] - - block_specs = [ - ("blocks", "transformer_blocks", config["num_layers"], True), - ("token_refiner.blocks", "token_refiner.refiner_blocks", config["num_refiner_layers"], False), - ] - for source_prefix, target_prefix, num_layers, has_adaln in block_specs: - for i in range(num_layers): - source = f"{source_prefix}.{i}" - target = f"{target_prefix}.{i}" - plan[f"{source}.norm1.weight"] = [(f"{target}.norm1.weight", [hidden_size])] - plan[f"{source}.norm2.weight"] = [(f"{target}.norm2.weight", [hidden_size])] - plan[f"{source}.attn.qkv_proj.weight"] = [ - (f"{target}.attn.to_q.weight", [inner_dim, hidden_size]), - (f"{target}.attn.to_k.weight", [inner_dim, hidden_size]), - (f"{target}.attn.to_v.weight", [inner_dim, hidden_size]), - ] - plan[f"{source}.attn.q_norm.weight"] = [(f"{target}.attn.norm_q.weight", [head_dim])] - plan[f"{source}.attn.k_norm.weight"] = [(f"{target}.attn.norm_k.weight", [head_dim])] - plan[f"{source}.attn.out_proj.weight"] = [(f"{target}.attn.to_out.0.weight", [hidden_size, inner_dim])] - # `fc1` stays fused, as diffusers' `SwiGLU` also fuses its two projections, but the halves are swapped - # from `[gate; value]` to `[value; gate]` (see `convert_transformer_key`). - plan[f"{source}.mlp.fc1.weight"] = [(f"{target}.ff.net.0.proj.weight", [2 * ffn_dim, hidden_size])] - plan[f"{source}.mlp.fc2.weight"] = [(f"{target}.ff.net.2.weight", [hidden_size, ffn_dim])] - if has_adaln: - plan[f"{source}.adaln_proj.linear.weight"] = [ - (f"{target}.adaln_proj.linear.weight", [6 * 3 * hidden_size, time_embed_dim]) - ] - plan[f"{source}.adaln_proj.linear.bias"] = [ - (f"{target}.adaln_proj.linear.bias", [6 * 3 * hidden_size]) - ] - - return plan - - -def convert_transformer_key( - source_key: str, tensor: torch.Tensor, config: dict[str, Any] -) -> list[tuple[str, torch.Tensor]]: - """Convert one original key/tensor pair into the diffusers key/tensor pair(s) it maps to.""" - if source_key in MINIMAX_H3_TRANSFORMER_DROPPED_KEYS: - return [] - - target_key = source_key - if target_key.startswith("token_refiner.blocks."): - target_key = target_key.replace("token_refiner.blocks.", "token_refiner.refiner_blocks.", 1) - elif target_key.startswith("blocks."): - target_key = target_key.replace("blocks.", "transformer_blocks.", 1) - target_key = target_key.replace("time_embedder.proj_in.", "time_embedder.linear_1.") - target_key = target_key.replace("time_embedder.proj_out.", "time_embedder.linear_2.") - target_key = target_key.replace("video_patch_proj.", "proj_in.") - target_key = target_key.replace("audio_patch_proj.", "audio_proj_in.") - target_key = target_key.replace("condition_proj.", "context_embedder.") - target_key = target_key.replace("final_layer.norm.", "norm_out.norm.") - target_key = target_key.replace("final_layer.adaln_proj.linear.", "norm_out.linear.") - target_key = target_key.replace("final_layer.video_out.", "proj_out.") - target_key = target_key.replace("final_layer.audio_out.", "audio_proj_out.") - target_key = target_key.replace(".attn.q_norm.", ".attn.norm_q.") - target_key = target_key.replace(".attn.k_norm.", ".attn.norm_k.") - target_key = target_key.replace(".attn.out_proj.", ".attn.to_out.0.") - - if target_key.endswith(".attn.qkv_proj.weight"): - # `convert_transformer_key` consumes tensors in the reference model's state-dict layout, where the fused QKV - # rows are already `[q_all; k_all; v_all]`. Raw checkpoint shards are per-head interleaved instead; the shard - # streamer (`convert_transformer`) normalizes them with `reorder_interleaved_qkv` before calling this. - query, key, value = split_fused_qkv(tensor, config["num_attention_heads"], config["attention_head_dim"]) - prefix = target_key.removesuffix("qkv_proj.weight") - return [(f"{prefix}to_q.weight", query), (f"{prefix}to_k.weight", key), (f"{prefix}to_v.weight", value)] - - if target_key.endswith(".mlp.fc1.weight"): - # The reference computes `fc2(silu(gate) * value)` from a fused `[gate; value]`; diffusers' `SwiGLU` computes - # `value * silu(gate)` from a fused `[value; gate]`, so the two halves swap places. Identical transform to the - # video VAE's `ff.w1` (see `convert_video_vae_key`). - gate, value = tensor.chunk(2, dim=0) - target_key = target_key.replace(".mlp.fc1.weight", ".ff.net.0.proj.weight") - return [(target_key, torch.cat([value, gate], dim=0).contiguous())] - - target_key = target_key.replace(".mlp.fc2.", ".ff.net.2.") - return [(target_key, tensor)] - - -# -# --------------------------------------------------------------------------------------------------------------- -# Video VAE -# --------------------------------------------------------------------------------------------------------------- -# - -# `AutoencoderKLMiniMaxH3` argument names. Field-for-field equal to `video_vae/source/config.json`, with the original -# names in the comments. The keys that only ever take one value in the release (`use_3d_conv`, `use_vit_decoder`, -# `causal_encoder`, `causal_decoder`, `use_t_isolated_gn`, `space_up` / `time_up`, `zq_ch_*`, `num_res_blocks_decoder`, -# `shift_factor` / `scaling_factor`) are baked into the port instead of being config knobs. -MINIMAX_H3_VIDEO_VAE_CONFIG = { - "in_channels": 3, - "out_channels": 3, # out_ch - "latent_channels": 24, # z_channels == embed_dim - "block_out_channels": [128, 256, 256, 512, 512, 1024], # ch * ch_mult - "layers_per_block": 2, # num_res_blocks - "spatial_downsample_factors": [2, 2, 2, 2, 1, 1], # space_down - "temporal_downsample_factors": [1, 2, 2, 1, 1, 1], # time_down - "norm_num_groups": 32, - "norm_eps": 1e-06, - "spatial_padding_mode": "reflect", # padding_mode - "decoder_num_layers": 36, # vit_decoder_kwargs.num_layers - "decoder_num_attention_heads": 32, # vit_decoder_kwargs.heads - "decoder_attention_head_dim": 64, # vit_decoder_kwargs.dim_head - "decoder_num_register_tokens": 4, # ViT3DDecoder default - "decoder_ffn_mult": 4, # FeedForward default - "decoder_rope_theta": 100.0, # vit_decoder_kwargs.rope_theta - "decoder_rope_dim_ratio": 0.75, # vit_decoder_kwargs.rope_dim_ratio - "decoder_norm_eps": 1e-05, # ViT3DDecoder eps - "clip_length": 17, # video_vae/config.json vae_clip_length - "token_drop": 3, # video_vae/config.json vae_token_drop -} - -# A tiny configuration with the checkpoint-tied dimensions (`latent_channels`, the temporal geometry and the rotary -# ratio) left intact, for building fixtures and for the CPU parity check. -MINIMAX_H3_TEST_VIDEO_VAE_CONFIG = { - **MINIMAX_H3_VIDEO_VAE_CONFIG, - "block_out_channels": [32, 64], - "layers_per_block": 1, - "spatial_downsample_factors": [2, 2], - "temporal_downsample_factors": [2, 2], - "decoder_num_layers": 4, - "decoder_num_attention_heads": 4, - "decoder_attention_head_dim": 32, -} - -# `decoder.mask_token` is an all-zero buffer belonging to the masked-autoencoding training objective; the released -# decoder never reads it, so the port does not carry the module and the conversion drops the key. -MINIMAX_H3_VIDEO_VAE_DROPPED_KEYS = ("decoder.mask_token",) - - -def convert_video_vae_key( - source_key: str, tensor: torch.Tensor, config: dict[str, Any] -) -> list[tuple[str, torch.Tensor]]: - """Convert one original video-VAE key/tensor pair into the diffusers key/tensor pair(s) it maps to. - - `quant_conv` / `post_quant_conv`, the encoder's `conv_in` / `norm_out` / `conv_out` and the ViT decoder's - `register_tokens` / `norm_out` / `proj_out` / `norm{1,2}` / `scale{1,2}` are pure pass-throughs. What moves: - - * the encoder's CNN levels are renamed from the original CompVis spelling onto the diffusers autoencoder idiom: - `down.{i}.block.{j}` -> `down_blocks.{i}.resnets.{j}`, `nin_shortcut` -> `conv_shortcut`, and - `down.{i}.downsample` -> `down_blocks.{i}.downsamplers.0`, - * the ViT decoder's `x_embedder` becomes `proj_in`, the counterpart of the `proj_out` it already ships, - * the fused per-head-interleaved `attn.to_qkv` is split into `attn.to_q` / `to_k` / `to_v`, - * `attn.to_out` becomes `attn.to_out.0` (diffusers wraps the output projection in an `nn.ModuleList`), - * the gated FFN's `w1` / `w2` become `ff.net.0.proj` / `ff.net.2`, and the two halves of `w1` are swapped because - diffusers' `SwiGLU` reads `[up; gate]` where the reference stores `[gate; up]`. - """ - if source_key in MINIMAX_H3_VIDEO_VAE_DROPPED_KEYS: - return [] - - if ".attn.to_qkv." in source_key: - # Same per-head interleave as the DiT: `[head0: q k v, head1: q k v, ...]`. - reordered = reorder_interleaved_qkv( - tensor, config["decoder_num_attention_heads"], config["decoder_attention_head_dim"] - ) - query, key, value = split_fused_qkv( - reordered, config["decoder_num_attention_heads"], config["decoder_attention_head_dim"] - ) - prefix, suffix = source_key.split(".attn.to_qkv.") - return [ - (f"{prefix}.attn.to_q.{suffix}", query), - (f"{prefix}.attn.to_k.{suffix}", key), - (f"{prefix}.attn.to_v.{suffix}", value), - ] - - target_key = _rename_video_vae_key(source_key) - - if ".ff.w1." in source_key: - gate, up = tensor.chunk(2, dim=0) - return [(target_key, torch.cat([up, gate], dim=0).contiguous())] - - return [(target_key, tensor)] - - -def _rename_video_vae_key(source_key: str) -> str: - """Rename one original video-VAE key onto its diffusers module path (no tensor transform).""" - target_key = source_key - if target_key.startswith("encoder.down."): - level, rest = target_key.removeprefix("encoder.down.").split(".", 1) - rest = rest.replace("block.", "resnets.", 1).replace("nin_shortcut.", "conv_shortcut.", 1) - rest = rest.replace("downsample.", "downsamplers.0.", 1) - target_key = f"encoder.down_blocks.{level}.{rest}" - target_key = target_key.replace("decoder.x_embedder.", "decoder.proj_in.") - target_key = target_key.replace(".attn.to_out.", ".attn.to_out.0.") - target_key = target_key.replace(".ff.w1.", ".ff.net.0.proj.") - target_key = target_key.replace(".ff.w2.", ".ff.net.2.") - return target_key - - -def get_video_vae_key_plan(config: dict[str, Any]) -> dict[str, list[str]]: - """Map every original video-VAE key to the diffusers key(s) it produces, derived from the config alone.""" - block_out_channels = config["block_out_channels"] - block_in_channels = [block_out_channels[0]] + block_out_channels[:-1] - plan: dict[str, list[str]] = {} - - def renamed(*keys: str) -> None: - """Register keys whose diffusers name is `_rename_video_vae_key(key)` and whose tensor is unchanged.""" - for key in keys: - plan[key] = [_rename_video_vae_key(key)] - - renamed("quant_conv.weight", "quant_conv.bias", "post_quant_conv.weight", "post_quant_conv.bias") - renamed("encoder.conv_in.weight", "encoder.conv_in.bias") - for level, (in_channels, out_channels) in enumerate(zip(block_in_channels, block_out_channels)): - for i in range(config["layers_per_block"]): - prefix = f"encoder.down.{level}.block.{i}" - for name in ("norm1", "conv1", "norm2", "conv2"): - renamed(f"{prefix}.{name}.weight", f"{prefix}.{name}.bias") - if (in_channels if i == 0 else out_channels) != out_channels: - renamed(f"{prefix}.nin_shortcut.weight", f"{prefix}.nin_shortcut.bias") - if config["spatial_downsample_factors"][level] * config["temporal_downsample_factors"][level] > 1: - renamed(f"encoder.down.{level}.downsample.conv.weight", f"encoder.down.{level}.downsample.conv.bias") - renamed("encoder.norm_out.weight", "encoder.norm_out.bias", "encoder.conv_out.weight", "encoder.conv_out.bias") - - renamed("decoder.x_embedder.weight", "decoder.x_embedder.bias", "decoder.register_tokens") - renamed("decoder.norm_out.weight", "decoder.norm_out.bias", "decoder.proj_out.weight", "decoder.proj_out.bias") - for i in range(config["decoder_num_layers"]): - prefix = f"decoder.transformer_blocks.{i}" - renamed(f"{prefix}.norm1.weight", f"{prefix}.norm2.weight", f"{prefix}.scale1", f"{prefix}.scale2") - for suffix in ("weight", "bias"): - plan[f"{prefix}.attn.to_qkv.{suffix}"] = [ - f"{prefix}.attn.to_q.{suffix}", - f"{prefix}.attn.to_k.{suffix}", - f"{prefix}.attn.to_v.{suffix}", - ] - plan[f"{prefix}.attn.to_out.{suffix}"] = [f"{prefix}.attn.to_out.0.{suffix}"] - plan[f"{prefix}.ff.w1.{suffix}"] = [f"{prefix}.ff.net.0.proj.{suffix}"] - plan[f"{prefix}.ff.w2.{suffix}"] = [f"{prefix}.ff.net.2.{suffix}"] - for key in MINIMAX_H3_VIDEO_VAE_DROPPED_KEYS: - plan[key] = [] - return plan - - -def convert_video_vae( - checkpoint_path: str, output_path: str, config: dict[str, Any], diffusers_version: str, max_shard_size: int -) -> None: - """Convert the video VAE and emit its config. - - The original weights live one level deeper than the rest of the checkpoint - (`video_vae/source/model.safetensors`, resolved by a hook in the reference); the diffusers layout flattens that to - `vae/`. `latents_mean` / `latents_std` and the tiling geometry come from `video_vae/config.json` — MiniMax-H3 - normalizes latents per channel instead of with a `scaling_factor`. - """ - source_dir = os.path.join(checkpoint_path, "video_vae") - with open(os.path.join(source_dir, "config.json")) as f: - wrapper_config = json.load(f) - for key in ("latents_mean", "latents_std"): - if key not in wrapper_config: - raise KeyError(f"{source_dir}/config.json does not carry `{key}`.") - - weights_path = os.path.join(source_dir, wrapper_config["source_path"], wrapper_config["source_safetensors_path"]) - plan = get_video_vae_key_plan(config) - - os.makedirs(output_path, exist_ok=True) - weight_map: dict[str, str] = {} - total_size = 0 - written: list[str] = [] - buffer: dict[str, torch.Tensor] = {} - buffer_size = 0 - seen_source_keys: set[str] = set() - - def flush() -> None: - nonlocal buffer, buffer_size - if not buffer: - return - path = os.path.join(output_path, f".tmp-shard-{len(written):05d}.safetensors") - save_file(buffer, path, metadata={"format": "pt"}) - for key in buffer: - weight_map[key] = path - written.append(path) - buffer = {} - buffer_size = 0 - - # `safe_open` memory-maps the file, so only the tensor being read is materialized. - with safe_open(weights_path, framework="pt", device="cpu") as f: - for source_key in f.keys(): - if source_key not in plan: - raise KeyError(f"Unexpected key in {weights_path}: {source_key}") - seen_source_keys.add(source_key) - for target_key, tensor in convert_video_vae_key(source_key, f.get_tensor(source_key), config): - if tensor.dtype != torch.float32: - raise ValueError(f"{source_key}: expected torch.float32, got {tensor.dtype}.") - buffer[target_key] = tensor - buffer_size += tensor.numel() * tensor.element_size() - total_size += tensor.numel() * tensor.element_size() - if buffer_size >= max_shard_size: - flush() - flush() - - missing = sorted(set(plan) - seen_source_keys) - if missing: - raise KeyError(f"{len(missing)} planned key(s) missing from {weights_path}, e.g. {missing[:5]}.") - - renames = { - path: os.path.join(output_path, f"diffusion_pytorch_model-{i + 1:05d}-of-{len(written):05d}.safetensors") - for i, path in enumerate(written) - } - for old, new in renames.items(): - os.rename(old, new) - index = { - "metadata": {"total_size": total_size}, - "weight_map": {key: os.path.basename(renames[path]) for key, path in weight_map.items()}, - } - with open(os.path.join(output_path, SAFE_WEIGHTS_INDEX_NAME), "w") as f: - json.dump(index, f, indent=2, sort_keys=True) - - with open(os.path.join(output_path, "config.json"), "w") as f: - json.dump( - { - "_class_name": "AutoencoderKLMiniMaxH3", - "_diffusers_version": diffusers_version, - **config, - "latents_mean": wrapper_config["latents_mean"], - "latents_std": wrapper_config["latents_std"], - }, - f, - indent=2, - ) - - print( - f"video_vae: {len(seen_source_keys)} original keys -> {len(weight_map)} diffusers keys in " - f"{len(written)} shard(s), {total_size / 1024**3:.2f} GiB " - f"(latents_mean/latents_std: {len(wrapper_config['latents_mean'])}/" - f"{len(wrapper_config['latents_std'])} channels; tiling {wrapper_config['vae_tile_size']}px / " - f"{wrapper_config['vae_tile_overlap_min']}px min overlap)." - ) - - -# Not present in `audio_vae/metadata.json`: the reference implementation hardcodes these in its DAC audio VAE and its -# attention projection, keyed off the sample rate. -MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG = { - "num_attention_heads": 8, - "resblock_kernel_sizes": [3, 7, 11], - "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], -} - - -def get_audio_vae_config(checkpoint_path: str) -> dict[str, Any]: - """Build the `AutoencoderKLMiniMaxH3Audio` config from the original audio-VAE metadata. - - `audio_vae/metadata.json` carries the constructor kwargs the checkpoint was built with, and - `audio_vae/config.json` carries the per-channel `latents_mean` / `latents_std` MiniMax-H3 normalizes with. The two - are cross-checked here because they duplicate the latent width and sample rate. - """ - source_dir = os.path.join(checkpoint_path, "audio_vae") - with open(os.path.join(source_dir, "metadata.json")) as f: - kwargs = json.load(f)["metadata"]["kwargs"] - with open(os.path.join(source_dir, "config.json")) as f: - wrapper_config = json.load(f) - - if kwargs["decoder_type"] != "bigvgan": - raise ValueError(f"Only the BigVGAN decoder is supported, got {kwargs['decoder_type']!r}.") - if not kwargs["attn_proj"]: - raise ValueError("The audio VAE is expected to carry the causal-attention latent projection.") - latent_channels = kwargs["vae_latent_channels"] - if wrapper_config["latent_channels"] != latent_channels: - raise ValueError( - f"latent width disagreement: metadata.json says {latent_channels}, " - f"config.json says {wrapper_config['latent_channels']}." - ) - if wrapper_config["sample_rate"] != kwargs["sample_rate"]: - raise ValueError( - f"sample rate disagreement: metadata.json says {kwargs['sample_rate']}, " - f"config.json says {wrapper_config['sample_rate']}." - ) - for key in ("latents_mean", "latents_std"): - if len(wrapper_config[key]) != latent_channels: - raise KeyError(f"{source_dir}/config.json `{key}` does not have {latent_channels} entries.") - - return { - "encoder_dim": kwargs["encoder_dim"], - "encoder_rates": kwargs["encoder_rates"], - "latent_dim": kwargs["latent_dim"], - "latent_channels": latent_channels, - "decoder_dim": kwargs["decoder_dim"], - "decoder_rates": kwargs["decoder_rates"], - # The reference's two hardcoded BigVGAN tables (16 kHz and 32 kHz) both pair rate `u` with kernel - # `2u` for even `u` and `2u - 1` for odd `u`, i.e. [5, 5, 2, ...] -> [9, 9, 4, ...]. - "decoder_kernel_sizes": [2 * rate - (rate % 2) for rate in kwargs["decoder_rates"]], - **MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG, - # Renamed from the original `sample_rate` to the diffusers audio convention. - "sampling_rate": kwargs["sample_rate"], - "latents_mean": wrapper_config["latents_mean"], - "latents_std": wrapper_config["latents_std"], - } - - -def convert_audio_vae(checkpoint_path: str, output_path: str, diffusers_version: str) -> None: - """Convert the audio VAE and emit its config. - - The mapping is an identity: `AutoencoderKLMiniMaxH3Audio` reproduces the original module tree name for name, - including `torch.nn.utils.weight_norm`'s `weight_g` / `weight_v` spelling and the Kaiser-window `filter` buffers of - the anti-aliased activations. The keys are therefore only *validated* against a freshly built model, not renamed. - """ - config = get_audio_vae_config(checkpoint_path) - expected_keys = set(AutoencoderKLMiniMaxH3Audio(**config).state_dict()) - - weights_path = os.path.join(checkpoint_path, "audio_vae", "model.safetensors") - state_dict: dict[str, torch.Tensor] = {} - total_size = 0 - with safe_open(weights_path, framework="pt", device="cpu") as f: - for key in f.keys(): - if key not in expected_keys: - raise KeyError(f"Unexpected key in {weights_path}: {key}") - tensor = f.get_tensor(key) - if tensor.dtype != torch.float32: - raise ValueError(f"{key}: expected torch.float32, got {tensor.dtype}.") - state_dict[key] = tensor - total_size += tensor.numel() * tensor.element_size() - - missing = sorted(expected_keys - set(state_dict)) - if missing: - raise KeyError(f"{len(missing)} key(s) missing from {weights_path}, e.g. {missing[:5]}.") - - os.makedirs(output_path, exist_ok=True) - save_file( - state_dict, - os.path.join(output_path, "diffusion_pytorch_model.safetensors"), - metadata={"format": "pt"}, - ) - with open(os.path.join(output_path, "config.json"), "w") as f: - json.dump( - {"_class_name": "AutoencoderKLMiniMaxH3Audio", "_diffusers_version": diffusers_version, **config}, - f, - indent=2, - ) - - print( - f"audio_vae: {len(state_dict)} keys carried over unchanged, {total_size / 1024**3:.2f} GiB " - f"({config['sampling_rate'] // math.prod(config['encoder_rates'])} latents/s, " - f"latents_mean/latents_std: {len(config['latents_mean'])}/{len(config['latents_std'])} channels)." - ) - - -def write_scheduler_configs(checkpoint_path: str, output_path: str, diffusers_version: str) -> None: - """Emit the two `MiniMaxH3Scheduler` configs, one per modality. - - The source `model_index.json` leaves `scheduler` null and instead carries the schedule constants in its - `_minimax_h3.sigma_shift_scales` block. The sigma shift is the only per-modality difference, so it becomes two - scheduler folders holding the same class at different `shift` values. - """ - with open(os.path.join(checkpoint_path, "model_index.json")) as f: - shift_scales = json.load(f)["_minimax_h3"]["sigma_shift_scales"] - - for folder, modality in (("scheduler", "video"), ("audio_scheduler", "audio")): - folder_path = os.path.join(output_path, folder) - os.makedirs(folder_path, exist_ok=True) - with open(os.path.join(folder_path, "scheduler_config.json"), "w") as f: - json.dump( - { - "_class_name": "MiniMaxH3Scheduler", - "_diffusers_version": diffusers_version, - "shift": float(shift_scales[modality]), - }, - f, - indent=2, - ) - print(f"scheduler: shift={shift_scales['video']} (video), audio_scheduler: shift={shift_scales['audio']} (audio).") - - -def read_safetensors_header(path: str) -> dict[str, Any]: - """Read the metadata header of a safetensors file without touching the tensor payload.""" - with open(path, "rb") as f: - header_size = struct.unpack(" None: - plan = get_transformer_key_plan(config) - - transformer_dir = os.path.join(checkpoint_path, "transformer") - shards = sorted(glob.glob(os.path.join(transformer_dir, "*.safetensors"))) - header: dict[str, Any] = {} - for shard in shards: - header.update(read_safetensors_header(shard)) - if shards: - print(f"Read headers of {len(shards)} shard(s) in {transformer_dir}: {len(header)} keys present.\n") - else: - print(f"No shards found under {transformer_dir}; validating the plan against the config only.\n") - - print(f"{'original key':<48} {'->':^4} {'diffusers key':<52} {'shape':<24} dtype") - print("-" * 150) - num_target_keys = 0 - shape_mismatches: list[str] = [] - for source_key, targets in plan.items(): - present = source_key in header - if not targets: - print(f"{source_key:<48} {'-x':^4} {'(dropped, recomputed by the port)':<52}") - continue - for index, (target_key, shape) in enumerate(targets): - num_target_keys += 1 - expected_dtype = "F32" if source_key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES) else "BF16" - if present: - actual_dtype = header[source_key]["dtype"] - actual_shape = header[source_key]["shape"] - if index == 0 and len(targets) == 1 and actual_shape != shape: - shape_mismatches.append(f"{source_key}: header {actual_shape} != planned {shape}") - if actual_dtype != expected_dtype: - shape_mismatches.append(f"{source_key}: header dtype {actual_dtype} != expected {expected_dtype}") - marker = "->" - else: - marker = "->?" - left = source_key if index == 0 else "" - print(f"{left:<48} {marker:^4} {target_key:<52} {str(shape):<24} {expected_dtype}") - - missing = [key for key in plan if key not in header] - unexpected = [key for key in header if key not in plan] - - print("\n" + "=" * 150) - print(f"planned original keys : {len(plan)}") - print(f"planned diffusers keys: {num_target_keys}") - print( - f"dropped original keys : {len(MINIMAX_H3_TRANSFORMER_DROPPED_KEYS)} {list(MINIMAX_H3_TRANSFORMER_DROPPED_KEYS)}" - ) - print(f"fp32 diffusers keys : {sum(1 for key in plan if key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES))}") - if shards: - print(f"keys present in shards: {len(header)}") - print(f"planned but absent : {len(missing)}" + (" (shards still downloading?)" if missing else "")) - print(f"present but unplanned : {len(unexpected)}") - if unexpected: - print(f" {unexpected}") - print(f"header disagreements : {len(shape_mismatches)}") - for line in shape_mismatches: - print(f" {line}") - total_bytes = sum( - (4 if source_key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES) else 2) * torch.Size(shape).numel() - for source_key, targets in plan.items() - for _, shape in targets - ) - print(f"total output bytes : {total_bytes} ({total_bytes / 1024**3:.2f} GiB)") - - -def convert_transformer(checkpoint_path: str, output_path: str, config: dict[str, Any], max_shard_size: int) -> None: - plan = get_transformer_key_plan(config) - transformer_dir = os.path.join(checkpoint_path, "transformer") - shards = sorted(glob.glob(os.path.join(transformer_dir, "*.safetensors"))) - if not shards: - raise FileNotFoundError(f"No `*.safetensors` shards found under {transformer_dir}.") - - os.makedirs(output_path, exist_ok=True) - weight_map: dict[str, str] = {} - total_size = 0 - written: list[str] = [] - buffer: dict[str, torch.Tensor] = {} - buffer_size = 0 - seen_source_keys: set[str] = set() - - def flush() -> None: - nonlocal buffer, buffer_size - if not buffer: - return - path = os.path.join(output_path, f".tmp-shard-{len(written):05d}.safetensors") - save_file(buffer, path, metadata={"format": "pt"}) - for key in buffer: - weight_map[key] = path - written.append(path) - buffer = {} - buffer_size = 0 - - for shard in shards: - # `safe_open` memory-maps the file, so only the tensor being read is materialized. - with safe_open(shard, framework="pt", device="cpu") as f: - for source_key in f.keys(): - if source_key not in plan: - raise KeyError(f"Unexpected key in {os.path.basename(shard)}: {source_key}") - seen_source_keys.add(source_key) - source_tensor = f.get_tensor(source_key) - if source_key.endswith(".attn.qkv_proj.weight"): - # Raw shards store fused QKV per-head interleaved; normalize to the reference's - # `[q_all; k_all; v_all]` layout (the same transform the reference applies at load time) so - # `convert_transformer_key` sees its state-dict-layout contract. The composition is bit-identical - # to de-interleaving the raw tensor directly. - source_tensor = reorder_interleaved_qkv( - source_tensor, config["num_attention_heads"], config["attention_head_dim"] - ) - for target_key, tensor in convert_transformer_key(source_key, source_tensor, config): - expected_dtype = ( - torch.float32 if source_key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES) else torch.bfloat16 - ) - if tensor.dtype != expected_dtype: - raise ValueError(f"{source_key}: expected {expected_dtype}, got {tensor.dtype}.") - buffer[target_key] = tensor - buffer_size += tensor.numel() * tensor.element_size() - total_size += tensor.numel() * tensor.element_size() - if buffer_size >= max_shard_size: - flush() - flush() - - missing = sorted(set(plan) - seen_source_keys) - if missing: - raise KeyError(f"{len(missing)} planned key(s) missing from the checkpoint, e.g. {missing[:5]}.") - - # The shard count is only known once every source shard has been streamed, so the files are written under - # provisional names and renamed here. - renames = { - path: os.path.join(output_path, f"diffusion_pytorch_model-{i + 1:05d}-of-{len(written):05d}.safetensors") - for i, path in enumerate(written) - } - for old, new in renames.items(): - os.rename(old, new) - index = { - "metadata": {"total_size": total_size}, - "weight_map": {key: os.path.basename(renames[path]) for key, path in weight_map.items()}, - } - with open(os.path.join(output_path, SAFE_WEIGHTS_INDEX_NAME), "w") as f: - json.dump(index, f, indent=2, sort_keys=True) - - print( - f"transformer: {len(seen_source_keys)} original keys -> {len(weight_map)} diffusers keys " - f"in {len(written)} shard(s), {total_size / 1024**3:.2f} GiB." - ) - - -def write_transformer_config(output_path: str, config: dict[str, Any], diffusers_version: str) -> None: - os.makedirs(output_path, exist_ok=True) - with open(os.path.join(output_path, "config.json"), "w") as f: - json.dump( - {"_class_name": "MiniMaxH3Transformer3DModel", "_diffusers_version": diffusers_version, **config}, - f, - indent=2, - ) - - -# The components a MiniMax-H3 repository holds, and the class each one loads as. `video_processor` is absent: the -# blocks create it from config rather than loading it. -MINIMAX_H3_COMPONENTS = { - # The source names a checkpoint-local wrapper class (`MiniMaxH3Qwen3VLHFEncoder`); the conditioner is the - # released Qwen3-VL, read at its 50th decoder layer with its language-model head unused. - "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"], - "tokenizer": ["transformers", "Qwen2TokenizerFast"], - "processor": ["transformers", "Qwen3VLProcessor"], - # Renamed to the diffusers audio/video VAE convention (see `LTX2Pipeline`). - "vae": ["diffusers", "AutoencoderKLMiniMaxH3"], - "audio_vae": ["diffusers", "AutoencoderKLMiniMaxH3Audio"], - "transformer": ["diffusers", "MiniMaxH3Transformer3DModel"], - # One repository holds both checkpoint partitions: `transformer/` serves `MiniMaxH3Blocks` (`t2va` / `fl2va`) and - # `transformer_ref/` serves `MiniMaxH3Ref2VABlocks`, while every other component is shared and converted once. - "transformer_ref": ["diffusers", "MiniMaxH3Transformer3DModel"], - # The source leaves `scheduler` null. MiniMax-H3 samples with Euler at eta=0 over shifted flow-matching sigmas, at - # a different shift per modality, so it needs two scheduler entries (see `write_scheduler_configs`). - "scheduler": ["diffusers", "MiniMaxH3Scheduler"], - "audio_scheduler": ["diffusers", "MiniMaxH3Scheduler"], -} - - -def write_model_index(output_path: str, repo_id: str, diffusers_version: str) -> None: - """Emit `modular_model_index.json`, the only index a MiniMax-H3 repository carries. - - MiniMax-H3 is integrated as Modular Diffusers blocks only, so there is no `model_index.json`: a modular repository - declares one entry per component with its full loading spec rather than just its class, and a blockset then fetches - exactly the subfolders it declares. That is what lets one repository hold both transformer partitions, and the - original checkpoint folders next to the converted ones, without either half pulling the rest down. - - `_class_name` and `_blocks_class_name` name the `t2va` / `fl2va` half, which is what - `ModularPipeline.from_pretrained` resolves to. The `ref2va` half reads the very same file through - `MiniMaxH3Ref2VABlocks().init_pipeline(repo_id)`. - - The component map is the static one above, so this needs no source checkpoint: an index can be regenerated for a - repository that is already published. - """ - modular_index = { - "_class_name": "MiniMaxH3ModularPipeline", - "_diffusers_version": diffusers_version, - "_blocks_class_name": "MiniMaxH3Blocks", - } - for name, (library, class_name) in MINIMAX_H3_COMPONENTS.items(): - modular_index[name] = [ - library, - class_name, - { - "type_hint": [library, class_name], - "pretrained_model_name_or_path": repo_id, - "subfolder": name, - "variant": None, - "revision": None, - }, - ] - os.makedirs(output_path, exist_ok=True) - with open(os.path.join(output_path, "modular_model_index.json"), "w") as f: - json.dump(modular_index, f, indent=2) - print(f"modular_model_index.json: {len(MINIMAX_H3_COMPONENTS)} components load from {repo_id}.") - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--checkpoint_path", - type=str, - required=True, - help="Local path to an original MiniMax-H3 variant folder (the one holding `transformer/`, `audio_vae/`, ...).", - ) - parser.add_argument("--output_path", type=str, required=True, help="Where the diffusers checkpoint is written.") - parser.add_argument( - "--modular_repo_id", - type=str, - default=None, - help=( - "Repository the component entries of `modular_model_index.json` point at. Defaults to `--output_path`, so " - "pass the Hub id the checkpoint is published under. Every entry carries its own loading spec, so a " - "blockset fetches exactly the subfolders it declares out of that repository." - ), - ) - parser.add_argument( - "--version", - type=str, - default="h3", - choices=["h3", "test"], - help="`test` emits the tiny config used for fixtures.", - ) - parser.add_argument( - "--max_shard_size", - type=int, - default=5 * 1024**3, - help="Maximum size of an output safetensors shard, in bytes.", - ) - parser.add_argument( - "--dry_run", - action="store_true", - help="Print the full planned key mapping (and cross-check any shard headers already present) without writing.", - ) - return parser.parse_args() - - -def main(args): - from diffusers import __version__ as diffusers_version - - config = MINIMAX_H3_TEST_TRANSFORMER_CONFIG if args.version == "test" else MINIMAX_H3_TRANSFORMER_CONFIG - - if args.dry_run: - dry_run(args.checkpoint_path, config) - return - - transformer_path = os.path.join(args.output_path, "transformer") - convert_transformer(args.checkpoint_path, transformer_path, config, args.max_shard_size) - write_transformer_config(transformer_path, config, diffusers_version) - video_vae_config = MINIMAX_H3_TEST_VIDEO_VAE_CONFIG if args.version == "test" else MINIMAX_H3_VIDEO_VAE_CONFIG - convert_video_vae( - args.checkpoint_path, - os.path.join(args.output_path, "vae"), - video_vae_config, - diffusers_version, - args.max_shard_size, - ) - convert_audio_vae(args.checkpoint_path, os.path.join(args.output_path, "audio_vae"), diffusers_version) - write_scheduler_configs(args.checkpoint_path, args.output_path, diffusers_version) - write_model_index(args.output_path, args.modular_repo_id or args.output_path, diffusers_version) - - -if __name__ == "__main__": - main(get_args()) diff --git a/scripts/convert_minimax_music3_to_diffusers.py b/scripts/convert_minimax_music3_to_diffusers.py deleted file mode 100644 index c6cef069c27e..000000000000 --- a/scripts/convert_minimax_music3_to_diffusers.py +++ /dev/null @@ -1,269 +0,0 @@ -# Conversion script for MiniMax Music 3 (https://huggingface.co/MiniMaxAI/MiniMax-Music3). -# -# Original checkpoint layout: -# flowmatching_vae.pth flow-matching DiT + condition projection -# dav.pth Flow-VAE (DAC-style) decoder -# qwen_7B/qwen_7B/ Qwen3 backbone + audio embedding + RVQ depth decoder (sharded safetensors) -# qwen_7B/qwen3-8B-tokenizer-music/ music tokenizer -# -# Usage: -# python scripts/convert_minimax_music3_to_diffusers.py \ -# --checkpoint_dir MiniMaxAI/MiniMax-Music3 --output_path ./minimax-music3-diffusers - -import argparse -import json -import os - -import torch -from safetensors.torch import load_file - -from diffusers import ( - FlowMatchEulerDiscreteScheduler, - MiniMaxMusic3ConditionEncoder, - MiniMaxMusic3RVQDepthDecoder, - MiniMaxMusic3Transformer1DModel, - MiniMaxMusic3Vocoder, -) - - -def load_dit_state_dict(checkpoint_dir: str) -> dict: - return torch.load(os.path.join(checkpoint_dir, "flowmatching_vae.pth"), map_location="cpu", weights_only=True) - - -def load_dav_state_dict(checkpoint_dir: str) -> dict: - return torch.load(os.path.join(checkpoint_dir, "dav.pth"), map_location="cpu", weights_only=True) - - -def load_qwen_state_dict(checkpoint_dir: str) -> dict: - qwen_dir = os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B") - with open(os.path.join(qwen_dir, "model.safetensors.index.json")) as f: - index = json.load(f) - state_dict = {} - for filename in sorted(set(index["weight_map"].values())): - state_dict.update(load_file(os.path.join(qwen_dir, filename), device="cpu")) - return state_dict - - -def convert_transformer(dit_state_dict: dict) -> MiniMaxMusic3Transformer1DModel: - prefix = "diffusion_transformer." - converted = { - "time_proj.weight": dit_state_dict[prefix + "timestep_features.weight"], - "time_embed.linear_1.weight": dit_state_dict[prefix + "to_timestep_embed.0.weight"], - "time_embed.linear_1.bias": dit_state_dict[prefix + "to_timestep_embed.0.bias"], - "time_embed.linear_2.weight": dit_state_dict[prefix + "to_timestep_embed.2.weight"], - "time_embed.linear_2.bias": dit_state_dict[prefix + "to_timestep_embed.2.bias"], - "preprocess_conv.weight": dit_state_dict[prefix + "preprocess_conv.weight"], - "postprocess_conv.weight": dit_state_dict[prefix + "postprocess_conv.weight"], - "proj_in.weight": dit_state_dict[prefix + "transformer.project_in.weight"], - "proj_out.weight": dit_state_dict[prefix + "transformer.project_out.weight"], - } - num_layers = 0 - while prefix + f"transformer.layers.{num_layers}.pre_norm.gamma" in dit_state_dict: - num_layers += 1 - for i in range(num_layers): - original = prefix + f"transformer.layers.{i}." - target = f"transformer_blocks.{i}." - converted[target + "norm1.weight"] = dit_state_dict[original + "pre_norm.gamma"] - converted[target + "norm1.bias"] = dit_state_dict[original + "pre_norm.beta"] - query, key, value = dit_state_dict[original + "self_attn.to_qkv.weight"].chunk(3, dim=0) - converted[target + "attn.to_q.weight"] = query - converted[target + "attn.to_k.weight"] = key - converted[target + "attn.to_v.weight"] = value - converted[target + "attn.to_out.0.weight"] = dit_state_dict[original + "self_attn.to_out.weight"] - converted[target + "norm2.weight"] = dit_state_dict[original + "ff_norm.gamma"] - converted[target + "norm2.bias"] = dit_state_dict[original + "ff_norm.beta"] - converted[target + "ff_in.weight"] = dit_state_dict[original + "ff.ff.0.proj.weight"] - converted[target + "ff_in.bias"] = dit_state_dict[original + "ff.ff.0.proj.bias"] - converted[target + "ff_out.weight"] = dit_state_dict[original + "ff.ff.2.weight"] - converted[target + "ff_out.bias"] = dit_state_dict[original + "ff.ff.2.bias"] - - with torch.device("meta"): - transformer = MiniMaxMusic3Transformer1DModel(num_layers=num_layers) - transformer.load_state_dict(converted, strict=True, assign=True) - return transformer - - -def convert_condition_encoder(dit_state_dict: dict) -> MiniMaxMusic3ConditionEncoder: - converted = { - "layer_weight_logits": dit_state_dict["cond_layer_logits"], - "layer_scale": dit_state_dict["cond_layer_scale"], - "proj.weight": dit_state_dict["latent_conditioners.0.weight"], - "proj.bias": dit_state_dict["latent_conditioners.0.bias"], - } - with torch.device("meta"): - condition_encoder = MiniMaxMusic3ConditionEncoder() - condition_encoder.load_state_dict(converted, strict=True, assign=True) - return condition_encoder - - -def convert_vocoder(dav_state_dict: dict) -> MiniMaxMusic3Vocoder: - converted = { - "dec_in_proj.weight": dav_state_dict["dec_in_proj.weight"], - "dec_in_proj.bias": dav_state_dict["dec_in_proj.bias"], - } - # The reference decoder is one nn.Sequential: [conv_in, block*4, snake, conv_out, tanh]. - for suffix in ("weight_g", "weight_v", "bias"): - converted[f"conv_in.{suffix}"] = dav_state_dict[f"decoder.model.0.{suffix}"] - converted[f"conv_out.{suffix}"] = dav_state_dict[f"decoder.model.6.{suffix}"] - converted["snake_out.alpha"] = dav_state_dict["decoder.model.5.alpha"] - for block_index in range(4): - original = f"decoder.model.{block_index + 1}.block." - target = f"blocks.{block_index}." - converted[target + "snake1.alpha"] = dav_state_dict[original + "0.alpha"] - for suffix in ("weight_g", "weight_v", "bias"): - converted[target + f"conv_t1.{suffix}"] = dav_state_dict[original + f"1.{suffix}"] - for unit_index, unit_name in ((2, "res_unit1"), (3, "res_unit2"), (4, "res_unit3")): - converted[target + f"{unit_name}.snake1.alpha"] = dav_state_dict[original + f"{unit_index}.block.0.alpha"] - converted[target + f"{unit_name}.snake2.alpha"] = dav_state_dict[original + f"{unit_index}.block.2.alpha"] - for suffix in ("weight_g", "weight_v", "bias"): - converted[target + f"{unit_name}.conv1.{suffix}"] = dav_state_dict[ - original + f"{unit_index}.block.1.{suffix}" - ] - converted[target + f"{unit_name}.conv2.{suffix}"] = dav_state_dict[ - original + f"{unit_index}.block.3.{suffix}" - ] - - vocoder = MiniMaxMusic3Vocoder() - vocoder.load_state_dict(converted, strict=True) - return vocoder - - -def convert_rvq_depth_decoder(qwen_state_dict: dict, model_config: dict) -> MiniMaxMusic3RVQDepthDecoder: - prefix = "model.audio_decoder." - converted = { - "audio_embeddings.weight": qwen_state_dict["model.audio_extra_embedding.weight"], - "projection.weight": qwen_state_dict[prefix + "projection.weight"], - "pos_embedding.weight": qwen_state_dict[prefix + "pos_embedding.weight"], - "norm.weight": qwen_state_dict[prefix + "norm.weight"], - } - num_codebooks = int(model_config["audio_num_codebooks"]) - for i in range(num_codebooks - 1): - converted[f"audio_heads.{i}.weight"] = qwen_state_dict[prefix + f"audio_heads.{i}.weight"] - num_layers = int(model_config["decoder_num_layers"]) - for i in range(num_layers): - original = prefix + f"layers.{i}." - target = f"layers.{i}." - converted[target + "input_layernorm.weight"] = qwen_state_dict[original + "input_layernorm.weight"] - converted[target + "post_attention_layernorm.weight"] = qwen_state_dict[ - original + "post_attention_layernorm.weight" - ] - converted[target + "attn.to_q.weight"] = qwen_state_dict[original + "self_attn.q_proj.weight"] - converted[target + "attn.to_k.weight"] = qwen_state_dict[original + "self_attn.k_proj.weight"] - converted[target + "attn.to_v.weight"] = qwen_state_dict[original + "self_attn.v_proj.weight"] - converted[target + "attn.to_out.weight"] = qwen_state_dict[original + "self_attn.o_proj.weight"] - for proj in ("gate_proj", "up_proj", "down_proj"): - converted[target + proj + ".weight"] = qwen_state_dict[original + f"mlp.{proj}.weight"] - - with torch.device("meta"): - rvq_depth_decoder = MiniMaxMusic3RVQDepthDecoder( - hidden_size=int(model_config["hidden_size"]), - num_layers=num_layers, - num_attention_heads=int(model_config["decoder_num_heads"]), - intermediate_size=int(model_config["decoder_intermediate_size"]), - audio_vocab_size=int(model_config["audio_vocab_size"]), - num_codebooks=num_codebooks, - ) - rvq_depth_decoder.load_state_dict(converted, strict=True, assign=True) - return rvq_depth_decoder - - -def convert_language_model(qwen_state_dict: dict, model_config: dict): - from transformers import Qwen3Config, Qwen3ForCausalLM - - config = Qwen3Config( - vocab_size=model_config["vocab_size"], - hidden_size=model_config["hidden_size"], - intermediate_size=model_config["intermediate_size"], - num_hidden_layers=model_config["num_hidden_layers"], - num_attention_heads=model_config["num_attention_heads"], - num_key_value_heads=model_config["num_key_value_heads"], - head_dim=model_config["head_dim"], - max_position_embeddings=model_config.get("max_position_embeddings", 10240), - rope_theta=model_config.get("rope_theta", 1000000), - rms_norm_eps=model_config.get("rms_norm_eps", 1e-6), - tie_word_embeddings=model_config.get("tie_word_embeddings", False), - ) - backbone_state_dict = { - key: value - for key, value in qwen_state_dict.items() - if not key.startswith(("model.audio_extra_embedding", "model.audio_decoder.")) - } - with torch.device("meta"): - language_model = Qwen3ForCausalLM(config) - language_model.load_state_dict(backbone_state_dict, strict=True, assign=True) - return language_model - - -def main(args): - checkpoint_dir = args.checkpoint_dir - if not os.path.isdir(checkpoint_dir): - from huggingface_hub import snapshot_download - - checkpoint_dir = snapshot_download(checkpoint_dir) - - with open(os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B", "config.json")) as f: - model_config = json.load(f) - - dit_state_dict = load_dit_state_dict(checkpoint_dir) - transformer = convert_transformer(dit_state_dict).to(args.dtype) - condition_encoder = convert_condition_encoder(dit_state_dict).to(args.dtype) - del dit_state_dict - vocoder = convert_vocoder(load_dav_state_dict(checkpoint_dir)).to(args.dtype) - - qwen_state_dict = load_qwen_state_dict(checkpoint_dir) - rvq_depth_decoder = convert_rvq_depth_decoder(qwen_state_dict, model_config).to(torch.bfloat16) - language_model = convert_language_model(qwen_state_dict, model_config) - del qwen_state_dict - - from transformers import AutoTokenizer - - from diffusers import MiniMaxMusic3Blocks - - tokenizer = AutoTokenizer.from_pretrained(os.path.join(checkpoint_dir, "qwen_7B", "qwen3-8B-tokenizer-music")) - # num_train_timesteps=1 keeps `scheduler.timesteps` equal to the flow-matching time in [0, 1] that the - # transformer's Fourier embedding expects. - scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0, invert_sigmas=True) - - pipeline = MiniMaxMusic3Blocks().init_pipeline() - pipeline.update_components( - language_model=language_model, - rvq_depth_decoder=rvq_depth_decoder, - condition_encoder=condition_encoder, - transformer=transformer, - vocoder=vocoder, - tokenizer=tokenizer, - scheduler=scheduler, - ) - pipeline.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") - - # save_pretrained bakes the local output path into the modular index's loading specs; point them at the - # Hub repo the components will be uploaded to instead. - index_path = os.path.join(args.output_path, "modular_model_index.json") - with open(index_path) as f: - index = json.load(f) - for entry in index.values(): - if isinstance(entry, list) and len(entry) == 3 and isinstance(entry[2], dict): - if entry[2].get("pretrained_model_name_or_path") == args.output_path: - entry[2]["pretrained_model_name_or_path"] = args.repo_id - with open(index_path, "w") as f: - json.dump(index, f, indent=2, sort_keys=True) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - "--checkpoint_dir", - type=str, - default="MiniMaxAI/MiniMax-Music3", - help="Local directory or Hugging Face Hub repo id of the original checkpoint.", - ) - parser.add_argument("--output_path", type=str, required=True) - parser.add_argument( - "--repo_id", - type=str, - default="MiniMaxAI/MiniMax-Music3", - help="Hub repo id the converted components will live in (written into the modular index loading specs).", - ) - parser.add_argument("--dtype", type=lambda name: getattr(torch, name), default="float32") - args = parser.parse_args() - main(args) diff --git a/scripts/convert_mochi_to_diffusers.py b/scripts/convert_mochi_to_diffusers.py deleted file mode 100644 index 64e4f69eac17..000000000000 --- a/scripts/convert_mochi_to_diffusers.py +++ /dev/null @@ -1,463 +0,0 @@ -import argparse -from contextlib import nullcontext - -import torch -from accelerate import init_empty_weights -from safetensors.torch import load_file -from transformers import T5EncoderModel, T5Tokenizer - -from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel -from diffusers.utils.import_utils import is_accelerate_available - - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - -TOKENIZER_MAX_LENGTH = 256 - -parser = argparse.ArgumentParser() -parser.add_argument("--transformer_checkpoint_path", default=None, type=str) -parser.add_argument("--vae_encoder_checkpoint_path", default=None, type=str) -parser.add_argument("--vae_decoder_checkpoint_path", default=None, type=str) -parser.add_argument("--output_path", required=True, type=str) -parser.add_argument("--push_to_hub", action="store_true", default=False, help="Whether to push to HF Hub after saving") -parser.add_argument("--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory") -parser.add_argument("--dtype", type=str, default=None) - -args = parser.parse_args() - - -# This is specific to `AdaLayerNormContinuous`: -# Diffusers implementation split the linear projection into the scale, shift while Mochi split it into shift, scale -def swap_scale_shift(weight, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def swap_proj_gate(weight): - proj, gate = weight.chunk(2, dim=0) - new_weight = torch.cat([gate, proj], dim=0) - return new_weight - - -def convert_mochi_transformer_checkpoint_to_diffusers(ckpt_path): - original_state_dict = load_file(ckpt_path, device="cpu") - new_state_dict = {} - - # Convert patch_embed - new_state_dict["patch_embed.proj.weight"] = original_state_dict.pop("x_embedder.proj.weight") - new_state_dict["patch_embed.proj.bias"] = original_state_dict.pop("x_embedder.proj.bias") - - # Convert time_embed - new_state_dict["time_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop("t_embedder.mlp.0.weight") - new_state_dict["time_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop("t_embedder.mlp.0.bias") - new_state_dict["time_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop("t_embedder.mlp.2.weight") - new_state_dict["time_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop("t_embedder.mlp.2.bias") - new_state_dict["time_embed.pooler.to_kv.weight"] = original_state_dict.pop("t5_y_embedder.to_kv.weight") - new_state_dict["time_embed.pooler.to_kv.bias"] = original_state_dict.pop("t5_y_embedder.to_kv.bias") - new_state_dict["time_embed.pooler.to_q.weight"] = original_state_dict.pop("t5_y_embedder.to_q.weight") - new_state_dict["time_embed.pooler.to_q.bias"] = original_state_dict.pop("t5_y_embedder.to_q.bias") - new_state_dict["time_embed.pooler.to_out.weight"] = original_state_dict.pop("t5_y_embedder.to_out.weight") - new_state_dict["time_embed.pooler.to_out.bias"] = original_state_dict.pop("t5_y_embedder.to_out.bias") - new_state_dict["time_embed.caption_proj.weight"] = original_state_dict.pop("t5_yproj.weight") - new_state_dict["time_embed.caption_proj.bias"] = original_state_dict.pop("t5_yproj.bias") - - # Convert transformer blocks - num_layers = 48 - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - old_prefix = f"blocks.{i}." - - # norm1 - new_state_dict[block_prefix + "norm1.linear.weight"] = original_state_dict.pop(old_prefix + "mod_x.weight") - new_state_dict[block_prefix + "norm1.linear.bias"] = original_state_dict.pop(old_prefix + "mod_x.bias") - if i < num_layers - 1: - new_state_dict[block_prefix + "norm1_context.linear.weight"] = original_state_dict.pop( - old_prefix + "mod_y.weight" - ) - new_state_dict[block_prefix + "norm1_context.linear.bias"] = original_state_dict.pop( - old_prefix + "mod_y.bias" - ) - else: - new_state_dict[block_prefix + "norm1_context.linear_1.weight"] = original_state_dict.pop( - old_prefix + "mod_y.weight" - ) - new_state_dict[block_prefix + "norm1_context.linear_1.bias"] = original_state_dict.pop( - old_prefix + "mod_y.bias" - ) - - # Visual attention - qkv_weight = original_state_dict.pop(old_prefix + "attn.qkv_x.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - new_state_dict[block_prefix + "attn1.to_q.weight"] = q - new_state_dict[block_prefix + "attn1.to_k.weight"] = k - new_state_dict[block_prefix + "attn1.to_v.weight"] = v - new_state_dict[block_prefix + "attn1.norm_q.weight"] = original_state_dict.pop( - old_prefix + "attn.q_norm_x.weight" - ) - new_state_dict[block_prefix + "attn1.norm_k.weight"] = original_state_dict.pop( - old_prefix + "attn.k_norm_x.weight" - ) - new_state_dict[block_prefix + "attn1.to_out.0.weight"] = original_state_dict.pop( - old_prefix + "attn.proj_x.weight" - ) - new_state_dict[block_prefix + "attn1.to_out.0.bias"] = original_state_dict.pop(old_prefix + "attn.proj_x.bias") - - # Context attention - qkv_weight = original_state_dict.pop(old_prefix + "attn.qkv_y.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - new_state_dict[block_prefix + "attn1.add_q_proj.weight"] = q - new_state_dict[block_prefix + "attn1.add_k_proj.weight"] = k - new_state_dict[block_prefix + "attn1.add_v_proj.weight"] = v - new_state_dict[block_prefix + "attn1.norm_added_q.weight"] = original_state_dict.pop( - old_prefix + "attn.q_norm_y.weight" - ) - new_state_dict[block_prefix + "attn1.norm_added_k.weight"] = original_state_dict.pop( - old_prefix + "attn.k_norm_y.weight" - ) - if i < num_layers - 1: - new_state_dict[block_prefix + "attn1.to_add_out.weight"] = original_state_dict.pop( - old_prefix + "attn.proj_y.weight" - ) - new_state_dict[block_prefix + "attn1.to_add_out.bias"] = original_state_dict.pop( - old_prefix + "attn.proj_y.bias" - ) - - # MLP - new_state_dict[block_prefix + "ff.net.0.proj.weight"] = swap_proj_gate( - original_state_dict.pop(old_prefix + "mlp_x.w1.weight") - ) - new_state_dict[block_prefix + "ff.net.2.weight"] = original_state_dict.pop(old_prefix + "mlp_x.w2.weight") - if i < num_layers - 1: - new_state_dict[block_prefix + "ff_context.net.0.proj.weight"] = swap_proj_gate( - original_state_dict.pop(old_prefix + "mlp_y.w1.weight") - ) - new_state_dict[block_prefix + "ff_context.net.2.weight"] = original_state_dict.pop( - old_prefix + "mlp_y.w2.weight" - ) - - # Output layers - new_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("final_layer.mod.weight"), dim=0 - ) - new_state_dict["norm_out.linear.bias"] = swap_scale_shift(original_state_dict.pop("final_layer.mod.bias"), dim=0) - new_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - new_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - - new_state_dict["pos_frequencies"] = original_state_dict.pop("pos_frequencies") - - print("Remaining Keys:", original_state_dict.keys()) - - return new_state_dict - - -def convert_mochi_vae_state_dict_to_diffusers(encoder_ckpt_path, decoder_ckpt_path): - encoder_state_dict = load_file(encoder_ckpt_path, device="cpu") - decoder_state_dict = load_file(decoder_ckpt_path, device="cpu") - new_state_dict = {} - - # ==== Decoder ===== - prefix = "decoder." - - # Convert conv_in - new_state_dict[f"{prefix}conv_in.weight"] = decoder_state_dict.pop("blocks.0.0.weight") - new_state_dict[f"{prefix}conv_in.bias"] = decoder_state_dict.pop("blocks.0.0.bias") - - # Convert block_in (MochiMidBlock3D) - for i in range(3): # layers_per_block[-1] = 3 - new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.0.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.0.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.weight"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.2.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.bias"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.2.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.3.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.3.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.weight"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.5.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.bias"] = decoder_state_dict.pop( - f"blocks.0.{i + 1}.stack.5.bias" - ) - - # Convert up_blocks (MochiUpBlock3D) - down_block_layers = [6, 4, 3] # layers_per_block[-2], layers_per_block[-3], layers_per_block[-4] - for block in range(3): - for i in range(down_block_layers[block]): - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.0.weight" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.0.bias" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.2.weight" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.2.bias" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.3.weight" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.3.bias" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.5.weight" - ) - new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"] = decoder_state_dict.pop( - f"blocks.{block + 1}.blocks.{i}.stack.5.bias" - ) - new_state_dict[f"{prefix}up_blocks.{block}.proj.weight"] = decoder_state_dict.pop( - f"blocks.{block + 1}.proj.weight" - ) - new_state_dict[f"{prefix}up_blocks.{block}.proj.bias"] = decoder_state_dict.pop( - f"blocks.{block + 1}.proj.bias" - ) - - # Convert block_out (MochiMidBlock3D) - for i in range(3): # layers_per_block[0] = 3 - new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.0.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.0.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.weight"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.2.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.bias"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.2.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.3.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.3.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.weight"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.5.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.bias"] = decoder_state_dict.pop( - f"blocks.4.{i}.stack.5.bias" - ) - - # Convert proj_out (Conv1x1 ~= nn.Linear) - new_state_dict[f"{prefix}proj_out.weight"] = decoder_state_dict.pop("output_proj.weight") - new_state_dict[f"{prefix}proj_out.bias"] = decoder_state_dict.pop("output_proj.bias") - - print("Remaining Decoder Keys:", decoder_state_dict.keys()) - - # ==== Encoder ===== - prefix = "encoder." - - new_state_dict[f"{prefix}proj_in.weight"] = encoder_state_dict.pop("layers.0.weight") - new_state_dict[f"{prefix}proj_in.bias"] = encoder_state_dict.pop("layers.0.bias") - - # Convert block_in (MochiMidBlock3D) - for i in range(3): # layers_per_block[0] = 3 - new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.0.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.0.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.weight"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.2.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.bias"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.2.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.3.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.3.bias" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.weight"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.5.weight" - ) - new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.bias"] = encoder_state_dict.pop( - f"layers.{i + 1}.stack.5.bias" - ) - - # Convert down_blocks (MochiDownBlock3D) - down_block_layers = [3, 4, 6] # layers_per_block[1], layers_per_block[2], layers_per_block[3] - for block in range(3): - new_state_dict[f"{prefix}down_blocks.{block}.conv_in.conv.weight"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.0.weight" - ) - new_state_dict[f"{prefix}down_blocks.{block}.conv_in.conv.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.0.bias" - ) - - for i in range(down_block_layers[block]): - # Convert resnets - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"] = ( - encoder_state_dict.pop(f"layers.{block + 4}.layers.{i + 1}.stack.0.weight") - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.0.bias" - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.2.weight" - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.2.bias" - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"] = ( - encoder_state_dict.pop(f"layers.{block + 4}.layers.{i + 1}.stack.3.weight") - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.3.bias" - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.5.weight" - ) - new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.stack.5.bias" - ) - - # Convert attentions - qkv_weight = encoder_state_dict.pop(f"layers.{block + 4}.layers.{i + 1}.attn_block.attn.qkv.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight"] = q - new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight"] = k - new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight"] = v - new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.attn_block.attn.out.weight" - ) - new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.attn_block.attn.out.bias" - ) - new_state_dict[f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.attn_block.norm.weight" - ) - new_state_dict[f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{block + 4}.layers.{i + 1}.attn_block.norm.bias" - ) - - # Convert block_out (MochiMidBlock3D) - for i in range(3): # layers_per_block[-1] = 3 - # Convert resnets - new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.0.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.0.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.2.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.2.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.3.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.3.bias" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.5.weight" - ) - new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.stack.5.bias" - ) - - # Convert attentions - qkv_weight = encoder_state_dict.pop(f"layers.{i + 7}.attn_block.attn.qkv.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - new_state_dict[f"{prefix}block_out.attentions.{i}.to_q.weight"] = q - new_state_dict[f"{prefix}block_out.attentions.{i}.to_k.weight"] = k - new_state_dict[f"{prefix}block_out.attentions.{i}.to_v.weight"] = v - new_state_dict[f"{prefix}block_out.attentions.{i}.to_out.0.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.attn_block.attn.out.weight" - ) - new_state_dict[f"{prefix}block_out.attentions.{i}.to_out.0.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.attn_block.attn.out.bias" - ) - new_state_dict[f"{prefix}block_out.norms.{i}.norm_layer.weight"] = encoder_state_dict.pop( - f"layers.{i + 7}.attn_block.norm.weight" - ) - new_state_dict[f"{prefix}block_out.norms.{i}.norm_layer.bias"] = encoder_state_dict.pop( - f"layers.{i + 7}.attn_block.norm.bias" - ) - - # Convert output layers - new_state_dict[f"{prefix}norm_out.norm_layer.weight"] = encoder_state_dict.pop("output_norm.weight") - new_state_dict[f"{prefix}norm_out.norm_layer.bias"] = encoder_state_dict.pop("output_norm.bias") - new_state_dict[f"{prefix}proj_out.weight"] = encoder_state_dict.pop("output_proj.weight") - - print("Remaining Encoder Keys:", encoder_state_dict.keys()) - - return new_state_dict - - -def main(args): - if args.dtype is None: - dtype = None - if args.dtype == "fp16": - dtype = torch.float16 - elif args.dtype == "bf16": - dtype = torch.bfloat16 - elif args.dtype == "fp32": - dtype = torch.float32 - else: - raise ValueError(f"Unsupported dtype: {args.dtype}") - - transformer = None - vae = None - - if args.transformer_checkpoint_path is not None: - converted_transformer_state_dict = convert_mochi_transformer_checkpoint_to_diffusers( - args.transformer_checkpoint_path - ) - transformer = MochiTransformer3DModel() - transformer.load_state_dict(converted_transformer_state_dict, strict=True) - if dtype is not None: - transformer = transformer.to(dtype=dtype) - - if args.vae_encoder_checkpoint_path is not None and args.vae_decoder_checkpoint_path is not None: - vae = AutoencoderKLMochi(latent_channels=12, out_channels=3) - converted_vae_state_dict = convert_mochi_vae_state_dict_to_diffusers( - args.vae_encoder_checkpoint_path, args.vae_decoder_checkpoint_path - ) - vae.load_state_dict(converted_vae_state_dict, strict=True) - if dtype is not None: - vae = vae.to(dtype=dtype) - - text_encoder_id = "google/t5-v1_1-xxl" - tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH) - text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir) - - # Apparently, the conversion does not work anymore without this :shrug: - for param in text_encoder.parameters(): - param.data = param.data.contiguous() - - pipe = MochiPipeline( - scheduler=FlowMatchEulerDiscreteScheduler(invert_sigmas=True), - vae=vae, - text_encoder=text_encoder, - tokenizer=tokenizer, - transformer=transformer, - ) - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", push_to_hub=args.push_to_hub) - - -if __name__ == "__main__": - main(args) diff --git a/scripts/convert_models_diffuser_to_diffusers.py b/scripts/convert_models_diffuser_to_diffusers.py deleted file mode 100644 index cc5321e33fe0..000000000000 --- a/scripts/convert_models_diffuser_to_diffusers.py +++ /dev/null @@ -1,100 +0,0 @@ -import json -import os - -import torch - -from diffusers import UNet1DModel - - -os.makedirs("hub/hopper-medium-v2/unet/hor32", exist_ok=True) -os.makedirs("hub/hopper-medium-v2/unet/hor128", exist_ok=True) - -os.makedirs("hub/hopper-medium-v2/value_function", exist_ok=True) - - -def unet(hor): - if hor == 128: - down_block_types = ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D") - block_out_channels = (32, 128, 256) - up_block_types = ("UpResnetBlock1D", "UpResnetBlock1D") - - elif hor == 32: - down_block_types = ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D") - block_out_channels = (32, 64, 128, 256) - up_block_types = ("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D") - model = torch.load(f"/Users/bglickenhaus/Documents/diffuser/temporal_unet-hopper-mediumv2-hor{hor}.torch") - state_dict = model.state_dict() - config = { - "down_block_types": down_block_types, - "block_out_channels": block_out_channels, - "up_block_types": up_block_types, - "layers_per_block": 1, - "use_timestep_embedding": True, - "out_block_type": "OutConv1DBlock", - "norm_num_groups": 8, - "downsample_each_block": False, - "in_channels": 14, - "out_channels": 14, - "extra_in_channels": 0, - "time_embedding_type": "positional", - "flip_sin_to_cos": False, - "freq_shift": 1, - "sample_size": 65536, - "mid_block_type": "MidResTemporalBlock1D", - "act_fn": "mish", - } - hf_value_function = UNet1DModel(**config) - print(f"length of state dict: {len(state_dict.keys())}") - print(f"length of value function dict: {len(hf_value_function.state_dict().keys())}") - mapping = dict(zip(model.state_dict().keys(), hf_value_function.state_dict().keys())) - for k, v in mapping.items(): - state_dict[v] = state_dict.pop(k) - hf_value_function.load_state_dict(state_dict) - - torch.save(hf_value_function.state_dict(), f"hub/hopper-medium-v2/unet/hor{hor}/diffusion_pytorch_model.bin") - with open(f"hub/hopper-medium-v2/unet/hor{hor}/config.json", "w") as f: - json.dump(config, f) - - -def value_function(): - config = { - "in_channels": 14, - "down_block_types": ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"), - "up_block_types": (), - "out_block_type": "ValueFunction", - "mid_block_type": "ValueFunctionMidBlock1D", - "block_out_channels": (32, 64, 128, 256), - "layers_per_block": 1, - "downsample_each_block": True, - "sample_size": 65536, - "out_channels": 14, - "extra_in_channels": 0, - "time_embedding_type": "positional", - "use_timestep_embedding": True, - "flip_sin_to_cos": False, - "freq_shift": 1, - "norm_num_groups": 8, - "act_fn": "mish", - } - - model = torch.load("/Users/bglickenhaus/Documents/diffuser/value_function-hopper-mediumv2-hor32.torch") - state_dict = model - hf_value_function = UNet1DModel(**config) - print(f"length of state dict: {len(state_dict.keys())}") - print(f"length of value function dict: {len(hf_value_function.state_dict().keys())}") - - mapping = dict(zip(state_dict.keys(), hf_value_function.state_dict().keys())) - for k, v in mapping.items(): - state_dict[v] = state_dict.pop(k) - - hf_value_function.load_state_dict(state_dict) - - torch.save(hf_value_function.state_dict(), "hub/hopper-medium-v2/value_function/diffusion_pytorch_model.bin") - with open("hub/hopper-medium-v2/value_function/config.json", "w") as f: - json.dump(config, f) - - -if __name__ == "__main__": - unet(32) - # unet(128) - value_function() diff --git a/scripts/convert_ms_text_to_video_to_diffusers.py b/scripts/convert_ms_text_to_video_to_diffusers.py deleted file mode 100644 index e150a491a0b0..000000000000 --- a/scripts/convert_ms_text_to_video_to_diffusers.py +++ /dev/null @@ -1,428 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the LDM checkpoints.""" - -import argparse - -import torch - -from diffusers import UNet3DConditionModel - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - weight = old_checkpoint[path["old"]] - names = ["proj_attn.weight"] - names_2 = ["proj_out.weight", "proj_in.weight"] - if any(k in new_path for k in names): - checkpoint[new_path] = weight[:, :, 0] - elif any(k in new_path for k in names_2) and len(weight.shape) > 2 and ".attentions." not in new_path: - checkpoint[new_path] = weight[:, :, 0] - else: - checkpoint[new_path] = weight - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_temp_conv_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - mapping.append({"old": old_item, "new": old_item}) - - return mapping - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - if "temopral_conv" not in old_item: - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has both EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - first_temp_attention = [v for v in unet_state_dict if v.startswith("input_blocks.0.1")] - paths = renew_attention_paths(first_temp_attention) - meta_path = {"old": "input_blocks.0.1", "new": "transformer_in"} - assign_to_checkpoint(paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config) - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - temp_attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.2" in key] - - if f"input_blocks.{i}.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temporal_convs = [key for key in resnets if "temopral_conv" in key] - paths = renew_temp_conv_paths(temporal_convs) - meta_path = { - "old": f"input_blocks.{i}.0.temopral_conv", - "new": f"down_blocks.{block_id}.temp_convs.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(temp_attentions): - paths = renew_attention_paths(temp_attentions) - meta_path = { - "old": f"input_blocks.{i}.2", - "new": f"down_blocks.{block_id}.temp_attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - temporal_convs_0 = [key for key in resnet_0 if "temopral_conv" in key] - attentions = middle_blocks[1] - temp_attentions = middle_blocks[2] - resnet_1 = middle_blocks[3] - temporal_convs_1 = [key for key in resnet_1 if "temopral_conv" in key] - - resnet_0_paths = renew_resnet_paths(resnet_0) - meta_path = {"old": "middle_block.0", "new": "mid_block.resnets.0"} - assign_to_checkpoint( - resnet_0_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - temp_conv_0_paths = renew_temp_conv_paths(temporal_convs_0) - meta_path = {"old": "middle_block.0.temopral_conv", "new": "mid_block.temp_convs.0"} - assign_to_checkpoint( - temp_conv_0_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - resnet_1_paths = renew_resnet_paths(resnet_1) - meta_path = {"old": "middle_block.3", "new": "mid_block.resnets.1"} - assign_to_checkpoint( - resnet_1_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - temp_conv_1_paths = renew_temp_conv_paths(temporal_convs_1) - meta_path = {"old": "middle_block.3.temopral_conv", "new": "mid_block.temp_convs.1"} - assign_to_checkpoint( - temp_conv_1_paths, new_checkpoint, unet_state_dict, config=config, additional_replacements=[meta_path] - ) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temp_attentions_paths = renew_attention_paths(temp_attentions) - meta_path = {"old": "middle_block.2", "new": "mid_block.temp_attentions.0"} - assign_to_checkpoint( - temp_attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - temp_attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.2" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - temporal_convs = [key for key in resnets if "temopral_conv" in key] - paths = renew_temp_conv_paths(temporal_convs) - meta_path = { - "old": f"output_blocks.{i}.0.temopral_conv", - "new": f"up_blocks.{block_id}.temp_convs.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(temp_attentions): - paths = renew_attention_paths(temp_attentions) - meta_path = { - "old": f"output_blocks.{i}.2", - "new": f"up_blocks.{block_id}.temp_attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - new_checkpoint[new_path] = unet_state_dict[old_path] - - temopral_conv_paths = [l for l in output_block_layers if "temopral_conv" in l] - for path in temopral_conv_paths: - pruned_path = path.split("temopral_conv.")[-1] - old_path = ".".join(["output_blocks", str(i), str(block_id), "temopral_conv", pruned_path]) - new_path = ".".join(["up_blocks", str(block_id), "temp_convs", str(layer_in_block_id), pruned_path]) - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - args = parser.parse_args() - - unet_checkpoint = torch.load(args.checkpoint_path, map_location="cpu") - unet = UNet3DConditionModel() - - converted_ckpt = convert_ldm_unet_checkpoint(unet_checkpoint, unet.config) - - diff_0 = set(unet.state_dict().keys()) - set(converted_ckpt.keys()) - diff_1 = set(converted_ckpt.keys()) - set(unet.state_dict().keys()) - - assert len(diff_0) == len(diff_1) == 0, "Converted weights don't match" - - # load state_dict - unet.load_state_dict(converted_ckpt) - - unet.save_pretrained(args.dump_path) - - # -- finish converting the unet -- diff --git a/scripts/convert_music_spectrogram_to_diffusers.py b/scripts/convert_music_spectrogram_to_diffusers.py deleted file mode 100644 index 1ac9ceddb57e..000000000000 --- a/scripts/convert_music_spectrogram_to_diffusers.py +++ /dev/null @@ -1,203 +0,0 @@ -#!/usr/bin/env python3 -import argparse -import os - -import jax as jnp -import numpy as onp -import torch -import torch.nn as nn -from music_spectrogram_diffusion import inference -from t5x import checkpoints - -from diffusers import DDPMScheduler, OnnxRuntimeModel, SpectrogramDiffusionPipeline -from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, T5FilmDecoder - - -MODEL = "base_with_context" - - -def load_notes_encoder(weights, model): - model.token_embedder.weight = nn.Parameter(torch.Tensor(weights["token_embedder"]["embedding"])) - model.position_encoding.weight = nn.Parameter(torch.Tensor(weights["Embed_0"]["embedding"]), requires_grad=False) - for lyr_num, lyr in enumerate(model.encoders): - ly_weight = weights[f"layers_{lyr_num}"] - lyr.layer[0].layer_norm.weight = nn.Parameter(torch.Tensor(ly_weight["pre_attention_layer_norm"]["scale"])) - - attention_weights = ly_weight["attention"] - lyr.layer[0].SelfAttention.q.weight = nn.Parameter(torch.Tensor(attention_weights["query"]["kernel"].T)) - lyr.layer[0].SelfAttention.k.weight = nn.Parameter(torch.Tensor(attention_weights["key"]["kernel"].T)) - lyr.layer[0].SelfAttention.v.weight = nn.Parameter(torch.Tensor(attention_weights["value"]["kernel"].T)) - lyr.layer[0].SelfAttention.o.weight = nn.Parameter(torch.Tensor(attention_weights["out"]["kernel"].T)) - - lyr.layer[1].layer_norm.weight = nn.Parameter(torch.Tensor(ly_weight["pre_mlp_layer_norm"]["scale"])) - - lyr.layer[1].DenseReluDense.wi_0.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_0"]["kernel"].T)) - lyr.layer[1].DenseReluDense.wi_1.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_1"]["kernel"].T)) - lyr.layer[1].DenseReluDense.wo.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wo"]["kernel"].T)) - - model.layer_norm.weight = nn.Parameter(torch.Tensor(weights["encoder_norm"]["scale"])) - return model - - -def load_continuous_encoder(weights, model): - model.input_proj.weight = nn.Parameter(torch.Tensor(weights["input_proj"]["kernel"].T)) - - model.position_encoding.weight = nn.Parameter(torch.Tensor(weights["Embed_0"]["embedding"]), requires_grad=False) - - for lyr_num, lyr in enumerate(model.encoders): - ly_weight = weights[f"layers_{lyr_num}"] - attention_weights = ly_weight["attention"] - - lyr.layer[0].SelfAttention.q.weight = nn.Parameter(torch.Tensor(attention_weights["query"]["kernel"].T)) - lyr.layer[0].SelfAttention.k.weight = nn.Parameter(torch.Tensor(attention_weights["key"]["kernel"].T)) - lyr.layer[0].SelfAttention.v.weight = nn.Parameter(torch.Tensor(attention_weights["value"]["kernel"].T)) - lyr.layer[0].SelfAttention.o.weight = nn.Parameter(torch.Tensor(attention_weights["out"]["kernel"].T)) - lyr.layer[0].layer_norm.weight = nn.Parameter(torch.Tensor(ly_weight["pre_attention_layer_norm"]["scale"])) - - lyr.layer[1].DenseReluDense.wi_0.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_0"]["kernel"].T)) - lyr.layer[1].DenseReluDense.wi_1.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_1"]["kernel"].T)) - lyr.layer[1].DenseReluDense.wo.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wo"]["kernel"].T)) - lyr.layer[1].layer_norm.weight = nn.Parameter(torch.Tensor(ly_weight["pre_mlp_layer_norm"]["scale"])) - - model.layer_norm.weight = nn.Parameter(torch.Tensor(weights["encoder_norm"]["scale"])) - - return model - - -def load_decoder(weights, model): - model.conditioning_emb[0].weight = nn.Parameter(torch.Tensor(weights["time_emb_dense0"]["kernel"].T)) - model.conditioning_emb[2].weight = nn.Parameter(torch.Tensor(weights["time_emb_dense1"]["kernel"].T)) - - model.position_encoding.weight = nn.Parameter(torch.Tensor(weights["Embed_0"]["embedding"]), requires_grad=False) - - model.continuous_inputs_projection.weight = nn.Parameter( - torch.Tensor(weights["continuous_inputs_projection"]["kernel"].T) - ) - - for lyr_num, lyr in enumerate(model.decoders): - ly_weight = weights[f"layers_{lyr_num}"] - lyr.layer[0].layer_norm.weight = nn.Parameter( - torch.Tensor(ly_weight["pre_self_attention_layer_norm"]["scale"]) - ) - - lyr.layer[0].FiLMLayer.scale_bias.weight = nn.Parameter( - torch.Tensor(ly_weight["FiLMLayer_0"]["DenseGeneral_0"]["kernel"].T) - ) - - attention_weights = ly_weight["self_attention"] - lyr.layer[0].attention.to_q.weight = nn.Parameter(torch.Tensor(attention_weights["query"]["kernel"].T)) - lyr.layer[0].attention.to_k.weight = nn.Parameter(torch.Tensor(attention_weights["key"]["kernel"].T)) - lyr.layer[0].attention.to_v.weight = nn.Parameter(torch.Tensor(attention_weights["value"]["kernel"].T)) - lyr.layer[0].attention.to_out[0].weight = nn.Parameter(torch.Tensor(attention_weights["out"]["kernel"].T)) - - attention_weights = ly_weight["MultiHeadDotProductAttention_0"] - lyr.layer[1].attention.to_q.weight = nn.Parameter(torch.Tensor(attention_weights["query"]["kernel"].T)) - lyr.layer[1].attention.to_k.weight = nn.Parameter(torch.Tensor(attention_weights["key"]["kernel"].T)) - lyr.layer[1].attention.to_v.weight = nn.Parameter(torch.Tensor(attention_weights["value"]["kernel"].T)) - lyr.layer[1].attention.to_out[0].weight = nn.Parameter(torch.Tensor(attention_weights["out"]["kernel"].T)) - lyr.layer[1].layer_norm.weight = nn.Parameter( - torch.Tensor(ly_weight["pre_cross_attention_layer_norm"]["scale"]) - ) - - lyr.layer[2].layer_norm.weight = nn.Parameter(torch.Tensor(ly_weight["pre_mlp_layer_norm"]["scale"])) - lyr.layer[2].film.scale_bias.weight = nn.Parameter( - torch.Tensor(ly_weight["FiLMLayer_1"]["DenseGeneral_0"]["kernel"].T) - ) - lyr.layer[2].DenseReluDense.wi_0.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_0"]["kernel"].T)) - lyr.layer[2].DenseReluDense.wi_1.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wi_1"]["kernel"].T)) - lyr.layer[2].DenseReluDense.wo.weight = nn.Parameter(torch.Tensor(ly_weight["mlp"]["wo"]["kernel"].T)) - - model.decoder_norm.weight = nn.Parameter(torch.Tensor(weights["decoder_norm"]["scale"])) - - model.spec_out.weight = nn.Parameter(torch.Tensor(weights["spec_out_dense"]["kernel"].T)) - - return model - - -def main(args): - t5_checkpoint = checkpoints.load_t5x_checkpoint(args.checkpoint_path) - t5_checkpoint = jnp.tree_util.tree_map(onp.array, t5_checkpoint) - - gin_overrides = [ - "from __gin__ import dynamic_registration", - "from music_spectrogram_diffusion.models.diffusion import diffusion_utils", - "diffusion_utils.ClassifierFreeGuidanceConfig.eval_condition_weight = 2.0", - "diffusion_utils.DiffusionConfig.classifier_free_guidance = @diffusion_utils.ClassifierFreeGuidanceConfig()", - ] - - gin_file = os.path.join(args.checkpoint_path, "..", "config.gin") - gin_config = inference.parse_training_gin_file(gin_file, gin_overrides) - synth_model = inference.InferenceModel(args.checkpoint_path, gin_config) - - scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", variance_type="fixed_large") - - notes_encoder = SpectrogramNotesEncoder( - max_length=synth_model.sequence_length["inputs"], - vocab_size=synth_model.model.module.config.vocab_size, - d_model=synth_model.model.module.config.emb_dim, - dropout_rate=synth_model.model.module.config.dropout_rate, - num_layers=synth_model.model.module.config.num_encoder_layers, - num_heads=synth_model.model.module.config.num_heads, - d_kv=synth_model.model.module.config.head_dim, - d_ff=synth_model.model.module.config.mlp_dim, - feed_forward_proj="gated-gelu", - ) - - continuous_encoder = SpectrogramContEncoder( - input_dims=synth_model.audio_codec.n_dims, - targets_context_length=synth_model.sequence_length["targets_context"], - d_model=synth_model.model.module.config.emb_dim, - dropout_rate=synth_model.model.module.config.dropout_rate, - num_layers=synth_model.model.module.config.num_encoder_layers, - num_heads=synth_model.model.module.config.num_heads, - d_kv=synth_model.model.module.config.head_dim, - d_ff=synth_model.model.module.config.mlp_dim, - feed_forward_proj="gated-gelu", - ) - - decoder = T5FilmDecoder( - input_dims=synth_model.audio_codec.n_dims, - targets_length=synth_model.sequence_length["targets_context"], - max_decoder_noise_time=synth_model.model.module.config.max_decoder_noise_time, - d_model=synth_model.model.module.config.emb_dim, - num_layers=synth_model.model.module.config.num_decoder_layers, - num_heads=synth_model.model.module.config.num_heads, - d_kv=synth_model.model.module.config.head_dim, - d_ff=synth_model.model.module.config.mlp_dim, - dropout_rate=synth_model.model.module.config.dropout_rate, - ) - - notes_encoder = load_notes_encoder(t5_checkpoint["target"]["token_encoder"], notes_encoder) - continuous_encoder = load_continuous_encoder(t5_checkpoint["target"]["continuous_encoder"], continuous_encoder) - decoder = load_decoder(t5_checkpoint["target"]["decoder"], decoder) - - melgan = OnnxRuntimeModel.from_pretrained("kashif/soundstream_mel_decoder") - - pipe = SpectrogramDiffusionPipeline( - notes_encoder=notes_encoder, - continuous_encoder=continuous_encoder, - decoder=decoder, - scheduler=scheduler, - melgan=melgan, - ) - if args.save: - pipe.save_pretrained(args.output_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--output_path", default=None, type=str, required=True, help="Path to the converted model.") - parser.add_argument( - "--save", default=True, type=bool, required=False, help="Whether to save the converted model or not." - ) - parser.add_argument( - "--checkpoint_path", - default=f"{MODEL}/checkpoint_500000", - type=str, - required=False, - help="Path to the original jax model checkpoint.", - ) - args = parser.parse_args() - - main(args) diff --git a/scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py b/scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py deleted file mode 100644 index bcab90e2a3db..000000000000 --- a/scripts/convert_ncsnpp_original_checkpoint_to_diffusers.py +++ /dev/null @@ -1,185 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the NCSNPP checkpoints.""" - -import argparse -import json - -import torch - -from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNet2DModel - - -def convert_ncsnpp_checkpoint(checkpoint, config): - """ - Takes a state dict and the path to - """ - new_model_architecture = UNet2DModel(**config) - new_model_architecture.time_proj.W.data = checkpoint["all_modules.0.W"].data - new_model_architecture.time_proj.weight.data = checkpoint["all_modules.0.W"].data - new_model_architecture.time_embedding.linear_1.weight.data = checkpoint["all_modules.1.weight"].data - new_model_architecture.time_embedding.linear_1.bias.data = checkpoint["all_modules.1.bias"].data - - new_model_architecture.time_embedding.linear_2.weight.data = checkpoint["all_modules.2.weight"].data - new_model_architecture.time_embedding.linear_2.bias.data = checkpoint["all_modules.2.bias"].data - - new_model_architecture.conv_in.weight.data = checkpoint["all_modules.3.weight"].data - new_model_architecture.conv_in.bias.data = checkpoint["all_modules.3.bias"].data - - new_model_architecture.conv_norm_out.weight.data = checkpoint[list(checkpoint.keys())[-4]].data - new_model_architecture.conv_norm_out.bias.data = checkpoint[list(checkpoint.keys())[-3]].data - new_model_architecture.conv_out.weight.data = checkpoint[list(checkpoint.keys())[-2]].data - new_model_architecture.conv_out.bias.data = checkpoint[list(checkpoint.keys())[-1]].data - - module_index = 4 - - def set_attention_weights(new_layer, old_checkpoint, index): - new_layer.query.weight.data = old_checkpoint[f"all_modules.{index}.NIN_0.W"].data.T - new_layer.key.weight.data = old_checkpoint[f"all_modules.{index}.NIN_1.W"].data.T - new_layer.value.weight.data = old_checkpoint[f"all_modules.{index}.NIN_2.W"].data.T - - new_layer.query.bias.data = old_checkpoint[f"all_modules.{index}.NIN_0.b"].data - new_layer.key.bias.data = old_checkpoint[f"all_modules.{index}.NIN_1.b"].data - new_layer.value.bias.data = old_checkpoint[f"all_modules.{index}.NIN_2.b"].data - - new_layer.proj_attn.weight.data = old_checkpoint[f"all_modules.{index}.NIN_3.W"].data.T - new_layer.proj_attn.bias.data = old_checkpoint[f"all_modules.{index}.NIN_3.b"].data - - new_layer.group_norm.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data - new_layer.group_norm.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.bias"].data - - def set_resnet_weights(new_layer, old_checkpoint, index): - new_layer.conv1.weight.data = old_checkpoint[f"all_modules.{index}.Conv_0.weight"].data - new_layer.conv1.bias.data = old_checkpoint[f"all_modules.{index}.Conv_0.bias"].data - new_layer.norm1.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.weight"].data - new_layer.norm1.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_0.bias"].data - - new_layer.conv2.weight.data = old_checkpoint[f"all_modules.{index}.Conv_1.weight"].data - new_layer.conv2.bias.data = old_checkpoint[f"all_modules.{index}.Conv_1.bias"].data - new_layer.norm2.weight.data = old_checkpoint[f"all_modules.{index}.GroupNorm_1.weight"].data - new_layer.norm2.bias.data = old_checkpoint[f"all_modules.{index}.GroupNorm_1.bias"].data - - new_layer.time_emb_proj.weight.data = old_checkpoint[f"all_modules.{index}.Dense_0.weight"].data - new_layer.time_emb_proj.bias.data = old_checkpoint[f"all_modules.{index}.Dense_0.bias"].data - - if new_layer.in_channels != new_layer.out_channels or new_layer.up or new_layer.down: - new_layer.conv_shortcut.weight.data = old_checkpoint[f"all_modules.{index}.Conv_2.weight"].data - new_layer.conv_shortcut.bias.data = old_checkpoint[f"all_modules.{index}.Conv_2.bias"].data - - for i, block in enumerate(new_model_architecture.downsample_blocks): - has_attentions = hasattr(block, "attentions") - for j in range(len(block.resnets)): - set_resnet_weights(block.resnets[j], checkpoint, module_index) - module_index += 1 - if has_attentions: - set_attention_weights(block.attentions[j], checkpoint, module_index) - module_index += 1 - - if hasattr(block, "downsamplers") and block.downsamplers is not None: - set_resnet_weights(block.resnet_down, checkpoint, module_index) - module_index += 1 - block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.Conv_0.weight"].data - block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.Conv_0.bias"].data - module_index += 1 - - set_resnet_weights(new_model_architecture.mid_block.resnets[0], checkpoint, module_index) - module_index += 1 - set_attention_weights(new_model_architecture.mid_block.attentions[0], checkpoint, module_index) - module_index += 1 - set_resnet_weights(new_model_architecture.mid_block.resnets[1], checkpoint, module_index) - module_index += 1 - - for i, block in enumerate(new_model_architecture.up_blocks): - has_attentions = hasattr(block, "attentions") - for j in range(len(block.resnets)): - set_resnet_weights(block.resnets[j], checkpoint, module_index) - module_index += 1 - if has_attentions: - set_attention_weights( - block.attentions[0], checkpoint, module_index - ) # why can there only be a single attention layer for up? - module_index += 1 - - if hasattr(block, "resnet_up") and block.resnet_up is not None: - block.skip_norm.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data - block.skip_norm.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data - module_index += 1 - block.skip_conv.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data - block.skip_conv.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data - module_index += 1 - set_resnet_weights(block.resnet_up, checkpoint, module_index) - module_index += 1 - - new_model_architecture.conv_norm_out.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data - new_model_architecture.conv_norm_out.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data - module_index += 1 - new_model_architecture.conv_out.weight.data = checkpoint[f"all_modules.{module_index}.weight"].data - new_model_architecture.conv_out.bias.data = checkpoint[f"all_modules.{module_index}.bias"].data - - return new_model_architecture.state_dict() - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", - default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_pytorch_model.bin", - type=str, - required=False, - help="Path to the checkpoint to convert.", - ) - - parser.add_argument( - "--config_file", - default="/Users/arthurzucker/Work/diffusers/ArthurZ/config.json", - type=str, - required=False, - help="The config json file corresponding to the architecture.", - ) - - parser.add_argument( - "--dump_path", - default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt", - type=str, - required=False, - help="Path to the output model.", - ) - - args = parser.parse_args() - - checkpoint = torch.load(args.checkpoint_path, map_location="cpu") - - with open(args.config_file) as f: - config = json.loads(f.read()) - - converted_checkpoint = convert_ncsnpp_checkpoint( - checkpoint, - config, - ) - - if "sde" in config: - del config["sde"] - - model = UNet2DModel(**config) - model.load_state_dict(converted_checkpoint) - - try: - scheduler = ScoreSdeVeScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) - - pipe = ScoreSdeVePipeline(unet=model, scheduler=scheduler) - pipe.save_pretrained(args.dump_path) - except: # noqa: E722 - model.save_pretrained(args.dump_path) diff --git a/scripts/convert_original_audioldm2_to_diffusers.py b/scripts/convert_original_audioldm2_to_diffusers.py deleted file mode 100644 index 2c0695ce5595..000000000000 --- a/scripts/convert_original_audioldm2_to_diffusers.py +++ /dev/null @@ -1,1135 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the AudioLDM2 checkpoints.""" - -import argparse -import re -from typing import List, Union - -import torch -import yaml -from transformers import ( - AutoFeatureExtractor, - AutoTokenizer, - ClapConfig, - ClapModel, - GPT2Config, - GPT2Model, - SpeechT5HifiGan, - SpeechT5HifiGanConfig, - T5Config, - T5EncoderModel, -) - -from diffusers import ( - AudioLDM2Pipeline, - AudioLDM2ProjectionModel, - AudioLDM2UNet2DConditionModel, - AutoencoderKL, - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, -) -from diffusers.utils import is_safetensors_available -from diffusers.utils.import_utils import BACKENDS_MAPPING - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.shave_segments -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_resnet_paths -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_resnet_paths -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_attention_paths -def renew_attention_paths(old_list): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["to_q.weight", "to_k.weight", "to_v.weight"] - proj_key = "to_out.0.weight" - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys or ".".join(key.split(".")[-3:]) == proj_key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key].squeeze() - - -def create_unet_diffusers_config(original_config, image_size: int): - """ - Creates a UNet config for diffusers based on the config of the original AudioLDM2 model. - """ - unet_params = original_config["model"]["params"]["unet_config"]["params"] - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - cross_attention_dim = list(unet_params["context_dim"]) if "context_dim" in unet_params else block_out_channels - if len(cross_attention_dim) > 1: - # require two or more cross-attention layers per-block, each of different dimension - cross_attention_dim = [cross_attention_dim for _ in range(len(block_out_channels))] - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "out_channels": unet_params["out_channels"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "transformer_layers_per_block": unet_params["transformer_depth"], - "cross_attention_dim": tuple(cross_attention_dim), - } - - return config - - -# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config -def create_vae_diffusers_config(original_config, checkpoint, image_size: int): - """ - Creates a VAE config for diffusers based on the config of the original AudioLDM2 model. Compared to the original - Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. - """ - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - "scaling_factor": float(scaling_factor), - } - return config - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_diffusers_schedular -def create_diffusers_schedular(original_config): - schedular = DDIMScheduler( - num_train_timesteps=original_config["model"]["params"]["timesteps"], - beta_start=original_config["model"]["params"]["linear_start"], - beta_end=original_config["model"]["params"]["linear_end"], - beta_schedule="scaled_linear", - ) - return schedular - - -def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted UNet checkpoint. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has both EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - # strip the unet prefix from the weight names - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}." in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}." in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}." in key] - for layer_id in range(num_output_blocks) - } - - # Check how many Transformer blocks we have per layer - if isinstance(config.get("cross_attention_dim"), (list, tuple)): - if isinstance(config["cross_attention_dim"][0], (list, tuple)): - # in this case we have multiple cross-attention layers per-block - num_attention_layers = len(config.get("cross_attention_dim")[0]) - else: - num_attention_layers = 1 - - if config.get("extra_self_attn_layer"): - num_attention_layers += 1 - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.0" not in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = [ - { - "old": f"input_blocks.{i}.{1 + layer_id}", - "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id * num_attention_layers + layer_id}", - } - for layer_id in range(num_attention_layers) - ] - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=meta_path, config=config - ) - - resnet_0 = middle_blocks[0] - resnet_1 = middle_blocks[num_middle_blocks - 1] - - resnet_0_paths = renew_resnet_paths(resnet_0) - meta_path = {"old": "middle_block.0", "new": "mid_block.resnets.0"} - assign_to_checkpoint( - resnet_0_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_1_paths = renew_resnet_paths(resnet_1) - meta_path = {"old": f"middle_block.{len(middle_blocks) - 1}", "new": "mid_block.resnets.1"} - assign_to_checkpoint( - resnet_1_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(1, num_middle_blocks - 1): - attentions = middle_blocks[i] - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": f"middle_block.{i}", "new": f"mid_block.attentions.{i - 1}"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.0" not in key] - - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - attentions.remove(f"output_blocks.{i}.{index}.conv.bias") - attentions.remove(f"output_blocks.{i}.{index}.conv.weight") - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = [ - { - "old": f"output_blocks.{i}.{1 + layer_id}", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id * num_attention_layers + layer_id}", - } - for layer_id in range(num_attention_layers) - ] - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=meta_path, config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - vae_key = "first_stage_model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -CLAP_KEYS_TO_MODIFY_MAPPING = { - "text_branch": "text_model", - "audio_branch": "audio_model.audio_encoder", - "attn": "attention.self", - "self.proj": "output.dense", - "attention.self_mask": "attn_mask", - "mlp.fc1": "intermediate.dense", - "mlp.fc2": "output.dense", - "norm1": "layernorm_before", - "norm2": "layernorm_after", - "bn0": "batch_norm", -} - -CLAP_KEYS_TO_IGNORE = [ - "text_transform", - "audio_transform", - "stft", - "logmel_extractor", - "tscam_conv", - "head", - "attn_mask", -] - -CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] - - -def convert_open_clap_checkpoint(checkpoint): - """ - Takes a state dict and returns a converted CLAP checkpoint. - """ - # extract state dict for CLAP text embedding model, discarding the audio component - model_state_dict = {} - model_key = "clap.model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(model_key): - model_state_dict[key.replace(model_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - sequential_layers_pattern = r".*sequential.(\d+).*" - text_projection_pattern = r".*_projection.(\d+).*" - - for key, value in model_state_dict.items(): - # check if key should be ignored in mapping - if so map it to a key name that we'll filter out at the end - for key_to_ignore in CLAP_KEYS_TO_IGNORE: - if key_to_ignore in key: - key = "spectrogram" - - # check if any key needs to be modified - for key_to_modify, new_key in CLAP_KEYS_TO_MODIFY_MAPPING.items(): - if key_to_modify in key: - key = key.replace(key_to_modify, new_key) - - if re.match(sequential_layers_pattern, key): - # replace sequential layers with list - sequential_layer = re.match(sequential_layers_pattern, key).group(1) - - key = key.replace(f"sequential.{sequential_layer}.", f"layers.{int(sequential_layer) // 3}.linear.") - elif re.match(text_projection_pattern, key): - projecton_layer = int(re.match(text_projection_pattern, key).group(1)) - - # Because in CLAP they use `nn.Sequential`... - transformers_projection_layer = 1 if projecton_layer == 0 else 2 - - key = key.replace(f"_projection.{projecton_layer}.", f"_projection.linear{transformers_projection_layer}.") - - if "audio" and "qkv" in key: - # split qkv into query key and value - mixed_qkv = value - qkv_dim = mixed_qkv.size(0) // 3 - - query_layer = mixed_qkv[:qkv_dim] - key_layer = mixed_qkv[qkv_dim : qkv_dim * 2] - value_layer = mixed_qkv[qkv_dim * 2 :] - - new_checkpoint[key.replace("qkv", "query")] = query_layer - new_checkpoint[key.replace("qkv", "key")] = key_layer - new_checkpoint[key.replace("qkv", "value")] = value_layer - elif key != "spectrogram": - new_checkpoint[key] = value - - return new_checkpoint - - -def create_transformers_vocoder_config(original_config): - """ - Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. - """ - vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] - - config = { - "model_in_dim": vocoder_params["num_mels"], - "sampling_rate": vocoder_params["sampling_rate"], - "upsample_initial_channel": vocoder_params["upsample_initial_channel"], - "upsample_rates": list(vocoder_params["upsample_rates"]), - "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), - "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), - "resblock_dilation_sizes": [ - list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] - ], - "normalize_before": False, - } - - return config - - -def extract_sub_model(checkpoint, key_prefix): - """ - Takes a state dict and returns the state dict for a particular sub-model. - """ - - sub_model_state_dict = {} - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(key_prefix): - sub_model_state_dict[key.replace(key_prefix, "")] = checkpoint.get(key) - - return sub_model_state_dict - - -def convert_hifigan_checkpoint(checkpoint, config): - """ - Takes a state dict and config, and returns a converted HiFiGAN vocoder checkpoint. - """ - # extract state dict for vocoder - vocoder_state_dict = extract_sub_model(checkpoint, key_prefix="first_stage_model.vocoder.") - - # fix upsampler keys, everything else is correct already - for i in range(len(config.upsample_rates)): - vocoder_state_dict[f"upsampler.{i}.weight"] = vocoder_state_dict.pop(f"ups.{i}.weight") - vocoder_state_dict[f"upsampler.{i}.bias"] = vocoder_state_dict.pop(f"ups.{i}.bias") - - if not config.normalize_before: - # if we don't set normalize_before then these variables are unused, so we set them to their initialised values - vocoder_state_dict["mean"] = torch.zeros(config.model_in_dim) - vocoder_state_dict["scale"] = torch.ones(config.model_in_dim) - - return vocoder_state_dict - - -def convert_projection_checkpoint(checkpoint): - projection_state_dict = {} - conditioner_state_dict = extract_sub_model(checkpoint, key_prefix="cond_stage_models.0.") - - projection_state_dict["sos_embed"] = conditioner_state_dict["start_of_sequence_tokens.weight"][0] - projection_state_dict["sos_embed_1"] = conditioner_state_dict["start_of_sequence_tokens.weight"][1] - - projection_state_dict["eos_embed"] = conditioner_state_dict["end_of_sequence_tokens.weight"][0] - projection_state_dict["eos_embed_1"] = conditioner_state_dict["end_of_sequence_tokens.weight"][1] - - projection_state_dict["projection.weight"] = conditioner_state_dict["input_sequence_embed_linear.0.weight"] - projection_state_dict["projection.bias"] = conditioner_state_dict["input_sequence_embed_linear.0.bias"] - - projection_state_dict["projection_1.weight"] = conditioner_state_dict["input_sequence_embed_linear.1.weight"] - projection_state_dict["projection_1.bias"] = conditioner_state_dict["input_sequence_embed_linear.1.bias"] - - return projection_state_dict - - -# Adapted from https://github.com/haoheliu/AudioLDM2/blob/81ad2c6ce015c1310387695e2dae975a7d2ed6fd/audioldm2/utils.py#L143 -DEFAULT_CONFIG = { - "model": { - "params": { - "linear_start": 0.0015, - "linear_end": 0.0195, - "timesteps": 1000, - "channels": 8, - "scale_by_std": True, - "unet_config": { - "target": "audioldm2.latent_diffusion.openaimodel.UNetModel", - "params": { - "context_dim": [None, 768, 1024], - "in_channels": 8, - "out_channels": 8, - "model_channels": 128, - "attention_resolutions": [8, 4, 2], - "num_res_blocks": 2, - "channel_mult": [1, 2, 3, 5], - "num_head_channels": 32, - "transformer_depth": 1, - }, - }, - "first_stage_config": { - "target": "audioldm2.variational_autoencoder.autoencoder.AutoencoderKL", - "params": { - "embed_dim": 8, - "ddconfig": { - "z_channels": 8, - "resolution": 256, - "in_channels": 1, - "out_ch": 1, - "ch": 128, - "ch_mult": [1, 2, 4], - "num_res_blocks": 2, - }, - }, - }, - "cond_stage_config": { - "crossattn_audiomae_generated": { - "target": "audioldm2.latent_diffusion.modules.encoders.modules.SequenceGenAudioMAECond", - "params": { - "sequence_gen_length": 8, - "sequence_input_embed_dim": [512, 1024], - }, - } - }, - "vocoder_config": { - "target": "audioldm2.first_stage_model.vocoder", - "params": { - "upsample_rates": [5, 4, 2, 2, 2], - "upsample_kernel_sizes": [16, 16, 8, 4, 4], - "upsample_initial_channel": 1024, - "resblock_kernel_sizes": [3, 7, 11], - "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "num_mels": 64, - "sampling_rate": 16000, - }, - }, - }, - }, -} - - -def load_pipeline_from_original_AudioLDM2_ckpt( - checkpoint_path: str, - original_config_file: str = None, - image_size: int = 1024, - prediction_type: str = None, - extract_ema: bool = False, - scheduler_type: str = "ddim", - cross_attention_dim: Union[List, List[List]] = None, - transformer_layers_per_block: int = None, - device: str = None, - from_safetensors: bool = False, -) -> AudioLDM2Pipeline: - """ - Load an AudioLDM2 pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. - - Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the - global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is - recommended that you override the default values and/or supply an `original_config_file` wherever possible. - - Args: - checkpoint_path (`str`): Path to `.ckpt` file. - original_config_file (`str`): - Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically - set to the AudioLDM2 base config. - image_size (`int`, *optional*, defaults to 1024): - The image size that the model was trained on. - prediction_type (`str`, *optional*): - The prediction type that the model was trained on. If `None`, will be automatically - inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. - scheduler_type (`str`, *optional*, defaults to 'ddim'): - Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", - "ddim"]`. - cross_attention_dim (`list`, *optional*, defaults to `None`): - The dimension of the cross-attention layers. If `None`, the cross-attention dimension will be - automatically inferred. Set to `[768, 1024]` for the base model, or `[768, 1024, None]` for the large model. - transformer_layers_per_block (`int`, *optional*, defaults to `None`): - The number of transformer layers in each transformer block. If `None`, number of layers will be " - "automatically inferred. Set to `1` for the base model, or `2` for the large model. - extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for - checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to - `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for - inference. Non-EMA weights are usually better to continue fine-tuning. - device (`str`, *optional*, defaults to `None`): - The device to use. Pass `None` to determine automatically. - from_safetensors (`str`, *optional*, defaults to `False`): - If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. - return: An AudioLDM2Pipeline object representing the passed-in `.ckpt`/`.safetensors` file. - """ - - if from_safetensors: - if not is_safetensors_available(): - raise ValueError(BACKENDS_MAPPING["safetensors"][1]) - - from safetensors import safe_open - - checkpoint = {} - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - checkpoint[key] = f.get_tensor(key) - else: - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - checkpoint = torch.load(checkpoint_path, map_location=device) - else: - checkpoint = torch.load(checkpoint_path, map_location=device) - - if "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - if original_config_file is None: - original_config = DEFAULT_CONFIG - else: - original_config = yaml.safe_load(original_config_file) - - if image_size is not None: - original_config["model"]["params"]["unet_config"]["params"]["image_size"] = image_size - - if cross_attention_dim is not None: - original_config["model"]["params"]["unet_config"]["params"]["context_dim"] = cross_attention_dim - - if transformer_layers_per_block is not None: - original_config["model"]["params"]["unet_config"]["params"]["transformer_depth"] = transformer_layers_per_block - - if ( - "parameterization" in original_config["model"]["params"] - and original_config["model"]["params"]["parameterization"] == "v" - ): - if prediction_type is None: - prediction_type = "v_prediction" - else: - if prediction_type is None: - prediction_type = "epsilon" - - num_train_timesteps = original_config["model"]["params"]["timesteps"] - beta_start = original_config["model"]["params"]["linear_start"] - beta_end = original_config["model"]["params"]["linear_end"] - - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type=prediction_type, - ) - # make sure scheduler works correctly with DDIM - scheduler.register_to_config(clip_sample=False) - - if scheduler_type == "pndm": - config = dict(scheduler.config) - config["skip_prk_steps"] = True - scheduler = PNDMScheduler.from_config(config) - elif scheduler_type == "lms": - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "heun": - scheduler = HeunDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler": - scheduler = EulerDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) - elif scheduler_type == "ddim": - scheduler = scheduler - else: - raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") - - # Convert the UNet2DModel - unet_config = create_unet_diffusers_config(original_config, image_size=image_size) - unet = AudioLDM2UNet2DConditionModel(**unet_config) - - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema - ) - - unet.load_state_dict(converted_unet_checkpoint) - - # Convert the VAE model - vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - - # Convert the joint audio-text encoding model - clap_config = ClapConfig.from_pretrained("laion/clap-htsat-unfused") - clap_config.audio_config.update( - { - "patch_embeds_hidden_size": 128, - "hidden_size": 1024, - "depths": [2, 2, 12, 2], - } - ) - # AudioLDM2 uses the same tokenizer and feature extractor as the original CLAP model - clap_tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") - clap_feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused") - - converted_clap_model = convert_open_clap_checkpoint(checkpoint) - clap_model = ClapModel(clap_config) - - missing_keys, unexpected_keys = clap_model.load_state_dict(converted_clap_model, strict=False) - # we expect not to have token_type_ids in our original state dict so let's ignore them - missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) - - if len(unexpected_keys) > 0: - raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") - - if len(missing_keys) > 0: - raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") - - # Convert the vocoder model - vocoder_config = create_transformers_vocoder_config(original_config) - vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) - converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) - - vocoder = SpeechT5HifiGan(vocoder_config) - vocoder.load_state_dict(converted_vocoder_checkpoint) - - # Convert the Flan-T5 encoder model: AudioLDM2 uses the same configuration and tokenizer as the original Flan-T5 large model - t5_config = T5Config.from_pretrained("google/flan-t5-large") - converted_t5_checkpoint = extract_sub_model(checkpoint, key_prefix="cond_stage_models.1.model.") - - t5_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large") - # hard-coded in the original implementation (i.e. not retrievable from the config) - t5_tokenizer.model_max_length = 128 - t5_model = T5EncoderModel(t5_config) - t5_model.load_state_dict(converted_t5_checkpoint) - - # Convert the GPT2 encoder model: AudioLDM2 uses the same configuration as the original GPT2 base model - gpt2_config = GPT2Config.from_pretrained("gpt2") - gpt2_model = GPT2Model(gpt2_config) - gpt2_model.config.max_new_tokens = original_config["model"]["params"]["cond_stage_config"][ - "crossattn_audiomae_generated" - ]["params"]["sequence_gen_length"] - - converted_gpt2_checkpoint = extract_sub_model(checkpoint, key_prefix="cond_stage_models.0.model.") - gpt2_model.load_state_dict(converted_gpt2_checkpoint) - - # Convert the extra embedding / projection layers - projection_model = AudioLDM2ProjectionModel(clap_config.projection_dim, t5_config.d_model, gpt2_config.n_embd) - - converted_projection_checkpoint = convert_projection_checkpoint(checkpoint) - projection_model.load_state_dict(converted_projection_checkpoint) - - # Instantiate the diffusers pipeline - pipe = AudioLDM2Pipeline( - vae=vae, - text_encoder=clap_model, - text_encoder_2=t5_model, - projection_model=projection_model, - language_model=gpt2_model, - tokenizer=clap_tokenizer, - tokenizer_2=t5_tokenizer, - feature_extractor=clap_feature_extractor, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - ) - - return pipe - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--original_config_file", - default=None, - type=str, - help="The YAML config file corresponding to the original architecture.", - ) - parser.add_argument( - "--cross_attention_dim", - default=None, - type=int, - nargs="+", - help="The dimension of the cross-attention layers. If `None`, the cross-attention dimension will be " - "automatically inferred. Set to `768+1024` for the base model, or `768+1024+640` for the large model", - ) - parser.add_argument( - "--transformer_layers_per_block", - default=None, - type=int, - help="The number of transformer layers in each transformer block. If `None`, number of layers will be " - "automatically inferred. Set to `1` for the base model, or `2` for the large model.", - ) - parser.add_argument( - "--scheduler_type", - default="ddim", - type=str, - help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", - ) - parser.add_argument( - "--image_size", - default=1048, - type=int, - help="The image size that the model was trained on.", - ) - parser.add_argument( - "--prediction_type", - default=None, - type=str, - help=("The prediction type that the model was trained on."), - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument( - "--from_safetensors", - action="store_true", - help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", - ) - parser.add_argument( - "--to_safetensors", - action="store_true", - help="Whether to store pipeline in safetensors format or not.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") - args = parser.parse_args() - - pipe = load_pipeline_from_original_AudioLDM2_ckpt( - checkpoint_path=args.checkpoint_path, - original_config_file=args.original_config_file, - image_size=args.image_size, - prediction_type=args.prediction_type, - extract_ema=args.extract_ema, - scheduler_type=args.scheduler_type, - cross_attention_dim=args.cross_attention_dim, - transformer_layers_per_block=args.transformer_layers_per_block, - from_safetensors=args.from_safetensors, - device=args.device, - ) - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/convert_original_audioldm_to_diffusers.py b/scripts/convert_original_audioldm_to_diffusers.py deleted file mode 100644 index 44183f1aea29..000000000000 --- a/scripts/convert_original_audioldm_to_diffusers.py +++ /dev/null @@ -1,1042 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the AudioLDM checkpoints.""" - -import argparse -import re - -import torch -import yaml -from transformers import ( - AutoTokenizer, - ClapTextConfig, - ClapTextModelWithProjection, - SpeechT5HifiGan, - SpeechT5HifiGanConfig, -) - -from diffusers import ( - AudioLDMPipeline, - AutoencoderKL, - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, - UNet2DConditionModel, -) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.shave_segments -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_resnet_paths -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_resnet_paths -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_attention_paths -def renew_attention_paths(old_list): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_attention_paths -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "query.weight") - new_item = new_item.replace("q.bias", "query.bias") - - new_item = new_item.replace("k.weight", "key.weight") - new_item = new_item.replace("k.bias", "key.bias") - - new_item = new_item.replace("v.weight", "value.weight") - new_item = new_item.replace("v.bias", "value.bias") - - new_item = new_item.replace("proj_out.weight", "proj_attn.weight") - new_item = new_item.replace("proj_out.bias", "proj_attn.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.assign_to_checkpoint -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.conv_attn_to_linear -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def create_unet_diffusers_config(original_config, image_size: int): - """ - Creates a UNet config for diffusers based on the config of the original AudioLDM model. - """ - unet_params = original_config["model"]["params"]["unet_config"]["params"] - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - cross_attention_dim = ( - unet_params["cross_attention_dim"] if "cross_attention_dim" in unet_params else block_out_channels - ) - - class_embed_type = "simple_projection" if "extra_film_condition_dim" in unet_params else None - projection_class_embeddings_input_dim = ( - unet_params["extra_film_condition_dim"] if "extra_film_condition_dim" in unet_params else None - ) - class_embeddings_concat = unet_params["extra_film_use_concat"] if "extra_film_use_concat" in unet_params else None - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "out_channels": unet_params["out_channels"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": cross_attention_dim, - "class_embed_type": class_embed_type, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "class_embeddings_concat": class_embeddings_concat, - } - - return config - - -# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config -def create_vae_diffusers_config(original_config, checkpoint, image_size: int): - """ - Creates a VAE config for diffusers based on the config of the original AudioLDM model. Compared to the original - Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. - """ - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - "scaling_factor": float(scaling_factor), - } - return config - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_diffusers_schedular -def create_diffusers_schedular(original_config): - schedular = DDIMScheduler( - num_train_timesteps=original_config["model"]["params"]["timesteps"], - beta_start=original_config["model"]["params"]["linear_start"], - beta_end=original_config["model"]["params"]["linear_end"], - beta_schedule="scaled_linear", - ) - return schedular - - -# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_unet_checkpoint -def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. Compared to the original Stable Diffusion - conversion, this function additionally converts the learnt film embedding linear layer. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has both EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - new_checkpoint["class_embedding.weight"] = unet_state_dict["film_emb.weight"] - new_checkpoint["class_embedding.bias"] = unet_state_dict["film_emb.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_vae_checkpoint -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - vae_key = "first_stage_model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -CLAP_KEYS_TO_MODIFY_MAPPING = { - "text_branch": "text_model", - "attn": "attention.self", - "self.proj": "output.dense", - "attention.self_mask": "attn_mask", - "mlp.fc1": "intermediate.dense", - "mlp.fc2": "output.dense", - "norm1": "layernorm_before", - "norm2": "layernorm_after", - "bn0": "batch_norm", -} - -CLAP_KEYS_TO_IGNORE = ["text_transform"] - -CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] - - -def convert_open_clap_checkpoint(checkpoint): - """ - Takes a state dict and returns a converted CLAP checkpoint. - """ - # extract state dict for CLAP text embedding model, discarding the audio component - model_state_dict = {} - model_key = "cond_stage_model.model.text_" - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(model_key): - model_state_dict[key.replace(model_key, "text_")] = checkpoint.get(key) - - new_checkpoint = {} - - sequential_layers_pattern = r".*sequential.(\d+).*" - text_projection_pattern = r".*_projection.(\d+).*" - - for key, value in model_state_dict.items(): - # check if key should be ignored in mapping - if key.split(".")[0] in CLAP_KEYS_TO_IGNORE: - continue - - # check if any key needs to be modified - for key_to_modify, new_key in CLAP_KEYS_TO_MODIFY_MAPPING.items(): - if key_to_modify in key: - key = key.replace(key_to_modify, new_key) - - if re.match(sequential_layers_pattern, key): - # replace sequential layers with list - sequential_layer = re.match(sequential_layers_pattern, key).group(1) - - key = key.replace(f"sequential.{sequential_layer}.", f"layers.{int(sequential_layer) // 3}.linear.") - elif re.match(text_projection_pattern, key): - projecton_layer = int(re.match(text_projection_pattern, key).group(1)) - - # Because in CLAP they use `nn.Sequential`... - transformers_projection_layer = 1 if projecton_layer == 0 else 2 - - key = key.replace(f"_projection.{projecton_layer}.", f"_projection.linear{transformers_projection_layer}.") - - if "audio" and "qkv" in key: - # split qkv into query key and value - mixed_qkv = value - qkv_dim = mixed_qkv.size(0) // 3 - - query_layer = mixed_qkv[:qkv_dim] - key_layer = mixed_qkv[qkv_dim : qkv_dim * 2] - value_layer = mixed_qkv[qkv_dim * 2 :] - - new_checkpoint[key.replace("qkv", "query")] = query_layer - new_checkpoint[key.replace("qkv", "key")] = key_layer - new_checkpoint[key.replace("qkv", "value")] = value_layer - else: - new_checkpoint[key] = value - - return new_checkpoint - - -def create_transformers_vocoder_config(original_config): - """ - Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. - """ - vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] - - config = { - "model_in_dim": vocoder_params["num_mels"], - "sampling_rate": vocoder_params["sampling_rate"], - "upsample_initial_channel": vocoder_params["upsample_initial_channel"], - "upsample_rates": list(vocoder_params["upsample_rates"]), - "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), - "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), - "resblock_dilation_sizes": [ - list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] - ], - "normalize_before": False, - } - - return config - - -def convert_hifigan_checkpoint(checkpoint, config): - """ - Takes a state dict and config, and returns a converted HiFiGAN vocoder checkpoint. - """ - # extract state dict for vocoder - vocoder_state_dict = {} - vocoder_key = "first_stage_model.vocoder." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vocoder_key): - vocoder_state_dict[key.replace(vocoder_key, "")] = checkpoint.get(key) - - # fix upsampler keys, everything else is correct already - for i in range(len(config.upsample_rates)): - vocoder_state_dict[f"upsampler.{i}.weight"] = vocoder_state_dict.pop(f"ups.{i}.weight") - vocoder_state_dict[f"upsampler.{i}.bias"] = vocoder_state_dict.pop(f"ups.{i}.bias") - - if not config.normalize_before: - # if we don't set normalize_before then these variables are unused, so we set them to their initialised values - vocoder_state_dict["mean"] = torch.zeros(config.model_in_dim) - vocoder_state_dict["scale"] = torch.ones(config.model_in_dim) - - return vocoder_state_dict - - -# Adapted from https://huggingface.co/spaces/haoheliu/audioldm-text-to-audio-generation/blob/84a0384742a22bd80c44e903e241f0623e874f1d/audioldm/utils.py#L72-L73 -DEFAULT_CONFIG = { - "model": { - "params": { - "linear_start": 0.0015, - "linear_end": 0.0195, - "timesteps": 1000, - "channels": 8, - "scale_by_std": True, - "unet_config": { - "target": "audioldm.latent_diffusion.openaimodel.UNetModel", - "params": { - "extra_film_condition_dim": 512, - "extra_film_use_concat": True, - "in_channels": 8, - "out_channels": 8, - "model_channels": 128, - "attention_resolutions": [8, 4, 2], - "num_res_blocks": 2, - "channel_mult": [1, 2, 3, 5], - "num_head_channels": 32, - }, - }, - "first_stage_config": { - "target": "audioldm.variational_autoencoder.autoencoder.AutoencoderKL", - "params": { - "embed_dim": 8, - "ddconfig": { - "z_channels": 8, - "resolution": 256, - "in_channels": 1, - "out_ch": 1, - "ch": 128, - "ch_mult": [1, 2, 4], - "num_res_blocks": 2, - }, - }, - }, - "vocoder_config": { - "target": "audioldm.first_stage_model.vocoder", - "params": { - "upsample_rates": [5, 4, 2, 2, 2], - "upsample_kernel_sizes": [16, 16, 8, 4, 4], - "upsample_initial_channel": 1024, - "resblock_kernel_sizes": [3, 7, 11], - "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "num_mels": 64, - "sampling_rate": 16000, - }, - }, - }, - }, -} - - -def load_pipeline_from_original_audioldm_ckpt( - checkpoint_path: str, - original_config_file: str = None, - image_size: int = 512, - prediction_type: str = None, - extract_ema: bool = False, - scheduler_type: str = "ddim", - num_in_channels: int = None, - model_channels: int = None, - num_head_channels: int = None, - device: str = None, - from_safetensors: bool = False, -) -> AudioLDMPipeline: - """ - Load an AudioLDM pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. - - Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the - global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is - recommended that you override the default values and/or supply an `original_config_file` wherever possible. - - Args: - checkpoint_path (`str`): Path to `.ckpt` file. - original_config_file (`str`): - Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically - set to the audioldm-s-full-v2 config. - image_size (`int`, *optional*, defaults to 512): - The image size that the model was trained on. - prediction_type (`str`, *optional*): - The prediction type that the model was trained on. If `None`, will be automatically - inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. - num_in_channels (`int`, *optional*, defaults to None): - The number of UNet input channels. If `None`, it will be automatically inferred from the config. - model_channels (`int`, *optional*, defaults to None): - The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override - to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large. - num_head_channels (`int`, *optional*, defaults to None): - The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override - to 32 for the small and medium checkpoints, and 64 for the large. - scheduler_type (`str`, *optional*, defaults to 'pndm'): - Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", - "ddim"]`. - extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for - checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to - `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for - inference. Non-EMA weights are usually better to continue fine-tuning. - device (`str`, *optional*, defaults to `None`): - The device to use. Pass `None` to determine automatically. - from_safetensors (`str`, *optional*, defaults to `False`): - If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. - return: An AudioLDMPipeline object representing the passed-in `.ckpt`/`.safetensors` file. - """ - - if from_safetensors: - from safetensors import safe_open - - checkpoint = {} - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - checkpoint[key] = f.get_tensor(key) - else: - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - checkpoint = torch.load(checkpoint_path, map_location=device) - else: - checkpoint = torch.load(checkpoint_path, map_location=device) - - if "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - if original_config_file is None: - original_config = DEFAULT_CONFIG - else: - original_config = yaml.safe_load(original_config_file) - - if num_in_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels - - if model_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["model_channels"] = model_channels - - if num_head_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = num_head_channels - - if ( - "parameterization" in original_config["model"]["params"] - and original_config["model"]["params"]["parameterization"] == "v" - ): - if prediction_type is None: - prediction_type = "v_prediction" - else: - if prediction_type is None: - prediction_type = "epsilon" - - if image_size is None: - image_size = 512 - - num_train_timesteps = original_config["model"]["params"]["timesteps"] - beta_start = original_config["model"]["params"]["linear_start"] - beta_end = original_config["model"]["params"]["linear_end"] - - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type=prediction_type, - ) - # make sure scheduler works correctly with DDIM - scheduler.register_to_config(clip_sample=False) - - if scheduler_type == "pndm": - config = dict(scheduler.config) - config["skip_prk_steps"] = True - scheduler = PNDMScheduler.from_config(config) - elif scheduler_type == "lms": - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "heun": - scheduler = HeunDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler": - scheduler = EulerDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) - elif scheduler_type == "ddim": - scheduler = scheduler - else: - raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") - - # Convert the UNet2DModel - unet_config = create_unet_diffusers_config(original_config, image_size=image_size) - unet = UNet2DConditionModel(**unet_config) - - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema - ) - - unet.load_state_dict(converted_unet_checkpoint) - - # Convert the VAE model - vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - - # Convert the text model - # AudioLDM uses the same configuration and tokenizer as the original CLAP model - config = ClapTextConfig.from_pretrained("laion/clap-htsat-unfused") - tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") - - converted_text_model = convert_open_clap_checkpoint(checkpoint) - text_model = ClapTextModelWithProjection(config) - - missing_keys, unexpected_keys = text_model.load_state_dict(converted_text_model, strict=False) - # we expect not to have token_type_ids in our original state dict so let's ignore them - missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) - - if len(unexpected_keys) > 0: - raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") - - if len(missing_keys) > 0: - raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") - - # Convert the vocoder model - vocoder_config = create_transformers_vocoder_config(original_config) - vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) - converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) - - vocoder = SpeechT5HifiGan(vocoder_config) - vocoder.load_state_dict(converted_vocoder_checkpoint) - - # Instantiate the diffusers pipeline - pipe = AudioLDMPipeline( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - ) - - return pipe - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--original_config_file", - default=None, - type=str, - help="The YAML config file corresponding to the original architecture.", - ) - parser.add_argument( - "--num_in_channels", - default=None, - type=int, - help="The number of input channels. If `None` number of input channels will be automatically inferred.", - ) - parser.add_argument( - "--model_channels", - default=None, - type=int, - help="The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override" - " to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large.", - ) - parser.add_argument( - "--num_head_channels", - default=None, - type=int, - help="The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override" - " to 32 for the small and medium checkpoints, and 64 for the large.", - ) - parser.add_argument( - "--scheduler_type", - default="ddim", - type=str, - help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", - ) - parser.add_argument( - "--image_size", - default=None, - type=int, - help=("The image size that the model was trained on."), - ) - parser.add_argument( - "--prediction_type", - default=None, - type=str, - help=("The prediction type that the model was trained on."), - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument( - "--from_safetensors", - action="store_true", - help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", - ) - parser.add_argument( - "--to_safetensors", - action="store_true", - help="Whether to store pipeline in safetensors format or not.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") - args = parser.parse_args() - - pipe = load_pipeline_from_original_audioldm_ckpt( - checkpoint_path=args.checkpoint_path, - original_config_file=args.original_config_file, - image_size=args.image_size, - prediction_type=args.prediction_type, - extract_ema=args.extract_ema, - scheduler_type=args.scheduler_type, - num_in_channels=args.num_in_channels, - model_channels=args.model_channels, - num_head_channels=args.num_head_channels, - from_safetensors=args.from_safetensors, - device=args.device, - ) - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/convert_original_musicldm_to_diffusers.py b/scripts/convert_original_musicldm_to_diffusers.py deleted file mode 100644 index 00836fde2592..000000000000 --- a/scripts/convert_original_musicldm_to_diffusers.py +++ /dev/null @@ -1,1056 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the MusicLDM checkpoints.""" - -import argparse -import re - -import torch -import yaml -from transformers import ( - AutoFeatureExtractor, - AutoTokenizer, - ClapConfig, - ClapModel, - SpeechT5HifiGan, - SpeechT5HifiGanConfig, -) - -from diffusers import ( - AutoencoderKL, - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - HeunDiscreteScheduler, - LMSDiscreteScheduler, - MusicLDMPipeline, - PNDMScheduler, - UNet2DConditionModel, -) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.shave_segments -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_resnet_paths -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_resnet_paths -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_attention_paths -def renew_attention_paths(old_list): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.assign_to_checkpoint -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["to_q.weight", "to_k.weight", "to_v.weight"] - proj_key = "to_out.0.weight" - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys or ".".join(key.split(".")[-3:]) == proj_key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key].squeeze() - - -def create_unet_diffusers_config(original_config, image_size: int): - """ - Creates a UNet config for diffusers based on the config of the original MusicLDM model. - """ - unet_params = original_config["model"]["params"]["unet_config"]["params"] - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - cross_attention_dim = ( - unet_params["cross_attention_dim"] if "cross_attention_dim" in unet_params else block_out_channels - ) - - class_embed_type = "simple_projection" if "extra_film_condition_dim" in unet_params else None - projection_class_embeddings_input_dim = ( - unet_params["extra_film_condition_dim"] if "extra_film_condition_dim" in unet_params else None - ) - class_embeddings_concat = unet_params["extra_film_use_concat"] if "extra_film_use_concat" in unet_params else None - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "out_channels": unet_params["out_channels"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": cross_attention_dim, - "class_embed_type": class_embed_type, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "class_embeddings_concat": class_embeddings_concat, - } - - return config - - -# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config -def create_vae_diffusers_config(original_config, checkpoint, image_size: int): - """ - Creates a VAE config for diffusers based on the config of the original MusicLDM model. Compared to the original - Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. - """ - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - "scaling_factor": float(scaling_factor), - } - return config - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_diffusers_schedular -def create_diffusers_schedular(original_config): - schedular = DDIMScheduler( - num_train_timesteps=original_config["model"]["params"]["timesteps"], - beta_start=original_config["model"]["params"]["linear_start"], - beta_end=original_config["model"]["params"]["linear_end"], - beta_schedule="scaled_linear", - ) - return schedular - - -def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. Compared to the original Stable Diffusion - conversion, this function additionally converts the learnt film embedding linear layer. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - print(f"Checkpoint {path} has both EMA and non-EMA weights.") - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - new_checkpoint["class_embedding.weight"] = unet_state_dict["film_emb.weight"] - new_checkpoint["class_embedding.bias"] = unet_state_dict["film_emb.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_vae_checkpoint -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - vae_key = "first_stage_model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -CLAP_KEYS_TO_MODIFY_MAPPING = { - "text_branch": "text_model", - "audio_branch": "audio_model.audio_encoder", - "attn": "attention.self", - "self.proj": "output.dense", - "attention.self_mask": "attn_mask", - "mlp.fc1": "intermediate.dense", - "mlp.fc2": "output.dense", - "norm1": "layernorm_before", - "norm2": "layernorm_after", - "bn0": "batch_norm", -} - -CLAP_KEYS_TO_IGNORE = [ - "text_transform", - "audio_transform", - "stft", - "logmel_extractor", - "tscam_conv", - "head", - "attn_mask", -] - -CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] - - -def convert_open_clap_checkpoint(checkpoint): - """ - Takes a state dict and returns a converted CLAP checkpoint. - """ - # extract state dict for CLAP text embedding model, discarding the audio component - model_state_dict = {} - model_key = "cond_stage_model.model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(model_key): - model_state_dict[key.replace(model_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - sequential_layers_pattern = r".*sequential.(\d+).*" - text_projection_pattern = r".*_projection.(\d+).*" - - for key, value in model_state_dict.items(): - # check if key should be ignored in mapping - if so map it to a key name that we'll filter out at the end - for key_to_ignore in CLAP_KEYS_TO_IGNORE: - if key_to_ignore in key: - key = "spectrogram" - - # check if any key needs to be modified - for key_to_modify, new_key in CLAP_KEYS_TO_MODIFY_MAPPING.items(): - if key_to_modify in key: - key = key.replace(key_to_modify, new_key) - - if re.match(sequential_layers_pattern, key): - # replace sequential layers with list - sequential_layer = re.match(sequential_layers_pattern, key).group(1) - - key = key.replace(f"sequential.{sequential_layer}.", f"layers.{int(sequential_layer) // 3}.linear.") - elif re.match(text_projection_pattern, key): - projecton_layer = int(re.match(text_projection_pattern, key).group(1)) - - # Because in CLAP they use `nn.Sequential`... - transformers_projection_layer = 1 if projecton_layer == 0 else 2 - - key = key.replace(f"_projection.{projecton_layer}.", f"_projection.linear{transformers_projection_layer}.") - - if "audio" and "qkv" in key: - # split qkv into query key and value - mixed_qkv = value - qkv_dim = mixed_qkv.size(0) // 3 - - query_layer = mixed_qkv[:qkv_dim] - key_layer = mixed_qkv[qkv_dim : qkv_dim * 2] - value_layer = mixed_qkv[qkv_dim * 2 :] - - new_checkpoint[key.replace("qkv", "query")] = query_layer - new_checkpoint[key.replace("qkv", "key")] = key_layer - new_checkpoint[key.replace("qkv", "value")] = value_layer - elif key != "spectrogram": - new_checkpoint[key] = value - - return new_checkpoint - - -def create_transformers_vocoder_config(original_config): - """ - Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. - """ - vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] - - config = { - "model_in_dim": vocoder_params["num_mels"], - "sampling_rate": vocoder_params["sampling_rate"], - "upsample_initial_channel": vocoder_params["upsample_initial_channel"], - "upsample_rates": list(vocoder_params["upsample_rates"]), - "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), - "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), - "resblock_dilation_sizes": [ - list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] - ], - "normalize_before": False, - } - - return config - - -def convert_hifigan_checkpoint(checkpoint, config): - """ - Takes a state dict and config, and returns a converted HiFiGAN vocoder checkpoint. - """ - # extract state dict for vocoder - vocoder_state_dict = {} - vocoder_key = "first_stage_model.vocoder." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vocoder_key): - vocoder_state_dict[key.replace(vocoder_key, "")] = checkpoint.get(key) - - # fix upsampler keys, everything else is correct already - for i in range(len(config.upsample_rates)): - vocoder_state_dict[f"upsampler.{i}.weight"] = vocoder_state_dict.pop(f"ups.{i}.weight") - vocoder_state_dict[f"upsampler.{i}.bias"] = vocoder_state_dict.pop(f"ups.{i}.bias") - - if not config.normalize_before: - # if we don't set normalize_before then these variables are unused, so we set them to their initialised values - vocoder_state_dict["mean"] = torch.zeros(config.model_in_dim) - vocoder_state_dict["scale"] = torch.ones(config.model_in_dim) - - return vocoder_state_dict - - -# Adapted from https://huggingface.co/spaces/haoheliu/MusicLDM-text-to-audio-generation/blob/84a0384742a22bd80c44e903e241f0623e874f1d/MusicLDM/utils.py#L72-L73 -DEFAULT_CONFIG = { - "model": { - "params": { - "linear_start": 0.0015, - "linear_end": 0.0195, - "timesteps": 1000, - "channels": 8, - "scale_by_std": True, - "unet_config": { - "target": "MusicLDM.latent_diffusion.openaimodel.UNetModel", - "params": { - "extra_film_condition_dim": 512, - "extra_film_use_concat": True, - "in_channels": 8, - "out_channels": 8, - "model_channels": 128, - "attention_resolutions": [8, 4, 2], - "num_res_blocks": 2, - "channel_mult": [1, 2, 3, 5], - "num_head_channels": 32, - }, - }, - "first_stage_config": { - "target": "MusicLDM.variational_autoencoder.autoencoder.AutoencoderKL", - "params": { - "embed_dim": 8, - "ddconfig": { - "z_channels": 8, - "resolution": 256, - "in_channels": 1, - "out_ch": 1, - "ch": 128, - "ch_mult": [1, 2, 4], - "num_res_blocks": 2, - }, - }, - }, - "vocoder_config": { - "target": "MusicLDM.first_stage_model.vocoder", - "params": { - "upsample_rates": [5, 4, 2, 2, 2], - "upsample_kernel_sizes": [16, 16, 8, 4, 4], - "upsample_initial_channel": 1024, - "resblock_kernel_sizes": [3, 7, 11], - "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], - "num_mels": 64, - "sampling_rate": 16000, - }, - }, - }, - }, -} - - -def load_pipeline_from_original_MusicLDM_ckpt( - checkpoint_path: str, - original_config_file: str = None, - image_size: int = 1024, - prediction_type: str = None, - extract_ema: bool = False, - scheduler_type: str = "ddim", - num_in_channels: int = None, - model_channels: int = None, - num_head_channels: int = None, - device: str = None, - from_safetensors: bool = False, -) -> MusicLDMPipeline: - """ - Load an MusicLDM pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. - - Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the - global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is - recommended that you override the default values and/or supply an `original_config_file` wherever possible. - - Args: - checkpoint_path (`str`): Path to `.ckpt` file. - original_config_file (`str`): - Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically - set to the MusicLDM-s-full-v2 config. - image_size (`int`, *optional*, defaults to 1024): - The image size that the model was trained on. - prediction_type (`str`, *optional*): - The prediction type that the model was trained on. If `None`, will be automatically - inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. - num_in_channels (`int`, *optional*, defaults to None): - The number of UNet input channels. If `None`, it will be automatically inferred from the config. - model_channels (`int`, *optional*, defaults to None): - The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override - to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large. - num_head_channels (`int`, *optional*, defaults to None): - The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override - to 32 for the small and medium checkpoints, and 64 for the large. - scheduler_type (`str`, *optional*, defaults to 'pndm'): - Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", - "ddim"]`. - extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for - checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to - `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for - inference. Non-EMA weights are usually better to continue fine-tuning. - device (`str`, *optional*, defaults to `None`): - The device to use. Pass `None` to determine automatically. - from_safetensors (`str`, *optional*, defaults to `False`): - If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. - return: An MusicLDMPipeline object representing the passed-in `.ckpt`/`.safetensors` file. - """ - if from_safetensors: - from safetensors import safe_open - - checkpoint = {} - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - checkpoint[key] = f.get_tensor(key) - else: - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - checkpoint = torch.load(checkpoint_path, map_location=device) - else: - checkpoint = torch.load(checkpoint_path, map_location=device) - - if "state_dict" in checkpoint: - checkpoint = checkpoint["state_dict"] - - if original_config_file is None: - original_config = DEFAULT_CONFIG - else: - original_config = yaml.safe_load(original_config_file) - - if num_in_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels - - if model_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["model_channels"] = model_channels - - if num_head_channels is not None: - original_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = num_head_channels - - if ( - "parameterization" in original_config["model"]["params"] - and original_config["model"]["params"]["parameterization"] == "v" - ): - if prediction_type is None: - prediction_type = "v_prediction" - else: - if prediction_type is None: - prediction_type = "epsilon" - - if image_size is None: - image_size = 512 - - num_train_timesteps = original_config["model"]["params"]["timesteps"] - beta_start = original_config["model"]["params"]["linear_start"] - beta_end = original_config["model"]["params"]["linear_end"] - - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type=prediction_type, - ) - # make sure scheduler works correctly with DDIM - scheduler.register_to_config(clip_sample=False) - - if scheduler_type == "pndm": - config = dict(scheduler.config) - config["skip_prk_steps"] = True - scheduler = PNDMScheduler.from_config(config) - elif scheduler_type == "lms": - scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "heun": - scheduler = HeunDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler": - scheduler = EulerDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) - elif scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) - elif scheduler_type == "ddim": - scheduler = scheduler - else: - raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") - - # Convert the UNet2DModel - unet_config = create_unet_diffusers_config(original_config, image_size=image_size) - unet = UNet2DConditionModel(**unet_config) - - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema - ) - - unet.load_state_dict(converted_unet_checkpoint) - - # Convert the VAE model - vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - - # Convert the text model - # MusicLDM uses the same tokenizer as the original CLAP model, but a slightly different configuration - config = ClapConfig.from_pretrained("laion/clap-htsat-unfused") - config.audio_config.update( - { - "patch_embeds_hidden_size": 128, - "hidden_size": 1024, - "depths": [2, 2, 12, 2], - } - ) - tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") - feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused") - - converted_text_model = convert_open_clap_checkpoint(checkpoint) - text_model = ClapModel(config) - - missing_keys, unexpected_keys = text_model.load_state_dict(converted_text_model, strict=False) - # we expect not to have token_type_ids in our original state dict so let's ignore them - missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) - - if len(unexpected_keys) > 0: - raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") - - if len(missing_keys) > 0: - raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") - - # Convert the vocoder model - vocoder_config = create_transformers_vocoder_config(original_config) - vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) - converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) - - vocoder = SpeechT5HifiGan(vocoder_config) - vocoder.load_state_dict(converted_vocoder_checkpoint) - - # Instantiate the diffusers pipeline - pipe = MusicLDMPipeline( - vae=vae, - text_encoder=text_model, - tokenizer=tokenizer, - unet=unet, - scheduler=scheduler, - vocoder=vocoder, - feature_extractor=feature_extractor, - ) - - return pipe - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--original_config_file", - default=None, - type=str, - help="The YAML config file corresponding to the original architecture.", - ) - parser.add_argument( - "--num_in_channels", - default=None, - type=int, - help="The number of input channels. If `None` number of input channels will be automatically inferred.", - ) - parser.add_argument( - "--model_channels", - default=None, - type=int, - help="The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override" - " to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large.", - ) - parser.add_argument( - "--num_head_channels", - default=None, - type=int, - help="The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override" - " to 32 for the small and medium checkpoints, and 64 for the large.", - ) - parser.add_argument( - "--scheduler_type", - default="ddim", - type=str, - help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", - ) - parser.add_argument( - "--image_size", - default=None, - type=int, - help=("The image size that the model was trained on."), - ) - parser.add_argument( - "--prediction_type", - default=None, - type=str, - help=("The prediction type that the model was trained on."), - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument( - "--from_safetensors", - action="store_true", - help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", - ) - parser.add_argument( - "--to_safetensors", - action="store_true", - help="Whether to store pipeline in safetensors format or not.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") - args = parser.parse_args() - - pipe = load_pipeline_from_original_MusicLDM_ckpt( - checkpoint_path=args.checkpoint_path, - original_config_file=args.original_config_file, - image_size=args.image_size, - prediction_type=args.prediction_type, - extract_ema=args.extract_ema, - scheduler_type=args.scheduler_type, - num_in_channels=args.num_in_channels, - model_channels=args.model_channels, - num_head_channels=args.num_head_channels, - from_safetensors=args.from_safetensors, - device=args.device, - ) - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/convert_original_t2i_adapter.py b/scripts/convert_original_t2i_adapter.py deleted file mode 100644 index e23a2431ce9e..000000000000 --- a/scripts/convert_original_t2i_adapter.py +++ /dev/null @@ -1,250 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Conversion script for the T2I-Adapter checkpoints. -""" - -import argparse - -import torch - -from diffusers import T2IAdapter - - -def convert_adapter(src_state, in_channels): - original_body_length = max([int(x.split(".")[1]) for x in src_state.keys() if "body." in x]) + 1 - - assert original_body_length == 8 - - # (0, 1) -> channels 1 - assert src_state["body.0.block1.weight"].shape == (320, 320, 3, 3) - - # (2, 3) -> channels 2 - assert src_state["body.2.in_conv.weight"].shape == (640, 320, 1, 1) - - # (4, 5) -> channels 3 - assert src_state["body.4.in_conv.weight"].shape == (1280, 640, 1, 1) - - # (6, 7) -> channels 4 - assert src_state["body.6.block1.weight"].shape == (1280, 1280, 3, 3) - - res_state = { - "adapter.conv_in.weight": src_state.pop("conv_in.weight"), - "adapter.conv_in.bias": src_state.pop("conv_in.bias"), - # 0.resnets.0 - "adapter.body.0.resnets.0.block1.weight": src_state.pop("body.0.block1.weight"), - "adapter.body.0.resnets.0.block1.bias": src_state.pop("body.0.block1.bias"), - "adapter.body.0.resnets.0.block2.weight": src_state.pop("body.0.block2.weight"), - "adapter.body.0.resnets.0.block2.bias": src_state.pop("body.0.block2.bias"), - # 0.resnets.1 - "adapter.body.0.resnets.1.block1.weight": src_state.pop("body.1.block1.weight"), - "adapter.body.0.resnets.1.block1.bias": src_state.pop("body.1.block1.bias"), - "adapter.body.0.resnets.1.block2.weight": src_state.pop("body.1.block2.weight"), - "adapter.body.0.resnets.1.block2.bias": src_state.pop("body.1.block2.bias"), - # 1 - "adapter.body.1.in_conv.weight": src_state.pop("body.2.in_conv.weight"), - "adapter.body.1.in_conv.bias": src_state.pop("body.2.in_conv.bias"), - # 1.resnets.0 - "adapter.body.1.resnets.0.block1.weight": src_state.pop("body.2.block1.weight"), - "adapter.body.1.resnets.0.block1.bias": src_state.pop("body.2.block1.bias"), - "adapter.body.1.resnets.0.block2.weight": src_state.pop("body.2.block2.weight"), - "adapter.body.1.resnets.0.block2.bias": src_state.pop("body.2.block2.bias"), - # 1.resnets.1 - "adapter.body.1.resnets.1.block1.weight": src_state.pop("body.3.block1.weight"), - "adapter.body.1.resnets.1.block1.bias": src_state.pop("body.3.block1.bias"), - "adapter.body.1.resnets.1.block2.weight": src_state.pop("body.3.block2.weight"), - "adapter.body.1.resnets.1.block2.bias": src_state.pop("body.3.block2.bias"), - # 2 - "adapter.body.2.in_conv.weight": src_state.pop("body.4.in_conv.weight"), - "adapter.body.2.in_conv.bias": src_state.pop("body.4.in_conv.bias"), - # 2.resnets.0 - "adapter.body.2.resnets.0.block1.weight": src_state.pop("body.4.block1.weight"), - "adapter.body.2.resnets.0.block1.bias": src_state.pop("body.4.block1.bias"), - "adapter.body.2.resnets.0.block2.weight": src_state.pop("body.4.block2.weight"), - "adapter.body.2.resnets.0.block2.bias": src_state.pop("body.4.block2.bias"), - # 2.resnets.1 - "adapter.body.2.resnets.1.block1.weight": src_state.pop("body.5.block1.weight"), - "adapter.body.2.resnets.1.block1.bias": src_state.pop("body.5.block1.bias"), - "adapter.body.2.resnets.1.block2.weight": src_state.pop("body.5.block2.weight"), - "adapter.body.2.resnets.1.block2.bias": src_state.pop("body.5.block2.bias"), - # 3.resnets.0 - "adapter.body.3.resnets.0.block1.weight": src_state.pop("body.6.block1.weight"), - "adapter.body.3.resnets.0.block1.bias": src_state.pop("body.6.block1.bias"), - "adapter.body.3.resnets.0.block2.weight": src_state.pop("body.6.block2.weight"), - "adapter.body.3.resnets.0.block2.bias": src_state.pop("body.6.block2.bias"), - # 3.resnets.1 - "adapter.body.3.resnets.1.block1.weight": src_state.pop("body.7.block1.weight"), - "adapter.body.3.resnets.1.block1.bias": src_state.pop("body.7.block1.bias"), - "adapter.body.3.resnets.1.block2.weight": src_state.pop("body.7.block2.weight"), - "adapter.body.3.resnets.1.block2.bias": src_state.pop("body.7.block2.bias"), - } - - assert len(src_state) == 0 - - adapter = T2IAdapter(in_channels=in_channels, adapter_type="full_adapter") - - adapter.load_state_dict(res_state) - - return adapter - - -def convert_light_adapter(src_state): - original_body_length = max([int(x.split(".")[1]) for x in src_state.keys() if "body." in x]) + 1 - - assert original_body_length == 4 - - res_state = { - # body.0.in_conv - "adapter.body.0.in_conv.weight": src_state.pop("body.0.in_conv.weight"), - "adapter.body.0.in_conv.bias": src_state.pop("body.0.in_conv.bias"), - # body.0.resnets.0 - "adapter.body.0.resnets.0.block1.weight": src_state.pop("body.0.body.0.block1.weight"), - "adapter.body.0.resnets.0.block1.bias": src_state.pop("body.0.body.0.block1.bias"), - "adapter.body.0.resnets.0.block2.weight": src_state.pop("body.0.body.0.block2.weight"), - "adapter.body.0.resnets.0.block2.bias": src_state.pop("body.0.body.0.block2.bias"), - # body.0.resnets.1 - "adapter.body.0.resnets.1.block1.weight": src_state.pop("body.0.body.1.block1.weight"), - "adapter.body.0.resnets.1.block1.bias": src_state.pop("body.0.body.1.block1.bias"), - "adapter.body.0.resnets.1.block2.weight": src_state.pop("body.0.body.1.block2.weight"), - "adapter.body.0.resnets.1.block2.bias": src_state.pop("body.0.body.1.block2.bias"), - # body.0.resnets.2 - "adapter.body.0.resnets.2.block1.weight": src_state.pop("body.0.body.2.block1.weight"), - "adapter.body.0.resnets.2.block1.bias": src_state.pop("body.0.body.2.block1.bias"), - "adapter.body.0.resnets.2.block2.weight": src_state.pop("body.0.body.2.block2.weight"), - "adapter.body.0.resnets.2.block2.bias": src_state.pop("body.0.body.2.block2.bias"), - # body.0.resnets.3 - "adapter.body.0.resnets.3.block1.weight": src_state.pop("body.0.body.3.block1.weight"), - "adapter.body.0.resnets.3.block1.bias": src_state.pop("body.0.body.3.block1.bias"), - "adapter.body.0.resnets.3.block2.weight": src_state.pop("body.0.body.3.block2.weight"), - "adapter.body.0.resnets.3.block2.bias": src_state.pop("body.0.body.3.block2.bias"), - # body.0.out_conv - "adapter.body.0.out_conv.weight": src_state.pop("body.0.out_conv.weight"), - "adapter.body.0.out_conv.bias": src_state.pop("body.0.out_conv.bias"), - # body.1.in_conv - "adapter.body.1.in_conv.weight": src_state.pop("body.1.in_conv.weight"), - "adapter.body.1.in_conv.bias": src_state.pop("body.1.in_conv.bias"), - # body.1.resnets.0 - "adapter.body.1.resnets.0.block1.weight": src_state.pop("body.1.body.0.block1.weight"), - "adapter.body.1.resnets.0.block1.bias": src_state.pop("body.1.body.0.block1.bias"), - "adapter.body.1.resnets.0.block2.weight": src_state.pop("body.1.body.0.block2.weight"), - "adapter.body.1.resnets.0.block2.bias": src_state.pop("body.1.body.0.block2.bias"), - # body.1.resnets.1 - "adapter.body.1.resnets.1.block1.weight": src_state.pop("body.1.body.1.block1.weight"), - "adapter.body.1.resnets.1.block1.bias": src_state.pop("body.1.body.1.block1.bias"), - "adapter.body.1.resnets.1.block2.weight": src_state.pop("body.1.body.1.block2.weight"), - "adapter.body.1.resnets.1.block2.bias": src_state.pop("body.1.body.1.block2.bias"), - # body.1.body.2 - "adapter.body.1.resnets.2.block1.weight": src_state.pop("body.1.body.2.block1.weight"), - "adapter.body.1.resnets.2.block1.bias": src_state.pop("body.1.body.2.block1.bias"), - "adapter.body.1.resnets.2.block2.weight": src_state.pop("body.1.body.2.block2.weight"), - "adapter.body.1.resnets.2.block2.bias": src_state.pop("body.1.body.2.block2.bias"), - # body.1.body.3 - "adapter.body.1.resnets.3.block1.weight": src_state.pop("body.1.body.3.block1.weight"), - "adapter.body.1.resnets.3.block1.bias": src_state.pop("body.1.body.3.block1.bias"), - "adapter.body.1.resnets.3.block2.weight": src_state.pop("body.1.body.3.block2.weight"), - "adapter.body.1.resnets.3.block2.bias": src_state.pop("body.1.body.3.block2.bias"), - # body.1.out_conv - "adapter.body.1.out_conv.weight": src_state.pop("body.1.out_conv.weight"), - "adapter.body.1.out_conv.bias": src_state.pop("body.1.out_conv.bias"), - # body.2.in_conv - "adapter.body.2.in_conv.weight": src_state.pop("body.2.in_conv.weight"), - "adapter.body.2.in_conv.bias": src_state.pop("body.2.in_conv.bias"), - # body.2.body.0 - "adapter.body.2.resnets.0.block1.weight": src_state.pop("body.2.body.0.block1.weight"), - "adapter.body.2.resnets.0.block1.bias": src_state.pop("body.2.body.0.block1.bias"), - "adapter.body.2.resnets.0.block2.weight": src_state.pop("body.2.body.0.block2.weight"), - "adapter.body.2.resnets.0.block2.bias": src_state.pop("body.2.body.0.block2.bias"), - # body.2.body.1 - "adapter.body.2.resnets.1.block1.weight": src_state.pop("body.2.body.1.block1.weight"), - "adapter.body.2.resnets.1.block1.bias": src_state.pop("body.2.body.1.block1.bias"), - "adapter.body.2.resnets.1.block2.weight": src_state.pop("body.2.body.1.block2.weight"), - "adapter.body.2.resnets.1.block2.bias": src_state.pop("body.2.body.1.block2.bias"), - # body.2.body.2 - "adapter.body.2.resnets.2.block1.weight": src_state.pop("body.2.body.2.block1.weight"), - "adapter.body.2.resnets.2.block1.bias": src_state.pop("body.2.body.2.block1.bias"), - "adapter.body.2.resnets.2.block2.weight": src_state.pop("body.2.body.2.block2.weight"), - "adapter.body.2.resnets.2.block2.bias": src_state.pop("body.2.body.2.block2.bias"), - # body.2.body.3 - "adapter.body.2.resnets.3.block1.weight": src_state.pop("body.2.body.3.block1.weight"), - "adapter.body.2.resnets.3.block1.bias": src_state.pop("body.2.body.3.block1.bias"), - "adapter.body.2.resnets.3.block2.weight": src_state.pop("body.2.body.3.block2.weight"), - "adapter.body.2.resnets.3.block2.bias": src_state.pop("body.2.body.3.block2.bias"), - # body.2.out_conv - "adapter.body.2.out_conv.weight": src_state.pop("body.2.out_conv.weight"), - "adapter.body.2.out_conv.bias": src_state.pop("body.2.out_conv.bias"), - # body.3.in_conv - "adapter.body.3.in_conv.weight": src_state.pop("body.3.in_conv.weight"), - "adapter.body.3.in_conv.bias": src_state.pop("body.3.in_conv.bias"), - # body.3.body.0 - "adapter.body.3.resnets.0.block1.weight": src_state.pop("body.3.body.0.block1.weight"), - "adapter.body.3.resnets.0.block1.bias": src_state.pop("body.3.body.0.block1.bias"), - "adapter.body.3.resnets.0.block2.weight": src_state.pop("body.3.body.0.block2.weight"), - "adapter.body.3.resnets.0.block2.bias": src_state.pop("body.3.body.0.block2.bias"), - # body.3.body.1 - "adapter.body.3.resnets.1.block1.weight": src_state.pop("body.3.body.1.block1.weight"), - "adapter.body.3.resnets.1.block1.bias": src_state.pop("body.3.body.1.block1.bias"), - "adapter.body.3.resnets.1.block2.weight": src_state.pop("body.3.body.1.block2.weight"), - "adapter.body.3.resnets.1.block2.bias": src_state.pop("body.3.body.1.block2.bias"), - # body.3.body.2 - "adapter.body.3.resnets.2.block1.weight": src_state.pop("body.3.body.2.block1.weight"), - "adapter.body.3.resnets.2.block1.bias": src_state.pop("body.3.body.2.block1.bias"), - "adapter.body.3.resnets.2.block2.weight": src_state.pop("body.3.body.2.block2.weight"), - "adapter.body.3.resnets.2.block2.bias": src_state.pop("body.3.body.2.block2.bias"), - # body.3.body.3 - "adapter.body.3.resnets.3.block1.weight": src_state.pop("body.3.body.3.block1.weight"), - "adapter.body.3.resnets.3.block1.bias": src_state.pop("body.3.body.3.block1.bias"), - "adapter.body.3.resnets.3.block2.weight": src_state.pop("body.3.body.3.block2.weight"), - "adapter.body.3.resnets.3.block2.bias": src_state.pop("body.3.body.3.block2.bias"), - # body.3.out_conv - "adapter.body.3.out_conv.weight": src_state.pop("body.3.out_conv.weight"), - "adapter.body.3.out_conv.bias": src_state.pop("body.3.out_conv.bias"), - } - - assert len(src_state) == 0 - - adapter = T2IAdapter(in_channels=3, channels=[320, 640, 1280], num_res_blocks=4, adapter_type="light_adapter") - - adapter.load_state_dict(res_state) - - return adapter - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--output_path", default=None, type=str, required=True, help="Path to the store the result checkpoint." - ) - parser.add_argument( - "--is_adapter_light", - action="store_true", - help="Is checkpoint come from Adapter-Light architecture. ex: color-adapter", - ) - parser.add_argument("--in_channels", required=False, type=int, help="Input channels for non-light adapter") - - args = parser.parse_args() - src_state = torch.load(args.checkpoint_path) - - if args.is_adapter_light: - adapter = convert_light_adapter(src_state) - else: - if args.in_channels is None: - raise ValueError("set `--in_channels=`") - adapter = convert_adapter(src_state, args.in_channels) - - adapter.save_pretrained(args.output_path) diff --git a/scripts/convert_ovis_image_to_diffusers.py b/scripts/convert_ovis_image_to_diffusers.py deleted file mode 100644 index 0d3d9cd44bf6..000000000000 --- a/scripts/convert_ovis_image_to_diffusers.py +++ /dev/null @@ -1,263 +0,0 @@ -import argparse -from contextlib import nullcontext - -import safetensors.torch -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download - -from diffusers import OvisImageTransformer2DModel -from diffusers.utils.import_utils import is_accelerate_available - - -""" -# Transformer - -python scripts/convert_ovis_image_to_diffusers.py \ ---original_state_dict_repo_id "AIDC-AI/Ovis-Image-7B" \ ---filename "ovis_image.safetensors" ---output_path "ovis-image" \ ---transformer -""" - - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - -parser = argparse.ArgumentParser() -parser.add_argument("--original_state_dict_repo_id", default=None, type=str) -parser.add_argument("--filename", default="ovis_image.safetensors", type=str) -parser.add_argument("--checkpoint_path", default=None, type=str) -parser.add_argument("--in_channels", type=int, default=64) -parser.add_argument("--out_channels", type=int, default=None) -parser.add_argument("--transformer", action="store_true") -parser.add_argument("--output_path", type=str) -parser.add_argument("--dtype", type=str, default="bf16") - -args = parser.parse_args() -dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - - -def load_original_checkpoint(args): - if args.original_state_dict_repo_id is not None: - ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=args.filename) - elif args.checkpoint_path is not None: - ckpt_path = args.checkpoint_path - else: - raise ValueError(" please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") - - original_state_dict = safetensors.torch.load_file(ckpt_path) - return original_state_dict - - -# in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; -# while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation -def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_ovis_image_transformer_checkpoint_to_diffusers( - original_state_dict, num_layers, num_single_layers, inner_dim, mlp_ratio=4.0 -): - converted_state_dict = {} - - ## time_text_embed.timestep_embedder <- time_in - converted_state_dict["timestep_embedder.linear_1.weight"] = original_state_dict.pop("time_in.in_layer.weight") - converted_state_dict["timestep_embedder.linear_1.bias"] = original_state_dict.pop("time_in.in_layer.bias") - converted_state_dict["timestep_embedder.linear_2.weight"] = original_state_dict.pop("time_in.out_layer.weight") - converted_state_dict["timestep_embedder.linear_2.bias"] = original_state_dict.pop("time_in.out_layer.bias") - - # context_embedder - converted_state_dict["context_embedder_norm.weight"] = original_state_dict.pop("semantic_txt_norm.weight") - converted_state_dict["context_embedder.weight"] = original_state_dict.pop("semantic_txt_in.weight") - converted_state_dict["context_embedder.bias"] = original_state_dict.pop("semantic_txt_in.bias") - - # x_embedder - converted_state_dict["x_embedder.weight"] = original_state_dict.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = original_state_dict.pop("img_in.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - # norms. - ## norm1 - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mod.lin.bias" - ) - ## norm1_context - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mod.lin.bias" - ) - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0 - ) - context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - # qk_norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.query_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.norm.key_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.query_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.norm.key_norm.weight" - ) - # ff img_mlp - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = torch.cat( - [ - original_state_dict.pop(f"double_blocks.{i}.img_mlp.up_proj.weight"), - original_state_dict.pop(f"double_blocks.{i}.img_mlp.gate_proj.weight"), - ], - dim=0, - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = torch.cat( - [ - original_state_dict.pop(f"double_blocks.{i}.img_mlp.up_proj.bias"), - original_state_dict.pop(f"double_blocks.{i}.img_mlp.gate_proj.bias"), - ], - dim=0, - ) - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.down_proj.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_mlp.down_proj.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = torch.cat( - [ - original_state_dict.pop(f"double_blocks.{i}.txt_mlp.up_proj.weight"), - original_state_dict.pop(f"double_blocks.{i}.txt_mlp.gate_proj.weight"), - ], - dim=0, - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = torch.cat( - [ - original_state_dict.pop(f"double_blocks.{i}.txt_mlp.up_proj.bias"), - original_state_dict.pop(f"double_blocks.{i}.txt_mlp.gate_proj.bias"), - ], - dim=0, - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.down_proj.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_mlp.down_proj.bias" - ) - # output projections. - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = original_state_dict.pop( - f"double_blocks.{i}.img_attn.proj.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = original_state_dict.pop( - f"double_blocks.{i}.txt_attn.proj.bias" - ) - - # single transformer blocks - for i in range(num_single_layers): - block_prefix = f"single_transformer_blocks.{i}." - # norm.linear <- single_blocks.0.modulation.lin - converted_state_dict[f"{block_prefix}norm.linear.weight"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm.linear.bias"] = original_state_dict.pop( - f"single_blocks.{i}.modulation.lin.bias" - ) - # Q, K, V, mlp - mlp_hidden_dim = int(inner_dim * mlp_ratio) - split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim * 2) - q, k, v, mlp = torch.split(original_state_dict.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0) - q_bias, k_bias, v_bias, mlp_bias = torch.split( - original_state_dict.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp]) - converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias]) - # qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.query_norm.weight" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = original_state_dict.pop( - f"single_blocks.{i}.norm.key_norm.weight" - ) - # output projections. - converted_state_dict[f"{block_prefix}proj_out.weight"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.weight" - ) - converted_state_dict[f"{block_prefix}proj_out.bias"] = original_state_dict.pop( - f"single_blocks.{i}.linear2.bias" - ) - - converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.weight") - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.bias") - ) - - return converted_state_dict - - -def main(args): - original_ckpt = load_original_checkpoint(args) - - if args.transformer: - num_layers = 6 - num_single_layers = 27 - inner_dim = 3072 - mlp_ratio = 4.0 - - converted_transformer_state_dict = convert_ovis_image_transformer_checkpoint_to_diffusers( - original_ckpt, num_layers, num_single_layers, inner_dim, mlp_ratio=mlp_ratio - ) - transformer = OvisImageTransformer2DModel(in_channels=args.in_channels, out_channels=args.out_channels) - transformer.load_state_dict(converted_transformer_state_dict, strict=True) - - print("Saving Ovis-Image Transformer in Diffusers format.") - transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") - - -if __name__ == "__main__": - main(args) diff --git a/scripts/convert_pixart_alpha_to_diffusers.py b/scripts/convert_pixart_alpha_to_diffusers.py deleted file mode 100644 index 228b479df090..000000000000 --- a/scripts/convert_pixart_alpha_to_diffusers.py +++ /dev/null @@ -1,198 +0,0 @@ -import argparse -import os - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, PixArtAlphaPipeline, Transformer2DModel - - -ckpt_id = "PixArt-alpha/PixArt-alpha" -# https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/scripts/inference.py#L125 -interpolation_scale = {256: 0.5, 512: 1, 1024: 2} - - -def main(args): - all_state_dict = torch.load(args.orig_ckpt_path, map_location="cpu") - state_dict = all_state_dict.pop("state_dict") - converted_state_dict = {} - - # Patch embeddings. - converted_state_dict["pos_embed.proj.weight"] = state_dict.pop("x_embedder.proj.weight") - converted_state_dict["pos_embed.proj.bias"] = state_dict.pop("x_embedder.proj.bias") - - # Caption projection. - converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") - - # AdaLN-single LN - converted_state_dict["adaln_single.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") - converted_state_dict["adaln_single.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") - - if args.image_size == 1024: - # Resolution. - converted_state_dict["adaln_single.emb.resolution_embedder.linear_1.weight"] = state_dict.pop( - "csize_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_1.bias"] = state_dict.pop( - "csize_embedder.mlp.0.bias" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_2.weight"] = state_dict.pop( - "csize_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_2.bias"] = state_dict.pop( - "csize_embedder.mlp.2.bias" - ) - # Aspect ratio. - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_1.weight"] = state_dict.pop( - "ar_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_1.bias"] = state_dict.pop( - "ar_embedder.mlp.0.bias" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_2.weight"] = state_dict.pop( - "ar_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_2.bias"] = state_dict.pop( - "ar_embedder.mlp.2.bias" - ) - # Shared norm. - converted_state_dict["adaln_single.linear.weight"] = state_dict.pop("t_block.1.weight") - converted_state_dict["adaln_single.linear.bias"] = state_dict.pop("t_block.1.bias") - - for depth in range(28): - # Transformer blocks. - converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( - f"blocks.{depth}.scale_shift_table" - ) - - # Attention is all you need 🤘 - - # Self attention. - q, k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.weight"), 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.bias"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.bias"] = v_bias - # Projection. - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.attn.proj.bias" - ) - - # Feed-forward. - converted_state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.2.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.2.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.fc2.bias" - ) - - # Cross-attention. - q = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.weight") - q_bias = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.bias") - k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.weight"), 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.bias"), 2, dim=0) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.bias" - ) - - # Final block. - converted_state_dict["proj_out.weight"] = state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = state_dict.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = state_dict.pop("final_layer.scale_shift_table") - - # DiT XL/2 - transformer = Transformer2DModel( - sample_size=args.image_size // 8, - num_layers=28, - attention_head_dim=72, - in_channels=4, - out_channels=8, - patch_size=2, - attention_bias=True, - num_attention_heads=16, - cross_attention_dim=1152, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - caption_channels=4096, - ) - transformer.load_state_dict(converted_state_dict, strict=True) - - assert transformer.pos_embed.pos_embed is not None - state_dict.pop("pos_embed") - state_dict.pop("y_embedder.y_embedding") - assert len(state_dict) == 0, f"State dict is not empty, {state_dict.keys()}" - - num_model_params = sum(p.numel() for p in transformer.parameters()) - print(f"Total number of transformer parameters: {num_model_params}") - - if args.only_transformer: - transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) - else: - scheduler = DPMSolverMultistepScheduler() - - vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="sd-vae-ft-ema") - - tokenizer = T5Tokenizer.from_pretrained(ckpt_id, subfolder="t5-v1_1-xxl") - text_encoder = T5EncoderModel.from_pretrained(ckpt_id, subfolder="t5-v1_1-xxl") - - pipeline = PixArtAlphaPipeline( - tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler - ) - - pipeline.save_pretrained(args.dump_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--orig_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--image_size", - default=1024, - type=int, - choices=[256, 512, 1024], - required=False, - help="Image size of pretrained model, either 512 or 1024.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") - parser.add_argument("--only_transformer", default=True, type=bool, required=True) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_pixart_sigma_to_diffusers.py b/scripts/convert_pixart_sigma_to_diffusers.py deleted file mode 100644 index 9572a83c0644..000000000000 --- a/scripts/convert_pixart_sigma_to_diffusers.py +++ /dev/null @@ -1,223 +0,0 @@ -import argparse -import os - -import torch -from transformers import T5EncoderModel, T5Tokenizer - -from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, PixArtSigmaPipeline, Transformer2DModel - - -ckpt_id = "PixArt-alpha" -# https://github.com/PixArt-alpha/PixArt-sigma/blob/dd087141864e30ec44f12cb7448dd654be065e88/scripts/inference.py#L158 -interpolation_scale = {256: 0.5, 512: 1, 1024: 2, 2048: 4} - - -def main(args): - all_state_dict = torch.load(args.orig_ckpt_path) - state_dict = all_state_dict.pop("state_dict") - converted_state_dict = {} - - # Patch embeddings. - converted_state_dict["pos_embed.proj.weight"] = state_dict.pop("x_embedder.proj.weight") - converted_state_dict["pos_embed.proj.bias"] = state_dict.pop("x_embedder.proj.bias") - - # Caption projection. - converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") - - # AdaLN-single LN - converted_state_dict["adaln_single.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") - converted_state_dict["adaln_single.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") - - if args.micro_condition: - # Resolution. - converted_state_dict["adaln_single.emb.resolution_embedder.linear_1.weight"] = state_dict.pop( - "csize_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_1.bias"] = state_dict.pop( - "csize_embedder.mlp.0.bias" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_2.weight"] = state_dict.pop( - "csize_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.resolution_embedder.linear_2.bias"] = state_dict.pop( - "csize_embedder.mlp.2.bias" - ) - # Aspect ratio. - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_1.weight"] = state_dict.pop( - "ar_embedder.mlp.0.weight" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_1.bias"] = state_dict.pop( - "ar_embedder.mlp.0.bias" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_2.weight"] = state_dict.pop( - "ar_embedder.mlp.2.weight" - ) - converted_state_dict["adaln_single.emb.aspect_ratio_embedder.linear_2.bias"] = state_dict.pop( - "ar_embedder.mlp.2.bias" - ) - # Shared norm. - converted_state_dict["adaln_single.linear.weight"] = state_dict.pop("t_block.1.weight") - converted_state_dict["adaln_single.linear.bias"] = state_dict.pop("t_block.1.bias") - - for depth in range(28): - # Transformer blocks. - converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( - f"blocks.{depth}.scale_shift_table" - ) - # Attention is all you need 🤘 - - # Self attention. - q, k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.weight"), 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.bias"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.bias"] = v_bias - # Projection. - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.attn.proj.bias" - ) - if args.qk_norm: - converted_state_dict[f"transformer_blocks.{depth}.attn1.q_norm.weight"] = state_dict.pop( - f"blocks.{depth}.attn.q_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.q_norm.bias"] = state_dict.pop( - f"blocks.{depth}.attn.q_norm.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.k_norm.weight"] = state_dict.pop( - f"blocks.{depth}.attn.k_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.k_norm.bias"] = state_dict.pop( - f"blocks.{depth}.attn.k_norm.bias" - ) - - # Feed-forward. - converted_state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.0.proj.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.2.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.net.2.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.fc2.bias" - ) - - # Cross-attention. - q = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.weight") - q_bias = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.bias") - k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.weight"), 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.bias"), 2, dim=0) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.bias" - ) - - # Final block. - converted_state_dict["proj_out.weight"] = state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = state_dict.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = state_dict.pop("final_layer.scale_shift_table") - - # PixArt XL/2 - transformer = Transformer2DModel( - sample_size=args.image_size // 8, - num_layers=28, - attention_head_dim=72, - in_channels=4, - out_channels=8, - patch_size=2, - attention_bias=True, - num_attention_heads=16, - cross_attention_dim=1152, - activation_fn="gelu-approximate", - num_embeds_ada_norm=1000, - norm_type="ada_norm_single", - norm_elementwise_affine=False, - norm_eps=1e-6, - caption_channels=4096, - interpolation_scale=interpolation_scale[args.image_size], - use_additional_conditions=args.micro_condition, - ) - transformer.load_state_dict(converted_state_dict, strict=True) - - assert transformer.pos_embed.pos_embed is not None - try: - state_dict.pop("y_embedder.y_embedding") - state_dict.pop("pos_embed") - except Exception as e: - print(f"Skipping {str(e)}") - pass - assert len(state_dict) == 0, f"State dict is not empty, {state_dict.keys()}" - - num_model_params = sum(p.numel() for p in transformer.parameters()) - print(f"Total number of transformer parameters: {num_model_params}") - - if args.only_transformer: - transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) - else: - # pixart-Sigma vae link: https://huggingface.co/PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers/tree/main/vae - vae = AutoencoderKL.from_pretrained(f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="vae") - - scheduler = DPMSolverMultistepScheduler() - - tokenizer = T5Tokenizer.from_pretrained(f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="tokenizer") - text_encoder = T5EncoderModel.from_pretrained( - f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="text_encoder" - ) - - pipeline = PixArtSigmaPipeline( - tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler - ) - - pipeline.save_pretrained(args.dump_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--micro_condition", action="store_true", help="If use Micro-condition in PixArtMS structure during training." - ) - parser.add_argument("--qk_norm", action="store_true", help="If use qk norm during training.") - parser.add_argument( - "--orig_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--image_size", - default=1024, - type=int, - choices=[256, 512, 1024, 2048], - required=False, - help="Image size of pretrained model, 256, 512, 1024, or 2048.", - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") - parser.add_argument("--only_transformer", default=True, type=bool, required=True) - - args = parser.parse_args() - main(args) diff --git a/scripts/convert_sana_controlnet_to_diffusers.py b/scripts/convert_sana_controlnet_to_diffusers.py deleted file mode 100644 index f7fcd7252576..000000000000 --- a/scripts/convert_sana_controlnet_to_diffusers.py +++ /dev/null @@ -1,216 +0,0 @@ -#!/usr/bin/env python -from __future__ import annotations - -import argparse -from contextlib import nullcontext - -import torch -from accelerate import init_empty_weights - -from diffusers import ( - SanaControlNetModel, -) -from diffusers.models.model_loading_utils import load_model_dict_into_meta -from diffusers.utils.import_utils import is_accelerate_available - - -CTX = init_empty_weights if is_accelerate_available else nullcontext - - -def main(args): - file_path = args.orig_ckpt_path - - all_state_dict = torch.load(file_path, weights_only=True) - state_dict = all_state_dict.pop("state_dict") - converted_state_dict = {} - - # Patch embeddings. - converted_state_dict["patch_embed.proj.weight"] = state_dict.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = state_dict.pop("x_embedder.proj.bias") - - # Caption projection. - converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") - - # AdaLN-single LN - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") - - # Shared norm. - converted_state_dict["time_embed.linear.weight"] = state_dict.pop("t_block.1.weight") - converted_state_dict["time_embed.linear.bias"] = state_dict.pop("t_block.1.bias") - - # y norm - converted_state_dict["caption_norm.weight"] = state_dict.pop("attention_y_norm.weight") - - # Positional embedding interpolation scale. - interpolation_scale = {512: None, 1024: None, 2048: 1.0, 4096: 2.0} - - # ControlNet Input Projection. - converted_state_dict["input_block.weight"] = state_dict.pop("controlnet.0.before_proj.weight") - converted_state_dict["input_block.bias"] = state_dict.pop("controlnet.0.before_proj.bias") - - for depth in range(7): - # Transformer blocks. - converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( - f"controlnet.{depth}.copied_block.scale_shift_table" - ) - - # Linear Attention is all you need 🤘 - # Self attention. - q, k, v = torch.chunk(state_dict.pop(f"controlnet.{depth}.copied_block.attn.qkv.weight"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - # Projection. - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( - f"controlnet.{depth}.copied_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( - f"controlnet.{depth}.copied_block.attn.proj.bias" - ) - - # Feed-forward. - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.weight"] = state_dict.pop( - f"controlnet.{depth}.copied_block.mlp.inverted_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.bias"] = state_dict.pop( - f"controlnet.{depth}.copied_block.mlp.inverted_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.weight"] = state_dict.pop( - f"controlnet.{depth}.copied_block.mlp.depth_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.bias"] = state_dict.pop( - f"controlnet.{depth}.copied_block.mlp.depth_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_point.weight"] = state_dict.pop( - f"controlnet.{depth}.copied_block.mlp.point_conv.conv.weight" - ) - - # Cross-attention. - q = state_dict.pop(f"controlnet.{depth}.copied_block.cross_attn.q_linear.weight") - q_bias = state_dict.pop(f"controlnet.{depth}.copied_block.cross_attn.q_linear.bias") - k, v = torch.chunk(state_dict.pop(f"controlnet.{depth}.copied_block.cross_attn.kv_linear.weight"), 2, dim=0) - k_bias, v_bias = torch.chunk( - state_dict.pop(f"controlnet.{depth}.copied_block.cross_attn.kv_linear.bias"), 2, dim=0 - ) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( - f"controlnet.{depth}.copied_block.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( - f"controlnet.{depth}.copied_block.cross_attn.proj.bias" - ) - - # ControlNet After Projection - converted_state_dict[f"controlnet_blocks.{depth}.weight"] = state_dict.pop( - f"controlnet.{depth}.after_proj.weight" - ) - converted_state_dict[f"controlnet_blocks.{depth}.bias"] = state_dict.pop(f"controlnet.{depth}.after_proj.bias") - - # ControlNet - with CTX(): - controlnet = SanaControlNetModel( - num_attention_heads=model_kwargs[args.model_type]["num_attention_heads"], - attention_head_dim=model_kwargs[args.model_type]["attention_head_dim"], - num_layers=model_kwargs[args.model_type]["num_layers"], - num_cross_attention_heads=model_kwargs[args.model_type]["num_cross_attention_heads"], - cross_attention_head_dim=model_kwargs[args.model_type]["cross_attention_head_dim"], - cross_attention_dim=model_kwargs[args.model_type]["cross_attention_dim"], - caption_channels=2304, - sample_size=args.image_size // 32, - interpolation_scale=interpolation_scale[args.image_size], - ) - - if is_accelerate_available(): - load_model_dict_into_meta(controlnet, converted_state_dict) - else: - controlnet.load_state_dict(converted_state_dict, strict=True, assign=True) - - num_model_params = sum(p.numel() for p in controlnet.parameters()) - print(f"Total number of controlnet parameters: {num_model_params}") - - controlnet = controlnet.to(weight_dtype) - controlnet.load_state_dict(converted_state_dict, strict=True) - - print(f"Saving Sana ControlNet in Diffusers format in {args.dump_path}.") - controlnet.save_pretrained(args.dump_path) - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - -VARIANT_MAPPING = { - "fp32": None, - "fp16": "fp16", - "bf16": "bf16", -} - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--orig_ckpt_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--image_size", - default=1024, - type=int, - choices=[512, 1024, 2048, 4096], - required=False, - help="Image size of pretrained model, 512, 1024, 2048 or 4096.", - ) - parser.add_argument( - "--model_type", - default="SanaMS_1600M_P1_ControlNet_D7", - type=str, - choices=["SanaMS_1600M_P1_ControlNet_D7", "SanaMS_600M_P1_ControlNet_D7"], - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") - parser.add_argument("--dtype", default="fp16", type=str, choices=["fp32", "fp16", "bf16"], help="Weight dtype.") - - args = parser.parse_args() - - model_kwargs = { - "SanaMS_1600M_P1_ControlNet_D7": { - "num_attention_heads": 70, - "attention_head_dim": 32, - "num_cross_attention_heads": 20, - "cross_attention_head_dim": 112, - "cross_attention_dim": 2240, - "num_layers": 7, - }, - "SanaMS_600M_P1_ControlNet_D7": { - "num_attention_heads": 36, - "attention_head_dim": 32, - "num_cross_attention_heads": 16, - "cross_attention_head_dim": 72, - "cross_attention_dim": 1152, - "num_layers": 7, - }, - } - - device = "cuda" if torch.cuda.is_available() else "cpu" - weight_dtype = DTYPE_MAPPING[args.dtype] - variant = VARIANT_MAPPING[args.dtype] - - main(args) diff --git a/scripts/convert_sd3_controlnet_to_diffusers.py b/scripts/convert_sd3_controlnet_to_diffusers.py deleted file mode 100644 index 171f40a7aa06..000000000000 --- a/scripts/convert_sd3_controlnet_to_diffusers.py +++ /dev/null @@ -1,185 +0,0 @@ -""" -A script to convert Stable Diffusion 3.5 ControlNet checkpoints to the Diffusers format. - -Example: - Convert a SD3.5 ControlNet checkpoint to Diffusers format using local file: - ```bash - python scripts/convert_sd3_controlnet_to_diffusers.py \ - --checkpoint_path "path/to/local/sd3.5_large_controlnet_canny.safetensors" \ - --output_path "output/sd35-controlnet-canny" \ - --dtype "fp16" # optional, defaults to fp32 - ``` - - Or download and convert from HuggingFace repository: - ```bash - python scripts/convert_sd3_controlnet_to_diffusers.py \ - --original_state_dict_repo_id "stabilityai/stable-diffusion-3.5-controlnets" \ - --filename "sd3.5_large_controlnet_canny.safetensors" \ - --output_path "/raid/yiyi/sd35-controlnet-canny-diffusers" \ - --dtype "fp32" # optional, defaults to fp32 - ``` - -Note: - The script supports the following ControlNet types from SD3.5: - - Canny edge detection - - Depth estimation - - Blur detection - - The checkpoint files can be downloaded from: - https://huggingface.co/stabilityai/stable-diffusion-3.5-controlnets -""" - -import argparse - -import safetensors.torch -import torch -from huggingface_hub import hf_hub_download - -from diffusers import SD3ControlNetModel - - -parser = argparse.ArgumentParser() -parser.add_argument("--checkpoint_path", type=str, default=None, help="Path to local checkpoint file") -parser.add_argument( - "--original_state_dict_repo_id", type=str, default=None, help="HuggingFace repo ID containing the checkpoint" -) -parser.add_argument("--filename", type=str, default=None, help="Filename of the checkpoint in the HF repo") -parser.add_argument("--output_path", type=str, required=True, help="Path to save the converted model") -parser.add_argument( - "--dtype", type=str, default="fp32", help="Data type for the converted model (fp16, bf16, or fp32)" -) - -args = parser.parse_args() - - -def load_original_checkpoint(args): - if args.original_state_dict_repo_id is not None: - if args.filename is None: - raise ValueError("When using `original_state_dict_repo_id`, `filename` must also be specified") - print(f"Downloading checkpoint from {args.original_state_dict_repo_id}/{args.filename}") - ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=args.filename) - elif args.checkpoint_path is not None: - print(f"Loading checkpoint from local path: {args.checkpoint_path}") - ckpt_path = args.checkpoint_path - else: - raise ValueError("Please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") - - original_state_dict = safetensors.torch.load_file(ckpt_path) - return original_state_dict - - -def convert_sd3_controlnet_checkpoint_to_diffusers(original_state_dict): - converted_state_dict = {} - - # Direct mappings for controlnet blocks - for i in range(19): # 19 controlnet blocks - converted_state_dict[f"controlnet_blocks.{i}.weight"] = original_state_dict[f"controlnet_blocks.{i}.weight"] - converted_state_dict[f"controlnet_blocks.{i}.bias"] = original_state_dict[f"controlnet_blocks.{i}.bias"] - - # Positional embeddings - converted_state_dict["pos_embed_input.proj.weight"] = original_state_dict["pos_embed_input.proj.weight"] - converted_state_dict["pos_embed_input.proj.bias"] = original_state_dict["pos_embed_input.proj.bias"] - - # Time and text embeddings - time_text_mappings = { - "time_text_embed.timestep_embedder.linear_1.weight": "time_text_embed.timestep_embedder.linear_1.weight", - "time_text_embed.timestep_embedder.linear_1.bias": "time_text_embed.timestep_embedder.linear_1.bias", - "time_text_embed.timestep_embedder.linear_2.weight": "time_text_embed.timestep_embedder.linear_2.weight", - "time_text_embed.timestep_embedder.linear_2.bias": "time_text_embed.timestep_embedder.linear_2.bias", - "time_text_embed.text_embedder.linear_1.weight": "time_text_embed.text_embedder.linear_1.weight", - "time_text_embed.text_embedder.linear_1.bias": "time_text_embed.text_embedder.linear_1.bias", - "time_text_embed.text_embedder.linear_2.weight": "time_text_embed.text_embedder.linear_2.weight", - "time_text_embed.text_embedder.linear_2.bias": "time_text_embed.text_embedder.linear_2.bias", - } - - for new_key, old_key in time_text_mappings.items(): - if old_key in original_state_dict: - converted_state_dict[new_key] = original_state_dict[old_key] - - # Transformer blocks - for i in range(19): - # Split QKV into separate Q, K, V - qkv_weight = original_state_dict[f"transformer_blocks.{i}.attn.qkv.weight"] - qkv_bias = original_state_dict[f"transformer_blocks.{i}.attn.qkv.bias"] - q, k, v = torch.chunk(qkv_weight, 3, dim=0) - q_bias, k_bias, v_bias = torch.chunk(qkv_bias, 3, dim=0) - - block_mappings = { - f"transformer_blocks.{i}.attn.to_q.weight": q, - f"transformer_blocks.{i}.attn.to_q.bias": q_bias, - f"transformer_blocks.{i}.attn.to_k.weight": k, - f"transformer_blocks.{i}.attn.to_k.bias": k_bias, - f"transformer_blocks.{i}.attn.to_v.weight": v, - f"transformer_blocks.{i}.attn.to_v.bias": v_bias, - # Output projections - f"transformer_blocks.{i}.attn.to_out.0.weight": original_state_dict[ - f"transformer_blocks.{i}.attn.proj.weight" - ], - f"transformer_blocks.{i}.attn.to_out.0.bias": original_state_dict[ - f"transformer_blocks.{i}.attn.proj.bias" - ], - # Feed forward - f"transformer_blocks.{i}.ff.net.0.proj.weight": original_state_dict[ - f"transformer_blocks.{i}.mlp.fc1.weight" - ], - f"transformer_blocks.{i}.ff.net.0.proj.bias": original_state_dict[f"transformer_blocks.{i}.mlp.fc1.bias"], - f"transformer_blocks.{i}.ff.net.2.weight": original_state_dict[f"transformer_blocks.{i}.mlp.fc2.weight"], - f"transformer_blocks.{i}.ff.net.2.bias": original_state_dict[f"transformer_blocks.{i}.mlp.fc2.bias"], - # Norms - f"transformer_blocks.{i}.norm1.linear.weight": original_state_dict[ - f"transformer_blocks.{i}.adaLN_modulation.1.weight" - ], - f"transformer_blocks.{i}.norm1.linear.bias": original_state_dict[ - f"transformer_blocks.{i}.adaLN_modulation.1.bias" - ], - } - converted_state_dict.update(block_mappings) - - return converted_state_dict - - -def main(args): - original_ckpt = load_original_checkpoint(args) - original_dtype = next(iter(original_ckpt.values())).dtype - - # Initialize dtype with fp32 as default - if args.dtype == "fp16": - dtype = torch.float16 - elif args.dtype == "bf16": - dtype = torch.bfloat16 - elif args.dtype == "fp32": - dtype = torch.float32 - else: - raise ValueError(f"Unsupported dtype: {args.dtype}. Must be one of: fp16, bf16, fp32") - - if dtype != original_dtype: - print( - f"Converting checkpoint from {original_dtype} to {dtype}. This can lead to unexpected results, proceed with caution." - ) - - converted_controlnet_state_dict = convert_sd3_controlnet_checkpoint_to_diffusers(original_ckpt) - - controlnet = SD3ControlNetModel( - patch_size=2, - in_channels=16, - num_layers=19, - attention_head_dim=64, - num_attention_heads=38, - joint_attention_dim=None, - caption_projection_dim=2048, - pooled_projection_dim=2048, - out_channels=16, - pos_embed_max_size=None, - pos_embed_type=None, - use_pos_embed=False, - force_zeros_for_pooled_projection=False, - ) - - controlnet.load_state_dict(converted_controlnet_state_dict, strict=True) - - print(f"Saving SD3 ControlNet in Diffusers format in {args.output_path}.") - controlnet.to(dtype).save_pretrained(args.output_path) - - -if __name__ == "__main__": - main(args) diff --git a/scripts/convert_sd3_to_diffusers.py b/scripts/convert_sd3_to_diffusers.py deleted file mode 100644 index 83cb436e6e32..000000000000 --- a/scripts/convert_sd3_to_diffusers.py +++ /dev/null @@ -1,351 +0,0 @@ -import argparse -from contextlib import nullcontext - -import safetensors.torch -import torch -from accelerate import init_empty_weights - -from diffusers import AutoencoderKL, SD3Transformer2DModel -from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint -from diffusers.models.model_loading_utils import load_model_dict_into_meta -from diffusers.utils.import_utils import is_accelerate_available - - -CTX = init_empty_weights if is_accelerate_available() else nullcontext - -parser = argparse.ArgumentParser() -parser.add_argument("--checkpoint_path", type=str) -parser.add_argument("--output_path", type=str) -parser.add_argument("--dtype", type=str) - -args = parser.parse_args() - - -def load_original_checkpoint(ckpt_path): - original_state_dict = safetensors.torch.load_file(ckpt_path) - keys = list(original_state_dict.keys()) - for k in keys: - if "model.diffusion_model." in k: - original_state_dict[k.replace("model.diffusion_model.", "")] = original_state_dict.pop(k) - - return original_state_dict - - -# in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; -# while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation -def swap_scale_shift(weight, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def convert_sd3_transformer_checkpoint_to_diffusers( - original_state_dict, num_layers, caption_projection_dim, dual_attention_layers, has_qk_norm -): - converted_state_dict = {} - - # Positional and patch embeddings. - converted_state_dict["pos_embed.pos_embed"] = original_state_dict.pop("pos_embed") - converted_state_dict["pos_embed.proj.weight"] = original_state_dict.pop("x_embedder.proj.weight") - converted_state_dict["pos_embed.proj.bias"] = original_state_dict.pop("x_embedder.proj.bias") - - # Timestep embeddings. - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "t_embedder.mlp.0.bias" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "t_embedder.mlp.2.bias" - ) - - # Context projections. - converted_state_dict["context_embedder.weight"] = original_state_dict.pop("context_embedder.weight") - converted_state_dict["context_embedder.bias"] = original_state_dict.pop("context_embedder.bias") - - # Pooled context projection. - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = original_state_dict.pop( - "y_embedder.mlp.0.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = original_state_dict.pop( - "y_embedder.mlp.0.bias" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = original_state_dict.pop( - "y_embedder.mlp.2.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = original_state_dict.pop( - "y_embedder.mlp.2.bias" - ) - - # Transformer blocks 🎸. - for i in range(num_layers): - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.x_block.attn.qkv.weight"), 3, dim=0 - ) - context_q, context_k, context_v = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.context_block.attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.x_block.attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.context_block.attn.qkv.bias"), 3, dim=0 - ) - - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.bias"] = torch.cat([sample_v_bias]) - - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn.norm_q.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_k.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn.ln_k.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_q.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_k.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.attn.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn.proj.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.attn.proj.bias" - ) - - # attn2 - if i in dual_attention_layers: - # Q, K, V - sample_q2, sample_k2, sample_v2 = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.x_block.attn2.qkv.weight"), 3, dim=0 - ) - sample_q2_bias, sample_k2_bias, sample_v2_bias = torch.chunk( - original_state_dict.pop(f"joint_blocks.{i}.x_block.attn2.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.weight"] = torch.cat([sample_q2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.bias"] = torch.cat([sample_q2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.weight"] = torch.cat([sample_k2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.bias"] = torch.cat([sample_k2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.weight"] = torch.cat([sample_v2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.bias"] = torch.cat([sample_v2_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_q.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn2.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_k.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn2.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn2.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.attn2.proj.bias" - ) - - # norms. - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias" - ) - else: - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = swap_scale_shift( - original_state_dict.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight"), - dim=caption_projection_dim, - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = swap_scale_shift( - original_state_dict.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias"), - dim=caption_projection_dim, - ) - - # ffs. - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.weight"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.bias"] = original_state_dict.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.bias" - ) - - # Final blocks. - converted_state_dict["proj_out.weight"] = original_state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = original_state_dict.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.weight"), dim=caption_projection_dim - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("final_layer.adaLN_modulation.1.bias"), dim=caption_projection_dim - ) - - return converted_state_dict - - -def is_vae_in_checkpoint(original_state_dict): - return ("first_stage_model.decoder.conv_in.weight" in original_state_dict) and ( - "first_stage_model.encoder.conv_in.weight" in original_state_dict - ) - - -def get_attn2_layers(state_dict): - attn2_layers = [] - for key in state_dict.keys(): - if "attn2." in key: - # Extract the layer number from the key - layer_num = int(key.split(".")[1]) - attn2_layers.append(layer_num) - return tuple(sorted(set(attn2_layers))) - - -def get_pos_embed_max_size(state_dict): - num_patches = state_dict["pos_embed"].shape[1] - pos_embed_max_size = int(num_patches**0.5) - return pos_embed_max_size - - -def get_caption_projection_dim(state_dict): - caption_projection_dim = state_dict["context_embedder.weight"].shape[0] - return caption_projection_dim - - -def main(args): - original_ckpt = load_original_checkpoint(args.checkpoint_path) - original_dtype = next(iter(original_ckpt.values())).dtype - - # Initialize dtype with a default value - dtype = None - - if args.dtype is None: - dtype = original_dtype - elif args.dtype == "fp16": - dtype = torch.float16 - elif args.dtype == "bf16": - dtype = torch.bfloat16 - elif args.dtype == "fp32": - dtype = torch.float32 - else: - raise ValueError(f"Unsupported dtype: {args.dtype}") - - if dtype != original_dtype: - print( - f"Checkpoint dtype {original_dtype} does not match requested dtype {dtype}. This can lead to unexpected results, proceed with caution." - ) - - 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 = get_caption_projection_dim(original_ckpt) - - # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5 - attn2_layers = get_attn2_layers(original_ckpt) - - # sd3.5 use qk norm("rms_norm") - has_qk_norm = any("ln_q" in key for key in original_ckpt.keys()) - - # sd3.5 2b use pox_embed_max_size=384 and sd3.0 and sd3.5 8b use 192 - pos_embed_max_size = get_pos_embed_max_size(original_ckpt) - - converted_transformer_state_dict = convert_sd3_transformer_checkpoint_to_diffusers( - original_ckpt, num_layers, caption_projection_dim, attn2_layers, has_qk_norm - ) - - with CTX(): - transformer = SD3Transformer2DModel( - sample_size=128, - patch_size=2, - in_channels=16, - joint_attention_dim=4096, - num_layers=num_layers, - caption_projection_dim=caption_projection_dim, - num_attention_heads=num_layers, - pos_embed_max_size=pos_embed_max_size, - qk_norm="rms_norm" if has_qk_norm else None, - dual_attention_layers=attn2_layers, - ) - 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__": - main(args) diff --git a/scripts/convert_shap_e_to_diffusers.py b/scripts/convert_shap_e_to_diffusers.py deleted file mode 100644 index ac6543667af9..000000000000 --- a/scripts/convert_shap_e_to_diffusers.py +++ /dev/null @@ -1,1080 +0,0 @@ -import argparse -import tempfile - -import torch -from accelerate import load_checkpoint_and_dispatch - -from diffusers.models.transformers.prior_transformer import PriorTransformer -from diffusers.pipelines.shap_e import ShapERenderer - - -""" -Example - From the diffusers root directory: - -Download weights: -```sh -$ wget "https://openaipublic.azureedge.net/main/shap-e/text_cond.pt" -``` - -Convert the model: -```sh -$ python scripts/convert_shap_e_to_diffusers.py \ - --prior_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/text_cond.pt \ - --prior_image_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/image_cond.pt \ - --transmitter_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/transmitter.pt\ - --dump_path /home/yiyi_huggingface_co/model_repo/shap-e-img2img/shap_e_renderer\ - --debug renderer -``` -""" - - -# prior - -PRIOR_ORIGINAL_PREFIX = "wrapped" - -PRIOR_CONFIG = { - "num_attention_heads": 16, - "attention_head_dim": 1024 // 16, - "num_layers": 24, - "embedding_dim": 1024, - "num_embeddings": 1024, - "additional_embeddings": 0, - "time_embed_act_fn": "gelu", - "norm_in_type": "layer", - "encoder_hid_proj_type": None, - "added_emb_type": None, - "time_embed_dim": 1024 * 4, - "embedding_proj_dim": 768, - "clip_embed_dim": 1024 * 2, -} - - -def prior_model_from_original_config(): - model = PriorTransformer(**PRIOR_CONFIG) - - return model - - -def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # .time_embed.c_fc -> .time_embedding.linear_1 - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.c_fc.weight"], - "time_embedding.linear_1.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.c_fc.bias"], - } - ) - - # .time_embed.c_proj -> .time_embedding.linear_2 - diffusers_checkpoint.update( - { - "time_embedding.linear_2.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.c_proj.weight"], - "time_embedding.linear_2.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.time_embed.c_proj.bias"], - } - ) - - # .input_proj -> .proj_in - diffusers_checkpoint.update( - { - "proj_in.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.input_proj.weight"], - "proj_in.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.input_proj.bias"], - } - ) - - # .clip_emb -> .embedding_proj - diffusers_checkpoint.update( - { - "embedding_proj.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_embed.weight"], - "embedding_proj.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.clip_embed.bias"], - } - ) - - # .pos_emb -> .positional_embedding - diffusers_checkpoint.update({"positional_embedding": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.pos_emb"][None, :]}) - - # .ln_pre -> .norm_in - diffusers_checkpoint.update( - { - "norm_in.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.ln_pre.weight"], - "norm_in.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.ln_pre.bias"], - } - ) - - # .backbone.resblocks. -> .transformer_blocks. - for idx in range(len(model.transformer_blocks)): - diffusers_transformer_prefix = f"transformer_blocks.{idx}" - original_transformer_prefix = f"{PRIOR_ORIGINAL_PREFIX}.backbone.resblocks.{idx}" - - # .attn -> .attn1 - diffusers_attention_prefix = f"{diffusers_transformer_prefix}.attn1" - original_attention_prefix = f"{original_transformer_prefix}.attn" - diffusers_checkpoint.update( - prior_attention_to_diffusers( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - original_attention_prefix=original_attention_prefix, - attention_head_dim=model.attention_head_dim, - ) - ) - - # .mlp -> .ff - diffusers_ff_prefix = f"{diffusers_transformer_prefix}.ff" - original_ff_prefix = f"{original_transformer_prefix}.mlp" - diffusers_checkpoint.update( - prior_ff_to_diffusers( - checkpoint, diffusers_ff_prefix=diffusers_ff_prefix, original_ff_prefix=original_ff_prefix - ) - ) - - # .ln_1 -> .norm1 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm1.weight": checkpoint[ - f"{original_transformer_prefix}.ln_1.weight" - ], - f"{diffusers_transformer_prefix}.norm1.bias": checkpoint[f"{original_transformer_prefix}.ln_1.bias"], - } - ) - - # .ln_2 -> .norm3 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm3.weight": checkpoint[ - f"{original_transformer_prefix}.ln_2.weight" - ], - f"{diffusers_transformer_prefix}.norm3.bias": checkpoint[f"{original_transformer_prefix}.ln_2.bias"], - } - ) - - # .ln_post -> .norm_out - diffusers_checkpoint.update( - { - "norm_out.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.ln_post.weight"], - "norm_out.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.ln_post.bias"], - } - ) - - # .output_proj -> .proj_to_clip_embeddings - diffusers_checkpoint.update( - { - "proj_to_clip_embeddings.weight": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.output_proj.weight"], - "proj_to_clip_embeddings.bias": checkpoint[f"{PRIOR_ORIGINAL_PREFIX}.output_proj.bias"], - } - ) - - return diffusers_checkpoint - - -def prior_attention_to_diffusers( - checkpoint, *, diffusers_attention_prefix, original_attention_prefix, attention_head_dim -): - diffusers_checkpoint = {} - - # .c_qkv -> .{to_q, to_k, to_v} - [q_weight, k_weight, v_weight], [q_bias, k_bias, v_bias] = split_attentions( - weight=checkpoint[f"{original_attention_prefix}.c_qkv.weight"], - bias=checkpoint[f"{original_attention_prefix}.c_qkv.bias"], - split=3, - chunk_size=attention_head_dim, - ) - - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_q.weight": q_weight, - f"{diffusers_attention_prefix}.to_q.bias": q_bias, - f"{diffusers_attention_prefix}.to_k.weight": k_weight, - f"{diffusers_attention_prefix}.to_k.bias": k_bias, - f"{diffusers_attention_prefix}.to_v.weight": v_weight, - f"{diffusers_attention_prefix}.to_v.bias": v_bias, - } - ) - - # .c_proj -> .to_out.0 - diffusers_checkpoint.update( - { - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{original_attention_prefix}.c_proj.weight"], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{original_attention_prefix}.c_proj.bias"], - } - ) - - return diffusers_checkpoint - - -def prior_ff_to_diffusers(checkpoint, *, diffusers_ff_prefix, original_ff_prefix): - diffusers_checkpoint = { - # .c_fc -> .net.0.proj - f"{diffusers_ff_prefix}.net.{0}.proj.weight": checkpoint[f"{original_ff_prefix}.c_fc.weight"], - f"{diffusers_ff_prefix}.net.{0}.proj.bias": checkpoint[f"{original_ff_prefix}.c_fc.bias"], - # .c_proj -> .net.2 - f"{diffusers_ff_prefix}.net.{2}.weight": checkpoint[f"{original_ff_prefix}.c_proj.weight"], - f"{diffusers_ff_prefix}.net.{2}.bias": checkpoint[f"{original_ff_prefix}.c_proj.bias"], - } - - return diffusers_checkpoint - - -# done prior - - -# prior_image (only slightly different from prior) - - -PRIOR_IMAGE_ORIGINAL_PREFIX = "wrapped" - -# Uses default arguments -PRIOR_IMAGE_CONFIG = { - "num_attention_heads": 8, - "attention_head_dim": 1024 // 8, - "num_layers": 24, - "embedding_dim": 1024, - "num_embeddings": 1024, - "additional_embeddings": 0, - "time_embed_act_fn": "gelu", - "norm_in_type": "layer", - "embedding_proj_norm_type": "layer", - "encoder_hid_proj_type": None, - "added_emb_type": None, - "time_embed_dim": 1024 * 4, - "embedding_proj_dim": 1024, - "clip_embed_dim": 1024 * 2, -} - - -def prior_image_model_from_original_config(): - model = PriorTransformer(**PRIOR_IMAGE_CONFIG) - - return model - - -def prior_image_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # .time_embed.c_fc -> .time_embedding.linear_1 - diffusers_checkpoint.update( - { - "time_embedding.linear_1.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.time_embed.c_fc.weight"], - "time_embedding.linear_1.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.time_embed.c_fc.bias"], - } - ) - - # .time_embed.c_proj -> .time_embedding.linear_2 - diffusers_checkpoint.update( - { - "time_embedding.linear_2.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.time_embed.c_proj.weight"], - "time_embedding.linear_2.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.time_embed.c_proj.bias"], - } - ) - - # .input_proj -> .proj_in - diffusers_checkpoint.update( - { - "proj_in.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.input_proj.weight"], - "proj_in.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.input_proj.bias"], - } - ) - - # .clip_embed.0 -> .embedding_proj_norm - diffusers_checkpoint.update( - { - "embedding_proj_norm.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.clip_embed.0.weight"], - "embedding_proj_norm.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.clip_embed.0.bias"], - } - ) - - # ..clip_embed.1 -> .embedding_proj - diffusers_checkpoint.update( - { - "embedding_proj.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.clip_embed.1.weight"], - "embedding_proj.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.clip_embed.1.bias"], - } - ) - - # .pos_emb -> .positional_embedding - diffusers_checkpoint.update( - {"positional_embedding": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.pos_emb"][None, :]} - ) - - # .ln_pre -> .norm_in - diffusers_checkpoint.update( - { - "norm_in.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.ln_pre.weight"], - "norm_in.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.ln_pre.bias"], - } - ) - - # .backbone.resblocks. -> .transformer_blocks. - for idx in range(len(model.transformer_blocks)): - diffusers_transformer_prefix = f"transformer_blocks.{idx}" - original_transformer_prefix = f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.backbone.resblocks.{idx}" - - # .attn -> .attn1 - diffusers_attention_prefix = f"{diffusers_transformer_prefix}.attn1" - original_attention_prefix = f"{original_transformer_prefix}.attn" - diffusers_checkpoint.update( - prior_attention_to_diffusers( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - original_attention_prefix=original_attention_prefix, - attention_head_dim=model.attention_head_dim, - ) - ) - - # .mlp -> .ff - diffusers_ff_prefix = f"{diffusers_transformer_prefix}.ff" - original_ff_prefix = f"{original_transformer_prefix}.mlp" - diffusers_checkpoint.update( - prior_ff_to_diffusers( - checkpoint, diffusers_ff_prefix=diffusers_ff_prefix, original_ff_prefix=original_ff_prefix - ) - ) - - # .ln_1 -> .norm1 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm1.weight": checkpoint[ - f"{original_transformer_prefix}.ln_1.weight" - ], - f"{diffusers_transformer_prefix}.norm1.bias": checkpoint[f"{original_transformer_prefix}.ln_1.bias"], - } - ) - - # .ln_2 -> .norm3 - diffusers_checkpoint.update( - { - f"{diffusers_transformer_prefix}.norm3.weight": checkpoint[ - f"{original_transformer_prefix}.ln_2.weight" - ], - f"{diffusers_transformer_prefix}.norm3.bias": checkpoint[f"{original_transformer_prefix}.ln_2.bias"], - } - ) - - # .ln_post -> .norm_out - diffusers_checkpoint.update( - { - "norm_out.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.ln_post.weight"], - "norm_out.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.ln_post.bias"], - } - ) - - # .output_proj -> .proj_to_clip_embeddings - diffusers_checkpoint.update( - { - "proj_to_clip_embeddings.weight": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.output_proj.weight"], - "proj_to_clip_embeddings.bias": checkpoint[f"{PRIOR_IMAGE_ORIGINAL_PREFIX}.output_proj.bias"], - } - ) - - return diffusers_checkpoint - - -# done prior_image - - -# renderer - -## create the lookup table for marching cubes method used in MeshDecoder - -MC_TABLE = [ - [], - [[0, 1, 0, 2, 0, 4]], - [[1, 0, 1, 5, 1, 3]], - [[0, 4, 1, 5, 0, 2], [1, 5, 1, 3, 0, 2]], - [[2, 0, 2, 3, 2, 6]], - [[0, 1, 2, 3, 0, 4], [2, 3, 2, 6, 0, 4]], - [[1, 0, 1, 5, 1, 3], [2, 6, 0, 2, 3, 2]], - [[3, 2, 2, 6, 3, 1], [3, 1, 2, 6, 1, 5], [1, 5, 2, 6, 0, 4]], - [[3, 1, 3, 7, 3, 2]], - [[0, 2, 0, 4, 0, 1], [3, 7, 2, 3, 1, 3]], - [[1, 5, 3, 7, 1, 0], [3, 7, 3, 2, 1, 0]], - [[2, 0, 0, 4, 2, 3], [2, 3, 0, 4, 3, 7], [3, 7, 0, 4, 1, 5]], - [[2, 0, 3, 1, 2, 6], [3, 1, 3, 7, 2, 6]], - [[1, 3, 3, 7, 1, 0], [1, 0, 3, 7, 0, 4], [0, 4, 3, 7, 2, 6]], - [[0, 1, 1, 5, 0, 2], [0, 2, 1, 5, 2, 6], [2, 6, 1, 5, 3, 7]], - [[0, 4, 1, 5, 3, 7], [0, 4, 3, 7, 2, 6]], - [[4, 0, 4, 6, 4, 5]], - [[0, 2, 4, 6, 0, 1], [4, 6, 4, 5, 0, 1]], - [[1, 5, 1, 3, 1, 0], [4, 6, 5, 4, 0, 4]], - [[5, 1, 1, 3, 5, 4], [5, 4, 1, 3, 4, 6], [4, 6, 1, 3, 0, 2]], - [[2, 0, 2, 3, 2, 6], [4, 5, 0, 4, 6, 4]], - [[6, 4, 4, 5, 6, 2], [6, 2, 4, 5, 2, 3], [2, 3, 4, 5, 0, 1]], - [[2, 6, 2, 0, 3, 2], [1, 0, 1, 5, 3, 1], [6, 4, 5, 4, 0, 4]], - [[1, 3, 5, 4, 1, 5], [1, 3, 4, 6, 5, 4], [1, 3, 3, 2, 4, 6], [3, 2, 2, 6, 4, 6]], - [[3, 1, 3, 7, 3, 2], [6, 4, 5, 4, 0, 4]], - [[4, 5, 0, 1, 4, 6], [0, 1, 0, 2, 4, 6], [7, 3, 2, 3, 1, 3]], - [[3, 2, 1, 0, 3, 7], [1, 0, 1, 5, 3, 7], [6, 4, 5, 4, 0, 4]], - [[3, 7, 3, 2, 1, 5], [3, 2, 6, 4, 1, 5], [1, 5, 6, 4, 5, 4], [3, 2, 2, 0, 6, 4]], - [[3, 7, 2, 6, 3, 1], [2, 6, 2, 0, 3, 1], [5, 4, 0, 4, 6, 4]], - [[1, 0, 1, 3, 5, 4], [1, 3, 2, 6, 5, 4], [1, 3, 3, 7, 2, 6], [5, 4, 2, 6, 4, 6]], - [[0, 1, 1, 5, 0, 2], [0, 2, 1, 5, 2, 6], [2, 6, 1, 5, 3, 7], [4, 5, 0, 4, 4, 6]], - [[6, 2, 4, 6, 4, 5], [4, 5, 5, 1, 6, 2], [6, 2, 5, 1, 7, 3]], - [[5, 1, 5, 4, 5, 7]], - [[0, 1, 0, 2, 0, 4], [5, 7, 1, 5, 4, 5]], - [[1, 0, 5, 4, 1, 3], [5, 4, 5, 7, 1, 3]], - [[4, 5, 5, 7, 4, 0], [4, 0, 5, 7, 0, 2], [0, 2, 5, 7, 1, 3]], - [[2, 0, 2, 3, 2, 6], [7, 5, 1, 5, 4, 5]], - [[2, 6, 0, 4, 2, 3], [0, 4, 0, 1, 2, 3], [7, 5, 1, 5, 4, 5]], - [[5, 7, 1, 3, 5, 4], [1, 3, 1, 0, 5, 4], [6, 2, 0, 2, 3, 2]], - [[3, 1, 3, 2, 7, 5], [3, 2, 0, 4, 7, 5], [3, 2, 2, 6, 0, 4], [7, 5, 0, 4, 5, 4]], - [[3, 7, 3, 2, 3, 1], [5, 4, 7, 5, 1, 5]], - [[0, 4, 0, 1, 2, 0], [3, 1, 3, 7, 2, 3], [4, 5, 7, 5, 1, 5]], - [[7, 3, 3, 2, 7, 5], [7, 5, 3, 2, 5, 4], [5, 4, 3, 2, 1, 0]], - [[0, 4, 2, 3, 0, 2], [0, 4, 3, 7, 2, 3], [0, 4, 4, 5, 3, 7], [4, 5, 5, 7, 3, 7]], - [[2, 0, 3, 1, 2, 6], [3, 1, 3, 7, 2, 6], [4, 5, 7, 5, 1, 5]], - [[1, 3, 3, 7, 1, 0], [1, 0, 3, 7, 0, 4], [0, 4, 3, 7, 2, 6], [5, 7, 1, 5, 5, 4]], - [[2, 6, 2, 0, 3, 7], [2, 0, 4, 5, 3, 7], [3, 7, 4, 5, 7, 5], [2, 0, 0, 1, 4, 5]], - [[4, 0, 5, 4, 5, 7], [5, 7, 7, 3, 4, 0], [4, 0, 7, 3, 6, 2]], - [[4, 6, 5, 7, 4, 0], [5, 7, 5, 1, 4, 0]], - [[1, 0, 0, 2, 1, 5], [1, 5, 0, 2, 5, 7], [5, 7, 0, 2, 4, 6]], - [[0, 4, 4, 6, 0, 1], [0, 1, 4, 6, 1, 3], [1, 3, 4, 6, 5, 7]], - [[0, 2, 4, 6, 5, 7], [0, 2, 5, 7, 1, 3]], - [[5, 1, 4, 0, 5, 7], [4, 0, 4, 6, 5, 7], [3, 2, 6, 2, 0, 2]], - [[2, 3, 2, 6, 0, 1], [2, 6, 7, 5, 0, 1], [0, 1, 7, 5, 1, 5], [2, 6, 6, 4, 7, 5]], - [[0, 4, 4, 6, 0, 1], [0, 1, 4, 6, 1, 3], [1, 3, 4, 6, 5, 7], [2, 6, 0, 2, 2, 3]], - [[3, 1, 2, 3, 2, 6], [2, 6, 6, 4, 3, 1], [3, 1, 6, 4, 7, 5]], - [[4, 6, 5, 7, 4, 0], [5, 7, 5, 1, 4, 0], [2, 3, 1, 3, 7, 3]], - [[1, 0, 0, 2, 1, 5], [1, 5, 0, 2, 5, 7], [5, 7, 0, 2, 4, 6], [3, 2, 1, 3, 3, 7]], - [[0, 1, 0, 4, 2, 3], [0, 4, 5, 7, 2, 3], [0, 4, 4, 6, 5, 7], [2, 3, 5, 7, 3, 7]], - [[7, 5, 3, 7, 3, 2], [3, 2, 2, 0, 7, 5], [7, 5, 2, 0, 6, 4]], - [[0, 4, 4, 6, 5, 7], [0, 4, 5, 7, 1, 5], [0, 2, 1, 3, 3, 7], [3, 7, 2, 6, 0, 2]], - [ - [3, 1, 7, 3, 6, 2], - [6, 2, 0, 1, 3, 1], - [6, 4, 0, 1, 6, 2], - [6, 4, 5, 1, 0, 1], - [6, 4, 7, 5, 5, 1], - ], - [ - [4, 0, 6, 4, 7, 5], - [7, 5, 1, 0, 4, 0], - [7, 3, 1, 0, 7, 5], - [7, 3, 2, 0, 1, 0], - [7, 3, 6, 2, 2, 0], - ], - [[7, 3, 6, 2, 6, 4], [7, 5, 7, 3, 6, 4]], - [[6, 2, 6, 7, 6, 4]], - [[0, 4, 0, 1, 0, 2], [6, 7, 4, 6, 2, 6]], - [[1, 0, 1, 5, 1, 3], [7, 6, 4, 6, 2, 6]], - [[1, 3, 0, 2, 1, 5], [0, 2, 0, 4, 1, 5], [7, 6, 4, 6, 2, 6]], - [[2, 3, 6, 7, 2, 0], [6, 7, 6, 4, 2, 0]], - [[4, 0, 0, 1, 4, 6], [4, 6, 0, 1, 6, 7], [6, 7, 0, 1, 2, 3]], - [[6, 4, 2, 0, 6, 7], [2, 0, 2, 3, 6, 7], [5, 1, 3, 1, 0, 1]], - [[1, 5, 1, 3, 0, 4], [1, 3, 7, 6, 0, 4], [0, 4, 7, 6, 4, 6], [1, 3, 3, 2, 7, 6]], - [[3, 2, 3, 1, 3, 7], [6, 4, 2, 6, 7, 6]], - [[3, 7, 3, 2, 1, 3], [0, 2, 0, 4, 1, 0], [7, 6, 4, 6, 2, 6]], - [[1, 5, 3, 7, 1, 0], [3, 7, 3, 2, 1, 0], [4, 6, 2, 6, 7, 6]], - [[2, 0, 0, 4, 2, 3], [2, 3, 0, 4, 3, 7], [3, 7, 0, 4, 1, 5], [6, 4, 2, 6, 6, 7]], - [[7, 6, 6, 4, 7, 3], [7, 3, 6, 4, 3, 1], [3, 1, 6, 4, 2, 0]], - [[0, 1, 4, 6, 0, 4], [0, 1, 6, 7, 4, 6], [0, 1, 1, 3, 6, 7], [1, 3, 3, 7, 6, 7]], - [[0, 2, 0, 1, 4, 6], [0, 1, 3, 7, 4, 6], [0, 1, 1, 5, 3, 7], [4, 6, 3, 7, 6, 7]], - [[7, 3, 6, 7, 6, 4], [6, 4, 4, 0, 7, 3], [7, 3, 4, 0, 5, 1]], - [[4, 0, 6, 2, 4, 5], [6, 2, 6, 7, 4, 5]], - [[2, 6, 6, 7, 2, 0], [2, 0, 6, 7, 0, 1], [0, 1, 6, 7, 4, 5]], - [[6, 7, 4, 5, 6, 2], [4, 5, 4, 0, 6, 2], [3, 1, 0, 1, 5, 1]], - [[2, 0, 2, 6, 3, 1], [2, 6, 4, 5, 3, 1], [2, 6, 6, 7, 4, 5], [3, 1, 4, 5, 1, 5]], - [[0, 2, 2, 3, 0, 4], [0, 4, 2, 3, 4, 5], [4, 5, 2, 3, 6, 7]], - [[0, 1, 2, 3, 6, 7], [0, 1, 6, 7, 4, 5]], - [[0, 2, 2, 3, 0, 4], [0, 4, 2, 3, 4, 5], [4, 5, 2, 3, 6, 7], [1, 3, 0, 1, 1, 5]], - [[5, 4, 1, 5, 1, 3], [1, 3, 3, 2, 5, 4], [5, 4, 3, 2, 7, 6]], - [[4, 0, 6, 2, 4, 5], [6, 2, 6, 7, 4, 5], [1, 3, 7, 3, 2, 3]], - [[2, 6, 6, 7, 2, 0], [2, 0, 6, 7, 0, 1], [0, 1, 6, 7, 4, 5], [3, 7, 2, 3, 3, 1]], - [[0, 1, 1, 5, 3, 7], [0, 1, 3, 7, 2, 3], [0, 4, 2, 6, 6, 7], [6, 7, 4, 5, 0, 4]], - [ - [6, 2, 7, 6, 5, 4], - [5, 4, 0, 2, 6, 2], - [5, 1, 0, 2, 5, 4], - [5, 1, 3, 2, 0, 2], - [5, 1, 7, 3, 3, 2], - ], - [[3, 1, 3, 7, 2, 0], [3, 7, 5, 4, 2, 0], [2, 0, 5, 4, 0, 4], [3, 7, 7, 6, 5, 4]], - [[1, 0, 3, 1, 3, 7], [3, 7, 7, 6, 1, 0], [1, 0, 7, 6, 5, 4]], - [ - [1, 0, 5, 1, 7, 3], - [7, 3, 2, 0, 1, 0], - [7, 6, 2, 0, 7, 3], - [7, 6, 4, 0, 2, 0], - [7, 6, 5, 4, 4, 0], - ], - [[7, 6, 5, 4, 5, 1], [7, 3, 7, 6, 5, 1]], - [[5, 7, 5, 1, 5, 4], [6, 2, 7, 6, 4, 6]], - [[0, 2, 0, 4, 1, 0], [5, 4, 5, 7, 1, 5], [2, 6, 7, 6, 4, 6]], - [[1, 0, 5, 4, 1, 3], [5, 4, 5, 7, 1, 3], [2, 6, 7, 6, 4, 6]], - [[4, 5, 5, 7, 4, 0], [4, 0, 5, 7, 0, 2], [0, 2, 5, 7, 1, 3], [6, 7, 4, 6, 6, 2]], - [[2, 3, 6, 7, 2, 0], [6, 7, 6, 4, 2, 0], [1, 5, 4, 5, 7, 5]], - [[4, 0, 0, 1, 4, 6], [4, 6, 0, 1, 6, 7], [6, 7, 0, 1, 2, 3], [5, 1, 4, 5, 5, 7]], - [[0, 2, 2, 3, 6, 7], [0, 2, 6, 7, 4, 6], [0, 1, 4, 5, 5, 7], [5, 7, 1, 3, 0, 1]], - [ - [5, 4, 7, 5, 3, 1], - [3, 1, 0, 4, 5, 4], - [3, 2, 0, 4, 3, 1], - [3, 2, 6, 4, 0, 4], - [3, 2, 7, 6, 6, 4], - ], - [[5, 4, 5, 7, 1, 5], [3, 7, 3, 2, 1, 3], [4, 6, 2, 6, 7, 6]], - [[1, 0, 0, 2, 0, 4], [1, 5, 5, 4, 5, 7], [3, 2, 1, 3, 3, 7], [2, 6, 7, 6, 4, 6]], - [[7, 3, 3, 2, 7, 5], [7, 5, 3, 2, 5, 4], [5, 4, 3, 2, 1, 0], [6, 2, 7, 6, 6, 4]], - [ - [0, 4, 2, 3, 0, 2], - [0, 4, 3, 7, 2, 3], - [0, 4, 4, 5, 3, 7], - [4, 5, 5, 7, 3, 7], - [6, 7, 4, 6, 2, 6], - ], - [[7, 6, 6, 4, 7, 3], [7, 3, 6, 4, 3, 1], [3, 1, 6, 4, 2, 0], [5, 4, 7, 5, 5, 1]], - [ - [0, 1, 4, 6, 0, 4], - [0, 1, 6, 7, 4, 6], - [0, 1, 1, 3, 6, 7], - [1, 3, 3, 7, 6, 7], - [5, 7, 1, 5, 4, 5], - ], - [ - [6, 7, 4, 6, 0, 2], - [0, 2, 3, 7, 6, 7], - [0, 1, 3, 7, 0, 2], - [0, 1, 5, 7, 3, 7], - [0, 1, 4, 5, 5, 7], - ], - [[4, 0, 6, 7, 4, 6], [4, 0, 7, 3, 6, 7], [4, 0, 5, 7, 7, 3], [4, 5, 5, 7, 4, 0]], - [[7, 5, 5, 1, 7, 6], [7, 6, 5, 1, 6, 2], [6, 2, 5, 1, 4, 0]], - [[0, 2, 1, 5, 0, 1], [0, 2, 5, 7, 1, 5], [0, 2, 2, 6, 5, 7], [2, 6, 6, 7, 5, 7]], - [[1, 3, 1, 0, 5, 7], [1, 0, 2, 6, 5, 7], [5, 7, 2, 6, 7, 6], [1, 0, 0, 4, 2, 6]], - [[2, 0, 6, 2, 6, 7], [6, 7, 7, 5, 2, 0], [2, 0, 7, 5, 3, 1]], - [[0, 4, 0, 2, 1, 5], [0, 2, 6, 7, 1, 5], [0, 2, 2, 3, 6, 7], [1, 5, 6, 7, 5, 7]], - [[7, 6, 5, 7, 5, 1], [5, 1, 1, 0, 7, 6], [7, 6, 1, 0, 3, 2]], - [ - [2, 0, 3, 2, 7, 6], - [7, 6, 4, 0, 2, 0], - [7, 5, 4, 0, 7, 6], - [7, 5, 1, 0, 4, 0], - [7, 5, 3, 1, 1, 0], - ], - [[7, 5, 3, 1, 3, 2], [7, 6, 7, 5, 3, 2]], - [[7, 5, 5, 1, 7, 6], [7, 6, 5, 1, 6, 2], [6, 2, 5, 1, 4, 0], [3, 1, 7, 3, 3, 2]], - [ - [0, 2, 1, 5, 0, 1], - [0, 2, 5, 7, 1, 5], - [0, 2, 2, 6, 5, 7], - [2, 6, 6, 7, 5, 7], - [3, 7, 2, 3, 1, 3], - ], - [ - [3, 7, 2, 3, 0, 1], - [0, 1, 5, 7, 3, 7], - [0, 4, 5, 7, 0, 1], - [0, 4, 6, 7, 5, 7], - [0, 4, 2, 6, 6, 7], - ], - [[2, 0, 3, 7, 2, 3], [2, 0, 7, 5, 3, 7], [2, 0, 6, 7, 7, 5], [2, 6, 6, 7, 2, 0]], - [ - [5, 7, 1, 5, 0, 4], - [0, 4, 6, 7, 5, 7], - [0, 2, 6, 7, 0, 4], - [0, 2, 3, 7, 6, 7], - [0, 2, 1, 3, 3, 7], - ], - [[1, 0, 5, 7, 1, 5], [1, 0, 7, 6, 5, 7], [1, 0, 3, 7, 7, 6], [1, 3, 3, 7, 1, 0]], - [[0, 2, 0, 1, 0, 4], [3, 7, 6, 7, 5, 7]], - [[7, 5, 7, 3, 7, 6]], - [[7, 3, 7, 5, 7, 6]], - [[0, 1, 0, 2, 0, 4], [6, 7, 3, 7, 5, 7]], - [[1, 3, 1, 0, 1, 5], [7, 6, 3, 7, 5, 7]], - [[0, 4, 1, 5, 0, 2], [1, 5, 1, 3, 0, 2], [6, 7, 3, 7, 5, 7]], - [[2, 6, 2, 0, 2, 3], [7, 5, 6, 7, 3, 7]], - [[0, 1, 2, 3, 0, 4], [2, 3, 2, 6, 0, 4], [5, 7, 6, 7, 3, 7]], - [[1, 5, 1, 3, 0, 1], [2, 3, 2, 6, 0, 2], [5, 7, 6, 7, 3, 7]], - [[3, 2, 2, 6, 3, 1], [3, 1, 2, 6, 1, 5], [1, 5, 2, 6, 0, 4], [7, 6, 3, 7, 7, 5]], - [[3, 1, 7, 5, 3, 2], [7, 5, 7, 6, 3, 2]], - [[7, 6, 3, 2, 7, 5], [3, 2, 3, 1, 7, 5], [4, 0, 1, 0, 2, 0]], - [[5, 7, 7, 6, 5, 1], [5, 1, 7, 6, 1, 0], [1, 0, 7, 6, 3, 2]], - [[2, 3, 2, 0, 6, 7], [2, 0, 1, 5, 6, 7], [2, 0, 0, 4, 1, 5], [6, 7, 1, 5, 7, 5]], - [[6, 2, 2, 0, 6, 7], [6, 7, 2, 0, 7, 5], [7, 5, 2, 0, 3, 1]], - [[0, 4, 0, 1, 2, 6], [0, 1, 5, 7, 2, 6], [2, 6, 5, 7, 6, 7], [0, 1, 1, 3, 5, 7]], - [[1, 5, 0, 2, 1, 0], [1, 5, 2, 6, 0, 2], [1, 5, 5, 7, 2, 6], [5, 7, 7, 6, 2, 6]], - [[5, 1, 7, 5, 7, 6], [7, 6, 6, 2, 5, 1], [5, 1, 6, 2, 4, 0]], - [[4, 5, 4, 0, 4, 6], [7, 3, 5, 7, 6, 7]], - [[0, 2, 4, 6, 0, 1], [4, 6, 4, 5, 0, 1], [3, 7, 5, 7, 6, 7]], - [[4, 6, 4, 5, 0, 4], [1, 5, 1, 3, 0, 1], [6, 7, 3, 7, 5, 7]], - [[5, 1, 1, 3, 5, 4], [5, 4, 1, 3, 4, 6], [4, 6, 1, 3, 0, 2], [7, 3, 5, 7, 7, 6]], - [[2, 3, 2, 6, 0, 2], [4, 6, 4, 5, 0, 4], [3, 7, 5, 7, 6, 7]], - [[6, 4, 4, 5, 6, 2], [6, 2, 4, 5, 2, 3], [2, 3, 4, 5, 0, 1], [7, 5, 6, 7, 7, 3]], - [[0, 1, 1, 5, 1, 3], [0, 2, 2, 3, 2, 6], [4, 5, 0, 4, 4, 6], [5, 7, 6, 7, 3, 7]], - [ - [1, 3, 5, 4, 1, 5], - [1, 3, 4, 6, 5, 4], - [1, 3, 3, 2, 4, 6], - [3, 2, 2, 6, 4, 6], - [7, 6, 3, 7, 5, 7], - ], - [[3, 1, 7, 5, 3, 2], [7, 5, 7, 6, 3, 2], [0, 4, 6, 4, 5, 4]], - [[1, 0, 0, 2, 4, 6], [1, 0, 4, 6, 5, 4], [1, 3, 5, 7, 7, 6], [7, 6, 3, 2, 1, 3]], - [[5, 7, 7, 6, 5, 1], [5, 1, 7, 6, 1, 0], [1, 0, 7, 6, 3, 2], [4, 6, 5, 4, 4, 0]], - [ - [7, 5, 6, 7, 2, 3], - [2, 3, 1, 5, 7, 5], - [2, 0, 1, 5, 2, 3], - [2, 0, 4, 5, 1, 5], - [2, 0, 6, 4, 4, 5], - ], - [[6, 2, 2, 0, 6, 7], [6, 7, 2, 0, 7, 5], [7, 5, 2, 0, 3, 1], [4, 0, 6, 4, 4, 5]], - [ - [4, 6, 5, 4, 1, 0], - [1, 0, 2, 6, 4, 6], - [1, 3, 2, 6, 1, 0], - [1, 3, 7, 6, 2, 6], - [1, 3, 5, 7, 7, 6], - ], - [ - [1, 5, 0, 2, 1, 0], - [1, 5, 2, 6, 0, 2], - [1, 5, 5, 7, 2, 6], - [5, 7, 7, 6, 2, 6], - [4, 6, 5, 4, 0, 4], - ], - [[5, 1, 4, 6, 5, 4], [5, 1, 6, 2, 4, 6], [5, 1, 7, 6, 6, 2], [5, 7, 7, 6, 5, 1]], - [[5, 4, 7, 6, 5, 1], [7, 6, 7, 3, 5, 1]], - [[7, 3, 5, 1, 7, 6], [5, 1, 5, 4, 7, 6], [2, 0, 4, 0, 1, 0]], - [[3, 1, 1, 0, 3, 7], [3, 7, 1, 0, 7, 6], [7, 6, 1, 0, 5, 4]], - [[0, 2, 0, 4, 1, 3], [0, 4, 6, 7, 1, 3], [1, 3, 6, 7, 3, 7], [0, 4, 4, 5, 6, 7]], - [[5, 4, 7, 6, 5, 1], [7, 6, 7, 3, 5, 1], [0, 2, 3, 2, 6, 2]], - [[1, 5, 5, 4, 7, 6], [1, 5, 7, 6, 3, 7], [1, 0, 3, 2, 2, 6], [2, 6, 0, 4, 1, 0]], - [[3, 1, 1, 0, 3, 7], [3, 7, 1, 0, 7, 6], [7, 6, 1, 0, 5, 4], [2, 0, 3, 2, 2, 6]], - [ - [2, 3, 6, 2, 4, 0], - [4, 0, 1, 3, 2, 3], - [4, 5, 1, 3, 4, 0], - [4, 5, 7, 3, 1, 3], - [4, 5, 6, 7, 7, 3], - ], - [[1, 5, 5, 4, 1, 3], [1, 3, 5, 4, 3, 2], [3, 2, 5, 4, 7, 6]], - [[1, 5, 5, 4, 1, 3], [1, 3, 5, 4, 3, 2], [3, 2, 5, 4, 7, 6], [0, 4, 1, 0, 0, 2]], - [[1, 0, 5, 4, 7, 6], [1, 0, 7, 6, 3, 2]], - [[2, 3, 0, 2, 0, 4], [0, 4, 4, 5, 2, 3], [2, 3, 4, 5, 6, 7]], - [[1, 3, 1, 5, 0, 2], [1, 5, 7, 6, 0, 2], [1, 5, 5, 4, 7, 6], [0, 2, 7, 6, 2, 6]], - [ - [5, 1, 4, 5, 6, 7], - [6, 7, 3, 1, 5, 1], - [6, 2, 3, 1, 6, 7], - [6, 2, 0, 1, 3, 1], - [6, 2, 4, 0, 0, 1], - ], - [[6, 7, 2, 6, 2, 0], [2, 0, 0, 1, 6, 7], [6, 7, 0, 1, 4, 5]], - [[6, 2, 4, 0, 4, 5], [6, 7, 6, 2, 4, 5]], - [[6, 7, 7, 3, 6, 4], [6, 4, 7, 3, 4, 0], [4, 0, 7, 3, 5, 1]], - [[1, 5, 1, 0, 3, 7], [1, 0, 4, 6, 3, 7], [1, 0, 0, 2, 4, 6], [3, 7, 4, 6, 7, 6]], - [[1, 0, 3, 7, 1, 3], [1, 0, 7, 6, 3, 7], [1, 0, 0, 4, 7, 6], [0, 4, 4, 6, 7, 6]], - [[6, 4, 7, 6, 7, 3], [7, 3, 3, 1, 6, 4], [6, 4, 3, 1, 2, 0]], - [[6, 7, 7, 3, 6, 4], [6, 4, 7, 3, 4, 0], [4, 0, 7, 3, 5, 1], [2, 3, 6, 2, 2, 0]], - [ - [7, 6, 3, 7, 1, 5], - [1, 5, 4, 6, 7, 6], - [1, 0, 4, 6, 1, 5], - [1, 0, 2, 6, 4, 6], - [1, 0, 3, 2, 2, 6], - ], - [ - [1, 0, 3, 7, 1, 3], - [1, 0, 7, 6, 3, 7], - [1, 0, 0, 4, 7, 6], - [0, 4, 4, 6, 7, 6], - [2, 6, 0, 2, 3, 2], - ], - [[3, 1, 7, 6, 3, 7], [3, 1, 6, 4, 7, 6], [3, 1, 2, 6, 6, 4], [3, 2, 2, 6, 3, 1]], - [[3, 2, 3, 1, 7, 6], [3, 1, 0, 4, 7, 6], [7, 6, 0, 4, 6, 4], [3, 1, 1, 5, 0, 4]], - [ - [0, 1, 2, 0, 6, 4], - [6, 4, 5, 1, 0, 1], - [6, 7, 5, 1, 6, 4], - [6, 7, 3, 1, 5, 1], - [6, 7, 2, 3, 3, 1], - ], - [[0, 1, 4, 0, 4, 6], [4, 6, 6, 7, 0, 1], [0, 1, 6, 7, 2, 3]], - [[6, 7, 2, 3, 2, 0], [6, 4, 6, 7, 2, 0]], - [ - [2, 6, 0, 2, 1, 3], - [1, 3, 7, 6, 2, 6], - [1, 5, 7, 6, 1, 3], - [1, 5, 4, 6, 7, 6], - [1, 5, 0, 4, 4, 6], - ], - [[1, 5, 1, 0, 1, 3], [4, 6, 7, 6, 2, 6]], - [[0, 1, 2, 6, 0, 2], [0, 1, 6, 7, 2, 6], [0, 1, 4, 6, 6, 7], [0, 4, 4, 6, 0, 1]], - [[6, 7, 6, 2, 6, 4]], - [[6, 2, 7, 3, 6, 4], [7, 3, 7, 5, 6, 4]], - [[7, 5, 6, 4, 7, 3], [6, 4, 6, 2, 7, 3], [1, 0, 2, 0, 4, 0]], - [[6, 2, 7, 3, 6, 4], [7, 3, 7, 5, 6, 4], [0, 1, 5, 1, 3, 1]], - [[2, 0, 0, 4, 1, 5], [2, 0, 1, 5, 3, 1], [2, 6, 3, 7, 7, 5], [7, 5, 6, 4, 2, 6]], - [[3, 7, 7, 5, 3, 2], [3, 2, 7, 5, 2, 0], [2, 0, 7, 5, 6, 4]], - [[3, 2, 3, 7, 1, 0], [3, 7, 6, 4, 1, 0], [3, 7, 7, 5, 6, 4], [1, 0, 6, 4, 0, 4]], - [[3, 7, 7, 5, 3, 2], [3, 2, 7, 5, 2, 0], [2, 0, 7, 5, 6, 4], [1, 5, 3, 1, 1, 0]], - [ - [7, 3, 5, 7, 4, 6], - [4, 6, 2, 3, 7, 3], - [4, 0, 2, 3, 4, 6], - [4, 0, 1, 3, 2, 3], - [4, 0, 5, 1, 1, 3], - ], - [[2, 3, 3, 1, 2, 6], [2, 6, 3, 1, 6, 4], [6, 4, 3, 1, 7, 5]], - [[2, 3, 3, 1, 2, 6], [2, 6, 3, 1, 6, 4], [6, 4, 3, 1, 7, 5], [0, 1, 2, 0, 0, 4]], - [[1, 0, 1, 5, 3, 2], [1, 5, 4, 6, 3, 2], [3, 2, 4, 6, 2, 6], [1, 5, 5, 7, 4, 6]], - [ - [0, 2, 4, 0, 5, 1], - [5, 1, 3, 2, 0, 2], - [5, 7, 3, 2, 5, 1], - [5, 7, 6, 2, 3, 2], - [5, 7, 4, 6, 6, 2], - ], - [[2, 0, 3, 1, 7, 5], [2, 0, 7, 5, 6, 4]], - [[4, 6, 0, 4, 0, 1], [0, 1, 1, 3, 4, 6], [4, 6, 1, 3, 5, 7]], - [[0, 2, 1, 0, 1, 5], [1, 5, 5, 7, 0, 2], [0, 2, 5, 7, 4, 6]], - [[5, 7, 4, 6, 4, 0], [5, 1, 5, 7, 4, 0]], - [[5, 4, 4, 0, 5, 7], [5, 7, 4, 0, 7, 3], [7, 3, 4, 0, 6, 2]], - [[0, 1, 0, 2, 4, 5], [0, 2, 3, 7, 4, 5], [4, 5, 3, 7, 5, 7], [0, 2, 2, 6, 3, 7]], - [[5, 4, 4, 0, 5, 7], [5, 7, 4, 0, 7, 3], [7, 3, 4, 0, 6, 2], [1, 0, 5, 1, 1, 3]], - [ - [1, 5, 3, 1, 2, 0], - [2, 0, 4, 5, 1, 5], - [2, 6, 4, 5, 2, 0], - [2, 6, 7, 5, 4, 5], - [2, 6, 3, 7, 7, 5], - ], - [[2, 3, 0, 4, 2, 0], [2, 3, 4, 5, 0, 4], [2, 3, 3, 7, 4, 5], [3, 7, 7, 5, 4, 5]], - [[3, 2, 7, 3, 7, 5], [7, 5, 5, 4, 3, 2], [3, 2, 5, 4, 1, 0]], - [ - [2, 3, 0, 4, 2, 0], - [2, 3, 4, 5, 0, 4], - [2, 3, 3, 7, 4, 5], - [3, 7, 7, 5, 4, 5], - [1, 5, 3, 1, 0, 1], - ], - [[3, 2, 1, 5, 3, 1], [3, 2, 5, 4, 1, 5], [3, 2, 7, 5, 5, 4], [3, 7, 7, 5, 3, 2]], - [[2, 6, 2, 3, 0, 4], [2, 3, 7, 5, 0, 4], [2, 3, 3, 1, 7, 5], [0, 4, 7, 5, 4, 5]], - [ - [3, 2, 1, 3, 5, 7], - [5, 7, 6, 2, 3, 2], - [5, 4, 6, 2, 5, 7], - [5, 4, 0, 2, 6, 2], - [5, 4, 1, 0, 0, 2], - ], - [ - [4, 5, 0, 4, 2, 6], - [2, 6, 7, 5, 4, 5], - [2, 3, 7, 5, 2, 6], - [2, 3, 1, 5, 7, 5], - [2, 3, 0, 1, 1, 5], - ], - [[2, 3, 2, 0, 2, 6], [1, 5, 7, 5, 4, 5]], - [[5, 7, 4, 5, 4, 0], [4, 0, 0, 2, 5, 7], [5, 7, 0, 2, 1, 3]], - [[5, 4, 1, 0, 1, 3], [5, 7, 5, 4, 1, 3]], - [[0, 2, 4, 5, 0, 4], [0, 2, 5, 7, 4, 5], [0, 2, 1, 5, 5, 7], [0, 1, 1, 5, 0, 2]], - [[5, 4, 5, 1, 5, 7]], - [[4, 6, 6, 2, 4, 5], [4, 5, 6, 2, 5, 1], [5, 1, 6, 2, 7, 3]], - [[4, 6, 6, 2, 4, 5], [4, 5, 6, 2, 5, 1], [5, 1, 6, 2, 7, 3], [0, 2, 4, 0, 0, 1]], - [[3, 7, 3, 1, 2, 6], [3, 1, 5, 4, 2, 6], [3, 1, 1, 0, 5, 4], [2, 6, 5, 4, 6, 4]], - [ - [6, 4, 2, 6, 3, 7], - [3, 7, 5, 4, 6, 4], - [3, 1, 5, 4, 3, 7], - [3, 1, 0, 4, 5, 4], - [3, 1, 2, 0, 0, 4], - ], - [[2, 0, 2, 3, 6, 4], [2, 3, 1, 5, 6, 4], [6, 4, 1, 5, 4, 5], [2, 3, 3, 7, 1, 5]], - [ - [0, 4, 1, 0, 3, 2], - [3, 2, 6, 4, 0, 4], - [3, 7, 6, 4, 3, 2], - [3, 7, 5, 4, 6, 4], - [3, 7, 1, 5, 5, 4], - ], - [ - [1, 3, 0, 1, 4, 5], - [4, 5, 7, 3, 1, 3], - [4, 6, 7, 3, 4, 5], - [4, 6, 2, 3, 7, 3], - [4, 6, 0, 2, 2, 3], - ], - [[3, 7, 3, 1, 3, 2], [5, 4, 6, 4, 0, 4]], - [[3, 1, 2, 6, 3, 2], [3, 1, 6, 4, 2, 6], [3, 1, 1, 5, 6, 4], [1, 5, 5, 4, 6, 4]], - [ - [3, 1, 2, 6, 3, 2], - [3, 1, 6, 4, 2, 6], - [3, 1, 1, 5, 6, 4], - [1, 5, 5, 4, 6, 4], - [0, 4, 1, 0, 2, 0], - ], - [[4, 5, 6, 4, 6, 2], [6, 2, 2, 3, 4, 5], [4, 5, 2, 3, 0, 1]], - [[2, 3, 6, 4, 2, 6], [2, 3, 4, 5, 6, 4], [2, 3, 0, 4, 4, 5], [2, 0, 0, 4, 2, 3]], - [[1, 3, 5, 1, 5, 4], [5, 4, 4, 6, 1, 3], [1, 3, 4, 6, 0, 2]], - [[1, 3, 0, 4, 1, 0], [1, 3, 4, 6, 0, 4], [1, 3, 5, 4, 4, 6], [1, 5, 5, 4, 1, 3]], - [[4, 6, 0, 2, 0, 1], [4, 5, 4, 6, 0, 1]], - [[4, 6, 4, 0, 4, 5]], - [[4, 0, 6, 2, 7, 3], [4, 0, 7, 3, 5, 1]], - [[1, 5, 0, 1, 0, 2], [0, 2, 2, 6, 1, 5], [1, 5, 2, 6, 3, 7]], - [[3, 7, 1, 3, 1, 0], [1, 0, 0, 4, 3, 7], [3, 7, 0, 4, 2, 6]], - [[3, 1, 2, 0, 2, 6], [3, 7, 3, 1, 2, 6]], - [[0, 4, 2, 0, 2, 3], [2, 3, 3, 7, 0, 4], [0, 4, 3, 7, 1, 5]], - [[3, 7, 1, 5, 1, 0], [3, 2, 3, 7, 1, 0]], - [[0, 4, 1, 3, 0, 1], [0, 4, 3, 7, 1, 3], [0, 4, 2, 3, 3, 7], [0, 2, 2, 3, 0, 4]], - [[3, 7, 3, 1, 3, 2]], - [[2, 6, 3, 2, 3, 1], [3, 1, 1, 5, 2, 6], [2, 6, 1, 5, 0, 4]], - [[1, 5, 3, 2, 1, 3], [1, 5, 2, 6, 3, 2], [1, 5, 0, 2, 2, 6], [1, 0, 0, 2, 1, 5]], - [[2, 3, 0, 1, 0, 4], [2, 6, 2, 3, 0, 4]], - [[2, 3, 2, 0, 2, 6]], - [[1, 5, 0, 4, 0, 2], [1, 3, 1, 5, 0, 2]], - [[1, 5, 1, 0, 1, 3]], - [[0, 2, 0, 1, 0, 4]], - [], -] - - -def create_mc_lookup_table(): - cases = torch.zeros(256, 5, 3, dtype=torch.long) - masks = torch.zeros(256, 5, dtype=torch.bool) - - edge_to_index = { - (0, 1): 0, - (2, 3): 1, - (4, 5): 2, - (6, 7): 3, - (0, 2): 4, - (1, 3): 5, - (4, 6): 6, - (5, 7): 7, - (0, 4): 8, - (1, 5): 9, - (2, 6): 10, - (3, 7): 11, - } - - for i, case in enumerate(MC_TABLE): - for j, tri in enumerate(case): - for k, (c1, c2) in enumerate(zip(tri[::2], tri[1::2])): - cases[i, j, k] = edge_to_index[(c1, c2) if c1 < c2 else (c2, c1)] - masks[i, j] = True - return cases, masks - - -RENDERER_CONFIG = {} - - -def renderer_model_from_original_config(): - model = ShapERenderer(**RENDERER_CONFIG) - - return model - - -RENDERER_MLP_ORIGINAL_PREFIX = "renderer.nerstf" - -RENDERER_PARAMS_PROJ_ORIGINAL_PREFIX = "encoder.params_proj" - - -def renderer_model_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - diffusers_checkpoint.update( - {f"mlp.{k}": checkpoint[f"{RENDERER_MLP_ORIGINAL_PREFIX}.{k}"] for k in model.mlp.state_dict().keys()} - ) - - diffusers_checkpoint.update( - { - f"params_proj.{k}": checkpoint[f"{RENDERER_PARAMS_PROJ_ORIGINAL_PREFIX}.{k}"] - for k in model.params_proj.state_dict().keys() - } - ) - - diffusers_checkpoint.update({"void.background": model.state_dict()["void.background"]}) - - cases, masks = create_mc_lookup_table() - - diffusers_checkpoint.update({"mesh_decoder.cases": cases}) - diffusers_checkpoint.update({"mesh_decoder.masks": masks}) - - return diffusers_checkpoint - - -# done renderer - - -# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) -def split_attentions(*, weight, bias, split, chunk_size): - weights = [None] * split - biases = [None] * split - - weights_biases_idx = 0 - - for starting_row_index in range(0, weight.shape[0], chunk_size): - row_indices = torch.arange(starting_row_index, starting_row_index + chunk_size) - - weight_rows = weight[row_indices, :] - bias_rows = bias[row_indices] - - if weights[weights_biases_idx] is None: - assert weights[weights_biases_idx] is None - weights[weights_biases_idx] = weight_rows - biases[weights_biases_idx] = bias_rows - else: - assert weights[weights_biases_idx] is not None - weights[weights_biases_idx] = torch.concat([weights[weights_biases_idx], weight_rows]) - biases[weights_biases_idx] = torch.concat([biases[weights_biases_idx], bias_rows]) - - weights_biases_idx = (weights_biases_idx + 1) % split - - return weights, biases - - -# done unet utils - - -# Driver functions - - -def prior(*, args, checkpoint_map_location): - print("loading prior") - - prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) - - prior_model = prior_model_from_original_config() - - prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint(prior_model, prior_checkpoint) - - del prior_checkpoint - - load_prior_checkpoint_to_model(prior_diffusers_checkpoint, prior_model) - - print("done loading prior") - - return prior_model - - -def prior_image(*, args, checkpoint_map_location): - print("loading prior_image") - - print(f"load checkpoint from {args.prior_image_checkpoint_path}") - prior_checkpoint = torch.load(args.prior_image_checkpoint_path, map_location=checkpoint_map_location) - - prior_model = prior_image_model_from_original_config() - - prior_diffusers_checkpoint = prior_image_original_checkpoint_to_diffusers_checkpoint(prior_model, prior_checkpoint) - - del prior_checkpoint - - load_prior_checkpoint_to_model(prior_diffusers_checkpoint, prior_model) - - print("done loading prior_image") - - return prior_model - - -def renderer(*, args, checkpoint_map_location): - print(" loading renderer") - - renderer_checkpoint = torch.load(args.transmitter_checkpoint_path, map_location=checkpoint_map_location) - - renderer_model = renderer_model_from_original_config() - - renderer_diffusers_checkpoint = renderer_model_original_checkpoint_to_diffusers_checkpoint( - renderer_model, renderer_checkpoint - ) - - del renderer_checkpoint - - load_checkpoint_to_model(renderer_diffusers_checkpoint, renderer_model, strict=True) - - print("done loading renderer") - - return renderer_model - - -# prior model will expect clip_mean and clip_std, which are missing from the state_dict -PRIOR_EXPECTED_MISSING_KEYS = ["clip_mean", "clip_std"] - - -def load_prior_checkpoint_to_model(checkpoint, model): - with tempfile.NamedTemporaryFile() as file: - torch.save(checkpoint, file.name) - del checkpoint - missing_keys, unexpected_keys = model.load_state_dict(torch.load(file.name), strict=False) - missing_keys = list(set(missing_keys) - set(PRIOR_EXPECTED_MISSING_KEYS)) - - if len(unexpected_keys) > 0: - raise ValueError(f"Unexpected keys when loading prior model: {unexpected_keys}") - if len(missing_keys) > 0: - raise ValueError(f"Missing keys when loading prior model: {missing_keys}") - - -def load_checkpoint_to_model(checkpoint, model, strict=False): - with tempfile.NamedTemporaryFile() as file: - torch.save(checkpoint, file.name) - del checkpoint - if strict: - model.load_state_dict(torch.load(file.name), strict=True) - else: - load_checkpoint_and_dispatch(model, file.name, device_map="auto") - - -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( - "--prior_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the prior checkpoint to convert.", - ) - - parser.add_argument( - "--prior_image_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the prior_image checkpoint to convert.", - ) - - parser.add_argument( - "--transmitter_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to the transmitter checkpoint to convert.", - ) - - parser.add_argument( - "--checkpoint_load_device", - default="cpu", - type=str, - required=False, - help="The device passed to `map_location` when loading checkpoints.", - ) - - parser.add_argument( - "--debug", - default=None, - type=str, - required=False, - help="Only run a specific stage of the convert script. Used for debugging", - ) - - args = parser.parse_args() - - print(f"loading checkpoints to {args.checkpoint_load_device}") - - checkpoint_map_location = torch.device(args.checkpoint_load_device) - - if args.debug is not None: - print(f"debug: only executing {args.debug}") - - if args.debug is None: - print("YiYi TO-DO") - elif args.debug == "prior": - prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) - prior_model.save_pretrained(args.dump_path) - elif args.debug == "prior_image": - prior_model = prior_image(args=args, checkpoint_map_location=checkpoint_map_location) - prior_model.save_pretrained(args.dump_path) - elif args.debug == "renderer": - renderer_model = renderer(args=args, checkpoint_map_location=checkpoint_map_location) - renderer_model.save_pretrained(args.dump_path) - else: - raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/convert_skyreelsv2_to_diffusers.py b/scripts/convert_skyreelsv2_to_diffusers.py deleted file mode 100644 index 29d5a7666a6e..000000000000 --- a/scripts/convert_skyreelsv2_to_diffusers.py +++ /dev/null @@ -1,637 +0,0 @@ -import argparse -import os -import pathlib -from typing import Any, Dict - -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download -from safetensors.torch import load_file -from transformers import AutoProcessor, AutoTokenizer, CLIPVisionModelWithProjection, UMT5EncoderModel - -from diffusers import ( - AutoencoderKLWan, - SkyReelsV2DiffusionForcingPipeline, - SkyReelsV2ImageToVideoPipeline, - SkyReelsV2Pipeline, - SkyReelsV2Transformer3DModel, - UniPCMultistepScheduler, -) - - -TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "ffn.0": "ffn.net.0.proj", - "ffn.2": "ffn.net.2", - "fps_projection.0": "fps_projection.net.0.proj", - "fps_projection.2": "fps_projection.net.2", - # Hack to swap the layer names - # The original model calls the norms in following order: norm1, norm3, norm2 - # We convert it to: norm1, norm2, norm3 - "norm2": "norm__placeholder", - "norm3": "norm2", - "norm__placeholder": "norm3", - # For the I2V model - "img_emb.proj.0": "condition_embedder.image_embedder.norm1", - "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", - "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", - "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # for the FLF2V model - "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", - # Add attention component mappings - "self_attn.q": "attn1.to_q", - "self_attn.k": "attn1.to_k", - "self_attn.v": "attn1.to_v", - "self_attn.o": "attn1.to_out.0", - "self_attn.norm_q": "attn1.norm_q", - "self_attn.norm_k": "attn1.norm_k", - "cross_attn.q": "attn2.to_q", - "cross_attn.k": "attn2.to_k", - "cross_attn.v": "attn2.to_v", - "cross_attn.o": "attn2.to_out.0", - "cross_attn.norm_q": "attn2.norm_q", - "cross_attn.norm_k": "attn2.norm_k", - "attn2.to_k_img": "attn2.add_k_proj", - "attn2.to_v_img": "attn2.add_v_proj", - "attn2.norm_k_img": "attn2.norm_added_k", -} - -TRANSFORMER_SPECIAL_KEYS_REMAP = {} - - -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def load_sharded_safetensors(dir: pathlib.Path): - if "720P" in str(dir): - file_paths = list(dir.glob("diffusion_pytorch_model*.safetensors")) - else: - file_paths = list(dir.glob("model*.safetensors")) - state_dict = {} - for path in file_paths: - state_dict.update(load_file(path)) - return state_dict - - -def get_transformer_config(model_type: str) -> dict[str, Any]: - if model_type == "SkyReels-V2-DF-1.3B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-DF-1.3B-540P", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 8960, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 12, - "inject_sample_info": True, - "num_layers": 30, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - elif model_type == "SkyReels-V2-DF-14B-720P": - config = { - "model_id": "Skywork/SkyReels-V2-DF-14B-720P", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - elif model_type == "SkyReels-V2-DF-14B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-DF-14B-540P", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - elif model_type == "SkyReels-V2-T2V-14B-720P": - config = { - "model_id": "Skywork/SkyReels-V2-T2V-14B-720P", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - elif model_type == "SkyReels-V2-T2V-14B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-T2V-14B-540P", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - elif model_type == "SkyReels-V2-I2V-1.3B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-1.3B-540P", - "diffusers_config": { - "added_kv_proj_dim": 1536, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 8960, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 12, - "inject_sample_info": False, - "num_layers": 30, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - }, - } - elif model_type == "SkyReels-V2-I2V-14B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-14B-540P", - "diffusers_config": { - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - }, - } - elif model_type == "SkyReels-V2-I2V-14B-720P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-14B-720P", - "diffusers_config": { - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - }, - } - elif model_type == "SkyReels-V2-FLF2V-1.3B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-1.3B-540P", - "diffusers_config": { - "added_kv_proj_dim": 1536, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 8960, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 12, - "inject_sample_info": False, - "num_layers": 30, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - "pos_embed_seq_len": 514, - }, - } - elif model_type == "SkyReels-V2-FLF2V-14B-540P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-14B-540P", - "diffusers_config": { - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - "pos_embed_seq_len": 514, - }, - } - elif model_type == "SkyReels-V2-FLF2V-14B-720P": - config = { - "model_id": "Skywork/SkyReels-V2-I2V-14B-720P", - "diffusers_config": { - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "inject_sample_info": False, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "image_dim": 1280, - "pos_embed_seq_len": 514, - }, - } - return config - - -def convert_transformer(model_type: str): - config = get_transformer_config(model_type) - diffusers_config = config["diffusers_config"] - model_id = config["model_id"] - - if "1.3B" in model_type: - original_state_dict = load_file(hf_hub_download(model_id, "model.safetensors")) - else: - os.makedirs(model_type, exist_ok=True) - model_dir = pathlib.Path(model_type) - if "720P" in model_type: - top_shard = 7 if "I2V" in model_type else 6 - zeros = "0" * (4 if "I2V" or "T2V" in model_type else 3) - model_name = "diffusion_pytorch_model" - elif "540P" in model_type: - top_shard = 14 if "I2V" in model_type else 12 - model_name = "model" - - for i in range(1, top_shard + 1): - shard_path = f"{model_name}-{i:05d}-of-{zeros}{top_shard}.safetensors" - hf_hub_download(model_id, shard_path, local_dir=model_dir) - original_state_dict = load_sharded_safetensors(model_dir) - - with init_empty_weights(): - transformer = SkyReelsV2Transformer3DModel.from_config(diffusers_config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - if "FLF2V" in model_type: - if ( - hasattr(transformer.condition_embedder, "image_embedder") - and hasattr(transformer.condition_embedder.image_embedder, "pos_embed") - and transformer.condition_embedder.image_embedder.pos_embed is not None - ): - pos_embed_shape = transformer.condition_embedder.image_embedder.pos_embed.shape - original_state_dict["condition_embedder.image_embedder.pos_embed"] = torch.zeros(pos_embed_shape) - - transformer.load_state_dict(original_state_dict, strict=True, assign=True) - return transformer - - -def convert_vae(): - vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") - old_state_dict = torch.load(vae_ckpt_path, weights_only=True) - new_state_dict = {} - - # Create mappings for specific components - middle_key_mapping = { - # Encoder middle block - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - # Create a mapping for attention blocks - attention_mapping = { - # Encoder middle attention - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - # Create a mapping for the head components - head_mapping = { - # Encoder head - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - # Create a mapping for the quant components - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - # Process each key in the state dict - for key, value in old_state_dict.items(): - # Handle middle block keys using the mapping - if key in middle_key_mapping: - new_key = middle_key_mapping[key] - new_state_dict[new_key] = value - # Handle attention blocks using the mapping - elif key in attention_mapping: - new_key = attention_mapping[key] - new_state_dict[new_key] = value - # Handle head keys using the mapping - elif key in head_mapping: - new_key = head_mapping[key] - new_state_dict[new_key] = value - # Handle quant keys using the mapping - elif key in quant_mapping: - new_key = quant_mapping[key] - new_state_dict[new_key] = value - # Handle encoder conv1 - elif key == "encoder.conv1.weight": - new_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - new_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 - elif key == "decoder.conv1.weight": - new_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - new_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples - elif key.startswith("encoder.downsamples."): - # Convert to down_blocks - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Convert residual block naming but keep the original structure - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - - new_state_dict[new_key] = value - - # Handle decoder upsamples - elif key.startswith("decoder.upsamples."): - # Convert to up_blocks - parts = key.split(".") - block_idx = int(parts[2]) - - # Group residual blocks - if "residual" in key: - if block_idx in [0, 1, 2]: - new_block_idx = 0 - resnet_idx = block_idx - elif block_idx in [4, 5, 6]: - new_block_idx = 1 - resnet_idx = block_idx - 4 - elif block_idx in [8, 9, 10]: - new_block_idx = 2 - resnet_idx = block_idx - 8 - elif block_idx in [12, 13, 14]: - new_block_idx = 3 - resnet_idx = block_idx - 12 - else: - # Keep as is for other blocks - new_state_dict[key] = value - continue - - # Convert residual block naming - if ".residual.0.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" - elif ".residual.2.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias" - elif ".residual.2.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight" - elif ".residual.3.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma" - elif ".residual.6.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias" - elif ".residual.6.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" - else: - new_key = key - - new_state_dict[new_key] = value - - # Handle shortcut connections - elif ".shortcut." in key: - if block_idx == 4: - new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") - new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - - new_state_dict[new_key] = value - - # Handle upsamplers - elif ".resample." in key or ".time_conv." in key: - if block_idx == 3: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") - elif block_idx == 7: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") - elif block_idx == 11: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - - new_state_dict[new_key] = value - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_state_dict[new_key] = value - else: - # Keep other keys unchanged - new_state_dict[key] = value - - with init_empty_weights(): - vae = AutoencoderKLWan() - vae.load_state_dict(new_state_dict, strict=True, assign=True) - return vae - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--model_type", type=str, default=None) - parser.add_argument("--output_path", type=str, required=True) - parser.add_argument("--dtype", default="fp32") - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -if __name__ == "__main__": - args = get_args() - - transformer = None - dtype = DTYPE_MAPPING[args.dtype] - - transformer = convert_transformer(args.model_type).to(dtype=dtype) - vae = convert_vae() - text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl") - tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") - scheduler = UniPCMultistepScheduler( - prediction_type="flow_prediction", - num_train_timesteps=1000, - use_flow_sigmas=True, - ) - - if "I2V" in args.model_type or "FLF2V" in args.model_type: - image_encoder = CLIPVisionModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") - image_processor = AutoProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") - pipe = SkyReelsV2ImageToVideoPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - elif "T2V" in args.model_type: - pipe = SkyReelsV2Pipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ) - elif "DF" in args.model_type: - pipe = SkyReelsV2DiffusionForcingPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ) - - pipe.save_pretrained( - args.output_path, - safe_serialization=True, - max_shard_size="5GB", - # push_to_hub=True, - # repo_id=f"/{args.model_type}-Diffusers", - ) diff --git a/scripts/convert_stable_audio.py b/scripts/convert_stable_audio.py deleted file mode 100644 index c3479fc6b2bb..000000000000 --- a/scripts/convert_stable_audio.py +++ /dev/null @@ -1,279 +0,0 @@ -# Run this script to convert the Stable Audio model weights to a diffusers pipeline. -import argparse -import json -import os -from contextlib import nullcontext - -import torch -from safetensors.torch import load_file -from transformers import ( - AutoTokenizer, - T5EncoderModel, -) - -from diffusers import ( - AutoencoderOobleck, - CosineDPMSolverMultistepScheduler, - StableAudioDiTModel, - StableAudioPipeline, - StableAudioProjectionModel, -) -from diffusers.models.model_loading_utils import load_model_dict_into_meta -from diffusers.utils import is_accelerate_available - - -if is_accelerate_available(): - from accelerate import init_empty_weights - - -def convert_stable_audio_state_dict_to_diffusers(state_dict, num_autoencoder_layers=5): - projection_model_state_dict = { - k.replace("conditioner.conditioners.", "").replace("embedder.embedding", "time_positional_embedding"): v - for (k, v) in state_dict.items() - if "conditioner.conditioners" in k - } - - # NOTE: we assume here that there's no projection layer from the text encoder to the latent space, script should be adapted a bit if there is. - for key, value in list(projection_model_state_dict.items()): - new_key = key.replace("seconds_start", "start_number_conditioner").replace( - "seconds_total", "end_number_conditioner" - ) - projection_model_state_dict[new_key] = projection_model_state_dict.pop(key) - - model_state_dict = {k.replace("model.model.", ""): v for (k, v) in state_dict.items() if "model.model." in k} - for key, value in list(model_state_dict.items()): - # attention layers - new_key = ( - key.replace("transformer.", "") - .replace("layers", "transformer_blocks") - .replace("self_attn", "attn1") - .replace("cross_attn", "attn2") - .replace("ff.ff", "ff.net") - ) - new_key = ( - new_key.replace("pre_norm", "norm1") - .replace("cross_attend_norm", "norm2") - .replace("ff_norm", "norm3") - .replace("to_out", "to_out.0") - ) - new_key = new_key.replace("gamma", "weight").replace("beta", "bias") # replace layernorm - - # other layers - new_key = ( - new_key.replace("project", "proj") - .replace("to_timestep_embed", "timestep_proj") - .replace("timestep_features", "time_proj") - .replace("to_global_embed", "global_proj") - .replace("to_cond_embed", "cross_attention_proj") - ) - - # we're using diffusers implementation of time_proj (GaussianFourierProjection) which creates a 1D tensor - if new_key == "time_proj.weight": - model_state_dict[key] = model_state_dict[key].squeeze(1) - - if "to_qkv" in new_key: - q, k, v = torch.chunk(model_state_dict.pop(key), 3, dim=0) - model_state_dict[new_key.replace("qkv", "q")] = q - model_state_dict[new_key.replace("qkv", "k")] = k - model_state_dict[new_key.replace("qkv", "v")] = v - elif "to_kv" in new_key: - k, v = torch.chunk(model_state_dict.pop(key), 2, dim=0) - model_state_dict[new_key.replace("kv", "k")] = k - model_state_dict[new_key.replace("kv", "v")] = v - else: - model_state_dict[new_key] = model_state_dict.pop(key) - - autoencoder_state_dict = { - k.replace("pretransform.model.", "").replace("coder.layers.0", "coder.conv1"): v - for (k, v) in state_dict.items() - if "pretransform.model." in k - } - - for key, _ in list(autoencoder_state_dict.items()): - new_key = key - if "coder.layers" in new_key: - # get idx of the layer - idx = int(new_key.split("coder.layers.")[1].split(".")[0]) - - new_key = new_key.replace(f"coder.layers.{idx}", f"coder.block.{idx - 1}") - - if "encoder" in new_key: - for i in range(3): - new_key = new_key.replace(f"block.{idx - 1}.layers.{i}", f"block.{idx - 1}.res_unit{i + 1}") - new_key = new_key.replace(f"block.{idx - 1}.layers.3", f"block.{idx - 1}.snake1") - new_key = new_key.replace(f"block.{idx - 1}.layers.4", f"block.{idx - 1}.conv1") - else: - for i in range(2, 5): - new_key = new_key.replace(f"block.{idx - 1}.layers.{i}", f"block.{idx - 1}.res_unit{i - 1}") - new_key = new_key.replace(f"block.{idx - 1}.layers.0", f"block.{idx - 1}.snake1") - new_key = new_key.replace(f"block.{idx - 1}.layers.1", f"block.{idx - 1}.conv_t1") - - new_key = new_key.replace("layers.0.beta", "snake1.beta") - new_key = new_key.replace("layers.0.alpha", "snake1.alpha") - new_key = new_key.replace("layers.2.beta", "snake2.beta") - new_key = new_key.replace("layers.2.alpha", "snake2.alpha") - new_key = new_key.replace("layers.1.bias", "conv1.bias") - new_key = new_key.replace("layers.1.weight_", "conv1.weight_") - new_key = new_key.replace("layers.3.bias", "conv2.bias") - new_key = new_key.replace("layers.3.weight_", "conv2.weight_") - - if idx == num_autoencoder_layers + 1: - new_key = new_key.replace(f"block.{idx - 1}", "snake1") - elif idx == num_autoencoder_layers + 2: - new_key = new_key.replace(f"block.{idx - 1}", "conv2") - - else: - new_key = new_key - - value = autoencoder_state_dict.pop(key) - if "snake" in new_key: - value = value.unsqueeze(0).unsqueeze(-1) - if new_key in autoencoder_state_dict: - raise ValueError(f"{new_key} already in state dict.") - autoencoder_state_dict[new_key] = value - - return model_state_dict, projection_model_state_dict, autoencoder_state_dict - - -parser = argparse.ArgumentParser(description="Convert Stable Audio 1.0 model weights to a diffusers pipeline") -parser.add_argument("--model_folder_path", type=str, help="Location of Stable Audio weights and config") -parser.add_argument("--use_safetensors", action="store_true", help="Use SafeTensors for conversion") -parser.add_argument( - "--save_directory", - type=str, - default="./tmp/stable-audio-1.0", - help="Directory to save a pipeline to. Will be created if it doesn't exist.", -) -parser.add_argument( - "--repo_id", - type=str, - default="stable-audio-1.0", - help="Hub organization to save the pipelines to", -) -parser.add_argument("--push_to_hub", action="store_true", help="Push to hub") -parser.add_argument("--variant", type=str, help="Set to bf16 to save bfloat16 weights") - -args = parser.parse_args() - -checkpoint_path = ( - os.path.join(args.model_folder_path, "model.safetensors") - if args.use_safetensors - else os.path.join(args.model_folder_path, "model.ckpt") -) -config_path = os.path.join(args.model_folder_path, "model_config.json") - -device = "cpu" -if args.variant == "bf16": - dtype = torch.bfloat16 -else: - dtype = torch.float32 - -with open(config_path) as f_in: - config_dict = json.load(f_in) - -conditioning_dict = { - conditioning["id"]: conditioning["config"] for conditioning in config_dict["model"]["conditioning"]["configs"] -} - -t5_model_config = conditioning_dict["prompt"] - -# T5 Text encoder -text_encoder = T5EncoderModel.from_pretrained(t5_model_config["t5_model_name"]) -tokenizer = AutoTokenizer.from_pretrained( - t5_model_config["t5_model_name"], truncation=True, model_max_length=t5_model_config["max_length"] -) - - -# scheduler -scheduler = CosineDPMSolverMultistepScheduler( - sigma_min=0.3, - sigma_max=500, - solver_order=2, - prediction_type="v_prediction", - sigma_data=1.0, - sigma_schedule="exponential", -) -ctx = init_empty_weights if is_accelerate_available() else nullcontext - - -if args.use_safetensors: - orig_state_dict = load_file(checkpoint_path, device=device) -else: - orig_state_dict = torch.load(checkpoint_path, map_location=device) - - -model_config = config_dict["model"]["diffusion"]["config"] - -model_state_dict, projection_model_state_dict, autoencoder_state_dict = convert_stable_audio_state_dict_to_diffusers( - orig_state_dict -) - - -with ctx(): - projection_model = StableAudioProjectionModel( - text_encoder_dim=text_encoder.config.d_model, - conditioning_dim=config_dict["model"]["conditioning"]["cond_dim"], - min_value=conditioning_dict["seconds_start"][ - "min_val" - ], # assume `seconds_start` and `seconds_total` have the same min / max values. - max_value=conditioning_dict["seconds_start"][ - "max_val" - ], # assume `seconds_start` and `seconds_total` have the same min / max values. - ) -if is_accelerate_available(): - load_model_dict_into_meta(projection_model, projection_model_state_dict) -else: - projection_model.load_state_dict(projection_model_state_dict) - -attention_head_dim = model_config["embed_dim"] // model_config["num_heads"] -with ctx(): - model = StableAudioDiTModel( - sample_size=int(config_dict["sample_size"]) - / int(config_dict["model"]["pretransform"]["config"]["downsampling_ratio"]), - in_channels=model_config["io_channels"], - num_layers=model_config["depth"], - attention_head_dim=attention_head_dim, - num_key_value_attention_heads=model_config["cond_token_dim"] // attention_head_dim, - num_attention_heads=model_config["num_heads"], - out_channels=model_config["io_channels"], - cross_attention_dim=model_config["cond_token_dim"], - time_proj_dim=256, - global_states_input_dim=model_config["global_cond_dim"], - cross_attention_input_dim=model_config["cond_token_dim"], - ) -if is_accelerate_available(): - load_model_dict_into_meta(model, model_state_dict) -else: - model.load_state_dict(model_state_dict) - - -autoencoder_config = config_dict["model"]["pretransform"]["config"] -with ctx(): - autoencoder = AutoencoderOobleck( - encoder_hidden_size=autoencoder_config["encoder"]["config"]["channels"], - downsampling_ratios=autoencoder_config["encoder"]["config"]["strides"], - decoder_channels=autoencoder_config["decoder"]["config"]["channels"], - decoder_input_channels=autoencoder_config["decoder"]["config"]["latent_dim"], - audio_channels=autoencoder_config["io_channels"], - channel_multiples=autoencoder_config["encoder"]["config"]["c_mults"], - sampling_rate=config_dict["sample_rate"], - ) - -if is_accelerate_available(): - load_model_dict_into_meta(autoencoder, autoencoder_state_dict) -else: - autoencoder.load_state_dict(autoencoder_state_dict) - - -# Prior pipeline -pipeline = StableAudioPipeline( - transformer=model, - tokenizer=tokenizer, - text_encoder=text_encoder, - scheduler=scheduler, - vae=autoencoder, - projection_model=projection_model, -) -pipeline.to(dtype).save_pretrained( - args.save_directory, repo_id=args.repo_id, push_to_hub=args.push_to_hub, variant=args.variant -) diff --git a/scripts/convert_svd_to_diffusers.py b/scripts/convert_svd_to_diffusers.py deleted file mode 100644 index e46410ccb3bd..000000000000 --- a/scripts/convert_svd_to_diffusers.py +++ /dev/null @@ -1,730 +0,0 @@ -from diffusers.utils import is_accelerate_available, logging - - -if is_accelerate_available(): - pass - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -def create_unet_diffusers_config(original_config, image_size: int, controlnet=False): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - if controlnet: - unet_params = original_config.model.params.control_stage_config.params - else: - if "unet_config" in original_config.model.params and original_config.model.params.unet_config is not None: - unet_params = original_config.model.params.unet_config.params - else: - unet_params = original_config.model.params.network_config.params - - vae_params = original_config.model.params.first_stage_config.params.encoder_config.params - - block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = ( - "CrossAttnDownBlockSpatioTemporal" - if resolution in unet_params.attention_resolutions - else "DownBlockSpatioTemporal" - ) - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = ( - "CrossAttnUpBlockSpatioTemporal" - if resolution in unet_params.attention_resolutions - else "UpBlockSpatioTemporal" - ) - up_block_types.append(block_type) - resolution //= 2 - - if unet_params.transformer_depth is not None: - transformer_layers_per_block = ( - unet_params.transformer_depth - if isinstance(unet_params.transformer_depth, int) - else list(unet_params.transformer_depth) - ) - else: - transformer_layers_per_block = 1 - - vae_scale_factor = 2 ** (len(vae_params.ch_mult) - 1) - - head_dim = unet_params.num_heads if "num_heads" in unet_params else None - use_linear_projection = ( - unet_params.use_linear_in_transformer if "use_linear_in_transformer" in unet_params else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim_mult = unet_params.model_channels // unet_params.num_head_channels - head_dim = [head_dim_mult * c for c in list(unet_params.channel_mult)] - - class_embed_type = None - addition_embed_type = None - addition_time_embed_dim = None - projection_class_embeddings_input_dim = None - context_dim = None - - if unet_params.context_dim is not None: - context_dim = ( - unet_params.context_dim if isinstance(unet_params.context_dim, int) else unet_params.context_dim[0] - ) - - if "num_classes" in unet_params: - if unet_params.num_classes == "sequential": - addition_time_embed_dim = 256 - assert "adm_in_channels" in unet_params - projection_class_embeddings_input_dim = unet_params.adm_in_channels - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params.in_channels, - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params.num_res_blocks, - "cross_attention_dim": context_dim, - "attention_head_dim": head_dim, - "use_linear_projection": use_linear_projection, - "class_embed_type": class_embed_type, - "addition_embed_type": addition_embed_type, - "addition_time_embed_dim": addition_time_embed_dim, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "transformer_layers_per_block": transformer_layers_per_block, - } - - if "disable_self_attentions" in unet_params: - config["only_cross_attention"] = unet_params.disable_self_attentions - - if "num_classes" in unet_params and isinstance(unet_params.num_classes, int): - config["num_class_embeds"] = unet_params.num_classes - - if controlnet: - config["conditioning_channels"] = unet_params.hint_channels - else: - config["out_channels"] = unet_params.out_channels - config["up_block_types"] = tuple(up_block_types) - - return config - - -def assign_to_checkpoint( - paths, - checkpoint, - old_checkpoint, - attention_paths_to_split=None, - additional_replacements=None, - config=None, - mid_block_suffix="", -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - if mid_block_suffix is not None: - mid_block_suffix = f".{mid_block_suffix}" - else: - mid_block_suffix = "" - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", f"mid_block.resnets.0{mid_block_suffix}") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", f"mid_block.resnets.1{mid_block_suffix}") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - if new_path == "mid_block.resnets.0.spatial_res_block.norm1.weight": - print("yeyy") - - # proj_attn.weight has to be converted from conv 1D to linear - is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) - shape = old_checkpoint[path["old"]].shape - if is_attn_weight and len(shape) == 3: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif is_attn_weight and len(shape) == 4: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - new_item = new_item.replace("time_stack", "temporal_transformer_blocks") - - new_item = new_item.replace("time_pos_embed.0.bias", "time_pos_embed.linear_1.bias") - new_item = new_item.replace("time_pos_embed.0.weight", "time_pos_embed.linear_1.weight") - new_item = new_item.replace("time_pos_embed.2.bias", "time_pos_embed.linear_2.bias") - new_item = new_item.replace("time_pos_embed.2.weight", "time_pos_embed.linear_2.weight") - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = new_item.replace("time_stack.", "") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def convert_ldm_unet_checkpoint( - checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False -): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - if skip_extract_state_dict: - unet_state_dict = checkpoint - else: - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning(f"Checkpoint {path} has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - # if config["addition_embed_type"] == "text_time": - new_checkpoint["add_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["add_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["add_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["add_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - spatial_resnets = [ - key - for key in input_blocks[i] - if f"input_blocks.{i}.0" in key - and ( - f"input_blocks.{i}.0.op" not in key - and f"input_blocks.{i}.0.time_stack" not in key - and f"input_blocks.{i}.0.time_mixer" not in key - ) - ] - temporal_resnets = [key for key in input_blocks[i] if f"input_blocks.{i}.0.time_stack" in key] - # import ipdb; ipdb.set_trace() - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(spatial_resnets) - meta_path = { - "old": f"input_blocks.{i}.0", - "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}.spatial_res_block", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - paths = renew_resnet_paths(temporal_resnets) - meta_path = { - "old": f"input_blocks.{i}.0", - "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}.temporal_res_block", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - # TODO resnet time_mixer.mix_factor - if f"input_blocks.{i}.0.time_mixer.mix_factor" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.resnets.{layer_in_block_id}.time_mixer.mix_factor"] = ( - unet_state_dict[f"input_blocks.{i}.0.time_mixer.mix_factor"] - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - # import ipdb; ipdb.set_trace() - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_spatial = [key for key in resnet_0 if "time_stack" not in key and "time_mixer" not in key] - resnet_0_paths = renew_resnet_paths(resnet_0_spatial) - # import ipdb; ipdb.set_trace() - assign_to_checkpoint( - resnet_0_paths, new_checkpoint, unet_state_dict, config=config, mid_block_suffix="spatial_res_block" - ) - - resnet_0_temporal = [key for key in resnet_0 if "time_stack" in key and "time_mixer" not in key] - resnet_0_paths = renew_resnet_paths(resnet_0_temporal) - assign_to_checkpoint( - resnet_0_paths, new_checkpoint, unet_state_dict, config=config, mid_block_suffix="temporal_res_block" - ) - - resnet_1_spatial = [key for key in resnet_1 if "time_stack" not in key and "time_mixer" not in key] - resnet_1_paths = renew_resnet_paths(resnet_1_spatial) - assign_to_checkpoint( - resnet_1_paths, new_checkpoint, unet_state_dict, config=config, mid_block_suffix="spatial_res_block" - ) - - resnet_1_temporal = [key for key in resnet_1 if "time_stack" in key and "time_mixer" not in key] - resnet_1_paths = renew_resnet_paths(resnet_1_temporal) - assign_to_checkpoint( - resnet_1_paths, new_checkpoint, unet_state_dict, config=config, mid_block_suffix="temporal_res_block" - ) - - new_checkpoint["mid_block.resnets.0.time_mixer.mix_factor"] = unet_state_dict[ - "middle_block.0.time_mixer.mix_factor" - ] - new_checkpoint["mid_block.resnets.1.time_mixer.mix_factor"] = unet_state_dict[ - "middle_block.2.time_mixer.mix_factor" - ] - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - spatial_resnets = [ - key - for key in output_blocks[i] - if f"output_blocks.{i}.0" in key - and (f"output_blocks.{i}.0.time_stack" not in key and "time_mixer" not in key) - ] - # import ipdb; ipdb.set_trace() - - temporal_resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0.time_stack" in key] - - paths = renew_resnet_paths(spatial_resnets) - meta_path = { - "old": f"output_blocks.{i}.0", - "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}.spatial_res_block", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - paths = renew_resnet_paths(temporal_resnets) - meta_path = { - "old": f"output_blocks.{i}.0", - "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}.temporal_res_block", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if f"output_blocks.{i}.0.time_mixer.mix_factor" in unet_state_dict: - new_checkpoint[f"up_blocks.{block_id}.resnets.{layer_in_block_id}.time_mixer.mix_factor"] = ( - unet_state_dict[f"output_blocks.{i}.0.time_mixer.mix_factor"] - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key and "conv" not in key] - if len(attentions): - paths = renew_attention_paths(attentions) - # import ipdb; ipdb.set_trace() - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - spatial_layers = [ - layer for layer in output_block_layers if "time_stack" not in layer and "time_mixer" not in layer - ] - resnet_0_paths = renew_resnet_paths(spatial_layers, n_shave_prefix_segments=1) - # import ipdb; ipdb.set_trace() - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join( - ["up_blocks", str(block_id), "resnets", str(layer_in_block_id), "spatial_res_block", path["new"]] - ) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - temporal_layers = [ - layer for layer in output_block_layers if "time_stack" in layer and "time_mixer" not in key - ] - resnet_0_paths = renew_resnet_paths(temporal_layers, n_shave_prefix_segments=1) - # import ipdb; ipdb.set_trace() - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join( - ["up_blocks", str(block_id), "resnets", str(layer_in_block_id), "temporal_res_block", path["new"]] - ) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - new_checkpoint["up_blocks.0.resnets.0.time_mixer.mix_factor"] = unet_state_dict[ - f"output_blocks.{str(i)}.0.time_mixer.mix_factor" - ] - - return new_checkpoint - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["to_q.weight", "to_k.weight", "to_v.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0, is_temporal=False): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # Temporal resnet - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = new_item.replace("time_stack.", "temporal_res_block.") - - # Spatial resnet - new_item = new_item.replace("conv1", "spatial_res_block.conv1") - new_item = new_item.replace("norm1", "spatial_res_block.norm1") - - new_item = new_item.replace("conv2", "spatial_res_block.conv2") - new_item = new_item.replace("norm2", "spatial_res_block.norm2") - - new_item = new_item.replace("nin_shortcut", "spatial_res_block.conv_shortcut") - - new_item = new_item.replace("mix_factor", "spatial_res_block.time_mixer.mix_factor") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - keys = list(checkpoint.keys()) - vae_key = "first_stage_model." if any(k.startswith("first_stage_model.") for k in keys) else "" - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - new_checkpoint["decoder.time_conv_out.weight"] = vae_state_dict["decoder.time_mix_conv.weight"] - new_checkpoint["decoder.time_conv_out.bias"] = vae_state_dict["decoder.time_mix_conv.bias"] - - # new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - # new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - # new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - # new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint diff --git a/scripts/convert_tiny_autoencoder_to_diffusers.py b/scripts/convert_tiny_autoencoder_to_diffusers.py deleted file mode 100644 index 9bb2df98a77a..000000000000 --- a/scripts/convert_tiny_autoencoder_to_diffusers.py +++ /dev/null @@ -1,71 +0,0 @@ -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) diff --git a/scripts/convert_unidiffuser_to_diffusers.py b/scripts/convert_unidiffuser_to_diffusers.py deleted file mode 100644 index 4c38172754f6..000000000000 --- a/scripts/convert_unidiffuser_to_diffusers.py +++ /dev/null @@ -1,786 +0,0 @@ -# Convert the original UniDiffuser checkpoints into diffusers equivalents. - -import argparse -from argparse import Namespace - -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextConfig, - CLIPTextModel, - CLIPTokenizer, - CLIPVisionConfig, - CLIPVisionModelWithProjection, - GPT2Tokenizer, -) - -from diffusers import ( - AutoencoderKL, - DPMSolverMultistepScheduler, - UniDiffuserModel, - UniDiffuserPipeline, - UniDiffuserTextDecoder, -) - - -SCHEDULER_CONFIG = Namespace( - **{ - "beta_start": 0.00085, - "beta_end": 0.012, - "beta_schedule": "scaled_linear", - "solver_order": 3, - } -) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.shave_segments -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_resnet_paths -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.renew_vae_attention_paths -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.conv_attn_to_linear -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -# Modified from diffusers.pipelines.stable_diffusion.convert_from_ckpt.assign_to_checkpoint -# config.num_head_channels => num_head_channels -def assign_to_checkpoint( - paths, - checkpoint, - old_checkpoint, - attention_paths_to_split=None, - additional_replacements=None, - num_head_channels=1, -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // num_head_channels // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) - shape = old_checkpoint[path["old"]].shape - if is_attn_weight and len(shape) == 3: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif is_attn_weight and len(shape) == 4: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def create_vae_diffusers_config(config_type): - # Hardcoded for now - if args.config_type == "test": - vae_config = create_vae_diffusers_config_test() - elif args.config_type == "big": - vae_config = create_vae_diffusers_config_big() - else: - raise NotImplementedError( - f"Config type {config_type} is not implemented, currently only config types" - " 'test' and 'big' are available." - ) - return vae_config - - -def create_unidiffuser_unet_config(config_type, version): - # Hardcoded for now - if args.config_type == "test": - unet_config = create_unidiffuser_unet_config_test() - elif args.config_type == "big": - unet_config = create_unidiffuser_unet_config_big() - else: - raise NotImplementedError( - f"Config type {config_type} is not implemented, currently only config types" - " 'test' and 'big' are available." - ) - # Unidiffuser-v1 uses data type embeddings - if version == 1: - unet_config["use_data_type_embedding"] = True - return unet_config - - -def create_text_decoder_config(config_type): - # Hardcoded for now - if args.config_type == "test": - text_decoder_config = create_text_decoder_config_test() - elif args.config_type == "big": - text_decoder_config = create_text_decoder_config_big() - else: - raise NotImplementedError( - f"Config type {config_type} is not implemented, currently only config types" - " 'test' and 'big' are available." - ) - return text_decoder_config - - -# Hardcoded configs for test versions of the UniDiffuser models, corresponding to those in the fast default tests. -def create_vae_diffusers_config_test(): - vae_config = { - "sample_size": 32, - "in_channels": 3, - "out_channels": 3, - "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"], - "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"], - "block_out_channels": [32, 64], - "latent_channels": 4, - "layers_per_block": 1, - } - return vae_config - - -def create_unidiffuser_unet_config_test(): - unet_config = { - "text_dim": 32, - "clip_img_dim": 32, - "num_text_tokens": 77, - "num_attention_heads": 2, - "attention_head_dim": 8, - "in_channels": 4, - "out_channels": 4, - "num_layers": 2, - "dropout": 0.0, - "norm_num_groups": 32, - "attention_bias": False, - "sample_size": 16, - "patch_size": 2, - "activation_fn": "gelu", - "num_embeds_ada_norm": 1000, - "norm_type": "layer_norm", - "block_type": "unidiffuser", - "pre_layer_norm": False, - "use_timestep_embedding": False, - "norm_elementwise_affine": True, - "use_patch_pos_embed": False, - "ff_final_dropout": True, - "use_data_type_embedding": False, - } - return unet_config - - -def create_text_decoder_config_test(): - text_decoder_config = { - "prefix_length": 77, - "prefix_inner_dim": 32, - "prefix_hidden_dim": 32, - "vocab_size": 1025, # 1024 + 1 for new EOS token - "n_positions": 1024, - "n_embd": 32, - "n_layer": 5, - "n_head": 4, - "n_inner": 37, - "activation_function": "gelu", - "resid_pdrop": 0.1, - "embd_pdrop": 0.1, - "attn_pdrop": 0.1, - "layer_norm_epsilon": 1e-5, - "initializer_range": 0.02, - } - return text_decoder_config - - -# Hardcoded configs for the UniDiffuser V1 model at https://huggingface.co/thu-ml/unidiffuser-v1 -# See also https://github.com/thu-ml/unidiffuser/blob/main/configs/sample_unidiffuser_v1.py -def create_vae_diffusers_config_big(): - vae_config = { - "sample_size": 256, - "in_channels": 3, - "out_channels": 3, - "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], - "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], - "block_out_channels": [128, 256, 512, 512], - "latent_channels": 4, - "layers_per_block": 2, - } - return vae_config - - -def create_unidiffuser_unet_config_big(): - unet_config = { - "text_dim": 64, - "clip_img_dim": 512, - "num_text_tokens": 77, - "num_attention_heads": 24, - "attention_head_dim": 64, - "in_channels": 4, - "out_channels": 4, - "num_layers": 30, - "dropout": 0.0, - "norm_num_groups": 32, - "attention_bias": False, - "sample_size": 64, - "patch_size": 2, - "activation_fn": "gelu", - "num_embeds_ada_norm": 1000, - "norm_type": "layer_norm", - "block_type": "unidiffuser", - "pre_layer_norm": False, - "use_timestep_embedding": False, - "norm_elementwise_affine": True, - "use_patch_pos_embed": False, - "ff_final_dropout": True, - "use_data_type_embedding": False, - } - return unet_config - - -# From https://huggingface.co/gpt2/blob/main/config.json, the GPT2 checkpoint used by UniDiffuser -def create_text_decoder_config_big(): - text_decoder_config = { - "prefix_length": 77, - "prefix_inner_dim": 768, - "prefix_hidden_dim": 64, - "vocab_size": 50258, # 50257 + 1 for new EOS token - "n_positions": 1024, - "n_embd": 768, - "n_layer": 12, - "n_head": 12, - "n_inner": 3072, - "activation_function": "gelu", - "resid_pdrop": 0.1, - "embd_pdrop": 0.1, - "attn_pdrop": 0.1, - "layer_norm_epsilon": 1e-5, - "initializer_range": 0.02, - } - return text_decoder_config - - -# Based on diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_vae_checkpoint -def convert_vae_to_diffusers(ckpt, diffusers_model, num_head_channels=1): - """ - Converts a UniDiffuser autoencoder_kl.pth checkpoint to a diffusers AutoencoderKL. - """ - # autoencoder_kl.pth ckpt is a torch state dict - vae_state_dict = torch.load(ckpt, map_location="cpu") - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - paths, - new_checkpoint, - vae_state_dict, - additional_replacements=[meta_path], - num_head_channels=num_head_channels, # not used in vae - ) - conv_attn_to_linear(new_checkpoint) - - missing_keys, unexpected_keys = diffusers_model.load_state_dict(new_checkpoint) - for missing_key in missing_keys: - print(f"Missing key: {missing_key}") - for unexpected_key in unexpected_keys: - print(f"Unexpected key: {unexpected_key}") - - return diffusers_model - - -def convert_uvit_block_to_diffusers_block( - uvit_state_dict, - new_state_dict, - block_prefix, - new_prefix="transformer.transformer_", - skip_connection=False, -): - """ - Maps the keys in a UniDiffuser transformer block (`Block`) to the keys in a diffusers transformer block - (`UTransformerBlock`/`UniDiffuserBlock`). - """ - prefix = new_prefix + block_prefix - if skip_connection: - new_state_dict[prefix + ".skip.skip_linear.weight"] = uvit_state_dict[block_prefix + ".skip_linear.weight"] - new_state_dict[prefix + ".skip.skip_linear.bias"] = uvit_state_dict[block_prefix + ".skip_linear.bias"] - new_state_dict[prefix + ".skip.norm.weight"] = uvit_state_dict[block_prefix + ".norm1.weight"] - new_state_dict[prefix + ".skip.norm.bias"] = uvit_state_dict[block_prefix + ".norm1.bias"] - - # Create the prefix string for out_blocks. - prefix += ".block" - - # Split up attention qkv.weight into to_q.weight, to_k.weight, to_v.weight - qkv = uvit_state_dict[block_prefix + ".attn.qkv.weight"] - new_attn_keys = [".attn1.to_q.weight", ".attn1.to_k.weight", ".attn1.to_v.weight"] - new_attn_keys = [prefix + key for key in new_attn_keys] - shape = qkv.shape[0] // len(new_attn_keys) - for i, attn_key in enumerate(new_attn_keys): - new_state_dict[attn_key] = qkv[i * shape : (i + 1) * shape] - - new_state_dict[prefix + ".attn1.to_out.0.weight"] = uvit_state_dict[block_prefix + ".attn.proj.weight"] - new_state_dict[prefix + ".attn1.to_out.0.bias"] = uvit_state_dict[block_prefix + ".attn.proj.bias"] - new_state_dict[prefix + ".norm1.weight"] = uvit_state_dict[block_prefix + ".norm2.weight"] - new_state_dict[prefix + ".norm1.bias"] = uvit_state_dict[block_prefix + ".norm2.bias"] - new_state_dict[prefix + ".ff.net.0.proj.weight"] = uvit_state_dict[block_prefix + ".mlp.fc1.weight"] - new_state_dict[prefix + ".ff.net.0.proj.bias"] = uvit_state_dict[block_prefix + ".mlp.fc1.bias"] - new_state_dict[prefix + ".ff.net.2.weight"] = uvit_state_dict[block_prefix + ".mlp.fc2.weight"] - new_state_dict[prefix + ".ff.net.2.bias"] = uvit_state_dict[block_prefix + ".mlp.fc2.bias"] - new_state_dict[prefix + ".norm3.weight"] = uvit_state_dict[block_prefix + ".norm3.weight"] - new_state_dict[prefix + ".norm3.bias"] = uvit_state_dict[block_prefix + ".norm3.bias"] - - return uvit_state_dict, new_state_dict - - -def convert_uvit_to_diffusers(ckpt, diffusers_model): - """ - Converts a UniDiffuser uvit_v*.pth checkpoint to a diffusers UniDiffusersModel. - """ - # uvit_v*.pth ckpt is a torch state dict - uvit_state_dict = torch.load(ckpt, map_location="cpu") - - new_state_dict = {} - - # Input layers - new_state_dict["vae_img_in.proj.weight"] = uvit_state_dict["patch_embed.proj.weight"] - new_state_dict["vae_img_in.proj.bias"] = uvit_state_dict["patch_embed.proj.bias"] - new_state_dict["clip_img_in.weight"] = uvit_state_dict["clip_img_embed.weight"] - new_state_dict["clip_img_in.bias"] = uvit_state_dict["clip_img_embed.bias"] - new_state_dict["text_in.weight"] = uvit_state_dict["text_embed.weight"] - new_state_dict["text_in.bias"] = uvit_state_dict["text_embed.bias"] - - new_state_dict["pos_embed"] = uvit_state_dict["pos_embed"] - - # Handle data type token embeddings for UniDiffuser-v1 - if "token_embedding.weight" in uvit_state_dict and diffusers_model.use_data_type_embedding: - new_state_dict["data_type_pos_embed_token"] = uvit_state_dict["pos_embed_token"] - new_state_dict["data_type_token_embedding.weight"] = uvit_state_dict["token_embedding.weight"] - - # Also initialize the PatchEmbedding in UTransformer2DModel with the PatchEmbedding from the checkpoint. - # This isn't used in the current implementation, so might want to remove. - new_state_dict["transformer.pos_embed.proj.weight"] = uvit_state_dict["patch_embed.proj.weight"] - new_state_dict["transformer.pos_embed.proj.bias"] = uvit_state_dict["patch_embed.proj.bias"] - - # Output layers - new_state_dict["transformer.norm_out.weight"] = uvit_state_dict["norm.weight"] - new_state_dict["transformer.norm_out.bias"] = uvit_state_dict["norm.bias"] - - new_state_dict["vae_img_out.weight"] = uvit_state_dict["decoder_pred.weight"] - new_state_dict["vae_img_out.bias"] = uvit_state_dict["decoder_pred.bias"] - new_state_dict["clip_img_out.weight"] = uvit_state_dict["clip_img_out.weight"] - new_state_dict["clip_img_out.bias"] = uvit_state_dict["clip_img_out.bias"] - new_state_dict["text_out.weight"] = uvit_state_dict["text_out.weight"] - new_state_dict["text_out.bias"] = uvit_state_dict["text_out.bias"] - - # in_blocks - in_blocks_prefixes = {".".join(layer.split(".")[:2]) for layer in uvit_state_dict if "in_blocks" in layer} - for in_block_prefix in list(in_blocks_prefixes): - convert_uvit_block_to_diffusers_block(uvit_state_dict, new_state_dict, in_block_prefix) - - # mid_block - # Assume there's only one mid block - convert_uvit_block_to_diffusers_block(uvit_state_dict, new_state_dict, "mid_block") - - # out_blocks - out_blocks_prefixes = {".".join(layer.split(".")[:2]) for layer in uvit_state_dict if "out_blocks" in layer} - for out_block_prefix in list(out_blocks_prefixes): - convert_uvit_block_to_diffusers_block(uvit_state_dict, new_state_dict, out_block_prefix, skip_connection=True) - - missing_keys, unexpected_keys = diffusers_model.load_state_dict(new_state_dict) - for missing_key in missing_keys: - print(f"Missing key: {missing_key}") - for unexpected_key in unexpected_keys: - print(f"Unexpected key: {unexpected_key}") - - return diffusers_model - - -def convert_caption_decoder_to_diffusers(ckpt, diffusers_model): - """ - Converts a UniDiffuser caption_decoder.pth checkpoint to a diffusers UniDiffuserTextDecoder. - """ - # caption_decoder.pth ckpt is a torch state dict - checkpoint_state_dict = torch.load(ckpt, map_location="cpu") - decoder_state_dict = {} - # Remove the "module." prefix, if necessary - caption_decoder_key = "module." - for key in checkpoint_state_dict: - if key.startswith(caption_decoder_key): - decoder_state_dict[key.replace(caption_decoder_key, "")] = checkpoint_state_dict.get(key) - else: - decoder_state_dict[key] = checkpoint_state_dict.get(key) - - new_state_dict = {} - - # Encoder and Decoder - new_state_dict["encode_prefix.weight"] = decoder_state_dict["encode_prefix.weight"] - new_state_dict["encode_prefix.bias"] = decoder_state_dict["encode_prefix.bias"] - new_state_dict["decode_prefix.weight"] = decoder_state_dict["decode_prefix.weight"] - new_state_dict["decode_prefix.bias"] = decoder_state_dict["decode_prefix.bias"] - - # Internal GPT2LMHeadModel transformer model - for key, val in decoder_state_dict.items(): - if key.startswith("gpt"): - suffix = key[len("gpt") :] - new_state_dict["transformer" + suffix] = val - - missing_keys, unexpected_keys = diffusers_model.load_state_dict(new_state_dict) - for missing_key in missing_keys: - print(f"Missing key: {missing_key}") - for unexpected_key in unexpected_keys: - print(f"Unexpected key: {unexpected_key}") - - return diffusers_model - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--caption_decoder_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to caption decoder checkpoint to convert.", - ) - parser.add_argument( - "--uvit_checkpoint_path", default=None, type=str, required=False, help="Path to U-ViT checkpoint to convert." - ) - parser.add_argument( - "--vae_checkpoint_path", - default=None, - type=str, - required=False, - help="Path to VAE checkpoint to convert.", - ) - parser.add_argument( - "--pipeline_output_path", - default=None, - type=str, - required=True, - help="Path to save the output pipeline to.", - ) - parser.add_argument( - "--config_type", - default="test", - type=str, - help=( - "Config type to use. Should be 'test' to create small models for testing or 'big' to convert a full" - " checkpoint." - ), - ) - parser.add_argument( - "--version", - default=0, - type=int, - help="The UniDiffuser model type to convert to. Should be 0 for UniDiffuser-v0 and 1 for UniDiffuser-v1.", - ) - parser.add_argument( - "--safe_serialization", - action="store_true", - help="Whether to use safetensors/safe seialization when saving the pipeline.", - ) - - args = parser.parse_args() - - # Convert the VAE model. - if args.vae_checkpoint_path is not None: - vae_config = create_vae_diffusers_config(args.config_type) - vae = AutoencoderKL(**vae_config) - vae = convert_vae_to_diffusers(args.vae_checkpoint_path, vae) - - # Convert the U-ViT ("unet") model. - if args.uvit_checkpoint_path is not None: - unet_config = create_unidiffuser_unet_config(args.config_type, args.version) - unet = UniDiffuserModel(**unet_config) - unet = convert_uvit_to_diffusers(args.uvit_checkpoint_path, unet) - - # Convert the caption decoder ("text_decoder") model. - if args.caption_decoder_checkpoint_path is not None: - text_decoder_config = create_text_decoder_config(args.config_type) - text_decoder = UniDiffuserTextDecoder(**text_decoder_config) - text_decoder = convert_caption_decoder_to_diffusers(args.caption_decoder_checkpoint_path, text_decoder) - - # Scheduler is the same for both the test and big models. - scheduler_config = SCHEDULER_CONFIG - scheduler = DPMSolverMultistepScheduler( - beta_start=scheduler_config.beta_start, - beta_end=scheduler_config.beta_end, - beta_schedule=scheduler_config.beta_schedule, - solver_order=scheduler_config.solver_order, - ) - - if args.config_type == "test": - # Make a small random CLIPTextModel - torch.manual_seed(0) - clip_text_encoder_config = CLIPTextConfig( - bos_token_id=0, - eos_token_id=2, - hidden_size=32, - intermediate_size=37, - layer_norm_eps=1e-05, - num_attention_heads=4, - num_hidden_layers=5, - pad_token_id=1, - vocab_size=1000, - ) - text_encoder = CLIPTextModel(clip_text_encoder_config) - clip_tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") - - # Make a small random CLIPVisionModel and accompanying CLIPImageProcessor - torch.manual_seed(0) - clip_image_encoder_config = CLIPVisionConfig( - image_size=32, - patch_size=2, - num_channels=3, - hidden_size=32, - projection_dim=32, - num_hidden_layers=5, - num_attention_heads=4, - intermediate_size=37, - dropout=0.1, - attention_dropout=0.1, - initializer_range=0.02, - ) - image_encoder = CLIPVisionModelWithProjection(clip_image_encoder_config) - image_processor = CLIPImageProcessor(crop_size=32, size=32) - - # Note that the text_decoder should already have its token embeddings resized. - text_tokenizer = GPT2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-GPT2Model") - eos = "<|EOS|>" - special_tokens_dict = {"eos_token": eos} - text_tokenizer.add_special_tokens(special_tokens_dict) - elif args.config_type == "big": - text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") - clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") - - image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") - image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32") - - # Note that the text_decoder should already have its token embeddings resized. - text_tokenizer = GPT2Tokenizer.from_pretrained("gpt2") - eos = "<|EOS|>" - special_tokens_dict = {"eos_token": eos} - text_tokenizer.add_special_tokens(special_tokens_dict) - else: - raise NotImplementedError( - f"Config type {args.config_type} is not implemented, currently only config types" - " 'test' and 'big' are available." - ) - - pipeline = UniDiffuserPipeline( - vae=vae, - text_encoder=text_encoder, - image_encoder=image_encoder, - clip_image_processor=image_processor, - clip_tokenizer=clip_tokenizer, - text_decoder=text_decoder, - text_tokenizer=text_tokenizer, - unet=unet, - scheduler=scheduler, - ) - pipeline.save_pretrained(args.pipeline_output_path, safe_serialization=args.safe_serialization) diff --git a/scripts/convert_vae_pt_to_diffusers.py b/scripts/convert_vae_pt_to_diffusers.py deleted file mode 100644 index 8c7dc71ddfd8..000000000000 --- a/scripts/convert_vae_pt_to_diffusers.py +++ /dev/null @@ -1,177 +0,0 @@ -import argparse -import io - -import requests -import torch -import yaml - -from diffusers import AutoencoderKL -from diffusers.pipelines.stable_diffusion.convert_from_ckpt import ( - assign_to_checkpoint, - conv_attn_to_linear, - create_vae_diffusers_config, - renew_vae_attention_paths, - renew_vae_resnet_paths, -) -from diffusers.utils.constants import DIFFUSERS_REQUEST_TIMEOUT - - -def custom_convert_ldm_vae_checkpoint(checkpoint, config): - vae_state_dict = checkpoint - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [ - key - for key in down_blocks[i] - if f"down.{i}" in key and f"down.{i}.downsample" not in key and "attn" not in key - ] - attentions = [key for key in down_blocks[i] if f"down.{i}.attn" in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - paths = renew_vae_attention_paths(attentions) - meta_path = {"old": f"down.{i}.attn", "new": f"down_blocks.{i}.attentions"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key - for key in up_blocks[block_id] - if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key and "attn" not in key - ] - attentions = [key for key in up_blocks[block_id] if f"up.{block_id}.attn" in key] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - paths = renew_vae_attention_paths(attentions) - meta_path = {"old": f"up.{block_id}.attn", "new": f"up_blocks.{i}.attentions"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -def vae_pt_to_vae_diffuser( - checkpoint_path: str, - output_path: str, -): - # Only support V1 - r = requests.get( - " https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml", - timeout=DIFFUSERS_REQUEST_TIMEOUT, - ) - io_obj = io.BytesIO(r.content) - - original_config = yaml.safe_load(io_obj) - image_size = 512 - device = "cuda" if torch.cuda.is_available() else "cpu" - if checkpoint_path.endswith("safetensors"): - from safetensors import safe_open - - checkpoint = {} - with safe_open(checkpoint_path, framework="pt", device="cpu") as f: - for key in f.keys(): - checkpoint[key] = f.get_tensor(key) - else: - checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"] - - # Convert the VAE model. - vae_config = create_vae_diffusers_config(original_config, image_size=image_size) - converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - vae.save_pretrained(output_path) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument("--vae_pt_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.") - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.") - - args = parser.parse_args() - - vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path) diff --git a/scripts/convert_versatile_diffusion_to_diffusers.py b/scripts/convert_versatile_diffusion_to_diffusers.py deleted file mode 100644 index ce68bb4c2e8c..000000000000 --- a/scripts/convert_versatile_diffusion_to_diffusers.py +++ /dev/null @@ -1,791 +0,0 @@ -# coding=utf-8 -# Copyright 2025 The HuggingFace Inc. team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Conversion script for the Versatile Stable Diffusion checkpoints.""" - -import argparse -from argparse import Namespace - -import torch -from transformers import ( - CLIPImageProcessor, - CLIPTextModelWithProjection, - CLIPTokenizer, - CLIPVisionModelWithProjection, -) - -from diffusers import ( - AutoencoderKL, - DDIMScheduler, - DPMSolverMultistepScheduler, - EulerAncestralDiscreteScheduler, - EulerDiscreteScheduler, - LMSDiscreteScheduler, - PNDMScheduler, - UNet2DConditionModel, - VersatileDiffusionPipeline, -) -from diffusers.pipelines.versatile_diffusion.modeling_text_unet import UNetFlatConditionModel - - -SCHEDULER_CONFIG = Namespace( - **{ - "beta_linear_start": 0.00085, - "beta_linear_end": 0.012, - "timesteps": 1000, - "scale_factor": 0.18215, - } -) - -IMAGE_UNET_CONFIG = Namespace( - **{ - "input_channels": 4, - "model_channels": 320, - "output_channels": 4, - "num_noattn_blocks": [2, 2, 2, 2], - "channel_mult": [1, 2, 4, 4], - "with_attn": [True, True, True, False], - "num_heads": 8, - "context_dim": 768, - "use_checkpoint": True, - } -) - -TEXT_UNET_CONFIG = Namespace( - **{ - "input_channels": 768, - "model_channels": 320, - "output_channels": 768, - "num_noattn_blocks": [2, 2, 2, 2], - "channel_mult": [1, 2, 4, 4], - "second_dim": [4, 4, 4, 4], - "with_attn": [True, True, True, False], - "num_heads": 8, - "context_dim": 768, - "use_checkpoint": True, - } -) - -AUTOENCODER_CONFIG = Namespace( - **{ - "double_z": True, - "z_channels": 4, - "resolution": 256, - "in_channels": 3, - "out_ch": 3, - "ch": 128, - "ch_mult": [1, 2, 4, 4], - "num_res_blocks": 2, - "attn_resolutions": [], - "dropout": 0.0, - } -) - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "query.weight") - new_item = new_item.replace("q.bias", "query.bias") - - new_item = new_item.replace("k.weight", "key.weight") - new_item = new_item.replace("k.bias", "key.bias") - - new_item = new_item.replace("v.weight", "value.weight") - new_item = new_item.replace("v.bias", "value.bias") - - new_item = new_item.replace("proj_out.weight", "proj_attn.weight") - new_item = new_item.replace("proj_out.bias", "proj_attn.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming - to them. It splits attention layers, and takes into account additional replacements - that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - if "proj_attn.weight" in new_path: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif path["old"] in old_checkpoint: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def create_image_unet_diffusers_config(unet_params): - """ - Creates a config for the diffusers based on the config of the VD model. - """ - - block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if unet_params.with_attn[i] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if unet_params.with_attn[-i - 1] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - if not all(n == unet_params.num_noattn_blocks[0] for n in unet_params.num_noattn_blocks): - raise ValueError("Not all num_res_blocks are equal, which is not supported in this script.") - - config = { - "sample_size": None, - "in_channels": unet_params.input_channels, - "out_channels": unet_params.output_channels, - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params.num_noattn_blocks[0], - "cross_attention_dim": unet_params.context_dim, - "attention_head_dim": unet_params.num_heads, - } - - return config - - -def create_text_unet_diffusers_config(unet_params): - """ - Creates a config for the diffusers based on the config of the VD model. - """ - - block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlockFlat" if unet_params.with_attn[i] else "DownBlockFlat" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlockFlat" if unet_params.with_attn[-i - 1] else "UpBlockFlat" - up_block_types.append(block_type) - resolution //= 2 - - if not all(n == unet_params.num_noattn_blocks[0] for n in unet_params.num_noattn_blocks): - raise ValueError("Not all num_res_blocks are equal, which is not supported in this script.") - - config = { - "sample_size": None, - "in_channels": (unet_params.input_channels, 1, 1), - "out_channels": (unet_params.output_channels, 1, 1), - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params.num_noattn_blocks[0], - "cross_attention_dim": unet_params.context_dim, - "attention_head_dim": unet_params.num_heads, - } - - return config - - -def create_vae_diffusers_config(vae_params): - """ - Creates a config for the diffusers based on the config of the VD model. - """ - - block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - config = { - "sample_size": vae_params.resolution, - "in_channels": vae_params.in_channels, - "out_channels": vae_params.out_ch, - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params.z_channels, - "layers_per_block": vae_params.num_res_blocks, - } - return config - - -def create_diffusers_scheduler(original_config): - schedular = DDIMScheduler( - num_train_timesteps=original_config.model.params.timesteps, - beta_start=original_config.model.params.linear_start, - beta_end=original_config.model.params.linear_end, - beta_schedule="scaled_linear", - ) - return schedular - - -def convert_vd_unet_checkpoint(checkpoint, config, unet_key, extract_ema=False): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100: - print("Checkpoint has both EMA and non-EMA weights.") - if extract_ema: - print( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - print( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = checkpoint["model.diffusion_model.time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = checkpoint["model.diffusion_model.time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = checkpoint["model.diffusion_model.time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = checkpoint["model.diffusion_model.time_embed.2.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - elif f"input_blocks.{i}.0.weight" in unet_state_dict: - # text_unet uses linear layers in place of downsamplers - shape = unet_state_dict[f"input_blocks.{i}.0.weight"].shape - if shape[0] != shape[1]: - continue - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if ["conv.weight", "conv.bias"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.weight", "conv.bias"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - elif f"output_blocks.{i}.1.weight" in unet_state_dict: - # text_unet uses linear layers in place of upsamplers - shape = unet_state_dict[f"output_blocks.{i}.1.weight"].shape - if shape[0] != shape[1]: - continue - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.weight"] = unet_state_dict.pop( - f"output_blocks.{i}.1.weight" - ) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.bias"] = unet_state_dict.pop( - f"output_blocks.{i}.1.bias" - ) - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - elif f"output_blocks.{i}.2.weight" in unet_state_dict: - # text_unet uses linear layers in place of upsamplers - shape = unet_state_dict[f"output_blocks.{i}.2.weight"].shape - if shape[0] != shape[1]: - continue - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.weight"] = unet_state_dict.pop( - f"output_blocks.{i}.2.weight" - ) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.bias"] = unet_state_dict.pop( - f"output_blocks.{i}.2.bias" - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - return new_checkpoint - - -def convert_vd_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - keys = list(checkpoint.keys()) - for key in keys: - vae_state_dict[key] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--unet_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--vae_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--optimus_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--scheduler_type", - default="pndm", - type=str, - help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - - args = parser.parse_args() - - scheduler_config = SCHEDULER_CONFIG - - num_train_timesteps = scheduler_config.timesteps - beta_start = scheduler_config.beta_linear_start - beta_end = scheduler_config.beta_linear_end - if args.scheduler_type == "pndm": - scheduler = PNDMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - skip_prk_steps=True, - steps_offset=1, - ) - elif args.scheduler_type == "lms": - scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear") - elif args.scheduler_type == "euler": - scheduler = EulerDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear") - elif args.scheduler_type == "euler-ancestral": - scheduler = EulerAncestralDiscreteScheduler( - beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear" - ) - elif args.scheduler_type == "dpm": - scheduler = DPMSolverMultistepScheduler( - beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear" - ) - elif args.scheduler_type == "ddim": - scheduler = DDIMScheduler( - beta_start=beta_start, - beta_end=beta_end, - beta_schedule="scaled_linear", - clip_sample=False, - set_alpha_to_one=False, - steps_offset=1, - ) - else: - raise ValueError(f"Scheduler of type {args.scheduler_type} doesn't exist!") - - # Convert the UNet2DConditionModel models. - if args.unet_checkpoint_path is not None: - # image UNet - image_unet_config = create_image_unet_diffusers_config(IMAGE_UNET_CONFIG) - checkpoint = torch.load(args.unet_checkpoint_path) - converted_image_unet_checkpoint = convert_vd_unet_checkpoint( - checkpoint, image_unet_config, unet_key="model.diffusion_model.unet_image.", extract_ema=args.extract_ema - ) - image_unet = UNet2DConditionModel(**image_unet_config) - image_unet.load_state_dict(converted_image_unet_checkpoint) - - # text UNet - text_unet_config = create_text_unet_diffusers_config(TEXT_UNET_CONFIG) - converted_text_unet_checkpoint = convert_vd_unet_checkpoint( - checkpoint, text_unet_config, unet_key="model.diffusion_model.unet_text.", extract_ema=args.extract_ema - ) - text_unet = UNetFlatConditionModel(**text_unet_config) - text_unet.load_state_dict(converted_text_unet_checkpoint) - - # Convert the VAE model. - if args.vae_checkpoint_path is not None: - vae_config = create_vae_diffusers_config(AUTOENCODER_CONFIG) - checkpoint = torch.load(args.vae_checkpoint_path) - converted_vae_checkpoint = convert_vd_vae_checkpoint(checkpoint, vae_config) - - vae = AutoencoderKL(**vae_config) - vae.load_state_dict(converted_vae_checkpoint) - - tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") - image_feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14") - text_encoder = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") - image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") - - pipe = VersatileDiffusionPipeline( - scheduler=scheduler, - tokenizer=tokenizer, - image_feature_extractor=image_feature_extractor, - text_encoder=text_encoder, - image_encoder=image_encoder, - image_unet=image_unet, - text_unet=text_unet, - vae=vae, - ) - pipe.save_pretrained(args.dump_path) diff --git a/scripts/convert_vq_diffusion_to_diffusers.py b/scripts/convert_vq_diffusion_to_diffusers.py deleted file mode 100644 index cd705e57b33c..000000000000 --- a/scripts/convert_vq_diffusion_to_diffusers.py +++ /dev/null @@ -1,916 +0,0 @@ -""" -This script ports models from VQ-diffusion (https://github.com/microsoft/VQ-Diffusion) to diffusers. - -It currently only supports porting the ITHQ dataset. - -ITHQ dataset: -```sh -# From the root directory of diffusers. - -# Download the VQVAE checkpoint -$ Refer to https://github.com/microsoft/VQ-Diffusion/blob/main/vqdiffusion_download_checkpoints.sh - -# Download the VQVAE config -# NOTE that in VQ-diffusion the documented file is `configs/ithq.yaml` but the target class -# `image_synthesis.modeling.codecs.image_codec.ema_vqvae.PatchVQVAE` -# loads `OUTPUT/pretrained_model/taming_dvae/config.yaml` -$ wget https://raw.githubusercontent.com/microsoft/VQ-Diffusion/main/OUTPUT/pretrained_model/taming_dvae/config.yaml -O ithq_vqvae.yaml - -# Download the main model checkpoint -$ wget https://facevcstandard.blob.core.windows.net/v-zhictang/Improved-VQ-Diffusion_model_release/ithq_learnable.pth?sv=2020-10-02&st=2022-05-30T10%3A22%3A06Z&se=2030-05-31T10%3A22%3A00Z&sr=b&sp=r&sig=GOE%2Bza02%2FPnGxYVOOPtwrTR4RA3%2F5NVgMxdW4kjaEZ8%3D -O ithq_learnable.pth # trufflehog:ignore - -# Download the main model config -$ wget https://raw.githubusercontent.com/microsoft/VQ-Diffusion/main/configs/ithq.yaml -O ithq.yaml - -# run the convert script -$ python ./scripts/convert_vq_diffusion_to_diffusers.py \ - --checkpoint_path ./ithq_learnable.pth \ - --original_config_file ./ithq.yaml \ - --vqvae_checkpoint_path ./ithq_vqvae.pth \ - --vqvae_original_config_file ./ithq_vqvae.yaml \ - --dump_path -``` -""" - -import argparse -import tempfile - -import torch -import yaml -from accelerate import init_empty_weights, load_checkpoint_and_dispatch -from transformers import CLIPTextModel, CLIPTokenizer -from yaml.loader import FullLoader - -from diffusers import Transformer2DModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel -from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings - - -# vqvae model - -PORTED_VQVAES = ["image_synthesis.modeling.codecs.image_codec.patch_vqgan.PatchVQGAN"] - - -def vqvae_model_from_original_config(original_config): - assert original_config["target"] in PORTED_VQVAES, ( - f"{original_config['target']} has not yet been ported to diffusers." - ) - - original_config = original_config["params"] - - original_encoder_config = original_config["encoder_config"]["params"] - original_decoder_config = original_config["decoder_config"]["params"] - - in_channels = original_encoder_config["in_channels"] - out_channels = original_decoder_config["out_ch"] - - down_block_types = get_down_block_types(original_encoder_config) - up_block_types = get_up_block_types(original_decoder_config) - - assert original_encoder_config["ch"] == original_decoder_config["ch"] - assert original_encoder_config["ch_mult"] == original_decoder_config["ch_mult"] - block_out_channels = tuple( - [original_encoder_config["ch"] * a_ch_mult for a_ch_mult in original_encoder_config["ch_mult"]] - ) - - assert original_encoder_config["num_res_blocks"] == original_decoder_config["num_res_blocks"] - layers_per_block = original_encoder_config["num_res_blocks"] - - assert original_encoder_config["z_channels"] == original_decoder_config["z_channels"] - latent_channels = original_encoder_config["z_channels"] - - num_vq_embeddings = original_config["n_embed"] - - # Hard coded value for ResnetBlock.GoupNorm(num_groups) in VQ-diffusion - norm_num_groups = 32 - - e_dim = original_config["embed_dim"] - - model = VQModel( - in_channels=in_channels, - out_channels=out_channels, - down_block_types=down_block_types, - up_block_types=up_block_types, - block_out_channels=block_out_channels, - layers_per_block=layers_per_block, - latent_channels=latent_channels, - num_vq_embeddings=num_vq_embeddings, - norm_num_groups=norm_num_groups, - vq_embed_dim=e_dim, - ) - - return model - - -def get_down_block_types(original_encoder_config): - attn_resolutions = coerce_attn_resolutions(original_encoder_config["attn_resolutions"]) - num_resolutions = len(original_encoder_config["ch_mult"]) - resolution = coerce_resolution(original_encoder_config["resolution"]) - - curr_res = resolution - down_block_types = [] - - for _ in range(num_resolutions): - if curr_res in attn_resolutions: - down_block_type = "AttnDownEncoderBlock2D" - else: - down_block_type = "DownEncoderBlock2D" - - down_block_types.append(down_block_type) - - curr_res = [r // 2 for r in curr_res] - - return down_block_types - - -def get_up_block_types(original_decoder_config): - attn_resolutions = coerce_attn_resolutions(original_decoder_config["attn_resolutions"]) - num_resolutions = len(original_decoder_config["ch_mult"]) - resolution = coerce_resolution(original_decoder_config["resolution"]) - - curr_res = [r // 2 ** (num_resolutions - 1) for r in resolution] - up_block_types = [] - - for _ in reversed(range(num_resolutions)): - if curr_res in attn_resolutions: - up_block_type = "AttnUpDecoderBlock2D" - else: - up_block_type = "UpDecoderBlock2D" - - up_block_types.append(up_block_type) - - curr_res = [r * 2 for r in curr_res] - - return up_block_types - - -def coerce_attn_resolutions(attn_resolutions): - attn_resolutions = list(attn_resolutions) - attn_resolutions_ = [] - for ar in attn_resolutions: - if isinstance(ar, (list, tuple)): - attn_resolutions_.append(list(ar)) - else: - attn_resolutions_.append([ar, ar]) - return attn_resolutions_ - - -def coerce_resolution(resolution): - if isinstance(resolution, int): - resolution = [resolution, resolution] # H, W - elif isinstance(resolution, (tuple, list)): - resolution = list(resolution) - else: - raise ValueError("Unknown type of resolution:", resolution) - return resolution - - -# done vqvae model - -# vqvae checkpoint - - -def vqvae_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - diffusers_checkpoint.update(vqvae_encoder_to_diffusers_checkpoint(model, checkpoint)) - - # quant_conv - - diffusers_checkpoint.update( - { - "quant_conv.weight": checkpoint["quant_conv.weight"], - "quant_conv.bias": checkpoint["quant_conv.bias"], - } - ) - - # quantize - diffusers_checkpoint.update({"quantize.embedding.weight": checkpoint["quantize.embedding"]}) - - # post_quant_conv - diffusers_checkpoint.update( - { - "post_quant_conv.weight": checkpoint["post_quant_conv.weight"], - "post_quant_conv.bias": checkpoint["post_quant_conv.bias"], - } - ) - - # decoder - diffusers_checkpoint.update(vqvae_decoder_to_diffusers_checkpoint(model, checkpoint)) - - return diffusers_checkpoint - - -def vqvae_encoder_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # conv_in - diffusers_checkpoint.update( - { - "encoder.conv_in.weight": checkpoint["encoder.conv_in.weight"], - "encoder.conv_in.bias": checkpoint["encoder.conv_in.bias"], - } - ) - - # down_blocks - for down_block_idx, down_block in enumerate(model.encoder.down_blocks): - diffusers_down_block_prefix = f"encoder.down_blocks.{down_block_idx}" - down_block_prefix = f"encoder.down.{down_block_idx}" - - # resnets - for resnet_idx, resnet in enumerate(down_block.resnets): - diffusers_resnet_prefix = f"{diffusers_down_block_prefix}.resnets.{resnet_idx}" - resnet_prefix = f"{down_block_prefix}.block.{resnet_idx}" - - diffusers_checkpoint.update( - vqvae_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - # downsample - - # do not include the downsample when on the last down block - # There is no downsample on the last down block - if down_block_idx != len(model.encoder.down_blocks) - 1: - # There's a single downsample in the original checkpoint but a list of downsamples - # in the diffusers model. - diffusers_downsample_prefix = f"{diffusers_down_block_prefix}.downsamplers.0.conv" - downsample_prefix = f"{down_block_prefix}.downsample.conv" - diffusers_checkpoint.update( - { - f"{diffusers_downsample_prefix}.weight": checkpoint[f"{downsample_prefix}.weight"], - f"{diffusers_downsample_prefix}.bias": checkpoint[f"{downsample_prefix}.bias"], - } - ) - - # attentions - - if hasattr(down_block, "attentions"): - for attention_idx, _ in enumerate(down_block.attentions): - diffusers_attention_prefix = f"{diffusers_down_block_prefix}.attentions.{attention_idx}" - attention_prefix = f"{down_block_prefix}.attn.{attention_idx}" - diffusers_checkpoint.update( - vqvae_attention_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - attention_prefix=attention_prefix, - ) - ) - - # mid block - - # mid block attentions - - # There is a single hardcoded attention block in the middle of the VQ-diffusion encoder - diffusers_attention_prefix = "encoder.mid_block.attentions.0" - attention_prefix = "encoder.mid.attn_1" - diffusers_checkpoint.update( - vqvae_attention_to_diffusers_checkpoint( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # mid block resnets - - for diffusers_resnet_idx, resnet in enumerate(model.encoder.mid_block.resnets): - diffusers_resnet_prefix = f"encoder.mid_block.resnets.{diffusers_resnet_idx}" - - # the hardcoded prefixes to `block_` are 1 and 2 - orig_resnet_idx = diffusers_resnet_idx + 1 - # There are two hardcoded resnets in the middle of the VQ-diffusion encoder - resnet_prefix = f"encoder.mid.block_{orig_resnet_idx}" - - diffusers_checkpoint.update( - vqvae_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - diffusers_checkpoint.update( - { - # conv_norm_out - "encoder.conv_norm_out.weight": checkpoint["encoder.norm_out.weight"], - "encoder.conv_norm_out.bias": checkpoint["encoder.norm_out.bias"], - # conv_out - "encoder.conv_out.weight": checkpoint["encoder.conv_out.weight"], - "encoder.conv_out.bias": checkpoint["encoder.conv_out.bias"], - } - ) - - return diffusers_checkpoint - - -def vqvae_decoder_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - # conv in - diffusers_checkpoint.update( - { - "decoder.conv_in.weight": checkpoint["decoder.conv_in.weight"], - "decoder.conv_in.bias": checkpoint["decoder.conv_in.bias"], - } - ) - - # up_blocks - - for diffusers_up_block_idx, up_block in enumerate(model.decoder.up_blocks): - # up_blocks are stored in reverse order in the VQ-diffusion checkpoint - orig_up_block_idx = len(model.decoder.up_blocks) - 1 - diffusers_up_block_idx - - diffusers_up_block_prefix = f"decoder.up_blocks.{diffusers_up_block_idx}" - up_block_prefix = f"decoder.up.{orig_up_block_idx}" - - # resnets - for resnet_idx, resnet in enumerate(up_block.resnets): - diffusers_resnet_prefix = f"{diffusers_up_block_prefix}.resnets.{resnet_idx}" - resnet_prefix = f"{up_block_prefix}.block.{resnet_idx}" - - diffusers_checkpoint.update( - vqvae_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - # upsample - - # there is no up sample on the last up block - if diffusers_up_block_idx != len(model.decoder.up_blocks) - 1: - # There's a single upsample in the VQ-diffusion checkpoint but a list of downsamples - # in the diffusers model. - diffusers_downsample_prefix = f"{diffusers_up_block_prefix}.upsamplers.0.conv" - downsample_prefix = f"{up_block_prefix}.upsample.conv" - diffusers_checkpoint.update( - { - f"{diffusers_downsample_prefix}.weight": checkpoint[f"{downsample_prefix}.weight"], - f"{diffusers_downsample_prefix}.bias": checkpoint[f"{downsample_prefix}.bias"], - } - ) - - # attentions - - if hasattr(up_block, "attentions"): - for attention_idx, _ in enumerate(up_block.attentions): - diffusers_attention_prefix = f"{diffusers_up_block_prefix}.attentions.{attention_idx}" - attention_prefix = f"{up_block_prefix}.attn.{attention_idx}" - diffusers_checkpoint.update( - vqvae_attention_to_diffusers_checkpoint( - checkpoint, - diffusers_attention_prefix=diffusers_attention_prefix, - attention_prefix=attention_prefix, - ) - ) - - # mid block - - # mid block attentions - - # There is a single hardcoded attention block in the middle of the VQ-diffusion decoder - diffusers_attention_prefix = "decoder.mid_block.attentions.0" - attention_prefix = "decoder.mid.attn_1" - diffusers_checkpoint.update( - vqvae_attention_to_diffusers_checkpoint( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # mid block resnets - - for diffusers_resnet_idx, resnet in enumerate(model.encoder.mid_block.resnets): - diffusers_resnet_prefix = f"decoder.mid_block.resnets.{diffusers_resnet_idx}" - - # the hardcoded prefixes to `block_` are 1 and 2 - orig_resnet_idx = diffusers_resnet_idx + 1 - # There are two hardcoded resnets in the middle of the VQ-diffusion decoder - resnet_prefix = f"decoder.mid.block_{orig_resnet_idx}" - - diffusers_checkpoint.update( - vqvae_resnet_to_diffusers_checkpoint( - resnet, checkpoint, diffusers_resnet_prefix=diffusers_resnet_prefix, resnet_prefix=resnet_prefix - ) - ) - - diffusers_checkpoint.update( - { - # conv_norm_out - "decoder.conv_norm_out.weight": checkpoint["decoder.norm_out.weight"], - "decoder.conv_norm_out.bias": checkpoint["decoder.norm_out.bias"], - # conv_out - "decoder.conv_out.weight": checkpoint["decoder.conv_out.weight"], - "decoder.conv_out.bias": checkpoint["decoder.conv_out.bias"], - } - ) - - return diffusers_checkpoint - - -def vqvae_resnet_to_diffusers_checkpoint(resnet, checkpoint, *, diffusers_resnet_prefix, resnet_prefix): - rv = { - # norm1 - f"{diffusers_resnet_prefix}.norm1.weight": checkpoint[f"{resnet_prefix}.norm1.weight"], - f"{diffusers_resnet_prefix}.norm1.bias": checkpoint[f"{resnet_prefix}.norm1.bias"], - # conv1 - f"{diffusers_resnet_prefix}.conv1.weight": checkpoint[f"{resnet_prefix}.conv1.weight"], - f"{diffusers_resnet_prefix}.conv1.bias": checkpoint[f"{resnet_prefix}.conv1.bias"], - # norm2 - f"{diffusers_resnet_prefix}.norm2.weight": checkpoint[f"{resnet_prefix}.norm2.weight"], - f"{diffusers_resnet_prefix}.norm2.bias": checkpoint[f"{resnet_prefix}.norm2.bias"], - # conv2 - f"{diffusers_resnet_prefix}.conv2.weight": checkpoint[f"{resnet_prefix}.conv2.weight"], - f"{diffusers_resnet_prefix}.conv2.bias": checkpoint[f"{resnet_prefix}.conv2.bias"], - } - - if resnet.conv_shortcut is not None: - rv.update( - { - f"{diffusers_resnet_prefix}.conv_shortcut.weight": checkpoint[f"{resnet_prefix}.nin_shortcut.weight"], - f"{diffusers_resnet_prefix}.conv_shortcut.bias": checkpoint[f"{resnet_prefix}.nin_shortcut.bias"], - } - ) - - return rv - - -def vqvae_attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix): - return { - # group_norm - f"{diffusers_attention_prefix}.group_norm.weight": checkpoint[f"{attention_prefix}.norm.weight"], - f"{diffusers_attention_prefix}.group_norm.bias": checkpoint[f"{attention_prefix}.norm.bias"], - # query - f"{diffusers_attention_prefix}.query.weight": checkpoint[f"{attention_prefix}.q.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.query.bias": checkpoint[f"{attention_prefix}.q.bias"], - # key - f"{diffusers_attention_prefix}.key.weight": checkpoint[f"{attention_prefix}.k.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.key.bias": checkpoint[f"{attention_prefix}.k.bias"], - # value - f"{diffusers_attention_prefix}.value.weight": checkpoint[f"{attention_prefix}.v.weight"][:, :, 0, 0], - f"{diffusers_attention_prefix}.value.bias": checkpoint[f"{attention_prefix}.v.bias"], - # proj_attn - f"{diffusers_attention_prefix}.proj_attn.weight": checkpoint[f"{attention_prefix}.proj_out.weight"][ - :, :, 0, 0 - ], - f"{diffusers_attention_prefix}.proj_attn.bias": checkpoint[f"{attention_prefix}.proj_out.bias"], - } - - -# done vqvae checkpoint - -# transformer model - -PORTED_DIFFUSIONS = ["image_synthesis.modeling.transformers.diffusion_transformer.DiffusionTransformer"] -PORTED_TRANSFORMERS = ["image_synthesis.modeling.transformers.transformer_utils.Text2ImageTransformer"] -PORTED_CONTENT_EMBEDDINGS = ["image_synthesis.modeling.embeddings.dalle_mask_image_embedding.DalleMaskImageEmbedding"] - - -def transformer_model_from_original_config( - original_diffusion_config, original_transformer_config, original_content_embedding_config -): - assert original_diffusion_config["target"] in PORTED_DIFFUSIONS, ( - f"{original_diffusion_config['target']} has not yet been ported to diffusers." - ) - assert original_transformer_config["target"] in PORTED_TRANSFORMERS, ( - f"{original_transformer_config['target']} has not yet been ported to diffusers." - ) - assert original_content_embedding_config["target"] in PORTED_CONTENT_EMBEDDINGS, ( - f"{original_content_embedding_config['target']} has not yet been ported to diffusers." - ) - - original_diffusion_config = original_diffusion_config["params"] - original_transformer_config = original_transformer_config["params"] - original_content_embedding_config = original_content_embedding_config["params"] - - inner_dim = original_transformer_config["n_embd"] - - n_heads = original_transformer_config["n_head"] - - # VQ-Diffusion gives dimension of the multi-headed attention layers as the - # number of attention heads times the sequence length (the dimension) of a - # single head. We want to specify our attention blocks with those values - # specified separately - assert inner_dim % n_heads == 0 - d_head = inner_dim // n_heads - - depth = original_transformer_config["n_layer"] - context_dim = original_transformer_config["condition_dim"] - - num_embed = original_content_embedding_config["num_embed"] - # the number of embeddings in the transformer includes the mask embedding. - # the content embedding (the vqvae) does not include the mask embedding. - num_embed = num_embed + 1 - - height = original_transformer_config["content_spatial_size"][0] - width = original_transformer_config["content_spatial_size"][1] - - assert width == height, "width has to be equal to height" - dropout = original_transformer_config["resid_pdrop"] - num_embeds_ada_norm = original_diffusion_config["diffusion_step"] - - model_kwargs = { - "attention_bias": True, - "cross_attention_dim": context_dim, - "attention_head_dim": d_head, - "num_layers": depth, - "dropout": dropout, - "num_attention_heads": n_heads, - "num_vector_embeds": num_embed, - "num_embeds_ada_norm": num_embeds_ada_norm, - "norm_num_groups": 32, - "sample_size": width, - "activation_fn": "geglu-approximate", - } - - model = Transformer2DModel(**model_kwargs) - return model - - -# done transformer model - -# transformer checkpoint - - -def transformer_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): - diffusers_checkpoint = {} - - transformer_prefix = "transformer.transformer" - - diffusers_latent_image_embedding_prefix = "latent_image_embedding" - latent_image_embedding_prefix = f"{transformer_prefix}.content_emb" - - # DalleMaskImageEmbedding - diffusers_checkpoint.update( - { - f"{diffusers_latent_image_embedding_prefix}.emb.weight": checkpoint[ - f"{latent_image_embedding_prefix}.emb.weight" - ], - f"{diffusers_latent_image_embedding_prefix}.height_emb.weight": checkpoint[ - f"{latent_image_embedding_prefix}.height_emb.weight" - ], - f"{diffusers_latent_image_embedding_prefix}.width_emb.weight": checkpoint[ - f"{latent_image_embedding_prefix}.width_emb.weight" - ], - } - ) - - # transformer blocks - for transformer_block_idx, transformer_block in enumerate(model.transformer_blocks): - diffusers_transformer_block_prefix = f"transformer_blocks.{transformer_block_idx}" - transformer_block_prefix = f"{transformer_prefix}.blocks.{transformer_block_idx}" - - # ada norm block - diffusers_ada_norm_prefix = f"{diffusers_transformer_block_prefix}.norm1" - ada_norm_prefix = f"{transformer_block_prefix}.ln1" - - diffusers_checkpoint.update( - transformer_ada_norm_to_diffusers_checkpoint( - checkpoint, diffusers_ada_norm_prefix=diffusers_ada_norm_prefix, ada_norm_prefix=ada_norm_prefix - ) - ) - - # attention block - diffusers_attention_prefix = f"{diffusers_transformer_block_prefix}.attn1" - attention_prefix = f"{transformer_block_prefix}.attn1" - - diffusers_checkpoint.update( - transformer_attention_to_diffusers_checkpoint( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # ada norm block - diffusers_ada_norm_prefix = f"{diffusers_transformer_block_prefix}.norm2" - ada_norm_prefix = f"{transformer_block_prefix}.ln1_1" - - diffusers_checkpoint.update( - transformer_ada_norm_to_diffusers_checkpoint( - checkpoint, diffusers_ada_norm_prefix=diffusers_ada_norm_prefix, ada_norm_prefix=ada_norm_prefix - ) - ) - - # attention block - diffusers_attention_prefix = f"{diffusers_transformer_block_prefix}.attn2" - attention_prefix = f"{transformer_block_prefix}.attn2" - - diffusers_checkpoint.update( - transformer_attention_to_diffusers_checkpoint( - checkpoint, diffusers_attention_prefix=diffusers_attention_prefix, attention_prefix=attention_prefix - ) - ) - - # norm block - diffusers_norm_block_prefix = f"{diffusers_transformer_block_prefix}.norm3" - norm_block_prefix = f"{transformer_block_prefix}.ln2" - - diffusers_checkpoint.update( - { - f"{diffusers_norm_block_prefix}.weight": checkpoint[f"{norm_block_prefix}.weight"], - f"{diffusers_norm_block_prefix}.bias": checkpoint[f"{norm_block_prefix}.bias"], - } - ) - - # feedforward block - diffusers_feedforward_prefix = f"{diffusers_transformer_block_prefix}.ff" - feedforward_prefix = f"{transformer_block_prefix}.mlp" - - diffusers_checkpoint.update( - transformer_feedforward_to_diffusers_checkpoint( - checkpoint, - diffusers_feedforward_prefix=diffusers_feedforward_prefix, - feedforward_prefix=feedforward_prefix, - ) - ) - - # to logits - - diffusers_norm_out_prefix = "norm_out" - norm_out_prefix = f"{transformer_prefix}.to_logits.0" - - diffusers_checkpoint.update( - { - f"{diffusers_norm_out_prefix}.weight": checkpoint[f"{norm_out_prefix}.weight"], - f"{diffusers_norm_out_prefix}.bias": checkpoint[f"{norm_out_prefix}.bias"], - } - ) - - diffusers_out_prefix = "out" - out_prefix = f"{transformer_prefix}.to_logits.1" - - diffusers_checkpoint.update( - { - f"{diffusers_out_prefix}.weight": checkpoint[f"{out_prefix}.weight"], - f"{diffusers_out_prefix}.bias": checkpoint[f"{out_prefix}.bias"], - } - ) - - return diffusers_checkpoint - - -def transformer_ada_norm_to_diffusers_checkpoint(checkpoint, *, diffusers_ada_norm_prefix, ada_norm_prefix): - return { - f"{diffusers_ada_norm_prefix}.emb.weight": checkpoint[f"{ada_norm_prefix}.emb.weight"], - f"{diffusers_ada_norm_prefix}.linear.weight": checkpoint[f"{ada_norm_prefix}.linear.weight"], - f"{diffusers_ada_norm_prefix}.linear.bias": checkpoint[f"{ada_norm_prefix}.linear.bias"], - } - - -def transformer_attention_to_diffusers_checkpoint(checkpoint, *, diffusers_attention_prefix, attention_prefix): - return { - # key - f"{diffusers_attention_prefix}.to_k.weight": checkpoint[f"{attention_prefix}.key.weight"], - f"{diffusers_attention_prefix}.to_k.bias": checkpoint[f"{attention_prefix}.key.bias"], - # query - f"{diffusers_attention_prefix}.to_q.weight": checkpoint[f"{attention_prefix}.query.weight"], - f"{diffusers_attention_prefix}.to_q.bias": checkpoint[f"{attention_prefix}.query.bias"], - # value - f"{diffusers_attention_prefix}.to_v.weight": checkpoint[f"{attention_prefix}.value.weight"], - f"{diffusers_attention_prefix}.to_v.bias": checkpoint[f"{attention_prefix}.value.bias"], - # linear out - f"{diffusers_attention_prefix}.to_out.0.weight": checkpoint[f"{attention_prefix}.proj.weight"], - f"{diffusers_attention_prefix}.to_out.0.bias": checkpoint[f"{attention_prefix}.proj.bias"], - } - - -def transformer_feedforward_to_diffusers_checkpoint(checkpoint, *, diffusers_feedforward_prefix, feedforward_prefix): - return { - f"{diffusers_feedforward_prefix}.net.0.proj.weight": checkpoint[f"{feedforward_prefix}.0.weight"], - f"{diffusers_feedforward_prefix}.net.0.proj.bias": checkpoint[f"{feedforward_prefix}.0.bias"], - f"{diffusers_feedforward_prefix}.net.2.weight": checkpoint[f"{feedforward_prefix}.2.weight"], - f"{diffusers_feedforward_prefix}.net.2.bias": checkpoint[f"{feedforward_prefix}.2.bias"], - } - - -# done transformer checkpoint - - -def read_config_file(filename): - # The yaml file contains annotations that certain values should - # loaded as tuples. - with open(filename) as f: - original_config = yaml.load(f, FullLoader) - - return original_config - - -# We take separate arguments for the vqvae because the ITHQ vqvae config file -# is separate from the config file for the rest of the model. -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--vqvae_checkpoint_path", - default=None, - type=str, - required=True, - help="Path to the vqvae checkpoint to convert.", - ) - - parser.add_argument( - "--vqvae_original_config_file", - default=None, - type=str, - required=True, - help="The YAML config file corresponding to the original architecture for the vqvae.", - ) - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - - parser.add_argument( - "--original_config_file", - default=None, - type=str, - required=True, - help="The YAML config file corresponding to the original architecture.", - ) - - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - - parser.add_argument( - "--checkpoint_load_device", - default="cpu", - type=str, - required=False, - help="The device passed to `map_location` when loading checkpoints.", - ) - - # See link for how ema weights are always selected - # https://github.com/microsoft/VQ-Diffusion/blob/3c98e77f721db7c787b76304fa2c96a36c7b00af/inference_VQ_Diffusion.py#L65 - parser.add_argument( - "--no_use_ema", - action="store_true", - required=False, - help=( - "Set to not use the ema weights from the original VQ-Diffusion checkpoint. You probably do not want to set" - " it as the original VQ-Diffusion always uses the ema weights when loading models." - ), - ) - - args = parser.parse_args() - - use_ema = not args.no_use_ema - - print(f"loading checkpoints to {args.checkpoint_load_device}") - - checkpoint_map_location = torch.device(args.checkpoint_load_device) - - # vqvae_model - - print(f"loading vqvae, config: {args.vqvae_original_config_file}, checkpoint: {args.vqvae_checkpoint_path}") - - vqvae_original_config = read_config_file(args.vqvae_original_config_file).model - vqvae_checkpoint = torch.load(args.vqvae_checkpoint_path, map_location=checkpoint_map_location)["model"] - - with init_empty_weights(): - vqvae_model = vqvae_model_from_original_config(vqvae_original_config) - - vqvae_diffusers_checkpoint = vqvae_original_checkpoint_to_diffusers_checkpoint(vqvae_model, vqvae_checkpoint) - - with tempfile.NamedTemporaryFile() as vqvae_diffusers_checkpoint_file: - torch.save(vqvae_diffusers_checkpoint, vqvae_diffusers_checkpoint_file.name) - del vqvae_diffusers_checkpoint - del vqvae_checkpoint - load_checkpoint_and_dispatch(vqvae_model, vqvae_diffusers_checkpoint_file.name, device_map="auto") - - print("done loading vqvae") - - # done vqvae_model - - # transformer_model - - print( - f"loading transformer, config: {args.original_config_file}, checkpoint: {args.checkpoint_path}, use ema:" - f" {use_ema}" - ) - - original_config = read_config_file(args.original_config_file).model - - diffusion_config = original_config["params"]["diffusion_config"] - transformer_config = original_config["params"]["diffusion_config"]["params"]["transformer_config"] - content_embedding_config = original_config["params"]["diffusion_config"]["params"]["content_emb_config"] - - pre_checkpoint = torch.load(args.checkpoint_path, map_location=checkpoint_map_location) - - if use_ema: - if "ema" in pre_checkpoint: - checkpoint = {} - for k, v in pre_checkpoint["model"].items(): - checkpoint[k] = v - - for k, v in pre_checkpoint["ema"].items(): - # The ema weights are only used on the transformer. To mimic their key as if they came - # from the state_dict for the top level model, we prefix with an additional "transformer." - # See the source linked in the args.use_ema config for more information. - checkpoint[f"transformer.{k}"] = v - else: - print("attempted to load ema weights but no ema weights are specified in the loaded checkpoint.") - checkpoint = pre_checkpoint["model"] - else: - checkpoint = pre_checkpoint["model"] - - del pre_checkpoint - - with init_empty_weights(): - transformer_model = transformer_model_from_original_config( - diffusion_config, transformer_config, content_embedding_config - ) - - diffusers_transformer_checkpoint = transformer_original_checkpoint_to_diffusers_checkpoint( - transformer_model, checkpoint - ) - - # classifier free sampling embeddings interlude - - # The learned embeddings are stored on the transformer in the original VQ-diffusion. We store them on a separate - # model, so we pull them off the checkpoint before the checkpoint is deleted. - - learnable_classifier_free_sampling_embeddings = diffusion_config["params"].learnable_cf - - if learnable_classifier_free_sampling_embeddings: - learned_classifier_free_sampling_embeddings_embeddings = checkpoint["transformer.empty_text_embed"] - else: - learned_classifier_free_sampling_embeddings_embeddings = None - - # done classifier free sampling embeddings interlude - - with tempfile.NamedTemporaryFile() as diffusers_transformer_checkpoint_file: - torch.save(diffusers_transformer_checkpoint, diffusers_transformer_checkpoint_file.name) - del diffusers_transformer_checkpoint - del checkpoint - load_checkpoint_and_dispatch(transformer_model, diffusers_transformer_checkpoint_file.name, device_map="auto") - - print("done loading transformer") - - # done transformer_model - - # text encoder - - print("loading CLIP text encoder") - - clip_name = "openai/clip-vit-base-patch32" - - # The original VQ-Diffusion specifies the pad value by the int used in the - # returned tokens. Each model uses `0` as the pad value. The transformers clip api - # specifies the pad value via the token before it has been tokenized. The `!` pad - # token is the same as padding with the `0` pad value. - pad_token = "!" - - tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto") - - assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0 - - text_encoder_model = CLIPTextModel.from_pretrained( - clip_name, - # `CLIPTextModel` does not support device_map="auto" - # device_map="auto" - ) - - print("done loading CLIP text encoder") - - # done text encoder - - # scheduler - - scheduler_model = VQDiffusionScheduler( - # the scheduler has the same number of embeddings as the transformer - num_vec_classes=transformer_model.num_vector_embeds - ) - - # done scheduler - - # learned classifier free sampling embeddings - - with init_empty_weights(): - learned_classifier_free_sampling_embeddings_model = LearnedClassifierFreeSamplingEmbeddings( - learnable_classifier_free_sampling_embeddings, - hidden_size=text_encoder_model.config.hidden_size, - length=tokenizer_model.model_max_length, - ) - - learned_classifier_free_sampling_checkpoint = { - "embeddings": learned_classifier_free_sampling_embeddings_embeddings.float() - } - - with tempfile.NamedTemporaryFile() as learned_classifier_free_sampling_checkpoint_file: - torch.save(learned_classifier_free_sampling_checkpoint, learned_classifier_free_sampling_checkpoint_file.name) - del learned_classifier_free_sampling_checkpoint - del learned_classifier_free_sampling_embeddings_embeddings - load_checkpoint_and_dispatch( - learned_classifier_free_sampling_embeddings_model, - learned_classifier_free_sampling_checkpoint_file.name, - device_map="auto", - ) - - # done learned classifier free sampling embeddings - - print(f"saving VQ diffusion model, path: {args.dump_path}") - - pipe = VQDiffusionPipeline( - vqvae=vqvae_model, - transformer=transformer_model, - tokenizer=tokenizer_model, - text_encoder=text_encoder_model, - learned_classifier_free_sampling_embeddings=learned_classifier_free_sampling_embeddings_model, - scheduler=scheduler_model, - ) - pipe.save_pretrained(args.dump_path) - - print("done writing VQ diffusion model") diff --git a/scripts/convert_wan_to_diffusers.py b/scripts/convert_wan_to_diffusers.py deleted file mode 100644 index eca1af61ae69..000000000000 --- a/scripts/convert_wan_to_diffusers.py +++ /dev/null @@ -1,1287 +0,0 @@ -import argparse -import pathlib -from typing import Any, Dict, Tuple - -import torch -from accelerate import init_empty_weights -from huggingface_hub import hf_hub_download, snapshot_download -from safetensors.torch import load_file -from transformers import ( - AutoProcessor, - AutoTokenizer, - CLIPImageProcessor, - CLIPVisionModel, - CLIPVisionModelWithProjection, - UMT5EncoderModel, -) - -from diffusers import ( - AutoencoderKLWan, - UniPCMultistepScheduler, - WanAnimatePipeline, - WanAnimateTransformer3DModel, - WanImageToVideoPipeline, - WanPipeline, - WanTransformer3DModel, - WanVACEPipeline, - WanVACETransformer3DModel, -) - - -TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "ffn.0": "ffn.net.0.proj", - "ffn.2": "ffn.net.2", - # Hack to swap the layer names - # The original model calls the norms in following order: norm1, norm3, norm2 - # We convert it to: norm1, norm2, norm3 - "norm2": "norm__placeholder", - "norm3": "norm2", - "norm__placeholder": "norm3", - # For the I2V model - "img_emb.proj.0": "condition_embedder.image_embedder.norm1", - "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", - "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", - "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # for the FLF2V model - "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", - # Add attention component mappings - "self_attn.q": "attn1.to_q", - "self_attn.k": "attn1.to_k", - "self_attn.v": "attn1.to_v", - "self_attn.o": "attn1.to_out.0", - "self_attn.norm_q": "attn1.norm_q", - "self_attn.norm_k": "attn1.norm_k", - "cross_attn.q": "attn2.to_q", - "cross_attn.k": "attn2.to_k", - "cross_attn.v": "attn2.to_v", - "cross_attn.o": "attn2.to_out.0", - "cross_attn.norm_q": "attn2.norm_q", - "cross_attn.norm_k": "attn2.norm_k", - "attn2.to_k_img": "attn2.add_k_proj", - "attn2.to_v_img": "attn2.add_v_proj", - "attn2.norm_k_img": "attn2.norm_added_k", -} - -VACE_TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "ffn.0": "ffn.net.0.proj", - "ffn.2": "ffn.net.2", - # Hack to swap the layer names - # The original model calls the norms in following order: norm1, norm3, norm2 - # We convert it to: norm1, norm2, norm3 - "norm2": "norm__placeholder", - "norm3": "norm2", - "norm__placeholder": "norm3", - # # For the I2V model - # "img_emb.proj.0": "condition_embedder.image_embedder.norm1", - # "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", - # "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", - # "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # # for the FLF2V model - # "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed", - # Add attention component mappings - "self_attn.q": "attn1.to_q", - "self_attn.k": "attn1.to_k", - "self_attn.v": "attn1.to_v", - "self_attn.o": "attn1.to_out.0", - "self_attn.norm_q": "attn1.norm_q", - "self_attn.norm_k": "attn1.norm_k", - "cross_attn.q": "attn2.to_q", - "cross_attn.k": "attn2.to_k", - "cross_attn.v": "attn2.to_v", - "cross_attn.o": "attn2.to_out.0", - "cross_attn.norm_q": "attn2.norm_q", - "cross_attn.norm_k": "attn2.norm_k", - "attn2.to_k_img": "attn2.add_k_proj", - "attn2.to_v_img": "attn2.add_v_proj", - "attn2.norm_k_img": "attn2.norm_added_k", - "before_proj": "proj_in", - "after_proj": "proj_out", -} - -ANIMATE_TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "ffn.0": "ffn.net.0.proj", - "ffn.2": "ffn.net.2", - # Hack to swap the layer names - # The original model calls the norms in following order: norm1, norm3, norm2 - # We convert it to: norm1, norm2, norm3 - "norm2": "norm__placeholder", - "norm3": "norm2", - "norm__placeholder": "norm3", - "img_emb.proj.0": "condition_embedder.image_embedder.norm1", - "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", - "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", - "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # Add attention component mappings - "self_attn.q": "attn1.to_q", - "self_attn.k": "attn1.to_k", - "self_attn.v": "attn1.to_v", - "self_attn.o": "attn1.to_out.0", - "self_attn.norm_q": "attn1.norm_q", - "self_attn.norm_k": "attn1.norm_k", - "cross_attn.q": "attn2.to_q", - "cross_attn.k": "attn2.to_k", - "cross_attn.v": "attn2.to_v", - "cross_attn.o": "attn2.to_out.0", - "cross_attn.norm_q": "attn2.norm_q", - "cross_attn.norm_k": "attn2.norm_k", - "cross_attn.k_img": "attn2.to_k_img", - "cross_attn.v_img": "attn2.to_v_img", - "cross_attn.norm_k_img": "attn2.norm_k_img", - # After cross_attn -> attn2 rename, we need to rename the img keys - "attn2.to_k_img": "attn2.add_k_proj", - "attn2.to_v_img": "attn2.add_v_proj", - "attn2.norm_k_img": "attn2.norm_added_k", - # Wan Animate-specific mappings (motion encoder, face encoder, face adapter) - # Motion encoder mappings - # The name mapping is complicated for the convolutional part so we handle that in its own function - "motion_encoder.enc.fc": "motion_encoder.motion_network", - "motion_encoder.dec.direction.weight": "motion_encoder.motion_synthesis_weight", - # Face encoder mappings - CausalConv1d has a .conv submodule that we need to flatten - "face_encoder.conv1_local.conv": "face_encoder.conv1_local", - "face_encoder.conv2.conv": "face_encoder.conv2", - "face_encoder.conv3.conv": "face_encoder.conv3", - # Face adapter mappings are handled in a separate function -} - - -# TODO: Verify this and simplify if possible. -def convert_animate_motion_encoder_weights(key: str, state_dict: Dict[str, Any], final_conv_idx: int = 8) -> None: - """ - Convert all motion encoder weights for Animate model. - - In the original model: - - All Linear layers in fc use EqualLinear - - All Conv2d layers in convs use EqualConv2d (except blur_conv which is initialized separately) - - Blur kernels are stored as buffers in Sequential modules - - ConvLayer is nn.Sequential with indices: [Blur (optional), EqualConv2d, FusedLeakyReLU (optional)] - - Conversion strategy: - 1. Drop .kernel buffers (blur kernels) - 2. Rename sequential indices to named components (e.g., 0 -> conv2d, 1 -> bias_leaky_relu) - """ - # Skip if not a weight, bias, or kernel - if ".weight" not in key and ".bias" not in key and ".kernel" not in key: - return - - # Handle Blur kernel buffers from original implementation. - # After renaming, these appear under: motion_encoder.res_blocks.*.conv{2,skip}.blur_kernel - # Diffusers constructs blur kernels as a non-persistent buffer so we must drop these keys - if ".kernel" in key and "motion_encoder" in key: - # Remove unexpected blur kernel buffers to avoid strict load errors - state_dict.pop(key, None) - return - - # Rename Sequential indices to named components in ConvLayer and ResBlock - if ".enc.net_app.convs." in key and (".weight" in key or ".bias" in key): - parts = key.split(".") - - # Find the sequential index (digit) after convs or after conv1/conv2/skip - # Examples: - # - enc.net_app.convs.0.0.weight -> conv_in.weight (initial conv layer weight) - # - enc.net_app.convs.0.1.bias -> conv_in.act_fn.bias (initial conv layer bias) - # - enc.net_app.convs.{n:1-7}.conv1.0.weight -> res_blocks.{(n-1):0-6}.conv1.weight (conv1 weight) - # - e.g. enc.net_app.convs.1.conv1.0.weight -> res_blocks.0.conv1.weight - # - enc.net_app.convs.{n:1-7}.conv1.1.bias -> res_blocks.{(n-1):0-6}.conv1.act_fn.bias (conv1 bias) - # - e.g. enc.net_app.convs.1.conv1.1.bias -> res_blocks.0.conv1.act_fn.bias - # - enc.net_app.convs.{n:1-7}.conv2.1.weight -> res_blocks.{(n-1):0-6}.conv2.weight (conv2 weight) - # - enc.net_app.convs.1.conv2.2.bias -> res_blocks.0.conv2.act_fn.bias (conv2 bias) - # - enc.net_app.convs.{n:1-7}.skip.1.weight -> res_blocks.{(n-1):0-6}.conv_skip.weight (skip conv weight) - # - enc.net_app.convs.8 -> conv_out (final conv layer) - - convs_idx = parts.index("convs") if "convs" in parts else -1 - if convs_idx >= 0 and len(parts) - convs_idx >= 2: - bias = False - # The nn.Sequential index will always follow convs - sequential_idx = int(parts[convs_idx + 1]) - if sequential_idx == 0: - if key.endswith(".weight"): - new_key = "motion_encoder.conv_in.weight" - elif key.endswith(".bias"): - new_key = "motion_encoder.conv_in.act_fn.bias" - bias = True - elif sequential_idx == final_conv_idx: - if key.endswith(".weight"): - new_key = "motion_encoder.conv_out.weight" - else: - # Intermediate .convs. layers, which get mapped to .res_blocks. - prefix = "motion_encoder.res_blocks." - - layer_name = parts[convs_idx + 2] - if layer_name == "skip": - layer_name = "conv_skip" - - if key.endswith(".weight"): - param_name = "weight" - elif key.endswith(".bias"): - param_name = "act_fn.bias" - bias = True - - suffix_parts = [str(sequential_idx - 1), layer_name, param_name] - suffix = ".".join(suffix_parts) - new_key = prefix + suffix - - param = state_dict.pop(key) - if bias: - param = param.squeeze() - state_dict[new_key] = param - return - return - return - - -def convert_animate_face_adapter_weights(key: str, state_dict: Dict[str, Any]) -> None: - """ - Convert face adapter weights for the Animate model. - - The original model uses a fused KV projection but the diffusers models uses separate K and V projections. - """ - # Skip if not a weight or bias - if ".weight" not in key and ".bias" not in key: - return - - prefix = "face_adapter." - if ".fuser_blocks." in key: - parts = key.split(".") - - module_list_idx = parts.index("fuser_blocks") if "fuser_blocks" in parts else -1 - if module_list_idx >= 0 and (len(parts) - 1) - module_list_idx == 3: - block_idx = parts[module_list_idx + 1] - layer_name = parts[module_list_idx + 2] - param_name = parts[module_list_idx + 3] - - if layer_name == "linear1_kv": - layer_name_k = "to_k" - layer_name_v = "to_v" - - suffix_k = ".".join([block_idx, layer_name_k, param_name]) - suffix_v = ".".join([block_idx, layer_name_v, param_name]) - new_key_k = prefix + suffix_k - new_key_v = prefix + suffix_v - - kv_proj = state_dict.pop(key) - k_proj, v_proj = torch.chunk(kv_proj, 2, dim=0) - state_dict[new_key_k] = k_proj - state_dict[new_key_v] = v_proj - return - else: - if layer_name == "q_norm": - new_layer_name = "norm_q" - elif layer_name == "k_norm": - new_layer_name = "norm_k" - elif layer_name == "linear1_q": - new_layer_name = "to_q" - elif layer_name == "linear2": - new_layer_name = "to_out" - - suffix_parts = [block_idx, new_layer_name, param_name] - suffix = ".".join(suffix_parts) - new_key = prefix + suffix - state_dict[new_key] = state_dict.pop(key) - return - return - - -TRANSFORMER_SPECIAL_KEYS_REMAP = {} -VACE_TRANSFORMER_SPECIAL_KEYS_REMAP = {} -ANIMATE_TRANSFORMER_SPECIAL_KEYS_REMAP = { - "motion_encoder": convert_animate_motion_encoder_weights, - "face_adapter": convert_animate_face_adapter_weights, -} - - -def update_state_dict_(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - -def load_sharded_safetensors(dir: pathlib.Path): - file_paths = list(dir.glob("diffusion_pytorch_model*.safetensors")) - state_dict = {} - for path in file_paths: - state_dict.update(load_file(path)) - return state_dict - - -def get_transformer_config(model_type: str) -> Tuple[Dict[str, Any], ...]: - if model_type == "Wan-T2V-1.3B": - config = { - "model_id": "StevenZhang/Wan2.1-T2V-1.3B-Diff", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 8960, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 12, - "num_layers": 30, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-T2V-14B": - config = { - "model_id": "StevenZhang/Wan2.1-T2V-14B-Diff", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-I2V-14B-480p": - config = { - "model_id": "StevenZhang/Wan2.1-I2V-14B-480P-Diff", - "diffusers_config": { - "image_dim": 1280, - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-I2V-14B-720p": - config = { - "model_id": "StevenZhang/Wan2.1-I2V-14B-720P-Diff", - "diffusers_config": { - "image_dim": 1280, - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-FLF2V-14B-720P": - config = { - "model_id": "ypyp/Wan2.1-FLF2V-14B-720P", # This is just a placeholder - "diffusers_config": { - "image_dim": 1280, - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "rope_max_seq_len": 1024, - "pos_embed_seq_len": 257 * 2, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-VACE-1.3B": - config = { - "model_id": "Wan-AI/Wan2.1-VACE-1.3B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 8960, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 12, - "num_layers": 30, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "vace_layers": [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], - "vace_in_channels": 96, - }, - } - RENAME_DICT = VACE_TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = VACE_TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan-VACE-14B": - config = { - "model_id": "Wan-AI/Wan2.1-VACE-14B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "vace_layers": [0, 5, 10, 15, 20, 25, 30, 35], - "vace_in_channels": 96, - }, - } - RENAME_DICT = VACE_TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = VACE_TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan2.2-VACE-Fun-14B": - config = { - "model_id": "alibaba-pai/Wan2.2-VACE-Fun-A14B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "vace_layers": [0, 5, 10, 15, 20, 25, 30, 35], - "vace_in_channels": 96, - }, - } - RENAME_DICT = VACE_TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = VACE_TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan2.2-I2V-14B-720p": - config = { - "model_id": "Wan-AI/Wan2.2-I2V-A14B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan2.2-T2V-A14B": - config = { - "model_id": "Wan-AI/Wan2.2-T2V-A14B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 16, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan2.2-TI2V-5B": - config = { - "model_id": "Wan-AI/Wan2.2-TI2V-5B", - "diffusers_config": { - "added_kv_proj_dim": None, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 14336, - "freq_dim": 256, - "in_channels": 48, - "num_attention_heads": 24, - "num_layers": 30, - "out_channels": 48, - "patch_size": [1, 2, 2], - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - }, - } - RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP - elif model_type == "Wan2.2-Animate-14B": - config = { - "model_id": "Wan-AI/Wan2.2-Animate-14B", - "diffusers_config": { - "image_dim": 1280, - "added_kv_proj_dim": 5120, - "attention_head_dim": 128, - "cross_attn_norm": True, - "eps": 1e-06, - "ffn_dim": 13824, - "freq_dim": 256, - "in_channels": 36, - "num_attention_heads": 40, - "num_layers": 40, - "out_channels": 16, - "patch_size": (1, 2, 2), - "qk_norm": "rms_norm_across_heads", - "text_dim": 4096, - "rope_max_seq_len": 1024, - "pos_embed_seq_len": None, - "motion_encoder_size": 512, # Start of Wan Animate-specific configs - "motion_style_dim": 512, - "motion_dim": 20, - "motion_encoder_dim": 512, - "face_encoder_hidden_dim": 1024, - "face_encoder_num_heads": 4, - "inject_face_latents_blocks": 5, - }, - } - RENAME_DICT = ANIMATE_TRANSFORMER_KEYS_RENAME_DICT - SPECIAL_KEYS_REMAP = ANIMATE_TRANSFORMER_SPECIAL_KEYS_REMAP - return config, RENAME_DICT, SPECIAL_KEYS_REMAP - - -def convert_transformer(model_type: str, stage: str = None): - config, RENAME_DICT, SPECIAL_KEYS_REMAP = get_transformer_config(model_type) - - diffusers_config = config["diffusers_config"] - model_id = config["model_id"] - model_dir = pathlib.Path(snapshot_download(model_id, repo_type="model")) - - if stage is not None: - model_dir = model_dir / stage - - original_state_dict = load_sharded_safetensors(model_dir) - - with init_empty_weights(): - if "Animate" in model_type: - transformer = WanAnimateTransformer3DModel.from_config(diffusers_config) - elif "VACE" in model_type: - transformer = WanVACETransformer3DModel.from_config(diffusers_config) - else: - transformer = WanTransformer3DModel.from_config(diffusers_config) - - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - # Load state dict into the meta model, which will materialize the tensors - transformer.load_state_dict(original_state_dict, strict=True, assign=True) - - # Move to CPU to ensure all tensors are materialized - transformer = transformer.to("cpu") - - return transformer - - -def convert_vae(): - vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") - old_state_dict = torch.load(vae_ckpt_path, weights_only=True) - new_state_dict = {} - - # Create mappings for specific components - middle_key_mapping = { - # Encoder middle block - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - # Create a mapping for attention blocks - attention_mapping = { - # Encoder middle attention - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - # Create a mapping for the head components - head_mapping = { - # Encoder head - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - # Create a mapping for the quant components - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - # Process each key in the state dict - for key, value in old_state_dict.items(): - # Handle middle block keys using the mapping - if key in middle_key_mapping: - new_key = middle_key_mapping[key] - new_state_dict[new_key] = value - # Handle attention blocks using the mapping - elif key in attention_mapping: - new_key = attention_mapping[key] - new_state_dict[new_key] = value - # Handle head keys using the mapping - elif key in head_mapping: - new_key = head_mapping[key] - new_state_dict[new_key] = value - # Handle quant keys using the mapping - elif key in quant_mapping: - new_key = quant_mapping[key] - new_state_dict[new_key] = value - # Handle encoder conv1 - elif key == "encoder.conv1.weight": - new_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - new_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 - elif key == "decoder.conv1.weight": - new_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - new_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples - elif key.startswith("encoder.downsamples."): - # Convert to down_blocks - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Convert residual block naming but keep the original structure - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - - new_state_dict[new_key] = value - - # Handle decoder upsamples - elif key.startswith("decoder.upsamples."): - # Convert to up_blocks - parts = key.split(".") - block_idx = int(parts[2]) - - # Group residual blocks - if "residual" in key: - if block_idx in [0, 1, 2]: - new_block_idx = 0 - resnet_idx = block_idx - elif block_idx in [4, 5, 6]: - new_block_idx = 1 - resnet_idx = block_idx - 4 - elif block_idx in [8, 9, 10]: - new_block_idx = 2 - resnet_idx = block_idx - 8 - elif block_idx in [12, 13, 14]: - new_block_idx = 3 - resnet_idx = block_idx - 12 - else: - # Keep as is for other blocks - new_state_dict[key] = value - continue - - # Convert residual block naming - if ".residual.0.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" - elif ".residual.2.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias" - elif ".residual.2.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight" - elif ".residual.3.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma" - elif ".residual.6.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias" - elif ".residual.6.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" - else: - new_key = key - - new_state_dict[new_key] = value - - # Handle shortcut connections - elif ".shortcut." in key: - if block_idx == 4: - new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") - new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - - new_state_dict[new_key] = value - - # Handle upsamplers - elif ".resample." in key or ".time_conv." in key: - if block_idx == 3: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") - elif block_idx == 7: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") - elif block_idx == 11: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - - new_state_dict[new_key] = value - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_state_dict[new_key] = value - else: - # Keep other keys unchanged - new_state_dict[key] = value - - with init_empty_weights(): - vae = AutoencoderKLWan() - vae.load_state_dict(new_state_dict, strict=True, assign=True) - return vae - - -vae22_diffusers_config = { - "base_dim": 160, - "z_dim": 48, - "is_residual": True, - "in_channels": 12, - "out_channels": 12, - "decoder_base_dim": 256, - "scale_factor_temporal": 4, - "scale_factor_spatial": 16, - "patch_size": 2, - "latents_mean": [ - -0.2289, - -0.0052, - -0.1323, - -0.2339, - -0.2799, - 0.0174, - 0.1838, - 0.1557, - -0.1382, - 0.0542, - 0.2813, - 0.0891, - 0.1570, - -0.0098, - 0.0375, - -0.1825, - -0.2246, - -0.1207, - -0.0698, - 0.5109, - 0.2665, - -0.2108, - -0.2158, - 0.2502, - -0.2055, - -0.0322, - 0.1109, - 0.1567, - -0.0729, - 0.0899, - -0.2799, - -0.1230, - -0.0313, - -0.1649, - 0.0117, - 0.0723, - -0.2839, - -0.2083, - -0.0520, - 0.3748, - 0.0152, - 0.1957, - 0.1433, - -0.2944, - 0.3573, - -0.0548, - -0.1681, - -0.0667, - ], - "latents_std": [ - 0.4765, - 1.0364, - 0.4514, - 1.1677, - 0.5313, - 0.4990, - 0.4818, - 0.5013, - 0.8158, - 1.0344, - 0.5894, - 1.0901, - 0.6885, - 0.6165, - 0.8454, - 0.4978, - 0.5759, - 0.3523, - 0.7135, - 0.6804, - 0.5833, - 1.4146, - 0.8986, - 0.5659, - 0.7069, - 0.5338, - 0.4889, - 0.4917, - 0.4069, - 0.4999, - 0.6866, - 0.4093, - 0.5709, - 0.6065, - 0.6415, - 0.4944, - 0.5726, - 1.2042, - 0.5458, - 1.6887, - 0.3971, - 1.0600, - 0.3943, - 0.5537, - 0.5444, - 0.4089, - 0.7468, - 0.7744, - ], - "clip_output": False, -} - - -def convert_vae_22(): - vae_ckpt_path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth") - old_state_dict = torch.load(vae_ckpt_path, weights_only=True) - new_state_dict = {} - - # Create mappings for specific components - middle_key_mapping = { - # Encoder middle block - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - # Create a mapping for attention blocks - attention_mapping = { - # Encoder middle attention - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - # Create a mapping for the head components - head_mapping = { - # Encoder head - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - # Create a mapping for the quant components - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - # Process each key in the state dict - for key, value in old_state_dict.items(): - # Handle middle block keys using the mapping - if key in middle_key_mapping: - new_key = middle_key_mapping[key] - new_state_dict[new_key] = value - # Handle attention blocks using the mapping - elif key in attention_mapping: - new_key = attention_mapping[key] - new_state_dict[new_key] = value - # Handle head keys using the mapping - elif key in head_mapping: - new_key = head_mapping[key] - new_state_dict[new_key] = value - # Handle quant keys using the mapping - elif key in quant_mapping: - new_key = quant_mapping[key] - new_state_dict[new_key] = value - # Handle encoder conv1 - elif key == "encoder.conv1.weight": - new_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - new_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 - elif key == "decoder.conv1.weight": - new_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - new_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples - elif key.startswith("encoder.downsamples."): - # Change encoder.downsamples to encoder.down_blocks - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Handle residual blocks - change downsamples to resnets and rename components - if "residual" in new_key or "shortcut" in new_key: - # Change the second downsamples to resnets - new_key = new_key.replace(".downsamples.", ".resnets.") - - # Rename residual components - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - - # Handle resample blocks - change downsamples to downsampler and remove index - elif "resample" in new_key or "time_conv" in new_key: - # Change the second downsamples to downsampler and remove the index - parts = new_key.split(".") - # Find the pattern: encoder.down_blocks.X.downsamples.Y.resample... - # We want to change it to: encoder.down_blocks.X.downsampler.resample... - if len(parts) >= 4 and parts[3] == "downsamples": - # Remove the index (parts[4]) and change downsamples to downsampler - new_parts = parts[:3] + ["downsampler"] + parts[5:] - new_key = ".".join(new_parts) - - new_state_dict[new_key] = value - - # Handle decoder upsamples - elif key.startswith("decoder.upsamples."): - # Change decoder.upsamples to decoder.up_blocks - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - - # Handle residual blocks - change upsamples to resnets and rename components - if "residual" in new_key or "shortcut" in new_key: - # Change the second upsamples to resnets - new_key = new_key.replace(".upsamples.", ".resnets.") - - # Rename residual components - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - - # Handle resample blocks - change upsamples to upsampler and remove index - elif "resample" in new_key or "time_conv" in new_key: - # Change the second upsamples to upsampler and remove the index - parts = new_key.split(".") - # Find the pattern: encoder.down_blocks.X.downsamples.Y.resample... - # We want to change it to: encoder.down_blocks.X.downsampler.resample... - if len(parts) >= 4 and parts[3] == "upsamples": - # Remove the index (parts[4]) and change upsamples to upsampler - new_parts = parts[:3] + ["upsampler"] + parts[5:] - new_key = ".".join(new_parts) - - new_state_dict[new_key] = value - else: - # Keep other keys unchanged - new_state_dict[key] = value - - with init_empty_weights(): - vae = AutoencoderKLWan(**vae22_diffusers_config) - vae.load_state_dict(new_state_dict, strict=True, assign=True) - return vae - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--model_type", type=str, default=None) - parser.add_argument("--output_path", type=str, required=True) - parser.add_argument("--dtype", default="fp32", choices=["fp32", "fp16", "bf16", "none"]) - return parser.parse_args() - - -DTYPE_MAPPING = { - "fp32": torch.float32, - "fp16": torch.float16, - "bf16": torch.bfloat16, -} - - -if __name__ == "__main__": - args = get_args() - - if "Wan2.2" in args.model_type and "TI2V" not in args.model_type and "Animate" not in args.model_type: - transformer = convert_transformer(args.model_type, stage="high_noise_model") - transformer_2 = convert_transformer(args.model_type, stage="low_noise_model") - else: - transformer = convert_transformer(args.model_type) - transformer_2 = None - - if "Wan2.2" in args.model_type and "TI2V" in args.model_type: - vae = convert_vae_22() - else: - vae = convert_vae() - - text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16) - tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") - if "FLF2V" in args.model_type: - flow_shift = 16.0 - elif "TI2V" in args.model_type or "Animate" in args.model_type: - flow_shift = 5.0 - else: - flow_shift = 3.0 - scheduler = UniPCMultistepScheduler( - prediction_type="flow_prediction", use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=flow_shift - ) - - # If user has specified "none", we keep the original dtypes of the state dict without any conversion - if args.dtype != "none": - dtype = DTYPE_MAPPING[args.dtype] - transformer.to(dtype) - if transformer_2 is not None: - transformer_2.to(dtype) - - if "Wan2.2" and "I2V" in args.model_type and "TI2V" not in args.model_type: - pipe = WanImageToVideoPipeline( - transformer=transformer, - transformer_2=transformer_2, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - boundary_ratio=0.9, - ) - elif "Wan2.2" and "T2V" in args.model_type: - pipe = WanPipeline( - transformer=transformer, - transformer_2=transformer_2, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - boundary_ratio=0.875, - ) - elif "Wan2.2" and "TI2V" in args.model_type: - pipe = WanPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - expand_timesteps=True, - ) - elif "I2V" in args.model_type or "FLF2V" in args.model_type: - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=torch.bfloat16 - ) - image_processor = AutoProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") - pipe = WanImageToVideoPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - elif "Wan2.2-VACE" in args.model_type: - pipe = WanVACEPipeline( - transformer=transformer, - transformer_2=transformer_2, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - boundary_ratio=0.875, - ) - elif "Wan-VACE" in args.model_type: - pipe = WanVACEPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ) - elif "Animate" in args.model_type: - image_encoder = CLIPVisionModel.from_pretrained( - "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=torch.bfloat16 - ) - image_processor = CLIPImageProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") - - pipe = WanAnimatePipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - image_encoder=image_encoder, - image_processor=image_processor, - ) - else: - pipe = WanPipeline( - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - vae=vae, - scheduler=scheduler, - ) - - pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") diff --git a/scripts/convert_wuerstchen.py b/scripts/convert_wuerstchen.py deleted file mode 100644 index 826b9b208181..000000000000 --- a/scripts/convert_wuerstchen.py +++ /dev/null @@ -1,115 +0,0 @@ -# Run inside root directory of official source code: https://github.com/dome272/wuerstchen/ -import os - -import torch -from transformers import AutoTokenizer, CLIPTextModel -from vqgan import VQModel - -from diffusers import ( - DDPMWuerstchenScheduler, - WuerstchenCombinedPipeline, - WuerstchenDecoderPipeline, - WuerstchenPriorPipeline, -) -from diffusers.pipelines.wuerstchen import PaellaVQModel, WuerstchenDiffNeXt, WuerstchenPrior - - -model_path = "models/" -device = "cpu" - -paella_vqmodel = VQModel() -state_dict = torch.load(os.path.join(model_path, "vqgan_f4_v1_500k.pt"), map_location=device)["state_dict"] -paella_vqmodel.load_state_dict(state_dict) - -state_dict["vquantizer.embedding.weight"] = state_dict["vquantizer.codebook.weight"] -state_dict.pop("vquantizer.codebook.weight") -vqmodel = PaellaVQModel(num_vq_embeddings=paella_vqmodel.codebook_size, latent_channels=paella_vqmodel.c_latent) -vqmodel.load_state_dict(state_dict) - -# Clip Text encoder and tokenizer -text_encoder = CLIPTextModel.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") -tokenizer = AutoTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") - -# Generator -gen_text_encoder = CLIPTextModel.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K").to("cpu") -gen_tokenizer = AutoTokenizer.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") - -orig_state_dict = torch.load(os.path.join(model_path, "model_v2_stage_b.pt"), map_location=device)["state_dict"] -state_dict = {} -for key in orig_state_dict.keys(): - if key.endswith("in_proj_weight"): - weights = orig_state_dict[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = orig_state_dict[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = orig_state_dict[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = orig_state_dict[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - else: - state_dict[key] = orig_state_dict[key] -decoder = WuerstchenDiffNeXt() -decoder.load_state_dict(state_dict) - -# Prior -orig_state_dict = torch.load(os.path.join(model_path, "model_v3_stage_c.pt"), map_location=device)["ema_state_dict"] -state_dict = {} -for key in orig_state_dict.keys(): - if key.endswith("in_proj_weight"): - weights = orig_state_dict[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = orig_state_dict[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = orig_state_dict[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = orig_state_dict[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - else: - state_dict[key] = orig_state_dict[key] -prior_model = WuerstchenPrior(c_in=16, c=1536, c_cond=1280, c_r=64, depth=32, nhead=24).to(device) -prior_model.load_state_dict(state_dict) - -# scheduler -scheduler = DDPMWuerstchenScheduler() - -# Prior pipeline -prior_pipeline = WuerstchenPriorPipeline( - prior=prior_model, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler -) - -prior_pipeline.save_pretrained("warp-ai/wuerstchen-prior") - -decoder_pipeline = WuerstchenDecoderPipeline( - text_encoder=gen_text_encoder, tokenizer=gen_tokenizer, vqgan=vqmodel, decoder=decoder, scheduler=scheduler -) -decoder_pipeline.save_pretrained("warp-ai/wuerstchen") - -# Wuerstchen pipeline -wuerstchen_pipeline = WuerstchenCombinedPipeline( - # Decoder - text_encoder=gen_text_encoder, - tokenizer=gen_tokenizer, - decoder=decoder, - scheduler=scheduler, - vqgan=vqmodel, - # Prior - prior_tokenizer=tokenizer, - prior_text_encoder=text_encoder, - prior=prior_model, - prior_scheduler=scheduler, -) -wuerstchen_pipeline.save_pretrained("warp-ai/WuerstchenCombinedPipeline") diff --git a/scripts/convert_zero123_to_diffusers.py b/scripts/convert_zero123_to_diffusers.py deleted file mode 100644 index b46633fae7ff..000000000000 --- a/scripts/convert_zero123_to_diffusers.py +++ /dev/null @@ -1,807 +0,0 @@ -""" -This script modified from -https://github.com/huggingface/diffusers/blob/bc691231360a4cbc7d19a58742ebb8ed0f05e027/scripts/convert_original_stable_diffusion_to_diffusers.py - -Convert original Zero1to3 checkpoint to diffusers checkpoint. - -# run the convert script -$ python convert_zero123_to_diffusers.py \ - --checkpoint_path /path/zero123/105000.ckpt \ - --dump_path ./zero1to3 \ - --original_config_file /path/zero123/configs/sd-objaverse-finetune-c_concat-256.yaml -``` -""" - -import argparse - -import torch -import yaml -from accelerate import init_empty_weights -from accelerate.utils import set_module_tensor_to_device -from pipeline_zero1to3 import CCProjection, Zero1to3StableDiffusionPipeline -from transformers import ( - CLIPImageProcessor, - CLIPVisionModelWithProjection, -) - -from diffusers.models import ( - AutoencoderKL, - UNet2DConditionModel, -) -from diffusers.schedulers import DDIMScheduler -from diffusers.utils import logging - - -logger = logging.get_logger(__name__) - - -def create_unet_diffusers_config(original_config, image_size: int, controlnet=False): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - if controlnet: - unet_params = original_config["model"]["params"]["control_stage_config"]["params"] - else: - if ( - "unet_config" in original_config["model"]["params"] - and original_config["model"]["params"]["unet_config"] is not None - ): - unet_params = original_config["model"]["params"]["unet_config"]["params"] - else: - unet_params = original_config["model"]["params"]["network_config"]["params"] - - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - - block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] - - down_block_types = [] - resolution = 1 - for i in range(len(block_out_channels)): - block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" - down_block_types.append(block_type) - if i != len(block_out_channels) - 1: - resolution *= 2 - - up_block_types = [] - for i in range(len(block_out_channels)): - block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" - up_block_types.append(block_type) - resolution //= 2 - - if unet_params["transformer_depth"] is not None: - transformer_layers_per_block = ( - unet_params["transformer_depth"] - if isinstance(unet_params["transformer_depth"], int) - else list(unet_params["transformer_depth"]) - ) - else: - transformer_layers_per_block = 1 - - vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) - - head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None - use_linear_projection = ( - unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False - ) - if use_linear_projection: - # stable diffusion 2-base-512 and 2-768 - if head_dim is None: - head_dim_mult = unet_params["model_channels"] // unet_params["num_head_channels"] - head_dim = [head_dim_mult * c for c in list(unet_params["channel_mult"])] - - class_embed_type = None - addition_embed_type = None - addition_time_embed_dim = None - projection_class_embeddings_input_dim = None - context_dim = None - - if unet_params["context_dim"] is not None: - context_dim = ( - unet_params["context_dim"] - if isinstance(unet_params["context_dim"], int) - else unet_params["context_dim"][0] - ) - - if "num_classes" in unet_params: - if unet_params["num_classes"] == "sequential": - if context_dim in [2048, 1280]: - # SDXL - addition_embed_type = "text_time" - addition_time_embed_dim = 256 - else: - class_embed_type = "projection" - assert "adm_in_channels" in unet_params - projection_class_embeddings_input_dim = unet_params["adm_in_channels"] - else: - raise NotImplementedError(f"Unknown conditional unet num_classes config: {unet_params['num_classes']}") - - config = { - "sample_size": image_size // vae_scale_factor, - "in_channels": unet_params["in_channels"], - "down_block_types": tuple(down_block_types), - "block_out_channels": tuple(block_out_channels), - "layers_per_block": unet_params["num_res_blocks"], - "cross_attention_dim": context_dim, - "attention_head_dim": head_dim, - "use_linear_projection": use_linear_projection, - "class_embed_type": class_embed_type, - "addition_embed_type": addition_embed_type, - "addition_time_embed_dim": addition_time_embed_dim, - "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, - "transformer_layers_per_block": transformer_layers_per_block, - } - - if controlnet: - config["conditioning_channels"] = unet_params["hint_channels"] - else: - config["out_channels"] = unet_params["out_channels"] - config["up_block_types"] = tuple(up_block_types) - - return config - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) - shape = old_checkpoint[path["old"]].shape - if is_attn_weight and len(shape) == 3: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif is_attn_weight and len(shape) == 4: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def convert_ldm_unet_checkpoint( - checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False -): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - - if skip_extract_state_dict: - unet_state_dict = checkpoint - else: - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - if controlnet: - unet_key = "control_model." - else: - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning(f"Checkpoint {path} has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint[flat_ema_key] - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint[key] - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - if config["addition_embed_type"] == "text_time": - new_checkpoint["add_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["add_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["add_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["add_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - if not controlnet: - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in output_block_list.items()} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - if controlnet: - # conditioning embedding - - orig_index = 0 - - new_checkpoint["controlnet_cond_embedding.conv_in.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_in.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - orig_index += 2 - - diffusers_index = 0 - - while diffusers_index < 6: - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - diffusers_index += 1 - orig_index += 2 - - new_checkpoint["controlnet_cond_embedding.conv_out.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_out.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - # down blocks - for i in range(num_input_blocks): - new_checkpoint[f"controlnet_down_blocks.{i}.weight"] = unet_state_dict.pop(f"zero_convs.{i}.0.weight") - new_checkpoint[f"controlnet_down_blocks.{i}.bias"] = unet_state_dict.pop(f"zero_convs.{i}.0.bias") - - # mid block - new_checkpoint["controlnet_mid_block.weight"] = unet_state_dict.pop("middle_block_out.0.weight") - new_checkpoint["controlnet_mid_block.bias"] = unet_state_dict.pop("middle_block_out.0.bias") - - return new_checkpoint - - -def create_vae_diffusers_config(original_config, image_size: int): - """ - Creates a config for the diffusers based on the config of the LDM model. - """ - vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] - _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] - - block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] - down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) - up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) - - config = { - "sample_size": image_size, - "in_channels": vae_params["in_channels"], - "out_channels": vae_params["out_ch"], - "down_block_types": tuple(down_block_types), - "up_block_types": tuple(up_block_types), - "block_out_channels": tuple(block_out_channels), - "latent_channels": vae_params["z_channels"], - "layers_per_block": vae_params["num_res_blocks"], - } - return config - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - vae_key = "first_stage_model." - keys = list(checkpoint.keys()) - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint - - -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - -def convert_from_original_zero123_ckpt(checkpoint_path, original_config_file, extract_ema, device): - ckpt = torch.load(checkpoint_path, map_location=device) - ckpt["global_step"] - checkpoint = ckpt["state_dict"] - del ckpt - torch.cuda.empty_cache() - - original_config = yaml.safe_load(original_config_file) - original_config["model"]["params"]["cond_stage_config"]["target"].split(".")[-1] - num_in_channels = 8 - original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels - prediction_type = "epsilon" - image_size = 256 - num_train_timesteps = getattr(original_config["model"]["params"], "timesteps", None) or 1000 - - beta_start = getattr(original_config["model"]["params"], "linear_start", None) or 0.02 - beta_end = getattr(original_config["model"]["params"], "linear_end", None) or 0.085 - scheduler = DDIMScheduler( - beta_end=beta_end, - beta_schedule="scaled_linear", - beta_start=beta_start, - num_train_timesteps=num_train_timesteps, - steps_offset=1, - clip_sample=False, - set_alpha_to_one=False, - prediction_type=prediction_type, - ) - scheduler.register_to_config(clip_sample=False) - - # Convert the UNet2DConditionModel model. - upcast_attention = None - unet_config = create_unet_diffusers_config(original_config, image_size=image_size) - unet_config["upcast_attention"] = upcast_attention - with init_empty_weights(): - unet = UNet2DConditionModel(**unet_config) - converted_unet_checkpoint = convert_ldm_unet_checkpoint( - checkpoint, unet_config, path=None, extract_ema=extract_ema - ) - for param_name, param in converted_unet_checkpoint.items(): - set_module_tensor_to_device(unet, param_name, "cpu", value=param) - - # Convert the VAE model. - vae_config = create_vae_diffusers_config(original_config, image_size=image_size) - converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) - - if ( - "model" in original_config - and "params" in original_config["model"] - and "scale_factor" in original_config["model"]["params"] - ): - vae_scaling_factor = original_config["model"]["params"]["scale_factor"] - else: - vae_scaling_factor = 0.18215 # default SD scaling factor - - vae_config["scaling_factor"] = vae_scaling_factor - - with init_empty_weights(): - vae = AutoencoderKL(**vae_config) - - for param_name, param in converted_vae_checkpoint.items(): - set_module_tensor_to_device(vae, param_name, "cpu", value=param) - - feature_extractor = CLIPImageProcessor.from_pretrained( - "lambdalabs/sd-image-variations-diffusers", subfolder="feature_extractor" - ) - image_encoder = CLIPVisionModelWithProjection.from_pretrained( - "lambdalabs/sd-image-variations-diffusers", subfolder="image_encoder" - ) - - cc_projection = CCProjection() - cc_projection.load_state_dict( - { - "projection.weight": checkpoint["cc_projection.weight"].cpu(), - "projection.bias": checkpoint["cc_projection.bias"].cpu(), - } - ) - - pipe = Zero1to3StableDiffusionPipeline( - vae, image_encoder, unet, scheduler, None, feature_extractor, cc_projection, requires_safety_checker=False - ) - - return pipe - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - - parser.add_argument( - "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." - ) - parser.add_argument( - "--original_config_file", - default=None, - type=str, - help="The YAML config file corresponding to the original architecture.", - ) - parser.add_argument( - "--extract_ema", - action="store_true", - help=( - "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" - " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" - " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." - ), - ) - parser.add_argument( - "--to_safetensors", - action="store_true", - help="Whether to store pipeline in safetensors format or not.", - ) - parser.add_argument("--half", action="store_true", help="Save weights in half precision.") - parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") - parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") - args = parser.parse_args() - - pipe = convert_from_original_zero123_ckpt( - checkpoint_path=args.checkpoint_path, - original_config_file=args.original_config_file, - extract_ema=args.extract_ema, - device=args.device, - ) - - if args.half: - pipe.to(dtype=torch.float16) - - pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/export_pipeline_checkpoint.py b/scripts/export_pipeline_checkpoint.py new file mode 100644 index 000000000000..bcaeb21cb68b --- /dev/null +++ b/scripts/export_pipeline_checkpoint.py @@ -0,0 +1,30 @@ +"""Package a local SD/SDXL pipeline into an original checkpoint using the shared component conversions.""" + +import argparse + +import torch + +from diffusers.loaders.conversion.pipeline import export_pipeline_checkpoint + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", required=True, help="Local Diffusers pipeline directory") + parser.add_argument("--output", required=True, help="New output checkpoint file") + parser.add_argument("--pipeline-format", choices=("sd", "sdxl"), help="Default: infer from component directories") + parser.add_argument("--output-format", choices=("safetensors", "pytorch"), default="safetensors") + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), help="Default: preserve source dtypes") + args = parser.parse_args() + print( + export_pipeline_checkpoint( + args.input, + args.output, + pipeline_format=args.pipeline_format, + output_format=args.output_format, + dtype=getattr(torch, args.dtype) if args.dtype else None, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/convert_stable_diffusion_checkpoint_to_onnx.py b/scripts/export_stable_diffusion_checkpoint_to_onnx.py similarity index 100% rename from scripts/convert_stable_diffusion_checkpoint_to_onnx.py rename to scripts/export_stable_diffusion_checkpoint_to_onnx.py diff --git a/scripts/convert_stable_diffusion_controlnet_to_onnx.py b/scripts/export_stable_diffusion_controlnet_to_onnx.py similarity index 99% rename from scripts/convert_stable_diffusion_controlnet_to_onnx.py rename to scripts/export_stable_diffusion_controlnet_to_onnx.py index 4af39b287836..82a17c5d9ee7 100644 --- a/scripts/convert_stable_diffusion_controlnet_to_onnx.py +++ b/scripts/export_stable_diffusion_controlnet_to_onnx.py @@ -232,14 +232,14 @@ def convert_models( Function to convert models in stable diffusion controlnet pipeline into ONNX format Example: - python convert_stable_diffusion_controlnet_to_onnx.py + python export_stable_diffusion_controlnet_to_onnx.py --model_path danbrown/RevAnimated-v1-2-2 --controlnet_path lllyasviel/control_v11f1e_sd15_tile ioclab/brightness-controlnet --output_path path-to-models-stable_diffusion/RevAnimated-v1-2-2 --fp16 Example for SD XL: - python convert_stable_diffusion_controlnet_to_onnx.py + python export_stable_diffusion_controlnet_to_onnx.py --model_path stabilityai/stable-diffusion-xl-base-1.0 --controlnet_path SargeZT/sdxl-controlnet-seg --output_path path-to-models-stable_diffusion/stable-diffusion-xl-base-1.0 diff --git a/scripts/convert_stable_diffusion_controlnet_to_tensorrt.py b/scripts/export_stable_diffusion_controlnet_to_tensorrt.py similarity index 97% rename from scripts/convert_stable_diffusion_controlnet_to_tensorrt.py rename to scripts/export_stable_diffusion_controlnet_to_tensorrt.py index 52ab02c221e9..89a6e6e04e5f 100644 --- a/scripts/convert_stable_diffusion_controlnet_to_tensorrt.py +++ b/scripts/export_stable_diffusion_controlnet_to_tensorrt.py @@ -9,14 +9,14 @@ def convert_models(onnx_path: str, num_controlnet: int, output_path: str, fp16: Function to convert models in stable diffusion controlnet pipeline into TensorRT format Example: - python convert_stable_diffusion_controlnet_to_tensorrt.py + python export_stable_diffusion_controlnet_to_tensorrt.py --onnx_path path-to-models-stable_diffusion/RevAnimated-v1-2-2/unet/model.onnx --output_path path-to-models-stable_diffusion/RevAnimated-v1-2-2/unet/model.engine --fp16 --num_controlnet 2 Example for SD XL: - python convert_stable_diffusion_controlnet_to_tensorrt.py + python export_stable_diffusion_controlnet_to_tensorrt.py --onnx_path path-to-models-stable_diffusion/stable-diffusion-xl-base-1.0/unet/model.onnx --output_path path-to-models-stable_diffusion/stable-diffusion-xl-base-1.0/unet/model.engine --fp16 diff --git a/scripts/convert_vae_diff_to_onnx.py b/scripts/export_vae_diff_to_onnx.py similarity index 100% rename from scripts/convert_vae_diff_to_onnx.py rename to scripts/export_vae_diff_to_onnx.py diff --git a/scripts/merge_lora.py b/scripts/merge_lora.py new file mode 100644 index 000000000000..b63dff300cbc --- /dev/null +++ b/scripts/merge_lora.py @@ -0,0 +1,26 @@ +"""Merge a LoRA adapter into a Diffusers pipeline using its supported adapter loader.""" + +import argparse +from pathlib import Path + +from diffusers import DiffusionPipeline + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model", required=True, help="Base Diffusers pipeline directory or Hub ID") + parser.add_argument("--adapter", required=True, help="LoRA file, directory or Hub ID") + parser.add_argument("--output", required=True, help="New output pipeline directory") + parser.add_argument("--scale", type=float, default=1.0) + args = parser.parse_args() + if Path(args.output).exists(): + raise FileExistsError(f"Output already exists: {args.output}") + pipeline = DiffusionPipeline.from_pretrained(args.model) + pipeline.load_lora_weights(args.adapter) + pipeline.fuse_lora(lora_scale=args.scale, safe_fusing=True) + pipeline.unload_lora_weights() + pipeline.save_pretrained(args.output, safe_serialization=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/recipes/README.md b/scripts/recipes/README.md new file mode 100644 index 000000000000..4cfda6da4b3e --- /dev/null +++ b/scripts/recipes/README.md @@ -0,0 +1,39 @@ +# Pipeline recipe dependencies + +Run a recipe through `python scripts/build_pipeline.py --help`. Install this checkout into the +environment that runs the recipe. The generic `convert_checkpoint.py` command does not need the original +model runtimes listed below; recipes use them to read original containers or assemble complete pipelines. + +These recipes require additional packages or source trees beyond the usual Diffusers dependencies: + +| Recipe | Additional dependency | +| --- | --- | +| `amused` | [`open-muse`](https://github.com/huggingface/open-muse), which provides `muse` | +| `blip_diffusion` | [`LAVIS`](https://github.com/salesforce/LAVIS) from source, including its BLIP Diffusion models | +| `dance_diffusion` | [`sample-generator`](https://github.com/Harmonai-org/sample-generator) and [`v-diffusion-pytorch`](https://github.com/crowsonkb/v-diffusion-pytorch) | +| `dit` | `torchvision`, matched to the installed PyTorch build | +| `k_upscaler` | `k-diffusion` | +| `music_spectrogram` | [`music-spectrogram-diffusion`](https://github.com/magenta/music-spectrogram-diffusion), T5X, T5, SeqIO, JAX/Flax, TensorFlow and TensorFlow Text | +| `sana`, `sana_video` | `termcolor` | +| `wuerstchen` | [`Wuerstchen`](https://github.com/dome272/Wuerstchen), which provides `vqgan`, and [`pytorch-tools`](https://github.com/pabloppp/pytorch-tools), which provides `torchtools` | +| `zero123` | This checkout's `examples/community` directory on `PYTHONPATH`, providing `pipeline_zero1to3` | + +Use separate environments for older reference implementations when their dependency constraints conflict. +Putting a source directory on `PYTHONPATH` does not install its dependencies. In particular, the similarly named +PyPI `torchtools` package is not the Wuerstchen dependency linked above. + +LAVIS uses older Transformers and Diffusers import paths. Updating these dependencies can require corresponding +changes in the external LAVIS checkout. Keep those compatibility changes separate from checkpoint tensor mappings. +The music recipe's TensorFlow Text dependency also needs a supported native wheel; use a compatible Linux +environment when that wheel is unavailable for the host platform. + +## Validation scope + +A successful `--help` check verifies startup imports and argument parsing. It does not exercise checkpoint +loading, model construction, downloads, or pipeline execution. Test those paths with the intended checkpoint +and its matching original runtime before relying on a complete pipeline conversion. + +For tensor-layout refactors, compare shared converter outputs against independent legacy converter outputs. +Compare the tensor names, shapes, dtypes and exact tensor bytes, or hash that same canonical representation. +Serialized checkpoint file hashes can differ because of container metadata and serialization order. Round trips +are useful but cannot detect two directions that consistently implement the same incorrect mapping. diff --git a/scripts/convert_ace_step_to_diffusers.py b/scripts/recipes/ace_step.py similarity index 85% rename from scripts/convert_ace_step_to_diffusers.py rename to scripts/recipes/ace_step.py index 252f5517f91b..6cb0d41b691c 100644 --- a/scripts/convert_ace_step_to_diffusers.py +++ b/scripts/recipes/ace_step.py @@ -1,7 +1,7 @@ # Run this script to convert ACE-Step model weights to a diffusers pipeline. # # Usage: -# python scripts/convert_ace_step_to_diffusers.py \ +# python scripts/recipes/ace_step.py \ # --checkpoint_dir /path/to/ACE-Step-1.5/checkpoints \ # --dit_config acestep-v15-turbo \ # --output_dir /path/to/output/ACE-Step-v1-5-turbo \ @@ -15,6 +15,8 @@ import torch from safetensors.torch import load_file +from diffusers.loaders.conversion import get_conversion + def convert_ace_step_weights(checkpoint_dir, dit_config, output_dir, dtype_str="bf16"): """ @@ -135,57 +137,17 @@ def convert_ace_step_weights(checkpoint_dir, dit_config, output_dir, dtype_str=" audio_token_detokenizer_sd = {} other_sd = {} - # Rename original ACE-Step attention keys to the diffusers `Attention` + - # `AttnProcessor` convention (`to_q`/`to_k`/`to_v`/`to_out.0`/`norm_q`/`norm_k`). - # Applies uniformly to both the DiT (self-attn and cross-attn) and the - # condition-encoder self-attention, since both use `AceStepAttention`. - _ATTN_KEY_RENAMES = [ - (".q_proj.", ".to_q."), - (".k_proj.", ".to_k."), - (".v_proj.", ".to_v."), - (".o_proj.", ".to_out.0."), - (".q_norm.", ".norm_q."), - (".k_norm.", ".norm_k."), - ] - - def _rename_attn_keys(key: str) -> str: - for old, new in _ATTN_KEY_RENAMES: - key = key.replace(old, new) - return key - for key, value in state_dict.items(): if key.startswith("decoder."): - # Strip "decoder." prefix for the transformer - new_key = key[len("decoder.") :] - # The original model uses nn.Sequential for proj_in/proj_out: - # proj_in = Sequential(Lambda, Conv1d, Lambda) - # proj_out = Sequential(Lambda, ConvTranspose1d, Lambda) - # Only the Conv1d/ConvTranspose1d (index 1) has parameters. - # In diffusers, we use standalone Conv1d/ConvTranspose1d named proj_in_conv/proj_out_conv. - new_key = new_key.replace("proj_in.1.", "proj_in_conv.") - new_key = new_key.replace("proj_out.1.", "proj_out_conv.") - new_key = _rename_attn_keys(new_key) - transformer_sd[new_key] = value.to(target_dtype) - elif key.startswith("encoder."): - # Strip "encoder." prefix for the condition encoder - new_key = key[len("encoder.") :] - new_key = _rename_attn_keys(new_key) - condition_encoder_sd[new_key] = value.to(target_dtype) - elif key == "null_condition_emb": - # Learned unconditional embedding (used by the base/SFT CFG path). - # Keep it co-located with the condition encoder since that is where the - # pipeline pulls unconditional sequences from. - condition_encoder_sd["null_condition_emb"] = value.to(target_dtype) + transformer_sd[key.removeprefix("decoder.")] = value + elif key.startswith("encoder.") or key == "null_condition_emb": + condition_encoder_sd[key] = value elif key.startswith("tokenizer."): - new_key = key[len("tokenizer.") :] - new_key = _rename_attn_keys(new_key) - audio_tokenizer_sd[new_key] = value.to(target_dtype) + audio_tokenizer_sd[key.removeprefix("tokenizer.")] = value elif key.startswith("detokenizer."): - new_key = key[len("detokenizer.") :] - new_key = _rename_attn_keys(new_key) - audio_token_detokenizer_sd[new_key] = value.to(target_dtype) + audio_token_detokenizer_sd[key.removeprefix("detokenizer.")] = value else: - other_sd[key] = value.to(target_dtype) + other_sd[key] = value print(f" Transformer keys: {len(transformer_sd)}") print(f" Condition encoder keys: {len(condition_encoder_sd)}") @@ -316,9 +278,7 @@ def _rename_attn_keys(key: str) -> str: silence_latent_src = os.path.join(dit_dir, "silence_latent.pt") if os.path.exists(silence_latent_src): silence_raw = torch.load(silence_latent_src, weights_only=True, map_location="cpu") - silence_latent = silence_raw.transpose(1, 2).to(target_dtype).contiguous() - print(f" silence_latent raw shape: {tuple(silence_raw.shape)} -> baked shape: {tuple(silence_latent.shape)}") - condition_encoder_sd["silence_latent"] = silence_latent + condition_encoder_sd["silence_latent"] = silence_raw # ========================================================================= # 4. Build the AceStepPipeline in memory and save via `save_pretrained`. @@ -351,19 +311,32 @@ def _rename_attn_keys(key: str) -> str: print("\nConstructing transformer ...") transformer = AceStepTransformer1DModel(**transformer_init_kwargs).to(target_dtype) - transformer.load_state_dict(transformer_sd, strict=True) + transformer.load_state_dict( + get_conversion("AceStepTransformer1DModel", dict(transformer.config)).to_diffusers(transformer_sd), strict=True + ) print("Constructing condition_encoder ...") condition_encoder = AceStepConditionEncoder(**condition_encoder_init_kwargs).to(target_dtype) - condition_encoder.load_state_dict(condition_encoder_sd, strict=True) + condition_encoder.load_state_dict( + get_conversion("AceStepConditionEncoder", dict(condition_encoder.config)).to_diffusers(condition_encoder_sd), + strict=True, + ) print("Constructing audio_tokenizer ...") audio_tokenizer = AceStepAudioTokenizer(**audio_tokenizer_init_kwargs).to(target_dtype) - audio_tokenizer.load_state_dict(audio_tokenizer_sd, strict=True) + audio_tokenizer.load_state_dict( + get_conversion("AceStepAudioTokenizer", dict(audio_tokenizer.config)).to_diffusers(audio_tokenizer_sd), + strict=True, + ) print("Constructing audio_token_detokenizer ...") audio_token_detokenizer = AceStepAudioTokenDetokenizer(**audio_token_detokenizer_init_kwargs).to(target_dtype) - audio_token_detokenizer.load_state_dict(audio_token_detokenizer_sd, strict=True) + audio_token_detokenizer.load_state_dict( + get_conversion("AceStepAudioTokenDetokenizer", dict(audio_token_detokenizer.config)).to_diffusers( + audio_token_detokenizer_sd + ), + strict=True, + ) print("Loading VAE ...") vae = AutoencoderOobleck.from_pretrained(vae_dir).to(target_dtype) diff --git a/scripts/recipes/amused.py b/scripts/recipes/amused.py new file mode 100644 index 000000000000..cc61e91271f7 --- /dev/null +++ b/scripts/recipes/amused.py @@ -0,0 +1,231 @@ +import inspect +import os +from argparse import ArgumentParser + +import numpy as np +import torch +from muse import MaskGiTUViT, VQGANModel +from muse import PipelineMuse as OldPipelineMuse +from transformers import CLIPTextModelWithProjection, CLIPTokenizer + +from diffusers import VQModel +from diffusers.loaders.conversion import get_conversion +from diffusers.models.attention_processor import AttnProcessor +from diffusers.models.unets.uvit_2d import UVit2DModel +from diffusers.pipelines.deprecated.amused.pipeline_amused import AmusedPipeline +from diffusers.schedulers import AmusedScheduler + + +torch.backends.cuda.enable_flash_sdp(False) +torch.backends.cuda.enable_mem_efficient_sdp(False) +torch.backends.cuda.enable_math_sdp(True) + +os.environ["CUDA_LAUNCH_BLOCKING"] = "1" +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" +torch.use_deterministic_algorithms(True) + +# Enable CUDNN deterministic mode +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False +torch.backends.cuda.matmul.allow_tf32 = False + +device = "cuda" + + +def main(): + args = ArgumentParser() + args.add_argument("--model_256", action="store_true") + args.add_argument("--write_to", type=str, required=False, default=None) + args.add_argument("--transformer_path", type=str, required=False, default=None) + args = args.parse_args() + + transformer_path = args.transformer_path + subfolder = "transformer" + + if transformer_path is None: + if args.model_256: + transformer_path = "openMUSE/muse-256" + else: + transformer_path = ( + "../research-run-512-checkpoints/research-run-512-with-downsample-checkpoint-554000/unwrapped_model/" + ) + subfolder = None + + old_transformer = MaskGiTUViT.from_pretrained(transformer_path, subfolder=subfolder) + + old_transformer.to(device) + + old_vae = VQGANModel.from_pretrained("openMUSE/muse-512", subfolder="vae") + old_vae.to(device) + + vqvae = make_vqvae(old_vae) + + tokenizer = CLIPTokenizer.from_pretrained("openMUSE/muse-512", subfolder="text_encoder") + + text_encoder = CLIPTextModelWithProjection.from_pretrained("openMUSE/muse-512", subfolder="text_encoder") + text_encoder.to(device) + + transformer = make_transformer(old_transformer, args.model_256) + + scheduler = AmusedScheduler(mask_token_id=old_transformer.config.mask_token_id) + + new_pipe = AmusedPipeline( + vqvae=vqvae, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, scheduler=scheduler + ) + + old_pipe = OldPipelineMuse( + vae=old_vae, transformer=old_transformer, text_encoder=text_encoder, tokenizer=tokenizer + ) + old_pipe.to(device) + + if args.model_256: + transformer_seq_len = 256 + orig_size = (256, 256) + else: + transformer_seq_len = 1024 + orig_size = (512, 512) + + old_out = old_pipe( + "dog", + generator=torch.Generator(device).manual_seed(0), + transformer_seq_len=transformer_seq_len, + orig_size=orig_size, + timesteps=12, + )[0] + + new_out = new_pipe("dog", generator=torch.Generator(device).manual_seed(0)).images[0] + + old_out = np.array(old_out) + new_out = np.array(new_out) + + diff = np.abs(old_out.astype(np.float64) - new_out.astype(np.float64)) + + # assert diff diff.sum() == 0 + print("skipping pipeline full equivalence check") + + print(f"max diff: {diff.max()}, diff.sum() / diff.size {diff.sum() / diff.size}") + + if args.model_256: + assert diff.max() <= 3 + assert diff.sum() / diff.size < 0.7 + else: + assert diff.max() <= 1 + assert diff.sum() / diff.size < 0.4 + + if args.write_to is not None: + new_pipe.save_pretrained(args.write_to) + + +def make_transformer(old_transformer, model_256): + args = dict(old_transformer.config) + force_down_up_sample = args["force_down_up_sample"] + + signature = inspect.signature(UVit2DModel.__init__) + + args_ = { + "downsample": force_down_up_sample, + "upsample": force_down_up_sample, + "block_out_channels": args["block_out_channels"][0], + "sample_size": 16 if model_256 else 32, + } + + for s in list(signature.parameters.keys()): + if s in ["self", "downsample", "upsample", "sample_size", "block_out_channels"]: + continue + + args_[s] = args[s] + + new_transformer = UVit2DModel(**args_) + new_transformer.to(device) + + new_transformer.set_attn_processor(AttnProcessor()) + + state_dict = old_transformer.state_dict() + + state_dict = get_conversion("UVit2DModel", dict(new_transformer.config)).to_diffusers(state_dict) + + new_transformer.load_state_dict(state_dict) + + input_ids = torch.randint(0, 10, (1, 32, 32), device=old_transformer.device) + encoder_hidden_states = torch.randn((1, 77, 768), device=old_transformer.device) + cond_embeds = torch.randn((1, 768), device=old_transformer.device) + micro_conds = torch.tensor([[512, 512, 0, 0, 6]], dtype=torch.float32, device=old_transformer.device) + + old_out = old_transformer(input_ids.reshape(1, -1), encoder_hidden_states, cond_embeds, micro_conds) + old_out = old_out.reshape(1, 32, 32, 8192).permute(0, 3, 1, 2) + + new_out = new_transformer(input_ids, encoder_hidden_states, cond_embeds, micro_conds) + + # NOTE: these differences are solely due to using the geglu block that has a single linear layer of + # double output dimension instead of two different linear layers + max_diff = (old_out - new_out).abs().max() + total_diff = (old_out - new_out).abs().sum() + print(f"Transformer max_diff: {max_diff} total_diff: {total_diff}") + assert max_diff < 0.01 + assert total_diff < 1500 + + return new_transformer + + +def make_vqvae(old_vae): + new_vae = VQModel( + act_fn="silu", + block_out_channels=[128, 256, 256, 512, 768], + down_block_types=[ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ], + in_channels=3, + latent_channels=64, + layers_per_block=2, + norm_num_groups=32, + num_vq_embeddings=8192, + out_channels=3, + sample_size=32, + up_block_types=[ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ], + mid_block_add_attention=False, + lookup_from_codebook=True, + ) + new_vae.to(device) + + new_state_dict = get_conversion("VQModel", dict(new_vae.config)).to_diffusers(old_vae.state_dict()) + + new_vae.load_state_dict(new_state_dict) + + input = torch.randn((1, 3, 512, 512), device=device) + input = input.clamp(-1, 1) + + old_encoder_output = old_vae.quant_conv(old_vae.encoder(input)) + new_encoder_output = new_vae.quant_conv(new_vae.encoder(input)) + assert (old_encoder_output == new_encoder_output).all() + + old_decoder_output = old_vae.decoder(old_vae.post_quant_conv(old_encoder_output)) + new_decoder_output = new_vae.decoder(new_vae.post_quant_conv(new_encoder_output)) + + # assert (old_decoder_output == new_decoder_output).all() + print("kipping vae decoder equivalence check") + print(f"vae decoder diff {(old_decoder_output - new_decoder_output).float().abs().sum()}") + + old_output = old_vae(input)[0] + new_output = new_vae(input)[0] + + # assert (old_output == new_output).all() + print("skipping full vae equivalence check") + print(f"vae full diff {(old_output - new_output).float().abs().sum()}") + + return new_vae + + # fmt: on + + +if __name__ == "__main__": + main() diff --git a/scripts/recipes/anima.py b/scripts/recipes/anima.py new file mode 100644 index 000000000000..fa2262b109c1 --- /dev/null +++ b/scripts/recipes/anima.py @@ -0,0 +1,159 @@ +""" +Convert Anima checkpoints to Diffusers format. + +Example: +```bash +python scripts/recipes/anima.py \ + --transformer_ckpt_path anima_model/anima-preview3-base.safetensors \ + --text_encoder_ckpt_path anima_model/qwen_3_06b_base.safetensors \ + --vae_ckpt_path anima_model/qwen_image_vae.safetensors \ + --qwen_tokenizer_path path/to/qwen25_tokenizer \ + --t5_tokenizer_path path/to/t5_tokenizer \ + --output_path anima_model/anima-preview3-diffusers \ + --save_pipeline +``` +""" + +import argparse +import pathlib +import sys + +import torch +from accelerate import init_empty_weights +from cosmos import convert_transformer +from safetensors.torch import load_file +from transformers import AutoTokenizer, Qwen3Model, T5TokenizerFast + +from diffusers import ( + AnimaAutoBlocks, + AnimaTextConditioner, + AutoencoderKLQwenImage, + FlowMatchEulerDiscreteScheduler, +) +from diffusers.loaders.conversion.configs.anima import infer_qwen3_config, infer_text_conditioner_config + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +def convert_qwen_image_vae(state_dict): + return AutoencoderKLQwenImage.from_single_file(state_dict, config={}) + + +def convert_text_conditioner(state_dict): + return AnimaTextConditioner.from_single_file(state_dict, config=infer_text_conditioner_config(state_dict)) + + +def convert_text_encoder(state_dict: dict[str, torch.Tensor]) -> Qwen3Model: + state_dict = {key.removeprefix("model."): value for key, value in state_dict.items()} + config = infer_qwen3_config(state_dict) + with init_empty_weights(): + text_encoder = Qwen3Model(config) + + expected_keys = set(text_encoder.state_dict().keys()) + converted_keys = set(state_dict.keys()) + missing_keys = expected_keys - converted_keys + unexpected_keys = converted_keys - expected_keys + if missing_keys or unexpected_keys: + if missing_keys: + print(f"ERROR: missing Qwen3 keys ({len(missing_keys)}):", file=sys.stderr) + for key in sorted(missing_keys): + print(key, file=sys.stderr) + if unexpected_keys: + print(f"ERROR: unexpected Qwen3 keys ({len(unexpected_keys)}):", file=sys.stderr) + for key in sorted(unexpected_keys): + print(key, file=sys.stderr) + sys.exit(1) + + text_encoder.load_state_dict(state_dict, strict=True, assign=True) + return text_encoder + + +def split_anima_transformer_checkpoint( + state_dict: dict[str, torch.Tensor], +) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + transformer_state_dict = {} + text_conditioner_state_dict = {} + adapter_prefix = "net.llm_adapter." + + for key, value in state_dict.items(): + if key.startswith(adapter_prefix): + text_conditioner_state_dict[key.removeprefix(adapter_prefix)] = value + else: + transformer_state_dict[key] = value + + return transformer_state_dict, text_conditioner_state_dict + + +def save_pipeline(args, transformer, text_conditioner, text_encoder, vae): + tokenizer = AutoTokenizer.from_pretrained(args.qwen_tokenizer_path) + t5_tokenizer = T5TokenizerFast.from_pretrained(args.t5_tokenizer_path) + scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0) + + pipe = AnimaAutoBlocks().init_pipeline() + pipe.update_components( + text_encoder=text_encoder, + tokenizer=tokenizer, + t5_tokenizer=t5_tokenizer, + text_conditioner=text_conditioner, + transformer=transformer, + vae=vae, + scheduler=scheduler, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size=args.max_shard_size) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--transformer_ckpt_path", type=str, required=True, help="Path to Anima DiT safetensors") + parser.add_argument("--text_encoder_ckpt_path", type=str, required=True, help="Path to Qwen3 text encoder") + parser.add_argument("--vae_ckpt_path", type=str, required=True, help="Path to Qwen-Image VAE safetensors") + parser.add_argument("--qwen_tokenizer_path", type=str, default=None) + parser.add_argument("--t5_tokenizer_path", type=str, default=None) + parser.add_argument("--output_path", type=str, required=True) + parser.add_argument("--save_pipeline", action="store_true") + parser.add_argument("--dtype", default="bf16", choices=list(DTYPE_MAPPING.keys())) + parser.add_argument("--max_shard_size", default="5GB") + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + output_path = pathlib.Path(args.output_path) + dtype = DTYPE_MAPPING[args.dtype] + + raw_transformer_state_dict = load_file(args.transformer_ckpt_path, device="cpu") + transformer_state_dict, text_conditioner_state_dict = split_anima_transformer_checkpoint( + raw_transformer_state_dict + ) + transformer = convert_transformer( + "Cosmos-2.0-Diffusion-2B-Text2Image", state_dict=transformer_state_dict, weights_only=True + ).to(dtype=dtype) + text_conditioner = convert_text_conditioner(text_conditioner_state_dict).to(dtype=dtype) + + text_encoder_state_dict = load_file(args.text_encoder_ckpt_path, device="cpu") + text_encoder = convert_text_encoder(text_encoder_state_dict).to(dtype=dtype) + + vae_state_dict = load_file(args.vae_ckpt_path, device="cpu") + vae = convert_qwen_image_vae(vae_state_dict).to(dtype=dtype) + + if args.save_pipeline: + if args.qwen_tokenizer_path is None or args.t5_tokenizer_path is None: + raise ValueError("`--qwen_tokenizer_path` and `--t5_tokenizer_path` are required with `--save_pipeline`.") + save_pipeline(args, transformer, text_conditioner, text_encoder, vae) + else: + output_path.mkdir(parents=True, exist_ok=True) + transformer.save_pretrained( + output_path / "transformer", safe_serialization=True, max_shard_size=args.max_shard_size + ) + text_conditioner.save_pretrained( + output_path / "text_conditioner", safe_serialization=True, max_shard_size=args.max_shard_size + ) + text_encoder.save_pretrained( + output_path / "text_encoder", safe_serialization=True, max_shard_size=args.max_shard_size + ) + vae.save_pretrained(output_path / "vae", safe_serialization=True, max_shard_size=args.max_shard_size) diff --git a/scripts/convert_anyflow_to_diffusers.py b/scripts/recipes/anyflow.py similarity index 99% rename from scripts/convert_anyflow_to_diffusers.py rename to scripts/recipes/anyflow.py index 229d286c4701..446c59c13c3e 100644 --- a/scripts/convert_anyflow_to_diffusers.py +++ b/scripts/recipes/anyflow.py @@ -26,7 +26,7 @@ Example: ```bash -python scripts/convert_anyflow_to_diffusers.py \\ +python scripts/recipes/anyflow.py \\ --variant AnyFlow-FAR-Wan2.1-1.3B-Diffusers \\ --ckpt /path/to/anyflow-checkpoint.pt \\ --output-dir /path/to/output/AnyFlow-FAR-Wan2.1-1.3B-Diffusers diff --git a/scripts/recipes/audioldm.py b/scripts/recipes/audioldm.py new file mode 100644 index 000000000000..e503a98ac1c3 --- /dev/null +++ b/scripts/recipes/audioldm.py @@ -0,0 +1,361 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the AudioLDM checkpoints.""" + +import argparse + +import torch +import yaml +from transformers import ( + AutoTokenizer, + ClapTextConfig, + ClapTextModelWithProjection, + SpeechT5HifiGan, + SpeechT5HifiGanConfig, +) + +from diffusers import ( + AudioLDMPipeline, + AutoencoderKL, + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, + UNet2DConditionModel, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.audioldm import ( + DEFAULT_CONFIG, + create_transformers_vocoder_config, + create_unet_diffusers_config, + create_vae_diffusers_config, +) + + +# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config + + +# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_unet_checkpoint +def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, **kwargs): + return convert_component_checkpoint(checkpoint, config, "UNet2DConditionModel", extract_ema=extract_ema) + + +def convert_ldm_vae_checkpoint(checkpoint, config): + return convert_component_checkpoint(checkpoint, config, "AutoencoderKL") + + +CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] + + +def convert_open_clap_checkpoint(checkpoint, config): + prefix = "cond_stage_model.model." + state = { + key.removeprefix(prefix): value + for key, value in checkpoint.items() + if key.startswith(prefix) and key.startswith(prefix + "text_") + } + return convert_component_checkpoint(state, config.to_dict(), "ClapTextModelWithProjection") + + +def convert_hifigan_checkpoint(checkpoint, config): + state = { + key.removeprefix("first_stage_model.vocoder."): value + for key, value in checkpoint.items() + if key.startswith("first_stage_model.vocoder.") + } + return convert_component_checkpoint(state, config.to_dict(), "SpeechT5HifiGan") + + +# Adapted from https://huggingface.co/spaces/haoheliu/audioldm-text-to-audio-generation/blob/84a0384742a22bd80c44e903e241f0623e874f1d/audioldm/utils.py#L72-L73 + + +def load_pipeline_from_original_audioldm_ckpt( + checkpoint_path: str, + original_config_file: str = None, + image_size: int = 512, + prediction_type: str = None, + extract_ema: bool = False, + scheduler_type: str = "ddim", + num_in_channels: int = None, + model_channels: int = None, + num_head_channels: int = None, + device: str = None, + from_safetensors: bool = False, +) -> AudioLDMPipeline: + """ + Load an AudioLDM pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. + + Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the + global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is + recommended that you override the default values and/or supply an `original_config_file` wherever possible. + + Args: + checkpoint_path (`str`): Path to `.ckpt` file. + original_config_file (`str`): + Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically + set to the audioldm-s-full-v2 config. + image_size (`int`, *optional*, defaults to 512): + The image size that the model was trained on. + prediction_type (`str`, *optional*): + The prediction type that the model was trained on. If `None`, will be automatically + inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. + num_in_channels (`int`, *optional*, defaults to None): + The number of UNet input channels. If `None`, it will be automatically inferred from the config. + model_channels (`int`, *optional*, defaults to None): + The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override + to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large. + num_head_channels (`int`, *optional*, defaults to None): + The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override + to 32 for the small and medium checkpoints, and 64 for the large. + scheduler_type (`str`, *optional*, defaults to 'pndm'): + Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", + "ddim"]`. + extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for + checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to + `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for + inference. Non-EMA weights are usually better to continue fine-tuning. + device (`str`, *optional*, defaults to `None`): + The device to use. Pass `None` to determine automatically. + from_safetensors (`str`, *optional*, defaults to `False`): + If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. + return: An AudioLDMPipeline object representing the passed-in `.ckpt`/`.safetensors` file. + """ + + if from_safetensors: + from safetensors import safe_open + + checkpoint = {} + with safe_open(checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + checkpoint[key] = f.get_tensor(key) + else: + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint = torch.load(checkpoint_path, map_location=device) + else: + checkpoint = torch.load(checkpoint_path, map_location=device) + + if "state_dict" in checkpoint: + checkpoint = checkpoint["state_dict"] + + if original_config_file is None: + original_config = DEFAULT_CONFIG + else: + original_config = yaml.safe_load(original_config_file) + + if num_in_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels + + if model_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["model_channels"] = model_channels + + if num_head_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = num_head_channels + + if ( + "parameterization" in original_config["model"]["params"] + and original_config["model"]["params"]["parameterization"] == "v" + ): + if prediction_type is None: + prediction_type = "v_prediction" + else: + if prediction_type is None: + prediction_type = "epsilon" + + if image_size is None: + image_size = 512 + + num_train_timesteps = original_config["model"]["params"]["timesteps"] + beta_start = original_config["model"]["params"]["linear_start"] + beta_end = original_config["model"]["params"]["linear_end"] + + scheduler = DDIMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + steps_offset=1, + clip_sample=False, + set_alpha_to_one=False, + prediction_type=prediction_type, + ) + # make sure scheduler works correctly with DDIM + scheduler.register_to_config(clip_sample=False) + + if scheduler_type == "pndm": + config = dict(scheduler.config) + config["skip_prk_steps"] = True + scheduler = PNDMScheduler.from_config(config) + elif scheduler_type == "lms": + scheduler = LMSDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "heun": + scheduler = HeunDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler": + scheduler = EulerDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler-ancestral": + scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "dpm": + scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) + elif scheduler_type == "ddim": + scheduler = scheduler + else: + raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") + + # Convert the UNet2DModel + unet_config = create_unet_diffusers_config(original_config, image_size=image_size) + unet = UNet2DConditionModel(**unet_config) + + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema + ) + + unet.load_state_dict(converted_unet_checkpoint) + + # Convert the VAE model + vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) + converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + + # Convert the text model + # AudioLDM uses the same configuration and tokenizer as the original CLAP model + config = ClapTextConfig.from_pretrained("laion/clap-htsat-unfused") + tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") + + converted_text_model = convert_open_clap_checkpoint(checkpoint, config) + text_model = ClapTextModelWithProjection(config) + + missing_keys, unexpected_keys = text_model.load_state_dict(converted_text_model, strict=False) + # we expect not to have token_type_ids in our original state dict so let's ignore them + missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) + + if len(unexpected_keys) > 0: + raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") + + if len(missing_keys) > 0: + raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") + + # Convert the vocoder model + vocoder_config = create_transformers_vocoder_config(original_config) + vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) + converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) + + vocoder = SpeechT5HifiGan(vocoder_config) + vocoder.load_state_dict(converted_vocoder_checkpoint) + + # Instantiate the diffusers pipeline + pipe = AudioLDMPipeline( + vae=vae, + text_encoder=text_model, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + vocoder=vocoder, + ) + + return pipe + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--original_config_file", + default=None, + type=str, + help="The YAML config file corresponding to the original architecture.", + ) + parser.add_argument( + "--num_in_channels", + default=None, + type=int, + help="The number of input channels. If `None` number of input channels will be automatically inferred.", + ) + parser.add_argument( + "--model_channels", + default=None, + type=int, + help="The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override" + " to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large.", + ) + parser.add_argument( + "--num_head_channels", + default=None, + type=int, + help="The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override" + " to 32 for the small and medium checkpoints, and 64 for the large.", + ) + parser.add_argument( + "--scheduler_type", + default="ddim", + type=str, + help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", + ) + parser.add_argument( + "--image_size", + default=None, + type=int, + help=("The image size that the model was trained on."), + ) + parser.add_argument( + "--prediction_type", + default=None, + type=str, + help=("The prediction type that the model was trained on."), + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument( + "--from_safetensors", + action="store_true", + help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", + ) + parser.add_argument( + "--to_safetensors", + action="store_true", + help="Whether to store pipeline in safetensors format or not.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") + args = parser.parse_args() + + pipe = load_pipeline_from_original_audioldm_ckpt( + checkpoint_path=args.checkpoint_path, + original_config_file=args.original_config_file, + image_size=args.image_size, + prediction_type=args.prediction_type, + extract_ema=args.extract_ema, + scheduler_type=args.scheduler_type, + num_in_channels=args.num_in_channels, + model_channels=args.model_channels, + num_head_channels=args.num_head_channels, + from_safetensors=args.from_safetensors, + device=args.device, + ) + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/recipes/audioldm2.py b/scripts/recipes/audioldm2.py new file mode 100644 index 000000000000..35c2e7fa7218 --- /dev/null +++ b/scripts/recipes/audioldm2.py @@ -0,0 +1,420 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the AudioLDM2 checkpoints.""" + +import argparse +from typing import List, Union + +import torch +import yaml +from transformers import ( + AutoFeatureExtractor, + AutoTokenizer, + ClapConfig, + ClapModel, + GPT2Config, + GPT2Model, + SpeechT5HifiGan, + SpeechT5HifiGanConfig, + T5Config, + T5EncoderModel, +) + +from diffusers import ( + AudioLDM2Pipeline, + AudioLDM2ProjectionModel, + AudioLDM2UNet2DConditionModel, + AutoencoderKL, + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.audioldm2 import ( + DEFAULT_CONFIG, + create_transformers_vocoder_config, + create_unet_diffusers_config, + create_vae_diffusers_config, +) +from diffusers.utils import is_safetensors_available +from diffusers.utils.import_utils import BACKENDS_MAPPING + + +# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config + + +def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, **kwargs): + return convert_component_checkpoint(checkpoint, config, "AudioLDM2UNet2DConditionModel", extract_ema=extract_ema) + + +def convert_ldm_vae_checkpoint(checkpoint, config): + return convert_component_checkpoint(checkpoint, config, "AutoencoderKL") + + +CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] + + +def convert_open_clap_checkpoint(checkpoint, config): + prefix = "clap.model." + state = {key.removeprefix(prefix): value for key, value in checkpoint.items() if key.startswith(prefix)} + return convert_component_checkpoint(state, config.to_dict(), "ClapModel") + + +def extract_sub_model(checkpoint, key_prefix): + """ + Takes a state dict and returns the state dict for a particular sub-model. + """ + + sub_model_state_dict = {} + keys = list(checkpoint.keys()) + for key in keys: + if key.startswith(key_prefix): + sub_model_state_dict[key.replace(key_prefix, "")] = checkpoint.get(key) + + return sub_model_state_dict + + +def convert_hifigan_checkpoint(checkpoint, config): + state = { + key.removeprefix("first_stage_model.vocoder."): value + for key, value in checkpoint.items() + if key.startswith("first_stage_model.vocoder.") + } + return convert_component_checkpoint(state, config.to_dict(), "SpeechT5HifiGan") + + +def convert_projection_checkpoint(checkpoint): + prefix = "cond_stage_models.0." + modules = ("start_of_sequence_tokens.", "end_of_sequence_tokens.", "input_sequence_embed_linear.") + state = { + key.removeprefix(prefix): value + for key, value in checkpoint.items() + if key.startswith(prefix) and key.removeprefix(prefix).startswith(modules) + } + return convert_component_checkpoint(state, {}, "AudioLDM2ProjectionModel") + + +# Adapted from https://github.com/haoheliu/AudioLDM2/blob/81ad2c6ce015c1310387695e2dae975a7d2ed6fd/audioldm2/utils.py#L143 + + +def load_pipeline_from_original_AudioLDM2_ckpt( + checkpoint_path: str, + original_config_file: str = None, + image_size: int = 1024, + prediction_type: str = None, + extract_ema: bool = False, + scheduler_type: str = "ddim", + cross_attention_dim: Union[List, List[List]] = None, + transformer_layers_per_block: int = None, + device: str = None, + from_safetensors: bool = False, +) -> AudioLDM2Pipeline: + """ + Load an AudioLDM2 pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. + + Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the + global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is + recommended that you override the default values and/or supply an `original_config_file` wherever possible. + + Args: + checkpoint_path (`str`): Path to `.ckpt` file. + original_config_file (`str`): + Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically + set to the AudioLDM2 base config. + image_size (`int`, *optional*, defaults to 1024): + The image size that the model was trained on. + prediction_type (`str`, *optional*): + The prediction type that the model was trained on. If `None`, will be automatically + inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. + scheduler_type (`str`, *optional*, defaults to 'ddim'): + Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", + "ddim"]`. + cross_attention_dim (`list`, *optional*, defaults to `None`): + The dimension of the cross-attention layers. If `None`, the cross-attention dimension will be + automatically inferred. Set to `[768, 1024]` for the base model, or `[768, 1024, None]` for the large model. + transformer_layers_per_block (`int`, *optional*, defaults to `None`): + The number of transformer layers in each transformer block. If `None`, number of layers will be " + "automatically inferred. Set to `1` for the base model, or `2` for the large model. + extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for + checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to + `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for + inference. Non-EMA weights are usually better to continue fine-tuning. + device (`str`, *optional*, defaults to `None`): + The device to use. Pass `None` to determine automatically. + from_safetensors (`str`, *optional*, defaults to `False`): + If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. + return: An AudioLDM2Pipeline object representing the passed-in `.ckpt`/`.safetensors` file. + """ + + if from_safetensors: + if not is_safetensors_available(): + raise ValueError(BACKENDS_MAPPING["safetensors"][1]) + + from safetensors import safe_open + + checkpoint = {} + with safe_open(checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + checkpoint[key] = f.get_tensor(key) + else: + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint = torch.load(checkpoint_path, map_location=device) + else: + checkpoint = torch.load(checkpoint_path, map_location=device) + + if "state_dict" in checkpoint: + checkpoint = checkpoint["state_dict"] + + if original_config_file is None: + original_config = DEFAULT_CONFIG + else: + original_config = yaml.safe_load(original_config_file) + + if image_size is not None: + original_config["model"]["params"]["unet_config"]["params"]["image_size"] = image_size + + if cross_attention_dim is not None: + original_config["model"]["params"]["unet_config"]["params"]["context_dim"] = cross_attention_dim + + if transformer_layers_per_block is not None: + original_config["model"]["params"]["unet_config"]["params"]["transformer_depth"] = transformer_layers_per_block + + if ( + "parameterization" in original_config["model"]["params"] + and original_config["model"]["params"]["parameterization"] == "v" + ): + if prediction_type is None: + prediction_type = "v_prediction" + else: + if prediction_type is None: + prediction_type = "epsilon" + + num_train_timesteps = original_config["model"]["params"]["timesteps"] + beta_start = original_config["model"]["params"]["linear_start"] + beta_end = original_config["model"]["params"]["linear_end"] + + scheduler = DDIMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + steps_offset=1, + clip_sample=False, + set_alpha_to_one=False, + prediction_type=prediction_type, + ) + # make sure scheduler works correctly with DDIM + scheduler.register_to_config(clip_sample=False) + + if scheduler_type == "pndm": + config = dict(scheduler.config) + config["skip_prk_steps"] = True + scheduler = PNDMScheduler.from_config(config) + elif scheduler_type == "lms": + scheduler = LMSDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "heun": + scheduler = HeunDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler": + scheduler = EulerDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler-ancestral": + scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "dpm": + scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) + elif scheduler_type == "ddim": + scheduler = scheduler + else: + raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") + + # Convert the UNet2DModel + unet_config = create_unet_diffusers_config(original_config, image_size=image_size) + unet = AudioLDM2UNet2DConditionModel(**unet_config) + + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema + ) + + unet.load_state_dict(converted_unet_checkpoint) + + # Convert the VAE model + vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) + converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + + # Convert the joint audio-text encoding model + clap_config = ClapConfig.from_pretrained("laion/clap-htsat-unfused") + clap_config.audio_config.update( + { + "patch_embeds_hidden_size": 128, + "hidden_size": 1024, + "depths": [2, 2, 12, 2], + } + ) + # AudioLDM2 uses the same tokenizer and feature extractor as the original CLAP model + clap_tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") + clap_feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused") + + converted_clap_model = convert_open_clap_checkpoint(checkpoint, clap_config) + clap_model = ClapModel(clap_config) + + missing_keys, unexpected_keys = clap_model.load_state_dict(converted_clap_model, strict=False) + # we expect not to have token_type_ids in our original state dict so let's ignore them + missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) + + if len(unexpected_keys) > 0: + raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") + + if len(missing_keys) > 0: + raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") + + # Convert the vocoder model + vocoder_config = create_transformers_vocoder_config(original_config) + vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) + converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) + + vocoder = SpeechT5HifiGan(vocoder_config) + vocoder.load_state_dict(converted_vocoder_checkpoint) + + # Convert the Flan-T5 encoder model: AudioLDM2 uses the same configuration and tokenizer as the original Flan-T5 large model + t5_config = T5Config.from_pretrained("google/flan-t5-large") + converted_t5_checkpoint = extract_sub_model(checkpoint, key_prefix="cond_stage_models.1.model.") + + t5_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large") + # hard-coded in the original implementation (i.e. not retrievable from the config) + t5_tokenizer.model_max_length = 128 + t5_model = T5EncoderModel(t5_config) + t5_model.load_state_dict(converted_t5_checkpoint) + + # Convert the GPT2 encoder model: AudioLDM2 uses the same configuration as the original GPT2 base model + gpt2_config = GPT2Config.from_pretrained("gpt2") + gpt2_model = GPT2Model(gpt2_config) + gpt2_model.config.max_new_tokens = original_config["model"]["params"]["cond_stage_config"][ + "crossattn_audiomae_generated" + ]["params"]["sequence_gen_length"] + + converted_gpt2_checkpoint = extract_sub_model(checkpoint, key_prefix="cond_stage_models.0.model.") + gpt2_model.load_state_dict(converted_gpt2_checkpoint) + + # Convert the extra embedding / projection layers + projection_model = AudioLDM2ProjectionModel(clap_config.projection_dim, t5_config.d_model, gpt2_config.n_embd) + + converted_projection_checkpoint = convert_projection_checkpoint(checkpoint) + projection_model.load_state_dict(converted_projection_checkpoint) + + # Instantiate the diffusers pipeline + pipe = AudioLDM2Pipeline( + vae=vae, + text_encoder=clap_model, + text_encoder_2=t5_model, + projection_model=projection_model, + language_model=gpt2_model, + tokenizer=clap_tokenizer, + tokenizer_2=t5_tokenizer, + feature_extractor=clap_feature_extractor, + unet=unet, + scheduler=scheduler, + vocoder=vocoder, + ) + + return pipe + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--original_config_file", + default=None, + type=str, + help="The YAML config file corresponding to the original architecture.", + ) + parser.add_argument( + "--cross_attention_dim", + default=None, + type=int, + nargs="+", + help="The dimension of the cross-attention layers. If `None`, the cross-attention dimension will be " + "automatically inferred. Set to `768+1024` for the base model, or `768+1024+640` for the large model", + ) + parser.add_argument( + "--transformer_layers_per_block", + default=None, + type=int, + help="The number of transformer layers in each transformer block. If `None`, number of layers will be " + "automatically inferred. Set to `1` for the base model, or `2` for the large model.", + ) + parser.add_argument( + "--scheduler_type", + default="ddim", + type=str, + help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", + ) + parser.add_argument( + "--image_size", + default=1048, + type=int, + help="The image size that the model was trained on.", + ) + parser.add_argument( + "--prediction_type", + default=None, + type=str, + help=("The prediction type that the model was trained on."), + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument( + "--from_safetensors", + action="store_true", + help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", + ) + parser.add_argument( + "--to_safetensors", + action="store_true", + help="Whether to store pipeline in safetensors format or not.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") + args = parser.parse_args() + + pipe = load_pipeline_from_original_AudioLDM2_ckpt( + checkpoint_path=args.checkpoint_path, + original_config_file=args.original_config_file, + image_size=args.image_size, + prediction_type=args.prediction_type, + extract_ema=args.extract_ema, + scheduler_type=args.scheduler_type, + cross_attention_dim=args.cross_attention_dim, + transformer_layers_per_block=args.transformer_layers_per_block, + from_safetensors=args.from_safetensors, + device=args.device, + ) + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/recipes/blip_diffusion.py b/scripts/recipes/blip_diffusion.py new file mode 100644 index 000000000000..e509c98afdc2 --- /dev/null +++ b/scripts/recipes/blip_diffusion.py @@ -0,0 +1,106 @@ +""" +This script requires you to build `LAVIS` from source, since the pip version doesn't have BLIP Diffusion. Follow instructions here: https://github.com/salesforce/LAVIS/tree/main. +""" + +import argparse +import os +import tempfile + +import torch +from lavis.models import load_model_and_preprocess +from transformers import CLIPTokenizer +from transformers.models.blip_2.configuration_blip_2 import Blip2Config + +from diffusers import ( + AutoencoderKL, + PNDMScheduler, + UNet2DConditionModel, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.blip_diffusion import BLIP2_CONFIG +from diffusers.pipelines import BlipDiffusionPipeline +from diffusers.pipelines.deprecated.blip_diffusion.blip_image_processing import BlipImageProcessor +from diffusers.pipelines.deprecated.blip_diffusion.modeling_blip2 import Blip2QFormerModel +from diffusers.pipelines.deprecated.blip_diffusion.modeling_ctx_clip import ContextCLIPTextModel + + +blip2config = Blip2Config(**BLIP2_CONFIG) + + +def qformer_model_from_original_config(): + qformer = Blip2QFormerModel(blip2config) + return qformer + + +def qformer_original_checkpoint_to_diffusers_checkpoint(model): + prefixes = ("blip.Qformer.bert.", "blip.query_tokens", "blip.visual_encoder.", "blip.ln_vision.", "proj_layer.") + state = {key: value for key, value in model.items() if key.startswith(prefixes)} + # The source position IDs are deterministic, and the shared conversion recreates them. + state.pop("blip.Qformer.bert.embeddings.position_ids", None) + return get_conversion("Blip2QFormerModel", blip2config.to_dict()).to_diffusers(state) + + +def get_qformer(model): + print("loading qformer") + + qformer = qformer_model_from_original_config() + qformer_diffusers_checkpoint = qformer_original_checkpoint_to_diffusers_checkpoint(model) + + load_checkpoint_to_model(qformer_diffusers_checkpoint, qformer) + + print("done loading qformer") + return qformer + + +def load_checkpoint_to_model(checkpoint, model): + with tempfile.NamedTemporaryFile(delete=False) as file: + torch.save(checkpoint, file.name) + del checkpoint + model.load_state_dict(torch.load(file.name), strict=False) + + os.remove(file.name) + + +def save_blip_diffusion_model(model, args): + qformer = get_qformer(model) + qformer.eval() + + text_encoder = ContextCLIPTextModel.from_pretrained( + "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="text_encoder" + ) + vae = AutoencoderKL.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="vae") + unet = UNet2DConditionModel.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="unet") + vae.eval() + text_encoder.eval() + scheduler = PNDMScheduler( + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + set_alpha_to_one=False, + skip_prk_steps=True, + ) + tokenizer = CLIPTokenizer.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="tokenizer") + image_processor = BlipImageProcessor() + blip_diffusion = BlipDiffusionPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + vae=vae, + unet=unet, + scheduler=scheduler, + qformer=qformer, + image_processor=image_processor, + ) + blip_diffusion.save_pretrained(args.checkpoint_path) + + +def main(args): + model, _, _ = load_model_and_preprocess("blip_diffusion", "base", device="cpu", is_eval=True) + save_blip_diffusion_model(model.state_dict(), args) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") + args = parser.parse_args() + + main(args) diff --git a/scripts/convert_cogvideox_to_diffusers.py b/scripts/recipes/cogvideox.py similarity index 56% rename from scripts/convert_cogvideox_to_diffusers.py rename to scripts/recipes/cogvideox.py index 5220aa0bde9d..c9538ee4ea13 100644 --- a/scripts/convert_cogvideox_to_diffusers.py +++ b/scripts/recipes/cogvideox.py @@ -11,130 +11,16 @@ CogVideoXPipeline, CogVideoXTransformer3DModel, ) +from diffusers.loaders.conversion.cogvideox import ( + convert_cogvideox_transformer_checkpoint_to_diffusers, + convert_cogvideox_vae_checkpoint_to_diffusers, +) +from diffusers.models.modeling_utils import load_state_dict -def reassign_query_key_value_inplace(key: str, state_dict: Dict[str, Any]): - to_q_key = key.replace("query_key_value", "to_q") - to_k_key = key.replace("query_key_value", "to_k") - to_v_key = key.replace("query_key_value", "to_v") - to_q, to_k, to_v = torch.chunk(state_dict[key], chunks=3, dim=0) - state_dict[to_q_key] = to_q - state_dict[to_k_key] = to_k - state_dict[to_v_key] = to_v - state_dict.pop(key) - - -def reassign_query_key_layernorm_inplace(key: str, state_dict: Dict[str, Any]): - layer_id, weight_or_bias = key.split(".")[-2:] - - if "query" in key: - new_key = f"transformer_blocks.{layer_id}.attn1.norm_q.{weight_or_bias}" - elif "key" in key: - new_key = f"transformer_blocks.{layer_id}.attn1.norm_k.{weight_or_bias}" - - state_dict[new_key] = state_dict.pop(key) - - -def reassign_adaln_norm_inplace(key: str, state_dict: Dict[str, Any]): - layer_id, _, weight_or_bias = key.split(".")[-3:] - - weights_or_biases = state_dict[key].chunk(12, dim=0) - norm1_weights_or_biases = torch.cat(weights_or_biases[0:3] + weights_or_biases[6:9]) - norm2_weights_or_biases = torch.cat(weights_or_biases[3:6] + weights_or_biases[9:12]) - - norm1_key = f"transformer_blocks.{layer_id}.norm1.linear.{weight_or_bias}" - state_dict[norm1_key] = norm1_weights_or_biases - - norm2_key = f"transformer_blocks.{layer_id}.norm2.linear.{weight_or_bias}" - state_dict[norm2_key] = norm2_weights_or_biases - - state_dict.pop(key) - - -def remove_keys_inplace(key: str, state_dict: Dict[str, Any]): - state_dict.pop(key) - - -def replace_up_keys_inplace(key: str, state_dict: Dict[str, Any]): - key_split = key.split(".") - layer_index = int(key_split[2]) - replace_layer_index = 4 - 1 - layer_index - - key_split[1] = "up_blocks" - key_split[2] = str(replace_layer_index) - new_key = ".".join(key_split) - - state_dict[new_key] = state_dict.pop(key) - - -TRANSFORMER_KEYS_RENAME_DICT = { - "transformer.final_layernorm": "norm_final", - "transformer": "transformer_blocks", - "attention": "attn1", - "mlp": "ff.net", - "dense_h_to_4h": "0.proj", - "dense_4h_to_h": "2", - ".layers": "", - "dense": "to_out.0", - "input_layernorm": "norm1.norm", - "post_attn1_layernorm": "norm2.norm", - "time_embed.0": "time_embedding.linear_1", - "time_embed.2": "time_embedding.linear_2", - "ofs_embed.0": "ofs_embedding.linear_1", - "ofs_embed.2": "ofs_embedding.linear_2", - "mixins.patch_embed": "patch_embed", - "mixins.final_layer.norm_final": "norm_out.norm", - "mixins.final_layer.linear": "proj_out", - "mixins.final_layer.adaLN_modulation.1": "norm_out.linear", - "mixins.pos_embed.pos_embedding": "patch_embed.pos_embedding", # Specific to CogVideoX-5b-I2V -} - -TRANSFORMER_SPECIAL_KEYS_REMAP = { - "query_key_value": reassign_query_key_value_inplace, - "query_layernorm_list": reassign_query_key_layernorm_inplace, - "key_layernorm_list": reassign_query_key_layernorm_inplace, - "adaln_layer.adaLN_modulations": reassign_adaln_norm_inplace, - "embed_tokens": remove_keys_inplace, - "freqs_sin": remove_keys_inplace, - "freqs_cos": remove_keys_inplace, - "position_embedding": remove_keys_inplace, -} - -VAE_KEYS_RENAME_DICT = { - "block.": "resnets.", - "down.": "down_blocks.", - "downsample": "downsamplers.0", - "upsample": "upsamplers.0", - "nin_shortcut": "conv_shortcut", - "encoder.mid.block_1": "encoder.mid_block.resnets.0", - "encoder.mid.block_2": "encoder.mid_block.resnets.1", - "decoder.mid.block_1": "decoder.mid_block.resnets.0", - "decoder.mid.block_2": "decoder.mid_block.resnets.1", -} - -VAE_SPECIAL_KEYS_REMAP = { - "loss": remove_keys_inplace, - "up.": replace_up_keys_inplace, -} - TOKENIZER_MAX_LENGTH = 226 -def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: - state_dict = saved_dict - if "model" in saved_dict.keys(): - state_dict = state_dict["model"] - if "module" in saved_dict.keys(): - state_dict = state_dict["module"] - if "state_dict" in saved_dict.keys(): - state_dict = state_dict["state_dict"] - return state_dict - - -def update_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) -> dict[str, Any]: - state_dict[new_key] = state_dict.pop(old_key) - - def convert_transformer( ckpt_path: str, num_layers: int, @@ -144,9 +30,7 @@ def convert_transformer( dtype: torch.dtype, init_kwargs: Dict[str, Any], ): - PREFIX_KEY = "model.diffusion_model." - - original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) + original_state_dict = load_state_dict(ckpt_path) transformer = CogVideoXTransformer3DModel( in_channels=32 if i2v else 16, num_layers=num_layers, @@ -157,19 +41,8 @@ def convert_transformer( **init_kwargs, ).to(dtype=dtype) - for key in list(original_state_dict.keys()): - new_key = key[len(PREFIX_KEY) :] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - transformer.load_state_dict(original_state_dict, strict=True) + converted = convert_cogvideox_transformer_checkpoint_to_diffusers(original_state_dict, transformer.config) + transformer.load_state_dict(converted, strict=True) return transformer @@ -178,22 +51,11 @@ def convert_vae(ckpt_path: str, scaling_factor: float, version: str, dtype: torc if version == "1.5": init_kwargs.update({"invert_scale_latents": True}) - original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True)) + original_state_dict = load_state_dict(ckpt_path) vae = AutoencoderKLCogVideoX(**init_kwargs).to(dtype=dtype) - for key in list(original_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_inplace(original_state_dict, key, new_key) - - for key in list(original_state_dict.keys()): - for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, original_state_dict) - - vae.load_state_dict(original_state_dict, strict=True) + converted = convert_cogvideox_vae_checkpoint_to_diffusers(original_state_dict, vae.config) + vae.load_state_dict(converted, strict=True) return vae diff --git a/scripts/convert_cogview3_to_diffusers.py b/scripts/recipes/cogview3.py similarity index 52% rename from scripts/convert_cogview3_to_diffusers.py rename to scripts/recipes/cogview3.py index 605555ebdbef..94bc2a2171b7 100644 --- a/scripts/convert_cogview3_to_diffusers.py +++ b/scripts/recipes/cogview3.py @@ -5,7 +5,7 @@ with the Diffusers library. Example usage: - python scripts/convert_cogview3_to_diffusers.py \ + python scripts/recipes/cogview3.py \ --transformer_checkpoint_path 'your path/cogview3plus_3b/1/mp_rank_00_model_states.pt' \ --vae_checkpoint_path 'your path/3plus_ae/imagekl_ch16.pt' \ --output_path "/raid/yiyi/cogview3_diffusers" \ @@ -25,19 +25,14 @@ """ import argparse -from contextlib import nullcontext import torch -from accelerate import init_empty_weights from transformers import T5EncoderModel, T5Tokenizer from diffusers import AutoencoderKL, CogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModel -from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint -from diffusers.utils.import_utils import is_accelerate_available +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint -CTX = init_empty_weights if is_accelerate_available() else nullcontext - TOKENIZER_MAX_LENGTH = 224 parser = argparse.ArgumentParser() @@ -53,106 +48,17 @@ # this is specific to `AdaLayerNormContinuous`: # diffusers implementation split the linear projection into the scale, shift while CogView3 split it tino shift, scale -def swap_scale_shift(weight, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight def convert_cogview3_transformer_checkpoint_to_diffusers(ckpt_path): - original_state_dict = torch.load(ckpt_path, map_location="cpu") - original_state_dict = original_state_dict["module"] - original_state_dict = {k.replace("model.diffusion_model.", ""): v for k, v in original_state_dict.items()} - - new_state_dict = {} - - # Convert patch_embed - new_state_dict["patch_embed.proj.weight"] = original_state_dict.pop("mixins.patch_embed.proj.weight") - new_state_dict["patch_embed.proj.bias"] = original_state_dict.pop("mixins.patch_embed.proj.bias") - new_state_dict["patch_embed.text_proj.weight"] = original_state_dict.pop("mixins.patch_embed.text_proj.weight") - new_state_dict["patch_embed.text_proj.bias"] = original_state_dict.pop("mixins.patch_embed.text_proj.bias") - - # Convert time_condition_embed - new_state_dict["time_condition_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "time_embed.0.weight" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "time_embed.0.bias" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "time_embed.2.weight" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "time_embed.2.bias" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_1.weight"] = original_state_dict.pop( - "label_emb.0.0.weight" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_1.bias"] = original_state_dict.pop( - "label_emb.0.0.bias" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_2.weight"] = original_state_dict.pop( - "label_emb.0.2.weight" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_2.bias"] = original_state_dict.pop( - "label_emb.0.2.bias" - ) - - # Convert transformer blocks - for i in range(30): - block_prefix = f"transformer_blocks.{i}." - old_prefix = f"transformer.layers.{i}." - adaln_prefix = f"mixins.adaln.adaln_modules.{i}." - - new_state_dict[block_prefix + "norm1.linear.weight"] = original_state_dict.pop(adaln_prefix + "1.weight") - new_state_dict[block_prefix + "norm1.linear.bias"] = original_state_dict.pop(adaln_prefix + "1.bias") - - qkv_weight = original_state_dict.pop(old_prefix + "attention.query_key_value.weight") - qkv_bias = original_state_dict.pop(old_prefix + "attention.query_key_value.bias") - q, k, v = qkv_weight.chunk(3, dim=0) - q_bias, k_bias, v_bias = qkv_bias.chunk(3, dim=0) - - new_state_dict[block_prefix + "attn1.to_q.weight"] = q - new_state_dict[block_prefix + "attn1.to_q.bias"] = q_bias - new_state_dict[block_prefix + "attn1.to_k.weight"] = k - new_state_dict[block_prefix + "attn1.to_k.bias"] = k_bias - new_state_dict[block_prefix + "attn1.to_v.weight"] = v - new_state_dict[block_prefix + "attn1.to_v.bias"] = v_bias - - new_state_dict[block_prefix + "attn1.to_out.0.weight"] = original_state_dict.pop( - old_prefix + "attention.dense.weight" - ) - new_state_dict[block_prefix + "attn1.to_out.0.bias"] = original_state_dict.pop( - old_prefix + "attention.dense.bias" - ) - - new_state_dict[block_prefix + "ff.net.0.proj.weight"] = original_state_dict.pop( - old_prefix + "mlp.dense_h_to_4h.weight" - ) - new_state_dict[block_prefix + "ff.net.0.proj.bias"] = original_state_dict.pop( - old_prefix + "mlp.dense_h_to_4h.bias" - ) - new_state_dict[block_prefix + "ff.net.2.weight"] = original_state_dict.pop( - old_prefix + "mlp.dense_4h_to_h.weight" - ) - new_state_dict[block_prefix + "ff.net.2.bias"] = original_state_dict.pop(old_prefix + "mlp.dense_4h_to_h.bias") - - # Convert final norm and projection - new_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("mixins.final_layer.adaln.1.weight"), dim=0 - ) - new_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("mixins.final_layer.adaln.1.bias"), dim=0 - ) - new_state_dict["proj_out.weight"] = original_state_dict.pop("mixins.final_layer.linear.weight") - new_state_dict["proj_out.bias"] = original_state_dict.pop("mixins.final_layer.linear.bias") - - return new_state_dict + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + state = state["module"] + return convert_component_checkpoint(state, {}, "CogView3PlusTransformer2DModel") def convert_cogview3_vae_checkpoint_to_diffusers(ckpt_path, vae_config): original_state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"] - return convert_ldm_vae_checkpoint(original_state_dict, vae_config) + return convert_component_checkpoint(original_state_dict, vae_config, "AutoencoderKL") def main(args): diff --git a/scripts/convert_cogview4_to_diffusers.py b/scripts/recipes/cogview4.py similarity index 51% rename from scripts/convert_cogview4_to_diffusers.py rename to scripts/recipes/cogview4.py index b6d01c797aeb..a7c9c99dfd25 100644 --- a/scripts/convert_cogview4_to_diffusers.py +++ b/scripts/recipes/cogview4.py @@ -6,7 +6,7 @@ with the Diffusers library. Example usage: - python scripts/convert_cogview4_to_diffusers.py \ + python scripts/recipes/cogview4.py \ --transformer_checkpoint_path 'your path/cogview4_6b/1/mp_rank_00_model_states.pt' \ --vae_checkpoint_path 'your path/cogview4_6b/imagekl_ch16.pt' \ --output_path "THUDM/CogView4-6B" \ @@ -26,19 +26,14 @@ """ import argparse -from contextlib import nullcontext import torch -from accelerate import init_empty_weights from transformers import GlmForCausalLM, PreTrainedTokenizerFast from diffusers import AutoencoderKL, CogView4Pipeline, CogView4Transformer2DModel, FlowMatchEulerDiscreteScheduler -from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint -from diffusers.utils.import_utils import is_accelerate_available +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint -CTX = init_empty_weights if is_accelerate_available() else nullcontext - parser = argparse.ArgumentParser() parser.add_argument("--transformer_checkpoint_path", default=None, type=str) parser.add_argument("--vae_checkpoint_path", default=None, type=str) @@ -52,115 +47,17 @@ # this is specific to `AdaLayerNormContinuous`: # diffusers implementation split the linear projection into the scale, shift while CogView4 split it tino shift, scale -def swap_scale_shift(weight, dim): - """ - Swap the scale and shift components in the weight tensor. - - Args: - weight (torch.Tensor): The original weight tensor. - dim (int): The dimension along which to split. - - Returns: - torch.Tensor: The modified weight tensor with scale and shift swapped. - """ - shift, scale = weight.chunk(2, dim=dim) - new_weight = torch.cat([scale, shift], dim=dim) - return new_weight def convert_cogview4_transformer_checkpoint_to_diffusers(ckpt_path): - original_state_dict = torch.load(ckpt_path, map_location="cpu") - original_state_dict = original_state_dict["module"] - original_state_dict = {k.replace("model.diffusion_model.", ""): v for k, v in original_state_dict.items()} - - new_state_dict = {} - - # Convert patch_embed - new_state_dict["patch_embed.proj.weight"] = original_state_dict.pop("mixins.patch_embed.proj.weight") - new_state_dict["patch_embed.proj.bias"] = original_state_dict.pop("mixins.patch_embed.proj.bias") - new_state_dict["patch_embed.text_proj.weight"] = original_state_dict.pop("mixins.patch_embed.text_proj.weight") - new_state_dict["patch_embed.text_proj.bias"] = original_state_dict.pop("mixins.patch_embed.text_proj.bias") - - # Convert time_condition_embed - new_state_dict["time_condition_embed.timestep_embedder.linear_1.weight"] = original_state_dict.pop( - "time_embed.0.weight" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_1.bias"] = original_state_dict.pop( - "time_embed.0.bias" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_2.weight"] = original_state_dict.pop( - "time_embed.2.weight" - ) - new_state_dict["time_condition_embed.timestep_embedder.linear_2.bias"] = original_state_dict.pop( - "time_embed.2.bias" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_1.weight"] = original_state_dict.pop( - "label_emb.0.0.weight" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_1.bias"] = original_state_dict.pop( - "label_emb.0.0.bias" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_2.weight"] = original_state_dict.pop( - "label_emb.0.2.weight" - ) - new_state_dict["time_condition_embed.condition_embedder.linear_2.bias"] = original_state_dict.pop( - "label_emb.0.2.bias" - ) - - # Convert transformer blocks, for cogview4 is 28 blocks - for i in range(28): - block_prefix = f"transformer_blocks.{i}." - old_prefix = f"transformer.layers.{i}." - adaln_prefix = f"mixins.adaln.adaln_modules.{i}." - new_state_dict[block_prefix + "norm1.linear.weight"] = original_state_dict.pop(adaln_prefix + "1.weight") - new_state_dict[block_prefix + "norm1.linear.bias"] = original_state_dict.pop(adaln_prefix + "1.bias") - - qkv_weight = original_state_dict.pop(old_prefix + "attention.query_key_value.weight") - qkv_bias = original_state_dict.pop(old_prefix + "attention.query_key_value.bias") - q, k, v = qkv_weight.chunk(3, dim=0) - q_bias, k_bias, v_bias = qkv_bias.chunk(3, dim=0) - - new_state_dict[block_prefix + "attn1.to_q.weight"] = q - new_state_dict[block_prefix + "attn1.to_q.bias"] = q_bias - new_state_dict[block_prefix + "attn1.to_k.weight"] = k - new_state_dict[block_prefix + "attn1.to_k.bias"] = k_bias - new_state_dict[block_prefix + "attn1.to_v.weight"] = v - new_state_dict[block_prefix + "attn1.to_v.bias"] = v_bias - - new_state_dict[block_prefix + "attn1.to_out.0.weight"] = original_state_dict.pop( - old_prefix + "attention.dense.weight" - ) - new_state_dict[block_prefix + "attn1.to_out.0.bias"] = original_state_dict.pop( - old_prefix + "attention.dense.bias" - ) - - new_state_dict[block_prefix + "ff.net.0.proj.weight"] = original_state_dict.pop( - old_prefix + "mlp.dense_h_to_4h.weight" - ) - new_state_dict[block_prefix + "ff.net.0.proj.bias"] = original_state_dict.pop( - old_prefix + "mlp.dense_h_to_4h.bias" - ) - new_state_dict[block_prefix + "ff.net.2.weight"] = original_state_dict.pop( - old_prefix + "mlp.dense_4h_to_h.weight" - ) - new_state_dict[block_prefix + "ff.net.2.bias"] = original_state_dict.pop(old_prefix + "mlp.dense_4h_to_h.bias") - - # Convert final norm and projection - new_state_dict["norm_out.linear.weight"] = swap_scale_shift( - original_state_dict.pop("mixins.final_layer.adaln.1.weight"), dim=0 - ) - new_state_dict["norm_out.linear.bias"] = swap_scale_shift( - original_state_dict.pop("mixins.final_layer.adaln.1.bias"), dim=0 - ) - new_state_dict["proj_out.weight"] = original_state_dict.pop("mixins.final_layer.linear.weight") - new_state_dict["proj_out.bias"] = original_state_dict.pop("mixins.final_layer.linear.bias") - - return new_state_dict + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + state = state["module"] + return convert_component_checkpoint(state, {}, "CogView4Transformer2DModel") def convert_cogview4_vae_checkpoint_to_diffusers(ckpt_path, vae_config): original_state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"] - return convert_ldm_vae_checkpoint(original_state_dict, vae_config) + return convert_component_checkpoint(original_state_dict, vae_config, "AutoencoderKL") def main(args): diff --git a/scripts/convert_cogview4_to_diffusers_megatron.py b/scripts/recipes/cogview4_megatron.py similarity index 58% rename from scripts/convert_cogview4_to_diffusers_megatron.py rename to scripts/recipes/cogview4_megatron.py index 8faeccb13888..32dfa425dfd4 100644 --- a/scripts/convert_cogview4_to_diffusers_megatron.py +++ b/scripts/recipes/cogview4_megatron.py @@ -2,7 +2,7 @@ Convert a CogView4 checkpoint from Megatron to the Diffusers format. Example usage: - python scripts/convert_cogview4_to_diffusers.py \ + python scripts/build_pipeline.py cogview4_megatron \ --transformer_checkpoint_path 'your path/cogview4_6b/mp_rank_00/model_optim_rng.pt' \ --vae_checkpoint_path 'your path/cogview4_6b/imagekl_ch16.pt' \ --output_path "THUDM/CogView4-6B" \ @@ -24,7 +24,6 @@ import argparse import torch -from tqdm import tqdm from transformers import GlmModel, PreTrainedTokenizerFast from diffusers import ( @@ -34,7 +33,7 @@ CogView4Transformer2DModel, FlowMatchEulerDiscreteScheduler, ) -from diffusers.loaders.single_file_utils import convert_ldm_vae_checkpoint +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint parser = argparse.ArgumentParser() @@ -128,125 +127,15 @@ args = parser.parse_args() -def swap_scale_shift(weight, dim): - """ - Swap the scale and shift components in the weight tensor. - - Args: - weight (torch.Tensor): The original weight tensor. - dim (int): The dimension along which to split. - - Returns: - torch.Tensor: The modified weight tensor with scale and shift swapped. - """ - shift, scale = weight.chunk(2, dim=dim) - new_weight = torch.cat([scale, shift], dim=dim) - return new_weight - - -def convert_megatron_transformer_checkpoint_to_diffusers( - ckpt_path: str, - num_layers: int, - num_heads: int, - hidden_size: int, -): - """ - Convert a Megatron Transformer checkpoint to Diffusers format. - - Args: - ckpt_path (str): Path to the Megatron Transformer checkpoint. - num_layers (int): Number of Transformer layers. - num_heads (int): Number of attention heads. - hidden_size (int): Hidden size of the Transformer. - - Returns: - dict: The converted state dictionary compatible with Diffusers. - """ - ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) - mega = ckpt["model"] - - new_state_dict = {} - - # Patch Embedding - new_state_dict["patch_embed.proj.weight"] = mega["encoder_expand_linear.weight"].reshape( - hidden_size, 128 if args.control else 64 - ) - new_state_dict["patch_embed.proj.bias"] = mega["encoder_expand_linear.bias"] - new_state_dict["patch_embed.text_proj.weight"] = mega["text_projector.weight"] - new_state_dict["patch_embed.text_proj.bias"] = mega["text_projector.bias"] - - # Time Condition Embedding - new_state_dict["time_condition_embed.timestep_embedder.linear_1.weight"] = mega[ - "time_embedding.time_embed.0.weight" - ] - new_state_dict["time_condition_embed.timestep_embedder.linear_1.bias"] = mega["time_embedding.time_embed.0.bias"] - new_state_dict["time_condition_embed.timestep_embedder.linear_2.weight"] = mega[ - "time_embedding.time_embed.2.weight" - ] - new_state_dict["time_condition_embed.timestep_embedder.linear_2.bias"] = mega["time_embedding.time_embed.2.bias"] - - new_state_dict["time_condition_embed.condition_embedder.linear_1.weight"] = mega[ - "label_embedding.label_embed.0.weight" - ] - new_state_dict["time_condition_embed.condition_embedder.linear_1.bias"] = mega[ - "label_embedding.label_embed.0.bias" - ] - new_state_dict["time_condition_embed.condition_embedder.linear_2.weight"] = mega[ - "label_embedding.label_embed.2.weight" - ] - new_state_dict["time_condition_embed.condition_embedder.linear_2.bias"] = mega[ - "label_embedding.label_embed.2.bias" - ] - - # Convert each Transformer layer - for i in tqdm(range(num_layers), desc="Converting layers (Megatron->Diffusers)"): - block_prefix = f"transformer_blocks.{i}." - - # AdaLayerNorm - new_state_dict[block_prefix + "norm1.linear.weight"] = mega[f"decoder.layers.{i}.adaln.weight"] - new_state_dict[block_prefix + "norm1.linear.bias"] = mega[f"decoder.layers.{i}.adaln.bias"] - qkv_weight = mega[f"decoder.layers.{i}.self_attention.linear_qkv.weight"] - qkv_bias = mega[f"decoder.layers.{i}.self_attention.linear_qkv.bias"] - - # Reshape to match SAT logic - qkv_weight = qkv_weight.view(num_heads, 3, hidden_size // num_heads, hidden_size) - qkv_weight = qkv_weight.permute(1, 0, 2, 3).reshape(3 * hidden_size, hidden_size) - - qkv_bias = qkv_bias.view(num_heads, 3, hidden_size // num_heads) - qkv_bias = qkv_bias.permute(1, 0, 2).reshape(3 * hidden_size) - - # Assign to Diffusers keys - q, k, v = torch.chunk(qkv_weight, 3, dim=0) - qb, kb, vb = torch.chunk(qkv_bias, 3, dim=0) - - new_state_dict[block_prefix + "attn1.to_q.weight"] = q - new_state_dict[block_prefix + "attn1.to_q.bias"] = qb - new_state_dict[block_prefix + "attn1.to_k.weight"] = k - new_state_dict[block_prefix + "attn1.to_k.bias"] = kb - new_state_dict[block_prefix + "attn1.to_v.weight"] = v - new_state_dict[block_prefix + "attn1.to_v.bias"] = vb - - # Attention Output - new_state_dict[block_prefix + "attn1.to_out.0.weight"] = mega[ - f"decoder.layers.{i}.self_attention.linear_proj.weight" - ] - new_state_dict[block_prefix + "attn1.to_out.0.bias"] = mega[ - f"decoder.layers.{i}.self_attention.linear_proj.bias" - ] - - # MLP - new_state_dict[block_prefix + "ff.net.0.proj.weight"] = mega[f"decoder.layers.{i}.mlp.linear_fc1.weight"] - new_state_dict[block_prefix + "ff.net.0.proj.bias"] = mega[f"decoder.layers.{i}.mlp.linear_fc1.bias"] - new_state_dict[block_prefix + "ff.net.2.weight"] = mega[f"decoder.layers.{i}.mlp.linear_fc2.weight"] - new_state_dict[block_prefix + "ff.net.2.bias"] = mega[f"decoder.layers.{i}.mlp.linear_fc2.bias"] - - # Final Layers - new_state_dict["norm_out.linear.weight"] = swap_scale_shift(mega["adaln_final.weight"], dim=0) - new_state_dict["norm_out.linear.bias"] = swap_scale_shift(mega["adaln_final.bias"], dim=0) - new_state_dict["proj_out.weight"] = mega["output_projector.weight"] - new_state_dict["proj_out.bias"] = mega["output_projector.bias"] - - return new_state_dict +def convert_megatron_transformer_checkpoint_to_diffusers(ckpt_path, num_layers, num_heads, hidden_size): + state = torch.load(ckpt_path, map_location="cpu", weights_only=False)["model"] + config = { + "num_layers": num_layers, + "num_attention_heads": num_heads, + "attention_head_dim": hidden_size // num_heads, + "original_format": "megatron", + } + return convert_component_checkpoint(state, config, "CogView4Transformer2DModel") def convert_cogview4_vae_checkpoint_to_diffusers(ckpt_path, vae_config): @@ -261,7 +150,7 @@ def convert_cogview4_vae_checkpoint_to_diffusers(ckpt_path, vae_config): dict: The converted VAE state dictionary compatible with Diffusers. """ original_state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False)["state_dict"] - return convert_ldm_vae_checkpoint(original_state_dict, vae_config) + return convert_component_checkpoint(original_state_dict, vae_config, "AutoencoderKL") def main(args): diff --git a/scripts/recipes/consistency.py b/scripts/recipes/consistency.py new file mode 100644 index 000000000000..7941fb5c3f9e --- /dev/null +++ b/scripts/recipes/consistency.py @@ -0,0 +1,87 @@ +import argparse +import os + +import torch + +from diffusers import ( + CMStochasticIterativeScheduler, + ConsistencyModelPipeline, + UNet2DModel, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.consistency import ( + CD_SCHEDULER_CONFIG, + CT_IMAGENET_64_SCHEDULER_CONFIG, + CT_LSUN_256_SCHEDULER_CONFIG, + IMAGENET_64_UNET_CONFIG, + LSUN_256_UNET_CONFIG, + TEST_UNET_CONFIG, +) + + +def str2bool(v): + """ + https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse + """ + if isinstance(v, bool): + return v + if v.lower() in ("yes", "true", "t", "y", "1"): + return True + elif v.lower() in ("no", "false", "f", "n", "0"): + return False + else: + raise argparse.ArgumentTypeError("boolean value expected") + + +def con_pt_to_diffuser(checkpoint_path, config): + state = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + return get_conversion("UNet2DModel", {**config, "original_format": "consistency"}).to_diffusers(state) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--unet_path", default=None, type=str, required=True, help="Path to the unet.pt to convert.") + parser.add_argument( + "--dump_path", default=None, type=str, required=True, help="Path to output the converted UNet model." + ) + parser.add_argument("--class_cond", default=True, type=str, help="Whether the model is class-conditional.") + + args = parser.parse_args() + args.class_cond = str2bool(args.class_cond) + + ckpt_name = os.path.basename(args.unet_path) + print(f"Checkpoint: {ckpt_name}") + + # Get U-Net config + if "imagenet64" in ckpt_name: + unet_config = IMAGENET_64_UNET_CONFIG + elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)): + unet_config = LSUN_256_UNET_CONFIG + elif "test" in ckpt_name: + unet_config = TEST_UNET_CONFIG + else: + raise ValueError(f"Checkpoint type {ckpt_name} is not currently supported.") + + if not args.class_cond: + unet_config["num_class_embeds"] = None + + converted_unet_ckpt = con_pt_to_diffuser(args.unet_path, unet_config) + + image_unet = UNet2DModel(**unet_config) + image_unet.load_state_dict(converted_unet_ckpt) + + # Get scheduler config + if "cd" in ckpt_name or "test" in ckpt_name: + scheduler_config = CD_SCHEDULER_CONFIG + elif "ct" in ckpt_name and "imagenet64" in ckpt_name: + scheduler_config = CT_IMAGENET_64_SCHEDULER_CONFIG + elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)): + scheduler_config = CT_LSUN_256_SCHEDULER_CONFIG + else: + raise ValueError(f"Checkpoint type {ckpt_name} is not currently supported.") + + cm_scheduler = CMStochasticIterativeScheduler(**scheduler_config) + + consistency_model = ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler) + consistency_model.save_pretrained(args.dump_path) diff --git a/scripts/recipes/consistency_decoder.py b/scripts/recipes/consistency_decoder.py new file mode 100644 index 000000000000..9b74929446fc --- /dev/null +++ b/scripts/recipes/consistency_decoder.py @@ -0,0 +1,36 @@ +"""Assemble a consistency decoder VAE from the published TorchScript decoder and a Diffusers SD encoder.""" + +import argparse + +import torch + +from diffusers import AutoencoderKL, ConsistencyDecoderVAE +from diffusers.loaders.conversion import get_conversion + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--decoder", required=True, help="Local published decoder.pt TorchScript file") + parser.add_argument("--encoder", required=True, help="Diffusers AutoencoderKL directory containing the SD encoder") + parser.add_argument("--output", required=True) + parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float16") + args = parser.parse_args() + encoder = AutoencoderKL.from_pretrained(args.encoder) + original_encoder = get_conversion("AutoencoderKL", dict(encoder.config)).to_original(encoder.state_dict()) + state = {key: value for key, value in original_encoder.items() if key.startswith(("encoder.", "quant_conv."))} + decoder = torch.jit.load(args.decoder, map_location="cpu") + state.update({"decoder." + key: value for key, value in decoder.state_dict().items()}) + config = { + "original_format": "consistency_decoder_jit", + "scaling_factor": encoder.config.scaling_factor, + "latent_channels": encoder.config.latent_channels, + "encoder_block_out_channels": encoder.config.block_out_channels, + "encoder_down_block_types": encoder.config.down_block_types, + "encoder_layers_per_block": encoder.config.layers_per_block, + } + model = ConsistencyDecoderVAE.from_single_file(state, config=config, torch_dtype=getattr(torch, args.dtype)) + model.save_pretrained(args.output) + + +if __name__ == "__main__": + main() diff --git a/scripts/convert_original_controlnet_to_diffusers.py b/scripts/recipes/controlnet.py similarity index 100% rename from scripts/convert_original_controlnet_to_diffusers.py rename to scripts/recipes/controlnet.py diff --git a/scripts/recipes/cosmos.py b/scripts/recipes/cosmos.py new file mode 100644 index 000000000000..df00159f14b1 --- /dev/null +++ b/scripts/recipes/cosmos.py @@ -0,0 +1,456 @@ +""" +# Cosmos 2 Predict + +Download checkpoint +```bash +hf download nvidia/Cosmos-Predict2-2B-Text2Image +``` + +convert checkpoint +```bash +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2-2B-Text2Image/snapshots/acdb5fde992a73ef0355f287977d002cbfd127e0/model.pt + +python scripts/recipes/cosmos.py \ + --transformer_ckpt_path $transformer_ckpt_path \ + --transformer_type Cosmos-2.0-Diffusion-2B-Text2Image \ + --text_encoder_path google-t5/t5-11b \ + --tokenizer_path google-t5/t5-11b \ + --vae_type wan2.1 \ + --output_path converted/cosmos-p2-t2i-2b \ + --save_pipeline +``` + +# Cosmos 2.5 Predict + +Download checkpoint +```bash +hf download nvidia/Cosmos-Predict2.5-2B +``` + +Convert checkpoint +```bash +# pre-trained +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/pre-trained/d20b7120-df3e-4911-919d-db6e08bad31c_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Predict-Base-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/2b/d20b7120-df3e-4911-919d-db6e08bad31c \ + --save_pipeline + +# post-trained +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-2B/snapshots/865baf084d4c9e850eac59a021277d5a9b9e8b63/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Predict-Base-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/2b/81edfebe-bd6a-4039-8c1d-737df1a790bf \ + --save_pipeline +``` + +## 14B + +```bash +hf download nvidia/Cosmos-Predict2.5-14B +``` + +```bash +# pre-trained +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/pre-trained/54937b8c-29de-4f04-862c-e67b04ec41e8_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Predict-Base-14B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/14b/54937b8c-29de-4f04-862c-e67b04ec41e8/ \ + --save_pipeline + +# post-trained +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Predict2.5-14B/snapshots/71ebf3e8af30ecfe440bf0481115975fcc052b46/base/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Predict-Base-14B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/14b/e21d2a49-4747-44c8-ba44-9f6f9243715f/ \ + --save_pipeline +``` + +# Cosmos 2.5 Transfer + +Download checkpoint +```bash +hf download nvidia/Cosmos-Transfer2.5-2B +``` + +Convert checkpoint +```bash +# depth +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/depth/626e6618-bfcd-4d9a-a077-1409e2ce353f_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/depth/pipeline \ + --save_pipeline + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/depth/models + +# edge +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/edge/61f5694b-0ad5-4ecd-8ad7-c8545627d125_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/edge/pipeline \ + --save_pipeline + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/edge/models + +# blur +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/blur/ba2f44f2-c726-4fe7-949f-597069d9b91c_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/blur/pipeline \ + --save_pipeline + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/blur/models + +# seg +transformer_ckpt_path=~/.cache/huggingface/hub/models--nvidia--Cosmos-Transfer2.5-2B/snapshots/eb5325b77d358944da58a690157dd2b8071bbf85/general/seg/5136ef49-6d8d-42e8-8abf-7dac722a304a_ema_bf16.pt + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/seg/pipeline \ + --save_pipeline + +python scripts/recipes/cosmos.py \ + --transformer_type Cosmos-2.5-Transfer-General-2B \ + --transformer_ckpt_path $transformer_ckpt_path \ + --vae_type wan2.1 \ + --output_path converted/transfer/2b/general/seg/models +``` +""" + +import argparse +import pathlib +from typing import Any, Dict + +import torch +from huggingface_hub import snapshot_download +from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration, T5EncoderModel, T5TokenizerFast + +from diffusers import ( + AutoencoderKLCosmos, + AutoencoderKLWan, + Cosmos2TextToImagePipeline, + Cosmos2VideoToWorldPipeline, + CosmosControlNetModel, + CosmosTextToWorldPipeline, + CosmosTransformer3DModel, + CosmosVideoToWorldPipeline, + EDMEulerScheduler, + FlowMatchEulerDiscreteScheduler, + UniPCMultistepScheduler, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.cosmos import CONTROLNET_CONFIGS, TRANSFORMER_CONFIGS, VAE_CONFIGS +from diffusers.pipelines.cosmos.pipeline_cosmos2_5_predict import Cosmos2_5_PredictBasePipeline +from diffusers.pipelines.cosmos.pipeline_cosmos2_5_transfer import Cosmos2_5_TransferPipeline + + +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: + state_dict = saved_dict + if "model" in saved_dict.keys(): + state_dict = state_dict["model"] + if "module" in saved_dict.keys(): + state_dict = state_dict["module"] + if "state_dict" in saved_dict.keys(): + state_dict = state_dict["state_dict"] + return state_dict + + +def convert_transformer(transformer_type: str, state_dict=None, weights_only=True): + config = dict(TRANSFORMER_CONFIGS[transformer_type]) + config["original_format"] = "cosmos1" if "Cosmos-1.0" in transformer_type else "cosmos2" + return CosmosTransformer3DModel.from_single_file(state_dict, config=config) + + +def convert_controlnet(transformer_type: str, control_state_dict, base_state_dict, weights_only=True): + config = CONTROLNET_CONFIGS[transformer_type] + conversion = get_conversion("CosmosControlNetModel", config) + base_config = {**TRANSFORMER_CONFIGS[transformer_type], "original_format": "cosmos2"} + base = get_conversion("CosmosTransformer3DModel", base_config).to_original(base_state_dict) + state = {key.removeprefix("net."): value for key, value in control_state_dict.items()} + # The original ControlNet omits modules shared with the separately loaded base transformer. + for key in conversion.original_keys: + if key.startswith("base."): + state[key] = base[key.removeprefix("base.")] + return CosmosControlNetModel.from_single_file(state, config=config) + + +def convert_vae(vae_type: str): + model_name = VAE_CONFIGS[vae_type]["name"] + snapshot_directory = snapshot_download(model_name, repo_type="model") + directory = pathlib.Path(snapshot_directory) + + autoencoder_file = directory / "autoencoder.jit" + mean_std_file = directory / "mean_std.pt" + + original_state_dict = torch.jit.load(autoencoder_file.as_posix()).state_dict() + if mean_std_file.exists(): + mean_std = torch.load(mean_std_file, map_location="cpu", weights_only=True) + else: + mean_std = (None, None) + + config = dict(VAE_CONFIGS[vae_type]["diffusers_config"]) + config.update( + { + "latents_mean": mean_std[0].detach().cpu().numpy().tolist(), + "latents_std": mean_std[1].detach().cpu().numpy().tolist(), + } + ) + return AutoencoderKLCosmos.from_single_file(original_state_dict, config=config) + + +def save_pipeline_cosmos_1_0(args, transformer, vae): + text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16) + tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path) + # The original code initializes EDM config with sigma_min=0.0002, but does not make use of it anywhere directly. + # So, the sigma_min values that is used is the default value of 0.002. + scheduler = EDMEulerScheduler( + sigma_min=0.002, + sigma_max=80, + sigma_data=0.5, + sigma_schedule="karras", + num_train_timesteps=1000, + prediction_type="epsilon", + rho=7.0, + final_sigmas_type="sigma_min", + ) + + pipe_cls = CosmosTextToWorldPipeline if "Text2World" in args.transformer_type else CosmosVideoToWorldPipeline + pipe = pipe_cls( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + vae=vae, + scheduler=scheduler, + safety_checker=lambda *args, **kwargs: None, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + +def save_pipeline_cosmos_2_0(args, transformer, vae): + text_encoder = T5EncoderModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.bfloat16) + tokenizer = T5TokenizerFast.from_pretrained(args.tokenizer_path) + + scheduler = FlowMatchEulerDiscreteScheduler(use_karras_sigmas=True) + + pipe_cls = Cosmos2TextToImagePipeline if "Text2Image" in args.transformer_type else Cosmos2VideoToWorldPipeline + pipe = pipe_cls( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + vae=vae, + scheduler=scheduler, + safety_checker=lambda *args, **kwargs: None, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + +def save_pipeline_cosmos2_5_predict(args, transformer, vae): + text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" + tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" + + text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + text_encoder_path, torch_dtype="auto", device_map="cpu" + ) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + scheduler = UniPCMultistepScheduler( + use_karras_sigmas=True, + use_flow_sigmas=True, + prediction_type="flow_prediction", + sigma_max=200.0, + sigma_min=0.01, + ) + + pipe = Cosmos2_5_PredictBasePipeline( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + vae=vae, + scheduler=scheduler, + safety_checker=lambda *args, **kwargs: None, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + +def save_pipeline_cosmos2_5_transfer(args, transformer, controlnet, vae): + text_encoder_path = args.text_encoder_path or "nvidia/Cosmos-Reason1-7B" + tokenizer_path = args.tokenizer_path or "Qwen/Qwen2.5-VL-7B-Instruct" + + text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + text_encoder_path, torch_dtype="auto", device_map="cpu" + ) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + + scheduler = UniPCMultistepScheduler( + use_karras_sigmas=True, + use_flow_sigmas=True, + prediction_type="flow_prediction", + sigma_max=200.0, + sigma_min=0.01, + ) + + pipe = Cosmos2_5_TransferPipeline( + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + controlnet=controlnet, + vae=vae, + scheduler=scheduler, + safety_checker=lambda *args, **kwargs: None, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--transformer_type", type=str, default=None, choices=list(TRANSFORMER_CONFIGS.keys())) + parser.add_argument( + "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" + ) + parser.add_argument( + "--vae_type", type=str, default="wan2.1", choices=["wan2.1", *list(VAE_CONFIGS.keys())], help="Type of VAE" + ) + parser.add_argument("--text_encoder_path", type=str, default=None) + parser.add_argument("--tokenizer_path", type=str, default=None) + parser.add_argument("--save_pipeline", action="store_true") + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") + parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.") + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +if __name__ == "__main__": + args = get_args() + + transformer = None + controlnet = None + dtype = DTYPE_MAPPING[args.dtype] + + if args.save_pipeline: + assert args.transformer_ckpt_path is not None + assert args.vae_type is not None + + raw_state_dict = None + if args.transformer_ckpt_path is not None: + weights_only = "Cosmos-1.0" in args.transformer_type + raw_state_dict = get_state_dict( + torch.load(args.transformer_ckpt_path, map_location="cpu", weights_only=weights_only) + ) + + if raw_state_dict is not None: + if "Transfer" in args.transformer_type: + base_state_dict = {} + control_state_dict = {} + for k, v in raw_state_dict.items(): + plain_key = k.removeprefix("net.") if k.startswith("net.") else k + if "control" in plain_key.lower(): + control_state_dict[k] = v + else: + base_state_dict[k] = v + assert len(base_state_dict.keys() & control_state_dict.keys()) == 0 + + # Convert transformer first to get the processed base state dict + transformer = convert_transformer( + args.transformer_type, state_dict=base_state_dict, weights_only=weights_only + ) + transformer = transformer.to(dtype=dtype) + + # Get converted transformer state dict to copy shared weights to controlnet + converted_base_state_dict = transformer.state_dict() + + # Convert controlnet with both control-specific and shared weights from transformer + controlnet = convert_controlnet( + args.transformer_type, control_state_dict, converted_base_state_dict, weights_only=weights_only + ) + controlnet = controlnet.to(dtype=dtype) + + if not args.save_pipeline: + transformer.save_pretrained( + pathlib.Path(args.output_path) / "transformer", safe_serialization=True, max_shard_size="5GB" + ) + controlnet.save_pretrained( + pathlib.Path(args.output_path) / "controlnet", safe_serialization=True, max_shard_size="5GB" + ) + else: + transformer = convert_transformer( + args.transformer_type, state_dict=raw_state_dict, weights_only=weights_only + ) + transformer = transformer.to(dtype=dtype) + if not args.save_pipeline: + transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.vae_type is not None: + if "Cosmos-1.0" in args.transformer_type: + vae = convert_vae(args.vae_type) + elif "Cosmos-2.0" in args.transformer_type or "Cosmos-2.5" in args.transformer_type: + vae = AutoencoderKLWan.from_pretrained( + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", subfolder="vae", torch_dtype=torch.float32 + ) + else: + raise AssertionError(f"{args.transformer_type} not supported") + + if not args.save_pipeline: + vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + else: + vae = None + + if args.save_pipeline: + if "Cosmos-1.0" in args.transformer_type: + assert args.text_encoder_path is not None + assert args.tokenizer_path is not None + save_pipeline_cosmos_1_0(args, transformer, vae) + elif "Cosmos-2.0" in args.transformer_type: + assert args.text_encoder_path is not None + assert args.tokenizer_path is not None + save_pipeline_cosmos_2_0(args, transformer, vae) + elif "Cosmos-2.5" in args.transformer_type: + if "Predict" in args.transformer_type: + save_pipeline_cosmos2_5_predict(args, transformer, vae) + elif "Transfer" in args.transformer_type: + save_pipeline_cosmos2_5_transfer(args, transformer, None, vae) + else: + raise AssertionError(f"{args.transformer_type} not supported") + else: + raise AssertionError(f"{args.transformer_type} not supported") diff --git a/scripts/recipes/dance_diffusion.py b/scripts/recipes/dance_diffusion.py new file mode 100644 index 000000000000..8dafb20e1aea --- /dev/null +++ b/scripts/recipes/dance_diffusion.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +import argparse +import math +import os +from copy import deepcopy + +import requests +import torch +from audio_diffusion.models import DiffusionAttnUnet1D +from diffusion import sampling +from torch import nn + +from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNet1DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.utils.constants import DIFFUSERS_REQUEST_TIMEOUT + + +MODELS_MAP = { + "gwf-440k": { + "url": "https://model-server.zqevans2.workers.dev/gwf-440k.ckpt", + "sample_rate": 48000, + "sample_size": 65536, + }, + "jmann-small-190k": { + "url": "https://model-server.zqevans2.workers.dev/jmann-small-190k.ckpt", + "sample_rate": 48000, + "sample_size": 65536, + }, + "jmann-large-580k": { + "url": "https://model-server.zqevans2.workers.dev/jmann-large-580k.ckpt", + "sample_rate": 48000, + "sample_size": 131072, + }, + "maestro-uncond-150k": { + "url": "https://model-server.zqevans2.workers.dev/maestro-uncond-150k.ckpt", + "sample_rate": 16000, + "sample_size": 65536, + }, + "unlocked-uncond-250k": { + "url": "https://model-server.zqevans2.workers.dev/unlocked-uncond-250k.ckpt", + "sample_rate": 16000, + "sample_size": 65536, + }, + "honk-140k": { + "url": "https://model-server.zqevans2.workers.dev/honk-140k.ckpt", + "sample_rate": 16000, + "sample_size": 65536, + }, +} + + +def alpha_sigma_to_t(alpha, sigma): + """Returns a timestep, given the scaling factors for the clean image and for + the noise.""" + return torch.atan2(sigma, alpha) / math.pi * 2 + + +def get_crash_schedule(t): + sigma = torch.sin(t * math.pi / 2) ** 2 + alpha = (1 - sigma**2) ** 0.5 + return alpha_sigma_to_t(alpha, sigma) + + +class Object(object): + pass + + +class DiffusionUncond(nn.Module): + def __init__(self, global_args): + super().__init__() + + self.diffusion = DiffusionAttnUnet1D(global_args, n_attn_layers=4) + self.diffusion_ema = deepcopy(self.diffusion) + self.rng = torch.quasirandom.SobolEngine(1, scramble=True) + + +def download(model_name): + url = MODELS_MAP[model_name]["url"] + r = requests.get(url, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT) + + local_filename = f"./{model_name}.ckpt" + with open(local_filename, "wb") as fp: + for chunk in r.iter_content(chunk_size=8192): + fp.write(chunk) + + return local_filename + + +def main(args): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model_name = args.model_path.split("/")[-1].split(".")[0] + if not os.path.isfile(args.model_path): + assert model_name == args.model_path, ( + f"Make sure to provide one of the official model names {MODELS_MAP.keys()}" + ) + args.model_path = download(model_name) + + sample_rate = MODELS_MAP[model_name]["sample_rate"] + sample_size = MODELS_MAP[model_name]["sample_size"] + + config = Object() + config.sample_size = sample_size + config.sample_rate = sample_rate + config.latent_dim = 0 + + diffusers_model = UNet1DModel(sample_size=sample_size, sample_rate=sample_rate) + + orig_model = DiffusionUncond(config) + orig_model.load_state_dict(torch.load(args.model_path, map_location=device)["state_dict"]) + orig_model = orig_model.diffusion_ema.eval() + orig_model_state_dict = orig_model.state_dict() + converted = convert_component_checkpoint(orig_model_state_dict, dict(diffusers_model.config), "UNet1DModel") + diffusers_model.load_state_dict(converted, strict=True) + + steps = 100 + seed = 33 + + diffusers_scheduler = IPNDMScheduler(num_train_timesteps=steps) + + generator = torch.manual_seed(seed) + noise = torch.randn([1, 2, config.sample_size], generator=generator).to(device) + + t = torch.linspace(1, 0, steps + 1, device=device)[:-1] + step_list = get_crash_schedule(t) + + pipe = DanceDiffusionPipeline(unet=diffusers_model, scheduler=diffusers_scheduler) + + generator = torch.manual_seed(33) + audio = pipe(num_inference_steps=steps, generator=generator).audios + + generated = sampling.iplms_sample(orig_model, noise, step_list, {}) + generated = generated.clamp(-1, 1) + + diff_sum = (generated - audio).abs().sum() + diff_max = (generated - audio).abs().max() + + if args.save: + pipe.save_pretrained(args.checkpoint_path) + + print("Diff sum", diff_sum) + print("Diff max", diff_max) + + assert diff_max < 1e-3, f"Diff max: {diff_max} is too much :-/" + + print(f"Conversion for {model_name} successful!") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.") + parser.add_argument( + "--save", default=True, type=bool, required=False, help="Whether to save the converted model or not." + ) + parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.") + args = parser.parse_args() + + main(args) diff --git a/scripts/recipes/ddpm.py b/scripts/recipes/ddpm.py new file mode 100644 index 000000000000..a737dda47df5 --- /dev/null +++ b/scripts/recipes/ddpm.py @@ -0,0 +1,66 @@ +import argparse +import json + +import torch + +from diffusers import AutoencoderKL, DDPMPipeline, DDPMScheduler, UNet2DModel, VQModel +from diffusers.loaders.conversion import get_conversion + + +def convert_ddpm_checkpoint(checkpoint, config): + return get_conversion("UNet2DModel", {**config, "original_format": "ddpm"}).to_diffusers(checkpoint) + + +def convert_vq_autoenc_checkpoint(checkpoint, config): + return get_conversion(config["_class_name"], config).to_diffusers(checkpoint) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + + parser.add_argument( + "--config_file", + default=None, + type=str, + required=True, + help="The config json file corresponding to the architecture.", + ) + + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + + args = parser.parse_args() + checkpoint = torch.load(args.checkpoint_path) + + with open(args.config_file) as f: + config = json.loads(f.read()) + + # unet case + key_prefix_set = {key.split(".")[0] for key in checkpoint.keys()} + if "encoder" in key_prefix_set and "decoder" in key_prefix_set: + converted_checkpoint = convert_vq_autoenc_checkpoint(checkpoint, config) + else: + converted_checkpoint = convert_ddpm_checkpoint(checkpoint, config) + + if "ddpm" in config: + del config["ddpm"] + + if config["_class_name"] == "VQModel": + model = VQModel(**config) + model.load_state_dict(converted_checkpoint) + model.save_pretrained(args.dump_path) + elif config["_class_name"] == "AutoencoderKL": + model = AutoencoderKL(**config) + model.load_state_dict(converted_checkpoint) + model.save_pretrained(args.dump_path) + else: + model = UNet2DModel(**config) + model.load_state_dict(converted_checkpoint) + + scheduler = DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) + + pipe = DDPMPipeline(unet=model, scheduler=scheduler) + pipe.save_pretrained(args.dump_path) diff --git a/scripts/recipes/deepfloyd_if.py b/scripts/recipes/deepfloyd_if.py new file mode 100644 index 000000000000..6c751cca8c96 --- /dev/null +++ b/scripts/recipes/deepfloyd_if.py @@ -0,0 +1,313 @@ +import argparse +import os + +import numpy as np +import torch +import yaml +from transformers import CLIPConfig, CLIPImageProcessor, CLIPVisionModelWithProjection, T5EncoderModel, T5Tokenizer + +from diffusers import DDPMScheduler, IFPipeline, IFSuperResolutionPipeline, UNet2DConditionModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.deepfloyd_if import ( + create_unet_diffusers_config, + superres_create_unet_diffusers_config, +) +from diffusers.pipelines.deepfloyd_if.safety_checker import IFSafetyChecker + + +def parse_args(): + parser = argparse.ArgumentParser() + + parser.add_argument("--dump_path", required=False, default=None, type=str) + + parser.add_argument("--dump_path_stage_2", required=False, default=None, type=str) + + parser.add_argument("--dump_path_stage_3", required=False, default=None, type=str) + + parser.add_argument("--unet_config", required=False, default=None, type=str, help="Path to unet config file") + + parser.add_argument( + "--unet_checkpoint_path", required=False, default=None, type=str, help="Path to unet checkpoint file" + ) + + parser.add_argument( + "--unet_checkpoint_path_stage_2", + required=False, + default=None, + type=str, + help="Path to stage 2 unet checkpoint file", + ) + + parser.add_argument( + "--unet_checkpoint_path_stage_3", + required=False, + default=None, + type=str, + help="Path to stage 3 unet checkpoint file", + ) + + parser.add_argument("--p_head_path", type=str, required=True) + + parser.add_argument("--w_head_path", type=str, required=True) + + args = parser.parse_args() + + return args + + +def main(args): + tokenizer = T5Tokenizer.from_pretrained("google/t5-v1_1-xxl") + text_encoder = T5EncoderModel.from_pretrained("google/t5-v1_1-xxl") + + feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14") + safety_checker = convert_safety_checker(p_head_path=args.p_head_path, w_head_path=args.w_head_path) + + if args.unet_config is not None and args.unet_checkpoint_path is not None and args.dump_path is not None: + convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args) + + if args.unet_checkpoint_path_stage_2 is not None and args.dump_path_stage_2 is not None: + convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=2) + + if args.unet_checkpoint_path_stage_3 is not None and args.dump_path_stage_3 is not None: + convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage=3) + + +def convert_stage_1_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args): + unet = get_stage_1_unet(args.unet_config, args.unet_checkpoint_path) + + scheduler = DDPMScheduler( + variance_type="learned_range", + beta_schedule="squaredcos_cap_v2", + prediction_type="epsilon", + thresholding=True, + dynamic_thresholding_ratio=0.95, + sample_max_value=1.5, + ) + + pipe = IFPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + unet=unet, + scheduler=scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + requires_safety_checker=True, + ) + + pipe.save_pretrained(args.dump_path) + + +def convert_super_res_pipeline(tokenizer, text_encoder, feature_extractor, safety_checker, args, stage): + if stage == 2: + unet_checkpoint_path = args.unet_checkpoint_path_stage_2 + sample_size = None + dump_path = args.dump_path_stage_2 + elif stage == 3: + unet_checkpoint_path = args.unet_checkpoint_path_stage_3 + sample_size = 1024 + dump_path = args.dump_path_stage_3 + else: + assert False + + unet = get_super_res_unet(unet_checkpoint_path, verify_param_count=False, sample_size=sample_size) + + image_noising_scheduler = DDPMScheduler( + beta_schedule="squaredcos_cap_v2", + ) + + scheduler = DDPMScheduler( + variance_type="learned_range", + beta_schedule="squaredcos_cap_v2", + prediction_type="epsilon", + thresholding=True, + dynamic_thresholding_ratio=0.95, + sample_max_value=1.0, + ) + + pipe = IFSuperResolutionPipeline( + tokenizer=tokenizer, + text_encoder=text_encoder, + unet=unet, + scheduler=scheduler, + image_noising_scheduler=image_noising_scheduler, + safety_checker=safety_checker, + feature_extractor=feature_extractor, + requires_safety_checker=True, + ) + + pipe.save_pretrained(dump_path) + + +def get_stage_1_unet(unet_config, unet_checkpoint_path): + original_unet_config = yaml.safe_load(unet_config) + original_unet_config = original_unet_config["params"] + + unet_diffusers_config = create_unet_diffusers_config(original_unet_config) + + unet = UNet2DConditionModel(**unet_diffusers_config) + + device = "cuda" if torch.cuda.is_available() else "cpu" + unet_checkpoint = torch.load(unet_checkpoint_path, map_location=device) + + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + unet_checkpoint, unet_diffusers_config, path=unet_checkpoint_path + ) + + unet.load_state_dict(converted_unet_checkpoint) + + return unet + + +def convert_safety_checker(p_head_path, w_head_path): + state = {} + for name, path in (("p_head", p_head_path), ("w_head", w_head_path)): + with np.load(path) as archive: + state.update({f"{name}.{key}": torch.from_numpy(archive[key]) for key in ("weights", "biases")}) + vision_model = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") + state.update({"vision_model." + key: value for key, value in vision_model.state_dict().items()}) + config = CLIPConfig.from_pretrained("openai/clip-vit-large-patch14") + model = IFSafetyChecker(config) + model.load_state_dict(get_conversion("IFSafetyChecker", config.to_dict()).to_diffusers(state), strict=True) + return model + + +def convert_ldm_unet_checkpoint(unet_state_dict, config, path=None, **kwargs): + return get_conversion("UNet2DConditionModel", config).to_diffusers(unet_state_dict) + + +# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) + + +# below is copy and pasted from original convert_if_stage_2.py script + + +def get_super_res_unet(unet_checkpoint_path, verify_param_count=True, sample_size=None): + orig_path = unet_checkpoint_path + + with open(os.path.join(orig_path, "config.yml")) as handle: + original_unet_config = yaml.safe_load(handle) + original_unet_config = original_unet_config["params"] + + unet_diffusers_config = superres_create_unet_diffusers_config(original_unet_config) + unet_diffusers_config["time_embedding_dim"] = original_unet_config["model_channels"] * int( + original_unet_config["channel_mult"].split(",")[-1] + ) + if original_unet_config["encoder_dim"] != original_unet_config["encoder_channels"]: + unet_diffusers_config["encoder_hid_dim"] = original_unet_config["encoder_dim"] + unet_diffusers_config["class_embed_type"] = "timestep" + unet_diffusers_config["addition_embed_type"] = "text" + + unet_diffusers_config["time_embedding_act_fn"] = "gelu" + unet_diffusers_config["resnet_skip_time_act"] = True + unet_diffusers_config["resnet_out_scale_factor"] = 1 / 0.7071 + unet_diffusers_config["mid_block_scale_factor"] = 1 / 0.7071 + unet_diffusers_config["only_cross_attention"] = ( + bool(original_unet_config["disable_self_attentions"]) + if ( + "disable_self_attentions" in original_unet_config + and isinstance(original_unet_config["disable_self_attentions"], int) + ) + else True + ) + + if sample_size is None: + unet_diffusers_config["sample_size"] = original_unet_config["image_size"] + else: + # The second upscaler unet's sample size is incorrectly specified + # in the config and is instead hardcoded in source + unet_diffusers_config["sample_size"] = sample_size + + unet_checkpoint = torch.load(os.path.join(unet_checkpoint_path, "pytorch_model.bin"), map_location="cpu") + + if verify_param_count: + # check that architecture matches - is a bit slow + verify_parameter_count(orig_path, unet_diffusers_config) + + converted_unet_checkpoint = superres_convert_ldm_unet_checkpoint( + unet_checkpoint, unet_diffusers_config, path=unet_checkpoint_path + ) + converted_keys = converted_unet_checkpoint.keys() + + model = UNet2DConditionModel(**unet_diffusers_config) + expected_weights = model.state_dict().keys() + + diff_c_e = set(converted_keys) - set(expected_weights) + diff_e_c = set(expected_weights) - set(converted_keys) + + assert len(diff_e_c) == 0, f"Expected, but not converted: {diff_e_c}" + assert len(diff_c_e) == 0, f"Converted, but not expected: {diff_c_e}" + + model.load_state_dict(converted_unet_checkpoint) + + return model + + +def superres_convert_ldm_unet_checkpoint(unet_state_dict, config, path=None, **kwargs): + return get_conversion("UNet2DConditionModel", config).to_diffusers(unet_state_dict) + + +def verify_parameter_count(orig_path, unet_diffusers_config): + if "-II-" in orig_path: + from deepfloyd_if.modules import IFStageII + + if_II = IFStageII(device="cpu", dir_or_name=orig_path) + elif "-III-" in orig_path: + from deepfloyd_if.modules import IFStageIII + + if_II = IFStageIII(device="cpu", dir_or_name=orig_path) + else: + assert f"Weird name. Should have -II- or -III- in path: {orig_path}" + + unet = UNet2DConditionModel(**unet_diffusers_config) + + # in params + assert_param_count(unet.time_embedding, if_II.model.time_embed) + assert_param_count(unet.conv_in, if_II.model.input_blocks[:1]) + + # downblocks + assert_param_count(unet.down_blocks[0], if_II.model.input_blocks[1:4]) + assert_param_count(unet.down_blocks[1], if_II.model.input_blocks[4:7]) + assert_param_count(unet.down_blocks[2], if_II.model.input_blocks[7:11]) + + if "-II-" in orig_path: + assert_param_count(unet.down_blocks[3], if_II.model.input_blocks[11:17]) + assert_param_count(unet.down_blocks[4], if_II.model.input_blocks[17:]) + if "-III-" in orig_path: + assert_param_count(unet.down_blocks[3], if_II.model.input_blocks[11:15]) + assert_param_count(unet.down_blocks[4], if_II.model.input_blocks[15:20]) + assert_param_count(unet.down_blocks[5], if_II.model.input_blocks[20:]) + + # mid block + assert_param_count(unet.mid_block, if_II.model.middle_block) + + # up block + if "-II-" in orig_path: + assert_param_count(unet.up_blocks[0], if_II.model.output_blocks[:6]) + assert_param_count(unet.up_blocks[1], if_II.model.output_blocks[6:12]) + assert_param_count(unet.up_blocks[2], if_II.model.output_blocks[12:16]) + assert_param_count(unet.up_blocks[3], if_II.model.output_blocks[16:19]) + assert_param_count(unet.up_blocks[4], if_II.model.output_blocks[19:]) + if "-III-" in orig_path: + assert_param_count(unet.up_blocks[0], if_II.model.output_blocks[:5]) + assert_param_count(unet.up_blocks[1], if_II.model.output_blocks[5:10]) + assert_param_count(unet.up_blocks[2], if_II.model.output_blocks[10:14]) + assert_param_count(unet.up_blocks[3], if_II.model.output_blocks[14:18]) + assert_param_count(unet.up_blocks[4], if_II.model.output_blocks[18:21]) + assert_param_count(unet.up_blocks[5], if_II.model.output_blocks[21:24]) + + # out params + assert_param_count(unet.conv_norm_out, if_II.model.out[0]) + assert_param_count(unet.conv_out, if_II.model.out[2]) + + # make sure all model architecture has same param count + assert_param_count(unet, if_II.model) + + +def assert_param_count(model_1, model_2): + count_1 = sum(p.numel() for p in model_1.parameters()) + count_2 = sum(p.numel() for p in model_2.parameters()) + assert count_1 == count_2, f"{model_1.__class__}: {count_1} != {model_2.__class__}: {count_2}" + + +if __name__ == "__main__": + main(parse_args()) diff --git a/scripts/recipes/dit.py b/scripts/recipes/dit.py new file mode 100644 index 000000000000..a13f40d23b15 --- /dev/null +++ b/scripts/recipes/dit.py @@ -0,0 +1,89 @@ +import argparse +import os + +import torch +from torchvision.datasets.utils import download_url + +from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, Transformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +pretrained_models = {512: "DiT-XL-2-512x512.pt", 256: "DiT-XL-2-256x256.pt"} + + +def download_model(model_name): + """ + Downloads a pre-trained DiT model from the web. + """ + local_path = f"pretrained_models/{model_name}" + if not os.path.isfile(local_path): + os.makedirs("pretrained_models", exist_ok=True) + web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{model_name}" + download_url(web_path, "pretrained_models") + model = torch.load(local_path, map_location=lambda storage, loc: storage) + return model + + +def main(args): + state_dict = download_model(pretrained_models[args.image_size]) + + # DiT XL/2 + transformer = Transformer2DModel( + sample_size=args.image_size // 8, + num_layers=28, + attention_head_dim=72, + in_channels=4, + out_channels=8, + patch_size=2, + attention_bias=True, + num_attention_heads=16, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + norm_type="ada_norm_zero", + norm_elementwise_affine=False, + ) + transformer.load_state_dict( + convert_component_checkpoint(state_dict, dict(transformer.config), "DiTTransformer2DModel"), strict=True + ) + + scheduler = DDIMScheduler( + num_train_timesteps=1000, + beta_schedule="linear", + prediction_type="epsilon", + clip_sample=False, + ) + + vae = AutoencoderKL.from_pretrained(args.vae_model) + + pipeline = DiTPipeline(transformer=transformer, vae=vae, scheduler=scheduler) + + if args.save: + pipeline.save_pretrained(args.checkpoint_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--image_size", + default=256, + type=int, + required=False, + help="Image size of pretrained model, either 256 or 512.", + ) + parser.add_argument( + "--vae_model", + default="stabilityai/sd-vae-ft-ema", + type=str, + required=False, + help="Path to pretrained VAE model, either stabilityai/sd-vae-ft-mse or stabilityai/sd-vae-ft-ema.", + ) + parser.add_argument( + "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." + ) + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the output pipeline." + ) + + args = parser.parse_args() + main(args) diff --git a/scripts/recipes/flux2.py b/scripts/recipes/flux2.py new file mode 100644 index 000000000000..59b7a1234c7a --- /dev/null +++ b/scripts/recipes/flux2.py @@ -0,0 +1,133 @@ +import argparse + +import safetensors.torch +import torch +from huggingface_hub import hf_hub_download +from transformers import AutoProcessor, GenerationConfig, Mistral3ForConditionalGeneration + +from diffusers import AutoencoderKLFlux2, FlowMatchEulerDiscreteScheduler, Flux2Pipeline, Flux2Transformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.flux2 import get_flux2_transformer_config + + +""" +# VAE + +python scripts/recipes/flux2.py \ +--original_state_dict_repo_id "diffusers-internal-dev/new-model-image" \ +--vae_filename "flux2-vae.sft" \ +--output_path "/raid/yiyi/dummy-flux2-diffusers" \ +--vae + +# DiT + +python scripts/recipes/flux2.py \ + --original_state_dict_repo_id diffusers-internal-dev/new-model-image \ + --dit_filename flux-dev-dummy.sft \ + --dit \ + --output_path . + +# Full pipe + +python scripts/recipes/flux2.py \ + --original_state_dict_repo_id diffusers-internal-dev/new-model-image \ + --dit_filename flux-dev-dummy.sft \ + --vae_filename "flux2-vae.sft" \ + --dit --vae --full_pipe \ + --output_path . +""" + + +parser = argparse.ArgumentParser() +parser.add_argument("--original_state_dict_repo_id", default=None, type=str) +parser.add_argument("--vae_filename", default="flux2-vae.sft", type=str) +parser.add_argument("--dit_filename", default="flux2-dev.safetensors", type=str) +parser.add_argument("--vae", action="store_true") +parser.add_argument("--dit", action="store_true") +parser.add_argument("--vae_dtype", type=str, default="fp32") +parser.add_argument("--dit_dtype", type=str, default="bf16") +parser.add_argument("--checkpoint_path", default=None, type=str) +parser.add_argument("--full_pipe", action="store_true") +parser.add_argument("--output_path", type=str) + +args = parser.parse_args() + + +def load_original_checkpoint(args, filename): + if args.original_state_dict_repo_id is not None: + ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename) + elif args.checkpoint_path is not None: + ckpt_path = args.checkpoint_path + else: + raise ValueError(" please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") + + original_state_dict = safetensors.torch.load_file(ckpt_path) + return original_state_dict + + +def convert_flux2_vae_checkpoint_to_diffusers(original_state_dict, config): + return convert_component_checkpoint(original_state_dict, config, "AutoencoderKLFlux2") + + +# in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; +# while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use +# diffusers implementation + + +def convert_flux2_transformer_to_diffusers(original_state_dict, model_type): + config = get_flux2_transformer_config(model_type)["diffusers_config"] + return Flux2Transformer2DModel.from_single_file(original_state_dict, config=config) + + +def main(args): + if args.vae: + original_vae_ckpt = load_original_checkpoint(args, filename=args.vae_filename) + vae = AutoencoderKLFlux2() + converted_vae_state_dict = convert_flux2_vae_checkpoint_to_diffusers(original_vae_ckpt, vae.config) + vae.load_state_dict(converted_vae_state_dict, strict=True) + if not args.full_pipe: + vae_dtype = torch.bfloat16 if args.vae_dtype == "bf16" else torch.float32 + vae.to(vae_dtype).save_pretrained(f"{args.output_path}/vae") + + if args.dit: + original_dit_ckpt = load_original_checkpoint(args, filename=args.dit_filename) + + if "klein-4b" in args.dit_filename: + model_type = "klein-4b" + elif "klein-9b" in args.dit_filename: + model_type = "klein-9b" + else: + model_type = "flux2-dev" + transformer = convert_flux2_transformer_to_diffusers(original_dit_ckpt, model_type) + if not args.full_pipe: + dit_dtype = torch.bfloat16 if args.dit_dtype == "bf16" else torch.float32 + transformer.to(dit_dtype).save_pretrained(f"{args.output_path}/transformer") + + if args.full_pipe: + tokenizer_id = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" + text_encoder_id = "mistralai/Mistral-Small-3.2-24B-Instruct-2506" + generate_config = GenerationConfig.from_pretrained(text_encoder_id) + generate_config.do_sample = True + text_encoder = Mistral3ForConditionalGeneration.from_pretrained( + text_encoder_id, generation_config=generate_config, torch_dtype=torch.bfloat16 + ) + tokenizer = AutoProcessor.from_pretrained(tokenizer_id) + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + "black-forest-labs/FLUX.1-dev", subfolder="scheduler" + ) + + if_distilled = "base" not in args.dit_filename + + pipe = Flux2Pipeline( + vae=vae, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + if_distilled=if_distilled, + ) + pipe.save_pretrained(args.output_path) + + +if __name__ == "__main__": + main(args) diff --git a/scripts/convert_flux_xlabs_ipadapter_to_diffusers.py b/scripts/recipes/flux_ip_adapter.py similarity index 53% rename from scripts/convert_flux_xlabs_ipadapter_to_diffusers.py rename to scripts/recipes/flux_ip_adapter.py index b701b7fb40b1..4ba864bd8238 100644 --- a/scripts/convert_flux_xlabs_ipadapter_to_diffusers.py +++ b/scripts/recipes/flux_ip_adapter.py @@ -1,11 +1,10 @@ import argparse -from contextlib import nullcontext import safetensors.torch -from accelerate import init_empty_weights from huggingface_hub import hf_hub_download -from diffusers.utils.import_utils import is_accelerate_available, is_transformers_available +from diffusers.loaders.conversion import get_conversion +from diffusers.utils.import_utils import is_transformers_available if is_transformers_available(): @@ -16,15 +15,13 @@ vision = False """ -python scripts/convert_flux_xlabs_ipadapter_to_diffusers.py \ +python scripts/recipes/flux_ip_adapter.py \ --original_state_dict_repo_id "XLabs-AI/flux-ip-adapter" \ --filename "flux-ip-adapter.safetensors" --output_path "flux-ip-adapter-hf/" """ -CTX = init_empty_weights if is_accelerate_available else nullcontext - parser = argparse.ArgumentParser() parser.add_argument("--original_state_dict_repo_id", default=None, type=str) parser.add_argument("--filename", default="flux.safetensors", type=str) @@ -48,35 +45,7 @@ def load_original_checkpoint(args): def convert_flux_ipadapter_checkpoint_to_diffusers(original_state_dict, num_layers): - converted_state_dict = {} - - # image_proj - ## norm - converted_state_dict["image_proj.norm.weight"] = original_state_dict.pop("ip_adapter_proj_model.norm.weight") - converted_state_dict["image_proj.norm.bias"] = original_state_dict.pop("ip_adapter_proj_model.norm.bias") - ## proj - converted_state_dict["image_proj.proj.weight"] = original_state_dict.pop("ip_adapter_proj_model.norm.weight") - converted_state_dict["image_proj.proj.bias"] = original_state_dict.pop("ip_adapter_proj_model.norm.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"ip_adapter.{i}." - # to_k_ip - converted_state_dict[f"{block_prefix}to_k_ip.bias"] = original_state_dict.pop( - f"double_blocks.{i}.processor.ip_adapter_double_stream_k_proj.bias" - ) - converted_state_dict[f"{block_prefix}to_k_ip.weight"] = original_state_dict.pop( - f"double_blocks.{i}.processor.ip_adapter_double_stream_k_proj.weight" - ) - # to_v_ip - converted_state_dict[f"{block_prefix}to_v_ip.bias"] = original_state_dict.pop( - f"double_blocks.{i}.processor.ip_adapter_double_stream_v_proj.bias" - ) - converted_state_dict[f"{block_prefix}to_k_ip.weight"] = original_state_dict.pop( - f"double_blocks.{i}.processor.ip_adapter_double_stream_v_proj.weight" - ) - - return converted_state_dict + return get_conversion("FluxIPAdapter", {"num_layers": num_layers}).to_diffusers(original_state_dict) def main(args): @@ -85,6 +54,9 @@ def main(args): num_layers = 19 converted_ip_adapter_state_dict = convert_flux_ipadapter_checkpoint_to_diffusers(original_ckpt, num_layers) + from pathlib import Path + + Path(args.output_path).mkdir(parents=True, exist_ok=True) print("Saving Flux IP-Adapter in Diffusers format.") safetensors.torch.save_file(converted_ip_adapter_state_dict, f"{args.output_path}/model.safetensors") diff --git a/scripts/recipes/gligen.py b/scripts/recipes/gligen.py new file mode 100644 index 000000000000..e27a5854ae3b --- /dev/null +++ b/scripts/recipes/gligen.py @@ -0,0 +1,190 @@ +import argparse + +import torch +import yaml +from transformers import ( + CLIPProcessor, + CLIPTextModel, + CLIPTokenizer, + CLIPVisionModelWithProjection, +) + +from diffusers import ( + AutoencoderKL, + DDIMScheduler, + StableDiffusionGLIGENPipeline, + StableDiffusionGLIGENTextImagePipeline, + UNet2DConditionModel, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.gligen import create_unet_config, create_vae_config + + +def convert_open_clip_checkpoint(checkpoint): + model = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") + state = checkpoint["text_encoder"] + config = model.config.to_dict() + if any(key.startswith("transformer.text_model.") for key in state): + state = {key.removeprefix("transformer."): value for key, value in state.items()} + config["original_format"] = "clip" + else: + config["original_format"] = "openclip" + state = dict(state) + state.pop("text_model.embeddings.position_ids", None) + converted = convert_component_checkpoint(state, config, "CLIPTextModel") + model.load_state_dict(converted, strict=True) + return model + + +def convert_gligen_vae_checkpoint(checkpoint, config): + return convert_component_checkpoint(checkpoint["autoencoder"], config, "AutoencoderKL") + + +def convert_gligen_unet_checkpoint(checkpoint, config, path=None, extract_ema=False): + return convert_component_checkpoint(checkpoint["model"], config, "UNet2DConditionModel", extract_ema=extract_ema) + + +def convert_gligen_to_diffusers( + checkpoint_path: str, + original_config_file: str, + attention_type: str, + image_size: int = 512, + extract_ema: bool = False, + num_in_channels: int = None, + device: str = None, +): + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint = torch.load(checkpoint_path, map_location=device) + else: + checkpoint = torch.load(checkpoint_path, map_location=device) + + if "global_step" in checkpoint: + checkpoint["global_step"] + else: + print("global_step key not found in model") + + original_config = yaml.safe_load(original_config_file) + + if num_in_channels is not None: + original_config["model"]["params"]["in_channels"] = num_in_channels + + num_train_timesteps = original_config["diffusion"]["params"]["timesteps"] + beta_start = original_config["diffusion"]["params"]["linear_start"] + beta_end = original_config["diffusion"]["params"]["linear_end"] + + scheduler = DDIMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + steps_offset=1, + clip_sample=False, + set_alpha_to_one=False, + prediction_type="epsilon", + ) + + # Convert the UNet2DConditionalModel model + unet_config = create_unet_config(original_config, image_size, attention_type) + unet = UNet2DConditionModel(**unet_config) + + converted_unet_checkpoint = convert_gligen_unet_checkpoint( + checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema + ) + + unet.load_state_dict(converted_unet_checkpoint) + + # Convert the VAE model + vae_config = create_vae_config(original_config, image_size) + converted_vae_checkpoint = convert_gligen_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + + # Convert the text model + text_encoder = convert_open_clip_checkpoint(checkpoint) + tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") + + if attention_type == "gated-text-image": + image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") + processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") + + pipe = StableDiffusionGLIGENTextImagePipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + image_encoder=image_encoder, + processor=processor, + unet=unet, + scheduler=scheduler, + safety_checker=None, + feature_extractor=None, + ) + elif attention_type == "gated": + pipe = StableDiffusionGLIGENPipeline( + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + safety_checker=None, + feature_extractor=None, + ) + + return pipe + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--original_config_file", + default=None, + type=str, + required=True, + help="The YAML config file corresponding to the gligen architecture.", + ) + parser.add_argument( + "--num_in_channels", + default=None, + type=int, + help="The number of input channels. If `None` number of input channels will be automatically inferred.", + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument( + "--attention_type", + default=None, + type=str, + required=True, + help="Type of attention ex: gated or gated-text-image", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--device", type=str, help="Device to use.") + parser.add_argument("--half", action="store_true", help="Save weights in half precision.") + + args = parser.parse_args() + + pipe = convert_gligen_to_diffusers( + checkpoint_path=args.checkpoint_path, + original_config_file=args.original_config_file, + attention_type=args.attention_type, + extract_ema=args.extract_ema, + num_in_channels=args.num_in_channels, + device=args.device, + ) + + if args.half: + pipe.to(dtype=torch.float16) + + pipe.save_pretrained(args.dump_path) diff --git a/scripts/recipes/hunyuan_dit.py b/scripts/recipes/hunyuan_dit.py new file mode 100644 index 000000000000..e023b6ea2f7c --- /dev/null +++ b/scripts/recipes/hunyuan_dit.py @@ -0,0 +1,81 @@ +import argparse + +import torch + +from diffusers import HunyuanDiT2DModel + + +def main(args): + state_dict = torch.load(args.pt_checkpoint_path, map_location="cpu") + + if args.load_key != "none": + try: + state_dict = state_dict[args.load_key] + except KeyError: + raise KeyError( + f"{args.load_key} not found in the checkpoint.Please load from the following keys:{state_dict.keys()}" + ) + + device = "cuda" + model_config = HunyuanDiT2DModel.load_config("Tencent-Hunyuan/HunyuanDiT-Diffusers", subfolder="transformer") + model_config["use_style_cond_and_image_meta_size"] = ( + args.use_style_cond_and_image_meta_size + ) ### version <= v1.1: True; version >= v1.2: False + + model = HunyuanDiT2DModel.from_single_file(state_dict, config=model_config).to(device) + + from diffusers import HunyuanDiTPipeline + + if args.use_style_cond_and_image_meta_size: + pipe = HunyuanDiTPipeline.from_pretrained( + "Tencent-Hunyuan/HunyuanDiT-Diffusers", transformer=model, torch_dtype=torch.float32 + ) + else: + pipe = HunyuanDiTPipeline.from_pretrained( + "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers", transformer=model, torch_dtype=torch.float32 + ) + pipe.to("cuda") + pipe.to(dtype=torch.float32) + + if args.save: + pipe.save_pretrained(args.output_checkpoint_path) + + # ### NOTE: HunyuanDiT supports both Chinese and English inputs + prompt = "一个宇航员在骑马" + # prompt = "An astronaut riding a horse" + generator = torch.Generator(device="cuda").manual_seed(0) + image = pipe( + height=1024, width=1024, prompt=prompt, generator=generator, num_inference_steps=25, guidance_scale=5.0 + ).images[0] + + image.save("img.png") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--save", default=True, type=bool, required=False, help="Whether to save the converted pipeline or not." + ) + parser.add_argument( + "--pt_checkpoint_path", default=None, type=str, required=True, help="Path to the .pt pretrained model." + ) + parser.add_argument( + "--output_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the output converted diffusers pipeline.", + ) + parser.add_argument( + "--load_key", default="none", type=str, required=False, help="The key to load from the pretrained .pt file" + ) + parser.add_argument( + "--use_style_cond_and_image_meta_size", + type=bool, + default=False, + help="version <= v1.1: True; version >= v1.2: False", + ) + + args = parser.parse_args() + main(args) diff --git a/scripts/recipes/hunyuan_image.py b/scripts/recipes/hunyuan_image.py new file mode 100644 index 000000000000..bf1a4dc01c5f --- /dev/null +++ b/scripts/recipes/hunyuan_image.py @@ -0,0 +1,312 @@ +import argparse +import logging + +import torch +from safetensors import safe_open + +from diffusers import AutoencoderKLHunyuanImage, AutoencoderKLHunyuanImageRefiner, HunyuanImageTransformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +logger = logging.getLogger(__name__) # pylint: disable=invalid-name + + +""" +Usage examples +============== + +python scripts/recipes/hunyuan_image.py \ + --model_type hunyuanimage2.1 \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1.safetensors" \ + --vae_checkpoint_path "HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ + --output_path "/raid/yiyi/test-hy21-diffusers" \ + --dtype fp32 + +python scripts/recipes/hunyuan_image.py \ + --model_type hunyuanimage2.1-distilled \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage2.1-distilled.safetensors" \ + --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_2_1/pytorch_model.ckpt" \ + --output_path "/raid/yiyi/test-hy21-distilled-diffusers" \ + --dtype fp32 + + +python scripts/recipes/hunyuan_image.py \ + --model_type hunyuanimage-refiner \ + --transformer_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/dit/hunyuanimage-refiner.safetensors" \ + --vae_checkpoint_path "/raid/yiyi/HunyuanImage-2.1/ckpts/vae/vae_refiner/pytorch_model.pt" \ + --output_path "/raid/yiyi/test-hy2-refiner-diffusers" \ + --dtype fp32 +""" + +parser = argparse.ArgumentParser() +parser.add_argument( + "--model_type", type=str, default=None +) # hunyuanimage2.1, hunyuanimage2.1-distilled, hunyuanimage-refiner +parser.add_argument("--transformer_checkpoint_path", default=None, type=str) # ckpts/dit/hunyuanimage2.1.safetensors +parser.add_argument("--vae_checkpoint_path", default=None, type=str) # ckpts/vae/vae_2_1/pytorch_model.ckpt +parser.add_argument("--output_path", type=str) +parser.add_argument("--dtype", type=str, default="fp32") + +args = parser.parse_args() +dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + + +# copied from https://github.com/Tencent-Hunyuan/HunyuanImage-2.1/hyimage/models/hunyuan/modules/hunyuanimage_dit.py#L21 + + +def load_original_vae_checkpoint(args): + # "ckpts/vae/vae_2_1/pytorch_model.ckpt" + state_dict = torch.load(args.vae_checkpoint_path) + + if "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + vae_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("vae."): + vae_state_dict[k.replace("vae.", "")] = v + + return vae_state_dict + + +def load_original_refiner_vae_checkpoint(args): + # "ckpts/vae/vae_refiner/pytorch_model.pt" + state_dict = torch.load(args.vae_checkpoint_path) + + if "state_dict" in state_dict: + state_dict = state_dict["state_dict"] + vae_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("vae."): + vae_state_dict[k.replace("vae.", "")] = v + return vae_state_dict + + +def load_original_transformer_checkpoint(args): + # ckpts/dit/hunyuanimage-refiner.safetensors" + # ckpts/dit/hunyuanimage2.1.safetensors" + state_dict = {} + with safe_open(args.transformer_checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + return state_dict + + +def convert_hunyuan_image_transformer_checkpoint_to_diffusers(original_state_dict, config): + return convert_component_checkpoint(original_state_dict, config, "HunyuanImageTransformer2DModel"), {} + + +def convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 +): + config = {"block_out_channels": block_out_channels, "layers_per_block": layers_per_block} + return convert_component_checkpoint(original_state_dict, config, "AutoencoderKLHunyuanImage"), {} + + +def convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( + original_state_dict, block_out_channels=[128, 256, 512, 1024, 1024], layers_per_block=2 +): + config = {"block_out_channels": block_out_channels, "layers_per_block": layers_per_block} + return convert_component_checkpoint(original_state_dict, config, "AutoencoderKLHunyuanImageRefiner"), {} + + +def main(args): + if args.model_type == "hunyuanimage2.1": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_vae_checkpoint(args) + + transformer_config = { + "in_channels": 64, + "out_channels": 64, + "num_attention_heads": 28, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "patch_size": (1, 1), + "qk_norm": "rms_norm", + "guidance_embeds": False, + "text_embed_dim": 3584, + "text_embed_2_dim": 1472, + "rope_theta": 256.0, + "rope_axes_dim": (64, 64), + } + + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, config=transformer_config + ) + ) + + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae_config_diffusers = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 64, + "block_out_channels": [128, 256, 512, 512, 1024, 1024], + "layers_per_block": 2, + "spatial_compression_ratio": 32, + "sample_size": 384, + "scaling_factor": 0.75289, + "downsample_match_channel": True, + "upsample_match_channel": True, + } + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for vae: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + elif args.model_type == "hunyuanimage2.1-distilled": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_vae_checkpoint(args) + + transformer_config = { + "in_channels": 64, + "out_channels": 64, + "num_attention_heads": 28, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "patch_size": (1, 1), + "qk_norm": "rms_norm", + "guidance_embeds": True, + "text_embed_dim": 3584, + "text_embed_2_dim": 1472, + "rope_theta": 256.0, + "rope_axes_dim": (64, 64), + "use_meanflow": True, + } + + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, config=transformer_config + ) + ) + + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae_config_diffusers = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 64, + "block_out_channels": [128, 256, 512, 512, 1024, 1024], + "layers_per_block": 2, + "spatial_compression_ratio": 32, + "sample_size": 384, + "scaling_factor": 0.75289, + "downsample_match_channel": True, + "upsample_match_channel": True, + } + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_vae_checkpoint_to_diffusers( + original_vae_state_dict, block_out_channels=[128, 256, 512, 512, 1024, 1024], layers_per_block=2 + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + vae = AutoencoderKLHunyuanImage(**vae_config_diffusers) + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + + if missing_keys: + logger.warning(f"Missing keys for vae: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + elif args.model_type == "hunyuanimage-refiner": + original_transformer_state_dict = load_original_transformer_checkpoint(args) + original_vae_state_dict = load_original_refiner_vae_checkpoint(args) + + transformer_config = { + "in_channels": 128, + "out_channels": 64, + "num_layers": 20, + "num_single_layers": 40, + "rope_axes_dim": [16, 56, 56], + "num_attention_heads": 26, + "attention_head_dim": 128, + "mlp_ratio": 4, + "patch_size": (1, 1, 1), + "text_embed_dim": 3584, + "guidance_embeds": True, + } + converted_transformer_state_dict, original_transformer_state_dict = ( + convert_hunyuan_image_transformer_checkpoint_to_diffusers( + original_transformer_state_dict, config=transformer_config + ) + ) + if original_transformer_state_dict: + logger.warning( + f"Unused {len(original_transformer_state_dict)} original keys for transformer: {list(original_transformer_state_dict.keys())}" + ) + + transformer = HunyuanImageTransformer2DModel(**transformer_config) + missing_keys, unexpected_key = transformer.load_state_dict(converted_transformer_state_dict, strict=True) + if missing_keys: + logger.warning(f"Missing keys for transformer: {missing_keys}") + if unexpected_key: + logger.warning(f"Unexpected keys for transformer: {unexpected_key}") + + transformer.to(dtype).save_pretrained(f"{args.output_path}/transformer") + + vae = AutoencoderKLHunyuanImageRefiner() + + converted_vae_state_dict, original_vae_state_dict = convert_hunyuan_image_refiner_vae_checkpoint_to_diffusers( + original_vae_state_dict + ) + if original_vae_state_dict: + logger.warning( + f"Unused {len(original_vae_state_dict)} original keys for vae: {list(original_vae_state_dict.keys())}" + ) + + missing_keys, unexpected_key = vae.load_state_dict(converted_vae_state_dict, strict=True) + logger.warning(f"Missing keys for vae: {missing_keys}") + logger.warning(f"Unexpected keys for vae: {unexpected_key}") + + vae.to(dtype).save_pretrained(f"{args.output_path}/vae") + + +if __name__ == "__main__": + main(args) diff --git a/scripts/recipes/hunyuan_video.py b/scripts/recipes/hunyuan_video.py new file mode 100644 index 000000000000..037d50296c4a --- /dev/null +++ b/scripts/recipes/hunyuan_video.py @@ -0,0 +1,132 @@ +import argparse +from typing import Any, Dict + +import torch +from transformers import ( + AutoModel, + AutoTokenizer, + CLIPImageProcessor, + CLIPTextModel, + CLIPTokenizer, + LlavaForConditionalGeneration, +) + +from diffusers import ( + AutoencoderKLHunyuanVideo, + FlowMatchEulerDiscreteScheduler, + HunyuanVideoImageToVideoPipeline, + HunyuanVideoPipeline, + HunyuanVideoTransformer3DModel, +) +from diffusers.loaders.conversion.configs.hunyuan_video import TRANSFORMER_CONFIGS + + +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: + state_dict = saved_dict + if "model" in saved_dict.keys(): + state_dict = state_dict["model"] + if "module" in saved_dict.keys(): + state_dict = state_dict["module"] + if "state_dict" in saved_dict.keys(): + state_dict = state_dict["state_dict"] + return state_dict + + +def convert_transformer(ckpt_path: str, transformer_type: str): + state = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True)) + return HunyuanVideoTransformer3DModel.from_single_file(state, config=TRANSFORMER_CONFIGS[transformer_type]) + + +def convert_vae(ckpt_path: str): + state = get_state_dict(torch.load(ckpt_path, map_location="cpu", weights_only=True)) + return AutoencoderKLHunyuanVideo.from_single_file(state, config={}) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" + ) + parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original VAE checkpoint") + parser.add_argument("--text_encoder_path", type=str, default=None, help="Path to original llama checkpoint") + parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to original llama tokenizer") + parser.add_argument("--text_encoder_2_path", type=str, default=None, help="Path to original clip checkpoint") + parser.add_argument("--save_pipeline", action="store_true") + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") + parser.add_argument("--dtype", default="bf16", help="Torch dtype to save the transformer in.") + parser.add_argument( + "--transformer_type", type=str, default="HYVideo-T/2-cfgdistill", choices=list(TRANSFORMER_CONFIGS.keys()) + ) + parser.add_argument("--flow_shift", type=float, default=7.0) + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +if __name__ == "__main__": + args = get_args() + + transformer = None + dtype = DTYPE_MAPPING[args.dtype] + + if args.save_pipeline: + assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None + assert args.text_encoder_path is not None + assert args.tokenizer_path is not None + assert args.text_encoder_2_path is not None + + if args.transformer_ckpt_path is not None: + transformer = convert_transformer(args.transformer_ckpt_path, args.transformer_type) + transformer = transformer.to(dtype=dtype) + if not args.save_pipeline: + transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.vae_ckpt_path is not None: + vae = convert_vae(args.vae_ckpt_path) + if not args.save_pipeline: + vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.save_pipeline: + if args.transformer_type == "HYVideo-T/2-cfgdistill": + text_encoder = AutoModel.from_pretrained(args.text_encoder_path, torch_dtype=torch.float16) + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, padding_side="right") + text_encoder_2 = CLIPTextModel.from_pretrained(args.text_encoder_2_path, torch_dtype=torch.float16) + tokenizer_2 = CLIPTokenizer.from_pretrained(args.text_encoder_2_path) + scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) + + pipe = HunyuanVideoPipeline( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + scheduler=scheduler, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + else: + text_encoder = LlavaForConditionalGeneration.from_pretrained( + args.text_encoder_path, torch_dtype=torch.float16 + ) + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path, padding_side="right") + text_encoder_2 = CLIPTextModel.from_pretrained(args.text_encoder_2_path, torch_dtype=torch.float16) + tokenizer_2 = CLIPTokenizer.from_pretrained(args.text_encoder_2_path) + scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) + image_processor = CLIPImageProcessor.from_pretrained(args.text_encoder_path) + + pipe = HunyuanVideoImageToVideoPipeline( + transformer=transformer, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + text_encoder_2=text_encoder_2, + tokenizer_2=tokenizer_2, + scheduler=scheduler, + image_processor=image_processor, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") diff --git a/scripts/recipes/hunyuan_video15.py b/scripts/recipes/hunyuan_video15.py new file mode 100644 index 000000000000..4bf6aa5dc94c --- /dev/null +++ b/scripts/recipes/hunyuan_video15.py @@ -0,0 +1,303 @@ +import argparse +import json +import os +import pathlib + +import torch +from huggingface_hub import hf_hub_download, snapshot_download +from safetensors.torch import load_file +from transformers import ( + AutoModel, + AutoTokenizer, + SiglipImageProcessor, + SiglipVisionModel, + T5EncoderModel, +) + +from diffusers import ( + AutoencoderKLHunyuanVideo15, + ClassifierFreeGuidance, + FlowMatchEulerDiscreteScheduler, + HunyuanVideo15ImageToVideoPipeline, + HunyuanVideo15Pipeline, + HunyuanVideo15Transformer3DModel, +) +from diffusers.loaders.conversion.configs.hunyuan_video15 import ( + GUIDANCE_CONFIGS, + SCHEDULER_CONFIGS, + TRANSFORMER_CONFIGS, +) + + +# to convert only transformer +""" +python scripts/recipes/hunyuan_video15.py \ + --original_state_dict_repo_id tencent/HunyuanVideo-1.5\ + --output_path /fsx/yiyi/HunyuanVideo-1.5-Diffusers/transformer\ + --transformer_type 480p_t2v +""" + +# to convert full pipeline +""" +python scripts/recipes/hunyuan_video15.py \ + --original_state_dict_repo_id tencent/HunyuanVideo-1.5\ + --output_path /fsx/yiyi/HunyuanVideo-1.5-Diffusers \ + --save_pipeline \ + --byt5_path /fsx/yiyi/hy15/text_encoder/Glyph-SDXL-v2\ + --transformer_type 480p_t2v +""" + + +def load_sharded_safetensors(path): + from diffusers.loaders.conversion.source import load_tensor_sources + + return load_tensor_sources(path) + + +def load_original_transformer_state_dict(args): + if args.original_state_dict_repo_id is not None: + model_dir = snapshot_download( + args.original_state_dict_repo_id, + repo_type="model", + allow_patterns="transformer/" + args.transformer_type + "/*", + ) + elif args.original_state_dict_folder is not None: + model_dir = pathlib.Path(args.original_state_dict_folder) + else: + raise ValueError("Please provide either `original_state_dict_repo_id` or `original_state_dict_folder`") + model_dir = pathlib.Path(model_dir) + model_dir = model_dir / "transformer" / args.transformer_type + return load_sharded_safetensors(model_dir) + + +def load_original_vae_state_dict(args): + if args.original_state_dict_repo_id is not None: + ckpt_path = hf_hub_download( + repo_id=args.original_state_dict_repo_id, filename="vae/diffusion_pytorch_model.safetensors" + ) + elif args.original_state_dict_folder is not None: + model_dir = pathlib.Path(args.original_state_dict_folder) + ckpt_path = model_dir / "vae/diffusion_pytorch_model.safetensors" + else: + raise ValueError("Please provide either `original_state_dict_repo_id` or `original_state_dict_folder`") + + original_state_dict = load_file(ckpt_path) + return original_state_dict + + +def convert_transformer(args): + state = load_original_transformer_state_dict(args) + return HunyuanVideo15Transformer3DModel.from_single_file(state, config=TRANSFORMER_CONFIGS[args.transformer_type]) + + +def convert_vae(args): + state = load_original_vae_state_dict(args) + return AutoencoderKLHunyuanVideo15.from_single_file(state, config={}) + + +def load_mllm(): + print(" loading from Qwen/Qwen2.5-VL-7B-Instruct") + text_encoder = AutoModel.from_pretrained( + "Qwen/Qwen2.5-VL-7B-Instruct", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True + ) + if hasattr(text_encoder, "language_model"): + text_encoder = text_encoder.language_model + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", padding_side="right") + return text_encoder, tokenizer + + +# copied from https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/910da2a829c484ea28982e8cff3bbc2cacdf1681/hyvideo/models/text_encoders/byT5/__init__.py#L89 +def add_special_token( + tokenizer, + text_encoder, + add_color=True, + add_font=True, + multilingual=True, + color_ann_path="assets/color_idx.json", + font_ann_path="assets/multilingual_10-lang_idx.json", +): + """ + Add special tokens for color and font to tokenizer and text encoder. + + Args: + tokenizer: Huggingface tokenizer. + text_encoder: Huggingface T5 encoder. + add_color (bool): Whether to add color tokens. + add_font (bool): Whether to add font tokens. + color_ann_path (str): Path to color annotation JSON. + font_ann_path (str): Path to font annotation JSON. + multilingual (bool): Whether to use multilingual font tokens. + """ + with open(font_ann_path, "r") as f: + idx_font_dict = json.load(f) + with open(color_ann_path, "r") as f: + idx_color_dict = json.load(f) + + if multilingual: + font_token = [f"<{font_code[:2]}-font-{idx_font_dict[font_code]}>" for font_code in idx_font_dict] + else: + font_token = [f"" for i in range(len(idx_font_dict))] + color_token = [f"" for i in range(len(idx_color_dict))] + additional_special_tokens = [] + if add_color: + additional_special_tokens += color_token + if add_font: + additional_special_tokens += font_token + + tokenizer.add_tokens(additional_special_tokens, special_tokens=True) + # Set mean_resizing=False to avoid PyTorch LAPACK dependency + text_encoder.resize_token_embeddings(len(tokenizer), mean_resizing=False) + + +def load_byt5(args): + """ + Load ByT5 encoder with Glyph-SDXL-v2 weights and save in HuggingFace format. + """ + + # 1. Load base tokenizer and encoder + tokenizer = AutoTokenizer.from_pretrained("google/byt5-small") + + # Load as T5EncoderModel + encoder = T5EncoderModel.from_pretrained("google/byt5-small") + + byt5_checkpoint_path = os.path.join(args.byt5_path, "checkpoints/byt5_model.pt") + color_ann_path = os.path.join(args.byt5_path, "assets/color_idx.json") + font_ann_path = os.path.join(args.byt5_path, "assets/multilingual_10-lang_idx.json") + + # 2. Add special tokens + add_special_token( + tokenizer=tokenizer, + text_encoder=encoder, + add_color=True, + add_font=True, + color_ann_path=color_ann_path, + font_ann_path=font_ann_path, + multilingual=True, + ) + + # 3. Load Glyph-SDXL-v2 checkpoint + print(f"\n3. Loading Glyph-SDXL-v2 checkpoint: {byt5_checkpoint_path}") + checkpoint = torch.load(byt5_checkpoint_path, map_location="cpu") + + # Handle different checkpoint formats + if "state_dict" in checkpoint: + state_dict = checkpoint["state_dict"] + else: + state_dict = checkpoint + + # add 'encoder.' prefix to the keys + # Remove 'module.text_tower.encoder.' prefix if present + cleaned_state_dict = {} + for key, value in state_dict.items(): + if key.startswith("module.text_tower.encoder."): + new_key = "encoder." + key[len("module.text_tower.encoder.") :] + cleaned_state_dict[new_key] = value + else: + new_key = "encoder." + key + cleaned_state_dict[new_key] = value + + # 4. Load weights + missing_keys, unexpected_keys = encoder.load_state_dict(cleaned_state_dict, strict=False) + if unexpected_keys: + raise ValueError(f"Unexpected keys: {unexpected_keys}") + if "shared.weight" in missing_keys: + print(" Missing shared.weight as expected") + missing_keys.remove("shared.weight") + if missing_keys: + raise ValueError(f"Missing keys: {missing_keys}") + + return encoder, tokenizer + + +def load_siglip(): + image_encoder = SiglipVisionModel.from_pretrained( + "black-forest-labs/FLUX.1-Redux-dev", subfolder="image_encoder", torch_dtype=torch.bfloat16 + ) + feature_extractor = SiglipImageProcessor.from_pretrained( + "black-forest-labs/FLUX.1-Redux-dev", subfolder="feature_extractor" + ) + return image_encoder, feature_extractor + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--original_state_dict_repo_id", type=str, default=None, help="Path to original hub_id for the model" + ) + parser.add_argument( + "--original_state_dict_folder", type=str, default=None, help="Local folder name of the original state dict" + ) + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model(s) should be saved") + parser.add_argument("--transformer_type", type=str, default="480p_i2v", choices=list(TRANSFORMER_CONFIGS.keys())) + parser.add_argument( + "--byt5_path", + type=str, + default=None, + help=( + "path to the downloaded byt5 checkpoint & assets. " + "Note: They use Glyph-SDXL-v2 as byt5 encoder. You can download from modelscope like: " + "`modelscope download --model AI-ModelScope/Glyph-SDXL-v2 --local_dir ./ckpts/text_encoder/Glyph-SDXL-v2` " + "or manually download following the instructions on " + "https://github.com/Tencent-Hunyuan/HunyuanVideo-1.5/blob/910da2a829c484ea28982e8cff3bbc2cacdf1681/checkpoints-download.md. " + "The path should point to the Glyph-SDXL-v2 folder which should contain an `assets` folder and a `checkpoints` folder, " + "like: Glyph-SDXL-v2/assets/... and Glyph-SDXL-v2/checkpoints/byt5_model.pt" + ), + ) + parser.add_argument("--save_pipeline", action="store_true") + return parser.parse_args() + + +if __name__ == "__main__": + args = get_args() + + if args.save_pipeline and args.byt5_path is None: + raise ValueError("Please provide --byt5_path when saving pipeline") + + transformer = None + + transformer = convert_transformer(args) + if not args.save_pipeline: + transformer.save_pretrained(args.output_path, safe_serialization=True) + else: + task_type = transformer.config.task_type + + vae = convert_vae(args) + + text_encoder, tokenizer = load_mllm() + text_encoder_2, tokenizer_2 = load_byt5(args) + + flow_shift = SCHEDULER_CONFIGS[args.transformer_type]["shift"] + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + + guidance_scale = GUIDANCE_CONFIGS[args.transformer_type]["guidance_scale"] + guider = ClassifierFreeGuidance(guidance_scale=guidance_scale) + + if task_type == "i2v": + image_encoder, feature_extractor = load_siglip() + pipeline = HunyuanVideo15ImageToVideoPipeline( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + transformer=transformer, + guider=guider, + scheduler=scheduler, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + ) + elif task_type == "t2v": + pipeline = HunyuanVideo15Pipeline( + vae=vae, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + transformer=transformer, + guider=guider, + scheduler=scheduler, + ) + else: + raise ValueError(f"Task type {task_type} is not supported") + + pipeline.save_pretrained(args.output_path, safe_serialization=True) diff --git a/scripts/recipes/i2vgen.py b/scripts/recipes/i2vgen.py new file mode 100644 index 000000000000..addadc87d9a5 --- /dev/null +++ b/scripts/recipes/i2vgen.py @@ -0,0 +1,95 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the LDM checkpoints.""" + +import argparse + +import torch +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer, CLIPVisionModelWithProjection + +from diffusers import DDIMScheduler, I2VGenXLPipeline, I2VGenXLUNet, StableDiffusionPipeline +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +CLIP_ID = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K" + + +def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, **kwargs): + return convert_component_checkpoint(checkpoint, config, "I2VGenXLUNet", extract_ema=extract_ema) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--unet_checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--push_to_hub", action="store_true") + args = parser.parse_args() + + # UNet + unet_checkpoint = torch.load(args.unet_checkpoint_path, map_location="cpu") + unet_checkpoint = unet_checkpoint["state_dict"] + unet = I2VGenXLUNet(sample_size=32) + + converted_ckpt = convert_ldm_unet_checkpoint(unet_checkpoint, unet.config) + + diff_0 = set(unet.state_dict().keys()) - set(converted_ckpt.keys()) + diff_1 = set(converted_ckpt.keys()) - set(unet.state_dict().keys()) + + assert len(diff_0) == len(diff_1) == 0, "Converted weights don't match" + + unet.load_state_dict(converted_ckpt, strict=True) + + # vae + temp_pipe = StableDiffusionPipeline.from_single_file( + "https://huggingface.co/ali-vilab/i2vgen-xl/blob/main/models/v2-1_512-ema-pruned.ckpt" + ) + vae = temp_pipe.vae + del temp_pipe + + # text encoder and tokenizer + text_encoder = CLIPTextModel.from_pretrained(CLIP_ID) + tokenizer = CLIPTokenizer.from_pretrained(CLIP_ID) + + # image encoder and feature extractor + image_encoder = CLIPVisionModelWithProjection.from_pretrained(CLIP_ID) + feature_extractor = CLIPImageProcessor.from_pretrained(CLIP_ID) + + # scheduler + # https://github.com/ali-vilab/i2vgen-xl/blob/main/configs/i2vgen_xl_train.yaml + scheduler = DDIMScheduler( + beta_schedule="squaredcos_cap_v2", + rescale_betas_zero_snr=True, + set_alpha_to_one=True, + clip_sample=False, + steps_offset=1, + timestep_spacing="leading", + prediction_type="v_prediction", + ) + + # final + pipeline = I2VGenXLPipeline( + unet=unet, + vae=vae, + image_encoder=image_encoder, + feature_extractor=feature_extractor, + text_encoder=text_encoder, + tokenizer=tokenizer, + scheduler=scheduler, + ) + + pipeline.save_pretrained(args.dump_path, push_to_hub=args.push_to_hub) diff --git a/scripts/recipes/joyimage.py b/scripts/recipes/joyimage.py new file mode 100644 index 000000000000..f1e339c198e1 --- /dev/null +++ b/scripts/recipes/joyimage.py @@ -0,0 +1,152 @@ +"""Convert JoyImage Edit / Edit Plus checkpoints to diffusers format. + +Supports both JoyImage-Edit (single-image editing) and JoyImage-Edit-Plus +(multi-image editing). The transformer weight layout is identical; only the +target model class and pipeline differ. + +Usage: + # Convert JoyImage Edit (default) + python recipes/joyimage.py \ + --transformer_ckpt_path /path/to/transformer.pt \ + --vae_ckpt_path /path/to/vae.pt \ + --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ + --output_path /path/to/output \ + --save_pipeline + + # Convert JoyImage Edit Plus + python recipes/joyimage.py \ + --model_type edit_plus \ + --transformer_ckpt_path /path/to/transformer.pt \ + --vae_ckpt_path /path/to/vae.pt \ + --text_encoder_path Qwen/Qwen3-VL-8B-Instruct \ + --output_path /path/to/output \ + --save_pipeline +""" + +import argparse + +import torch +from transformers import AutoProcessor, AutoTokenizer, Qwen3VLForConditionalGeneration + +from diffusers import ( + AutoencoderKLWan, + JoyImageEditPipeline, + JoyImageEditTransformer3DModel, +) +from diffusers.loaders.conversion.configs.joyimage import TRANSFORMER_CONFIG +from diffusers.models.transformers.transformer_joyimage_edit_plus import JoyImageEditPlusTransformer3DModel +from diffusers.pipelines.joyimage.pipeline_joyimage_edit_plus import JoyImageEditPlusPipeline +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) + + +def convert_vae(vae_ckpt_path): + return AutoencoderKLWan.from_single_file(vae_ckpt_path, config={}) + + +def convert_transformer(ckpt_path, model_type="edit"): + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True) + state = checkpoint.get("model", checkpoint) + cls = JoyImageEditPlusTransformer3DModel if model_type == "edit_plus" else JoyImageEditTransformer3DModel + return cls.from_single_file(state, config=TRANSFORMER_CONFIG) + + +def get_args(): + parser = argparse.ArgumentParser(description="Convert JoyImage Edit / Edit Plus checkpoints to diffusers format") + parser.add_argument( + "--model_type", + type=str, + choices=["edit", "edit_plus"], + default="edit", + help="Model type: 'edit' for JoyImage-Edit, 'edit_plus' for JoyImage-Edit-Plus", + ) + parser.add_argument( + "--transformer_ckpt_path", + type=str, + default=None, + help="Path to original transformer checkpoint", + ) + parser.add_argument( + "--vae_ckpt_path", + type=str, + default=None, + help="Path to original VAE checkpoint", + ) + parser.add_argument( + "--text_encoder_path", + type=str, + default=None, + help="Path to Qwen3-VL text encoder (e.g. Qwen/Qwen3-VL-8B-Instruct)", + ) + parser.add_argument("--save_pipeline", action="store_true") + parser.add_argument( + "--output_path", + type=str, + required=True, + help="Path where converted model should be saved", + ) + parser.add_argument("--dtype", default="bf16", help="Torch dtype (fp32, fp16, bf16)") + parser.add_argument("--flow_shift", type=float, default=1.5) + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +if __name__ == "__main__": + args = get_args() + transformer = None + vae = None + dtype = DTYPE_MAPPING[args.dtype] + + if args.save_pipeline: + assert args.transformer_ckpt_path is not None and args.vae_ckpt_path is not None + assert args.text_encoder_path is not None + + if args.transformer_ckpt_path is not None: + transformer = convert_transformer(args.transformer_ckpt_path, model_type=args.model_type) + transformer = transformer.to(dtype=dtype) + if not args.save_pipeline: + transformer.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.vae_ckpt_path is not None: + vae = convert_vae(args.vae_ckpt_path) + vae = vae.to(dtype=dtype) + if not args.save_pipeline: + vae.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.save_pipeline: + processor = AutoProcessor.from_pretrained(args.text_encoder_path) + text_encoder = Qwen3VLForConditionalGeneration.from_pretrained( + args.text_encoder_path, torch_dtype=torch.bfloat16 + ).to("cuda") + tokenizer = AutoTokenizer.from_pretrained(args.text_encoder_path) + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=args.flow_shift) + transformer = transformer.to("cuda") + vae = vae.to("cuda") + + if args.model_type == "edit_plus": + pipe = JoyImageEditPlusPipeline( + processor=processor, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ).to("cuda") + else: + pipe = JoyImageEditPipeline( + processor=processor, + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ).to("cuda") + + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + processor.save_pretrained(f"{args.output_path}/processor") diff --git a/scripts/recipes/k_upscaler.py b/scripts/recipes/k_upscaler.py new file mode 100644 index 000000000000..dc4d30f09115 --- /dev/null +++ b/scripts/recipes/k_upscaler.py @@ -0,0 +1,98 @@ +import argparse + +import huggingface_hub +import k_diffusion as K +import torch + +from diffusers import UNet2DConditionModel +from diffusers.loaders.conversion import get_conversion + + +UPSCALER_REPO = "pcuenq/k-upscaler" + + +def unet_to_diffusers_checkpoint(model, checkpoint): + return get_conversion("UNet2DConditionModel", dict(model.config)).to_diffusers(checkpoint) + + +def unet_model_from_original_config(original_config): + in_channels = original_config["input_channels"] + original_config["unet_cond_dim"] + out_channels = original_config["input_channels"] + (1 if original_config["has_variance"] else 0) + + block_out_channels = original_config["channels"] + + assert len(set(original_config["depths"])) == 1, ( + "UNet2DConditionModel currently do not support blocks with different number of layers" + ) + layers_per_block = original_config["depths"][0] + + class_labels_dim = original_config["mapping_cond_dim"] + cross_attention_dim = original_config["cross_cond_dim"] + + attn1_types = [] + attn2_types = [] + for s, c in zip(original_config["self_attn_depths"], original_config["cross_attn_depths"]): + if s: + a1 = "self" + a2 = "cross" if c else None + elif c: + a1 = "cross" + a2 = None + else: + a1 = None + a2 = None + attn1_types.append(a1) + attn2_types.append(a2) + + unet = UNet2DConditionModel( + in_channels=in_channels, + out_channels=out_channels, + down_block_types=("KDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D"), + mid_block_type=None, + up_block_types=("KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D"), + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn="gelu", + norm_num_groups=None, + cross_attention_dim=cross_attention_dim, + attention_head_dim=64, + time_cond_proj_dim=class_labels_dim, + resnet_time_scale_shift="scale_shift", + time_embedding_type="fourier", + timestep_post_act="gelu", + conv_in_kernel=1, + conv_out_kernel=1, + ) + + return unet + + +def main(args): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + orig_config_path = huggingface_hub.hf_hub_download(UPSCALER_REPO, "config_laion_text_cond_latent_upscaler_2.json") + orig_weights_path = huggingface_hub.hf_hub_download( + UPSCALER_REPO, "laion_text_cond_latent_upscaler_2_1_00470000_slim.pth" + ) + print(f"loading original model configuration from {orig_config_path}") + print(f"loading original model checkpoint from {orig_weights_path}") + + print("converting to diffusers unet") + orig_config = K.config.load_config(open(orig_config_path))["model"] + model = unet_model_from_original_config(orig_config) + + orig_checkpoint = torch.load(orig_weights_path, map_location=device)["model_ema"] + converted_checkpoint = unet_to_diffusers_checkpoint(model, orig_checkpoint) + + model.load_state_dict(converted_checkpoint, strict=True) + model.save_pretrained(args.dump_path) + print(f"saving converted unet model in {args.dump_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + args = parser.parse_args() + + main(args) diff --git a/scripts/recipes/kandinsky.py b/scripts/recipes/kandinsky.py new file mode 100644 index 000000000000..5865a42573c0 --- /dev/null +++ b/scripts/recipes/kandinsky.py @@ -0,0 +1,310 @@ +import argparse +import os +import tempfile + +import torch +from accelerate import load_checkpoint_and_dispatch + +from diffusers import UNet2DConditionModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.kandinsky import INPAINT_UNET_CONFIG, MOVQ_CONFIG, PRIOR_CONFIG, UNET_CONFIG +from diffusers.models.transformers.prior_transformer import PriorTransformer +from diffusers.models.vq_model import VQModel + + +""" +Example - From the diffusers root directory: + +Download weights: +```sh +$ wget https://huggingface.co/ai-forever/Kandinsky_2.1/blob/main/prior_fp16.ckpt +``` + +Convert the model: +```sh +python scripts/recipes/kandinsky.py \ + --prior_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/prior_fp16.ckpt \ + --clip_stat_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/ViT-L-14_stats.th \ + --text2img_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/decoder_fp16.ckpt \ + --inpaint_text2img_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/inpainting_fp16.ckpt \ + --movq_checkpoint_path /home/yiyi_huggingface_co/Kandinsky-2/checkpoints_Kandinsky_2.1/movq_final.ckpt \ + --dump_path /home/yiyi_huggingface_co/dump \ + --debug decoder +``` +""" + + +# prior + + +# Uses default arguments + + +def prior_model_from_original_config(): + model = PriorTransformer(**PRIOR_CONFIG) + + return model + + +def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint, clip_stats_checkpoint): + state = dict(checkpoint) + state["clip_stats.mean"], state["clip_stats.std"] = clip_stats_checkpoint + return get_conversion("PriorTransformer", {**dict(model.config), "original_format": "kandinsky"}).to_diffusers( + state + ) + + +# done prior + +# unet + +# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can +# update then. + + +def unet_model_from_original_config(): + model = UNet2DConditionModel(**UNET_CONFIG) + + return model + + +def unet_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + return get_conversion("UNet2DConditionModel", dict(model.config)).to_diffusers(checkpoint) + + +# done unet + +# inpaint unet + +# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can +# update then. + + +def inpaint_unet_model_from_original_config(): + model = UNet2DConditionModel(**INPAINT_UNET_CONFIG) + + return model + + +def inpaint_unet_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + return get_conversion("UNet2DConditionModel", dict(model.config)).to_diffusers(checkpoint) + + +# done inpaint unet + + +# unet utils + + +# .time_embed -> .time_embedding + + +# .input_blocks.0 -> .conv_in + + +# .out.0 -> .conv_norm_out + + +# .out.2 -> .conv_out + + +# .input_blocks -> .down_blocks + + +# .middle_block -> .mid_block + + +# .output_blocks -> .up_blocks + + +# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) + + +# done unet utils + + +def prior(*, args, checkpoint_map_location): + print("loading prior") + + prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) + + clip_stats_checkpoint = torch.load(args.clip_stat_path, map_location=checkpoint_map_location) + + prior_model = prior_model_from_original_config() + + prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint( + prior_model, prior_checkpoint, clip_stats_checkpoint + ) + + del prior_checkpoint + del clip_stats_checkpoint + + load_checkpoint_to_model(prior_diffusers_checkpoint, prior_model, strict=True) + + print("done loading prior") + + return prior_model + + +def text2img(*, args, checkpoint_map_location): + print("loading text2img") + + text2img_checkpoint = torch.load(args.text2img_checkpoint_path, map_location=checkpoint_map_location) + + unet_model = unet_model_from_original_config() + + unet_diffusers_checkpoint = unet_original_checkpoint_to_diffusers_checkpoint(unet_model, text2img_checkpoint) + + del text2img_checkpoint + + load_checkpoint_to_model(unet_diffusers_checkpoint, unet_model, strict=True) + + print("done loading text2img") + + return unet_model + + +def inpaint_text2img(*, args, checkpoint_map_location): + print("loading inpaint text2img") + + inpaint_text2img_checkpoint = torch.load( + args.inpaint_text2img_checkpoint_path, map_location=checkpoint_map_location + ) + + inpaint_unet_model = inpaint_unet_model_from_original_config() + + inpaint_unet_diffusers_checkpoint = inpaint_unet_original_checkpoint_to_diffusers_checkpoint( + inpaint_unet_model, inpaint_text2img_checkpoint + ) + + del inpaint_text2img_checkpoint + + load_checkpoint_to_model(inpaint_unet_diffusers_checkpoint, inpaint_unet_model, strict=True) + + print("done loading inpaint text2img") + + return inpaint_unet_model + + +# movq + + +def movq_model_from_original_config(): + movq = VQModel(**MOVQ_CONFIG) + return movq + + +def movq_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + return get_conversion("VQModel", dict(model.config)).to_diffusers(checkpoint) + + +def movq(*, args, checkpoint_map_location): + print("loading movq") + + movq_checkpoint = torch.load(args.movq_checkpoint_path, map_location=checkpoint_map_location) + + movq_model = movq_model_from_original_config() + + movq_diffusers_checkpoint = movq_original_checkpoint_to_diffusers_checkpoint(movq_model, movq_checkpoint) + + del movq_checkpoint + + load_checkpoint_to_model(movq_diffusers_checkpoint, movq_model, strict=True) + + print("done loading movq") + + return movq_model + + +def load_checkpoint_to_model(checkpoint, model, strict=False): + with tempfile.NamedTemporaryFile(delete=False) as file: + torch.save(checkpoint, file.name) + del checkpoint + if strict: + model.load_state_dict(torch.load(file.name), strict=True) + else: + load_checkpoint_and_dispatch(model, file.name, device_map="auto") + os.remove(file.name) + + +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( + "--prior_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the prior checkpoint to convert.", + ) + parser.add_argument( + "--clip_stat_path", + default=None, + type=str, + required=False, + help="Path to the clip stats checkpoint to convert.", + ) + parser.add_argument( + "--text2img_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the text2img checkpoint to convert.", + ) + parser.add_argument( + "--movq_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the text2img checkpoint to convert.", + ) + parser.add_argument( + "--inpaint_text2img_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the inpaint text2img checkpoint to convert.", + ) + parser.add_argument( + "--checkpoint_load_device", + default="cpu", + type=str, + required=False, + help="The device passed to `map_location` when loading checkpoints.", + ) + + parser.add_argument( + "--debug", + default=None, + type=str, + required=False, + help="Only run a specific stage of the convert script. Used for debugging", + ) + + args = parser.parse_args() + + print(f"loading checkpoints to {args.checkpoint_load_device}") + + checkpoint_map_location = torch.device(args.checkpoint_load_device) + + if args.debug is not None: + print(f"debug: only executing {args.debug}") + + if args.debug is None: + print("to-do") + elif args.debug == "prior": + prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) + prior_model.save_pretrained(args.dump_path) + elif args.debug == "text2img": + unet_model = text2img(args=args, checkpoint_map_location=checkpoint_map_location) + unet_model.save_pretrained(f"{args.dump_path}/unet") + elif args.debug == "inpaint_text2img": + inpaint_unet_model = inpaint_text2img(args=args, checkpoint_map_location=checkpoint_map_location) + inpaint_unet_model.save_pretrained(f"{args.dump_path}/inpaint_unet") + elif args.debug == "decoder": + decoder = movq(args=args, checkpoint_map_location=checkpoint_map_location) + decoder.save_pretrained(f"{args.dump_path}/decoder") + else: + raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/recipes/ldm.py b/scripts/recipes/ldm.py new file mode 100644 index 000000000000..0ed87c12a101 --- /dev/null +++ b/scripts/recipes/ldm.py @@ -0,0 +1,69 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the LDM checkpoints.""" + +import argparse +import json + +import torch + +from diffusers import DDPMScheduler, LDMPipeline, UNet2DModel, VQModel +from diffusers.loaders.conversion import get_conversion + + +def convert_ldm_checkpoint(checkpoint, config): + return get_conversion("UNet2DModel", {**config, "original_format": "ldm"}).to_diffusers(checkpoint) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + + parser.add_argument( + "--config_file", + default=None, + type=str, + required=True, + help="The config json file corresponding to the architecture.", + ) + + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + + args = parser.parse_args() + + checkpoint = torch.load(args.checkpoint_path) + + with open(args.config_file) as f: + config = json.loads(f.read()) + + converted_checkpoint = convert_ldm_checkpoint(checkpoint, config) + + if "ldm" in config: + del config["ldm"] + + model = UNet2DModel(**config) + model.load_state_dict(converted_checkpoint) + + try: + scheduler = DDPMScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) + vqvae = VQModel.from_pretrained("/".join(args.checkpoint_path.split("/")[:-1])) + + pipe = LDMPipeline(unet=model, scheduler=scheduler, vae=vqvae) + pipe.save_pretrained(args.dump_path) + except: # noqa: E722 + model.save_pretrained(args.dump_path) diff --git a/scripts/convert_longcat_audio_dit_to_diffusers.py b/scripts/recipes/longcat_audio.py similarity index 91% rename from scripts/convert_longcat_audio_dit_to_diffusers.py rename to scripts/recipes/longcat_audio.py index b7acee48675b..6c210f827ff2 100644 --- a/scripts/convert_longcat_audio_dit_to_diffusers.py +++ b/scripts/recipes/longcat_audio.py @@ -14,9 +14,9 @@ # limitations under the License. # Usage: -# python scripts/convert_longcat_audio_dit_to_diffusers.py --checkpoint_path /path/to/model --output_path /data/models -# python scripts/convert_longcat_audio_dit_to_diffusers.py --repo_id meituan-longcat/LongCat-AudioDiT-1B --output_path /data/models -# python scripts/convert_longcat_audio_dit_to_diffusers.py --checkpoint_path /path/to/model --output_path /data/models --dtype fp16 +# python scripts/recipes/longcat_audio.py --checkpoint_path /path/to/model --output_path /data/models +# python scripts/recipes/longcat_audio.py --repo_id meituan-longcat/LongCat-AudioDiT-1B --output_path /data/models +# python scripts/recipes/longcat_audio.py --checkpoint_path /path/to/model --output_path /data/models --dtype fp16 import argparse import json @@ -33,6 +33,7 @@ LongCatAudioDiTTransformer, LongCatAudioDiTVae, ) +from diffusers.loaders.conversion import get_conversion def find_checkpoint(input_dir: Path): @@ -133,13 +134,18 @@ def convert_longcat_audio_dit( use_latent_condition=config.get("dit_use_latent_condition", True), ff_mult=config.get("dit_ff_mult", 4), ) - transformer.load_state_dict(transformer_state_dict, strict=True) + transformer.load_state_dict( + get_conversion("LongCatAudioDiTTransformer", dict(transformer.config)).to_diffusers(transformer_state_dict), + strict=True, + ) transformer = transformer.to(dtype=torch_dtype) vae_config = dict(config["vae_config"]) vae_config.pop("model_type", None) vae = LongCatAudioDiTVae(**vae_config) - vae.load_state_dict(vae_state_dict, strict=True) + vae.load_state_dict( + get_conversion("LongCatAudioDiTVae", dict(vae.config)).to_diffusers(vae_state_dict), strict=True + ) vae = vae.to(dtype=torch_dtype) text_encoder_config = UMT5Config.from_dict(config["text_encoder_config"]) diff --git a/scripts/recipes/ltx.py b/scripts/recipes/ltx.py new file mode 100644 index 000000000000..fd91545e637a --- /dev/null +++ b/scripts/recipes/ltx.py @@ -0,0 +1,193 @@ +import argparse +from pathlib import Path +from typing import Any, Dict + +import torch +from safetensors.torch import load_file +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers import ( + AutoencoderKLLTXVideo, + FlowMatchEulerDiscreteScheduler, + LTXConditionPipeline, + LTXLatentUpsamplePipeline, + LTXPipeline, + LTXVideoTransformer3DModel, +) +from diffusers.loaders.conversion.configs.ltx import ( + get_spatial_latent_upsampler_config, + get_transformer_config, + get_vae_config, +) +from diffusers.pipelines.ltx.modeling_latent_upsampler import LTXLatentUpsamplerModel + + +TOKENIZER_MAX_LENGTH = 128 + + +def get_state_dict(saved_dict: Dict[str, Any]) -> dict[str, Any]: + state_dict = saved_dict + if "model" in saved_dict.keys(): + state_dict = state_dict["model"] + if "module" in saved_dict.keys(): + state_dict = state_dict["module"] + if "state_dict" in saved_dict.keys(): + state_dict = state_dict["state_dict"] + return state_dict + + +def convert_transformer(ckpt_path: str, config, dtype: torch.dtype): + state = get_state_dict(load_file(ckpt_path)) + return LTXVideoTransformer3DModel.from_single_file(state, config=config, torch_dtype=dtype) + + +def convert_vae(ckpt_path: str, config, dtype: torch.dtype): + state = get_state_dict(load_file(ckpt_path)) + return AutoencoderKLLTXVideo.from_single_file(state, config=config, torch_dtype=dtype) + + +def convert_spatial_latent_upsampler(ckpt_path: str, config, dtype: torch.dtype): + state = get_state_dict(load_file(ckpt_path)) + return LTXLatentUpsamplerModel.from_single_file(state, config=config, torch_dtype=dtype) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint" + ) + parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint") + parser.add_argument( + "--spatial_latent_upsampler_path", + type=str, + default=None, + help="Path to original spatial latent upsampler checkpoint", + ) + parser.add_argument( + "--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory" + ) + parser.add_argument( + "--typecast_text_encoder", + action="store_true", + default=False, + help="Whether or not to apply fp16/bf16 precision to text_encoder", + ) + parser.add_argument("--save_pipeline", action="store_true") + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") + parser.add_argument("--dtype", default="fp32", help="Torch dtype to save the model in.") + parser.add_argument( + "--version", + type=str, + default="0.9.0", + choices=["0.9.0", "0.9.1", "0.9.5", "0.9.7", "0.9.8"], + help="Version of the LTX model", + ) + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +VARIANT_MAPPING = { + "fp32": None, + "fp16": "fp16", + "bf16": "bf16", +} + + +if __name__ == "__main__": + args = get_args() + + transformer = None + dtype = DTYPE_MAPPING[args.dtype] + variant = VARIANT_MAPPING[args.dtype] + output_path = Path(args.output_path) + + if args.transformer_ckpt_path is not None: + config = get_transformer_config(args.version) + transformer: LTXVideoTransformer3DModel = convert_transformer(args.transformer_ckpt_path, config, dtype) + if not args.save_pipeline: + transformer.save_pretrained( + output_path / "transformer", safe_serialization=True, max_shard_size="5GB", variant=variant + ) + + if args.vae_ckpt_path is not None: + config = get_vae_config(args.version) + vae: AutoencoderKLLTXVideo = convert_vae(args.vae_ckpt_path, config, dtype) + if not args.save_pipeline: + vae.save_pretrained(output_path / "vae", safe_serialization=True, max_shard_size="5GB", variant=variant) + + if args.spatial_latent_upsampler_path is not None: + config = get_spatial_latent_upsampler_config(args.version) + latent_upsampler: LTXLatentUpsamplerModel = convert_spatial_latent_upsampler( + args.spatial_latent_upsampler_path, config, dtype + ) + if not args.save_pipeline: + latent_upsampler.save_pretrained( + output_path / "latent_upsampler", safe_serialization=True, max_shard_size="5GB", variant=variant + ) + + if args.save_pipeline: + text_encoder_id = "google/t5-v1_1-xxl" + tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH) + text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir) + + if args.typecast_text_encoder: + text_encoder = text_encoder.to(dtype=dtype) + + # Apparently, the conversion does not work anymore without this :shrug: + for param in text_encoder.parameters(): + param.data = param.data.contiguous() + + if args.version in ["0.9.5", "0.9.7"]: + scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=False) + else: + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, + base_shift=0.95, + max_shift=2.05, + base_image_seq_len=1024, + max_image_seq_len=4096, + shift_terminal=0.1, + ) + + if args.version in ["0.9.0", "0.9.1", "0.9.5"]: + pipe = LTXPipeline( + scheduler=scheduler, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + ) + pipe.save_pretrained( + output_path.as_posix(), safe_serialization=True, variant=variant, max_shard_size="5GB" + ) + elif args.version in ["0.9.7"]: + pipe = LTXConditionPipeline( + scheduler=scheduler, + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + ) + pipe_upsample = LTXLatentUpsamplePipeline( + vae=vae, + latent_upsampler=latent_upsampler, + ) + pipe.save_pretrained( + (output_path / "ltx_pipeline").as_posix(), + safe_serialization=True, + variant=variant, + max_shard_size="5GB", + ) + pipe_upsample.save_pretrained( + (output_path / "ltx_upsample_pipeline").as_posix(), + safe_serialization=True, + variant=variant, + max_shard_size="5GB", + ) + else: + raise ValueError(f"Unsupported version: {args.version}") diff --git a/scripts/recipes/ltx2.py b/scripts/recipes/ltx2.py new file mode 100644 index 000000000000..946572a990e9 --- /dev/null +++ b/scripts/recipes/ltx2.py @@ -0,0 +1,555 @@ +import argparse +import os +from typing import Any + +import safetensors.torch +import torch +from huggingface_hub import hf_hub_download +from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor, AutoTokenizer + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2LatentUpsamplePipeline, + LTX2Pipeline, + LTX2VideoDiffusionDecoderModel, + LTX2VideoTransformer3DModel, +) +from diffusers.loaders.conversion.configs.ltx2 import ( + get_ltx2_audio_vae_config, + get_ltx2_connectors_config, + get_ltx2_diffusion_video_vae_config, + get_ltx2_spatial_latent_upsampler_config, + get_ltx2_temporal_latent_upsampler_config, + get_ltx2_transformer_config, + get_ltx2_video_vae_config, + get_ltx2_vocoder_config, +) +from diffusers.pipelines.ltx2 import ( + LTX2DurationHead, + LTX2LatentUpsamplerModel, + LTX2TextConnectors, + LTX2Vocoder, + LTX2VocoderWithBWE, +) + + +# LTX-2.5's diffusion decoder replaces the conv decoder while keeping the same encoder, so only the +# `decoder.*` half of the VAE checkpoint is renamed with these rules. + +# Where a checkpoint carries static AdaLN gates, each is folded into the Linear it gates (W <- g * W) and +# dropped, because the decoder's residuals are ungated. Maps a renamed parameter to its gate's suffix. + + +def split_transformer_and_connector_state_dict(state_dict: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + connector_prefixes = ( + "video_embeddings_connector", + "audio_embeddings_connector", + "transformer_1d_blocks", + "text_embedding_projection", + "connectors.", + "video_connector", + "audio_connector", + "text_proj_in", + ) + + transformer_state_dict, connector_state_dict = {}, {} + for key, value in state_dict.items(): + if key.startswith(connector_prefixes): + connector_state_dict[key] = value + else: + transformer_state_dict[key] = value + + return transformer_state_dict, connector_state_dict + + +def convert_ltx2_transformer(original_state_dict, version): + config = get_ltx2_transformer_config(version)["diffusers_config"] + state, _ = split_transformer_and_connector_state_dict(original_state_dict) + return LTX2VideoTransformer3DModel.from_single_file(state, config=config) + + +def convert_ltx2_connectors(original_state_dict, version, gemma_text_config=None): + config = get_ltx2_connectors_config(version, gemma_text_config=gemma_text_config)["diffusers_config"] + _, state = split_transformer_and_connector_state_dict(original_state_dict) + return LTX2TextConnectors.from_single_file(state, config=config) + + +def convert_ltx2_duration_head(original_state_dict): + if not original_state_dict: + return None + state = original_state_dict + config = { + "video_cross_attention_dim": state["video_input_proj.weight"].shape[1], + "audio_cross_attention_dim": state["audio_input_proj.weight"].shape[1], + "pooler_hidden_dim": state["attention_pooler.cross_attn.in_proj_weight"].shape[1], + "num_queries": state["attention_pooler.query_tokens"].shape[0], + "mlp_hidden_dim": state["mlp_hidden.weight"].shape[0], + "num_pooler_heads": 4, + } + return LTX2DurationHead.from_single_file(state, config=config) + + +def convert_ltx2_video_vae(original_state_dict, version, timestep_conditioning): + config = get_ltx2_video_vae_config(version, timestep_conditioning)["diffusers_config"] + + return AutoencoderKLLTX2Video.from_single_file(original_state_dict, config=config) + + +def convert_ltx2_diffusion_video_vae(original_state_dict, version): + config = get_ltx2_diffusion_video_vae_config(version)["diffusers_config"] + + return LTX2VideoDiffusionDecoderModel.from_single_file(original_state_dict, config=config) + + +def convert_ltx2_audio_vae(original_state_dict, version): + config = get_ltx2_audio_vae_config(version)["diffusers_config"] + + return AutoencoderKLLTX2Audio.from_single_file(original_state_dict, config=config) + + +def convert_ltx2_vocoder(original_state_dict, version): + config = get_ltx2_vocoder_config(version)["diffusers_config"] + vocoder_cls = LTX2VocoderWithBWE if version in ("2.3", "2.5") else LTX2Vocoder + return vocoder_cls.from_single_file(original_state_dict, config=config) + + +def convert_ltx2_latent_upsampler(original_state_dict, config, dtype): + return LTX2LatentUpsamplerModel.from_single_file(original_state_dict, config=config, torch_dtype=dtype) + + +def load_original_checkpoint(args, filename: str | None) -> dict[str, Any]: + if args.original_state_dict_repo_id is not None: + ckpt_path = hf_hub_download(repo_id=args.original_state_dict_repo_id, filename=filename) + elif args.checkpoint_path is not None: + ckpt_path = args.checkpoint_path + else: + raise ValueError("Please provide either `original_state_dict_repo_id` or a local `checkpoint_path`") + + original_state_dict = safetensors.torch.load_file(ckpt_path) + return original_state_dict + + +def load_hub_or_local_checkpoint(repo_id: str | None = None, filename: str | None = None) -> dict[str, Any]: + if repo_id is None and filename is None: + raise ValueError("Please supply at least one of `repo_id` or `filename`") + + if repo_id is not None: + if filename is None: + raise ValueError("If repo_id is specified, filename must also be specified.") + ckpt_path = hf_hub_download(repo_id=repo_id, filename=filename) + else: + ckpt_path = filename + + _, ext = os.path.splitext(ckpt_path) + if ext in [".safetensors", ".sft"]: + state_dict = safetensors.torch.load_file(ckpt_path) + else: + state_dict = torch.load(ckpt_path, map_location="cpu") + + return state_dict + + +def get_model_state_dict_from_combined_ckpt(combined_ckpt: dict[str, Any], prefix: str) -> dict[str, Any]: + # Ensure that the key prefix ends with a dot (.) + if not prefix.endswith("."): + prefix = prefix + "." + + model_state_dict = {} + for param_name, param in combined_ckpt.items(): + if param_name.startswith(prefix): + model_state_dict[param_name.removeprefix(prefix)] = param + + if prefix == "model.diffusion_model.": + # Some checkpoints store the text connector projection outside the diffusion model prefix. + connector_prefixes = ["text_embedding_projection"] + for param_name, param in combined_ckpt.items(): + for prefix in connector_prefixes: + if param_name.startswith(prefix): + # Check to make sure we're not overwriting an existing key + if param_name not in model_state_dict: + model_state_dict[param_name] = combined_ckpt[param_name] + + return model_state_dict + + +def get_args(): + parser = argparse.ArgumentParser() + + def none_or_str(value: str): + if isinstance(value, str) and value.lower() == "none": + return None + return value + + parser.add_argument( + "--original_state_dict_repo_id", + default="Lightricks/LTX-2", + type=none_or_str, + help="HF Hub repo id with LTX 2.0 checkpoint", + ) + parser.add_argument( + "--checkpoint_path", + default=None, + type=str, + help="Local checkpoint path for LTX 2.0. Will be used if `original_state_dict_repo_id` is not specified.", + ) + parser.add_argument( + "--version", + type=str, + default="2.0", + choices=["test", "2.0", "2.3", "2.5"], + help="Version of the LTX 2.0 model", + ) + + parser.add_argument( + "--combined_filename", + default="ltx-2-19b-dev.safetensors", + type=none_or_str, + help="Filename for combined checkpoint with all LTX 2.0 models (VAE, DiT, etc.)", + ) + parser.add_argument("--vae_prefix", default="vae.", type=str) + parser.add_argument("--audio_vae_prefix", default="audio_vae.", type=str) + parser.add_argument("--dit_prefix", default="model.diffusion_model.", type=str) + parser.add_argument("--vocoder_prefix", default="vocoder.", type=str) + parser.add_argument("--duration_head_prefix", default="duration_head.", type=str) + + parser.add_argument("--vae_filename", default=None, type=str, help="VAE filename; overrides combined ckpt if set") + parser.add_argument( + "--audio_vae_filename", default=None, type=str, help="Audio VAE filename; overrides combined ckpt if set" + ) + parser.add_argument("--dit_filename", default=None, type=str, help="DiT filename; overrides combined ckpt if set") + parser.add_argument( + "--vocoder_filename", default=None, type=str, help="Vocoder filename; overrides combined ckpt if set" + ) + parser.add_argument( + "--text_encoder_model_id", + default="google/gemma-3-12b-it-qat-q4_0-unquantized", + type=none_or_str, + help=( + "HF Hub id for the text encoder model. Default is Gemma 3, used by LTX 2.0/2.3. LTX-2.5 requires a " + "Gemma 4 (`gemma4_unified`) checkpoint here instead -- passing the Gemma 3 default with `--version 2.5` " + "raises an error." + ), + ) + parser.add_argument( + "--tokenizer_id", + default="google/gemma-3-12b-it-qat-q4_0-unquantized", + type=none_or_str, + help="HF Hub id for the text tokenizer. Should match --text_encoder_model_id's family (Gemma 3 vs Gemma 4).", + ) + parser.add_argument( + "--prompt_enhancer_model_id", + default=None, + type=none_or_str, + help=( + "HF Hub id for the prompt-enhancer model (used with --add_processor). For LTX-2.0/2.3, defaults to " + "--text_encoder_model_id (the same Gemma 3 checkpoint serves both roles). LTX-2.5's fine-tuned text " + "encoder is not trained for enhancement, so this must be set explicitly for --version 2.5 -- e.g. to " + "google/gemma-4-E2B-it or google/gemma-4-E4B-it." + ), + ) + parser.add_argument( + "--temporal_latent_upsampler_filename", + default="ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors", + type=none_or_str, + help="Temporal x2 latent upsampler filename (LTX-2.5, used by the DFR pipeline's temporal refine rounds)", + ) + parser.add_argument( + "--latent_upsampler_filename", + default="ltx-2-spatial-upscaler-x2-1.0.safetensors", + type=none_or_str, + help="Latent upsampler filename", + ) + + parser.add_argument( + "--timestep_conditioning", action="store_true", help="Whether to add timestep condition to the video VAE model" + ) + parser.add_argument("--vae", action="store_true", help="Whether to convert the video VAE model") + parser.add_argument( + "--diffusion_vae", + action="store_true", + help=( + "Whether to convert the LTX-2.5 diffusion decoder, saved to a `diffusion_decoder` subfolder — the " + "component name `LTX2VideoDiffusionDecodePipeline` and the modular blocks resolve it by — so " + "`from_pretrained` keeps returning the conv decoder in `vae` by default" + ), + ) + parser.add_argument("--audio_vae", action="store_true", help="Whether to convert the audio VAE model") + parser.add_argument("--dit", action="store_true", help="Whether to convert the DiT model") + parser.add_argument("--connectors", action="store_true", help="Whether to convert the connector model") + parser.add_argument( + "--duration_head", + action="store_true", + help="Whether to convert the duration head (present in LTX-2.5 and later checkpoints only)", + ) + parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model") + parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder") + parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler") + parser.add_argument( + "--temporal_latent_upsampler", + action="store_true", + help="Whether to convert the temporal x2 latent upsampler (LTX-2.5)", + ) + parser.add_argument( + "--full_pipeline", + action="store_true", + help="Whether to save the pipeline. This will attempt to convert all models (e.g. vae, dit, etc.)", + ) + parser.add_argument( + "--upsample_pipeline", + action="store_true", + help="Whether to save a latent upsampling pipeline", + ) + parser.add_argument( + "--add_processor", + action="store_true", + help="Whether to add a text-encoder processor to the pipeline for prompt enhancement.", + ) + + parser.add_argument("--vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--audio_vae_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--dit_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--vocoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + parser.add_argument("--text_encoder_dtype", type=str, default="bf16", choices=["fp32", "fp16", "bf16"]) + + parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved") + parser.add_argument( + "--upsample_output_path", + type=str, + default=None, + help="Path where converted upsampling pipeline should be saved", + ) + + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +def main(args): + vae_dtype = DTYPE_MAPPING[args.vae_dtype] + audio_vae_dtype = DTYPE_MAPPING[args.audio_vae_dtype] + dit_dtype = DTYPE_MAPPING[args.dit_dtype] + vocoder_dtype = DTYPE_MAPPING[args.vocoder_dtype] + text_encoder_dtype = DTYPE_MAPPING[args.text_encoder_dtype] + + combined_ckpt = None + load_combined_models = any( + [ + args.vae, + args.diffusion_vae, + args.audio_vae, + args.dit, + args.vocoder, + args.connectors, + args.full_pipeline, + args.upsample_pipeline, + ] + ) + if args.combined_filename is not None and load_combined_models: + combined_ckpt = load_original_checkpoint(args, filename=args.combined_filename) + + # LTX-2.5 only works with a Gemma 4 (`gemma4_unified`) text encoder; --text_encoder_model_id defaults to + # Gemma 3 (for 2.0/2.3), so silently proceeding would pair a 2.5 checkpoint with the wrong text encoder. + gemma_text_config = None + if args.version == "2.5" and (args.text_encoder or args.connectors or args.full_pipeline): + gemma_config = AutoConfig.from_pretrained(args.text_encoder_model_id) + if gemma_config.model_type != "gemma4_unified": + raise ValueError( + f"LTX-2.5 requires a Gemma 4 (`gemma4_unified`) text encoder, but --text_encoder_model_id=" + f"{args.text_encoder_model_id!r} has model_type={gemma_config.model_type!r}. Pass " + "--text_encoder_model_id pointing at a Gemma 4 checkpoint (the default is Gemma 3, for 2.0/2.3)." + ) + gemma_text_config = gemma_config.text_config + + # LTX-2.5's fine-tuned text encoder is never a valid prompt enhancer (unlike LTX-2.0/2.3, where the same Gemma 3 + # checkpoint serves both roles) -- require an explicit, separate --prompt_enhancer_model_id instead of silently + # falling back to --text_encoder_model_id. + if ( + args.version == "2.5" + and args.add_processor + and (args.text_encoder or args.full_pipeline) + and args.prompt_enhancer_model_id is None + ): + raise ValueError( + "LTX-2.5's text encoder is not trained for prompt enhancement, so --prompt_enhancer_model_id must be " + "set explicitly when --add_processor is used with --version 2.5 -- e.g. to google/gemma-4-E2B-it or " + "google/gemma-4-E4B-it." + ) + + if args.vae or args.full_pipeline or args.upsample_pipeline: + if args.vae_filename is not None: + original_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename) + elif combined_ckpt is not None: + original_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix) + vae = convert_ltx2_video_vae( + original_vae_ckpt, version=args.version, timestep_conditioning=args.timestep_conditioning + ) + if not args.full_pipeline and not args.upsample_pipeline: + vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "vae")) + + if args.diffusion_vae: + if args.vae_filename is not None: + original_diffusion_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename) + elif combined_ckpt is not None: + original_diffusion_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix) + diffusion_vae = convert_ltx2_diffusion_video_vae(original_diffusion_vae_ckpt, version=args.version) + # "diffusion_decoder", not "vae_diffusion": pipeline-level `from_pretrained` resolves each component + # from the subfolder named after it, so this folder name must match the `diffusion_decoder` component + # of `LTX2VideoDiffusionDecodePipeline` (and the modular `ComponentSpec`) for those loads to work. + diffusion_vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "diffusion_decoder")) + + if args.audio_vae or args.full_pipeline: + if args.audio_vae_filename is not None: + original_audio_vae_ckpt = load_hub_or_local_checkpoint(filename=args.audio_vae_filename) + elif combined_ckpt is not None: + original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.audio_vae_prefix) + audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=args.version) + if not args.full_pipeline: + audio_vae.to(audio_vae_dtype).save_pretrained(os.path.join(args.output_path, "audio_vae")) + + if args.dit or args.full_pipeline: + if args.dit_filename is not None: + original_dit_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) + elif combined_ckpt is not None: + original_dit_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) + transformer = convert_ltx2_transformer(original_dit_ckpt, version=args.version) + if not args.full_pipeline: + transformer.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "transformer")) + + if args.connectors or args.full_pipeline: + if args.dit_filename is not None: + original_connectors_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename) + elif combined_ckpt is not None: + original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix) + connectors = convert_ltx2_connectors( + original_connectors_ckpt, version=args.version, gemma_text_config=gemma_text_config + ) + if not args.full_pipeline: + connectors.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "connectors")) + + duration_head = None + if args.duration_head or args.full_pipeline: + if combined_ckpt is not None: + original_duration_head_ckpt = get_model_state_dict_from_combined_ckpt( + combined_ckpt, args.duration_head_prefix + ) + duration_head = convert_ltx2_duration_head(original_duration_head_ckpt) + if duration_head is not None and not args.full_pipeline: + duration_head.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "duration_head")) + + if args.vocoder or args.full_pipeline: + if args.vocoder_filename is not None: + original_vocoder_ckpt = load_hub_or_local_checkpoint(filename=args.vocoder_filename) + elif combined_ckpt is not None: + original_vocoder_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vocoder_prefix) + vocoder = convert_ltx2_vocoder(original_vocoder_ckpt, version=args.version) + if not args.full_pipeline: + vocoder.to(vocoder_dtype).save_pretrained(os.path.join(args.output_path, "vocoder")) + + if args.text_encoder or args.full_pipeline: + text_encoder = AutoModelForImageTextToText.from_pretrained(args.text_encoder_model_id) + if not args.full_pipeline: + text_encoder.to(text_encoder_dtype).save_pretrained(os.path.join(args.output_path, "text_encoder")) + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + if not args.full_pipeline: + tokenizer.save_pretrained(os.path.join(args.output_path, "tokenizer")) + + processor = None + prompt_enhancer = None + if args.add_processor: + enhancer_model_id = args.prompt_enhancer_model_id or args.text_encoder_model_id + processor = AutoProcessor.from_pretrained(enhancer_model_id) + if not args.full_pipeline: + processor.save_pretrained(os.path.join(args.output_path, "processor")) + + if args.prompt_enhancer_model_id is not None: + # Separate, dedicated enhancer model (required for LTX-2.5); for LTX-2.0/2.3, the same + # `text_encoder` checkpoint already serves as its own enhancer, so nothing extra is saved. + prompt_enhancer = AutoModelForImageTextToText.from_pretrained(enhancer_model_id) + if not args.full_pipeline: + prompt_enhancer.to(text_encoder_dtype).save_pretrained( + os.path.join(args.output_path, "prompt_enhancer") + ) + + if args.latent_upsampler or args.upsample_pipeline: + original_latent_upsampler_ckpt = load_hub_or_local_checkpoint( + repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename + ) + latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version) + latent_upsampler = convert_ltx2_latent_upsampler( + original_latent_upsampler_ckpt, + latent_upsampler_config, + dtype=vae_dtype, + ) + if not args.full_pipeline and not args.upsample_pipeline: + latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler")) + + if args.temporal_latent_upsampler: + original_temporal_upsampler_ckpt = load_hub_or_local_checkpoint( + repo_id=args.original_state_dict_repo_id, filename=args.temporal_latent_upsampler_filename + ) + temporal_latent_upsampler = convert_ltx2_latent_upsampler( + original_temporal_upsampler_ckpt, + get_ltx2_temporal_latent_upsampler_config(args.version), + dtype=vae_dtype, + ) + temporal_latent_upsampler.save_pretrained(os.path.join(args.output_path, "temporal_latent_upsampler")) + + if args.full_pipeline: + is_distilled_ckpt = "distilled" in args.combined_filename + if is_distilled_ckpt: + # Disable dynamic shifting and terminal shift so that distilled sigmas are used as-is + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=False, + base_shift=0.95, + max_shift=2.05, + base_image_seq_len=1024, + max_image_seq_len=4096, + shift_terminal=None, + ) + else: + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, + base_shift=0.95, + max_shift=2.05, + base_image_seq_len=1024, + max_image_seq_len=4096, + shift_terminal=0.1, + ) + + pipe = LTX2Pipeline( + scheduler=scheduler, + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + processor=processor, + prompt_enhancer=prompt_enhancer, + ) + + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + if args.upsample_pipeline: + pipe = LTX2LatentUpsamplePipeline(vae=vae, latent_upsampler=latent_upsampler) + + # As two diffusers pipelines cannot be in the same directory, save the upsampling pipeline to its own directory + if args.upsample_output_path: + upsample_output_path = args.upsample_output_path + else: + upsample_output_path = args.output_path + pipe.save_pretrained(upsample_output_path, safe_serialization=True, max_shard_size="5GB") + + +if __name__ == "__main__": + args = get_args() + main(args) diff --git a/scripts/recipes/lumina.py b/scripts/recipes/lumina.py new file mode 100644 index 000000000000..7ab65522e90a --- /dev/null +++ b/scripts/recipes/lumina.py @@ -0,0 +1,72 @@ +import argparse +import os + +import torch +from safetensors.torch import load_file +from transformers import AutoModel, AutoTokenizer + +from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, LuminaNextDiT2DModel, LuminaPipeline +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +def main(args): + # checkpoint from https://huggingface.co/Alpha-VLLM/Lumina-Next-SFT or https://huggingface.co/Alpha-VLLM/Lumina-Next-T2I + all_sd = load_file(args.origin_ckpt_path, device="cpu") + transformer = LuminaNextDiT2DModel( + sample_size=128, + patch_size=2, + in_channels=4, + hidden_size=2304, + num_layers=24, + num_attention_heads=32, + num_kv_heads=8, + multiple_of=256, + ffn_dim_multiplier=None, + norm_eps=1e-5, + learn_sigma=True, + qk_norm=True, + cross_attention_dim=2048, + scaling_factor=1.0, + ) + transformer.load_state_dict( + convert_component_checkpoint(all_sd, dict(transformer.config), "LuminaNextDiT2DModel"), strict=True + ) + + num_model_params = sum(p.numel() for p in transformer.parameters()) + print(f"Total number of transformer parameters: {num_model_params}") + + if args.only_transformer: + transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) + else: + scheduler = FlowMatchEulerDiscreteScheduler() + + vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", torch_dtype=torch.float32) + + tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b") + text_encoder = AutoModel.from_pretrained("google/gemma-2b") + + pipeline = LuminaPipeline( + tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler + ) + pipeline.save_pretrained(args.dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--origin_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--image_size", + default=1024, + type=int, + choices=[256, 512, 1024], + required=False, + help="Image size of pretrained model, either 512 or 1024.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") + parser.add_argument("--only_transformer", default=True, type=bool, required=True) + + args = parser.parse_args() + main(args) diff --git a/scripts/recipes/minimax_h3.py b/scripts/recipes/minimax_h3.py new file mode 100644 index 000000000000..dc679f650fec --- /dev/null +++ b/scripts/recipes/minimax_h3.py @@ -0,0 +1,363 @@ +"""Convert an original MiniMax-H3 checkpoint into the diffusers layout. + +The transformer checkpoint is streamed shard by shard, so peak memory stays close to a single shard (~4.9 GiB) and +never approaches the 62 GiB of the full 33B DiT. + +Tensor layouts are converted by diffusers.loaders.conversion; this recipe supplies configuration and pipeline +metadata. The shared `minimax_h3_shards` layout handles per-head interleaved QKV and the FFN gate/value ordering. +Source preparation excludes the non-persistent `rope.inv_freq` buffer, which the model regenerates from config. + +The FL2VA and Ref2VA variants differ only in the transformer weights, so the variant is selected by pointing +`--checkpoint_path` at the corresponding folder. Both land in one repository, which carries a single +`modular_model_index.json`: MiniMax-H3 is integrated as Modular Diffusers blocks only, so no `model_index.json` is +written. + +Usage: + +```bash +# Validate the key mapping without any weights present. +python scripts/recipes/minimax_h3.py \ + --checkpoint_path /path/to/MiniMax-H3/FL2VA --output_path /tmp/h3-diffusers --dry_run + +# Convert, and point the component loading specs at the Hub id the result is published under. +python scripts/recipes/minimax_h3.py \ + --checkpoint_path /path/to/MiniMax-H3/FL2VA --output_path /tmp/h3-diffusers \ + --modular_repo_id MiniMaxAI/MiniMax-H3 +``` +""" + +import argparse +import glob +import json +import os +import struct +from typing import Any + +import torch + +from diffusers import MiniMaxH3Transformer3DModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.minimax_h3 import ( + MINIMAX_H3_TEST_TRANSFORMER_CONFIG, + MINIMAX_H3_TEST_VIDEO_VAE_CONFIG, + MINIMAX_H3_TRANSFORMER_CONFIG, + MINIMAX_H3_VIDEO_VAE_CONFIG, + get_audio_vae_config, +) +from diffusers.loaders.conversion.io import Checkpoint, convert_checkpoint +from diffusers.loaders.conversion.source import MergedCheckpoint + + +# `MiniMaxH3Transformer3DModel` argument names. The original config uses the sglang-native names listed in the +# comments; everything else in the original config (`adaln_out_features`, `final_adaln_out_features`) is derived. + +# A tiny configuration with the checkpoint-tied dimensions left intact, for building fixtures. + +# MiniMax-H3 ships a mixed-precision checkpoint. These *original* keys are float32; everything else is bfloat16 — +# including the AdaLN projections. +MINIMAX_H3_FP32_SOURCE_PREFIXES = ( + "video_patch_proj.", + "audio_patch_proj.", + "time_embedder.", + "final_layer.video_out.", + "final_layer.audio_out.", +) + +# `rope.inv_freq` is `1 / rope_theta ** (arange(0, 2 * rope_freq_dim, 2) / (2 * rope_freq_dim))`, which +# `MiniMaxH3RotaryPosEmbed` recomputes into a non-persistent buffer. The recomputed tensor is bitwise equal to the +# shipped one in both released variants, so the key is not carried into the diffusers checkpoint. +MINIMAX_H3_TRANSFORMER_DROPPED_KEYS = ("rope.inv_freq",) + + +def get_transformer_key_plan(config): + with torch.device("meta"): + model = MiniMaxH3Transformer3DModel(**config) + shapes = {key: tuple(value.shape) for key, value in model.state_dict().items()} + conversion = get_conversion("MiniMaxH3Transformer3DModel", config) + plan = {old: [(new, shapes[new])] for old, new in conversion.mapping.items()} + for rule in conversion.rules: + for old in rule.original: + plan[old] = [(new, shapes[new]) for new in rule.diffusers] + return plan + + +# +# --------------------------------------------------------------------------------------------------------------- +# Video VAE +# --------------------------------------------------------------------------------------------------------------- +# + +# `AutoencoderKLMiniMaxH3` argument names. Field-for-field equal to `video_vae/source/config.json`, with the original +# names in the comments. The keys that only ever take one value in the release (`use_3d_conv`, `use_vit_decoder`, +# `causal_encoder`, `causal_decoder`, `use_t_isolated_gn`, `space_up` / `time_up`, `zq_ch_*`, `num_res_blocks_decoder`, +# `shift_factor` / `scaling_factor`) are baked into the port instead of being config knobs. + +# A tiny configuration with the checkpoint-tied dimensions (`latent_channels`, the temporal geometry and the rotary +# ratio) left intact, for building fixtures and for the CPU parity check. + +# `decoder.mask_token` is an all-zero buffer belonging to the masked-autoencoding training objective; the released +# decoder never reads it, so the port does not carry the module and the conversion drops the key. + + +def convert_video_vae(checkpoint_path, output_path, config, diffusers_version, max_shard_size): + source_dir = os.path.join(checkpoint_path, "video_vae") + with open(os.path.join(source_dir, "config.json")) as handle: + wrapper = json.load(handle) + path = os.path.join(source_dir, wrapper["source_path"], wrapper["source_safetensors_path"]) + config = {**config, "latents_mean": wrapper["latents_mean"], "latents_std": wrapper["latents_std"]} + return convert_checkpoint( + path, output_path, config=config, model_class="AutoencoderKLMiniMaxH3", max_shard_size=max_shard_size + ) + + +# Not present in `audio_vae/metadata.json`: the reference implementation hardcodes these in its DAC audio VAE and its +# attention projection, keyed off the sample rate. + + +def convert_audio_vae(checkpoint_path, output_path, diffusers_version): + return convert_checkpoint( + os.path.join(checkpoint_path, "audio_vae", "model.safetensors"), + output_path, + config=get_audio_vae_config(checkpoint_path), + model_class="AutoencoderKLMiniMaxH3Audio", + ) + + +def write_scheduler_configs(checkpoint_path: str, output_path: str, diffusers_version: str) -> None: + """Emit the two `MiniMaxH3Scheduler` configs, one per modality. + + The source `model_index.json` leaves `scheduler` null and instead carries the schedule constants in its + `_minimax_h3.sigma_shift_scales` block. The sigma shift is the only per-modality difference, so it becomes two + scheduler folders holding the same class at different `shift` values. + """ + with open(os.path.join(checkpoint_path, "model_index.json")) as f: + shift_scales = json.load(f)["_minimax_h3"]["sigma_shift_scales"] + + for folder, modality in (("scheduler", "video"), ("audio_scheduler", "audio")): + folder_path = os.path.join(output_path, folder) + os.makedirs(folder_path, exist_ok=True) + with open(os.path.join(folder_path, "scheduler_config.json"), "w") as f: + json.dump( + { + "_class_name": "MiniMaxH3Scheduler", + "_diffusers_version": diffusers_version, + "shift": float(shift_scales[modality]), + }, + f, + indent=2, + ) + print(f"scheduler: shift={shift_scales['video']} (video), audio_scheduler: shift={shift_scales['audio']} (audio).") + + +def read_safetensors_header(path: str) -> dict[str, Any]: + """Read the metadata header of a safetensors file without touching the tensor payload.""" + with open(path, "rb") as f: + header_size = struct.unpack(" None: + plan = get_transformer_key_plan(config) + + transformer_dir = os.path.join(checkpoint_path, "transformer") + shards = sorted(glob.glob(os.path.join(transformer_dir, "*.safetensors"))) + header: dict[str, Any] = {} + for shard in shards: + header.update(read_safetensors_header(shard)) + if shards: + print(f"Read headers of {len(shards)} shard(s) in {transformer_dir}: {len(header)} keys present.\n") + else: + print(f"No shards found under {transformer_dir}; validating the plan against the config only.\n") + + print(f"{'original key':<48} {'->':^4} {'diffusers key':<52} {'shape':<24} dtype") + print("-" * 150) + num_target_keys = 0 + shape_mismatches: list[str] = [] + for source_key, targets in plan.items(): + present = source_key in header + if not targets: + print(f"{source_key:<48} {'-x':^4} {'(dropped, recomputed by the port)':<52}") + continue + for index, (target_key, shape) in enumerate(targets): + num_target_keys += 1 + expected_dtype = "F32" if source_key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES) else "BF16" + if present: + actual_dtype = header[source_key]["dtype"] + actual_shape = header[source_key]["shape"] + if index == 0 and len(targets) == 1 and actual_shape != shape: + shape_mismatches.append(f"{source_key}: header {actual_shape} != planned {shape}") + if actual_dtype != expected_dtype: + shape_mismatches.append(f"{source_key}: header dtype {actual_dtype} != expected {expected_dtype}") + marker = "->" + else: + marker = "->?" + left = source_key if index == 0 else "" + print(f"{left:<48} {marker:^4} {target_key:<52} {str(shape):<24} {expected_dtype}") + + missing = [key for key in plan if key not in header] + unexpected = [key for key in header if key not in plan] + + print("\n" + "=" * 150) + print(f"planned original keys : {len(plan)}") + print(f"planned diffusers keys: {num_target_keys}") + print( + f"dropped original keys : {len(MINIMAX_H3_TRANSFORMER_DROPPED_KEYS)} {list(MINIMAX_H3_TRANSFORMER_DROPPED_KEYS)}" + ) + print(f"fp32 diffusers keys : {sum(1 for key in plan if key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES))}") + if shards: + print(f"keys present in shards: {len(header)}") + print(f"planned but absent : {len(missing)}" + (" (shards still downloading?)" if missing else "")) + print(f"present but unplanned : {len(unexpected)}") + if unexpected: + print(f" {unexpected}") + print(f"header disagreements : {len(shape_mismatches)}") + for line in shape_mismatches: + print(f" {line}") + total_bytes = sum( + (4 if source_key.startswith(MINIMAX_H3_FP32_SOURCE_PREFIXES) else 2) * torch.Size(shape).numel() + for source_key, targets in plan.items() + for _, shape in targets + ) + print(f"total output bytes : {total_bytes} ({total_bytes / 1024**3:.2f} GiB)") + + +def convert_transformer(checkpoint_path, output_path, config, max_shard_size): + shards = sorted(glob.glob(os.path.join(checkpoint_path, "transformer", "*.safetensors"))) + source = MergedCheckpoint(("", Checkpoint(path)) for path in shards) + return convert_checkpoint( + source, + output_path, + config={**config, "original_format": "minimax_h3_shards"}, + model_class="MiniMaxH3Transformer3DModel", + max_shard_size=max_shard_size, + ) + + +# The components a MiniMax-H3 repository holds, and the class each one loads as. `video_processor` is absent: the +# blocks create it from config rather than loading it. +MINIMAX_H3_COMPONENTS = { + # The source names a checkpoint-local wrapper class (`MiniMaxH3Qwen3VLHFEncoder`); the conditioner is the + # released Qwen3-VL, read at its 50th decoder layer with its language-model head unused. + "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"], + "tokenizer": ["transformers", "Qwen2TokenizerFast"], + "processor": ["transformers", "Qwen3VLProcessor"], + # Renamed to the diffusers audio/video VAE convention (see `LTX2Pipeline`). + "vae": ["diffusers", "AutoencoderKLMiniMaxH3"], + "audio_vae": ["diffusers", "AutoencoderKLMiniMaxH3Audio"], + "transformer": ["diffusers", "MiniMaxH3Transformer3DModel"], + # One repository holds both checkpoint partitions: `transformer/` serves `MiniMaxH3Blocks` (`t2va` / `fl2va`) and + # `transformer_ref/` serves `MiniMaxH3Ref2VABlocks`, while every other component is shared and converted once. + "transformer_ref": ["diffusers", "MiniMaxH3Transformer3DModel"], + # The source leaves `scheduler` null. MiniMax-H3 samples with Euler at eta=0 over shifted flow-matching sigmas, at + # a different shift per modality, so it needs two scheduler entries (see `write_scheduler_configs`). + "scheduler": ["diffusers", "MiniMaxH3Scheduler"], + "audio_scheduler": ["diffusers", "MiniMaxH3Scheduler"], +} + + +def write_model_index(output_path: str, repo_id: str, diffusers_version: str) -> None: + """Emit `modular_model_index.json`, the only index a MiniMax-H3 repository carries. + + MiniMax-H3 is integrated as Modular Diffusers blocks only, so there is no `model_index.json`: a modular repository + declares one entry per component with its full loading spec rather than just its class, and a blockset then fetches + exactly the subfolders it declares. That is what lets one repository hold both transformer partitions, and the + original checkpoint folders next to the converted ones, without either half pulling the rest down. + + `_class_name` and `_blocks_class_name` name the `t2va` / `fl2va` half, which is what + `ModularPipeline.from_pretrained` resolves to. The `ref2va` half reads the very same file through + `MiniMaxH3Ref2VABlocks().init_pipeline(repo_id)`. + + The component map is the static one above, so this needs no source checkpoint: an index can be regenerated for a + repository that is already published. + """ + modular_index = { + "_class_name": "MiniMaxH3ModularPipeline", + "_diffusers_version": diffusers_version, + "_blocks_class_name": "MiniMaxH3Blocks", + } + for name, (library, class_name) in MINIMAX_H3_COMPONENTS.items(): + modular_index[name] = [ + library, + class_name, + { + "type_hint": [library, class_name], + "pretrained_model_name_or_path": repo_id, + "subfolder": name, + "variant": None, + "revision": None, + }, + ] + os.makedirs(output_path, exist_ok=True) + with open(os.path.join(output_path, "modular_model_index.json"), "w") as f: + json.dump(modular_index, f, indent=2) + print(f"modular_model_index.json: {len(MINIMAX_H3_COMPONENTS)} components load from {repo_id}.") + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint_path", + type=str, + required=True, + help="Local path to an original MiniMax-H3 variant folder (the one holding `transformer/`, `audio_vae/`, ...).", + ) + parser.add_argument("--output_path", type=str, required=True, help="Where the diffusers checkpoint is written.") + parser.add_argument( + "--modular_repo_id", + type=str, + default=None, + help=( + "Repository the component entries of `modular_model_index.json` point at. Defaults to `--output_path`, so " + "pass the Hub id the checkpoint is published under. Every entry carries its own loading spec, so a " + "blockset fetches exactly the subfolders it declares out of that repository." + ), + ) + parser.add_argument( + "--version", + type=str, + default="h3", + choices=["h3", "test"], + help="`test` emits the tiny config used for fixtures.", + ) + parser.add_argument( + "--max_shard_size", + type=int, + default=5 * 1024**3, + help="Maximum size of an output safetensors shard, in bytes.", + ) + parser.add_argument( + "--dry_run", + action="store_true", + help="Print the full planned key mapping (and cross-check any shard headers already present) without writing.", + ) + return parser.parse_args() + + +def main(args): + from diffusers import __version__ as diffusers_version + + config = MINIMAX_H3_TEST_TRANSFORMER_CONFIG if args.version == "test" else MINIMAX_H3_TRANSFORMER_CONFIG + + if args.dry_run: + dry_run(args.checkpoint_path, config) + return + + transformer_path = os.path.join(args.output_path, "transformer") + convert_transformer(args.checkpoint_path, transformer_path, config, args.max_shard_size) + video_vae_config = MINIMAX_H3_TEST_VIDEO_VAE_CONFIG if args.version == "test" else MINIMAX_H3_VIDEO_VAE_CONFIG + convert_video_vae( + args.checkpoint_path, + os.path.join(args.output_path, "vae"), + video_vae_config, + diffusers_version, + args.max_shard_size, + ) + convert_audio_vae(args.checkpoint_path, os.path.join(args.output_path, "audio_vae"), diffusers_version) + write_scheduler_configs(args.checkpoint_path, args.output_path, diffusers_version) + write_model_index(args.output_path, args.modular_repo_id or args.output_path, diffusers_version) + + +if __name__ == "__main__": + main(get_args()) diff --git a/scripts/recipes/minimax_music3.py b/scripts/recipes/minimax_music3.py new file mode 100644 index 000000000000..85fc877e4337 --- /dev/null +++ b/scripts/recipes/minimax_music3.py @@ -0,0 +1,182 @@ +# Conversion script for MiniMax Music 3 (https://huggingface.co/MiniMaxAI/MiniMax-Music3). +# +# Original checkpoint layout: +# flowmatching_vae.pth flow-matching DiT + condition projection +# dav.pth Flow-VAE (DAC-style) decoder +# qwen_7B/qwen_7B/ Qwen3 backbone + audio embedding + RVQ depth decoder (sharded safetensors) +# qwen_7B/qwen3-8B-tokenizer-music/ music tokenizer +# +# Usage: +# python scripts/recipes/minimax_music3.py \ +# --checkpoint_dir MiniMaxAI/MiniMax-Music3 --output_path ./minimax-music3-diffusers + +import argparse +import json +import os + +import torch +from safetensors.torch import load_file + +from diffusers import ( + FlowMatchEulerDiscreteScheduler, + MiniMaxMusic3ConditionEncoder, + MiniMaxMusic3RVQDepthDecoder, + MiniMaxMusic3Transformer1DModel, + MiniMaxMusic3Vocoder, +) + + +def load_dit_state_dict(checkpoint_dir: str) -> dict: + return torch.load(os.path.join(checkpoint_dir, "flowmatching_vae.pth"), map_location="cpu", weights_only=True) + + +def load_dav_state_dict(checkpoint_dir: str) -> dict: + return torch.load(os.path.join(checkpoint_dir, "dav.pth"), map_location="cpu", weights_only=True) + + +def load_qwen_state_dict(checkpoint_dir: str) -> dict: + qwen_dir = os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B") + with open(os.path.join(qwen_dir, "model.safetensors.index.json")) as f: + index = json.load(f) + state_dict = {} + for filename in sorted(set(index["weight_map"].values())): + state_dict.update(load_file(os.path.join(qwen_dir, filename), device="cpu")) + return state_dict + + +def convert_transformer(dit_state_dict): + prefix = "diffusion_transformer." + state = {key: value for key, value in dit_state_dict.items() if key.startswith(prefix)} + num_layers = 1 + max(int(key.split(".")[3]) for key in state if key.startswith(prefix + "transformer.layers.")) + return MiniMaxMusic3Transformer1DModel.from_single_file(state, config={"num_layers": num_layers}) + + +def convert_condition_encoder(dit_state_dict): + state = { + key: value for key, value in dit_state_dict.items() if key.startswith(("cond_layer_", "latent_conditioners.")) + } + return MiniMaxMusic3ConditionEncoder.from_single_file(state, config={}) + + +def convert_vocoder(dav_state_dict): + state = {key: value for key, value in dav_state_dict.items() if key.startswith(("dec_in_proj.", "decoder."))} + return MiniMaxMusic3Vocoder.from_single_file(state, config={}) + + +def convert_rvq_depth_decoder(qwen_state_dict, model_config): + state = { + key: value + for key, value in qwen_state_dict.items() + if key.startswith(("model.audio_decoder.", "model.audio_extra_embedding.")) + } + config = { + "hidden_size": int(model_config["hidden_size"]), + "num_layers": int(model_config["decoder_num_layers"]), + "num_attention_heads": int(model_config["decoder_num_heads"]), + "intermediate_size": int(model_config["decoder_intermediate_size"]), + "audio_vocab_size": int(model_config["audio_vocab_size"]), + "num_codebooks": int(model_config["audio_num_codebooks"]), + } + return MiniMaxMusic3RVQDepthDecoder.from_single_file(state, config=config) + + +def convert_language_model(qwen_state_dict: dict, model_config: dict): + from transformers import Qwen3Config, Qwen3ForCausalLM + + config = Qwen3Config( + vocab_size=model_config["vocab_size"], + hidden_size=model_config["hidden_size"], + intermediate_size=model_config["intermediate_size"], + num_hidden_layers=model_config["num_hidden_layers"], + num_attention_heads=model_config["num_attention_heads"], + num_key_value_heads=model_config["num_key_value_heads"], + head_dim=model_config["head_dim"], + max_position_embeddings=model_config.get("max_position_embeddings", 10240), + rope_theta=model_config.get("rope_theta", 1000000), + rms_norm_eps=model_config.get("rms_norm_eps", 1e-6), + tie_word_embeddings=model_config.get("tie_word_embeddings", False), + ) + backbone_state_dict = { + key: value + for key, value in qwen_state_dict.items() + if not key.startswith(("model.audio_extra_embedding", "model.audio_decoder.")) + } + with torch.device("meta"): + language_model = Qwen3ForCausalLM(config) + language_model.load_state_dict(backbone_state_dict, strict=True, assign=True) + return language_model + + +def main(args): + checkpoint_dir = args.checkpoint_dir + if not os.path.isdir(checkpoint_dir): + from huggingface_hub import snapshot_download + + checkpoint_dir = snapshot_download(checkpoint_dir) + + with open(os.path.join(checkpoint_dir, "qwen_7B", "qwen_7B", "config.json")) as f: + model_config = json.load(f) + + dit_state_dict = load_dit_state_dict(checkpoint_dir) + transformer = convert_transformer(dit_state_dict).to(args.dtype) + condition_encoder = convert_condition_encoder(dit_state_dict).to(args.dtype) + del dit_state_dict + vocoder = convert_vocoder(load_dav_state_dict(checkpoint_dir)).to(args.dtype) + + qwen_state_dict = load_qwen_state_dict(checkpoint_dir) + rvq_depth_decoder = convert_rvq_depth_decoder(qwen_state_dict, model_config).to(torch.bfloat16) + language_model = convert_language_model(qwen_state_dict, model_config) + del qwen_state_dict + + from transformers import AutoTokenizer + + from diffusers import MiniMaxMusic3Blocks + + tokenizer = AutoTokenizer.from_pretrained(os.path.join(checkpoint_dir, "qwen_7B", "qwen3-8B-tokenizer-music")) + # num_train_timesteps=1 keeps `scheduler.timesteps` equal to the flow-matching time in [0, 1] that the + # transformer's Fourier embedding expects. + scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1, shift=1.0, invert_sigmas=True) + + pipeline = MiniMaxMusic3Blocks().init_pipeline() + pipeline.update_components( + language_model=language_model, + rvq_depth_decoder=rvq_depth_decoder, + condition_encoder=condition_encoder, + transformer=transformer, + vocoder=vocoder, + tokenizer=tokenizer, + scheduler=scheduler, + ) + pipeline.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") + + # save_pretrained bakes the local output path into the modular index's loading specs; point them at the + # Hub repo the components will be uploaded to instead. + index_path = os.path.join(args.output_path, "modular_model_index.json") + with open(index_path) as f: + index = json.load(f) + for entry in index.values(): + if isinstance(entry, list) and len(entry) == 3 and isinstance(entry[2], dict): + if entry[2].get("pretrained_model_name_or_path") == args.output_path: + entry[2]["pretrained_model_name_or_path"] = args.repo_id + with open(index_path, "w") as f: + json.dump(index, f, indent=2, sort_keys=True) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--checkpoint_dir", + type=str, + default="MiniMaxAI/MiniMax-Music3", + help="Local directory or Hugging Face Hub repo id of the original checkpoint.", + ) + parser.add_argument("--output_path", type=str, required=True) + parser.add_argument( + "--repo_id", + type=str, + default="MiniMaxAI/MiniMax-Music3", + help="Hub repo id the converted components will live in (written into the modular index loading specs).", + ) + parser.add_argument("--dtype", type=lambda name: getattr(torch, name), default="float32") + args = parser.parse_args() + main(args) diff --git a/scripts/recipes/mochi.py b/scripts/recipes/mochi.py new file mode 100644 index 000000000000..056046ad9d6d --- /dev/null +++ b/scripts/recipes/mochi.py @@ -0,0 +1,95 @@ +import argparse + +import torch +from safetensors.torch import load_file +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +TOKENIZER_MAX_LENGTH = 256 + +parser = argparse.ArgumentParser() +parser.add_argument("--transformer_checkpoint_path", default=None, type=str) +parser.add_argument("--vae_encoder_checkpoint_path", default=None, type=str) +parser.add_argument("--vae_decoder_checkpoint_path", default=None, type=str) +parser.add_argument("--output_path", required=True, type=str) +parser.add_argument("--push_to_hub", action="store_true", default=False, help="Whether to push to HF Hub after saving") +parser.add_argument("--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory") +parser.add_argument("--dtype", type=str, default=None) + +args = parser.parse_args() + + +# This is specific to `AdaLayerNormContinuous`: +# Diffusers implementation split the linear projection into the scale, shift while Mochi split it into shift, scale + + +def convert_mochi_transformer_checkpoint_to_diffusers(ckpt_path): + state = load_file(ckpt_path, device="cpu") + + return convert_component_checkpoint(state, {}, "MochiTransformer3DModel") + + +def convert_mochi_vae_state_dict_to_diffusers(encoder_path, decoder_path): + encoder = load_file(encoder_path, device="cpu") + decoder = load_file(decoder_path, device="cpu") + state = {f"encoder.{key}": value for key, value in encoder.items()} + state.update({f"decoder.{key}": value for key, value in decoder.items()}) + return convert_component_checkpoint(state, {"latent_channels": 12, "out_channels": 3}, "AutoencoderKLMochi") + + +def main(args): + if args.dtype is None: + dtype = None + if args.dtype == "fp16": + dtype = torch.float16 + elif args.dtype == "bf16": + dtype = torch.bfloat16 + elif args.dtype == "fp32": + dtype = torch.float32 + else: + raise ValueError(f"Unsupported dtype: {args.dtype}") + + transformer = None + vae = None + + if args.transformer_checkpoint_path is not None: + converted_transformer_state_dict = convert_mochi_transformer_checkpoint_to_diffusers( + args.transformer_checkpoint_path + ) + transformer = MochiTransformer3DModel() + transformer.load_state_dict(converted_transformer_state_dict, strict=True) + if dtype is not None: + transformer = transformer.to(dtype=dtype) + + if args.vae_encoder_checkpoint_path is not None and args.vae_decoder_checkpoint_path is not None: + vae = AutoencoderKLMochi(latent_channels=12, out_channels=3) + converted_vae_state_dict = convert_mochi_vae_state_dict_to_diffusers( + args.vae_encoder_checkpoint_path, args.vae_decoder_checkpoint_path + ) + vae.load_state_dict(converted_vae_state_dict, strict=True) + if dtype is not None: + vae = vae.to(dtype=dtype) + + text_encoder_id = "google/t5-v1_1-xxl" + tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH) + text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir) + + # Apparently, the conversion does not work anymore without this :shrug: + for param in text_encoder.parameters(): + param.data = param.data.contiguous() + + pipe = MochiPipeline( + scheduler=FlowMatchEulerDiscreteScheduler(invert_sigmas=True), + vae=vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + transformer=transformer, + ) + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", push_to_hub=args.push_to_hub) + + +if __name__ == "__main__": + main(args) diff --git a/scripts/recipes/music_spectrogram.py b/scripts/recipes/music_spectrogram.py new file mode 100644 index 000000000000..0c023aa6b58f --- /dev/null +++ b/scripts/recipes/music_spectrogram.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +import argparse +import os + +import jax as jnp +import numpy as onp +from music_spectrogram_diffusion import inference +from t5x import checkpoints + +from diffusers import DDPMScheduler, OnnxRuntimeModel, SpectrogramDiffusionPipeline +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.source import flatten_tensor_tree +from diffusers.pipelines.deprecated.spectrogram_diffusion import ( + SpectrogramContEncoder, + SpectrogramNotesEncoder, + T5FilmDecoder, +) + + +MODEL = "base_with_context" + + +def load_notes_encoder(weights, model): + state = flatten_tensor_tree(weights) + model.load_state_dict( + get_conversion("SpectrogramNotesEncoder", dict(model.config)).to_diffusers(state), strict=True + ) + return model + + +def load_continuous_encoder(weights, model): + state = flatten_tensor_tree(weights) + model.load_state_dict( + get_conversion("SpectrogramContEncoder", dict(model.config)).to_diffusers(state), strict=True + ) + return model + + +def load_decoder(weights, model): + state = flatten_tensor_tree(weights) + model.load_state_dict(get_conversion("T5FilmDecoder", dict(model.config)).to_diffusers(state), strict=True) + return model + + +def main(args): + t5_checkpoint = checkpoints.load_t5x_checkpoint(args.checkpoint_path) + t5_checkpoint = jnp.tree_util.tree_map(onp.array, t5_checkpoint) + + gin_overrides = [ + "from __gin__ import dynamic_registration", + "from music_spectrogram_diffusion.models.diffusion import diffusion_utils", + "diffusion_utils.ClassifierFreeGuidanceConfig.eval_condition_weight = 2.0", + "diffusion_utils.DiffusionConfig.classifier_free_guidance = @diffusion_utils.ClassifierFreeGuidanceConfig()", + ] + + gin_file = os.path.join(args.checkpoint_path, "..", "config.gin") + gin_config = inference.parse_training_gin_file(gin_file, gin_overrides) + synth_model = inference.InferenceModel(args.checkpoint_path, gin_config) + + scheduler = DDPMScheduler(beta_schedule="squaredcos_cap_v2", variance_type="fixed_large") + + notes_encoder = SpectrogramNotesEncoder( + max_length=synth_model.sequence_length["inputs"], + vocab_size=synth_model.model.module.config.vocab_size, + d_model=synth_model.model.module.config.emb_dim, + dropout_rate=synth_model.model.module.config.dropout_rate, + num_layers=synth_model.model.module.config.num_encoder_layers, + num_heads=synth_model.model.module.config.num_heads, + d_kv=synth_model.model.module.config.head_dim, + d_ff=synth_model.model.module.config.mlp_dim, + feed_forward_proj="gated-gelu", + ) + + continuous_encoder = SpectrogramContEncoder( + input_dims=synth_model.audio_codec.n_dims, + targets_context_length=synth_model.sequence_length["targets_context"], + d_model=synth_model.model.module.config.emb_dim, + dropout_rate=synth_model.model.module.config.dropout_rate, + num_layers=synth_model.model.module.config.num_encoder_layers, + num_heads=synth_model.model.module.config.num_heads, + d_kv=synth_model.model.module.config.head_dim, + d_ff=synth_model.model.module.config.mlp_dim, + feed_forward_proj="gated-gelu", + ) + + decoder = T5FilmDecoder( + input_dims=synth_model.audio_codec.n_dims, + targets_length=synth_model.sequence_length["targets_context"], + max_decoder_noise_time=synth_model.model.module.config.max_decoder_noise_time, + d_model=synth_model.model.module.config.emb_dim, + num_layers=synth_model.model.module.config.num_decoder_layers, + num_heads=synth_model.model.module.config.num_heads, + d_kv=synth_model.model.module.config.head_dim, + d_ff=synth_model.model.module.config.mlp_dim, + dropout_rate=synth_model.model.module.config.dropout_rate, + ) + + notes_encoder = load_notes_encoder(t5_checkpoint["target"]["token_encoder"], notes_encoder) + continuous_encoder = load_continuous_encoder(t5_checkpoint["target"]["continuous_encoder"], continuous_encoder) + decoder = load_decoder(t5_checkpoint["target"]["decoder"], decoder) + + melgan = OnnxRuntimeModel.from_pretrained("kashif/soundstream_mel_decoder") + + pipe = SpectrogramDiffusionPipeline( + notes_encoder=notes_encoder, + continuous_encoder=continuous_encoder, + decoder=decoder, + scheduler=scheduler, + melgan=melgan, + ) + if args.save: + pipe.save_pretrained(args.output_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--output_path", default=None, type=str, required=True, help="Path to the converted model.") + parser.add_argument( + "--save", default=True, type=bool, required=False, help="Whether to save the converted model or not." + ) + parser.add_argument( + "--checkpoint_path", + default=f"{MODEL}/checkpoint_500000", + type=str, + required=False, + help="Path to the original jax model checkpoint.", + ) + args = parser.parse_args() + + main(args) diff --git a/scripts/recipes/musicldm.py b/scripts/recipes/musicldm.py new file mode 100644 index 000000000000..2399a9a218a7 --- /dev/null +++ b/scripts/recipes/musicldm.py @@ -0,0 +1,365 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the MusicLDM checkpoints.""" + +import argparse + +import torch +import yaml +from transformers import ( + AutoFeatureExtractor, + AutoTokenizer, + ClapConfig, + ClapModel, + SpeechT5HifiGan, + SpeechT5HifiGanConfig, +) + +from diffusers import ( + AutoencoderKL, + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + HeunDiscreteScheduler, + LMSDiscreteScheduler, + MusicLDMPipeline, + PNDMScheduler, + UNet2DConditionModel, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.musicldm import ( + DEFAULT_CONFIG, + create_transformers_vocoder_config, + create_unet_diffusers_config, + create_vae_diffusers_config, +) + + +# Adapted from diffusers.pipelines.stable_diffusion.convert_from_ckpt.create_vae_diffusers_config + + +def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, **kwargs): + return convert_component_checkpoint(checkpoint, config, "UNet2DConditionModel", extract_ema=extract_ema) + + +def convert_ldm_vae_checkpoint(checkpoint, config): + return convert_component_checkpoint(checkpoint, config, "AutoencoderKL") + + +CLAP_EXPECTED_MISSING_KEYS = ["text_model.embeddings.token_type_ids"] + + +def convert_open_clap_checkpoint(checkpoint, config): + prefix = "cond_stage_model.model." + state = {key.removeprefix(prefix): value for key, value in checkpoint.items() if key.startswith(prefix)} + return convert_component_checkpoint(state, config.to_dict(), "ClapModel") + + +def convert_hifigan_checkpoint(checkpoint, config): + state = { + key.removeprefix("first_stage_model.vocoder."): value + for key, value in checkpoint.items() + if key.startswith("first_stage_model.vocoder.") + } + return convert_component_checkpoint(state, config.to_dict(), "SpeechT5HifiGan") + + +# Adapted from https://huggingface.co/spaces/haoheliu/MusicLDM-text-to-audio-generation/blob/84a0384742a22bd80c44e903e241f0623e874f1d/MusicLDM/utils.py#L72-L73 + + +def load_pipeline_from_original_MusicLDM_ckpt( + checkpoint_path: str, + original_config_file: str = None, + image_size: int = 1024, + prediction_type: str = None, + extract_ema: bool = False, + scheduler_type: str = "ddim", + num_in_channels: int = None, + model_channels: int = None, + num_head_channels: int = None, + device: str = None, + from_safetensors: bool = False, +) -> MusicLDMPipeline: + """ + Load an MusicLDM pipeline object from a `.ckpt`/`.safetensors` file and (ideally) a `.yaml` config file. + + Although many of the arguments can be automatically inferred, some of these rely on brittle checks against the + global step count, which will likely fail for models that have undergone further fine-tuning. Therefore, it is + recommended that you override the default values and/or supply an `original_config_file` wherever possible. + + Args: + checkpoint_path (`str`): Path to `.ckpt` file. + original_config_file (`str`): + Path to `.yaml` config file corresponding to the original architecture. If `None`, will be automatically + set to the MusicLDM-s-full-v2 config. + image_size (`int`, *optional*, defaults to 1024): + The image size that the model was trained on. + prediction_type (`str`, *optional*): + The prediction type that the model was trained on. If `None`, will be automatically + inferred by looking for a key in the config. For the default config, the prediction type is `'epsilon'`. + num_in_channels (`int`, *optional*, defaults to None): + The number of UNet input channels. If `None`, it will be automatically inferred from the config. + model_channels (`int`, *optional*, defaults to None): + The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override + to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large. + num_head_channels (`int`, *optional*, defaults to None): + The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override + to 32 for the small and medium checkpoints, and 64 for the large. + scheduler_type (`str`, *optional*, defaults to 'pndm'): + Type of scheduler to use. Should be one of `["pndm", "lms", "heun", "euler", "euler-ancestral", "dpm", + "ddim"]`. + extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for + checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to + `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for + inference. Non-EMA weights are usually better to continue fine-tuning. + device (`str`, *optional*, defaults to `None`): + The device to use. Pass `None` to determine automatically. + from_safetensors (`str`, *optional*, defaults to `False`): + If `checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch. + return: An MusicLDMPipeline object representing the passed-in `.ckpt`/`.safetensors` file. + """ + if from_safetensors: + from safetensors import safe_open + + checkpoint = {} + with safe_open(checkpoint_path, framework="pt", device="cpu") as f: + for key in f.keys(): + checkpoint[key] = f.get_tensor(key) + else: + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + checkpoint = torch.load(checkpoint_path, map_location=device) + else: + checkpoint = torch.load(checkpoint_path, map_location=device) + + if "state_dict" in checkpoint: + checkpoint = checkpoint["state_dict"] + + if original_config_file is None: + original_config = DEFAULT_CONFIG + else: + original_config = yaml.safe_load(original_config_file) + + if num_in_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels + + if model_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["model_channels"] = model_channels + + if num_head_channels is not None: + original_config["model"]["params"]["unet_config"]["params"]["num_head_channels"] = num_head_channels + + if ( + "parameterization" in original_config["model"]["params"] + and original_config["model"]["params"]["parameterization"] == "v" + ): + if prediction_type is None: + prediction_type = "v_prediction" + else: + if prediction_type is None: + prediction_type = "epsilon" + + if image_size is None: + image_size = 512 + + num_train_timesteps = original_config["model"]["params"]["timesteps"] + beta_start = original_config["model"]["params"]["linear_start"] + beta_end = original_config["model"]["params"]["linear_end"] + + scheduler = DDIMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + steps_offset=1, + clip_sample=False, + set_alpha_to_one=False, + prediction_type=prediction_type, + ) + # make sure scheduler works correctly with DDIM + scheduler.register_to_config(clip_sample=False) + + if scheduler_type == "pndm": + config = dict(scheduler.config) + config["skip_prk_steps"] = True + scheduler = PNDMScheduler.from_config(config) + elif scheduler_type == "lms": + scheduler = LMSDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "heun": + scheduler = HeunDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler": + scheduler = EulerDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "euler-ancestral": + scheduler = EulerAncestralDiscreteScheduler.from_config(scheduler.config) + elif scheduler_type == "dpm": + scheduler = DPMSolverMultistepScheduler.from_config(scheduler.config) + elif scheduler_type == "ddim": + scheduler = scheduler + else: + raise ValueError(f"Scheduler of type {scheduler_type} doesn't exist!") + + # Convert the UNet2DModel + unet_config = create_unet_diffusers_config(original_config, image_size=image_size) + unet = UNet2DConditionModel(**unet_config) + + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + checkpoint, unet_config, path=checkpoint_path, extract_ema=extract_ema + ) + + unet.load_state_dict(converted_unet_checkpoint) + + # Convert the VAE model + vae_config = create_vae_diffusers_config(original_config, checkpoint=checkpoint, image_size=image_size) + converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + + # Convert the text model + # MusicLDM uses the same tokenizer as the original CLAP model, but a slightly different configuration + config = ClapConfig.from_pretrained("laion/clap-htsat-unfused") + config.audio_config.update( + { + "patch_embeds_hidden_size": 128, + "hidden_size": 1024, + "depths": [2, 2, 12, 2], + } + ) + tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused") + feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused") + + converted_text_model = convert_open_clap_checkpoint(checkpoint, config) + text_model = ClapModel(config) + + missing_keys, unexpected_keys = text_model.load_state_dict(converted_text_model, strict=False) + # we expect not to have token_type_ids in our original state dict so let's ignore them + missing_keys = list(set(missing_keys) - set(CLAP_EXPECTED_MISSING_KEYS)) + + if len(unexpected_keys) > 0: + raise ValueError(f"Unexpected keys when loading CLAP model: {unexpected_keys}") + + if len(missing_keys) > 0: + raise ValueError(f"Missing keys when loading CLAP model: {missing_keys}") + + # Convert the vocoder model + vocoder_config = create_transformers_vocoder_config(original_config) + vocoder_config = SpeechT5HifiGanConfig(**vocoder_config) + converted_vocoder_checkpoint = convert_hifigan_checkpoint(checkpoint, vocoder_config) + + vocoder = SpeechT5HifiGan(vocoder_config) + vocoder.load_state_dict(converted_vocoder_checkpoint) + + # Instantiate the diffusers pipeline + pipe = MusicLDMPipeline( + vae=vae, + text_encoder=text_model, + tokenizer=tokenizer, + unet=unet, + scheduler=scheduler, + vocoder=vocoder, + feature_extractor=feature_extractor, + ) + + return pipe + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--original_config_file", + default=None, + type=str, + help="The YAML config file corresponding to the original architecture.", + ) + parser.add_argument( + "--num_in_channels", + default=None, + type=int, + help="The number of input channels. If `None` number of input channels will be automatically inferred.", + ) + parser.add_argument( + "--model_channels", + default=None, + type=int, + help="The number of UNet model channels. If `None`, it will be automatically inferred from the config. Override" + " to 128 for the small checkpoints, 192 for the medium checkpoints and 256 for the large.", + ) + parser.add_argument( + "--num_head_channels", + default=None, + type=int, + help="The number of UNet head channels. If `None`, it will be automatically inferred from the config. Override" + " to 32 for the small and medium checkpoints, and 64 for the large.", + ) + parser.add_argument( + "--scheduler_type", + default="ddim", + type=str, + help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", + ) + parser.add_argument( + "--image_size", + default=None, + type=int, + help=("The image size that the model was trained on."), + ) + parser.add_argument( + "--prediction_type", + default=None, + type=str, + help=("The prediction type that the model was trained on."), + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument( + "--from_safetensors", + action="store_true", + help="If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.", + ) + parser.add_argument( + "--to_safetensors", + action="store_true", + help="Whether to store pipeline in safetensors format or not.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") + args = parser.parse_args() + + pipe = load_pipeline_from_original_MusicLDM_ckpt( + checkpoint_path=args.checkpoint_path, + original_config_file=args.original_config_file, + image_size=args.image_size, + prediction_type=args.prediction_type, + extract_ema=args.extract_ema, + scheduler_type=args.scheduler_type, + num_in_channels=args.num_in_channels, + model_channels=args.model_channels, + num_head_channels=args.num_head_channels, + from_safetensors=args.from_safetensors, + device=args.device, + ) + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/scripts/recipes/ncsnpp.py b/scripts/recipes/ncsnpp.py new file mode 100644 index 000000000000..fe297f0f77c1 --- /dev/null +++ b/scripts/recipes/ncsnpp.py @@ -0,0 +1,81 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the NCSNPP checkpoints.""" + +import argparse +import json + +import torch + +from diffusers import ScoreSdeVePipeline, ScoreSdeVeScheduler, UNet2DModel +from diffusers.loaders.conversion import get_conversion + + +def convert_ncsnpp_checkpoint(checkpoint, config): + return get_conversion("UNet2DModel", {**config, "original_format": "ncsnpp"}).to_diffusers(checkpoint) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", + default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_pytorch_model.bin", + type=str, + required=False, + help="Path to the checkpoint to convert.", + ) + + parser.add_argument( + "--config_file", + default="/Users/arthurzucker/Work/diffusers/ArthurZ/config.json", + type=str, + required=False, + help="The config json file corresponding to the architecture.", + ) + + parser.add_argument( + "--dump_path", + default="/Users/arthurzucker/Work/diffusers/ArthurZ/diffusion_model_new.pt", + type=str, + required=False, + help="Path to the output model.", + ) + + args = parser.parse_args() + + checkpoint = torch.load(args.checkpoint_path, map_location="cpu") + + with open(args.config_file) as f: + config = json.loads(f.read()) + + converted_checkpoint = convert_ncsnpp_checkpoint( + checkpoint, + config, + ) + + if "sde" in config: + del config["sde"] + + model = UNet2DModel(**config) + model.load_state_dict(converted_checkpoint) + + try: + scheduler = ScoreSdeVeScheduler.from_config("/".join(args.checkpoint_path.split("/")[:-1])) + + pipe = ScoreSdeVePipeline(unet=model, scheduler=scheduler) + pipe.save_pretrained(args.dump_path) + except: # noqa: E722 + model.save_pretrained(args.dump_path) diff --git a/scripts/convert_omnigen_to_diffusers.py b/scripts/recipes/omnigen.py similarity index 74% rename from scripts/convert_omnigen_to_diffusers.py rename to scripts/recipes/omnigen.py index 96bc935633f0..b35c41aa5cb0 100644 --- a/scripts/convert_omnigen_to_diffusers.py +++ b/scripts/recipes/omnigen.py @@ -7,6 +7,7 @@ from transformers import AutoTokenizer from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, OmniGenPipeline, OmniGenTransformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint def main(args): @@ -25,41 +26,6 @@ def main(args): ckpt = os.path.join(args.origin_ckpt_path, "model.safetensors") ckpt = load_file(ckpt, device="cpu") - mapping_dict = { - "pos_embed": "patch_embedding.pos_embed", - "x_embedder.proj.weight": "patch_embedding.output_image_proj.weight", - "x_embedder.proj.bias": "patch_embedding.output_image_proj.bias", - "input_x_embedder.proj.weight": "patch_embedding.input_image_proj.weight", - "input_x_embedder.proj.bias": "patch_embedding.input_image_proj.bias", - "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight", - "final_layer.adaLN_modulation.1.bias": "norm_out.linear.bias", - "final_layer.linear.weight": "proj_out.weight", - "final_layer.linear.bias": "proj_out.bias", - "time_token.mlp.0.weight": "time_token.linear_1.weight", - "time_token.mlp.0.bias": "time_token.linear_1.bias", - "time_token.mlp.2.weight": "time_token.linear_2.weight", - "time_token.mlp.2.bias": "time_token.linear_2.bias", - "t_embedder.mlp.0.weight": "t_embedder.linear_1.weight", - "t_embedder.mlp.0.bias": "t_embedder.linear_1.bias", - "t_embedder.mlp.2.weight": "t_embedder.linear_2.weight", - "t_embedder.mlp.2.bias": "t_embedder.linear_2.bias", - "llm.embed_tokens.weight": "embed_tokens.weight", - } - - converted_state_dict = {} - for k, v in ckpt.items(): - if k in mapping_dict: - converted_state_dict[mapping_dict[k]] = v - elif "qkv" in k: - to_q, to_k, to_v = v.chunk(3) - converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_q.weight"] = to_q - converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_k.weight"] = to_k - converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_v.weight"] = to_v - elif "o_proj" in k: - converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_out.0.weight"] = v - else: - converted_state_dict[k[4:]] = v - transformer = OmniGenTransformer2DModel( rope_scaling={ "long_factor": [ @@ -168,7 +134,9 @@ def main(args): in_channels=4, pos_embed_max_size=192, ) - transformer.load_state_dict(converted_state_dict, strict=True) + transformer.load_state_dict( + convert_component_checkpoint(ckpt, dict(transformer.config), "OmniGenTransformer2DModel"), strict=True + ) transformer.to(torch.bfloat16) num_model_params = sum(p.numel() for p in transformer.parameters()) diff --git a/scripts/recipes/pixart_alpha.py b/scripts/recipes/pixart_alpha.py new file mode 100644 index 000000000000..0029c43efaac --- /dev/null +++ b/scripts/recipes/pixart_alpha.py @@ -0,0 +1,79 @@ +import argparse +import os + +import torch +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, PixArtAlphaPipeline, Transformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +ckpt_id = "PixArt-alpha/PixArt-alpha" +# https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/scripts/inference.py#L125 +interpolation_scale = {256: 0.5, 512: 1, 1024: 2} + + +def main(args): + all_state_dict = torch.load(args.orig_ckpt_path, map_location="cpu") + state_dict = all_state_dict.pop("state_dict") + # DiT XL/2 + transformer = Transformer2DModel( + sample_size=args.image_size // 8, + num_layers=28, + attention_head_dim=72, + in_channels=4, + out_channels=8, + patch_size=2, + attention_bias=True, + num_attention_heads=16, + cross_attention_dim=1152, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + caption_channels=4096, + ) + transformer.load_state_dict( + convert_component_checkpoint(state_dict, dict(transformer.config), "PixArtTransformer2DModel"), strict=True + ) + + num_model_params = sum(p.numel() for p in transformer.parameters()) + print(f"Total number of transformer parameters: {num_model_params}") + + if args.only_transformer: + transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) + else: + scheduler = DPMSolverMultistepScheduler() + + vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="sd-vae-ft-ema") + + tokenizer = T5Tokenizer.from_pretrained(ckpt_id, subfolder="t5-v1_1-xxl") + text_encoder = T5EncoderModel.from_pretrained(ckpt_id, subfolder="t5-v1_1-xxl") + + pipeline = PixArtAlphaPipeline( + tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler + ) + + pipeline.save_pretrained(args.dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--orig_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--image_size", + default=1024, + type=int, + choices=[256, 512, 1024], + required=False, + help="Image size of pretrained model, either 512 or 1024.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") + parser.add_argument("--only_transformer", default=True, type=bool, required=True) + + args = parser.parse_args() + main(args) diff --git a/scripts/recipes/pixart_sigma.py b/scripts/recipes/pixart_sigma.py new file mode 100644 index 000000000000..22d01a0cb111 --- /dev/null +++ b/scripts/recipes/pixart_sigma.py @@ -0,0 +1,87 @@ +import argparse +import os + +import torch +from transformers import T5EncoderModel, T5Tokenizer + +from diffusers import AutoencoderKL, DPMSolverMultistepScheduler, PixArtSigmaPipeline, Transformer2DModel +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint + + +ckpt_id = "PixArt-alpha" +# https://github.com/PixArt-alpha/PixArt-sigma/blob/dd087141864e30ec44f12cb7448dd654be065e88/scripts/inference.py#L158 +interpolation_scale = {256: 0.5, 512: 1, 1024: 2, 2048: 4} + + +def main(args): + all_state_dict = torch.load(args.orig_ckpt_path) + state_dict = all_state_dict.pop("state_dict") + transformer = Transformer2DModel( + sample_size=args.image_size // 8, + num_layers=28, + attention_head_dim=72, + in_channels=4, + out_channels=8, + patch_size=2, + attention_bias=True, + num_attention_heads=16, + cross_attention_dim=1152, + activation_fn="gelu-approximate", + num_embeds_ada_norm=1000, + norm_type="ada_norm_single", + norm_elementwise_affine=False, + norm_eps=1e-6, + caption_channels=4096, + interpolation_scale=interpolation_scale[args.image_size], + use_additional_conditions=args.micro_condition, + ) + transformer.load_state_dict( + convert_component_checkpoint(state_dict, dict(transformer.config), "PixArtTransformer2DModel"), strict=True + ) + + num_model_params = sum(p.numel() for p in transformer.parameters()) + print(f"Total number of transformer parameters: {num_model_params}") + + if args.only_transformer: + transformer.save_pretrained(os.path.join(args.dump_path, "transformer")) + else: + # pixart-Sigma vae link: https://huggingface.co/PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers/tree/main/vae + vae = AutoencoderKL.from_pretrained(f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="vae") + + scheduler = DPMSolverMultistepScheduler() + + tokenizer = T5Tokenizer.from_pretrained(f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="tokenizer") + text_encoder = T5EncoderModel.from_pretrained( + f"{ckpt_id}/pixart_sigma_sdxlvae_T5_diffusers", subfolder="text_encoder" + ) + + pipeline = PixArtSigmaPipeline( + tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, scheduler=scheduler + ) + + pipeline.save_pretrained(args.dump_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--micro_condition", action="store_true", help="If use Micro-condition in PixArtMS structure during training." + ) + parser.add_argument("--qk_norm", action="store_true", help="If use qk norm during training.") + parser.add_argument( + "--orig_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--image_size", + default=1024, + type=int, + choices=[256, 512, 1024, 2048], + required=False, + help="Image size of pretrained model, 256, 512, 1024, or 2048.", + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") + parser.add_argument("--only_transformer", default=True, type=bool, required=True) + + args = parser.parse_args() + main(args) diff --git a/scripts/convert_prx_to_diffusers.py b/scripts/recipes/prx.py similarity index 73% rename from scripts/convert_prx_to_diffusers.py rename to scripts/recipes/prx.py index 7b1677d0d3ae..2a976a133c40 100644 --- a/scripts/convert_prx_to_diffusers.py +++ b/scripts/recipes/prx.py @@ -22,12 +22,12 @@ import json import os import sys -from dataclasses import asdict, dataclass -from typing import Dict, Optional, Tuple +from typing import Dict import torch from safetensors.torch import save_file +from diffusers.loaders.conversion.configs.prx import VARIANTS, build_config, create_scheduler_config from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel @@ -41,61 +41,6 @@ PIXEL_PROMPT_MAX_TOKENS = 256 -@dataclass(frozen=True) -class PRXBase: - context_in_dim: int = 2304 - hidden_size: int = 1792 - mlp_ratio: float = 3.5 - num_heads: int = 28 - depth: int = 16 - axes_dim: Tuple[int, int] = (32, 32) - theta: int = 10_000 - time_factor: float = 1000.0 - time_max_period: int = 10_000 - bottleneck_size: Optional[int] = None - resolution_embeds: bool = False - - -@dataclass(frozen=True) -class PRXFlux(PRXBase): - in_channels: int = 16 - patch_size: int = 2 - - -@dataclass(frozen=True) -class PRXDCAE(PRXBase): - in_channels: int = 32 - patch_size: int = 1 - - -@dataclass(frozen=True) -class PRXPixel(PRXBase): - # Pixel-space RGB diffusion (PRXPixel / 7B). - in_channels: int = 3 - patch_size: int = 16 - context_in_dim: int = 2048 # Qwen3-VL-Embedding-2B hidden size - hidden_size: int = 3584 - num_heads: int = 28 - depth: int = 24 - axes_dim: Tuple[int, int] = (64, 64) - bottleneck_size: int = 768 - resolution_embeds: bool = True - - -VARIANTS = {"flux": PRXFlux, "dc-ae": PRXDCAE, "pixel": PRXPixel} - - -def build_config(variant: str) -> dict: - if variant not in VARIANTS: - raise ValueError(f"Unsupported variant: {variant}. Choose from {list(VARIANTS)}") - config_dict = asdict(VARIANTS[variant]()) - config_dict["axes_dim"] = list(config_dict["axes_dim"]) - if config_dict["bottleneck_size"] is None: - # Keep config.json clean for variants that don't use the bottleneck. - config_dict.pop("bottleneck_size") - return config_dict - - # --------------------------------------------------------------------------- # Checkpoint loading # --------------------------------------------------------------------------- @@ -151,71 +96,17 @@ def load_denoiser_state_dict(checkpoint_path: str, prefix: str = DENOISER_PREFIX # --------------------------------------------------------------------------- # Parameter name remapping (research -> diffusers) # --------------------------------------------------------------------------- -def create_parameter_mapping(depth: int) -> dict: - """Map old parameter names (layers on PRXBlock) to diffusers names (layers on PRXAttention).""" - mapping = {} - for i in range(depth): - mapping[f"blocks.{i}.img_qkv_proj.weight"] = f"blocks.{i}.attention.img_qkv_proj.weight" - mapping[f"blocks.{i}.txt_kv_proj.weight"] = f"blocks.{i}.attention.txt_kv_proj.weight" - mapping[f"blocks.{i}.qk_norm.query_norm.scale"] = f"blocks.{i}.attention.norm_q.weight" - mapping[f"blocks.{i}.qk_norm.key_norm.scale"] = f"blocks.{i}.attention.norm_k.weight" - mapping[f"blocks.{i}.qk_norm.query_norm.weight"] = f"blocks.{i}.attention.norm_q.weight" - mapping[f"blocks.{i}.qk_norm.key_norm.weight"] = f"blocks.{i}.attention.norm_k.weight" - mapping[f"blocks.{i}.k_norm.scale"] = f"blocks.{i}.attention.norm_added_k.weight" - mapping[f"blocks.{i}.k_norm.weight"] = f"blocks.{i}.attention.norm_added_k.weight" - mapping[f"blocks.{i}.attn_out.weight"] = f"blocks.{i}.attention.to_out.0.weight" - return mapping - - -def convert_checkpoint_parameters(old_state_dict: Dict[str, torch.Tensor], depth: int) -> dict[str, torch.Tensor]: - """Apply the block remapping. Unmapped keys (img_in, time_in, txt_in, resolution_embedder, final_layer) - carry over unchanged.""" - mapping = create_parameter_mapping(depth) - converted = {} - num_mapped = 0 - for key, value in old_state_dict.items(): - new_key = mapping.get(key, key) - if new_key != key: - num_mapped += 1 - converted[new_key] = value - print(f"✓ Converted {len(converted)} parameters ({num_mapped} block keys remapped)") - return converted - - -def create_transformer_from_checkpoint(checkpoint_path: str, config: dict) -> PRXTransformer2DModel: - """Create and load a PRXTransformer2DModel from a research checkpoint.""" - state_dict = load_denoiser_state_dict(checkpoint_path) - converted = convert_checkpoint_parameters(state_dict, depth=int(config["depth"])) - - print("Creating PRXTransformer2DModel...") - transformer = PRXTransformer2DModel(**config) - - # Match the checkpoint dtype (research saves bf16). - param_dtype = next(iter(converted.values())).dtype - transformer = transformer.to(param_dtype) - - missing, unexpected = transformer.load_state_dict(converted, strict=False) - if missing: - print(f"⚠ Missing keys ({len(missing)}): {missing}") - if unexpected: - print(f"⚠ Unexpected keys ({len(unexpected)}): {unexpected}") - if not missing and not unexpected: - print("✓ All parameters loaded successfully (0 missing, 0 unexpected)!") - else: - raise RuntimeError("Checkpoint did not load cleanly; see missing/unexpected keys above.") - return transformer + + +def create_transformer_from_checkpoint(checkpoint_path, config): + state = load_denoiser_state_dict(checkpoint_path) + dtype = next(iter(state.values())).dtype + return PRXTransformer2DModel.from_single_file(state, config=config, torch_dtype=dtype) # --------------------------------------------------------------------------- # Auxiliary components # --------------------------------------------------------------------------- -def create_scheduler_config(output_path: str, shift: float): - scheduler_config = {"_class_name": "FlowMatchEulerDiscreteScheduler", "num_train_timesteps": 1000, "shift": shift} - scheduler_path = os.path.join(output_path, "scheduler") - os.makedirs(scheduler_path, exist_ok=True) - with open(os.path.join(scheduler_path, "scheduler_config.json"), "w") as f: - json.dump(scheduler_config, f, indent=2) - print("✓ Created scheduler config") def download_and_save_vae(variant: str, output_path: str): @@ -336,7 +227,10 @@ def main(args): print(f"✓ Saved transformer to {transformer_path} ({num_params:,} params)") # ---- scheduler ---- - create_scheduler_config(args.output_path, args.shift) + scheduler_path = os.path.join(args.output_path, "scheduler") + os.makedirs(scheduler_path, exist_ok=True) + with open(os.path.join(scheduler_path, "scheduler_config.json"), "w") as handle: + json.dump(create_scheduler_config(args.shift), handle, indent=2) # ---- vae (none for pixel) ---- if args.variant != "pixel" and not args.skip_vae: diff --git a/scripts/convert_rae_to_diffusers.py b/scripts/recipes/rae.py similarity index 76% rename from scripts/convert_rae_to_diffusers.py rename to scripts/recipes/rae.py index 0502e49ef30c..aa15b60c171e 100644 --- a/scripts/convert_rae_to_diffusers.py +++ b/scripts/recipes/rae.py @@ -3,32 +3,14 @@ from typing import Any import torch +from accelerate import init_empty_weights from huggingface_hub import HfApi, hf_hub_download from diffusers import AutoencoderRAE +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.rae import DECODER_CONFIGS -DECODER_CONFIGS = { - "ViTB": { - "decoder_hidden_size": 768, - "decoder_intermediate_size": 3072, - "decoder_num_attention_heads": 12, - "decoder_num_hidden_layers": 12, - }, - "ViTL": { - "decoder_hidden_size": 1024, - "decoder_intermediate_size": 4096, - "decoder_num_attention_heads": 16, - "decoder_num_hidden_layers": 24, - }, - "ViTXL": { - "decoder_hidden_size": 1152, - "decoder_intermediate_size": 4096, - "decoder_num_attention_heads": 16, - "decoder_num_hidden_layers": 28, - }, -} - ENCODER_DEFAULT_NAME_OR_PATH = { "dinov2": "facebook/dinov2-with-registers-base", "siglip2": "google/siglip2-base-patch16-256", @@ -110,27 +92,6 @@ def unwrap_state_dict(maybe_wrapped: dict[str, Any]) -> dict[str, Any]: return out -def remap_decoder_attention_keys_for_diffusers(state_dict: dict[str, Any]) -> dict[str, Any]: - """ - Map official RAE decoder attention key layout to diffusers Attention layout used by AutoencoderRAE decoder. - - Example mappings: - - `...attention.attention.query.*` -> `...attention.to_q.*` - - `...attention.attention.key.*` -> `...attention.to_k.*` - - `...attention.attention.value.*` -> `...attention.to_v.*` - - `...attention.output.dense.*` -> `...attention.to_out.0.*` - """ - remapped: dict[str, Any] = {} - for key, value in state_dict.items(): - new_key = key - new_key = new_key.replace(".attention.attention.query.", ".attention.to_q.") - new_key = new_key.replace(".attention.attention.key.", ".attention.to_k.") - new_key = new_key.replace(".attention.attention.value.", ".attention.to_v.") - new_key = new_key.replace(".attention.output.dense.", ".attention.to_out.0.") - remapped[new_key] = value - return remapped - - def resolve_decoder_file( accessor: RepoAccessor, encoder_type: str, variant: str, decoder_checkpoint: str | None ) -> str: @@ -218,8 +179,11 @@ def _load_hf_encoder_state_dict(encoder_type: str, encoder_name_or_path: str) -> # Our Siglip2Encoder wraps it inside SiglipVisionModel which nests it # under .vision_model, so we add the prefix to match the diffusers key layout. hf_model = SiglipModel.from_pretrained(encoder_name_or_path).vision_model - sd = {f"vision_model.{k}": v for k, v in hf_model.state_dict().items()} - return _strip_final_layernorm_affine(sd, prefix="vision_model.post_layernorm.") + import transformers + + prefix = "" if int(transformers.__version__.split(".")[0]) >= 5 else "vision_model." + sd = {prefix + key: value for key, value in hf_model.state_dict().items()} + return _strip_final_layernorm_affine(sd, prefix=prefix + "post_layernorm.") elif encoder_type == "mae": from transformers import ViTMAEForPreTraining @@ -249,7 +213,6 @@ def convert(args: argparse.Namespace) -> None: decoder_path = accessor.fetch(decoder_relpath) decoder_obj = torch.load(decoder_path, map_location="cpu") decoder_state_dict = unwrap_state_dict(decoder_obj) - decoder_state_dict = remap_decoder_attention_keys_for_diffusers(decoder_state_dict) latents_mean, latents_std = None, None if stats_relpath is not None: @@ -283,7 +246,7 @@ def convert(args: argparse.Namespace) -> None: encoder_state_dict = _load_hf_encoder_state_dict(args.encoder_type, encoder_name_or_path) # Build model on meta device to avoid double init overhead - with torch.device("meta"): + with init_empty_weights(include_buffers=False): model = AutoencoderRAE( encoder_type=args.encoder_type, encoder_hidden_size=encoder_hidden_size, @@ -303,45 +266,14 @@ def convert(args: argparse.Namespace) -> None: scaling_factor=args.scaling_factor, ) - # Assemble full state dict and load with assign=True - full_state_dict = {} - - # Encoder weights (prefixed with "encoder.") - for k, v in encoder_state_dict.items(): - full_state_dict[f"encoder.{k}"] = v - - # Decoder weights (prefixed with "decoder.") - for k, v in decoder_state_dict.items(): - full_state_dict[f"decoder.{k}"] = v - - # Buffers from config - full_state_dict["encoder_mean"] = torch.tensor(encoder_norm_mean, dtype=torch.float32).view(1, 3, 1, 1) - full_state_dict["encoder_std"] = torch.tensor(encoder_norm_std, dtype=torch.float32).view(1, 3, 1, 1) - if latents_mean is not None: - latents_mean_t = latents_mean if isinstance(latents_mean, torch.Tensor) else torch.tensor(latents_mean) - full_state_dict["_latents_mean"] = latents_mean_t - else: - full_state_dict["_latents_mean"] = torch.zeros(1) - if latents_std is not None: - latents_std_t = latents_std if isinstance(latents_std, torch.Tensor) else torch.tensor(latents_std) - full_state_dict["_latents_std"] = latents_std_t - else: - full_state_dict["_latents_std"] = torch.ones(1) - - model.load_state_dict(full_state_dict, strict=False, assign=True) - - # Verify no critical keys are missing - model_keys = {name for name, _ in model.named_parameters()} - model_keys |= {name for name, _ in model.named_buffers()} - loaded_keys = set(full_state_dict.keys()) - missing = model_keys - loaded_keys - # decoder_pos_embed is initialized in-model. trainable_cls_token is only - # allowed to be missing if it was absent in the source decoder checkpoint. - allowed_missing = {"decoder.decoder_pos_embed"} - if "trainable_cls_token" not in decoder_state_dict: - allowed_missing.add("decoder.trainable_cls_token") - if missing - allowed_missing: - print(f"Warning: missing keys after conversion: {sorted(missing - allowed_missing)}") + state = {f"encoder.{key}": value for key, value in encoder_state_dict.items()} + state.update({f"decoder.{key}": value for key, value in decoder_state_dict.items()}) + state["processor.image_mean"] = torch.tensor(encoder_norm_mean, dtype=torch.float32) + state["processor.image_std"] = torch.tensor(encoder_norm_std, dtype=torch.float32) + state["stats.latents_mean"] = torch.as_tensor(latents_mean) if latents_mean is not None else torch.zeros(1) + state["stats.latents_std"] = torch.as_tensor(latents_std) if latents_std is not None else torch.ones(1) + converted = convert_component_checkpoint(state, dict(model.config), "AutoencoderRAE") + model.load_state_dict(converted, strict=True, assign=True) output_path = Path(args.output_path) output_path.mkdir(parents=True, exist_ok=True) diff --git a/scripts/convert_sana_to_diffusers.py b/scripts/recipes/sana.py similarity index 55% rename from scripts/convert_sana_to_diffusers.py rename to scripts/recipes/sana.py index 833e64ba75e7..ace29adfeda0 100644 --- a/scripts/convert_sana_to_diffusers.py +++ b/scripts/recipes/sana.py @@ -20,7 +20,7 @@ SanaTransformer2DModel, SCMScheduler, ) -from diffusers.models.model_loading_utils import load_model_dict_into_meta +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint from diffusers.utils.import_utils import is_accelerate_available @@ -66,163 +66,8 @@ def main(args): print(colored(f"Loading checkpoint from {file_path}", "green", attrs=["bold"])) all_state_dict = torch.load(file_path, weights_only=True) state_dict = all_state_dict.pop("state_dict") - converted_state_dict = {} - - # Patch embeddings. - converted_state_dict["patch_embed.proj.weight"] = state_dict.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = state_dict.pop("x_embedder.proj.bias") - - # Caption projection. - converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") - - # Handle different time embedding structure based on model type - - if args.model_type in ["SanaSprint_1600M_P1_D20", "SanaSprint_600M_P1_D28"]: - # For Sana Sprint, the time embedding structure is different - converted_state_dict["time_embed.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") - - # Guidance embedder for Sana Sprint - converted_state_dict["time_embed.guidance_embedder.linear_1.weight"] = state_dict.pop( - "cfg_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.guidance_embedder.linear_1.bias"] = state_dict.pop("cfg_embedder.mlp.0.bias") - converted_state_dict["time_embed.guidance_embedder.linear_2.weight"] = state_dict.pop( - "cfg_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.guidance_embedder.linear_2.bias"] = state_dict.pop("cfg_embedder.mlp.2.bias") - else: - # Original Sana time embedding structure - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = state_dict.pop( - "t_embedder.mlp.0.bias" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = state_dict.pop( - "t_embedder.mlp.2.bias" - ) - - # Shared norm. - converted_state_dict["time_embed.linear.weight"] = state_dict.pop("t_block.1.weight") - converted_state_dict["time_embed.linear.bias"] = state_dict.pop("t_block.1.bias") - - # y norm - converted_state_dict["caption_norm.weight"] = state_dict.pop("attention_y_norm.weight") - - # scheduler - if args.image_size == 4096: - flow_shift = 6.0 - else: - flow_shift = 3.0 - - # model config - if args.model_type in ["SanaMS_1600M_P1_D20", "SanaSprint_1600M_P1_D20", "SanaMS1.5_1600M_P1_D20"]: - layer_num = 20 - elif args.model_type in ["SanaMS_600M_P1_D28", "SanaSprint_600M_P1_D28"]: - layer_num = 28 - elif args.model_type == "SanaMS_4800M_P1_D60": - layer_num = 60 - else: - raise ValueError(f"{args.model_type} is not supported.") - # Positional embedding interpolation scale. + flow_shift = 6.0 if args.image_size == 4096 else 3.0 interpolation_scale = {512: None, 1024: None, 2048: 1.0, 4096: 2.0} - qk_norm = ( - "rms_norm_across_heads" - if args.model_type - in ["SanaMS1.5_1600M_P1_D20", "SanaMS1.5_4800M_P1_D60", "SanaSprint_600M_P1_D28", "SanaSprint_1600M_P1_D20"] - else None - ) - - for depth in range(layer_num): - # Transformer blocks. - converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( - f"blocks.{depth}.scale_shift_table" - ) - - # Linear Attention is all you need 🤘 - # Self attention. - q, k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.weight"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - if qk_norm is not None: - # Add Q/K normalization for self-attention (attn1) - needed for Sana-Sprint and Sana-1.5 - converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_q.weight"] = state_dict.pop( - f"blocks.{depth}.attn.q_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_k.weight"] = state_dict.pop( - f"blocks.{depth}.attn.k_norm.weight" - ) - # Projection. - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.attn.proj.bias" - ) - - # Feed-forward. - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.inverted_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.inverted_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.depth_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.depth_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_point.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.point_conv.conv.weight" - ) - - # Cross-attention. - q = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.weight") - q_bias = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.bias") - k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.weight"), 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.bias"), 2, dim=0) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias - if qk_norm is not None: - # Add Q/K normalization for cross-attention (attn2) - needed for Sana-Sprint and Sana-1.5 - converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_q.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.q_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_k.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.k_norm.weight" - ) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.bias" - ) - - # Final block. - converted_state_dict["proj_out.weight"] = state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = state_dict.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = state_dict.pop("final_layer.scale_shift_table") # Transformer with CTX(): @@ -258,20 +103,8 @@ def main(args): transformer = SanaTransformer2DModel(**transformer_kwargs) - if is_accelerate_available(): - load_model_dict_into_meta(transformer, converted_state_dict) - else: - transformer.load_state_dict(converted_state_dict, strict=True, assign=True) - - try: - state_dict.pop("y_embedder.y_embedding") - state_dict.pop("pos_embed") - state_dict.pop("logvar_linear.weight") - state_dict.pop("logvar_linear.bias") - except KeyError: - print("y_embedder.y_embedding or pos_embed not found in the state_dict") - - assert len(state_dict) == 0, f"State dict is not empty, {state_dict.keys()}" + converted_state_dict = convert_component_checkpoint(state_dict, dict(transformer.config), "SanaTransformer2DModel") + transformer.load_state_dict(converted_state_dict, strict=True, assign=True) num_model_params = sum(p.numel() for p in transformer.parameters()) print(f"Total number of transformer parameters: {num_model_params}") diff --git a/scripts/recipes/sana_controlnet.py b/scripts/recipes/sana_controlnet.py new file mode 100644 index 000000000000..9646b463a767 --- /dev/null +++ b/scripts/recipes/sana_controlnet.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +from __future__ import annotations + +import argparse +from contextlib import nullcontext + +import torch +from accelerate import init_empty_weights + +from diffusers import ( + SanaControlNetModel, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.utils.import_utils import is_accelerate_available + + +CTX = init_empty_weights if is_accelerate_available else nullcontext + + +def main(args): + file_path = args.orig_ckpt_path + + all_state_dict = torch.load(file_path, weights_only=True) + state_dict = all_state_dict.pop("state_dict") + interpolation_scale = {512: None, 1024: None, 2048: 1.0, 4096: 2.0} + + # ControlNet + with CTX(): + controlnet = SanaControlNetModel( + num_attention_heads=model_kwargs[args.model_type]["num_attention_heads"], + attention_head_dim=model_kwargs[args.model_type]["attention_head_dim"], + num_layers=model_kwargs[args.model_type]["num_layers"], + num_cross_attention_heads=model_kwargs[args.model_type]["num_cross_attention_heads"], + cross_attention_head_dim=model_kwargs[args.model_type]["cross_attention_head_dim"], + cross_attention_dim=model_kwargs[args.model_type]["cross_attention_dim"], + caption_channels=2304, + sample_size=args.image_size // 32, + interpolation_scale=interpolation_scale[args.image_size], + ) + + converted_state_dict = convert_component_checkpoint(state_dict, dict(controlnet.config), "SanaControlNetModel") + controlnet.load_state_dict(converted_state_dict, strict=True, assign=True) + + num_model_params = sum(p.numel() for p in controlnet.parameters()) + print(f"Total number of controlnet parameters: {num_model_params}") + + controlnet = controlnet.to(weight_dtype) + + print(f"Saving Sana ControlNet in Diffusers format in {args.dump_path}.") + controlnet.save_pretrained(args.dump_path) + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + +VARIANT_MAPPING = { + "fp32": None, + "fp16": "fp16", + "bf16": "bf16", +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--orig_ckpt_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--image_size", + default=1024, + type=int, + choices=[512, 1024, 2048, 4096], + required=False, + help="Image size of pretrained model, 512, 1024, 2048 or 4096.", + ) + parser.add_argument( + "--model_type", + default="SanaMS_1600M_P1_ControlNet_D7", + type=str, + choices=["SanaMS_1600M_P1_ControlNet_D7", "SanaMS_600M_P1_ControlNet_D7"], + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.") + parser.add_argument("--dtype", default="fp16", type=str, choices=["fp32", "fp16", "bf16"], help="Weight dtype.") + + args = parser.parse_args() + + model_kwargs = { + "SanaMS_1600M_P1_ControlNet_D7": { + "num_attention_heads": 70, + "attention_head_dim": 32, + "num_cross_attention_heads": 20, + "cross_attention_head_dim": 112, + "cross_attention_dim": 2240, + "num_layers": 7, + }, + "SanaMS_600M_P1_ControlNet_D7": { + "num_attention_heads": 36, + "attention_head_dim": 32, + "num_cross_attention_heads": 16, + "cross_attention_head_dim": 72, + "cross_attention_dim": 1152, + "num_layers": 7, + }, + } + + device = "cuda" if torch.cuda.is_available() else "cpu" + weight_dtype = DTYPE_MAPPING[args.dtype] + variant = VARIANT_MAPPING[args.dtype] + + main(args) diff --git a/scripts/convert_sana_video_to_diffusers.py b/scripts/recipes/sana_video.py similarity index 56% rename from scripts/convert_sana_video_to_diffusers.py rename to scripts/recipes/sana_video.py index c6be52d455b8..ea8374bd8fc9 100644 --- a/scripts/convert_sana_video_to_diffusers.py +++ b/scripts/recipes/sana_video.py @@ -20,6 +20,7 @@ SanaVideoTransformer3DModel, UniPCMultistepScheduler, ) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint from diffusers.utils.import_utils import is_accelerate_available @@ -54,43 +55,13 @@ def main(args): print(colored(f"Loading checkpoint from {file_path}", "green", attrs=["bold"])) all_state_dict = torch.load(file_path, weights_only=True) state_dict = all_state_dict.pop("state_dict") - converted_state_dict = {} - - # Patch embeddings. - converted_state_dict["patch_embedding.weight"] = state_dict.pop("x_embedder.proj.weight") - converted_state_dict["patch_embedding.bias"] = state_dict.pop("x_embedder.proj.bias") - - # Caption projection. - converted_state_dict["caption_projection.linear_1.weight"] = state_dict.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = state_dict.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = state_dict.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = state_dict.pop("y_embedder.y_proj.fc2.bias") - - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = state_dict.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = state_dict.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = state_dict.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = state_dict.pop("t_embedder.mlp.2.bias") - - # Shared norm. - converted_state_dict["time_embed.linear.weight"] = state_dict.pop("t_block.1.weight") - converted_state_dict["time_embed.linear.bias"] = state_dict.pop("t_block.1.bias") - - # y norm - converted_state_dict["caption_norm.weight"] = state_dict.pop("attention_y_norm.weight") - # scheduler flow_shift = 8.0 if args.task == "i2v": assert args.scheduler_type == "flow-euler", "Scheduler type must be flow-euler for i2v task." # model config - layer_num = 20 # Positional embedding interpolation scale. - qk_norm = True # sample size if args.video_size == 480: @@ -112,87 +83,6 @@ def main(args): in_channels = 128 out_channels = 128 - for depth in range(layer_num): - # Transformer blocks. - converted_state_dict[f"transformer_blocks.{depth}.scale_shift_table"] = state_dict.pop( - f"blocks.{depth}.scale_shift_table" - ) - - # Linear Attention is all you need 🤘 - # Self attention. - q, k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.attn.qkv.weight"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_v.weight"] = v - if qk_norm is not None: - # Add Q/K normalization for self-attention (attn1) - needed for Sana-Sprint and Sana-1.5 - converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_q.weight"] = state_dict.pop( - f"blocks.{depth}.attn.q_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.norm_k.weight"] = state_dict.pop( - f"blocks.{depth}.attn.k_norm.weight" - ) - # Projection. - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn1.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.attn.proj.bias" - ) - - # Feed-forward. - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.inverted_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_inverted.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.inverted_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.depth_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_depth.bias"] = state_dict.pop( - f"blocks.{depth}.mlp.depth_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_point.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.point_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.ff.conv_temp.weight"] = state_dict.pop( - f"blocks.{depth}.mlp.t_conv.weight" - ) - - # Cross-attention. - q = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.weight") - q_bias = state_dict.pop(f"blocks.{depth}.cross_attn.q_linear.bias") - k, v = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.weight"), 2, dim=0) - k_bias, v_bias = torch.chunk(state_dict.pop(f"blocks.{depth}.cross_attn.kv_linear.bias"), 2, dim=0) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.weight"] = q - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_q.bias"] = q_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.weight"] = k - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_k.bias"] = k_bias - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.weight"] = v - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_v.bias"] = v_bias - if qk_norm is not None: - # Add Q/K normalization for cross-attention (attn2) - needed for Sana-Sprint and Sana-1.5 - converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_q.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.q_norm.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.norm_k.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.k_norm.weight" - ) - - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.weight"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{depth}.attn2.to_out.0.bias"] = state_dict.pop( - f"blocks.{depth}.cross_attn.proj.bias" - ) - - # Final block. - converted_state_dict["proj_out.weight"] = state_dict.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = state_dict.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = state_dict.pop("final_layer.scale_shift_table") - # Transformer with CTX(): transformer_kwargs = { @@ -217,18 +107,11 @@ def main(args): transformer = SanaVideoTransformer3DModel(**transformer_kwargs) + converted_state_dict = convert_component_checkpoint( + state_dict, dict(transformer.config), "SanaVideoTransformer3DModel" + ) transformer.load_state_dict(converted_state_dict, strict=True, assign=True) - try: - state_dict.pop("y_embedder.y_embedding") - state_dict.pop("pos_embed") - state_dict.pop("logvar_linear.weight") - state_dict.pop("logvar_linear.bias") - except KeyError: - print("y_embedder.y_embedding or pos_embed not found in the state_dict") - - assert len(state_dict) == 0, f"State dict is not empty, {state_dict.keys()}" - num_model_params = sum(p.numel() for p in transformer.parameters()) print(f"Total number of transformer parameters: {num_model_params}") diff --git a/scripts/recipes/shap_e.py b/scripts/recipes/shap_e.py new file mode 100644 index 000000000000..e8842fbb8134 --- /dev/null +++ b/scripts/recipes/shap_e.py @@ -0,0 +1,252 @@ +import argparse +import tempfile + +import torch +from accelerate import load_checkpoint_and_dispatch + +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.shap_e import PRIOR_CONFIG, PRIOR_IMAGE_CONFIG, RENDERER_CONFIG +from diffusers.models.transformers.prior_transformer import PriorTransformer +from diffusers.pipelines.shap_e import ShapERenderer + + +""" +Example - From the diffusers root directory: + +Download weights: +```sh +$ wget "https://openaipublic.azureedge.net/main/shap-e/text_cond.pt" +``` + +Convert the model: +```sh +$ python scripts/recipes/shap_e.py \ + --prior_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/text_cond.pt \ + --prior_image_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/image_cond.pt \ + --transmitter_checkpoint_path /home/yiyi_huggingface_co/shap-e/shap_e_model_cache/transmitter.pt\ + --dump_path /home/yiyi_huggingface_co/model_repo/shap-e-img2img/shap_e_renderer\ + --debug renderer +``` +""" + + +# prior + + +def prior_model_from_original_config(): + model = PriorTransformer(**PRIOR_CONFIG) + + return model + + +def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = {key: value for key, value in checkpoint.items() if key.startswith("wrapped.")} + return get_conversion("PriorTransformer", {**dict(model.config), "original_format": "shap_e"}).to_diffusers(state) + + +# done prior + + +# prior_image (only slightly different from prior) + + +# Uses default arguments + + +def prior_image_model_from_original_config(): + model = PriorTransformer(**PRIOR_IMAGE_CONFIG) + + return model + + +def prior_image_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = {key: value for key, value in checkpoint.items() if key.startswith("wrapped.")} + return get_conversion("PriorTransformer", {**dict(model.config), "original_format": "shap_e"}).to_diffusers(state) + + +# done prior_image + + +# renderer + +## create the lookup table for marching cubes method used in MeshDecoder + + +def renderer_model_from_original_config(): + model = ShapERenderer(**RENDERER_CONFIG) + + return model + + +def renderer_model_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = { + key: value + for key, value in checkpoint.items() + if key.startswith(("renderer.nerstf.mlp.", "encoder.params_proj.projections.")) + } + return get_conversion("ShapERenderer", dict(model.config)).to_diffusers(state) + + +# done renderer + + +# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) + + +# done unet utils + + +# Driver functions + + +def prior(*, args, checkpoint_map_location): + print("loading prior") + + prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) + + prior_model = prior_model_from_original_config() + + prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint(prior_model, prior_checkpoint) + + del prior_checkpoint + + load_prior_checkpoint_to_model(prior_diffusers_checkpoint, prior_model) + + print("done loading prior") + + return prior_model + + +def prior_image(*, args, checkpoint_map_location): + print("loading prior_image") + + print(f"load checkpoint from {args.prior_image_checkpoint_path}") + prior_checkpoint = torch.load(args.prior_image_checkpoint_path, map_location=checkpoint_map_location) + + prior_model = prior_image_model_from_original_config() + + prior_diffusers_checkpoint = prior_image_original_checkpoint_to_diffusers_checkpoint(prior_model, prior_checkpoint) + + del prior_checkpoint + + load_prior_checkpoint_to_model(prior_diffusers_checkpoint, prior_model) + + print("done loading prior_image") + + return prior_model + + +def renderer(*, args, checkpoint_map_location): + print(" loading renderer") + + renderer_checkpoint = torch.load(args.transmitter_checkpoint_path, map_location=checkpoint_map_location) + + renderer_model = renderer_model_from_original_config() + + renderer_diffusers_checkpoint = renderer_model_original_checkpoint_to_diffusers_checkpoint( + renderer_model, renderer_checkpoint + ) + + del renderer_checkpoint + + load_checkpoint_to_model(renderer_diffusers_checkpoint, renderer_model, strict=True) + + print("done loading renderer") + + return renderer_model + + +# prior model will expect clip_mean and clip_std, which are missing from the state_dict +PRIOR_EXPECTED_MISSING_KEYS = ["clip_mean", "clip_std"] + + +def load_prior_checkpoint_to_model(checkpoint, model): + with tempfile.NamedTemporaryFile() as file: + torch.save(checkpoint, file.name) + del checkpoint + missing_keys, unexpected_keys = model.load_state_dict(torch.load(file.name), strict=False) + missing_keys = list(set(missing_keys) - set(PRIOR_EXPECTED_MISSING_KEYS)) + + if len(unexpected_keys) > 0: + raise ValueError(f"Unexpected keys when loading prior model: {unexpected_keys}") + if len(missing_keys) > 0: + raise ValueError(f"Missing keys when loading prior model: {missing_keys}") + + +def load_checkpoint_to_model(checkpoint, model, strict=False): + with tempfile.NamedTemporaryFile() as file: + torch.save(checkpoint, file.name) + del checkpoint + if strict: + model.load_state_dict(torch.load(file.name), strict=True) + else: + load_checkpoint_and_dispatch(model, file.name, device_map="auto") + + +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( + "--prior_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the prior checkpoint to convert.", + ) + + parser.add_argument( + "--prior_image_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the prior_image checkpoint to convert.", + ) + + parser.add_argument( + "--transmitter_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to the transmitter checkpoint to convert.", + ) + + parser.add_argument( + "--checkpoint_load_device", + default="cpu", + type=str, + required=False, + help="The device passed to `map_location` when loading checkpoints.", + ) + + parser.add_argument( + "--debug", + default=None, + type=str, + required=False, + help="Only run a specific stage of the convert script. Used for debugging", + ) + + args = parser.parse_args() + + print(f"loading checkpoints to {args.checkpoint_load_device}") + + checkpoint_map_location = torch.device(args.checkpoint_load_device) + + if args.debug is not None: + print(f"debug: only executing {args.debug}") + + if args.debug is None: + print("YiYi TO-DO") + elif args.debug == "prior": + prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) + prior_model.save_pretrained(args.dump_path) + elif args.debug == "prior_image": + prior_model = prior_image(args=args, checkpoint_map_location=checkpoint_map_location) + prior_model.save_pretrained(args.dump_path) + elif args.debug == "renderer": + renderer_model = renderer(args=args, checkpoint_map_location=checkpoint_map_location) + renderer_model.save_pretrained(args.dump_path) + else: + raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/recipes/skyreels_v2.py b/scripts/recipes/skyreels_v2.py new file mode 100644 index 000000000000..60e1820779fe --- /dev/null +++ b/scripts/recipes/skyreels_v2.py @@ -0,0 +1,124 @@ +import argparse +import os +import pathlib + +import torch +from huggingface_hub import hf_hub_download +from safetensors.torch import load_file +from transformers import AutoProcessor, AutoTokenizer, CLIPVisionModelWithProjection, UMT5EncoderModel + +from diffusers import ( + AutoencoderKLWan, + SkyReelsV2DiffusionForcingPipeline, + SkyReelsV2ImageToVideoPipeline, + SkyReelsV2Pipeline, + SkyReelsV2Transformer3DModel, + UniPCMultistepScheduler, +) +from diffusers.loaders.conversion.configs.skyreels_v2 import get_transformer_config + + +def load_sharded_safetensors(path): + from diffusers.loaders.conversion.source import load_tensor_sources + + return load_tensor_sources(path) + + +def convert_transformer(model_type: str): + config = get_transformer_config(model_type) + diffusers_config = config["diffusers_config"] + model_id = config["model_id"] + + if "1.3B" in model_type: + original_state_dict = load_file(hf_hub_download(model_id, "model.safetensors")) + else: + os.makedirs(model_type, exist_ok=True) + model_dir = pathlib.Path(model_type) + if "720P" in model_type: + top_shard = 7 if "I2V" in model_type else 6 + zeros = "0" * (4 if "I2V" or "T2V" in model_type else 3) + model_name = "diffusion_pytorch_model" + elif "540P" in model_type: + top_shard = 14 if "I2V" in model_type else 12 + model_name = "model" + + for i in range(1, top_shard + 1): + shard_path = f"{model_name}-{i:05d}-of-{zeros}{top_shard}.safetensors" + hf_hub_download(model_id, shard_path, local_dir=model_dir) + original_state_dict = load_sharded_safetensors(model_dir) + + return SkyReelsV2Transformer3DModel.from_single_file(original_state_dict, config=diffusers_config) + + +def convert_vae(): + path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") + return AutoencoderKLWan.from_single_file(path, config={}) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_type", type=str, default=None) + parser.add_argument("--output_path", type=str, required=True) + parser.add_argument("--dtype", default="fp32") + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +if __name__ == "__main__": + args = get_args() + + transformer = None + dtype = DTYPE_MAPPING[args.dtype] + + transformer = convert_transformer(args.model_type).to(dtype=dtype) + vae = convert_vae() + text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl") + tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") + scheduler = UniPCMultistepScheduler( + prediction_type="flow_prediction", + num_train_timesteps=1000, + use_flow_sigmas=True, + ) + + if "I2V" in args.model_type or "FLF2V" in args.model_type: + image_encoder = CLIPVisionModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") + image_processor = AutoProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") + pipe = SkyReelsV2ImageToVideoPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + image_encoder=image_encoder, + image_processor=image_processor, + ) + elif "T2V" in args.model_type: + pipe = SkyReelsV2Pipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + elif "DF" in args.model_type: + pipe = SkyReelsV2DiffusionForcingPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + + pipe.save_pretrained( + args.output_path, + safe_serialization=True, + max_shard_size="5GB", + # push_to_hub=True, + # repo_id=f"/{args.model_type}-Diffusers", + ) diff --git a/scripts/recipes/stable_audio.py b/scripts/recipes/stable_audio.py new file mode 100644 index 000000000000..59930e1cc8a5 --- /dev/null +++ b/scripts/recipes/stable_audio.py @@ -0,0 +1,183 @@ +# Run this script to convert the Stable Audio model weights to a diffusers pipeline. +import argparse +import json +import os +from contextlib import nullcontext + +import torch +from safetensors.torch import load_file +from transformers import ( + AutoTokenizer, + T5EncoderModel, +) + +from diffusers import ( + AutoencoderOobleck, + CosineDPMSolverMultistepScheduler, + StableAudioDiTModel, + StableAudioPipeline, + StableAudioProjectionModel, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.utils import is_accelerate_available + + +if is_accelerate_available(): + from accelerate import init_empty_weights + + +parser = argparse.ArgumentParser(description="Convert Stable Audio 1.0 model weights to a diffusers pipeline") +parser.add_argument("--model_folder_path", type=str, help="Location of Stable Audio weights and config") +parser.add_argument("--use_safetensors", action="store_true", help="Use SafeTensors for conversion") +parser.add_argument( + "--save_directory", + type=str, + default="./tmp/stable-audio-1.0", + help="Directory to save a pipeline to. Will be created if it doesn't exist.", +) +parser.add_argument( + "--repo_id", + type=str, + default="stable-audio-1.0", + help="Hub organization to save the pipelines to", +) +parser.add_argument("--push_to_hub", action="store_true", help="Push to hub") +parser.add_argument("--variant", type=str, help="Set to bf16 to save bfloat16 weights") + +args = parser.parse_args() + +checkpoint_path = ( + os.path.join(args.model_folder_path, "model.safetensors") + if args.use_safetensors + else os.path.join(args.model_folder_path, "model.ckpt") +) +config_path = os.path.join(args.model_folder_path, "model_config.json") + +device = "cpu" +if args.variant == "bf16": + dtype = torch.bfloat16 +else: + dtype = torch.float32 + +with open(config_path) as f_in: + config_dict = json.load(f_in) + +conditioning_dict = { + conditioning["id"]: conditioning["config"] for conditioning in config_dict["model"]["conditioning"]["configs"] +} + +t5_model_config = conditioning_dict["prompt"] + +# T5 Text encoder +text_encoder = T5EncoderModel.from_pretrained(t5_model_config["t5_model_name"]) +tokenizer = AutoTokenizer.from_pretrained( + t5_model_config["t5_model_name"], truncation=True, model_max_length=t5_model_config["max_length"] +) + + +# scheduler +scheduler = CosineDPMSolverMultistepScheduler( + sigma_min=0.3, + sigma_max=500, + solver_order=2, + prediction_type="v_prediction", + sigma_data=1.0, + sigma_schedule="exponential", +) +ctx = init_empty_weights if is_accelerate_available() else nullcontext + + +if args.use_safetensors: + orig_state_dict = load_file(checkpoint_path, device=device) +else: + orig_state_dict = torch.load(checkpoint_path, map_location=device) + + +model_config = config_dict["model"]["diffusion"]["config"] + +model_state_dict = { + key.removeprefix("model.model."): value for key, value in orig_state_dict.items() if key.startswith("model.model.") +} +projection_model_state_dict = { + key.removeprefix("conditioner.conditioners."): value + for key, value in orig_state_dict.items() + if key.startswith("conditioner.conditioners.") +} +autoencoder_state_dict = { + key.removeprefix("pretransform.model."): value + for key, value in orig_state_dict.items() + if key.startswith("pretransform.model.") +} + + +with ctx(): + projection_model = StableAudioProjectionModel( + text_encoder_dim=text_encoder.config.d_model, + conditioning_dim=config_dict["model"]["conditioning"]["cond_dim"], + min_value=conditioning_dict["seconds_start"][ + "min_val" + ], # assume `seconds_start` and `seconds_total` have the same min / max values. + max_value=conditioning_dict["seconds_start"][ + "max_val" + ], # assume `seconds_start` and `seconds_total` have the same min / max values. + ) +projection_model.load_state_dict( + get_conversion("StableAudioProjectionModel", dict(projection_model.config)).to_diffusers( + projection_model_state_dict + ), + strict=True, + assign=True, +) + +attention_head_dim = model_config["embed_dim"] // model_config["num_heads"] +with ctx(): + model = StableAudioDiTModel( + sample_size=int(config_dict["sample_size"]) + / int(config_dict["model"]["pretransform"]["config"]["downsampling_ratio"]), + in_channels=model_config["io_channels"], + num_layers=model_config["depth"], + attention_head_dim=attention_head_dim, + num_key_value_attention_heads=model_config["cond_token_dim"] // attention_head_dim, + num_attention_heads=model_config["num_heads"], + out_channels=model_config["io_channels"], + cross_attention_dim=model_config["cond_token_dim"], + time_proj_dim=256, + global_states_input_dim=model_config["global_cond_dim"], + cross_attention_input_dim=model_config["cond_token_dim"], + ) +model.load_state_dict( + get_conversion("StableAudioDiTModel", dict(model.config)).to_diffusers(model_state_dict), strict=True, assign=True +) + + +autoencoder_config = config_dict["model"]["pretransform"]["config"] +with ctx(): + autoencoder = AutoencoderOobleck( + encoder_hidden_size=autoencoder_config["encoder"]["config"]["channels"], + downsampling_ratios=autoencoder_config["encoder"]["config"]["strides"], + decoder_channels=autoencoder_config["decoder"]["config"]["channels"], + decoder_input_channels=autoencoder_config["decoder"]["config"]["latent_dim"], + audio_channels=autoencoder_config["io_channels"], + channel_multiples=autoencoder_config["encoder"]["config"]["c_mults"], + sampling_rate=config_dict["sample_rate"], + ) + +autoencoder.load_state_dict( + get_conversion("AutoencoderOobleck", dict(autoencoder.config)).to_diffusers(autoencoder_state_dict), + strict=True, + assign=True, +) + + +# Prior pipeline +pipeline = StableAudioPipeline( + transformer=model, + tokenizer=tokenizer, + text_encoder=text_encoder, + scheduler=scheduler, + vae=autoencoder, + projection_model=projection_model, +) +pipeline.to(dtype).save_pretrained( + args.save_directory, repo_id=args.repo_id, push_to_hub=args.push_to_hub, variant=args.variant +) diff --git a/scripts/convert_stable_audio_3_to_diffusers.py b/scripts/recipes/stable_audio3.py similarity index 54% rename from scripts/convert_stable_audio_3_to_diffusers.py rename to scripts/recipes/stable_audio3.py index 1d940ea2159e..8b5ea77ae5d5 100644 --- a/scripts/convert_stable_audio_3_to_diffusers.py +++ b/scripts/recipes/stable_audio3.py @@ -41,7 +41,7 @@ rotary_pos_emb, per-block to_local_embed for inpainting) is converted. Usage: - python scripts/convert_stable_audio_3_to_diffusers.py \\ + python scripts/recipes/stable_audio3.py \\ --checkpoint_path stabilityai/stable-audio-3-medium \\ --model_config_path /path/to/model_config.json \\ --output_dir /path/to/output \\ @@ -51,14 +51,19 @@ import argparse import json -import math import sys from pathlib import Path -from typing import Optional import torch from safetensors.torch import load_file +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.stable_audio3 import ( + _infer_dit_config, + _infer_duration_embedder_config, + _infer_vae_config, +) + # Ensure UTF-8 stdout/stderr for Unicode output if hasattr(sys.stdout, "reconfigure"): @@ -71,56 +76,16 @@ # ────────────────────────────────────────────────────────────────────────────── -def _reorder_differential_qkv(weight: torch.Tensor) -> torch.Tensor: - """Reorder fused self-attention QKV weight from reference layout to diffusers layout. - - Reference layout (stable-audio-tools): - to_qkv output chunks: [q | k | v | q2 | k2] - Diffusers layout (StableAudio3SelfAttention): - to_qkv output chunks: [q1 | q2 | k1 | k2 | v] - """ - D = weight.shape[0] // 5 - q, k, v, q2, k2 = weight.split(D, dim=0) - return torch.cat([q, q2, k, k2, v], dim=0) - - # ────────────────────────────────────────────────────────────────────────────── # Key-transformation helpers # ────────────────────────────────────────────────────────────────────────────── # Inside a TransformerResamplingBlock's per-layer transformer blocks: # ref key → our key -_TRB_BLOCK_RENAMES = [ - ("pre_norm.", "norm_attn."), - ("self_attn.to_qkv.", "attn.to_qkv."), # fused; QKV reorder applied separately - ("self_attn.to_out.", "attn.to_out."), - ("self_attn.q_norm.", "attn.q_norm."), - ("self_attn.k_norm.", "attn.k_norm."), - ("ff_norm.", "norm_ff."), - ("ff.ff.0.proj.", "ff.proj_in."), # GLU → proj_in - ("ff.ff.2.", "ff.proj_out."), # linear_out → proj_out -] # Inside a DiT block (StableAudio3DiTBlock). The diffusers block keeps the reference norm and # attention names, so most entries are identity; only the SwiGLU feed-forward is renamed. # ref key → our key -_DIT_BLOCK_RENAMES = [ - ("ff.ff.0.proj.", "ff.proj_in."), - ("ff.ff.2.", "ff.proj_out."), -] - - -def _apply_renames(local_key: str, renames) -> str: - for old, new in renames: - if old in local_key: - local_key = local_key.replace(old, new, 1) - break - return local_key - - -def _is_differential_qkv_key(key: str) -> bool: - """Return True if this key is a fused QKV weight that needs reordering.""" - return "self_attn.to_qkv.weight" in key or "attn.to_qkv.weight" in key # ────────────────────────────────────────────────────────────────────────────── @@ -128,130 +93,10 @@ def _is_differential_qkv_key(key: str) -> bool: # ────────────────────────────────────────────────────────────────────────────── -def _infer_trb_depth(ref_sd: dict, base: str) -> int: - """Count how many TRB blocks exist under `base` (e.g. 'encoder.layers').""" - depth = 0 - while f"{base}.{depth}.new_tokens" in ref_sd: - depth += 1 - return depth - - -def _convert_trb_block( - ref_sd: dict, - ref_prefix: str, # e.g. "encoder.layers.0." - our_prefix: str, # e.g. "encoder.blocks.0." - differential: bool, - out: dict, - skipped: list, -): - """Copy all keys inside one TransformerResamplingBlock.""" - for key, val in ref_sd.items(): - if not key.startswith(ref_prefix): - continue - local = key[len(ref_prefix) :] # strip leading prefix - - # --- mapping (WNConv1d): weight_g, weight_v, bias copied as-is --- - if local.startswith("mapping."): - out[our_prefix + local] = val - continue - - # --- new_tokens --- - if local == "new_tokens": - out[our_prefix + local] = val - continue - - # --- transformers.{j}.* --- - if local.startswith("transformers."): - new_local = _apply_renames(local, _TRB_BLOCK_RENAMES) - new_key = our_prefix + new_local - if differential and _is_differential_qkv_key(new_key): - val = _reorder_differential_qkv(val) - out[new_key] = val - continue - - skipped.append(key) - - -def convert_vae(ref_sd: dict, differential: bool = True) -> dict: - """Build diffusers AutoencoderSAME state dict from reference checkpoint.""" - out = {} - skipped = [] - - ae_prefix = "pretransform.model." - - enc_base = ae_prefix + "encoder.layers" - dec_base = ae_prefix + "decoder.layers" - - enc_depth = _infer_trb_depth(ref_sd, enc_base) - - if enc_depth == 0: - print("WARNING: could not detect encoder TRB depth; defaulting to 1.") - enc_depth = 1 - - # ---- ENCODER TRBs ---- - # ref: encoder.layers.{i} (i = 0..enc_depth-1) → our: encoder.blocks.{i} - for i in range(enc_depth): - _convert_trb_block( - ref_sd, - f"{enc_base}.{i}.", - f"encoder.blocks.{i}.", - differential, - out, - skipped, - ) - - # ---- ENCODER projection (Linear after all TRBs) ---- - # Comes after enc_depth TRBs and one Transpose → index enc_depth+1 - # But Transpose has no params, so the Linear shows as layers.{enc_depth+1} - # (enc_depth TRBs, then 2 Transposes and 1 Linear: layers[enc_depth] = Transpose, - # layers[enc_depth+1] = Linear) - enc_linear_idx = enc_depth + 1 - enc_linear_prefix = f"{enc_base}.{enc_linear_idx}." - for key, val in ref_sd.items(): - if key.startswith(enc_linear_prefix): - suffix = key[len(enc_linear_prefix) :] - out[f"encoder.proj.{suffix}"] = val - - # ---- DECODER projection (Linear before all TRBs) ---- - # Decoder layers: [Transpose, Linear, Transpose, TRB_0, TRB_1, ...] - # Linear is always at index 1 - dec_linear_prefix = f"{dec_base}.1." - for key, val in ref_sd.items(): - if key.startswith(dec_linear_prefix): - suffix = key[len(dec_linear_prefix) :] - out[f"decoder.proj.{suffix}"] = val - - # ---- DECODER TRBs ---- - # Decoder TRBs start at index 3 in reference. - # Reference builds them in reverse order (channel_dims[depth]→[depth-1], …). - # Our SAMEDecoder.blocks also reverses: blocks[0] = largest stride, same order. - for i in range(enc_depth): # same depth as encoder - ref_idx = 3 + i - _convert_trb_block( - ref_sd, - f"{dec_base}.{ref_idx}.", - f"decoder.blocks.{i}.", - differential, - out, - skipped, - ) - - # ---- BOTTLENECK ---- - bn_prefix = ae_prefix + "bottleneck." - bottleneck_renames = { - "scaling_factor": "scale", - "bias": "bias", - "running_std": "running_std", - } - for ref_suffix, our_suffix in bottleneck_renames.items(): - full_key = bn_prefix + ref_suffix - if full_key in ref_sd: - out[f"bottleneck.{our_suffix}"] = ref_sd[full_key] - - if skipped: - print(f" VAE: skipped {len(skipped)} keys (pretransform.* patcher, noise_scaling_factor, ...)") - - return out +def convert_vae(ref_sd, config): + prefix = "pretransform.model." + state = {key.removeprefix(prefix): value for key, value in ref_sd.items() if key.startswith(prefix)} + return convert_component_checkpoint(state, config, "AutoencoderSAME") # ────────────────────────────────────────────────────────────────────────────── @@ -259,30 +104,10 @@ def convert_vae(ref_sd: dict, differential: bool = True) -> dict: # ────────────────────────────────────────────────────────────────────────────── -def convert_duration_embedder(ref_sd: dict, min_freq: float = 0.5, max_freq: float = 10000.0) -> dict: - """Build diffusers StableAudio3DurationEmbedder state dict.""" - out = {} - base = "conditioner.conditioners.seconds_total.embedder.embedding." - # index 0 = ExpoFourierFeatures (no learnable params) - # index 1 = nn.Linear - for suffix in ("weight", "bias"): - src = f"{base}1.{suffix}" - if src in ref_sd: - out[f"linear.{suffix}"] = ref_sd[src] - else: - print(f" DurationEmbedder: key not found: {src}") - - # Compute the freqs buffer (matches StableAudio3DurationEmbedder.__init__) - if "linear.weight" in out: - fourier_dim = out["linear.weight"].shape[1] - half = fourier_dim // 2 - ramp = torch.linspace(0.0, 1.0, half) - log_min = math.log(min_freq) - log_max = math.log(max_freq) - freqs = torch.exp(ramp * (log_max - log_min) + log_min) - out["freqs"] = freqs - - return out +def convert_duration_embedder(ref_sd): + state = {key: value for key, value in ref_sd.items() if key.startswith("conditioner.conditioners.seconds_total.")} + state.pop("conditioner.conditioners.seconds_total.embedder.embedding.0.weights", None) + return convert_component_checkpoint(state, _infer_duration_embedder_config(ref_sd), "StableAudio3DurationEmbedder") # ────────────────────────────────────────────────────────────────────────────── @@ -290,79 +115,13 @@ def convert_duration_embedder(ref_sd: dict, min_freq: float = 0.5, max_freq: flo # ────────────────────────────────────────────────────────────────────────────── -def convert_dit(ref_sd: dict, differential: bool = True) -> dict: - """Build diffusers StableAudio3DiTModel state dict.""" - out = {} - - # The DiT weights live under the "model.model." prefix in the checkpoint. - p = "model.model." - - # Top-level keys (outside the inner transformer). - top_renames = { - f"{p}to_timestep_embed.0.weight": "to_timestep_embed.0.weight", - f"{p}to_timestep_embed.0.bias": "to_timestep_embed.0.bias", - f"{p}to_timestep_embed.2.weight": "to_timestep_embed.2.weight", - f"{p}to_timestep_embed.2.bias": "to_timestep_embed.2.bias", - f"{p}to_cond_embed.0.weight": "to_cond_embed.0.weight", - f"{p}to_cond_embed.2.weight": "to_cond_embed.2.weight", - f"{p}to_global_embed.0.weight": "to_global_embed.0.weight", - f"{p}to_global_embed.2.weight": "to_global_embed.2.weight", - f"{p}preprocess_conv.weight": "preprocess_conv.weight", - f"{p}postprocess_conv.weight": "postprocess_conv.weight", - # Inner transformer module-level keys. - f"{p}transformer.project_in.weight": "proj_in.weight", - f"{p}transformer.project_out.weight": "proj_out.weight", - f"{p}transformer.memory_tokens": "memory_tokens", - f"{p}transformer.rotary_pos_emb.inv_freq": "rotary_pos_emb.inv_freq", - f"{p}transformer.global_cond_embedder.0.weight": "global_cond_embedder.0.weight", - f"{p}transformer.global_cond_embedder.0.bias": "global_cond_embedder.0.bias", - f"{p}transformer.global_cond_embedder.2.weight": "global_cond_embedder.2.weight", - f"{p}transformer.global_cond_embedder.2.bias": "global_cond_embedder.2.bias", +def convert_dit(ref_sd, config): + state = { + key: value + for key, value in ref_sd.items() + if key.startswith("model.model.") or key == "conditioner.conditioners.prompt.padding_embedding" } - - for ref_key, our_key in top_renames.items(): - if ref_key in ref_sd: - out[our_key] = ref_sd[ref_key] - else: - print(f" DiT: top-level key not found: {ref_key}") - - # Per-block keys: model.model.transformer.layers.{i}.* - block_prefix = f"{p}transformer.layers." - block_indices = set() - for k in ref_sd: - if k.startswith(block_prefix): - idx_str = k[len(block_prefix) :].split(".")[0] - if idx_str.isdigit(): - block_indices.add(int(idx_str)) - - for i in sorted(block_indices): - ref_blk = f"{block_prefix}{i}." - for key, val in ref_sd.items(): - if not key.startswith(ref_blk): - continue - local = key[len(ref_blk) :] - - new_local = _apply_renames(local, _DIT_BLOCK_RENAMES) - new_key = f"transformer_blocks.{i}.{new_local}" - - # Self-attention fused QKV needs reordering ([q|k|v|q2|k2] → [q1|q2|k1|k2|v]). - # Differential cross-attention to_q ([q|q2]) and to_kv ([k|k2|v]) already match the - # diffusers layout, so no cross-attention reorder is required. - if differential and _is_differential_qkv_key(new_key): - val = _reorder_differential_qkv(val) - - out[new_key] = val - - # The learned text-padding embedding lives on the conditioner in the reference checkpoint, - # but the diffusers DiT owns it: it replaces padded cross-attention positions with this vector - # and then attends to the full context (the reference disables the cross-attention mask). - pad_key = "conditioner.conditioners.prompt.padding_embedding" - if pad_key in ref_sd: - out["prompt_padding_embedding"] = ref_sd[pad_key] - else: - print(f" DiT: padding embedding not found: {pad_key}") - - return out + return convert_component_checkpoint(state, config, "StableAudio3DiTModel") # ────────────────────────────────────────────────────────────────────────────── @@ -385,194 +144,6 @@ def extract_text_encoder(ref_sd: dict) -> dict: # ────────────────────────────────────────────────────────────────────────────── -def _infer_vae_config(ref_sd: dict, model_config: Optional[dict] = None) -> dict: - """ - Infer AutoencoderSAME config from checkpoint tensor shapes. - - Some hyper-parameters (the TRB stride, the sliding-window width and the sinusoidal-FFN layer counts) are NOT - recoverable from the weights alone — with ``variable_stride`` the encoder/decoder ``new_tokens`` collapse to a - single shared token, so the stride is invisible. When ``model_config`` (the parsed ``model_config.json``) is - provided, those values are read directly from it; otherwise production SAME-L/S defaults are used. - """ - # Bottleneck scale shape: (1, latent_dim, 1) - latent_dim = ref_sd.get( - "pretransform.model.bottleneck.scaling_factor", - ref_sd.get("pretransform.model.bottleneck.scale", torch.zeros(1, 256, 1)), - ).shape[1] - - # Encoder linear weight: (latent_dim, channel_dims[-1]) - # We detect enc_depth to find the linear key - enc_base = "pretransform.model.encoder.layers" - enc_depth = _infer_trb_depth(ref_sd, enc_base) - if enc_depth == 0: - enc_depth = 1 - enc_linear_key = f"{enc_base}.{enc_depth + 1}.weight" - if enc_linear_key in ref_sd: - enc_final_ch = ref_sd[enc_linear_key].shape[1] # (latent_dim, enc_final_ch) - else: - enc_final_ch = 768 # SAME-S default - - # TRB 0 new_tokens shape (encoder): (1, 1, out_channels) - trb0_nt = ref_sd.get(f"{enc_base}.0.new_tokens") - enc_out_ch = trb0_nt.shape[2] if trb0_nt is not None else enc_final_ch - - # TRB 0 mapping.weight_v shape (encoder): (out_ch, in_ch, kernel) - mapping_wv = ref_sd.get(f"{enc_base}.0.mapping.weight_v") - if mapping_wv is not None: - patched_in = mapping_wv.shape[1] # = audio_channels * patch_size - else: - patched_in = 512 # default: 2ch * 256 patch - - # Infer transformer depth per TRB: - trb_trans_depth = sum( - 1 - for k in ref_sd - if k.startswith(f"{enc_base}.0.transformers.") - and k.endswith(".new_tokens") is False - and ".pre_norm.alpha" in k - ) - if trb_trans_depth == 0: - trb_trans_depth = 6 # default - - # Dim heads inferred from q_norm shape: (dim_heads,) - q_norm_key = f"{enc_base}.0.transformers.0.self_attn.q_norm.gamma" - dim_heads = ref_sd[q_norm_key].shape[0] if q_norm_key in ref_sd else 64 - - # channels base: choose so that enc_out_ch = channels * c_mults[0] - # We use c_mults = [6] for both SAME-S/L - c_mults = [6] - enc_channels_base = enc_out_ch // c_mults[0] - - # ── Weight-invisible hyper-parameters ──────────────────────────────────── - # These come from model_config.json when available (see docstring). - audio_channels = 2 - sliding_window = 1 - encoder_sinusoidal_blocks = [0] * enc_depth - decoder_sinusoidal_blocks = [0] * enc_depth - - if model_config is not None: - ae_cfg = model_config["model"]["pretransform"]["config"] - enc_cfg = ae_cfg["encoder"]["config"] - dec_cfg = ae_cfg["decoder"]["config"] - strides = list(enc_cfg["strides"]) - transformer_depths = list(enc_cfg["transformer_depths"]) - c_mults = list(enc_cfg["c_mults"]) - enc_channels_base = enc_cfg["channels"] - latent_dim = enc_cfg.get("latent_dim", latent_dim) - dim_heads = enc_cfg.get("dim_heads", dim_heads) - audio_channels = ae_cfg.get("io_channels", audio_channels) - downsampling_ratio = ae_cfg["downsampling_ratio"] - patch_size = downsampling_ratio // int(math.prod(strides)) - # sliding_window in the reference is a per-side list like [1, 1]; take the (symmetric) half-width. - sw = enc_cfg.get("sliding_window") or [sliding_window] - sliding_window = sw[0] - encoder_sinusoidal_blocks = list(enc_cfg.get("sinusoidal_blocks", encoder_sinusoidal_blocks)) - decoder_sinusoidal_blocks = list(dec_cfg.get("sinusoidal_blocks", decoder_sinusoidal_blocks)) - else: - # Stride is NOT recoverable from weights under variable_stride; assume the production value of 16. - strides = [16] * enc_depth - transformer_depths = [trb_trans_depth] * enc_depth - patch_size = patched_in // audio_channels - - return { - "audio_channels": audio_channels, - "patch_size": patch_size, - "encoder_channels": enc_channels_base, - "encoder_c_mults": c_mults, - "encoder_strides": strides, - "encoder_transformer_depths": transformer_depths, - "latent_dim": latent_dim, - "use_differential_attention": True, - "dim_heads": dim_heads, - "ff_mult": 3, - "sliding_window": sliding_window, - "encoder_sinusoidal_blocks": encoder_sinusoidal_blocks, - "decoder_sinusoidal_blocks": decoder_sinusoidal_blocks, - "sampling_rate": 44100, - } - - -def _infer_dit_config(ref_sd: dict) -> dict: - """Infer StableAudio3DiTModel config from checkpoint tensor shapes.""" - p = "model.model." - - # embed_dim from to_timestep_embed.0.weight shape (embed_dim, features_dim) - ts_w = ref_sd.get(f"{p}to_timestep_embed.0.weight") - embed_dim = ts_w.shape[0] if ts_w is not None else 1536 - timestep_features_dim = ts_w.shape[1] if ts_w is not None else 256 - - # depth: count transformer blocks (RMSNorm → pre_norm.gamma) - depth = 0 - while f"{p}transformer.layers.{depth}.pre_norm.gamma" in ref_sd: - depth += 1 - if depth == 0: - depth = 24 - - # num_heads: from self_attn.q_norm.gamma shape (dim_heads,) and embed_dim - q_norm_key = f"{p}transformer.layers.0.self_attn.q_norm.gamma" - dim_heads = ref_sd[q_norm_key].shape[0] if q_norm_key in ref_sd else 64 - num_heads = embed_dim // dim_heads - - # cond_token_dim / global_cond_dim from the projection in-weights (embed_dim, *) - cond_w = ref_sd.get(f"{p}to_cond_embed.0.weight") - cond_token_dim = cond_w.shape[1] if cond_w is not None else 768 - glob_w = ref_sd.get(f"{p}to_global_embed.0.weight") - global_cond_dim = glob_w.shape[1] if glob_w is not None else 768 - - # io_channels from preprocess_conv.weight (io_ch, io_ch, 1) - pc_w = ref_sd.get(f"{p}preprocess_conv.weight") - io_channels = pc_w.shape[0] if pc_w is not None else 256 - - # ff_mult from ff.ff.0.proj.weight (inner*2, embed_dim) - ff_w = ref_sd.get(f"{p}transformer.layers.0.ff.ff.0.proj.weight") - ff_mult = (ff_w.shape[0] // (2 * embed_dim)) if ff_w is not None else 4 - - # local_add_cond_dim from to_local_embed.0.weight (embed_dim, local_add_cond_dim) - loc_w = ref_sd.get(f"{p}transformer.layers.0.to_local_embed.0.weight") - local_add_cond_dim = loc_w.shape[1] if loc_w is not None else 257 - - # num_memory_tokens from memory_tokens (num_memory_tokens, embed_dim) - mem = ref_sd.get(f"{p}transformer.memory_tokens") - num_memory_tokens = mem.shape[0] if mem is not None else 64 - - # differential: self_attn.to_qkv rows = embed_dim*5 (differential) vs *3 (standard) - qkv_w = ref_sd.get(f"{p}transformer.layers.0.self_attn.to_qkv.weight") - use_differential = qkv_w is not None and qkv_w.shape[0] == embed_dim * 5 - - return { - "io_channels": io_channels, - "patch_size": 1, - "embed_dim": embed_dim, - "depth": depth, - "num_heads": num_heads, - "cond_token_dim": cond_token_dim, - "global_cond_dim": global_cond_dim, - "local_add_cond_dim": local_add_cond_dim, - "timestep_features_dim": timestep_features_dim, - "ff_mult": ff_mult, - "num_memory_tokens": num_memory_tokens, - "use_differential_attention": use_differential, - } - - -def _infer_duration_embedder_config(ref_sd: dict) -> dict: - """Infer StableAudio3DurationEmbedder config from checkpoint shapes.""" - w = ref_sd.get("conditioner.conditioners.seconds_total.embedder.embedding.1.weight") - if w is not None: - output_dim = w.shape[0] - fourier_dim = w.shape[1] - else: - output_dim, fourier_dim = 768, 256 - return { - "output_dim": output_dim, - "fourier_dim": fourier_dim, - "min_val": 0.0, - "max_val": 384.0, - "min_freq": 0.5, - "max_freq": 10000.0, - } - - # ────────────────────────────────────────────────────────────────────────────── # Main conversion # ────────────────────────────────────────────────────────────────────────────── @@ -635,7 +206,7 @@ def convert(args): vae_cfg = _infer_vae_config(ref_sd, model_config) print(f" Inferred VAE config: {vae_cfg}") - vae_sd = convert_vae(ref_sd, differential=vae_cfg["use_differential_attention"]) + vae_sd = convert_vae(ref_sd, config=vae_cfg) vae = AutoencoderSAME(**vae_cfg) missing, unexpected = vae.load_state_dict(vae_sd, strict=False) @@ -715,7 +286,7 @@ def convert(args): dit_cfg = _infer_dit_config(ref_sd) print(f" Inferred DiT config: {dit_cfg}") - dit_sd = convert_dit(ref_sd, differential=dit_cfg["use_differential_attention"]) + dit_sd = convert_dit(ref_sd, config=dit_cfg) transformer = StableAudio3DiTModel(**dit_cfg) dit_missing, dit_unexpected = transformer.load_state_dict(dit_sd, strict=False) diff --git a/scripts/convert_stable_cascade.py b/scripts/recipes/stable_cascade.py similarity index 95% rename from scripts/convert_stable_cascade.py rename to scripts/recipes/stable_cascade.py index 97ed18d9b4d4..49247aca8e13 100644 --- a/scripts/convert_stable_cascade.py +++ b/scripts/recipes/stable_cascade.py @@ -18,10 +18,10 @@ StableCascadeDecoderPipeline, StableCascadePriorPipeline, ) -from diffusers.loaders.single_file_utils import convert_stable_cascade_unet_single_file_to_diffusers +from diffusers.loaders.conversion import get_conversion from diffusers.models import StableCascadeUNet from diffusers.models.model_loading_utils import load_model_dict_into_meta -from diffusers.pipelines.wuerstchen import PaellaVQModel +from diffusers.pipelines.deprecated.wuerstchen import PaellaVQModel from diffusers.utils import is_accelerate_available @@ -96,8 +96,6 @@ else: prior_orig_state_dict = torch.load(prior_checkpoint_path, map_location=device) - prior_state_dict = convert_stable_cascade_unet_single_file_to_diffusers(prior_orig_state_dict) - with ctx(): prior_model = StableCascadeUNet( in_channels=16, @@ -125,6 +123,10 @@ timestep_conditioning_type=["sca", "crp"], switch_level=[False], ) + prior_state_dict = get_conversion("StableCascadeUNet", dict(prior_model.config)).to_diffusers( + prior_orig_state_dict + ) + if is_accelerate_available(): load_model_dict_into_meta(prior_model, prior_state_dict) else: @@ -150,7 +152,6 @@ else: decoder_orig_state_dict = torch.load(decoder_checkpoint_path, map_location=device) - decoder_state_dict = convert_stable_cascade_unet_single_file_to_diffusers(decoder_orig_state_dict) with ctx(): decoder = StableCascadeUNet( in_channels=4, @@ -180,6 +181,10 @@ timestep_conditioning_type=["sca"], ) + decoder_state_dict = get_conversion("StableCascadeUNet", dict(decoder.config)).to_diffusers( + decoder_orig_state_dict + ) + if is_accelerate_available(): load_model_dict_into_meta(decoder, decoder_state_dict) else: diff --git a/scripts/convert_stable_cascade_lite.py b/scripts/recipes/stable_cascade_lite.py similarity index 95% rename from scripts/convert_stable_cascade_lite.py rename to scripts/recipes/stable_cascade_lite.py index 5f4804e30f74..820c9498dd11 100644 --- a/scripts/convert_stable_cascade_lite.py +++ b/scripts/recipes/stable_cascade_lite.py @@ -18,10 +18,10 @@ StableCascadeDecoderPipeline, StableCascadePriorPipeline, ) -from diffusers.loaders.single_file_utils import convert_stable_cascade_unet_single_file_to_diffusers +from diffusers.loaders.conversion import get_conversion from diffusers.models import StableCascadeUNet from diffusers.models.model_loading_utils import load_model_dict_into_meta -from diffusers.pipelines.wuerstchen import PaellaVQModel +from diffusers.pipelines.deprecated.wuerstchen import PaellaVQModel from diffusers.utils import is_accelerate_available @@ -103,7 +103,6 @@ else: prior_orig_state_dict = torch.load(prior_checkpoint_path, map_location=device) - prior_state_dict = convert_stable_cascade_unet_single_file_to_diffusers(prior_orig_state_dict) with ctx(): prior_model = StableCascadeUNet( in_channels=16, @@ -132,6 +131,10 @@ switch_level=[False], ) + prior_state_dict = get_conversion("StableCascadeUNet", dict(prior_model.config)).to_diffusers( + prior_orig_state_dict + ) + if is_accelerate_available(): load_model_dict_into_meta(prior_model, prior_state_dict) else: @@ -157,8 +160,6 @@ else: decoder_orig_state_dict = torch.load(decoder_checkpoint_path, map_location=device) - decoder_state_dict = convert_stable_cascade_unet_single_file_to_diffusers(decoder_orig_state_dict) - with ctx(): decoder = StableCascadeUNet( in_channels=4, @@ -188,6 +189,10 @@ timestep_conditioning_type=["sca"], ) + decoder_state_dict = get_conversion("StableCascadeUNet", dict(decoder.config)).to_diffusers( + decoder_orig_state_dict + ) + if is_accelerate_available(): load_model_dict_into_meta(decoder, decoder_state_dict) else: diff --git a/scripts/convert_original_stable_diffusion_to_diffusers.py b/scripts/recipes/stable_diffusion.py similarity index 100% rename from scripts/convert_original_stable_diffusion_to_diffusers.py rename to scripts/recipes/stable_diffusion.py diff --git a/scripts/recipes/unclip.py b/scripts/recipes/unclip.py new file mode 100644 index 000000000000..dcf7e4b7701f --- /dev/null +++ b/scripts/recipes/unclip.py @@ -0,0 +1,431 @@ +import argparse +import tempfile + +import torch +from accelerate import load_checkpoint_and_dispatch +from transformers import CLIPTextModelWithProjection, CLIPTokenizer + +from diffusers import UnCLIPPipeline, UNet2DConditionModel, UNet2DModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.unclip import ( + DECODER_CONFIG, + PRIOR_CONFIG, + SUPER_RES_UNET_FIRST_STEPS_CONFIG, + SUPER_RES_UNET_LAST_STEP_CONFIG, +) +from diffusers.models.transformers.prior_transformer import PriorTransformer +from diffusers.pipelines.deprecated.unclip.text_proj import UnCLIPTextProjModel +from diffusers.schedulers.scheduling_unclip import UnCLIPScheduler + + +r""" +Example - From the diffusers root directory: + +Download weights: +```sh +$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/efdf6206d8ed593961593dc029a8affa/decoder-ckpt-step%3D01000000-of-01000000.ckpt +$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/4226b831ae0279020d134281f3c31590/improved-sr-ckpt-step%3D1.2M.ckpt +$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/85626483eaca9f581e2a78d31ff905ca/prior-ckpt-step%3D01000000-of-01000000.ckpt +$ wget https://arena.kakaocdn.net/brainrepo/models/karlo-public/v1.0.0.alpha/0b62380a75e56f073e2844ab5199153d/ViT-L-14_stats.th +``` + +Convert the model: +```sh +$ python scripts/recipes/unclip.py \ + --decoder_checkpoint_path ./decoder-ckpt-step\=01000000-of-01000000.ckpt \ + --super_res_unet_checkpoint_path ./improved-sr-ckpt-step\=1.2M.ckpt \ + --prior_checkpoint_path ./prior-ckpt-step\=01000000-of-01000000.ckpt \ + --clip_stat_path ./ViT-L-14_stats.th \ + --dump_path +``` +""" + + +# prior + + +# Uses default arguments + + +def prior_model_from_original_config(): + model = PriorTransformer(**PRIOR_CONFIG) + + return model + + +def prior_original_checkpoint_to_diffusers_checkpoint(model, checkpoint, clip_stats_checkpoint): + state = dict(checkpoint) + state["clip_stats.mean"], state["clip_stats.std"] = clip_stats_checkpoint + return get_conversion("PriorTransformer", {**dict(model.config), "original_format": "unclip"}).to_diffusers(state) + + +# done prior + + +# decoder + + +# We are hardcoding the model configuration for now. If we need to generalize to more model configurations, we can +# update then. + + +def decoder_model_from_original_config(): + model = UNet2DConditionModel(**DECODER_CONFIG) + + return model + + +def decoder_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + projection_prefixes = ("text_seq_proj.", "clip_tok_proj.", "text_feat_proj.", "clip_emb.", "cf_param") + state = {key.removeprefix("model."): value for key, value in checkpoint.items() if key.startswith("model.")} + state = {key: value for key, value in state.items() if not key.startswith(projection_prefixes)} + return get_conversion("UNet2DConditionModel", dict(model.config)).to_diffusers(state) + + +# done decoder + +# text proj + + +def text_proj_from_original_config(): + # From the conditional unet constructor where the dimension of the projected time embeddings is + # constructed + time_embed_dim = DECODER_CONFIG["block_out_channels"][0] * 4 + + cross_attention_dim = DECODER_CONFIG["cross_attention_dim"] + + model = UnCLIPTextProjModel(time_embed_dim=time_embed_dim, cross_attention_dim=cross_attention_dim) + + return model + + +# Note that the input checkpoint is the original decoder checkpoint +def text_proj_original_checkpoint_to_diffusers_checkpoint(checkpoint): + prefixes = ( + "model.text_seq_proj.", + "model.clip_tok_proj.", + "model.text_feat_proj.", + "model.clip_emb.", + "model.cf_param", + ) + state = {key.removeprefix("model."): value for key, value in checkpoint.items() if key.startswith(prefixes)} + return get_conversion("UnCLIPTextProjModel", {}).to_diffusers(state) + + +# done text proj + +# super res unet first steps + + +def super_res_unet_first_steps_model_from_original_config(): + model = UNet2DModel(**SUPER_RES_UNET_FIRST_STEPS_CONFIG) + + return model + + +def super_res_unet_first_steps_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = { + key.removeprefix("model_first_steps."): value + for key, value in checkpoint.items() + if key.startswith("model_first_steps.") + } + return get_conversion("UNet2DModel", {**dict(model.config), "original_format": "ldm"}).to_diffusers(state) + + +# done super res unet first steps + +# super res unet last step + + +def super_res_unet_last_step_model_from_original_config(): + model = UNet2DModel(**SUPER_RES_UNET_LAST_STEP_CONFIG) + + return model + + +def super_res_unet_last_step_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = { + key.removeprefix("model_last_step."): value + for key, value in checkpoint.items() + if key.startswith("model_last_step.") + } + return get_conversion("UNet2DModel", {**dict(model.config), "original_format": "ldm"}).to_diffusers(state) + + +# done super res unet last step + + +# unet utils + + +# .time_embed -> .time_embedding + + +# .input_blocks.0 -> .conv_in + + +# .out.0 -> .conv_norm_out + + +# .out.2 -> .conv_out + + +# .input_blocks -> .down_blocks + + +# .middle_block -> .mid_block + + +# .output_blocks -> .up_blocks + + +# TODO maybe document and/or can do more efficiently (build indices in for loop and extract once for each split?) + + +# done unet utils + + +# Driver functions + + +def text_encoder(): + print("loading CLIP text encoder") + + clip_name = "openai/clip-vit-large-patch14" + + # sets pad_value to 0 + pad_token = "!" + + tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto") + + assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0 + + text_encoder_model = CLIPTextModelWithProjection.from_pretrained( + clip_name, + # `CLIPTextModel` does not support device_map="auto" + # device_map="auto" + ) + + print("done loading CLIP text encoder") + + return text_encoder_model, tokenizer_model + + +def prior(*, args, checkpoint_map_location): + print("loading prior") + + prior_checkpoint = torch.load(args.prior_checkpoint_path, map_location=checkpoint_map_location) + prior_checkpoint = prior_checkpoint["state_dict"] + + clip_stats_checkpoint = torch.load(args.clip_stat_path, map_location=checkpoint_map_location) + + prior_model = prior_model_from_original_config() + + prior_diffusers_checkpoint = prior_original_checkpoint_to_diffusers_checkpoint( + prior_model, prior_checkpoint, clip_stats_checkpoint + ) + + del prior_checkpoint + del clip_stats_checkpoint + + load_checkpoint_to_model(prior_diffusers_checkpoint, prior_model, strict=True) + + print("done loading prior") + + return prior_model + + +def decoder(*, args, checkpoint_map_location): + print("loading decoder") + + decoder_checkpoint = torch.load(args.decoder_checkpoint_path, map_location=checkpoint_map_location) + decoder_checkpoint = decoder_checkpoint["state_dict"] + + decoder_model = decoder_model_from_original_config() + + decoder_diffusers_checkpoint = decoder_original_checkpoint_to_diffusers_checkpoint( + decoder_model, decoder_checkpoint + ) + + # text proj interlude + + # The original decoder implementation includes a set of parameters that are used + # for creating the `encoder_hidden_states` which are what the U-net is conditioned + # on. The diffusers conditional unet directly takes the encoder_hidden_states. We pull + # the parameters into the UnCLIPTextProjModel class + text_proj_model = text_proj_from_original_config() + + text_proj_checkpoint = text_proj_original_checkpoint_to_diffusers_checkpoint(decoder_checkpoint) + + load_checkpoint_to_model(text_proj_checkpoint, text_proj_model, strict=True) + + # done text proj interlude + + del decoder_checkpoint + + load_checkpoint_to_model(decoder_diffusers_checkpoint, decoder_model, strict=True) + + print("done loading decoder") + + return decoder_model, text_proj_model + + +def super_res_unet(*, args, checkpoint_map_location): + print("loading super resolution unet") + + super_res_checkpoint = torch.load(args.super_res_unet_checkpoint_path, map_location=checkpoint_map_location) + super_res_checkpoint = super_res_checkpoint["state_dict"] + + # model_first_steps + + super_res_first_model = super_res_unet_first_steps_model_from_original_config() + + super_res_first_steps_checkpoint = super_res_unet_first_steps_original_checkpoint_to_diffusers_checkpoint( + super_res_first_model, super_res_checkpoint + ) + + # model_last_step + super_res_last_model = super_res_unet_last_step_model_from_original_config() + + super_res_last_step_checkpoint = super_res_unet_last_step_original_checkpoint_to_diffusers_checkpoint( + super_res_last_model, super_res_checkpoint + ) + + del super_res_checkpoint + + load_checkpoint_to_model(super_res_first_steps_checkpoint, super_res_first_model, strict=True) + + load_checkpoint_to_model(super_res_last_step_checkpoint, super_res_last_model, strict=True) + + print("done loading super resolution unet") + + return super_res_first_model, super_res_last_model + + +def load_checkpoint_to_model(checkpoint, model, strict=False): + with tempfile.NamedTemporaryFile() as file: + torch.save(checkpoint, file.name) + del checkpoint + if strict: + model.load_state_dict(torch.load(file.name), strict=True) + else: + load_checkpoint_and_dispatch(model, file.name, device_map="auto") + + +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( + "--prior_checkpoint_path", + default=None, + type=str, + required=True, + help="Path to the prior checkpoint to convert.", + ) + + parser.add_argument( + "--decoder_checkpoint_path", + default=None, + type=str, + required=True, + help="Path to the decoder checkpoint to convert.", + ) + + parser.add_argument( + "--super_res_unet_checkpoint_path", + default=None, + type=str, + required=True, + help="Path to the super resolution checkpoint to convert.", + ) + + parser.add_argument( + "--clip_stat_path", default=None, type=str, required=True, help="Path to the clip stats checkpoint to convert." + ) + + parser.add_argument( + "--checkpoint_load_device", + default="cpu", + type=str, + required=False, + help="The device passed to `map_location` when loading checkpoints.", + ) + + parser.add_argument( + "--debug", + default=None, + type=str, + required=False, + help="Only run a specific stage of the convert script. Used for debugging", + ) + + args = parser.parse_args() + + print(f"loading checkpoints to {args.checkpoint_load_device}") + + checkpoint_map_location = torch.device(args.checkpoint_load_device) + + if args.debug is not None: + print(f"debug: only executing {args.debug}") + + if args.debug is None: + text_encoder_model, tokenizer_model = text_encoder() + + prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) + + decoder_model, text_proj_model = decoder(args=args, checkpoint_map_location=checkpoint_map_location) + + super_res_first_model, super_res_last_model = super_res_unet( + args=args, checkpoint_map_location=checkpoint_map_location + ) + + prior_scheduler = UnCLIPScheduler( + variance_type="fixed_small_log", + prediction_type="sample", + num_train_timesteps=1000, + clip_sample_range=5.0, + ) + + decoder_scheduler = UnCLIPScheduler( + variance_type="learned_range", + prediction_type="epsilon", + num_train_timesteps=1000, + ) + + super_res_scheduler = UnCLIPScheduler( + variance_type="fixed_small_log", + prediction_type="epsilon", + num_train_timesteps=1000, + ) + + print(f"saving Kakao Brain unCLIP to {args.dump_path}") + + pipe = UnCLIPPipeline( + prior=prior_model, + decoder=decoder_model, + text_proj=text_proj_model, + tokenizer=tokenizer_model, + text_encoder=text_encoder_model, + super_res_first=super_res_first_model, + super_res_last=super_res_last_model, + prior_scheduler=prior_scheduler, + decoder_scheduler=decoder_scheduler, + super_res_scheduler=super_res_scheduler, + ) + pipe.save_pretrained(args.dump_path) + + print("done writing Kakao Brain unCLIP") + elif args.debug == "text_encoder": + text_encoder_model, tokenizer_model = text_encoder() + elif args.debug == "prior": + prior_model = prior(args=args, checkpoint_map_location=checkpoint_map_location) + elif args.debug == "decoder": + decoder_model, text_proj_model = decoder(args=args, checkpoint_map_location=checkpoint_map_location) + elif args.debug == "super_res_unet": + super_res_first_model, super_res_last_model = super_res_unet( + args=args, checkpoint_map_location=checkpoint_map_location + ) + else: + raise ValueError(f"unknown debug value : {args.debug}") diff --git a/scripts/convert_unclip_txt2img_to_image_variation.py b/scripts/recipes/unclip_image_variation.py similarity index 100% rename from scripts/convert_unclip_txt2img_to_image_variation.py rename to scripts/recipes/unclip_image_variation.py diff --git a/scripts/recipes/unidiffuser.py b/scripts/recipes/unidiffuser.py new file mode 100644 index 000000000000..f12cd773e8c2 --- /dev/null +++ b/scripts/recipes/unidiffuser.py @@ -0,0 +1,220 @@ +# Convert the original UniDiffuser checkpoints into diffusers equivalents. + +import argparse + +import torch +from transformers import ( + CLIPImageProcessor, + CLIPTextConfig, + CLIPTextModel, + CLIPTokenizer, + CLIPVisionConfig, + CLIPVisionModelWithProjection, + GPT2Tokenizer, +) + +from diffusers import ( + AutoencoderKL, + DPMSolverMultistepScheduler, + UniDiffuserModel, + UniDiffuserPipeline, + UniDiffuserTextDecoder, +) +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.unidiffuser import ( + SCHEDULER_CONFIG, + create_text_decoder_config, + create_unidiffuser_unet_config, + create_vae_diffusers_config, +) + + +# Modified from diffusers.pipelines.stable_diffusion.convert_from_ckpt.assign_to_checkpoint +# config.num_head_channels => num_head_channels + + +# Hardcoded configs for test versions of the UniDiffuser models, corresponding to those in the fast default tests. + + +# Hardcoded configs for the UniDiffuser V1 model at https://huggingface.co/thu-ml/unidiffuser-v1 +# See also https://github.com/thu-ml/unidiffuser/blob/main/configs/sample_unidiffuser_v1.py + + +# From https://huggingface.co/gpt2/blob/main/config.json, the GPT2 checkpoint used by UniDiffuser + + +# Based on diffusers.pipelines.stable_diffusion.convert_from_ckpt.convert_ldm_vae_checkpoint +def convert_vae_to_diffusers(ckpt, diffusers_model, **kwargs): + state = torch.load(ckpt, map_location="cpu", weights_only=True) + if any(key.startswith("module.") for key in state): + state = {key.removeprefix("module."): value for key, value in state.items()} + converted = convert_component_checkpoint(state, dict(diffusers_model.config), "AutoencoderKL") + diffusers_model.load_state_dict(converted, strict=True) + return diffusers_model + + +def convert_uvit_to_diffusers(ckpt, diffusers_model, **kwargs): + state = torch.load(ckpt, map_location="cpu", weights_only=True) + if any(key.startswith("module.") for key in state): + state = {key.removeprefix("module."): value for key, value in state.items()} + converted = convert_component_checkpoint(state, dict(diffusers_model.config), "UniDiffuserModel") + diffusers_model.load_state_dict(converted, strict=True) + return diffusers_model + + +def convert_caption_decoder_to_diffusers(ckpt, diffusers_model, **kwargs): + state = torch.load(ckpt, map_location="cpu", weights_only=True) + if any(key.startswith("module.") for key in state): + state = {key.removeprefix("module."): value for key, value in state.items()} + converted = convert_component_checkpoint(state, dict(diffusers_model.config), "UniDiffuserTextDecoder") + diffusers_model.load_state_dict(converted, strict=True) + return diffusers_model + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--caption_decoder_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to caption decoder checkpoint to convert.", + ) + parser.add_argument( + "--uvit_checkpoint_path", default=None, type=str, required=False, help="Path to U-ViT checkpoint to convert." + ) + parser.add_argument( + "--vae_checkpoint_path", + default=None, + type=str, + required=False, + help="Path to VAE checkpoint to convert.", + ) + parser.add_argument( + "--pipeline_output_path", + default=None, + type=str, + required=True, + help="Path to save the output pipeline to.", + ) + parser.add_argument( + "--config_type", + default="test", + type=str, + help=( + "Config type to use. Should be 'test' to create small models for testing or 'big' to convert a full" + " checkpoint." + ), + ) + parser.add_argument( + "--version", + default=0, + type=int, + help="The UniDiffuser model type to convert to. Should be 0 for UniDiffuser-v0 and 1 for UniDiffuser-v1.", + ) + parser.add_argument( + "--safe_serialization", + action="store_true", + help="Whether to use safetensors/safe seialization when saving the pipeline.", + ) + + args = parser.parse_args() + + # Convert the VAE model. + if args.vae_checkpoint_path is not None: + vae_config = create_vae_diffusers_config(args.config_type) + vae = AutoencoderKL(**vae_config) + vae = convert_vae_to_diffusers(args.vae_checkpoint_path, vae) + + # Convert the U-ViT ("unet") model. + if args.uvit_checkpoint_path is not None: + unet_config = create_unidiffuser_unet_config(args.config_type, args.version) + unet = UniDiffuserModel(**unet_config) + unet = convert_uvit_to_diffusers(args.uvit_checkpoint_path, unet) + + # Convert the caption decoder ("text_decoder") model. + if args.caption_decoder_checkpoint_path is not None: + text_decoder_config = create_text_decoder_config(args.config_type) + text_decoder = UniDiffuserTextDecoder(**text_decoder_config) + text_decoder = convert_caption_decoder_to_diffusers(args.caption_decoder_checkpoint_path, text_decoder) + + # Scheduler is the same for both the test and big models. + scheduler_config = SCHEDULER_CONFIG + scheduler = DPMSolverMultistepScheduler( + beta_start=scheduler_config.beta_start, + beta_end=scheduler_config.beta_end, + beta_schedule=scheduler_config.beta_schedule, + solver_order=scheduler_config.solver_order, + ) + + if args.config_type == "test": + # Make a small random CLIPTextModel + torch.manual_seed(0) + clip_text_encoder_config = CLIPTextConfig( + bos_token_id=0, + eos_token_id=2, + hidden_size=32, + intermediate_size=37, + layer_norm_eps=1e-05, + num_attention_heads=4, + num_hidden_layers=5, + pad_token_id=1, + vocab_size=1000, + ) + text_encoder = CLIPTextModel(clip_text_encoder_config) + clip_tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip") + + # Make a small random CLIPVisionModel and accompanying CLIPImageProcessor + torch.manual_seed(0) + clip_image_encoder_config = CLIPVisionConfig( + image_size=32, + patch_size=2, + num_channels=3, + hidden_size=32, + projection_dim=32, + num_hidden_layers=5, + num_attention_heads=4, + intermediate_size=37, + dropout=0.1, + attention_dropout=0.1, + initializer_range=0.02, + ) + image_encoder = CLIPVisionModelWithProjection(clip_image_encoder_config) + image_processor = CLIPImageProcessor(crop_size=32, size=32) + + # Note that the text_decoder should already have its token embeddings resized. + text_tokenizer = GPT2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-GPT2Model") + eos = "<|EOS|>" + special_tokens_dict = {"eos_token": eos} + text_tokenizer.add_special_tokens(special_tokens_dict) + elif args.config_type == "big": + text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14") + clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") + + image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-base-patch32") + image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32") + + # Note that the text_decoder should already have its token embeddings resized. + text_tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + eos = "<|EOS|>" + special_tokens_dict = {"eos_token": eos} + text_tokenizer.add_special_tokens(special_tokens_dict) + else: + raise NotImplementedError( + f"Config type {args.config_type} is not implemented, currently only config types" + " 'test' and 'big' are available." + ) + + pipeline = UniDiffuserPipeline( + vae=vae, + text_encoder=text_encoder, + image_encoder=image_encoder, + clip_image_processor=image_processor, + clip_tokenizer=clip_tokenizer, + text_decoder=text_decoder, + text_tokenizer=text_tokenizer, + unet=unet, + scheduler=scheduler, + ) + pipeline.save_pretrained(args.pipeline_output_path, safe_serialization=args.safe_serialization) diff --git a/scripts/recipes/versatile_diffusion.py b/scripts/recipes/versatile_diffusion.py new file mode 100644 index 000000000000..7316bac939ee --- /dev/null +++ b/scripts/recipes/versatile_diffusion.py @@ -0,0 +1,183 @@ +# coding=utf-8 +# Copyright 2025 The HuggingFace Inc. team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Conversion script for the Versatile Stable Diffusion checkpoints.""" + +import argparse + +import torch +from transformers import ( + CLIPImageProcessor, + CLIPTextModelWithProjection, + CLIPTokenizer, + CLIPVisionModelWithProjection, +) + +from diffusers import ( + AutoencoderKL, + DDIMScheduler, + DPMSolverMultistepScheduler, + EulerAncestralDiscreteScheduler, + EulerDiscreteScheduler, + LMSDiscreteScheduler, + PNDMScheduler, + UNet2DConditionModel, + VersatileDiffusionPipeline, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs.versatile_diffusion import ( + AUTOENCODER_CONFIG, + IMAGE_UNET_CONFIG, + SCHEDULER_CONFIG, + TEXT_UNET_CONFIG, + create_image_unet_diffusers_config, + create_text_unet_diffusers_config, + create_vae_diffusers_config, +) +from diffusers.pipelines.deprecated.versatile_diffusion.modeling_text_unet import UNetFlatConditionModel + + +def convert_vd_unet_checkpoint(checkpoint, config, unet_key, extract_ema=False): + state = { + key: value + for key, value in checkpoint.items() + if key.startswith((unet_key, "model.diffusion_model.time_embed.")) + } + if extract_ema: + state = {key: checkpoint["model_ema." + "".join(key.split(".")[1:])] for key in state} + cls = "UNetFlatConditionModel" if "unet_text." in unet_key else "UNet2DConditionModel" + config = dict(config) + if cls == "UNet2DConditionModel": + config["original_format"] = "versatile_image" + return get_conversion(cls, config).to_diffusers(state) + + +def convert_vd_vae_checkpoint(checkpoint, config): + state = {key.removeprefix("first_stage_model."): value for key, value in checkpoint.items()} + return get_conversion("AutoencoderKL", config).to_diffusers(state) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--unet_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--vae_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--optimus_checkpoint_path", default=None, type=str, required=False, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--scheduler_type", + default="pndm", + type=str, + help="Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']", + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + + args = parser.parse_args() + + scheduler_config = SCHEDULER_CONFIG + + num_train_timesteps = scheduler_config.timesteps + beta_start = scheduler_config.beta_linear_start + beta_end = scheduler_config.beta_linear_end + if args.scheduler_type == "pndm": + scheduler = PNDMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + skip_prk_steps=True, + steps_offset=1, + ) + elif args.scheduler_type == "lms": + scheduler = LMSDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear") + elif args.scheduler_type == "euler": + scheduler = EulerDiscreteScheduler(beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear") + elif args.scheduler_type == "euler-ancestral": + scheduler = EulerAncestralDiscreteScheduler( + beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear" + ) + elif args.scheduler_type == "dpm": + scheduler = DPMSolverMultistepScheduler( + beta_start=beta_start, beta_end=beta_end, beta_schedule="scaled_linear" + ) + elif args.scheduler_type == "ddim": + scheduler = DDIMScheduler( + beta_start=beta_start, + beta_end=beta_end, + beta_schedule="scaled_linear", + clip_sample=False, + set_alpha_to_one=False, + steps_offset=1, + ) + else: + raise ValueError(f"Scheduler of type {args.scheduler_type} doesn't exist!") + + # Convert the UNet2DConditionModel models. + if args.unet_checkpoint_path is not None: + # image UNet + image_unet_config = create_image_unet_diffusers_config(IMAGE_UNET_CONFIG) + checkpoint = torch.load(args.unet_checkpoint_path) + converted_image_unet_checkpoint = convert_vd_unet_checkpoint( + checkpoint, image_unet_config, unet_key="model.diffusion_model.unet_image.", extract_ema=args.extract_ema + ) + image_unet = UNet2DConditionModel(**image_unet_config) + image_unet.load_state_dict(converted_image_unet_checkpoint) + + # text UNet + text_unet_config = create_text_unet_diffusers_config(TEXT_UNET_CONFIG) + converted_text_unet_checkpoint = convert_vd_unet_checkpoint( + checkpoint, text_unet_config, unet_key="model.diffusion_model.unet_text.", extract_ema=args.extract_ema + ) + text_unet = UNetFlatConditionModel(**text_unet_config) + text_unet.load_state_dict(converted_text_unet_checkpoint) + + # Convert the VAE model. + if args.vae_checkpoint_path is not None: + vae_config = create_vae_diffusers_config(AUTOENCODER_CONFIG) + checkpoint = torch.load(args.vae_checkpoint_path) + converted_vae_checkpoint = convert_vd_vae_checkpoint(checkpoint, vae_config) + + vae = AutoencoderKL(**vae_config) + vae.load_state_dict(converted_vae_checkpoint) + + tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14") + image_feature_extractor = CLIPImageProcessor.from_pretrained("openai/clip-vit-large-patch14") + text_encoder = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") + image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14") + + pipe = VersatileDiffusionPipeline( + scheduler=scheduler, + tokenizer=tokenizer, + image_feature_extractor=image_feature_extractor, + text_encoder=text_encoder, + image_encoder=image_encoder, + image_unet=image_unet, + text_unet=text_unet, + vae=vae, + ) + pipe.save_pretrained(args.dump_path) diff --git a/scripts/recipes/vq_diffusion.py b/scripts/recipes/vq_diffusion.py new file mode 100644 index 000000000000..7ce294499867 --- /dev/null +++ b/scripts/recipes/vq_diffusion.py @@ -0,0 +1,487 @@ +""" +This script ports models from VQ-diffusion (https://github.com/microsoft/VQ-Diffusion) to diffusers. + +It currently only supports porting the ITHQ dataset. + +ITHQ dataset: +```sh +# From the root directory of diffusers. + +# Download the VQVAE checkpoint +$ Refer to https://github.com/microsoft/VQ-Diffusion/blob/main/vqdiffusion_download_checkpoints.sh + +# Download the VQVAE config +# NOTE that in VQ-diffusion the documented file is `configs/ithq.yaml` but the target class +# `image_synthesis.modeling.codecs.image_codec.ema_vqvae.PatchVQVAE` +# loads `OUTPUT/pretrained_model/taming_dvae/config.yaml` +$ wget https://raw.githubusercontent.com/microsoft/VQ-Diffusion/main/OUTPUT/pretrained_model/taming_dvae/config.yaml -O ithq_vqvae.yaml + +# Download the main model checkpoint +$ wget https://facevcstandard.blob.core.windows.net/v-zhictang/Improved-VQ-Diffusion_model_release/ithq_learnable.pth?sv=2020-10-02&st=2022-05-30T10%3A22%3A06Z&se=2030-05-31T10%3A22%3A00Z&sr=b&sp=r&sig=GOE%2Bza02%2FPnGxYVOOPtwrTR4RA3%2F5NVgMxdW4kjaEZ8%3D -O ithq_learnable.pth # trufflehog:ignore + +# Download the main model config +$ wget https://raw.githubusercontent.com/microsoft/VQ-Diffusion/main/configs/ithq.yaml -O ithq.yaml + +# run the convert script +$ python ./scripts/recipes/vq_diffusion.py \ + --checkpoint_path ./ithq_learnable.pth \ + --original_config_file ./ithq.yaml \ + --vqvae_checkpoint_path ./ithq_vqvae.pth \ + --vqvae_original_config_file ./ithq_vqvae.yaml \ + --dump_path +``` +""" + +import argparse +import tempfile + +import torch +import yaml +from accelerate import init_empty_weights, load_checkpoint_and_dispatch +from transformers import CLIPTextModel, CLIPTokenizer +from yaml.loader import FullLoader + +from diffusers import Transformer2DModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel +from diffusers.loaders.conversion import get_conversion +from diffusers.pipelines.deprecated.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings + + +# vqvae model + +PORTED_VQVAES = ["image_synthesis.modeling.codecs.image_codec.patch_vqgan.PatchVQGAN"] + + +def vqvae_model_from_original_config(original_config): + assert original_config["target"] in PORTED_VQVAES, ( + f"{original_config['target']} has not yet been ported to diffusers." + ) + + original_config = original_config["params"] + + original_encoder_config = original_config["encoder_config"]["params"] + original_decoder_config = original_config["decoder_config"]["params"] + + in_channels = original_encoder_config["in_channels"] + out_channels = original_decoder_config["out_ch"] + + down_block_types = get_down_block_types(original_encoder_config) + up_block_types = get_up_block_types(original_decoder_config) + + assert original_encoder_config["ch"] == original_decoder_config["ch"] + assert original_encoder_config["ch_mult"] == original_decoder_config["ch_mult"] + block_out_channels = tuple( + [original_encoder_config["ch"] * a_ch_mult for a_ch_mult in original_encoder_config["ch_mult"]] + ) + + assert original_encoder_config["num_res_blocks"] == original_decoder_config["num_res_blocks"] + layers_per_block = original_encoder_config["num_res_blocks"] + + assert original_encoder_config["z_channels"] == original_decoder_config["z_channels"] + latent_channels = original_encoder_config["z_channels"] + + num_vq_embeddings = original_config["n_embed"] + + # Hard coded value for ResnetBlock.GoupNorm(num_groups) in VQ-diffusion + norm_num_groups = 32 + + e_dim = original_config["embed_dim"] + + model = VQModel( + in_channels=in_channels, + out_channels=out_channels, + down_block_types=down_block_types, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + latent_channels=latent_channels, + num_vq_embeddings=num_vq_embeddings, + norm_num_groups=norm_num_groups, + vq_embed_dim=e_dim, + ) + + return model + + +def get_down_block_types(original_encoder_config): + attn_resolutions = coerce_attn_resolutions(original_encoder_config["attn_resolutions"]) + num_resolutions = len(original_encoder_config["ch_mult"]) + resolution = coerce_resolution(original_encoder_config["resolution"]) + + curr_res = resolution + down_block_types = [] + + for _ in range(num_resolutions): + if curr_res in attn_resolutions: + down_block_type = "AttnDownEncoderBlock2D" + else: + down_block_type = "DownEncoderBlock2D" + + down_block_types.append(down_block_type) + + curr_res = [r // 2 for r in curr_res] + + return down_block_types + + +def get_up_block_types(original_decoder_config): + attn_resolutions = coerce_attn_resolutions(original_decoder_config["attn_resolutions"]) + num_resolutions = len(original_decoder_config["ch_mult"]) + resolution = coerce_resolution(original_decoder_config["resolution"]) + + curr_res = [r // 2 ** (num_resolutions - 1) for r in resolution] + up_block_types = [] + + for _ in reversed(range(num_resolutions)): + if curr_res in attn_resolutions: + up_block_type = "AttnUpDecoderBlock2D" + else: + up_block_type = "UpDecoderBlock2D" + + up_block_types.append(up_block_type) + + curr_res = [r * 2 for r in curr_res] + + return up_block_types + + +def coerce_attn_resolutions(attn_resolutions): + attn_resolutions = list(attn_resolutions) + attn_resolutions_ = [] + for ar in attn_resolutions: + if isinstance(ar, (list, tuple)): + attn_resolutions_.append(list(ar)) + else: + attn_resolutions_.append([ar, ar]) + return attn_resolutions_ + + +def coerce_resolution(resolution): + if isinstance(resolution, int): + resolution = [resolution, resolution] # H, W + elif isinstance(resolution, (tuple, list)): + resolution = list(resolution) + else: + raise ValueError("Unknown type of resolution:", resolution) + return resolution + + +# done vqvae model + +# vqvae checkpoint + + +def vqvae_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + return get_conversion("VQModel", dict(model.config)).to_diffusers(checkpoint) + + +# done vqvae checkpoint + +# transformer model + +PORTED_DIFFUSIONS = ["image_synthesis.modeling.transformers.diffusion_transformer.DiffusionTransformer"] +PORTED_TRANSFORMERS = ["image_synthesis.modeling.transformers.transformer_utils.Text2ImageTransformer"] +PORTED_CONTENT_EMBEDDINGS = ["image_synthesis.modeling.embeddings.dalle_mask_image_embedding.DalleMaskImageEmbedding"] + + +def transformer_model_from_original_config( + original_diffusion_config, original_transformer_config, original_content_embedding_config +): + assert original_diffusion_config["target"] in PORTED_DIFFUSIONS, ( + f"{original_diffusion_config['target']} has not yet been ported to diffusers." + ) + assert original_transformer_config["target"] in PORTED_TRANSFORMERS, ( + f"{original_transformer_config['target']} has not yet been ported to diffusers." + ) + assert original_content_embedding_config["target"] in PORTED_CONTENT_EMBEDDINGS, ( + f"{original_content_embedding_config['target']} has not yet been ported to diffusers." + ) + + original_diffusion_config = original_diffusion_config["params"] + original_transformer_config = original_transformer_config["params"] + original_content_embedding_config = original_content_embedding_config["params"] + + inner_dim = original_transformer_config["n_embd"] + + n_heads = original_transformer_config["n_head"] + + # VQ-Diffusion gives dimension of the multi-headed attention layers as the + # number of attention heads times the sequence length (the dimension) of a + # single head. We want to specify our attention blocks with those values + # specified separately + assert inner_dim % n_heads == 0 + d_head = inner_dim // n_heads + + depth = original_transformer_config["n_layer"] + context_dim = original_transformer_config["condition_dim"] + + num_embed = original_content_embedding_config["num_embed"] + # the number of embeddings in the transformer includes the mask embedding. + # the content embedding (the vqvae) does not include the mask embedding. + num_embed = num_embed + 1 + + height = original_transformer_config["content_spatial_size"][0] + width = original_transformer_config["content_spatial_size"][1] + + assert width == height, "width has to be equal to height" + dropout = original_transformer_config["resid_pdrop"] + num_embeds_ada_norm = original_diffusion_config["diffusion_step"] + + model_kwargs = { + "attention_bias": True, + "cross_attention_dim": context_dim, + "attention_head_dim": d_head, + "num_layers": depth, + "dropout": dropout, + "num_attention_heads": n_heads, + "num_vector_embeds": num_embed, + "num_embeds_ada_norm": num_embeds_ada_norm, + "norm_num_groups": 32, + "sample_size": width, + "activation_fn": "geglu-approximate", + } + + model = Transformer2DModel(**model_kwargs) + return model + + +# done transformer model + +# transformer checkpoint + + +def transformer_original_checkpoint_to_diffusers_checkpoint(model, checkpoint): + state = {key: value for key, value in checkpoint.items() if key.startswith("transformer.transformer.")} + return get_conversion("Transformer2DModel", dict(model.config)).to_diffusers(state) + + +# done transformer checkpoint + + +def read_config_file(filename): + # The yaml file contains annotations that certain values should + # loaded as tuples. + with open(filename) as f: + original_config = yaml.load(f, FullLoader) + + return original_config + + +# We take separate arguments for the vqvae because the ITHQ vqvae config file +# is separate from the config file for the rest of the model. +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--vqvae_checkpoint_path", + default=None, + type=str, + required=True, + help="Path to the vqvae checkpoint to convert.", + ) + + parser.add_argument( + "--vqvae_original_config_file", + default=None, + type=str, + required=True, + help="The YAML config file corresponding to the original architecture for the vqvae.", + ) + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + + parser.add_argument( + "--original_config_file", + default=None, + type=str, + required=True, + help="The YAML config file corresponding to the original architecture.", + ) + + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + + parser.add_argument( + "--checkpoint_load_device", + default="cpu", + type=str, + required=False, + help="The device passed to `map_location` when loading checkpoints.", + ) + + # See link for how ema weights are always selected + # https://github.com/microsoft/VQ-Diffusion/blob/3c98e77f721db7c787b76304fa2c96a36c7b00af/inference_VQ_Diffusion.py#L65 + parser.add_argument( + "--no_use_ema", + action="store_true", + required=False, + help=( + "Set to not use the ema weights from the original VQ-Diffusion checkpoint. You probably do not want to set" + " it as the original VQ-Diffusion always uses the ema weights when loading models." + ), + ) + + args = parser.parse_args() + + use_ema = not args.no_use_ema + + print(f"loading checkpoints to {args.checkpoint_load_device}") + + checkpoint_map_location = torch.device(args.checkpoint_load_device) + + # vqvae_model + + print(f"loading vqvae, config: {args.vqvae_original_config_file}, checkpoint: {args.vqvae_checkpoint_path}") + + vqvae_original_config = read_config_file(args.vqvae_original_config_file).model + vqvae_checkpoint = torch.load(args.vqvae_checkpoint_path, map_location=checkpoint_map_location)["model"] + + with init_empty_weights(): + vqvae_model = vqvae_model_from_original_config(vqvae_original_config) + + vqvae_diffusers_checkpoint = vqvae_original_checkpoint_to_diffusers_checkpoint(vqvae_model, vqvae_checkpoint) + + with tempfile.NamedTemporaryFile() as vqvae_diffusers_checkpoint_file: + torch.save(vqvae_diffusers_checkpoint, vqvae_diffusers_checkpoint_file.name) + del vqvae_diffusers_checkpoint + del vqvae_checkpoint + load_checkpoint_and_dispatch(vqvae_model, vqvae_diffusers_checkpoint_file.name, device_map="auto") + + print("done loading vqvae") + + # done vqvae_model + + # transformer_model + + print( + f"loading transformer, config: {args.original_config_file}, checkpoint: {args.checkpoint_path}, use ema:" + f" {use_ema}" + ) + + original_config = read_config_file(args.original_config_file).model + + diffusion_config = original_config["params"]["diffusion_config"] + transformer_config = original_config["params"]["diffusion_config"]["params"]["transformer_config"] + content_embedding_config = original_config["params"]["diffusion_config"]["params"]["content_emb_config"] + + pre_checkpoint = torch.load(args.checkpoint_path, map_location=checkpoint_map_location) + + if use_ema: + if "ema" in pre_checkpoint: + checkpoint = {} + for k, v in pre_checkpoint["model"].items(): + checkpoint[k] = v + + for k, v in pre_checkpoint["ema"].items(): + # The ema weights are only used on the transformer. To mimic their key as if they came + # from the state_dict for the top level model, we prefix with an additional "transformer." + # See the source linked in the args.use_ema config for more information. + checkpoint[f"transformer.{k}"] = v + else: + print("attempted to load ema weights but no ema weights are specified in the loaded checkpoint.") + checkpoint = pre_checkpoint["model"] + else: + checkpoint = pre_checkpoint["model"] + + del pre_checkpoint + + with init_empty_weights(): + transformer_model = transformer_model_from_original_config( + diffusion_config, transformer_config, content_embedding_config + ) + + diffusers_transformer_checkpoint = transformer_original_checkpoint_to_diffusers_checkpoint( + transformer_model, checkpoint + ) + + # classifier free sampling embeddings interlude + + # The learned embeddings are stored on the transformer in the original VQ-diffusion. We store them on a separate + # model, so we pull them off the checkpoint before the checkpoint is deleted. + + learnable_classifier_free_sampling_embeddings = diffusion_config["params"]["learnable_cf"] + + if learnable_classifier_free_sampling_embeddings: + learned_classifier_free_sampling_embeddings_embeddings = checkpoint["transformer.empty_text_embed"] + else: + learned_classifier_free_sampling_embeddings_embeddings = None + + # done classifier free sampling embeddings interlude + + with tempfile.NamedTemporaryFile() as diffusers_transformer_checkpoint_file: + torch.save(diffusers_transformer_checkpoint, diffusers_transformer_checkpoint_file.name) + del diffusers_transformer_checkpoint + del checkpoint + load_checkpoint_and_dispatch(transformer_model, diffusers_transformer_checkpoint_file.name, device_map="auto") + + print("done loading transformer") + + # done transformer_model + + # text encoder + + print("loading CLIP text encoder") + + clip_name = "openai/clip-vit-base-patch32" + + # The original VQ-Diffusion specifies the pad value by the int used in the + # returned tokens. Each model uses `0` as the pad value. The transformers clip api + # specifies the pad value via the token before it has been tokenized. The `!` pad + # token is the same as padding with the `0` pad value. + pad_token = "!" + + tokenizer_model = CLIPTokenizer.from_pretrained(clip_name, pad_token=pad_token, device_map="auto") + + assert tokenizer_model.convert_tokens_to_ids(pad_token) == 0 + + text_encoder_model = CLIPTextModel.from_pretrained( + clip_name, + # `CLIPTextModel` does not support device_map="auto" + # device_map="auto" + ) + + print("done loading CLIP text encoder") + + # done text encoder + + # scheduler + + scheduler_model = VQDiffusionScheduler( + # the scheduler has the same number of embeddings as the transformer + num_vec_classes=transformer_model.num_vector_embeds + ) + + # done scheduler + + # learned classifier free sampling embeddings + + with init_empty_weights(): + learned_classifier_free_sampling_embeddings_model = LearnedClassifierFreeSamplingEmbeddings( + learnable_classifier_free_sampling_embeddings, + hidden_size=text_encoder_model.config.hidden_size, + length=tokenizer_model.model_max_length, + ) + + source_embeddings = learned_classifier_free_sampling_embeddings_embeddings + if source_embeddings is None: + source_embeddings = torch.empty(0) + converted = get_conversion( + "LearnedClassifierFreeSamplingEmbeddings", dict(learned_classifier_free_sampling_embeddings_model.config) + ).to_diffusers({"transformer.empty_text_embed": source_embeddings}) + learned_classifier_free_sampling_embeddings_model.load_state_dict(converted, strict=True, assign=True) + + # done learned classifier free sampling embeddings + + print(f"saving VQ diffusion model, path: {args.dump_path}") + + pipe = VQDiffusionPipeline( + vqvae=vqvae_model, + transformer=transformer_model, + tokenizer=tokenizer_model, + text_encoder=text_encoder_model, + learned_classifier_free_sampling_embeddings=learned_classifier_free_sampling_embeddings_model, + scheduler=scheduler_model, + ) + pipe.save_pretrained(args.dump_path) + + print("done writing VQ diffusion model") diff --git a/scripts/recipes/wan.py b/scripts/recipes/wan.py new file mode 100644 index 000000000000..515cb61d81b2 --- /dev/null +++ b/scripts/recipes/wan.py @@ -0,0 +1,310 @@ +import argparse +import pathlib + +import torch +from huggingface_hub import hf_hub_download, snapshot_download +from transformers import ( + AutoProcessor, + AutoTokenizer, + CLIPImageProcessor, + CLIPVisionModel, + CLIPVisionModelWithProjection, + UMT5EncoderModel, +) + +from diffusers import ( + AutoencoderKLWan, + UniPCMultistepScheduler, + WanAnimatePipeline, + WanAnimateTransformer3DModel, + WanImageToVideoPipeline, + WanPipeline, + WanTransformer3DModel, + WanVACEPipeline, + WanVACETransformer3DModel, +) +from diffusers.loaders.conversion.configs.wan import get_transformer_config + + +# TODO: Verify this and simplify if possible. + + +def load_sharded_safetensors(path): + from diffusers.loaders.conversion.source import load_tensor_sources + + return load_tensor_sources(path) + + +def convert_transformer(model_type: str, stage: str = None): + spec = get_transformer_config(model_type) + model_dir = pathlib.Path(snapshot_download(spec["model_id"], repo_type="model")) + if stage is not None: + model_dir = model_dir / stage + state = load_sharded_safetensors(model_dir) + cls = ( + WanAnimateTransformer3DModel + if "Animate" in model_type + else (WanVACETransformer3DModel if "VACE" in model_type else WanTransformer3DModel) + ) + return cls.from_single_file(state, config=spec["diffusers_config"]) + + +def convert_vae(): + path = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth") + return AutoencoderKLWan.from_single_file(path, config={}) + + +vae22_diffusers_config = { + "base_dim": 160, + "z_dim": 48, + "is_residual": True, + "in_channels": 12, + "out_channels": 12, + "decoder_base_dim": 256, + "scale_factor_temporal": 4, + "scale_factor_spatial": 16, + "patch_size": 2, + "latents_mean": [ + -0.2289, + -0.0052, + -0.1323, + -0.2339, + -0.2799, + 0.0174, + 0.1838, + 0.1557, + -0.1382, + 0.0542, + 0.2813, + 0.0891, + 0.1570, + -0.0098, + 0.0375, + -0.1825, + -0.2246, + -0.1207, + -0.0698, + 0.5109, + 0.2665, + -0.2108, + -0.2158, + 0.2502, + -0.2055, + -0.0322, + 0.1109, + 0.1567, + -0.0729, + 0.0899, + -0.2799, + -0.1230, + -0.0313, + -0.1649, + 0.0117, + 0.0723, + -0.2839, + -0.2083, + -0.0520, + 0.3748, + 0.0152, + 0.1957, + 0.1433, + -0.2944, + 0.3573, + -0.0548, + -0.1681, + -0.0667, + ], + "latents_std": [ + 0.4765, + 1.0364, + 0.4514, + 1.1677, + 0.5313, + 0.4990, + 0.4818, + 0.5013, + 0.8158, + 1.0344, + 0.5894, + 1.0901, + 0.6885, + 0.6165, + 0.8454, + 0.4978, + 0.5759, + 0.3523, + 0.7135, + 0.6804, + 0.5833, + 1.4146, + 0.8986, + 0.5659, + 0.7069, + 0.5338, + 0.4889, + 0.4917, + 0.4069, + 0.4999, + 0.6866, + 0.4093, + 0.5709, + 0.6065, + 0.6415, + 0.4944, + 0.5726, + 1.2042, + 0.5458, + 1.6887, + 0.3971, + 1.0600, + 0.3943, + 0.5537, + 0.5444, + 0.4089, + 0.7468, + 0.7744, + ], + "clip_output": False, +} + + +def convert_vae_22(): + path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth") + return AutoencoderKLWan.from_single_file(path, config=vae22_diffusers_config) + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_type", type=str, default=None) + parser.add_argument("--output_path", type=str, required=True) + parser.add_argument("--dtype", default="fp32", choices=["fp32", "fp16", "bf16", "none"]) + return parser.parse_args() + + +DTYPE_MAPPING = { + "fp32": torch.float32, + "fp16": torch.float16, + "bf16": torch.bfloat16, +} + + +if __name__ == "__main__": + args = get_args() + + if "Wan2.2" in args.model_type and "TI2V" not in args.model_type and "Animate" not in args.model_type: + transformer = convert_transformer(args.model_type, stage="high_noise_model") + transformer_2 = convert_transformer(args.model_type, stage="low_noise_model") + else: + transformer = convert_transformer(args.model_type) + transformer_2 = None + + if "Wan2.2" in args.model_type and "TI2V" in args.model_type: + vae = convert_vae_22() + else: + vae = convert_vae() + + text_encoder = UMT5EncoderModel.from_pretrained("google/umt5-xxl", torch_dtype=torch.bfloat16) + tokenizer = AutoTokenizer.from_pretrained("google/umt5-xxl") + if "FLF2V" in args.model_type: + flow_shift = 16.0 + elif "TI2V" in args.model_type or "Animate" in args.model_type: + flow_shift = 5.0 + else: + flow_shift = 3.0 + scheduler = UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=flow_shift + ) + + # If user has specified "none", we keep the original dtypes of the state dict without any conversion + if args.dtype != "none": + dtype = DTYPE_MAPPING[args.dtype] + transformer.to(dtype) + if transformer_2 is not None: + transformer_2.to(dtype) + + if "Wan2.2" and "I2V" in args.model_type and "TI2V" not in args.model_type: + pipe = WanImageToVideoPipeline( + transformer=transformer, + transformer_2=transformer_2, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + boundary_ratio=0.9, + ) + elif "Wan2.2" and "T2V" in args.model_type: + pipe = WanPipeline( + transformer=transformer, + transformer_2=transformer_2, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + boundary_ratio=0.875, + ) + elif "Wan2.2" and "TI2V" in args.model_type: + pipe = WanPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + expand_timesteps=True, + ) + elif "I2V" in args.model_type or "FLF2V" in args.model_type: + image_encoder = CLIPVisionModelWithProjection.from_pretrained( + "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=torch.bfloat16 + ) + image_processor = AutoProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") + pipe = WanImageToVideoPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + image_encoder=image_encoder, + image_processor=image_processor, + ) + elif "Wan2.2-VACE" in args.model_type: + pipe = WanVACEPipeline( + transformer=transformer, + transformer_2=transformer_2, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + boundary_ratio=0.875, + ) + elif "Wan-VACE" in args.model_type: + pipe = WanVACEPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + elif "Animate" in args.model_type: + image_encoder = CLIPVisionModel.from_pretrained( + "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", torch_dtype=torch.bfloat16 + ) + image_processor = CLIPImageProcessor.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") + + pipe = WanAnimatePipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + image_encoder=image_encoder, + image_processor=image_processor, + ) + else: + pipe = WanPipeline( + transformer=transformer, + text_encoder=text_encoder, + tokenizer=tokenizer, + vae=vae, + scheduler=scheduler, + ) + + pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB") diff --git a/scripts/recipes/wuerstchen.py b/scripts/recipes/wuerstchen.py new file mode 100644 index 000000000000..c2d9a134e034 --- /dev/null +++ b/scripts/recipes/wuerstchen.py @@ -0,0 +1,87 @@ +# Run inside root directory of official source code: https://github.com/dome272/wuerstchen/ +import argparse +import os + +import torch +from transformers import AutoTokenizer, CLIPTextModel +from vqgan import VQModel + +from diffusers import ( + DDPMWuerstchenScheduler, + WuerstchenCombinedPipeline, + WuerstchenDecoderPipeline, + WuerstchenPriorPipeline, +) +from diffusers.loaders.conversion import get_conversion +from diffusers.pipelines.deprecated.wuerstchen import PaellaVQModel, WuerstchenDiffNeXt, WuerstchenPrior + + +def main(args): + model_path = args.model_path + device = "cpu" + + paella_vqmodel = VQModel() + state_dict = torch.load(os.path.join(model_path, "vqgan_f4_v1_500k.pt"), map_location=device)["state_dict"] + paella_vqmodel.load_state_dict(state_dict) + + vqmodel = PaellaVQModel(num_vq_embeddings=paella_vqmodel.codebook_size, latent_channels=paella_vqmodel.c_latent) + vqmodel.load_state_dict(get_conversion("PaellaVQModel", dict(vqmodel.config)).to_diffusers(state_dict)) + + # Clip Text encoder and tokenizer + text_encoder = CLIPTextModel.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") + tokenizer = AutoTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k") + + # Generator + gen_text_encoder = CLIPTextModel.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K").to("cpu") + gen_tokenizer = AutoTokenizer.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K") + + orig_state_dict = torch.load(os.path.join(model_path, "model_v2_stage_b.pt"), map_location=device)["state_dict"] + decoder = WuerstchenDiffNeXt() + decoder.load_state_dict(get_conversion("WuerstchenDiffNeXt", dict(decoder.config)).to_diffusers(orig_state_dict)) + + # Prior + orig_state_dict = torch.load(os.path.join(model_path, "model_v3_stage_c.pt"), map_location=device)[ + "ema_state_dict" + ] + prior_model = WuerstchenPrior(c_in=16, c=1536, c_cond=1280, c_r=64, depth=32, nhead=24).to(device) + prior_model.load_state_dict( + get_conversion("WuerstchenPrior", dict(prior_model.config)).to_diffusers(orig_state_dict) + ) + + # scheduler + scheduler = DDPMWuerstchenScheduler() + + # Prior pipeline + prior_pipeline = WuerstchenPriorPipeline( + prior=prior_model, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler + ) + + prior_pipeline.save_pretrained(os.path.join(args.output_path, "wuerstchen-prior")) + + decoder_pipeline = WuerstchenDecoderPipeline( + text_encoder=gen_text_encoder, tokenizer=gen_tokenizer, vqgan=vqmodel, decoder=decoder, scheduler=scheduler + ) + decoder_pipeline.save_pretrained(os.path.join(args.output_path, "wuerstchen")) + + # Wuerstchen pipeline + wuerstchen_pipeline = WuerstchenCombinedPipeline( + # Decoder + text_encoder=gen_text_encoder, + tokenizer=gen_tokenizer, + decoder=decoder, + scheduler=scheduler, + vqgan=vqmodel, + # Prior + prior_tokenizer=tokenizer, + prior_text_encoder=text_encoder, + prior=prior_model, + prior_scheduler=scheduler, + ) + wuerstchen_pipeline.save_pretrained(os.path.join(args.output_path, "WuerstchenCombinedPipeline")) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Assemble Wuerstchen from its original VAE, decoder and prior files.") + parser.add_argument("--model-path", default="models", help="Directory containing the three original checkpoints") + parser.add_argument("--output-path", default="warp-ai", help="Directory for the converted pipelines") + main(parser.parse_args()) diff --git a/scripts/recipes/zero123.py b/scripts/recipes/zero123.py new file mode 100644 index 000000000000..a77829c81337 --- /dev/null +++ b/scripts/recipes/zero123.py @@ -0,0 +1,173 @@ +""" +This script modified from +https://github.com/huggingface/diffusers/blob/bc691231360a4cbc7d19a58742ebb8ed0f05e027/scripts/convert_original_stable_diffusion_to_diffusers.py + +Convert original Zero1to3 checkpoint to diffusers checkpoint. + +# run the convert script +$ python recipes/zero123.py \ + --checkpoint_path /path/zero123/105000.ckpt \ + --dump_path ./zero1to3 \ + --original_config_file /path/zero123/configs/sd-objaverse-finetune-c_concat-256.yaml +``` +""" + +import argparse + +import torch +import yaml +from accelerate import init_empty_weights +from accelerate.utils import set_module_tensor_to_device +from pipeline_zero1to3 import CCProjection, Zero1to3StableDiffusionPipeline +from transformers import ( + CLIPImageProcessor, + CLIPVisionModelWithProjection, +) + +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.configs.zero123 import create_unet_diffusers_config, create_vae_diffusers_config +from diffusers.models import ( + AutoencoderKL, + UNet2DConditionModel, +) +from diffusers.schedulers import DDIMScheduler +from diffusers.utils import logging + + +logger = logging.get_logger(__name__) + + +def convert_ldm_unet_checkpoint(checkpoint, config, path=None, extract_ema=False, **kwargs): + return convert_component_checkpoint(checkpoint, config, "UNet2DConditionModel", extract_ema=extract_ema) + + +def convert_ldm_vae_checkpoint(checkpoint, config): + return convert_component_checkpoint(checkpoint, config, "AutoencoderKL") + + +def convert_from_original_zero123_ckpt(checkpoint_path, original_config_file, extract_ema, device): + ckpt = torch.load(checkpoint_path, map_location=device) + ckpt["global_step"] + checkpoint = ckpt["state_dict"] + del ckpt + torch.cuda.empty_cache() + + original_config = yaml.safe_load(original_config_file) + original_config["model"]["params"]["cond_stage_config"]["target"].split(".")[-1] + num_in_channels = 8 + original_config["model"]["params"]["unet_config"]["params"]["in_channels"] = num_in_channels + prediction_type = "epsilon" + image_size = 256 + num_train_timesteps = getattr(original_config["model"]["params"], "timesteps", None) or 1000 + + beta_start = getattr(original_config["model"]["params"], "linear_start", None) or 0.02 + beta_end = getattr(original_config["model"]["params"], "linear_end", None) or 0.085 + scheduler = DDIMScheduler( + beta_end=beta_end, + beta_schedule="scaled_linear", + beta_start=beta_start, + num_train_timesteps=num_train_timesteps, + steps_offset=1, + clip_sample=False, + set_alpha_to_one=False, + prediction_type=prediction_type, + ) + scheduler.register_to_config(clip_sample=False) + + # Convert the UNet2DConditionModel model. + upcast_attention = None + unet_config = create_unet_diffusers_config(original_config, image_size=image_size) + unet_config["upcast_attention"] = upcast_attention + with init_empty_weights(): + unet = UNet2DConditionModel(**unet_config) + converted_unet_checkpoint = convert_ldm_unet_checkpoint( + checkpoint, unet_config, path=None, extract_ema=extract_ema + ) + for param_name, param in converted_unet_checkpoint.items(): + set_module_tensor_to_device(unet, param_name, "cpu", value=param) + + # Convert the VAE model. + vae_config = create_vae_diffusers_config(original_config, image_size=image_size) + converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config) + + if ( + "model" in original_config + and "params" in original_config["model"] + and "scale_factor" in original_config["model"]["params"] + ): + vae_scaling_factor = original_config["model"]["params"]["scale_factor"] + else: + vae_scaling_factor = 0.18215 # default SD scaling factor + + vae_config["scaling_factor"] = vae_scaling_factor + + with init_empty_weights(): + vae = AutoencoderKL(**vae_config) + + for param_name, param in converted_vae_checkpoint.items(): + set_module_tensor_to_device(vae, param_name, "cpu", value=param) + + feature_extractor = CLIPImageProcessor.from_pretrained( + "lambdalabs/sd-image-variations-diffusers", subfolder="feature_extractor" + ) + image_encoder = CLIPVisionModelWithProjection.from_pretrained( + "lambdalabs/sd-image-variations-diffusers", subfolder="image_encoder" + ) + + cc_projection = CCProjection() + cc_projection.load_state_dict( + get_conversion("CCProjection", {}).to_diffusers( + {key: value.cpu() for key, value in checkpoint.items() if key.startswith("cc_projection.")} + ) + ) + + pipe = Zero1to3StableDiffusionPipeline( + vae, image_encoder, unet, scheduler, None, feature_extractor, cc_projection, requires_safety_checker=False + ) + + return pipe + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." + ) + parser.add_argument( + "--original_config_file", + default=None, + type=str, + help="The YAML config file corresponding to the original architecture.", + ) + parser.add_argument( + "--extract_ema", + action="store_true", + help=( + "Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights" + " or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield" + " higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning." + ), + ) + parser.add_argument( + "--to_safetensors", + action="store_true", + help="Whether to store pipeline in safetensors format or not.", + ) + parser.add_argument("--half", action="store_true", help="Save weights in half precision.") + parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.") + parser.add_argument("--device", type=str, help="Device to use (e.g. cpu, cuda:0, cuda:1, etc.)") + args = parser.parse_args() + + pipe = convert_from_original_zero123_ckpt( + checkpoint_path=args.checkpoint_path, + original_config_file=args.original_config_file, + extract_ema=args.extract_ema, + device=args.device, + ) + + if args.half: + pipe.to(dtype=torch.float16) + + pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors) diff --git a/setup.py b/setup.py index 60fde132fc6d..92e8524ddd96 100644 --- a/setup.py +++ b/setup.py @@ -284,7 +284,7 @@ def run(self): url="https://github.com/huggingface/diffusers", package_dir={"": "src"}, packages=find_packages("src"), - package_data={"diffusers": ["py.typed"]}, + package_data={"diffusers": ["py.typed", "loaders/conversion/configs/*.json", "loaders/conversion/presets/*.json"]}, include_package_data=True, python_requires=">=3.10.0", install_requires=list(install_requires), diff --git a/src/diffusers/loaders/conversion/README.md b/src/diffusers/loaders/conversion/README.md new file mode 100644 index 000000000000..1d90cd5e74b5 --- /dev/null +++ b/src/diffusers/loaders/conversion/README.md @@ -0,0 +1,230 @@ +# Reversible component conversions + +Every definition is a function of the matching Diffusers component config that returns +`Conversion(mapping=..., rules=...)`. `mapping` contains exact renames; grouped tensor +operations declare both directions in a `Transform`. Definitions never instantiate models +or infer architecture from checkpoint tensors. `get_conversion` supplies constructor defaults. + +```python +from diffusers.loaders.conversion import get_conversion + +conversion = get_conversion("FluxTransformer2DModel", config) +original = conversion.to_original(diffusers_state) +restored = conversion.to_diffusers(original) +``` + +Use `scripts/convert_checkpoint.py --list-models` from the repository root to inspect the +registry. The generic CLI reads local tensor files or validated shard indexes and writes a +new safetensors component directory or a PyTorch checkpoint file. See the [usage guide](../../../../docs/source/en/using-diffusers/other-formats.md#bidirectional-component-conversion). + +## Contract and format boundaries + +- A config must describe the actual checkpoint variant, including layer counts, optional + projections, and dimensions. The same config is needed when exporting a fresh Diffusers + checkpoint; a previous original checkpoint is unnecessary. +- Exact input coverage is required. Prefixes, PyTorch wrappers, external assets, auxiliary + training state, and tensor-parallel rank assembly belong outside a `Conversion`. + `checkpoint.py` shares known component prefixes, auxiliary-state filtering and declared missing buffers + between single-file loading and file conversion. `source.py` composes disjoint files or extracts module state; + pipeline recipes handle tensor-parallel assembly and external runtime preparation. `configs/` supplies reusable + original-config helpers and presets. None of these layers authors a second tensor mapping. +- Tensor operations preserve dtype/device and do not mutate inputs. Outputs may alias input + storage. The file writer clones views for safetensors and operates on CPU one rule at a + time, retaining at most one loaded PyTorch shard plus output shards and grouped tensors. + Writing a single PyTorch file materializes the converted component state dict. +- `lossless=False` identifies normalization. In the gated LTX2 diffusion decoder profile, + import folds linear gates into weights and export emits unit gates. Original gate values + are not recoverable. Weight tensors round-trip exactly for FP32/FP16/BF16; non-unit gate + folding follows the original importer's FP32 arithmetic. +- `MergeEqual` enforces equality when a source shared parameter was expanded into several + target parameters. Independently fine-tuned unequal copies cannot be represented by one + original tensor. `WithConstants` validates fixed buffers before removing/regenerating them. +- LoRA takes an explicit list of fully qualified `modules` in its config; underscore-joined + Kohya keys are not parsed heuristically. `include_alpha=True` retains alpha tensors as + `.alpha`, separate from PEFT factors; callers must pass network alphas when loading + adapters. DoRA is supported for Kohya/PEFT layouts. Adapter merging is not inverted. +- Some original formats already use Diffusers/Transformers names (for example Qwen Image, + Motif Video, T5, and Qwen3). Their identity keys are explicitly enumerated from config; + unknown tensors are still rejected. Transformers safetensors may omit tied T5/UMT5 or + configured Qwen3 embedding copies; file conversion reconstructs only those declared ties. +- RAE conversion expects the assembled encoder, decoder, statistics, and processor state + represented by its definition; a decoder-only training file does not contain all assets. + ConsistencyDecoderVAE supports both the original Python module graph and `consistency_decoder_jit` + archive keys. Use a TorchScript source manifest or the `consistency_decoder` pipeline recipe for archive extraction. +- CLAP conversion describes the complete component state, including persistent RoBERTa + buffers. Shared source preparation supplies declared position/token-type buffers missing from historical checkpoints. Tensor-parallel MiniMax H3 uses `minimax_h3_shards` after rank + assembly; the generic reader handles storage shards, not tensor-parallel concatenation. +- `Blip2QFormerModel` names the BLIP Diffusion composite in Diffusers' deprecated pipeline, + including its vision encoder and projection. It does not refer to Transformers' bare + Q-Former class with the same name. `ContextCLIPTextModel` includes its persistent position IDs. +- Output `conversion_config.json` records a Diffusers config, not original runtime settings. + Source formats inferred during import (such as Cosmos 1 or gated LTX2 decoders) are persisted as `original_format` + in the converted component config and single-file loaded model config, so a later export selects the same layout. + Tokenizers, schedulers, training loss modules, optimizer/EMA history, inference-irrelevant + original buffers, and full pipeline packaging are not reconstructed by tensor rules. + The shared PyTorch exporter accepts nested output wrappers and regenerates fixed CogVideoX positional embeddings. + Packed quantization, ONNX/TensorRT graph export and LoRA merging are separate operations. + +## Verifying a conversion + +For layout-only changes, compare original-importer and shared-converter outputs in the same target layout, including +key names, shapes, dtypes and exact tensor values. A SHA-256 digest of sorted tensor metadata and contiguous raw bytes +can verify the same equality without depending on serialization order or file metadata. Hashing the complete input and +output files is not meaningful when layouts, sharding or wrappers differ. + +Test both directions and include source-shaped fixtures, bundle selection, auxiliary state and format persistence. +A round trip through one definition alone cannot detect a consistently incorrect rename or permutation. Full-size +GPU inference is not required for a pure layout refactor when independent tensor parity is established. For numerical +normalizations such as gate folding, compare the specified arithmetic and use tolerances where rounding is expected. + +## Registered components + +An empty format cell means the component has one config-driven canonical layout. For rows +with named variants, omit `original_format` to use the definition's documented default or set +it explicitly. Architecture options within a variant still come from the component config. + +| Component class | Definition | Named `original_format` variants | +| --- | --- | --- | +| `AceStepAudioTokenDetokenizer` | [ace_step_detokenizer.py](ace_step_detokenizer.py) | | +| `AceStepAudioTokenizer` | [ace_step_tokenizer.py](ace_step_tokenizer.py) | | +| `AceStepConditionEncoder` | [ace_step_conditioner.py](ace_step_conditioner.py) | | +| `AceStepTransformer1DModel` | [ace_step.py](ace_step.py) | | +| `AnimaTextConditioner` | [anima_conditioner.py](anima_conditioner.py) | | +| `AnyFlowFARTransformer3DModel` | [anyflow_far.py](anyflow_far.py) | | +| `AnyFlowTransformer3DModel` | [anyflow.py](anyflow.py) | | +| `AsymmetricAutoencoderKL` | [asymmetric_vae.py](asymmetric_vae.py) | | +| `AudioLDM2ProjectionModel` | [audioldm2_projection.py](audioldm2_projection.py) | | +| `AudioLDM2UNet2DConditionModel` | [audioldm2_unet.py](audioldm2_unet.py) | | +| `AuraFlowTransformer2DModel` | [auraflow.py](auraflow.py) | | +| `AutoencoderDC` | [autoencoder_dc.py](autoencoder_dc.py) | | +| `AutoencoderKLCogVideoX` | [cogvideox.py](cogvideox.py) | | +| `AutoencoderKLCosmos` | [cosmos_vae.py](cosmos_vae.py) | | +| `AutoencoderKLFlux2` | [flux2_vae.py](flux2_vae.py) | | +| `AutoencoderKLHunyuanImageRefiner` | [hunyuan_image_refiner_vae.py](hunyuan_image_refiner_vae.py) | | +| `AutoencoderKLHunyuanImage` | [hunyuan_image_vae.py](hunyuan_image_vae.py) | `hunyuan_image_vae` (5D convolutions), `hunyuan_image_vae_2d` (pre-normalized 4D) | +| `AutoencoderKLHunyuanVideo15` | [hunyuan_video15_vae.py](hunyuan_video15_vae.py) | | +| `AutoencoderKLHunyuanVideo` | [hunyuan_video_vae.py](hunyuan_video_vae.py) | | +| `AutoencoderKLLTX2Audio` | [ltx2_audio_vae.py](ltx2_audio_vae.py) | | +| `AutoencoderKLLTX2Video` | [ltx2_vae.py](ltx2_vae.py) | | +| `AutoencoderKLLTXVideo` | [ltx_vae.py](ltx_vae.py) | | +| `AutoencoderKLMiniMaxH3Audio` | [minimax_h3_audio_vae.py](minimax_h3_audio_vae.py) | | +| `AutoencoderKLMiniMaxH3` | [minimax_h3_vae.py](minimax_h3_vae.py) | | +| `AutoencoderKLMochi` | [mochi_vae.py](mochi_vae.py) | | +| `AutoencoderKLQwenImage` | [qwen_image_vae.py](qwen_image_vae.py) | | +| `AutoencoderKLTemporalDecoder` | [svd_vae.py](svd_vae.py) | `svd`, `temporal_vae` | +| `AutoencoderKLWan` | [wan_vae.py](wan_vae.py) | | +| `AutoencoderKL` | [ldm_vae.py](ldm_vae.py) | | +| `AutoencoderOobleck` | [oobleck.py](oobleck.py) | | +| `AutoencoderRAE` | [rae.py](rae.py) | | +| `AutoencoderSAME` | [same.py](same.py) | | +| `AutoencoderTiny` | [tiny_vae.py](tiny_vae.py) | | +| `Blip2QFormerModel` | [blip_qformer.py](blip_qformer.py) | | +| `CCProjection` | [zero123_projection.py](zero123_projection.py) | | +| `CLIPTextModelWithProjection` | [clip.py](clip.py) | `clip`, `openclip` | +| `CLIPTextModel` | [clip.py](clip.py) | `clip`, `openclip` | +| `CLIPVisionModelWithProjection` | [clip_vision.py](clip_vision.py) | `clip`, `openclip` | +| `CLIPVisionModel` | [clip_vision.py](clip_vision.py) | `clip`, `openclip` | +| `ChromaTransformer2DModel` | [chroma.py](chroma.py) | | +| `ChronoEditTransformer3DModel` | [chronoedit.py](chronoedit.py) | | +| `ClapAudioModelWithProjection` | [clap_audio.py](clap_audio.py) | | +| `ClapAudioModel` | [clap_audio.py](clap_audio.py) | | +| `ClapModel` | [clap_audio.py](clap_audio.py) | | +| `ClapTextModelWithProjection` | [clap_text.py](clap_text.py) | | +| `CogVideoXTransformer3DModel` | [cogvideox.py](cogvideox.py) | | +| `CogView3PlusTransformer2DModel` | [cogview3plus.py](cogview3plus.py) | | +| `CogView4Transformer2DModel` | [cogview4.py](cogview4.py) | `cogview4`, `megatron` | +| `ConsistencyDecoderVAE` | [consistency_decoder.py](consistency_decoder.py) | `consistency_decoder_jit` | +| `ContextCLIPTextModel` | [clip.py](clip.py) | `clip`, `openclip` | +| `ControlNetModel` | [controlnet.py](controlnet.py) | | +| `CosmosControlNetModel` | [cosmos_controlnet.py](cosmos_controlnet.py) | | +| `CosmosTransformer3DModel` | [cosmos.py](cosmos.py) | `cosmos1`, `cosmos2` | +| `DiTTransformer2DModel` | [dit.py](dit.py) | | +| `ErnieImageTransformer2DModel` | [ernie_image.py](ernie_image.py) | | +| `Flux2Transformer2DModel` | [flux2.py](flux2.py) | | +| `FluxIPAdapter` | [flux_ip_adapter.py](flux_ip_adapter.py) | | +| `FluxTransformer2DModel` | [flux.py](flux.py) | | +| `HiDreamImageTransformer2DModel` | [hidream.py](hidream.py) | | +| `HunyuanDiT2DControlNetModel` | [hunyuan_dit_controlnet.py](hunyuan_dit_controlnet.py) | | +| `HunyuanDiT2DModel` | [hunyuan_dit.py](hunyuan_dit.py) | | +| `HunyuanImageTransformer2DModel` | [hunyuan_image.py](hunyuan_image.py) | `hunyuan_image_fused`, `hunyuan_image_split` | +| `HunyuanVideo15Transformer3DModel` | [hunyuan_video15.py](hunyuan_video15.py) | | +| `HunyuanVideoTransformer3DModel` | [hunyuan_video.py](hunyuan_video.py) | | +| `I2VGenXLUNet` | [i2vgen_xl.py](i2vgen_xl.py) | | +| `IFSafetyChecker` | [pipeline_components.py](pipeline_components.py) | | +| `JoyImageEditPlusTransformer3DModel` | [joy_image.py](joy_image.py) | | +| `JoyImageEditTransformer3DModel` | [joy_image.py](joy_image.py) | | +| `Kandinsky3UNet` | [kandinsky3.py](kandinsky3.py) | | +| `LDMBertModel` | [pipeline_components.py](pipeline_components.py) | | +| `LTX2DurationHead` | [ltx2_duration.py](ltx2_duration.py) | | +| `LTX2LatentUpsamplerModel` | [ltx2_upsampler.py](ltx2_upsampler.py) | | +| `LTX2TextConnectors` | [ltx2_connectors.py](ltx2_connectors.py) | | +| `LTX2VideoDiffusionDecoderModel` | [ltx2_diffusion_decoder.py](ltx2_diffusion_decoder.py) | `ltx2_diffusion_decoder`, `ltx2_diffusion_decoder_gated` | +| `LTX2VideoTransformer3DModel` | [ltx2.py](ltx2.py) | | +| `LTX2VocoderWithBWE` | [ltx2_vocoder.py](ltx2_vocoder.py) | | +| `LTX2Vocoder` | [ltx2_vocoder.py](ltx2_vocoder.py) | | +| `LTXLatentUpsamplerModel` | [ltx_upsampler.py](ltx_upsampler.py) | | +| `LTXVideoTransformer3DModel` | [ltx.py](ltx.py) | | +| `LearnedClassifierFreeSamplingEmbeddings` | [pipeline_components.py](pipeline_components.py) | | +| `LoRA` | [lora.py](lora.py) | `kohya`, `diffusers`, `diffusers_old`, `peft`, `animatediff` | +| `LongCatAudioDiTTransformer` | [longcat_audio.py](longcat_audio.py) | | +| `LongCatAudioDiTVae` | [longcat_audio_vae.py](longcat_audio_vae.py) | | +| `Lumina2Transformer2DModel` | [lumina2.py](lumina2.py) | | +| `LuminaNextDiT2DModel` | [lumina.py](lumina.py) | | +| `MiniMaxH3Transformer3DModel` | [minimax_h3.py](minimax_h3.py) | `minimax_h3`, `minimax_h3_shards` | +| `MiniMaxMusic3ConditionEncoder` | [minimax_music3_conditioner.py](minimax_music3_conditioner.py) | | +| `MiniMaxMusic3RVQDepthDecoder` | [minimax_music3_rvq.py](minimax_music3_rvq.py) | | +| `MiniMaxMusic3Transformer1DModel` | [minimax_music3.py](minimax_music3.py) | | +| `MiniMaxMusic3Vocoder` | [minimax_music3_vocoder.py](minimax_music3_vocoder.py) | | +| `MochiTransformer3DModel` | [mochi.py](mochi.py) | | +| `MotifVideoTransformer3DModel` | [motif_video.py](motif_video.py) | | +| `MotionAdapter` | [animatediff.py](animatediff.py) | | +| `OmniGenTransformer2DModel` | [omnigen.py](omnigen.py) | | +| `OvisImageTransformer2DModel` | [ovis_image.py](ovis_image.py) | | +| `PRXTransformer2DModel` | [prx.py](prx.py) | `prx`, `prx_weight_norm` | +| `PaellaVQModel` | [paella.py](paella.py) | | +| `PaintByExampleImageEncoder` | [pipeline_components.py](pipeline_components.py) | | +| `PixArtTransformer2DModel` | [pixart.py](pixart.py) | | +| `PriorTransformer` | [prior.py](prior.py) | `shap_e`, `unclip`, `kandinsky` | +| `Qwen3ForCausalLM` | [qwen3.py](qwen3.py) | | +| `Qwen3Model` | [qwen3.py](qwen3.py) | | +| `QwenImageTransformer2DModel` | [qwen_image.py](qwen_image.py) | | +| `SD3ControlNetModel` | [sd3_controlnet.py](sd3_controlnet.py) | | +| `SD3Transformer2DModel` | [sd3.py](sd3.py) | | +| `SanaControlNetModel` | [sana_controlnet.py](sana_controlnet.py) | | +| `SanaTransformer2DModel` | [sana.py](sana.py) | | +| `SanaVideoTransformer3DModel` | [sana_video.py](sana_video.py) | | +| `ShapERenderer` | [shap_e_renderer.py](shap_e_renderer.py) | | +| `SkyReelsV2Transformer3DModel` | [skyreels_v2.py](skyreels_v2.py) | | +| `SparseControlNetModel` | [sparse_controlnet.py](sparse_controlnet.py) | | +| `SpectrogramContEncoder` | [spectrogram_continuous.py](spectrogram_continuous.py) | | +| `SpectrogramNotesEncoder` | [spectrogram_notes.py](spectrogram_notes.py) | | +| `SpeechT5HifiGan` | [hifigan.py](hifigan.py) | | +| `StableAudio3DiTModel` | [stable_audio3.py](stable_audio3.py) | | +| `StableAudio3DurationEmbedder` | [stable_audio3_duration.py](stable_audio3_duration.py) | | +| `StableAudioDiTModel` | [stable_audio.py](stable_audio.py) | | +| `StableAudioProjectionModel` | [stable_audio_projection.py](stable_audio_projection.py) | | +| `StableCascadeUNet` | [stable_cascade.py](stable_cascade.py) | | +| `T2IAdapter` | [t2i_adapter.py](t2i_adapter.py) | | +| `T5EncoderModel` | [t5.py](t5.py) | | +| `T5FilmDecoder` | [t5_film.py](t5_film.py) | | +| `Transformer2DModel` | [vq_diffusion.py](vq_diffusion.py) | | +| `UMT5EncoderModel` | [umt5.py](umt5.py) | | +| `UNet1DModel` | [unet_1d.py](unet_1d.py) | `diffuser_rl`, `diffuser_rl_legacy` | +| `UNet2DConditionModel` | [ldm_unet.py](ldm_unet.py) | `ldm`, `versatile_image` | +| `UNet2DModel` | [unet_2d.py](unet_2d.py) | `ddpm`, `ldm`, `consistency`, `ncsnpp` | +| `UNet3DConditionModel` | [unet_3d.py](unet_3d.py) | | +| `UNetFlatConditionModel` | [versatile_text_unet.py](versatile_text_unet.py) | | +| `UNetSpatioTemporalConditionModel` | [svd.py](svd.py) | | +| `UVit2DModel` | [uvit.py](uvit.py) | | +| `UnCLIPTextProjModel` | [unclip_text_projection.py](unclip_text_projection.py) | | +| `UniDiffuserModel` | [unidiffuser.py](unidiffuser.py) | | +| `UniDiffuserTextDecoder` | [unidiffuser_text.py](unidiffuser_text.py) | | +| `VQModel` | [vq_model.py](vq_model.py) | | +| `WanAnimate2Transformer3DModel` | [wan_animate2.py](wan_animate2.py) | | +| `WanAnimateTransformer3DModel` | [wan_animate.py](wan_animate.py) | | +| `WanTransformer3DModel` | [wan.py](wan.py) | | +| `WanVACETransformer3DModel` | [wan_vace.py](wan_vace.py) | | +| `WuerstchenDiffNeXt` | [wuerstchen_decoder.py](wuerstchen_decoder.py) | | +| `WuerstchenPrior` | [wuerstchen_prior.py](wuerstchen_prior.py) | | +| `ZImageControlNetModel` | [z_image_controlnet.py](z_image_controlnet.py) | | +| `ZImageTransformer2DModel` | [z_image.py](z_image.py) | | diff --git a/src/diffusers/loaders/conversion/__init__.py b/src/diffusers/loaders/conversion/__init__.py new file mode 100644 index 000000000000..dd2ce26c7c0f --- /dev/null +++ b/src/diffusers/loaders/conversion/__init__.py @@ -0,0 +1,30 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .registry import get_conversion +from .transforms import ( + Chain, + FoldLinearGate, + Identity, + MergeEqual, + Permute, + ReorderChunks, + Reshape, + Reverse, + Split, + Squeeze, + Transform, + WithConstants, +) diff --git a/src/diffusers/loaders/conversion/ace_step.py b/src/diffusers/loaders/conversion/ace_step.py new file mode 100644 index 000000000000..c945c82537f4 --- /dev/null +++ b/src/diffusers/loaders/conversion/ace_step.py @@ -0,0 +1,73 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ace_step_conversion(config): + mapping = {"scale_shift_table": "scale_shift_table", "norm_out.weight": "norm_out.weight"} + modules = [ + ("proj_in.1", "proj_in_conv"), + ("proj_out.1", "proj_out_conv"), + ("condition_embedder", "condition_embedder"), + ] + modules.extend( + (f"{prefix}.{name}", f"{prefix}.{name}") + for prefix in ("time_embed", "time_embed_r") + for name in ("linear_1", "linear_2", "time_proj") + ) + for i in range(config["num_hidden_layers"]): + prefix = f"layers.{i}" + mapping[prefix + ".scale_shift_table"] = prefix + ".scale_shift_table" + mapping.update(_ace_step_attention_mapping(prefix + ".self_attn", config["attention_bias"])) + mapping.update(_ace_step_attention_mapping(prefix + ".cross_attn", config["attention_bias"])) + for name in ("self_attn_norm", "cross_attn_norm", "mlp_norm", "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj"): + mapping[f"{prefix}.{name}.weight"] = f"{prefix}.{name}.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) + + +def _ace_step_attention_mapping(prefix, bias): + mapping = {} + for old, new in (("q_proj", "to_q"), ("k_proj", "to_k"), ("v_proj", "to_v"), ("o_proj", "to_out.0")): + mapping.update( + {f"{prefix}.{old}.{p}": f"{prefix}.{new}.{p}" for p in (("weight", "bias") if bias else ("weight",))} + ) + mapping.update({f"{prefix}.{part}_norm.weight": f"{prefix}.norm_{part}.weight" for part in ("q", "k")}) + return mapping + + +def _ace_step_encoder_mapping(config, count): + mapping = { + "embed_tokens.weight": "embed_tokens.weight", + "embed_tokens.bias": "embed_tokens.bias", + "norm.weight": "norm.weight", + } + for i in range(count): + prefix = f"layers.{i}" + mapping.update(_ace_step_attention_mapping(prefix + ".self_attn", config["attention_bias"])) + mapping.update( + { + f"{prefix}.{name}.weight": f"{prefix}.{name}.weight" + for name in ( + "input_layernorm", + "post_attention_layernorm", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ) + } + ) + return mapping diff --git a/src/diffusers/loaders/conversion/ace_step_conditioner.py b/src/diffusers/loaders/conversion/ace_step_conditioner.py new file mode 100644 index 000000000000..6796f4d5485d --- /dev/null +++ b/src/diffusers/loaders/conversion/ace_step_conditioner.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .ace_step import _ace_step_encoder_mapping +from .core import Conversion, Rule +from .transforms import Permute + + +def ace_step_conditioner_conversion(config): + mapping = { + "encoder.text_projector.weight": "text_projector.weight", + "null_condition_emb": "null_condition_emb", + "encoder.timbre_encoder.special_token": "timbre_encoder.special_token", + } + for prefix, count in ( + ("lyric_encoder", config["num_lyric_encoder_hidden_layers"]), + ("timbre_encoder", config["num_timbre_encoder_hidden_layers"]), + ): + mapping.update( + { + f"encoder.{prefix}.{old}": f"{prefix}.{new}" + for old, new in _ace_step_encoder_mapping(config, count).items() + } + ) + return Conversion(mapping=mapping, rules=(Rule(("silence_latent",), ("silence_latent",), Permute((0, 2, 1))),)) diff --git a/src/diffusers/loaders/conversion/ace_step_detokenizer.py b/src/diffusers/loaders/conversion/ace_step_detokenizer.py new file mode 100644 index 000000000000..b8f66c33c7a9 --- /dev/null +++ b/src/diffusers/loaders/conversion/ace_step_detokenizer.py @@ -0,0 +1,24 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .ace_step import _ace_step_encoder_mapping +from .core import Conversion + + +def ace_step_detokenizer_conversion(config): + mapping = _ace_step_encoder_mapping(config, config["num_attention_pooler_hidden_layers"]) + mapping["special_tokens"] = "special_tokens" + mapping.update({f"proj_out.{p}": f"proj_out.{p}" for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ace_step_tokenizer.py b/src/diffusers/loaders/conversion/ace_step_tokenizer.py new file mode 100644 index 000000000000..fd5ee4d394f8 --- /dev/null +++ b/src/diffusers/loaders/conversion/ace_step_tokenizer.py @@ -0,0 +1,33 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .ace_step import _ace_step_encoder_mapping +from .core import Conversion + + +def ace_step_tokenizer_conversion(config): + mapping = { + f"attention_pooler.{old}": f"attention_pooler.{new}" + for old, new in _ace_step_encoder_mapping(config, config["num_attention_pooler_hidden_layers"]).items() + } + mapping["attention_pooler.special_token"] = "attention_pooler.special_token" + mapping.update( + { + f"{name}.{p}": f"{name}.{p}" + for name in ("audio_acoustic_proj", "quantizer.project_in", "quantizer.project_out") + for p in ("weight", "bias") + } + ) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/anima_conditioner.py b/src/diffusers/loaders/conversion/anima_conditioner.py new file mode 100644 index 000000000000..271565a2da07 --- /dev/null +++ b/src/diffusers/loaders/conversion/anima_conditioner.py @@ -0,0 +1,42 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def anima_conditioner_conversion(config): + keys = ["embed.weight", "norm.weight"] + modules = ["out_proj"] + if config["model_dim"] != config["target_dim"]: + modules.append("in_proj") + for i in range(config["num_layers"]): + prefix = f"blocks.{i}" + norms = ["norm_cross_attn", "norm_mlp"] + attentions = ["cross_attn"] + if config["use_self_attention"]: + norms.append("norm_self_attn") + attentions.append("self_attn") + for norm in norms: + keys.append(f"{prefix}.{norm}.weight") + if config["use_layer_norm"]: + keys.append(f"{prefix}.{norm}.bias") + for attn in attentions: + keys.extend( + f"{prefix}.{attn}.{name}.weight" + for name in ("q_proj", "k_proj", "v_proj", "o_proj", "q_norm", "k_norm") + ) + modules.extend(f"{prefix}.mlp.{j}" for j in (0, 2)) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/animatediff.py b/src/diffusers/loaders/conversion/animatediff.py new file mode 100644 index 000000000000..f669d93c1caa --- /dev/null +++ b/src/diffusers/loaders/conversion/animatediff.py @@ -0,0 +1,74 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import MergeEqual + + +def animatediff_conversion(config): + mapping, modules, rules = {}, [], [] + count = len(config["block_out_channels"]) + layers = config["motion_layers_per_block"] + layers = [layers] * count if isinstance(layers, int) else list(layers) + depths = config["motion_transformer_layers_per_block"] + depths = [depths] * count if isinstance(depths, int) else list(depths) + blocks = [(f"down_blocks.{i}", layers[i], depths[i]) for i in range(count)] + blocks.extend((f"up_blocks.{i}", layers[count - 1 - i] + 1, depths[count - 1 - i]) for i in range(count)) + if config["use_motion_mid_block"]: + blocks.append( + ( + "mid_block", + config["motion_mid_block_layers_per_block"], + config["motion_transformer_layers_per_mid_block"], + ) + ) + if config["conv_in_channels"]: + modules.append(("conv_in", "conv_in")) + for prefix, layers, depth in blocks: + for i in range(layers): + new = f"{prefix}.motion_modules.{i}" + old = new + ".temporal_transformer" + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("norm", "proj_in", "proj_out")) + n = depth[i] if isinstance(depth, (list, tuple)) else depth + for j in range(n): + a, b = f"{old}.transformer_blocks.{j}", f"{new}.transformer_blocks.{j}" + modules.extend( + (f"{a}.{source}", f"{b}.{target}") + for source, target in ( + ("norms.0", "norm1"), + ("norms.1", "norm2"), + ("ff_norm", "norm3"), + ("ff.net.0.proj", "ff.net.0.proj"), + ("ff.net.2", "ff.net.2"), + ("attention_blocks.0.to_out.0", "attn1.to_out.0"), + ("attention_blocks.1.to_out.0", "attn2.to_out.0"), + ) + ) + rules.append( + Rule( + tuple(f"{a}.attention_blocks.{k}.pos_encoder.pe" for k in (0, 1)), + (b + ".pos_embed.pe",), + MergeEqual(2), + ) + ) + mapping.update( + { + f"{a}.attention_blocks.{k}.to_{part}.weight": f"{b}.attn{k + 1}.to_{part}.weight" + for k in (0, 1) + for part in ("q", "k", "v") + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/anyflow.py b/src/diffusers/loaders/conversion/anyflow.py new file mode 100644 index 000000000000..d56773ff49a5 --- /dev/null +++ b/src/diffusers/loaders/conversion/anyflow.py @@ -0,0 +1,26 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .wan import wan_conversion + + +def anyflow_conversion(config): + base = wan_conversion( + {**config, "qk_norm": "rms_norm_across_heads", "added_kv_proj_dim": None, "pos_embed_seq_len": None} + ) + keys = set(base.diffusers_keys) + keys.update(f"condition_embedder.delta_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in sorted(keys)}) diff --git a/src/diffusers/loaders/conversion/anyflow_far.py b/src/diffusers/loaders/conversion/anyflow_far.py new file mode 100644 index 000000000000..6e6856d70e3c --- /dev/null +++ b/src/diffusers/loaders/conversion/anyflow_far.py @@ -0,0 +1,23 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .anyflow import anyflow_conversion +from .core import Conversion + + +def anyflow_far_conversion(config): + mapping = dict(anyflow_conversion(config).mapping) + mapping.update({f"far_patch_embedding.{p}": f"far_patch_embedding.{p}" for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/asymmetric_vae.py b/src/diffusers/loaders/conversion/asymmetric_vae.py new file mode 100644 index 000000000000..58fa58b3a227 --- /dev/null +++ b/src/diffusers/loaders/conversion/asymmetric_vae.py @@ -0,0 +1,51 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def asymmetric_vae_conversion(config): + encoder = ldm_vae_conversion( + { + **config, + "block_out_channels": config["down_block_out_channels"], + "layers_per_block": config["layers_per_down_block"], + "mid_block_add_attention": True, + } + ) + decoder = ldm_vae_conversion( + { + **config, + "block_out_channels": config["up_block_out_channels"], + "layers_per_block": config["layers_per_up_block"], + "mid_block_add_attention": True, + } + ) + mapping = {old: new for old, new in encoder.mapping.items() if new.startswith(("encoder.", "quant_conv."))} + mapping.update( + {old: new for old, new in decoder.mapping.items() if new.startswith(("decoder.", "post_quant_conv."))} + ) + mapping.update( + { + f"decoder.encoder.layers.{i}.{p}": f"decoder.condition_encoder.layers.{i}.{p}" + for i in range(5) + for p in ("weight", "bias") + } + ) + rules = tuple(rule for rule in encoder.rules if rule.diffusers[0].startswith("encoder.")) + tuple( + rule for rule in decoder.rules if rule.diffusers[0].startswith("decoder.") + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/audioldm2_projection.py b/src/diffusers/loaders/conversion/audioldm2_projection.py new file mode 100644 index 000000000000..f49ab054a635 --- /dev/null +++ b/src/diffusers/loaders/conversion/audioldm2_projection.py @@ -0,0 +1,31 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Reshape, Split + + +def audioldm2_projection_conversion(config): + modules = [("input_sequence_embed_linear.0", "projection"), ("input_sequence_embed_linear.1", "projection_1")] + mapping = {f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")} + dim = config["langauge_model_dim"] + transform = Chain((Reshape((2, dim), (2 * dim,)), Split((dim,) * 2))) + rules = tuple( + Rule((f"{old}_of_sequence_tokens.weight",), (f"{new}_embed", f"{new}_embed_1"), transform) + for old, new in (("start", "sos"), ("end", "eos")) + ) + if config["use_learned_position_embedding"] is not None: + mapping["learnable_positional_embedding"] = "learnable_positional_embedding" + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/audioldm2_unet.py b/src/diffusers/loaders/conversion/audioldm2_unet.py new file mode 100644 index 000000000000..bc58f86b06e3 --- /dev/null +++ b/src/diffusers/loaders/conversion/audioldm2_unet.py @@ -0,0 +1,58 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_unet import ldm_unet_conversion + + +def audioldm2_unet_conversion(config): + base = ldm_unet_conversion(config) + mapping = {} + dimensions = config["cross_attention_dim"] + if isinstance(dimensions, int): + dimensions = [dimensions] * len(config["block_out_channels"]) + counts = [len(value) if isinstance(value, (tuple, list)) else 1 for value in dimensions] + for old, new in base.mapping.items(): + if config["norm_num_groups"] is None and new.startswith("conv_norm_out."): + continue + if ".attentions." in new: + target_prefix, suffix = new.split(".attentions.", 1) + index, suffix = suffix.split(".", 1) + if target_prefix == "mid_block": + count = counts[-1] + source_prefix, source_suffix = old.split(".", 2)[0], old.split(".", 2)[2] + for j in range(count): + mapping[f"{source_prefix}.{j + 1}.{source_suffix}"] = f"mid_block.attentions.{j}.{suffix}" + else: + block_index = int(target_prefix.split(".")[1]) + if target_prefix.startswith("up_blocks"): + block_index = len(counts) - 1 - block_index + count = counts[block_index] + source_group, source_index, _, source_suffix = old.split(".", 3) + for j in range(count): + mapping[f"{source_group}.{source_index}.{j + 1}.{source_suffix}"] = ( + f"{target_prefix}.attentions.{int(index) * count + j}.{suffix}" + ) + elif new.startswith("mid_block.resnets.1."): + mapping[old.replace("middle_block.2.", f"middle_block.{counts[-1] + 1}.", 1)] = new + elif new.startswith("up_blocks.") and ".upsamplers." in new: + i = int(new.split(".")[1]) + if config["up_block_types"][i] == "CrossAttnUpBlock2D": + a, b, _, suffix = old.split(".", 3) + old = f"{a}.{b}.{counts[len(counts) - 1 - i] + 1}.{suffix}" + mapping[old] = new + else: + mapping[old] = new + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/auraflow.py b/src/diffusers/loaders/conversion/auraflow.py new file mode 100644 index 000000000000..5e7733dff99a --- /dev/null +++ b/src/diffusers/loaders/conversion/auraflow.py @@ -0,0 +1,65 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks + + +def auraflow_conversion(config): + mapping = {"register_tokens": "register_tokens", "positional_encoding": "pos_embed.pos_embed"} + modules = [("cond_seq_linear", "context_embedder"), ("final_linear", "proj_out")] + for old, new in ( + ("t_embedder.mlp.0", "time_step_proj.linear_1"), + ("t_embedder.mlp.2", "time_step_proj.linear_2"), + ("init_x_linear", "pos_embed.proj"), + ): + mapping.update({f"{old}.{p}": f"{new}.{p}" for p in ("weight", "bias")}) + for i in range(config["num_mmdit_layers"]): + old, new = f"double_layers.{i}", f"joint_transformer_blocks.{i}" + for source, target in (("mlpX", "ff"), ("mlpC", "ff_context")): + modules.extend( + (f"{old}.{source}.{a}", f"{new}.{target}.{b}") + for a, b in (("c_fc1", "linear_1"), ("c_fc2", "linear_2"), ("c_proj", "out_projection")) + ) + modules.extend( + (f"{old}.{a}.1", f"{new}.{b}.linear") for a, b in (("modX", "norm1"), ("modC", "norm1_context")) + ) + modules.extend( + (f"{old}.attn.{a}", f"{new}.attn.{b}") + for a, b in ( + ("w2q", "to_q"), + ("w2k", "to_k"), + ("w2v", "to_v"), + ("w2o", "to_out.0"), + ("w1q", "add_q_proj"), + ("w1k", "add_k_proj"), + ("w1v", "add_v_proj"), + ("w1o", "to_add_out"), + ) + ) + for i in range(config["num_single_dit_layers"]): + old, new = f"single_layers.{i}", f"single_transformer_blocks.{i}" + modules.extend( + (f"{old}.mlp.{a}", f"{new}.ff.{b}") + for a, b in (("c_fc1", "linear_1"), ("c_fc2", "linear_2"), ("c_proj", "out_projection")) + ) + modules.append((old + ".modCX.1", new + ".norm1.linear")) + modules.extend( + (f"{old}.attn.w1{a}", f"{new}.attn.{b}") + for a, b in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("o", "to_out.0")) + ) + mapping.update({old + ".weight": new + ".weight" for old, new in modules}) + return Conversion( + mapping=mapping, rules=(Rule(("modF.1.weight",), ("norm_out.linear.weight",), ReorderChunks((1, 0))),) + ) diff --git a/src/diffusers/loaders/conversion/autoencoder_dc.py b/src/diffusers/loaders/conversion/autoencoder_dc.py new file mode 100644 index 000000000000..393d9ffe1737 --- /dev/null +++ b/src/diffusers/loaders/conversion/autoencoder_dc.py @@ -0,0 +1,108 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Reshape, Split + + +def autoencoder_dc_conversion(config): + mapping, modules, rules = {}, [], [] + encoder_stem_downsample = config["encoder_layers_per_block"][0] == 0 + decoder_head_upsample = config["decoder_layers_per_block"][0] == 0 + encoder_original = "encoder.project_in.conv" + ( + ".conv" if encoder_stem_downsample and config["downsample_block_type"] == "pixel_unshuffle" else "" + ) + modules.append((encoder_original, "encoder.conv_in.conv" if encoder_stem_downsample else "encoder.conv_in")) + modules.append( + ( + "encoder.project_out." + ("main." if config["encoder_out_shortcut"] else "") + "op_list.0.conv", + "encoder.conv_out", + ) + ) + modules.append( + ("decoder.project_in." + ("main." if config["decoder_in_shortcut"] else "") + "conv", "decoder.conv_in") + ) + modules.append(("decoder.project_out.op_list.0", "decoder.norm_out")) + modules.append( + ( + "decoder.project_out.op_list.2.conv" + (".conv" if decoder_head_upsample else ""), + "decoder.conv_out.conv" if decoder_head_upsample else "decoder.conv_out", + ) + ) + for component, direction in (("encoder", "down"), ("decoder", "up")): + channels = config[f"{component}_block_out_channels"] + layers = config[f"{component}_layers_per_block"] + kinds = config[f"{component}_block_types"] + kinds = [kinds] * len(channels) if isinstance(kinds, str) else kinds + norms = "rms_norm" if component == "encoder" else config["decoder_norm_types"] + norms = [norms] * len(channels) if isinstance(norms, str) else norms + for i, channel in enumerate(channels): + resample = i < len(channels) - 1 and layers[i] > 0 + if resample: + j = layers[i] if component == "encoder" else 0 + old = f"{component}.stages.{i}.op_list.{j}.main.conv" + if component == "decoder" or config["downsample_block_type"] == "pixel_unshuffle": + old += ".conv" + modules.append((old, f"{component}.{direction}_blocks.{i}.{j}.conv")) + for j in range(layers[i]): + index = j + int(component == "decoder" and resample) + old = f"{component}.stages.{i}.op_list.{index}" + new = f"{component}.{direction}_blocks.{i}.{index}" + if kinds[i] == "ResBlock": + modules.append((old + ".main.conv1.conv", new + ".conv1")) + mapping[old + ".main.conv2.conv.weight"] = new + ".conv2.weight" + modules.append((old + ".main.conv2.norm", new + ".norm")) + elif kinds[i] == "EfficientViTBlock": + source, target = old + ".context_module.main", new + ".attn" + rules.append( + Rule( + (source + ".qkv.conv.weight",), + tuple(f"{target}.to_{part}.weight" for part in ("q", "k", "v")), + Chain( + (Reshape((3 * channel, channel, 1, 1), (3 * channel, channel)), Split((channel,) * 3)) + ), + ) + ) + scales = config[f"{component}_qkv_multiscales"][i] + for k in range(len(scales)): + mapping[f"{source}.aggreg.{k}.0.weight"] = f"{target}.to_qkv_multiscale.{k}.proj_in.weight" + mapping[f"{source}.aggreg.{k}.1.weight"] = f"{target}.to_qkv_multiscale.{k}.proj_out.weight" + width = channel * (len(scales) + 1) + rules.append( + Rule( + (source + ".proj.conv.weight",), + (target + ".to_out.weight",), + Reshape((channel, width, 1, 1), (channel, width)), + ) + ) + modules.append((source + ".proj.norm", target + ".norm_out")) + for a, b in ( + ("inverted_conv.conv", "conv_inverted"), + ("depth_conv.conv", "conv_depth"), + ("point_conv.norm", "norm"), + ): + modules.append((f"{old}.local_module.main.{a}", f"{new}.conv_out.{b}")) + mapping[old + ".local_module.main.point_conv.conv.weight"] = new + ".conv_out.conv_point.weight" + else: + raise ValueError(f"Unknown DC-AE block type {kinds[i]}.") + if norms[i] == "batch_norm": + source_norm = old + ( + ".main.conv2.norm" if kinds[i] == "ResBlock" else ".context_module.main.proj.norm" + ) + target_norm = new + (".norm" if kinds[i] == "ResBlock" else ".attn.norm_out") + for name in ("running_mean", "running_var", "num_batches_tracked"): + mapping[f"{source_norm}.{name}"] = f"{target_norm}.{name}" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/blip_qformer.py b/src/diffusers/loaders/conversion/blip_qformer.py new file mode 100644 index 000000000000..e42c071dec63 --- /dev/null +++ b/src/diffusers/loaders/conversion/blip_qformer.py @@ -0,0 +1,91 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from .core import Conversion, Rule +from .transforms import Reshape, WithConstants + + +def blip_qformer_conversion(config): + qformer, vision = config["qformer_config"], config["vision_config"] + mapping = { + "blip.query_tokens": "query_tokens", + "blip.Qformer.bert.embeddings.word_embeddings.weight": "embeddings.word_embeddings.weight", + "blip.Qformer.bert.embeddings.position_embeddings.weight": "embeddings.position_embeddings.weight", + "blip.visual_encoder.conv1.weight": "visual_encoder.embeddings.patch_embedding.weight", + } + modules = [ + ("blip.Qformer.bert.embeddings.LayerNorm", "embeddings.LayerNorm"), + ("blip.visual_encoder.ln_pre", "visual_encoder.pre_layernorm"), + ("blip.ln_vision", "visual_encoder.post_layernorm"), + ] + modules.extend(("proj_layer." + name, "proj_layer." + name) for name in ("dense1", "dense2", "LayerNorm")) + for i in range(qformer["num_hidden_layers"]): + old, new = f"blip.Qformer.bert.encoder.layer.{i}", f"encoder.layer.{i}" + attentions = ["attention"] + if i % qformer["cross_attention_frequency"] == 0: + attentions.append("crossattention") + for attn in attentions: + modules.extend( + (f"{old}.{attn}.self.{part}", f"{new}.{attn}.attention.{part}") for part in ("query", "key", "value") + ) + modules.extend( + (f"{old}.{attn}.output.{name}", f"{new}.{attn}.output.{name}") for name in ("dense", "LayerNorm") + ) + modules.extend( + (f"{old}.{name}", f"{new}.{name}") + for name in ( + "intermediate.dense", + "intermediate_query.dense", + "output.dense", + "output.LayerNorm", + "output_query.dense", + "output_query.LayerNorm", + ) + ) + for i in range(vision["num_hidden_layers"]): + old, new = f"blip.visual_encoder.transformer.resblocks.{i}", f"visual_encoder.encoder.layers.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ln_1", "layer_norm1"), + ("ln_2", "layer_norm2"), + ("attn.out_proj", "self_attn.projection"), + ("mlp.c_fc", "mlp.fc1"), + ("mlp.c_proj", "mlp.fc2"), + ) + ) + mapping.update({f"{old}.attn.in_proj_{p}": f"{new}.self_attn.qkv.{p}" for p in ("weight", "bias")}) + hidden = vision["hidden_size"] + positions = (vision["image_size"] // vision["patch_size"]) ** 2 + 1 + rules = ( + Rule( + ("blip.visual_encoder.class_embedding",), + ("visual_encoder.embeddings.class_embedding",), + Reshape((hidden,), (1, 1, hidden)), + ), + Rule( + ("blip.visual_encoder.positional_embedding",), + ("visual_encoder.embeddings.position_embedding",), + Reshape((positions, hidden), (1, positions, hidden)), + ), + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + position_ids = torch.arange(qformer["max_position_embeddings"], device="cpu").expand((1, -1)) + source = "blip.Qformer.bert.embeddings.word_embeddings.weight" + target = mapping.pop(source) + rules += (Rule((source,), (target, "embeddings.position_ids"), WithConstants((position_ids,))),) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/checkpoint.py b/src/diffusers/loaders/conversion/checkpoint.py new file mode 100644 index 000000000000..f08a89c6b7c4 --- /dev/null +++ b/src/diffusers/loaders/conversion/checkpoint.py @@ -0,0 +1,238 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Adapt published single-file component containers to the strict tensor conversion API.""" + +from collections.abc import Iterator, Mapping +from typing import Any + +import torch + +from ...utils import logging +from .cogvideox import ( + convert_cogvideox_transformer_checkpoint_to_diffusers, + convert_cogvideox_vae_checkpoint_to_diffusers, +) +from .core import Conversion +from .registry import get_conversion + + +logger = logging.get_logger(__name__) + + +class ComponentState(Mapping[str, Any]): + """Select a component without reading its tensors until the conversion needs them.""" + + def __init__(self, checkpoint: Mapping[str, Any], prefix: str = "") -> None: + self.checkpoint = checkpoint + self.keys_to_source = {key[len(prefix) :]: key for key in checkpoint if key.startswith(prefix)} + self.generated = {} + + def __getitem__(self, key: str) -> Any: + if key in self.generated: + return self.generated[key] + return self.checkpoint[self.keys_to_source[key]] + + def __iter__(self) -> Iterator[str]: + return iter(self.keys_to_source.keys() | self.generated.keys()) + + def __len__(self) -> int: + return len(self.keys_to_source.keys() | self.generated.keys()) + + +def convert_component_checkpoint( + checkpoint: Mapping[str, Any], + config: dict[str, Any], + model_class: str, + *, + extract_ema: bool = False, + return_config: bool = False, +) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Import a single-file component with its resolved Diffusers config, without mutating the checkpoint. + + Prefix selection and the known auxiliary keys formerly handled in `single_file_utils` live here. The conversion + definition itself remains a strict, reversible mapping of component tensors. Only explicitly known auxiliary state + is excluded; unknown keys inside the selected component raise instead of silently disappearing. With + `return_config=True`, also return the resolved config so callers can persist inferred original formats. + """ + config = dict(config) + if model_class == "CogVideoXTransformer3DModel": + converted = convert_cogvideox_transformer_checkpoint_to_diffusers(checkpoint, config=config) + elif model_class == "AutoencoderKLCogVideoX": + converted = convert_cogvideox_vae_checkpoint_to_diffusers(checkpoint, config=config) + else: + state, conversion, config = prepare_component_checkpoint( + checkpoint, config, model_class, extract_ema=extract_ema + ) + converted = conversion.to_diffusers(state) + return (converted, config) if return_config else converted + + +def prepare_component_checkpoint( + checkpoint: Mapping[str, Any], + config: dict[str, Any], + model_class: str, + *, + extract_ema: bool = False, +) -> tuple[Mapping[str, Any], Conversion, dict[str, Any]]: + """Return a lazy component view, its conversion, and the resolved config without mutating caller inputs.""" + + config = dict(config) + conversion = get_conversion(model_class, config) + if checkpoint.keys() == conversion.original_keys and not extract_ema: + return checkpoint, conversion, config + + # These definitions span several already-qualified namespaces. Keep the prefixes and reject unknown entries + # inside the selected component, while leaving sibling components in the bundle alone. + namespaces = () + if model_class == "AutoencoderKLTemporalDecoder" and config.get("original_format", "svd") == "svd": + namespaces = ("conditioner.embedders.3.encoder.", "first_stage_model.") + elif model_class == "UNetFlatConditionModel": + namespaces = ("model.diffusion_model.time_embed.", "model.diffusion_model.unet_text.") + elif model_class == "UNet2DConditionModel" and config.get("original_format") == "versatile_image": + namespaces = ("model.diffusion_model.time_embed.", "model.diffusion_model.unet_image.") + elif model_class == "PaintByExampleImageEncoder": + namespaces = ("cond_stage_model.", "proj_out.", "learnable_vector") + + prefixes = { + "UNet2DConditionModel": ("model.diffusion_model.",), + "ControlNetModel": ("control_model.",), + "AutoencoderKL": ("first_stage_model.", "vae."), + "AutoencoderKLHunyuanImage": ("vae.",), + "AutoencoderKLLTXVideo": ("vae.",), + "AutoencoderKLLTX2Video": ("vae.",), + "AutoencoderKLLTX2Audio": ("audio_vae.",), + "CosmosTransformer3DModel": ("net.", "model.diffusion_model."), + "MiniMaxMusic3Transformer1DModel": ("diffusion_transformer.",), + }.get(model_class, ("model.diffusion_model.",)) + prefix = "" if namespaces else next((p for p in prefixes if any(key.startswith(p) for key in checkpoint)), "") + state = ComponentState(checkpoint, prefix) + if namespaces: + state.keys_to_source = { + key: value for key, value in state.keys_to_source.items() if key.startswith(namespaces) + } + if extract_ema: + if not prefix and not namespaces: + raise ValueError("EMA selection requires a bundled LDM checkpoint with model.diffusion_model keys.") + ema_keys = {key: "model_ema." + "".join((prefix + key).split(".")[1:]) for key in state} + if any(key not in checkpoint for key in ema_keys.values()): + raise ValueError("The checkpoint does not contain a complete EMA copy of the UNet.") + state.keys_to_source = ema_keys + + if model_class == "CosmosTransformer3DModel" and "original_format" not in config: + config["original_format"] = "cosmos1" if any(key.startswith("blocks.block") for key in state) else "cosmos2" + if model_class == "LTX2VideoDiffusionDecoderModel": + if "original_format" not in config and any( + key.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")) for key in state + ): + config["original_format"] = "ltx2_diffusion_decoder_gated" + conversion = get_conversion(model_class, config) + if state.keys() == conversion.diffusers_keys: + return state, Conversion(mapping={key: key for key in state}), config + if model_class == "AutoencoderRAE" and "decoder.trainable_cls_token" not in state: + anchor = state["decoder.decoder_embed.weight"] + state.generated["decoder.trainable_cls_token"] = anchor.new_zeros((1, 1, anchor.shape[0])) + if ( + model_class == "SkyReelsV2Transformer3DModel" + and "img_emb.emb_pos" in conversion.original_keys + and "img_emb.emb_pos" not in state + ): + state.generated["img_emb.emb_pos"] = state["img_emb.proj.0.weight"].new_zeros( + 1, config["pos_embed_seq_len"], config["image_dim"] + ) + if model_class in ("ClapModel", "ClapTextModelWithProjection"): + text_config = config.get("text_config", config) + positions = text_config.get("max_position_embeddings", 514) + for name in ("position_ids", "token_type_ids"): + key = "text_branch.embeddings." + name + if key in conversion.original_keys and key not in state: + state.generated[key] = ( + torch.arange(positions, dtype=torch.long).unsqueeze(0) + if name == "position_ids" + else torch.zeros((1, positions), dtype=torch.long) + ) + if model_class == "LTX2VideoDiffusionDecoderModel": + # Distilled checkpoints can gate only a subset of blocks. An omitted gate is the identity. + for rule in conversion.rules: + gate = rule.original[0] + if gate.endswith((".gate_msa", ".gate_mlp", ".gate_ctx")) and gate not in state: + weight = state[rule.original[1]] + state.generated[gate] = weight.new_ones(weight.shape[0]) + auxiliary = set() + auxiliary_prefixes = () + if model_class in ("AutoencoderKL", "VQModel"): + auxiliary_prefixes = ("loss.", "vocoder.") + elif model_class in ("ClapModel", "ClapTextModelWithProjection"): + auxiliary_prefixes = ("text_transform.", "audio_transform.") + auxiliary = { + key + for key in state + if any(part in key.split(".") for part in ("stft", "logmel_extractor", "tscam_conv", "head", "attn_mask")) + } + elif model_class == "AsymmetricAutoencoderKL": + auxiliary_prefixes = ("loss.", "decoder.up_layers.") + elif model_class == "AutoencoderRAE": + auxiliary = {"decoder.decoder_pos_embed"} + elif model_class == "MiniMaxH3Transformer3DModel": + auxiliary = {"rope.inv_freq"} + elif model_class == "UNet1DModel": + auxiliary = {key for key in state if key.endswith(".kernel")} + elif model_class == "SanaControlNetModel": + auxiliary_prefixes = ("blocks.", "final_layer.") + auxiliary = {"pos_embed", "y_embedder.y_embedding", "logvar_linear.weight", "logvar_linear.bias"} + elif model_class in ( + "SanaTransformer2DModel", + "SanaVideoTransformer3DModel", + "PixArtTransformer2DModel", + "DiTTransformer2DModel", + "Transformer2DModel", + ): + auxiliary = {"pos_embed", "y_embedder.y_embedding"} + if model_class in ("SanaTransformer2DModel", "SanaVideoTransformer3DModel"): + auxiliary.update(("logvar_linear.weight", "logvar_linear.bias")) + elif model_class in ("AutoencoderKLLTXVideo", "AutoencoderKLLTX2Video"): + auxiliary = {"per_channel_statistics.channel", "per_channel_statistics.mean-of-stds"} + elif model_class == "LTX2VideoTransformer3DModel": + auxiliary_prefixes = ("video_embeddings_connector.", "audio_embeddings_connector.") + elif model_class == "LTX2VideoDiffusionDecoderModel": + auxiliary_prefixes = ("encoder.", "decoder.coarse_") + auxiliary = {"per_channel_statistics.channel", "per_channel_statistics.mean-of-stds"} + auxiliary.update(key for key in state if key.startswith("decoder.") and ".coarse_" in key) + elif model_class == "CosmosTransformer3DModel": + auxiliary = { + "logvar.0.freqs", + "logvar.0.phases", + "logvar.1.weight", + "pos_embedder.seq", + "pos_embedder.dim_spatial_range", + "pos_embedder.dim_temporal_range", + "_extra_state", + "accum_video_sample_counter", + "accum_image_sample_counter", + "accum_iteration", + "accum_train_in_hours", + } + elif model_class == "ZImageControlNetModel" and config.get("add_control_noise_refiner") == "control_layers": + auxiliary_prefixes = ("control_noise_refiner.",) + removed = [ + key + for key in state + if key not in conversion.original_keys and (key in auxiliary or key.startswith(auxiliary_prefixes)) + ] + if removed: + logger.info("Excluding known auxiliary state from %s: %s", model_class, sorted(removed)) + for key in removed: + state.keys_to_source.pop(key, None) + state.generated.pop(key, None) + return state, conversion, config diff --git a/src/diffusers/loaders/conversion/chroma.py b/src/diffusers/loaders/conversion/chroma.py new file mode 100644 index 000000000000..15225875e4d5 --- /dev/null +++ b/src/diffusers/loaders/conversion/chroma.py @@ -0,0 +1,72 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Split + + +def chroma_conversion(config): + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + parameters = ("weight", "bias") + modules = [("txt_in", "context_embedder"), ("img_in", "x_embedder"), ("final_layer.linear", "proj_out")] + mapping, rules = ({}, []) + for name in ("in_proj", "out_proj"): + modules.append(("distilled_guidance_layer." + name, "distilled_guidance_layer." + name)) + for i in range(config["approximator_layers"]): + prefix = f"distilled_guidance_layer.layers.{i}" + modules.extend([(prefix + ".in_layer", prefix + ".linear_1"), (prefix + ".out_layer", prefix + ".linear_2")]) + mapping[f"distilled_guidance_layer.norms.{i}.scale"] = f"distilled_guidance_layer.norms.{i}.weight" + for i in range(config["num_layers"]): + old, new = (f"double_blocks.{i}", f"transformer_blocks.{i}") + for modality, attention, mlp, norms in ( + ("img", ("to_q", "to_k", "to_v"), "ff", ("norm_q", "norm_k")), + ("txt", ("add_q_proj", "add_k_proj", "add_v_proj"), "ff_context", ("norm_added_q", "norm_added_k")), + ): + for parameter in parameters: + rules.append( + Rule( + (f"{old}.{modality}_attn.qkv.{parameter}",), + tuple((f"{new}.attn.{name}.{parameter}" for name in attention)), + Split((hidden_size,) * 3), + ) + ) + for source, target in zip(("query_norm", "key_norm"), norms): + mapping[f"{old}.{modality}_attn.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + modules.extend( + [ + (f"{old}.{modality}_mlp.0", f"{new}.{mlp}." + "net.0.proj"), + (f"{old}.{modality}_mlp.2", f"{new}.{mlp}." + "net.2"), + ( + f"{old}.{modality}_attn.proj", + f"{new}.attn." + ("to_out.0" if modality == "img" else "to_add_out"), + ), + ] + ) + for i in range(config["num_single_layers"]): + old, new = (f"single_blocks.{i}", f"single_transformer_blocks.{i}") + for parameter in parameters: + rules.append( + Rule( + (f"{old}.linear1.{parameter}",), + tuple( + (f"{new}.{name}.{parameter}" for name in ("attn.to_q", "attn.to_k", "attn.to_v", "proj_mlp")) + ), + Split((hidden_size,) * 3 + (4 * hidden_size,)), + ) + ) + modules.append((old + ".linear2", new + ".proj_out")) + for source, target in (("query_norm", "norm_q"), ("key_norm", "norm_k")): + mapping[f"{old}.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in parameters}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/chronoedit.py b/src/diffusers/loaders/conversion/chronoedit.py new file mode 100644 index 000000000000..fc4b4a66a439 --- /dev/null +++ b/src/diffusers/loaders/conversion/chronoedit.py @@ -0,0 +1,21 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .wan import wan_conversion + + +def chronoedit_conversion(config): + base = wan_conversion(config) + return Conversion(mapping=base.mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/clap_audio.py b/src/diffusers/loaders/conversion/clap_audio.py new file mode 100644 index 000000000000..ac7de4d59e73 --- /dev/null +++ b/src/diffusers/loaders/conversion/clap_audio.py @@ -0,0 +1,87 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .clap_text import clap_text_conversion +from .core import Conversion, Rule +from .transforms import Split + + +def clap_audio_conversion(config): + full = config["_class_name"] == "ClapModel" + cfg = config["audio_config"] if full else config + projected = config["_class_name"] != "ClapAudioModel" + prefix = "audio_model.audio_encoder" if projected else "audio_encoder" + mapping, modules, rules = ( + {}, + [ + ("audio_branch.patch_embed.proj", f"{prefix}.patch_embed.proj"), + ("audio_branch.norm", f"{prefix}.norm"), + ("audio_branch.bn0", f"{prefix}.batch_norm"), + ], + [], + ) + if cfg["enable_patch_layer_norm"]: + modules.append(("audio_branch.patch_embed.norm", f"{prefix}.patch_embed.norm")) + for p in ("running_mean", "running_var", "num_batches_tracked"): + mapping[f"audio_branch.bn0.{p}"] = f"{prefix}.batch_norm.{p}" + if cfg["enable_fusion"]: + modules.append(("audio_branch.patch_embed.mel_conv2d", f"{prefix}.patch_embed.mel_conv2d")) + for kind, indices, bns in (("local_att", (0, 1, 3, 4), (1, 4)), ("global_att", (1, 2, 4, 5), (2, 5))): + for i in indices: + old, new = ( + f"audio_branch.patch_embed.fusion_model.{kind}.{i}", + f"{prefix}.patch_embed.fusion_model.{kind}.{i}", + ) + modules.append((old, new)) + if i in bns: + mapping.update( + {f"{old}.{p}": f"{new}.{p}" for p in ("running_mean", "running_var", "num_batches_tracked")} + ) + for i, depth in enumerate(cfg["depths"]): + width = cfg["patch_embeds_hidden_size"] * 2**i + for j in range(depth): + old, new = f"audio_branch.layers.{i}.blocks.{j}", f"{prefix}.layers.{i}.blocks.{j}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("norm1", "layernorm_before"), + ("norm2", "layernorm_after"), + ("attn.proj", "attention.output.dense"), + ("mlp.fc1", "intermediate.dense"), + ("mlp.fc2", "output.dense"), + ) + ) + for name in ("relative_position_bias_table", "relative_position_index"): + mapping[f"{old}.attn.{name}"] = f"{new}.attention.self.{name}" + for p in ("weight", "bias") if cfg["qkv_bias"] else ("weight",): + rules.append( + Rule( + (f"{old}.attn.qkv.{p}",), + tuple(f"{new}.attention.self.{part}.{p}" for part in ("query", "key", "value")), + Split((width,) * 3), + ) + ) + if i < len(cfg["depths"]) - 1: + old, new = f"audio_branch.layers.{i}.downsample", f"{prefix}.layers.{i}.downsample" + modules.append((f"{old}.norm", f"{new}.norm")) + mapping[f"{old}.reduction.weight"] = f"{new}.reduction.weight" + if projected: + modules.extend((f"audio_projection.{i}", f"audio_projection.linear{j}") for i, j in ((0, 1), (2, 2))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + if full: + text = clap_text_conversion(config["text_config"]) + mapping.update(text.mapping) + rules.extend(text.rules) + mapping.update({name: name for name in ("logit_scale_a", "logit_scale_t")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/clap_text.py b/src/diffusers/loaders/conversion/clap_text.py new file mode 100644 index 000000000000..c646e05255ea --- /dev/null +++ b/src/diffusers/loaders/conversion/clap_text.py @@ -0,0 +1,51 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def clap_text_conversion(config): + mapping = { + f"text_branch.embeddings.{name}": f"text_model.embeddings.{name}" + for name in ( + "position_ids", + "token_type_ids", + "word_embeddings.weight", + "position_embeddings.weight", + "token_type_embeddings.weight", + ) + } + modules = [ + ("text_branch.embeddings.LayerNorm", "text_model.embeddings.LayerNorm"), + ("text_branch.pooler.dense", "text_model.pooler.dense"), + ("text_projection.0", "text_projection.linear1"), + ("text_projection.2", "text_projection.linear2"), + ] + for i in range(config["num_hidden_layers"]): + modules.extend( + (f"text_branch.encoder.layer.{i}.{name}", f"text_model.encoder.layer.{i}.{name}") + for name in ( + "attention.self.query", + "attention.self.key", + "attention.self.value", + "attention.output.dense", + "attention.output.LayerNorm", + "intermediate.dense", + "output.dense", + "output.LayerNorm", + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/clip.py b/src/diffusers/loaders/conversion/clip.py new file mode 100644 index 000000000000..c44787307d89 --- /dev/null +++ b/src/diffusers/loaders/conversion/clip.py @@ -0,0 +1,104 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from .core import Conversion, Rule +from .transforms import Permute, Split, WithConstants + + +def clip_conversion(config): + original_format = config.get( + "original_format", "openclip" if config["_class_name"] == "CLIPTextModelWithProjection" else "clip" + ) + if original_format == "openclip": + return openclip_conversion(config) + if original_format != "clip": + raise ValueError("CLIP original_format must be 'clip' or 'openclip'.") + modules = ["text_model.final_layer_norm"] + keys = ["text_model.embeddings.token_embedding.weight", "text_model.embeddings.position_embedding.weight"] + for i in range(config["num_hidden_layers"]): + prefix = f"text_model.encoder.layers.{i}" + modules.extend( + f"{prefix}.{name}" + for name in ( + "layer_norm1", + "layer_norm2", + "mlp.fc1", + "mlp.fc2", + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.out_proj", + ) + ) + if config["_class_name"] == "CLIPTextModelWithProjection": + keys.append("text_projection.weight") + if config["_class_name"] == "ContextCLIPTextModel": + keys.append("text_model.embeddings.position_ids") + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + prefix = "" if _uses_flat_text_model(config) else "text_model." + return Conversion( + mapping={ + key: prefix + key.removeprefix("text_model.") if key.startswith("text_model.") else key for key in keys + } + ) + + +def openclip_conversion(config): + prefix = "" if _uses_flat_text_model(config) else "text_model." + mapping = { + "token_embedding.weight": prefix + "embeddings.token_embedding.weight", + "positional_embedding": prefix + "embeddings.position_embedding.weight", + } + modules, rules = [("ln_final", prefix + "final_layer_norm")], [] + for i in range(config["num_hidden_layers"]): + old, new = f"transformer.resblocks.{i}", f"{prefix}encoder.layers.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ln_1", "layer_norm1"), + ("ln_2", "layer_norm2"), + ("mlp.c_fc", "mlp.fc1"), + ("mlp.c_proj", "mlp.fc2"), + ("attn.out_proj", "self_attn.out_proj"), + ) + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.attn.in_proj_{p}",), + tuple(f"{new}.self_attn.{part}_proj.{p}" for part in ("q", "k", "v")), + Split((config["hidden_size"],) * 3), + ) + ) + if config["_class_name"] == "CLIPTextModelWithProjection": + rules.append(Rule(("text_projection",), ("text_projection.weight",), Permute((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + if config["_class_name"] == "ContextCLIPTextModel": + target = mapping.pop("token_embedding.weight") + rules.append( + Rule( + ("token_embedding.weight",), + (target, prefix + "embeddings.position_ids"), + WithConstants((torch.arange(config["max_position_embeddings"]).unsqueeze(0),)), + ) + ) + return Conversion(mapping=mapping, rules=tuple(rules)) + + +def _uses_flat_text_model(config): + version = (config.get("transformers_version") or "4").split(".", 1)[0] + return config["_class_name"] == "CLIPTextModel" and int(version) >= 5 diff --git a/src/diffusers/loaders/conversion/clip_vision.py b/src/diffusers/loaders/conversion/clip_vision.py new file mode 100644 index 000000000000..ea84eaa21fad --- /dev/null +++ b/src/diffusers/loaders/conversion/clip_vision.py @@ -0,0 +1,67 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Permute, Split + + +def clip_vision_conversion(config): + """Convert CLIP/HF or OpenCLIP visual tower weights, including optional projection.""" + projected = config["_class_name"] == "CLIPVisionModelWithProjection" + prefix = ( + "" + if not projected and int((config.get("transformers_version") or "4").split(".")[0]) >= 5 + else "vision_model." + ) + original_format = config.get("original_format", "clip") + if original_format not in ("clip", "openclip"): + raise ValueError("CLIP vision original_format must be clip or openclip.") + modules = [("ln_pre", "pre_layrnorm"), ("ln_post", "post_layernorm")] + pairs = [ + ("class_embedding", "embeddings.class_embedding"), + ("positional_embedding", "embeddings.position_embedding.weight"), + ("conv1.weight", "embeddings.patch_embedding.weight"), + ] + rules = [] + for i in range(config["num_hidden_layers"]): + old, new = f"transformer.resblocks.{i}", f"encoder.layers.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ln_1", "layer_norm1"), + ("ln_2", "layer_norm2"), + ("mlp.c_fc", "mlp.fc1"), + ("mlp.c_proj", "mlp.fc2"), + ("attn.out_proj", "self_attn.out_proj"), + ) + ) + for p in ("weight", "bias"): + if original_format == "openclip": + rules.append( + Rule( + (f"{old}.attn.in_proj_{p}",), + tuple(f"{prefix}{new}.self_attn.{part}_proj.{p}" for part in ("q", "k", "v")), + Split((config["hidden_size"],) * 3), + ) + ) + else: + pairs.extend(("", f"{new}.self_attn.{part}_proj.{p}") for part in ("q", "k", "v")) + pairs.extend((f"{old}.{p}", f"{new}.{p}") for old, new in modules for p in ("weight", "bias")) + mapping = {(old if original_format == "openclip" else "vision_model." + new): prefix + new for old, new in pairs} + if projected: + if original_format == "openclip": + rules.append(Rule(("proj",), ("visual_projection.weight",), Permute((1, 0)))) + else: + mapping["visual_projection.weight"] = "visual_projection.weight" + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/cogvideox.py b/src/diffusers/loaders/conversion/cogvideox.py new file mode 100644 index 000000000000..a11172dfa32f --- /dev/null +++ b/src/diffusers/loaders/conversion/cogvideox.py @@ -0,0 +1,239 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from ...utils import logging +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +logger = logging.get_logger(__name__) + +COGVIDEOX_TRANSFORMER_PREFIX = "model.diffusion_model." + + +def cogvideox_fixed_position_embedding(config, dtype): + """Recreate the non-persistent Diffusers position buffer saved by the original SAT runtime.""" + from ...models.embeddings import get_3d_sincos_pos_embed + + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + width = config["sample_width"] // config["patch_size"] + height = config["sample_height"] // config["patch_size"] + frames = (config["sample_frames"] - 1) // config["temporal_compression_ratio"] + 1 + positions = get_3d_sincos_pos_embed( + hidden_size, + (width, height), + frames, + config["spatial_interpolation_scale"], + config["temporal_interpolation_scale"], + output_type="pt", + ).flatten(0, 1) + text_positions = positions.new_zeros(config["max_text_seq_length"], hidden_size) + return torch.cat((text_positions, positions)).unsqueeze(0).to(dtype=dtype) + + +@dataclass(frozen=True) +class CogVideoXAdaLN: + """Separate SAT's interleaved attention/MLP modulation into the two Diffusers norms.""" + + hidden_size: int + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + tensors = ReorderChunks((0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11)).forward(tensors) + return Split((6 * self.hidden_size, 6 * self.hidden_size)).forward(tensors) + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + tensors = Split((6 * self.hidden_size, 6 * self.hidden_size)).inverse(tensors) + return ReorderChunks((0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11)).inverse(tensors) + + +def cogvideox_transformer_conversion(config: Mapping[str, Any]) -> Conversion: + """Build SAT <-> Diffusers transformer weight rules from a `CogVideoXTransformer3DModel.config`. + + The original keys are relative to `model.diffusion_model.`. Non-learned positional state and training-only + checkpoint entries are outside this component conversion. No prior import or model allocation is required. + """ + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + modules = [ + ("time_embed.0", "time_embedding.linear_1"), + ("time_embed.2", "time_embedding.linear_2"), + ("mixins.patch_embed.text_proj", "patch_embed.text_proj"), + ("mixins.final_layer.linear", "proj_out"), + ("mixins.final_layer.adaLN_modulation.1", "norm_out.linear"), + ] + if config.get("ofs_embed_dim") is not None: + modules.extend([("ofs_embed.0", "ofs_embedding.linear_1"), ("ofs_embed.2", "ofs_embedding.linear_2")]) + if config.get("norm_elementwise_affine", True): + modules.extend( + [("transformer.final_layernorm", "norm_final"), ("mixins.final_layer.norm_final", "norm_out.norm")] + ) + + mapping = {"mixins.patch_embed.proj.weight": "patch_embed.proj.weight"} + # The 1.5 linear patch embed has a bias even when patch_bias=False (that flag configures the 1.0 convolution). + if config.get("patch_bias", True) or config.get("patch_size_t") is not None: + mapping["mixins.patch_embed.proj.bias"] = "patch_embed.proj.bias" + if config.get("use_learned_positional_embeddings", False): + mapping["mixins.pos_embed.pos_embedding"] = "patch_embed.pos_embedding" + + rules = [] + for i in range(config["num_layers"]): + original = f"transformer.layers.{i}" + diffusers = f"transformer_blocks.{i}" + modules.extend( + [ + (f"{original}.attention.dense", f"{diffusers}.attn1.to_out.0"), + (f"{original}.mlp.dense_h_to_4h", f"{diffusers}.ff.net.0.proj"), + (f"{original}.mlp.dense_4h_to_h", f"{diffusers}.ff.net.2"), + (f"mixins.adaln_layer.query_layernorm_list.{i}", f"{diffusers}.attn1.norm_q"), + (f"mixins.adaln_layer.key_layernorm_list.{i}", f"{diffusers}.attn1.norm_k"), + ] + ) + if config.get("norm_elementwise_affine", True): + modules.extend( + [ + (f"{original}.input_layernorm", f"{diffusers}.norm1.norm"), + (f"{original}.post_attention_layernorm", f"{diffusers}.norm2.norm"), + ] + ) + for parameter in ("weight", "bias"): + if parameter == "weight" or config.get("attention_bias", True): + rules.append( + Rule( + (f"{original}.attention.query_key_value.{parameter}",), + tuple(f"{diffusers}.attn1.to_{name}.{parameter}" for name in "qkv"), + Split((hidden_size, hidden_size, hidden_size)), + ) + ) + rules.append( + Rule( + (f"mixins.adaln_layer.adaLN_modulations.{i}.1.{parameter}",), + (f"{diffusers}.norm1.linear.{parameter}", f"{diffusers}.norm2.linear.{parameter}"), + CogVideoXAdaLN(hidden_size), + ) + ) + + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping, tuple(rules)) + + +def cogvideox_vae_conversion(config: Mapping[str, Any]) -> Conversion: + """Build SAT <-> Diffusers VAE weight rules from an `AutoencoderKLCogVideoX.config`.""" + channels = config["block_out_channels"] + layers = config["layers_per_block"] + modules = [] + resnets = [] + for component in ("encoder", "decoder"): + modules.extend((f"{component}.{name}", f"{component}.{name}") for name in ("conv_in.conv", "conv_out.conv")) + for i in range(2): + resnets.append( + (f"{component}.mid.block_{i + 1}", f"{component}.mid_block.resnets.{i}", channels[-1], channels[-1]) + ) + block_channels = channels if component == "encoder" else tuple(reversed(channels)) + previous_channels = block_channels[0] + for i, out_channels in enumerate(block_channels): + if component == "encoder": + original, diffusers = f"encoder.down.{i}", f"encoder.down_blocks.{i}" + sample_original, sample_diffusers = "downsample", "downsamplers.0" + else: + original, diffusers = f"decoder.up.{len(channels) - 1 - i}", f"decoder.up_blocks.{i}" + sample_original, sample_diffusers = "upsample", "upsamplers.0" + for j in range(layers + (component == "decoder")): + resnets.append((f"{original}.block.{j}", f"{diffusers}.resnets.{j}", previous_channels, out_channels)) + previous_channels = out_channels + if i < len(channels) - 1: + modules.append((f"{original}.{sample_original}.conv", f"{diffusers}.{sample_diffusers}.conv")) + + for original, diffusers, in_channels, out_channels in resnets: + modules.extend((f"{original}.{name}", f"{diffusers}.{name}") for name in ("conv1.conv", "conv2.conv")) + norm_names = ("norm1", "norm2") + if original.startswith("decoder."): + norm_names = tuple( + f"{norm}.{name}" for norm in norm_names for name in ("norm_layer", "conv_y.conv", "conv_b.conv") + ) + modules.extend((f"{original}.{name}", f"{diffusers}.{name}") for name in norm_names) + if in_channels != out_channels: + modules.append((f"{original}.nin_shortcut", f"{diffusers}.conv_shortcut")) + + modules.append(("encoder.norm_out", "encoder.norm_out")) + modules.extend( + (f"decoder.norm_out.{name}", f"decoder.norm_out.{name}") + for name in ("norm_layer", "conv_y.conv", "conv_b.conv") + ) + for name in ("quant_conv", "post_quant_conv"): + if config.get(f"use_{name}", False): + modules.append((name, name)) + return Conversion({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + + +def unwrap_cogvideox_checkpoint(checkpoint: Mapping[str, Any]) -> Mapping[str, Any]: + """Select the tensor dictionary from the nested wrappers used by SAT and converted component checkpoints.""" + while True: + wrappers = [key for key in ("model", "module", "state_dict") if isinstance(checkpoint.get(key), Mapping)] + if not wrappers: + return checkpoint + if len(wrappers) != 1: + raise ValueError(f"Ambiguous CogVideoX checkpoint wrappers: {wrappers}.") + checkpoint = checkpoint[wrappers[0]] + + +def convert_cogvideox_transformer_checkpoint_to_diffusers(checkpoint, config): + """Import a SAT transformer component, selecting its prefix and reporting omitted auxiliary keys.""" + checkpoint = unwrap_cogvideox_checkpoint(checkpoint) + conversion = cogvideox_transformer_conversion(config) + if any(key.startswith(COGVIDEOX_TRANSFORMER_PREFIX) for key in checkpoint): + if conversion.original_keys.intersection(checkpoint): + raise ValueError("CogVideoX checkpoint mixes prefixed and unprefixed transformer weights.") + checkpoint = { + key.removeprefix(COGVIDEOX_TRANSFORMER_PREFIX): tensor + for key, tensor in checkpoint.items() + if key.startswith(COGVIDEOX_TRANSFORMER_PREFIX) + } + else: + checkpoint = dict(checkpoint) + + omitted = sorted(cogvideox_transformer_auxiliary_keys(config).intersection(checkpoint)) + for key in omitted: + checkpoint.pop(key) + if omitted: + logger.info(f"Omitting CogVideoX auxiliary embeddings computed or unused by Diffusers: {omitted}") + return conversion.to_diffusers(checkpoint) + + +def cogvideox_transformer_auxiliary_keys(config): + """Known SAT embeddings computed or unused by the Diffusers component.""" + auxiliary_keys = { + "transformer.embed_tokens.weight", + "transformer.word_embeddings.weight", + "transformer.position_embeddings.weight", + "mixins.pos_embed.freqs_sin", + "mixins.pos_embed.freqs_cos", + } + if not config.get("use_learned_positional_embeddings", False): + auxiliary_keys.add("mixins.pos_embed.pos_embedding") + return auxiliary_keys + + +def convert_cogvideox_vae_checkpoint_to_diffusers(checkpoint, config): + """Import SAT VAE weights, excluding the training loss module without changing the input checkpoint.""" + checkpoint = unwrap_cogvideox_checkpoint(checkpoint) + omitted = sorted(key for key in checkpoint if key.startswith("loss.")) + if omitted: + logger.info(f"Omitting CogVideoX VAE training loss state: {omitted}") + checkpoint = {key: tensor for key, tensor in checkpoint.items() if not key.startswith("loss.")} + return cogvideox_vae_conversion(config).to_diffusers(checkpoint) diff --git a/src/diffusers/loaders/conversion/cogview3plus.py b/src/diffusers/loaders/conversion/cogview3plus.py new file mode 100644 index 000000000000..9f7af1a41f75 --- /dev/null +++ b/src/diffusers/loaders/conversion/cogview3plus.py @@ -0,0 +1,55 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def cogview3plus_conversion(config): + modules = [ + ("mixins.patch_embed.proj", "patch_embed.proj"), + ("mixins.patch_embed.text_proj", "patch_embed.text_proj"), + ("mixins.final_layer.linear", "proj_out"), + ] + modules.extend( + (f"{source}.{i}", f"time_condition_embed.{target}.linear_{j}") + for source, target in (("time_embed", "timestep_embedder"), ("label_emb.0", "condition_embedder")) + for i, j in ((0, 1), (2, 2)) + ) + rules = [] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + for i in range(config["num_layers"]): + old, new = f"transformer.layers.{i}", f"transformer_blocks.{i}" + modules.append((f"mixins.adaln.adaln_modules.{i}.1", new + ".norm1.linear")) + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("attention.dense", "attn1.to_out.0"), + ("mlp.dense_h_to_4h", "ff.net.0.proj"), + ("mlp.dense_4h_to_h", "ff.net.2"), + ) + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.attention.query_key_value.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + for p in ("weight", "bias"): + rules.append(Rule((f"mixins.final_layer.adaln.1.{p}",), (f"norm_out.linear.{p}",), ReorderChunks((1, 0)))) + mapping = {f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")} + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/cogview4.py b/src/diffusers/loaders/conversion/cogview4.py new file mode 100644 index 000000000000..50e17784c8c2 --- /dev/null +++ b/src/diffusers/loaders/conversion/cogview4.py @@ -0,0 +1,27 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .cogview3plus import cogview3plus_conversion +from .core import Conversion + + +def cogview4_conversion(config): + if config.get("original_format") == "megatron": + from .cogview4_megatron import cogview4_megatron_conversion + + return cogview4_megatron_conversion(config) + # The SAT parameter layout is shared with CogView3Plus; the configuration selects the dimensions and depth. + base = cogview3plus_conversion(config) + return Conversion(mapping=base.mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/cogview4_megatron.py b/src/diffusers/loaders/conversion/cogview4_megatron.py new file mode 100644 index 000000000000..5962d84b54cb --- /dev/null +++ b/src/diffusers/loaders/conversion/cogview4_megatron.py @@ -0,0 +1,63 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .cogview3plus import cogview3plus_conversion +from .core import Conversion, Rule +from .transforms import Chain, Permute, Reshape + + +def cogview4_megatron_conversion(config): + base = cogview3plus_conversion(config) + replacements = { + "mixins.patch_embed.proj": "encoder_expand_linear", + "mixins.patch_embed.text_proj": "text_projector", + "mixins.final_layer.linear": "output_projector", + "time_embed.": "time_embedding.time_embed.", + "label_emb.0.": "label_embedding.label_embed.", + "mixins.final_layer.adaln.1": "adaln_final", + } + for i in range(config["num_layers"]): + replacements[f"mixins.adaln.adaln_modules.{i}.1"] = f"decoder.layers.{i}.adaln" + for old, new in ( + ("attention.query_key_value", "self_attention.linear_qkv"), + ("attention.dense", "self_attention.linear_proj"), + ("mlp.dense_h_to_4h", "mlp.linear_fc1"), + ("mlp.dense_4h_to_h", "mlp.linear_fc2"), + ): + replacements[f"transformer.layers.{i}.{old}"] = f"decoder.layers.{i}.{new}" + + def rename(key): + for old, new in replacements.items(): + if key.startswith(old): + return new + key[len(old) :] + raise ValueError(f"Unknown CogView Megatron source key {key}.") + + mapping = {rename(old): new for old, new in base.mapping.items()} + rules = [] + heads, dim = config["num_attention_heads"], config["attention_head_dim"] + width = heads * dim + for rule in base.rules: + transform = rule.transform + if ".query_key_value." in rule.original[0]: + trailing = (width,) if rule.original[0].endswith("weight") else () + transform = Chain( + ( + Reshape((3 * width,) + trailing, (heads, 3, dim) + trailing), + Permute((1, 0, 2, 3) if trailing else (1, 0, 2)), + Reshape((3, heads, dim) + trailing, (3 * width,) + trailing), + transform, + ) + ) + rules.append(Rule(tuple(rename(key) for key in rule.original), rule.diffusers, transform)) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/configs/__init__.py b/src/diffusers/loaders/conversion/configs/__init__.py new file mode 100644 index 000000000000..9783c7a926f4 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/__init__.py @@ -0,0 +1,61 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reusable original configuration preparation, separate from the tensor conversion definitions.""" + +import copy +import importlib +import json +from collections.abc import Mapping +from pathlib import Path + + +def list_config_presets(): + directory = Path(__file__).resolve().parent + catalog = json.loads((directory / "catalog.json").read_text(encoding="utf-8")) + return sorted( + [path.stem for path in (directory.parent / "presets").glob("*.json")] + + [f"{module}.{name}" for module, names in catalog.items() for name in names] + ) + + +def get_config_preset(name, *, arguments=None, component=None): + """Load a built-in config or call a named original-config helper with explicit keyword arguments. + + `component` selects an entry in a collection, such as Cosmos' `TRANSFORMER_CONFIGS`. Helpers returning original + download metadata expose their `diffusers_config` field automatically. Config preparation never authors tensor + mappings; pass the resulting dictionary to `get_conversion` or the checkpoint CLI. + """ + if name not in list_config_presets(): + raise ValueError(f"Unknown config preset {name!r}; use --list-presets to see available names.") + directory = Path(__file__).resolve().parent + arguments = dict(arguments or {}) + if "." in name and not (directory.parent / "presets" / (name + ".json")).is_file(): + module, attribute = name.split(".", 1) + value = getattr(importlib.import_module(f"{__name__}.{module}"), attribute) + if callable(value): + value = value(**arguments) + elif arguments: + raise ValueError("Dictionary presets do not take arguments; use component to select a variant.") + else: + if arguments: + raise ValueError("JSON presets do not take arguments.") + value = json.loads((directory.parent / "presets" / (name + ".json")).read_text(encoding="utf-8")) + if component is not None: + value = value[component] + if hasattr(value, "to_dict"): + value = value.to_dict() + if not isinstance(value, Mapping): + raise ValueError("The selected preset did not produce a configuration dictionary.") + return copy.deepcopy(dict(value.get("diffusers_config", value))) diff --git a/src/diffusers/loaders/conversion/configs/anima.py b/src/diffusers/loaders/conversion/configs/anima.py new file mode 100644 index 000000000000..7710e1cc8f00 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/anima.py @@ -0,0 +1,64 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the anima assembly recipe.""" + +from typing import Any + +import torch +from transformers import Qwen3Config + + +def infer_text_conditioner_config(state_dict: dict[str, torch.Tensor]) -> dict[str, Any]: + model_dim = state_dict["blocks.0.self_attn.q_proj.weight"].shape[0] + source_dim = state_dict["blocks.0.cross_attn.k_proj.weight"].shape[1] + target_vocab_size, target_dim = state_dict["embed.weight"].shape + attention_head_dim = state_dict["blocks.0.self_attn.q_norm.weight"].shape[0] + num_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("blocks.")) + + return { + "source_dim": source_dim, + "target_dim": target_dim, + "model_dim": model_dim, + "num_layers": num_layers, + "num_attention_heads": model_dim // attention_head_dim, + "target_vocab_size": target_vocab_size, + } + + +def infer_qwen3_config(state_dict: dict[str, torch.Tensor]) -> Qwen3Config: + vocab_size, hidden_size = state_dict["embed_tokens.weight"].shape + intermediate_size = state_dict["layers.0.mlp.gate_proj.weight"].shape[0] + num_hidden_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("layers.")) + head_dim = state_dict["layers.0.self_attn.q_norm.weight"].shape[0] + num_attention_heads = state_dict["layers.0.self_attn.q_proj.weight"].shape[0] // head_dim + num_key_value_heads = state_dict["layers.0.self_attn.k_proj.weight"].shape[0] // head_dim + + return Qwen3Config( + vocab_size=vocab_size, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_hidden_layers=num_hidden_layers, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + max_position_embeddings=32768, + rms_norm_eps=1e-6, + rope_theta=1000000.0, + head_dim=head_dim, + attention_bias=False, + tie_word_embeddings=False, + ) + + +__all__ = ["infer_qwen3_config", "infer_text_conditioner_config"] diff --git a/src/diffusers/loaders/conversion/configs/audioldm.py b/src/diffusers/loaders/conversion/configs/audioldm.py new file mode 100644 index 000000000000..372ef595f9d5 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/audioldm.py @@ -0,0 +1,178 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the audioldm assembly recipe.""" + + +def create_unet_diffusers_config(original_config, image_size: int): + """ + Creates a UNet config for diffusers based on the config of the original AudioLDM model. + """ + unet_params = original_config["model"]["params"]["unet_config"]["params"] + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + + block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) + + cross_attention_dim = ( + unet_params["cross_attention_dim"] if "cross_attention_dim" in unet_params else block_out_channels + ) + + class_embed_type = "simple_projection" if "extra_film_condition_dim" in unet_params else None + projection_class_embeddings_input_dim = ( + unet_params["extra_film_condition_dim"] if "extra_film_condition_dim" in unet_params else None + ) + class_embeddings_concat = unet_params["extra_film_use_concat"] if "extra_film_use_concat" in unet_params else None + + config = { + "sample_size": image_size // vae_scale_factor, + "in_channels": unet_params["in_channels"], + "out_channels": unet_params["out_channels"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params["num_res_blocks"], + "cross_attention_dim": cross_attention_dim, + "class_embed_type": class_embed_type, + "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, + "class_embeddings_concat": class_embeddings_concat, + } + + return config + + +def create_vae_diffusers_config(original_config, checkpoint, image_size: int): + """ + Creates a VAE config for diffusers based on the config of the original AudioLDM model. Compared to the original + Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. + """ + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] + + block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 + + config = { + "sample_size": image_size, + "in_channels": vae_params["in_channels"], + "out_channels": vae_params["out_ch"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params["z_channels"], + "layers_per_block": vae_params["num_res_blocks"], + "scaling_factor": float(scaling_factor), + } + return config + + +def create_transformers_vocoder_config(original_config): + """ + Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. + """ + vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] + + config = { + "model_in_dim": vocoder_params["num_mels"], + "sampling_rate": vocoder_params["sampling_rate"], + "upsample_initial_channel": vocoder_params["upsample_initial_channel"], + "upsample_rates": list(vocoder_params["upsample_rates"]), + "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), + "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), + "resblock_dilation_sizes": [ + list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] + ], + "normalize_before": False, + } + + return config + + +DEFAULT_CONFIG = { + "model": { + "params": { + "linear_start": 0.0015, + "linear_end": 0.0195, + "timesteps": 1000, + "channels": 8, + "scale_by_std": True, + "unet_config": { + "target": "audioldm.latent_diffusion.openaimodel.UNetModel", + "params": { + "extra_film_condition_dim": 512, + "extra_film_use_concat": True, + "in_channels": 8, + "out_channels": 8, + "model_channels": 128, + "attention_resolutions": [8, 4, 2], + "num_res_blocks": 2, + "channel_mult": [1, 2, 3, 5], + "num_head_channels": 32, + }, + }, + "first_stage_config": { + "target": "audioldm.variational_autoencoder.autoencoder.AutoencoderKL", + "params": { + "embed_dim": 8, + "ddconfig": { + "z_channels": 8, + "resolution": 256, + "in_channels": 1, + "out_ch": 1, + "ch": 128, + "ch_mult": [1, 2, 4], + "num_res_blocks": 2, + }, + }, + }, + "vocoder_config": { + "target": "audioldm.first_stage_model.vocoder", + "params": { + "upsample_rates": [5, 4, 2, 2, 2], + "upsample_kernel_sizes": [16, 16, 8, 4, 4], + "upsample_initial_channel": 1024, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "num_mels": 64, + "sampling_rate": 16000, + }, + }, + }, + }, +} + +__all__ = [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config", +] diff --git a/src/diffusers/loaders/conversion/configs/audioldm2.py b/src/diffusers/loaders/conversion/configs/audioldm2.py new file mode 100644 index 000000000000..d7495fec7d80 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/audioldm2.py @@ -0,0 +1,180 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the audioldm2 assembly recipe.""" + + +def create_unet_diffusers_config(original_config, image_size: int): + """ + Creates a UNet config for diffusers based on the config of the original AudioLDM2 model. + """ + unet_params = original_config["model"]["params"]["unet_config"]["params"] + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + + block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) + + cross_attention_dim = list(unet_params["context_dim"]) if "context_dim" in unet_params else block_out_channels + if len(cross_attention_dim) > 1: + # require two or more cross-attention layers per-block, each of different dimension + cross_attention_dim = [cross_attention_dim for _ in range(len(block_out_channels))] + + config = { + "sample_size": image_size // vae_scale_factor, + "in_channels": unet_params["in_channels"], + "out_channels": unet_params["out_channels"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params["num_res_blocks"], + "transformer_layers_per_block": unet_params["transformer_depth"], + "cross_attention_dim": tuple(cross_attention_dim), + } + + return config + + +def create_vae_diffusers_config(original_config, checkpoint, image_size: int): + """ + Creates a VAE config for diffusers based on the config of the original AudioLDM2 model. Compared to the original + Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. + """ + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] + + block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 + + config = { + "sample_size": image_size, + "in_channels": vae_params["in_channels"], + "out_channels": vae_params["out_ch"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params["z_channels"], + "layers_per_block": vae_params["num_res_blocks"], + "scaling_factor": float(scaling_factor), + } + return config + + +def create_transformers_vocoder_config(original_config): + """ + Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. + """ + vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] + + config = { + "model_in_dim": vocoder_params["num_mels"], + "sampling_rate": vocoder_params["sampling_rate"], + "upsample_initial_channel": vocoder_params["upsample_initial_channel"], + "upsample_rates": list(vocoder_params["upsample_rates"]), + "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), + "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), + "resblock_dilation_sizes": [ + list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] + ], + "normalize_before": False, + } + + return config + + +DEFAULT_CONFIG = { + "model": { + "params": { + "linear_start": 0.0015, + "linear_end": 0.0195, + "timesteps": 1000, + "channels": 8, + "scale_by_std": True, + "unet_config": { + "target": "audioldm2.latent_diffusion.openaimodel.UNetModel", + "params": { + "context_dim": [None, 768, 1024], + "in_channels": 8, + "out_channels": 8, + "model_channels": 128, + "attention_resolutions": [8, 4, 2], + "num_res_blocks": 2, + "channel_mult": [1, 2, 3, 5], + "num_head_channels": 32, + "transformer_depth": 1, + }, + }, + "first_stage_config": { + "target": "audioldm2.variational_autoencoder.autoencoder.AutoencoderKL", + "params": { + "embed_dim": 8, + "ddconfig": { + "z_channels": 8, + "resolution": 256, + "in_channels": 1, + "out_ch": 1, + "ch": 128, + "ch_mult": [1, 2, 4], + "num_res_blocks": 2, + }, + }, + }, + "cond_stage_config": { + "crossattn_audiomae_generated": { + "target": "audioldm2.latent_diffusion.modules.encoders.modules.SequenceGenAudioMAECond", + "params": { + "sequence_gen_length": 8, + "sequence_input_embed_dim": [512, 1024], + }, + } + }, + "vocoder_config": { + "target": "audioldm2.first_stage_model.vocoder", + "params": { + "upsample_rates": [5, 4, 2, 2, 2], + "upsample_kernel_sizes": [16, 16, 8, 4, 4], + "upsample_initial_channel": 1024, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "num_mels": 64, + "sampling_rate": 16000, + }, + }, + }, + }, +} + +__all__ = [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config", +] diff --git a/src/diffusers/loaders/conversion/configs/blip_diffusion.py b/src/diffusers/loaders/conversion/configs/blip_diffusion.py new file mode 100644 index 000000000000..44b427e27398 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/blip_diffusion.py @@ -0,0 +1,35 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the blip_diffusion assembly recipe.""" + +BLIP2_CONFIG = { + "vision_config": { + "hidden_size": 1024, + "num_hidden_layers": 23, + "num_attention_heads": 16, + "image_size": 224, + "patch_size": 14, + "intermediate_size": 4096, + "hidden_act": "quick_gelu", + }, + "qformer_config": { + "cross_attention_frequency": 1, + "encoder_hidden_size": 1024, + "vocab_size": 30523, + }, + "num_query_tokens": 16, +} + +__all__ = ["BLIP2_CONFIG"] diff --git a/src/diffusers/loaders/conversion/configs/catalog.json b/src/diffusers/loaders/conversion/configs/catalog.json new file mode 100644 index 000000000000..090bdeea5d31 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/catalog.json @@ -0,0 +1,145 @@ +{ + "anima": [ + "infer_qwen3_config", + "infer_text_conditioner_config" + ], + "audioldm": [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config" + ], + "audioldm2": [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config" + ], + "blip_diffusion": [ + "BLIP2_CONFIG" + ], + "consistency": [ + "CD_SCHEDULER_CONFIG", + "CT_IMAGENET_64_SCHEDULER_CONFIG", + "CT_LSUN_256_SCHEDULER_CONFIG", + "IMAGENET_64_UNET_CONFIG", + "LSUN_256_UNET_CONFIG", + "TEST_UNET_CONFIG" + ], + "cosmos": [ + "CONTROLNET_CONFIGS", + "TRANSFORMER_CONFIGS", + "VAE_CONFIGS" + ], + "deepfloyd_if": [ + "create_unet_diffusers_config", + "superres_create_unet_diffusers_config" + ], + "flux2": [ + "get_flux2_transformer_config" + ], + "gligen": [ + "create_unet_config", + "create_vae_config" + ], + "hunyuan_video": [ + "TRANSFORMER_CONFIGS" + ], + "hunyuan_video15": [ + "GUIDANCE_CONFIGS", + "SCHEDULER_CONFIGS", + "TRANSFORMER_CONFIGS" + ], + "joyimage": [ + "TRANSFORMER_CONFIG" + ], + "kandinsky": [ + "INPAINT_UNET_CONFIG", + "MOVQ_CONFIG", + "PRIOR_CONFIG", + "UNET_CONFIG" + ], + "ltx": [ + "get_spatial_latent_upsampler_config", + "get_transformer_config", + "get_vae_config" + ], + "ltx2": [ + "get_ltx2_audio_vae_config", + "get_ltx2_connectors_config", + "get_ltx2_diffusion_video_vae_config", + "get_ltx2_spatial_latent_upsampler_config", + "get_ltx2_temporal_latent_upsampler_config", + "get_ltx2_transformer_config", + "get_ltx2_video_vae_config", + "get_ltx2_vocoder_config" + ], + "minimax_h3": [ + "MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG", + "MINIMAX_H3_TEST_TRANSFORMER_CONFIG", + "MINIMAX_H3_TEST_VIDEO_VAE_CONFIG", + "MINIMAX_H3_TRANSFORMER_CONFIG", + "MINIMAX_H3_VIDEO_VAE_CONFIG", + "get_audio_vae_config" + ], + "musicldm": [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config" + ], + "prx": [ + "build_config", + "create_scheduler_config" + ], + "rae": [ + "DECODER_CONFIGS" + ], + "shap_e": [ + "PRIOR_CONFIG", + "PRIOR_IMAGE_CONFIG", + "RENDERER_CONFIG" + ], + "skyreels_v2": [ + "get_transformer_config" + ], + "stable_audio3": [ + "_infer_dit_config", + "_infer_duration_embedder_config", + "_infer_vae_config" + ], + "unclip": [ + "DECODER_CONFIG", + "PRIOR_CONFIG", + "SUPER_RES_UNET_FIRST_STEPS_CONFIG", + "SUPER_RES_UNET_LAST_STEP_CONFIG" + ], + "unidiffuser": [ + "SCHEDULER_CONFIG", + "create_text_decoder_config", + "create_text_decoder_config_big", + "create_text_decoder_config_test", + "create_unidiffuser_unet_config", + "create_unidiffuser_unet_config_big", + "create_unidiffuser_unet_config_test", + "create_vae_diffusers_config", + "create_vae_diffusers_config_big", + "create_vae_diffusers_config_test" + ], + "versatile_diffusion": [ + "AUTOENCODER_CONFIG", + "IMAGE_UNET_CONFIG", + "SCHEDULER_CONFIG", + "TEXT_UNET_CONFIG", + "create_image_unet_diffusers_config", + "create_text_unet_diffusers_config", + "create_vae_diffusers_config" + ], + "wan": [ + "get_transformer_config" + ], + "zero123": [ + "create_unet_diffusers_config", + "create_vae_diffusers_config" + ] +} diff --git a/src/diffusers/loaders/conversion/configs/consistency.py b/src/diffusers/loaders/conversion/configs/consistency.py new file mode 100644 index 000000000000..faa44d6847d9 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/consistency.py @@ -0,0 +1,119 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the consistency assembly recipe.""" + +TEST_UNET_CONFIG = { + "sample_size": 32, + "in_channels": 3, + "out_channels": 3, + "layers_per_block": 2, + "num_class_embeds": 1000, + "block_out_channels": [32, 64], + "attention_head_dim": 8, + "down_block_types": [ + "ResnetDownsampleBlock2D", + "AttnDownBlock2D", + ], + "up_block_types": [ + "AttnUpBlock2D", + "ResnetUpsampleBlock2D", + ], + "resnet_time_scale_shift": "scale_shift", + "attn_norm_num_groups": 32, + "upsample_type": "resnet", + "downsample_type": "resnet", +} + +IMAGENET_64_UNET_CONFIG = { + "sample_size": 64, + "in_channels": 3, + "out_channels": 3, + "layers_per_block": 3, + "num_class_embeds": 1000, + "block_out_channels": [192, 192 * 2, 192 * 3, 192 * 4], + "attention_head_dim": 64, + "down_block_types": [ + "ResnetDownsampleBlock2D", + "AttnDownBlock2D", + "AttnDownBlock2D", + "AttnDownBlock2D", + ], + "up_block_types": [ + "AttnUpBlock2D", + "AttnUpBlock2D", + "AttnUpBlock2D", + "ResnetUpsampleBlock2D", + ], + "resnet_time_scale_shift": "scale_shift", + "attn_norm_num_groups": 32, + "upsample_type": "resnet", + "downsample_type": "resnet", +} + +LSUN_256_UNET_CONFIG = { + "sample_size": 256, + "in_channels": 3, + "out_channels": 3, + "layers_per_block": 2, + "num_class_embeds": None, + "block_out_channels": [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4], + "attention_head_dim": 64, + "down_block_types": [ + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "AttnDownBlock2D", + "AttnDownBlock2D", + "AttnDownBlock2D", + ], + "up_block_types": [ + "AttnUpBlock2D", + "AttnUpBlock2D", + "AttnUpBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + ], + "resnet_time_scale_shift": "default", + "upsample_type": "resnet", + "downsample_type": "resnet", +} + +CD_SCHEDULER_CONFIG = { + "num_train_timesteps": 40, + "sigma_min": 0.002, + "sigma_max": 80.0, +} + +CT_IMAGENET_64_SCHEDULER_CONFIG = { + "num_train_timesteps": 201, + "sigma_min": 0.002, + "sigma_max": 80.0, +} + +CT_LSUN_256_SCHEDULER_CONFIG = { + "num_train_timesteps": 151, + "sigma_min": 0.002, + "sigma_max": 80.0, +} + +__all__ = [ + "CD_SCHEDULER_CONFIG", + "CT_IMAGENET_64_SCHEDULER_CONFIG", + "CT_LSUN_256_SCHEDULER_CONFIG", + "IMAGENET_64_UNET_CONFIG", + "LSUN_256_UNET_CONFIG", + "TEST_UNET_CONFIG", +] diff --git a/src/diffusers/loaders/conversion/configs/cosmos.py b/src/diffusers/loaders/conversion/configs/cosmos.py new file mode 100644 index 000000000000..121bf78c99bd --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/cosmos.py @@ -0,0 +1,266 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the cosmos assembly recipe.""" + +TRANSFORMER_CONFIGS = { + "Cosmos-1.0-Diffusion-7B-Text2World": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 32, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (2.0, 1.0, 1.0), + "concat_padding_mask": True, + "extra_pos_embed_type": "learnable", + }, + "Cosmos-1.0-Diffusion-7B-Video2World": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 32, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (2.0, 1.0, 1.0), + "concat_padding_mask": True, + "extra_pos_embed_type": "learnable", + }, + "Cosmos-1.0-Diffusion-14B-Text2World": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 40, + "attention_head_dim": 128, + "num_layers": 36, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (2.0, 2.0, 2.0), + "concat_padding_mask": True, + "extra_pos_embed_type": "learnable", + }, + "Cosmos-1.0-Diffusion-14B-Video2World": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 40, + "attention_head_dim": 128, + "num_layers": 36, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (2.0, 2.0, 2.0), + "concat_padding_mask": True, + "extra_pos_embed_type": "learnable", + }, + "Cosmos-2.0-Diffusion-2B-Text2Image": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 16, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 4.0, 4.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + }, + "Cosmos-2.0-Diffusion-14B-Text2Image": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 40, + "attention_head_dim": 128, + "num_layers": 36, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 4.0, 4.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + }, + "Cosmos-2.0-Diffusion-2B-Video2World": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 16, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 3.0, 3.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + }, + "Cosmos-2.0-Diffusion-14B-Video2World": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 40, + "attention_head_dim": 128, + "num_layers": 36, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (20 / 24, 2.0, 2.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + }, + "Cosmos-2.5-Predict-Base-2B": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 16, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 3.0, 3.0), + "concat_padding_mask": True, + # NOTE: source config has pos_emb_learnable: 'True' - but params are missing + "extra_pos_embed_type": None, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + }, + "Cosmos-2.5-Predict-Base-14B": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 40, + "attention_head_dim": 128, + "num_layers": 36, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 3.0, 3.0), + "concat_padding_mask": True, + # NOTE: source config has pos_emb_learnable: 'True' - but params are missing + "extra_pos_embed_type": None, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + }, + "Cosmos-2.5-Transfer-General-2B": { + "in_channels": 16 + 1, + "out_channels": 16, + "num_attention_heads": 16, + "attention_head_dim": 128, + "num_layers": 28, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "max_size": (128, 240, 240), + "patch_size": (1, 2, 2), + "rope_scale": (1.0, 3.0, 3.0), + "concat_padding_mask": True, + "extra_pos_embed_type": None, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + "controlnet_block_every_n": 7, + "img_context_dim_in": 1152, + "img_context_dim_out": 2048, + "img_context_num_tokens": 256, + }, +} + +CONTROLNET_CONFIGS = { + "Cosmos-2.5-Transfer-General-2B": { + "n_controlnet_blocks": 4, + "model_channels": 2048, + "in_channels": 130, + "latent_channels": 18, # (16 latent + 1 condition_mask) + 1 padding_mask = 18 + "num_attention_heads": 16, + "attention_head_dim": 128, + "mlp_ratio": 4.0, + "text_embed_dim": 1024, + "adaln_lora_dim": 256, + "patch_size": (1, 2, 2), + "max_size": (128, 240, 240), + "rope_scale": (1.0, 3.0, 3.0), + "extra_pos_embed_type": None, + "img_context_dim_in": 1152, + "img_context_dim_out": 2048, + "use_crossattn_projection": True, + "crossattn_proj_in_channels": 100352, + "encoder_hidden_states_channels": 1024, + }, +} + +VAE_CONFIGS = { + "CV8x8x8-0.1": { + "name": "nvidia/Cosmos-0.1-Tokenizer-CV8x8x8", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 16, + "encoder_block_out_channels": (128, 256, 512, 512), + "decode_block_out_channels": (256, 512, 512, 512), + "attention_resolutions": (32,), + "resolution": 1024, + "num_layers": 2, + "patch_size": 4, + "patch_type": "haar", + "scaling_factor": 1.0, + "spatial_compression_ratio": 8, + "temporal_compression_ratio": 8, + "latents_mean": None, + "latents_std": None, + }, + }, + "CV8x8x8-1.0": { + "name": "nvidia/Cosmos-1.0-Tokenizer-CV8x8x8", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 16, + "encoder_block_out_channels": (128, 256, 512, 512), + "decode_block_out_channels": (256, 512, 512, 512), + "attention_resolutions": (32,), + "resolution": 1024, + "num_layers": 2, + "patch_size": 4, + "patch_type": "haar", + "scaling_factor": 1.0, + "spatial_compression_ratio": 8, + "temporal_compression_ratio": 8, + "latents_mean": None, + "latents_std": None, + }, + }, +} + +__all__ = ["CONTROLNET_CONFIGS", "TRANSFORMER_CONFIGS", "VAE_CONFIGS"] diff --git a/src/diffusers/loaders/conversion/configs/deepfloyd_if.py b/src/diffusers/loaders/conversion/configs/deepfloyd_if.py new file mode 100644 index 000000000000..a65647e24bf9 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/deepfloyd_if.py @@ -0,0 +1,200 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the deepfloyd_if assembly recipe.""" + + +def create_unet_diffusers_config(original_unet_config, class_embed_type=None): + attention_resolutions = parse_list(original_unet_config["attention_resolutions"]) + attention_resolutions = [original_unet_config["image_size"] // int(res) for res in attention_resolutions] + + channel_mult = parse_list(original_unet_config["channel_mult"]) + block_out_channels = [original_unet_config["model_channels"] * mult for mult in channel_mult] + + down_block_types = [] + resolution = 1 + + for i in range(len(block_out_channels)): + if resolution in attention_resolutions: + block_type = "SimpleCrossAttnDownBlock2D" + elif original_unet_config["resblock_updown"]: + block_type = "ResnetDownsampleBlock2D" + else: + block_type = "DownBlock2D" + + down_block_types.append(block_type) + + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + if resolution in attention_resolutions: + block_type = "SimpleCrossAttnUpBlock2D" + elif original_unet_config["resblock_updown"]: + block_type = "ResnetUpsampleBlock2D" + else: + block_type = "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + head_dim = original_unet_config["num_head_channels"] + + use_linear_projection = ( + original_unet_config["use_linear_in_transformer"] + if "use_linear_in_transformer" in original_unet_config + else False + ) + if use_linear_projection: + # stable diffusion 2-base-512 and 2-768 + if head_dim is None: + head_dim = [5, 10, 20, 20] + + projection_class_embeddings_input_dim = None + + if class_embed_type is None: + if "num_classes" in original_unet_config: + if original_unet_config["num_classes"] == "sequential": + class_embed_type = "projection" + assert "adm_in_channels" in original_unet_config + projection_class_embeddings_input_dim = original_unet_config["adm_in_channels"] + else: + raise NotImplementedError( + f"Unknown conditional unet num_classes config: {original_unet_config['num_classes']}" + ) + + config = { + "sample_size": original_unet_config["image_size"], + "in_channels": original_unet_config["in_channels"], + "down_block_types": tuple(down_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": original_unet_config["num_res_blocks"], + "cross_attention_dim": original_unet_config["encoder_channels"], + "attention_head_dim": head_dim, + "use_linear_projection": use_linear_projection, + "class_embed_type": class_embed_type, + "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, + "out_channels": original_unet_config["out_channels"], + "up_block_types": tuple(up_block_types), + "upcast_attention": False, # TODO: guessing + "cross_attention_norm": "group_norm", + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "addition_embed_type": "text", + "act_fn": "gelu", + } + + if original_unet_config["use_scale_shift_norm"]: + config["resnet_time_scale_shift"] = "scale_shift" + + if "encoder_dim" in original_unet_config: + config["encoder_hid_dim"] = original_unet_config["encoder_dim"] + + return config + + +def parse_list(value): + if isinstance(value, str): + value = value.split(",") + value = [int(v) for v in value] + elif isinstance(value, list): + pass + else: + raise ValueError(f"Can't parse list for type: {type(value)}") + + return value + + +def superres_create_unet_diffusers_config(original_unet_config): + attention_resolutions = parse_list(original_unet_config["attention_resolutions"]) + attention_resolutions = [original_unet_config["image_size"] // int(res) for res in attention_resolutions] + + channel_mult = parse_list(original_unet_config["channel_mult"]) + block_out_channels = [original_unet_config["model_channels"] * mult for mult in channel_mult] + + down_block_types = [] + resolution = 1 + + for i in range(len(block_out_channels)): + if resolution in attention_resolutions: + block_type = "SimpleCrossAttnDownBlock2D" + elif original_unet_config["resblock_updown"]: + block_type = "ResnetDownsampleBlock2D" + else: + block_type = "DownBlock2D" + + down_block_types.append(block_type) + + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + if resolution in attention_resolutions: + block_type = "SimpleCrossAttnUpBlock2D" + elif original_unet_config["resblock_updown"]: + block_type = "ResnetUpsampleBlock2D" + else: + block_type = "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + head_dim = original_unet_config["num_head_channels"] + use_linear_projection = ( + original_unet_config["use_linear_in_transformer"] + if "use_linear_in_transformer" in original_unet_config + else False + ) + if use_linear_projection: + # stable diffusion 2-base-512 and 2-768 + if head_dim is None: + head_dim = [5, 10, 20, 20] + + class_embed_type = None + projection_class_embeddings_input_dim = None + + if "num_classes" in original_unet_config: + if original_unet_config["num_classes"] == "sequential": + class_embed_type = "projection" + assert "adm_in_channels" in original_unet_config + projection_class_embeddings_input_dim = original_unet_config["adm_in_channels"] + else: + raise NotImplementedError( + f"Unknown conditional unet num_classes config: {original_unet_config['num_classes']}" + ) + + config = { + "in_channels": original_unet_config["in_channels"], + "down_block_types": tuple(down_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": tuple(original_unet_config["num_res_blocks"]), + "cross_attention_dim": original_unet_config["encoder_channels"], + "attention_head_dim": head_dim, + "use_linear_projection": use_linear_projection, + "class_embed_type": class_embed_type, + "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, + "out_channels": original_unet_config["out_channels"], + "up_block_types": tuple(up_block_types), + "upcast_attention": False, # TODO: guessing + "cross_attention_norm": "group_norm", + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "act_fn": "gelu", + } + + if original_unet_config["use_scale_shift_norm"]: + config["resnet_time_scale_shift"] = "scale_shift" + + return config + + +__all__ = ["create_unet_diffusers_config", "superres_create_unet_diffusers_config"] diff --git a/src/diffusers/loaders/conversion/configs/flux2.py b/src/diffusers/loaders/conversion/configs/flux2.py new file mode 100644 index 000000000000..4bd3d7824074 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/flux2.py @@ -0,0 +1,85 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the flux2 assembly recipe.""" + +from typing import Any, Dict + + +def get_flux2_transformer_config(model_type: str) -> Dict[str, Any]: + if model_type == "flux2-dev": + config = { + "model_id": "black-forest-labs/FLUX.2-dev", + "diffusers_config": { + "patch_size": 1, + "in_channels": 128, + "num_layers": 8, + "num_single_layers": 48, + "attention_head_dim": 128, + "num_attention_heads": 48, + "joint_attention_dim": 15360, + "timestep_guidance_channels": 256, + "mlp_ratio": 3.0, + "axes_dims_rope": (32, 32, 32, 32), + "rope_theta": 2000, + "eps": 1e-6, + }, + } + elif model_type == "klein-4b": + config = { + "model_id": "diffusers-internal-dev/dummy0115", + "diffusers_config": { + "patch_size": 1, + "in_channels": 128, + "num_layers": 5, + "num_single_layers": 20, + "attention_head_dim": 128, + "num_attention_heads": 24, + "joint_attention_dim": 7680, + "timestep_guidance_channels": 256, + "mlp_ratio": 3.0, + "axes_dims_rope": (32, 32, 32, 32), + "rope_theta": 2000, + "eps": 1e-6, + "guidance_embeds": False, + }, + } + + elif model_type == "klein-9b": + config = { + "model_id": "diffusers-internal-dev/dummy0115", + "diffusers_config": { + "patch_size": 1, + "in_channels": 128, + "num_layers": 8, + "num_single_layers": 24, + "attention_head_dim": 128, + "num_attention_heads": 32, + "joint_attention_dim": 12288, + "timestep_guidance_channels": 256, + "mlp_ratio": 3.0, + "axes_dims_rope": (32, 32, 32, 32), + "rope_theta": 2000, + "eps": 1e-6, + "guidance_embeds": False, + }, + } + + else: + raise ValueError(f"Unknown model_type: {model_type}. Choose from: flux2-dev, klein-4b, klein-9b") + + return config + + +__all__ = ["get_flux2_transformer_config"] diff --git a/src/diffusers/loaders/conversion/configs/gligen.py b/src/diffusers/loaders/conversion/configs/gligen.py new file mode 100644 index 000000000000..0a4c07a06432 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/gligen.py @@ -0,0 +1,85 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the gligen assembly recipe.""" + + +def create_vae_config(original_config, image_size: int): + vae_params = original_config["autoencoder"]["params"]["ddconfig"] + _ = original_config["autoencoder"]["params"]["embed_dim"] + + block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + config = { + "sample_size": image_size, + "in_channels": vae_params["in_channels"], + "out_channels": vae_params["out_ch"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params["z_channels"], + "layers_per_block": vae_params["num_res_blocks"], + } + + return config + + +def create_unet_config(original_config, image_size: int, attention_type): + unet_params = original_config["model"]["params"] + vae_params = original_config["autoencoder"]["params"]["ddconfig"] + + block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) + + head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None + use_linear_projection = ( + unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False + ) + if use_linear_projection: + if head_dim is None: + head_dim = [5, 10, 20, 20] + + config = { + "sample_size": image_size // vae_scale_factor, + "in_channels": unet_params["in_channels"], + "down_block_types": tuple(down_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params["num_res_blocks"], + "cross_attention_dim": unet_params["context_dim"], + "attention_head_dim": head_dim, + "use_linear_projection": use_linear_projection, + "attention_type": attention_type, + } + + return config + + +__all__ = ["create_unet_config", "create_vae_config"] diff --git a/src/diffusers/loaders/conversion/configs/hunyuan_video.py b/src/diffusers/loaders/conversion/configs/hunyuan_video.py new file mode 100644 index 000000000000..b46f286c9be4 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/hunyuan_video.py @@ -0,0 +1,77 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the hunyuan_video assembly recipe.""" + +TRANSFORMER_CONFIGS = { + "HYVideo-T/2-cfgdistill": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 24, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "mlp_ratio": 4.0, + "patch_size": 2, + "patch_size_t": 1, + "qk_norm": "rms_norm", + "guidance_embeds": True, + "text_embed_dim": 4096, + "pooled_projection_dim": 768, + "rope_theta": 256.0, + "rope_axes_dim": (16, 56, 56), + "image_condition_type": None, + }, + "HYVideo-T/2-I2V-33ch": { + "in_channels": 16 * 2 + 1, + "out_channels": 16, + "num_attention_heads": 24, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "mlp_ratio": 4.0, + "patch_size": 2, + "patch_size_t": 1, + "qk_norm": "rms_norm", + "guidance_embeds": False, + "text_embed_dim": 4096, + "pooled_projection_dim": 768, + "rope_theta": 256.0, + "rope_axes_dim": (16, 56, 56), + "image_condition_type": "latent_concat", + }, + "HYVideo-T/2-I2V-16ch": { + "in_channels": 16, + "out_channels": 16, + "num_attention_heads": 24, + "attention_head_dim": 128, + "num_layers": 20, + "num_single_layers": 40, + "num_refiner_layers": 2, + "mlp_ratio": 4.0, + "patch_size": 2, + "patch_size_t": 1, + "qk_norm": "rms_norm", + "guidance_embeds": True, + "text_embed_dim": 4096, + "pooled_projection_dim": 768, + "rope_theta": 256.0, + "rope_axes_dim": (16, 56, 56), + "image_condition_type": "token_replace", + }, +} + +__all__ = ["TRANSFORMER_CONFIGS"] diff --git a/src/diffusers/loaders/conversion/configs/hunyuan_video15.py b/src/diffusers/loaders/conversion/configs/hunyuan_video15.py new file mode 100644 index 000000000000..d672b054cd4e --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/hunyuan_video15.py @@ -0,0 +1,103 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the hunyuan_video15 assembly recipe.""" + +TRANSFORMER_CONFIGS = { + "480p_t2v": { + "target_size": 640, + "task_type": "i2v", + }, + "720p_t2v": { + "target_size": 960, + "task_type": "t2v", + }, + "720p_i2v": { + "target_size": 960, + "task_type": "i2v", + }, + "480p_t2v_distilled": { + "target_size": 640, + "task_type": "t2v", + }, + "480p_i2v_distilled": { + "target_size": 640, + "task_type": "i2v", + }, + "720p_i2v_distilled": { + "target_size": 960, + "task_type": "i2v", + }, + "480p_i2v_step_distilled": { + "target_size": 640, + "task_type": "i2v", + "use_meanflow": True, + }, +} + +SCHEDULER_CONFIGS = { + "480p_t2v": { + "shift": 5.0, + }, + "480p_i2v": { + "shift": 5.0, + }, + "720p_t2v": { + "shift": 9.0, + }, + "720p_i2v": { + "shift": 7.0, + }, + "480p_t2v_distilled": { + "shift": 5.0, + }, + "480p_i2v_distilled": { + "shift": 5.0, + }, + "720p_i2v_distilled": { + "shift": 7.0, + }, + "480p_i2v_step_distilled": { + "shift": 7.0, + }, +} + +GUIDANCE_CONFIGS = { + "480p_t2v": { + "guidance_scale": 6.0, + }, + "480p_i2v": { + "guidance_scale": 6.0, + }, + "720p_t2v": { + "guidance_scale": 6.0, + }, + "720p_i2v": { + "guidance_scale": 6.0, + }, + "480p_t2v_distilled": { + "guidance_scale": 1.0, + }, + "480p_i2v_distilled": { + "guidance_scale": 1.0, + }, + "720p_i2v_distilled": { + "guidance_scale": 1.0, + }, + "480p_i2v_step_distilled": { + "guidance_scale": 1.0, + }, +} + +__all__ = ["GUIDANCE_CONFIGS", "SCHEDULER_CONFIGS", "TRANSFORMER_CONFIGS"] diff --git a/src/diffusers/loaders/conversion/configs/joyimage.py b/src/diffusers/loaders/conversion/configs/joyimage.py new file mode 100644 index 000000000000..566a66c4c370 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/joyimage.py @@ -0,0 +1,30 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the joyimage assembly recipe.""" + +TRANSFORMER_CONFIG = { + "hidden_size": 4096, + "in_channels": 16, + "num_attention_heads": 32, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "rope_dim_list": [16, 56, 56], + "text_dim": 4096, + "rope_type": "rope", + "theta": 10000, +} + +__all__ = ["TRANSFORMER_CONFIG"] diff --git a/src/diffusers/loaders/conversion/configs/kandinsky.py b/src/diffusers/loaders/conversion/configs/kandinsky.py new file mode 100644 index 000000000000..aec53ce9b7c9 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/kandinsky.py @@ -0,0 +1,142 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the kandinsky assembly recipe.""" + +PRIOR_CONFIG = {} + +UNET_CONFIG = { + "act_fn": "silu", + "addition_embed_type": "text_image", + "addition_embed_type_num_heads": 64, + "attention_head_dim": 64, + "block_out_channels": [384, 768, 1152, 1536], + "center_input_sample": False, + "class_embed_type": None, + "class_embeddings_concat": False, + "conv_in_kernel": 3, + "conv_out_kernel": 3, + "cross_attention_dim": 768, + "cross_attention_norm": None, + "down_block_types": [ + "ResnetDownsampleBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + ], + "downsample_padding": 1, + "dual_cross_attention": False, + "encoder_hid_dim": 1024, + "encoder_hid_dim_type": "text_image_proj", + "flip_sin_to_cos": True, + "freq_shift": 0, + "in_channels": 4, + "layers_per_block": 3, + "mid_block_only_cross_attention": None, + "mid_block_scale_factor": 1, + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "norm_eps": 1e-05, + "norm_num_groups": 32, + "num_class_embeds": None, + "only_cross_attention": False, + "out_channels": 8, + "projection_class_embeddings_input_dim": None, + "resnet_out_scale_factor": 1.0, + "resnet_skip_time_act": False, + "resnet_time_scale_shift": "scale_shift", + "sample_size": 64, + "time_cond_proj_dim": None, + "time_embedding_act_fn": None, + "time_embedding_dim": None, + "time_embedding_type": "positional", + "timestep_post_act": None, + "up_block_types": [ + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "ResnetUpsampleBlock2D", + ], + "upcast_attention": False, + "use_linear_projection": False, +} + +INPAINT_UNET_CONFIG = { + "act_fn": "silu", + "addition_embed_type": "text_image", + "addition_embed_type_num_heads": 64, + "attention_head_dim": 64, + "block_out_channels": [384, 768, 1152, 1536], + "center_input_sample": False, + "class_embed_type": None, + "class_embeddings_concat": None, + "conv_in_kernel": 3, + "conv_out_kernel": 3, + "cross_attention_dim": 768, + "cross_attention_norm": None, + "down_block_types": [ + "ResnetDownsampleBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + ], + "downsample_padding": 1, + "dual_cross_attention": False, + "encoder_hid_dim": 1024, + "encoder_hid_dim_type": "text_image_proj", + "flip_sin_to_cos": True, + "freq_shift": 0, + "in_channels": 9, + "layers_per_block": 3, + "mid_block_only_cross_attention": None, + "mid_block_scale_factor": 1, + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "norm_eps": 1e-05, + "norm_num_groups": 32, + "num_class_embeds": None, + "only_cross_attention": False, + "out_channels": 8, + "projection_class_embeddings_input_dim": None, + "resnet_out_scale_factor": 1.0, + "resnet_skip_time_act": False, + "resnet_time_scale_shift": "scale_shift", + "sample_size": 64, + "time_cond_proj_dim": None, + "time_embedding_act_fn": None, + "time_embedding_dim": None, + "time_embedding_type": "positional", + "timestep_post_act": None, + "up_block_types": [ + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "ResnetUpsampleBlock2D", + ], + "upcast_attention": False, + "use_linear_projection": False, +} + +MOVQ_CONFIG = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "down_block_types": ("DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "AttnDownEncoderBlock2D"), + "up_block_types": ("AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"), + "num_vq_embeddings": 16384, + "block_out_channels": (128, 256, 256, 512), + "vq_embed_dim": 4, + "layers_per_block": 2, + "norm_type": "spatial", +} + +__all__ = ["INPAINT_UNET_CONFIG", "MOVQ_CONFIG", "PRIOR_CONFIG", "UNET_CONFIG"] diff --git a/src/diffusers/loaders/conversion/configs/ltx.py b/src/diffusers/loaders/conversion/configs/ltx.py new file mode 100644 index 000000000000..25a31223a30a --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/ltx.py @@ -0,0 +1,208 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the ltx assembly recipe.""" + +from typing import Any + + +def get_transformer_config(version: str) -> dict[str, Any]: + if version == "0.9.7": + config = { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 128, + "cross_attention_dim": 4096, + "num_layers": 48, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 4096, + "attention_bias": True, + "attention_out_bias": True, + } + else: + config = { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 64, + "cross_attention_dim": 2048, + "num_layers": 28, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 4096, + "attention_bias": True, + "attention_out_bias": True, + } + return config + + +def get_vae_config(version: str) -> dict[str, Any]: + if version in ["0.9.0"]: + config = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (128, 256, 512, 512), + "down_block_types": ( + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + ), + "decoder_block_out_channels": (128, 256, 512, 512), + "layers_per_block": (4, 3, 3, 3, 4), + "decoder_layers_per_block": (4, 3, 3, 3, 4), + "spatio_temporal_scaling": (True, True, True, False), + "decoder_spatio_temporal_scaling": (True, True, True, False), + "decoder_inject_noise": (False, False, False, False, False), + "downsample_type": ("conv", "conv", "conv", "conv"), + "upsample_residual": (False, False, False, False), + "upsample_factor": (1, 1, 1, 1), + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "scaling_factor": 1.0, + "encoder_causal": True, + "decoder_causal": False, + "timestep_conditioning": False, + } + elif version in ["0.9.1"]: + config = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (128, 256, 512, 512), + "down_block_types": ( + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + "LTXVideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 3, 3, 3, 4), + "decoder_layers_per_block": (5, 6, 7, 8), + "spatio_temporal_scaling": (True, True, True, False), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (True, True, True, False), + "downsample_type": ("conv", "conv", "conv", "conv"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": True, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "scaling_factor": 1.0, + "encoder_causal": True, + "decoder_causal": False, + } + elif version in ["0.9.5"]: + config = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (128, 256, 512, 1024, 2048), + "down_block_types": ( + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": True, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "scaling_factor": 1.0, + "encoder_causal": True, + "decoder_causal": False, + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + } + elif version in ["0.9.7"]: + config = { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (128, 256, 512, 1024, 2048), + "down_block_types": ( + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + "LTXVideo095DownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": True, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "scaling_factor": 1.0, + "encoder_causal": True, + "decoder_causal": False, + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + } + return config + + +def get_spatial_latent_upsampler_config(version: str) -> dict[str, Any]: + if version == "0.9.7": + config = { + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + } + elif version == "0.9.8": + config = { + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + } + else: + raise ValueError(f"Unsupported version: {version}") + return config + + +__all__ = ["get_spatial_latent_upsampler_config", "get_transformer_config", "get_vae_config"] diff --git a/src/diffusers/loaders/conversion/configs/ltx2.py b/src/diffusers/loaders/conversion/configs/ltx2.py new file mode 100644 index 000000000000..7198c945c320 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/ltx2.py @@ -0,0 +1,737 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the ltx2 assembly recipe.""" + +from typing import Any + + +def get_ltx2_transformer_config(version: str) -> dict[str, Any]: + if version == "test": + # Produces a transformer of the same size as used in test_models_transformer_ltx2.py + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "in_channels": 4, + "out_channels": 4, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 2, + "attention_head_dim": 8, + "cross_attention_dim": 16, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "audio_in_channels": 4, + "audio_out_channels": 4, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 2, + "audio_attention_head_dim": 4, + "audio_cross_attention_dim": 8, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "num_layers": 2, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 16, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": False, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1, + }, + } + elif version == "2.0": + config = { + "model_id": "Lightricks/LTX-2", + "diffusers_config": { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 128, + "cross_attention_dim": 4096, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "gated_attn": False, + "cross_attn_mod": False, + "audio_in_channels": 128, + "audio_out_channels": 128, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 32, + "audio_attention_head_dim": 64, + "audio_cross_attention_dim": 2048, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "audio_gated_attn": False, + "audio_cross_attn_mod": False, + "num_layers": 48, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 3840, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1000, + "rope_type": "split", + "use_prompt_embeddings": True, + "perturbed_attn": False, + }, + } + elif version == "2.3": + config = { + "model_id": "Lightricks/LTX-2.3", + "diffusers_config": { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 128, + "cross_attention_dim": 4096, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "gated_attn": True, + "cross_attn_mod": True, + "audio_in_channels": 128, + "audio_out_channels": 128, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 32, + "audio_attention_head_dim": 64, + "audio_cross_attention_dim": 2048, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "audio_gated_attn": True, + "audio_cross_attn_mod": True, + "num_layers": 48, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 3840, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1000, + "rope_type": "split", + "use_prompt_embeddings": False, + "perturbed_attn": True, + }, + } + elif version == "2.5": + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + "in_channels": 128, + "out_channels": 128, + "patch_size": 1, + "patch_size_t": 1, + "num_attention_heads": 32, + "attention_head_dim": 128, + "cross_attention_dim": 4096, + "vae_scale_factors": (8, 32, 32), + "pos_embed_max_pos": 20, + "base_height": 2048, + "base_width": 2048, + "gated_attn": True, + "cross_attn_mod": True, + "audio_in_channels": 128, + "audio_out_channels": 128, + "audio_patch_size": 1, + "audio_patch_size_t": 1, + "audio_num_attention_heads": 32, + "audio_attention_head_dim": 64, + "audio_cross_attention_dim": 2048, + "audio_scale_factor": 4, + "audio_pos_embed_max_pos": 20, + "audio_sampling_rate": 16000, + "audio_hop_length": 160, + "audio_gated_attn": True, + "audio_cross_attn_mod": True, + "num_layers": 48, + "activation_fn": "gelu-approximate", + "qk_norm": "rms_norm_across_heads", + "norm_elementwise_affine": False, + "norm_eps": 1e-6, + "caption_channels": 3840, + "attention_bias": True, + "attention_out_bias": True, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_offset": 1, + "timestep_scale_multiplier": 1000, + "cross_attn_timestep_scale_multiplier": 1000, + "rope_type": "split", + "use_prompt_embeddings": False, + "perturbed_attn": True, + # The only transformer-level deltas from 2.3: the video FFN drops its bias (audio_ff_bias and + # use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5 carries a + # learned keyframe absolute-position embedding. + "ff_bias": False, + "use_keyframes_abs_pos_embedding": True, + }, + } + return config + + +def get_ltx2_connectors_config( + version: str, gemma_text_config: Any | None = None +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "test": + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "caption_channels": 16, + "text_proj_in_factor": 3, + "video_connector_num_attention_heads": 4, + "video_connector_attention_head_dim": 8, + "video_connector_num_layers": 1, + "video_connector_num_learnable_registers": None, + "audio_connector_num_attention_heads": 4, + "audio_connector_attention_head_dim": 8, + "audio_connector_num_layers": 1, + "audio_connector_num_learnable_registers": None, + "connector_rope_base_seq_len": 32, + "rope_theta": 10000.0, + "rope_double_precision": False, + "causal_temporal_positioning": False, + }, + } + elif version == "2.0": + config = { + "model_id": "Lightricks/LTX-2", + "diffusers_config": { + "caption_channels": 3840, + "text_proj_in_factor": 49, + "video_connector_num_attention_heads": 30, + "video_connector_attention_head_dim": 128, + "video_connector_num_layers": 2, + "video_connector_num_learnable_registers": 128, + "video_gated_attn": False, + "audio_connector_num_attention_heads": 30, + "audio_connector_attention_head_dim": 128, + "audio_connector_num_layers": 2, + "audio_connector_num_learnable_registers": 128, + "audio_gated_attn": False, + "connector_rope_base_seq_len": 4096, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_temporal_positioning": False, + "rope_type": "split", + "per_modality_projections": False, + "proj_bias": False, + }, + } + elif version == "2.3": + config = { + "model_id": "Lightricks/LTX-2.3", + "diffusers_config": { + "caption_channels": 3840, + "text_proj_in_factor": 49, + "video_connector_num_attention_heads": 32, + "video_connector_attention_head_dim": 128, + "video_connector_num_layers": 8, + "video_connector_num_learnable_registers": 128, + "video_gated_attn": True, + "audio_connector_num_attention_heads": 32, + "audio_connector_attention_head_dim": 64, + "audio_connector_num_layers": 8, + "audio_connector_num_learnable_registers": 128, + "audio_gated_attn": True, + "connector_rope_base_seq_len": 4096, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_temporal_positioning": False, + "rope_type": "split", + "per_modality_projections": True, + "video_hidden_dim": 4096, + "audio_hidden_dim": 2048, + "proj_bias": True, + }, + } + elif version == "2.5": + if gemma_text_config is None: + raise ValueError("gemma_text_config is required to derive connector dims for LTX-2.5.") + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + # Derived from the Gemma 4 text config rather than hardcoded, since (unlike Gemma-3-12B) the + # 2.5 text encoder isn't a single fixed checkpoint. Formula matches the reference + # (`encoder_configurator._create_feature_extractor`): hidden_size, and num_hidden_layers + 1 + # for the embedding layer. + "caption_channels": gemma_text_config.hidden_size, + "text_proj_in_factor": gemma_text_config.num_hidden_layers + 1, + "video_connector_num_attention_heads": 32, + "video_connector_attention_head_dim": 128, + "video_connector_num_layers": 8, + "video_connector_num_learnable_registers": 128, + "video_gated_attn": True, + "audio_connector_num_attention_heads": 32, + "audio_connector_attention_head_dim": 64, + "audio_connector_num_layers": 8, + "audio_connector_num_learnable_registers": 128, + "audio_gated_attn": True, + "connector_rope_base_seq_len": 4096, + "rope_theta": 10000.0, + "rope_double_precision": True, + "causal_temporal_positioning": False, + "rope_type": "split", + "per_modality_projections": True, + "video_hidden_dim": 4096, + "audio_hidden_dim": 2048, + "proj_bias": True, + }, + } + + return config + + +def get_ltx2_video_vae_config( + version: str, timestep_conditioning: bool = False +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "test": + config = { + "model_id": "diffusers-internal-dev/dummy-ltx2", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 2048), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + elif version == "2.0": + config = { + "model_id": "Lightricks/LTX-2", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 2048), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 1024), + "layers_per_block": (4, 6, 6, 2, 2), + "decoder_layers_per_block": (5, 5, 5, 5), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True), + "decoder_inject_noise": (False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_type": ("spatiotemporal", "spatiotemporal", "spatiotemporal"), + "upsample_residual": (True, True, True), + "upsample_factor": (2, 2, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + elif version == "2.3": + config = { + "model_id": "Lightricks/LTX-2.3", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 1024), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 512, 1024), + "layers_per_block": (4, 6, 4, 2, 2), + "decoder_layers_per_block": (4, 6, 4, 2, 2), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True, True), + "decoder_inject_noise": (False, False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + "upsample_residual": (False, False, False, False), + "upsample_factor": (2, 2, 1, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "zeros", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + elif version == "2.5": + # Same block structure as 2.3 (32x32x8 compression); confirmed against the checkpoint's + # config["vae"]["encoder_blocks"]/["decoder_blocks"] metadata, which is byte-identical to 2.3's. + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 128, + "block_out_channels": (256, 512, 1024, 1024), + "down_block_types": ( + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + "LTX2VideoDownBlock3D", + ), + "decoder_block_out_channels": (256, 512, 512, 1024), + "layers_per_block": (4, 6, 4, 2, 2), + "decoder_layers_per_block": (4, 6, 4, 2, 2), + "spatio_temporal_scaling": (True, True, True, True), + "decoder_spatio_temporal_scaling": (True, True, True, True), + "decoder_inject_noise": (False, False, False, False, False), + "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"), + "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"), + "upsample_residual": (False, False, False, False), + "upsample_factor": (2, 2, 1, 2), + "timestep_conditioning": timestep_conditioning, + "patch_size": 4, + "patch_size_t": 1, + "resnet_norm_eps": 1e-6, + "encoder_causal": True, + "decoder_causal": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "zeros", + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + return config + + +def get_ltx2_diffusion_video_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version != "2.5": + raise ValueError( + f"The diffusion decoder was introduced in LTX-2.5, which the converter handles under " + f"`--version 2.5`; got version {version!r}." + ) + # The encoder half is 2.5's conv VAE encoder, unchanged (its weights are byte identical to 2.3's), so + # those entries must stay in sync with `get_ltx2_video_vae_config("2.5")`. + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + "out_channels": 3, + "latent_channels": 128, + "patch_size": 4, + "decoder_head_dim": 64, + "decoder_stage_channels": (2048, 1024, 512, 512, 256), + "decoder_stage_depths": (4, 6, 4, 2, 8), + "decoder_stage_kernels": ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)), + "decoder_upsample_strides": ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)), + "decoder_upsample_channel_reductions": (2, 2, 1, 2), + "decoder_stage5_kernel": (11, 11, 11), + "decoder_t_emb_dim": 384, + "decoder_timestep_scale_multiplier": 1000.0, + "decoder_model_output_type": "x0", + "decoder_num_inference_steps": 1, + "spatial_compression_ratio": 32, + "temporal_compression_ratio": 8, + }, + } + return config + + +def get_ltx2_audio_vae_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "2.0": + config = { + "model_id": "Lightricks/LTX-2", + "diffusers_config": { + "base_channels": 128, + "output_channels": 2, + "ch_mult": (1, 2, 4), + "num_res_blocks": 2, + "attn_resolutions": None, + "in_channels": 2, + "resolution": 256, + "latent_channels": 8, + "norm_type": "pixel", + "causality_axis": "height", + "dropout": 0.0, + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "is_causal": True, + "mel_bins": 64, + "double_z": True, + }, + } + elif version == "2.3": + config = { + "model_id": "Lightricks/LTX-2.3", + "diffusers_config": { + "base_channels": 128, + "output_channels": 2, + "ch_mult": (1, 2, 4), + "num_res_blocks": 2, + "attn_resolutions": None, + "in_channels": 2, + "resolution": 256, + "latent_channels": 8, + "norm_type": "pixel", + "causality_axis": "height", + "dropout": 0.0, + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "is_causal": True, + "mel_bins": 64, + "double_z": True, + }, # Same config as LTX-2.0 + } + elif version == "2.5": + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + "base_channels": 128, + "output_channels": 2, + "ch_mult": (1, 2, 4), + "num_res_blocks": 2, + "attn_resolutions": None, + "in_channels": 2, + "resolution": 256, + "latent_channels": 8, + "norm_type": "pixel", + "causality_axis": "height", + "dropout": 0.0, + "mid_block_add_attention": False, + "sample_rate": 16000, + "mel_hop_length": 160, + "is_causal": True, + "mel_bins": 64, + "double_z": True, + }, # Same config as LTX-2.0 / 2.3 + } + return config + + +def get_ltx2_vocoder_config(version: str) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if version == "2.0": + config = { + "model_id": "Lightricks/LTX-2", + "diffusers_config": { + "in_channels": 128, + "hidden_channels": 1024, + "out_channels": 2, + "upsample_kernel_sizes": [16, 15, 8, 4, 4], + "upsample_factors": [6, 5, 2, 2, 2], + "resnet_kernel_sizes": [3, 7, 11], + "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "act_fn": "leaky_relu", + "leaky_relu_negative_slope": 0.1, + "antialias": False, + "final_act_fn": "tanh", + "final_bias": True, + "output_sampling_rate": 24000, + }, + } + elif version == "2.3": + config = { + "model_id": "Lightricks/LTX-2.3", + "diffusers_config": { + "in_channels": 128, + "hidden_channels": 1536, + "out_channels": 2, + "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4], + "upsample_factors": [5, 2, 2, 2, 2, 2], + "resnet_kernel_sizes": [3, 7, 11], + "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "act_fn": "snakebeta", + "leaky_relu_negative_slope": 0.1, + "antialias": True, + "antialias_ratio": 2, + "antialias_kernel_size": 12, + "final_act_fn": None, + "final_bias": False, + "bwe_in_channels": 128, + "bwe_hidden_channels": 512, + "bwe_out_channels": 2, + "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4], + "bwe_upsample_factors": [6, 5, 2, 2, 2], + "bwe_resnet_kernel_sizes": [3, 7, 11], + "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "bwe_act_fn": "snakebeta", + "bwe_leaky_relu_negative_slope": 0.1, + "bwe_antialias": True, + "bwe_antialias_ratio": 2, + "bwe_antialias_kernel_size": 12, + "bwe_final_act_fn": None, + "bwe_final_bias": False, + "filter_length": 512, + "hop_length": 80, + "window_length": 512, + "num_mel_channels": 64, + "input_sampling_rate": 16000, + "output_sampling_rate": 48000, + }, + } + elif version == "2.5": + config = { + "model_id": "Lightricks/LTX-2.5", + "diffusers_config": { + "in_channels": 128, + "hidden_channels": 1536, + "out_channels": 2, + "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4], + "upsample_factors": [5, 2, 2, 2, 2, 2], + "resnet_kernel_sizes": [3, 7, 11], + "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "act_fn": "snakebeta", + "leaky_relu_negative_slope": 0.1, + "antialias": True, + "antialias_ratio": 2, + "antialias_kernel_size": 12, + "final_act_fn": None, + "final_bias": False, + "bwe_in_channels": 128, + "bwe_hidden_channels": 512, + "bwe_out_channels": 2, + "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4], + "bwe_upsample_factors": [6, 5, 2, 2, 2], + "bwe_resnet_kernel_sizes": [3, 7, 11], + "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "bwe_act_fn": "snakebeta", + "bwe_leaky_relu_negative_slope": 0.1, + "bwe_antialias": True, + "bwe_antialias_ratio": 2, + "bwe_antialias_kernel_size": 12, + "bwe_final_act_fn": None, + "bwe_final_bias": False, + "filter_length": 512, + "hop_length": 80, + "window_length": 512, + "num_mel_channels": 64, + "input_sampling_rate": 16000, + "output_sampling_rate": 48000, + }, # Same config as LTX-2.3 + } + return config + + +def get_ltx2_spatial_latent_upsampler_config(version: str): + if version == "2.0": + config = { + "in_channels": 128, + "mid_channels": 1024, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + "rational_spatial_scale": 2.0, + "use_rational_resampler": True, + } + elif version in ("2.3", "2.5"): + config = { + "in_channels": 128, + "mid_channels": 1024, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": True, + "temporal_upsample": False, + "rational_spatial_scale": 2.0, + "use_rational_resampler": False, + } + else: + raise ValueError(f"Unsupported version: {version}") + return config + + +def get_ltx2_temporal_latent_upsampler_config(version: str): + if version != "2.5": + raise ValueError(f"Unsupported version: {version}") + # The temporal x2 upsampler is narrower than its spatial sibling and pixel-shuffles along time only. + return { + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": False, + "temporal_upsample": True, + } + + +__all__ = [ + "get_ltx2_audio_vae_config", + "get_ltx2_connectors_config", + "get_ltx2_diffusion_video_vae_config", + "get_ltx2_spatial_latent_upsampler_config", + "get_ltx2_temporal_latent_upsampler_config", + "get_ltx2_transformer_config", + "get_ltx2_video_vae_config", + "get_ltx2_vocoder_config", +] diff --git a/src/diffusers/loaders/conversion/configs/minimax_h3.py b/src/diffusers/loaders/conversion/configs/minimax_h3.py new file mode 100644 index 000000000000..c7dff5119a50 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/minimax_h3.py @@ -0,0 +1,156 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the minimax_h3 assembly recipe.""" + +import json +import os +from typing import Any + + +MINIMAX_H3_TRANSFORMER_CONFIG = { + "num_attention_heads": 56, + "attention_head_dim": 128, + "hidden_size": 5376, + "num_layers": 50, + "num_refiner_layers": 2, # token_refiner_num_layers + "ffn_dim": 14336, # ffn_hidden_size + "in_channels": 24, # latents_dim + "audio_in_channels": 32, # audio_latents_dim + "patch_size": [1, 2, 2], + "text_dim": 5120, + "freq_dim": 256, # timestep_input_dim + "time_embed_hidden_dim": 5376, # time_embed_hidden_size + "time_embed_dim": 2688, + "rope_freq_dim": 16, # rope_inv_freq_len + "rope_theta": 10000.0, + "norm_eps": 1e-05, + "qk_norm_eps": 1e-05, + "final_norm_eps": 1e-05, +} + +MINIMAX_H3_TEST_TRANSFORMER_CONFIG = { + **MINIMAX_H3_TRANSFORMER_CONFIG, + "num_attention_heads": 2, + "attention_head_dim": 32, + "hidden_size": 64, + "num_layers": 2, + "num_refiner_layers": 2, + "ffn_dim": 128, + "text_dim": 48, + "freq_dim": 16, + "time_embed_hidden_dim": 64, + "time_embed_dim": 32, + "rope_freq_dim": 4, +} + +MINIMAX_H3_VIDEO_VAE_CONFIG = { + "in_channels": 3, + "out_channels": 3, # out_ch + "latent_channels": 24, # z_channels == embed_dim + "block_out_channels": [128, 256, 256, 512, 512, 1024], # ch * ch_mult + "layers_per_block": 2, # num_res_blocks + "spatial_downsample_factors": [2, 2, 2, 2, 1, 1], # space_down + "temporal_downsample_factors": [1, 2, 2, 1, 1, 1], # time_down + "norm_num_groups": 32, + "norm_eps": 1e-06, + "spatial_padding_mode": "reflect", # padding_mode + "decoder_num_layers": 36, # vit_decoder_kwargs.num_layers + "decoder_num_attention_heads": 32, # vit_decoder_kwargs.heads + "decoder_attention_head_dim": 64, # vit_decoder_kwargs.dim_head + "decoder_num_register_tokens": 4, # ViT3DDecoder default + "decoder_ffn_mult": 4, # FeedForward default + "decoder_rope_theta": 100.0, # vit_decoder_kwargs.rope_theta + "decoder_rope_dim_ratio": 0.75, # vit_decoder_kwargs.rope_dim_ratio + "decoder_norm_eps": 1e-05, # ViT3DDecoder eps + "clip_length": 17, # video_vae/config.json vae_clip_length + "token_drop": 3, # video_vae/config.json vae_token_drop +} + +MINIMAX_H3_TEST_VIDEO_VAE_CONFIG = { + **MINIMAX_H3_VIDEO_VAE_CONFIG, + "block_out_channels": [32, 64], + "layers_per_block": 1, + "spatial_downsample_factors": [2, 2], + "temporal_downsample_factors": [2, 2], + "decoder_num_layers": 4, + "decoder_num_attention_heads": 4, + "decoder_attention_head_dim": 32, +} + +MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG = { + "num_attention_heads": 8, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], +} + + +def get_audio_vae_config(checkpoint_path: str) -> dict[str, Any]: + """Build the `AutoencoderKLMiniMaxH3Audio` config from the original audio-VAE metadata. + + `audio_vae/metadata.json` carries the constructor kwargs the checkpoint was built with, and `audio_vae/config.json` + carries the per-channel `latents_mean` / `latents_std` MiniMax-H3 normalizes with. The two are cross-checked here + because they duplicate the latent width and sample rate. + """ + source_dir = os.path.join(checkpoint_path, "audio_vae") + with open(os.path.join(source_dir, "metadata.json")) as f: + kwargs = json.load(f)["metadata"]["kwargs"] + with open(os.path.join(source_dir, "config.json")) as f: + wrapper_config = json.load(f) + + if kwargs["decoder_type"] != "bigvgan": + raise ValueError(f"Only the BigVGAN decoder is supported, got {kwargs['decoder_type']!r}.") + if not kwargs["attn_proj"]: + raise ValueError("The audio VAE is expected to carry the causal-attention latent projection.") + latent_channels = kwargs["vae_latent_channels"] + if wrapper_config["latent_channels"] != latent_channels: + raise ValueError( + f"latent width disagreement: metadata.json says {latent_channels}, " + f"config.json says {wrapper_config['latent_channels']}." + ) + if wrapper_config["sample_rate"] != kwargs["sample_rate"]: + raise ValueError( + f"sample rate disagreement: metadata.json says {kwargs['sample_rate']}, " + f"config.json says {wrapper_config['sample_rate']}." + ) + for key in ("latents_mean", "latents_std"): + if len(wrapper_config[key]) != latent_channels: + raise KeyError(f"{source_dir}/config.json `{key}` does not have {latent_channels} entries.") + + return { + "encoder_dim": kwargs["encoder_dim"], + "encoder_rates": kwargs["encoder_rates"], + "latent_dim": kwargs["latent_dim"], + "latent_channels": latent_channels, + "decoder_dim": kwargs["decoder_dim"], + "decoder_rates": kwargs["decoder_rates"], + # The reference's two hardcoded BigVGAN tables (16 kHz and 32 kHz) both pair rate `u` with kernel + # `2u` for even `u` and `2u - 1` for odd `u`, i.e. [5, 5, 2, ...] -> [9, 9, 4, ...]. + "decoder_kernel_sizes": [2 * rate - (rate % 2) for rate in kwargs["decoder_rates"]], + **MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG, + # Renamed from the original `sample_rate` to the diffusers audio convention. + "sampling_rate": kwargs["sample_rate"], + "latents_mean": wrapper_config["latents_mean"], + "latents_std": wrapper_config["latents_std"], + } + + +__all__ = [ + "MINIMAX_H3_AUDIO_VAE_FIXED_CONFIG", + "MINIMAX_H3_TEST_TRANSFORMER_CONFIG", + "MINIMAX_H3_TEST_VIDEO_VAE_CONFIG", + "MINIMAX_H3_TRANSFORMER_CONFIG", + "MINIMAX_H3_VIDEO_VAE_CONFIG", + "get_audio_vae_config", +] diff --git a/src/diffusers/loaders/conversion/configs/musicldm.py b/src/diffusers/loaders/conversion/configs/musicldm.py new file mode 100644 index 000000000000..c8fc8f247de9 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/musicldm.py @@ -0,0 +1,178 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the musicldm assembly recipe.""" + + +def create_unet_diffusers_config(original_config, image_size: int): + """ + Creates a UNet config for diffusers based on the config of the original MusicLDM model. + """ + unet_params = original_config["model"]["params"]["unet_config"]["params"] + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + + block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) + + cross_attention_dim = ( + unet_params["cross_attention_dim"] if "cross_attention_dim" in unet_params else block_out_channels + ) + + class_embed_type = "simple_projection" if "extra_film_condition_dim" in unet_params else None + projection_class_embeddings_input_dim = ( + unet_params["extra_film_condition_dim"] if "extra_film_condition_dim" in unet_params else None + ) + class_embeddings_concat = unet_params["extra_film_use_concat"] if "extra_film_use_concat" in unet_params else None + + config = { + "sample_size": image_size // vae_scale_factor, + "in_channels": unet_params["in_channels"], + "out_channels": unet_params["out_channels"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params["num_res_blocks"], + "cross_attention_dim": cross_attention_dim, + "class_embed_type": class_embed_type, + "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, + "class_embeddings_concat": class_embeddings_concat, + } + + return config + + +def create_vae_diffusers_config(original_config, checkpoint, image_size: int): + """ + Creates a VAE config for diffusers based on the config of the original MusicLDM model. Compared to the original + Stable Diffusion conversion, this function passes a *learnt* VAE scaling factor to the diffusers VAE. + """ + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] + + block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + scaling_factor = checkpoint["scale_factor"] if "scale_by_std" in original_config["model"]["params"] else 0.18215 + + config = { + "sample_size": image_size, + "in_channels": vae_params["in_channels"], + "out_channels": vae_params["out_ch"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params["z_channels"], + "layers_per_block": vae_params["num_res_blocks"], + "scaling_factor": float(scaling_factor), + } + return config + + +def create_transformers_vocoder_config(original_config): + """ + Creates a config for transformers SpeechT5HifiGan based on the config of the vocoder model. + """ + vocoder_params = original_config["model"]["params"]["vocoder_config"]["params"] + + config = { + "model_in_dim": vocoder_params["num_mels"], + "sampling_rate": vocoder_params["sampling_rate"], + "upsample_initial_channel": vocoder_params["upsample_initial_channel"], + "upsample_rates": list(vocoder_params["upsample_rates"]), + "upsample_kernel_sizes": list(vocoder_params["upsample_kernel_sizes"]), + "resblock_kernel_sizes": list(vocoder_params["resblock_kernel_sizes"]), + "resblock_dilation_sizes": [ + list(resblock_dilation) for resblock_dilation in vocoder_params["resblock_dilation_sizes"] + ], + "normalize_before": False, + } + + return config + + +DEFAULT_CONFIG = { + "model": { + "params": { + "linear_start": 0.0015, + "linear_end": 0.0195, + "timesteps": 1000, + "channels": 8, + "scale_by_std": True, + "unet_config": { + "target": "MusicLDM.latent_diffusion.openaimodel.UNetModel", + "params": { + "extra_film_condition_dim": 512, + "extra_film_use_concat": True, + "in_channels": 8, + "out_channels": 8, + "model_channels": 128, + "attention_resolutions": [8, 4, 2], + "num_res_blocks": 2, + "channel_mult": [1, 2, 3, 5], + "num_head_channels": 32, + }, + }, + "first_stage_config": { + "target": "MusicLDM.variational_autoencoder.autoencoder.AutoencoderKL", + "params": { + "embed_dim": 8, + "ddconfig": { + "z_channels": 8, + "resolution": 256, + "in_channels": 1, + "out_ch": 1, + "ch": 128, + "ch_mult": [1, 2, 4], + "num_res_blocks": 2, + }, + }, + }, + "vocoder_config": { + "target": "MusicLDM.first_stage_model.vocoder", + "params": { + "upsample_rates": [5, 4, 2, 2, 2], + "upsample_kernel_sizes": [16, 16, 8, 4, 4], + "upsample_initial_channel": 1024, + "resblock_kernel_sizes": [3, 7, 11], + "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]], + "num_mels": 64, + "sampling_rate": 16000, + }, + }, + }, + }, +} + +__all__ = [ + "DEFAULT_CONFIG", + "create_transformers_vocoder_config", + "create_unet_diffusers_config", + "create_vae_diffusers_config", +] diff --git a/src/diffusers/loaders/conversion/configs/prx.py b/src/diffusers/loaders/conversion/configs/prx.py new file mode 100644 index 000000000000..7f3c5d4dac56 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/prx.py @@ -0,0 +1,80 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the prx assembly recipe.""" + +from dataclasses import asdict, dataclass +from typing import Optional, Tuple + + +@dataclass(frozen=True) +class PRXBase: + context_in_dim: int = 2304 + hidden_size: int = 1792 + mlp_ratio: float = 3.5 + num_heads: int = 28 + depth: int = 16 + axes_dim: Tuple[int, int] = (32, 32) + theta: int = 10_000 + time_factor: float = 1000.0 + time_max_period: int = 10_000 + bottleneck_size: Optional[int] = None + resolution_embeds: bool = False + + +@dataclass(frozen=True) +class PRXFlux(PRXBase): + in_channels: int = 16 + patch_size: int = 2 + + +@dataclass(frozen=True) +class PRXDCAE(PRXBase): + in_channels: int = 32 + patch_size: int = 1 + + +@dataclass(frozen=True) +class PRXPixel(PRXBase): + # Pixel-space RGB diffusion (PRXPixel / 7B). + in_channels: int = 3 + patch_size: int = 16 + context_in_dim: int = 2048 # Qwen3-VL-Embedding-2B hidden size + hidden_size: int = 3584 + num_heads: int = 28 + depth: int = 24 + axes_dim: Tuple[int, int] = (64, 64) + bottleneck_size: int = 768 + resolution_embeds: bool = True + + +VARIANTS = {"flux": PRXFlux, "dc-ae": PRXDCAE, "pixel": PRXPixel} + + +def build_config(variant: str) -> dict: + if variant not in VARIANTS: + raise ValueError(f"Unsupported variant: {variant}. Choose from {list(VARIANTS)}") + config_dict = asdict(VARIANTS[variant]()) + config_dict["axes_dim"] = list(config_dict["axes_dim"]) + if config_dict["bottleneck_size"] is None: + # Keep config.json clean for variants that don't use the bottleneck. + config_dict.pop("bottleneck_size") + return config_dict + + +def create_scheduler_config(shift: float): + return {"_class_name": "FlowMatchEulerDiscreteScheduler", "num_train_timesteps": 1000, "shift": shift} + + +__all__ = ["build_config", "create_scheduler_config"] diff --git a/src/diffusers/loaders/conversion/configs/rae.py b/src/diffusers/loaders/conversion/configs/rae.py new file mode 100644 index 000000000000..0f878e79e19a --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/rae.py @@ -0,0 +1,38 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the rae assembly recipe.""" + +DECODER_CONFIGS = { + "ViTB": { + "decoder_hidden_size": 768, + "decoder_intermediate_size": 3072, + "decoder_num_attention_heads": 12, + "decoder_num_hidden_layers": 12, + }, + "ViTL": { + "decoder_hidden_size": 1024, + "decoder_intermediate_size": 4096, + "decoder_num_attention_heads": 16, + "decoder_num_hidden_layers": 24, + }, + "ViTXL": { + "decoder_hidden_size": 1152, + "decoder_intermediate_size": 4096, + "decoder_num_attention_heads": 16, + "decoder_num_hidden_layers": 28, + }, +} + +__all__ = ["DECODER_CONFIGS"] diff --git a/src/diffusers/loaders/conversion/configs/shap_e.py b/src/diffusers/loaders/conversion/configs/shap_e.py new file mode 100644 index 000000000000..554c1cc31c39 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/shap_e.py @@ -0,0 +1,52 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the shap_e assembly recipe.""" + +PRIOR_CONFIG = { + "num_attention_heads": 16, + "attention_head_dim": 1024 // 16, + "num_layers": 24, + "embedding_dim": 1024, + "num_embeddings": 1024, + "additional_embeddings": 0, + "time_embed_act_fn": "gelu", + "norm_in_type": "layer", + "encoder_hid_proj_type": None, + "added_emb_type": None, + "time_embed_dim": 1024 * 4, + "embedding_proj_dim": 768, + "clip_embed_dim": 1024 * 2, +} + +PRIOR_IMAGE_CONFIG = { + "num_attention_heads": 8, + "attention_head_dim": 1024 // 8, + "num_layers": 24, + "embedding_dim": 1024, + "num_embeddings": 1024, + "additional_embeddings": 0, + "time_embed_act_fn": "gelu", + "norm_in_type": "layer", + "embedding_proj_norm_type": "layer", + "encoder_hid_proj_type": None, + "added_emb_type": None, + "time_embed_dim": 1024 * 4, + "embedding_proj_dim": 1024, + "clip_embed_dim": 1024 * 2, +} + +RENDERER_CONFIG = {} + +__all__ = ["PRIOR_CONFIG", "PRIOR_IMAGE_CONFIG", "RENDERER_CONFIG"] diff --git a/src/diffusers/loaders/conversion/configs/skyreels_v2.py b/src/diffusers/loaders/conversion/configs/skyreels_v2.py new file mode 100644 index 000000000000..4d0a6e37d351 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/skyreels_v2.py @@ -0,0 +1,253 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the skyreels_v2 assembly recipe.""" + +from typing import Any + + +def get_transformer_config(model_type: str) -> dict[str, Any]: + if model_type == "SkyReels-V2-DF-1.3B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-DF-1.3B-540P", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 8960, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 12, + "inject_sample_info": True, + "num_layers": 30, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "SkyReels-V2-DF-14B-720P": + config = { + "model_id": "Skywork/SkyReels-V2-DF-14B-720P", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "SkyReels-V2-DF-14B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-DF-14B-540P", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "SkyReels-V2-T2V-14B-720P": + config = { + "model_id": "Skywork/SkyReels-V2-T2V-14B-720P", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "SkyReels-V2-T2V-14B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-T2V-14B-540P", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "SkyReels-V2-I2V-1.3B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-1.3B-540P", + "diffusers_config": { + "added_kv_proj_dim": 1536, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 8960, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 12, + "inject_sample_info": False, + "num_layers": 30, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + }, + } + elif model_type == "SkyReels-V2-I2V-14B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-14B-540P", + "diffusers_config": { + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + }, + } + elif model_type == "SkyReels-V2-I2V-14B-720P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-14B-720P", + "diffusers_config": { + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + }, + } + elif model_type == "SkyReels-V2-FLF2V-1.3B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-1.3B-540P", + "diffusers_config": { + "added_kv_proj_dim": 1536, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 8960, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 12, + "inject_sample_info": False, + "num_layers": 30, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + "pos_embed_seq_len": 514, + }, + } + elif model_type == "SkyReels-V2-FLF2V-14B-540P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-14B-540P", + "diffusers_config": { + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + "pos_embed_seq_len": 514, + }, + } + elif model_type == "SkyReels-V2-FLF2V-14B-720P": + config = { + "model_id": "Skywork/SkyReels-V2-I2V-14B-720P", + "diffusers_config": { + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "inject_sample_info": False, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "image_dim": 1280, + "pos_embed_seq_len": 514, + }, + } + return config + + +__all__ = ["get_transformer_config"] diff --git a/src/diffusers/loaders/conversion/configs/stable_audio3.py b/src/diffusers/loaders/conversion/configs/stable_audio3.py new file mode 100644 index 000000000000..c55f37810120 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/stable_audio3.py @@ -0,0 +1,219 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the stable_audio3 assembly recipe.""" + +import math +from typing import Optional + +import torch + + +def _infer_trb_depth(ref_sd: dict, base: str) -> int: + """Count how many TRB blocks exist under `base` (e.g. 'encoder.layers').""" + depth = 0 + while f"{base}.{depth}.new_tokens" in ref_sd: + depth += 1 + return depth + + +def _infer_vae_config(ref_sd: dict, model_config: Optional[dict] = None) -> dict: + """ + Infer AutoencoderSAME config from checkpoint tensor shapes. + + Some hyper-parameters (the TRB stride, the sliding-window width and the sinusoidal-FFN layer counts) are NOT + recoverable from the weights alone — with ``variable_stride`` the encoder/decoder ``new_tokens`` collapse to a + single shared token, so the stride is invisible. When ``model_config`` (the parsed ``model_config.json``) is + provided, those values are read directly from it; otherwise production SAME-L/S defaults are used. + """ + # Bottleneck scale shape: (1, latent_dim, 1) + latent_dim = ref_sd.get( + "pretransform.model.bottleneck.scaling_factor", + ref_sd.get("pretransform.model.bottleneck.scale", torch.zeros(1, 256, 1)), + ).shape[1] + + # Encoder linear weight: (latent_dim, channel_dims[-1]) + # We detect enc_depth to find the linear key + enc_base = "pretransform.model.encoder.layers" + enc_depth = _infer_trb_depth(ref_sd, enc_base) + if enc_depth == 0: + enc_depth = 1 + enc_linear_key = f"{enc_base}.{enc_depth + 1}.weight" + if enc_linear_key in ref_sd: + enc_final_ch = ref_sd[enc_linear_key].shape[1] # (latent_dim, enc_final_ch) + else: + enc_final_ch = 768 # SAME-S default + + # TRB 0 new_tokens shape (encoder): (1, 1, out_channels) + trb0_nt = ref_sd.get(f"{enc_base}.0.new_tokens") + enc_out_ch = trb0_nt.shape[2] if trb0_nt is not None else enc_final_ch + + # TRB 0 mapping.weight_v shape (encoder): (out_ch, in_ch, kernel) + mapping_wv = ref_sd.get(f"{enc_base}.0.mapping.weight_v") + if mapping_wv is not None: + patched_in = mapping_wv.shape[1] # = audio_channels * patch_size + else: + patched_in = 512 # default: 2ch * 256 patch + + # Infer transformer depth per TRB: + trb_trans_depth = sum( + 1 + for k in ref_sd + if k.startswith(f"{enc_base}.0.transformers.") + and k.endswith(".new_tokens") is False + and ".pre_norm.alpha" in k + ) + if trb_trans_depth == 0: + trb_trans_depth = 6 # default + + # Dim heads inferred from q_norm shape: (dim_heads,) + q_norm_key = f"{enc_base}.0.transformers.0.self_attn.q_norm.gamma" + dim_heads = ref_sd[q_norm_key].shape[0] if q_norm_key in ref_sd else 64 + + # channels base: choose so that enc_out_ch = channels * c_mults[0] + # We use c_mults = [6] for both SAME-S/L + c_mults = [6] + enc_channels_base = enc_out_ch // c_mults[0] + + # ── Weight-invisible hyper-parameters ──────────────────────────────────── + # These come from model_config.json when available (see docstring). + audio_channels = 2 + sliding_window = 1 + encoder_sinusoidal_blocks = [0] * enc_depth + decoder_sinusoidal_blocks = [0] * enc_depth + + if model_config is not None: + ae_cfg = model_config["model"]["pretransform"]["config"] + enc_cfg = ae_cfg["encoder"]["config"] + dec_cfg = ae_cfg["decoder"]["config"] + strides = list(enc_cfg["strides"]) + transformer_depths = list(enc_cfg["transformer_depths"]) + c_mults = list(enc_cfg["c_mults"]) + enc_channels_base = enc_cfg["channels"] + latent_dim = enc_cfg.get("latent_dim", latent_dim) + dim_heads = enc_cfg.get("dim_heads", dim_heads) + audio_channels = ae_cfg.get("io_channels", audio_channels) + downsampling_ratio = ae_cfg["downsampling_ratio"] + patch_size = downsampling_ratio // int(math.prod(strides)) + # sliding_window in the reference is a per-side list like [1, 1]; take the (symmetric) half-width. + sw = enc_cfg.get("sliding_window") or [sliding_window] + sliding_window = sw[0] + encoder_sinusoidal_blocks = list(enc_cfg.get("sinusoidal_blocks", encoder_sinusoidal_blocks)) + decoder_sinusoidal_blocks = list(dec_cfg.get("sinusoidal_blocks", decoder_sinusoidal_blocks)) + else: + # Stride is NOT recoverable from weights under variable_stride; assume the production value of 16. + strides = [16] * enc_depth + transformer_depths = [trb_trans_depth] * enc_depth + patch_size = patched_in // audio_channels + + return { + "audio_channels": audio_channels, + "patch_size": patch_size, + "encoder_channels": enc_channels_base, + "encoder_c_mults": c_mults, + "encoder_strides": strides, + "encoder_transformer_depths": transformer_depths, + "latent_dim": latent_dim, + "use_differential_attention": True, + "dim_heads": dim_heads, + "ff_mult": 3, + "sliding_window": sliding_window, + "encoder_sinusoidal_blocks": encoder_sinusoidal_blocks, + "decoder_sinusoidal_blocks": decoder_sinusoidal_blocks, + "sampling_rate": 44100, + } + + +def _infer_dit_config(ref_sd: dict) -> dict: + """Infer StableAudio3DiTModel config from checkpoint tensor shapes.""" + p = "model.model." + + # embed_dim from to_timestep_embed.0.weight shape (embed_dim, features_dim) + ts_w = ref_sd.get(f"{p}to_timestep_embed.0.weight") + embed_dim = ts_w.shape[0] if ts_w is not None else 1536 + timestep_features_dim = ts_w.shape[1] if ts_w is not None else 256 + + # depth: count transformer blocks (RMSNorm → pre_norm.gamma) + depth = 0 + while f"{p}transformer.layers.{depth}.pre_norm.gamma" in ref_sd: + depth += 1 + if depth == 0: + depth = 24 + + # num_heads: from self_attn.q_norm.gamma shape (dim_heads,) and embed_dim + q_norm_key = f"{p}transformer.layers.0.self_attn.q_norm.gamma" + dim_heads = ref_sd[q_norm_key].shape[0] if q_norm_key in ref_sd else 64 + num_heads = embed_dim // dim_heads + + # cond_token_dim / global_cond_dim from the projection in-weights (embed_dim, *) + cond_w = ref_sd.get(f"{p}to_cond_embed.0.weight") + cond_token_dim = cond_w.shape[1] if cond_w is not None else 768 + glob_w = ref_sd.get(f"{p}to_global_embed.0.weight") + global_cond_dim = glob_w.shape[1] if glob_w is not None else 768 + + # io_channels from preprocess_conv.weight (io_ch, io_ch, 1) + pc_w = ref_sd.get(f"{p}preprocess_conv.weight") + io_channels = pc_w.shape[0] if pc_w is not None else 256 + + # ff_mult from ff.ff.0.proj.weight (inner*2, embed_dim) + ff_w = ref_sd.get(f"{p}transformer.layers.0.ff.ff.0.proj.weight") + ff_mult = (ff_w.shape[0] // (2 * embed_dim)) if ff_w is not None else 4 + + # local_add_cond_dim from to_local_embed.0.weight (embed_dim, local_add_cond_dim) + loc_w = ref_sd.get(f"{p}transformer.layers.0.to_local_embed.0.weight") + local_add_cond_dim = loc_w.shape[1] if loc_w is not None else 257 + + # num_memory_tokens from memory_tokens (num_memory_tokens, embed_dim) + mem = ref_sd.get(f"{p}transformer.memory_tokens") + num_memory_tokens = mem.shape[0] if mem is not None else 64 + + # differential: self_attn.to_qkv rows = embed_dim*5 (differential) vs *3 (standard) + qkv_w = ref_sd.get(f"{p}transformer.layers.0.self_attn.to_qkv.weight") + use_differential = qkv_w is not None and qkv_w.shape[0] == embed_dim * 5 + + return { + "io_channels": io_channels, + "patch_size": 1, + "embed_dim": embed_dim, + "depth": depth, + "num_heads": num_heads, + "cond_token_dim": cond_token_dim, + "global_cond_dim": global_cond_dim, + "local_add_cond_dim": local_add_cond_dim, + "timestep_features_dim": timestep_features_dim, + "ff_mult": ff_mult, + "num_memory_tokens": num_memory_tokens, + "use_differential_attention": use_differential, + } + + +def _infer_duration_embedder_config(ref_sd: dict) -> dict: + """Infer StableAudio3DurationEmbedder config from checkpoint shapes.""" + w = ref_sd.get("conditioner.conditioners.seconds_total.embedder.embedding.1.weight") + if w is not None: + output_dim = w.shape[0] + fourier_dim = w.shape[1] + else: + output_dim, fourier_dim = 768, 256 + return { + "output_dim": output_dim, + "fourier_dim": fourier_dim, + "min_val": 0.0, + "max_val": 384.0, + "min_freq": 0.5, + "max_freq": 10000.0, + } + + +__all__ = ["_infer_dit_config", "_infer_duration_embedder_config", "_infer_vae_config"] diff --git a/src/diffusers/loaders/conversion/configs/unclip.py b/src/diffusers/loaders/conversion/configs/unclip.py new file mode 100644 index 000000000000..c861885282d6 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/unclip.py @@ -0,0 +1,86 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the unclip assembly recipe.""" + +PRIOR_CONFIG = {} + +DECODER_CONFIG = { + "sample_size": 64, + "layers_per_block": 3, + "down_block_types": ( + "ResnetDownsampleBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + "SimpleCrossAttnDownBlock2D", + ), + "up_block_types": ( + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "SimpleCrossAttnUpBlock2D", + "ResnetUpsampleBlock2D", + ), + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "block_out_channels": (320, 640, 960, 1280), + "in_channels": 3, + "out_channels": 6, + "cross_attention_dim": 1536, + "class_embed_type": "identity", + "attention_head_dim": 64, + "resnet_time_scale_shift": "scale_shift", +} + +SUPER_RES_UNET_FIRST_STEPS_CONFIG = { + "sample_size": 256, + "layers_per_block": 3, + "down_block_types": ( + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + ), + "up_block_types": ( + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + ), + "block_out_channels": (320, 640, 960, 1280), + "in_channels": 6, + "out_channels": 3, + "add_attention": False, +} + +SUPER_RES_UNET_LAST_STEP_CONFIG = { + "sample_size": 256, + "layers_per_block": 3, + "down_block_types": ( + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + "ResnetDownsampleBlock2D", + ), + "up_block_types": ( + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + "ResnetUpsampleBlock2D", + ), + "block_out_channels": (320, 640, 960, 1280), + "in_channels": 6, + "out_channels": 3, + "add_attention": False, +} + +__all__ = ["DECODER_CONFIG", "PRIOR_CONFIG", "SUPER_RES_UNET_FIRST_STEPS_CONFIG", "SUPER_RES_UNET_LAST_STEP_CONFIG"] diff --git a/src/diffusers/loaders/conversion/configs/unidiffuser.py b/src/diffusers/loaders/conversion/configs/unidiffuser.py new file mode 100644 index 000000000000..507ec95b05d9 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/unidiffuser.py @@ -0,0 +1,214 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the unidiffuser assembly recipe.""" + +from argparse import Namespace + + +SCHEDULER_CONFIG = Namespace( + **{ + "beta_start": 0.00085, + "beta_end": 0.012, + "beta_schedule": "scaled_linear", + "solver_order": 3, + } +) + + +def create_vae_diffusers_config(config_type): + # Hardcoded for now + if config_type == "test": + vae_config = create_vae_diffusers_config_test() + elif config_type == "big": + vae_config = create_vae_diffusers_config_big() + else: + raise NotImplementedError( + f"Config type {config_type} is not implemented, currently only config types" + " 'test' and 'big' are available." + ) + return vae_config + + +def create_unidiffuser_unet_config(config_type, version): + # Hardcoded for now + if config_type == "test": + unet_config = create_unidiffuser_unet_config_test() + elif config_type == "big": + unet_config = create_unidiffuser_unet_config_big() + else: + raise NotImplementedError( + f"Config type {config_type} is not implemented, currently only config types" + " 'test' and 'big' are available." + ) + # Unidiffuser-v1 uses data type embeddings + if version == 1: + unet_config["use_data_type_embedding"] = True + return unet_config + + +def create_text_decoder_config(config_type): + # Hardcoded for now + if config_type == "test": + text_decoder_config = create_text_decoder_config_test() + elif config_type == "big": + text_decoder_config = create_text_decoder_config_big() + else: + raise NotImplementedError( + f"Config type {config_type} is not implemented, currently only config types" + " 'test' and 'big' are available." + ) + return text_decoder_config + + +def create_vae_diffusers_config_test(): + vae_config = { + "sample_size": 32, + "in_channels": 3, + "out_channels": 3, + "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D"], + "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D"], + "block_out_channels": [32, 64], + "latent_channels": 4, + "layers_per_block": 1, + } + return vae_config + + +def create_unidiffuser_unet_config_test(): + unet_config = { + "text_dim": 32, + "clip_img_dim": 32, + "num_text_tokens": 77, + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "dropout": 0.0, + "norm_num_groups": 32, + "attention_bias": False, + "sample_size": 16, + "patch_size": 2, + "activation_fn": "gelu", + "num_embeds_ada_norm": 1000, + "norm_type": "layer_norm", + "block_type": "unidiffuser", + "pre_layer_norm": False, + "use_timestep_embedding": False, + "norm_elementwise_affine": True, + "use_patch_pos_embed": False, + "ff_final_dropout": True, + "use_data_type_embedding": False, + } + return unet_config + + +def create_text_decoder_config_test(): + text_decoder_config = { + "prefix_length": 77, + "prefix_inner_dim": 32, + "prefix_hidden_dim": 32, + "vocab_size": 1025, # 1024 + 1 for new EOS token + "n_positions": 1024, + "n_embd": 32, + "n_layer": 5, + "n_head": 4, + "n_inner": 37, + "activation_function": "gelu", + "resid_pdrop": 0.1, + "embd_pdrop": 0.1, + "attn_pdrop": 0.1, + "layer_norm_epsilon": 1e-5, + "initializer_range": 0.02, + } + return text_decoder_config + + +def create_vae_diffusers_config_big(): + vae_config = { + "sample_size": 256, + "in_channels": 3, + "out_channels": 3, + "down_block_types": ["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"], + "up_block_types": ["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"], + "block_out_channels": [128, 256, 512, 512], + "latent_channels": 4, + "layers_per_block": 2, + } + return vae_config + + +def create_unidiffuser_unet_config_big(): + unet_config = { + "text_dim": 64, + "clip_img_dim": 512, + "num_text_tokens": 77, + "num_attention_heads": 24, + "attention_head_dim": 64, + "in_channels": 4, + "out_channels": 4, + "num_layers": 30, + "dropout": 0.0, + "norm_num_groups": 32, + "attention_bias": False, + "sample_size": 64, + "patch_size": 2, + "activation_fn": "gelu", + "num_embeds_ada_norm": 1000, + "norm_type": "layer_norm", + "block_type": "unidiffuser", + "pre_layer_norm": False, + "use_timestep_embedding": False, + "norm_elementwise_affine": True, + "use_patch_pos_embed": False, + "ff_final_dropout": True, + "use_data_type_embedding": False, + } + return unet_config + + +def create_text_decoder_config_big(): + text_decoder_config = { + "prefix_length": 77, + "prefix_inner_dim": 768, + "prefix_hidden_dim": 64, + "vocab_size": 50258, # 50257 + 1 for new EOS token + "n_positions": 1024, + "n_embd": 768, + "n_layer": 12, + "n_head": 12, + "n_inner": 3072, + "activation_function": "gelu", + "resid_pdrop": 0.1, + "embd_pdrop": 0.1, + "attn_pdrop": 0.1, + "layer_norm_epsilon": 1e-5, + "initializer_range": 0.02, + } + return text_decoder_config + + +__all__ = [ + "SCHEDULER_CONFIG", + "create_text_decoder_config", + "create_text_decoder_config_big", + "create_text_decoder_config_test", + "create_unidiffuser_unet_config", + "create_unidiffuser_unet_config_big", + "create_unidiffuser_unet_config_test", + "create_vae_diffusers_config", + "create_vae_diffusers_config_big", + "create_vae_diffusers_config_test", +] diff --git a/src/diffusers/loaders/conversion/configs/versatile_diffusion.py b/src/diffusers/loaders/conversion/configs/versatile_diffusion.py new file mode 100644 index 000000000000..b8d9e0f4aeec --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/versatile_diffusion.py @@ -0,0 +1,182 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the versatile_diffusion assembly recipe.""" + +from argparse import Namespace + + +SCHEDULER_CONFIG = Namespace( + **{ + "beta_linear_start": 0.00085, + "beta_linear_end": 0.012, + "timesteps": 1000, + "scale_factor": 0.18215, + } +) + +IMAGE_UNET_CONFIG = Namespace( + **{ + "input_channels": 4, + "model_channels": 320, + "output_channels": 4, + "num_noattn_blocks": [2, 2, 2, 2], + "channel_mult": [1, 2, 4, 4], + "with_attn": [True, True, True, False], + "num_heads": 8, + "context_dim": 768, + "use_checkpoint": True, + } +) + +TEXT_UNET_CONFIG = Namespace( + **{ + "input_channels": 768, + "model_channels": 320, + "output_channels": 768, + "num_noattn_blocks": [2, 2, 2, 2], + "channel_mult": [1, 2, 4, 4], + "second_dim": [4, 4, 4, 4], + "with_attn": [True, True, True, False], + "num_heads": 8, + "context_dim": 768, + "use_checkpoint": True, + } +) + +AUTOENCODER_CONFIG = Namespace( + **{ + "double_z": True, + "z_channels": 4, + "resolution": 256, + "in_channels": 3, + "out_ch": 3, + "ch": 128, + "ch_mult": [1, 2, 4, 4], + "num_res_blocks": 2, + "attn_resolutions": [], + "dropout": 0.0, + } +) + + +def create_image_unet_diffusers_config(unet_params): + """ + Creates a config for the diffusers based on the config of the VD model. + """ + + block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if unet_params.with_attn[i] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if unet_params.with_attn[-i - 1] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + if not all(n == unet_params.num_noattn_blocks[0] for n in unet_params.num_noattn_blocks): + raise ValueError("Not all num_res_blocks are equal, which is not supported in this script.") + + config = { + "sample_size": None, + "in_channels": unet_params.input_channels, + "out_channels": unet_params.output_channels, + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params.num_noattn_blocks[0], + "cross_attention_dim": unet_params.context_dim, + "attention_head_dim": unet_params.num_heads, + } + + return config + + +def create_text_unet_diffusers_config(unet_params): + """ + Creates a config for the diffusers based on the config of the VD model. + """ + + block_out_channels = [unet_params.model_channels * mult for mult in unet_params.channel_mult] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlockFlat" if unet_params.with_attn[i] else "DownBlockFlat" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlockFlat" if unet_params.with_attn[-i - 1] else "UpBlockFlat" + up_block_types.append(block_type) + resolution //= 2 + + if not all(n == unet_params.num_noattn_blocks[0] for n in unet_params.num_noattn_blocks): + raise ValueError("Not all num_res_blocks are equal, which is not supported in this script.") + + config = { + "sample_size": None, + "in_channels": (unet_params.input_channels, 1, 1), + "out_channels": (unet_params.output_channels, 1, 1), + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params.num_noattn_blocks[0], + "cross_attention_dim": unet_params.context_dim, + "attention_head_dim": unet_params.num_heads, + } + + return config + + +def create_vae_diffusers_config(vae_params): + """ + Creates a config for the diffusers based on the config of the VD model. + """ + + block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + config = { + "sample_size": vae_params.resolution, + "in_channels": vae_params.in_channels, + "out_channels": vae_params.out_ch, + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params.z_channels, + "layers_per_block": vae_params.num_res_blocks, + } + return config + + +__all__ = [ + "AUTOENCODER_CONFIG", + "IMAGE_UNET_CONFIG", + "SCHEDULER_CONFIG", + "TEXT_UNET_CONFIG", + "create_image_unet_diffusers_config", + "create_text_unet_diffusers_config", + "create_vae_diffusers_config", +] diff --git a/src/diffusers/loaders/conversion/configs/wan.py b/src/diffusers/loaders/conversion/configs/wan.py new file mode 100644 index 000000000000..4c28aa55a683 --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/wan.py @@ -0,0 +1,273 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the wan assembly recipe.""" + +from typing import Any, Dict + + +def get_transformer_config(model_type: str) -> Dict[str, Any]: + if model_type == "Wan-T2V-1.3B": + config = { + "model_id": "StevenZhang/Wan2.1-T2V-1.3B-Diff", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 8960, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 12, + "num_layers": 30, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan-T2V-14B": + config = { + "model_id": "StevenZhang/Wan2.1-T2V-14B-Diff", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan-I2V-14B-480p": + config = { + "model_id": "StevenZhang/Wan2.1-I2V-14B-480P-Diff", + "diffusers_config": { + "image_dim": 1280, + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan-I2V-14B-720p": + config = { + "model_id": "StevenZhang/Wan2.1-I2V-14B-720P-Diff", + "diffusers_config": { + "image_dim": 1280, + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan-FLF2V-14B-720P": + config = { + "model_id": "ypyp/Wan2.1-FLF2V-14B-720P", # This is just a placeholder + "diffusers_config": { + "image_dim": 1280, + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "rope_max_seq_len": 1024, + "pos_embed_seq_len": 257 * 2, + }, + } + elif model_type == "Wan-VACE-1.3B": + config = { + "model_id": "Wan-AI/Wan2.1-VACE-1.3B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 8960, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 12, + "num_layers": 30, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "vace_layers": [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28], + "vace_in_channels": 96, + }, + } + elif model_type == "Wan-VACE-14B": + config = { + "model_id": "Wan-AI/Wan2.1-VACE-14B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "vace_layers": [0, 5, 10, 15, 20, 25, 30, 35], + "vace_in_channels": 96, + }, + } + elif model_type == "Wan2.2-VACE-Fun-14B": + config = { + "model_id": "alibaba-pai/Wan2.2-VACE-Fun-A14B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "vace_layers": [0, 5, 10, 15, 20, 25, 30, 35], + "vace_in_channels": 96, + }, + } + elif model_type == "Wan2.2-I2V-14B-720p": + config = { + "model_id": "Wan-AI/Wan2.2-I2V-A14B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan2.2-T2V-A14B": + config = { + "model_id": "Wan-AI/Wan2.2-T2V-A14B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 16, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan2.2-TI2V-5B": + config = { + "model_id": "Wan-AI/Wan2.2-TI2V-5B", + "diffusers_config": { + "added_kv_proj_dim": None, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 14336, + "freq_dim": 256, + "in_channels": 48, + "num_attention_heads": 24, + "num_layers": 30, + "out_channels": 48, + "patch_size": [1, 2, 2], + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + }, + } + elif model_type == "Wan2.2-Animate-14B": + config = { + "model_id": "Wan-AI/Wan2.2-Animate-14B", + "diffusers_config": { + "image_dim": 1280, + "added_kv_proj_dim": 5120, + "attention_head_dim": 128, + "cross_attn_norm": True, + "eps": 1e-06, + "ffn_dim": 13824, + "freq_dim": 256, + "in_channels": 36, + "num_attention_heads": 40, + "num_layers": 40, + "out_channels": 16, + "patch_size": (1, 2, 2), + "qk_norm": "rms_norm_across_heads", + "text_dim": 4096, + "rope_max_seq_len": 1024, + "pos_embed_seq_len": None, + "motion_encoder_size": 512, # Start of Wan Animate-specific configs + "motion_style_dim": 512, + "motion_dim": 20, + "motion_encoder_dim": 512, + "face_encoder_hidden_dim": 1024, + "face_encoder_num_heads": 4, + "inject_face_latents_blocks": 5, + }, + } + return config + + +__all__ = ["get_transformer_config"] diff --git a/src/diffusers/loaders/conversion/configs/zero123.py b/src/diffusers/loaders/conversion/configs/zero123.py new file mode 100644 index 000000000000..35563b9e038c --- /dev/null +++ b/src/diffusers/loaders/conversion/configs/zero123.py @@ -0,0 +1,147 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Original configuration helpers and model presets for the zero123 assembly recipe.""" + + +def create_unet_diffusers_config(original_config, image_size: int, controlnet=False): + """ + Creates a config for the diffusers based on the config of the LDM model. + """ + if controlnet: + unet_params = original_config["model"]["params"]["control_stage_config"]["params"] + else: + if ( + "unet_config" in original_config["model"]["params"] + and original_config["model"]["params"]["unet_config"] is not None + ): + unet_params = original_config["model"]["params"]["unet_config"]["params"] + else: + unet_params = original_config["model"]["params"]["network_config"]["params"] + + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + + block_out_channels = [unet_params["model_channels"] * mult for mult in unet_params["channel_mult"]] + + down_block_types = [] + resolution = 1 + for i in range(len(block_out_channels)): + block_type = "CrossAttnDownBlock2D" if resolution in unet_params["attention_resolutions"] else "DownBlock2D" + down_block_types.append(block_type) + if i != len(block_out_channels) - 1: + resolution *= 2 + + up_block_types = [] + for i in range(len(block_out_channels)): + block_type = "CrossAttnUpBlock2D" if resolution in unet_params["attention_resolutions"] else "UpBlock2D" + up_block_types.append(block_type) + resolution //= 2 + + if unet_params["transformer_depth"] is not None: + transformer_layers_per_block = ( + unet_params["transformer_depth"] + if isinstance(unet_params["transformer_depth"], int) + else list(unet_params["transformer_depth"]) + ) + else: + transformer_layers_per_block = 1 + + vae_scale_factor = 2 ** (len(vae_params["ch_mult"]) - 1) + + head_dim = unet_params["num_heads"] if "num_heads" in unet_params else None + use_linear_projection = ( + unet_params["use_linear_in_transformer"] if "use_linear_in_transformer" in unet_params else False + ) + if use_linear_projection: + # stable diffusion 2-base-512 and 2-768 + if head_dim is None: + head_dim_mult = unet_params["model_channels"] // unet_params["num_head_channels"] + head_dim = [head_dim_mult * c for c in list(unet_params["channel_mult"])] + + class_embed_type = None + addition_embed_type = None + addition_time_embed_dim = None + projection_class_embeddings_input_dim = None + context_dim = None + + if unet_params["context_dim"] is not None: + context_dim = ( + unet_params["context_dim"] + if isinstance(unet_params["context_dim"], int) + else unet_params["context_dim"][0] + ) + + if "num_classes" in unet_params: + if unet_params["num_classes"] == "sequential": + if context_dim in [2048, 1280]: + # SDXL + addition_embed_type = "text_time" + addition_time_embed_dim = 256 + else: + class_embed_type = "projection" + assert "adm_in_channels" in unet_params + projection_class_embeddings_input_dim = unet_params["adm_in_channels"] + else: + raise NotImplementedError(f"Unknown conditional unet num_classes config: {unet_params['num_classes']}") + + config = { + "sample_size": image_size // vae_scale_factor, + "in_channels": unet_params["in_channels"], + "down_block_types": tuple(down_block_types), + "block_out_channels": tuple(block_out_channels), + "layers_per_block": unet_params["num_res_blocks"], + "cross_attention_dim": context_dim, + "attention_head_dim": head_dim, + "use_linear_projection": use_linear_projection, + "class_embed_type": class_embed_type, + "addition_embed_type": addition_embed_type, + "addition_time_embed_dim": addition_time_embed_dim, + "projection_class_embeddings_input_dim": projection_class_embeddings_input_dim, + "transformer_layers_per_block": transformer_layers_per_block, + } + + if controlnet: + config["conditioning_channels"] = unet_params["hint_channels"] + else: + config["out_channels"] = unet_params["out_channels"] + config["up_block_types"] = tuple(up_block_types) + + return config + + +def create_vae_diffusers_config(original_config, image_size: int): + """ + Creates a config for the diffusers based on the config of the LDM model. + """ + vae_params = original_config["model"]["params"]["first_stage_config"]["params"]["ddconfig"] + _ = original_config["model"]["params"]["first_stage_config"]["params"]["embed_dim"] + + block_out_channels = [vae_params["ch"] * mult for mult in vae_params["ch_mult"]] + down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels) + up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels) + + config = { + "sample_size": image_size, + "in_channels": vae_params["in_channels"], + "out_channels": vae_params["out_ch"], + "down_block_types": tuple(down_block_types), + "up_block_types": tuple(up_block_types), + "block_out_channels": tuple(block_out_channels), + "latent_channels": vae_params["z_channels"], + "layers_per_block": vae_params["num_res_blocks"], + } + return config + + +__all__ = ["create_unet_diffusers_config", "create_vae_diffusers_config"] diff --git a/src/diffusers/loaders/conversion/consistency_decoder.py b/src/diffusers/loaders/conversion/consistency_decoder.py new file mode 100644 index 000000000000..dd03d6d80f47 --- /dev/null +++ b/src/diffusers/loaders/conversion/consistency_decoder.py @@ -0,0 +1,100 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def consistency_decoder_conversion(config): + """Map the original Python decoder graph plus the LDM encoder/quantization component.""" + encoder_config = { + "block_out_channels": config["encoder_block_out_channels"], + "layers_per_block": config["encoder_layers_per_block"], + "mid_block_add_attention": True, + "down_block_types": config["encoder_down_block_types"], + } + encoder = ldm_vae_conversion(encoder_config) + mapping = { + old: new for old, new in encoder.mapping.items() if new.startswith("encoder.") or new.startswith("quant_conv.") + } + rules = [rule for rule in encoder.rules if all(key.startswith("encoder.") for key in rule.diffusers)] + modules = [ + ("embed_image.f", "conv_in"), + ("output.gn", "conv_norm_out"), + ("output.f", "conv_out"), + ("embed_time.f_1", "time_embedding.linear_1"), + ("embed_time.f_2", "time_embedding.linear_2"), + ] + mapping["decoder.embed_time.emb.weight"] = "decoder_unet.time_proj.weight" + channels, layers = config["decoder_block_out_channels"], config["decoder_layers_per_block"] + count = len(channels) + resnets = [(f"mid.{j}", f"mid_block.resnets.{j}", False) for j in range(2)] + previous = channels[0] + for i, channel in enumerate(channels): + for j in range(layers): + resnets.append((f"down.{i}.{j}", f"down_blocks.{i}.resnets.{j}", previous != channel)) + previous = channel + if i < count - 1: + resnets.append((f"down.{i}.{layers}", f"down_blocks.{i}.downsamplers.0", False)) + previous = channels[-1] + for i, channel in enumerate(reversed(channels)): + level = count - 1 - i + for j in range(layers + 1): + skip = channels[max(level - 1, 0)] if j == layers else channel + resnets.append((f"up.{level}.{j}", f"up_blocks.{i}.resnets.{j}", previous + skip != channel)) + previous = channel + if i < count - 1: + resnets.append((f"up.{level}.{layers + 1}", f"up_blocks.{i}.upsamplers.0", False)) + for old, new, shortcut in resnets: + pairs = [("gn_1", "norm1"), ("gn_2", "norm2"), ("f_1", "conv1"), ("f_2", "conv2"), ("f_t", "time_emb_proj")] + if shortcut: + pairs.append(("f_s", "conv_shortcut")) + modules.extend((f"{old}.{a}", f"{new}.{b}") for a, b in pairs) + if config["decoder_add_attention"]: + raise ValueError("The original consistency decoder has no attention layers.") + mapping.update( + {f"decoder.{old}.{p}": f"decoder_unet.{new}.{p}" for old, new in modules for p in ("weight", "bias")} + ) + if config.get("original_format") == "consistency_decoder_jit": + + def jit_key(key): + if not key.startswith("decoder."): + return key + key = key.removeprefix("decoder.") + if key == "embed_time.emb.weight": + return "decoder.embed_time.weight" + match = re.match(r"(down|up)\.(\d+)\.(\d+)\.(.*)", key) + if match: + direction, stage, index, tail = match.groups() + sampler_index = layers if direction == "down" else layers + 1 + block = ( + ("downsamp" if direction == "down" else "upsamp") + if int(index) == sampler_index + else "conv_" + index + ) + key = f"{direction}_{stage}_{block}.{tail}" + else: + key = re.sub(r"^mid\.(\d+)\.", r"mid_\1.", key) + prefix, leaf = key.rsplit(".", 1) + module = prefix.rsplit(".", 1)[-1] + if leaf == "bias": + leaf = "b" + elif leaf == "weight": + leaf = "g" if module.startswith("gn") else "w" + return f"decoder.blocks.{prefix}.{leaf}" + + mapping = {jit_key(old): new for old, new in mapping.items()} + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/controlnet.py b/src/diffusers/loaders/conversion/controlnet.py new file mode 100644 index 000000000000..f198d1b73919 --- /dev/null +++ b/src/diffusers/loaders/conversion/controlnet.py @@ -0,0 +1,21 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_unet import _ldm_unet_mapping + + +def controlnet_conversion(config): + return Conversion(mapping=_ldm_unet_mapping(config, controlnet=True)) diff --git a/src/diffusers/loaders/conversion/core.py b/src/diffusers/loaders/conversion/core.py new file mode 100644 index 000000000000..683c91b468ee --- /dev/null +++ b/src/diffusers/loaders/conversion/core.py @@ -0,0 +1,119 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from collections.abc import Mapping +from dataclasses import dataclass, field +from types import MappingProxyType + +import torch + +from .transforms import Identity, Transform + + +@dataclass(frozen=True) +class Rule: + """Map an ordered group of original keys to Diffusers keys through a reversible tensor operation.""" + + original: tuple[str, ...] + diffusers: tuple[str, ...] + transform: Transform = field(default_factory=Identity) + + def __post_init__(self): + for name in ("original", "diffusers"): + keys = getattr(self, name) + if isinstance(keys, str) or not keys: + raise ValueError(f"Rule.{name} must be a nonempty sequence of keys, not a string.") + keys = tuple(keys) + object.__setattr__(self, name, keys) + if not keys or any(not isinstance(key, str) or not key for key in keys): + raise ValueError(f"Rule.{name} keys must be nonempty strings.") + if len(set(keys)) != len(keys): + raise ValueError(f"Repeated {name} key in rule: {keys}.") + + +@dataclass(frozen=True) +class Conversion: + """Convert a component state dict in either direction using a single definition. + + `mapping` is shorthand for exact one-to-one renames; `rules` describe grouped tensor operations. Every source key + must be consumed once and every destination key written once. Known auxiliary keys and component prefixes must be + handled before conversion. Both methods leave the input dictionary and tensors unchanged, but outputs may share + storage with inputs. No model construction, file I/O, device movement, or dtype conversion is performed here. + """ + + mapping: Mapping[str, str] = field(default_factory=dict) + rules: tuple[Rule, ...] = () + original_keys: frozenset[str] = field(init=False) + diffusers_keys: frozenset[str] = field(init=False) + _rules: tuple[Rule, ...] = field(init=False, repr=False) + lossless: bool = field(init=False) + + def __post_init__(self): + object.__setattr__(self, "mapping", MappingProxyType(dict(self.mapping))) + object.__setattr__(self, "rules", tuple(self.rules)) + rules = tuple(Rule((old,), (new,)) for old, new in self.mapping.items()) + self.rules + original_keys, diffusers_keys = set(), set() + for rule in rules: + repeated_inputs = original_keys.intersection(rule.original) + repeated_outputs = diffusers_keys.intersection(rule.diffusers) + if repeated_inputs or repeated_outputs: + raise ValueError( + f"Conversion keys must be unique: repeated original keys {sorted(repeated_inputs)}, " + f"repeated Diffusers keys {sorted(repeated_outputs)}." + ) + original_keys.update(rule.original) + diffusers_keys.update(rule.diffusers) + object.__setattr__(self, "original_keys", frozenset(original_keys)) + object.__setattr__(self, "diffusers_keys", frozenset(diffusers_keys)) + object.__setattr__(self, "_rules", rules) + object.__setattr__(self, "lossless", all(getattr(rule.transform, "lossless", True) for rule in rules)) + + def to_diffusers(self, state_dict: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Convert original component weights to Diffusers weights without modifying the inputs.""" + return self._convert(state_dict, reverse=False) + + def to_original(self, state_dict: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Convert Diffusers component weights to original weights without modifying the inputs.""" + return self._convert(state_dict, reverse=True) + + def _convert(self, state_dict: Mapping[str, torch.Tensor], reverse: bool) -> dict[str, torch.Tensor]: + return dict(self.iter_converted(state_dict, reverse=reverse)) + + def iter_converted(self, state_dict: Mapping[str, torch.Tensor], *, reverse: bool = False): + """Yield converted key/tensor pairs one rule at a time, after validating complete input coverage. + + File readers can supply a lazy mapping and writers can flush shards between rules. A grouped operation still + needs all of its input and output tensors in memory. As with the dictionary methods, tensors may alias inputs. + """ + source_keys = self.diffusers_keys if reverse else self.original_keys + missing = sorted(source_keys - state_dict.keys()) + unexpected = sorted(state_dict.keys() - source_keys) + direction = "Diffusers -> original" if reverse else "original -> Diffusers" + if missing or unexpected: + raise ValueError(f"Cannot convert {direction}: missing keys {missing}; unexpected keys {unexpected}.") + + for rule in self._rules: + source = rule.diffusers if reverse else rule.original + target = rule.original if reverse else rule.diffusers + transform = rule.transform.inverse if reverse else rule.transform.forward + try: + tensors = tuple(state_dict[key] for key in source) + if any(not isinstance(tensor, torch.Tensor) for tensor in tensors): + raise ValueError("Every converted parameter must be a tensor.") + outputs = transform(tensors) + if len(outputs) != len(target): + raise ValueError(f"Expected {len(target)} output tensors, got {len(outputs)}.") + except (ValueError, RuntimeError) as error: + raise ValueError(f"Failed to convert {direction}, {source} -> {target}: {error}") from error + yield from zip(target, outputs) diff --git a/src/diffusers/loaders/conversion/cosmos.py b/src/diffusers/loaders/conversion/cosmos.py new file mode 100644 index 000000000000..69c3d3271b62 --- /dev/null +++ b/src/diffusers/loaders/conversion/cosmos.py @@ -0,0 +1,83 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def cosmos_conversion(config): + version = config.get("original_format", "cosmos2") + if version not in ("cosmos1", "cosmos2"): + raise ValueError("original_format must be 'cosmos1' or 'cosmos2'.") + first = version == "cosmos1" + modules = [ + ("x_embedder.proj.1", "patch_embed.proj"), + ("t_embedder.1.linear_1", "time_embed.t_embedder.linear_1"), + ("t_embedder.1.linear_2", "time_embed.t_embedder.linear_2"), + ("affline_norm" if first else "t_embedding_norm", "time_embed.norm"), + ("final_layer.linear", "proj_out"), + ] + for j in (1, 2): + source = f"final_layer.{'adaLN_modulation' if first else 'adaln_modulation'}.{j}" + modules.append((source, f"norm_out.linear_{j}")) + mapping = {} + if config["extra_pos_embed_type"] == "learnable": + for axis in ("t", "h", "w"): + source = "extra_pos_embedder" if first else "learnable_pos_embed" + mapping[f"{source}.pos_emb_{axis}"] = f"learnable_pos_embed.pos_emb_{axis}" + for i in range(config["num_layers"]): + old = f"blocks.block{i}" if first else f"blocks.{i}" + new = f"transformer_blocks.{i}" + for j, name in enumerate(("self_attn", "cross_attn", "mlp")): + old_norm = f"{old}.blocks.{j}.adaLN_modulation" if first else f"{old}.adaln_modulation_{name}" + for layer in (1, 2): + modules.append((f"{old_norm}.{layer}", f"{new}.norm{j + 1}.linear_{layer}")) + for j, name in enumerate(("self_attn", "cross_attn")): + source = f"{old}.blocks.{j}.block.attn" if first else f"{old}.{name}" + target = f"{new}.attn{j + 1}" + pairs = ( + ( + ("to_q.0", "to_q"), + ("to_q.1", "norm_q"), + ("to_k.0", "to_k"), + ("to_k.1", "norm_k"), + ("to_v.0", "to_v"), + ("to_out.0", "to_out.0"), + ) + if first + else ( + ("q_proj", "to_q"), + ("q_norm", "norm_q"), + ("k_proj", "to_k"), + ("k_norm", "norm_k"), + ("v_proj", "to_v"), + ("output_proj", "to_out.0"), + ) + ) + modules.extend((f"{source}.{a}", f"{target}.{b}") for a, b in pairs) + if j == 1 and config["img_context_dim_in"]: + modules.extend( + (f"{source}.{name}", f"{target}.{name}") + for name in ("q_img", "k_img", "v_img", "q_img_norm", "k_img_norm") + ) + source = f"{old}.blocks.2.block" if first else f"{old}.mlp" + modules.extend((f"{source}.layer{j}", f"{new}.ff.{name}") for j, name in ((1, "net.0.proj"), (2, "net.2"))) + mapping.update({old + ".weight": new + ".weight" for old, new in modules}) + for name, enabled in ( + ("crossattn_proj.0", config["use_crossattn_projection"]), + ("img_context_proj.0", config["img_context_dim_in"]), + ): + if enabled: + mapping.update({f"{name}.{p}": f"{name}.{p}" for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/cosmos_controlnet.py b/src/diffusers/loaders/conversion/cosmos_controlnet.py new file mode 100644 index 000000000000..c90e5b90a910 --- /dev/null +++ b/src/diffusers/loaders/conversion/cosmos_controlnet.py @@ -0,0 +1,41 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .cosmos import cosmos_conversion + + +def cosmos_controlnet_conversion(config): + # Shared modules come from the base checkpoint; keep their source namespace distinct from control weights. + base = cosmos_conversion({**config, "num_layers": config["n_controlnet_blocks"], "original_format": "cosmos2"}) + mapping = {} + for old, new in base.mapping.items(): + if old.startswith("final_layer."): + continue + if old.startswith("blocks."): + mapping[old.replace("blocks.", "control_blocks.", 1)] = new.replace( + "transformer_blocks.", "control_blocks.", 1 + ) + else: + mapping["base." + old] = new.replace("patch_embed.", "patch_embed_base.", 1) + mapping["control_embedder.proj.1.weight"] = "patch_embed.proj.weight" + for i in range(config["n_controlnet_blocks"]): + for p in ("weight", "bias"): + key = f"control_blocks.{i}.after_proj.{p}" + mapping[key] = key + if i == 0: + key = f"control_blocks.{i}.before_proj.{p}" + mapping[key] = key + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/cosmos_vae.py b/src/diffusers/loaders/conversion/cosmos_vae.py new file mode 100644 index 000000000000..43980231da7b --- /dev/null +++ b/src/diffusers/loaders/conversion/cosmos_vae.py @@ -0,0 +1,78 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +from .core import Conversion + + +def cosmos_vae_conversion(config): + modules = [("quant_conv.conv3d", "quant_conv"), ("post_quant_conv.conv3d", "post_quant_conv")] + projections, resnets, attentions = [], [], [] + spatial = int(math.log2(config["spatial_compression_ratio"])) - int(math.log2(config["patch_size"])) + temporal = int(math.log2(config["temporal_compression_ratio"])) - int(math.log2(config["patch_size"])) + for component, field in (("encoder", "encoder_block_out_channels"), ("decoder", "decode_block_out_channels")): + channels = list(config[field]) + if component == "decoder": + channels.reverse() + count = len(channels) - 1 + modules.append((f"{component}.norm_out.norm", f"{component}.norm_out.norm")) + projections.extend((f"{component}.{name}", f"{component}.{name}") for name in ("conv_in", "conv_out")) + resnets.extend( + (f"{component}.mid.block_{j + 1}", f"{component}.mid_block.resnets.{j}", False) for j in range(2) + ) + attentions.append((f"{component}.mid.attn_1", f"{component}.mid_block", 0)) + resolution = config["resolution"] // config["patch_size"] + if component == "decoder": + resolution //= 2 ** (count - 1) + for i in range(count): + encoder = component == "encoder" + original_index = i if encoder else count - 1 - i + direction = "down" if encoder else "up" + old, new = f"{component}.{direction}.{original_index}", f"{component}.{direction}_blocks.{i}" + depth = config["num_layers"] + int(not encoder) + for j in range(depth): + resnets.append((f"{old}.block.{j}", f"{new}.resnets.{j}", j == 0 and channels[i] != channels[i + 1])) + if resolution in config["attention_resolutions"]: + attentions.append((f"{old}.attn.{j}", new, j)) + if i < count - 1: + sampler = direction + "sample" + if encoder: + enabled = (i < spatial, i < temporal) + resolution //= 2 + else: + time_up = 0 < i < temporal + 1 + space_up = time_up or (i < spatial and spatial > temporal) + enabled = (time_up, space_up) + resolution *= 2 + for j, active in enumerate((*enabled, any(enabled)), 1): + if active: + modules.append((f"{old}.{sampler}.conv{j}.conv3d", f"{new}.{sampler}rs.0.conv{j}")) + for old, new, shortcut in resnets: + modules.extend((f"{old}.{norm}.norm", f"{new}.{norm}.norm") for norm in ("norm1", "norm2")) + projections.extend((f"{old}.{conv}", f"{new}.{conv}") for conv in ("conv1", "conv2")) + if shortcut: + modules.append((old + ".nin_shortcut.conv3d", new + ".conv_shortcut")) + for old, new in projections: + modules.extend([(old + ".0.conv3d", new + ".conv_s"), (old + ".1.conv3d", new + ".conv_t")]) + for old, new, index in attentions: + for j, kind in enumerate(("attentions", "temp_attentions")): + a, b = f"{old}.{j}", f"{new}.{kind}.{index}" + modules.append((a + ".norm.norm", b + ".norm.norm")) + modules.extend( + (f"{a}.{source}.conv3d", f"{b}.{target}") + for source, target in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("proj_out", "to_out.0")) + ) + return Conversion(mapping={f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) diff --git a/src/diffusers/loaders/conversion/dance_diffusion.py b/src/diffusers/loaders/conversion/dance_diffusion.py new file mode 100644 index 000000000000..71397facd827 --- /dev/null +++ b/src/diffusers/loaders/conversion/dance_diffusion.py @@ -0,0 +1,104 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Reshape, Split + + +def dance_diffusion_conversion(config): + count = len(config["down_block_types"]) + channels = config["block_out_channels"] + modules, resnets, attentions = [], [], [] + mapping = {} + fourier_dim = (config["time_embedding_dim"] or channels[0] * 2) // 2 + rules = [Rule(("timestep_embed.weight",), ("time_proj.weight",), Reshape((fourier_dim, 1), (fourier_dim,)))] + + def source_prefix(depth): + return "net" if depth == 0 else "net.3." + "main.7." * (depth - 1) + "main" + + previous = config["in_channels"] + config["extra_in_channels"] + for i, channel in enumerate(channels): + old = source_prefix(i) + for j in range(3): + resnets.append( + (f"{old}.{j if i == 0 else 2 * j + 1}", f"down_blocks.{i}.resnets.{j}", previous != channel, False) + ) + previous = channel + if "Attn" in config["down_block_types"][i]: + attentions.append((f"{old}.{2 * j + 2}", f"down_blocks.{i}.attentions.{j}", channel)) + if i > 0: + mapping[old + ".0.kernel"] = f"down_blocks.{i}.down.kernel" + old = source_prefix(count) + mapping[old + ".0.kernel"] = "mid_block.down.kernel" + mapping[old + ".14.kernel"] = "mid_block.up.kernel" + for j, source in enumerate((1, 3, 5, 8, 10, 12)): + resnets.append((f"{old}.{source}", f"mid_block.resnets.{j}", False, False)) + attentions.append((f"{old}.{source + 1}", f"mid_block.attentions.{j}", channels[-1])) + previous = channels[-1] + for i, kind in enumerate(config["up_block_types"]): + depth = count - i - 1 + old = source_prefix(depth) + output = channels[depth - 1] if i < len(config["up_block_types"]) - 1 else config["out_channels"] + middle = output if "Attn" in kind else previous + for j in range(3): + input_channel = 2 * previous if j == 0 else middle + output_channel = output if j == 2 else middle + resnets.append( + ( + f"{old}.{j + 4 if depth == 0 else 2 * j + 8}", + f"up_blocks.{i}.resnets.{j}", + input_channel != output_channel, + depth == 0 and j == 2, + ) + ) + if "Attn" in kind: + attentions.append((f"{old}.{2 * j + 9}", f"up_blocks.{i}.attentions.{j}", output_channel)) + previous = output + if depth > 0: + mapping[old + ".14.kernel"] = f"up_blocks.{i}.up.kernel" + for old, new, shortcut, last in resnets: + modules.extend( + (f"{old}.main.{i}", f"{new}.{name}") for i, name in ((0, "conv_1"), (1, "group_norm_1"), (3, "conv_2")) + ) + if not last: + modules.append((old + ".main.4", new + ".group_norm_2")) + if shortcut: + mapping[old + ".skip.weight"] = new + ".conv_skip.weight" + for old, new, channel in attentions: + modules.append((old + ".norm", new + ".group_norm")) + rules.append( + Rule( + (old + ".qkv_proj.weight",), + tuple(f"{new}.{part}.weight" for part in ("query", "key", "value")), + Chain((Reshape((3 * channel, channel, 1), (3 * channel, channel)), Split((channel,) * 3))), + ) + ) + rules.append( + Rule( + (old + ".qkv_proj.bias",), + tuple(f"{new}.{part}.bias" for part in ("query", "key", "value")), + Split((channel,) * 3), + ) + ) + rules.append( + Rule( + (old + ".out_proj.weight",), + (new + ".proj_attn.weight",), + Reshape((channel, channel, 1), (channel, channel)), + ) + ) + mapping[old + ".out_proj.bias"] = new + ".proj_attn.bias" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/diffuser_rl.py b/src/diffusers/loaders/conversion/diffuser_rl.py new file mode 100644 index 000000000000..4b1ab1dd5304 --- /dev/null +++ b/src/diffusers/loaders/conversion/diffuser_rl.py @@ -0,0 +1,80 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def diffuser_rl_conversion(config): + if config["layers_per_block"] != 1 or not config["use_timestep_embedding"]: + raise ValueError("Diffuser temporal checkpoints require two residual blocks per level and timestep embedding.") + value = config["out_block_type"] == "ValueFunction" + original_format = config.get("original_format", "diffuser_rl") + if original_format not in ("diffuser_rl", "diffuser_rl_legacy"): + raise ValueError("Diffuser original_format must be 'diffuser_rl' or 'diffuser_rl_legacy'.") + sampler_index = 2 if value or original_format == "diffuser_rl_legacy" else 3 + modules = [("time_mlp.1", "time_mlp.linear_1"), ("time_mlp.3", "time_mlp.linear_2")] + resnets = [] + channels = config["block_out_channels"] + previous = config["in_channels"] + config["extra_in_channels"] + for i, channel in enumerate(channels): + old = f"{'blocks' if value else 'downs'}.{i}" + for j in range(2): + resnets.append((f"{old}.{j}", f"down_blocks.{i}.resnets.{j}", previous != channel)) + previous = channel + if i < len(channels) - 1 or config["downsample_each_block"]: + modules.append((f"{old}.{sampler_index}.conv", f"down_blocks.{i}.downsample.conv")) + if value: + resnets.extend([("mid_block1", "mid_block.res1", True), ("mid_block2", "mid_block.res2", True)]) + modules.extend( + [ + ("mid_down1.conv", "mid_block.down1.conv"), + ("mid_down2.conv", "mid_block.down2.conv"), + ("final_block.0", "out_block.final_block.0"), + ("final_block.2", "out_block.final_block.2"), + ] + ) + else: + resnets.extend([("mid_block1", "mid_block.resnets.0", False), ("mid_block2", "mid_block.resnets.1", False)]) + previous = channels[-1] + for i in range(len(config["up_block_types"])): + output = channels[-i - 2] if i < len(config["up_block_types"]) - 1 else channels[0] + resnets.extend( + [ + (f"ups.{i}.0", f"up_blocks.{i}.resnets.0", 2 * previous != output), + (f"ups.{i}.1", f"up_blocks.{i}.resnets.1", False), + ] + ) + if i < len(channels) - 1: + modules.append((f"ups.{i}.{sampler_index}.conv", f"up_blocks.{i}.upsample.conv")) + previous = output + modules.extend( + [ + ("final_conv.0.block.0", "out_block.final_conv1d_1"), + ("final_conv.0.block.2", "out_block.final_conv1d_gn"), + ("final_conv.1", "out_block.final_conv1d_2"), + ] + ) + for old, new, shortcut in resnets: + modules.append((old + ".time_mlp.1", new + ".time_emb")) + for i, name in enumerate(("conv_in", "conv_out")): + modules.extend( + [ + (f"{old}.blocks.{i}.block.0", f"{new}.{name}.conv1d"), + (f"{old}.blocks.{i}.block.2", f"{new}.{name}.group_norm"), + ] + ) + if shortcut: + modules.append((old + ".residual_conv", new + ".residual_conv")) + return Conversion(mapping={f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) diff --git a/src/diffusers/loaders/conversion/dit.py b/src/diffusers/loaders/conversion/dit.py new file mode 100644 index 000000000000..469b559d7d9e --- /dev/null +++ b/src/diffusers/loaders/conversion/dit.py @@ -0,0 +1,71 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import MergeEqual, Reverse, Split + + +def dit_conversion(config): + modules = [ + ("x_embedder.proj", "pos_embed.proj"), + ("final_layer.adaLN_modulation.1", "proj_out_1"), + ("final_layer.linear", "proj_out_2"), + ] + mapping, rules = {}, [] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + embeddings = [ + ("t_embedder.mlp.0", "timestep_embedder.linear_1"), + ("t_embedder.mlp.2", "timestep_embedder.linear_2"), + ] + for old, new in embeddings: + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.{p}",), + tuple(f"transformer_blocks.{i}.norm1.emb.{new}.{p}" for i in range(config["num_layers"])), + Reverse(MergeEqual(config["num_layers"])), + ) + ) + rules.append( + Rule( + ("y_embedder.embedding_table.weight",), + tuple( + f"transformer_blocks.{i}.norm1.emb.class_embedder.embedding_table.weight" + for i in range(config["num_layers"]) + ), + Reverse(MergeEqual(config["num_layers"])), + ) + ) + for i in range(config["num_layers"]): + old, new = f"blocks.{i}", f"transformer_blocks.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("adaLN_modulation.1", "norm1.linear"), + ("attn.proj", "attn1.to_out.0"), + ("mlp.fc1", "ff.net.0.proj"), + ("mlp.fc2", "ff.net.2"), + ) + ) + for p in ("weight", "bias") if config["attention_bias"] else ("weight",): + rules.append( + Rule( + (f"{old}.attn.qkv.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/ernie_image.py b/src/diffusers/loaders/conversion/ernie_image.py new file mode 100644 index 000000000000..d540eccd4849 --- /dev/null +++ b/src/diffusers/loaders/conversion/ernie_image.py @@ -0,0 +1,48 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ernie_image_conversion(config): + modules = [ + "x_embedder.proj", + "time_embedding.linear_1", + "time_embedding.linear_2", + "adaLN_modulation.1", + "final_norm.linear", + "final_linear", + ] + keys = [] + if config["text_in_dim"] != config["hidden_size"]: + keys.append("text_proj.weight") + for i in range(config["num_layers"]): + prefix = f"layers.{i}" + names = [ + "adaLN_sa_ln", + "adaLN_mlp_ln", + "self_attention.to_q", + "self_attention.to_k", + "self_attention.to_v", + "self_attention.to_out.0", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.linear_fc2", + ] + if config["qk_layernorm"]: + names.extend(["self_attention.norm_q", "self_attention.norm_k"]) + keys.extend(f"{prefix}.{name}.weight" for name in names) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/flux.py b/src/diffusers/loaders/conversion/flux.py new file mode 100644 index 000000000000..6467e5094906 --- /dev/null +++ b/src/diffusers/loaders/conversion/flux.py @@ -0,0 +1,87 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def flux_conversion(config): + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + parameters = ("weight", "bias") + modules = [("txt_in", "context_embedder"), ("img_in", "x_embedder"), ("final_layer.linear", "proj_out")] + mapping, rules = ({}, []) + embedding = "time_text_embed" + embeddings = [("time_in", "timestep_embedder")] + embeddings.append(("vector_in", "text_embedder")) + if config.get("guidance_embeds", False): + embeddings.append(("guidance_in", "guidance_embedder")) + for old, new in embeddings: + modules.extend( + [(old + ".in_layer", f"{embedding}.{new}.linear_1"), (old + ".out_layer", f"{embedding}.{new}.linear_2")] + ) + for parameter in parameters: + rules.append( + Rule( + (f"final_layer.adaLN_modulation.1.{parameter}",), + (f"norm_out.linear.{parameter}",), + ReorderChunks((1, 0)), + ) + ) + for i in range(config["num_layers"]): + old, new = (f"double_blocks.{i}", f"transformer_blocks.{i}") + modules.extend( + [(old + ".img_mod.lin", new + ".norm1.linear"), (old + ".txt_mod.lin", new + ".norm1_context.linear")] + ) + for modality, attention, mlp, norms in ( + ("img", ("to_q", "to_k", "to_v"), "ff", ("norm_q", "norm_k")), + ("txt", ("add_q_proj", "add_k_proj", "add_v_proj"), "ff_context", ("norm_added_q", "norm_added_k")), + ): + for parameter in parameters: + rules.append( + Rule( + (f"{old}.{modality}_attn.qkv.{parameter}",), + tuple((f"{new}.attn.{name}.{parameter}" for name in attention)), + Split((hidden_size,) * 3), + ) + ) + for source, target in zip(("query_norm", "key_norm"), norms): + mapping[f"{old}.{modality}_attn.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + modules.extend( + [ + (f"{old}.{modality}_mlp.0", f"{new}.{mlp}." + "net.0.proj"), + (f"{old}.{modality}_mlp.2", f"{new}.{mlp}." + "net.2"), + ( + f"{old}.{modality}_attn.proj", + f"{new}.attn." + ("to_out.0" if modality == "img" else "to_add_out"), + ), + ] + ) + for i in range(config["num_single_layers"]): + old, new = (f"single_blocks.{i}", f"single_transformer_blocks.{i}") + modules.append((old + ".modulation.lin", new + ".norm.linear")) + for parameter in parameters: + rules.append( + Rule( + (f"{old}.linear1.{parameter}",), + tuple( + (f"{new}.{name}.{parameter}" for name in ("attn.to_q", "attn.to_k", "attn.to_v", "proj_mlp")) + ), + Split((hidden_size,) * 3 + (4 * hidden_size,)), + ) + ) + modules.append((old + ".linear2", new + ".proj_out")) + for source, target in (("query_norm", "norm_q"), ("key_norm", "norm_k")): + mapping[f"{old}.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in parameters}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/flux2.py b/src/diffusers/loaders/conversion/flux2.py new file mode 100644 index 000000000000..14406912d9a3 --- /dev/null +++ b/src/diffusers/loaders/conversion/flux2.py @@ -0,0 +1,74 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def flux2_conversion(config): + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + parameters = ("weight",) + modules = [("txt_in", "context_embedder"), ("img_in", "x_embedder"), ("final_layer.linear", "proj_out")] + mapping, rules = ({}, []) + embedding = "time_guidance_embed" + embeddings = [("time_in", "timestep_embedder")] + if config.get("guidance_embeds", True): + embeddings.append(("guidance_in", "guidance_embedder")) + for old, new in embeddings: + modules.extend( + [(old + ".in_layer", f"{embedding}.{new}.linear_1"), (old + ".out_layer", f"{embedding}.{new}.linear_2")] + ) + for parameter in parameters: + rules.append( + Rule( + (f"final_layer.adaLN_modulation.1.{parameter}",), + (f"norm_out.linear.{parameter}",), + ReorderChunks((1, 0)), + ) + ) + for name in ("double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation"): + modules.append((name + ".lin", name + ".linear")) + for i in range(config["num_layers"]): + old, new = (f"double_blocks.{i}", f"transformer_blocks.{i}") + for modality, attention, mlp, norms in ( + ("img", ("to_q", "to_k", "to_v"), "ff", ("norm_q", "norm_k")), + ("txt", ("add_q_proj", "add_k_proj", "add_v_proj"), "ff_context", ("norm_added_q", "norm_added_k")), + ): + for parameter in parameters: + rules.append( + Rule( + (f"{old}.{modality}_attn.qkv.{parameter}",), + tuple((f"{new}.attn.{name}.{parameter}" for name in attention)), + Split((hidden_size,) * 3), + ) + ) + for source, target in zip(("query_norm", "key_norm"), norms): + mapping[f"{old}.{modality}_attn.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + modules.extend( + [ + (f"{old}.{modality}_mlp.0", f"{new}.{mlp}." + "linear_in"), + (f"{old}.{modality}_mlp.2", f"{new}.{mlp}." + "linear_out"), + ( + f"{old}.{modality}_attn.proj", + f"{new}.attn." + ("to_out.0" if modality == "img" else "to_add_out"), + ), + ] + ) + for i in range(config["num_single_layers"]): + old, new = (f"single_blocks.{i}", f"single_transformer_blocks.{i}") + modules.extend([(old + ".linear1", new + ".attn.to_qkv_mlp_proj"), (old + ".linear2", new + ".attn.to_out")]) + for source, target in (("query_norm", "norm_q"), ("key_norm", "norm_k")): + mapping[f"{old}.norm.{source}.scale"] = f"{new}.attn.{target}.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in parameters}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/flux2_vae.py b/src/diffusers/loaders/conversion/flux2_vae.py new file mode 100644 index 000000000000..754ae5b1b7eb --- /dev/null +++ b/src/diffusers/loaders/conversion/flux2_vae.py @@ -0,0 +1,40 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def flux2_vae_conversion(config): + encoder = ldm_vae_conversion(config) + decoder = ldm_vae_conversion( + {**config, "block_out_channels": config["decoder_block_out_channels"] or config["block_out_channels"]} + ) + mapping = {} + for old, new in encoder.mapping.items(): + if new.startswith("encoder."): + mapping[old] = new + elif new.startswith("quant_conv."): + mapping["encoder." + old] = new + for old, new in decoder.mapping.items(): + if new.startswith("decoder."): + mapping[old] = new + elif new.startswith("post_quant_conv."): + mapping["decoder." + old] = new + mapping.update({f"bn.{name}": f"bn.{name}" for name in ("running_mean", "running_var", "num_batches_tracked")}) + rules = tuple(rule for rule in encoder.rules if rule.diffusers[0].startswith("encoder.")) + tuple( + rule for rule in decoder.rules if rule.diffusers[0].startswith("decoder.") + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/flux_ip_adapter.py b/src/diffusers/loaders/conversion/flux_ip_adapter.py new file mode 100644 index 000000000000..e4d6ee03e85c --- /dev/null +++ b/src/diffusers/loaders/conversion/flux_ip_adapter.py @@ -0,0 +1,30 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def flux_ip_adapter_conversion(config): + mapping = { + f"ip_adapter_proj_model.{name}.{p}": f"image_proj.{name}.{p}" + for name in ("norm", "proj") + for p in ("weight", "bias") + } + for i in range(config["num_layers"]): + for part in ("k", "v"): + for p in ("weight", "bias"): + mapping[f"double_blocks.{i}.processor.ip_adapter_double_stream_{part}_proj.{p}"] = ( + f"ip_adapter.{i}.to_{part}_ip.{p}" + ) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/gligen.py b/src/diffusers/loaders/conversion/gligen.py new file mode 100644 index 000000000000..4b6830c81358 --- /dev/null +++ b/src/diffusers/loaders/conversion/gligen.py @@ -0,0 +1,44 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .ldm_unet import _ldm_unet_mapping + + +def gligen_conversion(config): + mapping = _ldm_unet_mapping(config, controlnet=False) + blocks = [ + (old.removesuffix(".norm1.weight"), new.removesuffix(".norm1.weight")) + for old, new in mapping.items() + if ".transformer_blocks." in new and new.endswith(".norm1.weight") + ] + for old, new in blocks: + for name in ("linear", "norm1", "norm2", "attn.to_out.0", "ff.net.0.proj", "ff.net.2"): + mapping.update({f"{old}.fuser.{name}.{p}": f"{new}.fuser.{name}.{p}" for p in ("weight", "bias")}) + for part in ("q", "k", "v"): + mapping[f"{old}.fuser.attn.to_{part}.weight"] = f"{new}.fuser.attn.to_{part}.weight" + for name in ("alpha_attn", "alpha_dense"): + mapping[f"{old}.fuser.{name}"] = f"{new}.fuser.{name}" + image = config["attention_type"] == "gated-text-image" + names = ("linears_text", "linears_image") if image else ("linears",) + for name in names: + for i in (0, 2, 4): + for p in ("weight", "bias"): + key = f"position_net.{name}.{i}.{p}" + mapping[key] = key + for name in (("null_text_feature", "null_image_feature") if image else ("null_positive_feature",)) + ( + "null_position_feature", + ): + mapping[f"position_net.{name}"] = f"position_net.{name}" + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/hidream.py b/src/diffusers/loaders/conversion/hidream.py new file mode 100644 index 000000000000..c5f09ab11d8c --- /dev/null +++ b/src/diffusers/loaders/conversion/hidream.py @@ -0,0 +1,45 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def hidream_conversion(config): + modules = ["x_embedder.proj", "final_layer.linear", "final_layer.adaLN_modulation.1"] + modules.extend( + f"{prefix}.linear_{i}" + for prefix in ("t_embedder.timestep_embedder", "p_embedder.pooled_embedder") + for i in (1, 2) + ) + keys = [ + f"caption_projection.{i}.linear.weight" for i in range(config["num_layers"] + config["num_single_layers"] + 1) + ] + for group, count in ( + ("double_stream_blocks", config["num_layers"]), + ("single_stream_blocks", config["num_single_layers"]), + ): + for i in range(count): + prefix = f"{group}.{i}.block" + modules.append(prefix + ".adaLN_modulation.1") + for suffix in ("", "_t") if group == "double_stream_blocks" else ("",): + modules.extend(f"{prefix}.attn1.{name}{suffix}" for name in ("to_q", "to_k", "to_v", "to_out")) + keys.extend(f"{prefix}.attn1.{name}{suffix}.weight" for name in ("q_rms_norm", "k_rms_norm")) + experts = ["ff_i.shared_experts"] + [f"ff_i.experts.{j}" for j in range(config["num_routed_experts"])] + if group == "double_stream_blocks": + experts.append("ff_t") + keys.extend(f"{prefix}.{expert}.w{j}.weight" for expert in experts for j in (1, 2, 3)) + keys.append(prefix + ".ff_i.gate.weight") + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/hifigan.py b/src/diffusers/loaders/conversion/hifigan.py new file mode 100644 index 000000000000..880217fb17cf --- /dev/null +++ b/src/diffusers/loaders/conversion/hifigan.py @@ -0,0 +1,47 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from .core import Conversion, Rule +from .transforms import WithConstants + + +def hifigan_conversion(config): + modules = [("conv_pre", "conv_pre"), ("conv_post", "conv_post")] + count = len(config["upsample_rates"]) + modules.extend((f"ups.{i}", f"upsampler.{i}") for i in range(count)) + for i in range(count): + for j, dilations in enumerate(config["resblock_dilation_sizes"]): + index = i * len(config["resblock_kernel_sizes"]) + j + modules.extend( + (f"resblocks.{index}.convs{part}.{k}", f"resblocks.{index}.convs{part}.{k}") + for part in (1, 2) + for k in range(len(dilations)) + ) + mapping = {f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")} + rules = () + if config["normalize_before"]: + mapping.update({"mean": "mean", "scale": "scale"}) + else: + del mapping["conv_pre.weight"] + rules = ( + Rule( + ("conv_pre.weight",), + ("conv_pre.weight", "mean", "scale"), + WithConstants((torch.zeros(config["model_in_dim"]), torch.ones(config["model_in_dim"]))), + ), + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_dit.py b/src/diffusers/loaders/conversion/hunyuan_dit.py new file mode 100644 index 000000000000..8348ead6e126 --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_dit.py @@ -0,0 +1,94 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def hunyuan_dit_conversion(config): + mapping, rules = _hunyuan_dit_rules(config, controlnet=False) + return Conversion(mapping=mapping, rules=rules) + + +def _hunyuan_dit_rules(config, controlnet): + modules = [("x_embedder.proj", "pos_embed.proj")] + mapping = { + "text_embedding_padding": "text_embedding_padding", + "pooler.positional_embedding": "time_extra_emb.pooler.positional_embedding", + } + modules.extend( + (f"pooler.{name}", f"time_extra_emb.pooler.{name}") for name in ("q_proj", "k_proj", "v_proj", "c_proj") + ) + modules.extend( + (f"{source}.{i}", f"{target}.linear_{j}") + for source, target in ( + ("t_embedder.mlp", "time_extra_emb.timestep_embedder"), + ("mlp_t5", "text_embedder"), + ("extra_embedder", "time_extra_emb.extra_embedder"), + ) + for i, j in ((0, 1), (2, 2)) + ) + if config["use_style_cond_and_image_meta_size"]: + mapping["style_embedder.weight"] = "time_extra_emb.style_embedder.weight" + count = config["transformer_num_layers"] // 2 - 1 if controlnet else config["num_layers"] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + rules = [] + for i in range(count): + prefix = f"blocks.{i}" + modules.extend( + (f"{prefix}.{old}", f"{prefix}.{new}") + for old, new in ( + ("norm1", "norm1.norm"), + ("default_modulation.1", "norm1.linear"), + ("norm2", "norm3"), + ("norm3", "norm2"), + ("mlp.fc1", "ff.net.0.proj"), + ("mlp.fc2", "ff.net.2"), + ("attn2.q_proj", "attn2.to_q"), + ) + ) + for attn in ("attn1", "attn2"): + modules.extend( + (f"{prefix}.{attn}.{old}", f"{prefix}.{attn}.{new}") + for old, new in (("q_norm", "norm_q"), ("k_norm", "norm_k"), ("out_proj", "to_out.0")) + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attn1.Wqkv.{p}",), + tuple(f"{prefix}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + rules.append( + Rule( + (f"{prefix}.attn2.kv_proj.{p}",), + tuple(f"{prefix}.attn2.to_{part}.{p}" for part in ("k", "v")), + Split((hidden,) * 2), + ) + ) + if not controlnet and i > count // 2: + modules.extend((f"{prefix}.{name}", f"{prefix}.{name}") for name in ("skip_norm", "skip_linear")) + if controlnet: + modules.append(("input_block", "input_block")) + modules.extend((f"controlnet_blocks.{i}", f"controlnet_blocks.{i}") for i in range(count)) + else: + modules.append(("final_layer.linear", "proj_out")) + for p in ("weight", "bias"): + rules.append( + Rule((f"final_layer.adaLN_modulation.1.{p}",), (f"norm_out.linear.{p}",), ReorderChunks((1, 0))) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return mapping, tuple(rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_dit_controlnet.py b/src/diffusers/loaders/conversion/hunyuan_dit_controlnet.py new file mode 100644 index 000000000000..8c8d7a3233dc --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_dit_controlnet.py @@ -0,0 +1,22 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .hunyuan_dit import _hunyuan_dit_rules + + +def hunyuan_dit_controlnet_conversion(config): + mapping, rules = _hunyuan_dit_rules(config, controlnet=True) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_image.py b/src/diffusers/loaders/conversion/hunyuan_image.py new file mode 100644 index 000000000000..068cd7394d0d --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_image.py @@ -0,0 +1,83 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .hunyuan_video import hunyuan_video_conversion +from .transforms import Split + + +def hunyuan_image_conversion(config): + base = hunyuan_video_conversion(config) + mapping = {} + for old, new in base.mapping.items(): + if old.startswith("vector_in."): + continue + if old.startswith(("time_in.", "guidance_in.")): + new = new.replace("time_text_embed.", "time_guidance_embed.") + if old.startswith("single_blocks."): + old = old.replace(".linear2.", ".linear2.fc.") + mapping[old] = new + rules = [rule for rule in base.rules if not rule.original[0].startswith(("double_blocks.", "single_blocks."))] + modules = [] + if config["text_embed_2_dim"] is not None: + modules.append(("byt5_in.layernorm", "context_embedder_2.norm")) + modules.extend((f"byt5_in.fc{i}", f"context_embedder_2.linear_{i}") for i in (1, 2, 3)) + if config["use_meanflow"]: + modules.extend( + (f"time_r_in.mlp.{i}", f"time_guidance_embed.timestep_embedder_r.linear_{j}") for i, j in ((0, 1), (2, 2)) + ) + for i in range(config["num_layers"]): + for source, targets in ( + ("img", ("to_q", "to_k", "to_v")), + ("txt", ("add_q_proj", "add_k_proj", "add_v_proj")), + ): + modules.extend( + (f"double_blocks.{i}.{source}_attn_{part}", f"transformer_blocks.{i}.attn.{target}") + for part, target in zip(("q", "k", "v"), targets) + ) + for i in range(config["num_single_layers"]): + modules.extend( + (f"single_blocks.{i}.linear1_{part}", f"single_transformer_blocks.{i}.{target}") + for part, target in (("q", "attn.to_q"), ("k", "attn.to_k"), ("v", "attn.to_v"), ("mlp", "proj_mlp")) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + original_format = config.get( + "original_format", "hunyuan_image_split" if config["use_meanflow"] else "hunyuan_image_fused" + ) + if original_format == "hunyuan_image_fused": + width = config["num_attention_heads"] * config["attention_head_dim"] + for i in range(config["num_layers"]): + for modality in ("img", "txt"): + for p in ("weight", "bias"): + targets = tuple( + mapping.pop(f"double_blocks.{i}.{modality}_attn_{part}.{p}") for part in ("q", "k", "v") + ) + rules.append(Rule((f"double_blocks.{i}.{modality}_attn_qkv.{p}",), targets, Split((width,) * 3))) + for i in range(config["num_single_layers"]): + for p in ("weight", "bias"): + targets = tuple( + mapping.pop(f"single_blocks.{i}.linear1_{part}.{p}") for part in ("q", "k", "v", "mlp") + ) + rules.append( + Rule( + (f"single_blocks.{i}.linear1.{p}",), + targets, + Split((width, width, width, int(width * config["mlp_ratio"]))), + ) + ) + mapping[f"single_blocks.{i}.linear2.{p}"] = mapping.pop(f"single_blocks.{i}.linear2.fc.{p}") + elif original_format != "hunyuan_image_split": + raise ValueError("Hunyuan Image original_format must be hunyuan_image_fused or hunyuan_image_split.") + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_image_refiner_vae.py b/src/diffusers/loaders/conversion/hunyuan_image_refiner_vae.py new file mode 100644 index 000000000000..746159b7fd3c --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_image_refiner_vae.py @@ -0,0 +1,22 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .hunyuan_video15_vae import hunyuan_video15_vae_conversion + + +def hunyuan_image_refiner_vae_conversion(config): + base = hunyuan_video15_vae_conversion(config) + return Conversion(mapping=base.mapping) diff --git a/src/diffusers/loaders/conversion/hunyuan_image_vae.py b/src/diffusers/loaders/conversion/hunyuan_image_vae.py new file mode 100644 index 000000000000..d1a949f94c43 --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_image_vae.py @@ -0,0 +1,74 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +from .core import Conversion, Rule +from .transforms import Squeeze + + +def hunyuan_image_vae_conversion(config: dict) -> Conversion: + """Map the original image VAE, whose convolutions retain a singleton temporal axis.""" + modules, resnets = [], [] + for component in ("encoder", "decoder"): + encoder = component == "encoder" + direction = "down" if encoder else "up" + channels = list(config["block_out_channels"]) + if not encoder: + channels.reverse() + modules.extend((f"{component}.{name}", f"{component}.{name}") for name in ("conv_in", "conv_out", "norm_out")) + resnets.extend( + (f"{component}.mid.block_{i + 1}", f"{component}.mid_block.resnets.{i}", False) for i in range(2) + ) + old, new = f"{component}.mid.attn_1", f"{component}.mid_block.attentions.0" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in (("norm", "norm"), ("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("proj_out", "proj")) + ) + previous, index = channels[0], 0 + for i, channel in enumerate(channels): + for j in range(config["layers_per_block"] + int(not encoder)): + resnets.append( + ( + f"{component}.{direction}.{i}.block.{j}", + f"{component}.{direction}_blocks.{index}", + previous != channel, + ) + ) + index += 1 + previous = channel + if i < int(math.log2(config["spatial_compression_ratio"])) and i < len(channels) - 1: + modules.append( + ( + f"{component}.{direction}.{i}.{direction}sample.conv", + f"{component}.{direction}_blocks.{index}.conv", + ) + ) + index += 1 + if config["downsample_match_channel" if encoder else "upsample_match_channel"]: + previous = channels[i + 1] + for old, new, shortcut in resnets: + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("norm1", "norm2", "conv1", "conv2")) + if shortcut: + modules.append((old + ".nin_shortcut", new + ".conv_shortcut")) + mapping, rules = {}, [] + temporal = config.get("original_format", "hunyuan_image_vae") == "hunyuan_image_vae" + for old, new in modules: + mapping[old + ".bias"] = new + ".bias" + if temporal and "norm" not in old.rsplit(".", 1)[-1]: + rules.append(Rule((old + ".weight",), (new + ".weight",), Squeeze(dim=2, ndim=5))) + else: + mapping[old + ".weight"] = new + ".weight" + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_video.py b/src/diffusers/loaders/conversion/hunyuan_video.py new file mode 100644 index 000000000000..ba44859c5451 --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_video.py @@ -0,0 +1,115 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def hunyuan_video_conversion(config): + hidden = config["num_attention_heads"] * config["attention_head_dim"] + modules = [ + ("img_in.proj", "x_embedder.proj"), + ("txt_in.input_embedder", "context_embedder.proj_in"), + ("final_layer.linear", "proj_out"), + ] + mapping, rules = {}, [] + for old, new in ( + ("time_in.mlp.0", "timestep_embedder.linear_1"), + ("time_in.mlp.2", "timestep_embedder.linear_2"), + ("vector_in.in_layer", "text_embedder.linear_1"), + ("vector_in.out_layer", "text_embedder.linear_2"), + ): + modules.append((old, "time_text_embed." + new)) + if config["guidance_embeds"]: + modules.extend( + (f"guidance_in.mlp.{i}", f"time_text_embed.guidance_embedder.linear_{j}") for i, j in ((0, 1), (2, 2)) + ) + for old, new in ( + ("t_embedder.mlp.0", "timestep_embedder.linear_1"), + ("t_embedder.mlp.2", "timestep_embedder.linear_2"), + ("c_embedder.linear_1", "text_embedder.linear_1"), + ("c_embedder.linear_2", "text_embedder.linear_2"), + ): + modules.append(("txt_in." + old, "context_embedder.time_text_embed." + new)) + for i in range(config["num_refiner_layers"]): + old, new = f"txt_in.individual_token_refiner.blocks.{i}", f"context_embedder.token_refiner.refiner_blocks.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("norm1", "norm1"), + ("norm2", "norm2"), + ("self_attn_proj", "attn.to_out.0"), + ("mlp.fc1", "ff.net.0.proj"), + ("mlp.fc2", "ff.net.2"), + ("adaLN_modulation.1", "norm_out.linear"), + ) + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.self_attn_qkv.{p}",), + tuple(f"{new}.attn.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + for i in range(config["num_layers"]): + old, new = f"double_blocks.{i}", f"transformer_blocks.{i}" + for source, norm, ff, projections, output, norms in ( + ("img", "norm1", "ff", ("to_q", "to_k", "to_v"), "to_out.0", ("norm_q", "norm_k")), + ( + "txt", + "norm1_context", + "ff_context", + ("add_q_proj", "add_k_proj", "add_v_proj"), + "to_add_out", + ("norm_added_q", "norm_added_k"), + ), + ): + modules.extend( + [ + (f"{old}.{source}_mod.linear", f"{new}.{norm}.linear"), + (f"{old}.{source}_mlp.fc1", f"{new}.{ff}.net.0.proj"), + (f"{old}.{source}_mlp.fc2", f"{new}.{ff}.net.2"), + (f"{old}.{source}_attn_proj", f"{new}.attn.{output}"), + ] + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.{source}_attn_qkv.{p}",), + tuple(f"{new}.attn.{part}.{p}" for part in projections), + Split((hidden,) * 3), + ) + ) + if config["qk_norm"] is not None: + for part, target in zip(("q", "k"), norms): + mapping[f"{old}.{source}_attn_{part}_norm.weight"] = f"{new}.attn.{target}.weight" + for i in range(config["num_single_layers"]): + old, new = f"single_blocks.{i}", f"single_transformer_blocks.{i}" + modules.extend([(old + ".modulation.linear", new + ".norm.linear"), (old + ".linear2", new + ".proj_out")]) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.linear1.{p}",), + tuple(f"{new}.{part}.{p}" for part in ("attn.to_q", "attn.to_k", "attn.to_v", "proj_mlp")), + Split((hidden,) * 3 + (int(hidden * config["mlp_ratio"]),)), + ) + ) + if config["qk_norm"] is not None: + for part in ("q", "k"): + mapping[f"{old}.{part}_norm.weight"] = f"{new}.attn.norm_{part}.weight" + for p in ("weight", "bias"): + rules.append(Rule((f"final_layer.adaLN_modulation.1.{p}",), (f"norm_out.linear.{p}",), ReorderChunks((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/hunyuan_video15.py b/src/diffusers/loaders/conversion/hunyuan_video15.py new file mode 100644 index 000000000000..0b58ee0e2e23 --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_video15.py @@ -0,0 +1,51 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .hunyuan_video import hunyuan_video_conversion + + +def hunyuan_video15_conversion(config): + base = hunyuan_video_conversion({**config, "guidance_embeds": False, "num_single_layers": 0}) + mapping = { + old: new.replace("time_text_embed.timestep_embedder.", "time_embed.timestep_embedder.") + if old.startswith("time_in.") + else new + for old, new in base.mapping.items() + if not old.startswith("vector_in.") + } + rules = tuple(rule for rule in base.rules if not rule.original[0].startswith("double_blocks.")) + mapping["cond_type_embedding.weight"] = "cond_type_embed.weight" + modules = [("byt5_in.layernorm", "context_embedder_2.norm")] + modules.extend((f"byt5_in.fc{i}", f"context_embedder_2.linear_{i}") for i in (1, 2, 3)) + modules.extend( + (f"vision_in.proj.{i}", f"image_embedder.{name}") + for i, name in ((0, "norm_in"), (1, "linear_1"), (3, "linear_2"), (4, "norm_out")) + ) + if config["use_meanflow"]: + modules.extend( + (f"time_r_in.mlp.{i}", f"time_embed.timestep_embedder_r.linear_{j}") for i, j in ((0, 1), (2, 2)) + ) + for i in range(config["num_layers"]): + for source, targets in ( + ("img", ("to_q", "to_k", "to_v")), + ("txt", ("add_q_proj", "add_k_proj", "add_v_proj")), + ): + modules.extend( + (f"double_blocks.{i}.{source}_attn_{part}", f"transformer_blocks.{i}.attn.{target}") + for part, target in zip(("q", "k", "v"), targets) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/hunyuan_video15_vae.py b/src/diffusers/loaders/conversion/hunyuan_video15_vae.py new file mode 100644 index 000000000000..bfde64c8bf03 --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_video15_vae.py @@ -0,0 +1,59 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +from .core import Conversion + + +def hunyuan_video15_vae_conversion(config): + mapping, modules, resnets = {}, [], [] + spatial = int(math.log2(config["spatial_compression_ratio"])) + temporal = int(math.log2(config["temporal_compression_ratio"])) + for component in ("encoder", "decoder"): + encoder = component == "encoder" + direction = "down" if encoder else "up" + channels = list(config["block_out_channels"]) + if not encoder: + channels.reverse() + modules.extend((f"{component}.{name}.conv", f"{component}.{name}.conv") for name in ("conv_in", "conv_out")) + mapping[f"{component}.norm_out.gamma"] = f"{component}.norm_out.gamma" + resnets.extend( + (f"{component}.mid.block_{j + 1}", f"{component}.mid_block.resnets.{j}", False) for j in range(2) + ) + old, new = f"{component}.mid.attn_1", f"{component}.mid_block.attentions.0" + mapping[old + ".norm.gamma"] = new + ".norm.gamma" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("proj_out", "proj_out")) + ) + previous = channels[0] + for i, channel in enumerate(channels): + old, new = f"{component}.{direction}.{i}", f"{component}.{direction}_blocks.{i}" + for j in range(config["layers_per_block"] + int(not encoder)): + resnets.append((f"{old}.block.{j}", f"{new}.resnets.{j}", previous != channel)) + previous = channel + resampling = i < spatial if encoder else i < max(spatial, temporal) + if resampling: + modules.append((f"{old}.{direction}sample.conv.conv", f"{new}.{direction}samplers.0.conv.conv")) + if config["downsample_match_channel" if encoder else "upsample_match_channel"]: + previous = channels[i + 1] + for old, new, shortcut in resnets: + mapping.update({f"{old}.{norm}.gamma": f"{new}.{norm}.gamma" for norm in ("norm1", "norm2")}) + modules.extend((f"{old}.{conv}.conv", f"{new}.{conv}.conv") for conv in ("conv1", "conv2")) + if shortcut: + modules.append((old + ".nin_shortcut", new + ".conv_shortcut")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/hunyuan_video_vae.py b/src/diffusers/loaders/conversion/hunyuan_video_vae.py new file mode 100644 index 000000000000..9c43c419aa2c --- /dev/null +++ b/src/diffusers/loaders/conversion/hunyuan_video_vae.py @@ -0,0 +1,46 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def hunyuan_video_vae_conversion(config): + base = ldm_vae_conversion(config) + keys = [] + for key in sorted(base.diffusers_keys): + if ".downsamplers." in key or ".upsamplers." in key: + continue + parts = key.split(".") + if parts[-2] in ("conv_in", "conv_out", "conv1", "conv2", "conv_shortcut"): + parts.insert(-1, "conv") + keys.append(".".join(parts)) + count = len(config["block_out_channels"]) + spatial = int(math.log2(config["spatial_compression_ratio"])) + temporal = int(math.log2(config["temporal_compression_ratio"])) + for i in range(count): + resampling = i < spatial or (i >= count - 1 - temporal and i < count - 1) + if resampling: + keys.extend( + f"{component}.{direction}_blocks.{i}.{sampler}.0.conv.conv.{p}" + for component, direction, sampler in ( + ("encoder", "down", "downsamplers"), + ("decoder", "up", "upsamplers"), + ) + for p in ("weight", "bias") + ) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/i2vgen_xl.py b/src/diffusers/loaders/conversion/i2vgen_xl.py new file mode 100644 index 000000000000..c18a41349dcd --- /dev/null +++ b/src/diffusers/loaders/conversion/i2vgen_xl.py @@ -0,0 +1,49 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split +from .unet_3d import unet_3d_conversion + + +def i2vgen_xl_conversion(config): + base = unet_3d_conversion(config) + mapping, rules = dict(base.mapping), list(base.rules) + modules = [] + for old, new, indices in ( + ("local_image_concat", "image_latents_proj_in", (0, 2, 4)), + ("local_image_embedding", "image_latents_context_embedding", (0, 3, 5)), + ("context_embedding", "context_embedding", (0, 2)), + ("fps_embedding", "fps_embedding", (0, 2)), + ): + modules.extend((f"{old}.{i}", f"{new}.{i}") for i in indices) + old, new = "local_temporal_encoder.layers.0", "image_latents_temporal_encoder" + modules.extend( + [ + (old + ".0.norm", new + ".norm1"), + (old + ".0.fn.to_out.0", new + ".attn1.to_out.0"), + (old + ".1.net.0.0", new + ".ff.net.0.proj"), + (old + ".1.net.2", new + ".ff.net.2"), + ] + ) + rules.append( + Rule( + (old + ".0.fn.to_qkv.weight",), + tuple(f"{new}.attn1.to_{part}.weight" for part in ("q", "k", "v")), + Split((2 * config["in_channels"],) * 3), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/if_unet.py b/src/diffusers/loaders/conversion/if_unet.py new file mode 100644 index 000000000000..8d0dec20f825 --- /dev/null +++ b/src/diffusers/loaders/conversion/if_unet.py @@ -0,0 +1,77 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Reshape +from .unclip_unet import unclip_unet_conversion + + +def if_unet_conversion(config): + base = unclip_unet_conversion(config) + mapping, rules = dict(base.mapping), [] + modules = [] + if config.get("encoder_hid_dim") is not None: + modules.append(("encoder_proj", "encoder_hid_proj")) + if config.get("addition_embed_type") == "text": + modules.extend( + [ + ("encoder_pooling.0", "add_embedding.norm1"), + ("encoder_pooling.2", "add_embedding.proj"), + ("encoder_pooling.3", "add_embedding.norm2"), + ] + ) + modules.extend( + (f"encoder_pooling.1.{part}_proj", f"add_embedding.pool.{part}_proj") for part in ("q", "k", "v") + ) + mapping["encoder_pooling.1.positional_embedding"] = "add_embedding.pool.positional_embedding" + if config.get("class_embed_type") in ("timestep", "projection"): + modules.extend([("label_emb.0.0", "class_embedding.linear_1"), ("label_emb.0.2", "class_embedding.linear_2")]) + for rule in base.rules: + source = rule.original[0] + if ".qkv." not in source: + rules.append(rule) + continue + old = source.rsplit(".qkv.", 1)[0] + new = rule.diffusers[0].rsplit(".to_q.", 1)[0] + p = source.rsplit(".", 1)[1] + parts = new.split(".") + count = len(config["block_out_channels"]) + index = ( + count - 1 + if parts[0] == "mid_block" + else (count - 1 - int(parts[1]) if parts[0] == "up_blocks" else int(parts[1])) + ) + only_cross = config.get("only_cross_attention", False) + only_cross = only_cross[index] if isinstance(only_cross, (tuple, list)) else only_cross + if parts[0] == "mid_block": + only_cross = ( + config.get("mid_block_only_cross_attention") + if config.get("mid_block_only_cross_attention") is not None + else only_cross + ) + if only_cross: + channel = config["block_out_channels"][index] + dim = config["attention_head_dim"] + dim = dim[index] if isinstance(dim, (tuple, list)) else dim + inner = channel // dim * dim + if p == "weight": + rules.append(Rule((source,), (f"{new}.to_q.weight",), Reshape((inner, channel, 1), (inner, channel)))) + else: + mapping[source] = f"{new}.to_q.bias" + else: + rules.append(rule) + if p == "weight" and config.get("cross_attention_norm") is not None: + modules.append((f"{old}.norm_encoder", f"{new}.norm_cross")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/io.py b/src/diffusers/loaders/conversion/io.py new file mode 100644 index 000000000000..9e527f335fd5 --- /dev/null +++ b/src/diffusers/loaders/conversion/io.py @@ -0,0 +1,277 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Local tensor checkpoint I/O, separate from the reversible component definitions.""" + +import json +import tempfile +from collections.abc import Iterator, Mapping, Sequence +from pathlib import Path +from typing import Any, Literal + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from .registry import get_conversion + + +class Checkpoint(Mapping[str, Any]): + """Read a component checkpoint lazily, retaining at most one deserialized PyTorch shard. + + Safetensors are memory mapped per access. PyTorch files are loaded on CPU with `weights_only=True`. An explicit + wrapper selects a nested state dict; a prefix selects one component without discarding unexpected keys inside it. + Safe auxiliary metadata is retained until component preparation; tensor rules validate the values they consume. + Indexes must describe every tensor in every referenced shard exactly once, with paths inside the index directory. + """ + + def __init__(self, path: str | Path, *, prefix: str = "", wrapper: Sequence[str] = ()) -> None: + path = Path(path) + self.prefix, self.wrapper = prefix, tuple(wrapper) + self._aliases = {} + self._cached_path, self._cached_state = None, None + if path.is_dir(): + names = ( + "diffusion_pytorch_model.safetensors", + "diffusion_pytorch_model.safetensors.index.json", + "model.safetensors", + "model.safetensors.index.json", + "diffusion_pytorch_model.bin", + "diffusion_pytorch_model.bin.index.json", + "pytorch_model.bin", + "pytorch_model.bin.index.json", + ) + path = next((path / name for name in names if (path / name).is_file()), None) + if path is None: + raise FileNotFoundError("No tensor checkpoint or shard index found in the component directory.") + if path.name.endswith(".index.json"): + index = json.loads(path.read_text(encoding="utf-8")) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError("A checkpoint index needs a nonempty weight_map.") + root = path.parent.resolve() + paths = {} + for name in set(weight_map.values()): + shard = (root / name).resolve() + if Path(name).is_absolute() or not shard.is_relative_to(root): + raise ValueError(f"Checkpoint shard escapes the index directory: {name}.") + paths[name] = shard + actual = {} + for name, shard in paths.items(): + for key in self._keys(shard): + if key in actual: + raise ValueError(f"Duplicate key across checkpoint shards: {key}.") + actual[key] = name + if actual != weight_map: + raise ValueError("Checkpoint index does not match the keys and locations in its shards.") + self._paths = {key: paths[name] for key, name in weight_map.items()} + else: + self._paths = dict.fromkeys(self._keys(path), path) + self._paths = {key[len(prefix) :]: path for key, path in self._paths.items() if key.startswith(prefix)} + if not self._paths: + raise ValueError(f"No tensor keys match prefix {prefix!r}.") + + def _load(self, path: Path) -> Mapping[str, Any]: + if self._cached_path != path: + state = torch.load(path, map_location="cpu", weights_only=True) + for key in self.wrapper: + if not isinstance(state, Mapping) or key not in state: + raise ValueError(f"Missing checkpoint wrapper {self.wrapper} in {path}.") + state = state[key] + if not isinstance(state, Mapping) or any(not isinstance(key, str) for key in state): + raise ValueError("Expected a string-keyed state dict; select nested weights with wrapper.") + if not any(isinstance(value, torch.Tensor) for value in state.values()): + raise ValueError("Expected a tensor state dict; select nested weights with wrapper.") + self._cached_path, self._cached_state = path, state + return self._cached_state + + def _keys(self, path: Path) -> Sequence[str]: + if path.suffix == ".safetensors": + if self.wrapper: + raise ValueError("Safetensors do not have nested checkpoint wrappers.") + with safe_open(path, framework="pt", device="cpu") as handle: + return handle.keys() + return tuple(self._load(path)) + + def __getitem__(self, key: str) -> Any: + key = self._aliases.get(key, key) + path = self._paths[key] + if path.suffix == ".safetensors": + with safe_open(path, framework="pt", device="cpu") as handle: + return handle.get_tensor(self.prefix + key) + return self._load(path)[self.prefix + key] + + def __iter__(self) -> Iterator[str]: + return iter(self._paths) + + def __len__(self) -> int: + return len(self._paths) + + +def convert_checkpoint( + input_path: str | Path | Mapping[str, Any], + output_path: str | Path, + *, + config: dict[str, Any], + model_class: str | None = None, + reverse: bool = False, + input_prefix: str = "", + input_wrapper: Sequence[str] = (), + output_prefix: str = "", + output_format: Literal["safetensors", "pytorch"] = "safetensors", + output_wrapper: Sequence[str] = (), + max_shard_size: int = 5_000_000_000, +) -> Path: + """Write a new safetensors component directory or a single PyTorch checkpoint. + + `config` is the matching Diffusers component configuration in either direction. Original inputs must use the + selected definition's canonical component layout. The output contains safetensors plus configuration/format + metadata; original runtime configuration, pipeline assets and training state are not reconstructed. Existing output + directories are never overwritten. Conversion failures leave no partially published output directory. + """ + if max_shard_size <= 0: + raise ValueError("max_shard_size must be positive.") + if output_format not in ("safetensors", "pytorch"): + raise ValueError("output_format must be safetensors or pytorch.") + if output_wrapper and output_format != "pytorch": + raise ValueError("Nested output wrappers require PyTorch serialization.") + if config.get("quantization_config") is not None: + raise ValueError("Conversion requires unpacked, unquantized tensor weights.") + model_class = model_class or config.get("_class_name") + if model_class is None and len(config.get("architectures", [])) == 1: + model_class = config["architectures"][0] + conversion = get_conversion(model_class, config) + if isinstance(input_path, Mapping): + if input_prefix or input_wrapper: + raise ValueError("Select prefixes and wrappers when composing the input tensor mapping.") + checkpoint = input_path + else: + checkpoint = Checkpoint(input_path, prefix=input_prefix, wrapper=input_wrapper) + if not reverse and model_class == "CogVideoXTransformer3DModel": + from .checkpoint import ComponentState + from .cogvideox import cogvideox_transformer_auxiliary_keys + + checkpoint = ComponentState(checkpoint) + for key in cogvideox_transformer_auxiliary_keys(config): + checkpoint.keys_to_source.pop(key, None) + elif not reverse and model_class == "AutoencoderKLCogVideoX": + from .checkpoint import ComponentState + + checkpoint = ComponentState(checkpoint) + checkpoint.keys_to_source = { + key: value for key, value in checkpoint.keys_to_source.items() if not key.startswith("loss.") + } + # Transformers omits duplicate storage when saving tied embeddings as safetensors. Restore only declared ties, + # never infer an arbitrary missing parameter from another tensor with the same shape. + tied_groups = [] + if model_class in ("T5EncoderModel", "UMT5EncoderModel"): + tied_groups.append(("shared.weight", "encoder.embed_tokens.weight")) + if model_class == "Qwen3ForCausalLM" and config.get("tie_word_embeddings", False): + tied_groups.append(("model.embed_tokens.weight", "lm_head.weight")) + for group in tied_groups: + source = next((key for key in group if key in checkpoint), None) + if source is not None: + for key in group: + if key not in checkpoint: + if isinstance(checkpoint, Checkpoint): + checkpoint._aliases[key] = source + checkpoint._paths[key] = checkpoint._paths[source] + else: + from .checkpoint import ComponentState + + if not isinstance(checkpoint, ComponentState): + checkpoint = ComponentState(checkpoint) + checkpoint.keys_to_source[key] = source + if not reverse and model_class not in ("CogVideoXTransformer3DModel", "AutoencoderKLCogVideoX"): + from .checkpoint import prepare_component_checkpoint + + checkpoint, conversion, config = prepare_component_checkpoint(checkpoint, config, model_class) + output_path = Path(output_path) + if output_path.exists(): + raise FileExistsError(f"Output already exists: {output_path}.") + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix=".conversion-", dir=output_path.parent) as temporary: + if output_format == "pytorch": + converted = dict(conversion.iter_converted(checkpoint, reverse=reverse)) + if ( + reverse + and model_class == "CogVideoXTransformer3DModel" + and not config.get("use_rotary_positional_embeddings", False) + ): + from .cogvideox import cogvideox_fixed_position_embedding + + converted["mixins.pos_embed.pos_embedding"] = cogvideox_fixed_position_embedding( + config, + checkpoint["patch_embed.proj.weight"].dtype, + ) + converted = {output_prefix + key: tensor for key, tensor in converted.items()} + for wrapper in reversed(output_wrapper): + converted = {wrapper: converted} + staging = Path(temporary) / "checkpoint.pt" + torch.save(converted, staging) + staging.rename(output_path) + return output_path + staging = Path(temporary) / "component" + staging.mkdir() + shards, current, size, total_size = [], {}, 0, 0 + + def flush(): + if not current: + return + filename = f"part-{len(shards) + 1:05d}.safetensors" + save_file(current, staging / filename, metadata={"format": "pt"}) + shards.append((filename, tuple(current))) + current.clear() + + for key, tensor in conversion.iter_converted(checkpoint, reverse=reverse): + tensor_size = tensor.numel() * tensor.element_size() + if size and size + tensor_size > max_shard_size: + flush() + size = 0 + # Split views and tied parameters must have independent contiguous storage for safetensors. + current[output_prefix + key] = tensor.detach().cpu().contiguous().clone() + size += tensor_size + total_size += tensor_size + flush() + base = "model" if reverse or config.get("model_type") else "diffusion_pytorch_model" + weight_map = {} + for i, (filename, keys) in enumerate(shards, 1): + final_name = ( + f"{base}.safetensors" if len(shards) == 1 else f"{base}-{i:05d}-of-{len(shards):05d}.safetensors" + ) + (staging / filename).rename(staging / final_name) + weight_map.update(dict.fromkeys(keys, final_name)) + if len(shards) > 1: + (staging / f"{base}.safetensors.index.json").write_text( + json.dumps({"metadata": {"total_size": total_size}, "weight_map": weight_map}, indent=2) + "\n", + encoding="utf-8", + ) + saved_config = dict(config, _class_name=model_class) + config_name = "conversion_config.json" if reverse else "config.json" + (staging / config_name).write_text(json.dumps(saved_config, indent=2) + "\n", encoding="utf-8") + (staging / "conversion.json").write_text( + json.dumps( + { + "model_class": model_class, + "direction": "original" if reverse else "diffusers", + "lossless": conversion.lossless, + "output_prefix": output_prefix, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + staging.rename(output_path) + return output_path diff --git a/src/diffusers/loaders/conversion/joy_image.py b/src/diffusers/loaders/conversion/joy_image.py new file mode 100644 index 000000000000..49606d308bc9 --- /dev/null +++ b/src/diffusers/loaders/conversion/joy_image.py @@ -0,0 +1,42 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def joy_image_conversion(config): + names = [ + "img_in", + "proj_out", + "condition_embedder.time_embedder.linear_1", + "condition_embedder.time_embedder.linear_2", + "condition_embedder.time_proj", + "condition_embedder.text_embedder.linear_1", + "condition_embedder.text_embedder.linear_2", + ] + mapping = {} + for i in range(config["num_layers"]): + prefix = f"double_blocks.{i}" + for modality in ("img", "txt"): + mapping[f"{prefix}.{modality}_mod.modulate_table"] = f"{prefix}.{modality}_mod.modulate_table" + names.extend(f"{prefix}.{modality}_mlp.net.{leaf}" for leaf in ("0.proj", "2")) + for leaf in ("qkv", "proj"): + for p in ("weight", "bias"): + mapping[f"{prefix}.{modality}_attn_{leaf}.{p}"] = f"{prefix}.attn.{modality}_attn_{leaf}.{p}" + for part in ("q", "k"): + mapping[f"{prefix}.{modality}_attn_{part}_norm.weight"] = ( + f"{prefix}.attn.{modality}_attn_{part}_norm.weight" + ) + mapping.update({f"{name}.{p}": f"{name}.{p}" for name in names for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/k_upscaler.py b/src/diffusers/loaders/conversion/k_upscaler.py new file mode 100644 index 000000000000..2eec04a7b244 --- /dev/null +++ b/src/diffusers/loaders/conversion/k_upscaler.py @@ -0,0 +1,134 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Chain, Reshape, Split + + +def k_upscaler_conversion(config): + if config["time_embedding_type"] != "fourier" or config["mid_block_type"] is not None: + raise ValueError("Original k-upscaler uses Fourier time embeddings and no separate middle block.") + channels = config["block_out_channels"] + count = len(channels) + time_dim = config.get("time_embedding_dim") or channels[0] * 2 + modules = [ + ("inner_model.proj_in", "conv_in"), + ("inner_model.proj_out", "conv_out"), + ("inner_model.mapping.0", "time_embedding.linear_1"), + ("inner_model.mapping.2", "time_embedding.linear_2"), + ] + mapping = {} + if config.get("time_cond_proj_dim") is not None: + mapping["inner_model.mapping_cond.weight"] = "time_embedding.cond_proj.weight" + rules = [ + Rule( + ("inner_model.timestep_embed.weight",), + ("time_proj.weight",), + Reshape((time_dim // 2, 1), (time_dim // 2,)), + ) + ] + layers = config["layers_per_block"] + layers = [layers] * count if isinstance(layers, int) else layers + for direction in ("down", "up"): + for i, kind in enumerate(config[f"{direction}_block_types"]): + level = i if direction == "down" else count - 1 - i + channel = channels[level] + other = channels[max(level - 1, 0)] + first_up = direction == "up" and other == channel == time_dim + self_attn = i == count - 1 if direction == "down" else first_up + attention = "CrossAttn" in kind + stride = (3 if self_attn else 2) if attention else 1 + source = f"inner_model.u_net.{'d' if direction == 'down' else 'u'}_blocks.{i}" + for j in range(layers[level]): + index = stride * j + int(direction == "down") + old, new = f"{source}.{index}", f"{direction}_blocks.{i}.resnets.{j}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("main.0.mapper", "norm1.linear"), + ("main.2", "conv1"), + ("main.4.mapper", "norm2.linear"), + ("main.6", "conv2"), + ) + ) + if direction == "down": + input_width = other if j == 0 else channel + output_width = channel + else: + input_width = channel * (1 if first_up else 2) if j == 0 else channel + output_width = other if j == layers[level] - 1 else channel + if input_width != output_width: + mapping[f"{old}.skip.weight"] = f"{new}.conv_shortcut.weight" + if not attention: + continue + target = f"{direction}_blocks.{i}.attentions.{j}" + width = channel if direction == "down" else output_width + head_dim = config["attention_head_dim"] + head_dim = head_dim[level] if isinstance(head_dim, (tuple, list)) else head_dim + inner = width // head_dim * head_dim + context = config["cross_attention_dim"] + context = context[level] if isinstance(context, (tuple, list)) else context + if self_attn: + old = f"{source}.{index + 1}" + modules.append((f"{old}.norm_in.mapper", f"{target}.norm1.linear")) + rules.append( + Rule( + (f"{old}.qkv_proj.weight",), + tuple(f"{target}.attn1.to_{part}.weight" for part in ("q", "k", "v")), + Chain((Reshape((3 * inner, width, 1, 1), (3 * inner, width)), Split((inner,) * 3))), + ) + ) + rules.append( + Rule( + (f"{old}.qkv_proj.bias",), + tuple(f"{target}.attn1.to_{part}.bias" for part in ("q", "k", "v")), + Split((inner,) * 3), + ) + ) + rules.append( + Rule( + (f"{old}.out_proj.weight",), + (f"{target}.attn1.to_out.0.weight",), + Reshape((width, inner, 1, 1), (width, inner)), + ) + ) + mapping[f"{old}.out_proj.bias"] = f"{target}.attn1.to_out.0.bias" + old = f"{source}.{index + 1 + int(self_attn)}" + modules.extend( + [ + (f"{old}.norm_dec.mapper", f"{target}.norm2.linear"), + (f"{old}.norm_enc", f"{target}.attn2.norm_cross"), + ] + ) + rules.append( + Rule( + (f"{old}.kv_proj.weight",), + (f"{target}.attn2.to_k.weight", f"{target}.attn2.to_v.weight"), + Split((inner,) * 2), + ) + ) + rules.append( + Rule( + (f"{old}.kv_proj.bias",), + (f"{target}.attn2.to_k.bias", f"{target}.attn2.to_v.bias"), + Split((inner,) * 2), + ) + ) + for a, b, shape in (("q_proj", "to_q", (inner, width)), ("out_proj", "to_out.0", (width, inner))): + rules.append( + Rule((f"{old}.{a}.weight",), (f"{target}.attn2.{b}.weight",), Reshape(shape + (1, 1), shape)) + ) + mapping[f"{old}.{a}.bias"] = f"{target}.attn2.{b}.bias" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/kandinsky3.py b/src/diffusers/loaders/conversion/kandinsky3.py new file mode 100644 index 000000000000..9e82f03ecdcf --- /dev/null +++ b/src/diffusers/loaders/conversion/kandinsky3.py @@ -0,0 +1,115 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def kandinsky3_conversion(config): + modules = [ + ("to_time_embed.1", "time_embedding.linear_1"), + ("to_time_embed.3", "time_embedding.linear_2"), + ("in_layer", "conv_in"), + ("out_layer.0", "conv_norm_out"), + ("out_layer.2", "conv_out"), + ("projection_ln", "encoder_hid_proj.projection_norm"), + ] + weights = [("projection_lin", "encoder_hid_proj.projection_linear")] + attentions = [("feature_pooling", "add_time_condition", False)] + resnets = [] + channels = [config["block_out_channels"][0] // 2] + list(config["block_out_channels"]) + count, layers = len(channels) - 1, config["layers_per_block"] + flags = (False, True, True, True) + for i in range(count): + old, new = f"down_samples.{i}", f"down_blocks.{i}" + if flags[i]: + attentions.append((old + ".self_attention_block", new + ".attentions.0", True)) + for j in range(layers): + input_channel = channels[i] if j == 0 else channels[i + 1] + resnets.append( + ( + f"{old}.resnet_attn_blocks.{j}.0", + f"{new}.resnets_in.{j}", + input_channel != channels[i + 1], + False, + False, + ) + ) + resnets.append( + ( + f"{old}.resnet_attn_blocks.{j}.2", + f"{new}.resnets_out.{j}", + False, + False, + j == layers - 1 and i < count - 1, + ) + ) + if flags[i]: + attentions.append((f"{old}.resnet_attn_blocks.{j}.1", f"{new}.attentions.{j + 1}", True)) + for i in range(count): + level = count - 1 - i + input_channel, output_channel = channels[level + 1], channels[level] + cat = 0 if i == 0 else input_channel + old, new = f"up_samples.{i}", f"up_blocks.{i}" + if flags[-i - 1]: + attentions.append((old + ".self_attention_block", new + ".attentions.0", True)) + pairs = ( + [(input_channel + cat, input_channel)] + + [(input_channel, input_channel)] * (layers - 2) + + [(input_channel, output_channel)] + ) + for j, (a, b) in zip(range(layers), pairs): + resnets.append( + (f"{old}.resnet_attn_blocks.{j}.0", f"{new}.resnets_in.{j}", False, j == 0 and i > 0, False) + ) + resnets.append((f"{old}.resnet_attn_blocks.{j}.2", f"{new}.resnets_out.{j}", a != b, False, False)) + if flags[-i - 1]: + attentions.append((f"{old}.resnet_attn_blocks.{j}.1", f"{new}.attentions.{j + 1}", True)) + for old, new, shortcut, up, down in resnets: + for i in range(4): + modules.extend( + [ + ( + f"{old}.resnet_blocks.{i}.group_norm.context_mlp.1", + f"{new}.resnet_blocks.{i}.group_norm.context_mlp.1", + ), + (f"{old}.resnet_blocks.{i}.projection", f"{new}.resnet_blocks.{i}.projection"), + ] + ) + if shortcut: + modules.append((old + ".shortcut_projection", new + ".shortcut_projection")) + for active, index, direction in ((up, 1, "up"), (down, 2, "down")): + if active: + modules.extend( + [ + ( + f"{old}.resnet_blocks.{index}.{direction}_sample", + f"{new}.resnet_blocks.{index}.{direction}_sample", + ), + (f"{old}.shortcut_{direction}_sample", f"{new}.shortcut_{direction}_sample"), + ] + ) + for old, new, conditioned in attentions: + weights.extend( + (f"{old}.attention.{a}", f"{new}.attention.{b}") + for a, b in (("to_query", "to_q"), ("to_key", "to_k"), ("to_value", "to_v"), ("output_layer", "to_out.0")) + ) + if conditioned: + modules.extend( + (f"{old}.{name}.context_mlp.1", f"{new}.{name}.context_mlp.1") for name in ("in_norm", "out_norm") + ) + weights.extend((f"{old}.feed_forward.{i}", f"{new}.feed_forward.{i}") for i in (0, 2)) + mapping = {old + ".weight": new + ".weight" for old, new in weights} + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ldm_unet.py b/src/diffusers/loaders/conversion/ldm_unet.py new file mode 100644 index 000000000000..be7eda67d3e0 --- /dev/null +++ b/src/diffusers/loaders/conversion/ldm_unet.py @@ -0,0 +1,149 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ldm_unet_conversion(config): + if any(kind.startswith("K") for kind in config["down_block_types"]): + from .k_upscaler import k_upscaler_conversion + + return k_upscaler_conversion(config) + if config.get("addition_embed_type") == "text" and any( + "SimpleCrossAttn" in kind for kind in config["down_block_types"] + ): + from .if_unet import if_unet_conversion + + return if_unet_conversion(config) + if config.get("attention_type") in ("gated", "gated-text-image"): + from .gligen import gligen_conversion + + return gligen_conversion(config) + if config.get("original_format") == "versatile_image": + from .versatile_image_unet import versatile_image_unet_conversion + + return versatile_image_unet_conversion(config) + if any("SimpleCrossAttn" in kind for kind in config["down_block_types"]): + from .unclip_unet import unclip_unet_conversion + + return unclip_unet_conversion(config) + return Conversion(mapping=_ldm_unet_mapping(config, controlnet=False)) + + +def _ldm_unet_mapping(config, controlnet): + modules = [ + ("time_embed.0", "time_embedding.linear_1"), + ("time_embed.2", "time_embedding.linear_2"), + ("input_blocks.0.0", "conv_in"), + ] + mapping, resnets, attentions = {}, [], [] + channels = config["block_out_channels"] + count = len(channels) + layers = config["layers_per_block"] + layers = [layers] * count if isinstance(layers, int) else layers + depths = config["transformer_layers_per_block"] + depths = [depths] * count if isinstance(depths, int) else depths + if config.get("class_embed_type") in ("timestep", "projection"): + modules.extend([("label_emb.0.0", "class_embedding.linear_1"), ("label_emb.0.2", "class_embedding.linear_2")]) + elif config.get("class_embed_type") == "simple_projection": + modules.append(("film_emb", "class_embedding")) + elif config.get("num_class_embeds") is not None: + mapping["label_emb.weight"] = "class_embedding.weight" + if config.get("addition_embed_type") == "text_time": + modules.extend([("label_emb.0.0", "add_embedding.linear_1"), ("label_emb.0.2", "add_embedding.linear_2")]) + idx, previous = 1, channels[0] + for i, channel in enumerate(channels): + for j in range(layers[i]): + resnets.append((f"input_blocks.{idx}.0", f"down_blocks.{i}.resnets.{j}", previous != channel)) + previous = channel + if config["down_block_types"][i] == "CrossAttnDownBlock2D": + depth = depths[i][j] if isinstance(depths[i], (list, tuple)) else depths[i] + attentions.append((f"input_blocks.{idx}.1", f"down_blocks.{i}.attentions.{j}", depth)) + idx += 1 + if i < count - 1: + modules.append((f"input_blocks.{idx}.0.op", f"down_blocks.{i}.downsamplers.0.conv")) + idx += 1 + if config.get("mid_block_type", "UNetMidBlock2DCrossAttn") is not None: + resnets.extend((f"middle_block.{i * 2}", f"mid_block.resnets.{i}", False) for i in range(2)) + depth = depths[-1][-1] if isinstance(depths[-1], (list, tuple)) else depths[-1] + attentions.append(("middle_block.1", "mid_block.attentions.0", depth)) + if controlnet: + modules.extend((f"zero_convs.{i}.0", f"controlnet_down_blocks.{i}") for i in range(idx)) + modules.append(("middle_block_out.0", "controlnet_mid_block")) + cond_layers = 2 * (len(config["conditioning_embedding_out_channels"]) - 1) + modules.append(("input_hint_block.0", "controlnet_cond_embedding.conv_in")) + modules.extend( + (f"input_hint_block.{2 * (i + 1)}", f"controlnet_cond_embedding.blocks.{i}") for i in range(cond_layers) + ) + modules.append((f"input_hint_block.{2 * (cond_layers + 1)}", "controlnet_cond_embedding.conv_out")) + else: + modules.extend([("out.0", "conv_norm_out"), ("out.2", "conv_out")]) + up_channels = list(reversed(channels)) + up_layers = list(reversed(layers)) + up_depths = config.get("reverse_transformer_layers_per_block") or list(reversed(depths)) + previous, idx = up_channels[0], 0 + for i, channel in enumerate(up_channels): + input_channel = up_channels[min(i + 1, count - 1)] + attention = config["up_block_types"][i] == "CrossAttnUpBlock2D" + for j in range(up_layers[i] + 1): + skip_channel = input_channel if j == up_layers[i] else channel + resnets.append( + (f"output_blocks.{idx}.0", f"up_blocks.{i}.resnets.{j}", previous + skip_channel != channel) + ) + previous = channel + if attention: + depth = up_depths[i][j] if isinstance(up_depths[i], (list, tuple)) else up_depths[i] + attentions.append((f"output_blocks.{idx}.1", f"up_blocks.{i}.attentions.{j}", depth)) + if j == up_layers[i] and i < count - 1: + modules.append( + (f"output_blocks.{idx}.{2 if attention else 1}.conv", f"up_blocks.{i}.upsamplers.0.conv") + ) + idx += 1 + for old, new, shortcut in resnets: + pairs = [ + ("in_layers.0", "norm1"), + ("in_layers.2", "conv1"), + ("out_layers.0", "norm2"), + ("out_layers.3", "conv2"), + ("emb_layers.1", "time_emb_proj"), + ] + if shortcut: + pairs.append(("skip_connection", "conv_shortcut")) + modules.extend((f"{old}.{a}", f"{new}.{b}") for a, b in pairs) + for old, new, depth in attentions: + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("norm", "proj_in", "proj_out")) + for i in range(depth): + a, b = f"{old}.transformer_blocks.{i}", f"{new}.transformer_blocks.{i}" + modules.extend( + (f"{a}.{name}", f"{b}.{name}") + for name in ( + "norm1", + "norm2", + "norm3", + "ff.net.0.proj", + "ff.net.2", + "attn1.to_out.0", + "attn2.to_out.0", + ) + ) + mapping.update( + { + f"{a}.{attn}.to_{part}.weight": f"{b}.{attn}.to_{part}.weight" + for attn in ("attn1", "attn2") + for part in ("q", "k", "v") + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return mapping diff --git a/src/diffusers/loaders/conversion/ldm_vae.py b/src/diffusers/loaders/conversion/ldm_vae.py new file mode 100644 index 000000000000..eef99d5ae11a --- /dev/null +++ b/src/diffusers/loaders/conversion/ldm_vae.py @@ -0,0 +1,89 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Reshape + + +def ldm_vae_conversion(config): + modules, resnets, rules, attentions = [], [], [], [] + channels = config["block_out_channels"] + layers = config["layers_per_block"] + for name in ("quant_conv", "post_quant_conv"): + if config.get("use_" + name, True): + modules.append((name, name)) + for component in ("encoder", "decoder"): + modules.extend( + (f"{component}.{old}", f"{component}.{new}") + for old, new in (("conv_in", "conv_in"), ("conv_out", "conv_out"), ("norm_out", "conv_norm_out")) + ) + resnets.extend( + (f"{component}.mid.block_{i + 1}", f"{component}.mid_block.resnets.{i}", False) for i in range(2) + ) + if config["mid_block_add_attention"]: + attentions.append((f"{component}.mid.attn_1", f"{component}.mid_block.attentions.0", channels[-1])) + previous = channels[0] + for i, channel in enumerate(channels): + for j in range(layers): + resnets.append( + (f"encoder.down.{i}.block.{j}", f"encoder.down_blocks.{i}.resnets.{j}", previous != channel) + ) + previous = channel + if i < len(channels) - 1: + modules.append((f"encoder.down.{i}.downsample.conv", f"encoder.down_blocks.{i}.downsamplers.0.conv")) + previous = channels[-1] + for i, channel in enumerate(reversed(channels)): + original_index = len(channels) - 1 - i + for j in range(layers + 1): + resnets.append( + (f"decoder.up.{original_index}.block.{j}", f"decoder.up_blocks.{i}.resnets.{j}", previous != channel) + ) + previous = channel + if i < len(channels) - 1: + modules.append((f"decoder.up.{original_index}.upsample.conv", f"decoder.up_blocks.{i}.upsamplers.0.conv")) + for old, new, shortcut in resnets: + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("norm1", "norm2", "conv1", "conv2")) + if shortcut: + modules.append((old + ".nin_shortcut", new + ".conv_shortcut")) + for component, direction in (("encoder", "down"), ("decoder", "up")): + block_types = config.get(f"{direction}_block_types", ()) + for i, kind in enumerate(block_types): + if not kind.startswith("Attn"): + continue + original_index = i if component == "encoder" else len(channels) - 1 - i + channel = channels[original_index] + for j in range(layers + int(component == "decoder")): + attentions.append( + ( + f"{component}.{direction}.{original_index}.attn.{j}", + f"{component}.{direction}_blocks.{i}.attentions.{j}", + channel, + ) + ) + for old, new, width in attentions: + modules.append((old + ".norm", new + ".group_norm")) + for a, b in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("proj_out", "to_out.0")): + rules.append( + Rule((f"{old}.{a}.weight",), (f"{new}.{b}.weight",), Reshape((width, width, 1, 1), (width, width))) + ) + modules.append((f"{old}.{a}", f"{new}.{b}")) + transformed = {key for rule in rules for key in rule.original} + mapping = { + f"{old}.{p}": f"{new}.{p}" + for old, new in modules + for p in ("weight", "bias") + if f"{old}.{p}" not in transformed + } + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/longcat_audio.py b/src/diffusers/loaders/conversion/longcat_audio.py new file mode 100644 index 000000000000..d7c13492419f --- /dev/null +++ b/src/diffusers/loaders/conversion/longcat_audio.py @@ -0,0 +1,54 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def longcat_audio_conversion(config): + modules = [("proj_out", True), ("norm_out.linear", config["bias"])] + modules.extend((f"time_embed.time_mlp.{i}", True) for i in (0, 2)) + embeddings = ["input_embed", "text_embed"] + if config["use_latent_condition"]: + embeddings.extend(["latent_embed", "latent_cond_embedder"]) + modules.extend((f"{name}.proj.{i}", True) for name in embeddings for i in (0, 2)) + if config["adaln_type"] == "global": + modules.append(("adaln_global_mlp.mlp.1", True)) + keys = [] + if config["text_conv"]: + for i in range(4): + prefix = f"text_conv_layer.{i}" + modules.append((prefix + ".norm", True)) + modules.extend((f"{prefix}.{name}", config["bias"]) for name in ("dwconv", "pwconv1", "pwconv2")) + keys.extend(f"{prefix}.grn.{name}" for name in ("gamma", "beta")) + for i in range(config["dit_depth"]): + prefix = f"blocks.{i}" + if config["adaln_type"] == "global": + keys.append(prefix + ".adaln_scale_shift") + elif config["adaln_type"] == "local": + modules.append((prefix + ".adaln_mlp.mlp.1", True)) + attentions = ["self_attn"] + if config["cross_attn"]: + attentions.append("cross_attn") + if config["cross_attn_norm"]: + modules.extend((f"{prefix}.{name}", True) for name in ("cross_attn_norm", "cross_attn_norm_c")) + for attn in attentions: + modules.extend( + (f"{prefix}.{attn}.{name}", config["bias"]) for name in ("to_q", "to_k", "to_v", "to_out.0") + ) + if config["qk_norm"]: + keys.extend(f"{prefix}.{attn}.{part}_norm.weight" for part in ("q", "k")) + modules.extend((f"{prefix}.ffn.ff.{j}", config["bias"]) for j in (0, 3)) + keys.extend(f"{name}.{p}" for name, bias in modules for p in (("weight", "bias") if bias else ("weight",))) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/longcat_audio_vae.py b/src/diffusers/loaders/conversion/longcat_audio_vae.py new file mode 100644 index 000000000000..89492ea905f8 --- /dev/null +++ b/src/diffusers/loaders/conversion/longcat_audio_vae.py @@ -0,0 +1,46 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def longcat_audio_vae_conversion(config): + count = len(config["c_mults"] or (1, 2, 4, 8, 16)) + snake = config["act_fn"] == "snake" if config["act_fn"] is not None else config["use_snake"] is not False + convolutions = [ + ("encoder.layers.0", True), + (f"encoder.layers.{count + 1}", True), + ("decoder.layers.0", True), + (f"decoder.layers.{count + 2}", False), + ] + snakes = [f"decoder.layers.{count + 1}"] + for component in ("encoder", "decoder"): + encoder = component == "encoder" + for i in range(count): + prefix = f"{component}.layers.{i + 1}.layers" + snakes.append(prefix + (".3" if encoder else ".0")) + convolutions.append((prefix + (".4" if encoder else ".1"), True)) + for j in range(3): + unit = f"{prefix}.{j if encoder else j + 2}.layers" + snakes.extend([unit + ".0", unit + ".2"]) + convolutions.extend([(unit + ".1", True), (unit + ".3", True)]) + keys = [ + f"{name}.{p}" + for name, bias in convolutions + for p in (("weight_g", "weight_v", "bias") if bias else ("weight_g", "weight_v")) + ] + if snake: + keys.extend(f"{name}.{p}" for name in snakes for p in ("alpha", "beta")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/lora.py b/src/diffusers/loaders/conversion/lora.py new file mode 100644 index 000000000000..2c01daaaee02 --- /dev/null +++ b/src/diffusers/loaders/conversion/lora.py @@ -0,0 +1,99 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def lora_conversion(config): + """Convert adapter layouts using explicit qualified target modules from the adapter configuration. + + The Diffusers side uses PEFT A/B names. Optional alpha tensors remain separate `.alpha` entries; they must + be passed as network alphas when loading, never multiplied into the weights. This does not merge an adapter into a + base model. Explicit module names make underscore-delimited Kohya keys unambiguous. + """ + original_format = config.get("original_format", "kohya") + if original_format not in ("kohya", "diffusers", "diffusers_old", "peft", "animatediff"): + raise ValueError("Unknown LoRA original_format.") + modules = config["modules"] + if isinstance(modules, str) or not modules or len(set(modules)) != len(modules): + raise ValueError("LoRA modules must be a nonempty sequence of distinct, fully qualified module names.") + name = config.get("adapter_name") + adapter = f".{name}" if name else "" + mapping = {} + original_modules = set() + for module in modules: + component, _, path = module.partition(".") + if not path: + raise ValueError("LoRA module names must include their component, for example unet.conv_in.") + if original_format == "kohya": + prefixes = { + "unet": "lora_unet", + "text_encoder": "lora_te1", + "text_encoder_2": "lora_te2", + "transformer": "lora_transformer", + } + prefixes.update(config.get("component_prefixes", {})) + if component not in prefixes: + raise ValueError(f"No Kohya component prefix configured for {component}.") + original_module = prefixes[component] + "_" + path.replace(".", "_") + suffixes = ("lora_down.weight", "lora_up.weight") + elif original_format == "peft": + original_module = module + suffixes = (f"lora_A{adapter}.weight", f"lora_B{adapter}.weight") + elif original_format == "diffusers": + original_module = module + suffix = "lora" if component in ("unet", "transformer") else "lora_linear_layer" + suffixes = (f"{suffix}.down.weight", f"{suffix}.up.weight") + else: + original_module = path if original_format == "animatediff" else module + if original_format == "animatediff": + parts = original_module.split(".") + index = parts.index("motion_modules") + 2 + parts.insert(index, "temporal_transformer") + original_module = ( + ".".join(parts) + .replace(".norm1", ".norms.0") + .replace(".norm2", ".norms.1") + .replace(".norm3", ".ff_norm") + .replace(".attn1", ".attention_blocks.0") + .replace(".attn2", ".attention_blocks.1") + ) + replacements = { + "to_q": "to_q_lora", + "to_k": "to_k_lora", + "to_v": "to_v_lora", + "to_out.0": "to_out_lora", + "q_proj": "to_q_lora", + "k_proj": "to_k_lora", + "v_proj": "to_v_lora", + "out_proj": "to_out_lora", + } + matched = next((leaf for leaf in replacements if original_module.endswith("." + leaf)), None) + if matched is None: + raise ValueError(f"Legacy attention LoRA has no layout for module {module}.") + original_module = original_module[: -len(matched)] + replacements[matched] + suffixes = ("down.weight", "up.weight") + if original_module in original_modules: + raise ValueError(f"Original LoRA module name collision: {original_module}.") + original_modules.add(original_module) + for suffix, part in zip(suffixes, ("A", "B")): + mapping[f"{original_module}.{suffix}"] = f"{module}.lora_{part}{adapter}.weight" + if config.get("include_alpha", False): + mapping[f"{original_module}.alpha"] = f"{module}.alpha" + if config.get("use_dora", False): + if original_format not in ("kohya", "peft"): + raise ValueError("DoRA conversion requires kohya or peft format.") + suffix = "dora_scale" if original_format == "kohya" else f"lora_magnitude_vector{adapter}.weight" + mapping[f"{original_module}.{suffix}"] = f"{module}.lora_magnitude_vector{adapter}.weight" + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx.py b/src/diffusers/loaders/conversion/ltx.py new file mode 100644 index 000000000000..bc41417620b9 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx.py @@ -0,0 +1,44 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ltx_conversion(config): + mapping = {"scale_shift_table": "scale_shift_table"} + modules = [("patchify_proj", "proj_in"), ("proj_out", "proj_out")] + modules.extend( + ("adaln_single." + name, "time_embed." + name) + for name in ("emb.timestep_embedder.linear_1", "emb.timestep_embedder.linear_2", "linear") + ) + modules.extend(("caption_projection." + name, "caption_projection." + name) for name in ("linear_1", "linear_2")) + for i in range(config["num_layers"]): + prefix = f"transformer_blocks.{i}" + mapping[prefix + ".scale_shift_table"] = prefix + ".scale_shift_table" + modules.extend((prefix + ".ff." + name, prefix + ".ff." + name) for name in ("net.0.proj", "net.2")) + for attn in ("attn1", "attn2"): + for name in ("to_q", "to_k", "to_v", "to_out.0"): + bias = config["attention_out_bias"] if name == "to_out.0" else config["attention_bias"] + for p in ("weight", "bias") if bias else ("weight",): + key = f"{prefix}.{attn}.{name}.{p}" + mapping[key] = key + if config["qk_norm"] is not None: + for part in ("q", "k"): + mapping[f"{prefix}.{attn}.{part}_norm.weight"] = f"{prefix}.{attn}.norm_{part}.weight" + if config["norm_elementwise_affine"]: + for name in ("norm1", "norm2"): + mapping[f"{prefix}.{name}.weight"] = f"{prefix}.{name}.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx2.py b/src/diffusers/loaders/conversion/ltx2.py new file mode 100644 index 000000000000..e2bd64c77178 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2.py @@ -0,0 +1,93 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ltx2_conversion(config): + mapping = {} + modules = [] + for modality in ("", "audio_"): + mapping[modality + "scale_shift_table"] = modality + "scale_shift_table" + modules.extend( + [(modality + "patchify_proj", modality + "proj_in"), (modality + "proj_out", modality + "proj_out")] + ) + if config["use_prompt_embeddings"]: + modules.extend( + (modality + "caption_projection." + name, modality + "caption_projection." + name) + for name in ("linear_1", "linear_2") + ) + embeddings = [ + ("adaln_single", "time_embed"), + ("audio_adaln_single", "audio_time_embed"), + ("av_ca_video_scale_shift_adaln_single", "av_cross_attn_video_scale_shift"), + ("av_ca_audio_scale_shift_adaln_single", "av_cross_attn_audio_scale_shift"), + ("av_ca_a2v_gate_adaln_single", "av_cross_attn_video_a2v_gate"), + ("av_ca_v2a_gate_adaln_single", "av_cross_attn_audio_v2a_gate"), + ] + modulated = config["cross_attn_mod"] or config["audio_cross_attn_mod"] + if modulated and config["use_prompt_adaln_single"]: + embeddings.extend([("prompt_adaln", "prompt_adaln"), ("audio_prompt_adaln", "audio_prompt_adaln")]) + for old, new in embeddings: + modules.extend( + (old + "." + name, new + "." + name) + for name in ("emb.timestep_embedder.linear_1", "emb.timestep_embedder.linear_2", "linear") + ) + if config["use_keyframes_abs_pos_embedding"]: + mapping["keyframes_abs_pos_embedding"] = "keyframes_abs_pos_embedding" + for i in range(config["num_layers"]): + prefix = f"transformer_blocks.{i}" + for old, new in ( + ("scale_shift_table", "scale_shift_table"), + ("audio_scale_shift_table", "audio_scale_shift_table"), + ("scale_shift_table_a2v_ca_video", "video_a2v_cross_attn_scale_shift_table"), + ("scale_shift_table_a2v_ca_audio", "audio_a2v_cross_attn_scale_shift_table"), + ): + mapping[f"{prefix}.{old}"] = f"{prefix}.{new}" + if modulated: + for name in ("prompt_scale_shift_table", "audio_prompt_scale_shift_table"): + mapping[f"{prefix}.{name}"] = f"{prefix}.{name}" + for ff, bias in (("ff", config["ff_bias"]), ("audio_ff", config["audio_ff_bias"])): + for layer in ("net.0.proj", "net.2"): + for p in ("weight", "bias") if bias else ("weight",): + key = f"{prefix}.{ff}.{layer}.{p}" + mapping[key] = key + for attn in ("attn1", "attn2", "audio_attn1", "audio_attn2", "audio_to_video_attn", "video_to_audio_attn"): + for name in ("to_q", "to_k", "to_v", "to_out.0"): + bias = config["attention_out_bias"] if name == "to_out.0" else config["attention_bias"] + for p in ("weight", "bias") if bias else ("weight",): + key = f"{prefix}.{attn}.{name}.{p}" + mapping[key] = key + for part in ("q", "k"): + mapping[f"{prefix}.{attn}.{part}_norm.weight"] = f"{prefix}.{attn}.norm_{part}.weight" + audio = attn in ("audio_attn1", "audio_attn2", "video_to_audio_attn") + if config["audio_gated_attn" if audio else "gated_attn"]: + name = f"{prefix}.{attn}.to_gate_logits" + modules.append((name, name)) + if config["norm_elementwise_affine"]: + for norm in ( + "norm1", + "norm2", + "norm3", + "audio_norm1", + "audio_norm2", + "audio_norm3", + "audio_to_video_norm", + "video_to_audio_norm", + ): + key = f"{prefix}.{norm}.weight" + mapping[key] = key + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx2_audio_vae.py b/src/diffusers/loaders/conversion/ltx2_audio_vae.py new file mode 100644 index 000000000000..6d568bf3fc40 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_audio_vae.py @@ -0,0 +1,63 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ltx2_audio_vae_conversion(config): + mapping = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", + } + modules, resnets, attentions = [], [], [] + suffix = ".conv" if config["causality_axis"] is not None else "" + channels = [config["base_channels"] * m for m in config["ch_mult"]] + for component in ("encoder", "decoder"): + modules.extend(f"{component}.{name}{suffix}" for name in ("conv_in", "conv_out")) + if config["norm_type"] == "group": + modules.append(component + ".norm_out") + resnets.extend((f"{component}.mid.block_{i}", False) for i in (1, 2)) + if config["mid_block_add_attention"]: + attentions.append(component + ".mid.attn_1") + previous = channels[0] + for i, channel in enumerate(channels): + for j in range(config["num_res_blocks"]): + resnets.append((f"encoder.down.{i}.block.{j}", previous != channel)) + previous = channel + if config["resolution"] // 2**i in (config["attn_resolutions"] or ()): + attentions.append(f"encoder.down.{i}.attn.{j}") + if i < len(channels) - 1: + modules.append(f"encoder.down.{i}.downsample.conv") + previous = channels[-1] + for i in reversed(range(len(channels))): + for j in range(config["num_res_blocks"] + 1): + resnets.append((f"decoder.up.{i}.block.{j}", previous != channels[i])) + previous = channels[i] + if config["resolution"] // 2**i in (config["attn_resolutions"] or ()): + attentions.append(f"decoder.up.{i}.attn.{j}") + if i > 0: + modules.append(f"decoder.up.{i}.upsample.conv{suffix}") + for prefix, shortcut in resnets: + modules.extend(prefix + "." + name + suffix for name in ("conv1", "conv2")) + if shortcut: + modules.append(prefix + ".nin_shortcut" + suffix) + if config["norm_type"] == "group": + modules.extend(prefix + "." + name for name in ("norm1", "norm2")) + for prefix in attentions: + modules.extend(prefix + "." + name for name in ("q", "k", "v", "proj_out")) + if config["norm_type"] == "group": + modules.append(prefix + ".norm") + mapping.update({f"{name}.{p}": f"{name}.{p}" for name in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx2_connectors.py b/src/diffusers/loaders/conversion/ltx2_connectors.py new file mode 100644 index 000000000000..cbe8c1ecec53 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_connectors.py @@ -0,0 +1,42 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def ltx2_connectors_conversion(config): + mapping = {} + modalities = ("video", "audio") if config["per_modality_projections"] else ("",) + for modality in modalities: + prefix = f"{modality}_" if modality else "" + for parameter in ("weight", "bias") if config["proj_bias"] else ("weight",): + mapping[f"text_embedding_projection.{prefix}aggregate_embed.{parameter}"] = ( + f"{prefix}text_proj_in.{parameter}" + ) + for modality in ("video", "audio"): + old = f"{modality}_embeddings_connector" + new = f"{modality}_connector" + if config[f"{modality}_connector_num_learnable_registers"] is not None: + mapping[f"{old}.learnable_registers"] = f"{new}.learnable_registers" + for i in range(config[f"{modality}_connector_num_layers"]): + source, target = f"{old}.transformer_1d_blocks.{i}", f"{new}.transformer_blocks.{i}" + modules = ["attn1.to_q", "attn1.to_k", "attn1.to_v", "attn1.to_out.0", "ff.net.0.proj", "ff.net.2"] + if config[f"{modality}_gated_attn"]: + modules.append("attn1.to_gate_logits") + mapping.update( + {f"{source}.{name}.{p}": f"{target}.{name}.{p}" for name in modules for p in ("weight", "bias")} + ) + for kind in ("q", "k"): + mapping[f"{source}.attn1.{kind}_norm.weight"] = f"{target}.attn1.norm_{kind}.weight" + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx2_diffusion_decoder.py b/src/diffusers/loaders/conversion/ltx2_diffusion_decoder.py new file mode 100644 index 000000000000..ba5929386f5b --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_diffusion_decoder.py @@ -0,0 +1,75 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import FoldLinearGate, Split + + +def ltx2_diffusion_decoder_conversion(config): + """Convert the selected decoder and statistics; gated originals export with canonical unit gates.""" + mapping = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", + } + modules = [(name, name) for name in ("conv_in", "conv_in_x_t", "conv_out", "shared_adaln.proj")] + modules.extend((f"t_embedder.mlp.{i}", f"t_embedder.timestep_embedder.linear_{j}") for i, j in ((0, 1), (2, 2))) + modules.extend( + (f"upsamples.{i}.proj", f"upsamples.{i}.proj") for i in range(len(config["decoder_upsample_strides"])) + ) + mapping["decoder.norm_out.weight"] = "decoder.norm_out.weight" + rules = [] + gated = config.get("original_format", "ltx2_diffusion_decoder") == "ltx2_diffusion_decoder_gated" + blocks = [ + (f"det_stages.{i}.{j}", channels, False) + for i, channels in enumerate(config["decoder_stage_channels"][:-1]) + for j in range(config["decoder_stage_depths"][i]) + ] + blocks.extend( + (f"diff_blocks.{i}", config["decoder_stage_channels"][-1], True) + for i in range(config["decoder_stage_depths"][-1]) + ) + for name, width, diffusion in blocks: + prefix = f"decoder.{name}" + for norm in ("norm1", "norm2"): + mapping[f"{prefix}.{norm}.weight"] = f"{prefix}.{norm}.weight" + for kind in ("q", "k"): + mapping[f"{prefix}.attn.{kind}_norm.weight"] = f"{prefix}.attn.norm_{kind}.weight" + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attn.qkv.{p}",), + tuple(f"{prefix}.attn.to_{kind}.{p}" for kind in ("q", "k", "v")), + Split((width,) * 3), + ) + ) + for name_mlp in ("w_up", "w_gate"): + mapping[f"{prefix}.mlp.{name_mlp}.weight"] = f"{prefix}.mlp.{name_mlp}.weight" + leaves = [("attn.proj", "attn.to_out.0", "gate_msa", True), ("mlp.w_down", "mlp.w_down", "gate_mlp", False)] + if diffusion: + mapping[f"{prefix}.scale_shift_table"] = f"{prefix}.scale_shift_table" + leaves.append(("context_proj", "context_proj", "gate_ctx", True)) + for old, new, gate, bias in leaves: + parameters = ("weight", "bias") if bias else ("weight",) + if gated: + rules.append( + Rule( + (f"{prefix}.{gate}", *(f"{prefix}.{old}.{p}" for p in parameters)), + tuple(f"{prefix}.{new}.{p}" for p in parameters), + FoldLinearGate(bias=bias), + ) + ) + else: + mapping.update({f"{prefix}.{old}.{p}": f"{prefix}.{new}.{p}" for p in parameters}) + mapping.update({f"decoder.{old}.{p}": f"decoder.{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/ltx2_duration.py b/src/diffusers/loaders/conversion/ltx2_duration.py new file mode 100644 index 000000000000..284b1d7cff3a --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_duration.py @@ -0,0 +1,34 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Split + + +def ltx2_duration_conversion(config): + names = ["video_input_proj", "audio_input_proj", "mlp_hidden", "mlp_out"] + mapping = {f"{name}.{p}": f"{name}.{p}" for name in names for p in ("weight", "bias")} + for name in ("video_modality_emb", "audio_modality_emb", "attention_pooler.query_tokens"): + mapping[name] = name + rules = [] + for p in ("weight", "bias"): + mapping[f"attention_pooler.cross_attn.out_proj.{p}"] = f"attention_pooler.to_out.{p}" + rules.append( + Rule( + (f"attention_pooler.cross_attn.in_proj_{p}",), + tuple(f"attention_pooler.to_{kind}.{p}" for kind in ("q", "k", "v")), + Split((config["pooler_hidden_dim"],) * 3), + ) + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/ltx2_upsampler.py b/src/diffusers/loaders/conversion/ltx2_upsampler.py new file mode 100644 index 000000000000..4b86eafcd504 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_upsampler.py @@ -0,0 +1,31 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def ltx2_upsampler_conversion(config): + modules = ["initial_conv", "initial_norm", "final_conv"] + modules.extend( + f"{group}.{i}.{name}" + for group in ("res_blocks", "post_upsample_res_blocks") + for i in range(config["num_blocks_per_stage"]) + for name in ("conv1", "conv2", "norm1", "norm2") + ) + rational = config["spatial_upsample"] and not config["temporal_upsample"] and config["use_rational_resampler"] + modules.append("upsampler.conv" if rational else "upsampler.0") + mapping = {f"{name}.{p}": f"{name}.{p}" for name in modules for p in ("weight", "bias")} + if rational: + mapping["upsampler.blur_down.kernel"] = "upsampler.blur_down.kernel" + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx2_vae.py b/src/diffusers/loaders/conversion/ltx2_vae.py new file mode 100644 index 000000000000..5e6586032c3e --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_vae.py @@ -0,0 +1,21 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ltx_vae import _ltx_vae_mapping + + +def ltx2_vae_conversion(config): + return Conversion(mapping=_ltx_vae_mapping(config, ltx2=True)) diff --git a/src/diffusers/loaders/conversion/ltx2_vocoder.py b/src/diffusers/loaders/conversion/ltx2_vocoder.py new file mode 100644 index 000000000000..5fcaccb23c31 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx2_vocoder.py @@ -0,0 +1,54 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def ltx2_vocoder_conversion(config): + mapping = {} + bwe = config["_class_name"] == "LTX2VocoderWithBWE" + for prefix in ("vocoder.", "bwe_generator.") if bwe else ("",): + cfg = ( + {key.removeprefix("bwe_"): value for key, value in config.items() if key.startswith("bwe_")} + if prefix == "bwe_generator." + else config + ) + modules = [("conv_pre", "conv_in"), ("conv_post", "conv_out")] + for i in range(len(cfg["upsample_factors"])): + modules.append((f"ups.{i}", f"upsamplers.{i}")) + for j, dilations in enumerate(cfg["resnet_dilations"]): + index = i * len(cfg["resnet_dilations"]) + j + for n in range(len(dilations)): + for side in (1, 2): + modules.append((f"resblocks.{index}.convs{side}.{n}", f"resnets.{index}.convs{side}.{n}")) + old, new = f"resblocks.{index}.acts{side}.{n}", f"resnets.{index}.acts{side}.{n}" + if cfg["antialias"]: + mapping[f"{prefix}{old}.upsample.filter"] = f"{prefix}{new}.upsample.filter" + mapping[f"{prefix}{old}.downsample.lowpass.filter"] = f"{prefix}{new}.downsample.filter" + old, new = old + ".act", new + ".act" + if cfg["act_fn"] in ("snake", "snakebeta"): + for p in ("alpha", "beta") if cfg["act_fn"] == "snakebeta" else ("alpha",): + mapping[f"{prefix}{old}.{p}"] = f"{prefix}{new}.{p}" + for old, new in modules: + for p in ("weight",) if new == "conv_out" and not cfg["final_bias"] else ("weight", "bias"): + mapping[f"{prefix}{old}.{p}"] = f"{prefix}{new}.{p}" + if cfg["act_fn"] in ("snake", "snakebeta"): + for p in ("alpha", "beta"): + mapping[f"{prefix}act_post.act.{p}"] = f"{prefix}act_out.act.{p}" + mapping[f"{prefix}act_post.upsample.filter"] = f"{prefix}act_out.upsample.filter" + mapping[f"{prefix}act_post.downsample.lowpass.filter"] = f"{prefix}act_out.downsample.filter" + if bwe: + for key in ("mel_stft.mel_basis", "mel_stft.stft_fn.forward_basis", "mel_stft.stft_fn.inverse_basis"): + mapping[key] = key + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/ltx_upsampler.py b/src/diffusers/loaders/conversion/ltx_upsampler.py new file mode 100644 index 000000000000..fb7ea95da5ee --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx_upsampler.py @@ -0,0 +1,21 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .ltx2_upsampler import ltx2_upsampler_conversion + + +def ltx_upsampler_conversion(config): + base = ltx2_upsampler_conversion({**config, "use_rational_resampler": False}) + return Conversion(mapping=base.mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/ltx_vae.py b/src/diffusers/loaders/conversion/ltx_vae.py new file mode 100644 index 000000000000..f6f3aa9aa317 --- /dev/null +++ b/src/diffusers/loaders/conversion/ltx_vae.py @@ -0,0 +1,103 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def ltx_vae_conversion(config): + return Conversion(mapping=_ltx_vae_mapping(config, ltx2=False)) + + +def _ltx_vae_mapping(config, ltx2): + mapping = { + "per_channel_statistics.mean-of-means": "latents_mean", + "per_channel_statistics.std-of-means": "latents_std", + } + modules, resnets, embeddings = [], [], [] + for component in ("encoder", "decoder"): + modules.extend((f"{component}.{name}.conv", f"{component}.{name}.conv") for name in ("conv_in", "conv_out")) + channels = config["block_out_channels"] + modern = ltx2 or config["down_block_types"][-1] == "LTXVideo095DownBlock3D" + count = len(channels) - (1 if modern and not ltx2 else 0) + source_index = 0 + for i in range(count): + old, new = f"encoder.down_blocks.{source_index}", f"encoder.down_blocks.{i}" + resnets.extend( + (f"{old}.res_blocks.{j}", f"{new}.resnets.{j}", False, False, False) + for j in range(config["layers_per_block"][i]) + ) + source_index += 1 + if config["spatio_temporal_scaling"][i]: + suffix = "conv" if not modern or config["downsample_type"][i] == "conv" else "conv.conv" + modules.append((f"encoder.down_blocks.{source_index}.{suffix}", f"{new}.downsamplers.0.{suffix}")) + source_index += 1 + if not modern and i + 1 < count and channels[i] != channels[i + 1]: + resnets.append((f"encoder.down_blocks.{source_index}", new + ".conv_out", True, False, False)) + source_index += 1 + resnets.extend( + (f"encoder.down_blocks.{source_index}.res_blocks.{j}", f"encoder.mid_block.resnets.{j}", False, False, False) + for j in range(config["layers_per_block"][-1]) + ) + channels = list(reversed(config["decoder_block_out_channels"] or config["block_out_channels"])) + layers = list(reversed(config["decoder_layers_per_block"] or config["layers_per_block"])) + scaling = list(reversed(config["decoder_spatio_temporal_scaling"] or config["spatio_temporal_scaling"])) + noise = list(reversed(config["decoder_inject_noise"])) + factors = list(reversed(config["upsample_factor"])) + timed = config["timestep_conditioning"] + resnets.extend( + (f"decoder.up_blocks.0.res_blocks.{j}", f"decoder.mid_block.resnets.{j}", False, noise[0], timed) + for j in range(layers[0]) + ) + if timed: + embeddings.append(("decoder.up_blocks.0.time_embedder", "decoder.mid_block.time_embedder")) + embeddings.append(("decoder.time_embedder" if ltx2 else "decoder.last_time_embedder", "decoder.time_embedder")) + mapping["decoder.scale_shift_table" if ltx2 else "decoder.last_scale_shift_table"] = ( + "decoder.scale_shift_table" + ) + mapping["decoder.timestep_scale_multiplier"] = "decoder.timestep_scale_multiplier" + source_index = 1 + previous = channels[0] + for i, out_channels in enumerate(channels): + new = f"decoder.up_blocks.{i}" + if previous != out_channels: + resnets.append((f"decoder.up_blocks.{source_index}", new + ".conv_in", True, noise[i + 1], timed)) + source_index += 1 + if scaling[i]: + modules.append((f"decoder.up_blocks.{source_index}.conv.conv", new + ".upsamplers.0.conv.conv")) + source_index += 1 + old = f"decoder.up_blocks.{source_index}" + resnets.extend( + (f"{old}.res_blocks.{j}", f"{new}.resnets.{j}", False, noise[i + 1], timed) for j in range(layers[i + 1]) + ) + if timed: + embeddings.append((old + ".time_embedder", new + ".time_embedder")) + source_index += 1 + previous = out_channels // factors[i] + for old, new, shortcut, inject, timed in resnets: + modules.extend((f"{old}.{name}.conv", f"{new}.{name}.conv") for name in ("conv1", "conv2")) + if shortcut: + modules.append((old + (".conv_shortcut.conv" if ltx2 else ".conv_shortcut"), new + ".conv_shortcut.conv")) + modules.append((old + (".norm3" if ltx2 else ".norm3.norm"), new + ".norm3")) + if inject: + for name in ("per_channel_scale1", "per_channel_scale2"): + mapping[f"{old}.{name}"] = f"{new}.{name}" + if timed: + mapping[old + ".scale_shift_table"] = new + ".scale_shift_table" + for old, new in embeddings: + modules.extend( + (f"{old}.timestep_embedder.{name}", f"{new}.timestep_embedder.{name}") for name in ("linear_1", "linear_2") + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return mapping diff --git a/src/diffusers/loaders/conversion/lumina.py b/src/diffusers/loaders/conversion/lumina.py new file mode 100644 index 000000000000..dc90538f98e3 --- /dev/null +++ b/src/diffusers/loaders/conversion/lumina.py @@ -0,0 +1,76 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import MergeEqual, Reverse + + +def lumina_conversion(config): + mapping = {"pad_token": "pad_token"} + modules = [ + ("x_embedder", "patch_embedder.proj"), + ("cap_embedder.0", "time_caption_embed.caption_embedder.0"), + ("cap_embedder.1", "time_caption_embed.caption_embedder.1"), + ("final_layer.linear", "norm_out.linear_2"), + ("final_layer.adaLN_modulation.1", "norm_out.linear_1"), + ] + modules.extend( + (f"t_embedder.mlp.{i}", f"time_caption_embed.timestep_embedder.linear_{j}") for i, j in ((0, 1), (2, 2)) + ) + rules = [] + for i in range(config["num_layers"]): + prefix = f"layers.{i}" + mapping[prefix + ".attention.gate"] = prefix + ".gate" + modules.append((prefix + ".adaLN_modulation.1", prefix + ".norm1.linear")) + rules.append( + Rule( + (prefix + ".attention.wq.weight",), + (prefix + ".attn1.to_q.weight", prefix + ".attn2.to_q.weight"), + Reverse(MergeEqual(2)), + ) + ) + for old, new in ( + ("attention.wk", "attn1.to_k"), + ("attention.wv", "attn1.to_v"), + ("attention.wk_y", "attn2.to_k"), + ("attention.wv_y", "attn2.to_v"), + ("attention.wo", "attn2.to_out.0"), + ("attention_norm1", "norm1.norm"), + ("attention_norm2", "norm2"), + ("attention_y_norm", "norm1_context"), + ("feed_forward.w1", "feed_forward.linear_1"), + ("feed_forward.w2", "feed_forward.linear_2"), + ("feed_forward.w3", "feed_forward.linear_3"), + ("ffn_norm1", "ffn_norm1"), + ("ffn_norm2", "ffn_norm2"), + ): + mapping[f"{prefix}.{old}.weight"] = f"{prefix}.{new}.weight" + if config["qk_norm"]: + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attention.q_norm.{p}",), + (f"{prefix}.attn1.norm_q.{p}", f"{prefix}.attn2.norm_q.{p}"), + Reverse(MergeEqual(2)), + ) + ) + modules.extend( + [ + (prefix + ".attention.k_norm", prefix + ".attn1.norm_k"), + (prefix + ".attention.ky_norm", prefix + ".attn2.norm_k"), + ] + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/lumina2.py b/src/diffusers/loaders/conversion/lumina2.py new file mode 100644 index 000000000000..71e3f76a1de4 --- /dev/null +++ b/src/diffusers/loaders/conversion/lumina2.py @@ -0,0 +1,62 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Split + + +def lumina2_conversion(config): + modules = [ + ("x_embedder", "x_embedder"), + ("t_embedder.mlp.0", "time_caption_embed.timestep_embedder.linear_1"), + ("t_embedder.mlp.2", "time_caption_embed.timestep_embedder.linear_2"), + ("cap_embedder.1", "time_caption_embed.caption_embedder.1"), + ("final_layer.adaLN_modulation.1", "norm_out.linear_1"), + ("final_layer.linear", "norm_out.linear_2"), + ] + mapping = {"cap_embedder.0.weight": "time_caption_embed.caption_embedder.0.weight"} + rules = [] + hidden = config["hidden_size"] + kv = hidden // config["num_attention_heads"] * config["num_kv_heads"] + for group, count in ( + ("noise_refiner", config["num_refiner_layers"]), + ("context_refiner", config["num_refiner_layers"]), + ("layers", config["num_layers"]), + ): + for i in range(count): + prefix = f"{group}.{i}" + rules.append( + Rule( + (prefix + ".attention.qkv.weight",), + tuple(f"{prefix}.attn.to_{part}.weight" for part in ("q", "k", "v")), + Split((hidden, kv, kv)), + ) + ) + for old, new in ( + ("attention.q_norm", "attn.norm_q"), + ("attention.k_norm", "attn.norm_k"), + ("attention.out", "attn.to_out.0"), + ("feed_forward.w1", "feed_forward.linear_1"), + ("feed_forward.w2", "feed_forward.linear_2"), + ("feed_forward.w3", "feed_forward.linear_3"), + ("attention_norm1", "norm1" if group == "context_refiner" else "norm1.norm"), + ("attention_norm2", "norm2"), + ("ffn_norm1", "ffn_norm1"), + ("ffn_norm2", "ffn_norm2"), + ): + mapping[f"{prefix}.{old}.weight"] = f"{prefix}.{new}.weight" + if group != "context_refiner": + modules.append((prefix + ".adaLN_modulation.1", prefix + ".norm1.linear")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/minimax_h3.py b/src/diffusers/loaders/conversion/minimax_h3.py new file mode 100644 index 000000000000..9235e001f52f --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_h3.py @@ -0,0 +1,81 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Permute, ReorderChunks, Reshape, Split + + +def minimax_h3_conversion(config): + mapping = { + "token_refiner.final_norm.weight": "token_refiner.final_norm.weight", + "final_layer.norm.weight": "norm_out.norm.weight", + } + modules = [ + ("video_patch_proj", "proj_in"), + ("audio_patch_proj", "audio_proj_in"), + ("condition_proj", "context_embedder"), + ("time_embedder.proj_in", "time_embedder.linear_1"), + ("time_embedder.proj_out", "time_embedder.linear_2"), + ("final_layer.adaln_proj.linear", "norm_out.linear"), + ("final_layer.video_out", "proj_out"), + ("final_layer.audio_out", "audio_proj_out"), + ] + heads, head_dim = config["num_attention_heads"], config["attention_head_dim"] + inner, hidden = heads * head_dim, config["hidden_size"] + original_format = config.get("original_format", "minimax_h3") + if original_format not in ("minimax_h3", "minimax_h3_shards"): + raise ValueError("MiniMax H3 original_format must be 'minimax_h3' or 'minimax_h3_shards'.") + split = Split((inner,) * 3) + if original_format == "minimax_h3_shards": + split = Chain( + ( + Reshape((3 * inner, hidden), (heads, 3, head_dim, hidden)), + Permute((1, 0, 2, 3)), + Reshape((3, heads, head_dim, hidden), (3 * inner, hidden)), + split, + ) + ) + rules = [] + for old_group, new_group, count, modulated in ( + ("blocks", "transformer_blocks", config["num_layers"], True), + ("token_refiner.blocks", "token_refiner.refiner_blocks", config["num_refiner_layers"], False), + ): + for i in range(count): + old, new = f"{old_group}.{i}", f"{new_group}.{i}" + mapping.update( + { + f"{old}.{a}.weight": f"{new}.{b}.weight" + for a, b in ( + ("norm1", "norm1"), + ("norm2", "norm2"), + ("attn.q_norm", "attn.norm_q"), + ("attn.k_norm", "attn.norm_k"), + ("attn.out_proj", "attn.to_out.0"), + ("mlp.fc2", "ff.net.2"), + ) + } + ) + rules.append( + Rule( + (old + ".attn.qkv_proj.weight",), + tuple(f"{new}.attn.to_{part}.weight" for part in ("q", "k", "v")), + split, + ) + ) + rules.append(Rule((old + ".mlp.fc1.weight",), (new + ".ff.net.0.proj.weight",), ReorderChunks((1, 0)))) + if modulated: + modules.append((old + ".adaln_proj.linear", new + ".adaln_proj.linear")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/minimax_h3_audio_vae.py b/src/diffusers/loaders/conversion/minimax_h3_audio_vae.py new file mode 100644 index 000000000000..47bedbd52c57 --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_h3_audio_vae.py @@ -0,0 +1,76 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def minimax_h3_audio_vae_conversion(config): + modules = [ + "mean_proj", + "logs_proj", + "dec_in_proj", + "pre_block.norm1", + "pre_block.norm2", + "pre_block.norm3", + "pre_block.proj", + "pre_block.attn.proj", + "pre_block.mlp.norm", + "pre_block.mlp.w0", + "pre_block.mlp.w1", + "pre_block.mlp.w2", + ] + keys = [ + "pre_block.attn.qkv.weight", + "pre_block.attn.q_bias", + "pre_block.attn.v_bias", + "pre_block.attn.zero_k_bias", + ] + count = len(config["encoder_rates"]) + convolutions = [ + ("encoder.block.0", True), + (f"encoder.block.{count + 2}", True), + ("decoder.conv_pre", True), + ("decoder.conv_post", False), + ] + keys.append(f"encoder.block.{count + 1}.alpha") + for i in range(count): + prefix = f"encoder.block.{i + 1}.block" + keys.append(prefix + ".3.alpha") + convolutions.append((prefix + ".4", True)) + for j in range(3): + unit = f"{prefix}.{j}.block" + keys.extend(unit + f".{k}.alpha" for k in (0, 2)) + convolutions.extend((unit + f".{k}", True) for k in (1, 3)) + activations = ["decoder.activation_post"] + kernels = len(config["resblock_kernel_sizes"]) + for i in range(len(config["decoder_rates"])): + convolutions.append((f"decoder.ups.{i}.0", True)) + for j in range(kernels): + prefix = f"decoder.resblocks.{i * kernels + j}" + depth = len(config["resblock_dilation_sizes"][j]) + convolutions.extend((f"{prefix}.convs{part}.{k}", True) for part in (1, 2) for k in range(depth)) + activations.extend(f"{prefix}.activations.{k}" for k in range(2 * depth)) + keys.extend( + f"{prefix}.{name}" + for prefix in activations + for name in ("act.alpha", "act.beta", "upsample.filter", "downsample.lowpass.filter") + ) + keys.extend( + f"{name}.{p}" + for name, bias in convolutions + for p in (("weight_g", "weight_v", "bias") if bias else ("weight_g", "weight_v")) + ) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/minimax_h3_vae.py b/src/diffusers/loaders/conversion/minimax_h3_vae.py new file mode 100644 index 000000000000..bdac53e80948 --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_h3_vae.py @@ -0,0 +1,75 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Permute, ReorderChunks, Reshape, Split + + +def minimax_h3_vae_conversion(config): + modules = [ + (name, name) + for name in ( + "quant_conv", + "post_quant_conv", + "encoder.conv_in", + "encoder.norm_out", + "encoder.conv_out", + "decoder.norm_out", + "decoder.proj_out", + ) + ] + modules.append(("decoder.x_embedder", "decoder.proj_in")) + mapping = {"decoder.register_tokens": "decoder.register_tokens"} + previous = config["block_out_channels"][0] + for i, channel in enumerate(config["block_out_channels"]): + for j in range(config["layers_per_block"]): + old, new = f"encoder.down.{i}.block.{j}", f"encoder.down_blocks.{i}.resnets.{j}" + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("norm1", "conv1", "norm2", "conv2")) + if previous != channel: + modules.append((old + ".nin_shortcut", new + ".conv_shortcut")) + previous = channel + if config["spatial_downsample_factors"][i] * config["temporal_downsample_factors"][i] > 1: + modules.append((f"encoder.down.{i}.downsample.conv", f"encoder.down_blocks.{i}.downsamplers.0.conv")) + heads, head_dim = config["decoder_num_attention_heads"], config["decoder_attention_head_dim"] + hidden = heads * head_dim + rules = [] + for i in range(config["decoder_num_layers"]): + prefix = f"decoder.transformer_blocks.{i}" + mapping.update( + {f"{prefix}.{name}": f"{prefix}.{name}" for name in ("norm1.weight", "norm2.weight", "scale1", "scale2")} + ) + modules.extend( + [(prefix + ".attn.to_out", prefix + ".attn.to_out.0"), (prefix + ".ff.w2", prefix + ".ff.net.2")] + ) + for p in ("weight", "bias"): + trailing = (hidden,) if p == "weight" else () + transform = Chain( + ( + Reshape((3 * hidden,) + trailing, (heads, 3, head_dim) + trailing), + Permute((1, 0, 2, 3) if p == "weight" else (1, 0, 2)), + Reshape((3, heads, head_dim) + trailing, (3 * hidden,) + trailing), + Split((hidden,) * 3), + ) + ) + rules.append( + Rule( + (f"{prefix}.attn.to_qkv.{p}",), + tuple(f"{prefix}.attn.to_{part}.{p}" for part in ("q", "k", "v")), + transform, + ) + ) + rules.append(Rule((f"{prefix}.ff.w1.{p}",), (f"{prefix}.ff.net.0.proj.{p}",), ReorderChunks((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/minimax_music3.py b/src/diffusers/loaders/conversion/minimax_music3.py new file mode 100644 index 000000000000..0f6a4c7039b2 --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_music3.py @@ -0,0 +1,46 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def minimax_music3_conversion(config): + mapping = { + "timestep_features.weight": "time_proj.weight", + "preprocess_conv.weight": "preprocess_conv.weight", + "postprocess_conv.weight": "postprocess_conv.weight", + "transformer.project_in.weight": "proj_in.weight", + "transformer.project_out.weight": "proj_out.weight", + } + modules = [("to_timestep_embed.0", "time_embed.linear_1"), ("to_timestep_embed.2", "time_embed.linear_2")] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + rules = [] + for i in range(config["num_layers"]): + old, new = f"transformer.layers.{i}", f"transformer_blocks.{i}" + for a, b in (("pre_norm", "norm1"), ("ff_norm", "norm2")): + mapping[f"{old}.{a}.gamma"] = f"{new}.{b}.weight" + mapping[f"{old}.{a}.beta"] = f"{new}.{b}.bias" + rules.append( + Rule( + (old + ".self_attn.to_qkv.weight",), + tuple(f"{new}.attn.to_{part}.weight" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + mapping[old + ".self_attn.to_out.weight"] = new + ".attn.to_out.0.weight" + modules.extend([(old + ".ff.ff.0.proj", new + ".ff_in"), (old + ".ff.ff.2", new + ".ff_out")]) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/minimax_music3_conditioner.py b/src/diffusers/loaders/conversion/minimax_music3_conditioner.py new file mode 100644 index 000000000000..847741fa9043 --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_music3_conditioner.py @@ -0,0 +1,27 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def minimax_music3_conditioner_conversion(config): + return Conversion( + mapping={ + "cond_layer_logits": "layer_weight_logits", + "cond_layer_scale": "layer_scale", + "latent_conditioners.0.weight": "proj.weight", + "latent_conditioners.0.bias": "proj.bias", + } + ) diff --git a/src/diffusers/loaders/conversion/minimax_music3_rvq.py b/src/diffusers/loaders/conversion/minimax_music3_rvq.py new file mode 100644 index 000000000000..da74a075527c --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_music3_rvq.py @@ -0,0 +1,42 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def minimax_music3_rvq_conversion(config): + mapping = {"model.audio_extra_embedding.weight": "audio_embeddings.weight"} + modules = [("model.audio_decoder." + name, name) for name in ("projection", "pos_embedding", "norm")] + modules.extend( + (f"model.audio_decoder.audio_heads.{i}", f"audio_heads.{i}") for i in range(config["num_codebooks"] - 1) + ) + for i in range(config["num_layers"]): + old, new = f"model.audio_decoder.layers.{i}", f"layers.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("input_layernorm", "input_layernorm"), + ("post_attention_layernorm", "post_attention_layernorm"), + ("self_attn.q_proj", "attn.to_q"), + ("self_attn.k_proj", "attn.to_k"), + ("self_attn.v_proj", "attn.to_v"), + ("self_attn.o_proj", "attn.to_out"), + ("mlp.gate_proj", "gate_proj"), + ("mlp.up_proj", "up_proj"), + ("mlp.down_proj", "down_proj"), + ) + ) + mapping.update({old + ".weight": new + ".weight" for old, new in modules}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/minimax_music3_vocoder.py b/src/diffusers/loaders/conversion/minimax_music3_vocoder.py new file mode 100644 index 000000000000..2aa7f7100d16 --- /dev/null +++ b/src/diffusers/loaders/conversion/minimax_music3_vocoder.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def minimax_music3_vocoder_conversion(config): + count = len(config["upsampling_ratios"]) + mapping = { + "dec_in_proj.weight": "dec_in_proj.weight", + "dec_in_proj.bias": "dec_in_proj.bias", + f"decoder.model.{count + 1}.alpha": "snake_out.alpha", + } + convolutions = [("decoder.model.0", "conv_in"), (f"decoder.model.{count + 2}", "conv_out")] + for i in range(count): + old, new = f"decoder.model.{i + 1}.block", f"blocks.{i}" + mapping[old + ".0.alpha"] = new + ".snake1.alpha" + convolutions.append((old + ".1", new + ".conv_t1")) + for j in range(3): + a, b = f"{old}.{j + 2}.block", f"{new}.res_unit{j + 1}" + mapping[a + ".0.alpha"] = b + ".snake1.alpha" + mapping[a + ".2.alpha"] = b + ".snake2.alpha" + convolutions.extend([(a + ".1", b + ".conv1"), (a + ".3", b + ".conv2")]) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in convolutions for p in ("weight_g", "weight_v", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/mochi.py b/src/diffusers/loaders/conversion/mochi.py new file mode 100644 index 000000000000..96bab5a56d08 --- /dev/null +++ b/src/diffusers/loaders/conversion/mochi.py @@ -0,0 +1,70 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def mochi_conversion(config): + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + mapping = {"pos_frequencies": "pos_frequencies"} + modules = [ + ("x_embedder.proj", "patch_embed.proj"), + ("t_embedder.mlp.0", "time_embed.timestep_embedder.linear_1"), + ("t_embedder.mlp.2", "time_embed.timestep_embedder.linear_2"), + ("t5_yproj", "time_embed.caption_proj"), + ("final_layer.linear", "proj_out"), + ] + modules.extend(("t5_y_embedder." + name, "time_embed.pooler." + name) for name in ("to_q", "to_kv", "to_out")) + rules = [] + for i in range(config["num_layers"]): + old, new = f"blocks.{i}", f"transformer_blocks.{i}" + modules.extend( + [ + (old + ".mod_x", new + ".norm1.linear"), + ( + old + ".mod_y", + new + (".norm1_context.linear_1" if i == config["num_layers"] - 1 else ".norm1_context.linear"), + ), + ] + ) + for branch, attention, output, norms, ff in ( + ("x", ("to_q", "to_k", "to_v"), "to_out.0", ("norm_q", "norm_k"), "ff"), + ( + "y", + ("add_q_proj", "add_k_proj", "add_v_proj"), + "to_add_out", + ("norm_added_q", "norm_added_k"), + "ff_context", + ), + ): + rules.append( + Rule( + (f"{old}.attn.qkv_{branch}.weight",), + tuple(f"{new}.attn1.{name}.weight" for name in attention), + Split((hidden_size,) * 3), + ) + ) + for source, target in zip(("q", "k"), norms): + mapping[f"{old}.attn.{source}_norm_{branch}.weight"] = f"{new}.attn1.{target}.weight" + if branch == "x" or i < config["num_layers"] - 1: + modules.append((f"{old}.attn.proj_{branch}", f"{new}.attn1.{output}")) + rules.append( + Rule((f"{old}.mlp_{branch}.w1.weight",), (f"{new}.{ff}.net.0.proj.weight",), ReorderChunks((1, 0))) + ) + mapping[f"{old}.mlp_{branch}.w2.weight"] = f"{new}.{ff}.net.2.weight" + for p in ("weight", "bias"): + rules.append(Rule((f"final_layer.mod.{p}",), (f"norm_out.linear.{p}",), ReorderChunks((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/mochi_vae.py b/src/diffusers/loaders/conversion/mochi_vae.py new file mode 100644 index 000000000000..9f1d51becca9 --- /dev/null +++ b/src/diffusers/loaders/conversion/mochi_vae.py @@ -0,0 +1,82 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def mochi_vae_conversion(config): + modules = [ + ("decoder.blocks.0.0", "decoder.conv_in"), + ("decoder.output_proj", "decoder.proj_out"), + ("encoder.layers.0", "encoder.proj_in"), + ("encoder.output_norm", "encoder.norm_out.norm_layer"), + ] + mapping = {"encoder.output_proj.weight": "encoder.proj_out.weight"} + resnets, attentions = [], [] + layers = config["layers_per_block"] + decoder_count = len(config["decoder_block_out_channels"]) - 1 + resnets.extend((f"decoder.blocks.0.{i + 1}", f"decoder.block_in.resnets.{i}") for i in range(layers[-1])) + for i in range(decoder_count): + resnets.extend( + (f"decoder.blocks.{i + 1}.blocks.{j}", f"decoder.up_blocks.{i}.resnets.{j}") for j in range(layers[-i - 2]) + ) + modules.append((f"decoder.blocks.{i + 1}.proj", f"decoder.up_blocks.{i}.proj")) + resnets.extend( + (f"decoder.blocks.{decoder_count + 1}.{i}", f"decoder.block_out.resnets.{i}") for i in range(layers[0]) + ) + for i in range(layers[0]): + old = f"encoder.layers.{i + 1}" + resnets.append((old, f"encoder.block_in.resnets.{i}")) + if config["add_attention_block"][0]: + attentions.append((old + ".attn_block", "encoder.block_in", i, config["encoder_block_out_channels"][0])) + encoder_count = len(config["encoder_block_out_channels"]) - 1 + offset = 1 + layers[0] + for i in range(encoder_count): + old, new = f"encoder.layers.{offset + i}", f"encoder.down_blocks.{i}" + modules.append((old + ".layers.0", new + ".conv_in.conv")) + for j in range(layers[i + 1]): + resnets.append((f"{old}.layers.{j + 1}", f"{new}.resnets.{j}")) + if config["add_attention_block"][i + 1]: + attentions.append( + (f"{old}.layers.{j + 1}.attn_block", new, j, config["encoder_block_out_channels"][i + 1]) + ) + for i in range(layers[-1]): + old = f"encoder.layers.{offset + encoder_count + i}" + resnets.append((old, f"encoder.block_out.resnets.{i}")) + if config["add_attention_block"][-1]: + attentions.append((old + ".attn_block", "encoder.block_out", i, config["encoder_block_out_channels"][-1])) + for old, new in resnets: + modules.extend( + (f"{old}.stack.{i}", f"{new}.{name}") + for i, name in ((0, "norm1.norm_layer"), (2, "conv1.conv"), (3, "norm2.norm_layer"), (5, "conv2.conv")) + ) + rules = [] + for old, new, index, hidden in attentions: + modules.extend( + [ + (old + ".norm", f"{new}.norms.{index}.norm_layer"), + (old + ".attn.out", f"{new}.attentions.{index}.to_out.0"), + ] + ) + rules.append( + Rule( + (old + ".attn.qkv.weight",), + tuple(f"{new}.attentions.{index}.to_{part}.weight" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/motif_video.py b/src/diffusers/loaders/conversion/motif_video.py new file mode 100644 index 000000000000..439eae9ec50f --- /dev/null +++ b/src/diffusers/loaders/conversion/motif_video.py @@ -0,0 +1,69 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def motif_video_conversion(config): + """Validate the native Motif Video checkpoint layout used by its single-file loader.""" + modules = [ + "x_embedder.proj", + "context_embedder.linear_1", + "context_embedder.linear_2", + "time_text_embed.timestep_embedder.linear_1", + "time_text_embed.timestep_embedder.linear_2", + "norm_out.linear", + "proj_out", + ] + keys = [] + if config["image_embed_dim"] is not None: + modules.extend(f"image_embedder.{name}" for name in ("norm_in", "linear_1", "linear_2", "norm_out")) + for group, count in ( + ("transformer_blocks", config["num_layers"]), + ("single_transformer_blocks", config["num_single_layers"]), + ): + dual = group == "transformer_blocks" + for i in range(count): + prefix = f"{group}.{i}" + names = ( + ( + "norm1.linear", + "norm1_context.linear", + "ff.net.0.proj", + "ff.net.2", + "ff_context.net.0.proj", + "ff_context.net.2", + ) + if dual + else ("norm.linear", "proj_mlp", "proj_out") + ) + modules.extend(f"{prefix}.{name}" for name in names) + cross = ( + config["enable_text_cross_attention_dual"] + if dual + else (config["enable_text_cross_attention_single"] and i < count - config["num_decoder_layers"]) + ) + for attention in ("attn", "cross_attn") if cross else ("attn",): + names = ["to_q", "to_k", "to_v"] + if dual or attention == "cross_attn": + names.append("to_out.0") + if dual and attention == "attn": + names.extend(("add_q_proj", "add_k_proj", "add_v_proj", "to_add_out")) + keys.extend(f"{prefix}.attn.norm_added_{part}.weight" for part in ("q", "k")) + modules.extend(f"{prefix}.{attention}.{name}" for name in names) + if config["qk_norm"] in ("rms_norm", "layer_norm"): + parameters = ("weight", "bias") if config["qk_norm"] == "layer_norm" else ("weight",) + keys.extend(f"{prefix}.{attention}.norm_{part}.{p}" for part in ("q", "k") for p in parameters) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/movq.py b/src/diffusers/loaders/conversion/movq.py new file mode 100644 index 000000000000..0d8c652be319 --- /dev/null +++ b/src/diffusers/loaders/conversion/movq.py @@ -0,0 +1,32 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def movq_conversion(config): + base = ldm_vae_conversion(config) + mapping = {} + for old, new in base.mapping.items(): + module, _, parameter = new.rpartition(".") + if new.startswith("decoder.") and module.endswith((".norm1", ".norm2", ".group_norm", ".conv_norm_out")): + source = old.rsplit(".", 1)[0] + module = module.removesuffix("group_norm") + "spatial_norm" if module.endswith("group_norm") else module + for name in ("norm_layer", "conv_y", "conv_b"): + mapping[f"{source}.{name}.{parameter}"] = f"{module}.{name}.{parameter}" + else: + mapping[old] = new + mapping["quantize.embedding.weight"] = "quantize.embedding.weight" + return Conversion(mapping=mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/ncsnpp.py b/src/diffusers/loaders/conversion/ncsnpp.py new file mode 100644 index 000000000000..1a8acd8517c9 --- /dev/null +++ b/src/diffusers/loaders/conversion/ncsnpp.py @@ -0,0 +1,90 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Permute + + +def ncsnpp_conversion(config): + if not all("Skip" in kind for kind in (*config["down_block_types"], *config["up_block_types"])): + raise ValueError("The NCSN++ checkpoint format requires SkipDown/SkipUp block types.") + mapping = {"all_modules.0.W": "time_proj.weight"} + modules = [ + ("all_modules.1", "time_embedding.linear_1"), + ("all_modules.2", "time_embedding.linear_2"), + ("all_modules.3", "conv_in"), + ] + resnets, attentions = [], [] + index = 4 + channels, layers = config["block_out_channels"], config["layers_per_block"] + previous = channels[0] + for i, channel in enumerate(channels): + for j in range(layers): + resnets.append((index, f"down_blocks.{i}.resnets.{j}", previous != channel)) + index += 1 + previous = channel + if "Attn" in config["down_block_types"][i]: + attentions.append((index, f"down_blocks.{i}.attentions.{j}")) + index += 1 + if i < len(channels) - 1: + resnets.append((index, f"down_blocks.{i}.resnet_down", True)) + index += 1 + modules.append((f"all_modules.{index}.Conv_0", f"down_blocks.{i}.skip_conv")) + index += 1 + resnets.append((index, "mid_block.resnets.0", False)) + attentions.append((index + 1, "mid_block.attentions.0")) + resnets.append((index + 2, "mid_block.resnets.1", False)) + index += 3 + reversed_channels = list(reversed(channels)) + previous = reversed_channels[0] + for i, channel in enumerate(reversed_channels): + next_channel = reversed_channels[min(i + 1, len(channels) - 1)] + for j in range(layers + 1): + skip = next_channel if j == layers else channel + resnets.append((index, f"up_blocks.{i}.resnets.{j}", previous + skip != channel)) + previous = channel + index += 1 + if "Attn" in config["up_block_types"][i]: + attentions.append((index, f"up_blocks.{i}.attentions.0")) + index += 1 + if i < len(channels) - 1: + modules.append((f"all_modules.{index}", f"up_blocks.{i}.skip_norm")) + modules.append((f"all_modules.{index + 1}", f"up_blocks.{i}.skip_conv")) + resnets.append((index + 2, f"up_blocks.{i}.resnet_up", True)) + index += 3 + modules.extend([(f"all_modules.{index}", "conv_norm_out"), (f"all_modules.{index + 1}", "conv_out")]) + for index, new, shortcut in resnets: + old = f"all_modules.{index}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("Conv_0", "conv1"), + ("Conv_1", "conv2"), + ("GroupNorm_0", "norm1"), + ("GroupNorm_1", "norm2"), + ("Dense_0", "time_emb_proj"), + ) + ) + if shortcut: + modules.append((old + ".Conv_2", new + ".conv_shortcut")) + rules = [] + for index, new in attentions: + old = f"all_modules.{index}" + modules.append((old + ".GroupNorm_0", new + ".group_norm")) + for i, projection in enumerate(("to_q", "to_k", "to_v", "to_out.0")): + mapping[f"{old}.NIN_{i}.b"] = f"{new}.{projection}.bias" + rules.append(Rule((f"{old}.NIN_{i}.W",), (f"{new}.{projection}.weight",), Permute((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/omnigen.py b/src/diffusers/loaders/conversion/omnigen.py new file mode 100644 index 000000000000..6b232e47beb3 --- /dev/null +++ b/src/diffusers/loaders/conversion/omnigen.py @@ -0,0 +1,57 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def omnigen_conversion(config): + mapping = { + "pos_embed": "patch_embedding.pos_embed", + "llm.embed_tokens.weight": "embed_tokens.weight", + "llm.norm.weight": "norm.weight", + } + modules = [ + ("x_embedder.proj", "patch_embedding.output_image_proj"), + ("input_x_embedder.proj", "patch_embedding.input_image_proj"), + ("final_layer.adaLN_modulation.1", "norm_out.linear"), + ("final_layer.linear", "proj_out"), + ] + modules.extend( + (f"{prefix}.mlp.{i}", f"{prefix}.linear_{j}") + for prefix in ("time_token", "t_embedder") + for i, j in ((0, 1), (2, 2)) + ) + hidden = config["hidden_size"] + kv = hidden // config["num_attention_heads"] * config["num_key_value_heads"] + rules = [] + for i in range(config["num_layers"]): + old, new = f"llm.layers.{i}", f"layers.{i}" + rules.append( + Rule( + (old + ".self_attn.qkv_proj.weight",), + tuple(f"{new}.self_attn.to_{part}.weight" for part in ("q", "k", "v")), + Split((hidden, kv, kv)), + ) + ) + mapping[old + ".self_attn.o_proj.weight"] = new + ".self_attn.to_out.0.weight" + mapping.update( + { + f"{old}.{name}.weight": f"{new}.{name}.weight" + for name in ("input_layernorm", "post_attention_layernorm", "mlp.gate_up_proj", "mlp.down_proj") + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/oobleck.py b/src/diffusers/loaders/conversion/oobleck.py new file mode 100644 index 000000000000..9c852f4c260a --- /dev/null +++ b/src/diffusers/loaders/conversion/oobleck.py @@ -0,0 +1,62 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Reshape + + +def oobleck_conversion(config): + mapping, rules, convolutions, snakes = {}, [], [], [] + count = len(config["downsampling_ratios"]) + multiples = [1] + list(config["channel_multiples"]) + for component in ("encoder", "decoder"): + base = config["encoder_hidden_size"] if component == "encoder" else config["decoder_channels"] + convolutions.extend( + [ + (f"{component}.layers.0", f"{component}.conv1", True), + (f"{component}.layers.{count + 2}", f"{component}.conv2", component == "encoder"), + ] + ) + snakes.append( + ( + f"{component}.layers.{count + 1}", + f"{component}.snake1", + base * (multiples[-1] if component == "encoder" else 1), + ) + ) + for i in range(count): + old, new = f"{component}.layers.{i + 1}", f"{component}.block.{i}" + encoder = component == "encoder" + input_dim = base * multiples[i if encoder else count - i] + unit_dim = input_dim if encoder else base * multiples[count - i - 1] + snakes.append((old + (".layers.3" if encoder else ".layers.0"), new + ".snake1", input_dim)) + convolutions.append( + (old + (".layers.4" if encoder else ".layers.1"), new + (".conv1" if encoder else ".conv_t1"), True) + ) + for j in range(3): + a, b = f"{old}.layers.{j if encoder else j + 2}", f"{new}.res_unit{j + 1}" + snakes.extend([(a + ".layers.0", b + ".snake1", unit_dim), (a + ".layers.2", b + ".snake2", unit_dim)]) + convolutions.extend([(a + ".layers.1", b + ".conv1", True), (a + ".layers.3", b + ".conv2", True)]) + for old, new, bias in convolutions: + mapping.update( + { + f"{old}.{p}": f"{new}.{p}" + for p in (("weight_g", "weight_v", "bias") if bias else ("weight_g", "weight_v")) + } + ) + for old, new, channels in snakes: + for p in ("alpha", "beta"): + rules.append(Rule((f"{old}.{p}",), (f"{new}.{p}",), Reshape((channels,), (1, channels, 1)))) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/ovis_image.py b/src/diffusers/loaders/conversion/ovis_image.py new file mode 100644 index 000000000000..1bf7a62a2290 --- /dev/null +++ b/src/diffusers/loaders/conversion/ovis_image.py @@ -0,0 +1,55 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .flux import flux_conversion +from .transforms import Reverse, Split + + +def ovis_image_conversion(config): + base = flux_conversion({**config, "guidance_embeds": False}) + mapping = {} + for old, new in base.mapping.items(): + if old.startswith("vector_in.") or "_mlp.0." in old: + continue + old = ( + old.replace("txt_in.", "semantic_txt_in.") + .replace(".scale", ".weight") + .replace("_mlp.2.", "_mlp.down_proj.") + ) + new = new.replace("time_text_embed.timestep_embedder.", "timestep_embedder.") + mapping[old] = new + mapping["semantic_txt_norm.weight"] = "context_embedder_norm.weight" + hidden = config["num_attention_heads"] * config["attention_head_dim"] + mlp = 4 * hidden + rules = [] + for rule in base.rules: + if rule.original[0].startswith("single_blocks.") and ".linear1." in rule.original[0]: + rule = Rule(rule.original, rule.diffusers, Split((hidden,) * 3 + (2 * mlp,))) + rules.append(rule) + for i in range(config["num_layers"]): + for source, target in (("img", "ff"), ("txt", "ff_context")): + for p in ("weight", "bias"): + rules.append( + Rule( + ( + f"double_blocks.{i}.{source}_mlp.up_proj.{p}", + f"double_blocks.{i}.{source}_mlp.gate_proj.{p}", + ), + (f"transformer_blocks.{i}.{target}.net.0.proj.{p}",), + Reverse(Split((mlp,) * 2)), + ) + ) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/paella.py b/src/diffusers/loaders/conversion/paella.py new file mode 100644 index 000000000000..e094e94594c6 --- /dev/null +++ b/src/diffusers/loaders/conversion/paella.py @@ -0,0 +1,50 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def paella_conversion(config): + modules = ["in_block.1", "up_blocks.0.0", "out_block.0"] + residuals = [] + for i in range(config["levels"]): + if i > 0: + modules.append(f"down_blocks.{2 * i - 1}") + residuals.append(f"down_blocks.{2 * i}") + output = 2 * config["levels"] - 1 + mapping = { + "vquantizer.codebook.weight": "vquantizer.embedding.weight", + f"down_blocks.{output}.0.weight": f"down_blocks.{output}.0.weight", + } + modules.append(f"down_blocks.{output}.1") + mapping.update( + { + f"down_blocks.{output}.1.{name}": f"down_blocks.{output}.1.{name}" + for name in ("running_mean", "running_var", "num_batches_tracked") + } + ) + index = 1 + for i in range(config["levels"]): + for _ in range(config["bottleneck_blocks"] if i == 0 else 1): + residuals.append(f"up_blocks.{index}") + index += 1 + if i < config["levels"] - 1: + modules.append(f"up_blocks.{index}") + index += 1 + for prefix in residuals: + mapping[prefix + ".gammas"] = prefix + ".gammas" + modules.extend(f"{prefix}.{name}" for name in ("depthwise.1", "channelwise.0", "channelwise.2")) + mapping.update({f"{name}.{p}": f"{name}.{p}" for name in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/pipeline.py b/src/diffusers/loaders/conversion/pipeline.py new file mode 100644 index 000000000000..7e5543d75599 --- /dev/null +++ b/src/diffusers/loaders/conversion/pipeline.py @@ -0,0 +1,97 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Package converted components into the original Stable Diffusion checkpoint container.""" + +import json +import tempfile +from pathlib import Path + +import torch +from safetensors.torch import save_file + +from .io import Checkpoint +from .registry import get_conversion + + +def export_pipeline_checkpoint( + model_path, output_path, *, pipeline_format=None, output_format="safetensors", dtype=None +): + """Export SD/SDXL weights using the component conversions and original component namespaces. + + Tokenizers, schedulers, safety checkers and runtime configuration are separate assets. A single output file must + hold all component tensors during serialization; use the component CLI when sharded output is desired. + """ + model_path, output_path = Path(model_path), Path(output_path) + if output_path.exists(): + raise FileExistsError(f"Output already exists: {output_path}") + if output_format not in ("safetensors", "pytorch"): + raise ValueError("output_format must be safetensors or pytorch") + if pipeline_format is None: + pipeline_format = "sdxl" if (model_path / "text_encoder_2").is_dir() else "sd" + if pipeline_format not in ("sd", "sdxl"): + raise ValueError("pipeline_format must be sd or sdxl") + + components = [ + ("unet", "UNet2DConditionModel", "model.diffusion_model.", None), + ("vae", "AutoencoderKL", "first_stage_model.", None), + ] + if pipeline_format == "sdxl": + if (model_path / "text_encoder").is_dir(): + components.append(("text_encoder", "CLIPTextModel", "conditioner.embedders.0.transformer.", "clip")) + second_index = 1 + else: + second_index = 0 + components.append( + ( + "text_encoder_2", + "CLIPTextModelWithProjection", + f"conditioner.embedders.{second_index}.model.", + "openclip", + ) + ) + else: + text_config = json.loads((model_path / "text_encoder/config.json").read_text(encoding="utf-8")) + v2 = text_config.get("hidden_size", 768) == 1024 + components.append( + ( + "text_encoder", + "CLIPTextModel", + "cond_stage_model.model." if v2 else "cond_stage_model.transformer.", + "openclip" if v2 else "clip", + ) + ) + + tensors = {} + for directory, cls, prefix, original_format in components: + path = model_path / directory + config = json.loads((path / "config.json").read_text(encoding="utf-8")) + if config.get("quantization_config") is not None: + raise ValueError("Pipeline export requires unpacked, unquantized tensor weights.") + if original_format: + config["original_format"] = original_format + conversion = get_conversion(cls, config) + for key, tensor in conversion.iter_converted(Checkpoint(path), reverse=True): + if dtype is not None and tensor.is_floating_point(): + tensor = tensor.to(dtype) + tensors[prefix + key] = tensor.detach().cpu().contiguous().clone() + output_path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory(prefix=".pipeline-export-", dir=output_path.parent) as temporary: + staging = Path(temporary) / "checkpoint" + if output_format == "safetensors": + save_file(tensors, staging, metadata={"format": "pt"}) + else: + torch.save({"state_dict": tensors}, staging) + staging.rename(output_path) + return output_path diff --git a/src/diffusers/loaders/conversion/pipeline_components.py b/src/diffusers/loaders/conversion/pipeline_components.py new file mode 100644 index 000000000000..7d69620077d6 --- /dev/null +++ b/src/diffusers/loaders/conversion/pipeline_components.py @@ -0,0 +1,93 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reversible definitions for encoders and small components embedded in legacy pipeline converters.""" + +from .clip_vision import clip_vision_conversion +from .core import Conversion, Rule +from .transforms import Reshape, Split + + +def ldm_bert_conversion(config): + mapping = { + "transformer.token_emb.weight": "model.embed_tokens.weight", + "transformer.pos_emb.emb.weight": "model.embed_positions.weight", + } + modules = [("transformer.norm", "model.layer_norm"), ("transformer.to_logits", "to_logits")] + for i in range(config["encoder_layers"]): + source, target = "transformer.attn_layers.layers", f"model.layers.{i}" + modules.extend( + [ + (f"{source}.{2 * i}.0", target + ".self_attn_layer_norm"), + (f"{source}.{2 * i + 1}.0", target + ".final_layer_norm"), + (f"{source}.{2 * i}.1.to_out", target + ".self_attn.out_proj"), + (f"{source}.{2 * i + 1}.1.net.0.0", target + ".fc1"), + (f"{source}.{2 * i + 1}.1.net.2", target + ".fc2"), + ] + ) + for part in ("q", "k", "v"): + mapping[f"{source}.{2 * i}.1.to_{part}.weight"] = f"{target}.self_attn.{part}_proj.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) + + +def paint_by_example_conversion(config): + vision = clip_vision_conversion({**config, "_class_name": "CLIPVisionModel", "original_format": "clip"}) + mapping = {"cond_stage_model.transformer." + old: "model." + new for old, new in vision.mapping.items()} + mapping["learnable_vector"] = "uncond_vector" + modules = [("cond_stage_model.final_ln", "final_layer_norm"), ("proj_out", "proj_out")] + rules = [] + for i in range((config["num_hidden_layers"] + 1) // 5): + old, new = f"cond_stage_model.mapper.resblocks.{i}", f"mapper.blocks.{i}" + modules.extend( + (old + "." + a, new + "." + b) + for a, b in ( + ("attn.c_proj", "attn1.to_out.0"), + ("ln_1", "norm1"), + ("ln_2", "norm3"), + ("mlp.c_fc", "ff.net.0.proj"), + ("mlp.c_proj", "ff.net.2"), + ) + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (old + ".attn.c_qkv." + p,), + tuple(new + f".attn1.to_{part}." + p for part in ("q", "k", "v")), + Split((config["hidden_size"],) * 3), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) + + +def if_safety_checker_conversion(config): + cfg = {**config["vision_config"], "_class_name": "CLIPVisionModelWithProjection", "original_format": "clip"} + base = clip_vision_conversion(cfg) + mapping = {"vision_model." + old: "vision_model." + new for old, new in base.mapping.items()} + rules = [] + for name in ("p_head", "w_head"): + rules.append( + Rule( + (name + ".weights",), + (name + ".weight",), + Reshape((cfg["projection_dim"],), (1, cfg["projection_dim"])), + ) + ) + rules.append(Rule((name + ".biases",), (name + ".bias",), Reshape((), (1,)))) + return Conversion(mapping=mapping, rules=rules) + + +def learned_classifier_free_conversion(config): + return Conversion(mapping={"transformer.empty_text_embed": "embeddings"}) diff --git a/src/diffusers/loaders/conversion/pixart.py b/src/diffusers/loaders/conversion/pixart.py new file mode 100644 index 000000000000..8a8d278599ef --- /dev/null +++ b/src/diffusers/loaders/conversion/pixart.py @@ -0,0 +1,70 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def pixart_conversion(config): + modules = [ + ("x_embedder.proj", "pos_embed.proj"), + ("y_embedder.y_proj.fc1", "caption_projection.linear_1"), + ("y_embedder.y_proj.fc2", "caption_projection.linear_2"), + ("t_block.1", "adaln_single.linear"), + ("final_layer.linear", "proj_out"), + ] + if config["caption_channels"] is None: + modules = [(old, new) for old, new in modules if not new.startswith("caption_projection.")] + mapping = {"final_layer.scale_shift_table": "scale_shift_table"} + rules = [] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + embeddings = [("t_embedder", "timestep_embedder")] + additional = config.get("use_additional_conditions") + if additional is None: + additional = config["sample_size"] == 128 + if additional: + embeddings.extend([("csize_embedder", "resolution_embedder"), ("ar_embedder", "aspect_ratio_embedder")]) + for old, new in embeddings: + modules.extend((f"{old}.mlp.{i}", f"adaln_single.emb.{new}.linear_{j}") for i, j in ((0, 1), (2, 2))) + for i in range(config["num_layers"]): + old, new = f"blocks.{i}", f"transformer_blocks.{i}" + mapping[old + ".scale_shift_table"] = new + ".scale_shift_table" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("attn.proj", "attn1.to_out.0"), + ("cross_attn.q_linear", "attn2.to_q"), + ("cross_attn.proj", "attn2.to_out.0"), + ("mlp.fc1", "ff.net.0.proj"), + ("mlp.fc2", "ff.net.2"), + ) + ) + for p in ("weight", "bias") if config["attention_bias"] else ("weight",): + rules.append( + Rule( + (f"{old}.attn.qkv.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + rules.append( + Rule( + (f"{old}.cross_attn.kv_linear.{p}",), + tuple(f"{new}.attn2.to_{part}.{p}" for part in ("k", "v")), + Split((hidden,) * 2), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/presets/asymmetric-vae-1.5.json b/src/diffusers/loaders/conversion/presets/asymmetric-vae-1.5.json new file mode 100644 index 000000000000..7092cd9ce34e --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/asymmetric-vae-1.5.json @@ -0,0 +1,36 @@ +{ + "_class_name": "AsymmetricAutoencoderKL", + "in_channels": 3, + "out_channels": 3, + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "down_block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "layers_per_down_block": 2, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "up_block_out_channels": [ + 192, + 384, + 768, + 768 + ], + "layers_per_up_block": 3, + "act_fn": "silu", + "latent_channels": 4, + "norm_num_groups": 32, + "sample_size": 256, + "scaling_factor": 0.18215 +} diff --git a/src/diffusers/loaders/conversion/presets/asymmetric-vae-2.json b/src/diffusers/loaders/conversion/presets/asymmetric-vae-2.json new file mode 100644 index 000000000000..057beafc108f --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/asymmetric-vae-2.json @@ -0,0 +1,36 @@ +{ + "_class_name": "AsymmetricAutoencoderKL", + "in_channels": 3, + "out_channels": 3, + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "down_block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "layers_per_down_block": 2, + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "up_block_out_channels": [ + 256, + 512, + 1024, + 1024 + ], + "layers_per_up_block": 5, + "act_fn": "silu", + "latent_channels": 4, + "norm_num_groups": 32, + "sample_size": 256, + "scaling_factor": 0.18215 +} diff --git a/src/diffusers/loaders/conversion/presets/diffuser-rl-128.json b/src/diffusers/loaders/conversion/presets/diffuser-rl-128.json new file mode 100644 index 000000000000..9dea61907660 --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/diffuser-rl-128.json @@ -0,0 +1,32 @@ +{ + "_class_name": "UNet1DModel", + "original_format": "diffuser_rl_legacy", + "down_block_types": [ + "DownResnetBlock1D", + "DownResnetBlock1D", + "DownResnetBlock1D" + ], + "block_out_channels": [ + 32, + 128, + 256 + ], + "up_block_types": [ + "UpResnetBlock1D", + "UpResnetBlock1D" + ], + "layers_per_block": 1, + "use_timestep_embedding": true, + "out_block_type": "OutConv1DBlock", + "norm_num_groups": 8, + "downsample_each_block": false, + "in_channels": 14, + "out_channels": 14, + "extra_in_channels": 0, + "time_embedding_type": "positional", + "flip_sin_to_cos": false, + "freq_shift": 1, + "sample_size": 65536, + "mid_block_type": "MidResTemporalBlock1D", + "act_fn": "mish" +} diff --git a/src/diffusers/loaders/conversion/presets/diffuser-rl-32.json b/src/diffusers/loaders/conversion/presets/diffuser-rl-32.json new file mode 100644 index 000000000000..d7eedd9b3504 --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/diffuser-rl-32.json @@ -0,0 +1,35 @@ +{ + "_class_name": "UNet1DModel", + "original_format": "diffuser_rl_legacy", + "down_block_types": [ + "DownResnetBlock1D", + "DownResnetBlock1D", + "DownResnetBlock1D", + "DownResnetBlock1D" + ], + "block_out_channels": [ + 32, + 64, + 128, + 256 + ], + "up_block_types": [ + "UpResnetBlock1D", + "UpResnetBlock1D", + "UpResnetBlock1D" + ], + "layers_per_block": 1, + "use_timestep_embedding": true, + "out_block_type": "OutConv1DBlock", + "norm_num_groups": 8, + "downsample_each_block": false, + "in_channels": 14, + "out_channels": 14, + "extra_in_channels": 0, + "time_embedding_type": "positional", + "flip_sin_to_cos": false, + "freq_shift": 1, + "sample_size": 65536, + "mid_block_type": "MidResTemporalBlock1D", + "act_fn": "mish" +} diff --git a/src/diffusers/loaders/conversion/presets/diffuser-rl-value.json b/src/diffusers/loaders/conversion/presets/diffuser-rl-value.json new file mode 100644 index 000000000000..dbae8e310d11 --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/diffuser-rl-value.json @@ -0,0 +1,31 @@ +{ + "_class_name": "UNet1DModel", + "original_format": "diffuser_rl_legacy", + "in_channels": 14, + "down_block_types": [ + "DownResnetBlock1D", + "DownResnetBlock1D", + "DownResnetBlock1D", + "DownResnetBlock1D" + ], + "up_block_types": [], + "out_block_type": "ValueFunction", + "mid_block_type": "ValueFunctionMidBlock1D", + "block_out_channels": [ + 32, + 64, + 128, + 256 + ], + "layers_per_block": 1, + "downsample_each_block": true, + "sample_size": 65536, + "out_channels": 14, + "extra_in_channels": 0, + "time_embedding_type": "positional", + "use_timestep_embedding": true, + "flip_sin_to_cos": false, + "freq_shift": 1, + "norm_num_groups": 8, + "act_fn": "mish" +} diff --git a/src/diffusers/loaders/conversion/presets/modelscope-text-to-video.json b/src/diffusers/loaders/conversion/presets/modelscope-text-to-video.json new file mode 100644 index 000000000000..a4fd44774b3c --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/modelscope-text-to-video.json @@ -0,0 +1,3 @@ +{ + "_class_name": "UNet3DConditionModel" +} diff --git a/src/diffusers/loaders/conversion/presets/sd-vae.json b/src/diffusers/loaders/conversion/presets/sd-vae.json new file mode 100644 index 000000000000..404d7bba2140 --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/sd-vae.json @@ -0,0 +1,28 @@ +{ + "_class_name": "AutoencoderKL", + "in_channels": 3, + "out_channels": 3, + "block_out_channels": [ + 128, + 256, + 512, + 512 + ], + "down_block_types": [ + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D" + ], + "up_block_types": [ + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D" + ], + "layers_per_block": 2, + "latent_channels": 4, + "sample_size": 512, + "norm_num_groups": 32, + "scaling_factor": 0.18215 +} diff --git a/src/diffusers/loaders/conversion/presets/tiny-vae.json b/src/diffusers/loaders/conversion/presets/tiny-vae.json new file mode 100644 index 000000000000..a19261de8fc4 --- /dev/null +++ b/src/diffusers/loaders/conversion/presets/tiny-vae.json @@ -0,0 +1,3 @@ +{ + "_class_name": "AutoencoderTiny" +} diff --git a/src/diffusers/loaders/conversion/prior.py b/src/diffusers/loaders/conversion/prior.py new file mode 100644 index 000000000000..ba57ca5ff7d2 --- /dev/null +++ b/src/diffusers/loaders/conversion/prior.py @@ -0,0 +1,105 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch + +from .core import Conversion, Rule +from .transforms import Chain, Permute, Reshape, Split, WithConstants + + +def prior_conversion(config): + original_format = config.get("original_format", "shap_e" if config["encoder_hid_proj_type"] is None else "unclip") + if original_format not in ("shap_e", "unclip", "kandinsky"): + raise ValueError("Prior original_format must be 'unclip', 'kandinsky', or 'shap_e'.") + shape = original_format == "shap_e" + prefix = "wrapped" if shape else "model" + hidden = config["num_attention_heads"] * config["attention_head_dim"] + heads, dim = config["num_attention_heads"], config["attention_head_dim"] + clip_dim = config["clip_embed_dim"] or config["embedding_dim"] + mapping, rules = {}, [] + modules = [ + (f"{prefix}.time_embed.{a}", f"time_embedding.linear_{b}") + for a, b in (("c_fc" if shape else "0", 1), ("c_proj" if shape else "2", 2)) + ] + modules.extend( + [ + (prefix + (".input_proj" if shape else ".clip_img_proj"), "proj_in"), + (prefix + (".ln_post" if shape else ".final_ln"), "norm_out"), + (prefix + (".output_proj" if shape else ".out_proj"), "proj_to_clip_embeddings"), + ] + ) + projection = ".clip_embed.1" if config["embedding_proj_norm_type"] else ".clip_embed" + modules.append((prefix + (projection if shape else ".text_emb_proj"), "embedding_proj")) + if config["embedding_proj_norm_type"]: + modules.append((prefix + ".clip_embed.0", "embedding_proj_norm")) + if config["encoder_hid_proj_type"]: + modules.append((prefix + ".text_enc_proj", "encoder_hidden_states_proj")) + if config["norm_in_type"]: + modules.append((prefix + ".ln_pre", "norm_in")) + if config["added_emb_type"]: + mapping[prefix + ".prd_emb"] = "prd_embedding" + if shape: + tokens = config["num_embeddings"] + config["additional_embeddings"] + rules.append( + Rule((prefix + ".pos_emb",), ("positional_embedding",), Reshape((tokens, hidden), (1, tokens, hidden))) + ) + else: + mapping[prefix + ".positional_embedding"] = "positional_embedding" + rules.extend( + Rule((f"clip_stats.{name}",), (f"clip_{name}",), Reshape((clip_dim,), (1, clip_dim))) + for name in ("mean", "std") + ) + for i in range(config["num_layers"]): + old = f"{prefix}.{'backbone' if shape else 'transformer'}.resblocks.{i}" + new = f"transformer_blocks.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ln_1", "norm1"), + ("ln_2", "norm3"), + ("attn.c_proj", "attn1.to_out.0"), + ("mlp.c_fc", "ff.net.0.proj"), + ("mlp.c_proj", "ff.net.2"), + ) + ) + for p in ("weight", "bias"): + trailing = (hidden,) if p == "weight" else () + transform = Chain( + ( + Reshape((3 * hidden,) + trailing, (heads, 3, dim) + trailing), + Permute((1, 0, 2, 3) if trailing else (1, 0, 2)), + Reshape((3, heads, dim) + trailing, (3 * hidden,) + trailing), + Split((hidden,) * 3), + ) + ) + rules.append( + Rule( + (f"{old}.attn.c_qkv.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + transform, + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + if shape: + anchor = prefix + ".input_proj.weight" + del mapping[anchor] + rules.append( + Rule( + (anchor,), + ("proj_in.weight", "clip_mean", "clip_std"), + WithConstants((torch.zeros(1, clip_dim), torch.zeros(1, clip_dim))), + ) + ) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/prx.py b/src/diffusers/loaders/conversion/prx.py new file mode 100644 index 000000000000..1322aa740851 --- /dev/null +++ b/src/diffusers/loaders/conversion/prx.py @@ -0,0 +1,66 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def prx_conversion(config): + modules = [ + (name, name) + for name in ( + "txt_in", + "time_in.in_layer", + "time_in.out_layer", + "final_layer.linear", + "final_layer.adaLN_modulation.1", + ) + ] + modules.extend( + (name, name) for name in (("img_in.0", "img_in.1") if config["bottleneck_size"] is not None else ("img_in",)) + ) + if config["resolution_embeds"]: + modules.extend( + (f"resolution_embedder.mlp.{name}", f"resolution_embedder.mlp.{name}") + for name in ("in_layer", "out_layer") + ) + mapping = {} + original_format = config.get("original_format", "prx") + if original_format not in ("prx", "prx_weight_norm"): + raise ValueError("PRX original_format must be 'prx' or 'prx_weight_norm'.") + norm_parameter = "weight" if original_format == "prx_weight_norm" else "scale" + for i in range(config["depth"]): + prefix = f"blocks.{i}" + modules.append((prefix + ".modulation.lin", prefix + ".modulation.lin")) + for old, new in ( + ("img_qkv_proj", "attention.img_qkv_proj"), + ("txt_kv_proj", "attention.txt_kv_proj"), + ("attn_out", "attention.to_out.0"), + ("gate_proj", "gate_proj"), + ("up_proj", "up_proj"), + ("down_proj", "down_proj"), + ): + mapping[f"{prefix}.{old}.weight"] = f"{prefix}.{new}.weight" + mapping.update( + { + f"{prefix}.{old}.{norm_parameter}": f"{prefix}.attention.{new}.weight" + for old, new in ( + ("qk_norm.query_norm", "norm_q"), + ("qk_norm.key_norm", "norm_k"), + ("k_norm", "norm_added_k"), + ) + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/qwen3.py b/src/diffusers/loaders/conversion/qwen3.py new file mode 100644 index 000000000000..543683c67d2f --- /dev/null +++ b/src/diffusers/loaders/conversion/qwen3.py @@ -0,0 +1,45 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def qwen3_conversion(config): + causal = config["_class_name"] == "Qwen3ForCausalLM" + prefix = "model." if causal else "" + keys = [prefix + "embed_tokens.weight", prefix + "norm.weight"] + if causal: + keys.append("lm_head.weight") + for i in range(config["num_hidden_layers"]): + block = f"{prefix}layers.{i}" + keys.extend( + f"{block}.{name}.weight" + for name in ( + "input_layernorm", + "post_attention_layernorm", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "self_attn.q_norm", + "self_attn.k_norm", + ) + ) + if config["attention_bias"]: + keys.extend(f"{block}.self_attn.{name}.bias" for name in ("q_proj", "k_proj", "v_proj", "o_proj")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/qwen_image.py b/src/diffusers/loaders/conversion/qwen_image.py new file mode 100644 index 000000000000..87f21dfeb1a5 --- /dev/null +++ b/src/diffusers/loaders/conversion/qwen_image.py @@ -0,0 +1,44 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def qwen_image_conversion(config): + """Validate the published Qwen Image layout, which already uses Diffusers parameter names.""" + modules = [ + "img_in", + "txt_in", + "time_text_embed.timestep_embedder.linear_1", + "time_text_embed.timestep_embedder.linear_2", + "norm_out.linear", + "proj_out", + ] + keys = ["txt_norm.weight"] + if config["use_additional_t_cond"]: + keys.append("time_text_embed.addition_t_embedding.weight") + for i in range(config["num_layers"]): + prefix = f"transformer_blocks.{i}" + modules.extend( + f"{prefix}.{modality}_{name}" + for modality in ("img", "txt") + for name in ("mod.1", "mlp.net.0.proj", "mlp.net.2") + ) + modules.extend( + f"{prefix}.attn.{name}" + for name in ("to_q", "to_k", "to_v", "to_out.0", "add_q_proj", "add_k_proj", "add_v_proj", "to_add_out") + ) + keys.extend(f"{prefix}.attn.{name}.weight" for name in ("norm_q", "norm_k", "norm_added_q", "norm_added_k")) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/qwen_image_vae.py b/src/diffusers/loaders/conversion/qwen_image_vae.py new file mode 100644 index 000000000000..00154a14fd12 --- /dev/null +++ b/src/diffusers/loaders/conversion/qwen_image_vae.py @@ -0,0 +1,23 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .wan_vae import wan_vae_conversion + + +def qwen_image_vae_conversion(config): + # Qwen-Image retains Wan 2.1's parameter layout; its execution and input channels differ. + base = wan_vae_conversion({**config, "is_residual": False, "decoder_base_dim": None}) + return Conversion(mapping=base.mapping) diff --git a/src/diffusers/loaders/conversion/rae.py b/src/diffusers/loaders/conversion/rae.py new file mode 100644 index 000000000000..cdd504cbcddf --- /dev/null +++ b/src/diffusers/loaders/conversion/rae.py @@ -0,0 +1,128 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Reshape + + +def rae_conversion(config): + mapping = { + "stats.latents_mean": "_latents_mean", + "stats.latents_std": "_latents_std", + "decoder.trainable_cls_token": "decoder.trainable_cls_token", + } + modules = [("decoder." + name, "decoder." + name) for name in ("decoder_embed", "decoder_norm", "decoder_pred")] + rules = tuple( + Rule((f"processor.image_{name}",), (f"encoder_{name}",), Reshape((3,), (1, 3, 1, 1))) + for name in ("mean", "std") + ) + for i in range(config["decoder_num_hidden_layers"]): + prefix = f"decoder.decoder_layers.{i}" + modules.extend( + (f"{prefix}.{a}", f"{prefix}.{b}") + for a, b in ( + ("attention.attention.query", "attention.to_q"), + ("attention.attention.key", "attention.to_k"), + ("attention.attention.value", "attention.to_v"), + ("attention.output.dense", "attention.to_out.0"), + ("intermediate.dense", "intermediate.dense"), + ("output.dense", "output.dense"), + ("layernorm_before", "layernorm_before"), + ("layernorm_after", "layernorm_after"), + ) + ) + encoder_type = config["encoder_type"] + modern = int(config.get("transformers_version", "4").split(".")[0]) >= 5 + if encoder_type in ("dinov2", "mae"): + tokens = ["cls_token", "position_embeddings"] + if encoder_type == "dinov2": + tokens.extend(["mask_token", "register_tokens"]) + mapping.update({f"encoder.embeddings.{name}": f"encoder.embeddings.{name}" for name in tokens}) + name = "encoder.embeddings.patch_embeddings.projection" + modules.append((name, name)) + for i in range(config["encoder_num_hidden_layers"]): + if encoder_type == "dinov2": + prefix = f"encoder.encoder.layer.{i}" + names = ( + "norm1", + "norm2", + "attention.attention.query", + "attention.attention.key", + "attention.attention.value", + "attention.output.dense", + "mlp.fc1", + "mlp.fc2", + ) + for j in (1, 2): + key = f"{prefix}.layer_scale{j}.lambda1" + mapping[key] = key + elif modern: + prefix = f"encoder.layers.{i}" + names = ( + "attention.q_proj", + "attention.k_proj", + "attention.v_proj", + "attention.o_proj", + "layernorm_before", + "layernorm_after", + "mlp.fc1", + "mlp.fc2", + ) + else: + prefix = f"encoder.encoder.layer.{i}" + names = ( + "attention.attention.query", + "attention.attention.key", + "attention.attention.value", + "attention.output.dense", + "layernorm_before", + "layernorm_after", + "intermediate.dense", + "output.dense", + ) + modules.extend((f"{prefix}.{name}", f"{prefix}.{name}") for name in names) + elif encoder_type == "siglip2": + prefix = "encoder." if modern else "encoder.vision_model." + name = prefix + "embeddings.position_embedding.weight" + mapping[name] = name + modules.append((prefix + "embeddings.patch_embedding", prefix + "embeddings.patch_embedding")) + for i in range(config["encoder_num_hidden_layers"]): + modules.extend( + (f"{prefix}encoder.layers.{i}.{name}", f"{prefix}encoder.layers.{i}.{name}") + for name in ( + "layer_norm1", + "layer_norm2", + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.out_proj", + "mlp.fc1", + "mlp.fc2", + ) + ) + mapping.update( + { + prefix + "head." + name: prefix + "head." + name + for name in ("probe", "attention.in_proj_weight", "attention.in_proj_bias") + } + ) + modules.extend( + (prefix + "head." + name, prefix + "head." + name) + for name in ("attention.out_proj", "layernorm", "mlp.fc1", "mlp.fc2") + ) + else: + raise ValueError(f"Unknown RAE encoder type {encoder_type}.") + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/registry.py b/src/diffusers/loaders/conversion/registry.py new file mode 100644 index 000000000000..d0338481819f --- /dev/null +++ b/src/diffusers/loaders/conversion/registry.py @@ -0,0 +1,433 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Resolve a component's reversible conversion from its class and Diffusers configuration.""" + +import importlib +import inspect + +from .ace_step import ace_step_conversion +from .ace_step_conditioner import ace_step_conditioner_conversion +from .ace_step_detokenizer import ace_step_detokenizer_conversion +from .ace_step_tokenizer import ace_step_tokenizer_conversion +from .anima_conditioner import anima_conditioner_conversion +from .animatediff import animatediff_conversion +from .anyflow import anyflow_conversion +from .anyflow_far import anyflow_far_conversion +from .asymmetric_vae import asymmetric_vae_conversion +from .audioldm2_projection import audioldm2_projection_conversion +from .audioldm2_unet import audioldm2_unet_conversion +from .auraflow import auraflow_conversion +from .autoencoder_dc import autoencoder_dc_conversion +from .blip_qformer import blip_qformer_conversion +from .chroma import chroma_conversion +from .chronoedit import chronoedit_conversion +from .clap_audio import clap_audio_conversion +from .clap_text import clap_text_conversion +from .clip import clip_conversion +from .clip_vision import clip_vision_conversion +from .cogvideox import cogvideox_transformer_conversion, cogvideox_vae_conversion +from .cogview3plus import cogview3plus_conversion +from .cogview4 import cogview4_conversion +from .consistency_decoder import consistency_decoder_conversion +from .controlnet import controlnet_conversion +from .cosmos import cosmos_conversion +from .cosmos_controlnet import cosmos_controlnet_conversion +from .cosmos_vae import cosmos_vae_conversion +from .dit import dit_conversion +from .ernie_image import ernie_image_conversion +from .flux import flux_conversion +from .flux2 import flux2_conversion +from .flux2_vae import flux2_vae_conversion +from .flux_ip_adapter import flux_ip_adapter_conversion +from .hidream import hidream_conversion +from .hifigan import hifigan_conversion +from .hunyuan_dit import hunyuan_dit_conversion +from .hunyuan_dit_controlnet import hunyuan_dit_controlnet_conversion +from .hunyuan_image import hunyuan_image_conversion +from .hunyuan_image_refiner_vae import hunyuan_image_refiner_vae_conversion +from .hunyuan_image_vae import hunyuan_image_vae_conversion +from .hunyuan_video import hunyuan_video_conversion +from .hunyuan_video15 import hunyuan_video15_conversion +from .hunyuan_video15_vae import hunyuan_video15_vae_conversion +from .hunyuan_video_vae import hunyuan_video_vae_conversion +from .i2vgen_xl import i2vgen_xl_conversion +from .joy_image import joy_image_conversion +from .kandinsky3 import kandinsky3_conversion +from .ldm_unet import ldm_unet_conversion +from .ldm_vae import ldm_vae_conversion +from .longcat_audio import longcat_audio_conversion +from .longcat_audio_vae import longcat_audio_vae_conversion +from .lora import lora_conversion +from .ltx import ltx_conversion +from .ltx2 import ltx2_conversion +from .ltx2_audio_vae import ltx2_audio_vae_conversion +from .ltx2_connectors import ltx2_connectors_conversion +from .ltx2_diffusion_decoder import ltx2_diffusion_decoder_conversion +from .ltx2_duration import ltx2_duration_conversion +from .ltx2_upsampler import ltx2_upsampler_conversion +from .ltx2_vae import ltx2_vae_conversion +from .ltx2_vocoder import ltx2_vocoder_conversion +from .ltx_upsampler import ltx_upsampler_conversion +from .ltx_vae import ltx_vae_conversion +from .lumina import lumina_conversion +from .lumina2 import lumina2_conversion +from .minimax_h3 import minimax_h3_conversion +from .minimax_h3_audio_vae import minimax_h3_audio_vae_conversion +from .minimax_h3_vae import minimax_h3_vae_conversion +from .minimax_music3 import minimax_music3_conversion +from .minimax_music3_conditioner import minimax_music3_conditioner_conversion +from .minimax_music3_rvq import minimax_music3_rvq_conversion +from .minimax_music3_vocoder import minimax_music3_vocoder_conversion +from .mochi import mochi_conversion +from .mochi_vae import mochi_vae_conversion +from .motif_video import motif_video_conversion +from .omnigen import omnigen_conversion +from .oobleck import oobleck_conversion +from .ovis_image import ovis_image_conversion +from .paella import paella_conversion +from .pipeline_components import ( + if_safety_checker_conversion, + ldm_bert_conversion, + learned_classifier_free_conversion, + paint_by_example_conversion, +) +from .pixart import pixart_conversion +from .prior import prior_conversion +from .prx import prx_conversion +from .qwen3 import qwen3_conversion +from .qwen_image import qwen_image_conversion +from .qwen_image_vae import qwen_image_vae_conversion +from .rae import rae_conversion +from .same import same_conversion +from .sana import sana_conversion +from .sana_controlnet import sana_controlnet_conversion +from .sana_video import sana_video_conversion +from .sd3 import sd3_conversion +from .sd3_controlnet import sd3_controlnet_conversion +from .shap_e_renderer import shap_e_renderer_conversion +from .skyreels_v2 import skyreels_v2_conversion +from .sparse_controlnet import sparse_controlnet_conversion +from .spectrogram_continuous import spectrogram_continuous_conversion +from .spectrogram_notes import spectrogram_notes_conversion +from .stable_audio import stable_audio_conversion +from .stable_audio3 import stable_audio3_conversion +from .stable_audio3_duration import stable_audio3_duration_conversion +from .stable_audio_projection import stable_audio_projection_conversion +from .stable_cascade import stable_cascade_conversion +from .svd import svd_conversion +from .svd_vae import svd_vae_conversion +from .t2i_adapter import t2i_adapter_conversion +from .t5 import t5_conversion +from .t5_film import t5_film_conversion +from .tiny_vae import tiny_vae_conversion +from .umt5 import umt5_conversion +from .unclip_text_projection import unclip_text_projection_conversion +from .unet_1d import unet_1d_conversion +from .unet_2d import unet_2d_conversion +from .unet_3d import unet_3d_conversion +from .unidiffuser import unidiffuser_conversion +from .unidiffuser_text import unidiffuser_text_conversion +from .uvit import uvit_conversion +from .versatile_text_unet import versatile_text_unet_conversion +from .vq_diffusion import vq_diffusion_conversion +from .vq_model import vq_model_conversion +from .wan import wan_conversion +from .wan_animate import wan_animate_conversion +from .wan_animate2 import wan_animate2_conversion +from .wan_vace import wan_vace_conversion +from .wan_vae import wan_vae_conversion +from .wuerstchen_decoder import wuerstchen_decoder_conversion +from .wuerstchen_prior import wuerstchen_prior_conversion +from .z_image import z_image_conversion +from .z_image_controlnet import z_image_controlnet_conversion +from .zero123_projection import zero123_projection_conversion + + +CONVERSION_BUILDERS = { + "LDMBertModel": ldm_bert_conversion, + "PaintByExampleImageEncoder": paint_by_example_conversion, + "IFSafetyChecker": if_safety_checker_conversion, + "LearnedClassifierFreeSamplingEmbeddings": learned_classifier_free_conversion, + "CogVideoXTransformer3DModel": cogvideox_transformer_conversion, + "AutoencoderKLCogVideoX": cogvideox_vae_conversion, + "ChronoEditTransformer3DModel": chronoedit_conversion, + "QwenImageTransformer2DModel": qwen_image_conversion, + "MotifVideoTransformer3DModel": motif_video_conversion, + "LoRA": lora_conversion, + "LTXLatentUpsamplerModel": ltx_upsampler_conversion, + "SkyReelsV2Transformer3DModel": skyreels_v2_conversion, + "ClapModel": clap_audio_conversion, + "ClapAudioModel": clap_audio_conversion, + "ClapAudioModelWithProjection": clap_audio_conversion, + "ConsistencyDecoderVAE": consistency_decoder_conversion, + "FluxIPAdapter": flux_ip_adapter_conversion, + "ShapERenderer": shap_e_renderer_conversion, + "JoyImageEditTransformer3DModel": joy_image_conversion, + "JoyImageEditPlusTransformer3DModel": joy_image_conversion, + "CLIPVisionModel": clip_vision_conversion, + "CLIPVisionModelWithProjection": clip_vision_conversion, + "Transformer2DModel": vq_diffusion_conversion, + "UniDiffuserModel": unidiffuser_conversion, + "UniDiffuserTextDecoder": unidiffuser_text_conversion, + "CCProjection": zero123_projection_conversion, + "LTX2DurationHead": ltx2_duration_conversion, + "LTX2LatentUpsamplerModel": ltx2_upsampler_conversion, + "LTX2TextConnectors": ltx2_connectors_conversion, + "LTX2Vocoder": ltx2_vocoder_conversion, + "LTX2VocoderWithBWE": ltx2_vocoder_conversion, + "LTX2VideoDiffusionDecoderModel": ltx2_diffusion_decoder_conversion, + "WuerstchenPrior": wuerstchen_prior_conversion, + "WuerstchenDiffNeXt": wuerstchen_decoder_conversion, + "PaellaVQModel": paella_conversion, + "SparseControlNetModel": sparse_controlnet_conversion, + "T2IAdapter": t2i_adapter_conversion, + "AsymmetricAutoencoderKL": asymmetric_vae_conversion, + "Blip2QFormerModel": blip_qformer_conversion, + "ContextCLIPTextModel": clip_conversion, + "AutoencoderKLFlux2": flux2_vae_conversion, + "UNetFlatConditionModel": versatile_text_unet_conversion, + "UNetSpatioTemporalConditionModel": svd_conversion, + "AutoencoderKLTemporalDecoder": svd_vae_conversion, + "I2VGenXLUNet": i2vgen_xl_conversion, + "UNet3DConditionModel": unet_3d_conversion, + "UNet1DModel": unet_1d_conversion, + "Kandinsky3UNet": kandinsky3_conversion, + "ClapTextModelWithProjection": clap_text_conversion, + "SpeechT5HifiGan": hifigan_conversion, + "AudioLDM2UNet2DConditionModel": audioldm2_unet_conversion, + "AudioLDM2ProjectionModel": audioldm2_projection_conversion, + "UnCLIPTextProjModel": unclip_text_projection_conversion, + "PriorTransformer": prior_conversion, + "UNet2DModel": unet_2d_conversion, + "SD3ControlNetModel": sd3_controlnet_conversion, + "SanaVideoTransformer3DModel": sana_video_conversion, + "SanaControlNetModel": sana_controlnet_conversion, + "CosmosControlNetModel": cosmos_controlnet_conversion, + "AutoencoderRAE": rae_conversion, + "AutoencoderKLMiniMaxH3Audio": minimax_h3_audio_vae_conversion, + "AutoencoderKLHunyuanImage": hunyuan_image_vae_conversion, + "AutoencoderKLHunyuanImageRefiner": hunyuan_image_refiner_vae_conversion, + "AutoencoderKLMiniMaxH3": minimax_h3_vae_conversion, + "HunyuanImageTransformer2DModel": hunyuan_image_conversion, + "MiniMaxH3Transformer3DModel": minimax_h3_conversion, + "PRXTransformer2DModel": prx_conversion, + "HunyuanVideo15Transformer3DModel": hunyuan_video15_conversion, + "AutoencoderKLHunyuanVideo15": hunyuan_video15_vae_conversion, + "AutoencoderKLCosmos": cosmos_vae_conversion, + "AutoencoderKLHunyuanVideo": hunyuan_video_vae_conversion, + "AutoencoderKLMochi": mochi_vae_conversion, + "AutoencoderTiny": tiny_vae_conversion, + "SpectrogramNotesEncoder": spectrogram_notes_conversion, + "SpectrogramContEncoder": spectrogram_continuous_conversion, + "T5FilmDecoder": t5_film_conversion, + "Qwen3Model": qwen3_conversion, + "Qwen3ForCausalLM": qwen3_conversion, + "UMT5EncoderModel": umt5_conversion, + "UVit2DModel": uvit_conversion, + "VQModel": vq_model_conversion, + "AceStepTransformer1DModel": ace_step_conversion, + "AceStepConditionEncoder": ace_step_conditioner_conversion, + "AceStepAudioTokenizer": ace_step_tokenizer_conversion, + "AceStepAudioTokenDetokenizer": ace_step_detokenizer_conversion, + "LongCatAudioDiTTransformer": longcat_audio_conversion, + "LongCatAudioDiTVae": longcat_audio_vae_conversion, + "AutoencoderSAME": same_conversion, + "StableAudio3DiTModel": stable_audio3_conversion, + "StableAudio3DurationEmbedder": stable_audio3_duration_conversion, + "AutoencoderOobleck": oobleck_conversion, + "StableAudioProjectionModel": stable_audio_projection_conversion, + "MiniMaxMusic3ConditionEncoder": minimax_music3_conditioner_conversion, + "MiniMaxMusic3Vocoder": minimax_music3_vocoder_conversion, + "MiniMaxMusic3RVQDepthDecoder": minimax_music3_rvq_conversion, + "StableAudioDiTModel": stable_audio_conversion, + "MiniMaxMusic3Transformer1DModel": minimax_music3_conversion, + "OvisImageTransformer2DModel": ovis_image_conversion, + "AnyFlowTransformer3DModel": anyflow_conversion, + "AnyFlowFARTransformer3DModel": anyflow_far_conversion, + "AutoencoderKLQwenImage": qwen_image_vae_conversion, + "AnimaTextConditioner": anima_conditioner_conversion, + "OmniGenTransformer2DModel": omnigen_conversion, + "LuminaNextDiT2DModel": lumina_conversion, + "HunyuanDiT2DModel": hunyuan_dit_conversion, + "HunyuanDiT2DControlNetModel": hunyuan_dit_controlnet_conversion, + "DiTTransformer2DModel": dit_conversion, + "PixArtTransformer2DModel": pixart_conversion, + "CogView3PlusTransformer2DModel": cogview3plus_conversion, + "CogView4Transformer2DModel": cogview4_conversion, + "WanAnimateTransformer3DModel": wan_animate_conversion, + "WanVACETransformer3DModel": wan_vace_conversion, + "WanAnimate2Transformer3DModel": wan_animate2_conversion, + "ZImageControlNetModel": z_image_controlnet_conversion, + "AutoencoderKLLTX2Audio": ltx2_audio_vae_conversion, + "CLIPTextModel": clip_conversion, + "CLIPTextModelWithProjection": clip_conversion, + "T5EncoderModel": t5_conversion, + "AutoencoderDC": autoencoder_dc_conversion, + "StableCascadeUNet": stable_cascade_conversion, + "ErnieImageTransformer2DModel": ernie_image_conversion, + "HiDreamImageTransformer2DModel": hidream_conversion, + "MotionAdapter": animatediff_conversion, + "UNet2DConditionModel": ldm_unet_conversion, + "ControlNetModel": controlnet_conversion, + "AutoencoderKL": ldm_vae_conversion, + "AutoencoderKLWan": wan_vae_conversion, + "AutoencoderKLLTXVideo": ltx_vae_conversion, + "AutoencoderKLLTX2Video": ltx2_vae_conversion, + "LTX2VideoTransformer3DModel": ltx2_conversion, + "CosmosTransformer3DModel": cosmos_conversion, + "LTXVideoTransformer3DModel": ltx_conversion, + "WanTransformer3DModel": wan_conversion, + "ZImageTransformer2DModel": z_image_conversion, + "AuraFlowTransformer2DModel": auraflow_conversion, + "Lumina2Transformer2DModel": lumina2_conversion, + "HunyuanVideoTransformer3DModel": hunyuan_video_conversion, + "SanaTransformer2DModel": sana_conversion, + "FluxTransformer2DModel": flux_conversion, + "Flux2Transformer2DModel": flux2_conversion, + "ChromaTransformer2DModel": chroma_conversion, + "SD3Transformer2DModel": sd3_conversion, + "MochiTransformer3DModel": mochi_conversion, +} + + +CONVERSION_FORMATS = { + "ConsistencyDecoderVAE": ("consistency_decoder", "consistency_decoder_jit"), + "CLIPTextModel": ("clip", "openclip"), + "CLIPTextModelWithProjection": ("clip", "openclip"), + "ContextCLIPTextModel": ("clip", "openclip"), + "CLIPVisionModel": ("clip", "openclip"), + "CLIPVisionModelWithProjection": ("clip", "openclip"), + "CosmosTransformer3DModel": ("cosmos1", "cosmos2"), + "CogView4Transformer2DModel": ("cogview4", "megatron"), + "UNet1DModel": ("diffuser_rl", "diffuser_rl_legacy"), + "UNet2DModel": ("ddpm", "ldm", "consistency", "ncsnpp"), + "UNet2DConditionModel": ("ldm", "versatile_image"), + "HunyuanImageTransformer2DModel": ("hunyuan_image_fused", "hunyuan_image_split"), + "LoRA": ("kohya", "diffusers", "diffusers_old", "peft", "animatediff"), + "LTX2VideoDiffusionDecoderModel": ("ltx2_diffusion_decoder", "ltx2_diffusion_decoder_gated"), + "MiniMaxH3Transformer3DModel": ("minimax_h3", "minimax_h3_shards"), + "PriorTransformer": ("shap_e", "unclip", "kandinsky"), + "PRXTransformer2DModel": ("prx", "prx_weight_norm"), + "AutoencoderKLTemporalDecoder": ("svd", "temporal_vae"), + "AutoencoderKLHunyuanImage": ("hunyuan_image_vae", "hunyuan_image_vae_2d"), +} + + +def get_conversion(model_class: str, config: dict): + """Build a component conversion using its Diffusers config, without constructing a model. + + Each definition returns `Conversion(mapping=..., rules=...)`. Constructor defaults fill omitted config fields; no + original checkpoint or previous import is required. File wrappers, prefixes, and auxiliary state are separate. + """ + original_format = config.get("original_format") + if original_format is not None and original_format not in CONVERSION_FORMATS.get(model_class, ()): + raise ValueError(f"Unsupported original_format {original_format!r} for {model_class}.") + if model_class == "LoRA": + return lora_conversion(config) + if model_class == "FluxIPAdapter": + return flux_ip_adapter_conversion(config) + if model_class == "CCProjection": + return zero123_projection_conversion(config) + if model_class == "Transformer2DModel": + if config.get("norm_type") == "ada_norm_zero": + model_class = "DiTTransformer2DModel" + elif config.get("norm_type") == "ada_norm_single": + model_class = "PixArtTransformer2DModel" + if model_class == "CogVideoXTransformer3DModel": + return cogvideox_transformer_conversion(config) + if model_class == "AutoencoderKLCogVideoX": + return cogvideox_vae_conversion(config) + if model_class not in CONVERSION_BUILDERS: + raise ValueError(f"No reversible conversion registered for {model_class}.") + if model_class in ("LDMBertModel", "PaintByExampleImageEncoder", "IFSafetyChecker"): + if model_class == "LDMBertModel": + module = importlib.import_module("diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion") + resolved_config = module.LDMBertConfig(**config).to_dict() + else: + module = importlib.import_module("transformers") + config_class = ( + module.CLIPVisionConfig if model_class == "PaintByExampleImageEncoder" else module.CLIPConfig + ) + resolved_config = config_class(**config).to_dict() + resolved_config["transformers_version"] = module.__version__ + resolved_config.update(config) + return CONVERSION_BUILDERS[model_class](resolved_config) + if model_class in ( + "CLIPVisionModel", + "CLIPVisionModelWithProjection", + "CLIPTextModel", + "CLIPTextModelWithProjection", + "ContextCLIPTextModel", + "Blip2QFormerModel", + "T5EncoderModel", + "UMT5EncoderModel", + "Qwen3Model", + "Qwen3ForCausalLM", + "ClapModel", + "ClapAudioModel", + "ClapAudioModelWithProjection", + "ClapTextModelWithProjection", + "SpeechT5HifiGan", + ): + transformers = importlib.import_module("transformers") + config_name = { + "CLIPVisionModel": "CLIPVisionConfig", + "CLIPVisionModelWithProjection": "CLIPVisionConfig", + "T5EncoderModel": "T5Config", + "Blip2QFormerModel": "Blip2Config", + "UMT5EncoderModel": "UMT5Config", + "Qwen3Model": "Qwen3Config", + "Qwen3ForCausalLM": "Qwen3Config", + "ClapModel": "ClapConfig", + "ClapAudioModel": "ClapAudioConfig", + "ClapAudioModelWithProjection": "ClapAudioConfig", + "ClapTextModelWithProjection": "ClapTextConfig", + "SpeechT5HifiGan": "SpeechT5HifiGanConfig", + }.get(model_class, "CLIPTextConfig") + config_class = getattr(transformers, config_name) + resolved_config = config_class(**config).to_dict() + resolved_config.update(config) + resolved_config["_class_name"] = model_class + return CONVERSION_BUILDERS[model_class](resolved_config) + internal_modules = { + "LearnedClassifierFreeSamplingEmbeddings": "diffusers.pipelines.deprecated.vq_diffusion.pipeline_vq_diffusion", + "LTXLatentUpsamplerModel": "diffusers.pipelines.ltx.modeling_latent_upsampler", + "ShapERenderer": "diffusers.pipelines.shap_e.renderer", + "UniDiffuserModel": "diffusers.pipelines.deprecated.unidiffuser.modeling_uvit", + "UniDiffuserTextDecoder": "diffusers.pipelines.deprecated.unidiffuser.modeling_text_decoder", + "LTX2TextConnectors": "diffusers.pipelines.ltx2.connectors", + "LTX2Vocoder": "diffusers.pipelines.ltx2.vocoder", + "LTX2VocoderWithBWE": "diffusers.pipelines.ltx2.vocoder", + "LTX2DurationHead": "diffusers.pipelines.ltx2.duration_head", + "LTX2LatentUpsamplerModel": "diffusers.pipelines.ltx2.latent_upsampler", + "WuerstchenPrior": "diffusers.pipelines.deprecated.wuerstchen.modeling_wuerstchen_prior", + "WuerstchenDiffNeXt": "diffusers.pipelines.deprecated.wuerstchen.modeling_wuerstchen_diffnext", + "PaellaVQModel": "diffusers.pipelines.deprecated.wuerstchen.modeling_paella_vq_model", + "UNetFlatConditionModel": "diffusers.pipelines.deprecated.versatile_diffusion.modeling_text_unet", + "UnCLIPTextProjModel": "diffusers.pipelines.deprecated.unclip.text_proj", + "SpectrogramNotesEncoder": "diffusers.pipelines.deprecated.spectrogram_diffusion.notes_encoder", + "SpectrogramContEncoder": "diffusers.pipelines.deprecated.spectrogram_diffusion.continuous_encoder", + } + cls = getattr(importlib.import_module(internal_modules.get(model_class, "diffusers")), model_class) + resolved_config = { + name: parameter.default + for name, parameter in inspect.signature(cls.__init__).parameters.items() + if parameter.default is not inspect.Parameter.empty + } + resolved_config.update(config) + if model_class == "AutoencoderRAE" and not resolved_config.get("transformers_version"): + resolved_config["transformers_version"] = importlib.import_module("transformers").__version__ + resolved_config["_class_name"] = model_class + return CONVERSION_BUILDERS[model_class](resolved_config) diff --git a/src/diffusers/loaders/conversion/same.py b/src/diffusers/loaders/conversion/same.py new file mode 100644 index 000000000000..8055ebaf8d66 --- /dev/null +++ b/src/diffusers/loaders/conversion/same.py @@ -0,0 +1,56 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import ReorderChunks + + +def same_conversion(config): + mapping = { + "bottleneck.scaling_factor": "bottleneck.scale", + "bottleneck.bias": "bottleneck.bias", + "bottleneck.running_std": "bottleneck.running_std", + } + count = len(config["encoder_strides"]) + modules = [(f"encoder.layers.{count + 1}", "encoder.proj"), ("decoder.layers.1", "decoder.proj")] + channels = [config["audio_channels"] * config["patch_size"]] + [ + config["encoder_channels"] * c for c in config["encoder_c_mults"] + ] + rules = [] + for component in ("encoder", "decoder"): + for i in range(count): + level = i if component == "encoder" else count - 1 - i + old, new = f"{component}.layers.{i if component == 'encoder' else i + 3}", f"{component}.blocks.{i}" + mapping[old + ".new_tokens"] = new + ".new_tokens" + if channels[level] != channels[level + 1]: + mapping.update({f"{old}.mapping.{p}": f"{new}.mapping.{p}" for p in ("weight_g", "weight_v", "bias")}) + for j in range(config["encoder_transformer_depths"][level]): + a, b = f"{old}.transformers.{j}", f"{new}.transformers.{j}" + for source, target in ( + ("pre_norm", "norm_attn"), + ("ff_norm", "norm_ff"), + ("self_attn.q_norm", "attn.q_norm"), + ("self_attn.k_norm", "attn.k_norm"), + ): + mapping.update({f"{a}.{source}.{p}": f"{b}.{target}.{p}" for p in ("alpha", "gamma", "beta")}) + source, target = a + ".self_attn.to_qkv.weight", b + ".attn.to_qkv.weight" + if config["use_differential_attention"]: + rules.append(Rule((source,), (target,), ReorderChunks((0, 3, 1, 4, 2)))) + else: + mapping[source] = target + mapping[a + ".self_attn.to_out.weight"] = b + ".attn.to_out.weight" + modules.extend([(a + ".ff.ff.0.proj", b + ".ff.proj_in"), (a + ".ff.ff.2", b + ".ff.proj_out")]) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/sana.py b/src/diffusers/loaders/conversion/sana.py new file mode 100644 index 000000000000..69b716210e77 --- /dev/null +++ b/src/diffusers/loaders/conversion/sana.py @@ -0,0 +1,90 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import Split + + +def sana_conversion(config): + mapping = {"attention_y_norm.weight": "caption_norm.weight", "final_layer.scale_shift_table": "scale_shift_table"} + modules = [ + ("x_embedder.proj", "patch_embed.proj"), + ("t_block.1", "time_embed.linear"), + ("y_embedder.y_proj.fc1", "caption_projection.linear_1"), + ("y_embedder.y_proj.fc2", "caption_projection.linear_2"), + ("final_layer.linear", "proj_out"), + ] + timed = "time_embed" if config.get("guidance_embeds", False) else "time_embed.emb" + modules.extend((f"t_embedder.mlp.{i}", f"{timed}.timestep_embedder.linear_{j}") for i, j in ((0, 1), (2, 2))) + if config.get("guidance_embeds", False): + modules.extend( + (f"cfg_embedder.mlp.{i}", f"time_embed.guidance_embedder.linear_{j}") for i, j in ((0, 1), (2, 2)) + ) + rules = [] + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + has_cross = config["cross_attention_dim"] is not None + cross_size = config["num_cross_attention_heads"] * config["cross_attention_head_dim"] if has_cross else None + for i in range(config["num_layers"]): + old, new = f"blocks.{i}", f"transformer_blocks.{i}" + mapping[old + ".scale_shift_table"] = new + ".scale_shift_table" + for p in ("weight", "bias") if config["attention_bias"] else ("weight",): + rules.append( + Rule( + (f"{old}.attn.qkv.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden_size,) * 3), + ) + ) + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("attn.proj", "attn1.to_out.0"), + ("mlp.inverted_conv.conv", "ff.conv_inverted"), + ("mlp.depth_conv.conv", "ff.conv_depth"), + ) + ) + mapping[old + ".mlp.point_conv.conv.weight"] = new + ".ff.conv_point.weight" + if has_cross: + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.cross_attn.kv_linear.{p}",), + tuple(f"{new}.attn2.to_{part}.{p}" for part in ("k", "v")), + Split((cross_size,) * 2), + ) + ) + modules.extend( + [ + (old + ".cross_attn.q_linear", new + ".attn2.to_q"), + (old + ".cross_attn.proj", new + ".attn2.to_out.0"), + ] + ) + if config["norm_elementwise_affine"]: + modules.append((old + ".norm2", new + ".norm2")) + if config.get("qk_norm") not in (None, "l2"): + norm_parameters = ( + ("weight", "bias") + if config["qk_norm"] in ("layer_norm", "fp32_layer_norm", "layer_norm_across_heads") + else ("weight",) + ) + for source, target in [("attn", "attn1")] + ([("cross_attn", "attn2")] if has_cross else []): + mapping.update( + { + f"{old}.{source}.{part}_norm.{p}": f"{new}.{target}.norm_{part}.{p}" + for part in ("q", "k") + for p in norm_parameters + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/sana_controlnet.py b/src/diffusers/loaders/conversion/sana_controlnet.py new file mode 100644 index 000000000000..8b86f22675c0 --- /dev/null +++ b/src/diffusers/loaders/conversion/sana_controlnet.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .sana import sana_conversion + + +def sana_controlnet_conversion(config): + base = sana_conversion(config) + + def control_key(key): + if not key.startswith("blocks."): + return key + _, index, name = key.split(".", 2) + return f"controlnet.{index}.copied_block.{name}" + + mapping = {control_key(old): new for old, new in base.mapping.items() if not old.startswith("final_layer.")} + rules = tuple( + Rule(tuple(control_key(key) for key in rule.original), rule.diffusers, rule.transform) for rule in base.rules + ) + for p in ("weight", "bias"): + mapping[f"controlnet.0.before_proj.{p}"] = f"input_block.{p}" + mapping.update( + {f"controlnet.{i}.after_proj.{p}": f"controlnet_blocks.{i}.{p}" for i in range(config["num_layers"])} + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/sana_video.py b/src/diffusers/loaders/conversion/sana_video.py new file mode 100644 index 000000000000..6b2960e90ad8 --- /dev/null +++ b/src/diffusers/loaders/conversion/sana_video.py @@ -0,0 +1,28 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .sana import sana_conversion + + +def sana_video_conversion(config): + base = sana_conversion(config) + mapping = {old: new.replace("patch_embed.proj.", "patch_embedding.") for old, new in base.mapping.items()} + mapping.update( + { + f"blocks.{i}.mlp.t_conv.weight": f"transformer_blocks.{i}.ff.conv_temp.weight" + for i in range(config["num_layers"]) + } + ) + return Conversion(mapping=mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/sd3.py b/src/diffusers/loaders/conversion/sd3.py new file mode 100644 index 000000000000..a0464361ac56 --- /dev/null +++ b/src/diffusers/loaders/conversion/sd3.py @@ -0,0 +1,89 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import ReorderChunks, Split + + +def sd3_conversion(config): + hidden_size = config["num_attention_heads"] * config["attention_head_dim"] + mapping = {"pos_embed": "pos_embed.pos_embed"} + modules = [ + ("x_embedder.proj", "pos_embed.proj"), + ("context_embedder", "context_embedder"), + ("final_layer.linear", "proj_out"), + ] + rules = [] + for old, new in (("t_embedder", "timestep_embedder"), ("y_embedder", "text_embedder")): + modules.extend( + [ + (old + ".mlp.0", "time_text_embed." + new + ".linear_1"), + (old + ".mlp.2", "time_text_embed." + new + ".linear_2"), + ] + ) + for i in range(config["num_layers"]): + old, new = f"joint_blocks.{i}", f"transformer_blocks.{i}" + branches = [ + ("x_block.attn", "attn", ("to_q", "to_k", "to_v"), "to_out.0", ("norm_q", "norm_k")), + ( + "context_block.attn", + "attn", + ("add_q_proj", "add_k_proj", "add_v_proj"), + "to_add_out", + ("norm_added_q", "norm_added_k"), + ), + ] + if i in config.get("dual_attention_layers", ()): + branches.append(("x_block.attn2", "attn2", ("to_q", "to_k", "to_v"), "to_out.0", ("norm_q", "norm_k"))) + for branch, attn, targets, output, norms in branches: + for parameter in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.{branch}.qkv.{parameter}",), + tuple(f"{new}.{attn}.{name}.{parameter}" for name in targets), + Split((hidden_size,) * 3), + ) + ) + if branch != "context_block.attn" or i != config["num_layers"] - 1: + modules.append((f"{old}.{branch}.proj", f"{new}.{attn}.{output}")) + if config.get("qk_norm") is not None: + norm_parameters = ( + ("weight", "bias") if config["qk_norm"] in ("layer_norm", "fp32_layer_norm") else ("weight",) + ) + for source, target in zip(("ln_q", "ln_k"), norms): + mapping.update( + {f"{old}.{branch}.{source}.{p}": f"{new}.{attn}.{target}.{p}" for p in norm_parameters} + ) + for branch, norm, ff in (("x_block", "norm1", "ff"), ("context_block", "norm1_context", "ff_context")): + if branch == "context_block" and i == config["num_layers"] - 1: + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{old}.{branch}.adaLN_modulation.1.{p}",), + (f"{new}.{norm}.linear.{p}",), + ReorderChunks((1, 0)), + ) + ) + else: + modules.extend( + [ + (f"{old}.{branch}.adaLN_modulation.1", f"{new}.{norm}.linear"), + (f"{old}.{branch}.mlp.fc1", f"{new}.{ff}.net.0.proj"), + (f"{old}.{branch}.mlp.fc2", f"{new}.{ff}.net.2"), + ] + ) + for p in ("weight", "bias"): + rules.append(Rule((f"final_layer.adaLN_modulation.1.{p}",), (f"norm_out.linear.{p}",), ReorderChunks((1, 0)))) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/sd3_controlnet.py b/src/diffusers/loaders/conversion/sd3_controlnet.py new file mode 100644 index 000000000000..4cc5ca10bedb --- /dev/null +++ b/src/diffusers/loaders/conversion/sd3_controlnet.py @@ -0,0 +1,63 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .sd3 import sd3_conversion +from .transforms import Split + + +def sd3_controlnet_conversion(config): + modules = [("pos_embed_input.proj", "pos_embed_input.proj")] + modules.extend( + (f"time_text_embed.{kind}.linear_{i}", f"time_text_embed.{kind}.linear_{i}") + for kind in ("timestep_embedder", "text_embedder") + for i in (1, 2) + ) + modules.extend((f"controlnet_blocks.{i}", f"controlnet_blocks.{i}") for i in range(config["num_layers"])) + mapping, rules = {}, [] + if config["joint_attention_dim"] is not None: + # Joint-stream ControlNets were trained directly with Diffusers and already use its parameter names. + base = sd3_conversion({**config, "num_layers": config["num_layers"] + 1}) + for key in sorted(base.diffusers_keys): + if key.startswith("transformer_blocks.") and int(key.split(".")[1]) < config["num_layers"]: + mapping[key] = key + modules.append(("context_embedder", "context_embedder")) + else: + hidden = config["num_attention_heads"] * config["attention_head_dim"] + for i in range(config["num_layers"]): + prefix = f"transformer_blocks.{i}" + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attn.qkv.{p}",), + tuple(f"{prefix}.attn.to_{part}.{p}" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + modules.extend( + (f"{prefix}.{a}", f"{prefix}.{b}") + for a, b in ( + ("attn.proj", "attn.to_out.0"), + ("mlp.fc1", "ff.net.0.proj"), + ("mlp.fc2", "ff.net.2"), + ("adaLN_modulation.1", "norm1.linear"), + ) + ) + if config["use_pos_embed"]: + modules.append(("pos_embed.proj", "pos_embed.proj")) + if config["pos_embed_max_size"] is not None and config["pos_embed_type"] is not None: + mapping["pos_embed.pos_embed"] = "pos_embed.pos_embed" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/shap_e_renderer.py b/src/diffusers/loaders/conversion/shap_e_renderer.py new file mode 100644 index 000000000000..39cf6ad74f56 --- /dev/null +++ b/src/diffusers/loaders/conversion/shap_e_renderer.py @@ -0,0 +1,40 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from .core import Conversion, Rule +from .shap_e_tables import create_mc_lookup_table +from .transforms import WithConstants + + +def shap_e_renderer_conversion(config): + modules = [(f"renderer.nerstf.mlp.{i}", f"mlp.mlp.{i}") for i in range(config["n_hidden_layers"] + 1)] + for name in config["param_names"]: + name = name.replace(".", "__") + modules.extend( + (f"encoder.params_proj.projections.{name}.{leaf}", f"params_proj.projections.{name}.{leaf}") + for leaf in ("proj", "norm") + ) + mapping = {f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")} + anchor = "renderer.nerstf.mlp.0.weight" + target = mapping.pop(anchor) + cases, masks = create_mc_lookup_table() + background = torch.tensor(config["background"], device="cpu") / 255.0 + rule = Rule( + (anchor,), + (target, "void.background", "mesh_decoder.cases", "mesh_decoder.masks"), + WithConstants((background, cases, masks)), + ) + return Conversion(mapping=mapping, rules=(rule,)) diff --git a/src/diffusers/loaders/conversion/shap_e_tables.py b/src/diffusers/loaders/conversion/shap_e_tables.py new file mode 100644 index 000000000000..09780ab13747 --- /dev/null +++ b/src/diffusers/loaders/conversion/shap_e_tables.py @@ -0,0 +1,494 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + + +MC_TABLE = [ + [], + [[0, 1, 0, 2, 0, 4]], + [[1, 0, 1, 5, 1, 3]], + [[0, 4, 1, 5, 0, 2], [1, 5, 1, 3, 0, 2]], + [[2, 0, 2, 3, 2, 6]], + [[0, 1, 2, 3, 0, 4], [2, 3, 2, 6, 0, 4]], + [[1, 0, 1, 5, 1, 3], [2, 6, 0, 2, 3, 2]], + [[3, 2, 2, 6, 3, 1], [3, 1, 2, 6, 1, 5], [1, 5, 2, 6, 0, 4]], + [[3, 1, 3, 7, 3, 2]], + [[0, 2, 0, 4, 0, 1], [3, 7, 2, 3, 1, 3]], + [[1, 5, 3, 7, 1, 0], [3, 7, 3, 2, 1, 0]], + [[2, 0, 0, 4, 2, 3], [2, 3, 0, 4, 3, 7], [3, 7, 0, 4, 1, 5]], + [[2, 0, 3, 1, 2, 6], [3, 1, 3, 7, 2, 6]], + [[1, 3, 3, 7, 1, 0], [1, 0, 3, 7, 0, 4], [0, 4, 3, 7, 2, 6]], + [[0, 1, 1, 5, 0, 2], [0, 2, 1, 5, 2, 6], [2, 6, 1, 5, 3, 7]], + [[0, 4, 1, 5, 3, 7], [0, 4, 3, 7, 2, 6]], + [[4, 0, 4, 6, 4, 5]], + [[0, 2, 4, 6, 0, 1], [4, 6, 4, 5, 0, 1]], + [[1, 5, 1, 3, 1, 0], [4, 6, 5, 4, 0, 4]], + [[5, 1, 1, 3, 5, 4], [5, 4, 1, 3, 4, 6], [4, 6, 1, 3, 0, 2]], + [[2, 0, 2, 3, 2, 6], [4, 5, 0, 4, 6, 4]], + [[6, 4, 4, 5, 6, 2], [6, 2, 4, 5, 2, 3], [2, 3, 4, 5, 0, 1]], + [[2, 6, 2, 0, 3, 2], [1, 0, 1, 5, 3, 1], [6, 4, 5, 4, 0, 4]], + [[1, 3, 5, 4, 1, 5], [1, 3, 4, 6, 5, 4], [1, 3, 3, 2, 4, 6], [3, 2, 2, 6, 4, 6]], + [[3, 1, 3, 7, 3, 2], [6, 4, 5, 4, 0, 4]], + [[4, 5, 0, 1, 4, 6], [0, 1, 0, 2, 4, 6], [7, 3, 2, 3, 1, 3]], + [[3, 2, 1, 0, 3, 7], [1, 0, 1, 5, 3, 7], [6, 4, 5, 4, 0, 4]], + [[3, 7, 3, 2, 1, 5], [3, 2, 6, 4, 1, 5], [1, 5, 6, 4, 5, 4], [3, 2, 2, 0, 6, 4]], + [[3, 7, 2, 6, 3, 1], [2, 6, 2, 0, 3, 1], [5, 4, 0, 4, 6, 4]], + [[1, 0, 1, 3, 5, 4], [1, 3, 2, 6, 5, 4], [1, 3, 3, 7, 2, 6], [5, 4, 2, 6, 4, 6]], + [[0, 1, 1, 5, 0, 2], [0, 2, 1, 5, 2, 6], [2, 6, 1, 5, 3, 7], [4, 5, 0, 4, 4, 6]], + [[6, 2, 4, 6, 4, 5], [4, 5, 5, 1, 6, 2], [6, 2, 5, 1, 7, 3]], + [[5, 1, 5, 4, 5, 7]], + [[0, 1, 0, 2, 0, 4], [5, 7, 1, 5, 4, 5]], + [[1, 0, 5, 4, 1, 3], [5, 4, 5, 7, 1, 3]], + [[4, 5, 5, 7, 4, 0], [4, 0, 5, 7, 0, 2], [0, 2, 5, 7, 1, 3]], + [[2, 0, 2, 3, 2, 6], [7, 5, 1, 5, 4, 5]], + [[2, 6, 0, 4, 2, 3], [0, 4, 0, 1, 2, 3], [7, 5, 1, 5, 4, 5]], + [[5, 7, 1, 3, 5, 4], [1, 3, 1, 0, 5, 4], [6, 2, 0, 2, 3, 2]], + [[3, 1, 3, 2, 7, 5], [3, 2, 0, 4, 7, 5], [3, 2, 2, 6, 0, 4], [7, 5, 0, 4, 5, 4]], + [[3, 7, 3, 2, 3, 1], [5, 4, 7, 5, 1, 5]], + [[0, 4, 0, 1, 2, 0], [3, 1, 3, 7, 2, 3], [4, 5, 7, 5, 1, 5]], + [[7, 3, 3, 2, 7, 5], [7, 5, 3, 2, 5, 4], [5, 4, 3, 2, 1, 0]], + [[0, 4, 2, 3, 0, 2], [0, 4, 3, 7, 2, 3], [0, 4, 4, 5, 3, 7], [4, 5, 5, 7, 3, 7]], + [[2, 0, 3, 1, 2, 6], [3, 1, 3, 7, 2, 6], [4, 5, 7, 5, 1, 5]], + [[1, 3, 3, 7, 1, 0], [1, 0, 3, 7, 0, 4], [0, 4, 3, 7, 2, 6], [5, 7, 1, 5, 5, 4]], + [[2, 6, 2, 0, 3, 7], [2, 0, 4, 5, 3, 7], [3, 7, 4, 5, 7, 5], [2, 0, 0, 1, 4, 5]], + [[4, 0, 5, 4, 5, 7], [5, 7, 7, 3, 4, 0], [4, 0, 7, 3, 6, 2]], + [[4, 6, 5, 7, 4, 0], [5, 7, 5, 1, 4, 0]], + [[1, 0, 0, 2, 1, 5], [1, 5, 0, 2, 5, 7], [5, 7, 0, 2, 4, 6]], + [[0, 4, 4, 6, 0, 1], [0, 1, 4, 6, 1, 3], [1, 3, 4, 6, 5, 7]], + [[0, 2, 4, 6, 5, 7], [0, 2, 5, 7, 1, 3]], + [[5, 1, 4, 0, 5, 7], [4, 0, 4, 6, 5, 7], [3, 2, 6, 2, 0, 2]], + [[2, 3, 2, 6, 0, 1], [2, 6, 7, 5, 0, 1], [0, 1, 7, 5, 1, 5], [2, 6, 6, 4, 7, 5]], + [[0, 4, 4, 6, 0, 1], [0, 1, 4, 6, 1, 3], [1, 3, 4, 6, 5, 7], [2, 6, 0, 2, 2, 3]], + [[3, 1, 2, 3, 2, 6], [2, 6, 6, 4, 3, 1], [3, 1, 6, 4, 7, 5]], + [[4, 6, 5, 7, 4, 0], [5, 7, 5, 1, 4, 0], [2, 3, 1, 3, 7, 3]], + [[1, 0, 0, 2, 1, 5], [1, 5, 0, 2, 5, 7], [5, 7, 0, 2, 4, 6], [3, 2, 1, 3, 3, 7]], + [[0, 1, 0, 4, 2, 3], [0, 4, 5, 7, 2, 3], [0, 4, 4, 6, 5, 7], [2, 3, 5, 7, 3, 7]], + [[7, 5, 3, 7, 3, 2], [3, 2, 2, 0, 7, 5], [7, 5, 2, 0, 6, 4]], + [[0, 4, 4, 6, 5, 7], [0, 4, 5, 7, 1, 5], [0, 2, 1, 3, 3, 7], [3, 7, 2, 6, 0, 2]], + [ + [3, 1, 7, 3, 6, 2], + [6, 2, 0, 1, 3, 1], + [6, 4, 0, 1, 6, 2], + [6, 4, 5, 1, 0, 1], + [6, 4, 7, 5, 5, 1], + ], + [ + [4, 0, 6, 4, 7, 5], + [7, 5, 1, 0, 4, 0], + [7, 3, 1, 0, 7, 5], + [7, 3, 2, 0, 1, 0], + [7, 3, 6, 2, 2, 0], + ], + [[7, 3, 6, 2, 6, 4], [7, 5, 7, 3, 6, 4]], + [[6, 2, 6, 7, 6, 4]], + [[0, 4, 0, 1, 0, 2], [6, 7, 4, 6, 2, 6]], + [[1, 0, 1, 5, 1, 3], [7, 6, 4, 6, 2, 6]], + [[1, 3, 0, 2, 1, 5], [0, 2, 0, 4, 1, 5], [7, 6, 4, 6, 2, 6]], + [[2, 3, 6, 7, 2, 0], [6, 7, 6, 4, 2, 0]], + [[4, 0, 0, 1, 4, 6], [4, 6, 0, 1, 6, 7], [6, 7, 0, 1, 2, 3]], + [[6, 4, 2, 0, 6, 7], [2, 0, 2, 3, 6, 7], [5, 1, 3, 1, 0, 1]], + [[1, 5, 1, 3, 0, 4], [1, 3, 7, 6, 0, 4], [0, 4, 7, 6, 4, 6], [1, 3, 3, 2, 7, 6]], + [[3, 2, 3, 1, 3, 7], [6, 4, 2, 6, 7, 6]], + [[3, 7, 3, 2, 1, 3], [0, 2, 0, 4, 1, 0], [7, 6, 4, 6, 2, 6]], + [[1, 5, 3, 7, 1, 0], [3, 7, 3, 2, 1, 0], [4, 6, 2, 6, 7, 6]], + [[2, 0, 0, 4, 2, 3], [2, 3, 0, 4, 3, 7], [3, 7, 0, 4, 1, 5], [6, 4, 2, 6, 6, 7]], + [[7, 6, 6, 4, 7, 3], [7, 3, 6, 4, 3, 1], [3, 1, 6, 4, 2, 0]], + [[0, 1, 4, 6, 0, 4], [0, 1, 6, 7, 4, 6], [0, 1, 1, 3, 6, 7], [1, 3, 3, 7, 6, 7]], + [[0, 2, 0, 1, 4, 6], [0, 1, 3, 7, 4, 6], [0, 1, 1, 5, 3, 7], [4, 6, 3, 7, 6, 7]], + [[7, 3, 6, 7, 6, 4], [6, 4, 4, 0, 7, 3], [7, 3, 4, 0, 5, 1]], + [[4, 0, 6, 2, 4, 5], [6, 2, 6, 7, 4, 5]], + [[2, 6, 6, 7, 2, 0], [2, 0, 6, 7, 0, 1], [0, 1, 6, 7, 4, 5]], + [[6, 7, 4, 5, 6, 2], [4, 5, 4, 0, 6, 2], [3, 1, 0, 1, 5, 1]], + [[2, 0, 2, 6, 3, 1], [2, 6, 4, 5, 3, 1], [2, 6, 6, 7, 4, 5], [3, 1, 4, 5, 1, 5]], + [[0, 2, 2, 3, 0, 4], [0, 4, 2, 3, 4, 5], [4, 5, 2, 3, 6, 7]], + [[0, 1, 2, 3, 6, 7], [0, 1, 6, 7, 4, 5]], + [[0, 2, 2, 3, 0, 4], [0, 4, 2, 3, 4, 5], [4, 5, 2, 3, 6, 7], [1, 3, 0, 1, 1, 5]], + [[5, 4, 1, 5, 1, 3], [1, 3, 3, 2, 5, 4], [5, 4, 3, 2, 7, 6]], + [[4, 0, 6, 2, 4, 5], [6, 2, 6, 7, 4, 5], [1, 3, 7, 3, 2, 3]], + [[2, 6, 6, 7, 2, 0], [2, 0, 6, 7, 0, 1], [0, 1, 6, 7, 4, 5], [3, 7, 2, 3, 3, 1]], + [[0, 1, 1, 5, 3, 7], [0, 1, 3, 7, 2, 3], [0, 4, 2, 6, 6, 7], [6, 7, 4, 5, 0, 4]], + [ + [6, 2, 7, 6, 5, 4], + [5, 4, 0, 2, 6, 2], + [5, 1, 0, 2, 5, 4], + [5, 1, 3, 2, 0, 2], + [5, 1, 7, 3, 3, 2], + ], + [[3, 1, 3, 7, 2, 0], [3, 7, 5, 4, 2, 0], [2, 0, 5, 4, 0, 4], [3, 7, 7, 6, 5, 4]], + [[1, 0, 3, 1, 3, 7], [3, 7, 7, 6, 1, 0], [1, 0, 7, 6, 5, 4]], + [ + [1, 0, 5, 1, 7, 3], + [7, 3, 2, 0, 1, 0], + [7, 6, 2, 0, 7, 3], + [7, 6, 4, 0, 2, 0], + [7, 6, 5, 4, 4, 0], + ], + [[7, 6, 5, 4, 5, 1], [7, 3, 7, 6, 5, 1]], + [[5, 7, 5, 1, 5, 4], [6, 2, 7, 6, 4, 6]], + [[0, 2, 0, 4, 1, 0], [5, 4, 5, 7, 1, 5], [2, 6, 7, 6, 4, 6]], + [[1, 0, 5, 4, 1, 3], [5, 4, 5, 7, 1, 3], [2, 6, 7, 6, 4, 6]], + [[4, 5, 5, 7, 4, 0], [4, 0, 5, 7, 0, 2], [0, 2, 5, 7, 1, 3], [6, 7, 4, 6, 6, 2]], + [[2, 3, 6, 7, 2, 0], [6, 7, 6, 4, 2, 0], [1, 5, 4, 5, 7, 5]], + [[4, 0, 0, 1, 4, 6], [4, 6, 0, 1, 6, 7], [6, 7, 0, 1, 2, 3], [5, 1, 4, 5, 5, 7]], + [[0, 2, 2, 3, 6, 7], [0, 2, 6, 7, 4, 6], [0, 1, 4, 5, 5, 7], [5, 7, 1, 3, 0, 1]], + [ + [5, 4, 7, 5, 3, 1], + [3, 1, 0, 4, 5, 4], + [3, 2, 0, 4, 3, 1], + [3, 2, 6, 4, 0, 4], + [3, 2, 7, 6, 6, 4], + ], + [[5, 4, 5, 7, 1, 5], [3, 7, 3, 2, 1, 3], [4, 6, 2, 6, 7, 6]], + [[1, 0, 0, 2, 0, 4], [1, 5, 5, 4, 5, 7], [3, 2, 1, 3, 3, 7], [2, 6, 7, 6, 4, 6]], + [[7, 3, 3, 2, 7, 5], [7, 5, 3, 2, 5, 4], [5, 4, 3, 2, 1, 0], [6, 2, 7, 6, 6, 4]], + [ + [0, 4, 2, 3, 0, 2], + [0, 4, 3, 7, 2, 3], + [0, 4, 4, 5, 3, 7], + [4, 5, 5, 7, 3, 7], + [6, 7, 4, 6, 2, 6], + ], + [[7, 6, 6, 4, 7, 3], [7, 3, 6, 4, 3, 1], [3, 1, 6, 4, 2, 0], [5, 4, 7, 5, 5, 1]], + [ + [0, 1, 4, 6, 0, 4], + [0, 1, 6, 7, 4, 6], + [0, 1, 1, 3, 6, 7], + [1, 3, 3, 7, 6, 7], + [5, 7, 1, 5, 4, 5], + ], + [ + [6, 7, 4, 6, 0, 2], + [0, 2, 3, 7, 6, 7], + [0, 1, 3, 7, 0, 2], + [0, 1, 5, 7, 3, 7], + [0, 1, 4, 5, 5, 7], + ], + [[4, 0, 6, 7, 4, 6], [4, 0, 7, 3, 6, 7], [4, 0, 5, 7, 7, 3], [4, 5, 5, 7, 4, 0]], + [[7, 5, 5, 1, 7, 6], [7, 6, 5, 1, 6, 2], [6, 2, 5, 1, 4, 0]], + [[0, 2, 1, 5, 0, 1], [0, 2, 5, 7, 1, 5], [0, 2, 2, 6, 5, 7], [2, 6, 6, 7, 5, 7]], + [[1, 3, 1, 0, 5, 7], [1, 0, 2, 6, 5, 7], [5, 7, 2, 6, 7, 6], [1, 0, 0, 4, 2, 6]], + [[2, 0, 6, 2, 6, 7], [6, 7, 7, 5, 2, 0], [2, 0, 7, 5, 3, 1]], + [[0, 4, 0, 2, 1, 5], [0, 2, 6, 7, 1, 5], [0, 2, 2, 3, 6, 7], [1, 5, 6, 7, 5, 7]], + [[7, 6, 5, 7, 5, 1], [5, 1, 1, 0, 7, 6], [7, 6, 1, 0, 3, 2]], + [ + [2, 0, 3, 2, 7, 6], + [7, 6, 4, 0, 2, 0], + [7, 5, 4, 0, 7, 6], + [7, 5, 1, 0, 4, 0], + [7, 5, 3, 1, 1, 0], + ], + [[7, 5, 3, 1, 3, 2], [7, 6, 7, 5, 3, 2]], + [[7, 5, 5, 1, 7, 6], [7, 6, 5, 1, 6, 2], [6, 2, 5, 1, 4, 0], [3, 1, 7, 3, 3, 2]], + [ + [0, 2, 1, 5, 0, 1], + [0, 2, 5, 7, 1, 5], + [0, 2, 2, 6, 5, 7], + [2, 6, 6, 7, 5, 7], + [3, 7, 2, 3, 1, 3], + ], + [ + [3, 7, 2, 3, 0, 1], + [0, 1, 5, 7, 3, 7], + [0, 4, 5, 7, 0, 1], + [0, 4, 6, 7, 5, 7], + [0, 4, 2, 6, 6, 7], + ], + [[2, 0, 3, 7, 2, 3], [2, 0, 7, 5, 3, 7], [2, 0, 6, 7, 7, 5], [2, 6, 6, 7, 2, 0]], + [ + [5, 7, 1, 5, 0, 4], + [0, 4, 6, 7, 5, 7], + [0, 2, 6, 7, 0, 4], + [0, 2, 3, 7, 6, 7], + [0, 2, 1, 3, 3, 7], + ], + [[1, 0, 5, 7, 1, 5], [1, 0, 7, 6, 5, 7], [1, 0, 3, 7, 7, 6], [1, 3, 3, 7, 1, 0]], + [[0, 2, 0, 1, 0, 4], [3, 7, 6, 7, 5, 7]], + [[7, 5, 7, 3, 7, 6]], + [[7, 3, 7, 5, 7, 6]], + [[0, 1, 0, 2, 0, 4], [6, 7, 3, 7, 5, 7]], + [[1, 3, 1, 0, 1, 5], [7, 6, 3, 7, 5, 7]], + [[0, 4, 1, 5, 0, 2], [1, 5, 1, 3, 0, 2], [6, 7, 3, 7, 5, 7]], + [[2, 6, 2, 0, 2, 3], [7, 5, 6, 7, 3, 7]], + [[0, 1, 2, 3, 0, 4], [2, 3, 2, 6, 0, 4], [5, 7, 6, 7, 3, 7]], + [[1, 5, 1, 3, 0, 1], [2, 3, 2, 6, 0, 2], [5, 7, 6, 7, 3, 7]], + [[3, 2, 2, 6, 3, 1], [3, 1, 2, 6, 1, 5], [1, 5, 2, 6, 0, 4], [7, 6, 3, 7, 7, 5]], + [[3, 1, 7, 5, 3, 2], [7, 5, 7, 6, 3, 2]], + [[7, 6, 3, 2, 7, 5], [3, 2, 3, 1, 7, 5], [4, 0, 1, 0, 2, 0]], + [[5, 7, 7, 6, 5, 1], [5, 1, 7, 6, 1, 0], [1, 0, 7, 6, 3, 2]], + [[2, 3, 2, 0, 6, 7], [2, 0, 1, 5, 6, 7], [2, 0, 0, 4, 1, 5], [6, 7, 1, 5, 7, 5]], + [[6, 2, 2, 0, 6, 7], [6, 7, 2, 0, 7, 5], [7, 5, 2, 0, 3, 1]], + [[0, 4, 0, 1, 2, 6], [0, 1, 5, 7, 2, 6], [2, 6, 5, 7, 6, 7], [0, 1, 1, 3, 5, 7]], + [[1, 5, 0, 2, 1, 0], [1, 5, 2, 6, 0, 2], [1, 5, 5, 7, 2, 6], [5, 7, 7, 6, 2, 6]], + [[5, 1, 7, 5, 7, 6], [7, 6, 6, 2, 5, 1], [5, 1, 6, 2, 4, 0]], + [[4, 5, 4, 0, 4, 6], [7, 3, 5, 7, 6, 7]], + [[0, 2, 4, 6, 0, 1], [4, 6, 4, 5, 0, 1], [3, 7, 5, 7, 6, 7]], + [[4, 6, 4, 5, 0, 4], [1, 5, 1, 3, 0, 1], [6, 7, 3, 7, 5, 7]], + [[5, 1, 1, 3, 5, 4], [5, 4, 1, 3, 4, 6], [4, 6, 1, 3, 0, 2], [7, 3, 5, 7, 7, 6]], + [[2, 3, 2, 6, 0, 2], [4, 6, 4, 5, 0, 4], [3, 7, 5, 7, 6, 7]], + [[6, 4, 4, 5, 6, 2], [6, 2, 4, 5, 2, 3], [2, 3, 4, 5, 0, 1], [7, 5, 6, 7, 7, 3]], + [[0, 1, 1, 5, 1, 3], [0, 2, 2, 3, 2, 6], [4, 5, 0, 4, 4, 6], [5, 7, 6, 7, 3, 7]], + [ + [1, 3, 5, 4, 1, 5], + [1, 3, 4, 6, 5, 4], + [1, 3, 3, 2, 4, 6], + [3, 2, 2, 6, 4, 6], + [7, 6, 3, 7, 5, 7], + ], + [[3, 1, 7, 5, 3, 2], [7, 5, 7, 6, 3, 2], [0, 4, 6, 4, 5, 4]], + [[1, 0, 0, 2, 4, 6], [1, 0, 4, 6, 5, 4], [1, 3, 5, 7, 7, 6], [7, 6, 3, 2, 1, 3]], + [[5, 7, 7, 6, 5, 1], [5, 1, 7, 6, 1, 0], [1, 0, 7, 6, 3, 2], [4, 6, 5, 4, 4, 0]], + [ + [7, 5, 6, 7, 2, 3], + [2, 3, 1, 5, 7, 5], + [2, 0, 1, 5, 2, 3], + [2, 0, 4, 5, 1, 5], + [2, 0, 6, 4, 4, 5], + ], + [[6, 2, 2, 0, 6, 7], [6, 7, 2, 0, 7, 5], [7, 5, 2, 0, 3, 1], [4, 0, 6, 4, 4, 5]], + [ + [4, 6, 5, 4, 1, 0], + [1, 0, 2, 6, 4, 6], + [1, 3, 2, 6, 1, 0], + [1, 3, 7, 6, 2, 6], + [1, 3, 5, 7, 7, 6], + ], + [ + [1, 5, 0, 2, 1, 0], + [1, 5, 2, 6, 0, 2], + [1, 5, 5, 7, 2, 6], + [5, 7, 7, 6, 2, 6], + [4, 6, 5, 4, 0, 4], + ], + [[5, 1, 4, 6, 5, 4], [5, 1, 6, 2, 4, 6], [5, 1, 7, 6, 6, 2], [5, 7, 7, 6, 5, 1]], + [[5, 4, 7, 6, 5, 1], [7, 6, 7, 3, 5, 1]], + [[7, 3, 5, 1, 7, 6], [5, 1, 5, 4, 7, 6], [2, 0, 4, 0, 1, 0]], + [[3, 1, 1, 0, 3, 7], [3, 7, 1, 0, 7, 6], [7, 6, 1, 0, 5, 4]], + [[0, 2, 0, 4, 1, 3], [0, 4, 6, 7, 1, 3], [1, 3, 6, 7, 3, 7], [0, 4, 4, 5, 6, 7]], + [[5, 4, 7, 6, 5, 1], [7, 6, 7, 3, 5, 1], [0, 2, 3, 2, 6, 2]], + [[1, 5, 5, 4, 7, 6], [1, 5, 7, 6, 3, 7], [1, 0, 3, 2, 2, 6], [2, 6, 0, 4, 1, 0]], + [[3, 1, 1, 0, 3, 7], [3, 7, 1, 0, 7, 6], [7, 6, 1, 0, 5, 4], [2, 0, 3, 2, 2, 6]], + [ + [2, 3, 6, 2, 4, 0], + [4, 0, 1, 3, 2, 3], + [4, 5, 1, 3, 4, 0], + [4, 5, 7, 3, 1, 3], + [4, 5, 6, 7, 7, 3], + ], + [[1, 5, 5, 4, 1, 3], [1, 3, 5, 4, 3, 2], [3, 2, 5, 4, 7, 6]], + [[1, 5, 5, 4, 1, 3], [1, 3, 5, 4, 3, 2], [3, 2, 5, 4, 7, 6], [0, 4, 1, 0, 0, 2]], + [[1, 0, 5, 4, 7, 6], [1, 0, 7, 6, 3, 2]], + [[2, 3, 0, 2, 0, 4], [0, 4, 4, 5, 2, 3], [2, 3, 4, 5, 6, 7]], + [[1, 3, 1, 5, 0, 2], [1, 5, 7, 6, 0, 2], [1, 5, 5, 4, 7, 6], [0, 2, 7, 6, 2, 6]], + [ + [5, 1, 4, 5, 6, 7], + [6, 7, 3, 1, 5, 1], + [6, 2, 3, 1, 6, 7], + [6, 2, 0, 1, 3, 1], + [6, 2, 4, 0, 0, 1], + ], + [[6, 7, 2, 6, 2, 0], [2, 0, 0, 1, 6, 7], [6, 7, 0, 1, 4, 5]], + [[6, 2, 4, 0, 4, 5], [6, 7, 6, 2, 4, 5]], + [[6, 7, 7, 3, 6, 4], [6, 4, 7, 3, 4, 0], [4, 0, 7, 3, 5, 1]], + [[1, 5, 1, 0, 3, 7], [1, 0, 4, 6, 3, 7], [1, 0, 0, 2, 4, 6], [3, 7, 4, 6, 7, 6]], + [[1, 0, 3, 7, 1, 3], [1, 0, 7, 6, 3, 7], [1, 0, 0, 4, 7, 6], [0, 4, 4, 6, 7, 6]], + [[6, 4, 7, 6, 7, 3], [7, 3, 3, 1, 6, 4], [6, 4, 3, 1, 2, 0]], + [[6, 7, 7, 3, 6, 4], [6, 4, 7, 3, 4, 0], [4, 0, 7, 3, 5, 1], [2, 3, 6, 2, 2, 0]], + [ + [7, 6, 3, 7, 1, 5], + [1, 5, 4, 6, 7, 6], + [1, 0, 4, 6, 1, 5], + [1, 0, 2, 6, 4, 6], + [1, 0, 3, 2, 2, 6], + ], + [ + [1, 0, 3, 7, 1, 3], + [1, 0, 7, 6, 3, 7], + [1, 0, 0, 4, 7, 6], + [0, 4, 4, 6, 7, 6], + [2, 6, 0, 2, 3, 2], + ], + [[3, 1, 7, 6, 3, 7], [3, 1, 6, 4, 7, 6], [3, 1, 2, 6, 6, 4], [3, 2, 2, 6, 3, 1]], + [[3, 2, 3, 1, 7, 6], [3, 1, 0, 4, 7, 6], [7, 6, 0, 4, 6, 4], [3, 1, 1, 5, 0, 4]], + [ + [0, 1, 2, 0, 6, 4], + [6, 4, 5, 1, 0, 1], + [6, 7, 5, 1, 6, 4], + [6, 7, 3, 1, 5, 1], + [6, 7, 2, 3, 3, 1], + ], + [[0, 1, 4, 0, 4, 6], [4, 6, 6, 7, 0, 1], [0, 1, 6, 7, 2, 3]], + [[6, 7, 2, 3, 2, 0], [6, 4, 6, 7, 2, 0]], + [ + [2, 6, 0, 2, 1, 3], + [1, 3, 7, 6, 2, 6], + [1, 5, 7, 6, 1, 3], + [1, 5, 4, 6, 7, 6], + [1, 5, 0, 4, 4, 6], + ], + [[1, 5, 1, 0, 1, 3], [4, 6, 7, 6, 2, 6]], + [[0, 1, 2, 6, 0, 2], [0, 1, 6, 7, 2, 6], [0, 1, 4, 6, 6, 7], [0, 4, 4, 6, 0, 1]], + [[6, 7, 6, 2, 6, 4]], + [[6, 2, 7, 3, 6, 4], [7, 3, 7, 5, 6, 4]], + [[7, 5, 6, 4, 7, 3], [6, 4, 6, 2, 7, 3], [1, 0, 2, 0, 4, 0]], + [[6, 2, 7, 3, 6, 4], [7, 3, 7, 5, 6, 4], [0, 1, 5, 1, 3, 1]], + [[2, 0, 0, 4, 1, 5], [2, 0, 1, 5, 3, 1], [2, 6, 3, 7, 7, 5], [7, 5, 6, 4, 2, 6]], + [[3, 7, 7, 5, 3, 2], [3, 2, 7, 5, 2, 0], [2, 0, 7, 5, 6, 4]], + [[3, 2, 3, 7, 1, 0], [3, 7, 6, 4, 1, 0], [3, 7, 7, 5, 6, 4], [1, 0, 6, 4, 0, 4]], + [[3, 7, 7, 5, 3, 2], [3, 2, 7, 5, 2, 0], [2, 0, 7, 5, 6, 4], [1, 5, 3, 1, 1, 0]], + [ + [7, 3, 5, 7, 4, 6], + [4, 6, 2, 3, 7, 3], + [4, 0, 2, 3, 4, 6], + [4, 0, 1, 3, 2, 3], + [4, 0, 5, 1, 1, 3], + ], + [[2, 3, 3, 1, 2, 6], [2, 6, 3, 1, 6, 4], [6, 4, 3, 1, 7, 5]], + [[2, 3, 3, 1, 2, 6], [2, 6, 3, 1, 6, 4], [6, 4, 3, 1, 7, 5], [0, 1, 2, 0, 0, 4]], + [[1, 0, 1, 5, 3, 2], [1, 5, 4, 6, 3, 2], [3, 2, 4, 6, 2, 6], [1, 5, 5, 7, 4, 6]], + [ + [0, 2, 4, 0, 5, 1], + [5, 1, 3, 2, 0, 2], + [5, 7, 3, 2, 5, 1], + [5, 7, 6, 2, 3, 2], + [5, 7, 4, 6, 6, 2], + ], + [[2, 0, 3, 1, 7, 5], [2, 0, 7, 5, 6, 4]], + [[4, 6, 0, 4, 0, 1], [0, 1, 1, 3, 4, 6], [4, 6, 1, 3, 5, 7]], + [[0, 2, 1, 0, 1, 5], [1, 5, 5, 7, 0, 2], [0, 2, 5, 7, 4, 6]], + [[5, 7, 4, 6, 4, 0], [5, 1, 5, 7, 4, 0]], + [[5, 4, 4, 0, 5, 7], [5, 7, 4, 0, 7, 3], [7, 3, 4, 0, 6, 2]], + [[0, 1, 0, 2, 4, 5], [0, 2, 3, 7, 4, 5], [4, 5, 3, 7, 5, 7], [0, 2, 2, 6, 3, 7]], + [[5, 4, 4, 0, 5, 7], [5, 7, 4, 0, 7, 3], [7, 3, 4, 0, 6, 2], [1, 0, 5, 1, 1, 3]], + [ + [1, 5, 3, 1, 2, 0], + [2, 0, 4, 5, 1, 5], + [2, 6, 4, 5, 2, 0], + [2, 6, 7, 5, 4, 5], + [2, 6, 3, 7, 7, 5], + ], + [[2, 3, 0, 4, 2, 0], [2, 3, 4, 5, 0, 4], [2, 3, 3, 7, 4, 5], [3, 7, 7, 5, 4, 5]], + [[3, 2, 7, 3, 7, 5], [7, 5, 5, 4, 3, 2], [3, 2, 5, 4, 1, 0]], + [ + [2, 3, 0, 4, 2, 0], + [2, 3, 4, 5, 0, 4], + [2, 3, 3, 7, 4, 5], + [3, 7, 7, 5, 4, 5], + [1, 5, 3, 1, 0, 1], + ], + [[3, 2, 1, 5, 3, 1], [3, 2, 5, 4, 1, 5], [3, 2, 7, 5, 5, 4], [3, 7, 7, 5, 3, 2]], + [[2, 6, 2, 3, 0, 4], [2, 3, 7, 5, 0, 4], [2, 3, 3, 1, 7, 5], [0, 4, 7, 5, 4, 5]], + [ + [3, 2, 1, 3, 5, 7], + [5, 7, 6, 2, 3, 2], + [5, 4, 6, 2, 5, 7], + [5, 4, 0, 2, 6, 2], + [5, 4, 1, 0, 0, 2], + ], + [ + [4, 5, 0, 4, 2, 6], + [2, 6, 7, 5, 4, 5], + [2, 3, 7, 5, 2, 6], + [2, 3, 1, 5, 7, 5], + [2, 3, 0, 1, 1, 5], + ], + [[2, 3, 2, 0, 2, 6], [1, 5, 7, 5, 4, 5]], + [[5, 7, 4, 5, 4, 0], [4, 0, 0, 2, 5, 7], [5, 7, 0, 2, 1, 3]], + [[5, 4, 1, 0, 1, 3], [5, 7, 5, 4, 1, 3]], + [[0, 2, 4, 5, 0, 4], [0, 2, 5, 7, 4, 5], [0, 2, 1, 5, 5, 7], [0, 1, 1, 5, 0, 2]], + [[5, 4, 5, 1, 5, 7]], + [[4, 6, 6, 2, 4, 5], [4, 5, 6, 2, 5, 1], [5, 1, 6, 2, 7, 3]], + [[4, 6, 6, 2, 4, 5], [4, 5, 6, 2, 5, 1], [5, 1, 6, 2, 7, 3], [0, 2, 4, 0, 0, 1]], + [[3, 7, 3, 1, 2, 6], [3, 1, 5, 4, 2, 6], [3, 1, 1, 0, 5, 4], [2, 6, 5, 4, 6, 4]], + [ + [6, 4, 2, 6, 3, 7], + [3, 7, 5, 4, 6, 4], + [3, 1, 5, 4, 3, 7], + [3, 1, 0, 4, 5, 4], + [3, 1, 2, 0, 0, 4], + ], + [[2, 0, 2, 3, 6, 4], [2, 3, 1, 5, 6, 4], [6, 4, 1, 5, 4, 5], [2, 3, 3, 7, 1, 5]], + [ + [0, 4, 1, 0, 3, 2], + [3, 2, 6, 4, 0, 4], + [3, 7, 6, 4, 3, 2], + [3, 7, 5, 4, 6, 4], + [3, 7, 1, 5, 5, 4], + ], + [ + [1, 3, 0, 1, 4, 5], + [4, 5, 7, 3, 1, 3], + [4, 6, 7, 3, 4, 5], + [4, 6, 2, 3, 7, 3], + [4, 6, 0, 2, 2, 3], + ], + [[3, 7, 3, 1, 3, 2], [5, 4, 6, 4, 0, 4]], + [[3, 1, 2, 6, 3, 2], [3, 1, 6, 4, 2, 6], [3, 1, 1, 5, 6, 4], [1, 5, 5, 4, 6, 4]], + [ + [3, 1, 2, 6, 3, 2], + [3, 1, 6, 4, 2, 6], + [3, 1, 1, 5, 6, 4], + [1, 5, 5, 4, 6, 4], + [0, 4, 1, 0, 2, 0], + ], + [[4, 5, 6, 4, 6, 2], [6, 2, 2, 3, 4, 5], [4, 5, 2, 3, 0, 1]], + [[2, 3, 6, 4, 2, 6], [2, 3, 4, 5, 6, 4], [2, 3, 0, 4, 4, 5], [2, 0, 0, 4, 2, 3]], + [[1, 3, 5, 1, 5, 4], [5, 4, 4, 6, 1, 3], [1, 3, 4, 6, 0, 2]], + [[1, 3, 0, 4, 1, 0], [1, 3, 4, 6, 0, 4], [1, 3, 5, 4, 4, 6], [1, 5, 5, 4, 1, 3]], + [[4, 6, 0, 2, 0, 1], [4, 5, 4, 6, 0, 1]], + [[4, 6, 4, 0, 4, 5]], + [[4, 0, 6, 2, 7, 3], [4, 0, 7, 3, 5, 1]], + [[1, 5, 0, 1, 0, 2], [0, 2, 2, 6, 1, 5], [1, 5, 2, 6, 3, 7]], + [[3, 7, 1, 3, 1, 0], [1, 0, 0, 4, 3, 7], [3, 7, 0, 4, 2, 6]], + [[3, 1, 2, 0, 2, 6], [3, 7, 3, 1, 2, 6]], + [[0, 4, 2, 0, 2, 3], [2, 3, 3, 7, 0, 4], [0, 4, 3, 7, 1, 5]], + [[3, 7, 1, 5, 1, 0], [3, 2, 3, 7, 1, 0]], + [[0, 4, 1, 3, 0, 1], [0, 4, 3, 7, 1, 3], [0, 4, 2, 3, 3, 7], [0, 2, 2, 3, 0, 4]], + [[3, 7, 3, 1, 3, 2]], + [[2, 6, 3, 2, 3, 1], [3, 1, 1, 5, 2, 6], [2, 6, 1, 5, 0, 4]], + [[1, 5, 3, 2, 1, 3], [1, 5, 2, 6, 3, 2], [1, 5, 0, 2, 2, 6], [1, 0, 0, 2, 1, 5]], + [[2, 3, 0, 1, 0, 4], [2, 6, 2, 3, 0, 4]], + [[2, 3, 2, 0, 2, 6]], + [[1, 5, 0, 4, 0, 2], [1, 3, 1, 5, 0, 2]], + [[1, 5, 1, 0, 1, 3]], + [[0, 2, 0, 1, 0, 4]], + [], +] + + +def create_mc_lookup_table(): + cases = torch.zeros(256, 5, 3, dtype=torch.long, device="cpu") + masks = torch.zeros(256, 5, dtype=torch.bool, device="cpu") + + edge_to_index = { + (0, 1): 0, + (2, 3): 1, + (4, 5): 2, + (6, 7): 3, + (0, 2): 4, + (1, 3): 5, + (4, 6): 6, + (5, 7): 7, + (0, 4): 8, + (1, 5): 9, + (2, 6): 10, + (3, 7): 11, + } + + for i, case in enumerate(MC_TABLE): + for j, tri in enumerate(case): + for k, (c1, c2) in enumerate(zip(tri[::2], tri[1::2])): + cases[i, j, k] = edge_to_index[(c1, c2) if c1 < c2 else (c2, c1)] + masks[i, j] = True + return cases, masks diff --git a/src/diffusers/loaders/conversion/skyreels_v2.py b/src/diffusers/loaders/conversion/skyreels_v2.py new file mode 100644 index 000000000000..c4b446cab57f --- /dev/null +++ b/src/diffusers/loaders/conversion/skyreels_v2.py @@ -0,0 +1,31 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion +from .wan import wan_conversion + + +def skyreels_v2_conversion(config): + base = wan_conversion(config) + mapping = dict(base.mapping) + if config["inject_sample_info"]: + mapping["fps_embedding.weight"] = "fps_embedding.weight" + mapping.update( + { + f"fps_projection.{i}.{p}": f"fps_projection.{name}.{p}" + for i, name in ((0, "net.0.proj"), (2, "net.2")) + for p in ("weight", "bias") + } + ) + return Conversion(mapping=mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/source.py b/src/diffusers/loaders/conversion/source.py new file mode 100644 index 000000000000..061f92d35126 --- /dev/null +++ b/src/diffusers/loaders/conversion/source.py @@ -0,0 +1,132 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compose source files and unpack tensor trees before applying a component conversion.""" + +import json +from collections.abc import Mapping +from pathlib import Path + +import torch + + +class MergedCheckpoint(Mapping): + """Merge disjoint tensor dictionaries lazily, with optional component namespaces.""" + + def __init__(self, sources): + self.sources = tuple(sources) + self.locations = {} + for index, (prefix, source) in enumerate(self.sources): + for key in source: + target = prefix + key + if target in self.locations: + raise ValueError(f"Duplicate source tensor: {target}.") + self.locations[target] = (index, key) + if not self.locations: + raise ValueError("No tensors found in the checkpoint sources.") + + def __getitem__(self, key): + index, source_key = self.locations[key] + return self.sources[index][1][source_key] + + def __iter__(self): + return iter(self.locations) + + def __len__(self): + return len(self.locations) + + +def load_source_manifest(path): + """Load JSON `sources` entries containing path, input_prefix, output_prefix and wrapper fields. + + Paths are local and relative to the manifest directory unless absolute. Prefixes select or namespace complete + components; they do not rename model parameters. Tensor conversion remains the registered `Conversion`. + """ + from .io import Checkpoint + + path = Path(path) + manifest = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(manifest, dict) or set(manifest) != {"sources"} or not isinstance(manifest["sources"], list): + raise ValueError("A source manifest must contain a sources list.") + sources = [] + for entry in manifest["sources"]: + allowed = {"path", "input_prefix", "output_prefix", "wrapper", "format"} + if not isinstance(entry, dict) or "path" not in entry or set(entry) - allowed: + raise ValueError("Source entries require path and may specify input_prefix, output_prefix and wrapper.") + for key in ("path", "input_prefix", "output_prefix"): + if not isinstance(entry.get(key, ""), str): + raise ValueError(f"Source {key} must be a string.") + wrapper = entry.get("wrapper", []) + if not isinstance(wrapper, list) or not all(isinstance(key, str) for key in wrapper): + raise ValueError("Source wrapper must be a list of nested dictionary keys.") + source_path = path.parent / entry["path"] + source_format = entry.get("format", "tensors") + if source_format in ("torchscript", "python-model"): + prefix = entry.get("input_prefix", "") + if source_format == "torchscript": + if wrapper: + raise ValueError("TorchScript sources do not use dictionary wrappers.") + model = torch.jit.load(str(source_path), map_location="cpu") + else: + # Explicit opt-in for legacy research checkpoints containing a pickled nn.Module. + # Loading these executes pickle code and requires the original Python model package. + model = torch.load(source_path, map_location="cpu", weights_only=False) + for key in wrapper: + model = model[key] + state = model.state_dict() + checkpoint = {key.removeprefix(prefix): value for key, value in state.items() if key.startswith(prefix)} + elif source_format == "tensors": + checkpoint = Checkpoint(source_path, prefix=entry.get("input_prefix", ""), wrapper=wrapper) + else: + raise ValueError("Source format must be tensors, torchscript or python-model.") + sources.append((entry.get("output_prefix", ""), checkpoint)) + return MergedCheckpoint(sources) + + +def flatten_tensor_tree(weights, separator="."): + """Flatten nested original NumPy/JAX tensor dictionaries without changing tensor axes or values.""" + result = {} + + def visit(tree, prefix): + for name, value in tree.items(): + key = prefix + str(name) + if isinstance(value, Mapping): + visit(value, key + separator) + else: + if key in result: + raise ValueError(f"Duplicate flattened tensor: {key}.") + result[key] = torch.as_tensor(value) + + visit(weights, "") + return result + + +def load_tensor_sources(path): + """Read a tensor file, indexed component directory, or directory of disjoint safetensors shards.""" + from .io import Checkpoint + + path = Path(path) + if not path.is_dir(): + return Checkpoint(path) + if any(path.glob("*.index.json")) or any( + (path / name).is_file() + for name in ( + "model.safetensors", + "diffusion_pytorch_model.safetensors", + "pytorch_model.bin", + "diffusion_pytorch_model.bin", + ) + ): + return Checkpoint(path) + return MergedCheckpoint(("", Checkpoint(shard)) for shard in sorted(path.glob("*.safetensors"))) diff --git a/src/diffusers/loaders/conversion/sparse_controlnet.py b/src/diffusers/loaders/conversion/sparse_controlnet.py new file mode 100644 index 000000000000..6dd72b8abd5e --- /dev/null +++ b/src/diffusers/loaders/conversion/sparse_controlnet.py @@ -0,0 +1,64 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .animatediff import animatediff_conversion +from .controlnet import controlnet_conversion +from .core import Conversion, Rule + + +def sparse_controlnet_conversion(config): + spatial_config = { + **config, + "down_block_types": [kind.replace("Motion", "2D") for kind in config["down_block_types"]], + "conditioning_embedding_out_channels": config["conditioning_embedding_out_channels"] or (16, 32, 96, 256), + } + base = controlnet_conversion(spatial_config) + mapping = { + key: key + for key in sorted(base.diffusers_keys) + if not (config["use_simplified_condition_embedding"] and key.startswith("controlnet_cond_embedding.")) + } + if config["use_simplified_condition_embedding"]: + mapping.update( + {f"controlnet_cond_embedding.{p}": f"controlnet_cond_embedding.{p}" for p in ("weight", "bias")} + ) + if config.get("transformer_layers_per_mid_block") is not None: + mid = controlnet_conversion( + {**spatial_config, "transformer_layers_per_block": config["transformer_layers_per_mid_block"]} + ) + mapping = {old: new for old, new in mapping.items() if not new.startswith("mid_block.")} + mapping.update({key: key for key in sorted(mid.diffusers_keys) if key.startswith("mid_block.")}) + motion = animatediff_conversion( + { + "block_out_channels": config["block_out_channels"], + "motion_layers_per_block": config["layers_per_block"], + "motion_transformer_layers_per_block": config["temporal_transformer_layers_per_block"], + "use_motion_mid_block": False, + "conv_in_channels": None, + } + ) + mapping.update( + { + old: new + for old, new in motion.mapping.items() + if new.startswith("down_blocks.") and ".attn2." not in new and ".norm2." not in new + } + ) + rules = tuple( + Rule((rule.original[0],), rule.diffusers) + for rule in motion.rules + if rule.diffusers[0].startswith("down_blocks.") + ) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/spectrogram_continuous.py b/src/diffusers/loaders/conversion/spectrogram_continuous.py new file mode 100644 index 000000000000..4a199b368e27 --- /dev/null +++ b/src/diffusers/loaders/conversion/spectrogram_continuous.py @@ -0,0 +1,25 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .spectrogram_notes import _spectrogram_encoder_rules +from .transforms import Permute + + +def spectrogram_continuous_conversion(config): + mapping, rules = _spectrogram_encoder_rules(config) + return Conversion( + mapping=mapping, rules=rules + (Rule(("input_proj.kernel",), ("input_proj.weight",), Permute((1, 0))),) + ) diff --git a/src/diffusers/loaders/conversion/spectrogram_notes.py b/src/diffusers/loaders/conversion/spectrogram_notes.py new file mode 100644 index 000000000000..36df164b2c9b --- /dev/null +++ b/src/diffusers/loaders/conversion/spectrogram_notes.py @@ -0,0 +1,40 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Permute + + +def spectrogram_notes_conversion(config): + mapping, rules = _spectrogram_encoder_rules(config) + mapping["token_embedder.embedding"] = "token_embedder.weight" + return Conversion(mapping=mapping, rules=rules) + + +def _spectrogram_encoder_rules(config): + mapping = {"Embed_0.embedding": "position_encoding.weight", "encoder_norm.scale": "layer_norm.weight"} + modules = [] + for i in range(config["num_layers"]): + old, new = f"layers_{i}", f"encoders.{i}" + mapping[old + ".pre_attention_layer_norm.scale"] = new + ".layer.0.layer_norm.weight" + mapping[old + ".pre_mlp_layer_norm.scale"] = new + ".layer.1.layer_norm.weight" + modules.extend( + (f"{old}.attention.{a}", f"{new}.layer.0.SelfAttention.{b}") + for a, b in (("query", "q"), ("key", "k"), ("value", "v"), ("out", "o")) + ) + names = ("wi_0", "wi_1", "wo") if config["feed_forward_proj"].startswith("gated-") else ("wi", "wo") + modules.extend((f"{old}.mlp.{name}", f"{new}.layer.1.DenseReluDense.{name}") for name in names) + rules = tuple(Rule((old + ".kernel",), (new + ".weight",), Permute((1, 0))) for old, new in modules) + return mapping, rules diff --git a/src/diffusers/loaders/conversion/stable_audio.py b/src/diffusers/loaders/conversion/stable_audio.py new file mode 100644 index 000000000000..9b868a82797e --- /dev/null +++ b/src/diffusers/loaders/conversion/stable_audio.py @@ -0,0 +1,65 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Reshape, Split + + +def stable_audio_conversion(config): + mapping = { + "preprocess_conv.weight": "preprocess_conv.weight", + "postprocess_conv.weight": "postprocess_conv.weight", + "transformer.project_in.weight": "proj_in.weight", + "transformer.project_out.weight": "proj_out.weight", + } + modules = [("to_timestep_embed.0", "timestep_proj.0"), ("to_timestep_embed.2", "timestep_proj.2")] + for source, target in (("to_global_embed", "global_proj"), ("to_cond_embed", "cross_attention_proj")): + mapping.update({f"{source}.{i}.weight": f"{target}.{i}.weight" for i in (0, 2)}) + fourier = config["time_proj_dim"] // 2 + rules = [Rule(("timestep_features.weight",), ("time_proj.weight",), Reshape((fourier, 1), (fourier,)))] + hidden = config["num_attention_heads"] * config["attention_head_dim"] + kv = config["num_key_value_attention_heads"] * config["attention_head_dim"] + for i in range(config["num_layers"]): + old, new = f"transformer.layers.{i}", f"transformer_blocks.{i}" + for a, b in (("pre_norm", "norm1"), ("cross_attend_norm", "norm2"), ("ff_norm", "norm3")): + mapping[f"{old}.{a}.gamma"] = f"{new}.{b}.weight" + mapping[f"{old}.{a}.beta"] = f"{new}.{b}.bias" + rules.append( + Rule( + (old + ".self_attn.to_qkv.weight",), + tuple(f"{new}.attn1.to_{part}.weight" for part in ("q", "k", "v")), + Split((hidden,) * 3), + ) + ) + rules.append( + Rule( + (old + ".cross_attn.to_kv.weight",), + tuple(f"{new}.attn2.to_{part}.weight" for part in ("k", "v")), + Split((kv,) * 2), + ) + ) + mapping.update( + { + f"{old}.{a}.weight": f"{new}.{b}.weight" + for a, b in ( + ("self_attn.to_out", "attn1.to_out.0"), + ("cross_attn.to_q", "attn2.to_q"), + ("cross_attn.to_out", "attn2.to_out.0"), + ) + } + ) + modules.extend((f"{old}.ff.ff.{name}", f"{new}.ff.net.{name}") for name in ("0.proj", "2")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/stable_audio3.py b/src/diffusers/loaders/conversion/stable_audio3.py new file mode 100644 index 000000000000..a1de1f766703 --- /dev/null +++ b/src/diffusers/loaders/conversion/stable_audio3.py @@ -0,0 +1,74 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import ReorderChunks + + +def stable_audio3_conversion(config): + mapping = { + "conditioner.conditioners.prompt.padding_embedding": "prompt_padding_embedding", + "model.model.transformer.rotary_pos_emb.inv_freq": "rotary_pos_emb.inv_freq", + } + modules = [] + for name in ("to_timestep_embed.0", "to_timestep_embed.2"): + modules.append(("model.model." + name, name)) + for name in ( + "to_cond_embed.0", + "to_cond_embed.2", + "to_global_embed.0", + "to_global_embed.2", + "preprocess_conv", + "postprocess_conv", + ): + mapping[f"model.model.{name}.weight"] = name + ".weight" + for i in (0, 2): + modules.append((f"model.model.transformer.global_cond_embedder.{i}", f"global_cond_embedder.{i}")) + for name in ("in", "out"): + mapping[f"model.model.transformer.project_{name}.weight"] = f"proj_{name}.weight" + if config["num_memory_tokens"] > 0: + mapping["model.model.transformer.memory_tokens"] = "memory_tokens" + rules = [] + for i in range(config["depth"]): + old, new = f"model.model.transformer.layers.{i}", f"transformer_blocks.{i}" + mapping[old + ".to_scale_shift_gate"] = new + ".to_scale_shift_gate" + for name in ( + "pre_norm", + "cross_attend_norm", + "ff_norm", + "self_attn.q_norm", + "self_attn.k_norm", + "cross_attn.q_norm", + "cross_attn.k_norm", + ): + mapping[f"{old}.{name}.gamma"] = f"{new}.{name}.gamma" + for name in ("self_attn.to_out", "cross_attn.to_q", "cross_attn.to_kv", "cross_attn.to_out"): + mapping[f"{old}.{name}.weight"] = f"{new}.{name}.weight" + source, target = old + ".self_attn.to_qkv.weight", new + ".self_attn.to_qkv.weight" + if config["use_differential_attention"]: + rules.append(Rule((source,), (target,), ReorderChunks((0, 3, 1, 4, 2)))) + else: + mapping[source] = target + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ff.ff.0.proj", "ff.proj_in"), + ("ff.ff.2", "ff.proj_out"), + ("to_local_embed.0", "to_local_embed.0"), + ("to_local_embed.2", "to_local_embed.2"), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/stable_audio3_duration.py b/src/diffusers/loaders/conversion/stable_audio3_duration.py new file mode 100644 index 000000000000..cfa33fc08e32 --- /dev/null +++ b/src/diffusers/loaders/conversion/stable_audio3_duration.py @@ -0,0 +1,33 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +import torch + +from .core import Conversion, Rule +from .transforms import WithConstants + + +def stable_audio3_duration_conversion(config): + prefix = "conditioner.conditioners.seconds_total.embedder.embedding.1" + ramp = torch.linspace(0, 1, config["fourier_dim"] // 2, dtype=torch.float32) + freqs = torch.exp( + ramp * (math.log(config["max_freq"]) - math.log(config["min_freq"])) + math.log(config["min_freq"]) + ) + return Conversion( + mapping={prefix + ".bias": "linear.bias"}, + rules=(Rule((prefix + ".weight",), ("linear.weight", "freqs"), WithConstants((freqs,))),), + ) diff --git a/src/diffusers/loaders/conversion/stable_audio_projection.py b/src/diffusers/loaders/conversion/stable_audio_projection.py new file mode 100644 index 000000000000..a4a999f3317f --- /dev/null +++ b/src/diffusers/loaders/conversion/stable_audio_projection.py @@ -0,0 +1,28 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def stable_audio_projection_conversion(config): + mapping = {} + for old, new in (("seconds_start", "start_number_conditioner"), ("seconds_total", "end_number_conditioner")): + mapping[f"{old}.embedder.embedding.0.weights"] = f"{new}.time_positional_embedding.0.weights" + mapping.update( + {f"{old}.embedder.embedding.1.{p}": f"{new}.time_positional_embedding.1.{p}" for p in ("weight", "bias")} + ) + if config["text_encoder_dim"] != config["conditioning_dim"]: + mapping.update({f"prompt.proj_out.{p}": f"text_projection.{p}" for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/stable_cascade.py b/src/diffusers/loaders/conversion/stable_cascade.py new file mode 100644 index 000000000000..7ee43e70cfd0 --- /dev/null +++ b/src/diffusers/loaders/conversion/stable_cascade.py @@ -0,0 +1,78 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def stable_cascade_conversion(config): + modules = [("embedding.1", "embedding.1"), ("clf.1", "clf.1")] + pooled = "clip_txt_pooled_mapper" if config["clip_text_in_channels"] is not None else "clip_mapper" + modules.append((pooled, "clip_txt_pooled_mapper")) + for field, name in (("clip_text_in_channels", "clip_txt_mapper"), ("clip_image_in_channels", "clip_img_mapper")): + if config[field] is not None: + modules.append((name, name)) + for field, name in (("effnet_in_channels", "effnet_mapper"), ("pixel_mapper_in_channels", "pixels_mapper")): + if config[field] is not None: + modules.extend((f"{name}.{i}", f"{name}.{i}") for i in (0, 2)) + mapping, rules = {}, [] + count = len(config["block_out_channels"]) + for direction in ("down", "up"): + for i in range(count): + config_index = i if direction == "down" else count - 1 - i + channels = config["block_out_channels"][config_index] + if (direction == "down" and i > 0) or (direction == "up" and i < count - 1): + name = f"{direction}_{'downscalers' if direction == 'down' else 'upscalers'}.{i}.1" + if config["switch_level"] is not None: + name += ".blocks.0" if direction == "down" else ".blocks.1" + modules.append((name, name)) + repeats = config[f"{direction}_blocks_repeat_mappers"] + for j in range(repeats[i] - 1): + name = f"{direction}_repeat_mappers.{i}.{j}" + modules.append((name, name)) + block_types = config["block_types_per_layer"][config_index] + for j in range(config[f"{direction}_num_layers_per_block"][i]): + for k, block_type in enumerate(block_types): + prefix = f"{direction}_blocks.{i}.{j * len(block_types) + k}" + if block_type == "SDCascadeResBlock": + modules.extend( + (f"{prefix}.{name}", f"{prefix}.{name}") + for name in ("depthwise", "channelwise.0", "channelwise.4") + ) + for name in ("gamma", "beta"): + key = f"{prefix}.channelwise.2.{name}" + mapping[key] = key + elif block_type == "SDCascadeTimestepBlock": + for name in ["mapper"] + [f"mapper_{cond}" for cond in config["timestep_conditioning_type"]]: + modules.append((f"{prefix}.{name}", f"{prefix}.{name}")) + elif block_type == "SDCascadeAttnBlock": + modules.extend( + [ + (prefix + ".kv_mapper.1", prefix + ".kv_mapper.1"), + (prefix + ".attention.attn.out_proj", prefix + ".attention.to_out.0"), + ] + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attention.attn.in_proj_{p}",), + tuple(f"{prefix}.attention.to_{part}.{p}" for part in ("q", "k", "v")), + Split((channels,) * 3), + ) + ) + else: + raise ValueError(f"Unknown Stable Cascade block type {block_type}.") + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/svd.py b/src/diffusers/loaders/conversion/svd.py new file mode 100644 index 000000000000..a46efb3d2c67 --- /dev/null +++ b/src/diffusers/loaders/conversion/svd.py @@ -0,0 +1,76 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_unet import ldm_unet_conversion + + +def svd_conversion(config): + spatial_config = { + **config, + "down_block_types": [name.replace("SpatioTemporal", "2D") for name in config["down_block_types"]], + "up_block_types": [name.replace("SpatioTemporal", "2D") for name in config["up_block_types"]], + "addition_embed_type": "text_time", + } + base = ldm_unet_conversion(spatial_config) + mapping, resnets, attention_prefixes = {}, [], set() + for old, new in base.mapping.items(): + if ".resnets." in new: + prefix, rest = new.split(".resnets.") + index, suffix = rest.split(".", 1) + parent = f"{prefix}.resnets.{index}" + mapping[old] = parent + ".spatial_res_block." + suffix + if suffix == "norm1.weight": + resnets.append((old.removesuffix(".in_layers.0.weight"), parent)) + else: + mapping[old] = new + if ".attentions." in new and ".transformer_blocks." in new: + mapping[old.replace(".transformer_blocks.", ".time_stack.")] = new.replace( + ".transformer_blocks.", ".temporal_transformer_blocks." + ) + if ".attentions." in new and new.endswith(".proj_in.bias"): + attention_prefixes.add((old.removesuffix(".proj_in.bias"), new.removesuffix(".proj_in.bias"))) + for old, new in resnets: + mapping[old + ".time_mixer.mix_factor"] = new + ".time_mixer.mix_factor" + for a, b in ( + ("in_layers.0", "norm1"), + ("in_layers.2", "conv1"), + ("out_layers.0", "norm2"), + ("out_layers.3", "conv2"), + ("emb_layers.1", "time_emb_proj"), + ): + mapping.update( + {f"{old}.time_stack.{a}.{p}": f"{new}.temporal_res_block.{b}.{p}" for p in ("weight", "bias")} + ) + for old, new in sorted(attention_prefixes): + mapping[old + ".time_mixer.mix_factor"] = new + ".time_mixer.mix_factor" + for i, j in ((0, 1), (2, 2)): + mapping.update( + {f"{old}.time_pos_embed.{i}.{p}": f"{new}.time_pos_embed.linear_{j}.{p}" for p in ("weight", "bias")} + ) + indices = { + int(key.removeprefix(new + ".transformer_blocks.").split(".")[0]) + for key in base.diffusers_keys + if key.startswith(new + ".transformer_blocks.") + } + for i in sorted(indices): + mapping.update( + { + f"{old}.time_stack.{i}.{name}.{p}": f"{new}.temporal_transformer_blocks.{i}.{name}.{p}" + for name in ("norm_in", "ff_in.net.0.proj", "ff_in.net.2") + for p in ("weight", "bias") + } + ) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/svd_vae.py b/src/diffusers/loaders/conversion/svd_vae.py new file mode 100644 index 000000000000..f35ca8d30708 --- /dev/null +++ b/src/diffusers/loaders/conversion/svd_vae.py @@ -0,0 +1,78 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .ldm_vae import ldm_vae_conversion + + +def svd_vae_conversion(config): + original_format = config.get("original_format", "svd") + if original_format not in ("svd", "temporal_vae"): + raise ValueError("Temporal VAE original_format must be 'svd' or 'temporal_vae'.") + encoder_prefix = "conditioner.embedders.3.encoder." if original_format == "svd" else "" + decoder_prefix = "first_stage_model." if original_format == "svd" else "" + base = ldm_vae_conversion( + { + **config, + "up_block_types": ["UpDecoderBlock2D"] * len(config["block_out_channels"]), + "mid_block_add_attention": True, + "use_post_quant_conv": False, + } + ) + mapping, rules, resnets = {}, [], [] + for old, new in base.mapping.items(): + prefix = decoder_prefix if new.startswith("decoder.") else encoder_prefix + if new.startswith("decoder.mid_block.resnets."): + continue + if new.startswith("decoder.") and ".resnets." in new: + parent, suffix = new.rsplit(".resnets.", 1) + index, suffix = suffix.split(".", 1) + parent += ".resnets." + index + mapping[prefix + old] = parent + ".spatial_res_block." + suffix + if suffix == "norm1.weight": + resnets.append((prefix + old.removesuffix(".norm1.weight"), parent)) + else: + mapping[prefix + old] = new + for i in range(config["layers_per_block"]): + old, new = f"{decoder_prefix}decoder.mid.block_{i + 1}", f"decoder.mid_block.resnets.{i}" + resnets.append((old, new)) + mapping.update( + { + f"{old}.{name}.{p}": f"{new}.spatial_res_block.{name}.{p}" + for name in ("norm1", "norm2", "conv1", "conv2") + for p in ("weight", "bias") + } + ) + for rule in base.rules: + prefix = decoder_prefix if rule.diffusers[0].startswith("decoder.") else encoder_prefix + rules.append(Rule(tuple(prefix + key for key in rule.original), rule.diffusers, rule.transform)) + for old, new in resnets: + mapping[old + ".mix_factor"] = new + ".time_mixer.mix_factor" + for a, b in ( + ("in_layers.0", "norm1"), + ("in_layers.2", "conv1"), + ("out_layers.0", "norm2"), + ("out_layers.3", "conv2"), + ): + mapping.update( + {f"{old}.time_stack.{a}.{p}": f"{new}.temporal_res_block.{b}.{p}" for p in ("weight", "bias")} + ) + mapping.update( + { + f"{decoder_prefix}decoder.conv_out.time_mix_conv.{p}": f"decoder.time_conv_out.{p}" + for p in ("weight", "bias") + } + ) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/t2i_adapter.py b/src/diffusers/loaders/conversion/t2i_adapter.py new file mode 100644 index 000000000000..2897caf9407e --- /dev/null +++ b/src/diffusers/loaders/conversion/t2i_adapter.py @@ -0,0 +1,39 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def t2i_adapter_conversion(config): + modules = [] + channels, depth = config["channels"], config["num_res_blocks"] + light = config["adapter_type"] == "light_adapter" + if not light: + modules.append(("conv_in", "adapter.conv_in")) + count = len(channels) + int(light) + for i in range(count): + if light: + modules.extend((f"body.{i}.{name}", f"adapter.body.{i}.{name}") for name in ("in_conv", "out_conv")) + else: + input_channel = channels[i - 1] if i > 0 else channels[0] + if config["adapter_type"] == "full_adapter_xl" and i not in (1, 2): + input_channel = channels[i] + if input_channel != channels[i]: + modules.append((f"body.{i * depth}.in_conv", f"adapter.body.{i}.in_conv")) + for j in range(depth): + old = f"body.{i}.body.{j}" if light else f"body.{i * depth + j}" + new = f"adapter.body.{i}.resnets.{j}" + modules.extend((f"{old}.{name}", f"{new}.{name}") for name in ("block1", "block2")) + return Conversion(mapping={f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) diff --git a/src/diffusers/loaders/conversion/t5.py b/src/diffusers/loaders/conversion/t5.py new file mode 100644 index 000000000000..e04796c43b09 --- /dev/null +++ b/src/diffusers/loaders/conversion/t5.py @@ -0,0 +1,29 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def t5_conversion(config): + keys = ["shared.weight", "encoder.embed_tokens.weight", "encoder.final_layer_norm.weight"] + for i in range(config["num_layers"]): + prefix = f"encoder.block.{i}" + keys.extend(f"{prefix}.layer.0.SelfAttention.{part}.weight" for part in ("q", "k", "v", "o")) + keys.extend(f"{prefix}.layer.{j}.layer_norm.weight" for j in (0, 1)) + if i == 0: + keys.append(prefix + ".layer.0.SelfAttention.relative_attention_bias.weight") + names = ("wi_0", "wi_1", "wo") if config["feed_forward_proj"].startswith("gated-") else ("wi", "wo") + keys.extend(f"{prefix}.layer.1.DenseReluDense.{name}.weight" for name in names) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/t5_film.py b/src/diffusers/loaders/conversion/t5_film.py new file mode 100644 index 000000000000..817264007e7b --- /dev/null +++ b/src/diffusers/loaders/conversion/t5_film.py @@ -0,0 +1,49 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Permute + + +def t5_film_conversion(config): + mapping = {"Embed_0.embedding": "position_encoding.weight", "decoder_norm.scale": "decoder_norm.weight"} + modules = [ + ("time_emb_dense0", "conditioning_emb.0"), + ("time_emb_dense1", "conditioning_emb.2"), + ("continuous_inputs_projection", "continuous_inputs_projection"), + ("spec_out_dense", "spec_out"), + ] + for i in range(config["num_layers"]): + old, new = f"layers_{i}", f"decoders.{i}" + for j, norm in enumerate( + ("pre_self_attention_layer_norm", "pre_cross_attention_layer_norm", "pre_mlp_layer_norm") + ): + mapping[f"{old}.{norm}.scale"] = f"{new}.layer.{j}.layer_norm.weight" + for j, attn in enumerate(("self_attention", "MultiHeadDotProductAttention_0")): + modules.extend( + (f"{old}.{attn}.{a}", f"{new}.layer.{j}.attention.{b}") + for a, b in (("query", "to_q"), ("key", "to_k"), ("value", "to_v"), ("out", "to_out.0")) + ) + modules.extend( + [ + (old + ".FiLMLayer_0.DenseGeneral_0", new + ".layer.0.FiLMLayer.scale_bias"), + (old + ".FiLMLayer_1.DenseGeneral_0", new + ".layer.2.film.scale_bias"), + ] + ) + modules.extend( + (f"{old}.mlp.{name}", f"{new}.layer.2.DenseReluDense.{name}") for name in ("wi_0", "wi_1", "wo") + ) + rules = tuple(Rule((old + ".kernel",), (new + ".weight",), Permute((1, 0))) for old, new in modules) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/tiny_vae.py b/src/diffusers/loaders/conversion/tiny_vae.py new file mode 100644 index 000000000000..8ee251dd8bf3 --- /dev/null +++ b/src/diffusers/loaders/conversion/tiny_vae.py @@ -0,0 +1,48 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def tiny_vae_conversion(config): + modules = [] + index = 0 + for i, count in enumerate(config["num_encoder_blocks"]): + modules.append((f"encoder.{index}", f"encoder.layers.{index}", i == 0)) + index += 1 + for _ in range(count): + modules.extend((f"encoder.{index}.conv.{j}", f"encoder.layers.{index}.conv.{j}", True) for j in (0, 2, 4)) + index += 1 + modules.append((f"encoder.{index}", f"encoder.layers.{index}", True)) + modules.append(("decoder.1", "decoder.layers.0", True)) + index = 2 + for i, count in enumerate(config["num_decoder_blocks"]): + for _ in range(count): + modules.extend( + (f"decoder.{index + 1}.conv.{j}", f"decoder.layers.{index}.conv.{j}", True) for j in (0, 2, 4) + ) + index += 1 + final = i == len(config["num_decoder_blocks"]) - 1 + if not final: + index += 1 + modules.append((f"decoder.{index + 1}", f"decoder.layers.{index}", final)) + index += 1 + return Conversion( + mapping={ + f"{old}.{p}": f"{new}.{p}" + for old, new, bias in modules + for p in (("weight", "bias") if bias else ("weight",)) + } + ) diff --git a/src/diffusers/loaders/conversion/transforms.py b/src/diffusers/loaders/conversion/transforms.py new file mode 100644 index 000000000000..164d2d25a190 --- /dev/null +++ b/src/diffusers/loaders/conversion/transforms.py @@ -0,0 +1,327 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from math import prod +from typing import Protocol + +import torch + + +class Transform(Protocol): + """A reversible operation on an ordered group of tensors. + + Implementations must preserve dtype/device and must not modify their inputs. Outputs may share storage with inputs. + Configuration needed by either direction belongs on the transform, so export does not depend on an earlier import. + """ + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: ... + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: ... + + +@dataclass(frozen=True) +class Identity: + """Return the tensors unchanged, including their storage.""" + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + return tensors + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + return tensors + + +@dataclass(frozen=True) +class FoldLinearGate: + """Fold a static output-channel gate into a linear layer, exporting a canonical unit gate. + + Original gate factorization is lost. Diffusers -> original -> Diffusers preserves float32/float16/bfloat16 weights, + while the other round trip produces equivalent folded parameters. Multiplication follows the source importer's + float32 arithmetic. + """ + + bias: bool = True + lossless = False + + def forward(self, tensors): + if len(tensors) != 2 + int(self.bias): + raise ValueError("Expected gate, weight, and optional bias.") + gate, weight, *bias = tensors + if weight.ndim != 2 or gate.shape != weight.shape[:1]: + raise ValueError("Gate must be a vector matching the linear output channels.") + if bias and bias[0].shape != gate.shape: + raise ValueError("Bias must match the gate shape.") + if any(t.device != weight.device or t.dtype != weight.dtype for t in tensors): + raise ValueError("Gate and linear parameters must have the same dtype and device.") + if not weight.is_floating_point(): + raise ValueError("Gate folding requires floating-point parameters.") + outputs = [(gate.float().unsqueeze(1) * weight.float()).to(weight.dtype)] + if bias: + outputs.append((gate.float() * bias[0].float()).to(weight.dtype)) + return tuple(outputs) + + def inverse(self, tensors): + if len(tensors) != 1 + int(self.bias) or tensors[0].ndim != 2: + raise ValueError("Expected linear weight and optional bias.") + weight = tensors[0] + if self.bias and ( + tensors[1].shape != weight.shape[:1] + or tensors[1].dtype != weight.dtype + or tensors[1].device != weight.device + ): + raise ValueError("Linear weight and bias must have matching shape, dtype and device.") + return (weight.new_ones(weight.shape[0]), *tensors) + + +@dataclass(frozen=True) +class Split: + """Split one tensor into explicit sizes along `dim`; concatenate those pieces in the inverse direction.""" + + sizes: tuple[int, ...] + dim: int = 0 + + def __post_init__(self): + object.__setattr__(self, "sizes", tuple(self.sizes)) + if not self.sizes or any(size <= 0 for size in self.sizes): + raise ValueError("Split sizes must be positive.") + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + if len(tensors) != 1: + raise ValueError(f"Split expects one tensor, got {len(tensors)}.") + (tensor,) = tensors + if not -tensor.ndim <= self.dim < tensor.ndim: + raise ValueError(f"Invalid split dimension {self.dim} for shape {tuple(tensor.shape)}.") + if tensor.shape[self.dim] != sum(self.sizes): + raise ValueError(f"Split sizes {self.sizes} do not match dimension {self.dim} of {tuple(tensor.shape)}.") + return tensor.split(self.sizes, dim=self.dim) + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + if len(tensors) != len(self.sizes): + raise ValueError(f"Expected {len(self.sizes)} split pieces, got {len(tensors)}.") + first = tensors[0] + if not -first.ndim <= self.dim < first.ndim: + raise ValueError(f"Invalid split dimension {self.dim} for shape {tuple(first.shape)}.") + dim = self.dim % first.ndim + for tensor, size in zip(tensors, self.sizes): + expected_shape = list(first.shape) + expected_shape[dim] = size + if tuple(tensor.shape) != tuple(expected_shape): + raise ValueError(f"Expected split piece shape {tuple(expected_shape)}, got {tuple(tensor.shape)}.") + if tensor.dtype != first.dtype or tensor.device != first.device: + raise ValueError( + "Split pieces must have the same dtype and device; implicit casting is not supported." + ) + return (torch.cat(tensors, dim=dim),) + + +@dataclass(frozen=True) +class ReorderChunks: + """Permute equal-sized chunks of one tensor. `order[i]` identifies the source chunk at output position `i`.""" + + order: tuple[int, ...] + dim: int = 0 + + def __post_init__(self): + object.__setattr__(self, "order", tuple(self.order)) + if not self.order or sorted(self.order) != list(range(len(self.order))): + raise ValueError("Chunk order must be a permutation of consecutive indices starting at zero.") + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + if len(tensors) != 1: + raise ValueError(f"ReorderChunks expects one tensor, got {len(tensors)}.") + (tensor,) = tensors + if not -tensor.ndim <= self.dim < tensor.ndim: + raise ValueError(f"Invalid chunk dimension {self.dim} for shape {tuple(tensor.shape)}.") + if tensor.shape[self.dim] == 0 or tensor.shape[self.dim] % len(self.order): + raise ValueError( + f"Cannot divide dimension {self.dim} of {tuple(tensor.shape)} into {len(self.order)} chunks." + ) + chunks = tensor.chunk(len(self.order), dim=self.dim) + return (torch.cat(tuple(chunks[i] for i in self.order), dim=self.dim),) + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + inverse_order = tuple(self.order.index(i) for i in range(len(self.order))) + return ReorderChunks(inverse_order, self.dim).forward(tensors) + + +@dataclass(frozen=True) +class Reshape: + """Change a tensor's shape with both shapes specified, including singleton dimensions.""" + + original_shape: tuple[int, ...] + diffusers_shape: tuple[int, ...] + + def __post_init__(self): + object.__setattr__(self, "original_shape", tuple(self.original_shape)) + object.__setattr__(self, "diffusers_shape", tuple(self.diffusers_shape)) + if any(size < 0 for size in self.original_shape + self.diffusers_shape): + raise ValueError("Reshape requires explicit, nonnegative dimensions.") + if prod(self.original_shape) != prod(self.diffusers_shape): + raise ValueError("Reshape must preserve the number of elements.") + + def forward(self, tensors): + (tensor,) = tensors + if tuple(tensor.shape) != self.original_shape: + raise ValueError(f"Expected shape {self.original_shape}, got {tuple(tensor.shape)}.") + return (tensor.reshape(self.diffusers_shape),) + + def inverse(self, tensors): + return Reshape(self.diffusers_shape, self.original_shape).forward(tensors) + + +@dataclass(frozen=True) +class Squeeze: + """Remove one declared singleton axis and restore it on export, preserving all other dimensions.""" + + dim: int + ndim: int + + def __post_init__(self): + if not 0 <= self.dim < self.ndim: + raise ValueError("The singleton dimension must be within the declared tensor rank.") + + def forward(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + (tensor,) = tensors + if tensor.ndim != self.ndim or tensor.shape[self.dim] != 1: + raise ValueError( + f"Expected rank {self.ndim} with a singleton axis at {self.dim}, got {tuple(tensor.shape)}." + ) + return (tensor.squeeze(self.dim),) + + def inverse(self, tensors: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + (tensor,) = tensors + if tensor.ndim != self.ndim - 1: + raise ValueError(f"Expected rank {self.ndim - 1}, got {tensor.ndim}.") + return (tensor.unsqueeze(self.dim),) + + +@dataclass(frozen=True) +class Permute: + """Permute a tensor's axes; reverse with the inverse axis permutation.""" + + dims: tuple[int, ...] + + def __post_init__(self): + object.__setattr__(self, "dims", tuple(self.dims)) + if sorted(self.dims) != list(range(len(self.dims))): + raise ValueError("Axes must be a permutation of consecutive indices starting at zero.") + + def forward(self, tensors): + (tensor,) = tensors + if tensor.ndim != len(self.dims): + raise ValueError(f"Expected rank {len(self.dims)}, got {tensor.ndim}.") + return (tensor.permute(self.dims),) + + def inverse(self, tensors): + return Permute(tuple(self.dims.index(i) for i in range(len(self.dims)))).forward(tensors) + + +@dataclass(frozen=True) +class Chain: + """Compose reversible operations, reversing their order for export.""" + + transforms: tuple[Transform, ...] + + def __post_init__(self): + object.__setattr__(self, "transforms", tuple(self.transforms)) + + @property + def lossless(self): + return all(getattr(transform, "lossless", True) for transform in self.transforms) + + def forward(self, tensors): + for transform in self.transforms: + tensors = transform.forward(tensors) + return tensors + + def inverse(self, tensors): + for transform in reversed(self.transforms): + tensors = transform.inverse(tensors) + return tensors + + +@dataclass(frozen=True) +class MergeEqual: + """Merge equal copies of a buffer; restore all copies on export without losing information.""" + + copies: int + + def __post_init__(self): + if self.copies < 1: + raise ValueError("The number of copies must be positive.") + + def forward(self, tensors): + if len(tensors) != self.copies: + raise ValueError(f"Expected {self.copies} equal tensors, got {len(tensors)}.") + first = tensors[0] + for tensor in tensors[1:]: + if tensor.dtype != first.dtype or tensor.device != first.device or not torch.equal(tensor, first): + raise ValueError("Cannot merge unequal tensors without losing information.") + return (first,) + + def inverse(self, tensors): + (tensor,) = tensors + return (tensor,) * self.copies + + +@dataclass(frozen=True) +class Reverse: + """Use an existing transform in the opposite direction.""" + + transform: Transform + + @property + def lossless(self): + return getattr(self.transform, "lossless", True) + + def forward(self, tensors): + return self.transform.inverse(tensors) + + def inverse(self, tensors): + return self.transform.forward(tensors) + + +@dataclass(frozen=True, eq=False) +class WithConstants: + """Attach config-derived buffers to one tensor and verify those buffers before removing them on export. + + Floating buffers use the anchor tensor's dtype; integer and boolean buffers retain their own dtype. + """ + + values: tuple[torch.Tensor, ...] + + def __post_init__(self): + object.__setattr__(self, "values", tuple(value.detach().clone() for value in self.values)) + + def forward(self, tensors): + (anchor,) = tensors + return (anchor,) + tuple( + value.to(device=anchor.device, dtype=anchor.dtype if value.is_floating_point() else value.dtype) + for value in self.values + ) + + def inverse(self, tensors): + if len(tensors) != len(self.values) + 1: + raise ValueError(f"Expected an anchor and {len(self.values)} constant buffers.") + expected = self.forward((tensors[0],)) + for actual, reference in zip(tensors[1:], expected[1:]): + if ( + actual.dtype != reference.dtype + or actual.device != reference.device + or not torch.equal(actual, reference) + ): + raise ValueError( + "A config-derived buffer was changed and cannot be represented in the original format." + ) + return (tensors[0],) diff --git a/src/diffusers/loaders/conversion/umt5.py b/src/diffusers/loaders/conversion/umt5.py new file mode 100644 index 000000000000..3940a5bb9c3e --- /dev/null +++ b/src/diffusers/loaders/conversion/umt5.py @@ -0,0 +1,25 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .t5 import t5_conversion + + +def umt5_conversion(config): + mapping = dict(t5_conversion(config).mapping) + for i in range(1, config["num_layers"]): + key = f"encoder.block.{i}.layer.0.SelfAttention.relative_attention_bias.weight" + mapping[key] = key + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/unclip_text_projection.py b/src/diffusers/loaders/conversion/unclip_text_projection.py new file mode 100644 index 000000000000..3ef0a9f34f79 --- /dev/null +++ b/src/diffusers/loaders/conversion/unclip_text_projection.py @@ -0,0 +1,29 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def unclip_text_projection_conversion(config): + modules = [ + ("text_seq_proj.0", "encoder_hidden_states_proj"), + ("text_seq_proj.1", "text_encoder_hidden_states_norm"), + ("clip_tok_proj", "clip_extra_context_tokens_proj"), + ("text_feat_proj", "embedding_proj"), + ("clip_emb", "clip_image_embeddings_project_to_time_embeddings"), + ] + mapping = {"cf_param": "learned_classifier_free_guidance_embeddings"} + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/unclip_unet.py b/src/diffusers/loaders/conversion/unclip_unet.py new file mode 100644 index 000000000000..eb8ed5e182e8 --- /dev/null +++ b/src/diffusers/loaders/conversion/unclip_unet.py @@ -0,0 +1,74 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Permute, Reshape, Split +from .unet_2d import unet_2d_conversion + + +def unclip_unet_conversion(config): + base = unet_2d_conversion( + {**config, "original_format": "ldm", "downsample_type": "resnet", "upsample_type": "resnet"} + ) + mapping, rules = dict(base.mapping), list(base.rules) + modules = [] + if config.get("addition_embed_type") == "text_image": + modules.extend( + [ + ("ln_model_n", "add_embedding.text_norm"), + ("proj_n", "add_embedding.text_proj"), + ("img_layer", "add_embedding.image_proj"), + ] + ) + if config.get("encoder_hid_dim_type") == "text_image_proj": + modules.extend( + [("clip_to_seq", "encoder_hid_proj.image_embeds"), ("to_model_dim_n", "encoder_hid_proj.text_proj")] + ) + for rule in base.rules: + source = rule.original[0] + if not source.endswith(".qkv.weight"): + continue + old = source.removesuffix(".qkv.weight") + new = rule.diffusers[0].removesuffix(".to_q.weight") + parts = new.split(".") + if parts[0] == "mid_block": + index = len(config["block_out_channels"]) - 1 + elif parts[0] == "up_blocks": + index = len(config["block_out_channels"]) - 1 - int(parts[1]) + else: + index = int(parts[1]) + channel = config["block_out_channels"][index] + dim = config["attention_head_dim"] + dim = dim[index] if isinstance(dim, (tuple, list)) else dim + heads = channel // dim + inner = heads * dim + context = config["cross_attention_dim"] + context = context[index] if isinstance(context, (tuple, list)) else context + for p in ("weight", "bias"): + trailing = (context,) if p == "weight" else () + original_shape = (2 * inner, context, 1) if p == "weight" else (2 * inner,) + transform = Chain( + ( + Reshape(original_shape, (heads, 2, dim) + trailing), + Permute((1, 0, 2, 3) if trailing else (1, 0, 2)), + Reshape((2, heads, dim) + trailing, (2 * inner,) + trailing), + Split((inner,) * 2), + ) + ) + rules.append( + Rule((f"{old}.encoder_kv.{p}",), (f"{new}.add_k_proj.{p}", f"{new}.add_v_proj.{p}"), transform) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/unet_1d.py b/src/diffusers/loaders/conversion/unet_1d.py new file mode 100644 index 000000000000..3dfd1c48ae18 --- /dev/null +++ b/src/diffusers/loaders/conversion/unet_1d.py @@ -0,0 +1,24 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .dance_diffusion import dance_diffusion_conversion + + +def unet_1d_conversion(config): + if config["down_block_types"][0] == "DownResnetBlock1D": + from .diffuser_rl import diffuser_rl_conversion + + return diffuser_rl_conversion(config) + return dance_diffusion_conversion(config) diff --git a/src/diffusers/loaders/conversion/unet_2d.py b/src/diffusers/loaders/conversion/unet_2d.py new file mode 100644 index 000000000000..35d0dc216700 --- /dev/null +++ b/src/diffusers/loaders/conversion/unet_2d.py @@ -0,0 +1,174 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Chain, Permute, Reshape, Split + + +def unet_2d_conversion(config): + original_format = config.get("original_format") + if original_format is None: + original_format = "ncsnpp" if config["time_embedding_type"] == "fourier" else "ddpm" + if original_format == "ncsnpp": + from .ncsnpp import ncsnpp_conversion + + return ncsnpp_conversion(config) + if original_format not in ("ddpm", "ldm", "consistency"): + raise ValueError("UNet2D original_format must be 'ddpm', 'ldm', 'consistency', or 'ncsnpp'.") + ddpm = original_format == "ddpm" + mapping, modules, resnets, attentions = {}, [], [], [] + modules.extend( + [ + ("temb.dense.0" if ddpm else "time_embed.0", "time_embedding.linear_1"), + ("temb.dense.1" if ddpm else "time_embed.2", "time_embedding.linear_2"), + ("conv_in" if ddpm else "input_blocks.0.0", "conv_in"), + ("norm_out" if ddpm else "out.0", "conv_norm_out"), + ("conv_out" if ddpm else "out.2", "conv_out"), + ] + ) + if config.get("num_class_embeds") is not None: + mapping["label_emb.weight"] = "class_embedding.weight" + channels, layers = config["block_out_channels"], config["layers_per_block"] + count = len(channels) + layers = [layers] * count if isinstance(layers, int) else layers + previous, index = channels[0], 1 + for i, channel in enumerate(channels): + attention = "Attn" in config["down_block_types"][i] + for j in range(layers[i]): + old = f"down.{i}.block.{j}" if ddpm else f"input_blocks.{index}.0" + resnets.append((old, f"down_blocks.{i}.resnets.{j}", previous != channel)) + if attention: + old = f"down.{i}.attn.{j}" if ddpm else f"input_blocks.{index}.1" + attentions.append((old, f"down_blocks.{i}.attentions.{j}", channel)) + previous = channel + index += 1 + if i < count - 1: + old = f"down.{i}.downsample" if ddpm else f"input_blocks.{index}.0" + new = f"down_blocks.{i}.downsamplers.0" + residual = ( + config["down_block_types"][i] == "ResnetDownsampleBlock2D" or config.get("downsample_type") == "resnet" + ) + if residual: + resnets.append((old, new, False)) + else: + modules.append((old + ".op", new + ".conv")) + index += 1 + if config.get("mid_block_type", "UNetMidBlock2D") is not None: + resnets.extend( + (f"mid.block_{j + 1}" if ddpm else f"middle_block.{2 * j}", f"mid_block.resnets.{j}", False) + for j in range(2) + ) + if config.get("add_attention", True): + attentions.append(("mid.attn_1" if ddpm else "middle_block.1", "mid_block.attentions.0", channels[-1])) + reversed_channels = list(reversed(channels)) + previous, index = reversed_channels[0], 0 + for i, channel in enumerate(reversed_channels): + attention = "Attn" in config["up_block_types"][i] + next_channel = reversed_channels[min(i + 1, count - 1)] + for j in range(layers[count - 1 - i] + 1): + old = f"up.{count - 1 - i}.block.{j}" if ddpm else f"output_blocks.{index}.0" + skip = next_channel if j == layers[count - 1 - i] else channel + resnets.append((old, f"up_blocks.{i}.resnets.{j}", previous + skip != channel)) + if attention: + old = f"up.{count - 1 - i}.attn.{j}" if ddpm else f"output_blocks.{index}.1" + attentions.append((old, f"up_blocks.{i}.attentions.{j}", channel)) + previous = channel + index += 1 + if i < count - 1: + old = f"up.{count - 1 - i}.upsample" if ddpm else f"output_blocks.{index - 1}.{2 if attention else 1}" + new = f"up_blocks.{i}.upsamplers.0" + residual = ( + config["up_block_types"][i] == "ResnetUpsampleBlock2D" or config.get("upsample_type") == "resnet" + ) + if residual: + resnets.append((old, new, False)) + else: + modules.append((old + ".conv", new + ".conv")) + for old, new, shortcut in resnets: + pairs = ( + [ + ("norm1", "norm1"), + ("conv1", "conv1"), + ("norm2", "norm2"), + ("conv2", "conv2"), + ("temb_proj", "time_emb_proj"), + ] + if ddpm + else [ + ("in_layers.0", "norm1"), + ("in_layers.2", "conv1"), + ("out_layers.0", "norm2"), + ("out_layers.3", "conv2"), + ("emb_layers.1", "time_emb_proj"), + ] + ) + if shortcut: + pairs.append(("nin_shortcut" if ddpm else "skip_connection", "conv_shortcut")) + modules.extend((f"{old}.{a}", f"{new}.{b}") for a, b in pairs) + rules = [] + for old, new, channel in attentions: + modules.append((old + ".norm", new + ".group_norm")) + head_dim = config["attention_head_dim"] + if isinstance(head_dim, (tuple, list)): + parts = new.split(".") + index = ( + count - 1 + if parts[0] == "mid_block" + else (count - 1 - int(parts[1]) if parts[0] == "up_blocks" else int(parts[1])) + ) + head_dim = head_dim[index] + dim = head_dim or channel + heads = channel // dim + inner = heads * dim + if inner == 0: + raise ValueError("Attention head dimension must not exceed the block width.") + if ddpm: + for source, target in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("proj_out", "to_out.0")): + shape = (channel, inner) if source == "proj_out" else (inner, channel) + rules.append( + Rule((f"{old}.{source}.weight",), (f"{new}.{target}.weight",), Reshape(shape + (1, 1), shape)) + ) + mapping[f"{old}.{source}.bias"] = f"{new}.{target}.bias" + else: + for p in ("weight", "bias"): + trailing = (channel,) if p == "weight" else () + original_shape = (3 * inner, channel, 1) if p == "weight" else (3 * inner,) + transforms = [Reshape(original_shape, (3 * inner,) + trailing)] + if original_format == "ldm": + transforms.extend( + [ + Reshape((3 * inner,) + trailing, (heads, 3, dim) + trailing), + Permute((1, 0, 2, 3) if trailing else (1, 0, 2)), + Reshape((3, heads, dim) + trailing, (3 * inner,) + trailing), + ] + ) + transforms.append(Split((inner,) * 3)) + rules.append( + Rule( + (f"{old}.qkv.{p}",), + tuple(f"{new}.to_{part}.{p}" for part in ("q", "k", "v")), + Chain(tuple(transforms)), + ) + ) + rules.append( + Rule( + (old + ".proj_out.weight",), + (new + ".to_out.0.weight",), + Reshape((channel, inner, 1), (channel, inner)), + ) + ) + mapping[old + ".proj_out.bias"] = new + ".to_out.0.bias" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/unet_3d.py b/src/diffusers/loaders/conversion/unet_3d.py new file mode 100644 index 000000000000..054429889af6 --- /dev/null +++ b/src/diffusers/loaders/conversion/unet_3d.py @@ -0,0 +1,97 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .ldm_unet import ldm_unet_conversion +from .transforms import Reshape + + +def unet_3d_conversion(config): + spatial_config = { + **config, + "down_block_types": [name.replace("3D", "2D") for name in config["down_block_types"]], + "up_block_types": [name.replace("3D", "2D") for name in config["up_block_types"]], + "transformer_layers_per_block": 1, + } + base = ldm_unet_conversion(spatial_config) + mapping, rules, resnets, attention_prefixes = {}, [], [], [] + channels = config["block_out_channels"] + for old, new in base.mapping.items(): + if new.startswith("conv_norm_out.") and config["norm_num_groups"] is None: + continue + if new.startswith("mid_block.resnets.1."): + old = old.replace("middle_block.2.", "middle_block.3.") + elif ".downsamplers." in new: + old = old.replace(".0.op.", ".op.") + elif ".upsamplers." in new and config["up_block_types"][int(new.split(".")[1])] == "CrossAttnUpBlock3D": + old = old.replace(".2.conv.", ".3.conv.") + mapping[old] = new + if ".resnets." in new and new.endswith(".norm1.weight"): + resnets.append( + ( + old.removesuffix(".in_layers.0.weight"), + new.removesuffix(".norm1.weight").replace(".resnets.", ".temp_convs."), + ) + ) + if ".attentions." in new and new.endswith(".proj_in.bias"): + source = old.removesuffix(".proj_in.bias") + target = new.removesuffix(".proj_in.bias") + source = source.rsplit(".", 1)[0] + ".2" + if target.startswith("mid_block"): + channel = channels[-1] + else: + index = int(target.split(".")[1]) + channel = channels[-index - 1] if target.startswith("up_blocks") else channels[index] + attention_prefixes.append((source, target.replace(".attentions.", ".temp_attentions."), channel, channel)) + for old, new in resnets: + for i in range(1, 5): + for layer in (0, 2 if i == 1 else 3): + mapping.update( + { + f"{old}.temopral_conv.conv{i}.{layer}.{p}": f"{new}.conv{i}.{layer}.{p}" + for p in ("weight", "bias") + } + ) + attention_prefixes.append(("input_blocks.0.1", "transformer_in", channels[0], 8 * config["attention_head_dim"])) + for old, new, channel, inner in attention_prefixes: + block_mapping, block_rules = _temporal_transformer_rules(old, new, channel, inner) + mapping.update(block_mapping) + rules.extend(block_rules) + if config.get("time_cond_proj_dim") is not None: + mapping["time_embed.cond_proj.weight"] = "time_embedding.cond_proj.weight" + return Conversion(mapping=mapping, rules=tuple(rules)) + + +def _temporal_transformer_rules(old, new, channel, inner): + modules = [(old + ".norm", new + ".norm")] + mapping = {f"{old}.{name}.bias": f"{new}.{name}.bias" for name in ("proj_in", "proj_out")} + rules = [ + Rule((old + ".proj_in.weight",), (new + ".proj_in.weight",), Reshape((inner, channel, 1), (inner, channel))), + Rule((old + ".proj_out.weight",), (new + ".proj_out.weight",), Reshape((channel, inner, 1), (channel, inner))), + ] + a, b = old + ".transformer_blocks.0", new + ".transformer_blocks.0" + modules.extend( + (f"{a}.{name}", f"{b}.{name}") + for name in ("norm1", "norm2", "norm3", "ff.net.0.proj", "ff.net.2", "attn1.to_out.0", "attn2.to_out.0") + ) + mapping.update( + { + f"{a}.{attn}.to_{part}.weight": f"{b}.{attn}.to_{part}.weight" + for attn in ("attn1", "attn2") + for part in ("q", "k", "v") + } + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return mapping, rules diff --git a/src/diffusers/loaders/conversion/unidiffuser.py b/src/diffusers/loaders/conversion/unidiffuser.py new file mode 100644 index 000000000000..cf61c73a0703 --- /dev/null +++ b/src/diffusers/loaders/conversion/unidiffuser.py @@ -0,0 +1,78 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion, Rule +from .transforms import MergeEqual, Reverse, Split + + +def unidiffuser_conversion(config): + if ( + config["norm_type"] != "layer_norm" + or config["cross_attention_dim"] is not None + or config["block_type"] != "unidiffuser" + ): + raise ValueError("Original UniDiffuser uses layer_norm, self attention, and unidiffuser blocks.") + width = config["num_attention_heads"] * config["attention_head_dim"] + mapping = {"pos_embed": "pos_embed"} + modules = [ + ("clip_img_embed", "clip_img_in"), + ("text_embed", "text_in"), + ("decoder_pred", "vae_img_out"), + ("norm", "transformer.norm_out"), + ("clip_img_out", "clip_img_out"), + ("text_out", "text_out"), + ] + rules = [ + Rule( + (f"patch_embed.proj.{p}",), + (f"vae_img_in.proj.{p}", f"transformer.pos_embed.proj.{p}"), + Reverse(MergeEqual(2)), + ) + for p in ("weight", "bias") + ] + if config["use_data_type_embedding"]: + mapping.update( + { + "pos_embed_token": "data_type_pos_embed_token", + "token_embedding.weight": "data_type_token_embedding.weight", + } + ) + if config["use_timestep_embedding"]: + modules.extend( + (f"time_{modality}_embed.{i}", f"timestep_{modality}_embed.linear_{j}") + for modality in ("img", "text") + for i, j in ((0, 1), (2, 2)) + ) + blocks = ["mid_block"] + [f"{side}_blocks.{i}" for side in ("in", "out") for i in range(config["num_layers"] // 2)] + for old in blocks: + new = f"transformer.transformer_{old}" + if old.startswith("out_"): + modules.extend([(f"{old}.skip_linear", f"{new}.skip.skip_linear"), (f"{old}.norm1", f"{new}.skip.norm")]) + new += ".block" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in (("attn.proj", "attn1.to_out.0"), ("mlp.fc1", "ff.net.0.proj"), ("mlp.fc2", "ff.net.2")) + ) + if config["norm_elementwise_affine"]: + modules.extend([(f"{old}.norm2", f"{new}.norm1"), (f"{old}.norm3", f"{new}.norm3")]) + for p in ("weight", "bias") if config["attention_bias"] else ("weight",): + rules.append( + Rule( + (f"{old}.attn.qkv.{p}",), + tuple(f"{new}.attn1.to_{part}.{p}" for part in ("q", "k", "v")), + Split((width,) * 3), + ) + ) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/unidiffuser_text.py b/src/diffusers/loaders/conversion/unidiffuser_text.py new file mode 100644 index 000000000000..0f32de796b5d --- /dev/null +++ b/src/diffusers/loaders/conversion/unidiffuser_text.py @@ -0,0 +1,33 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def unidiffuser_text_conversion(config): + mapping = {} + if config["prefix_hidden_dim"] is not None: + mapping.update( + {f"{name}.{p}": f"{name}.{p}" for name in ("encode_prefix", "decode_prefix") for p in ("weight", "bias")} + ) + keys = ["transformer.wte.weight", "transformer.wpe.weight", "lm_head.weight"] + modules = ["transformer.ln_f"] + for i in range(config["n_layer"]): + modules.extend( + f"transformer.h.{i}.{name}" + for name in ("ln_1", "ln_2", "attn.c_attn", "attn.c_proj", "mlp.c_fc", "mlp.c_proj") + ) + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + mapping.update({f"gpt.{key}": f"transformer.{key}" for key in keys}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/uvit.py b/src/diffusers/loaders/conversion/uvit.py new file mode 100644 index 000000000000..285c4728af90 --- /dev/null +++ b/src/diffusers/loaders/conversion/uvit.py @@ -0,0 +1,117 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Reverse, Split + + +def uvit_conversion(config): + mapping = {"embed.embeddings.weight": "embed.embeddings.weight"} + modules = [ + (name, name) + for name in ( + "encoder_proj", + "embed.conv", + "project_to_hidden", + "project_from_hidden", + "mlm_layer.conv1", + "mlm_layer.conv2", + ) + ] + modules.extend([("cond_embed.0", "cond_embed.linear_1"), ("cond_embed.2", "cond_embed.linear_2")]) + norms = [ + (name, name) + for name in ( + "encoder_proj_layer_norm", + "embed.layer_norm", + "project_to_hidden_norm", + "project_from_hidden_norm", + ) + ] + norms.append(("mlm_layer.layer_norm.norm", "mlm_layer.layer_norm")) + rules = [] + for i in range(config["num_hidden_layers"]): + prefix = f"transformer_layers.{i}" + for source, target in ( + ("attn_layer_norm", "norm1.norm"), + ("crossattn_layer_norm", "norm2.norm"), + ("ffn.pre_mlp_layer_norm", "norm3.norm"), + ): + norms.append((f"{prefix}.{source}", f"{prefix}.{target}")) + modules.extend( + (f"{prefix}.{source}.mapper", f"{prefix}.norm{j}.linear") + for j, source in ( + (1, "self_attn_adaLN_modulation"), + (2, "cross_attn_adaLN_modulation"), + (3, "ffn.adaLN_modulation"), + ) + ) + for old, new in (("attention", "attn1"), ("crossattention", "attn2")): + modules.extend( + (f"{prefix}.{old}.{a}", f"{prefix}.{new}.{b}") + for a, b in (("query", "to_q"), ("key", "to_k"), ("value", "to_v"), ("out", "to_out.0")) + ) + modules.append((prefix + ".ffn.wo", prefix + ".ff.net.2")) + for p in ("weight", "bias") if config["use_bias"] else ("weight",): + rules.append( + Rule( + (f"{prefix}.ffn.wi_1.{p}", f"{prefix}.ffn.wi_0.{p}"), + (f"{prefix}.ff.net.0.proj.{p}",), + Reverse(Split((config["intermediate_size"],) * 2)), + ) + ) + for direction in ("down", "up"): + old, new = f"{direction}_blocks.0", f"{direction}_block" + resampling = direction + "sample" + if config[resampling]: + modules.append((f"{old}.{resampling}.1", f"{new}.{resampling}.conv")) + norms.append((f"{old}.{resampling}.0.norm", f"{new}.{resampling}.norm")) + for i in range(config["num_res_blocks"]): + a, b = f"{old}.res_blocks.{i}", f"{new}.res_blocks.{i}" + modules.extend( + (f"{a}.{source}", f"{b}.{target}") + for source, target in ( + ("depthwise", "depthwise"), + ("channelwise.0", "channelwise_linear_1"), + ("channelwise.4", "channelwise_linear_2"), + ("adaLN_modulation.mapper", "cond_embeds_mapper"), + ) + ) + norms.append((a + ".norm.norm", b + ".norm")) + mapping.update({f"{a}.channelwise.2.{p}": f"{b}.channelwise_norm.{p}" for p in ("gamma", "beta")}) + a, b = f"{old}.attention_blocks.{i}", f"{new}.attention_blocks.{i}" + mapping.update( + { + f"{a}.{source}.weight": f"{b}.{target}.weight" + for source, target in (("attn_layer_norm", "norm1"), ("crossattn_layer_norm", "norm2")) + } + ) + if config["hidden_size"] != config["block_out_channels"]: + modules.append((a + ".kv_mapper", b + ".kv_mapper")) + for source, target in (("attention", "attn1"), ("crossattention", "attn2")): + modules.extend( + (f"{a}.{source}.{x}", f"{b}.{target}.{y}") + for x, y in (("query", "to_q"), ("key", "to_k"), ("value", "to_v"), ("out", "to_out.0")) + ) + mapping.update( + { + f"{old}.{p}": f"{new}.{p}" + for old, new in modules + for p in (("weight", "bias") if config["use_bias"] else ("weight",)) + } + ) + if config["ln_elementwise_affine"]: + mapping.update({old + ".weight": new + ".weight" for old, new in norms}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/versatile_image_unet.py b/src/diffusers/loaders/conversion/versatile_image_unet.py new file mode 100644 index 000000000000..76cba4a72123 --- /dev/null +++ b/src/diffusers/loaders/conversion/versatile_image_unet.py @@ -0,0 +1,25 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_unet import _ldm_unet_mapping + + +def versatile_image_unet_conversion(config): + mapping = {} + for old, new in _ldm_unet_mapping(config, controlnet=False).items(): + prefix = "model.diffusion_model." if old.startswith("time_embed.") else "model.diffusion_model.unet_image." + mapping[prefix + old] = new + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/versatile_text_unet.py b/src/diffusers/loaders/conversion/versatile_text_unet.py new file mode 100644 index 000000000000..bf5051bddbdb --- /dev/null +++ b/src/diffusers/loaders/conversion/versatile_text_unet.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_unet import _ldm_unet_mapping + + +def versatile_text_unet_conversion(config): + spatial_config = { + **config, + "down_block_types": [kind.replace("Flat", "2D") for kind in config["down_block_types"]], + "up_block_types": [kind.replace("Flat", "2D") for kind in config["up_block_types"]], + } + base = _ldm_unet_mapping(spatial_config, controlnet=False) + mapping = {} + for old, new in base.items(): + if ".downsamplers." in new: + old = old.replace(".0.op.", ".0.") + new = new.replace(".downsamplers.0.conv.", ".downsamplers.0.") + elif ".upsamplers." in new: + old = old.replace(".conv.", ".") + new = new.replace(".upsamplers.0.conv.", ".upsamplers.0.") + prefix = "model.diffusion_model." if old.startswith("time_embed.") else "model.diffusion_model.unet_text." + mapping[prefix + old] = new + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/vq_diffusion.py b/src/diffusers/loaders/conversion/vq_diffusion.py new file mode 100644 index 000000000000..188548b90e9f --- /dev/null +++ b/src/diffusers/loaders/conversion/vq_diffusion.py @@ -0,0 +1,45 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def vq_diffusion_conversion(config): + prefix = "transformer.transformer" + mapping = { + f"{prefix}.content_emb.{name}.weight": f"latent_image_embedding.{name}.weight" + for name in ("emb", "height_emb", "width_emb") + } + modules = [(f"{prefix}.to_logits.0", "norm_out"), (f"{prefix}.to_logits.1", "out")] + for i in range(config["num_layers"]): + old, new = f"{prefix}.blocks.{i}", f"transformer_blocks.{i}" + modules.extend( + (f"{old}.{a}", f"{new}.{b}") + for a, b in ( + ("ln1.linear", "norm1.linear"), + ("ln1_1.linear", "norm2.linear"), + ("ln2", "norm3"), + ("mlp.0", "ff.net.0.proj"), + ("mlp.2", "ff.net.2"), + ) + ) + for a, b in (("ln1", "norm1"), ("ln1_1", "norm2")): + mapping[f"{old}.{a}.emb.weight"] = f"{new}.{b}.emb.weight" + for attn in ("attn1", "attn2"): + modules.append((f"{old}.{attn}.proj", f"{new}.{attn}.to_out.0")) + for a, b in (("query", "to_q"), ("key", "to_k"), ("value", "to_v")): + for p in ("weight", "bias") if config["attention_bias"] else ("weight",): + mapping[f"{old}.{attn}.{a}.{p}"] = f"{new}.{attn}.{b}.{p}" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/vq_model.py b/src/diffusers/loaders/conversion/vq_model.py new file mode 100644 index 000000000000..0546b5ab10f2 --- /dev/null +++ b/src/diffusers/loaders/conversion/vq_model.py @@ -0,0 +1,28 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .ldm_vae import ldm_vae_conversion + + +def vq_model_conversion(config): + if config.get("norm_type") == "spatial": + from .movq import movq_conversion + + return movq_conversion(config) + base = ldm_vae_conversion(config) + mapping = dict(base.mapping) + mapping["quantize.embedding.weight"] = "quantize.embedding.weight" + return Conversion(mapping=mapping, rules=base.rules) diff --git a/src/diffusers/loaders/conversion/wan.py b/src/diffusers/loaders/conversion/wan.py new file mode 100644 index 000000000000..a6e0e97e57a5 --- /dev/null +++ b/src/diffusers/loaders/conversion/wan.py @@ -0,0 +1,59 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def wan_conversion(config): + mapping = {"head.modulation": "scale_shift_table"} + modules = [ + ("patch_embedding", "patch_embedding"), + ("head.head", "proj_out"), + ("time_projection.1", "condition_embedder.time_proj"), + ] + for source, target in (("time_embedding", "time_embedder"), ("text_embedding", "text_embedder")): + modules.extend((f"{source}.{i}", f"condition_embedder.{target}.linear_{j}") for i, j in ((0, 1), (2, 2))) + if config["image_dim"] is not None: + modules.extend( + (f"img_emb.proj.{i}", f"condition_embedder.image_embedder.{name}") + for i, name in ((0, "norm1"), (1, "ff.net.0.proj"), (3, "ff.net.2"), (4, "norm2")) + ) + if config["pos_embed_seq_len"] is not None: + mapping["img_emb.emb_pos"] = "condition_embedder.image_embedder.pos_embed" + for i in range(config["num_layers"]): + prefix = f"blocks.{i}" + mapping[prefix + ".modulation"] = prefix + ".scale_shift_table" + modules.extend((f"{prefix}.ffn.{j}", f"{prefix}.ffn.{name}") for j, name in ((0, "net.0.proj"), (2, "net.2"))) + if config["cross_attn_norm"]: + modules.append((prefix + ".norm3", prefix + ".norm2")) + for source, target in (("self_attn", "attn1"), ("cross_attn", "attn2")): + modules.extend( + (f"{prefix}.{source}.{a}", f"{prefix}.{target}.{b}") + for a, b in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("o", "to_out.0")) + ) + if config["qk_norm"] is not None: + mapping.update( + { + f"{prefix}.{source}.norm_{part}.weight": f"{prefix}.{target}.norm_{part}.weight" + for part in ("q", "k") + } + ) + if source == "cross_attn" and config["added_kv_proj_dim"] is not None: + modules.extend( + (f"{prefix}.{source}.{part}_img", f"{prefix}.{target}.add_{part}_proj") for part in ("k", "v") + ) + mapping[f"{prefix}.{source}.norm_k_img.weight"] = f"{prefix}.{target}.norm_added_k.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/wan_animate.py b/src/diffusers/loaders/conversion/wan_animate.py new file mode 100644 index 000000000000..e92645ae2d47 --- /dev/null +++ b/src/diffusers/loaders/conversion/wan_animate.py @@ -0,0 +1,77 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import math + +from .core import Conversion, Rule +from .transforms import Reshape, Split +from .wan import wan_conversion + + +def wan_animate_conversion(config): + mapping = dict(wan_conversion(config).mapping) + modules = [("pose_patch_embedding", "pose_patch_embedding"), ("face_encoder.out_proj", "face_encoder.out_proj")] + mapping["face_encoder.padding_tokens"] = "face_encoder.padding_tokens" + mapping["motion_encoder.dec.direction.weight"] = "motion_encoder.motion_synthesis_weight" + mapping["motion_encoder.enc.net_app.convs.0.0.weight"] = "motion_encoder.conv_in.weight" + channels = config["motion_encoder_channel_sizes"] or { + "4": 512, + "8": 512, + "16": 512, + "32": 512, + "64": 256, + "128": 128, + "256": 64, + "512": 32, + "1024": 16, + } + size = config["motion_encoder_size"] + count = int(math.log2(size)) - 2 + rules = [ + Rule( + ("motion_encoder.enc.net_app.convs.0.1.bias",), + ("motion_encoder.conv_in.act_fn.bias",), + Reshape((1, channels[str(size)], 1, 1), (channels[str(size)],)), + ) + ] + for i in range(count): + old, new = f"motion_encoder.enc.net_app.convs.{i + 1}", f"motion_encoder.res_blocks.{i}" + mapping.update( + { + f"{old}.{a}.weight": f"{new}.{b}.weight" + for a, b in (("conv1.0", "conv1"), ("conv2.1", "conv2"), ("skip.1", "conv_skip")) + } + ) + for a, b, channel in ( + ("conv1.1", "conv1", channels[str(size // 2**i)]), + ("conv2.2", "conv2", channels[str(size // 2 ** (i + 1))]), + ): + rules.append( + Rule((f"{old}.{a}.bias",), (f"{new}.{b}.act_fn.bias",), Reshape((1, channel, 1, 1), (channel,))) + ) + mapping[f"motion_encoder.enc.net_app.convs.{count + 1}.weight"] = "motion_encoder.conv_out.weight" + modules.extend((f"motion_encoder.enc.fc.{i}", f"motion_encoder.motion_network.{i}") for i in range(5)) + modules.extend((f"face_encoder.{name}.conv", f"face_encoder.{name}") for name in ("conv1_local", "conv2", "conv3")) + hidden = config["num_attention_heads"] * config["attention_head_dim"] + for i in range(config["num_layers"] // config["inject_face_latents_blocks"]): + old, new = f"face_adapter.fuser_blocks.{i}", f"face_adapter.{i}" + modules.extend([(old + ".linear1_q", new + ".to_q"), (old + ".linear2", new + ".to_out")]) + for p in ("weight", "bias"): + rules.append( + Rule((f"{old}.linear1_kv.{p}",), (f"{new}.to_k.{p}", f"{new}.to_v.{p}"), Split((hidden,) * 2)) + ) + mapping.update({f"{old}.{part}_norm.weight": f"{new}.norm_{part}.weight" for part in ("q", "k")}) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/wan_animate2.py b/src/diffusers/loaders/conversion/wan_animate2.py new file mode 100644 index 000000000000..bab08190c477 --- /dev/null +++ b/src/diffusers/loaders/conversion/wan_animate2.py @@ -0,0 +1,53 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def wan_animate2_conversion(config): + mapping = {"head.modulation": "head.modulation"} + modules = [ + (name, name) + for name in ( + "patch_embedding", + "head.head", + "time_projection.1", + "time_embedding.0", + "time_embedding.2", + "text_embedding.0", + "text_embedding.2", + ) + ] + if config["use_img_emb"]: + modules.extend((f"img_emb.proj.{i}", f"img_emb.proj.{i}") for i in (0, 1, 3, 4)) + for i in range(config["num_layers"]): + old, new = f"blocks.{i}.block", f"blocks.{i}" + mapping[old + ".modulation"] = new + ".modulation" + modules.extend((f"{old}.ffn.{j}", f"{new}.ffn.{j}") for j in (0, 2)) + if config["cross_attn_norm"]: + modules.append((old + ".norm3", new + ".norm3")) + for attn in ("self_attn", "cross_attn"): + modules.extend( + (f"{old}.{attn}.{a}", f"{new}.{attn}.{b}") + for a, b in (("q", "to_q"), ("k", "to_k"), ("v", "to_v"), ("o", "to_out.0")) + ) + mapping.update( + {f"{old}.{attn}.norm_{part}.weight": f"{new}.{attn}.norm_{part}.weight" for part in ("q", "k")} + ) + if attn == "cross_attn" and config["use_img_emb"]: + modules.extend((f"{old}.{attn}.{part}_img", f"{new}.{attn}.add_{part}_proj") for part in ("k", "v")) + mapping[f"{old}.{attn}.norm_k_img.weight"] = f"{new}.{attn}.norm_added_k.weight" + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/wan_vace.py b/src/diffusers/loaders/conversion/wan_vace.py new file mode 100644 index 000000000000..b21280bd155f --- /dev/null +++ b/src/diffusers/loaders/conversion/wan_vace.py @@ -0,0 +1,37 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .wan import wan_conversion + + +def wan_vace_conversion(config): + base = wan_conversion(config) + mapping = dict(base.mapping) + for p in ("weight", "bias"): + mapping[f"vace_patch_embedding.{p}"] = f"vace_patch_embedding.{p}" + block_mapping = { + old.removeprefix("blocks.0."): new.removeprefix("blocks.0.") + for old, new in base.mapping.items() + if old.startswith("blocks.0.") + } + for i in range(len(config["vace_layers"])): + prefix = f"vace_blocks.{i}" + mapping.update({f"{prefix}.{old}": f"{prefix}.{new}" for old, new in block_mapping.items()}) + for p in ("weight", "bias"): + mapping[f"{prefix}.after_proj.{p}"] = f"{prefix}.proj_out.{p}" + if i == 0: + mapping[f"{prefix}.before_proj.{p}"] = f"{prefix}.proj_in.{p}" + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/wan_vae.py b/src/diffusers/loaders/conversion/wan_vae.py new file mode 100644 index 000000000000..46ccbd4f378e --- /dev/null +++ b/src/diffusers/loaders/conversion/wan_vae.py @@ -0,0 +1,93 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def wan_vae_conversion(config): + residual = config["is_residual"] + mapping = {} + modules = [("conv1", "quant_conv"), ("conv2", "post_quant_conv")] + resnets, attentions = [], [] + for component in ("encoder", "decoder"): + modules.extend( + [(f"{component}.conv1", f"{component}.conv_in"), (f"{component}.head.2", f"{component}.conv_out")] + ) + mapping[f"{component}.head.0.gamma"] = f"{component}.norm_out.gamma" + resnets.extend((f"{component}.middle.{i * 2}", f"{component}.mid_block.resnets.{i}", False) for i in range(2)) + attentions.append((f"{component}.middle.1", f"{component}.mid_block.attentions.0")) + dims = [config["base_dim"] * n for n in [1] + list(config["dim_mult"])] + idx, scale = 0, 1.0 + for i, (input_dim, output_dim) in enumerate(zip(dims, dims[1:])): + for j in range(config["num_res_blocks"]): + if residual: + old, new = f"encoder.downsamples.{i}.downsamples.{j}", f"encoder.down_blocks.{i}.resnets.{j}" + else: + old, new = f"encoder.downsamples.{idx}", f"encoder.down_blocks.{idx}" + resnets.append((old, new, input_dim != output_dim)) + input_dim = output_dim + idx += 1 + if not residual and scale in config["attn_scales"]: + attentions.append((f"encoder.downsamples.{idx}", f"encoder.down_blocks.{idx}")) + idx += 1 + if i < len(dims) - 2: + if residual: + old = f"encoder.downsamples.{i}.downsamples.{config['num_res_blocks']}" + new = f"encoder.down_blocks.{i}.downsampler" + else: + old, new = f"encoder.downsamples.{idx}", f"encoder.down_blocks.{idx}" + modules.append((old + ".resample.1", new + ".resample.1")) + if config["temperal_downsample"][i]: + modules.append((old + ".time_conv", new + ".time_conv")) + idx += 1 + scale /= 2 + decoder_dim = config["decoder_base_dim"] or config["base_dim"] + dims = [decoder_dim * n for n in [config["dim_mult"][-1]] + list(reversed(config["dim_mult"]))] + temporal = list(reversed(config["temperal_downsample"])) + idx = 0 + for i, (input_dim, output_dim) in enumerate(zip(dims, dims[1:])): + if i > 0 and not residual: + input_dim //= 2 + for j in range(config["num_res_blocks"] + 1): + resnets.append( + ( + f"decoder.upsamples.{i}.upsamples.{j}" if residual else f"decoder.upsamples.{idx}", + f"decoder.up_blocks.{i}.resnets.{j}", + input_dim != output_dim, + ) + ) + input_dim = output_dim + idx += 1 + if i < len(dims) - 2: + if residual: + old = f"decoder.upsamples.{i}.upsamples.{config['num_res_blocks'] + 1}" + new = f"decoder.up_blocks.{i}.upsampler" + else: + old, new = f"decoder.upsamples.{idx}", f"decoder.up_blocks.{i}.upsamplers.0" + modules.append((old + ".resample.1", new + ".resample.1")) + if temporal[i]: + modules.append((old + ".time_conv", new + ".time_conv")) + idx += 1 + for old, new, shortcut in resnets: + mapping[old + ".residual.0.gamma"] = new + ".norm1.gamma" + mapping[old + ".residual.3.gamma"] = new + ".norm2.gamma" + modules.extend([(old + ".residual.2", new + ".conv1"), (old + ".residual.6", new + ".conv2")]) + if shortcut: + modules.append((old + ".shortcut", new + ".conv_shortcut")) + for old, new in attentions: + mapping[old + ".norm.gamma"] = new + ".norm.gamma" + modules.extend((old + "." + name, new + "." + name) for name in ("to_qkv", "proj")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping) diff --git a/src/diffusers/loaders/conversion/wuerstchen_decoder.py b/src/diffusers/loaders/conversion/wuerstchen_decoder.py new file mode 100644 index 000000000000..685ce220761f --- /dev/null +++ b/src/diffusers/loaders/conversion/wuerstchen_decoder.py @@ -0,0 +1,42 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion +from .wuerstchen_prior import _wuerstchen_block_rules + + +def wuerstchen_decoder_conversion(config): + modules = ["clip_mapper", "embedding.1", "clf.1"] + injections = list(config["inject_effnet"]) + list(reversed(config["inject_effnet"])) + modules.extend(f"effnet_mappers.{i}" for i, active in enumerate(injections) if active) + blocks = [] + count = len(config["c_hidden"]) + for direction in ("down", "up"): + for i in range(count): + level = i if direction == "down" else count - 1 - i + offset = int(direction == "down" and level > 0) + kinds = config["level_config"][level] + for j in range(config["blocks"][level]): + blocks.extend( + (f"{direction}_blocks.{i}.{offset + j * len(kinds) + k}", kind, config["c_hidden"][level]) + for k, kind in enumerate(kinds) + ) + if direction == "down" and level > 0: + modules.append(f"down_blocks.{i}.0.1") + elif direction == "up" and level > 0: + modules.append(f"up_blocks.{i}.{config['blocks'][level] * len(kinds)}.1") + mapping, rules = _wuerstchen_block_rules(blocks) + mapping.update({f"{name}.{p}": f"{name}.{p}" for name in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=rules) diff --git a/src/diffusers/loaders/conversion/wuerstchen_prior.py b/src/diffusers/loaders/conversion/wuerstchen_prior.py new file mode 100644 index 000000000000..7563221210f2 --- /dev/null +++ b/src/diffusers/loaders/conversion/wuerstchen_prior.py @@ -0,0 +1,63 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def wuerstchen_prior_conversion(config): + blocks = [ + (f"blocks.{3 * i + j}", kind, config["c"]) for i in range(config["depth"]) for j, kind in enumerate("CTA") + ] + mapping, rules = _wuerstchen_block_rules(blocks) + mapping.update( + { + f"{name}.{p}": f"{name}.{p}" + for name in ("projection", "cond_mapper.0", "cond_mapper.2", "out.1") + for p in ("weight", "bias") + } + ) + return Conversion(mapping=mapping, rules=rules) + + +def _wuerstchen_block_rules(blocks): + mapping, modules, rules = {}, [], [] + for prefix, kind, channels in blocks: + if kind == "C": + modules.extend( + (f"{prefix}.{name}", f"{prefix}.{name}") for name in ("depthwise", "channelwise.0", "channelwise.4") + ) + mapping.update({f"{prefix}.channelwise.2.{p}": f"{prefix}.channelwise.2.{p}" for p in ("gamma", "beta")}) + elif kind == "T": + modules.append((prefix + ".mapper", prefix + ".mapper")) + elif kind == "A": + modules.extend( + [ + (prefix + ".kv_mapper.1", prefix + ".kv_mapper.1"), + (prefix + ".attention.attn.out_proj", prefix + ".attention.to_out.0"), + ] + ) + for p in ("weight", "bias"): + rules.append( + Rule( + (f"{prefix}.attention.attn.in_proj_{p}",), + tuple(f"{prefix}.attention.to_{part}.{p}" for part in ("q", "k", "v")), + Split((channels,) * 3), + ) + ) + else: + raise ValueError(f"Unknown Wuerstchen block type {kind}.") + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return mapping, tuple(rules) diff --git a/src/diffusers/loaders/conversion/z_image.py b/src/diffusers/loaders/conversion/z_image.py new file mode 100644 index 000000000000..d20c1f1f321f --- /dev/null +++ b/src/diffusers/loaders/conversion/z_image.py @@ -0,0 +1,80 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion, Rule +from .transforms import Split + + +def z_image_conversion(config): + mapping = { + "x_pad_token": "x_pad_token", + "cap_pad_token": "cap_pad_token", + "cap_embedder.0.weight": "cap_embedder.0.weight", + } + modules = [ + ("t_embedder.mlp.0", "t_embedder.mlp.0"), + ("t_embedder.mlp.2", "t_embedder.mlp.2"), + ("cap_embedder.1", "cap_embedder.1"), + ] + patches = list(zip(config["all_patch_size"], config["all_f_patch_size"])) + if patches != [(2, 1)]: + raise ValueError("The original Z-Image single-patch checkpoint format requires patch sizes (2, 1).") + modules.extend( + [ + ("x_embedder", "all_x_embedder.2-1"), + ("final_layer.linear", "all_final_layer.2-1.linear"), + ("final_layer.adaLN_modulation.1", "all_final_layer.2-1.adaLN_modulation.1"), + ] + ) + groups = [ + ("noise_refiner", config["n_refiner_layers"]), + ("context_refiner", config["n_refiner_layers"]), + ("layers", config["n_layers"]), + ] + if config["siglip_feat_dim"] is not None: + mapping.update( + {"siglip_pad_token": "siglip_pad_token", "siglip_embedder.0.weight": "siglip_embedder.0.weight"} + ) + modules.append(("siglip_embedder.1", "siglip_embedder.1")) + groups.append(("siglip_refiner", config["n_refiner_layers"])) + rules = [] + for group, count in groups: + for i in range(count): + prefix = f"{group}.{i}" + for name in ( + "feed_forward.w1", + "feed_forward.w2", + "feed_forward.w3", + "attention_norm1", + "attention_norm2", + "ffn_norm1", + "ffn_norm2", + ): + mapping[f"{prefix}.{name}.weight"] = f"{prefix}.{name}.weight" + mapping[prefix + ".attention.out.weight"] = prefix + ".attention.to_out.0.weight" + if config["qk_norm"]: + for part in ("q", "k"): + mapping[f"{prefix}.attention.{part}_norm.weight"] = f"{prefix}.attention.norm_{part}.weight" + rules.append( + Rule( + (prefix + ".attention.qkv.weight",), + tuple(f"{prefix}.attention.to_{part}.weight" for part in ("q", "k", "v")), + Split((config["dim"],) * 3), + ) + ) + if group in ("noise_refiner", "layers"): + modules.append((prefix + ".adaLN_modulation.0", prefix + ".adaLN_modulation.0")) + mapping.update({f"{old}.{p}": f"{new}.{p}" for old, new in modules for p in ("weight", "bias")}) + return Conversion(mapping=mapping, rules=tuple(rules)) diff --git a/src/diffusers/loaders/conversion/z_image_controlnet.py b/src/diffusers/loaders/conversion/z_image_controlnet.py new file mode 100644 index 000000000000..e0af0a0bdc2f --- /dev/null +++ b/src/diffusers/loaders/conversion/z_image_controlnet.py @@ -0,0 +1,54 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from .core import Conversion + + +def z_image_controlnet_conversion(config): + keys, modules = [], [] + modules.extend( + f"control_all_x_embedder.{p}-{f}" for p, f in zip(config["all_patch_size"], config["all_f_patch_size"]) + ) + blocks = [(f"control_layers.{i}", True, index == 0) for i, index in enumerate(config["control_layers_places"])] + mode = config["add_control_noise_refiner"] + if mode != "control_layers": + blocks.extend( + (f"control_noise_refiner.{i}", mode == "control_noise_refiner", i == 0) + for i in range(config["n_refiner_layers"]) + ) + for prefix, controlled, first in blocks: + names = [ + "attention.to_q", + "attention.to_k", + "attention.to_v", + "attention.to_out.0", + "feed_forward.w1", + "feed_forward.w2", + "feed_forward.w3", + "attention_norm1", + "attention_norm2", + "ffn_norm1", + "ffn_norm2", + ] + if config["qk_norm"]: + names.extend(["attention.norm_q", "attention.norm_k"]) + keys.extend(f"{prefix}.{name}.weight" for name in names) + modules.append(prefix + ".adaLN_modulation.0") + if controlled: + modules.append(prefix + ".after_proj") + if first: + modules.append(prefix + ".before_proj") + keys.extend(f"{name}.{p}" for name in modules for p in ("weight", "bias")) + return Conversion(mapping={key: key for key in keys}) diff --git a/src/diffusers/loaders/conversion/zero123_projection.py b/src/diffusers/loaders/conversion/zero123_projection.py new file mode 100644 index 000000000000..d62aaf9163e4 --- /dev/null +++ b/src/diffusers/loaders/conversion/zero123_projection.py @@ -0,0 +1,19 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .core import Conversion + + +def zero123_projection_conversion(config): + return Conversion(mapping={f"cc_projection.{p}": f"projection.{p}" for p in ("weight", "bias")}) diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index a07657159d36..f85639f4f0d1 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -11,10 +11,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -import importlib import inspect import re +from collections.abc import Mapping from contextlib import nullcontext +from typing import Any import torch from huggingface_hub.utils import validate_hf_hub_args @@ -29,36 +30,10 @@ from ..quantizers import DiffusersAutoQuantizer from ..utils import deprecate, is_accelerate_available, is_torch_version, logging from ..utils.torch_utils import empty_device_cache +from .conversion.registry import CONVERSION_BUILDERS from .single_file_utils import ( SingleFileComponentError, - convert_animatediff_checkpoint_to_diffusers, - convert_auraflow_transformer_checkpoint_to_diffusers, - convert_autoencoder_dc_checkpoint_to_diffusers, - convert_chroma_transformer_checkpoint_to_diffusers, - convert_controlnet_checkpoint, - convert_cosmos_transformer_checkpoint_to_diffusers, - convert_ernie_image_transformer_checkpoint_to_diffusers, - convert_flux2_transformer_checkpoint_to_diffusers, - convert_flux_transformer_checkpoint_to_diffusers, - convert_hidream_transformer_to_diffusers, - convert_hunyuan_video_transformer_to_diffusers, - convert_ldm_unet_checkpoint, - convert_ldm_vae_checkpoint, - convert_ltx2_audio_vae_to_diffusers, - convert_ltx2_transformer_to_diffusers, - convert_ltx2_vae_to_diffusers, - convert_ltx_transformer_checkpoint_to_diffusers, - convert_ltx_vae_checkpoint_to_diffusers, - convert_lumina2_to_diffusers, - convert_mochi_transformer_checkpoint_to_diffusers, - convert_sana_transformer_to_diffusers, - convert_sd3_transformer_checkpoint_to_diffusers, - convert_stable_cascade_unet_single_file_to_diffusers, - convert_wan_animate_2_transformer_to_diffusers, - convert_wan_transformer_to_diffusers, - convert_wan_vae_to_diffusers, - convert_z_image_controlnet_checkpoint_to_diffusers, - convert_z_image_transformer_checkpoint_to_diffusers, + convert_model_checkpoint, create_controlnet_diffusers_config_from_ldm, create_unet_diffusers_config_from_ldm, create_vae_diffusers_config_from_ldm, @@ -81,12 +56,11 @@ else: _LOW_CPU_MEM_USAGE_DEFAULT = False -SINGLE_FILE_LOADABLE_CLASSES = { - "StableCascadeUNet": { - "checkpoint_mapping_fn": convert_stable_cascade_unet_single_file_to_diffusers, - }, +# Models with established original-config mappings or checkpoint-based config inference. +# Other registered conversions require an explicit Diffusers component config. +SINGLE_FILE_CONFIGS = { + "StableCascadeUNet": {}, "UNet2DConditionModel": { - "checkpoint_mapping_fn": convert_ldm_unet_checkpoint, "config_mapping_fn": create_unet_diffusers_config_from_ldm, "default_subfolder": "unet", "legacy_kwargs": { @@ -94,130 +68,95 @@ }, }, "AutoencoderKL": { - "checkpoint_mapping_fn": convert_ldm_vae_checkpoint, "config_mapping_fn": create_vae_diffusers_config_from_ldm, "default_subfolder": "vae", }, "ControlNetModel": { - "checkpoint_mapping_fn": convert_controlnet_checkpoint, "config_mapping_fn": create_controlnet_diffusers_config_from_ldm, }, "SD3Transformer2DModel": { - "checkpoint_mapping_fn": convert_sd3_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, - "MotionAdapter": { - "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers, - }, - "SparseControlNetModel": { - "checkpoint_mapping_fn": convert_animatediff_checkpoint_to_diffusers, - }, + "MotionAdapter": {}, + "SparseControlNetModel": {}, "FluxTransformer2DModel": { - "checkpoint_mapping_fn": convert_flux_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "ChromaTransformer2DModel": { - "checkpoint_mapping_fn": convert_chroma_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "ErnieImageTransformer2DModel": { - "checkpoint_mapping_fn": convert_ernie_image_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "LTXVideoTransformer3DModel": { - "checkpoint_mapping_fn": convert_ltx_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "AutoencoderKLLTXVideo": { - "checkpoint_mapping_fn": convert_ltx_vae_checkpoint_to_diffusers, "default_subfolder": "vae", }, - "AutoencoderDC": {"checkpoint_mapping_fn": convert_autoencoder_dc_checkpoint_to_diffusers}, + "AutoencoderDC": {}, "MochiTransformer3DModel": { - "checkpoint_mapping_fn": convert_mochi_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "HunyuanVideoTransformer3DModel": { - "checkpoint_mapping_fn": convert_hunyuan_video_transformer_to_diffusers, "default_subfolder": "transformer", }, "AuraFlowTransformer2DModel": { - "checkpoint_mapping_fn": convert_auraflow_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "Lumina2Transformer2DModel": { - "checkpoint_mapping_fn": convert_lumina2_to_diffusers, "default_subfolder": "transformer", }, "SanaTransformer2DModel": { - "checkpoint_mapping_fn": convert_sana_transformer_to_diffusers, "default_subfolder": "transformer", }, "SkyReelsV2Transformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, "ChronoEditTransformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, "WanTransformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, "WanVACETransformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, "WanAnimateTransformer3DModel": { - "checkpoint_mapping_fn": convert_wan_transformer_to_diffusers, "default_subfolder": "transformer", }, "WanAnimate2Transformer3DModel": { - "checkpoint_mapping_fn": convert_wan_animate_2_transformer_to_diffusers, "default_subfolder": "transformer", }, "AutoencoderKLWan": { - "checkpoint_mapping_fn": convert_wan_vae_to_diffusers, "default_subfolder": "vae", }, "HiDreamImageTransformer2DModel": { - "checkpoint_mapping_fn": convert_hidream_transformer_to_diffusers, "default_subfolder": "transformer", }, "CosmosTransformer3DModel": { - "checkpoint_mapping_fn": convert_cosmos_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "QwenImageTransformer2DModel": { - "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint, "default_subfolder": "transformer", }, "Flux2Transformer2DModel": { - "checkpoint_mapping_fn": convert_flux2_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, "ZImageTransformer2DModel": { - "checkpoint_mapping_fn": convert_z_image_transformer_checkpoint_to_diffusers, "default_subfolder": "transformer", }, - "ZImageControlNetModel": { - "checkpoint_mapping_fn": convert_z_image_controlnet_checkpoint_to_diffusers, - }, + "ZImageControlNetModel": {}, "LTX2VideoTransformer3DModel": { - "checkpoint_mapping_fn": convert_ltx2_transformer_to_diffusers, "default_subfolder": "transformer", }, "AutoencoderKLLTX2Video": { - "checkpoint_mapping_fn": convert_ltx2_vae_to_diffusers, "default_subfolder": "vae", }, "AutoencoderKLLTX2Audio": { - "checkpoint_mapping_fn": convert_ltx2_audio_vae_to_diffusers, "default_subfolder": "audio_vae", }, "MotifVideoTransformer3DModel": { - "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint, "default_subfolder": "transformer", }, } @@ -232,14 +171,15 @@ def _should_convert_state_dict_to_diffusers(model_state_dict, checkpoint_state_d def _get_single_file_loadable_mapping_class(cls): - diffusers_module = importlib.import_module(__name__.split(".")[0]) - for loadable_class_str in SINGLE_FILE_LOADABLE_CLASSES: - loadable_class = getattr(diffusers_module, loadable_class_str) - - if issubclass(cls, loadable_class): - return loadable_class_str - - return None + # Follow the MRO so a specialized model keeps its own conversion before considering a base class. + return next( + ( + base.__name__ + for base in cls.__mro__ + if base.__module__.startswith("diffusers.") and base.__name__ in CONVERSION_BUILDERS + ), + None, + ) def _get_mapping_function_kwargs(mapping_fn, **kwargs): @@ -260,7 +200,9 @@ class FromOriginalModelMixin: @classmethod @validate_hf_hub_args - def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = None, **kwargs) -> Self: + def from_single_file( + cls, pretrained_model_link_or_path_or_dict: str | Mapping[str, Any] | None = None, **kwargs + ) -> Self: r""" Instantiate a model from pretrained weights saved in the original `.ckpt` or `.safetensors` format. The model is set in evaluation mode (`model.eval()`) by default. @@ -272,11 +214,12 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No `"https://huggingface.co//blob/main/.safetensors"`) on the Hub. - A path to a local *file* containing the weights of the component model. - A state dict containing the component model weights. - config (`str`, *optional*): + config (`str` or `dict`, *optional*): - A string, the *repo id* (for example `CompVis/ldm-text2im-large-256`) of a pretrained pipeline hosted on the Hub. - A path to a *directory* (for example `./my_pipeline_directory/`) containing the pipeline component configs in Diffusers format. + - A dictionary containing the matching Diffusers component configuration. subfolder (`str`, *optional*, defaults to `""`): The subfolder location of a model file within a larger model repository on the Hub or locally. original_config (`str`, *optional*): @@ -325,11 +268,8 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No """ mapping_class_name = _get_single_file_loadable_mapping_class(cls) - # if class_name not in SINGLE_FILE_LOADABLE_CLASSES: if mapping_class_name is None: - raise ValueError( - f"FromOriginalModelMixin is currently only compatible with {', '.join(SINGLE_FILE_LOADABLE_CLASSES.keys())}" - ) + raise ValueError(f"No original checkpoint conversion is registered for {cls.__name__}.") pretrained_model_link_or_path = kwargs.get("pretrained_model_link_or_path", None) if pretrained_model_link_or_path is not None: @@ -379,7 +319,7 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No f"Passed `torch_dtype` {torch_dtype} is not a `torch.dtype`. Defaulting to `torch.float32`." ) - if isinstance(pretrained_model_link_or_path_or_dict, dict): + if isinstance(pretrained_model_link_or_path_or_dict, Mapping): checkpoint = pretrained_model_link_or_path_or_dict else: checkpoint = load_single_file_checkpoint( @@ -401,12 +341,17 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No else: hf_quantizer = None - mapping_functions = SINGLE_FILE_LOADABLE_CLASSES[mapping_class_name] + loading_config = SINGLE_FILE_CONFIGS.get(mapping_class_name, {"config_required": True}) + + if loading_config.get("config_required", False) and config is None and original_config is None: + raise ValueError( + f"{mapping_class_name} requires an explicit Diffusers `config` directory or repository when loading " + "a single file. The original weights do not specify all settings of the model variant." + ) - checkpoint_mapping_fn = mapping_functions["checkpoint_mapping_fn"] if original_config is not None: - if "config_mapping_fn" in mapping_functions: - config_mapping_fn = mapping_functions["config_mapping_fn"] + if "config_mapping_fn" in loading_config: + config_mapping_fn = loading_config["config_mapping_fn"] else: config_mapping_fn = None @@ -430,7 +375,10 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No **config_mapping_kwargs, ) else: - if config is not None: + config_is_mapping = isinstance(config, Mapping) + if config_is_mapping: + diffusers_model_config = dict(config) + elif config is not None: if isinstance(config, str): default_pretrained_model_config_name = config else: @@ -445,25 +393,26 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No config = fetch_diffusers_config(checkpoint) default_pretrained_model_config_name = config["pretrained_model_name_or_path"] - if "default_subfolder" in mapping_functions: - subfolder = mapping_functions["default_subfolder"] + if "default_subfolder" in loading_config: + subfolder = loading_config["default_subfolder"] subfolder = subfolder or config.pop( "subfolder", None ) # some configs contain a subfolder key, e.g. StableCascadeUNet - diffusers_model_config = cls.load_config( - pretrained_model_name_or_path=default_pretrained_model_config_name, - subfolder=subfolder, - local_files_only=local_files_only, - token=token, - revision=config_revision, - ) + if not config_is_mapping: + diffusers_model_config = cls.load_config( + pretrained_model_name_or_path=default_pretrained_model_config_name, + subfolder=subfolder, + local_files_only=local_files_only, + token=token, + revision=config_revision, + ) expected_kwargs, optional_kwargs = cls._get_signature_keys(cls) # Map legacy kwargs to new kwargs - if "legacy_kwargs" in mapping_functions: - legacy_kwargs = mapping_functions["legacy_kwargs"] + if "legacy_kwargs" in loading_config: + legacy_kwargs = loading_config["legacy_kwargs"] for legacy_key, new_key in legacy_kwargs.items(): if legacy_key in kwargs: kwargs[new_key] = kwargs.pop(legacy_key) @@ -495,14 +444,17 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No expanded_device_map = _expand_device_map(device_map, model_state_dict.keys()) _caching_allocator_warmup(model, expanded_device_map, torch_dtype, hf_quantizer) - checkpoint_mapping_kwargs = _get_mapping_function_kwargs(checkpoint_mapping_fn, **kwargs) + checkpoint_mapping_kwargs = _get_mapping_function_kwargs(convert_model_checkpoint, **kwargs) + checkpoint_mapping_kwargs["return_config"] = True if _should_convert_state_dict_to_diffusers(model_state_dict, checkpoint): - diffusers_format_checkpoint = checkpoint_mapping_fn( - config=diffusers_model_config, + diffusers_format_checkpoint, conversion_config = convert_model_checkpoint( + config=dict(diffusers_model_config, _class_name=mapping_class_name), checkpoint=checkpoint, **checkpoint_mapping_kwargs, ) + if "original_format" in conversion_config: + model.register_to_config(original_format=conversion_config["original_format"]) else: diffusers_format_checkpoint = checkpoint diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index ac707947d491..821756bd94b9 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -12,13 +12,15 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Conversion script for the Stable Diffusion checkpoints.""" +"""Checkpoint discovery, configuration and component loading for single-file pipelines.""" import copy import os import re +from collections.abc import Mapping from contextlib import nullcontext from io import BytesIO +from typing import Any from urllib.parse import urlparse import requests @@ -59,6 +61,24 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +def convert_model_checkpoint( + checkpoint: Mapping[str, Any], + config: dict[str, Any], + extract_ema: bool = False, + return_config: bool = False, +) -> dict[str, torch.Tensor] | tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Convert a resolved single-file model using the shared bidirectional component definition. + + Model loading and original-format export share the config-driven definitions in `loaders.conversion`. + """ + from .conversion.checkpoint import convert_component_checkpoint + + return convert_component_checkpoint( + checkpoint, config, config["_class_name"], extract_ema=extract_ema, return_config=return_config + ) + + CHECKPOINT_KEY_NAMES = { "v1": "model.diffusion_model.output_blocks.11.0.skip_connection.weight", "v2": "model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight", @@ -260,115 +280,6 @@ } -DIFFUSERS_TO_LDM_MAPPING = { - "unet": { - "layers": { - "time_embedding.linear_1.weight": "time_embed.0.weight", - "time_embedding.linear_1.bias": "time_embed.0.bias", - "time_embedding.linear_2.weight": "time_embed.2.weight", - "time_embedding.linear_2.bias": "time_embed.2.bias", - "conv_in.weight": "input_blocks.0.0.weight", - "conv_in.bias": "input_blocks.0.0.bias", - "conv_norm_out.weight": "out.0.weight", - "conv_norm_out.bias": "out.0.bias", - "conv_out.weight": "out.2.weight", - "conv_out.bias": "out.2.bias", - }, - "class_embed_type": { - "class_embedding.linear_1.weight": "label_emb.0.0.weight", - "class_embedding.linear_1.bias": "label_emb.0.0.bias", - "class_embedding.linear_2.weight": "label_emb.0.2.weight", - "class_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - "addition_embed_type": { - "add_embedding.linear_1.weight": "label_emb.0.0.weight", - "add_embedding.linear_1.bias": "label_emb.0.0.bias", - "add_embedding.linear_2.weight": "label_emb.0.2.weight", - "add_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - }, - "controlnet": { - "layers": { - "time_embedding.linear_1.weight": "time_embed.0.weight", - "time_embedding.linear_1.bias": "time_embed.0.bias", - "time_embedding.linear_2.weight": "time_embed.2.weight", - "time_embedding.linear_2.bias": "time_embed.2.bias", - "conv_in.weight": "input_blocks.0.0.weight", - "conv_in.bias": "input_blocks.0.0.bias", - "controlnet_cond_embedding.conv_in.weight": "input_hint_block.0.weight", - "controlnet_cond_embedding.conv_in.bias": "input_hint_block.0.bias", - "controlnet_cond_embedding.conv_out.weight": "input_hint_block.14.weight", - "controlnet_cond_embedding.conv_out.bias": "input_hint_block.14.bias", - }, - "class_embed_type": { - "class_embedding.linear_1.weight": "label_emb.0.0.weight", - "class_embedding.linear_1.bias": "label_emb.0.0.bias", - "class_embedding.linear_2.weight": "label_emb.0.2.weight", - "class_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - "addition_embed_type": { - "add_embedding.linear_1.weight": "label_emb.0.0.weight", - "add_embedding.linear_1.bias": "label_emb.0.0.bias", - "add_embedding.linear_2.weight": "label_emb.0.2.weight", - "add_embedding.linear_2.bias": "label_emb.0.2.bias", - }, - }, - "vae": { - "encoder.conv_in.weight": "encoder.conv_in.weight", - "encoder.conv_in.bias": "encoder.conv_in.bias", - "encoder.conv_out.weight": "encoder.conv_out.weight", - "encoder.conv_out.bias": "encoder.conv_out.bias", - "encoder.conv_norm_out.weight": "encoder.norm_out.weight", - "encoder.conv_norm_out.bias": "encoder.norm_out.bias", - "decoder.conv_in.weight": "decoder.conv_in.weight", - "decoder.conv_in.bias": "decoder.conv_in.bias", - "decoder.conv_out.weight": "decoder.conv_out.weight", - "decoder.conv_out.bias": "decoder.conv_out.bias", - "decoder.conv_norm_out.weight": "decoder.norm_out.weight", - "decoder.conv_norm_out.bias": "decoder.norm_out.bias", - "quant_conv.weight": "quant_conv.weight", - "quant_conv.bias": "quant_conv.bias", - "post_quant_conv.weight": "post_quant_conv.weight", - "post_quant_conv.bias": "post_quant_conv.bias", - }, - "openclip": { - "layers": { - "text_model.embeddings.position_embedding.weight": "positional_embedding", - "text_model.embeddings.token_embedding.weight": "token_embedding.weight", - "text_model.final_layer_norm.weight": "ln_final.weight", - "text_model.final_layer_norm.bias": "ln_final.bias", - "text_projection.weight": "text_projection", - }, - "transformer": { - "text_model.encoder.layers.": "resblocks.", - "layer_norm1": "ln_1", - "layer_norm2": "ln_2", - ".fc1.": ".c_fc.", - ".fc2.": ".c_proj.", - ".self_attn": ".attn", - "transformer.text_model.final_layer_norm.": "ln_final.", - "transformer.text_model.embeddings.token_embedding.weight": "token_embedding.weight", - "transformer.text_model.embeddings.position_embedding.weight": "positional_embedding", - }, - }, -} - -SD_2_TEXT_ENCODER_KEYS_TO_IGNORE = [ - "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_bias", - "cond_stage_model.model.transformer.resblocks.23.attn.in_proj_weight", - "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.bias", - "cond_stage_model.model.transformer.resblocks.23.attn.out_proj.weight", - "cond_stage_model.model.transformer.resblocks.23.ln_1.bias", - "cond_stage_model.model.transformer.resblocks.23.ln_1.weight", - "cond_stage_model.model.transformer.resblocks.23.ln_2.bias", - "cond_stage_model.model.transformer.resblocks.23.ln_2.weight", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.bias", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_fc.weight", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.bias", - "cond_stage_model.model.transformer.resblocks.23.mlp.c_proj.weight", - "cond_stage_model.model.text_projection", -] - # To support legacy scheduler_type argument SCHEDULER_DEFAULT_CONFIG = { "beta_schedule": "scaled_linear", @@ -384,16 +295,12 @@ "timestep_spacing": "leading", } -LDM_VAE_KEYS = ["first_stage_model.", "vae."] LDM_VAE_DEFAULT_SCALING_FACTOR = 0.18215 PLAYGROUND_VAE_SCALING_FACTOR = 0.5 -LDM_UNET_KEY = "model.diffusion_model." -LDM_CONTROLNET_KEY = "control_model." LDM_CLIP_PREFIX_TO_REMOVE = [ "cond_stage_model.transformer.", "conditioner.embedders.0.transformer.", ] -LDM_OPEN_CLIP_TEXT_PROJECTION_DIM = 1024 SCHEDULER_LEGACY_KWARGS = ["prediction_type", "scheduler_type"] VALID_URL_PREFIXES = ["https://huggingface.co/", "huggingface.co/", "hf.co/", "https://hf.co/"] @@ -848,19 +755,6 @@ def set_image_size(checkpoint, image_size=None): return image_size -# Copied from diffusers.pipelines.stable_diffusion.convert_from_ckpt.conv_attn_to_linear -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - def create_unet_diffusers_config_from_ldm( original_config, checkpoint, image_size=None, upcast_attention=None, num_in_channels=None ): @@ -1075,513 +969,6 @@ def create_vae_diffusers_config_from_ldm(original_config, checkpoint, image_size return config -def update_unet_resnet_ldm_to_diffusers(ldm_keys, new_checkpoint, checkpoint, mapping=None): - for ldm_key in ldm_keys: - diffusers_key = ( - ldm_key.replace("in_layers.0", "norm1") - .replace("in_layers.2", "conv1") - .replace("out_layers.0", "norm2") - .replace("out_layers.3", "conv2") - .replace("emb_layers.1", "time_emb_proj") - .replace("skip_connection", "conv_shortcut") - ) - if mapping: - diffusers_key = diffusers_key.replace(mapping["old"], mapping["new"]) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_unet_attention_ldm_to_diffusers(ldm_keys, new_checkpoint, checkpoint, mapping): - for ldm_key in ldm_keys: - diffusers_key = ldm_key.replace(mapping["old"], mapping["new"]) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_vae_resnet_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ldm_key.replace(mapping["old"], mapping["new"]).replace("nin_shortcut", "conv_shortcut") - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - -def update_vae_attentions_ldm_to_diffusers(keys, new_checkpoint, checkpoint, mapping): - for ldm_key in keys: - diffusers_key = ( - ldm_key.replace(mapping["old"], mapping["new"]) - .replace("norm.weight", "group_norm.weight") - .replace("norm.bias", "group_norm.bias") - .replace("q.weight", "to_q.weight") - .replace("q.bias", "to_q.bias") - .replace("k.weight", "to_k.weight") - .replace("k.bias", "to_k.bias") - .replace("v.weight", "to_v.weight") - .replace("v.bias", "to_v.bias") - .replace("proj_out.weight", "to_out.0.weight") - .replace("proj_out.bias", "to_out.0.bias") - ) - new_checkpoint[diffusers_key] = checkpoint.get(ldm_key) - - # proj_attn.weight has to be converted from conv 1D to linear - shape = new_checkpoint[diffusers_key].shape - - if len(shape) == 3: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0] - elif len(shape) == 4: - new_checkpoint[diffusers_key] = new_checkpoint[diffusers_key][:, :, 0, 0] - - -def convert_stable_cascade_unet_single_file_to_diffusers(checkpoint, **kwargs): - is_stage_c = "clip_txt_mapper.weight" in checkpoint - - if is_stage_c: - state_dict = {} - for key in checkpoint.keys(): - if key.endswith("in_proj_weight"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - else: - state_dict[key] = checkpoint[key] - else: - state_dict = {} - for key in checkpoint.keys(): - if key.endswith("in_proj_weight"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_weight", "to_q.weight")] = weights[0] - state_dict[key.replace("attn.in_proj_weight", "to_k.weight")] = weights[1] - state_dict[key.replace("attn.in_proj_weight", "to_v.weight")] = weights[2] - elif key.endswith("in_proj_bias"): - weights = checkpoint[key].chunk(3, 0) - state_dict[key.replace("attn.in_proj_bias", "to_q.bias")] = weights[0] - state_dict[key.replace("attn.in_proj_bias", "to_k.bias")] = weights[1] - state_dict[key.replace("attn.in_proj_bias", "to_v.bias")] = weights[2] - elif key.endswith("out_proj.weight"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.weight", "to_out.0.weight")] = weights - elif key.endswith("out_proj.bias"): - weights = checkpoint[key] - state_dict[key.replace("attn.out_proj.bias", "to_out.0.bias")] = weights - # rename clip_mapper to clip_txt_pooled_mapper - elif key.endswith("clip_mapper.weight"): - weights = checkpoint[key] - state_dict[key.replace("clip_mapper.weight", "clip_txt_pooled_mapper.weight")] = weights - elif key.endswith("clip_mapper.bias"): - weights = checkpoint[key] - state_dict[key.replace("clip_mapper.bias", "clip_txt_pooled_mapper.bias")] = weights - else: - state_dict[key] = checkpoint[key] - - return state_dict - - -def convert_ldm_unet_checkpoint(checkpoint, config, extract_ema=False, **kwargs): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - unet_key = LDM_UNET_KEY - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning("Checkpoint has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.get(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - ldm_unet_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["layers"] - for diffusers_key, ldm_key in ldm_unet_keys.items(): - if ldm_key not in unet_state_dict: - continue - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - if ("class_embed_type" in config) and (config["class_embed_type"] in ["timestep", "projection"]): - class_embed_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["class_embed_type"] - for diffusers_key, ldm_key in class_embed_keys.items(): - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - if ("addition_embed_type" in config) and (config["addition_embed_type"] == "text_time"): - addition_embed_keys = DIFFUSERS_TO_LDM_MAPPING["unet"]["addition_embed_type"] - for diffusers_key, ldm_key in addition_embed_keys.items(): - new_checkpoint[diffusers_key] = unet_state_dict[ldm_key] - - # Relevant to StableDiffusionUpscalePipeline - if "num_class_embeds" in config: - if (config["num_class_embeds"] is not None) and ("label_emb.weight" in unet_state_dict): - new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - # Down blocks - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - unet_state_dict, - {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.get( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.get( - f"input_blocks.{i}.0.op.bias" - ) - - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - unet_state_dict, - {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - # Mid blocks - for key in middle_blocks.keys(): - diffusers_key = max(key - 1, 0) - if key % 2 == 0: - update_unet_resnet_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - unet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.resnets.{diffusers_key}"}, - ) - else: - update_unet_attention_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - unet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.attentions.{diffusers_key}"}, - ) - - # Up Blocks - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - - resnets = [ - key for key in output_blocks[i] if f"output_blocks.{i}.0" in key and f"output_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - unet_state_dict, - {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - attentions = [ - key for key in output_blocks[i] if f"output_blocks.{i}.1" in key and f"output_blocks.{i}.1.conv" not in key - ] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - unet_state_dict, - {"old": f"output_blocks.{i}.1", "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - if f"output_blocks.{i}.1.conv.weight" in unet_state_dict: - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.1.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.1.conv.bias" - ] - if f"output_blocks.{i}.2.conv.weight" in unet_state_dict: - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.2.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.2.conv.bias" - ] - - return new_checkpoint - - -def convert_controlnet_checkpoint( - checkpoint, - config, - **kwargs, -): - # Return checkpoint if it's already been converted - if "time_embedding.linear_1.weight" in checkpoint: - return checkpoint - # Some controlnet ckpt files are distributed independently from the rest of the - # model components i.e. https://huggingface.co/thibaud/controlnet-sd21/ - if "time_embed.0.weight" in checkpoint: - controlnet_state_dict = checkpoint - - else: - controlnet_state_dict = {} - keys = list(checkpoint.keys()) - controlnet_key = LDM_CONTROLNET_KEY - for key in keys: - if key.startswith(controlnet_key): - controlnet_state_dict[key.replace(controlnet_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - ldm_controlnet_keys = DIFFUSERS_TO_LDM_MAPPING["controlnet"]["layers"] - for diffusers_key, ldm_key in ldm_controlnet_keys.items(): - if ldm_key not in controlnet_state_dict: - continue - new_checkpoint[diffusers_key] = controlnet_state_dict[ldm_key] - - # Retrieves the keys for the input blocks only - num_input_blocks = len( - {".".join(layer.split(".")[:2]) for layer in controlnet_state_dict if "input_blocks" in layer} - ) - input_blocks = { - layer_id: [key for key in controlnet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Down blocks - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - update_unet_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - controlnet_state_dict, - {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"}, - ) - - if f"input_blocks.{i}.0.op.weight" in controlnet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = controlnet_state_dict.get( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = controlnet_state_dict.get( - f"input_blocks.{i}.0.op.bias" - ) - - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - if attentions: - update_unet_attention_ldm_to_diffusers( - attentions, - new_checkpoint, - controlnet_state_dict, - {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"}, - ) - - # controlnet down blocks - for i in range(num_input_blocks): - new_checkpoint[f"controlnet_down_blocks.{i}.weight"] = controlnet_state_dict.get(f"zero_convs.{i}.0.weight") - new_checkpoint[f"controlnet_down_blocks.{i}.bias"] = controlnet_state_dict.get(f"zero_convs.{i}.0.bias") - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len( - {".".join(layer.split(".")[:2]) for layer in controlnet_state_dict if "middle_block" in layer} - ) - middle_blocks = { - layer_id: [key for key in controlnet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Mid blocks - for key in middle_blocks.keys(): - diffusers_key = max(key - 1, 0) - if key % 2 == 0: - update_unet_resnet_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - controlnet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.resnets.{diffusers_key}"}, - ) - else: - update_unet_attention_ldm_to_diffusers( - middle_blocks[key], - new_checkpoint, - controlnet_state_dict, - mapping={"old": f"middle_block.{key}", "new": f"mid_block.attentions.{diffusers_key}"}, - ) - - # mid block - new_checkpoint["controlnet_mid_block.weight"] = controlnet_state_dict.get("middle_block_out.0.weight") - new_checkpoint["controlnet_mid_block.bias"] = controlnet_state_dict.get("middle_block_out.0.bias") - - # controlnet cond embedding blocks - cond_embedding_blocks = { - ".".join(layer.split(".")[:2]) - for layer in controlnet_state_dict - if "input_hint_block" in layer and ("input_hint_block.0" not in layer) and ("input_hint_block.14" not in layer) - } - num_cond_embedding_blocks = len(cond_embedding_blocks) - - for idx in range(1, num_cond_embedding_blocks + 1): - diffusers_idx = idx - 1 - cond_block_id = 2 * idx - - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_idx}.weight"] = controlnet_state_dict.get( - f"input_hint_block.{cond_block_id}.weight" - ) - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_idx}.bias"] = controlnet_state_dict.get( - f"input_hint_block.{cond_block_id}.bias" - ) - - return new_checkpoint - - -def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - # remove the LDM_VAE_KEY prefix from the ldm checkpoint keys so that it is easier to map them to diffusers keys - vae_state_dict = {} - keys = list(checkpoint.keys()) - vae_key = "" - for ldm_vae_key in LDM_VAE_KEYS: - if any(k.startswith(ldm_vae_key) for k in keys): - vae_key = ldm_vae_key - - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - vae_diffusers_ldm_map = DIFFUSERS_TO_LDM_MAPPING["vae"] - for diffusers_key, ldm_key in vae_diffusers_ldm_map.items(): - if ldm_key not in vae_state_dict: - continue - new_checkpoint[diffusers_key] = vae_state_dict[ldm_key] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len(config["down_block_types"]) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}, - ) - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.get( - f"encoder.down.{i}.downsample.conv.bias" - ) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len(config["up_block_types"]) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}, - ) - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - update_vae_resnet_ldm_to_diffusers( - resnets, - new_checkpoint, - vae_state_dict, - mapping={"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}, - ) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - update_vae_attentions_ldm_to_diffusers( - mid_attentions, new_checkpoint, vae_state_dict, mapping={"old": "mid.attn_1", "new": "mid_block.attentions.0"} - ) - conv_attn_to_linear(new_checkpoint) - - return new_checkpoint - - def convert_ldm_clip_checkpoint(checkpoint, remove_prefix=None): keys = list(checkpoint.keys()) text_model_dict = {} @@ -1600,70 +987,21 @@ def convert_ldm_clip_checkpoint(checkpoint, remove_prefix=None): return text_model_dict -def convert_open_clip_checkpoint( - text_model, - checkpoint, - prefix="cond_stage_model.model.", -): - text_model_dict = {} - text_proj_key = prefix + "text_projection" - - if text_proj_key in checkpoint: - text_proj_dim = int(checkpoint[text_proj_key].shape[0]) - elif hasattr(text_model.config, "hidden_size"): - text_proj_dim = text_model.config.hidden_size - else: - text_proj_dim = LDM_OPEN_CLIP_TEXT_PROJECTION_DIM - - keys = list(checkpoint.keys()) - keys_to_ignore = SD_2_TEXT_ENCODER_KEYS_TO_IGNORE - - openclip_diffusers_ldm_map = DIFFUSERS_TO_LDM_MAPPING["openclip"]["layers"] - for diffusers_key, ldm_key in openclip_diffusers_ldm_map.items(): - ldm_key = prefix + ldm_key - if ldm_key not in checkpoint: - continue - if ldm_key in keys_to_ignore: - continue - if ldm_key.endswith("text_projection"): - text_model_dict[diffusers_key] = checkpoint[ldm_key].T.contiguous() - else: - text_model_dict[diffusers_key] = checkpoint[ldm_key] - - for key in keys: - if key in keys_to_ignore: - continue - - if not key.startswith(prefix + "transformer."): - continue - - diffusers_key = key.replace(prefix + "transformer.", "") - transformer_diffusers_to_ldm_map = DIFFUSERS_TO_LDM_MAPPING["openclip"]["transformer"] - for new_key, old_key in transformer_diffusers_to_ldm_map.items(): - diffusers_key = ( - diffusers_key.replace(old_key, new_key).replace(".in_proj_weight", "").replace(".in_proj_bias", "") - ) - - if key.endswith(".in_proj_weight"): - weight_value = checkpoint.get(key) - - text_model_dict[diffusers_key + ".q_proj.weight"] = weight_value[:text_proj_dim, :].clone().detach() - text_model_dict[diffusers_key + ".k_proj.weight"] = ( - weight_value[text_proj_dim : text_proj_dim * 2, :].clone().detach() - ) - text_model_dict[diffusers_key + ".v_proj.weight"] = weight_value[text_proj_dim * 2 :, :].clone().detach() - - elif key.endswith(".in_proj_bias"): - weight_value = checkpoint.get(key) - text_model_dict[diffusers_key + ".q_proj.bias"] = weight_value[:text_proj_dim].clone().detach() - text_model_dict[diffusers_key + ".k_proj.bias"] = ( - weight_value[text_proj_dim : text_proj_dim * 2].clone().detach() - ) - text_model_dict[diffusers_key + ".v_proj.bias"] = weight_value[text_proj_dim * 2 :].clone().detach() - else: - text_model_dict[diffusers_key] = checkpoint.get(key) +def convert_open_clip_checkpoint(text_model, checkpoint, prefix="cond_stage_model.model."): + """Select the OpenCLIP text component and apply its shared layout definition.""" + from .conversion import get_conversion - return text_model_dict + config = text_model.config.to_dict() + config["original_format"] = "openclip" + model_class = type(text_model).__name__ + state = {key.removeprefix(prefix): tensor for key, tensor in checkpoint.items() if key.startswith(prefix)} + state.pop("logit_scale", None) + if model_class != "CLIPTextModelWithProjection": + state.pop("text_projection", None) + if prefix == "cond_stage_model.model." and config["num_hidden_layers"] == 23: + # Stable Diffusion 2 conditions on the penultimate layer of the 24-layer OpenCLIP encoder. + state = {key: value for key, value in state.items() if not key.startswith("transformer.resblocks.23.")} + return get_conversion(model_class, config).to_diffusers(state) def create_diffusers_clip_model_from_ldm( @@ -1958,226 +1296,6 @@ def _legacy_load_safety_checker(local_files_only, torch_dtype): # in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; # while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation -def swap_scale_shift(weight, dim): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - -def swap_proj_gate(weight): - proj, gate = weight.chunk(2, dim=0) - new_weight = torch.cat([gate, proj], dim=0) - return new_weight - - -def get_attn2_layers(state_dict): - attn2_layers = [] - for key in state_dict.keys(): - if "attn2." in key: - # Extract the layer number from the key - layer_num = int(key.split(".")[1]) - attn2_layers.append(layer_num) - - return tuple(sorted(set(attn2_layers))) - - -def get_caption_projection_dim(state_dict): - caption_projection_dim = state_dict["context_embedder.weight"].shape[0] - return caption_projection_dim - - -def convert_sd3_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "joint_blocks" in k))[-1] + 1 # noqa: C401 - dual_attention_layers = get_attn2_layers(checkpoint) - - caption_projection_dim = get_caption_projection_dim(checkpoint) - has_qk_norm = any("ln_q" in key for key in checkpoint.keys()) - - # Positional and patch embeddings. - converted_state_dict["pos_embed.pos_embed"] = checkpoint.pop("pos_embed") - converted_state_dict["pos_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["pos_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Timestep embeddings. - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - - # Context projections. - converted_state_dict["context_embedder.weight"] = checkpoint.pop("context_embedder.weight") - converted_state_dict["context_embedder.bias"] = checkpoint.pop("context_embedder.bias") - - # Pooled context projection. - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = checkpoint.pop("y_embedder.mlp.0.weight") - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = checkpoint.pop("y_embedder.mlp.0.bias") - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = checkpoint.pop("y_embedder.mlp.2.weight") - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = checkpoint.pop("y_embedder.mlp.2.bias") - - # Transformer blocks 🎸. - for i in range(num_layers): - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn.qkv.weight"), 3, dim=0 - ) - context_q, context_k, context_v = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.context_block.attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.context_block.attn.qkv.bias"), 3, dim=0 - ) - - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.to_v.bias"] = torch.cat([sample_v_bias]) - - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"transformer_blocks.{i}.attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn.norm_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.ln_k.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.norm_added_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_out.0.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn.proj.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn.to_add_out.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.attn.proj.bias" - ) - - if i in dual_attention_layers: - # Q, K, V - sample_q2, sample_k2, sample_v2 = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn2.qkv.weight"), 3, dim=0 - ) - sample_q2_bias, sample_k2_bias, sample_v2_bias = torch.chunk( - checkpoint.pop(f"joint_blocks.{i}.x_block.attn2.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.weight"] = torch.cat([sample_q2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.bias"] = torch.cat([sample_q2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.weight"] = torch.cat([sample_k2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.bias"] = torch.cat([sample_k2_bias]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.weight"] = torch.cat([sample_v2]) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.bias"] = torch.cat([sample_v2_bias]) - - # qk norm - if has_qk_norm: - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_q.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.ln_q.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.norm_k.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.ln_k.weight" - ) - - # output projections. - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.attn2.proj.bias" - ) - - # norms. - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1.linear.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.adaLN_modulation.1.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias" - ) - else: - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.weight"] = swap_scale_shift( - checkpoint.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.weight"), - dim=caption_projection_dim, - ) - converted_state_dict[f"transformer_blocks.{i}.norm1_context.linear.bias"] = swap_scale_shift( - checkpoint.pop(f"joint_blocks.{i}.context_block.adaLN_modulation.1.bias"), - dim=caption_projection_dim, - ) - - # ffs. - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.0.proj.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.weight"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.net.2.bias"] = checkpoint.pop( - f"joint_blocks.{i}.x_block.mlp.fc2.bias" - ) - if not (i == num_layers - 1): - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.0.proj.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc1.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.weight"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff_context.net.2.bias"] = checkpoint.pop( - f"joint_blocks.{i}.context_block.mlp.fc2.bias" - ) - - # Final blocks. - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.weight"), dim=caption_projection_dim - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.bias"), dim=caption_projection_dim - ) - - return converted_state_dict def is_t5_in_single_file(checkpoint): @@ -2241,1984 +1359,3 @@ def create_diffusers_t5_model_from_checkpoint( param.data = param.data.to(torch.float32) return model - - -def convert_animatediff_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - for k, v in checkpoint.items(): - if "pos_encoder" in k: - continue - - else: - converted_state_dict[ - k.replace(".norms.0", ".norm1") - .replace(".norms.1", ".norm2") - .replace(".ff_norm", ".norm3") - .replace(".attention_blocks.0", ".attn1") - .replace(".attention_blocks.1", ".attn2") - .replace(".temporal_transformer", "") - ] = v - - return converted_state_dict - - -def convert_flux_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - keys = list(checkpoint.keys()) - - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "double_blocks." in k))[-1] + 1 # noqa: C401 - num_single_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "single_blocks." in k))[-1] + 1 # noqa: C401 - mlp_ratio = 4.0 - inner_dim = 3072 - - # in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; - # while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation - def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - ## time_text_embed.timestep_embedder <- time_in - converted_state_dict["time_text_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "time_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("time_in.in_layer.bias") - converted_state_dict["time_text_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "time_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("time_in.out_layer.bias") - - ## time_text_embed.text_embedder <- vector_in - converted_state_dict["time_text_embed.text_embedder.linear_1.weight"] = checkpoint.pop("vector_in.in_layer.weight") - converted_state_dict["time_text_embed.text_embedder.linear_1.bias"] = checkpoint.pop("vector_in.in_layer.bias") - converted_state_dict["time_text_embed.text_embedder.linear_2.weight"] = checkpoint.pop( - "vector_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.text_embedder.linear_2.bias"] = checkpoint.pop("vector_in.out_layer.bias") - - # guidance - has_guidance = any("guidance" in k for k in checkpoint) - if has_guidance: - converted_state_dict["time_text_embed.guidance_embedder.linear_1.weight"] = checkpoint.pop( - "guidance_in.in_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_1.bias"] = checkpoint.pop( - "guidance_in.in_layer.bias" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.weight"] = checkpoint.pop( - "guidance_in.out_layer.weight" - ) - converted_state_dict["time_text_embed.guidance_embedder.linear_2.bias"] = checkpoint.pop( - "guidance_in.out_layer.bias" - ) - - # context_embedder - converted_state_dict["context_embedder.weight"] = checkpoint.pop("txt_in.weight") - converted_state_dict["context_embedder.bias"] = checkpoint.pop("txt_in.bias") - - # x_embedder - converted_state_dict["x_embedder.weight"] = checkpoint.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = checkpoint.pop("img_in.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - # norms. - ## norm1 - converted_state_dict[f"{block_prefix}norm1.linear.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1.linear.bias"] = checkpoint.pop( - f"double_blocks.{i}.img_mod.lin.bias" - ) - ## norm1_context - converted_state_dict[f"{block_prefix}norm1_context.linear.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mod.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm1_context.linear.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mod.lin.bias" - ) - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk(checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0) - context_q, context_k, context_v = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - # qk_norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.norm.key_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.norm.key_norm.scale" - ) - # ff img_mlp - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.0.bias") - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.weight") - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.bias") - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.0.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.2.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.2.bias" - ) - # output projections. - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.proj.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.proj.bias" - ) - - # single transformer blocks - for i in range(num_single_layers): - block_prefix = f"single_transformer_blocks.{i}." - # norm.linear <- single_blocks.0.modulation.lin - converted_state_dict[f"{block_prefix}norm.linear.weight"] = checkpoint.pop( - f"single_blocks.{i}.modulation.lin.weight" - ) - converted_state_dict[f"{block_prefix}norm.linear.bias"] = checkpoint.pop( - f"single_blocks.{i}.modulation.lin.bias" - ) - # Q, K, V, mlp - mlp_hidden_dim = int(inner_dim * mlp_ratio) - split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) - q, k, v, mlp = torch.split(checkpoint.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0) - q_bias, k_bias, v_bias, mlp_bias = torch.split( - checkpoint.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp]) - converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias]) - # qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( - f"single_blocks.{i}.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( - f"single_blocks.{i}.norm.key_norm.scale" - ) - # output projections. - converted_state_dict[f"{block_prefix}proj_out.weight"] = checkpoint.pop(f"single_blocks.{i}.linear2.weight") - converted_state_dict[f"{block_prefix}proj_out.bias"] = checkpoint.pop(f"single_blocks.{i}.linear2.bias") - - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.weight") - ) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift( - checkpoint.pop("final_layer.adaLN_modulation.1.bias") - ) - - return converted_state_dict - - -def convert_ltx_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if "vae" not in key} - - TRANSFORMER_KEYS_RENAME_DICT = { - "model.diffusion_model.": "", - "patchify_proj": "proj_in", - "adaln_single": "time_embed", - "q_norm": "norm_q", - "k_norm": "norm_k", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = {} - - for key in list(converted_state_dict.keys()): - new_key = key - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_ltx_vae_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys()) if "vae." in key} - - def remove_keys_(key: str, state_dict): - state_dict.pop(key) - - VAE_KEYS_RENAME_DICT = { - # common - "vae.": "", - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0", - "up_blocks.2": "up_blocks.1.upsamplers.0", - "up_blocks.3": "up_blocks.1", - "up_blocks.4": "up_blocks.2.conv_in", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.conv_in", - "up_blocks.8": "up_blocks.3.upsamplers.0", - "up_blocks.9": "up_blocks.3", - # encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.0.conv_out", - "down_blocks.3": "down_blocks.1", - "down_blocks.4": "down_blocks.1.downsamplers.0", - "down_blocks.5": "down_blocks.1.conv_out", - "down_blocks.6": "down_blocks.2", - "down_blocks.7": "down_blocks.2.downsamplers.0", - "down_blocks.8": "down_blocks.3", - "down_blocks.9": "mid_block", - # common - "conv_shortcut": "conv_shortcut.conv", - "res_blocks": "resnets", - "norm3.norm": "norm3", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", - } - - VAE_091_RENAME_DICT = { - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", - } - - VAE_095_RENAME_DICT = { - # decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - "up_blocks.7": "up_blocks.3.upsamplers.0", - "up_blocks.8": "up_blocks.3", - # encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.1", - "down_blocks.3": "down_blocks.1.downsamplers.0", - "down_blocks.4": "down_blocks.2", - "down_blocks.5": "down_blocks.2.downsamplers.0", - "down_blocks.6": "down_blocks.3", - "down_blocks.7": "down_blocks.3.downsamplers.0", - "down_blocks.8": "mid_block", - # common - "last_time_embedder": "time_embedder", - "last_scale_shift_table": "scale_shift_table", - } - - VAE_SPECIAL_KEYS_REMAP = { - "per_channel_statistics.channel": remove_keys_, - "per_channel_statistics.mean-of-means": remove_keys_, - "per_channel_statistics.mean-of-stds": remove_keys_, - } - - if converted_state_dict["vae.encoder.conv_out.conv.weight"].shape[1] == 2048: - VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT) - elif "vae.decoder.last_time_embedder.timestep_embedder.linear_1.weight" in converted_state_dict: - VAE_KEYS_RENAME_DICT.update(VAE_091_RENAME_DICT) - - for key in list(converted_state_dict.keys()): - new_key = key - for replace_key, rename_key in VAE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_autoencoder_dc_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - def remap_qkv_(key: str, state_dict): - qkv = state_dict.pop(key) - q, k, v = torch.chunk(qkv, 3, dim=0) - parent_module, _, _ = key.rpartition(".qkv.conv.weight") - state_dict[f"{parent_module}.to_q.weight"] = q.squeeze() - state_dict[f"{parent_module}.to_k.weight"] = k.squeeze() - state_dict[f"{parent_module}.to_v.weight"] = v.squeeze() - - def remap_proj_conv_(key: str, state_dict): - parent_module, _, _ = key.rpartition(".proj.conv.weight") - state_dict[f"{parent_module}.to_out.weight"] = state_dict.pop(key).squeeze() - - AE_KEYS_RENAME_DICT = { - # common - "main.": "", - "op_list.": "", - "context_module": "attn", - "local_module": "conv_out", - # NOTE: The below two lines work because scales in the available configs only have a tuple length of 1 - # If there were more scales, there would be more layers, so a loop would be better to handle this - "aggreg.0.0": "to_qkv_multiscale.0.proj_in", - "aggreg.0.1": "to_qkv_multiscale.0.proj_out", - "depth_conv.conv": "conv_depth", - "inverted_conv.conv": "conv_inverted", - "point_conv.conv": "conv_point", - "point_conv.norm": "norm", - "conv.conv.": "conv.", - "conv1.conv": "conv1", - "conv2.conv": "conv2", - "conv2.norm": "norm", - "proj.norm": "norm_out", - # encoder - "encoder.project_in.conv": "encoder.conv_in", - "encoder.project_out.0.conv": "encoder.conv_out", - "encoder.stages": "encoder.down_blocks", - # decoder - "decoder.project_in.conv": "decoder.conv_in", - "decoder.project_out.0": "decoder.norm_out", - "decoder.project_out.2.conv": "decoder.conv_out", - "decoder.stages": "decoder.up_blocks", - } - - AE_F32C32_F64C128_F128C512_KEYS = { - "encoder.project_in.conv": "encoder.conv_in.conv", - "decoder.project_out.2.conv": "decoder.conv_out.conv", - } - - AE_SPECIAL_KEYS_REMAP = { - "qkv.conv.weight": remap_qkv_, - "proj.conv.weight": remap_proj_conv_, - } - if "encoder.project_in.conv.bias" not in converted_state_dict: - AE_KEYS_RENAME_DICT.update(AE_F32C32_F64C128_F128C512_KEYS) - - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in AE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in AE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_mochi_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Comfy checkpoints add this prefix - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - # Convert patch_embed - converted_state_dict["patch_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Convert time_embed - converted_state_dict["time_embed.timestep_embedder.linear_1.weight"] = checkpoint.pop("t_embedder.mlp.0.weight") - converted_state_dict["time_embed.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.timestep_embedder.linear_2.weight"] = checkpoint.pop("t_embedder.mlp.2.weight") - converted_state_dict["time_embed.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - converted_state_dict["time_embed.pooler.to_kv.weight"] = checkpoint.pop("t5_y_embedder.to_kv.weight") - converted_state_dict["time_embed.pooler.to_kv.bias"] = checkpoint.pop("t5_y_embedder.to_kv.bias") - converted_state_dict["time_embed.pooler.to_q.weight"] = checkpoint.pop("t5_y_embedder.to_q.weight") - converted_state_dict["time_embed.pooler.to_q.bias"] = checkpoint.pop("t5_y_embedder.to_q.bias") - converted_state_dict["time_embed.pooler.to_out.weight"] = checkpoint.pop("t5_y_embedder.to_out.weight") - converted_state_dict["time_embed.pooler.to_out.bias"] = checkpoint.pop("t5_y_embedder.to_out.bias") - converted_state_dict["time_embed.caption_proj.weight"] = checkpoint.pop("t5_yproj.weight") - converted_state_dict["time_embed.caption_proj.bias"] = checkpoint.pop("t5_yproj.bias") - - # Convert transformer blocks - num_layers = 48 - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - old_prefix = f"blocks.{i}." - - # norm1 - converted_state_dict[block_prefix + "norm1.linear.weight"] = checkpoint.pop(old_prefix + "mod_x.weight") - converted_state_dict[block_prefix + "norm1.linear.bias"] = checkpoint.pop(old_prefix + "mod_x.bias") - if i < num_layers - 1: - converted_state_dict[block_prefix + "norm1_context.linear.weight"] = checkpoint.pop( - old_prefix + "mod_y.weight" - ) - converted_state_dict[block_prefix + "norm1_context.linear.bias"] = checkpoint.pop( - old_prefix + "mod_y.bias" - ) - else: - converted_state_dict[block_prefix + "norm1_context.linear_1.weight"] = checkpoint.pop( - old_prefix + "mod_y.weight" - ) - converted_state_dict[block_prefix + "norm1_context.linear_1.bias"] = checkpoint.pop( - old_prefix + "mod_y.bias" - ) - - # Visual attention - qkv_weight = checkpoint.pop(old_prefix + "attn.qkv_x.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - converted_state_dict[block_prefix + "attn1.to_q.weight"] = q - converted_state_dict[block_prefix + "attn1.to_k.weight"] = k - converted_state_dict[block_prefix + "attn1.to_v.weight"] = v - converted_state_dict[block_prefix + "attn1.norm_q.weight"] = checkpoint.pop( - old_prefix + "attn.q_norm_x.weight" - ) - converted_state_dict[block_prefix + "attn1.norm_k.weight"] = checkpoint.pop( - old_prefix + "attn.k_norm_x.weight" - ) - converted_state_dict[block_prefix + "attn1.to_out.0.weight"] = checkpoint.pop( - old_prefix + "attn.proj_x.weight" - ) - converted_state_dict[block_prefix + "attn1.to_out.0.bias"] = checkpoint.pop(old_prefix + "attn.proj_x.bias") - - # Context attention - qkv_weight = checkpoint.pop(old_prefix + "attn.qkv_y.weight") - q, k, v = qkv_weight.chunk(3, dim=0) - - converted_state_dict[block_prefix + "attn1.add_q_proj.weight"] = q - converted_state_dict[block_prefix + "attn1.add_k_proj.weight"] = k - converted_state_dict[block_prefix + "attn1.add_v_proj.weight"] = v - converted_state_dict[block_prefix + "attn1.norm_added_q.weight"] = checkpoint.pop( - old_prefix + "attn.q_norm_y.weight" - ) - converted_state_dict[block_prefix + "attn1.norm_added_k.weight"] = checkpoint.pop( - old_prefix + "attn.k_norm_y.weight" - ) - if i < num_layers - 1: - converted_state_dict[block_prefix + "attn1.to_add_out.weight"] = checkpoint.pop( - old_prefix + "attn.proj_y.weight" - ) - converted_state_dict[block_prefix + "attn1.to_add_out.bias"] = checkpoint.pop( - old_prefix + "attn.proj_y.bias" - ) - - # MLP - converted_state_dict[block_prefix + "ff.net.0.proj.weight"] = swap_proj_gate( - checkpoint.pop(old_prefix + "mlp_x.w1.weight") - ) - converted_state_dict[block_prefix + "ff.net.2.weight"] = checkpoint.pop(old_prefix + "mlp_x.w2.weight") - if i < num_layers - 1: - converted_state_dict[block_prefix + "ff_context.net.0.proj.weight"] = swap_proj_gate( - checkpoint.pop(old_prefix + "mlp_y.w1.weight") - ) - converted_state_dict[block_prefix + "ff_context.net.2.weight"] = checkpoint.pop( - old_prefix + "mlp_y.w2.weight" - ) - - # Output layers - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(checkpoint.pop("final_layer.mod.weight"), dim=0) - converted_state_dict["norm_out.linear.bias"] = swap_scale_shift(checkpoint.pop("final_layer.mod.bias"), dim=0) - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - - converted_state_dict["pos_frequencies"] = checkpoint.pop("pos_frequencies") - - return converted_state_dict - - -def convert_hunyuan_video_transformer_to_diffusers(checkpoint, **kwargs): - def remap_norm_scale_shift_(key, state_dict): - weight = state_dict.pop(key) - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - state_dict[key.replace("final_layer.adaLN_modulation.1", "norm_out.linear")] = new_weight - - def remap_txt_in_(key, state_dict): - def rename_key(key): - new_key = key.replace("individual_token_refiner.blocks", "token_refiner.refiner_blocks") - new_key = new_key.replace("adaLN_modulation.1", "norm_out.linear") - new_key = new_key.replace("txt_in", "context_embedder") - new_key = new_key.replace("t_embedder.mlp.0", "time_text_embed.timestep_embedder.linear_1") - new_key = new_key.replace("t_embedder.mlp.2", "time_text_embed.timestep_embedder.linear_2") - new_key = new_key.replace("c_embedder", "time_text_embed.text_embedder") - new_key = new_key.replace("mlp", "ff") - return new_key - - if "self_attn_qkv" in key: - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_q"))] = to_q - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_k"))] = to_k - state_dict[rename_key(key.replace("self_attn_qkv", "attn.to_v"))] = to_v - else: - state_dict[rename_key(key)] = state_dict.pop(key) - - def remap_img_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("img_attn_qkv", "attn.to_q")] = to_q - state_dict[key.replace("img_attn_qkv", "attn.to_k")] = to_k - state_dict[key.replace("img_attn_qkv", "attn.to_v")] = to_v - - def remap_txt_attn_qkv_(key, state_dict): - weight = state_dict.pop(key) - to_q, to_k, to_v = weight.chunk(3, dim=0) - state_dict[key.replace("txt_attn_qkv", "attn.add_q_proj")] = to_q - state_dict[key.replace("txt_attn_qkv", "attn.add_k_proj")] = to_k - state_dict[key.replace("txt_attn_qkv", "attn.add_v_proj")] = to_v - - def remap_single_transformer_blocks_(key, state_dict): - hidden_size = 3072 - - if "linear1.weight" in key: - linear1_weight = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_weight.size(0) - 3 * hidden_size) - q, k, v, mlp = torch.split(linear1_weight, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.weight") - state_dict[f"{new_key}.attn.to_q.weight"] = q - state_dict[f"{new_key}.attn.to_k.weight"] = k - state_dict[f"{new_key}.attn.to_v.weight"] = v - state_dict[f"{new_key}.proj_mlp.weight"] = mlp - - elif "linear1.bias" in key: - linear1_bias = state_dict.pop(key) - split_size = (hidden_size, hidden_size, hidden_size, linear1_bias.size(0) - 3 * hidden_size) - q_bias, k_bias, v_bias, mlp_bias = torch.split(linear1_bias, split_size, dim=0) - new_key = key.replace("single_blocks", "single_transformer_blocks").removesuffix(".linear1.bias") - state_dict[f"{new_key}.attn.to_q.bias"] = q_bias - state_dict[f"{new_key}.attn.to_k.bias"] = k_bias - state_dict[f"{new_key}.attn.to_v.bias"] = v_bias - state_dict[f"{new_key}.proj_mlp.bias"] = mlp_bias - - else: - new_key = key.replace("single_blocks", "single_transformer_blocks") - new_key = new_key.replace("linear2", "proj_out") - new_key = new_key.replace("q_norm", "attn.norm_q") - new_key = new_key.replace("k_norm", "attn.norm_k") - state_dict[new_key] = state_dict.pop(key) - - TRANSFORMER_KEYS_RENAME_DICT = { - "img_in": "x_embedder", - "time_in.mlp.0": "time_text_embed.timestep_embedder.linear_1", - "time_in.mlp.2": "time_text_embed.timestep_embedder.linear_2", - "guidance_in.mlp.0": "time_text_embed.guidance_embedder.linear_1", - "guidance_in.mlp.2": "time_text_embed.guidance_embedder.linear_2", - "vector_in.in_layer": "time_text_embed.text_embedder.linear_1", - "vector_in.out_layer": "time_text_embed.text_embedder.linear_2", - "double_blocks": "transformer_blocks", - "img_attn_q_norm": "attn.norm_q", - "img_attn_k_norm": "attn.norm_k", - "img_attn_proj": "attn.to_out.0", - "txt_attn_q_norm": "attn.norm_added_q", - "txt_attn_k_norm": "attn.norm_added_k", - "txt_attn_proj": "attn.to_add_out", - "img_mod.linear": "norm1.linear", - "img_norm1": "norm1.norm", - "img_norm2": "norm2", - "img_mlp": "ff", - "txt_mod.linear": "norm1_context.linear", - "txt_norm1": "norm1.norm", - "txt_norm2": "norm2_context", - "txt_mlp": "ff_context", - "self_attn_proj": "attn.to_out.0", - "modulation.linear": "norm.linear", - "pre_norm": "norm.norm", - "final_layer.norm_final": "norm_out.norm", - "final_layer.linear": "proj_out", - "fc1": "net.0.proj", - "fc2": "net.2", - "input_embedder": "proj_in", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - "txt_in": remap_txt_in_, - "img_attn_qkv": remap_img_attn_qkv_, - "txt_attn_qkv": remap_txt_attn_qkv_, - "single_blocks": remap_single_transformer_blocks_, - "final_layer.adaLN_modulation.1": remap_norm_scale_shift_, - } - - def update_state_dict_(state_dict, old_key, new_key): - state_dict[new_key] = state_dict.pop(old_key) - - for key in list(checkpoint.keys()): - new_key = key[:] - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - update_state_dict_(checkpoint, key, new_key) - - for key in list(checkpoint.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, checkpoint) - - return checkpoint - - -def convert_auraflow_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - state_dict_keys = list(checkpoint.keys()) - - # Handle register tokens and positional embeddings - converted_state_dict["register_tokens"] = checkpoint.pop("register_tokens", None) - - # Handle time step projection - converted_state_dict["time_step_proj.linear_1.weight"] = checkpoint.pop("t_embedder.mlp.0.weight", None) - converted_state_dict["time_step_proj.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias", None) - converted_state_dict["time_step_proj.linear_2.weight"] = checkpoint.pop("t_embedder.mlp.2.weight", None) - converted_state_dict["time_step_proj.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias", None) - - # Handle context embedder - converted_state_dict["context_embedder.weight"] = checkpoint.pop("cond_seq_linear.weight", None) - - # Calculate the number of layers - def calculate_layers(keys, key_prefix): - layers = set() - for k in keys: - if key_prefix in k: - layer_num = int(k.split(".")[1]) # get the layer number - layers.add(layer_num) - return len(layers) - - mmdit_layers = calculate_layers(state_dict_keys, key_prefix="double_layers") - single_dit_layers = calculate_layers(state_dict_keys, key_prefix="single_layers") - - # MMDiT blocks - for i in range(mmdit_layers): - # Feed-forward - path_mapping = {"mlpX": "ff", "mlpC": "ff_context"} - weight_mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for orig_k, diffuser_k in path_mapping.items(): - for k, v in weight_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.{v}.weight"] = checkpoint.pop( - f"double_layers.{i}.{orig_k}.{k}.weight", None - ) - - # Norms - path_mapping = {"modX": "norm1", "modC": "norm1_context"} - for orig_k, diffuser_k in path_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.{diffuser_k}.linear.weight"] = checkpoint.pop( - f"double_layers.{i}.{orig_k}.1.weight", None - ) - - # Attentions - x_attn_mapping = {"w2q": "to_q", "w2k": "to_k", "w2v": "to_v", "w2o": "to_out.0"} - context_attn_mapping = {"w1q": "add_q_proj", "w1k": "add_k_proj", "w1v": "add_v_proj", "w1o": "to_add_out"} - for attn_mapping in [x_attn_mapping, context_attn_mapping]: - for k, v in attn_mapping.items(): - converted_state_dict[f"joint_transformer_blocks.{i}.attn.{v}.weight"] = checkpoint.pop( - f"double_layers.{i}.attn.{k}.weight", None - ) - - # Single-DiT blocks - for i in range(single_dit_layers): - # Feed-forward - mapping = {"c_fc1": "linear_1", "c_fc2": "linear_2", "c_proj": "out_projection"} - for k, v in mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.ff.{v}.weight"] = checkpoint.pop( - f"single_layers.{i}.mlp.{k}.weight", None - ) - - # Norms - converted_state_dict[f"single_transformer_blocks.{i}.norm1.linear.weight"] = checkpoint.pop( - f"single_layers.{i}.modCX.1.weight", None - ) - - # Attentions - x_attn_mapping = {"w1q": "to_q", "w1k": "to_k", "w1v": "to_v", "w1o": "to_out.0"} - for k, v in x_attn_mapping.items(): - converted_state_dict[f"single_transformer_blocks.{i}.attn.{v}.weight"] = checkpoint.pop( - f"single_layers.{i}.attn.{k}.weight", None - ) - # Final blocks - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_linear.weight", None) - - # Handle the final norm layer - norm_weight = checkpoint.pop("modF.1.weight", None) - if norm_weight is not None: - converted_state_dict["norm_out.linear.weight"] = swap_scale_shift(norm_weight, dim=None) - else: - converted_state_dict["norm_out.linear.weight"] = None - - converted_state_dict["pos_embed.pos_embed"] = checkpoint.pop("positional_encoding") - converted_state_dict["pos_embed.proj.weight"] = checkpoint.pop("init_x_linear.weight") - converted_state_dict["pos_embed.proj.bias"] = checkpoint.pop("init_x_linear.bias") - - return converted_state_dict - - -def convert_lumina2_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Original Lumina-Image-2 has an extra norm parameter that is unused - # We just remove it here - checkpoint.pop("norm_final.weight", None) - - # Comfy checkpoints add this prefix - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - LUMINA_KEY_MAP = { - "cap_embedder": "time_caption_embed.caption_embedder", - "t_embedder.mlp.0": "time_caption_embed.timestep_embedder.linear_1", - "t_embedder.mlp.2": "time_caption_embed.timestep_embedder.linear_2", - "attention": "attn", - ".out.": ".to_out.0.", - "k_norm": "norm_k", - "q_norm": "norm_q", - "w1": "linear_1", - "w2": "linear_2", - "w3": "linear_3", - "adaLN_modulation.1": "norm1.linear", - } - ATTENTION_NORM_MAP = { - "attention_norm1": "norm1.norm", - "attention_norm2": "norm2", - } - CONTEXT_REFINER_MAP = { - "context_refiner.0.attention_norm1": "context_refiner.0.norm1", - "context_refiner.0.attention_norm2": "context_refiner.0.norm2", - "context_refiner.1.attention_norm1": "context_refiner.1.norm1", - "context_refiner.1.attention_norm2": "context_refiner.1.norm2", - } - FINAL_LAYER_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear_1", - "final_layer.linear": "norm_out.linear_2", - } - - def convert_lumina_attn_to_diffusers(tensor, diffusers_key): - q_dim = 2304 - k_dim = v_dim = 768 - - to_q, to_k, to_v = torch.split(tensor, [q_dim, k_dim, v_dim], dim=0) - - return { - diffusers_key.replace("qkv", "to_q"): to_q, - diffusers_key.replace("qkv", "to_k"): to_k, - diffusers_key.replace("qkv", "to_v"): to_v, - } - - for key in keys: - diffusers_key = key - for k, v in CONTEXT_REFINER_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in FINAL_LAYER_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in ATTENTION_NORM_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - for k, v in LUMINA_KEY_MAP.items(): - diffusers_key = diffusers_key.replace(k, v) - - if "qkv" in diffusers_key: - converted_state_dict.update(convert_lumina_attn_to_diffusers(checkpoint.pop(key), diffusers_key)) - else: - converted_state_dict[diffusers_key] = checkpoint.pop(key) - - return converted_state_dict - - -def convert_sana_transformer_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "blocks" in k))[-1] + 1 # noqa: C401 - - # Positional and patch embeddings. - checkpoint.pop("pos_embed") - converted_state_dict["patch_embed.proj.weight"] = checkpoint.pop("x_embedder.proj.weight") - converted_state_dict["patch_embed.proj.bias"] = checkpoint.pop("x_embedder.proj.bias") - - # Timestep embeddings. - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.weight"] = checkpoint.pop( - "t_embedder.mlp.0.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_1.bias"] = checkpoint.pop("t_embedder.mlp.0.bias") - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.weight"] = checkpoint.pop( - "t_embedder.mlp.2.weight" - ) - converted_state_dict["time_embed.emb.timestep_embedder.linear_2.bias"] = checkpoint.pop("t_embedder.mlp.2.bias") - converted_state_dict["time_embed.linear.weight"] = checkpoint.pop("t_block.1.weight") - converted_state_dict["time_embed.linear.bias"] = checkpoint.pop("t_block.1.bias") - - # Caption Projection. - checkpoint.pop("y_embedder.y_embedding") - converted_state_dict["caption_projection.linear_1.weight"] = checkpoint.pop("y_embedder.y_proj.fc1.weight") - converted_state_dict["caption_projection.linear_1.bias"] = checkpoint.pop("y_embedder.y_proj.fc1.bias") - converted_state_dict["caption_projection.linear_2.weight"] = checkpoint.pop("y_embedder.y_proj.fc2.weight") - converted_state_dict["caption_projection.linear_2.bias"] = checkpoint.pop("y_embedder.y_proj.fc2.bias") - converted_state_dict["caption_norm.weight"] = checkpoint.pop("attention_y_norm.weight") - - for i in range(num_layers): - converted_state_dict[f"transformer_blocks.{i}.scale_shift_table"] = checkpoint.pop( - f"blocks.{i}.scale_shift_table" - ) - - # Self-Attention - sample_q, sample_k, sample_v = torch.chunk(checkpoint.pop(f"blocks.{i}.attn.qkv.weight"), 3, dim=0) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_v.weight"] = torch.cat([sample_v]) - - # Output Projections - converted_state_dict[f"transformer_blocks.{i}.attn1.to_out.0.weight"] = checkpoint.pop( - f"blocks.{i}.attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn1.to_out.0.bias"] = checkpoint.pop( - f"blocks.{i}.attn.proj.bias" - ) - - # Cross-Attention - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.weight"] = checkpoint.pop( - f"blocks.{i}.cross_attn.q_linear.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_q.bias"] = checkpoint.pop( - f"blocks.{i}.cross_attn.q_linear.bias" - ) - - linear_sample_k, linear_sample_v = torch.chunk( - checkpoint.pop(f"blocks.{i}.cross_attn.kv_linear.weight"), 2, dim=0 - ) - linear_sample_k_bias, linear_sample_v_bias = torch.chunk( - checkpoint.pop(f"blocks.{i}.cross_attn.kv_linear.bias"), 2, dim=0 - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.weight"] = linear_sample_k - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.weight"] = linear_sample_v - converted_state_dict[f"transformer_blocks.{i}.attn2.to_k.bias"] = linear_sample_k_bias - converted_state_dict[f"transformer_blocks.{i}.attn2.to_v.bias"] = linear_sample_v_bias - - # Output Projections - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.weight"] = checkpoint.pop( - f"blocks.{i}.cross_attn.proj.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.attn2.to_out.0.bias"] = checkpoint.pop( - f"blocks.{i}.cross_attn.proj.bias" - ) - - # MLP - converted_state_dict[f"transformer_blocks.{i}.ff.conv_inverted.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.inverted_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_inverted.bias"] = checkpoint.pop( - f"blocks.{i}.mlp.inverted_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_depth.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.depth_conv.conv.weight" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_depth.bias"] = checkpoint.pop( - f"blocks.{i}.mlp.depth_conv.conv.bias" - ) - converted_state_dict[f"transformer_blocks.{i}.ff.conv_point.weight"] = checkpoint.pop( - f"blocks.{i}.mlp.point_conv.conv.weight" - ) - - # Final layer - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - converted_state_dict["scale_shift_table"] = checkpoint.pop("final_layer.scale_shift_table") - - return converted_state_dict - - -def convert_wan_transformer_to_diffusers(checkpoint, **kwargs): - def generate_motion_encoder_mappings(): - mappings = { - "motion_encoder.dec.direction.weight": "motion_encoder.motion_synthesis_weight", - "motion_encoder.enc.net_app.convs.0.0.weight": "motion_encoder.conv_in.weight", - "motion_encoder.enc.net_app.convs.0.1.bias": "motion_encoder.conv_in.act_fn.bias", - "motion_encoder.enc.net_app.convs.8.weight": "motion_encoder.conv_out.weight", - "motion_encoder.enc.fc": "motion_encoder.motion_network", - } - - for i in range(7): - conv_idx = i + 1 - mappings.update( - { - f"motion_encoder.enc.net_app.convs.{conv_idx}.conv1.0.weight": f"motion_encoder.res_blocks.{i}.conv1.weight", - f"motion_encoder.enc.net_app.convs.{conv_idx}.conv1.1.bias": f"motion_encoder.res_blocks.{i}.conv1.act_fn.bias", - f"motion_encoder.enc.net_app.convs.{conv_idx}.conv2.1.weight": f"motion_encoder.res_blocks.{i}.conv2.weight", - f"motion_encoder.enc.net_app.convs.{conv_idx}.conv2.2.bias": f"motion_encoder.res_blocks.{i}.conv2.act_fn.bias", - f"motion_encoder.enc.net_app.convs.{conv_idx}.skip.1.weight": f"motion_encoder.res_blocks.{i}.conv_skip.weight", - } - ) - - return mappings - - def generate_face_adapter_mappings(): - return { - "face_adapter.fuser_blocks": "face_adapter", - ".k_norm.": ".norm_k.", - ".q_norm.": ".norm_q.", - ".linear1_q.": ".to_q.", - ".linear2.": ".to_out.", - "conv1_local.conv": "conv1_local", - "conv2.conv": "conv2", - "conv3.conv": "conv3", - } - - def split_tensor_handler(key, state_dict, split_pattern, target_keys): - tensor = state_dict.pop(key) - split_idx = tensor.shape[0] // 2 - - new_key_1 = key.replace(split_pattern, target_keys[0]) - new_key_2 = key.replace(split_pattern, target_keys[1]) - - state_dict[new_key_1] = tensor[:split_idx] - state_dict[new_key_2] = tensor[split_idx:] - - def reshape_bias_handler(key, state_dict): - if "motion_encoder.enc.net_app.convs." in key and ".bias" in key: - state_dict[key] = state_dict[key][0, :, 0, 0] - - converted_state_dict = {} - - # Strip model.diffusion_model prefix - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - # Base transformer mappings - TRANSFORMER_KEYS_RENAME_DICT = { - "time_embedding.0": "condition_embedder.time_embedder.linear_1", - "time_embedding.2": "condition_embedder.time_embedder.linear_2", - "text_embedding.0": "condition_embedder.text_embedder.linear_1", - "text_embedding.2": "condition_embedder.text_embedder.linear_2", - "time_projection.1": "condition_embedder.time_proj", - "cross_attn": "attn2", - "self_attn": "attn1", - ".o.": ".to_out.0.", - ".q.": ".to_q.", - ".k.": ".to_k.", - ".v.": ".to_v.", - ".k_img.": ".add_k_proj.", - ".v_img.": ".add_v_proj.", - ".norm_k_img.": ".norm_added_k.", - "head.modulation": "scale_shift_table", - "head.head": "proj_out", - "modulation": "scale_shift_table", - "ffn.0": "ffn.net.0.proj", - "ffn.2": "ffn.net.2", - # Hack to swap the layer names - "norm2": "norm__placeholder", - "norm3": "norm2", - "norm__placeholder": "norm3", - # I2V model - "img_emb.proj.0": "condition_embedder.image_embedder.norm1", - "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj", - "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2", - "img_emb.proj.4": "condition_embedder.image_embedder.norm2", - # VACE model - "before_proj": "proj_in", - "after_proj": "proj_out", - } - - SPECIAL_KEYS_HANDLERS = {} - if any("face_adapter" in k for k in checkpoint.keys()): - TRANSFORMER_KEYS_RENAME_DICT.update(generate_face_adapter_mappings()) - SPECIAL_KEYS_HANDLERS[".linear1_kv."] = (split_tensor_handler, [".to_k.", ".to_v."]) - - if any("motion_encoder" in k for k in checkpoint.keys()): - TRANSFORMER_KEYS_RENAME_DICT.update(generate_motion_encoder_mappings()) - - for key in list(checkpoint.keys()): - reshape_bias_handler(key, checkpoint) - - for key in list(checkpoint.keys()): - new_key = key - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = checkpoint.pop(key) - - for key in list(converted_state_dict.keys()): - for pattern, (handler_fn, target_keys) in SPECIAL_KEYS_HANDLERS.items(): - if pattern not in key: - continue - handler_fn(key, converted_state_dict, pattern, target_keys) - break - - return converted_state_dict - - -def convert_wan_animate_2_transformer_to_diffusers(checkpoint, **kwargs): - r""" - Converts the state dict of the Wan-Animate-2 transformer from the official checkpoint format to the diffusers - format. - """ - attention_renames = { - ".q.": ".to_q.", - ".k.": ".to_k.", - ".v.": ".to_v.", - ".o.": ".to_out.0.", - ".k_img.": ".add_k_proj.", - ".v_img.": ".add_v_proj.", - ".norm_k_img.": ".norm_added_k.", - } - - converted_state_dict = {} - for key in list(checkpoint.keys()): - new_key = key.replace("model.diffusion_model.", "") - # The official checkpoint wraps every transformer block in an in-context module the - # diffusers layout does not have: `blocks.N.block.X` -> `blocks.N.X`. - if new_key.startswith("blocks."): - new_key = new_key.replace(".block.", ".", 1) - if ".self_attn." in new_key or ".cross_attn." in new_key: - for old, new in attention_renames.items(): - if old in new_key: - new_key = new_key.replace(old, new) - break - converted_state_dict[new_key] = checkpoint.pop(key) - - return converted_state_dict - - -def convert_wan_vae_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - - # Create mappings for specific components - middle_key_mapping = { - # Encoder middle block - "encoder.middle.0.residual.0.gamma": "encoder.mid_block.resnets.0.norm1.gamma", - "encoder.middle.0.residual.2.bias": "encoder.mid_block.resnets.0.conv1.bias", - "encoder.middle.0.residual.2.weight": "encoder.mid_block.resnets.0.conv1.weight", - "encoder.middle.0.residual.3.gamma": "encoder.mid_block.resnets.0.norm2.gamma", - "encoder.middle.0.residual.6.bias": "encoder.mid_block.resnets.0.conv2.bias", - "encoder.middle.0.residual.6.weight": "encoder.mid_block.resnets.0.conv2.weight", - "encoder.middle.2.residual.0.gamma": "encoder.mid_block.resnets.1.norm1.gamma", - "encoder.middle.2.residual.2.bias": "encoder.mid_block.resnets.1.conv1.bias", - "encoder.middle.2.residual.2.weight": "encoder.mid_block.resnets.1.conv1.weight", - "encoder.middle.2.residual.3.gamma": "encoder.mid_block.resnets.1.norm2.gamma", - "encoder.middle.2.residual.6.bias": "encoder.mid_block.resnets.1.conv2.bias", - "encoder.middle.2.residual.6.weight": "encoder.mid_block.resnets.1.conv2.weight", - # Decoder middle block - "decoder.middle.0.residual.0.gamma": "decoder.mid_block.resnets.0.norm1.gamma", - "decoder.middle.0.residual.2.bias": "decoder.mid_block.resnets.0.conv1.bias", - "decoder.middle.0.residual.2.weight": "decoder.mid_block.resnets.0.conv1.weight", - "decoder.middle.0.residual.3.gamma": "decoder.mid_block.resnets.0.norm2.gamma", - "decoder.middle.0.residual.6.bias": "decoder.mid_block.resnets.0.conv2.bias", - "decoder.middle.0.residual.6.weight": "decoder.mid_block.resnets.0.conv2.weight", - "decoder.middle.2.residual.0.gamma": "decoder.mid_block.resnets.1.norm1.gamma", - "decoder.middle.2.residual.2.bias": "decoder.mid_block.resnets.1.conv1.bias", - "decoder.middle.2.residual.2.weight": "decoder.mid_block.resnets.1.conv1.weight", - "decoder.middle.2.residual.3.gamma": "decoder.mid_block.resnets.1.norm2.gamma", - "decoder.middle.2.residual.6.bias": "decoder.mid_block.resnets.1.conv2.bias", - "decoder.middle.2.residual.6.weight": "decoder.mid_block.resnets.1.conv2.weight", - } - - # Create a mapping for attention blocks - attention_mapping = { - # Encoder middle attention - "encoder.middle.1.norm.gamma": "encoder.mid_block.attentions.0.norm.gamma", - "encoder.middle.1.to_qkv.weight": "encoder.mid_block.attentions.0.to_qkv.weight", - "encoder.middle.1.to_qkv.bias": "encoder.mid_block.attentions.0.to_qkv.bias", - "encoder.middle.1.proj.weight": "encoder.mid_block.attentions.0.proj.weight", - "encoder.middle.1.proj.bias": "encoder.mid_block.attentions.0.proj.bias", - # Decoder middle attention - "decoder.middle.1.norm.gamma": "decoder.mid_block.attentions.0.norm.gamma", - "decoder.middle.1.to_qkv.weight": "decoder.mid_block.attentions.0.to_qkv.weight", - "decoder.middle.1.to_qkv.bias": "decoder.mid_block.attentions.0.to_qkv.bias", - "decoder.middle.1.proj.weight": "decoder.mid_block.attentions.0.proj.weight", - "decoder.middle.1.proj.bias": "decoder.mid_block.attentions.0.proj.bias", - } - - # Create a mapping for the head components - head_mapping = { - # Encoder head - "encoder.head.0.gamma": "encoder.norm_out.gamma", - "encoder.head.2.bias": "encoder.conv_out.bias", - "encoder.head.2.weight": "encoder.conv_out.weight", - # Decoder head - "decoder.head.0.gamma": "decoder.norm_out.gamma", - "decoder.head.2.bias": "decoder.conv_out.bias", - "decoder.head.2.weight": "decoder.conv_out.weight", - } - - # Create a mapping for the quant components - quant_mapping = { - "conv1.weight": "quant_conv.weight", - "conv1.bias": "quant_conv.bias", - "conv2.weight": "post_quant_conv.weight", - "conv2.bias": "post_quant_conv.bias", - } - - # Process each key in the state dict - for key, value in checkpoint.items(): - # Handle middle block keys using the mapping - if key in middle_key_mapping: - new_key = middle_key_mapping[key] - converted_state_dict[new_key] = value - # Handle attention blocks using the mapping - elif key in attention_mapping: - new_key = attention_mapping[key] - converted_state_dict[new_key] = value - # Handle head keys using the mapping - elif key in head_mapping: - new_key = head_mapping[key] - converted_state_dict[new_key] = value - # Handle quant keys using the mapping - elif key in quant_mapping: - new_key = quant_mapping[key] - converted_state_dict[new_key] = value - # Handle encoder conv1 - elif key == "encoder.conv1.weight": - converted_state_dict["encoder.conv_in.weight"] = value - elif key == "encoder.conv1.bias": - converted_state_dict["encoder.conv_in.bias"] = value - # Handle decoder conv1 - elif key == "decoder.conv1.weight": - converted_state_dict["decoder.conv_in.weight"] = value - elif key == "decoder.conv1.bias": - converted_state_dict["decoder.conv_in.bias"] = value - # Handle encoder downsamples - elif key.startswith("encoder.downsamples."): - # Convert to down_blocks - new_key = key.replace("encoder.downsamples.", "encoder.down_blocks.") - - # Convert residual block naming but keep the original structure - if ".residual.0.gamma" in new_key: - new_key = new_key.replace(".residual.0.gamma", ".norm1.gamma") - elif ".residual.2.bias" in new_key: - new_key = new_key.replace(".residual.2.bias", ".conv1.bias") - elif ".residual.2.weight" in new_key: - new_key = new_key.replace(".residual.2.weight", ".conv1.weight") - elif ".residual.3.gamma" in new_key: - new_key = new_key.replace(".residual.3.gamma", ".norm2.gamma") - elif ".residual.6.bias" in new_key: - new_key = new_key.replace(".residual.6.bias", ".conv2.bias") - elif ".residual.6.weight" in new_key: - new_key = new_key.replace(".residual.6.weight", ".conv2.weight") - elif ".shortcut.bias" in new_key: - new_key = new_key.replace(".shortcut.bias", ".conv_shortcut.bias") - elif ".shortcut.weight" in new_key: - new_key = new_key.replace(".shortcut.weight", ".conv_shortcut.weight") - - converted_state_dict[new_key] = value - - # Handle decoder upsamples - elif key.startswith("decoder.upsamples."): - # Convert to up_blocks - parts = key.split(".") - block_idx = int(parts[2]) - - # Group residual blocks - if "residual" in key: - if block_idx in [0, 1, 2]: - new_block_idx = 0 - resnet_idx = block_idx - elif block_idx in [4, 5, 6]: - new_block_idx = 1 - resnet_idx = block_idx - 4 - elif block_idx in [8, 9, 10]: - new_block_idx = 2 - resnet_idx = block_idx - 8 - elif block_idx in [12, 13, 14]: - new_block_idx = 3 - resnet_idx = block_idx - 12 - else: - # Keep as is for other blocks - converted_state_dict[key] = value - continue - - # Convert residual block naming - if ".residual.0.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm1.gamma" - elif ".residual.2.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.bias" - elif ".residual.2.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv1.weight" - elif ".residual.3.gamma" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.norm2.gamma" - elif ".residual.6.bias" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.bias" - elif ".residual.6.weight" in key: - new_key = f"decoder.up_blocks.{new_block_idx}.resnets.{resnet_idx}.conv2.weight" - else: - new_key = key - - converted_state_dict[new_key] = value - - # Handle shortcut connections - elif ".shortcut." in key: - if block_idx == 4: - new_key = key.replace(".shortcut.", ".resnets.0.conv_shortcut.") - new_key = new_key.replace("decoder.upsamples.4", "decoder.up_blocks.1") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - new_key = new_key.replace(".shortcut.", ".conv_shortcut.") - - converted_state_dict[new_key] = value - - # Handle upsamplers - elif ".resample." in key or ".time_conv." in key: - if block_idx == 3: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.0.upsamplers.0") - elif block_idx == 7: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.1.upsamplers.0") - elif block_idx == 11: - new_key = key.replace(f"decoder.upsamples.{block_idx}", "decoder.up_blocks.2.upsamplers.0") - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - - converted_state_dict[new_key] = value - else: - new_key = key.replace("decoder.upsamples.", "decoder.up_blocks.") - converted_state_dict[new_key] = value - else: - # Keep other keys unchanged - converted_state_dict[key] = value - - return converted_state_dict - - -def convert_hidream_transformer_to_diffusers(checkpoint, **kwargs): - keys = list(checkpoint.keys()) - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - return checkpoint - - -def convert_chroma_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {} - keys = list(checkpoint.keys()) - - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - num_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "double_blocks." in k))[-1] + 1 # noqa: C401 - num_single_layers = list(set(int(k.split(".", 2)[1]) for k in checkpoint if "single_blocks." in k))[-1] + 1 # noqa: C401 - num_guidance_layers = ( - list(set(int(k.split(".", 3)[2]) for k in checkpoint if "distilled_guidance_layer.layers." in k))[-1] + 1 # noqa: C401 - ) - mlp_ratio = 4.0 - inner_dim = 3072 - - # in SD3 original implementation of AdaLayerNormContinuous, it split linear projection output into shift, scale; - # while in diffusers it split into scale, shift. Here we swap the linear projection weights in order to be able to use diffusers implementation - def swap_scale_shift(weight): - shift, scale = weight.chunk(2, dim=0) - new_weight = torch.cat([scale, shift], dim=0) - return new_weight - - # guidance - converted_state_dict["distilled_guidance_layer.in_proj.bias"] = checkpoint.pop( - "distilled_guidance_layer.in_proj.bias" - ) - converted_state_dict["distilled_guidance_layer.in_proj.weight"] = checkpoint.pop( - "distilled_guidance_layer.in_proj.weight" - ) - converted_state_dict["distilled_guidance_layer.out_proj.bias"] = checkpoint.pop( - "distilled_guidance_layer.out_proj.bias" - ) - converted_state_dict["distilled_guidance_layer.out_proj.weight"] = checkpoint.pop( - "distilled_guidance_layer.out_proj.weight" - ) - for i in range(num_guidance_layers): - block_prefix = f"distilled_guidance_layer.layers.{i}." - converted_state_dict[f"{block_prefix}linear_1.bias"] = checkpoint.pop( - f"distilled_guidance_layer.layers.{i}.in_layer.bias" - ) - converted_state_dict[f"{block_prefix}linear_1.weight"] = checkpoint.pop( - f"distilled_guidance_layer.layers.{i}.in_layer.weight" - ) - converted_state_dict[f"{block_prefix}linear_2.bias"] = checkpoint.pop( - f"distilled_guidance_layer.layers.{i}.out_layer.bias" - ) - converted_state_dict[f"{block_prefix}linear_2.weight"] = checkpoint.pop( - f"distilled_guidance_layer.layers.{i}.out_layer.weight" - ) - converted_state_dict[f"distilled_guidance_layer.norms.{i}.weight"] = checkpoint.pop( - f"distilled_guidance_layer.norms.{i}.scale" - ) - - # context_embedder - converted_state_dict["context_embedder.weight"] = checkpoint.pop("txt_in.weight") - converted_state_dict["context_embedder.bias"] = checkpoint.pop("txt_in.bias") - - # x_embedder - converted_state_dict["x_embedder.weight"] = checkpoint.pop("img_in.weight") - converted_state_dict["x_embedder.bias"] = checkpoint.pop("img_in.bias") - - # double transformer blocks - for i in range(num_layers): - block_prefix = f"transformer_blocks.{i}." - # Q, K, V - sample_q, sample_k, sample_v = torch.chunk(checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.weight"), 3, dim=0) - context_q, context_k, context_v = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.weight"), 3, dim=0 - ) - sample_q_bias, sample_k_bias, sample_v_bias = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.img_attn.qkv.bias"), 3, dim=0 - ) - context_q_bias, context_k_bias, context_v_bias = torch.chunk( - checkpoint.pop(f"double_blocks.{i}.txt_attn.qkv.bias"), 3, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([sample_q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([sample_q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([sample_k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([sample_k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([sample_v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([sample_v_bias]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.weight"] = torch.cat([context_q]) - converted_state_dict[f"{block_prefix}attn.add_q_proj.bias"] = torch.cat([context_q_bias]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.weight"] = torch.cat([context_k]) - converted_state_dict[f"{block_prefix}attn.add_k_proj.bias"] = torch.cat([context_k_bias]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.weight"] = torch.cat([context_v]) - converted_state_dict[f"{block_prefix}attn.add_v_proj.bias"] = torch.cat([context_v_bias]) - # qk_norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.norm.key_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_q.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_added_k.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.norm.key_norm.scale" - ) - # ff img_mlp - converted_state_dict[f"{block_prefix}ff.net.0.proj.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff.net.0.proj.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.0.bias") - converted_state_dict[f"{block_prefix}ff.net.2.weight"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.weight") - converted_state_dict[f"{block_prefix}ff.net.2.bias"] = checkpoint.pop(f"double_blocks.{i}.img_mlp.2.bias") - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.0.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.0.proj.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.0.bias" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.2.weight" - ) - converted_state_dict[f"{block_prefix}ff_context.net.2.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_mlp.2.bias" - ) - # output projections. - converted_state_dict[f"{block_prefix}attn.to_out.0.weight"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_out.0.bias"] = checkpoint.pop( - f"double_blocks.{i}.img_attn.proj.bias" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.weight"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.proj.weight" - ) - converted_state_dict[f"{block_prefix}attn.to_add_out.bias"] = checkpoint.pop( - f"double_blocks.{i}.txt_attn.proj.bias" - ) - - # single transformer blocks - for i in range(num_single_layers): - block_prefix = f"single_transformer_blocks.{i}." - # Q, K, V, mlp - mlp_hidden_dim = int(inner_dim * mlp_ratio) - split_size = (inner_dim, inner_dim, inner_dim, mlp_hidden_dim) - q, k, v, mlp = torch.split(checkpoint.pop(f"single_blocks.{i}.linear1.weight"), split_size, dim=0) - q_bias, k_bias, v_bias, mlp_bias = torch.split( - checkpoint.pop(f"single_blocks.{i}.linear1.bias"), split_size, dim=0 - ) - converted_state_dict[f"{block_prefix}attn.to_q.weight"] = torch.cat([q]) - converted_state_dict[f"{block_prefix}attn.to_q.bias"] = torch.cat([q_bias]) - converted_state_dict[f"{block_prefix}attn.to_k.weight"] = torch.cat([k]) - converted_state_dict[f"{block_prefix}attn.to_k.bias"] = torch.cat([k_bias]) - converted_state_dict[f"{block_prefix}attn.to_v.weight"] = torch.cat([v]) - converted_state_dict[f"{block_prefix}attn.to_v.bias"] = torch.cat([v_bias]) - converted_state_dict[f"{block_prefix}proj_mlp.weight"] = torch.cat([mlp]) - converted_state_dict[f"{block_prefix}proj_mlp.bias"] = torch.cat([mlp_bias]) - # qk norm - converted_state_dict[f"{block_prefix}attn.norm_q.weight"] = checkpoint.pop( - f"single_blocks.{i}.norm.query_norm.scale" - ) - converted_state_dict[f"{block_prefix}attn.norm_k.weight"] = checkpoint.pop( - f"single_blocks.{i}.norm.key_norm.scale" - ) - # output projections. - converted_state_dict[f"{block_prefix}proj_out.weight"] = checkpoint.pop(f"single_blocks.{i}.linear2.weight") - converted_state_dict[f"{block_prefix}proj_out.bias"] = checkpoint.pop(f"single_blocks.{i}.linear2.bias") - - converted_state_dict["proj_out.weight"] = checkpoint.pop("final_layer.linear.weight") - converted_state_dict["proj_out.bias"] = checkpoint.pop("final_layer.linear.bias") - - return converted_state_dict - - -def convert_cosmos_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - def remove_keys_(key: str, state_dict): - state_dict.pop(key) - - def rename_transformer_blocks_(key: str, state_dict): - block_index = int(key.split(".")[1].removeprefix("block")) - new_key = key - old_prefix = f"blocks.block{block_index}" - new_prefix = f"transformer_blocks.{block_index}" - new_key = new_prefix + new_key.removeprefix(old_prefix) - state_dict[new_key] = state_dict.pop(key) - - TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 = { - "t_embedder.1": "time_embed.t_embedder", - "affline_norm": "time_embed.norm", - ".blocks.0.block.attn": ".attn1", - ".blocks.1.block.attn": ".attn2", - ".blocks.2.block": ".ff", - ".blocks.0.adaLN_modulation.1": ".norm1.linear_1", - ".blocks.0.adaLN_modulation.2": ".norm1.linear_2", - ".blocks.1.adaLN_modulation.1": ".norm2.linear_1", - ".blocks.1.adaLN_modulation.2": ".norm2.linear_2", - ".blocks.2.adaLN_modulation.1": ".norm3.linear_1", - ".blocks.2.adaLN_modulation.2": ".norm3.linear_2", - "to_q.0": "to_q", - "to_q.1": "norm_q", - "to_k.0": "to_k", - "to_k.1": "norm_k", - "to_v.0": "to_v", - "layer1": "net.0.proj", - "layer2": "net.2", - "proj.1": "proj", - "x_embedder": "patch_embed", - "extra_pos_embedder": "learnable_pos_embed", - "final_layer.adaLN_modulation.1": "norm_out.linear_1", - "final_layer.adaLN_modulation.2": "norm_out.linear_2", - "final_layer.linear": "proj_out", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 = { - "blocks.block": rename_transformer_blocks_, - "logvar.0.freqs": remove_keys_, - "logvar.0.phases": remove_keys_, - "logvar.1.weight": remove_keys_, - "pos_embedder.seq": remove_keys_, - } - - TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 = { - "t_embedder.1": "time_embed.t_embedder", - "t_embedding_norm": "time_embed.norm", - "blocks": "transformer_blocks", - "adaln_modulation_self_attn.1": "norm1.linear_1", - "adaln_modulation_self_attn.2": "norm1.linear_2", - "adaln_modulation_cross_attn.1": "norm2.linear_1", - "adaln_modulation_cross_attn.2": "norm2.linear_2", - "adaln_modulation_mlp.1": "norm3.linear_1", - "adaln_modulation_mlp.2": "norm3.linear_2", - "self_attn": "attn1", - "cross_attn": "attn2", - "q_proj": "to_q", - "k_proj": "to_k", - "v_proj": "to_v", - "output_proj": "to_out.0", - "q_norm": "norm_q", - "k_norm": "norm_k", - "mlp.layer1": "ff.net.0.proj", - "mlp.layer2": "ff.net.2", - "x_embedder.proj.1": "patch_embed.proj", - "final_layer.adaln_modulation.1": "norm_out.linear_1", - "final_layer.adaln_modulation.2": "norm_out.linear_2", - "final_layer.linear": "proj_out", - } - - TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 = { - "accum_video_sample_counter": remove_keys_, - "accum_image_sample_counter": remove_keys_, - "accum_iteration": remove_keys_, - "accum_train_in_hours": remove_keys_, - "pos_embedder.seq": remove_keys_, - "pos_embedder.dim_spatial_range": remove_keys_, - "pos_embedder.dim_temporal_range": remove_keys_, - "_extra_state": remove_keys_, - } - - PREFIX_KEY = "net." - if "net.blocks.block1.blocks.0.block.attn.to_q.0.weight" in checkpoint: - TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_1_0 - TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_1_0 - else: - TRANSFORMER_KEYS_RENAME_DICT = TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0 - TRANSFORMER_SPECIAL_KEYS_REMAP = TRANSFORMER_SPECIAL_KEYS_REMAP_COSMOS_2_0 - - state_dict_keys = list(converted_state_dict.keys()) - for key in state_dict_keys: - new_key = key[:] - if new_key.startswith(PREFIX_KEY): - new_key = new_key.removeprefix(PREFIX_KEY) - for replace_key, rename_key in TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - converted_state_dict[new_key] = converted_state_dict.pop(key) - - state_dict_keys = list(converted_state_dict.keys()) - for key in state_dict_keys: - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_flux2_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - FLUX2_TRANSFORMER_KEYS_RENAME_DICT = { - # Image and text input projections - "img_in": "x_embedder", - "txt_in": "context_embedder", - # Timestep and guidance embeddings - "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", - "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", - "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", - "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", - # Modulation parameters - "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", - "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", - "single_stream_modulation.lin": "single_stream_modulation.linear", - # Final output layer - # "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params - "final_layer.linear": "proj_out", - } - - FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear", - } - - FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = { - # Handle fused QKV projections separately as we need to break into Q, K, V projections - "img_attn.norm.query_norm": "attn.norm_q", - "img_attn.norm.key_norm": "attn.norm_k", - "img_attn.proj": "attn.to_out.0", - "img_mlp.0": "ff.linear_in", - "img_mlp.2": "ff.linear_out", - "txt_attn.norm.query_norm": "attn.norm_added_q", - "txt_attn.norm.key_norm": "attn.norm_added_k", - "txt_attn.proj": "attn.to_add_out", - "txt_mlp.0": "ff_context.linear_in", - "txt_mlp.2": "ff_context.linear_out", - } - - FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = { - "linear1": "attn.to_qkv_mlp_proj", - "linear2": "attn.to_out", - "norm.query_norm": "attn.norm_q", - "norm.key_norm": "attn.norm_k", - } - - def convert_flux2_single_stream_blocks(key: str, state_dict: dict[str, object]) -> None: - # Skip if not a weight, bias, or scale - if ".weight" not in key and ".bias" not in key and ".scale" not in key: - return - - # Mapping: - # - single_blocks.{N}.linear1 --> single_transformer_blocks.{N}.attn.to_qkv_mlp_proj - # - single_blocks.{N}.linear2 --> single_transformer_blocks.{N}.attn.to_out - # - single_blocks.{N}.norm.query_norm.scale --> single_transformer_blocks.{N}.attn.norm_q.weight - # - single_blocks.{N}.norm.key_norm.scale --> single_transformer_blocks.{N}.attn.norm_k.weight - new_prefix = "single_transformer_blocks" - if "single_blocks." in key: - parts = key.split(".") - block_idx = parts[1] - within_block_name = ".".join(parts[2:-1]) - param_type = parts[-1] - - if param_type == "scale": - param_type = "weight" - - new_within_block_name = FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP[within_block_name] - new_key = ".".join([new_prefix, block_idx, new_within_block_name, param_type]) - - param = state_dict.pop(key) - state_dict[new_key] = param - - return - - def convert_ada_layer_norm_weights(key: str, state_dict: dict[str, object]) -> None: - # Skip if not a weight - if ".weight" not in key: - return - - # If adaLN_modulation is in the key, swap scale and shift parameters - # Original implementation is (shift, scale); diffusers implementation is (scale, shift) - if "adaLN_modulation" in key: - key_without_param_type, param_type = key.rsplit(".", maxsplit=1) - # Assume all such keys are in the AdaLayerNorm key map - new_key_without_param_type = FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP[key_without_param_type] - new_key = ".".join([new_key_without_param_type, param_type]) - - swapped_weight = swap_scale_shift(state_dict.pop(key), 0) - state_dict[new_key] = swapped_weight - - return - - def convert_flux2_double_stream_blocks(key: str, state_dict: dict[str, object]) -> None: - # Skip if not a weight, bias, or scale - if ".weight" not in key and ".bias" not in key and ".scale" not in key: - return - - new_prefix = "transformer_blocks" - if "double_blocks." in key: - parts = key.split(".") - block_idx = parts[1] - modality_block_name = parts[2] # img_attn, img_mlp, txt_attn, txt_mlp - within_block_name = ".".join(parts[2:-1]) - param_type = parts[-1] - - if param_type == "scale": - param_type = "weight" - - if "qkv" in within_block_name: - fused_qkv_weight = state_dict.pop(key) - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - if "img" in modality_block_name: - # double_blocks.{N}.img_attn.qkv --> transformer_blocks.{N}.attn.{to_q|to_k|to_v} - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - new_q_name = "attn.to_q" - new_k_name = "attn.to_k" - new_v_name = "attn.to_v" - elif "txt" in modality_block_name: - # double_blocks.{N}.txt_attn.qkv --> transformer_blocks.{N}.attn.{add_q_proj|add_k_proj|add_v_proj} - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - new_q_name = "attn.add_q_proj" - new_k_name = "attn.add_k_proj" - new_v_name = "attn.add_v_proj" - new_q_key = ".".join([new_prefix, block_idx, new_q_name, param_type]) - new_k_key = ".".join([new_prefix, block_idx, new_k_name, param_type]) - new_v_key = ".".join([new_prefix, block_idx, new_v_name, param_type]) - state_dict[new_q_key] = to_q_weight - state_dict[new_k_key] = to_k_weight - state_dict[new_v_key] = to_v_weight - else: - new_within_block_name = FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP[within_block_name] - new_key = ".".join([new_prefix, block_idx, new_within_block_name, param_type]) - - param = state_dict.pop(key) - state_dict[new_key] = param - return - - def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - "adaLN_modulation": convert_ada_layer_norm_weights, - "double_blocks": convert_flux2_double_stream_blocks, - "single_blocks": convert_flux2_single_stream_blocks, - } - - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in FLUX2_TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - update_state_dict(converted_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_z_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - Z_IMAGE_KEYS_RENAME_DICT = { - "final_layer.": "all_final_layer.2-1.", - "x_embedder.": "all_x_embedder.2-1.", - ".attention.out.bias": ".attention.to_out.0.bias", - ".attention.k_norm.weight": ".attention.norm_k.weight", - ".attention.q_norm.weight": ".attention.norm_q.weight", - ".attention.out.weight": ".attention.to_out.0.weight", - "model.diffusion_model.": "", - } - - def convert_z_image_fused_attention(key: str, state_dict: dict[str, object]) -> None: - if ".attention.qkv.weight" not in key: - return - - fused_qkv_weight = state_dict.pop(key) - to_q_weight, to_k_weight, to_v_weight = torch.chunk(fused_qkv_weight, 3, dim=0) - new_q_name = key.replace(".attention.qkv.weight", ".attention.to_q.weight") - new_k_name = key.replace(".attention.qkv.weight", ".attention.to_k.weight") - new_v_name = key.replace(".attention.qkv.weight", ".attention.to_v.weight") - - state_dict[new_q_name] = to_q_weight - state_dict[new_k_name] = to_k_weight - state_dict[new_v_name] = to_v_weight - return - - TRANSFORMER_SPECIAL_KEYS_REMAP = { - ".attention.qkv.weight": convert_z_image_fused_attention, - } - - def update_state_dict(state_dict: dict[str, object], old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - # Handle single file --> diffusers key remapping via the remap dict - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in Z_IMAGE_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - update_state_dict(converted_state_dict, key, new_key) - - if "norm_final.weight" in converted_state_dict.keys(): - _ = converted_state_dict.pop("norm_final.weight") - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_z_image_controlnet_checkpoint_to_diffusers(checkpoint, config, **kwargs): - if config["add_control_noise_refiner"] is None: - return checkpoint - elif config["add_control_noise_refiner"] == "control_noise_refiner": - return checkpoint - elif config["add_control_noise_refiner"] == "control_layers": - converted_state_dict = { - key: checkpoint.pop(key) for key in list(checkpoint.keys()) if not key.startswith("control_noise_refiner.") - } - return converted_state_dict - else: - raise ValueError("Unknown Z-Image Turbo ControlNet type.") - - -def convert_ltx2_transformer_to_diffusers(checkpoint, **kwargs): - LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT = { - # Transformer prefix - "model.diffusion_model.": "", - # Input Patchify Projections - "patchify_proj": "proj_in", - "audio_patchify_proj": "audio_proj_in", - # Modulation Parameters - # Handle adaln_single --> time_embed, audioln_single --> audio_time_embed separately as the original keys are - # substrings of the other modulation parameters below - "av_ca_video_scale_shift_adaln_single": "av_cross_attn_video_scale_shift", - "av_ca_a2v_gate_adaln_single": "av_cross_attn_video_a2v_gate", - "av_ca_audio_scale_shift_adaln_single": "av_cross_attn_audio_scale_shift", - "av_ca_v2a_gate_adaln_single": "av_cross_attn_audio_v2a_gate", - # Transformer Blocks - # Per-Block Cross Attention Modulation Parameters - "scale_shift_table_a2v_ca_video": "video_a2v_cross_attn_scale_shift_table", - "scale_shift_table_a2v_ca_audio": "audio_a2v_cross_attn_scale_shift_table", - # Attention QK Norms - "q_norm": "norm_q", - "k_norm": "norm_k", - } - - def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - def remove_keys_inplace(key: str, state_dict) -> None: - state_dict.pop(key) - - def convert_ltx2_transformer_adaln_single(key: str, state_dict) -> None: - # Skip if not a weight, bias - if ".weight" not in key and ".bias" not in key: - return - - if key.startswith("adaln_single."): - new_key = key.replace("adaln_single.", "time_embed.") - param = state_dict.pop(key) - state_dict[new_key] = param - - if key.startswith("audio_adaln_single."): - new_key = key.replace("audio_adaln_single.", "audio_time_embed.") - param = state_dict.pop(key) - state_dict[new_key] = param - - return - - LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP = { - "video_embeddings_connector": remove_keys_inplace, - "audio_embeddings_connector": remove_keys_inplace, - "adaln_single": convert_ltx2_transformer_adaln_single, - } - - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in LTX_2_0_TRANSFORMER_KEYS_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - update_state_dict_inplace(converted_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in LTX_2_0_TRANSFORMER_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_ltx2_vae_to_diffusers(checkpoint, **kwargs): - LTX_2_0_VIDEO_VAE_RENAME_DICT = { - # Video VAE prefix - "vae.": "", - # Encoder - "down_blocks.0": "down_blocks.0", - "down_blocks.1": "down_blocks.0.downsamplers.0", - "down_blocks.2": "down_blocks.1", - "down_blocks.3": "down_blocks.1.downsamplers.0", - "down_blocks.4": "down_blocks.2", - "down_blocks.5": "down_blocks.2.downsamplers.0", - "down_blocks.6": "down_blocks.3", - "down_blocks.7": "down_blocks.3.downsamplers.0", - "down_blocks.8": "mid_block", - # Decoder - "up_blocks.0": "mid_block", - "up_blocks.1": "up_blocks.0.upsamplers.0", - "up_blocks.2": "up_blocks.0", - "up_blocks.3": "up_blocks.1.upsamplers.0", - "up_blocks.4": "up_blocks.1", - "up_blocks.5": "up_blocks.2.upsamplers.0", - "up_blocks.6": "up_blocks.2", - # Common - # For all 3D ResNets - "res_blocks": "resnets", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", - } - - def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - def remove_keys_inplace(key: str, state_dict) -> None: - state_dict.pop(key) - - LTX_2_0_VAE_SPECIAL_KEYS_REMAP = { - "per_channel_statistics.channel": remove_keys_inplace, - "per_channel_statistics.mean-of-stds": remove_keys_inplace, - } - - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in LTX_2_0_VIDEO_VAE_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - update_state_dict_inplace(converted_state_dict, key, new_key) - - # Handle any special logic which can't be expressed by a simple 1:1 remapping with the handlers in - # special_keys_remap - for key in list(converted_state_dict.keys()): - for special_key, handler_fn_inplace in LTX_2_0_VAE_SPECIAL_KEYS_REMAP.items(): - if special_key not in key: - continue - handler_fn_inplace(key, converted_state_dict) - - return converted_state_dict - - -def convert_ltx2_audio_vae_to_diffusers(checkpoint, **kwargs): - LTX_2_0_AUDIO_VAE_RENAME_DICT = { - # Audio VAE prefix - "audio_vae.": "", - "per_channel_statistics.mean-of-means": "latents_mean", - "per_channel_statistics.std-of-means": "latents_std", - } - - def update_state_dict_inplace(state_dict, old_key: str, new_key: str) -> None: - state_dict[new_key] = state_dict.pop(old_key) - - converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())} - - # Handle official code --> diffusers key remapping via the remap dict - for key in list(converted_state_dict.keys()): - new_key = key[:] - for replace_key, rename_key in LTX_2_0_AUDIO_VAE_RENAME_DICT.items(): - new_key = new_key.replace(replace_key, rename_key) - - update_state_dict_inplace(converted_state_dict, key, new_key) - - return converted_state_dict - - -def convert_ernie_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs): - keys = list(checkpoint.keys()) - - for k in keys: - if "model.diffusion_model." in k: - checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k) - - return checkpoint diff --git a/src/diffusers/models/adapter.py b/src/diffusers/models/adapter.py index 2072749c65ae..29ec8a19db91 100644 --- a/src/diffusers/models/adapter.py +++ b/src/diffusers/models/adapter.py @@ -18,6 +18,7 @@ import torch.nn as nn from ..configuration_utils import ConfigMixin, register_to_config +from ..loaders import FromOriginalModelMixin from ..utils import logging from .modeling_utils import ModelMixin @@ -217,7 +218,7 @@ def from_pretrained(cls, pretrained_model_path: str | os.PathLike | None, **kwar return cls(adapters) -class T2IAdapter(ModelMixin, ConfigMixin): +class T2IAdapter(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" A simple ResNet-like model that accepts images containing control signals such as keyposes and depth. The model generates multiple feature maps that are used as additional conditioning in [`UNet2DConditionModel`]. The model's diff --git a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py index bf13a4b3b134..d34040ed4a03 100644 --- a/src/diffusers/models/autoencoders/autoencoder_asym_kl.py +++ b/src/diffusers/models/autoencoders/autoencoder_asym_kl.py @@ -15,13 +15,14 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils.accelerate_utils import apply_forward_hook from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin from .vae import AutoencoderMixin, DecoderOutput, DiagonalGaussianDistribution, Encoder, MaskConditionDecoder -class AsymmetricAutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin): +class AsymmetricAutoencoderKL(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" Designing a Better Asymmetric VQGAN for StableDiffusion https://huggingface.co/papers/2306.04632 . A VAE model with KL loss for encoding images into latents and decoding latent representations into images. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py index 362df0bd96a2..2ca2026988d2 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import get_logger from ...utils.accelerate_utils import apply_forward_hook from ..modeling_outputs import AutoencoderKLOutput @@ -876,7 +877,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLCosmos(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLCosmos(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" Autoencoder used in [Cosmos](https://huggingface.co/papers/2501.03575). diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py index fece756ebec6..79765caf902b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuan_video.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation @@ -622,7 +623,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLHunyuanVideo(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLHunyuanVideo(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Introduced in [HunyuanVideo](https://huggingface.co/papers/2412.03603). diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py index 5297e3c850ba..0ad1c9801a3f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanimage_refiner.py @@ -20,6 +20,7 @@ import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation @@ -590,7 +591,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLHunyuanImageRefiner(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLHunyuanImageRefiner(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for HunyuanImage-2.1 Refiner. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py index dec20aacb7d5..497d88c1a72b 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_hunyuanvideo15.py @@ -20,6 +20,7 @@ import torch.utils.checkpoint from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation @@ -631,7 +632,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states -class AutoencoderKLHunyuanVideo15(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLHunyuanVideo15(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. Used for HunyuanVideo-1.5. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py index fb773dbdc01e..cb0d45ced348 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_ltx2_audio.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils.accelerate_utils import apply_forward_hook from ..modeling_outputs import AutoencoderKLOutput from ..modeling_utils import ModelMixin @@ -665,7 +666,7 @@ def forward( return decoded_output -class AutoencoderKLLTX2Audio(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLLTX2Audio(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" LTX2 audio VAE for encoding and decoding audio latent representations. """ diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py index 23ad2c725c00..af35237686e7 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.accelerate_utils import apply_forward_hook from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -498,7 +499,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) -class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, AutoencoderMixin): +class AutoencoderKLMiniMaxH3(ModelMixin, ConfigMixin, AttentionMixin, AutoencoderMixin, FromOriginalModelMixin): r""" A VAE model with a causal 3D CNN encoder and a non-causal ViT decoder, used in [MiniMax-H3](https://huggingface.co/MiniMaxAI). diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py index a70c947712ae..c6caf71b43e6 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_minimax_h3_audio.py @@ -43,6 +43,7 @@ from torch.nn.utils import weight_norm from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ...utils.torch_utils import randn_tensor @@ -486,7 +487,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return torch.clamp(hidden_states, min=-1.0, max=1.0) -class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin): +class AutoencoderKLMiniMaxH3Audio(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): r""" The audio autoencoder used by [MiniMax-H3](https://huggingface.co/MiniMaxAI): a DAC-lineage convolutional encoder and a BigVGAN decoder, operating directly on mono 32 kHz waveforms. diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py index bb447015c54d..000a9dbd41fb 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py @@ -19,6 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.accelerate_utils import apply_forward_hook from ..activations import get_activation @@ -652,7 +653,7 @@ def forward(self, hidden_states: torch.Tensor, conv_cache: dict[str, torch.Tenso return hidden_states, new_conv_cache -class AutoencoderKLMochi(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLMochi(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. Used in [Mochi 1 preview](https://github.com/genmoai/models). diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py index 8b0e5806d8ef..3aacf049f0ea 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py +++ b/src/diffusers/models/autoencoders/autoencoder_kl_temporal_decoder.py @@ -17,6 +17,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils.accelerate_utils import apply_forward_hook from ..attention import AttentionMixin from ..attention_processor import CROSS_ATTENTION_PROCESSORS, AttnProcessor @@ -135,7 +136,7 @@ def forward( return sample -class AutoencoderKLTemporalDecoder(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderKLTemporalDecoder(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VAE model with KL loss for encoding images into latents and decoding latent representations into images. diff --git a/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py b/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py index 3b5e81d814c0..cc3455608a17 100644 --- a/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py +++ b/src/diffusers/models/autoencoders/autoencoder_longcat_audio_dit.py @@ -24,6 +24,7 @@ from torch.nn.utils import weight_norm from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ...utils.torch_utils import randn_tensor @@ -293,7 +294,7 @@ class LongCatAudioDiTVaeDecoderOutput(BaseOutput): sample: torch.Tensor -class LongCatAudioDiTVae(ModelMixin, AutoencoderMixin, ConfigMixin): +class LongCatAudioDiTVae(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): _supports_group_offloading = False @register_to_config diff --git a/src/diffusers/models/autoencoders/autoencoder_oobleck.py b/src/diffusers/models/autoencoders/autoencoder_oobleck.py index d4251fd9f1a9..d126856ffe97 100644 --- a/src/diffusers/models/autoencoders/autoencoder_oobleck.py +++ b/src/diffusers/models/autoencoders/autoencoder_oobleck.py @@ -20,6 +20,7 @@ from torch.nn.utils import weight_norm from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ...utils.torch_utils import randn_tensor @@ -291,7 +292,7 @@ def forward(self, hidden_state): return hidden_state -class AutoencoderOobleck(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderOobleck(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" An autoencoder for encoding waveforms into latents and decoding latent representations into waveforms. First introduced in Stable Audio. diff --git a/src/diffusers/models/autoencoders/autoencoder_rae.py b/src/diffusers/models/autoencoders/autoencoder_rae.py index 35a96e6f67bc..e5de1fb86f00 100644 --- a/src/diffusers/models/autoencoders/autoencoder_rae.py +++ b/src/diffusers/models/autoencoders/autoencoder_rae.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput, logging from ...utils.accelerate_utils import apply_forward_hook from ...utils.import_utils import is_transformers_available @@ -110,8 +111,10 @@ def _build_encoder( ) model = SiglipVisionModel(config) # See dinov2 comment above. - model.vision_model.post_layernorm.weight = None - model.vision_model.post_layernorm.bias = None + # Transformers 5 exposes the vision backbone directly on SiglipVisionModel. + vision_model = getattr(model, "vision_model", model) + vision_model.post_layernorm.weight = None + vision_model.post_layernorm.bias = None elif encoder_type == "mae": config = ViTMAEConfig( hidden_size=hidden_size, @@ -390,7 +393,7 @@ def forward( return RAEDecoderOutput(logits=logits) -class AutoencoderRAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderRAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" Representation Autoencoder (RAE) model for encoding images to latents and decoding latents to images. diff --git a/src/diffusers/models/autoencoders/autoencoder_same.py b/src/diffusers/models/autoencoders/autoencoder_same.py index 23dccf7732d9..9e971178c925 100644 --- a/src/diffusers/models/autoencoders/autoencoder_same.py +++ b/src/diffusers/models/autoencoders/autoencoder_same.py @@ -33,6 +33,7 @@ from torch.nn.utils import weight_norm from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput, logging from ...utils.accelerate_utils import apply_forward_hook from ..attention import AttentionMixin, AttentionModuleMixin @@ -605,7 +606,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # ────────────────────────────────────────────────────────────────────────────── -class AutoencoderSAME(ModelMixin, ConfigMixin, AttentionMixin): +class AutoencoderSAME(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): r""" Semantically-Aligned Music Encoder (SAME) autoencoder from *Stable Audio 3* (`arXiv 2605.17991 `_). diff --git a/src/diffusers/models/autoencoders/autoencoder_tiny.py b/src/diffusers/models/autoencoders/autoencoder_tiny.py index 5647203e02e1..1aa27a40e31f 100644 --- a/src/diffusers/models/autoencoders/autoencoder_tiny.py +++ b/src/diffusers/models/autoencoders/autoencoder_tiny.py @@ -18,6 +18,7 @@ import torch from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ..modeling_utils import ModelMixin @@ -37,7 +38,7 @@ class AutoencoderTinyOutput(BaseOutput): latents: torch.Tensor -class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin): +class AutoencoderTiny(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A tiny distilled VAE model for encoding images into latents and decoding latent representations into images. diff --git a/src/diffusers/models/autoencoders/consistency_decoder_vae.py b/src/diffusers/models/autoencoders/consistency_decoder_vae.py index dbe0f4c30541..fd3bb991206d 100644 --- a/src/diffusers/models/autoencoders/consistency_decoder_vae.py +++ b/src/diffusers/models/autoencoders/consistency_decoder_vae.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...schedulers import ConsistencyDecoderScheduler from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook @@ -48,7 +49,7 @@ class ConsistencyDecoderVAEOutput(BaseOutput): latent_dist: "DiagonalGaussianDistribution" -class ConsistencyDecoderVAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin): +class ConsistencyDecoderVAE(ModelMixin, AttentionMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" The consistency decoder used with DALL-E 3. diff --git a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py index 41388991e0b4..6e3495894b0d 100644 --- a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py +++ b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import is_kernels_available, logging from ...utils.accelerate_utils import apply_forward_hook from ...utils.constants import DIFFUSERS_DISABLE_REMOTE_CODE @@ -697,7 +698,7 @@ def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)] -class LTX2VideoDiffusionDecoderModel(ModelMixin, AttentionMixin, ConfigMixin): +class LTX2VideoDiffusionDecoderModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): r""" The LTX-2 diffusion video decoder, introduced in LTX-2.5. diff --git a/src/diffusers/models/autoencoders/minimax_music3_vocoder.py b/src/diffusers/models/autoencoders/minimax_music3_vocoder.py index bf325634e451..6aff731fc24d 100644 --- a/src/diffusers/models/autoencoders/minimax_music3_vocoder.py +++ b/src/diffusers/models/autoencoders/minimax_music3_vocoder.py @@ -19,6 +19,7 @@ from torch.nn.utils import weight_norm from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ..modeling_utils import ModelMixin @@ -68,7 +69,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return self.res_unit3(hidden_states) -class MiniMaxMusic3Vocoder(ModelMixin, ConfigMixin): +class MiniMaxMusic3Vocoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" The Flow-VAE waveform decoder of MiniMax Music 3 (a DAC-style decoder). It decodes flow-matched latents of shape `(batch, latent_channels, length)` into stereo waveforms at `sampling_rate`; the two audio channels are decoded as diff --git a/src/diffusers/models/autoencoders/vq_model.py b/src/diffusers/models/autoencoders/vq_model.py index 619327dde417..fb6a18704e01 100644 --- a/src/diffusers/models/autoencoders/vq_model.py +++ b/src/diffusers/models/autoencoders/vq_model.py @@ -17,6 +17,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.accelerate_utils import apply_forward_hook from ..autoencoders.vae import Decoder, DecoderOutput, Encoder, VectorQuantizer @@ -37,7 +38,7 @@ class VQEncoderOutput(BaseOutput): latents: torch.Tensor -class VQModel(ModelMixin, AutoencoderMixin, ConfigMixin): +class VQModel(ModelMixin, AutoencoderMixin, ConfigMixin, FromOriginalModelMixin): r""" A VQ-VAE model for decoding latent representations. diff --git a/src/diffusers/models/condition_embedders/condition_embedder_anima.py b/src/diffusers/models/condition_embedders/condition_embedder_anima.py index 40fda447ec68..1d1ae4c2b04d 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_anima.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_anima.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ..attention import AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..modeling_utils import ModelMixin @@ -226,7 +226,7 @@ def forward( return hidden_states -class AnimaTextConditioner(ModelMixin, ConfigMixin, PeftAdapterMixin): +class AnimaTextConditioner(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): r""" Text conditioner used by Anima to map Qwen3 hidden states and T5 token ids to Cosmos text embeddings. diff --git a/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py b/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py index 2ee3b51fa0e7..62343b175f52 100644 --- a/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py +++ b/src/diffusers/models/condition_embedders/condition_embedder_minimax_music3.py @@ -17,10 +17,11 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ..modeling_utils import ModelMixin -class MiniMaxMusic3ConditionEncoder(ModelMixin, ConfigMixin): +class MiniMaxMusic3ConditionEncoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" Projects the per-frame hidden states of the autoregressive stage onto the Flow-VAE latent timeline. diff --git a/src/diffusers/models/controlnets/controlnet_hunyuan.py b/src/diffusers/models/controlnets/controlnet_hunyuan.py index 6ef92d78dd6e..b35886907fdc 100644 --- a/src/diffusers/models/controlnets/controlnet_hunyuan.py +++ b/src/diffusers/models/controlnets/controlnet_hunyuan.py @@ -17,6 +17,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput, logging from ..attention_processor import AttentionProcessor from ..embeddings import ( @@ -37,7 +38,7 @@ class HunyuanControlNetOutput(BaseOutput): controlnet_block_samples: tuple[torch.Tensor] -class HunyuanDiT2DControlNetModel(ModelMixin, ConfigMixin): +class HunyuanDiT2DControlNetModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/models/controlnets/controlnet_sana.py b/src/diffusers/models/controlnets/controlnet_sana.py index 4b6e3010ec67..b207b1515766 100644 --- a/src/diffusers/models/controlnets/controlnet_sana.py +++ b/src/diffusers/models/controlnets/controlnet_sana.py @@ -19,7 +19,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from ..attention import AttentionMixin from ..embeddings import PatchEmbed, PixArtAlphaTextProjection @@ -38,7 +38,7 @@ class SanaControlNetOutput(BaseOutput): controlnet_block_samples: tuple[torch.Tensor] -class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin): +class SanaControlNetModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): _supports_gradient_checkpointing = True _no_split_modules = ["SanaTransformerBlock", "PatchEmbed"] _skip_layerwise_casting_patterns = ["patch_embed", "norm"] diff --git a/src/diffusers/models/transformers/ace_step_transformer.py b/src/diffusers/models/transformers/ace_step_transformer.py index 821c7ad1491a..d29091c71cc3 100644 --- a/src/diffusers/models/transformers/ace_step_transformer.py +++ b/src/diffusers/models/transformers/ace_step_transformer.py @@ -21,7 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale, logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import ( @@ -429,7 +429,9 @@ def forward( # --------------------------------------------------------------------------- # -class AceStepTransformer1DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin, CacheMixin): +class AceStepTransformer1DModel( + ModelMixin, ConfigMixin, PeftAdapterMixin, AttentionMixin, CacheMixin, FromOriginalModelMixin +): """Diffusion Transformer for ACE-Step 1.5 music generation. Generates audio latents conditioned on text, lyrics, and timbre. Uses 1D patch embedding (`Conv1d` with stride diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py index 08299f05e1b8..68a1d7faa71b 100644 --- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py +++ b/src/diffusers/models/transformers/cogvideox_transformer_3d.py @@ -19,7 +19,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import Attention, AttentionMixin, FeedForward @@ -157,7 +157,9 @@ def forward( return hidden_states, encoder_hidden_states -class CogVideoXTransformer3DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, CacheMixin): +class CogVideoXTransformer3DModel( + ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, CacheMixin, FromOriginalModelMixin +): """ A Transformer model for video-like data in [CogVideoX](https://github.com/THUDM/CogVideo). diff --git a/src/diffusers/models/transformers/dit_transformer_2d.py b/src/diffusers/models/transformers/dit_transformer_2d.py index 0457acf77108..8e4b66902e7e 100644 --- a/src/diffusers/models/transformers/dit_transformer_2d.py +++ b/src/diffusers/models/transformers/dit_transformer_2d.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import BasicTransformerBlock from ..embeddings import PatchEmbed @@ -28,7 +29,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class DiTTransformer2DModel(ModelMixin, ConfigMixin): +class DiTTransformer2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" A 2D Transformer model as introduced in DiT (https://huggingface.co/papers/2212.09748). diff --git a/src/diffusers/models/transformers/hunyuan_transformer_2d.py b/src/diffusers/models/transformers/hunyuan_transformer_2d.py index 83b3797c4fc3..d52e35e6454b 100644 --- a/src/diffusers/models/transformers/hunyuan_transformer_2d.py +++ b/src/diffusers/models/transformers/hunyuan_transformer_2d.py @@ -15,6 +15,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, FeedForward @@ -198,7 +199,7 @@ def forward( return hidden_states -class HunyuanDiT2DModel(ModelMixin, AttentionMixin, ConfigMixin): +class HunyuanDiT2DModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): """ HunYuanDiT: Diffusion model with a Transformer backbone. diff --git a/src/diffusers/models/transformers/lumina_nextdit2d.py b/src/diffusers/models/transformers/lumina_nextdit2d.py index 73468b5d853f..626315fda44c 100644 --- a/src/diffusers/models/transformers/lumina_nextdit2d.py +++ b/src/diffusers/models/transformers/lumina_nextdit2d.py @@ -18,6 +18,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import LuminaFeedForward from ..attention_processor import Attention, LuminaAttnProcessor2_0 @@ -175,7 +176,7 @@ def forward( return hidden_states -class LuminaNextDiT2DModel(ModelMixin, ConfigMixin): +class LuminaNextDiT2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ LuminaNextDiT: Diffusion model with a Transformer backbone. diff --git a/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py b/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py index 97bee84933c5..fe59aca34b01 100644 --- a/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py +++ b/src/diffusers/models/transformers/minimax_music3_rvq_depth_decoder.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ..attention import AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..modeling_utils import ModelMixin @@ -88,7 +89,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states + self.down_proj(F.silu(self.gate_proj(norm_states)) * self.up_proj(norm_states)) -class MiniMaxMusic3RVQDepthDecoder(ModelMixin, ConfigMixin): +class MiniMaxMusic3RVQDepthDecoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" The local language model of MiniMax Music 3. Within each audio frame it autoregressively predicts the seven residual RVQ codebooks (c1..c7) from the global language model's hidden state and the frame's semantic code, and diff --git a/src/diffusers/models/transformers/pixart_transformer_2d.py b/src/diffusers/models/transformers/pixart_transformer_2d.py index e5e6178eaf4a..ededb33d33e2 100644 --- a/src/diffusers/models/transformers/pixart_transformer_2d.py +++ b/src/diffusers/models/transformers/pixart_transformer_2d.py @@ -17,6 +17,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import AttentionMixin, BasicTransformerBlock from ..attention_processor import Attention, AttnProcessor, FusedAttnProcessor2_0 @@ -29,7 +30,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class PixArtTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin): +class PixArtTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): r""" A 2D Transformer model as introduced in PixArt family of models (https://huggingface.co/papers/2310.00426, https://huggingface.co/papers/2403.04692). diff --git a/src/diffusers/models/transformers/prior_transformer.py b/src/diffusers/models/transformers/prior_transformer.py index f3890446e28e..3fbd56e92cba 100644 --- a/src/diffusers/models/transformers/prior_transformer.py +++ b/src/diffusers/models/transformers/prior_transformer.py @@ -5,7 +5,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin from ...utils import BaseOutput from ..attention import AttentionMixin, BasicTransformerBlock from ..attention_processor import ( @@ -31,7 +31,9 @@ class PriorTransformerOutput(BaseOutput): predicted_image_embedding: torch.Tensor -class PriorTransformer(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): +class PriorTransformer( + ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin, FromOriginalModelMixin +): """ A Prior Transformer model. diff --git a/src/diffusers/models/transformers/stable_audio_transformer.py b/src/diffusers/models/transformers/stable_audio_transformer.py index f4974926ec72..6f065ed446ab 100644 --- a/src/diffusers/models/transformers/stable_audio_transformer.py +++ b/src/diffusers/models/transformers/stable_audio_transformer.py @@ -18,6 +18,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import AttentionMixin, FeedForward @@ -180,7 +181,7 @@ def forward( return hidden_states -class StableAudioDiTModel(ModelMixin, AttentionMixin, ConfigMixin): +class StableAudioDiTModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): """ The Diffusion Transformer model introduced in Stable Audio. diff --git a/src/diffusers/models/transformers/t5_film_transformer.py b/src/diffusers/models/transformers/t5_film_transformer.py index 547e72089990..adb6c8d92a6a 100644 --- a/src/diffusers/models/transformers/t5_film_transformer.py +++ b/src/diffusers/models/transformers/t5_film_transformer.py @@ -17,12 +17,13 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ..attention_processor import Attention from ..embeddings import get_timestep_embedding from ..modeling_utils import ModelMixin -class T5FilmDecoder(ModelMixin, ConfigMixin): +class T5FilmDecoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" T5 style decoder with FiLM conditioning. diff --git a/src/diffusers/models/transformers/transformer_2d.py b/src/diffusers/models/transformers/transformer_2d.py index 6714383b77ab..4c3261d1c315 100644 --- a/src/diffusers/models/transformers/transformer_2d.py +++ b/src/diffusers/models/transformers/transformer_2d.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import LegacyConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import deprecate, logging from ..attention import BasicTransformerBlock from ..embeddings import ImagePositionalEmbeddings, PatchEmbed, PixArtAlphaTextProjection @@ -36,7 +37,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) -class Transformer2DModel(LegacyModelMixin, LegacyConfigMixin): +class Transformer2DModel(LegacyModelMixin, LegacyConfigMixin, FromOriginalModelMixin): """ A 2D Transformer model for image-like data. diff --git a/src/diffusers/models/transformers/transformer_cogview3plus.py b/src/diffusers/models/transformers/transformer_cogview3plus.py index ad6a442acbcc..fd2a6aae1d9a 100644 --- a/src/diffusers/models/transformers/transformer_cogview3plus.py +++ b/src/diffusers/models/transformers/transformer_cogview3plus.py @@ -17,6 +17,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import AttentionMixin, FeedForward from ..attention_processor import Attention, CogVideoXAttnProcessor2_0 @@ -123,7 +124,7 @@ def forward( return hidden_states, encoder_hidden_states -class CogView3PlusTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin): +class CogView3PlusTransformer2DModel(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): r""" The Transformer model introduced in [CogView3: Finer and Faster Text-to-Image Generation via Relay Diffusion](https://huggingface.co/papers/2403.05121). diff --git a/src/diffusers/models/transformers/transformer_cogview4.py b/src/diffusers/models/transformers/transformer_cogview4.py index 2856fffd2a63..cf1f56f85b93 100644 --- a/src/diffusers/models/transformers/transformer_cogview4.py +++ b/src/diffusers/models/transformers/transformer_cogview4.py @@ -19,7 +19,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale, logging from ...utils.torch_utils import maybe_allow_in_graph from ..attention import FeedForward @@ -612,7 +612,7 @@ def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torc return x -class CogView4Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, CacheMixin): +class CogView4Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, CacheMixin, FromOriginalModelMixin): r""" Args: patch_size (`int`, defaults to `2`): diff --git a/src/diffusers/models/transformers/transformer_joyimage.py b/src/diffusers/models/transformers/transformer_joyimage.py index b17ddb05f799..ca2cc8de7777 100644 --- a/src/diffusers/models/transformers/transformer_joyimage.py +++ b/src/diffusers/models/transformers/transformer_joyimage.py @@ -20,6 +20,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -369,7 +370,7 @@ def forward( # --------------------------------------------------------------------------- -class JoyImageEditTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin): +class JoyImageEditTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): """JoyImage Transformer model for image generation / editing. Dual-stream DiT architecture with WAN-style conditioning embeddings and custom rotary position embeddings. diff --git a/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py index 4a13845faad3..5587e47a5710 100644 --- a/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py +++ b/src/diffusers/models/transformers/transformer_joyimage_edit_plus.py @@ -20,6 +20,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward from ..attention_dispatch import dispatch_attention_fn @@ -314,7 +315,7 @@ def forward( return temb, timestep_proj, encoder_hidden_states -class JoyImageEditPlusTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin): +class JoyImageEditPlusTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): r""" JoyImage Edit Plus Transformer for multi-image editing. diff --git a/src/diffusers/models/transformers/transformer_longcat_audio_dit.py b/src/diffusers/models/transformers/transformer_longcat_audio_dit.py index 9b8c0b4bf147..7362e54ab2d5 100644 --- a/src/diffusers/models/transformers/transformer_longcat_audio_dit.py +++ b/src/diffusers/models/transformers/transformer_longcat_audio_dit.py @@ -23,6 +23,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ...utils.torch_utils import lru_cache_unless_export, maybe_allow_in_graph from ..attention import AttentionModuleMixin @@ -452,7 +453,7 @@ def forward( return hidden_states -class LongCatAudioDiTTransformer(ModelMixin, ConfigMixin): +class LongCatAudioDiTTransformer(ModelMixin, ConfigMixin, FromOriginalModelMixin): _supports_gradient_checkpointing = False _repeated_blocks = ["AudioDiTBlock"] diff --git a/src/diffusers/models/transformers/transformer_minimax_h3.py b/src/diffusers/models/transformers/transformer_minimax_h3.py index f49cdaca2eb6..17c1d461903b 100644 --- a/src/diffusers/models/transformers/transformer_minimax_h3.py +++ b/src/diffusers/models/transformers/transformer_minimax_h3.py @@ -19,7 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import BaseOutput, apply_lora_scale, logging from .._modeling_parallel import ContextParallelInput, ContextParallelOutput from ..attention import AttentionMixin, AttentionModuleMixin, FeedForward @@ -373,7 +373,9 @@ def forward( return hidden_states -class MiniMaxH3Transformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin, CacheMixin): +class MiniMaxH3Transformer3DModel( + ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin, CacheMixin, FromOriginalModelMixin +): r""" A Transformer model for joint video + audio generation, introduced in MiniMax-H3. diff --git a/src/diffusers/models/transformers/transformer_minimax_music3.py b/src/diffusers/models/transformers/transformer_minimax_music3.py index 6ff174fe5e87..d36dd540b54a 100644 --- a/src/diffusers/models/transformers/transformer_minimax_music3.py +++ b/src/diffusers/models/transformers/transformer_minimax_music3.py @@ -19,6 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils.torch_utils import lru_cache_unless_export from ..attention import AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn @@ -144,7 +145,7 @@ def forward(self, hidden_states: torch.Tensor, rotary_emb: Tuple[torch.Tensor, t return hidden_states -class MiniMaxMusic3Transformer1DModel(ModelMixin, ConfigMixin): +class MiniMaxMusic3Transformer1DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" The flow-matching diffusion transformer of MiniMax Music 3. It denoises Flow-VAE audio latents conditioned on per-frame hidden states produced by the autoregressive language-model stage. diff --git a/src/diffusers/models/transformers/transformer_omnigen.py b/src/diffusers/models/transformers/transformer_omnigen.py index f860f5d5ab3e..415c9fb470ef 100644 --- a/src/diffusers/models/transformers/transformer_omnigen.py +++ b/src/diffusers/models/transformers/transformer_omnigen.py @@ -19,6 +19,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ..attention_processor import Attention from ..embeddings import TimestepEmbedding, Timesteps, get_2d_sincos_pos_embed @@ -280,7 +281,7 @@ def forward( return hidden_states -class OmniGenTransformer2DModel(ModelMixin, ConfigMixin): +class OmniGenTransformer2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ The Transformer model introduced in OmniGen (https://huggingface.co/papers/2409.11340). diff --git a/src/diffusers/models/transformers/transformer_prx.py b/src/diffusers/models/transformers/transformer_prx.py index 33e1a44dad92..40072c1b158a 100644 --- a/src/diffusers/models/transformers/transformer_prx.py +++ b/src/diffusers/models/transformers/transformer_prx.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import logging from ...utils.torch_utils import maybe_adjust_dtype_for_device from ..attention import AttentionMixin, AttentionModuleMixin @@ -631,7 +632,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te return seq -class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): +class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): r""" Transformer-based 2D model for text to image generation. diff --git a/src/diffusers/models/transformers/transformer_stable_audio3.py b/src/diffusers/models/transformers/transformer_stable_audio3.py index aab80b744d8f..b6f319e40671 100644 --- a/src/diffusers/models/transformers/transformer_stable_audio3.py +++ b/src/diffusers/models/transformers/transformer_stable_audio3.py @@ -21,6 +21,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput, logging from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn @@ -400,7 +401,7 @@ def forward( return hidden_states -class StableAudio3DiTModel(ModelMixin, ConfigMixin, AttentionMixin): +class StableAudio3DiTModel(ModelMixin, ConfigMixin, AttentionMixin, FromOriginalModelMixin): r""" The Diffusion Transformer (DiT) backbone of [Stable Audio 3](https://stability.ai/news/stable-audio-3). diff --git a/src/diffusers/models/unets/unet_1d.py b/src/diffusers/models/unets/unet_1d.py index 959e82e9d7cd..3160a3804f8d 100644 --- a/src/diffusers/models/unets/unet_1d.py +++ b/src/diffusers/models/unets/unet_1d.py @@ -18,6 +18,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ..embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin @@ -37,7 +38,7 @@ class UNet1DOutput(BaseOutput): sample: torch.Tensor -class UNet1DModel(ModelMixin, ConfigMixin): +class UNet1DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" A 1D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. diff --git a/src/diffusers/models/unets/unet_2d.py b/src/diffusers/models/unets/unet_2d.py index 4bbe0535e94a..273b0b896025 100644 --- a/src/diffusers/models/unets/unet_2d.py +++ b/src/diffusers/models/unets/unet_2d.py @@ -17,6 +17,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput from ..embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps from ..modeling_utils import ModelMixin @@ -36,7 +37,7 @@ class UNet2DOutput(BaseOutput): sample: torch.Tensor -class UNet2DModel(ModelMixin, ConfigMixin): +class UNet2DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" A 2D UNet model that takes a noisy sample and a timestep and returns a sample shaped output. diff --git a/src/diffusers/models/unets/unet_3d_condition.py b/src/diffusers/models/unets/unet_3d_condition.py index 0d15e93da68f..cc1bde9df04f 100644 --- a/src/diffusers/models/unets/unet_3d_condition.py +++ b/src/diffusers/models/unets/unet_3d_condition.py @@ -20,7 +20,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin +from ...loaders import FromOriginalModelMixin, UNet2DConditionLoadersMixin from ...utils import BaseOutput, logging from ...utils.torch_utils import maybe_adjust_dtype_for_device from ..activations import get_activation @@ -59,7 +59,9 @@ class UNet3DConditionOutput(BaseOutput): sample: torch.Tensor -class UNet3DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class UNet3DConditionModel( + ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, FromOriginalModelMixin +): r""" A conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output. diff --git a/src/diffusers/models/unets/unet_i2vgen_xl.py b/src/diffusers/models/unets/unet_i2vgen_xl.py index 9e7841f95e58..31606e68b875 100644 --- a/src/diffusers/models/unets/unet_i2vgen_xl.py +++ b/src/diffusers/models/unets/unet_i2vgen_xl.py @@ -18,7 +18,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin +from ...loaders import FromOriginalModelMixin, UNet2DConditionLoadersMixin from ...utils import logging from ...utils.torch_utils import maybe_adjust_dtype_for_device from ..activations import get_activation @@ -93,7 +93,7 @@ def forward( return hidden_states -class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class I2VGenXLUNet(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, FromOriginalModelMixin): r""" I2VGenXL UNet. It is a conditional 3D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample-shaped output. diff --git a/src/diffusers/models/unets/unet_kandinsky3.py b/src/diffusers/models/unets/unet_kandinsky3.py index 790d255101a4..629b3484ec14 100644 --- a/src/diffusers/models/unets/unet_kandinsky3.py +++ b/src/diffusers/models/unets/unet_kandinsky3.py @@ -18,6 +18,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...utils import BaseOutput, logging from ..attention import AttentionMixin from ..attention_processor import Attention, AttnProcessor @@ -45,7 +46,7 @@ def forward(self, x): return x -class Kandinsky3UNet(ModelMixin, AttentionMixin, ConfigMixin): +class Kandinsky3UNet(ModelMixin, AttentionMixin, ConfigMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/models/unets/unet_spatio_temporal_condition.py b/src/diffusers/models/unets/unet_spatio_temporal_condition.py index 7d6b861df82b..6147f3723883 100644 --- a/src/diffusers/models/unets/unet_spatio_temporal_condition.py +++ b/src/diffusers/models/unets/unet_spatio_temporal_condition.py @@ -4,7 +4,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin +from ...loaders import FromOriginalModelMixin, UNet2DConditionLoadersMixin from ...utils import BaseOutput, logging from ...utils.torch_utils import maybe_adjust_dtype_for_device from ..attention import AttentionMixin @@ -30,7 +30,9 @@ class UNetSpatioTemporalConditionOutput(BaseOutput): sample: torch.Tensor = None -class UNetSpatioTemporalConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class UNetSpatioTemporalConditionModel( + ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, FromOriginalModelMixin +): r""" A conditional Spatio-Temporal UNet model that takes a noisy video frames, conditional state, and a timestep and returns a sample shaped output. diff --git a/src/diffusers/models/unets/uvit_2d.py b/src/diffusers/models/unets/uvit_2d.py index 317abe80b1eb..2687417b33c3 100644 --- a/src/diffusers/models/unets/uvit_2d.py +++ b/src/diffusers/models/unets/uvit_2d.py @@ -20,7 +20,7 @@ from torch.utils.checkpoint import checkpoint from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...utils import apply_lora_scale from ..attention import AttentionMixin, BasicTransformerBlock, SkipFFTransformerBlock from ..attention_processor import ( @@ -35,7 +35,7 @@ from ..resnet import Downsample2D, Upsample2D -class UVit2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin): +class UVit2DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin): _supports_gradient_checkpointing = True @register_to_config diff --git a/src/diffusers/pipelines/ace_step/modeling_ace_step.py b/src/diffusers/pipelines/ace_step/modeling_ace_step.py index 769b07044420..7a0a4ffae1d3 100644 --- a/src/diffusers/pipelines/ace_step/modeling_ace_step.py +++ b/src/diffusers/pipelines/ace_step/modeling_ace_step.py @@ -27,6 +27,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin from ...models.normalization import RMSNorm from ...models.transformers.ace_step_transformer import ( @@ -562,7 +563,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states.reshape(batch_size, num_patches, -1) -class AceStepAudioTokenDetokenizer(ModelMixin, ConfigMixin): +class AceStepAudioTokenDetokenizer(ModelMixin, ConfigMixin, FromOriginalModelMixin): """Expands ACE-Step 5 Hz audio tokens back to 25 Hz acoustic conditioning.""" _supports_gradient_checkpointing = True @@ -662,7 +663,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return hidden_states.reshape(batch_size, num_tokens * self.pool_window_size, -1) -class AceStepAudioTokenizer(ModelMixin, ConfigMixin): +class AceStepAudioTokenizer(ModelMixin, ConfigMixin, FromOriginalModelMixin): """Converts 25 Hz acoustic latents to ACE-Step 5 Hz audio tokens.""" _supports_gradient_checkpointing = True @@ -749,7 +750,7 @@ def tokenize( # --------------------------------------------------------------------------- # -class AceStepConditionEncoder(ModelMixin, ConfigMixin): +class AceStepConditionEncoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): """Fuses text + lyric + timbre conditioning into the packed sequence used by the DiT's cross-attention. """ diff --git a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py index dc2aa1df66a7..a795cbe0c5ae 100644 --- a/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py +++ b/src/diffusers/pipelines/audioldm2/modeling_audioldm2.py @@ -19,7 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import UNet2DConditionLoadersMixin +from ...loaders import FromOriginalModelMixin, UNet2DConditionLoadersMixin from ...models.activations import get_activation from ...models.attention import AttentionMixin from ...models.attention_processor import ( @@ -75,7 +75,7 @@ class AudioLDM2ProjectionModelOutput(BaseOutput): attention_mask: torch.LongTensor | None = None -class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin): +class AudioLDM2ProjectionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ A simple linear projection model to map two text embeddings to a shared latent space. It also inserts learned embedding vectors at the start and end of each text embedding sequence respectively. Each variable appended with @@ -160,7 +160,9 @@ def forward( ) -class AudioLDM2UNet2DConditionModel(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin): +class AudioLDM2UNet2DConditionModel( + ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, FromOriginalModelMixin +): r""" A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output. Compared to the vanilla [`UNet2DConditionModel`], this variant optionally includes an additional diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py index 51716e6acc65..4f9cb82a0e08 100644 --- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py +++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/continuous_encoder.py @@ -23,10 +23,11 @@ ) from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin -class SpectrogramContEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): +class SpectrogramContEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py index a0b0e37e86d4..2144556efed4 100644 --- a/src/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py +++ b/src/diffusers/pipelines/deprecated/spectrogram_diffusion/notes_encoder.py @@ -19,10 +19,11 @@ from transformers.models.t5.modeling_t5 import T5Block, T5Config, T5LayerNorm from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin -class SpectrogramNotesEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): +class SpectrogramNotesEncoder(ModelMixin, ConfigMixin, ModuleUtilsMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/pipelines/deprecated/unclip/text_proj.py b/src/diffusers/pipelines/deprecated/unclip/text_proj.py index 5493df794acc..dfa6f80a3efa 100644 --- a/src/diffusers/pipelines/deprecated/unclip/text_proj.py +++ b/src/diffusers/pipelines/deprecated/unclip/text_proj.py @@ -16,10 +16,11 @@ from torch import nn from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin -class UnCLIPTextProjModel(ModelMixin, ConfigMixin): +class UnCLIPTextProjModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Utility class for CLIP embeddings. Used to combine the image and text embeddings into a format usable by the decoder. diff --git a/src/diffusers/pipelines/deprecated/unidiffuser/modeling_text_decoder.py b/src/diffusers/pipelines/deprecated/unidiffuser/modeling_text_decoder.py index a068f99c6368..f407e37d9800 100644 --- a/src/diffusers/pipelines/deprecated/unidiffuser/modeling_text_decoder.py +++ b/src/diffusers/pipelines/deprecated/unidiffuser/modeling_text_decoder.py @@ -5,11 +5,12 @@ from transformers.modeling_utils import ModuleUtilsMixin from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin # Modified from ClipCaptionModel in https://github.com/thu-ml/unidiffuser/blob/main/libs/caption_decoder.py -class UniDiffuserTextDecoder(ModelMixin, ConfigMixin, ModuleUtilsMixin): +class UniDiffuserTextDecoder(ModelMixin, ConfigMixin, ModuleUtilsMixin, FromOriginalModelMixin): """ Text decoder model for a image-text [UniDiffuser](https://huggingface.co/papers/2303.06555) model. This is used to generate text from the UniDiffuser image-text embedding. diff --git a/src/diffusers/pipelines/deprecated/unidiffuser/modeling_uvit.py b/src/diffusers/pipelines/deprecated/unidiffuser/modeling_uvit.py index 80af8e0fe801..de239e343c24 100644 --- a/src/diffusers/pipelines/deprecated/unidiffuser/modeling_uvit.py +++ b/src/diffusers/pipelines/deprecated/unidiffuser/modeling_uvit.py @@ -4,6 +4,7 @@ from torch import nn from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin from ....models.attention import FeedForward from ....models.attention_processor import Attention @@ -829,7 +830,7 @@ def forward( return Transformer2DModelOutput(sample=output) -class UniDiffuserModel(ModelMixin, ConfigMixin): +class UniDiffuserModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Transformer model for a image-text [UniDiffuser](https://huggingface.co/papers/2303.06555) model. This is a modification of [`UTransformer2DModel`] with input and output heads for the VAE-embedded latent image, the diff --git a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py index 7be159d77af5..ecf6db3cda66 100644 --- a/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py +++ b/src/diffusers/pipelines/deprecated/versatile_diffusion/modeling_text_unet.py @@ -8,6 +8,7 @@ from diffusers.utils import deprecate from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models import ModelMixin from ....models.activations import get_activation from ....models.attention_processor import ( @@ -268,7 +269,7 @@ def forward( return objs -class UNetFlatConditionModel(ModelMixin, ConfigMixin): +class UNetFlatConditionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output. diff --git a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py index 8f3b9512e888..004eecfb5405 100644 --- a/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py +++ b/src/diffusers/pipelines/deprecated/vq_diffusion/pipeline_vq_diffusion.py @@ -18,6 +18,7 @@ from transformers import CLIPTextModel, CLIPTokenizer from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders.single_file_model import FromOriginalModelMixin from ....models import ModelMixin, Transformer2DModel, VQModel from ....schedulers import VQDiffusionScheduler from ....utils import logging @@ -27,7 +28,7 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin): +class LearnedClassifierFreeSamplingEmbeddings(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Utility class for storing learned text embeddings for classifier free sampling """ diff --git a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_paella_vq_model.py b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_paella_vq_model.py index a18ab8eed6ec..f61b9c652f70 100644 --- a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_paella_vq_model.py +++ b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_paella_vq_model.py @@ -18,6 +18,7 @@ import torch.nn as nn from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models.autoencoders.vae import DecoderOutput, VectorQuantizer from ....models.modeling_utils import ModelMixin from ....models.vq_model import VQEncoderOutput @@ -54,7 +55,7 @@ def forward(self, x): return x -class PaellaVQModel(ModelMixin, ConfigMixin): +class PaellaVQModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): r"""VQ-VAE model from Paella model. This model inherits from [`ModelMixin`]. Check the superclass documentation for the generic methods the library diff --git a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_diffnext.py b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_diffnext.py index 5c6c12264da9..25d48ce9cc37 100644 --- a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_diffnext.py +++ b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_diffnext.py @@ -20,11 +20,12 @@ import torch.nn as nn from ....configuration_utils import ConfigMixin, register_to_config +from ....loaders import FromOriginalModelMixin from ....models.modeling_utils import ModelMixin from .modeling_wuerstchen_common import AttnBlock, GlobalResponseNorm, TimestepBlock, WuerstchenLayerNorm -class WuerstchenDiffNeXt(ModelMixin, ConfigMixin): +class WuerstchenDiffNeXt(ModelMixin, ConfigMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_prior.py b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_prior.py index f19eec46ec3b..e7d8af9954e6 100644 --- a/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_prior.py +++ b/src/diffusers/pipelines/deprecated/wuerstchen/modeling_wuerstchen_prior.py @@ -19,7 +19,7 @@ import torch.nn as nn from ....configuration_utils import ConfigMixin, register_to_config -from ....loaders import PeftAdapterMixin, UNet2DConditionLoadersMixin +from ....loaders import FromOriginalModelMixin, PeftAdapterMixin, UNet2DConditionLoadersMixin from ....models.attention import AttentionMixin from ....models.attention_processor import ( ADDED_KV_ATTENTION_PROCESSORS, @@ -31,7 +31,9 @@ from .modeling_wuerstchen_common import AttnBlock, ResBlock, TimestepBlock, WuerstchenLayerNorm -class WuerstchenPrior(ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin): +class WuerstchenPrior( + ModelMixin, AttentionMixin, ConfigMixin, UNet2DConditionLoadersMixin, PeftAdapterMixin, FromOriginalModelMixin +): unet_name = "prior" _supports_gradient_checkpointing = True diff --git a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py index f579cf00dbe7..d020b1242647 100644 --- a/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py +++ b/src/diffusers/pipelines/ltx/modeling_latent_upsampler.py @@ -15,6 +15,7 @@ import torch from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin @@ -73,7 +74,7 @@ def forward(self, x): return x.unflatten(1, (-1, *self.upscale_factors[:1])).permute(0, 1, 3, 2, 4, 5).flatten(2, 3) -class LTXLatentUpsamplerModel(ModelMixin, ConfigMixin): +class LTXLatentUpsamplerModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Model to spatially upsample VAE latents. diff --git a/src/diffusers/pipelines/ltx2/connectors.py b/src/diffusers/pipelines/ltx2/connectors.py index 0b0b333209aa..5bcf3b2135ea 100644 --- a/src/diffusers/pipelines/ltx2/connectors.py +++ b/src/diffusers/pipelines/ltx2/connectors.py @@ -4,7 +4,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config -from ...loaders import PeftAdapterMixin +from ...loaders import FromOriginalModelMixin, PeftAdapterMixin from ...models.attention import FeedForward from ...models.modeling_utils import ModelMixin from ...models.transformers.transformer_ltx2 import LTX2Attention, LTX2AudioVideoAttnProcessor @@ -332,7 +332,7 @@ def forward( return hidden_states, attention_mask -class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin): +class LTX2TextConnectors(ModelMixin, PeftAdapterMixin, ConfigMixin, FromOriginalModelMixin): """ Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio streams. diff --git a/src/diffusers/pipelines/ltx2/duration_head.py b/src/diffusers/pipelines/ltx2/duration_head.py index a396e0cd628c..6bb994dbc12f 100644 --- a/src/diffusers/pipelines/ltx2/duration_head.py +++ b/src/diffusers/pipelines/ltx2/duration_head.py @@ -16,6 +16,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.attention import AttentionModuleMixin from ...models.attention_dispatch import dispatch_attention_fn from ...models.modeling_utils import ModelMixin @@ -78,7 +79,7 @@ def forward(self, tokens: torch.Tensor) -> torch.Tensor: return self.processor(self, tokens) -class LTX2DurationHead(ModelMixin, ConfigMixin): +class LTX2DurationHead(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Predicts the natural duration of the shot implied by a caption, from the LTX-2 text connector outputs. diff --git a/src/diffusers/pipelines/ltx2/latent_upsampler.py b/src/diffusers/pipelines/ltx2/latent_upsampler.py index 329ced36d45b..203e4b611c42 100644 --- a/src/diffusers/pipelines/ltx2/latent_upsampler.py +++ b/src/diffusers/pipelines/ltx2/latent_upsampler.py @@ -18,6 +18,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin @@ -167,7 +168,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin): +class LTX2LatentUpsamplerModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Model to spatially upsample VAE latents. diff --git a/src/diffusers/pipelines/ltx2/vocoder.py b/src/diffusers/pipelines/ltx2/vocoder.py index f0004f2ec02d..0f37c04c7307 100644 --- a/src/diffusers/pipelines/ltx2/vocoder.py +++ b/src/diffusers/pipelines/ltx2/vocoder.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin @@ -276,7 +277,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -class LTX2Vocoder(ModelMixin, ConfigMixin): +class LTX2Vocoder(ModelMixin, ConfigMixin, FromOriginalModelMixin): r""" LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms. """ @@ -476,7 +477,7 @@ def forward(self, waveform: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, t return log_mel, magnitude, phase, energy -class LTX2VocoderWithBWE(ModelMixin, ConfigMixin): +class LTX2VocoderWithBWE(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ LTX-2.X vocoder with bandwidth extension (BWE) upsampling. The vocoder and the BWE module run in sequence, with the BWE module upsampling the vocoder output waveform to a higher sampling rate. The BWE module itself has the same diff --git a/src/diffusers/pipelines/shap_e/renderer.py b/src/diffusers/pipelines/shap_e/renderer.py index 0c2058c887fc..de4e6ffd289c 100644 --- a/src/diffusers/pipelines/shap_e/renderer.py +++ b/src/diffusers/pipelines/shap_e/renderer.py @@ -21,6 +21,7 @@ from torch import nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models import ModelMixin from ...utils import BaseOutput from .camera import create_pan_cameras @@ -780,7 +781,7 @@ def forward(self, x: torch.Tensor): return out -class ShapERenderer(ModelMixin, ConfigMixin): +class ShapERenderer(ModelMixin, ConfigMixin, FromOriginalModelMixin): @register_to_config def __init__( self, diff --git a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py index d40269411bc0..402fc4d2a469 100644 --- a/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py +++ b/src/diffusers/pipelines/stable_audio/modeling_stable_audio.py @@ -19,6 +19,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin from ...utils import BaseOutput, logging @@ -111,7 +112,7 @@ def forward( return float_embeds -class StableAudioProjectionModel(ModelMixin, ConfigMixin): +class StableAudioProjectionModel(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ A simple linear projection model to map the conditioning values to a shared latent space. diff --git a/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py b/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py index 52def1d8871d..9b86ad3d8ac8 100644 --- a/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py +++ b/src/diffusers/pipelines/stable_audio_3/modeling_stable_audio_3.py @@ -33,6 +33,7 @@ import torch.nn as nn from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import FromOriginalModelMixin from ...models.modeling_utils import ModelMixin from ...utils import logging @@ -40,7 +41,7 @@ logger = logging.get_logger(__name__) -class StableAudio3DurationEmbedder(ModelMixin, ConfigMixin): +class StableAudio3DurationEmbedder(ModelMixin, ConfigMixin, FromOriginalModelMixin): """ Embeds a duration value (in seconds) into a global conditioning vector for the Stable Audio 3 DiT (used as the ``global_hidden_states`` AdaLN input). diff --git a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py index 0c8fd842fcba..6b9f5ace6ea0 100644 --- a/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py +++ b/src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py @@ -14,7 +14,6 @@ # limitations under the License. """Conversion script for the Stable Diffusion checkpoints.""" -import re from contextlib import nullcontext from io import BytesIO @@ -33,6 +32,8 @@ CLIPVisionModelWithProjection, ) +from ...loaders.conversion import get_conversion +from ...loaders.conversion.checkpoint import convert_component_checkpoint from ...models import ( AutoencoderKL, ControlNetModel, @@ -67,172 +68,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -def shave_segments(path, n_shave_prefix_segments=1): - """ - Removes segments. Positive values shave the first segments, negative shave the last segments. - """ - if n_shave_prefix_segments >= 0: - return ".".join(path.split(".")[n_shave_prefix_segments:]) - else: - return ".".join(path.split(".")[:n_shave_prefix_segments]) - - -def renew_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item.replace("in_layers.0", "norm1") - new_item = new_item.replace("in_layers.2", "conv1") - - new_item = new_item.replace("out_layers.0", "norm2") - new_item = new_item.replace("out_layers.3", "conv2") - - new_item = new_item.replace("emb_layers.1", "time_emb_proj") - new_item = new_item.replace("skip_connection", "conv_shortcut") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside resnets to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("nin_shortcut", "conv_shortcut") - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - # new_item = new_item.replace('norm.weight', 'group_norm.weight') - # new_item = new_item.replace('norm.bias', 'group_norm.bias') - - # new_item = new_item.replace('proj_out.weight', 'proj_attn.weight') - # new_item = new_item.replace('proj_out.bias', 'proj_attn.bias') - - # new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0): - """ - Updates paths inside attentions to the new naming scheme (local renaming) - """ - mapping = [] - for old_item in old_list: - new_item = old_item - - new_item = new_item.replace("norm.weight", "group_norm.weight") - new_item = new_item.replace("norm.bias", "group_norm.bias") - - new_item = new_item.replace("q.weight", "to_q.weight") - new_item = new_item.replace("q.bias", "to_q.bias") - - new_item = new_item.replace("k.weight", "to_k.weight") - new_item = new_item.replace("k.bias", "to_k.bias") - - new_item = new_item.replace("v.weight", "to_v.weight") - new_item = new_item.replace("v.bias", "to_v.bias") - - new_item = new_item.replace("proj_out.weight", "to_out.0.weight") - new_item = new_item.replace("proj_out.bias", "to_out.0.bias") - - new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments) - - mapping.append({"old": old_item, "new": new_item}) - - return mapping - - -def assign_to_checkpoint( - paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None -): - """ - This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits - attention layers, and takes into account additional replacements that may arise. - - Assigns the weights to the new checkpoint. - """ - assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." - - # Splits the attention layers into three variables. - if attention_paths_to_split is not None: - for path, path_map in attention_paths_to_split.items(): - old_tensor = old_checkpoint[path] - channels = old_tensor.shape[0] // 3 - - target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) - - num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 - - old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) - query, key, value = old_tensor.split(channels // num_heads, dim=1) - - checkpoint[path_map["query"]] = query.reshape(target_shape) - checkpoint[path_map["key"]] = key.reshape(target_shape) - checkpoint[path_map["value"]] = value.reshape(target_shape) - - for path in paths: - new_path = path["new"] - - # These have already been assigned - if attention_paths_to_split is not None and new_path in attention_paths_to_split: - continue - - # Global renaming happens here - new_path = new_path.replace("middle_block.0", "mid_block.resnets.0") - new_path = new_path.replace("middle_block.1", "mid_block.attentions.0") - new_path = new_path.replace("middle_block.2", "mid_block.resnets.1") - - if additional_replacements is not None: - for replacement in additional_replacements: - new_path = new_path.replace(replacement["old"], replacement["new"]) - - # proj_attn.weight has to be converted from conv 1D to linear - is_attn_weight = "proj_attn.weight" in new_path or ("attentions" in new_path and "to_" in new_path) - shape = old_checkpoint[path["old"]].shape - if is_attn_weight and len(shape) == 3: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0] - elif is_attn_weight and len(shape) == 4: - checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0, 0] - else: - checkpoint[new_path] = old_checkpoint[path["old"]] - - -def conv_attn_to_linear(checkpoint): - keys = list(checkpoint.keys()) - attn_keys = ["query.weight", "key.weight", "value.weight"] - for key in keys: - if ".".join(key.split(".")[-2:]) in attn_keys: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0, 0] - elif "proj_attn.weight" in key: - if checkpoint[key].ndim > 2: - checkpoint[key] = checkpoint[key][:, :, 0] - - def create_unet_diffusers_config(original_config, image_size: int, controlnet=False): """ Creates a config for the diffusers based on the config of the LDM model. @@ -395,406 +230,27 @@ def create_ldm_bert_config(original_config): def convert_ldm_unet_checkpoint( checkpoint, config, path=None, extract_ema=False, controlnet=False, skip_extract_state_dict=False ): - """ - Takes a state dict and a config, and returns a converted checkpoint. - """ - + cls = "ControlNetModel" if controlnet else "UNet2DConditionModel" if skip_extract_state_dict: - unet_state_dict = checkpoint - else: - # extract state_dict for UNet - unet_state_dict = {} - keys = list(checkpoint.keys()) - - if controlnet: - unet_key = "control_model." - else: - unet_key = "model.diffusion_model." - - # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA - if sum(k.startswith("model_ema") for k in keys) > 100 and extract_ema: - logger.warning(f"Checkpoint {path} has both EMA and non-EMA weights.") - logger.warning( - "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA" - " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag." - ) - for key in keys: - if key.startswith("model.diffusion_model"): - flat_ema_key = "model_ema." + "".join(key.split(".")[1:]) - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(flat_ema_key) - else: - if sum(k.startswith("model_ema") for k in keys) > 100: - logger.warning( - "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA" - " weights (usually better for inference), please make sure to add the `--extract_ema` flag." - ) - - for key in keys: - if key.startswith(unet_key): - unet_state_dict[key.replace(unet_key, "")] = checkpoint.pop(key) - - new_checkpoint = {} - - new_checkpoint["time_embedding.linear_1.weight"] = unet_state_dict["time_embed.0.weight"] - new_checkpoint["time_embedding.linear_1.bias"] = unet_state_dict["time_embed.0.bias"] - new_checkpoint["time_embedding.linear_2.weight"] = unet_state_dict["time_embed.2.weight"] - new_checkpoint["time_embedding.linear_2.bias"] = unet_state_dict["time_embed.2.bias"] - - if config["class_embed_type"] is None: - # No parameters to port - ... - elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection": - new_checkpoint["class_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["class_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["class_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["class_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - else: - raise NotImplementedError(f"Not implemented `class_embed_type`: {config['class_embed_type']}") - - if config["addition_embed_type"] == "text_time": - new_checkpoint["add_embedding.linear_1.weight"] = unet_state_dict["label_emb.0.0.weight"] - new_checkpoint["add_embedding.linear_1.bias"] = unet_state_dict["label_emb.0.0.bias"] - new_checkpoint["add_embedding.linear_2.weight"] = unet_state_dict["label_emb.0.2.weight"] - new_checkpoint["add_embedding.linear_2.bias"] = unet_state_dict["label_emb.0.2.bias"] - - # Relevant to StableDiffusionUpscalePipeline - if "num_class_embeds" in config: - if (config["num_class_embeds"] is not None) and ("label_emb.weight" in unet_state_dict): - new_checkpoint["class_embedding.weight"] = unet_state_dict["label_emb.weight"] - - new_checkpoint["conv_in.weight"] = unet_state_dict["input_blocks.0.0.weight"] - new_checkpoint["conv_in.bias"] = unet_state_dict["input_blocks.0.0.bias"] - - if not controlnet: - new_checkpoint["conv_norm_out.weight"] = unet_state_dict["out.0.weight"] - new_checkpoint["conv_norm_out.bias"] = unet_state_dict["out.0.bias"] - new_checkpoint["conv_out.weight"] = unet_state_dict["out.2.weight"] - new_checkpoint["conv_out.bias"] = unet_state_dict["out.2.bias"] - - # Retrieves the keys for the input blocks only - num_input_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "input_blocks" in layer}) - input_blocks = { - layer_id: [key for key in unet_state_dict if f"input_blocks.{layer_id}" in key] - for layer_id in range(num_input_blocks) - } - - # Retrieves the keys for the middle blocks only - num_middle_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "middle_block" in layer}) - middle_blocks = { - layer_id: [key for key in unet_state_dict if f"middle_block.{layer_id}" in key] - for layer_id in range(num_middle_blocks) - } - - # Retrieves the keys for the output blocks only - num_output_blocks = len({".".join(layer.split(".")[:2]) for layer in unet_state_dict if "output_blocks" in layer}) - output_blocks = { - layer_id: [key for key in unet_state_dict if f"output_blocks.{layer_id}" in key] - for layer_id in range(num_output_blocks) - } - - for i in range(1, num_input_blocks): - block_id = (i - 1) // (config["layers_per_block"] + 1) - layer_in_block_id = (i - 1) % (config["layers_per_block"] + 1) - - resnets = [ - key for key in input_blocks[i] if f"input_blocks.{i}.0" in key and f"input_blocks.{i}.0.op" not in key - ] - attentions = [key for key in input_blocks[i] if f"input_blocks.{i}.1" in key] - - if f"input_blocks.{i}.0.op.weight" in unet_state_dict: - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.weight"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.weight" - ) - new_checkpoint[f"down_blocks.{block_id}.downsamplers.0.conv.bias"] = unet_state_dict.pop( - f"input_blocks.{i}.0.op.bias" - ) - - paths = renew_resnet_paths(resnets) - meta_path = {"old": f"input_blocks.{i}.0", "new": f"down_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - if len(attentions): - paths = renew_attention_paths(attentions) - - meta_path = {"old": f"input_blocks.{i}.1", "new": f"down_blocks.{block_id}.attentions.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - resnet_0 = middle_blocks[0] - attentions = middle_blocks[1] - resnet_1 = middle_blocks[2] - - resnet_0_paths = renew_resnet_paths(resnet_0) - assign_to_checkpoint(resnet_0_paths, new_checkpoint, unet_state_dict, config=config) - - resnet_1_paths = renew_resnet_paths(resnet_1) - assign_to_checkpoint(resnet_1_paths, new_checkpoint, unet_state_dict, config=config) - - attentions_paths = renew_attention_paths(attentions) - meta_path = {"old": "middle_block.1", "new": "mid_block.attentions.0"} - assign_to_checkpoint( - attentions_paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - for i in range(num_output_blocks): - block_id = i // (config["layers_per_block"] + 1) - layer_in_block_id = i % (config["layers_per_block"] + 1) - output_block_layers = [shave_segments(name, 2) for name in output_blocks[i]] - output_block_list = {} - - for layer in output_block_layers: - layer_id, layer_name = layer.split(".")[0], shave_segments(layer, 1) - if layer_id in output_block_list: - output_block_list[layer_id].append(layer_name) - else: - output_block_list[layer_id] = [layer_name] - - if len(output_block_list) > 1: - resnets = [key for key in output_blocks[i] if f"output_blocks.{i}.0" in key] - attentions = [key for key in output_blocks[i] if f"output_blocks.{i}.1" in key] - - resnet_0_paths = renew_resnet_paths(resnets) - paths = renew_resnet_paths(resnets) - - meta_path = {"old": f"output_blocks.{i}.0", "new": f"up_blocks.{block_id}.resnets.{layer_in_block_id}"} - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - - output_block_list = {k: sorted(v) for k, v in sorted(output_block_list.items())} - if ["conv.bias", "conv.weight"] in output_block_list.values(): - index = list(output_block_list.values()).index(["conv.bias", "conv.weight"]) - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.weight"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.weight" - ] - new_checkpoint[f"up_blocks.{block_id}.upsamplers.0.conv.bias"] = unet_state_dict[ - f"output_blocks.{i}.{index}.conv.bias" - ] - - # Clear attentions as they have been attributed above. - if len(attentions) == 2: - attentions = [] - - if len(attentions): - paths = renew_attention_paths(attentions) - meta_path = { - "old": f"output_blocks.{i}.1", - "new": f"up_blocks.{block_id}.attentions.{layer_in_block_id}", - } - assign_to_checkpoint( - paths, new_checkpoint, unet_state_dict, additional_replacements=[meta_path], config=config - ) - else: - resnet_0_paths = renew_resnet_paths(output_block_layers, n_shave_prefix_segments=1) - for path in resnet_0_paths: - old_path = ".".join(["output_blocks", str(i), path["old"]]) - new_path = ".".join(["up_blocks", str(block_id), "resnets", str(layer_in_block_id), path["new"]]) - - new_checkpoint[new_path] = unet_state_dict[old_path] - - if controlnet: - # conditioning embedding - - orig_index = 0 - - new_checkpoint["controlnet_cond_embedding.conv_in.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_in.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - orig_index += 2 - - diffusers_index = 0 - - while diffusers_index < 6: - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint[f"controlnet_cond_embedding.blocks.{diffusers_index}.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - diffusers_index += 1 - orig_index += 2 - - new_checkpoint["controlnet_cond_embedding.conv_out.weight"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.weight" - ) - new_checkpoint["controlnet_cond_embedding.conv_out.bias"] = unet_state_dict.pop( - f"input_hint_block.{orig_index}.bias" - ) - - # down blocks - for i in range(num_input_blocks): - new_checkpoint[f"controlnet_down_blocks.{i}.weight"] = unet_state_dict.pop(f"zero_convs.{i}.0.weight") - new_checkpoint[f"controlnet_down_blocks.{i}.bias"] = unet_state_dict.pop(f"zero_convs.{i}.0.bias") - - # mid block - new_checkpoint["controlnet_mid_block.weight"] = unet_state_dict.pop("middle_block_out.0.weight") - new_checkpoint["controlnet_mid_block.bias"] = unet_state_dict.pop("middle_block_out.0.bias") - - return new_checkpoint + return get_conversion(cls, config).to_diffusers(checkpoint) + return convert_component_checkpoint(checkpoint, config, cls, extract_ema=extract_ema) def convert_ldm_vae_checkpoint(checkpoint, config): - # extract state dict for VAE - vae_state_dict = {} - keys = list(checkpoint.keys()) - vae_key = "first_stage_model." if any(k.startswith("first_stage_model.") for k in keys) else "" - for key in keys: - if key.startswith(vae_key): - vae_state_dict[key.replace(vae_key, "")] = checkpoint.get(key) - - new_checkpoint = {} - - new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"] - new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"] - new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"] - new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"] - new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"] - new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"] - - new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"] - new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"] - new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"] - new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"] - new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"] - new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"] - - new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"] - new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"] - new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"] - new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"] - - # Retrieves the keys for the encoder down blocks only - num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}) - down_blocks = { - layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks) - } - - # Retrieves the keys for the decoder up blocks only - num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}) - up_blocks = { - layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks) - } - - for i in range(num_down_blocks): - resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key] - - if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict: - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.weight" - ) - new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop( - f"encoder.down.{i}.downsample.conv.bias" - ) - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - - for i in range(num_up_blocks): - block_id = num_up_blocks - 1 - i - resnets = [ - key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key - ] - - if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict: - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.weight" - ] - new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[ - f"decoder.up.{block_id}.upsample.conv.bias" - ] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key] - num_mid_res_blocks = 2 - for i in range(1, num_mid_res_blocks + 1): - resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key] - - paths = renew_vae_resnet_paths(resnets) - meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - - mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key] - paths = renew_vae_attention_paths(mid_attentions) - meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"} - assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config) - conv_attn_to_linear(new_checkpoint) - return new_checkpoint + return convert_component_checkpoint(checkpoint, config, "AutoencoderKL") def convert_ldm_bert_checkpoint(checkpoint, config): - def _copy_attn_layer(hf_attn_layer, pt_attn_layer): - hf_attn_layer.q_proj.weight.data = pt_attn_layer.to_q.weight - hf_attn_layer.k_proj.weight.data = pt_attn_layer.to_k.weight - hf_attn_layer.v_proj.weight.data = pt_attn_layer.to_v.weight - - hf_attn_layer.out_proj.weight = pt_attn_layer.to_out.weight - hf_attn_layer.out_proj.bias = pt_attn_layer.to_out.bias - - def _copy_linear(hf_linear, pt_linear): - hf_linear.weight = pt_linear.weight - hf_linear.bias = pt_linear.bias - - def _copy_layer(hf_layer, pt_layer): - # copy layer norms - _copy_linear(hf_layer.self_attn_layer_norm, pt_layer[0][0]) - _copy_linear(hf_layer.final_layer_norm, pt_layer[1][0]) - - # copy attn - _copy_attn_layer(hf_layer.self_attn, pt_layer[0][1]) - - # copy MLP - pt_mlp = pt_layer[1][1] - _copy_linear(hf_layer.fc1, pt_mlp.net[0][0]) - _copy_linear(hf_layer.fc2, pt_mlp.net[2]) - - def _copy_layers(hf_layers, pt_layers): - for i, hf_layer in enumerate(hf_layers): - if i != 0: - i += i - pt_layer = pt_layers[i : i + 2] - _copy_layer(hf_layer, pt_layer) - - hf_model = LDMBertModel(config).eval() - - # copy embeds - hf_model.model.embed_tokens.weight = checkpoint.transformer.token_emb.weight - hf_model.model.embed_positions.weight.data = checkpoint.transformer.pos_emb.emb.weight - - # copy layer norm - _copy_linear(hf_model.model.layer_norm, checkpoint.transformer.norm) - - # copy hidden layers - _copy_layers(hf_model.model.layers, checkpoint.transformer.attn_layers.layers) - - _copy_linear(hf_model.to_logits, checkpoint.transformer.to_logits) - - return hf_model + state = checkpoint.state_dict() if hasattr(checkpoint, "state_dict") else checkpoint + if any(key.startswith("cond_stage_model.") for key in state): + state = { + key.removeprefix("cond_stage_model."): value + for key, value in state.items() + if key.startswith("cond_stage_model.") + } + model = LDMBertModel(config).eval() + model.load_state_dict(get_conversion("LDMBertModel", config.to_dict()).to_diffusers(state), strict=True) + return model def convert_ldm_clip_checkpoint(checkpoint, local_files_only=False, text_encoder=None): @@ -836,95 +292,17 @@ def convert_ldm_clip_checkpoint(checkpoint, local_files_only=False, text_encoder return text_model -textenc_conversion_lst = [ - ("positional_embedding", "text_model.embeddings.position_embedding.weight"), - ("token_embedding.weight", "text_model.embeddings.token_embedding.weight"), - ("ln_final.weight", "text_model.final_layer_norm.weight"), - ("ln_final.bias", "text_model.final_layer_norm.bias"), - ("text_projection", "text_projection.weight"), -] -textenc_conversion_map = {x[0]: x[1] for x in textenc_conversion_lst} - -textenc_transformer_conversion_lst = [ - # (stable-diffusion, HF Diffusers) - ("resblocks.", "text_model.encoder.layers."), - ("ln_1", "layer_norm1"), - ("ln_2", "layer_norm2"), - (".c_fc.", ".fc1."), - (".c_proj.", ".fc2."), - (".attn", ".self_attn"), - ("ln_final.", "transformer.text_model.final_layer_norm."), - ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"), - ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"), -] -protected = {re.escape(x[0]): x[1] for x in textenc_transformer_conversion_lst} -textenc_pattern = re.compile("|".join(protected.keys())) - - def convert_paint_by_example_checkpoint(checkpoint, local_files_only=False): config = CLIPVisionConfig.from_pretrained("openai/clip-vit-large-patch14", local_files_only=local_files_only) model = PaintByExampleImageEncoder(config) - - keys = list(checkpoint.keys()) - - text_model_dict = {} - - for key in keys: - if key.startswith("cond_stage_model.transformer"): - text_model_dict[key[len("cond_stage_model.transformer.") :]] = checkpoint[key] - - # load clip vision - model.model.load_state_dict(text_model_dict) - - # load mapper - keys_mapper = { - k[len("cond_stage_model.mapper.res") :]: v - for k, v in checkpoint.items() - if k.startswith("cond_stage_model.mapper") + state = { + key: value + for key, value in checkpoint.items() + if key.startswith(("cond_stage_model.", "proj_out.")) or key == "learnable_vector" } - - MAPPING = { - "attn.c_qkv": ["attn1.to_q", "attn1.to_k", "attn1.to_v"], - "attn.c_proj": ["attn1.to_out.0"], - "ln_1": ["norm1"], - "ln_2": ["norm3"], - "mlp.c_fc": ["ff.net.0.proj"], - "mlp.c_proj": ["ff.net.2"], - } - - mapped_weights = {} - for key, value in keys_mapper.items(): - prefix = key[: len("blocks.i")] - suffix = key.split(prefix)[-1].split(".")[-1] - name = key.split(prefix)[-1].split(suffix)[0][1:-1] - mapped_names = MAPPING[name] - - num_splits = len(mapped_names) - for i, mapped_name in enumerate(mapped_names): - new_name = ".".join([prefix, mapped_name, suffix]) - shape = value.shape[0] // num_splits - mapped_weights[new_name] = value[i * shape : (i + 1) * shape] - - model.mapper.load_state_dict(mapped_weights) - - # load final layer norm - model.final_layer_norm.load_state_dict( - { - "bias": checkpoint["cond_stage_model.final_ln.bias"], - "weight": checkpoint["cond_stage_model.final_ln.weight"], - } - ) - - # load final proj - model.proj_out.load_state_dict( - { - "bias": checkpoint["proj_out.bias"], - "weight": checkpoint["proj_out.weight"], - } + model.load_state_dict( + get_conversion("PaintByExampleImageEncoder", config.to_dict()).to_diffusers(state), strict=True ) - - # load uncond vector - model.uncond_vector.data = torch.nn.Parameter(checkpoint["learnable_vector"]) return model @@ -951,62 +329,10 @@ def convert_open_clip_checkpoint( with ctx(): text_model = CLIPTextModelWithProjection(config) if has_projection else CLIPTextModel(config) - keys = list(checkpoint.keys()) - - keys_to_ignore = [] - if config_name == "stabilityai/stable-diffusion-2" and config.num_hidden_layers == 23: - # make sure to remove all keys > 22 - keys_to_ignore += [k for k in keys if k.startswith("cond_stage_model.model.transformer.resblocks.23")] - keys_to_ignore += ["cond_stage_model.model.text_projection"] - - text_model_dict = {} - - if prefix + "text_projection" in checkpoint: - d_model = int(checkpoint[prefix + "text_projection"].shape[0]) - else: - d_model = 1024 - - text_model_dict["text_model.embeddings.position_ids"] = text_model.text_model.embeddings.get_buffer("position_ids") - - for key in keys: - if key in keys_to_ignore: - continue - if key[len(prefix) :] in textenc_conversion_map: - if key.endswith("text_projection"): - value = checkpoint[key].T.contiguous() - else: - value = checkpoint[key] - - text_model_dict[textenc_conversion_map[key[len(prefix) :]]] = value - - if key.startswith(prefix + "transformer."): - new_key = key[len(prefix + "transformer.") :] - if new_key.endswith(".in_proj_weight"): - new_key = new_key[: -len(".in_proj_weight")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.weight"] = checkpoint[key][:d_model, :] - text_model_dict[new_key + ".k_proj.weight"] = checkpoint[key][d_model : d_model * 2, :] - text_model_dict[new_key + ".v_proj.weight"] = checkpoint[key][d_model * 2 :, :] - elif new_key.endswith(".in_proj_bias"): - new_key = new_key[: -len(".in_proj_bias")] - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - text_model_dict[new_key + ".q_proj.bias"] = checkpoint[key][:d_model] - text_model_dict[new_key + ".k_proj.bias"] = checkpoint[key][d_model : d_model * 2] - text_model_dict[new_key + ".v_proj.bias"] = checkpoint[key][d_model * 2 :] - else: - new_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], new_key) - - text_model_dict[new_key] = checkpoint[key] - - if is_accelerate_available(): - for param_name, param in text_model_dict.items(): - set_module_tensor_to_device(text_model, param_name, "cpu", value=param) - else: - if not (hasattr(text_model, "embeddings") and hasattr(text_model.embeddings.position_ids)): - text_model_dict.pop("text_model.embeddings.position_ids", None) - - text_model.load_state_dict(text_model_dict) + from ...loaders.single_file_utils import convert_open_clip_checkpoint as convert_text_component + state = convert_text_component(text_model, checkpoint, prefix=prefix) + text_model.load_state_dict(state, strict=True, assign=True) return text_model diff --git a/tests/others/test_checkpoint_conversion.py b/tests/others/test_checkpoint_conversion.py new file mode 100644 index 000000000000..590fa6d24960 --- /dev/null +++ b/tests/others/test_checkpoint_conversion.py @@ -0,0 +1,245 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.loaders.conversion import Conversion, Identity, ReorderChunks, Rule, Split + + +class TestCheckpointConversion: + def test_zero123_projection_layout(self): + from diffusers.loaders.conversion import get_conversion + + state = {"projection.weight": torch.randn(8, 16), "projection.bias": torch.randn(8)} + conversion = get_conversion("CCProjection", {}) + original = conversion.to_original(state) + assert set(original) == {"cc_projection.weight", "cc_projection.bias"} + restored = conversion.to_diffusers(original) + for key in state: + torch.testing.assert_close(restored[key], state[key], rtol=0, atol=0) + + @pytest.mark.parametrize("key", ["", 2, None]) + def test_invalid_key_name(self, key): + with pytest.raises(ValueError, match="nonempty strings"): + Conversion(mapping={"original": key}) + + def test_rename_and_key_swap_do_not_mutate_inputs(self): + mapping = {"a": "b", "b": "a"} + conversion = Conversion(mapping) + mapping["a"] = "changed" + original = {"a": torch.tensor([1, 2]), "b": torch.tensor([3, 4])} + converted = conversion.to_diffusers(original) + assert converted["b"] is original["a"] + assert converted["a"] is original["b"] + assert list(original) == ["a", "b"] + assert original["a"].tolist() == [1, 2] + restored = conversion.to_original(converted) + for key in original: + assert restored[key] is original[key] + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("dim", [0, -1]) + def test_unequal_split_in_both_directions(self, dtype, dim): + original = torch.arange(40, dtype=dtype).reshape(10, 4) + if dim == -1: + original = original.T + conversion = Conversion(rules=(Rule(("qkv",), ("q", "k", "v"), Split((6, 2, 2), dim)),)) + converted = conversion.to_diffusers({"qkv": original}) + for name, expected in zip(("q", "k", "v"), original.split((6, 2, 2), dim=dim)): + torch.testing.assert_close(converted[name], expected, rtol=0, atol=0) + restored = conversion.to_original(converted) + torch.testing.assert_close(restored["qkv"], original, rtol=0, atol=0) + # Fresh destination weights can be exported without an earlier import. + fresh = {"q": converted["q"] + 1, "k": converted["k"] + 2, "v": converted["v"] + 3} + reimported = conversion.to_diffusers(conversion.to_original(fresh)) + for name in fresh: + torch.testing.assert_close(reimported[name], fresh[name], rtol=0, atol=0) + + def test_non_involutive_permutation(self): + original = torch.arange(10).reshape(5, 2) + transform = ReorderChunks((0, 3, 1, 4, 2)) + converted = transform.forward((original,)) + assert converted[0].tolist() == [[0, 1], [6, 7], [2, 3], [8, 9], [4, 5]] + torch.testing.assert_close(transform.inverse(converted)[0], original, rtol=0, atol=0) + assert not torch.equal(transform.forward(converted)[0], original) + + def test_custom_many_to_many_transform(self): + class SwapAndTranspose: + def forward(self, tensors): + first, second = tensors + return second.T, first.T + + def inverse(self, tensors): + second, first = tensors + return first.T, second.T + + conversion = Conversion(rules=(Rule(("a", "b"), ("c", "d"), SwapAndTranspose()),)) + original = {"a": torch.arange(6).reshape(2, 3), "b": torch.arange(12).reshape(4, 3)} + converted = conversion.to_diffusers(original) + torch.testing.assert_close(converted["c"], original["b"].T) + torch.testing.assert_close(converted["d"], original["a"].T) + for key, value in conversion.to_original(converted).items(): + torch.testing.assert_close(value, original[key]) + + @pytest.mark.parametrize( + "mapping,rules", + [ + ({"a": "c", "b": "c"}, ()), + ({"a": "b"}, (Rule(("a",), ("c",)),)), + ({"a": "b"}, (Rule(("c",), ("b",)),)), + ], + ) + def test_collisions_are_rejected_before_conversion(self, mapping, rules): + with pytest.raises(ValueError, match="keys must be unique"): + Conversion(mapping, rules) + + @pytest.mark.parametrize("keys", [("a", "a"), (), "a"]) + def test_invalid_rule_keys(self, keys): + with pytest.raises(ValueError): + Rule(keys, ("b",)) + + def test_missing_group_member_and_unknown_key(self): + conversion = Conversion(rules=(Rule(("qkv",), ("q", "k", "v"), Split((2, 2, 2))),)) + with pytest.raises(ValueError, match=r"Diffusers -> original: missing keys \['v'\]"): + conversion.to_original({"q": torch.zeros(2), "k": torch.zeros(2)}) + with pytest.raises(ValueError, match=r"unexpected keys \['typo'\]"): + conversion.to_diffusers({"qkv": torch.zeros(6), "typo": torch.zeros(1)}) + + def test_shape_error_identifies_rule_and_direction(self): + conversion = Conversion(rules=(Rule(("qkv",), ("q", "k", "v"), Split((2, 2, 2))),)) + with pytest.raises(ValueError, match=r"original -> Diffusers.*qkv.*Split sizes"): + conversion.to_diffusers({"qkv": torch.zeros(7)}) + with pytest.raises(ValueError, match=r"Diffusers -> original.*Expected split piece shape"): + conversion.to_original({"q": torch.zeros(2, 4), "k": torch.zeros(2, 3), "v": torch.zeros(2, 4)}) + + def test_wrong_number_of_transform_outputs(self): + conversion = Conversion(rules=(Rule(("a",), ("b", "c"), Identity()),)) + with pytest.raises(ValueError, match="Expected 2 output tensors, got 1"): + conversion.to_diffusers({"a": torch.ones(1)}) + + @pytest.mark.parametrize("second", [torch.ones(2, dtype=torch.float16), torch.empty(2, device="meta")]) + def test_concat_rejects_implicit_dtype_or_device_changes(self, second): + with pytest.raises(ValueError, match="same dtype and device"): + Split((2, 2)).inverse((torch.ones(2), second)) + + @pytest.mark.parametrize("sizes", [(), (0, 2), (-1, 2)]) + def test_invalid_split_sizes(self, sizes): + with pytest.raises(ValueError, match="positive"): + Split(sizes) + + @pytest.mark.parametrize("order", [(), (0, 0), (1, 2)]) + def test_invalid_permutation(self, order): + with pytest.raises(ValueError, match="permutation"): + ReorderChunks(order) + + @pytest.mark.parametrize("shape,dim", [((), 0), ((2, 3), 2), ((2, 3), -3)]) + def test_invalid_dimension(self, shape, dim): + with pytest.raises(ValueError, match="dimension"): + Split((1, 1), dim).forward((torch.zeros(shape),)) + with pytest.raises(ValueError, match="dimension"): + ReorderChunks((1, 0), dim).forward((torch.zeros(shape),)) + + def test_reorder_requires_equal_nonempty_chunks(self): + for size in (0, 5): + with pytest.raises(ValueError, match="Cannot divide"): + ReorderChunks((1, 0)).forward((torch.zeros(size),)) + + def test_meta_tensor_conversion(self): + conversion = Conversion(rules=(Rule(("qkv",), ("q", "k", "v"), Split((6, 2, 2))),)) + original = {"qkv": torch.empty(10, 4, device="meta")} + converted = conversion.to_diffusers(original) + assert [tensor.shape for tensor in converted.values()] == [(6, 4), (2, 4), (2, 4)] + assert conversion.to_original(converted)["qkv"].shape == (10, 4) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("bias", [False, True]) +def test_linear_gate_has_a_canonical_inverse(dtype, bias): + from diffusers.loaders.conversion import FoldLinearGate + + operation = FoldLinearGate(bias=bias) + gate = torch.tensor([0.0, -0.25, 1.75], dtype=dtype) + weight = torch.randn(3, 5, dtype=dtype) + original = (gate, weight, torch.randn(3, dtype=dtype)) if bias else (gate, weight) + copies = tuple(t.clone() for t in original) + folded = operation.forward(original) + expected = (gate.float()[:, None] * weight.float()).to(dtype) + torch.testing.assert_close(folded[0], expected, rtol=0, atol=0) + canonical = operation.inverse(folded) + torch.testing.assert_close(canonical[0], torch.ones_like(gate), rtol=0, atol=0) + for actual, wanted in zip(operation.forward(canonical), folded): + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + for actual, wanted in zip(original, copies): + torch.testing.assert_close(actual, wanted, rtol=0, atol=0) + original_keys = ("gate", "weight", "bias") if bias else ("gate", "weight") + converted_keys = ("weight", "bias") if bias else ("weight",) + conversion = Conversion(rules=(Rule(original_keys, converted_keys, operation),)) + assert not conversion.lossless + + +def test_composition_reports_normalization_and_freezes_steps(): + from diffusers.loaders.conversion import Chain, FoldLinearGate, Reverse + + steps = [FoldLinearGate(bias=False)] + operation = Chain(steps) + steps.clear() + assert len(operation.transforms) == 1 + assert not operation.lossless + assert not Reverse(operation).lossless + assert Conversion(mapping={"a": "b"}).lossless + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_lora_layout_preserves_scaling_and_factorization(dtype): + from diffusers.loaders.conversion import get_conversion + + module = "unet.down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q" + config = {"modules": [module], "original_format": "kohya", "include_alpha": True, "use_dora": True} + stem = "lora_unet_down_blocks_0_attentions_0_transformer_blocks_0_attn1_to_q" + original = { + stem + ".lora_down.weight": torch.randn(2, 8, dtype=dtype), + stem + ".lora_up.weight": torch.randn(8, 2, dtype=dtype), + stem + ".alpha": torch.tensor(0.5, dtype=torch.float32), + stem + ".dora_scale": torch.randn(8, dtype=dtype), + } + conversion = get_conversion("LoRA", config) + converted = conversion.to_diffusers(original) + assert converted[module + ".alpha"].dtype == torch.float32 + assert torch.equal(converted[module + ".lora_A.weight"], original[stem + ".lora_down.weight"]) + for key, tensor in conversion.to_original(converted).items(): + torch.testing.assert_close(tensor, original[key], rtol=0, atol=0) + + +def test_lora_underscore_collision_is_rejected(): + from diffusers.loaders.conversion import get_conversion + + with pytest.raises(ValueError, match="collision"): + get_conversion("LoRA", {"modules": ["unet.a_b.c", "unet.a.b_c"]}) + + +def test_flux_ip_adapter_keeps_distinct_projection_and_value_weights(): + from diffusers.loaders.conversion import get_conversion + + conversion = get_conversion("FluxIPAdapter", {"num_layers": 2}) + original = {key: torch.tensor([i], dtype=torch.float32) for i, key in enumerate(sorted(conversion.original_keys))} + converted = conversion.to_diffusers(original) + assert torch.equal(converted["image_proj.proj.weight"], original["ip_adapter_proj_model.proj.weight"]) + assert torch.equal( + converted["ip_adapter.1.to_v_ip.weight"], + original["double_blocks.1.processor.ip_adapter_double_stream_v_proj.weight"], + ) + assert not torch.equal(converted["ip_adapter.1.to_v_ip.weight"], converted["ip_adapter.1.to_k_ip.weight"]) + for key, tensor in conversion.to_original(converted).items(): + torch.testing.assert_close(tensor, original[key], rtol=0, atol=0) diff --git a/tests/others/test_conversion_io.py b/tests/others/test_conversion_io.py new file mode 100644 index 000000000000..608aa2e06a12 --- /dev/null +++ b/tests/others/test_conversion_io.py @@ -0,0 +1,418 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib +import json + +import pytest +import torch +from safetensors.torch import save_file + +from diffusers import FluxTransformer2DModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.checkpoint import convert_component_checkpoint +from diffusers.loaders.conversion.io import Checkpoint, convert_checkpoint + + +def checkpoint_digest(state): + """Hash names, shapes, dtypes and raw tensor bytes, independently of checkpoint serialization.""" + digest = hashlib.sha256() + for key in sorted(state): + value = state[key].detach().cpu().contiguous() + digest.update(json.dumps([key, list(value.shape), str(value.dtype)]).encode()) + digest.update(value.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def test_svd_vae_selects_qualified_namespaces_from_bundle(tmp_path): + from tests.models.autoencoders.test_models_autoencoder_kl_temporal_decoder import ( + AutoencoderKLTemporalDecoderTesterConfig, + ) + + tester = AutoencoderKLTemporalDecoderTesterConfig() + model = tester.model_class(**tester.get_init_dict()) + config = dict(model.config) + original = get_conversion(type(model).__name__, config).to_original(model.state_dict()) + bundle = { + **original, + "model.diffusion_model.input_blocks.0.0.weight": torch.zeros(1), + "conditioner.embedders.0.model.visual.conv1.weight": torch.zeros(1), + } + loaded = type(model).from_single_file(bundle, config=config, local_files_only=True) + assert checkpoint_digest(loaded.state_dict()) == checkpoint_digest(model.state_dict()) + source = tmp_path / "svd.safetensors" + save_file(bundle, source) + output = convert_checkpoint(source, tmp_path / "vae", config=config, model_class=type(model).__name__) + assert checkpoint_digest(Checkpoint(output)) == checkpoint_digest(model.state_dict()) + bundle["first_stage_model.decoder.typo.weight"] = torch.zeros(1) + with pytest.raises(ValueError, match="unexpected keys.*typo"): + type(model).from_single_file(bundle, config=config, local_files_only=True) + + +@pytest.mark.parametrize("component", ["unet_text", "unet_image"]) +def test_versatile_components_keep_qualified_namespaces(component): + from diffusers import UNet2DConditionModel + from diffusers.pipelines.deprecated.versatile_diffusion.modeling_text_unet import UNetFlatConditionModel + from tests.single_file.test_auxiliary_conversions import CASES + + cls = UNetFlatConditionModel if component == "unet_text" else UNet2DConditionModel + if component == "unet_image": + from tests.models.unets.test_models_unet_2d_condition import UNet2DConditionTesterConfig + + config = UNet2DConditionTesterConfig().get_init_dict() + config["original_format"] = "versatile_image" + else: + config = dict(next(config for name, config in CASES if name == cls.__name__)) + model = cls(**{key: value for key, value in config.items() if key != "original_format"}) + resolved_config = {**dict(model.config), **config} + original = get_conversion(cls.__name__, resolved_config).to_original(model.state_dict()) + other = "unet_image" if component == "unet_text" else "unet_text" + bundle = {**original, f"model.diffusion_model.{other}.unrelated.weight": torch.ones(1)} + loaded = cls.from_single_file(bundle, config=resolved_config) + assert checkpoint_digest(loaded.state_dict()) == checkpoint_digest(model.state_dict()) + bundle[f"model.diffusion_model.{component}.typo.weight"] = torch.ones(1) + with pytest.raises(ValueError, match="unexpected keys.*typo"): + cls.from_single_file(bundle, config=resolved_config) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_hunyuan_image_vae_original_temporal_axis(dtype): + from diffusers import AutoencoderKLHunyuanImage + + model = AutoencoderKLHunyuanImage( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(32, 64), + layers_per_block=1, + spatial_compression_ratio=2, + sample_size=16, + ).to(dtype) + config = dict(model.config) + spatial = get_conversion(type(model).__name__, {**config, "original_format": "hunyuan_image_vae_2d"}) + # Construct the original 5D source independently of the temporal-axis transform's inverse. + original = { + key: value.unsqueeze(2) if key.endswith(".weight") and value.ndim == 4 else value + for key, value in spatial.to_original(model.state_dict()).items() + } + conversion = get_conversion(type(model).__name__, config) + restored = conversion.to_diffusers(original) + assert checkpoint_digest(restored) == checkpoint_digest(model.state_dict()) + assert checkpoint_digest(conversion.to_original(restored)) == checkpoint_digest(original) + assert conversion.to_original(restored)["encoder.conv_in.weight"].shape == (32, 3, 1, 3, 3) + loaded = type(model).from_single_file( + {"vae." + k: v for k, v in original.items()}, config=config, torch_dtype=dtype + ) + assert checkpoint_digest(loaded.state_dict()) == checkpoint_digest(model.state_dict()) + invalid = {**original, "encoder.conv_in.weight": original["encoder.conv_in.weight"].expand(-1, -1, 2, -1, -1)} + with pytest.raises(ValueError, match="singleton axis"): + conversion.to_diffusers(invalid) + + +@pytest.mark.parametrize("family", ["cosmos", "ltx2"]) +def test_inferred_original_format_survives_save_and_export(tmp_path, family): + if family == "cosmos": + from tests.models.transformers.test_models_transformer_cosmos import CosmosTransformerTesterConfig + + tester = CosmosTransformerTesterConfig() + model = tester.model_class(**tester.get_init_dict()) + original_format = "cosmos1" + else: + from diffusers import LTX2VideoDiffusionDecoderModel + + model = LTX2VideoDiffusionDecoderModel( + latent_channels=8, + decoder_head_dim=8, + decoder_stage_channels=[32, 16, 8, 8, 8], + decoder_stage_depths=[1, 1, 1, 1, 1], + decoder_upsample_channel_reductions=[2, 2, 1, 1], + decoder_t_emb_dim=16, + ) + original_format = "ltx2_diffusion_decoder_gated" + config = dict(model.config) + conversion = get_conversion(type(model).__name__, {**config, "original_format": original_format}) + original = conversion.to_original(model.state_dict()) + source = tmp_path / "original.safetensors" + save_file(original, source) + output = convert_checkpoint(source, tmp_path / "model", config=config, model_class=type(model).__name__) + saved_config = json.loads((output / "config.json").read_text()) + assert saved_config["original_format"] == original_format + exported = convert_checkpoint(output, tmp_path / "exported", config=saved_config, reverse=True) + assert checkpoint_digest(Checkpoint(exported)) == checkpoint_digest(original) + loaded = type(model).from_single_file(original, config=config) + loaded.save_config(tmp_path / "loader") + assert json.loads((tmp_path / "loader/config.json").read_text())["original_format"] == original_format + assert "original_format" not in config + + +def test_pytorch_auxiliary_metadata_matches_mapping_import(tmp_path): + from tests.models.transformers.test_models_transformer_cosmos import CosmosTransformerTesterConfig + + tester = CosmosTransformerTesterConfig() + model = tester.model_class(**tester.get_init_dict()) + config = dict(model.config) + original = get_conversion(type(model).__name__, config).to_original(model.state_dict()) + source = {**original, "_extra_state": {"train_info": 1}} + path = tmp_path / "original.pt" + torch.save({"state_dict": source}, path) + output = convert_checkpoint( + path, tmp_path / "model", config=config, model_class=type(model).__name__, input_wrapper=["state_dict"] + ) + assert checkpoint_digest(Checkpoint(output)) == checkpoint_digest( + convert_component_checkpoint(source, config, type(model).__name__) + ) + source["unknown_metadata"] = {"train_info": 1} + torch.save({"state_dict": source}, path) + with pytest.raises(ValueError, match="unexpected keys.*unknown_metadata"): + convert_checkpoint( + path, tmp_path / "invalid", config=config, model_class=type(model).__name__, input_wrapper=["state_dict"] + ) + del source["unknown_metadata"] + source["x_embedder.proj.1.weight"] = {"not": "a tensor"} + torch.save({"state_dict": source}, path) + with pytest.raises(ValueError, match="must be a tensor"): + convert_checkpoint( + path, tmp_path / "invalid", config=config, model_class=type(model).__name__, input_wrapper=["state_dict"] + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_sharded_file_round_trip(tmp_path, dtype): + model = FluxTransformer2DModel( + num_layers=1, + num_single_layers=1, + num_attention_heads=2, + attention_head_dim=8, + in_channels=4, + joint_attention_dim=8, + pooled_projection_dim=8, + axes_dims_rope=(2, 2, 4), + ).to(dtype=dtype) + source = tmp_path / "source" + model.save_pretrained(source, max_shard_size="10KB") + original = convert_checkpoint( + source, + tmp_path / "original", + config=dict(model.config), + reverse=True, + output_prefix="model.diffusion_model.", + max_shard_size=3000, + ) + restored = convert_checkpoint( + original, + tmp_path / "restored", + config=dict(model.config), + input_prefix="model.diffusion_model.", + max_shard_size=2000, + ) + loaded = FluxTransformer2DModel.from_pretrained(restored, torch_dtype=dtype, local_files_only=True) + for key, value in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], value, rtol=0, atol=0) + assert json.loads((original / "conversion.json").read_text())["lossless"] + assert (original / "conversion_config.json").is_file() + + +def test_explicit_pytorch_wrapper_and_component(tmp_path): + path = tmp_path / "weights.pt" + torch.save({"training_step": 10, "state_dict": {"component.x": torch.arange(3), "other.y": torch.ones(1)}}, path) + with pytest.raises(ValueError, match="tensor state dict"): + Checkpoint(path) + state = Checkpoint(path, wrapper=("state_dict",), prefix="component.") + assert set(state) == {"x"} + torch.testing.assert_close(state["x"], torch.arange(3)) + + +def test_tied_text_embeddings_survive_file_conversion(tmp_path): + from transformers import T5Config, T5EncoderModel + + model = T5EncoderModel(T5Config(vocab_size=16, d_model=8, d_ff=16, d_kv=4, num_heads=2, num_layers=1)) + source = tmp_path / "source" + model.save_pretrained(source) + config = model.config.to_dict() + original = convert_checkpoint( + source, tmp_path / "original", config=config, model_class="T5EncoderModel", reverse=True + ) + restored = convert_checkpoint(original, tmp_path / "restored", config=config, model_class="T5EncoderModel") + loaded = T5EncoderModel.from_pretrained(restored, local_files_only=True) + for key, value in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], value, rtol=0, atol=0) + + +@pytest.mark.parametrize("failure", ["outside", "wrong_shard", "missing_key", "extra_key", "duplicate"]) +def test_invalid_shard_index_is_rejected(tmp_path, failure): + save_file({"a": torch.ones(2)}, tmp_path / "one.safetensors") + save_file({"b": torch.zeros(2)}, tmp_path / "two.safetensors") + weight_map = {"a": "one.safetensors", "b": "two.safetensors"} + if failure == "outside": + weight_map["b"] = "../outside.safetensors" + elif failure == "wrong_shard": + weight_map = {"a": "two.safetensors", "b": "one.safetensors"} + elif failure == "missing_key": + weight_map["c"] = "one.safetensors" + elif failure == "extra_key": + save_file({"a": torch.ones(2), "c": torch.ones(2)}, tmp_path / "one.safetensors") + else: + save_file({"a": torch.ones(2), "b": torch.ones(2)}, tmp_path / "two.safetensors") + index = tmp_path / "model.safetensors.index.json" + index.write_text(json.dumps({"weight_map": weight_map})) + with pytest.raises(ValueError, match="escapes|does not match|Duplicate"): + Checkpoint(index) + + +def test_failed_conversion_does_not_publish_partial_output(tmp_path): + source = tmp_path / "input.safetensors" + save_file({"projection.weight": torch.ones(2)}, source) + output = tmp_path / "output" + with pytest.raises(ValueError, match="missing keys"): + convert_checkpoint(source, output, model_class="CCProjection", config={}, reverse=True) + assert not output.exists() + assert not list(tmp_path.glob(".conversion-*")) + output.mkdir() + with pytest.raises(FileExistsError): + convert_checkpoint(source, output, model_class="CCProjection", config={}, reverse=True) + + +def test_pytorch_output_wrappers_round_trip(tmp_path): + source = tmp_path / "weights.safetensors" + state = {"projection.weight": torch.randn(8, 16), "projection.bias": torch.randn(8)} + save_file(state, source) + output = convert_checkpoint( + source, + tmp_path / "original.pt", + config={}, + model_class="CCProjection", + reverse=True, + output_format="pytorch", + output_wrapper=("model", "state_dict"), + output_prefix="component.", + ) + restored = convert_checkpoint( + output, + tmp_path / "restored", + config={}, + model_class="CCProjection", + input_wrapper=("model", "state_dict"), + input_prefix="component.", + ) + for key, value in Checkpoint(restored).items(): + torch.testing.assert_close(value, state[key], rtol=0, atol=0) + + +def test_single_file_component_selection_is_strict_and_nonmutating(): + model = FluxTransformer2DModel( + num_layers=1, + num_single_layers=1, + num_attention_heads=2, + attention_head_dim=8, + in_channels=4, + joint_attention_dim=8, + pooled_projection_dim=8, + axes_dims_rope=(2, 2, 4), + ) + conversion = get_conversion(type(model).__name__, dict(model.config)) + original = conversion.to_original(model.state_dict()) + torch.testing.assert_close( + original["double_blocks.0.img_attn.qkv.weight"], + torch.cat([model.state_dict()[f"transformer_blocks.0.attn.to_{part}.weight"] for part in ("q", "k", "v")]), + rtol=0, + atol=0, + ) + checkpoint = {"model.diffusion_model." + key: value for key, value in original.items()} + checkpoint["first_stage_model.some_weight"] = torch.ones(1) + keys = set(checkpoint) + loaded = convert_component_checkpoint(checkpoint, dict(model.config), type(model).__name__) + assert set(checkpoint) == keys + for key, value in model.state_dict().items(): + torch.testing.assert_close(loaded[key], value, rtol=0, atol=0) + checkpoint["model.diffusion_model.typo"] = torch.ones(1) + with pytest.raises(ValueError, match="unexpected keys.*typo"): + convert_component_checkpoint(checkpoint, dict(model.config), type(model).__name__) + + +def test_asymmetric_loader_excludes_training_aliases(): + from diffusers import AsymmetricAutoencoderKL + + model = AsymmetricAutoencoderKL( + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + down_block_out_channels=(32,), + up_block_out_channels=(32,), + layers_per_down_block=1, + layers_per_up_block=1, + ) + config = dict(model.config) + original = get_conversion("AsymmetricAutoencoderKL", config).to_original(model.state_dict()) + original["decoder.up_layers.0.weight"] = torch.ones(1) + original["loss.discriminator.weight"] = torch.ones(1) + for restored in ( + AsymmetricAutoencoderKL.from_single_file(original, config=config, local_files_only=True).state_dict(), + convert_component_checkpoint(original, config, "AsymmetricAutoencoderKL"), + ): + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(restored[key], tensor, rtol=0, atol=0) + + +def test_single_file_accepts_config_mapping_without_mutating_it(): + model = FluxTransformer2DModel( + num_layers=1, + num_single_layers=1, + num_attention_heads=2, + attention_head_dim=8, + in_channels=4, + joint_attention_dim=8, + pooled_projection_dim=8, + axes_dims_rope=(2, 2, 4), + ) + config = dict(model.config) + before = dict(config) + original = get_conversion(type(model).__name__, config).to_original(model.state_dict()) + loaded = type(model).from_single_file(original, config=config, local_files_only=True) + assert config == before + for key, value in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], value, rtol=0, atol=0) + + +def test_ltx2_decoder_partial_gates_and_auxiliary_state(tmp_path): + from diffusers import LTX2VideoDiffusionDecoderModel + + model = LTX2VideoDiffusionDecoderModel( + latent_channels=8, + decoder_head_dim=8, + decoder_stage_channels=[32, 16, 8, 8, 8], + decoder_stage_depths=[1, 1, 1, 1, 1], + decoder_upsample_channel_reductions=[2, 2, 1, 1], + decoder_t_emb_dim=16, + ) + config = dict(model.config) + conversion = get_conversion(type(model).__name__, {**config, "original_format": "ltx2_diffusion_decoder_gated"}) + source = conversion.to_original(model.state_dict()) + gates = sorted(key for key in source if key.endswith((".gate_msa", ".gate_mlp", ".gate_ctx"))) + for key in gates[1:]: + del source[key] + source[gates[0]] = torch.full_like(source[gates[0]], 0.5) + source["encoder.unused.weight"] = torch.ones(1) + source["decoder.coarse_head.weight"] = torch.ones(1) + source["decoder.diff_blocks.0.coarse_proj.weight"] = torch.ones(1) + expected = dict(model.state_dict()) + rule = next(rule for rule in conversion.rules if rule.original[0] == gates[0]) + for key in rule.diffusers: + expected[key] = (expected[key].float() * 0.5).to(expected[key].dtype) + original_keys = set(source) + path = tmp_path / "original.safetensors" + save_file(source, path) + output = convert_checkpoint(path, tmp_path / "converted", config=config, model_class=type(model).__name__) + for actual in (convert_component_checkpoint(source, config, type(model).__name__), Checkpoint(output)): + for key, value in expected.items(): + torch.testing.assert_close(actual[key], value, rtol=0, atol=0) + assert set(source) == original_keys diff --git a/tests/others/test_pipeline_conversion.py b/tests/others/test_pipeline_conversion.py new file mode 100644 index 000000000000..86d5bf7666c9 --- /dev/null +++ b/tests/others/test_pipeline_conversion.py @@ -0,0 +1,199 @@ +import importlib +import json + +import pytest +import torch +from safetensors.torch import save_file +from transformers import CLIPConfig, CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPVisionConfig + +from diffusers import AutoencoderKL, AutoencoderTiny, UNet2DConditionModel +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.configs import get_config_preset, list_config_presets +from diffusers.loaders.conversion.io import Checkpoint, convert_checkpoint +from diffusers.loaders.conversion.pipeline import export_pipeline_checkpoint +from diffusers.loaders.conversion.source import MergedCheckpoint, load_source_manifest +from diffusers.pipelines.deepfloyd_if.safety_checker import IFSafetyChecker +from diffusers.pipelines.deprecated.paint_by_example.image_encoder import PaintByExampleImageEncoder +from diffusers.pipelines.deprecated.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings +from diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion import LDMBertConfig, LDMBertModel + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("kind", ["bert", "paint", "if", "learned", "empty"]) +def test_legacy_pipeline_component_conversion(kind, dtype): + if kind == "bert": + config = LDMBertConfig( + vocab_size=32, + max_position_embeddings=16, + d_model=16, + encoder_layers=2, + encoder_ffn_dim=32, + encoder_attention_heads=2, + head_dim=8, + ) + model = LDMBertModel(config) + elif kind == "paint": + config = CLIPVisionConfig( + hidden_size=16, + intermediate_size=32, + num_hidden_layers=4, + num_attention_heads=2, + projection_dim=8, + image_size=16, + patch_size=8, + ) + model = PaintByExampleImageEncoder(config) + elif kind == "if": + config = CLIPConfig( + vision_config={ + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "projection_dim": 8, + "image_size": 16, + "patch_size": 8, + } + ) + model = IFSafetyChecker(config) + else: + config = {"learnable": kind == "learned", "hidden_size": 8, "length": 4} + model = LearnedClassifierFreeSamplingEmbeddings(**config) + model = model.to(dtype) + config = config.to_dict() if hasattr(config, "to_dict") else config + conversion = get_conversion(type(model).__name__, config) + state = model.state_dict() + assert conversion.diffusers_keys == state.keys() + original = conversion.to_original(state) + restored = conversion.to_diffusers(original) + for key, value in state.items(): + torch.testing.assert_close(restored[key], value, rtol=0, atol=0) + + +def test_tiny_vae_source_manifest_and_duplicates(tmp_path): + model = AutoencoderTiny( + encoder_block_out_channels=(8, 8), + decoder_block_out_channels=(8, 8), + num_encoder_blocks=(1, 1), + num_decoder_blocks=(1, 1), + ) + config = dict(model.config) + original = get_conversion("AutoencoderTiny", config).to_original(model.state_dict()) + sources = [] + for part in ("encoder", "decoder"): + state = {key.removeprefix(part + "."): value for key, value in original.items() if key.startswith(part + ".")} + save_file(state, tmp_path / f"{part}.safetensors") + sources.append({"path": f"{part}.safetensors", "output_prefix": part + "."}) + path = tmp_path / "sources.json" + path.write_text(json.dumps({"sources": sources})) + merged = load_source_manifest(path) + output = convert_checkpoint(merged, tmp_path / "output", config=config, model_class="AutoencoderTiny") + for key, value in Checkpoint(output).items(): + torch.testing.assert_close(value, model.state_dict()[key], rtol=0, atol=0) + with pytest.raises(ValueError, match="Duplicate source tensor"): + MergedCheckpoint([("", merged), ("", merged)]) + + +@pytest.mark.parametrize("pipeline_format", ["sd", "sdxl"]) +@pytest.mark.parametrize("output_format", ["safetensors", "pytorch"]) +def test_pipeline_checkpoint_export(tmp_path, pipeline_format, output_format): + unet = UNet2DConditionModel( + block_out_channels=(32,), + down_block_types=("CrossAttnDownBlock2D",), + up_block_types=("CrossAttnUpBlock2D",), + layers_per_block=1, + cross_attention_dim=32, + ) + vae = AutoencoderKL( + block_out_channels=(32,), + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + layers_per_block=1, + ) + text_config = CLIPTextConfig( + vocab_size=32, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + projection_dim=8, + max_position_embeddings=8, + ) + text = CLIPTextModel(text_config) + components = {"unet": unet, "vae": vae, "text_encoder": text} + specs = { + "unet": ("model.diffusion_model.", None), + "vae": ("first_stage_model.", None), + "text_encoder": ( + "cond_stage_model.transformer." if pipeline_format == "sd" else "conditioner.embedders.0.transformer.", + "clip", + ), + } + if pipeline_format == "sdxl": + components["text_encoder_2"] = CLIPTextModelWithProjection(text_config) + specs["text_encoder_2"] = ("conditioner.embedders.1.model.", "openclip") + for name, model in components.items(): + model.save_pretrained(tmp_path / name) + output = export_pipeline_checkpoint( + tmp_path, + tmp_path / ("original.safetensors" if output_format == "safetensors" else "original.pt"), + pipeline_format=pipeline_format, + output_format=output_format, + ) + for name, model in components.items(): + prefix, fmt = specs[name] + config = model.config.to_dict() if hasattr(model.config, "to_dict") else dict(model.config) + if fmt: + config["original_format"] = fmt + state = Checkpoint(output, prefix=prefix, wrapper=("state_dict",) if output_format == "pytorch" else ()) + restored = get_conversion(type(model).__name__, config).to_diffusers(state) + for key, value in model.state_dict().items(): + torch.testing.assert_close(restored[key], value, rtol=0, atol=0) + + +def test_configuration_helpers_are_importable_and_reusable(): + names = list_config_presets() + modules = {name.split(".", 1)[0] for name in names if "." in name and not name.startswith("asymmetric-")} + for name in modules: + importlib.import_module("diffusers.loaders.conversion.configs." + name) + tiny = get_config_preset("tiny-vae") + assert tiny == {"_class_name": "AutoencoderTiny"} + first = get_config_preset("asymmetric-vae-1.5") + first["up_block_out_channels"].append(1) + assert len(get_config_preset("asymmetric-vae-1.5")["up_block_out_channels"]) == 4 + config = get_config_preset("prx.build_config", arguments={"variant": "flux"}) + assert config["in_channels"] == 16 + config = get_config_preset( + "unidiffuser.create_unidiffuser_unet_config", arguments={"config_type": "test", "version": 1} + ) + assert config["use_data_type_embedding"] is True + config = get_config_preset("wan.get_transformer_config", arguments={"model_type": "Wan-T2V-1.3B"}) + assert config["num_attention_heads"] == 12 + config = get_config_preset("prx.create_scheduler_config", arguments={"shift": 3.0}) + assert config["shift"] == 3.0 + + +@pytest.mark.parametrize("source_format", ["torchscript", "python-model"]) +def test_module_source_manifest(tmp_path, source_format): + model = torch.nn.Linear(3, 2) + checkpoint = tmp_path / "model.pt" + wrapper = [] + if source_format == "torchscript": + torch.jit.script(model).save(str(checkpoint)) + else: + torch.save({"ema": model}, checkpoint) + wrapper = ["ema"] + manifest = tmp_path / "sources.json" + manifest.write_text( + json.dumps( + { + "sources": [ + {"path": "model.pt", "format": source_format, "wrapper": wrapper, "output_prefix": "decoder."} + ] + } + ) + ) + state = load_source_manifest(manifest) + assert set(state) == {"decoder.weight", "decoder.bias"} + for name, value in model.state_dict().items(): + torch.testing.assert_close(state["decoder." + name], value, rtol=0, atol=0) diff --git a/tests/single_file/test_auxiliary_conversions.py b/tests/single_file/test_auxiliary_conversions.py new file mode 100644 index 000000000000..04735436ef66 --- /dev/null +++ b/tests/single_file/test_auxiliary_conversions.py @@ -0,0 +1,1044 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import json + +import pytest +import torch + +import diffusers +from diffusers.loaders.conversion import get_conversion + + +CASES = [ + ( + "AudioLDM2ProjectionModel", + { + "text_encoder_dim": 8, + "text_encoder_1_dim": 16, + "langauge_model_dim": 32, + "use_learned_position_embedding": True, + "max_seq_length": 16, + }, + ), + ( + "AudioLDM2UNet2DConditionModel", + { + "block_out_channels": (32, 64), + "layers_per_block": 1, + "down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"), + "up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D"), + "cross_attention_dim": ((8, 16), (8, 16)), + }, + ), + ( + "AutoencoderKLHunyuanImage", + { + "in_channels": 3, + "out_channels": 3, + "latent_channels": 4, + "block_out_channels": (32, 64), + "layers_per_block": 1, + "spatial_compression_ratio": 2, + "sample_size": 16, + }, + ), + ( + "AutoencoderKLHunyuanImageRefiner", + { + "latent_channels": 4, + "block_out_channels": (32, 64), + "layers_per_block": 1, + "spatial_compression_ratio": 2, + "temporal_compression_ratio": 2, + }, + ), + ( + "HunyuanDiT2DControlNetModel", + { + "num_attention_heads": 2, + "attention_head_dim": 16, + "in_channels": 4, + "patch_size": 2, + "sample_size": 8, + "hidden_size": 32, + "transformer_num_layers": 4, + "cross_attention_dim": 16, + "cross_attention_dim_t5": 32, + "pooled_projection_dim": 16, + "text_len": 4, + "text_len_t5": 8, + }, + ), + ( + "Kandinsky3UNet", + { + "block_out_channels": (32, 64, 64, 64), + "time_embedding_dim": 128, + "layers_per_block": 1, + "attention_head_dim": 32, + "cross_attention_dim": 32, + "encoder_hid_dim": 32, + "groups": 8, + }, + ), + ( + "SD3ControlNetModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "num_layers": 2, + "in_channels": 4, + "out_channels": 4, + "joint_attention_dim": 16, + "caption_projection_dim": 16, + "pooled_projection_dim": 16, + "pos_embed_max_size": 8, + "qk_norm": "rms_norm", + "dual_attention_layers": (0,), + }, + ), + ( + "SanaControlNetModel", + { + "num_attention_heads": 2, + "attention_head_dim": 16, + "num_layers": 2, + "num_cross_attention_heads": 2, + "cross_attention_head_dim": 16, + "cross_attention_dim": 32, + "caption_channels": 32, + "in_channels": 4, + "out_channels": 4, + }, + ), + ( + "UNetFlatConditionModel", + { + "block_out_channels": (32, 64), + "layers_per_block": 1, + "down_block_types": ("CrossAttnDownBlockFlat", "DownBlockFlat"), + "up_block_types": ("UpBlockFlat", "CrossAttnUpBlockFlat"), + "cross_attention_dim": 16, + }, + ), + ( + "UnCLIPTextProjModel", + {"clip_extra_context_tokens": 2, "clip_embeddings_dim": 8, "time_embed_dim": 16, "cross_attention_dim": 16}, + ), + ( + "SpectrogramContEncoder", + { + "input_dims": 8, + "targets_context_length": 4, + "d_model": 16, + "dropout_rate": 0, + "num_layers": 1, + "num_heads": 2, + "d_kv": 8, + "d_ff": 32, + "feed_forward_proj": "gated-gelu", + }, + ), + ( + "SpectrogramNotesEncoder", + { + "max_length": 8, + "vocab_size": 16, + "d_model": 16, + "dropout_rate": 0, + "num_layers": 1, + "num_heads": 2, + "d_kv": 8, + "d_ff": 32, + "feed_forward_proj": "gated-gelu", + }, + ), + ( + "AutoencoderRAE", + { + "encoder_type": "siglip2", + "encoder_hidden_size": 64, + "encoder_num_hidden_layers": 1, + "encoder_patch_size": 16, + "encoder_input_size": 32, + "decoder_hidden_size": 32, + "decoder_num_hidden_layers": 1, + "decoder_num_attention_heads": 2, + "decoder_intermediate_size": 64, + }, + ), + ("UNet1DModel", {}), + ( + "ChronoEditTransformer3DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "num_layers": 1, + }, + ), + ( + "QwenImageTransformer2DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "joint_attention_dim": 16, + "num_layers": 1, + "axes_dims_rope": [2, 2, 4], + "use_additional_t_cond": True, + }, + ), + ( + "MotifVideoTransformer3DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_embed_dim": 16, + "num_layers": 1, + "num_single_layers": 2, + "num_decoder_layers": 1, + "rope_axes_dim": [2, 2, 4], + }, + ), + ( + "MotifVideoTransformer3DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_embed_dim": 16, + "image_embed_dim": 8, + "num_layers": 1, + "num_single_layers": 2, + "num_decoder_layers": 1, + "rope_axes_dim": [2, 2, 4], + "qk_norm": "layer_norm", + "enable_text_cross_attention_dual": True, + "enable_text_cross_attention_single": True, + }, + ), + ("LTXLatentUpsamplerModel", {"in_channels": 4, "mid_channels": 32, "num_blocks_per_stage": 1}), + ( + "SkyReelsV2Transformer3DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "num_layers": 1, + "inject_sample_info": True, + }, + ), + ( + "SkyReelsV2Transformer3DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_dim": 16, + "freq_dim": 16, + "ffn_dim": 32, + "num_layers": 1, + "image_dim": 16, + "added_kv_proj_dim": 16, + "pos_embed_seq_len": 4, + }, + ), + ( + "HunyuanImageTransformer2DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "num_layers": 1, + "num_single_layers": 1, + "num_refiner_layers": 1, + "in_channels": 4, + "out_channels": 4, + "text_embed_dim": 16, + "rope_axes_dim": [4, 4], + }, + ), + ( + "HunyuanImageTransformer2DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "num_layers": 1, + "num_single_layers": 1, + "num_refiner_layers": 1, + "in_channels": 4, + "out_channels": 4, + "text_embed_dim": 16, + "text_embed_2_dim": 8, + "rope_axes_dim": [4, 4], + "use_meanflow": True, + "guidance_embeds": True, + }, + ), + ( + "SparseControlNetModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["CrossAttnDownBlockMotion", "DownBlockMotion"], + "layers_per_block": 1, + "cross_attention_dim": 16, + "motion_num_attention_heads": 2, + }, + ), + ( + "SparseControlNetModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["CrossAttnDownBlockMotion", "DownBlockMotion"], + "layers_per_block": 1, + "cross_attention_dim": 16, + "motion_num_attention_heads": 2, + "use_simplified_condition_embedding": False, + "conditioning_embedding_out_channels": [8, 16], + "transformer_layers_per_mid_block": 2, + }, + ), + ( + "CogView4Transformer2DModel", + { + "num_layers": 2, + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "text_embed_dim": 16, + "condition_dim": 8, + "original_format": "megatron", + }, + ), + ( + "ConsistencyDecoderVAE", + { + "encoder_block_out_channels": [32, 64], + "encoder_down_block_types": ["DownEncoderBlock2D"] * 2, + "encoder_layers_per_block": 1, + "decoder_block_out_channels": [32, 64], + "decoder_down_block_types": ["ResnetDownsampleBlock2D"] * 2, + "decoder_up_block_types": ["ResnetUpsampleBlock2D"] * 2, + "decoder_layers_per_block": 1, + }, + ), + ( + "UNet2DConditionModel", + { + "block_out_channels": [32, 64, 64], + "down_block_types": ["KDownBlock2D", "KCrossAttnDownBlock2D", "KCrossAttnDownBlock2D"], + "up_block_types": ["KCrossAttnUpBlock2D", "KCrossAttnUpBlock2D", "KUpBlock2D"], + "mid_block_type": None, + "layers_per_block": 2, + "cross_attention_dim": 16, + "attention_head_dim": 8, + "time_embedding_type": "fourier", + "norm_num_groups": None, + "time_cond_proj_dim": 16, + "conv_in_kernel": 1, + "conv_out_kernel": 1, + }, + ), + ( + "UNet2DConditionModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["SimpleCrossAttnDownBlock2D", "ResnetDownsampleBlock2D"], + "up_block_types": ["ResnetUpsampleBlock2D", "SimpleCrossAttnUpBlock2D"], + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "layers_per_block": 1, + "cross_attention_dim": 32, + "attention_head_dim": 8, + "encoder_hid_dim": 16, + "addition_embed_type": "text", + "addition_embed_type_num_heads": 2, + "cross_attention_norm": "group_norm", + }, + ), + ( + "UNet2DConditionModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["SimpleCrossAttnDownBlock2D", "ResnetDownsampleBlock2D"], + "up_block_types": ["ResnetUpsampleBlock2D", "SimpleCrossAttnUpBlock2D"], + "mid_block_type": "UNetMidBlock2DSimpleCrossAttn", + "layers_per_block": 1, + "cross_attention_dim": 32, + "attention_head_dim": 8, + "encoder_hid_dim": 16, + "addition_embed_type": "text", + "addition_embed_type_num_heads": 2, + "cross_attention_norm": "group_norm", + "only_cross_attention": True, + "class_embed_type": "timestep", + }, + ), + ( + "JoyImageEditTransformer3DModel", + {"hidden_size": 32, "num_attention_heads": 4, "text_dim": 16, "num_layers": 1, "rope_dim_list": [2, 2, 4]}, + ), + ( + "JoyImageEditPlusTransformer3DModel", + {"hidden_size": 32, "num_attention_heads": 4, "text_dim": 16, "num_layers": 1, "rope_dim_list": [2, 2, 4]}, + ), + ( + "Transformer2DModel", + { + "num_attention_heads": 2, + "attention_head_dim": 8, + "num_layers": 1, + "num_vector_embeds": 16, + "sample_size": 4, + "cross_attention_dim": 16, + "num_embeds_ada_norm": 10, + "norm_type": "ada_norm", + "attention_bias": True, + "activation_fn": "gelu", + }, + ), + ( + "VQModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"], + "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D"], + "layers_per_block": 1, + "latent_channels": 4, + "num_vq_embeddings": 16, + }, + ), + ( + "ShapERenderer", + { + "d_latent": 16, + "d_hidden": 16, + "param_names": ["nerstf.mlp.0.weight"], + "param_shapes": [[16, 93]], + "n_hidden_layers": 2, + "insert_direction_at": 1, + }, + ), + ( + "LTX2DurationHead", + { + "video_cross_attention_dim": 16, + "audio_cross_attention_dim": 8, + "pooler_hidden_dim": 16, + "mlp_hidden_dim": 16, + }, + ), + ("LTX2LatentUpsamplerModel", {"in_channels": 4, "mid_channels": 32, "num_blocks_per_stage": 1}), + ( + "LTX2LatentUpsamplerModel", + { + "in_channels": 4, + "mid_channels": 32, + "num_blocks_per_stage": 1, + "spatial_upsample": False, + "temporal_upsample": True, + }, + ), + ( + "UNet2DConditionModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["CrossAttnDownBlock2D", "DownBlock2D"], + "up_block_types": ["UpBlock2D", "CrossAttnUpBlock2D"], + "layers_per_block": 1, + "cross_attention_dim": 32, + "attention_type": "gated", + }, + ), + ( + "UNet2DConditionModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["CrossAttnDownBlock2D", "DownBlock2D"], + "up_block_types": ["UpBlock2D", "CrossAttnUpBlock2D"], + "layers_per_block": 1, + "cross_attention_dim": 32, + "attention_type": "gated-text-image", + }, + ), + ( + "VQModel", + { + "block_out_channels": [32, 64], + "down_block_types": ["DownEncoderBlock2D"] * 2, + "up_block_types": ["UpDecoderBlock2D"] * 2, + "layers_per_block": 1, + "norm_type": "spatial", + "latent_channels": 4, + "num_vq_embeddings": 16, + }, + ), + ( + "UniDiffuserModel", + { + "text_dim": 8, + "clip_img_dim": 8, + "num_attention_heads": 2, + "attention_head_dim": 8, + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "sample_size": 8, + "patch_size": 2, + "activation_fn": "gelu", + "use_data_type_embedding": True, + }, + ), + ( + "UniDiffuserTextDecoder", + { + "prefix_length": 4, + "prefix_inner_dim": 8, + "prefix_hidden_dim": 8, + "vocab_size": 32, + "n_positions": 16, + "n_embd": 16, + "n_layer": 1, + "n_head": 2, + }, + ), + ( + "LTX2TextConnectors", + { + "caption_channels": 16, + "text_proj_in_factor": 2, + "video_connector_num_attention_heads": 2, + "video_connector_attention_head_dim": 8, + "video_connector_num_layers": 1, + "audio_connector_num_attention_heads": 2, + "audio_connector_attention_head_dim": 8, + "audio_connector_num_layers": 1, + }, + ), + ( + "LTX2TextConnectors", + { + "caption_channels": 16, + "text_proj_in_factor": 2, + "video_connector_num_attention_heads": 2, + "video_connector_attention_head_dim": 8, + "video_connector_num_layers": 1, + "audio_connector_num_attention_heads": 2, + "audio_connector_attention_head_dim": 8, + "audio_connector_num_layers": 1, + "per_modality_projections": True, + "video_hidden_dim": 16, + "audio_hidden_dim": 16, + "video_gated_attn": True, + "audio_gated_attn": True, + "proj_bias": True, + }, + ), + ( + "LTX2Vocoder", + { + "in_channels": 8, + "hidden_channels": 32, + "upsample_kernel_sizes": [4, 4], + "upsample_factors": [2, 2], + "resnet_kernel_sizes": [3], + "resnet_dilations": [[1, 3]], + }, + ), + ( + "LTX2VocoderWithBWE", + { + "in_channels": 8, + "hidden_channels": 32, + "upsample_kernel_sizes": [4, 4], + "upsample_factors": [2, 2], + "resnet_kernel_sizes": [3], + "resnet_dilations": [[1, 3]], + "bwe_in_channels": 8, + "bwe_hidden_channels": 32, + "bwe_upsample_kernel_sizes": [4], + "bwe_upsample_factors": [2], + "bwe_resnet_kernel_sizes": [3], + "bwe_resnet_dilations": [[1, 3]], + "filter_length": 16, + "window_length": 16, + "num_mel_channels": 8, + }, + ), + ( + "LTX2VideoDiffusionDecoderModel", + { + "latent_channels": 8, + "decoder_head_dim": 8, + "decoder_stage_channels": [32, 16, 8, 8, 8], + "decoder_stage_depths": [1, 1, 1, 1, 1], + "decoder_upsample_channel_reductions": [2, 2, 1, 1], + "decoder_t_emb_dim": 16, + }, + ), + ( + "LTX2VideoDiffusionDecoderModel", + { + "latent_channels": 8, + "decoder_head_dim": 8, + "decoder_stage_channels": [32, 16, 8, 8, 8], + "decoder_stage_depths": [1, 1, 1, 1, 1], + "decoder_upsample_channel_reductions": [2, 2, 1, 1], + "decoder_t_emb_dim": 16, + "original_format": "ltx2_diffusion_decoder_gated", + }, + ), + ("WuerstchenPrior", {"c_in": 4, "c": 16, "c_cond": 8, "c_r": 8, "depth": 2, "nhead": 2}), + ( + "WuerstchenDiffNeXt", + { + "c_hidden": [16, 32], + "nhead": [2, 2], + "blocks": [1, 2], + "level_config": ["CT", "CTA"], + "inject_effnet": [False, True], + "c_cond": 16, + "clip_embd": 8, + "effnet_embd": 4, + }, + ), + ("PaellaVQModel", {"levels": 2, "bottleneck_blocks": 2, "embed_dim": 16, "num_vq_embeddings": 16}), + ("T2IAdapter", {"channels": [16, 32], "num_res_blocks": 2, "adapter_type": "full_adapter"}), + ("T2IAdapter", {"channels": [16, 32], "num_res_blocks": 2, "adapter_type": "light_adapter"}), + ("T2IAdapter", {"channels": [16, 32, 32, 32], "num_res_blocks": 2, "adapter_type": "full_adapter_xl"}), + ( + "AutoencoderKLFlux2", + { + "block_out_channels": [32, 64], + "decoder_block_out_channels": [32, 32], + "layers_per_block": 1, + "down_block_types": ["DownEncoderBlock2D"] * 2, + "up_block_types": ["UpDecoderBlock2D"] * 2, + }, + ), + ( + "I2VGenXLUNet", + { + "block_out_channels": [32, 64], + "down_block_types": ["CrossAttnDownBlock3D", "DownBlock3D"], + "up_block_types": ["UpBlock3D", "CrossAttnUpBlock3D"], + "layers_per_block": 1, + "cross_attention_dim": 32, + "attention_head_dim": 8, + }, + ), + ( + "UNet2DModel", + { + "block_out_channels": (32, 64), + "down_block_types": ("DownBlock2D", "AttnDownBlock2D"), + "up_block_types": ("AttnUpBlock2D", "UpBlock2D"), + "layers_per_block": 1, + }, + ), + ( + "UNet2DModel", + { + "block_out_channels": (32, 64), + "down_block_types": ("DownBlock2D", "AttnDownBlock2D"), + "up_block_types": ("AttnUpBlock2D", "UpBlock2D"), + "layers_per_block": 1, + "original_format": "ldm", + }, + ), + ( + "UNet2DModel", + { + "block_out_channels": (32, 64), + "down_block_types": ("ResnetDownsampleBlock2D", "AttnDownBlock2D"), + "up_block_types": ("AttnUpBlock2D", "ResnetUpsampleBlock2D"), + "layers_per_block": 1, + "downsample_type": "resnet", + "upsample_type": "resnet", + "original_format": "consistency", + }, + ), + ( + "UNet2DModel", + { + "block_out_channels": (32, 64), + "down_block_types": ("SkipDownBlock2D", "AttnSkipDownBlock2D"), + "up_block_types": ("AttnSkipUpBlock2D", "SkipUpBlock2D"), + "layers_per_block": 1, + "time_embedding_type": "fourier", + }, + ), + ( + "AutoencoderRAE", + { + "encoder_type": "dinov2", + "encoder_hidden_size": 64, + "encoder_num_hidden_layers": 1, + "encoder_patch_size": 16, + "encoder_input_size": 32, + "decoder_hidden_size": 32, + "decoder_num_hidden_layers": 1, + "decoder_num_attention_heads": 2, + "decoder_intermediate_size": 64, + }, + ), + ( + "AutoencoderRAE", + { + "encoder_type": "mae", + "encoder_hidden_size": 64, + "encoder_num_hidden_layers": 1, + "encoder_patch_size": 16, + "encoder_input_size": 32, + "decoder_hidden_size": 32, + "decoder_num_hidden_layers": 1, + "decoder_num_attention_heads": 2, + "decoder_intermediate_size": 64, + }, + ), + ( + "AutoencoderKLMochi", + { + "encoder_block_out_channels": (32, 64), + "decoder_block_out_channels": (32, 64), + "layers_per_block": (1, 1, 1), + "temporal_expansions": (2,), + "spatial_expansions": (2,), + "add_attention_block": (False, True, True), + }, + ), + ( + "AutoencoderKLHunyuanVideo", + { + "block_out_channels": (32, 64), + "down_block_types": ("HunyuanVideoDownBlock3D",) * 2, + "up_block_types": ("HunyuanVideoUpBlock3D",) * 2, + "layers_per_block": 1, + "spatial_compression_ratio": 2, + }, + ), + ( + "AutoencoderKLHunyuanVideo15", + { + "block_out_channels": (8, 16, 16), + "layers_per_block": 1, + "latent_channels": 4, + "spatial_compression_ratio": 4, + "temporal_compression_ratio": 2, + }, + ), + ( + "AutoencoderKLCosmos", + { + "encoder_block_out_channels": (8, 16, 16), + "decode_block_out_channels": (8, 16, 16), + "num_layers": 1, + "resolution": 16, + "patch_size": 2, + "attention_resolutions": (8, 4), + "spatial_compression_ratio": 4, + "temporal_compression_ratio": 4, + }, + ), + ( + "HunyuanVideo15Transformer3DModel", + { + "in_channels": 4, + "out_channels": 4, + "num_attention_heads": 2, + "attention_head_dim": 8, + "num_layers": 2, + "num_refiner_layers": 1, + "text_embed_dim": 16, + "text_embed_2_dim": 16, + "image_embed_dim": 16, + "rope_axes_dim": (2, 2, 4), + }, + ), + ( + "AnimaTextConditioner", + { + "source_dim": 8, + "target_dim": 8, + "model_dim": 16, + "num_layers": 2, + "num_attention_heads": 2, + "target_vocab_size": 32, + }, + ), + ("AutoencoderKLQwenImage", {"base_dim": 4, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 1}), + ( + "ControlNetModel", + { + "block_out_channels": (32, 64), + "down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"), + "layers_per_block": 1, + "cross_attention_dim": 32, + }, + ), + ("MotionAdapter", {"block_out_channels": (32, 64), "motion_layers_per_block": 1}), + ( + "StableCascadeUNet", + { + "block_out_channels": (16, 32), + "conditioning_dim": 16, + "num_attention_heads": (2, 4), + "down_num_layers_per_block": (1, 2), + "up_num_layers_per_block": (2, 1), + }, + ), + ( + "ZImageControlNetModel", + { + "control_layers_places": [0, 2], + "control_in_dim": 4, + "dim": 16, + "n_refiner_layers": 1, + "n_heads": 2, + "n_kv_heads": 2, + }, + ), + ( + "StableAudioDiTModel", + { + "in_channels": 4, + "out_channels": 4, + "num_layers": 2, + "num_attention_heads": 2, + "num_key_value_attention_heads": 1, + "attention_head_dim": 8, + "cross_attention_dim": 16, + "time_proj_dim": 16, + "global_states_input_dim": 16, + "cross_attention_input_dim": 16, + }, + ), + ("StableAudioProjectionModel", {"text_encoder_dim": 16, "conditioning_dim": 16, "min_value": 0, "max_value": 60}), + ("StableAudio3DurationEmbedder", {"output_dim": 8, "fourier_dim": 16}), + ( + "AutoencoderSAME", + { + "audio_channels": 2, + "patch_size": 4, + "encoder_channels": 8, + "encoder_c_mults": (2, 2), + "encoder_strides": (2, 2), + "encoder_transformer_depths": (1, 2), + "latent_dim": 4, + "dim_heads": 4, + }, + ), + ( + "LongCatAudioDiTTransformer", + {"dit_dim": 16, "dit_depth": 2, "dit_heads": 2, "dit_text_dim": 16, "latent_dim": 4}, + ), + ( + "LongCatAudioDiTVae", + {"channels": 8, "c_mults": [1, 2], "strides": [2, 2], "latent_dim": 4, "encoder_latent_dim": 8}, + ), + ("MiniMaxMusic3ConditionEncoder", {"condition_hidden_dim": 8, "num_condition_layers": 2, "out_dim": 16}), + ( + "MiniMaxMusic3Vocoder", + {"latent_channels": 8, "decoder_input_dim": 16, "decoder_hidden_dim": 32, "upsampling_ratios": (2, 2)}, + ), + ( + "MiniMaxMusic3RVQDepthDecoder", + { + "hidden_size": 16, + "num_layers": 2, + "num_attention_heads": 2, + "intermediate_size": 32, + "audio_vocab_size": 16, + "num_codebooks": 3, + }, + ), + ( + "AceStepTransformer1DModel", + { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + "in_channels": 12, + "audio_acoustic_hidden_dim": 4, + }, + ), + ( + "AceStepConditionEncoder", + { + "hidden_size": 16, + "intermediate_size": 32, + "text_hidden_dim": 8, + "timbre_hidden_dim": 4, + "num_lyric_encoder_hidden_layers": 1, + "num_timbre_encoder_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + }, + ), + ( + "AceStepAudioTokenizer", + { + "hidden_size": 16, + "intermediate_size": 32, + "audio_acoustic_hidden_dim": 4, + "fsq_dim": 16, + "num_attention_pooler_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + }, + ), + ( + "AceStepAudioTokenDetokenizer", + { + "hidden_size": 16, + "intermediate_size": 32, + "audio_acoustic_hidden_dim": 4, + "num_attention_pooler_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + }, + ), + ( + "UVit2DModel", + { + "hidden_size": 16, + "cond_embed_dim": 8, + "micro_cond_embed_dim": 16, + "encoder_hidden_size": 16, + "vocab_size": 32, + "codebook_size": 32, + "in_channels": 16, + "block_out_channels": 16, + "num_res_blocks": 1, + "block_num_heads": 2, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "intermediate_size": 32, + }, + ), + ( + "VQModel", + { + "block_out_channels": (16, 32), + "down_block_types": ("DownEncoderBlock2D",) * 2, + "up_block_types": ("UpDecoderBlock2D",) * 2, + "norm_num_groups": 8, + "latent_channels": 4, + "num_vq_embeddings": 16, + }, + ), + ( + "AutoencoderTiny", + { + "encoder_block_out_channels": (8, 8), + "decoder_block_out_channels": (8, 8), + "num_encoder_blocks": (1, 2), + "num_decoder_blocks": (2, 1), + }, + ), + ( + "T5FilmDecoder", + {"input_dims": 4, "targets_length": 16, "d_model": 16, "num_layers": 2, "num_heads": 2, "d_kv": 8, "d_ff": 32}, + ), +] + + +CASES.extend( + (name, {**config, "original_format": "consistency_decoder_jit"}) + for name, config in list(CASES) + if name == "ConsistencyDecoderVAE" +) + + +@pytest.mark.parametrize("model_class,config", CASES, ids=[name for name, _ in CASES]) +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16], ids=["float32", "float16", "bfloat16"] +) +def test_auxiliary_model_conversion(model_class, config, dtype, tmp_path): + model_config = { + key: value for key, value in config.items() if key not in ("original_format", "transformers_version") + } + internal_modules = { + "UNetFlatConditionModel": "diffusers.pipelines.deprecated.versatile_diffusion.modeling_text_unet", + "UnCLIPTextProjModel": "diffusers.pipelines.deprecated.unclip.text_proj", + "SpectrogramContEncoder": "diffusers.pipelines.deprecated.spectrogram_diffusion.continuous_encoder", + "SpectrogramNotesEncoder": "diffusers.pipelines.deprecated.spectrogram_diffusion.notes_encoder", + "LTXLatentUpsamplerModel": "diffusers.pipelines.ltx.modeling_latent_upsampler", + "ShapERenderer": "diffusers.pipelines.shap_e.renderer", + "UniDiffuserModel": "diffusers.pipelines.deprecated.unidiffuser.modeling_uvit", + "UniDiffuserTextDecoder": "diffusers.pipelines.deprecated.unidiffuser.modeling_text_decoder", + "LTX2TextConnectors": "diffusers.pipelines.ltx2.connectors", + "LTX2Vocoder": "diffusers.pipelines.ltx2.vocoder", + "LTX2VocoderWithBWE": "diffusers.pipelines.ltx2.vocoder", + "LTX2DurationHead": "diffusers.pipelines.ltx2.duration_head", + "LTX2LatentUpsamplerModel": "diffusers.pipelines.ltx2.latent_upsampler", + "WuerstchenPrior": "diffusers.pipelines.deprecated.wuerstchen.modeling_wuerstchen_prior", + "WuerstchenDiffNeXt": "diffusers.pipelines.deprecated.wuerstchen.modeling_wuerstchen_diffnext", + "PaellaVQModel": "diffusers.pipelines.deprecated.wuerstchen.modeling_paella_vq_model", + } + module = importlib.import_module(internal_modules[model_class]) if model_class in internal_modules else diffusers + model = getattr(module, model_class)(**model_config).to(dtype=dtype) + state = model.state_dict() + if model_class == "ShapERenderer": + from diffusers.loaders.conversion.shap_e_tables import create_mc_lookup_table + + state["mesh_decoder.cases"], state["mesh_decoder.masks"] = create_mc_lookup_table() + conversion_config = dict(model.config) + conversion_config.update( + {key: value for key, value in config.items() if key in ("original_format", "transformers_version")} + ) + conversion = get_conversion(model_class, conversion_config) + assert conversion.diffusers_keys == set(state) + if model_class == "UniDiffuserModel": + for p in ("weight", "bias"): + state[f"transformer.pos_embed.proj.{p}"] = state[f"vae_img_in.proj.{p}"].clone() + original = conversion.to_original(state) + restored = conversion.to_diffusers(original) + for key in state: + torch.testing.assert_close(restored[key], state[key], rtol=0, atol=0) + if model_class == "AutoencoderRAE" and config["encoder_type"] == "siglip2" and dtype == torch.float32: + model.eval() + sample = torch.randn(1, 3, 32, 32) + with torch.no_grad(): + expected = model.encode(sample).latent + model.load_state_dict(restored) + actual = model.encode(sample).latent + assert actual.shape == (1, 64, 2, 2) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + if dtype == torch.float32 and hasattr(type(model), "from_single_file"): + (tmp_path / "config.json").write_text(json.dumps(conversion_config), encoding="utf-8") + loaded = type(model).from_single_file(original, config=str(tmp_path), local_files_only=True) + for key in state: + torch.testing.assert_close(loaded.state_dict()[key], state[key], rtol=0, atol=0) diff --git a/tests/single_file/test_cogvideox_conversion.py b/tests/single_file/test_cogvideox_conversion.py new file mode 100644 index 000000000000..c0edb3a42cc0 --- /dev/null +++ b/tests/single_file/test_cogvideox_conversion.py @@ -0,0 +1,308 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest +import torch + +from diffusers import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel +from diffusers.loaders.conversion.cogvideox import ( + COGVIDEOX_TRANSFORMER_PREFIX, + CogVideoXAdaLN, + cogvideox_transformer_conversion, + cogvideox_vae_conversion, + convert_cogvideox_transformer_checkpoint_to_diffusers, + convert_cogvideox_vae_checkpoint_to_diffusers, + unwrap_cogvideox_checkpoint, +) +from diffusers.loaders.conversion.io import convert_checkpoint + + +def export_component(model_path, output_path, component): + from pathlib import Path + + model_path = Path(model_path) + if not (model_path / "config.json").is_file(): + model_path = model_path / component + config = json.loads((model_path / "config.json").read_text()) + expected = "CogVideoXTransformer3DModel" if component == "transformer" else "AutoencoderKLCogVideoX" + if config["_class_name"] != expected: + raise ValueError(f"Expected {expected}") + return convert_checkpoint( + model_path, + output_path, + config=config, + reverse=True, + output_format="pytorch", + output_prefix=COGVIDEOX_TRANSFORMER_PREFIX if component == "transformer" else "", + output_wrapper=("module" if component == "transformer" else "state_dict",), + ) + + +def make_transformer(variant="1.0-t2v", **kwargs): + torch.manual_seed(0) + return CogVideoXTransformer3DModel( + num_layers=2, + num_attention_heads=2, + attention_head_dim=8, + in_channels=8 if variant.endswith("i2v") else 4, + out_channels=4, + time_embed_dim=16, + text_embed_dim=8, + sample_height=4, + sample_width=4, + sample_frames=5, + max_text_seq_length=4, + patch_size_t=2 if variant.startswith("1.5") else None, + patch_bias=variant.startswith("1.0"), + use_rotary_positional_embeddings=variant != "1.0-t2v", + use_learned_positional_embeddings=variant == "1.0-i2v", + ofs_embed_dim=8 if variant == "1.5-i2v" else None, + **kwargs, + ) + + +def make_vae(channels=(8, 16), quant_conv=False): + torch.manual_seed(0) + return AutoencoderKLCogVideoX( + block_out_channels=channels, + down_block_types=("CogVideoXDownBlock3D",) * len(channels), + up_block_types=("CogVideoXUpBlock3D",) * len(channels), + layers_per_block=1, + norm_num_groups=4, + latent_channels=4, + sample_height=16, + sample_width=16, + use_quant_conv=quant_conv, + use_post_quant_conv=quant_conv, + ) + + +class TestCogVideoXConversion: + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("parameter", ["weight", "bias"]) + def test_adaln_known_layout(self, dtype, parameter): + # SAT groups: image attention, image MLP, text attention, text MLP; three tensors per group. + original = torch.arange(24, dtype=dtype) + if parameter == "weight": + original = original[:, None].expand(-1, 3).clone() + norm1, norm2 = CogVideoXAdaLN(2).forward((original,)) + expected1 = torch.cat((original[0:6], original[12:18])) + expected2 = torch.cat((original[6:12], original[18:24])) + torch.testing.assert_close(norm1, expected1, rtol=0, atol=0) + torch.testing.assert_close(norm2, expected2, rtol=0, atol=0) + torch.testing.assert_close(CogVideoXAdaLN(2).inverse((norm1, norm2))[0], original, rtol=0, atol=0) + + @pytest.mark.parametrize("variant", ["1.0-t2v", "1.0-i2v", "1.5-t2v", "1.5-i2v"]) + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_transformer_complete_roundtrip(self, variant, dtype): + model = make_transformer(variant).to(dtype=dtype) + conversion = cogvideox_transformer_conversion(model.config) + diffusers = model.state_dict() + original = conversion.to_original(diffusers) + # Verify names and ordering independently of the inverse executor. + torch.testing.assert_close( + original["transformer.layers.1.attention.query_key_value.weight"], + torch.cat([diffusers[f"transformer_blocks.1.attn1.to_{name}.weight"] for name in "qkv"]), + rtol=0, + atol=0, + ) + assert ( + original["transformer.layers.0.post_attention_layernorm.weight"] + is diffusers["transformer_blocks.0.norm2.norm.weight"] + ) + assert ( + original["mixins.adaln_layer.query_layernorm_list.0.bias"] + is diffusers["transformer_blocks.0.attn1.norm_q.bias"] + ) + restored = conversion.to_diffusers(original) + for key in diffusers: + torch.testing.assert_close(restored[key], diffusers[key], rtol=0, atol=0) + for key, tensor in conversion.to_original(restored).items(): + torch.testing.assert_close(tensor, original[key], rtol=0, atol=0) + if variant == "1.0-i2v": + assert "mixins.pos_embed.pos_embedding" in original + if variant.startswith("1.5"): + assert "mixins.patch_embed.proj.bias" in original + + def test_transformer_without_optional_biases_and_affine_norms(self): + model = make_transformer(attention_bias=False, norm_elementwise_affine=False) + conversion = cogvideox_transformer_conversion(model.config) + original = conversion.to_original(model.state_dict()) + assert "transformer.layers.0.attention.query_key_value.bias" not in original + assert "transformer.layers.0.input_layernorm.weight" not in original + restored = conversion.to_diffusers(original) + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(restored[key], tensor, rtol=0, atol=0) + + @pytest.mark.parametrize("channels", [(8, 8), (8, 16, 16, 32)]) + @pytest.mark.parametrize("quant_conv", [False, True]) + def test_vae_complete_roundtrip_and_up_block_numbering(self, channels, quant_conv): + model = make_vae(channels, quant_conv) + conversion = cogvideox_vae_conversion(model.config) + diffusers = model.state_dict() + original = conversion.to_original(diffusers) + for i in range(len(channels)): + assert ( + original[f"decoder.up.{len(channels) - 1 - i}.block.0.conv1.conv.weight"] + is diffusers[f"decoder.up_blocks.{i}.resnets.0.conv1.conv.weight"] + ) + assert ( + original["encoder.mid.block_2.conv1.conv.weight"] + is diffusers["encoder.mid_block.resnets.1.conv1.conv.weight"] + ) + restored = conversion.to_diffusers(original) + for key in diffusers: + torch.testing.assert_close(restored[key], diffusers[key], rtol=0, atol=0) + + def test_nested_checkpoint_and_component_selection_are_non_mutating(self): + model = make_transformer() + original = cogvideox_transformer_conversion(model.config).to_original(model.state_dict()) + prefixed = {COGVIDEOX_TRANSFORMER_PREFIX + key: tensor for key, tensor in original.items()} + prefixed[COGVIDEOX_TRANSFORMER_PREFIX + "mixins.pos_embed.freqs_sin"] = torch.zeros(1) + prefixed["conditioner.sibling.weight"] = torch.ones(1) + wrapped = {"model": {"module": {"state_dict": prefixed}}, "iteration": 10} + original_keys = set(prefixed) + converted = convert_cogvideox_transformer_checkpoint_to_diffusers(wrapped, model.config) + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(converted[key], tensor, rtol=0, atol=0) + assert set(prefixed) == original_keys + assert wrapped["iteration"] == 10 + assert prefixed["conditioner.sibling.weight"].item() == 1 + + def test_unknown_auxiliary_key_is_not_silently_ignored(self): + model = make_transformer() + original = cogvideox_transformer_conversion(model.config).to_original(model.state_dict()) + original["mixins.pos_embed.freqs_sin_typo"] = torch.zeros(1) + with pytest.raises(ValueError, match="unexpected keys.*freqs_sin_typo"): + convert_cogvideox_transformer_checkpoint_to_diffusers(original, model.config) + + def test_mixed_prefixes_are_rejected(self): + model = make_transformer() + original = cogvideox_transformer_conversion(model.config).to_original(model.state_dict()) + original[COGVIDEOX_TRANSFORMER_PREFIX + "time_embed.0.weight"] = original["time_embed.0.weight"] + with pytest.raises(ValueError, match="mixes prefixed and unprefixed"): + convert_cogvideox_transformer_checkpoint_to_diffusers(original, model.config) + + def test_ambiguous_wrappers_are_rejected(self): + with pytest.raises(ValueError, match="Ambiguous"): + unwrap_cogvideox_checkpoint({"model": {}, "module": {}}) + + def test_vae_loss_is_excluded_without_mutation(self): + model = make_vae() + original = cogvideox_vae_conversion(model.config).to_original(model.state_dict()) + original["loss.discriminator.weight"] = torch.ones(1) + converted = convert_cogvideox_vae_checkpoint_to_diffusers({"state_dict": original}, model.config) + assert "loss.discriminator.weight" in original + assert set(converted) == set(model.state_dict()) + + @pytest.mark.parametrize("model_class", [CogVideoXTransformer3DModel, AutoencoderKLCogVideoX]) + def test_single_file_requires_explicit_config(self, model_class): + with pytest.raises(ValueError, match="requires an explicit Diffusers `config`"): + model_class.from_single_file({}) + + @pytest.mark.parametrize("component", ["transformer", "vae"]) + @pytest.mark.parametrize("safe_serialization", [False, True]) + @pytest.mark.parametrize("sharded", [False, True]) + def test_export_and_single_file_loading(self, tmp_path, component, safe_serialization, sharded): + model = make_transformer("1.5-i2v") if component == "transformer" else make_vae() + model = model.to(dtype=torch.bfloat16) + component_path = tmp_path / "pipeline" / component + model.save_pretrained( + component_path, + safe_serialization=safe_serialization, + max_shard_size="20KB" if sharded else "5GB", + ) + output_path = tmp_path / "original.pt" + export_component(str(component_path.parent), str(output_path), component) + checkpoint = torch.load(output_path, weights_only=True) + original = unwrap_cogvideox_checkpoint(checkpoint) + assert all(tensor.dtype == torch.bfloat16 for tensor in original.values()) + loaded = type(model).from_single_file( + str(output_path), config=str(component_path), torch_dtype=torch.bfloat16, local_files_only=True + ) + restored_path = tmp_path / "shared-converter" + convert_checkpoint( + output_path, + restored_path, + config=dict(model.config), + input_wrapper=("module" if component == "transformer" else "state_dict",), + input_prefix=COGVIDEOX_TRANSFORMER_PREFIX if component == "transformer" else "", + ) + restored_model = type(model).from_pretrained(restored_path, torch_dtype=torch.bfloat16, local_files_only=True) + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], tensor, rtol=0, atol=0) + torch.testing.assert_close(restored_model.state_dict()[key], tensor, rtol=0, atol=0) + + # Conversion can create disjoint tensor views sharing storage. They must survive safetensors serialization. + roundtrip_path = tmp_path / "roundtrip" + loaded.save_pretrained(roundtrip_path) + reloaded = type(model).from_pretrained(roundtrip_path, torch_dtype=torch.bfloat16, local_files_only=True) + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(reloaded.state_dict()[key], tensor, rtol=0, atol=0) + + @pytest.mark.parametrize("low_cpu_mem_usage", [False, True]) + def test_transformer_forward_after_single_file_loading(self, tmp_path, low_cpu_mem_usage): + model = make_transformer().eval() + model.save_config(tmp_path) + original = cogvideox_transformer_conversion(model.config).to_original(model.state_dict()) + loaded = ( + type(model) + .from_single_file( + original, config=str(tmp_path), low_cpu_mem_usage=low_cpu_mem_usage, local_files_only=True + ) + .eval() + ) + generator = torch.Generator().manual_seed(42) + inputs = { + "hidden_states": torch.randn(1, 2, 4, 4, 4, generator=generator), + "encoder_hidden_states": torch.randn(1, 4, 8, generator=generator), + "timestep": torch.tensor([1]), + } + with torch.no_grad(): + expected = model(**inputs).sample + actual = loaded(**inputs).sample + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) + def test_older_transformer_config_and_fixed_position_embedding(self, tmp_path, dtype): + model = make_transformer().to(dtype=dtype) + model.save_pretrained(tmp_path) + config_path = tmp_path / "config.json" + config = json.loads(config_path.read_text()) + # These options were introduced after the original 1.0 checkpoints. + for key in ("ofs_embed_dim", "patch_size_t", "patch_bias", "use_learned_positional_embeddings"): + config.pop(key) + config_path.write_text(json.dumps(config)) + output_path = tmp_path / "original.pt" + export_component(str(tmp_path), str(output_path), "transformer") + checkpoint = torch.load(output_path, weights_only=True) + original_positions = checkpoint["module"][COGVIDEOX_TRANSFORMER_PREFIX + "mixins.pos_embed.pos_embedding"] + torch.testing.assert_close(original_positions, model.patch_embed.pos_embedding, rtol=0, atol=0) + loaded = type(model).from_single_file( + str(output_path), config=str(tmp_path), local_files_only=True, torch_dtype=dtype + ) + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], tensor, rtol=0, atol=0) + + def test_quantized_export_is_rejected(self, tmp_path): + model = make_transformer() + model.save_config(tmp_path) + config_path = tmp_path / "config.json" + config = json.loads(config_path.read_text()) + config["quantization_config"] = {"quant_method": "bitsandbytes"} + config_path.write_text(json.dumps(config)) + with pytest.raises(ValueError, match="requires unpacked, unquantized tensor weights"): + export_component(str(tmp_path), str(tmp_path / "original.pt"), "transformer") diff --git a/tests/single_file/test_model_conversions.py b/tests/single_file/test_model_conversions.py new file mode 100644 index 000000000000..a7563e16af96 --- /dev/null +++ b/tests/single_file/test_model_conversions.py @@ -0,0 +1,198 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import importlib + +import pytest +import torch + +from diffusers.loaders.conversion import get_conversion +from diffusers.loaders.conversion.transforms import MergeEqual, Reverse + + +def test_single_file_config_defaults_have_bidirectional_definitions(): + from diffusers.loaders.conversion.registry import CONVERSION_BUILDERS + from diffusers.loaders.single_file_model import SINGLE_FILE_CONFIGS + + assert SINGLE_FILE_CONFIGS.keys() <= CONVERSION_BUILDERS.keys() + + +def test_new_single_file_model_requires_config_and_resolves_subclasses(): + from diffusers import AutoencoderTiny, DiTTransformer2DModel + from diffusers.loaders.single_file_model import _get_single_file_loadable_mapping_class + + class DerivedDiT(DiTTransformer2DModel): + pass + + assert _get_single_file_loadable_mapping_class(DerivedDiT) == "DiTTransformer2DModel" + with pytest.raises(ValueError, match="requires an explicit Diffusers"): + AutoencoderTiny.from_single_file({}, local_files_only=True) + + +CASES = [ + ("tests.models.autoencoders.test_models_asymmetric_autoencoder_kl", "AsymmetricAutoencoderKLTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_cosmos", "AutoencoderKLCosmosTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_dc", "AutoencoderDCTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_hunyuan_video", "AutoencoderKLHunyuanVideoTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_kl", "AutoencoderKLTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_kl_ltx2_audio", "AutoencoderKLLTX2AudioTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_kl_minimax_h3", "AutoencoderKLMiniMaxH3TesterConfig"), + ( + "tests.models.autoencoders.test_models_autoencoder_kl_minimax_h3_audio", + "AutoencoderKLMiniMaxH3AudioTesterConfig", + ), + ( + "tests.models.autoencoders.test_models_autoencoder_kl_temporal_decoder", + "AutoencoderKLTemporalDecoderTesterConfig", + ), + ("tests.models.autoencoders.test_models_autoencoder_ltx2_video", "AutoencoderKLLTX2VideoTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_ltx_video", "AutoencoderKLLTXVideo090TesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_ltx_video", "AutoencoderKLLTXVideo091TesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_mochi", "AutoencoderKLMochiTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_oobleck", "AutoencoderOobleckTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_same", "AutoencoderSAMETesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_tiny", "AutoencoderTinyTesterConfig"), + ("tests.models.autoencoders.test_models_autoencoder_wan", "AutoencoderKLWanTesterConfig"), + ("tests.models.autoencoders.test_models_vq", "VQModelTesterConfig"), + ("tests.models.controlnets.test_models_controlnet_cosmos", "CosmosControlNetModelTesterConfig"), + ("tests.models.transformers.test_models_dit_transformer2d", "DiTTransformer2DTesterConfig"), + ("tests.models.transformers.test_models_pixart_transformer2d", "PixArtTransformer2DTesterConfig"), + ("tests.models.transformers.test_models_prior", "PriorTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_ace_step", "AceStepTransformer1DModelTesterConfig"), + ("tests.models.transformers.test_models_transformer_anyflow", "AnyFlowTransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_anyflow_far", "AnyFlowFARTransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_aura_flow", "AuraFlowTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_chroma", "ChromaTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_cogview3plus", "CogView3PlusTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_cogview4", "CogView4TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_cosmos", "CosmosTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_cosmos", "CosmosTransformerVideoToWorldTesterConfig"), + ("tests.models.transformers.test_models_transformer_ernie_image", "ErnieImageTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_flux", "FluxTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_flux2", "Flux2TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_flux2", "Flux2TransformerKVCacheTesterConfig"), + ("tests.models.transformers.test_models_transformer_hidream", "HiDreamTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_hunyuan_1_5", "HunyuanVideo15TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_hunyuan_dit", "HunyuanDiTTesterConfig"), + ("tests.models.transformers.test_models_transformer_hunyuan_video", "HunyuanVideoTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_hunyuan_video", "HunyuanVideoI2VTransformerTesterConfig"), + ( + "tests.models.transformers.test_models_transformer_hunyuan_video", + "HunyuanVideoTokenReplaceTransformerTesterConfig", + ), + ("tests.models.transformers.test_models_transformer_longcat_audio_dit", "LongCatAudioDiTTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_ltx", "LTXTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_ltx2", "LTX2TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_lumina", "LuminaNextDiTTesterConfig"), + ("tests.models.transformers.test_models_transformer_lumina2", "Lumina2TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_minimax_h3", "MiniMaxH3TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_minimax_music3", "MiniMaxMusic3Transformer1DTesterConfig"), + ("tests.models.transformers.test_models_transformer_mochi", "MochiTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_omnigen", "OmniGenTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_ovis_image", "OvisImageTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_prx", "PRXTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_sana", "SanaTransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_sana_video", "SanaVideoTransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_sd3", "SD3TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_sd3", "SD35TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_stable_audio3", "StableAudio3DiTTesterConfig"), + ("tests.models.transformers.test_models_transformer_wan", "WanTransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_wan_animate", "WanAnimateTransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_wan_animate_2", "WanAnimate2TransformerTesterConfig"), + ("tests.models.transformers.test_models_transformer_wan_vace", "WanVACETransformer3DTesterConfig"), + ("tests.models.transformers.test_models_transformer_z_image", "ZImageTransformerTesterConfig"), + ("tests.models.unets.test_models_unet_1d", "UNet1DModelTesterConfig"), + ("tests.models.unets.test_models_unet_1d", "UNetRLModelTesterConfig"), + ("tests.models.unets.test_models_unet_2d", "Unet2DModelTesterConfig"), + ("tests.models.unets.test_models_unet_2d", "UNetLDMModelTesterConfig"), + ("tests.models.unets.test_models_unet_2d", "NCSNppModelTesterConfig"), + ("tests.models.unets.test_models_unet_2d_condition", "UNet2DConditionTesterConfig"), + ("tests.models.unets.test_models_unet_3d_condition", "UNet3DConditionModelTesterConfig"), + ("tests.models.unets.test_models_unet_spatiotemporal", "UNetSpatioTemporalConditionModelTesterConfig"), +] + + +@pytest.mark.parametrize("module_name,tester_name", CASES) +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16], ids=["float32", "float16", "bfloat16"] +) +def test_model_conversion_covers_parameters_and_round_trips(module_name, tester_name, dtype, tmp_path): + module = importlib.import_module(module_name) + tester = getattr(module, tester_name)() + model = tester.model_class(**tester.get_init_dict()).to(dtype=dtype) + state = { + key: torch.randn_like(value) if value.is_floating_point() else value.clone() + for key, value in model.state_dict().items() + } + conversion = get_conversion(type(model).__name__, dict(model.config)) + assert conversion.diffusers_keys == set(state) + # DiT and Lumina imports replicate shared original parameters into separate Diffusers modules. + for rule in conversion.rules: + if isinstance(rule.transform, Reverse) and isinstance(rule.transform.transform, MergeEqual): + for key in rule.diffusers[1:]: + state[key] = state[rule.diffusers[0]].clone() + original = conversion.to_original(state) + restored = conversion.to_diffusers(original) + assert set(restored) == set(state) + for key in state: + torch.testing.assert_close(restored[key], state[key], rtol=0, atol=0) + original_restored = conversion.to_original(restored) + for key in original: + torch.testing.assert_close(original_restored[key], original[key], rtol=0, atol=0) + if dtype == torch.float32: + if hasattr(type(model), "from_single_file"): + model.save_config(tmp_path) + loaded = type(model).from_single_file(original, config=str(tmp_path), local_files_only=True) + for key in state: + torch.testing.assert_close(loaded.state_dict()[key], state[key], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "module_name,tester_name,prefix", + [ + ( + "tests.models.unets.test_models_unet_2d_condition", + "UNet2DConditionTesterConfig", + "model.diffusion_model.", + ), + ( + "tests.models.autoencoders.test_models_autoencoder_kl", + "AutoencoderKLTesterConfig", + "first_stage_model.", + ), + ], +) +def test_ldm_component_selection_and_ema(module_name, tester_name, prefix): + from diffusers.loaders import single_file_utils + + tester = getattr(importlib.import_module(module_name), tester_name)() + model = tester.model_class(**tester.get_init_dict()) + state = model.state_dict() + config = dict(model.config, _class_name=type(model).__name__) + conversion = get_conversion(type(model).__name__, config) + original = conversion.to_original(state) + checkpoint = {prefix + key: value for key, value in original.items()} + legacy = single_file_utils.convert_model_checkpoint(checkpoint.copy(), config) + assert set(legacy) == set(state) + for key in state: + torch.testing.assert_close(legacy[key], state[key], rtol=0, atol=0) + if prefix == "model.diffusion_model.": + checkpoint.update( + {"model_ema." + "".join(key.split(".")[1:]): value + 1 for key, value in checkpoint.copy().items()} + ) + expected = conversion.to_diffusers({key: value + 1 for key, value in original.items()}) + selected = single_file_utils.convert_model_checkpoint(checkpoint, config, extract_ema=True) + for key in state: + torch.testing.assert_close(selected[key], expected[key], rtol=0, atol=0) diff --git a/tests/single_file/test_text_encoder_conversions.py b/tests/single_file/test_text_encoder_conversions.py new file mode 100644 index 000000000000..ab56ae4a0c1c --- /dev/null +++ b/tests/single_file/test_text_encoder_conversions.py @@ -0,0 +1,224 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import patch + +import pytest +import torch +import transformers + +from diffusers.loaders.conversion import get_conversion + + +@pytest.mark.parametrize("projection", [False, True]) +def test_pipeline_openclip_loading_uses_shared_conversion(projection): + from diffusers.loaders.single_file_utils import create_diffusers_clip_model_from_ldm + + cls = transformers.CLIPTextModelWithProjection if projection else transformers.CLIPTextModel + config = transformers.CLIPTextConfig( + hidden_size=8, + intermediate_size=16, + num_attention_heads=2, + num_hidden_layers=23 if not projection else 2, + vocab_size=16, + max_position_embeddings=8, + projection_dim=8, + bos_token_id=0, + eos_token_id=1, + ) + model = cls(config) + conversion = get_conversion(cls.__name__, dict(config.to_dict(), original_format="openclip")) + original = conversion.to_original(model.state_dict()) + prefix = "cond_stage_model.model." + checkpoint = {prefix + key: value for key, value in original.items()} + checkpoint[prefix + "logit_scale"] = torch.ones(1) + if not projection: + checkpoint[prefix + "text_projection"] = torch.randn(8, 8) + checkpoint[prefix + "transformer.resblocks.23.attn.in_proj_weight"] = torch.randn(24, 8) + keys = set(checkpoint) + with patch.object(cls.config_class, "from_pretrained", return_value=config): + loaded = create_diffusers_clip_model_from_ldm(cls, checkpoint, config="unused", local_files_only=True) + assert set(checkpoint) == keys + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(loaded.state_dict()[key], tensor, rtol=0, atol=0) + + +CLIP_CONFIG = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "vocab_size": 32, + "max_position_embeddings": 16, + "image_size": 16, + "patch_size": 8, + "projection_dim": 8, +} +CLAP_TEXT_CONFIG = { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "vocab_size": 32, + "max_position_embeddings": 16, + "projection_dim": 8, +} +CLAP_AUDIO_CONFIG = { + "hidden_size": 16, + "patch_embeds_hidden_size": 8, + "depths": [1, 2], + "num_attention_heads": [2, 2], + "spec_size": 16, + "num_mel_bins": 4, + "patch_size": 2, + "patch_stride": [2, 2], + "window_size": 2, + "projection_dim": 8, +} +CASES = [ + (model, "CLIPVisionConfig" if "Vision" in model else "CLIPTextConfig", CLIP_CONFIG, fmt) + for model in ("CLIPTextModel", "CLIPTextModelWithProjection", "CLIPVisionModel", "CLIPVisionModelWithProjection") + for fmt in ("clip", "openclip") +] + [ + ("ContextCLIPTextModel", "CLIPTextConfig", CLIP_CONFIG, "clip"), + ("ContextCLIPTextModel", "CLIPTextConfig", CLIP_CONFIG, "openclip"), + ( + "UMT5EncoderModel", + "UMT5Config", + {"d_model": 16, "d_ff": 32, "d_kv": 8, "num_heads": 2, "num_layers": 2, "vocab_size": 32}, + None, + ), + ( + "SpeechT5HifiGan", + "SpeechT5HifiGanConfig", + { + "model_in_dim": 8, + "upsample_initial_channel": 16, + "upsample_rates": [2, 2], + "upsample_kernel_sizes": [4, 4], + "resblock_kernel_sizes": [3], + "resblock_dilation_sizes": [[1, 3, 5]], + }, + None, + ), + ("ClapTextModelWithProjection", "ClapTextConfig", CLAP_TEXT_CONFIG, None), + ("ClapAudioModel", "ClapAudioConfig", CLAP_AUDIO_CONFIG, None), + ("ClapAudioModelWithProjection", "ClapAudioConfig", CLAP_AUDIO_CONFIG, None), + ( + "ClapModel", + "ClapConfig", + {"text_config": CLAP_TEXT_CONFIG, "audio_config": CLAP_AUDIO_CONFIG, "projection_dim": 8}, + None, + ), + ( + "ClapModel", + "ClapConfig", + { + "text_config": CLAP_TEXT_CONFIG, + "audio_config": {**CLAP_AUDIO_CONFIG, "enable_fusion": True}, + "projection_dim": 8, + }, + None, + ), + ( + "T5EncoderModel", + "T5Config", + {"vocab_size": 32, "d_model": 16, "d_ff": 32, "d_kv": 8, "num_layers": 2, "num_heads": 2}, + None, + ), + ( + "Qwen3Model", + "Qwen3Config", + { + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + }, + None, + ), + ( + "Qwen3ForCausalLM", + "Qwen3Config", + { + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 8, + }, + None, + ), + ( + "Blip2QFormerModel", + "Blip2Config", + { + "qformer_config": { + "vocab_size": 32, + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "encoder_hidden_size": 16, + }, + "vision_config": { + "hidden_size": 16, + "intermediate_size": 32, + "num_hidden_layers": 1, + "num_attention_heads": 2, + "image_size": 16, + "patch_size": 8, + }, + }, + None, + ), +] + + +@pytest.mark.parametrize( + "name,config_name,values,original_format", CASES, ids=[f"{name}-{fmt}" for name, _, _, fmt in CASES] +) +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16], ids=["float32", "float16", "bfloat16"] +) +def test_external_encoder_conversion(name, config_name, values, original_format, dtype): + config = getattr(transformers, config_name)(**values) + if name == "Blip2QFormerModel": + from diffusers.pipelines.deprecated.blip_diffusion.modeling_blip2 import Blip2QFormerModel + + # Tokenization is unused by state-dict conversion; avoid downloading a vocabulary to construct the model. + with patch("diffusers.pipelines.deprecated.blip_diffusion.modeling_blip2.BertTokenizer.from_pretrained"): + model = Blip2QFormerModel(config) + elif name == "ContextCLIPTextModel": + from diffusers.pipelines.deprecated.blip_diffusion.modeling_ctx_clip import ContextCLIPTextModel + + model = ContextCLIPTextModel(config) + else: + model = getattr(transformers, name)(config) + model.to(dtype=dtype) + state = model.state_dict() + config = model.config.to_dict() + if original_format is not None: + config["original_format"] = original_format + conversion = get_conversion(name, config) + assert conversion.diffusers_keys == set(state) + original = conversion.to_original(state) + restored = conversion.to_diffusers(original) + for key in state: + torch.testing.assert_close(restored[key], state[key], rtol=0, atol=0)