diff --git a/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.spec.tsx b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.spec.tsx new file mode 100644 --- /dev/null +++ b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.spec.tsx @@ -0,0 +1,45 @@ +import { render, cleanup } from "@testing-library/react"; +import { afterEach, expect, it, vi } from "vitest"; +import { CommonAspectRatio } from "@storyteller/model-list"; +import type { PageDrawAdapter } from "@storyteller/ui-pagedraw"; +import PageDraw from "./PageDraw"; + +const mocks = vi.hoisted(() => ({ + adapter: undefined as PageDrawAdapter | undefined, + generate: vi.fn(async (_request: { aspect_ratio?: string }) => ({ status: "success" })), +})); +vi.mock("@storyteller/ui-pagedraw", () => ({ + PageDraw: ({ adapter }: { adapter: PageDrawAdapter }) => { + mocks.adapter = adapter; + return null; + }, + useSceneStore: { getState: vi.fn() }, +})); +vi.mock("@storyteller/tauri-api", () => ({ + GenerateImage: mocks.generate, + EnqueueImageBgRemoval: vi.fn(), + useCanvasBgRemovedEvent: vi.fn(), +})); +vi.mock("@storyteller/tauri-events", () => ({ useTextToImageGenerationCompleteEvent: vi.fn() })); +vi.mock("@storyteller/api", () => ({ UploadImageMedia: vi.fn() })); +vi.mock("./BaseImageSelector", () => ({ BaseImageSelector: () => null })); +afterEach(() => { cleanup(); mocks.generate.mockClear(); }); + +it.each(["enqueueEditImage", "enqueueInpaint"] as const)( + "%s sends the resolved source ratio to generation", async (method) => { + render(); + await mocks.adapter![method]({ commonAspectRatio: CommonAspectRatio.TallNineBySixteen }); + expect(mocks.generate).toHaveBeenCalledWith(expect.objectContaining({ + aspect_ratio: CommonAspectRatio.TallNineBySixteen, + })); + }, +); +it.each(["enqueueEditImage", "enqueueInpaint"] as const)( + "%s preserves explicit ratios and omits unsupported Auto", async (method) => { + render(); + await mocks.adapter![method]({ aspectRatio: "square" }); + expect(mocks.generate.mock.calls.at(-1)![0]).toMatchObject({ aspect_ratio: CommonAspectRatio.Square }); + await mocks.adapter![method]({}); + expect(mocks.generate.mock.calls.at(-1)![0].aspect_ratio).toBeUndefined(); + }, +); diff --git a/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx --- a/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx +++ b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx @@ -105,7 +105,7 @@ const useTauriAdapter = (): PageDrawAdapter => { ? !req.disableSystemPrompt : undefined, batch_size: req.imageCount, - aspect_ratio: mapAspectRatio(req.aspectRatio), + aspect_ratio: req.commonAspectRatio ?? mapAspectRatio(req.aspectRatio), resolution: mapResolution(req.imageResolution), frontend_caller: req.frontendCaller, frontend_subscriber_id: req.frontendSubscriberId, @@ -119,6 +119,7 @@ const useTauriAdapter = (): PageDrawAdapter => { model: req.model, image_media_tokens: req.imageMediaToken ? [req.imageMediaToken] : undefined, inpainting_mask_image_raw_bytes: req.maskImageRawBytes, + aspect_ratio: req.commonAspectRatio ?? mapAspectRatio(req.aspectRatio), prompt: req.prompt, batch_size: req.imageCount, frontend_caller: req.frontendCaller, diff --git a/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx b/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx --- a/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx +++ b/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx @@ -35,6 +35,7 @@ import { GenerationProvider } from "@storyteller/api-enums"; import { HistoryStack } from "./HistoryStack"; import { type BaseSelectorImage } from "./types"; import { EncodeImageBitmapToBase64 } from "./utilities/EncodeImageBitmapToBase64"; +import { editAspectRatioFields } from "./utilities/editAspectRatio"; import { compositeInWorker, maskInWorker } from "./utilities/generatePipeline"; import { RefImage, usePrompt2DStore } from "@storyteller/ui-promptbox"; import { PromptsApi } from "@storyteller/api"; @@ -406,6 +407,11 @@ const PageDraw = ({ redo, } = useSceneStore(useShallow(selector)); + const sourceAspectRatioType = useSceneStore((state) => state.aspectRatioType); + useEffect(() => { + usePrompt2DStore.setState({ sourceImageDimensions: getAspectRatioDimensions() }); + }, [baseImageBitmap, baseImageInfo, sourceAspectRatioType, getAspectRatioDimensions]); + const promptStoreProvider = usePrompt2DStore; const generationCount = promptStoreProvider((state) => state.generationCount); const setGenerationCount = promptStoreProvider( @@ -869,6 +875,12 @@ const PageDraw = ({ timestamp: Date.now(), }); + const aspectRatioFields = editAspectRatioFields( + options?.aspectRatio, + selectedImageModel, + getAspectRatioDimensions(), + ); + let result; if (selectedImageModel?.editingIsInpainting) { @@ -878,6 +890,7 @@ const PageDraw = ({ model: selectedImageModel, imageMediaToken: editedImageToken, maskImageRawBytes: arrayBuffer, + ...aspectRatioFields, prompt: prompt, imageCount: generationCount, frontendCaller: "image_editor", @@ -914,7 +927,7 @@ const PageDraw = ({ imageCount: generationCount, frontendCaller: "image_editor", frontendSubscriberId: subscriberId, - aspectRatio: options?.aspectRatio, + ...aspectRatioFields, imageResolution: options?.resolution, provider: options?.selectedProvider, }); @@ -948,7 +961,7 @@ const PageDraw = ({ imageCount: generationCount, frontendCaller: "image_editor", frontendSubscriberId: subscriberId, - aspectRatio: options?.aspectRatio, + ...aspectRatioFields, imageResolution: options?.resolution, disableSystemPrompt: !useSystemPrompt, provider: options?.selectedProvider, @@ -966,6 +979,7 @@ const PageDraw = ({ generationCount, getCompositeFile, getMaskArrayBuffer, + getAspectRatioDimensions, selectedImageModel, adapter, baseImageBitmap, diff --git a/frontend/libs/components/pagedraw/src/lib/adapter.ts b/frontend/libs/components/pagedraw/src/lib/adapter.ts --- a/frontend/libs/components/pagedraw/src/lib/adapter.ts +++ b/frontend/libs/components/pagedraw/src/lib/adapter.ts @@ -1,5 +1,5 @@ import type { ReactNode } from "react"; -import type { ImageModel } from "@storyteller/model-list"; +import type { CommonAspectRatio, ImageModel } from "@storyteller/model-list"; import type { GenerationProvider } from "@storyteller/api-enums"; import type { UploadMediaFn } from "@storyteller/api"; import type { BaseSelectorImage } from "./types"; @@ -12,15 +12,21 @@ export interface PageDrawEditRequest { prompt?: string; disableSystemPrompt?: boolean; imageCount?: number; - /** "auto" | "wide" | "tall" | "square" */ + /** "wide" | "tall" | "square". Unset when the picker is on "auto"; see `commonAspectRatio`. */ aspectRatio?: string; + /** + * Set instead of `aspectRatio` when the picker is on "auto": the value for + * `model` that keeps the base image's aspect ratio. Hosts that do not read + * it send no aspect ratio, which leaves the choice to the model. + */ + commonAspectRatio?: CommonAspectRatio; /** "1k" | "2k" | "4k" */ imageResolution?: string; frontendCaller?: string; frontendSubscriberId?: string; } -export interface PageDrawInpaintRequest { +export interface PageDrawInpaintRequest extends Pick { model?: ImageModel; provider?: GenerationProvider; imageMediaToken?: string; diff --git a/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.spec.ts b/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.spec.ts new file mode 100644 --- /dev/null +++ b/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.spec.ts @@ -0,0 +1,72 @@ +import { describe, expect, it } from "vitest"; +import { + buildImageModelsFromListing, + CommonAspectRatio, +} from "@storyteller/model-list"; +import { editAspectRatioFields } from "./editAspectRatio"; + +const LANDSCAPE = { width: 1600, height: 1200 }; + +const [NANO_BANANA_PRO, GPT_IMAGE_1P5, FUTURE_MODEL] = buildImageModelsFromListing( + [], + [ + { + model: "nano_banana_pro", + aspect_ratio_options: ["auto", "square", "wide_four_by_three"], + }, + { + model: "gpt_image_1p5", + aspect_ratio_options: ["square", "wide", "tall"], + }, + { model: "future_edit_model" }, + ], +); + +describe("editAspectRatioFields", () => { + describe("a ratio picked by the user", () => { + it.each(["wide", "tall", "square"])("is sent as picked (%s)", (picked) => { + const fields = editAspectRatioFields(picked, NANO_BANANA_PRO, LANDSCAPE); + expect(fields.aspectRatio).toBe(picked); + expect(fields.commonAspectRatio).toBeUndefined(); + }); + + it("is left unset when nothing is picked", () => { + const fields = editAspectRatioFields(undefined, NANO_BANANA_PRO, LANDSCAPE); + expect(fields.aspectRatio).toBeUndefined(); + expect(fields.commonAspectRatio).toBeUndefined(); + }); + }); + + describe("auto", () => { + it("asks a model that has an auto option for it", () => { + const fields = editAspectRatioFields("auto", NANO_BANANA_PRO, LANDSCAPE); + expect(fields.commonAspectRatio).toBe(CommonAspectRatio.Auto); + }); + + it("asks other models for the listed ratio closest to the base image", () => { + expect( + editAspectRatioFields("auto", GPT_IMAGE_1P5, { width: 1920, height: 1080 }) + .commonAspectRatio, + ).toBe(CommonAspectRatio.Wide); + expect( + editAspectRatioFields("auto", GPT_IMAGE_1P5, { width: 1080, height: 1920 }) + .commonAspectRatio, + ).toBe(CommonAspectRatio.Tall); + }); + + it("never sends the word auto itself", () => { + for (const model of [NANO_BANANA_PRO, GPT_IMAGE_1P5, FUTURE_MODEL, undefined]) { + expect(editAspectRatioFields("auto", model, LANDSCAPE).aspectRatio).toBeUndefined(); + } + }); + + it("sends no ratio when the model has no aspect ratio control", () => { + expect( + editAspectRatioFields("auto", FUTURE_MODEL, LANDSCAPE).commonAspectRatio, + ).toBeUndefined(); + expect( + editAspectRatioFields("auto", undefined, LANDSCAPE).commonAspectRatio, + ).toBeUndefined(); + }); + }); +}); diff --git a/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.ts b/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.ts new file mode 100644 --- /dev/null +++ b/frontend/libs/components/pagedraw/src/lib/utilities/editAspectRatio.ts @@ -0,0 +1,29 @@ +import { + matchSourceAspectRatio, + type ImageModel, +} from "@storyteller/model-list"; +import type { PageDrawEditRequest } from "../adapter"; + +/** + * The aspect ratio fields of an edit request for the value picked in the + * prompt box. + * + * "auto" means keep the base image's aspect ratio. Only the model's own + * options can express that, so it is resolved here, where the base image size + * is known, instead of being sent as the word "auto". + */ +export function editAspectRatioFields( + picked: string | undefined, + model: ImageModel | undefined, + baseImage: { width: number; height: number }, +): Pick { + if (picked !== "auto") return { aspectRatio: picked }; + + return { + commonAspectRatio: matchSourceAspectRatio( + model, + baseImage.width, + baseImage.height, + ), + }; +} diff --git a/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.spec.tsx b/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.spec.tsx new file mode 100644 --- /dev/null +++ b/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.spec.tsx @@ -0,0 +1,108 @@ +import { renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, expect, it, vi } from "vitest"; +import { ModelPage } from "@storyteller/ui-model-selector"; +import { + buildImageModelsFromListing, + type ImageModel, +} from "@storyteller/model-list"; +import { useImageCostEstimate } from "./useImageCostEstimate"; + +const mocks = vi.hoisted(() => ({ + estimate: vi.fn(), + setCredits: vi.fn(), + canvas2D: { + sourceImageDimensions: undefined as { width: number; height: number } | undefined, + aspectRatio: "auto", + resolution: "1k", + referenceImages: [], + generationCount: 1, + }, +})); + +vi.mock("@storyteller/ui-promptbox", () => ({ + usePromptImageStore: (select: any) => + select({ referenceImages: [], generationCount: 1 }), + usePrompt2DStore: (select: any) => select(mocks.canvas2D), + usePrompt3DStore: (select: any) => select({ resolution: "1k", referenceImages: [] }), + usePromptEditStore: (select: any) => + select({ aspectRatio: "auto", resolution: "1k", referenceImages: [] }), +})); +vi.mock("@storyteller/ui-model-selector", () => ({ + ModelPage: { + TextToImage: "text-to-image", + Canvas2D: "canvas-2d", + Stage3D: "stage-3d", + ImageEditor: "image-editor", + Angles: "angles", + }, +})); +vi.mock("@storyteller/tauri-api", () => ({ + EstimateImageCost: mocks.estimate, + isEstimateImageCostSuccess: (result: any) => result.status === "success", +})); +vi.mock("./cost-breakdown-modal-store", () => ({ + useCostBreakdownModalStore: (select: any) => + select({ setEstimatedCreditsForPage: mocks.setCredits }), +})); + +const [NANO_BANANA_PRO, SEEDREAM_4, GPT_IMAGE_1P5] = buildImageModelsFromListing([], [ + { model: "nano_banana_pro", aspect_ratio_options: ["auto", "square", "wide_sixteen_by_nine"] }, + { model: "seedream_4", aspect_ratio_options: ["auto_2k", "auto_4k", "square"] }, + { model: "gpt_image_1p5", aspect_ratio_options: ["square", "wide", "tall"] }, +]); + +beforeEach(() => { + mocks.canvas2D.aspectRatio = "auto"; + mocks.canvas2D.sourceImageDimensions = undefined; + mocks.estimate.mockReset(); + mocks.setCredits.mockReset(); + mocks.estimate.mockResolvedValue({ + status: "success", + payload: { cost_in_credits: 12 }, + }); +}); + +it.each([ + ["a model with an auto option", NANO_BANANA_PRO, "auto"], + ["a model whose auto option also picks a resolution", SEEDREAM_4, "auto_2k"], + ["a model with no auto option, which leaves it to the API", GPT_IMAGE_1P5, undefined], +])("quotes an edit that keeps the base image's shape for %s", async (_name, model, expected) => { + await quoteCanvas2D(model); + expect(mocks.estimate).toHaveBeenCalledWith( + expect.objectContaining({ aspect_ratio: expected }), + ); +}); + +it.each([ + ["wide", "wide"], + ["tall", "tall"], + ["square", "square"], +])("quotes the %s ratio picked by the user as before", async (picked, expected) => { + mocks.canvas2D.aspectRatio = picked; + await quoteCanvas2D(GPT_IMAGE_1P5); + expect(mocks.estimate).toHaveBeenCalledWith( + expect.objectContaining({ aspect_ratio: expected }), + ); +}); + +async function quoteCanvas2D(model: ImageModel) { + renderHook(() => useImageCostEstimate(ModelPage.Canvas2D, model, "artcraft")); + await waitFor(() => + expect(mocks.setCredits).toHaveBeenLastCalledWith(ModelPage.Canvas2D, 12), + ); +} + +it("quotes the same fixed ratio as generation for a portrait canvas", async () => { + mocks.canvas2D.sourceImageDimensions = { width: 900, height: 1600 }; + await quoteCanvas2D(GPT_IMAGE_1P5); + expect(mocks.estimate).toHaveBeenCalledWith(expect.objectContaining({ aspect_ratio: "tall" })); +}); + +it("updates the quote when the source canvas changes shape", async () => { + mocks.canvas2D.sourceImageDimensions = { width: 900, height: 1600 }; + const { rerender } = renderHook(() => useImageCostEstimate(ModelPage.Canvas2D, GPT_IMAGE_1P5, "artcraft")); + await waitFor(() => expect(mocks.estimate).toHaveBeenLastCalledWith(expect.objectContaining({ aspect_ratio: "tall" }))); + mocks.canvas2D.sourceImageDimensions = { width: 1600, height: 900 }; + rerender(); + await waitFor(() => expect(mocks.estimate).toHaveBeenLastCalledWith(expect.objectContaining({ aspect_ratio: "wide" }))); +}); diff --git a/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.ts b/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.ts --- a/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.ts +++ b/frontend/libs/components/pricing-modal/src/lib/useImageCostEstimate.ts @@ -1,6 +1,6 @@ import { useEffect, useState } from "react"; import { ModelPage } from "@storyteller/ui-model-selector"; -import { Model, ImageModel } from "@storyteller/model-list"; +import { Model, ImageModel, autoAspectRatio, matchSourceAspectRatio } from "@storyteller/model-list"; import { GenerationProvider } from "@storyteller/api-enums"; import { usePromptImageStore, @@ -45,6 +45,7 @@ export function useImageCostEstimate( const imageQuality = usePromptImageStore((s) => s.commonQuality); // Canvas2D store + const sourceImageDimensions = usePrompt2DStore((s) => s.sourceImageDimensions); const prompt2DAspectRatio = usePrompt2DStore((s) => s.aspectRatio); const prompt2DResolution = usePrompt2DStore((s) => s.resolution); const prompt2DReferenceImages = usePrompt2DStore((s) => s.referenceImages); @@ -84,7 +85,14 @@ export function useImageCostEstimate( generationCount = imageGenerationCount; break; case ModelPage.Canvas2D: - legacyAspectRatioStr = prompt2DAspectRatio; + if (prompt2DAspectRatio === "auto") { + // Use the same source dimensions and model options as generation. + aspectRatioStr = sourceImageDimensions + ? matchSourceAspectRatio(imageModel, sourceImageDimensions.width, sourceImageDimensions.height) + : autoAspectRatio(imageModel); + } else { + legacyAspectRatioStr = prompt2DAspectRatio; + } resolutionStr = prompt2DResolution; referenceImageCount = prompt2DReferenceImages.length + 1; generationCount = prompt2DGenerationCount; @@ -161,6 +169,7 @@ export function useImageCostEstimate( imageReferenceImages.length, imageGenerationCount, imageQuality, + sourceImageDimensions, prompt2DAspectRatio, prompt2DResolution, prompt2DReferenceImages.length, diff --git a/frontend/libs/components/promptbox/src/lib/PromptBox2D.spec.tsx b/frontend/libs/components/promptbox/src/lib/PromptBox2D.spec.tsx new file mode 100644 --- /dev/null +++ b/frontend/libs/components/promptbox/src/lib/PromptBox2D.spec.tsx @@ -0,0 +1,104 @@ +import { cleanup, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, expect, it, vi } from "vitest"; +import { buildImageModelsFromListing } from "@storyteller/model-list"; +import { PromptBox2D } from "./PromptBox2D"; +import { usePrompt2DStore } from "./promptStore"; + +const PROMPT = "make the sky stormy"; + +const [MODEL] = buildImageModelsFromListing([], [ + { + model: "nano_banana_pro", + aspect_ratio_options: ["auto", "square", "wide_sixteen_by_nine"], + }, +]); + +beforeEach(() => { + vi.stubGlobal("ResizeObserver", class { + observe() {} + disconnect() {} + unobserve() {} + }); + // jsdom has no layout; give popovers a visible rect. + vi.spyOn(HTMLElement.prototype, "getBoundingClientRect").mockReturnValue( + new DOMRect(100, 100, 280, 40), + ); + usePrompt2DStore.setState(usePrompt2DStore.getInitialState(), true); +}); + +afterEach(() => { + cleanup(); + vi.restoreAllMocks(); + vi.unstubAllGlobals(); +}); + +it("keeps the base image's aspect ratio unless the user picks another", async () => { + const { user, generate } = await renderPromptBox(); + + await user.type(screen.getByPlaceholderText("Describe your image..."), PROMPT); + await user.click(screen.getByRole("button", { name: /Generate/ })); + + await waitFor(() => expect(generate).toHaveBeenCalledTimes(1)); + expect(generate).toHaveBeenCalledWith( + PROMPT, + expect.objectContaining({ aspectRatio: "auto" }), + ); +}); + +it("offers auto next to wide, tall and square", async () => { + const { user } = await renderPromptBox(); + + await user.click(aspectRatioButton("Auto")); + + for (const label of ["Auto", "Wide", "Tall", "Square"]) { + expect(await screen.findAllByText(label)).not.toHaveLength(0); + } +}); + +it.each([ + ["Wide", "wide"], + ["Tall", "tall"], + ["Square", "square"], +])("sends %s once picked, and auto again when switched back", async (label, expected) => { + const { user, generate } = await renderPromptBox(); + await user.type(screen.getByPlaceholderText("Describe your image..."), PROMPT); + + await user.click(aspectRatioButton("Auto")); + await user.click(await screen.findByRole("button", { name: label })); + await user.click(screen.getByRole("button", { name: /Generate/ })); + await waitFor(() => expect(generate).toHaveBeenCalledTimes(1)); + expect(generate).toHaveBeenLastCalledWith( + PROMPT, + expect.objectContaining({ aspectRatio: expected }), + ); + + await user.click(aspectRatioButton(label)); + await user.click(await screen.findByRole("button", { name: "Auto" })); + await user.click(screen.getByRole("button", { name: /Generate/ })); + await waitFor(() => expect(generate).toHaveBeenCalledTimes(2)); + expect(generate).toHaveBeenLastCalledWith( + PROMPT, + expect.objectContaining({ aspectRatio: "auto" }), + ); +}); + +async function renderPromptBox() { + const user = userEvent.setup(); + const generate = vi.fn(); + render( + ""} + onGenerateClick={generate} + />, + ); + return { user, generate }; +} + +// The picker's trigger shows the selected ratio's name, as does that ratio's +// row in the open menu. Only the trigger reports whether the menu is expanded. +function aspectRatioButton(name: string) { + return screen.getByRole("button", { name, expanded: false }); +} diff --git a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx b/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx --- a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx @@ -4,7 +4,7 @@ import { PopoverMenu, PopoverItem } from "@storyteller/ui-popover"; import { Tooltip } from "@storyteller/ui-tooltip"; import { Button, GenerateButton } from "@storyteller/ui-button"; import { Modal } from "@storyteller/ui-modal"; -import { ChevronDownIcon, ChevronUpIcon, FrameIcon, MaximizeIcon, RectangleHorizontalIcon, RectangleVerticalIcon, SquareIcon } from "lucide-react"; +import { ChevronDownIcon, ChevronUpIcon, FrameIcon, MaximizeIcon, RectangleHorizontalIcon, RectangleVerticalIcon, SparklesIcon, SquareIcon } from "lucide-react"; import type { LucideIcon } from "lucide-react"; import { DynamicIcon } from "@storyteller/icons"; @@ -27,7 +27,8 @@ import { PromptClearAllButton } from "./PromptClearAllButton"; import { StoreApi, UseBoundStore } from "zustand"; import { toast } from "@storyteller/ui-toaster"; -export type AspectRatio = "wide" | "tall" | "square"; +// "auto" keeps the aspect ratio of the image being edited. +export type AspectRatio = "auto" | "wide" | "tall" | "square"; export interface PromptBox2DProps { uploadImage?: UploadImageFn; @@ -106,6 +107,11 @@ export const PromptBox2D = ({ const enterToGenerate = useEnterToGenerateStore((s) => s.enabled); const [showImagePrompts, setShowImagePrompts] = useState(false); const [aspectRatioList, setAspectRatioList] = useState([ + { + label: "Auto", + selected: aspectRatio === "auto", + icon: , + }, { label: "Wide", selected: aspectRatio === "wide", diff --git a/frontend/libs/components/promptbox/src/lib/promptStore.ts b/frontend/libs/components/promptbox/src/lib/promptStore.ts --- a/frontend/libs/components/promptbox/src/lib/promptStore.ts +++ b/frontend/libs/components/promptbox/src/lib/promptStore.ts @@ -32,17 +32,20 @@ export interface RefAudio { // ----- 2D Prompt Box Store ----- type AspectRatio = "wide" | "tall" | "square"; +// Editors can also keep the shape of the image being edited. +type EditAspectRatio = "auto" | AspectRatio; type Resolution = "1k" | "2k" | "4k"; export interface Prompt2DStore { + sourceImageDimensions?: { width: number; height: number }; prompt: string; - aspectRatio: AspectRatio; + aspectRatio: EditAspectRatio; resolution: Resolution; useSystemPrompt: boolean; referenceImages: RefImage[]; generationCount: number; setPrompt: (prompt: string) => void; - setAspectRatio: (ratio: AspectRatio) => void; + setAspectRatio: (ratio: EditAspectRatio) => void; setResolution: (resolution: Resolution) => void; setUseSystemPrompt: (value: boolean) => void; setReferenceImages: (images: RefImage[]) => void; @@ -51,7 +54,7 @@ export interface Prompt2DStore { export const usePrompt2DStore = create()((set) => ({ prompt: "", - aspectRatio: "wide", + aspectRatio: "auto", resolution: "1k", useSystemPrompt: true, referenceImages: [], @@ -285,8 +288,6 @@ export const usePromptAudioStore = create()((set) => ({ })); // ----- Edit Prompt Box Store ----- -type EditAspectRatio = "auto" | "wide" | "tall" | "square"; - interface PromptEditStore { referenceImages: RefImage[]; aspectRatio: EditAspectRatio; diff --git a/frontend/libs/model-list/src/index.ts b/frontend/libs/model-list/src/index.ts --- a/frontend/libs/model-list/src/index.ts +++ b/frontend/libs/model-list/src/index.ts @@ -27,5 +27,6 @@ export * from "./lib/lists/Object3DModels.js"; export * from "./lib/loader/buildModelsFromListing.js"; export * from "./lib/classes/videoOptions.js"; +export * from "./lib/classes/matchSourceAspectRatio.js"; export { imageModelForProvider } from "./lib/classes/imageModelForProvider.js"; diff --git a/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.spec.ts b/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.spec.ts new file mode 100644 --- /dev/null +++ b/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.spec.ts @@ -0,0 +1,202 @@ +// @vitest-environment jsdom +import { describe, expect, it } from "vitest"; +import { buildImageModelsFromListing } from "../loader/buildModelsFromListing.js"; +import { + autoAspectRatio, + matchSourceAspectRatio, +} from "./matchSourceAspectRatio.js"; +import { CommonAspectRatio as Ratio } from "./properties/CommonAspectRatio.js"; + +// Same lists the API serves for these models. +const NANO_BANANA_PRO = [ + Ratio.Auto, + Ratio.Square, + Ratio.WideFiveByFour, + Ratio.WideFourByThree, + Ratio.WideThreeByTwo, + Ratio.WideSixteenByNine, + Ratio.WideTwentyOneByNine, + Ratio.TallFourByFive, + Ratio.TallThreeByFour, + Ratio.TallTwoByThree, + Ratio.TallNineBySixteen, +]; +const SEEDREAM_4 = [ + Ratio.Auto2k, + Ratio.Auto4k, + Ratio.Square, + Ratio.SquareHd, + Ratio.WideFourByThree, + Ratio.WideSixteenByNine, + Ratio.TallThreeByFour, + Ratio.TallNineBySixteen, +]; +const FLUX_PRO_1P1 = [ + Ratio.Square, + Ratio.WideFourByThree, + Ratio.WideSixteenByNine, + Ratio.TallThreeByFour, + Ratio.TallNineBySixteen, +]; +const GPT_IMAGE_1P5 = [Ratio.Square, Ratio.Wide, Ratio.Tall]; + +describe("matchSourceAspectRatio", () => { + describe("models with an auto option", () => { + it.each([ + ["square", 1024, 1024], + ["landscape", 1920, 1080], + ["portrait", 1080, 1920], + ["an unlisted shape", 2400, 1000], + ])("asks for auto with a %s source", (_name, width, height) => { + expect(matchSourceAspectRatio(model(NANO_BANANA_PRO), width, height)).toBe( + Ratio.Auto, + ); + }); + + it("asks for auto when the source size is unknown", () => { + expect(matchSourceAspectRatio(model(NANO_BANANA_PRO), 0, 0)).toBe(Ratio.Auto); + }); + + it("prefers plain auto over the auto variants that pick a resolution", () => { + const options = [Ratio.Auto4k, Ratio.Auto2k, Ratio.Auto, Ratio.Square]; + expect(matchSourceAspectRatio(model(options), 1920, 1080)).toBe(Ratio.Auto); + }); + + it.each([ + [SEEDREAM_4, Ratio.Auto2k], + [[Ratio.Auto4k, Ratio.Square], Ratio.Auto4k], + [[Ratio.Square, Ratio.Auto3k, Ratio.Auto4k], Ratio.Auto3k], + ])("falls back to the auto variant a model lists (%#)", (options, expected) => { + expect(matchSourceAspectRatio(model(options), 1920, 1080)).toBe(expected); + }); + }); + + describe("models without an auto option", () => { + it.each([ + // [options, width, height, expected] + [FLUX_PRO_1P1, 1024, 1024, Ratio.Square], + [FLUX_PRO_1P1, 1200, 900, Ratio.WideFourByThree], + [FLUX_PRO_1P1, 900, 1200, Ratio.TallThreeByFour], + [FLUX_PRO_1P1, 1920, 1080, Ratio.WideSixteenByNine], + [FLUX_PRO_1P1, 1080, 1920, Ratio.TallNineBySixteen], + // 3:2 sits between 4:3 and 16:9; 4:3 is nearer + [FLUX_PRO_1P1, 3000, 2000, Ratio.WideFourByThree], + [FLUX_PRO_1P1, 2000, 3000, Ratio.TallThreeByFour], + // beyond the widest and tallest options + [FLUX_PRO_1P1, 4000, 1000, Ratio.WideSixteenByNine], + [FLUX_PRO_1P1, 1000, 4000, Ratio.TallNineBySixteen], + // not quite square + [FLUX_PRO_1P1, 1000, 1040, Ratio.Square], + // shapes are compared as ratios, not differences: 1.16 is nearer 4:3 than 1:1 + [FLUX_PRO_1P1, 1160, 1000, Ratio.WideFourByThree], + [FLUX_PRO_1P1, 1000, 1160, Ratio.TallThreeByFour], + ])("picks the closest listed ratio (%#)", (options, width, height, expected) => { + expect(matchSourceAspectRatio(model(options), width, height)).toBe(expected); + }); + + it("scales the source before comparing, so size does not matter", () => { + const small = matchSourceAspectRatio(model(FLUX_PRO_1P1), 40, 30); + const large = matchSourceAspectRatio(model(FLUX_PRO_1P1), 4000, 3000); + expect(small).toBe(Ratio.WideFourByThree); + expect(large).toBe(small); + }); + + it("breaks ties toward the option listed first", () => { + expect( + matchSourceAspectRatio(model([Ratio.Square, Ratio.SquareHd]), 512, 512), + ).toBe(Ratio.Square); + expect( + matchSourceAspectRatio(model([Ratio.SquareHd, Ratio.Square]), 512, 512), + ).toBe(Ratio.SquareHd); + }); + + it.each([ + [1024, 1024, Ratio.Square], + [1920, 1080, Ratio.Wide], + [1080, 1920, Ratio.Tall], + ])("reads wide and tall as 3:2 and 2:3 (%ix%i)", (width, height, expected) => { + expect(matchSourceAspectRatio(model(GPT_IMAGE_1P5), width, height)).toBe( + expected, + ); + }); + + it("uses square, wide and tall for a model that lists no ratios", () => { + const unlisted = model([]); + expect(matchSourceAspectRatio(unlisted, 1000, 1000)).toBe(Ratio.Square); + expect(matchSourceAspectRatio(unlisted, 1600, 900)).toBe(Ratio.Wide); + expect(matchSourceAspectRatio(unlisted, 900, 1600)).toBe(Ratio.Tall); + }); + + it("skips options this build does not know", () => { + const future = "future_aspect" as Ratio; + expect( + matchSourceAspectRatio(model([future, Ratio.WideFourByThree]), 1600, 900), + ).toBe(Ratio.WideFourByThree); + expect(matchSourceAspectRatio(model([future]), 1600, 900)).toBeUndefined(); + }); + + it.each([ + [0, 1080], + [1920, 0], + [-1920, 1080], + [Number.NaN, 1080], + [Number.POSITIVE_INFINITY, 1080], + ])("returns nothing for an unusable source size (%f x %f)", (width, height) => { + expect(matchSourceAspectRatio(model(FLUX_PRO_1P1), width, height)).toBeUndefined(); + }); + }); + + describe("models that cannot change aspect ratio", () => { + it("leaves the choice to the model, even when ratios are listed", () => { + const fixed = model(NANO_BANANA_PRO, { canChangeAspectRatio: false }); + expect(matchSourceAspectRatio(fixed, 1920, 1080)).toBeUndefined(); + }); + + it("leaves the choice to the model when none is selected yet", () => { + expect(matchSourceAspectRatio(undefined, 1920, 1080)).toBeUndefined(); + }); + }); + + describe("models built from the API listing", () => { + it.each([ + ["nano_banana_pro", ["auto", "square", "wide_sixteen_by_nine"], Ratio.Auto], + ["seedream_4", ["auto_2k", "auto_4k", "square"], Ratio.Auto2k], + ["flux_pro_1p1", ["square", "wide_four_by_three", "tall_three_by_four"], Ratio.WideFourByThree], + ])("%s keeps the source shape using the ratios the API lists", (id, aspectRatioOptions, expected) => { + const [hydrated] = buildImageModelsFromListing([], [ + { model: id, aspect_ratio_options: aspectRatioOptions }, + ]); + expect(matchSourceAspectRatio(hydrated, 1200, 900)).toBe(expected); + }); + + it("sends nothing for a model the API gives no ratios", () => { + const [hydrated] = buildImageModelsFromListing([], [{ model: "future_edit_model" }]); + expect(matchSourceAspectRatio(hydrated, 1200, 900)).toBeUndefined(); + }); + }); +}); + +describe("autoAspectRatio", () => { + it.each([ + [NANO_BANANA_PRO, Ratio.Auto], + [SEEDREAM_4, Ratio.Auto2k], + [FLUX_PRO_1P1, undefined], + [GPT_IMAGE_1P5, undefined], + ])("finds the option that works the shape out from the source (%#)", (options, expected) => { + expect(autoAspectRatio(model(options))).toBe(expected); + }); + + it("finds nothing for a model that cannot change aspect ratio", () => { + expect( + autoAspectRatio(model(NANO_BANANA_PRO, { canChangeAspectRatio: false })), + ).toBeUndefined(); + expect(autoAspectRatio(undefined)).toBeUndefined(); + }); +}); + +function model( + aspectRatios: Ratio[], + overrides: { canChangeAspectRatio?: boolean } = {}, +) { + return { canChangeAspectRatio: true, aspectRatios, ...overrides }; +} diff --git a/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.ts b/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.ts new file mode 100644 --- /dev/null +++ b/frontend/libs/model-list/src/lib/classes/matchSourceAspectRatio.ts @@ -0,0 +1,110 @@ +import type { ImageModel } from "./ImageModel.js"; +import { CommonAspectRatio } from "./properties/CommonAspectRatio.js"; + +// Options that leave the shape to the model, which works it out from the input +// image. The `auto_*` variants also fix a resolution (Seedream), so plain +// `auto` wins when a model offers both. +const AUTO_ASPECT_RATIOS = [ + CommonAspectRatio.Auto, + CommonAspectRatio.Auto2k, + CommonAspectRatio.Auto3k, + CommonAspectRatio.Auto4k, +]; + +// Models that can change aspect ratio but do not say which ones they accept. +// These are the values the image editor has always sent. +const LEGACY_ASPECT_RATIOS = [ + CommonAspectRatio.Square, + CommonAspectRatio.Wide, + CommonAspectRatio.Tall, +]; + +// width / height of every option that names a fixed shape. `auto*` and values +// from a newer server than this build have no entry, so they are never chosen +// by closeness. +const ASPECT_RATIO_SHAPES: Partial> = { + [CommonAspectRatio.Square]: 1, + [CommonAspectRatio.SquareHd]: 1, + + [CommonAspectRatio.WideFiveByFour]: 5 / 4, + [CommonAspectRatio.WideFourByThree]: 4 / 3, + [CommonAspectRatio.WideThreeByTwo]: 3 / 2, + [CommonAspectRatio.WideSixteenByNine]: 16 / 9, + [CommonAspectRatio.WideTwentyOneByNine]: 21 / 9, + + [CommonAspectRatio.TallFourByFive]: 4 / 5, + [CommonAspectRatio.TallThreeByFour]: 3 / 4, + [CommonAspectRatio.TallTwoByThree]: 2 / 3, + [CommonAspectRatio.TallNineBySixteen]: 9 / 16, + [CommonAspectRatio.TallNineByTwentyOne]: 9 / 21, + + // Imprecise labels whose exact shape varies by model (4:3 to 16:9); 3:2 and + // 2:3 are the middle. + [CommonAspectRatio.Wide]: 3 / 2, + [CommonAspectRatio.Tall]: 2 / 3, +}; + +type ModelAspectRatios = Pick; + +/** + * The aspect ratio to request so an edit keeps the shape of its source image. + * + * That is the model's own auto option when it has one, since it reads the + * source image itself. Otherwise it is the listed ratio closest to the + * source's shape. Returns `undefined` when the model has no aspect ratio + * control or lists nothing usable; leave the field out and the model decides. + */ +export function matchSourceAspectRatio( + model: ModelAspectRatios | undefined, + sourceWidth: number, + sourceHeight: number, +): CommonAspectRatio | undefined { + if (!model?.canChangeAspectRatio) return undefined; + const options = listedOrLegacyAspectRatios(model); + return ( + findAutoAspectRatio(options) ?? + closestAspectRatio(options, sourceWidth / sourceHeight) + ); +} + +/** The model's "work it out from the input image" option, if it lists one. */ +export function autoAspectRatio( + model: ModelAspectRatios | undefined, +): CommonAspectRatio | undefined { + if (!model?.canChangeAspectRatio) return undefined; + return findAutoAspectRatio(listedOrLegacyAspectRatios(model)); +} + +function listedOrLegacyAspectRatios( + model: ModelAspectRatios, +): readonly CommonAspectRatio[] { + return model.aspectRatios.length > 0 ? model.aspectRatios : LEGACY_ASPECT_RATIOS; +} + +function findAutoAspectRatio( + options: readonly CommonAspectRatio[], +): CommonAspectRatio | undefined { + return AUTO_ASPECT_RATIOS.find((ratio) => options.includes(ratio)); +} + +function closestAspectRatio( + options: readonly CommonAspectRatio[], + sourceShape: number, +): CommonAspectRatio | undefined { + if (!Number.isFinite(sourceShape) || sourceShape <= 0) return undefined; + + let closest: CommonAspectRatio | undefined; + let closestDistance = Infinity; + for (const option of options) { + const shape = ASPECT_RATIO_SHAPES[option]; + if (shape === undefined) continue; + // Compare in log space so 2:1 and 1:2 are equally far from square. + const distance = Math.abs(Math.log(sourceShape / shape)); + // Strictly closer: on a tie the option listed first wins. + if (distance < closestDistance) { + closest = option; + closestDistance = distance; + } + } + return closest; +} diff --git a/netlify.toml b/netlify.toml new file mode 100644 --- /dev/null +++ b/netlify.toml @@ -0,0 +1,7 @@ +# The `storyteller-docs` Netlify site still builds this repository, but its +# build command (./script/netlify_build_docs.sh) and the docs themselves now +# live in storytold/artcraft-services, so every deploy preview here fails. +# Skip those builds until the Netlify site is repointed or disconnected, then +# delete this file. +[build] + ignore = "exit 0"