Wan2.1 i2v running issue

[luren55]

Describe the bug

When I was trying to run the script which is python cache-dit/examples/pipelinerun_wan_2.2.py by the weight of Wan2.1-I2V-14B-720P-Diffusers, an error displayed.

Traceback (most recent call last):
File "/home/mzb/cache-dit/examples/pipeline/run_wan.py", line 60, in
video = pipe(
^^^^^
File "/root/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.11/site-packages/diffusers/pipelines/wan/pipeline_wan.py", line 613, in call
latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.11/site-packages/diffusers/schedulers/scheduling_unipc_multistep.py", line 992, in step
model_output_convert = self.convert_model_output(model_output, sample=sample)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/root/miniconda3/lib/python3.11/site-packages/diffusers/schedulers/scheduling_unipc_multistep.py", line 628, in convert_model_output
x0_pred = sample - sigma_t * model_output
~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (36) must match the size of tensor b (16) at non-singleton dimension 1

Reproduction

See description

Logs

System Info

• 🤗 Diffusers version: 0.35.2
• Platform: Linux-5.15.0-25-generic-x86_64-with-glibc2.35
• Running on Google Colab?: No
• Python version: 3.11.11
• PyTorch version (GPU?): 2.9.0+cu128 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.36.0
• Transformers version: 4.57.1
• Accelerate version: 1.11.0
• PEFT version: not installed
• Bitsandbytes version: not installed
• Safetensors version: 0.6.2
• xFormers version: not installed
• Accelerator: NVIDIA L20, 46068 MiB

Who can help?

No response

[asomoza]

Hi, can you test it by installing diffusers from main?

pip uninstall diffusers
pip install git+https://github.com/huggingface/diffusers.git

[zhangtao0408]

I tested in diffusers main branch, still not working, got same error.

Code

import os
import sys

sys.path.append("..")

import time
import torch
import diffusers
from diffusers import AutoModel, WanPipeline, AutoencoderKLWan, ContextParallelConfig
from diffusers.utils import export_to_video
from diffusers.schedulers.scheduling_unipc_multistep import (
    UniPCMultistepScheduler,
)

height, width = 1024, 1024
model_path = os.environ.get(
                            "WAN_DIR",
                            "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
                        )

# Initialization 
torch.distributed.init_process_group("nccl")
rank = torch.distributed.get_rank()
device = torch.device("cuda", rank % torch.cuda.device_count())
torch.cuda.set_device(device)

pipe = WanPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16, device_map="cuda")
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=4))

# flow shift should be 3.0 for 480p images, 5.0 for 720p images
if hasattr(pipe, "scheduler") and pipe.scheduler is not None:
    # Use the UniPCMultistepScheduler with the specified flow shift
    flow_shift = 3.0 if height == 480 else 5.0
    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config,
        flow_shift=flow_shift,
    )

start = time.time()
video = pipe(
    prompt=(
        "An astronaut dancing vigorously on the moon with earth "
        "flying past in the background, hyperrealistic"
    ),
    negative_prompt="",
    height=height,
    width=width,
    num_frames=49,
    num_inference_steps=35,
    generator=torch.Generator("cpu").manual_seed(0),
).frames[0]
end = time.time()

time_cost = end - start
save_path = f"wan.mp4"
print(f"Time cost: {time_cost:.2f}s")
print(f"Saving video to {save_path}")
export_to_video(video, save_path, fps=16)

ERROR

  0%|                                                                                                                                                                                        | 0/35 [00:14<?, ?it/s]
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/vipshop/wan_test.py", line 67, in <module>
[rank1]:     video = pipe(
[rank1]:             ^^^^^
[rank1]:   File "/root/miniconda3/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/vipshop/diffusers/src/diffusers/pipelines/wan/pipeline_wan.py", line 613, in __call__
[rank1]:     latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/vipshop/diffusers/src/diffusers/schedulers/scheduling_unipc_multistep.py", line 992, in step
[rank1]:     model_output_convert = self.convert_model_output(model_output, sample=sample)
[rank1]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/vipshop/diffusers/src/diffusers/schedulers/scheduling_unipc_multistep.py", line 628, in convert_model_output
[rank1]:     x0_pred = sample - sigma_t * model_output
[rank1]:               ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
[rank1]: RuntimeError: The size of tensor a (36) must match the size of tensor b (16) at non-singleton dimension 1

[asomoza]

@zhangtao0408 the issue doesn't have any code so it's hard to guess what's the issue, so I was just making sure that this issue is not something that is already fixed in the main branch (there was also another issue from cache-dit)

Your problem seems similar to the one in this PR.

[asomoza]

@luren55 can you please update the issue with a minimal reproducible code with just diffusers that can reproduce your error please?