Fixes #332: Adds getImageModel method for image generation

src/gen-ai.ts

@@ -21,6 +21,7 @@ import {
 } from "./errors";
 import { CachedContent, ModelParams, RequestOptions } from "../types";
 import { GenerativeModel } from "./models/generative-model";
+import { ImageModel } from "./models/image-model";
 
 export { ChatSession } from "./methods/chat-session";
 export { GenerativeModel };
@@ -48,6 +49,20 @@ export class GoogleGenerativeAI {
     return new GenerativeModel(this.apiKey, modelParams, requestOptions);
   }
 
+  getImageModel(
+    modelParams: ModelParams,
+    requestOptions?: RequestOptions,
+  ): ImageModel {
+    if (!modelParams.model) {
+      throw new GoogleGenerativeAIError(
+        `Must provide a model name. ` +
+          `Example: genai.getGenerativeModel({ model: 'my-model-name' })`,
+      );
+    }
+    return new ImageModel(this.apiKey, modelParams, requestOptions);
+  }
+
+
   /**
    * Creates a {@link GenerativeModel} instance from provided content cache.
    */

src/methods/generate-image.ts

@@ -0,0 +1,20 @@
+import { GenerateImageRequest, GenerateImagesResult, SingleRequestOptions } from "../../types";
+import { Task, makeModelRequest } from "../requests/request";
+
+  export async function generateImages(
+    apiKey: string,
+    model: string,
+    params: GenerateImageRequest,
+    requestOptions: SingleRequestOptions,
+  ): Promise<GenerateImagesResult> {
+    const response = await makeModelRequest(
+      model,
+      Task.GENERATE_IMAGES,
+      apiKey,
+      /* stream */ false,
+      JSON.stringify(params),
+      requestOptions,
+    );
+    return response.json();
+  }
+

src/models/generative-model.ts

@@ -32,6 +32,7 @@ import {
   GenerateContentResult,
   GenerateContentStreamResult,
   GenerationConfig,
+  ImageModelParams,
   ModelParams,
   Part,
   RequestOptions,
@@ -85,7 +86,8 @@ export class GenerativeModel {
     );
     this.cachedContent = modelParams.cachedContent;
   }
-
+
+
   /**
    * Makes a single non-streaming call to the model
    * and returns an object containing a single {@link GenerateContentResponse}.

src/models/image-model.ts

@@ -0,0 +1,45 @@
+import { BaseImageParams, GenerateImagesResult, ImageModelParams, SingleRequestOptions } from "../../types";
+import { RequestOptions } from "../server";
+import {generateImages} from "../methods/generate-image"
+export class ImageModel{
+    model:string;
+    modelParams:ImageModelParams;
+
+    constructor(
+        public apiKey: string,
+        modelParams: ImageModelParams,
+        private _requestOptions: RequestOptions = {},
+      ) 
+      {
+        if (modelParams.model.includes("/")) {
+            // Models may be named "models/model-name" or "tunedModels/model-name"
+            this.model = modelParams.model;
+          } else {
+            // If path is not included, assume it's a non-tuned model.
+            this.model = `models/${modelParams.model}`;
+          }
+        this.modelParams=modelParams;
+      }
+
+
+ async generateImages(prompt:string,requestConfig?:BaseImageParams,requestOptions?:SingleRequestOptions
+      ):Promise<GenerateImagesResult>{
+        const Params: ImageModelParams={
+            model:this.model,
+            ...requestConfig,
+        }
+
+        const generativeModelRequestOptions: SingleRequestOptions = {
+            ...this._requestOptions,
+            ...requestOptions,
+          };
+
+        return generateImages(this.apiKey,this.model,{
+            instances:[{prompt}],
+            paramaters: {
+                ...this.modelParams,
+                ...Params, 
+            }, 
+        },generativeModelRequestOptions);
+      }
+}

src/requests/request-helpers.ts

@@ -29,6 +29,7 @@ import {
   GoogleGenerativeAIError,
   GoogleGenerativeAIRequestInputError,
 } from "../errors";
+import { generateImages } from "../methods/generate-image";
 
 export function formatSystemInstruction(
   input?: string | Part | Content,

src/requests/request.ts

@@ -40,6 +40,8 @@ export enum Task {
   COUNT_TOKENS = "countTokens",
   EMBED_CONTENT = "embedContent",
   BATCH_EMBED_CONTENTS = "batchEmbedContents",
+  GENERATE_IMAGES="predict",
+
 }
 
 export class RequestUrl {

types/requests.ts

@@ -34,6 +34,44 @@ export interface BaseParams {
   generationConfig?: GenerationConfig;
 }
 
+export interface BaseImageParams{
+
+  guidanceScale?:number;
+  seed?:number;
+  safetyFilterLevel?:safetyFilterLevel;
+  personGeneration?:PersonGeneration;
+  includeSafetyAttributes?:boolean;
+  includeRaiReason?:boolean;
+  language?:ImagePromptLanguage;
+  outputMimeType?:string;
+  outputCompressionQuality?:number;
+  addWatermark?:boolean;
+  enhancePrompt?:boolean;
+}
+enum safetyFilterLevel{
+  BLOCK_LOW_AND_ABOVE='BLOCK_LOW_AND_ABOVE',
+  BLOCK_NONE='BLOCK_NONE',
+  BLOCK_ONLY_HIGH='BLOCK_ONLY_HIGH',
+  BLOCK_MEDIUM_AND_ABOVE='BLOCK_MEDIUM_AND_ABOVE'
+
+}
+
+enum PersonGeneration{
+  DONT_ALLOW='DONT_ALLOW',
+  ALLOW_ADULT='ALLOW_ADULT',
+  ALLOW_ALL='ALLOW_ALL'
+}
+enum ImagePromptLanguage{
+  auto ='auto',
+  en='en',
+  ja='ja',
+  ko='ko',
+  hi='hi'
+}
+
+export interface ImageModelParams extends BaseImageParams{
+  model: string;
+}
 /**
  * Params passed to {@link GoogleGenerativeAI.getGenerativeModel}.
  * @public
@@ -61,6 +99,16 @@ export interface GenerateContentRequest extends BaseParams {
   cachedContent?: string;
 }
 
+
+/**
+ * Request sent to `generateImage` endpoint.
+ * @public
+ */
+export interface GenerateImageRequest{
+  instances:Array<{prompt:string}>;
+    paramaters:BaseImageParams; 
+}
+
 /**
  * Request sent to `generateContent` endpoint.
  * @internal
@@ -153,6 +201,12 @@ export interface CountTokensRequest {
   contents?: Content[];
 }
 
+export interface ImageModelParams {
+  model: string;
+  aspectRatio?: string; 
+}
+
+
 /**
  * Params for calling {@link GenerativeModel.countTokens}
  * @internal

types/responses.ts

@@ -33,6 +33,24 @@ export interface GenerateContentResult {
   response: EnhancedGenerateContentResponse;
 }
 
+interface Image{
+  gcsUri?:string;
+  bytesBase64Encoded:string;
+  mimeType:string;
+}
+
+
+interface  GeneratedImage{
+  image?:Image[];
+  raiFilterReason?:string[];
+  enhancedPrompt?:string[];
+
+}
+
+export interface GenerateImagesResult{
+  predictions?:GeneratedImage[]; 
+ }
+
 /**
  * Result object returned from generateContentStream() call.
  * Iterate over `stream` to get chunks as they come in and/or