Splitting the AI Test Kitchen into molecules and connecting to Imagen

6 min readMay 31, 2023

One fine day I decided to see what queries Google’s AI Test Kitchen mobile app sends and came across a working API of a closed neural network for Imagen image generation.

I’m sharing how I discovered it, what else interesting I found, and a link to the GitHub repository of the implementation.

Foreword

Things may have changed since the publication of this article, so if you’re reading these lines a few weeks/months/years later, there’s a 99% chance you won’t be able to repeat all my steps.

The AI Test Kitchen (hereinafter referred to as ATK) mobile app for Android, using Burp Suite Proxy, will be investigated.

At the time this article was published, there was one demo available on the ATK app, MusicLM. Imagen was not officially available, so it is surprising that the developers left Imagen as an available API method.

Analyzing the application

We open the ATK and start analyzing the traffic. From the very beginning we find a request to https://aitestkitchen.withgoogle.com/ to get the code of the latest version of the React application on Next.js. From this it follows that the application is not native, which has quite a big impact on performance.

Let’s try to generate music through the MusicLM demo.

The application sends a json request to another host https://content-aisandbox-pa.googleapis.co to the soundDemo method. A simple Bearer token is used for authorization

https://content-aisandbox-pa.googleapis.com/v1:soundDemo?alt=json
{
  "generationCount": 2,
  "input": {
    "textInput": "Ambient soundscape, light, eery, and dreary"
  },
  "sessionId": "c66c7436-f66f-4d65-b481-103c734b102a",
  "soundLengthSeconds": 30
}

The response we get is the generated tracks.

Notably, the soundLengthSeconds parameter has no effect on track length, but it is mandatory.

Going Deeper

In addition, another, more interesting GET request was sent during generation. It opens before us available api-methods and as a consequence, inaccessible through the demo interface.

https://content-aisandbox-pa.googleapis.com/$discovery/rest?key=AIzaSyBtrm0o5ab1c-Ec8ZuLcGt3oJAA5VWt3pY&pp=0&fields=fields%5B%22kind%22%5D%2Cfields%5B%22name%22%5D%2Cfields%5B%22version%22%5D%2Cfields%5B%22rootUrl%22%5D%2Cfields%5B%22servicePath%22%5D%2Cfields%5B%22resources%22%5D%2Cfields%5B%22parameters%22%5D%2Cfields%5B%22methods%22%5D%2Cfields%5B%22batchPath%22%5D%2Cfields%5B%22id%22%5D

If we look at this, we see this:

rate — Rate a LaMDA response.
soundDemo — Reproduces sound on given SoundDemoRequest.
deleteSession — Delete all data for a session.
upscaleImage — Upscales a given image as per the request params. This can be an expensive call depending on the flow — whether a model is triggered or previously stored image data is returned.
imageDemo — Generates a set of images for the query and specs in the request. Currently limits the maximum #images in the response to 8. If request sets a param >8, it’s capped to 8 and RPC will not fail on the limit validation.

What interested me the most is the imageDemo method. Could it be that Imagen demo?

I could send a request to the method right away, but there was a big problem — I didn’t know the parameters that imageDemo was accepting.

{
  "httpMethod": "POST",
  "description": "Generates a set of images for the query and specs in the request. Currently limits the maximum #images in the response to 8. If request sets a param >8, it's capped to 8 and RPC will not fail on the limit validation.",
  "flatPath": "v1:imageDemo",
  "request": {
    "$ref": "ImageDemoRequest"
  },
  "parameters": {},
  "response": {
    "$ref": "ImageDemoResponse"
  },
  "id": "aisandbox_pa.imageDemo",
  "parameterOrder": [],
  "path": "v1:imageDemo"
}

Digging into the frontend

If the respected developers for some unknown reason have not removed many unused API methods, then perhaps in WebPack chunks you can find a lot of interesting things too. Spoiler alert: yes.

After filtering out the chunks I didn’t need, I found the most suitable one — https://aitestkitchen.withgoogle.com/51ebd516-8ab4-450e-91f4-24539f887bdf/app/_next/static/chunks/955-89f740984868e27c.js

Demo list

