# ComfyUI-LensTurbo

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
[![ComfyUI](https://img.shields.io/badge/ComfyUI-custom%20node-purple)](https://github.com/comfyanonymous/ComfyUI)
[![Models](https://img.shields.io/badge/HF-microsoft%2FLens--Turbo-orange)](https://huggingface.co/microsoft/Lens-Turbo)

> ComfyUI nodes + standalone CLI for **Microsoft Lens / Lens-Turbo / Lens-Base** — a 3.8B-parameter text-to-image diffusion model that pairs a GPT-OSS-20B text encoder with a FLUX.2 VAE.

The killer feature: a **hybrid CPU/GPU placement** strategy that fits the whole pipeline on a single 24 GB consumer GPU. The 40 GB dequantized text encoder stays on CPU; the transformer and VAE run on GPU. Measured **11.6 s per 1024² image at 4 steps**, peak **11.11 GB VRAM** on an RTX 3090.

## Table of contents

- [Features](#features)
- [Variants](#variants)
- [Hardware requirements](#hardware-requirements)
- [Install](#install)
- [Usage — ComfyUI](#usage--comfyui)
- [Usage — CLI](#usage--cli)
- [Usage — Batch CLI](#usage--batch-cli)
- [Offload modes](#offload-modes)
- [Performance](#performance)
- [Troubleshooting](#troubleshooting)
- [Acknowledgments](#acknowledgments)
- [License](#license)

## Features

- **Two nodes**: `LensTurboLoader` (with cache + evict + force-reload) and `LensTurboSampler` (T2I, all the usual knobs).
- **Single-prompt CLI** (`cli/lens_turbo_cli.py`) — submit one prompt to a running ComfyUI server.
- **Batch CLI** (`cli/lens_turbo_batch.py`) — JSON-driven, resumable, retry-aware. Drop a `[{word, prompt}, ...]` list, get a directory of named PNGs.
- **24 GB single-card support** out of the box via the `cpu_text_encoder` offload mode (default).
- **Three monkey-patches** quietly handle the cross-device communication that vanilla LensPipeline doesn't.
- **Self-verifying banner** on import — `grep ComfyUI-LensTurbo comfy.log` tells you which copy is loaded.
- **Bash scripts** for upstream vendoring, dep install, HF login check, and model download.

## Variants

| Model | `--repo-id` | Steps | CFG | Use case |
|---|---|---|---|---|
| Lens-Turbo *(default)* | `microsoft/Lens-Turbo` | 4 | 1.0 | Fast preview, 4-step distilled |
| Lens (RL) | `microsoft/Lens` | 20 | 5.0 | Best quality, RL-tuned |
| Lens-Base | `microsoft/Lens-Base` | 50 | 5.0 | Research baseline, no RL/distill |

## Hardware requirements

| Resource | Minimum | Recommended |
|---|---|---|
| GPU | 24 GB VRAM (RTX 3090 / 4090) | 48 GB+ (A100 / H100) |
| System RAM | **48 GB** | 64 GB+ |
| Disk | 80 GB free (for HF cache) | NVMe SSD |
| Python | 3.10+ | 3.13 |
| CUDA | 12.0+ | 12.6 |

**Why 48 GB RAM**: GPT-OSS-20B's MoE layers are natively MXFP4-quantized, but Ampere (sm_86) cards can't run MXFP4 kernels — transformers dequantizes to bf16 (~40 GB), which lives in host memory when the encoder isn't actively running.

## Install

```bash
cd $COMFY_ROOT/custom_nodes
git clone https://github.com/aadebuger/ComfyUI-LensTurbo
cd ComfyUI-LensTurbo

# 1. Vendor the upstream `lens` Python package
bash scripts/fetch_upstream.sh

# 2. Install Python deps (handles Py3.13 override + kernels uninstall)
bash scripts/install_deps.sh

# 3. Check HF gated access
bash scripts/hf_login_check.sh
hf auth login   # if needed
```

**Agree to the licenses** on the HF web pages before downloading:

- <https://huggingface.co/microsoft/Lens-Turbo> (gated, research-only)
- <https://huggingface.co/black-forest-labs/FLUX.2-dev> (gated, FLUX Non-Commercial)
- `openai/gpt-oss-20b` is Apache-2.0, no agreement needed.

```bash
# 4. Pre-download model weights (~58 GB, sequential, resumable).
#    Run inside tmux/screen if you want to detach.
bash scripts/download_models.sh

# 5. Restart ComfyUI
pkill -9 -f main.py; sleep 5
export HF_HOME=$HOME/hf_cache
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
python main.py --listen 0.0.0.0 --port 8188
```

You should see two banner lines in the ComfyUI log:

```
[ComfyUI-LensTurbo] v0.1.6 loaded from .../custom_nodes/ComfyUI-LensTurbo/__init__.py
[ComfyUI-LensTurbo/nodes] module loaded from .../nodes (lens_source=OK)
```

If `lens_source=MISSING`, `scripts/fetch_upstream.sh` didn't run.

## Usage — ComfyUI

Load `example_workflows/lens_turbo_t2i.json` in the ComfyUI UI:

```
LensTurboLoader  →  LensTurboSampler  →  SaveImage
```

Default sampler knobs match Lens-Turbo's recipe: `base_resolution=1024`, `aspect_ratio=1:1`, `num_inference_steps=4`, `guidance_scale=1.0`.

## Usage — CLI

Submit a single prompt to a running ComfyUI server:

```bash
python cli/lens_turbo_cli.py \
    "a photorealistic photo of an astronaut riding a horse on mars" \
    --ip 192.168.1.10 --gpu 0 \
    --base-resolution 1024 --aspect-ratio 1:1 \
    --steps 4 --cfg 1.0 --seed 42 \
    --out ./out
```

See [`cli/README.md`](cli/README.md) for all flags.

## Usage — Batch CLI

Drop a JSON file like:

```json
[
  {"word": "apple",  "prompt": "a fresh red apple, studio light, 4k"},
  {"word": "sunset", "prompt": "warm sunset over the ocean, cinematic"}
]
```

Then:

```bash
python cli/lens_turbo_batch.py vocab.json \
    --ip 192.168.1.10 --gpu 0 \
    --out-dir ./vocab_images
```

Outputs → `<out-dir>/<sanitized-word>.png`. Existing files are skipped on re-run (use `--force` to overwrite). CJK preserved in filenames; spaces and punctuation become `_`.

The JSON root may also be `{"entries": [...]}`, `{"data": [...]}`, or `{"items": [...]}` — auto-detected.

## Offload modes

The `offload_mode` knob on `LensTurboLoader` (also `--offload-mode` on both CLIs):

| Mode | VRAM | Speed | Notes |
|---|---|---|---|
| `cpu_text_encoder` *(default)* | ~11 GB | ~12 s/image | **24 GB-safe**. Text encoder permanently on CPU, transformer + VAE on GPU. Three monkey-patches bridge device crossings. |
| `model` | ~50 GB | ~6 s/image | `enable_model_cpu_offload()`. Needs ≥48 GB VRAM to hold the dequantized text encoder when active. |
| `none` | ~50 GB | ~5 s/image | `pipe.to(device)`, no offload. ≥48 GB VRAM. Fastest. |
| `sequential` | — | — | **KNOWN BROKEN** on accelerate ≥1.x — meta tensor bug in `pre_forward` + custom Lens TextEncoder forward. Kept for completeness, don't use. See INSTALL_NOTES L6. |

## Performance

**Measured on RTX 3090 (24 GB), `cpu_text_encoder` mode, Lens-Turbo:**

| Metric | Value |
|---|---|
| Peak VRAM | **11.11 GB** (46% of card) |
| Cold load (`from_pretrained`) | ~117 s |
| Warm sample @ 1024² / 4 steps | **11.6 s** total |
| - text encoding (CPU, bf16) | ~5-7 s |
| - transformer 4 steps | 4 s (≈1 s/step) |
| - VAE decode | <1 s |

After the cold load, the `LensTurboLoader` cache makes subsequent images hit instantly. Batch jobs amortize the cold load over the whole run — 16 images take ~3-4 minutes including cold start.

## Troubleshooting

| Symptom | Likely cause | Fix |
|---|---|---|
| `Illegal header value b'kernels/0.14.0; ...'` | `kernels==0.14.0` has a user_agent header bug | `uv pip uninstall kernels`. RTX 3090 can't run MXFP4 native anyway. |
| `OutOfMemoryError` during sampling | Wrong offload mode for your VRAM | Use `offload_mode="cpu_text_encoder"` (default) |
| `Cannot copy out of meta tensor; no data!` | `enable_sequential_cpu_offload` + Lens custom forward | Switch to `offload_mode="cpu_text_encoder"` |
| `Cannot generate cpu tensor from cuda generator` | `_execution_device` resolved to CPU because text_encoder is there | Use `cpu_text_encoder` mode (auto-patches the property) |
| `Expected all tensors to be on the same device` | encoder outputs not moved to GPU | Use `cpu_text_encoder` mode (auto-wraps encode_layers) |
| Nodes don't appear in the UI | `__init__.py` errored at import, or `.bak` directory is shadowing | `grep ComfyUI-LensTurbo comfy.log`, check `IMPORT FAILED`, move any `.bak` directories out of `custom_nodes/` |
| First sample times out via CLI | Cold load > 600 s default WS timeout | `--ws-recv-timeout 1800` |

See [INSTALL_NOTES.md](INSTALL_NOTES.md) for the full deployment story including 10 pitfalls discovered while bringing this online.

## Acknowledgments

- **Lens / Lens-Turbo / Lens-Base** — research by Microsoft. Model weights at [microsoft/Lens-Turbo](https://huggingface.co/microsoft/Lens-Turbo). Architecture pairs a large LLM-based text encoder (GPT-OSS-20B) with a diffusion transformer and FLUX.2 VAE.
- **GPT-OSS-20B** — text encoder by OpenAI ([openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b), Apache-2.0).
- **FLUX.2 VAE** — by Black Forest Labs ([black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev)).
- **ComfyUI** — the host runtime by comfyanonymous.

## License

MIT for this wrapper — see [LICENSE](LICENSE).

Model weights are released by Microsoft for **research purposes only**, not for product or service deployment. Consult the upstream HF model card for the full terms. This repository does not redistribute weights.