Check a Local Image Model Memory Budget in Node.js

When an image-model worker starts, its checkpoint is only one of the things that can occupy GPU memory. A scratch buffer or a second resident component can push an otherwise reasonable setup over its budget. The first useful check is whether the allocations you already know about fit together.

This tutorial builds a small Node.js preflight for a single device and a single phase of a job. It reads an explicit allocation plan, converts amounts to bytes, and returns signed headroom after a reserve.

1. Define which allocations coexist

Start with the free-memory snapshot from before these planned allocations are made. Give each allocation expected to coexist during the phase a unique ID. If a buffer is already accounted for in the free-memory snapshot, leave it out of the planned allocations.

Use separate plans for phases whose buffers are released in between. For example, you might keep a weights entry in both the encoding and denoising plans while changing the workspace entry. Review the smallest headroom across those plans.

Keep device, model variant, resolution, batch size, and measurement time alongside your estimates. Measurements from a different configuration belong in a separate plan.

2. Write a small fixture

Save the following example as fitting.json. Its numbers are fictional. The example uses an available snapshot of 10 GiB and budgets a further 1 GiB of reserve.

{
  "available": "10 GiB",
  "reserve": "1 GiB",
  "allocations": [
    {
      "id": "weights",
      "amount": "7 GiB"
    },
    {
      "id": "workspace",
      "amount": "1536 MiB"
    }
  ]
}

Amounts are strings with a whole number of up to 20 digits, one space, and B, MiB, or GiB. One MiB is 1,048,576 bytes and one GiB is 1,073,741,824 bytes. Convert decimal units using the data source’s convention before passing them in.

The reserve is a policy choice for headroom you want to keep free. The workspace is an expected allocation. Budget each quantity once.

3. Sum exact byte counts

Save the script as memory-budget.mjs and run it with Node.js 24.20.0. The check rejects ambiguous units, empty allocation lists, repeated IDs, and zero-sized allocation entries so those inputs cannot quietly produce a pass.

import { readFileSync } from "node:fs";

const unitBytes = { B: 1n, MiB: 1024n ** 2n, GiB: 1024n ** 3n };

function parseBytes(value, label) {
  const match = typeof value === "string"
    ? /^(0|[1-9][0-9]{0,19}) (B|MiB|GiB)$/u.exec(value)
    : null;
  if (!match || match[0] !== value) {
    throw new Error(label + " must use a whole amount and B, MiB, or GiB");
  }
  return BigInt(match[1]) * unitBytes[match[2]];
}

export function checkMemoryBudget(plan) {
  if (!plan || !Array.isArray(plan.allocations) || !plan.allocations.length) {
    throw new Error("allocations must be a nonempty array");
  }
  const available = parseBytes(plan.available, "available");
  const reserve = parseBytes(plan.reserve, "reserve");
  const seen = new Set();
  let allocated = 0n;

  for (const [index, item] of plan.allocations.entries()) {
    if (!item || typeof item.id !== "string" ||
        !item.id.trim() || item.id.trim() !== item.id) {
      throw new Error("allocation " + (index + 1) + " needs a trimmed ID");
    }
    if (seen.has(item.id)) throw new Error("duplicate allocation: " + item.id);
    seen.add(item.id);
    const bytes = parseBytes(item.amount, item.id);
    if (bytes === 0n) throw new Error(item.id + " must be positive");
    allocated += bytes;
  }

  const headroom = available - reserve - allocated;
  return {
    fits: headroom >= 0n,
    availableBytes: available.toString(),
    reserveBytes: reserve.toString(),
    allocatedBytes: allocated.toString(),
    headroomBytes: headroom.toString(),
  };
}

if (import.meta.main) {
  try {
    if (process.argv.length !== 3) {
      throw new Error("Usage: node memory-budget.mjs plan.json");
    }
    const plan = JSON.parse(readFileSync(process.argv[2], "utf8"));
    const result = checkMemoryBudget(plan);
    console.log(JSON.stringify(result, null, 2));
    process.exitCode = result.fits ? 0 : 1;
  } catch (error) {
    console.error(error.message);
    process.exitCode = 2;
  }
}

Byte counts stay as integers throughout. The result serializes them as decimal strings so consumers can retain precision even beyond JavaScript’s safe Number range. MDN BigInt reference.

4. Run a passing and a failing plan

Make a copy named over.json and change weights from 7 GiB to 8 GiB. Run both checks, then run without arguments to exercise invalid input.

node memory-budget.mjs fitting.json
node memory-budget.mjs over.json
node memory-budget.mjs

The first command prints this report.

{
  "fits": true,
  "availableBytes": "10737418240",
  "reserveBytes": "1073741824",
  "allocatedBytes": "9126805504",
  "headroomBytes": "536870912"
}

The initial plan uses 8,704 MiB. With 1,024 MiB reserved, that leaves 512 MiB. Increasing weights by 1 GiB changes headroom to minus 512 MiB.

5. Handle the three exit codes

Exit 0 means the listed allocations and reserve fit the snapshot. Exactly zero headroom passes because the reserve is already included in the calculation.

Exit 1 means a valid plan exceeds the budget. Use the signed headroom to see the shortfall in bytes. Exit 2 means the script could not evaluate the plan, so use the message on standard error to fix the input.

If you run several phases in a deployment check, keep the input and result together for each phase. A missing plan or stale snapshot needs its own handling in the calling workflow.

6. Compare against a measured peak

Use the first real run to measure the chosen workload under representative conditions. If allocation lifetimes overlap differently from what the plan assumed, revise the plan and repeat the check with a fresh snapshot.

A passing result describes the supplied allocations. The runtime can still need additional memory for allocator behavior, temporary work, or another process that starts after the snapshot.

AI-assisted writing and code drafting. The examples were checked with synthetic local JSON on Node.js 24.20.0. No GPU workload or model run was measured.

This tutorial develops the local hardware planning topic in the PixMind source article. Flux AI Image Generator Review and Complete Guide.