batch matching fix

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<HTMLDivElement>(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<number>(
+    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]"
     >
-      <div className="h-full w-full p-4">
+      <div className="relative h-full w-full p-16">
         <div className="flex h-full w-full flex-col items-center justify-center rounded-md pb-12">
-          <div className="relative z-20 mb-12 flex flex-col items-center justify-center text-center drop-shadow-xl">
-            <span className="text-7xl font-bold">Generate Image</span>
-            <span className="pt-2 text-xl opacity-80">
-              Add a prompt, then generate
-            </span>
-          </div>
+          {!showPromptAtBottom && (
+            <div className="relative z-20 mb-52 flex flex-col items-center justify-center text-center drop-shadow-xl">
+              <span className="text-7xl font-bold">Generate Image</span>
+              <span className="pt-2 text-xl opacity-80">
+                Add a prompt, then generate
+              </span>
+            </div>
+          )}
 
-          <PromptBoxImage
-            useJobContext={() => {
-              console.log("Using job context");
-              return jobContext;
-            }}
-            model={selectedModel}
-            modelInfo={selectedModelInfo}
-            imageMediaId={imageMediaId}
-            url={imageUrl ?? undefined}
-            onEnqueuePressed={async () => {
-              console.log("Enqueue pressed");
-            }}
-          />
+          {hasAnyBatches && (
+            <div className="h-full w-full overflow-y-auto pb-40">
+              <div className="mx-auto flex max-w-screen-2xl flex-col gap-8 pr-2">
+                {batches.map((batch) => (
+                  <div key={batch.id} className="flex items-start gap-4">
+                    <div>
+                      <div className="glass inline-block w-[320px] shrink-0 rounded-xl px-4 py-3 text-left text-sm text-white/90">
+                        <div>{batch.prompt}</div>
+                      </div>
+                      <div className="mt-2 flex justify-end">
+                        <Badge
+                          label={batch.modelLabel}
+                          className="px-2 py-1 text-xs opacity-70"
+                        />
+                      </div>
+                    </div>
+                    <div className="grid flex-1 grid-cols-4 gap-4">
+                      {batch.status === "pending" && batch.images.length === 0
+                        ? Array.from({
+                            length: Math.max(
+                              1,
+                              Math.min(4, batch.requestedCount ?? 4),
+                            ),
+                          }).map((_, i) => (
+                            <div
+                              key={`sk-${batch.id}-${i}`}
+                              className="aspect-square w-full animate-pulse rounded-lg bg-white/5"
+                            />
+                          ))
+                        : batch.images.slice(0, 4).map((img) => (
+                            <button
+                              key={img.media_token}
+                              onClick={() => {
+                                galleryModalLightboxImage.value = {
+                                  id: img.media_token,
+                                  label: batch.prompt || "Generated Image",
+                                  thumbnail:
+                                    img.maybe_thumbnail_template || img.cdn_url,
+                                  fullImage: img.cdn_url,
+                                  createdAt: new Date(
+                                    batch.createdAt,
+                                  ).toISOString(),
+                                };
+                                galleryModalLightboxVisible.value = true;
+                              }}
+                              className="aspect-square w-full overflow-hidden rounded-lg"
+                            >
+                              <img
+                                src={img.cdn_url}
+                                alt="Generated"
+                                className="h-full w-full object-cover"
+                              />
+                            </button>
+                          ))}
+                    </div>
+                  </div>
+                ))}
+              </div>
+            </div>
+          )}
 
-          <BackgroundGallery />
+          <animated.div
+            className="fixed left-1/2 z-20 w-[730px] -translate-x-1/2"
+            style={promptAnim}
+          >
+            {showPromptAtBottom && batches.length > 0 && (
+              <div className="absolute -top-9 flex w-full justify-end">
+                <button
+                  onClick={() => resetBatches()}
+                  className="rounded-md bg-red/20 px-3 py-1 text-xs text-white/70 transition-colors hover:bg-red/30"
+                >
+                  Clear session
+                </button>
+              </div>
+            )}
+            <PromptBoxImage
+              useJobContext={() => {
+                return jobContext;
+              }}
+              model={selectedModel}
+              modelInfo={selectedModelInfo}
+              imageMediaId={imageMediaId}
+              url={imageUrl ?? undefined}
+              onEnqueuePressed={async (prompt, count, subscriberId) => {
+                startBatch(prompt, count, selectedModel, subscriberId);
+              }}
+            />
+          </animated.div>
+
+          {!showPromptAtBottom && <BackgroundGallery />}
 
           <div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
             <ModelSelector

frontend/apps/editor3d/app/src/pages/PageImage/TextToImageStore.ts

@@ -0,0 +1,96 @@
+import { create } from "zustand";
+import type { GeneratedImage } from "@storyteller/tauri-events";
+
+export type TextToImageBatch = {
+  id: string;
+  prompt: string;
+  status: "pending" | "complete";
+  images: GeneratedImage[];
+  createdAt: number;
+  requestedCount: number;
+  modelLabel: string;
+  subscriberId: string;
+};
+
+type TextToImageState = {
+  batches: TextToImageBatch[];
+  startBatch: (
+    prompt: string,
+    requestedCount: number,
+    modelLabel: string,
+    subscriberId?: string,
+  ) => string;
+  completeBatch: (
+    images: GeneratedImage[],
+    maybeSubscriberId?: string,
+    maybePrompt?: string,
+  ) => void;
+  reset: () => void;
+};
+
+export const useTextToImageStore = create<TextToImageState>((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: [] }),
+}));

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<void>;
+  onEnqueuePressed?: (
+    prompt: string,
+    count: number,
+    subscriberId: string
+  ) => void | Promise<void>;
   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: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
     }));
     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);
   };
 

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
 }

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<string, ModelConfig> = 
-  new Map(ALL_MODELS.map(model => [model.id, model]));
+export const ALL_MODELS_BY_ID: Map<string, ModelConfig> = new Map(
+  ALL_MODELS.map((model) => [model.id, model])
+);
 
 export const getAllModels = (): ModelConfig[] => ALL_MODELS;
 

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';
+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";