Skip to content

feat: add minimal memory profile - #22

Open
CloudRipple wants to merge 1 commit into
OpenMOSS:mainfrom
CloudRipple:main
Open

CloudRipple wants to merge 1 commit into
OpenMOSS:mainfrom
CloudRipple:main

Conversation

@CloudRipple

Copy link
Copy Markdown

Summary

Add a minimal profile for inference with limited GPU and system memory. DiT weights stay in CPU RAM across sampling steps, while the text encoder and VAE load parameters on demand from the existing safetensors files using Accelerate offload hooks.

  • Add Low memory to Gradio's memory modes and select it automatically below 12 GiB of GPU budget.
  • Disable DiT compilation for this mode and preserve the selection in generation metadata.
  • Include the new runtime module in the package and document usage and measured memory consumption in both READMEs.

Usage

python gradio_app.py --profile minimal --model OpenMOSS-Team/MOSS-Image-Base

Validation

Two consecutive 1024×1024, 50-step Base generations through ModelRunner, reusing one pipeline, matched the resident reference's float RGB outputs under matched attention settings.

Measured on an H200 with a 4 GiB PyTorch allocator budget and a 14 GiB task RAM limit:

Metric Result
Peak PyTorch allocated / reserved GPU memory 3.65 / 3.85 GiB
Peak task RAM, including file cache 14.00 GiB
Disk reads during denoising 0 for both requests
Generation time, excluding loading 177.44 s / 173.42 s

CUDA runtime memory is additional to PyTorch allocator usage. These measurements were collected on the H200, rather than a physical low-memory GPU.

Also checked configuration and metadata roundtrips, Gradio construction and mode switching, model unloading, wheel packaging, Ruff lint, and diff whitespace.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant