Add Flow-CPS and matched sampler pipeline

Add a Diffusers custom image-to-video pipeline for CPS 0.1/0.3/0.7/0.9, FlowMatch Euler, and UniPC. Update the model card with runnable loading examples, paper-matched settings, and sampler results.

README.md

@@ -22,7 +22,7 @@ tags:
 
 # VBVR-Pro Wan2.2 TI2V-5B — Rule-RL
 
-This repository contains a complete Diffusers pipeline for the VBVR-Pro
+This repository contains a complete Diffusers checkpoint for the VBVR-Pro
 Wan2.2 TI2V-5B model optimized with task-specific rule-based reinforcement
 learning. It is derived from
 [`Wan-AI/Wan2.2-TI2V-5B-Diffusers`](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers)
@@ -30,7 +30,27 @@ and is intended for research on image-conditioned video generation and visual
 reasoning.
 
 The repository includes the transformer, text encoder, tokenizer, VAE, and
-scheduler required by `WanPipeline.from_pretrained`.
+scheduler. It now also includes [`pipeline.py`](./pipeline.py), a custom
+image-to-video pipeline exposing all six inference configurations evaluated in
+the VBVR-Pro paper.
+
+## Inference samplers
+
+The sampler is selected per call; the model weights do not change.
+
+| `sampler` | Inference method | CPS coefficient | Paper overall score |
+|---|---|---:|---:|
+| `cps-0.1` | Flow-CPS | 0.1 | 0.509 |
+| `cps-0.3` | Flow-CPS | 0.3 | 0.526 |
+| `cps-0.7` | Flow-CPS | 0.7 | **0.548** |
+| `cps-0.9` | Flow-CPS | 0.9 | 0.539 |
+| `euler` | FlowMatch Euler ODE | — | 0.522 |
+| `unipc` | UniPC ODE | — | 0.522 |
+
+These are the aggregate VBVR-Pro-Bench results reported in Table 8 under the
+matched settings below. The model was trained with Flow-CPS coefficient 0.7.
+Reported scores are evaluation results, not guarantees for other prompts or
+runtime configurations.
 
 ## Recommended evaluation settings
 
@@ -40,32 +60,35 @@ scheduler required by `WanPipeline.from_pretrained`.
 - Inference steps: 30
 - Guidance scale: 1.0
 
