diff --git a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx b/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx index b04fd95ef1..537b686374 100644 --- a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx +++ b/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx @@ -1,4 +1,4 @@ -import { useRef } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { JobContextType } from "@storyteller/common"; import { PromptBoxImage } from "@storyteller/ui-promptbox"; import BackgroundGallery from "./BackgroundGallery"; @@ -15,20 +15,31 @@ interface TextToImageProps { } import { ModelInfo } from "@storyteller/model-list"; import { useTextToImageGenerationCompleteEvent } from "@storyteller/tauri-events"; +import { useTextToImageStore } from "./TextToImageStore"; +import { animated, useSpring } from "@react-spring/web"; +import { + galleryModalLightboxImage, + galleryModalLightboxVisible, +} from "@storyteller/ui-gallery-modal"; +import { Badge } from "@storyteller/ui-badge"; -const PAGE_ID : ModelPage = ModelPage.TextToImage; +const PAGE_ID: ModelPage = ModelPage.TextToImage; const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { const containerRef = useRef(null); const { selectedModels } = useModelSelectorStore(); + const batches = useTextToImageStore((s) => s.batches); + const startBatch = useTextToImageStore((s) => s.startBatch); + const completeBatch = useTextToImageStore((s) => s.completeBatch); + const resetBatches = useTextToImageStore((s) => s.reset); const selectedModel = - selectedModels[PAGE_ID] || - TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.label; + selectedModels[PAGE_ID] || TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.label; - const selectedModelInfo: ModelInfo | undefined = TEXT_TO_IMAGE_PAGE_MODEL_LIST.find( - (m) => m.label === selectedModel, - )?.modelInfo; + const selectedModelInfo: ModelInfo | undefined = + TEXT_TO_IMAGE_PAGE_MODEL_LIST.find( + (m) => m.label === selectedModel, + )?.modelInfo; const jobContext: JobContextType = { jobTokens: [], @@ -38,17 +49,31 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { }; useTextToImageGenerationCompleteEvent(async (event) => { - console.log("Text to image generation complete:", event); + completeBatch( + event.generated_images || [], + event.maybe_frontend_subscriber_id, + ); + }); - // TODO: Logic to show edit history - for (const generatedImage of event.generated_images) { - console.log("Generated image:", generatedImage.media_token); - console.log("Generated image url:", generatedImage.cdn_url); - console.log( - "Generated image thumbnail template:", - generatedImage.maybe_thumbnail_template, - ); - } + const hasAnyBatches = batches.length > 0; + const showPromptAtBottom = useMemo(() => hasAnyBatches, [hasAnyBatches]); + + const [vh, setVh] = useState( + typeof window !== "undefined" ? window.innerHeight : 800, + ); + useEffect(() => { + const onResize = () => setVh(window.innerHeight); + window.addEventListener("resize", onResize); + return () => window.removeEventListener("resize", onResize); + }, []); + + const bottomOffsetPx = 138; + const targetTop = showPromptAtBottom + ? Math.max(0, vh - bottomOffsetPx) + : Math.floor(vh / 2); + const promptAnim = useSpring({ + top: targetTop, + config: { tension: 200, friction: 28, mass: 1.1 }, }); return ( @@ -56,30 +81,107 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { ref={containerRef} className="flex h-[calc(100vh-56px)] w-full bg-[#121212]" > -
+
-
- Generate Image - - Add a prompt, then generate - -
+ {!showPromptAtBottom && ( +
+ Generate Image + + Add a prompt, then generate + +
+ )} - { - console.log("Using job context"); - return jobContext; - }} - model={selectedModel} - modelInfo={selectedModelInfo} - imageMediaId={imageMediaId} - url={imageUrl ?? undefined} - onEnqueuePressed={async () => { - console.log("Enqueue pressed"); - }} - /> + {hasAnyBatches && ( +
+
+ {batches.map((batch) => ( +
+
+
+
{batch.prompt}
+
+
+ +
+
+
+ {batch.status === "pending" && batch.images.length === 0 + ? Array.from({ + length: Math.max( + 1, + Math.min(4, batch.requestedCount ?? 4), + ), + }).map((_, i) => ( +
+ )) + : batch.images.slice(0, 4).map((img) => ( + + ))} +
+
+ ))} +
+
+ )} - + + {showPromptAtBottom && batches.length > 0 && ( +
+ +
+ )} + { + return jobContext; + }} + model={selectedModel} + modelInfo={selectedModelInfo} + imageMediaId={imageMediaId} + url={imageUrl ?? undefined} + onEnqueuePressed={async (prompt, count, subscriberId) => { + startBatch(prompt, count, selectedModel, subscriberId); + }} + /> +
+ + {!showPromptAtBottom && }
string; + completeBatch: ( + images: GeneratedImage[], + maybeSubscriberId?: string, + maybePrompt?: string, + ) => void; + reset: () => void; +}; + +export const useTextToImageStore = create((set, get) => ({ + batches: [], + startBatch: ( + prompt: string, + requestedCount: number, + modelLabel: string, + subscriberId?: string, + ) => { + const id = subscriberId + ? subscriberId + : crypto.randomUUID + ? crypto.randomUUID() + : Math.random().toString(36).slice(2); + const batch: TextToImageBatch = { + id, + prompt, + status: "pending", + images: [], + createdAt: Date.now(), + requestedCount, + modelLabel, + subscriberId: id, + }; + set((s) => ({ batches: [...s.batches, batch] })); + return id; + }, + completeBatch: ( + images: GeneratedImage[], + maybeSubscriberId?: string, + maybePrompt?: string, + ) => { + const pending = maybeSubscriberId + ? get().batches.find((b) => b.subscriberId === maybeSubscriberId) + : get().batches.find((b) => b.status === "pending"); + const prompt = pending?.prompt ?? maybePrompt ?? ""; + const modelLabel = pending?.modelLabel ?? ""; + // Mark the most recent pending batch complete, or create one if none exists + set((s) => { + const idx = pending + ? s.batches.findIndex((b) => b.id === pending.id) + : -1; + if (idx === -1) { + const id = Math.random().toString(36).slice(2); + const batch: TextToImageBatch = { + id, + prompt, + status: "complete", + images: images.slice(0, 4), + createdAt: Date.now(), + requestedCount: Math.max(1, Math.min(4, images.length || 4)), + modelLabel, + subscriberId: id, + }; + return { batches: [...s.batches, batch] }; + } + const updated = [...s.batches]; + updated[idx] = { + ...updated[idx], + status: "complete", + images: images.slice(0, 4), + }; + return { batches: updated }; + }); + }, + reset: () => set({ batches: [] }), +})); diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx index cf4eab25cc..f9fca84e78 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx @@ -34,7 +34,11 @@ import { gtagEvent } from "@storyteller/google-analytics"; interface PromptBoxImageProps { useJobContext: () => JobContextType; - onEnqueuePressed?: () => void | Promise; + onEnqueuePressed?: ( + prompt: string, + count: number, + subscriberId: string + ) => void | Promise; model: string; modelInfo?: ModelInfo; imageMediaId?: string; @@ -130,22 +134,22 @@ export const PromptBoxImage = ({ useEffect(() => { // Build generation count options based on selected model const caps = getCapabilitiesForModel(modelInfo); + const defaultCount = Math.min( + Math.max(1, caps.defaultGenerationCount ?? 1), + caps.maxGenerationCount + ); + + setGenerationCount(defaultCount); + const items: PopoverItem[] = Array.from( { length: caps.maxGenerationCount }, (_, i) => i + 1 ).map((count) => ({ label: String(count), - selected: count === generationCount, + selected: count === defaultCount, icon: , })); setGenerationCountList(items); - - // Clamp selection - if (generationCount < 1 || generationCount > caps.maxGenerationCount) { - setGenerationCount( - Math.min(Math.max(1, generationCount), caps.maxGenerationCount) - ); - } // eslint-disable-next-line react-hooks/exhaustive-deps }, [modelInfo]); @@ -264,6 +268,12 @@ export const PromptBoxImage = ({ gtagEvent("enqueue_image"); + const subscriberId = crypto.randomUUID + ? crypto.randomUUID() + : Math.random().toString(36).slice(2); + + onEnqueuePressed?.(prompt, generationCount, subscriberId); + setTimeout(() => { // TODO(bt,2025-05-08): This is a hack so we don't accidentally wind up with a permanently disabled prompt box if // the backend hangs on a given request. @@ -279,12 +289,11 @@ export const PromptBoxImage = ({ aspect_ratio: aspectRatio, number_images: generationCount, frontend_caller: "text_to_image", + frontend_subscriber_id: subscriberId, }); console.log("PromptBoxImage - generateResponse", generateResponse); - onEnqueuePressed?.(); - setIsEnqueueing(false); }; diff --git a/frontend/libs/model-list/src/lib/ModelConfig.ts b/frontend/libs/model-list/src/lib/ModelConfig.ts index 7718c1e916..ca8d65d730 100644 --- a/frontend/libs/model-list/src/lib/ModelConfig.ts +++ b/frontend/libs/model-list/src/lib/ModelConfig.ts @@ -8,6 +8,7 @@ export type ModelCategory = "image" | "video"; // TODO: add more capabilities here - BFlat export interface ModelCapabilities { maxGenerationCount: number; + defaultGenerationCount?: number; } export interface ModelConfig { @@ -18,5 +19,5 @@ export interface ModelConfig { category: ModelCategory; info: ModelInfo; capabilities: ModelCapabilities; - tags?: (ModelTag|string)[]; // optional tags, e.g. ["instructiveEdit"] - for filtering + tags?: (ModelTag | string)[]; // optional tags, e.g. ["instructiveEdit"] - for filtering } diff --git a/frontend/libs/model-list/src/lib/Models.ts b/frontend/libs/model-list/src/lib/Models.ts index ab32322c5c..390944181a 100644 --- a/frontend/libs/model-list/src/lib/Models.ts +++ b/frontend/libs/model-list/src/lib/Models.ts @@ -1,12 +1,17 @@ import { ModelCreator } from "./ModelCreator.js"; import { ModelInfo } from "./ModelInfo.js"; import { ModelTag } from "./ModelTag.js"; -import { ModelConfig, ModelCapabilities, ModelCategory } from "./ModelConfig.js"; +import { + ModelConfig, + ModelCapabilities, + ModelCategory, +} from "./ModelConfig.js"; const mc = ModelCreator; const DEFAULT_CAPABILITIES: ModelCapabilities = { maxGenerationCount: 1, + defaultGenerationCount: 1, }; const cfg = ( @@ -38,7 +43,7 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "Incredible style and quality", badges: [{ label: "15 sec." }], - capabilities: { maxGenerationCount: 4 }, + capabilities: { maxGenerationCount: 4, defaultGenerationCount: 4 }, }), cfg({ id: "flux_pro_1_1_ultra", @@ -50,7 +55,7 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "High quality model", badges: [{ label: "15 sec." }], - capabilities: { maxGenerationCount: 4 }, + capabilities: { maxGenerationCount: 4, defaultGenerationCount: 4 }, }), cfg({ id: "flux_pro_1_1", @@ -62,7 +67,7 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "High quality model", badges: [{ label: "15 sec." }], - capabilities: { maxGenerationCount: 4 }, + capabilities: { maxGenerationCount: 4, defaultGenerationCount: 4 }, }), cfg({ id: "flux_1_dev", @@ -74,7 +79,7 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "High quality model", badges: [{ label: "15 sec." }], - capabilities: { maxGenerationCount: 4 }, + capabilities: { maxGenerationCount: 4, defaultGenerationCount: 4 }, }), cfg({ id: "flux_1_schnell", @@ -86,7 +91,7 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "High quality model", badges: [{ label: "15 sec." }], - capabilities: { maxGenerationCount: 4 }, + capabilities: { maxGenerationCount: 4, defaultGenerationCount: 4 }, }), cfg({ id: "gpt_image_1", @@ -111,13 +116,10 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "Fast and high-quality model", badges: [{ label: "20 sec." }], - capabilities: { + capabilities: { maxGenerationCount: 4, }, - tags: [ - ModelTag.InstructiveEdit, - ModelTag.NonMaskedInpainting, - ], + tags: [ModelTag.InstructiveEdit, ModelTag.NonMaskedInpainting], }), cfg({ id: "flux_pro_inpaint", @@ -129,8 +131,8 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "Fast and high-quality model", badges: [{ label: "20 sec." }], - capabilities: { - maxGenerationCount: 1 // NB: For some reason Fal only supports ONE image! + capabilities: { + maxGenerationCount: 1, // NB: For some reason Fal only supports ONE image! }, tags: [ModelTag.MaskedInpainting], }), @@ -144,8 +146,8 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "Fast and high-quality model", badges: [{ label: "20 sec." }], - capabilities: { - maxGenerationCount: 4 + capabilities: { + maxGenerationCount: 4, }, tags: [ModelTag.MaskedInpainting], }), @@ -211,8 +213,9 @@ export const ALL_MODELS: ModelConfig[] = [ }), ]; -export const ALL_MODELS_BY_ID : Map = - new Map(ALL_MODELS.map(model => [model.id, model])); +export const ALL_MODELS_BY_ID: Map = new Map( + ALL_MODELS.map((model) => [model.id, model]) +); export const getAllModels = (): ModelConfig[] => ALL_MODELS; diff --git a/frontend/libs/tauri-events/src/index.ts b/frontend/libs/tauri-events/src/index.ts index c3b74d74b8..6daaa12322 100644 --- a/frontend/libs/tauri-events/src/index.ts +++ b/frontend/libs/tauri-events/src/index.ts @@ -1,10 +1,10 @@ -export * from './lib/common/BasicEventWrapper'; -export * from './lib/events/functional/ImageEditCompleteEvent'; -export * from './lib/events/functional/RefreshAccountStateEvent'; -export * from './lib/events/functional/ShowProviderLoginModalEvent'; -export * from './lib/events/functional/TextToImageGenerationCompleteEvent'; -export * from './lib/events/generation/useGenerationCompleteEvent'; -export * from './lib/events/generation/useGenerationEnqueueFailureEvent'; -export * from './lib/events/generation/useGenerationEnqueueSuccessEvent'; -export * from './lib/events/generation/useGenerationFailedEvent'; -export * from './lib/events/warning/FlashUserInputErrorEvent'; \ No newline at end of file +export * from "./lib/common/BasicEventWrapper"; +export * from "./lib/events/functional/ImageEditCompleteEvent"; +export * from "./lib/events/functional/RefreshAccountStateEvent"; +export * from "./lib/events/functional/ShowProviderLoginModalEvent"; +export * from "./lib/events/functional/TextToImageGenerationCompleteEvent"; +export * from "./lib/events/generation/useGenerationCompleteEvent"; +export * from "./lib/events/generation/useGenerationEnqueueFailureEvent"; +export * from "./lib/events/generation/useGenerationEnqueueSuccessEvent"; +export * from "./lib/events/generation/useGenerationFailedEvent"; +export * from "./lib/events/warning/FlashUserInputErrorEvent";