{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ComfyScript Runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nodes: 172\n"
     ]
    }
   ],
   "source": [
    "from comfy_script.runtime import *\n",
    "\n",
    "# load('http://127.0.0.1:8188/')\n",
    "load()\n",
    "\n",
    "# Nodes can only be imported after load()\n",
    "from comfy_script.runtime.nodes import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Queue remaining: 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████| 20/20\n",
      "Queue remaining: 0\n"
     ]
    }
   ],
   "source": [
    "with Workflow():\n",
    "    model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "    conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "    conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "    latent = EmptyLatentImage(512, 512, 1)\n",
    "    latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "    image = VAEDecode(latent, vae)\n",
    "    SaveImage(image, 'ComfyUI')\n",
    "\n",
    "    # To retrieve `image` instead of saving it, replace `SaveImage` with:\n",
    "    # images = util.get_images(image)\n",
    "    # `images` is of type `list[PIL.Image.Image]`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To control the queue:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "queue.cancel_current()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "queue.cancel_remaining()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "queue.cancel_all()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To wait for the task:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Queue remaining: 1\n",
      "Queue remaining: 0\n"
     ]
    }
   ],
   "source": [
    "with Workflow(wait=True):\n",
    "    model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "    conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "    conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "    latent = EmptyLatentImage(512, 512, 1)\n",
    "    latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "    image = VAEDecode(latent, vae)\n",
    "    SaveImage(image, 'ComfyUI')\n",
    "\n",
    "# Or:\n",
    "with Workflow() as wf:\n",
    "    model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "    conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "    conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "    latent = EmptyLatentImage(512, 512, 1)\n",
    "    latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "    image = VAEDecode(latent, vae)\n",
    "    SaveImage(image, 'ComfyUI')\n",
    "wf.task.wait()\n",
    "\n",
    "# Or with await:\n",
    "# await wf.task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Without `Workflow`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task 39 (315e8ccd-63eb-48ad-8818-8f22d15520de)\n"
     ]
    }
   ],
   "source": [
    "model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "latent = EmptyLatentImage(512, 512, 1)\n",
    "latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "image = VAEDecode(latent, vae)\n",
    "save = SaveImage(image, 'ComfyUI')\n",
    "\n",
    "task = queue.put(save)\n",
    "print(task)\n",
    "\n",
    "# To wait the task:\n",
    "# task.wait()\n",
    "# or:\n",
    "# await task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or use `Workflow` but not use `with`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task 40 (7fa2acca-0ecb-41d1-be73-5450200805e5)\n"
     ]
    }
   ],
   "source": [
    "wf = Workflow()\n",
    "model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "latent = EmptyLatentImage(512, 512, 1)\n",
    "latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "image = VAEDecode(latent, vae)\n",
    "wf += SaveImage(image, 'ComfyUI')\n",
    "\n",
    "task = queue.put(wf)\n",
    "print(task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or just `await` or `wait()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "latent = EmptyLatentImage(512, 512, 1)\n",
    "latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "image = VAEDecode(latent, vae)\n",
    "await SaveImage(image, 'ComfyUI')\n",
    "# or:\n",
    "# SaveImage(image, 'ComfyUI').wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Select and process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "queue.watch_display(False)\n",
    "\n",
    "latents = []\n",
    "image_batches = []\n",
    "with Workflow():\n",
    "    seed = 0\n",
    "    pos = 'sky, 1girl, smile'\n",
    "    neg = 'embedding:easynegative'\n",
    "    model, clip, vae = CheckpointLoaderSimple(Checkpoints.AOM3A1B_orangemixs)\n",
    "    model2, clip2, vae2 = CheckpointLoaderSimple(Checkpoints.CounterfeitV25_25)\n",
    "    for color in 'red', 'green', 'blue':\n",
    "        latent = EmptyLatentImage(440, 640)\n",
    "        latent = KSampler(model, seed, steps=15, cfg=6, sampler_name='uni_pc',\n",
    "                          positive=CLIPTextEncode(f'{color}, {pos}', clip), negative=CLIPTextEncode(neg, clip),\n",
    "                          latent_image=latent)\n",
    "        latents.append(latent)\n",
    "        image_batches.append(SaveImage(VAEDecode(latent, vae), f'{seed} {color}'))\n",
    "\n",
    "grid = widgets.GridspecLayout(1, len(image_batches))\n",
    "for i, image_batch in enumerate(image_batches):\n",
    "    image_batch = image_batch.wait()\n",
    "    image = widgets.Image(value=image_batch[0]._repr_png_())\n",
    "\n",
    "    button = widgets.Button(description=f'Hires fix {i}')\n",
    "    def hiresfix(button, i=i):\n",
    "        print(f'Image {i} is chosen')\n",
    "        with Workflow():\n",
    "            latent = LatentUpscaleBy(latents[i], scale_by=2)\n",
    "            latent = KSampler(model2, seed, steps=15, cfg=6, sampler_name='uni_pc',\n",
    "                            positive=CLIPTextEncode(pos, clip2), negative=CLIPTextEncode(neg, clip2),\n",
    "                            latent_image=latent, denoise=0.6)\n",
    "            image_batch = SaveImage(VAEDecode(latent, vae2), f'{seed} hires')\n",
    "        display(image_batch.wait())\n",
    "    button.on_click(hiresfix)\n",
    "\n",
    "    grid[0, i] = widgets.VBox(children=(image, button))\n",
    "display(grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Workflow generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"0\": {\n",
      "    \"inputs\": {\n",
      "      \"ckpt_name\": \"v1-5-pruned-emaonly.ckpt\"\n",
      "    },\n",
      "    \"class_type\": \"CheckpointLoaderSimple\"\n",
      "  },\n",
      "  \"1\": {\n",
      "    \"inputs\": {\n",
      "      \"text\": \"beautiful scenery nature glass bottle landscape, , purple galaxy bottle,\",\n",
      "      \"clip\": [\n",
      "        \"0\",\n",
      "        1\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"CLIPTextEncode\"\n",
      "  },\n",
      "  \"2\": {\n",
      "    \"inputs\": {\n",
      "      \"text\": \"text, watermark\",\n",
      "      \"clip\": [\n",
      "        \"0\",\n",
      "        1\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"CLIPTextEncode\"\n",
      "  },\n",
      "  \"3\": {\n",
      "    \"inputs\": {\n",
      "      \"width\": 512,\n",
      "      \"height\": 512,\n",
      "      \"batch_size\": 1\n",
      "    },\n",
      "    \"class_type\": \"EmptyLatentImage\"\n",
      "  },\n",
      "  \"4\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 1,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"3\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"5\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"4\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"6\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"0\",\n",
      "      \"images\": [\n",
      "        \"5\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  },\n",
      "  \"7\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 0.8,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"4\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"8\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"7\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"9\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"0\",\n",
      "      \"images\": [\n",
      "        \"8\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  },\n",
      "  \"10\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 0.8,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"7\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"11\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"10\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"12\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"1\",\n",
      "      \"images\": [\n",
      "        \"11\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  },\n",
      "  \"13\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 0.8,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"10\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"14\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"13\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"15\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"2\",\n",
      "      \"images\": [\n",
      "        \"14\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  },\n",
      "  \"16\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 0.8,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"13\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"17\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"16\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"18\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"3\",\n",
      "      \"images\": [\n",
      "        \"17\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  },\n",
      "  \"19\": {\n",
      "    \"inputs\": {\n",
      "      \"seed\": 123,\n",
      "      \"steps\": 20,\n",
      "      \"cfg\": 8,\n",
      "      \"sampler_name\": \"euler\",\n",
      "      \"scheduler\": \"normal\",\n",
      "      \"denoise\": 0.8,\n",
      "      \"model\": [\n",
      "        \"0\",\n",
      "        0\n",
      "      ],\n",
      "      \"positive\": [\n",
      "        \"1\",\n",
      "        0\n",
      "      ],\n",
      "      \"negative\": [\n",
      "        \"2\",\n",
      "        0\n",
      "      ],\n",
      "      \"latent_image\": [\n",
      "        \"16\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"KSampler\"\n",
      "  },\n",
      "  \"20\": {\n",
      "    \"inputs\": {\n",
      "      \"samples\": [\n",
      "        \"19\",\n",
      "        0\n",
      "      ],\n",
      "      \"vae\": [\n",
      "        \"0\",\n",
      "        2\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"VAEDecode\"\n",
      "  },\n",
      "  \"21\": {\n",
      "    \"inputs\": {\n",
      "      \"filename_prefix\": \"4\",\n",
      "      \"images\": [\n",
      "        \"20\",\n",
      "        0\n",
      "      ]\n",
      "    },\n",
      "    \"class_type\": \"SaveImage\"\n",
      "  }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "with Workflow(queue=False) as wf:\n",
    "    model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "    conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "    conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "    latent = EmptyLatentImage(512, 512, 1)\n",
    "    latent = KSampler(model, 123, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "    SaveImage(VAEDecode(latent, vae), '0')\n",
    "    for i in range(5):\n",
    "        latent = KSampler(model, 123, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 0.8)\n",
    "        SaveImage(VAEDecode(latent, vae), f'{i}')\n",
    "\n",
    "json = wf.api_format_json()\n",
    "with open('prompt.json', 'w') as f:\n",
    "    f.write(json)\n",
    "\n",
    "print(json)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real mode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from comfy_script.runtime.real import *\n",
    "load()\n",
    "from comfy_script.runtime.real.nodes import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with Workflow():\n",
    "    model, clip, vae = CheckpointLoaderSimple('v1-5-pruned-emaonly.ckpt')\n",
    "    conditioning = CLIPTextEncode('beautiful scenery nature glass bottle landscape, , purple galaxy bottle,', clip)\n",
    "    conditioning2 = CLIPTextEncode('text, watermark', clip)\n",
    "    latent = EmptyLatentImage(512, 512, 1)\n",
    "    latent = KSampler(model, 156680208700286, 20, 8, 'euler', 'normal', conditioning, conditioning2, latent, 1)\n",
    "    image = VAEDecode(latent, vae)\n",
    "    SaveImage(image, 'ComfyUI')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}