refactor(flux2): simplify offload lifecycle

Author: Watebear

lightx2v/common/offload/event_manager.py

@@ -28,13 +28,14 @@ def __init__(self, offload_granularity, load_stream=None, compute_stream=None):
         self.device_module = torch_device_module
         self._ready_events = [torch_device_module.Event() for _ in range(self._EVENT_SLOT_COUNT)]
         self._free_events = [torch_device_module.Event() for _ in range(self._EVENT_SLOT_COUNT)]
-        self._reset_slot_state()
+        self.reset_slots()
 
     @property
     def slot_count(self):
         return self._EVENT_SLOT_COUNT
 
-    def _reset_slot_state(self):
+    def reset_slots(self):
+        """Reset slot bookkeeping; synchronize pending device work first."""
         self._slot_pending = [False] * self.slot_count
         self._slot_ready_waited = [False] * self.slot_count
         self._slot_free_recorded = [False] * self.slot_count
@@ -117,13 +118,6 @@ def record_free(self, slot_idx, stream=None):
         self._slot_ready_waited[slot_idx] = False
         self._slot_free_recorded[slot_idx] = True
 
-    def flush(self):
-        """Drain the offload streams and reset the slot state."""
-        self.cuda_load_stream.synchronize()
-        if self.compute_stream is not self.cuda_load_stream:
-            self.compute_stream.synchronize()
-        self._reset_slot_state()
-
     def init_block_slabs(self, block_slabs, staging_raw=None):
         """Prepare the shared device buffer used for block-slab copies."""
         block_slabs = dict(block_slabs or {})

lightx2v/models/networks/flux2/model.py

@@ -31,8 +31,7 @@ def __init__(self, config, model_path, device):
         # event offload, with block-level CPU offload, BF16 Default weights, and no LoRA,
         # lazy loading, or tensor parallelism.
         self.use_block_slab_offload = self.config.get("offload_use_block_slab", False)
-        self._offload_weights_loaded = False
-        self._offload_weights_preparing = False
+        self._offload_weights_active = False
         self.in_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 64))
         self.attention_kwargs = {}
         self._combined_img_ids_cache = None
@@ -298,61 +297,43 @@ def _init_offload_manager(self):
             if single_slabs:
                 self.transformer_infer.offload_manager_single.init_block_slabs(single_slabs)
 
-    def _prepare_infer_weights(self):
-        if not self.cpu_offload or self._offload_weights_loaded:
+    def prepare_offload_weights(self):
+        """Load weights kept resident for one runner invocation."""
+        if not self.cpu_offload:
             return
-        if self._offload_weights_preparing:
-            raise RuntimeError("Flux2 offload weights are already being prepared")
+        if self._offload_weights_active:
+            raise RuntimeError("Flux2 offload weights are already active")
 
-        self._offload_weights_preparing = True
-        try:
-            if self.offload_granularity == "model":
-                self.to_cuda()
-            else:
-                # These weights are used on every diffusion step, so keep them
-                # resident for the complete denoising loop.
-                preserve_weight_module_cpu_tensors(self.pre_weight)
-                preserve_weight_module_cpu_tensors(self.post_weight)
-                self.pre_weight.to_cuda()
-                self.post_weight.to_cuda()
-                self.transformer_weights.non_block_weights_to_cuda()
-                self.transformer_weights.resident_blocks_to_cuda()
-        except BaseException as prepare_error:
-            try:
-                self.force_cleanup_offload_weights()
-            except BaseException as cleanup_error:
-                if hasattr(prepare_error, "add_note"):
-                    prepare_error.add_note(f"Flux2 offload cleanup also failed: {cleanup_error!r}")
-            raise
+        # Mark the model active before moving weights so runner cleanup also
+        # handles a partially completed preparation.
+        self._offload_weights_active = True
+        if self.offload_granularity == "model":
+            self.to_cuda()
         else:
