feat(prompt): show phase + token progress for LLM prompt expansion

[Pfannkuchensack]

Summary

Both /utilities/expand-prompt and /utilities/image-to-prompt were blocking calls with only a spinner — users had no signal whether the model was still loading or already generating, which is rough on larger LLMs.

The pipelines now stream tokens via HuggingFace's TextIteratorStreamer, and the routes emit new llm_task_progress / complete / error socket events correlated to a client-supplied task_id and routed privately to the requesting user. The Expand Prompt and Image-to-Prompt popovers render a phase label ("Loading model…" / "Generating…") plus a token progress bar (current / max_tokens) while the request is in flight.

Related Issues / Discussions

https://discord.com/channels/1020123559063990373/1149506274971631688/1506024129730707497

QA Instructions

1. Install a small Text LLM (e.g. Qwen2.5-1.5B-Instruct starter model) and a LLaVA OneVision model.
2. Open the generation UI, click the ✨ Expand Prompt button, pick the Text LLM, type a short prompt, click Expand.
   • First call: popover shows "Loading model…" with an indeterminate bar, then "Generating… X / 300" with a filling bar, then closes with the expanded prompt written to the field.
   • Second call (same model): skips straight to "Generating…" because the model is already in VRAM.
3. Repeat with the 🖼️ Image to Prompt button using a LLaVA model — same phase + progress UX, counter goes up to 400.
4. Open DevTools → Network → WS → frame inspector and confirm llm_task_progress events arrive with monotonically increasing current_tokens and a matching task_id.
5. Trigger an error (e.g. delete the model mid-request, or send a model_key for a non-LLM type) and confirm the popover shows the error message and the spinner clears.
6. Multi-user mode (if applicable): a second logged-in user should NOT receive progress events for the first user's request — events are routed to user:{user_id} rooms only.
7. Backend unit tests: pytest tests/backend/text_llm/ — all 7 should pass (one new test covers the per-chunk progress callback).
8. Frontend checks from invokeai/frontend/web/: pnpm lint and pnpm test:no-watch should pass (the one pre-existing timer flake in navigation-api.test.ts is unrelated).

Merge Plan

No DB migration. The new socket events are additive — old frontends ignore them. Generated OpenAPI schema (schema.ts) should be regenerated after merge so the LLM task event payload types come from S[...] instead of the hand-typed fallbacks in services/events/types.ts.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[lstein]

Code Review

Overview

Nice feature — the design is sound overall. I verified FastAPIEventService.dispatch uses call_soon_threadsafe, so emitting from the asyncio.to_thread worker is safe; the user:{user_id} + admin room routing exactly mirrors the existing private-event handlers in sockets.py; and the min(token_count, max_new_tokens) clamp correctly prevents a pydantic le=1 validation crash inside the callback. task_id being client-supplied is not a spoofing surface, since routing is by the authenticated user_id.

However, there is one issue I'd consider a merge blocker.

Critical

1. Deadlock when generate() raises in the producer thread

Affects both text_llm_pipeline.py and llava_onevision_pipeline.py.

TextIteratorStreamer is constructed without a timeout, so its __next__ blocks indefinitely on queue.get(). In transformers (verified against both the locked 4.56 and 5.5.4), streamer.end() is only called at the normal exit of the generation loop in _sample — it is not in a finally. If model.generate() raises (CUDA OOM is the realistic case, on exactly the larger LLMs this PR is motivated by), _generate captures the exception into generation_error, the thread dies, end() is never called, and the consuming for chunk in streamer: loop hangs forever. The thread.join() / re-raise code below it is unreachable.

Consequences: the asyncio.to_thread worker leaks, the HTTP request never resolves, the loaded_model.model_on_device() context never exits (model stays pinned on device), and the frontend spinner is stuck permanently.

The transformers docstring calls this out explicitly: the timeout param is "useful to handle exceptions in .generate(), when it is called in a separate thread." Suggested fix in both pipelines:

streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=STREAM_TIMEOUT)

def _generate() -> None:
    try:
        self._model.generate(**generation_kwargs)
    except BaseException as e:
        generation_error.append(e)
        streamer.end()  # unblock the consumer

Calling streamer.end() in the except block is the primary fix; a generous timeout= (with queue.Empty handled) is a good belt-and-braces addition since generate() can also hang without raising. Note that first-token latency can be long, so the timeout needs to be generous.

Please also add a test for this path — the existing FakeStreamer can't reproduce it; a test with model.generate.side_effect = RuntimeError(...) and a real streamer would have caught this.

(FYI: upgrading to transformers 5.5.4 does not mitigate this — the streamer and _sample code are unchanged in the relevant way, and 5.x even adds a new early-raise path: generate() raises NotImplementedError when a streamer is combined with continuous batching, which would trigger the same hang.)

Medium

2. Frontend store leak on the event/HTTP race

