Compare commits

...
Author SHA1 Message Date
bflatastic 3201333056 centralized models refactor 2025-08-13 23:42:56 +07:00
11 changed files with 360 additions and 248 deletions
@@ -2,144 +2,52 @@ import type { PopoverItem } from "@storyteller/ui-popover";
import { faClock, faFilm, faImage } from "@fortawesome/pro-solid-svg-icons";
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import {
IMAGE_MODELS,
VIDEO_MODELS,
getCreatorIcon,
ModelCreator,
getModelsByCategory,
getInstructiveImageEditModels,
} from "@storyteller/model-list";
export type ModelList = Omit<PopoverItem, "selected">[];
const withIcon = (creatorIcon: any, fallback: any) => creatorIcon || fallback;
const buildItems = (
models: {
label: string;
description?: string;
badges?: { label: string }[];
info: { creator: any };
}[],
fallbackIcon: any
) =>
models.map((m) => ({
label: m.label,
icon: withIcon(getCreatorIcon((m as any).info.creator), fallbackIcon),
description: (m as any).description,
badges: (m as any).badges?.map((b: any) => ({
label: b.label,
icon: <FontAwesomeIcon icon={faClock} />,
})),
modelInfo: (m as any).info,
}));
const imageModels = getModelsByCategory("image");
const videoModels = getModelsByCategory("video");
const instructiveModels = getInstructiveImageEditModels();
export const allModels = {
video: [
{
label: "Kling 1.6 Pro",
icon: getCreatorIcon(ModelCreator.Kling) || (
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
description: "Good quality model",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: VIDEO_MODELS.kling_1_6_pro,
},
{
label: "Kling 2.1 Pro",
icon: getCreatorIcon(ModelCreator.Kling) || (
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: VIDEO_MODELS.kling_2_1_pro,
},
{
label: "Kling 2.1 Master",
icon: getCreatorIcon(ModelCreator.Kling) || (
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
description: "Master quality model ($$)",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: VIDEO_MODELS.kling_2_1_master,
},
{
label: "Seedance 1.0 Lite",
icon: getCreatorIcon(ModelCreator.Bytedance) || (
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
description: "Fast and high-quality model",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: VIDEO_MODELS.seedance_1_0_lite,
},
{
label: "Google Veo 2",
icon: getCreatorIcon(ModelCreator.Google) || (
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
description: "Fast and high-quality model",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: VIDEO_MODELS.veo_2,
},
],
image: [
{
label: "Flux Pro 1.1 Ultra",
icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.flux_pro_1_1_ultra,
},
{
label: "Flux Pro 1.1",
icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.flux_pro_1_1,
},
{
label: "Flux 1 Dev",
icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.flux_1_dev,
},
{
label: "Flux 1 Schnell",
icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.flux_1_schnell,
},
{
label: "GPT Image 1 (GPT-4o)",
icon: getCreatorIcon(ModelCreator.OpenAi) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "Slow, ultra instructive model",
badges: [{ label: "45 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.gpt_image_1,
},
//{
// id: "todo",
// label: "Recraft 3",
// icon: <FontAwesomeIcon icon={faImage} className="h-4 w-4" />,
// description: "Fast and high-quality model",
// badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
//},
//{
// id: "todo",
// label: "Flux.1 Kontext",
// icon: <FontAwesomeIcon icon={faImage} className="h-4 w-4" />,
// description: "New model with advanced context",
// badges: [{ label: "25 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
//},
],
// These models are used for 2D and 3D editing and scene building.
instructiveImageEdits: [
{
label: "GPT-4o",
icon: getCreatorIcon(ModelCreator.OpenAi) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "High quality model",
badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.gpt_image_1,
},
{
label: "FLUX.1 Pro Kontext Max",
icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
description: "Fast and high-quality model",
badges: [{ label: "20 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
modelInfo: IMAGE_MODELS.flux_pro_kontext_max,
},
],
video: buildItems(
videoModels as any,
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
),
image: buildItems(
imageModels as any,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
instructiveImageEdits: buildItems(
instructiveModels as any,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
),
};
export const videoGenerationModels: ModelList = allModels.video;
@@ -22,6 +22,7 @@ import {
faUpload,
faImages,
faFrame,
faCopy,
} from "@fortawesome/pro-solid-svg-icons";
import {
faRectangleVertical,
@@ -46,7 +47,7 @@ import {
// import { invoke } from "@tauri-apps/api/core";
import { usePrompt2DStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
import { ModelInfo } from "@storyteller/model-list";
import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
export type AspectRatio = "1:1" | "3:2" | "2:3";
@@ -92,6 +93,8 @@ export const PromptBox2D = ({
const setUseSystemPrompt = usePrompt2DStore((s) => s.setUseSystemPrompt);
const aspectRatio = usePrompt2DStore((s) => s.aspectRatio);
const setAspectRatio = usePrompt2DStore((s) => s.setAspectRatio);
const generationCount = usePrompt2DStore((s) => s.generationCount);
const setGenerationCount = usePrompt2DStore((s) => s.setGenerationCount);
const [isEnqueueing, setIsEnqueueing] = useState(false);
const [selectedGalleryImages, setSelectedGalleryImages] = useState<string[]>(
@@ -120,6 +123,41 @@ export const PromptBox2D = ({
},
]);
const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
[]
);
// Update generation count options from selected model capabilities
useEffect(() => {
const caps = getCapabilitiesForModel(selectedModelInfo);
const items: PopoverItem[] = Array.from(
{ length: caps.maxGenerationCount },
(_, i) => i + 1
).map((count) => ({
label: String(count),
selected: count === generationCount,
}));
setGenerationCountList(items);
// Clamp selection to allowed range
if (generationCount < 1 || generationCount > caps.maxGenerationCount) {
setGenerationCount(
Math.min(Math.max(1, generationCount), caps.maxGenerationCount)
);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [selectedModelInfo]);
// Keep UI selection in sync when store value changes
useEffect(() => {
setGenerationCountList((prev) =>
prev.map((item) => ({
...item,
selected: item.label === String(generationCount),
}))
);
}, [generationCount]);
const fileInputRef = useRef<HTMLInputElement>(null);
const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -249,6 +287,17 @@ export const PromptBox2D = ({
}
};
const handleGenerationCountSelect = (selectedItem: PopoverItem) => {
const count = parseInt(selectedItem.label, 10);
setGenerationCount(count);
setGenerationCountList((prev) =>
prev.map((item) => ({
...item,
selected: item.label === selectedItem.label,
}))
);
};
const handlePaste = (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
e.preventDefault();
e.stopPropagation();
@@ -327,7 +376,7 @@ export const PromptBox2D = ({
image_media_tokens: referenceImages.map((image) => image.mediaToken),
disable_system_prompt: !useSystemPrompt,
prompt: prompt,
image_count: 1,
image_count: generationCount,
aspect_ratio: aspectRatio,
});
@@ -564,8 +613,8 @@ export const PromptBox2D = ({
onChange={handleChange}
onPaste={handlePaste}
onKeyDown={handleKeyDown}
onFocus={() => { }}
onBlur={() => { }}
onFocus={() => {}}
onBlur={() => {}}
/>
</div>
<div className="mt-2 flex items-center justify-between gap-2">
@@ -626,6 +675,23 @@ export const PromptBox2D = ({
</Tooltip>
</div>
<div className="flex items-center gap-2">
<Tooltip
content="Number of generations"
position="top"
className="z-50"
closeOnClick={true}
>
<PopoverMenu
items={generationCountList}
onSelect={handleGenerationCountSelect}
mode="toggle"
panelTitle="No. of images"
triggerIcon={
<FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
}
buttonClassName="h-9"
/>
</Tooltip>
<Button
className="flex items-center border-none bg-primary px-3 text-sm text-white disabled:cursor-not-allowed disabled:opacity-50"
icon={!isEnqueueing ? faSparkles : undefined}
@@ -28,7 +28,7 @@ import {
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import { IconProp } from "@fortawesome/fontawesome-svg-core";
import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
import { ModelInfo } from "@storyteller/model-list";
import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
import { usePromptImageStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
@@ -102,28 +102,7 @@ export const PromptBoxImage = ({
},
]);
const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
[
{
label: "1",
selected: generationCount === 1,
icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
},
{
label: "2",
selected: generationCount === 2,
icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
},
{
label: "3",
selected: generationCount === 3,
icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
},
{
label: "4",
selected: generationCount === 4,
icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
},
]
[]
);
const fileInputRef = useRef<HTMLInputElement>(null);
@@ -148,6 +127,37 @@ export const PromptBoxImage = ({
}
}, [imageMediaId, url]);
useEffect(() => {
// Build generation count options based on selected model
const caps = getCapabilitiesForModel(modelInfo);
const items: PopoverItem[] = Array.from(
{ length: caps.maxGenerationCount },
(_, i) => i + 1
).map((count) => ({
label: String(count),
selected: count === generationCount,
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]);
useEffect(() => {
setGenerationCountList((prev) =>
prev.map((item) => ({
...item,
selected: item.label === String(generationCount),
}))
);
}, [generationCount]);
useEffect(() => {
setAspectRatioList((prev) =>
prev.map((item) => ({
@@ -157,15 +167,6 @@ export const PromptBoxImage = ({
);
}, [aspectRatio]);
useEffect(() => {
setGenerationCountList((prev) =>
prev.map((item) => ({
...item,
selected: item.label === generationCount.toString(),
}))
);
}, [generationCount]);
const handleAspectRatioSelect = (selectedItem: PopoverItem) => {
setAspectRatio(selectedItem.label as any);
setAspectRatioList((prev) =>
@@ -15,10 +15,12 @@ interface Prompt2DStore {
aspectRatio: AspectRatio;
useSystemPrompt: boolean;
referenceImages: RefImage[];
generationCount: number;
setPrompt: (prompt: string) => void;
setAspectRatio: (ratio: AspectRatio) => void;
setUseSystemPrompt: (value: boolean) => void;
setReferenceImages: (images: RefImage[]) => void;
setGenerationCount: (count: number) => void;
}
export const usePrompt2DStore = create<Prompt2DStore>()((set) => ({
@@ -26,10 +28,12 @@ export const usePrompt2DStore = create<Prompt2DStore>()((set) => ({
aspectRatio: "3:2",
useSystemPrompt: true,
referenceImages: [],
generationCount: 1,
setPrompt: (prompt) => set({ prompt }),
setAspectRatio: (aspectRatio) => set({ aspectRatio }),
setUseSystemPrompt: (useSystemPrompt) => set({ useSystemPrompt }),
setReferenceImages: (referenceImages) => set({ referenceImages }),
setGenerationCount: (generationCount) => set({ generationCount }),
}));
export { usePrompt2DStore as usePromptStore };
+2 -4
View File
@@ -1,7 +1,5 @@
export * from "./lib/model-list.js";
export * from "./lib/ModelInfo.js";
export * from "./lib/ModelCreator.js";
export * from "./lib/ImageModels.js";
export * from "./lib/VideoModels.js";
export * from "./lib/ModelCreatorIcons.js";
export * from "./lib/ModelTypeMapping.js";
export * from "./lib/ModelMapping.js";
export * from "./lib/Models.js";
@@ -1,35 +0,0 @@
import { ModelCreator } from "./ModelCreator.js";
import { ModelInfo } from "./ModelInfo.js";
export const IMAGE_MODELS: Record<string, ModelInfo> = {
flux_pro_1_1_ultra: {
name: "Flux Pro 1.1 Ultra",
tauri_id: "flux_pro_11_ultra",
creator: ModelCreator.BlackForestLabs,
},
flux_pro_1_1: {
name: "Flux Pro 1.1",
tauri_id: "flux_pro_11",
creator: ModelCreator.BlackForestLabs,
},
flux_1_dev: {
name: "Flux 1 Dev",
tauri_id: "flux_1_dev",
creator: ModelCreator.BlackForestLabs,
},
flux_1_schnell: {
name: "Flux 1 Schnell",
tauri_id: "flux_1_schnell",
creator: ModelCreator.BlackForestLabs,
},
gpt_image_1: {
name: "GPT Image 1 (GPT-4o)",
tauri_id: "gpt_image_1",
creator: ModelCreator.OpenAi,
},
flux_pro_kontext_max: {
name: "Flux Pro Kontext Max",
tauri_id: "flux_pro_kontext_max",
creator: ModelCreator.BlackForestLabs,
},
};
@@ -94,12 +94,13 @@ export const getModelDisplayName = (modelType: string): string => {
hunyuan_3d_2p1: "Hunyuan 3D 2.1",
hunyuan_3d: "Hunyuan 3D",
flux_pro_kontext_max: "Flux Pro Kontext Max",
// TODO: add more models here - BFlat
};
return displayNames[modelType] || modelType;
};
// Convert provider string to human-readable display name
// Convert provider string to human-readable display name (this is for the provider priority in settings)
export const getProviderDisplayName = (provider: string): string => {
const displayNames: Record<string, string> = {
artcraft: "ArtCraft",
+209
View File
@@ -0,0 +1,209 @@
import { ModelCreator } from "./ModelCreator.js";
import { ModelInfo } from "./ModelInfo.js";
export type ModelCategory = "image" | "video";
export interface ModelConfig {
id: string;
label: string; // UI label
description?: string;
badges?: { label: string }[];
category: ModelCategory;
info: ModelInfo;
capabilities: ModelCapabilities;
tags?: string[]; // optional tags, e.g. ["instructiveEdit"] - for filtering
}
const mc = ModelCreator;
// Centralized model capability definition
// For setting options for the model
// TODO: add more capabilities here - BFlat
export interface ModelCapabilities {
maxGenerationCount: number;
}
const DEFAULT_CAPABILITIES: ModelCapabilities = {
maxGenerationCount: 1,
};
const cfg = (
m: Partial<ModelConfig> & {
id: string;
category: ModelCategory;
info: ModelInfo;
}
): ModelConfig => ({
label: m.label ?? m.info.name,
description: m.description,
badges: m.badges,
capabilities: m.capabilities ?? DEFAULT_CAPABILITIES,
tags: m.tags ?? [],
...m,
});
export const ALL_MODELS: ModelConfig[] = [
//////////////////////////////
// Image models
//////////////////////////////
cfg({
id: "flux_pro_1_1_ultra",
category: "image",
info: {
name: "Flux Pro 1.1 Ultra",
tauri_id: "flux_pro_11_ultra",
creator: mc.BlackForestLabs,
},
description: "High quality model",
badges: [{ label: "15 sec." }],
capabilities: { maxGenerationCount: 4 },
tags: ["instructiveEdit"],
}),
cfg({
id: "flux_pro_1_1",
category: "image",
info: {
name: "Flux Pro 1.1",
tauri_id: "flux_pro_11",
creator: mc.BlackForestLabs,
},
description: "High quality model",
badges: [{ label: "15 sec." }],
capabilities: { maxGenerationCount: 4 },
tags: ["instructiveEdit"],
}),
cfg({
id: "flux_1_dev",
category: "image",
info: {
name: "Flux 1 Dev",
tauri_id: "flux_1_dev",
creator: mc.BlackForestLabs,
},
description: "High quality model",
badges: [{ label: "15 sec." }],
capabilities: { maxGenerationCount: 4 },
}),
cfg({
id: "flux_1_schnell",
category: "image",
info: {
name: "Flux 1 Schnell",
tauri_id: "flux_1_schnell",
creator: mc.BlackForestLabs,
},
description: "High quality model",
badges: [{ label: "15 sec." }],
capabilities: { maxGenerationCount: 4 },
}),
cfg({
id: "gpt_image_1",
category: "image",
info: {
name: "GPT Image 1 (GPT-4o)",
tauri_id: "gpt_image_1",
creator: mc.OpenAi,
},
description: "Slow, ultra instructive model",
badges: [{ label: "45 sec." }],
capabilities: { maxGenerationCount: 1 },
tags: ["instructiveEdit"],
}),
cfg({
id: "flux_pro_kontext_max",
category: "image",
info: {
name: "Flux Pro Kontext Max",
tauri_id: "flux_pro_kontext_max",
creator: mc.BlackForestLabs,
},
description: "Fast and high-quality model",
badges: [{ label: "20 sec." }],
capabilities: { maxGenerationCount: 4 },
tags: ["instructiveEdit"],
}),
//////////////////////////////
// Video models
//////////////////////////////
cfg({
id: "kling_1_6_pro",
category: "video",
info: {
name: "Kling 1.6 Pro",
tauri_id: "kling_1.6_pro",
creator: mc.Kling,
},
description: "Good quality model",
badges: [{ label: "2 min." }],
capabilities: { maxGenerationCount: 1 },
}),
cfg({
id: "kling_2_1_pro",
category: "video",
info: {
name: "Kling 2.1 Pro",
tauri_id: "kling_2.1_pro",
creator: mc.Kling,
},
description: "High quality model",
badges: [{ label: "2 min." }],
capabilities: { maxGenerationCount: 1 },
}),
cfg({
id: "kling_2_1_master",
category: "video",
info: {
name: "Kling 2.1 Master",
tauri_id: "kling_2.1_master",
creator: mc.Kling,
},
description: "Master quality model ($$)",
badges: [{ label: "2 min." }],
capabilities: { maxGenerationCount: 1 },
}),
cfg({
id: "seedance_1_0_lite",
category: "video",
info: {
name: "Seedance 1.0 Lite",
tauri_id: "seedance_1.0_lite",
creator: mc.Bytedance,
},
description: "Fast and high-quality model",
badges: [{ label: "2 min." }],
capabilities: { maxGenerationCount: 1 },
}),
cfg({
id: "veo_2",
category: "video",
info: { name: "Google Veo 2", tauri_id: "veo_2", creator: mc.Google },
description: "Fast and high-quality model",
badges: [{ label: "2 min." }],
capabilities: { maxGenerationCount: 1 },
}),
];
export const getAllModels = (): ModelConfig[] => ALL_MODELS;
export const getModelsByCategory = (category: ModelCategory): ModelConfig[] =>
ALL_MODELS.filter((m) => m.category === category);
export const getInstructiveImageEditModels = (): ModelConfig[] =>
ALL_MODELS.filter(
(m) => m.category === "image" && m.tags?.includes("instructiveEdit")
);
export const lookupModelByTauriId = (
tauriId: string
): ModelConfig | undefined =>
ALL_MODELS.find((m) => m.info.tauri_id === tauriId);
// Single exported capability resolver so callers never need to touch anything else
export const getCapabilitiesForModel = (
model?: ModelInfo
): ModelCapabilities => {
if (!model) return DEFAULT_CAPABILITIES;
const cfg = lookupModelByTauriId(model.tauri_id);
return cfg?.capabilities ?? DEFAULT_CAPABILITIES;
};
@@ -1,30 +0,0 @@
import { ModelCreator } from "./ModelCreator.js";
import { ModelInfo } from "./ModelInfo.js";
export const VIDEO_MODELS: Record<string, ModelInfo> = {
kling_1_6_pro: {
name: "Kling 1.6 Pro",
tauri_id: "kling_1.6_pro",
creator: ModelCreator.Kling,
},
kling_2_1_pro: {
name: "Kling 2.1 Pro",
tauri_id: "kling_2.1_pro",
creator: ModelCreator.Kling,
},
kling_2_1_master: {
name: "Kling 2.1 Master",
tauri_id: "kling_2.1_master",
creator: ModelCreator.Kling,
},
seedance_1_0_lite: {
name: "Seedance 1.0 Lite",
tauri_id: "seedance_1.0_lite",
creator: ModelCreator.Bytedance,
},
veo_2: {
name: "Google Veo 2",
tauri_id: "veo_2",
creator: ModelCreator.Google,
},
};
@@ -1,7 +0,0 @@
import { modelList } from './model-list.js';
describe('modelList', () => {
it('should work', () => {
expect(modelList()).toEqual('model-list');
})
})
@@ -1,3 +0,0 @@
export function modelList(): string {
return "model-list";
}