diff --git a/README.md b/README.md index 5b7c432..5373287 100644 --- a/README.md +++ b/README.md @@ -64,3 +64,9 @@ Exact memory/lifetime options: - `attention: "kj_head_sliced"` slices attention heads and runs the slice backend from `H3_HEAD_SLICE_BACKEND` (`sage2` by default) with `H3_HEAD_SLICE_SIZE` heads per slice (`8` by default). - `--mlp-chunks N` on `tools/serve_hot_runtime.py` or `tools/direct_t2v_preview.py` chunks H3 SwiGLU rows exactly to reduce peak activation memory. Default is `1` (disabled). + +Approximate cache options are opt-in and must be quality-gated per prompt: + +- `cache_mode: "easycache"` reuses cached denoised deltas while cumulative latent input change stays below `cache_threshold`. +- `cache_mode: "h3_cache"` reuses cached denoised deltas when the current per-step latent input change is below `cache_threshold`. +- Both modes accept `cache_start_percent`, `cache_end_percent`, and `cache_subsample_factor` in `POST /generate`; the CLI exposes equivalent `--cache-*` flags. diff --git a/src/h3_blackwell_runtime/runtime.py b/src/h3_blackwell_runtime/runtime.py index ac3a638..513648b 100644 --- a/src/h3_blackwell_runtime/runtime.py +++ b/src/h3_blackwell_runtime/runtime.py @@ -141,8 +141,14 @@ class H3HotRuntime: mux_audio: bool = True, ffmpeg_loglevel: str = "error", save_latent: str | Path | None = None, + cache_mode: str | None = None, + cache_threshold: float = 0.0, + cache_start_percent: float = 0.0, + cache_end_percent: float = 1.0, + cache_subsample_factor: int = 2, ) -> dict: stages: list[dict] = [] + cache_stats: dict = {} def timed(stage: str, fn): _sync() @@ -172,6 +178,12 @@ class H3HotRuntime: audio, steps=steps, return_audio=mux_audio, + cache_mode=cache_mode, + cache_threshold=cache_threshold, + cache_start_percent=cache_start_percent, + cache_end_percent=cache_end_percent, + cache_subsample_factor=cache_subsample_factor, + cache_stats=cache_stats, ), ) if mux_audio: @@ -241,5 +253,6 @@ class H3HotRuntime: "vae_dtype": self.config.vae_dtype, "vae_tile_size": self.config.vae_tile_size, "stages": stages, + "cache": cache_stats, "request_seconds": sum(stage["seconds"] for stage in stages), } diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index a8528d9..90a9e3e 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/src/h3_blackwell_runtime/sampler.py @@ -57,6 +57,14 @@ def _decode_audio_latent(audio_carried: torch.Tensor, *, shift_video: float = 12 return audio_carried * (shift_audio / shift_video) +def _cache_sample(x: torch.Tensor, factor: int) -> torch.Tensor: + if factor <= 1: + return x + if x.ndim == 5: + return x[..., ::factor, ::factor] + return x[..., ::factor] + + @torch.inference_mode() def sample_video_res_multistep( model, @@ -69,6 +77,12 @@ def sample_video_res_multistep( model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None, return_audio: bool = False, progress: bool = False, + cache_mode: str | None = None, + cache_threshold: float = 0.0, + cache_start_percent: float = 0.0, + cache_end_percent: float = 1.0, + cache_subsample_factor: int = 2, + cache_stats: dict | None = None, ) -> torch.Tensor: """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics.""" sigmas = beta_sigmas(steps, device=video.device) @@ -77,25 +91,68 @@ def sample_video_res_multistep( video_history_sigma = audio_history_sigma = None total_steps = len(sigmas) - 1 started = time.perf_counter() + cache_mode = None if cache_mode in {None, "", "disabled", "none"} else cache_mode + if cache_mode not in {None, "easycache", "h3_cache"}: + raise ValueError(f"Unsupported cache mode: {cache_mode}") + if cache_stats is not None: + cache_stats.update({"mode": cache_mode, "threshold": cache_threshold, "skipped_steps": 0, "rates": []}) + cache = { + "video_diff": None, + "audio_diff": None, + "video_prev": None, + "audio_prev": None, + "prev_norm": None, + "cumulative_rate": 0.0, + } for index, sigma in enumerate(sigmas[:-1], start=1): step_started = time.perf_counter() previous_index = index - 1 sigma_down = sigmas[index] - sigma_audio = _audio_sigma(sigma) - carry = sigma_audio / sigma - native_audio = audio_carried.to(torch.bfloat16) * carry - step_timesteps = None if model_timesteps is None else model_timesteps[previous_index] - hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps) - raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment) - raw_video = raw_video.to(torch.bfloat16).float() - raw_audio = raw_audio.to(torch.bfloat16) - velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1]) - velocity_audio = ( - (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16)) - + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio)) - ).float() - video_denoised = video - sigma * velocity_video - audio_denoised = audio_carried - sigma * velocity_audio + current_percent = previous_index / total_steps + can_cache = cache_mode is not None and cache_threshold > 0 and cache_start_percent <= current_percent <= cache_end_percent and cache["video_diff"] is not None + skipped = False + if can_cache: + video_now = _cache_sample(video, cache_subsample_factor) + audio_now = _cache_sample(audio_carried, cache_subsample_factor) + input_change = (video_now - cache["video_prev"]).flatten().abs().mean() + (audio_now - cache["audio_prev"]).flatten().abs().mean() + input_norm = cache["prev_norm"].clamp_min(1e-8) + rate = (input_change / input_norm).item() + if cache_mode == "easycache": + cache["cumulative_rate"] += rate + skipped = cache["cumulative_rate"] < cache_threshold + if not skipped: + cache["cumulative_rate"] = 0.0 + else: + skipped = rate < cache_threshold + if cache_stats is not None: + cache_stats["rates"].append({"step": previous_index, "rate": rate, "skipped": skipped}) + if skipped: + video_denoised = video + cache["video_diff"] + audio_denoised = audio_carried + cache["audio_diff"] + if cache_stats is not None: + cache_stats["skipped_steps"] += 1 + else: + sigma_audio = _audio_sigma(sigma) + carry = sigma_audio / sigma + native_audio = audio_carried.to(torch.bfloat16) * carry + step_timesteps = None if model_timesteps is None else model_timesteps[previous_index] + hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, _model_sigma(sigma), step_timesteps) + raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment) + raw_video = raw_video.to(torch.bfloat16).float() + raw_audio = raw_audio.to(torch.bfloat16) + velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1]) + velocity_audio = ( + (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16)) + + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio)) + ).float() + video_denoised = video - sigma * velocity_video + audio_denoised = audio_carried - sigma * velocity_audio + if cache_mode is not None: + cache["video_diff"] = (video_denoised - video).detach() + cache["audio_diff"] = (audio_denoised - audio_carried).detach() + cache["video_prev"] = _cache_sample(video, cache_subsample_factor).detach().clone() + cache["audio_prev"] = _cache_sample(audio_carried, cache_subsample_factor).detach().clone() + cache["prev_norm"] = video.flatten().abs().mean() + audio_carried.flatten().abs().mean() previous_sigma = sigmas[previous_index - 1] if previous_index else None video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, video_history_sigma, previous_sigma) audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, audio_history_sigma, previous_sigma) diff --git a/tools/direct_t2v_preview.py b/tools/direct_t2v_preview.py index f9d2ae4..378c5be 100644 --- a/tools/direct_t2v_preview.py +++ b/tools/direct_t2v_preview.py @@ -48,6 +48,11 @@ parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), d parser.add_argument("--vae-tile-size", type=int, default=int(os.getenv("H3_VAE_TILE_SIZE", "256"))) parser.add_argument("--mlp-chunks", type=int, default=int(os.getenv("H3_MLP_CHUNKS", "1"))) parser.add_argument("--mlp-chunk-threshold", type=int, default=int(os.getenv("H3_MLP_CHUNK_THRESHOLD", "4096"))) +parser.add_argument("--cache-mode", choices=("disabled", "easycache", "h3_cache"), default="disabled") +parser.add_argument("--cache-threshold", type=float, default=0.0) +parser.add_argument("--cache-start-percent", type=float, default=0.0) +parser.add_argument("--cache-end-percent", type=float, default=1.0) +parser.add_argument("--cache-subsample-factor", type=int, default=2) args = parser.parse_args() started = time.perf_counter() last_report = started @@ -101,7 +106,26 @@ if args.model_timesteps_capture is not None: for index in range(args.steps) ] want_audio = args.save_audio_latent is not None or args.audio_output is not None or args.mux_audio -sampled = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps, return_audio=want_audio, progress=args.progress) +cache_stats = {} +sampled = sample_video_res_multistep( + model, + H3PromptPacker(checkpoint), + text, + video, + audio, + steps=args.steps, + model_timesteps=model_timesteps, + return_audio=want_audio, + progress=args.progress, + cache_mode=args.cache_mode, + cache_threshold=args.cache_threshold, + cache_start_percent=args.cache_start_percent, + cache_end_percent=args.cache_end_percent, + cache_subsample_factor=args.cache_subsample_factor, + cache_stats=cache_stats, +) +if cache_stats: + report({"cache": cache_stats}) if want_audio: latent, audio_latent = sampled else: diff --git a/tools/serve_hot_runtime.py b/tools/serve_hot_runtime.py index fd12624..1ebffcd 100644 --- a/tools/serve_hot_runtime.py +++ b/tools/serve_hot_runtime.py @@ -92,6 +92,11 @@ class Handler(BaseHTTPRequestHandler): mux_audio = bool(payload.get("mux_audio", True)) ffmpeg_loglevel = payload.get("ffmpeg_loglevel", "error") save_latent = payload.get("save_latent") + cache_mode = payload.get("cache_mode") + cache_threshold = float(payload.get("cache_threshold", 0.0)) + cache_start_percent = float(payload.get("cache_start_percent", 0.0)) + cache_end_percent = float(payload.get("cache_end_percent", 1.0)) + cache_subsample_factor = int(payload.get("cache_subsample_factor", 2)) started = time.perf_counter() with runtime_lock: result = runtime.generate( @@ -106,6 +111,11 @@ class Handler(BaseHTTPRequestHandler): mux_audio=mux_audio, ffmpeg_loglevel=ffmpeg_loglevel, save_latent=save_latent, + cache_mode=cache_mode, + cache_threshold=cache_threshold, + cache_start_percent=cache_start_percent, + cache_end_percent=cache_end_percent, + cache_subsample_factor=cache_subsample_factor, ) result["wall_seconds"] = time.perf_counter() - started write_json(self, 200, result)