reuse: support cross-worker and prefix-segment reuse (#1428)

- Persist encoder outputs and request metadata in shared disk caches for
Wan,
  Qwen-Image, and InfiniteTalk.
- Add InfiniteTalk prefix-segment reuse with cached motion boundaries
and
  previous-result prefix merging.
- Publish reuse caches only after successful generation and reject reuse
for
  streaming or tensor-returning requests.
- Add reuse capability checks, request schema wiring, cache
configuration, and
  a request example.

Author: zhangbilang

configs/infinitetalk/infinitetalk_480p_multi_reuse.json

@@ -32,5 +32,6 @@
     "norm_output_audio": true,
     "mxfp8_fuse_enable": false,
     "use_timestep_transform": true,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/infinitetalk_480p_multi_reuse"
 }

configs/qwen_image/qwen_image_i2i_2511_reuse.json

@@ -10,5 +10,6 @@
     "CONDITION_IMAGE_SIZE": 147456,
     "USE_IMAGE_ID_IN_PROMPT": true,
     "use_compile": true,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/qwen_image_i2i_2511"
 }

configs/qwen_image/qwen_image_t2i_2512_reuse.json

@@ -8,5 +8,6 @@
     "enable_cfg": true,
     "sample_guide_scale": 4.0,
     "use_compile": true,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/qwen_image_t2i_2512"
 }

configs/wan/wan_i2v_reuse.json

@@ -10,5 +10,6 @@
     "sample_shift": 3,
     "enable_cfg": true,
     "cpu_offload": false,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/wan_i2v"
 }

configs/wan/wan_t2v_reuse.json

@@ -11,5 +11,6 @@
     "sample_shift": 8,
     "enable_cfg": true,
     "cpu_offload": false,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/wan_t2v"
 }

configs/wan22/wan_moe_i2v_reuse.json

@@ -17,5 +17,6 @@
     "offload_granularity": "model",
     "boundary": 0.900,
     "use_image_encoder": false,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/wan22_moe_i2v"
 }

configs/wan22/wan_moe_t2v_reuse.json

@@ -18,5 +18,6 @@
     "t5_cpu_offload": false,
     "vae_cpu_offload": false,
     "boundary": 0.875,
-    "enable_reuse": true
+    "enable_reuse": true,
+    "reuse_cache_path": "save_results/reuse_cache/wan22_moe_t2v"
 }

lightx2v/models/runners/base_runner.py

@@ -22,7 +22,7 @@ def __init__(self, config):
         self.input_info = None
         self.enable_reuse = config.get("enable_reuse", False)
         self.reuse = False
-        self._reuse_cache = None
+        self.reuse_prefix_segments = 0
         self._gc_frozen = False  # one-shot guard for _maybe_freeze_gc()
         self._init_modules_depth = 0
         self._warmup_done = False
@@ -68,24 +68,25 @@ def warmup(self):
         if self.config.get("warmup", False):
             raise NotImplementedError(f"Warmup is not supported for {type(self).__name__}")
 
-    def set_reuse(self, reuse):
+    def set_reuse(self, reuse, reuse_prefix_segments=0):
         if reuse and not self.enable_reuse:
             raise ValueError(f"This {type(self).__name__} service does not enable reuse")
         if reuse and self.config.get("disagg_mode"):
             raise NotImplementedError(f"{type(self).__name__} reuse does not support disaggregated inference")
+        if reuse_prefix_segments and not reuse:
+            raise ValueError("reuse_prefix_segments requires reuse=true")
+        if reuse:
+            self.check_reuse_support()
+        if reuse_prefix_segments:
+            self.check_segment_reuse_support()
         self.reuse = reuse
+        self.reuse_prefix_segments = reuse_prefix_segments
 
-    def _reuse_key(self):
-        raise NotImplementedError
-
-    def _get_reused_inputs(self):
-        reuse_cache = self._reuse_cache
-        if reuse_cache is None:
-            raise RuntimeError("No previous successful request is available for reuse")
-        if self._reuse_key() != reuse_cache["reuse_key"]:
-            raise ValueError("Reuse inputs must match the previous successful request")
-        logger.info("[Reuse] Reusing the previous request's input encoder output")
-        return reuse_cache["inputs"]
+    def check_reuse_support(self):
+        raise NotImplementedError(f"{type(self).__name__} does not support reuse")
+
+    def check_segment_reuse_support(self):
+        raise NotImplementedError(f"{type(self).__name__} does not support segment reuse")
 
     def _maybe_freeze_gc(self):
         """Move the steady-state object graph into the GC's permanent generation once."""

lightx2v/models/runners/default_runner.py

@@ -1,5 +1,7 @@
 import gc
+import json
 import os
+import shutil
 
 import numpy as np
 import requests
