87 lines
3.8 KiB
Python
87 lines
3.8 KiB
Python
from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
class WanTransformer3DModel(DiffusersWanTransformer3DModel, OstrisModelMixin):
|
|
"""Wan 2.1/2.2 video DiT (wan22 loads two of these into its dual wrapper)."""
|
|
|
|
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))
|