Extensions/ComfyUI-MacOffload-TE

ComfyUI Extension

ComfyUI-MacOffload-TE

Experimental text encoder offloading for Apple Silicon / Mac setups.

By j0eyhu·Created 8 days ago·Updated a day ago· 0

j0eyhu/ComfyUI-MacOffload-TE

Nodes2

On cloudLocal install

Category⚡ MacOffload/Conditioning

Stars0

Updateda day ago

Nodes (2)

Remote CLIP Text Encode (Flux Klein) ⚡

Klein's Text Encoder Won't Fit on 8GB, So Let the Mac Encode the Prompt

⚡ MacOffload/Conditioning

Remote Text Encode Qwen Image 2.1 ⚡

Text, Sixteen Reference Images, and a Latent — Encoded on Your Mac Instead

⚡ MacOffload/Conditioning

Readme

ComfyUI-MacOffload-TE ⚡

High-Performance Cross-Platform Conditioning Bridge for ComfyUI
Offload heavy conditioning models (Flux.2 Klein, Flux.3 Klein, Qwen-Image-2.1) from your VRAM-constrained GPU (e.g. RTX 4060 8GB) to a macOS machine with Apple Silicon Unified Memory (MPS) over your local network — eliminating PCIe model swapping and freeing local VRAM for the diffusion backbone.

1. Overview & Motivation

Running modern diffusion pipelines (such as FLUX.2 Klein, FLUX.3 Klein, or Qwen-Image-2.1) on consumer GPUs with 8GB VRAM (such as an RTX 4060 8GB) presents a severe memory bottleneck. The text and vision encoders for these models demand 8GB to 18GB of memory on their own. On a single-GPU setup, ComfyUI must constantly swap models back and forth between GPU VRAM and system RAM over the PCIe bus, introducing noticeable delays between the conditioning and sampling stages, and frequently triggering out-of-memory (OOM) errors.

ComfyUI-MacOffload-TE resolves this by offloading the entire conditioning forward pass to a remote macOS device (such as a MacBook, Mac mini, or Mac Studio) equipped with Apple Silicon Unified Memory:

0 MB Local VRAM Footprint: The conditioning stage runs entirely on the remote Mac (Apple Silicon MPS), requiring zero local GPU VRAM.
Eliminates PCIe Swapping Delays: Bypasses the continuous unloading and reloading of multi-gigabyte models over the PCIe bus. The local GPU remains primed with the diffusion model, launching sampling immediately once conditioning embeddings arrive.
More Headroom for Diffusion & Higher Resolutions: Freeing up 8GB to 18GB of local VRAM allows your GPU to dedicate its full memory capacity to the diffusion transformer (DiT) and VAE, enabling direct generation at higher resolutions without OOM crashes.
Hardware Synergy Across Devices:
- Local PC: Dedicates 100% of GPU compute and VRAM to diffusion sampling and VAE decoding.
- Remote Mac: Uses high-capacity Apple Silicon Unified Memory to keep heavy encoders resident, returning conditioning tensors over the local network.

2. Architecture & Communication Flow

Text Architecture Overview

+-------------------------------------------------------------+
|                      Windows Client                         |
|  [ComfyUI]                                                  |
|   └── Custom Node (Facade + Helpers)                        |
|        ├── Read server address from local disk config       |
|        ├── Pack prompt text & optional reference images     |
|        ├── Serialize to safe, data-only safetensors stream  |
|        └── HTTP POST (application/octet-stream)             |
+------------------------------+------------------------------+
                               |
                   Local Area Network (LAN)
                   Low-latency safetensors Stream
                               |
+------------------------------v------------------------------+
|                     macOS Server (MPS)                      |
|  [FastAPI Bridge]                                           |
|   ├── Route: /v1/encode/{model_name}                        |
|   ├── Parse safetensors payload (no pickle / no RCE risk)   |
|   ├── Dynamic Model Pool (Auto-discovery + LRU Eviction)    |
|   ├── Direct Accelerated MPS Loading (safetensors mmap)     |
|   └── Return serialized Conditioning tensors (safetensors)  |
+-------------------------------------------------------------+

