Полнофункциональное приложение для преобразования текста в изображение с искусственным интеллектом, которое можно создать за полдня

В качестве забавного проекта для жарких летних выходных я создал минимальное полнофункциональное веб-приложение для генеративного искусственного интеллекта с преобразованием текста в изображение с моделью Стабильная диффузия (развертывается через Amazon SageMaker JumpStart), FastAPI для веб-сервера и React для фронтенда.

Здесь позвольте мне провести вас шаг за шагом 😆.

Развертывание конечной точки модели

Чтобы развернуть модель Stable Diffusion, я использовал ярлык SageMaker JumpStart.

Во-первых, найдите базу Stable Diffusion 2.1, предварительно обученную на LAION-5B в рамках задачи машинного обучения SageMaker JumpStart Foundation Models: Image Generation.

Его можно развернуть одним щелчком мыши, так что давайте сделаем это!

Когда конечная точка будет готова к работе, вы можете нажать «Открыть блокнот», чтобы просмотреть пример кода и поэкспериментировать с ним. Записная книжка откроется в SageMaker Studio, но если вам удобнее играть на локальном сервере VSCode и Jupyter (как я), просто настройте свои учетные данные AWS локально и загрузите записную книжку.

В образце записной книжки есть три ключевые вспомогательные функции, которые позволяют отправить приглашение в модель Stable Diffusion и получить обратно сгенерированное изображение. Просто используйте их и поменяйте местами свое собственное приглашение, чтобы повеселиться:

# response = query_endpoint("cottage in impressionist style")
response = query_endpoint("a cat astronaut fighting aliens in space, realistic, high res")

img, prmpt = parse_response(response)

# Display hallucinated image
display_image(img,prmpt)

Создайте API

Быстрый и минимальный старт

Теперь, когда конечная точка модели работает должным образом, давайте создадим минимальный веб-API для пересылки запросов, которые будет отправлять интерфейс нашего веб-приложения.

Поскольку это минимальный проект, в корне проекта я просто создал одну папку для API (/api) и одну для пользовательского интерфейса (/ui).

Что касается API, я давно хотел попробовать FastAPI, поскольку он поставляется с пользовательским интерфейсом Swagger, который документирует ваши маршруты API из коробки. Получим, следуя официальному руководству:

# Install fastapi as well as the ASGI server Uvicorn
$ pip3 install fastapi
$ pip3 install uvicorn[standard]

Затем создайте папку нашего проекта с main.py файлом:

from typing import Union

from fastapi import FastAPI

app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.get("/items/{item_id}")
def read_item(item_id: int, q: Union[str, None] = None):
    return {"item_id": item_id, "q": q}

Запустите сервер разработки uvicorn в папке /api:

$ uvicorn main:app --reload

Как и было обещано, API работает в http://127.0.0.1:8000, а пользовательский интерфейс документации API Swagger мгновенно доступен в http://127.0.0.1:8000/docs 🚀!

Подключить API к модели

Теперь пришло время подключить наш API и конечную точку модели. Давайте добавим маршрут и его обработчик в main.py под примерами маршрутов:

# api/main.py
@app.get("/generate-image")
def generate_image(prompt: str):
    image, prmpt = utils.parse_response(utils.query_endpoint(prompt))
    print(image)
    return {"out": "yeah"}

Чтобы заставить его работать, я сначала взял примеры функций query_endpoint() и parse_response() из примера блокнота SageMaker JumpStart и упаковал их в файл utils.py, а затем импортировал и использовал их в обработчике маршрута. Поскольку я не знаю, в каком формате будет отправлено сгенерированное изображение, я сначала просто вывожу его на консоль и отправляю фиктивный JSON в качестве ответа моего API, который выкрикивает «да».

После сохранения файла main.py мы видим, что новый маршрут /generate-image автоматически появляется в пользовательском интерфейсе документа Swagger. Поскольку наш обработчик ожидает строковый параметр prompt, Swagger удобно предоставляет пользовательский интерфейс для ввода текста, чтобы мы могли его попробовать!