-            self._offload_weights_loaded = True
-            self._offload_weights_preparing = False
-
-    def _release_infer_weights(self):
-        if not self.cpu_offload or self.scheduler.step_index != self.scheduler.infer_steps - 1:
-            return
-        self.force_cleanup_offload_weights()
-
-    def _flush_event_offload_managers(self):
-        if not getattr(self.transformer_infer, "use_event_offload", False):
-            return
-        for manager_name in ("offload_manager_double", "offload_manager_single"):
-            manager = getattr(self.transformer_infer, manager_name, None)
-            if manager is not None and hasattr(manager, "flush"):
-                manager.flush()
+            # These weights are used on every diffusion step, so keep them
+            # resident for the complete denoising loop.
+            preserve_weight_module_cpu_tensors(self.pre_weight)
+            preserve_weight_module_cpu_tensors(self.post_weight)
+            self.pre_weight.to_cuda()
+            self.post_weight.to_cuda()
+            self.transformer_weights.non_block_weights_to_cuda()
+            self.transformer_weights.resident_blocks_to_cuda()
 
     def force_cleanup_offload_weights(self):
         """Release loaded offload weights and reset event-slot state.
 
         This method is intentionally idempotent and may be called from a
         runner ``finally`` block after a short run, cancellation, or error.
         """
-        if not self.cpu_offload or not (self._offload_weights_loaded or self._offload_weights_preparing):
-            return False
+        if not self.cpu_offload or not self._offload_weights_active:
+            return
 
-        self._flush_event_offload_managers()
         # Device execution and non-blocking H2D copies are asynchronous.
         self._sync_device()
+        if getattr(self.transformer_infer, "use_event_offload", False):
+            self.transformer_infer.offload_manager_double.reset_slots()
+            self.transformer_infer.offload_manager_single.reset_slots()
+
         if self.offload_granularity == "model":
             self.to_cpu()
         else:
@@ -361,9 +342,7 @@ def force_cleanup_offload_weights(self):
             self.transformer_weights.release_non_block_weights()
             self.transformer_weights.release_resident_blocks()
 
-        self._offload_weights_loaded = False
-        self._offload_weights_preparing = False
-        return True
+        self._offload_weights_active = False
 
     def _get_combined_img_ids(self, img_ids, input_image_ids):
         cached = self._combined_img_ids_cache
@@ -459,8 +438,6 @@ def _init_infer_class(self):
     @compiled_method()
     @torch.no_grad()
     def infer(self, inputs):
-        self._prepare_infer_weights()
-
         latents = self.scheduler.latents
         do_cfg = self.config.get("enable_cfg", True) and self.config.get("sample_guide_scale", 1.0) > 1.0
 
@@ -544,8 +521,6 @@ def infer(self, inputs):
             )
             self.scheduler.noise_pred = noise_pred
 
-        self._release_infer_weights()
-
 
 class Flux2DevTransformerModel(_Flux2TransformerModelBase):
     """Flux2 Dev transformer: single forward pass with embedded guidance (no CFG)."""
@@ -571,8 +546,6 @@ def _init_infer_class(self):
     @compiled_method()
     @torch.no_grad()
     def infer(self, inputs):
-        self._prepare_infer_weights()
-
         latents = self.scheduler.latents
         txt_ids = inputs["text_encoder_output"].get("text_ids", None)
         img_ids = getattr(self.scheduler, "latent_image_ids", None)
@@ -585,5 +558,3 @@ def infer(self, inputs):
             img_ids=img_ids,
         )
         self.scheduler.noise_pred = noise_pred
-
-        self._release_infer_weights()

lightx2v/models/runners/flux2/flux2_runner.py

@@ -247,7 +247,11 @@ def run(self, total_steps=None):
             total_steps = self.model.scheduler.infer_steps
         run_error = None
         try:
-            for step_index in range(total_steps):
+            step_indices = range(total_steps)
+            if step_indices:
+                self.model.prepare_offload_weights()
+
+            for step_index in step_indices:
                 logger.info(f"==> step_index: {step_index + 1} / {total_steps}")
 
                 with ProfilingContext4DebugL1("step_pre"):
@@ -267,15 +271,13 @@ def run(self, total_steps=None):
             run_error = caught_error
             raise
         finally:
-            cleanup = getattr(self.model, "force_cleanup_offload_weights", None)
-            if cleanup is not None:
-                try:
-                    cleanup()
-                except BaseException as cleanup_error:
-                    if run_error is None:
-                        raise
-                    if hasattr(run_error, "add_note"):
-                        run_error.add_note(f"Flux2 offload cleanup also failed: {cleanup_error!r}")
+            try:
+                self.model.force_cleanup_offload_weights()
+            except BaseException as cleanup_error:
+                if run_error is None:
+                    raise
+                if hasattr(run_error, "add_note"):
+                    run_error.add_note(f"Flux2 offload cleanup also failed: {cleanup_error!r}")
 
     def get_custom_shape(self):
         default_aspect_ratios = {