Flow Sequence

sequenceDiagram
    autonumber
    participant Win as Windows ComfyUI Client
    participant Srv as macOS FastAPI Server
    participant MPS as Apple Silicon MPS Pool

    Win->>Srv: POST /v1/encode/{model_name} (safetensors Payload)
    Note over Srv: Dynamic Model Manager checks resident pool
    alt Model Not In Unified Memory
        Srv->>MPS: Direct accelerated load into Unified Memory (bfloat16)
    end
    Srv->>MPS: Execute Forward Pass (Tokenization & Embeddings)
    MPS-->>Srv: Return Conditioning / Latent Tensors
    Srv-->>Win: HTTP 200 (Serialized safetensors Binary Stream)
    Note over Win: Inject conditioning into diffusion sampling

3. Supported Models

| Model Folder Name | DiT Model | Text Encoder Architecture | Hugging Face Repository | Summary | | :--- | :--- | :--- | :--- | :--- | | flux2-klein-4b-uncensored-text-encoder | FLUX.2 Klein (4B) | Qwen3-4B | ponpoke/flux2-klein-4b-uncensored-text-encoder | Uncensored text encoder variant for FLUX.2 Klein 4B DiT. Processed to remove refusal vectors to prevent safety filter refusals. | | flux2-klein-9b-uncensored-text-encoder | FLUX.2 Klein (9B) | Qwen3-8B | ponpoke/flux2-klein-9b-uncensored-text-encoder | Uncensored text encoder variant for FLUX.2 Klein 9B DiT with higher parameter capacity for nuanced prompt comprehension without safety filter refusals. | | Qwen3-VL-4b-Heretic | FLUX.3 Klein (5.6B) | Qwen3-VL-4B | DreamFast/Qwen3-VL-4b-Heretic | Abliterated version of Qwen3-VL-4B-Instruct processed using the Heretic optimizer to remove refusal behaviors. Eliminates safety filter refusals while preserving underlying language comprehension and vision capabilities. | | Qwen-Image-2.1-Text-Encoder-Heretic | Qwen-Image-2.1 | Qwen3-VL-8B | pottokao/Qwen-Image-2.1-Text-Encoder-Heretic | Abliterated "Heretic" version of the Qwen-Image-2.1 text encoder based on the Qwen3-VL-8B architecture. Removes safety alignment refusals using directional ablation on o_proj and down_proj layers while preserving functional vision-language performance. |

4. Installation & Setup

4.1 Client Setup (Windows ComfyUI)

Clone or copy this repository into your ComfyUI custom_nodes/ directory:
```
cd ComfyUI/custom_nodes
git clone <repo-url> ComfyUI-MacOffload-TE
```
Restart ComfyUI. The nodes will be available under category: ⚡ MacOffload/Conditioning
Configuring Mac Server Address:
- The Mac server address is read strictly from local disk or environment variable, and is never embedded inside shared workflow files (protecting users from SSRF and untrusted redirects).
- Configure it in either of two ways:
  - Config file: Edit macoffload_config.json inside the ComfyUI-MacOffload-TE custom node directory:
```
{
  "server_url": "http://<your-mac-lan-ip>:8288"
}
```
  - Environment variable: Set MACOFFLOAD_SERVER_URL in your environment (e.g. export MACOFFLOAD_SERVER_URL="http://<your-mac-lan-ip>:8288" or via your ComfyUI startup batch script).
Node Controls & UI Features:
- 🔄 Fetch Models from Server: Queries the remote Mac server using your saved address and populates the model_name dropdown.
- Saved Workflow Memory: The node preserves selected models across workflow reloads.
- ❓ Help Badge: Click the ? icon in the node title bar (or right-click the node and select "About this node ❓") to view diagnostics and usage information.

4.2 Server Setup (macOS / Apple Silicon)

