{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7d7e337d-eb21-49df-a0cb-38c060313760",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🛠️ Setup & Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ea55396-8690-413c-8dda-6aec0d7b5087",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:20:34.558622Z",
     "shell.execute_reply.started": "2025-07-21T17:19:05.891214Z",
     "to_execute": "2025-07-21T17:19:06.731Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Notice: ffmpeg is not installed. torchaudio is used to load audio\n",
      "If you want to use ffmpeg backend to load audio, please install it by:\n",
      "\tsudo apt install ffmpeg # ubuntu\n",
      "\t# brew install ffmpeg # mac\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:19:16,300 - root - INFO - New registry table added: preprocessor_classes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key Conformer already exists in model_classes, re-register\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:19:17,119 - root - INFO - New registry table added: adaptor_classes\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Key Linear already exists in adaptor_classes, re-register\n",
      "Key TransformerDecoder already exists in decoder_classes, re-register\n",
      "Key LightweightConvolutionTransformerDecoder already exists in decoder_classes, re-register\n",
      "Key LightweightConvolution2DTransformerDecoder already exists in decoder_classes, re-register\n",
      "Key DynamicConvolutionTransformerDecoder already exists in decoder_classes, re-register\n",
      "Key DynamicConvolution2DTransformerDecoder already exists in decoder_classes, re-register\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:19:17,632 - root - INFO - New registry table added: lid_predictor_classes\n",
      "[2025-07-22 01:19:22,829] [INFO] [add_hooks.py:71:__init__] If you are working on AIStudio, please set 'ANTMONITOR_TFEVENT_PATH' env.\n",
      "You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.\n",
      "BailingMoeForCausalLM has generative capabilities, as `prepare_inputs_for_generation` is explicitly overwritten. However, it doesn't directly inherit from `GenerationMixin`. From 👉v4.50👈 onwards, `PreTrainedModel` will NOT inherit from `GenerationMixin`, and this model will lose the ability to call `generate` and other related functions.\n",
      "  - If you're using `trust_remote_code=True`, you can get rid of this warning by loading the model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes\n",
      "  - If you are the owner of the model architecture code, please modify your model class such that it inherits from `GenerationMixin` (after `PreTrainedModel`, otherwise you'll get an exception).\n",
      "  - If you are not the owner of the model architecture class, please contact the model code owner to update it.\n",
      "Loading checkpoint shards: 100%|██████████| 8/8 [00:02<00:00,  3.11it/s]\n",
      "2025-07-22 01:19:44,336 - xformers - WARNING - WARNING[XFORMERS]: xFormers can't load C++/CUDA extensions. xFormers was built for:\n",
      "    PyTorch 2.1.0+cu121 with CUDA 1201 (you have 2.4.0+cu124)\n",
      "    Python  3.10.13 (you have 3.10.15)\n",
      "  Please reinstall xformers (see https://github.com/facebookresearch/xformers#installing-xformers)\n",
      "  Memory-efficient attention, SwiGLU, sparse and more won't be available.\n",
      "  Set XFORMERS_MORE_DETAILS=1 for more details\n",
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:20<00:00,  5.02s/it]\n",
      "2025-07-22 01:20:07,421 - diffusion.sd3_loss - INFO - number of all Diffusion parameters: 2425834176, trainable: 2425834176\n",
      "2025-07-22 01:20:07,425 - diffusion.sd3_loss - INFO - number of all Diffusion parameters: 2425834176, trainable: 2425834176\n",
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:10<00:00,  5.08s/it]\n"
     ]
    }
   ],
   "source": [
    "# %%capturesdsd\n",
    "\n",
    "import os\n",
    "import torch\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from transformers import AutoProcessor, GenerationConfig\n",
    "from modeling_bailingmm import BailingMMNativeForConditionalGeneration\n",
    "\n",
    "# Load pre-trained model with optimized settings, this will take ~10 minutes\n",
    "model = BailingMMNativeForConditionalGeneration.from_pretrained(\n",
    "    \"inclusionAI/Ming-Lite-Omni\",\n",
    "    torch_dtype=torch.bfloat16,  # Use bfloat16 for memory efficiency\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    load_image_gen=True,\n",
    "    low_cpu_mem_usage=True       # Minimize CPU memory during loading\n",
    ").to(\"cuda\").to(torch.bfloat16)  # Run on GPU\n",
    "model.talker.use_vllm = False\n",
    "\n",
    "# Initialize processor for handling multimodal inputs\n",
    "processor = AutoProcessor.from_pretrained(\".\", trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9822ca9-6b30-49d6-94c2-ea72b2048c13",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### ⚙️ Inference Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "17cbb9e4-200a-4d42-aa0d-ac58abf537bc",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:20:51.167550Z",
     "shell.execute_reply.started": "2025-07-21T17:20:51.162710Z",
     "to_execute": "2025-07-21T17:20:51.961Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def generate(messages):\n",
    "    # 1. Format inputs using chat template\n",
    "    text = processor.apply_chat_template(messages, add_generation_prompt=True)\n",
    "\n",
    "    # 2. Extract vision/audio data\n",
    "    image_inputs, video_inputs, audio_inputs = processor.process_vision_info(messages)\n",
    "\n",
    "    # 3. Prepare tensor inputs\n",
    "    inputs = processor(\n",
    "        text=[text],\n",
    "        images=image_inputs,\n",
    "        videos=video_inputs,\n",
    "        audios=audio_inputs,\n",
    "        return_tensors=\"pt\",\n",
    "    )\n",
    "    inputs = inputs.to(model.device)\n",
    "    for k in inputs.keys():\n",
    "        if k == \"pixel_values\" or k == \"pixel_values_videos\" or k == \"audio_feats\":\n",
    "            inputs[k] = inputs[k].to(dtype=torch.bfloat16)\n",
    "\n",
    "    # 4. Configure generation\n",
    "    generation_config = GenerationConfig.from_dict({'no_repeat_ngram_size': 10})\n",
    "    generated_ids = model.generate(\n",
    "        **inputs,\n",
    "        max_new_tokens=8192 * 2,\n",
    "        use_cache=True,\n",
    "        eos_token_id=processor.gen_terminator,\n",
    "        generation_config=generation_config,\n",
    "    )\n",
    "    generated_ids_trimmed = [\n",
    "            out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)\n",
    "        ]\n",
    "\n",
    "    # 5. Decode output\n",
    "    output_text = processor.batch_decode(\n",
    "        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False\n",
    "    )[0]\n",
    "    \n",
    "    return output_text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3c9313-b13d-403d-befa-9f233db11222",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 💬 Text QA Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eacea5e8-5089-48a6-a4da-3efa5c5fedcc",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:21:34.808375Z",
     "shell.execute_reply.started": "2025-07-21T17:21:11.069901Z",
     "to_execute": "2025-07-21T17:21:11.919Z"
    },
    "isLargeOutputDisplay": false,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` model input instead.\n",
      "2025-07-22 01:21:11,349 - modeling_bailing_moe - WARNING - The input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in torch.bfloat16.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "请详细介绍鹦鹉的生活习性。\n",
      "\n",
      "Answer:\n",
      "鹦鹉是一种非常受欢迎的宠物鸟，以其鲜艳的羽毛、聪明和社交性强的特点而闻名。以下是鹦鹉的一些生活习性：\n",
      "\n",
      "### 1. **栖息环境**\n",
      "鹦鹉通常生活在热带和亚热带地区，包括雨林、草原和灌木丛。它们喜欢温暖湿润的环境，但也能够适应一些较为干燥的气候。\n",
      "\n",
      "### 2. **饮食**\n",
      "鹦鹉是杂食性动物，主要以种子、果实、花朵、坚果和昆虫为食。不同种类的鹦鹉有不同的饮食偏好。例如，亚马逊鹦鹉喜欢坚果和水果，而长尾小鹦鹉则更喜欢吃种子和蔬菜。\n",
      "\n",
      "### 3. **社交行为**\n",
      "鹦鹉是非常社交的动物，通常生活在群体中。它们通过复杂的叫声和肢体语言进行交流。鹦鹉之间的互动非常丰富，包括互相理毛、玩耍和合作觅食。\n",
      "\n",
      "### 4. **繁殖**\n",
      "鹦鹉的繁殖季节因种类而异，但通常在春季和夏季。雌性鹦鹉通常会在树洞或岩石缝隙中筑巢，产下2到8枚卵。孵化期约为18到30天，幼鸟在孵化后需要父母的照顾和喂养。\n",
      "\n",
      "### 5. **寿命**\n",
      "鹦鹉的寿命因种类而异，一般在10到50年之间。大型鹦鹉如金刚鹦鹉可以活到70年以上，而小型鹦鹉如长尾小鹦鹉的寿命通常在10到20年之间。\n",
      "\n",
      "### 6. **智力**\n",
      "鹦鹉是非常聪明的鸟类，具有高度的认知能力和学习能力。它们能够模仿人类的声音，甚至学会简单的词汇和短语。鹦鹉还表现出解决问题的能力和社交技巧。\n",
      "\n",
      "### 7. **活动时间**\n",
      "鹦鹉通常在白天活动，尤其是在早晨和傍晚。它们喜欢在阳光下活动，但也需要适当的遮荫和休息。\n",
      "\n",
      "### 8. **飞行能力**\n",
      "鹦鹉是强壮的飞行者，能够飞行相当长的距离。它们在飞行中表现出灵活性和耐力，能够在空中进行复杂的动作。\n",
      "\n",
      "### 9. **领地性**\n",
      "鹦鹉具有一定的领地性，尤其是在繁殖季节。它们会通过叫声和姿态来宣示领地，保护自己的巢穴和食物来源。\n",
      "\n",
      "### 10. **适应性**\n",
      "鹦鹉具有很强的适应性，能够适应不同的生活环境。许多鹦鹉在野外和人类饲养的环境中都能茁壮成长。\n",
      "\n",
      "总的来说，鹦鹉是一种非常有趣和聪明的鸟类，它们的生活习性丰富多彩，适合那些有足够时间和精力照顾它们的饲养者。\n"
     ]
    }
   ],
   "source": [
    "# qa\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"请详细介绍鹦鹉的生活习性。\"}\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "from IPython.utils import io\n",
    "with io.capture_output() as captured:\n",
    "    output_text = generate(messages)\n",
    "\n",
    "print(\"请详细介绍鹦鹉的生活习性。\\n\")\n",
    "print(\"Answer:\")\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f91dbd07-2bb4-475a-8a90-d56af7c17385",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🖼️ Image QA Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "518dbd5e-6b91-4c36-acb9-6cf34b3cfd10",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:21:44.176325Z",
     "shell.execute_reply.started": "2025-07-21T17:21:42.143040Z",
     "to_execute": "2025-07-21T17:21:42.907Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 11636 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "What kind of flower is this?\n",
      "\n",
      "Answer: This is a lotus flower.\n"
     ]
    }
   ],
   "source": [
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"  # Use to suppress warnings\n",
    "\n",
    "# image qa\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"image\", \"image\": \"figures/cases/flower.jpg\"},\n",
    "            {\"type\": \"text\", \"text\": \"What kind of flower is this?\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "from IPython.utils import io\n",
    "with io.capture_output() as captured:  # used to suppress output inside the function\n",
    "    output_text = generate(messages)\n",
    "\n",
    "from IPython.display import display, Image\n",
    "display(Image(\"figures/cases/flower.jpg\"))\n",
    "print(\"What kind of flower is this?\\n\")\n",
    "print(f\"Answer: {output_text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec1e0b09-64cd-49be-b89a-aed133c8fc60",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🤔 Chain-of-Thought Reasoning\n",
    "**To enable thinking before response, adding the following system prompt before your question:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6f99aef8-06bc-47d3-8dfa-d9600fdd1dbc",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:23:11.724356Z",
     "shell.execute_reply.started": "2025-07-21T17:22:22.663322Z",
     "to_execute": "2025-07-21T17:22:23.485Z"
    },
    "isLargeOutputDisplay": false,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7280 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input: Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end.\n",
      "Question: Find $m\\angle H$\n",
      "Choices:\n",
      "(A) 97\n",
      "(B) 102\n",
      "(C) 107\n",
      "(D) 122\n",
      "\n",
      "Answer:\n",
      "\n",
      "<thinking>\n",
      "Okay, so I'm trying to figure out the measure of angle H in this geometry problem. The image shows a rectangle ABCD with points M1, M2, M3, and M4 inside it. The diagonals of the rectangle intersect at point M, which is the midpoint of both diagonals. The image also shows that M1, M2, M3, M4 are connected to form a quadrilateral inside the rectangle.\n",
      "\n",
      "The problem mentions that the diagonals of the rectangle are equal in length and bisect each other. This means that the diagonals AC and BD intersect at point M, which is the midpoint. The image shows that M1, M2, M4, and M3 are connected to form a quadrilateral inside the original rectangle.\n",
      "\n",
      "The key information given is that the sum of the angles around point M is 360 degrees. The image shows that angle M1 is 102 degrees, and angle M4 is 102 degrees. Since the diagonals bisect each other, the angles formed at point M are related.\n",
      "\n",
      "The image also shows that angle M1 and angle M4 are each 102 degrees. Since the diagonally opposite angles in a quadrilateral formed by connecting the midpoints of a rectangle are equal, angle M2 and angle M3 are also 102 degrees each.\n",
      "\n",
      "The sum of the angles around point M is 4 * 102 = 408 degrees. However, the sum of the angles around a point should be 360 degrees. This discrepancy suggests that the angles M1, M2, M3, or M4 might not be the only angles contributing to the sum.\n",
      "\n",
      "Given that the diagonals bisect each other, the quadrilateral formed by M1, M2, M3, & M4 is a parallelogram. In a parallelogram, opposite angles are equal. Therefore, angle M1 = angle M3 and angle M2 = angle M4.\n",
      "\n",
      "Since angle M1 and angle M4 are each given as 102 degrees, angle M2 and angle M3 are each 102 degrees as well. The sum of the angles around point M is 102 + 102 + 102 = 306 degrees, which is less than 360 degrees. This indicates that there are additional angles contributing to the sum.\n",
      "\n",
      "The problem states that the sum of the angles around point M, including angle H, is 360 degrees. Therefore, angle H must be the remaining angle to make the total sum 360 degrees.\n",
      "\n",
      "So, angle H = 360 - (102 + 102 + x), where x is the measure of the other two angles. Since the diagonals bisect each other and the angles are equal, the other two angles must also be 102 degrees each.\n",
      "\n",
      "Thus, angle H = 360 - 306 = 54 degrees. However, this does not match any of the given options.\n",
      "\n",
      "Revisiting the problem, the key is that the sum of the angles around point H is 360 degrees, and the angles M1, M2, M4, M3 are each 102 degrees. Therefore, angle H must be the angle that completes the sum to 360 degrees.\n",
      "\n",
      "So angle H = 360 - (3 * 102) = 360 - 304 = 56 degrees. This still does not match any of the given options. Rechecking the problem, the correct approach is to consider the angles formed by the diagonals and the sides of the rectangle.\n",
      "\n",
      "The correct approach is to recognize that the angles formed by the diagonals and the rectangle's sides are supplementary to the angles at the vertices of the rectangle. Since the diagonals bisect each other at 90 degrees, the angles at the vertices of the rectangle are 90 degrees.\n",
      "\n",
      "Therefore, the angle H, which is the angle formed by the diagonals and the sides of one of the triangles formed by the diagonals, is 180 - 90 - 102 = 88 degrees. However, this does not match any given options.\n",
      "\n",
      "Revisiting the options, the correct answer is 102 degrees, as it is the angle formed by the diagonals and one of the sides of the rectangle.\n",
      "</thinking>\n",
      "<answer>\n",
      "\\boxed{B}\n",
      "</answer>\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n",
    "\n",
    "\n",
    "cot_prompt = \"SYSTEM: You are a helpful assistant. When the user asks a question, your response must include two parts: first, the reasoning process enclosed in <thinking>...</thinking> tags, then the final answer enclosed in <answer>...</answer> tags. The critical answer or key result should be placed within \\\\boxed{}.\\n\"\n",
    "# And your input message should be like this:\n",
    "question = \"Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end.\\nQuestion: Find $m\\\\angle H$\\nChoices:\\n(A) 97\\n(B) 102\\n(C) 107\\n(D) 122\"\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"image\", \"image\": \"figures/cases/reasoning.png\"},\n",
    "            {\"type\": \"text\", \"text\": cot_prompt + question},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "from IPython.utils import io\n",
    "with io.capture_output() as captured: # used to suppress output inside the function\n",
    "    output_text = generate(messages)\n",
    "\n",
    "from IPython.display import display, Image\n",
    "display(Image(\"figures/cases/reasoning.png\"))\n",
    "print(f\"Input: {question}\\n\")\n",
    "print(\"Answer:\\n\")\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e462690-0407-4ab0-ab3d-e79704b9195a",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🎥 Video QA Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4d26e5bd-5381-4640-9130-15186e58da1c",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:23:35.516352Z",
     "shell.execute_reply.started": "2025-07-21T17:23:33.907430Z",
     "to_execute": "2025-07-21T17:23:34.412Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:23:34,087 - bailingmm_utils - INFO - decord:  video_path='figures/cases/yoga.mp4', total_frames=300, video_fps=30.0, time=0.082s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<video alt=\"test\" controls>\n",
       "                <source src=\"data:video/mp4;base64,[encoded data, 1034912 characters]\" type=\"video/mp4\" />\n",
       "             </video>\n",
       "             <p>What is the woman doing in the video?</p>\n",
       "             <p>Answer: The woman is performing a yoga sequence.</p>\n",
       "          "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# video qa\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"video\", \"video\": \"figures/cases/yoga.mp4\", \"max_frames\": 64, \"sample\": \"uniform\"},\n",
    "            {\"type\": \"text\", \"text\": \"What is the woman doing?\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "from IPython.utils import io\n",
    "with io.capture_output() as captured: # used to suppress output inside the function\n",
    "    output_text = generate(messages)\n",
    "\n",
    "import io\n",
    "import base64\n",
    "from IPython.display import HTML\n",
    "\n",
    "video = io.open('figures/cases/yoga.mp4', 'r+b').read()\n",
    "encoded = base64.b64encode(video)\n",
    "HTML(data='''<video alt=\"test\" controls>\n",
    "                <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
    "             </video>\n",
    "             <p>What is the woman doing in the video?</p>\n",
    "             <p>Answer: {1}</p>\n",
    "          '''.format(encoded.decode('ascii'), output_text))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8232eef-f42b-4b3d-80cc-a92c3ee5ece8",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 💬 Multi-turn Conversation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e2c4a985-fb06-4345-b4d5-525166f909f4",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:23:53.681648Z",
     "shell.execute_reply.started": "2025-07-21T17:23:53.072337Z",
     "to_execute": "2025-07-21T17:23:53.865Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: 中国的首都是哪里？\n",
      "Assistant: 北京\n",
      "Human: 它有哪些著名景点？列出3个。\n",
      "\n",
      "Answer:\n",
      "故宫、长城、颐和园。\n"
     ]
    }
   ],
   "source": [
    "# multi-turn chat\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"中国的首都是哪里？\"},\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"ASSISTANT\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"北京\"},\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"它有哪些著名景点？列出3个。\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "from IPython.utils import io\n",
    "with io.capture_output() as captured: # used to suppress output inside the function\n",
    "    output_text = generate(messages)\n",
    "\n",
    "print(\"Human: 中国的首都是哪里？\")\n",
    "print(\"Assistant: 北京\")\n",
    "print(\"Human: 它有哪些著名景点？列出3个。\\n\")\n",
    "\n",
    "print(\"Answer:\")\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b06bd343-27f6-4fe6-8757-1244ef9f0146",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🔊 Audio Tasks\n",
    "\n",
    "For detailed usage for ASR, SpeechQA, and TTS tasks, please refer to `test_audio_tasks.py`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e77a90ef-774c-4562-8ba9-f72595631b96",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:24:25.607199Z",
     "shell.execute_reply.started": "2025-07-21T17:24:01.582982Z",
     "to_execute": "2025-07-21T17:24:02.379Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:24:07,546 - root - INFO - input frame rate=50\n",
      "2025-07-22 01:24:17,154 - root - INFO - input frame rate=50\n"
     ]
    }
   ],
   "source": [
    "from hyperpyyaml import load_hyperpyyaml\n",
    "from audio_detokenizer.cli.frontend import TTSFrontEnd\n",
    "from modeling_bailing_talker import AudioDetokenizer\n",
    "\n",
    "model_path = 'inclusionAI/Ming-Lite-Omni'\n",
    "\n",
    "spk_info = {\n",
    "        'luna': torch.load('data/spks/luna_v2.pt'),\n",
    "        'eng': torch.load('data/spks/eng_v2.pt'),\n",
    "    }\n",
    "\n",
    "with open(f\"{model_path}/talker/audio_detokenizer_stream.yaml\", \"r\") as f:\n",
    "    configs = load_hyperpyyaml(f)\n",
    "\n",
    "audio_detokenizer = AudioDetokenizer(\n",
    "    f\"{model_path}/talker/audio_detokenizer_stream.yaml\",\n",
    "    flow_model_path=f\"{model_path}/talker/flow_stream.pt\",\n",
    "    hifigan_model_path=f\"{model_path}/talker/hift_v2.pt\",\n",
    "    spk_info=spk_info,\n",
    ")\n",
    "# new mel\n",
    "audio_frontend = TTSFrontEnd(\n",
    "    configs[\"feat_extractor\"],\n",
    "    f\"{model_path}/talker/campplus.onnx\",\n",
    "    f\"{model_path}/talker/speech_tokenizer_v1.onnx\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d8be6821-b107-4d50-ac0a-a10bf7fa7979",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:24:36.122632Z",
     "shell.execute_reply.started": "2025-07-21T17:24:34.509888Z",
     "to_execute": "2025-07-21T17:24:35.289Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing ASR...\n",
      "Chinese\t也正是看中了中国消费者的消费潜力\n",
      "Generate time: 0.72s\n",
      "Input audio:\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/x-wav;base64,[encoded data, 192060 characters]\" type=\"audio/x-wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output text: Chinese\t也正是看中了中国消费者的消费潜力\n"
     ]
    }
   ],
   "source": [
    "# ASR\n",
    "from IPython.display import display, Audio\n",
    "from test_audio_tasks import generate\n",
    "print(\"Testing ASR...\")\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\",\n",
    "                \"text\": \"Please recognize the language of this speech and transcribe it. Format: oral.\",\n",
    "            },\n",
    "            {\"type\": \"audio\", \"audio\": \"data/wavs/BAC009S0915W0283.wav\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "output_text = generate(messages=messages, processor=processor, model=model)\n",
    "\n",
    "print(\"Input audio:\")\n",
    "display(Audio('data/wavs/BAC009S0915W0283.wav'))\n",
    "print(f\"Output text: {output_text}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4b256a09-e83a-4c3e-b182-81c069a52564",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:25:34.402201Z",
     "shell.execute_reply.started": "2025-07-21T17:25:23.995255Z",
     "to_execute": "2025-07-21T17:25:24.445Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing Speech QA...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "From v4.47 onwards, when a model cache is to be returned, `generate` will return a `Cache` instance instead by default (as opposed to the legacy tuple of tuples format). If you want to keep returning the legacy format, please set `return_legacy_cache=True`.\n",
      "Starting from v4.46, the `logits` model output will have the same type as the model (except at train time, where it will always be FP32)\n",
      "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class (https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate time: 10.11s\n",
      "Input audio:\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/x-wav;base64,[encoded data, 68828 characters]\" type=\"audio/x-wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output audio: 你好！我是一个名为百灵原生多模态的多模态模型，由蚂蚁集团研发。我可以回答你的问题，帮助你解决疑惑。请问有什么我可以帮助你的吗？。\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/x-wav;base64,[encoded data, 1346668 characters]\" type=\"audio/x-wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Speech QA\n",
    "from test_audio_tasks import generate_e2e\n",
    "\n",
    "print(\"Testing Speech QA...\")\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"audio\", \"audio\": \"data/wavs/speechQA_sample.wav\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "GENERATION_CONFIG = {\n",
    "    \"output_hidden_states\": True,\n",
    "    \"return_dict_in_generate\": True,\n",
    "    \"no_repeat_ngram_size\": 10,\n",
    "}\n",
    "output_audio_path = 'data/wavs/out_speechqa.wav'\n",
    "\n",
    "text_response, audio_wave = generate_e2e(\n",
    "    messages=messages,\n",
    "    processor=processor,\n",
    "    model=model,\n",
    "    audio_detokenizer=audio_detokenizer,\n",
    "    output_audio_path=output_audio_path,\n",
    "    spk_info=spk_info,\n",
    "    generation_config=GENERATION_CONFIG,\n",
    "    stream=False,\n",
    ")\n",
    "\n",
    "print(\"Input audio:\")\n",
    "display(Audio('data/wavs/speechQA_sample.wav'))\n",
    "\n",
    "print(f\"Output audio: {text_response}\")\n",
    "display(Audio(output_audio_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "055851d7-8e18-4097-b7ac-ff0dc6efcfdb",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:26:06.344349Z",
     "shell.execute_reply.started": "2025-07-21T17:26:01.551140Z",
     "to_execute": "2025-07-21T17:26:02.376Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing TTS...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-22 01:26:04,502 - root - INFO - for last sentence\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate time: 4.72s\n",
      "input text: 这是一条测试语句。\n",
      "Output audio:\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" >\n",
       "                    <source src=\"data:audio/x-wav;base64,[encoded data, 225388 characters]\" type=\"audio/x-wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# zero-shot TTS\n",
    "from test_audio_tasks import generate_tts\n",
    "\n",
    "output_audio_path = 'data/wavs/out_tts.wav'\n",
    "# TTS Generation\n",
    "print(\"Testing TTS...\")\n",
    "tts_text=\"这是一条测试语句。\"\n",
    "tts_wave = generate_tts(\n",
    "    tts_text=tts_text,\n",
    "    prompt_text=\"感谢你的认可。\",\n",
    "    prompt_wav_path=\"data/spks/prompt.wav\",\n",
    "    audio_frontend=audio_frontend,\n",
    "    model=model,\n",
    "    audio_detokenizer=audio_detokenizer,\n",
    "    output_audio_path=output_audio_path,\n",
    "    stream=False,\n",
    ")\n",
    "\n",
    "print(f\"input text: {tts_text}\")\n",
    "print(f\"Output audio:\")\n",
    "display(Audio(output_audio_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "550037e0-7a94-4db7-b3f1-dc3c2f265be1",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "### 🎨 Image Generation & Editing\n",
    "\n",
    "Ming-omni natively supports image generation and image editing. To use this function, you only need to add the corresponding parameters in the generate function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25d1f221-3bcd-4d80-9176-b9cb7adae55d",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "Text-to-image generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b301fd40-9853-49b8-8786-c33bab751167",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:27:20.746210Z",
     "shell.execute_reply.started": "2025-07-21T17:27:17.643696Z",
     "to_execute": "2025-07-21T17:27:18.475Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Hyper-Param]: \n",
      "   do_img_cfg: False \n",
      "   guidance_scale: 5.0 \n",
      "   img_cfg: 1.0 \n",
      "   cfg_mode: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:02<00:00, 12.41it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instruction: Draw a beautiful girl with short black hair and red dress.\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 26304 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Image generation mode currently limits the range of input pixels.\n",
    "gen_input_pixels = 451584\n",
    "processor.image_processor.max_pixels = gen_input_pixels\n",
    "processor.image_processor.min_pixels = gen_input_pixels\n",
    "\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"text\", \"text\": \"Draw a beautiful girl with short black hair and red dress.\"},\n",
    "        ],\n",
    "    }\n",
    "]\n",
    "\n",
    "text = processor.apply_chat_template(messages, add_generation_prompt=True)\n",
    "image_inputs, video_inputs, audio_inputs = processor.process_vision_info(messages)\n",
    "\n",
    "\n",
    "inputs = processor(\n",
    "    text=[text],\n",
    "    images=image_inputs,\n",
    "    videos=video_inputs,\n",
    "    audios=audio_inputs,\n",
    "    return_tensors=\"pt\",\n",
    ").to(model.device)\n",
    "\n",
    "for k in inputs.keys():\n",
    "    if k in [\"pixel_values\", \"pixel_values_videos\", \"audio_feats\", \"pixel_values_reference\"]:\n",
    "        inputs[k] = inputs[k].to(dtype=torch.bfloat16)\n",
    "\n",
    "# set `image_gen=True` to enable image generation\n",
    "image = model.generate(\n",
    "    **inputs,\n",
    "    image_gen=True,\n",
    ")\n",
    "\n",
    "image.save(\"figures/cases/t2i.jpg\")\n",
    "\n",
    "print(\"Instruction: Draw a beautiful girl with short black hair and red dress.\")\n",
    "from IPython.display import display, Image\n",
    "display(Image(\"figures/cases/t2i.jpg\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dc27595-b61a-40e6-92bb-59189883348a",
   "metadata": {
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "source": [
    "Image Editing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0cb482d1-f1e2-4ef7-9c0f-651f575036f6",
   "metadata": {
    "execution": {
     "shell.execute_reply.end": "2025-07-21T17:27:56.863249Z",
     "shell.execute_reply.started": "2025-07-21T17:27:51.774234Z",
     "to_execute": "2025-07-21T17:27:52.640Z"
    },
    "isLargeOutputDisplay": true,
    "libroFormatter": "formatter-string",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Hyper-Param]: \n",
      "   do_img_cfg: False \n",
      "   guidance_scale: 5.0 \n",
      "   img_cfg: 1.0 \n",
      "   cfg_mode: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 30/30 [00:04<00:00,  7.32it/s]\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 37536 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instruction: add candles in the middle\n",
      "\n",
      "Output image:\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 39600 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "messages = [\n",
    "    {\n",
    "        \"role\": \"HUMAN\",\n",
    "        \"content\": [\n",
    "            {\"type\": \"image\", \"image\": \"figures/cases/cake.jpg\"},\n",
    "            {\"type\": \"text\", \"text\": \"add candles in the middle\"},\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "text = processor.apply_chat_template(messages, add_generation_prompt=True)\n",
    "image_inputs, video_inputs, audio_inputs = processor.process_vision_info(messages)\n",
    "\n",
    "\n",
    "inputs = processor(\n",
    "    text=[text],\n",
    "    images=image_inputs,\n",
    "    videos=video_inputs,\n",
    "    audios=audio_inputs,\n",
    "    return_tensors=\"pt\",\n",
    ").to(model.device)\n",
    "\n",
    "for k in inputs.keys():\n",
    "    if k in [\"pixel_values\", \"pixel_values_videos\", \"audio_feats\", \"pixel_values_reference\"]:\n",
    "        inputs[k] = inputs[k].to(dtype=torch.bfloat16)\n",
    "\n",
    "image = model.generate(\n",
    "    **inputs,\n",
    "    image_gen=True,\n",
    ")\n",
    "image.save(\"figures/cases/cake_edit.jpg\")\n",
    "\n",
    "from IPython.display import display, Image\n",
    "display(Image(\"figures/cases/cake.jpg\"))\n",
    "print(\"Instruction: add candles in the middle\\n\")\n",
    "\n",
    "print(\"Output image:\")\n",
    "display(Image(\"figures/cases/cake_edit.jpg\"))\n",
    "     "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (myenv)",
   "name": "myenv"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
