diff --git a/example_workflows/Wan2.1_I2V_StartEndFrame_Workflow.json b/example_workflows/Wan2.1_I2V_StartEndFrame_Workflow.json index c1a2c31..0d4386d 100644 --- a/example_workflows/Wan2.1_I2V_StartEndFrame_Workflow.json +++ b/example_workflows/Wan2.1_I2V_StartEndFrame_Workflow.json @@ -1,536 +1,7 @@ { - "last_node_id": 76, - "last_link_id": 85, + "last_node_id": 87, + "last_link_id": 107, "nodes": [ - { - "id": 42, - "type": "Note", - "pos": [ - -580, - -760 - ], - "size": [ - 314.96246337890625, - 152.77333068847656 - ], - "flags": {}, - "order": 0, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "Adjust the blocks to swap based on your VRAM, this is a tradeoff between speed and memory usage.\n\nAlternatively there's option to use VRAM management introduced in DiffSynt-Studios. This is usually slower, but saves even more VRAM compared to BlockSwap" - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 45, - "type": "WanVideoVRAMManagement", - "pos": [ - -210, - -580 - ], - "size": [ - 315, - 58 - ], - "flags": {}, - "order": 1, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "vram_management_args", - "type": "VRAM_MANAGEMENTARGS", - "links": [] - } - ], - "properties": { - "Node name for S&R": "WanVideoVRAMManagement" - }, - "widgets_values": [ - 1 - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 36, - "type": "Note", - "pos": [ - 160, - -1010 - ], - "size": [ - 374.3061828613281, - 171.9547576904297 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "fp8_fast seems to cause huge quality degradation\n\nfp_16_fast enables \"Full FP16 Accmumulation in FP16 GEMMs\" feature available in the very latest pytorch nightly, this is around 20% speed boost. \n\nSageattn if you have it installed can be used for almost double inference speed" - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 33, - "type": "Note", - "pos": [ - 170, - -1150 - ], - "size": [ - 359.0753479003906, - 88 - ], - "flags": {}, - "order": 3, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "Models:\nhttps://huggingface.co/Kijai/WanVideo_comfy/tree/main" - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 60, - "type": "Note", - "pos": [ - -432.5627136230469, - -224.5513458251953 - ], - "size": [ - 253.16725158691406, - 88 - ], - "flags": {}, - "order": 4, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "You can use either the original clip vision or the normal comfyui clip vision loader, they are the same model in the end." - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 52, - "type": "WanVideoTeaCache", - "pos": [ - 1307.6705322265625, - -787.4303588867188 - ], - "size": [ - 315, - 154 - ], - "flags": {}, - "order": 5, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "teacache_args", - "type": "TEACACHEARGS", - "links": [ - 76 - ] - } - ], - "properties": { - "Node name for S&R": "WanVideoTeaCache" - }, - "widgets_values": [ - 0.25, - 1, - -1, - "offload_device", - "true" - ] - }, - { - "id": 54, - "type": "Note", - "pos": [ - 961.6879272460938, - -580.803466796875 - ], - "size": [ - 327.61932373046875, - 88 - ], - "flags": {}, - "order": 6, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "Enhance-a-video can increase the fidelity of the results, too high values lead to noisy results." - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 53, - "type": "Note", - "pos": [ - 960.3718872070312, - -810.77099609375 - ], - "size": [ - 324.64129638671875, - 159.47401428222656 - ], - "flags": {}, - "order": 7, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "TeaCache could be considered to be sort of an automated step skipper \n\nThe relative l1 threshold -value determines how aggressive this is, higher values are faster but quality suffers more. Very first steps should NEVER be skipped with this model or it kills the motion. When using the pre-calculated coefficients, the treshold value should be much higher than with the default coefficients." - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 44, - "type": "Note", - "pos": [ - -620.9041137695312, - -1049.732421875 - ], - "size": [ - 303.0501403808594, - 88 - ], - "flags": {}, - "order": 8, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "If you have Triton installed, connect this for ~30% speed increase" - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 39, - "type": "WanVideoBlockSwap", - "pos": [ - -210, - -760 - ], - "size": [ - 315, - 106 - ], - "flags": {}, - "order": 9, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "block_swap_args", - "type": "BLOCKSWAPARGS", - "links": [ - 50 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoBlockSwap" - }, - "widgets_values": [ - 20, - false, - false - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 38, - "type": "WanVideoVAELoader", - "pos": [ - 169.25408935546875, - -322.9471740722656 - ], - "size": [ - 372.7727966308594, - 82 - ], - "flags": {}, - "order": 10, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "vae", - "type": "WANVAE", - "links": [ - 67, - 71 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoVAELoader" - }, - "widgets_values": [ - "wan_2.1_vae.safetensors", - "bf16" - ], - "color": "#322", - "bgcolor": "#533" - }, - { - "id": 11, - "type": "LoadWanVideoT5TextEncoder", - "pos": [ - 161.7229461669922, - -501.2225036621094 - ], - "size": [ - 377.1661376953125, - 130 - ], - "flags": {}, - "order": 11, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "wan_t5_model", - "type": "WANTEXTENCODER", - "links": [ - 15 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "LoadWanVideoT5TextEncoder" - }, - "widgets_values": [ - "ft_local/umt5-xxl-enc-bf16.safetensors", - "bf16", - "offload_device", - "disabled" - ], - "color": "#332922", - "bgcolor": "#593930" - }, - { - "id": 55, - "type": "WanVideoEnhanceAVideo", - "pos": [ - 1283.0704345703125, - -561.0289306640625 - ], - "size": [ - 315, - 106 - ], - "flags": {}, - "order": 12, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "feta_args", - "type": "FETAARGS", - "links": [ - 75 - ] - } - ], - "properties": { - "Node name for S&R": "WanVideoEnhanceAVideo" - }, - "widgets_values": [ - 2, - 0, - 1 - ] - }, - { - "id": 65, - "type": "WanVideoStartEndDecode", - "pos": [ - 1772.6585693359375, - -387.48602294921875 - ], - "size": [ - 327.5999755859375, - 214 - ], - "flags": {}, - "order": 24, - "mode": 0, - "inputs": [ - { - "name": "vae", - "type": "WANVAE", - "link": 71 - }, - { - "name": "samples", - "type": "LATENT", - "link": 77 - }, - { - "name": "start", - "type": "IMAGE", - "link": 82 - }, - { - "name": "end", - "type": "IMAGE", - "link": 83 - } - ], - "outputs": [ - { - "name": "images", - "type": "IMAGE", - "links": [ - 85 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoStartEndDecode" - }, - "widgets_values": [ - true, - 272, - 272, - 144, - 128 - ] - }, - { - "id": 62, - "type": "Note", - "pos": [ - 315.68389892578125, - 215.8892364501953 - ], - "size": [ - 268.73455810546875, - 90.03050994873047 - ], - "flags": {}, - "order": 13, - "mode": 0, - "inputs": [], - "outputs": [], - "properties": {}, - "widgets_values": [ - "The original code had automatic resolution adjustment based on input image total pixels and aspect ratio. If you want to set it manually, disable the adjust_resolution" - ], - "color": "#432", - "bgcolor": "#653" - }, - { - "id": 35, - "type": "WanVideoTorchCompileSettings", - "pos": [ - -276.8500671386719, - -1050.6326904296875 - ], - "size": [ - 390.5999755859375, - 178 - ], - "flags": {}, - "order": 14, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "torch_compile_args", - "type": "WANCOMPILEARGS", - "links": [], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoTorchCompileSettings" - }, - "widgets_values": [ - "inductor", - false, - "default", - false, - 64, - true - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 72, - "type": "WanVideoLoraSelect", - "pos": [ - 680.5396118164062, - -725.003662109375 - ], - "size": [ - 315, - 102 - ], - "flags": {}, - "order": 15, - "mode": 0, - "inputs": [ - { - "name": "prev_lora", - "type": "WANVIDLORA", - "shape": 7, - "link": null - }, - { - "name": "blocks", - "type": "SELECTEDBLOCKS", - "shape": 7, - "link": null - } - ], - "outputs": [ - { - "name": "lora", - "type": "WANVIDLORA", - "links": null - } - ], - "properties": { - "Node name for S&R": "WanVideoLoraSelect" - }, - "widgets_values": [ - "11heros_threeview_1536_768-000036.safetensors", - 1 - ] - }, { "id": 30, "type": "VHS_VideoCombine", @@ -543,13 +14,13 @@ 453.9615478515625 ], "flags": {}, - "order": 25, + "order": 14, "mode": 0, "inputs": [ { "name": "images", "type": "IMAGE", - "link": 85 + "link": 102 }, { "name": "audio", @@ -598,46 +69,69 @@ "hidden": false, "paused": false, "params": { - "filename": "WanVideoWrapper_I2V_00597.mp4", + "filename": "WanVideoWrapper_I2V_00598.mp4", "subfolder": "", "type": "output", "format": "video/h264-mp4", "frame_rate": 16, - "workflow": "WanVideoWrapper_I2V_00597.png", - "fullpath": "/cfs-cq/sophiejiang/ComfyUI/output/WanVideoWrapper_I2V_00597.mp4" + "workflow": "WanVideoWrapper_I2V_00598.png", + "fullpath": "/cfs-cq/sophiejiang/ComfyUI/output/WanVideoWrapper_I2V_00598.mp4" } } } }, { - "id": 68, - "type": "WanVideoStartEndSampler", + "id": 60, + "type": "Note", "pos": [ - 894.9513549804688, - 135.1397247314453 + -1197.6680908203125, + 96.31451416015625 ], "size": [ - 340.20001220703125, + 253.16725158691406, + 88 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "You can use either the original clip vision or the normal comfyui clip vision loader, they are the same model in the end." + ], + "color": "#432", + "bgcolor": "#653" + }, + { + "id": 81, + "type": "WanVideoStartEndSampler_2frames", + "pos": [ + 926.7401733398438, + 369.2915954589844 + ], + "size": [ + 453.5999755859375, 390 ], "flags": {}, - "order": 23, + "order": 12, "mode": 0, "inputs": [ { "name": "model", "type": "WANVIDEOMODEL", - "link": 73 + "link": 107 }, { "name": "text_embeds", "type": "WANVIDEOTEXTEMBEDS", - "link": 78 + "link": 95 }, { "name": "image_embeds", "type": "WANVIDIMAGE_EMBEDS", - "link": 74 + "link": 91 }, { "name": "samples", @@ -649,7 +143,7 @@ "name": "feta_args", "type": "FETAARGS", "shape": 7, - "link": 75 + "link": 92 }, { "name": "context_options", @@ -661,7 +155,7 @@ "name": "teacache_args", "type": "TEACACHEARGS", "shape": 7, - "link": 76 + "link": 93 }, { "name": "flowedit_args", @@ -675,19 +169,19 @@ "name": "samples", "type": "LATENT", "links": [ - 77 + 97 ], "slot_index": 0 } ], "properties": { - "Node name for S&R": "WanVideoStartEndSampler" + "Node name for S&R": "WanVideoStartEndSampler_2frames" }, "widgets_values": [ 30, 6, 5, - 894520695993030, + 1001789441376328, "randomize", true, "unipc", @@ -699,15 +193,15 @@ "id": 58, "type": "LoadImage", "pos": [ - -1091.252197265625, - 44.508827209472656 + -1756.3026123046875, + 440.4684753417969 ], "size": [ 413.10479736328125, 498.3180847167969 ], "flags": {}, - "order": 16, + "order": 1, "mode": 0, "inputs": [], "outputs": [ @@ -715,8 +209,8 @@ "name": "IMAGE", "type": "IMAGE", "links": [ - 66, - 82 + 87, + 99 ], "slot_index": 0 }, @@ -738,19 +232,397 @@ "color": "#2a363b", "bgcolor": "#3f5159" }, + { + "id": 82, + "type": "WanVideoStartEndDecode_2frames", + "pos": [ + 1786.1859130859375, + -151.71971130371094 + ], + "size": [ + 441, + 214 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "WANVAE", + "link": 104 + }, + { + "name": "samples", + "type": "LATENT", + "link": 97 + }, + { + "name": "start", + "type": "IMAGE", + "link": 99 + }, + { + "name": "end", + "type": "IMAGE", + "link": 100 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 102 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoStartEndDecode_2frames" + }, + "widgets_values": [ + true, + 272, + 272, + 144, + 128 + ] + }, + { + "id": 16, + "type": "WanVideoTextEncode", + "pos": [ + 107.43799591064453, + 305.5315246582031 + ], + "size": [ + 420.30511474609375, + 261.5306701660156 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "t5", + "type": "WANTEXTENCODER", + "link": 15 + }, + { + "name": "model_to_offload", + "type": "WANVIDEOMODEL", + "shape": 7, + "link": null + } + ], + "outputs": [ + { + "name": "text_embeds", + "type": "WANVIDEOTEXTEMBEDS", + "links": [ + 95 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoTextEncode" + }, + "widgets_values": [ + "两个角色搀扶着往前走，并看向彼此", + "色调艳丽，过曝，细节模糊不清，字幕，风格，作品，画作，画面，运动，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", + true + ], + "color": "#332922", + "bgcolor": "#593930" + }, + { + "id": 85, + "type": "LoadWanVideoClipTextEncoder_2frames", + "pos": [ + -829.1338500976562, + 162.15457153320312 + ], + "size": [ + 441, + 106 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "wan_clip_vision", + "type": "CLIP_VISION", + "links": [ + 105 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LoadWanVideoClipTextEncoder_2frames" + }, + "widgets_values": [ + "ft_local/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors", + "fp32", + "offload_device" + ] + }, + { + "id": 83, + "type": "WanVideoVAELoader_2frames", + "pos": [ + 125.07241821289062, + -263.2586975097656 + ], + "size": [ + 352.79998779296875, + 82 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "vae", + "type": "WANVAE", + "links": [ + 103, + 104 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoVAELoader_2frames" + }, + "widgets_values": [ + "wan_2.1_vae.safetensors", + "bf16" + ] + }, + { + "id": 11, + "type": "LoadWanVideoT5TextEncoder", + "pos": [ + -524.9284057617188, + -29.823974609375 + ], + "size": [ + 377.1661376953125, + 130 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "wan_t5_model", + "type": "WANTEXTENCODER", + "links": [ + 15 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LoadWanVideoT5TextEncoder" + }, + "widgets_values": [ + "ft_local/umt5-xxl-enc-bf16.safetensors", + "bf16", + "offload_device", + "disabled" + ], + "color": "#332922", + "bgcolor": "#593930" + }, + { + "id": 39, + "type": "WanVideoBlockSwap", + "pos": [ + -1330.6507568359375, + -418.16363525390625 + ], + "size": [ + 315, + 130 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "block_swap_args", + "type": "BLOCKSWAPARGS", + "links": [ + 106 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoBlockSwap" + }, + "widgets_values": [ + 20, + false, + false, + true + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 87, + "type": "WanVideoModelLoader_2frames", + "pos": [ + -588.1751098632812, + -436.261962890625 + ], + "size": [ + 378, + 214 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "type": "WANCOMPILEARGS", + "shape": 7, + "link": null + }, + { + "name": "block_swap_args", + "type": "BLOCKSWAPARGS", + "shape": 7, + "link": 106 + }, + { + "name": "lora", + "type": "WANVIDLORA", + "shape": 7, + "link": null + }, + { + "name": "vram_management_args", + "type": "VRAM_MANAGEMENTARGS", + "shape": 7, + "link": null + } + ], + "outputs": [ + { + "name": "model", + "type": "WANVIDEOMODEL", + "links": [ + 107 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoModelLoader_2frames" + }, + "widgets_values": [ + "ft_local/wan2.1_i2v_720p_14B_fp16.safetensors", + "fp16", + "disabled", + "offload_device", + "sdpa" + ] + }, + { + "id": 55, + "type": "WanVideoEnhanceAVideo", + "pos": [ + 887.888427734375, + -256.91546630859375 + ], + "size": [ + 315, + 106 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "feta_args", + "type": "FETAARGS", + "links": [ + 92 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoEnhanceAVideo" + }, + "widgets_values": [ + 2, + 0, + 1 + ] + }, + { + "id": 52, + "type": "WanVideoTeaCache", + "pos": [ + 885.9960327148438, + -476.69366455078125 + ], + "size": [ + 315, + 154 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "teacache_args", + "type": "TEACACHEARGS", + "links": [ + 93 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "WanVideoTeaCache" + }, + "widgets_values": [ + 0.25, + 1, + -1, + "offload_device", + "true" + ] + }, { "id": 69, "type": "LoadImage", "pos": [ - -1078.102783203125, - 623.2489013671875 + -1740.4964599609375, + 1026.31640625 ], "size": [ 413.10479736328125, 498.3180847167969 ], "flags": {}, - "order": 17, + "order": 8, "mode": 0, "inputs": [], "outputs": [ @@ -758,8 +630,8 @@ "name": "IMAGE", "type": "IMAGE", "links": [ - 79, - 83 + 88, + 100 ], "slot_index": 0 }, @@ -782,150 +654,39 @@ "bgcolor": "#3f5159" }, { - "id": 16, - "type": "WanVideoTextEncode", + "id": 80, + "type": "WanVideoStartEndImageClipEncode_2frames", "pos": [ - 1362.827880859375, - 74.02191925048828 + -358.740966796875, + 1187.232177734375 ], "size": [ - 420.30511474609375, - 261.5306701660156 - ], - "flags": {}, - "order": 21, - "mode": 0, - "inputs": [ - { - "name": "t5", - "type": "WANTEXTENCODER", - "link": 15 - } - ], - "outputs": [ - { - "name": "text_embeds", - "type": "WANVIDEOTEXTEMBEDS", - "links": [ - 78 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoTextEncode" - }, - "widgets_values": [ - "两个角色搀扶着往前走，并看向彼此", - "色调艳丽，过曝，细节模糊不清，字幕，风格，作品，画作，画面，运动，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", - true - ], - "color": "#332922", - "bgcolor": "#593930" - }, - { - "id": 56, - "type": "LoadWanVideoClipTextEncoder", - "pos": [ - -357.8293762207031, - -83.7756118774414 - ], - "size": [ - 510.6601257324219, - 106 - ], - "flags": {}, - "order": 18, - "mode": 0, - "inputs": [], - "outputs": [ - { - "name": "wan_clip_vision", - "type": "CLIP_VISION", - "links": [ - 65 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "LoadWanVideoClipTextEncoder" - }, - "widgets_values": [ - "ft_local/open-clip-xlm-roberta-large-vit-huge-14_visual_fp32.safetensors", - "fp32", - "offload_device" - ], - "color": "#2a363b", - "bgcolor": "#3f5159" - }, - { - "id": 59, - "type": "CLIPVisionLoader", - "pos": [ - -158.17127990722656, - -210.2847442626953 - ], - "size": [ - 315, - 58 - ], - "flags": {}, - "order": 19, - "mode": 2, - "inputs": [], - "outputs": [ - { - "name": "CLIP_VISION", - "type": "CLIP_VISION", - "links": null - } - ], - "properties": { - "cnr_id": "comfy-core", - "ver": "0.3.18", - "Node name for S&R": "CLIPVisionLoader" - }, - "widgets_values": [ - "model-sdxl.safetensors" - ], - "color": "#2a363b", - "bgcolor": "#3f5159" - }, - { - "id": 66, - "type": "WanVideoStartEndImageClipEncode", - "pos": [ - 437.09112548828125, - -110.37893676757812 - ], - "size": [ - 453.5999755859375, + 567, 286 ], "flags": {}, - "order": 22, + "order": 11, "mode": 0, "inputs": [ { "name": "clip_vision", "type": "CLIP_VISION", - "link": 65 + "link": 105 }, { "name": "start", "type": "IMAGE", - "link": 66 + "link": 87 }, { "name": "end", "type": "IMAGE", - "link": 79 + "link": 88 }, { "name": "vae", "type": "WANVAE", - "link": 67 + "link": 103 } ], "outputs": [ @@ -933,87 +694,24 @@ "name": "image_embeds", "type": "WANVIDIMAGE_EMBEDS", "links": [ - 74 + 91 ], "slot_index": 0 } ], "properties": { - "Node name for S&R": "WanVideoStartEndImageClipEncode" + "Node name for S&R": "WanVideoStartEndImageClipEncode_2frames" }, "widgets_values": [ - 1280, - 720, - 33, + 832, + 480, + 81, true, - 0.3, + 0, 1, 1, true ] - }, - { - "id": 22, - "type": "WanVideoModelLoader", - "pos": [ - -379.0900573730469, - -526.811279296875 - ], - "size": [ - 477.4410095214844, - 226.43276977539062 - ], - "flags": {}, - "order": 20, - "mode": 0, - "inputs": [ - { - "name": "compile_args", - "type": "WANCOMPILEARGS", - "shape": 7, - "link": null - }, - { - "name": "block_swap_args", - "type": "BLOCKSWAPARGS", - "shape": 7, - "link": 50 - }, - { - "name": "lora", - "type": "WANVIDLORA", - "shape": 7, - "link": null - }, - { - "name": "vram_management_args", - "type": "VRAM_MANAGEMENTARGS", - "shape": 7, - "link": null - } - ], - "outputs": [ - { - "name": "model", - "type": "WANVIDEOMODEL", - "links": [ - 73 - ], - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "WanVideoModelLoader" - }, - "widgets_values": [ - "ft_local/wan2.1_i2v_720p_14B_fp16.safetensors", - "fp16", - "disabled", - "offload_device", - "sdpa" - ], - "color": "#223", - "bgcolor": "#335" } ], "links": [ @@ -1026,135 +724,135 @@ "WANTEXTENCODER" ], [ - 50, - 39, - 0, - 22, - 1, - "BLOCKSWAPARGS" - ], - [ - 65, - 56, - 0, - 66, - 0, - "CLIP_VISION" - ], - [ - 66, + 87, 58, 0, - 66, + 80, 1, "IMAGE" ], [ - 67, - 38, + 88, + 69, 0, - 66, - 3, - "WANVAE" + 80, + 2, + "IMAGE" ], [ - 71, - 38, + 91, + 80, 0, - 65, - 0, - "WANVAE" - ], - [ - 73, - 22, - 0, - 68, - 0, - "WANVIDEOMODEL" - ], - [ - 74, - 66, - 0, - 68, + 81, 2, "WANVIDIMAGE_EMBEDS" ], [ - 75, + 92, 55, 0, - 68, + 81, 4, "FETAARGS" ], [ - 76, + 93, 52, 0, - 68, + 81, 6, "TEACACHEARGS" ], [ - 77, - 68, - 0, - 65, - 1, - "LATENT" - ], - [ - 78, + 95, 16, 0, - 68, + 81, 1, "WANVIDEOTEXTEMBEDS" ], [ - 79, - 69, + 97, + 81, 0, - 66, - 2, - "IMAGE" + 82, + 1, + "LATENT" ], [ - 82, + 99, 58, 0, - 65, + 82, 2, "IMAGE" ], [ - 83, + 100, 69, 0, - 65, + 82, 3, "IMAGE" ], [ - 85, - 65, + 102, + 82, 0, 30, 0, "IMAGE" + ], + [ + 103, + 83, + 0, + 80, + 3, + "WANVAE" + ], + [ + 104, + 83, + 0, + 82, + 0, + "WANVAE" + ], + [ + 105, + 85, + 0, + 80, + 0, + "CLIP_VISION" + ], + [ + 106, + 39, + 0, + 87, + 1, + "BLOCKSWAPARGS" + ], + [ + 107, + 87, + 0, + 81, + 0, + "WANVIDEOMODEL" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 0.7972024500000027, - "offset": [ - 573.1721782843497, - 641.3784590059872 - ] + "scale": 3.0273749771965517, + "offset": { + "0": -1524.655517578125, + "1": -825.6109008789062 + } }, "node_versions": { "ComfyUI-WanVideoWrapper": "721cd65e7b5224c70a3d20446d9d561f1732216b", diff --git a/nodes.py b/nodes.py index a0b06ee..b554260 100644 --- a/nodes.py +++ b/nodes.py @@ -35,7 +35,7 @@ def add_noise_to_reference_video(image, ratio=None): image = image + image_noise return image -class WanVideoBlockSwap: +class WanVideoBlockSwap_2frames: @classmethod def INPUT_TYPES(s): return { @@ -54,7 +54,7 @@ class WanVideoBlockSwap: def setargs(self, **kwargs): return (kwargs, ) -class WanVideoVRAMManagement: +class WanVideoVRAMManagement_2frames: @classmethod def INPUT_TYPES(s): return { @@ -71,7 +71,7 @@ class WanVideoVRAMManagement: def setargs(self, **kwargs): return (kwargs, ) -class WanVideoTeaCache: +class WanVideoTeaCache_2frames: @classmethod def INPUT_TYPES(s): return { @@ -127,7 +127,7 @@ Official recommended values https://github.com/ali-vilab/TeaCache/tree/main/TeaC return (teacache_args,) -class WanVideoModel(comfy.model_base.BaseModel): +class WanVideoModel_2frames(comfy.model_base.BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.pipeline = {} @@ -146,7 +146,7 @@ except: #for backwards compatibility from comfy.latent_formats import HunyuanVideo latent_format = HunyuanVideo -class WanVideoModelConfig: +class WanVideoModelConfig_2frames: def __init__(self, dtype): self.unet_config = {} self.unet_extra_config = {} @@ -204,7 +204,7 @@ def standardize_lora_key_format(lora_sd): new_sd[k] = v return new_sd -class WanVideoEnhanceAVideo: +class WanVideoEnhanceAVideo_2frames: @classmethod def INPUT_TYPES(s): return { @@ -223,7 +223,7 @@ class WanVideoEnhanceAVideo: def setargs(self, **kwargs): return (kwargs, ) -class WanVideoLoraSelect: +class WanVideoLoraSelect_2frames: @classmethod def INPUT_TYPES(s): return { @@ -259,7 +259,7 @@ class WanVideoLoraSelect: loras_list.append(lora) return (loras_list,) -class WanVideoLoraBlockEdit: +class WanVideoLoraBlockEdit_2frames: def __init__(self): self.loaded_lora = None @@ -286,7 +286,7 @@ class WanVideoLoraBlockEdit: return (selected_blocks,) #region Model loading -class WanVideoModelLoader: +class WanVideoModelLoader_2frames: @classmethod def INPUT_TYPES(s): return { @@ -407,8 +407,8 @@ class WanVideoModelLoader: transformer = WanModel(**TRANSFORMER_CONFIG) transformer.eval() - comfy_model = WanVideoModel( - WanVideoModelConfig(base_dtype), + comfy_model = WanVideoModel_2frames( + WanVideoModelConfig_2frames(base_dtype), model_type=comfy.model_base.ModelType.FLOW, device=device, ) @@ -586,7 +586,7 @@ class WanVideoModelLoader: #region load VAE -class WanVideoVAELoader: +class WanVideoVAELoader_2frames: @classmethod def INPUT_TYPES(s): return { @@ -607,7 +607,7 @@ class WanVideoVAELoader: DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'" def loadmodel(self, model_name, precision): - from .wanvideo.wan_video_vae import WanVideoVAE + from .wanvideo.wan_video_vae_2frames import WanVideoVAE device = mm.get_torch_device() offload_device = mm.unet_offload_device() @@ -632,7 +632,7 @@ class WanVideoVAELoader: -class WanVideoTorchCompileSettings: +class WanVideoTorchCompileSettings_2frames: @classmethod def INPUT_TYPES(s): return { @@ -667,7 +667,7 @@ class WanVideoTorchCompileSettings: #region TextEncode -class LoadWanVideoT5TextEncoder: +class LoadWanVideoT5TextEncoder_2frames: @classmethod def INPUT_TYPES(s): return { @@ -718,7 +718,7 @@ class LoadWanVideoT5TextEncoder: return (text_encoder,) -class LoadWanVideoClipTextEncoder: +class LoadWanVideoClipTextEncoder_2frames: @classmethod def INPUT_TYPES(s): return { @@ -757,7 +757,7 @@ class LoadWanVideoClipTextEncoder: return (clip_model,) -class WanVideoTextEncode: +class WanVideoTextEncode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -806,7 +806,7 @@ class WanVideoTextEncode: } return (prompt_embeds_dict,) -class WanVideoTextEmbedBridge: +class WanVideoTextEmbedBridge_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -830,7 +830,7 @@ class WanVideoTextEmbedBridge: return (prompt_embeds_dict,) #region clip image encode -class WanVideoImageClipEncode: +class WanVideoImageClipEncode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -959,7 +959,7 @@ class WanVideoImageClipEncode: return (image_embeds,) -class WanVideoStartEndImageClipEncode: +class WanVideoStartEndImageClipEncode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -1102,7 +1102,7 @@ class WanVideoStartEndImageClipEncode: -class WanVideoEmptyEmbeds: +class WanVideoEmptyEmbeds_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -1141,7 +1141,7 @@ class WanVideoEmptyEmbeds: #region Sampler -class WanVideoContextOptions: +class WanVideoContextOptions_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -1172,7 +1172,7 @@ class WanVideoContextOptions: return (context_options,) -class WanVideoFlowEdit: +class WanVideoFlowEdit_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -1197,7 +1197,7 @@ class WanVideoFlowEdit: def process(self, **kwargs): return (kwargs,) -class WanVideoSampler: +class WanVideoSampler_2frames: @classmethod def INPUT_TYPES(s): return { @@ -1347,7 +1347,7 @@ class WanVideoSampler: context_stride = context_options["context_stride"] // 4 context_overlap = context_options["context_overlap"] // 4 - self.window_tracker = WindowTracker(verbose=context_options["verbose"]) + self.window_tracker = WindowTracker_2frames(verbose=context_options["verbose"]) # Get total number of prompts num_prompts = len(text_embeds["prompt_embeds"]) @@ -1821,7 +1821,7 @@ class WanVideoSampler: }, ) -class WanVideoStartEndSampler: +class WanVideoStartEndSampler_2frames: @classmethod def INPUT_TYPES(s): return { @@ -1972,7 +1972,7 @@ class WanVideoStartEndSampler: context_stride = context_options["context_stride"] // 4 context_overlap = context_options["context_overlap"] // 4 - self.window_tracker = WindowTracker(verbose=context_options["verbose"]) + self.window_tracker = WindowTracker_2frames(verbose=context_options["verbose"]) # Get total number of prompts num_prompts = len(text_embeds["prompt_embeds"]) @@ -2460,7 +2460,7 @@ class WanVideoStartEndSampler: return ({ "samples": x0.unsqueeze(0).cpu() },) -class WindowTracker: +class WindowTracker_2frames: def __init__(self, verbose=False): self.window_map = {} # Maps frame sequence to persistent ID self.next_id = 0 @@ -2484,7 +2484,7 @@ class WindowTracker: return self.teacache_states[window_id] #region VideoDecode -class WanVideoDecode: +class WanVideoDecode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -2527,7 +2527,7 @@ class WanVideoDecode: return (image,) -class WanVideoStartEndDecode: +class WanVideoStartEndDecode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -2580,7 +2580,7 @@ class WanVideoStartEndDecode: #region VideoEncode -class WanVideoEncode: +class WanVideoEncode_2frames: @classmethod def INPUT_TYPES(s): return {"required": { @@ -2650,7 +2650,7 @@ class WanVideoEncode: return ({"samples": latents, "mask": mask},) -class WanVideoLatentPreview: +class WanVideoLatentPreview_2frames: @classmethod def INPUT_TYPES(s): return { @@ -2730,50 +2730,50 @@ class WanVideoLatentPreview: return (latent_images.float().cpu(), out_factors) NODE_CLASS_MAPPINGS = { - "WanVideoTextEncode": WanVideoTextEncode, - "WanVideoModelLoader": WanVideoModelLoader, - "WanVideoVAELoader": WanVideoVAELoader, - "LoadWanVideoT5TextEncoder": LoadWanVideoT5TextEncoder, - "WanVideoImageClipEncode": WanVideoImageClipEncode, - "LoadWanVideoClipTextEncoder": LoadWanVideoClipTextEncoder, - "WanVideoEncode": WanVideoEncode, - "WanVideoBlockSwap": WanVideoBlockSwap, - "WanVideoTorchCompileSettings": WanVideoTorchCompileSettings, - "WanVideoLatentPreview": WanVideoLatentPreview, - "WanVideoEmptyEmbeds": WanVideoEmptyEmbeds, - "WanVideoLoraSelect": WanVideoLoraSelect, - "WanVideoLoraBlockEdit": WanVideoLoraBlockEdit, - "WanVideoEnhanceAVideo": WanVideoEnhanceAVideo, - "WanVideoContextOptions": WanVideoContextOptions, - "WanVideoTeaCache": WanVideoTeaCache, - "WanVideoVRAMManagement": WanVideoVRAMManagement, - "WanVideoTextEmbedBridge": WanVideoTextEmbedBridge, - "WanVideoFlowEdit": WanVideoFlowEdit, - "WanVideoStartEndImageClipEncode":WanVideoStartEndImageClipEncode, - "WanVideoStartEndDecode": WanVideoStartEndDecode, - "WanVideoStartEndSampler":WanVideoStartEndSampler, + "WanVideoTextEncode_2frames": WanVideoTextEncode_2frames, + "WanVideoModelLoader_2frames": WanVideoModelLoader_2frames, + "WanVideoVAELoader_2frames": WanVideoVAELoader_2frames, + "LoadWanVideoT5TextEncoder_2frames": LoadWanVideoT5TextEncoder_2frames, + "WanVideoImageClipEncode_2frames": WanVideoImageClipEncode_2frames, + "LoadWanVideoClipTextEncoder_2frames": LoadWanVideoClipTextEncoder_2frames, + "WanVideoEncode_2frames": WanVideoEncode_2frames, + "WanVideoBlockSwap_2frames": WanVideoBlockSwap_2frames, + "WanVideoTorchCompileSettings_2frames": WanVideoTorchCompileSettings_2frames, + "WanVideoLatentPreview_2frames": WanVideoLatentPreview_2frames, + "WanVideoEmptyEmbeds_2frames": WanVideoEmptyEmbeds_2frames, + "WanVideoLoraSelect_2frames": WanVideoLoraSelect_2frames, + "WanVideoLoraBlockEdit_2frames": WanVideoLoraBlockEdit_2frames, + "WanVideoEnhanceAVideo_2frames": WanVideoEnhanceAVideo_2frames, + "WanVideoContextOptions_2frames": WanVideoContextOptions_2frames, + "WanVideoTeaCache_2frames": WanVideoTeaCache_2frames, + "WanVideoVRAMManagement_2frames": WanVideoVRAMManagement_2frames, + "WanVideoTextEmbedBridge_2frames": WanVideoTextEmbedBridge_2frames, + "WanVideoFlowEdit_2frames": WanVideoFlowEdit_2frames, + "WanVideoStartEndImageClipEncode_2frames":WanVideoStartEndImageClipEncode_2frames, + "WanVideoStartEndDecode_2frames": WanVideoStartEndDecode_2frames, + "WanVideoStartEndSampler_2frames":WanVideoStartEndSampler_2frames, } NODE_DISPLAY_NAME_MAPPINGS = { - "WanVideoTextEncode": "WanVideo TextEncode(2Frames)", - "WanVideoModelLoader": "WanVideo Model Loader(2Frames)", - "WanVideoVAELoader": "WanVideo VAE Loader(2Frames)", - "LoadWanVideoT5TextEncoder": "Load WanVideo T5 TextEncoder(2Frames)", - "WanVideoImageClipEncode": "WanVideo ImageClip Encode(2Frames)", - "LoadWanVideoClipTextEncoder": "Load WanVideo Clip Encoder(2Frames)", - "WanVideoEncode": "WanVideo Encode(2Frames)", - "WanVideoBlockSwap": "WanVideo BlockSwap(2Frames)", - "WanVideoTorchCompileSettings": "WanVideo Torch Compile Settings(2Frames)", - "WanVideoLatentPreview": "WanVideo Latent Preview(2Frames)", - "WanVideoEmptyEmbeds": "WanVideo Empty Embeds(2Frames)", - "WanVideoLoraSelect": "WanVideo Lora Select(2Frames)", - "WanVideoLoraBlockEdit": "WanVideo Lora Block Edit(2Frames)", - "WanVideoEnhanceAVideo": "WanVideo Enhance-A-Video(2Frames)", - "WanVideoContextOptions": "WanVideo Context Options(2Frames)", - "WanVideoTeaCache": "WanVideo TeaCache(2Frames)", - "WanVideoVRAMManagement": "WanVideo VRAM Management(2Frames)", - "WanVideoTextEmbedBridge": "WanVideo TextEmbed Bridge(2Frames)", - "WanVideoFlowEdit": "WanVideo FlowEdit(2Frames)", - "WanVideoStartEndImageClipEncode":"WanVideo Start End Image Clip Encode(2Frames)", - "WanVideoStartEndDecode": "Wan Video Start End Decode(2Frames)", - "WanVideoStartEndSampler": "Wan Video Start End Sampler(2Frames)", + "WanVideoTextEncode_2frames": "WanVideo TextEncode(2Frames)", + "WanVideoModelLoader_2frames": "WanVideo Model Loader(2Frames)", + "WanVideoVAELoader_2frames": "WanVideo VAE Loader(2Frames)", + "LoadWanVideoT5TextEncoder_2frames": "Load WanVideo T5 TextEncoder(2Frames)", + "WanVideoImageClipEncode_2frames": "WanVideo ImageClip Encode(2Frames)", + "LoadWanVideoClipTextEncoder_2frames": "Load WanVideo Clip Encoder(2Frames)", + "WanVideoEncode_2frames": "WanVideo Encode(2Frames)", + "WanVideoBlockSwap_2frames": "WanVideo BlockSwap(2Frames)", + "WanVideoTorchCompileSettings_2frames": "WanVideo Torch Compile Settings(2Frames)", + "WanVideoLatentPreview_2frames": "WanVideo Latent Preview(2Frames)", + "WanVideoEmptyEmbeds_2frames": "WanVideo Empty Embeds(2Frames)", + "WanVideoLoraSelect_2frames": "WanVideo Lora Select(2Frames)", + "WanVideoLoraBlockEdit_2frames": "WanVideo Lora Block Edit(2Frames)", + "WanVideoEnhanceAVideo_2frames": "WanVideo Enhance-A-Video(2Frames)", + "WanVideoContextOptions_2frames": "WanVideo Context Options(2Frames)", + "WanVideoTeaCache_2frames": "WanVideo TeaCache(2Frames)", + "WanVideoVRAMManagement_2frames": "WanVideo VRAM Management(2Frames)", + "WanVideoTextEmbedBridge_2frames": "WanVideo TextEmbed Bridge(2Frames)", + "WanVideoFlowEdit_2frames": "WanVideo FlowEdit(2Frames)", + "WanVideoStartEndImageClipEncode_2frames":"WanVideo Start End Image Clip Encode(2Frames)", + "WanVideoStartEndDecode_2frames": "Wan Video Start End Decode(2Frames)", + "WanVideoStartEndSampler_2frames": "Wan Video Start End Sampler(2Frames)", } diff --git a/wanvideo/wan_video_vae_2frames.py b/wanvideo/wan_video_vae_2frames.py new file mode 100644 index 0000000..ae5c438 --- /dev/null +++ b/wanvideo/wan_video_vae_2frames.py @@ -0,0 +1,903 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _, device = *x.size(), x.device + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if check_is_instance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0,): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + def forward(self, x): + mu, log_var = self.encode(x) + z = self.reparameterize(mu, log_var) + x_recon = self.decode(z) + return x_recon, mu, log_var + + + def encode_2(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 2 + (t - 2) // 4 + + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + elif i== iter_-1: + out_ = self.encoder(x[:, :, -1:, :, :], + feat_cache=[None] * self._enc_conv_num, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + else: + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + out_head = out[:, :, :iter_ - 1, :, :] + out_tail = out[:, :, -1, :, :].unsqueeze(2) + mu, log_var = torch.cat([self.conv1(out_head), self.conv1(out_tail)], dim=2).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + + def encode(self, x, scale): + self.clear_cache() + ## cache + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + for i in range(iter_): + self._enc_conv_idx = [0] + if i == 0: + out = self.encoder(x[:, :, :1, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + else: + out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=self._enc_feat_map, + feat_idx=self._enc_conv_idx) + out = torch.cat([out, out_], 2) + mu, log_var = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + + + def decode_2(self, z, scale): + print("decode_2222222222") + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + z_head=z[:,:,:-1,:,:] + z_tail=z[:,:,-1,:,:].unsqueeze(2) + x = torch.cat([self.conv2(z_head), self.conv2(z_tail)], dim=2) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + + elif i==iter_-1: + out_ = self.decoder(x[:, :, -1, :, :], + feat_cache=[None] * self._enc_conv_num, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) # may add tensor offload + else: + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) # may add tensor offload + print(out.shape,'out') + return out + + + + def decode(self, z, scale): + self.clear_cache() + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + for i in range(iter_): + self._conv_idx = [0] + if i == 0: + out = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + else: + out_ = self.decoder(x[:, :, i:i + 1, :, :], + feat_cache=self._feat_map, + feat_idx=self._conv_idx) + out = torch.cat([out, out_], 2) # may add tensor offload + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + def clear_cache(self): + self._conv_num = count_conv3d(self.decoder) + self._conv_idx = [0] + self._feat_map = [None] * self._conv_num + # cache encode + self._enc_conv_num = count_conv3d(self.encoder) + self._enc_conv_idx = [0] + self._enc_feat_map = [None] * self._enc_conv_num + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16, dtype=torch.float32): + super().__init__() + + self.dtype = dtype + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean) + self.std = torch.tensor(std) + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + + + def build_1d_mask(self, length, left_bound, right_bound, border_width): + x = torch.ones((length,)) + if not left_bound: + x[:border_width] = (torch.arange(border_width) + 1) / border_width + if not right_bound: + x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0]) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, device, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = T * 4 - 3 + weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype, device=data_device) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.float().clamp_(-1, 1) + return values + + + def tiled_encode(self, video, device, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + data_device = "cpu" + computation_device = device + + out_T = (T + 3) // 4 + weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + values = torch.zeros((1, 16, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype, device=data_device) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.float() + return values + + + def single_encode(self, video, device): + video = video.to(device) + x = self.model.encode(video, self.scale) + return x.float() + + + def single_decode(self, hidden_state, device): + print("single_decode") + hidden_state = hidden_state.to(device) + video = self.model.decode(hidden_state, self.scale) + return video.float().clamp_(-1, 1) + + def double_encode(self, video, device): + print('double_encode') + video = video.to(device) + x = self.model.encode_2(video, self.scale) + return x.float() + + + def double_decode(self, hidden_state, device): + print('double_decode') + hidden_state = hidden_state.to(device) + video = self.model.decode_2(hidden_state, self.scale) + return video.float().clamp_(-1, 1) + + def encode(self, videos, device, tiled=False,end_=False, tile_size=(34, 34), tile_stride=(18, 16)): + + videos = [video.to("cpu") for video in videos] + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0], tile_size[1]) + tile_stride = (tile_stride[0], tile_stride[1]) + hidden_state = self.tiled_encode(video, device, tile_size, tile_stride) + else: + if end_: + hidden_state = self.double_encode(video, device) + else: + hidden_state = self.single_encode(video, device) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + + def decode(self, hidden_states, device, tiled=False, end_=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states] + videos = [] + for hidden_state in hidden_states: + hidden_state = hidden_state.unsqueeze(0) + if tiled: + video = self.tiled_decode(hidden_state, device, tile_size, tile_stride) + else: + if end_: + video = self.double_decode(hidden_state, device) + else: + video = self.single_decode(hidden_state, device) + video = video.squeeze(0) + videos.append(video) + return videos + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_