incorporate InvokeAI into start_pipeline.py

3d-generation-pipeline/notebooks/local_image_generation.ipynb

@@ -0,0 +1,122 @@
+{
+ "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
+}