From 38ab136c7cfc667e697e10ab95c5d59a3f58a9d2 Mon Sep 17 00:00:00 2001 From: Aryan Kumar Date: Fri, 30 Oct 2026 12:46:23 -0700 Subject: [PATCH] [bugfix]: move an offloaded transformer to the GPU and back around the Regenerate-2K refine With dit_cpu_offload and no layerwise offload, the refine stage never staged the transformer like the stock H3 stage does, so it could stay resident through the 2K VAE decode. It now moves it to the device for the refine and back to the CPU in finally, and releases layerwise-offload buffers there too. --- .../open_h3_regenerate_2k/stages/denoising.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py index 4ffef9cab..8e20a67ac 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py @@ -60,8 +60,21 @@ class Regenerate2KDenoisingStage(PipelineStage): device = get_local_torch_device() sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma) state = batch.latents.to(device) - with self._attention_kernel(device): - state = self._euler_loop(batch, layout, state, sigmas, device, transformer) + layerwise = bool(getattr(fastvideo_args, "dit_layerwise_offload", False)) + full_cpu_offload = (bool(getattr(fastvideo_args, "dit_cpu_offload", False)) and not layerwise + and not bool(getattr(fastvideo_args, "use_fsdp_inference", False))) + try: + if full_cpu_offload: + transformer.to(device) + with self._attention_kernel(device): + state = self._euler_loop(batch, layout, state, sigmas, device, transformer) + finally: + if layerwise: + manager = getattr(transformer, "_layerwise_offload_manager", None) + if manager is not None and getattr(manager, "enabled", False): + manager.release_all() + if full_cpu_offload: + transformer.to("cpu") batch.latents = state return batch