update default flow shift

Signed-off-by: Olivia Stoner <245287810+o-stoner@users.noreply.github.com>

Author: o-stoner

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py

@@ -499,12 +499,17 @@ def forward(
         resolved_flow_shift = (
             flow_shift if flow_shift is not None else self._default_flow_shift(height, width)
         )
-        if self.scheduler.config.shift != resolved_flow_shift:
+
+        sched_cfg = self.scheduler.config
+        shift_key = (
+            "shift" if "shift" in sched_cfg else "flow_shift" if "flow_shift" in sched_cfg else None
+        )
+        if shift_key is not None and sched_cfg[shift_key] != resolved_flow_shift:
             logger.info(
-                f"flow_shift: {self.scheduler.config.shift} -> {resolved_flow_shift} "
+                f"flow_shift: {sched_cfg[shift_key]} -> {resolved_flow_shift} "
                 f"({'user' if flow_shift is not None else 'variant default'})"
             )
-            self.scheduler.config.shift = resolved_flow_shift
+            self.scheduler.register_to_config(**{shift_key: resolved_flow_shift})
 
         self.scheduler.set_timesteps(num_inference_steps, device=self.device)
 

tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py

@@ -550,12 +550,17 @@ def forward(
         resolved_flow_shift = (
             flow_shift if flow_shift is not None else self._default_flow_shift(height, width)
         )
-        if self.scheduler.config.shift != resolved_flow_shift:
+
+        sched_cfg = self.scheduler.config
+        shift_key = (
+            "shift" if "shift" in sched_cfg else "flow_shift" if "flow_shift" in sched_cfg else None
+        )
+        if shift_key is not None and sched_cfg[shift_key] != resolved_flow_shift:
             logger.info(
-                f"flow_shift: {self.scheduler.config.shift} -> {resolved_flow_shift} "
+                f"flow_shift: {sched_cfg[shift_key]} -> {resolved_flow_shift} "
                 f"({'user' if flow_shift is not None else 'variant default'})"
             )
-            self.scheduler.config.shift = resolved_flow_shift
+            self.scheduler.register_to_config(**{shift_key: resolved_flow_shift})
 
         self.scheduler.set_timesteps(num_inference_steps, device=self.device)