Одним прекрасным днём я решил посмотреть, какие запросы отправляет мобильное Google-приложение AI Test Kitchen и наткнулся на пока-что рабочий API закрытой нейросети для генерации изображений Imagen.

Делюсь тем, как я это обнаружил, что еще интересного нашел, и ссылкой на GitHub-репозиторий реализации.

Предисловие

С момента публикации этой статьи все могло измениться, поэтому если вы читаете эти строки несколько недель/месяцев/лет спустя, то с вероятностью 99% вы не сможете повторить все мои действия.

Исследовать мы будет мобильное приложение AI Test Kitchen (далее - ATK) под Android с помощью Burp Suite Proxy.

На момент публикации этой статьи в приложении ATK была доступна одна демонстрация - MusicLM. Imagen не был официально доступен, поэтому удивительно, что разработчики оставили Imagen в качестве доступного метода API.

Анализируем приложение

Открываем ATK и начинаем анализировать трафик. С самого начала мы обнаруживаем запрос к https://aitestkitchen.withgoogle.com/ для получение кода последней версии React-приложения на Next.js. Из этого следует, что приложение не является нативным, что довольно сильно сказывается на производительности.

Давайте попробуем сгенерировать музыку через демо MusicLM.

Приложение отправляет json-запрос на другой хост https://content-aisandbox-pa.googleapis.co к методу soundDemo. Для авторизации используется простой Bearer-токен

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
}

Ответом мы получаем сгенерированные треки.

Примечательно, что параметр soundLengthSeconds никак не влияет на длинную треков, но он является обязательным.

Углубляемся

Помимо этого, при генерации был отправлен другой, более интересный GET-запрос. Он открывает перед нами доступные api-методы и как следствие, недоступные через интерфейс демо.

Длинный ответ

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

