# Directory Structure
```
├── .env.example
├── .gitignore
├── assets
│ └── Claude-desktop.png
├── Dockerfile
├── LICENSE
├── package.json
├── README.md
├── smithery.yaml
├── src
│ └── index.ts
└── tsconfig.json
```
# Files
--------------------------------------------------------------------------------
/.env.example:
--------------------------------------------------------------------------------
```
1 | PIAPI_API_KEY=your-api-key
2 |
```
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
```
1 | .env
2 | package-lock.json
3 | node_modules
4 | dist
```
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
```markdown
1 | # piapi-mcp-server
2 |
3 | [](https://piapi.ai)
4 | [](https://piapi.ai/docs)
5 | [](https://discord.gg/qRRvcGa7Wb)
6 |
7 | [](https://smithery.ai/server/piapi-mcp-server)
8 |
9 | A TypeScript implementation of a Model Context Protocol (MCP) server that integrates with PiAPI's API. PiAPI makes user able to generate media content with Midjourney/Flux/Kling/LumaLabs/Udio/Chrip/Trellis directly from Claude or any other MCP-compatible apps.
10 |
11 | <a href="https://glama.ai/mcp/servers/ywvke8xruo"><img width="380" height="200" src="https://glama.ai/mcp/servers/ywvke8xruo/badge" alt="PiAPI-Server MCP server" /></a>
12 |
13 | ## Features (more coming soon)
14 |
15 | Note: Time-consuming tools like video generation may not complete due to Claude's timeout limitations
16 |
17 | - [x] Base Image toolkit
18 | - [x] Base Video toolkit
19 | - [x] Flux Image generation from text/image prompt
20 | - [x] Hunyuan Video generation from text/image prompt
21 | - [x] Skyreels Video generation from image prompt
22 | - [x] Wan Video generation from text/image prompt
23 | - [x] MMAudio Music generation from video
24 | - [x] TTS Zero-Shot voice generation
25 | - [ ] Midjourney Image generation
26 | - [x] imagine
27 | - [ ] other
28 | - [x] Kling Video and Effects generation
29 | - [x] Luma Dream Machine video generation
30 | - [x] Suno Music generation
31 | - [ ] Suno Lyrics generation
32 | - [ ] Udio Music and Lyrics generation
33 | - [x] Trellis 3D model generation from image
34 | - [ ] Workflow planning inside LLMs
35 |
36 | ## Working with Claude Desktop Example
37 |
38 | 
39 |
40 | ## Prerequisites
41 |
42 | - Node.js 16.x or higher
43 | - npm or yarn
44 | - A PiAPI API key (get one at [piapi.ai](https://piapi.ai/workspace/key))
45 |
46 | ## Installation
47 |
48 | ### Installing via Smithery
49 |
50 | To install PiAPI MCP Server for Claude Desktop automatically via [Smithery](https://smithery.ai/server/piapi-mcp-server):
51 |
52 | ```bash
53 | npx -y @smithery/cli install piapi-mcp-server --client claude
54 | ```
55 |
56 | ### Manual Installation
57 | 1. Clone the repository:
58 |
59 | ```bash
60 | git clone https://github.com/apinetwork/piapi-mcp-server
61 | cd piapi-mcp-server
62 | ```
63 |
64 | 2. Install dependencies:
65 |
66 | ```bash
67 | npm install
68 | ```
69 |
70 | 3. Build the project:
71 |
72 | ```bash
73 | npm run build
74 | ```
75 |
76 | After building, a `dist/index.js` file will be generated. You can then configure this file with Claude Desktop and other applications. For detailed configuration instructions, please refer to the Usage section.
77 |
78 | 4. (Optional) Test server with MCP Inspector:
79 |
80 | First, create a `.env` file in the project root directory with your API key:
81 |
82 | ```bash
83 | PIAPI_API_KEY=your_api_key_here
84 | ```
85 |
86 | Then run the following command to start the MCP Inspector:
87 |
88 | ```bash
89 | npm run inspect
90 | ```
91 |
92 | After running the command, MCP Inspector will be available at http://localhost:5173 (default port: 5173). Open this URL in your browser to start testing. The default timeout for inspector operations is 10000ms (10 seconds), which may not be sufficient for image generation tasks. It's recommended to increase the timeout when testing image generation or other time-consuming operations. You can adjust the timeout by adding a timeout parameter to the URL, for example: http://localhost:5173?timeout=60000 (sets timeout to 60 seconds)
93 |
94 | The MCP Inspector is a powerful development tool that helps you test and debug your MCP server implementation. Key features include:
95 |
96 | - **Interactive Testing**: Test your server's functions directly through a web interface
97 | - **Real-time Feedback**: See immediate results of your function calls and any errors that occur
98 | - **Request/Response Inspection**: View detailed information about requests and responses
99 | - **Function Documentation**: Browse available functions and their parameters
100 | - **Custom Parameters**: Set custom timeout values and other configuration options
101 | - **History Tracking**: Keep track of your previous function calls and their results
102 |
103 | For detailed information about using the MCP Inspector and its features, visit the [official MCP documentation](https://modelcontextprotocol.io/docs/tools/inspector).
104 |
105 | ## Usage
106 |
107 | ### Connecting to Claude Desktop
108 |
109 | Add this to your Claude Desktop configuration file (`~/Library/Application Support/Claude/claude_desktop_config.json` on macOS or `%APPDATA%\Claude\claude_desktop_config.json` on Windows):
110 |
111 | ```json
112 | {
113 | "mcpServers": {
114 | "piapi": {
115 | "command": "node",
116 | "args": ["/absolute/path/to/piapi-mcp-server/dist/index.js"],
117 | "env": {
118 | "PIAPI_API_KEY": "your_api_key_here"
119 | }
120 | }
121 | }
122 | }
123 | ```
124 |
125 | After updating your configuration file, you need to restart Claude for Desktop. Upon restarting, you should see a hammer icon in the bottom right corner of the input box.
126 | For more detailed information, visit the [official MCP documentation](https://modelcontextprotocol.io/quickstart/user)
127 |
128 | ### Connecting to Cursor
129 |
130 | Note: Following guide is based on Cursor 0.47.5. Features and behaviors may vary in different versions.
131 |
132 | To configure the MCP server:
133 |
134 | 1. Navigate to: File > Preferences > Cursor Settings, or use the shortcut key `Ctrl+Shift+J`
135 | 2. Select "MCP" tab on the left panel
136 | 3. Click "Add new global MCP server" button in the top right
137 | 4. Add your configuration in the opened mcp.json file
138 |
139 | ```json
140 | {
141 | "mcpServers": {
142 | "piapi": {
143 | "command": "node",
144 | "args": ["/absolute/path/to/piapi-mcp-server/dist/index.js"],
145 | "env": {
146 | "PIAPI_API_KEY": "your_api_key_here"
147 | }
148 | }
149 | }
150 | }
151 | ```
152 |
153 | 5. After configuration, you'll see a "piapi" entry in MCP Servers page
154 | 6. Click the Refresh button on the entry or restart Cursor to connect to the piapi server
155 |
156 | To test the piapi image generation:
157 |
158 | 1. Open and select "Agent mode" in Cursor Chat, or use the shortcut key `Ctrl+I`
159 | 2. Enter a test prompt, for example: "generate image of a dog"
160 | 3. The image will be generated based on your prompt using piapi server
161 |
162 | To disable the piapi server:
163 |
164 | 1. Navigate to the MCP Servers page in Cursor Settings
165 | 2. Find the "piapi" entry in the server list
166 | 3. Click the "Enabled" toggle button to switch it to "Disabled"
167 |
168 | ## Development
169 |
170 | ### Project Structure
171 |
172 | ```
173 | piapi-mcp-server/
174 | ├── assets/
175 | ├── src/
176 | │ ├── index.ts # Main server entry point
177 | ├── package.json
178 | ├── tsconfig.json
179 | └── .env.example
180 | ```
181 |
182 | ## License
183 |
184 | MIT
185 |
```
--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "compilerOptions": {
3 | "target": "ES2022",
4 | "module": "ES2022",
5 | "moduleResolution": "node",
6 | "outDir": "./dist",
7 | "rootDir": "./src",
8 | "strict": true,
9 | "esModuleInterop": true,
10 | "skipLibCheck": true
11 | },
12 | "include": ["src/**/*"],
13 | "exclude": ["node_modules"]
14 | }
```
--------------------------------------------------------------------------------
/package.json:
--------------------------------------------------------------------------------
```json
1 | {
2 | "name": "piapi-mcp-server",
3 | "version": "1.0.0",
4 | "description": "MCP Server for PiAPI image generation",
5 | "type": "module",
6 | "main": "dist/index.js",
7 | "engines": {
8 | "node": ">=18.15.0"
9 | },
10 | "scripts": {
11 | "build": "tsc",
12 | "inspect": "npx fastmcp inspect dist/index.js"
13 | },
14 | "dependencies": {
15 | "dotenv": "^16.3.1",
16 | "fastmcp": "^1.1.0",
17 | "zod": "^3.24.1"
18 | },
19 | "devDependencies": {
20 | "@types/node": "^18.19.0",
21 | "typescript": "^5.3.3"
22 | }
23 | }
24 |
```
--------------------------------------------------------------------------------
/smithery.yaml:
--------------------------------------------------------------------------------
```yaml
1 | # Smithery configuration file: https://smithery.ai/docs/config#smitheryyaml
2 |
3 | startCommand:
4 | type: stdio
5 | configSchema:
6 | # JSON Schema defining the configuration options for the MCP.
7 | type: object
8 | required:
9 | - piapiApiKey
10 | properties:
11 | piapiApiKey:
12 | type: string
13 | description: The API key for accessing the PiAPI service.
14 | commandFunction:
15 | # A function that produces the CLI command to start the MCP on stdio.
16 | |-
17 | (config) => ({ command: 'node', args: ['dist/index.js'], env: { PIAPI_API_KEY: config.piapiApiKey } })
```
--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
```dockerfile
1 | # Generated by https://smithery.ai. See: https://smithery.ai/docs/config#dockerfile
2 | # Use the official Node.js image.
3 | FROM node:16-alpine AS builder
4 |
5 | # Set the working directory in the container
6 | WORKDIR /app
7 |
8 | # Copy the package.json and package-lock.json files
9 | COPY package*.json ./
10 |
11 | # Install dependencies
12 | RUN npm install
13 |
14 | # Copy the rest of the application source code
15 | COPY . .
16 |
17 | # Build the application
18 | RUN npm run build
19 |
20 | # Use a separate runtime image for the final build
21 | FROM node:16-alpine AS runner
22 |
23 | # Set the working directory in the container
24 | WORKDIR /app
25 |
26 | # Copy built files and node_modules from the builder stage
27 | COPY --from=builder /app/dist ./dist
28 | COPY --from=builder /app/node_modules ./node_modules
29 | COPY --from=builder /app/package.json ./
30 |
31 | # Copy the .env file (if any) - ensure you have a .dockerignore to exclude any sensitive files
32 | COPY .env ./
33 |
34 | # Expose port (if the server listens on a specific port, specify it here)
35 | # EXPOSE 3000
36 |
37 | # Set environment variables if required
38 | ENV NODE_ENV=production
39 |
40 | # Run the application
41 | CMD ["node", "dist/index.js"]
```
--------------------------------------------------------------------------------
/src/index.ts:
--------------------------------------------------------------------------------
```typescript
1 | import { config } from "dotenv";
2 | import { FastMCP, imageContent, Progress, UserError, Content } from "fastmcp";
3 | import { z } from "zod";
4 | // Load environment variables
5 | config();
6 |
7 | if (!process.env.PIAPI_API_KEY) {
8 | console.error("Error: PIAPI API key not set");
9 | process.exit(1);
10 | }
11 |
12 | // Parse command line arguments for environment
13 | const args = process.argv.slice(2);
14 | const envArg = args.find(arg => arg.startsWith('--env='));
15 | const envValue = envArg ? envArg.split('=')[1] : process.env.NODE_ENV;
16 |
17 | const apiKey: string = process.env.PIAPI_API_KEY;
18 | const isProduction = envValue === 'production';
19 |
20 | // Configure logging levels based on environment
21 | const logger = {
22 | debug: (msg: string) => {
23 | if (!isProduction) process.stderr.write(`[DEBUG] ${msg}\n`);
24 | },
25 | info: (msg: string) => {
26 | if (!isProduction) process.stderr.write(`[INFO] ${msg}\n`);
27 | },
28 | warn: (msg: string) => process.stderr.write(`[WARN] ${msg}\n`),
29 | error: (msg: string) => process.stderr.write(`[ERROR] ${msg}\n`),
30 | };
31 |
32 | // Log environment information
33 | logger.info(`Running in ${isProduction ? 'production' : 'development'} mode`);
34 |
35 | const server = new FastMCP({
36 | name: "piapi",
37 | version: "1.0.0",
38 | });
39 |
40 | registerTools(server);
41 |
42 | // Start the server
43 | async function main() {
44 | try {
45 | await server.start({
46 | transportType: "stdio",
47 | });
48 | } catch (error) {
49 | console.error("Failed to start server:", error);
50 | process.exit(1);
51 | }
52 | }
53 |
54 | main().catch((error) => {
55 | console.error("Fatal error in main():", error);
56 | process.exit(1);
57 | });
58 |
59 | // Register Tools
60 | function registerTools(server: FastMCP) {
61 | registerGeneralTool(server);
62 | registerImageTool(server);
63 | registerVideoTool(server);
64 | registerFluxTool(server);
65 | registerHunyuanTool(server);
66 | registerSkyreelsTool(server);
67 | registerWanTool(server);
68 | registerMMAudioTool(server);
69 | registerTTSTool(server);
70 | registerMidjourneyTool(server);
71 | registerKlingTool(server);
72 | registerLumaTool(server);
73 | registerSunoTool(server);
74 | registerTrellisTool(server);
75 | registerHailuoTool(server);
76 | }
77 |
78 | // Tool Definitions
79 |
80 | function registerGeneralTool(server: FastMCP) {
81 | server.addTool({
82 | name: "show_image",
83 | description:
84 | "Show an image with pixels less than 768*1024 due to Claude limitation",
85 | parameters: z.object({
86 | url: z.string().url().describe("The URL of the image to show"),
87 | }),
88 | execute: async (args) => {
89 | return imageContent({ url: args.url });
90 | },
91 | });
92 | }
93 |
94 | interface BaseConfig {
95 | maxAttempts: number;
96 | timeout: number; // in seconds
97 | }
98 |
99 | const IMAGE_TOOL_CONFIG: Record<string, BaseConfig> = {
100 | faceswap: { maxAttempts: 30, timeout: 60 },
101 | rmbg: { maxAttempts: 30, timeout: 60 },
102 | segment: { maxAttempts: 30, timeout: 60 },
103 | upscale: { maxAttempts: 30, timeout: 60 },
104 | };
105 |
106 | function registerImageTool(server: FastMCP) {
107 | server.addTool({
108 | name: "image_faceswap",
109 | description: "Faceswap an image",
110 | parameters: z.object({
111 | swapImage: z.string().url().describe("The URL of the image to swap"),
112 | targetImage: z.string().url().describe("The URL of the target image"),
113 | }),
114 | execute: async (args, { log, reportProgress }) => {
115 | // Create image generation task
116 | if (!args.swapImage || !args.targetImage) {
117 | throw new UserError("Swap image and target image are required");
118 | }
119 | const config = IMAGE_TOOL_CONFIG["faceswap"];
120 |
121 | const requestBody = JSON.stringify({
122 | model: "Qubico/image-toolkit",
123 | task_type: "face-swap",
124 | input: {
125 | swap_image: args.swapImage,
126 | target_image: args.targetImage,
127 | },
128 | });
129 |
130 | const { taskId, usage, output } = await handleTask(
131 | log,
132 | reportProgress,
133 | requestBody,
134 | config
135 | );
136 |
137 | const urls = parseImageOutput(taskId, output, log);
138 | return {
139 | content: [
140 | {
141 | type: "text",
142 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
143 | "\n"
144 | )}`,
145 | },
146 | ],
147 | };
148 | },
149 | });
150 | server.addTool({
151 | name: "image_rmbg",
152 | description: "Remove the background of an image",
153 | parameters: z.object({
154 | image: z
155 | .string()
156 | .url()
157 | .describe("The URL of the image to remove the background"),
158 | }),
159 | execute: async (args, { log, reportProgress }) => {
160 | // Create image generation task
161 | if (!args.image) {
162 | throw new UserError("Image URL is required");
163 | }
164 | const config = IMAGE_TOOL_CONFIG["rmbg"];
165 |
166 | const requestBody = JSON.stringify({
167 | model: "Qubico/image-toolkit",
168 | task_type: "background-remove",
169 | input: {
170 | image: args.image,
171 | },
172 | });
173 |
174 | const { taskId, usage, output } = await handleTask(
175 | log,
176 | reportProgress,
177 | requestBody,
178 | config
179 | );
180 |
181 | const urls = parseImageOutput(taskId, output, log);
182 | return {
183 | content: [
184 | {
185 | type: "text",
186 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
187 | "\n"
188 | )}`,
189 | },
190 | ],
191 | };
192 | },
193 | });
194 | server.addTool({
195 | name: "image_segment",
196 | description: "Segment an image",
197 | parameters: z.object({
198 | image: z.string().url().describe("The URL of the image to segment"),
199 | prompt: z.string().describe("The prompt to segment the image"),
200 | negativePrompt: z
201 | .string()
202 | .optional()
203 | .describe("The negative prompt to segment the image"),
204 | segmentFactor: z
205 | .number()
206 | .optional()
207 | .default(-15)
208 | .describe("The factor to segment the image"),
209 | }),
210 | execute: async (args, { log, reportProgress }) => {
211 | // Create image generation task
212 | if (!args.image || !args.prompt) {
213 | throw new UserError("Image URL and prompt are required");
214 | }
215 | const config = IMAGE_TOOL_CONFIG["segment"];
216 |
217 | const requestBody = JSON.stringify({
218 | model: "Qubico/image-toolkit",
219 | task_type: "segment",
220 | input: {
221 | image: args.image,
222 | prompt: args.prompt,
223 | negative_prompt: args.negativePrompt,
224 | segment_factor: args.segmentFactor,
225 | },
226 | });
227 |
228 | const { taskId, usage, output } = await handleTask(
229 | log,
230 | reportProgress,
231 | requestBody,
232 | config
233 | );
234 |
235 | const urls = parseImageOutput(taskId, output, log);
236 | return {
237 | content: [
238 | {
239 | type: "text",
240 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
241 | "\n"
242 | )}`,
243 | },
244 | ],
245 | };
246 | },
247 | });
248 | server.addTool({
249 | name: "image_upscale",
250 | description: "Upscale an image to a higher resolution",
251 | parameters: z.object({
252 | image: z.string().url().describe("The URL of the image to upscale"),
253 | scale: z
254 | .number()
255 | .pipe(z.number().min(2).max(10))
256 | .optional()
257 | .default(2)
258 | .describe("The scale of the image to upscale, defaults to 2"),
259 | faceEnhance: z
260 | .boolean()
261 | .optional()
262 | .default(false)
263 | .describe("Whether to enhance the face of the image"),
264 | }),
265 | execute: async (args, { log, reportProgress }) => {
266 | // Create image generation task
267 | if (!args.image) {
268 | throw new UserError("Image URL is required");
269 | }
270 | const config = IMAGE_TOOL_CONFIG["upscale"];
271 |
272 | const requestBody = JSON.stringify({
273 | model: "Qubico/image-toolkit",
274 | task_type: "upscale",
275 | input: {
276 | image: args.image,
277 | scale: args.scale,
278 | face_enhance: args.faceEnhance,
279 | },
280 | });
281 |
282 | const { taskId, usage, output } = await handleTask(
283 | log,
284 | reportProgress,
285 | requestBody,
286 | config
287 | );
288 |
289 | const urls = parseImageOutput(taskId, output, log);
290 | return {
291 | content: [
292 | {
293 | type: "text",
294 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
295 | "\n"
296 | )}`,
297 | },
298 | ],
299 | };
300 | },
301 | });
302 | }
303 |
304 | const VIDEO_TOOL_CONFIG: Record<string, BaseConfig> = {
305 | faceswap: { maxAttempts: 30, timeout: 600 },
306 | upscale: { maxAttempts: 30, timeout: 300 },
307 | };
308 |
309 | function registerVideoTool(server: FastMCP) {
310 | server.addTool({
311 | name: "video_faceswap",
312 | description: "Faceswap a video",
313 | parameters: z.object({
314 | swapImage: z.string().url().describe("The URL of the image to swap"),
315 | targetVideo: z
316 | .string()
317 | .url()
318 | .describe("The URL of the video to faceswap"),
319 | }),
320 | execute: async (args, { log, reportProgress }) => {
321 | // Create image generation task
322 | if (!args.swapImage || !args.targetVideo) {
323 | throw new UserError("Swap image and target video are required");
324 | }
325 | const config = VIDEO_TOOL_CONFIG["faceswap"];
326 |
327 | const requestBody = JSON.stringify({
328 | model: "Qubico/video-toolkit",
329 | task_type: "face-swap",
330 | input: {
331 | swap_image: args.swapImage,
332 | target_video: args.targetVideo,
333 | },
334 | });
335 |
336 | const { taskId, usage, output } = await handleTask(
337 | log,
338 | reportProgress,
339 | requestBody,
340 | config
341 | );
342 |
343 | const url = parseVideoOutput(taskId, output, log);
344 | return {
345 | content: [
346 | {
347 | type: "text",
348 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
349 | },
350 | ],
351 | };
352 | },
353 | });
354 | server.addTool({
355 | name: "video_upscale",
356 | description: "Upscale video resolution to 2x",
357 | parameters: z.object({
358 | video: z.string().url().describe("The URL of the video to upscale"),
359 | }),
360 | execute: async (args, { log, reportProgress }) => {
361 | // Create image generation task
362 | if (!args.video) {
363 | throw new UserError("Video URL is required");
364 | }
365 | const config = VIDEO_TOOL_CONFIG["upscale"];
366 |
367 | const requestBody = JSON.stringify({
368 | model: "Qubico/video-toolkit",
369 | task_type: "upscale",
370 | input: {
371 | video: args.video,
372 | },
373 | });
374 |
375 | const { taskId, usage, output } = await handleTask(
376 | log,
377 | reportProgress,
378 | requestBody,
379 | config
380 | );
381 |
382 | const url = parseVideoOutput(taskId, output, log);
383 | return {
384 | content: [
385 | {
386 | type: "text",
387 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
388 | },
389 | ],
390 | };
391 | },
392 | });
393 | }
394 |
395 | interface FluxConfig extends BaseConfig {
396 | defaultSteps: number;
397 | maxSteps: number;
398 | }
399 |
400 | const FLUX_MODEL_CONFIG: Record<string, FluxConfig> = {
401 | schnell: { defaultSteps: 4, maxSteps: 10, maxAttempts: 30, timeout: 60 },
402 | dev: { defaultSteps: 25, maxSteps: 40, maxAttempts: 30, timeout: 120 },
403 | inpaint: { defaultSteps: 25, maxSteps: 40, maxAttempts: 30, timeout: 120 },
404 | outpaint: { defaultSteps: 25, maxSteps: 40, maxAttempts: 30, timeout: 120 },
405 | variation: { defaultSteps: 25, maxSteps: 40, maxAttempts: 30, timeout: 120 },
406 | controlnet: { defaultSteps: 25, maxSteps: 40, maxAttempts: 30, timeout: 180 },
407 | };
408 |
409 | function registerFluxTool(server: FastMCP) {
410 | server.addTool({
411 | name: "generate_image",
412 | description: "Generate a image using Qubico Flux",
413 | parameters: z.object({
414 | prompt: z.string().describe("The prompt to generate an image from"),
415 | negativePrompt: z
416 | .string()
417 | .optional()
418 | .default("chaos, bad photo, low quality, low resolution")
419 | .describe("The negative prompt to generate an image from"),
420 | referenceImage: z
421 | .string()
422 | .url()
423 | .optional()
424 | .describe(
425 | "The reference image to generate an image from, must be a valid image url"
426 | ),
427 | width: z
428 | .union([z.string(), z.number()])
429 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
430 | .pipe(z.number().min(128).max(1024))
431 | .optional()
432 | .default(1024)
433 | .describe(
434 | "The width of the image to generate, must be between 128 and 1024, defaults to 1024"
435 | ),
436 | height: z
437 | .union([z.string(), z.number()])
438 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
439 | .pipe(z.number().min(128).max(1024))
440 | .optional()
441 | .default(1024)
442 | .describe(
443 | "The height of the image to generate, must be between 128 and 1024, defaults to 1024"
444 | ),
445 | steps: z
446 | .union([z.string(), z.number()])
447 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
448 | .optional()
449 | .default(0)
450 | .describe("The number of steps to generate the image"),
451 | lora: z
452 | .enum([
453 | "",
454 | "mystic-realism",
455 | "ob3d-isometric-3d-room",
456 | "remes-abstract-poster-style",
457 | "paper-quilling-and-layering-style",
458 | ])
459 | .optional()
460 | .default("")
461 | .describe(
462 | "The lora to use for image generation, only available for 'dev' model, defaults to ''"
463 | ),
464 | model: z
465 | .enum(["schnell", "dev"])
466 | .optional()
467 | .default("schnell")
468 | .describe(
469 | "The model to use for image generation, 'schnell' is faster and cheaper but less detailed, 'dev' is slower but more detailed"
470 | ),
471 | }),
472 | execute: async (args, { log, reportProgress }) => {
473 | // Create image generation task
474 | if (!args.prompt) {
475 | throw new UserError("Prompt is required");
476 | }
477 | const config = FLUX_MODEL_CONFIG[args.model];
478 | let steps = args.steps || config.defaultSteps;
479 | steps = Math.min(steps, config.maxSteps);
480 |
481 | let requestBody = "";
482 | if (args.lora !== "") {
483 | requestBody = JSON.stringify({
484 | model: "Qubico/flux1-dev-advanced",
485 | task_type: args.referenceImage ? "img2img-lora" : "txt2img-lora",
486 | input: {
487 | prompt: args.prompt,
488 | negative_prompt: args.negativePrompt,
489 | image: args.referenceImage,
490 | width: args.width,
491 | height: args.height,
492 | steps: steps,
493 | lora_settings: [
494 | {
495 | lora_type: args.lora,
496 | },
497 | ],
498 | },
499 | });
500 | } else {
501 | requestBody = JSON.stringify({
502 | model:
503 | args.model === "schnell"
504 | ? "Qubico/flux1-schnell"
505 | : "Qubico/flux1-dev",
506 | task_type: args.referenceImage ? "img2img" : "txt2img",
507 | input: {
508 | prompt: args.prompt,
509 | negative_prompt: args.negativePrompt,
510 | image: args.referenceImage,
511 | width: args.width,
512 | height: args.height,
513 | steps: steps,
514 | },
515 | });
516 | }
517 |
518 | const { taskId, usage, output } = await handleTask(
519 | log,
520 | reportProgress,
521 | requestBody,
522 | config
523 | );
524 |
525 | const urls = parseImageOutput(taskId, output, log);
526 | return {
527 | content: [
528 | {
529 | type: "text",
530 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
531 | "\n"
532 | )}`,
533 | },
534 | ],
535 | };
536 | },
537 | });
538 | server.addTool({
539 | name: "modify_image",
540 | description: "Modify a image using Qubico Flux, inpaint or outpaint",
541 | parameters: z.object({
542 | prompt: z.string().describe("The prompt to modify an image from"),
543 | negativePrompt: z
544 | .string()
545 | .optional()
546 | .default("chaos, bad photo, low quality, low resolution")
547 | .describe("The negative prompt to modify an image from"),
548 | referenceImage: z
549 | .string()
550 | .url()
551 | .describe(
552 | "The reference image to modify an image from, must be a valid image url"
553 | ),
554 | paddingLeft: z
555 | .union([z.string(), z.number()])
556 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
557 | .optional()
558 | .default(0)
559 | .describe("The padding left of the image, only available for outpaint"),
560 | paddingRight: z
561 | .union([z.string(), z.number()])
562 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
563 | .optional()
564 | .default(0)
565 | .describe(
566 | "The padding right of the image, only available for outpaint"
567 | ),
568 | paddingTop: z
569 | .union([z.string(), z.number()])
570 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
571 | .optional()
572 | .default(0)
573 | .describe("The padding top of the image, only available for outpaint"),
574 | paddingBottom: z
575 | .union([z.string(), z.number()])
576 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
577 | .optional()
578 | .default(0)
579 | .describe(
580 | "The padding bottom of the image, only available for outpaint"
581 | ),
582 | steps: z
583 | .union([z.string(), z.number()])
584 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
585 | .optional()
586 | .default(0)
587 | .describe("The number of steps to generate the image"),
588 | model: z
589 | .enum(["inpaint", "outpaint"])
590 | .describe("The model to use for image modification"),
591 | }),
592 | execute: async (args, { log, reportProgress }) => {
593 | // Create image generation task
594 | if (!args.prompt) {
595 | throw new UserError("Prompt is required");
596 | } else if (!args.referenceImage) {
597 | throw new UserError("Reference image is required");
598 | }
599 | const config = FLUX_MODEL_CONFIG[args.model];
600 | let steps = args.steps || config.defaultSteps;
601 | steps = Math.min(steps, config.maxSteps);
602 |
603 | let requestBody = "";
604 | if (args.model === "inpaint") {
605 | requestBody = JSON.stringify({
606 | model: "Qubico/flux1-dev-advanced",
607 | task_type: "fill-inpaint",
608 | input: {
609 | prompt: args.prompt,
610 | negative_prompt: args.negativePrompt,
611 | image: args.referenceImage,
612 | steps: steps,
613 | },
614 | });
615 | } else {
616 | requestBody = JSON.stringify({
617 | model: "Qubico/flux1-dev-advanced",
618 | task_type: "fill-outpaint",
619 | input: {
620 | prompt: args.prompt,
621 | negative_prompt: args.negativePrompt,
622 | image: args.referenceImage,
623 | steps: steps,
624 | custom_settings: [
625 | {
626 | setting_type: "outpaint",
627 | outpaint_left: args.paddingLeft,
628 | outpaint_right: args.paddingRight,
629 | outpaint_top: args.paddingTop,
630 | outpaint_bottom: args.paddingBottom,
631 | },
632 | ],
633 | },
634 | });
635 | }
636 |
637 | const { taskId, usage, output } = await handleTask(
638 | log,
639 | reportProgress,
640 | requestBody,
641 | config
642 | );
643 |
644 | const urls = parseImageOutput(taskId, output, log);
645 | return {
646 | content: [
647 | {
648 | type: "text",
649 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
650 | "\n"
651 | )}`,
652 | },
653 | ],
654 | };
655 | },
656 | });
657 | server.addTool({
658 | name: "derive_image",
659 | description: "Derive a image using Qubico Flux, variation",
660 | parameters: z.object({
661 | prompt: z.string().describe("The prompt to derive an image from"),
662 | negativePrompt: z
663 | .string()
664 | .optional()
665 | .default("chaos, bad photo, low quality, low resolution")
666 | .describe("The negative prompt to derive an image from"),
667 | referenceImage: z
668 | .string()
669 | .url()
670 | .describe(
671 | "The reference image to derive an image from, must be a valid image url"
672 | ),
673 | width: z
674 | .union([z.string(), z.number()])
675 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
676 | .pipe(z.number().min(128).max(1024))
677 | .optional()
678 | .default(1024)
679 | .describe(
680 | "The width of the image to generate, must be between 128 and 1024, defaults to 1024"
681 | ),
682 | height: z
683 | .union([z.string(), z.number()])
684 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
685 | .pipe(z.number().min(128).max(1024))
686 | .optional()
687 | .default(1024)
688 | .describe(
689 | "The height of the image to generate, must be between 128 and 1024, defaults to 1024"
690 | ),
691 | steps: z
692 | .union([z.string(), z.number()])
693 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
694 | .optional()
695 | .default(0)
696 | .describe("The number of steps to generate the image"),
697 | }),
698 | execute: async (args, { log, reportProgress }) => {
699 | // Create image generation task
700 | if (!args.prompt) {
701 | throw new UserError("Prompt is required");
702 | } else if (!args.referenceImage) {
703 | throw new UserError("Reference image is required");
704 | }
705 | const config = FLUX_MODEL_CONFIG["variation"];
706 | let steps = args.steps || config.defaultSteps;
707 | steps = Math.min(steps, config.maxSteps);
708 |
709 | const requestBody = JSON.stringify({
710 | model: "Qubico/flux1-dev-advanced",
711 | task_type: "redux-variation",
712 | input: {
713 | prompt: args.prompt,
714 | negative_prompt: args.negativePrompt,
715 | image: args.referenceImage,
716 | width: args.width,
717 | height: args.height,
718 | steps: steps,
719 | },
720 | });
721 |
722 | const { taskId, usage, output } = await handleTask(
723 | log,
724 | reportProgress,
725 | requestBody,
726 | config
727 | );
728 |
729 | const urls = parseImageOutput(taskId, output, log);
730 | return {
731 | content: [
732 | {
733 | type: "text",
734 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
735 | "\n"
736 | )}`,
737 | },
738 | ],
739 | };
740 | },
741 | });
742 | server.addTool({
743 | name: "generate_image_controlnet",
744 | description: "Generate a image using Qubico Flux with ControlNet",
745 | parameters: z.object({
746 | prompt: z.string().describe("The prompt to generate an image from"),
747 | negativePrompt: z
748 | .string()
749 | .optional()
750 | .default("chaos, bad photo, low quality, low resolution")
751 | .describe("The negative prompt to generate an image from"),
752 | referenceImage: z
753 | .string()
754 | .url()
755 | .describe(
756 | "The reference image to generate an image from, must be a valid image url"
757 | ),
758 | width: z
759 | .union([z.string(), z.number()])
760 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
761 | .pipe(z.number().min(128).max(1024))
762 | .optional()
763 | .default(1024)
764 | .describe(
765 | "The width of the image to generate, must be between 128 and 1024, defaults to 1024"
766 | ),
767 | height: z
768 | .union([z.string(), z.number()])
769 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
770 | .pipe(z.number().min(128).max(1024))
771 | .optional()
772 | .default(1024)
773 | .describe(
774 | "The height of the image to generate, must be between 128 and 1024, defaults to 1024"
775 | ),
776 | steps: z
777 | .union([z.string(), z.number()])
778 | .transform((val) => (typeof val === "string" ? parseInt(val) : val))
779 | .optional()
780 | .default(0)
781 | .describe("The number of steps to generate the image"),
782 | lora: z
783 | .enum([
784 | "",
785 | "mystic-realism",
786 | "ob3d-isometric-3d-room",
787 | "remes-abstract-poster-style",
788 | "paper-quilling-and-layering-style",
789 | ])
790 | .optional()
791 | .default("")
792 | .describe("The lora to use for image generation"),
793 | controlType: z
794 | .enum(["depth", "canny", "hed", "openpose"])
795 | .optional()
796 | .default("depth")
797 | .describe("The control type to use for image generation"),
798 | }),
799 | execute: async (args, { log, reportProgress }) => {
800 | // Create image generation task
801 | if (!args.prompt) {
802 | throw new UserError("Prompt is required");
803 | } else if (!args.referenceImage) {
804 | throw new UserError("Reference image is required");
805 | }
806 | const config = FLUX_MODEL_CONFIG["controlnet"];
807 | let steps = args.steps || config.defaultSteps;
808 | steps = Math.min(steps, config.maxSteps);
809 |
810 | const requestBody = JSON.stringify({
811 | model: "Qubico/flux1-dev-advanced",
812 | task_type: "controlnet-lora",
813 | input: {
814 | prompt: args.prompt,
815 | negative_prompt: args.negativePrompt,
816 | width: args.width,
817 | height: args.height,
818 | steps: steps,
819 | lora_settings: args.lora !== "" ? [{ lora_type: args.lora }] : [],
820 | control_net_settings: [
821 | {
822 | control_type: args.controlType,
823 | control_image: args.referenceImage,
824 | },
825 | ],
826 | },
827 | });
828 |
829 | const { taskId, usage, output } = await handleTask(
830 | log,
831 | reportProgress,
832 | requestBody,
833 | config
834 | );
835 |
836 | const urls = parseImageOutput(taskId, output, log);
837 | return {
838 | content: [
839 | {
840 | type: "text",
841 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
842 | "\n"
843 | )}`,
844 | },
845 | ],
846 | };
847 | },
848 | });
849 | }
850 |
851 | interface HunyuanConfig extends BaseConfig {
852 | taskType: string;
853 | }
854 |
855 | const HUNYUAN_MODEL_CONFIG: Record<string, HunyuanConfig> = {
856 | hunyuan: { maxAttempts: 60, timeout: 900, taskType: "txt2video" },
857 | fastHunyuan: { maxAttempts: 60, timeout: 600, taskType: "fast-txt2video" },
858 | hunyuanConcat: {
859 | maxAttempts: 60,
860 | timeout: 900,
861 | taskType: "img2video-concat",
862 | },
863 | hunyuanReplace: {
864 | maxAttempts: 60,
865 | timeout: 900,
866 | taskType: "img2video-replace",
867 | },
868 | };
869 |
870 | function registerHunyuanTool(server: FastMCP) {
871 | server.addTool({
872 | name: "generate_video_hunyuan",
873 | description: "Generate a video using Qubico Hunyuan",
874 | parameters: z.object({
875 | prompt: z.string().describe("The prompt to generate a video from"),
876 | negativePrompt: z
877 | .string()
878 | .describe("The negative prompt to generate a video from")
879 | .optional()
880 | .default("chaos, bad video, low quality, low resolution"),
881 | referenceImage: z
882 | .string()
883 | .url()
884 | .optional()
885 | .describe(
886 | "The reference image to generate a video from, must be a valid image url"
887 | ),
888 | aspectRatio: z
889 | .enum(["16:9", "1:1", "9:16"])
890 | .optional()
891 | .default("16:9")
892 | .describe(
893 | "The aspect ratio of the video to generate, must be either '16:9', '1:1', or '9:16', defaults to '16:9'"
894 | ),
895 | model: z
896 | .enum(["hunyuan", "fastHunyuan", "hunyuanConcat", "hunyuanReplace"])
897 | .optional()
898 | .default("hunyuan")
899 | .describe(
900 | "The model to use for video generation, 'hunyuan' is slower but more detailed, 'fastHunyuan' is faster but less detailed, both for txt2video. 'hunyuanReplace' sticks to reference image, and 'hunyuanConcat' allows for more creative movement, both for img2video"
901 | ),
902 | }),
903 | execute: async (args, { log, reportProgress }) => {
904 | // Create video generation task
905 | if (!args.prompt) {
906 | throw new UserError("Prompt is required");
907 | }
908 | if (
909 | args.referenceImage &&
910 | (args.model === "hunyuan" || args.model === "fastHunyuan")
911 | ) {
912 | log.warn(
913 | "Reference image is not supported for 'hunyuan' or 'fastHunyuan' model, using 'hunyuanConcat' as default"
914 | );
915 | args.model = "hunyuanConcat";
916 | }
917 | const config = HUNYUAN_MODEL_CONFIG[args.model];
918 |
919 | const requestBody = JSON.stringify({
920 | model: "Qubico/hunyuan",
921 | task_type: config.taskType,
922 | input: {
923 | image: args.referenceImage,
924 | prompt: args.prompt,
925 | negative_prompt: args.negativePrompt,
926 | aspect_ratio: args.aspectRatio,
927 | },
928 | });
929 | const { taskId, usage, output } = await handleTask(
930 | log,
931 | reportProgress,
932 | requestBody,
933 | config
934 | );
935 |
936 | const url = parseVideoOutput(taskId, output, log);
937 | return {
938 | content: [
939 | {
940 | type: "text",
941 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
942 | },
943 | ],
944 | };
945 | },
946 | });
947 | }
948 |
949 | const SKYREELS_MODEL_CONFIG: Record<string, BaseConfig> = {
950 | skyreels: { maxAttempts: 30, timeout: 300 },
951 | };
952 |
953 | function registerSkyreelsTool(server: FastMCP) {
954 | server.addTool({
955 | name: "generate_video_skyreels",
956 | description: "Generate a video using Qubico Skyreels",
957 | parameters: z.object({
958 | prompt: z.string().describe("The prompt to generate a video from"),
959 | negativePrompt: z
960 | .string()
961 | .describe("The negative prompt to generate a video from")
962 | .optional()
963 | .default("chaos, bad video, low quality, low resolution"),
964 | aspectRatio: z
965 | .enum(["16:9", "1:1", "9:16"])
966 | .optional()
967 | .default("16:9")
968 | .describe(
969 | "The aspect ratio of the video to generate, must be either '16:9', '1:1', or '9:16', defaults to '16:9'"
970 | ),
971 | referenceImage: z
972 | .string()
973 | .url()
974 | .describe(
975 | "The reference image to generate a video from, must be a valid image url, only available for 'wan14b' model"
976 | ),
977 | }),
978 | execute: async (args, { log, reportProgress }) => {
979 | // Create video generation task
980 | if (!args.prompt || !args.referenceImage) {
981 | throw new UserError("Prompt and reference image are required");
982 | }
983 | const config = SKYREELS_MODEL_CONFIG["skyreels"];
984 |
985 | const requestBody = JSON.stringify({
986 | model: "Qubico/skyreels",
987 | task_type: "img2video",
988 | input: {
989 | prompt: args.prompt,
990 | negative_prompt: args.negativePrompt,
991 | aspect_ratio: args.aspectRatio,
992 | image: args.referenceImage,
993 | },
994 | });
995 | const { taskId, usage, output } = await handleTask(
996 | log,
997 | reportProgress,
998 | requestBody,
999 | config
1000 | );
1001 |
1002 | const url = parseVideoOutput(taskId, output, log);
1003 | return {
1004 | content: [
1005 | {
1006 | type: "text",
1007 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
1008 | },
1009 | ],
1010 | };
1011 | },
1012 | });
1013 | }
1014 |
1015 | const WAN_MODEL_CONFIG: Record<string, BaseConfig> = {
1016 | wan1_3b: { maxAttempts: 30, timeout: 300 },
1017 | wan14b: { maxAttempts: 60, timeout: 900 },
1018 | };
1019 |
1020 | function registerWanTool(server: FastMCP) {
1021 | server.addTool({
1022 | name: "generate_video_wan",
1023 | description: "Generate a video using Qubico Wan",
1024 | parameters: z.object({
1025 | prompt: z.string().describe("The prompt to generate a video from"),
1026 | negativePrompt: z
1027 | .string()
1028 | .describe("The negative prompt to generate a video from")
1029 | .optional()
1030 | .default("chaos, bad video, low quality, low resolution"),
1031 | aspectRatio: z
1032 | .enum(["16:9", "1:1", "9:16"])
1033 | .optional()
1034 | .default("16:9")
1035 | .describe(
1036 | "The aspect ratio of the video to generate, must be either '16:9', '1:1', or '9:16', defaults to '16:9'"
1037 | ),
1038 | referenceImage: z
1039 | .string()
1040 | .url()
1041 | .optional()
1042 | .describe(
1043 | "The reference image to generate a video from, must be a valid image url, only available for 'wan14b' model"
1044 | ),
1045 | model: z
1046 | .enum(["wan1_3b", "wan14b"])
1047 | .optional()
1048 | .default("wan1_3b")
1049 | .describe(
1050 | "The model to use for video generation, must be either 'wan1_3b' or 'wan14b', 'wan1_3b' is faster but less detailed, 'wan14b' is slower but more detailed"
1051 | ),
1052 | }),
1053 | execute: async (args, { log, reportProgress }) => {
1054 | // Create video generation task
1055 | if (!args.prompt) {
1056 | throw new UserError("Prompt is required");
1057 | }
1058 | let taskType =
1059 | args.model === "wan1_3b" ? "txt2video-1.3b" : "txt2video-14b";
1060 | if (args.referenceImage) {
1061 | args.model = "wan14b";
1062 | taskType = "img2video-14b";
1063 | }
1064 | const config = WAN_MODEL_CONFIG[args.model];
1065 |
1066 | const requestBody = JSON.stringify({
1067 | model: "Qubico/wanx",
1068 | task_type: taskType,
1069 | input: {
1070 | prompt: args.prompt,
1071 | negative_prompt: args.negativePrompt,
1072 | aspect_ratio: args.aspectRatio,
1073 | image: args.referenceImage,
1074 | },
1075 | });
1076 | const { taskId, usage, output } = await handleTask(
1077 | log,
1078 | reportProgress,
1079 | requestBody,
1080 | config
1081 | );
1082 |
1083 | const url = parseVideoOutput(taskId, output, log);
1084 | return {
1085 | content: [
1086 | {
1087 | type: "text",
1088 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
1089 | },
1090 | ],
1091 | };
1092 | },
1093 | });
1094 | }
1095 |
1096 | const MMAUDIO_MODEL_CONFIG: Record<string, BaseConfig> = {
1097 | mmaudio: { maxAttempts: 30, timeout: 600 },
1098 | };
1099 |
1100 | function registerMMAudioTool(server: FastMCP) {
1101 | server.addTool({
1102 | name: "generate_music_for_video",
1103 | description: "Generate a music for a video using Qubico MMAudio",
1104 | parameters: z.object({
1105 | prompt: z.string().describe("The prompt to generate a music from"),
1106 | negativePrompt: z
1107 | .string()
1108 | .describe("The negative prompt to generate a music from")
1109 | .optional()
1110 | .default("chaos, bad music"),
1111 | video: z.string().url().describe("The video to generate a music from"),
1112 | }),
1113 | execute: async (args, { log, reportProgress }) => {
1114 | // Create video generation task
1115 | if (!args.prompt || !args.video) {
1116 | throw new UserError("Prompt and video are required");
1117 | }
1118 | const config = MMAUDIO_MODEL_CONFIG["mmaudio"];
1119 |
1120 | const requestBody = JSON.stringify({
1121 | model: "Qubico/mmaudio",
1122 | task_type: "video2audio",
1123 | input: {
1124 | prompt: args.prompt,
1125 | negative_prompt: args.negativePrompt,
1126 | video: args.video,
1127 | },
1128 | });
1129 | const { taskId, usage, output } = await handleTask(
1130 | log,
1131 | reportProgress,
1132 | requestBody,
1133 | config
1134 | );
1135 |
1136 | const url = parseAudioOutput(taskId, output, log);
1137 | return {
1138 | content: [
1139 | {
1140 | type: "text",
1141 | text: `TaskId: ${taskId}\nMusic generated successfully!\nUsage: ${usage} tokens\nMusic url:\n${url}`,
1142 | },
1143 | ],
1144 | };
1145 | },
1146 | });
1147 | }
1148 |
1149 | const TTS_MODEL_CONFIG: Record<string, BaseConfig> = {
1150 | zeroShot: { maxAttempts: 30, timeout: 600 },
1151 | };
1152 |
1153 | function registerTTSTool(server: FastMCP) {
1154 | server.addTool({
1155 | name: "tts_zero_shot",
1156 | description: "Zero-shot TTS using Qubico f5-tts",
1157 | parameters: z.object({
1158 | genText: z.string().describe("The text to generate a speech from"),
1159 | refText: z
1160 | .string()
1161 | .optional()
1162 | .describe(
1163 | "The reference text to generate a speech from, auto detect from refAudio if not provided"
1164 | ),
1165 | refAudio: z
1166 | .string()
1167 | .url()
1168 | .describe("The reference audio to generate a speech from"),
1169 | }),
1170 | execute: async (args, { log, reportProgress }) => {
1171 | // Create video generation task
1172 | if (!args.genText || !args.refAudio) {
1173 | throw new UserError("genText and refAudio are required");
1174 | }
1175 | const config = TTS_MODEL_CONFIG["zeroShot"];
1176 |
1177 | const requestBody = JSON.stringify({
1178 | model: "Qubico/tts",
1179 | task_type: "zero-shot",
1180 | input: {
1181 | gen_text: args.genText,
1182 | ref_text: args.refText,
1183 | ref_audio: args.refAudio,
1184 | },
1185 | });
1186 | const { taskId, usage, output } = await handleTask(
1187 | log,
1188 | reportProgress,
1189 | requestBody,
1190 | config
1191 | );
1192 |
1193 | const url = parseAudioOutput(taskId, output, log);
1194 | return {
1195 | content: [
1196 | {
1197 | type: "text",
1198 | text: `TaskId: ${taskId}\nSpeech generated successfully!\nUsage: ${usage} tokens\nSpeech url:\n${url}`,
1199 | },
1200 | ],
1201 | };
1202 | },
1203 | });
1204 | }
1205 |
1206 | const MIDJOURNEY_MODEL_CONFIG: Record<string, BaseConfig> = {
1207 | imagine: { maxAttempts: 30, timeout: 900 },
1208 | };
1209 |
1210 | function registerMidjourneyTool(server: FastMCP) {
1211 | server.addTool({
1212 | name: "midjourney_imagine",
1213 | description: "Generate a image using Midjourney Imagine",
1214 | parameters: z.object({
1215 | prompt: z.string().describe("The prompt to generate a image from"),
1216 | aspectRatio: z
1217 | .string()
1218 | .optional()
1219 | .describe("The aspect ratio of the image"),
1220 | }),
1221 | execute: async (args, { log, reportProgress }) => {
1222 | // Create video generation task
1223 | if (!args.prompt) {
1224 | throw new UserError("Prompt is required");
1225 | }
1226 | const config = MIDJOURNEY_MODEL_CONFIG["imagine"];
1227 |
1228 | const requestBody = JSON.stringify({
1229 | model: "midjourney",
1230 | task_type: "imagine",
1231 | input: {
1232 | prompt: args.prompt,
1233 | aspect_ratio: args.aspectRatio,
1234 | process_mode: "fast",
1235 | },
1236 | });
1237 | const { taskId, usage, output } = await handleTask(
1238 | log,
1239 | reportProgress,
1240 | requestBody,
1241 | config
1242 | );
1243 |
1244 | const urls = parseImageOutput(taskId, output, log);
1245 | return {
1246 | content: [
1247 | {
1248 | type: "text",
1249 | text: `TaskId: ${taskId}\nImage generated successfully!\nUsage: ${usage} tokens\nImage urls:\n${urls.join(
1250 | "\n"
1251 | )}`,
1252 | },
1253 | ],
1254 | };
1255 | },
1256 | });
1257 | }
1258 |
1259 | const KLING_MODEL_CONFIG: Record<string, BaseConfig> = {
1260 | video: { maxAttempts: 30, timeout: 900 },
1261 | effect: { maxAttempts: 30, timeout: 900 },
1262 | };
1263 |
1264 | function registerKlingTool(server: FastMCP) {
1265 | server.addTool({
1266 | name: "generate_video_kling",
1267 | description: "Generate a video using Kling",
1268 | parameters: z.object({
1269 | prompt: z.string().describe("The prompt to generate a video from"),
1270 | negativePrompt: z
1271 | .string()
1272 | .describe("The negative prompt to generate a video from")
1273 | .optional()
1274 | .default("chaos, bad video, low quality, low resolution"),
1275 | referenceImage: z
1276 | .string()
1277 | .url()
1278 | .optional()
1279 | .describe("The reference image to generate a video with"),
1280 | aspectRatio: z
1281 | .enum(["16:9", "1:1", "9:16"])
1282 | .optional()
1283 | .default("16:9")
1284 | .describe(
1285 | "The aspect ratio of the video to generate, must be either '16:9', '1:1', or '9:16', defaults to '16:9'"
1286 | ),
1287 | duration: z
1288 | .enum(["5s", "10s"])
1289 | .optional()
1290 | .default("5s")
1291 | .describe(
1292 | "The duration of the video to generate, defaults to 5 seconds"
1293 | ),
1294 | }),
1295 | execute: async (args, { log, reportProgress }) => {
1296 | // Create video generation task
1297 | if (!args.prompt) {
1298 | throw new UserError("Prompt is required");
1299 | }
1300 | const config = KLING_MODEL_CONFIG["video"];
1301 |
1302 | const requestBody = JSON.stringify({
1303 | model: "kling",
1304 | task_type: "video_generation",
1305 | input: {
1306 | prompt: args.prompt,
1307 | negative_prompt: args.negativePrompt,
1308 | aspect_ratio: args.aspectRatio,
1309 | image_url: args.referenceImage,
1310 | duration: args.duration === "5s" ? 5 : 10,
1311 | },
1312 | });
1313 | const { taskId, usage, output } = await handleTask(
1314 | log,
1315 | reportProgress,
1316 | requestBody,
1317 | config
1318 | );
1319 |
1320 | const urls = parseKlingOutput(taskId, output, log);
1321 | return {
1322 | content: [
1323 | {
1324 | type: "text",
1325 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo urls:\n${urls.join(
1326 | "\n"
1327 | )}`,
1328 | },
1329 | ],
1330 | };
1331 | },
1332 | });
1333 | server.addTool({
1334 | name: "generate_video_effect_kling",
1335 | description: "Generate a video effect using Kling",
1336 | parameters: z.object({
1337 | image: z
1338 | .string()
1339 | .url()
1340 | .describe("The reference image to generate a video effect from"),
1341 | effectName: z
1342 | .enum(["squish", "expansion"])
1343 | .optional()
1344 | .default("squish")
1345 | .describe(
1346 | "The effect name to generate, must be either 'squish' or 'expansion', defaults to 'squish'"
1347 | ),
1348 | }),
1349 | execute: async (args, { log, reportProgress }) => {
1350 | // Create video generation task
1351 | if (!args.image) {
1352 | throw new UserError("Image is required");
1353 | }
1354 | const config = KLING_MODEL_CONFIG["effect"];
1355 |
1356 | const requestBody = JSON.stringify({
1357 | model: "kling",
1358 | task_type: "effects",
1359 | input: {
1360 | image_url: args.image,
1361 | effect: args.effectName,
1362 | },
1363 | });
1364 | const { taskId, usage, output } = await handleTask(
1365 | log,
1366 | reportProgress,
1367 | requestBody,
1368 | config
1369 | );
1370 |
1371 | const urls = parseKlingOutput(taskId, output, log);
1372 | return {
1373 | content: [
1374 | {
1375 | type: "text",
1376 | text: `TaskId: ${taskId}\nVideo effect generated successfully!\nUsage: ${usage} tokens\nVideo urls:\n${urls.join(
1377 | "\n"
1378 | )}`,
1379 | },
1380 | ],
1381 | };
1382 | },
1383 | });
1384 | }
1385 |
1386 | const SUNO_MODEL_CONFIG: Record<string, BaseConfig> = {
1387 | music: { maxAttempts: 30, timeout: 900 },
1388 | };
1389 |
1390 | function registerSunoTool(server: FastMCP) {
1391 | server.addTool({
1392 | name: "generate_music_suno",
1393 | description: "Generate music using Suno",
1394 | parameters: z.object({
1395 | prompt: z
1396 | .string()
1397 | .max(3000)
1398 | .describe(
1399 | "The prompt to generate a music from, limited to 3000 characters"
1400 | ),
1401 | makeInstrumental: z
1402 | .boolean()
1403 | .optional()
1404 | .default(false)
1405 | .describe(
1406 | "Whether to make the music instrumental, defaults to false. Not compatible with title, tags, negativeTags"
1407 | ),
1408 | title: z
1409 | .string()
1410 | .max(80)
1411 | .optional()
1412 | .describe("The title of the music, limited to 80 characters"),
1413 | tags: z
1414 | .string()
1415 | .max(200)
1416 | .optional()
1417 | .describe("The tags of the music, limited to 200 characters"),
1418 | negativeTags: z
1419 | .string()
1420 | .max(200)
1421 | .optional()
1422 | .describe("The negative tags of the music, limited to 200 characters"),
1423 | }),
1424 | execute: async (args, { log, reportProgress }) => {
1425 | // Create video generation task
1426 | if (!args.prompt) {
1427 | throw new UserError("Prompt is required");
1428 | }
1429 | const config = SUNO_MODEL_CONFIG["music"];
1430 |
1431 | let requestBody = "";
1432 | if (args.title || args.tags || args.negativeTags) {
1433 | if (args.makeInstrumental) {
1434 | throw new UserError(
1435 | "makeInstrumental is not compatible with title, tags, negativeTags, please remove them if you want to make the music instrumental"
1436 | );
1437 | }
1438 | requestBody = JSON.stringify({
1439 | model: "music-s",
1440 | task_type: "generate_music_custom",
1441 | input: {
1442 | prompt: args.prompt,
1443 | title: args.title,
1444 | tags: args.tags,
1445 | negative_tags: args.negativeTags,
1446 | },
1447 | });
1448 | } else {
1449 | requestBody = JSON.stringify({
1450 | model: "music-s",
1451 | task_type: "generate_music",
1452 | input: {
1453 | prompt: args.prompt,
1454 | make_instrumental: args.makeInstrumental,
1455 | },
1456 | });
1457 | }
1458 |
1459 | const { taskId, usage, output } = await handleTask(
1460 | log,
1461 | reportProgress,
1462 | requestBody,
1463 | config
1464 | );
1465 |
1466 | const clips = parseSunoMusicOutput(taskId, output, log);
1467 | let content: Content[] = [];
1468 | content.push({
1469 | type: "text",
1470 | text: `TaskId: ${taskId}\nMusic generated successfully!\nUsage: ${usage} tokens`,
1471 | });
1472 | for (const clip of clips) {
1473 | content.push({
1474 | type: "text",
1475 | text: `Audio url: ${clip.audio_url}\nImage url: ${clip.image_url}`,
1476 | });
1477 | }
1478 | return {
1479 | content,
1480 | };
1481 | },
1482 | });
1483 | }
1484 |
1485 | const LUMA_MODEL_CONFIG: Record<string, BaseConfig> = {
1486 | luma: { maxAttempts: 30, timeout: 900 },
1487 | };
1488 |
1489 | function registerLumaTool(server: FastMCP) {
1490 | server.addTool({
1491 | name: "generate_video_luma",
1492 | description: "Generate a video using Luma",
1493 | parameters: z.object({
1494 | prompt: z.string().describe("The prompt to generate a video from"),
1495 | duration: z
1496 | .enum(["5s", "10s"])
1497 | .optional()
1498 | .default("5s")
1499 | .describe(
1500 | "The duration of the video, defaults to 5s. If keyFrame is provided, only 5s is supported"
1501 | ),
1502 | aspectRatio: z
1503 | .string()
1504 | .optional()
1505 | .describe("The aspect ratio of the video, defaults to 16:9"),
1506 | keyFrame: z
1507 | .string()
1508 | .url()
1509 | .optional()
1510 | .describe("The key frame to generate a video with"),
1511 | }),
1512 | execute: async (args, { log, reportProgress }) => {
1513 | // Create video generation task
1514 | if (!args.prompt) {
1515 | throw new UserError("Prompt is required");
1516 | }
1517 | const config = LUMA_MODEL_CONFIG["luma"];
1518 |
1519 | const requestBody = JSON.stringify({
1520 | model: "luma",
1521 | task_type: "video_generation",
1522 | input: {
1523 | prompt: args.prompt,
1524 | duration: args.duration === "5s" ? 5 : 10,
1525 | aspect_ratio: args.aspectRatio,
1526 | key_frames: {
1527 | frame0: {
1528 | type: args.keyFrame ? "image" : "",
1529 | url: args.keyFrame,
1530 | },
1531 | },
1532 | },
1533 | });
1534 | const { taskId, usage, output } = await handleTask(
1535 | log,
1536 | reportProgress,
1537 | requestBody,
1538 | config
1539 | );
1540 |
1541 | const [video_raw, last_frame] = parseLumaOutput(taskId, output, log);
1542 | return {
1543 | content: [
1544 | {
1545 | type: "text",
1546 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${video_raw.url}\nVideo resolution: ${video_raw.width}x${video_raw.height}\nLast frame url:\n${last_frame.url}\nLast frame resolution: ${last_frame.width}x${last_frame.height}`,
1547 | },
1548 | ],
1549 | };
1550 | },
1551 | });
1552 | }
1553 |
1554 | const TRELLIS_MODEL_CONFIG: Record<string, BaseConfig> = {
1555 | trellis: { maxAttempts: 30, timeout: 600 },
1556 | };
1557 |
1558 | const HAILUO_MODEL_CONFIG: Record<string, BaseConfig> = {
1559 | hailuo: { maxAttempts: 60, timeout: 900 },
1560 | };
1561 |
1562 | function registerTrellisTool(server: FastMCP) {
1563 | server.addTool({
1564 | name: "generate_3d_model",
1565 | description: "Generate a 3d model using Qubico Trellis",
1566 | parameters: z.object({
1567 | image: z.string().url().describe("The image to generate a 3d model from"),
1568 | }),
1569 | execute: async (args, { log, reportProgress }) => {
1570 | // Create 3d model generation task
1571 | if (!args.image) {
1572 | throw new UserError("Image is required");
1573 | }
1574 | const config = TRELLIS_MODEL_CONFIG["trellis"];
1575 |
1576 | const requestBody = JSON.stringify({
1577 | model: "Qubico/trellis",
1578 | task_type: "image-to-3d",
1579 | input: {
1580 | image: args.image,
1581 | },
1582 | });
1583 | const { taskId, usage, output } = await handleTask(
1584 | log,
1585 | reportProgress,
1586 | requestBody,
1587 | config
1588 | );
1589 |
1590 | const [imageUrl, videoUrl, modelFileUrl] = parseTrellisOutput(
1591 | taskId,
1592 | output,
1593 | log
1594 | );
1595 | return {
1596 | content: [
1597 | {
1598 | type: "text",
1599 | text: `TaskId: ${taskId}\n3d model generated successfully!\nUsage: ${usage} tokens\nImage url:\n${imageUrl}\nVideo url:\n${videoUrl}\nModel file url:\n${modelFileUrl}`,
1600 | },
1601 | ],
1602 | };
1603 | },
1604 | });
1605 | }
1606 |
1607 | function registerHailuoTool(server: FastMCP) {
1608 | server.addTool({
1609 | name: "generate_video_hailuo",
1610 | description: "Generate a video using Hailuo",
1611 | parameters: z.object({
1612 | prompt: z
1613 | .string()
1614 | .max(2000)
1615 | .describe("The prompt to generate a video from (max 2000 characters)"),
1616 | imageUrl: z
1617 | .string()
1618 | .url()
1619 | .optional()
1620 | .describe("The image URL for image-to-video models"),
1621 | expandPrompt: z
1622 | .boolean()
1623 | .optional()
1624 | .default(false)
1625 | .describe("Whether to expand the prompt"),
1626 | model: z
1627 | .enum(["t2v-01", "t2v-01-director", "i2v-01", "i2v-01-live", "i2v-01-director", "s2v-01"])
1628 | .optional()
1629 | .default("t2v-01")
1630 | .describe("The model to use for video generation. t2v models are text-to-video, i2v models are image-to-video, s2v-01 requires human face detection"),
1631 | }),
1632 | execute: async (args, { log, reportProgress }) => {
1633 | if (!args.prompt) {
1634 | throw new UserError("Prompt is required");
1635 | }
1636 |
1637 | // Validate model requirements
1638 | const isImageToVideo = args.model.startsWith("i2v") || args.model === "s2v-01";
1639 | if (isImageToVideo && !args.imageUrl) {
1640 | throw new UserError(`Image URL is required for ${args.model} model`);
1641 | }
1642 | if (!isImageToVideo && args.imageUrl) {
1643 | log.warn(`Image URL provided but ${args.model} is a text-to-video model`);
1644 | }
1645 |
1646 | const config = HAILUO_MODEL_CONFIG["hailuo"];
1647 |
1648 | const requestBody = JSON.stringify({
1649 | model: args.model,
1650 | task_type: "video_generation",
1651 | input: {
1652 | prompt: args.prompt,
1653 | image_url: args.imageUrl,
1654 | expand_prompt: args.expandPrompt,
1655 | },
1656 | });
1657 |
1658 | const { taskId, usage, output } = await handleTask(
1659 | log,
1660 | reportProgress,
1661 | requestBody,
1662 | config
1663 | );
1664 |
1665 | const url = parseVideoOutput(taskId, output);
1666 | return {
1667 | content: [
1668 | {
1669 | type: "text",
1670 | text: `TaskId: ${taskId}\nVideo generated successfully!\nUsage: ${usage} tokens\nVideo url:\n${url}`,
1671 | },
1672 | ],
1673 | };
1674 | },
1675 | });
1676 | }
1677 |
1678 | // Task handler
1679 | async function handleTask(
1680 | log: any,
1681 | reportProgress: (progress: Progress) => Promise<void>,
1682 | requestBody: string,
1683 | config: BaseConfig
1684 | ): Promise<{ taskId: string; usage: string; output: unknown }> {
1685 | const taskId = await createTask(requestBody);
1686 | log.info(`Task created with ID: ${taskId}`);
1687 | return await getTaskResult(
1688 | log,
1689 | reportProgress,
1690 | taskId,
1691 | config.maxAttempts,
1692 | config.timeout
1693 | );
1694 | }
1695 |
1696 | async function createTask(requestBody: string) {
1697 | const createResponse = await fetch("https://api.piapi.ai/api/v1/task", {
1698 | method: "POST",
1699 | headers: {
1700 | "X-API-Key": apiKey,
1701 | Accept: "application/json",
1702 | "Content-Type": "application/json",
1703 | },
1704 | body: requestBody,
1705 | });
1706 |
1707 | const createData = await createResponse.json();
1708 |
1709 | if (createData.code !== 200) {
1710 | throw new UserError(`Task creation failed: ${createData.message}`);
1711 | }
1712 |
1713 | return createData.data.task_id;
1714 | }
1715 |
1716 | async function getTaskResult(
1717 | log: any,
1718 | reportProgress: (progress: Progress) => Promise<void>,
1719 | taskId: string,
1720 | maxAttempts: number,
1721 | timeout: number
1722 | ): Promise<{ taskId: string; usage: string; output: unknown }> {
1723 | for (let attempt = 0; attempt < maxAttempts; attempt++) {
1724 | // Use environment-specific logger, fallback to provided log if exists
1725 | const useLogger = log || logger;
1726 | useLogger.info(`Checking task ${taskId} status (attempt ${attempt + 1}/${maxAttempts})...`);
1727 |
1728 | reportProgress({
1729 | progress: (attempt / maxAttempts) * 100,
1730 | total: 100,
1731 | });
1732 |
1733 | const statusResponse = await fetch(
1734 | `https://api.piapi.ai/api/v1/task/${taskId}`,
1735 | {
1736 | headers: {
1737 | "X-API-Key": apiKey,
1738 | },
1739 | }
1740 | );
1741 |
1742 | const statusData = await statusResponse.json();
1743 |
1744 | if (statusData.code !== 200) {
1745 | useLogger.error(`Status check failed for task ${taskId}: ${statusData.message}`);
1746 | throw new UserError(
1747 | `TaskId: ${taskId}, Status check failed: ${statusData.message}`
1748 | );
1749 | }
1750 |
1751 | const { status, output, error } = statusData.data;
1752 |
1753 | useLogger.info(`Task ${taskId} status: ${status}`);
1754 |
1755 | // Safely check if progress property exists
1756 | if (status === "in_progress" && statusData.data.progress !== undefined) {
1757 | useLogger.info(`Task ${taskId} progress: ${statusData.data.progress}%`);
1758 | }
1759 |
1760 | if (status === "completed") {
1761 | if (!output) {
1762 | useLogger.error(`Task ${taskId} completed but no output found`);
1763 | throw new UserError(
1764 | `TaskId: ${taskId}, Task completed but no output found`
1765 | );
1766 | }
1767 | const usage = statusData.data.meta?.usage?.consume || "unknown";
1768 | useLogger.info(`Task ${taskId} completed successfully. Usage: ${usage}`);
1769 |
1770 | // Don't log huge JSON objects that might crash the console
1771 | try {
1772 | const outputStr = JSON.stringify(output);
1773 | if (outputStr.length < 1000) {
1774 | useLogger.debug(`Task ${taskId} output: ${outputStr}`);
1775 | } else {
1776 | useLogger.debug(`Task ${taskId} output: [Large output, length: ${outputStr.length} chars]`);
1777 | }
1778 | } catch (err: any) {
1779 | useLogger.debug(`Task ${taskId} output: [Could not stringify output: ${err.message}]`);
1780 | }
1781 |
1782 | return { taskId, usage, output };
1783 | }
1784 |
1785 | if (status === "failed") {
1786 | useLogger.error(`Task ${taskId} failed: ${error?.message || "Unknown error"}`);
1787 | throw new UserError(
1788 | `TaskId: ${taskId}, Generation failed: ${error?.message || "Unknown error"}`
1789 | );
1790 | }
1791 |
1792 | await new Promise((resolve) =>
1793 | setTimeout(resolve, (timeout * 1000) / maxAttempts)
1794 | );
1795 | }
1796 |
1797 | logger.error(`Task ${taskId} timed out after ${timeout} seconds`);
1798 | throw new UserError(
1799 | `TaskId: ${taskId}, Generation timed out after ${timeout} seconds`
1800 | );
1801 | }
1802 |
1803 | // Result parser
1804 |
1805 | const ImageOutputSchema = z
1806 | .object({
1807 | image_url: z.string().optional(),
1808 | image_urls: z.array(z.string()).nullable().optional(),
1809 | temporary_image_urls: z.array(z.string()).nullable().optional(),
1810 | })
1811 | .refine(
1812 | (data) =>
1813 | data.image_url ||
1814 | (data.image_urls && data.image_urls.length > 0) ||
1815 | (data.temporary_image_urls && data.temporary_image_urls.length > 0),
1816 | {
1817 | message: "At least one image URL must be provided",
1818 | path: ["image_url", "image_urls", "temporary_image_urls"],
1819 | }
1820 | );
1821 |
1822 | function parseImageOutput(taskId: string, output: unknown, log?: any): string[] {
1823 | const useLogger = log || logger;
1824 |
1825 | useLogger.info(`Parsing image output for task ${taskId}`);
1826 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
1827 |
1828 | const result = ImageOutputSchema.safeParse(output);
1829 |
1830 | if (!result.success) {
1831 | useLogger.error(`Invalid image output format for task ${taskId}: ${result.error.message}`);
1832 | throw new UserError(
1833 | `TaskId: ${taskId}, Invalid image output format: ${result.error.message}`
1834 | );
1835 | }
1836 |
1837 | const imageOutput = result.data;
1838 | useLogger.debug(`Image URLs found - image_url: ${imageOutput.image_url || 'none'}, image_urls count: ${imageOutput.image_urls?.length || 0}, temporary_image_urls count: ${imageOutput.temporary_image_urls?.length || 0}`);
1839 |
1840 | // Determine if this is a Midjourney response (has temporary_image_urls but null image_urls)
1841 | const isMidjourney = Array.isArray(imageOutput.temporary_image_urls) &&
1842 | imageOutput.temporary_image_urls.length > 0 &&
1843 | imageOutput.image_urls === null;
1844 |
1845 | const imageUrls = [
1846 | ...(imageOutput.image_url ? [imageOutput.image_url] : []),
1847 | ...(!isMidjourney && imageOutput.image_urls ? imageOutput.image_urls : []),
1848 | ...(imageOutput.temporary_image_urls || []),
1849 | ].filter(Boolean);
1850 |
1851 | if (imageUrls.length === 0) {
1852 | useLogger.error(`No image URLs found for task ${taskId}`);
1853 | throw new UserError(
1854 | `TaskId: ${taskId}, Task completed but no image URLs found`
1855 | );
1856 | }
1857 |
1858 | useLogger.info(`Found ${imageUrls.length} image URLs for task ${taskId}`);
1859 | return imageUrls;
1860 | }
1861 |
1862 | const AudioOutputSchema = z
1863 | .object({
1864 | audio_url: z.string(),
1865 | })
1866 | .refine((data) => data.audio_url, {
1867 | message: "At least one audio URL must be provided",
1868 | path: ["audio_url"],
1869 | });
1870 |
1871 | function parseAudioOutput(taskId: string, output: unknown, log?: any): string {
1872 | const useLogger = log || logger;
1873 |
1874 | useLogger.info(`Parsing audio output for task ${taskId}`);
1875 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
1876 |
1877 | const result = AudioOutputSchema.safeParse(output);
1878 |
1879 | if (!result.success) {
1880 | useLogger.error(`Invalid audio output format for task ${taskId}: ${result.error.message}`);
1881 | throw new UserError(
1882 | `TaskId: ${taskId}, Invalid audio output format: ${result.error.message}`
1883 | );
1884 | }
1885 |
1886 | const audioUrl = result.data.audio_url;
1887 |
1888 | if (!audioUrl) {
1889 | useLogger.error(`Task ${taskId} completed but no audio URL found`);
1890 | throw new UserError(
1891 | `TaskId: ${taskId}, Task completed but no audio URL found`
1892 | );
1893 | }
1894 |
1895 | useLogger.info(`Found audio URL for task ${taskId}: ${audioUrl}`);
1896 | return audioUrl;
1897 | }
1898 |
1899 | const VideoOutputSchema = z
1900 | .object({
1901 | video_url: z.string(),
1902 | })
1903 | .refine((data) => data.video_url, {
1904 | message: "At least one video URL must be provided",
1905 | path: ["video_url"],
1906 | });
1907 |
1908 | function parseVideoOutput(taskId: string, output: unknown, log?: any): string {
1909 | const useLogger = log || logger;
1910 |
1911 | useLogger.info(`Parsing video output for task ${taskId}`);
1912 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
1913 |
1914 | const result = VideoOutputSchema.safeParse(output);
1915 |
1916 | if (!result.success) {
1917 | useLogger.error(`Invalid video output format for task ${taskId}: ${result.error.message}`);
1918 | throw new UserError(
1919 | `TaskId: ${taskId}, Invalid video output format: ${result.error.message}`
1920 | );
1921 | }
1922 |
1923 | const videoUrl = result.data.video_url;
1924 |
1925 | if (!videoUrl) {
1926 | useLogger.error(`Task ${taskId} completed but no video URL found`);
1927 | throw new UserError(
1928 | `TaskId: ${taskId}, Task completed but no video URL found`
1929 | );
1930 | }
1931 |
1932 | useLogger.info(`Found video URL for task ${taskId}: ${videoUrl}`);
1933 | return videoUrl;
1934 | }
1935 |
1936 | const KlingOutputSchema = z.object({
1937 | video_url: z.string(),
1938 | works: z.array(z.object({
1939 | video: z.object({
1940 | resource_without_watermark: z.string(),
1941 | // height: z.number(),
1942 | // width: z.number(),
1943 | // duration: z.number(),
1944 | })
1945 | }))
1946 | })
1947 |
1948 | function parseKlingOutput(taskId: string, output: unknown, log?: any): string[] {
1949 | const useLogger = log || logger;
1950 |
1951 | useLogger.info(`Parsing Kling output for task ${taskId}`);
1952 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
1953 |
1954 | const result = KlingOutputSchema.safeParse(output);
1955 |
1956 | if (!result.success) {
1957 | useLogger.error(`Invalid kling output format for task ${taskId}: ${result.error.message}`);
1958 | throw new UserError(
1959 | `TaskId: ${taskId}, Invalid kling output format: ${result.error.message}`
1960 | );
1961 | }
1962 |
1963 | let urls: string[] = [];
1964 | urls.push(result.data.video_url);
1965 | for (const work of result.data.works) {
1966 | urls.push(work.video.resource_without_watermark);
1967 | }
1968 |
1969 | if (urls.length === 0) {
1970 | useLogger.error(`Task ${taskId} completed but no video/work URLs found`);
1971 | throw new UserError(
1972 | `TaskId: ${taskId}, Task completed but no video/work URLs found`
1973 | );
1974 | }
1975 |
1976 | useLogger.info(`Found ${urls.length} Kling URLs for task ${taskId}`);
1977 | return urls;
1978 | }
1979 |
1980 | interface LumaResult {
1981 | url: string;
1982 | width: number;
1983 | height: number;
1984 | }
1985 |
1986 | const LumaOutputSchema = z
1987 | .object({
1988 | video_raw: z.object({
1989 | url: z.string(),
1990 | width: z.number(),
1991 | height: z.number(),
1992 | }),
1993 | last_frame: z.object({
1994 | url: z.string(),
1995 | width: z.number(),
1996 | height: z.number(),
1997 | }),
1998 | })
1999 | .refine((data) => data.video_raw && data.last_frame, {
2000 | message: "At least one video URL must be provided",
2001 | path: ["video_raw", "last_frame"],
2002 | });
2003 |
2004 | function parseLumaOutput(
2005 | taskId: string,
2006 | output: unknown,
2007 | log?: any
2008 | ): [LumaResult, LumaResult] {
2009 | const useLogger = log || logger;
2010 |
2011 | useLogger.info(`Parsing Luma output for task ${taskId}`);
2012 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
2013 |
2014 | const result = LumaOutputSchema.safeParse(output);
2015 |
2016 | if (!result.success) {
2017 | useLogger.error(`Invalid luma output format for task ${taskId}: ${result.error.message}`);
2018 | throw new UserError(
2019 | `TaskId: ${taskId}, Invalid luma output format: ${result.error.message}`
2020 | );
2021 | }
2022 |
2023 | useLogger.info(`Found Luma video and last frame for task ${taskId}`);
2024 | return [result.data.video_raw, result.data.last_frame];
2025 | }
2026 |
2027 | interface SunoMusicClip {
2028 | audio_url: string;
2029 | image_url: string;
2030 | }
2031 |
2032 | const SunoMusicOutputSchema = z.object({
2033 | clips: z.map(
2034 | z.string(),
2035 | z.object({
2036 | audio_url: z.string(),
2037 | image_url: z.string(),
2038 | })
2039 | ),
2040 | });
2041 |
2042 | function parseSunoMusicOutput(
2043 | taskId: string,
2044 | output: unknown,
2045 | log?: any
2046 | ): SunoMusicClip[] {
2047 | const useLogger = log || logger;
2048 |
2049 | useLogger.info(`Parsing Suno music output for task ${taskId}`);
2050 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
2051 |
2052 | const result = SunoMusicOutputSchema.safeParse(output);
2053 |
2054 | if (!result.success) {
2055 | useLogger.error(`Invalid suno music output format for task ${taskId}: ${result.error.message}`);
2056 | throw new UserError(
2057 | `TaskId: ${taskId}, Invalid suno music output format: ${result.error.message}`
2058 | );
2059 | }
2060 |
2061 | const results: SunoMusicClip[] = [];
2062 | for (const [key, value] of Object.entries(result.data.clips)) {
2063 | results.push({
2064 | audio_url: value.audio_url,
2065 | image_url: value.image_url,
2066 | });
2067 | }
2068 |
2069 | if (results.length === 0) {
2070 | useLogger.error(`Task ${taskId} completed but no audio/image URLs found`);
2071 | throw new UserError(
2072 | `TaskId: ${taskId}, Task completed but no audio/image URLs found`
2073 | );
2074 | }
2075 |
2076 | useLogger.info(`Found ${results.length} Suno music clips for task ${taskId}`);
2077 | return results;
2078 | }
2079 |
2080 | const TrellisOutputSchema = z
2081 | .object({
2082 | no_background_image: z.string(),
2083 | combined_video: z.string(),
2084 | model_file: z.string(),
2085 | })
2086 | .refine(
2087 | (data) =>
2088 | data.no_background_image && data.combined_video && data.model_file,
2089 | {
2090 | message: "At least one image/video/model file URL must be provided",
2091 | path: ["no_background_image", "combined_video", "model_file"],
2092 | }
2093 | );
2094 |
2095 | function parseTrellisOutput(
2096 | taskId: string,
2097 | output: unknown,
2098 | log?: any
2099 | ): [string, string, string] {
2100 | const useLogger = log || logger;
2101 |
2102 | useLogger.info(`Parsing Trellis output for task ${taskId}`);
2103 | useLogger.debug(`Raw output: ${JSON.stringify(output)}`);
2104 |
2105 | const result = TrellisOutputSchema.safeParse(output);
2106 |
2107 | if (!result.success) {
2108 | useLogger.error(`Invalid trellis output format for task ${taskId}: ${result.error.message}`);
2109 | throw new UserError(
2110 | `TaskId: ${taskId}, Invalid trellis output format: ${result.error.message}`
2111 | );
2112 | }
2113 |
2114 | const imageUrl = result.data.no_background_image;
2115 | const videoUrl = result.data.combined_video;
2116 | const modelFileUrl = result.data.model_file;
2117 |
2118 | if (!imageUrl || !videoUrl || !modelFileUrl) {
2119 | useLogger.error(`Task ${taskId} completed but no image/video/model file URL found`);
2120 | throw new UserError(
2121 | `TaskId: ${taskId}, Task completed but no image/video/model file URL found`
2122 | );
2123 | }
2124 |
2125 | useLogger.info(`Found Trellis outputs for task ${taskId}`);
2126 | return [imageUrl, videoUrl, modelFileUrl];
2127 | }
2128 |
```