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

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inpainting Inference using Online Endpoints\n",
    "\n",
    "This sample shows how to deploy `inpainting` type stable diffusion models to an online endpoint for inference.\n",
    "\n",
    "### Task\n",
    "`inpainting` task takes an original image, a text prompt and a mask image as input. The model generates inpainted image by modifying the original image.\n",
    "\n",
    " \n",
    "### Model\n",
    "Models that can perform the `inpainting` task are tagged with `text-to-image`. We will use the `runwayml-stable-diffusion-inpainting` 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 prompt, original image and mask image.\n",
    "5. Clean up resources - delete the online endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setup pre-requisites\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": [
    "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",
    ")"
   ]
  },
  {
   "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-inpainting` 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": [
    "# Name of the inpainting model to be deployed\n",
    "model_name = \"runwayml-stable-diffusion-inpainting\"\n",
    "\n",
    "try:\n",
    "    model = registry_ml_client.models.get(name=model_name, label=\"latest\")\n",
    "except Exception as ex:\n",
    "    print(\n",
    "        f\"No model named {model_name} found in registry. \"\n",
    "        \"Please check model name present in Azure model catalog\"\n",
    "    )\n",
    "    raise ex\n",
    "\n",
    "print(\n",
    "    f\"\\n\\nUsing model name: {model.name}, version: {model.version}, id: {model.id} for generating images from text.\"\n",
    ")"
   ]
  },
  {
   "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 uuid\n",
    "from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment\n",
    "\n",
    "# Endpoint names need to be unique in a region, hence using uuid (first 8 character) to create unique endpoint name\n",
    "online_endpoint_name = (\n",
    "    \"inpainting-\" + str(uuid.uuid4())[:8]\n",
    ")  # Replace with your endpoint name\n",
    "# Create an online endpoint\n",
    "endpoint = ManagedOnlineEndpoint(\n",
    "    name=online_endpoint_name,\n",
    "    description=\"Online endpoint for \" + model.name + \", for inpainting task\",\n",
    "    auth_mode=\"key\",\n",
    ")\n",
    "workspace_ml_client.begin_create_or_update(endpoint).wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import OnlineRequestSettings, ProbeSettings\n",
    "\n",
    "deployment_name = \"inpainting-deploy\"\n",
    "\n",
    "# Create a deployment\n",
    "demo_deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    model=model.id,\n",
    "    instance_type=\"Standard_NC6s_v3\",  # Use GPU instance type like Standard_NC6s_v3 or above\n",
    "    instance_count=1,\n",
    "    request_settings=OnlineRequestSettings(\n",
    "        max_concurrent_requests_per_instance=1,\n",
    "        request_timeout_ms=90000,\n",
    "        max_queue_wait_ms=500,\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()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Test the endpoint\n",
    "\n",
    "We will fetch some sample data from the test dataset and submit to online endpoint for inference.\n",
    "\n",
    "### Supported Parameters\n",
    "\n",
    "- negative_prompt: The prompt to guide what to not include in image generation. Ignored when not using guidance (`guidance_scale < 1`).\n",
    "- num_inference_steps: The number of de-noising steps. More de-noising steps usually lead to a higher quality image at the expense of slower inference, defaults to 50.\n",
    "- guidance_scale: A higher guidance scale value encourages the model to generate images closely linked to the text `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`, defaults to 7.5.\n",
    "\n",
    "> These `parameters` are optional inputs. If you need support for new parameters, please file a support ticket.\n",
    "\n",
    "The sample of input schema for inpainting task:\n",
    "```json\n",
    "{\n",
    "   \"input_data\": {\n",
    "        \"columns\": [\"prompt\", \"image\", \"mask_image\", \"negative_prompt\"],\n",
    "        \"data\": [\n",
    "            {\n",
    "                \"prompt\": \"Face of a yellow cat, high resolution, sitting on a park bench\",\n",
    "                \"image\": \"image1\",\n",
    "                \"mask_image\": \"mask1\",\n",
    "                \"negative_prompt\": \"blurry; cartoonish\"\n",
    "            },\n",
    "            {\n",
    "                \"prompt\": \"Face of a green cat, high resolution, sitting on a park bench\",\n",
    "                \"image\": \"image2\",\n",
    "                \"mask_image\": \"mask2\",\n",
    "                \"negative_prompt\": \"blurry; cartoonish\"\n",
    "            }\n",
    "        ],\n",
    "        \"index\": [0, 1],\n",
    "        \"parameters\": {\n",
    "            \"num_inference_steps\": 50,\n",
    "            \"guidance_scale\": 7.5\n",
    "        }\n",
    "    }\n",
    "}\n",
    "```\n",
    "> - The base and mask images (1 and 2) strings should be in base64 format or publicly accessible urls.\n",
    "> - The mask structure is white for inpainting and black for keeping as is\n",
    "\n",
    "The sample of output schema for inpainting task:\n",
    "```json\n",
    "[\n",
    "    {\n",
    "        \"generated_image\": \"image1\",\n",
    "        \"nsfw_content_detected\": False\n",
    "    },\n",
    "    {\n",
    "        \"generated_image\": \"image2\",\n",
    "        \"nsfw_content_detected\": True\n",
    "    }\n",
    "]\n",
    "```\n",
    "> - If \"nsfw_content_detected\" is True then generated image will be totally black.\n",
    "> - Generated images \"image1\" and \"image2\" strings are in base64 format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create request json\n",
    "import base64\n",
    "import json\n",
    "\n",
    "\n",
    "def read_image(image_path: str) -> bytes:\n",
    "    \"\"\"Reads an image from a file path into a byte array.\"\"\"\n",
    "    with open(image_path, \"rb\") as f:\n",
    "        return f.read()\n",
    "\n",
    "\n",
    "base_image = \"inpainting_data/images/dog_on_bench.png\"\n",
    "mask_image = \"inpainting_data/masks/dog_on_bench.png\"\n",
    "\n",
    "request_json = {\n",
    "    \"input_data\": {\n",
    "        \"columns\": [\"image\", \"mask_image\", \"prompt\"],\n",
    "        \"index\": [0],\n",
    "        \"data\": [\n",
    "            {\n",
    "                \"image\": base64.encodebytes(read_image(base_image)).decode(\"utf-8\"),\n",
    "                \"mask_image\": base64.encodebytes(read_image(mask_image)).decode(\"utf-8\"),\n",
    "                \"prompt\": \"A yellow cat, high resolution, sitting on a park bench\",\n",
    "            }\n",
    "        ],\n",
    "    }\n",
    "}\n",
    "\n",
    "request_file_name = \"sample_request_data.json\"\n",
    "\n",
    "with open(request_file_name, \"w\") as request_file:\n",
    "    json.dump(request_json, request_file)"
   ]
  },
  {
   "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",
    ")\n",
    "\n",
    "import io\n",
    "import base64\n",
    "from PIL import Image\n",
    "\n",
    "generations = json.loads(response)\n",
    "for generation in generations:\n",
    "    print(f\"nsfw content detected: \", generation[\"nsfw_content_detected\"])\n",
    "    img = Image.open(io.BytesIO(base64.b64decode(generation[\"generated_image\"])))\n",
    "    display(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "request_json = {\n",
    "    \"input_data\": {\n",
    "        \"columns\": [\"image\", \"mask_image\", \"prompt\", \"negative_prompt\"],\n",
    "        \"data\": [\n",
    "            {\n",
    "                \"image\": base64.encodebytes(read_image(base_image)).decode(\"utf-8\"),\n",
    "                \"mask_image\": base64.encodebytes(read_image(mask_image)).decode(\"utf-8\"),\n",
    "                \"prompt\": \"A yellow cat, high resolution, sitting on a park bench\",\n",
    "                \"negative_prompt\": \"blurry; cartoonish\",\n",
    "            }\n",
    "        ],\n",
    "        \"parameters\": {\n",
    "            \"num_inference_steps\": 50,\n",
    "            \"guidance_scale\": 7.5\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "request_file_name = \"sample_request_data.json\"\n",
    "\n",
    "with open(request_file_name, \"w\") as request_file:\n",
    "    json.dump(request_json, request_file)"
   ]
  },
  {
   "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",
    ")\n",
    "\n",
    "generations = json.loads(response)\n",
    "for generation in generations:\n",
    "    print(f\"nsfw content detected: \", generation[\"nsfw_content_detected\"])\n",
    "    img = Image.open(io.BytesIO(base64.b64decode(generation[\"generated_image\"])))\n",
    "    display(img)"
   ]
  },
  {
   "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
}