{
  "id": "aisandbox_pa:v1",
  "version": "v1",
  "kind": "discovery#restDescription",
  "servicePath": "",
  "batchPath": "batch",
  "rootUrl": "https://aisandbox-pa.googleapis.com/",
  "parameters": {
    "upload_protocol": {
      "description": "Upload protocol for media (e.g. \"raw\", \"multipart\").",
      "type": "string",
      "location": "query"
    },
    "fields": {
      "location": "query",
      "type": "string",
      "description": "Selector specifying which fields to include in a partial response."
    },
    "oauth_token": {
      "description": "OAuth 2.0 token for the current user.",
      "location": "query",
      "type": "string"
    },
    "alt": {
      "location": "query",
      "description": "Data format for response.",
      "default": "json",
      "type": "string",
      "enum": [
        "json",
        "media",
        "proto"
      ],
      "enumDescriptions": [
        "Responses with Content-Type of application/json",
        "Media download with context-dependent Content-Type",
        "Responses with Content-Type of application/x-protobuf"
      ]
    },
    "callback": {
      "type": "string",
      "description": "JSONP",
      "location": "query"
    },
    "key": {
      "type": "string",
      "description": "API key. Your API key identifies your project and provides you with API access, quota, and reports. Required unless you provide an OAuth 2.0 token.",
      "location": "query"
    },
    "quotaUser": {
      "type": "string",
      "location": "query",
      "description": "Available to use for quota purposes for server-side applications. Can be any arbitrary string assigned to a user, but should not exceed 40 characters."
    },
    "prettyPrint": {
      "location": "query",
      "type": "boolean",
      "description": "Returns response with indentations and line breaks.",
      "default": "true"
    },
    "access_token": {
      "type": "string",
      "location": "query",
      "description": "OAuth access token."
    },
    "$.xgafv": {
      "enum": [
        "1",
        "2"
      ],
      "location": "query",
      "description": "V1 error format.",
      "type": "string",
      "enumDescriptions": [
        "v1 error format",
        "v2 error format"
      ]
    },
    "uploadType": {
      "location": "query",
      "description": "Legacy upload protocol for media (e.g. \"media\", \"multipart\").",
      "type": "string"
    }
  },
  "name": "aisandbox_pa",
  "resources": {
    "waitlistEntries": {
      "methods": {
        "delete": {
          "id": "aisandbox_pa.waitlistEntries.delete",
          "parameterOrder": [
            "name"
          ],
          "httpMethod": "DELETE",
          "flatPath": "v1/waitlistEntries/{waitlistEntriesId}",
          "path": "v1/{+name}",
          "parameters": {
            "name": {
              "required": true,
              "pattern": "^waitlistEntries/[^/]+$",
              "type": "string",
              "location": "path",
              "description": "Required. The name of the waitlist entry to delete. This should always be \"waitlistEntries/me\""
            }
          },
          "response": {
            "$ref": "Empty"
          },
          "description": "Deletes the waitlist entry for the current user if there is one. Errors with NOT_FOUND if no waitlist entry exists for the current user. The name should always be \"waitlistEntries/me\"."
        },
        "create": {
          "request": {
            "$ref": "WaitlistEntry"
          },
          "path": "v1/waitlistEntries",
          "id": "aisandbox_pa.waitlistEntries.create",
          "flatPath": "v1/waitlistEntries",
          "parameterOrder": [],
          "httpMethod": "POST",
          "response": {
            "$ref": "WaitlistEntry"
          },
          "description": "Creates a waitlist entry for the current user if it doesn't already exist. If a waitlist entry already exists for the current user, then this call errors with ALREADY_EXISTS. See [1] below for details. None of the fields can be set by the client, so an empty entry should be passed as the body. [1] The waitlist entry is always based on the authenticated user's ID, so this behaves like a standard Create method with a User-specified id as described here: https://google.aip.dev/133#user-specified-ids",
          "parameters": {
            "frontendCommitHash": {
              "location": "query",
              "description": "The commit hash of the web frontend. This is used for tracking the version of the user consent text/UI at the time the user joined the waitlist. Represented as a hex-encoded string version of the SHA1 commit hash.",
              "type": "string"
            }
          }
        },
        "get": {
          "parameterOrder": [
            "name"
          ],
          "description": "Gets the waitlist entry for the current user if there is one. Errors with NOT_FOUND if no waitlist entry exists for the current user. The name should always be \"waitlistEntries/me\".",
          "response": {
            "$ref": "WaitlistEntry"
          },
          "parameters": {
            "name": {
              "location": "path",
              "type": "string",
              "description": "Required. The name of the waitlist entry to retrieve. This should always be \"waitlistEntries/me\"",
              "required": true,
              "pattern": "^waitlistEntries/[^/]+$"
            }
          },
          "path": "v1/{+name}",
          "httpMethod": "GET",
          "flatPath": "v1/waitlistEntries/{waitlistEntriesId}",
          "id": "aisandbox_pa.waitlistEntries.get"
        }
      }
    },
    "v1": {
      "methods": {
        "rate": {
          "response": {
            "$ref": "RateLamdaResponseResponse"
          },
          "httpMethod": "POST",
          "parameterOrder": [],
          "path": "v1:rate",
          "id": "aisandbox_pa.rate",
          "request": {
            "$ref": "RateLamdaResponseRequest"
          },
          "flatPath": "v1:rate",
          "parameters": {},
          "description": "Rate a LaMDA response."
        },
        "soundDemo": {
          "path": "v1:soundDemo",
          "id": "aisandbox_pa.soundDemo",
          "httpMethod": "POST",
          "parameterOrder": [],
          "response": {
            "$ref": "SoundDemoResponse"
          },
          "request": {
            "$ref": "SoundDemoRequest"
          },
          "flatPath": "v1:soundDemo",
          "description": "Reproduces sound on given SoundDemoRequest.",
          "parameters": {}
        },
        "deleteSession": {
          "path": "v1:deleteSession",
          "id": "aisandbox_pa.deleteSession",
          "flatPath": "v1:deleteSession",
          "parameters": {},
          "httpMethod": "POST",
          "request": {
            "$ref": "DeleteSessionRequest"
          },
          "parameterOrder": [],
          "response": {
            "$ref": "DeleteSessionResponse"
          },
          "description": "Delete all data for a session."
        },
        "demo": {
          "path": "v1:demo",
          "request": {
            "$ref": "DemoRequest"
          },
          "id": "aisandbox_pa.demo",
          "flatPath": "v1:demo",
          "response": {
            "$ref": "DemoResponse"
          },
          "parameters": {},
          "httpMethod": "POST",
          "description": "Invoke a specified demo.",
          "parameterOrder": []
        },
        "upscaleImage": {
          "response": {
            "$ref": "UpscaleImageResponse"
          },
          "request": {
            "$ref": "UpscaleImageRequest"
          },
          "description": "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.",
          "path": "v1:upscaleImage",
          "flatPath": "v1:upscaleImage",
          "parameters": {},
          "httpMethod": "POST",
          "parameterOrder": [],
          "id": "aisandbox_pa.upscaleImage"
        },
        "checkAppAvailability": {
          "parameters": {},
          "request": {
            "$ref": "CheckAppAvailabilityRequest"
          },
          "description": "Checks whether AI Test Kitchen is available",
          "path": "v1:checkAppAvailability",
          "id": "aisandbox_pa.checkAppAvailability",
          "httpMethod": "POST",
          "flatPath": "v1:checkAppAvailability",
          "response": {
            "$ref": "CheckAppAvailabilityResponse"
          },
          "parameterOrder": []
        },
        "addUserAcknowledgement": {
          "path": "v1:addUserAcknowledgement",
          "response": {
            "$ref": "AddUserAcknowledgementResponse"
          },
          "parameters": {},
          "flatPath": "v1:addUserAcknowledgement",
          "httpMethod": "POST",
          "parameterOrder": [],
          "request": {
            "$ref": "AddUserAcknowledgementRequest"
          },
          "id": "aisandbox_pa.addUserAcknowledgement",
          "description": "Creates a record of user acknowledgement."
        },
        "checkUserAcknowledgement": {
          "id": "aisandbox_pa.checkUserAcknowledgement",
          "response": {
            "$ref": "CheckUserAcknowledgementResponse"
          },
          "description": "Check a record of user acknowledgement.",
          "flatPath": "v1:checkUserAcknowledgement",
          "path": "v1:checkUserAcknowledgement",
          "httpMethod": "POST",
          "parameterOrder": [],
          "request": {
            "$ref": "CheckUserAcknowledgementRequest"
          },
          "parameters": {}
        },
        "shareMedia": {
          "path": "v1:shareMedia",
          "httpMethod": "POST",
          "parameterOrder": [],
          "flatPath": "v1:shareMedia",
          "id": "aisandbox_pa.shareMedia",
          "description": "By default, shares the media with \"PUBLIC\", accessible to all users. Errors with NOT_FOUND if the media file doesn't exist or the requesting user is not authorized to share it.",
          "response": {
            "$ref": "Empty"
          },
          "request": {
            "$ref": "ShareMediaRequest"
          },
          "parameters": {}
        },
        "imageDemo": {
          "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"
        }
      }
    },
    "media": {
      "methods": {
        "get": {
          "id": "aisandbox_pa.media.get",
          "response": {
            "$ref": "Media"
          },
          "httpMethod": "GET",
          "parameterOrder": [
            "name"
          ],
          "parameters": {
            "name": {
              "type": "string",
              "pattern": "^media/[^/]+$",
              "description": "Required. Unique media_key to be used as the media URL param and for fetching from the server. E.g. \"media/4b4e6936-f1a2-49ba-91bc-3cf2252a264f\".",
              "location": "path",
              "required": true
            }
          },
          "description": "Fetches the media for the requested key. Does not require the user to be authenticated. Errors with NOT_FOUND if the media file doesn't exist or is not shared with \"PUBLIC\".",
          "flatPath": "v1/media/{mediaId}",
          "path": "v1/{+name}"
        }
      }
    }
  }
}

