diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 27b5160a3..447d54789 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/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) diff --git a/fastvideo/pipelines/stages/encoding.py b/fastvideo/pipelines/stages/encoding.py index 13f2c1feb..f8cee3840 100644 --- a/fastvideo/pipelines/stages/encoding.py +++ b/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: diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 6ebb6db86..c4ecbcb64 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/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) diff --git a/fastvideo/tests/stages/test_vae_tiling_reset.py b/fastvideo/tests/stages/test_vae_tiling_reset.py new file mode 100644 index 000000000..05bca2b35 --- /dev/null +++ b/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)