sdk/python/foundation-models/system/inference/text-to-image/text-to-image-online-endpoint.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text to Image Inference using Online Endpoints\n",
    "\n",
    "This sample shows how to deploy `text-to-image` type stable diffusion models to an online endpoint for inference.\n",
    "\n",
    "### Task\n",
    "`text-to-image` tasks generates image as output based on text prompt given in input.\n",
    " \n",
    "### Model\n",
    "Models that can perform the `text-to-image` task are tagged with `text-to-image`. We will use the `runwayml-stable-diffusion-v1-5` model in this notebook. If you opened this notebook from a specific model card, remember to replace the specific model name.\n",
    "\n",
    "\n",
    "### Outline\n",
    "1. Setup pre-requisites\n",
    "2. Pick a model to deploy\n",
    "3. Deploy the model to an online endpoint for real time inference\n",
    "4. Test the endpoint using sample text prompts\n",
    "5. Clean up resources - delete the online endpoint"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setup pre-requisites\n",
    "* Install dependencies\n",
    "* Connect to AzureML Workspace. Learn more at [set up SDK authentication](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication?tabs=sdk). Replace  `<WORKSPACE_NAME>`, `<RESOURCE_GROUP>` and `<SUBSCRIPTION_ID>` below.\n",
    "* Connect to `azureml` system registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the required packages\n",
    "%pip install azure-identity==1.13.0\n",
    "%pip install azure-mgmt-cognitiveservices==13.4.0\n",
    "%pip install azure-ai-ml==1.11.1\n",
    "%pip install azure-mgmt-msi==7.0.0\n",
    "%pip install azure-mgmt-authorization==3.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "import time\n",
    "\n",
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    credential = InteractiveBrowserCredential()\n",
    "\n",
    "try:\n",
    "    workspace_ml_client = MLClient.from_config(credential)\n",
    "    subscription_id = workspace_ml_client.subscription_id\n",
    "    resource_group = workspace_ml_client.resource_group_name\n",
    "    workspace_name = workspace_ml_client.workspace_name\n",
    "except Exception as ex:\n",
    "    print(ex)\n",
    "    # Enter details of your AML workspace\n",
    "    subscription_id = \"<SUBSCRIPTION_ID>\"\n",
    "    resource_group = \"<RESOURCE_GROUP>\"\n",
    "    workspace_name = \"<AML_WORKSPACE_NAME>\"\n",
    "workspace_ml_client = MLClient(\n",
    "    credential, subscription_id, resource_group, workspace_name\n",
    ")\n",
    "\n",
    "# The models, fine tuning pipelines and environments are available in the AzureML system registry, \"azureml\"\n",
    "registry_ml_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id,\n",
    "    resource_group,\n",
    "    registry_name=\"azureml\",\n",
    ")\n",
    "# Generating a unique timestamp that can be used for names and versions that need to be unique\n",
    "timestamp = str(int(time.time()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Pick a model to deploy\n",
    "\n",
    "Browse models in the Model Catalog in the AzureML Studio, filtering by the `text-to-image` task. In this example, we use the `runwayml-stable-diffusion-v1-5` model. If you have opened this notebook for a different model, replace the model name accordingly. This is a pre-trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"runwayml-stable-diffusion-v1-5\"\n",
    "diffusion_model = registry_ml_client.models.get(name=model_name, label=\"latest\")\n",
    "print(\n",
    "    f\"\\n\\nUsing model name: {diffusion_model.name}, version: {diffusion_model.version}, id: {diffusion_model.id} for generating images from text.\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Deploy the model to an online endpoint for real time inference\n",
    "Online endpoints give a durable REST API that can be used to integrate with applications that need to use the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time, sys\n",
    "from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment\n",
    "\n",
    "# Endpoint names need to be unique in a region, hence using timestamp to create unique endpoint name\n",
    "timestamp = int(time.time())\n",
    "online_endpoint_name = \"text-to-image-\" + str(timestamp)\n",
    "# Create an online endpoint\n",
    "endpoint = ManagedOnlineEndpoint(\n",
    "    name=online_endpoint_name,\n",
    "    description=\"Online endpoint for \"\n",
    "    + diffusion_model.name\n",
    "    + \", for text-to-image task\",\n",
    "    auth_mode=\"key\",\n",
    ")\n",
    "workspace_ml_client.begin_create_or_update(endpoint).wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1 Setup Deployment Parameters\n",
    "\n",
    "We utilize an optimized __foundation-model-inference__ container for model scoring. This container is designed to deliver high throughput and low latency using <a href=\"https://github.com/microsoft/DeepSpeed-MII\" target=\"_blank\">  Deepspeed-mii </a>. In this section, we introduce several environment variables that can be adjusted to customize a deployment for either high throughput or low latency scenarios.\n",
    "\n",
    "- __WORKER_COUNT__: The number of workers to use for inferencing. This is used as a proxy for the number of concurrent requests that the server should handle.\n",
    "- __TENSOR_PARALLEL__: The number of GPUs to use for tensor parallelism.\n",
    "- __NUM_REPLICAS__: The number of model instances to load for the deployment. This is used to increase throughput by loading multiple models on multiple GPUs, if the model is small enough to fit.\n",
    "\n",
    "`NUM_REPLICAS` and `TENSOR_PARALLEL` work hand-in-hand to determine the most optimal configuration to increase the throughput for the deployment without degrading too much on the latency. The total number of GPUs used for inference will be `NUM_REPLICAS` * `TENSOR_PARALLEL`. For example, if `NUM_REPLICAS` = 2 and `TENSOR_PARALLEL` = 2, then 4 GPUs will be used for inference. Ensure that the model you are deploying is small enough to fit on the number of GPUs you are using, specified by `TENSOR_PARALLEL`. For instance, if there are 4 GPUs available, and `TENSOR_PARALLEL` = 2, then the model must be small enough to fit on 2 GPUs. If the model is too large, then the deployment will fail.\n",
    "\n",
    "For stable diffusion model, the scoring script uses default `TENSOR_PARALLEL` = 1 and `NUM_REPLICAS` = number of GPUs in SKU for optimal balance of latency and throughput."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_CONCURRENT_REQUESTS = (\n",
    "    2  # the maximum number of concurrent requests supported by the endpoint\n",
    ")\n",
    "REQUEST_TIMEOUT_MS = 90000  # the timeout for each request in milliseconds\n",
    "\n",
    "deployment_env_vars = {\n",
    "    \"WORKER_COUNT\": MAX_CONCURRENT_REQUESTS,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Note__: We have set the value of `MAX_CONCURRENT_REQUESTS` to 2, as we are utilizing the `Standard_NC6s_v3` SKU for deployment, which has one GPU. If you are using a larger SKU, please increase this value to get the maximum performance.\n",
    "\n",
    "For model `stabilityai-stable-diffusion-xl-base-1-0`, set the value of `MAX_CONCURRENT_REQUESTS` to 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.2 Deploy the model\n",
    "This step may take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import OnlineRequestSettings, ProbeSettings\n",
    "\n",
    "deployment_name = \"text-to-image-mlflow-deploy\"\n",
    "sku_name = \"STANDARD_NC6S_V3\"  # Use GPU instance type like Standard_NC6s_v3 or above\n",
    "\n",
    "# Create a deployment\n",
    "demo_deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    model=diffusion_model.id,\n",
    "    instance_type=sku_name,\n",
    "    instance_count=1,\n",
    "    environment_variables=deployment_env_vars,\n",
    "    request_settings=OnlineRequestSettings(\n",
    "        max_concurrent_requests_per_instance=MAX_CONCURRENT_REQUESTS,\n",
    "        request_timeout_ms=REQUEST_TIMEOUT_MS,\n",
    "    ),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=49,\n",
    "        success_threshold=1,\n",
    "        timeout=299,\n",
    "        period=180,\n",
    "        initial_delay=180,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=10,\n",
    "        success_threshold=1,\n",
    "        timeout=10,\n",
    "        period=10,\n",
    "        initial_delay=10,\n",
    "    ),\n",
    ")\n",
    "workspace_ml_client.online_deployments.begin_create_or_update(demo_deployment).wait()\n",
    "endpoint.traffic = {deployment_name: 100}\n",
    "workspace_ml_client.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Test the endpoint\n",
    "\n",
    "We will submit request to online endpoint with some sample text prompts for image generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_deployment = workspace_ml_client.online_deployments.get(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=online_endpoint_name,\n",
    ")\n",
    "\n",
    "# Get the details for online endpoint\n",
    "endpoint = workspace_ml_client.online_endpoints.get(name=online_endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "request_json = {\n",
    "    \"input_data\": {\n",
    "        \"columns\": [\"prompt\"],\n",
    "        \"data\": [\"a photograph of an astronaut riding a horse\"],\n",
    "        \"index\": [0],\n",
    "        \"parameters\": {\"num_inference_steps\": 50, \"guidance_scale\": 7.5},\n",
    "    }\n",
    "}\n",
    "\n",
    "# Create request json\n",
    "request_file_name = \"sample_request_data.json\"\n",
    "with open(request_file_name, \"w\") as request_file:\n",
    "    json.dump(request_json, request_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> The `parameters` are optional inputs that support `num_inference_steps` and `guidance_scale`. If you need support for new parameters, please file a support ticket.\n",
    "\n",
    "For model `stabilityai-stable-diffusion-xl-base-1-0`, the input `parameters` is not supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = workspace_ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    deployment_name=demo_deployment.name,\n",
    "    request_file=request_file_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import base64\n",
    "from PIL import Image\n",
    "\n",
    "generations = json.loads(response)\n",
    "for generation in generations:\n",
    "    img = Image.open(io.BytesIO(base64.b64decode(generation[\"generated_image\"])))\n",
    "    display(img)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Clean up resources - delete the online endpoint\n",
    "Don't forget to delete the online endpoint, else you will leave the billing meter running for the compute used by the endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace_ml_client.online_endpoints.begin_delete(name=online_endpoint_name).wait()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "ipython"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}