Когда я набрал подсказку «Астронавт-единорог» и нажал «Выполнить», модель отправила обратно сгенерированное изображение. Но как это выглядит? Ну, это массив значений канала RGB для каждого пикселя изображения! Не то, что наше веб-приложение может показать, но это изображение!

Обработайте и сохраните изображение

Чтобы превратить магические числа пикселей в изображение, которое может видеть человек, я использовал numpy и PIL (Pillow). Для быстрого теста (и поскольку мы еще не построили наш интерфейс) я просто добавил еще одну служебную функцию для преобразования пикселей в массив numpy, а затем объект изображения, который затем сохраняется на диск.

# api/utils.py
from PIL import Image
import numpy as np

# ...

def save_image(pixels):
    arr = np.array(pixels, dtype=np.uint8)
    img = Image.fromarray(arr)
    img.save("new.png")

В обработчике нашего маршрута /generate-image теперь мы можем просто подключить массив пикселей изображения из ответа на эту новую служебную функцию и на данный момент отправить приглашение в качестве ответа API.

# api/main.py
# ...
@app.get("/generate-image")
def generate_image(prompt: str):
    image, prmpt = utils.parse_response(utils.query_endpoint(prompt))
    utils.save_image(image)
    return {"prompt": prmpt}

Давайте проверим это в Swagger с новой подсказкой «Космонавт-единорог в космосе, все тело сбоку». И вот что я нашел в своей папке API — Stable Diffusion дал мне единорога в космосе!

Поддержка КОРС

Прежде чем мы перейдем к внешнему интерфейсу, нам нужно сделать еще одну вещь в API: поддержать CORS, чтобы браузер позволял нашему интерфейсному приложению выполнять вызовы AJAX на наш сервер API.

Самый простой способ сделать это — добавить в наше приложение API FastAPI CORSMiddleware:

# api/main.py
from fastapi.middleware.cors import CORSMiddleware

# Support CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

Создайте пользовательский интерфейс

Теперь пришло время для лица нашего минимального приложения!

строительные леса

Я не прикасался к React в течение многих лет, но слышал, что это все еще популярный интерфейсный фреймворк. На этот раз, однако, вместо классической CRA я использую более легкую и быструю альтернативу Vite.

Для компонентов пользовательского интерфейса я попробовал Чакру по ее приятному гайду по работе с Vite.

$ npm i @chakra-ui/react @emotion/react @emotion/styled framer-motion

После установки просто добавьте теги <ChakraProvider> в корень React:

// ui/src/main.tsx
import { ChakraProvider } from "@chakra-ui/react";

ReactDOM.createRoot(document.getElementById("root")!).render(
  <React.StrictMode>
    <ChakraProvider>
      <App />
    </ChakraProvider>
  </React.StrictMode>
);

Минималистичный пользовательский интерфейс

Наше веб-приложение делает для пользователя одну простую вещь: принимает текстовое приглашение и показывает сгенерированное изображение. Итак, минимум, который нам нужен, — это три компонента пользовательского интерфейса:

1. ввод текста для ввода подсказки
2. кнопка отправки для отправки подсказки
3. компонент изображения для рендеринга сгенерированного изображения

С готовыми компонентами Chakra-UI Input, Button и Image сделать красивый пользовательский интерфейс совсем несложно. Для макета я использовал Container и InputGroup, затем немного подправил поля здесь и там. Конечно, нам также нужен хороший Heading просто для полноты картины.

// ui/src/App.tsx
// ...

return (
    <Container maxWidth={"2xl"} marginTop={30} centerContent>
      <Heading margin={8}>Image Generator</Heading>
      <InputGroup>
        <Input
          pr="4.5rem"
          value={prompt}
          placeholder="Enter your prompt"
          onChange={onInputChange}
        />
        <InputRightElement width={"6rem"}>
          <Button onClick={onButtonClick} isDisabled={isLoading}>
            Generate!
          </Button>
        </InputRightElement>
      </InputGroup>
      {imgSrc ? (
        <Box boxSize={"l"} marginTop={5}>
          <Image src="https://picsum.photos/640" />
        </Box>
      ) : null}
    </Container> 
);

