Skip to content

Commit

Permalink
Davidmotson.imagen java (#6618)
Browse files Browse the repository at this point in the history
Co-authored-by: David Motsonashvili <davidmotson@google.com>
  • Loading branch information
davidmotson and David Motsonashvili committed Jan 30, 2025
1 parent be2313d commit e881958
Showing 1 changed file with 75 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/*
* Copyright 2025 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package com.google.firebase.vertexai.java

import androidx.concurrent.futures.SuspendToFutureAdapter
import com.google.common.util.concurrent.ListenableFuture
import com.google.firebase.vertexai.ImagenModel
import com.google.firebase.vertexai.type.ImagenGCSImage
import com.google.firebase.vertexai.type.ImagenGenerationResponse
import com.google.firebase.vertexai.type.ImagenInlineImage

/**
* Wrapper class providing Java compatible methods for [ImagenModel].
*
* @see [ImagenModel]
*/
public abstract class ImagenModelFutures internal constructor() {
/**
* Generates an image, returning the result directly to the caller.
*
* @param prompt The main text prompt from which the image is generated.
*/
public abstract fun generateImages(
prompt: String,
): ListenableFuture<ImagenGenerationResponse<ImagenInlineImage>>

/**
* Generates an image, storing the result in Google Cloud Storage and returning a URL
*
* @param prompt The main text prompt from which the image is generated.
* @param gcsUri Specifies the GCS bucket in which to store the image.
*/
public abstract fun generateImages(
prompt: String,
gcsUri: String,
): ListenableFuture<ImagenGenerationResponse<ImagenGCSImage>>

/** Returns the [ImagenModel] object wrapped by this object. */
public abstract fun getImageModel(): ImagenModel

private class FuturesImpl(private val model: ImagenModel) : ImagenModelFutures() {
override fun generateImages(
prompt: String,
): ListenableFuture<ImagenGenerationResponse<ImagenInlineImage>> =
SuspendToFutureAdapter.launchFuture { model.generateImages(prompt) }

override fun generateImages(
prompt: String,
gcsUri: String,
): ListenableFuture<ImagenGenerationResponse<ImagenGCSImage>> =
SuspendToFutureAdapter.launchFuture { model.generateImages(prompt, gcsUri) }

override fun getImageModel(): ImagenModel = model
}

public companion object {

/** @return a [ImagenModelFutures] created around the provided [ImagenModel] */
@JvmStatic public fun from(model: ImagenModel): ImagenModelFutures = FuturesImpl(model)
}
}

0 comments on commit e881958

Please sign in to comment.