wehub-resource-sync
457 lines
15 KiB
Plaintext
457 lines
15 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "597c13c0",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Serving a Stable Diffusion Model with Ray Serve\n",
|
|
"\n",
|
|
"| Template Specification | Description |\n",
|
|
"| ---------------------- | ----------- |\n",
|
|
"| Summary | This template loads a pretrained stable diffusion model from HuggingFace and serves it to a local endpoint as a [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) deployment. |\n",
|
|
"| Time to Run | Around 2 minutes to setup the models and generate your first image(s). Less than 10 seconds for every subsequent round of image generation (depending on the image size). |\n",
|
|
"| Minimum Compute Requirements | At least 1 GPU node. The default is 4 nodes, each with 1 NVIDIA T4 GPU. |\n",
|
|
"| Cluster Environment | This template uses a custom docker image built on top of the Anyscale-provided Ray image using Python 3.9: [`anyscale/ray:latest-py39-cu118`](https://docs.anyscale.com/reference/base-images/overview). See the appendix in the `README` for more details. |\n",
|
|
"\n",
|
|
"By the end, we'll have an application that generates images using stable diffusion for a given prompt!\n",
|
|
"\n",
|
|
"The application will look something like this:\n",
|
|
"\n",
|
|
"```text\n",
|
|
"Enter a prompt (or 'q' to quit): twin peaks sf in basquiat painting style\n",
|
|
"\n",
|
|
"Generating image(s)...\n",
|
|
"\n",
|
|
"Generated 4 image(s) in 8.75 seconds to the directory: 58b298d9\n",
|
|
"```\n",
|
|
"\n",
|
|
"\n",
|
|
"> Slot in your code below wherever you see the ✂️ icon to build off of this template!\n",
|
|
">\n",
|
|
"> The framework and data format used in this template can be easily replaced to suit your own application!\n",
|
|
"\n",
|
|
"We'll start with some imports and initialize Ray:"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "59842da3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from fastapi import FastAPI\n",
|
|
"from fastapi.responses import Response\n",
|
|
"from io import BytesIO\n",
|
|
"import matplotlib.pyplot as plt\n",
|
|
"import numpy as np\n",
|
|
"import os\n",
|
|
"import requests\n",
|
|
"import time\n",
|
|
"import uuid\n",
|
|
"\n",
|
|
"import ray\n",
|
|
"from ray import serve\n",
|
|
"\n",
|
|
"ray.init()\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "520ef4d7",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Deploy the Ray Serve application locally\n",
|
|
"\n",
|
|
"First, we define the Ray Serve application with the model loading and inference logic. This includes setting up:\n",
|
|
"- The `/imagine` API endpoint that we query to generate the image.\n",
|
|
"- The stable diffusion model loaded inside a Ray Serve Deployment.\n",
|
|
" We'll specify the *number of model replicas* to keep active in our Ray cluster. These model replicas can process incoming requests concurrently.\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "de6318ac",
|
|
"metadata": {},
|
|
"source": [
|
|
"> ✂️ Replace these values to change the number of model replicas to serve, as well as the GPU resources required by each replica.\n",
|
|
">\n",
|
|
"> With more model replicas, more images can be generated in parallel!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "90eca147",
|
|
"metadata": {
|
|
"tags": []
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"NUM_REPLICAS: int = 4\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "be41ca9e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"if NUM_REPLICAS > ray.available_resources()[\"GPU\"]:\n",
|
|
" print(\n",
|
|
" \"Your cluster does not currently have enough resources to run with these settings. \"\n",
|
|
" \"Consider decreasing the number of workers, or decreasing the resources needed \"\n",
|
|
" \"per worker. Ignore this if your cluster auto-scales.\"\n",
|
|
" )\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "89eb3e2c",
|
|
"metadata": {},
|
|
"source": [
|
|
"First, we define the Ray Serve Deployment, which will load a stable diffusion model and perform inference with it.\n",
|
|
"\n",
|
|
"> ✂️ Modify this block to load your own model, and change the `generate` method to perform your own online inference logic!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f203efd4",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"@serve.deployment(\n",
|
|
" ray_actor_options={\"num_gpus\": 1},\n",
|
|
" num_replicas=NUM_REPLICAS,\n",
|
|
")\n",
|
|
"class StableDiffusionV2:\n",
|
|
" def __init__(self):\n",
|
|
" # <Replace with your own model loading logic>\n",
|
|
" import torch\n",
|
|
" from diffusers import EulerDiscreteScheduler, StableDiffusionPipeline\n",
|
|
"\n",
|
|
" model_id = \"stabilityai/stable-diffusion-2\"\n",
|
|
" scheduler = EulerDiscreteScheduler.from_pretrained(\n",
|
|
" model_id, subfolder=\"scheduler\"\n",
|
|
" )\n",
|
|
" self.pipe = StableDiffusionPipeline.from_pretrained(\n",
|
|
" model_id, scheduler=scheduler, revision=\"fp16\", torch_dtype=torch.float16\n",
|
|
" )\n",
|
|
" self.pipe = self.pipe.to(\"cuda\")\n",
|
|
"\n",
|
|
" def generate(self, prompt: str, img_size: int = 776):\n",
|
|
" # <Replace with your own model inference logic>\n",
|
|
" assert len(prompt), \"prompt parameter cannot be empty\"\n",
|
|
" image = self.pipe(prompt, height=img_size, width=img_size).images[0]\n",
|
|
" return image\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "0134aa54",
|
|
"metadata": {},
|
|
"source": [
|
|
"Next, we'll define the actual API endpoint to live at `/imagine`.\n",
|
|
"\n",
|
|
"> ✂️ Modify this block to change the endpoint URL, response schema, and add any post-processing logic needed from your model output!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6f80fee2",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"app = FastAPI()\n",
|
|
"\n",
|
|
"\n",
|
|
"@serve.deployment(num_replicas=1)\n",
|
|
"@serve.ingress(app)\n",
|
|
"class APIIngress:\n",
|
|
" def __init__(self, diffusion_model_handle) -> None:\n",
|
|
" self.handle = diffusion_model_handle\n",
|
|
"\n",
|
|
" @app.get(\n",
|
|
" \"/imagine\",\n",
|
|
" responses={200: {\"content\": {\"image/png\": {}}}},\n",
|
|
" response_class=Response,\n",
|
|
" )\n",
|
|
" async def generate(self, prompt: str, img_size: int = 776):\n",
|
|
" assert len(prompt), \"prompt parameter cannot be empty\"\n",
|
|
"\n",
|
|
" image = await self.handle.generate.remote(prompt, img_size=img_size)\n",
|
|
"\n",
|
|
" file_stream = BytesIO()\n",
|
|
" image.save(file_stream, \"PNG\")\n",
|
|
" return Response(content=file_stream.getvalue(), media_type=\"image/png\")\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "61b8916d",
|
|
"metadata": {},
|
|
"source": [
|
|
"Now, we deploy the Ray Serve application locally at `http://localhost:8000`!"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "dfc2e244",
|
|
"metadata": {
|
|
"scrolled": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"entrypoint = APIIngress.bind(StableDiffusionV2.bind())\n",
|
|
"\n",
|
|
"# Shutdown any existing Serve replicas, if they're still around.\n",
|
|
"serve.shutdown()\n",
|
|
"serve.run(entrypoint, name=\"serving_stable_diffusion_template\")\n",
|
|
"print(\"Done setting up replicas! Now accepting requests...\")\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "757678cc",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Make requests to the endpoint\n",
|
|
"\n",
|
|
"Next, we'll build a simple client to submit prompts as HTTP requests to the local endpoint at `http://localhost:8000/imagine`."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "008976b5",
|
|
"metadata": {},
|
|
"source": [
|
|
"Start the client script in the next few cells, and generate your first image!\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "67ad095b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"endpoint = \"http://localhost:8000/imagine\"\n",
|
|
"\n",
|
|
"\n",
|
|
"@ray.remote(num_cpus=0)\n",
|
|
"def generate_image(prompt, image_size):\n",
|
|
" req = {\"prompt\": prompt, \"img_size\": image_size}\n",
|
|
" resp = requests.get(endpoint, params=req)\n",
|
|
" return resp.content\n",
|
|
"\n",
|
|
"\n",
|
|
"def show_images(filenames):\n",
|
|
" fig, axs = plt.subplots(1, len(filenames), figsize=(4 * len(filenames), 4))\n",
|
|
" for i, filename in enumerate(filenames):\n",
|
|
" ax = axs if len(filenames) == 1 else axs[i]\n",
|
|
" ax.imshow(plt.imread(filename))\n",
|
|
" ax.axis(\"off\")\n",
|
|
" plt.show()\n",
|
|
"\n",
|
|
"\n",
|
|
"def main(\n",
|
|
" interactive: bool = False,\n",
|
|
" prompt: str = \"twin peaks sf in basquiat painting style\",\n",
|
|
" num_images: int = 4,\n",
|
|
" image_size: int = 640,\n",
|
|
"):\n",
|
|
" try:\n",
|
|
" requests.get(endpoint, timeout=0.1)\n",
|
|
" except Exception as e:\n",
|
|
" raise RuntimeWarning(\n",
|
|
" \"Did you setup the Ray Serve model replicas with `serve.run` \"\n",
|
|
" \"in a previous cell?\"\n",
|
|
" ) from e\n",
|
|
"\n",
|
|
" generation_times = []\n",
|
|
" while True:\n",
|
|
" prompt = (\n",
|
|
" prompt\n",
|
|
" if not interactive\n",
|
|
" else input(f\"\\nEnter a prompt (or 'q' to quit): \")\n",
|
|
" )\n",
|
|
" if prompt.lower() == \"q\":\n",
|
|
" break\n",
|
|
"\n",
|
|
" print(\"\\nGenerating image(s)...\\n\")\n",
|
|
" start = time.time()\n",
|
|
"\n",
|
|
" # Make `num_images` requests to the endpoint at once!\n",
|
|
" images = ray.get(\n",
|
|
" [generate_image.remote(prompt, image_size) for _ in range(num_images)]\n",
|
|
" )\n",
|
|
"\n",
|
|
" dirname = f\"{uuid.uuid4().hex[:8]}\"\n",
|
|
" os.makedirs(dirname)\n",
|
|
" filenames = []\n",
|
|
" for i, image in enumerate(images):\n",
|
|
" filename = os.path.join(dirname, f\"{i}.png\")\n",
|
|
" with open(filename, \"wb\") as f:\n",
|
|
" f.write(image)\n",
|
|
" filenames.append(filename)\n",
|
|
"\n",
|
|
" elapsed = time.time() - start\n",
|
|
" generation_times.append(elapsed)\n",
|
|
" print(\n",
|
|
" f\"\\nGenerated {len(images)} image(s) in {elapsed:.2f} seconds to \"\n",
|
|
" f\"the directory: {dirname}\\n\"\n",
|
|
" )\n",
|
|
" show_images(filenames)\n",
|
|
" if not interactive:\n",
|
|
" break\n",
|
|
" return np.mean(generation_times) if generation_times else -1\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "c8949cc7",
|
|
"metadata": {},
|
|
"source": [
|
|
"Once the stable diffusion model finishes generating your image(s), it will be included in the HTTP response body.\n",
|
|
"The client saves all the images in a local directory for you to view, and they'll also show up in the notebook cell!"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "3e29193b",
|
|
"metadata": {},
|
|
"source": [
|
|
"> ✂️ Replace this value to change the number of images to generate per prompt.\n",
|
|
">\n",
|
|
"> Each image will be generated starting from a different set of random noise,\n",
|
|
"> so you'll be able to see multiple options per prompt!\n",
|
|
">\n",
|
|
"> Try starting with `NUM_IMAGES_PER_PROMPT` equal to `NUM_REPLICAS` from earlier.\n",
|
|
">\n",
|
|
"> You can choose to run this interactively, or submit a single `PROMPT`."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "dd20a52d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"NUM_IMAGES_PER_PROMPT: int = NUM_REPLICAS\n",
|
|
"\n",
|
|
"# Control the output size: (IMAGE_SIZE, IMAGE_SIZE)\n",
|
|
"# The stable diffusion model requires `IMAGE_SIZE` to be a multiple of 8.\n",
|
|
"# NOTE: Generated image quality degrades rapidly if you reduce the size too much.\n",
|
|
"IMAGE_SIZE: int = 640\n",
|
|
"\n",
|
|
"INTERACTIVE: bool = False\n",
|
|
"PROMPT = \"twin peaks sf in basquiat painting style\"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"mean_generation_time = main(\n",
|
|
" interactive=INTERACTIVE,\n",
|
|
" prompt=PROMPT,\n",
|
|
" num_images=NUM_IMAGES_PER_PROMPT,\n",
|
|
" image_size=IMAGE_SIZE,\n",
|
|
")\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "fb124968",
|
|
"metadata": {},
|
|
"source": [
|
|
"You've successfully served a stable diffusion model!\n",
|
|
"You can modify this template and iterate your model deployment directly on your cluster within your Anyscale Workspace,\n",
|
|
"testing with the local endpoint."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9e360cf9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Shut down the model replicas once you're done!\n",
|
|
"serve.shutdown()\n"
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "880c2d6f",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Summary\n",
|
|
"\n",
|
|
"This template used [Ray Serve](https://docs.ray.io/en/latest/serve/index.html) to serve many replicas of a stable diffusion model. \n",
|
|
"\n",
|
|
"At a high level, this template showed how to:\n",
|
|
"1. Define a Ray Serve deployment to load a HuggingFace model and perform inference.\n",
|
|
"2. Set up a local endpoint to accept and route requests to the different model replicas.\n",
|
|
"3. Make multiple requests in parallel to generate many images at a time.\n",
|
|
"\n",
|
|
"See this [getting started guide](https://docs.ray.io/en/latest/serve/getting_started.html) for a more detailed walkthrough of Ray Serve."
|
|
]
|
|
},
|
|
{
|
|
"attachments": {},
|
|
"cell_type": "markdown",
|
|
"id": "bcc69b2d",
|
|
"metadata": {},
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "ray_dev_py38",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.8.13"
|
|
},
|
|
"vscode": {
|
|
"interpreter": {
|
|
"hash": "265d195fda5292fe8f69c6e37c435a5634a1ed3b6799724e66a975f68fa21517"
|
|
}
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|