1
0
forked from cgvr/DeltaVR

123 lines
3.9 KiB
Plaintext

{
"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
}