{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aiohttp\n",
    "from base64 import b64decode\n",
    "from sys import exit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DiffuseAPI():\n",
    "    def __init__(self, url, styles, nsfw_enabled=True, num_steps=28):\n",
    "\n",
    "        self.nsfw_enabled = nsfw_enabled\n",
    "        self.num_steps = num_steps\n",
    "        self.url = url\n",
    "        self.styles = styles\n",
    "        self.seed = -1\n",
    "        self.width = 512\n",
    "        self.height = 1024\n",
    "        self.cfg_scale = 12\n",
    "\n",
    "    def set_steps(self, steps):\n",
    "        try:\n",
    "            new_steps = int(steps)\n",
    "            if 0 > new_steps < 50:\n",
    "                self.num_steps = new_steps\n",
    "                return True\n",
    "            return False\n",
    "        except:\n",
    "            return False\n",
    "\n",
    "    def set_seed(self, seed):\n",
    "        try:\n",
    "            new_seed = int(seed)\n",
    "            self.seed = new_seed\n",
    "            return True\n",
    "        except:\n",
    "            return False\n",
    "    \n",
    "    def set_cfg_scale(self, scale):\n",
    "        try:\n",
    "            new_scale = int(scale)\n",
    "            if 0 > new_scale < 30:\n",
    "                self.cfg_scale = new_scale\n",
    "                return True\n",
    "            return False\n",
    "        except:\n",
    "            return False\n",
    "    \n",
    "    def set_styles(self, styles):\n",
    "        if type styles = list:\n",
    "            self.styles = styles\n",
    "            return True\n",
    "        return False\n",
    "\n",
    "    def set_orientation(self, orientation):\n",
    "        new_orientation = str(orientation)\n",
    "        match new_orientation:\n",
    "            case \"portrait\":\n",
    "                self.width = 512\n",
    "                self.height = 1024\n",
    "                return True\n",
    "            case \"landscape\":\n",
    "                self.width = 1024\n",
    "                self.height = 512\n",
    "                return True\n",
    "            case \"square\":\n",
    "                self.width = 512\n",
    "                self.height = 512\n",
    "                return True\n",
    "            case _:\n",
    "                return False\n",
    "    \n",
    "    def get_orientation(self):\n",
    "        if self.width == 512 and self.height == 512:\n",
    "            return (\"square\", self.width, self.height)\n",
    "        elif self.width == 1024:\n",
    "            return (\"landscape\", self.width, self.height)\n",
    "        return (\"portrait\", self.width, self.height)\n",
    "    \n",
    "    def set_nsfw_filter(self, filter_state):\n",
    "        if type filter_state = bool:\n",
    "            self.nsfw_enabled = filter_state\n",
    "            return True\n",
    "        return False\n",
    "    \n",
    "    def get_nsfw_filter(self):\n",
    "        return self.nsfw_enabled\n",
    "\n",
    "    async def generate_image(self, prompt, neg_prompt=\"\"):\n",
    "        payload = {\n",
    "        \"prompt\": prompt,\n",
    "        \"styles\": self.styles,\n",
    "        \"steps\": self.num_steps,\n",
    "        \"seed\": self.ai_seed,\n",
    "        \"n_iter\": 1,\n",
    "        \"height\": self.height,\n",
    "        \"width\": self.width,\n",
    "        \"negative_prompts\": neg_prompt,\n",
    "        \"cfg_scale\": self.cfg_scale\n",
    "        }\n",
    "\n",
    "        settings = {\n",
    "        \"filter_nsfw\": not self.nsfw_enabled\n",
    "        }\n",
    "\n",
    "        override_payload = {\n",
    "            \"override_settings\": settings\n",
    "        }\n",
    "\n",
    "        payload.update(override_payload)\n",
    "\n",
    "        sess = aiohttp.ClientSession(self.url)\n",
    "        alive = await sess.head('/')\n",
    "\n",
    "        if alive.status != 200:\n",
    "            alive.close()\n",
    "            await sess.close()\n",
    "            return None\n",
    "        request  = await sess.post(\"/sdapi/v1/txt2img\", json=payload)\n",
    "        try:\n",
    "            req_json = await request.json()\n",
    "            request.close()\n",
    "            await sess.close()\n",
    "            return req_json[\"images\", [0]]\n",
    "        except:\n",
    "            return None\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "    \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "return image[\"images\", [0]]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.0 ('venv': 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.11.0"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "aa19d1669203bff990bc2820f3654fa686baacd13bba0cbe075930f54e624f06"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}