Fix Wan I2V Quality

[chengzeyi]

According to: https://github.com/Wan-Video/Wan2.1/blob/b58b7c573776b76b6fe8d36086590e033173f9b1/wan/modules/vae.py#L537

[yiyixuxu]

I think this is actually should be same as sample_mode == "argmax" , no?

diffusers/src/diffusers/models/autoencoders/vae.py

Line 743 in 1001425

def mode(self) -> torch.Tensor:

[chengzeyi]

Yeah, I think it is just without any extra computation cost.

[yiyixuxu]

got it, let's :

1. still use sample_mode="argmax" so that this is argument is consistent acorss all pipelines
2. remove the code that's not needed for wan (I think all of the other mode lol

[yiyixuxu]

got it, let's :

1. still use sample_mode="argmax" so that this is argument is consistent acorss all pipelines
2. remove the code that's not needed for wan (I think all of the other mode lol

[yiyixuxu]

cc @hlky here, I think we can refactor &simplify this code a bit here just for wan, this is not a #Copied from method anymore, it does not have to have same pattern as the original one

[hlky]

I think now retrieve_latents can be Copied from.

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
):
    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
        return encoder_output.latent_dist.sample(generator)
    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
        return encoder_output.latent_dist.mode()
    elif hasattr(encoder_output, "latents"):
        return encoder_output.latents
    else:
        raise AttributeError("Could not access latents of provided encoder_output")

We pass sample_mode="argmax" then do the scaling

        if isinstance(generator, list):
            latent_condition = [
                retrieve_latents(self.vae.encode(video_condition), g, sample_mode="argmax") for g in generator
            ]
            latent_condition = torch.cat(latent_condition)
        else:
            latent_condition = retrieve_latents(self.vae.encode(video_condition), generator, sample_mode="argmax")
            latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)

        latent_condition = (latent_condition - latents_mean) * latents_std

[hlky]

Output is good with #11087 (review)

output.23.mp4

[hlky]

@chengzeyi Can we apply these changes #11087 (review) instead

[chengzeyi]

@chengzeyi Can we apply these changes #11087 (review) instead

From my part I think these changes are OK.👌
And obviously no need to pass the generator in such cases.

[hlky]

The changes are ok but the other changes allow us to keep copied from for retrieve_latents and only add 1 line latent_condition = (latent_condition - latents_mean) * latents_std to prepare_latents and modify 2 others to add sample_mode="argmax", generator can also be removed from the function call as you mentioned.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

thanks a ton @chengzeyi !!!!