Missing ai-toolkit/wan2.1-vae when using local WAN 2.2 (14B)

[JoseCovaOS]

This is for bugs only

Did you already ask in the discord?

No

You verified that this is a bug and not a feature request or question by asking in the discord?

Yes, it is a bug

Describe the bug

When trying to execute a job selecting WAN 2.2 (14B) and pointing it to a directory containing the required files (transformers, text_encoder, tokenizer and even the vae from your hugging face page) an error occurs that says the following:

OSErrorOSError: : ai-toolkit/wan2.1-vae does not appear to have a file named config.json.

After having a look at the code, it seems there's a bug in toolkit/models/wan21/wan21.py in the block from lines 436-442. In this if-else block, self._wan_vae_path is always not None therefore, entering in the if block and using the wrong path for the Autoencoder when for a local install, it should enter in the else block and try to use the vae from the local path.

self._wan_vae_path comes from this line in extensions_built_in/wan22/wan22_14b_model.py where path is hardcoded and the reason why code always enters in the if block.

[VextraTiger]

To Fix Modify Line 167 of wan22_14b_model.py

_wan_vae_path = "ai-toolkit/wan2.1-vae" change to _wan_vae_path = None

This will ensure the code enters the else block in toolkit/models/wan21/wan21.py (lines 436-442), allowing it to load the VAE from your local directory instead of the invalid remote path.

[2O48]

Hello, where can I find a place to load my vae models in ai-toolkit training job page?

[JoseCovaOS]

Hi @2O48

In the Name or Path field. This path should point to a directory that must contain the following: diffusers, UMT5 encoder and vae. I downloaded everything from here.

Structure should be like:

point_to_this_directory/
├─ transformer/
├─ transformer_2/
├─ text_encoder/
├─ tokenizer/
├─ vae/
│  ├─ (here just only the files from wan2.1-vae)

Because of the bug I describe in this issue, you need to modify the line 167 of wan22_14b_model.py like @VextraTiger suggest in order for this to work.

By doing this I was able to run a WAN2.2 job.

[2O48]

thanks