Feat; Add support for Wan/Qwen TAEHV decoding #937
Conversation
|
Sorry for the unrelated whitespace changes and the debug spam, will fix later |
|
Oh a new version of the taew2.1 weights just came out, coincidentally.
|
|
Now tae decoding for the outputs of Wan2.1 models (and Wan2.2 A14B) works in txt2img mode. Video decoding is running as well, but the results are obviously incorrect (flashing lights warning) If someone can see what I'm doing wrong when decoding videos, let me know. |
|
After fixing the three bugs mentioned in review, image results look correct (tested on GH200 with diffsdiff --git a/tae.hpp b/tae.hpp
index ad0bd37..6a7951f 100644
--- a/tae.hpp
+++ b/tae.hpp
@@ -224,7 +224,7 @@ public:
h = conv1->forward(ctx, h);
h = ggml_relu_inplace(ctx->ggml_ctx, h);
h = conv2->forward(ctx, h);
- h = ggml_relu_inplace(ctx->ggml_ctx, h);
+ // h = ggml_relu_inplace(ctx->ggml_ctx, h);
auto skip = x;
if (has_skip_conv) {
@@ -323,7 +323,7 @@ public:
for (int i = 0; i < num_layers; i++) {
for (int j = 0; j < num_blocks; j++) {
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
- auto mem = ggml_pad(ctx->ggml_ctx, h, 0, 0, 0, 1);
+ auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
h = block->forward(ctx, h, mem);
}
@@ -341,7 +341,7 @@ public:
h = last_conv->forward(ctx, h);
// shape(W, H, 3, T+3) => shape(W, H, 3, T)
- h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - 3, h->nb[1], h->nb[2], h->nb[3], 0);
+ h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - 3, h->nb[1], h->nb[2], h->nb[3], 3*h->nb[3]);
return h; |
|
Video is still completely broken, but image decoding works very well now. |
|
Results for taew2.2 are quite interesting for now. output.mp4 |
|
Wan 2.2 and Hunyuan 1.5 have 2x2 pixelshuffle on the input/output |
|
@madebyollin Yes I saw that when looking at the VAE code in |
output.mp4 |
output.mp4 |
|
The Wan 2.2 TI2V results still look broken. There's a scaling issue on ~L3600 where output_with_disabled_process_latent.mp4This gray-mush issue happens with the default VAE on output_with_official_vae_on_8f05f5bc6ee9d6aba9d1ff2be7739a5a3cf1586d.mp4 |
Co-authored-by: Ollin Boer Bohan <madebyollin@gmail.com>
Co-authored-by: Ollin Boer Bohan <madebyollin@gmail.com>
|
@madebyollin yes I figured it was probably something that after noticing how much worse the img2vid results were compared to txt2vid. I get no "gray-mush" on my end with this fix though. |
|
@stduhpf I used taehv and get results very close to the results of wan vae. Maybe this PR can be merged now? |
|
I think so too. I haven't tested every possible use case though (for example VACE). |
I tested I2V with this image and successfully enabled the TAE feature on AMD RX 7800 XT using my ROCm 6.4.2 build. Thanks for implementing it! ❤️
20251216_005604.mp4 |
|
Thank you for your contributions. |
|
How is the VRAM requirement reduction? |
Reesults on ROCm 6.2, Windows, with AMD Radeon RX 6800 (gfx1030) 16GB Without
|
|
Thanks! I'm using wan2.2 5b, and the default vae requires 30GiB VRAM, so now generating video with 16GB RAM becomes possible |
|
@madebyollin is this expected for taew2_1 to be slower and save less VRAM than taew2_2? |
|
It seems that this implementation has a maximum VRAM. When I used longer video (960x720x81), an error occured: |
|
@CarlGao4 Was this with |
|
@stduhpf For a given image size, TAEW2.2 should be ~1/4 the cost of TAEW2.1, because TAEW2.2 uses 2x2 patchification which halves the height and width of all intermediate tensors. The cost of 2x2 patchification is a reduced ability to represent fine details. I pushed a new checkpoint today that should improve TAEW2.2 quality slightly, but it's still less sharp/stable than TAEW2.1 in my testing. |
|
@stduhpf With |
Hard (moving) edges are now noticeably less pixelated/squiggley, especially every non-first frame. Second frame:
|
|
@madebyollin Can you update the safetensors as well? |
|
@CarlGao4 Good point, updated the safetensors and added |
6 participants
https://github.com/madebyollin/taehv
Model weights:
.\bin\Release\sd.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\qwen-image-Q8_0.gguf --vae ..\..\ComfyUI\models\vae\qwen_image_vae.safetensors --qwen2vl ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct-Q8_0.gguf -p '一个穿着"QWEN"标志的T恤的中国美女正拿着黑色的马克笔面相镜头微笑。她身后的玻璃板上手写体写着 “一、Qwen-Image的技术路线： 探索视觉生成基础模型的极限，开创理解与生成一体化的未来。二、Qwen-Image的模型特色：1、复杂文字渲染。支持中英渲染、自动布局； 2、精准图像编辑。支持文字编辑、物体增减、风格变换。三、Qwen-Image的未来愿景：赋能专业内容创作、助力生成式AI发展。”' --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu -H 1024 -W 1024 --diffusion-fa --flow-shift 3 --tae ..\ComfyUI\models\vae_approx\taew2_1.pth --vae-conv-direct.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model '..\..\ComfyUI\models\unet\Wan2.2-TI2V-5B-Q8_0.gguf' --t5xxl ..\..\ComfyUI\models\clip\t5\umt5-xxl-encoder-Q8_0.gguf --tae ..\..\ComfyUI\models\vae_approx\taew2_2.pth -p "The woman drops the marker, and then she starts laughing a bit" -n "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" --cfg-scale 5.0 --sampling-method euler -v -W 768 -H 768 --color --video-frames 49 -i .\image.png --vae-conv-direct --scheduler smoothstep --steps 50 --fps 24 --diffusion-faoutput.mp4
Speedup and memory saving aren't that impressive yet, maybe it can be improved further?