stores.ts / setEventListeners.tsx: the server emits llm_task_complete before returning the HTTP response, but socket delivery is independent of the HTTP channel. If the socket event lands after the mutation's finally has run clearLLMTaskState, the handler re-creates the entry ({status: 'complete'} — or a late progress event re-creates a progress entry), and nothing ever deletes it. $llmTaskStates grows by one orphan per request for the session.

Simplest fix: make the llm_task_complete handler delete the key instead of storing a 'complete' state (the component treats complete as render-null anyway), keeping clearLLMTaskState as the backstop for progress/error stragglers.

3. The error branch of LLMTaskProgressDisplay is effectively dead

The component is rendered only while isLoading is true, and the mutation's finally clears the store entry the moment the request settles — so state.status === 'error' can practically never be displayed (errors surface via the RTK Query toast instead). Either drop the error branch and the llm_task_error state storage, or keep the display mounted briefly on failure so QA step 5's behavior actually comes from this path.

This was confirmed empirically by deleting the model while it was loading. The progress bar stopped, but the error message was not displayed.

Minor / suggestions

• Per-chunk cost: each chunk re-encodes the full accumulated text (O(n²) overall) and emits a socket event through call_soon_threadsafe + socket.io. Fine at the 300/400 defaults, but max_tokens can be 2048; consider throttling emissions (every N tokens or ~100 ms), analogous to per-step invocation progress.
• types.ts hand-typed payloads are already unnecessary: this PR's own schema.ts diff contains LLMTaskProgressEvent etc., so the socket payloads can use S['LLMTaskProgressEvent'] now, matching the bulk_download_* convention — no need to defer to a post-merge regen as the PR body suggests.
• Duplication: the ~15-line progress_callback construction is copy-pasted between _run_expand_prompt and _run_image_to_prompt; a small _make_progress_callback(events, task_id, user_id) helper in utilities.py would keep them in sync.
• Test coverage: the text pipeline callback test is good; there's no test for the LLaVA pipeline changes (which also changed output parsing from the split("assistant\n") hack to skip_prompt=True — a behavior-relevant rewrite), and none for the exception path in finding 1.

Verdict

Well-integrated and follows the codebase's event/room conventions closely. Finding 1 is the blocker: the code path added specifically to surface errors instead deadlocks the request (and pins the model on device) on the failure modes this PR is meant to make less painful. The fix is small and localized to the two _generate wrappers.

Note two conflicts with main.

🤖 Generated with Claude Code and edited by @lstein

[lstein]

See review comment.

[lstein]

Pushed fixes for the review findings + merged main

I pushed two commits to this branch (fast-forward, nothing of yours overwritten): a main merge and a fixes commit (40caf66). Here's what changed.

Merge with main

Resolved three conflicts, all additive:

• utilities.py — kept the task_id progress parameters.
• sockets.py / events/types.ts — kept both the LLM-task and the new workflow event registrations.

Heads-up: the merge pulled in main's frontend dependency bumps (TypeScript 5.9, Vite 8, Vitest 4), so a fresh pnpm install is needed after pulling.

Critical — streamer deadlock on generation failure

If model.generate() raised (CUDA OOM being the realistic case on the larger LLMs this feature targets), transformers never called streamer.end(), so the consuming for chunk in streamer: loop blocked forever — leaking the worker thread, never resolving the HTTP request, and pinning the model on device.

Both text_llm_pipeline.py and llava_onevision_pipeline.py now:

• call streamer.end() in the worker's except block so the consumer is unblocked, and
• pass a generous STREAM_TIMEOUT to TextIteratorStreamer as a backstop for a hang that never raises; a queue.Empty is surfaced as an error instead of blocking on thread.join().

Added regression tests that drive generate().side_effect through the real streamer (the old FakeStreamer couldn't reproduce this) for both pipelines.

Medium — frontend store leak + dead error branch

• Socket delivery is independent of the HTTP response, so a llm_task_complete/llm_task_error event arriving after the mutation's finally cleared the entry re-created an orphan (one per request). Those handlers now delete the store key instead of writing a terminal state.
• The error branch of LLMTaskProgressDisplay could never render (the component unmounts the moment isLoading flips false), so I removed it — load/generation errors already surface via the RTK Query toast. The LLMTaskState union collapses to the in-flight progress case.

Minor

• Throttle progress emissions to ≥100 ms with a guaranteed final emit, to bound the O(n²) re-encode + socket cost when max_tokens is large.
• Switched events/types.ts to the generated S['LLMTask*Event'] schema types (they're already in schema.ts) instead of the hand-typed payloads.
• Extracted a shared _make_progress_callback helper so the expand-prompt and image-to-prompt endpoints stay in sync.
• Added LLaVA pipeline test coverage (tests/backend/llava_onevision/).

Verification

21 backend + 28 frontend event tests pass; tsc, eslint, knip, prettier, and ruff are all clean.

One judgment call: I removed the dead inline error UI rather than keeping it mounted on failure, since errors already toast. If you'd prefer the error shown inline in the popover instead, that's an easy follow-up — happy to switch it.