[bugfix] flux2 adachace status consistency issue

Author: wangshankun

lightx2v/models/runners/flux2/flux2_runner.py

@@ -368,18 +368,20 @@ def run_pipeline(self, input_info):
 
         latents, generator = self.run_dit()
         images = self.run_vae_decoder(latents)
+        self.end_run()
 
         if not input_info.return_result_tensor and is_main_process():
             image = images[0]
             image.save(input_info.save_result_path)
             logger.info(f"Image saved: {input_info.save_result_path}")
 
+        del latents, generator
         torch_device_module.empty_cache()
         gc.collect()
 
-        result = {"images": images} if input_info.return_result_tensor else {"images": None}
-        self.end_run()
-        return result
+        if input_info.return_result_tensor:
+            return {"images": images}
+        return {"images": None}
 
 
 @RUNNER_REGISTER("flux2_klein")

lightx2v/models/schedulers/flux2/feature_caching/scheduler.py

@@ -11,11 +11,11 @@ def _refresh_caching_records(self):
         self.caching_records_2 = [True] * self.infer_steps
 
     def set_timesteps(self):
-        self.clear()
         super().set_timesteps()
         self._refresh_caching_records()
 
     def clear(self):
+        super().clear()
         if self.transformer_infer is not None:
             self.transformer_infer.clear()
 

lightx2v/models/schedulers/flux2/scheduler.py

@@ -229,6 +229,16 @@ def prepare_i2i(self, input_info, input_image, vae, inpaint_mask=None):
         self.input_image_latents = packed_latents
         self.input_image_ids = image_latent_ids
 
+    def clear(self):
+        self.generator = None
+        self.latents = None
+        self.inpaint_mask = None
+        self.input_latents = None
+        self.input_image_latents = None
+        self.input_image_ids = None
+        self.noise_pred = None
+        self.infer_condition = True
+
 
 class Flux2DevScheduler(Flux2Scheduler):
     """Scheduler for Flux2 Dev: adds guidance_proj pre-computation."""