544 lines
19 KiB
Python
544 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = [
|
|
# "websocket-client>=1.6",
|
|
# "requests>=2.31",
|
|
# "pillow>=10.0",
|
|
# ]
|
|
# ///
|
|
"""
|
|
age_shift_cli.py — detect current age and generate target-age portraits
|
|
via ComfyUI + ComfyUI-AgeShift + ComfyUI-PuLID-Flux.
|
|
|
|
Usage:
|
|
# Detect only
|
|
uv run python age_shift_cli.py face.jpg --detect-only --ip 10.8.0.12
|
|
|
|
# Single target age (e.g. age this person up to 65)
|
|
uv run python age_shift_cli.py face.jpg --target-age 65 --ip 10.8.0.12
|
|
|
|
# Age sequence (5/15/25/45/65)
|
|
uv run python age_shift_cli.py face.jpg \\
|
|
--target-ages 5,15,25,45,65 \\
|
|
--gender male --out ./out --ip 10.8.0.12
|
|
|
|
# Dump workflow JSONs for debugging
|
|
uv run python age_shift_cli.py face.jpg --target-age 65 \\
|
|
--dump-workflow /tmp/wf.json --ip 10.8.0.12
|
|
"""
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List
|
|
|
|
import requests
|
|
from websocket import create_connection
|
|
|
|
|
|
# =============================================================================
|
|
# ComfyClient
|
|
# =============================================================================
|
|
class ComfyClient:
|
|
def __init__(self, server: str = "127.0.0.1:8188"):
|
|
self.server = server
|
|
self.client_id = str(uuid.uuid4())
|
|
self.http = f"http://{server}"
|
|
self.ws = f"ws://{server}/ws?clientId={self.client_id}"
|
|
|
|
def upload_image(self, file_path: Path, overwrite: bool = True) -> str:
|
|
with open(file_path, "rb") as f:
|
|
files = {"image": (file_path.name, f, "application/octet-stream")}
|
|
data = {"overwrite": "true" if overwrite else "false", "type": "input"}
|
|
r = requests.post(f"{self.http}/upload/image", files=files, data=data, timeout=60)
|
|
r.raise_for_status()
|
|
return r.json().get("name", file_path.name)
|
|
|
|
def submit(self, workflow: dict) -> str:
|
|
payload = {"prompt": workflow, "client_id": self.client_id}
|
|
r = requests.post(f"{self.http}/prompt", json=payload, timeout=30)
|
|
if r.status_code != 200:
|
|
raise RuntimeError(f"/prompt {r.status_code}: {r.text[:1000]}")
|
|
return r.json()["prompt_id"]
|
|
|
|
def history(self, prompt_id: str) -> dict:
|
|
r = requests.get(f"{self.http}/history/{prompt_id}", timeout=30)
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
return data.get(prompt_id, {})
|
|
|
|
def view(self, filename: str, subfolder: str = "", type_: str = "output") -> bytes:
|
|
params = {"filename": filename, "subfolder": subfolder, "type": type_}
|
|
r = requests.get(f"{self.http}/view", params=params, timeout=120)
|
|
r.raise_for_status()
|
|
return r.content
|
|
|
|
def run(self, workflow: dict, timeout_s: float = 600) -> dict:
|
|
ws = create_connection(self.ws, timeout=timeout_s)
|
|
try:
|
|
pid = self.submit(workflow)
|
|
print(f" prompt_id={pid}")
|
|
t0 = time.time()
|
|
last_node = None
|
|
while True:
|
|
if time.time() - t0 > timeout_s:
|
|
raise TimeoutError(f"Workflow timed out after {timeout_s}s")
|
|
msg = ws.recv()
|
|
if isinstance(msg, bytes):
|
|
continue
|
|
evt = json.loads(msg)
|
|
t = evt.get("type")
|
|
d = evt.get("data", {})
|
|
if d.get("prompt_id") and d["prompt_id"] != pid:
|
|
continue
|
|
if t == "progress":
|
|
cur, tot = d.get("value", 0), d.get("max", 1)
|
|
print(f"\r ▸ {cur:>4}/{tot} ", end="", flush=True)
|
|
elif t == "executing":
|
|
n = d.get("node")
|
|
if n is None:
|
|
print() # finalize line
|
|
break
|
|
if n != last_node:
|
|
print(f"\r ▸ node {str(n):<32}", end="", flush=True)
|
|
last_node = n
|
|
elif t == "execution_error":
|
|
raise RuntimeError(
|
|
f"Error in node {d.get('node_id')}: {d.get('exception_message')}"
|
|
)
|
|
return self.history(pid)
|
|
finally:
|
|
ws.close()
|
|
|
|
|
|
# =============================================================================
|
|
# Workflow builders (API format — node ids are strings)
|
|
# =============================================================================
|
|
def build_detect_workflow(
|
|
image_name: str,
|
|
*,
|
|
model_id: str = "nateraw/vit-age-classifier",
|
|
device: str = "auto",
|
|
dtype: str = "fp32",
|
|
) -> Dict[str, Any]:
|
|
return {
|
|
"load": {
|
|
"class_type": "LoadImage",
|
|
"inputs": {"image": image_name, "upload": "image"},
|
|
},
|
|
"loader": {
|
|
"class_type": "AgeDetectorLoader",
|
|
"inputs": {"model_id": model_id, "device": device, "dtype": dtype},
|
|
},
|
|
"detect": {
|
|
"class_type": "AgeDetect",
|
|
"inputs": {
|
|
"detector": ["loader", 0],
|
|
"image": ["load", 0],
|
|
"method": "weighted_mean",
|
|
},
|
|
},
|
|
"show_age": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["detect", 0]},
|
|
},
|
|
"show_band": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["detect", 1]},
|
|
},
|
|
"show_conf": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["detect", 2]},
|
|
},
|
|
"show_summary": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["detect", 3]},
|
|
},
|
|
}
|
|
|
|
|
|
def build_shift_workflow(
|
|
image_name: str,
|
|
target_age: int,
|
|
*,
|
|
# flux model file names (must be present in ComfyUI/models/*)
|
|
flux_unet: str,
|
|
flux_clip_l: str,
|
|
flux_t5: str,
|
|
flux_vae: str,
|
|
pulid_weights: str,
|
|
insightface_provider: str = "CUDA",
|
|
# generation knobs
|
|
width: int = 1024,
|
|
height: int = 1024,
|
|
steps: int = 20,
|
|
cfg: float = 1.0,
|
|
guidance: float = 3.5,
|
|
sampler: str = "euler",
|
|
scheduler: str = "simple",
|
|
seed: int = 0,
|
|
# pulid knobs
|
|
pulid_weight: float = 0.9,
|
|
pulid_start: float = 0.0,
|
|
pulid_end: float = 1.0,
|
|
# prompt knobs
|
|
subject: str = "person",
|
|
gender: str = "unspecified",
|
|
style: str = "photorealistic portrait, sharp focus, natural studio lighting",
|
|
extra_positive: str = "same identity, same facial structure, same eye shape",
|
|
negative_base: str = (
|
|
"cartoon, illustration, painting, anime, 3d render, "
|
|
"deformed, distorted, blurry, low quality, plastic skin, watermark, text"
|
|
),
|
|
) -> Dict[str, Any]:
|
|
"""Build a full PuLID-Flux + AgeShift workflow (API format).
|
|
|
|
Requires ComfyUI-PuLID-Flux installed alongside this package.
|
|
"""
|
|
return {
|
|
# Inputs
|
|
"load_face": {
|
|
"class_type": "LoadImage",
|
|
"inputs": {"image": image_name, "upload": "image"},
|
|
},
|
|
# Age-conditioned prompt
|
|
"prompt": {
|
|
"class_type": "AgePromptBuilder",
|
|
"inputs": {
|
|
"target_age": int(target_age),
|
|
"subject": subject,
|
|
"gender": gender,
|
|
"style": style,
|
|
"extra_positive": extra_positive,
|
|
"negative_base": negative_base,
|
|
"include_negative_age_drift": True,
|
|
},
|
|
},
|
|
# Flux backbone
|
|
"unet": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": flux_unet, "weight_dtype": "default"},
|
|
},
|
|
"dual_clip": {
|
|
"class_type": "DualCLIPLoader",
|
|
"inputs": {
|
|
"clip_name1": flux_clip_l,
|
|
"clip_name2": flux_t5,
|
|
"type": "flux",
|
|
},
|
|
},
|
|
"vae": {"class_type": "VAELoader", "inputs": {"vae_name": flux_vae}},
|
|
# PuLID-Flux identity injection
|
|
"pulid_model": {
|
|
"class_type": "PulidFluxModelLoader",
|
|
"inputs": {"pulid_file": pulid_weights},
|
|
},
|
|
"eva_clip": {"class_type": "PulidFluxEvaClipLoader", "inputs": {}},
|
|
"insightface": {
|
|
"class_type": "PulidFluxInsightFaceLoader",
|
|
"inputs": {"provider": insightface_provider},
|
|
},
|
|
"apply_pulid": {
|
|
"class_type": "ApplyPulidFlux",
|
|
"inputs": {
|
|
"model": ["unet", 0],
|
|
"pulid_flux": ["pulid_model", 0],
|
|
"eva_clip": ["eva_clip", 0],
|
|
"face_analysis": ["insightface", 0],
|
|
"image": ["load_face", 0],
|
|
"weight": pulid_weight,
|
|
"start_at": pulid_start,
|
|
"end_at": pulid_end,
|
|
"fusion": "mean",
|
|
"fusion_weight_max": 1.0,
|
|
"fusion_weight_min": 0.0,
|
|
"train_step": 1000,
|
|
"use_gray": True,
|
|
},
|
|
},
|
|
# Text → conditioning
|
|
"pos_encode": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"clip": ["dual_clip", 0], "text": ["prompt", 0]},
|
|
},
|
|
"neg_encode": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"clip": ["dual_clip", 0], "text": ["prompt", 1]},
|
|
},
|
|
"guidance_node": {
|
|
"class_type": "FluxGuidance",
|
|
"inputs": {"conditioning": ["pos_encode", 0], "guidance": guidance},
|
|
},
|
|
# Latent + sample
|
|
"latent": {
|
|
"class_type": "EmptySD3LatentImage",
|
|
"inputs": {"width": width, "height": height, "batch_size": 1},
|
|
},
|
|
"sampler_node": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["apply_pulid", 0],
|
|
"positive": ["guidance_node", 0],
|
|
"negative": ["neg_encode", 0],
|
|
"latent_image": ["latent", 0],
|
|
"seed": int(seed),
|
|
"steps": int(steps),
|
|
"cfg": float(cfg),
|
|
"sampler_name": sampler,
|
|
"scheduler": scheduler,
|
|
"denoise": 1.0,
|
|
},
|
|
},
|
|
"decode": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["sampler_node", 0], "vae": ["vae", 0]},
|
|
},
|
|
"save": {
|
|
"class_type": "SaveImage",
|
|
"inputs": {
|
|
"images": ["decode", 0],
|
|
"filename_prefix": f"ageshift_age{int(target_age):03d}",
|
|
},
|
|
},
|
|
"show_prompt_summary": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["prompt", 2]},
|
|
},
|
|
}
|
|
|
|
|
|
# =============================================================================
|
|
# History helpers
|
|
# =============================================================================
|
|
def extract_text(history: dict, node_id: str) -> str:
|
|
outs = history.get("outputs", {})
|
|
if node_id not in outs:
|
|
return ""
|
|
no = outs[node_id]
|
|
if "text" in no:
|
|
v = no["text"]
|
|
return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v)
|
|
parts = []
|
|
for v in no.values():
|
|
if isinstance(v, list):
|
|
parts.extend(str(x) for x in v)
|
|
return "\n".join(parts)
|
|
|
|
|
|
def extract_images(history: dict, node_id: str) -> List[Dict[str, str]]:
|
|
outs = history.get("outputs", {})
|
|
if node_id not in outs:
|
|
return []
|
|
return outs[node_id].get("images", [])
|
|
|
|
|
|
# =============================================================================
|
|
# Main
|
|
# =============================================================================
|
|
def main():
|
|
ap = argparse.ArgumentParser(
|
|
description="AgeShift: detect current age and generate target-age portraits "
|
|
"via ComfyUI + PuLID-Flux",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
ap.add_argument("image", type=Path, help="Source face image (jpg/png)")
|
|
|
|
# server
|
|
ap.add_argument("--ip", default="127.0.0.1")
|
|
ap.add_argument("--port", type=int, default=8188)
|
|
ap.add_argument("--server", default=None, help="full host:port (overrides --ip --port)")
|
|
|
|
# mode
|
|
ap.add_argument("--detect-only", action="store_true", help="only run age detection")
|
|
ap.add_argument("--target-age", type=int, default=None, help="single target age to generate")
|
|
ap.add_argument(
|
|
"--target-ages",
|
|
default=None,
|
|
help="csv list of target ages, e.g. 5,15,25,45,65",
|
|
)
|
|
|
|
# generation knobs
|
|
ap.add_argument("--subject", default="person")
|
|
ap.add_argument(
|
|
"--gender",
|
|
default="unspecified",
|
|
choices=["unspecified", "male", "female", "non-binary"],
|
|
)
|
|
ap.add_argument(
|
|
"--style",
|
|
default="photorealistic portrait, sharp focus, natural studio lighting, neutral background",
|
|
)
|
|
ap.add_argument(
|
|
"--extra-positive",
|
|
default="same identity, same facial structure, same eye shape",
|
|
)
|
|
ap.add_argument("--width", type=int, default=1024)
|
|
ap.add_argument("--height", type=int, default=1024)
|
|
ap.add_argument("--steps", type=int, default=20)
|
|
ap.add_argument("--cfg", type=float, default=1.0)
|
|
ap.add_argument("--guidance", type=float, default=3.5)
|
|
ap.add_argument("--sampler", default="euler")
|
|
ap.add_argument("--scheduler", default="simple")
|
|
ap.add_argument("--seed", type=int, default=0, help="0 = random per age")
|
|
ap.add_argument("--pulid-weight", type=float, default=0.9)
|
|
ap.add_argument("--pulid-start", type=float, default=0.0)
|
|
ap.add_argument("--pulid-end", type=float, default=1.0)
|
|
|
|
# model files (must already exist on the ComfyUI server)
|
|
ap.add_argument("--flux-unet", default="flux1-dev.safetensors")
|
|
ap.add_argument("--flux-clip-l", default="clip_l.safetensors")
|
|
ap.add_argument("--flux-t5", default="t5xxl_fp8_e4m3fn.safetensors")
|
|
ap.add_argument("--flux-vae", default="ae.safetensors")
|
|
ap.add_argument("--pulid-weights", default="pulid_flux_v0.9.1.safetensors")
|
|
ap.add_argument(
|
|
"--insightface-provider",
|
|
default="CUDA",
|
|
choices=["CUDA", "CPU", "ROCM"],
|
|
)
|
|
|
|
# detector
|
|
ap.add_argument("--detector-model", default="nateraw/vit-age-classifier")
|
|
ap.add_argument(
|
|
"--detector-device",
|
|
default="auto",
|
|
choices=["auto", "cuda", "cuda:0", "cuda:1", "cpu"],
|
|
)
|
|
ap.add_argument("--detector-dtype", default="fp32", choices=["fp32", "fp16", "bf16"])
|
|
|
|
# output / debug
|
|
ap.add_argument("--out", type=Path, default=Path("./out"))
|
|
ap.add_argument(
|
|
"--dump-workflow",
|
|
type=Path,
|
|
default=None,
|
|
help="dump generated workflow JSON(s) for debugging",
|
|
)
|
|
ap.add_argument("--timeout", type=float, default=600.0)
|
|
|
|
args = ap.parse_args()
|
|
server = args.server or f"{args.ip}:{args.port}"
|
|
args.out.mkdir(parents=True, exist_ok=True)
|
|
|
|
if not args.image.exists():
|
|
print(f"ERROR: image not found: {args.image}", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
client = ComfyClient(server)
|
|
|
|
# ---- upload ----
|
|
print(f"[1/3] Uploading {args.image.name} → {server} ...")
|
|
server_name = client.upload_image(args.image)
|
|
print(f" uploaded as: {server_name}")
|
|
|
|
# ---- detect ----
|
|
print(f"[2/3] Detecting age ...")
|
|
detect_wf = build_detect_workflow(
|
|
server_name,
|
|
model_id=args.detector_model,
|
|
device=args.detector_device,
|
|
dtype=args.detector_dtype,
|
|
)
|
|
if args.dump_workflow:
|
|
dump = args.dump_workflow.with_suffix(".detect.json")
|
|
dump.parent.mkdir(parents=True, exist_ok=True)
|
|
dump.write_text(json.dumps(detect_wf, indent=2, ensure_ascii=False))
|
|
print(f" wrote {dump}")
|
|
hist = client.run(detect_wf, timeout_s=args.timeout)
|
|
summary = extract_text(hist, "show_summary") or "(no summary)"
|
|
age_str = extract_text(hist, "show_age").strip()
|
|
try:
|
|
detected_age = int(float(age_str))
|
|
except (ValueError, TypeError):
|
|
detected_age = None
|
|
print(f" → {summary}")
|
|
|
|
if args.detect_only:
|
|
print("Done (detect only).")
|
|
return
|
|
|
|
# ---- shift ----
|
|
targets: List[int] = []
|
|
if args.target_age is not None:
|
|
targets.append(args.target_age)
|
|
if args.target_ages:
|
|
for s in args.target_ages.split(","):
|
|
s = s.strip()
|
|
if s:
|
|
try:
|
|
targets.append(int(s))
|
|
except ValueError:
|
|
print(f" ! skip invalid age: {s!r}", file=sys.stderr)
|
|
if not targets:
|
|
print(
|
|
"No --target-age / --target-ages given; nothing to generate.",
|
|
file=sys.stderr,
|
|
)
|
|
sys.exit(0)
|
|
# dedup, keep order
|
|
seen = set()
|
|
targets = [a for a in targets if not (a in seen or seen.add(a))]
|
|
|
|
print(f"[3/3] Generating {len(targets)} age-shifted image(s) ...")
|
|
for i, age in enumerate(targets, 1):
|
|
print(f" ({i}/{len(targets)}) target age = {age}")
|
|
seed = args.seed if args.seed else int(time.time() * 1000) % (2**31 - 1) + i
|
|
wf = build_shift_workflow(
|
|
server_name,
|
|
age,
|
|
flux_unet=args.flux_unet,
|
|
flux_clip_l=args.flux_clip_l,
|
|
flux_t5=args.flux_t5,
|
|
flux_vae=args.flux_vae,
|
|
pulid_weights=args.pulid_weights,
|
|
insightface_provider=args.insightface_provider,
|
|
width=args.width,
|
|
height=args.height,
|
|
steps=args.steps,
|
|
cfg=args.cfg,
|
|
guidance=args.guidance,
|
|
sampler=args.sampler,
|
|
scheduler=args.scheduler,
|
|
seed=seed,
|
|
pulid_weight=args.pulid_weight,
|
|
pulid_start=args.pulid_start,
|
|
pulid_end=args.pulid_end,
|
|
subject=args.subject,
|
|
gender=args.gender,
|
|
style=args.style,
|
|
extra_positive=args.extra_positive,
|
|
)
|
|
if args.dump_workflow:
|
|
dump = args.dump_workflow.with_suffix(f".shift_age{age}.json")
|
|
dump.write_text(json.dumps(wf, indent=2, ensure_ascii=False))
|
|
print(f" wrote {dump}")
|
|
try:
|
|
hist = client.run(wf, timeout_s=args.timeout)
|
|
except Exception as e:
|
|
print(f" ! failed: {e}", file=sys.stderr)
|
|
continue
|
|
prompt_summary = extract_text(hist, "show_prompt_summary").strip()
|
|
if prompt_summary:
|
|
print(f" prompt: {prompt_summary}")
|
|
imgs = extract_images(hist, "save")
|
|
if not imgs:
|
|
print(f" ! no output images for age {age}", file=sys.stderr)
|
|
continue
|
|
for img in imgs:
|
|
data = client.view(
|
|
img["filename"],
|
|
img.get("subfolder", ""),
|
|
img.get("type", "output"),
|
|
)
|
|
out_path = args.out / f"age_{age:03d}_{img['filename']}"
|
|
out_path.write_bytes(data)
|
|
print(f" → {out_path} ({len(data)//1024} KB)")
|
|
|
|
print("All done.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|