Starting with line 6936, we can see a description of all the demos that were available in ATK.

Pretend — Name a place and LaMDA will suggest ways to develop your imagination.
List — Name a goal or topic and see how LaMDA can break it down into multiple lists of subtasks.
Puppies — Keep the conversation going and see where it leads. It’s just a fun, weird, open-ended conversation.

These three demos were available to everyone at the very initial launch of ATK. They demonstrated the capabilities of the LaMDA neural network, which is currently running in Bard, but is no longer available in ATK.

City — Describe an imaginary building and our text-image model will show you what it might look like.
Wobble — Imagine a monster and describe what it’s wearing. Using 2D-3D animation techniques, “wobble” it to make it dance.
Brainstorm — Imagine a scene, choose a style, and our text-to-image model will show you what it might look like.

These three demos demonstrated the work of Imagen, which we’ll connect to a little later. It’s weird that you can’t interact with these demos through the official interface, but their API is available to anyone.

Soundsmith — Describe a musical idea and hear how it comes to life using artificial intelligence

This is the MusicLM demo, which is available to anyone at the time of publication of this article.

Looking for imageDemo parameters

The search did not take long. The first result of the search leads to this code (line 9714):

gapi.client.aisandbox_pa.imageDemo({
    demo: "CITY_DREAMER",
    imageCount: 4,
    query: e,
    sessionId: t,
    imageSpec: {
        inlineUpscale: !0
    },
})

Looking for parameters of the upscaleImage method
Line 3711.

gapi.client.aisandbox_pa.upscaleImage({
 imageId: t
})

That looks pretty clear. That’s all the frontend needed. Thank you!

Forming requests

I used the soundDemo method as the basis and by changing the address and parameters I was able to successfully generate an imageDemo image:

{
  "demo": "BRAINSTORM",
  "imageCount": 4,
  "query": "Photo of a dog in a sunglasses. Oil",
  "sessionId": "by08oc78-8f67-4fb3-98e6-cd2332b87597",
  "imageSpec": {
    "inlineUpscale": false
  }
}

The response we get is base64 images and their id for upscale:

Of interest:

Censorship. A lot. A lot. Censorship.
The sessionId can be anything you want.
You can use different demos BRAINSTORM, WOBBLE, CITY_DREAMER.

Writing a python implementation

class Imagen:
    def __init__(self, bearer) -> None:
        self.bearer = bearer
        self.domain = "https://content-aisandbox-pa.googleapis.com"
        self.headers = {
            "Authorization": f"Bearer {self.bearer}",
            "Content-Type": "application/json",
        }

    def generate(self, prompt: str, demo: str, count: int) -> list:
        data = {
            "demo": demo,
            "imageCount": count,
            "query": prompt,
            "sessionId": str(uuid.uuid4()),
            "imageSpec": {"inlineUpscale": False},
        }
        response = requests.post(
            f"{self.domain}/v1:imageDemo?alt=json", headers=self.headers, json=data
        )
        if response.status_code == 200:
            return [
                Img(image["imageData"], image["imageId"], self)
                for image in response.json()["images"]
            ]
        else:
            raise Exception(f"Error: {response.json()['error']['message']}")

The same for increasing the resolution to 1024x1024:

class Img:
    def __init__(self, base: str, id: str, imagen: Imagen) -> None:
        self.id = id
        self.image = Image.open(BytesIO(base64.b64decode(base)))
        self.imagen = imagen

    def save(self, path: str) -> None:
        self.image.save(path)

    def show(self) -> None:
        self.image.show()

    def upscale(self) -> None:
        response = requests.post(
            f"{self.imagen.domain}/v1:upscaleImage?alt=json",
            headers=self.imagen.headers,
            json={
                "imageId": self.id,
            },
        )
        if response.status_code == 200:
            self.image = Image.open(
                BytesIO(base64.b64decode(response.json()["image"]["imageData"]))
            )
        else:
            raise Exception(f"Error: {response.json()['error']['message']}")

How do I get a Bearer token?

Sign up for https://aitestkitchen.withgoogle.com/ via your computer. Go to any available demo and start interacting with enabled developer tools (F12 key)