On 4 October a Pixal3D run on my Mac went into its texture stage, and within ten seconds the process grew from 27 to 44 GB. Available memory fell to 6 % and swap grew by 12 GB, so the guard I run heavy jobs under stopped it. The Mac has 48 GB. Before I had that guard, two Pixal3D runs, one of them next to other heavy work, ended with the Mac rebooting.
Pixal3D turns one image into a textured 3D mesh, and Pixal3D-mac is its port to Apple Silicon. The 1024_cascade pipeline runs seven models in turn: a flow model and a decoder for the coarse structure, two flow models for the shape at 512 and 1024, a shape decoder, and a flow model and a decoder for the texture. Around them work four image feature extractors, a background remover and a camera estimator. Each stage needs one or two of the seven models, and the port keeps all of them in memory from the first second to the last.
Why the memory does not come back
On Apple Silicon the CPU and the GPU share one pool of memory. Pixal3D's low-VRAM mode, which is on by default, keeps idle models on the CPU and copies the active one to the GPU. On a PC with a separate graphics card, that saves video memory. On a Mac the CPU copy and the GPU copy sit in the same RAM, so the mode keeps every model resident and adds a second copy of the active one. The four feature extractors also load the same DINOv3 backbone, four times.
Loading one stage at a time
stageload is a small library I wrote for this. It swaps the pipeline's dictionary of models for one that loads a model the first time a stage asks for it, and that releases every model the next stage does not list. Releasing moves each tensor of the module to PyTorch's meta device. That frees the storage even while other code still holds a reference to the module, and calling the module afterwards raises an error that names it and the stage that released it, instead of failing somewhere inside PyTorch.
The port's own run() executes unchanged. stageload wraps four of the pipeline's methods to learn where each stage begins: background removal, the first image conditioning, the conditioning of each shape and texture pass, and the final decoding. The four feature extractors are built so that they share one backbone.
One part of this raises no error when it goes wrong. Pixal3D seeds once at the start of a run and then draws the noise of every stage from the CPU random generator. Building a model runs its random initialisation, which draws from that same generator. A model built lazily in the middle of the run would therefore change the noise of every stage after it, and with it the mesh. stageload saves the generators' state before each load and restores it afterwards: torch's CPU, MPS and CUDA generators, Python's random and NumPy's global one. On a stand-in pipeline that seeds and draws noise the way the port does, a test checks that staged and eager runs produce the same latents bit for bit, and a control that turns the restore off sees them drift apart.
Not only Pixal3D
Nothing in the core of stageload knows about Pixal3D. A pipeline needs three things from it: a function that builds each model on its own, a list of the models each stage uses, and a call where each stage begins. In code you own, that call is one line, models.enter("decode"). In someone else's code, like this port, on_call attaches it to a method that already runs at that point, so the pipeline itself stays as it is. A text-to-image pipeline has the same shape as Pixal3D: a text encoder, a denoiser and an image decoder run one after another. The memory meter and the guard work for any program on a Mac. Pixal3D is the only pipeline I have measured so far; the README shows the core on a small two-model pipeline, and that example runs as a test.
What it changed
On this M5 Pro with 48 GB, I ran the port's example image through 1024_cascade with seed 7 and 2048-pixel textures, two runs per mode, each in its own process under the guard. "Eager" is the port's own loading and "staged" is stageload. Eager runs do not get through the texture stage on this machine within the guard's limits, so the two eager runs stop where it begins. The texture number for eager comes from the run on 4 October that went into it.
| stage | eager, peak footprint | staged, peak footprint |
|---|---|---|
| setup | 20.4 GB | 2.9 GB |
| camera | 23.2 GB | 6.7 GB |
| structure | 22.2 GB | 10.9 GB |
| shape_512 | 26.6 GB | 11.4 to 11.5 GB |
| shape_1024 | 30.5 GB | 13.1 GB |
| texture | stopped by the guard at 44.2 GB | 30.1 GB |
| decode | not reached | 17.9 to 18.4 GB |
| export | not reached | 69.6 to 70.2 GB |
In every stage both modes ran, staged held 11.3 to 17.5 GB less. The eight loads, the seven models plus the shape decoder a second time for decoding, read 13.2 GB in 24 s out of about ten minutes per run.
Most of the texture stage's 30 GB is working memory: the texture model itself is 2.6 GB. When decoding began, stageload released it together with the feature extractors and emptied the MPS allocator's cache, and within two seconds the footprint was down to 4.5 GB in one run and 8.0 GB in the other.
The highest staged number, about 70 GB, comes in the export, after every model weight has been released. It is the port's mesh processing (remeshing, UV unwrapping, baking), which stageload does not touch. The footprint is the number Activity Monitor shows, and it counts memory macOS has compressed, so it can be larger than the 48 GB of RAM.
What I could not show
The first is whether staged loading changes the result. Every run saved a hash and a copy of each sampler's output. The first latent, the coarse structure, is the same bit for bit in all four runs. From the 512 shape stage on, the port does not repeat itself on MPS even with the same seed: the two eager runs differ by up to 1.77 in the shape latent, and a staged run differs from an eager one by 1.24 to 1.86. So the comparison says any change from staged loading is no larger than the port's own variation, measured on that single pair of eager runs. The texture latent and the final mesh were not compared at all, because eager runs never reach them here.
Building on the CPU did not help. stageload builds each model directly on the GPU, and that changes one tensor the checkpoint does not cover: the rotary frequencies of the sparse attention, computed when a model is built, come out up to 6e-8 apart. To see whether that explained the differences, I ran a staged pass that built every model on the CPU and moved it to the GPU, the way the port does. Its latents differed from the eager runs by 1.84 and 1.12, the same order as every other pair. Loading took 82 s instead of 24 s, and three stages peaked 2.6 to 3.8 GB higher while two came out lower. I went back to building on the GPU.
My own summary had a bug that inflated one number. It took each stage's peak from the moment the stage was recorded as started, but releasing the previous stage's models takes up to 1.7 s. One memory sample in that gap put the decode peak of the CPU-built run at 29.5 GB; counted from the end of the releases, it is 16.7. A code review caught it. A stage's window now begins when the releases its start set off have finished.
The guard
The runs went through a second part of stageload, stageload guard, because other heavy jobs run on this Mac too. It starts one command only when enough memory is available and no process matching a list of heavy jobs has been seen for a while. While the command runs, it stops the command's whole process group if swap grows past a budget or available memory stays under 10 %. It needs no root, and the only processes it ever stops are the ones it started.
The wait for quiet came from a failure. A chain of video jobs leaves short gaps between them, a run started in one of those gaps, and the next video job took swap past the budget within 20 s. In the bench every run started only after three minutes with no heavy job in sight.
On another image
The same staged run on a character from a game I am making took 416 s and loaded the same 13.2 GB in 37 s. Its texture stage peaked at 30.1 GB, as on the example image. Here is the input and the mesh it produced, rendered in Blender:
The code, the traces of every run and the scripts that drew these charts are in the repository. The Pixal3D adapter is tested against one commit of Pixal3D-mac.
Top comments (0)