chore: import upstream snapshot with attribution

scripts/calibrate_qwen_vae_working_memory.py

@@ -0,0 +1,305 @@
+"""Calibrate the Qwen Image VAE working-memory estimate against measured peak CUDA/HIP memory.
+
+Background
+----------
+``estimate_vae_working_memory_qwen_image`` models peak working memory as a linear function of
+spatial area::
+
+    working_memory = h * w * element_size * scaling_constant
+
+This script measures the *actual* peak reserved memory the VAE consumes during decode/encode across
+a grid of resolutions so the ``scaling_constant`` can be fit from several points instead of one, and
+so we can check whether the pure-linear model holds or whether a super-linear (attention) term
+appears at high resolution.
+
+The estimate is consumed by the model cache via ``free >= estimate`` to decide whether to evict, so
+it MUST be an upper bound: we measure peak *reserved* (not just allocated) memory, the conservative
+quantity that includes caching-allocator overhead and kernel scratch/workspace.
+
+Portability
+-----------
+Backend-agnostic: uses only ``torch.cuda.*``, which works on both NVIDIA/CUDA and AMD/ROCm (HIP)
+builds of PyTorch. Run the SAME script on each backend and compare the ``implied_constant`` columns
+-- the curve *shape* (linear vs. super-linear) is architectural and should match, but the absolute
+constant can differ (cuDNN vs. MIOpen conv workspaces, flash-attention availability, allocator
+rounding). Ship ``max`` across backends plus headroom.
+
+Each (operation, resolution) point is measured in a FRESH SUBPROCESS so the caching allocator's
+fragmentation history from earlier points cannot contaminate the reserved-delta reading. A point
+that OOMs is recorded as ``oom`` rather than aborting the run, so the grid can probe up to the
+card's ceiling safely.
+
+Usage
+-----
+    python scripts/calibrate_qwen_vae_working_memory.py [--vae /path/to/vae_dir] [--csv out.csv]
+
+If ``--vae`` is omitted, the script auto-discovers an ``AutoencoderKLQwenImage`` under
+``$INVOKEAI_ROOT/models``.
+"""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage
+
+LATENT_SCALE_FACTOR = 8
+
+# (height, width) pixel-space resolutions. Squares to test linearity in area, plus non-square
+# points (incl. the original 1248x832 calibration point) to confirm area = h*w is the right
+# predictor rather than max(h, w) or perimeter. Subprocess isolation + OOM capture means we can
+# list aggressive resolutions; ones that don't fit are simply recorded as oom.
+DEFAULT_RESOLUTIONS = [
+    (512, 512),
+    (768, 768),
+    (832, 1248),  # original single calibration point (as HxW)
+    (1024, 1024),
+    (1088, 1920),
+    (1280, 1280),
+    (1536, 1024),
+    (1536, 1536),
+    (1792, 1792),
+    (2048, 2048),
+]
+
+
+def discover_vae() -> Path:
+    """Find an AutoencoderKLQwenImage VAE directory under $INVOKEAI_ROOT/models."""
+    root = os.environ.get("INVOKEAI_ROOT")
+    if not root:
+        raise SystemExit("INVOKEAI_ROOT not set; pass --vae explicitly.")
+    models = Path(root) / "models"
+    for config_path in models.glob("*/vae/config.json"):
+        try:
+            cfg = json.loads(config_path.read_text())
+        except Exception:
+            continue
+        if cfg.get("_class_name") == "AutoencoderKLQwenImage":
+            return config_path.parent
+    raise SystemExit(f"No AutoencoderKLQwenImage VAE found under {models}; pass --vae explicitly.")
+
+
+DTYPES = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}
+
+
+def _load_vae(vae_path: str, dtype: torch.dtype) -> AutoencoderKLQwenImage:
+    """Load an AutoencoderKLQwenImage from either a diffusers directory or a single .safetensors file.
+
+    Directory: standard ``from_pretrained``.
+    Single file: ``AutoencoderKLQwenImage`` has no single-file converter registered in diffusers,
+    so we instantiate with the default config and load the state dict directly. Two on-disk layouts
+    exist: the diffusers layout (``encoder.conv_in`` / ``down_blocks`` / ``mid_block`` keys, e.g. the
+    weights InvokeAI's VAELoader consumes) and the original Qwen-Image/Wan release layout
+    (``encoder.conv1`` / ``downsamples`` / ``middle`` / ``time_conv`` keys). We try a direct strict
+    load first, and on a key mismatch fall back to diffusers' Wan VAE converter -- the Qwen-Image VAE
+    shares the Wan VAE key structure -- before retrying.
+    """
+    path = Path(vae_path)
+    if not path.is_file():
+        return AutoencoderKLQwenImage.from_pretrained(vae_path, local_files_only=True, torch_dtype=dtype)
+
+    from safetensors.torch import load_file
+
+    sd = load_file(str(path))
+    for k in list(sd.keys()):
+        if sd[k].is_floating_point():
+            sd[k] = sd[k].to(dtype)
+
+    vae = AutoencoderKLQwenImage()
+    try:
+        # diffusers-layout checkpoint: keys already match the model. State dict was converted to
+        # `dtype` above and is assigned in place, so params carry the correct dtype.
+        vae.load_state_dict(sd, strict=True, assign=True)
+    except RuntimeError:
+        # Original Qwen-Image/Wan release layout: convert keys to the diffusers layout, then retry.
+        from diffusers.loaders.single_file_utils import convert_wan_vae_to_diffusers
+
+        converted = convert_wan_vae_to_diffusers(sd)
+        for k in list(converted.keys()):
+            if converted[k].is_floating_point():
+                converted[k] = converted[k].to(dtype)
+        vae.load_state_dict(converted, strict=True, assign=True)
+    return vae
+
+
+def _build_input(operation: str, h: int, w: int, z_dim: int, dtype: torch.dtype) -> torch.Tensor:
+    """Construct the 5D (B, C, num_frames, H, W) input the invocation feeds the VAE.
+
+    decode: latents at latent resolution (H/8, W/8) with z_dim channels.
+    encode: image at pixel resolution (H, W) with 3 channels.
+    These mirror QwenImageLatentsToImageInvocation / QwenImageImageToLatentsInvocation exactly.
+    """
+    device = torch.device("cuda")
+    if operation == "decode":
+        return torch.randn(1, z_dim, 1, h // LATENT_SCALE_FACTOR, w // LATENT_SCALE_FACTOR, device=device, dtype=dtype)
+    return torch.randn(1, 3, 1, h, w, device=device, dtype=dtype)
+
+
+@torch.inference_mode()
+def measure_one(vae_path: str, operation: str, h: int, w: int, dtype: torch.dtype) -> dict:
+    """Measure peak reserved-memory growth for a single decode/encode. Runs in a child process."""
+    vae = _load_vae(vae_path, dtype)
+    vae.to("cuda")
+    vae.disable_tiling()  # Qwen invocations never tile; match that.
+
+    param = next(vae.parameters())
+    dtype = param.dtype
+    element_size = param.element_size()
+    z_dim = int(vae.config.z_dim)
+
+    x = _build_input(operation, h, w, z_dim, dtype)
+
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+    baseline_reserved = torch.cuda.memory_reserved()
+
+    # Measure the COLD first call -- it includes conv-algorithm-search / attention workspace
+    # allocation, which is exactly what the real (single-shot) invocation pays.
+    try:
+        if operation == "decode":
+            vae.decode(x, return_dict=False)
+        else:
+            vae.encode(x).latent_dist.mode()
+        torch.cuda.synchronize()
+    except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+        if "out of memory" not in str(e).lower():
+            raise
+        return {"operation": operation, "h": h, "w": w, "oom": True}
+
+    peak_reserved = torch.cuda.max_memory_reserved()
+    peak_allocated = torch.cuda.max_memory_allocated()
+    reserved_delta = peak_reserved - baseline_reserved
+
+    area = h * w
+    return {
+        "operation": operation,
+        "h": h,
+        "w": w,
+        "area": area,
+        "element_size": element_size,
+        "dtype": str(dtype),
+        "reserved_delta": reserved_delta,
+        "allocated_peak": peak_allocated,
+        "reserved_baseline": baseline_reserved,
+        # The constant as the estimator parameterizes it: mem = area * element_size * k
+        "implied_constant": reserved_delta / (area * element_size),
+        "oom": False,
+    }
+
+
+def run_grid(vae_path: str, resolutions: list[tuple[int, int]], dtype_name: str, csv_path: Path | None) -> None:
+    rows: list[dict] = []
+    print(f"VAE: {vae_path}")
+    print(
+        f"torch {torch.__version__} | device {torch.cuda.get_device_name(0)} | hip={torch.version.hip} | dtype={dtype_name}\n"
+    )
+    print(f"{'op':6} {'HxW':>11} {'area':>10} {'reserved(GiB)':>14} {'alloc(GiB)':>11} {'implied_k':>10}")
+    print("-" * 70)
+
+    for operation in ("decode", "encode"):
+        for h, w in resolutions:
+            # Fresh subprocess per point for an uncontaminated reserved-memory reading.
+            proc = subprocess.run(
+                [
+                    sys.executable,
+                    __file__,
+                    "--single",
+                    operation,
+                    str(h),
+                    str(w),
+                    "--vae",
+                    vae_path,
+                    "--dtype",
+                    dtype_name,
+                ],
+                capture_output=True,
+                text=True,
+            )
+            line = proc.stdout.strip().splitlines()[-1] if proc.stdout.strip() else ""
+            try:
+                row = json.loads(line)
+            except Exception:
+                print(f"{operation:6} {f'{h}x{w}':>11}  FAILED: {proc.stderr.strip().splitlines()[-1:]}")
+                continue
+            rows.append(row)
+            if row.get("oom"):
+                print(f"{operation:6} {f'{h}x{w}':>11} {h * w:>10}  {'OOM':>14}")
+                continue
+            gib = 1024**3
+            print(
+                f"{operation:6} {f'{h}x{w}':>11} {row['area']:>10} "
+                f"{row['reserved_delta'] / gib:>14.3f} {row['allocated_peak'] / gib:>11.3f} "
+                f"{row['implied_constant']:>10.1f}"
+            )
+
+    # Summary: the shippable constant is the MAX implied constant over fitting points (upper bound).
+    print("\n=== summary (max implied constant = candidate scaling_constant, before headroom) ===")
+    for operation in ("decode", "encode"):
+        ks = [r["implied_constant"] for r in rows if r["operation"] == operation and not r.get("oom")]
+        if ks:
+            print(
+                f"{operation:6}: n={len(ks)}  min_k={min(ks):.1f}  max_k={max(ks):.1f}  "
+                f"-> use >= {max(ks):.0f} (+headroom)"
+            )
+
+    if csv_path:
+        import csv
+
+        fieldnames = [
+            "operation",
+            "h",
+            "w",
+            "area",
+            "element_size",
+            "dtype",
+            "reserved_delta",
+            "allocated_peak",
+            "reserved_baseline",
+            "implied_constant",
+            "oom",
+        ]
+        with csv_path.open("w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
+            writer.writeheader()
+            for r in rows:
+                writer.writerow(r)
+        print(f"\nWrote {csv_path}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "--vae",
+        type=str,
+        default=None,
+        help="Path to an AutoencoderKLQwenImage diffusers dir OR a single .safetensors checkpoint.",
+    )
+    parser.add_argument("--csv", type=str, default=None, help="Optional path to write the raw results as CSV.")
+    parser.add_argument(
+        "--dtype",
+        choices=list(DTYPES),
+        default="float16",
+        help="Compute dtype. Default float16 to match InvokeAI's default precision on CUDA/ROCm.",
+    )
+    # Internal: measure a single point in this process and print one JSON line.
+    parser.add_argument("--single", nargs=3, metavar=("OP", "H", "W"), default=None, help=argparse.SUPPRESS)
+    args = parser.parse_args()
+
+    vae_path = args.vae or str(discover_vae())
+    dtype = DTYPES[args.dtype]
+
+    if args.single:
+        op, h, w = args.single[0], int(args.single[1]), int(args.single[2])
+        print(json.dumps(measure_one(vae_path, op, h, w, dtype)))
+        return
+
+    run_grid(vae_path, DEFAULT_RESOLUTIONS, args.dtype, Path(args.csv) if args.csv else None)
+
+
+if __name__ == "__main__":
+    main()