Comfy-Org / Comfy-Org/ComfyUI

Add Qwen2VL Support

Open
#5,777 0 comments 8 reactions 0 assignees View on GitHub
Feature
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.