Если мы посмотрим на methods, то увидим вот это:

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.

Больше всего меня заинтересовал метод imageDemo. Неужели это та самая демо-версия Imagen?

Можно было бы сразу отправить запрос к методу, но была большая проблема - я не знал параметров, которые принимал imageDemo.

{
  "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"
}

Копаемся во фронтенде

Если уважаемые разработчики по непонятным причинам не убрали многие не используемые методы API, то, возможно в WebPack чанках тоже можно будет найти много интересного. Спойлер: да.

Отфильтровав не нужные чанки, я нашёл самый подходящий - https://aitestkitchen.withgoogle.com/51ebd516-8ab4-450e-91f4-24539f887bdf/app/_next/static/chunks/955-89f740984868e27c.js

Список демо

Начиная со строки 6936, мы можем увидеть описание всех демо, которые были доступны в ATK.

Много кода

switch (e) {
    case "Pretend":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Explore your imagination.",
                info: [{
                        title: "What should I do?",
                        description: "Name a place and LaMDA will offer paths to explore your imagination.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If LaMDA generates interesting scene descriptions that are relevant to your idea.",
                    },
                ],
        };
    case "List":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Break down a complex goal or topic.",
                info: [{
                        title: "What should I do?",
                        description: "Name a goal or topic and see how much LaMDA can break it down into multiple lists of subtasks.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If LaMDA generates useful lists of subtasks, some of which you might not have thought of.",
                    },
                ],
        };
    case "Puppies":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Talk to a tennis ball about dogs.",
                info: [{
                        title: "What should I do?",
                        description: "Roll with the conversation and see where it goes. It\u2019s just a fun, kinda-weird, open-ended chat.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If LaMDA, no matter what you ask it, keeps the conversation going while bringing the topic back to dogs.",
                    },
                ],
        };
    case "City":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Dream up a new city.",
                info: [{
                        title: "What should I do?",
                        description: "Describe an imaginary building, and our text-to-image model will show you what it might look like.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If our text-to-image model generates an interesting and accurate building based on your description.",
                    },
                ],
                backgroundImage: io[ao],
        };
    case "Wobble":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Imagine a monster and make it wobble.",
                info: [{
                        title: "What should I do?",
                        description: "Imagine a monster and describe what it's wearing. Using 2D-to-3D animation techniques, \u201cwobble\u201d it to make it dance.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If our text-to-image model generates monsters that are relevant to your description.",
                    },
                ],
                backgroundImage: so[co],
        };
    case "Swim":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Search What I Mean",
                info: [{
                        title: "What should I do?",
                        description: "Have a topic in mind? Search anything related to the topic, it will remember the context you are in and provide some helpful hints.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If our model is able to correctly infer what you are asking at each step.",
                    },
                ],
        };
    case "Brainstorm":
        return {
            buttonLabel: "Launch demo",
                finalTitle: "Dream of anything.",
                info: [{
                        title: "What should I do?",
                        description: "Imagine a scene, pick a style, and our text-to-image model will show you what it might look like.",
                    },
                    {
                        title: "What should I give feedback on?",
                        description: "If our text-to-image model generates scenes that are relevant to your description.",
                    },
                ],
                backgroundImage: uo[lo],
        };
    case "Soundsmith":
        return {
            buttonLabel: "Launch demo",
                finalTitle:
                "Describe a musical idea and hear it come to life with AI",
                info: [{
                        title: "How to make a good prompt",
                        html: (0, i.jsxs)("ul", {
                            children: [
                                (0, i.jsx)("li", {
                                    children: "Be very descriptive, include electronic or classical instruments",
                                }),
                                (0, i.jsx)("li", {
                                    children: "Mention the vibe, mood or emotion you want to create",
                                }),
                                (0, i.jsx)("li", {
                                    children: "Certain queries that mention specific artists or include vocals will not be generated",
                                }),
                            ],
                        }),
                    },
                    {
                        title: "Improve the model by giving out trophies",
                        html: (0, i.jsx)("ul", {
                            children: (0, i.jsxs)(go, {
                                children: [
                                    (0, i.jsx)("div", {
                                        children: "Which track is better? Give it a trophy. This feedback will improve MusicLM for everyone.",
                                    }),
                                    (0, i.jsx)(mo, {
                                        src: to(
                                            "/website/season-3/musiclm-feedback.svg"
                                        ),
                                        alt: "Trophy icon buttons to indicate how to leave user feedback.",
                                        width: 116,
                                        height: 64,
                                    }),
                                ],
                            }),
                        }),
                    },
                    {
                        title: "Remember to save your creations",
                        html: (0, i.jsx)("ul", {
                            children: (0, i.jsx)("li", {
                                children: "Audio tracks are not automatically saved. If you want to save your tracks, click the more icon to download",
                            }),
                        }),
                    },
                ],
                backgroundImage: fo[0],
        };

