QwenImageTEModel missing memory_estimation_function causes OOM during inference
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Problem
When running workflows using `QwenImageTEModel` (Qwen2.5-7B based text encoder, e.g. WAN2.1 EditPlus) on 24GB GPUs, CUDA OOM can occur during the forward pass of `TextEncodeQwenImageEditPlus`, even though there appears to be enough GPU memory for model loading.
Typical scenario: a previous model (e.g. WanVAE) is cached on GPU. When the text encoder loads, `load_models_gpu` does not free enough cached memory, leaving insufficient room for inference activation memory.
Example log sequence:
```
[prev task] QwenImage cached on GPU (20.43GB total, ~12.82GB remains after partial unload)
[new task] QwenImageTEModel_ loaded fully: 7.75GB
[new task] OOM: tried to allocate 2.03GB, only 2.02GB free, 21.48GB in-use
```
The OOM happens during `encode()` (forward pass), not during model loading.
## Root Cause
`QwenImageTEModel` does not implement `memory_estimation_function`. In `sd.py:418-422`, `load_model()` checks for this method to estimate inference activation memory:
```python
def load_model(self, tokens={}):
memory_used = 0
if hasattr(self.cond_stage_model, "memory_estimation_function"):
memory_used = self.cond_stage_model.memory_estimation_function(tokens, device=...)
model_management.load_models_gpu([self.patcher], memory_required=memory_used)
```
Without it, `memory_required=0`, so `extra_mem` in `load_models_gpu` falls back to `max(minimum_inference_memory, 0 + extra_reserved_memory())` ≈ 1.2GB. For Qwen2.5-7B (hidden_size=3584, 28 layers), the actual activation memory can exceed 2GB, causing OOM.
`ace15.py` and `lt.py` already implement this method, but `qwen_image.py` does not.
## Suggested Fix
Add `memory_estimation_function` to `QwenImageTEModel` in `comfy/text_encoders/qwen_image.py`, following the same pattern as `LTXAVTEModel` in `lt.py`:
```python
def memory_estimation_function(self, token_weight_pairs, device=None):
constant = 6.0
if comfy.model_management.should_use_bf16(device):
constant /= 2.0
token_weight_pairs = token_weight_pairs.get("qwen25_7b", [])
num_tokens = sum(map(lambda a: len(a), token_weight_pairs))
num_tokens = max(num_tokens, 642)
return num_tokens * constant * 1024 * 1024
```
## Other TE models potentially affected
Other large-LLM-based text encoders also lack `memory_estimation_function` and may have similar issues on 24GB GPUs:
- `HunyuanImageTEModel` (`hunyuan_image.py`) — Qwen2.5-7B + BytT5
- `LongCatImageTEModel` (`longcat_image.py`) — Qwen2.5-7B
- `Kandinsky5TEModel` (`kandinsky5.py`) — Qwen2.5-7B + CLIP-L
- `Flux2TEModel` (`flux.py`) — Mistral3-24B
- `HiDreamTEModel` (`hidream.py`) — T5-XXL + Llama2
---
*Disclosure: The diagnosis and root cause analysis of this issue was assisted by [Claude Code](https://claude.com/claude-code) (AI coding assistant). If AI-assisted issues are not welcome in this project, please feel free to close this directly.*
Contributor guide
Assessment
This issue has not been assessed yet.