diff --git a/docs/source/en/api/pipelines/anyflow.md b/docs/source/en/api/pipelines/anyflow.md index 9e496a61113f..93e0b570a17d 100644 --- a/docs/source/en/api/pipelines/anyflow.md +++ b/docs/source/en/api/pipelines/anyflow.md @@ -167,6 +167,29 @@ export_to_video(video, "anyflow_far_v2v.mp4", fps=16) +### Locking motion priors with PhaseLock + +`AnyFlowPipeline.__call__` accepts an optional `phase_lock` argument that mitigates the loss of physical consistency in longer denoising trajectories. It implements PhaseLock from [Physics in 2-Steps: Locking Motion Priors Before Visual Refinement Erases Them](https://arxiv.org/abs/2606.06361): the *phase* of the latent spectrum carries the physically-consistent motion prior and erodes over many steps, while the *magnitude* (visual fidelity) stays stable. `PhaseLockGuidance` snapshots the phase from an early, few-step latent and re-imposes it on the per-step latent throughout the rest of the trajectory. It is training-free, has negligible overhead, and leaves the default (`phase_lock=None`) path unchanged. + +```py +import torch +from diffusers import AnyFlowPipeline +from diffusers.pipelines.anyflow import PhaseLockGuidance +from diffusers.utils import export_to_video + +pipe = AnyFlowPipeline.from_pretrained( + "nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers", torch_dtype=torch.bfloat16 +).to("cuda") + +# Capture the motion prior after 2 steps, then lock its low-frequency phase onto the +# high-fidelity latents. `freq_cutoff_ratio` keeps fine appearance detail untouched. +phase_lock = PhaseLockGuidance(prior_step=1, lock_strength=0.5, freq_cutoff_ratio=0.25) + +prompt = "A red panda eating bamboo in a forest, cinematic lighting" +video = pipe(prompt, num_inference_steps=50, num_frames=81, phase_lock=phase_lock).frames[0] +export_to_video(video, "anyflow_phase_lock.mp4", fps=16) +``` + ## Notes - Classifier-free guidance is fused into the released checkpoints, so inference does not run a second guided forward pass. Keep the default `guidance_scale=1.0` unless your own checkpoint requires otherwise. diff --git a/src/diffusers/pipelines/anyflow/__init__.py b/src/diffusers/pipelines/anyflow/__init__.py index 10603cdedc3b..7206cc56f265 100644 --- a/src/diffusers/pipelines/anyflow/__init__.py +++ b/src/diffusers/pipelines/anyflow/__init__.py @@ -22,6 +22,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: + _import_structure["phase_lock"] = ["PhaseLockGuidance"] _import_structure["pipeline_anyflow"] = ["AnyFlowPipeline"] _import_structure["pipeline_anyflow_far"] = ["AnyFlowFARPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -32,6 +33,7 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * else: + from .phase_lock import PhaseLockGuidance from .pipeline_anyflow import AnyFlowPipeline from .pipeline_anyflow_far import AnyFlowFARPipeline else: diff --git a/src/diffusers/pipelines/anyflow/phase_lock.py b/src/diffusers/pipelines/anyflow/phase_lock.py new file mode 100644 index 000000000000..f7d42fd80919 --- /dev/null +++ b/src/diffusers/pipelines/anyflow/phase_lock.py @@ -0,0 +1,184 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Training-free motion-phase locking for image-to-video diffusion, adapted from +# "Physics in 2-Steps: Locking Motion Priors Before Visual Refinement Erases Them" +# (PhaseLock, https://arxiv.org/abs/2606.06361). The paper observes that the *phase* +# of the latent spectrum carries the physically-consistent motion prior and erodes +# (~18%) over a long denoising trajectory, while the *magnitude* (visual fidelity) +# stays stable. PhaseLock captures the phase from an early, few-step latent and +# re-imposes it on later high-fidelity latents, preserving motion without retraining. + +from typing import Optional, Tuple + +import torch + + +def _spectral_dims(ndim: int) -> Tuple[int, ...]: + """Spectral axes that carry motion: temporal + spatial, never batch or channel. + + AnyFlow runs the denoising loop in the ``(B, T, C, H, W)`` layout, so the motion + axes are ``(T, H, W)`` -> ``(-4, -2, -1)``. Plain ``(B, C, H, W)`` image latents + fall back to the spatial axes. + """ + if ndim >= 5: + return (-4, -2, -1) + if ndim == 4: + return (-2, -1) + return tuple(range(1, ndim)) + + +def extract_motion_phase(latents: torch.Tensor) -> torch.Tensor: + """Return the phase angle of the latent spectrum over the motion axes. + + This is the "motion prior" PhaseLock locks onto: the argument of the complex FFT, + which encodes structure/motion independently of amplitude (visual fidelity). + """ + dims = _spectral_dims(latents.ndim) + spectrum = torch.fft.fftn(latents.float(), dim=dims) + return torch.angle(spectrum) + + +def apply_phase_lock( + latents: torch.Tensor, + prior_phase: torch.Tensor, + lock_strength: float, + freq_cutoff_ratio: Optional[float] = None, +) -> torch.Tensor: + """Re-impose ``prior_phase`` onto ``latents`` while keeping the current magnitude. + + Implements PhaseLock's Latent Delta Guidance: the current latent's magnitude + spectrum (visual fidelity) is preserved exactly, and only the phase (motion) is + nudged along the shortest arc toward the early-step prior by ``lock_strength``. + + Args: + latents (`torch.Tensor`): Current high-fidelity latent to refine in place. + prior_phase (`torch.Tensor`): Phase captured from the few-step motion prior. + lock_strength (`float`): Fraction of the phase delta to apply, in ``[0, 1]``. + ``0.0`` is an exact no-op; ``1.0`` snaps the phase onto the prior. + freq_cutoff_ratio (`float`, *optional*): When set in ``(0, 1]``, only the + lowest fraction of spectral coefficients (where physical motion lives) is + locked, leaving high-frequency appearance detail untouched. + + Returns: + `torch.Tensor`: The phase-locked latent, cast back to the input dtype. + """ + if lock_strength <= 0.0: + return latents + + dims = _spectral_dims(latents.ndim) + spectrum = torch.fft.fftn(latents.float(), dim=dims) + magnitude = torch.abs(spectrum) + current_phase = torch.angle(spectrum) + + # Shortest-arc delta toward the prior phase, robust to 2*pi wrap-around. + delta = torch.angle(torch.exp(1j * (prior_phase - current_phase))) + + weight = lock_strength + if freq_cutoff_ratio is not None: + weight = lock_strength * _low_pass_mask(spectrum.shape, dims, freq_cutoff_ratio, latents.device) + + locked_phase = current_phase + weight * delta + locked = magnitude * torch.exp(1j * locked_phase) + refined = torch.fft.ifftn(locked, dim=dims).real + return refined.to(latents.dtype) + + +def _low_pass_mask( + shape: torch.Size, + dims: Tuple[int, ...], + cutoff_ratio: float, + device: torch.device, +) -> torch.Tensor: + """Build a broadcastable low-pass mask that is 1 on low frequencies, 0 elsewhere.""" + radius_sq = torch.zeros(1, device=device) + view_shape = [1] * len(shape) + for d in dims: + n = shape[d] + freq = torch.fft.fftfreq(n, device=device) # in [-0.5, 0.5) + axis_shape = list(view_shape) + axis_shape[d] = n + radius_sq = radius_sq + (freq.view(axis_shape) * 2.0) ** 2 # normalize to [-1, 1] + radius = torch.sqrt(radius_sq) + return (radius <= cutoff_ratio).to(radius.dtype) + + +class PhaseLockGuidance: + """Stateful, training-free PhaseLock guidance for an I2V denoising loop. + + Captures the motion-prior phase from an early ("few-step") latent, then re-imposes + it on every subsequent high-fidelity latent. Designed to be invoked once per + denoising step from a pipeline's per-step latent hook — exactly the I/O exposed by + `~AnyFlowPipeline`'s ``callback_on_step_end`` contract. + + Args: + prior_step (`int`, defaults to `1`): Zero-based step index at which to capture + the motion prior. The paper's "2-step" prior corresponds to capturing after + two scheduler updates. + lock_strength (`float`, defaults to `0.5`): Fraction of the phase delta applied + per step. ``0.0`` disables the effect. + freq_cutoff_ratio (`float`, *optional*): Lock only the lowest fraction of + spectral coefficients (motion), preserving high-frequency appearance. + cutoff_step_ratio (`float`, defaults to `1.0`): Stop applying the lock after + this fraction of the trajectory, letting the final steps refine freely. + """ + + def __init__( + self, + prior_step: int = 1, + lock_strength: float = 0.5, + freq_cutoff_ratio: Optional[float] = None, + cutoff_step_ratio: float = 1.0, + ): + if not 0.0 <= lock_strength <= 1.0: + raise ValueError(f"`lock_strength` must be in [0, 1], got {lock_strength}.") + if freq_cutoff_ratio is not None and not 0.0 < freq_cutoff_ratio <= 1.0: + raise ValueError(f"`freq_cutoff_ratio` must be in (0, 1], got {freq_cutoff_ratio}.") + self.prior_step = prior_step + self.lock_strength = lock_strength + self.freq_cutoff_ratio = freq_cutoff_ratio + self.cutoff_step_ratio = cutoff_step_ratio + self._prior_phase: Optional[torch.Tensor] = None + + def reset(self) -> None: + """Drop the captured prior so the guidance can be reused across generations.""" + self._prior_phase = None + + @property + def has_prior(self) -> bool: + return self._prior_phase is not None + + def __call__(self, latents: torch.Tensor, step_index: int, num_inference_steps: int) -> torch.Tensor: + """Capture the prior or lock the latent's phase for ``step_index``. + + Returns ``latents`` unchanged until the prior has been captured, and during any + steps past ``cutoff_step_ratio``; otherwise returns the phase-locked latent. + """ + if step_index < self.prior_step: + return latents + + if self._prior_phase is None: + # The early, physically-consistent latent: snapshot its phase as the prior. + self._prior_phase = extract_motion_phase(latents) + return latents + + if step_index >= self.cutoff_step_ratio * num_inference_steps: + return latents + + return apply_phase_lock( + latents, + self._prior_phase, + self.lock_strength, + freq_cutoff_ratio=self.freq_cutoff_ratio, + ) diff --git a/src/diffusers/pipelines/anyflow/pipeline_anyflow.py b/src/diffusers/pipelines/anyflow/pipeline_anyflow.py index c3e1dbf3a459..d7928b803e11 100644 --- a/src/diffusers/pipelines/anyflow/pipeline_anyflow.py +++ b/src/diffusers/pipelines/anyflow/pipeline_anyflow.py @@ -29,6 +29,7 @@ from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline +from .phase_lock import PhaseLockGuidance from .pipeline_output import AnyFlowPipelineOutput @@ -405,6 +406,7 @@ def __call__( callback_on_step_end_tensor_inputs: List[str] = ["latents"], max_sequence_length: int = 512, use_mean_velocity: bool = True, + phase_lock: Optional[PhaseLockGuidance] = None, ): r""" The call function to the pipeline for generation. @@ -473,6 +475,11 @@ def __call__( When `True`, the flow-map model is conditioned on both the source timestep `t` and the target timestep `r` to predict a mean velocity, matching the training-time behavior. Disable to mirror raw Euler stepping (`r = t`). + phase_lock ([`~pipelines.anyflow.phase_lock.PhaseLockGuidance`], *optional*): + Training-free PhaseLock guidance. When provided, the motion-prior phase captured from an early + denoising step is re-imposed on the per-step latent throughout the trajectory, mitigating the phase + erosion that degrades physical consistency of long-horizon generations. Negligible overhead and no + effect on the default `None` path. Examples: @@ -568,6 +575,9 @@ def __call__( self.scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas, timesteps=timesteps) timesteps = self.scheduler.timesteps # length N; `step` resolves the next sigma internally. + if phase_lock is not None: + phase_lock.reset() + with self.progress_bar(total=len(timesteps)) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: @@ -610,6 +620,9 @@ def __call__( latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if phase_lock is not None: + latents = phase_lock(latents, i, len(timesteps)) + if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs or []: diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 0786186dff53..b860218d2a94 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3152,6 +3152,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class PhaseLockGuidance(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class PIAPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/pipelines/anyflow/test_anyflow_phase_lock.py b/tests/pipelines/anyflow/test_anyflow_phase_lock.py new file mode 100644 index 000000000000..f809d966bb25 --- /dev/null +++ b/tests/pipelines/anyflow/test_anyflow_phase_lock.py @@ -0,0 +1,190 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import unittest + +import torch +from transformers import AutoConfig, AutoTokenizer, T5EncoderModel + +from diffusers import ( + AnyFlowPipeline, + AnyFlowTransformer3DModel, + AutoencoderKLWan, + FlowMapEulerDiscreteScheduler, +) +from diffusers.pipelines.anyflow import PhaseLockGuidance +from diffusers.pipelines.anyflow.phase_lock import apply_phase_lock, extract_motion_phase + +from ...testing_utils import enable_full_determinism + + +enable_full_determinism() + + +def _dummy_components(): + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + torch.manual_seed(0) + scheduler = FlowMapEulerDiscreteScheduler(num_train_timesteps=1000, shift=5.0) + config = AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5") + text_encoder = T5EncoderModel(config) + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + torch.manual_seed(0) + transformer = AnyFlowTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=2, + attention_head_dim=12, + in_channels=16, + out_channels=16, + text_dim=32, + freq_dim=256, + ffn_dim=32, + num_layers=2, + cross_attn_norm=True, + rope_max_seq_len=32, + gate_value=0.25, + deltatime_type="r", + ) + return { + "transformer": transformer, + "vae": vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + } + + +def _dummy_inputs(device, seed=0): + generator = torch.Generator(device=device).manual_seed(seed) + return { + "prompt": "dance monkey", + "negative_prompt": "negative", + "generator": generator, + "num_inference_steps": 2, + "guidance_scale": 6.0, + "height": 16, + "width": 16, + "num_frames": 9, + "max_sequence_length": 16, + "output_type": "latent", + } + + +class PhaseLockGuidanceUnitTests(unittest.TestCase): + """Spectral guarantees of PhaseLock, exercised in isolation.""" + + def test_zero_strength_is_identity(self): + latents = torch.randn(1, 5, 16, 8, 8) + prior_phase = extract_motion_phase(torch.randn_like(latents)) + out = apply_phase_lock(latents, prior_phase, lock_strength=0.0) + self.assertTrue(torch.equal(out, latents)) + + def test_full_strength_preserves_magnitude_and_snaps_phase(self): + # At full strength the constructed spectrum stays conjugate-symmetric, so the + # inverse FFT is exactly real: magnitude (fidelity) is preserved bit-for-bit + # while the phase (motion) snaps onto the prior. + latents = torch.randn(1, 5, 16, 8, 8) + prior_phase = extract_motion_phase(torch.randn_like(latents)) + locked = apply_phase_lock(latents, prior_phase, lock_strength=1.0) + + mag_before = torch.abs(torch.fft.fftn(latents.float(), dim=(-4, -2, -1))) + mag_after = torch.abs(torch.fft.fftn(locked.float(), dim=(-4, -2, -1))) + self.assertTrue(torch.allclose(mag_before, mag_after, atol=1e-4)) + self.assertTrue(torch.allclose(extract_motion_phase(locked), prior_phase, atol=1e-4)) + + def test_partial_lock_keeps_magnitude_stable_and_moves_phase(self): + # PhaseLock's central claim: magnitude (visual fidelity) stays relatively stable + # while the phase (motion) is steered toward the early-step prior. + latents = torch.randn(1, 5, 16, 8, 8) + prior_phase = extract_motion_phase(torch.randn_like(latents)) + locked = apply_phase_lock(latents, prior_phase, lock_strength=0.7) + + mag_before = torch.abs(torch.fft.fftn(latents.float(), dim=(-4, -2, -1))) + mag_after = torch.abs(torch.fft.fftn(locked.float(), dim=(-4, -2, -1))) + rel_mag_err = torch.norm(mag_after - mag_before) / torch.norm(mag_before) + self.assertLess(rel_mag_err.item(), 0.1) + + # The locked phase must sit closer to the prior than the original did. + def dist_to_prior(x): + return torch.angle(torch.exp(1j * (prior_phase - extract_motion_phase(x)))).abs().mean() + + self.assertLess(dist_to_prior(locked).item(), dist_to_prior(latents).item()) + + def test_guidance_captures_then_applies(self): + guidance = PhaseLockGuidance(prior_step=0, lock_strength=0.5) + latents = torch.randn(1, 5, 16, 8, 8) + + self.assertFalse(guidance.has_prior) + captured = guidance(latents, step_index=0, num_inference_steps=4) + self.assertTrue(guidance.has_prior) + self.assertTrue(torch.equal(captured, latents)) # capture step is a no-op + + applied = guidance(torch.randn_like(latents), step_index=1, num_inference_steps=4) + self.assertFalse(torch.allclose(applied, latents)) + + +class AnyFlowPhaseLockIntegrationTests(unittest.TestCase): + """PhaseLock wired through the real AnyFlow denoising loop call site.""" + + def test_call_signature_exposes_phase_lock(self): + sig = inspect.signature(AnyFlowPipeline.__call__) + self.assertIn("phase_lock", sig.parameters) + + def test_pipeline_runs_and_phase_lock_alters_output(self): + device = "cpu" + pipe = AnyFlowPipeline(**_dummy_components()).to(device) + pipe.set_progress_bar_config(disable=True) + + baseline = pipe(**_dummy_inputs(device), phase_lock=None).frames + + guided = pipe( + **_dummy_inputs(device), + phase_lock=PhaseLockGuidance(prior_step=0, lock_strength=0.8), + ).frames + + self.assertEqual(guided.shape, baseline.shape) + # The guidance fires after the scheduler step at the wired call site, so the + # locked trajectory diverges from the untouched baseline by far more than the + # pipeline's own CPU run-to-run noise (~1e-3). + self.assertGreater((guided - baseline).abs().mean().item(), 1e-2) + + def test_zero_strength_is_a_noop_at_call_site(self): + device = "cpu" + pipe = AnyFlowPipeline(**_dummy_components()).to(device) + pipe.set_progress_bar_config(disable=True) + + baseline = pipe(**_dummy_inputs(device), phase_lock=None).frames + # Establish the pipeline's intrinsic run-to-run noise floor on this hardware. + noise = (pipe(**_dummy_inputs(device), phase_lock=None).frames - baseline).abs().max().item() + + noop = pipe( + **_dummy_inputs(device), + phase_lock=PhaseLockGuidance(prior_step=0, lock_strength=0.0), + ).frames + # A zero-strength lock returns the latent untouched, so the output may only + # differ from the baseline within that intrinsic noise. + self.assertLessEqual((noop - baseline).abs().max().item(), max(noise * 2.0, 1e-3)) + + +if __name__ == "__main__": + unittest.main()