-The bundled scheduler is UniPC for direct Diffusers compatibility. The
-Flow-CPS samplers used in the VBVR-Pro experiments are available in
-[`pufanyi/vbvr-rl`](https://github.com/pufanyi/vbvr-rl).
+## Usage with all six samplers
 
-## Usage
+Use Diffusers 0.37.1 or newer. Because this loads Python code from the model
+repository, review `pipeline.py`, pass `trust_remote_code=True`, and pin a
+reviewed `revision` in production.
 
 ```python
 import torch
-from diffusers import AutoencoderKLWan, WanPipeline
+from diffusers import AutoencoderKLWan, DiffusionPipeline
 from diffusers.utils import export_to_video, load_image
 
 model_id = "pufanyi/VBVR-Pro-Wan2.2-TI2V-5B-Rule-RL"
 
+# Wan's VAE is kept in float32 for stable decoding.
 vae = AutoencoderKLWan.from_pretrained(
     model_id,
     subfolder="vae",
     torch_dtype=torch.float32,
 )
-pipe = WanPipeline.from_pretrained(
+pipe = DiffusionPipeline.from_pretrained(
     model_id,
+    custom_pipeline="pipeline",
+    trust_remote_code=True,
     vae=vae,
     torch_dtype=torch.bfloat16,
 )
 pipe.enable_model_cpu_offload()
 
-image = load_image("input.png")
+image = load_image("input.png").convert("RGB")
 frames = pipe(
     image=image,
     prompt="Move the marked object to the matching target.",
@@ -74,15 +97,60 @@ frames = pipe(
     num_frames=81,
     num_inference_steps=30,
     guidance_scale=1.0,
-    generator=torch.Generator(device="cpu").manual_seed(0),
+    sampler="cps-0.7",  # cps-0.1, cps-0.3, cps-0.7, cps-0.9, euler, or unipc
+    generator=torch.Generator(device="cuda").manual_seed(0),
 ).frames[0]
 
 export_to_video(frames, "output.mp4", fps=16)
 ```
 
-Use Diffusers 0.37.1 or newer. Loading the complete pipeline requires
-substantial CPU and accelerator memory; CPU offloading is recommended on
-smaller GPUs.
+The generic form `sampler="cps", cps_eta=<value>` accepts any finite
+coefficient from 0 to 1. `generator` controls the initial latent and, by
+default, the fresh Flow-CPS transition noise. Pass a separate
+`cps_generator` when the two random streams must be controlled independently.
+
+Loading the complete pipeline requires substantial CPU and accelerator
+memory. CPU offloading is recommended on smaller GPUs.
+
+## Standard Diffusers compatibility
+
+The bundled scheduler remains UniPC and `model_index.json` is unchanged. Users
+who only need the standard deterministic path can load the checkpoint without
+remote custom code:
+
+```python
+import torch
+from diffusers import AutoencoderKLWan, WanImageToVideoPipeline
+
+model_id = "pufanyi/VBVR-Pro-Wan2.2-TI2V-5B-Rule-RL"
+
+vae = AutoencoderKLWan.from_pretrained(
+    model_id,
+    subfolder="vae",
+    torch_dtype=torch.float32,
+)
+pipe = WanImageToVideoPipeline.from_pretrained(
+    model_id,
+    vae=vae,
+    torch_dtype=torch.bfloat16,
+)
+```
+
+Use `WanImageToVideoPipeline`, not the text-to-video `WanPipeline`: the latter
+does not accept the first-frame `image` argument in Diffusers 0.37.1.
+
+## Implementation and reproducibility notes
+
+- Flow-CPS uses the training-time shifted `linspace(1, 0, T + 1)` sigma grid
+  and preserves the released scheduler's `flow_shift: 5.0`.
+- CPS updates are evaluated in float32 and cast back to the transformer latent
+  dtype. Euler and UniPC retain their native Diffusers latent precision and
+  solver grids.
+- The custom class subclasses `WanImageToVideoPipeline`, preserving the
+  official first-frame VAE conditioning and TI2V-5B expanded-timestep mask.
+- The release training/evaluation repository remains the source of truth for
+  formal score provenance. Exact output bytes can vary with PyTorch,
+  Diffusers, attention backend, dtype, and device.
 
 ## Training summary
 

pipeline.py

@@ -0,0 +1,384 @@
+"""Diffusers pipeline for the six VBVR-Pro RL inference configurations.
+
+The released TI2V-5B checkpoints contain ordinary flow-prediction weights.
+They are not tied to the bundled UniPC scheduler.  This module keeps the
+official Diffusers image-conditioning path and adds the Flow-CPS transition
+used for RL rollouts, alongside matched FlowMatch Euler and UniPC ODE modes.
+
+The file is intentionally self-contained so it can also be copied to a model
+repository as a Diffusers custom pipeline.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler, WanImageToVideoPipeline
+from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteSchedulerOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+PAPER_SAMPLER_PRESETS = (
+    "cps-0.1",
+    "cps-0.3",
+    "cps-0.7",
+    "cps-0.9",
+    "euler",
+    "unipc",
+)
+
+_CPS_PRESETS = {
+    "cps-0.1": 0.1,
+    "cps-0.3": 0.3,
+    "cps-0.7": 0.7,
+    "cps-0.9": 0.9,
+    "cps0p1": 0.1,
+    "cps0p3": 0.3,
+    "cps0p7": 0.7,
+    "cps0p9": 0.9,
+}
+
+_SAMPLER_ALIASES = {
+    "cps": "cps",
+    "flow-cps": "cps",
+    "flowcps": "cps",
+    "euler": "euler",
+    "euler-ode": "euler",
+    "flow-match-euler": "euler",
+    "flowmatch-euler": "euler",
+    "unipc": "unipc",
+    "unipc-ode": "unipc",
+}
+
+
+def _config_value(config: Any, name: str, default: Any = None) -> Any:
+    if hasattr(config, "get"):
+        return config.get(name, default)
+    return getattr(config, name, default)
+
+
+def _normalize_sampler_name(sampler: str) -> str:
+    if not isinstance(sampler, str) or not sampler.strip():
+        raise ValueError(f"sampler must be a non-empty string, got {sampler!r}")
+    return sampler.strip().lower().replace("_", "-").replace(" ", "-")
+
+
+def resolve_vbvr_sampler(sampler: str, cps_eta: float | None = None) -> tuple[str, float | None]:
+    """Resolve a public sampler name to ``(kind, eta)``.
+
+    ``kind`` is one of ``cps``, ``euler``, and ``unipc``.  The four CPS names
+    in :data:`PAPER_SAMPLER_PRESETS` carry their paper coefficient, while the
+    generic ``cps`` name defaults to the training coefficient 0.7.
+    """
+
+    normalized = _normalize_sampler_name(sampler)
+    preset_key = normalized.replace("cps-0p", "cps0p")
+    if normalized in _CPS_PRESETS or preset_key in _CPS_PRESETS:
+        preset_eta = _CPS_PRESETS.get(normalized, _CPS_PRESETS[preset_key])
+        if cps_eta is not None and float(cps_eta) != preset_eta:
+            raise ValueError(f"sampler={sampler!r} fixes cps_eta={preset_eta}, but cps_eta={cps_eta} was also passed")
+        return "cps", preset_eta
+
+    try:
+        kind = _SAMPLER_ALIASES[normalized]
+    except KeyError as exc:
+        choices = ", ".join(PAPER_SAMPLER_PRESETS)
+        raise ValueError(f"Unknown sampler {sampler!r}; paper presets are: {choices}") from exc
+
+    if kind != "cps":
+        if cps_eta is not None:
+            raise ValueError(f"cps_eta only applies to Flow-CPS, not sampler={sampler!r}")
+        return kind, None
+
+    if cps_eta is None:
+        cps_eta = 0.7
+    if isinstance(cps_eta, bool):
+        raise ValueError(f"cps_eta must be a finite number in [0, 1], got {cps_eta!r}")
+    eta = float(cps_eta)
+    if not math.isfinite(eta) or not 0.0 <= eta <= 1.0:
+        raise ValueError(f"cps_eta must be a finite number in [0, 1], got {cps_eta!r}")
+    return kind, eta
+
+
+def _flow_shift(config: Any) -> float:
+    value = _config_value(config, "flow_shift", _config_value(config, "shift", 1.0))
+    shift = float(value)
+    if not math.isfinite(shift) or shift <= 0.0:
+        raise ValueError(f"The flow scheduler shift must be positive and finite, got {value!r}")
+    return shift
+
+
+def _new_flow_scheduler(scheduler_cls: type[FlowMatchEulerDiscreteScheduler], config: Any):
+    """Build a FlowMatch scheduler without silently dropping UniPC's ``flow_shift``."""
+
+    return scheduler_cls(
+        num_train_timesteps=int(_config_value(config, "num_train_timesteps", 1000)),
+        shift=_flow_shift(config),
+        use_dynamic_shifting=bool(_config_value(config, "use_dynamic_shifting", False)),
+        base_shift=_config_value(config, "base_shift", 0.5),
+        max_shift=_config_value(config, "max_shift", 1.15),
+        base_image_seq_len=int(_config_value(config, "base_image_seq_len", 256)),
+        max_image_seq_len=int(_config_value(config, "max_image_seq_len", 4096)),
+        invert_sigmas=bool(_config_value(config, "invert_sigmas", False)),
+        shift_terminal=_config_value(config, "shift_terminal"),
+        use_karras_sigmas=bool(_config_value(config, "use_karras_sigmas", False)),
+        use_exponential_sigmas=bool(_config_value(config, "use_exponential_sigmas", False)),
+        use_beta_sigmas=bool(_config_value(config, "use_beta_sigmas", False)),
+        time_shift_type=str(_config_value(config, "time_shift_type", "exponential")),
+        stochastic_sampling=False,
+    )
+
+
+class FlowMatchCPSScheduler(FlowMatchEulerDiscreteScheduler):
+    """FlowMatch sigma schedule with the coefficient-preserving transition.
+
+    At ``cps_eta=0`` this transition is the first-order rectified-flow Euler
+    update.  Positive values rotate part of the scheduler-prescribed noise
+    coefficient into freshly sampled Gaussian noise while preserving its total
+    magnitude.
+    """
+
+    _cps_eta = 0.7
+    _cps_generator: torch.Generator | list[torch.Generator] | None = None
+
+    @property
+    def cps_eta(self) -> float:
+        return self._cps_eta
+
+    def set_cps_eta(self, value: float) -> None:
+        _, eta = resolve_vbvr_sampler("cps", value)
+        assert eta is not None
+        self._cps_eta = eta
+
+    def set_cps_generator(self, generator: torch.Generator | list[torch.Generator] | None) -> None:
+        self._cps_generator = generator
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device | None = None,
+        sigmas: list[float] | None = None,
+        mu: float | None = None,
+        timesteps: list[float] | None = None,
+    ) -> None:
+        """Install the exact training-time ``T+1`` sigma grid.
+
+        Diffusers' FlowMatch Euler scheduler spaces ``T`` source sigmas from
+        raw time 1 to ``1 / num_train_timesteps`` and then appends zero. The
+        Flow-CPS rollout used for VBVR-Pro instead shifts ``linspace(1, 0,
+        T+1)``. Keeping this distinction is required to reproduce the CPS
+        inference contract.
+        """
+
+        if sigmas is not None or timesteps is not None:
+            raise ValueError("FlowMatchCPSScheduler only supports its training-time automatic sigma grid")
+        if isinstance(num_inference_steps, bool) or not isinstance(num_inference_steps, int):
+            raise ValueError(f"num_inference_steps must be a positive integer, got {num_inference_steps!r}")
+        if num_inference_steps < 1:
+            raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}")
+        if self.config.use_dynamic_shifting and mu is None:
+            raise ValueError("mu is required when use_dynamic_shifting=True")
+        if any(
+            (
+                self.config.invert_sigmas,
+                self.config.use_karras_sigmas,
+                self.config.use_exponential_sigmas,
+                self.config.use_beta_sigmas,
+            )
+        ):
+            raise ValueError("FlowMatchCPSScheduler requires the standard non-inverted linear flow sigma family")
+
+        raw_sigmas = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device, dtype=torch.float32)
+        if self.config.use_dynamic_shifting:
+            shifted_sigmas = self.time_shift(mu, 1.0, raw_sigmas)
+        else:
+            shifted_sigmas = self.shift * raw_sigmas / (1.0 + (self.shift - 1.0) * raw_sigmas)
+        if self.config.shift_terminal:
+            shifted_sigmas = self.stretch_shift_to_terminal(shifted_sigmas)
+
+        self.num_inference_steps = num_inference_steps
+        self.sigmas = shifted_sigmas
+        self.timesteps = shifted_sigmas[:-1] * self.config.num_train_timesteps
+        self._step_index = None
+        self._begin_index = None
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        s_churn: float = 0.0,
+        s_tmin: float = 0.0,
+        s_tmax: float = float("inf"),
+        s_noise: float = 1.0,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        per_token_timesteps: torch.Tensor | None = None,
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.Tensor]:
+        """Advance one Flow-CPS step using the scheduler's current sigma pair."""
+
+        if per_token_timesteps is not None:
+            raise ValueError("FlowMatchCPSScheduler does not support per_token_timesteps")
+        if s_churn != 0.0 or s_tmin != 0.0 or s_tmax != float("inf") or s_noise != 1.0:
+            raise ValueError("FlowMatchCPSScheduler does not support Euler churn parameters")
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError("Pass a value from scheduler.timesteps, not an integer step index")
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        sigma = self.sigmas[self.step_index].to(device=sample.device, dtype=torch.float32)
+        sigma_next = self.sigmas[self.step_index + 1].to(device=sample.device, dtype=torch.float32)
+        sample_fp32 = sample.to(torch.float32)
+        model_output_fp32 = model_output.to(torch.float32)
+
+        predicted_clean = sample_fp32 - sigma * model_output_fp32
+        predicted_noise = sample_fp32 + (1.0 - sigma) * model_output_fp32
+        angle = self.cps_eta * math.pi / 2.0
+        prev_sample = (1.0 - sigma_next) * predicted_clean + sigma_next * math.cos(angle) * predicted_noise
+
+        noise_std = sigma_next * math.sin(angle)
+        if self.cps_eta > 0.0 and bool(noise_std > 0.0):
+            noise = randn_tensor(
+                sample.shape,
+                generator=generator if generator is not None else self._cps_generator,
+                device=sample.device,
+                dtype=torch.float32,
+            )
+            prev_sample = prev_sample + noise_std * noise
+
+        self._step_index += 1
+        prev_sample = prev_sample.to(sample.dtype)
+        if not return_dict:
+            return (prev_sample,)
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+
+def create_vbvr_scheduler(config: Any, sampler: str, cps_eta: float | None = None):
+    """Create one of the schedulers used in the paper's matched sampler matrix."""
+
+    kind, eta = resolve_vbvr_sampler(sampler, cps_eta)
+    if kind == "cps":
+        scheduler = _new_flow_scheduler(FlowMatchCPSScheduler, config)
+        assert eta is not None
+        scheduler.set_cps_eta(eta)
+        return scheduler
+    if kind == "euler":
+        return _new_flow_scheduler(FlowMatchEulerDiscreteScheduler, config)
+
+    return UniPCMultistepScheduler.from_config(
+        config,
+        solver_order=2,
+        solver_type="bh2",
+        prediction_type="flow_prediction",
+        use_flow_sigmas=True,
+        flow_shift=_flow_shift(config),
+        predict_x0=True,
+        lower_order_final=True,
+        final_sigmas_type="zero",
+    )
+
+
+class VBVRWanPipeline(WanImageToVideoPipeline):
+    """Wan TI2V/I2V pipeline with the six VBVR-Pro paper samplers.
+
+    ``sampler`` accepts the values in :data:`PAPER_SAMPLER_PRESETS`, plus the
+    generic ``cps`` name with an arbitrary ``cps_eta`` in ``[0, 1]``.
+    """
+
+    @property
+    def active_sampler(self) -> str | None:
+        return getattr(self, "_vbvr_active_sampler", None)
+
+    @property
+    def active_cps_eta(self) -> float | None:
+        return getattr(self, "_vbvr_active_cps_eta", None)
+
+    def _original_scheduler_config(self) -> dict[str, Any]:
+        config = getattr(self, "_vbvr_original_scheduler_config", None)
+        if config is None:
+            config = dict(self.scheduler.config)
+            self._vbvr_original_scheduler_config = config
+        return config
+
+    def set_sampler(self, sampler: str, *, cps_eta: float | None = None):
+        """Install a fresh scheduler for one inference call and return it."""
+
+        kind, eta = resolve_vbvr_sampler(sampler, cps_eta)
+        scheduler = create_vbvr_scheduler(self._original_scheduler_config(), sampler, cps_eta)
+        self.scheduler = scheduler
+        self._vbvr_active_sampler = kind
+        self._vbvr_active_cps_eta = eta
+        return scheduler
+
+    def prepare_latents(
+        self,
+        image: Any,
+        batch_size: int,
+        num_channels_latents: int = 16,
+        height: int = 480,
+        width: int = 832,
+        num_frames: int = 81,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        last_image: torch.Tensor | None = None,
+    ):
+        # The training-time Flow-CPS rollout starts and transitions in the
+        # transformer dtype.  Keep that path while leaving standard Diffusers
+        # UniPC/Euler latent precision untouched.
+        if self.active_sampler == "cps":
+            transformer = self.transformer if self.transformer is not None else self.transformer_2
+            dtype = transformer.dtype
+        return super().prepare_latents(
+            image=image,
+            batch_size=batch_size,
+            num_channels_latents=num_channels_latents,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            dtype=dtype,
+            device=device,
+            generator=generator,
+            latents=latents,
+            last_image=last_image,
+        )
+
+    def __call__(
+        self,
+        *args: Any,
+        sampler: str = "cps",
+        cps_eta: float | None = None,
+        cps_generator: torch.Generator | list[torch.Generator] | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        **kwargs: Any,
+    ):
+        """Run the official I2V pipeline with a selected VBVR-Pro sampler.
+
+        ``generator`` controls the initial latent.  ``cps_generator`` can be
+        supplied separately to control per-step CPS noise, which is useful for
+        exact experiment seed contracts.  If omitted, CPS reuses ``generator``.
+        """
+
+        scheduler = self.set_sampler(sampler, cps_eta=cps_eta)
+        if not isinstance(scheduler, FlowMatchCPSScheduler):
+            if cps_generator is not None:
+                raise ValueError("cps_generator only applies to a Flow-CPS sampler")
+            return super().__call__(*args, generator=generator, **kwargs)
+
+        scheduler.set_cps_generator(cps_generator if cps_generator is not None else generator)
+        try:
+            return super().__call__(*args, generator=generator, **kwargs)
+        finally:
+            scheduler.set_cps_generator(None)
+
+
+__all__ = [
+    "FlowMatchCPSScheduler",
+    "PAPER_SAMPLER_PRESETS",
+    "VBVRWanPipeline",
+    "create_vbvr_scheduler",
+    "resolve_vbvr_sampler",
+]