@@ -16,7 +18,7 @@
 from lightx2v.utils.generate_task_id import generate_task_id
 from lightx2v.utils.global_paras import CALIB
 from lightx2v.utils.profiler import *
-from lightx2v.utils.utils import fixed_shape_resize, get_optimal_patched_size_with_sp, isotropic_crop_resize, mux_audio_from_video, save_to_image, save_to_video, wan_vae_to_comfy
+from lightx2v.utils.utils import fixed_shape_resize, get_optimal_patched_size_with_sp, is_main_process, isotropic_crop_resize, mux_audio_from_video, save_to_image, save_to_video, wan_vae_to_comfy
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 torch_device_module = getattr(torch, AI_DEVICE)
@@ -88,6 +90,14 @@ def __init__(self, config):
         super().__init__(config)
         self.has_prompt_enhancer = False
         self.progress_callback = None
+        self.reuse_cache_path = self.config.get("reuse_cache_path")
+        if self.enable_reuse and not self.reuse_cache_path:
+            raise ValueError("enable_reuse requires reuse_cache_path")
+        self.reuse_cache_dir = None
+        self.reuse_cache_stage_dir = None
+        self.final_result_path = None
+        self.previous_result_path = None
+        self.work_result_path = None
         if self.config["task"] == "t2v" and self.config.get("sub_servers", {}).get("prompt_enhancer") is not None:
             self.has_prompt_enhancer = True
             if not self.check_sub_servers("prompt_enhancer"):
@@ -98,6 +108,116 @@ def __init__(self, config):
         self.set_init_device()
         self.init_scheduler()
 
+    def reuse_key(self):
+        raise NotImplementedError
+
+    def reuse_inputs_path(self, cache_dir):
+        rank = dist.get_rank() if dist.is_initialized() else 0
+        return os.path.join(cache_dir, f"inputs_rank_{rank:05d}.pt")
+
+    def reuse_input_info(self):
+        return {}
+
+    def load_reuse_state(self, map_location=AI_DEVICE):
+        manifest_path = os.path.join(self.reuse_cache_dir, "manifest.json")
+        if not os.path.isfile(manifest_path):
+            raise RuntimeError(f"No previous successful {type(self).__name__} request is available for reuse")
+        with open(manifest_path, encoding="utf-8") as f:
+            manifest = json.load(f)
+        if manifest["reuse_key"] != self.reuse_key():
+            raise ValueError("Reuse inputs must match the previous successful request")
+        self.previous_result_path = manifest["result_path"]
+        return torch.load(self.reuse_inputs_path(self.reuse_cache_dir), map_location=map_location, weights_only=True)
+
+    def load_reused_inputs(self):
+        cached = self.load_reuse_state()
+        for name, value in cached["input_info"].items():
+            setattr(self.input_info, name, value)
+        logger.info("[Reuse] Loaded the previous request's input encoder output from disk")
+        return cached["inputs"]
+
+    def save_reuse_inputs(self):
+        torch.save(
+            {"inputs": self.inputs, "input_info": self.reuse_input_info()},
+            self.reuse_inputs_path(self.reuse_cache_stage_dir),
+        )
+
+    def prepare_reuse_output(self):
+        self.reuse_cache_dir = None
+        self.reuse_cache_stage_dir = None
+        self.final_result_path = None
+        self.previous_result_path = None
+        self.work_result_path = None
+
+        output_path = self.input_info.save_result_path
+        local_output = bool(output_path) and not output_path.startswith(("http://", "https://", "rtmp://"))
+        reuse_cache_enabled = self.enable_reuse and local_output and not self.input_info.return_result_tensor
+        if self.reuse and not reuse_cache_enabled:
+            raise ValueError(f"{type(self).__name__} reuse requires a local output and return_result_tensor=false")
+        if not reuse_cache_enabled:
+            return
+
+        self.final_result_path = os.path.abspath(os.path.expanduser(output_path))
+        self.reuse_cache_dir = os.path.abspath(os.path.expanduser(self.reuse_cache_path))
+        self.reuse_cache_stage_dir = f"{self.reuse_cache_dir}.tmp"
+
+    def stage_reuse_cache(self):
+        if self.reuse_cache_dir is None:
+            return
+
+        if is_main_process():
+            shutil.rmtree(self.reuse_cache_stage_dir, ignore_errors=True)
+            os.makedirs(self.reuse_cache_stage_dir)
+        if dist.is_initialized():
+            dist.barrier()
+
+        if self.reuse:
+            shutil.copy2(
+                self.reuse_inputs_path(self.reuse_cache_dir),
+                self.reuse_inputs_path(self.reuse_cache_stage_dir),
+            )
+        else:
+            self.save_reuse_inputs()
+
+        if dist.is_initialized():
+            dist.barrier()
+        if is_main_process():
+            with open(os.path.join(self.reuse_cache_stage_dir, "manifest.json"), "w", encoding="utf-8") as f:
+                json.dump(
+                    {"reuse_key": self.reuse_key(), "result_path": self.final_result_path},
+                    f,
+                    ensure_ascii=False,
+                    sort_keys=True,
+                )
+
+    def commit_reuse_result(self):
+        if self.reuse_cache_dir is None or not is_main_process():
+            return
+
+        cache_backup_dir = f"{self.reuse_cache_dir}.old"
+        shutil.rmtree(cache_backup_dir, ignore_errors=True)
+        cache_backed_up = os.path.isdir(self.reuse_cache_dir)
+        if cache_backed_up:
+            os.replace(self.reuse_cache_dir, cache_backup_dir)
+        try:
+            os.replace(self.reuse_cache_stage_dir, self.reuse_cache_dir)
+            if self.work_result_path is not None:
+                os.replace(self.work_result_path, self.final_result_path)
+        except Exception:
+            shutil.rmtree(self.reuse_cache_dir, ignore_errors=True)
+            if cache_backed_up:
+                os.replace(cache_backup_dir, self.reuse_cache_dir)
+            raise
+        shutil.rmtree(cache_backup_dir, ignore_errors=True)
+
+    def discard_reuse_result(self):
+        if not is_main_process():
+            return
+        if self.reuse_cache_stage_dir:
+            shutil.rmtree(self.reuse_cache_stage_dir, ignore_errors=True)
+        if self.work_result_path and os.path.exists(self.work_result_path):
+            os.remove(self.work_result_path)
+
     def warmup(self):
         if not self.config.get("warmup", False):
             return

