├── .gitignore ├── README.md ├── install.py ├── preload.py ├── requirements.txt └── scripts └── storage.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib64/ 18 | parts/ 19 | sdist/ 20 | var/ 21 | wheels/ 22 | pip-wheel-metadata/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | 53 | # Translations 54 | *.mo 55 | *.pot 56 | 57 | # Django stuff: 58 | *.log 59 | local_settings.py 60 | db.sqlite3 61 | db.sqlite3-journal 62 | 63 | # Flask stuff: 64 | instance/ 65 | .webassets-cache 66 | 67 | # Scrapy stuff: 68 | .scrapy 69 | 70 | # Sphinx documentation 71 | docs/_build/ 72 | 73 | # PyBuilder 74 | target/ 75 | 76 | # Jupyter Notebook 77 | .ipynb_checkpoints 78 | 79 | # IPython 80 | profile_default/ 81 | ipython_config.py 82 | 83 | # pyenv 84 | .python-version 85 | 86 | # pipenv 87 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 88 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 89 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 90 | # install all needed dependencies. 91 | #Pipfile.lock 92 | 93 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 94 | __pypackages__/ 95 | 96 | # Celery stuff 97 | celerybeat-schedule 98 | celerybeat.pid 99 | 100 | # SageMath parsed files 101 | *.sage.py 102 | 103 | # Environments 104 | .env 105 | .venv 106 | env/ 107 | venv/ 108 | ENV/ 109 | env.bak/ 110 | venv.bak/ 111 | 112 | # Spyder project settings 113 | .spyderproject 114 | .spyproject 115 | 116 | # Rope project settings 117 | .ropeproject 118 | 119 | # mkdocs documentation 120 | /site 121 | 122 | # mypy 123 | .mypy_cache/ 124 | .dmypy.json 125 | dmypy.json 126 | 127 | # Pyre type checker 128 | .pyre/ 129 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Introduction 2 | 3 | **db-storage1111 is an extension for [AUTOMATIC1111's Stable Diffusion Web UI](https://github.com/AUTOMATIC1111/stable-diffusion-webui).** 4 | 5 | It allows to store pictures to databases.At the moment it only supports MongoDB. 6 | ## Features 7 | 8 | - **Store images on a MongoDB instance** 9 | 10 | ## Installation 11 | 12 | 13 | 1. Visit the **Extensions** tab of Automatic's WebUI. 14 | 2. Visit the **Install from URL** subtab. 15 | 3. Paste this repo's URL into the first field: `https://github.com/takoyaro/db-storage1111 16 | 4. Click **Install**. 17 | 18 | ## Usage 19 | Set environment variables if needed before starting the app: 20 | | Variable | Default | 21 | |----------|---------------| 22 | | `DB_HOST` | `'localhost'` | 23 | | `DB_PORT` | `27017` | 24 | | `DB_USER` | `""` | 25 | | `DB_PASS` | `""` | 26 | 27 | Then, simply check the `Save to DB` checkbox and generate! 28 | 29 | 30 | ## Contributing 31 | I barely write any python, I have no doubt this extension could be improved and optimized. Feel free to submit PRs! 32 | -------------------------------------------------------------------------------- /install.py: -------------------------------------------------------------------------------- 1 | import launch 2 | 3 | if not launch.is_installed("pymongo"): 4 | launch.run_pip("install pymongo", "requirements for db-storage1111") 5 | 6 | if not launch.is_installed("tssplit"): 7 | launch.run_pip("install tssplit", "requirements for db-storage1111") -------------------------------------------------------------------------------- /preload.py: -------------------------------------------------------------------------------- 1 | def preload(parser): 2 | parser.add_argument("--db-provider", type=str, help="Database Provider", default="mongodb") -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pymongo -------------------------------------------------------------------------------- /scripts/storage.py: -------------------------------------------------------------------------------- 1 | import base64 2 | from io import BytesIO 3 | import os 4 | import re 5 | import modules.scripts as scripts 6 | import gradio as gr 7 | from pymongo import MongoClient 8 | import shlex 9 | from tssplit import tssplit 10 | 11 | mongo_host = os.environ.get('DB_HOST', 'localhost') 12 | mongo_port = int(os.environ.get('DB_PORT', 27017)) 13 | mongo_username = os.environ.get('DB_USER', '') 14 | mongo_password = os.environ.get('DB_PASS', '') 15 | 16 | creds = f"{mongo_username}:{mongo_password}@" if mongo_username and mongo_password else "" 17 | client = MongoClient(f"mongodb://{creds}{mongo_host}:{mongo_port}/") 18 | 19 | 20 | def get_collection(database_name, collection_name): 21 | db = client[database_name] 22 | collection = db[collection_name] 23 | return collection 24 | 25 | 26 | class Scripts(scripts.Script): 27 | def title(self): 28 | return "Mongo Storage" 29 | 30 | def show(self, is_img2img): 31 | return scripts.AlwaysVisible 32 | 33 | def ui(self, is_img2img): 34 | checkbox_save_to_db = gr.inputs.Checkbox(label="Save to DB", default=False) 35 | database_name = gr.inputs.Textbox(label="Database Name", default="StableDiffusion") 36 | collection_name = gr.inputs.Textbox(label="Collection Name", default="Automatic1111") 37 | return [checkbox_save_to_db, database_name, collection_name] 38 | 39 | def postprocess(self, p, processed, checkbox_save_to_db, database_name, collection_name): 40 | collection = get_collection(database_name, collection_name) if checkbox_save_to_db else None 41 | if collection is None: 42 | return True 43 | 44 | for i in range(len(processed.images)): 45 | # Extract image information 46 | regex = r"Steps:.*$" 47 | seed = processed.seed 48 | prompt = processed.prompt 49 | neg_prompt = processed.negative_prompt 50 | info = re.findall(regex, processed.info, re.M)[0] 51 | input_dict = dict(tssplit(item, quote='"', delimiter=':', trim=' ') for item in 52 | tssplit(info, quote='"', delimiter=',', trim=' ', quote_keep=True)) 53 | steps = int(input_dict["Steps"]) 54 | seed = int(input_dict["Seed"]) 55 | sampler = input_dict["Sampler"] 56 | cfg_scale = float(input_dict["CFG scale"]) 57 | size = tuple(map(int, input_dict["Size"].split("x"))) 58 | model_hash = input_dict["Model hash"] 59 | model = input_dict["Model"] 60 | lora_hashes = input_dict["Lora hashes"] 61 | 62 | image = processed.images[i] 63 | buffer = BytesIO() 64 | image.save(buffer, "png") 65 | image_bytes = buffer.getvalue() 66 | 67 | collection.insert_one({ 68 | "prompt": prompt, 69 | "negative_prompt": neg_prompt, 70 | "steps": int(steps), 71 | "seed": int(seed), 72 | "sampler": sampler, 73 | "cfg_scale": float(cfg_scale), 74 | "size": size, 75 | "model_hash": model_hash, 76 | "model": model, 77 | "lora_hashes": lora_hashes, 78 | "image": image_bytes 79 | }) 80 | return True 81 | --------------------------------------------------------------------------------