diff --git a/server.mjs b/server.mjs index 9a4accf..b5835a7 100644 --- a/server.mjs +++ b/server.mjs @@ -431,6 +431,30 @@ function registerHandlers(srv) { required: ['prompt'], }, }, + { + name: 'upload_image', + description: + 'Upload a reference image to the server for use in image_urls. ' + + 'Accepts a base64 data URL (e.g. "data:image/png;base64,iVBOR...") or ' + + 'an HTTP(S) URL. Returns a download_url that can be passed to generate_image ' + + 'or generate_image_transparent as image_urls entry. ' + + 'The uploaded file persists only while the container is alive.', + inputSchema: { + type: 'object', + properties: { + data: { + type: 'string', + description: 'Base64-encoded image data (raw base64 string, not including data: URI prefix) OR a full data:image URL OR an http(s):// URL to fetch.', + }, + filename: { + type: 'string', + description: 'Optional filename hint (e.g. "ref.png"). Auto-generated if omitted.', + default: '', + }, + }, + required: ['data'], + }, + }, { name: 'get_image', description: @@ -494,9 +518,9 @@ function registerHandlers(srv) { download_url: `https://img-gen-ai.dynv6.net/images/${filename}`, }, null, 2) }, ], - }; - } - if (name === 'get_image') { + } + } + if (name === 'get_image') { const filename = String(args?.filename || '').trim(); if (!filename) return errResp('filename required'); if (!/^[\w.-]+$/.test(filename)) return errResp('filename must be alphanumeric with dots/hyphens only'); @@ -513,14 +537,45 @@ function registerHandlers(srv) { { type: 'text', text: JSON.stringify({ filename, filepath, bytes: buf.length, mime, download_url: `https://img-gen-ai.dynv6.net/images/${filename}` }, null, 2) }, ], }; - } - return errResp(`Unknown tool: ${name}`); - } catch (e) { - process.stderr.write(`[gen-ERR] ${e.stack || e}\n`); - return { content: [{ type: 'text', text: `Error: ${e.message}` }], isError: true }; - } - }); // end CallToolRequestSchema -} // end registerHandlers + } + if (name === 'upload_image') { + const rawData = String(args?.data || '').trim(); + if (!rawData) return errResp('data is required'); + let buf, ext = 'png'; + if (/^https?:\/\//i.test(rawData)) { + const fetched = await fetch(rawData, { signal: AbortSignal.timeout(60_000) }); + if (!fetched.ok) return errResp(`URL fetch ${fetched.status}`); + const ct = fetched.headers.get('content-type') || ''; + if (ct.includes('png')) ext = 'png'; + else if (ct.includes('jpeg') || ct.includes('jpg')) ext = 'jpg'; + else if (ct.includes('webp')) ext = 'webp'; + buf = Buffer.from(await fetched.arrayBuffer()); + } else { + // data: URL or raw base64 + const base64 = rawData.replace(/^data:image\/\w+;base64,/, ''); + try { buf = Buffer.from(base64, 'base64'); } catch (e) { return errResp('Invalid base64 data'); } + } + await fs.mkdir(OUTPUT_DIR, { recursive: true }); + const filename = args?.filename || `upload-${Date.now()}.${ext}`; + if (!/^[\w.-]+$/.test(filename)) return errResp('Invalid filename'); + const filepath = path.join(OUTPUT_DIR, filename); + await fs.writeFile(filepath, buf); + const downloadUrl = `https://img-gen-ai.dynv6.net/images/${filename}`; + process.stderr.write(`[upload] ${filename} ${buf.length}b -> ${downloadUrl}\n`); + return { + content: [ + { type: 'image', mimeType: `image/${ext}`, data: buf.toString('base64') }, + { type: 'text', text: JSON.stringify({ filename, filepath, bytes: buf.length, download_url: downloadUrl }, null, 2) }, + ], + }; + } + return errResp(`Unknown tool: ${name}`); + } catch (e) { + process.stderr.write(`[gen-ERR] ${e.stack || e}\n`); + return { content: [{ type: 'text', text: `Error: ${e.message}` }], isError: true }; + } + }); // end CallToolRequestSchema + } // end registerHandlers // ---------- Create MCP server instance ---------- const mcpServer = new Server(