[bugfix]: PAN2 follows the reference sampler: checkpoint flow shift, first-latent CFG skip for image-to-video (#1961)

docs/inference/pan2.md

@@ -18,8 +18,9 @@ from the `pan2_t2v` and `pan2_i2v` presets:
 |---|---|
 | Resolution | 704x1248 (height and width must be multiples of 32) |
 | Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) |
-| Steps | 50, uniform sigmas from 1 to 0 with flow shift 7 |
-| Guidance | 3.0, classifier-free guidance with an empty negative prompt |
+| Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) |
+| Guidance | 3.0, classifier-free guidance. Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) |
+| Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps |
 
 PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw
 prompts work, with weaker results.

fastvideo/configs/pipelines/pan2.py

@@ -24,7 +24,12 @@ class PAN2PipelineConfig(PipelineConfig):
 
     dit_precision: str = "bf16"
     vae_precision: str = "fp32"
-    flow_shift: float | None = 7.0
+    # Unset: the checkpoint's scheduler config decides, which is shift 9 for IFM/PAN2-TEST, the reference sampler's
+    # fixed value. A fixed default here would override the released shift with whatever a test checkpoint declares.
+    flow_shift: float | None = None
+    # Image-to-video: the first latent frame (the conditioning image) takes the positive prediction instead of the
+    # guided one, as the reference sampler does by default (``--i2v_skip_first_latent_cfg``).
+    i2v_skip_first_latent_cfg: bool = True
     # PAN2 samples on fp32 latents; keep the Euler update out of autocast.
     scheduler_step_in_fp32: bool = True
 

fastvideo/pipelines/basic/pan2/presets.py

@@ -10,8 +10,9 @@ _DENOISE_STAGE = PresetStageSpec(
     allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}),
 )
 
-# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. PAN2 has no default negative prompt; an empty one is
-# encoded as a single space.
+# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. The preset sends no negative prompt (an empty one is encoded
+# as a single space); the reference sampler sends the release's structured negative caption, so pass it as the
+# request's ``negative_prompt`` to follow that recipe.
 _SHARED_DEFAULTS = {
     "height": 704,
     "width": 1248,

fastvideo/pipelines/basic/pan2/stages.py

@@ -196,7 +196,16 @@ class PAN2ConditioningStage(PipelineStage):
 
 
 class PAN2DenoisingStage(DenoisingStage):
-    """Keeps the latents in fp32 across the loop; they are cast to the DiT dtype for each forward."""
+    """Keeps the latents in fp32 across the loop; they are cast to the DiT dtype for each forward.
+
+    Image-to-video leaves the conditioning image's latent frame out of classifier-free guidance: that frame takes the
+    positive prediction, as in the reference sampler.
+    """
+
+    def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor:
+        if batch.pil_image is not None and fastvideo_args.pipeline_config.i2v_skip_first_latent_cfg:
+            noise_pred[:, :, :1] = noise_pred_text[:, :, :1]
+        return noise_pred
 
     def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
         return noise_pred.float()

fastvideo/pipelines/stages/denoising.py

@@ -464,6 +464,8 @@ class DenoisingStage(PipelineStage):
                                 noise_pred = noise_pred_uncond + current_guidance_scale * (noise_pred_text -
                                                                                            noise_pred_uncond)
 
+                        noise_pred = self.finish_guidance(noise_pred, noise_pred_text, batch, fastvideo_args)
+
                         # Apply guidance rescale if needed
                         if batch.guidance_rescale > 0.0:
                             # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
@@ -606,6 +608,10 @@ class DenoisingStage(PipelineStage):
     def prepare_timestep(self, timestep, latent_model_input, batch, state) -> torch.Tensor:
         return timestep.repeat(latent_model_input.shape[0])
 
+    def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor:
+        """The guided prediction handed on to the scheduler; a model overrides it to keep part of the output out of CFG."""
+        return noise_pred
+
     def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor:
         """The model output handed to the scheduler step, which returns the next latents in its dtype."""
         return noise_pred

fastvideo/tests/stages/test_pan2_first_latent_cfg.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""PAN2 image-to-video keeps the first latent frame out of classifier-free guidance; no weights or CUDA needed."""
+
+import pytest
+import torch
+
+from fastvideo.tests.stages._denoising_fixtures import (
+    NullProgressBar, TinyDenoiser, TinyScheduler, _patch_denoising_module, _tiny_args, _tiny_batch,
+)
+
+
+def _one_step(monkeypatch, env_overrides, *, image_to_video: bool, skip_first_latent: bool):
+    _patch_denoising_module(monkeypatch, env_overrides, None)
+    from fastvideo.pipelines.basic.pan2.stages import PAN2DenoisingStage
+
+    stage = PAN2DenoisingStage(TinyDenoiser(), TinyScheduler())
+    stage.progress_bar = lambda iterable=None, total=None: NullProgressBar()
+    batch = _tiny_batch()
+    batch.latents = torch.zeros(1, 1, 3, 1, 1)
+    batch.raw_latent_shape = (1, 1, 3, 1, 1)
+    batch.timesteps = torch.tensor([0.0])
+    batch.num_inference_steps = 1
+    batch.pil_image = object() if image_to_video else None
+    args = _tiny_args()
+    args.pipeline_config.i2v_skip_first_latent_cfg = skip_first_latent
+    args.pipeline_config.scheduler_step_in_fp32 = True
+    return stage.forward(batch, args).latents.flatten().tolist()
+
+
+def test_image_to_video_takes_the_positive_prediction_for_the_first_latent_frame(monkeypatch, env_overrides):
+    # With a zero latent and timestep the cond prediction is 0.5 and the uncond -0.5, so CFG 2 guides to 1.5.
+    latents = _one_step(monkeypatch, env_overrides, image_to_video=True, skip_first_latent=True)
+
+    assert latents == pytest.approx([-0.005, -0.015, -0.015])
+
+
+@pytest.mark.parametrize("image_to_video,skip_first_latent", [(True, False), (False, True)])
+def test_every_frame_is_guided_without_a_conditioning_image_or_with_the_skip_off(monkeypatch, env_overrides,
+                                                                                image_to_video, skip_first_latent):
+    latents = _one_step(monkeypatch, env_overrides, image_to_video=image_to_video, skip_first_latent=skip_first_latent)
+
+    assert latents == pytest.approx([-0.015, -0.015, -0.015])