Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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 | |
| ] | |
| 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"), | |
| ] | |
| 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 | |