Navigate to the server/ directory on your Mac:
```
cd ComfyUI-MacOffload-TE/server
```

Create and activate a Python virtual environment:

python3 -m venv .venv
source .venv/bin/activate

Install server dependencies:
```
pip install -r requirements.txt
```

Configure model paths: Copy the .env.example file to .env:

cp .env.example .env

Edit .env to point MODELS_DIR to your model directory (defaults to ~/models):

MODELS_DIR=/Users/<your-username>/models
BRIDGE_PORT=8288
MAX_MODELS=1

Place your model folders inside MODELS_DIR:

~/models/
├── Qwen-Image-2.1-Text-Encoder-Heretic/
├── Qwen3-VL-4b-Heretic/
├── flux2-klein-4b-uncensored-text-encoder/
└── flux2-klein-9b-uncensored-text-encoder/

Launch the server:
```
python server.py
```

5. Server Startup Arguments & Options

The server supports command-line arguments, .env file, and environment variables:

Example CLI Commands:

# Run with default configuration (.env or ~/models)
python server.py

# Custom port and custom models directory
python server.py --port 8288 --models-dir /Volumes/ExternalSSD/models

# Keep at most 1 resident model with 300s idle auto-unload
python server.py --max-models 1 --idle-timeout 300

6. Node Reference

1. `Remote CLIP Text Encode (Flux Klein) ⚡`

Class Identifier: MACOFFLOAD_RemoteCLIPTextEncode
Category: ⚡ MacOffload/Conditioning
Description: Drop-in replacement for native CLIPTextEncode in Flux.2 Klein and Flux.3 Klein workflows.
Inputs:
- model_name (COMBO): Target model folder name (e.g. flux2-klein-4b-uncensored-text-encoder, Qwen3-VL-4b-Heretic).
- text (STRING): Positive or negative text prompt.
Outputs:
- conditioning (CONDITIONING): Formatted conditioning tuple containing stacked hidden states ([batch, seq_len, 12288] for Flux.2, [batch, seq_len, 20480] for Flux.3) and pooled representation.

2. `Remote Text Encode Qwen Image 2.1 ⚡`

Class Identifier: MACOFFLOAD_RemoteTextEncodeQwenImage21
Category: ⚡ MacOffload/Conditioning
Description: Multimodal text and multi-image reference encoder for Qwen-Image-2.1.
Inputs:
- model_name (COMBO): Target model folder name (e.g. Qwen-Image-2.1-Text-Encoder-Heretic).
- prompt (STRING): Positive text prompt.
- negative_prompt (STRING): Negative text prompt.
- resolution (INT, default: 1024): Target dimension for reference image scaling (multiples of 32).
- vae (VAE, optional): WanVAE model used to encode reference images into reference latents.
- image_1 ~ image_16 (IMAGE, optional): Up to 16 input reference images.
Outputs:
- positive (CONDITIONING): Positive conditioning with sliced hidden states and vision slots.
- negative (CONDITIONING): Negative conditioning.
- latent (LATENT): WanVAE-compatible 64-channel latent tensor [1, 64, H//16, W//16].

7. Server REST API Reference

List Available Models

Endpoint: GET /v1/models
Description: Scans MODELS_DIR and returns discovered model folder names, their encoder categories, and resident status.

Response:

{
  "models": [
    "Qwen-Image-2.1-Text-Encoder-Heretic",
    "Qwen3-VL-4b-Heretic",
    "flux2-klein-4b-uncensored-text-encoder",
    "flux2-klein-9b-uncensored-text-encoder"
  ],
  "model_types": {
    "Qwen-Image-2.1-Text-Encoder-Heretic": "qwen_image",
    "Qwen3-VL-4b-Heretic": "flux3_klein",
    "flux2-klein-4b-uncensored-text-encoder": "flux_klein",
    "flux2-klein-9b-uncensored-text-encoder": "flux_klein"
  },
  "resident_models": ["Qwen3-VL-4b-Heretic"]
}

Health Probe

Endpoint: GET /health