Support sleep, wake_up and load_weights for Omni Diffusion

[knlnguyen1802]

Purpose

Fix #316

This will support load and offload weight for Diffusion Model

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[SamitHuang]

test plan and results are missing. can you add unit tests for the new interfaces?

[knlnguyen1802]

test plan and results are missing. can you add unit tests for the new interfaces?

Add unit test

[ZJY0516]

Do we need to add a function in engine to call this and expose an interface to user?

[knlnguyen1802]

Yes for the new interface of OmniStage after #391. And it is also fix by #355

[ZJY0516]

I'd like to discuss one more point. Since vLLM-Omni uses libraries like Transformers and Diffusers to load components (e.g., text encoder and VAE), does the current sleep method also handle memory allocated by these external libraries? @knlnguyen1802

[knlnguyen1802]

I'd like to discuss one more point. Since vLLM-Omni uses libraries like Transformers and Diffusers to load components (e.g., text encoder and VAE), does the current sleep method also handle memory allocated by these external libraries? @knlnguyen1802

The answer is no since sleep only handle memory that wrap in the context of _maybe_get_memory_pool_context

[ZJY0516]

Could we just offload model to cpu like model.to(cpu) since diffusion models don't have cuda graph and kv cache

[knlnguyen1802]

Could we just offload model to cpu like model.to(cpu) since diffusion models don't have cuda graph and kv cache

For model, it already work as in this PR because I wrap the model loader in context of _maybe_get_memory_pool_context.
For other part like Cache it might need to do that way.
Noted: Sorry that I misunderstanding your question at the begining

[ZJY0516]

Could we just offload model to cpu like model.to(cpu) since diffusion models don't have cuda graph and kv cache

For model, it already work as in this PR because I wrap the model loader in context of _maybe_get_memory_pool_context. For other part like Cache it might need to do that way. Noted: Sorry that I misunderstanding your question at the begining

Being able to call model.to('cpu') directly would make the current PR, with its new _maybe_get_memory_pool_context, seem unnecessarily complex.

[knlnguyen1802]

Being able to call model.to('cpu') directly would make the current PR, with its new _maybe_get_memory_pool_context, seem unnecessarily complex.

Yes, but the "CuMemAllocator" is a well define class that help to easier track how many GB of memory if offload and move back to GPU when wakeup, also keep track of time to handle these operation. Also it can have some optimization.

[hsliuustc0106]

lgtm

[ZJY0516]

Could you please add some doc for sleep mode? @knlnguyen1802

[knlnguyen1802]

Could you please add some doc for sleep mode? @knlnguyen1802

Got it I'll add in a new PR