Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+extensions-builtin/sd-webui-agent-scheduler/docs/images/walkthrough.png filter=lfs diff=lfs merge=lfs -text
+extensions-builtin/stable-diffusion-webui-rembg/preview.png filter=lfs diff=lfs merge=lfs -text
+javascript/notosans-nerdfont-regular.ttf filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,68 @@
+# defaults
+__pycache__
+.ruff_cache
+/cache.json
+/*.json
+/*.yaml
+/params.txt
+/styles.csv
+/user.css
+/webui-user.bat
+/webui-user.sh
+/html/extensions.json
+/html/themes.json
+node_modules
+pnpm-lock.yaml
+package-lock.json
+venv
+.history
+cache
+**/.DS_Store
+
+# all models and temp files
+*.log
+*.log.*
+*.bak
+*.ckpt
+*.safetensors
+*.pth
+*.pt
+*.bin
+*.optim
+*.lock
+*.zip
+*.rar
+*.7z
+*.pyc
+/*.bat
+/*.sh
+/*.txt
+/*.mp3
+/*.lnk
+!webui.bat
+!webui.sh
+!package.json
+
+# all dynamic stuff
+/extensions/**/*
+/outputs/**/*
+/embeddings/**/*
+/models/**/*
+/interrogate/**/*
+/train/log/**/*
+/textual_inversion/**/*
+/detected_maps/**/*
+/tmp
+/log
+/cert
+.vscode/
+.idea/
+/localizations
+
+.*/
+
+# force included
+!/models/VAE-approx
+!/models/VAE-approx/model.pt
+!/models/Reference
+!/models/Reference/**/*

.gitmodules

@@ -0,0 +1,32 @@
+[submodule "wiki"]
+  path = wiki
+  url = https://github.com/vladmandic/automatic.wiki
+  ignore = dirty
+[submodule "modules/k-diffusion"]
+  path = modules/k-diffusion
+  url = https://github.com/crowsonkb/k-diffusion
+  ignore = dirty
+[submodule "extensions-builtin/sd-extension-system-info"]
+  path = extensions-builtin/sd-extension-system-info
+  url = https://github.com/vladmandic/sd-extension-system-info
+  ignore = dirty
+[submodule "extensions-builtin/sd-extension-chainner"]
+  path = extensions-builtin/sd-extension-chainner
+  url = https://github.com/vladmandic/sd-extension-chainner
+  ignore = dirty
+[submodule "extensions-builtin/stable-diffusion-webui-rembg"]
+  path = extensions-builtin/stable-diffusion-webui-rembg
+  url = https://github.com/vladmandic/sd-extension-rembg
+  ignore = dirty
+[submodule "extensions-builtin/stable-diffusion-webui-images-browser"]
+  path = extensions-builtin/stable-diffusion-webui-images-browser
+  url = https://github.com/AlUlkesh/stable-diffusion-webui-images-browser
+  ignore = dirty
+[submodule "extensions-builtin/sd-webui-controlnet"]
+  path = extensions-builtin/sd-webui-controlnet
+  url = https://github.com/Mikubill/sd-webui-controlnet
+  ignore = dirty
+[submodule "extensions-builtin/sd-webui-agent-scheduler"]
+  path = extensions-builtin/sd-webui-agent-scheduler
+  url = https://github.com/ArtVentureX/sd-webui-agent-scheduler
+  ignore = dirty

.pylintrc

@@ -0,0 +1,237 @@
+[MAIN]
+analyse-fallback-blocks=no
+clear-cache-post-run=no
+#enable-all-extensions=
+#errors-only=
+#exit-zero=
+extension-pkg-allow-list=
+extension-pkg-whitelist=
+fail-on=
+fail-under=10
+ignore=CVS
+ignore-paths=/usr/lib/.*$,
+             ^repositories/.*$,
+             ^extensions/.*$,
+             ^extensions-builtin/.*$,
+             ^modules/dml/.*$,
+             ^modules/tcd/.*$,
+             ^modules/xadapters/.*$,
+ignore-patterns=
+ignored-modules=
+jobs=0
+limit-inference-results=100
+load-plugins=
+persistent=yes
+py-version=3.10
+recursive=no
+source-roots=
+suggestion-mode=yes
+unsafe-load-any-extension=no
+#verbose=
+
+[BASIC]
+argument-naming-style=snake_case
+#argument-rgx=
+attr-naming-style=snake_case
+#attr-rgx=
+bad-names=foo, bar, baz, toto, tutu, tata
+bad-names-rgxs=
+class-attribute-naming-style=any
+class-const-naming-style=UPPER_CASE
+#class-const-rgx=
+class-naming-style=PascalCase
+#class-rgx=
+const-naming-style=snake_case
+#const-rgx=
+docstring-min-length=-1
+function-naming-style=snake_case
+#function-rgx=
+# Good variable names which should always be accepted, separated by a comma.
+good-names=i,j,k,e,ex,ok,p
+good-names-rgxs=
+include-naming-hint=no
+inlinevar-naming-style=any
+#inlinevar-rgx=
+method-naming-style=snake_case
+#method-rgx=
+module-naming-style=snake_case
+#module-rgx=
+name-group=
+no-docstring-rgx=^_
+property-classes=abc.abstractproperty
+#typealias-rgx=
+#typevar-rgx=
+variable-naming-style=snake_case
+#variable-rgx=
+
+[CLASSES]
+check-protected-access-in-special-methods=no
+defining-attr-methods=__init__,
+                      __new__,
+                      setUp,
+                      asyncSetUp,
+                      __post_init__
+exclude-protected=_asdict,_fields,_replace,_source,_make,os._exit
+valid-classmethod-first-arg=cls
+valid-metaclass-classmethod-first-arg=mcs
+
+[DESIGN]
+exclude-too-few-public-methods=
+ignored-parents=
+max-args=99
+max-attributes=99
+max-bool-expr=99
+max-branches=99
+max-locals=99
+max-parents=99
+max-public-methods=99
+max-returns=99
+max-statements=199
+min-public-methods=1
+
+[EXCEPTIONS]
+overgeneral-exceptions=builtins.BaseException,builtins.Exception
+
+[FORMAT]
+expected-line-ending-format=
+ignore-long-lines=^\s*(# )?<?https?://\S+>?$
+indent-after-paren=4
+indent-string='    '
+max-line-length=200
+max-module-lines=9999
+single-line-class-stmt=no
+single-line-if-stmt=no
+
+[IMPORTS]
+allow-any-import-level=
+allow-reexport-from-package=no
+allow-wildcard-with-all=no
+deprecated-modules=
+ext-import-graph=
+import-graph=
+int-import-graph=
+known-standard-library=
+known-third-party=enchant
+preferred-modules=
+
+[LOGGING]
+logging-format-style=new
+logging-modules=logging
+
+[MESSAGES CONTROL]
+confidence=HIGH,
+           CONTROL_FLOW,
+           INFERENCE,
+           INFERENCE_FAILURE,
+           UNDEFINED
+# disable=C,R,W
+disable=bad-inline-option,
+        bare-except,
+        broad-exception-caught,
+        chained-comparison,
+        consider-iterating-dictionary,
+        consider-using-dict-items,
+        consider-using-generator,
+        consider-using-enumerate,
+        consider-using-sys-exit,
+        consider-using-from-import,
+        consider-using-get,
+        consider-using-in,
+        consider-using-min-builtin,
+        dangerous-default-value,
+        deprecated-pragma,
+        duplicate-code,
+        file-ignored,
+        import-error,
+        import-outside-toplevel,
+        invalid-name,
+        line-too-long,
+        locally-disabled,
+        logging-fstring-interpolation,
+        missing-class-docstring,
+        missing-function-docstring,
+        missing-module-docstring,
+        no-else-return,
+        not-callable,
+        pointless-string-statement,
+        raw-checker-failed,
+        simplifiable-if-expression,
+        suppressed-message,
+        too-many-nested-blocks,
+        too-few-public-methods,
+        too-many-statements,
+        too-many-locals,
+        too-many-instance-attributes,
+        unnecessary-dunder-call,
+        unnecessary-lambda,
+        use-dict-literal,
+        use-symbolic-message-instead,
+        useless-suppression,
+        unidiomatic-typecheck,
+        wrong-import-position
+enable=c-extension-no-member
+
+[METHOD_ARGS]
+timeout-methods=requests.api.delete,requests.api.get,requests.api.head,requests.api.options,requests.api.patch,requests.api.post,requests.api.put,requests.api.request
+
+[MISCELLANEOUS]
+notes=FIXME,
+      XXX,
+      TODO
+notes-rgx=
+
+[REFACTORING]
+max-nested-blocks=5
+never-returning-functions=sys.exit,argparse.parse_error
+
+[REPORTS]
+evaluation=max(0, 0 if fatal else 10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10))
+msg-template=
+#output-format=
+reports=no
+score=no
+
+[SIMILARITIES]
+ignore-comments=yes
+ignore-docstrings=yes
+ignore-imports=yes
+ignore-signatures=yes
+min-similarity-lines=4
+
+[SPELLING]
+max-spelling-suggestions=4
+spelling-dict=
+spelling-ignore-comment-directives=fmt: on,fmt: off,noqa:,noqa,nosec,isort:skip,mypy:
+spelling-ignore-words=
+spelling-private-dict-file=
+spelling-store-unknown-words=no
+
+[STRING]
+check-quote-consistency=no
+check-str-concat-over-line-jumps=no
+
+[TYPECHECK]
+contextmanager-decorators=contextlib.contextmanager
+generated-members=numpy.*,logging.*,torch.*,cv2.*
+ignore-none=yes
+ignore-on-opaque-inference=yes
+ignored-checks-for-mixins=no-member,
+                          not-async-context-manager,
+                          not-context-manager,
+                          attribute-defined-outside-init
+ignored-classes=optparse.Values,thread._local,_thread._local,argparse.Namespace
+missing-member-hint=yes
+missing-member-hint-distance=1
+missing-member-max-choices=1
+mixin-class-rgx=.*[Mm]ixin
+signature-mutators=
+
+[VARIABLES]
+additional-builtins=
+allow-global-unused-variables=yes
+allowed-redefined-builtins=
+callbacks=cb_,
+dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_
+ignored-argument-names=_.*|^ignored_|^unused_
+init-import=no
+redefining-builtins-modules=six.moves,past.builtins,future.builtins,builtins,io

CITATION.cff

@@ -0,0 +1,28 @@
+cff-version: 1.2.0
+title: SD.Next
+url: 'https://github.com/vladmandic/automatic'
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file
+type: software
+authors:
+  - given-names: Vladimir
+    name-particle: Vlado
+    family-names: Mandic
+    orcid: 'https://orcid.org/0009-0003-4592-5074'
+identifiers:
+  - type: url
+    value: 'https://github.com/vladmandic'
+    description: GitHub
+  - type: url
+    value: 'https://www.linkedin.com/in/cyan051/'
+    description: LinkedIn
+repository-code: 'https://github.com/vladmandic/automatic'
+abstract: >-
+  SD.Next: Advanced Implementation of Stable Diffusion and
+  other diffusion models for text, image and video
+  generation
+keywords:
+  - stablediffusion diffusers sdnext
+license: AGPL-3.0
+date-released: 2022-12-24

README.md

@@ -1,12 +1,281 @@
 ---
-title: Test
-emoji: 👁
-colorFrom: green
-colorTo: purple
+title: test
+app_file: webui.py
 sdk: gradio
-sdk_version: 4.31.1
-app_file: app.py
-pinned: false
+sdk_version: 4.29.0
 ---
