Skip to content

feat: enable dynamic cpu-offloading for qwen image based models#791

Open
Cyclica-bin wants to merge 1 commit intonunchaku-ai:mainfrom
Cyclica-bin:feat/qwen-dynamic-offload
Open

feat: enable dynamic cpu-offloading for qwen image based models#791
Cyclica-bin wants to merge 1 commit intonunchaku-ai:mainfrom
Cyclica-bin:feat/qwen-dynamic-offload

Conversation

@Cyclica-bin
Copy link

Motivation

This PR automates CPU offloading in NunchakuQwenImageDiTLoader, replacing the manual block-count parameter with a dynamic offloading system.

Currently, users must manually specify the number of blocks to load into VRAM. This leads to:

  • Trial and error: Users have to guess how many blocks fit their specific GPU.
  • Crashes: If other processes (or browser tabs) consume VRAM, the static block count causes OOM (Out of Memory) errors.
  • Speeds up the workflow execution even without --lowvram flag if user has sufficient RAM.

Modifications

  1. Removal of num_blocks_on_gpu: Removed the manual input parameter from the ComfyUI node since users no longer need to guess how many transformer blocks their free VRAM can hold.
  2. API Update on load_model: Removed the same parameter from the argument list of load_model method.
  3. Uses Comfy UI's model management utilities to
    • Calculate the available free memory of GPU
    • The size of transformer per block calculation
    • Identifies the blocks that must be reserved in the memory that are variable size based on nature ( minor perf improvements )
  4. Added additional clamping logic just in case there are architectural changes to base model revision to avoid unexpected crashes.

Checklist

  • Code is formatted using Pre-Commit hooks (run pre-commit run --all-files).
  • Tested on a hardware with 8GB VRAM with 32GB RAM.
  • Tested on a hardware with more than 8GB VRAM.

Note: Other checklists are not applicable as this is an extension to existing node.

  • For reviewers: If you're only helping merge the main branch and haven't contributed code to this PR, please remove yourself as a co-author when merging.
  • Please feel free to join our Discord orWeChat to discuss your PR.

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