Add Qwen2VL Support
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Feature Idea
Link to ChatGPT on steps to develop a solution:
https://chatgpt.com/share/674537c1-b728-800f-9381-ce85f916b3ef
Main script to analyze:
https://raw.githubusercontent.com/erwold/qwen2vl-flux/refs/heads/main/model.py
**Quick Summary of Changes Needed**
Let's review the integration plan with a focus on utilizing the existing ComfyUI infrastructure, specifically the `CLIP` class, `DualClipLoader`, and `CLIPTextEncode` nodes. Our goal is to integrate the Qwen2-VL model into ComfyUI's pipeline without unnecessary reloading or reinitialization.
---
### **Understanding the Existing CLIP Infrastructure**
**1. `DualClipLoader` Node:**
- **Function**: Loads two CLIP models (e.g., `clip_l` and `t5xxl`) and returns a `CLIP` object.
- **Process**:
- Calls `comfy.sd.load_clip`, which in turn calls `load_text_encoder_state_dicts`.
- Constructs a `CLIP` object that contains the text encoder model (`cond_stage_model`) and tokenizer.
**2. `CLIP` Class:**
- **Components**:
- **`cond_stage_model`**: The text encoder model used for encoding text.
- **`tokenizer`**: Tokenizer used for converting text into tokens.
- **Methods**:
- **`tokenize`**: Tokenizes input text.
- **`encode_from_tokens`**: Encodes tokens into embeddings.
**3. `CLIPTextEncode` Node:**
- **Function**: Uses the `CLIP` object to encode text into embeddings.
- **Process**:
- Calls `clip.tokenize` to tokenize the text.
- Calls `clip.encode_from_tokens` to obtain embeddings.
---
### **Integrating Qwen2-VL into the CLIP Pipeline**
**Objective**: Seamlessly integrate the Qwen2-VL model into the existing `CLIP` infrastructure so that it can be used with the `CLIPTextEncode` node without reloading or re-initializing models.
---
### **Step-by-Step Integration Plan**
#### **1. Extend the `TEModel` Enumeration**
First, we need to allow the system to recognize the Qwen2-VL model type.
```python
class TEModel(Enum):
CLIP_L = 1
CLIP_G = 2
T5_XXL = 3
# ... other models ...
QWEN2_VL = 10 # Assign an appropriate unique value
```
#### **2. Update the `detect_te_model` Function**
Modify `detect_te_model` to detect the Qwen2-VL model based on specific keys in its state dictionary.
```python
def detect_te_model(sd):
if "qwen2vl_specific_key" in sd:
return TEModel.QWEN2_VL
# ... existing detection logic ...
```
*Note*: Replace `"qwen2vl_specific_key"` with an actual unique key from the Qwen2-VL model's state dict.
#### **3. Modify `load_text_encoder_state_dicts`**
Update `load_text_encoder_state_dicts` to handle the Qwen2-VL model:
```python
def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
# ... existing code ...
if len(clip_data) == 1:
te_model = detect_te_model(clip_data[0])
if te_model == TEModel.QWEN2_VL:
clip_target.clip = Qwen2VLClipModel
clip_target.tokenizer = Qwen2VLTokenizer
# ... existing code ...
```
#### **4. Implement `Qwen2VLClipModel`**
Create a class `Qwen2VLClipModel` that conforms to the expected interface of `cond_stage_model`.
```python
class Qwen2VLClipModel:
def __init__(self, **params):
self.qwen2vl_model = params.get('qwen2vl_model')
self.t5_context_embedder = params.get('t5_context_embedder')
self.device = params['device']
self.dtype = params['dtype']
# Ensure models are on the correct device and dtype
self.qwen2vl_model.to(self.dtype).to(self.device)
self.t5_context_embedder.to(self.dtype).to(self.device)
def encode_token_weights(self, tokens):
text_input_ids = tokens.input_ids.to(self.device)
# Obtain embeddings from the Qwen2-VL model
with torch.no_grad():
prompt_embeds = self.qwen2vl_model.text_encoder(text_input_ids)[0]
# Pass embeddings through the T5 context embedder
prompt_embeds = self.t5_context_embedder(prompt_embeds)
# Return embeddings in the expected format
cond = prompt_embeds
pooled = None # Modify if pooled output is required
return cond, pooled
```
#### **5. Implement `Qwen2VLTokenizer`**
Create a `Qwen2VLTokenizer` class that conforms to the tokenizer interface expected by the `CLIP` class.
```python
class Qwen2VLTokenizer:
def __init__(self, embedding_directory=None, tokenizer_data={}):
self.tokenizer = T5TokenizerFast.from_pretrained('t5-xxl')
def tokenize_with_weights(self, text, return_word_ids=False):
tokens = self.tokenizer(
text,
padding="max_length",
max_length=256,
truncation=True,
return_tensors="pt",
)
return tokens
```
#### **6. Update the `CLIP` Class Initialization**
Ensure that the `CLIP` class can accept and initialize with the Qwen2-VL model and tokenizer.
In `load_text_encoder_state_dicts`, pass the pre-loaded models to the `CLIP` class:
```python
clip_target = EmptyClass()
clip_target.params = {
'qwen2vl_model': qwen2vl_model, # Pass the pre-loaded Qwen2-VL model
't5_context_embedder': t5_context_embedder, # Pass the T5 context embedder
# ... other parameters ...
}
clip_target.clip = Qwen2VLClipModel
clip_target.tokenizer = Qwen2VLTokenizer
```
#### **7. Adjust the `DualClipLoader` Node**
Modify the `DualClipLoader` to recognize when the Qwen2-VL model is being loaded:
```python
class DualCLIPLoader:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"clip_name1": (folder_paths.get_filename_list("text_encoders"), ),
"clip_name2": (folder_paths.get_filename_list("text_encoders"), ),
"type": (["sdxl", "sd3", "flux", "qwen2vl"], ), # Added "qwen2vl"
}
}
# ... existing code ...
def load_clip(self, clip_name1, clip_name2, type):
# ... existing code ...
if type == "qwen2vl":
clip_type = CLIPType.QWEN2_VL
# ... existing code ...
```
*Note*: You'll need to define `CLIPType.QWEN2_VL` similarly to other types.
#### **8. Ensure Models are Pre-Loaded and Passed Correctly**
Load the Qwen2-VL model and `t5_context_embedder` once, possibly in the `DualClipLoader` node or another appropriate loader node, and pass them into the `CLIP` class via `params`.
```python
def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
clip_data = []
qwen2vl_model = None
t5_context_embedder = None
for p in ckpt_paths:
if clip_type == CLIPType.QWEN2_VL:
# Load the Qwen2-VL model and t5_context_embedder here
qwen2vl_model = Qwen2VLSimplifiedModel.from_pretrained(p)
t5_context_embedder = load_t5_context_embedder()
else:
clip_data.append(comfy.utils.load_torch_file(p, safe_load=True))
# ... existing code ...
```
#### **9. Modify `CLIPTextEncode` Node if Necessary**
Since the `CLIP` object now contains the Qwen2-VL model and tokenizer, the `CLIPTextEncode` node should work as is. It calls `clip.tokenize` and `clip.encode_from_tokens`, which are implemented by `Qwen2VLTokenizer` and `Qwen2VLClipModel`.
#### **10. Device and Dtype Consistency**
Ensure that all models and tensors are on the correct device and have consistent data types throughout the pipeline.
```python
# In the CLIP class or wherever models are initialized
params['device'] = model_management.text_encoder_device()
params['dtype'] = model_management.text_encoder_dtype(params['device'])
```
---
### **Avoiding Model Reloading**
By ensuring that the Qwen2-VL model and `t5_context_embedder` are loaded once and passed through the pipeline, we avoid unnecessary re-initialization.
- **Loading Once**: Load models in the loader node (`DualClipLoader` or a custom loader node).
- **Passing Through**: Pass the models via `params` to the `CLIP` class and ensure they are used in the `cond_stage_model`.
---
### **Final Notes**
- **Conformity to Interfaces**: Ensure that the custom classes (`Qwen2VLClipModel`, `Qwen2VLTokenizer`) conform to the interfaces expected by the `CLIP` class.
- **Flexibility**: This approach leverages the existing ComfyUI infrastructure, maintaining consistency and flexibility.
- **Minimal Changes**: By extending existing classes and functions, we keep changes minimal and focused.
---
### Existing Solutions
_No response_
### Other
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.