commit 69cfed368f36107d846659dfb0ea84eb6f3a19ec Author: molbal Date: Wed Jan 21 12:04:53 2026 +0100 Initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..91217e1 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.idea +/__pycache__ diff --git a/README.md b/README.md new file mode 100644 index 0000000..fdd86d0 --- /dev/null +++ b/README.md @@ -0,0 +1,36 @@ +# LLM-CLIP Custom Nodes for ComfyUI + +These nodes allow you to use a distilled Qwen3-4B model as your text encoder in Stable Diffusion XL. + +## Installation + +0. Download https://huggingface.co/molbal/qwen-clip-resampler-adapter/blob/main/resampler.pth and place it in the **text encoders** directory +1. Navigate to your ComfyUI custom nodes folder. +2. Clone this repository. +3. Install the required Python packages. +```bash +pip install transformers accelerate peft bitsandbytes +``` + +## Nodes + +### LLM CLIP Loader + +This node loads the LLM and the Resampler weights. + +* **base_repo_id:** The Hugging Face ID of the base LLM (use `unsloth/Qwen3-4B-unsloth-bnb-4bit`). +* **resampler_name:** The local `.pth` file containing your trained Resampler weights. +* **lora_repo_id:** The Hugging Face ID or local path of your trained LoRA adapter. (use `molbal/qwen-clip-resampler-adapter`) + +### LLM CLIP Text Encode + +This node converts your text prompt into SDXL conditioning. + +* **text:** Multiline input for your prompt. +* **clip_bundle:** Connection from the Loader node. + +## Usage +![img.png](img.png) + + +Connect the `CONDITIONING` output from the **LLM CLIP Text Encode** node to the `positive` or `negative` inputs of a KSampler. This replaces the standard CLIP Text Encode nodes. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..66b426d --- /dev/null +++ b/__init__.py @@ -0,0 +1,11 @@ +from .llm_clip_nodes import LLMClipLoader, LLMClipTextEncode + +NODE_CLASS_MAPPINGS = { + "LLMClipLoader": LLMClipLoader, + "LLMClipTextEncode": LLMClipTextEncode +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "LLMClipLoader": "LLM CLIP Loader (Qwen)", + "LLMClipTextEncode": "LLM CLIP Text Encode" +} \ No newline at end of file diff --git a/img.png b/img.png new file mode 100644 index 0000000..419192b Binary files /dev/null and b/img.png differ diff --git a/llm_clip_nodes.py b/llm_clip_nodes.py new file mode 100644 index 0000000..a5ea9b5 --- /dev/null +++ b/llm_clip_nodes.py @@ -0,0 +1,110 @@ +import torch +import os +import folder_paths +import comfy.model_management +from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig +from peft import PeftModel +from .resampler import SDXLResampler + + +class LLMClipLoader: + @classmethod + def INPUT_TYPES(s): + file_list = folder_paths.get_filename_list("clip") + return { + "required": { + "base_repo_id": ("STRING", {"default": "unsloth/Qwen3-4B-unsloth-bnb-4bit"}), + "resampler_name": (file_list,), + }, + "optional": { + "lora_repo_id": ("STRING", {"default": "molbal/qwen-clip-resampler-adapter"}), + } + } + + RETURN_TYPES = ("LLM_CLIP_BUNDLE",) + FUNCTION = "load_llm_clip" + CATEGORY = "LLM_CLIP" + TITLE = "LLM CLIP Loader (Full HF)" + + def load_llm_clip(self, base_repo_id, resampler_name, lora_repo_id=None): + device = comfy.model_management.text_encoder_device() + resampler_path = folder_paths.get_full_path("clip", resampler_name) + + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch.float16, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + llm_int8_enable_fp32_cpu_offload=True + ) + + print(f"LLM CLIP: Loading base {base_repo_id}...") + tokenizer = AutoTokenizer.from_pretrained(base_repo_id, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained( + base_repo_id, + quantization_config=bnb_config, + trust_remote_code=True, + device_map="auto" + ) + + if lora_repo_id and lora_repo_id.strip() != "": + print(f"LLM CLIP: Applying LoRA from {lora_repo_id}...") + model = PeftModel.from_pretrained(model, lora_repo_id) + + # 1. Initialize Resampler + hidden_size = model.config.hidden_size + # 2. CRITICAL: Move to device AND cast to HALF immediately + resampler = SDXLResampler(llm_dim=hidden_size, sdxl_dim=2048).to(device).half() + + resampler.load_state_dict(torch.load(resampler_path, map_location=device)) + resampler.eval() + + torch.cuda.empty_cache() + + return ({"model": model, "tokenizer": tokenizer, "resampler": resampler},) + + +class LLMClipTextEncode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "clip_bundle": ("LLM_CLIP_BUNDLE",), + "text": ("STRING", {"multiline": True}), + } + } + + RETURN_TYPES = ("CONDITIONING",) + FUNCTION = "encode" + CATEGORY = "LLM_CLIP" + + def encode(self, clip_bundle, text): + model = clip_bundle["model"] + tokenizer = clip_bundle["tokenizer"] + resampler = clip_bundle["resampler"] + + model_device = next(model.parameters()).device + resampler_device = next(resampler.parameters()).device + + inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256) + inputs = {k: v.to(model_device) for k, v in inputs.items()} + + with torch.no_grad(): + outputs = model(**inputs, output_hidden_states=True) + + # 3. CRITICAL: Ensure hidden states are cast to match resampler (Half) + last_hidden = outputs.hidden_states[-1].to(resampler_device).half() + + # Run resampler + cond, pooled = resampler(last_hidden) + + torch.cuda.empty_cache() + + # Final output tensors for ComfyUI need to be Float32 (Comfy's standard) + return ([[cond.float(), {"pooled_output": pooled.float()}]],) + + +NODE_CLASS_MAPPINGS = { + "LLMClipLoader": LLMClipLoader, + "LLMClipTextEncode": LLMClipTextEncode +} \ No newline at end of file diff --git a/resampler.py b/resampler.py new file mode 100644 index 0000000..a3f7a64 --- /dev/null +++ b/resampler.py @@ -0,0 +1,58 @@ +import torch +import torch.nn as nn + +class SDXLResampler(nn.Module): + def __init__(self, llm_dim=2048, sdxl_dim=2048, num_queries=77, depth=4, heads=16): + super().__init__() + + # 1. The "Buckets" (77 fixed vectors that query the LLM output) + self.latents = nn.Parameter(torch.randn(1, num_queries, sdxl_dim) * 0.02) + + # 2. Projector + self.proj_in = nn.Linear(llm_dim, sdxl_dim) if llm_dim != sdxl_dim else nn.Identity() + + # 3. Attention Layers + self.layers = nn.ModuleList([ + nn.TransformerDecoderLayer( + d_model=sdxl_dim, + nhead=heads, + dim_feedforward=sdxl_dim * 4, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True + ) + for _ in range(depth) + ]) + + # 4. Output Norms + self.norm_out = nn.LayerNorm(sdxl_dim) + + # 5. Pooled Head (For the "Vibe" vector [1280]) + self.pool_proj = nn.Linear(sdxl_dim, 1280) + + def forward(self, llm_hidden_states): + """ + Input: [Batch, Seq_Len, LLM_Dim] (Qwen Output) + Output: + - prompt_embeds: [Batch, 77, 2048] + - pooled_embeds: [Batch, 1280] + """ + batch_size = llm_hidden_states.shape[0] + + # Project Qwen dimension to SDXL dimension + x = self.proj_in(llm_hidden_states) + + # Expand latents for the batch + latents = self.latents.repeat(batch_size, 1, 1) + + # Apply Cross-Attention (Latents attend to Qwen Output) + for layer in self.layers: + latents = layer(tgt=latents, memory=x) + + latents = self.norm_out(latents) + + # Create Pooled Vector (Average of all 77 tokens -> projected) + pooled = self.pool_proj(latents.mean(dim=1)) + + return latents, pooled diff --git a/workflow.json b/workflow.json new file mode 100644 index 0000000..7b273e0 --- /dev/null +++ b/workflow.json @@ -0,0 +1,465 @@ +{ + "id": "9240f5b3-ee50-48f6-89a3-0edcba07e8f8", + "revision": 0, + "last_node_id": 21, + "last_link_id": 26, + "nodes": [ + { + "id": 13, + "type": "LLMClipLoader", + "pos": [ + -1072, + -32 + ], + "size": [ + 416, + 160 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LLM_CLIP_BUNDLE", + "type": "LLM_CLIP_BUNDLE", + "links": [ + 14, + 15 + ] + } + ], + "properties": { + "Node name for S&R": "LLMClipLoader", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "repo_id": true, + "resampler_name": true, + "lora_name": true + } + } + }, + "widgets_values": [ + "unsloth/Qwen3-4B-unsloth-bnb-4bit", + "resampler.pth", + "molbal/qwen-clip-resampler-adapter" + ] + }, + { + "id": 1, + "type": "CheckpointLoaderSimple", + "pos": [ + -1072, + -192 + ], + "size": [ + 416, + 98 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 3 + ] + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [] + }, + { + "name": "VAE", + "type": "VAE", + "links": [ + 8 + ] + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "ckpt_name": true + } + } + }, + "widgets_values": [ + "albedobaseXL_v21.safetensors" + ] + }, + { + "id": 7, + "type": "VAEDecode", + "pos": [ + -96, + 304 + ], + "size": [ + 272, + 48 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 9 + ] + } + ], + "properties": { + "Node name for S&R": "VAEDecode", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": {} + } + }, + "widgets_values": [] + }, + { + "id": 6, + "type": "EmptyLatentImage", + "pos": [ + -464, + -224 + ], + "size": [ + 270, + 106 + ], + "flags": { + "collapsed": false + }, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 6 + ] + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "width": true, + "height": true, + "batch_size": true + } + } + }, + "widgets_values": [ + 720, + 1280, + 1 + ] + }, + { + "id": 3, + "type": "LLMClipTextEncode", + "pos": [ + -592, + -48 + ], + "size": [ + 400, + 200 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "clip_bundle", + "type": "LLM_CLIP_BUNDLE", + "link": 14 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 4 + ] + } + ], + "properties": { + "Node name for S&R": "LLMClipTextEncode", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "text": true + } + } + }, + "widgets_values": [ + "A woman standing in green grass, in golden hour." + ] + }, + { + "id": 5, + "type": "KSampler", + "pos": [ + -96, + -224 + ], + "size": [ + 270, + 474 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 3 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 4 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 5 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 6 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 7 + ] + } + ], + "properties": { + "Node name for S&R": "KSampler", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "seed": true, + "steps": true, + "cfg": true, + "sampler_name": true, + "scheduler": true, + "denoise": true + } + } + }, + "widgets_values": [ + 1, + "fixed", + 30, + 8, + "euler", + "normal", + 1 + ] + }, + { + "id": 8, + "type": "PreviewImage", + "pos": [ + 224, + -224 + ], + "size": [ + 432, + 576 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 9 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "PreviewImage", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": {} + } + }, + "widgets_values": [] + }, + { + "id": 4, + "type": "LLMClipTextEncode", + "pos": [ + -592, + 192 + ], + "size": [ + 400, + 160 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "clip_bundle", + "type": "LLM_CLIP_BUNDLE", + "link": 15 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 5 + ] + } + ], + "properties": { + "Node name for S&R": "LLMClipTextEncode", + "ue_properties": { + "version": "7.0.1", + "widget_ue_connectable": { + "text": true + } + } + }, + "widgets_values": [ + "Worst quality, bad quality, deformed" + ] + } + ], + "links": [ + [ + 3, + 1, + 0, + 5, + 0, + "MODEL" + ], + [ + 4, + 3, + 0, + 5, + 1, + "CONDITIONING" + ], + [ + 5, + 4, + 0, + 5, + 2, + "CONDITIONING" + ], + [ + 6, + 6, + 0, + 5, + 3, + "LATENT" + ], + [ + 7, + 5, + 0, + 7, + 0, + "LATENT" + ], + [ + 8, + 1, + 2, + 7, + 1, + "VAE" + ], + [ + 9, + 7, + 0, + 8, + 0, + "IMAGE" + ], + [ + 14, + 13, + 0, + 3, + 0, + "LLM_CLIP_BUNDLE" + ], + [ + 15, + 13, + 0, + 4, + 0, + "LLM_CLIP_BUNDLE" + ] + ], + "groups": [], + "config": {}, + "extra": { + "workflowRendererVersion": "LG", + "ue_links": [], + "ds": { + "scale": 1.0152559799477274, + "offset": [ + 1183.1848515598294, + 1123.4393206522357 + ] + }, + "node_versions": { + "comfy-core": "0.9.2" + }, + "links_added_by_ue": [], + "frontendVersion": "1.36.14" + }, + "version": 0.4 +} \ No newline at end of file