@genkit-ai/googleai
Version:
Genkit AI framework plugin for Google AI APIs, including Gemini APIs.
1 lines • 7.49 kB
Source Map (JSON)
{"version":3,"sources":["../src/imagen.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport { GenkitError, MessageData, z, type Genkit } from 'genkit';\nimport {\n getBasicUsageStats,\n modelRef,\n type GenerateRequest,\n type ModelAction,\n type ModelInfo,\n type ModelReference,\n} from 'genkit/model';\nimport { getApiKeyFromEnvVar } from './common.js';\nimport { predictModel } from './predict.js';\n\n/**\n * @deprecated\n */\nexport type KNOWN_IMAGEN_MODELS = 'imagen-3.0-generate-002';\n\n/**\n * See https://ai.google.dev/gemini-api/docs/image-generation#imagen-model\n * @deprecated\n */\nexport const ImagenConfigSchema = z\n .object({\n numberOfImages: z\n .number()\n .describe(\n 'The number of images to generate, from 1 to 4 (inclusive). The default is 1.'\n )\n .optional(),\n aspectRatio: z\n .enum(['1:1', '9:16', '16:9', '3:4', '4:3'])\n .describe('Desired aspect ratio of the output image.')\n .optional(),\n personGeneration: z\n .enum(['dont_allow', 'allow_adult', 'allow_all'])\n .describe(\n 'Control if/how images of people will be generated by the model.'\n )\n .optional(),\n })\n .passthrough();\n\ninterface ImagenParameters {\n sampleCount?: number;\n aspectRatio?: string;\n personGeneration?: string;\n}\n\nfunction toParameters(\n request: GenerateRequest<typeof ImagenConfigSchema>\n): ImagenParameters {\n const out = {\n sampleCount: request.config?.numberOfImages ?? 1,\n ...request?.config,\n };\n\n for (const k in out) {\n if (!out[k]) delete out[k];\n }\n\n return out;\n}\n\nfunction extractText(request: GenerateRequest) {\n return request.messages\n .at(-1)!\n .content.map((c) => c.text || '')\n .join('');\n}\n\nfunction extractBaseImage(request: GenerateRequest): string | undefined {\n return request.messages\n .at(-1)\n ?.content.find((p) => !!p.media)\n ?.media?.url.split(',')[1];\n}\n\ninterface ImagenPrediction {\n predictions: { bytesBase64Encoded: string; mimeType: string }[];\n}\n\ninterface ImagenInstance {\n prompt: string;\n image?: { bytesBase64Encoded: string };\n mask?: { image?: { bytesBase64Encoded: string } };\n}\n\n/**\n * @deprecated\n */\nexport const GENERIC_IMAGEN_INFO = {\n label: `Google AI - Generic Imagen`,\n supports: {\n media: true,\n multiturn: false,\n tools: false,\n systemRole: false,\n output: ['media'],\n },\n} as ModelInfo;\n\n/**\n * @deprecated\n */\nexport function defineImagenModel(\n ai: Genkit,\n name: string,\n apiKey?: string | false\n): ModelAction {\n if (apiKey !== false) {\n apiKey = apiKey || getApiKeyFromEnvVar();\n if (!apiKey) {\n throw new GenkitError({\n status: 'FAILED_PRECONDITION',\n message:\n 'Please pass in the API key or set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable.\\n' +\n 'For more details see https://genkit.dev/docs/plugins/google-genai',\n });\n }\n }\n const modelName = `googleai/${name}`;\n const model: ModelReference<z.ZodTypeAny> = modelRef({\n name: modelName,\n info: {\n ...GENERIC_IMAGEN_INFO,\n label: `Google AI - ${name}`,\n },\n configSchema: ImagenConfigSchema,\n });\n\n return ai.defineModel(\n {\n name: modelName,\n ...model.info,\n configSchema: ImagenConfigSchema,\n },\n async (request) => {\n const instance: ImagenInstance = {\n prompt: extractText(request),\n };\n const baseImage = extractBaseImage(request);\n if (baseImage) {\n instance.image = { bytesBase64Encoded: baseImage };\n }\n\n const predictClient = predictModel<\n ImagenInstance,\n ImagenPrediction,\n ImagenParameters\n >(model.version || name, apiKey as string, 'predict');\n const response = await predictClient([instance], toParameters(request));\n\n if (!response.predictions || response.predictions.length == 0) {\n throw new Error(\n 'Model returned no predictions. Possibly due to content filters.'\n );\n }\n\n const message = {\n role: 'model',\n content: [],\n } as MessageData;\n\n response.predictions.forEach((p, i) => {\n const b64data = p.bytesBase64Encoded;\n const mimeType = p.mimeType;\n message.content.push({\n media: {\n url: `data:${mimeType};base64,${b64data}`,\n contentType: mimeType,\n },\n });\n });\n return {\n finishReason: 'stop',\n message,\n usage: getBasicUsageStats(request.messages, message),\n custom: response,\n };\n }\n );\n}\n"],"mappings":"AAgBA,SAAS,aAA0B,SAAsB;AACzD;AAAA,EACE;AAAA,EACA;AAAA,OAKK;AACP,SAAS,2BAA2B;AACpC,SAAS,oBAAoB;AAWtB,MAAM,qBAAqB,EAC/B,OAAO;AAAA,EACN,gBAAgB,EACb,OAAO,EACP;AAAA,IACC;AAAA,EACF,EACC,SAAS;AAAA,EACZ,aAAa,EACV,KAAK,CAAC,OAAO,QAAQ,QAAQ,OAAO,KAAK,CAAC,EAC1C,SAAS,2CAA2C,EACpD,SAAS;AAAA,EACZ,kBAAkB,EACf,KAAK,CAAC,cAAc,eAAe,WAAW,CAAC,EAC/C;AAAA,IACC;AAAA,EACF,EACC,SAAS;AACd,CAAC,EACA,YAAY;AAQf,SAAS,aACP,SACkB;AAClB,QAAM,MAAM;AAAA,IACV,aAAa,QAAQ,QAAQ,kBAAkB;AAAA,IAC/C,GAAG,SAAS;AAAA,EACd;AAEA,aAAW,KAAK,KAAK;AACnB,QAAI,CAAC,IAAI,CAAC,EAAG,QAAO,IAAI,CAAC;AAAA,EAC3B;AAEA,SAAO;AACT;AAEA,SAAS,YAAY,SAA0B;AAC7C,SAAO,QAAQ,SACZ,GAAG,EAAE,EACL,QAAQ,IAAI,CAAC,MAAM,EAAE,QAAQ,EAAE,EAC/B,KAAK,EAAE;AACZ;AAEA,SAAS,iBAAiB,SAA8C;AACtE,SAAO,QAAQ,SACZ,GAAG,EAAE,GACJ,QAAQ,KAAK,CAAC,MAAM,CAAC,CAAC,EAAE,KAAK,GAC7B,OAAO,IAAI,MAAM,GAAG,EAAE,CAAC;AAC7B;AAeO,MAAM,sBAAsB;AAAA,EACjC,OAAO;AAAA,EACP,UAAU;AAAA,IACR,OAAO;AAAA,IACP,WAAW;AAAA,IACX,OAAO;AAAA,IACP,YAAY;AAAA,IACZ,QAAQ,CAAC,OAAO;AAAA,EAClB;AACF;AAKO,SAAS,kBACd,IACA,MACA,QACa;AACb,MAAI,WAAW,OAAO;AACpB,aAAS,UAAU,oBAAoB;AACvC,QAAI,CAAC,QAAQ;AACX,YAAM,IAAI,YAAY;AAAA,QACpB,QAAQ;AAAA,QACR,SACE;AAAA,MAEJ,CAAC;AAAA,IACH;AAAA,EACF;AACA,QAAM,YAAY,YAAY,IAAI;AAClC,QAAM,QAAsC,SAAS;AAAA,IACnD,MAAM;AAAA,IACN,MAAM;AAAA,MACJ,GAAG;AAAA,MACH,OAAO,eAAe,IAAI;AAAA,IAC5B;AAAA,IACA,cAAc;AAAA,EAChB,CAAC;AAED,SAAO,GAAG;AAAA,IACR;AAAA,MACE,MAAM;AAAA,MACN,GAAG,MAAM;AAAA,MACT,cAAc;AAAA,IAChB;AAAA,IACA,OAAO,YAAY;AACjB,YAAM,WAA2B;AAAA,QAC/B,QAAQ,YAAY,OAAO;AAAA,MAC7B;AACA,YAAM,YAAY,iBAAiB,OAAO;AAC1C,UAAI,WAAW;AACb,iBAAS,QAAQ,EAAE,oBAAoB,UAAU;AAAA,MACnD;AAEA,YAAM,gBAAgB,aAIpB,MAAM,WAAW,MAAM,QAAkB,SAAS;AACpD,YAAM,WAAW,MAAM,cAAc,CAAC,QAAQ,GAAG,aAAa,OAAO,CAAC;AAEtE,UAAI,CAAC,SAAS,eAAe,SAAS,YAAY,UAAU,GAAG;AAC7D,cAAM,IAAI;AAAA,UACR;AAAA,QACF;AAAA,MACF;AAEA,YAAM,UAAU;AAAA,QACd,MAAM;AAAA,QACN,SAAS,CAAC;AAAA,MACZ;AAEA,eAAS,YAAY,QAAQ,CAAC,GAAG,MAAM;AACrC,cAAM,UAAU,EAAE;AAClB,cAAM,WAAW,EAAE;AACnB,gBAAQ,QAAQ,KAAK;AAAA,UACnB,OAAO;AAAA,YACL,KAAK,QAAQ,QAAQ,WAAW,OAAO;AAAA,YACvC,aAAa;AAAA,UACf;AAAA,QACF,CAAC;AAAA,MACH,CAAC;AACD,aAAO;AAAA,QACL,cAAc;AAAA,QACd;AAAA,QACA,OAAO,mBAAmB,QAAQ,UAAU,OAAO;AAAA,QACnD,QAAQ;AAAA,MACV;AAAA,IACF;AAAA,EACF;AACF;","names":[]}