{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50e24baa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from invokeai_mcp_server import create_text2img_graph, enqueue_graph, wait_for_completion, get_image_url\n",
    "from urllib.parse import urljoin\n",
    "\n",
    "import asyncio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0407cd9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "INVOKEAI_BASE_URL = \"http://127.0.0.1:9090\"\n",
    "\n",
    "\n",
    "async def generate_image(arguments: dict):\n",
    "\n",
    "    # Extract parameters\n",
    "    prompt = arguments[\"prompt\"]\n",
    "    negative_prompt = arguments.get(\"negative_prompt\", \"\")\n",
    "    width = arguments.get(\"width\", 512)\n",
    "    height = arguments.get(\"height\", 512)\n",
    "    steps = arguments.get(\"steps\", 30)\n",
    "    cfg_scale = arguments.get(\"cfg_scale\", 7.5)\n",
    "    scheduler = arguments.get(\"scheduler\", \"euler\")\n",
    "    seed = arguments.get(\"seed\")\n",
    "    model_key = arguments.get(\"model_key\")\n",
    "    lora_key = arguments.get(\"lora_key\")\n",
    "    lora_weight = arguments.get(\"lora_weight\", 1.0)\n",
    "    vae_key = arguments.get(\"vae_key\")\n",
    "\n",
    "    #logger.info(f\"Generating image with prompt: {prompt[:50]}...\")\n",
    "\n",
    "    # Create graph\n",
    "    graph = await create_text2img_graph(\n",
    "        prompt=prompt,\n",
    "        negative_prompt=negative_prompt,\n",
    "        model_key=model_key,\n",
    "        lora_key=lora_key,\n",
    "        lora_weight=lora_weight,\n",
    "        vae_key=vae_key,\n",
    "        width=width,\n",
    "        height=height,\n",
    "        steps=steps,\n",
    "        cfg_scale=cfg_scale,\n",
    "        scheduler=scheduler,\n",
    "        seed=seed\n",
    "    )\n",
    "\n",
    "    # Enqueue and wait for completion\n",
    "    result = await enqueue_graph(graph)\n",
    "    batch_id = result[\"batch\"][\"batch_id\"]\n",
    "\n",
    "    #logger.info(f\"Enqueued batch {batch_id}, waiting for completion...\")\n",
    "\n",
    "    completed = await wait_for_completion(batch_id)\n",
    "\n",
    "    # Extract image name from result\n",
    "    if \"result\" in completed and \"outputs\" in completed[\"result\"]:\n",
    "        outputs = completed[\"result\"][\"outputs\"]\n",
    "        # Find the image output\n",
    "        for node_id, output in outputs.items():\n",
    "            if output.get(\"type\") == \"image_output\":\n",
    "                image_name = output[\"image\"][\"image_name\"]\n",
    "                image_url = await get_image_url(image_name)\n",
    "\n",
    "                text=f\"Image generated successfully!\\n\\nImage Name: {image_name}\\nImage URL: {image_url}\\n\\nYou can view the image at: {urljoin(INVOKEAI_BASE_URL, f'/api/v1/images/i/{image_name}/full')}\"\n",
    "                print(text)\n",
    "\n",
    "    # Fallback if we couldn't find image output\n",
    "    #text=f\"Image generation completed but output format was unexpected. Batch ID: {batch_id}\\n\\nResult: {json.dumps(completed, indent=2)}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cf9d879",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def main():\n",
    "    args = {\n",
    "        \"prompt\": \"a golden katana with a fancy pommel\"\n",
    "    }\n",
    "    await generate_image(args)\n",
    "\n",
    "asyncio.run(main())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}