Чтобы настроить размер и положение изображения, я использовал фиктивный сервис изображений https://picsum.photos/, запросив у него динамическое изображение шириной 640 пикселей. И вот все в сборе:

Состояния пользовательского интерфейса, обработчики событий и вызов API

Теперь пришло время заставить наш пользовательский интерфейс работать! В нашем минимальном пользовательском интерфейсе есть три части динамических состояний:

1. текст prompt, который привязан к значению ввода текста
2. imgSrc, часть данных, которая должна поступать из бэкенда, а также определяет, следует ли вообще отображать компонент пользовательского интерфейса изображения.
3. состояние isLoading для лучшего взаимодействия с пользователем. Когда для него установлено значение true, я хочу отключить кнопку отправки и показать загрузчик, чтобы пользователь знал, что наше приложение усердно работает над созданием этого изображения!

// ui/src/App.tsx
function App() {
  // states
  const [prompt, setPrompt] = useState("");
  const [isLoading, setIsLoading] = useState(false);
  const [imgSrc, setImgSrc] = useState("https://picsum.photos/640"); // we'll replace the initial value to null later so that the image component is not rendered when there is no imgSrc
//...
}

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

1. событие change ввода текста, которое обновляет состояние prompt каждый раз, когда пользователь что-либо меняет в вводе текста
2. событие click кнопки отправки, которое, по-видимому, запускает API, но также соответствующим образом обновляет состояние isLoading

Для первого мы определяем обработчик события onInputChange следующим образом:

// ui/src/App.tsx
const onInputChange = (e: ChangeEvent<HTMLInputElement>) => {
  setPrompt(e.target.value);
};

Для последнего мы определяем обработчик события onButtonClick. Это асинхронно, так как ему нужно вызывать API и ждать ответа. Я также обернул вызов API некоторой обработкой ошибок и добавил логику для обновления состояния isLoading в нужное время:

// ui/src/App.tsx
const onButtonClick = async () => {
  setIsLoading(true);
  try {
    const response = await (
      await fetch(`http://127.0.0.1:8000/generate-image?prompt=${prompt}`)
    ).json();
    console.log(response);
  } catch (err: any) {
    // catch any runtime error
    console.log(err.message);
  } finally {
    setIsLoading(false);
  }
};

И чтобы показать пользователю, когда мы загружаемся, давайте добавим счетчик под текстовым вводом:

// ui/src/App.tsx
// ...
    </InputGroup>    
    {isLoading ? (
        <Spinner
          thickness="4px"
          speed="0.65s"
          emptyColor="gray.200"
          color="blue.500"
          size="xl"
          marginTop={6}
        />
      ) : imgSrc ? (
        <Box boxSize={"l"} marginTop={5}>
          {/* <Image src={`data:image/png;base64,${imageData}`} /> */}
          <Image src={imgSrc} />
        </Box>
      ) : null}
  </Container>

На этом этапе отправка запроса в пользовательском интерфейсе (нажатие кнопки «Создать!») фактически вызывает API и запускает создание изображения. Однако наш API еще не отправляет изображение обратно в ответ. Вместо этого он сохраняет изображение для себя!

Итак, давайте это исправим.

Бэкэнд и Фронтенд интеграция

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

Отправить обратно изображение в ответе API

Чтобы получить изображение из API, мы отправим данные изображения в виде строки в кодировке base64. Тогда наш ответ API в формате JSON будет выглядеть так:

{
    "prompt": prmpt, // the prompt string from user 
    "img_base64": img_str // the generated image as a base64 data string
}

Я все равно хотел бы сохранить сгенерированное изображение где-нибудь на сервере, поэтому я переставил функции утилиты следующим образом:

Во-первых, общая служебная функция для преобразования массивов пикселей из модели стабильной диффузии в изображение PIL:

