205 lines
9.1 KiB
Python
205 lines
9.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from diffusers import WanTransformer3DModel as DiffusersWanTransformer3DModel
|
|
|
|
from .._mixin import OstrisModelMixin
|
|
|
|
|
|
def patch_wan_patch_embedding(transformer) -> bool:
|
|
"""Wan's patch_embedding is a Conv3d whose kernel == stride (non-overlapping
|
|
patchify), i.e. a plain linear projection of each flattened patch. bf16
|
|
Conv3d has no fast cuDNN kernel and falls back to a slow, GPU-underutilizing
|
|
path — swap the forward for the exactly-equivalent F.linear (a GEMM),
|
|
returning the conv's own [b, out, t', h', w'] layout so it stays a drop-in.
|
|
The weight is read lazily so this survives later .to(device)/dtype moves."""
|
|
conv = getattr(transformer, "patch_embedding", None)
|
|
if not isinstance(conv, torch.nn.Conv3d):
|
|
return False
|
|
if (
|
|
tuple(conv.kernel_size) != tuple(conv.stride)
|
|
or any(p != 0 for p in conv.padding)
|
|
or any(d != 1 for d in conv.dilation)
|
|
or conv.groups != 1
|
|
):
|
|
return False
|
|
|
|
def fast_forward(x, _conv=conv):
|
|
weight = _conv.weight
|
|
out_c = weight.shape[0]
|
|
kt, kh, kw = _conv.kernel_size
|
|
b, c, t, h, w = x.shape
|
|
tp, hp, wp = t // kt, h // kh, w // kw
|
|
# patchify with (c, kt, kh, kw) fastest — matches weight.reshape below
|
|
x = x.reshape(b, c, tp, kt, hp, kh, wp, kw)
|
|
x = x.permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(b, tp * hp * wp, -1)
|
|
w2d = weight.reshape(out_c, -1)
|
|
x = F.linear(x.to(w2d.dtype), w2d, _conv.bias)
|
|
return x.transpose(1, 2).reshape(b, out_c, tp, hp, wp)
|
|
|
|
conv.forward = fast_forward
|
|
return True
|
|
|
|
|
|
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",
|
|
],
|
|
# the UI defaults point at the ai-toolkit bf16 repacks; same comfy files
|
|
("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer"): [
|
|
"split_files/diffusion_models/wan2.2_t2v_high_noise_14B_fp8_scaled.safetensors",
|
|
],
|
|
("ai-toolkit/Wan2.2-T2V-A14B-Diffusers-bf16", "transformer_2"): [
|
|
"split_files/diffusion_models/wan2.2_t2v_low_noise_14B_fp8_scaled.safetensors",
|
|
],
|
|
("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "transformer"): [
|
|
"split_files/diffusion_models/wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
|
|
],
|
|
("ai-toolkit/Wan2.2-I2V-A14B-Diffusers-bf16", "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",
|
|
],
|
|
},
|
|
}
|
|
|
|
# the comfy fp8/scaled_fp8 wan files mix fp16 conv biases with fp32
|
|
# tables; wan's forward assumes one uniform dtype, so cast the
|
|
# non-quantized tensors to the load dtype (the old holder did this with a
|
|
# blanket .to(device, dtype) after load)
|
|
aitk_cast_quantized_load = True
|
|
|
|
@classmethod
|
|
def get_transformer_block_names(cls):
|
|
return ["blocks"]
|
|
|
|
def aitk_post_load(self, **kwargs):
|
|
# slow bf16 Conv3d patchify -> equivalent GEMM, before quantize/offload
|
|
patch_wan_patch_embedding(self)
|
|
return super().aitk_post_load(**kwargs)
|
|
|
|
def get_offload_ignore_modules(self):
|
|
# fp32 modulation tables must stay resident on the compute device
|
|
return [self.scale_shift_table] + [
|
|
block.scale_shift_table for block in self.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))
|
|
|
|
# inverse of diffusers' rename table, for the base/t2v/i2v variants (no
|
|
# vace/animate: their extra mappings collide in reverse). Order matters:
|
|
# longest/most-specific first, and the norm2/norm3 swap uses a placeholder.
|
|
_SAVE_RENAMES = [
|
|
("condition_embedder.time_embedder.linear_1", "time_embedding.0"),
|
|
("condition_embedder.time_embedder.linear_2", "time_embedding.2"),
|
|
("condition_embedder.text_embedder.linear_1", "text_embedding.0"),
|
|
("condition_embedder.text_embedder.linear_2", "text_embedding.2"),
|
|
("condition_embedder.time_proj", "time_projection.1"),
|
|
("condition_embedder.image_embedder.norm1", "img_emb.proj.0"),
|
|
("condition_embedder.image_embedder.ff.net.0.proj", "img_emb.proj.1"),
|
|
("condition_embedder.image_embedder.ff.net.2", "img_emb.proj.3"),
|
|
("condition_embedder.image_embedder.norm2", "img_emb.proj.4"),
|
|
("ffn.net.0.proj", "ffn.0"),
|
|
("ffn.net.2", "ffn.2"),
|
|
(".norm_added_k.", ".norm_k_img."),
|
|
(".add_k_proj.", ".k_img."),
|
|
(".add_v_proj.", ".v_img."),
|
|
(".to_out.0.", ".o."),
|
|
(".to_q.", ".q."),
|
|
(".to_k.", ".k."),
|
|
(".to_v.", ".v."),
|
|
("attn2", "cross_attn"),
|
|
("attn1", "self_attn"),
|
|
# norm2 <-> norm3 swap back
|
|
("norm3", "norm__placeholder"),
|
|
("norm2", "norm3"),
|
|
("norm__placeholder", "norm2"),
|
|
("proj_out", "head.head"),
|
|
]
|
|
|
|
@classmethod
|
|
def convert_state_dict_on_save(cls, state_dict):
|
|
"""Diffusers layout back to the original/comfy wan key layout
|
|
(rename-only, so quantized weight/scale/marker keys ride along)."""
|
|
if any(".self_attn." in k or ".cross_attn." in k for k in state_dict):
|
|
return state_dict # already original layout
|
|
new_sd = {}
|
|
for key, value in state_dict.items():
|
|
k = key
|
|
# scale_shift_table: blocks keep the name as `modulation`, the
|
|
# top-level one belongs to the output head
|
|
if k == "scale_shift_table":
|
|
k = "head.modulation"
|
|
elif k.endswith(".scale_shift_table"):
|
|
k = k[: -len("scale_shift_table")] + "modulation"
|
|
for src, dst in cls._SAVE_RENAMES:
|
|
k = k.replace(src, dst)
|
|
new_sd[k] = value
|
|
return new_sd
|