[Wan2.2] pipeline_wan miss the 'shift' parameter which used by Wan2.2-A14B-diffusers.

[yvmilir]

Firstly, I found that the quality of output using diffusers is poor
Later, I found that the pipeline_wan in diffusers[0.34.0] did not support two-stage processing. I noticed that the community had already updated it, so I installed diffusers[0.35.0-dev] by source code and it worked.

Then I found that the scheduler in diffusers does not support the parameter "shift", but "sample_shift" is an important parameter generated by Wan2.2, which may also lead to differences from the official inference code of Wan2.2. Therefore, the video effect may still be inferior to the original inference code.

Wan-Video/Wan2.2#69

What I need
Can the community provide the UniPCMultistepScheduler and DPMSolverMultistepScheduler that support the 'shift' parameter? Or can it be adapted in pipeline_wan so that the shift parameter can be used.

Or is there something wrong with my understanding? How can I correctly use the shift parameter when using diffusers?

Thanks!!
cc @yiyixuxu @a-r-r-o-w

[lopho]

# pipe.scheduler is UniPCMultistepScheduler
pipe.scheduler.config.flow_shift = 6.9 # nice

if I'm mistaken, and you mean another shift parameter, could you please point to the exact location of this parameter in the original WAN2.2 source?

[yiyixuxu]

@yvmilir thanks for the issue
I agree with @lopho, I think flow_shift might be what you look for here

so, to change a scheduler along with shift parameter

pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
flow_shift = 6.9
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
pipe.to("cuda")

[yvmilir]

pipe.scheduler is UniPCMultistepScheduler

pipe.scheduler.config.flow_shift = 6.9 # nice
if I'm mistaken, and you mean another shift parameter, could you please point to the exact location of this parameter in the original WAN2.2 source?

Thanks for your reply!

The parameter 'shift' I am referring to is the one used in sample_scheduler.set_timesteps. https://github.com/Wan-Video/Wan2.2/blob/main/wan/text2video.py#L309
This parameter is declared during generation and has a default configuration of 'sample_sthift', which affects the generated timesteps

I am a beginner, if my understanding is incorrect, please point it out.

Thanks!
cc @yiyixuxu

[lopho]

Yes, that is the same parameter. flow_shift in diffusers is shift in Wan2.2.

Original Wan2.2 source:
https://github.com/Wan-Video/Wan2.2/blob/031a9be56cec91e86d140d3d3a74280fb05a9b1c/wan/utils/fm_solvers_unipc.py#L116

sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)

diffusers:

diffusers/src/diffusers/schedulers/scheduling_unipc_multistep.py

Line 391 in 03c3f69

sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()

sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()

[yvmilir]

Yes, that is the same parameter. flow_shift in diffusers is shift in Wan2.2.

Original Wan2.2 source: https://github.com/Wan-Video/Wan2.2/blob/031a9be56cec91e86d140d3d3a74280fb05a9b1c/wan/utils/fm_solvers_unipc.py#L116

sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
diffusers:

diffusers/src/diffusers/schedulers/scheduling_unipc_multistep.py

Line 391 in 03c3f69

sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()
sigmas = np.flip(self.config.flow_shift * sigmas / (1 + (self.config.flow_shift - 1) * sigmas))[:-1].copy()

So in the implementation of diffusers, flow_stift is determined when declaring the pipe.

I saw in Wan2.2 that shift affects the timesteps, so I thought the timesteps in the__call__method of the diffusers' pipeline can also be changed based on the shift value.

https://github.com/Wan-Video/Wan2.2/blob/031a9be56cec91e86d140d3d3a74280fb05a9b1c/wan/text2video.py#L313

              sample_scheduler.set_timesteps(sampling_steps, device=self.device, shift=shift)
              timesteps = sample_scheduler

I am trying to use a diffuser to provide model generation services and want to load model weights at once, so if I want to use different offset values, I can only declare different pipelines, right?

Thank you for your reply!

[lopho]

You can change the shift parameter even after initializing the pipe. You do not have to recreate the pipe to change the scheduler or the shift parameter.

# load the pipeline
pipe = WanPipeline.from_pretrained(model_id)
# change shift of the UniPC scheduler
pipe.scheduler.config.flow_shift = 6.9
# run the pipeline
frames = pipe(prompt="funny banana doing a backflip").frames[0]
# change the shift again
pipe.scheduler.config.flow_shift = 5.02
# run the pipeline with the new shift applied
frames2 = pipe(prompt="funny banana doing a backflip, insano mode").frames[0]

maybe wrap it in a function so it behaves more closely to what you expected:

def gen(prompt, neg_prompt, steps, ..., shift):
    pipe.scheduler.config.flow_shift = shift
    return pipe(prompt, ...)

[yvmilir]

You can change the shift parameter even after initializing the pipe. You do not have to recreate the pipe to change the scheduler or the shift parameter.

load the pipeline

pipe = WanPipeline.from_pretrained(model_id)

change shift of the UniPC scheduler

pipe.scheduler.config.flow_shift = 6.9

run the pipeline

frames = pipe(prompt="funny banana doing a backflip").frames[0]

change the shift again

pipe.scheduler.config.flow_shift = 5.02

run the pipeline with the new shift applied

frames2 = pipe(prompt="funny banana doing a backflip, insano mode").frames[0]
maybe wrap it in a function so it behaves more closely to what you expected:

def gen(prompt, neg_prompt, steps, ..., shift):
pipe.scheduler.config.flow_shift = shift
return pipe(prompt, ...)

Thanks, this is very helpful for me