如何使用Jina來(lái)構(gòu)建一個(gè)使用Stable Diffusion生成圖像的Executor呢

什么是 Executor

Executor 是一個(gè)獨(dú)立的 gRPC 微服務(wù)，它可以在 DocumentArray 上執(zhí)行任務(wù)。Executor Hub 上有大量預(yù)構(gòu)建的 Executor 可供選擇，包含了各種常見(jiàn)的任務(wù)，如文本分類，圖像識(shí)別、目標(biāo)檢測(cè)等。

ExecutorHub:cloud.jina.ai/executors

為了讓你能夠更輕松地部署和管理微服務(wù)，我們正將 Executor 從 Flow 中分離出來(lái)。同時(shí)，也方便你更好地利用 Jina 的其他強(qiáng)大功能，比如：

利用 gRPC 和 protobuf 實(shí)現(xiàn)高效的網(wǎng)絡(luò)傳輸和序列化，更快地處理數(shù)據(jù)和模型之間的通信；

使用 DocArray 更準(zhǔn)確、靈活地表示多模態(tài)數(shù)據(jù)，以滿足不同場(chǎng)景下的需求；

“Array-first”概念，將輸入數(shù)據(jù)分批進(jìn)行模型推理，可以大幅提高吞吐量，使你的模型處理更加高效；

輕松地將 ML 模型部署到生產(chǎn)環(huán)境中，享受云原生所帶來(lái)的便利和絲滑。

此外，請(qǐng)密切關(guān)注即將推出的 Jina AI Cloud(cloud.jina.ai)，在 Jina Cloud 上免費(fèi)運(yùn)行模型部署。

Jina 吉祥物

前段時(shí)間，我們?cè)谥芪逡黄鸷染屏奶斓臅r(shí)候，突然聊到要不要給 Jina 選一個(gè)可愛(ài)的吉祥物，就像米其林輪胎人一樣。

酒后頭腦風(fēng)暴之后，我們最終的決定是「彩虹獨(dú)角獸蝴蝶小貓」，考慮到基因改造工程的復(fù)雜度，要怎么真正創(chuàng)造出這樣一個(gè)神奇的新生物呢？我們決定先動(dòng)手畫張圖：

但想一想，還有什么比使用 Jina 本身更好的方式來(lái)生成 Jina 的吉祥物呢？考慮到這一點(diǎn)，我們立馬開(kāi)發(fā)了一個(gè)圖像生成的 Executor 和 Deployment。因此，在這篇文章中，我們將介紹如何將模型構(gòu)建成 Executor、部署它、擴(kuò)展它以及與全世界共享它。

構(gòu)建 Executor

需要一個(gè) GPU 才能在本地運(yùn)行和部署這個(gè) Executor。但你也可以調(diào)整代碼，使用 Executor Hub Sandbox 版本，托管在 Jina AI Cloud上。

在 Jina 中部署服務(wù)時(shí)總是以 Executor 的形式進(jìn)行。Executor是一個(gè)Python類，用于轉(zhuǎn)換和處理 Document?？梢詫⑽谋?圖像編碼為向量、OCR、從 PDF 中提取表格等等，不僅限于圖像生成。

當(dāng)然如果你只是想把它用起來(lái)，而不是從頭開(kāi)始構(gòu)建它，可以直接跳到 Executor Hub 部分。

在本教程中，我們將重點(diǎn)關(guān)注 Executor 和 Deployment，而不會(huì)深入研究 Stable Diffusion 模型的復(fù)雜性。我們希望本教程適用于任何微服務(wù)或模型，而不是只適用特定用例。

以下大致就是我們希望 Executor 看起來(lái)的樣子。用戶傳入提示詞，Executor 使用該提示詞生成圖像，然后將該圖像傳回給用戶：

先決條件

您需要安裝 Hugging Face Diffusers，pip install diffusers[torch]。

現(xiàn)在，讓我們從整體上看一下 Executor 代碼，然后逐節(jié)分析：

我們將從創(chuàng)建開(kāi)始 text_to_image.py：

from docarray import DocumentArray
from jina import Executor, requests
import numpy as np

class TextToImage(Executor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        import torch
        from diffusers import StableDiffusionPipeline

        self.pipe = StableDiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16
        ).to("cuda")

    @requests
    def generate_image(self, docs: DocumentArray, **kwargs):
    # image here is in PIL format
        images = self.pipe(docs.texts).images
        
        for i, doc in enumerate(docs):
            doc.tensor = np.array(images[i])

Imports

from docarray import DocumentArray
from jina import Executor, requests
import numpy as np

注：Documents 和 DocumentArrays 是 Jina 的原生 IO 格式。

Executor 類

