{"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":[]}