diff --git a/docs/minimax_remover_usage.md b/docs/minimax_remover_usage.md index 944c5132..8abe2db7 100644 --- a/docs/minimax_remover_usage.md +++ b/docs/minimax_remover_usage.md @@ -12,7 +12,13 @@ kernelized inference on Blackwell SM120. Two precision paths ship under Both reuse the **generic** FlashRT kernels for the transformer denoise path (quantized GEMMs, fused norm/gate/residual/gelu ops, kernel attention via -FA2 / SageAttention). The VAE encode/decode is additionally accelerated by +FA2 / SageAttention). On top of the kernelised transformer, the FP8 path +also enables **training-free TeaCache step caching** (default +`--teacache-skip 3,5,7,9`): 4 of 12 interior denoise steps reuse the cached +noise prediction instead of running the transformer — mirroring the +Motus/Cosmos3/Wan2.2 TeaCache mechanism and cutting ~0.9 s with no +measurable PSNR loss (the flow-matching velocity is highly cacheable). +The VAE encode/decode is additionally accelerated by **model-specific fp16 fused CUDA kernels** in the standalone `flash_rt_minimax_remover` module (opt-in build, see [Build](#build)): `fp16_rms_norm_ncdhw` (single-pass RMSNorm, fp16-native), @@ -288,7 +294,13 @@ new pipeline so the scale is re-calibrated.** at load time; activation quantised with a calibrated static scale); - per-block LayerNorm + adaLN modulation + gate-residual fused into Triton kernels (fp32 statistics); -- `torch.nn.functional.scaled_dot_product_attention` -> FA2 / SageAttention. +- `torch.nn.functional.scaled_dot_product_attention` -> FA2 / SageAttention; +- **TeaCache step caching** (opt-in via `skip_steps=[...]` / quickstart + `--teacache-skip`): zeroth-order reuse of the cached noise prediction at + the listed denoise steps. On the first (calibration) call it is forwarded + to the diffusers reference loop; on steady-state calls it is forwarded to + the FP8 manual denoise (and baked into the captured CUDA Graph). Step 0 + (calibration) and the last step are never skipped. ### NVFP4 — `MiniMaxRemoverPipeline` (small-region only) @@ -336,8 +348,8 @@ python3 examples/minimax_remover_quickstart.py \ --masks-dir ./object_removal_data/ \ --output-dir ./out # FP8 + VAE opt (default) python3 examples/minimax_remover_quickstart.py ... --no-vae-opt # FP8, no VAE kernels -python3 examples/minimax_remover_quickstart.py ... --use-fp4 # NVFP4 -python3 examples/minimax_remover_quickstart.py ... --no-flashrt # fp16 reference +python3 examples/minimax_remover_quickstart.py ... --nvfp4-transformer # NVFP4 transformer (small-region only) +python3 examples/minimax_remover_quickstart.py ... --no-flashrt # fp16 reference (disables ALL opts) ``` Wall time is a single end-to-end segment (load -> encode -> denoise loop -> @@ -359,11 +371,12 @@ All rows compare against the non-FlashRT `--no-flashrt` fp16 reference on the | tennis (70 frames, 432x240) | FlashRT FP8 + VAE opt + CL + FP8 conv3d + fused FFN epilogue + fused bias-gate residual | 7.57 s | 2.30x | 40.0 / 36.2 dB | 0.99981 | | tennis (70 frames, 432x240) | **… + fused adaLN+quant (shared-scale QKV) + fused RMSNorm+RoPE + fused norm2+quant (FP8-only stack, `--no-nvfp4-vae`)** | **7.18 s** | **2.41x** | **40.0 / 36.0 dB** | 0.99981 | | tennis (70 frames, 432x240) | **… + NVFP4 VAE (38 conv layers → W4A4, FP4 cache)** | **6.71 s** | **2.58x** | **34.7 / 30.6 dB** | 0.99981 | -| tennis (70 frames, 432x240) | **… + NVFP4 fused norm+silu+quant + FP4 cache reuse (#1) + FP8 fused norm+silu+amax+quant (#2) + Q/K bias fused into rmsnorm+rope+quant (#3, current default)** | **6.56 s** | **2.64x** | **35.2 / 31.7 dB** | 0.99918 | -| tennis (70 frames, 432x240) | FlashRT NVFP4 transformer (`--use-fp4`) | 9.52 s | 1.82x | 7.0 / 6.2 dB | 0.00000 (broken — transformer FP4 error accumulates over 12 denoise steps) | +| tennis (70 frames, 432x240) | **… + NVFP4 fused norm+silu+quant + FP4 cache reuse (#1) + FP8 fused norm+silu+amax+quant (#2) + Q/K bias fused into rmsnorm+rope+quant (#3)** | **6.56 s** | **2.64x** | **35.2 / 31.7 dB** | 0.99918 | +| tennis (70 frames, 432x240) | **… + TeaCache {3,5,7,9} step caching (current default, `--teacache-skip 3,5,7,9`)** | **5.66 s** | **3.06x** | **35.1 / 31.5 dB** | 0.99918 | +| tennis (70 frames, 432x240) | FlashRT NVFP4 transformer (`--nvfp4-transformer`) | 9.52 s | 1.82x | 7.0 / 6.2 dB | 0.00000 (broken — transformer FP4 error accumulates over 12 denoise steps) | | bmx-trees (80 frames, 432x240) | fp16 reference (`--no-flashrt`) | 19.76 s | 1.0x | — | — | | bmx-trees (80 frames, 432x240) | FlashRT FP8 (default) | 13.24 s | **1.49x** | 35.1 / 32.0 dB | 0.99912 | -| bmx-trees (80 frames, 432x240) | FlashRT NVFP4 transformer (`--use-fp4`) | 10.72 s | 1.84x | 7.3 / 7.0 dB | 0.00000 (broken — transformer FP4 error accumulates over 12 denoise steps) | +| bmx-trees (80 frames, 432x240) | FlashRT NVFP4 transformer (`--nvfp4-transformer`) | 10.72 s | 1.84x | 7.3 / 7.0 dB | 0.00000 (broken — transformer FP4 error accumulates over 12 denoise steps) | > Tennis numbers are from a same-session serial A/B (`--no-flashrt > --no-vae-opt` vs default FlashRT FP8 stack) measured on RTX 5060 Ti, @@ -377,9 +390,15 @@ All rows compare against the non-FlashRT `--no-flashrt` fp16 reference on the Takeaways: - **The current default** (all of the below + NVFP4 VAE + #1/#2/#3 fused - norm+quant/bias) is **2.64× faster** than the fp16 reference (17.34 → - 6.56 s) with PSNR **35.2 dB** mean over 69 inpainted frames, peak VRAM - 2.57 GB (vs 3.67 GB fp16 ref, –30%). Add `--no-nvfp4-vae` and set + norm+quant/bias + **TeaCache {3,5,7,9} step caching**) is **3.06× faster** + than the fp16 reference (17.33 → 5.66 s) with PSNR **35.1 dB** mean over + 69 inpainted frames, peak VRAM 2.57 GB (vs 3.67 GB fp16 ref, –30%). The + TeaCache layer (`--teacache-skip 3,5,7,9`, on by default) reuses the + cached noise prediction at 4 of 12 interior denoise steps — MiniMax-Remover's + flow-matching velocity is highly cacheable, so this contributes ~0.9 s + (6.56 → 5.66 s) with no measurable PSNR change (35.2 → 35.1 dB). Pass + `--teacache-skip ''` to disable, or `--teacache-skip 4,7` for a + conservative 2-step schedule. Add `--no-nvfp4-vae` and set `FLASHRT_NVFP4_FUSED_NORMQUANT=0 FLASHRT_FP8_FUSED_NORMQUANT=0` for the higher-precision FP8-only stack (~40 dB, ~7.2 s) when absolute fidelity matters more than speed. The fused FFN epilogue kernel @@ -442,12 +461,12 @@ Takeaways: Combined steady-state 5.86 → 5.73 s; PSNR 35.16 dB mean / 31.73 worst (vs 35.11 at the pre-fusion baseline) — the fp32 bias add in #3 is slightly *more* precise than the fp16 bias-add it replaces. -- **NVFP4 transformer (`--use-fp4`) is unusable on full-frame latents**: cosine collapses to +- **NVFP4 transformer (`--nvfp4-transformer`) is unusable on full-frame latents**: cosine collapses to ~0.0 and PSNR to ~7 dB (median per-pixel deviation ~85/255). The FP4 quantisation error accumulates over the 12-step denoise loop in the transformer and the output drifts to black. NVFP4 transformer is only appropriate for small cropped regions, where its per-block error stays bounded. The NVFP4 - **VAE** path (above) is unaffected because the VAE is a single-pass + **VAE** path (default, always on) is unaffected because the VAE is a single-pass encoder/decoder with no iterative error accumulation. ### Transformer GEMM (NVFP4 vs fp16 matmul, single layer) @@ -503,9 +522,10 @@ the steady state; the FP8 path calibrates on the first call then freezes. |------|---------|--------| | `--vae-opt` / `--no-vae-opt` | **enabled** | Install FlashRT fp16 fused VAE kernels (`fp16_rms_norm_ncdhw`, `fp16_rms_silu_ncdhw`, NDHWC variants, fused norm+silu+amax) + channels-last 3D pipeline + running-max amax sharing between norm and conv + FP8 implicit-GEMM conv3d + `WanUpsample` cast elimination. Requires the `flash_rt_minimax_remover` module. | | `--fp8-conv` / `--no-fp8-conv` | **enabled** | Use FP8 implicit-GEMM conv3d kernel for applicable 3×3×3 causal convs (requires `--vae-opt`). Trades ~1 dB PSNR for ~14% decode speedup over channels-last-only cuDNN. | -| `--use-fp4` | off | Use NVFP4 (W4A4) transformer instead of FP8 (W8A8). Small-region only — broken on full-frame. | +| `--nvfp4-transformer` | off | Switch the **transformer** to NVFP4 (W4A4) instead of the default FP8 (W8A8). Note: the default path already uses NVFP4 for the **VAE** (single-pass, error-free); this flag adds NVFP4 to the 12-step iterative transformer denoise, where FP4 error accumulates and breaks full-frame outputs (cosine ~0.0, drifts to black). Only calibrated for small cropped regions (bbox crop). | | `--no-nvfp4-vae` | off | Disable NVFP4 W4A4 VAE conv3d (default: enabled). Uses purpose-built NVFP4 MMA kernel for eligible VAE conv layers (Ci>=192) for ~16% VAE speedup. | -| `--no-flashrt` | off | Run the reference diffusers fp16 path (no FlashRT). | +| `--teacache-skip` | `3,5,7,9` | Training-free TeaCache step caching for the transformer denoise loop: comma-separated 0-indexed step indices to SKIP (reuse the cached noise prediction from the last computed step, matching the Motus/Cosmos3/Wan2.2 TeaCache mechanism). Step 0 (FP8 calibration) and the last step are never skipped. The default `3,5,7,9` skips 4 of 12 interior steps and is **quality-neutral** on full-frame inpainting (PSNR within noise of the no-skip path, ~225 ms saved per skipped step). Use `''` to disable, or `'4,7'` for a conservative 2-step schedule. Only applies to the FlashRT paths. | +| `--no-flashrt` | off | Master "pure reference" switch: run the vanilla diffusers fp16 path **and disable ALL FlashRT optimisations** (VAE fused kernels, FP8/NVFP4 conv, NVFP4 VAE, TeaCache). This is the ground-truth baseline for PSNR/timing A/B — no need to also pass `--no-vae-opt`. | ## Environment variables @@ -545,10 +565,18 @@ output = pipeline( num_frames=len(frames), height=720, width=1280, num_inference_steps=12, + skip_steps=[3, 5, 7, 9], # optional TeaCache: skip 4 interior steps (~0.9 s faster, PSNR-neutral) ) video = output.frames ``` +`skip_steps` (optional list of int) enables training-free TeaCache step +caching: the listed denoise steps reuse the cached noise prediction instead +of running the transformer, mirroring the Motus/Cosmos3/Wan2.2 mechanism. +Step 0 (FP8 calibration) and the last step are never skipped. The default +`[3, 5, 7, 9]` (also the quickstart `--teacache-skip` default) is +quality-neutral on full-frame inpainting; pass `None` or `[]` to disable. + ### FP8 + NVFP4 VAE (recommended default — full-frame, fastest) The NVFP4 VAE optimization is **enabled by default** when using the FP8 diff --git a/examples/minimax_remover_quickstart.py b/examples/minimax_remover_quickstart.py index 14487320..58ef7e5d 100644 --- a/examples/minimax_remover_quickstart.py +++ b/examples/minimax_remover_quickstart.py @@ -11,7 +11,7 @@ It wraps a loaded diffusers MiniMax-Remover pipeline with ``flash_rt.models.minimax_remover.MiniMaxRemoverPipelineFP8`` (default) or -``MiniMaxRemoverPipeline`` (NVFP4, --use-fp4). The FP8 path rewrites the +``MiniMaxRemoverPipeline`` (NVFP4 transformer, --nvfp4-transformer). The FP8 path rewrites the transformer denoise path onto FP8 (W8A8) GEMMs with static calibration, fused norm/RoPE/gelu kernels and kernel attention; the NVFP4 path adds a graph-captured manual flow-matching loop. FP8 stays close to the fp16 @@ -59,9 +59,17 @@ ------------------------------------------------------------------ Precision note ------------------------------------------------------------------ -NVFP4 (W4A4) is calibrated for the model's large GEMMs. The fused -norm/RoPE kernels keep the precision-critical path fp32-stat. For -large full-frame latents, 4-bit weight/activation quantisation can +The default path uses FP8 (W8A8) for the transformer and NVFP4 (W4A4) +for the VAE conv layers. The NVFP4 VAE is safe because the VAE is a +single-pass encoder/decoder (no error accumulation). + +`--nvfp4-transformer` additionally switches the TRANSFORMER to NVFP4 +(W4A4). The 12-step iterative denoise accumulates FP4 error, so +full-frame latents drift to black (cosine ~0.0). Only use it for small +cropped regions (bbox crop) where per-block error stays bounded. + +The fused norm/RoPE kernels keep the precision-critical path fp32-stat. +For large full-frame latents, 4-bit weight/activation quantisation can accumulate small per-block error; if you observe colour drift on a high-resolution full-frame job, prefer the bbox-cropped regime (crop to the mask region before inference) where the quantisation grid is @@ -94,6 +102,13 @@ import torch.nn.functional as F from PIL import Image +# OpenCV's PNG/JPEG encoders are faster than PIL (notably ~20% on PNG) and +# release the GIL, so prefer cv2 when available; fall back to PIL otherwise. +# cv2's own import is ~0.19s, so only probe for its presence here and defer +# the actual import to _save_one (startup-critical path stays untaxed). +import importlib.util as _ilu +_CV2_SPEC = _ilu.find_spec("cv2") + # Make the flash_rt package importable when run directly from the repo root. ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: @@ -123,6 +138,7 @@ NUM_INFERENCE_STEPS = 12 PIPE_ITERATIONS = 6 RANDOM_SEED = 42 +TEACACHE_SKIP_DEFAULT = "3,5,7,9" # ===================================================================== @@ -397,7 +413,8 @@ def __call__(self, height: int = 720, width: int = 1280, num_frames: int = 81, images: Optional[torch.Tensor] = None, masks: Optional[torch.Tensor] = None, latents: Optional[torch.Tensor] = None, - output_type: Optional[str] = "np", iterations: int = 16): + output_type: Optional[str] = "np", iterations: int = 16, + skip_steps: Optional[List[int]] = None): device = self._execution_device batch_size = 1 transformer_dtype = torch.float16 @@ -432,14 +449,28 @@ def __call__(self, height: int = 720, width: int = 1280, num_frames: int = 81, masks_latents = (masks_latents - latents_mean) * latents_std self._num_timesteps = len(timesteps) + skip_set = set(skip_steps) if skip_steps else set() + skip_set.discard(0) + if len(timesteps) > 1: + skip_set.discard(len(timesteps) - 1) + cached_noise_pred = None + n_skipped = 0 for i, t in enumerate(timesteps): - latent_model_input = latents.to(transformer_dtype) - latent_model_input = torch.cat( - [latent_model_input, masked_latents, masks_latents], dim=1) - timestep = t.expand(latents.shape[0]) - noise_pred = self.transformer( - hidden_states=latent_model_input.half(), timestep=timestep)[0] + if i in skip_set and cached_noise_pred is not None: + noise_pred = cached_noise_pred + n_skipped += 1 + else: + latent_model_input = latents.to(transformer_dtype) + latent_model_input = torch.cat( + [latent_model_input, masked_latents, masks_latents], dim=1) + timestep = t.expand(latents.shape[0]) + noise_pred = self.transformer( + hidden_states=latent_model_input.half(), timestep=timestep)[0] + cached_noise_pred = noise_pred latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + if n_skipped: + print(f" [teacache] reused cached noise_pred at {n_skipped}/{len(timesteps)} " + f"steps (training-free step caching)") latents = latents.half() / latents_std + latents_mean video = self.vae.decode(latents, return_dict=False)[0] @@ -585,13 +616,14 @@ def parse_args() -> argparse.Namespace: help="Denoise steps (default 12).") p.add_argument("--no-flashrt", action="store_true", help="Run the reference diffusers path instead of FlashRT (for diffing).") - p.add_argument("--use-fp4", action="store_true", - help="Use NVFP4 (W4A4) instead of the default FP8 (W8A8). " - "NVFP4 is only calibrated for small cropped regions " - "(bbox crop); full-frame inpainting will produce " - "black/drift outputs. FP8 (default) is recommended " - "for full-frame inpainting (end-to-end cosine >= 0.999, " - "PSNR ~35-41 dB vs fp16).") + p.add_argument("--nvfp4-transformer", action="store_true", + help="Switch the TRANSFORMER to NVFP4 (W4A4) instead of the " + "default FP8 (W8A8). Note: the default path already uses " + "NVFP4 for the VAE (single-pass, error-free); this flag " + "adds NVFP4 to the 12-step iterative transformer denoise, " + "where FP4 error accumulates and breaks full-frame " + "outputs (cosine ~0.0, drifts to black). Only calibrated " + "for small cropped regions (bbox crop). (default: off)") p.add_argument("--vae-opt", action=argparse.BooleanOptionalAction, default=True, help="Apply FlashRT VAE optimisations: fused fp16 RMS_norm " @@ -609,6 +641,31 @@ def parse_args() -> argparse.Namespace: help="Disable NVFP4 W4A4 VAE conv3d (default: enabled). " "Uses purpose-built NVFP4 MMA kernel for eligible " "decode conv layers (Ci>=192) for ~3-7%% VAE speedup.") + p.add_argument("--teacache-skip", type=str, default=TEACACHE_SKIP_DEFAULT, + help="Training-free TeaCache step caching: comma-separated " + "0-indexed denoise step indices to SKIP (reuse the " + "cached noise prediction from the last computed step, " + "matching the Motus/Cosmos3 TeaCache mechanism). Step 0 " + "and the last step are never skipped. The default " + "'3,5,7,9' skips 4 of 12 interior steps (= 8 forward " + "passes) and is quality-neutral on full-frame " + "inpainting. Use '' to disable. (default: '3,5,7,9')") + p.add_argument("--universal-scale", action=argparse.BooleanOptionalAction, + default=True, + help="Persist FP8 activation scales to disk " + "(~/.flash_rt/calibration/) after the first run and " + "reuse them across ALL resolutions on subsequent runs " + "(margin=--universal-margin). Eliminates the per-call " + "FP8 calibration step + Triton JIT cold-start for a " + "~23%% cold-call speedup. PSNR >= 36 dB vs fp16 across " + "288x160..480x272. (default: enabled; use " + "--no-universal-scale to calibrate per-resolution)") + p.add_argument("--universal-margin", type=float, default=1.3, + help="FP8 act_scale margin for the universal-scale path " + "(default 1.3). Cross-resolution amax varies <5%% in " + "median but ~3%% of layers deviate >20%%; the enlarged " + "margin safely covers this. Lower (1.1) = higher " + "fidelity but saturation risk at unseen resolutions.") return p.parse_args() @@ -663,14 +720,21 @@ def main() -> None: pipe = build_pipeline(model_dir) - if args.vae_opt: + # --no-flashrt is the master "pure reference" switch: it disables ALL + # FlashRT optimisations (VAE fused kernels, FP8/NVFP4 conv, step + # caching) so the run is a vanilla fp16 diffusers ground truth — the + # baseline for PSNR/timing A/B. Users need not also pass --no-vae-opt. + vae_opt = args.vae_opt and not args.no_flashrt + nvfp4_vae = (vae_opt and not args.nvfp4_transformer and not args.no_nvfp4_vae) + + if vae_opt: from flash_rt.models.minimax_remover._vae_opt import install_vae_optimizations stats = install_vae_optimizations(pipe.vae, use_fp8_conv=args.fp8_conv) print(f" VAE optimised: {stats}") # NVFP4 VAE conv3d (default ON for FlashRT FP8 path; --no-nvfp4-vae to disable) - if args.vae_opt and not args.no_flashrt and not args.use_fp4 and not args.no_nvfp4_vae: + if nvfp4_vae: from flash_rt.models.minimax_remover._vae_nvfp4 import install_vae_nvfp4 nvfp4_stats = install_vae_nvfp4(pipe.vae) if nvfp4_stats.get('enabled'): @@ -679,13 +743,16 @@ def main() -> None: if args.no_flashrt: runner = pipe tag = "reference (diffusers fp16)" - elif args.use_fp4: + elif args.nvfp4_transformer: from flash_rt.models.minimax_remover import MiniMaxRemoverPipeline runner = MiniMaxRemoverPipeline(pipe) - tag = "FlashRT NVFP4 W4A4 (small-region only)" + tag = "FlashRT NVFP4 W4A4 transformer (small-region only)" else: from flash_rt.models.minimax_remover import MiniMaxRemoverPipelineFP8 - runner = MiniMaxRemoverPipelineFP8(pipe) + runner = MiniMaxRemoverPipelineFP8( + pipe, + use_universal_scale=args.universal_scale, + universal_margin=args.universal_margin) tag = "FlashRT FP8 W8A8 (full-frame)" t, h, w, _ = frames_padded.shape @@ -693,15 +760,27 @@ def main() -> None: images_tensor = images_tensor / 127.5 - 1.0 masks_infer = torch.from_numpy(masks_padded.astype(np.float32)) + skip_steps = None + # TeaCache only applies to the FlashRT FP8 path. The `--no-flashrt` + # reference is the fp16 ground truth (full N-step denoise) used for + # PSNR/timing A/B, so never skip steps there even if --teacache-skip + # carries its default value. The `--nvfp4-transformer` NVFP4 wrapper + # does not accept skip_steps, so it is also excluded. + if args.teacache_skip.strip() and not args.no_flashrt and not args.nvfp4_transformer: + skip_steps = sorted({int(s) for s in args.teacache_skip.split(",") if s.strip()}) + print(f" running inference [{tag}] (all frames at once)...") torch.cuda.reset_peak_memory_stats() torch.cuda.empty_cache() t0 = time.time() - out = runner( + call_kwargs = dict( images=images_tensor, masks=masks_infer, num_frames=t, height=h, width=w, num_inference_steps=args.num_inference_steps, generator=torch.Generator(device=DEVICE).manual_seed(RANDOM_SEED), - iterations=args.iterations).frames[0] + iterations=args.iterations) + if skip_steps is not None: + call_kwargs["skip_steps"] = skip_steps + out = runner(**call_kwargs).frames[0] torch.cuda.synchronize() elapsed = time.time() - t0 print(f" inference wall time: {elapsed:.2f}s") @@ -710,19 +789,52 @@ def main() -> None: proc_len = out_u8.shape[0] output_dir.mkdir(parents=True, exist_ok=True) - saved = inpainted = uncovered = 0 + # Pre-build the per-frame output arrays, then fan out the (CPU-bound) + # PNG/JPEG encoding + disk writes across a thread pool. zlib/jpeg release + # the GIL while compressing, so this scales near-linearly with cores -- + # the sequential loop below was a major fraction of end-to-end wall time. + tasks: List[Tuple[Path, np.ndarray, str]] = [] + inpainted = uncovered = 0 for local_i, path in enumerate(image_paths): - dst = output_dir / path.name + ext = path.suffix.lower() if local_i < proc_len and has_region[local_i]: - crop = out_u8[local_i, :orig_h, :orig_w, :] - Image.fromarray(crop, mode="RGB").save(dst) + tasks.append((output_dir / path.name, + np.ascontiguousarray(out_u8[local_i, :orig_h, :orig_w, :]), + ext)) inpainted += 1 - elif local_i >= proc_len and has_region[local_i]: - Image.fromarray(frames[local_i], mode="RGB").save(dst) - uncovered += 1 else: - Image.fromarray(frames[local_i], mode="RGB").save(dst) - saved += 1 + if local_i >= proc_len and has_region[local_i]: + uncovered += 1 + tasks.append((output_dir / path.name, + np.ascontiguousarray(frames[local_i]), ext)) + + def _save_one(task: Tuple[Path, np.ndarray, str]) -> None: + dst, arr, ext = task + if _CV2_SPEC is not None: + # Lazy import: cv2 is ~0.19s to import, paid once on first save + # rather than at process startup. Python caches it, so the cost + # is only incurred by the first thread; the import lock keeps + # concurrent threads safe. + import cv2 + # cv2 expects BGR; our arrays are RGB. + bgr = np.ascontiguousarray(arr[:, :, ::-1]) + if ext == ".png": + cv2.imwrite(str(dst), bgr, [cv2.IMWRITE_PNG_COMPRESSION, 3]) + else: + cv2.imwrite(str(dst), bgr, [cv2.IMWRITE_JPEG_QUALITY, 95]) + elif ext == ".png": + Image.fromarray(arr).save(dst, compress_level=3) + else: + Image.fromarray(arr).save(dst, quality=95) + + max_workers = max(1, min(8, (os.cpu_count() or 4), len(tasks))) + if len(tasks) <= 1 or max_workers <= 1: + for task in tasks: + _save_one(task) + else: + with ThreadPoolExecutor(max_workers=max_workers) as ex: + list(ex.map(_save_one, tasks)) + saved = len(tasks) peak_alloc = torch.cuda.max_memory_allocated() / 1024 ** 3 peak_reserved = torch.cuda.max_memory_reserved() / 1024 ** 3 diff --git a/flash_rt/models/minimax_remover/_attention.py b/flash_rt/models/minimax_remover/_attention.py index 88f34dcd..dbaebcd8 100644 --- a/flash_rt/models/minimax_remover/_attention.py +++ b/flash_rt/models/minimax_remover/_attention.py @@ -27,6 +27,7 @@ only. """ +import logging import math import os @@ -34,6 +35,14 @@ from ._kernels import rms_norm_fp32stat, rope_apply_bshd, freqs_to_cos_sin +logger = logging.getLogger(__name__) + +# Sequence lengths where SageAttention's fused int8-quant sm89 kernel +# crashed (query_scale shape mismatch). Once blacklisted, those shapes +# use the standard fp16 path (rmsnorm+rope → sage high-level API, which +# handles any seq len). Populated at runtime via try/except fallback. +_sage_blacklist: set = set() + # Optional MiniMax-Remover fused kernels — used when both are available. _fvk = None try: @@ -224,10 +233,13 @@ def __call__(self, attn, hidden_states, rotary_emb=None, and (Dd & 7) == 0) # ── Fully-fused path: RMSNorm + RoPE + int8 quant → sm89 attn ── # Eliminates the fp16 intermediate between norm+rope and QK quantize. + # Skipped for seq lengths where SageAttention's sm89 kernel is known + # to crash (query_scale shape mismatch at certain S values). _fuse_quant = (_fuse_qk and _has_fused_rmsnorm_rope_quant and mode in ("sage_fp8", "sage2") and _get_sm89() is not None - and _get_per_channel_fp8() is not None) + and _get_per_channel_fp8() is not None + and S not in _sage_blacklist) # Q/K GEMM: when the Q/K bias will be fused into the downstream # rmsnorm+rope+quant kernel (pre-norm add), fetch Q/K WITHOUT bias @@ -315,17 +327,39 @@ def __call__(self, attn, hidden_states, rotary_emb=None, out = torch.empty(B, S, H, Dd, device=q.device, dtype=q.dtype) sm89 = _get_sm89() - sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( - q_int8, k_int8, v_fp8, out, q_scale, k_scale, v_scale, - 0, 0, 2, scale, 0) - - hidden_states = out.view(B, S, H * Dd) - to_out0 = attn.to_out[0] - if no_out_bias and hasattr(to_out0, "gemm_no_bias"): - hidden_states = to_out0.gemm_no_bias(hidden_states) + try: + sm89.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( + q_int8, k_int8, v_fp8, out, q_scale, k_scale, v_scale, + 0, 0, 2, scale, 0) + except RuntimeError as e: + if "query_scale" not in str(e) and "must have shape" not in str(e): + raise + # SageAttention sm89 tiling limitation at this seq len. + # Blacklist S and fall back to the standard rmsnorm+rope + # path, which routes through sage's high-level API (handles + # any seq len correctly). + _sage_blacklist.add(S) + logger.warning( + "SageAttention fused-quant crashed at S=%d (%s); " + "falling back to standard path for this seq len " + "(~5%% slower, same precision)", S, e) + # Redo Q/K WITH bias (the nobias GEMM skipped it). + if _qk_nobias: + if _use_fp8: + q = attn.to_q.gemm_from_fp8_ext(fp8_hidden, fp8_scale) + k = attn.to_k.gemm_from_fp8_ext(fp8_hidden, fp8_scale) + else: + q = attn.to_q(hidden_states) + k = attn.to_k(hidden_states) + # Fall through to _fuse_qk standard path below. else: - hidden_states = to_out0(hidden_states) - return hidden_states + hidden_states = out.view(B, S, H * Dd) + to_out0 = attn.to_out[0] + if no_out_bias and hasattr(to_out0, "gemm_no_bias"): + hidden_states = to_out0.gemm_no_bias(hidden_states) + else: + hidden_states = to_out0(hidden_states) + return hidden_states if _fuse_qk: if not q.is_contiguous(): diff --git a/flash_rt/models/minimax_remover/_fp8_manual_denoise.py b/flash_rt/models/minimax_remover/_fp8_manual_denoise.py index f2426490..16578fdb 100644 --- a/flash_rt/models/minimax_remover/_fp8_manual_denoise.py +++ b/flash_rt/models/minimax_remover/_fp8_manual_denoise.py @@ -60,8 +60,7 @@ import torch -from ._kernels import (ada_layernorm_fp16_io, euler_step_inplace, mask_mul, - latent_normalize, latent_denormalize) +from ._kernels import ada_layernorm_fp16_io logger = logging.getLogger(__name__) @@ -165,22 +164,45 @@ def _precompute_static(self, latents, num_steps): # The graph-capturable N-step loop. # # ------------------------------------------------------------------ # def _denoise_loop_body(self, lat_buf, masked_buf, masks_buf, concat_buf, - tproj_all, mod_out, dt_all, rotary_emb, num_steps): + tproj_all, mod_out, dt_all, rotary_emb, num_steps, + skip_steps=None): + """Zeroth-order TeaCache: at skip steps reuse the cached noise_pred + (the velocity from the last COMPUTED step) and only run the cheap + euler step, mirroring the Motus/Cosmos3 TeaCache mechanism.""" C = lat_buf.shape[1] tr = self.transformer eps = self.eps + skip_set = set(skip_steps) if skip_steps else set() + skip_set.discard(0) + if num_steps > 1: + skip_set.discard(num_steps - 1) + cached_noise_pred = None for step in range(num_steps): - concat_buf[:, :C].copy_(lat_buf) - concat_buf[:, C:2 * C].copy_(masked_buf) - concat_buf[:, 2 * C:3 * C].copy_(masks_buf) - noise_pred = transformer_forward_fp8( - tr, concat_buf, tproj_all[step], mod_out[step], - rotary_emb, eps) - euler_step_inplace(lat_buf, noise_pred, dt_all[step]) + if step in skip_set and cached_noise_pred is not None: + noise_pred = cached_noise_pred + else: + concat_buf[:, :C].copy_(lat_buf) + concat_buf[:, C:2 * C].copy_(masked_buf) + concat_buf[:, 2 * C:3 * C].copy_(masks_buf) + noise_pred = transformer_forward_fp8( + tr, concat_buf, tproj_all[step], mod_out[step], + rotary_emb, eps) + cached_noise_pred = noise_pred + # Euler flow-matching step: latents += dt * noise_pred. + # PyTorch in-place add_ replaces the Triton ``euler_step_inplace`` + # kernel to avoid Triton JIT cold-start (~0.3s on the first call). + # The fp16 accumulation difference is negligible (verified + # opt-vs-Triton PSNR >= 46 dB over 8 denoise steps). + lat_buf.add_(noise_pred, alpha=dt_all[step]) def _capture_graph(self, latents, masked_latents, masks_latents, - tproj_all, mod_out, dt_all, rotary_emb, num_steps): - """Capture the full N-step denoise loop as one CUDA Graph.""" + tproj_all, mod_out, dt_all, rotary_emb, num_steps, + skip_steps=None): + """Capture the full N-step denoise loop as one CUDA Graph. + + The TeaCache skip schedule is baked into the captured graph (the + Python branch is resolved at capture time), matching how Motus + bakes its skip set before graph capture.""" device = latents.device C = latents.shape[1] B, _, T, H, W = latents.shape @@ -194,7 +216,8 @@ def _capture_graph(self, latents, masked_latents, masks_latents, def denoise(): self._denoise_loop_body( lat_buf, masked_buf, masks_buf, concat_buf, - tproj_all, mod_out, dt_all, rotary_emb, num_steps) + tproj_all, mod_out, dt_all, rotary_emb, num_steps, + skip_steps=skip_steps) # Warmup on a side stream to compile all kernels / init cuBLASLt # workspaces before capture. The FP8 scales are already frozen, so @@ -216,8 +239,9 @@ def denoise(): denoise() logger.info("MiniMax-Remover FP8: CUDA Graph captured for shape %s " - "(%d steps, warmup=%d)", tuple(latents.shape), - num_steps, n_warmup) + "(%d steps, warmup=%d, skip=%s)", tuple(latents.shape), + num_steps, n_warmup, + sorted(skip_steps) if skip_steps else None) return (graph, lat_buf, masked_buf, masks_buf, tproj_all, mod_out, dt_all, rotary_emb) @@ -225,12 +249,19 @@ def denoise(): # Public entry: run the denoise, capture-or-replay per shape. # # ------------------------------------------------------------------ # def denoise(self, latents, masked_latents, masks_latents, num_steps, - use_graph=True): + use_graph=True, skip_steps=None): """Run the N-step denoise; capture+replay a graph when use_graph. + ``skip_steps`` (optional list of int) enables zeroth-order + TeaCache: the listed steps reuse the cached noise prediction + instead of running the transformer forward. Step 0 and the last + step are never skipped. When ``use_graph`` the skip schedule is + baked into the captured graph (per-shape, per-schedule cache). + Returns the final latents tensor (same shape/dtype as ``latents``). """ - shape_key = tuple(latents.shape) + (num_steps,) + skip_key = tuple(sorted(skip_steps)) if skip_steps else () + shape_key = tuple(latents.shape) + (num_steps,) + skip_key entry = self._graphs.get(shape_key) if use_graph else None if entry is None: @@ -239,7 +270,8 @@ def denoise(self, latents, masked_latents, masks_latents, num_steps, if use_graph: entry = self._capture_graph( latents, masked_latents, masks_latents, - tproj_all, mod_out, dt_all, rotary_emb, num_steps) + tproj_all, mod_out, dt_all, rotary_emb, num_steps, + skip_steps=skip_steps) self._graphs[shape_key] = entry else: # Eager manual path (no graph) -- still avoids condition_embedder @@ -251,7 +283,8 @@ def denoise(self, latents, masked_latents, masks_latents, num_steps, lat_buf = latents.clone() self._denoise_loop_body( lat_buf, masked_latents, masks_latents, concat_buf, - tproj_all, mod_out, dt_all, rotary_emb, num_steps) + tproj_all, mod_out, dt_all, rotary_emb, num_steps, + skip_steps=skip_steps) return lat_buf (graph, lat_buf, masked_buf, masks_buf, diff --git a/flash_rt/models/minimax_remover/_fp8_pipeline.py b/flash_rt/models/minimax_remover/_fp8_pipeline.py index d9a91aac..fc4ba45f 100644 --- a/flash_rt/models/minimax_remover/_fp8_pipeline.py +++ b/flash_rt/models/minimax_remover/_fp8_pipeline.py @@ -5,21 +5,34 @@ to the fp16 reference: end-to-end cosine >= 0.999 and PSNR ~35-41 dB vs fp16 on full-frame clips. -Uses static calibration: the first inference call runs in dynamic-FP8 -calibration mode (accumulating activation amax on GPU), then freezes to a -static act_scale for all subsequent calls (zero CPU sync overhead in the -steady state). +Universal-scale fast path (default, ``use_universal_scale=True``): + The FP8 activation scales (``act_amax_max`` per Linear) are calibrated + ONCE at a representative resolution, persisted to disk + (``~/.flash_rt/calibration/``), and reused across ALL resolutions with + an enlarged margin (``universal_margin=1.3``) that absorbs + cross-resolution activation variance. This lets the very first call + skip the dynamic-FP8 calibration step entirely (saves ~0.15s) and, + combined with PyTorch-native elementwise ops (no Triton JIT cold-start), + yields a ~23% cold-call speedup for one-shot / arbitrary-resolution use. + Measured PSNR >= 36 dB vs fp16 reference across 288×160 … 480×272. + +Per-resolution calibration path (``use_universal_scale=False``): + The first ``__call__`` runs in dynamic-FP8 calibration mode + (accumulating activation amax on GPU), then freezes to a static + act_scale for all subsequent calls. """ +import hashlib +import json import logging import os +from pathlib import Path import torch logger = logging.getLogger(__name__) from flash_rt.models.minimax_remover._utils import load_fp8_kernels -from flash_rt.models.minimax_remover._kernels import mask_mul def _import_runtime_fp8(): @@ -58,11 +71,26 @@ class MiniMaxRemoverPipelineFP8: num_inference_steps: denoise steps (12) fp8_target: "all" or "ffn_only" use_bf16: run transformer in bf16 (default False, keeps fp16) - calib_margin: act_scale margin multiplier (1.1) + calib_margin: act_scale margin multiplier for per-resolution + calibration mode (1.1). Ignored when ``use_universal_scale`` + is True and a cached scale exists. + use_universal_scale: if True (default), load/persist FP8 + ``act_amax_max`` from disk so the first call skips the + dynamic-FP8 calibration step entirely. The scale is + calibrated once at a representative resolution and reused + across all resolutions with an enlarged margin + (``universal_margin``). Set False to calibrate + per-resolution every run (higher fidelity, slower first call). + universal_margin: act_scale margin for the universal-scale path + (default 1.3). Cross-resolution activation amax varies by + <5% in median, but ~3% of layers deviate >20%; the enlarged + margin safely covers this. Ignored when + ``use_universal_scale`` is False. """ def __init__(self, pipe, num_inference_steps=12, fp8_target="all", - use_bf16=False, calib_margin=1.1): + use_bf16=False, calib_margin=1.1, + use_universal_scale=True, universal_margin=1.3): self.fvk = load_fp8_kernels() (install_flashrt_fp8, set_calibration, freeze_calibration, install_fused_blocks, install_fa2_attention) = _import_runtime_fp8() @@ -71,7 +99,10 @@ def __init__(self, pipe, num_inference_steps=12, fp8_target="all", self.transformer = pipe.transformer self.num_inference_steps = num_inference_steps self.calib_margin = calib_margin + self.use_universal_scale = use_universal_scale + self.universal_margin = universal_margin self._calibrated = False + self._scale_dirty = False # need to dump scales after calibration self._set_calibration = lambda on: set_calibration(self.transformer, on) self._freeze_calibration = lambda: freeze_calibration( @@ -83,6 +114,18 @@ def __init__(self, pipe, num_inference_steps=12, fp8_target="all", logger.info("MiniMax-Remover FP8: target=%r, %d Linears -> FP8 W8A8 GEMM", fp8_target_env, n_lin) + # Try loading universal scales from disk (skip calibration on first call). + if self.use_universal_scale: + if self._load_universal_scales(): + self._calibrated = True + logger.info("MiniMax-Remover FP8: universal scale loaded " + "(margin=%.2f) — calibration skipped", + self.universal_margin) + else: + self._scale_dirty = True + logger.info("MiniMax-Remover FP8: no universal-scale cache; " + "will calibrate on first call then persist") + if use_bf16: self.transformer.to(torch.bfloat16) logger.info("MiniMax-Remover FP8: transformer -> bf16") @@ -110,6 +153,85 @@ def __init__(self, pipe, num_inference_steps=12, fp8_target="all", self._dtype = torch.bfloat16 if use_bf16 else torch.float16 self._vae_dtype = next(self.pipe.vae.parameters()).dtype + # ------------------------------------------------------------------ # + # Universal-scale disk cache (cross-resolution FP8 calibration). # + # ------------------------------------------------------------------ # + _FP8_MAX = 448.0 + + def _model_fingerprint(self): + """Hash the transformer architecture for the scale-cache key. + + Uses config + all FP8 Linear shapes so a different checkpoint or + architecture yields a different key (stale cache is impossible). + """ + from flash_rt.models.minimax_remover._fp8_linear import FlashRTFp8Linear + parts = [str(dict(self.transformer.config))] + for name, m in self.transformer.named_modules(): + if isinstance(m, FlashRTFp8Linear): + parts.append(f"{name}:{m.in_features}x{m.out_features}") + return hashlib.md5("|".join(parts).encode()).hexdigest()[:16] + + def _scale_cache_path(self): + fp = self._model_fingerprint() + d = Path.home() / ".flash_rt" / "calibration" + return d / f"minimax_remover_fp8_{fp}.json" + + def _load_universal_scales(self): + """Load persisted ``act_amax_max`` and inject frozen act_scales. + + Returns True if scales were loaded and injected, False if no cache + exists (caller should calibrate then call ``_dump_universal_scales``). + """ + from flash_rt.models.minimax_remover._fp8_linear import FlashRTFp8Linear + cache = self._scale_cache_path() + if not cache.is_file(): + return False + try: + data = json.loads(cache.read_text()) + except (json.JSONDecodeError, OSError): + logger.warning("MiniMax-Remover FP8: universal-scale cache " + "corrupt — ignoring") + return False + amax_list = data.get("amax_max", []) + margin = float(data.get("margin", self.universal_margin)) + layers = [m for m in self.transformer.modules() + if isinstance(m, FlashRTFp8Linear)] + if len(amax_list) != len(layers): + logger.warning("MiniMax-Remover FP8: universal-scale cache size " + "mismatch (%d vs %d layers) — ignoring", + len(amax_list), len(layers)) + return False + with torch.no_grad(): + for m, amax in zip(layers, amax_list): + a = float(amax) + if a <= 0: + a = float(m.weight_scale.item()) * self._FP8_MAX + sc = max(a * margin / self._FP8_MAX, 1e-12) + m.act_scale.data = torch.tensor( + [sc], dtype=torch.float32, device=m.weight_fp8.device) + m.calibrating = False + return True + + def _dump_universal_scales(self): + """Persist ``act_amax_max`` from all FP8 Linears to disk. + + Called once after the first calibration call. The raw amax is + margin-neutral — the universal margin is applied at load time. + """ + from flash_rt.models.minimax_remover._fp8_linear import FlashRTFp8Linear + amax_list = [float(m.act_amax_max.item()) + for m in self.transformer.modules() + if isinstance(m, FlashRTFp8Linear)] + cache = self._scale_cache_path() + cache.parent.mkdir(parents=True, exist_ok=True) + cache.write_text(json.dumps({ + "version": 1, + "amax_max": amax_list, + "n_layers": len(amax_list), + })) + logger.info("MiniMax-Remover FP8: universal scale persisted (%d " + "layers -> %s)", len(amax_list), cache) + @torch.no_grad() def __call__(self, *args, **kwargs): """Run the wrapped pipe, calibrating FP8 scales on the first call. @@ -121,6 +243,13 @@ def __call__(self, *args, **kwargs): benefits from the fused path instead of only multi-call ones. The cost is a single CPU sync (~1 ms) after step 1. + ``skip_steps`` (optional list of int) enables training-free + TeaCache step caching: the listed denoise steps reuse the cached + noise prediction instead of running the transformer, mirroring + the Motus/Cosmos3 TeaCache mechanism. On the first call it is + forwarded to the diffusers reference loop; on steady-state calls + it is forwarded to the FP8 manual denoise loop. + When ``FLASHRT_FP8_GRAPH=1`` and scales are frozen (call 2+), the denoise loop runs via the manual graph-capturable path (``_manual_call`` -> ``FP8ManualDenoise``). The first call always @@ -128,6 +257,10 @@ def __call__(self, *args, **kwargs): the second call and replayed thereafter. """ use_graph = os.environ.get("FLASHRT_FP8_GRAPH", "0") == "1" + skip_steps = kwargs.pop("skip_steps", None) + fwd_kwargs = dict(kwargs) + if skip_steps is not None: + fwd_kwargs["skip_steps"] = skip_steps if not self._calibrated: logger.info("MiniMax-Remover FP8: calibration mode " "(first call, dynamic FP8 + amax accumulation; " @@ -150,27 +283,33 @@ def _freeze_after_step1(_module, _inp, _out): handle = self.transformer.register_forward_hook( _freeze_after_step1) try: - result = self._orig_pipe_call(*args, **kwargs) + result = self._orig_pipe_call(*args, **fwd_kwargs) finally: handle.remove() + # Persist universal scales for future cold-start speedup. + if self._scale_dirty: + self._dump_universal_scales() + self._scale_dirty = False elif use_graph: # Frozen scales + graph requested: manual graph-capturable path. - result = self._manual_call(*args, use_graph=True, **kwargs) + result = self._manual_call(*args, use_graph=True, + skip_steps=skip_steps, **kwargs) elif os.environ.get("FLASHRT_FP8_EAGER_MANUAL", "1") == "1": # Steady-state: eager manual denoise (avoids the per-step # torch.cat of [latents, masked, masks] and the scheduler.step # CPU sync of the diffusers path). masked/masks latents are # constant across steps; _denoise_loop_body copies only the # changing latents slice into a persistent concat buffer. - result = self._manual_call(*args, use_graph=False, **kwargs) + result = self._manual_call(*args, use_graph=False, + skip_steps=skip_steps, **kwargs) else: - result = self._orig_pipe_call(*args, **kwargs) + result = self._orig_pipe_call(*args, **fwd_kwargs) return result @torch.no_grad() def _manual_call(self, images, masks, num_frames, height, width, num_inference_steps=12, generator=None, iterations=16, - output_type="np", use_graph=False): + output_type="np", use_graph=False, skip_steps=None): """Manual encode + graph-denoise + decode (mirrors the diffusers ``MinimaxRemoverPipeline.__call__`` but replaces the denoise loop with the CUDA-graph-capturable ``FP8ManualDenoise``). Requires @@ -196,7 +335,7 @@ def _manual_call(self, images, masks, num_frames, height, width, from einops import rearrange images_t = rearrange(images, "f h w c -> c f h w") images_t = pipe.resize(images_t[None, ...], height, width).to(device).to(self._vae_dtype) - masked_images = mask_mul(images_t, masks_t) + masked_images = images_t * (1.0 - masks_t) latents_mean = (torch.tensor(pipe.vae.config.latents_mean) .view(1, pipe.vae.config.z_dim, 1, 1, 1) @@ -214,7 +353,7 @@ def _manual_call(self, images, masks, num_frames, height, width, result_latents = self._graph_denoise.denoise( latents, masked_latents, masks_latents, num_inference_steps, - use_graph=use_graph) + use_graph=use_graph, skip_steps=skip_steps) result_latents = (result_latents.to(self._vae_dtype) / latents_std + latents_mean) diff --git a/tests/test_minimax_remover_smoke.py b/tests/test_minimax_remover_smoke.py index 83021d9a..070bfa0a 100644 --- a/tests/test_minimax_remover_smoke.py +++ b/tests/test_minimax_remover_smoke.py @@ -384,7 +384,11 @@ def install_fa2_attention(_transformer): pipe2 = _CallablePipe("pipe2") original_call = _CallablePipe.__call__ - wrapped = _fp8_pipeline.MiniMaxRemoverPipelineFP8(pipe1) + # use_universal_scale=False keeps this class-isolation test off the + # cross-resolution scale-cache path, which needs a realistic dict-like + # transformer.config and would write to ~/.flash_rt/calibration/. + wrapped = _fp8_pipeline.MiniMaxRemoverPipelineFP8( + pipe1, use_universal_scale=False) assert _CallablePipe.__call__ is original_call assert pipe2("unwrapped") == ("pipe2", ("unwrapped",), {}) @@ -408,6 +412,207 @@ def install_fa2_attention(_transformer): assert freeze_calls == [1.1] +# ── 5. TeaCache step-caching plumbing (skip_steps) ── + +def test_fp8_pipeline_call_forwards_skip_steps_to_pipe(monkeypatch): + """skip_steps reaches the wrapped pipe's __call__ on the calibration call. + + The single-call quickstart runs the first (calibration) call through the + diffusers reference ``__call__``, so the TeaCache schedule must be + forwarded there for it to take effect without a warm-up pass. + """ + import torch + + from flash_rt.models.minimax_remover import _fp8_pipeline + + # Exercise the delegation path (orig pipe __call__) rather than the + # eager-manual denoise default, so the stub does not need a real + # transformer/scheduler. + monkeypatch.setenv("FLASHRT_FP8_EAGER_MANUAL", "0") + monkeypatch.setattr(_fp8_pipeline, "load_fp8_kernels", lambda: object()) + + def _fake_runtime(): + def install_flashrt_fp8(_t, verbose=True, target="all"): + return 0 + + def set_calibration(_t, on): + return None + + def freeze_calibration(_t, margin=1.1): + return 3 + + def install_fused_blocks(_t): + return 0 + + def install_fa2_attention(_t): + return 0 + + return (install_flashrt_fp8, set_calibration, freeze_calibration, + install_fused_blocks, install_fa2_attention) + + monkeypatch.setattr(_fp8_pipeline, "_import_runtime_fp8", _fake_runtime) + + class _Param: + dtype = torch.float16 + + class _Transformer: + def __init__(self): + self.config = types.SimpleNamespace(eps=1e-6) + self._hooks = [] + + def to(self, _d): + return self + + def parameters(self): + return iter([_Param()]) + + def register_forward_hook(self, fn): + self._hooks.append(fn) + + class _Handle: + def remove(_self): + pass + + return _Handle() + + def _fire_hooks(self): + for fn in list(self._hooks): + fn(self, None, None) + + class _Vae: + def parameters(self): + return iter([_Param()]) + + received = [] + + class _Pipe: + def __init__(self): + self.transformer = _Transformer() + self.vae = _Vae() + + def __call__(self, *args, **kwargs): + received.append(dict(kwargs)) + # Fire the one-shot calibration freeze hook on the first call. + self.transformer._fire_hooks() + return ("ok", args, kwargs) + + # use_universal_scale=False avoids the disk-backed scale cache so this + # skip_steps-forwarding test stays hermetic (no ~/.flash_rt writes). + wrapped = _fp8_pipeline.MiniMaxRemoverPipelineFP8( + _Pipe(), use_universal_scale=False) + out = wrapped(num_frames=5, skip_steps=[3, 5, 7, 9]) + + assert out[0] == "ok" + assert received, "wrapped pipe __call__ was never invoked" + assert received[0].get("skip_steps") == [3, 5, 7, 9], ( + "skip_steps must be forwarded to the reference __call__ on the " + "calibration (first) call") + + +def test_fp8_manual_denoise_supports_skip_steps(): + """The FP8 manual denoise carries skip_steps through every entry point.""" + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + src = (root / "flash_rt/models/minimax_remover/_fp8_manual_denoise.py").read_text() + + # Zeroth-order TeaCache reuse logic (skip step reuses cached noise_pred). + assert "cached_noise_pred" in src + assert "if step in skip_set and cached_noise_pred is not None:" in src + # The public denoise() / _denoise_loop_body() / _capture_graph() all + # carry the parameter. + assert src.count("skip_steps=None") >= 3 + # The captured-graph cache key is per skip-schedule, because the skip + # set is baked into the graph at capture time (like Motus). + assert "skip_key" in src + assert "skip_steps=skip_steps" in src + + +def test_fp8_pipeline_threads_skip_steps_to_manual_call(): + """_manual_call carries skip_steps and forwards it to FP8ManualDenoise.""" + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + src = (root / "flash_rt/models/minimax_remover/_fp8_pipeline.py").read_text() + + # __call__ pops skip_steps out of the public kwargs... + assert 'skip_steps = kwargs.pop("skip_steps", None)' in src + # ...forwards it to the diffusers reference loop on the calibration call... + assert 'fwd_kwargs["skip_steps"] = skip_steps' in src + # ...and threads it into _manual_call on the steady-state branches. + assert "skip_steps=skip_steps, **kwargs" in src + # _manual_call signature carries the parameter... + assert "skip_steps=None):" in src + # ...and forwards it to the FP8ManualDenoise.denoise(). + assert "use_graph=use_graph, skip_steps=skip_steps" in src + + +def test_quickstart_teacache_default_and_reference_guard(): + """The quickstart defaults TeaCache on and keeps --no-flashrt a pure ref.""" + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + src = (root / "examples/minimax_remover_quickstart.py").read_text() + + # TeaCache default schedule is quality-neutral on full-frame inpainting. + # The quickstart factors the default into a named constant used by the + # --teacache-skip argument. + assert 'TEACACHE_SKIP_DEFAULT = "3,5,7,9"' in src + assert "default=TEACACHE_SKIP_DEFAULT" in src + # The reference __call__ carries the parameter... + assert "skip_steps: Optional[List[int]] = None" in src + # ...with zeroth-order reuse in the denoise loop. + assert "cached_noise_pred" in src + assert "if i in skip_set and cached_noise_pred is not None:" in src + # --no-flashrt is the master "pure reference" switch: it disables ALL + # FlashRT optimisations (no need to also pass --no-vae-opt). + assert "vae_opt = args.vae_opt and not args.no_flashrt" in src + # TeaCache is never applied on the reference path (ground truth = full + # N-step denoise for PSNR/timing A/B) nor on the NVFP4 transformer + # path (its wrapper does not accept skip_steps). + assert ("if args.teacache_skip.strip() and not args.no_flashrt " + "and not args.nvfp4_transformer:") in src + + +# ── 6. --nvfp4-transformer naming (replaces the confusing --use-fp4) ── + +def test_quickstart_nvfp4_transformer_flag_naming(): + """The transformer-NVFP4 flag is named --nvfp4-transformer, not --use-fp4. + + The old ``--use-fp4`` name was misleading because the default path + *already* uses NVFP4 for the VAE. The renamed flag makes explicit that + it switches the **transformer** (the 12-step iterative denoise, where + FP4 error accumulates) — distinct from the always-on NVFP4 VAE. + """ + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + src = (root / "examples/minimax_remover_quickstart.py").read_text() + + # New flag is present... + assert '"--nvfp4-transformer"' in src + assert "args.nvfp4_transformer" in src + # ...and the old confusing name is gone from the quickstart. + assert '"--use-fp4"' not in src + assert "args.use_fp4" not in src + # The NVFP4 VAE install is gated on the new attribute. + assert "not args.nvfp4_transformer and not args.no_nvfp4_vae" in src + + +def test_quickstart_nvfp4_transformer_excludes_nvfp4_vae(): + """--nvfp4-transformer disables the NVFP4 VAE (the two NVFP4 paths are + mutually exclusive: VAE NVFP4 is validated only on the FP8 transformer + path, and the NVFP4 transformer path is a standalone small-region + experiment).""" + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + src = (root / "examples/minimax_remover_quickstart.py").read_text() + + # nvfp4_vae is False when nvfp4_transformer is set. + assert "nvfp4_vae = (vae_opt and not args.nvfp4_transformer" in src + + # ── 4. Gated build: required symbols present and callable ── def _get_kernels_or_skip():