class TextToImage(Executor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        import torch
        from diffusers import StableDiffusionPipeline

        self.pipe = StableDiffusionPipeline.from_pretrained(
            "CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16
        ).to("cuda")

所有 Executor 都是從 Jina 的 Executor 類創(chuàng)建的。用戶可定義的參數(shù)是方法中定義的參數(shù)__init__()。

Requests 裝飾器

@requests
def generate_image(self, docs: DocumentArray, **kwargs):
    # image here is in PIL format
    images = self.pipe(docs.texts).images

    for i, doc in enumerate(docs):
        doc.tensor = np.array(images[i])

當(dāng)你在一個(gè) Executor 中定義了一個(gè)方法，并使用了@requests 裝飾器，那么你就可以通過(guò)端點(diǎn)調(diào)用這個(gè)方法。當(dāng)你沒(méi)有指定端點(diǎn)路徑時(shí)，那么這個(gè)方法會(huì)成為默認(rèn)處理程序。

這里我們沒(méi)有像@requests(on='/foo') 這樣指定端點(diǎn)路徑，只是使用了裸的@requests，所以當(dāng)調(diào)用 Executor 時(shí)，generate_image() 方法會(huì)成為默認(rèn)處理程序。

部署我們的微服務(wù)

那么，現(xiàn)在我們有了 Executor，下一步當(dāng)然就是部署啦！通過(guò) Deployment，你可以運(yùn)行和擴(kuò)展 Executor，添加 replicas(副本), shards(分片) 和 dynamic batching(動(dòng)態(tài)批處理)。此外，部署到 Kubernetes 或 Docker Compose 也很容易，我們將在本文后面介紹。

1. Deploy via Python API

運(yùn)行 python deployment.py

from jina import Deployment
from text_to_image import TextToImage

with Deployment(uses=TextToImage, timeout_ready=-1, install_requirements=True) as dep:
    dep.block()

2. Deploy via YAML

使用 CLI 運(yùn)行 YAML 部署：jina deployment --uses deployment.yml

jtype: Deployment
with:
  port: 12345
  uses: TextToImage
  py_modules:
    - text_to_image.py # name of the module containing Executor
  timeout_ready: -1
install_requirements: True
And run the YAML Deployment with the CLI: jina deployment --uses deployment.yml

無(wú)論您運(yùn)行哪種 Deployment，您都會(huì)看到以下輸出：

────────────────────────────────────────  Deployment is ready to serve! ─────────────────────────────────────────
╭──────────────  Endpoint ───────────────╮
│        Protocol                   GRPC │
│          Local          0.0.0.0:12345  │
│        Private      172.28.0.12:12345  │
│         Public    35.230.97.208:12345  │
╰──────────────────────────────────────────╯

與我們的微服務(wù)通信

我們可以使用 Jina Client 通過(guò) gRPC 向我們的服務(wù)發(fā)送請(qǐng)求。如前所述，我們使用 Document 作為基本的 IO 格式:

運(yùn)行 client.py 獲得我們的夢(mèng)中精靈貓。

from docarray import Document
from jina import Client

image_text = Document(text='rainbow unicorn butterfly kitten')

client = Client(port=12345)  # use port from output above
response = client.post(on='/', inputs=[image_text])

response[0].display()

擴(kuò)展我們的微服務(wù)

Jina 有開(kāi)箱即用的可擴(kuò)展功能，如副本、分片和動(dòng)態(tài)批處理。這使您可以輕松增加應(yīng)用程序的吞吐量。

讓我們部署 Deployment，并使用副本和動(dòng)態(tài)批處理對(duì)其進(jìn)行擴(kuò)展。我們將:

創(chuàng)建兩個(gè)副本，每個(gè)副本分配一個(gè) GPU。

啟用動(dòng)態(tài)批處理以并行處理傳入同一模型的請(qǐng)求。

這是 Deployment 的原始（未擴(kuò)展）deployment.yml：

jtype: Deployment
with:
  timeout_ready: -1
  uses: jinaai://jina-ai/TextToImage
  install_requirements: true

現(xiàn)在讓我們擴(kuò)大規(guī)模：

jtype: Deployment
with:
  timeout_ready: -1
  uses: jinaai://jina-ai/TextToImage
  install_requirements: true
  env:
   CUDA_VISIBLE_DEVICES: RR
  replicas: 2
  uses_dynamic_batching: # configure dynamic batching
    /default:
      preferred_batch_size: 10
      timeout: 200

我們通過(guò) CUDA_VISIBLE_DEVICES添加了 GPU 支持，使用了兩個(gè)副本(每個(gè)副本分配一個(gè) GPU)和動(dòng)態(tài)批處理，可以累積并批處理請(qǐng)求，再發(fā)送到 Executor。

假設(shè)您的機(jī)器有兩個(gè) GPU，使用擴(kuò)展后的 Deployment YAML 會(huì)比普通部署獲得更高的吞吐量。

感謝 YAML 語(yǔ)法，你可以直接注入部署配置，不用修改 Executor 代碼。當(dāng)然了，所有這些也可以通過(guò) Python API 完成。

Kubernetes, Docker Compose andOpenTelemetry

使用 Kubernetes 和 Jina 很容易

jina export kubernetes deployment.yml ./my-k8s
kubectl apply -R -f my-k8s

同樣的，Docker Compose 也很容易

jina export docker-compose deployment.yml docker-compose.yml
docker-compose up

甚至，使用 OpenTelemetry 進(jìn)行 tracing(跟蹤) 和 monitoring (監(jiān)視) 也很簡(jiǎn)單。

from docarray import DocumentArray
from jina import Executor, requests

class Encoder(Executor):
    @requests
    def encode(self, docs: DocumentArray, **kwargs):
        with self.tracer.start_as_current_span(
            'encode', context=tracing_context
        ) as span:
            with self.monitor(
                'preprocessing_seconds', 'Time preprocessing the requests'
            ):
                docs.tensors = preprocessing(docs)
            with self.monitor(
                'model_inference_seconds', 'Time doing inference the requests'
            ):
                docs.embedding = model_inference(docs.tensors)

您可以集成 Jaeger 或任何其他分布式跟蹤工具，來(lái)收集和可視化請(qǐng)求級(jí)別和應(yīng)用級(jí)別的服務(wù)操作屬性。這有助于分析請(qǐng)求-響應(yīng)生命周期、應(yīng)用程序行為和性能。要使用 Grafana，你可以下載這個(gè) JSON 文件并導(dǎo)入 Grafana:

共享你的 Executor

使用 Executor Hub 共享你的 Executors 或使用公共/私有 Executors，幾乎不需要擔(dān)心依賴關(guān)系。

創(chuàng)建 Hub-ready Executor:

jina hub new

將其推送到 Executor Hub：

jina hub push

您可以通過(guò) Python 在 Deployment 中使用 Hub Executor：

Deployment(uses='jinaai://jina-ai/TextToImage', install_requirements=True)  # download and run locally
Deployment(uses='jinaai+docker://jina-ai/TextToImage')  # run in Docker container
Deployment(uses='jinaai+sandbox://jina-ai/TextToImage')  # run in hosted sandbox

或者 YAML：

uses: jinaai://jina-ai/TextToImage  # download and run locally
install_requirements: true

uses: jinaai+docker://jina-ai/TextToImage  # run in Docker container
uses: jinaai+sandbox://jina-ai/TextToImage  # run in hosted sandbox

Executor Hub 管理后端的所有內(nèi)容，包括：

云端自動(dòng)構(gòu)建；

高效且經(jīng)濟(jì)地存儲(chǔ)、部署和交付 Executor；

自動(dòng)解決版本沖突和依賴；

通過(guò) Sandbox 即時(shí)交付任何 Executor，而無(wú)需將任何內(nèi)容 pull 到本地；

將微服務(wù)串聯(lián)成pipeline 中

有時(shí)你可能希望將微服務(wù)串聯(lián)成一個(gè) pipeline。這就是 Flow 的用武之地。我們將在以后的博客中更深入地介紹如何構(gòu)建 Flow，目前您可以查看我們的 README。

Readme: get.jina.ai

總結(jié)

正如用 Executor 和 Deployment 包裝微服務(wù)或模型一樣，我們也必須總結(jié)這篇文章?？偨Y(jié)一下我們所涵蓋的內(nèi)容：

使用 Jina，你可以將模型包裝為 Executor，通過(guò) Deployment 可以直接部署這些 Executor，或者將他們串聯(lián)成 pipeline 作為 Flow 去部署。

Jina 與 Kubernetes、Docker Compose 和 OpenTelemetry 集成輕而易舉。

你可以在 Executor Hub 輕松找到和共享所有內(nèi)容。

如果您想繼續(xù)深入了解，請(qǐng)查看我們的文檔以獲取有關(guān) Executors 和 Deployments 的更多信息，或者使用 Jina Flows 查看 pipeline。您還可以聯(lián)系我們的 Slack 社區(qū)jina.ai/community。

多多和我們互動(dòng)吧！這樣我們才更加有動(dòng)力分享出更多好文章，未來(lái)我們將發(fā)布更多文章深入探討作為人人可用的多模態(tài)數(shù)據(jù)平臺(tái)，如何利用 Jina 地云原生，MLOps 和 LMOps 技術(shù)，讓每個(gè)企業(yè)和開(kāi)發(fā)者都能享受到最好的搜索和生成技術(shù)。

Jina 吉祥物之彩虹獨(dú)角獸蝴蝶小貓排行榜

像所有圖像生成一樣，我們花了很長(zhǎng)時(shí)間才生成一只完美的可愛(ài)小貓。





審核編輯：劉清