Add pipeline recipe profiling scripts for optimization benchmarking

examples/profiling/README.md

@@ -46,8 +46,10 @@ Follow the flux-fast pattern: **annotate key pipeline methods** with `torch.prof
 ### New Files
 
 ```bash
+pipeline_registry.py    # Shared pipeline configs used by the profiling CLIs
 profiling_utils.py       # Annotation helper + profiler setup
 profiling_pipelines.py   # CLI entry point with pipeline configs
+profile_pipeline_recipes.py  # Recipe benchmark CLI that emits JSON + Markdown reports
 run_profiling.sh         # Bulk launch runs for multiple pipelines
 ```
 
@@ -121,7 +123,45 @@ All configs use `output_type="latent"` by default (skip VAE decode for cleaner d
 
 **Output:** `{output_dir}/{pipeline}_{mode}.json` Chrome trace + stdout summary.
 
-### Step 3: Known Sync Issues to Validate
+### Step 3: `profile_pipeline_recipes.py` — Benchmark optimization recipes
+
+This companion CLI focuses on end-to-end tradeoffs instead of timeline traces. It runs real pipeline calls under a
+set of optimization recipes and emits both a machine-readable JSON file and a Markdown table that can be attached to
+an issue or PR.
+
+Example:
+
+```bash
+python examples/profiling/profile_pipeline_recipes.py --pipeline flux
+python examples/profiling/profile_pipeline_recipes.py --pipeline wan --full_decode
+python examples/profiling/profile_pipeline_recipes.py \
+  --pipeline qwenimage \
+  --recipe baseline \
+  --recipe layerwise_casting+attention:_native_flash \
+  --recipe compile+attention:sage
+```
+
+Supported actions:
+
+- `baseline`
+- `model_cpu_offload`
+- `group_offload_leaf`
+- `layerwise_casting`
+- `attention:<backend>`
+- `compile`
+- `vae_tiling`
+- `channels_last`
+
+Important notes:
+
+- The recipe runner is capability-checked. If a pipeline does not expose the required component or method, that recipe fails with an explicit note in the report.
+- `vae_tiling` is only meaningful with `--full_decode` because the default profiling configs use `output_type="latent"`.
+- `channels_last` is primarily useful for UNet-based pipelines and is therefore not part of the default recipe set for the current transformer-heavy registry.
+- The built-in default recipe set is intentionally conservative: `baseline`, `model_cpu_offload`, `group_offload_leaf`, `layerwise_casting`, and `compile`.
+
+**Output:** `{output_dir}/{pipeline}_recipes.json` and `{output_dir}/{pipeline}_recipes.md`.
+
+### Step 4: Known Sync Issues to Validate
 
 The profiling should surface these known/suspected issues:
 

examples/profiling/pipeline_registry.py

@@ -0,0 +1,101 @@
+import torch
+
+try:
+    from .profiling_utils import PipelineProfilingConfig
+except ImportError:
+    from profiling_utils import PipelineProfilingConfig
+
+
+PROMPT = "A cat holding a sign that says hello world"
+
+
+def build_registry():
+    """Build the pipeline config registry. Imports are deferred to avoid loading all pipelines upfront."""
+    from diffusers import Flux2KleinPipeline, FluxPipeline, LTX2Pipeline, QwenImagePipeline, WanPipeline
+
+    return {
+        "flux": PipelineProfilingConfig(
+            name="flux",
+            pipeline_cls=FluxPipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "black-forest-labs/FLUX.1-dev",
+                "torch_dtype": torch.bfloat16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "height": 1024,
+                "width": 1024,
+                "num_inference_steps": 4,
+                "guidance_scale": 3.5,
+                "output_type": "latent",
+            },
+        ),
+        "flux2": PipelineProfilingConfig(
+            name="flux2",
+            pipeline_cls=Flux2KleinPipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "black-forest-labs/FLUX.2-klein-base-9B",
+                "torch_dtype": torch.bfloat16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "height": 1024,
+                "width": 1024,
+                "num_inference_steps": 4,
+                "guidance_scale": 3.5,
+                "output_type": "latent",
+            },
+        ),
+        "wan": PipelineProfilingConfig(
+            name="wan",
+            pipeline_cls=WanPipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "Wan-AI/Wan2.1-T2V-14B-Diffusers",
+                "torch_dtype": torch.bfloat16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "negative_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
+                "height": 480,
+                "width": 832,
+                "num_frames": 81,
+                "num_inference_steps": 4,
+                "output_type": "latent",
+            },
+        ),
+        "ltx2": PipelineProfilingConfig(
+            name="ltx2",
+            pipeline_cls=LTX2Pipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "Lightricks/LTX-2",
+                "torch_dtype": torch.bfloat16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted",
+                "height": 512,
+                "width": 768,
+                "num_frames": 121,
+                "num_inference_steps": 4,
+                "guidance_scale": 4.0,
+                "output_type": "latent",
+            },
+        ),
+        "qwenimage": PipelineProfilingConfig(
+            name="qwenimage",
+            pipeline_cls=QwenImagePipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "Qwen/Qwen-Image",
+                "torch_dtype": torch.bfloat16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "negative_prompt": " ",
+                "height": 1024,
+                "width": 1024,
+                "num_inference_steps": 4,
+                "true_cfg_scale": 4.0,
+                "output_type": "latent",
+            },
+        ),
+    }

examples/profiling/profile_diffusion_gemma_quantized.py

@@ -0,0 +1,146 @@
+import argparse
+import json
+import sys
+import time
+import traceback
+
+import torch
+from transformers import AutoProcessor, BitsAndBytesConfig, DiffusionGemmaForBlockDiffusion
+
+from diffusers import BlockRefinementScheduler, DiffusionGemmaPipeline
+
+
+def _emit_stage(stage: str):
+    print(f"[profile] stage={stage}", file=sys.stderr, flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark a quantized DiffusionGemmaPipeline run")
+    parser.add_argument("--model_id", default="google/diffusiongemma-26B-A4B-it")
+    parser.add_argument("--prompt", default="Why is the sky blue?")
+    parser.add_argument("--gen_length", type=int, default=64)
+    parser.add_argument("--num_inference_steps", type=int, default=8)
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--cache_implementation", choices=["dynamic", "static"], default="dynamic")
+    parser.add_argument("--compile_decoder", action="store_true")
+    parser.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto")
+    parser.add_argument("--disable_4bit", action="store_true")
+    parser.add_argument("--max_new_tokens_report", type=int, default=120)
+    args = parser.parse_args()
+
+    if args.device == "auto":
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    else:
+        device = args.device
+
+    if device == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA requested but not available")
+
+    load_in_4bit = device == "cuda" and not args.disable_4bit
+    model_kwargs = {
+        "dtype": torch.bfloat16 if device == "cuda" else torch.float32,
+    }
+
+    if load_in_4bit:
+        model_kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+        )
+        model_kwargs["device_map"] = "auto"
+
+    current_stage = "setup"
+    result = {
+        "model_id": args.model_id,
+        "status": "fail",
+        "stage": current_stage,
+        "device": device,
+        "load_in_4bit": load_in_4bit,
+        "cache_implementation": args.cache_implementation,
+        "compile_decoder": args.compile_decoder,
+        "gen_length": args.gen_length,
+        "num_inference_steps": args.num_inference_steps,
+        "temperature": args.temperature,
+        "wall_time_s": None,
+        "tokens_returned": None,
+        "tokens_per_second": None,
+        "peak_vram_gb": round(torch.cuda.max_memory_allocated() / (1024**3), 2) if device == "cuda" else None,
+        "text_preview": None,
+        "error_type": None,
+        "error_message": None,
+    }
+
+    started_at = time.perf_counter()
+    try:
+        current_stage = "load_model"
+        result["stage"] = current_stage
+        _emit_stage(current_stage)
+        model = DiffusionGemmaForBlockDiffusion.from_pretrained(args.model_id, **model_kwargs).eval()
+        if device == "cpu":
+            model.to("cpu")
+
+        current_stage = "load_processor"
+        result["stage"] = current_stage
+        _emit_stage(current_stage)
+        processor = AutoProcessor.from_pretrained(args.model_id)
+        scheduler = BlockRefinementScheduler()
+        pipe = DiffusionGemmaPipeline(model=model, scheduler=scheduler, processor=processor)
+
+        if args.compile_decoder:
+            current_stage = "compile_decoder"
+            result["stage"] = current_stage
+            _emit_stage(current_stage)
+            if device != "cuda":
+                raise RuntimeError("Decoder compile benchmark only supports CUDA")
+            pipe.model.model.decoder = torch.compile(
+                pipe.model.model.decoder,
+                mode="reduce-overhead",
+                fullgraph=True,
+            )
+
+        if device == "cuda":
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+        current_stage = "generate"
+        result["stage"] = current_stage
+        _emit_stage(current_stage)
+        output = pipe(
+            prompt=args.prompt,
+            gen_length=args.gen_length,
+            num_inference_steps=args.num_inference_steps,
+            temperature=args.temperature,
+            cache_implementation=args.cache_implementation,
+        )
+        if device == "cuda":
+            torch.cuda.synchronize()
+        elapsed_s = time.perf_counter() - started_at
+        text = output.texts[0]
+        generated_tokens = int(output.sequences.shape[-1])
+        result.update(
+            {
+                "status": "pass",
+                "stage": current_stage,
+                "wall_time_s": round(elapsed_s, 2),
+                "tokens_returned": generated_tokens,
+                "tokens_per_second": round(generated_tokens / elapsed_s, 2) if elapsed_s > 0 else None,
+                "peak_vram_gb": round(torch.cuda.max_memory_allocated() / (1024**3), 2) if device == "cuda" else None,
+                "text_preview": text[: args.max_new_tokens_report],
+            }
+        )
+    except Exception as error:  # noqa: BLE001 - profiling should report failures instead of crashing
+        elapsed_s = time.perf_counter() - started_at
+        result.update(
+            {
+                "wall_time_s": round(elapsed_s, 2),
+                "peak_vram_gb": round(torch.cuda.max_memory_allocated() / (1024**3), 2) if device == "cuda" else None,
+                "error_type": type(error).__name__,
+                "error_message": str(error),
+                "traceback_tail": traceback.format_exc().strip().splitlines()[-1],
+            }
+        )
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

examples/profiling/profile_diffusion_gemma_quantized_colab_run.py

@@ -0,0 +1,9 @@
+import runpy
+import sys
+sys.argv = [
+    "profile_diffusion_gemma_quantized.py",
+    "--gen_length", "32",
+    "--num_inference_steps", "4",
+    "--cache_implementation", "dynamic",
+]
+runpy.run_path("/content/profile_diffusion_gemma_quantized.py", run_name="__main__")