diff --git a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx index c8f10f6dd8..57f8aa22f6 100644 --- a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx +++ b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx @@ -5,6 +5,8 @@ import { EnqueueImageInpaint, EnqueueEditImageSize, EnqueueEditImageResolution, + EnqueueEditImageRequest, + EnqueueImageInpaintRequest, } from "@storyteller/tauri-api"; import { PromptsApi } from "@storyteller/api"; import { RefImage, usePromptEditStore } from "@storyteller/ui-promptbox"; @@ -21,6 +23,7 @@ import MarkerToolControlBar from "./MarkerToolControlBar"; import { ClassyModelSelector, useSelectedImageModel, + useSelectedProviderForModel, IMAGE_EDITOR_PAGE_MODEL_LIST, ModelPage, //ProviderSelector, @@ -29,6 +32,8 @@ import { import { ImageModel } from "@storyteller/model-list"; import { HistoryStack, ImageBundle } from "./HistoryStack"; import { TutorialModalButton } from "@storyteller/ui-tutorial-modal"; +import { GenerationProvider } from "@storyteller/api-enums"; + const PAGE_ID: ModelPage = ModelPage.ImageEditor; @@ -38,6 +43,9 @@ const PageEdit = () => { const selectedImageModel: ImageModel | undefined = useSelectedImageModel(PAGE_ID); + const selectedProvider : GenerationProvider | undefined = + useSelectedProviderForModel(PAGE_ID, selectedImageModel?.id); + // State for canvas dimensions const canvasWidth = useRef(1024); const canvasHeight = useRef(1024); @@ -322,6 +330,7 @@ const PageEdit = () => { aspectRatio?: string; resolution?: string; images?: RefImage[]; + selectedProvider?: GenerationProvider; }, ) => { const editedImageToken = store.baseImageInfo?.mediaToken; @@ -374,7 +383,7 @@ const PageEdit = () => { if (selectedImageModel?.editingIsInpainting) { const arrayBuffer = await getMaskArrayBuffer(); - result = await EnqueueImageInpaint({ + let request : EnqueueImageInpaintRequest = { model: selectedImageModel, image_media_token: editedImageToken, mask_image_raw_bytes: arrayBuffer, @@ -382,7 +391,11 @@ const PageEdit = () => { image_count: generationCount, frontend_caller: "image_editor", frontend_subscriber_id: subscriberId, - }); + }; + if (!!options?.selectedProvider) { + request.provider = options.selectedProvider; + } + result = await EnqueueImageInpaint(request); } else if (selectedImageModel?.isNanoBananaModel()) { const compositeFile = await getCompositeCanvasFile(); if (!compositeFile) { @@ -398,7 +411,7 @@ const PageEdit = () => { return; } const imgs = options?.images || []; - result = await EnqueueEditImage({ + let request: EnqueueEditImageRequest = { model: selectedImageModel, scene_image_media_token: snapshotResult.data, image_media_tokens: imgs @@ -410,10 +423,14 @@ const PageEdit = () => { frontend_subscriber_id: subscriberId, aspect_ratio: mapAspectRatio(options?.aspectRatio), image_resolution: mapResolution(options?.resolution), - }); + }; + if (!!options?.selectedProvider) { + request.provider = options.selectedProvider; + } + result = await EnqueueEditImage(request); } else { const imgs = options?.images || []; - result = await EnqueueEditImage({ + let request : EnqueueEditImageRequest = { model: selectedImageModel, image_media_tokens: [ editedImageToken, @@ -428,10 +445,14 @@ const PageEdit = () => { frontend_subscriber_id: subscriberId, aspect_ratio: mapAspectRatio(options?.aspectRatio), image_resolution: mapResolution(options?.resolution), - }); + }; + if (!!options?.selectedProvider) { + request.provider = options.selectedProvider; + } + result = await EnqueueEditImage(request); } - if (result.status === "success") { + if (result?.status === "success") { setPendingGenerations((prev) => [ ...prev, { id: subscriberId as string, count: generationCount }, @@ -609,6 +630,7 @@ const PageEdit = () => { > { store.setActiveTool(mode as ActiveEditTool); }} diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx index 63931acad7..bddd273154 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx @@ -29,6 +29,7 @@ import { ImageModel, getCapabilitiesForModel } from "@storyteller/model-list"; import { twMerge } from "tailwind-merge"; import { ImagePromptRow, type UploadImageFn } from "./ImagePromptRow"; import { RefImage, usePromptEditStore } from "./promptStore"; +import { GenerationProvider } from "@storyteller/api-enums"; export interface PromptBoxEditProps { onModeChange?: (mode: string) => void; @@ -39,11 +40,13 @@ export interface PromptBoxEditProps { aspectRatio?: string; resolution?: string; images?: RefImage[]; + selectedProvider?: GenerationProvider; }, ) => void | Promise; isDisabled?: boolean; onFitPressed?: () => void | Promise; selectedImageModel?: ImageModel; + selectedProvider?: GenerationProvider; generationCount?: number; onGenerationCountChange?: (count: number) => void; supportsMaskedInpainting?: boolean; @@ -62,6 +65,7 @@ export const PromptBoxEdit = ({ isDisabled, onFitPressed, selectedImageModel, + selectedProvider, generationCount: generationCountProp, onGenerationCountChange, supportsMaskedInpainting = false, @@ -277,6 +281,7 @@ export const PromptBoxEdit = ({ aspectRatio, resolution, images: referenceImages, + selectedProvider: selectedProvider, }), ); } finally { diff --git a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueImageInpaint.ts b/frontend/libs/tauri-api/src/lib/enqueue/EnqueueImageInpaint.ts index 91c3c2c241..22b0eca42f 100644 --- a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueImageInpaint.ts +++ b/frontend/libs/tauri-api/src/lib/enqueue/EnqueueImageInpaint.ts @@ -1,11 +1,15 @@ import { invoke } from "@tauri-apps/api/core"; import { CommandResult } from "../common/CommandStatus"; import { ImageModel, ModelInfo } from "@storyteller/model-list"; +import { GenerationProvider } from "@storyteller/api-enums"; export interface EnqueueImageInpaintRequest { // The model to use. model?: ImageModel | ModelInfo | EnqueueImageInpaintModel; + // The provider to use. + provider?: GenerationProvider; + // The image we're editing. image_media_token?: string; @@ -30,6 +34,7 @@ export interface EnqueueImageInpaintRequest { // (We do some transformations from the public-facing request object.) interface RawEnqueueImageInpaintRequest { model?: EnqueueImageInpaintModel | string; + provider?: GenerationProvider; image_media_token?: string; mask_image_raw_bytes?: Uint8Array; prompt?: string; @@ -89,6 +94,10 @@ export const EnqueueImageInpaint = async (request: EnqueueImageInpaintRequest) : model: modelName, prompt: request.prompt, }; + + if (!!request.provider) { + mutableRequest.provider = request.provider; + } if (!!request.image_media_token) { mutableRequest.image_media_token = request.image_media_token;