From 9e6e17e98a9db76d081518ed446b5aae70936435 Mon Sep 17 00:00:00 2001 From: Marcelo Duarte Date: Mon, 31 Aug 2026 16:28:50 -0300 Subject: [PATCH 1/2] fix(ui): keep base model selection consistent across training flows --- ui/src/Components/BaseModelDropdown.jsx | 39 ++- ui/src/Components/BaseModelDropdownHelper.js | 58 ++++ .../BaseModelDropdownHelper.test.js | 101 ++++++ .../CreateEditModelTrainingHelper.js | 18 +- .../CreateEditModelTrainingModal.jsx | 8 +- .../Components/LabelingTool/LabelingTool.jsx | 10 + .../LabelingTool/LabelingToolRightPanel.jsx | 7 +- ui/src/Components/LabelingTool/LeftPanel.jsx | 317 ------------------ ui/src/Components/LabelingTool/RightPanel.jsx | 247 -------------- 9 files changed, 208 insertions(+), 597 deletions(-) create mode 100644 ui/src/Components/BaseModelDropdownHelper.js create mode 100644 ui/src/Components/BaseModelDropdownHelper.test.js delete mode 100644 ui/src/Components/LabelingTool/LeftPanel.jsx delete mode 100644 ui/src/Components/LabelingTool/RightPanel.jsx diff --git a/ui/src/Components/BaseModelDropdown.jsx b/ui/src/Components/BaseModelDropdown.jsx index 7ee893f6..02c8a1f2 100644 --- a/ui/src/Components/BaseModelDropdown.jsx +++ b/ui/src/Components/BaseModelDropdown.jsx @@ -1,5 +1,10 @@ import React from "react"; +import PropTypes from "prop-types"; import { Dropdown, Option, Field } from "@fluentui/react-components"; +import { + applyBaseModelSelection, + normalizeBaseModelOptions, +} from "./BaseModelDropdownHelper"; const styles = { container: { @@ -57,20 +62,12 @@ const renderOption = (option) => { function BaseModelDropdown({ componentState, setComponentState, - onFormChange, - }) { const { baseModelId, baseModelIdError, cataloguedModels = [] } = componentState; - const options = React.useMemo(() => { - return cataloguedModels.map((m) => ({ - key: m.key || "none", - baseModelName: m.value.baseModelName || "", - description: m.value.description.substring(0, 30) + "..." || "", - checkpointFilePath: m.value.checkpointFilePath || "", - eventTypes: m.value.eventTypes || [], - imagerySource: m.value.imagerySource || "" - })); - }, [cataloguedModels]); + const options = React.useMemo( + () => normalizeBaseModelOptions(cataloguedModels), + [cataloguedModels] + ); const selectedOption = options.find( (o) => String(o.key) === String(baseModelId) @@ -78,11 +75,8 @@ function BaseModelDropdown({ const handleOptionSelect = (_ev, data) => { const picked = options.find((o) => String(o.key) === data.optionValue); - onFormChange( - picked ? picked.checkpointFilePath : "", - "initialWeightsUrl", - setComponentState, - componentState + setComponentState((currentState) => + applyBaseModelSelection(currentState, picked) ); }; @@ -106,4 +100,13 @@ function BaseModelDropdown({ ); } -export default BaseModelDropdown; \ No newline at end of file +BaseModelDropdown.propTypes = { + componentState: PropTypes.shape({ + baseModelId: PropTypes.oneOfType([PropTypes.string, PropTypes.number]), + baseModelIdError: PropTypes.string, + cataloguedModels: PropTypes.array, + }).isRequired, + setComponentState: PropTypes.func.isRequired, +}; + +export default BaseModelDropdown; diff --git a/ui/src/Components/BaseModelDropdownHelper.js b/ui/src/Components/BaseModelDropdownHelper.js new file mode 100644 index 00000000..ec92664a --- /dev/null +++ b/ui/src/Components/BaseModelDropdownHelper.js @@ -0,0 +1,58 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +export function buildModelCatalogEndpoint(imageLayer = {}, eventTypes) { + const params = new URLSearchParams(); + const definedEventTypes = Array.isArray(eventTypes) + ? eventTypes.filter(Boolean) + : []; + + if (definedEventTypes.length > 0) { + params.set("eventTypes", definedEventTypes.join(",")); + } + + const imagerySource = imageLayer?.sourceTypePostEvent; + if (imagerySource !== "" && imagerySource != null) { + params.set("imagerySource", imagerySource); + } + + const query = params.toString(); + return query ? `GetModelCatalog?${query}` : "GetModelCatalog"; +} + +export function normalizeBaseModelOptions(cataloguedModels = []) { + return cataloguedModels.map((model) => { + const value = model?.value ?? {}; + const baseModelName = value.baseModelName || ""; + const description = value.description == null + ? "" + : String(value.description); + + return { + key: model?.key || baseModelName, + baseModelName, + description: description ? `${description.substring(0, 30)}...` : "", + checkpointFilePath: value.checkpointFilePath || "", + eventTypes: value.eventTypes || [], + imagerySource: value.imagerySource || "", + }; + }); +} + +export function resolveBaseModelId(cataloguedModels, initialWeightsUrl) { + if (!initialWeightsUrl) return ""; + + const matchingOption = normalizeBaseModelOptions(cataloguedModels).find( + (option) => option.checkpointFilePath === initialWeightsUrl + ); + return matchingOption?.key || ""; +} + +export function applyBaseModelSelection(componentState, selectedOption) { + return { + ...componentState, + baseModelId: selectedOption?.key || "", + baseModelIdError: "", + initialWeightsUrl: selectedOption?.checkpointFilePath || "", + }; +} diff --git a/ui/src/Components/BaseModelDropdownHelper.test.js b/ui/src/Components/BaseModelDropdownHelper.test.js new file mode 100644 index 00000000..4a667bdb --- /dev/null +++ b/ui/src/Components/BaseModelDropdownHelper.test.js @@ -0,0 +1,101 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + applyBaseModelSelection, + buildModelCatalogEndpoint, + normalizeBaseModelOptions, + resolveBaseModelId, +} from "./BaseModelDropdownHelper.js"; + +test("builds a catalog query with event types and imagery source", () => { + const endpoint = buildModelCatalogEndpoint( + { sourceTypePostEvent: "Planet" }, + ["Hurricane", "Flood"] + ); + + assert.equal( + endpoint, + "GetModelCatalog?eventTypes=Hurricane%2CFlood&imagerySource=Planet" + ); +}); + +test("omits absent catalog filters instead of stringifying them", () => { + assert.equal(buildModelCatalogEndpoint({}, undefined), "GetModelCatalog"); + assert.equal( + buildModelCatalogEndpoint( + { sourceTypePostEvent: "Planet" }, + undefined + ), + "GetModelCatalog?imagerySource=Planet" + ); +}); + +test("encodes catalog filter values", () => { + assert.equal( + buildModelCatalogEndpoint( + { sourceTypePostEvent: "World View" }, + ["Severe Storm"] + ), + "GetModelCatalog?eventTypes=Severe+Storm&imagerySource=World+View" + ); +}); + +test("normalizes null descriptions and uses model names as fallback keys", () => { + const options = normalizeBaseModelOptions([ + { + value: { + baseModelName: "External checkpoint A", + description: null, + checkpointFilePath: "models/external-a.pt", + }, + }, + { + value: { + baseModelName: "External checkpoint B", + description: null, + checkpointFilePath: "models/external-b.pt", + }, + }, + ]); + + assert.deepEqual( + options.map((option) => option.key), + ["External checkpoint A", "External checkpoint B"] + ); + assert.equal(options[0].description, ""); + assert.equal(options[1].description, ""); +}); + +test("resolves an existing checkpoint URL to its catalog key", () => { + const cataloguedModels = [ + { + key: "model-1", + value: { + baseModelName: "Base model", + checkpointFilePath: "models/base.pt", + }, + }, + ]; + + assert.equal( + resolveBaseModelId(cataloguedModels, "models/base.pt"), + "model-1" + ); + assert.equal(resolveBaseModelId(cataloguedModels, "models/other.pt"), ""); +}); + +test("applies the selected model id and checkpoint in one state update", () => { + const currentState = { name: "Training model", baseModelIdError: "Required" }; + const selectedOption = { + key: "model-1", + checkpointFilePath: "models/base.pt", + }; + + assert.deepEqual(applyBaseModelSelection(currentState, selectedOption), { + name: "Training model", + baseModelId: "model-1", + baseModelIdError: "", + initialWeightsUrl: "models/base.pt", + }); +}); diff --git a/ui/src/Components/CreateEditModelTrainingHelper.js b/ui/src/Components/CreateEditModelTrainingHelper.js index bd427a1f..e71685d4 100644 --- a/ui/src/Components/CreateEditModelTrainingHelper.js +++ b/ui/src/Components/CreateEditModelTrainingHelper.js @@ -2,26 +2,16 @@ // Licensed under the MIT License. import { apiGet } from "../util/api"; +import { buildModelCatalogEndpoint } from "./BaseModelDropdownHelper"; export async function fetchModelCatalog(imageLayer, eventTypes) { const cataloguedModels = []; try { - let eventTypesL = "eventTypes=" + eventTypes; - - if ( - imageLayer.sourceTypePostEvent !== "" && - imageLayer.sourceTypePostEvent !== null && - imageLayer.sourceTypePostEvent !== undefined - ) { - const concatChar = eventTypesL === "" ? "" : "&"; - eventTypesL += concatChar + "imagerySource=" + imageLayer.sourceTypePostEvent; - } - - await apiGet(`GetModelCatalog?${eventTypesL}`) + await apiGet(buildModelCatalogEndpoint(imageLayer, eventTypes)) .then((response) => { cataloguedModels.push( ...response.modelCatalog.map((model) => ({ - key: model.modelId, + key: model.modelId || model.baseModelName, text: model.baseModelName, value: model, })) @@ -65,6 +55,8 @@ export function createComponentDefaultState(modelToEdit, imageLayer, projectId) batchSizeError: "", maxEpochs: "3", maxEpochsError: "", + baseModelId: "", + baseModelIdError: "", initialWeightsUrl: "", cataloguedModels: [], catalogLoading: true, diff --git a/ui/src/Components/CreateEditModelTrainingModal.jsx b/ui/src/Components/CreateEditModelTrainingModal.jsx index 9599f6cd..b2fe7640 100644 --- a/ui/src/Components/CreateEditModelTrainingModal.jsx +++ b/ui/src/Components/CreateEditModelTrainingModal.jsx @@ -18,6 +18,7 @@ import { validateEmptyOrInvalid, validateInt, validateFloat } from "../util/vali import { useNavigate } from "react-router-dom"; import { initGuidedTourState, setGuidedTourState } from "./GuidedTourHelper"; import BaseModelDropdown from "./BaseModelDropdown"; +import { resolveBaseModelId } from "./BaseModelDropdownHelper"; import { createComponentDefaultState, @@ -65,8 +66,14 @@ const CreateEditModelTrainingModal = ({ projectId ); const cataloguedModels = await fetchModelCatalog(imageLayer, eventTypes); + const baseModelId = resolveBaseModelId( + cataloguedModels, + baseState.initialWeightsUrl + ); setComponentState({ ...baseState, + baseModelId, + baseModelIdError: "", cataloguedModels, catalogLoading: false, }); @@ -227,7 +234,6 @@ const CreateEditModelTrainingModal = ({ diff --git a/ui/src/Components/LabelingTool/LabelingTool.jsx b/ui/src/Components/LabelingTool/LabelingTool.jsx index 121da742..ab6e4d9b 100644 --- a/ui/src/Components/LabelingTool/LabelingTool.jsx +++ b/ui/src/Components/LabelingTool/LabelingTool.jsx @@ -48,6 +48,8 @@ const LabelingTool = ({ setModalComponent }) => { const [isMapReady, setIsMapReady] = useState(false); const [selectedShape, setSelectedShape] = useState(null); const [hasUnsavedChanges, setHasUnsavedChanges] = useState(false); + const [eventTypes, setEventTypes] = useState([]); + const [imageLayer, setImageLayer] = useState(null); const { undo, redo } = useDrawingUndoRedo(drawingManager, mapRef); useEffect(() => { @@ -181,6 +183,12 @@ const LabelingTool = ({ setModalComponent }) => { const projectDetails = await apiGet( "GetProjectDetails?projectId=" + projectId ); + setEventTypes(projectDetails.eventTypes || []); + setImageLayer( + projectDetails.imageLayer?.find( + (layer) => layer.imageLayerId === imageLayerId + ) || null + ); const map = new window.atlas.Map(mapRef.current, { center: [0, 0], @@ -310,6 +318,8 @@ const LabelingTool = ({ setModalComponent }) => { setDrawingCount={setDrawingCount} selectedShape={selectedShape} imageLayerId={imageLayerId} + imageLayer={imageLayer} + eventTypes={eventTypes} undo={undo} redo={redo} /> diff --git a/ui/src/Components/LabelingTool/LabelingToolRightPanel.jsx b/ui/src/Components/LabelingTool/LabelingToolRightPanel.jsx index f027ebad..9ca25394 100644 --- a/ui/src/Components/LabelingTool/LabelingToolRightPanel.jsx +++ b/ui/src/Components/LabelingTool/LabelingToolRightPanel.jsx @@ -37,6 +37,8 @@ const LabelingToolRightPanel = ({ setDrawingCount, selectedShape, imageLayerId, + imageLayer, + eventTypes, undo, redo, }) => { @@ -57,6 +59,8 @@ const LabelingToolRightPanel = ({ setDrawingCount: PropType.func.isRequired, selectedShape: PropType.object, imageLayerId: PropType.string.isRequired, + imageLayer: PropType.object.isRequired, + eventTypes: PropType.array.isRequired, undo: PropType.func.isRequired, redo: PropType.func.isRequired, }; @@ -162,9 +166,10 @@ const LabelingToolRightPanel = ({ setModalComponent(null)} projectId={projectId} - imageLayer={labelingToolDataRef.current} + imageLayer={imageLayer} guidedTour="createEditModelTrainingModalGuide" autoLaunchGuidedTour={true} + eventTypes={eventTypes} /> ); } diff --git a/ui/src/Components/LabelingTool/LeftPanel.jsx b/ui/src/Components/LabelingTool/LeftPanel.jsx deleted file mode 100644 index 40330df3..00000000 --- a/ui/src/Components/LabelingTool/LeftPanel.jsx +++ /dev/null @@ -1,317 +0,0 @@ -import { - Button, - Slider, - Switch, - Label, - Field, -} from "@fluentui/react-components"; -import { FluentIcon } from "../../util/icons"; -import { useState, useEffect } from "react"; -import { saveLabels } from "./LabelingToolHelper"; - -import PropType from "prop-types"; -import { useNavigate } from "react-router-dom"; - -const LeftPanel = ({ - mapRef, - drawingCount, - preImageryRef, - postImageryRef, - hasUnsavedChanges, - setDialog, - setIsLoading, - drawingManager, - imageLayerId, - labelingToolDataRef, - setHasUnsavedChanges, -}) => { - LeftPanel.propTypes = { - mapRef: PropType.object.isRequired, - drawingCount: PropType.number.isRequired, - preImageryRef: PropType.object.isRequired, - postImageryRef: PropType.object.isRequired, - hasUnsavedChanges: PropType.bool.isRequired, - setDialog: PropType.func.isRequired, - setIsLoading: PropType.func.isRequired, - drawingManager: PropType.object.isRequired, - imageLayerId: PropType.string.isRequired, - labelingToolDataRef: PropType.object.isRequired, - setHasUnsavedChanges: PropType.func.isRequired, - }; - - const [eventImageryVisibilityState, setEventImageryVisibilityState] = - useState(true); - - const [imageryValues, setImageryValues] = useState({ - opacity: 1, - contrast: 0, - hueRotation: 0, - saturation: 0, - }); - - const updateValues = (key, value) => { - try { - switch (eventImageryVisibilityState) { - case true: - postImageryRef.current.setOptions({ - [key]: value, - }); - break; - case false: - if (preImageryRef.current == null) return; - preImageryRef.current.setOptions({ - [key]: value, - }); - break; - default: - break; - } - - setImageryValues({ - ...imageryValues, - [key]: value, - }); - } catch (error) { - console.error("Error updating imagery values:", error); - } - }; - - const resetControls = () => { - try { - setImageryValues({ - opacity: 1, - contrast: 0, - hueRotation: 0, - saturation: 0, - }); - - switch (eventImageryVisibilityState) { - case true: - postImageryRef.current.setOptions({ - opacity: 1, - contrast: 0, - hueRotation: 0, - saturation: 0, - }); - break; - case false: - preImageryRef.current.setOptions({ - opacity: 1, - contrast: 0, - hueRotation: 0, - saturation: 0, - }); - break; - default: - break; - } - } catch (error) { - console.error("Error resetting controls:", error); - } - }; - - function getLayerById(currentMap, customId) { - const layers = currentMap.current.layers.getLayers(); - return layers.find((layer) => layer.customId === customId); - } - - function togglePostEventLayerVisibility(customId, isVisible) { - const layer = getLayerById(mapRef, customId); - if (layer) { - layer.setOptions({ visible: isVisible }); - } - } - - useEffect(() => { - if (preImageryRef.current || postImageryRef.current) { - try { - - // Layer visibility toggle - if (preImageryRef.current !== null) { - togglePostEventLayerVisibility( - "preEventImageryLayer", - !eventImageryVisibilityState - ); - } - - togglePostEventLayerVisibility( - "postEventImageryLayer", - eventImageryVisibilityState - ); - - // Set imagery options based on the current state. - if (eventImageryVisibilityState) { - postImageryRef.current.setOptions({ - opacity: imageryValues.opacity, - contrast: imageryValues.contrast, - hueRotation: imageryValues.hueRotation, - saturation: imageryValues.saturation, - }); - } else { - if (preImageryRef.current !== null) { - preImageryRef.current.setOptions({ - opacity: imageryValues.opacity, - contrast: imageryValues.contrast, - hueRotation: imageryValues.hueRotation, - saturation: imageryValues.saturation, - }); - } - } - } catch (error) { - console.error("Error toggling layer visibility:", error); - } - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [eventImageryVisibilityState]); - - const navigate = useNavigate(); - - const handleBackNavigation = () => { - if (hasUnsavedChanges) { - setDialog("Important", `Do you want to save changes before leaving?`, [ - { - type: "primary", - key: "yes", - text: "Yes", - onClick: saveAndLeave, - }, - { - type: "default", - key: "no", - text: "No", - onClick: () => (setDialog(), navigate(-1)), - }, - { - type: "default", - key: "cancel", - text: "Cancel", - onClick: () => setDialog(), - }, - ]); - } else { - navigate(-1); - } - }; - - const saveAndLeave = async () => { - setDialog(); - const isSaved = await saveLabels( - drawingManager, - labelingToolDataRef, - setIsLoading, - setHasUnsavedChanges - ); - if (isSaved) { - navigate(-1); - } - }; - - return ( - <> -
-
- - -
- - updateValues("opacity", data.value)} - value={imageryValues.opacity} - /> - - - - updateValues("contrast", data.value)} - value={imageryValues.contrast} - /> - - - - updateValues("hueRotation", data.value)} - value={imageryValues.hueRotation} - /> - - - - updateValues("saturation", data.value)} - value={imageryValues.saturation} - /> - - - - -
- - - setEventImageryVisibilityState(data.checked) - } - /> -
-
-
-
- -
- Number of labels: {drawingCount} -
- - ); -}; - -export default LeftPanel; diff --git a/ui/src/Components/LabelingTool/RightPanel.jsx b/ui/src/Components/LabelingTool/RightPanel.jsx deleted file mode 100644 index 9b7102e1..00000000 --- a/ui/src/Components/LabelingTool/RightPanel.jsx +++ /dev/null @@ -1,247 +0,0 @@ -import { useEffect, useState } from "react"; -import { - RadioGroup, - Radio, - Field, - SplitButton, - Menu, - MenuTrigger, - MenuPopover, - MenuList, - MenuItem, -} from "@fluentui/react-components"; -import { FluentIcon } from "../../util/icons"; -import CreateEditModelTrainingModal from "../CreateEditModelTrainingModal"; -import { saveLabels } from "./LabelingToolHelper"; -import DrawingToolbar from "./DrawingToolbar"; -import PropType from "prop-types"; -import { updateShape } from "./LabelingToolHelper"; - -const RightPanel = ({ - primaryClasses, - selectedPrimaryClass, - setSelectedPrimaryClass, - drawingManager, - setDrawingManager, - setDialog, - labelingToolDataRef, - setIsLoading, - setHasUnsavedChanges, - setModalComponent, - projectId, - drawingCount, - selectedShape, -}) => { - RightPanel.propTypes = { - primaryClasses: PropType.array.isRequired, - selectedPrimaryClass: PropType.string.isRequired, - setSelectedPrimaryClass: PropType.func.isRequired, - drawingManager: PropType.object.isRequired, - setDrawingManager: PropType.func.isRequired, - setDialog: PropType.func.isRequired, - labelingToolDataRef: PropType.object.isRequired, - setIsLoading: PropType.func.isRequired, - setHasUnsavedChanges: PropType.func.isRequired, - setModalComponent: PropType.func.isRequired, - projectId: PropType.string.isRequired, - drawingCount: PropType.number.isRequired, - selectedShape: PropType.object, - }; - - useEffect(() => { - if (selectedShape !== null) { - updateShape( - drawingManager, - selectedShape, - selectedPrimaryClass, - ); - setHasUnsavedChanges(true); - } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [selectedPrimaryClass]); - - async function handleSave() { - const isSaved = await saveLabels( - drawingManager, - labelingToolDataRef, - setIsLoading, - setHasUnsavedChanges - ); - - var message = "Labels saved successfully"; - - if (!isSaved) { - message = "There was an error saving the labels"; - } - - const buttons = [ - { - type: "primary", - key: "close", - text: "Close", - onClick: () => { - setDialog("", "", []); - }, - }, - ]; - - setDialog("Important", message, buttons); - } - - function handlePrimaryClassChange(newClass) { - if(selectedShape !== null) { - const buttons = [ - { - type: "primary", - key: "yes", - text: "Yes", - onClick: () => { - setSelectedPrimaryClass(newClass); - setDialog("", "", []); - setTimeout(() => { - document.getElementById("pointer").click(); - }, 100); - } - }, - { - type: "default", - key: "no", - text: "No", - onClick: () => { - setDialog("", "", []); - }, - }, - ]; - setDialog("Important", "You are about to change the primary class of the selected label. Do you want to continue?", buttons); - }else{ - setSelectedPrimaryClass(newClass); - } - } - - async function handleSaveAndTrain() { - const isSaved = await saveLabels( - drawingManager, - labelingToolDataRef, - setIsLoading, - setHasUnsavedChanges - ); - - if (!isSaved) { - const buttons = [ - { - type: "primary", - key: "close", - text: "Close", - onClick: () => { - setDialog("", "", []); - }, - }, - ]; - - setDialog("Important", "There was an error saving the labels", buttons); - } else { - setModalComponent( - setModalComponent(null)} - projectId={projectId} - imageLayer={labelingToolDataRef.current} - guidedTour="createEditModelTrainingModalGuide" - autoLaunchGuidedTour={true} - /> - ); - } - } - - const labelSavingMenuOptions = () => ({ - items: [ - { - key: "saveAndTrain", - text: "Save and Train", - iconProps: { iconName: "SaveAndClose" }, - title: drawingCount === 0 && "Training requires at least one label", - disabled: drawingCount === 0, - onClick: () => { - handleSaveAndTrain(); - }, - }, - ], - }); - - return ( - <> - -
-
- - { - handlePrimaryClassChange(data.value); - }} - > - {(primaryClasses && primaryClasses.length > 0 - ? primaryClasses - : [] - ).map((option) => ( - - ))} - - -
-
- - - {(triggerProps) => ( - } - menuButton={triggerProps} - primaryActionButton={{ - onClick: async () => { - await handleSave(); - }, - }} - > - Save - - )} - - - - {labelSavingMenuOptions().items.map((item) => ( - } - disabled={item.disabled} - onClick={item.onClick} - > - {item.text} - - ))} - - - -
-
- - ); -}; - -export default RightPanel; From d52ab4e4264394f4bf5fb4c4732f3c310fd4f280 Mon Sep 17 00:00:00 2001 From: Marcelo Duarte Date: Mon, 31 Aug 2026 16:51:12 -0300 Subject: [PATCH 2/2] fix(ui): prevent base model dropdown key collisions --- ui/src/Components/BaseModelDropdownHelper.js | 8 +++++- .../BaseModelDropdownHelper.test.js | 28 +++++++++++++++---- .../CreateEditModelTrainingHelper.js | 7 +++-- 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/ui/src/Components/BaseModelDropdownHelper.js b/ui/src/Components/BaseModelDropdownHelper.js index ec92664a..9cb6500e 100644 --- a/ui/src/Components/BaseModelDropdownHelper.js +++ b/ui/src/Components/BaseModelDropdownHelper.js @@ -20,6 +20,12 @@ export function buildModelCatalogEndpoint(imageLayer = {}, eventTypes) { return query ? `GetModelCatalog?${query}` : "GetModelCatalog"; } +export function buildBaseModelOptionKey(model) { + return model?.modelId + ? `modelId:${model.modelId}` + : `baseModelName:${model?.baseModelName || ""}`; +} + export function normalizeBaseModelOptions(cataloguedModels = []) { return cataloguedModels.map((model) => { const value = model?.value ?? {}; @@ -29,7 +35,7 @@ export function normalizeBaseModelOptions(cataloguedModels = []) { : String(value.description); return { - key: model?.key || baseModelName, + key: model?.key || `baseModelName:${baseModelName}`, baseModelName, description: description ? `${description.substring(0, 30)}...` : "", checkpointFilePath: value.checkpointFilePath || "", diff --git a/ui/src/Components/BaseModelDropdownHelper.test.js b/ui/src/Components/BaseModelDropdownHelper.test.js index 4a667bdb..ff6baf06 100644 --- a/ui/src/Components/BaseModelDropdownHelper.test.js +++ b/ui/src/Components/BaseModelDropdownHelper.test.js @@ -3,6 +3,7 @@ import assert from "node:assert/strict"; import { applyBaseModelSelection, + buildBaseModelOptionKey, buildModelCatalogEndpoint, normalizeBaseModelOptions, resolveBaseModelId, @@ -41,6 +42,20 @@ test("encodes catalog filter values", () => { ); }); +test("keeps model ID and fallback name keys in disjoint namespaces", () => { + const hasteModelKey = buildBaseModelOptionKey({ + modelId: "3516", + baseModelName: "HASTE model", + }); + const externalModelKey = buildBaseModelOptionKey({ + baseModelName: "3516", + }); + + assert.equal(hasteModelKey, "modelId:3516"); + assert.equal(externalModelKey, "baseModelName:3516"); + assert.notEqual(hasteModelKey, externalModelKey); +}); + test("normalizes null descriptions and uses model names as fallback keys", () => { const options = normalizeBaseModelOptions([ { @@ -61,7 +76,10 @@ test("normalizes null descriptions and uses model names as fallback keys", () => assert.deepEqual( options.map((option) => option.key), - ["External checkpoint A", "External checkpoint B"] + [ + "baseModelName:External checkpoint A", + "baseModelName:External checkpoint B", + ] ); assert.equal(options[0].description, ""); assert.equal(options[1].description, ""); @@ -70,7 +88,7 @@ test("normalizes null descriptions and uses model names as fallback keys", () => test("resolves an existing checkpoint URL to its catalog key", () => { const cataloguedModels = [ { - key: "model-1", + key: "modelId:model-1", value: { baseModelName: "Base model", checkpointFilePath: "models/base.pt", @@ -80,7 +98,7 @@ test("resolves an existing checkpoint URL to its catalog key", () => { assert.equal( resolveBaseModelId(cataloguedModels, "models/base.pt"), - "model-1" + "modelId:model-1" ); assert.equal(resolveBaseModelId(cataloguedModels, "models/other.pt"), ""); }); @@ -88,13 +106,13 @@ test("resolves an existing checkpoint URL to its catalog key", () => { test("applies the selected model id and checkpoint in one state update", () => { const currentState = { name: "Training model", baseModelIdError: "Required" }; const selectedOption = { - key: "model-1", + key: "modelId:model-1", checkpointFilePath: "models/base.pt", }; assert.deepEqual(applyBaseModelSelection(currentState, selectedOption), { name: "Training model", - baseModelId: "model-1", + baseModelId: "modelId:model-1", baseModelIdError: "", initialWeightsUrl: "models/base.pt", }); diff --git a/ui/src/Components/CreateEditModelTrainingHelper.js b/ui/src/Components/CreateEditModelTrainingHelper.js index e71685d4..81893793 100644 --- a/ui/src/Components/CreateEditModelTrainingHelper.js +++ b/ui/src/Components/CreateEditModelTrainingHelper.js @@ -2,7 +2,10 @@ // Licensed under the MIT License. import { apiGet } from "../util/api"; -import { buildModelCatalogEndpoint } from "./BaseModelDropdownHelper"; +import { + buildBaseModelOptionKey, + buildModelCatalogEndpoint, +} from "./BaseModelDropdownHelper"; export async function fetchModelCatalog(imageLayer, eventTypes) { const cataloguedModels = []; @@ -11,7 +14,7 @@ export async function fetchModelCatalog(imageLayer, eventTypes) { .then((response) => { cataloguedModels.push( ...response.modelCatalog.map((model) => ({ - key: model.modelId || model.baseModelName, + key: buildBaseModelOptionKey(model), text: model.baseModelName, value: model, }))