Wan 2.2 VAE forward fails

[a-r-r-o-w]

Hey,

With Wan2.1 we was able to pass just a RGB PIl image. With 2.2 i get

def _conv_forward(self, input: Tensor, weight: Tensor, bias: Optional[Tensor]):
        if self.padding_mode != "zeros":
            return F.conv3d(
                F.pad(
                    input, self._reversed_padding_repeated_twice, mode=self.padding_mode
                ),
                weight,
                bias,
                self.stride,
                _triple(0),
                self.dilation,
                self.groups,
            )
>       return F.conv3d(
            input, weight, bias, self.stride, self.padding, self.dilation, self.groups
        )
E       RuntimeError: Given groups=1, weight of size [160, 12, 3, 3, 3], expected input[1, 3, 3, 258, 258] to have 12 channels, but got 3 channels instead

/opt/conda/lib/python3.11/site-packages/torch/nn/modules/conv.py:720: RuntimeError

Cheers,
Joe

Originally posted by @JoeGaffney in #12006 (comment)

[a-r-r-o-w]

@JoeGaffney, could you provide us with a minimal reproducible snippet that triggers this error?

[JoeGaffney]

Hey @a-r-r-o-w

I will aim to get one to you and double check my logic.

Cheers,
Joe

[JoeGaffney]

Hey,

Narrowed it down to vae.enable_tiling()

import torch
from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
from diffusers.utils import export_to_video, load_image
from transformers import UMT5EncoderModel

from common.logger import logger
from common.pipeline_helpers import get_quantized_model
from utils.utils import get_16_9_resolution, resize_image

UMT_T5_MODEL_PATH = "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers"


def get_pipeline(model_id, torch_dtype=torch.bfloat16) -> WanImageToVideoPipeline:
    text_encoder = get_quantized_model(
        model_id=UMT_T5_MODEL_PATH,
        subfolder="text_encoder",
        model_class=UMT5EncoderModel,
        target_precision=8,
        torch_dtype=torch_dtype,
    )

    vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)

    pipe = WanImageToVideoPipeline.from_pretrained(
        model_id,
        text_encoder=text_encoder,
        torch_dtype=torch_dtype,
        vae=vae,
    )

    # NOTE works with wan2.1 but not with wan2.2
    # try:
    #     pipe.vae.enable_tiling()  # Enable VAE tiling to improve memory efficiency
    # except:
    #     pass

    pipe.enable_model_cpu_offload()
    return pipe


def main():
    pipe = get_pipeline(model_id="Wan-AI/Wan2.2-TI2V-5B-Diffusers", torch_dtype=torch.bfloat16)
    image = load_image(
        "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/guitar-man.png"
    )
    logger.info(f"Image size: {image.size}")

    width, height = get_16_9_resolution("720p")
    image = resize_image(image, 16, 1.0, width, height)

    prompt = "A man with short gray hair plays a red electric guitar."
    # Wan gives better results with a default negative prompt
    negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"

    num_frames = 48
    num_inference_steps = 5
    guidance_scale = 5.0

    output = pipe(
        width=image.size[0],
        height=image.size[1],
        image=image,
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        num_frames=num_frames,
        guidance_scale=guidance_scale,
    ).frames[0]

    return export_to_video(output, "output.mp4", fps=24)

My tests now run with this off for both 2.1 and 2.2.

Maybe I shouldn't be setting this now?

Cheers,
Joe

[JoeGaffney]

I do get a very canvas looking like result even when raising steps to 20 or 50

This wasn't really here with WAN 2.1

wan-2-2-5b_image_to_video.mp4

[OLDDDDDDD]

I do get a very canvas looking like result even when raising steps to 20 or 50

This wasn't really here with WAN 2.1

wan-2-2-5b_image_to_video.mp4

I got like this too when step is 50, I will check this in 20 steps

[a-r-r-o-w]

@JoeGaffney Could you try to pull the latest main branch commit and see if that fixes it for you? I believe this PR should have fixed the problem: #12041

[wenqsun]

@a-r-r-o-w I still have similar problem when I use the latest vae code on the main brach

[a-r-r-o-w]

cc @yiyixuxu Is this known issue?

[JoeGaffney]

@OLDDDDDDD

Try disabling pipe.vae.enable_tiling()

If you had it on. That fixed it for me.

@a-r-r-o-w
Thanks I tried the latest from main and that fixed the canvas issue.

Although I'm not that impressed with the 5b model. It seems its not that great for people at least so will probably abandon more testing for now.

The 5b is quite fast though I will give it that, but its sort of comparable with LTX looks wise. Which is even faster.

Cheers,
Joe

[a-r-r-o-w]

Thanks for testing @JoeGaffney! So, if I understand correctly, you get poor outputs when you enable vae tiling, but the outputs look good visibly without it?

Many latest video models require prompting in special ways (sometimes LLM augmented) to be able to generate good quality videos. This is something part of their training procedure, and so probably nothing can be done on our end regarding that. I would recommend using a small prompt enhancer model for your needs, because IME it improves the quality quite a lot when I'm lazy to type longer and more descriptive prompts

[JoeGaffney]

hey @a-r-r-o-w

Sorry I wasn't so clear the vae tiling was causing the channel issue. The canvas look was fixed with using the latest from main.

Yes of course there is nothing you can do about the actually model. I've switch to using the 14b one as there is generally less than positive feedback across the web about the 5b model. Just wanted to give it ago myself.

You can close the ticket if you want i'm good from my end.

Cheers,
Joe

[a-r-r-o-w]

Awesome, thanks for clarifying! I believe this issue can be closed then