# utils.py
def pixel_to_image(pixel_array):
   arr = np.array(pixel_array, dtype=np.uint8)
   img = Image.fromarray(arr)
   return img

Во-вторых, функция сохранения изображения в папку generated_image на сервере:

# utils.py
def save_image(img, filePath="generated_images/new.png"):
   img.save(filePath)

В-третьих, функция для преобразования изображения PIL в строку данных изображения в кодировке base64, которую мы можем отправить в нашем ответе для отображения интерфейса:

# utils.py
def image_to_base64_str(img, format="PNG"):
   buffered = BytesIO()
   img.save(buffered, format=format)
   img_str = base64.b64encode(buffered.getvalue())
   return img_str

Теперь с помощью этих служебных функций наш обработчик маршрутов API может выполнять свою работу:

# main.py
@app.get("/generate-image")
def generate_image(prompt: str):
    pixel_array, prmpt = utils.parse_response(utils.query_endpoint(prompt))
    image = utils.pixel_to_image(pixel_array)

    utils.save_image(
        image, filePath=f"generated_images/{str(datetime.datetime.now())}.png"
    )
    img_str = utils.image_to_base64_str(image)

    return {"prompt": prmpt, "img_base64": img_str}

Чтобы не перезаписывать new.png каждый раз, я использовал метку времени сохранения изображения в качестве имени изображения.

Показать изображение во внешнем интерфейсе

Теперь, когда наш API отправляет обратно данные изображения, давайте обновим логику обработки API внешнего интерфейса, чтобы извлечь их и передать в состояние imgSrc:

// App.tsx
const onButtonClick = async () => {
    setIsLoading(true);
    try {
      const response = await (
        await fetch(`http://127.0.0.1:8000/generate-image?prompt=${prompt}`)
      ).json();
      console.log(response);
      const lastPrompt = response["prompt"];
      const imgBase64 = response["img_base64"];
      setImgSrc(`data:image/png;base64, ${imgBase64}`);
    } catch (err: any) {
      // catch any runtime error
      console.log(err.message);
    } finally {
      setIsLoading(false);
    }
};

Использование строки данных изображения base64 в качестве src HTML-тега <img> не только работает так же хорошо, как использование URI изображения, но также экономит пользователю один сетевой путь туда и обратно, чтобы получить изображение 😌.

Теперь давайте попросим модель сгенерировать изображение семейства лам в дикой природе.

Вуаля! 🦙

Отказ от ответственности: для простоты я поместил все маршруты и обработчики API в main.py, а также все компоненты и логику пользовательского интерфейса в App.tsx . Для любого проекта производственного уровня, который вам необходимо поддерживать, вы должны правильно организовать и разбить код на модули.

Заключение

Здесь я рассказал вам, как создать минимальное веб-приложение для генерирования текста в изображения, состоящее из модели Stable Diffusion, развернутой с помощью SageMaker JumpStart, серверной части Python API, созданной с помощью FastAPI, и внешнего веб-приложения, созданного с помощью React (Vite + Чакра-УИ). Довольно приятно получить полнофункциональное веб-приложение для генеративного ИИ и запустить его так быстро и легко, не так ли?

Вы можете найти весь код приложения в моем репозитории GitHub.

Однако и бэкенд, и внешний интерфейс работают только на локальных серверах разработки. Более того, конечная точка модели, развернутая в один клик, оплачивается посекундно независимо от того, вызываете вы ее или нет. Поэтому просто удалите конечную точку модели, а в будущих статьях я хотел бы показать вам, как:

• собрать и выпустить интерфейс (пока не знаю, куда)
• упаковать и выпустить локальный API либо в контейнерное приложение, либо в бессерверную серверную часть с помощью AWS Lambda и Amazon API Gateway.
• разверните конечную точку бессерверной модели, которая взимает плату только тогда, когда вы ее фактически используете

Подтверждение

Этот проект был вдохновлен этим удивительным YouTube видео. Салют гениальному творцу Николя Ренотт!