Initial commit

2026-01-21 12:04:53 +01:00
commit 69cfed368f
7 changed files with 682 additions and 0 deletions
@@ -0,0 +1,2 @@
+.idea
+/__pycache__
@@ -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.
@@ -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"
+}
@@ -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
+}
@@ -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
@@ -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
+}