{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1bbe66b-3784-4da2-a5ca-5406a4d0c843",
   "metadata": {},
   "source": [
    "# Generate Dalle Images With Conversable Agents"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa341ad5-32f0-4e4b-8869-e19f19471d81",
   "metadata": {},
   "source": [
    "This notebook illustrates how to add the image generation capability to a conversable agent. \n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "Some extra dependencies are needed for this notebook, which can be installed via pip:\n",
    "\n",
    "```bash\n",
    "pip install autogen-agentchat[lmm]~=0.2\n",
    "```\n",
    "\n",
    "For more information, please refer to the [installation guide](/docs/installation/).\n",
    ":::\n",
    "````\n",
    "First, let's import all the required modules to run this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3951ffc-ab4e-4bf6-bc36-f0d7f84c5a40",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "from typing import Dict, Optional\n",
    "\n",
    "from IPython.display import display\n",
    "from PIL.Image import Image\n",
    "\n",
    "import autogen\n",
    "from autogen.agentchat.contrib import img_utils\n",
    "from autogen.agentchat.contrib.capabilities import generate_images\n",
    "from autogen.cache import Cache\n",
    "from autogen.oai import openai_utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "156c6db3-fb95-4ff4-803b-81c105a095e2",
   "metadata": {},
   "source": [
    "Let's define our LLM configs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "104db7e8-e30a-4012-916d-5f1622e5da02",
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt_config = {\n",
    "    \"config_list\": [{\"model\": \"gpt-4-turbo-preview\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}],\n",
    "    \"timeout\": 120,\n",
    "    \"temperature\": 0.7,\n",
    "}\n",
    "gpt_vision_config = {\n",
    "    \"config_list\": [{\"model\": \"gpt-4-vision-preview\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}],\n",
    "    \"timeout\": 120,\n",
    "    \"temperature\": 0.7,\n",
    "}\n",
    "dalle_config = {\n",
    "    \"config_list\": [{\"model\": \"dall-e-3\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}],\n",
    "    \"timeout\": 120,\n",
    "    \"temperature\": 0.7,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57cec727-c237-4205-8681-e8a792f0a242",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](/docs/topics/llm_configuration).\n",
    ":::\n",
    "````\n",
    "\n",
    "Our system will consist of 2 main agents:\n",
    "1. Image generator agent.\n",
    "2. Critic agent.\n",
    "\n",
    "The image generator agent will carry a conversation with the critic, and generate images based on the critic's requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "132e57ea-a041-4c4d-99ca-b0616505b6be",
   "metadata": {},
   "outputs": [],
   "source": [
    "CRITIC_SYSTEM_MESSAGE = \"\"\"You need to improve the prompt of the figures you saw.\n",
    "How to create an image that is better in terms of color, shape, text (clarity), and other things.\n",
    "Reply with the following format:\n",
    "\n",
    "CRITICS: the image needs to improve...\n",
    "PROMPT: here is the updated prompt!\n",
    "\n",
    "If you have no critique or a prompt, just say TERMINATE\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "427233a0-190b-4acd-80a3-c5951609b337",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _is_termination_message(msg) -> bool:\n",
    "    # Detects if we should terminate the conversation\n",
    "    if isinstance(msg.get(\"content\"), str):\n",
    "        return msg[\"content\"].rstrip().endswith(\"TERMINATE\")\n",
    "    elif isinstance(msg.get(\"content\"), list):\n",
    "        for content in msg[\"content\"]:\n",
    "            if isinstance(content, dict) and \"text\" in content:\n",
    "                return content[\"text\"].rstrip().endswith(\"TERMINATE\")\n",
    "    return False\n",
    "\n",
    "\n",
    "def critic_agent() -> autogen.ConversableAgent:\n",
    "    return autogen.ConversableAgent(\n",
    "        name=\"critic\",\n",
    "        llm_config=gpt_vision_config,\n",
    "        system_message=CRITIC_SYSTEM_MESSAGE,\n",
    "        max_consecutive_auto_reply=3,\n",
    "        human_input_mode=\"NEVER\",\n",
    "        is_termination_msg=lambda msg: _is_termination_message(msg),\n",
    "    )\n",
    "\n",
    "\n",
    "def image_generator_agent() -> autogen.ConversableAgent:\n",
    "    # Create the agent\n",
    "    agent = autogen.ConversableAgent(\n",
    "        name=\"dalle\",\n",
    "        llm_config=gpt_vision_config,\n",
    "        max_consecutive_auto_reply=3,\n",
    "        human_input_mode=\"NEVER\",\n",
    "        is_termination_msg=lambda msg: _is_termination_message(msg),\n",
    "    )\n",
    "\n",
    "    # Add image generation ability to the agent\n",
    "    dalle_gen = generate_images.DalleImageGenerator(llm_config=dalle_config)\n",
    "    image_gen_capability = generate_images.ImageGeneration(\n",
    "        image_generator=dalle_gen, text_analyzer_llm_config=gpt_config\n",
    "    )\n",
    "\n",
    "    image_gen_capability.add_to_agent(agent)\n",
    "    return agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38f65910-b811-43f9-b6a0-986b52c83d94",
   "metadata": {},
   "source": [
    "We'll define `extract_img` to help us extract the image generated by the image generator agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ff882a30-0f1f-4c6a-8951-9e11df48bc26",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_images(sender: autogen.ConversableAgent, recipient: autogen.ConversableAgent) -> Image:\n",
    "    images = []\n",
    "    all_messages = sender.chat_messages[recipient]\n",
    "\n",
    "    for message in reversed(all_messages):\n",
    "        # The GPT-4V format, where the content is an array of data\n",
    "        contents = message.get(\"content\", [])\n",
    "        for content in contents:\n",
    "            if isinstance(content, str):\n",
    "                continue\n",
    "            if content.get(\"type\", \"\") == \"image_url\":\n",
    "                img_data = content[\"image_url\"][\"url\"]\n",
    "                images.append(img_utils.get_pil_image(img_data))\n",
    "\n",
    "    if not images:\n",
    "        raise ValueError(\"No image data found in messages.\")\n",
    "\n",
    "    return images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d804f33b-3fb7-4756-ada8-01da371994bb",
   "metadata": {},
   "source": [
    "Start the converstion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ed01470b-ad59-483d-990e-772d27e0764a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mdalle\u001b[0m (to critic):\n",
      "\n",
      "A happy dog wearing a shirt saying 'I Love AutoGen'. Make sure the text is clear.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcritic\u001b[0m (to dalle):\n",
      "\n",
      "CRITICS: the image needs to improve the contrast and size of the text to enhance its clarity, and the shirt's color should not clash with the dog's fur color to maintain a harmonious color scheme.\n",
      "\n",
      "PROMPT: here is the updated prompt!\n",
      "Create an image of a joyful dog with a coat of a contrasting color to its fur, wearing a shirt with bold, large text saying 'I Love AutoGen' for clear readability.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mdalle\u001b[0m (to critic):\n",
      "\n",
      "I generated an image with the prompt: Joyful dog, contrasting coat color to its fur, shirt with bold, large text \"I Love AutoGen\" for clear readability.<image>\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[33mcritic\u001b[0m (to dalle):\n",
      "\n",
      "CRITICS: the image effectively showcases a joyful dog with a contrasting shirt color, and the text 'I Love AutoGen' is large and bold, ensuring clear readability.\n",
      "\n",
      "PROMPT: TERMINATE\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "dalle = image_generator_agent()\n",
    "critic = critic_agent()\n",
    "\n",
    "img_prompt = \"A happy dog wearing a shirt saying 'I Love AutoGen'. Make sure the text is clear.\"\n",
    "# img_prompt = \"Ask me how I'm doing\"\n",
    "\n",
    "result = dalle.initiate_chat(critic, message=img_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85846d75-e954-48ba-ab84-9a1a5e1f9204",
   "metadata": {},
   "source": [
    "Let's display all the images that was generated by Dalle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "dfa623e5-9444-4868-bdaf-ccae85d3ecd5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 19000 characters]",
      "image/png": "[encoded data, 210160 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=300x300>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "images = extract_images(dalle, critic)\n",
    "\n",
    "for image in reversed(images):\n",
    "    display(image.resize((300, 300)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99994e50-543e-425b-803c-2a7263f16a36",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Generate images with conversable agents.",
   "tags": [
    "capability",
    "multimodal"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