+<div align="center">
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SD.Next
+
+**Stable Diffusion implementation with advanced features**
+
+[![Sponsors](https://img.shields.io/static/v1?label=Sponsor&message=%E2%9D%A4&logo=GitHub&color=%23fe8e86)](https://github.com/sponsors/vladmandic)
+![Last Commit](https://img.shields.io/github/last-commit/vladmandic/automatic?svg=true)
+![License](https://img.shields.io/github/license/vladmandic/automatic?svg=true)
+[![Discord](https://img.shields.io/discord/1101998836328697867?logo=Discord&svg=true)](https://discord.gg/VjvR2tabEX)
+
+[Wiki](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.gg/VjvR2tabEX) | [Changelog](CHANGELOG.md)
+
+</div>
+</br>
+
+## Notable features
+
+All individual features are not listed here, instead check [ChangeLog](CHANGELOG.md) for full list of changes
+- Multiple backends!  
+  ▹ **Diffusers | Original**  
+- Multiple diffusion models!  
+  ▹ **Stable Diffusion 1.5/2.1 | SD-XL | LCM | Segmind | Kandinsky | Pixart-α | Stable Cascade | Würstchen | aMUSEd | DeepFloyd IF | UniDiffusion | SD-Distilled | BLiP Diffusion | KOALA | etc.**
+- Built-in Control for Text, Image, Batch and video processing!  
+  ▹ **ControlNet | ControlNet XS | Control LLLite | T2I Adapters | IP Adapters**  
+- Multiplatform!  
+ ▹ **Windows | Linux | MacOS with CPU | nVidia | AMD | IntelArc | DirectML | OpenVINO | ONNX+Olive | ZLUDA**
+- Platform specific autodetection and tuning performed on install
+- Optimized processing with latest `torch` developments with built-in support for `torch.compile`  
+  and multiple compile backends: *Triton, ZLUDA, StableFast, DeepCache, OpenVINO, NNCF, IPEX*  
+- Improved prompt parser  
+- Enhanced *Lora*/*LoCon*/*Lyco* code supporting latest trends in training  
+- Built-in queue management  
+- Enterprise level logging and hardened API  
+- Built in installer with automatic updates and dependency management  
+- Modernized UI with theme support and number of built-in themes *(dark and light)*  
+
+<br>
+
+*Main text2image interface*:  
+![Screenshot-Dark](html/screenshot-text2image.jpg)
+
+For screenshots and informations on other available themes, see [Themes Wiki](https://github.com/vladmandic/automatic/wiki/Themes)
+
+<br>
+
+## Backend support
+
+**SD.Next** supports two main backends: *Diffusers* and *Original*:
+
+- **Diffusers**: Based on new [Huggingface Diffusers](https://huggingface.co/docs/diffusers/index) implementation  
+  Supports *all* models listed below  
+  This backend is set as default for new installations  
+  See [wiki article](https://github.com/vladmandic/automatic/wiki/Diffusers) for more information  
+- **Original**: Based on [LDM](https://github.com/Stability-AI/stablediffusion) reference implementation and significantly expanded on by [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui)  
+  This backend and is fully compatible with most existing functionality and extensions written for *A1111 SDWebUI*  
+  Supports **SD 1.x** and **SD 2.x** models  
+  All other model types such as *SD-XL, LCM, PixArt, Segmind, Kandinsky, etc.* require backend **Diffusers**  
+
+## Model support
+
+Additional models will be added as they become available and there is public interest in them
+
+- [RunwayML Stable Diffusion](https://github.com/Stability-AI/stablediffusion/) 1.x and 2.x *(all variants)*  
+- [StabilityAI Stable Diffusion XL](https://github.com/Stability-AI/generative-models)  
+- [StabilityAI Stable Video Diffusion](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid) Base, XT 1.0, XT 1.1
+- [LCM: Latent Consistency Models](https://github.com/openai/consistency_models)  
+- [Playground](https://huggingface.co/playgroundai/playground-v2-256px-base) *v1, v2 256, v2 512, v2 1024 and latest v2.5*  
+- [Stable Cascade](https://github.com/Stability-AI/StableCascade) *Full* and *Lite*
+- [aMUSEd 256](https://huggingface.co/amused/amused-256) 256 and 512
+- [Segmind Vega](https://huggingface.co/segmind/Segmind-Vega)  
+- [Segmind SSD-1B](https://huggingface.co/segmind/SSD-1B)  
+- [Segmind SegMoE](https://github.com/segmind/segmoe) *SD and SD-XL*  
+- [Kandinsky](https://github.com/ai-forever/Kandinsky-2) *2.1 and 2.2 and latest 3.0*  
+- [PixArt-α XL 2](https://github.com/PixArt-alpha/PixArt-alpha) *Medium and Large*  
+- [Warp Wuerstchen](https://huggingface.co/blog/wuertschen)  
+- [Tsinghua UniDiffusion](https://github.com/thu-ml/unidiffuser)
+- [DeepFloyd IF](https://github.com/deep-floyd/IF) *Medium and Large*
+- [ModelScope T2V](https://huggingface.co/damo-vilab/text-to-video-ms-1.7b)
+- [Segmind SD Distilled](https://huggingface.co/blog/sd_distillation) *(all variants)*
+- [BLIP-Diffusion](https://dxli94.github.io/BLIP-Diffusion-website/)  
+- [KOALA 700M](https://github.com/youngwanLEE/sdxl-koala)
+- [VGen](https://huggingface.co/ali-vilab/i2vgen-xl)  
+
+
+Also supported are modifiers such as:
+- **LCM** and **Turbo** (*adversarial diffusion distillation*) networks
+- All **LoRA** types such as LoCon, LyCORIS, HADA, IA3, Lokr, OFT
+- **IP-Adapters** for SD 1.5 and SD-XL
+- **InstantID**, **FaceSwap**, **FaceID**, **PhotoMerge**  
+- **AnimateDiff** for SD 1.5
+
+## Examples
+
+*IP Adapters*:
+![Screenshot-IPAdapter](html/screenshot-ipadapter.jpg)
+
+*Color grading*:  
+![Screenshot-Color](html/screenshot-color.jpg)
+
+*InstantID*:  
+![Screenshot-InstantID](html/screenshot-instantid.jpg)
+
+> [!IMPORTANT]
+> - Loading any model other than standard SD 1.x / SD 2.x requires use of backend **Diffusers**  
+> - Loading any other models using **Original** backend is not supported  
+> - Loading manually download model `.safetensors` files is supported for specified models only (typically SD 1.x / SD 2.x / SD-XL models only)  
+> - For all other model types, use backend **Diffusers** and use built in Model downloader or  
+  select model from Networks -> Models -> Reference list in which case it will be auto-downloaded and loaded  
+
+## Platform support
+
+- *nVidia* GPUs using **CUDA** libraries on both *Windows and Linux*  
+- *AMD* GPUs using **ROCm** libraries on *Linux*  
+  Support will be extended to *Windows* once AMD releases ROCm for Windows  
+- *Intel Arc* GPUs using **OneAPI** with *IPEX XPU* libraries on both *Windows and Linux*  
+- Any GPU compatible with *DirectX* on *Windows* using **DirectML** libraries  
+  This includes support for AMD GPUs that are not supported by native ROCm libraries  
+- Any GPU or device compatible with **OpenVINO** libraries on both *Windows and Linux*  
+- *Apple M1/M2* on *OSX* using built-in support in Torch with **MPS** optimizations  
+- *ONNX/Olive*  
+
+## Install
+
+- [Step-by-step install guide](https://github.com/vladmandic/automatic/wiki/Installation)
+- [Advanced install notes](https://github.com/vladmandic/automatic/wiki/Advanced-Install)
+- [Common installation errors](https://github.com/vladmandic/automatic/discussions/1627)  
+- [FAQ](https://github.com/vladmandic/automatic/discussions/1011)
+- If you can't run us locally, try our friends at [RunDuffusion!](https://rundiffusion.com?utm_source=github&utm_medium=referral&utm_campaign=SDNext)
+
+> [!TIP]
+> - Server can run with or without virtual environment,  
+  Recommended to use `VENV` to avoid library version conflicts with other applications  
+> - **nVidia/CUDA** / **AMD/ROCm** / **Intel/OneAPI** are auto-detected if present and available,  
+  For any other use case such as **DirectML**, **ONNX/Olive**, **OpenVINO** specify required parameter explicitly  
+  or wrong packages may be installed as installer will assume CPU-only environment  
+> - Full startup sequence is logged in `sdnext.log`,  
+  so if you encounter any issues, please check it first  
+
+### Run
+
+Once SD.Next is installed, simply run `webui.ps1` or `webui.bat` (*Windows*) or `webui.sh` (*Linux or MacOS*)
+
+List of available parameters, run `webui --help` for the full & up-to-date list:
+
+    Server options:
+      --config CONFIG                                    Use specific server configuration file, default: config.json
+      --ui-config UI_CONFIG                              Use specific UI configuration file, default: ui-config.json
+      --medvram                                          Split model stages and keep only active part in VRAM, default: False
+      --lowvram                                          Split model components and keep only active part in VRAM, default: False
+      --ckpt CKPT                                        Path to model checkpoint to load immediately, default: None
+      --vae VAE                                          Path to VAE checkpoint to load immediately, default: None
+      --data-dir DATA_DIR                                Base path where all user data is stored, default:
+      --models-dir MODELS_DIR                            Base path where all models are stored, default: models
+      --allow-code                                       Allow custom script execution, default: False
+      --share                                            Enable UI accessible through Gradio site, default: False
+      --insecure                                         Enable extensions tab regardless of other options, default: False
+      --use-cpu USE_CPU [USE_CPU ...]                    Force use CPU for specified modules, default: []
+      --listen                                           Launch web server using public IP address, default: False
+      --port PORT                                        Launch web server with given server port, default: 7860
+      --freeze                                           Disable editing settings
+      --auth AUTH                                        Set access authentication like "user:pwd,user:pwd""
+      --auth-file AUTH_FILE                              Set access authentication using file, default: None
+      --autolaunch                                       Open the UI URL in the system's default browser upon launch
+      --docs                                             Mount API docs, default: False
+      --api-only                                         Run in API only mode without starting UI
+      --api-log                                          Enable logging of all API requests, default: False
+      --device-id DEVICE_ID                              Select the default CUDA device to use, default: None
+      --cors-origins CORS_ORIGINS                        Allowed CORS origins as comma-separated list, default: None
+      --cors-regex CORS_REGEX                            Allowed CORS origins as regular expression, default: None
+      --tls-keyfile TLS_KEYFILE                          Enable TLS and specify key file, default: None
+      --tls-certfile TLS_CERTFILE                        Enable TLS and specify cert file, default: None
+      --tls-selfsign                                     Enable TLS with self-signed certificates, default: False
+      --server-name SERVER_NAME                          Sets hostname of server, default: None
+      --no-hashing                                       Disable hashing of checkpoints, default: False
+      --no-metadata                                      Disable reading of metadata from models, default: False
+      --disable-queue                                    Disable queues, default: False
+      --subpath SUBPATH                                  Customize the URL subpath for usage with reverse proxy
+      --backend {original,diffusers}                     force model pipeline type
+      --allowed-paths ALLOWED_PATHS [ALLOWED_PATHS ...]  add additional paths to paths allowed for web access
+
+    Setup options:
+      --reset                                            Reset main repository to latest version, default: False
+      --upgrade                                          Upgrade main repository to latest version, default: False
+      --requirements                                     Force re-check of requirements, default: False
+      --quick                                            Bypass version checks, default: False
+      --use-directml                                     Use DirectML if no compatible GPU is detected, default: False
+      --use-openvino                                     Use Intel OpenVINO backend, default: False
+      --use-ipex                                         Force use Intel OneAPI XPU backend, default: False
+      --use-cuda                                         Force use nVidia CUDA backend, default: False
+      --use-rocm                                         Force use AMD ROCm backend, default: False
+      --use-zluda                                        Force use ZLUDA, AMD GPUs only, default: False
+      --use-xformers                                     Force use xFormers cross-optimization, default: False
+      --skip-requirements                                Skips checking and installing requirements, default: False
+      --skip-extensions                                  Skips running individual extension installers, default: False
+      --skip-git                                         Skips running all GIT operations, default: False
+      --skip-torch                                       Skips running Torch checks, default: False
+      --skip-all                                         Skips running all checks, default: False
+      --skip-env                                         Skips setting of env variables during startup, default: False
+      --experimental                                     Allow unsupported versions of libraries, default: False
+      --reinstall                                        Force reinstallation of all requirements, default: False
+      --test                                             Run test only and exit
+      --version                                          Print version information
+      --ignore                                           Ignore any errors and attempt to continue
+      --safe                                             Run in safe mode with no user extensions
+
+    Logging options:
+      --log LOG                                          Set log file, default: None
+      --debug                                            Run installer with debug logging, default: False
+      --profile                                          Run profiler, default: False
+
+## Notes
+
+### Control
+
+**SD.Next** comes with built-in control for all types of text2image, image2image, video2video and batch processing
+
+*Control interface*:  
+![Screenshot-Control](html/screenshot-control.jpg)
+
+*Control processors*:  
+![Screenshot-Process](html/screenshot-processors.jpg)
+
+*Masking*:
+![Screenshot-Mask](html/screenshot-mask.jpg)
+
+### **Extensions**
+
+SD.Next comes with several extensions pre-installed:
+
+- [ControlNet](https://github.com/Mikubill/sd-webui-controlnet) (*active in backend: original only*)
+- [Agent Scheduler](https://github.com/ArtVentureX/sd-webui-agent-scheduler)
+- [Image Browser](https://github.com/AlUlkesh/stable-diffusion-webui-images-browser)
+
+### **Collab**
+
+- We'd love to have additional maintainers (with comes with full repo rights). If you're interested, ping us!  
+- In addition to general cross-platform code, desire is to have a lead for each of the main platforms  
+This should be fully cross-platform, but we'd really love to have additional contributors and/or maintainers to join and help lead the efforts on different platforms  
+
+### **Credits**
+
+- Main credit goes to [Automatic1111 WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui) for original codebase  
+- Additional credits are listed in [Credits](https://github.com/AUTOMATIC1111/stable-diffusion-webui/#credits)  
+- Licenses for modules are listed in [Licenses](html/licenses.html)  
+
+### **Evolution**
+
+<a href="https://star-history.com/#vladmandic/automatic&Date">
+  <picture width=640>
+    <source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=vladmandic/automatic&type=Date&theme=dark" />
+    <img src="https://api.star-history.com/svg?repos=vladmandic/automatic&type=Date" alt="starts" width="320">
+  </picture>
+</a>
+
+- [OSS Stats](https://ossinsight.io/analyze/vladmandic/automatic#overview)
+
+### **Docs**
+
+If you're unsure how to use a feature, best place to start is [Wiki](https://github.com/vladmandic/automatic/wiki) and if its not there,  
+check [ChangeLog](CHANGELOG.md) for when feature was first introduced as it will always have a short note on how to use it  
+
+- [Wiki](https://github.com/vladmandic/automatic/wiki)
+- [ReadMe](README.md)
+- [ToDo](TODO.md)  
+- [ChangeLog](CHANGELOG.md)
+- [CLI Tools](cli/README.md)
+
+### **Sponsors**
+
+<div align="center">
+<!-- sponsors --><a href="https://github.com/allangrant"><img src="https://github.com/allangrant.png" width="60px" alt="Allan Grant" /></a><a href="https://github.com/BrentOzar"><img src="https://github.com/BrentOzar.png" width="60px" alt="Brent Ozar" /></a><a href="https://github.com/inktomi"><img src="https://github.com/inktomi.png" width="60px" alt="Matthew Runo" /></a><a href="https://github.com/4joeknight4"><img src="https://github.com/4joeknight4.png" width="60px" alt="" /></a><a href="https://github.com/SaladTechnologies"><img src="https://github.com/SaladTechnologies.png" width="60px" alt="Salad Technologies" /></a><a href="https://github.com/mantzaris"><img src="https://github.com/mantzaris.png" width="60px" alt="a.v.mantzaris" /></a><a href="https://github.com/CurseWave"><img src="https://github.com/CurseWave.png" width="60px" alt="" /></a><!-- sponsors -->
+</div>
+
+<br>

SECURITY.md

@@ -0,0 +1,36 @@
+# Security & Privacy Policy
+
+<br>
+
+## Issues
+
+All issues are tracked publicly on GitHub: <https://github.com/vladmandic/automatic/issues>
+
+<br>
+
+## Vulnerabilities
+
+`SD.Next` code base and included dependencies are automatically scanned against known security vulnerabilities
+
+Any code commit is validated before merge
+
+- [Dependencies](https://github.com/vladmandic/automatic/security/dependabot)
+- [Scanning Alerts](https://github.com/vladmandic/automatic/security/code-scanning)
+
+<br>
+
+## Privacy
+
+`SD.Next` app:
+
+- Is fully self-contained and does not send or share data of any kind with external targets
+- Does not store any user or system data tracking, user provided inputs (images, video) or detection results
+- Does not utilize any analytic services (such as Google Analytics)
+
+`SD.Next` library can establish external connections *only* for following purposes and *only* when explicitly configured by user:
+
+- Download extensions and themes indexes from automatically updated indexes  
+- Download required packages and repositories from GitHub during installation/upgrade
+- Download installed/enabled extensions
+- Download models from CivitAI and/or Huggingface when instructed by user
+- Submit benchmark info upon user interaction  

TODO.md

@@ -0,0 +1,20 @@
+# TODO
+
+Main ToDo list can be found at [GitHub projects](https://github.com/users/vladmandic/projects)
+
+## Candidates for next release
+
+- defork
+- stable diffusion 3.0
+- ipadapter masking: <https://github.com/huggingface/diffusers/pull/6847>
+- x-adapter: <https://github.com/showlab/X-Adapter>
+- async lowvram: <https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14855>
+- init latents: variations, img2img
+- diffusers public callbacks  
+- remove builtin: controlnet
+- remove builtin: image-browser
+
+## Control missing features
+
+- second pass: <https://github.com/vladmandic/automatic/issues/2783>  
+- control api  

cli/README.md

@@ -0,0 +1,108 @@
+# Stable-Diffusion Productivity Scripts
+
+Note: All scripts have built-in `--help` parameter that can be used to get more information
+
+<br>
+
+## Main Scripts
+
+### Generate
+
+Text-to-image with all of the possible parameters  
+Supports upsampling, face restoration and grid creation  
+> python generate.py
+
+By default uses parameters from  `generate.json`
+
+Parameters that are not specified will be randomized:
+
+- Prompt will be dynamically created from template of random samples: `random.json`
+- Sampler/Scheduler will be randomly picked from available ones
+- CFG Scale set to 5-10
+
+### Train
+
+Combined pipeline for **embeddings**, **lora**, **lycoris**, **dreambooth** and **hypernetwork**  
+Optionally runs several image processing steps before training:
+
+- keep original image
+- detect and extract face
+- detect and extract body
+- detect blur
+- detect dynamic range
+- attempt to upscale low resolution images
+- attempt to restore quality of low quality images
+- automatically generate captions using interrogate
+- resize image
+- square image
+- run image segmentation to remove background
+
+> python train.py
+
+<br>
+
+## Auxiliary Scripts
+
+### Benchmark
+
+> python run-benchmark.py
+
+### Create Previews
+
+Create previews for **embeddings**, **lora**, **lycoris**, **dreambooth** and **hypernetwork**
+
+> python create-previews.py
+
+## Image Grid
+
+> python image-grid.py
+
+### Image Watermark
+
+Create invisible image watermark and remove existing EXIF tags  
+
+> python image-watermark.py
+
+### Image Interrogate
+
+Runs CLiP and Booru image interrogation  
+
+> python image-interrogate.py
+
+### Palette Extract
+
+Extract color palette from image(s)  
+
+> python image-palette.py
+
+### Prompt Ideas
+
+Generate complex prompt ideas
+
+> python prompt-ideas.py
+
+### Prompt Promptist
+
+Attempts to beautify the provided prompt  
+
+> python prompt-promptist.py
+
+### Video Extract
+
+Extract frames from video files  
+
+> python video-extract.py
+
+<br>
+
+## Utility Scripts
+
+### SDAPI
+
+Utility module that handles async communication to Automatic API endpoints  
+Note: Requires SD API  
+
+Can be used to manually execute specific commands:
+> python sdapi.py progress  
+> python sdapi.py interrupt
+> python sdapi.py shutdown

cli/clone.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+import os
+import logging
+import git
+from rich import console, progress
+
+
+class GitRemoteProgress(git.RemoteProgress):
+    OP_CODES = ["BEGIN", "CHECKING_OUT", "COMPRESSING", "COUNTING", "END", "FINDING_SOURCES", "RECEIVING", "RESOLVING", "WRITING"]
+    OP_CODE_MAP = { getattr(git.RemoteProgress, _op_code): _op_code for _op_code in OP_CODES }
+
+    def __init__(self, url, folder) -> None:
+        super().__init__()
+        self.url = url
+        self.folder = folder
+        self.progressbar = progress.Progress(
+            progress.SpinnerColumn(),
+            progress.TextColumn("[cyan][progress.description]{task.description}"),
+            progress.BarColumn(),
+            progress.TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            progress.TimeRemainingColumn(),
+            progress.TextColumn("[yellow]<{task.fields[url]}>"),
+            progress.TextColumn("{task.fields[message]}"),
+            console=console.Console(),
+            transient=False,
+        )
+        self.progressbar.start()
+        self.active_task = None
+
+    def __del__(self) -> None:
+        self.progressbar.stop()
+
+    @classmethod
+    def get_curr_op(cls, op_code: int) -> str:
+        op_code_masked = op_code & cls.OP_MASK
+        return cls.OP_CODE_MAP.get(op_code_masked, "?").title()
+
+    def update(self, op_code: int, cur_count: str | float, max_count: str | float | None = None, message: str | None = "") -> None:
+        if op_code & self.BEGIN:
+            self.curr_op = self.get_curr_op(op_code) # pylint: disable=attribute-defined-outside-init
+            self.active_task = self.progressbar.add_task(description=self.curr_op, total=max_count, message=message, url=self.url)
+        self.progressbar.update(task_id=self.active_task, completed=cur_count, message=message)
+        if op_code & self.END:
+            self.progressbar.update(task_id=self.active_task, message=f"[bright_black]{message}")
+
+
+def clone(url: str, folder: str):
+    git.Repo.clone_from(
+        url=url,
+        to_path=folder,
+        progress=GitRemoteProgress(url=url, folder=folder),
+        multi_options=['--config core.compression=0', '--config core.loosecompression=0', '--config pack.window=0'],
+        allow_unsafe_options=True,
+        depth=1,
+        )
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description = 'downloader')
+    parser.add_argument('--url', required=True, help="download url, required")
+    parser.add_argument('--folder', required=False, help="output folder, default: autodetect")
+    args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
+    log = logging.getLogger(__name__)
+    try:
+        if not args.url.startswith('http'):
+            raise ValueError(f'invalid url: {args.url}')
+        f = args.url.split('/')[-1].split('.')[0] if args.folder is None else args.folder
+        if os.path.exists(f):
+            raise FileExistsError(f'folder already exists: {f}')
+        log.info(f'Clone start: url={args.url} folder={f}')
+        clone(url=args.url, folder=f)
+        log.info(f'Clone complete: url={args.url} folder={f}')
+    except KeyboardInterrupt:
+        log.warning(f'Clone cancelled: url={args.url} folder={f}')
+    except Exception as e:
+        log.error(f'Clone: url={args.url} {e}')

cli/create-previews.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python
+# pylint: disable=no-member
+import os
+import re
+import json
+import time
+import logging
+import importlib
+import asyncio
+import argparse
+from pathlib import Path
+from util import Map, log
+from sdapi import get, post, close
+from generate import generate # pylint: disable=import-error
+grid = importlib.import_module('image-grid').grid
+
+
+options = Map({
+    # used by extra networks
+    'prompt': 'photo of <keyword> <embedding>, photograph, posing, pose, high detailed, intricate, elegant, sharp focus, skin texture, looking forward, facing camera, 135mm, shot on dslr, canon 5d, 4k, modelshoot style, cinematic lighting',
+    # used by models
+    'prompts': [
+        ('photo citiscape', 'cityscape during night, photorealistic, high detailed, sharp focus, depth of field, 4k'),
+        ('photo car', 'photo of a sports car, high detailed, sharp focus, dslr, cinematic lighting, realistic'),
+        ('photo woman', 'portrait photo of beautiful woman, high detailed, dslr, 35mm'),
+        ('photo naked', 'full body photo of beautiful sexy naked woman, high detailed, dslr, 35mm'),
+
+        ('photo taylor', 'portrait photo of beautiful woman taylor swift, high detailed, sharp focus, depth of field, dslr, 35mm <lora:taylor-swift:1>'),
+        ('photo ti-mia', 'portrait photo of beautiful woman "ti-mia", naked, high detailed, dslr, 35mm'),
+        ('photo ti-vlado', 'portrait photo of man "ti-vlado", high detailed, dslr, 35mm'),
+        ('photo lora-vlado', 'portrait photo of man vlado, high detailed, dslr, 35mm <lora:vlado-original:1>'),
+
+        ('wlop', 'a stunning portrait of sexy teen girl in a wet t-shirt, vivid color palette, digital painting, octane render, highly detailed, particles, light effect, volumetric lighting, art by wlop'),
+        ('greg rutkowski', 'beautiful woman, high detailed, sharp focus, depth of field, 4k, art by greg rutkowski'),
+        ('carne griffiths', 'beautiful woman taylor swift, high detailed, sharp focus, depth of field, art by carne griffiths <lora:taylor-swift:1>'),
+        ('carne griffiths', 'man vlado, high detailed, sharp focus, depth of field, art by carne griffiths <lora:vlado-full:1>'),
+    ],
+    # save format
+    'format': '.jpg',
+    # used by generate script
+    'paths': {
+        "root": "/mnt/c/Users/mandi/OneDrive/Generative/Generate",
+        "generate": "image",
+        "upscale": "upscale",
+        "grid": "grid",
+    },
+    # generate params
+    'generate': {
+        'restore_faces': True,
+        'prompt': '',
+        'negative_prompt': 'foggy, blurry, blurred, duplicate, ugly, mutilated, mutation, mutated, out of frame, bad anatomy, disfigured, deformed, censored, low res, low resolution, watermark, text, poorly drawn face, poorly drawn hands, signature',
+        'steps': 20,
+        'batch_size': 2,
+        'n_iter': 1,
+        'seed': -1,
+        'sampler_name': 'UniPC',
+        'cfg_scale': 6,
+        'width': 512,
+        'height': 512,
+    },
+    'lora': {
+        'strength': 1.0,
+    },
+    'hypernetwork': {
+        'keyword': '',
+        'strength': 1.0,
+    },
+})
+
+
+def preview_exists(folder, model):
+    model = os.path.splitext(model)[0]
+    for suffix in ['', '.preview']:
+        for ext in ['.jpg', '.png', '.webp']:
+            fn = os.path.join(folder, f'{model}{suffix}{ext}')
+            if os.path.exists(fn):
+                return True
+    return False
+
+
+async def preview_models(params):
+    data = await get('/sdapi/v1/sd-models')
+    allmodels = [m['title'] for m in data]
+    models = []
+    excluded = []
+    for m in allmodels: # loop through all registered models
+        ok = True
+        for e in params.exclude: # check if model is excluded
+            if e in m:
+                excluded.append(m)
+                ok = False
+                break
+        if ok:
+            short = m.split(' [')[0]
+            short = short.replace('.ckpt', '').replace('.safetensors', '')
+            models.append(short)
+    if len(params.input) > 0: # check if model is included in cmd line
+        filtered = []
+        for m in params.input:
+            if m in models:
+                filtered.append(m)
+            else:
+                log.error({ 'model not found': m })
+                return
+        models = filtered
+    log.info({ 'models preview' })
+    log.info({ 'models': len(models), 'excluded': len(excluded) })
+    opt = await get('/sdapi/v1/options')
+    log.info({ 'total jobs': len(models) * options.generate.batch_size, 'per-model': options.generate.batch_size })
+    log.info(json.dumps(options, indent=2))
+    for model in models:
+        if preview_exists(opt['ckpt_dir'], model) and len(params.input) == 0: # if model preview exists and not manually included
+            log.info({ 'model preview exists': model })
+            continue
+        fn = os.path.join(opt['ckpt_dir'], os.path.splitext(model)[0] + options.format)
+        log.info({ 'model load': model })
+
+        opt['sd_model_checkpoint'] = model
+        del opt['sd_lora']
+        del opt['sd_lyco']
+        await post('/sdapi/v1/options', opt)
+        opt = await get('/sdapi/v1/options')
+        images = []
+        labels = []
+        t0 = time.time()
+        for label, p in options.prompts:
+            options.generate.prompt = p
+            log.info({ 'model generating': model, 'label': label, 'prompt': options.generate.prompt })
+            data = await generate(options = options, quiet=True)
+            if 'image' in data:
+                for img in data['image']:
+                    images.append(img)
+                    labels.append(label)
+            else:
+                log.error({ 'model': model, 'error': data })
+        t1 = time.time()
+        if len(images) == 0:
+            log.error({ 'model': model, 'error': 'no images generated' })
+            continue
+        image = grid(images = images, labels = labels, border = 8)
+        log.info({ 'saving preview': fn, 'images': len(images), 'size': [image.width, image.height] })
+        image.save(fn)
+        t = t1 - t0
+        its = 1.0 * options.generate.steps * len(images) / t
+        log.info({ 'model preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
+
+    opt = await get('/sdapi/v1/options')
+    if opt['sd_model_checkpoint'] != params.model:
+        log.info({ 'model set default': params.model })
+        opt['sd_model_checkpoint'] = params.model
+        del opt['sd_lora']
+        del opt['sd_lyco']
+        await post('/sdapi/v1/options', opt)
+
+
+async def lora(params):
+    opt = await get('/sdapi/v1/options')
+    folder = opt['lora_dir']
+    if not os.path.exists(folder):
+        log.error({ 'lora directory not found': folder })
+        return
+    models1 = list(Path(folder).glob('**/*.safetensors'))
+    models2 = list(Path(folder).glob('**/*.ckpt'))
+    models = [os.path.splitext(f)[0] for f in models1 + models2]
+    log.info({ 'loras': len(models) })
+    for model in models:
+        if preview_exists('', model) and len(params.input) == 0: # if model preview exists and not manually included
+            log.info({ 'lora preview exists': model })
+            continue
+        fn = model + options.format
+        model = os.path.basename(model)
+        images = []
+        labels = []
+        t0 = time.time()
+        keywords = re.sub(r'\d', '', model)
+        keywords = keywords.replace('-v', ' ').replace('-', ' ').strip().split(' ')
+        keyword = '\"' + '\" \"'.join(keywords) + '\"'
+        options.generate.prompt = options.prompt.replace('<keyword>', keyword)
+        options.generate.prompt = options.generate.prompt.replace('<embedding>', '')
+        options.generate.prompt += f' <lora:{model}:{options.lora.strength}>'
+        log.info({ 'lora generating': model, 'keyword': keyword, 'prompt': options.generate.prompt })
+        data = await generate(options = options, quiet=True)
+        if 'image' in data:
+            for img in data['image']:
+                images.append(img)
+                labels.append(keyword)
+        else:
+            log.error({ 'lora': model, 'keyword': keyword, 'error': data })
+        t1 = time.time()
+        if len(images) == 0:
+            log.error({ 'model': model, 'error': 'no images generated' })
+            continue
+        image = grid(images = images, labels = labels, border = 8)
+        log.info({ 'saving preview': fn, 'images': len(images), 'size': [image.width, image.height] })
+        image.save(fn)
+        t = t1 - t0
+        its = 1.0 * options.generate.steps * len(images) / t
+        log.info({ 'lora preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
+
+
+async def lyco(params):
+    opt = await get('/sdapi/v1/options')
+    folder = opt['lyco_dir']
+    if not os.path.exists(folder):
+        log.error({ 'lyco directory not found': folder })
+        return
+    models1 = list(Path(folder).glob('**/*.safetensors'))
+    models2 = list(Path(folder).glob('**/*.ckpt'))
+    models = [os.path.splitext(f)[0] for f in models1 + models2]
+    log.info({ 'lycos': len(models) })
+    for model in models:
+        if preview_exists('', model) and len(params.input) == 0: # if model preview exists and not manually included
+            log.info({ 'lyco preview exists': model })
+            continue
+        fn = model + options.format
+        model = os.path.basename(model)
+        images = []
+        labels = []
+        t0 = time.time()
+        keywords = re.sub(r'\d', '', model)
+        keywords = keywords.replace('-v', ' ').replace('-', ' ').strip().split(' ')
+        keyword = '\"' + '\" \"'.join(keywords) + '\"'
+        options.generate.prompt = options.prompt.replace('<keyword>', keyword)
+        options.generate.prompt = options.generate.prompt.replace('<embedding>', '')
+        options.generate.prompt += f' <lyco:{model}:{options.lora.strength}>'
+        log.info({ 'lyco generating': model, 'keyword': keyword, 'prompt': options.generate.prompt })
+        data = await generate(options = options, quiet=True)
+        if 'image' in data:
+            for img in data['image']:
+                images.append(img)
+                labels.append(keyword)
+        else:
+            log.error({ 'lyco': model, 'keyword': keyword, 'error': data })
+        t1 = time.time()
+        if len(images) == 0:
+            log.error({ 'model': model, 'error': 'no images generated' })
+            continue
+        image = grid(images = images, labels = labels, border = 8)
+        log.info({ 'saving preview': fn, 'images': len(images), 'size': [image.width, image.height] })
+        image.save(fn)
+        t = t1 - t0
+        its = 1.0 * options.generate.steps * len(images) / t
+        log.info({ 'lyco preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
+
+
+async def hypernetwork(params):
+    opt = await get('/sdapi/v1/options')
+    folder = opt['hypernetwork_dir']
+    if not os.path.exists(folder):
+        log.error({ 'hypernetwork directory not found': folder })
+        return
+    models = [os.path.splitext(f)[0] for f in Path(folder).glob('**/*.pt')]
+    log.info({ 'hypernetworks': len(models) })
+    for model in models:
+        if preview_exists(folder, model) and len(params.input) == 0: # if model preview exists and not manually included
+            log.info({ 'hypernetwork preview exists': model })
+            continue
+        fn = os.path.join(folder, model + options.format)
+        images = []
+        labels = []
+        t0 = time.time()
+        keyword = options.hypernetwork.keyword
+        options.generate.prompt = options.prompt.replace('<keyword>', options.hypernetwork.keyword)
+        options.generate.prompt = options.generate.prompt.replace('<embedding>', '')
+        options.generate.prompt = f' <hypernet:{model}:{options.hypernetwork.strength}> ' + options.generate.prompt
+        log.info({ 'hypernetwork generating': model, 'keyword': keyword, 'prompt': options.generate.prompt })
+        data = await generate(options = options, quiet=True)
+        if 'image' in data:
+            for img in data['image']:
+                images.append(img)
+                labels.append(keyword)
+        else:
+            log.error({ 'hypernetwork': model, 'keyword': keyword, 'error': data })
+        t1 = time.time()
+        if len(images) == 0:
+            log.error({ 'model': model, 'error': 'no images generated' })
+            continue
+        image = grid(images = images, labels = labels, border = 8)
+        log.info({ 'saving preview': fn, 'images': len(images), 'size': [image.width, image.height] })
+        image.save(fn)
+        t = t1 - t0
+        its = 1.0 * options.generate.steps * len(images) / t
+        log.info({ 'hypernetwork preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
+
+
+async def embedding(params):
+    opt = await get('/sdapi/v1/options')
+    folder = opt['embeddings_dir']
+    if not os.path.exists(folder):
+        log.error({ 'embeddings directory not found': folder })
+        return
+    models = [os.path.splitext(f)[0] for f in Path(folder).glob('**/*.pt')]
+    log.info({ 'embeddings': len(models) })
+    for model in models:
+        if preview_exists(folder, model) and len(params.input) == 0: # if model preview exists and not manually included
+            log.info({ 'embedding preview exists': model })
+            continue
+        fn = os.path.join(folder, model + '.preview' + options.format)
+        images = []
+        labels = []
+        t0 = time.time()
+        keyword = '\"' + re.sub(r'\d', '', model) + '\"'
+        options.generate.batch_size = 4
+        options.generate.prompt = options.prompt.replace('<keyword>', keyword)
+        options.generate.prompt = options.generate.prompt.replace('<embedding>', '')
+        log.info({ 'embedding generating': model, 'keyword': keyword, 'prompt': options.generate.prompt })
+        data = await generate(options = options, quiet=True)
+        if 'image' in data:
+            for img in data['image']:
+                images.append(img)
+                labels.append(keyword)
+        else:
+            log.error({ 'embeding': model, 'keyword': keyword, 'error': data })
+        t1 = time.time()
+        if len(images) == 0:
+            log.error({ 'model': model, 'error': 'no images generated' })
+            continue
+        image = grid(images = images, labels = labels, border = 8)
+        log.info({ 'saving preview': fn, 'images': len(images), 'size': [image.width, image.height] })
+        image.save(fn)
+        t = t1 - t0
+        its = 1.0 * options.generate.steps * len(images) / t
+        log.info({ 'embeding preview created': model, 'image': fn, 'images': len(images), 'grid': [image.width, image.height], 'time': round(t, 2), 'its': round(its, 2) })
+
+
+async def create_previews(params):
+    await preview_models(params)
+    await lora(params)
+    await lyco(params)
+    await hypernetwork(params)
+    await embedding(params)
+    await close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'generate model previews')
+    parser.add_argument('--model', default='best/icbinp-icantbelieveIts-final.safetensors [73f48afbdc]', help="model used to create extra network previews")
+    parser.add_argument('--exclude', default=['sd-v20', 'sd-v21', 'inpainting', 'pix2pix'], help="exclude models with keywords")
+    parser.add_argument('--debug', default = False, action='store_true', help = 'print extra debug information')
+    parser.add_argument('input', type = str, nargs = '*')
+    args = parser.parse_args()
+    if args.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': args.__dict__ })
+    asyncio.run(create_previews(args))

cli/download.py

@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+import os
+import time
+import argparse
+import tempfile
+import urllib
+import requests
+import urllib3
+import rich.progress as p
+from rich import print # pylint: disable=redefined-builtin
+
+
+pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn())
+headers = {
+    'Content-type': 'application/json',
+    'User-Agent': 'Mozilla/5.0',
+}
+
+
+def get_filename(args, res):
+    content_fn = (res.headers.get('content-disposition', '').split('filename=')[1]).strip().strip('\"') if 'filename=' in res.headers.get('content-disposition', '') else None
+    return args.file or content_fn or next(tempfile._get_candidate_names()) # pylint: disable=protected-access
+
+
+def download_requests(args):
+    res = requests.get(args.url, timeout=30, headers=headers, verify=False, allow_redirects=True, stream=True)
+    content_length = int(res.headers.get('content-length', 0))
+    fn = get_filename(args, res)
+    print(f'downloading: url={args.url} file={fn} size={content_length if content_length > 0 else "unknown"} lib=requests block={args.block}')
+    with open(fn, 'wb') as f:
+        with pbar:
+            task = pbar.add_task(description="Download starting", total=content_length)
+            for data in res.iter_content(args.block):
+                f.write(data)
+                pbar.update(task, advance=args.block, description="Downloading")
+    return fn
+
+
+def download_urllib(args):
+    fn = ''
+    req = urllib.request.Request(args.url, headers=headers)
+    res = urllib.request.urlopen(req)
+    res.getheader('content-length')
+    content_length = int(res.getheader('content-length') or 0)
+    fn = get_filename(args, res)
+    print(f'downloading: url={args.url} file={fn} size={content_length if content_length > 0 else "unknown"} lib=urllib block={args.block}')
+    with open(fn, 'wb') as f:
+        with pbar:
+            task = pbar.add_task(description="Download starting", total=content_length)
+            while True:
+                buf = res.read(args.block)
+                if not buf:
+                    break
+                f.write(buf)
+                pbar.update(task, advance=args.block, description="Downloading")
+    return fn
+
+
+def download_urllib3(args):
+    http_pool = urllib3.PoolManager()
+    res = http_pool.request('GET', args.url, preload_content=False, headers=headers)
+    fn = get_filename(args, res)
+    content_length = int(res.headers.get('content-length', 0))
+    print(f'downloading: url={args.url} file={fn} size={content_length if content_length > 0 else "unknown"} lib=urllib3 block={args.block}')
+    with open(fn, 'wb') as f:
+        with pbar:
+            task = pbar.add_task(description="Download starting", total=content_length)
+            while True:
+                buf = res.read(args.block)
+                if not buf:
+                    break
+                f.write(buf)
+                pbar.update(task, advance=args.block, description="Downloading")
+    return fn
+
+
+def download_httpx(args):
+    try:
+        import httpx
+    except ImportError:
+        print('httpx is not installed')
+        return None
+    with httpx.stream("GET", args.url, headers=headers, verify=False, follow_redirects=True) as res:
+        fn = get_filename(args, res)
+        content_length = int(res.headers.get('content-length', 0))
+        print(f'downloading: url={args.url} file={fn} size={content_length if content_length > 0 else "unknown"} lib=httpx block=internal')
+        with open(fn, 'wb') as f:
+            with pbar:
+                task = pbar.add_task(description="Download starting", total=content_length)
+                for buf in res.iter_bytes():
+                    f.write(buf)
+                    pbar.update(task, advance=args.block, description="Downloading")
+        return fn
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'downloader')
+    parser.add_argument('--url', required=True, help="download url, required")
+    parser.add_argument('--file', required=False, help="output file, default: autodetect")
+    parser.add_argument('--lib', required=False, default='requests', choices=['urllib', 'urllib3', 'requests', 'httpx'], help="download mode, default: %(default)s")
+    parser.add_argument('--block', required=False, type=int, default=16384, help="download block size, default: %(default)s")
+    parsed = parser.parse_args()
+    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+    try:
+        t0 = time.time()
+        if parsed.lib == 'requests':
+            filename = download_requests(parsed)
+        elif parsed.lib == 'urllib':
+            filename = download_urllib(parsed)
+        elif parsed.lib == 'urllib3':
+            filename = download_urllib3(parsed)
+        elif parsed.lib == 'httpx':
+            filename = download_httpx(parsed)
+        else:
+            print(f'unknown download library: {parsed.lib}')
+            exit(1)
+        t1 = time.time()
+        if filename is None:
+            print(f'download error: args={parsed}')
+            exit(1)
+        speed = round(os.path.getsize(filename) / (t1 - t0) / 1024 / 1024, 3)
+        print(f'download complete: url={parsed.url} file={filename} speed={speed} mb/s')
+    except KeyboardInterrupt:
+        print(f'download cancelled: args={parsed}')
+    except Exception as e:
+        print(f'download error: args={parsed} {e}')

cli/gen-styles.py

@@ -0,0 +1,79 @@
+#!/bin/env python
+
+import io
+import json
+import base64
+import argparse
+import requests
+from PIL import Image
+
+
+options = {
+    "negative_prompt": "",
+    "steps": 20,
+    "batch_size": 1,
+    "n_iter": 1,
+    "seed": -1,
+    "sampler_name": "UniPC",
+    "cfg_scale": 6,
+    "width": 512,
+    "height": 512,
+    "save_images": False,
+    "send_images": True,
+}
+styles = []
+
+
+def pil_to_b64(img: Image, size: int, quality: int):
+    img = img.convert('RGB')
+    img = img.resize((size, size))
+    buffer = io.BytesIO()
+    img.save(buffer, format="JPEG", quality=quality)
+    b64encoded = base64.b64encode(buffer.getvalue()).decode("utf-8")
+    return f'data:image/jpeg;base64,{b64encoded}'
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(endpoint, json = dct, timeout=300, verify=False)
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'gen-styles.py')
+    parser.add_argument('--input', type=str, required=True, help="input text file with one line per prompt")
+    parser.add_argument('--output', type=str, required=True, help="output json file")
+    parser.add_argument('--nopreviews', default=False, action='store_true', help = 'generate previews')
+    parser.add_argument('--prompt', type=str, required=False, default='girl walking in a city', help="applied prompt when generating previews")
+    parser.add_argument('--size', type=int, default=128, help="image size for previews")
+    parser.add_argument('--quality', type=int, default=35, help="image quality for previews")
+    parser.add_argument('--url', type=str, required=False, default='http://127.0.0.1:7860', help="sd.next server url")
+    args = parser.parse_args()
+    with open(args.input, encoding='utf-8') as f:
+        lines = f.readlines()
+    for line in lines:
+        line = line.strip().replace('\n', '')
+        if len(line) == 0:
+            continue
+        print(f'processing: {line}')
+        if not args.nopreviews:
+            options['prompt'] = f'{line} {args.prompt}'
+            data = post(f'{args.url}/sdapi/v1/txt2img', options)
+            if 'error' in data:
+                print(f'error: {data}')
+                continue
+            b64str = data['images'][0].split(',',1)[0]
+            image = Image.open(io.BytesIO(base64.b64decode(b64str)))
+        else:
+            image = None
+        styles.append({
+            'name': line,
+            'prompt': line + ' {prompt}',
+            'negative': '',
+            'extra': '',
+            'preview': pil_to_b64(image, args.size, args.quality) if image is not None else '',
+        })
+        with open(args.output, 'w', encoding='utf-8') as outfile:
+            json.dump(styles, outfile, indent=2)

cli/generate.json

@@ -0,0 +1,38 @@
+{
+  "paths":
+  {
+      "root": "/mnt/c/Users/mandi/OneDrive/Generative/Generate",
+      "generate": "image",
+      "upscale": "upscale",
+      "grid": "grid"
+  },
+  "generate":
+  {
+      "restore_faces": true,
+      "prompt": "dynamic",
+      "negative_prompt": "foggy, blurry, blurred, duplicate, ugly, mutilated, mutation, mutated, out of frame, bad anatomy, disfigured, deformed, censored, low res, watermark, text, poorly drawn face, signature",
+      "steps": 30,
+      "batch_size": 2,
+      "n_iter": 1,
+      "seed": -1,
+      "sampler_name": "DPM2 Karras",
+      "cfg_scale": 6,
+      "width": 512,
+      "height": 512
+  },
+  "upscale":
+  {
+      "upscaler_1": "SwinIR_4x",
+      "upscaler_2": "None",
+      "upscale_first": false,
+      "upscaling_resize": 0,
+      "gfpgan_visibility": 0,
+      "codeformer_visibility": 0,
+      "codeformer_weight": 0.5
+  },
+  "options":
+  {
+      "sd_model_checkpoint": "sd-v15-runwayml",
+      "sd_vae": "vae-ft-mse-840000-ema-pruned.ckpt"
+  }
+}

cli/generate.py

@@ -0,0 +1,373 @@
+#!/usr/bin/env python
+# pylint: disable=no-member
+"""generate batches of images from prompts and upscale them
+
+params: run with `--help`
+
+default workflow runs infinite loop and prints stats when interrupted:
+1. choose random scheduler lookup all available and pick one
+2. generate dynamic prompt based on styles, embeddings, places, artists, suffixes
+3. beautify prompt
+4. generate 3x3 images
+5. create image grid
+6. upscale images with face restoration
+"""
+
+import argparse
+import asyncio
+import base64
+import io
+import json
+import logging
+import math
+import os
+import pathlib
+import secrets
+import time
+import sys
+import importlib
+
+from random import randrange
+from PIL import Image
+from PIL.ExifTags import TAGS
+from PIL.TiffImagePlugin import ImageFileDirectory_v2
+
+from sdapi import close, get, interrupt, post, session
+from util import Map, log, safestring
+
+
+sd = {}
+random = {}
+stats = Map({ 'images': 0, 'wall': 0, 'generate': 0, 'upscale': 0 })
+avg = {}
+
+
+def grid(data):
+    if len(data.image) > 1:
+        w, h = data.image[0].size
+        rows = round(math.sqrt(len(data.image)))
+        cols = math.ceil(len(data.image) / rows)
+        image = Image.new('RGB', size = (cols * w, rows * h), color = 'black')
+        for i, img in enumerate(data.image):
+            image.paste(img, box=(i % cols * w, i // cols * h))
+        short = data.info.prompt[:min(len(data.info.prompt), 96)] # limit prompt part of filename to 96 chars
+        name = '{seed:0>9} {short}'.format(short = short, seed = data.info.all_seeds[0]) # pylint: disable=consider-using-f-string
+        name = safestring(name) + '.jpg'
+        f = os.path.join(sd.paths.root, sd.paths.grid, name)
+        log.info({ 'grid': { 'name': f, 'size': image.size, 'images': len(data.image) } })
+        image.save(f, 'JPEG', exif = exif(data.info, None, 'grid'), optimize = True, quality = 70)
+        return image
+    return data.image
+
+
+def exif(info, i = None, op = 'generate'):
+    seed = [info.all_seeds[i]] if len(info.all_seeds) > 0 and i is not None else info.all_seeds # always returns list
+    seed = ', '.join([str(x) for x in seed]) # int list to str list to single str
+    template = '{prompt} | negative {negative_prompt} | seed {s} | steps {steps} | cfgscale {cfg_scale} | sampler {sampler_name} | batch {batch_size} | timestamp {job_timestamp} | model {model} | vae {vae}'.format(s = seed, model = sd.options['sd_model_checkpoint'], vae = sd.options['sd_vae'], **info) # pylint: disable=consider-using-f-string
+    if op == 'upscale':
+        template += ' | faces gfpgan' if sd.upscale.gfpgan_visibility > 0 else ''
+        template += ' | faces codeformer' if sd.upscale.codeformer_visibility > 0 else ''
+        template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_1) if sd.upscale.upscaler_1 != 'None' else '' # pylint: disable=consider-using-f-string
+        template += ' | upscale {resize}x {upscaler}'.format(resize = sd.upscale.upscaling_resize, upscaler = sd.upscale.upscaler_2) if sd.upscale.upscaler_2 != 'None' else '' # pylint: disable=consider-using-f-string
+    if op == 'grid':
+        template += ' | grid {num}'.format(num = sd.generate.batch_size * sd.generate.n_iter) # pylint: disable=consider-using-f-string
+    ifd = ImageFileDirectory_v2()
+    exif_stream = io.BytesIO()
+    _TAGS = {v: k for k, v in TAGS.items()} # enumerate possible exif tags
+    ifd[_TAGS['ImageDescription']] = template
+    ifd.save(exif_stream)
+    val = b'Exif\x00\x00' + exif_stream.getvalue()
+    return val
+
+
+def randomize(lst):
+    if len(lst) > 0:
+        return secrets.choice(lst)
+    else:
+        return ''
+
+
+def prompt(params): # generate dynamic prompt or use one if provided
+    sd.generate.prompt = params.prompt if params.prompt != 'dynamic' else randomize(random.prompts)
+    sd.generate.negative_prompt = params.negative if params.negative != 'dynamic' else randomize(random.negative)
+    embedding = params.embedding if params.embedding != 'random' else randomize(random.embeddings)
+    sd.generate.prompt = sd.generate.prompt.replace('<embedding>', embedding)
+    artist = params.artist if params.artist != 'random' else randomize(random.artists)
+    sd.generate.prompt = sd.generate.prompt.replace('<artist>', artist)
+    style = params.style if params.style != 'random' else randomize(random.styles)
+    sd.generate.prompt = sd.generate.prompt.replace('<style>', style)
+    suffix = params.suffix if params.suffix != 'random' else randomize(random.suffixes)
+    sd.generate.prompt = sd.generate.prompt.replace('<suffix>', suffix)
+    place = params.suffix if params.suffix != 'random' else randomize(random.places)
+    sd.generate.prompt = sd.generate.prompt.replace('<place>', place)
+    if params.prompts or params.debug:
+        log.info({ 'random initializers': random })
+    if params.prompt == 'dynamic':
+        log.info({ 'dynamic prompt': sd.generate.prompt })
+    return sd.generate.prompt
+
+
+def sampler(params, options): # find sampler
+    if params.sampler == 'random':
+        sd.generate.sampler_name = randomize(options.samplers)
+        log.info({ 'random sampler': sd.generate.sampler_name })
+    else:
+        found = [i for i in options.samplers if i.startswith(params.sampler)]
+        if len(found) == 0:
+            log.error({ 'sampler error': sd.generate.sampler_name, 'available': options.samplers})
+            exit()
+        sd.generate.sampler_name = found[0]
+    return sd.generate.sampler_name
+
+
+async def generate(prompt = None, options = None, quiet = False): # pylint: disable=redefined-outer-name
+    global sd # pylint: disable=global-statement
+    if options:
+        sd = Map(options)
+    if prompt is not None:
+        sd.generate.prompt = prompt
+    if not quiet:
+        log.info({ 'generate': sd.generate })
+    if sd.get('options', None) is None:
+        sd['options'] = await get('/sdapi/v1/options')
+    names = []
+    b64s = []
+    images = []
+    info = Map({})
+    data = await post('/sdapi/v1/txt2img', sd.generate)
+    if 'error' in data:
+        log.error({ 'generate': data['error'], 'reason': data['reason'] })
+        return Map({})
+    info = Map(json.loads(data['info']))
+    log.debug({ 'info': info })
+    images = data['images']
+    short = info.prompt[:min(len(info.prompt), 96)] # limit prompt part of filename to 64 chars
+    for i in range(len(images)):
+        b64s.append(images[i])
+        images[i] = Image.open(io.BytesIO(base64.b64decode(images[i].split(',',1)[0])))
+        name = '{seed:0>9} {short}'.format(short = short, seed = info.all_seeds[i]) # pylint: disable=consider-using-f-string
+        name = safestring(name) + '.jpg'
+        f = os.path.join(sd.paths.root, sd.paths.generate, name)
+        names.append(f)
+        if not quiet:
+            log.info({ 'image': { 'name': f, 'size': images[i].size } })
+        images[i].save(f, 'JPEG', exif = exif(info, i), optimize = True, quality = 70)
+    return Map({ 'name': names, 'image': images, 'b64': b64s, 'info': info })
+
+
+async def upscale(data):
+    data.upscaled = []
+    if sd.upscale.upscaling_resize <=1:
+        return data
+    sd.upscale.image = ''
+    log.info({ 'upscale': sd.upscale })
+    for i in range(len(data.image)):
+        f = data.name[i].replace(sd.paths.generate, sd.paths.upscale)
+        sd.upscale.image = data.b64[i]
+        res = await post('/sdapi/v1/extra-single-image', sd.upscale)
+        image = Image.open(io.BytesIO(base64.b64decode(res['image'].split(',',1)[0])))
+        data.upscaled.append(image)
+        log.info({ 'image': { 'name': f, 'size': image.size } })
+        image.save(f, 'JPEG', exif = exif(data.info, i, 'upscale'), optimize = True, quality = 70)
+    return data
+
+
+async def init():
+    '''
+    import torch
+    log.info({ 'torch': torch.__version__, 'available': torch.cuda.is_available() })
+    current_device = torch.cuda.current_device()
+    mem_free, mem_total = torch.cuda.mem_get_info()
+    log.info({ 'cuda': torch.version.cuda, 'available': torch.cuda.is_available(), 'arch': torch.cuda.get_arch_list(), 'device': torch.cuda.get_device_name(current_device), 'memory': { 'free': round(mem_free / 1024 / 1024), 'total': (mem_total / 1024 / 1024) } })
+    '''
+    options = Map({})
+    options.flags = await get('/sdapi/v1/cmd-flags')
+    log.debug({ 'flags': options.flags })
+    data = await get('/sdapi/v1/sd-models')
+    options.models = [obj['title'] for obj in data]
+    log.debug({ 'registered models': options.models })
+    found = sd.options.sd_model_checkpoint if sd.options.sd_model_checkpoint in options.models else None
+    if found is None:
+        found = [i for i in options.models if i.startswith(sd.options.sd_model_checkpoint)]
+    if len(found) == 0:
+        log.error({ 'model error': sd.generate.sd_model_checkpoint, 'available': options.models})
+        exit()
+    sd.options.sd_model_checkpoint = found[0]
+    data = await get('/sdapi/v1/samplers')
+    options.samplers = [obj['name'] for obj in data]
+    log.debug({ 'registered samplers': options.samplers })
+    data = await get('/sdapi/v1/upscalers')
+    options.upscalers = [obj['name'] for obj in data]
+    log.debug({ 'registered upscalers': options.upscalers })
+    data = await get('/sdapi/v1/face-restorers')
+    options.restorers = [obj['name'] for obj in data]
+    log.debug({ 'registered face restorers': options.restorers })
+    await interrupt()
+    await post('/sdapi/v1/options', sd.options)
+    options.options = await get('/sdapi/v1/options')
+    log.info({ 'target models': { 'diffuser': options.options['sd_model_checkpoint'], 'vae': options.options['sd_vae'] } })
+    log.info({ 'paths': sd.paths })
+    options.queue = await get('/queue/status')
+    log.info({ 'queue': options.queue })
+    pathlib.Path(sd.paths.root).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.generate)).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.upscale)).mkdir(parents = True, exist_ok = True)
+    pathlib.Path(os.path.join(sd.paths.root, sd.paths.grid)).mkdir(parents = True, exist_ok = True)
+    return options
+
+
+def args(): # parse cmd arguments
+    global sd # pylint: disable=global-statement
+    global random # pylint: disable=global-statement
+    parser = argparse.ArgumentParser(description = 'sd pipeline')
+    parser.add_argument('--config', type = str, default = 'generate.json', required = False, help = 'configuration file')
+    parser.add_argument('--random', type = str, default = 'random.json', required = False, help = 'prompt file with randomized sections')
+    parser.add_argument('--max', type = int, default = 1, required = False, help = 'maximum number of generated images')
+    parser.add_argument('--prompt', type = str, default = 'dynamic', required = False, help = 'prompt')
+    parser.add_argument('--negative', type = str, default = 'dynamic', required = False, help = 'negative prompt')
+    parser.add_argument('--artist', type = str, default = 'random', required = False, help = 'artist style, used to guide dynamic prompt when prompt is not provided')
+    parser.add_argument('--embedding', type = str, default = 'random', required = False, help = 'use embedding, used to guide dynamic prompt when prompt is not provided')
+    parser.add_argument('--style', type = str, default = 'random', required = False, help = 'image style, used to guide dynamic prompt when prompt is not provided')
+    parser.add_argument('--suffix', type = str, default = 'random', required = False, help = 'style suffix, used to guide dynamic prompt when prompt is not provided')
+    parser.add_argument('--place', type = str, default = 'random', required = False, help = 'place locator, used to guide dynamic prompt when prompt is not provided')
+    parser.add_argument('--faces', default = False, action='store_true', help = 'restore faces during upscaling')
+    parser.add_argument('--steps', type = int, default = 0, required = False, help = 'number of steps')
+    parser.add_argument('--batch', type = int, default = 0, required = False, help = 'batch size, limited by gpu vram')
+    parser.add_argument('--n', type = int, default = 0, required = False, help = 'number of iterations')
+    parser.add_argument('--cfg', type = int, default = 0, required = False, help = 'classifier free guidance scale')
+    parser.add_argument('--sampler', type = str, default = 'random', required = False, help = 'sampler')
+    parser.add_argument('--seed', type = int, default = 0, required = False, help = 'seed, default is random')
+    parser.add_argument('--upscale', type = int, default = 0, required = False, help = 'upscale factor, disabled if 0')
+    parser.add_argument('--model', type = str, default = '', required = False, help = 'diffusion model')
+    parser.add_argument('--vae', type = str, default = '', required = False, help = 'vae model')
+    parser.add_argument('--path', type = str, default = '', required = False, help = 'output path')
+    parser.add_argument('--width', type = int, default = 0, required = False, help = 'width')
+    parser.add_argument('--height', type = int, default = 0, required = False, help = 'height')
+    parser.add_argument('--beautify', default = False, action='store_true', help = 'beautify prompt')
+    parser.add_argument('--prompts', default = False, action='store_true', help = 'print dynamic prompt templates')
+    parser.add_argument('--debug', default = False, action='store_true', help = 'print extra debug information')
+    params = parser.parse_args()
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    home = pathlib.Path(sys.argv[0]).parent
+    if os.path.isfile(params.config):
+        try:
+            with open(params.config, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                sd = Map(data)
+                log.debug({ 'config': sd })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    elif os.path.isfile(os.path.join(home, params.config)):
+        try:
+            with open(os.path.join(home, params.config), 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                sd = Map(data)
+                log.debug({ 'config': sd })
+        except Exception as e:
+            log.error({ 'config error': params.config, 'exception': e })
+            exit()
+    else:
+        log.error({ 'config file not found': params.config})
+        exit()
+    if params.prompt == 'dynamic':
+        log.info({ 'prompt template': params.random })
+        if os.path.isfile(params.random):
+            try:
+                with open(params.random, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    random = Map(data)
+                    log.debug({ 'random template': sd })
+            except Exception:
+                log.error({ 'random template error': params.random})
+                exit()
+        elif os.path.isfile(os.path.join(home, params.random)):
+            try:
+                with open(os.path.join(home, params.random), 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    random = Map(data)
+                    log.debug({ 'random template': sd })
+            except Exception:
+                log.error({ 'random template error': params.random})
+                exit()
+        else:
+            log.error({ 'random template file not found': params.random})
+            exit()
+        _dynamic = prompt(params)
+
+    sd.paths.root = params.path if params.path != '' else sd.paths.root
+    sd.generate.restore_faces = params.faces if params.faces is not None else sd.generate.restore_faces
+    sd.generate.seed = params.seed if params.seed > 0 else sd.generate.seed
+    sd.generate.sampler_name = params.sampler if params.sampler != 'random' else sd.generate.sampler_name
+    sd.generate.batch_size = params.batch if params.batch > 0 else sd.generate.batch_size
+    sd.generate.cfg_scale = params.cfg if params.cfg > 0 else sd.generate.cfg_scale
+    sd.generate.n_iter = params.n if params.n > 0 else sd.generate.n_iter
+    sd.generate.width = params.width if params.width > 0 else sd.generate.width
+    sd.generate.height = params.height if params.height > 0 else sd.generate.height
+    sd.generate.steps = params.steps if params.steps > 0 else sd.generate.steps
+    sd.upscale.upscaling_resize = params.upscale if params.upscale > 0 else sd.upscale.upscaling_resize
+    sd.upscale.codeformer_visibility = 1 if params.faces else sd.upscale.codeformer_visibility
+    sd.options.sd_vae = params.vae if params.vae != '' else sd.options.sd_vae
+    sd.options.sd_model_checkpoint = params.model if params.model != '' else sd.options.sd_model_checkpoint
+    sd.upscale.upscaler_1 = 'SwinIR_4x' if params.upscale > 1 else sd.upscale.upscaler_1
+    if sd.generate.cfg_scale == 0:
+        sd.generate.cfg_scale = randrange(5, 10)
+    return params
+
+
+async def main():
+    params = args()
+    sess = await session()
+    if sess is None:
+        await close()
+        exit()
+    options = await init()
+    iteration = 0
+    while True:
+        iteration += 1
+        log.info('')
+        log.info({ 'iteration': iteration, 'batch': sd.generate.batch_size, 'n': sd.generate.n_iter, 'total': sd.generate.n_iter * sd.generate.batch_size })
+        dynamic = prompt(params)
+        if params.beautify:
+            try:
+                promptist = importlib.import_module('modules.promptist')
+                sd.generate.prompt = promptist.beautify(dynamic)
+            except Exception as e:
+                log.error({ 'beautify': e })
+        scheduler = sampler(params, options)
+        t0 = time.perf_counter()
+        data = await generate() # generate returns list of images
+        if 'image' not in data:
+            break
+        stats.images += len(data.image)
+        t1 = time.perf_counter()
+        if len(data.image) > 0:
+            avg[scheduler] = (t1 - t0) / len(data.image)
+        stats.generate += t1 - t0
+        _image = grid(data)
+        data = await upscale(data)
+        t2 = time.perf_counter()
+        stats.upscale += t2 - t1
+        stats.wall += t2 - t0
+        its = sd.generate.steps / ((t1 - t0) / len(data.image)) if len(data.image) > 0 else 0
+        avg_time = round((t1 - t0) / len(data.image)) if len(data.image) > 0 else 0
+        log.info({ 'time' : { 'wall': round(t1 - t0), 'average': avg_time, 'upscale': round(t2 - t1), 'its': round(its, 2) } })
+        log.info({ 'generated': stats.images, 'max': params.max, 'progress': round(100 * stats.images / params.max, 1) })
+        if params.max != 0 and stats.images >= params.max:
+            break
+
+
+if __name__ == '__main__':
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        asyncio.run(interrupt())
+        asyncio.run(close())
+        log.info({ 'interrupt': True })
+    finally:
+        log.info({ 'sampler performance': avg })
+        log.info({ 'stats' : stats })
+        asyncio.run(close())

cli/hf-convert.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+
+import os
+import sys
+import logging
+import torch
+import diffusers
+import safetensors
+import safetensors.torch as sf
+
+log = logging.getLogger("sd")
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s | %(message)s')
+
+
+def convert(model_id, output_name):
+    if os.path.exists(output_name):
+        log.error(f'Output already exists: {output_name}')
+        return
+    pipe = diffusers.DiffusionPipeline.from_pretrained(model_id)
+    metadata = { 'model_id': model_id }
+    model = {}
+    model['state_dict'] = vars(pipe)['_internal_dict']
+    for k in model['state_dict'].keys():
+        # print(k, getattr(pipe, k))
+        model[k] = getattr(pipe, k)
+    sf.save_model(model, output_name, metadata=metadata)
+    # log.info(f'Saved model: {output_name}')
+
+if __name__ == "__main__":
+    sys.argv.pop(0)
+    if len(sys.argv) < 2:
+        log.info('Usage: hf-convert.py <model_id> <output_name>')
+        sys.exit(1)
+    log.debug(f'Packages: torch={torch.__version__} diffusers={diffusers.__version__} safetensors={safetensors.__version__}')
+    convert(sys.argv[0], sys.argv[1])

cli/hf-search.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+
+import sys
+import huggingface_hub as hf
+from rich import print # pylint: disable=redefined-builtin
+
+if __name__ == "__main__":
+    sys.argv.pop(0)
+    keyword = sys.argv[0] if len(sys.argv) > 0 else ''
+    hf_api = hf.HfApi()
+    model_filter = hf.ModelFilter(
+        model_name=keyword,
+        # task='text-to-image',
+        library=['diffusers'],
+    )
+    res = hf_api.list_models(filter=model_filter, full=True, limit=50, sort="downloads", direction=-1)
+    models = [{ 'name': m.modelId, 'downloads': m.downloads, 'mtime': m.lastModified, 'url': f'https://huggingface.co/{m.modelId}', 'pipeline': m.pipeline_tag, 'tags': m.tags } for m in res]
+    print(models)

cli/idle.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python
+
+import os
+import time
+import datetime
+import logging
+import urllib3
+import requests
+
+class Dot(dict):
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+opts = Dot({
+    "timeout": 3600,
+    "frequency": 60,
+    "action": "sudo shutdown now",
+    "url": "https://127.0.0.1:7860",
+    "user": "",
+    "password": "",
+})
+
+log_format = '%(asctime)s %(levelname)s: %(message)s'
+logging.basicConfig(level = logging.INFO, format = log_format)
+log = logging.getLogger("sd")
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+status = None
+
+def progress():
+    auth = requests.auth.HTTPBasicAuth(opts.user, opts.password) if opts.user is not None and len(opts.user) > 0 and opts.password is not None and len(opts.password) > 0 else None
+    req = requests.get(f'{opts.url}/sdapi/v1/progress?skip_current_image=true', verify=False, auth=auth, timeout=60)
+    if req.status_code != 200:
+        log.error({ 'url': req.url, 'request': req.status_code, 'reason': req.reason })
+        return status
+    else:
+        res = Dot(req.json())
+        log.debug({ 'url': req.url, 'request': req.status_code, 'result': res })
+        return res
+
+log.info(f'sdnext monitor started: {opts}')
+while True:
+    try:
+        status = progress()
+        state = status.get('state', {})
+        last_job = state.get('job_timestamp', None)
+        if last_job is None:
+            log.warning(f'sdnext montoring cannot get last job info: {status}')
+        else:
+            last_job = datetime.datetime.strptime(last_job, "%Y%m%d%H%M%S")
+            elapsed = datetime.datetime.now() - last_job
+            timeout = round(opts.timeout - elapsed.total_seconds())
+            log.info(f'sdnext: last_job={last_job} elapsed={elapsed} timeout={timeout}')
+            if timeout < 0:
+                log.warning(f'sdnext reached: timeout={opts.timeout} action={opts.action}')
+                os.system(opts.action)
+    except Exception as e:
+        log.error(f'sdnext monitor error: {e}')
+    finally:
+        time.sleep(opts.frequency)

cli/image-exif.py

@@ -0,0 +1,163 @@
+#!/bin/env python
+
+import os
+import io
+import re
+import sys
+import json
+from PIL import Image, ExifTags, TiffImagePlugin, PngImagePlugin
+from rich import print # pylint: disable=redefined-builtin
+
+
+def unquote(text):
+    if len(text) == 0 or text[0] != '"' or text[-1] != '"':
+        return text
+    try:
+        return json.loads(text)
+    except Exception:
+        return text
+
+
+def parse_generation_parameters(infotext): # copied from modules.generation_parameters_copypaste
+    if not isinstance(infotext, str):
+        return {}
+
+    re_param = re.compile(r'\s*([\w ]+):\s*("(?:\\"[^,]|\\"|\\|[^\"])+"|[^,]*)(?:,|$)') # multi-word: value
+    re_size = re.compile(r"^(\d+)x(\d+)$") # int x int
+    sanitized = infotext.replace('prompt:', 'Prompt:').replace('negative prompt:', 'Negative prompt:').replace('Negative Prompt', 'Negative prompt') # cleanup everything in brackets so re_params can work
+    sanitized = re.sub(r'<[^>]*>', lambda match: ' ' * len(match.group()), sanitized)
+    sanitized = re.sub(r'\([^)]*\)', lambda match: ' ' * len(match.group()), sanitized)
+    sanitized = re.sub(r'\{[^}]*\}', lambda match: ' ' * len(match.group()), sanitized)
+
+    params = dict(re_param.findall(sanitized))
+    params = { k.strip():params[k].strip() for k in params if k.lower() not in ['hashes', 'lora', 'embeddings', 'prompt', 'negative prompt']} # remove some keys
+    first_param = next(iter(params)) if params else None
+    params_idx = sanitized.find(f'{first_param}:') if first_param else -1
+    negative_idx = infotext.find("Negative prompt:")
+
+    prompt = infotext[:params_idx] if negative_idx == -1 else infotext[:negative_idx] # prompt can be with or without negative prompt
+    negative = infotext[negative_idx:params_idx] if negative_idx >= 0 else ''
+
+    for k, v in params.copy().items(): # avoid dict-has-changed
+        if len(v) > 0 and v[0] == '"' and v[-1] == '"':
+            v = unquote(v)
+        m = re_size.match(v)
+        if v.replace('.', '', 1).isdigit():
+            params[k] = float(v) if '.' in v else int(v)
+        elif v == "True":
+            params[k] = True
+        elif v == "False":
+            params[k] = False
+        elif m is not None:
+            params[f"{k}-1"] = int(m.group(1))
+            params[f"{k}-2"] = int(m.group(2))
+        elif k == 'VAE' and v == 'TAESD':
+            params["Full quality"] = False
+        else:
+            params[k] = v
+    params["Prompt"] = prompt.replace('Prompt:', '').strip()
+    params["Negative prompt"] = negative.replace('Negative prompt:', '').strip()
+    return params
+
+
+class Exif: # pylint: disable=single-string-used-for-slots
+    __slots__ = ('__dict__') # pylint: disable=superfluous-parens
+    def __init__(self, image = None):
+        super(Exif, self).__setattr__('exif', Image.Exif()) # pylint: disable=super-with-arguments
+        self.pnginfo = PngImagePlugin.PngInfo()
+        self.tags = {**dict(ExifTags.TAGS.items()), **dict(ExifTags.GPSTAGS.items())}
+        self.ids = {**{v: k for k, v in ExifTags.TAGS.items()}, **{v: k for k, v in ExifTags.GPSTAGS.items()}}
+        if image is not None:
+            self.load(image)
+
+    def __getattr__(self, attr):
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        return self.exif.get(attr, None)
+
+    def load(self, img: Image):
+        img.load() # exif may not be ready
+        exif_dict = {}
+        try:
+            exif_dict = dict(img._getexif().items()) # pylint: disable=protected-access
+        except Exception:
+            exif_dict = dict(img.info.items())
+        for key, val in exif_dict.items():
+            if isinstance(val, bytes): # decode bytestring
+                val = self.decode(val)
+            if val is not None:
+                if isinstance(key, str):
+                    self.exif[key] = val
+                    self.pnginfo.add_text(key, str(val), zip=False)
+                elif isinstance(key, int) and key in ExifTags.TAGS: # add known tags
+                    if self.tags[key] in ['ExifOffset']:
+                        continue
+                    self.exif[self.tags[key]] = val
+                    self.pnginfo.add_text(self.tags[key], str(val), zip=False)
+                    # if self.tags[key] == 'UserComment': # add geninfo from UserComment
+                        # self.geninfo = val
+                else:
+                    print('metadata unknown tag:', key, val)
+        for key, val in self.exif.items():
+            if isinstance(val, bytes): # decode bytestring
+                self.exif[key] = self.decode(val)
+
+    def decode(self, s: bytes):
+        remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment
+        for encoding in ['utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings
+            try:
+                s = remove_prefix(s, b'UNICODE')
+                s = remove_prefix(s, b'ASCII')
+                s = remove_prefix(s, b'\x00')
+                val = s.decode(encoding, errors="strict")
+                val = re.sub(r'[\x00-\x09]', '', val).strip() # remove remaining special characters
+                if len(val) == 0: # remove empty strings
+                    val = None
+                return val
+            except Exception:
+                pass
+        return None
+
+    def parse(self):
+        x = self.exif.pop('parameters', None) or self.exif.pop('UserComment', None)
+        res = parse_generation_parameters(x)
+        return res
+
+    def get_bytes(self):
+        ifd = TiffImagePlugin.ImageFileDirectory_v2()
+        exif_stream = io.BytesIO()
+        for key, val in self.exif.items():
+            if key in self.ids:
+                ifd[self.ids[key]] = val
+            else:
+                print('metadata unknown exif tag:', key, val)
+        ifd.save(exif_stream)
+        raw = b'Exif\x00\x00' + exif_stream.getvalue()
+        return raw
+
+
+def read_exif(filename: str):
+    if filename.lower().endswith('.heic'):
+        from pi_heif import register_heif_opener
+        register_heif_opener()
+    try:
+        image = Image.open(filename)
+        exif = Exif(image)
+        print('image:', filename, 'format:', image)
+        print('exif:', vars(exif.exif)['_data'])
+        print('info:', exif.parse())
+    except Exception as e:
+        print('metadata error reading:', filename, e)
+
+
+if __name__ == '__main__':
+    sys.argv.pop(0)
+    if len(sys.argv) == 0:
+        print('metadata:', 'no files specified')
+    for fn in sys.argv:
+        if os.path.isfile(fn):
+            read_exif(fn)
+        elif os.path.isdir(fn):
+            for root, _dirs, files in os.walk(fn):
+                for file in files:
+                    read_exif(os.path.join(root, file))

cli/image-grid.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+"""
+Create image grid
+"""
+
+import os
+import argparse
+import math
+import logging
+from pathlib import Path
+import filetype
+from PIL import Image, ImageDraw, ImageFont
+from util import log
+
+
+params = None
+
+
+def wrap(text: str, font: ImageFont.ImageFont, length: int):
+    lines = ['']
+    for word in text.split():
+        line = f'{lines[-1]} {word}'.strip()
+        if font.getlength(line) <= length:
+            lines[-1] = line
+        else:
+            lines.append(word)
+    return '\n'.join(lines)
+
+
+def grid(images, labels = None, width = 0, height = 0, border = 0, square = False, horizontal = False, vertical = False): # pylint: disable=redefined-outer-name
+    if horizontal:
+        rows = 1
+    elif vertical:
+        rows = len(images)
+    elif square:
+        rows = round(math.sqrt(len(images)))
+    else:
+        rows = math.floor(math.sqrt(len(images)))
+    cols = math.ceil(len(images) / rows)
+    size = [0, 0]
+    if width == 0:
+        w = max([i.size[0] for i in images])
+        size[0] = cols * w + cols * border
+    else:
+        size[0] = width
+        w = round(width / cols)
+    if height == 0:
+        h = max([i.size[1] for i in images])
+        size[1] = rows * h + rows * border
+    else:
+        size[1] = height
+        h = round(height / rows)
+    size = tuple(size)
+    image = Image.new('RGB', size = size, color = 'black') # pylint: disable=redefined-outer-name
+    font = ImageFont.truetype('DejaVuSansMono', round(w / 40))
+    for i, img in enumerate(images): # pylint: disable=redefined-outer-name
+        x = (i % cols * w) + (i % cols * border)
+        y = (i // cols * h) + (i // cols * border)
+        img.thumbnail((w, h), Image.Resampling.HAMMING)
+        image.paste(img, box=(x + int(border / 2), y + int(border / 2)))
+        if labels is not None and len(images) == len(labels):
+            ctx = ImageDraw.Draw(image)
+            label = wrap(labels[i], font, w)
+            ctx.text((x + 1 + round(w / 200), y + 1 + round(w / 200)), label, font = font, fill = (0, 0, 0))
+            ctx.text((x, y), label, font = font, fill = (255, 255, 255))
+    log.info({ 'grid': { 'images': len(images), 'rows': rows, 'cols': cols, 'cell': [w, h] } })
+    return image
+
+
+if __name__ == '__main__':
+    log.info({ 'create grid' })
+    parser = argparse.ArgumentParser(description='image grid utility')
+    parser.add_argument("--square", default = False, action='store_true', help = "create square grid")
+    parser.add_argument("--horizontal", default = False, action='store_true', help = "create horizontal grid")
+    parser.add_argument("--vertical", default = False, action='store_true', help = "create vertical grid")
+    parser.add_argument("--width", type = int, default = 0, required = False, help = "fixed grid width")
+    parser.add_argument("--height", type = int, default = 0, required = False, help = "fixed grid height")
+    parser.add_argument("--border", type = int, default = 0, required = False, help = "image border")
+    parser.add_argument('--nolabels', default = False, action='store_true', help = "do not print image labels")
+    parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information")
+    parser.add_argument('output', type = str)
+    parser.add_argument('input', type = str, nargs = '*')
+    params = parser.parse_args()
+    output = params.output if params.output.lower().endswith('.jpg') else params.output + '.jpg'
+    if params.debug:
+        log.setLevel(logging.DEBUG)
+        log.debug({ 'debug': True })
+    log.debug({ 'args': params.__dict__ })
+    images = []
+    labels = []
+    for f in params.input:
+        path = Path(f)
+        if path.is_dir():
+            files = [os.path.join(f, file) for file in os.listdir(f) if os.path.isfile(os.path.join(f, file))]
+        elif path.is_file():
+            files = [f]
+        else:
+            log.warning({ 'grid not a valid file/folder', f})
+            continue
+        files.sort()
+        for file in files:
+            if not filetype.is_image(file):
+                continue
+            if file.lower().endswith('.heic'):
+                from pi_heif import register_heif_opener
+                register_heif_opener()
+            log.debug(file)
+            img = Image.open(file)
+            # img.verify()
+            images.append(img)
+            fp = Path(file)
+            if not params.nolabels:
+                labels.append(fp.stem)
+    # log.info({ 'folder': path.parent, 'labels': labels })
+    if len(images) > 0:
+        image = grid(
+            images = images,
+            labels = labels,
+            width = params.width,
+            height = params.height,
+            border = params.border,
+            square = params.square,
+            horizontal = params.horizontal,
+            vertical = params.vertical)
+        image.save(output, 'JPEG', optimize = True, quality = 60)
+        log.info({ 'grid': { 'file': output, 'size': list(image.size) } })
+    else:
+        log.info({ 'grid': 'nothing to do' })

cli/image-interrogate.py

@@ -0,0 +1,109 @@
+#!/usr/bin/env python
+"""
+use clip to interrogate image(s)
+"""
+
+import io
+import base64
+import sys
+import os
+import asyncio
+import filetype
+from PIL import Image
+from util import log, Map
+import sdapi
+
+
+stats = { 'captions': {}, 'keywords': {} }
+exclude = ['a', 'in', 'on', 'out', 'at', 'the', 'and', 'with', 'next', 'to', 'it', 'for', 'of', 'into', 'that']
+
+
+def decode(encoding):
+    if encoding.startswith("data:image/"):
+        encoding = encoding.split(";")[1].split(",")[1]
+    return Image.open(io.BytesIO(base64.b64decode(encoding)))
+
+
+def encode(f):
+    image = Image.open(f)
+    exif = image.getexif()
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
+    with io.BytesIO() as stream:
+        image.save(stream, 'JPEG', exif = exif)
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+def print_summary():
+    captions = dict(sorted(stats['captions'].items(), key=lambda x:x[1], reverse=True))
+    log.info({ 'caption stats': captions })
+    keywords = dict(sorted(stats['keywords'].items(), key=lambda x:x[1], reverse=True))
+    log.info({ 'keyword stats': keywords })
+
+
+async def interrogate(f):
+    if not filetype.is_image(f):
+        log.info({ 'interrogate skip': f })
+        return
+    json = Map({ 'image': encode(f) })
+    log.info({ 'interrogate': f })
+    # run clip
+    json.model = 'clip'
+    res = await sdapi.post('/sdapi/v1/interrogate', json)
+    caption = ""
+    style = ""
+    if 'caption' in res:
+        caption = res.caption
+        log.info({ 'interrogate caption': caption })
+        if ', by' in caption:
+            style = caption.split(', by')[1].strip()
+            log.info({ 'interrogate style': style })
+        for word in caption.split(' '):
+            if word not in exclude:
+                stats['captions'][word] = stats['captions'][word] + 1 if word in stats['captions'] else 1
+    else:
+        log.error({ 'interrogate clip error': res })
+    # run booru
+    json.model = 'deepdanbooru'
+    res = await sdapi.post('/sdapi/v1/interrogate', json)
+    keywords = {}
+    if 'caption' in res:
+        for term in res.caption.split(', '):
+            term = term.replace('(', '').replace(')', '').replace('\\', '').split(':')
+            if len(term) < 2:
+                continue
+            keywords[term[0]] = term[1]
+        keywords = dict(sorted(keywords.items(), key=lambda x:x[1], reverse=True))
+        for word in keywords.items():
+            stats['keywords'][word[0]] = stats['keywords'][word[0]] + 1 if word[0] in stats['keywords'] else 1
+        log.info({ 'interrogate keywords': keywords })
+    else:
+        log.error({ 'interrogate booru error': res })
+    return caption, keywords, style
+
+
+async def main():
+    sys.argv.pop(0)
+    await sdapi.session()
+    if len(sys.argv) == 0:
+        log.error({ 'interrogate': 'no files specified' })
+    for arg in sys.argv:
+        if os.path.exists(arg):
+            if os.path.isfile(arg):
+                await interrogate(arg)
+            elif os.path.isdir(arg):
+                for root, _dirs, files in os.walk(arg):
+                    for f in files:
+                        _caption, _keywords, _style = await interrogate(os.path.join(root, f))
+            else:
+                log.error({ 'interrogate unknown file type': arg })
+        else:
+            log.error({ 'interrogate file missing': arg })
+    await sdapi.close()
+    print_summary()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cli/image-palette.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python
+# based on <https://towardsdatascience.com/image-color-extraction-with-python-in-4-steps-8d9370d9216e>
+
+import os
+import io
+import pathlib
+import argparse
+import importlib
+import pandas as pd
+import numpy as np
+import extcolors
+import filetype
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from matplotlib.offsetbox import OffsetImage, AnnotationBbox
+from colormap import rgb2hex
+from PIL import Image
+from util import log
+grid = importlib.import_module('image-grid').grid
+
+
+def color_to_df(param):
+    colors_pre_list = str(param).replace('([(','').split(', (')[0:-1]
+    df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list]
+    df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list]
+    #convert RGB to HEX code
+    df_color_up = [rgb2hex(int(i.split(", ")[0].replace("(","")),
+                           int(i.split(", ")[1]),
+                           int(i.split(", ")[2].replace(")",""))) for i in df_rgb]
+    df = pd.DataFrame(zip(df_color_up, df_percent), columns = ['c_code','occurence'])
+    return df
+
+
+def palette(img, params, output):
+    size = 1024
+    img.thumbnail((size, size), Image.HAMMING)
+
+    #crate dataframe
+    colors_x = extcolors.extract_from_image(img, tolerance = params.color, limit = 13)
+    df_color = color_to_df(colors_x)
+
+    #annotate text
+    list_color = list(df_color['c_code'])
+    list_precent = [int(i) for i in list(df_color['occurence'])]
+    text_c = [c + ' ' + str(round(p * 100 / sum(list_precent), 1)) +'%' for c, p in zip(list_color, list_precent)]
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(120,60), dpi=10)
+    fig.set_facecolor('black')
+
+    #donut plot
+    wedges, _text = ax1.pie(list_precent, labels= text_c, labeldistance= 1.05, colors = list_color, textprops={'fontsize': 100, 'color':'white'})
+    plt.setp(wedges, width=0.3)
+
+    #add image in the center of donut plot
+    data = np.asarray(img)
+    imagebox = OffsetImage(data, zoom=2.5)
+    ab = AnnotationBbox(imagebox, (0, 0))
+    ax1.add_artist(ab)
+
+    #color palette
+    x_posi, y_posi, y_posi2 = 160, -260, -260
+    for c in list_color:
+        if list_color.index(c) <= 5:
+            y_posi += 240
+            rect = patches.Rectangle((x_posi, y_posi), 540, 230, facecolor = c)
+            ax2.add_patch(rect)
+            ax2.text(x = x_posi + 100, y = y_posi + 140, s = c, fontdict={'fontsize': 140}, color = 'white')
+        else:
+            y_posi2 += 240
+            rect = patches.Rectangle((x_posi + 600, y_posi2), 540, 230, facecolor = c)
+            ax2.add_artist(rect)
+            ax2.text(x = x_posi + 700, y = y_posi2 + 140, s = c, fontdict={'fontsize': 140}, color = 'white')
+
+    # add background to force layout
+    fig.set_facecolor('black')
+    ax2.axis('off')
+    tmp = Image.new('RGB', (2000, 1400), (0, 0, 0))
+    plt.imshow(tmp)
+    plt.tight_layout(rect = (-0.08, -0.2, 1.18, 1.05))
+
+    # save image
+    if output is not None:
+        buf = io.BytesIO()
+        plt.savefig(buf, format='png')
+        pltimg = Image.open(buf)
+        pltimg = pltimg.convert('RGB')
+        pltimg.save(output)
+        buf.close()
+        log.info({ 'palette created': output })
+
+    plt.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'extract image color palette')
+    parser.add_argument('--color', type=int, default=20, help="color tolerance threshdold")
+    parser.add_argument('--output', type=str, required=False, default='', help='folder to store images')
+    parser.add_argument('--suffix', type=str, required=False, default='pallete', help='add suffix to image name')
+    parser.add_argument('--grid', default=False, action='store_true', help = "create grid of images before processing")
+    parser.add_argument('input', type=str, nargs='*')
+    args = parser.parse_args()
+    log.info({ 'palette args': vars(args) })
+    if args.output != '':
+        pathlib.Path(args.output).mkdir(parents = True, exist_ok = True)
+    if not args.grid:
+        for arg in args.input:
+            if os.path.isfile(arg) and filetype.is_image(arg):
+                image = Image.open(arg)
+                fn = os.path.join(args.output, pathlib.Path(arg).stem + '-' + args.suffix + '.jpg')
+                palette(image, args, fn)
+            elif os.path.isdir(arg):
+                for root, _dirs, files in os.walk(arg):
+                    for f in files:
+                        if filetype.is_image(os.path.join(root, f)):
+                            image = Image.open(os.path.join(root, f))
+                            fn = os.path.join(args.output, pathlib.Path(f).stem + '-' + args.suffix + '.jpg')
+                            palette(image, args, fn)
+    else:
+        images = []
+        for arg in args.input:
+            if os.path.isfile(arg) and filetype.is_image(arg):
+                images.append(Image.open(arg))
+            elif os.path.isdir(arg):
+                for root, _dirs, files in os.walk(arg):
+                    for f in files:
+                        if filetype.is_image(os.path.join(root, f)):
+                            images.append(Image.open(os.path.join(root, f)))
+        image = grid(images)
+        fn = os.path.join(args.output, args.suffix + '.jpg')
+        palette(image, args, fn)

cli/image-watermark.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python
+import os
+import io
+import pathlib
+import argparse
+import filetype
+import numpy as np
+from imwatermark import WatermarkEncoder, WatermarkDecoder
+from PIL import Image
+from PIL.ExifTags import TAGS
+from PIL.TiffImagePlugin import ImageFileDirectory_v2
+from util import log, Map
+import piexif
+import piexif.helper
+
+
+options = Map({ 'method': 'dwtDctSvd', 'type': 'bytes' })
+
+
+def get_exif(image):
+    # using piexif
+    res1 = {}
+    try:
+        exif = piexif.load(image.info["exif"])
+        exif = exif.get("Exif", {})
+        for k, v in exif.items():
+            key = list(vars(piexif.ExifIFD).keys())[list(vars(piexif.ExifIFD).values()).index(k)]
+            res1[key] = piexif.helper.UserComment.load(v)
+    except Exception:
+        pass
+    # using pillow
+    res2 = {}
+    try:
+        res2 = { TAGS[k]: v for k, v in image.getexif().items() if k in TAGS }
+    except Exception:
+        pass
+    return {**res1, **res2}
+
+
+def set_exif(d: dict):
+    ifd = ImageFileDirectory_v2()
+    _TAGS = {v: k for k, v in TAGS.items()} # enumerate possible exif tags
+    for k, v in d.items():
+        ifd[_TAGS[k]] = v
+    exif_stream = io.BytesIO()
+    ifd.save(exif_stream)
+    encoded = b'Exif\x00\x00' + exif_stream.getvalue()
+    return encoded
+
+
+def get_watermark(image, params):
+    data = np.asarray(image)
+    decoder = WatermarkDecoder(options.type, params.length)
+    decoded = decoder.decode(data, options.method)
+    wm = decoded.decode(encoding='ascii', errors='ignore')
+    return wm
+
+
+def set_watermark(image, params):
+    data = np.asarray(image)
+    encoder = WatermarkEncoder()
+    length = params.length // 8
+    text = f"{params.wm:<{length}}"[:length]
+    bytearr = text.encode(encoding='ascii', errors='ignore')
+    encoder.set_watermark(options.type, bytearr)
+    encoded = encoder.encode(data, options.method)
+    image = Image.fromarray(encoded)
+    return image
+
+
+def watermark(params, file):
+    if not os.path.exists(file):
+        log.error({ 'watermark': 'file not found' })
+        return
+    if not filetype.is_image(file):
+        log.error({ 'watermark': 'file is not an image' })
+        return
+    image = Image.open(file)
+    if image.width * image.height < 256 * 256:
+        log.error({ 'watermark': 'image too small' })
+        return
+
+    exif = get_exif(image)
+
+    if params.command == 'read':
+        fn = params.input
+        wm = get_watermark(image, params)
+
+    elif params.command == 'write':
+        metadata = b'' if params.strip else set_exif(exif)
+        if params.output != '':
+            pathlib.Path(params.output).mkdir(parents = True, exist_ok = True)
+        image=set_watermark(image, params)
+        fn = os.path.join(params.output, file)
+        image.save(fn, exif=metadata)
+
+        if params.verify:
+            image = Image.open(fn)
+            data = np.asarray(image)
+            decoder = WatermarkDecoder(options.type, params.length)
+            decoded = decoder.decode(data, options.method)
+            wm = decoded.decode(encoding='ascii', errors='ignore')
+        else:
+            wm = params.wm
+
+    log.info({ 'file': fn })
+    log.info({ 'resolution': f'{image.width}x{image.height}' })
+    log.info({ 'watermark': wm })
+    log.info({ 'exif': None if params.strip else exif })
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'image watermarking')
+    parser.add_argument('command', choices = ['read', 'write'])
+    parser.add_argument('--wm', type=str, required=False, default='sdnext', help='watermark string')
+    parser.add_argument('--strip', default=False, action='store_true', help = "strip existing exif data")
+    parser.add_argument('--verify', default=False, action='store_true', help = "verify watermark during write")
+    parser.add_argument('--length', type=int, default=32, help="watermark length in bits")
+    parser.add_argument('--output', type=str, required=False, default='', help='folder to store images, default is overwrite in-place')
+    parser.add_argument('input', type=str, nargs='*')
+    args = parser.parse_args()
+    # log.info({ 'watermark args': vars(args), 'options': options })
+    for arg in args.input:
+        if os.path.isfile(arg):
+            watermark(args, arg)
+        elif os.path.isdir(arg):
+            for root, _dirs, files in os.walk(arg):
+                for f in files:
+                    watermark(args, os.path.join(root, f))

cli/install-sf.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+import os
+import re
+import sys
+
+torch_supported = ['211', '212','220','221']
+cuda_supported = ['cu118', 'cu121']
+python_supported = ['39', '310', '311']
+repo_url = 'https://github.com/chengzeyi/stable-fast'
+api_url = 'https://api.github.com/repos/chengzeyi/stable-fast/releases/tags/nightly'
+path_url = '/releases/download/nightly'
+
+
+def install_pip(arg: str):
+    import subprocess
+    cmd = f'"{sys.executable}" -m pip install -U {arg}'
+    print(f'Running: {cmd}')
+    result = subprocess.run(cmd, shell=True, check=False, env=os.environ)
+    return result.returncode == 0
+
+
+def get_nightly():
+    import requests
+    r = requests.get(api_url, timeout=10)
+    if r.status_code != 200:
+        print('Failed to get nightly version')
+        return None
+    json = r.json()
+    assets = json.get('assets', [])
+    if len(assets) == 0:
+        print('Failed to get nightly version')
+        return None
+    asset = assets[0].get('name', '')
+    pattern = r"-(.+?)\+"
+    match = re.search(pattern, asset)
+    if match:
+        ver = match.group(1)
+        print(f'Nightly version: {ver}')
+        return ver
+    else:
+        print('Failed to get nightly version')
+        return None
+
+
+def install_stable_fast():
+    import torch
+
+    python_ver = f'{sys.version_info.major}{sys.version_info.minor}'
+    if python_ver not in python_supported:
+        raise ValueError(f'StableFast unsupported python: {python_ver} required {python_supported}')
+    if sys.platform == 'linux':
+        bin_url = 'manylinux2014_x86_64.whl'
+    elif sys.platform == 'win32':
+        bin_url = 'win_amd64.whl'
+    else:
+        raise ValueError(f'StableFast unsupported platform: {sys.platform}')
+
+    torch_ver, cuda_ver = torch.__version__.split('+')
+    torch_ver = torch_ver.replace('.', '')
+    sf_ver = get_nightly()
+
+    if torch_ver not in torch_supported:
+        print(f'StableFast unsupported torch: {torch_ver} required {torch_supported}')
+        print('Installing from source...')
+        url = 'git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast'
+    elif cuda_ver not in cuda_supported:
+        print(f'StableFast unsupported CUDA: {cuda_ver} required {cuda_supported}')
+        print('Installing from source...')
+        url = 'git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast'
+    elif sf_ver is None:
+        print('StableFast cannot determine version')
+        print('Installing from source...')
+        url = 'git+https://github.com/chengzeyi/stable-fast.git@main#egg=stable-fast'
+    else:
+        print('Installing wheel...')
+        file_url = f'stable_fast-{sf_ver}+torch{torch_ver}{cuda_ver}-cp{python_ver}-cp{python_ver}-{bin_url}'
+        url = f'{repo_url}/{path_url}/{file_url}'
+
+    ok = install_pip(url)
+    if ok:
+        import sfast
+        print(f'StableFast installed: {sfast.__version__}')
+    else:
+        print('StableFast install failed')
+
+if __name__ == '__main__':
+    install_stable_fast()

cli/latents.py

@@ -0,0 +1,170 @@
+#!/usr/bin/env python
+
+import os
+import sys
+import json
+import pathlib
+import argparse
+import warnings
+
+import cv2
+import numpy as np
+import torch
+from PIL import Image
+from torchvision import transforms
+from tqdm import tqdm
+from util import Map
+
+from rich.pretty import install as pretty_install
+from rich.traceback import install as traceback_install
+from rich.console import Console
+
+console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+pretty_install(console=console)
+traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False)
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'modules', 'lora'))
+import library.model_util as model_util
+import library.train_util as train_util
+
+warnings.filterwarnings('ignore')
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+options = Map({
+  'batch': 1,
+  'input': '',
+  'json': '',
+  'max': 1024,
+  'min': 256,
+  'noupscale': False,
+  'precision': 'fp32',
+  'resolution': '512,512',
+  'steps': 64,
+  'vae': 'stabilityai/sd-vae-ft-mse'
+})
+vae = None
+
+
+def get_latents(local_vae, images, weight_dtype):
+    image_transforms = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.5], [0.5]) ])
+    img_tensors = [image_transforms(image) for image in images]
+    img_tensors = torch.stack(img_tensors)
+    img_tensors = img_tensors.to(device, weight_dtype)
+    with torch.no_grad():
+        latents = local_vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
+    return latents, [images[0].shape[0], images[0].shape[1]]
+
+
+def get_npz_filename_wo_ext(data_dir, image_key):
+    return os.path.join(data_dir, os.path.splitext(os.path.basename(image_key))[0])
+
+
+def create_vae_latents(local_params):
+    args = Map({**options, **local_params})
+    console.log(f'create vae latents args: {args}')
+    image_paths = train_util.glob_images(args.input)
+    if os.path.exists(args.json):
+        with open(args.json, 'rt', encoding='utf-8') as f:
+            metadata = json.load(f)
+    else:
+        return
+    if args.precision == 'fp16':
+        weight_dtype = torch.float16
+    elif args.precision == 'bf16':
+        weight_dtype = torch.bfloat16
+    else:
+        weight_dtype = torch.float32
+    global vae # pylint: disable=global-statement
+    if vae is None:
+        vae = model_util.load_vae(args.vae, weight_dtype)
+        vae.eval()
+        vae.to(device, dtype=weight_dtype)
+    max_reso = tuple([int(t) for t in args.resolution.split(',')])
+    assert len(max_reso) == 2, f'illegal resolution: {args.resolution}'
+    bucket_manager = train_util.BucketManager(args.noupscale, max_reso, args.min, args.max, args.steps)
+    if not args.noupscale:
+        bucket_manager.make_buckets()
+    img_ar_errors = []
+    def process_batch(is_last):
+        for bucket in bucket_manager.buckets:
+            if (is_last and len(bucket) > 0) or len(bucket) >= args.batch:
+                latents, original_size = get_latents(vae, [img for _, img in bucket], weight_dtype)
+                assert latents.shape[2] == bucket[0][1].shape[0] // 8 and latents.shape[3] == bucket[0][1].shape[1] // 8, f'latent shape {latents.shape}, {bucket[0][1].shape}'
+                for (image_key, _), latent in zip(bucket, latents):
+                    npz_file_name = get_npz_filename_wo_ext(args.input, image_key)
+                    # np.savez(npz_file_name, latent)
+                    kwargs = {}
+                    np.savez(
+                        npz_file_name,
+                        latents=latent,
+                        original_size=np.array(original_size),
+                        crop_ltrb=np.array([0, 0]),
+                        **kwargs,
+                    )
+                bucket.clear()
+    data = [[(None, ip)] for ip in image_paths]
+    bucket_counts = {}
+    for data_entry in tqdm(data, smoothing=0.0):
+        if data_entry[0] is None:
+            continue
+        img_tensor, image_path = data_entry[0]
+        if img_tensor is not None:
+            image = transforms.functional.to_pil_image(img_tensor)
+        else:
+            image = Image.open(image_path)
+        image_key = os.path.basename(image_path)
+        image_key = os.path.join(os.path.basename(pathlib.Path(image_path).parent), pathlib.Path(image_path).stem)
+        if image_key not in metadata:
+            metadata[image_key] = {}
+        reso, resized_size, ar_error = bucket_manager.select_bucket(image.width, image.height)
+        img_ar_errors.append(abs(ar_error))
+        bucket_counts[reso] = bucket_counts.get(reso, 0) + 1
+        metadata[image_key]['train_resolution'] = (reso[0] - reso[0] % 8, reso[1] - reso[1] % 8)
+        if not args.noupscale:
+            assert resized_size[0] == reso[0] or resized_size[1] == reso[1], f'internal error, resized size not match: {reso}, {resized_size}, {image.width}, {image.height}'
+            assert resized_size[0] >= reso[0] and resized_size[1] >= reso[1], f'internal error, resized size too small: {reso}, {resized_size}, {image.width}, {image.height}'
+        assert resized_size[0] >= reso[0] and resized_size[1] >= reso[1], f'internal error resized size is small: {resized_size}, {reso}'
+        image = np.array(image)
+        if resized_size[0] != image.shape[1] or resized_size[1] != image.shape[0]:
+            image = cv2.resize(image, resized_size, interpolation=cv2.INTER_AREA)
+        if resized_size[0] > reso[0]:
+            trim_size = resized_size[0] - reso[0]
+            image = image[:, trim_size//2:trim_size//2 + reso[0]]
+        if resized_size[1] > reso[1]:
+            trim_size = resized_size[1] - reso[1]
+            image = image[trim_size//2:trim_size//2 + reso[1]]
+        assert image.shape[0] == reso[1] and image.shape[1] == reso[0], f'internal error, illegal trimmed size: {image.shape}, {reso}'
+        bucket_manager.add_image(reso, (image_key, image))
+        process_batch(False)
+
+    process_batch(True)
+    vae.to('cpu')
+
+    bucket_manager.sort()
+    img_ar_errors = np.array(img_ar_errors)
+    for i, reso in enumerate(bucket_manager.resos):
+        count = bucket_counts.get(reso, 0)
+        if count > 0:
+            console.log(f'vae latents bucket: {i+1}/{len(bucket_manager.resos)} resolution: {reso} images: {count} mean-ar-error: {np.mean(img_ar_errors)}')
+    with open(args.json, 'wt', encoding='utf-8') as f:
+        json.dump(metadata, f, indent=2)
+
+
+def unload_vae():
+    global vae # pylint: disable=global-statement
+    vae = None
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('input', type=str, help='directory for train images')
+    parser.add_argument('--json', type=str, required=True, help='metadata file to input')
+    parser.add_argument('--vae', type=str, required=True, help='model name or path to encode latents')
+    parser.add_argument('--batch', type=int, default=1, help='batch size in inference')
+    parser.add_argument('--resolution', type=str, default='512,512', help='max resolution in fine tuning (width,height)')
+    parser.add_argument('--min', type=int, default=256, help='minimum resolution for buckets')
+    parser.add_argument('--max', type=int, default=1024, help='maximum resolution for buckets')
+    parser.add_argument('--steps', type=int, default=64, help='steps of resolution for buckets, divisible by 8')
+    parser.add_argument('--noupscale', action='store_true', help='make bucket for each image without upscaling')
+    parser.add_argument('--precision', type=str, default='fp32', choices=['fp32', 'fp16', 'bf16'], help='use precision')
+    params = parser.parse_args()
+    create_vae_latents(vars(params))

cli/lcm-convert.py

@@ -0,0 +1,55 @@
+import os
+import argparse
+import torch
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, AutoPipelineForText2Image, LCMScheduler
+
+parser = argparse.ArgumentParser("lcm_convert")
+parser.add_argument("--name", help="Name of the new LCM model", type=str)
+parser.add_argument("--model", help="A model to convert", type=str)
+parser.add_argument("--lora-scale", default=1.0, help="Strenght of the LCM", type=float)
+parser.add_argument("--huggingface", action="store_true", help="Use Hugging Face models instead of safetensors models")
+parser.add_argument("--upload", action="store_true", help="Upload the new LCM model to Hugging Face")
+parser.add_argument("--no-half", action="store_true", help="Convert the new LCM model to FP32")
+parser.add_argument("--no-save", action="store_true", help="Don't save the new LCM model to local disk")
+parser.add_argument("--sdxl", action="store_true", help="Use SDXL models")
+parser.add_argument("--ssd-1b", action="store_true", help="Use SSD-1B models")
+
+args = parser.parse_args()
+
+if args.huggingface:
+    pipeline = AutoPipelineForText2Image.from_pretrained(args.model, torch_dtype=torch.float16, variant="fp16")
+else:
+    if args.sdxl or args.ssd_1b:
+        pipeline = StableDiffusionXLPipeline.from_single_file(args.model)
+    else:
+        pipeline = StableDiffusionPipeline.from_single_file(args.model)
+
+pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
+if args.sdxl:
+    pipeline.load_lora_weights("latent-consistency/lcm-lora-sdxl")
+elif args.ssd_1b:
+    pipeline.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
+else:
+    pipeline.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
+pipeline.fuse_lora(lora_scale=args.lora_scale)
+
+#components = pipeline.components
+#pipeline = LatentConsistencyModelPipeline(**components)
+
+if args.no_half:
+    pipeline = pipeline.to(dtype=torch.float32)
+else:
+    pipeline = pipeline.to(dtype=torch.float16)
+print(pipeline)
+
+if not args.no_save:
+    os.makedirs(f"models--local--{args.name}/snapshots")
+    if args.no_half:
+        pipeline.save_pretrained(f"models--local--{args.name}/snapshots/{args.name}")
+    else:
+        pipeline.save_pretrained(f"models--local--{args.name}/snapshots/{args.name}", variant="fp16")
+if args.upload:
+    if args.no_half:
+        pipeline.push_to_hub(args.name)
+    else:
+        pipeline.push_to_hub(args.name, variant="fp16")

cli/model-jit.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python
+import os
+import time
+import functools
+import argparse
+import logging
+import warnings
+from dataclasses import dataclass
+
+logging.getLogger("DeepSpeed").disabled = True
+warnings.filterwarnings(action="ignore", category=FutureWarning)
+warnings.filterwarnings(action="ignore", category=DeprecationWarning)
+
+import torch
+import diffusers
+
+n_warmup = 5
+n_traces = 10
+n_runs = 100
+args = {}
+pipe = None
+log = logging.getLogger("sd")
+
+
+def setup_logging():
+    from rich.theme import Theme
+    from rich.logging import RichHandler
+    from rich.console import Console
+    from rich.traceback import install
+    log.setLevel(logging.DEBUG)
+    console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({ "traceback.border": "black", "traceback.border.syntax_error": "black", "inspect.value.border": "black" }))
+    logging.basicConfig(level=logging.ERROR, format='%(asctime)s | %(name)s | %(levelname)s | %(module)s | %(message)s', handlers=[logging.NullHandler()]) # redirect default logger to null
+    rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=logging.DEBUG, console=console)
+    rh.setLevel(logging.DEBUG)
+    log.addHandler(rh)
+    logging.getLogger("diffusers").setLevel(logging.ERROR)
+    logging.getLogger("torch").setLevel(logging.ERROR)
+    warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning)
+    install(console=console, extra_lines=1, max_frames=10, width=console.width, word_wrap=False, indent_guides=False, suppress=[])
+
+
+def generate_inputs():
+    if args.type == 'sd15':
+        sample = torch.randn(2, 4, 64, 64).half().cuda()
+        timestep = torch.rand(1).half().cuda() * 999
+        encoder_hidden_states = torch.randn(2, 77, 768).half().cuda()
+        return sample, timestep, encoder_hidden_states
+    if args.type == 'sdxl':
+        sample = torch.randn(2, 4, 64, 64).half().cuda()
+        timestep = torch.rand(1).half().cuda() * 999
+        encoder_hidden_states = torch.randn(2, 77, 768).half().cuda()
+        text_embeds = torch.randn(1, 77, 2048).half().cuda()
+        return sample, timestep, encoder_hidden_states, text_embeds
+
+
+def load_model():
+    log.info(f'versions: torch={torch.__version__} diffusers={diffusers.__version__}')
+    diffusers_load_config = {
+        "low_cpu_mem_usage": True,
+        "torch_dtype": torch.float16,
+        "safety_checker": None,
+        "requires_safety_checker": False,
+        "load_safety_checker": False,
+        "load_connected_pipeline": True,
+        "use_safetensors": True,
+    }
+    pipeline = diffusers.StableDiffusionPipeline if args.type == 'sd15' else diffusers.StableDiffusionXLPipeline
+    global pipe # pylint: disable=global-statement
+    t0 = time.time()
+    pipe = pipeline.from_single_file(args.model, **diffusers_load_config).to('cuda')
+    size = os.path.getsize(args.model)
+    log.info(f'load: model={args.model} type={args.type} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb')
+
+
+def load_trace(fn: str):
+
+    @dataclass
+    class UNet2DConditionOutput:
+        sample: torch.FloatTensor
+
+    class TracedUNet(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.in_channels = pipe.unet.in_channels
+            self.device = pipe.unet.device
+
+        def forward(self, latent_model_input, t, encoder_hidden_states):
+            sample = unet_traced(latent_model_input, t, encoder_hidden_states)[0]
+            return UNet2DConditionOutput(sample=sample)
+
+    t0 = time.time()
+    unet_traced = torch.jit.load(fn)
+    pipe.unet = TracedUNet()
+    size = os.path.getsize(fn)
+    log.info(f'load: optimized={fn} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb')
+
+
+def trace_model():
+    log.info(f'tracing model: {args.model}')
+    torch.set_grad_enabled(False)
+    unet = pipe.unet
+    unet.eval()
+    # unet.to(memory_format=torch.channels_last)  # use channels_last memory format
+    unet.forward = functools.partial(unet.forward, return_dict=False)  # set return_dict=False as default
+
+    # warmup
+    t0 = time.time()
+    for _ in range(n_warmup):
+        with torch.inference_mode():
+            inputs = generate_inputs()
+            _output = unet(*inputs)
+    log.info(f'warmup: time={time.time() - t0:.3f}s passes={n_warmup}')
+
+    # trace
+    t0 = time.time()
+    unet_traced = torch.jit.trace(unet, inputs, check_trace=True)
+    unet_traced.eval()
+    log.info(f'trace: time={time.time() - t0:.3f}s')
+
+    # optimize graph
+    t0 = time.time()
+    for _ in range(n_traces):
+        with torch.inference_mode():
+            inputs = generate_inputs()
+            _output = unet_traced(*inputs)
+    log.info(f'optimize: time={time.time() - t0:.3f}s passes={n_traces}')
+
+    # save the model
+    if args.save:
+        t0 = time.time()
+        basename, _ext = os.path.splitext(args.model)
+        fn = f"{basename}.pt"
+        unet_traced.save(fn)
+        size = os.path.getsize(fn)
+        log.info(f'save: optimized={fn} time={time.time() - t0:.3f}s size={size / 1024 / 1024:.3f}mb')
+        return fn
+
+    pipe.unet = unet_traced
+    return None
+
+
+def benchmark_model(msg: str):
+    with torch.inference_mode():
+        inputs = generate_inputs()
+        torch.cuda.synchronize()
+        for n in range(n_runs):
+            if n > n_runs / 10:
+                t0 = time.time()
+            _output = pipe.unet(*inputs)
+        torch.cuda.synchronize()
+        t1 = time.time()
+        log.info(f"benchmark unet: {t1 - t0:.3f}s passes={n_runs} type={msg}")
+        return t1 - t0
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description = 'SD.Next')
+    parser.add_argument('--model', type=str, default='', required=True, help='model path')
+    parser.add_argument('--type', type=str, default='sd15', choices=['sd15', 'sdxl'], required=False, help='model type, default: %(default)s')
+    parser.add_argument('--benchmark', default = False, action='store_true', help = "run benchmarks, default: %(default)s")
+    parser.add_argument('--trace', default = True, action='store_true', help = "run jit tracing, default: %(default)s")
+    parser.add_argument('--save', default = False, action='store_true', help = "save optimized unet, default: %(default)s")
+    args = parser.parse_args()
+    setup_logging()
+    log.info('sdnext model jit tracing')
+    if not os.path.isfile(args.model):
+        log.error(f"invalid model path: {args.model}")
+        exit(1)
+    load_model()
+    if args.benchmark:
+        time0 = benchmark_model('original')
+    unet_saved = trace_model()
+    if unet_saved is not None:
+        load_trace(unet_saved)
+    if args.benchmark:
+        time1 = benchmark_model('traced')
+        log.info(f'benchmark speedup: {100 * (time0 - time1) / time0:.3f}%')

cli/model-metadata.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+import os
+import sys
+import json
+from rich import print # pylint: disable=redefined-builtin
+
+
+def read_metadata(fn):
+    res = {}
+    with open(fn, mode="rb") as f:
+        metadata_len = f.read(8)
+        metadata_len = int.from_bytes(metadata_len, "little")
+        json_start = f.read(2)
+        if metadata_len <= 2 or json_start not in (b'{"', b"{'"):
+            print(f"Not a valid safetensors file: {fn}")
+        json_data = json_start + f.read(metadata_len-2)
+        json_obj = json.loads(json_data)
+        for k, v in json_obj.get("__metadata__", {}).items():
+            res[k] = v
+            if isinstance(v, str) and v[0:1] == '{':
+                try:
+                    res[k] = json.loads(v)
+                except Exception:
+                    pass
+    print(f"{fn}: {json.dumps(res, indent=4)}")
+
+
+def main():
+    if len(sys.argv) == 0:
+        print('metadata:', 'no files specified')
+    for fn in sys.argv:
+        if os.path.isfile(fn):
+            read_metadata(fn)
+        elif os.path.isdir(fn):
+            for root, _dirs, files in os.walk(fn):
+                for file in files:
+                    read_metadata(os.path.join(root, file))
+
+if __name__ == '__main__':
+    sys.argv.pop(0)
+    main()

cli/nvidia-smi.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+import os
+import json
+import shutil
+import subprocess
+import xmltodict
+from rich import print # pylint: disable=redefined-builtin
+from util import log, Map
+
+
+def get_nvidia_smi(output='dict'):
+    smi = shutil.which('nvidia-smi')
+    if smi is None:
+        log.error("nvidia-smi not found")
+        return None
+    result = subprocess.run(f'"{smi}" -q -x', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    xml = result.stdout.decode(encoding="utf8", errors="ignore")
+    d = xmltodict.parse(xml)
+    if 'nvidia_smi_log' in d:
+        d = d['nvidia_smi_log']
+    if 'gpu' in d and 'supported_clocks' in d['gpu']:
+        del d['gpu']['supported_clocks']
+    if output == 'dict':
+        return d
+    elif output == 'class' or output == 'map':
+        d = Map(d)
+        return d
+    elif output == 'json':
+        return json.dumps(d, indent=4)
+    return None
+
+
+if __name__ == "__main__":
+    res = get_nvidia_smi(output='dict')
+    print(type(res), res)

cli/options.py

@@ -0,0 +1,141 @@
+from util import Map
+
+embedding = Map({
+    "id_task": 0,
+    "embedding_name": "",
+    "learn_rate": -1,
+    "batch_size": 1,
+    "steps": 500,
+    "data_root": "",
+    "log_directory": "train/log",
+    "template_filename": "subject_filewords.txt",
+    "gradient_step": 20,
+    "training_width": 512,
+    "training_height": 512,
+    "shuffle_tags": False,
+    "tag_drop_out": 0,
+    "clip_grad_mode": "disabled",
+    "clip_grad_value": "0.1",
+    "latent_sampling_method": "deterministic",
+    "create_image_every": 0,
+    "save_embedding_every": 0,
+    "save_image_with_stored_embedding": False,
+    "preview_from_txt2img": False,
+    "preview_prompt": "",
+    "preview_negative_prompt": "blurry, duplicate, ugly, deformed, low res, watermark, text",
+    "preview_steps": 20,
+    "preview_sampler_index": 0,
+    "preview_cfg_scale": 6,
+    "preview_seed": -1,
+    "preview_width": 512,
+    "preview_height": 512,
+    "varsize": False,
+    "use_weight": False,
+})
+
+lora = Map({
+    "bucket_no_upscale": False,
+    "bucket_reso_steps": 64,
+    "cache_latents": True,
+    "caption_dropout_every_n_epochs": None,
+    "caption_dropout_rate": 0.0,
+    "caption_extension": ".txt",
+    "caption_extention": ".txt",
+    "caption_tag_dropout_rate": 0.0,
+    "clip_skip": None,
+    "color_aug": False,
+    "dataset_repeats": 1,
+    "debug_dataset": False,
+    "enable_bucket": False,
+    "face_crop_aug_range": None,
+    "flip_aug": False,
+    "full_fp16": False,
+    "gradient_accumulation_steps": 1,
+    "gradient_checkpointing": False,
+    "in_json": "",
+    "keep_tokens": None,
+    "learning_rate": 5e-05,
+    "log_prefix": None,
+    "logging_dir": None,
+    "lr_scheduler_num_cycles": 1,
+    "lr_scheduler_power": 1,
+    "lr_scheduler": "cosine",
+    "lr_warmup_steps": 0,
+    "max_bucket_reso": 1024,
+    "max_data_loader_n_workers": 8,
+    "max_grad_norm": 0.0,
+    "max_token_length": None,
+    "max_train_epochs": None,
+    "max_train_steps": 2500,
+    "mem_eff_attn": False,
+    "min_bucket_reso": 256,
+    "mixed_precision": "fp16",
+    "network_alpha": 1.0,
+    "network_args": None,
+    "network_dim": 16,
+    "network_module": "networks.lora",
+    "network_train_text_encoder_only": False,
+    "network_train_unet_only": False,
+    "network_weights": None,
+    "no_metadata": False,
+    "output_dir": "",
+    "output_name": "",
+    "persistent_data_loader_workers": False,
+    "pretrained_model_name_or_path": "",
+    "prior_loss_weight": 1.0,
+    "random_crop": False,
+    "reg_data_dir": None,
+    "resolution": "512,512",
+    "resume": None,
+    "save_every_n_epochs": None,
+    "save_last_n_epochs_state": None,
+    "save_last_n_epochs": None,
+    "save_model_as": "ckpt",
+    "save_n_epoch_ratio": None,
+    "save_precision": "fp16",
+    "save_state": False,
+    "seed": 42,
+    "shuffle_caption": False,
+    "text_encoder_lr": 5e-05,
+    "train_batch_size": 1,
+    "train_data_dir": "",
+    "training_comment": "",
+    "unet_lr": 1e-04,
+    "use_8bit_adam": False,
+    "v_parameterization": False,
+    "v2": False,
+    "vae": None,
+    "xformers": False,
+})
+
+process = Map({
+    # general settings, do not modify
+    'format': '.jpg', # image format
+    'target_size': 512, # target resolution
+    'segmentation_model': 0, # segmentation model 0/general 1/landscape
+    'segmentation_background': (192, 192, 192), # segmentation background color
+    'blur_score': 1.8, # max score for face blur detection
+    'blur_samplesize': 60, # sample size to use for blur detection
+    'similarity_score': 0.8, # maximum similarity score before image is discarded
+    'similarity_size': 64, # base similarity detection on reduced images
+    'range_score': 0.15, # min score for face color dynamicrange detection
+    # face processing settings
+    'face_score': 0.7, # min face detection score
+    'face_pad': 0.1, # pad face image percentage
+    'face_model': 1, # which face model to use 0/close-up 1/standard
+    # body processing settings
+    'body_score': 0.9, # min body detection score
+    'body_visibility': 0.5, # min visibility score for each detected body part
+    'body_parts': 15, # min number of detected body parts with sufficient visibility
+    'body_pad': 0.2,  # pad body image percentage
+    'body_model': 2, # body model to use 0/low 1/medium 2/high
+    # similarity detection settings
+    # interrogate settings
+    'interrogate': False, # interrogate images
+    'interrogate_model': ['clip', 'deepdanbooru'], # interrogate models
+    'tag_limit': 5, # number of tags to extract
+    # validations
+    # tbd
+    'face_segmentation': False, # segmentation enabled
+    'body_segmentation': False, # segmentation enabled
+})

cli/process.py

@@ -0,0 +1,327 @@
+ # pylint: disable=global-statement
+import os
+import io
+import math
+import base64
+import numpy as np
+import mediapipe as mp
+from PIL import Image, ImageOps
+from pi_heif import register_heif_opener
+from skimage.metrics import structural_similarity as ssim
+from scipy.stats import beta
+
+import util
+import sdapi
+import options
+
+face_model = None
+body_model = None
+segmentation_model = None
+all_images = []
+all_images_by_type = {}
+
+
+class Result():
+    def __init__(self, typ: str, fn: str, tag: str = None, requested: list = []):
+        self.type = typ
+        self.input = fn
+        self.output = ''
+        self.basename = ''
+        self.message = ''
+        self.image = None
+        self.caption = ''
+        self.tag = tag
+        self.tags = []
+        self.ops = []
+        self.steps = requested
+
+
+def detect_blur(image: Image):
+    # based on <https://github.com/karthik9319/Blur-Detection/>
+    bw = ImageOps.grayscale(image)
+    cx, cy = image.size[0] // 2, image.size[1] // 2
+    fft = np.fft.fft2(bw)
+    fftShift = np.fft.fftshift(fft)
+    fftShift[cy - options.process.blur_samplesize: cy + options.process.blur_samplesize, cx - options.process.blur_samplesize: cx + options.process.blur_samplesize] = 0
+    fftShift = np.fft.ifftshift(fftShift)
+    recon = np.fft.ifft2(fftShift)
+    magnitude = np.log(np.abs(recon))
+    mean = round(np.mean(magnitude), 2)
+    return mean
+
+
+def detect_dynamicrange(image: Image):
+    # based on <https://towardsdatascience.com/measuring-enhancing-image-quality-attributes-234b0f250e10>
+    data = np.asarray(image)
+    image = np.float32(data)
+    RGB = [0.299, 0.587, 0.114]
+    height, width = image.shape[:2] # pylint: disable=unsubscriptable-object
+    brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2]) # pylint: disable=unsubscriptable-object
+    hist, _ = np.histogram(brightness_image, bins=256, range=(0, 255))
+    img_brightness_pmf = hist / (height * width)
+    dist = beta(2, 2)
+    ys = dist.pdf(np.linspace(0, 1, 256))
+    ref_pmf = ys / np.sum(ys)
+    dot_product = np.dot(ref_pmf, img_brightness_pmf)
+    squared_dist_a = np.sum(ref_pmf ** 2)
+    squared_dist_b = np.sum(img_brightness_pmf ** 2)
+    res = dot_product / math.sqrt(squared_dist_a * squared_dist_b)
+    return round(res, 2)
+
+
+def detect_simmilar(image: Image):
+    img = image.resize((options.process.similarity_size, options.process.similarity_size))
+    img = ImageOps.grayscale(img)
+    data = np.array(img)
+    similarity = 0
+    for i in all_images:
+        val = ssim(data, i, data_range=255, channel_axis=None, gradient=False, full=False)
+        if val > similarity:
+            similarity = val
+    all_images.append(data)
+    return similarity
+
+
+def segmentation(res: Result):
+    global segmentation_model
+    if segmentation_model is None:
+        segmentation_model = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=options.process.segmentation_model)
+    data = np.array(res.image)
+    results = segmentation_model.process(data)
+    condition = np.stack((results.segmentation_mask,) * 3, axis=-1) > 0.1
+    background = np.zeros(data.shape, dtype=np.uint8)
+    background[:] = options.process.segmentation_background
+    data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine
+    segmented = Image.fromarray(data)
+    res.image = segmented
+    res.ops.append('segmentation')
+    return res
+
+
+def unload():
+    global face_model
+    if face_model is not None:
+        face_model = None
+    global body_model
+    if body_model is not None:
+        body_model = None
+    global segmentation_model
+    if segmentation_model is not None:
+        segmentation_model = None
+
+
+def encode(img):
+    with io.BytesIO() as stream:
+        img.save(stream, 'JPEG')
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+def reset():
+    unload()
+    global all_images_by_type
+    all_images_by_type = {}
+    global all_images
+    all_images = []
+
+
+def upscale_restore_image(res: Result, upscale: bool = False, restore: bool = False):
+    kwargs = util.Map({
+        'image': encode(res.image),
+        'codeformer_visibility': 0.0,
+        'codeformer_weight': 0.0,
+    })
+    if res.image.width >= options.process.target_size and res.image.height >= options.process.target_size:
+        upscale = False
+    if upscale:
+        kwargs.upscaler_1 = 'SwinIR_4x'
+        kwargs.upscaling_resize = 2
+        res.ops.append('upscale')
+    if restore:
+        kwargs.codeformer_visibility = 1.0
+        kwargs.codeformer_weight = 0.2
+        res.ops.append('restore')
+    if upscale or restore:
+        result = sdapi.postsync('/sdapi/v1/extra-single-image', kwargs)
+        if 'image' not in result:
+            res.message = 'failed to upscale/restore image'
+        else:
+            res.image = Image.open(io.BytesIO(base64.b64decode(result['image'])))
+    return res
+
+
+def interrogate_image(res: Result, tag: str = None):
+    caption = ''
+    tags = []
+    for model in options.process.interrogate_model:
+        json = util.Map({ 'image': encode(res.image), 'model': model })
+        result = sdapi.postsync('/sdapi/v1/interrogate', json)
+        if model == 'clip':
+            caption = result.caption if 'caption' in result else ''
+            caption = caption.split(',')[0].replace(' a ', ' ').strip()
+            if tag is not None:
+                caption = res.tag + ', ' + caption
+        if model == 'deepdanbooru':
+            tag = result.caption if 'caption' in result else ''
+            tags = tag.split(',')
+            tags = [t.replace('(', '').replace(')', '').replace('\\', '').split(':')[0].strip() for t in tags]
+            if tag is not None:
+                for t in res.tag.split(',')[::-1]:
+                    tags.insert(0, t.strip())
+    pos = 0 if len(tags) == 0 else 1
+    tags.insert(pos, caption.split(' ')[1])
+    tags = [t for t in tags if len(t) > 2]
+    if len(tags) > options.process.tag_limit:
+        tags = tags[:options.process.tag_limit]
+    res.caption = caption
+    res.tags = tags
+    res.ops.append('interrogate')
+    return res
+
+
+def resize_image(res: Result):
+    resized = res.image
+    resized.thumbnail((options.process.target_size, options.process.target_size), Image.Resampling.HAMMING)
+    res.image = resized
+    res.ops.append('resize')
+    return res
+
+
+def square_image(res: Result):
+    size = max(res.image.width, res.image.height)
+    squared = Image.new('RGB', (size, size))
+    squared.paste(res.image, ((size - res.image.width) // 2, (size - res.image.height) // 2))
+    res.image = squared
+    res.ops.append('square')
+    return res
+
+
+def process_face(res: Result):
+    res.ops.append('face')
+    global face_model
+    if face_model is None:
+        face_model = mp.solutions.face_detection.FaceDetection(min_detection_confidence=options.process.face_score, model_selection=options.process.face_model)
+    results = face_model.process(np.array(res.image))
+    if results.detections is None:
+        res.message = 'no face detected'
+        res.image = None
+        return res
+    box = results.detections[0].location_data.relative_bounding_box
+    if box.xmin < 0 or box.ymin < 0 or (box.width - box.xmin) > 1 or (box.height - box.ymin) > 1:
+        res.message = 'face out of frame'
+        res.image = None
+        return res
+    x = max(0, (box.xmin - options.process.face_pad / 2) * res.image.width)
+    y = max(0, (box.ymin - options.process.face_pad / 2)* res.image.height)
+    w = min(res.image.width, (box.width + options.process.face_pad) * res.image.width)
+    h = min(res.image.height, (box.height + options.process.face_pad) * res.image.height)
+    x = max(0, x)
+    res.image = res.image.crop((x, y, x + w, y + h))
+    return res
+
+
+def process_body(res: Result):
+    res.ops.append('body')
+    global body_model
+    if body_model is None:
+        body_model = mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=options.process.body_score, model_complexity=options.process.body_model)
+    results = body_model.process(np.array(res.image))
+    if results.pose_landmarks is None:
+        res.message = 'no body detected'
+        res.image = None
+        return res
+    x0 = [res.image.width * (i.x - options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    y0 = [res.image.height * (i.y - options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    x1 = [res.image.width * (i.x + options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    y1 = [res.image.height * (i.y + options.process.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > options.process.body_visibility]
+    if len(x0) < options.process.body_parts:
+        res.message = f'insufficient body parts detected: {len(x0)}'
+        res.image = None
+        return res
+    res.image = res.image.crop((max(0, min(x0)), max(0, min(y0)), min(res.image.width, max(x1)), min(res.image.height, max(y1))))
+    return res
+
+
+def process_original(res: Result):
+    res.ops.append('original')
+    return res
+
+
+def save_image(res: Result, folder: str):
+    if res.image is None or folder is None:
+        return res
+    all_images_by_type[res.type] = all_images_by_type.get(res.type, 0) + 1
+    res.basename = os.path.basename(res.input).split('.')[0]
+    res.basename = str(all_images_by_type[res.type]).rjust(3, '0') + '-' + res.type + '-' + res.basename
+    res.basename = os.path.join(folder, res.basename)
+    res.output = res.basename + options.process.format
+    res.image.save(res.output)
+    res.image.close()
+    res.ops.append('save')
+    return res
+
+
+def file(filename: str, folder: str, tag = None, requested = []):
+    # initialize result dict
+    res = Result(fn = filename, typ='unknown', tag=tag, requested = requested)
+    # open image
+    try:
+        register_heif_opener()
+        res.image = Image.open(filename)
+        if res.image.mode == 'RGBA':
+            res.image = res.image.convert('RGB')
+        res.image = ImageOps.exif_transpose(res.image) # rotate image according to EXIF orientation
+    except Exception as e:
+        res.message = f'error opening: {e}'
+        return res
+    # primary steps
+    if 'face' in requested:
+        res.type = 'face'
+        res = process_face(res)
+    elif 'body' in requested:
+        res.type = 'body'
+        res = process_body(res)
+    elif 'original' in requested:
+        res.type = 'original'
+        res = process_original(res)
+    # validation steps
+    if res.image is None:
+        return res
+    if 'blur' in requested:
+        res.ops.append('blur')
+        val = detect_blur(res.image)
+        if val > options.process.blur_score:
+            res.message = f'blur check failed: {val}'
+            res.image = None
+    if 'range' in requested:
+        res.ops.append('range')
+        val = detect_dynamicrange(res.image)
+        if val < options.process.range_score:
+            res.message = f'dynamic range check failed: {val}'
+            res.image = None
+    if 'similarity' in requested:
+        res.ops.append('similarity')
+        val = detect_simmilar(res.image)
+        if val > options.process.similarity_score:
+            res.message = f'dynamic range check failed: {val}'
+            res.image = None
+    if res.image is None:
+        return res
+    # post processing steps
+    res = upscale_restore_image(res, 'upscale' in requested, 'restore' in requested)
+    if res.image.width < options.process.target_size or res.image.height < options.process.target_size:
+        res.message = f'low resolution: [{res.image.width}, {res.image.height}]'
+        res.image = None
+        return res
+    if 'interrogate' in requested:
+        res = interrogate_image(res, tag)
+    if 'resize' in requested:
+        res = resize_image(res)
+    if 'square' in requested:
+        res = square_image(res)
+    if 'segment' in requested:
+        res = segmentation(res)
+    # finally save image
+    res = save_image(res, folder)
+    return res

cli/random.json

@@ -0,0 +1,31 @@
+{
+  "prompts": [
+    "<style> of <embedding> <place>, high detailed, by <artist>, <suffix>"
+  ],
+  "negative": [
+    "watermark, fog, clouds, blurry, duplicate, deformed, mutation"
+  ],
+  "places": [
+    "standing in the city", "on a spaceship", "in fantasy landscape", "on a shore", "in a forest", "in winter wonderland"
+  ],
+  "embeddings": [
+    "man", "man next to a beautiful girl", "man next to a car", "beautiful girl", "sexy naked girl", "cute girl holding a flower", "beautiful robot",
+    "young korean girl with medium-length white hair", "monster", "pin up girl",
+    "man vlado", "beutiful girl ana", "man lee", "beautiful girl abby"
+  ],
+  "artists": [
+    "John Salminen", "Greg Rutkowski", "Akihiko Yoshida", "Alejandro Burdisio", "Artgerm", "Patrick Brown", "Walt Disney", "Neal Adams", "Jeremy Chong",
+    "Chris Rallis", "Roy Lichtenstein", "Claude Monet", "Jon Whitcomb", "Pablo Picasso", "Raymond Leech", "Tom Lovell", "Noriyoshi Ohrai", "Shingei",
+    "Helmut Newton", "Maciej Kuciara", "Daniel F. Gerhartz", "Stephan Martinière", "Magali Villeneuve", "Carne Griffiths", "Alberto Seveso",
+    "Vincent Van Gogh", "WLOP", "Frank Xavier Leyendecker", "Peter Lindbergh", "Nick Gentry", "Howard Chandler Christy", "Raphael", "Henri Matisse"
+  ],
+  "styles": [
+    "illustration", "painting", "portrait", "photograph", "drawing", "sketch", "pencil sketch", "3d render", "cartoon", "anime", "scribbles", "pop art",
+    "ink painting", "steampunk illustration", "dc comics illustration", "marvel comics", "vray render", "photoillustration", "pixar", "marble sculpture",
+    "bronze sculpture", "christmas theme"
+  ],
+  "suffixes": [
+    "cinematic lighting", "artstation", "fineart", "cinematic", "photorealistic", "soft light", "sharp focus", "bokeh", "dreamlike", "semirealism",
+    "colorful", "black and white", "intricate", "elegant"
+  ]
+}

cli/requirements.txt

@@ -0,0 +1,7 @@
+aiohttp
+mediapipe
+extcolors
+colormap
+filetype
+albumentations
+matplotlib

cli/run-benchmark.py

@@ -0,0 +1,149 @@
+#!/usr/bin/env python
+"""
+sd api txt2img benchmark
+"""
+import os
+import asyncio
+import base64
+import io
+import json
+import time
+import argparse
+from PIL import Image
+import sdapi
+from util import Map, log
+
+
+oom = 0
+args = None
+options = None
+
+
+async def txt2img():
+    t0 = time.perf_counter()
+    data = {}
+    try:
+        data = await sdapi.post('/sdapi/v1/txt2img', options)
+    except Exception:
+        return -1
+    if 'error' in data:
+        return -1
+    if 'info' in data:
+        info = Map(json.loads(data['info']))
+    else:
+        return 0
+    log.debug({ 'info': info })
+    if options['batch_size'] != len(data['images']):
+        log.error({ 'requested': options['batch_size'], 'received': len(data['images']) })
+        return 0
+    for i in range(len(data['images'])):
+        data['images'][i] = Image.open(io.BytesIO(base64.b64decode(data['images'][i].split(',',1)[0])))
+        if args.save:
+            fn = os.path.join(args.save, f'benchmark-{i}-{len(data["images"])}.png')
+            data["images"][i].save(fn)
+            log.debug({ 'save': fn })
+    log.debug({ "images": data["images"] })
+    t1 = time.perf_counter()
+    return t1 - t0
+
+
+def memstats():
+    mem = sdapi.getsync('/sdapi/v1/memory')
+    cpu = mem.get('ram', 'unavailable')
+    gpu = mem.get('cuda', 'unavailable')
+    if 'active' in gpu:
+        gpu['session'] = gpu.pop('active')
+    if 'reserved' in gpu:
+        gpu.pop('allocated')
+        gpu.pop('reserved')
+        gpu.pop('inactive')
+    if 'events' in gpu:
+        global oom # pylint: disable=global-statement
+        oom = gpu['events']['oom']
+        gpu.pop('events')
+    return cpu, gpu
+
+
+def gb(val: float):
+    return round(val / 1024 / 1024 / 1024, 2)
+
+
+async def main():
+    sdapi.quiet = True
+    await sdapi.session()
+    await sdapi.interrupt()
+    ver = await sdapi.get("/sdapi/v1/version")
+    log.info({ 'version': ver})
+    platform = await sdapi.get("/sdapi/v1/platform")
+    log.info({ 'platform': platform })
+    opts = await sdapi.get('/sdapi/v1/options')
+    opts = Map(opts)
+    log.info({ 'model': opts.sd_model_checkpoint })
+    cpu, gpu = memstats()
+    log.info({ 'system': { 'cpu': cpu, 'gpu': gpu }})
+    batch = [1, 1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256]
+    batch = [b for b in batch if b <= args.maxbatch]
+    log.info({"batch-sizes": batch})
+    for i in range(len(batch)):
+        if oom > 0:
+            continue
+        options['batch_size'] = batch[i]
+        warmup = await txt2img()
+        ts = await txt2img()
+        if i == 0:
+            ts += warmup
+        if ts > 0.01: # cannot be faster than 10ms per run
+            await asyncio.sleep(0)
+            cpu, gpu = memstats()
+            if i == 0:
+                log.info({ 'warmup': round(ts, 2) })
+            else:
+                peak = gpu['system']['used'] # gpu['session']['peak'] if 'session' in gpu else 0
+                log.info({ 'batch': batch[i], 'its': round(options.steps / (ts / batch[i]), 2), 'img': round(ts / batch[i], 2), 'wall': round(ts, 2), 'peak': gb(peak), 'oom': oom > 0 })
+        else:
+            await asyncio.sleep(10)
+            cpu, gpu = memstats()
+            log.info({ 'batch': batch[i], 'result': 'error', 'gpu': gpu, 'oom': oom > 0 })
+            break
+    if oom > 0:
+        log.info({ 'benchmark': 'ended with oom so you should probably restart your automatic server now' })
+    await sdapi.close()
+
+
+if __name__ == '__main__':
+    log.info({ 'run-benchmark' })
+    parser = argparse.ArgumentParser(description = 'run-benchmark')
+    parser.add_argument("--steps", type=int, default=50, required=False, help="steps")
+    parser.add_argument("--sampler", type=str, default='Euler a', required=False, help="Use specific sampler")
+    parser.add_argument("--prompt", type=str, default='photo of two dice on a table', required=False, help="prompt")
+    parser.add_argument("--negative", type=str, default='foggy, blurry', required=False, help="prompt")
+    parser.add_argument("--maxbatch", type=int, default=16, required=False, help="max batch size")
+    parser.add_argument("--width", type=int, default=512, required=False, help="width")
+    parser.add_argument("--height", type=int, default=512, required=False, help="height")
+    parser.add_argument('--debug', default = False, action='store_true', help = 'debug logging')
+    parser.add_argument('--taesd', default = False, action='store_true', help = 'use taesd as vae')
+    parser.add_argument("--save", type=str, default='', required=False, help="save images to folder")
+    args = parser.parse_args()
+    if args.debug:
+        log.setLevel('DEBUG')
+    options = Map(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative,
+            "steps": args.steps,
+            "sampler_name": args.sampler,
+            "width": args.width,
+            "height": args.height,
+            "full_quality": not args.taesd,
+            "cfg_scale": 0,
+            "batch_size": 1,
+            "n_iter": 1,
+            "seed": -1,
+        }
+    )
+    log.info({"options": options})
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        log.warning({ 'interrupted': 'keyboard request' })
+        sdapi.interruptsync()

cli/sdapi.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python
+#pylint: disable=redefined-outer-name
+"""
+helper methods that creates HTTP session with managed connection pool
+provides async HTTP get/post methods and several helper methods
+"""
+
+import io
+import os
+import sys
+import ssl
+import base64
+import asyncio
+import logging
+import aiohttp
+import requests
+import urllib3
+from PIL import Image
+from util import Map, log
+from rich import print # pylint: disable=redefined-builtin
+
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860") # api url root
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+use_session = True
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+ssl.create_default_context = ssl._create_unverified_context # pylint: disable=protected-access
+timeout = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+sess = None
+quiet = False
+BaseThreadPolicy = asyncio.WindowsSelectorEventLoopPolicy if sys.platform == "win32" and hasattr(asyncio, "WindowsSelectorEventLoopPolicy") else asyncio.DefaultEventLoopPolicy
+
+
+class AnyThreadEventLoopPolicy(BaseThreadPolicy):
+    def get_event_loop(self) -> asyncio.AbstractEventLoop:
+        try:
+            return super().get_event_loop()
+        except (RuntimeError, AssertionError):
+            loop = self.new_event_loop()
+            self.set_event_loop(loop)
+            return loop
+
+asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy())
+
+
+def authsync():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return aiohttp.BasicAuth(sd_username, sd_password)
+    return None
+
+
+async def result(req):
+    if req.status != 200:
+        if not quiet:
+            log.error({ 'request error': req.status, 'reason': req.reason, 'url': req.url })
+        if not use_session and sess is not None:
+            await sess.close()
+        return Map({ 'error': req.status, 'reason': req.reason, 'url': req.url })
+    else:
+        json = await req.json()
+        if isinstance(json, list):
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        log.debug({ 'request': req.status, 'url': req.url, 'reason': req.reason })
+        return res
+
+
+def resultsync(req: requests.Response):
+    if req.status_code != 200:
+        if not quiet:
+            log.error({ 'request error': req.status_code, 'reason': req.reason, 'url': req.url })
+        return Map({ 'error': req.status_code, 'reason': req.reason, 'url': req.url })
+    else:
+        json = req.json()
+        if isinstance(json, list):
+            res = json
+        elif json is None:
+            res = {}
+        else:
+            res = Map(json)
+        log.debug({ 'request': req.status_code, 'url': req.url, 'reason': req.reason })
+        return res
+
+
+async def get(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    sess = sess if sess is not None else await session()
+    try:
+        async with sess.get(url=endpoint, json=json, verify_ssl=False) as req:
+            res = await result(req)
+            return res
+    except Exception as err:
+        log.error({ 'session': err })
+        return {}
+
+
+def getsync(endpoint: str, json: dict = None):
+    try:
+        req = requests.get(f'{sd_url}{endpoint}', json=json, verify=False, auth=authsync()) # pylint: disable=missing-timeout
+        res = resultsync(req)
+        return res
+    except Exception as err:
+        log.error({ 'session': err })
+        return {}
+
+
+async def post(endpoint: str, json: dict = None):
+    global sess # pylint: disable=global-statement
+    # sess = sess if sess is not None else await session()
+    if sess and not sess.closed:
+        await sess.close()
+    sess = await session()
+    try:
+        async with sess.post(url=endpoint, json=json, verify_ssl=False) as req:
+            res = await result(req)
+            return res
+    except Exception as err:
+        log.error({ 'session': err })
+        return {}
+
+
+def postsync(endpoint: str, json: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json=json, verify=False, auth=authsync()) # pylint: disable=missing-timeout
+    res = resultsync(req)
+    return res
+
+
+async def interrupt():
+    res = await get('/sdapi/v1/progress?skip_current_image=true')
+    if 'state' in res and res.state.job_count > 0:
+        log.debug({ 'interrupt': res.state })
+        res = await post('/sdapi/v1/interrupt')
+        await asyncio.sleep(1)
+        return res
+    else:
+        log.debug({ 'interrupt': 'idle' })
+        return { 'interrupt': 'idle' }
+
+
+def interruptsync():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    if 'state' in res and res.state.job_count > 0:
+        log.debug({ 'interrupt': res.state })
+        res = postsync('/sdapi/v1/interrupt')
+        return res
+    else:
+        log.debug({ 'interrupt': 'idle' })
+        return { 'interrupt': 'idle' }
+
+
+async def progress():
+    res = await get('/sdapi/v1/progress?skip_current_image=false')
+    try:
+        if res is not None and res.get('current_image', None) is not None:
+            res.current_image = Image.open(io.BytesIO(base64.b64decode(res['current_image'])))
+    except Exception:
+        pass
+    log.debug({ 'progress': res })
+    return res
+
+
+def progresssync():
+    res = getsync('/sdapi/v1/progress?skip_current_image=true')
+    log.debug({ 'progress': res })
+    return res
+
+
+def get_log():
+    res = getsync('/sdapi/v1/log')
+    for line in res:
+        log.debug(line)
+    return res
+
+
+def get_info():
+    import time
+    t0 = time.time()
+    res = getsync('/sdapi/v1/system-info/status?full=true&refresh=true')
+    t1 = time.time()
+    print({ 'duration': 1000 * round(t1-t0, 3), **res })
+    return res
+
+
+def options():
+    opts = getsync('/sdapi/v1/options')
+    flags = getsync('/sdapi/v1/cmd-flags')
+    return { 'options': opts, 'flags': flags }
+
+
+def shutdown():
+    try:
+        postsync('/sdapi/v1/shutdown')
+    except Exception as e:
+        log.info({ 'shutdown': e })
+
+
+async def session():
+    global sess # pylint: disable=global-statement
+    time = aiohttp.ClientTimeout(total = None, sock_connect = 10, sock_read = None) # default value is 5 minutes, we need longer for training
+    sess = aiohttp.ClientSession(timeout = time, base_url = sd_url, auth=auth())
+    log.debug({ 'sdapi': 'session created', 'endpoint': sd_url })
+    """
+    sess = await aiohttp.ClientSession(timeout = timeout).__aenter__()
+    try:
+        async with sess.get(url = f'{sd_url}/') as req:
+            log.debug({ 'sdapi': 'session created', 'endpoint': sd_url })
+    except Exception as e:
+        log.error({ 'sdapi': e })
+        await asyncio.sleep(0)
+        await sess.__aexit__(None, None, None)
+        sess = None
+    return sess
+    """
+    return sess
+
+
+async def close():
+    if sess is not None:
+        await asyncio.sleep(0)
+        await sess.close()
+        await sess.__aexit__(None, None, None)
+        log.debug({ 'sdapi': 'session closed', 'endpoint': sd_url })
+
+
+if __name__ == "__main__":
+    sys.argv.pop(0)
+    log.setLevel(logging.DEBUG)
+    if 'interrupt' in sys.argv:
+        asyncio.run(interrupt())
+    elif 'progress' in sys.argv:
+        asyncio.run(progress())
+    elif 'progresssync' in sys.argv:
+        progresssync()
+    elif 'options' in sys.argv:
+        opt = options()
+        log.debug({ 'options' })
+        import json
+        print(json.dumps(opt['options'], indent = 2))
+        log.debug({ 'cmd-flags' })
+        print(json.dumps(opt['flags'], indent = 2))
+    elif 'log' in sys.argv:
+        get_log()
+    elif 'info' in sys.argv:
+        get_info()
+    elif 'shutdown' in sys.argv:
+        shutdown()
+    else:
+        res = getsync(sys.argv[0])
+        print(res)
+    asyncio.run(close(), debug=True)
+    asyncio.run(asyncio.sleep(0.5))

cli/simple-img2img.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+import os
+import io
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+options = {
+    "save_images": False,
+    "send_images": True,
+}
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def encode(f):
+    image = Image.open(f)
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
+    with io.BytesIO() as stream:
+        image.save(stream, 'JPEG')
+        image.close()
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+def generate(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    if args.model is not None:
+        post('/sdapi/v1/options', { 'sd_model_checkpoint': args.model })
+        post('/sdapi/v1/reload-checkpoint') # needed if running in api-only to trigger new model load
+    options['prompt'] = args.prompt
+    options['negative_prompt'] = args.negative
+    options['steps'] = int(args.steps)
+    options['seed'] = int(args.seed)
+    options['sampler_name'] = args.sampler
+    options['init_images'] = [encode(args.init)]
+    image = Image.open(args.init)
+    options['width'] = image.width
+    options['height'] = image.height
+    image.close()
+    if args.mask is not None:
+        options['mask'] = encode(args.mask)
+    data = post('/sdapi/v1/img2img', options)
+    t1 = time.time()
+    if 'images' in data:
+        for i in range(len(data['images'])):
+            b64 = data['images'][i].split(',',1)[0]
+            info = data['info']
+            image = Image.open(io.BytesIO(base64.b64decode(b64)))
+            log.info(f'received image: size={image.size} time={t1-t0:.2f} info="{info}"')
+            if args.output:
+                image.save(args.output)
+                log.info(f'image saved: size={image.size} filename={args.output}')
+
+    else:
+        log.warning(f'no images received: {data}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-img2img')
+    parser.add_argument('--init', required=True, help='init image')
+    parser.add_argument('--mask', required=False, help='mask image')
+    parser.add_argument('--prompt', required=False, default='', help='prompt text')
+    parser.add_argument('--negative', required=False, default='', help='negative prompt text')
+    parser.add_argument('--steps', required=False, default=20, help='number of steps')
+    parser.add_argument('--seed', required=False, default=-1, help='initial seed')
+    parser.add_argument('--sampler', required=False, default='Euler a', help='sampler name')
+    parser.add_argument('--output', required=False, default=None, help='output image file')
+    parser.add_argument('--model', required=False, help='model name')
+    args = parser.parse_args()
+    log.info(f'img2img: {args}')
+    generate(args)

cli/simple-info.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python
+import os
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def get(endpoint: str, dct: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json=dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def info(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    with open(args.input, 'rb') as f:
+        content = f.read()
+    data = post('/sdapi/v1/png-info', { 'image': base64.b64encode(content).decode() })
+    t1 = time.time()
+    log.info(f'received: {data} time={t1-t0:.2f}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-info')
+    parser.add_argument('--input', required=True, help='input image')
+    args = parser.parse_args()
+    log.info(f'info: {args}')
+    info(args)

cli/simple-mask.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+import io
+import os
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def get(endpoint: str, dct: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json=dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def info(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    with open(args.input, 'rb') as f:
+        image = base64.b64encode(f.read()).decode()
+    if args.mask:
+        with open(args.mask, 'rb') as f:
+            mask = base64.b64encode(f.read()).decode()
+    else:
+        mask = None
+    options = get('/sdapi/v1/masking')
+    log.info(f'options: {options}')
+    req = {
+        'image': image,
+        'mask': mask,
+        'type': args.type or 'Composite',
+        'params': { 'auto_mask': 'Grayscale' if mask is None else None },
+    }
+    data = post('/sdapi/v1/mask', req)
+    t1 = time.time()
+    if 'mask' in data:
+        b64 = data['mask'].split(',',1)[0]
+        image = Image.open(io.BytesIO(base64.b64decode(b64)))
+        log.info(f'received image: size={image.size} time={t1-t0:.2f}')
+        if args.output:
+            image.save(args.output)
+            log.info(f'saved image: fn={args.output}')
+    else:
+        log.info(f'received: {data} time={t1-t0:.2f}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-info')
+    parser.add_argument('--input', required=True, help='input image')
+    parser.add_argument('--mask', required=False, help='input mask')
+    parser.add_argument('--type', required=False, help='output mask type')
+    parser.add_argument('--output', required=False, help='output image')
+    args = parser.parse_args()
+    log.info(f'info: {args}')
+    info(args)

cli/simple-preprocess.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+import io
+import os
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def get(endpoint: str, dct: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json=dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def info(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    with open(args.input, 'rb') as f:
+        content = f.read()
+    models = get('/sdapi/v1/preprocessors')
+    log.info(f'models: {models}')
+    req = {
+        'model': args.model or 'Canny',
+        'image': base64.b64encode(content).decode(),
+        'config': { 'low_threshold': 50 },
+    }
+    data = post('/sdapi/v1/preprocess', req)
+    t1 = time.time()
+    if 'image' in data:
+        b64 = data['image'].split(',',1)[0]
+        image = Image.open(io.BytesIO(base64.b64decode(b64)))
+        log.info(f'received image: size={image.size} time={t1-t0:.2f}')
+        if args.output:
+            image.save(args.output)
+            log.info(f'saved image: fn={args.output}')
+    else:
+        log.info(f'received: {data} time={t1-t0:.2f}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-info')
+    parser.add_argument('--input', required=True, help='input image')
+    parser.add_argument('--model', required=True, help='preprocessing model')
+    parser.add_argument('--output', required=False, help='output image')
+    args = parser.parse_args()
+    log.info(f'info: {args}')
+    info(args)

cli/simple-txt2img.js

@@ -0,0 +1,63 @@
+#!/usr/bin/env node
+
+// simple nodejs script to test sdnext api
+
+const fs = require('fs'); // eslint-disable-line no-undef
+const process = require('process'); // eslint-disable-line no-undef
+
+const sd_url = process.env.SDAPI_URL || 'http://127.0.0.1:7860';
+const sd_username = process.env.SDAPI_USR;
+const sd_password = process.env.SDAPI_PWD;
+const sd_options = {
+  // first pass
+  prompt: 'city at night',
+  negative_prompt: 'foggy, blurry',
+  sampler_name: 'UniPC',
+  seed: -1,
+  steps: 20,
+  batch_size: 1,
+  n_iter: 1,
+  cfg_scale: 6,
+  width: 512,
+  height: 512,
+  // enable second pass
+  enable_hr: true,
+  // second pass: upscale
+  hr_upscaler: 'SCUNet GAN',
+  hr_scale: 2.0,
+  // second pass: hires
+  hr_force: true,
+  hr_second_pass_steps: 20,
+  hr_sampler_name: 'UniPC',
+  denoising_strength: 0.5,
+  // second pass: refiner
+  refiner_steps: 5,
+  refiner_start: 0.8,
+  refiner_prompt: '',
+  refiner_negative: '',
+  // api return options
+  save_images: false,
+  send_images: true,
+};
+
+async function main() {
+  const method = 'POST';
+  const headers = new Headers();
+  const body = JSON.stringify(sd_options);
+  headers.set('Content-Type', 'application/json');
+  if (sd_username && sd_password) headers.set({ Authorization: `Basic ${btoa('sd_username:sd_password')}` });
+  const res = await fetch(`${sd_url}/sdapi/v1/txt2img`, { method, headers, body });
+  if (res.status !== 200) {
+    console.log('Error', res.status);
+  } else {
+    const json = await res.json();
+    console.log('result:', json.info);
+    for (const i in json.images) { // eslint-disable-line guard-for-in
+      const f = `/tmp/test-{${i}.jpg`;
+      fs.writeFileSync(f, atob(json.images[i]), 'binary');
+      console.log('image saved:', f);
+    }
+  }
+}
+
+main();

cli/simple-txt2img.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+import io
+import os
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+options = {
+    "save_images": False,
+    "send_images": True,
+}
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def generate(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    if args.model is not None:
+        post('/sdapi/v1/options', { 'sd_model_checkpoint': args.model })
+        post('/sdapi/v1/reload-checkpoint') # needed if running in api-only to trigger new model load
+    options['prompt'] = args.prompt
+    options['negative_prompt'] = args.negative
+    options['steps'] = int(args.steps)
+    options['seed'] = int(args.seed)
+    options['sampler_name'] = args.sampler
+    options['width'] = int(args.width)
+    options['height'] = int(args.height)
+    data = post('/sdapi/v1/txt2img', options)
+    t1 = time.time()
+    if 'images' in data:
+        for i in range(len(data['images'])):
+            b64 = data['images'][i].split(',',1)[0]
+            image = Image.open(io.BytesIO(base64.b64decode(b64)))
+            info = data['info']
+            log.info(f'image received: size={image.size} time={t1-t0:.2f} info="{info}"')
+            if args.output:
+                image.save(args.output)
+                log.info(f'image saved: size={image.size} filename={args.output}')
+    else:
+        log.warning(f'no images received: {data}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-txt2img')
+    parser.add_argument('--prompt', required=False, default='', help='prompt text')
+    parser.add_argument('--negative', required=False, default='', help='negative prompt text')
+    parser.add_argument('--width', required=False, default=512, help='image width')
+    parser.add_argument('--height', required=False, default=512, help='image height')
+    parser.add_argument('--steps', required=False, default=20, help='number of steps')
+    parser.add_argument('--seed', required=False, default=-1, help='initial seed')
+    parser.add_argument('--sampler', required=False, default='Euler a', help='sampler name')
+    parser.add_argument('--output', required=False, default=None, help='output image file')
+    parser.add_argument('--model', required=False, help='model name')
+    args = parser.parse_args()
+    log.info(f'txt2img: {args}')
+    generate(args)

cli/simple-upscale.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+import os
+import io
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def get(endpoint: str, dct: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json=dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def encode(f):
+    image = Image.open(f)
+    if image.mode == 'RGBA':
+        image = image.convert('RGB')
+    log.info(f'encoding image: {image}')
+    with io.BytesIO() as stream:
+        image.save(stream, 'JPEG')
+        image.close()
+        values = stream.getvalue()
+        encoded = base64.b64encode(values).decode()
+        return encoded
+
+
+def upscale(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    # options['mask'] = encode(args.mask)
+    upscalers = get('/sdapi/v1/upscalers')
+    upscalers = [u['name'] for u in upscalers]
+    log.info(f'upscalers: {upscalers}')
+    options = {
+        "save_images": False,
+        "send_images": True,
+        'image': encode(args.input),
+        'upscaler_1': args.upscaler,
+        'resize_mode': 0, # rescale_by
+        'upscaling_resize': args.scale,
+
+    }
+    data = post('/sdapi/v1/extra-single-image', options)
+    t1 = time.time()
+    if 'image' in data:
+        b64 = data['image'].split(',',1)[0]
+        image = Image.open(io.BytesIO(base64.b64decode(b64)))
+        image.save(args.output)
+        log.info(f'received: image={image} file={args.output} time={t1-t0:.2f}')
+    else:
+        log.warning(f'no images received: {data}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-upscale')
+    parser.add_argument('--input', required=True, help='input image')
+    parser.add_argument('--output', required=True, help='output image')
+    parser.add_argument('--upscaler', required=False, default='Nearest', help='upscaler name')
+    parser.add_argument('--scale', required=False, default=2, help='upscaler scale')
+    args = parser.parse_args()
+    log.info(f'upscale: {args}')
+    upscale(args)

cli/torch-compile.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python
+# pylint: disable=cell-var-from-loop
+"""
+Test Torch Dynamo functionality and backends
+"""
+import json
+import warnings
+
+import numpy as np
+import torch
+from torchvision.models import resnet18
+
+
+print('torch:', torch.__version__)
+try:
+    # must be imported explicitly or namespace is not found
+    import torch._dynamo as dynamo # pylint: disable=ungrouped-imports
+except Exception as err:
+    print('torch without dynamo support', err)
+
+
+N_ITERS = 20
+torch._dynamo.config.verbose=True # pylint: disable=protected-access
+warnings.filterwarnings('ignore', category=UserWarning) # disable those for now as many backends reports tons
+# torch.set_float32_matmul_precision('high') # enable to test in fp32
+
+
+def timed(fn): # returns the result of running `fn()` and the time it took for `fn()` to run in ms using CUDA events
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    result = fn()
+    end.record()
+    torch.cuda.synchronize()
+    return result, start.elapsed_time(end)
+
+
+def generate_data(b):
+    return (
+        torch.randn(b, 3, 128, 128).to(torch.float32).cuda(),
+        torch.randint(1000, (b,)).cuda(),
+    )
+
+
+def init_model():
+    return resnet18().to(torch.float32).cuda()
+
+
+def evaluate(mod, val):
+    return mod(val)
+
+
+if __name__ == '__main__':
+    # first pass, dynamo is going to be slower as it compiles
+    model = init_model()
+    inp = generate_data(16)[0]
+
+    # repeat test
+    results = {}
+    times = []
+    print('eager initial eval:', timed(lambda: evaluate(model, inp))[1])
+    for _i in range(N_ITERS):
+        inp = generate_data(16)[0]
+        _res, time = timed(lambda: evaluate(model, inp)) # noqa: B023
+        times.append(time)
+    results['default'] = np.median(times)
+
+    print('dynamo available backends:', dynamo.list_backends())
+    for backend in dynamo.list_backends():
+        try:
+            # required before changing backends
+            torch._dynamo.reset() # pylint: disable=protected-access
+            eval_dyn = dynamo.optimize(backend)(evaluate)
+            print('dynamo initial eval:', backend, timed(lambda: eval_dyn(model, inp))[1]) # noqa: B023
+            times = []
+            for _i in range(N_ITERS):
+                inp = generate_data(16)[0]
+                _res, time = timed(lambda: eval_dyn(model, inp)) # noqa: B023
+                times.append(time)
+            results[backend] = np.median(times)
+        except Exception as err:
+            lines = str(err).split('\n')
+            print('dyanmo backend failed:', backend, lines[0]) # print just first error line as backtraces can be quite long
+            results[backend] = 'error'
+
+    # print stats
+    print(json.dumps(results, indent = 4))
+
+"""
+Reference: <https://github.com/pytorch/pytorch/blob/4f4b62e4a255708e928445b6502139d5962974fa/docs/source/dynamo/get-started.rst>
+Training & Inference backends:
+    dynamo.optimize("inductor") - Uses TorchInductor backend with AotAutograd and cudagraphs by leveraging codegened Triton kernels
+    dynamo.optimize("aot_nvfuser") - nvFuser with AotAutograd
+    dynamo.optimize("aot_cudagraphs") - cudagraphs with AotAutograd
+Inference-only backends:
+    dynamo.optimize("ofi") - Uses Torchscript optimize_for_inference
+    dynamo.optimize("fx2trt") - Uses Nvidia TensorRT for inference optimizations
+    dynamo.optimize("onnxrt") - Uses ONNXRT for inference on CPU/GPU
+"""

cli/train.py

@@ -0,0 +1,443 @@
+#!/usr/bin/env python
+
+"""
+Examples:
+- sd15:    train.py --type lora --tag girl --comments sdnext --input ~/generative/Input/mia --process original,interrogate,resize --name mia
+- sdxl:    train.py --type lora --tag girl --comments sdnext --input ~/generative/Input/mia --process original,interrogate,resize --precision fp32 --optimizer Adafactor --sdxl --name miaxl
+- offline: train.py --type lora --tag girl --comments sdnext --input ~/generative/Input/mia --model /home/vlado/dev/sdnext/models/Stable-diffusion/sdxl/miaanimeSFWNSFWSDXL_v40.safetensors --dir /home/vlado/dev/sdnext/models/Lora/ --precision fp32 --optimizer Adafactor --sdxl --name miaxl
+"""
+
+# system imports
+import os
+import re
+import gc
+import sys
+import json
+import shutil
+import pathlib
+import asyncio
+import logging
+import tempfile
+import argparse
+
+# local imports
+import util
+import sdapi
+import options
+
+
+# globals
+args = None
+log = logging.getLogger('train')
+valid_steps = ['original', 'face', 'body', 'blur', 'range', 'upscale', 'restore', 'interrogate', 'resize', 'square', 'segment']
+log_file = os.path.join(os.path.dirname(__file__), 'train.log')
+server_ok = False
+
+# methods
+
+def setup_logging():
+    from rich.theme import Theme
+    from rich.logging import RichHandler
+    from rich.console import Console
+    from rich.pretty import install as pretty_install
+    from rich.traceback import install as traceback_install
+    console = Console(log_time=True, log_time_format='%H:%M:%S-%f', theme=Theme({
+        "traceback.border": "black",
+        "traceback.border.syntax_error": "black",
+        "inspect.value.border": "black",
+    }))
+    # logging.getLogger("urllib3").setLevel(logging.ERROR)
+    # logging.getLogger("httpx").setLevel(logging.ERROR)
+    level = logging.DEBUG if args.debug else logging.INFO
+    logging.basicConfig(level=logging.ERROR, format='%(asctime)s | %(name)s | %(levelname)s | %(module)s | %(message)s', filename=log_file, filemode='a', encoding='utf-8', force=True)
+    log.setLevel(logging.DEBUG) # log to file is always at level debug for facility `sd`
+    pretty_install(console=console)
+    traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, suppress=[])
+    rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=level, console=console)
+    rh.set_name(level)
+    while log.hasHandlers() and len(log.handlers) > 0:
+        log.removeHandler(log.handlers[0])
+    log.addHandler(rh)
+
+
+def mem_stats():
+    gc.collect()
+    import torch
+    if torch.cuda.is_available():
+        with torch.no_grad():
+            torch.cuda.empty_cache()
+        with torch.cuda.device('cuda'):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+    mem = util.get_memory()
+    peak = { 'active': mem['gpu-active']['peak'], 'allocated': mem['gpu-allocated']['peak'], 'reserved': mem['gpu-reserved']['peak'] }
+    log.debug(f"memory cpu: {mem.ram} gpu current: {mem.gpu} gpu peak: {peak}")
+
+
+def parse_args():
+    global args # pylint: disable=global-statement
+    parser = argparse.ArgumentParser(description = 'SD.Next Train')
+
+    group_server = parser.add_argument_group('Server')
+    group_server.add_argument('--server', type=str, default='http://127.0.0.1:7860', required=False, help='server url, default: %(default)s')
+    group_server.add_argument('--user', type=str, default=None, required=False, help='server url, default: %(default)s')
+    group_server.add_argument('--password', type=str, default=None, required=False, help='server url, default: %(default)s')
+    group_server.add_argument('--dir', type=str, default=None, required=False, help='folder with trained networks, default: use server setting')
+
+    group_main = parser.add_argument_group('Main')
+    group_main.add_argument('--type', type=str, choices=['embedding', 'ti', 'lora', 'lyco', 'dreambooth', 'hypernetwork'], default=None, required=True, help='training type')
+    group_main.add_argument('--model', type=str, default='', required=False, help='base model to use for training, default: current loaded model')
+    group_main.add_argument('--name', type=str, default=None, required=True, help='output filename')
+    group_main.add_argument('--tag', type=str, default='person', required=False, help='primary tags, default: %(default)s')
+    group_main.add_argument('--comments', type=str, default='', required=False, help='comments to be added to trained model metadata, default: %(default)s')
+
+    group_data = parser.add_argument_group('Dataset')
+    group_data.add_argument('--input', type=str, default=None, required=True, help='input folder with training images')
+    group_data.add_argument('--interim', type=str, default='', required=False, help='where to store processed images, default is system temp/train')
+    group_data.add_argument('--process', type=str, default='original,interrogate,resize,square', required=False, help=f'list of possible processing steps: {valid_steps}, default: %(default)s')
+
+    group_train = parser.add_argument_group('Train')
+    group_train.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
+    group_train.add_argument('--steps', type=int, default=2500, required=False, help='training steps, default: %(default)s')
+    group_train.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
+    group_train.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
+    group_train.add_argument('--dim', type=int, default=32, required=False, help='network dimension or number of vectors, default: %(default)s')
+
+    # lora params
+    group_train.add_argument('--repeats', type=int, default=1, required=False, help='number of repeats per image, default: %(default)s')
+    group_train.add_argument('--alpha', type=float, default=0, required=False, help='lora/lyco alpha for weights scaling, default: dim/2')
+    group_train.add_argument('--algo', type=str, default=None, choices=['locon', 'loha', 'lokr', 'ia3'], required=False, help='alternative lyco algoritm, default: %(default)s')
+    group_train.add_argument('--args', type=str, default=None, required=False, help='lora/lyco additional network arguments, default: %(default)s')
+    group_train.add_argument('--optimizer', type=str, default='AdamW', required=False, help='optimizer type, default: %(default)s')
+    group_train.add_argument('--precision', type=str, choices=['fp16', 'fp32'], default='fp16', required=False, help='training precision, default: %(default)s')
+    group_train.add_argument('--sdxl', default = False, action='store_true', help = "run sdxl training, default: %(default)s")
+    # AdamW (default), AdamW8bit, PagedAdamW8bit, Lion8bit, PagedLion8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation(DAdaptAdamPreprint), DAdaptAdaGrad, DAdaptAdam, DAdaptAdan, DAdaptAdanIP, DAdaptLion, DAdaptSGD, AdaFactor
+
+    group_other = parser.add_argument_group('Other')
+    group_other.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing training, default: %(default)s")
+    group_other.add_argument('--experimental', default = False, action='store_true', help = "enable experimental options, default: %(default)s")
+    group_other.add_argument('--debug', default = False, action='store_true', help = "enable debug level logging, default: %(default)s")
+
+    args = parser.parse_args()
+
+
+def prepare_server():
+    global server_ok # pylint: disable=global-statement
+    try:
+        server_status = util.Map(sdapi.progresssync())
+        server_state = server_status['state']
+        server_ok = True
+    except Exception:
+        log.warning(f'sdnext server error: {server_status}')
+        server_ok = False
+    if server_ok and server_state['job_count'] > 0:
+        log.error(f'sdnext server not idle: {server_state}')
+        exit(1)
+    if server_ok:
+        server_options = util.Map(sdapi.options())
+        server_options.options.save_training_settings_to_txt = False
+        server_options.options.training_enable_tensorboard = False
+        server_options.options.training_tensorboard_save_images = False
+        server_options.options.pin_memory = True
+        server_options.options.save_optimizer_state = False
+        server_options.options.training_image_repeats_per_epoch = args.repeats
+        server_options.options.training_write_csv_every = 0
+        sdapi.postsync('/sdapi/v1/options', server_options.options)
+        log.info('updated server options')
+
+
+def verify_args():
+    server_options = util.Map(sdapi.options())
+    if args.model != '':
+        if not os.path.isfile(args.model):
+            log.error(f'cannot find loaded model: {args.model}')
+            exit(1)
+        if server_ok:
+            server_options.options.sd_model_checkpoint = args.model
+            sdapi.postsync('/sdapi/v1/options', server_options.options)
+    elif server_ok:
+        args.model = server_options.options.sd_model_checkpoint.split(' [')[0]
+        if args.sdxl and (server_options.sd_backend != 'diffusers' or server_options.diffusers_pipeline != 'Stable Diffusion XL'):
+            log.warning('server checkpoint is not sdxl')
+    else:
+        log.error('no model specified')
+        exit(1)
+    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    if args.type == 'lora' and not server_ok and not args.dir:
+        log.error('offline lora training requires --dir <lora folder>')
+        exit(1)
+    if args.type == 'lora':
+        import transformers
+        if transformers.__version__ != '4.30.2':
+            log.error(f'lora training requires specific transformers version: current {transformers.__version__} required transformers==4.30.2')
+            exit(1)
+    args.lora_dir = server_options.options.lora_dir or args.dir
+    if not os.path.isabs(args.lora_dir):
+        args.lora_dir = os.path.join(base_dir, args.lora_dir)
+    args.lyco_dir = server_options.options.lyco_dir or args.dir
+    if not os.path.isabs(args.lyco_dir):
+        args.lyco_dir = os.path.join(base_dir, args.lyco_dir)
+    args.embeddings_dir = server_options.options.embeddings_dir or args.dir
+    if not os.path.isfile(args.model):
+        args.ckpt_dir = server_options.options.ckpt_dir
+        if not os.path.isabs(args.ckpt_dir):
+            args.ckpt_dir = os.path.join(base_dir, args.ckpt_dir)
+        attempt = os.path.abspath(os.path.join(args.ckpt_dir, args.model))
+        args.model = attempt if os.path.isfile(attempt) else args.model
+    if not os.path.isfile(args.model):
+        attempt = os.path.abspath(os.path.join(args.ckpt_dir, args.model + '.safetensors'))
+        args.model = attempt if os.path.isfile(attempt) else args.model
+    if not os.path.isfile(args.model):
+        log.error(f'cannot find loaded model: {args.model}')
+        exit(1)
+    if not os.path.exists(args.input) or not os.path.isdir(args.input):
+        log.error(f'cannot find training folder: {args.input}')
+        exit(1)
+    if not os.path.exists(args.lora_dir) or not os.path.isdir(args.lora_dir):
+        log.error(f'cannot find lora folder: {args.lora_dir}')
+        exit(1)
+    if not os.path.exists(args.lyco_dir) or not os.path.isdir(args.lyco_dir):
+        log.error(f'cannot find lyco folder: {args.lyco_dir}')
+        exit(1)
+    if args.interim != '':
+        args.process_dir = args.interim
+    else:
+        args.process_dir = os.path.join(tempfile.gettempdir(), 'train', args.name)
+    log.debug(f'args: {vars(args)}')
+    log.debug(f'server flags: {server_options.flags}')
+    log.debug(f'server options: {server_options.options}')
+
+
+async def training_loop():
+    async def async_train():
+        res = await sdapi.post('/sdapi/v1/train/embedding', options.embedding)
+        log.info(f'train embedding result: {res}')
+
+    async def async_monitor():
+        from tqdm.rich import tqdm
+        await asyncio.sleep(3)
+        res = util.Map(sdapi.progress())
+        with tqdm(desc='train embedding', total=res.state.job_count) as pbar:
+            while res.state.job_no < res.state.job_count and not res.state.interrupted and not res.state.skipped:
+                await asyncio.sleep(2)
+                prev_job = res.state.job_no
+                res = util.Map(sdapi.progress())
+                loss = re.search(r"Loss: (.*?)(?=\<)", res.textinfo)
+                if loss:
+                    pbar.set_postfix({ 'loss': loss.group(0) })
+                    pbar.update(res.state.job_no - prev_job)
+
+    a = asyncio.create_task(async_train())
+    b = asyncio.create_task(async_monitor())
+    await asyncio.gather(a, b) # wait for both pipeline and monitor to finish
+
+
+def train_embedding():
+    log.info(f'{args.type} options: {options.embedding}')
+    create_options = util.Map({
+        "name": args.name,
+        "num_vectors_per_token": args.dim,
+        "overwrite_old": False,
+        "init_text": args.tag,
+    })
+    fn = os.path.join(args.embeddings_dir, args.name) + '.pt'
+    if os.path.exists(fn) and args.overwrite:
+        log.warning(f'delete existing embedding {fn}')
+        os.remove(fn)
+    else:
+        log.error(f'embedding exists {fn}')
+        return
+    log.info(f'create embedding {create_options}')
+    res = sdapi.postsync('/sdapi/v1/create/embedding', create_options)
+    if 'info' in res and 'error' in res['info']: # formatted error
+        log.error(res.info)
+    elif 'info' in res: # no error
+        asyncio.run(training_loop())
+    else: # unknown error
+        log.error(f'create embedding error {res}')
+
+
+def train_lora():
+    fn = os.path.join(options.lora.output_dir, args.name)
+    for ext in ['.ckpt', '.pt', '.safetensors']:
+        if os.path.exists(fn + ext):
+            if args.overwrite:
+                log.warning(f'delete existing lora: {fn + ext}')
+                os.remove(fn + ext)
+            else:
+                log.error(f'lora exists: {fn + ext}')
+                return
+    log.info(f'{args.type} options: {options.lora}')
+    # lora imports
+    lora_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lora'))
+    lycoris_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lycoris'))
+    sys.path.append(lora_path)
+    if args.type == 'lyco':
+        sys.path.append(lycoris_path)
+    log.debug('importing lora lib')
+    if not args.sdxl:
+        import train_network
+        trainer = train_network.NetworkTrainer()
+        trainer.train(options.lora)
+    else:
+        import sdxl_train_network
+        trainer = sdxl_train_network.SdxlNetworkTrainer()
+        trainer.train(options.lora)
+    if args.type == 'lyco':
+        log.debug('importing lycoris lib')
+        import importlib
+        _network_module = importlib.import_module(options.lora.network_module)
+
+
+def prepare_options():
+    if args.type == 'embedding':
+        log.info('train embedding')
+        options.lora.in_json = None
+    if args.type == 'dreambooth':
+        log.info('train using dreambooth style training')
+        options.lora.vae_batch_size = args.batch
+        options.lora.in_json = None
+    if args.type == 'lora':
+        log.info('train using lora style training')
+        options.lora.output_dir = args.lora_dir
+        options.lora.in_json = os.path.join(args.process_dir, args.name + '.json')
+    if args.type == 'lyco':
+        log.info('train using lycoris network')
+        options.lora.output_dir = args.lora_dir
+        options.lora.network_module = 'lycoris.kohya'
+        options.lora.in_json = os.path.join(args.process_dir, args.name + '.json')
+    # lora specific
+    options.lora.save_model_as = 'safetensors'
+    options.lora.pretrained_model_name_or_path = args.model
+    options.lora.output_name = args.name
+    options.lora.max_train_steps = args.steps
+    options.lora.network_dim = args.dim
+    options.lora.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
+    options.lora.network_args = []
+    options.lora.training_comment = args.comments
+    options.lora.sdpa = True
+    options.lora.optimizer_type = args.optimizer
+    if args.algo is not None:
+        options.lora.network_args.append(f'algo={args.algo}')
+    if args.args is not None:
+        for net_arg in args.args:
+            options.lora.network_args.append(net_arg)
+    options.lora.gradient_accumulation_steps = args.gradient
+    options.lora.learning_rate = args.lr
+    options.lora.train_batch_size = args.batch
+    options.lora.train_data_dir = args.process_dir
+    options.lora.no_half_vae = args.precision == 'fp16'
+    # embedding specific
+    options.embedding.embedding_name = args.name
+    options.embedding.learn_rate = str(args.lr)
+    options.embedding.batch_size = args.batch
+    options.embedding.steps = args.steps
+    options.embedding.data_root = args.process_dir
+    options.embedding.log_directory = os.path.join(args.process_dir, 'log')
+    options.embedding.gradient_step = args.gradient
+
+
+def process_inputs():
+    import process
+    import filetype
+    pathlib.Path(args.process_dir).mkdir(parents=True, exist_ok=True)
+    processing_options = args.process.split(',') if isinstance(args.process, str) else args.process
+    processing_options = [opt.strip() for opt in re.split(',| ', args.process)]
+    log.info(f'processing steps: {processing_options}')
+    for step in processing_options:
+        if step not in valid_steps:
+            log.error(f'invalid processing step: {[step]}')
+            exit(1)
+    for root, _sub_dirs, folder in os.walk(args.input):
+        files = [os.path.join(root, f) for f in folder if filetype.is_image(os.path.join(root, f))]
+    log.info(f'processing input images: {len(files)}')
+    if os.path.exists(args.process_dir):
+        if args.overwrite:
+            log.warning(f'removing existing processed folder: {args.process_dir}')
+            shutil.rmtree(args.process_dir, ignore_errors=True)
+        else:
+            log.info(f'processed folder exists: {args.process_dir}')
+    steps = [step for step in processing_options if step in ['face', 'body', 'original']]
+    process.reset()
+    options.process.target_size = 1024 if args.sdxl else 512
+    metadata = {}
+    for step in steps:
+        if step == 'face':
+            opts = [step for step in processing_options if step not in ['body', 'original']]
+        if step == 'body':
+            opts = [step for step in processing_options if step not in ['face', 'original', 'upscale', 'restore']] # body does not perform upscale or restore
+        if step == 'original':
+            opts = [step for step in processing_options if step not in ['face', 'body', 'upscale', 'restore', 'blur', 'range', 'segment']] # original does not perform most steps
+        log.info(f'processing current step: {opts}')
+        tag = step
+        if tag == 'original' and args.tag is not None:
+            concept = args.tag.split(',')[0].strip()
+        else:
+            concept = step
+        if args.type in ['lora', 'lyco', 'dreambooth']:
+            folder = os.path.join(args.process_dir, str(args.repeats) + '_' + concept) # separate concepts per folder
+        if args.type in ['embedding']:
+            folder = os.path.join(args.process_dir) # everything into same folder
+        log.info(f'processing concept: {concept}')
+        log.info(f'processing output folder: {folder}')
+        pathlib.Path(folder).mkdir(parents=True, exist_ok=True)
+        results = {}
+        if server_ok:
+            for f in files:
+                res = process.file(filename = f, folder = folder, tag = args.tag, requested = opts)
+                if res.image: # valid result
+                    results[res.type] = results.get(res.type, 0) + 1
+                    results['total'] = results.get('total', 0) + 1
+                    rel_path = res.basename.replace(os.path.commonpath([res.basename, args.process_dir]), '')
+                    if rel_path.startswith(os.path.sep):
+                        rel_path = rel_path[1:]
+                    metadata[rel_path] = { 'caption': res.caption, 'tags': ','.join(res.tags) }
+                    if options.lora.in_json is None:
+                        with open(res.output.replace(options.process.format, '.txt'), "w", encoding='utf-8') as outfile:
+                            outfile.write(res.caption)
+                log.info(f"processing {'saved' if res.image is not None else 'skipped'}: {f} => {res.output} {res.ops} {res.message}")
+        else:
+            log.info('processing skipped: offline')
+    folders = [os.path.join(args.process_dir, folder) for folder in os.listdir(args.process_dir) if os.path.isdir(os.path.join(args.process_dir, folder))]
+    log.info(f'input datasets {folders}')
+    if options.lora.in_json is not None:
+        with open(options.lora.in_json, "w", encoding='utf-8') as outfile: # write json at the end only
+            outfile.write(json.dumps(metadata, indent=2))
+        for folder in folders: # create latents
+            import latents
+            latents.create_vae_latents(util.Map({ 'input': folder, 'json': options.lora.in_json }))
+            latents.unload_vae()
+    r = { 'inputs': len(files), 'outputs': results, 'metadata': options.lora.in_json }
+    log.info(f'processing steps result: {r}')
+    if args.gradient < 0:
+        log.info(f"setting gradient accumulation to number of images: {results['total']}")
+        options.lora.gradient_accumulation_steps = results['total']
+        options.embedding.gradient_step = results['total']
+    process.unload()
+
+
+if __name__ == '__main__':
+    parse_args()
+    setup_logging()
+    log.info('SD.Next Train')
+    sdapi.sd_url = args.server
+    if args.user is not None:
+        sdapi.sd_username = args.user
+    if args.password is not None:
+        sdapi.sd_password = args.password
+    prepare_server()
+    verify_args()
+    prepare_options()
+    mem_stats()
+    process_inputs()
+    mem_stats()
+    try:
+        if args.type == 'embedding':
+            train_embedding()
+        if args.type == 'lora' or args.type == 'lyco' or args.type == 'dreambooth':
+            train_lora()
+    except KeyboardInterrupt:
+        log.error('interrupt requested')
+        sdapi.interrupt()
+    mem_stats()
+    log.info('done')

cli/util.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python
+"""
+generic helper methods
+"""
+
+import os
+import string
+import logging
+import warnings
+
+log_format = '%(asctime)s %(levelname)s: %(message)s'
+logging.basicConfig(level = logging.INFO, format = log_format)
+warnings.filterwarnings(action="ignore", category=DeprecationWarning)
+warnings.filterwarnings(action="ignore", category=FutureWarning)
+warnings.filterwarnings(action="ignore", category=UserWarning)
+log = logging.getLogger("sd")
+
+
+def set_logfile(logfile):
+    fh = logging.FileHandler(logfile)
+    formatter = logging.Formatter(log_format)
+    fh.setLevel(log.getEffectiveLevel())
+    fh.setFormatter(formatter)
+    log.addHandler(fh)
+    log.info({ 'log file': logfile })
+
+
+def safestring(text: str):
+    lines = []
+    for line in text.splitlines():
+        lines.append(line.translate(str.maketrans('', '', string.punctuation)).strip())
+    res = ', '.join(lines)
+    return res[:1000]
+
+
+def get_memory():
+    def gb(val: float):
+        return round(val / 1024 / 1024 / 1024, 2)
+    mem = {}
+    try:
+        import psutil
+        process = psutil.Process(os.getpid())
+        res = process.memory_info()
+        ram_total = 100 * res.rss / process.memory_percent()
+        ram = { 'free': gb(ram_total - res.rss), 'used': gb(res.rss), 'total': gb(ram_total) }
+        mem.update({ 'ram': ram })
+    except Exception as e:
+        mem.update({ 'ram': e })
+    try:
+        import torch
+        if torch.cuda.is_available():
+            s = torch.cuda.mem_get_info()
+            gpu = { 'free': gb(s[0]), 'used': gb(s[1] - s[0]), 'total': gb(s[1]) }
+            s = dict(torch.cuda.memory_stats('cuda'))
+            allocated = { 'current': gb(s['allocated_bytes.all.current']), 'peak': gb(s['allocated_bytes.all.peak']) }
+            reserved = { 'current': gb(s['reserved_bytes.all.current']), 'peak': gb(s['reserved_bytes.all.peak']) }
+            active = { 'current': gb(s['active_bytes.all.current']), 'peak': gb(s['active_bytes.all.peak']) }
+            inactive = { 'current': gb(s['inactive_split_bytes.all.current']), 'peak': gb(s['inactive_split_bytes.all.peak']) }
+            events = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
+            mem.update({
+                'gpu': gpu,
+                'gpu-active': active,
+                'gpu-allocated': allocated,
+                'gpu-reserved': reserved,
+                'gpu-inactive': inactive,
+                'events': events,
+            })
+    except Exception:
+        pass
+    return Map(mem)
+
+
+class Map(dict): # pylint: disable=C0205
+    __slots__ = ('__dict__') # pylint: disable=superfluous-parens
+    def __init__(self, *args, **kwargs):
+        super(Map, self).__init__(*args, **kwargs) # pylint: disable=super-with-arguments
+        for arg in args:
+            if isinstance(arg, dict):
+                for k, v in arg.items():
+                    if isinstance(v, dict):
+                        v = Map(v)
+                    if isinstance(v, list):
+                        self.__convert(v)
+                    self[k] = v
+        if kwargs:
+            for k, v in kwargs.items():
+                if isinstance(v, dict):
+                    v = Map(v)
+                elif isinstance(v, list):
+                    self.__convert(v)
+                self[k] = v
+    def __convert(self, v):
+        for elem in range(0, len(v)): # pylint: disable=consider-using-enumerate
+            if isinstance(v[elem], dict):
+                v[elem] = Map(v[elem])
+            elif isinstance(v[elem], list):
+                self.__convert(v[elem])
+    def __getattr__(self, attr):
+        return self.get(attr)
+    def __setattr__(self, key, value):
+        self.__setitem__(key, value)
+    def __setitem__(self, key, value):
+        super(Map, self).__setitem__(key, value) # pylint: disable=super-with-arguments
+        self.__dict__.update({key: value})
+    def __delattr__(self, item):
+        self.__delitem__(item)
+    def __delitem__(self, key):
+        super(Map, self).__delitem__(key) # pylint: disable=super-with-arguments
+        del self.__dict__[key]
+
+
+if __name__ == "__main__":
+    pass

cli/validate-locale.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+
+import os
+import sys
+import json
+from rich import print # pylint: disable=redefined-builtin
+
+if __name__ == "__main__":
+    sys.argv.pop(0)
+    fn = sys.argv[0] if len(sys.argv) > 0 else 'locale_en.json'
+    if not os.path.isfile(fn):
+        print(f'File not found: {fn}')
+        sys.exit(1)
+    with open(fn, 'r', encoding="utf-8") as f:
+        data = json.load(f)
+    keys = []
+    t_names = 0
+    t_hints = 0
+    t_localized = 0
+    t_long = 0
+    for k in data.keys():
+        names = len(data[k])
+        t_names += names
+        hints = len([k for k in data[k] if k["hint"] != ""])
+        t_hints += hints
+        localized = len([k for k in data[k] if k["localized"] != ""])
+        t_localized += localized
+        missing = names - hints
+        long = 0
+        for v in data[k]:
+            if v['label'] in keys:
+                print(f'  Duplicate: {k}.{v["label"]}')
+            else:
+                if len(v['label']) > 63:
+                    long += 1
+                    print(f'  Long label: {k}.{v["label"]}')
+                keys.append(v['label'])
+        t_long += long
+        print(f'Section: [bold magenta]{k.ljust(20)}[/bold magenta] entries={names} localized={"[bold green]" + str(localized) + "[/bold green]" if localized > 0 else "0"} long={"[bold red]" + str(long) + "[/bold red]" if long > 0 else "0"} hints={hints} missing={"[bold red]" + str(missing) + "[/bold red]" if missing > 0 else "[bold green]0[/bold green]"}')
+    print(f'Totals: entries={t_names} localized={localized} long={t_long} hints={t_hints} missing={t_names - t_hints}')

cli/video-extract.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+"""
+use ffmpeg for animation processing
+"""
+import os
+import json
+import subprocess
+import pathlib
+import argparse
+import filetype
+from util import log, Map
+
+
+def probe(src: str):
+    cmd = f"ffprobe -hide_banner -loglevel 0 -print_format json -show_format -show_streams \"{src}\""
+    result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
+    data = json.loads(result.stdout)
+    stream = [x for x in data['streams'] if x["codec_type"] == "video"][0]
+    fmt = data['format'] if 'format' in data else {}
+    res = {**stream, **fmt}
+    video = Map({
+        'codec': res.get('codec_name', 'unknown') + '/' + res.get('codec_tag_string', ''),
+        'resolution': [int(res.get('width', 0)), int(res.get('height', 0))],
+        'duration': float(res.get('duration', 0)),
+        'frames': int(res.get('nb_frames', 0)),
+        'bitrate': round(float(res.get('bit_rate', 0)) / 1024),
+    })
+    return video
+
+
+def extract(src: str, dst: str, rate: float = 0.015, fps: float = 0, start = 0, end = 0):
+    images = []
+    if not os.path.isfile(src) or not filetype.is_video(src):
+        log.error({ 'extract': 'input is not movie file' })
+        return 0
+    dst = dst if dst.endswith('/') else dst + '/'
+
+    video = probe(src)
+    log.info({ 'extract': { 'source': src, **video } })
+
+    ssstart = f' -ss {start}' if start > 0 else ''
+    ssend = f' -to {video.duration - end}' if start > 0 else ''
+    filename = pathlib.Path(src).stem
+    if rate > 0:
+        cmd = f"ffmpeg -hide_banner -y -loglevel info {ssstart} {ssend} -i \"{src}\" -filter:v \"select='gt(scene,{rate})',metadata=print\" -vsync vfr -frame_pts 1 \"{dst}{filename}-%05d.jpg\""
+    elif fps > 0:
+        cmd = f"ffmpeg -hide_banner -y -loglevel info {ssstart} {ssend} -i \"{src}\" -r {fps} -vsync vfr -frame_pts 1 \"{dst}{filename}-%05d.jpg\""
+    else:
+        log.error({ 'extract': 'requires either rate or fps' })
+        return 0
+    log.debug({ 'extract': cmd })
+    pathlib.Path(dst).mkdir(parents = True, exist_ok = True)
+    result = subprocess.run(cmd, shell = True, capture_output = True, text = True, check = True)
+    for line in result.stderr.split('\n'):
+        if 'pts_time' in line:
+            log.debug({ 'extract': { 'keyframe': line.strip().split(' ')[-1].split(':')[-1] } })
+    images = next(os.walk(dst))[2]
+    log.info({ 'extract': { 'destination': dst, 'keyframes': len(images), 'rate': rate, 'fps': fps } })
+    return len(images)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="ffmpeg pipeline")
+    parser.add_argument("--input", type = str, required = True, help="input")
+    parser.add_argument("--output", type = str, required = True, help="output")
+    parser.add_argument("--rate", type = float, default = 0, required = False, help="extraction change rate threshold")
+    parser.add_argument("--fps", type = float, default = 0, required = False, help="extraction frames per second")
+    parser.add_argument("--skipstart", type = float, default = 1, required = False, help="skip time from start of video")
+    parser.add_argument("--skipend", type = float, default = 1, required = False, help="skip time to end of video")
+    params = parser.parse_args()
+    extract(src = params.input, dst = params.output, rate = params.rate, fps = params.fps, start = params.skipstart, end = params.skipend)

configs/alt-diffusion-inference.yaml

@@ -0,0 +1,72 @@
+model:
+  base_learning_rate: 1.0e-04
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 10000 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: modules.xlmr.BertSeriesModelWithTransformation
+      params:
+        name: "XLMR-Large"

configs/instruct-pix2pix.yaml

@@ -0,0 +1,98 @@
+# File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion).
+# See more details in LICENSE.
+
+model:
+  base_learning_rate: 1.0e-04
+  target: modules.hijack.ddpm_edit.LatentDiffusion
+  params:
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: edited
+    cond_stage_key: edit
+    # image_size: 64
+    # image_size: 32
+    image_size: 16
+    channels: 4
+    cond_stage_trainable: false   # Note: different from the one we trained before
+    conditioning_key: hybrid
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: false
+
+    scheduler_config: # 10000 warmup steps
+      target: ldm.lr_scheduler.LambdaLinearScheduler
+      params:
+        warm_up_steps: [ 0 ]
+        cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases
+        f_start: [ 1.e-6 ]
+        f_max: [ 1. ]
+        f_min: [ 1. ]
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        image_size: 32 # unused
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_heads: 8
+        use_spatial_transformer: True
+        transformer_depth: 1
+        context_dim: 768
+        use_checkpoint: True
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenCLIPEmbedder
+
+data:
+  target: main.DataModuleFromConfig
+  params:
+    batch_size: 128
+    num_workers: 1
+    wrap: false
+    validation:
+      target: edit_dataset.EditDataset
+      params:
+        path: data/clip-filtered-dataset
+        cache_dir:  data/
+        cache_name: data_10k
+        split: val
+        min_text_sim: 0.2
+        min_image_sim: 0.75
+        min_direction_sim: 0.2
+        max_samples_per_prompt: 1
+        min_resize_res: 512
+        max_resize_res: 512
+        crop_res: 512
+        output_as_edit: False
+        real_input: True