Pretend - Назовите место, и LaMDA предложит пути для развития вашего воображения.
List - Назовите цель или тему и посмотрите, насколько LaMDA может разбить ее на несколько списков подзадач.
Puppies - Поддерживайте разговор и смотрите, куда он приведет. Это просто веселый, странный, открытый разговор.

Эти три демо были доступны всем при самом начальном запуске ATK. Они демонстрировали возможности нейросети LaMDA, которая в настоящее время работает в Bard, но более не доступной в ATK.

City - Опишите воображаемое здание, и наша текстово-изобразительная модель покажет вам, как оно может выглядеть.
Wobble - Представьте себе монстра и опишите, во что он одет. Используя технику 2D-3D анимации, "покачайте" его, чтобы заставить танцевать.
Brainstorm - Представьте себе сцену, выберите стиль, и наша модель преобразования текста в изображение покажет вам, как это может выглядеть.

Эти три демо демонстрировали работу Imagen, к которому мы подключимся чуть позже. Весьма странно, что через официальный интерфейс невозможно взаимодействовать с этими демо, но при том их API доступен любому желающему.

Soundsmith - Опишите музыкальную идею и услышьте, как она воплощается в жизнь с помощью искусственного интеллекта

Это демо MusicLM, которое доступно всем желающим на момент публикации статьи.

