github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Phase 2

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 15:37:13 -06:00
parent 9ed2e0b8e7
commit 9113420b61
13 changed files with 799 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
toolkit/models/v2/diffusion_models/wan.py
View File

@@ -8,6 +8,79 @@ class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
aitk_subfolder = "transformer"
aitk_comfy_repo = "Comfy-Org/Wan_2.2_ComfyUI_Repackaged"
# comfy wan files use the original key layout; convert_state_dict_on_load
# renames them (pure substring renames, so the legacy scaled_fp8 weight /
# scale keys ride along with their modules). wan2.2 A14B checkpoints hold
# two DiTs, keyed by (repo, subfolder): transformer = high noise,
# transformer_2 = low noise. wan2.1 entries override the comfy repo.
aitk_comfy_weight_names = {
"Wan-AI/Wan2.2-TI2V-5B-Diffusers": [
"split_files/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors",
],
("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer"): [
"split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
],
("Wan-AI/Wan2.2-T2V-A14B-Diffusers", "transformer_2"): [
"split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
],
("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer"): [
"split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
],
("Wan-AI/Wan2.2-I2V-A14B-Diffusers", "transformer_2"): [
"split_files/diffusion_models/wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
],
"Wan-AI/Wan2.1-T2V-1.3B-Diffusers": {
"repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
"files": [
"split_files/diffusion_models/wan2.1_t2v_1.3B_bf16.safetensors",
"split_files/diffusion_models/wan2.1_t2v_1.3B_fp16.safetensors",
],
},
"Wan-AI/Wan2.1-T2V-14B-Diffusers": {
"repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
"files": [
"split_files/diffusion_models/wan2.1_t2v_14B_fp8_scaled.safetensors",
"split_files/diffusion_models/wan2.1_t2v_14B_bf16.safetensors",
"split_files/diffusion_models/wan2.1_t2v_14B_fp16.safetensors",
],
},
"Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": {
"repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
"files": [
"split_files/diffusion_models/wan2.1_i2v_480p_14B_fp8_scaled.safetensors",
"split_files/diffusion_models/wan2.1_i2v_480p_14B_bf16.safetensors",
"split_files/diffusion_models/wan2.1_i2v_480p_14B_fp16.safetensors",
],
},
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": {
"repo": "Comfy-Org/Wan_2.1_ComfyUI_repackaged",
"files": [
"split_files/diffusion_models/wan2.1_i2v_720p_14B_fp8_scaled.safetensors",
"split_files/diffusion_models/wan2.1_i2v_720p_14B_bf16.safetensors",
"split_files/diffusion_models/wan2.1_i2v_720p_14B_fp16.safetensors",
],
},
}
@classmethod
def get_transformer_block_names(cls):
return ["blocks"]
@classmethod
def convert_state_dict_on_load(cls, state_dict):
# original/comfy wan keys -> diffusers layout via diffusers' own
# converter (rename-only for the base/i2v/vace variants)
is_original = any(
".self_attn." in k
or ".cross_attn." in k
or k.startswith("model.diffusion_model.")
for k in state_dict
)
if not is_original:
return state_dict
from diffusers.loaders.single_file_utils import (
convert_wan_transformer_to_diffusers,
)
return convert_wan_transformer_to_diffusers(dict(state_dict))