Files
2024-10-03 02:03:29 +02:00

215 lines
7.7 KiB
Markdown

# ComfyUI GLM-4 Wrapper
This repository contains custom nodes for ComfyUI, specifically designed to enhance and infer prompts using the GLM-4 model on local hardware.
The nodes leverage the GLM-4 model to generate detailed and descriptive image/video captions or enhance user-provided prompts, among regular inference.
Prompts and inference can be combined with image if `THUDM/glm-4v-9b` model is used.
All models will be downloaded automatically through HuggingFace.co. `THUDM/glm-4v-9b` will hold ~26 GB of hdd space and `THUDM/glm-4-9b` will hold ~18 GB of hdd space.
The nodes containes an "unload_model" option which frees up VRAM space and makes it suitable for workflows that requires larger VRAM space, like FLUX.1-dev and CogVideoX-5b(-I2V).
The prompt enhancer is based on this example from THUDM [convert_demo.py](https://github.com/THUDM/CogVideo/blob/main/inference/convert_demo.py).
Thier demo is only for usage through OpenAI API and I wanted to build something local.
Hope you will enjoy your enhanced prompts and inference capabilities of these models. They are great!
## Update 2024-10-03
Added support for quantized models. They are performing exceptionally well. Check metrics below.
Model `alexwww94/glm-4v-9b-gptq-4bit` is significatly more lightweight than the original and will hold ~8.5 GB of hdd space.
Model `alexwww94/glm-4v-9b-gptq-3bit` is even more lightweight and will hold ~7.6 GB of hdd space.
[Link to metrics](https://huggingface.co/alexwww94/glm-4v-9b-gptq-4bit#metrics)
## Features
- **GLM-4 Model Loader**: Load various GLM-4 models with different precision and quantization settings.
- **GLM-4 Prompt Enhancer**: Enhances base prompts using the GLM-4 model.
- **GLM-4 Inferencing**: Performs inference using various GLM-4 models.
## Installation
1. Navigate to ComfyUI custom nodes:
```bash
cd /your/path/to/ComfyUI/ComfyUI/custom_nodes/
```
2. Clone the repository:
```bash
git clone https://github.com/Nojahhh/ComfyUI_GLM4_Wrapper.git
```
3. Navigate to the cloned directory:
```bash
cd ComfyUI_GLM4_Wrapper
```
4. Install the required dependencies:
```bash
../../python_embeded python.exe -m pip install -r requirements.txt
```
## Usage
### GLM-4 Model Loader
The `GLM4ModelLoader` class is responsible for loading GLM-4 models. It supports various models and precision settings.
#### Input Types
- **model**: Choose the GLM-4 model to load. Model will download automatically from HuggingFace.co
- **precision**: Precision type (`fp16`, `fp32`, `bf16`). `THUDM/glm-4v-9b` requires `bf16` and is set to run in 8-bit by default based on it's size. `alexwww94/glm-4v-9b-gptq-4bit` requires `bf16` and is set to run in 4-bit by default. `alexwww94/glm-4v-9b-gptq-3bit` requires `bf16` and is set to run in 3-bit by default.
- **quantization**: Set the number of bits for quantization (`4`, `8`, `16`).
#### Output
- **GLM4Pipeline**: The GLM-4 pipeline.
### GLM-4 Prompt Enhancer
Enhances a given prompt using the GLM-4 model.
#### Input Parameters
- **GLMPipeline**: Provide a GLM-4 pipeline.
- **prompt**: Base prompt to enhance.
- **max_tokens**: Maximum number of output tokens.
- **temperature**: Temperature parameter for sampling.
- **top_k**: Top-k parameter for sampling.
- **top_p**: Top-p parameter for sampling.
- **repetition_penalty**: Repetition penalty for sampling.
- **image** (optional): Image to enhance the prompt. Only works with `THUDM/glm-4v-9b` and `alexwww94/glm-4v-9b-gptq-4bit` and `alexwww94/glm-4v-9b-gptq-3bit`.
- **unload_model**: Unload the model after use.
#### Output
- **enhanced_prompt**: The enhanced prompt.
### GLM-4 Inferencing
Performs inference using the GLM-4 model.
#### Input Parameters
- **GLMPipeline**: Provide a GLM-4 pipeline.
- **system_prompt**: System prompt for inferencing.
- **user_prompt**: User prompt for inferencing.
- **max_tokens**: Maximum number of output tokens.
- **temperature**: Temperature parameter for sampling.
- **top_k**: Top-k parameter for sampling.
- **top_p**: Top-p parameter for sampling.
- **repetition_penalty**: Repetition penalty for sampling.
- **image** (optional): Image to use as input for inferencing. Only works with `THUDM/glm-4v-9b` and `alexwww94/glm-4v-9b-gptq-4bit` and `alexwww94/glm-4v-9b-gptq-3bit`.
- **unload_model**: Unload the model after use.
#### Output
- **output_text**: The generated text from the model.
## Node Class Mappings
- **GLM-4 Model Loader**: `GLM4ModelLoader`
- **GLM-4 Prompt Enhancer**: `GLM4PromptEnhancer`
- **GLM-4 Inferencing**: `GLM4Inference`
## Node Display Name Mappings
- **GLM-4ModelLoader**: "GLM-4 Model Loader"
- **GLM-4PromptEnhancer**: "GLM-4 Prompt Enhancer"
- **GLM-4Inference**: "GLM-4 Inferencing"
## Supported Models
The following GLM-4 models are supported by this wrapper:
| Model Name | Size | Recommended Precision |
|---------------------------------|-------|------------------------|
| `alexwww94/glm-4v-9b-gptq-4bit` | 9B | `bf16` (4-bit quant) |
| `alexwww94/glm-4v-9b-gptq-3bit` | 9B | `bf16` (3-bit quant) |
| `THUDM/glm-4v-9b` | 9B | `bf16` (4/8-bit quant) |
| `THUDM/glm-4-9b` | 9B | `fp16`, `fp32`, `bf16` |
| `THUDM/glm-4-9b-chat` | 9B | `fp16`, `fp32`, `bf16` |
| `THUDM/glm-4-9b-chat-1m` | 9B | `fp16`, `fp32`, `bf16` |
| `THUDM/LongCite-glm4-9b` | 9B | `fp16`, `fp32`, `bf16` |
| `THUDM/LongWriter-glm4-9b` | 9B | `fp16`, `fp32`, `bf16` |
### Notes:
- `THUDM/glm-4v-9b` requires `bf16` precision and is default 8-bit quantization due to its size and the typical VRAM limitations of consumer-grade GPUs (often 24GB or less).
- `alexwww94/glm-4v-9b-gptq-4bit` requires `bf16` and is default 4-bit.
- `alexwww94/glm-4v-9b-gptq-3bit` requires `bf16` and is default 3-bit.
- Only `THUDM/glm-4v-9b` and `alexwww94/glm-4v-9b-gptq-4bit` and `alexwww94/glm-4v-9b-gptq-3bit` models are able to handle image input.
## Example Usage
Below is an example of how to use the GLM-4 Prompt Enhancer and GLM-4 Inferencing nodes in your code:
### GLM-4 Prompt Enhancer
```python
from comfyui_glm4_wrapper import GLM4ModelLoader, GLM4PromptEnhancer, GLM4Inference
# Load the model
model_loader = GLM4ModelLoader()
pipeline = model_loader.gen(model="THUDM/glm-4v-9b", precision="bf16", quantization="8")[0]
# Enhance the prompt
enhancer = GLM4PromptEnhancer()
enhanced_prompt = enhancer.enhance_prompt(
GLMPipeline=pipeline,
prompt="A beautiful sunrise over the mountains",
max_tokens=200,
temperature=0.1,
top_k=40,
top_p=0.7,
repetition_penalty=1.1,
image=None, # PIL Image
unload_model=True
)
print(enhanced_prompt)
```
### GLM-4 Inferencing
```python
from comfyui_glm4_wrapper import GLM4ModelLoader, GLM4PromptEnhancer, GLM4Inference
# Load the model
model_loader = GLM4ModelLoader()
pipeline = model_loader.gen(model="THUDM/glm-4v-9b", precision="bf16", quantization="8")[0]
# Perform inference
inference = GLM4Inference()
output_text = inference.infer(
GLMPipeline=pipeline,
system_prompt="Describe the scene in detail:",
user_prompt="A bustling city street at night",
max_tokens=250,
temperature=0.7,
top_k=50,
top_p=1,
repetition_penalty=1.0,
image=None,
unload_model=True
)
print(output_text)
```
For more detailed examples and advanced usage, please refer to the documentation or the example scripts provided in the repository.
## License
This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for details.
## Acknowledgements
- [ComfyUI](https://github.com/comfyui)
- [Transformers](https://github.com/huggingface/transformers)
- [THUDM](https://github.com/THUDM)
## Contact
For any questions or feedback, please open an issue on GitHub or contact me at [mellin.johan@gmail.com](mailto:mellin.johan@gmail.com).