Skip to content

Commit 571d286

Browse files
Enhance LoRA picker to default to current base model architecture
Co-authored-by: kent <kent@invoke.ai> Enhance LoRA picker to default filter by current base model architecture Co-authored-by: kent <kent@invoke.ai>
1 parent 1320a2c commit 571d286

File tree

3 files changed

+36
-5
lines changed

3 files changed

+36
-5
lines changed

invokeai/frontend/web/src/common/components/Picker/Picker.tsx

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -198,6 +198,10 @@ type PickerProps<T extends object> = {
198198
* Whether the picker should be searchable. If true, renders a search input.
199199
*/
200200
searchable?: boolean;
201+
/**
202+
* Initial state for group toggles. If provided, groups will start with these states instead of all being disabled.
203+
*/
204+
initialGroupStates?: Record<string, boolean>;
201205
};
202206

203207
export type PickerContextState<T extends object> = {
@@ -312,7 +316,10 @@ const flattenOptions = <T extends object>(options: OptionOrGroup<T>[]): T[] => {
312316

313317
type GroupStatusMap = Record<string, boolean>;
314318

315-
const useTogglableGroups = <T extends object>(options: OptionOrGroup<T>[]) => {
319+
const useTogglableGroups = <T extends object>(
320+
options: OptionOrGroup<T>[],
321+
initialGroupStates?: Record<string, boolean>
322+
) => {
316323
const groupsWithOptions = useMemo(() => {
317324
const ids: string[] = [];
318325
for (const optionOrGroup of options) {
@@ -332,14 +339,16 @@ const useTogglableGroups = <T extends object>(options: OptionOrGroup<T>[]) => {
332339
const groupStatusMap = $groupStatusMap.get();
333340
const newMap: GroupStatusMap = {};
334341
for (const id of groupsWithOptions) {
335-
if (newMap[id] === undefined) {
336-
newMap[id] = false;
342+
if (initialGroupStates && initialGroupStates[id] !== undefined) {
343+
newMap[id] = initialGroupStates[id];
337344
} else if (groupStatusMap[id] !== undefined) {
338345
newMap[id] = groupStatusMap[id];
346+
} else {
347+
newMap[id] = false;
339348
}
340349
}
341350
$groupStatusMap.set(newMap);
342-
}, [groupsWithOptions, $groupStatusMap]);
351+
}, [groupsWithOptions, $groupStatusMap, initialGroupStates]);
343352

344353
const toggleGroup = useCallback(
345354
(idToToggle: string) => {
@@ -511,10 +520,14 @@ export const Picker = typedMemo(<T extends object>(props: PickerProps<T>) => {
511520
OptionComponent = DefaultOptionComponent,
512521
NextToSearchBar,
513522
searchable,
523+
initialGroupStates,
514524
} = props;
515525
const rootRef = useRef<HTMLDivElement>(null);
516526
const inputRef = useRef<HTMLInputElement>(null);
517-
const { $groupStatusMap, $areAllGroupsDisabled, toggleGroup } = useTogglableGroups(optionsOrGroups);
527+
const { $groupStatusMap, $areAllGroupsDisabled, toggleGroup } = useTogglableGroups(
528+
optionsOrGroups,
529+
initialGroupStates
530+
);
518531
const $activeOptionId = useAtom(getFirstOptionId(optionsOrGroups, getOptionId));
519532
const $compactView = useAtom(true);
520533
const $optionsOrGroups = useAtom(optionsOrGroups);

invokeai/frontend/web/src/features/lora/components/LoRASelect.tsx

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { useRelatedGroupedModelCombobox } from 'common/hooks/useRelatedGroupedMo
66
import { loraAdded, selectLoRAsSlice } from 'features/controlLayers/store/lorasSlice';
77
import { selectBase } from 'features/controlLayers/store/paramsSlice';
88
import { ModelPicker } from 'features/parameters/components/ModelPicker';
9+
import { API_BASE_MODELS } from 'features/parameters/types/constants';
910
import { memo, useCallback, useMemo } from 'react';
1011
import { useTranslation } from 'react-i18next';
1112
import { useLoRAModels } from 'services/api/hooks/modelsByType';
@@ -58,6 +59,19 @@ const LoRASelect = () => {
5859
return t('models.addLora');
5960
}, [isLoading, options.length, t]);
6061

62+
// Calculate initial group states to default to the current base model architecture
63+
const initialGroupStates = useMemo(() => {
64+
if (!currentBaseModel) {
65+
return undefined;
66+
}
67+
68+
// Determine the group ID for the current base model
69+
const groupId = API_BASE_MODELS.includes(currentBaseModel) ? 'api' : currentBaseModel;
70+
71+
// Return a map with only the current base model group enabled
72+
return { [groupId]: true };
73+
}, [currentBaseModel]);
74+
6175
return (
6276
<FormControl gap={2}>
6377
<InformationalPopover feature="lora">
@@ -72,6 +86,7 @@ const LoRASelect = () => {
7286
placeholder={placeholder}
7387
getIsOptionDisabled={getIsDisabled}
7488
noOptionsText={t('models.noLoRAsInstalled')}
89+
initialGroupStates={initialGroupStates}
7590
/>
7691
</FormControl>
7792
);

invokeai/frontend/web/src/features/parameters/components/ModelPicker.tsx

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ export const ModelPicker = typedMemo(
125125
isInvalid,
126126
className,
127127
noOptionsText,
128+
initialGroupStates,
128129
}: {
129130
modelConfigs: T[];
130131
selectedModelConfig: T | undefined;
@@ -137,6 +138,7 @@ export const ModelPicker = typedMemo(
137138
isInvalid?: boolean;
138139
className?: string;
139140
noOptionsText?: string;
141+
initialGroupStates?: Record<string, boolean>;
140142
}) => {
141143
const { t } = useTranslation();
142144
const options = useMemo<T[] | Group<T>[]>(() => {
@@ -244,6 +246,7 @@ export const ModelPicker = typedMemo(
244246
NextToSearchBar={<NavigateToModelManagerButton />}
245247
getIsOptionDisabled={getIsOptionDisabled}
246248
searchable
249+
initialGroupStates={initialGroupStates}
247250
/>
248251
</PopoverBody>
249252
</PopoverContent>

0 commit comments

Comments
 (0)