[bugfix] Make HYWorld memory retrieval deterministic

Author: zjn20030811

fastvideo/models/dits/hyworld/retrieval_context.py

@@ -22,17 +22,52 @@
 import math
 
 
-def generate_points_in_sphere(n_points: int, radius: float) -> torch.Tensor:
+def make_retrieval_generator(
+    generator: torch.Generator | list[torch.Generator] | tuple[torch.Generator, ...] | None = None,
+    seed: int | None = None,
+) -> torch.Generator:
+    """Create an isolated CPU generator for Monte Carlo memory retrieval.
+
+    The source generator is inspected only for its initial seed, so memory
+    retrieval does not advance either the diffusion generator or the global
+    RNG state. HYWorld currently supports one trajectory per batch, therefore
+    a generator sequence follows its first trajectory.
+    """
+    if isinstance(generator, (list, tuple)):
+        if not generator:
+            raise ValueError("generator list must not be empty")
+        generator = generator[0]
+
+    if generator is not None:
+        if not isinstance(generator, torch.Generator):
+            raise TypeError("generator must be a torch.Generator or a sequence of them")
+        retrieval_seed = generator.initial_seed()
+    elif seed is not None:
+        retrieval_seed = seed
+    else:
+        # Reading the initial seed does not consume the global RNG state.
+        retrieval_seed = torch.initial_seed()
+
+    return torch.Generator(device="cpu").manual_seed(retrieval_seed)
+
+
+def generate_points_in_sphere(
+    n_points: int,
+    radius: float,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
     """
     Uniformly sample points within a sphere of a specified radius.
 
     :param n_points: The number of points to generate.
     :param radius: The radius of the sphere.
+    :param generator: Optional generator used for all three random draws.
     :return: A tensor of shape (n_points, 3), representing the (x, y, z) coordinates of the points.
     """
-    samples_r = torch.rand(n_points)
-    samples_phi = torch.rand(n_points)
-    samples_u = torch.rand(n_points)
+    sampling_device = generator.device if generator is not None else None
+    samples_r = torch.rand(n_points, generator=generator, device=sampling_device)
+    samples_phi = torch.rand(n_points, generator=generator, device=sampling_device)
+    samples_u = torch.rand(n_points, generator=generator, device=sampling_device)
 
     r = radius * torch.pow(samples_r, 1 / 3)
     phi = 2 * math.pi * samples_phi

fastvideo/pipelines/stages/hyworld_denoising.py

@@ -18,7 +18,8 @@
 from fastvideo.pipelines.stages.validators import StageValidators as V
 from fastvideo.pipelines.stages.validators import VerificationResult
 from fastvideo.utils import dict_to_3d_list
-from fastvideo.models.dits.hyworld.retrieval_context import (generate_points_in_sphere, select_aligned_memory_frames)
+from fastvideo.models.dits.hyworld.retrieval_context import (generate_points_in_sphere, make_retrieval_generator,
+                                                             select_aligned_memory_frames)
 from fastvideo.models.dits.hyworld.pose import pose_to_input, compute_latent_num
 
 logger = init_logger(__name__)
@@ -156,7 +157,8 @@ def forward(
 
         # Generate local points if not provided
         if points_local is None:
-            points_local = generate_points_in_sphere(50000, 8.0).to(device)
+            retrieval_generator = make_retrieval_generator(batch.generator, batch.seed)
+            points_local = generate_points_in_sphere(50000, 8.0, generator=retrieval_generator).to(device)
         else:
             points_local = points_local.to(device)
 

fastvideo/tests/stages/test_hyworld_retrieval_rng.py

@@ -0,0 +1,120 @@
+import numpy as np
+import pytest
+import torch
+
+from fastvideo.models.dits.hyworld.retrieval_context import (
+    generate_points_in_sphere,
+    make_retrieval_generator,
+    select_aligned_memory_frames,
+)
+
+
+def _w2c_at_x(x: float) -> np.ndarray:
+    c2w = np.eye(4, dtype=np.float64)
+    c2w[0, 3] = x
+    return np.linalg.inv(c2w)
+
+
+def _symmetric_camera_path() -> np.ndarray:
+    poses = np.stack([_w2c_at_x(0.0) for _ in range(36)])
+    poses[4:8] = _w2c_at_x(1.0)
+    poses[8:12] = _w2c_at_x(-1.0)
+    poses[12:16] = _w2c_at_x(4.0)
+    return poses
+
+
+def _select_memory(points: torch.Tensor) -> list[int]:
+    return select_aligned_memory_frames(
+        _symmetric_camera_path(),
+        current_frame_idx=28,
+        memory_frames=20,
+        temporal_context_size=12,
+        pred_latent_size=4,
+        device="cpu",
+        points_local=points,
+    )
+
+
+def test_explicit_seed_is_independent_of_global_rng() -> None:
+    torch.manual_seed(11)
+    torch.rand(37)
+    points_a = generate_points_in_sphere(4096, 8.0, generator=make_retrieval_generator(seed=1234))
+
+    torch.manual_seed(99)
+    torch.rand(113)
+    points_b = generate_points_in_sphere(4096, 8.0, generator=make_retrieval_generator(seed=1234))
+
+    torch.testing.assert_close(points_a, points_b, rtol=0, atol=0)
+    assert _select_memory(points_a) == _select_memory(points_b)
+
+
+def test_retrieval_does_not_advance_source_or_global_rng() -> None:
+    source = torch.Generator(device="cpu").manual_seed(4321)
+    torch.rand(17, generator=source)
+    source_state = source.get_state().clone()
+
+    torch.manual_seed(8765)
+    global_state = torch.random.get_rng_state().clone()
+
+    points = generate_points_in_sphere(256, 2.0, generator=make_retrieval_generator(source))
+
+    assert points.shape == (256, 3)
+    assert torch.equal(source.get_state(), source_state)
+    assert torch.equal(torch.random.get_rng_state(), global_state)
+
+
+def test_same_initial_seed_ignores_source_state_progress() -> None:
+    first = torch.Generator(device="cpu").manual_seed(31415)
+    second = torch.Generator(device="cpu").manual_seed(31415)
+    torch.rand(13, generator=first)
+    torch.rand(79, generator=second)
+
+    points_first = generate_points_in_sphere(512, 3.0, generator=make_retrieval_generator(first))
+    points_second = generate_points_in_sphere(512, 3.0, generator=make_retrieval_generator(second))
+
+    torch.testing.assert_close(points_first, points_second, rtol=0, atol=0)
+
+
+def test_source_generator_takes_precedence_over_seed() -> None:
+    source = torch.Generator(device="cpu").manual_seed(17)
+
+    from_source = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator(source, seed=999))
+    expected = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator(seed=17))
+
+    torch.testing.assert_close(from_source, expected, rtol=0, atol=0)
+
+
+def test_generator_list_uses_first_trajectory_seed() -> None:
+    first = torch.Generator(device="cpu").manual_seed(7)
+    second = torch.Generator(device="cpu").manual_seed(9)
+
+    from_list = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator([first, second]))
+    from_first = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator(first))
+
+    torch.testing.assert_close(from_list, from_first, rtol=0, atol=0)
+
+
+def test_global_seed_fallback_is_reproducible_without_advancing_rng() -> None:
+    torch.manual_seed(2468)
+    global_state = torch.random.get_rng_state().clone()
+
+    fallback_points = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator())
+    explicit_points = generate_points_in_sphere(128, 1.0, generator=make_retrieval_generator(seed=2468))
+
+    torch.testing.assert_close(fallback_points, explicit_points, rtol=0, atol=0)
+    assert torch.equal(torch.random.get_rng_state(), global_state)
+
+
+def test_known_seeds_can_change_selected_symmetric_chunk() -> None:
+    points_zero = generate_points_in_sphere(4096, 8.0, generator=make_retrieval_generator(seed=0))
+    points_two = generate_points_in_sphere(4096, 8.0, generator=make_retrieval_generator(seed=2))
+
+    assert _select_memory(points_zero)[4:8] == [4, 5, 6, 7]
+    assert _select_memory(points_two)[4:8] == [8, 9, 10, 11]
+
+
+def test_invalid_generator_sequences_are_rejected() -> None:
+    with pytest.raises(ValueError, match="must not be empty"):
+        make_retrieval_generator([])
+    with pytest.raises(TypeError, match="torch.Generator"):
+        make_retrieval_generator("not-a-generator")