lightx2v/models/runners/qwen_image/qwen_image_runner.py

@@ -73,6 +73,9 @@ def __init__(self, config):
         if self.text_encoder_type in ["lightllm_service", "lightllm_kernel"]:
             logger.info(f"Using LightLLM text encoder: {self.text_encoder_type}")
 
+    def check_reuse_support(self):
+        pass
+
     @ProfilingContext4DebugL1("Warmup")
     def run_warmup(self):
         task = self.config.get("task")
@@ -567,29 +570,44 @@ def _finalize_pipeline_outputs(self, input_info, images, latents=None, generator
         elif input_info.save_result_path is not None:
             return {"images": None}
 
-    def _reuse_key(self):
-        reuse_key = (self.input_info.prompt, self.input_info.negative_prompt)
+    def reuse_key(self):
+        reuse_key = {
+            "prompt": self.input_info.prompt,
+            "negative_prompt": self.input_info.negative_prompt,
+        }
         if self.config["task"] == "i2i":
-            reuse_key += (tuple(self.input_info.image_path.split(",")),)
+            reuse_key["image_path"] = self.input_info.image_path.split(",")
         return reuse_key
 
+    def reuse_input_info(self):
+        input_info = {"txt_seq_lens": list(self.input_info.txt_seq_lens)}
+        if self.config["task"] == "i2i":
+            input_info["original_size"] = list(self.input_info.original_size)
+        return input_info
+
     def _run_pipeline_local(self, input_info):
-        if self.reuse:
-            self.inputs = self._get_reused_inputs()
-        else:
-            self.inputs = self.run_input_encoder()
-            if self.enable_reuse:
-                self._reuse_cache = {"reuse_key": self._reuse_key(), "inputs": self.inputs}
-        if self.config["task"] == "i2i" and "image_encoder_output" in self.inputs:
-            self.input_info.image_encoder_output = self.inputs["image_encoder_output"]
-        self.set_target_shape()
-        self.set_img_shapes()
-        logger.info(f"input_info: {self.input_info}")
-        latents, generator = self.run_dit()
-        images = self.run_vae_decoder(latents)
-        self.end_run()
-        self._save_images(images, input_info, log_prefix="Image saved")
-        return self._finalize_pipeline_outputs(input_info, images, latents=latents, generator=generator)
+        self.prepare_reuse_output()
+        try:
+            self.inputs = self.load_reused_inputs() if self.reuse else self.run_input_encoder()
+            self.stage_reuse_cache()
+            if self.config["task"] == "i2i" and "image_encoder_output" in self.inputs:
+                self.input_info.image_encoder_output = self.inputs["image_encoder_output"]
+            self.set_target_shape()
+            self.set_img_shapes()
+            logger.info(f"input_info: {self.input_info}")
+            latents, generator = self.run_dit()
+            images = self.run_vae_decoder(latents)
+            self.end_run()
+            self._save_images(images, input_info, log_prefix="Image saved")
+            result = self._finalize_pipeline_outputs(input_info, images, latents=latents, generator=generator)
+            self.commit_reuse_result()
+            return result
+        except Exception:
+            self.discard_reuse_result()
+            raise
+        finally:
+            if self.input_info is not None:
+                self.end_run()
 
     def _run_pipeline_disagg_encoder(self):
         self.inputs = self.run_input_encoder()