[
  {
    "path": ".gitignore",
    "content": "token.txt\n__pycache__"
  },
  {
    "path": "LICENSE",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\nTERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n1.  Definitions.\n\n    \"License\" shall mean the terms and conditions for use, reproduction,\n    and distribution as defined by Sections 1 through 9 of this document.\n\n    \"Licensor\" shall mean the copyright owner or entity authorized by\n    the copyright owner that is granting the License.\n\n    \"Legal Entity\" shall mean the union of the acting entity and all\n    other entities that control, are controlled by, or are under common\n    control with that entity. For the purposes of this definition,\n    \"control\" means (i) the power, direct or indirect, to cause the\n    direction or management of such entity, whether by contract or\n    otherwise, or (ii) ownership of fifty percent (50%) or more of the\n    outstanding shares, or (iii) beneficial ownership of such entity.\n\n    \"You\" (or \"Your\") shall mean an individual or Legal Entity\n    exercising permissions granted by this License.\n\n    \"Source\" form shall mean the preferred form for making modifications,\n    including but not limited to software source code, documentation\n    source, and configuration files.\n\n    \"Object\" form shall mean any form resulting from mechanical\n    transformation or translation of a Source form, including but\n    not limited to compiled object code, generated documentation,\n    and conversions to other media types.\n\n    \"Work\" shall mean the work of authorship, whether in Source or\n    Object form, made available under the License, as indicated by a\n    copyright notice that is included in or attached to the work\n    (an example is provided in the Appendix below).\n\n    \"Derivative Works\" shall mean any work, whether in Source or Object\n    form, that is based on (or derived from) the Work and for which the\n    editorial revisions, annotations, elaborations, or other modifications\n    represent, as a whole, an original work of authorship. For the purposes\n    of this License, Derivative Works shall not include works that remain\n    separable from, or merely link (or bind by name) to the interfaces of,\n    the Work and Derivative Works thereof.\n\n    \"Contribution\" shall mean any work of authorship, including\n    the original version of the Work and any modifications or additions\n    to that Work or Derivative Works thereof, that is intentionally\n    submitted to Licensor for inclusion in the Work by the copyright owner\n    or by an individual or Legal Entity authorized to submit on behalf of\n    the copyright owner. For the purposes of this definition, \"submitted\"\n    means any form of electronic, verbal, or written communication sent\n    to the Licensor or its representatives, including but not limited to\n    communication on electronic mailing lists, source code control systems,\n    and issue tracking systems that are managed by, or on behalf of, the\n    Licensor for the purpose of discussing and improving the Work, but\n    excluding communication that is conspicuously marked or otherwise\n    designated in writing by the copyright owner as \"Not a Contribution.\"\n\n    \"Contributor\" shall mean Licensor and any individual or Legal Entity\n    on behalf of whom a Contribution has been received by Licensor and\n    subsequently incorporated within the Work.\n\n2.  Grant of Copyright License. Subject to the terms and conditions of\n    this License, each Contributor hereby grants to You a perpetual,\n    worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n    copyright license to reproduce, prepare Derivative Works of,\n    publicly display, publicly perform, sublicense, and distribute the\n    Work and such Derivative Works in Source or Object form.\n\n3.  Grant of Patent License. Subject to the terms and conditions of\n    this License, each Contributor hereby grants to You a perpetual,\n    worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n    (except as stated in this section) patent license to make, have made,\n    use, offer to sell, sell, import, and otherwise transfer the Work,\n    where such license applies only to those patent claims licensable\n    by such Contributor that are necessarily infringed by their\n    Contribution(s) alone or by combination of their Contribution(s)\n    with the Work to which such Contribution(s) was submitted. If You\n    institute patent litigation against any entity (including a\n    cross-claim or counterclaim in a lawsuit) alleging that the Work\n    or a Contribution incorporated within the Work constitutes direct\n    or contributory patent infringement, then any patent licenses\n    granted to You under this License for that Work shall terminate\n    as of the date such litigation is filed.\n\n4.  Redistribution. You may reproduce and distribute copies of the\n    Work or Derivative Works thereof in any medium, with or without\n    modifications, and in Source or Object form, provided that You\n    meet the following conditions:\n\n    (a) You must give any other recipients of the Work or\n    Derivative Works a copy of this License; and\n\n    (b) You must cause any modified files to carry prominent notices\n    stating that You changed the files; and\n\n    (c) You must retain, in the Source form of any Derivative Works\n    that You distribute, all copyright, patent, trademark, and\n    attribution notices from the Source form of the Work,\n    excluding those notices that do not pertain to any part of\n    the Derivative Works; and\n\n    (d) If the Work includes a \"NOTICE\" text file as part of its\n    distribution, then any Derivative Works that You distribute must\n    include a readable copy of the attribution notices contained\n    within such NOTICE file, excluding those notices that do not\n    pertain to any part of the Derivative Works, in at least one\n    of the following places: within a NOTICE text file distributed\n    as part of the Derivative Works; within the Source form or\n    documentation, if provided along with the Derivative Works; or,\n    within a display generated by the Derivative Works, if and\n    wherever such third-party notices normally appear. The contents\n    of the NOTICE file are for informational purposes only and\n    do not modify the License. You may add Your own attribution\n    notices within Derivative Works that You distribute, alongside\n    or as an addendum to the NOTICE text from the Work, provided\n    that such additional attribution notices cannot be construed\n    as modifying the License.\n\n    You may add Your own copyright statement to Your modifications and\n    may provide additional or different license terms and conditions\n    for use, reproduction, or distribution of Your modifications, or\n    for any such Derivative Works as a whole, provided Your use,\n    reproduction, and distribution of the Work otherwise complies with\n    the conditions stated in this License.\n\n5.  Submission of Contributions. Unless You explicitly state otherwise,\n    any Contribution intentionally submitted for inclusion in the Work\n    by You to the Licensor shall be under the terms and conditions of\n    this License, without any additional terms or conditions.\n    Notwithstanding the above, nothing herein shall supersede or modify\n    the terms of any separate license agreement you may have executed\n    with Licensor regarding such Contributions.\n\n6.  Trademarks. This License does not grant permission to use the trade\n    names, trademarks, service marks, or product names of the Licensor,\n    except as required for reasonable and customary use in describing the\n    origin of the Work and reproducing the content of the NOTICE file.\n\n7.  Disclaimer of Warranty. Unless required by applicable law or\n    agreed to in writing, Licensor provides the Work (and each\n    Contributor provides its Contributions) on an \"AS IS\" BASIS,\n    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n    implied, including, without limitation, any warranties or conditions\n    of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n    PARTICULAR PURPOSE. You are solely responsible for determining the\n    appropriateness of using or redistributing the Work and assume any\n    risks associated with Your exercise of permissions under this License.\n\n8.  Limitation of Liability. In no event and under no legal theory,\n    whether in tort (including negligence), contract, or otherwise,\n    unless required by applicable law (such as deliberate and grossly\n    negligent acts) or agreed to in writing, shall any Contributor be\n    liable to You for damages, including any direct, indirect, special,\n    incidental, or consequential damages of any character arising as a\n    result of this License or out of the use or inability to use the\n    Work (including but not limited to damages for loss of goodwill,\n    work stoppage, computer failure or malfunction, or any and all\n    other commercial damages or losses), even if such Contributor\n    has been advised of the possibility of such damages.\n\n9.  Accepting Warranty or Additional Liability. While redistributing\n    the Work or Derivative Works thereof, You may choose to offer,\n    and charge a fee for, acceptance of support, warranty, indemnity,\n    or other liability obligations and/or rights consistent with this\n    License. However, in accepting such obligations, You may act only\n    on Your own behalf and on Your sole responsibility, not on behalf\n    of any other Contributor, and only if You agree to indemnify,\n    defend, and hold each Contributor harmless for any liability\n    incurred by, or claims asserted against, such Contributor by reason\n    of your accepting any such warranty or additional liability.\n\nEND OF TERMS AND CONDITIONS\n\nAPPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\nCopyright [yyyy] [name of copyright owner]\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n"
  },
  {
    "path": "README.md",
    "content": "# Stable Diffusion API Server\n\nA local inference REST API server for the [Stable Diffusion Photoshop plugin](https://christiancantrell.com/#ai-ml). (Also a generic Stable Diffusion REST API for whatever you want.)\n\nThe API server currently supports:\n\n1. Stable Diffusion weights automatically downloaded from Hugging Face.\n1. Custom fine-tuned models in the Hugging Face diffusers file format like those created with [DreamBooth](https://github.com/XavierXiao/Dreambooth-Stable-Diffusion).\n\n(Note that loading checkpoint files directly is not currently supported, but you can easily convert `.ckpt` files into the diffusers format using the aptly named [`convert_original_stable_diffusion_to_diffusers.py`](https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py) script.)\n\nThe server will run on Windows and Linux machines with NVIDIA GPUs, and on M1 Macs. M1 Mac support using MPS (Metal Performance Shaders) is highly experimental (and not easy to configure) but it does work, and it will get better over time.\n\nIf you can swing it, for best results, use a dedicated Linux box. Performance on Windows is also very good, but I recommend a dedicated machine with no other apps running. You can run Photoshop on the same machine if you have to, but you will be giving up some of your GPU memory which is good for the Photoshop user experience, but bad for optimal local inference.\n\n**Note that this project uses the content safety filter.**\n\n## Installation\n\n🤞 If anyone wants to make a detailed installation video, I would love to embed it right here. 🤞\n\n### Windows, Linux, and Mac Instructions\n\n1. Install Python.\n1. Install [Conda](https://conda.io/projects/conda/en/latest/user-guide/install/download.html).\n1. Download this repo.\n1. cd into the repo's directory.\n1. Set up a [Conda](https://conda.io) environment named `sd-api-server` by running the following command:\n\nWindows and Linux:\n\n**(Note that the '%' character below is meant to denote the command prompt; do not include it when copying and pasting.)**\n\n```\n% conda env create -f environment.yaml\n```\n\nM1 Macs:\n\n```\n% conda env create -f environment-m1.yaml\n```\n\nThen activate the Conda environment:\n\n```\n% conda activate sd-api-server\n```\n\nIf you are updating the server, make sure to update your Conda environment (using the platform-specific `yaml` file):\n\n```\n% conda env update -f environment.yaml\n% conda activate sd-api-server\n```\n\nIf you want to remove an old environment and create it from scratch (using the platform-specific `yaml` file):\n\n```\n% conda env remove -n sd-api-server\n% conda env create -f environment.yaml\n% conda activate sd-api-server\n```\n\n### Hugging Face Configuration\n\nThere are two things you need to configure with Hugging Face in order to run the Stable Diffusion model locally:\n\n1. You need to [agree to share your username and email address with Hugging Face](https://huggingface.co/CompVis/stable-diffusion-v1-4) in order to access the model.\n1. You also need to set up [a Hugging Face token](https://huggingface.co/settings/tokens). Once you've created a read-only token, copy and paste it into the `config.json` file as the value to the `hf_token` key (and don't forget to save the file).\n\nWindows and Linux users, you're good to go! All you have to do now is start the server:\n\n```\n% python3 server.py\n```\n\n### M1 Mac Additional Instructions\n\nNote that this is highly experimental, and may not work for you. But it will probably get easier with the next release of [PyTorch](https://pytorch.org/).\n\n#### Method 1: Nightly Builds\n\nIn Terminal, at the Conda prompt, **with the `sd-api-server` environment activated**:\n\n```\n% conda install pytorch torchvision -c pytorch-nightly\n% conda deactivate\n% conda activate sd-api-server\n% python3 server.py\n```\n\nYou might have noticed that you just installed a nightly build of PyTorch and Torchvision. Nightly builds come with neither warranties nor guarantees. If your server starts and you can generate images, you just won the nightly Lottery! If not, you can play again tomorrow. This is a temporary situation and probably won't be necessary with the next release of PyTorch.\n\n#### Method 2: Environment Variables\n\nIf the nightly build didn't work for you — or if you're simply allergic to nightly builds — you can tell PyTorch to use use the CPU in addition to MPS. If you already installed the nightly build, remove your Conda environment using the command above, start all over again, skip the nightly build step, and try this (with the `sd-api-server` environment active):\n\n```\n% conda env config vars set PYTORCH_ENABLE_MPS_FALLBACK=1\n% conda activate sd-api-server\n% python3 server.py\n```\n\nIf you get the message `ModuleNotFoundError: No module named 'flask'`, it probably means you're using the wrong version of Python. If you used `python3` then try `python`. If you used `python` then try `python3`. (These are the joys of old versions of Python being preinstalled on Macs.)\n\n## Configuring Custom Models\n\nIf you want to use the server (and the Photoshop plugin) with custom-trained models, the first thing you need are the custom-trained models themselves. Instructions for how to do so are beyond the scope of this README, but here are some resources:\n\n- [My custom fork of the DreamBooth repo](https://github.com/cantrell/Dreambooth-Stable-Diffusion-Tweaked) (dramatically simplified).\n- [A DreamBooth Stable Diffusion Colab notebook](https://colab.research.google.com/github/ShivamShrirao/diffusers/blob/main/examples/dreambooth/DreamBooth_Stable_Diffusion.ipynb) (much easier than training locally).\n- [A good YouTube tutorial on using the Colab notebook](https://www.youtube.com/watch?v=FaLTztGGueQ).\n- [The original DreamBooth paper](https://arxiv.org/abs/2208.12242).\n\nLoading checkpoint files directly is not currently supported, but you can easily convert `.ckpt` files into the diffusers format using the aptly named [`convert_original_stable_diffusion_to_diffusers.py`](https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py) script.\n\nOnce you have the models trained, the rest is easy. All you have to do is:\n\n1. Replace your `config.json` file with the `config-custom-models.json` template (rename `config-custom-models.json` to `config.json`).\n1. Make sure you copy and paste your Hugging Face token into the new `config.json` file.\n1. Fill in the `custom_model` array of the config file appropriately.\n\nHere's an explanation of what the key/value pairs mean:\n\n- `model_path`: The full path to the directory which contains the `model_index.json` file (just the directory; don't include the file itself). **Do not** escape spaces, but **do** escape backslashes with backslashes (e.g. `G:\\\\My Drive\\\\stable_diffusion_weights\\\\MyCustomModelOutput`).\n- `ui_label`: The name of the model as you want it to appear in the Photoshop plugin.\n- `url_path`: A unique, URL-friendly value that will be used as the endpoint path (see the REST API section below).\n- `requires_safety_checker`: Whether or not your custom model expects the safety checker. For models in the Hugging Face diffusers file format, this will be true; for models compiled from checkpoint files into the diffusers file format, this will probably be false.\n\nOnce your config file is ready, (re)start the server. If the Photoshop plugin is already loaded, you may need to restart it (or you can just click on the 'Reload Plugin' link in the lower right-hand corner of the 'Generate' tab).\n\nNote that the `custom_model` section of the `config.json` file is an array. That means you can include as many custom models as you want. Here's what it should look like for more than one custom-trained model:\n\n```\n{\n  \"hf_token\": \"your_hugging_face_token\",\n  \"custom_models\": [\n    {\n      \"model_path\": \"/path/to/directory/containing/model_index.json\",\n      \"ui_label\": \"My First Model\",\n      \"url_path\": \"my_first_model\",\n      \"requires_safety_checker\": true\n    }\n  ],\n  [\n    {\n      \"model_path\": \"/path/to/another/directory/containing/model_index.json\",\n      \"ui_label\": \"My Second Model\",\n      \"url_path\": \"my_second_model\",\n      \"requires_safety_checker\": true\n    }\n  ]\n}\n```\n\nTo see your custom models in the Generate tab of the Stable Diffusion Photoshop plugin, make sure you've configured your local inference server in the API Key tab.\n\n## REST API\n\nNote that all `POST` requests use the `application/x-www-form-urlencoded` content type, and all images are base64 encoded strings.\n\n`GET /ping`\n\n#### Response\n\n```\n{'status':'success'}\n```\n\n`GET /custom_models`\n\n#### Response\n\n```\n[\n  {\n    \"model_path\": \"/path/to/directory/containing/model_index.json\",\n    \"ui_label\": \"My First Model\",\n    \"url_path\": \"my_first_model\",\n    \"requires_safety_checker\": \"true | false\"\n  }\n], [...]\n```\n\n(If no custom models are configured, you will get back an empty array.)\n\n`POST /txt2img`\n\nParameters:\n\n- `prompt`: A text description.\n- `seed`: A numeric seed.\n- `num_outputs`: The number of images you want to get back.\n- `width`: The width of your results.\n- `height`: The height of your results.\n- `num_inference_steps`: The number of steps (more steps mean higher quality).\n- `guidance_scale`: Prompt strength.\n\n#### Response\n\n```\n{\n  'status':'success | failure',\n  'message':'Only if there was a failure',\n  'images': [\n    {\n      'base64': 'base64EncodedImage==',\n      'seed': 123456789,\n      'mimetype': 'image/png',\n      'nsfw': true | false\n    }\n  ]\n}\n```\n\n`POST /img2img`\n\nParameters:\n\n- `prompt`: A text description.\n- `seed`: A numeric seed.\n- `num_outputs`: The number of images you want to get back.\n- `num_inference_steps`: The number of steps (more steps mean higher quality).\n- `guidance_scale`: Prompt strength.\n- `init_image`: The initial input image.\n- `strength`: The image strength.\n\n#### Response\n\n```\n{\n  'status':'success | failure',\n  'message':'Only if there was a failure',\n  'images': [\n    {\n      'base64': 'base64EncodedImage==',\n      'seed': 123456789,\n      'mimetype': 'image/png',\n      'nsfw': true | false\n    }\n  ]\n}\n```\n\n`POST /masking`\n\nParameters:\n\n- `prompt`: A text description.\n- `seed`: A numeric seed.\n- `num_outputs`: The number of images you want to get back.\n- `num_inference_steps`: The number of steps (more steps mean higher quality).\n- `guidance_scale`: Prompt strength.\n- `init_image`: The initial input image.\n- `strength`: The image strength.\n- `mask_image`: A mask representing the pixels to replace.\n\n#### Response\n\n```\n{\n  'status':'success | failure',\n  'message':'Only if there was a failure',\n  'images': [\n    {\n      'base64': 'base64EncodedImage==',\n      'seed': 123456789,\n      'mimetype': 'image/png',\n      'nsfw': true | false\n    }\n  ]\n}\n```\n\n`POST /custom/<url_path>`\n\n`url_path` refers to the `url_path` key/value pair you defined in your `config.json` file.\n\nParameters:\n\n- `prompt`: A text description.\n- `seed`: A numeric seed.\n- `num_outputs`: The number of images you want to get back.\n- `width`: The width of your results.\n- `height`: The height of your results.\n- `num_inference_steps`: The number of steps (more steps mean higher quality).\n- `guidance_scale`: Prompt strength.\n\n#### Response\n\n```\n{\n  'status':'success | failure',\n  'message':'Only if there was a failure',\n  'images': [\n    {\n      'base64': 'base64EncodedImage==',\n      'seed': 123456789,\n      'mimetype': 'image/png',\n      'nsfw': true | false\n    }\n  ]\n}\n```\n"
  },
  {
    "path": "client_test.py",
    "content": "import json\nimport requests\nimport base64\nfrom PIL import Image\nfrom io import BytesIO\nimport matplotlib.pyplot as plt\n\ndef load_image_from_path(img_path):\n    img = Image.open( img_path )\n    return img\n\ndef load_image_from_url(img_url):\n    res = requests.get( img_url )\n    img = Image.open( BytesIO( res.content ) )\n    return img\n\ndef resize_image_preserve_aspect(img_pil, w):\n    wp = ( w / float( img_pil.size[0] ) )\n    hs = int( float( img_pil.size[1] ) * float( wp ) )\n    return img_pil.resize( ( w, hs ), Image.ANTIALIAS )\n\ndef pil_to_b64(input):\n    buffer = BytesIO()\n    input.save( buffer, 'PNG' )\n    output = base64.b64encode( buffer.getvalue() ).decode( 'utf-8' ).replace( '\\n', '' )\n    buffer.close()\n    return output\n\ndef b64_to_pil(input):\n    output = Image.open( BytesIO( base64.b64decode( input ) ) )\n    return output\n\ndef test_txt2img():\n    ENDPOINT = \"http://localhost:1337/txt2img\"\n    \n    data = {\n        'prompt':'a photo of a dog sitting on a bench',\n        'width':str( 512 ),\n        'height':str( 512 ),\n        'num_inference_steps':str( 100 ),\n        'guidance_scale':str( 7.5 ),\n        'num_outputs':str( 2 ),\n        'seed':str( 0 ),\n    }\n\n    response = json.loads( requests.post( url=ENDPOINT, data=data ).text )\n\n    def b64_to_pil(input):\n        output = Image.open( BytesIO( base64.b64decode( input ) ) )\n        return output\n\n    if 'status' in response:\n        if response[ 'status' ] == 'success':\n            images = response[ 'images' ]\n            for i, image in enumerate( images ):\n                plt.imshow( b64_to_pil( image['base64'] ) )\n                plt.show( block=True )\n                plt.pause( 10 )\n                plt.close()\n\ndef test_img2img():\n    ENDPOINT = \"http://localhost:1337/img2img\"    \n    IMG_URL  = 'https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg'\n\n    #IMG_URL  = 'https://minitravellers.co.uk/wp-content/uploads/2019/05/40612080213_81852c19fc_k.jpg'\n\n    data = {\n        'prompt':'a family of pixar characters on vacation in new york',\n        'init_image':pil_to_b64( resize_image_preserve_aspect( load_image_from_url( IMG_URL ).convert( 'RGB' ), 512 ) ),\n        'num_inference_steps':str( 100 ),\n        'guidance_scale':str( 7.5 ),\n        'num_outputs':str( 2 ),\n        'seed':str( 0 ),\n        'strength':str( 0.5 ),\n        'eta':str( 0.0 ),\n    }\n\n    response = json.loads( requests.post( url=ENDPOINT, data=data ).text )\n\n    if 'status' in response:\n        if response[ 'status' ] == 'success':\n            images = response[ 'images' ]\n            for i, image in enumerate( images ):\n                plt.imshow( b64_to_pil( image['base64'] ) )\n                plt.show( block=True )\n                plt.pause( 10 )\n                plt.close()\n\ndef test_inpaint():\n    ENDPOINT = \"http://localhost:1337/masking\"    \n    IMG_URL  = 'https://raw.githubusercontent.com/CompVis/stable-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png'\n    MSK_URL  = 'https://raw.githubusercontent.com/CompVis/stable-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png'\n\n    data = {\n        'prompt':'a cat sitting on a bench',\n        'init_image':pil_to_b64( resize_image_preserve_aspect( load_image_from_url( IMG_URL ).convert( 'RGB' ), 512 ) ),\n        'mask_image':pil_to_b64( resize_image_preserve_aspect( load_image_from_url( MSK_URL ).convert( 'RGB' ), 512 ) ),\n        'num_inference_steps':str( 100 ),\n        'guidance_scale':str( 7.5 ),\n        'num_outputs':str( 2 ),\n        'seed':str( 0 ),\n        'strength':str( 0.8 ),\n        'eta':str( 0.0 ),\n    }\n\n    response = json.loads( requests.post( url=ENDPOINT, data=data ).text )\n\n    if 'status' in response:\n        if response[ 'status' ] == 'success':\n            images = response[ 'images' ]\n            for i, image in enumerate( images ):\n                plt.imshow( b64_to_pil( image['base64'] ) )\n                plt.show( block=True )\n                plt.pause( 10 )\n                plt.close()\n\n\n# Run tests:\n\n# test_txt2img()\n# test_img2img()\ntest_inpaint()"
  },
  {
    "path": "config-custom-models.json",
    "content": "{\n  \"hf_token\": \"your_hugging_face_token\",\n  \"custom_models\": [\n    {\n      \"model_path\": \"/path/to/directory/containing/model_index.json\",\n      \"ui_label\": \"My Custom Model\",\n      \"url_path\": \"my_custom_model\",\n      \"requires_safety_checker\": true\n    }\n  ]\n}\n"
  },
  {
    "path": "config.json",
    "content": "{\n  \"hf_token\": \"your_hugging_face_token\"\n}\n"
  },
  {
    "path": "environment-m1.yaml",
    "content": "name: sd-api-server\nchannels:\n  - pytorch\n  - defaults\ndependencies:\n  - python=3.8.5\n  - pip=20.3\n  - pytorch=1.12.1\n  - torchvision=0.13.1\n  - numpy=1.19.2\n  - pip:\n      - diffusers\n      - transformers==4.19.2\n      - Pillow\n      - Flask\n"
  },
  {
    "path": "environment.yaml",
    "content": "name: sd-api-server\nchannels:\n  - pytorch\n  - defaults\ndependencies:\n  - python=3.8.5\n  - pip=20.3\n  - cudatoolkit=11.3\n  - pytorch=1.12.1\n  - torchvision=0.13.1\n  - numpy=1.19.2\n  - pip:\n      - diffusers\n      - transformers==4.19.2\n      - Pillow\n      - Flask\n"
  },
  {
    "path": "server.py",
    "content": "import re\nimport time\nimport inspect\nimport json\nimport flask\nimport sys\nimport base64\nfrom PIL import Image\nfrom io import BytesIO\n\nimport torch\nimport diffusers\n\n\n##################################################\n# Utils\n\ndef retrieve_param(key, data, cast, default):\n    if key in data:\n        value = flask.request.form[ key ]\n        value = cast( value )\n        return value\n    return default\n\ndef pil_to_b64(input):\n    buffer = BytesIO()\n    input.save( buffer, 'PNG' )\n    output = base64.b64encode( buffer.getvalue() ).decode( 'utf-8' ).replace( '\\n', '' )\n    buffer.close()\n    return output\n\ndef b64_to_pil(input):\n    output = Image.open( BytesIO( base64.b64decode( input ) ) )\n    return output\n\ndef get_compute_platform(context):\n    try:\n        import torch\n        if torch.cuda.is_available():\n            return 'cuda'\n        elif torch.backends.mps.is_available() and context == 'engine':\n            return 'mps'\n        else:\n            return 'cpu'\n    except ImportError:\n        return 'cpu'\n\n##################################################\n# Engines\n\nclass Engine(object):\n    def __init__(self):\n        pass\n\n    def process(self, kwargs):\n        return []\n\nclass EngineStableDiffusion(Engine):\n    def __init__(self, pipe, sibling=None, custom_model_path=None, requires_safety_checker=True):\n        super().__init__()\n        if sibling == None:\n            self.engine = pipe.from_pretrained( 'runwayml/stable-diffusion-v1-5', use_auth_token=hf_token.strip() )\n        elif custom_model_path:\n            if requires_safety_checker:\n                self.engine = diffusers.StableDiffusionPipeline.from_pretrained(custom_model_path,\n                                                                                safety_checker=sibling.engine.safety_checker,\n                                                                                feature_extractor=sibling.engine.feature_extractor)\n            else:\n                self.engine = diffusers.StableDiffusionPipeline.from_pretrained(custom_model_path,\n                                                                                feature_extractor=sibling.engine.feature_extractor)\n        else:\n            self.engine = pipe(\n                vae=sibling.engine.vae,\n                text_encoder=sibling.engine.text_encoder,\n                tokenizer=sibling.engine.tokenizer,\n                unet=sibling.engine.unet,\n                scheduler=sibling.engine.scheduler,\n                safety_checker=sibling.engine.safety_checker,\n                feature_extractor=sibling.engine.feature_extractor\n            )\n        self.engine.to( get_compute_platform('engine') )\n\n    def process(self, kwargs):\n        output = self.engine( **kwargs )\n        return {'image': output.images[0], 'nsfw':output.nsfw_content_detected[0]}\n\nclass EngineManager(object):\n    def __init__(self):\n        self.engines = {}\n\n    def has_engine(self, name):\n        return ( name in self.engines )\n\n    def add_engine(self, name, engine):\n        if self.has_engine( name ):\n            return False\n        self.engines[ name ] = engine\n        return True\n\n    def get_engine(self, name):\n        if not self.has_engine( name ):\n            return None\n        engine = self.engines[ name ]\n        return engine\n\n##################################################\n# App\n\n# Load and parse the config file:\ntry:\n    config_file = open ('config.json', 'r')\nexcept:\n    sys.exit('config.json not found.')\n\nconfig = json.loads(config_file.read())\n\nhf_token = config['hf_token']\n\nif (hf_token == None):\n    sys.exit('No Hugging Face token found in config.json.')\n\ncustom_models = config['custom_models'] if 'custom_models' in config else []\n\n# Initialize app:\napp = flask.Flask( __name__ )\n\n# Initialize engine manager:\nmanager = EngineManager()\n\n# Add supported engines to manager:\nmanager.add_engine( 'txt2img', EngineStableDiffusion( diffusers.StableDiffusionPipeline,        sibling=None ) )\nmanager.add_engine( 'img2img', EngineStableDiffusion( diffusers.StableDiffusionImg2ImgPipeline, sibling=manager.get_engine( 'txt2img' ) ) )\nmanager.add_engine( 'masking', EngineStableDiffusion( diffusers.StableDiffusionInpaintPipeline, sibling=manager.get_engine( 'txt2img' ) ) )\nfor custom_model in custom_models:\n    manager.add_engine( custom_model['url_path'],\n                        EngineStableDiffusion( diffusers.StableDiffusionPipeline, sibling=manager.get_engine( 'txt2img' ),\n                        custom_model_path=custom_model['model_path'],\n                        requires_safety_checker=custom_model['requires_safety_checker'] ) )\n\n# Define routes:\n@app.route('/ping', methods=['GET'])\ndef stable_ping():\n    return flask.jsonify( {'status':'success'} )\n\n@app.route('/custom_models', methods=['GET'])\ndef stable_custom_models():\n    if custom_models == None:\n        return flask.jsonify( [] )\n    else:\n        return custom_models\n\n@app.route('/txt2img', methods=['POST'])\ndef stable_txt2img():\n    return _generate('txt2img')\n\n@app.route('/img2img', methods=['POST'])\ndef stable_img2img():\n    return _generate('img2img')\n\n@app.route('/masking', methods=['POST'])\ndef stable_masking():\n    return _generate('masking')\n\n@app.route('/custom/<path:model>', methods=['POST'])\ndef stable_custom(model):\n    return _generate('txt2img', model)\n\ndef _generate(task, engine=None):\n    # Retrieve engine:\n    if engine == None:\n        engine = task\n\n    engine = manager.get_engine( engine )\n\n    # Prepare output container:\n    output_data = {}\n\n    # Handle request:\n    try:\n        seed = retrieve_param( 'seed', flask.request.form, int, 0 )\n        count = retrieve_param( 'num_outputs', flask.request.form, int,   1 )\n        total_results = []\n        for i in range( count ):\n            if (seed == 0):\n                generator = torch.Generator( device=get_compute_platform('generator') )\n            else:\n                generator = torch.Generator( device=get_compute_platform('generator') ).manual_seed( seed )\n            new_seed = generator.seed()\n            prompt = flask.request.form[ 'prompt' ]\n            args_dict = {\n                'prompt' : [ prompt ],\n                'num_inference_steps' : retrieve_param( 'num_inference_steps', flask.request.form, int,   100 ),\n                'guidance_scale' : retrieve_param( 'guidance_scale', flask.request.form, float, 7.5 ),\n                'eta' : retrieve_param( 'eta', flask.request.form, float, 0.0 ),\n                'generator' : generator\n            }\n            if (task == 'txt2img'):\n                args_dict[ 'width' ] = retrieve_param( 'width', flask.request.form, int,   512 )\n                args_dict[ 'height' ] = retrieve_param( 'height', flask.request.form, int,   512 )\n            if (task == 'img2img' or task == 'masking'):\n                init_img_b64 = flask.request.form[ 'init_image' ]\n                init_img_b64 = re.sub( '^data:image/png;base64,', '', init_img_b64 )\n                init_img_pil = b64_to_pil( init_img_b64 )\n                args_dict[ 'init_image' ] = init_img_pil\n                args_dict[ 'strength' ] = retrieve_param( 'strength', flask.request.form, float, 0.7 )\n            if (task == 'masking'):\n                mask_img_b64 = flask.request.form[ 'mask_image' ]\n                mask_img_b64 = re.sub( '^data:image/png;base64,', '', mask_img_b64 )\n                mask_img_pil = b64_to_pil( mask_img_b64 )\n                args_dict[ 'mask_image' ] = mask_img_pil\n            # Perform inference:\n            pipeline_output = engine.process( args_dict )\n            pipeline_output[ 'seed' ] = new_seed\n            total_results.append( pipeline_output )\n        # Prepare response\n        output_data[ 'status' ] = 'success'\n        images = []\n        for result in total_results:\n            images.append({\n                'base64' : pil_to_b64( result['image'].convert( 'RGB' ) ),\n                'seed' : result['seed'],\n                'mime_type': 'image/png',\n                'nsfw': result['nsfw']\n            })\n        output_data[ 'images' ] = images        \n    except RuntimeError as e:\n        output_data[ 'status' ] = 'failure'\n        output_data[ 'message' ] = 'A RuntimeError occurred. You probably ran out of GPU memory. Check the server logs for more details.'\n        print(str(e))\n    return flask.jsonify( output_data )\n\nif __name__ == '__main__':\n    app.run( host='0.0.0.0', port=1337, debug=False )"
  }
]