{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DALL-E\n",
    "\n",
    "This notebook gives examples of using each of the endpoints in the DALL-E API Private Preview\n",
    "\n",
    "Currently these consist of:\n",
    "- **Generation:** An API that allows you to generate an image or images with a supplied caption\n",
    "- **Edits:** An API that alows you to edit or extend an existing image\n",
    "- **Variations:** An API that generates variations of a supplied image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Here we perform the following steps:\n",
    "- Import the packages we need\n",
    "- Import OpenAI API key\n",
    "- Set urls for DALL-E endpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import requests\n",
    "import os\n",
    "from PIL import Image as PILImg\n",
    "\n",
    "openai.api_key = os.environ.get('OPENAI_API_KEY')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the urls for our API endpoints\n",
    "\n",
    "dalle_generation_url = 'https://api.openai.com/v1/images/generations'\n",
    "dalle_variations_url = 'https://api.openai.com/v1/images/variations'\n",
    "dalle_edits_url = 'https://api.openai.com/v1/images/edits'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a directory to write our dalle images to\n",
    "\n",
    "image_dir = os.path.join(os.curdir,'images')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation\n",
    "\n",
    "This is the base endpoint for DALLE and allows you to generate an image with a given caption\n",
    "\n",
    "**Required:**\n",
    "- caption: str describing the desired image\n",
    "\n",
    "**Optional:**\n",
    "- model: str ID of the model to use.\n",
    "- num_images: 1-10 generated images\n",
    "- size: The size of image to generate 64x64, 256x256 or 1024x1024. Smaller images are faster to generate.\n",
    "- response_format:\"url\" or b64_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "generation_headers = {\n",
    "            \"Authorization\": f\"Bearer {openai.api_key}\",\n",
    "            \"Content-Type\": \"application/json\"\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_images(generation_headers,caption,model=\"image-alpha-001\",num_images=1,size=\"1024x1024\",response_format='url'):\n",
    "    generation_json = {\n",
    "                          \"model\": model,\n",
    "                          \"caption\": caption,\n",
    "                          \"num_images\": num_images,\n",
    "                          \"size\":size , \n",
    "                          \"response_format\": response_format\n",
    "                        }\n",
    "    \n",
    "    response = requests.post(dalle_generation_url,headers=generation_headers,json=generation_json)\n",
    "    \n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_gen_response = generate_images(generation_headers,\"A monkey hacker dreaming of beautiful bananas, cyberpunk style\").json()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_image = requests.get(image_gen_response['data'][0]['url']).content\n",
    "with open(os.path.join(image_dir,'image_name.png'), 'wb') as handler:\n",
    "    handler.write(generated_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2277412 characters]\n",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PILImg.open(os.path.join(image_dir,'image_name.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variations\n",
    "\n",
    "Here we use the image generated in the previous step to generate a bunch of variations using the DALLE API\n",
    "\n",
    "**Required:**\n",
    "- image: The image for variations in binary format\n",
    "\n",
    "**Optional:**\n",
    "- model: str ID of the model to use.\n",
    "- num_images: 1-10 generated variations of the original image\n",
    "- size: The size of image to generate 64x64, 256x256 or 1024x1024. Smaller images are faster to generate.\n",
    "- response_format: url, b64_json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "variation_headers = {\n",
    "            \"Authorization\": f\"Bearer {openai.api_key}\"\n",
    "        }\n",
    "\n",
    "image_path = os.path.join(image_dir,'image_name.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_variations(headers\n",
    "                        ,image_path\n",
    "                        ,model=\"image-alpha-001\"\n",
    "                        ,size=\"1024x1024\"\n",
    "                        ,num_images=5\n",
    "                        ,response_format=\"url\"):\n",
    "    variation_files = {\n",
    "    \"image\": open(image_path, 'rb'),\n",
    "    }\n",
    "\n",
    "    variation_body = {\n",
    "        \"model\": model,\n",
    "        \"size\": size,\n",
    "        \"num_images\": num_images,\n",
    "        \"response_format\": response_format\n",
    "    }\n",
    "    variations_response = requests.post(dalle_variations_url,headers=headers,data=variation_body,files=variation_files)\n",
    "    \n",
    "    return variations_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "variations_response = generate_variations(headers=variation_headers,image_path=image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_variation_list = variations_response.json()['data']\n",
    "counter = 0\n",
    "for image_variation in image_variation_list:\n",
    "    counter += 1\n",
    "    url_image = requests.get(image_variation['url']).content\n",
    "    with open(os.path.join(image_dir,f'edited_image_name_{counter}.png'), 'wb') as handler:\n",
    "        handler.write(url_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2455008 characters]\n",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PILImg.open(os.path.join(image_dir,'edited_image_name_2.png'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Edits\n",
    "\n",
    "The edit function allows you to have DALL-E generate new edits of selected parts of an existing image. This is done using a \"mask\" to set the part of the original image that you want edited, along with the entire image and a caption\n",
    "\n",
    "**Required params:** \n",
    "- caption: str describing the desired image.\n",
    "- image: The image for editing in binary format\n",
    "- mask: The mask for editing as a JSON serialized encoded numpy array.\n",
    "\n",
    "**Optional params:**\n",
    "- model: str ID of the model to use.\n",
    "- num_images: 1-10 generated images\n",
    "- size: The size of image to generate 64x64, 256x256 or 1024x1024. Smaller images are faster to generate.\n",
    "- Response_format: url, b64_json\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Edit Area\n",
    "\n",
    "To perform an edit you need a \"mask\", which is a JSON serialised numpy array representation of your original image with pixels that you want to be edited replaced with 0.\n",
    "\n",
    "To do this we'll follow these steps:\n",
    "- Read in our image\n",
    "- Create a masked_image zeroing out the lower half of the image to be replaced with an edit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = PILImg.open(os.path.join(image_dir,\"image_name.png\"))\n",
    "img = img.convert('RGBA')\n",
    "width, height = img.size\n",
    "\n",
    "for x in range(width):\n",
    "    for y in range(height//2, height):\n",
    "        # Get the current RGBA color\n",
    "        r, g, b, a = img.getpixel((x, y))\n",
    "        \n",
    "        # Make the pixel transparent by setting 'a' (alpha) to zero\n",
    "        img.putpixel((x, y), (r, g, b, 0))\n",
    "\n",
    "img.save(os.path.join(image_dir,'image_name_transparent.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2655104 characters]\n",
      "text/plain": [
       "<PIL.Image.Image image mode=RGBA size=1024x1024>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform Edit\n",
    "\n",
    "Now we supply our image, caption and mask to the API to get 5 examples of edits to our image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_headers = {\n",
    "            \"Authorization\": f\"Bearer {openai.api_key}\"\n",
    "    #,\"Content-Type\": 'application/json'#\"multipart/form-data\"\n",
    "        }\n",
    "\n",
    "edit_image_path = os.path.join(image_dir,'image_name.png')\n",
    "edit_mask_path = os.path.join(image_dir,'image_name_transparent.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_edits(headers\n",
    "                        ,caption\n",
    "                        ,image_path\n",
    "                        ,mask_path\n",
    "                        ,model=\"image-alpha-001\"\n",
    "                        ,size=\"1024x1024\"\n",
    "                        ,num_images=5\n",
    "                        ,response_format=\"url\"):\n",
    "    edit_files = {\n",
    "    \"image\": open(edit_image_path, 'rb'),\n",
    "    \"mask\": open(edit_mask_path, 'rb'),\n",
    "    }\n",
    "\n",
    "    edit_body = {\n",
    "        \"caption\": caption,\n",
    "        \"model\": model,\n",
    "        \"size\": size,\n",
    "        \"num_images\": num_images,\n",
    "        \"response_format\": response_format\n",
    "    }\n",
    "    edit_response = requests.post(dalle_edits_url,headers=headers,data=edit_body,files=edit_files)\n",
    "    \n",
    "    return edit_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "edit_response = generate_edits(headers=edit_headers,caption='A field of computers made of bananas, modern art',image_path=edit_image_path,mask_path=edit_mask_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_edit_list = edit_response.json()['data']\n",
    "counter = 0\n",
    "for image_edit in image_edit_list:\n",
    "    counter += 1\n",
    "    url_image = requests.get(image_edit['url']).content\n",
    "    with open(os.path.join(image_dir,f'image_name_edit_{counter}.png'), 'wb') as handler:\n",
    "        handler.write(url_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2236992 characters]\n",
      "text/plain": [
       "<PIL.PngImagePlugin.PngImageFile image mode=RGB size=1024x1024>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "PILImg.open(os.path.join(image_dir,f'image_name_edit_3.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
