diff --git a/src/app/api/media-generate/route.ts b/src/app/api/media-generate/route.ts index e99814c9..a23fac21 100644 --- a/src/app/api/media-generate/route.ts +++ b/src/app/api/media-generate/route.ts @@ -2,7 +2,6 @@ import { NextRequest, NextResponse } from "next/server"; import { auth } from "@/lib/auth"; import { db } from "@/lib/db"; import { - initializeMediaGenerators, getMediaGeneratorPlugin, getAvailableModels, isMediaGenerationAvailable, @@ -15,8 +14,6 @@ export async function GET() { return NextResponse.json({ error: "Unauthorized" }, { status: 401 }); } - initializeMediaGenerators(); - const available = isMediaGenerationAvailable(); const imageModels = getAvailableModels("image"); const videoModels = getAvailableModels("video"); @@ -90,8 +87,6 @@ export async function POST(request: NextRequest) { ); } - initializeMediaGenerators(); - const plugin = getMediaGeneratorPlugin(provider); if (!plugin) { diff --git a/src/components/prompts/media-generator.tsx b/src/components/prompts/media-generator.tsx index 2a69b9d7..5c333a45 100644 --- a/src/components/prompts/media-generator.tsx +++ b/src/components/prompts/media-generator.tsx @@ -40,14 +40,14 @@ const ASPECT_RATIOS: { value: AspectRatio; label: string }[] = [ { value: "3:2", label: "3:2 (Photo)" }, { value: "2:3", label: "2:3 (Portrait)" }, ]; -import { wiroGeneratorPlugin } from "@/lib/plugins/media-generators/wiro"; -import { falGeneratorPlugin } from "@/lib/plugins/media-generators/fal"; +import { getProviderWebSocketHandler } from "@/lib/plugins/media-generators"; interface MediaGeneratorModel { id: string; name: string; type: "image" | "video"; provider: string; + providerName: string; } interface MediaGeneratorProps { @@ -76,17 +76,6 @@ function fillPromptVariables(prompt: string): string { }); } -// Get WebSocket handler for a provider -function getProviderHandler(provider: string): WebSocketHandler { - switch (provider) { - case "wiro": - return wiroGeneratorPlugin.webSocketHandler; - case "fal": - return falGeneratorPlugin.webSocketHandler; - default: - throw new Error(`Unknown provider: ${provider}`); - } -} export function MediaGenerator({ prompt, @@ -172,7 +161,7 @@ export function MediaGenerator({ const { socketAccessToken, webSocketUrl, provider } = await response.json(); // Get provider-specific handler - const handler = getProviderHandler(provider); + const handler = getProviderWebSocketHandler(provider); // Connect to WebSocket for progress tracking setStatus("queued"); @@ -271,10 +260,7 @@ export function MediaGenerator({ return acc; }, {} as Record); - const providerNames: Record = { - wiro: "Wiro.ai", - fal: "Fal.ai", - }; + const providerDisplayName = selectedModel?.providerName || selectedModel?.provider || ""; return ( <> @@ -300,7 +286,7 @@ export function MediaGenerator({ {Object.entries(modelsByProvider).map(([provider, providerModels]) => (
- {providerNames[provider] || provider} + {providerModels[0]?.providerName || provider} {providerModels.map((model) => ( {t("confirmGeneration")} {t("confirmGenerationDescription", { - provider: providerNames[selectedModel?.provider || ""] || selectedModel?.provider || "", + provider: providerDisplayName, model: selectedModel?.name || "" })} @@ -403,7 +389,7 @@ export function MediaGenerator({ - {t("generatingMedia", { provider: selectedModel?.provider || "" })} + {t("generatingMedia", { provider: providerDisplayName })} {t("doNotCloseWindow")} diff --git a/src/lib/plugins/media-generators/index.ts b/src/lib/plugins/media-generators/index.ts index b4994971..e4c7bbb9 100644 --- a/src/lib/plugins/media-generators/index.ts +++ b/src/lib/plugins/media-generators/index.ts @@ -2,6 +2,10 @@ * Media Generators Plugin Registry * * Manages AI-powered media generation plugins for images and videos. + * + * To add a new plugin: + * 1. Create your plugin file (e.g., ./my-provider.ts) implementing MediaGeneratorPlugin + * 2. Import and add it to the `plugins` array below */ import { wiroGeneratorPlugin } from "./wiro"; @@ -10,23 +14,33 @@ import type { MediaGeneratorPlugin, MediaGeneratorModel, MediaType, WebSocketHan export * from "./types"; -// Export individual plugins for direct access to handlers -export { wiroGeneratorPlugin } from "./wiro"; -export { falGeneratorPlugin } from "./fal"; +/** + * Register all plugins here - just add your plugin to this array + */ +const plugins: MediaGeneratorPlugin[] = [ + wiroGeneratorPlugin, + falGeneratorPlugin, + // Add new plugins here: + // myNewPlugin, +]; const mediaGeneratorRegistry = new Map(); let initialized = false; -export function registerMediaGeneratorPlugin(plugin: MediaGeneratorPlugin): void { - mediaGeneratorRegistry.set(plugin.id, plugin); +function initializeMediaGenerators(): void { + if (initialized) return; + plugins.forEach((plugin) => mediaGeneratorRegistry.set(plugin.id, plugin)); + initialized = true; } export function getMediaGeneratorPlugin(id: string): MediaGeneratorPlugin | undefined { + initializeMediaGenerators(); return mediaGeneratorRegistry.get(id); } export function getAllMediaGeneratorPlugins(): MediaGeneratorPlugin[] { + initializeMediaGenerators(); return Array.from(mediaGeneratorRegistry.values()); } @@ -34,22 +48,13 @@ export function getEnabledMediaGeneratorPlugins(): MediaGeneratorPlugin[] { return getAllMediaGeneratorPlugins().filter((p) => p.isEnabled()); } -export function initializeMediaGenerators(): void { - if (initialized) return; - - registerMediaGeneratorPlugin(wiroGeneratorPlugin); - registerMediaGeneratorPlugin(falGeneratorPlugin); - - initialized = true; -} - /** * Get all available models from enabled generators */ -export function getAvailableModels(type?: MediaType): Array { +export function getAvailableModels(type?: MediaType): Array { initializeMediaGenerators(); - const models: Array = []; + const models: Array = []; for (const plugin of getEnabledMediaGeneratorPlugins()) { const pluginModels = plugin.getModels(); @@ -58,6 +63,7 @@ export function getAvailableModels(type?: MediaType): Array 0; } + +/** + * Get WebSocket handler for a provider + */ +export function getProviderWebSocketHandler(providerId: string): WebSocketHandler { + initializeMediaGenerators(); + const plugin = getMediaGeneratorPlugin(providerId); + if (!plugin) { + throw new Error(`Unknown provider: ${providerId}`); + } + return plugin.webSocketHandler; +}