# ComfyUI-AgeShift

[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Python](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/)
[![ComfyUI](https://img.shields.io/badge/ComfyUI-custom%20nodes-7e5cad.svg)](https://github.com/comfyanonymous/ComfyUI)

> **Detect a face's age, then generate identity-preserved portraits at any target age — from infant to 90+.**

ComfyUI nodes for face age detection (ViT classifier) and age-conditioned
prompt building. Pairs with [ComfyUI-PuLID-Flux](https://github.com/balazik/ComfyUI-PuLID-Flux)
(or [InstantID](https://github.com/cubiq/ComfyUI_InstantID)) to produce
identity-preserving age-shifted portraits.

```text
source face ──► AgeDetect ──► current age = 40
                                 │
target age (e.g. 65) ──► AgePromptBuilder ──► positive + negative prompt
                                 │
                                 ▼
                      PuLID-Flux + Flux 1 dev
                                 │
                                 ▼
                       aged portrait (65 y/o)
```

---

## ✨ Features

- 🎂 **Continuous age estimation** — ViT classifier with weighted-mean
  decoding over bin midpoints; produces smooth integer ages (`32`, `47`)
  rather than snapping to bin centers.
- 🎂 **Concrete-feature prompts** — translates `target_age=65` into
  `"deep facial wrinkles, gray or white hair often thinning, sagging jowls, age spots"`,
  which Flux follows much more reliably than numeric age alone.
- 🎂 **Anti-drift negatives** — auto-injects contextual negatives so a 40-year-old
  source doesn't leak adult features into the "age 5" generation.
- 🎂 **Identity-agnostic** — works with PuLID-Flux, InstantID, IPAdapter FaceID,
  or any face-conditioning backbone. We don't reimplement ID injection; we just
  drive the prompt side.
- 🎂 **Batch CLI** — generate an age progression sequence (5 → 15 → 25 → 45 → 65)
  in one command from your laptop, talking to a remote ComfyUI server.

---

## 📦 Installation

### 1. Clone the node package

```bash
cd /path/to/ComfyUI/custom_nodes
git clone https://github.com/aadebuger/ComfyUI-AgeShift.git
cd ComfyUI-AgeShift

# in the ComfyUI venv:
source /path/to/ComfyUI/.venv/bin/activate
uv pip install -r requirements.txt
```

Dependencies are minimal (`transformers`, `pillow`, `numpy`) — almost
always already in a ComfyUI venv.

### 2. (Optional but recommended) Companion packages

The age-detect side works standalone. For identity-preserving generation
you also need PuLID-Flux + Flux:

```bash
cd /path/to/ComfyUI/custom_nodes
git clone https://github.com/balazik/ComfyUI-PuLID-Flux.git
cd ComfyUI-PuLID-Flux
uv pip install -r requirements.txt
```

| Asset | Where | Size | Notes |
|---|---|---|---|
| Flux 1 dev fp8 | `models/unet/flux1-dev-fp8.safetensors` | ~12 GB | from `Kijai/flux-fp8` (no license required) |
| CLIP-L | `models/clip/clip_l.safetensors` | ~250 MB | standard Flux |
| T5-XXL | `models/clip/t5xxl_fp8_*.safetensors` | ~5 GB | scaled or e4m3fn both work |
| VAE | `models/vae/ae.safetensors` | ~330 MB | standard Flux |
| PuLID-Flux weights | `models/pulid/pulid_flux_v0.9.1.safetensors` | ~1.1 GB | from `guozinan/PuLID` |
| InsightFace antelopev2 | `models/insightface/models/antelopev2/*.onnx` | ~360 MB | **manual placement required** — see below |
| EVA02-CLIP-L | auto-downloaded | ~430 MB | downloads on first run |

> ⚠️ **InsightFace `antelopev2` does NOT auto-download** (unlike the default
> `buffalo_l`). You must place the 5 ONNX files manually at
> `models/insightface/models/antelopev2/` (note the double `models` —
> that's how the InsightFace lib expects it). See
> [INSTALL_NOTES.md](INSTALL_NOTES.md) for one-liner downloads.

### 3. Restart ComfyUI

Search the node palette for `🎂` — four nodes should appear.

Quick verification from the CLI:

```bash
curl -s http://127.0.0.1:8188/object_info | python3 -c "
import json,sys; d = json.load(sys.stdin)
for n in ['AgeDetectorLoader','AgeDetect','AgePromptBuilder','AgeRangeBatch',
         'PulidFluxModelLoader','ApplyPulidFlux']:
    print(f'  {n}:', '✓' if n in d else '✗')"
```

---

## 🚀 Quick start

### Detect only (UI)

Open `example_workflows/age_detect_only.json` in ComfyUI, drop a portrait
into the LoadImage node, queue. The summary `PreviewAny` shows:

```text
Estimated age: 40 (adult) | top bin: 30-39 (46.8%) | top-3: 30-39=0.47, 40-49=0.43, 50-59=0.05
```

### Full pipeline (UI)

Open `example_workflows/age_shift_pulid_flux.json`, set the `target_age`
on the AgePromptBuilder node, queue.

### CLI

```bash
# Detect only — runs in ~2 seconds on CPU
uv run --script cli/age_shift_cli.py face.jpg --detect-only --ip <server>

# Single target age
uv run --script cli/age_shift_cli.py face.jpg --target-age 65 --ip <server>

# Age progression sequence
uv run --script cli/age_shift_cli.py face.jpg \
    --target-ages 5,15,25,45,65 \
    --gender male --subject man \
    --pulid-weight 0.95 --guidance 3.0 --steps 28 --seed 42 \
    --flux-unet flux1-dev-fp8.safetensors \
    --flux-t5 t5xxl_fp8_scaled.safetensors \
    --ip <server> --out ./out
```

Output: `./out/age_005_*.png`, `./out/age_015_*.png`, ...

See [cli/README.md](cli/README.md) for the full argument list.

---

## 🧠 Node reference

### 🎂 Age Detector Loader

Loads a HF ViT age-classification model, with in-process cache.

| Input | Type | Default |
|---|---|---|
| `model_id` | STRING | `nateraw/vit-age-classifier` |
| `device` | enum | `auto` |
| `dtype` | enum | `fp32` |
| `force_reload` | BOOL | False |

Output: `AGE_DETECTOR`.

### 🎂 Age Detect

| Input | Type | Default |
|---|---|---|
| `detector` | AGE_DETECTOR | — |
| `image` | IMAGE | — |
| `method` | enum | `weighted_mean` |

Outputs: `age:INT`, `age_band:STRING`, `confidence:FLOAT`, `summary:STRING`.

> Note: this is whole-image classification. For non-portrait input, run a
> face-crop step upstream (e.g. via [ComfyUI-Crystools](https://github.com/crystian/ComfyUI-Crystools)
> or InsightFace).

### 🎂 Age Prompt Builder

| Input | Type | Default |
|---|---|---|
| `target_age` | INT | 30 |
| `subject` | STRING | `person` |
| `gender` | enum | `unspecified` |
| `style` | STRING | photorealistic portrait... |
| `extra_positive` | STRING | same identity... |
| `negative_base` | STRING | cartoon, illustration... |
| `include_negative_age_drift` | BOOL | True |

Outputs: `positive:STRING`, `negative:STRING`, `summary:STRING`, `target_age:INT`.

📖 **See [PROMPTS.md](PROMPTS.md) for the complete list of generated prompts
across all 20 age ranges + examples at age 5/15/25/45/65/80.**

### 🎂 Age Range Batch

CLI-oriented. Emits JSONL `{age, band, prompt}` for a comma-separated list
of target ages.

---

## 🎯 Identity preservation tuning

The default settings work for most cases. If the generated faces look like
*different people*, walk this checklist:

### 1. Lock the seed across ages

Different seeds = different sub-identities. Pass `--seed 42` (or any fixed
value) so all ages share the same noise schedule.

### 2. Bump PuLID weight

```bash
--pulid-weight 0.95     # default 0.9; up to 1.0
```

### 3. Lower the prompt guidance

```bash
--guidance 3.0          # default 3.5
```

The age prompt is a strong signal — lowering guidance leaves more room for
the identity to assert itself.

### 4. More steps

```bash
--steps 28              # default 20
```

PuLID continues injecting identity throughout sampling, so more steps =
more chances for ID to converge.

### 5. Crop the source to a clean portrait

PuLID extracts identity via InsightFace, which works best on a centered
frontal face. If the input has background / side angle / multiple faces,
crop first.

### 6. Strengthen identity anchors in the prompt

```bash
--extra-positive "exact same person, identical facial bone structure, same eye distance, same nose shape, same lip shape"
```

### Direction-specific tips

| Direction | Tip |
|---|---|
| **Aging up** (40 → 70) | Default settings work well. Optionally `--style "photorealistic portrait, soft natural lighting"` (avoid sharp focus, which over-emphasizes wrinkles) |
| **Aging down** (40 → 20) | Default settings work well. |
| **Going to child** (adult → 5–12) | Drop `--pulid-weight 0.75–0.80` — full strength preserves adult bone structure and produces unsettling "small adult" results. |
| **Going to infant** (adult → 0–3) | Hard. Consider face-landmark ControlNet weighted on upper face only, or accept that the result is "what this person's son/daughter might look like" rather than a literal regression. |

---

## 🐞 Troubleshooting

### `PulidFluxModelLoader not found`

PuLID-Flux either isn't installed or failed to import. Check:

```bash
grep -B2 -A8 -iE "pulid|IMPORT FAILED" ~/comfy_start.log | tail -40
```

Most often `facexlib` / `onnxruntime-gpu` / `insightface` is missing:

```bash
cd /path/to/ComfyUI/custom_nodes/ComfyUI-PuLID-Flux
uv pip install -r requirements.txt --no-cache-dir
```

Then restart ComfyUI.

### `assert 'detection' in self.models` (AssertionError from InsightFace)

The `antelopev2` model files aren't at the expected path. Required layout
(note the **double `models`** directory):

```text
/path/to/ComfyUI/models/insightface/models/antelopev2/
├── 1k3d68.onnx
├── 2d106det.onnx
├── genderage.onnx
├── glintr100.onnx
└── scrfd_10g_bnkps.onnx
```

Fix:

```bash
cd /path/to/ComfyUI/models/insightface
mkdir -p models && cd models
wget "https://github.com/deepinsight/insightface/releases/download/v0.7/antelopev2.zip"
unzip antelopev2.zip && rm antelopev2.zip
```

### `No module named 'websocket'`

The CLI needs `websocket-client` (NOT `websocket`, which is a separate
abandoned package):

```bash
uv pip install websocket-client
```

Or just use `uv run --script cli/age_shift_cli.py ...` — the script's
PEP 723 header declares deps, uv handles them automatically.

### Faces all look like different people

See the [identity preservation tuning](#-identity-preservation-tuning)
section. Most often: missing `--seed`, or `--pulid-weight` too low.

### "child" generation produces creepy mini-adults

PuLID weight is too high for child direction. Try `--pulid-weight 0.75`
for child-direction runs; keep it at 0.9–0.95 for aging-up runs.

### Hangs on first generation

PuLID's first run downloads EVA02-CLIP-L (~430 MB) and (if you skipped the
manual antelopev2 step) tries to set up InsightFace. Watch the ComfyUI server
log — once it prints `apply_pulid done`, generations become fast.

---

## 🧰 Project layout

```text
ComfyUI-AgeShift/
├── __init__.py                    # NODE_CLASS_MAPPINGS
├── nodes/
│   ├── _types.py                  # age → features mapping (the brain)
│   ├── detector_loader.py         # _MODEL_CACHE pattern
│   ├── detect.py                  # ViT classifier execution
│   └── prompt_builder.py          # AgePromptBuilder + AgeRangeBatch
├── cli/
│   ├── README.md
│   └── age_shift_cli.py           # PEP 723 inline-deps script
├── example_workflows/
│   ├── age_detect_only.json       # minimal UI workflow
│   └── age_shift_pulid_flux.json  # full pipeline UI workflow
├── PROMPTS.md                     # complete prompt reference
├── INSTALL_NOTES.md               # server-side install record
├── CHANGELOG.md
├── CONTRIBUTING.md
├── pyproject.toml
└── requirements.txt
```

---

## 🛣️ Roadmap

- [ ] Built-in face-crop node (so users don't need a separate package)
- [ ] CLI flag `--pulid-fusion {mean,concat}` and `--pulid-train-step` for advanced ID tuning
- [ ] InstantID + SDXL alternative workflow (for users without Flux)
- [ ] ControlNet face-landmark integration for stronger child-direction results
- [ ] Side-by-side grid output (concatenated PNG) for sequence runs

---

## 🤝 Contributing

PRs welcome. See [CONTRIBUTING.md](CONTRIBUTING.md). The contribution most
likely to help right now: testing different age-feature phrasings on Flux
and submitting improvements to `_types.py` with before/after comparison
images.

---

## 📜 License

MIT — see [LICENSE](LICENSE).

The age classifier (`nateraw/vit-age-classifier`) is MIT-licensed.
PuLID, Flux, and InsightFace weights have their own licenses — see those
projects for details.

---

## 🙏 Acknowledgments

- [PuLID](https://github.com/ToTheBeginning/PuLID) — Liu et al., the actual
  identity-preservation magic.
- [ComfyUI-PuLID-Flux](https://github.com/balazik/ComfyUI-PuLID-Flux) — balazik,
  the ComfyUI port we depend on.
- [nateraw/vit-age-classifier](https://huggingface.co/nateraw/vit-age-classifier) —
  the age classifier we ship as default.
- [Flux](https://github.com/black-forest-labs/flux) — Black Forest Labs.