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