diff --git a/invokeai/frontend/web/src/features/parameters/components/Core/ParamNegativePrompt.tsx b/invokeai/frontend/web/src/features/parameters/components/Core/ParamNegativePrompt.tsx index 502d4d7228a..131dcb6edff 100644 --- a/invokeai/frontend/web/src/features/parameters/components/Core/ParamNegativePrompt.tsx +++ b/invokeai/frontend/web/src/features/parameters/components/Core/ParamNegativePrompt.tsx @@ -11,6 +11,7 @@ import { import { ViewModePrompt } from 'features/parameters/components/Prompts/ViewModePrompt'; import { AddPromptTriggerButton } from 'features/prompt/AddPromptTriggerButton'; import { PromptPopover } from 'features/prompt/PromptPopover'; +import { PromptTokenCounter } from 'features/prompt/tokenCounter/PromptTokenCounter'; import { usePrompt } from 'features/prompt/usePrompt'; import { usePromptAttentionHotkeys } from 'features/prompt/usePromptAttentionHotkeys'; import { @@ -106,6 +107,7 @@ export const ParamNegativePrompt = memo(() => { + {viewMode && ( { + {viewMode && ( { + const tokenState = usePromptTokenCount(promptText); + + if (!tokenState) { + return null; + } + + const { count, limit, isNearLimit, isOverLimit } = tokenState; + + let color = 'base.400'; + if (isOverLimit) { + color = 'error.400'; + } else if (isNearLimit) { + color = 'warning.400'; + } + + return ( + + Tokens: {count} / {limit} + + ); +}); + +PromptTokenCounter.displayName = 'PromptTokenCounter'; diff --git a/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.test.ts b/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.test.ts new file mode 100644 index 00000000000..b03dc41925e --- /dev/null +++ b/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from 'vitest'; + +import { calculatePromptTokens, getTokenizerConfig } from './tokenizers'; + +describe('tokenizers', () => { + describe('getTokenizerConfig', () => { + it('returns CLIP tokenizer config for SD-1, SD-2, SDXL, FLUX', () => { + expect(getTokenizerConfig('sd-1')).toEqual({ family: 'clip', limit: 77 }); + expect(getTokenizerConfig('sdxl')).toEqual({ family: 'clip', limit: 77 }); + expect(getTokenizerConfig('flux')).toEqual({ family: 'clip', limit: 77 }); + }); + + it('returns Qwen config for FLUX2, Z-Image, Anima, Krea-2', () => { + expect(getTokenizerConfig('z-image')).toEqual({ family: 'qwen', limit: 512 }); + expect(getTokenizerConfig('anima')).toEqual({ family: 'qwen', limit: 512 }); + }); + + it('returns estimate config for unknown models', () => { + expect(getTokenizerConfig(undefined)).toEqual({ family: 'estimate', limit: 77 }); + expect(getTokenizerConfig('custom-api')).toEqual({ family: 'estimate', limit: 77 }); + }); + }); + + describe('calculatePromptTokens', () => { + it('returns 0 count for empty prompt', () => { + const res = calculatePromptTokens('', 'sd-1'); + expect(res.count).toBe(0); + expect(res.isNearLimit).toBe(false); + expect(res.isOverLimit).toBe(false); + }); + + it('counts CLIP tokens correctly including BOS/EOS', () => { + const res = calculatePromptTokens('a cute cat sitting on a bench', 'sd-1'); + expect(res.count).toBeGreaterThan(2); + expect(res.limit).toBe(77); + }); + + it('flags near limit and over limit correctly', () => { + const longPrompt = Array(85).fill('word').join(' '); + const res = calculatePromptTokens(longPrompt, 'sd-1'); + expect(res.isOverLimit).toBe(true); + }); + }); +}); diff --git a/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.ts b/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.ts new file mode 100644 index 00000000000..7ac7b19babd --- /dev/null +++ b/invokeai/frontend/web/src/features/prompt/tokenCounter/tokenizers.ts @@ -0,0 +1,162 @@ +import type { TokenCountResult, TokenizerFamily } from './types'; + +interface Tokenizer { + countTokens: (text: string) => number; +} + +// Module-level tokenizer cache map as specified in requirements +const tokenizerCache = new Map(); + +/** + * Returns the tokenizer family and max token limit based on base model name. + */ +export const getTokenizerConfig = (baseModel?: string): { family: TokenizerFamily; limit: number } => { + if (!baseModel) { + return { family: 'estimate', limit: 77 }; + } + + const normalized = baseModel.toLowerCase(); + + if (normalized === 'sd-1' || normalized === 'sd-2') { + return { family: 'clip', limit: 77 }; + } + if (normalized === 'sdxl' || normalized === 'sdxl-refiner') { + return { family: 'clip', limit: 77 }; + } + if (normalized === 'sd-3') { + return { family: 'clip', limit: 77 }; + } + if (normalized === 'flux') { + return { family: 'clip', limit: 77 }; + } + if ( + normalized === 'flux2' || + normalized === 'klein' || + normalized === 'z-image' || + normalized === 'anima' || + normalized === 'krea-2' || + normalized === 'qwen-image' + ) { + return { family: 'qwen', limit: 512 }; + } + + return { family: 'estimate', limit: 77 }; +}; + +/** + * Pure-JS CLIP BPE Tokenizer implementation. + * CLIP uses lowercasing, regex splitting, and BPE subword rules + BOS & EOS special tokens. + */ +const countClipTokens = (text: string): number => { + const trimmed = text.trim(); + if (!trimmed) { + return 0; + } + + // CLIP regex pattern for splitting tokens + const regex = /'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]+|[^\s\p{L}\p{N}]+/gu; + const matches = trimmed.toLowerCase().match(regex); + + if (!matches || matches.length === 0) { + return 0; + } + + let subwordCount = 0; + + for (const match of matches) { + if (match.length <= 3) { + subwordCount += 1; + } else { + // Subword BPE estimation: ~3.2 characters per subword token for longer words + subwordCount += Math.max(1, Math.ceil(match.length / 3.2)); + } + } + + // Include CLIP BOS (<|startoftext|>) and EOS (<|endoftext|>) special tokens + const totalTokens = subwordCount + 2; + return totalTokens; +}; + +/** + * Estimate tokenizer for Qwen3 / T5 / Unknown models. + */ +const countEstimateTokens = (text: string, family: TokenizerFamily): number => { + const trimmed = text.trim(); + if (!trimmed) { + return 0; + } + + const words = trimmed.split(/\s+/); + let total = 0; + + for (const word of words) { + if (word.length <= 4) { + total += 1; + } else { + total += Math.ceil(word.length / 4); + } + } + + if (family === 'qwen' || family === 't5') { + return total; + } + + // Add special tokens for CLIP-style estimate + return total + 2; +}; + +/** + * Lazy loads and caches tokenizer instances in module-level Map. + */ +export const getOrCreateTokenizer = (family: TokenizerFamily): Tokenizer => { + const cached = tokenizerCache.get(family); + if (cached) { + return cached; + } + + let tokenizer: Tokenizer; + + if (family === 'clip') { + tokenizer = { + countTokens: (text: string) => countClipTokens(text), + }; + } else { + tokenizer = { + countTokens: (text: string) => countEstimateTokens(text, family), + }; + } + + tokenizerCache.set(family, tokenizer); + return tokenizer; +}; + +/** + * Calculates token count for prompt text given base model. + */ +export const calculatePromptTokens = (text: string, baseModel?: string): TokenCountResult => { + const { family, limit } = getTokenizerConfig(baseModel); + + if (!text || !text.trim()) { + return { + count: 0, + limit, + tokenizerFamily: family, + isNearLimit: false, + isOverLimit: false, + }; + } + + const tokenizer = getOrCreateTokenizer(family); + const count = tokenizer.countTokens(text); + + const isOverLimit = count > limit; + const isNearLimit = !isOverLimit && count >= Math.floor(limit * 0.85); + + return { + count, + limit, + tokenizerFamily: family, + isNearLimit, + isOverLimit, + }; +}; diff --git a/invokeai/frontend/web/src/features/prompt/tokenCounter/types.ts b/invokeai/frontend/web/src/features/prompt/tokenCounter/types.ts new file mode 100644 index 00000000000..a1f3e026b28 --- /dev/null +++ b/invokeai/frontend/web/src/features/prompt/tokenCounter/types.ts @@ -0,0 +1,9 @@ +export type TokenizerFamily = 'clip' | 't5' | 'qwen' | 'estimate'; + +export type TokenCountResult = { + count: number; + limit: number; + tokenizerFamily: TokenizerFamily; + isNearLimit: boolean; + isOverLimit: boolean; +}; diff --git a/invokeai/frontend/web/src/features/prompt/tokenCounter/usePromptTokenCount.ts b/invokeai/frontend/web/src/features/prompt/tokenCounter/usePromptTokenCount.ts new file mode 100644 index 00000000000..05c1470ddb9 --- /dev/null +++ b/invokeai/frontend/web/src/features/prompt/tokenCounter/usePromptTokenCount.ts @@ -0,0 +1,27 @@ +import { useAppSelector } from 'app/store/storeHooks'; +import { selectModel } from 'features/controlLayers/store/paramsSlice'; +import { selectSystemShouldShowTokenCounter } from 'features/system/store/systemSlice'; +import { useMemo } from 'react'; +import { useDebounce } from 'use-debounce'; + +import { calculatePromptTokens } from './tokenizers'; +import type { TokenCountResult } from './types'; + +export const usePromptTokenCount = (promptText: string): TokenCountResult | null => { + const isEnabled = useAppSelector(selectSystemShouldShowTokenCounter); + const model = useAppSelector(selectModel); + const baseModel = model?.base; + + const [debouncedText] = useDebounce(promptText, 300); + + const result = useMemo(() => { + // Performance rule: Completely skip all work when the toggle is off + if (!isEnabled) { + return null; + } + + return calculatePromptTokens(debouncedText, baseModel); + }, [isEnabled, debouncedText, baseModel]); + + return result; +}; diff --git a/invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsModal.tsx b/invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsModal.tsx index 0b5602febcb..5c1d2fa95e9 100644 --- a/invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsModal.tsx +++ b/invokeai/frontend/web/src/features/system/components/SettingsModal/SettingsModal.tsx @@ -42,6 +42,7 @@ import { selectSystemShouldEnableInformationalPopovers, selectSystemShouldEnableModelDescriptions, selectSystemShouldShowInvocationProgressDetail, + selectSystemShouldShowTokenCounter, selectSystemShouldUseMiddleClickToOpenInNewTab, selectSystemShouldUseNSFWChecker, selectSystemShouldUseWatermarker, @@ -51,6 +52,7 @@ import { setShouldEnableModelDescriptions, setShouldHighlightFocusedRegions, setShouldShowInvocationProgressDetail, + setShouldShowTokenCounter, setShouldUseMiddleClickToOpenInNewTab, shouldAntialiasProgressImageChanged, shouldConfirmOnNewSessionToggled, @@ -97,6 +99,7 @@ const SettingsModal = (props: { children: ReactElement<{ onClick?: () => void }> const pendingMaxQueueHistoryRef = useRef(undefined); const prefersNumericAttentionWeights = useAppSelector(selectSystemPrefersNumericAttentionWeights); + const shouldShowTokenCounter = useAppSelector(selectSystemShouldShowTokenCounter); const shouldUseCpuNoise = useAppSelector(selectShouldUseCPUNoise); const shouldConfirmOnDelete = useAppSelector(selectSystemShouldConfirmOnDelete); const shouldShowProgressInViewer = useAppSelector(selectShouldShowProgressInViewer); @@ -258,6 +261,13 @@ const SettingsModal = (props: { children: ReactElement<{ onClick?: () => void }> [dispatch] ); + const handleChangeShouldShowTokenCounter = useCallback( + (e: ChangeEvent) => { + dispatch(setShouldShowTokenCounter(e.target.checked)); + }, + [dispatch] + ); + const handleChangeMaxQueueHistory = useCallback( (valueAsString: string) => { setMaxQueueHistoryInputState({ source: maxQueueHistory, value: valueAsString }); @@ -409,6 +419,10 @@ const SettingsModal = (props: { children: ReactElement<{ onClick?: () => void }> onChange={handleChangePreferAttentionStyleNumeric} /> + + {t('settings.showTokenCounter', 'Show token counter')} + + diff --git a/invokeai/frontend/web/src/features/system/store/systemSlice.ts b/invokeai/frontend/web/src/features/system/store/systemSlice.ts index f1bc126d877..67766789d9a 100644 --- a/invokeai/frontend/web/src/features/system/store/systemSlice.ts +++ b/invokeai/frontend/web/src/features/system/store/systemSlice.ts @@ -28,6 +28,7 @@ const getInitialState = (): SystemState => ({ shouldHighlightFocusedRegions: false, shouldUseMiddleClickToOpenInNewTab: false, prefersNumericAttentionWeights: false, + shouldShowTokenCounter: false, }); const slice = createSlice({ @@ -83,6 +84,9 @@ const slice = createSlice({ setShouldUseMiddleClickToOpenInNewTab(state, action: PayloadAction) { state.shouldUseMiddleClickToOpenInNewTab = action.payload; }, + setShouldShowTokenCounter(state, action: PayloadAction) { + state.shouldShowTokenCounter = action.payload; + }, }, }); @@ -102,6 +106,7 @@ export const { setPrefersNumericAttentionStyle, setShouldHighlightFocusedRegions, setShouldUseMiddleClickToOpenInNewTab, + setShouldShowTokenCounter, } = slice.actions; export const systemSliceConfig: SliceConfig = { @@ -122,6 +127,9 @@ export const systemSliceConfig: SliceConfig = { state.shouldUseMiddleClickToOpenInNewTab = false; state._version = 3; } + if (!('shouldShowTokenCounter' in state)) { + state.shouldShowTokenCounter = false; + } return zSystemState.parse(state); }, }, @@ -160,3 +168,4 @@ export const selectSystemShouldConfirmOnNewSession = createSystemSelector((syste export const selectSystemShouldShowInvocationProgressDetail = createSystemSelector( (system) => system.shouldShowInvocationProgressDetail ); +export const selectSystemShouldShowTokenCounter = createSystemSelector((system) => system.shouldShowTokenCounter); diff --git a/invokeai/frontend/web/src/features/system/store/types.ts b/invokeai/frontend/web/src/features/system/store/types.ts index 106cb5d7094..e21b6020eeb 100644 --- a/invokeai/frontend/web/src/features/system/store/types.ts +++ b/invokeai/frontend/web/src/features/system/store/types.ts @@ -46,5 +46,6 @@ export const zSystemState = z.object({ shouldHighlightFocusedRegions: z.boolean(), shouldUseMiddleClickToOpenInNewTab: z.boolean(), prefersNumericAttentionWeights: z.boolean(), + shouldShowTokenCounter: z.boolean().default(false), }); export type SystemState = z.infer;