From 6975e6c5332be4921de0e720ff35fc4fa843532c Mon Sep 17 00:00:00 2001 From: Pawel Date: Wed, 8 Jul 2026 11:40:37 -0700 Subject: [PATCH] Add profiling recipe benchmarks --- examples/profiling/README.md | 42 +- examples/profiling/pipeline_registry.py | 101 +++++ .../profile_diffusion_gemma_quantized.py | 146 +++++++ ...ile_diffusion_gemma_quantized_colab_run.py | 9 + .../profiling/profile_pipeline_recipes.py | 395 ++++++++++++++++++ examples/profiling/profiling_pipelines.py | 102 +---- 6 files changed, 698 insertions(+), 97 deletions(-) create mode 100644 examples/profiling/pipeline_registry.py create mode 100644 examples/profiling/profile_diffusion_gemma_quantized.py create mode 100644 examples/profiling/profile_diffusion_gemma_quantized_colab_run.py create mode 100644 examples/profiling/profile_pipeline_recipes.py diff --git a/examples/profiling/README.md b/examples/profiling/README.md index 38b35772d03d..9f06b6926d50 100644 --- a/examples/profiling/README.md +++ b/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:` +- `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: diff --git a/examples/profiling/pipeline_registry.py b/examples/profiling/pipeline_registry.py new file mode 100644 index 000000000000..af40e6d5a65e --- /dev/null +++ b/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", + }, + ), + } diff --git a/examples/profiling/profile_diffusion_gemma_quantized.py b/examples/profiling/profile_diffusion_gemma_quantized.py new file mode 100644 index 000000000000..ebebbfa592a8 --- /dev/null +++ b/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() diff --git a/examples/profiling/profile_diffusion_gemma_quantized_colab_run.py b/examples/profiling/profile_diffusion_gemma_quantized_colab_run.py new file mode 100644 index 000000000000..0102d05a0d45 --- /dev/null +++ b/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__") diff --git a/examples/profiling/profile_pipeline_recipes.py b/examples/profiling/profile_pipeline_recipes.py new file mode 100644 index 000000000000..2b1bfc105862 --- /dev/null +++ b/examples/profiling/profile_pipeline_recipes.py @@ -0,0 +1,395 @@ +""" +Benchmark end-to-end pipeline optimization recipes and emit JSON + Markdown reports. + +Example usage: + 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 +""" + +import argparse +import copy +import gc +import json +import logging +import os +from dataclasses import asdict, dataclass +from datetime import datetime, timezone + +import torch + +try: + from .pipeline_registry import build_registry +except ImportError: + from pipeline_registry import build_registry + + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") +logger = logging.getLogger(__name__) + +ACTION_ORDER = { + "baseline": 0, + "layerwise_casting": 1, + "attention": 2, + "vae_tiling": 3, + "channels_last": 4, + "model_cpu_offload": 5, + "group_offload_leaf": 6, + "compile": 7, +} +DEFAULT_RECIPES = [ + "baseline", + "model_cpu_offload", + "group_offload_leaf", + "layerwise_casting", + "compile", +] + + +def flush(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.reset_max_memory_allocated() + torch.cuda.reset_peak_memory_stats() + + +@dataclass +class RecipeBenchmarkResult: + pipeline: str + recipe: str + status: str + mean_ms: float | None + std_ms: float | None + peak_vram_gb: float | None + num_runs: int + num_warmups: int + notes: str + + +def parse_recipe(recipe: str) -> list[tuple[str, str | None]]: + actions = [] + for raw_token in recipe.split("+"): + token = raw_token.strip() + if not token: + continue + if ":" in token: + name, value = token.split(":", 1) + else: + name, value = token, None + if name not in ACTION_ORDER: + raise ValueError(f"Unknown recipe action '{name}' in recipe '{recipe}'") + actions.append((name, value)) + + if not actions: + raise ValueError("Recipe cannot be empty") + + deduped = [] + seen = set() + for name, value in sorted(actions, key=lambda item: ACTION_ORDER[item[0]]): + key = (name, value) + if key not in seen and name != "baseline": + deduped.append((name, value)) + seen.add(key) + + return deduped or [("baseline", None)] + + +def benchmark_pipeline_call(pipe, call_kwargs, num_runs: int, num_warmups: int) -> tuple[float, float, float]: + for _ in range(num_warmups): + pipe(**call_kwargs) + torch.cuda.synchronize() + + times = [] + peak_memories = [] + + for _ in range(num_runs): + torch.cuda.reset_peak_memory_stats() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + pipe(**call_kwargs) + end.record() + torch.cuda.synchronize() + + times.append(start.elapsed_time(end)) + peak_memories.append(torch.cuda.max_memory_allocated() / (1024**3)) + + mean_ms = sum(times) / len(times) + variance = sum((t - mean_ms) ** 2 for t in times) / len(times) + return mean_ms, variance**0.5, max(peak_memories) + + +def resolve_vae_tiling(pipe): + if hasattr(pipe, "enable_vae_tiling"): + return pipe.enable_vae_tiling + vae = getattr(pipe, "vae", None) + if vae is not None and hasattr(vae, "enable_tiling"): + return vae.enable_tiling + return None + + +def apply_recipe(pipe, actions, *, call_kwargs, compile_kwargs, compile_regional): + notes = [] + uses_offloading = False + + for action_name, action_value in actions: + if action_name == "baseline": + continue + + if action_name == "layerwise_casting": + transformer = getattr(pipe, "transformer", None) + if transformer is None or not hasattr(transformer, "enable_layerwise_casting"): + raise RuntimeError("layerwise_casting requires pipe.transformer.enable_layerwise_casting()") + compute_dtype = getattr(transformer, "dtype", None) or getattr(pipe, "dtype", None) + if compute_dtype is None: + raise RuntimeError("layerwise_casting could not infer a compute dtype for this transformer") + transformer.enable_layerwise_casting( + storage_dtype=torch.float8_e4m3fn, + compute_dtype=compute_dtype, + ) + continue + + if action_name == "attention": + backend = action_value or "native" + transformer = getattr(pipe, "transformer", None) + if transformer is None or not hasattr(transformer, "set_attention_backend"): + raise RuntimeError("attention backend switching requires pipe.transformer.set_attention_backend()") + transformer.set_attention_backend(backend) + notes.append(f"attention_backend={backend}") + continue + + if action_name == "vae_tiling": + if call_kwargs.get("output_type") == "latent": + notes.append("vae_tiling requested with latent output; VAE decode is skipped") + continue + enable_vae_tiling = resolve_vae_tiling(pipe) + if enable_vae_tiling is None: + raise RuntimeError("vae_tiling is not supported by this pipeline") + enable_vae_tiling() + continue + + if action_name == "channels_last": + if hasattr(pipe, "unet") and pipe.unet is not None: + pipe.unet.to(memory_format=torch.channels_last) + continue + raise RuntimeError("channels_last is only supported for pipelines exposing pipe.unet") + + if action_name == "model_cpu_offload": + pipe.enable_model_cpu_offload() + uses_offloading = True + continue + + if action_name == "group_offload_leaf": + pipe.enable_group_offload( + onload_device=torch.device("cuda"), + offload_device=torch.device("cpu"), + offload_type="leaf_level", + ) + uses_offloading = True + continue + + if not uses_offloading: + pipe.to("cuda") + + for action_name, action_value in actions: + if action_name != "compile": + continue + transformer = getattr(pipe, "transformer", None) + if transformer is None: + raise RuntimeError("compile requires pipe.transformer") + if compile_regional and hasattr(transformer, "compile_repeated_blocks"): + transformer.compile_repeated_blocks(**compile_kwargs) + notes.append("compile=regional") + elif hasattr(transformer, "compile"): + transformer.compile(**compile_kwargs) + notes.append("compile=full") + else: + raise RuntimeError("compile is not supported by this transformer") + + return notes + + +def load_pipeline(config): + pipe = config.pipeline_cls.from_pretrained(**config.pipeline_init_kwargs) + pipe.set_progress_bar_config(disable=True) + return pipe + + +def run_recipe(config, recipe: str, *, num_runs: int, num_warmups: int, compile_kwargs, compile_regional): + actions = parse_recipe(recipe) + pipe = None + + try: + flush() + pipe = load_pipeline(config) + notes = apply_recipe( + pipe, + actions, + call_kwargs=config.pipeline_call_kwargs, + compile_kwargs=compile_kwargs, + compile_regional=compile_regional, + ) + mean_ms, std_ms, peak_vram_gb = benchmark_pipeline_call( + pipe, + config.pipeline_call_kwargs, + num_runs=num_runs, + num_warmups=num_warmups, + ) + return RecipeBenchmarkResult( + pipeline=config.name, + recipe=recipe, + status="pass", + mean_ms=round(mean_ms, 1), + std_ms=round(std_ms, 1), + peak_vram_gb=round(peak_vram_gb, 2), + num_runs=num_runs, + num_warmups=num_warmups, + notes="; ".join(notes), + ) + except Exception as error: + return RecipeBenchmarkResult( + pipeline=config.name, + recipe=recipe, + status="fail", + mean_ms=None, + std_ms=None, + peak_vram_gb=None, + num_runs=num_runs, + num_warmups=num_warmups, + notes=str(error), + ) + finally: + if pipe is not None: + try: + pipe.to("cpu") + except Exception: + pass + del pipe + flush() + + +def write_json_report(path: str, payload: dict): + with open(path, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + + +def write_markdown_report(path: str, payload: dict): + lines = [ + f"# Pipeline recipe report: {payload['pipeline']}", + "", + f"- Generated: {payload['generated_at_utc']}", + f"- Device: {payload['device_name']}", + f"- Torch: {payload['torch_version']}", + f"- Call kwargs: `{json.dumps(payload['call_kwargs'], sort_keys=True)}`", + "", + "| Recipe | Status | Mean (ms) | Std (ms) | Peak VRAM (GB) | Notes |", + "|---|---|---:|---:|---:|---|", + ] + + for result in payload["results"]: + lines.append( + "| {recipe} | {status} | {mean} | {std} | {peak} | {notes} |".format( + recipe=result["recipe"], + status=result["status"], + mean=result["mean_ms"] if result["mean_ms"] is not None else "—", + std=result["std_ms"] if result["std_ms"] is not None else "—", + peak=result["peak_vram_gb"] if result["peak_vram_gb"] is not None else "—", + notes=result["notes"] or "—", + ) + ) + + with open(path, "w", encoding="utf-8") as handle: + handle.write("\n".join(lines) + "\n") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark pipeline optimization recipes") + parser.add_argument( + "--pipeline", + choices=["flux", "flux2", "wan", "ltx2", "qwenimage"], + required=True, + help="Which pipeline to benchmark", + ) + parser.add_argument( + "--recipe", + action="append", + default=None, + help=( + "Recipe string to run. Combine actions with '+', for example " + "'group_offload_leaf+layerwise_casting' or 'compile+attention:_native_flash'. " + "Defaults to a conservative built-in recipe set." + ), + ) + parser.add_argument("--output_dir", default="profiling_results", help="Directory for report output") + parser.add_argument("--num_steps", type=int, default=None, help="Override num_inference_steps") + parser.add_argument("--full_decode", action="store_true", help="Set output_type='pil' so VAE recipes are meaningful") + parser.add_argument("--num_runs", type=int, default=5, help="Number of timed runs") + parser.add_argument("--num_warmups", type=int, default=2, help="Number of warmup runs") + parser.add_argument( + "--compile_mode", + default="default", + choices=["default", "reduce-overhead", "max-autotune"], + help="torch.compile mode for recipes that include compile", + ) + parser.add_argument("--compile_fullgraph", action="store_true", help="Use fullgraph=True for torch.compile") + parser.add_argument( + "--compile_regional", + action="store_true", + help="Use compile_repeated_blocks() when available instead of full transformer compilation", + ) + args = parser.parse_args() + + registry = build_registry() + config = copy.deepcopy(registry[args.pipeline]) + + if args.num_steps is not None: + config.pipeline_call_kwargs["num_inference_steps"] = args.num_steps + if args.full_decode: + config.pipeline_call_kwargs["output_type"] = "pil" + + recipes = args.recipe or list(DEFAULT_RECIPES) + compile_kwargs = {"fullgraph": args.compile_fullgraph, "mode": args.compile_mode} + + os.makedirs(args.output_dir, exist_ok=True) + results = [] + for recipe in recipes: + logger.info("Running recipe %s for pipeline %s", recipe, config.name) + results.append( + run_recipe( + config, + recipe, + num_runs=args.num_runs, + num_warmups=args.num_warmups, + compile_kwargs=compile_kwargs, + compile_regional=args.compile_regional, + ) + ) + + device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu" + payload = { + "pipeline": config.name, + "generated_at_utc": datetime.now(timezone.utc).isoformat(), + "device_name": device_name, + "torch_version": torch.__version__, + "call_kwargs": config.pipeline_call_kwargs, + "results": [asdict(result) for result in results], + } + + json_path = os.path.join(args.output_dir, f"{config.name}_recipes.json") + md_path = os.path.join(args.output_dir, f"{config.name}_recipes.md") + write_json_report(json_path, payload) + write_markdown_report(md_path, payload) + + logger.info("Wrote JSON report to %s", json_path) + logger.info("Wrote Markdown report to %s", md_path) + + +if __name__ == "__main__": + main() diff --git a/examples/profiling/profiling_pipelines.py b/examples/profiling/profiling_pipelines.py index 5a0b4bfe938b..b37e64881e26 100644 --- a/examples/profiling/profiling_pipelines.py +++ b/examples/profiling/profiling_pipelines.py @@ -18,107 +18,17 @@ import copy import logging -import torch -from profiling_utils import PipelineProfiler, PipelineProfilingConfig +try: + from .pipeline_registry import build_registry + from .profiling_utils import PipelineProfiler +except ImportError: + from pipeline_registry import build_registry + from profiling_utils import PipelineProfiler logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") logger = logging.getLogger(__name__) -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", - }, - ), - } - def main(): parser = argparse.ArgumentParser(description="Profile diffusers pipelines with torch.profiler")