[bugfix] Reset VAE tiling for each request

fastvideo/pipelines/stages/decoding.py

@@ -180,6 +180,8 @@ class DecodingStage(PipelineStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             # if fastvideo_args.vae_sp:
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:
@@ -249,6 +251,8 @@ class DecodingStage(PipelineStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             if not vae_autocast_enabled:
                 latents = latents.to(vae_dtype)
             image, cache = self.vae.streaming_decode(latents, cache, is_first_chunk)

fastvideo/pipelines/stages/encoding.py

@@ -77,6 +77,8 @@ class EncodingStage(PipelineStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             # if fastvideo_args.vae_sp:
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:

fastvideo/pipelines/stages/image_encoding.py

@@ -179,14 +179,11 @@ class Hy15ImageEncodingStage(ImageEncodingStage):
 
         self.vae = self.vae.to(device)
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
-            # The decoding stage turns tiling on and nothing turns it back off,
-            # and both stages hold the same VAE. Without this the first
-            # generation in a process encodes the reference image untiled and
-            # every later one encodes it tiled, so the same image and seed give
-            # different conditioning. Decide from the config, as the other
-            # VAE-encoding stages in this file do.
+            # Encoding and decoding share a VAE; apply the current request in both directions.
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             cond_latents = self.vae.encode(pixels).mode()
             cond_latents = cond_latents * self.vae.config.scaling_factor
         if fastvideo_args.vae_cpu_offload:
@@ -560,6 +557,8 @@ class ImageVAEEncodingStage(PipelineStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             # if fastvideo_args.vae_sp:
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:
@@ -736,6 +735,8 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             if not vae_autocast_enabled:
                 video_condition = video_condition.to(vae_dtype)
             encoder_output = self.vae.encode(video_condition)
@@ -924,6 +925,8 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             if not vae_autocast_enabled:
                 video_condition = video_condition.to(vae_dtype)
             encoder_output = self.vae.encode(video_condition)
@@ -1041,6 +1044,8 @@ class MatrixGame3ImageVAEEncodingStage(ImageVAEEncodingStage):
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
             if fastvideo_args.pipeline_config.vae_tiling:
                 self.vae.enable_tiling()
+            else:
+                self.vae.disable_tiling()
             if not vae_autocast_enabled:
                 video_condition = video_condition.to(vae_dtype)
             encoder_output = self.vae.encode(video_condition)

fastvideo/tests/stages/test_vae_tiling_reset.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+"""A shared VAE must follow this request's tiling flag, including True -> False."""
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.pipelines.stages import decoding, encoding, image_encoding
+from fastvideo.fastvideo_args import ExecutionMode
+
+
+class StatefulVAE:
+    handles_latent_denorm = True
+    spatial_compression_ratio = 1
+    scaling_factor = 1.0
+
+    def __init__(self, initial):
+        self.use_tiling = initial
+        self.use_temporal_tiling = True
+        self.tile_sample_min_height = 256
+        self.calls = []
+
+    def to(self, *args, **kwargs):
+        return self
+
+    def enable_tiling(self):
+        self.use_tiling = True
+
+    def disable_tiling(self):
+        self.use_tiling = False
+
+    def decode(self, x):
+        self.calls.append(self.use_tiling)
+        return x + (0.25 if self.use_tiling else 0)
+
+    def encode(self, x):
+        return SimpleNamespace(mean=self.decode(x))
+
+    def get_streaming_cache(self):
+        return []
+
+    def streaming_decode(self, x, cache, is_first_chunk):
+        return self.decode(x), cache
+
+
+@pytest.mark.parametrize("initial", [False, True])
+@pytest.mark.parametrize("boundary", ["decode", "streaming_decode", "encode", "image_encode"])
+def test_request_resets_vae_tiling(monkeypatch, initial, boundary):
+    for module in (decoding, encoding, image_encoding):
+        monkeypatch.setattr(module, "get_local_torch_device", lambda: torch.device("cpu"))
+    vae = StatefulVAE(initial)
+    if boundary == "image_encode":
+        stage = image_encoding.ImageVAEEncodingStage(vae)
+    else:
+        stage = encoding.EncodingStage(vae) if boundary == "encode" else decoding.DecodingStage(vae)
+    args = SimpleNamespace(
+        pipeline_config=SimpleNamespace(vae_precision="fp32", vae_decode_precision=None, vae_tiling=False),
+        mode=ExecutionMode.PREPROCESS,
+        disable_autocast=True,
+        vae_cpu_offload=False,
+    )
+    outputs = []
+    for flag in (False, True, False):
+        args.pipeline_config.vae_tiling = flag
+        x = torch.full((1, 3, 1, 2, 2), 0.5)
+        if boundary == "image_encode":
+            batch = SimpleNamespace(pil_image=x[:, :, 0], height=[2], width=[2], num_frames=[1])
+            output = stage.forward(batch, args).image_latent
+        elif boundary == "encode":
+            output = stage.forward(SimpleNamespace(latents=x), args).latents
+        elif boundary == "decode":
+            output = stage.decode(x, args)
+        else:
+            output, _ = stage.streaming_decode(x, args, is_first_chunk=True)
+        outputs.append(output)
+        assert vae.use_tiling is flag
+        # Resetting the request must preserve tile geometry and secondary configuration.
+        assert vae.use_temporal_tiling
+        assert vae.tile_sample_min_height == 256
+    assert vae.calls == [False, True, False]
+    torch.testing.assert_close(outputs[0], outputs[2], rtol=0, atol=0)