Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/master' into settings-improvements

Browse Source
This commit is contained in:
Alex Heller
2023-06-05 02:52:29 +02:00
parent d182699f82 9726b4d23c
commit 1b90460a3f
237 changed files with 11268 additions and 11707 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.eslintignore
+4
View File
@@ -0,0 +1,4 @@
extensions
extensions-disabled
repositories
venv
.eslintrc.json
+89
View File
@@ -0,0 +1,89 @@
{
"env": {
"browser": true,
"commonjs": false,
"node": false,
"jquery": false,
"es2020": true
},
"parserOptions": { "ecmaVersion": 2020 },
"plugins": [],
"extends": ["eslint:recommended", "airbnb-base"],
"rules": {
"arrow-spacing": "error",
"block-spacing": "error",
"brace-style": "error",
"comma-dangle": ["error", "only-multiline"],
"comma-spacing": "error",
"comma-style": ["error", "last"],
"curly": ["error", "multi-line", "consistent"],
"eol-last": "error",
"func-call-spacing": "error",
"function-call-argument-newline": ["error", "consistent"],
"function-paren-newline": ["error", "consistent"],
"indent": ["error", 4],
"key-spacing": "error",
"keyword-spacing": "error",
"linebreak-style": ["error", "unix"],
"no-extra-semi": "error",
"no-mixed-spaces-and-tabs": "error",
"no-multi-spaces": "error",
"no-redeclare": ["error", { "builtinGlobals": false }],
"no-trailing-spaces": "error",
"no-unused-vars": "off",
"no-whitespace-before-property": "error",
"object-curly-newline": ["error", { "consistent": true, "multiline": true }],
"object-curly-spacing": ["error", "never"],
"operator-linebreak": ["error", "after"],
"quote-props": ["error", "consistent-as-needed"],
"semi": ["error", "always"],
"semi-spacing": "error",
"semi-style": ["error", "last"],
"space-before-blocks": "error",
"space-before-function-paren": ["error", "never"],
"space-in-parens": ["error", "never"],
"space-infix-ops": "error",
"space-unary-ops": "error",
"switch-colon-spacing": "error",
"template-curly-spacing": ["error", "never"],
"unicode-bom": "error"
},
"globals": {
//script.js
"gradioApp": "readonly",
"onUiLoaded": "readonly",
"onUiUpdate": "readonly",
"onOptionsChanged": "readonly",
"uiCurrentTab": "writable",
"uiElementIsVisible": "readonly",
"uiElementInSight": "readonly",
"executeCallbacks": "readonly",
//ui.js
"opts": "writable",
"all_gallery_buttons": "readonly",
"selected_gallery_button": "readonly",
"selected_gallery_index": "readonly",
"switch_to_txt2img": "readonly",
"switch_to_img2img_tab": "readonly",
"switch_to_img2img": "readonly",
"switch_to_sketch": "readonly",
"switch_to_inpaint": "readonly",
"witch_to_inpaint_sketch": "readonly",
"switch_to_extras": "readonly",
"get_tab_index": "readonly",
"create_submit_args": "readonly",
"restart_reload": "readonly",
"updateInput": "readonly",
//extraNetworks.js
"requestGet": "readonly",
"popup": "readonly",
// from python
"localization": "readonly",
// progrssbar.js
"randomId": "readonly",
"requestProgress": "readonly",
// imageviewer.js
"modalPrevImage": "readonly",
"modalNextImage": "readonly"
}
}
.github/ISSUE_TEMPLATE/bug_report.yml
+5 -8
View File
@@ -1,7 +1,7 @@
name: Issue Report
description: Something is broken
title: "[Issue]: "
labels: ["issue"]
labels: []
body:
- type: textarea
@@ -23,22 +23,19 @@ body:
Look at console log and copy the version string from there
For example: `Version: f256fb8b Fri Apr 14 17:41:30 2023 -0400`
Additionally provide any relevant platorm information (OS, browser, versions)
Additionally provide any relevant platorm information:
- OS details, GPU information, browser used, app versions
- type: markdown
attributes:
value: |
If issue is setup, installation or startup related, please check `setup.log` before reporting
If issue is setup, installation or startup related, please check `webui.log` before reporting
And when posting console logs, please use code blocks ( \`\`\` ) to format them insead of uploading screenshots
- type: markdown
attributes:
value: |
Verify that issue persists after restarting server from command line
Restarting from UI or reloading browser is insufficient to create reproducible environment
- type: markdown
attributes:
value: |
If you have additional extensions installed, try to reproduce the issue with user extensions disabled
And if the issue is with compatibility with specific extension, mark it as such when creating the issue
Try running with `--safe` command line flag with disables loading of user-installed extensions
- type: markdown
attributes:
value: |
.github/ISSUE_TEMPLATE/config.yml
+1 -1
View File
@@ -1,5 +1,5 @@
blank_issues_enabled: false
contact_links:
- name: WebUI Community Support
- name: SD.Next Community Support
url: https://github.com/vladmandic/automatic/discussions
about: Please ask and answer questions here.
.github/ISSUE_TEMPLATE/extension_report.yml
+33
View File
@@ -0,0 +1,33 @@
name: Extension Compatibility Report
description: Extension is not working as expected
title: "[Extension]: "
labels: ["extension"]
body:
- type: textarea
id: description
attributes:
label: Issue Description
description: Tell us what happened in a very clear and simple way
value: Please fill this form with as much information as possible
- type: textarea
id: platform
attributes:
label: Version Platform Description
description: Describe your platform (program version, OS, browser)
value:
- type: textarea
id: url
attributes:
label: URL link of the extension
description: URL link of the extension
value:
- type: markdown
attributes:
value: |
Any issues without version information and actual link to extension will be closed
- type: markdown
attributes:
value: |
If issue is extension installation or startup related, please check `webui.log` before reporting
And when posting console logs, please use code blocks ( \`\`\` ) to format them insead of uploading screenshots
.gitignore
+10 -4
View File
@@ -1,18 +1,22 @@
# defaults
__pycache__
/params.txt
/cache.json
/config.json
/ui-config.json
/setup.log
/params.txt
/styles.csv
/ui-config.json
/user.css
/webui-user.bat
/webui-user.sh
/javascript/themes.json
/html/extensions.json
/html/themes.json
node_modules
pnpm-lock.yaml
package-lock.json
venv
# all models and temp files
*.log
*.bak
*.ckpt
*.safetensors
@@ -27,6 +31,8 @@ venv
/*.bat
/*.sh
/*.txt
/*.mp3
/*.lnk
!webui.bat
!webui.sh
.gitmodules
+9 -14
View File
@@ -4,24 +4,10 @@
[submodule "extensions-builtin/sd-extension-system-info"]
path = extensions-builtin/sd-extension-system-info
url = https://github.com/vladmandic/sd-extension-system-info
[submodule "extensions-builtin/sd-extension-aesthetic-scorer"]
path = extensions-builtin/sd-extension-aesthetic-scorer
url = https://github.com/vladmandic/sd-extension-aesthetic-scorer
[submodule "extensions-builtin/sd-extension-steps-animation"]
path = extensions-builtin/sd-extension-steps-animation
url = https://github.com/vladmandic/sd-extension-steps-animation
[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/seed_travel"]
path = extensions-builtin/seed_travel
url = https://github.com/yownas/seed_travel
ignore = dirty
[submodule "extensions-builtin/sd-webui-model-converter"]
path = extensions-builtin/sd-webui-model-converter
url = https://github.com/Akegarasu/sd-webui-model-converter
ignore = dirty
[submodule "extensions-builtin/sd-dynamic-thresholding"]
path = extensions-builtin/sd-dynamic-thresholding
url = https://github.com/mcmonkeyprojects/sd-dynamic-thresholding
@@ -33,18 +19,27 @@
[submodule "extensions-builtin/clip-interrogator-ext"]
path = extensions-builtin/clip-interrogator-ext
url = https://github.com/pharmapsychotic/clip-interrogator-ext.git
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 "modules/lycoris"]
path = modules/lycoris
url = https://github.com/KohakuBlueleaf/LyCORIS
ignore = dirty
[submodule "extensions-builtin/stable-diffusion-webui-rembg"]
path = extensions-builtin/stable-diffusion-webui-rembg
url = https://github.com/AUTOMATIC1111/stable-diffusion-webui-rembg
ignore = dirty
[submodule "extensions-builtin/a1111-sd-webui-lycoris"]
path = extensions-builtin/a1111-sd-webui-lycoris
url = https://github.com/KohakuBlueleaf/a1111-sd-webui-lycoris
ignore = dirty
[submodule "extensions-builtin/multidiffusion-upscaler-for-automatic1111"]
path = extensions-builtin/multidiffusion-upscaler-for-automatic1111
url = https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111
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
.pylintrc
+28 -25
View File
@@ -75,16 +75,16 @@ valid-metaclass-classmethod-first-arg=mcs
[DESIGN]
exclude-too-few-public-methods=
ignored-parents=
max-args=5
max-attributes=7
max-bool-expr=5
max-branches=12
max-locals=15
max-parents=7
max-args=99
max-attributes=99
max-bool-expr=99
max-branches=99
max-locals=99
max-parents=99
max-public-methods=20
max-returns=6
max-statements=50
min-public-methods=2
max-returns=99
max-statements=99
min-public-methods=1
[EXCEPTIONS]
overgeneral-exceptions=builtins.BaseException,builtins.Exception
@@ -124,26 +124,29 @@ confidence=HIGH,
# disable=C,R,W
disable=raw-checker-failed,
bad-inline-option,
locally-disabled,
file-ignored,
suppressed-message,
useless-suppression,
deprecated-pragma,
use-symbolic-message-instead,
line-too-long,
missing-function-docstring,
missing-module-docstring,
missing-class-docstring,
logging-fstring-interpolation,
import-outside-toplevel,
bare-except,
broad-exception-caught,
consider-iterating-dictionary,
wrong-import-position,
unnecessary-lambda,
consider-using-dict-items,
dangerous-default-value,
unnecessary-dunder-call,
deprecated-pragma,
file-ignored,
import-outside-toplevel,
invalid-name,
R0801,
line-too-long,
locally-disabled,
logging-fstring-interpolation,
missing-class-docstring,
missing-function-docstring,
missing-module-docstring,
duplicate-code,
suppressed-message,
unnecessary-dunder-call,
unnecessary-lambda,
use-symbolic-message-instead,
too-many-nested-blocks,
useless-suppression,
wrong-import-position,
enable=c-extension-no-member
[METHOD_ARGS]
CHANGELOG.md
+395
View File
@@ -0,0 +1,395 @@
# Change Log for SD.Next
## Update for 06/03/2023
- new vae decode method to help with larger batch sizes, thanks @bigdog
- profiling of scripts/extensions callbacks
- additional exception handling so bad exception does not crash main app
- additional background removal models
## Update for 06/02/2023
Some quality-of-life improvements while working on larger stuff in the background...
- redesign action box to be uniform accross all themes
- add **pause** option next to stop/skip
- redesigned progress bar
- add new built-in extension: **agent-scheduler**
very elegant way to getting full queing capabilies, thank @artventurdev
- enable more image formats
note: not all are understood by browser so previews and images may appear as blank
unless you have some browser extensions that can handle them
but they are saved correctly. and cant beat raw quality of 32-bit `tiff` or `psd` :)
- change in behavior: `xformers` will be uninstalled on startup if they are not active
if you do have `xformers` selected as your desired cross-optimization method, then they will be used
reason is that a lot of libaries try to blindy import xformers even if they are not selected or not functional
## Update for 05/30/2023
Another bigger one...And more to come in the next few days...
- new live preview mode: taesd
i really like this one, so its enabled as default for new installs
- settings search feature
- new sampler: sde++ 2m sde
- fully common save/zip/delete (new) options in all tabs
which (again) meant rework of process image tab
- system info tab: live gpu utilization/memory graphs for nvidia gpus
- updated controlnet interface
- minor style changes
- updated lora, swinir, scunet and ldsr code from upstream
- start of merge from a1111 v1.3
## Update for 05/26/2023
Some quality-of-life improvements...
- updated [README](https://github.com/vladmandic/automatic/blob/master/README.md)
- created [CHANGELOG](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md)
this will be the source for all info about new things moving forward
and cross-posted to [Discussions#99](https://github.com/vladmandic/automatic/discussions/99) as well as discord [announcements](https://discord.com/channels/1101998836328697867/1109953953396957286)
- optimize model loading on startup
this should reduce startup time significantly
- set default cross-optimization method for each platform backend
applicable for new installs only
- `cuda` => Scaled-Dot-Product
- `rocm` => Sub-quadratic
- `directml` => Sub-quadratic
- `ipex` => InvokeAI's
- `mps` => Doggettx's
- `cpu` => Doggettx's
- optimize logging
- optimize profiling
now includes startup profiling as well as `cuda` profiling during generate
- minor lightbox improvements
- bugfixes...i don't recall when was a release with at least several of those
other than that - first stage of [Diffusers](https://github.com/huggingface/diffusers) integration is now in master branch
i don't recommend anyone to try it (and dont even think reporting issues for it)
but if anyone wants to contribute, take a look at [project page](https://github.com/users/vladmandic/projects/1/views/1)
## Update for 05/23/2023
Major internal work with perhaps not that much user-facing to show for it ;)
- update core repos: **stability-ai**, **taming-transformers**, **k-diffusion, blip**, **codeformer**
note: to avoid disruptions, this is applicable for new installs only
- tested with **torch 2.1**, **cuda 12.1**, **cudnn 8.9**
(production remains on torch2.0.1+cuda11.8+cudnn8.8)
- fully extend support of `--data-dir`
allows multiple installations to share pretty much everything, not just models
especially useful if you want to run in a stateless container or cloud instance
- redo api authentication
now api authentication will use same user/pwd (if specified) for ui and strictly enforce it using httpbasicauth
new authentication is also fully supported in combination with ssl for both sync and async calls
if you want to use api programatically, see examples in `cli/sdapi.py`
- add dark/light theme mode toggle
- redo some `clip-skip` functionality
- better matching for vae vs model
- update to `xyz grid` to allow creation of large number of images without creating grid itself
- update `gradio` (again)
- more prompt parser optimizations
- better error handling when importing image settings which are not compatible with current install
for example, when upscaler or sampler originally used is not available
- fixes...amazing how many issues were introduced by porting a1111 v1.20 code without adding almost no new functionality
next one is v1.30 (still in dev) which does bring a lot of new features
## Update for 05/17/2023
This is a massive one due to huge number of changes,
but hopefully it will go ok...
- new **prompt parsers**
select in UI -> Settings -> Stable Diffusion
- **Full**: my new implementation
- **A1111**: for backward compatibility
- **Compel**: as used in ComfyUI and InvokeAI (a.k.a *Temporal Weighting*)
- **Fixed**: for really old backward compatibility
- monitor **extensions** install/startup and
log if they modify any packages/requirements
this is a *deep-experimental* python hack, but i think its worth it as extensions modifying requirements
is one of most common causes of issues
- added `--safe` command line flag mode which skips loading user extensions
please try to use it before opening new issue
- reintroduce `--api-only` mode to start server without ui
- port *all* upstream changes from [A1111](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
up to today - commit hash `89f9faa`
## Update for 05/15/2023
- major work on **prompt parsing**
this can cause some differences in results compared to what you're used to, but its all about fixes & improvements
- prompt parser was adding commas and spaces as separate words and tokens and/or prefixes
- negative prompt weight using `[word:weight]` was ignored, it was always `0.909`
- bracket matching was anything but correct. complex nested attention brackets are now working.
- btw, if you run with `--debug` flag, you'll now actually see parsed prompt & schedule
- updated all scripts in `/cli`
- add option in settings to force different **latent sampler** instead of using primary only
- add **interrupt/skip** capabilities to process images
## Update for 05/13/2023
This is mostly about optimizations...
- improved `torch-directml` support
especially interesting for **amd** users on **windows** where **torch+rocm** is not yet available
dont forget to run using `--use-directml` or default is **cpu**
- improved compatibility with **nvidia** rtx 1xxx/2xxx series gpus
- fully working `torch.compile` with **torch 2.0.1**
using `inductor` compile takes a while on first run, but does result in 5-10% performance increase
- improved memory handling
for highest performance, you can also disable aggressive **gc** in settings
- improved performance
especially *after* generate as image handling has been moved to separate thread
- allow per-extension updates in extension manager
- option to reset configuration in settings
## Update for 05/11/2023
- brand new **extension manager**
this is pretty much a complete rewrite, so new issues are possible
- support for `torch` 2.0.1
note that if you are experiencing frequent hangs, this may be a worth a try
- updated `gradio` to 3.29.0
- added `--reinstall` flag to force reinstall of all packages
- auto-recover & re-attempt when `--upgrade` is requested but fails
- check for duplicate extensions
## Update for 05/08/2023
Back online with few updates:
- bugfixes. yup, quite a lot of those
- auto-detect some cpu/gpu capabilities on startup
this should reduce need to tweak and tune settings like no-half, no-half-vae, fp16 vs fp32, etc
- configurable order of top level tabs
- configurable order of scripts in txt2img and img2img
for both, see sections in ui-> settings -> user interface
## Update for 05/04/2023
Again, few days later...
- reviewed/ported **all** commits from **A1111** upstream
some a few are not applicable as i already have alternative implementations
and very few i choose not to implement (save/restore last-known-good-config is a bad hack)
otherwise, we're fully up to date (its doesn't show on fork status as code merges were mostly manual due to conflicts)
but...due to sheer size of the updates, this may introduce some temporary issues
- redesigned server restart function
now available and working in ui
actually, since server restart is now a true restart and not ui restart, it can be used much more flexibly
- faster model load
plus support for slower devices via stream-load function (in ui settings)
- better logging
this includes new `--debug` flag for more verbose logging when troubleshooting
## Update for 05/01/2023
Been a bit quieter for last few days as changes were quite significant, but finally here we are...
- Updated core libraries: Gradio, Diffusers, Transformers
- Added support for **Intel ARC** GPUs via Intel OneAPI IPEX (auto-detected)
- Added support for **TorchML** (set by default when running on non-compatible GPU or on CPU)
- Enhanced support for AMD GPUs with **ROCm**
- Enhanced support for Apple **M1/M2**
- Redesigned command params: run `webui --help` for details
- Redesigned API and script processing
- Experimental support for multiple **Torch compile** options
- Improved sampler support
- Google Colab: <https://colab.research.google.com/drive/126cDNwHfifCyUpCCQF9IHpEdiXRfHrLN>
Maintained by <https://github.com/Linaqruf/sd-notebook-collection>
- Fixes, fixes, fixes...
To take advantage of new out-of-the-box tunings, its recommended to delete your `config.json` so new defaults are applied. Its not necessary, but otherwise you may need to play with UI Settings to get the best of Intel ARC, TorchML, ROCm or Apple M1/M2.
## Update for 04/27/2023
a bit shorter list as:
- i've been busy with buxfixing
there are a lot of them, not going to list each here.
but seems like critical issues backlog is quieting down and soon i can focus on new features development.
- i've started collaboration with couple of major projects,
hopefully this will accelerate future development.
what's new:
- ability to view/add/edit model description shown in extra networks cards
- add option to specify fallback sampler if primary sampler is not compatible with desired operation
- make clip skip a local parameter
- remove obsolete items from UI settings
- set defaults for AMD ROCm
if you have issues, you may want to start with a fresh install so configuration can be created from scratch
- set defaults for Apple M1/M2
if you have issues, you may want to start with a fresh install so configuration can be created from scratch
## Update for 04/25/2023
- update process image -> info
- add VAE info to metadata
- update GPU utility search paths for better GPU type detection
- update git flags for wider compatibility
- update environment tuning
- update ti training defaults
- update VAE search paths
- add compatibility opts for some old extensions
- validate script args for always-on scripts
fixes: deforum with controlnet
## Update for 04/24/2023
- identify race condition where generate locks up while fetching preview
- add pulldowns to x/y/z script
- add VAE rollback feature in case of NaNs
- use samples format for live preview
- add token merging
- use **Approx NN** for live preview
- create default `styles.csv`
- fix setup not installing `tensorflow` dependencies
- update default git flags to reduce number of warnings
## Update for 04/23/2023
- fix VAE dtype
should fix most issues with NaN or black images
- add built-in Gradio themes
- reduce requirements
- more AMD specific work
- initial work on Apple platform support
- additional PR merges
- handle torch cuda crashing in setup
- fix setup race conditions
- fix ui lightbox
- mark tensorflow as optional
- add additional image name templates
## Update for 04/22/2023
- autodetect which system libs should be installed
this is a first pass of autoconfig for **nVidia** vs **AMD** environments
- fix parse cmd line args from extensions
- only install `xformers` if actually selected as desired cross-attention method
- do not attempt to use `xformers` or `sdp` if running on cpu
- merge tomesd token merging
- merge 23 PRs pending from a1111 backlog (!!)
*expect shorter updates for the next few days as i'll be partically ooo*
## Update for 04/20/2023
- full CUDA tuning section in UI Settings
- improve exif/pnginfo metadata parsing
it can now handle 3rd party images or images edited in external software
- optimized setup performance and logging
- improve compatibility with some 3rd party extensions
for example handle extensions that install packages directly from github urls
- fix initial model download if no models found
- fix vae not found issues
- fix multiple git issues
note: if you previously had command line optimizations such as --no-half, those are now ignored and moved to ui settings
## Update for 04/19/2023
- fix live preview
- fix model merge
- fix handling of user-defined temp folders
- fix submit benchmark
- option to override `torch` and `xformers` installer
- separate benchmark data for system-info extension
- minor css fixes
- created initial merge backlog from pending prs on a1111 repo
see #258 for details
## Update for 04/18/2023
- reconnect ui to active session on browser restart
this is one of most frequently asked for items, finally figured it out
works for text and image generation, but not for process as there is no progress bar reported there to start with
- force unload `xformers` when not used
improves compatibility with AMD/M1 platforms
- add `styles.csv` to UI settings to allow customizing path
- add `--skip-git` to cmd flags for power users that want
to skip all git checks and operations and perform manual updates
- add `--disable-queue` to cmd flags that disables Gradio queues (experimental)
this forces it to use HTTP instead of WebSockets and can help on unreliable network connections
- set scripts & extensions loading priority and allow custom priorities
fixes random extension issues:
`ScuNet` upscaler dissapearing, `Additional Networks` not showing up on XYZ axis, etc.
- improve html loading order
- remove some `asserts` causing runtime errors and replace with user-friendly messages
- update README.md
- update TODO.md
## Update for 04/17/2023
- **themes** are now dynamic and discovered from list of available gradio themes on huggingface
its quite a list of 30+ supported themes so far
- added option to see **theme preview** without the need to apply it or restart server
- integrated **image info** functionality into **process image** tab and removed separate **image info** tab
- more installer improvements
- fix urls
- updated github integration
- make model download as optional if no models found
## Update for 04/16/2023
- support for ui themes! to to *settings* -> *user interface* -> "ui theme*
includes 12 predefined themes
- ability to restart server from ui
- updated requirements
- removed `styles.csv` from repo, its now fully under user control
- removed model-keyword extension as overly aggresive
- rewrite of the fastapi middleware handlers
- install bugfixes, hopefully new installer is now ok \
i really want to focus on features and not troubleshooting installer
## Update for 04/15/2023
- update default values
- remove `ui-config.json` from repo, its not fully under user control
- updated extensions mangager
- updated locon/lycoris plugin
- enable quick launch by default
- add multidiffusion upscaler extensions
- add model keyword extension
- enable strong linting
- fix circular imports
- fix extensions updated
- fix git update issues
- update github templates
## Update for 04/14/2023
- handle duplicate extensions
- redo exception handler
- fix generate forever
- enable cmdflags compatibility
- change default css font
- fix ti previews on initial start
- enhance tracebacks
- pin transformers version to last known good version
- fix extension loader
## Update for 04/12/2023
This has been pending for a while, but finally uploaded some massive changes
- New launcher
- `webui.bat` and `webui.sh`:
Platform specific wrapper scripts that starts `launch.py` in Python virtual environment
*Note*: Server can run without virtual environment, but it is recommended to use it
This is carry-over from original repo
**If you're unsure which launcher to use, this is the one you want**
- `launch.py`:
Main startup script
Can be used directly to start server in manually activated `venv` or to run it without `venv`
- `installer.py`:
Main installer, used by `launch.py`
- `webui.py`:
Main server script
- New logger
- New exception handler
- Built-in performance profiler
- New requirements handling
- Move of most of command line flags into UI Settings
README.md
+76 -138
View File
@@ -3,163 +3,101 @@
![License](https://img.shields.io/github/license/vladmandic/human?style=flat-square&svg=true)
![GitHub Status Checks](https://img.shields.io/github/checks-status/vladmandic/human/main?style=flat-square&svg=true)
# Stable Diffusion - Automatic
*Heavily opinionated custom fork of* <https://github.com/AUTOMATIC1111/stable-diffusion-webui>
# SD.Next
Fork is as close as up-to-date with origin as time allows
All code changes are merged upstream whenever possible
**Stable Diffusion implementation with modern UI and advanced features**
The idea behind the fork is to enable latest technologies and advances in text-to-image generation
*Sometimes this is not the same as "as simple as possible to use"*
If you are looking an amazing simple-to-use Stable Diffusion tool, I'd suggest [InvokeAI](https://invoke-ai.github.io/InvokeAI/) specifically due to its automated installer and ease of use
This project started as a form from [Automatic1111 WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui/) and it grew siginificantly since then, but although it diverged significanly, any substantial features to original work is ported to this repository as well
<br>
Individual features are not listed here, instead check [Changelog](CHANGELOG.md) for full list of changes
### Follow [Development updates](https://github.com/vladmandic/automatic/discussions/99) for daily updates on new features/fixes
## Platform support
<br>
![screenshot](javascript/black-orange.jpg)
<br>
## Notes
### Fork does differ in few things
- New installer
- Advanced CUDA tuning
Available in UI Settings
- Advanced environment tuning
- Optimized startup and models lazy-loading
- Built-in performance profiler
- Updated libraries to latest known compatible versions
- Includes opinionated **System** and **Options** configuration
- Does not rely on `Accelerate` as it only affects distributed systems
Gradio web server will be initialized much earlier which model load is done in the background
Faster model loading plus ability to fallback on corrupt models
- Uses simplified folder structure
e.g. `/train`, `/outputs/*`, `/models/*`, etc.
- Enhanced training templates
- Built-in `LoRA`, `LyCORIS`, `Custom Diffusion`, `Dreambooth` training
- Majority of settings configurable via UI without the need for command line flags
e.g, cross-optimization methods, system folders, etc.
- New logger
- New error and exception handlers
### Optimizations
- Optimized for `Torch` 2.0
- Runs with `SDP` memory attention enabled by default if supported by system
*Note*: `xFormers` and other cross-optimization methods are still available
- Auto-adjust parameters when running on **CPU** or **CUDA**
*Note:* AMD and M1 platforms are supported, but without out-of-the-box optimizations
### Integrated Extensions
Hand-picked list of extensions that are deeply integrated into core workflows:
- [System Info](https://github.com/vladmandic/sd-extension-system-info)
- [ControlNet](https://github.com/Mikubill/sd-webui-controlnet)
- [Image Browser](https://github.com/AlUlkesh/stable-diffusion-webui-images-browser)
- [LORA](https://github.com/kohya-ss/sd-scripts) *(both training and inference)*
- [LyCORIS](https://github.com/KohakuBlueleaf/LyCORIS) *(both training and inference)*
- [Model Converter](https://github.com/Akegarasu/sd-webui-model-converter)
- [CLiP Interrogator](https://github.com/pharmapsychotic/clip-interrogator-ext)
- [Dynamic Thresholding](https://github.com/mcmonkeyprojects/sd-dynamic-thresholding)
- [Steps Animation](https://github.com/vladmandic/sd-extension-steps-animation)
- [Seed Travel](https://github.com/yownas/seed_travel)
- [Multi-Diffusion Upscaler](https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111)
### User Interface
- Includes updated **UI**: reskinned and reorganized
Black and orange dark theme with fixed width options panels and larger previews
- Includes support for **Gradio themes**
*Settings* -> *User interface* -> *UI theme*
Link to themes list & previews: <https://huggingface.co/spaces/gradio/theme-gallery>
### Removed
- Drops compatibility with older versions of `python` and requires **3.9** or **3.10**
- Drops localizations
### Integrated CLI/API tools
Fork adds extra functionality:
- New skin and UI layout
- Ships with set of **CLI** tools that rely on *SD API* for execution:
e.g. `generate`, `train`, `bench`, etc.
[Full list](<cli/>)
<br>
- **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
- Any GPU compatibile with **DirectX** on *Windows* using **DirectML** libraries
This includes support for AMD GPUs that are not supported by native ROCm libraries
- **Intel Arc** GPUs using Intel OneAPI **Ipex/XPU** libraries
- **Apple M1/M2** on *OSX* using built-in support in Torch with **MPS** optimizations
## Install
1. Install first:
**Python** & **Git**
2. If you have nVidia GPU, install nVidia CUDA toolkit:
<https://developer.nvidia.com/cuda-downloads>
3. Clone repository
2. Clone repository
`git clone https://github.com/vladmandic/automatic`
3. Run launcher
`webui.bat` or `webui.sh`:
- Platform specific wrapper scripts For Windows, Linux and OSX
- Starts `launch.py` in a Python virtual environment (`venv`)
- Uses `install.py` to handle all actual requirements and dependencies
- *Note*: Server can run without virtual environment, but it is recommended to use it to avoid library version conflicts with other applications
## Run
*Note*: **nVidia/CUDA** and **AMD/ROCm** are auto-detected is present and available, but for any other use case specify required parameter explicitly or wrong packages may be installed as installer will assume CPU-only environment
Run desired startup script to install dependencies and extensions and start server:
Full startup sequence is logged in `webui.log`, so if you encounter any issues, please check it first
- `webui.bat` and `webui.sh`:
Platform specific wrapper scripts For Windows, Linux and OSX
Starts `launch.py` in a Python virtual environment (venv)
*Note*: Server can run without virtual environment, but it is recommended to use it to avoid library version conflicts with other applications
**If you're unsure which launcher to use, this is the one you want**
- `launch.py`:
Main startup script
Can be used directly to start server in a manually activated `venv` or to run server without `venv`
- `setup.py`:
Main installer, used by `launch.py`
Can also be used directly to update repository or extensions
If running manually, make sure to activate `venv` first (if used)
- `webui.py`:
Main server script
Below is partial list of all available parameters, run `webui --help` for the full list:
Any of the above scripts can be used with `--help` to display detailed usage information and available parameters
For example:
> webui.bat --help
Setup options:
--use-ipex Use Intel OneAPI XPU backend, default: False
--use-directml Use DirectML if no compatible GPU is detected, default: False
--use-cuda Force use nVidia CUDA backend, default: False
--use-rocm Force use AMD ROCm backend, default: False
--skip-update Skip update of extensions and submodules, 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
--reinstall Force reinstallation of all requirements, default: False
--debug Run installer with debug logging, default: False
--reset Reset main repository to latest version, default: False
--upgrade Upgrade main repository to latest version, default: False
--safe Run in safe mode with no user extensions
Full startup sequence is logged in `setup.log`, so if you encounter any issues, please check it first
<br>![screenshot](html/black-orange.jpg)<br>
## Update
## Notes
The launcher can perform automatic update of main repository, requirements, extensions and submodules:
### **Collab**
- **Main repository**:
Update is *not* performed by default, enable with `--upgrade` flag
- **Requirements**:
Check is performed on each startup and missing requirements are auto-installed
Can be skipped with `--skip-requirements` flag
- **Extensions and submodules**:
Update is performed on each startup and installer for each extension is started
Can be skipped with `--skip-extensions` flag
- **Quick mode**: Automatically enabled if timestamp of last sucessful setup is newer than actual repository version or version of newest extension
<br>
## Other
### Scripts
This repository comes with a large collection of scripts that can be used to process inputs, train, generate, and benchmark models
As well as number of auxiliary scripts that do not rely on **WebUI**, but can be used for end-to-end solutions such as extract frames from videos, etc.
For full details see [Docs](cli/README.md)
<br>
### Docs
- Scripts are in [Scripts](cli/README.md)
- Everything else is in [Wiki](https://github.com/vladmandic/automatic/wiki)
- Except my current [TODO](TODO.md)
- To avoid having this repo rely just on me, I'd love to have additional maintainers with full admin rights. If you're interested, ping me!
- 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 I would really love to have additional contibutors and/or maintainers to join and help lead the effords on different platforms
### **Goals**
The idea behind the fork is to enable latest technologies and advances in text-to-image generation
*Sometimes this is not the same as "as simple as possible to use"*
If you are looking an amazing simple-to-use Stable Diffusion tool, I'd suggest [InvokeAI](https://invoke-ai.github.io/InvokeAI/) specifically due to its automated installer and ease of use
General goals:
- Cross-platform
- Create uniform experience while automatically managing any platform specific differences
- Performance
- Enable best possible performance on all platforms
- Ease-of-Use
- Automatically handle all requirements, dependencies, flags regardless of platform
- Integrate all best options for uniform out-of-the-box experience without the need to tweak anything manually
- Look-and-Feel
- Create modern, intuitive and clean UI
- Up-to-Date
- Keep code up to date with latest advanced in text-to-image generation
## Credits
- Main credit goes to [Automatic1111 WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui)
- Additional credits are listed in [Credits](https://github.com/AUTOMATIC1111/stable-diffusion-webui/#credits)
- Licenses for modules are listed in [Licenses](html/licenses.html)
### **Docs**
- [Radme](README.md)
- [ToDo](TODO.md)
- [Changelog](CHANGELOG.md)
- [CLI Tools](cli/README.md)
<br>
TODO.md
+17 -18
View File
@@ -4,35 +4,26 @@
Stuff to be fixed...
- Move Restart Server from WebUI to Launch and reload modules
- Mdularize `cli` scripts
## Features
Stuff to be added...
- Update README
- Add Gradio theme maker
- Transformers version
- Create new GitHub hooks/actions for CI/CD
- Redo Extensions tab: see <https://vladmandic.github.io/sd-extension-manager/pages/extensions.html>
- Stream-load models as option for slow storage
- Auto-test `torch.layer_norm` for FP16
- Update `Wiki`
- Create new `GitHub` hooks/actions for CI/CD
- Import core repos
- Update `train.py` to use `interrogator`
- Update `train.py` to use `rembg`
## Investigate
Stuff to be investigated...
- Torch Compile
- `Torch-DirectML`
- `TensorRT`
- [Temporal Weighing](https://github.com/comfyanonymous/ComfyUI/discussions/473)
## Merge PRs
Pick & merge PRs from main repo...
- Merge backlog: <https://github.com/vladmandic/automatic/pulls>
- Compare commits: <https://github.com/vladmandic/automatic/compare/e59ebe25..b6289d5>
## Models
@@ -50,10 +41,18 @@ Tech that can be integrated as part of the core workflow...
- [Null-text inversion](https://github.com/ouhenio/null-text-inversion-colab)
- [Custom diffusion](https://github.com/guaneec/custom-diffusion-webui), [Custom diffusion](https://www.cs.cmu.edu/~custom-diffusion/)
- [Dream artist](https://github.com/7eu7d7/DreamArtist-sd-webui-extension)
- [QuickEmbedding](https://github.com/ethansmith2000/QuickEmbedding)
- `TensorRT`
## Random
- Bunch of stuff: <https://pharmapsychotic.com/tools.html>
### Pending Code Updates
- <https://towardsdatascience.com/mastering-memoization-in-python-dcdd8b435189>
- <https://github.com/AUTOMATIC1111/stable-diffusion-webui/compare/89f9faa...20ae71f>
- <https://github.com/vladmandic/automatic/discussions/1246>
- shared.info
- hints
- localization
- docker
- port `p.all_hr_prompts`
- test `lyco_patch_lora`
cli/.pylintrc
-3
View File
@@ -1,3 +0,0 @@
# See https://pylint.pycqa.org/en/latest/user_guide/messages/message_control.html
[MESSAGES CONTROL]
enable=C,R,W,E,I
cli/README.md
+30 -85
View File
@@ -1,9 +1,6 @@
# Stable-Diffusion Productivity Scripts
*Notes*:
- Offline scripts can be used with or without **Automatic WebUI**
- Online scripts rely on **Automatic WebUI** API which should be started with `--api` parameter
- All scripts have built-in `--help` parameter that can be used to get more information
Note: All scripts have built-in `--help` parameter that can be used to get more information
<br>
@@ -18,32 +15,29 @@ Supports upsampling, face restoration and grid creation
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
Textual inversion embedding training
> python train-ti.py
Combined pipeline for **embeddings**, **lora**, **lycoris**, **dreambooth** and **hypernetwork**
Optionally runs several image processing steps before training:
Combined pipeline:
1. Creates embedding
2. Extracts images if input is movie
3. Preprocesses images
4. Runs 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
LoRA training
> python train-lora.py
Combined pipeline:
1. Preprocesses images
2. Runs training
[Detailed documentation](https://github.com/vladmandic/automatic/wiki/Process.md)
LoRA extract from model
> python moidules/lora-extract.py
> python train.py
<br>
@@ -51,107 +45,58 @@ LoRA extract from model
### Benchmark
Benchmark your **Automatic WebUI**
Note: Requires SD API
> python run-benchmark.py
> python modules/bench.py
### Create Previews
### Embedding Previews
Create previews for **embeddings**, **lora**, **lycoris**, **dreambooth** and **hypernetwork**
Create previews of embeddings using preview templates
Note: Requires SD API
> python create-previews.py
> python modules/preview-embeddings.py
## Image Grid
## Grid
Create flexible image grids from any number of images
Note: Offline tool
> python modiles/grid.py
> python image-grid.py
### Image Watermark
Create invisible image watermark and remove existing EXIF tags
Note: Offline tool
> python modules/image-watermark.py
> python image-watermark.py
### Interrogate
### Image Interrogate
Runs CLiP and Booru image interrogation
Note: Requires SD API
> python modules/interrogate.py
### Interrogate-Offline
Standalone implementation of GiT, CLiP and ViT image interrogation
Note: Offline tool
> python modules/interrogate-offline.py
### Models Previews
Create previews of models using built-in templates
Note: Requires SD API
> python modules/preview-models.py
> python image-interrogate.py
### Palette Extract
Extract color palette from image(s)
Note: Offline tool
> python modules/palette-extract.py
### Image Process
Run image processing to extract face/body segments and run resolution/blur/dynamic-range checks
Note: Offline except for interrogate to generate caption files which requires SD API
> python modules/process.py
[Detailed documentation](https://github.com/vladmandic/automatic/wiki/Process.md)
> python image-palette.py
### Prompt Ideas
Generate complex prompt ideas
Note: Offline tool
> python modules/prompt-ideas.py
> python prompt-ideas.py
### Prompt Promptist
Attempts to beautify the provided prompt
Note: Offline tool
> python modules/promptist.py
### Training Loss-Chart
Create loss-chart from training log
Note: Offline tool, may require adjustment to train paths if used with other repos
> python modules/train-losschart.py
### Training Loss-Rate
Create customizable loss rate to be used in training
Note: Offline tool
> python modules/train-lossrate.py
> python prompt-promptist.py
### Video Extract
Extract frames from video files
Note: Offline tool
> python modules/video-extract.py
> python video-extract.py
<br>
## Utility Scripts
### SDAPI
Utility module that handles async communication to Automatic API endpoints
cli/modules/preview-models.py → cli/create-previews.py
+118 -38
View File
@@ -1,17 +1,15 @@
#!/bin/env python
#!/usr/bin/env python
import os
import sys
import json
import time
import importlib
import asyncio
import argparse
from pathlib import Path
from util import Map, log
from sdapi import get, post, close
from grid import grid
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from generate import sd, generate
from generate import generate # pylint: disable=import-error
grid = importlib.import_module('image-grid').grid
default = 'sd-v15-runwayml.ckpt [cc6cb27103]'
@@ -63,7 +61,7 @@ options = Map({
"sd_vae": "vae-ft-mse-840000-ema-pruned.ckpt",
},
'lora': {
'strength': 0.9,
'strength': 1.0,
},
'hypernetwork': {
'keyword': 'beautiful sexy woman',
@@ -72,13 +70,12 @@ options = Map({
})
async def models(params):
global sd
async def preview_models(params):
data = await get('/sdapi/v1/sd-models')
all = [m['title'] for m in data]
allmodels = [m['title'] for m in data]
models = []
excluded = []
for m in all: # loop through all registered models
for m in allmodels: # loop through all registered models
ok = True
for e in exclude: # check if model is excluded
if e in m:
@@ -100,30 +97,31 @@ async def models(params):
models = filtered
log.info({ 'models preview' })
log.info({ 'models': len(models), 'excluded': len(excluded) })
cmdflags = await get('/sdapi/v1/cmd-flags')
opt = await get('/sdapi/v1/options')
if params.output != '':
dir = params.output
folder = params.output
else:
dir = os.path.abspath(os.path.join(cmdflags['hypernetwork_dir'], '..', 'Stable-diffusion'))
log.info({ 'output directory': dir })
folder = os.path.abspath(os.path.join(opt['hypernetwork_dir'], '..', 'Stable-diffusion'))
log.info({ 'output directory': folder })
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:
fn = os.path.join(dir, os.path.basename(model) + options.format)
fn = os.path.join(folder, os.path.basename(model) + options.format)
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
log.info({ 'model preview exists': model })
continue
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, prompt in prompts:
options.generate.prompt = prompt
for label, p in 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:
@@ -139,26 +137,28 @@ async def models(params):
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'] != default and not params.fixed:
log.info({ 'model set default': default })
opt['sd_model_checkpoint'] = default
del opt['sd_lora']
del opt['sd_lyco']
await post('/sdapi/v1/options', opt)
async def lora(params):
cmdflags = await get('/sdapi/v1/cmd-flags')
dir = cmdflags['lora_dir']
if not os.path.exists(dir):
log.error({ 'lora directory not found': dir })
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 = [f for f in Path(dir).glob('*.safetensors')]
models2 = [f for f in Path(dir).glob('*.ckpt')]
models1 = [f for f in Path(folder).glob('*.safetensors')]
models2 = [f for f in Path(folder).glob('*.ckpt')]
models = [f.stem for f in models1 + models2]
log.info({ 'loras': len(models) })
for model in models:
fn = os.path.join(dir, model + options.format)
fn = os.path.join(folder, model + options.format)
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
log.info({ 'lora preview exists': model })
continue
@@ -166,7 +166,7 @@ async def lora(params):
labels = []
t0 = time.time()
import re
keywords = re.sub('\d', '', model)
keywords = re.sub(r'\d', '', model)
keywords = keywords.replace('-v', ' ').replace('-', ' ').strip().split(' ')
keyword = '\"' + '\" \"'.join(keywords) + '\"'
options.generate.prompt = prompt.replace('<keyword>', keyword)
@@ -188,16 +188,57 @@ async def lora(params):
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 hypernetwork(params):
cmdflags = await get('/sdapi/v1/cmd-flags')
dir = cmdflags['hypernetwork_dir']
if not os.path.exists(dir):
log.error({ 'hypernetwork directory not found': dir })
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
models = [f.stem for f in Path(dir).glob('*.pt')]
log.info({ 'loras': len(models) })
models1 = [f for f in Path(folder).glob('*.safetensors')]
models2 = [f for f in Path(folder).glob('*.ckpt')]
models = [f.stem for f in models1 + models2]
log.info({ 'lycos': len(models) })
for model in models:
fn = os.path.join(dir, model + options.format)
fn = os.path.join(folder, model + options.format)
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
log.info({ 'lyco preview exists': model })
continue
images = []
labels = []
t0 = time.time()
import re
keywords = re.sub(r'\d', '', model)
keywords = keywords.replace('-v', ' ').replace('-', ' ').strip().split(' ')
keyword = '\"' + '\" \"'.join(keywords) + '\"'
options.generate.prompt = 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()
image = grid(images = images, labels = labels, border = 8)
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 = [f.stem for f in Path(folder).glob('*.pt')]
log.info({ 'hypernetworks': len(models) })
for model in models:
fn = os.path.join(folder, model + options.format)
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
log.info({ 'hypernetwork preview exists': model })
continue
@@ -224,10 +265,49 @@ async def hypernetwork(params):
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 = [f.stem for f in Path(folder).glob('*.pt')]
log.info({ 'embeddings': len(models) })
for model in models:
fn = os.path.join(folder, model + '.preview' + options.format)
if os.path.exists(fn) and len(params.input) == 0: # if model preview exists and not manually included
log.info({ 'embedding preview exists': model })
continue
images = []
labels = []
t0 = time.time()
import re
keyword = '\"' + re.sub(r'\d', '', model) + '\"'
options.generate.batch_size = 4
options.generate.prompt = 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({ 'lyco': model, 'keyword': keyword, 'error': data })
t1 = time.time()
image = grid(images = images, labels = labels, border = 8)
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 models(params)
await preview_models(params)
await lora(params)
await lyco(params)
await hypernetwork(params)
await embedding(params)
await close()
@@ -236,5 +316,5 @@ if __name__ == '__main__':
parser.add_argument('--output', type = str, default = '', required = False, help = 'output directory')
parser.add_argument('--fixed', default = False, action='store_true', help = "do not change model")
parser.add_argument('input', type = str, nargs = '*')
params = parser.parse_args()
asyncio.run(create_previews(params))
args = parser.parse_args()
asyncio.run(create_previews(args))
cli/generate.py
+4 -5
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
# pylint: disable=no-member
"""generate batches of images from prompts and upscale them
@@ -32,9 +32,8 @@ from PIL import Image
from PIL.ExifTags import TAGS
from PIL.TiffImagePlugin import ImageFileDirectory_v2
sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
from modules.sdapi import close, get, interrupt, post, session
from modules.util import Map, log, safestring
from sdapi import close, get, interrupt, post, session
from util import Map, log, safestring
sd = {}
@@ -121,7 +120,7 @@ def sampler(params, options): # find sampler
async def generate(prompt = None, options = None, quiet = False): # pylint: disable=redefined-outer-name
global sd
global sd # pylint: disable=global-statement
if options:
sd = Map(options)
if prompt is not None:
cli/hfsearch.py
+17
View File
@@ -0,0 +1,17 @@
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',
tags='stable-diffusion',
library=['diffusers', 'stable-diffusion'],
)
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}' } for m in res]
print('Online', models)
cli/modules/grid.py → cli/image-grid.py
+1 -1
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
Create image grid
"""
cli/modules/interrogate.py → cli/image-interrogate.py
+3 -3
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
use clip to interrogate image(s)
"""
@@ -11,7 +11,7 @@ import asyncio
import filetype
from PIL import Image
from util import log, Map
import sdapi as sdapi
import sdapi
stats = { 'captions': {}, 'keywords': {} }
@@ -96,7 +96,7 @@ async def main():
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))
_caption, _keywords, _style = await interrogate(os.path.join(root, f))
else:
log.error({ 'interrogate unknown file type': arg })
else:
cli/modules/palette-extract.py → cli/image-palette.py
+21 -21
View File
@@ -1,27 +1,27 @@
#!/bin/env python
#!/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
import matplotlib.image as mpimg
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from colormap import rgb2hex
from PIL import Image
from util import log
from grid import grid
grid = importlib.import_module('image-grid').grid
def color_to_df(input):
colors_pre_list = str(input).replace('([(','').split(', (')[0:-1]
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]
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]),
@@ -30,14 +30,14 @@ def color_to_df(input):
return df
def palette(img, args, output):
def palette(img, params, output):
size = 1024
img.thumbnail((size, size), Image.HAMMING)
#crate dataframe
colors_x = extcolors.extract_from_image(img, tolerance = args.color, limit = 13)
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'])]
@@ -54,7 +54,7 @@ def palette(img, args, output):
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:
@@ -100,20 +100,20 @@ if __name__ == '__main__':
args = parser.parse_args()
log.info({ 'palette args': vars(args) })
if args.output != '':
pathlib.Path(args.output).mkdir(parents = True, exist_ok = True)
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):
img = Image.open(arg)
output = os.path.join(args.output, pathlib.Path(arg).stem + '-' + args.suffix + '.jpg')
palette(img, args, output)
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)):
img = Image.open(os.path.join(root, f))
output = os.path.join(args.output, pathlib.Path(f).stem + '-' + args.suffix + '.jpg')
palette(img, args, output)
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:
@@ -124,6 +124,6 @@ if __name__ == '__main__':
for f in files:
if filetype.is_image(os.path.join(root, f)):
images.append(Image.open(os.path.join(root, f)))
img = grid(images)
output = os.path.join(args.output, args.suffix + '.jpg')
palette(img, args, output)
image = grid(images)
fn = os.path.join(args.output, args.suffix + '.jpg')
palette(image, args, fn)
cli/modules/image-watermark.py → cli/image-watermark.py
+29 -29
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
import os
import io
import pathlib
@@ -44,31 +44,31 @@ def set_exif(d: dict):
ifd[_TAGS[k]] = v
exif_stream = io.BytesIO()
ifd.save(exif_stream)
bytes = b'Exif\x00\x00' + exif_stream.getvalue()
return bytes
encoded = b'Exif\x00\x00' + exif_stream.getvalue()
return encoded
def get_watermark(image, args):
def get_watermark(image, params):
data = np.asarray(image)
decoder = WatermarkDecoder(options.type, args.length)
bytes = decoder.decode(data, options.method)
decoder = WatermarkDecoder(options.type, params.length)
decoded = decoder.decode(data, options.method)
try:
watermark = str(bytes, 'UTF-8').replace('\x00', '')
s = str(decoded, 'UTF-8').replace('\x00', '')
except:
watermark = ''
return watermark
s = ''
return s
def set_watermark(image, args):
def set_watermark(image, params):
data = np.asarray(image)
encoder = WatermarkEncoder()
encoder.set_watermark(options.type, args.wm.encode('utf-8'))
encoder.set_watermark(options.type, params.wm.encode('utf-8'))
encoded = encoder.encode(data, options.method)
image = Image.fromarray(encoded)
return image
def watermark(args, file):
def watermark(params, file):
if not os.path.exists(file):
log.error({ 'watermark': 'file not found' })
return
@@ -82,30 +82,30 @@ def watermark(args, file):
exif = get_exif(image)
if args.command == 'read':
watermark = get_watermark(image, args)
log.info({ 'file': file, 'watermark': watermark, 'exif': exif, 'resolution': f'{image.width}x{image.height}' })
if params.command == 'read':
wm = get_watermark(image, params)
log.info({ 'file': file, 'watermark': wm, 'exif': exif, 'resolution': f'{image.width}x{image.height}' })
elif args.command == 'write':
metadata = b'' if args.strip else set_exif(exif)
if args.output != '':
pathlib.Path(args.output).mkdir(parents = True, exist_ok = True)
image=set_watermark(image, args)
fn = os.path.join(args.output, file)
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 args.verify:
if params.verify:
data = np.asarray(image)
decoder = WatermarkDecoder(options.type, args.length)
bytes = decoder.decode(data, options.method)
if bytes.startswith(b'\xff'):
watermark = ''
decoder = WatermarkDecoder(options.type, params.length)
decoded = decoder.decode(data, options.method)
if decoded.startswith(b'\xff'):
wm = ''
else:
watermark = str(bytes, 'UTF-8').replace('\x00', '')
wm = str(decoded, 'UTF-8').replace('\x00', '')
else:
watermark = args.wm
wm = params.wm
log.info({ 'file': fn, 'watermark': watermark, 'exif': None if args.strip else exif, 'resolution': f'{image.width}x{image.height}' })
log.info({ 'file': fn, 'watermark': wm, 'exif': None if params.strip else exif, 'resolution': f'{image.width}x{image.height}' })
if __name__ == '__main__':
cli/train/latents.py → cli/latents.py
+8 -8
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
import os
import sys
@@ -23,7 +23,7 @@ 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'))
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
@@ -44,13 +44,13 @@ options = Map({
vae = None
def get_latents(vae, images, weight_dtype):
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 = vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
latents = local_vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
return latents
@@ -58,8 +58,8 @@ 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(params):
args = Map({**options, **params})
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):
@@ -73,7 +73,7 @@ def create_vae_latents(params):
weight_dtype = torch.bfloat16
else:
weight_dtype = torch.float32
global vae
global vae # pylint: disable=global-statement
if vae is None:
vae = model_util.load_vae(args.vae, weight_dtype)
vae.eval()
@@ -142,7 +142,7 @@ def create_vae_latents(params):
def unload_vae():
global vae
global vae # pylint: disable=global-statement
vae = None
cli/modules/interrogate-offline.py
-166
View File
@@ -1,166 +0,0 @@
#!/bin/env python
import os
import gc
import json
import time
import argparse
import torch
import filetype
from PIL import Image
import transformers
from transformers import AutoProcessor, AutoModelForCausalLM
from transformers import BlipProcessor, BlipForConditionalGeneration
from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
from util import log, Map
model = None
processor = None
extractor = None
dtype = torch.float32
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
options = Map({
'input': '',
'min': 8,
'max': 256,
'beams': 1,
'json': '',
'txt': False,
'tag': '',
'git': True,
'blip': True,
'precision': 'fp16',
'model': 'git',
})
def cleanup(s: str):
s = s.split('"')[0].split('.')[0].split(' that')[0]
s = s.split(' with a letter')[0].split(' with the number')[0].split(' with the word')[0]
s = s.replace('arafed image of ', '')
return s.replace('a ', '')
def load_model(args):
global model
global processor
global extractor
transformers.logging.set_verbosity_error()
if args.model == 'git':
model_name = "microsoft/git-large-textcaps"
if model is None:
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=dtype)
model.to(device)
processor = AutoProcessor.from_pretrained(model_name, torch_dtype=dtype)
log.info( { 'interrogate loaded model': model_name })
elif args.model == 'blip':
model_name = "Salesforce/blip-image-captioning-large"
if model is None:
model = BlipForConditionalGeneration.from_pretrained(model_name, torch_dtype=dtype)
model.to(device)
processor = BlipProcessor.from_pretrained(model_name, torch_dtype=dtype)
log.info( { 'interrogate loaded model': model_name })
elif args.model == 'vit':
model_name = "nlpconnect/vit-gpt2-image-captioning"
if model is None:
model = VisionEncoderDecoderModel.from_pretrained(model_name, torch_dtype=dtype)
model.to(device)
extractor = ViTFeatureExtractor.from_pretrained(model_name, torch_dtype=dtype)
processor = AutoTokenizer.from_pretrained(model_name, torch_dtype=dtype)
log.info( { 'interrogate loaded model': model_name })
else:
log.info( { 'interrogate unknown model': args.model })
def interrogate_files(params, files):
args = Map({**options, **params})
data = [f for f in files if filetype.is_image(f)]
log.info({ 'interrogate files': len(files), 'images': len(data), 'args': args })
load_model(args)
metadata = {}
for image_path in data:
image = Image.open(image_path).convert('RGB')
caption = ''
if args.model == 'git':
inputs = processor(images=[image], return_tensors="pt").to(device)
ids = model.generate(pixel_values=inputs.pixel_values, num_beams=args.beams, min_length=args.min, max_length=args.max)
caption = processor.batch_decode(ids, skip_special_tokens=True)[0]
elif args.model == 'blip':
inputs = processor(image, return_tensors="pt").to(device, dtype)
ids = model.generate(**inputs, num_beams=args.beams, min_length=args.min, max_length=args.max)
caption = processor.decode(ids[0], skip_special_tokens=True)
elif args.model == 'vit':
inputs = extractor(images=[image], return_tensors="pt").pixel_values.to(device)
ids = model.generate(inputs, num_beams=args.beams, min_length=args.min, max_length=args.max)
caption = processor.batch_decode(ids, skip_special_tokens=True)[0]
else:
log.error({ 'interrogate unknown model': args.model })
caption = cleanup(caption)
tags = ''
if args.tag != '':
tags += args.tag + ','
tags += caption.split(' ')[0]
if args.txt:
with open(os.path.splitext(image_path)[0] + '.txt', "wt", encoding='utf-8') as f:
f.write(caption + "\n")
metadata[image_path] = { 'caption': caption, 'tags': tags }
log.info({ 'interrogate image': image_path, 'moodel': args.model, 'caption': caption, 'tags': tags })
if args.json != '':
with open(args.json, "wt", encoding='utf-8') as f:
f.write(json.dumps(metadata, indent=2) + "\n")
return metadata
def unload_model():
global processor
global model
global extractor
if model is not None:
del model
model = None
if processor is not None:
del processor
processor = None
if extractor is not None:
del extractor
extractor = None
gc.collect()
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()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description = 'image interrogate')
parser.add_argument('input', type=str, nargs='*', help='input file or directory')
parser.add_argument('--model', default = 'git', choices = ['git', 'blip', 'vit'], help = "which model to use")
parser.add_argument("--min", type=int, default=8, help="min length of caption")
parser.add_argument("--max", type=int, default=256, help="max length of caption")
parser.add_argument("--beams", type=int, default=1, help="number of beams to use")
parser.add_argument("--json", type=str, default='', help="output json file")
parser.add_argument("--tag", type=str, default='', help="append tag")
parser.add_argument('--txt', default = False, action='store_true', help = "write captions to text files")
params = parser.parse_args()
log.info({ 'interrogate args': vars(params) })
if len(params.input) == 0:
parser.print_help()
exit(1)
files = []
for loc in params.input:
if os.path.isfile(loc):
files.append(loc)
elif os.path.isdir(loc):
for root, _sub_dirs, dir in os.walk(loc):
files = [os.path.join(root, f) for f in dir]
t0 = time.time()
metadata = interrogate_files(vars(params), files)
t1 = time.time()
log.info({ 'interrogate files': len(files), 'time': round(t1 - t0, 2) })
unload_model()
cli/modules/lora-extract.py
-144
View File
@@ -1,144 +0,0 @@
#!/bin/env python
"""
Extract approximating LoRA by SVD from two SD models
Based on: <https://github.com/kohya-ss/sd-scripts/blob/main/networks/extract_lora_from_models.py>
"""
import os
import sys
import time
import argparse
import torch
import transformers
from tqdm import tqdm
from util import log
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..', 'modules', 'lora'))
import library.model_util as model_util
import networks.lora as lora
def svd(args): # pylint: disable=redefined-outer-name
device = 'cuda' if torch.cuda.is_available() and args.device == 'cuda' else 'cpu'
transformers.logging.set_verbosity_error()
CLAMP_QUANTILE = 0.99
MIN_DIFF = 1e-6
if args.precision == 'fp32':
save_dtype = torch.float
elif args.precision == 'fp16':
save_dtype = torch.float16
elif args.precision == 'bf16':
save_dtype = torch.bfloat16
else:
save_dtype = None
t0 = time.time()
log.info({ 'loading model': args.original })
text_encoder_o, _, unet_o = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.original)
log.info({ 'loading model': args.tuned })
text_encoder_t, _, unet_t = model_util.load_models_from_stable_diffusion_checkpoint(args.v2, args.tuned)
with torch.no_grad():
torch.cuda.empty_cache()
# create LoRA network to extract weights: Use dim (rank) as alpha
lora_network_o = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_o, unet_o)
lora_network_t = lora.create_network(1.0, args.dim, args.dim, None, text_encoder_t, unet_t)
assert len(lora_network_o.text_encoder_loras) == len(lora_network_t.text_encoder_loras), 'model version is different'
# get diffs
diffs = {}
text_encoder_different = False
for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.text_encoder_loras, lora_network_t.text_encoder_loras)):
lora_name = lora_o.lora_name
module_o = lora_o.org_module
module_t = lora_t.org_module
diff = module_t.weight - module_o.weight
# Text Encoder might be same
if torch.max(torch.abs(diff)) > MIN_DIFF:
text_encoder_different = True
diff = diff.float()
diffs[lora_name] = diff
if not text_encoder_different:
log.info({ 'lora': 'text encoder is same, extract U-Net only' })
lora_network_o.text_encoder_loras = []
diffs = {}
for i, (lora_o, lora_t) in enumerate(zip(lora_network_o.unet_loras, lora_network_t.unet_loras)):
lora_name = lora_o.lora_name
module_o = lora_o.org_module
module_t = lora_t.org_module
diff = module_t.weight - module_o.weight
diff = diff.float()
diff = diff.to(device)
diffs[lora_name] = diff
t1 = time.time()
log.info({ 'lora models': 'ready', 'time': round(t1 - t0, 2) })
# make LoRA with svd
log.info({ 'lora': 'calculating by svd' })
rank = args.dim
lora_weights = {}
with torch.no_grad():
for lora_name, mat in tqdm(list(diffs.items())):
conv2d = len(mat.size()) == 4
if conv2d:
mat = mat.squeeze()
U, S, Vh = torch.linalg.svd(mat)
U = U[:, :rank]
S = S[:rank]
U = U @ torch.diag(S)
Vh = Vh[:rank, :]
dist = torch.cat([U.flatten(), Vh.flatten()])
hi_val = torch.quantile(dist, CLAMP_QUANTILE)
low_val = -hi_val
U = U.clamp(low_val, hi_val)
Vh = Vh.clamp(low_val, hi_val)
lora_weights[lora_name] = (U, Vh)
t2 = time.time()
# make state dict for LoRA
lora_network_o.apply_to(text_encoder_o, unet_o, text_encoder_different, True)
lora_sd = lora_network_o.state_dict()
log.info({ 'lora extracted weights': len(lora_sd), 'time': round(t2 - t1, 2) })
for key in list(lora_sd.keys()):
if 'alpha' in key:
continue
lora_name = key.split('.')[0]
i = 0 if 'lora_up' in key else 1
weights = lora_weights[lora_name][i]
# print(key, i, weights.size(), lora_sd[key].size())
if len(lora_sd[key].size()) == 4: # pylint: disable=unsubscriptable-object
weights = weights.unsqueeze(2).unsqueeze(3)
assert weights.size() == lora_sd[key].size(), f'size unmatch: {key}' # pylint: disable=unsubscriptable-object
lora_sd[key] = weights # pylint: disable=unsupported-assignment-operation
# load state dict to LoRA and save it
info = lora_network_o.load_state_dict(lora_sd)
log.info({ 'lora loading extracted weights': info })
dir_name = os.path.dirname(args.save)
if dir_name and not os.path.exists(dir_name):
os.makedirs(dir_name, exist_ok=True)
# minimum metadata
metadata = {'ss_network_dim': str(args.dim), 'ss_network_alpha': str(args.dim)}
lora_network_o.save_weights(args.save, save_dtype, metadata)
t3 = time.time()
log.info({ 'lora saved weights': args.save, 'time': round(t3 - t2, 2) })
if __name__ == '__main__':
parser = argparse.ArgumentParser(description = 'extract lora weights')
parser.add_argument('--v2', action='store_true', help='load Stable Diffusion v2.x model / Stable Diffusion')
parser.add_argument('--precision', type=str, default='fp16', choices=[None, 'fp32', 'fp16', 'bf16'], help='precision in saving, same to merging if omitted')
parser.add_argument('--device', type=str, default='cuda', choices=['cpu', 'cuda'], help='use cpu or cuda if available')
parser.add_argument('--original', type=str, default=None, required=True, help='Stable Diffusion original model: ckpt or safetensors file')
parser.add_argument('--tuned', type=str, default=None, required=True, help='Stable Diffusion tuned model, LoRA is difference of `original to tuned`: ckpt or safetensors file')
parser.add_argument('--save', type=str, default=None, required=True, help='destination file name: ckpt or safetensors file')
parser.add_argument('--dim', type=int, default=4, help='dimension (rank) of LoRA')
args = parser.parse_args()
log.info({ 'extract lora args': vars(args) })
if not os.path.exists(args.original) or not os.path.exists(args.tuned):
log.error({ 'models not found': [args.original, args.tuned] })
else:
svd(args)
cli/modules/lora-latents.py
-160
View File
@@ -1,160 +0,0 @@
#!/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 log, Map
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(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 = vae.encode(img_tensors).latent_dist.sample().float().to('cpu').numpy()
return latents
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(params):
args = Map({**options, **params})
log.info({ 'latents args': args })
if args.steps % 8 > 0:
log.warning({ 'latents': 'resolution is not multiple of 8' })
image_paths = train_util.glob_images(args.input)
if os.path.exists(args.json):
log.info({ 'latents metadata': args.json, 'images': len(image_paths) })
with open(args.json, 'rt', encoding='utf-8') as f:
metadata = json.load(f)
else:
log.error({ 'latents metadata missing': args.json })
return
if args.precision == 'fp16':
weight_dtype = torch.float16
elif args.precision == 'bf16':
weight_dtype = torch.bfloat16
else:
weight_dtype = torch.float32
global vae
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()
else:
log.warning({ 'latents': 'min and max are ignored if noupscale is set' })
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 = 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)
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:
log.info({ 'latents bucket': i, 'resolution': reso, 'count': 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
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/modules/models-diff.py
-74
View File
@@ -1,74 +0,0 @@
#!/bin/env python
# based on <https://huggingface.co/JosephusCheung/ASimilarityCalculatior>
import safetensors
import sys
import torch
from pathlib import Path
import torch.nn as nn
import torch.nn.functional as F
import warnings
from util import log
warnings.filterwarnings("ignore", category=UserWarning)
def cal_cross_attn(to_q, to_k, to_v, rand_input):
hidden_dim, embed_dim = to_q.shape
attn_to_q = nn.Linear(hidden_dim, embed_dim, bias=False)
attn_to_k = nn.Linear(hidden_dim, embed_dim, bias=False)
attn_to_v = nn.Linear(hidden_dim, embed_dim, bias=False)
attn_to_q.load_state_dict({"weight": to_q})
attn_to_k.load_state_dict({"weight": to_k})
attn_to_v.load_state_dict({"weight": to_v})
return torch.einsum(
"ik, jk -> ik",
F.softmax(torch.einsum("ij, kj -> ik", attn_to_q(rand_input), attn_to_k(rand_input)), dim=-1),
attn_to_v(rand_input)
)
def load_model(path):
if path.suffix == ".safetensors":
return safetensors.torch.load_file(path, device="cpu")
else:
ckpt = torch.load(path, map_location="cpu")
return ckpt["state_dict"] if "state_dict" in ckpt else ckpt
def eval(model, n, input):
qk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"
uk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_k.weight"
vk = f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_v.weight"
atoq, atok, atov = model[qk], model[uk], model[vk]
attn = cal_cross_attn(atoq, atok, atov, input)
return attn
def main():
file1 = Path(sys.argv[1])
files = sys.argv[2:]
seed = 114514
torch.manual_seed(seed)
model_a = load_model(file1)
log.info(f"base: {file1.name}")
map_attn_a = {}
map_rand_input = {}
for n in range(3, 11):
hidden_dim, embed_dim = model_a[f"model.diffusion_model.output_blocks.{n}.1.transformer_blocks.0.attn1.to_q.weight"].shape
rand_input = torch.randn([embed_dim, hidden_dim])
map_attn_a[n] = eval(model_a, n, rand_input)
map_rand_input[n] = rand_input
del model_a
for file2 in files:
file2 = Path(file2)
model_b = load_model(file2)
sims = []
for n in range(3, 11):
attn_a = map_attn_a[n]
attn_b = eval(model_b, n, map_rand_input[n])
sim = torch.mean(torch.cosine_similarity(attn_a, attn_b))
sims.append(sim)
log.info(f"{file2}: {torch.mean(torch.stack(sims)) * 1e2:.2f}%")
if __name__ == "__main__":
main()
cli/modules/preview-embeddings.py
-101
View File
@@ -1,101 +0,0 @@
#!/bin/env python
"""
create preview images from embeddings
"""
import os
import io
import sys
import json
import base64
import argparse
from pathlib import Path
from PIL import Image
from inspect import getsourcefile
from util import Map, log
from sdapi import getsync, postsync
from grid import grid
template = 'photo of "{name}", {suffix}, high detailed, skin texture, looking forward, facing camera, 135mm, shot on dslr, 4k, modelshoot style'
img2img_options = Map({
'prompt': None,
'negative_prompt': 'cartoon, drawing, cgi, sketch, comic, disfigured, deformed',
'init_images': [],
'sampler_name': 'DPM2 Karras',
'batch_size': 4,
'n_iter': 1,
'steps': 30,
'cfg_scale': 6,
'width': 512,
'height': 512,
'restore_faces': False
})
def encode(f):
img = Image.open(f)
with io.BytesIO() as stream:
img.save(stream, 'JPEG')
values = stream.getvalue()
encoded = base64.b64encode(values).decode()
return encoded
def create_preview(name: str, suffix: str):
options = getsync('/sdapi/v1/options')
cmdflags = getsync('/sdapi/v1/cmd-flags')
img2img_options['prompt'] = template.format(name = name, suffix = suffix)
log.debug({ 'preview options': img2img_options })
if len(img2img_options['init_images']) == 0:
for i in range(img2img_options.batch_size):
mask = os.path.join(os.path.dirname(getsourcefile(lambda:0)), 'preview-template'+ str(i+1) +'.jpg')
if (not os.path.isfile(mask)):
log.error({ 'preview': 'missing preview mask' })
return
img2img_options['init_images'].append(encode(mask))
data = postsync('/sdapi/v1/img2img', img2img_options)
if 'error' in data:
log.error({ 'preview': data['error'], 'reason': data['reason'] })
return
info = Map(json.loads(data['info']))
if not 'images' in data:
log.error({ 'preview': 'no images' })
return
fn = os.path.join(cmdflags.embeddings_dir, name + '.preview.png')
log.info({ 'preview': { 'name': fn, 'model': options.sd_model_checkpoint, 'seed': info.seed } })
images = []
for b64 in data['images']:
images.append(Image.open(io.BytesIO(base64.b64decode(b64.split(",",1)[0]))))
image = grid(images, None, square=True)
image.save(fn)
if __name__ == "__main__":
log.info({ 'preview': 'start' })
cmdflags = getsync('/sdapi/v1/cmd-flags')
parser = argparse.ArgumentParser(description = 'generate embeddings previews')
parser.add_argument('--overwrite', default = False, action='store_true', help = 'overwrite existing previews')
parser.add_argument('input', type=str, nargs='*')
params = parser.parse_args()
if len(params.input) == 0:
files = list(Path(cmdflags.embeddings_dir).glob('*.pt'))
else:
files = list(os.path.join(cmdflags.embeddings_dir, a + '.pt') for a in params.input if os.path.isfile(os.path.join(cmdflags.embeddings_dir, a + '.pt')))
candidates = [str(f) for f in files]
candidates.sort(key=os.path.getctime, reverse=True)
files = []
for f in candidates:
fn = f.replace('.pt', '.preview.png')
if os.path.isfile(f.replace('.pt', '.preview.png')):
if params.overwrite:
log.info({ 'preview add': fn })
files.append(f)
else:
log.info({ 'preview skip': fn })
else:
log.info({ 'preview add': fn })
files.append(f)
log.info({ 'preview embeddings': len(files) })
for f in files:
name = Path(f).stem
create_preview(name, 'person')
cli/modules/preview-template1.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 7.1 KiB

cli/modules/preview-template2.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 8.0 KiB

cli/modules/preview-template3.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 7.6 KiB

cli/modules/preview-template4.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 9.1 KiB

cli/modules/process.py
-500
View File
@@ -1,500 +0,0 @@
#!/bin/env python
"""
process people images
- check image resolution
- runs detection of face and body
- extracts crop and performs checks:
- visible: is face or body detected
- in frame: for face based on box, for body based on number of visible keypoints
- resolution: is cropped image still of sufficient resolution
- optionaly upsample and restore face quality
- blur: is image sharp enough
- dynamic range: is image bright enough
- similarity: compares image to all previously processed images to see if its unique enough
- images are resized and optionally squared
- face additionally runs through semantic segmentation to remove background
- if image passes checks
image padded and saved as extracted image
- body requires that face is detected and in-frame,
but does not have to pass all other checks as body performs its own checks
- runs clip interrogation on extracted images to generate filewords
"""
import os
import sys
import io
import math
import base64
import pathlib
import argparse
import logging
import filetype
import numpy as np
import mediapipe as mp
from PIL import Image, ImageOps
from skimage.metrics import structural_similarity as ssim
from scipy.stats import beta
sys.path.append(os.path.join(os.path.dirname(__file__)))
from util import log, Map
from sdapi import postsync
params = Map({
# general settings, do not modify
'src': '', # source folder
'dst': '', # destination folder
'clear_dst': True, # remove all files from destination at the start
'format': '.jpg', # image format
'target_size': 512, # target resolution
'square_images': True, # should output images be squared
'segmentation_model': 0, # segmentation model 0/general 1/landscape
'segmentation_background': (192, 192, 192), # segmentation background color
'blur_samplesize': 60, # sample size to use for blur detection
'similarity_size': 64, # base similarity detection on reduced images
# original image processing settings
'keep_original': False, # keep original image
# face processing settings
'extract_face': False, # extract face from image
'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
'face_blur': False, # check for body blur
'face_blur_score': 1.5, # max score for face blur detection
'face_range': False, # check for body blur
'face_range_score': 0.15, # min score for face dynamic range detection
'face_restore': False, # attempt to restore face quality
'face_upscale': False, # attempt to scale small faces
'face_segmentation': False, # segmentation enabled
# body processing settings
'extract_body': False, # extract body from image
'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
'body_blur': False, # check for body blur
'body_blur_score': 1.8, # max score for body blur detection
'body_range': False, # check for body blur
'body_range_score': 0.15, # min score for body dynamic range detection
'body_segmentation': False, # segmentation enabled
# similarity detection settings
'similarity_score': 0.8, # maximum similarity score before image is discarded
# interrogate settings
'interrogate_model': ['clip', 'deepdanbooru'], # interrogate models
'interrogate_captions': True, # write captions to file
'tag_limit': 5, # number of tags to extract
})
face_model = None
body_model = None
segmentation_model = None
def detect_blur(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 - params.blur_samplesize: cy + params.blur_samplesize, cx - params.blur_samplesize: cx + params.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):
# 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]
brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2])
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)
images = []
def detect_simmilar(image):
img = image.resize((params.similarity_size, params.similarity_size))
img = ImageOps.grayscale(img)
data = np.array(img)
similarity = 0
for i in images:
val = ssim(data, i, data_range=255, channel_axis=None, gradient=False, full=False)
if val > similarity:
similarity = val
images.append(data)
return similarity
def segmentation(image):
global segmentation_model
if segmentation_model is None:
segmentation_model = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=params.segmentation_model)
data = np.array(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[:] = params.segmentation_background
data = np.where(condition, data, background) # consider using a joint bilateral filter instead of pure combine
segmented = Image.fromarray(data)
return segmented
def extract_face(img):
if not params.extract_face:
return None, True
if img.mode == 'RGBA':
img = img.convert('RGB')
scale = max(img.size[0], img.size[1]) / params.target_size
resized = img.copy()
resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
global face_model
if face_model is None:
face_model = mp.solutions.face_detection.FaceDetection(min_detection_confidence=params.face_score, model_selection=params.face_model)
results = face_model.process(np.array(resized))
if results.detections is None:
return None, False
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:
log.info({ 'process face skip': 'out of frame' })
return None, False
x = (box.xmin - params.face_pad / 2) * resized.width
y = (box.ymin - params.face_pad / 2)* resized.height
w = (box.width + params.face_pad) * resized.width
h = (box.height + params.face_pad) * resized.height
cx = x + w / 2
cy = y + h / 2
l = max(w, h) / 2
square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)]
cropped = img.crop(tuple(square))
upscale = 1
if params.face_restore or params.face_upscale:
if (cropped.size[0] < params.target_size or cropped.size[1] < params.target_size) and params.face_upscale:
upscale = 2
kwargs = Map({
'image': encode(cropped),
'upscaler_1': 'SwinIR_4x' if params.face_upscale else None,
'codeformer_visibility': 1.0 if params.face_restore else 0.0,
'codeformer_weight': 0.15 if params.face_restore else 0.0,
'upscaling_resize': upscale,
})
original = [cropped.size[0], cropped.size[1]]
res = postsync('/sdapi/v1/extra-single-image', kwargs)
if 'image' not in res:
log.error({ 'process face': 'upscale failed' })
raise ValueError('upscale failed')
cropped = Image.open(io.BytesIO(base64.b64decode(res['image'])))
kwargs.image = [cropped.size[0], cropped.size[1]]
upscaled = [cropped.size[0], cropped.size[1]]
upscale = False if upscale == 1 else { 'original': original, 'upscaled': upscaled }
log.info({ 'process face restore': params.face_restore, 'upscale': upscale })
if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size:
log.info({ 'process face skip': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
return None, True
cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
if params.square_images:
squared = Image.new('RGB', (params.target_size, params.target_size))
squared.paste(cropped, ((params.target_size - cropped.width) // 2, (params.target_size - cropped.height) // 2))
if params.face_segmentation:
squared = segmentation(squared)
else:
squared = cropped
if params.face_blur:
blur = detect_blur(squared)
if blur > params.face_blur_score:
log.info({ 'process face skip': 'blur check fail', 'blur': blur })
return None, True
else:
log.debug({ 'process face blur': blur })
if params.face_range:
range = detect_dynamicrange(squared)
if range < params.face_range_score:
log.info({ 'process face skip': 'dynamic range check fail', 'range': range })
return None, True
else:
log.debug({ 'process face dynamic range': range })
similarity = detect_simmilar(squared)
if similarity > params.similarity_score:
log.info({ 'process face skip': 'similarity check fail', 'score': round(similarity, 2) })
return None, True
return squared, True
def extract_body(img):
if not params.extract_body:
return None, True
if img.mode == 'RGBA':
img = img.convert('RGB')
scale = max(img.size[0], img.size[1]) / params.target_size
resized = img.copy()
resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
global body_model
if body_model is None:
body_model = mp.solutions.pose.Pose(static_image_mode=True, min_detection_confidence=params.body_score, model_complexity=params.body_model)
results = body_model.process(np.array(resized))
if results.pose_landmarks is None:
return None, False
x = [resized.width * (i.x - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility]
y = [resized.height * (i.y - params.body_pad / 2) for i in results.pose_landmarks.landmark if i.visibility > params.body_visibility]
if len(x) < params.body_parts:
log.info({ 'process body skip': 'insufficient body parts', 'detected': len(x) })
return None, True
w = max(x) - min(x) + resized.width * params.body_pad
h = max(y) - min(y) + resized.height * params.body_pad
cx = min(x) + w / 2
cy = min(y) + h / 2
l = max(w, h) / 2
square = [scale * (cx - l), scale * (cy - l), scale * (cx + l), scale * (cy + l)]
square = [max(square[0], 0), max(square[1], 0), min(square[2], img.width), min(square[3], img.height)]
cropped = img.crop(tuple(square))
if cropped.size[0] < params.target_size and cropped.size[1] < params.target_size:
log.info({ 'process body skip': 'low resolution', 'size': [cropped.size[0], cropped.size[1]] })
return None, True
cropped.thumbnail((params.target_size, params.target_size), Image.HAMMING)
if params.square_images:
squared = Image.new('RGB', (params.target_size, params.target_size))
squared.paste(cropped, ((params.target_size - cropped.width) // 2, (params.target_size - cropped.height) // 2))
if params.body_segmentation:
squared = segmentation(squared)
else:
squared = cropped
if params.body_blur:
blur = detect_blur(squared)
if blur > params.body_blur_score:
log.info({ 'process body skip': 'blur check fail', 'blur': blur })
return None, True
else:
log.debug({ 'process body blur': blur })
if params.body_range:
range = detect_dynamicrange(squared)
if range < params.body_range_score:
log.info({ 'process body skip': 'dynamic range check fail', 'range': range })
return None, True
else:
log.debug({ 'process body dynamic range': range })
similarity = detect_simmilar(squared)
if similarity > params.similarity_score:
log.info({ 'process body skip': 'similarity check fail', 'score': round(similarity, 2) })
return None, True
return squared, True
def save_original(img):
if img.mode == 'RGBA':
img = img.convert('RGB')
resized = img.copy()
resized.thumbnail((params.target_size, params.target_size), Image.HAMMING)
if params.square_images:
squared = Image.new('RGB', (params.target_size, params.target_size))
squared.paste(resized, ((params.target_size - resized.width) // 2, (params.target_size - resized.height) // 2))
else:
squared = resized
return squared
def encode(img):
with io.BytesIO() as stream:
img.save(stream, 'JPEG')
values = stream.getvalue()
encoded = base64.b64encode(values).decode()
return encoded
def interrogate(img, fn, intag = None):
if len(params.interrogate_model) == 0:
return
caption = ''
tags = []
for model in params.interrogate_model:
json = Map({ 'image': encode(img), 'model': model })
res = postsync('/sdapi/v1/interrogate', json)
if model == 'clip':
caption = res.caption if 'caption' in res else ''
caption = caption.split(',')[0].replace('a ', '')
if intag is not None:
caption = intag + ', ' + caption
if model == 'deepdanbooru':
tag = res.caption if 'caption' in res else ''
tags = tag.split(',')
tags = [t.replace('(', '').replace(')', '').replace('\\', '').split(':')[0].strip() for t in tags]
if intag is not None:
for t in intag.split(',')[::-1]:
tags.insert(0, t.strip())
if params.interrogate_captions:
file = fn.replace(params.format, '.txt')
f = open(file, 'w')
f.write(caption)
f.close()
pos = 0 if len(tags) == 0 else 1
tags.insert(pos, caption.split(' ')[1])
if len(tags) > params.tag_limit:
tags = tags[:params.tag_limit]
log.info({ 'interrogate': caption, 'tags': tags })
return caption, tags
i = {}
metadata = Map({})
# entry point when used as module
def process_file(f: str, dst: str = None, preview: bool = False, offline: bool = False, txt = None, tag = None, opts = []):
def save(img, f, what):
i[what] = i.get(what, 0) + 1
if dst is None:
dir = os.path.dirname(f)
else:
dir = dst
base = os.path.basename(f).split('.')[0]
parent = os.path.basename(pathlib.Path(dir))
basename = str(i[what]).rjust(3, '0') + '-' + what + '-' + base
fn = basename + params.format
# log.debug({ 'save': fn })
caption = ''
tags = ''
if not preview:
img.save(os.path.join(dir, fn))
if not offline:
caption, tags = interrogate(img, os.path.join(dir, fn), tag)
metadata[os.path.join(parent, basename)] = { 'caption': caption, 'tags': ','.join(tags) }
return fn
# overrides
if len(opts) > 0:
params.keep_original = True if 'original' in opts else False
params.extract_face = True if 'face' in opts else False
params.extract_body = True if 'body' in opts else False
params.face_blur = True if 'blur' in opts else False
params.body_blur = True if 'blur' in opts else False
params.face_range = True if 'range' in opts else False
params.body_range = True if 'range' in opts else False
params.face_upscale = True if 'upscale' in opts else False
params.face_restore = True if 'restore' in opts else False
log.info({ 'processing': f })
try:
image = Image.open(f)
except Exception as err:
log.error({ 'image': f, 'error': err })
return 0, {}
image = ImageOps.exif_transpose(image) # rotate image according to EXIF orientation
if txt is not None:
params.interrogate_captions = txt
if image.width < 512 or image.height < 512:
log.info({ 'process skip': 'low resolution', 'resolution': [image.width, image.height] })
return 0, {}
log.debug({ 'resolution': [image.width, image.height], 'mp': round((image.width * image.height) / 1024 / 1024, 1) })
face, ok = extract_face(image)
if face is not None:
fn = save(face, f, 'face')
log.info({ 'extract face': fn })
else:
log.debug({ 'no face': f })
if not ok:
return 0, {}
body, ok = extract_body(image)
if body is not None:
fn = save(body, f, 'body')
log.info({ 'extract body': fn })
else:
log.debug({ 'no body': f })
if params.keep_original:
resized = save_original(image)
fn = save(resized, f, 'original')
log.info({ 'original': fn })
image.close()
return i, metadata
def process_images(src: str, dst: str, args = None):
params.src = src
params.dst = dst
if args is not None:
params.update(args)
log.info({ 'processing': params })
if not os.path.isdir(src):
log.error({ 'process': 'not a folder', 'src': src })
else:
if os.path.isdir(dst) and params.clear_dst:
log.info({ 'clear dst': dst })
i = [os.path.join(dst, f) for f in os.listdir(dst) if os.path.isfile(os.path.join(dst, f)) and filetype.is_image(os.path.join(dst, f))]
for f in i:
os.remove(f)
pathlib.Path(dst).mkdir(parents=True, exist_ok=True)
for root, _sub_dirs, files in os.walk(src):
for f in files:
i, _metadata = process_file(os.path.join(root, f), dst)
return i
def unload_models():
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
if __name__ == '__main__':
# log.setLevel(logging.DEBUG)
parser = argparse.ArgumentParser(description = 'dataset processor')
parser.add_argument('--output', type=str, required=True, help='folder to store images')
parser.add_argument('--preview', default=False, action='store_true', help = "run processing but do not store results")
parser.add_argument('--offline', default=False, action='store_true', help = "run only processing steps that do not require running server")
parser.add_argument('--debug', default=False, action='store_true', help = "enable debug logging")
parser.add_argument('input', type=str, nargs='*')
args = parser.parse_args()
params.dst = args.output
if args.debug:
log.setLevel(logging.DEBUG)
log.debug({ 'debug': True })
log.info({ 'processing': params })
if not os.path.exists(params.dst) and not args.preview:
pathlib.Path(params.dst).mkdir(parents=True, exist_ok=True)
files = []
for loc in args.input:
if os.path.isfile(loc):
files.append(loc)
elif os.path.isdir(loc):
for root, _sub_dirs, dir in os.walk(loc):
for f in dir:
files.append(os.path.join(root, f))
for f in files:
process_file(f, params.dst, args.preview, args.offline)
log.info({ 'processed': i, 'inputs': len(files) })
# print(json.dumps(metadata, indent=2))
cli/modules/train-losschart.py
-191
View File
@@ -1,191 +0,0 @@
#!/bin/env python
import io
import os
import sys
import json
import pathlib
import logging
import torch
import numpy as np
from PIL import Image, ImageFont, ImageDraw
from matplotlib import pyplot as plt
from util import log, Map
def settings(logdir: str, name: str):
filename = os.path.join(logdir, name, 'settings.json')
with open(filename, 'r', encoding='utf-8') as f:
data = json.load(f)
data = Map(data)
log.debug({ 'settings': data })
return data
def plot(logdir: str, name: str):
f = os.path.join(logdir, name, 'train.csv')
if not os.path.isfile(f):
log.debug({ 'train log missing': f })
return
name = pathlib.Path(f).parent.name
img = os.path.join(logdir, f"{name}.train.png")
step, loss, rate = plt.np.loadtxt(f, delimiter = ',', skiprows = 1, usecols = [0, 3, 4], unpack = True)
d = settings(logdir, name)
# window = d.get('gradient_step', 1) * d.get('batch_size', 1)
window = d.get('save_embedding_every', 1)
try:
log.debug({ 'loss plot': name, 'output': img, 'data': f, 'records': len(step) })
except:
return # no data
if len(step) < 5:
return
plt.rcParams.update({'font.variant':'small-caps'})
plt.rc('axes', edgecolor='gray')
plt.rc('font', size=10)
plt.rc('font', variant='small-caps')
plt.grid(color='gray', linewidth=1, axis='both', alpha=0.5)
plt.rcParams['figure.figsize'] = [14, 14]
plt.rcParams['figure.facecolor'] = 'black'
figure, axis = plt.subplots(2, 1)
# create top graph
ax0 = axis[0]
ax0.set_facecolor(color = (0.1, 0.1, 0.1, 0.5))
ax0.tick_params(axis='x', labelcolor='white')
ax0.set_xlabel('step'.upper(), color='white')
ax0.set_xlim(0, d.steps)
ax0.set_ylim(0, 0.5)
ax0.set_axisbelow(True)
ax0.xaxis.grid(color='gray', linestyle='dashed')
ax0.yaxis.grid(color='gray', linestyle='dashed')
# loss values
ax0.plot(step, loss, color='gray', label='loss value')
ax0.set_ylabel('loss value'.upper(), color='#CE6400')
ax0.tick_params(axis='y', labelcolor='#CE6400')
# trendline
z = np.polyfit(step, loss, 1)
p = np.poly1d(z)
ax0.plot(step, p(loss), color='#5020F0', linewidth=3, label='loss trendline')
# moving average
if len(loss) > window:
maval = []
minval = []
maxval = []
for ind in range(window - 1):
maval.insert(0, np.nan)
minval.insert(0, np.nan)
maxval.insert(0, np.nan)
for ind in range(len(loss) - window + 1):
maval.append(np.mean(loss[ind:ind+window]))
minval.append(np.min(loss[ind:ind+window]))
maxval.append(np.max(loss[ind:ind+window]))
ax0.plot(step, maval, color='#CE6400', linewidth=5, label='average loss value')
ax0.plot(step, minval, color='#500010', linewidth=5, label='min loss per epoch')
ax0.plot(step, maxval, color='#005010', linewidth=5, label='max loss per epoch')
# learning rate
ax0_right = ax0.twinx()
ax0_right.set_ylabel('learn rate'.upper(), color='cyan')
ax0_right.plot(step, rate, color='cyan', linewidth=3, linestyle='dashed', label='learn rate')
ax0_right.tick_params(axis='y', labelcolor='cyan')
# axis legend
handles0, labels0 = ax0.get_legend_handles_labels() # because ax2 is twin, both are included
ax0.legend(handles0, labels0, loc="best")
# embeddings
ax1 = axis[1]
ax1.set_facecolor(color = (0.1, 0.1, 0.1, 0.5))
ax1.set_ylabel('vector average'.upper(), color=(1, 0.2, 0.5))
ax1.tick_params(axis='y', labelcolor=(1, 0.2, 0.5))
ax1.set_xlim(0, d.steps)
ax1_right = ax1.twinx()
ax1_right.set_ylabel('vector norm'.upper(), color=(0.2, 1.0, 0.5))
ax1_right.tick_params(axis='y', labelcolor=(0.2, 1.0, 0.5))
ax1.xaxis.grid(color='gray', linestyle='dashed')
ax1.yaxis.grid(color='gray', linestyle='dashed')
x = []
avg = []
norm = []
avg_v = [[] for y in range(d.num_vectors_per_token)]
norm_v = [[] for y in range(d.num_vectors_per_token)]
embedding_files = sorted(pathlib.Path(os.path.join(logdir, name, 'embeddings')).glob('*.pt'), key=os.path.getmtime)
for f in embedding_files:
embed = torch.load(f, map_location=torch.device("cpu")) # pylint: disable=no-member
x.append(embed["step"] + 1)
token = list(embed["string_to_token"].keys())[0]
tensors = embed["string_to_param"][token]
val = tensors.detach().numpy()
data = val.flatten()
avg.append(np.average(np.abs(data)))
norm.append(np.linalg.norm(data))
for i in range(val.shape[0]):
avg_v[i].append(np.average(np.abs(val[i])))
norm_v[i].append(np.linalg.norm(val[i]))
ax1.plot(x, avg, color=(1, 0.2, 0.5), linewidth=3, label='all vectors average value')
ax1_right.plot(x, norm, color=(0.2, 1, 0.5), linewidth=3, label='all vectors norm value', linestyle='dashed')
for i in range(d.num_vectors_per_token):
ax1.plot(x, avg_v[i], color= (i / (d.num_vectors_per_token + 1), 0.2, 0.5), linewidth=1, label=f'vector={i} average value')
ax1_right.plot(x, norm_v[i], color= (0.2, i / (d.num_vectors_per_token + 1), 0.5), linewidth=1, label=f'vector={i} norm value', linestyle='dashed')
# axis legend
handles1, labels1 = ax1.get_legend_handles_labels() # because ax2 is twin, both are included
ax1.legend(handles1, labels1, loc="upper left")
ax1_right.legend(loc="upper right")
# create chart and convert to pil
figure.tight_layout()
buf = io.BytesIO()
plt.savefig(buf, format='png')
pltimg = Image.open(buf)
size = (pltimg.size[0], pltimg.size[1] + 240)
image = Image.new('RGB', size = size, color = (206, 100, 0))
font = ImageFont.truetype('DejaVuSansMono', 18)
image.paste(pltimg, box=(0, 240))
buf.close()
plt.close()
# text
textl = f"""NAME: {d.embedding_name.upper()}
IMAGES: {d.num_of_dataset_images}
VECTORS: {d.num_vectors_per_token}
STEPS: {d.steps}
BATCH-SIZE: {d.batch_size}
GRADIENT-STEP: {d.gradient_step}
SAMPLING-METHOD: {d.latent_sampling_method}
MODEL: {d.model_name.upper()}
LEARN-RATE: {d.learn_rate.replace(' ', '')}
"""
minval = f"{round(np.min(loss), 4)} @ {round(step[np.argmin(loss)])}"
maxval = f"{round(np.max(loss), 4)} @ {round(step[np.argmax(loss)])}"
textr = f"""{d.datetime}
LOSS: {round(loss[-1], 4)}
MIN: {minval}
MAX: {maxval}
TREND: {z[0]:.5f}
"""
if len(avg) > 0:
textr += f"VECTOR AVG: {avg[-1]:.3f}\n"
if len(norm) > 0:
textr += f"VECTOR NORM: {norm[-1]:.3f}\n"
ctx = ImageDraw.Draw(image)
ctx.text((8, 8), textl, font = font, fill = (255, 255, 255), spacing = 8)
ctx.text((image.size[0] - 220, 8), textr, font = font, fill = (255, 255, 255))
image.save(img)
if __name__ == "__main__":
log.setLevel(logging.DEBUG)
if len(sys.argv) == 2:
arg = sys.argv[1]
log.debug({ 'args': arg })
plot(os.path.dirname(arg), os.path.basename(arg))
else:
log.error({ 'loss chart': 'specify embedding name'})
cli/modules/train-lossrate.py
-132
View File
@@ -1,132 +0,0 @@
#!/bin/env python
"""
auto-generate learn-rate
"""
import io
import math
import logging
import numpy as np
from PIL import Image, ImageFont, ImageDraw
from matplotlib import pyplot as plt
from util import log, Map
loss_types = ['linear', 'log', 'linalg', 'power']
def gen_steps(steps, step):
return [x for x in range(1, steps + step) if x % step == 0]
def gen_loss_rate(steps: int, step: int, loss_start: float, loss_end: float, loss_type: loss_types, power: int = 3):
def norm(val):
return ((loss_start - loss_end) * val) / val.max() + loss_end
steps_val = gen_steps(steps, step)
if loss_type == 'linear':
loss_val = np.interp(steps_val, [steps_val[0], steps_val[-1]], [loss_start, loss_end])
elif loss_type == 'log':
loss_val = np.logspace(loss_start, 0, num=len(steps_val), base=math.e)
loss_val = norm(loss_val - loss_val.min())
elif loss_type == 'linalg':
loss_val = np.array(steps_val[::-1], dtype='float')
loss_val = norm(loss_val / np.linalg.norm(loss_val))
elif loss_type == 'power':
loss_val = np.array([math.pow(x, power) for x in range(len(steps_val))][::-1])
loss_val = norm(loss_val)
else:
return []
return loss_val
def gen_loss_rate_str(steps: int, step: int, loss_start: float, loss_end: float, loss_type: loss_types, power: int = 3):
steps_val = gen_steps(steps, step)
loss_val = gen_loss_rate(steps, step, loss_start, loss_end, loss_type, power)
loss_rate = [f"{loss_val[i]:.4f}:{steps_val[i]}" for i in range(len(steps_val))]
loss_rate = ', '.join(loss_rate)
log.debug({ 'loss_rate': loss_rate, 'function': loss_type, 'power': power })
return loss_rate
def example_plot(steps: int, step: int, loss_start: float, loss_end: float):
plt.rcParams.update({'font.variant':'small-caps'})
plt.rc('axes', edgecolor='gray')
plt.rc('font', size=10)
plt.rc('font', variant='small-caps')
plt.grid(color='gray', linewidth=1, axis='both', alpha=0.5)
plt.rcParams['figure.figsize'] = [14, 6]
plt.figure(facecolor='black')
loss_rates = []
ax1 = plt.subplot(1, 2, 1)
ax1.set_facecolor('grey')
ax1.set_xlabel('step'.upper(), color='white')
ax1.set_ylabel('loss value'.upper(), color='white')
ax1.xaxis.grid(color='gray', linestyle='dashed')
ax1.tick_params(axis='x', labelcolor='white')
ax1.tick_params(axis='y', labelcolor='white')
ax1.legend(loc="best")
for loss_type in [x for x in loss_types if x != 'power']:
col = (np.random.random(), np.random.random(), np.random.random())
x = gen_steps(steps, step)
y = gen_loss_rate(loss_type = loss_type, steps = steps, step = step, loss_start = loss_start, loss_end = loss_end)
ax1.plot(x, y, label=loss_type, color = col)
loss = gen_loss_rate_str(loss_type = loss_type, steps = steps, step = step, loss_start = loss_start, loss_end = loss_end)
loss_rates.append(f"LOSS {loss} TYPE {loss_type}")
handles, labels = ax1.get_legend_handles_labels()
ax1.legend(handles, labels)
ax2 = plt.subplot(1, 2, 2)
ax2.set_facecolor('grey')
ax2.set_xlabel('step'.upper(), color='white')
ax2.set_ylabel('loss value'.upper(), color='white')
ax2.xaxis.grid(color='gray', linestyle='dashed')
ax2.tick_params(axis='x', labelcolor='white')
ax2.tick_params(axis='y', labelcolor='white')
ax2.legend(loc="best")
for power in range(1, 11):
col = (np.random.random(), np.random.random(), np.random.random())
x = gen_steps(steps, step)
y = gen_loss_rate(loss_type = 'power', power = power, steps = steps, step = step, loss_start = loss_start, loss_end = loss_end)
ax2.plot(x, y, label=f"power={pow}", color = col)
loss = gen_loss_rate_str(loss_type = 'power', power = power, steps = steps, step = step, loss_start = loss_start, loss_end = loss_end)
loss_rates.append(f"LOSS {loss} TYPE power={power}")
handles, labels = ax2.get_legend_handles_labels()
ax2.legend(handles, labels)
plt.tight_layout()
buf = io.BytesIO()
plt.savefig(buf, format='png')
pltimg = Image.open(buf)
size = (pltimg.size[0], pltimg.size[1] + 300)
image = Image.new('RGB', size = size, color = (206, 100, 0))
font = ImageFont.truetype('DejaVuSansMono', 14)
image.paste(pltimg, box=(0, 300))
buf.close()
# text
rates = "\n".join(loss_rates)
text = f"STEPS {steps} STEP {step} LOSS-START {loss_start} LOSS-END {loss_end}\n" + rates
ctx = ImageDraw.Draw(image)
ctx.text((8, 8), text, font = font, fill = (255, 255, 255), spacing = 8)
image.save('lossrate.jpg')
if __name__ == "__main__":
log.setLevel(logging.DEBUG)
arg = Map({
"steps": 500,
"step": 50,
"loss_start": 0.01,
"loss_end": 0.001
})
log.debug({ 'options': arg })
example_plot(**arg)
cli/nvidia-smi.py
Executable
+32
View File
@@ -0,0 +1,32 @@
#!/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)
if __name__ == "__main__":
res = get_nvidia_smi(output='dict')
print(type(res), res)
cli/train/options.py → cli/options.py
View File
cli/train/process.py → cli/process.py
+7 -9
View File
@@ -1,15 +1,13 @@
# pylint: disable=global-statement
import os
import sys
import io
import math
import base64
import pathlib
import numpy as np
import mediapipe as mp
from PIL import Image, ImageOps
from skimage.metrics import structural_similarity as ssim
from scipy.stats import beta
sys.path.append(os.path.join(os.path.dirname(__file__)))
import util
import sdapi
@@ -23,9 +21,9 @@ all_images_by_type = {}
class Result(object):
def __init__(self, type: str, input: str, tag: str = None, requested: list = []):
self.type = type
self.input = input
def __init__(self, typ: str, fn: str, tag: str = None, requested: list = []):
self.type = typ
self.input = fn
self.output = ''
self.basename = ''
self.message = ''
@@ -56,8 +54,8 @@ def detect_dynamicrange(image: Image):
data = np.asarray(image)
image = np.float32(data)
RGB = [0.299, 0.587, 0.114]
height, width = image.shape[:2]
brightness_image = np.sqrt(image[..., 0] ** 2 * RGB[0] + image[..., 1] ** 2 * RGB[1] + image[..., 2] ** 2 * RGB[2])
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)
@@ -264,7 +262,7 @@ def save_image(res: Result, folder: str):
def file(filename: str, folder: str, tag = None, requested = []):
# initialize result dict
res = Result(input = filename, type='unknown', tag=tag, requested = requested)
res = Result(fn = filename, typ='unknown', tag=tag, requested = requested)
# open image
try:
res.image = Image.open(filename)
cli/modules/prompt-ideas.py → cli/prompt-ideas.py
+6 -11
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
generate prompt ideas
model from: <https://huggingface.co/FredZhang7/distilgpt2-stable-diffusion-v2>
@@ -10,18 +10,13 @@ from transformers import GPT2Tokenizer, GPT2LMHeadModel
from util import log
tokenizer = None
model = None
def prompt(text: str, temp: float = 0.9, top: int = 8, penalty: float = 1.2, alpha: float = 0.6, num: int = 5, length: int = 80):
global tokenizer, model # pylint: disable=global-statement
if tokenizer is None:
tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
if model is None:
model = GPT2LMHeadModel.from_pretrained('FredZhang7/distilgpt2-stable-diffusion-v2')
log.info({ 'loading': 'tokenizer' })
tokenizer = GPT2Tokenizer.from_pretrained('distilgpt2')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
input_ids = tokenizer(text, return_tensors='pt').input_ids
log.info({ 'loading': 'model' })
model = GPT2LMHeadModel.from_pretrained('FredZhang7/distilgpt2-stable-diffusion-v2')
output = model.generate(input_ids,
do_sample = True,
temperature = temp,
cli/modules/prompt-promptist.py → cli/prompt-promptist.py
+12 -9
View File
@@ -1,28 +1,31 @@
#!/bin/env python
#!/usr/bin/env python
"""
use microsoft promptist to beautify prompt
- <https://huggingface.co/spaces/microsoft/Promptist>
"""
import sys
from transformers import AutoModelForCausalLM, AutoTokenizer
from util import log
def load_prompter():
def load_model():
log.info({ 'loading': 'model' })
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("microsoft/Promptist") # pylint: disable=redefined-outer-name
return model
def load_tokenizer():
log.info({ 'loading': 'tokenizer' })
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2") # pylint: disable=redefined-outer-name
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
return model, tokenizer
model, tokenizer = load_prompter()
return tokenizer
def beautify(plain_text):
tokenizer = load_tokenizer()
input_ids = tokenizer(plain_text.strip() + " Rephrase:", return_tensors = "pt").input_ids
eos_id = tokenizer.eos_token_id
model = load_model()
outputs = model.generate(input_ids, do_sample = False, max_new_tokens = 75, num_beams = 8, num_return_sequences = 8, eos_token_id = eos_id, pad_token_id = eos_id, length_penalty = -1.0)
output_texts = tokenizer.batch_decode(outputs, skip_special_tokens = True)
texts = []
cli/random/detectmodel.py
-43
View File
@@ -1,43 +0,0 @@
#!/bin/env python
"""
Detect model type
Works for v1 and v2-base (EPS models), both standard inference and inpainting
But looking at model dumps between EPS and V type models, its only about parametrization, there are no differences in actual model (its just weighted differently without any structural difference)
So i don't see easy way to auto-detect if model should be run in `EPS` or `V` mode
Only difference are some calculations in `ldm/models/diffusion/ddpm.py` and by then we already need to know which code-path to trigger
(maaaybe there could be a way by looking at some cherry-picked base tensors min/max range, but I dont see that as reliable)
"""
import os
import sys
import torch
def signature(model):
if model is None:
return None
try:
size = model['state_dict']['model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight'].shape[1]
unet = model['state_dict']['model.diffusion_model.input_blocks.0.0.weight'].shape[1]
except:
return 'unknown'
guess = 'v1' if size == 768 else 'v2' # 768 for v1 and 1024 for v2
guess += '-inference' if unet == 4 else '-inpainting' # inference models have shorter inputs, 4 for inference 9 for inpainting
return guess
def load(file: str):
try:
model = torch.load(file, map_location='cpu')
return model
except Exception as err:
print(f"Error loading {f}: {err}")
if __name__ == "__main__":
sys.argv.pop(0)
for f in sys.argv:
if os.path.isfile(f):
print(f"Model {f} is of type {signature(load(f))}")
else:
print(f"{f} is not a file")
cli/random/versions.py
-33
View File
@@ -1,33 +0,0 @@
#!/bin/env python
"""
print module versions
"""
import importlib
import pkg_resources
modules = [
'diffusers', 'xformers', 'tokenizers', 'accelerate', 'safetensors'
]
def get_torch():
try:
torch = importlib.import_module('torch')
print('torch:', { 'version': torch.__version__ })
print('cuda:', { 'available': torch.cuda.is_available(), 'version': torch.version.cuda, 'arch': torch.cuda.get_arch_list() })
print('device:', { 'name': torch.cuda.get_device_name(torch.cuda.current_device()) })
except Exception as err:
print('torch:', { 'error': err })
def version(name: str):
try:
ver = pkg_resources.get_distribution(name).version
print(f"{name}: {ver}")
except Exception as err:
print(f"{name} error: {err}")
if __name__ == "__main__": # create & train test embedding when used from cli
get_torch()
for module in modules:
version(module)
cli/requirements.txt
+1
View File
@@ -2,3 +2,4 @@ mediapipe
colormap
invisible-watermark
filetype
albumentations
cli/modules/bench.py → cli/run-benchmark.py
+1 -2
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
sd api txt2img benchmark
"""
@@ -87,7 +87,6 @@ async def main():
'vae': opts.sd_vae,
'hypernetwork': opts.sd_hypernetwork,
'sampler': options.sampler_name,
'clip-stop': opts.CLIP_stop_at_last_layers,
'preview': opts.show_progress_every_n_steps
} })
cpu, gpu = memstats()
cli/modules/sdapi.py → cli/sdapi.py
+53 -13
View File
@@ -1,23 +1,59 @@
#!/bin/env python
#!/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 os
import sys
import json
import aiohttp
import ssl
import asyncio
import logging
import aiohttp
import requests
import urllib3
from util import Map, log
sd_url = "http://127.0.0.1:7860" # automatic1111 api url root
sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860") # automatic1111 api url root
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():
sd_username = os.environ.get('SDAPI_USR', None)
sd_password = os.environ.get('SDAPI_PWD', None)
if sd_username is not None and sd_password is not None:
return requests.auth.HTTPBasicAuth(sd_username, sd_password)
return None
def auth():
sd_username = os.environ.get('SDAPI_USR', None)
sd_password = os.environ.get('SDAPI_PWD', None)
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):
@@ -60,7 +96,7 @@ 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) as req:
async with sess.get(url=endpoint, json=json, verify_ssl=False) as req:
res = await result(req)
return res
except Exception as err:
@@ -70,13 +106,12 @@ async def get(endpoint: str, json: dict = None):
def getsync(endpoint: str, json: dict = None):
try:
req = requests.get(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
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):
@@ -86,7 +121,7 @@ async def post(endpoint: str, json: dict = None):
await sess.close()
sess = await session()
try:
async with sess.post(url = endpoint, json = json) as req:
async with sess.post(url=endpoint, json=json, verify_ssl=False) as req:
res = await result(req)
return res
except Exception as err:
@@ -95,7 +130,7 @@ async def post(endpoint: str, json: dict = None):
def postsync(endpoint: str, json: dict = None):
req = requests.post(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
req = requests.post(f'{sd_url}{endpoint}', json=json, verify=False, auth=authsync()) # pylint: disable=missing-timeout
res = resultsync(req)
return res
@@ -136,9 +171,9 @@ def progresssync():
def options():
options = getsync('/sdapi/v1/options')
opts = getsync('/sdapi/v1/options')
flags = getsync('/sdapi/v1/cmd-flags')
return { 'options': options, 'flags': flags }
return { 'options': opts, 'flags': flags }
def shutdown():
@@ -151,7 +186,7 @@ def shutdown():
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)
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__()
@@ -171,6 +206,7 @@ async def session():
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 })
@@ -181,12 +217,16 @@ if __name__ == "__main__":
asyncio.run(interrupt())
if 'progress' in sys.argv:
asyncio.run(progress())
if 'progresssync' in sys.argv:
progresssync()
if '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))
if 'shutdown' in sys.argv:
shutdown()
asyncio.run(close())
asyncio.run(close(), debug=True)
asyncio.run(asyncio.sleep(0.5))
cli/simple-img2txt.py
Executable
+62
View File
@@ -0,0 +1,62 @@
#!/usr/bin/env python
import io
import sys
import base64
import logging
import requests
from PIL import Image
logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
log = logging.getLogger(__name__)
sd_url = "http://127.0.0.1:7860"
options = {
"init_images": [],
"prompt": "city at night",
"negative_prompt": "foggy, blurry",
"steps": 1,
"batch_size": 1,
"n_iter": 1,
"seed": -1,
"sampler_name": "Euler a",
"cfg_scale": 6,
"width": 512,
"height": 512,
"save_images": False,
"send_images": True,
}
def post(endpoint: str, dct: dict = None):
req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300)
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')
values = stream.getvalue()
encoded = base64.b64encode(values).decode()
return encoded
def generate(num: int = 0):
log.info(f'sending generate request: {num+1} {options}')
options['init_images'] = [encode('../html/logo.png')]
data = post('/sdapi/v1/img2img', options)
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)))
log.info(f'received image: {image.size}')
else:
log.warning(f'no images received: {data}')
if __name__ == "__main__":
sys.argv.pop(0)
repeats = int(''.join(sys.argv) or '1')
log.info(f'repeats: {repeats}')
for n in range(repeats):
generate(n)
cli/simple-txt2txt.py
Executable
+50
View File
@@ -0,0 +1,50 @@
#!/usr/bin/env python
import io
import sys
import base64
import logging
import requests
from PIL import Image
logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
log = logging.getLogger(__name__)
sd_url = "http://127.0.0.1:7860"
options = {
"prompt": "city at night",
"negative_prompt": "foggy, blurry",
"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,
}
def post(endpoint: str, dct: dict = None):
req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300)
if req.status_code != 200:
return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
else:
return req.json()
def generate(num: int = 0):
log.info(f'sending generate request: {num+1} {options}')
data = post('/sdapi/v1/txt2img', options)
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)))
log.info(f'received image: {image.size}')
else:
log.warning(f'no images received: {data}')
if __name__ == "__main__":
sys.argv.pop(0)
repeats = int(''.join(sys.argv) or '1')
log.info(f'repeats: {repeats}')
for n in range(repeats):
generate(n)
cli/random/dynamotest.py → cli/torch-compile.py
+12 -9
View File
@@ -1,4 +1,5 @@
#!/bin/env python
#!/usr/bin/env python
# pylint: disable=cell-var-from-loop
"""
Test Torch Dynamo functionality and backends
"""
@@ -12,13 +13,14 @@ from torchvision.models import resnet18
print('torch:', torch.__version__)
try:
import torch._dynamo as dynamo # must be imported explicitly or namespace is not found
# 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
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
@@ -44,8 +46,8 @@ def init_model():
return resnet18().to(torch.float32).cuda()
def eval(mod, inp):
return mod(inp)
def evaluate(mod, val):
return mod(val)
if __name__ == '__main__':
@@ -56,18 +58,19 @@ if __name__ == '__main__':
# repeat test
results = {}
times = []
print('eager initial eval:', timed(lambda: eval(model, inp))[1])
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: eval(model, inp))
_res, time = timed(lambda: evaluate(model, inp))
times.append(time)
results['default'] = np.median(times)
print('dynamo available backends:', dynamo.list_backends())
for backend in dynamo.list_backends():
try:
torch._dynamo.reset() # required before changing backends
eval_dyn = dynamo.optimize(backend)(eval)
# 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])
times = []
for i in range(N_ITERS):
cli/train-lora.py
-274
View File
@@ -1,274 +0,0 @@
#!/bin/env python
"""
Extract approximating LoRA by SVD from two SD models
Based on: <https://github.com/kohya-ss/sd-scripts/blob/main/networks/train_network.py>
Train LoRA with custom preprocessing, tagging and bucketing
Disabled/broken:
- `accelerate` with *dynamo* enabled
- `xformers` due to *faketensors* requirement
- `mem_eff_attn` due to *forwardfunc* mismatch
- 'use_8bit_adam` due to *bitsandbyttes* CUDA errors
"""
import os
import re
import gc
import sys
import json
import time
import shutil
import argparse
import tempfile
import torch
import logging
import importlib
import transformers
from pathlib import Path
from modules.util import log, Map, get_memory
import modules.process
import modules.sdapi
latents = importlib.import_module('modules.lora-latents')
lora_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lora'))
sys.path.append(lora_path)
lycoris_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lycoris'))
sys.path.append(lycoris_path)
from train_network import train
options = 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": 5000,
"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": "mood-magic",
"unet_lr": 0.001,
"use_8bit_adam": False,
"v_parameterization": False,
"v2": False,
"vae": None,
"xformers": False,
})
def mem_stats():
gc.collect()
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 = get_memory()
log.info({ 'memory': { 'ram': mem.ram, 'gpu': mem.gpu } })
if __name__ == '__main__':
parser = argparse.ArgumentParser(description = 'train lora')
parser.add_argument('--model', type=str, default=None, required=False, help='original model to use a base for training, default: active model')
parser.add_argument('--input', '--dataset', type=str, default=None, required=True, help='input folder with training images')
parser.add_argument('--output', '--lora', type=str, default=None, required=True, help='lora name')
parser.add_argument('--tag', type=str, default=None, required=False, help='primary tag')
parser.add_argument('--dir', type=str, default=None, required=False, help='folder containing lora checkpoints')
parser.add_argument('--interim', type=int, default=0, help = 'save interim checkpoints after n epoch')
parser.add_argument('--process', type=str, default='original', required=False, help='list of processing steps: original,face,body,blur,range,upscale,restore')
parser.add_argument('--noprocess', default = False, action='store_true', help = 'skip processing and use existing input data')
parser.add_argument('--notrain', default = False, action='store_true', help = 'just run processing and skip training')
parser.add_argument('--nocaptions', default = False, action='store_true', help = 'skip creating captions and tags')
parser.add_argument('--nolatents', default = False, action='store_true', help = 'skip generating vae latents')
parser.add_argument('--offline', default = False, action='store_true', help = 'do not use webui server for processing')
parser.add_argument('--shutdown', default = False, action='store_true', help = 'shutdown webui server')
parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
parser.add_argument('--steps', type=int, default=4000, required=False, help='training steps, default: %(default)s')
parser.add_argument('--dim', type=int, default=40, required=False, help='network dimension, default: %(default)s')
parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
parser.add_argument('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: half of dim')
parser.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
parser.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
parser.add_argument('--unetlr', type=float, default=1e-04, required=False, help='unet learning rate, default: %(default)s')
parser.add_argument('--textlr', type=float, default=5e-05, required=False, help='text encoder learning rate, default: %(default)s')
parser.add_argument('--dreambooth', default=False, action='store_true', help = "use dreambooth style training")
parser.add_argument('--lycoris', default=False, action='store_true', help = "use lycoris style training")
parser.add_argument('--debug', default=False, action='store_true', help = "enable debug logging")
args = parser.parse_args()
defaults = Map({ 'options': {}, 'flags': {} }) if args.offline else Map(modules.sdapi.options())
if args.debug:
log.setLevel(logging.DEBUG)
log.debug({ 'debug': True })
if args.model is None:
args.model = defaults.options.get('sd_model_checkpoint', None)
args.model = args.model.split(' [')[0] if args.model is not None else None
if args.dir is None:
args.dir = defaults.flags.get('lora_dir', None)
if not os.path.isabs(args.model) and args.dir is not None and not os.path.exists(args.model):
args.model = os.path.abspath(os.path.join(args.dir, os.pardir, 'Stable-diffusion', args.model))
if args.dir is None:
args.dir = os.path.join(args.input, 'lora')
if not os.path.exists(args.model) or not os.path.isfile(args.model):
log.error({ 'lora cannot find model': args.model })
exit(1)
if not os.path.exists(args.input) or not os.path.isdir(args.input):
log.error({ 'lora cannot find training dir': args.input })
exit(1)
if not os.path.exists(args.dir) or not os.path.isdir(args.dir):
log.error({ 'lora cannot find training dir': args.dir })
exit(1)
options.pretrained_model_name_or_path = args.model
options.output_dir = args.dir
options.output_name = args.output
options.max_train_steps = args.steps
options.network_dim = args.dim
options.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
options.gradient_accumulation_steps = args.gradient
options.save_every_n_epochs = args.interim if args.interim > 0 else None
options.learning_rate = args.lr
options.unet_lr = args.unetlr
options.text_encoder_lr = args.textlr
options.train_batch_size = args.batch
log.info({ 'train lora args': vars(options) })
transformers.logging.set_verbosity_error()
mem_stats()
json_file = os.path.join(tempfile.gettempdir(), args.output, args.output + '.json')
base = os.path.join(tempfile.gettempdir(), args.output)
options.train_data_dir = base
res = None
if args.dreambooth:
log.info({ 'using dreambooth style training': True })
options.in_json = None
else:
options.in_json = json_file
for root, _sub_dirs, folder in os.walk(args.input):
files = [os.path.join(root, f) for f in folder]
if not args.noprocess:
# preprocess
processing_options = args.process.split(',')
processing_options = [opt.strip() for opt in re.split(',| ', args.process)]
log.info({ 'processing steps': processing_options })
if os.path.exists(json_file):
os.remove(json_file)
steps = [step for step in processing_options if step in ['face', 'body', 'original']]
for step in steps:
# processing_options = [step for step in processing_options if step not in ['face', 'body', 'original']].append(step)
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']]
if step == 'original':
opts = [step for step in processing_options if step not in ['face', 'body', 'upscale', 'restore', 'blur', 'range']]
log.info({ 'processing step': opts })
concept = step
if concept == 'original' and args.tag is not None:
concept = args.tag.split(',')[0].strip()
dir = os.path.join(base, str(args.repeats) + '_' + concept)
if os.path.exists(dir):
shutil.rmtree(dir, ignore_errors=True)
Path(dir).mkdir(parents=True, exist_ok=True)
for f in files:
try:
res, metadata = modules.process.process_file(f = f, dst = dir, preview = False, offline = args.offline, txt = args.dreambooth, tag = args.tag, opts = opts)
if not args.dreambooth:
with open(json_file, "w") as outfile:
outfile.write(json.dumps(metadata, indent=2))
except ValueError as e:
exit(1)
log.info({ 'processed step': step, 'outputs': res, 'inputs': len(files), 'metadata': json_file, 'path': dir })
modules.process.unload_models()
mem_stats()
dirs = [os.path.join(base, dir) for dir in os.listdir(base) if os.path.isdir(os.path.join(base, dir))]
log.info({ 'input datasets': dirs, 'metadata': json_file })
if not args.nolatents and not args.dreambooth:
# create latents
for dir in dirs:
latents.create_vae_latents(Map({ 'input': dir, 'json': json_file }))
latents.unload_vae()
mem_stats()
else:
log.info({ 'skip processing': len(files), 'metadata': json_file, 'path': dir })
if args.shutdown:
log.info({ 'server shutdown required': True })
modules.sdapi.shutdown()
time.sleep(1)
if args.lycoris:
log.info({ 'using lycoris network': True })
options.network_module = 'lycoris.kohya'
if not args.notrain:
train(options)
mem_stats()
cli/train-ti.py
-591
View File
@@ -1,591 +0,0 @@
#!/bin/env python
# pylint: disable=no-member
"""
simple implementation of training api: `/sdapi/v1/train`
- supports: create embedding, image preprocess, train embedding (with all known parameters)
- does not (yet) support: create hyper-network, train hyper-network
- compatible with progress api: `/sdapi/v1/progress`
- if interrupted, auto-continues from last known step
- create and preprocess executed as sync jobs
- train is executed as async job with progress monitoring
"""
import argparse
import asyncio
import logging
import math
import os
import sys
import time
import importlib
from pathlib import Path, PurePath
import filetype
from PIL import Image
sys.path.append(os.path.join(os.path.dirname(__file__), 'modules'))
from modules.util import Map, log, set_logfile
from modules.sdapi import close, get, interrupt, post, progress, session
from modules.process import process_images
from modules.grid import grid
create_preview = importlib.import_module('modules.preview-embeddings').create_preview
plot = importlib.import_module('modules.train-losschart').plot
extract = importlib.import_module('modules.video-extract').extract
gen_loss_rate_str = importlib.import_module('modules.train-lossrate').gen_loss_rate_str
images = []
args = {}
options = None
cmdflags = None
args = Map({
"training_model": "sd-v15-runwayml.ckpt",
"extract_video": {
"rate": 0,
"fps": 5,
"vstart": 0,
"vend": 0
},
"create_embedding": {
"name": "test",
"num_vectors_per_token": 1,
"overwrite_old": False,
"init_text": "*"
},
"preprocess": {
"id_task": 0,
"process_src": "",
"process_dst": "",
"process_width": 512,
"process_height": 512,
"process_flip": False,
"process_split": False,
"process_caption": True,
"process_caption_deepbooru": False,
"preprocess_txt_action": "ignore",
"process_focal_crop": True,
"process_focal_crop_face_weight": 0.9,
"process_focal_crop_entropy_weight": 0.3,
"process_focal_crop_edges_weight": 0.5,
"process_focal_crop_debug": False,
"split_threshold": 0.5,
"overlap_ratio": 0.2,
"process_multicrop": None,
"process_multicrop_mindim": None,
"process_multicrop_maxdim": None,
"process_multicrop_minarea": None,
"process_multicrop_maxarea": None,
"process_multicrop_objective": None,
"process_multicrop_threshold": None,
},
"train_embedding": {
"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": "once",
"create_image_every": -1,
"save_embedding_every": -1,
"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,
},
})
async def plotloss(params):
logdir = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log'))
try:
plot(logdir, params.name)
except Exception as err:
log.warning({ 'loss chart error': err })
async def captions(docs: list):
exclude = ['a', 'in', 'on', 'out', 'at', 'the', 'and', 'with', 'next', 'to', 'it', 'for', 'of', 'into', 'that']
d = dict()
for f in docs:
text = open(f, 'r', encoding='utf-8')
for line in text:
line = line.strip()
line = line.lower()
words = line.split(" ")
for word in words:
if word in exclude:
continue
d[word] = d[word] + 1 if word in d else 1
pairs = ((value, key) for (key,value) in d.items())
sort = sorted(pairs, reverse = True)
if len(sort) > 10:
del sort[10:]
d = {k: v for v, k in sort}
log.info({ 'top captions': d })
async def preprocess_cleanup(params):
log.info({ 'preprocess cleanup': params.dst })
for f in Path(params.dst).glob('*.png'):
f.unlink()
for f in Path(params.dst).glob('*.jpg'):
f.unlink()
for f in Path(params.dst).glob('*.txt'):
f.unlink()
try:
if os.path.isdir(params.dst):
Path(params.dst).rmdir()
except Exception as err:
log.warning({ 'preprocess cleanup': params.dst, 'error': err })
async def preprocess_builtin(params):
global images # pylint: disable=global-statement
log.debug({ 'preprocess start' })
files = [os.path.join(params.src, f) for f in os.listdir(params.src) if os.path.isfile(os.path.join(params.src, f))]
candidates = [f for f in files if filetype.is_image(f)]
not_images = [f for f in files if (not filetype.is_image(f) and not f.endswith('.txt'))]
images = []
low_res = []
for f in candidates:
img = Image.open(f)
mp = (img.size[0] * img.size[1]) / 1024 / 1024
if mp < 1 or img.size[0] < 512 or img.size[1] < 512:
low_res.append(f)
os.rename(f, f + '.skip')
else:
images.append(f)
log.debug({ 'preprocess skipping': not_images })
log.debug({ 'preprocess low res': low_res })
args.preprocess.process_src = params.src
args.preprocess.process_dst = params.dst
log.debug({ 'preprocess args': args.preprocess })
_res = await post('/sdapi/v1/preprocess', json = args.preprocess)
processed = [os.path.join(params.dst, f) for f in os.listdir(params.dst) if os.path.isfile(os.path.join(params.dst, f))]
processed_imgs = [f for f in processed if f.endswith('.png')]
processed_docs = [f for f in processed if f.endswith('.txt')]
log.info({ 'preprocess': {
'source': params.src,
'destination': params.dst,
'files': len(files),
'images': len(images),
'processed': len(processed_imgs),
'captions': len(processed_docs),
'skipped': len(not_images),
'low-res': len(low_res) }
})
if len(processed_docs) > 0:
await captions(processed_docs)
return len(processed_imgs)
async def preprocess(params):
global images # pylint: disable=global-statement
res = 0
if os.path.isfile(params.src):
if not filetype.is_video(params.src):
kind = filetype.guess(params.src)
log.error({ 'preprocess error': { 'not a valid movie file': params.src, 'guess': kind } })
else:
extract_dst = os.path.join(params.dst, 'extract')
log.debug({ 'preprocess args': args.extract_video })
images = extract(params.src, extract_dst, rate = args.extract_video.rate, fps = args.extract_video.fps, start = args.extract_video.vstart, end = args.extract_video.vend) # extract keyframes from movie
if images > 0:
params.src = extract_dst
res = await preprocess(params) # call again but now with keyframes
else:
log.error({ 'preprocess video extract': 'no images' })
elif os.path.isdir(params.src):
if params.overwrite:
await preprocess_cleanup(params)
elif os.path.isdir(params.dst):
log.error({ 'preprocess output folder already exists': params.dst })
return 0
if params.preprocess == 'builtin':
res = await preprocess_builtin(params)
i = [os.path.join(params.dst, f) for f in os.listdir(params.dst) if os.path.isfile(os.path.join(params.dst, f)) and filetype.is_image(os.path.join(params.dst, f))]
images = [Image.open(img) for img in i]
res = len(images)
elif params.preprocess == 'custom':
t0 = time.perf_counter()
args.preprocess.process_src = params.src
args.preprocess.process_dst = params.dst
process_images(src = params.src, dst = params.dst)
i = [os.path.join(params.dst, f) for f in os.listdir(params.dst) if os.path.isfile(os.path.join(params.dst, f)) and filetype.is_image(os.path.join(params.dst, f))]
images = [Image.open(img) for img in i]
t1 = time.perf_counter()
log.info({ 'preprocess': { 'source': params.src, 'destination': params.dst, 'images': len(images), 'time': round(t1 - t0, 2) } })
res = len(images)
else:
args.preprocess.process_dst = params.src
i = [os.path.join(params.src, f) for f in os.listdir(params.src) if os.path.isfile(os.path.join(params.src, f)) and filetype.is_image(os.path.join(params.src, f))]
images = [Image.open(img) for img in i]
res = len(images)
else:
log.error({ 'preprocess error': { 'not a valid input': params.src } })
if len(images) > 0:
img = grid(images, labels = None, width = 2048, height = 2048, border = 8, square = True)
logdir = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log'))
Path(logdir).mkdir(parents = True, exist_ok = True)
fn = os.path.join(logdir, params.name + '.inputs.jpg')
img.save(fn)
log.info({ 'preprocess input grid': fn })
return res
async def check(params):
global options # pylint: disable=global-statement
options = await get('/sdapi/v1/options')
global cmdflags # pylint: disable=global-statement
cmdflags = await get('/sdapi/v1/cmd-flags')
logdir = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log', params.name))
logfile = os.path.abspath(os.path.join(cmdflags['embeddings_dir'], '../train/log', params.name + '.train.log'))
set_logfile(logfile)
log.info({ 'checking server options' })
options['training_image_repeats_per_epoch'] = 1
if params.skipmodel:
log.info({ 'using model': options['sd_model_checkpoint'] })
else:
log.debug({ 'check model': args.training_model })
if len(args.training_model) > 0 and not options['sd_model_checkpoint'].startswith(args.training_model):
models = await get('/sdapi/v1/sd-models')
models = [obj["title"] for obj in models]
found = [i for i in models if i.startswith(args.training_model)]
if len(found) == 0:
log.error({ 'model not found': args.training_model, 'available': models })
exit()
else:
log.warning({ 'switching model': found[0] })
options['sd_model_checkpoint'] = found[0]
log.debug({ 'check embedding': params.name })
lst = os.path.join(cmdflags['embeddings_dir'])
log.debug({ 'embeddings folder': lst })
path = Path(cmdflags['embeddings_dir']).glob(f'{params.name}.pt*')
matches = [f for f in path]
for match in matches:
if params.overwrite:
log.info({ 'delete embedding': match.name })
os.remove(os.path.join(cmdflags['embeddings_dir'], match.name))
else:
log.error({ 'embedding exists': match.name })
await close()
exit()
f = os.path.join(logdir, 'train.csv')
if os.path.isfile(f):
if params.overwrite:
log.info({ 'delete training log': f })
os.remove(os.path.join(logdir, 'train.csv'))
else:
log.warning({ 'training log exists': f })
f = os.path.join(logdir, '..', params.name, '.png')
if os.path.isfile(f):
if params.overwrite:
log.info({ 'delete training graph': f })
os.remove(f)
log.debug({ 'options': 'update' })
await post('/sdapi/v1/options', options)
return
async def create(params):
log.debug({ 'create start' })
if not os.path.isdir(args.preprocess.process_dst):
log.error({ 'train source not found': args.preprocess.process_dst })
exit()
if params.vectors == -1: # dynamically determine number of vectors depending on number of input images
if len(images) <= 20:
vectors = 2
elif len(images) <= 100:
vectors = 4
else:
vectors = 6
else:
vectors = params.vectors
if os.path.exists(params.name) and os.path.isfile(params.name):
log.info({ 'deleting existing embedding': { 'name': params.name } })
os.remove(params.name)
args.create_embedding.name = params.name
words = params.init.split(',')
if len(words) > vectors:
params.init = ','.join(words[:vectors])
log.warning({ 'create embedding init words cut': params.init })
args.create_embedding.init_text = params.init
args.create_embedding.num_vectors_per_token = vectors
log.debug({ 'create args': args.create_embedding })
res = await post('/sdapi/v1/create/embedding', args.create_embedding)
if 'info' in res:
log.info({ 'create embedding': { 'name': params.name, 'init': params.init, 'vectors': vectors, 'message': res.info } })
else:
log.error({ 'create failed:', res })
return None
log.debug({ 'create end' })
return params.name
async def train(params):
log.debug({ 'train start' })
args.train_embedding.embedding_name = params.name
imgs = [f for f in os.listdir(args.preprocess.process_dst) if os.path.isfile(os.path.join(args.preprocess.process_dst, f)) and filetype.is_image(os.path.join(args.preprocess.process_dst, f))]
args.train_embedding.data_root = args.preprocess.process_dst
if len(imgs) == 0:
log.error({ 'train no input images in folder': args.preprocess.process_dst })
return
if params.grad == -1:
args.train_embedding.gradient_step = len(imgs) // args.train_embedding.batch_size
divisor = args.train_embedding.gradient_step // 60
args.train_embedding.gradient_step = args.train_embedding.gradient_step // (1 + divisor)
log.info({ 'dynamic gradient step': args.train_embedding.gradient_step })
if params.steps == -1:
args.train_embedding.steps = params.maxsteps // args.train_embedding.gradient_step
log.info({ 'dynamic steps': args.train_embedding.steps, 'estimated total steps': args.train_embedding.steps * args.train_embedding.gradient_step * args.train_embedding.batch_size })
epoch_size = args.train_embedding.batch_size * args.train_embedding.gradient_step
if args.train_embedding.create_image_every == -1:
args.train_embedding.create_image_every = args.train_embedding.steps // 10
if args.train_embedding.save_embedding_every == -1:
args.train_embedding.save_embedding_every = args.train_embedding.steps // 10
if args.train_embedding.learn_rate == -1:
loss_args = {
"steps": args.train_embedding.steps,
"step": args.train_embedding.create_image_every,
"loss_start": params.rstart,
"loss_end": params.rend,
"loss_type": 'power',
"power": params.rdescend
}
args.train_embedding.learn_rate = gen_loss_rate_str(**loss_args)
log.info({ 'dynamic learn-rate': loss_args })
log.debug({ 'learn rate': args.train_embedding.learn_rate, 'params': loss_args })
log.info({ 'train embedding': {
'name': params.name,
'source': args.preprocess.process_dst,
'images': len(imgs),
'steps': args.train_embedding.steps,
'batch': args.train_embedding.batch_size,
'gradient-step': args.train_embedding.gradient_step,
'sampling': args.train_embedding.latent_sampling_method,
'epoch-size': epoch_size }
})
log.info({ 'learn rate': args.train_embedding.learn_rate })
log.debug({ 'train args': args.train_embedding })
t0 = time.time()
res = await post('/sdapi/v1/train/embedding', args.train_embedding)
log.info({ 'train result': res })
t1 = time.time()
log.info({ 'train embedding finished': { 'name': params.name, 'time': round(t1 - t0) } })
log.debug({ 'train end': res.info if 'info' in res else res })
return
async def pipeline(params):
log.debug({ 'pipeline start' })
# interrupt
await interrupt()
# preprocess
num = await preprocess(params)
if num == 0:
log.warning({ 'preprocess': 'no resulting images'})
return
# create embedding
name = await create(params)
if not params.name in name:
log.error({ 'create embedding failed': name })
return
# train embedding
await train(params)
await plotloss(params)
# create_preview(params.name, params.init)
log.debug({ 'pipeline end' })
return
async def monitor(params):
step = 0
t0 = time.perf_counter()
t1 = time.perf_counter()
log.info({' starting monitor': t0 })
finished = 0
while True:
await asyncio.sleep(params.monitor)
res = await progress()
if not 'state' in res:
log.info({ 'monitor disconnected': res })
break
if (res.state.job_count == params.steps and res.state.job_no >= res.state.job_count) or (res.eta_relative < 0) or (res.interrupted) or (res.state.job_count == 0): # need exit case if interrupted or failed
if res.interrupted:
log.info({ 'monitor interrupted': { 'embedding': params.name } })
break # exit for monitor job
else:
finished += 1
if finished >= 2: # do it more than once since preprocessing job can finish just in time for monitor to finish
log.info({ 'monitor finished': { 'embedding': params.name } })
break
else:
if res.state.job_no == 0:
step = 0
t0 = time.perf_counter()
t1 = time.perf_counter()
try:
if 'Loss:' in res.textinfo:
text = res.textinfo.split('<br/>')[0].split()
loss = float(text[-1])
else:
loss = -1
except:
loss = -1
if math.isnan(loss):
log.error({ 'monitor': { 'progress': round(100 * res.progress), 'embedding': params.name, 'eta': round(res.eta_relative), 'step': res.state.job_no, 'steps': res.state.job_count, 'loss': 'nan' } })
await interrupt()
else:
elapsed = t1 - t0
log.info({ 'monitor': {
'job': res.state.job,
'progress': round(100 * res.progress),
'embedding': params.name,
'epoch': (1 + res.state.job_no // len(images)) if len(images) > 0 else 'n/a',
'step': res.state.job_no,
'steps': res.state.job_count,
'loss': loss if loss > -1 else 'n/a',
'total': round(1.0 * elapsed * res.state.job_count / res.state.job_no) if res.state.job_no > 0 and t1 != t0 else 'n/a',
'elapsed': round(elapsed),
'remaining': round(res.eta_relative),
'it/s': round((res.state.job_no - step) / (time.perf_counter() - t1), 2) }
})
if step % 10 == 0:
await plotloss(params)
step = res.state.job_no
t1 = time.perf_counter()
return
async def main():
parser = argparse.ArgumentParser(description="sd train ti pipeline")
parser.add_argument("--name", type = str, required = True, help = "embedding name, set to auto to use src folder name")
parser.add_argument("--src", type = str, required = True, help = "source image folder or movie file")
parser.add_argument("--init", type = str, default = "person", required = False, help = "initialization class, default: %(default)s")
parser.add_argument("--dst", type = str, default = "/tmp", required = False, help = "destination image folder for processed images, default: %(default)s")
parser.add_argument("--steps", type = int, default = -1, required = False, help = "training steps, default: %(default)s")
parser.add_argument("--maxsteps", type = int, default = 5000, required = False, help = "max training steps used when dynamic gradient is active, default: %(default)s")
parser.add_argument("--vectors", type = int, default = -1, required = False, help = "number of vectors per token, default: dynamic based on number of input images")
parser.add_argument("--batch", type = int, default = 1, required = False, help = "batch size, default: %(default)s")
parser.add_argument("--rate", type = str, default = "", required = False, help = "learn rate, default: dynamic")
parser.add_argument("--rstart", type = float, default = 0.02, required = False, help = "starting learn rate if using dynamic rate, default: %(default)s")
parser.add_argument("--rend", type = float, default = 0.0005, required = False, help = "ending learn rate if using dynamic rate, default: %(default)s")
parser.add_argument("--rdescend", type = float, default = 2, required = False, help = "learn rate descend power when using dynamic rate, default: %(default)s")
parser.add_argument("--grad", type = int, default = -1, required = False, help = "accumulate gradient over n images, default: : %(default)s")
parser.add_argument("--type", type = str, default = 'subject', required = False, help = "training type: subject/style/unknown, default: %(default)s")
parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing embedding, default: %(default)s")
parser.add_argument("--vstart", type = float, default = 0, required = False, help = "if processing video skip first n seconds, default: %(default)s")
parser.add_argument("--vend", type = float, default = 0, required = False, help = "if processing video skip last n seconds, default: %(default)s")
parser.add_argument('--skipcaption', default = False, action='store_true', help = "do not auto-generate captions, default: %(default)s")
parser.add_argument('--skipmodel', default = False, action='store_true', help = "skip model validation and switch, default: %(default)s")
parser.add_argument('--preprocess', type = str, choices=['builtin', 'custom', 'none'], default = 'custom', help = "preprocessing type, default: %(default)s")
parser.add_argument('--nocleanup', default = False, action='store_true', help = "skip cleanup after completion, default: %(default)s")
parser.add_argument("--monitor", type = int, default = 30, required = False, help = "progress monitor frequency, default: : %(default)s")
parser.add_argument('--debug', default = False, action='store_true', help = "print extra debug information, default: %(default)s")
params = parser.parse_args()
if params.debug:
log.setLevel(logging.DEBUG)
log.debug({ 'debug': True })
log.debug({ 'args': params.__dict__ })
home = Path(sys.argv[0]).parent
global args # pylint: disable=global-statement
if params.vstart > 0:
args.extract_video.vstart = params.vstart
if params.vend > 0:
args.extract_video.vend = params.vend
if params.steps > -1:
args.train_embedding.steps = params.steps
if params.batch > -1:
args.train_embedding.batch_size = params.batch
if params.rate != '':
args.train_embedding.learn_rate = params.rate
if params.grad > -1:
args.train_embedding.gradient_step = params.grad
if params.type == 'subject':
if params.skipcaption:
args.train_embedding.template_filename = 'subject.txt'
args.preprocess.process_caption = False
else:
args.train_embedding.template_filename = 'subject_filewords.txt'
elif params.type == 'style':
if params.skipcaption:
args.train_embedding.template_filename = 'style.txt'
args.preprocess.process_caption = False
else:
args.train_embedding.template_filename = 'style_filewords.txt'
else:
if params.skipcaption:
args.train_embedding.template_filename = 'unknown.txt'
args.preprocess.process_caption = False
else:
args.train_embedding.template_filename = 'unknown_filewords.txt'
if params.name == 'auto':
params.name = PurePath(params.src).name
log.info({ 'training name': params.name })
if params.dst == "/tmp":
params.dst = os.path.join("/tmp/train", params.name)
log.debug({ 'args': params.__dict__ })
params.src = os.path.abspath(params.src)
params.dst = os.path.abspath(params.dst)
try:
await session()
await check(params)
a = asyncio.create_task(pipeline(params))
b = asyncio.create_task(monitor(params))
await asyncio.gather(a, b) # wait for both pipeline and monitor to finish
except Exception as e:
log.error({ 'exception': e })
finally:
if not params.nocleanup:
await preprocess_cleanup(params)
await close()
return
if __name__ == "__main__":
log.info({ 'train textual inversion' })
try:
asyncio.run(main())
except KeyboardInterrupt:
log.warning({ 'interrupted': 'keyboard request' })
# asyncio.run(interrupt())
cli/train.py
Executable
+395
View File
@@ -0,0 +1,395 @@
#!/usr/bin/env python
# system imports
import os
import re
import gc
import sys
import json
import time
import shutil
import pathlib
import asyncio
import logging
import tempfile
import argparse
import warnings
warnings.filterwarnings(action="ignore", category=DeprecationWarning)
warnings.filterwarnings(action="ignore", category=UserWarning)
warnings.filterwarnings(action="ignore", category=FutureWarning)
# 3rd party imports
import filetype
import torch
from tqdm.rich import tqdm
# local imports
import util
import sdapi
import options
import process
import latents
# 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')
# methods
def setup_logging(clean=False):
try:
if clean and os.path.isfile(log_file):
os.remove(log_file)
time.sleep(0.1) # prevent race condition
except:
pass
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()
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_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_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('--output', 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=10, 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_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('--debug', default = False, action='store_true', help = "enable debug level logging, default: %(default)s")
args = parser.parse_args()
def prepare_server():
try:
server_status = util.Map(sdapi.progresssync())
server_state = server_status['state']
except:
log.error(f'server error: {server_status}')
exit(1)
if server_state['job_count'] > 0:
log.error(f'server not idle: {server_state}')
exit(1)
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)
server_options.options.sd_model_checkpoint = args.model
sdapi.postsync('/sdapi/v1/options', server_options.options)
else:
args.model = server_options.options.sd_model_checkpoint.split(' [')[0]
args.lora_dir = server_options.options.lora_dir
args.lyco_dir = server_options.options.lyco_dir
args.ckpt_dir = server_options.options.ckpt_dir
args.embeddings_dir = server_options.options.embeddings_dir
if not os.path.isfile(args.model):
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))
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.ckpt_dir) or not os.path.isdir(args.ckpt_dir):
log.error(f'cannot find models folder: {args.ckpt_dir}')
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.output != '':
args.process_dir = args.output
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():
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')
import train_network
train_network.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.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.lyco_dir
options.lora.network_module = 'lycoris.kohya'
options.lora.in_json = os.path.join(args.process_dir, args.name + '.json')
# lora specific
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.netwoork_args = []
if args.algo is not None:
options.lora.netwoork_args.append(f'algo={args.algo}')
if args.args is not None:
for net_arg in args.args:
options.lora.netwoork_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
# 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():
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()
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 = {}
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}")
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
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__':
log.info('SD.Next train script')
parse_args()
setup_logging()
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 as e:
log.error('interrupt requested')
sdapi.interrupt()
mem_stats()
log.info('done')
cli/train/console.py
View File
cli/train/sdapi.py
-111
View File
@@ -1,111 +0,0 @@
import sys
import json
import aiohttp
import asyncio
import requests
from util import Map
sd_url = "http://127.0.0.1:7860" # automatic1111 api url root
use_session = True
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
async def result(req):
if req.status != 200:
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 type(json) == list:
res = json
elif json is None:
res = {}
else:
res = Map(json)
return res
def resultsync(req: requests.Response):
if req.status_code != 200:
return Map({ 'error': req.status_code, 'reason': req.reason, 'url': req.url })
else:
json = req.json()
if type(json) == list:
res = json
elif json is None:
res = {}
else:
res = Map(json)
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()
async with sess.get(url = endpoint, json = json) as req:
res = await result(req)
return res
def getsync(endpoint: str, json: dict = None):
req = requests.get(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
res = resultsync(req)
return res
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()
async with sess.post(url = endpoint, json = json) as req:
res = await result(req)
return res
def postsync(endpoint: str, json: dict = None):
req = requests.post(f'{sd_url}{endpoint}', json = json) # pylint: disable=missing-timeout
res = resultsync(req)
return res
def interrupt():
res = getsync('/sdapi/v1/progress?skip_current_image=true')
if 'state' in res and res.state.job_count > 0:
res = postsync('/sdapi/v1/interrupt')
return res
else:
return { 'interrupt': 'idle' }
def progress():
res = getsync('/sdapi/v1/progress?skip_current_image=true')
return res
def options():
options = getsync('/sdapi/v1/options')
flags = getsync('/sdapi/v1/cmd-flags')
return { 'options': options, 'flags': flags }
def shutdown():
postsync('/sdapi/v1/shutdown')
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)
return sess
async def close():
if sess is not None:
await asyncio.sleep(0)
await sess.__aexit__(None, None, None)
cli/train/train.py
-322
View File
@@ -1,322 +0,0 @@
#!/bin/env python
# system imports
import os
import re
import gc
import sys
import json
import shutil
import pathlib
import asyncio
import tempfile
import argparse
# 3rd party imports
import filetype
from tqdm.rich import tqdm
# local imports
import util
import sdapi
import process
import latents
import options
# console handler
from rich import print # pylint: disable=redefined-builtin
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)
import torch, accelerate, diffusers, requests, urllib3, http
traceback_install(console=console, extra_lines=1, width=console.width, word_wrap=False, indent_guides=False, suppress=[torch,accelerate,diffusers,asyncio,http,urllib3,requests])
# lora imports
lora_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lora'))
sys.path.append(lora_path)
lycoris_path = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'modules', 'lycoris'))
sys.path.append(lycoris_path)
import train_network
# globals
args = None
valid_steps = ['original', 'face', 'body', 'blur', 'range', 'upscale', 'restore', 'interrogate', 'resize', 'square', 'segment']
# methods
def mem_stats():
gc.collect()
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'] }
console.log(f"memory cpu: {mem.ram} gpu current: {mem.gpu} gpu peak: {peak}")
def parse_args():
global args
parser = argparse.ArgumentParser(description = 'train lora')
# basic section
parser.add_argument('--output', '--name', type=str, default=None, required=True, help='output filename')
parser.add_argument('--type', type=str, choices=['embedding', 'lora', 'lycoris', 'dreambooth'], default=None, required=True, help='training type')
parser.add_argument('--tag', type=str, default='person', required=False, help='primary tag, default: %(default)s')
parser.add_argument('--process', type=str, default='original,interrogate,resize,square', required=False, help=f'list of possible processing steps: {valid_steps}, default: %(default)s')
parser.add_argument('--dir', type=str, default='', required=False, help='where to store processed images, default is system temp/train')
parser.add_argument('--input', '--dataset', type=str, default=None, required=True, help='input folder with training images')
parser.add_argument('--overwrite', default = False, action='store_true', help = "overwrite existing training, default: %(default)s")
# global params
parser.add_argument('--gradient', type=int, default=1, required=False, help='gradient accumulation steps, default: %(default)s')
parser.add_argument('--steps', type=int, default=2500, required=False, help='training steps, default: %(default)s')
parser.add_argument('--batch', type=int, default=1, required=False, help='batch size, default: %(default)s')
parser.add_argument('--lr', type=float, default=1e-04, required=False, help='model learning rate, default: %(default)s')
parser.add_argument('--dim', '--vectors', type=int, default=40, required=False, help='network dimension, default: %(default)s')
# lora params
parser.add_argument('--repeats', type=int, default=10, required=False, help='number of repeats per image, default: %(default)s')
parser.add_argument('--alpha', type=float, default=0, required=False, help='alpha for weights scaling, default: half of dim')
args = parser.parse_args()
def prepare_server():
try:
server_status = util.Map(sdapi.progress())
server_state = server_status['state']
except:
console.log('server error:', server_status)
exit(1)
if server_state['job_count'] > 0:
console.log('server not idle:', server_state)
exit(1)
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)
console.log(f'updated server options')
def verify_args():
global args
server_options = util.Map(sdapi.options())
args.model = server_options.options['sd_model_checkpoint'].split(' [')[0]
args.lora_dir = server_options.flags['lora_dir']
if not os.path.isabs(args.model) and not os.path.exists(args.model):
args.model = os.path.abspath(os.path.join(args.lora_dir, os.pardir, 'Stable-diffusion', args.model))
if not os.path.exists(args.model) or not os.path.isfile(args.model):
console.log('cannot find model:', args.model)
exit(1)
if not os.path.exists(args.input) or not os.path.isdir(args.input):
console.log('cannot find training folder:', args.input)
exit(1)
if not os.path.exists(args.lora_dir) or not os.path.isdir(args.lora_dir):
console.log('cannot find lora folder:', args.dir)
exit(1)
if args.dir != '':
args.process_dir = args.dir
else:
args.process_dir = os.path.join(tempfile.gettempdir(), 'train', args.output)
console.log(f'args: {vars(args)}')
async def training_loop():
async def async_train():
res = await sdapi.post('/sdapi/v1/train/embedding', options.embedding)
console.log(f'train embedding result: {res}')
async def async_monitor():
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():
console.log(f'{args.type} options: {options.embedding}')
create_options = util.Map({
"name": args.output,
"num_vectors_per_token": args.dim,
"overwrite_old": False,
"init_text": args.tag,
})
server_options = util.Map(sdapi.options())
fn = os.path.join(server_options.flags.embeddings_dir, args.output) + '.pt'
if os.path.exists(fn) and args.overwrite:
console.log(f'delete existing embedding {fn}')
os.remove(fn)
else:
console.log(f'embedding exists {fn}')
return
console.log(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
console.log(res.info)
elif 'info' in res: # no error
asyncio.run(training_loop())
else: # unknown error
console.log(f'create embedding error {res}')
def train_lora():
fn = os.path.join(args.lora_dir, args.output)
for ext in ['.ckpt', '.pt', '.safetensors']:
if os.path.exists(fn + ext):
if args.overwrite:
console.log(f'delete existing lora: {fn + ext}')
os.remove(fn + ext)
else:
console.log(f'lora exists: {fn + ext}')
return
console.log(f'{args.type} options: {options.lora}')
train_network.train(options.lora)
def prepare_options():
# lora specific
options.lora.pretrained_model_name_or_path = args.model
options.lora.output_dir = args.lora_dir
options.lora.output_name = args.output
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.gradient_accumulation_steps = args.gradient
options.lora.learning_rate = args.lr
options.lora.train_batch_size = args.batch
options.lora.network_alpha = args.dim // 2 if args.alpha == 0 else args.alpha
options.lora.train_data_dir = args.process_dir
if args.type == 'lycoris':
console.log('train using lycoris network')
options.lora.network_module = 'lycoris.kohya'
options.lora.in_json = os.path.join(args.process_dir, args.output + '.json')
if args.type == 'dreambooth':
console.log('train using dreambooth style training')
options.lora.in_json = None
if args.type == 'lora':
console.log('train using lora style training')
options.lora.in_json = os.path.join(args.process_dir, args.output + '.json')
if args.type == 'embedding':
console.log('train embedding')
options.lora.in_json = None
pass
# embedding specific
options.embedding.embedding_name = args.output
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():
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)]
console.log(f'processing steps: {processing_options}')
for step in processing_options:
if step not in valid_steps:
console.log(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))]
console.log(f'processing input images: {len(files)}')
if os.path.exists(args.process_dir):
console.log('removing existing processed folder:', args.process_dir)
shutil.rmtree(args.process_dir, ignore_errors=True)
steps = [step for step in processing_options if step in ['face', 'body', 'original']]
process.reset()
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
console.log(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', 'lycoris', 'dreambooth']:
dir = os.path.join(args.process_dir, str(args.repeats) + '_' + concept) # separate concepts per folder
if args.type in ['embedding']:
dir = os.path.join(args.process_dir) # everything into same folder
console.log('processing concept:', concept)
console.log('processing output folder:', dir)
pathlib.Path(dir).mkdir(parents=True, exist_ok=True)
results = {}
for f in files:
res = process.file(filename = f, folder = dir, 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") as outfile:
outfile.write(res.caption)
console.log(f"processing {'saved' if res.image is not None else 'skipped'}: {f} => {res.output} {res.ops} {res.message}")
dirs = [os.path.join(args.process_dir, dir) for dir in os.listdir(args.process_dir) if os.path.isdir(os.path.join(args.process_dir, dir))]
console.log(f'input datasets {dirs}')
if options.lora.in_json is not None:
with open(options.lora.in_json, "w") as outfile: # write json at the end only
outfile.write(json.dumps(metadata, indent=2))
for dir in dirs: # create latents
latents.create_vae_latents(util.Map({ 'input': dir, 'json': options.lora.in_json }))
latents.unload_vae()
r = { 'inputs': len(files), 'outputs': results, 'metadata': options.lora.in_json }
console.log(f'processing steps result: {r}')
if args.gradient < 0:
console.log(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__':
console.log('train script for stable diffusion')
parse_args()
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 == 'lycoris' or args.type == 'dreambooth':
train_lora()
except KeyboardInterrupt as e:
console.log('interrupt requested')
sdapi.interrupt()
mem_stats()
console.log('done')
cli/train/util.py
-85
View File
@@ -1,85 +0,0 @@
#!/bin/env python
import os
import transformers
transformers.logging.set_verbosity_error()
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']) }
warnings = { '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': warnings,
})
except:
pass
return Map(mem)
class Map(dict):
__slots__ = ('__dict__')
def __init__(self, *args, **kwargs):
super(Map, self).__init__(*args, **kwargs)
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)
self.__dict__.update({key: value})
def __delattr__(self, item):
self.__delitem__(item)
def __delitem__(self, key):
super(Map, self).__delitem__(key)
del self.__dict__[key]
if __name__ == "__main__":
pass
cli/modules/util.py → cli/util.py
+9 -5
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
generic helper methods
"""
@@ -6,9 +6,13 @@ 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")
@@ -52,22 +56,22 @@ def get_memory():
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']) }
warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] }
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': warnings,
'events': events,
})
except:
pass
return Map(mem)
class Map(dict):
__slots__ = ('__dict__')
class Map(dict): # pylint: disable=C0205
__slots__ = ('__dict__') # pylint: disable=C0325
def __init__(self, *args, **kwargs):
super(Map, self).__init__(*args, **kwargs)
for arg in args:
cli/modules/video-extract.py → cli/video-extract.py
+3 -3
View File
@@ -1,4 +1,4 @@
#!/bin/env python
#!/usr/bin/env python
"""
use ffmpeg for animation processing
"""
@@ -16,8 +16,8 @@ def probe(src: str):
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]
format = data['format'] if 'format' in data else {}
res = {**stream, **format}
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))],
cli/xformers.sh
-12
View File
@@ -1,12 +0,0 @@
#/bin/env bash
echo "Installing xformers"
NVCC_FLAGS="--use_fast_math"
FORCE_CUDA="1"
TORCH_CUDA_ARCH_LIST="8.6"
pip install ninja -q
pip uninstall xformers -y 2>/dev/null
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
pip show torch
pip show xformers
python -m xformers.info
extensions-builtin/LDSR/ldsr_model_arch.py
+7 -8
View File
@@ -40,7 +40,7 @@ class LDSR:
model = model.to(shared.device)
if half_attention:
model = model.half()
if shared.opts.opt_channelslast:
if shared.cmd_opts.opt_channelslast:
model = model.to(memory_format=torch.channels_last)
sd_hijack.model_hijack.hijack(model) # apply optimization
@@ -88,7 +88,7 @@ class LDSR:
x_t = None
logs = None
for n in range(n_runs):
for _ in range(n_runs):
if custom_shape is not None:
x_t = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device)
x_t = repeat(x_t, '1 c h w -> b c h w', b=custom_shape[0])
@@ -110,7 +110,6 @@ class LDSR:
diffusion_steps = int(steps)
eta = 1.0
down_sample_method = 'Lanczos'
gc.collect()
if torch.cuda.is_available:
@@ -131,11 +130,11 @@ class LDSR:
im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS)
else:
print(f"Down sample rate is 1 from {target_scale} / 4 (Not downsampling)")
# pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts
pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size
im_padded = Image.fromarray(np.pad(np.array(im_og), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge'))
logs = self.run(model["model"], im_padded, diffusion_steps, eta)
sample = logs["sample"]
@@ -158,7 +157,7 @@ class LDSR:
def get_cond(selected_path):
example = dict()
example = {}
up_f = 4
c = selected_path.convert('RGB')
c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0)
@@ -196,7 +195,7 @@ def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_s
@torch.no_grad()
def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize_x0=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None,
corrector_kwargs=None, x_T=None, ddim_use_x0_pred=False):
log = dict()
log = {}
z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key,
return_first_stage_outputs=True,
@@ -244,7 +243,7 @@ def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize
x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True)
log["sample_noquant"] = x_sample_noquant
log["sample_diff"] = torch.abs(x_sample_noquant - x_sample)
except:
except Exception:
pass
log["sample"] = x_sample
extensions-builtin/LDSR/preload.py
+6
View File
@@ -0,0 +1,6 @@
import os
from modules import paths
def preload(parser):
parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(paths.models_path, 'LDSR'))
extensions-builtin/LDSR/scripts/ldsr_model.py
+15 -8
View File
@@ -7,7 +7,8 @@ from basicsr.utils.download_util import load_file_from_url
from modules.upscaler import Upscaler, UpscalerData
from ldsr_model_arch import LDSR
from modules import shared, script_callbacks
import sd_hijack_autoencoder, sd_hijack_ddpm_v1
import sd_hijack_autoencoder # noqa: F401
import sd_hijack_ddpm_v1 # noqa: F401
class UpscalerLDSR(Upscaler):
@@ -25,22 +26,28 @@ class UpscalerLDSR(Upscaler):
yaml_path = os.path.join(self.model_path, "project.yaml")
old_model_path = os.path.join(self.model_path, "model.pth")
new_model_path = os.path.join(self.model_path, "model.ckpt")
safetensors_model_path = os.path.join(self.model_path, "model.safetensors")
local_model_paths = self.find_models(ext_filter=[".ckpt", ".safetensors"])
local_ckpt_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.ckpt")]), None)
local_safetensors_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.safetensors")]), None)
local_yaml_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("project.yaml")]), None)
if os.path.exists(yaml_path):
statinfo = os.stat(yaml_path)
if statinfo.st_size >= 10485760:
print("Removing invalid LDSR YAML file.")
os.remove(yaml_path)
if os.path.exists(old_model_path):
print("Renaming model from model.pth to model.ckpt")
os.rename(old_model_path, new_model_path)
if os.path.exists(safetensors_model_path):
model = safetensors_model_path
if local_safetensors_path is not None and os.path.exists(local_safetensors_path):
model = local_safetensors_path
else:
model = load_file_from_url(url=self.model_url, model_dir=self.model_path,
file_name="model.ckpt", progress=True)
yaml = load_file_from_url(url=self.yaml_url, model_dir=self.model_path,
file_name="project.yaml", progress=True)
model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="model.ckpt", progress=True)
yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml", progress=True)
try:
return LDSR(model, yaml)
extensions-builtin/LDSR/sd_hijack_autoencoder.py
+17 -11
View File
@@ -1,16 +1,21 @@
# The content of this file comes from the ldm/models/autoencoder.py file of the compvis/stable-diffusion repo
# The VQModel & VQModelInterface were subsequently removed from ldm/models/autoencoder.py when we moved to the stability-ai/stablediffusion repo
# As the LDSR upscaler relies on VQModel & VQModelInterface, the hijack aims to put them back into the ldm.models.autoencoder
import numpy as np
import torch
import pytorch_lightning as pl
import torch.nn.functional as F
from contextlib import contextmanager
from torch.optim.lr_scheduler import LambdaLR
from ldm.modules.ema import LitEma
from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
from ldm.modules.diffusionmodules.model import Encoder, Decoder
from ldm.util import instantiate_from_config
import ldm.models.autoencoder
from packaging import version
class VQModel(pl.LightningModule):
def __init__(self,
@@ -19,7 +24,7 @@ class VQModel(pl.LightningModule):
n_embed,
embed_dim,
ckpt_path=None,
ignore_keys=[],
ignore_keys=None,
image_key="image",
colorize_nlabels=None,
monitor=None,
@@ -57,7 +62,7 @@ class VQModel(pl.LightningModule):
print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
if ckpt_path is not None:
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or [])
self.scheduler_config = scheduler_config
self.lr_g_factor = lr_g_factor
@@ -76,11 +81,11 @@ class VQModel(pl.LightningModule):
if context is not None:
print(f"{context}: Restored training weights")
def init_from_ckpt(self, path, ignore_keys=list()):
def init_from_ckpt(self, path, ignore_keys=None):
sd = torch.load(path, map_location="cpu")["state_dict"]
keys = list(sd.keys())
for k in keys:
for ik in ignore_keys:
for ik in ignore_keys or []:
if k.startswith(ik):
print("Deleting key {} from state_dict.".format(k))
del sd[k]
@@ -165,7 +170,7 @@ class VQModel(pl.LightningModule):
def validation_step(self, batch, batch_idx):
log_dict = self._validation_step(batch, batch_idx)
with self.ema_scope():
log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
self._validation_step(batch, batch_idx, suffix="_ema")
return log_dict
def _validation_step(self, batch, batch_idx, suffix=""):
@@ -232,7 +237,7 @@ class VQModel(pl.LightningModule):
return self.decoder.conv_out.weight
def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
log = dict()
log = {}
x = self.get_input(batch, self.image_key)
x = x.to(self.device)
if only_inputs:
@@ -249,7 +254,8 @@ class VQModel(pl.LightningModule):
if plot_ema:
with self.ema_scope():
xrec_ema, _ = self(x)
if x.shape[1] > 3: xrec_ema = self.to_rgb(xrec_ema)
if x.shape[1] > 3:
xrec_ema = self.to_rgb(xrec_ema)
log["reconstructions_ema"] = xrec_ema
return log
@@ -264,7 +270,7 @@ class VQModel(pl.LightningModule):
class VQModelInterface(VQModel):
def __init__(self, embed_dim, *args, **kwargs):
super().__init__(embed_dim=embed_dim, *args, **kwargs)
super().__init__(*args, embed_dim=embed_dim, **kwargs)
self.embed_dim = embed_dim
def encode(self, x):
@@ -282,5 +288,5 @@ class VQModelInterface(VQModel):
dec = self.decoder(quant)
return dec
setattr(ldm.models.autoencoder, "VQModel", VQModel)
setattr(ldm.models.autoencoder, "VQModelInterface", VQModelInterface)
ldm.models.autoencoder.VQModel = VQModel
ldm.models.autoencoder.VQModelInterface = VQModelInterface
extensions-builtin/LDSR/sd_hijack_ddpm_v1.py
+30 -36
View File
@@ -48,7 +48,7 @@ class DDPMV1(pl.LightningModule):
beta_schedule="linear",
loss_type="l2",
ckpt_path=None,
ignore_keys=[],
ignore_keys=None,
load_only_unet=False,
monitor="val/loss",
use_ema=True,
@@ -100,7 +100,7 @@ class DDPMV1(pl.LightningModule):
if monitor is not None:
self.monitor = monitor
if ckpt_path is not None:
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or [], only_model=load_only_unet)
self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
@@ -182,13 +182,13 @@ class DDPMV1(pl.LightningModule):
if context is not None:
print(f"{context}: Restored training weights")
def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
def init_from_ckpt(self, path, ignore_keys=None, only_model=False):
sd = torch.load(path, map_location="cpu")
if "state_dict" in list(sd.keys()):
sd = sd["state_dict"]
keys = list(sd.keys())
for k in keys:
for ik in ignore_keys:
for ik in ignore_keys or []:
if k.startswith(ik):
print("Deleting key {} from state_dict.".format(k))
del sd[k]
@@ -375,7 +375,7 @@ class DDPMV1(pl.LightningModule):
@torch.no_grad()
def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
log = dict()
log = {}
x = self.get_input(batch, self.first_stage_key)
N = min(x.shape[0], N)
n_row = min(x.shape[0], n_row)
@@ -383,7 +383,7 @@ class DDPMV1(pl.LightningModule):
log["inputs"] = x
# get diffusion row
diffusion_row = list()
diffusion_row = []
x_start = x[:n_row]
for t in range(self.num_timesteps):
@@ -444,13 +444,13 @@ class LatentDiffusionV1(DDPMV1):
conditioning_key = None
ckpt_path = kwargs.pop("ckpt_path", None)
ignore_keys = kwargs.pop("ignore_keys", [])
super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
super().__init__(*args, conditioning_key=conditioning_key, **kwargs)
self.concat_mode = concat_mode
self.cond_stage_trainable = cond_stage_trainable
self.cond_stage_key = cond_stage_key
try:
self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
except:
except Exception:
self.num_downs = 0
if not scale_by_std:
self.scale_factor = scale_factor
@@ -460,7 +460,7 @@ class LatentDiffusionV1(DDPMV1):
self.instantiate_cond_stage(cond_stage_config)
self.cond_stage_forward = cond_stage_forward
self.clip_denoised = False
self.bbox_tokenizer = None
self.bbox_tokenizer = None
self.restarted_from_ckpt = False
if ckpt_path is not None:
@@ -792,7 +792,7 @@ class LatentDiffusionV1(DDPMV1):
z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
# 2. apply model loop over last dim
if isinstance(self.first_stage_model, VQModelInterface):
if isinstance(self.first_stage_model, VQModelInterface):
output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
force_not_quantize=predict_cids or force_not_quantize)
for i in range(z.shape[-1])]
@@ -877,16 +877,6 @@ class LatentDiffusionV1(DDPMV1):
c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
return self.p_losses(x, c, t, *args, **kwargs)
def _rescale_annotations(self, bboxes, crop_coordinates): # TODO: move to dataset
def rescale_bbox(bbox):
x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
w = min(bbox[2] / crop_coordinates[2], 1 - x0)
h = min(bbox[3] / crop_coordinates[3], 1 - y0)
return x0, y0, w, h
return [rescale_bbox(b) for b in bboxes]
def apply_model(self, x_noisy, t, cond, return_ids=False):
if isinstance(cond, dict):
@@ -900,7 +890,7 @@ class LatentDiffusionV1(DDPMV1):
if hasattr(self, "split_input_params"):
assert len(cond) == 1 # todo can only deal with one conditioning atm
assert not return_ids
assert not return_ids
ks = self.split_input_params["ks"] # eg. (128, 128)
stride = self.split_input_params["stride"] # eg. (64, 64)
@@ -1126,7 +1116,7 @@ class LatentDiffusionV1(DDPMV1):
if cond is not None:
if isinstance(cond, dict):
cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
[x[:batch_size] for x in cond[key]] for key in cond}
else:
cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
@@ -1157,8 +1147,10 @@ class LatentDiffusionV1(DDPMV1):
if i % log_every_t == 0 or i == timesteps - 1:
intermediates.append(x0_partial)
if callback: callback(i)
if img_callback: img_callback(img, i)
if callback:
callback(i)
if img_callback:
img_callback(img, i)
return img, intermediates
@torch.no_grad()
@@ -1205,8 +1197,10 @@ class LatentDiffusionV1(DDPMV1):
if i % log_every_t == 0 or i == timesteps - 1:
intermediates.append(img)
if callback: callback(i)
if img_callback: img_callback(img, i)
if callback:
callback(i)
if img_callback:
img_callback(img, i)
if return_intermediates:
return img, intermediates
@@ -1221,7 +1215,7 @@ class LatentDiffusionV1(DDPMV1):
if cond is not None:
if isinstance(cond, dict):
cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
[x[:batch_size] for x in cond[key]] for key in cond}
else:
cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
return self.p_sample_loop(cond,
@@ -1253,7 +1247,7 @@ class LatentDiffusionV1(DDPMV1):
use_ddim = ddim_steps is not None
log = dict()
log = {}
z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
return_first_stage_outputs=True,
force_c_encode=True,
@@ -1280,7 +1274,7 @@ class LatentDiffusionV1(DDPMV1):
if plot_diffusion_rows:
# get diffusion row
diffusion_row = list()
diffusion_row = []
z_start = z[:n_row]
for t in range(self.num_timesteps):
if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
@@ -1322,7 +1316,7 @@ class LatentDiffusionV1(DDPMV1):
if inpaint:
# make a simple center square
b, h, w = z.shape[0], z.shape[2], z.shape[3]
h, w = z.shape[2], z.shape[3]
mask = torch.ones(N, h, w).to(self.device)
# zeros will be filled in
mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
@@ -1424,10 +1418,10 @@ class Layout2ImgDiffusionV1(LatentDiffusionV1):
# TODO: move all layout-specific hacks to this class
def __init__(self, cond_stage_key, *args, **kwargs):
assert cond_stage_key == 'coordinates_bbox', 'Layout2ImgDiffusion only for cond_stage_key="coordinates_bbox"'
super().__init__(cond_stage_key=cond_stage_key, *args, **kwargs)
super().__init__(*args, cond_stage_key=cond_stage_key, **kwargs)
def log_images(self, batch, N=8, *args, **kwargs):
logs = super().log_images(batch=batch, N=N, *args, **kwargs)
logs = super().log_images(*args, batch=batch, N=N, **kwargs)
key = 'train' if self.training else 'validation'
dset = self.trainer.datamodule.datasets[key]
@@ -1443,7 +1437,7 @@ class Layout2ImgDiffusionV1(LatentDiffusionV1):
logs['bbox_image'] = cond_img
return logs
setattr(ldm.models.diffusion.ddpm, "DDPMV1", DDPMV1)
setattr(ldm.models.diffusion.ddpm, "LatentDiffusionV1", LatentDiffusionV1)
setattr(ldm.models.diffusion.ddpm, "DiffusionWrapperV1", DiffusionWrapperV1)
setattr(ldm.models.diffusion.ddpm, "Layout2ImgDiffusionV1", Layout2ImgDiffusionV1)
ldm.models.diffusion.ddpm.DDPMV1 = DDPMV1
ldm.models.diffusion.ddpm.LatentDiffusionV1 = LatentDiffusionV1
ldm.models.diffusion.ddpm.DiffusionWrapperV1 = DiffusionWrapperV1
ldm.models.diffusion.ddpm.Layout2ImgDiffusionV1 = Layout2ImgDiffusionV1
extensions-builtin/Lora/extra_networks_lora.py
+19
View File
@@ -1,6 +1,7 @@
from modules import extra_networks, shared
import lora
class ExtraNetworkLora(extra_networks.ExtraNetwork):
def __init__(self):
super().__init__('lora')
@@ -22,5 +23,23 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
lora.load_loras(names, multipliers)
if shared.opts.lora_add_hashes_to_infotext:
lora_hashes = []
for item in lora.loaded_loras:
shorthash = item.lora_on_disk.shorthash
if not shorthash:
continue
alias = item.mentioned_name
if not alias:
continue
alias = alias.replace(":", "").replace(",", "")
lora_hashes.append(f"{alias}: {shorthash}")
if lora_hashes:
p.extra_generation_params["Lora hashes"] = ", ".join(lora_hashes)
def deactivate(self, p):
pass
extensions-builtin/Lora/lora.py
+181 -57
View File
@@ -1,10 +1,9 @@
import glob
import os
import re
from typing import Union
import torch
from typing import Union
from modules import shared, devices, sd_models, errors
from modules import shared, devices, sd_models, errors, scripts, sd_hijack, hashes
metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20}
@@ -77,9 +76,9 @@ class LoraOnDisk:
self.name = name
self.filename = filename
self.metadata = {}
self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors"
_, ext = os.path.splitext(filename)
if ext.lower() == ".safetensors":
if self.is_safetensors:
try:
self.metadata = sd_models.read_metadata_from_safetensors(filename)
except Exception as e:
@@ -93,15 +92,45 @@ class LoraOnDisk:
self.metadata = m
self.ssmd_cover_images = self.metadata.pop('ssmd_cover_images', None) # those are cover images and they are too big to display in UI as text
self.alias = self.metadata.get('ss_output_name', self.name)
self.hash = None
self.shorthash = None
self.set_hash(
self.metadata.get('sshs_model_hash') or
hashes.sha256_from_cache(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or
''
)
def set_hash(self, v):
self.hash = v
self.shorthash = self.hash[0:12]
if self.shorthash:
available_lora_hash_lookup[self.shorthash] = self
def read_hash(self):
if not self.hash:
self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
def get_alias(self):
if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in forbidden_lora_aliases:
return self.name
else:
return self.alias
class LoraModule:
def __init__(self, name):
def __init__(self, name, lora_on_disk: LoraOnDisk):
self.name = name
self.lora_on_disk = lora_on_disk
self.multiplier = 1.0
self.modules = {}
self.mtime = None
self.mentioned_name = None
"""the text that was used to add lora to prompt - can be either name or an alias"""
class LoraUpDownModule:
def __init__(self):
@@ -126,20 +155,23 @@ def assign_lora_names_to_compvis_modules(sd_model):
sd_model.lora_layer_mapping = lora_layer_mapping
def load_lora(name, filename):
lora = LoraModule(name)
lora.mtime = os.path.getmtime(filename)
def load_lora(name, lora_on_disk):
lora = LoraModule(name, lora_on_disk)
lora.mtime = os.path.getmtime(lora_on_disk.filename)
sd = sd_models.read_state_dict(filename)
sd = sd_models.read_state_dict(lora_on_disk.filename)
# this should not be needed but is here as an emergency fix for an unknown error people are experiencing in 1.2.0
if not hasattr(shared.sd_model, 'lora_layer_mapping'):
assign_lora_names_to_compvis_modules(shared.sd_model)
keys_failed_to_match = {}
is_sd2 = 'model_transformer_resblocks' in shared.sd_model.lora_layer_mapping
for key_diffusers, weight in sd.items():
lora_key_parts = key_diffusers.split(".", 1)
key_diffusers_without_lora_parts = lora_key_parts[0]
lora_key = lora_key_parts[1] if len(lora_key_parts) > 1 else ""
key_diffusers_without_lora_parts, lora_key = key_diffusers.split(".", 1)
key = convert_diffusers_name_to_compvis(key_diffusers_without_lora_parts, is_sd2)
sd_module = shared.sd_model.lora_layer_mapping.get(key, None)
if sd_module is None:
@@ -166,11 +198,14 @@ def load_lora(name, filename):
module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
elif type(sd_module) == torch.nn.MultiheadAttention:
module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)
elif type(sd_module) == torch.nn.Conv2d:
module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (weight.shape[2], weight.shape[3]), bias=False)
elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (1, 1):
module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False)
elif type(sd_module) == torch.nn.Conv2d and weight.shape[2:] == (3, 3):
module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (3, 3), bias=False)
else:
print(f'Lora layer {key_diffusers} matched a layer with unsupported type: {type(sd_module).__name__}')
continue
raise AssertionError(f"Lora layer {key_diffusers} matched a layer with unsupported type: {type(sd_module).__name__}")
with torch.no_grad():
module.weight.copy_(weight)
@@ -182,10 +217,10 @@ def load_lora(name, filename):
elif lora_key == "lora_down.weight":
lora_module.down = module
else:
assert False, f'Bad Lora layer name: {key_diffusers} - must end in lora_up.weight, lora_down.weight or alpha'
raise AssertionError(f"Bad Lora layer name: {key_diffusers} - must end in lora_up.weight, lora_down.weight or alpha")
if len(keys_failed_to_match) > 0:
print(f"Failed to match keys when loading Lora {filename}: {keys_failed_to_match}")
print(f"Failed to match keys when loading Lora {lora_on_disk.filename}: {keys_failed_to_match}")
return lora
@@ -199,27 +234,42 @@ def load_loras(names, multipliers=None):
loaded_loras.clear()
loras_on_disk = [available_loras.get(name, None) for name in names]
if any([x is None for x in loras_on_disk]):
loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
if any(x is None for x in loras_on_disk):
list_available_loras()
loras_on_disk = [available_loras.get(name, None) for name in names]
loras_on_disk = [available_lora_aliases.get(name, None) for name in names]
failed_to_load_loras = []
for i, name in enumerate(names):
lora = already_loaded.get(name, None)
lora_on_disk = loras_on_disk[i]
if lora_on_disk is not None:
if lora is None or os.path.getmtime(lora_on_disk.filename) > lora.mtime:
lora = load_lora(name, lora_on_disk.filename)
try:
lora = load_lora(name, lora_on_disk)
except Exception as e:
errors.display(e, f"loading Lora {lora_on_disk.filename}")
continue
lora.mentioned_name = name
lora_on_disk.read_hash()
if lora is None:
failed_to_load_loras.append(name)
print(f"Couldn't find Lora with name {name}")
continue
lora.multiplier = multipliers[i] if multipliers else 1.0
loaded_loras.append(lora)
if len(failed_to_load_loras) > 0:
sd_hijack.model_hijack.comments.append("Failed to find Loras: " + ", ".join(failed_to_load_loras))
def lora_calc_updown(lora, module, target):
with torch.no_grad():
@@ -228,31 +278,29 @@ def lora_calc_updown(lora, module, target):
if up.shape[2:] == (1, 1) and down.shape[2:] == (1, 1):
updown = (up.squeeze(2).squeeze(2) @ down.squeeze(2).squeeze(2)).unsqueeze(2).unsqueeze(3)
elif up.shape[2:] == (3, 3) or down.shape[2:] == (3, 3):
updown = torch.nn.functional.conv2d(down.permute(1, 0, 2, 3), up).permute(1, 0, 2, 3)
else:
permute, h, w = False, 1, 1
if len(up.shape) == 4 and len(down.shape) == 4:
if up.shape[2:] == (1, 1):
up = up.squeeze(2).squeeze(2)
else:
n, c, h, w = up.shape
up = up.view(n, c, -1).permute(2, 0, 1)
permute = True
if down.shape[2:] == (1, 1):
down = down.squeeze(2).squeeze(2)
else:
n, c, h, w = down.shape
down = down.view(n, c, -1).permute(2, 0, 1)
permute = True
updown = up @ down
if permute:
nh, nw = updown.shape[1:]
updown = updown.permute(1, 2, 0).view(nh, nw, h, w)
updown = updown * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
return updown
def lora_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
weights_backup = getattr(self, "lora_weights_backup", None)
if weights_backup is None:
return
if isinstance(self, torch.nn.MultiheadAttention):
self.in_proj_weight.copy_(weights_backup[0])
self.out_proj.weight.copy_(weights_backup[1])
else:
self.weight.copy_(weights_backup)
def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.MultiheadAttention]):
"""
Applies the currently selected set of Loras to the weights of torch layer self.
@@ -277,12 +325,7 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu
self.lora_weights_backup = weights_backup
if current_names != wanted_names:
if weights_backup is not None:
if isinstance(self, torch.nn.MultiheadAttention):
self.in_proj_weight.copy_(weights_backup[0])
self.out_proj.weight.copy_(weights_backup[1])
else:
self.weight.copy_(weights_backup)
lora_restore_weights_from_backup(self)
for lora in loaded_loras:
module = lora.modules.get(lora_layer_name, None)
@@ -310,18 +353,51 @@ def lora_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.Mu
print(f'failed to calculate lora weights for layer {lora_layer_name}')
setattr(self, "lora_current_names", wanted_names)
self.lora_current_names = wanted_names
def lora_forward(module, input, original_forward):
"""
Old way of applying Lora by executing operations during layer's forward.
Stacking many loras this way results in big performance degradation.
"""
if len(loaded_loras) == 0:
return original_forward(module, input)
input = devices.cond_cast_unet(input)
lora_restore_weights_from_backup(module)
lora_reset_cached_weight(module)
res = original_forward(module, input)
lora_layer_name = getattr(module, 'lora_layer_name', None)
for lora in loaded_loras:
module = lora.modules.get(lora_layer_name, None)
if module is None:
continue
module.up.to(device=devices.device)
module.down.to(device=devices.device)
res = res + module.up(module.down(input)) * lora.multiplier * (module.alpha / module.up.weight.shape[1] if module.alpha else 1.0)
return res
def lora_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]):
setattr(self, "lora_current_names", ())
setattr(self, "lora_weights_backup", None)
self.lora_current_names = ()
self.lora_weights_backup = None
def lora_Linear_forward(self, lora_input):
def lora_Linear_forward(self, input):
if shared.opts.lora_functional:
return lora_forward(self, input, torch.nn.Linear_forward_before_lora)
lora_apply_weights(self)
return torch.nn.Linear_forward_before_lora(self, lora_input)
return torch.nn.Linear_forward_before_lora(self, input)
def lora_Linear_load_state_dict(self, *args, **kwargs):
@@ -330,10 +406,13 @@ def lora_Linear_load_state_dict(self, *args, **kwargs):
return torch.nn.Linear_load_state_dict_before_lora(self, *args, **kwargs)
def lora_Conv2d_forward(self, lora_input):
def lora_Conv2d_forward(self, input):
if shared.opts.lora_functional:
return lora_forward(self, input, torch.nn.Conv2d_forward_before_lora)
lora_apply_weights(self)
return torch.nn.Conv2d_forward_before_lora(self, lora_input)
return torch.nn.Conv2d_forward_before_lora(self, input)
def lora_Conv2d_load_state_dict(self, *args, **kwargs):
@@ -356,23 +435,68 @@ def lora_MultiheadAttention_load_state_dict(self, *args, **kwargs):
def list_available_loras():
available_loras.clear()
available_lora_aliases.clear()
forbidden_lora_aliases.clear()
available_lora_hash_lookup.clear()
forbidden_lora_aliases.update({"none": 1, "Addams": 1})
os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)
candidates = \
glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.pt'), recursive=True) + \
glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.safetensors'), recursive=True) + \
glob.glob(os.path.join(shared.cmd_opts.lora_dir, '**/*.ckpt'), recursive=True)
candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))
for filename in sorted(candidates, key=str.lower):
if os.path.isdir(filename):
continue
name = os.path.splitext(os.path.basename(filename))[0]
entry = LoraOnDisk(name, filename)
available_loras[name] = LoraOnDisk(name, filename)
available_loras[name] = entry
if entry.alias in available_lora_aliases:
forbidden_lora_aliases[entry.alias.lower()] = 1
available_lora_aliases[name] = entry
available_lora_aliases[entry.alias] = entry
re_lora_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
def infotext_pasted(infotext, params):
if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]:
return # if the other extension is active, it will handle those fields, no need to do anything
added = []
for k in params:
if not k.startswith("AddNet Model "):
continue
num = k[13:]
if params.get("AddNet Module " + num) != "LoRA":
continue
name = params.get("AddNet Model " + num)
if name is None:
continue
m = re_lora_name.match(name)
if m:
name = m.group(1)
multiplier = params.get("AddNet Weight A " + num, "1.0")
added.append(f"<lora:{name}:{multiplier}>")
if added:
params["Prompt"] += "\n" + "".join(added)
available_loras = {}
available_lora_aliases = {}
available_lora_hash_lookup = {}
forbidden_lora_aliases = {}
loaded_loras = []
list_available_loras()
extensions-builtin/Lora/preload.py
+6
View File
@@ -0,0 +1,6 @@
import os
from modules import paths
def preload(parser):
parser.add_argument("--lora-dir", type=str, help="Path to directory with Lora networks.", default=os.path.join(paths.models_path, 'Lora'))
extensions-builtin/Lora/scripts/lora_script.py
+57 -2
View File
@@ -1,12 +1,14 @@
import re
import torch
import gradio as gr
from fastapi import FastAPI
import lora
import extra_networks_lora
import ui_extra_networks_lora
from modules import script_callbacks, ui_extra_networks, extra_networks, shared
def unload():
torch.nn.Linear.forward = torch.nn.Linear_forward_before_lora
torch.nn.Linear._load_from_state_dict = torch.nn.Linear_load_state_dict_before_lora
@@ -49,8 +51,61 @@ torch.nn.MultiheadAttention._load_from_state_dict = lora.lora_MultiheadAttention
script_callbacks.on_model_loaded(lora.assign_lora_names_to_compvis_modules)
script_callbacks.on_script_unloaded(unload)
script_callbacks.on_before_ui(before_ui)
script_callbacks.on_infotext_pasted(lora.infotext_pasted)
shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), {
"sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in lora.available_loras]}, refresh=lora.list_available_loras),
"sd_lora": shared.OptionInfo("None", "Add Lora to prompt", gr.Dropdown, lambda: {"choices": ["None", *lora.available_loras]}, refresh=lora.list_available_loras),
"lora_preferred_name": shared.OptionInfo("Alias from file", "When adding to prompt, refer to Lora by", gr.Radio, {"choices": ["Alias from file", "Filename"]}),
"lora_add_hashes_to_infotext": shared.OptionInfo(True, "Add Lora hashes to infotext"),
}))
def create_lora_json(obj: lora.LoraOnDisk):
return {
"name": obj.name,
"alias": obj.alias,
"path": obj.filename,
"metadata": obj.metadata,
}
def api_loras(_: gr.Blocks, app: FastAPI):
@app.get("/sdapi/v1/loras")
async def get_loras():
return [create_lora_json(obj) for obj in lora.available_loras.values()]
@app.post("/sdapi/v1/refresh-loras")
async def refresh_loras():
return lora.list_available_loras()
script_callbacks.on_app_started(api_loras)
re_lora = re.compile("<lora:([^:]+):")
def infotext_pasted(infotext, d):
hashes = d.get("Lora hashes")
if not hashes:
return
hashes = [x.strip().split(':', 1) for x in hashes.split(",")]
hashes = {x[0].strip().replace(",", ""): x[1].strip() for x in hashes}
def lora_replacement(m):
alias = m.group(1)
shorthash = hashes.get(alias)
if shorthash is None:
return m.group(0)
lora_on_disk = lora.available_lora_hash_lookup.get(shorthash)
if lora_on_disk is None:
return m.group(0)
return f'<lora:{lora_on_disk.get_alias()}:'
d["Prompt"] = re.sub(re_lora, lora_replacement, d["Prompt"])
script_callbacks.on_infotext_pasted(infotext_pasted)
extensions-builtin/Lora/ui_extra_networks_lora.py
+6 -2
View File
@@ -15,16 +15,20 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage):
def list_items(self):
for name, lora_on_disk in lora.available_loras.items():
path, ext = os.path.splitext(lora_on_disk.filename)
alias = lora_on_disk.get_alias()
yield {
"name": name,
"filename": path,
"preview": self.find_preview(path),
"description": self.find_description(path),
"search_term": self.search_terms_from_path(lora_on_disk.filename),
"prompt": json.dumps(f"<lora:{name}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
"prompt": json.dumps(f"<lora:{alias}:") + " + opts.extra_networks_default_multiplier + " + json.dumps(">"),
"local_preview": f"{path}.{shared.opts.samples_format}",
"metadata": json.dumps(lora_on_disk.metadata, indent=4) if lora_on_disk.metadata else None,
}
def allowed_directories_for_previews(self):
return [shared.opts.lora_dir]
return [shared.cmd_opts.lora_dir]
extensions-builtin/ScuNET/preload.py
+6
View File
@@ -0,0 +1,6 @@
import os
from modules import paths
def preload(parser):
parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(paths.models_path, 'ScuNET'))
extensions-builtin/ScuNET/scripts/scunet_model.py
+79 -17
View File
@@ -5,11 +5,14 @@ import traceback
import PIL.Image
import numpy as np
import torch
from tqdm import tqdm
from basicsr.utils.download_util import load_file_from_url
import modules.upscaler
from modules import devices, modelloader
from modules import devices, modelloader, script_callbacks
from scunet_model_arch import SCUNet as net
from modules.shared import opts
class UpscalerScuNET(modules.upscaler.Upscaler):
@@ -42,34 +45,83 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
scalers.append(scaler_data2)
self.scalers = scalers
def do_upscale(self, img: PIL.Image, selected_file):
@staticmethod
@torch.no_grad()
def tiled_inference(img, model):
# test the image tile by tile
h, w = img.shape[2:]
tile = opts.SCUNET_tile
tile_overlap = opts.SCUNET_tile_overlap
if tile == 0:
return model(img)
device = devices.get_device_for('scunet')
assert tile % 8 == 0, "tile size should be a multiple of window_size"
sf = 1
stride = tile - tile_overlap
h_idx_list = list(range(0, h - tile, stride)) + [h - tile]
w_idx_list = list(range(0, w - tile, stride)) + [w - tile]
E = torch.zeros(1, 3, h * sf, w * sf, dtype=img.dtype, device=device)
W = torch.zeros_like(E, dtype=devices.dtype, device=device)
with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="ScuNET tiles") as pbar:
for h_idx in h_idx_list:
for w_idx in w_idx_list:
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)
E[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch)
W[
..., h_idx * sf: (h_idx + tile) * sf, w_idx * sf: (w_idx + tile) * sf
].add_(out_patch_mask)
pbar.update(1)
output = E.div_(W)
return output
def do_upscale(self, img: PIL.Image.Image, selected_file):
torch.cuda.empty_cache()
model = self.load_model(selected_file)
if model is None:
print(f"ScuNET: Unable to load model from {selected_file}", file=sys.stderr)
return img
device = devices.get_device_for('scunet')
img = np.array(img)
img = img[:, :, ::-1]
img = np.moveaxis(img, 2, 0) / 255
img = torch.from_numpy(img).float()
img = img.unsqueeze(0).to(device)
tile = opts.SCUNET_tile
h, w = img.height, img.width
np_img = np.array(img)
np_img = np_img[:, :, ::-1] # RGB to BGR
np_img = np_img.transpose((2, 0, 1)) / 255 # HWC to CHW
torch_img = torch.from_numpy(np_img).float().unsqueeze(0).to(device) # type: ignore
with torch.no_grad():
output = model(img)
output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
output = 255. * np.moveaxis(output, 0, 2)
output = output.astype(np.uint8)
output = output[:, :, ::-1]
if tile > h or tile > w:
_img = torch.zeros(1, 3, max(h, tile), max(w, tile), dtype=torch_img.dtype, device=torch_img.device)
_img[:, :, :h, :w] = torch_img # pad image
torch_img = _img
torch_output = self.tiled_inference(torch_img, model).squeeze(0)
torch_output = torch_output[:, :h * 1, :w * 1] # remove padding, if any
np_output: np.ndarray = torch_output.float().cpu().clamp_(0, 1).numpy()
del torch_img, torch_output
torch.cuda.empty_cache()
return PIL.Image.fromarray(output, 'RGB')
output = np_output.transpose((1, 2, 0)) # CHW to HWC
output = output[:, :, ::-1] # BGR to RGB
return PIL.Image.fromarray((output * 255).astype(np.uint8))
def load_model(self, path: str):
device = devices.get_device_for('scunet')
if "http" in path:
filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
progress=True)
filename = load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="%s.pth" % self.name, progress=True)
else:
filename = path
if not os.path.exists(os.path.join(self.model_path, filename)) or filename is None:
@@ -79,9 +131,19 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
model = net(in_nc=3, config=[4, 4, 4, 4, 4, 4, 4], dim=64)
model.load_state_dict(torch.load(filename), strict=True)
model.eval()
for k, v in model.named_parameters():
for _, v in model.named_parameters():
v.requires_grad = False
model = model.to(device)
return model
def on_ui_settings():
import gradio as gr
from modules import shared
shared.opts.add_option("SCUNET_tile", shared.OptionInfo(256, "Tile size for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling")).info("0 = no tiling"))
shared.opts.add_option("SCUNET_tile_overlap", shared.OptionInfo(8, "Tile overlap for SCUNET upscalers.", gr.Slider, {"minimum": 0, "maximum": 64, "step": 1}, section=('upscaling', "Upscaling")).info("Low values = visible seam"))
script_callbacks.on_ui_settings(on_ui_settings)
extensions-builtin/ScuNET/scunet_model_arch.py
+7 -4
View File
@@ -61,7 +61,9 @@ class WMSA(nn.Module):
Returns:
output: tensor shape [b h w c]
"""
if self.type != 'W': x = torch.roll(x, shifts=(-(self.window_size // 2), -(self.window_size // 2)), dims=(1, 2))
if self.type != 'W':
x = torch.roll(x, shifts=(-(self.window_size // 2), -(self.window_size // 2)), dims=(1, 2))
x = rearrange(x, 'b (w1 p1) (w2 p2) c -> b w1 w2 p1 p2 c', p1=self.window_size, p2=self.window_size)
h_windows = x.size(1)
w_windows = x.size(2)
@@ -85,8 +87,9 @@ class WMSA(nn.Module):
output = self.linear(output)
output = rearrange(output, 'b (w1 w2) (p1 p2) c -> b (w1 p1) (w2 p2) c', w1=h_windows, p1=self.window_size)
if self.type != 'W': output = torch.roll(output, shifts=(self.window_size // 2, self.window_size // 2),
dims=(1, 2))
if self.type != 'W':
output = torch.roll(output, shifts=(self.window_size // 2, self.window_size // 2), dims=(1, 2))
return output
def relative_embedding(self):
@@ -262,4 +265,4 @@ class SCUNet(nn.Module):
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.LayerNorm):
nn.init.constant_(m.bias, 0)
nn.init.constant_(m.weight, 1.0)
nn.init.constant_(m.weight, 1.0)
extensions-builtin/SwinIR/preload.py
+6
View File
@@ -0,0 +1,6 @@
import os
from modules import paths
def preload(parser):
parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(paths.models_path, 'SwinIR'))
extensions-builtin/SwinIR/scripts/swinir_model.py
+18 -19
View File
@@ -1,13 +1,13 @@
import os
import numpy as np
import torch
from PIL import Image
from basicsr.utils.download_util import load_file_from_url
from tqdm import tqdm
from rich import print, progress # pylint: disable=redefined-builtin
from modules import modelloader, devices, script_callbacks, shared
from modules.shared import cmd_opts, opts, state
from modules.shared import opts, state
from swinir_model_arch import SwinIR as net
from swinir_model_arch_v2 import Swin2SR as net2
from modules.upscaler import Upscaler, UpscalerData
@@ -44,31 +44,31 @@ class UpscalerSwinIR(Upscaler):
img = upscale(img, model)
try:
torch.cuda.empty_cache()
except:
except Exception:
pass
return img
def load_model(self, path, scale=4):
if "http" in path:
dl_name = "%s%s" % (self.model_name.replace(" ", "_"), ".pth")
filename = load_file_from_url(url=path, model_dir=self.model_path, file_name=dl_name, progress=True)
filename = load_file_from_url(url=path, model_dir=self.model_download_path, file_name=dl_name, progress=True)
else:
filename = path
if filename is None or not os.path.exists(filename):
return None
if filename.endswith(".v2.pth"):
model = net2(
upscale=scale,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv",
upscale=scale,
in_chans=3,
img_size=64,
window_size=8,
img_range=1.0,
depths=[6, 6, 6, 6, 6, 6],
embed_dim=180,
num_heads=[6, 6, 6, 6, 6, 6],
mlp_ratio=2,
upsampler="nearest+conv",
resi_connection="1conv",
)
params = None
else:
@@ -87,9 +87,8 @@ class UpscalerSwinIR(Upscaler):
)
params = "params_ema"
with progress.open(filename, 'rb', description=f'Loading weights: [cyan]{filename}', auto_refresh=True) as f:
pretrained_model = torch.load(filename)
if params is not None and params in pretrained_model:
pretrained_model = torch.load(filename)
if params is not None:
model.load_state_dict(pretrained_model[params], strict=True)
else:
model.load_state_dict(pretrained_model, strict=True)
@@ -151,7 +150,7 @@ def inference(img, model, tile, tile_overlap, window_size, scale):
for w_idx in w_idx_list:
if state.interrupted or state.skipped:
break
in_patch = img[..., h_idx: h_idx + tile, w_idx: w_idx + tile]
out_patch = model(in_patch)
out_patch_mask = torch.ones_like(out_patch)
extensions-builtin/SwinIR/swinir_model_arch.py
+3 -3
View File
@@ -644,7 +644,7 @@ class SwinIR(nn.Module):
"""
def __init__(self, img_size=64, patch_size=1, in_chans=3,
embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],
embed_dim=96, depths=(6, 6, 6, 6), num_heads=(6, 6, 6, 6),
window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
@@ -805,7 +805,7 @@ class SwinIR(nn.Module):
def forward(self, x):
H, W = x.shape[2:]
x = self.check_image_size(x)
self.mean = self.mean.type_as(x)
x = (x - self.mean) * self.img_range
@@ -844,7 +844,7 @@ class SwinIR(nn.Module):
H, W = self.patches_resolution
flops += H * W * 3 * self.embed_dim * 9
flops += self.patch_embed.flops()
for i, layer in enumerate(self.layers):
for layer in self.layers:
flops += layer.flops()
flops += H * W * 3 * self.embed_dim * self.embed_dim
flops += self.upsample.flops()
extensions-builtin/SwinIR/swinir_model_arch_v2.py
+29 -29
View File
@@ -74,7 +74,7 @@ class WindowAttention(nn.Module):
"""
def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.,
pretrained_window_size=[0, 0]):
pretrained_window_size=(0, 0)):
super().__init__()
self.dim = dim
@@ -241,7 +241,7 @@ class SwinTransformerBlock(nn.Module):
attn_mask = None
self.register_buffer("attn_mask", attn_mask)
def calculate_mask(self, x_size):
# calculate attention mask for SW-MSA
H, W = x_size
@@ -263,7 +263,7 @@ class SwinTransformerBlock(nn.Module):
attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
return attn_mask
return attn_mask
def forward(self, x, x_size):
H, W = x_size
@@ -288,7 +288,7 @@ class SwinTransformerBlock(nn.Module):
attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C
else:
attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))
# merge windows
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C
@@ -369,7 +369,7 @@ class PatchMerging(nn.Module):
H, W = self.input_resolution
flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
flops += H * W * self.dim // 2
return flops
return flops
class BasicLayer(nn.Module):
""" A basic Swin Transformer layer for one stage.
@@ -447,7 +447,7 @@ class BasicLayer(nn.Module):
nn.init.constant_(blk.norm1.weight, 0)
nn.init.constant_(blk.norm2.bias, 0)
nn.init.constant_(blk.norm2.weight, 0)
class PatchEmbed(nn.Module):
r""" Image to Patch Embedding
Args:
@@ -492,7 +492,7 @@ class PatchEmbed(nn.Module):
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
if self.norm is not None:
flops += Ho * Wo * self.embed_dim
return flops
return flops
class RSTB(nn.Module):
"""Residual Swin Transformer Block (RSTB).
@@ -531,7 +531,7 @@ class RSTB(nn.Module):
num_heads=num_heads,
window_size=window_size,
mlp_ratio=mlp_ratio,
qkv_bias=qkv_bias,
qkv_bias=qkv_bias,
drop=drop, attn_drop=attn_drop,
drop_path=drop_path,
norm_layer=norm_layer,
@@ -622,7 +622,7 @@ class Upsample(nn.Sequential):
else:
raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
super(Upsample, self).__init__(*m)
class Upsample_hf(nn.Sequential):
"""Upsample module.
@@ -642,7 +642,7 @@ class Upsample_hf(nn.Sequential):
m.append(nn.PixelShuffle(3))
else:
raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')
super(Upsample_hf, self).__init__(*m)
super(Upsample_hf, self).__init__(*m)
class UpsampleOneStep(nn.Sequential):
@@ -667,8 +667,8 @@ class UpsampleOneStep(nn.Sequential):
H, W = self.input_resolution
flops = H * W * self.num_feat * 3 * 9
return flops
class Swin2SR(nn.Module):
r""" Swin2SR
@@ -698,8 +698,8 @@ class Swin2SR(nn.Module):
"""
def __init__(self, img_size=64, patch_size=1, in_chans=3,
embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],
window_size=7, mlp_ratio=4., qkv_bias=True,
embed_dim=96, depths=(6, 6, 6, 6), num_heads=(6, 6, 6, 6),
window_size=7, mlp_ratio=4., qkv_bias=True,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',
@@ -764,7 +764,7 @@ class Swin2SR(nn.Module):
num_heads=num_heads[i_layer],
window_size=window_size,
mlp_ratio=self.mlp_ratio,
qkv_bias=qkv_bias,
qkv_bias=qkv_bias,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results
norm_layer=norm_layer,
@@ -776,7 +776,7 @@ class Swin2SR(nn.Module):
)
self.layers.append(layer)
if self.upsampler == 'pixelshuffle_hf':
self.layers_hf = nn.ModuleList()
for i_layer in range(self.num_layers):
@@ -787,7 +787,7 @@ class Swin2SR(nn.Module):
num_heads=num_heads[i_layer],
window_size=window_size,
mlp_ratio=self.mlp_ratio,
qkv_bias=qkv_bias,
qkv_bias=qkv_bias,
drop=drop_rate, attn_drop=attn_drop_rate,
drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results
norm_layer=norm_layer,
@@ -799,7 +799,7 @@ class Swin2SR(nn.Module):
)
self.layers_hf.append(layer)
self.norm = norm_layer(self.num_features)
# build the last conv layer in deep feature extraction
@@ -829,10 +829,10 @@ class Swin2SR(nn.Module):
self.conv_aux = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
self.conv_after_aux = nn.Sequential(
nn.Conv2d(3, num_feat, 3, 1, 1),
nn.LeakyReLU(inplace=True))
nn.LeakyReLU(inplace=True))
self.upsample = Upsample(upscale, num_feat)
self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
elif self.upsampler == 'pixelshuffle_hf':
self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
nn.LeakyReLU(inplace=True))
@@ -846,7 +846,7 @@ class Swin2SR(nn.Module):
nn.Conv2d(embed_dim, num_feat, 3, 1, 1),
nn.LeakyReLU(inplace=True))
self.conv_last_hf = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
elif self.upsampler == 'pixelshuffledirect':
# for lightweight SR (to save parameters)
self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,
@@ -905,7 +905,7 @@ class Swin2SR(nn.Module):
x = self.patch_unembed(x, x_size)
return x
def forward_features_hf(self, x):
x_size = (x.shape[2], x.shape[3])
x = self.patch_embed(x)
@@ -919,7 +919,7 @@ class Swin2SR(nn.Module):
x = self.norm(x) # B L C
x = self.patch_unembed(x, x_size)
return x
return x
def forward(self, x):
H, W = x.shape[2:]
@@ -951,7 +951,7 @@ class Swin2SR(nn.Module):
x = self.conv_after_body(self.forward_features(x)) + x
x_before = self.conv_before_upsample(x)
x_out = self.conv_last(self.upsample(x_before))
x_hf = self.conv_first_hf(x_before)
x_hf = self.conv_after_body_hf(self.forward_features_hf(x_hf)) + x_hf
x_hf = self.conv_before_upsample_hf(x_hf)
@@ -977,15 +977,15 @@ class Swin2SR(nn.Module):
x_first = self.conv_first(x)
res = self.conv_after_body(self.forward_features(x_first)) + x_first
x = x + self.conv_last(res)
x = x / self.img_range + self.mean
if self.upsampler == "pixelshuffle_aux":
return x[:, :, :H*self.upscale, :W*self.upscale], aux
elif self.upsampler == "pixelshuffle_hf":
x_out = x_out / self.img_range + self.mean
return x_out[:, :, :H*self.upscale, :W*self.upscale], x[:, :, :H*self.upscale, :W*self.upscale], x_hf[:, :, :H*self.upscale, :W*self.upscale]
else:
return x[:, :, :H*self.upscale, :W*self.upscale]
@@ -994,7 +994,7 @@ class Swin2SR(nn.Module):
H, W = self.patches_resolution
flops += H * W * 3 * self.embed_dim * 9
flops += self.patch_embed.flops()
for i, layer in enumerate(self.layers):
for layer in self.layers:
flops += layer.flops()
flops += H * W * 3 * self.embed_dim * self.embed_dim
flops += self.upsample.flops()
@@ -1014,4 +1014,4 @@ if __name__ == '__main__':
x = torch.randn((1, 3, height, width))
x = model(x)
print(x.shape)
print(x.shape)
extensions-builtin/a1111-sd-webui-lycoris
+1 -1
Submodule extensions-builtin/a1111-sd-webui-lycoris updated: ce584a0ff8...21e9ea0f58
extensions-builtin/multidiffusion-upscaler-for-automatic1111
+1 -1
Submodule extensions-builtin/multidiffusion-upscaler-for-automatic1111 updated: 6931b89cb4...70b3c5ea3c
extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js
-103
View File
@@ -1,103 +0,0 @@
// Stable Diffusion WebUI - Bracket checker
// Version 1.0
// By Hingashi no Florin/Bwin4L
// Counts open and closed brackets (round, square, curly) in the prompt and negative prompt text boxes in the txt2img and img2img tabs.
// If there's a mismatch, the keyword counter turns red and if you hover on it, a tooltip tells you what's wrong.
function checkBrackets(evt, textArea, counterElt) {
errorStringParen = '(...) - Different number of opening and closing parentheses detected.\n';
errorStringSquare = '[...] - Different number of opening and closing square brackets detected.\n';
errorStringCurly = '{...} - Different number of opening and closing curly brackets detected.\n';
openBracketRegExp = /\(/g;
closeBracketRegExp = /\)/g;
openSquareBracketRegExp = /\[/g;
closeSquareBracketRegExp = /\]/g;
openCurlyBracketRegExp = /\{/g;
closeCurlyBracketRegExp = /\}/g;
totalOpenBracketMatches = 0;
totalCloseBracketMatches = 0;
totalOpenSquareBracketMatches = 0;
totalCloseSquareBracketMatches = 0;
totalOpenCurlyBracketMatches = 0;
totalCloseCurlyBracketMatches = 0;
openBracketMatches = textArea.value.match(openBracketRegExp);
if(openBracketMatches) {
totalOpenBracketMatches = openBracketMatches.length;
}
closeBracketMatches = textArea.value.match(closeBracketRegExp);
if(closeBracketMatches) {
totalCloseBracketMatches = closeBracketMatches.length;
}
openSquareBracketMatches = textArea.value.match(openSquareBracketRegExp);
if(openSquareBracketMatches) {
totalOpenSquareBracketMatches = openSquareBracketMatches.length;
}
closeSquareBracketMatches = textArea.value.match(closeSquareBracketRegExp);
if(closeSquareBracketMatches) {
totalCloseSquareBracketMatches = closeSquareBracketMatches.length;
}
openCurlyBracketMatches = textArea.value.match(openCurlyBracketRegExp);
if(openCurlyBracketMatches) {
totalOpenCurlyBracketMatches = openCurlyBracketMatches.length;
}
closeCurlyBracketMatches = textArea.value.match(closeCurlyBracketRegExp);
if(closeCurlyBracketMatches) {
totalCloseCurlyBracketMatches = closeCurlyBracketMatches.length;
}
if(totalOpenBracketMatches != totalCloseBracketMatches) {
if(!counterElt.title.includes(errorStringParen)) {
counterElt.title += errorStringParen;
}
} else {
counterElt.title = counterElt.title.replace(errorStringParen, '');
}
if(totalOpenSquareBracketMatches != totalCloseSquareBracketMatches) {
if(!counterElt.title.includes(errorStringSquare)) {
counterElt.title += errorStringSquare;
}
} else {
counterElt.title = counterElt.title.replace(errorStringSquare, '');
}
if(totalOpenCurlyBracketMatches != totalCloseCurlyBracketMatches) {
if(!counterElt.title.includes(errorStringCurly)) {
counterElt.title += errorStringCurly;
}
} else {
counterElt.title = counterElt.title.replace(errorStringCurly, '');
}
if(counterElt.title != '') {
counterElt.classList.add('error');
} else {
counterElt.classList.remove('error');
}
}
function setupBracketChecking(id_prompt, id_counter){
var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea");
var counter = gradioApp().getElementById(id_counter)
textarea.addEventListener("input", function(evt){
checkBrackets(evt, textarea, counter)
});
}
onUiLoaded(function(){
setupBracketChecking('txt2img_prompt', 'txt2img_token_counter')
setupBracketChecking('txt2img_neg_prompt', 'txt2img_negative_token_counter')
setupBracketChecking('img2img_prompt', 'img2img_token_counter')
setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter')
})
extensions-builtin/sd-dynamic-thresholding
+1 -1
Submodule extensions-builtin/sd-dynamic-thresholding updated: 5d3465c4b2...f02cacfc92
extensions-builtin/sd-extension-aesthetic-scorer
-1
Submodule extensions-builtin/sd-extension-aesthetic-scorer deleted from 896b6aa034
extensions-builtin/sd-extension-steps-animation
-1
Submodule extensions-builtin/sd-extension-steps-animation deleted from 90663eb745
extensions-builtin/sd-extension-system-info
+1 -1
Submodule extensions-builtin/sd-extension-system-info updated: 70ab5cf312...064c856aca
extensions-builtin/sd-webui-agent-scheduler
Submodule
+1
Submodule extensions-builtin/sd-webui-agent-scheduler added at b070319583
extensions-builtin/sd-webui-controlnet
+1 -1
Submodule extensions-builtin/sd-webui-controlnet updated: 4d4b1f8c00...d8551e447d
extensions-builtin/sd-webui-model-converter
-1
Submodule extensions-builtin/sd-webui-model-converter deleted from d19e281682
extensions-builtin/seed_travel
-1
Submodule extensions-builtin/seed_travel deleted from 4bc8b2f10d
extensions-builtin/stable-diffusion-webui-images-browser
+1 -1
Submodule extensions-builtin/stable-diffusion-webui-images-browser updated: 2c988c08c7...5795886bee
extensions-builtin/stable-diffusion-webui-rembg
+1 -1
Submodule extensions-builtin/stable-diffusion-webui-rembg updated: 64821f0476...3d9eedbbf0
javascript/black-orange.jpg → html/black-orange.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 151 KiB

After

Width:  |  Height:  |  Size: 151 KiB

html/card-no-preview.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 82 KiB

After

Width:  |  Height:  |  Size: 2.3 KiB

javascript/gradio-base.jpg → html/gradio-base.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 101 KiB

After

Width:  |  Height:  |  Size: 101 KiB

javascript/gradio-default.jpg → html/gradio-default.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 96 KiB

After

Width:  |  Height:  |  Size: 96 KiB

javascript/gradio-glass.jpg → html/gradio-glass.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 97 KiB

After

Width:  |  Height:  |  Size: 97 KiB

javascript/gradio-monochrome.jpg → html/gradio-monochrome.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 89 KiB

After

Width:  |  Height:  |  Size: 89 KiB

javascript/gradio-soft.jpg → html/gradio-soft.jpg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 215 KiB

After

Width:  |  Height:  |  Size: 215 KiB

html/licenses.html
+26
View File
@@ -661,4 +661,30 @@ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
</pre>
<h2><a href="https://github.com/madebyollin/taesd/blob/main/LICENSE">TAESD</a></h2>
<small>Tiny AutoEncoder for Stable Diffusion option for live previews</small>
<pre>
MIT License
Copyright (c) 2023 Ollin Boer Bohan
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
</pre>
automatic.ico → html/logo.ico
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 23 KiB

After

Width:  |  Height:  |  Size: 23 KiB

automatic.png → html/logo.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 84 KiB

After

Width:  |  Height:  |  Size: 84 KiB

automatic.svg → html/logo.svg
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 3.6 KiB

After

Width:  |  Height:  |  Size: 3.6 KiB

html/notification.mp3
BIN
View File
Binary file not shown.
javascript/roboto.ttf → html/roboto.ttf
View File
installer.py
+782
View File
@@ -0,0 +1,782 @@
import os
import sys
import json
import time
import shutil
import logging
import platform
import subprocess
import io
import pstats
import cProfile
import pkg_resources
try:
from modules.cmd_args import parser
except:
import argparse
parser = argparse.ArgumentParser(description="SD.Next", conflict_handler='resolve', formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=55, indent_increment=2, width=200))
class Dot(dict): # dot notation access to dictionary attributes
__getattr__ = dict.get
__setattr__ = dict.__setitem__
__delattr__ = dict.__delitem__
log = logging.getLogger("sd")
log_file = os.path.join(os.path.dirname(__file__), 'webui.log')
quick_allowed = True
errors = 0
opts = {}
args = Dot({
'debug': False,
'reset': False,
'upgrade': False,
'skip_update': False,
'skip_extensions': False,
'skip_requirements': False,
'skip_git': False,
'skip_torch': False,
'use_directml': False,
'use_ipex': False,
'use_cuda': False,
'use_rocm': False,
'experimental': False,
'test': False,
'tls_selfsign': False,
'reinstall': False,
'version': False,
'ignore': False,
})
git_commit = "unknown"
# setup console and file logging
def setup_logging(clean=False):
try:
if clean and os.path.isfile(log_file):
os.remove(log_file)
time.sleep(0.1) # prevent race condition
except:
pass
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 print_profile(profile: cProfile.Profile, msg: str):
try:
from rich import print # pylint: disable=redefined-builtin
except:
pass
profile.disable()
stream = io.StringIO()
ps = pstats.Stats(profile, stream=stream)
ps.sort_stats(pstats.SortKey.CUMULATIVE).print_stats(15)
profile = None
lines = stream.getvalue().split('\n')
lines = [l for l in lines if '<frozen' not in l and '{built-in' not in l and '/logging' not in l and '/rich' not in l]
print(f'Profile {msg}:', '\n'.join(lines))
# check if package is installed
def installed(package, friendly: str = None):
ok = True
try:
if friendly:
pkgs = friendly.split()
else:
pkgs = [p for p in package.split() if not p.startswith('-') and not p.startswith('=')]
pkgs = [p.split('/')[-1] for p in pkgs] # get only package name if installing from url
for pkg in pkgs:
if '>=' in pkg:
p = pkg.split('>=')
else:
p = pkg.split('==')
spec = pkg_resources.working_set.by_key.get(p[0], None) # more reliable than importlib
if spec is None:
spec = pkg_resources.working_set.by_key.get(p[0].lower(), None) # check name variations
if spec is None:
spec = pkg_resources.working_set.by_key.get(p[0].replace('_', '-'), None) # check name variations
ok = ok and spec is not None
if ok:
version = pkg_resources.get_distribution(p[0]).version
log.debug(f"Package version found: {p[0]} {version}")
if len(p) > 1:
ok = ok and version == p[1]
if not ok:
log.warning(f"Package wrong version: {p[0]} {version} required {p[1]}")
else:
log.debug(f"Package version not found: {p[0]}")
return ok
except ModuleNotFoundError:
log.debug(f"Package not installed: {pkgs}")
return False
def pip(arg: str, ignore: bool = False, quiet: bool = False):
arg = arg.replace('>=', '==')
if not quiet:
log.info(f'Installing package: {arg.replace("install", "").replace("--upgrade", "").replace("--no-deps", "").replace("--force", "").replace(" ", " ").strip()}')
log.debug(f"Running pip: {arg}")
result = subprocess.run(f'"{sys.executable}" -m pip {arg}', shell=True, check=False, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
txt = result.stdout.decode(encoding="utf8", errors="ignore")
if len(result.stderr) > 0:
txt += ('\n' if len(txt) > 0 else '') + result.stderr.decode(encoding="utf8", errors="ignore")
txt = txt.strip()
if result.returncode != 0 and not ignore:
global errors # pylint: disable=global-statement
errors += 1
log.error(f'Error running pip: {arg}')
log.debug(f'Pip output: {txt}')
return txt
# install package using pip if not already installed
def install(package, friendly: str = None, ignore: bool = False):
if args.reinstall:
global quick_allowed # pylint: disable=global-statement
quick_allowed = False
if args.use_ipex and package == "pytorch_lightning==1.9.4":
package = "pytorch_lightning==1.8.6"
if args.reinstall or not installed(package, friendly):
pip(f"install --upgrade {package}", ignore=ignore)
# execute git command
def git(arg: str, folder: str = None, ignore: bool = False):
if args.skip_git:
return ''
git_cmd = os.environ.get('GIT', "git")
result = subprocess.run(f'"{git_cmd}" {arg}', check=False, shell=True, env=os.environ, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=folder or '.')
txt = result.stdout.decode(encoding="utf8", errors="ignore")
if len(result.stderr) > 0:
txt += ('\n' if len(txt) > 0 else '') + result.stderr.decode(encoding="utf8", errors="ignore")
txt = txt.strip()
if result.returncode != 0 and not ignore:
global errors # pylint: disable=global-statement
errors += 1
log.error(f'Error running git: {folder} / {arg}')
if 'or stash them' in txt:
log.error(f'Local changes detected: check log for details: {log_file}')
log.debug(f'Git output: {txt}')
return txt
# update switch to main branch as head can get detached and update repository
def update(folder):
if not os.path.exists(os.path.join(folder, '.git')):
return
branch = git('branch', folder)
if 'main' in branch:
branch = 'main'
elif 'master' in branch:
branch = 'master'
else:
branch = branch.split('\n')[0].replace('*', '').strip()
# log.debug(f'Setting branch: {folder} / {branch}')
git(f'checkout {branch}', folder)
if branch is None:
git('pull --autostash --rebase --force', folder)
else:
git(f'pull origin {branch} --autostash --rebase --force', folder)
# branch = git('branch', folder)
# clone git repository
def clone(url, folder, commithash=None):
if os.path.exists(folder):
if args.skip_update:
return
if commithash is None:
update(folder)
else:
current_hash = git('rev-parse HEAD', folder).strip()
if current_hash != commithash:
git('fetch', folder)
git(f'checkout {commithash}', folder)
return
else:
log.info(f'Cloning repository: {url}')
git(f'clone "{url}" "{folder}"')
if commithash is not None:
git(f'-C "{folder}" checkout {commithash}')
# check python version
def check_python():
supported_minors = [9, 10]
if args.experimental:
supported_minors.append(11)
log.info(f'Python {platform.python_version()} on {platform.system()}')
if not (int(sys.version_info.major) == 3 and int(sys.version_info.minor) in supported_minors):
log.error(f"Incompatible Python version: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
if not args.ignore:
sys.exit(1)
if not args.skip_git:
git_cmd = os.environ.get('GIT', "git")
if shutil.which(git_cmd) is None:
log.error('Git not found')
if not args.ignore:
sys.exit(1)
else:
git_version = git('--version', folder=None, ignore=False)
log.debug(f'Git {git_version.replace("git version", "").strip()}')
# check torch version
def check_torch():
if args.profile:
pr = cProfile.Profile()
pr.enable()
allow_cuda = not (args.use_rocm or args.use_directml or args.use_ipex)
allow_rocm = not (args.use_cuda or args.use_directml or args.use_ipex)
allow_ipex = not (args.use_cuda or args.use_rocm or args.use_directml)
allow_directml = not (args.use_cuda or args.use_rocm or args.use_ipex)
log.debug(f'Torch overrides: cuda={args.use_cuda} rocm={args.use_rocm} ipex={args.use_ipex} diml={args.use_directml}')
log.debug(f'Torch allowed: cuda={allow_cuda} rocm={allow_rocm} ipex={allow_ipex} diml={allow_directml}')
torch_command = os.environ.get('TORCH_COMMAND', '')
if torch_command != '':
pass
elif allow_cuda and (shutil.which('nvidia-smi') is not None or os.path.exists(os.path.join(os.environ.get('SystemRoot') or r'C:\Windows', 'System32', 'nvidia-smi.exe'))):
log.info('nVidia CUDA toolkit detected')
torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/cu118')
xformers_package = os.environ.get('XFORMERS_PACKAGE', 'xformers==0.0.17' if opts.get('cross_attention_optimization', '') == 'xFormers' else 'none')
elif allow_rocm and (shutil.which('rocminfo') is not None or os.path.exists('/opt/rocm/bin/rocminfo') or os.path.exists('/dev/kfd')):
log.info('AMD ROCm toolkit detected')
os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', '10.3.0')
os.environ.setdefault('PYTORCH_HIP_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2 --index-url https://download.pytorch.org/whl/rocm5.4.2')
xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
elif allow_ipex and args.use_ipex and shutil.which('sycl-ls') is not None:
log.info('Intel OneAPI Toolkit detected')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0 torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu')
xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
else:
machine = platform.machine()
if sys.platform == 'darwin':
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.0.1 torchvision==0.15.2')
elif allow_directml and args.use_directml and ('arm' not in machine and 'aarch' not in machine):
log.info('Using DirectML Backend')
torch_command = os.environ.get('TORCH_COMMAND', 'torch-directml')
xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
if 'torch' in torch_command and not args.version:
install(torch_command, 'torch torchvision')
else:
log.info('Using CPU-only Torch')
torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none')
if 'torch' in torch_command and not args.version:
install(torch_command, 'torch torchvision')
if args.skip_torch:
log.info('Skipping Torch tests')
else:
try:
import torch
log.info(f'Torch {torch.__version__}')
if args.use_ipex and allow_ipex:
import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
log.info(f'Torch backend: Intel IPEX {ipex.__version__}')
log.info(f'{os.popen("icpx --version").read().rstrip()}')
for device in [torch.xpu.device(i) for i in range(torch.xpu.device_count())]:
log.info(f'Torch detected GPU: {torch.xpu.get_device_name(device)} VRAM {round(torch.xpu.get_device_properties(device).total_memory / 1024 / 1024)}')
elif torch.cuda.is_available() and (allow_cuda or allow_rocm):
# log.debug(f'Torch allocator: {torch.cuda.get_allocator_backend()}')
if torch.version.cuda and allow_cuda:
log.info(f'Torch backend: nVidia CUDA {torch.version.cuda} cuDNN {torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else "N/A"}')
elif torch.version.hip and allow_rocm:
log.info(f'Torch backend: AMD ROCm HIP {torch.version.hip}')
else:
log.warning('Unknown Torch backend')
for device in [torch.cuda.device(i) for i in range(torch.cuda.device_count())]:
log.info(f'Torch detected GPU: {torch.cuda.get_device_name(device)} VRAM {round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024)} Arch {torch.cuda.get_device_capability(device)} Cores {torch.cuda.get_device_properties(device).multi_processor_count}')
else:
try:
if args.use_directml and allow_directml:
import torch_directml # pylint: disable=import-error
version = pkg_resources.get_distribution("torch-directml")
log.info(f'Torch backend: DirectML ({version})')
for i in range(0, torch_directml.device_count()):
log.info(f'Torch detected GPU: {torch_directml.device_name(i)}')
except:
log.warning("Torch reports CUDA not available")
except Exception as e:
log.error(f'Could not load torch: {e}')
if not args.ignore:
sys.exit(1)
if args.version:
return
try:
if 'xformers' in xformers_package:
install(f'--no-deps {xformers_package}', ignore=True)
else:
x = pkg_resources.working_set.by_key.get('xformers', None)
if x is not None:
log.warning(f'Not used, uninstalling: {x}')
pip('uninstall xformers --yes --quiet', ignore=True, quiet=True)
except Exception as e:
log.debug(f'Cannot install xformers package: {e}')
try:
tensorflow_package = os.environ.get('TENSORFLOW_PACKAGE', 'tensorflow==2.12.0')
install(tensorflow_package, 'tensorflow', ignore=True)
except Exception as e:
log.debug(f'Cannot install tensorflow package: {e}')
if opts.get('cuda_compile_mode', '') == 'hidet':
install('hidet', 'hidet')
if args.profile:
print_profile(pr, 'Torch')
# install required packages
def install_packages():
if args.profile:
pr = cProfile.Profile()
pr.enable()
log.info('Installing packages')
# gfpgan_package = os.environ.get('GFPGAN_PACKAGE', "git+https://github.com/TencentARC/GFPGAN.git@8d2447a2d918f8eba5a4a01463fd48e45126a379")
# openclip_package = os.environ.get('OPENCLIP_PACKAGE', "git+https://github.com/mlfoundations/open_clip.git@bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b")
# install(gfpgan_package, 'gfpgan')
# install(openclip_package, 'open-clip-torch')
clip_package = os.environ.get('CLIP_PACKAGE', "git+https://github.com/openai/CLIP.git")
install(clip_package, 'clip')
install('onnxruntime==1.14.0', 'onnxruntime', ignore=True)
if args.profile:
print_profile(pr, 'Packages')
# clone required repositories
def install_repositories():
if args.profile:
pr = cProfile.Profile()
pr.enable()
def d(name):
return os.path.join(os.path.dirname(__file__), 'repositories', name)
log.info('Installing repositories')
os.makedirs(os.path.join(os.path.dirname(__file__), 'repositories'), exist_ok=True)
stable_diffusion_repo = os.environ.get('STABLE_DIFFUSION_REPO', "https://github.com/Stability-AI/stablediffusion.git")
# stable_diffusion_commit = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', "cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf")
stable_diffusion_commit = os.environ.get('STABLE_DIFFUSION_COMMIT_HASH', None)
clone(stable_diffusion_repo, d('stable-diffusion-stability-ai'), stable_diffusion_commit)
taming_transformers_repo = os.environ.get('TAMING_TRANSFORMERS_REPO', "https://github.com/CompVis/taming-transformers.git")
# taming_transformers_commit = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', "3ba01b241669f5ade541ce990f7650a3b8f65318")
taming_transformers_commit = os.environ.get('TAMING_TRANSFORMERS_COMMIT_HASH', None)
clone(taming_transformers_repo, d('taming-transformers'), taming_transformers_commit)
k_diffusion_repo = os.environ.get('K_DIFFUSION_REPO', 'https://github.com/crowsonkb/k-diffusion.git')
# k_diffusion_commit = os.environ.get('K_DIFFUSION_COMMIT_HASH', "b43db16749d51055f813255eea2fdf1def801919")
k_diffusion_commit = os.environ.get('K_DIFFUSION_COMMIT_HASH', None)
clone(k_diffusion_repo, d('k-diffusion'), k_diffusion_commit)
codeformer_repo = os.environ.get('CODEFORMER_REPO', 'https://github.com/sczhou/CodeFormer.git')
# codeformer_commit = os.environ.get('CODEFORMER_COMMIT_HASH', "c5b4593074ba6214284d6acd5f1719b6c5d739af")
codeformer_commit = os.environ.get('CODEFORMER_COMMIT_HASH', "7a584fd")
clone(codeformer_repo, d('CodeFormer'), codeformer_commit)
blip_repo = os.environ.get('BLIP_REPO', 'https://github.com/salesforce/BLIP.git')
# blip_commit = os.environ.get('BLIP_COMMIT_HASH', "48211a1594f1321b00f14c9f7a5b4813144b2fb9")
blip_commit = os.environ.get('BLIP_COMMIT_HASH', None)
clone(blip_repo, d('BLIP'), blip_commit)
if args.profile:
print_profile(pr, 'Repositories')
# run extension installer
def run_extension_installer(folder):
path_installer = os.path.join(folder, "install.py")
if not os.path.isfile(path_installer):
return
try:
log.debug(f"Running extension installer: {path_installer}")
env = os.environ.copy()
env['PYTHONPATH'] = os.path.abspath(".")
result = subprocess.run(f'"{sys.executable}" "{path_installer}"', shell=True, env=env, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=folder)
if result.returncode != 0:
global errors # pylint: disable=global-statement
errors += 1
txt = result.stdout.decode(encoding="utf8", errors="ignore")
if len(result.stderr) > 0:
txt = txt + '\n' + result.stderr.decode(encoding="utf8", errors="ignore")
log.error(f'Error running extension installer: {path_installer}')
log.debug(txt)
except Exception as e:
log.error(f'Exception running extension installer: {e}')
# get list of all enabled extensions
def list_extensions(folder):
disabled_extensions = opts.get('disable_all_extensions', 'none')
if disabled_extensions != 'none':
log.debug(f'Disabled extensions: {disabled_extensions}')
return []
disabled_extensions = set(opts.get('disabled_extensions', []))
if len(disabled_extensions) > 0:
log.debug(f'Disabled extensions: {disabled_extensions}')
return [x for x in os.listdir(folder) if x not in disabled_extensions and not x.startswith('.')]
# run installer for each installed and enabled extension and optionally update them
def install_extensions():
if args.profile:
pr = cProfile.Profile()
pr.enable()
pkg_resources._initialize_master_working_set() # pylint: disable=protected-access
pkgs = [f'{p.project_name}=={p._version}' for p in pkg_resources.working_set] # pylint: disable=protected-access,not-an-iterable
log.debug(f'Installed packages: {len(pkgs)}')
from modules.paths_internal import extensions_builtin_dir, extensions_dir
extensions_duplicates = []
extensions_enabled = []
extension_folders = [extensions_builtin_dir] if args.safe else [extensions_builtin_dir, extensions_dir]
for folder in extension_folders:
if not os.path.isdir(folder):
continue
extensions = list_extensions(folder)
log.debug(f'Extensions all: {extensions}')
for ext in extensions:
if ext in extensions_enabled:
extensions_duplicates.append(ext)
continue
extensions_enabled.append(ext)
if not args.skip_update:
try:
update(os.path.join(folder, ext))
except:
log.error(f'Error updating extension: {os.path.join(folder, ext)}')
if not args.skip_extensions:
run_extension_installer(os.path.join(folder, ext))
pkg_resources._initialize_master_working_set() # pylint: disable=protected-access
updated = [f'{p.project_name}=={p._version}' for p in pkg_resources.working_set] # pylint: disable=protected-access,not-an-iterable
diff = [x for x in updated if x not in pkgs]
pkgs = updated
if len(diff) > 0:
log.info(f'Extension installed packages: {ext} {diff}')
log.info(f'Extensions enabled: {extensions_enabled}')
if len(extensions_duplicates) > 0:
log.warning(f'Extensions duplicates: {extensions_duplicates}')
if args.profile:
print_profile(pr, 'Extensions')
# initialize and optionally update submodules
def install_submodules():
if args.profile:
pr = cProfile.Profile()
pr.enable()
log.info('Installing submodules')
txt = git('submodule')
log.debug(f'Submodules list: {txt}')
if 'no submodule mapping found' in txt:
log.warning('Attempting repository recover')
git('add .')
git('stash')
git('merge --abort', folder=None, ignore=True)
git('fetch --all')
git('reset --hard origin/master')
git('checkout master')
txt = git('submodule')
log.info('Continuing setup')
git('submodule --quiet update --init --recursive')
if not args.skip_update:
log.info('Updating submodules')
submodules = txt.splitlines()
for submodule in submodules:
try:
name = submodule.split()[1].strip()
update(name)
except:
log.error(f'Error updating submodule: {submodule}')
if args.profile:
print_profile(pr, 'Submodule')
def ensure_base_requirements():
try:
import rich # pylint: disable=unused-import
except ImportError:
install('rich', 'rich')
def install_requirements():
if args.profile:
pr = cProfile.Profile()
pr.enable()
if args.skip_requirements:
return
log.info('Verifying requirements')
with open('requirements.txt', 'r', encoding='utf8') as f:
lines = [line.strip() for line in f.readlines() if line.strip() != '' and not line.startswith('#') and line is not None]
for line in lines:
install(line)
if args.profile:
print_profile(pr, 'Requirements')
# set environment variables controling the behavior of various libraries
def set_environment():
log.info('Setting environment tuning')
os.environ.setdefault('USE_TORCH', '1')
os.environ.setdefault('TF_CPP_MIN_LOG_LEVEL', '2')
os.environ.setdefault('ACCELERATE', 'True')
os.environ.setdefault('FORCE_CUDA', '1')
os.environ.setdefault('ATTN_PRECISION', 'fp16')
os.environ.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'garbage_collection_threshold:0.8,max_split_size_mb:512')
os.environ.setdefault('CUDA_LAUNCH_BLOCKING', '0')
os.environ.setdefault('CUDA_CACHE_DISABLE', '0')
os.environ.setdefault('CUDA_AUTO_BOOST', '1')
os.environ.setdefault('CUDA_MODULE_LOADING', 'LAZY')
os.environ.setdefault('CUDA_DEVICE_DEFAULT_PERSISTING_L2_CACHE_PERCENTAGE_LIMIT', '0')
os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
os.environ.setdefault('SAFETENSORS_FAST_GPU', '1')
os.environ.setdefault('NUMEXPR_MAX_THREADS', '16')
os.environ.setdefault('PYTHONHTTPSVERIFY', '0')
os.environ.setdefault('HF_HUB_DISABLE_TELEMETRY', '1')
if sys.platform == 'darwin':
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK', '1')
def check_extensions():
newest_all = os.path.getmtime('requirements.txt')
from modules.paths_internal import extensions_builtin_dir, extensions_dir
extension_folders = [extensions_builtin_dir] if args.safe else [extensions_builtin_dir, extensions_dir]
for folder in extension_folders:
if not os.path.isdir(folder):
continue
extensions = list_extensions(folder)
for ext in extensions:
newest = 0
extension_dir = os.path.join(folder, ext)
if not os.path.isdir(extension_dir):
log.debug(f'Extension listed as installed but folder missing: {extension_dir}')
continue
for f in os.listdir(extension_dir):
if '.json' in f or '.csv' in f or '__pycache__' in f:
continue
ts = os.path.getmtime(os.path.join(extension_dir, f))
newest = max(newest, ts)
newest_all = max(newest_all, newest)
log.debug(f'Extension version: {time.ctime(newest)} {folder}{os.pathsep}{ext}')
return round(newest_all)
# check version of the main repo and optionally upgrade it
def check_version(offline=False, reset=True): # pylint: disable=unused-argument
if not os.path.exists('.git'):
log.error('Not a git repository')
if not args.ignore:
sys.exit(1)
# status = git('status')
# if 'branch' not in status:
# log.error('Cannot get git repository status')
# sys.exit(1)
ver = git('log -1 --pretty=format:"%h %ad"')
log.info(f'Version: {ver}')
if args.version:
return
commit = git('rev-parse HEAD')
global git_commit # pylint: disable=global-statement
git_commit = commit[:7]
try:
import requests
except ImportError:
return
commits = None
try:
commits = requests.get('https://api.github.com/repos/vladmandic/automatic/branches/master', timeout=10).json()
if commits['commit']['sha'] != commit:
if args.upgrade:
global quick_allowed # pylint: disable=global-statement
quick_allowed = False
log.info('Updating main repository')
try:
git('add .')
git('stash')
update('.')
# git('git stash pop')
ver = git('log -1 --pretty=format:"%h %ad"')
log.info(f'Upgraded to version: {ver}')
except Exception:
if not reset:
log.error('Error during repository upgrade')
else:
log.warning('Retrying repository upgrade...')
git_reset()
check_version(offline=offline, reset=False)
else:
log.info(f'Latest published version: {commits["commit"]["sha"]} {commits["commit"]["commit"]["author"]["date"]}')
except Exception as e:
log.error(f'Failed to check version: {e} {commits}')
def update_wiki():
if not args.skip_update:
log.info('Updating Wiki')
try:
update(os.path.join(os.path.dirname(__file__), "wiki"))
update(os.path.join(os.path.dirname(__file__), "wiki", "origin-wiki"))
except:
log.error('Error updating wiki')
# check if we can run setup in quick mode
def check_timestamp():
if not quick_allowed or not os.path.isfile(log_file):
return False
if args.skip_git:
return True
ok = True
setup_time = -1
with open(log_file, 'r', encoding='utf8') as f:
lines = f.readlines()
for line in lines:
if 'Setup complete without errors' in line:
setup_time = int(line.split(' ')[-1])
try:
version_time = int(git('log -1 --pretty=format:"%at"'))
except Exception as e:
log.error(f'Error getting local repository version: {e}')
if not args.ignore:
sys.exit(1)
log.debug(f'Repository update time: {time.ctime(int(version_time))}')
if setup_time == -1:
return False
log.debug(f'Previous setup time: {time.ctime(setup_time)}')
if setup_time < version_time:
ok = False
extension_time = check_extensions()
log.debug(f'Latest extensions time: {time.ctime(extension_time)}')
if setup_time < extension_time:
ok = False
log.debug(f'Timestamps: version:{version_time} setup:{setup_time} extension:{extension_time}')
if args.reinstall:
ok = False
return ok
def add_args():
group = parser.add_argument_group('Setup options')
group.add_argument('--debug', default = False, action='store_true', help = "Run installer with debug logging, default: %(default)s")
group.add_argument('--reset', default = False, action='store_true', help = "Reset main repository to latest version, default: %(default)s")
group.add_argument('--upgrade', default = False, action='store_true', help = "Upgrade main repository to latest version, default: %(default)s")
group.add_argument("--use-ipex", default = False, action='store_true', help="Use Intel OneAPI XPU backend, default: %(default)s")
group.add_argument('--use-directml', default = False, action='store_true', help = "Use DirectML if no compatible GPU is detected, default: %(default)s")
group.add_argument("--use-cuda", default=False, action='store_true', help="Force use nVidia CUDA backend, default: %(default)s")
group.add_argument("--use-rocm", default=False, action='store_true', help="Force use AMD ROCm backend, default: %(default)s")
group.add_argument('--skip-update', default = False, action='store_true', help = "Skip update of extensions and submodules, default: %(default)s")
group.add_argument('--skip-requirements', default = False, action='store_true', help = "Skips checking and installing requirements, default: %(default)s")
group.add_argument('--skip-extensions', default = False, action='store_true', help = "Skips running individual extension installers, default: %(default)s")
group.add_argument('--skip-git', default = False, action='store_true', help = "Skips running all GIT operations, default: %(default)s")
group.add_argument('--skip-torch', default = False, action='store_true', help = "Skips running Torch checks, default: %(default)s")
group.add_argument('--experimental', default = False, action='store_true', help = "Allow unsupported versions of libraries, default: %(default)s")
group.add_argument('--reinstall', default = False, action='store_true', help = "Force reinstallation of all requirements, default: %(default)s")
group.add_argument('--test', default = False, action='store_true', help = "Run test only and exit")
group.add_argument('--version', default = False, action='store_true', help = "Print version information")
group.add_argument('--ignore', default = False, action='store_true', help = "Ignore any errors and attempt to continue")
group.add_argument('--safe', default = False, action='store_true', help = "Run in safe mode with no user extensions")
def parse_args():
# command line args
global args # pylint: disable=global-statement
args = parser.parse_args()
def extensions_preload(force = False):
if args.profile:
pr = cProfile.Profile()
pr.enable()
setup_time = 0
if not force:
if os.path.isfile(log_file):
with open(log_file, 'r', encoding='utf8') as f:
lines = f.readlines()
for line in lines:
if 'Setup complete without errors' in line:
setup_time = int(line.split(' ')[-1])
if setup_time > 0 or force:
log.info('Running extension preloading')
if args.safe:
log.info('Running in safe mode without user extensions')
try:
from modules.script_loading import preload_extensions
from modules.paths_internal import extensions_builtin_dir, extensions_dir
extension_folders = [extensions_builtin_dir] if args.safe else [extensions_builtin_dir, extensions_dir]
for ext_dir in extension_folders:
t0 = time.time()
preload_extensions(ext_dir, parser)
t1 = time.time()
log.debug(f'Extension preload: {round(t1 - t0, 1)}s {ext_dir}')
except:
log.error('Error running extension preloading')
if args.profile:
print_profile(pr, 'Preload')
def git_reset():
log.warning('Running GIT reset')
global quick_allowed # pylint: disable=global-statement
quick_allowed = False
git('merge --abort')
git('fetch --all')
git('reset --hard origin/master')
git('checkout master')
log.info('GIT reset complete')
def read_options():
global opts # pylint: disable=global-statement
if os.path.isfile(args.config):
with open(args.config, "r", encoding="utf8") as file:
opts = json.load(file)
# entry method when used as module
def run_setup():
# setup_logging(args.upgrade)
log.info('Starting SD.Next')
read_options()
check_python()
if args.reset:
git_reset()
if args.skip_git:
log.info('Skipping GIT operations')
check_version()
set_environment()
if args.reinstall:
log.info('Forcing reinstall of all packages')
check_torch()
install_requirements()
install_packages()
if check_timestamp():
log.info('No changes detected: Quick launch active')
return
log.info("Running setup")
log.debug(f"Args: {vars(args)}")
install_repositories()
install_submodules()
install_extensions()
update_wiki()
if errors == 0:
log.debug(f'Setup complete without errors: {round(time.time())}')
else:
log.warning(f'Setup complete with errors: {errors}')
log.warning(f'See log file for more details: {log_file}')
if __name__ == "__main__":
add_args()
ensure_base_requirements()
parse_args()
run_setup()

Some files were not shown because too many files have changed in this diff Show More