Ищем параметры imageDemo

Искать долго не пришлось. Первый же результат поиска выводит на вот этот код (строка 9714):

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

Иищем параметры метода upscaleImage

Строка 3711.

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

Выглядит весьма понятно. Это всё, что требовалось от фронтенда. Спасибо!

Формируем запросы

За основу я использовал метод soundDemo, поменяв адрес и параметры, я смог успешно сгенерировать изображения imageDemo:

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

Ответом мы получаем base64 изображения и их id для upscale:

Из интересного:

Цензура. Очень. Много. Цензуры.
sessionId можно указывать абсолютно любой.
Можно использовать разные демо BRAINSTORM, WOBBLE, CITY_DREAMER.

Пишем реализацию на python

Здесь нет ничего сложного.

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']}")

Тоже самое для увеличения разрешения до 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']}")

Отлично. Написали.

Как получить Bearer токен?

Регистрируемся на https://aitestkitchen.withgoogle.com/ через компьютер. Переходим на любое доступное демо и начинаем взаимодействовать с включёнными инструментами разработчика (клавиша F12)

Берём любой запрос к API и видим Bearer токен в заголовках.

Заключение

Как установить, как подключиться - 0x7o/AITestKitchen

Больше примеров изображений - 0x7o/AITestKitchen/tree/main/examples

Что думаете Вы?

Расщепляем AI Test Kitchen на молекулы и подключаемся к Imagen