[Bugfix][Wan2.2] Apply UniPC flow shift once (#7576)

Signed-off-by: HAAZZZEEEE <wang-taicheng@qq.com>

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -59,9 +59,11 @@ FASTWAN_DMD_SCHEDULER_SHIFT = 8.0
 
 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any:
     if sample_solver == "unipc":
+        # Keep native Wan's unshifted training endpoints; set_timesteps applies
+        # the requested shift to the interpolated inference sigmas.
         return FlowUniPCMultistepScheduler(
             num_train_timesteps=1000,
-            shift=flow_shift,
+            shift=1.0,
             prediction_type="flow_prediction",
         )
     if sample_solver == "euler":
@@ -774,7 +776,10 @@ class Wan22Pipeline(
                 self._sample_solver = sample_solver
                 self._flow_shift = flow_shift
 
-            self.scheduler.set_timesteps(num_steps, device=device)
+            if sample_solver == "unipc":
+                self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift)
+            else:
+                self.scheduler.set_timesteps(num_steps, device=device)
             timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
         boundary_timestep = None

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -629,7 +629,10 @@ class Wan22I2VPipeline(
             self._flow_shift = flow_shift
 
         # Timesteps
-        self.scheduler.set_timesteps(num_steps, device=device)
+        if sample_solver == "unipc":
+            self.scheduler.set_timesteps(num_steps, device=device, shift=flow_shift)
+        else:
+            self.scheduler.set_timesteps(num_steps, device=device)
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -545,10 +545,10 @@ class Wan22S2VPipeline(
             self._init_original_format(model_path, dtype)
 
         # -- Scheduler --
-        flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0  # S2V default
+        self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 3.0  # S2V default
         self.scheduler = FlowUniPCMultistepScheduler(
             num_train_timesteps=1000,
-            shift=flow_shift,
+            shift=1.0,
             prediction_type="flow_prediction",
         )
 
@@ -1342,7 +1342,7 @@ class Wan22S2VPipeline(
                 )
 
             # -- Scheduler --
-            self.scheduler.set_timesteps(num_steps, device=device)
+            self.scheduler.set_timesteps(num_steps, device=device, shift=self._flow_shift)
             timesteps = self.scheduler.timesteps
             self._num_timesteps = len(timesteps)
 

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py

@@ -690,7 +690,10 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput):
         )
 
         # Set up scheduler
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        if self._sample_solver == "unipc":
+            self.scheduler.set_timesteps(num_inference_steps, device=device, shift=self._flow_shift)
+        else:
+            self.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = self.scheduler.timesteps
         self._num_timesteps = len(timesteps)