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

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image to Image Inference using Online Endpoints\n",
    "\n",
    "This sample shows how to deploy `image to image` type stable diffusion models to an online endpoint for inference.\n",
    "\n",
    "### Task\n",
    "`image to image` task takes an original image and a text prompt as input. The model generates image by modifying the original image.\n",
    "\n",
    " \n",
    "### Model\n",
    "Models that can perform the `image to image` task are tagged with `image-to-image`. We will use the `stabilityai-stable-diffusion-xl-refiner-1-0` 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 and original 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_name = \"azureml\"\n",
    "registry_ml_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id,\n",
    "    resource_group,\n",
    "    registry_name=registry_name,\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 `image-to-image` task. In this example, we use the `stabilityai-stable-diffusion-xl-refiner-1-0` 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 image to image model to be deployed\n",
    "model_name = \"stabilityai-stable-diffusion-xl-refiner-1-0\"\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 image and 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",
    "    \"image-to-image-\" + 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 image to image 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 = \"image-to-image-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",
    "The sample of input schema for image to image task:\n",
    "```json\n",
    "{\n",
    "   \"input_data\": {\n",
    "        \"columns\": [\"prompt\", \"image\"],\n",
    "        \"data\": [\n",
    "            {\n",
    "                \"prompt\": \"sample prompt\",\n",
    "                \"image\": \"base image1\"\n",
    "            },\n",
    "            {\n",
    "                \"prompt\": \"sample prompt\",\n",
    "                \"image\": \"base image2\"            }\n",
    "        ],\n",
    "        \"index\": [0, 1]\n",
    "    }\n",
    "}\n",
    "```\n",
    "> - The base image (1, 2) strings should be in base64 format or publicly accessible urls.\n",
    "\n",
    "The sample of output schema for image to image task:\n",
    "```json\n",
    "[\n",
    "    {\n",
    "        \"generated_image\": \"image1\",\n",
    "        \"nsfw_content_detected\": None\n",
    "    },\n",
    "    {\n",
    "        \"generated_image\": \"image2\",\n",
    "        \"nsfw_content_detected\": None\n",
    "    }\n",
    "]\n",
    "```\n",
    "> - \"nsfw_content_detected\" is not supported for this model.\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",
    "\n",
    "request_json = {\n",
    "    \"input_data\": {\n",
    "        \"columns\": [\"image\", \"prompt\"],\n",
    "        \"index\": [0],\n",
    "        \"data\": [\n",
    "            {\n",
    "                \"image\": base64.encodebytes(read_image(base_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",
    ")"
   ]
  },
  {
   "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",
    "    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
}