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

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create an Azure AI Content Safety __(AACS)__ enabled Inpainting online endpoint (Preview)\n",
    "### This notebook is under preview.\n",
    "\n",
    "### Steps to create an __AACS__ enabled __inpainting__ online endpoint\n",
    "1. Create an __AACS__ resource for moderating the request from user and response from the __inpainting__ online endpoint.\n",
    "2. Create a new __AACS__ enabled __inpainting__ online endpoint with a custom [score_online.py](./aacs-scoring-files/score/score_online.py) which will integrate with the __AACS__ resource to moderate the response from the __text-to-image__ inpainting model and the request from the user, but to make the custom [score_online.py](./aacs-scoring-files/score/score_online.py) to successfully authenticated to the __AACS__ resource, we have 2 options:\n",
    "    1. __UAI__, recommended but more complex approach, is to create a User Assigned Identity (UAI) and assign appropriate roles to the UAI. Then, the custom [score_online.py](./aacs-scoring-files/score/score_online.py) can obtain the access token of the UAI from the AAD server to access the AACS resource. Use [this notebook](aacs-prepare-uai.ipynb) to create UAI account for step 3 below\n",
    "    2. __Environment variable__, simpler but less secure approach, is to just pass the access key of the AACS resource to the custom [score_online.py](./aacs-scoring-files/score/score_online.py) via environment variable, then the custom [score_online.py](./aacs-scoring-files/score/score_online.py) can use the key directly to access the AACS resource, this option is less secure than the first option, if someone in your org has access to the endpoint, he/she can get the access key from the environment variable and use it to access the AACS resource.\n",
    "\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",
    "### Outline\n",
    "1. Setup pre-requisites\n",
    "2. Create AACS resource\n",
    "3. Pick a model to deploy\n",
    "4. Deploy the model to an online endpoint for real time inference\n",
    "5. Test the endpoint\n",
    "6. Clean up resources - delete the online endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setup pre-requisites\n",
    "* Check List\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": "markdown",
   "metadata": {},
   "source": [
    "> [x] The identity you are using to execute this notebook(yourself or your VM) need to have the __Contributor__ role on the resource group where the AML Workspace your specified is located, because this notebook will create an AACS resource using that identity."
   ]
  },
  {
   "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",
    "\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",
    "print(f\"Connected to workspace {workspace_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The models, fine tuning pipelines and environments are available in the AzureML system registry, \"azureml\"\n",
    "\n",
    "registry_name = \"azureml\"\n",
    "\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. Create AACS resource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1 Assign variables for Azure Content Safety\n",
    "Currently, AACS is available in a limited set of regions:\n",
    "\n",
    "\n",
    "__NOTE__: before you choose the region to deploy the AACS, please be aware that your data will be transferred to the region you choose and by selecting a region outside your current location, you may be allowing the transmission of your data to regions outside your jurisdiction. It is important to note that data protection and privacy laws may vary between jurisdictions. Before proceeding, we strongly advise you to familiarize yourself with the local laws and regulations governing data transfer and ensure that you are legally permitted to transmit your data to an overseas location for processing. By continuing with the selection of a different region, you acknowledge that you have understood and accepted any potential risks associated with such data transmission. Please proceed with caution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The severity level that will trigger response be blocked\n",
    "# Please reference Azure AI content documentation for more details\n",
    "# https://learn.microsoft.com/en-us/azure/cognitive-services/content-safety/concepts/harm-categories\n",
    "content_severity_threshold = \"2\"\n",
    "\n",
    "# UAI to be used for endpoint if you choose to use UAI as authentication method.\n",
    "# Use default name \"aacs-uai\" as used in prepare uai notebook\n",
    "\n",
    "# uai_name = \"aacs-uai\"\n",
    "\n",
    "# If you choose environment variables for authentication of AACS resource, then assign empty (\"\") value to uai_name\n",
    "uai_name = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient\n",
    "\n",
    "aacs_client = CognitiveServicesManagementClient(credential, subscription_id)\n",
    "\n",
    "\n",
    "# settings for the Azure AI Content Safety (AACS) resource\n",
    "# we will choose existing AACS resource if it exists, otherwise create a new one\n",
    "# name of AACS resource, has to be unique\n",
    "\n",
    "aacs_name = f\"aacs-inpainting-{str(uuid4())[:8]}\"\n",
    "available_aacs_locations = [\"east us\", \"west europe\"]\n",
    "\n",
    "# create a new Cognitive Services Account\n",
    "kind = \"ContentSafety\"\n",
    "aacs_sku_name = \"S0\"\n",
    "aacs_location = available_aacs_locations[0]\n",
    "\n",
    "\n",
    "print(\"Available SKUs:\")\n",
    "aacs_skus = aacs_client.resource_skus.list()\n",
    "print(\"SKU Name\\tSKU Tier\\tLocations\")\n",
    "for sku in aacs_skus:\n",
    "    if sku.kind == \"ContentSafety\":\n",
    "        locations = \",\".join(sku.locations)\n",
    "        print(sku.name + \"\\t\\t\" + sku.tier + \"\\t\\t\" + locations)\n",
    "\n",
    "print(f\"Choose a new AACS resource in {aacs_location} with SKU {aacs_sku_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2 Create AACS Resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.mgmt.cognitiveservices.models import Account, Sku, AccountProperties\n",
    "\n",
    "parameters = Account(\n",
    "    sku=Sku(name=aacs_sku_name),\n",
    "    kind=kind,\n",
    "    location=aacs_location,\n",
    "    properties=AccountProperties(\n",
    "        custom_sub_domain_name=aacs_name, public_network_access=\"Enabled\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "def find_acs(accounts):\n",
    "    return next(\n",
    "        x\n",
    "        for x in accounts\n",
    "        if x.kind == \"ContentSafety\"\n",
    "        and x.location == aacs_location\n",
    "        and x.sku.name == aacs_sku_name\n",
    "    )\n",
    "\n",
    "\n",
    "try:\n",
    "    # check if AACS exists\n",
    "    aacs = aacs_client.accounts.get(resource_group, aacs_name)\n",
    "    print(f\"Found existing AACS Account {aacs.name}.\")\n",
    "except:\n",
    "    try:\n",
    "        # check if there is an existing AACS resource within same resource group\n",
    "        aacs = find_acs(aacs_client.accounts.list_by_resource_group(resource_group))\n",
    "        print(\n",
    "            f\"Found existing AACS Account {aacs.name} in resource group {resource_group}.\"\n",
    "        )\n",
    "    except:\n",
    "        print(f\"Creating AACS Account {aacs_name}.\")\n",
    "        aacs_client.accounts.begin_create(resource_group, aacs_name, parameters).wait()\n",
    "        print(\"Resource created.\")\n",
    "        aacs = aacs_client.accounts.get(resource_group, aacs_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aacs_endpoint = aacs.properties.endpoint\n",
    "aacs_resource_id = aacs.id\n",
    "aacs_name = aacs.name\n",
    "print(\n",
    "    f\"AACS name is {aacs.name} .\\nUse this name in UAI preparation notebook to create UAI.\"\n",
    ")\n",
    "print(f\"AACS endpoint is {aacs_endpoint}\")\n",
    "print(f\"AACS ResourceId is {aacs_resource_id}\")\n",
    "\n",
    "aacs_access_key = aacs_client.accounts.list_keys(\n",
    "    resource_group_name=resource_group, account_name=aacs.name\n",
    ").key1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.3 Check if UAI is used (Required for using UAI authentication method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uai_id = \"\"\n",
    "uai_client_id = \"\"\n",
    "if uai_name != \"\":\n",
    "    from azure.mgmt.msi import ManagedServiceIdentityClient\n",
    "    from azure.mgmt.msi.models import Identity\n",
    "\n",
    "    try:\n",
    "        msi_client = ManagedServiceIdentityClient(\n",
    "            subscription_id=subscription_id,\n",
    "            credential=credential,\n",
    "        )\n",
    "        uai_resource = msi_client.user_assigned_identities.get(resource_group, uai_name)\n",
    "        uai_id = uai_resource.id\n",
    "        uai_client_id = uai_resource.client_id\n",
    "    except Exception as ex:\n",
    "        print(\"Please run aacs-prepare-uai.ipynb notebook and re-run the cell.\")\n",
    "        raise ex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. 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": [
    "### 4. 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": "markdown",
   "metadata": {},
   "source": [
    "Create an online endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Endpoint names need to be unique in a region,\n",
    "# hence using uuid (first 8 character) to create unique endpoint name\n",
    "\n",
    "endpoint_name = f\"safe-inpainting-{str(uuid4())[:8]}\"  # Replace with your endpoint name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import (\n",
    "    ManagedOnlineEndpoint,\n",
    "    IdentityConfiguration,\n",
    "    ManagedIdentityConfiguration,\n",
    ")\n",
    "\n",
    "# Check if the endpoint already exists in the workspace\n",
    "try:\n",
    "    endpoint = workspace_ml_client.online_endpoints.get(endpoint_name)\n",
    "    print(\"---Endpoint already exists---\")\n",
    "except:\n",
    "    # Create an online endpoint if it doesn't exist\n",
    "\n",
    "    # Define the endpoint\n",
    "    endpoint = ManagedOnlineEndpoint(\n",
    "        name=endpoint_name,\n",
    "        description=f\"Test endpoint for {model.name}\",\n",
    "        identity=IdentityConfiguration(\n",
    "            type=\"user_assigned\",\n",
    "            user_assigned_identities=[ManagedIdentityConfiguration(resource_id=uai_id)],\n",
    "        )\n",
    "        if uai_id != \"\"\n",
    "        else None,\n",
    "    )\n",
    "\n",
    "    # Trigger the endpoint creation\n",
    "    try:\n",
    "        workspace_ml_client.begin_create_or_update(endpoint).wait()\n",
    "        print(\"\\n---Endpoint created successfully---\\n\")\n",
    "    except Exception as err:\n",
    "        raise RuntimeError(\n",
    "            f\"Endpoint creation failed. Detailed Response:\\n{err}\"\n",
    "        ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 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": [
    "# Initialize deployment parameters\n",
    "\n",
    "deployment_name = \"inpainting-deploy\"\n",
    "sku_name = \"STANDARD_NC6S_V3\"  # Name of the sku(instance type). Check the model card in catalog to get the most optimal sku for model.\n",
    "\n",
    "REQUEST_TIMEOUT_MS = 90000\n",
    "\n",
    "acs_env_vars = {\n",
    "    \"CONTENT_SAFETY_ACCOUNT_NAME\": aacs_name,\n",
    "    \"CONTENT_SAFETY_ENDPOINT\": aacs_endpoint,\n",
    "    \"CONTENT_SAFETY_KEY\": aacs_access_key if uai_client_id == \"\" else None,\n",
    "    \"CONTENT_SAFETY_THRESHOLD\": content_severity_threshold,\n",
    "    \"SUBSCRIPTION_ID\": subscription_id,\n",
    "    \"RESOURCE_GROUP_NAME\": resource_group,\n",
    "    \"UAI_CLIENT_ID\": uai_client_id,\n",
    "}\n",
    "\n",
    "MAX_CONCURRENT_REQUESTS = (\n",
    "    2  # the maximum number of concurrent requests supported by the endpoint\n",
    ")\n",
    "\n",
    "fm_container_default_env_vars = {\n",
    "    \"WORKER_COUNT\": MAX_CONCURRENT_REQUESTS,\n",
    "}\n",
    "\n",
    "deployment_env_vars = {**fm_container_default_env_vars, **acs_env_vars}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import (\n",
    "    OnlineRequestSettings,\n",
    "    ManagedOnlineDeployment,\n",
    "    ProbeSettings,\n",
    ")\n",
    "\n",
    "deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=endpoint_name,\n",
    "    model=model.id,\n",
    "    instance_type=sku_name,\n",
    "    instance_count=1,\n",
    "    environment_variables=deployment_env_vars,\n",
    "    request_settings=OnlineRequestSettings(request_timeout_ms=REQUEST_TIMEOUT_MS),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=30,\n",
    "        success_threshold=1,\n",
    "        period=100,\n",
    "        initial_delay=500,\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Trigger the deployment creation\n",
    "try:\n",
    "    workspace_ml_client.begin_create_or_update(deployment).wait()\n",
    "    print(\"\\n---Deployment created successfully---\\n\")\n",
    "except Exception as err:\n",
    "    raise RuntimeError(\n",
    "        f\"Deployment creation failed. Detailed Response:\\n{err}\"\n",
    "    ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. 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",
    "\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": "markdown",
   "metadata": {},
   "source": [
    "#### 5.1.1 Sample input for safe prompt."
   ]
  },
  {
   "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(\n",
    "                    \"utf-8\"\n",
    "                ),\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": [
    "# Invoke the endpoint\n",
    "\n",
    "response = workspace_ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=endpoint.name,\n",
    "    deployment_name=deployment.name,\n",
    "    request_file=request_file_name,\n",
    ")\n",
    "\n",
    "# Visualize the model output\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": "markdown",
   "metadata": {},
   "source": [
    "#### 5.1.2 Sample input for safe prompt."
   ]
  },
  {
   "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(\n",
    "                    \"utf-8\"\n",
    "                ),\n",
    "                \"prompt\": \"A yellow cat, high resolution, sitting on a park bench\",\n",
    "                \"negative_prompt\": \"blurry; cartoonish\",\n",
    "            }\n",
    "        ],\n",
    "        \"parameters\": {\"num_inference_steps\": 50, \"guidance_scale\": 7.5},\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=endpoint.name,\n",
    "    deployment_name=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.2 Sample input for un-safe prompt"
   ]
  },
  {
   "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):\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(\n",
    "                    \"utf-8\"\n",
    "                ),\n",
    "                \"prompt\": \"a dog with severed leg and bleeding profusely from deep laceration to the lower extremities, exposing tissues\",\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": [
    "# Invoke the endpoint\n",
    "\n",
    "response = workspace_ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=endpoint.name,\n",
    "    deployment_name=deployment.name,\n",
    "    request_file=request_file_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model response should be empty because it is blocked by the Azure AI Content Safety (AACS) service.\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. 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=endpoint.name).wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK v2",
   "language": "python",
   "name": "python310-sdkv2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}