Files
pi-imgen/extensions/imgen/commands/retouch.ts
T
mozempk aee92d284c cleanup: clear stale TODOs, tighten types, translate art-prompt hints
Every 'being written in parallel' / 'ASSUMED SIGNATURE' TODO is gone: each
referenced module now exists and each assumption was checked against it
rather than the comment simply deleted.

- logo/presets/retouch: renderer typed as the real Renderer contract instead
  of unknown / (...args: never[]).
- director: the ModelRegistry assumption is VERIFIED against the installed
  @earendil-works types (find() on ModelRegistry, getModel() on ModelRuntime).
- strings: the img2imgBug TODO had it backwards -- drawthings raises a graded
  message a static string cannot express, and this is the fallback.
- poster: an Italian hint was being spliced into an English art prompt, which
  degrades these models. New director.refineArtPrompt() folds the note in via
  the LLM and falls back to the old splice on any failure, so it can only
  improve on the previous behaviour. scrubArtPrompt still runs either way.

Remaining TODOs are one category only (Italian strings living in per-file
tables rather than ui/strings.ts) and are accurate, not stale.
2026-08-27 10:30:54 +02:00

960 lines
37 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
/**
* `/retouch` — work on an image he already has.
*
* Five things, in Italian, in this order of trustworthiness:
*
* 1. rimuovi lo sfondo -> cpu.removeBackground() (Apple Vision / rembg)
* 2. ingrandisci -> cpu.upscaleRealesrgan() (ncnn/Vulkan)
* 3. ritaglia per i social -> contrast.smartCrop() (sharp)
* 4. usa come sfondo di una locandina -> hands off to the poster flow, art.source="photo"
* 5. modifica con l'AI -> backend.generate({ initImage }) <- the fragile one
*
* 1–3 are CPU work: no diffusion model, no resident gRPC server, no Draw Things. They are
* deliberately reachable without ever touching `deps.backend`, because Draw Things upstream
* issue #121 crashes img2img (EXC_BREAKPOINT) on 16GB Macs — which is exactly this machine.
* When option 5 dies that way we say so in Italian, name it as a known bug of the program
* rather than something he did wrong, and repeat that 1–4 still work.
*
* The command is thin on purpose: it resolves an image and an action, then calls
* `runRetouch()`. The LLM-facing tool calls the same `runRetouch()`, so there is exactly
* one implementation of every operation.
*/
import { stat } from "node:fs/promises";
import { basename, extname, isAbsolute, join, resolve } from "node:path";
import { homedir } from "node:os";
import sharp from "sharp";
import { Type } from "typebox";
import { StringEnum } from "@earendil-works/pi-ai";
import type { ImageContent, TextContent } from "@earendil-works/pi-ai";
import { defineTool, withFileMutationQueue } from "@earendil-works/pi-coding-agent";
import type {
ExtensionAPI,
ExtensionCommandContext,
ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { ImgenConfig } from "../config.ts";
import { BackendError, type Backend, type ProgressFn } from "../backends/types.ts";
import type {
CpuOpResult,
RemoveBackgroundOptions,
UpscaleOptions,
} from "../backends/cpu.ts";
import { FORMATS, type Format } from "../design/spec.ts";
import { ART_GENERATION_MAX_PX, FORMATS_GEOMETRY, roundToMultiple } from "../render/formats.ts";
import type { Brief, DirectResult, DirectorContext } from "../design/director.ts";
import { createJob, openInFinder, slugify } from "../job.ts";
import { S, bullets, errorText, fill, menu } from "../ui/strings.ts";
import type { Renderer } from "./poster.ts";
// ---------------------------------------------------------------------------
// Italian strings local to /retouch
//
// TODO: move to `S.retouch` in ui/strings.ts (which currently only carries
// `S.commands.retouch`). These live here until it grows one; every user-visible
// string in this file comes from here or from `S`.
// ---------------------------------------------------------------------------
export const R = {
askImage: "Quale immagine?",
askImagePlaceholder: "Trascina qui il file, oppure incolla il percorso",
askAction: "Cosa ci faccio?",
actions: {
background: "Togli lo sfondo",
upscale: "Ingrandisci",
crop: "Ritaglia per i social",
poster: "Usa come sfondo di una locandina",
edit: "Modifica con l'AI",
},
order: ["background", "upscale", "crop", "poster", "edit"] as const,
hints: {
background: "Ritaglio il soggetto e ti lascio un PNG trasparente.",
upscale: "Rifaccio l'immagine più grande e più nitida, senza inventare niente.",
crop: "Taglio la stessa immagine nei formati di Instagram, Facebook e YouTube.",
poster: "Passo alla locandina vera e propria, usando questa foto al posto del disegno.",
edit: "Chiedo al modello di ridisegnare partendo da questa immagine.",
},
askFactor: "Quanto la ingrandisco?",
factors: { x2: "Il doppio", x4: "Quattro volte" },
askFormats: "Per dove?",
allFormats: "Tutti quanti",
askEdit: "Cosa cambio nell'immagine?",
askEditPlaceholder: "di sera invece che di giorno, colori più caldi, senza persone",
askCropSubject: "Ritaglio anche intorno al soggetto?",
askCropSubjectHint: "Utile per un logo, sbagliato se poi devi sovrapporlo a qualcos'altro.",
needEdit: "Non mi hai detto cosa cambiare, quindi non tocco niente.",
notAnImage: "Questo non sembra un file di immagine.",
working: "Sto lavorando su «{file}»…",
ready: "Fatto.",
thinkingEdit: "Sto capendo cosa vuoi cambiare…",
notifyDone: "Immagine pronta.",
producedOne: "Ho salvato un file:",
producedMany: "Ho salvato {n} file:",
nothingProduced: "Non è uscito nessun file.",
usage: [
"Posso lavorare su un'immagine che hai già. Scrivimi il percorso del file e cosa vuoi:",
"",
" /retouch togli lo sfondo a ~/Immagini/logo.png",
" /retouch ingrandisci ~/Immagini/foto.jpg",
" /retouch ritaglia per i social ~/Immagini/foto.jpg",
" /retouch usa ~/Immagini/foto.jpg come sfondo di una locandina",
" /retouch su ~/Immagini/foto.jpg: rendila notturna",
].join("\n"),
unknownAction: "Non ho capito cosa vuoi che faccia. Ecco cosa so fare su un'immagine:",
handoffPoster: "Va bene: passo alla locandina e uso «{file}» come immagine di sfondo.",
handoffManual:
"Uso «{file}» come sfondo. Se non parte da solo, scrivi «/poster» e dimmi di usare questa foto.",
stillWorks: "Le altre cose che so fare su questa immagine funzionano lo stesso:",
} as const;
// ---------------------------------------------------------------------------
// Dependencies — index.ts owns construction, this file owns orchestration only
// ---------------------------------------------------------------------------
/** The two CPU entry points `/retouch` needs. Structural, so tests can pass fakes. */
export interface RetouchCpuOps {
removeBackground(input: string, out: string, opts?: RemoveBackgroundOptions): Promise<CpuOpResult>;
upscaleRealesrgan(input: string, out: string, factor: number, opts?: UpscaleOptions): Promise<CpuOpResult>;
}
/**
* The crop slice of `render/contrast.ts`. Matches its exported `smartCrop`, which writes
* a saliency-aware crop (sharp's attention strategy) and returns the size read back from
* the written file rather than the size that was requested.
*/
export interface RetouchCropper {
smartCrop(
input: string,
out: string,
size: { width: number; height: number },
opts?: { signal?: AbortSignal; onProgress?: ProgressFn },
): Promise<{ path: string; width: number; height: number }>;
}
/** Only the one director entry point this command uses. */
export interface RetouchDirector {
direct(brief: Brief, cfg: ImgenConfig, ctx: DirectorContext): Promise<DirectResult>;
}
/**
* `/retouch` never typesets anything itself; the renderer is carried here only so
* `startPoster` can be built with it in index.ts and the extension is wired from one
* place. Aliased to the real contract so a signature change cannot drift unnoticed.
*/
export type RetouchRenderer = Renderer;
/** Hand-off into the poster flow with an existing photo as the artwork. */
export type PosterHandoff = (
opts: { photoPath: string; freeText?: string },
ctx: ExtensionCommandContext,
) => Promise<void>;
export interface RetouchDeps {
config: ImgenConfig;
backend: Backend;
cpu: RetouchCpuOps;
crop: RetouchCropper;
/** Turns his Italian instruction into an English art prompt for the generative edit. */
director?: RetouchDirector;
/** Unused by retouch itself; see RetouchRenderer. */
renderer?: RetouchRenderer;
/** When absent, the hand-off degrades to typing `/poster` for him. */
startPoster?: PosterHandoff;
}
// ---------------------------------------------------------------------------
// Public shapes
// ---------------------------------------------------------------------------
export const RETOUCH_ACTIONS = ["background", "upscale", "crop", "poster", "edit"] as const;
export type RetouchAction = (typeof RETOUCH_ACTIONS)[number];
export interface RetouchRequest {
action: RetouchAction;
/** Absolute path to the source image. */
input: string;
/** Defaults to a job folder under `config.outputDir`. */
outDir?: string;
/** upscale only. 2 or 4; anything in (1, 8] is accepted. */
factor?: number;
/** crop only. Defaults to every screen format. */
formats?: Format[];
/** background only: also crop to the subject's bounding box. */
cropToSubject?: boolean;
/** edit / poster: what he asked for, in Italian. */
instruction?: string;
/** edit only. How far from the original, 0..1. */
strength?: number;
}
export interface RetouchOutput {
path: string;
/** Italian, user-visible. */
label: string;
width: number;
height: number;
}
export interface RetouchResult {
action: RetouchAction;
input: string;
dir: string;
outputs: RetouchOutput[];
/** Italian notes about anything we degraded silently. */
notes: string[];
tool: string;
elapsedMs: number;
/** Set for action "poster": nothing was produced, the poster flow takes over. */
handoff?: { kind: "poster"; photoPath: string; freeText?: string };
}
/** Every failure leaving this module carries an Italian `message`. */
export class RetouchError extends Error {
readonly italian: string;
readonly detail?: string;
constructor(italian: string, detail?: string) {
super(italian);
this.name = "RetouchError";
this.italian = italian;
this.detail = detail;
}
}
// ---------------------------------------------------------------------------
// Image paths — he drags files into the terminal, so be generous
// ---------------------------------------------------------------------------
const IMAGE_EXTENSIONS = [
".png", ".jpg", ".jpeg", ".heic", ".heif", ".webp", ".tif", ".tiff", ".bmp", ".gif",
];
export function looksLikeImage(path: string): boolean {
return IMAGE_EXTENSIONS.includes(extname(path).toLowerCase());
}
/**
* Normalises whatever the terminal handed us: surrounding quotes, `file://` URLs,
* percent-escapes, backslash-escaped spaces (how macOS pastes a dragged file), `~`.
*/
export function normaliseImagePath(raw: string, cwd: string): string {
let p = raw.trim();
if (p.length >= 2 && ((p.startsWith('"') && p.endsWith('"')) || (p.startsWith("'") && p.endsWith("'")))) {
p = p.slice(1, -1);
}
if (p.startsWith("file://")) {
p = p.slice("file://".length);
try {
p = decodeURIComponent(p);
} catch {
/* not percent-encoded after all */
}
}
p = p.replace(/\\(.)/g, "$1").trim();
if (p === "~") p = homedir();
else if (p.startsWith("~/")) p = join(homedir(), p.slice(2));
return isAbsolute(p) ? p : resolve(cwd, p);
}
/**
* A single token -> the image path it denotes, or undefined. Sentence punctuation glued to
* the end ("su ~/foto.jpg: rendila notturna") is shaved off before giving up.
*/
export function imageCandidate(token: string, cwd: string): string | undefined {
const direct = normaliseImagePath(token, cwd);
if (looksLikeImage(direct)) return direct;
const trimmed = token.replace(/[),;:.!?]+$/, "");
if (trimmed === token) return undefined;
const stripped = normaliseImagePath(trimmed, cwd);
return looksLikeImage(stripped) ? stripped : undefined;
}
/**
* Tokenises free text the way a dragged file survives it: quoted strings stay whole, and a
* backslash-escaped space (what macOS inserts when you drop a file into a terminal) does
* NOT split a path in two.
*/
const TOKEN_RE = /(["'])(?:\\.|(?!\1).)+\1|(?:\\.|\S)+/g;
/** Pulls the first thing that looks like an image path out of free text. */
export function findImagePathInText(text: string, cwd: string): string | undefined {
for (const token of text.match(TOKEN_RE) ?? []) {
const p = imageCandidate(token, cwd);
if (p) return p;
}
return undefined;
}
async function assertReadableImage(path: string): Promise<void> {
let info;
try {
info = await stat(path);
} catch {
throw new RetouchError(errorText(S.errors.photoMissing(path)), path);
}
if (!info.isFile() || info.size === 0) {
throw new RetouchError(errorText(S.errors.photoMissing(path)), path);
}
if (!looksLikeImage(path)) {
throw new RetouchError(`${R.notAnImage} (${basename(path)})`, path);
}
}
// ---------------------------------------------------------------------------
// Free-text understanding — offline, deterministic, Italian
// ---------------------------------------------------------------------------
/**
* Order matters: "usa come SFONDO di una locandina" and "togli lo SFONDO" share a word, so
* the poster rule has to be consulted before the background rule.
*/
const ACTION_RULES: readonly { action: RetouchAction; re: RegExp }[] = [
{ action: "poster", re: /locandin|manifest|poster|volantin|sfondo\s+(?:di|per|della|di una)\s+(?:una\s+)?locandin/i },
{ action: "background", re: /sfond|scontorn|trasparen|ritagli\w*\s+il\s+soggett|soggetto\s+ritagli|rimuov\w*\s+lo\s+sfond|png\s+trasparent/i },
{ action: "crop", re: /ritagl|taglia|social|instagram|storia|storie|stories|facebook|copertin|youtube|miniatur|anteprim|formato|quadrat/i },
{ action: "upscale", re: /ingrand|upscal|nitid|risoluzione|più\s+grande|piu\s+grande|sgranat|pixel|allarg|zoom/i },
{ action: "edit", re: /modific|cambia|ridisegn|trasform|rendila|rendilo|generativ|con\s+l'?ai|intelligenza\s+artificiale|reimmagin|variante/i },
];
const FORMAT_RULES: readonly { format: Format; re: RegExp }[] = [
{ format: "ig-story", re: /storia|storie|stories|verticale\s+intero|9\s*[:x]\s*16/i },
{ format: "ig-post", re: /instagram|\bpost\b|\big\b|4\s*[:x]\s*5/i },
{ format: "fb-cover", re: /facebook|\bfb\b|copertin/i },
{ format: "yt-thumb", re: /youtube|\byt\b|miniatur|anteprim|thumb/i },
];
export interface ParsedRetouchArgs {
action?: RetouchAction;
input?: string;
formats?: Format[];
factor?: number;
cropToSubject?: boolean;
/** What is left once the path is removed — the instruction for edit / poster. */
rest: string;
}
export function parseRetouchArgs(args: string, cwd: string): ParsedRetouchArgs {
const text = args.trim();
if (!text) return { rest: "" };
const input = findImagePathInText(text, cwd);
const rest = (input
? text.replace(TOKEN_RE, (tok) => (imageCandidate(tok, cwd) === input ? "" : tok))
: text)
.replace(/\s+/g, " ")
.replace(/^[\s,;:.-]+|[\s,;:.-]+$/g, "")
.replace(/^(?:su|a|ad|di|dal|dalla|per|con|questa|questo|l'immagine|la foto|il file)\b[\s,;:.-]*/gi, "")
.replace(/^[\s,;:.-]+/, "")
.trim();
const action = ACTION_RULES.find((r) => r.re.test(text))?.action;
const formats = FORMAT_RULES.filter((r) => r.re.test(text)).map((r) => r.format);
const factorMatch = text.match(/(?:\b([234])\s*x\b|\bx\s*([234])\b|\b(doppio|quadrupl\w*)\b)/i);
let factor: number | undefined;
if (factorMatch) {
if (factorMatch[1] || factorMatch[2]) factor = Number(factorMatch[1] ?? factorMatch[2]);
else factor = /doppio/i.test(factorMatch[3] ?? "") ? 2 : 4;
}
const cropToSubject = /ritagli\w*\s+(?:anche\s+)?(?:intorno|attorno|al soggett|il soggett)|\bbounding\b|stretto\s+sul\s+soggett/i.test(text)
? true
: undefined;
return {
...(action ? { action } : {}),
...(input ? { input } : {}),
...(formats.length ? { formats } : {}),
...(factor ? { factor } : {}),
...(cropToSubject !== undefined ? { cropToSubject } : {}),
rest,
};
}
// ---------------------------------------------------------------------------
// Draw Things issue #121
// ---------------------------------------------------------------------------
/**
* True for the one failure this command has to narrate correctly: img2img crashing on a
* 16GB Mac. drawthings.ts already builds the Italian sentence; we only add the reminder
* that nothing else here is affected.
*/
export function isIssue121(e: unknown): e is BackendError {
if (!(e instanceof BackendError) && (e as { name?: string })?.name !== "BackendError") return false;
const err = e as BackendError;
return /issue\s*#?121/i.test(`${err.message} ${err.italian ?? ""}`);
}
function issue121Message(e: BackendError): string {
const opening = e.italian?.trim() || errorText(S.errors.img2imgBug);
return [
opening,
"",
R.stillWorks,
bullets([R.actions.background, R.actions.upscale, R.actions.crop, R.actions.poster]),
].join("\n");
}
// ---------------------------------------------------------------------------
// The one implementation, shared by the command and the tool
// ---------------------------------------------------------------------------
const SCREEN_FORMATS: readonly Format[] = ["ig-post", "ig-story", "fb-cover", "yt-thumb"];
const DEFAULT_UPSCALE_FACTOR = 2;
const DEFAULT_EDIT_STRENGTH = 0.55;
export interface RunRetouchOptions {
onProgress?: ProgressFn;
signal?: AbortSignal;
/** Passed through to the director for the generative edit. */
directorContext?: DirectorContext;
}
export async function runRetouch(
deps: RetouchDeps,
req: RetouchRequest,
opts: RunRetouchOptions = {},
): Promise<RetouchResult> {
const started = Date.now();
const input = normaliseImagePath(req.input, process.cwd());
await assertReadableImage(input);
const base = basename(input, extname(input));
if (req.action === "poster") {
// Nothing to produce here: the poster flow owns the artwork from now on.
return {
action: "poster",
input,
dir: req.outDir ?? deps.config.outputDir,
outputs: [],
notes: [],
tool: "handoff",
elapsedMs: Date.now() - started,
handoff: {
kind: "poster",
photoPath: input,
...(req.instruction?.trim() ? { freeText: req.instruction.trim() } : {}),
},
};
}
const dir = req.outDir ?? createJob(deps.config, slugify(`ritocco ${base}`), { reuse: true }).dir;
const notes: string[] = [];
const outputs: RetouchOutput[] = [];
let tool = "";
switch (req.action) {
case "background": {
const out = join(dir, `${base}-senza-sfondo.png`);
opts.onProgress?.(S.progress.removingBackground);
const res = await withFileMutationQueue(out, () =>
deps.cpu.removeBackground(input, out, {
...(req.cropToSubject !== undefined ? { crop: req.cropToSubject } : {}),
...(opts.signal ? { signal: opts.signal } : {}),
...(opts.onProgress ? { onProgress: opts.onProgress } : {}),
config: deps.config,
}),
).catch(rethrowItalian(S.errors.backgroundRemovalFailed));
tool = res.tool;
notes.push(...res.notes);
outputs.push({ path: res.path, label: "Immagine senza sfondo (PNG trasparente)", width: res.width, height: res.height });
break;
}
case "upscale": {
const factor = clampFactor(req.factor ?? DEFAULT_UPSCALE_FACTOR);
const out = join(dir, `${base}-x${factor}.png`);
opts.onProgress?.(S.progress.upscaling);
const res = await withFileMutationQueue(out, () =>
upscaleWithFallback(deps, input, out, factor, notes, opts),
);
tool = res.tool;
outputs.push({ path: res.path, label: fill("Immagine ingrandita {n} volte", { n: factor }), width: res.width, height: res.height });
break;
}
case "crop": {
const formats = (req.formats?.length ? req.formats : SCREEN_FORMATS).filter((f) => FORMATS.includes(f));
if (!formats.length) throw new RetouchError(errorText(S.errors.exportFailed(R.actions.crop)));
opts.onProgress?.(S.progress.cropping);
tool = "smart-crop";
for (const format of formats) {
const geo = FORMATS_GEOMETRY[format];
const out = join(dir, `${base}-${format}.png`);
const size = { width: geo.pxWidth, height: geo.pxHeight };
try {
const res = await withFileMutationQueue(out, () =>
deps.crop.smartCrop(input, out, size, {
...(opts.signal ? { signal: opts.signal } : {}),
...(opts.onProgress ? { onProgress: opts.onProgress } : {}),
}),
);
outputs.push({ path: res.path, label: S.formats.labels[format], width: res.width, height: res.height });
} catch (e) {
// One bad format must not cost him the other three.
notes.push(
e instanceof BackendError && e.italian
? e.italian
: errorText(S.errors.exportFailed(S.formats.labels[format])),
);
}
}
if (!outputs.length) throw new RetouchError(errorText(S.errors.exportFailed(R.actions.crop)));
const src = await imageSize(input);
if (src) {
const tooSmall = formats.some((f) => FORMATS_GEOMETRY[f].pxWidth > src.width * 1.5);
if (tooSmall) notes.push(S.warnings.photoLowRes);
notes.push(S.warnings.photoCropped);
}
break;
}
case "edit": {
const instruction = req.instruction?.trim();
if (!instruction) throw new RetouchError(R.needEdit);
const out = join(dir, `${base}-modificata.png`);
const res = await withFileMutationQueue(out, () =>
generativeEdit(deps, { input, out, instruction, ...(req.strength !== undefined ? { strength: req.strength } : {}) }, notes, opts),
);
tool = res.model;
outputs.push({ path: res.path, label: "Immagine modificata", width: res.width, height: res.height });
break;
}
}
return { action: req.action, input, dir, outputs, notes, tool, elapsedMs: Date.now() - started };
}
function clampFactor(n: number): number {
if (!Number.isFinite(n) || n <= 1) return DEFAULT_UPSCALE_FACTOR;
return Math.min(8, Math.round(n));
}
function rethrowItalian(msg: { message: string; fix?: string }) {
return (e: unknown): never => {
if (e instanceof RetouchError) throw e;
const detail = e instanceof Error ? e.message : String(e);
const italian = e instanceof BackendError && e.italian ? e.italian : errorText(msg);
throw new RetouchError(italian, detail);
};
}
async function imageSize(path: string): Promise<{ width: number; height: number } | undefined> {
try {
const m = await sharp(path).metadata();
if (m.width && m.height) return { width: m.width, height: m.height };
} catch {
/* unreadable by sharp — the caller already knows the file exists */
}
return undefined;
}
/**
* Real-ESRGAN first, because it is the reliable path and needs no diffusion model at all.
* Only if the ncnn binary is missing do we fall back to the backend, which degrades to a
* Lanczos resample of its own accord rather than failing.
*/
async function upscaleWithFallback(
deps: RetouchDeps,
input: string,
out: string,
factor: number,
notes: string[],
opts: RunRetouchOptions,
): Promise<{ path: string; width: number; height: number; tool: string }> {
try {
const res = await deps.cpu.upscaleRealesrgan(input, out, factor, {
...(opts.signal ? { signal: opts.signal } : {}),
...(opts.onProgress ? { onProgress: opts.onProgress } : {}),
config: deps.config,
});
notes.push(...res.notes);
return { path: res.path, width: res.width, height: res.height, tool: res.tool };
} catch (e) {
if (opts.signal?.aborted) throw e;
if (e instanceof BackendError && e.italian) notes.push(e.italian);
opts.onProgress?.(S.progress.upscaling);
try {
const res = await deps.backend.upscale(input, out, factor, opts.onProgress, opts.signal);
return { path: res.path, width: res.width, height: res.height, tool: res.model };
} catch (e2) {
if (isIssue121(e2)) throw new RetouchError(issue121Message(e2), e2.detail);
throw rethrowItalian(S.errors.upscaleFailed)(e2);
}
}
}
/**
* The generative edit. The prompt handed to the model must be ENGLISH and must never ask
* for lettering, so his Italian instruction goes through the art director first; without a
* director we still refuse to invent text and fall back to the instruction plus the
* mandatory negative.
*/
async function generativeEdit(
deps: RetouchDeps,
args: { input: string; out: string; instruction: string; strength?: number },
notes: string[],
opts: RunRetouchOptions,
): Promise<{ path: string; width: number; height: number; model: string }> {
const src = (await imageSize(args.input)) ?? { width: ART_GENERATION_MAX_PX, height: ART_GENERATION_MAX_PX };
const scale = Math.min(1, ART_GENERATION_MAX_PX / Math.max(src.width, src.height));
const width = Math.max(256, roundToMultiple(Math.round(src.width * scale)));
const height = Math.max(256, roundToMultiple(Math.round(src.height * scale)));
let prompt = args.instruction;
let negative = "text, letters, words, typography, watermark, signature";
let seed = Math.floor(Math.random() * 2_147_483_647);
if (deps.director) {
opts.onProgress?.(R.thinkingEdit);
try {
const brief: Brief = {
kind: "poster",
title: basename(args.input, extname(args.input)),
freeText: args.instruction,
photoPath: args.input,
};
const directed = await deps.director.direct(brief, deps.config, {
...(opts.directorContext ?? {}),
...(opts.signal ? { signal: opts.signal } : {}),
});
prompt = directed.spec.art.prompt;
negative = directed.spec.art.negative;
seed = directed.spec.art.seed;
} catch (e) {
// Not being able to translate the instruction is not a reason to refuse the edit.
notes.push(errorText(S.errors.directorFailed(e instanceof Error ? e.message : undefined)));
}
}
opts.onProgress?.(S.progress.generating);
try {
const res = await deps.backend.generate(
{
prompt,
negative,
seed,
width,
height,
steps: 24,
tier: "final",
initImage: args.input,
strength: args.strength ?? DEFAULT_EDIT_STRENGTH,
outPath: args.out,
},
opts.onProgress,
opts.signal,
);
return { path: res.path, width: res.width, height: res.height, model: res.model };
} catch (e) {
if (opts.signal?.aborted) throw e;
if (isIssue121(e)) throw new RetouchError(issue121Message(e), e.detail);
throw rethrowItalian(S.errors.generationFailed())(e);
}
}
// ---------------------------------------------------------------------------
// Presentation
// ---------------------------------------------------------------------------
const PREVIEW_MAX_PX = 1200;
/** A downscaled PNG copy so the terminal can render it inline without hauling 20MB around. */
export async function previewImage(path: string): Promise<ImageContent | undefined> {
try {
const data = await sharp(path)
.resize({ width: PREVIEW_MAX_PX, height: PREVIEW_MAX_PX, fit: "inside", withoutEnlargement: true })
.png()
.toBuffer();
return { type: "image", data: data.toString("base64"), mimeType: "image/png" };
} catch {
return undefined;
}
}
export function describeResult(res: RetouchResult): string {
if (res.handoff) return fill(R.handoffPoster, { file: basename(res.handoff.photoPath) });
if (!res.outputs.length) return R.nothingProduced;
const header = res.outputs.length === 1 ? R.producedOne : fill(R.producedMany, { n: res.outputs.length });
const lines = res.outputs.map((o) => `${o.label} — ${basename(o.path)} (${o.width}×${o.height})`);
const parts = [R.ready, "", header, bullets(lines), "", fill(S.done.folder, { cartella: res.dir })];
if (res.notes.length) parts.push("", bullets(res.notes));
return parts.join("\n");
}
// ---------------------------------------------------------------------------
// The command
// ---------------------------------------------------------------------------
const STATUS_KEY = "imgen-retouch";
const MESSAGE_TYPE = "imgen-retouch";
const ENTRY_TYPE = "imgen-retouch";
/** Dialogs exist in TUI and RPC; `ctx.ui.custom()` would additionally need mode === "tui". */
function canPrompt(ctx: ExtensionContext): boolean {
return ctx.hasUI && (ctx.mode === "tui" || ctx.mode === "rpc");
}
function say(pi: ExtensionAPI, content: string | (TextContent | ImageContent)[]): void {
pi.sendMessage({ customType: MESSAGE_TYPE, content, display: true });
}
export function registerRetouch(pi: ExtensionAPI, deps: RetouchDeps): void {
pi.registerCommand("retouch", {
description: S.commands.retouch.description,
handler: (args, ctx) => handleRetouch(pi, deps, args, ctx),
});
pi.registerTool(createRetouchTool(deps));
}
async function handleRetouch(
pi: ExtensionAPI,
deps: RetouchDeps,
args: string,
ctx: ExtensionCommandContext,
): Promise<void> {
const parsed = parseRetouchArgs(args, ctx.cwd);
const interactive = canPrompt(ctx);
try {
// -- 1. which image ----------------------------------------------------
let input = parsed.input;
if (!input) {
if (!interactive) {
say(pi, `${R.usage}\n\n${S.commands.retouch.usage}`);
return;
}
const typed = await ctx.ui.input(R.askImage, R.askImagePlaceholder);
if (typed === undefined || !typed.trim()) {
say(pi, S.common.cancelled);
return;
}
input = normaliseImagePath(typed, ctx.cwd);
}
await assertReadableImage(input);
// -- 2. which action ---------------------------------------------------
let action = parsed.action;
if (!action) {
if (!interactive) {
say(pi, [R.unknownAction, bullets(R.order.map((k) => R.actions[k])), "", R.usage].join("\n"));
return;
}
const m = menu(R.actions, R.order);
action = m.pick(await ctx.ui.select(R.askAction, m.options));
if (!action) {
say(pi, S.common.cancelled);
return;
}
}
// -- 3. the few extra questions each action needs -----------------------
const req: RetouchRequest = { action, input };
if (action === "upscale") {
let factor = parsed.factor;
if (factor === undefined && interactive) {
const m = menu(R.factors, ["x2", "x4"] as const);
const picked = m.pick(await ctx.ui.select(R.askFactor, m.options));
if (!picked) {
say(pi, S.common.cancelled);
return;
}
factor = picked === "x4" ? 4 : 2;
}
req.factor = factor ?? DEFAULT_UPSCALE_FACTOR;
}
if (action === "crop") {
let formats = parsed.formats;
if (!formats?.length && interactive) {
const labels = SCREEN_FORMATS.map((f) => S.formats.labels[f]);
const chosen = await ctx.ui.select(R.askFormats, [R.allFormats, ...labels]);
if (chosen === undefined) {
say(pi, S.common.cancelled);
return;
}
const idx = labels.indexOf(chosen);
formats = idx >= 0 ? [SCREEN_FORMATS[idx] as Format] : [...SCREEN_FORMATS];
}
req.formats = formats?.length ? formats : [...SCREEN_FORMATS];
}
if (action === "background") {
let cropToSubject = parsed.cropToSubject;
if (cropToSubject === undefined && interactive) {
cropToSubject = await ctx.ui.confirm(R.askCropSubject, R.askCropSubjectHint);
}
req.cropToSubject = cropToSubject ?? false;
}
if (action === "edit" || action === "poster") {
let instruction = parsed.rest;
if (action === "edit" && !instruction && interactive) {
const typed = await ctx.ui.input(R.askEdit, R.askEditPlaceholder);
if (typed === undefined) {
say(pi, S.common.cancelled);
return;
}
instruction = typed.trim();
}
if (action === "edit" && !instruction) throw new RetouchError(R.needEdit);
if (instruction) req.instruction = instruction;
}
if (action === "poster" && interactive) {
const ok = await ctx.ui.confirm(
S.confirm.usePhoto.title,
fill(S.confirm.usePhoto.message, { file: basename(input) }),
);
if (!ok) {
say(pi, S.common.cancelled);
return;
}
}
// -- 4. do the work ----------------------------------------------------
ctx.ui.setStatus(STATUS_KEY, fill(R.working, { file: basename(input) }));
const result = await runRetouch(deps, req, {
onProgress: (msg) => ctx.ui.setStatus(STATUS_KEY, msg),
...(ctx.signal ? { signal: ctx.signal } : {}),
directorContext: { modelRegistry: ctx.modelRegistry, model: ctx.model },
});
pi.appendEntry(ENTRY_TYPE, {
action: result.action,
input: result.input,
dir: result.dir,
outputs: result.outputs.map((o) => o.path),
tool: result.tool,
});
// -- 5. hand off, or show what came out --------------------------------
if (result.handoff) {
say(pi, describeResult(result));
if (deps.startPoster) {
await deps.startPoster(
{
photoPath: result.handoff.photoPath,
...(result.handoff.freeText ? { freeText: result.handoff.freeText } : {}),
},
ctx,
);
} else {
// No hand-off callback in Deps: type the command for him and let /poster's
// free-text parser pick the path back up. Same destination, one extra hop.
say(pi, fill(R.handoffManual, { file: basename(result.handoff.photoPath) }));
pi.sendUserMessage(
`/poster foto: ${result.handoff.photoPath}${result.handoff.freeText ? ` ${result.handoff.freeText}` : ""}`,
{ expandPromptTemplates: true },
);
}
return;
}
const content: (TextContent | ImageContent)[] = [];
const first = result.outputs[0];
if (first) {
const image = await previewImage(first.path);
if (image) content.push(image);
}
content.push({ type: "text", text: describeResult(result) });
say(pi, content);
ctx.ui.notify(R.notifyDone, "info");
// -- 6. the folder -----------------------------------------------------
if (interactive && result.outputs.length) {
const open = await ctx.ui.confirm(S.confirm.openFolder.title, S.confirm.openFolder.message);
if (open && !openInFinder(result.dir)) {
say(pi, errorText(S.errors.openFolderFailed(result.dir)));
}
}
} catch (e) {
const italian = e instanceof RetouchError
? e.italian
: e instanceof BackendError && e.italian
? e.italian
: errorText(S.errors.unknown(e instanceof Error ? e.message : String(e)));
ctx.ui.notify(S.notify.failed, "error");
throw new Error(italian);
} finally {
// A status line that outlives the command is worse than no status line at all.
ctx.ui.setStatus(STATUS_KEY, undefined);
}
}
// ---------------------------------------------------------------------------
// The tool — same implementation, no dialogs
// ---------------------------------------------------------------------------
const RetouchParams = Type.Object({
image_path: Type.String({ description: "Absolute path to the image to work on." }),
action: StringEnum([...RETOUCH_ACTIONS], {
description:
"background = cut the subject out (CPU); upscale = enlarge (CPU); crop = social crops (CPU); " +
"poster = hand the photo to the poster flow as artwork; edit = generative img2img " +
"(UNRELIABLE on this 16GB Mac — Draw Things issue #121).",
}),
instruction: Type.Optional(Type.String({ description: "Italian, required for action 'edit'." })),
factor: Type.Optional(Type.Number({ minimum: 1.1, maximum: 8, description: "Upscale factor. Default 2." })),
formats: Type.Optional(Type.Array(StringEnum([...FORMATS]), { description: "Crop targets. Default: every screen format." })),
crop_to_subject: Type.Optional(Type.Boolean({ description: "background only: crop to the subject's bounding box." })),
strength: Type.Optional(Type.Number({ minimum: 0.05, maximum: 1, description: "edit only. Default 0.55." })),
});
export function createRetouchTool(deps: RetouchDeps) {
return defineTool({
name: "imgen_retouch",
label: S.toolLabels.retouch,
description:
"Operate on an image the user already has: remove its background, upscale it, crop it " +
"for social formats, adopt it as poster artwork, or edit it generatively. The first " +
"three are CPU-only and always available even when the diffusion backend is broken.",
parameters: RetouchParams,
// Two concurrent diffusion runs OOM a 16GB machine.
executionMode: "sequential",
async execute(_id, params, signal, onUpdate, ctx) {
const req: RetouchRequest = {
action: params.action as RetouchAction,
input: normaliseImagePath(params.image_path, ctx.cwd),
};
if (params.instruction) req.instruction = params.instruction;
if (params.factor !== undefined) req.factor = params.factor;
if (params.formats?.length) req.formats = params.formats as Format[];
if (params.crop_to_subject !== undefined) req.cropToSubject = params.crop_to_subject;
if (params.strength !== undefined) req.strength = params.strength;
const result = await runRetouch(deps, req, {
onProgress: (msg) => onUpdate?.({ content: [{ type: "text", text: msg }], details: undefined }),
...(signal ? { signal } : {}),
directorContext: { modelRegistry: ctx.modelRegistry, model: ctx.model },
});
const content: (TextContent | ImageContent)[] = [];
const first = result.outputs[0];
if (first) {
const image = await previewImage(first.path);
if (image) content.push(image);
}
const text = result.handoff
? `Photo adopted as poster artwork: ${result.handoff.photoPath}. Continue with the poster flow ` +
`(art.source="photo", art.photoPath=${JSON.stringify(result.handoff.photoPath)}).`
: [
describeResult(result),
"",
...result.outputs.map((o) => `${o.path} (${o.width}x${o.height})`),
].join("\n");
content.push({ type: "text", text });
return { content, details: result };
},
});
}
export default registerRetouch;