changes to dogettx optimizations to run on m1

* Author @any-winter-4079
* Author @dogettx
Thanks to many individuals who contributed time and hardware to
benchmarking and debugging these changes.

.github/workflows/cache-model.yml

@@ -0,0 +1,64 @@
+name: Cache Model
+on:
+  workflow_dispatch
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ macos-12 ]
+    name: Create Caches using ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v3
+      - name: Cache model
+        id: cache-sd-v1-4
+        uses: actions/cache@v3
+        env:
+          cache-name: cache-sd-v1-4
+        with:
+          path: models/ldm/stable-diffusion-v1/model.ckpt
+          key: ${{ env.cache-name }}
+          restore-keys: |
+            ${{ env.cache-name }}
+      - name: Download Stable Diffusion v1.4 model
+        if: ${{ steps.cache-sd-v1-4.outputs.cache-hit != 'true' }}
+        continue-on-error: true
+        run: |
+          if [ ! -e models/ldm/stable-diffusion-v1 ]; then
+              mkdir -p models/ldm/stable-diffusion-v1
+          fi
+          if [ ! -e models/ldm/stable-diffusion-v1/model.ckpt ]; then
+            curl -o models/ldm/stable-diffusion-v1/model.ckpt ${{ secrets.SD_V1_4_URL }}
+          fi
+# Uncomment this when we no longer make changes to environment-mac.yaml
+#       - name: Cache environment
+#         id: cache-conda-env-ldm
+#         uses: actions/cache@v3
+#         env:
+#           cache-name: cache-conda-env-ldm
+#         with:
+#           path: ~/.conda/envs/ldm
+#           key: ${{ env.cache-name }}
+#           restore-keys: |
+#             ${{ env.cache-name }}
+      - name: Install dependencies
+#         if: ${{ steps.cache-conda-env-ldm.outputs.cache-hit != 'true' }}
+        run: |
+          conda env create -f environment-mac.yaml
+      - name: Cache hugginface and torch models
+        id: cache-hugginface-torch
+        uses: actions/cache@v3
+        env:
+          cache-name: cache-hugginface-torch
+        with:
+          path: ~/.cache
+          key: ${{ env.cache-name }}
+          restore-keys: |
+            ${{ env.cache-name }}
+      - name: Download Huggingface and Torch models
+        if: ${{ steps.cache-hugginface-torch.outputs.cache-hit != 'true' }}
+        continue-on-error: true
+        run: |
+          export PYTHON_BIN=/usr/local/miniconda/envs/ldm/bin/python
+          $PYTHON_BIN scripts/preload_models.py

.github/workflows/macos12-miniconda.yml

@@ -0,0 +1,80 @@
+name: Build
+on: 
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ macos-12 ]
+    name: Build on ${{ matrix.os }} miniconda
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v3
+      - name: Cache model
+        id: cache-sd-v1-4
+        uses: actions/cache@v3
+        env:
+          cache-name: cache-sd-v1-4
+        with:
+          path: models/ldm/stable-diffusion-v1/model.ckpt
+          key: ${{ env.cache-name }}
+          restore-keys: |
+            ${{ env.cache-name }}
+      - name: Download Stable Diffusion v1.4 model
+        if: ${{ steps.cache-sd-v1-4.outputs.cache-hit != 'true' }}
+        continue-on-error: true
+        run: |
+          if [ ! -e models/ldm/stable-diffusion-v1 ]; then
+              mkdir -p models/ldm/stable-diffusion-v1
+          fi
+          if [ ! -e models/ldm/stable-diffusion-v1/model.ckpt ]; then
+            curl -o models/ldm/stable-diffusion-v1/model.ckpt ${{ secrets.SD_V1_4_URL }}
+          fi
+# Uncomment this when we no longer make changes to environment-mac.yaml
+#       - name: Cache environment
+#         id: cache-conda-env-ldm
+#         uses: actions/cache@v3
+#         env:
+#           cache-name: cache-conda-env-ldm
+#         with:
+#           path: ~/.conda/envs/ldm
+#           key: ${{ env.cache-name }}
+#           restore-keys: |
+#             ${{ env.cache-name }}
+      - name: Install dependencies
+#         if: ${{ steps.cache-conda-env-ldm.outputs.cache-hit != 'true' }}
+        run: |
+          conda env create -f environment-mac.yaml
+      - name: Cache hugginface and torch models
+        id: cache-hugginface-torch
+        uses: actions/cache@v3
+        env:
+          cache-name: cache-hugginface-torch
+        with:
+          path: ~/.cache
+          key: ${{ env.cache-name }}
+          restore-keys: |
+            ${{ env.cache-name }}
+      - name: Download Huggingface and Torch models
+        if: ${{ steps.cache-hugginface-torch.outputs.cache-hit != 'true' }}
+        continue-on-error: true
+        run: |
+          export PYTHON_BIN=/usr/local/miniconda/envs/ldm/bin/python
+          $PYTHON_BIN scripts/preload_models.py
+      - name: Run the tests
+        run: |
+          # Note, can't "activate" via automation, and activation is just env vars and path
+          export PYTHON_BIN=/usr/local/miniconda/envs/ldm/bin/python
+          export PYTORCH_ENABLE_MPS_FALLBACK=1
+          $PYTHON_BIN scripts/preload_models.py
+          mkdir -p outputs/img-samples
+          time $PYTHON_BIN scripts/dream.py --from_file tests/prompts.txt </dev/null 2> outputs/img-samples/err.log > outputs/img-samples/out.log
+      - name: Archive results
+        uses: actions/upload-artifact@v3
+        with:
+          name: results
+          path: outputs/img-samples

.gitignore

@@ -77,6 +77,9 @@ db.sqlite3-journal
 instance/
 .webassets-cache
 
+# WebUI temp files:
+img2img-tmp.png
+
 # Scrapy stuff:
 .scrapy
 

README-Mac-MPS.md

@@ -32,6 +32,23 @@ While that is downloading, open Terminal and run the following commands one at a
 # install brew (and Xcode command line tools):
 /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
 
+#
+# Now there are two different routes to get the Python (miniconda) environment up and running:
+# 1. Alongside pyenv
+# 2. No pyenv
+#
+# If you don't know what we are talking about, choose 2.
+# 
+# NOW EITHER DO
+# 1. Installing alongside pyenv 
+
+brew install pyenv-virtualenv # you might have this from before, no problem
+pyenv install anaconda3-latest
+pyenv virtualenv anaconda3-latest lstein-stable-diffusion
+pyenv activate lstein-stable-diffusion
+
+# OR, 
+# 2. Installing standalone
 # install python 3, git, cmake, protobuf:
 brew install cmake protobuf rust
 
@@ -39,6 +56,10 @@ brew install cmake protobuf rust
 curl https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -o Miniconda3-latest-MacOSX-arm64.sh
 /bin/bash Miniconda3-latest-MacOSX-arm64.sh
 
+
+# EITHER WAY,
+# continue from here
+
 # clone the repo
 git clone https://github.com/lstein/stable-diffusion.git
 cd stable-diffusion
@@ -320,3 +341,20 @@ something that depends on it-- Rosetta can translate some Intel instructions but
 not the specialized ones here. To avoid this, make sure to use the environment
 variable `CONDA_SUBDIR=osx-arm64`, which restricts the Conda environment to only
 use ARM packages, and use `nomkl` as described above.
+
+### input types 'tensor<2x1280xf32>' and 'tensor<*xf16>' are not broadcast compatible
+
+May appear when just starting to generate, e.g.:
+
+```
+dream> clouds
+Generating:   0%|                                                              | 0/1 [00:00<?, ?it/s]/Users/[...]/dev/stable-diffusion/ldm/modules/embedding_manager.py:152: UserWarning: The operator 'aten::nonzero' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1662016319283/work/aten/src/ATen/mps/MPSFallback.mm:11.)
+  placeholder_idx = torch.where(
+                                                                                                    loc("mps_add"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/20d6c351-ee94-11ec-bcaf-7247572f23b4/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":219:0)): error: input types 'tensor<2x1280xf32>' and 'tensor<*xf16>' are not broadcast compatible
+LLVM ERROR: Failed to infer result type(s).
+Abort trap: 6
+/Users/[...]/opt/anaconda3/envs/ldm/lib/python3.9/multiprocessing/resource_tracker.py:216: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
+  warnings.warn('resource_tracker: There appear to be %d '
+  ```
+
+Macs do not support autocast/mixed-precision. Supply `--full_precision` to use float32 everywhere.

README.md

@@ -1,7 +1,7 @@
 <h1 align='center'><b>Stable Diffusion Dream Script</b></h1>
 
 <p align='center'>
-<img src="static/logo_temp.png"/>
+<img src="static/logo.png"/>
 </p>
 
 <p align="center">
@@ -22,22 +22,24 @@ text-to-image generator. This fork supports:
    generating images in your browser.
 
 3. Support for img2img in which you provide a seed image to guide the
-      image creation. (inpainting & masking coming soon)
+   image creation. (inpainting & masking coming soon)
 
-4. A notebook for running the code on Google Colab.
+4. Preliminary inpainting support.
 
-5. Upscaling and face fixing using the optional ESRGAN and GFPGAN
+5. A notebook for running the code on Google Colab.
+
+6. Upscaling and face fixing using the optional ESRGAN and GFPGAN
    packages.
 
-6. Weighted subprompts for prompt tuning.
+7. Weighted subprompts for prompt tuning.
 
 7. [Image variations](VARIATIONS.md) which allow you to systematically
-generate variations of an image you like and combine two or more
-images together to combine the best features of both.
+   generate variations of an image you like and combine two or more
+   images together to combine the best features of both.
 
-8. Textual inversion for customization of the prompt language and images.
+9. Textual inversion for customization of the prompt language and images.
 
-8. ...and more!
+10. ...and more!
 
 This fork is rapidly evolving, so use the Issues panel to report bugs
 and make feature requests, and check back periodically for
@@ -75,9 +77,10 @@ log file of image names and prompts to the selected output directory.
 In addition, as of version 1.02, it also writes the prompt into the PNG
 file's metadata where it can be retrieved using scripts/images2prompt.py
 
-The script is confirmed to work on Linux and Windows systems. It should
-work on MacOSX as well, but this is not confirmed. Note that this script
-runs from the command-line (CMD or Terminal window), and does not have a GUI.
+The script is confirmed to work on Linux, Windows and Mac
+systems. Note that this script runs from the command-line or can be used
+as a Web application. The Web GUI is currently rudimentary, but a much
+better replacement is on its way.
 
 ```
 (ldm) ~/stable-diffusion$ python3 ./scripts/dream.py
@@ -97,7 +100,7 @@ dream> "there's a fly in my soup" -n6 -g
 dream> q
 
 # this shows how to retrieve the prompt stored in the saved image's metadata
-(ldm) ~/stable-diffusion$ python3 ./scripts/images2prompt.py outputs/img_samples/*.png
+(ldm) ~/stable-diffusion$ python ./scripts/images2prompt.py outputs/img_samples/*.png
 00009.png: "ashley judd riding a camel" -s150 -S 416354203
 00010.png: "ashley judd riding a camel" -s150 -S 1362479620
 00011.png: "there's a fly in my soup" -n6 -g -S 2685670268
@@ -118,26 +121,72 @@ The script itself also recognizes a series of command-line switches
 that will change important global defaults, such as the directory for
 image outputs and the location of the model weight files.
 
+## Hardware Requirements
+
+You will need one of:
+
+1. An NVIDIA-based graphics card with 8 GB or more of VRAM memory*.
+
+2. An Apple computer with an M1 chip.**
+
+3. At least 12 GB of main memory RAM.
+
+4. At least 6 GB of free disk space for the machine learning model,
+python, and all its dependencies.
+
+* If you are have a Nvidia 10xx series card (e.g. the 1080ti), please
+run the dream script in full-precision mode as shown below.
+
+** Similarly, specify full-precision mode on Apple M1 hardware.
+
+To run in full-precision mode, start dream.py with the
+--full_precision flag:
+
+~~~~
+(ldm) ~/stable-diffusion$ python scripts/dream.py --full_precision
+~~~~
+
 ## Image-to-Image
 
 This script also provides an img2img feature that lets you seed your
-creations with a drawing or photo. This is a really cool feature that tells
-stable diffusion to build the prompt on top of the image you provide, preserving
-the original's basic shape and layout. To use it, provide the --init_img
-option as shown here:
+creations with an initial drawing or photo. This is a really cool
+feature that tells stable diffusion to build the prompt on top of the
+image you provide, preserving the original's basic shape and
+layout. To use it, provide the --init_img option as shown here:
 
 ```
 dream> "waterfall and rainbow" --init_img=./init-images/crude_drawing.png --strength=0.5 -s100 -n4
 ```
 
-The --init_img (-I) option gives the path to the seed picture. --strength (-f) controls how much
-the original will be modified, ranging from 0.0 (keep the original intact), to 1.0 (ignore the original
-completely). The default is 0.75, and ranges from 0.25-0.75 give interesting results.
+The --init_img (-I) option gives the path to the seed
+picture. --strength (-f) controls how much the original will be
+modified, ranging from 0.0 (keep the original intact), to 1.0 (ignore
+the original completely). The default is 0.75, and ranges from
+0.25-0.75 give interesting results.
 
-You may also pass a -v<count> option to generate count variants on the original image. This is done by
-passing the first generated image back into img2img the requested number of times. It generates interesting
+You may also pass a -v<count> option to generate count variants on the
+original image. This is done by passing the first generated image back
+into img2img the requested number of times. It generates interesting
 variants.
 
+If the initial image contains transparent regions, then Stable
+Diffusion will only draw within the transparent regions, a process
+called "inpainting". However, for this to work correctly, the color
+information underneath the transparent needs to be preserved, not
+erased. See [Creating Transparent Images for
+Inpainting](#creating-transparent-images-for-inpainting) for details.
+
+## Seamless Tiling
+
+The seamless tiling mode causes generated images to seamlessly tile
+with itself. To use it, add the --seamless option when starting the
+script which will result in all generated images to tile, or for each
+dream> prompt as shown here:
+
+```
+dream> "pond garden with lotus by claude monet" --seamless -s100 -n4
+```
+
 ## GFPGAN and Real-ESRGAN Support
 
 The script also provides the ability to do face restoration and
@@ -338,9 +387,8 @@ and introducing a new vocabulary to the fixed model.
 
 To train, prepare a folder that contains images sized at 512x512 and execute the following:
 
-
 WINDOWS: As the default backend is not available on Windows, if you're using that platform, set the environment variable `PL_TORCH_DISTRIBUTED_BACKEND=gloo`
-    
+
 ```
 (ldm) ~/stable-diffusion$ python3 ./main.py --base ./configs/stable-diffusion/v1-finetune.yaml \
                                             -t \
@@ -400,11 +448,18 @@ repository and associated paper for details and limitations.
 
 # Latest Changes
 
-- v1.13 (in process)
+- v1.14 (In progress)
 
+  - Add "seamless mode" for circular tiling of image. Generates beautiful effects. ([prixt](https://github.com/prixt))
+
+- v1.13 (3 September 2022
+
+  - Support image variations (see [VARIATIONS](VARIATIONS.md) ([Kevin Gibbons](https://github.com/bakkot) and many contributors and reviewers)
   - Supports a Google Colab notebook for a standalone server running on Google hardware [Arturo Mendivil](https://github.com/artmen1516)
   - WebUI supports GFPGAN/ESRGAN facial reconstruction and upscaling [Kevin Gibbons](https://github.com/bakkot)
   - WebUI supports incremental display of in-progress images during generation [Kevin Gibbons](https://github.com/bakkot)
+  - A new configuration file scheme that allows new models (including upcoming stable-diffusion-v1.5)
+    to be added without altering the code. ([David Wager](https://github.com/maddavid12))
   - Can specify --grid on dream.py command line as the default.
   - Miscellaneous internal bug and stability fixes.
   - Works on M1 Apple hardware.
@@ -424,10 +479,12 @@ There are separate installation walkthroughs for [Linux](#linux), [Windows](#win
 - Python (version 3.8.5 recommended; higher may work)
 - git
 
-2. Install the Python Anaconda environment manager using pip3.
+2. Install the Python Anaconda environment manager.
 
 ```
-~$ pip3 install anaconda
+~$  wget https://repo.anaconda.com/archive/Anaconda3-2022.05-Linux-x86_64.sh
+~$  chmod +x Anaconda3-2022.05-Linux-x86_64.sh
+~$  ./Anaconda3-2022.05-Linux-x86_64.sh
 ```
 
 After installing anaconda, you should log out of your system and log back in. If the installation
@@ -616,9 +673,9 @@ python scripts\dream.py
 ```
 
 10. Subsequently, to relaunch the script, first activate the Anaconda
-command window (step 3), enter the stable-diffusion directory (step 5,
-"cd \path\to\stable-diffusion"), run "conda activate ldm" (step 6b),
-and then launch the dream script (step 9).
+    command window (step 3), enter the stable-diffusion directory (step 5,
+    "cd \path\to\stable-diffusion"), run "conda activate ldm" (step 6b),
+    and then launch the dream script (step 9).
 
 **Note:** Tildebyte has written an alternative ["Easy peasy Windows
 install"](https://github.com/lstein/stable-diffusion/wiki/Easy-peasy-Windows-install)
@@ -758,6 +815,49 @@ of branch>
 You will need to go through the install procedure again, but it should
 be fast because all the dependencies are already loaded.
 
+# Creating Transparent Regions for Inpainting
+
+Inpainting is really cool. To do it, you start with an initial image
+and use a photoeditor to make one or more regions transparent
+(i.e. they have a "hole" in them). You then provide the path to this
+image at the dream> command line using the -I switch. Stable Diffusion
+will only paint within the transparent region.
+
+There's a catch. In the current implementation, you have to prepare
+the initial image correctly so that the underlying colors are
+preserved under the transparent area. Many imaging editing
+applications will by default erase the color information under the
+transparent pixels and replace them with white or black, which will
+lead to suboptimal inpainting. You also must take care to export the
+PNG file in such a way that the color information is preserved.
+
+If your photoeditor is erasing the underlying color information,
+dream.py will give you a big fat warning. If you can't find a way to
+coax your photoeditor to retain color values under transparent areas,
+then you can combine the -I and -M switches to provide both the
+original unedited image and the masked (partially transparent) image:
+
+~~~~
+dream> man with cat on shoulder -I./images/man.png -M./images/man-transparent.png
+~~~~
+
+We are hoping to get rid of the need for this workaround in an
+upcoming release.
+
+## Recipe for GIMP
+
+GIMP is a popular Linux photoediting tool.
+
+1. Open image in GIMP.
+2. Layer->Transparency->Add Alpha Channel
+2. Use lasoo tool to select region to mask
+3. Choose Select -> Float to create a floating selection
+4. Open the Layers toolbar (^L) and select "Floating Selection"
+5. Set opacity to 0%
+6. Export as PNG
+7. In the export dialogue, Make sure the "Save colour values from
+transparent pixels" checkbox is selected.
+
 # Contributing
 
 Anyone who wishes to contribute to this project, whether

Stable-Diffusion-local-Windows.ipynb

@@ -65,25 +65,31 @@
     "imageio-ffmpeg==0.4.2\n",
     "imageio==2.9.0\n",
     "kornia==0.6.0\n",
+    "# pip will resolve the version which matches torch\n",
+    "numpy\n",
     "omegaconf==2.1.1\n",
     "opencv-python==4.6.0.66\n",
     "pillow==9.2.0\n",
+    "pip>=22\n",
     "pudb==2019.2\n",
     "pytorch-lightning==1.4.2\n",
     "streamlit==1.12.0\n",
-    "# Regular \"taming-transformers\" doesn't seem to work\n",
+    "# \"CompVis/taming-transformers\" doesn't work\n",
+    "# ldm\\models\\autoencoder.py\", line 6, in <module>\n",
+    "#    from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer\n",
+    "# ModuleNotFoundError\n",
     "taming-transformers-rom1504==0.0.6\n",
     "test-tube>=0.7.5\n",
     "torch-fidelity==0.3.0\n",
     "torchmetrics==0.6.0\n",
-    "torchvision==0.12.0\n",
     "transformers==4.19.2\n",
     "git+https://github.com/openai/CLIP.git@main#egg=clip\n",
     "git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion\n",
     "# No CUDA in PyPi builds\n",
-    "torch@https://download.pytorch.org/whl/cu113/torch-1.11.0%2Bcu113-cp310-cp310-win_amd64.whl\n",
-    "# No MKL in PyPi builds (faster, more robust than OpenBLAS)\n",
-    "numpy@https://download.lfd.uci.edu/pythonlibs/archived/numpy-1.22.4+mkl-cp310-cp310-win_amd64.whl\n",
+    "--extra-index-url https://download.pytorch.org/whl/cu113 --trusted-host https://download.pytorch.org\n",
+    "torch==1.11.0\n",
+    "# Same as numpy - let pip do its thing\n",
+    "torchvision\n",
     "-e .\n"
    ]
   },

Stable_Diffusion_AI_Notebook.ipynb

@@ -3,7 +3,6 @@
   "nbformat_minor": 0,
   "metadata": {
     "colab": {
-      "name": "Stable_Diffusion_AI_Notebook.ipynb",
       "provenance": [],
       "collapsed_sections": [],
       "private_outputs": true
@@ -22,18 +21,18 @@
     {
       "cell_type": "markdown",
       "source": [
-        "# Stable Diffusion AI Notebook\n",
+        "# Stable Diffusion AI Notebook (Release 1.13)\n",
         "\n",
         "<img src=\"https://user-images.githubusercontent.com/60411196/186547976-d9de378a-9de8-4201-9c25-c057a9c59bad.jpeg\" alt=\"stable-diffusion-ai\" width=\"170px\"/> <br>\n",
         "#### Instructions:\n",
         "1. Execute each cell in order to mount a Dream bot and create images from text. <br>\n",
-        "2. Once cells 1-8 were run correctly you'll be executing a terminal in cell #9, you'll to enter `pipenv run scripts/dream.py` command to run Dream bot.<br> \n",
+        "2. Once cells 1-8 were run correctly you'll be executing a terminal in cell #9, you'll need to enter `python scripts/dream.py` command to run Dream bot.<br> \n",
         "3. After launching dream bot, you'll see: <br> `Dream > ` in terminal. <br> Insert a command, eg. `Dream > Astronaut floating in a distant galaxy`, or type `-h` for help.\n",
-        "3. After completion you'll see your generated images in path `stable-diffusion/outputs/img-samples/`, you can also display images in cell #10.\n",
+        "3. After completion you'll see your generated images in path `stable-diffusion/outputs/img-samples/`, you can also show last generated images in cell #10.\n",
         "4. To quit Dream bot use `q` command. <br> \n",
         "---\n",
         "<font color=\"red\">Note:</font> It takes some time to load, but after installing all dependencies you can use the bot all time you want while colab instance is up. <br>\n",
-        "<font color=\"red\">Requirements:</font> For this notebook to work you need to have [Stable-Diffusion-v-1-4](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original) stored in your Google Drive, it will be needed in cell #6\n",
+        "<font color=\"red\">Requirements:</font> For this notebook to work you need to have [Stable-Diffusion-v-1-4](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original) stored in your Google Drive, it will be needed in cell #7\n",
         "##### For more details visit Github repository: [lstein/stable-diffusion](https://github.com/lstein/stable-diffusion)\n",
         "---\n"
       ],
@@ -41,6 +40,15 @@
         "id": "ycYWcsEKc6w7"
       }
     },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## ◢ Installation"
+      ],
+      "metadata": {
+        "id": "dr32VLxlnouf"
+      }
+    },
     {
       "cell_type": "code",
       "source": [
@@ -68,43 +76,28 @@
         "from os.path import exists\n",
         "\n",
         "if exists(\"/content/stable-diffusion/\")==True:\n",
+        "  %cd /content/stable-diffusion/\n",
         "  print(\"Already downloaded repo\")\n",
         "else:\n",
         "  !git clone --quiet https://github.com/lstein/stable-diffusion.git # Original repo\n",
-        "  %cd stable-diffusion/\n",
-        "  !git checkout --quiet tags/release-1.09\n",
-        "  "
+        "  %cd /content/stable-diffusion/\n",
+        "  !git checkout --quiet tags/release-1.13"
       ]
     },
     {
       "cell_type": "code",
       "source": [
-        "#@title 3. Install Python 3.8 \n",
-        "%%capture --no-stderr\n",
+        "#@title 3. Install dependencies\n",
         "import gc\n",
-        "!apt-get -qq install python3.8\n",
-        "gc.collect()"
-      ],
-      "metadata": {
-        "id": "daHlozvwKesj",
-        "cellView": "form"
-      },
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "source": [
-        "#@title 4. Install dependencies from file in a VirtualEnv\n",
-        "#@markdown Be patient, it takes ~ 5 - 7min <br>\n",
-        "%%capture --no-stderr\n",
-        "#Virtual environment\n",
-        "!pip install pipenv -q\n",
+        "\n",
+        "if exists(\"/content/stable-diffusion/requirements-colab.txt\")==True:\n",
+        "  %cd /content/stable-diffusion/\n",
+        "  print(\"Already downloaded requirements file\")\n",
+        "else:\n",
+        "  !wget https://raw.githubusercontent.com/lstein/stable-diffusion/development/requirements-colab.txt\n",
         "!pip install colab-xterm\n",
-        "%load_ext colabxterm\n",
-        "!pipenv --python 3.8\n",
-        "!pipenv install -r requirements.txt --skip-lock\n",
-        "gc.collect()\n"
+        "!pip install -r requirements-colab.txt\n",
+        "gc.collect()"
       ],
       "metadata": {
         "cellView": "form",
@@ -116,7 +109,44 @@
     {
       "cell_type": "code",
       "source": [
-        "#@title 5. Mount google Drive\n",
+        "#@title 4. Load small ML models required\n",
+        "%cd /content/stable-diffusion/\n",
+        "!python scripts/preload_models.py\n",
+        "gc.collect()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "ChIDWxLVHGGJ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title 5. Restart Runtime\n",
+        "exit()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "8rSMhgnAttQa"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## ◢ Configuration"
+      ],
+      "metadata": {
+        "id": "795x1tMoo8b1"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title 6. Mount google Drive\n",
         "from google.colab import drive\n",
         "drive.mount('/content/drive')"
       ],
@@ -130,7 +160,7 @@
     {
       "cell_type": "code",
       "source": [
-        "#@title 6. Drive Path to model\n",
+        "#@title 7. Drive Path to model\n",
         "#@markdown Path should start with /content/drive/path-to-your-file <br>\n",
         "#@markdown <font color=\"red\">Note:</font> Model should be downloaded from https://huggingface.co <br>\n",
         "#@markdown Lastest release: [Stable-Diffusion-v-1-4](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original)\n",
@@ -152,7 +182,7 @@
     {
       "cell_type": "code",
       "source": [
-        "#@title 7. Symlink to model\n",
+        "#@title 8. Symlink to model\n",
         "\n",
         "from os.path import exists\n",
         "import os \n",
@@ -181,32 +211,27 @@
       "outputs": []
     },
     {
-      "cell_type": "code",
+      "cell_type": "markdown",
       "source": [
-        "#@title 8. Load small ML models required\n",
-        "%%capture --no-stderr\n",
-        "!pipenv run scripts/preload_models.py\n",
-        "gc.collect()"
+        "## ◢ Execution"
       ],
       "metadata": {
-        "cellView": "form",
-        "id": "ChIDWxLVHGGJ"
-      },
-      "execution_count": null,
-      "outputs": []
+        "id": "Mc28N0_NrCQH"
+      }
     },
     {
       "cell_type": "code",
       "source": [
         "#@title 9. Run Terminal and Execute Dream bot\n",
         "#@markdown <font color=\"blue\">Steps:</font> <br>\n",
-        "#@markdown 1. Execute command `pipenv run scripts/dream.py` to run dream bot.<br>\n",
+        "#@markdown 1. Execute command `python scripts/dream.py` to run dream bot.<br>\n",
         "#@markdown 2. After initialized you'll see `Dream>` line.<br>\n",
         "#@markdown 3. Example text: `Astronaut floating in a distant galaxy` <br>\n",
         "#@markdown 4. To quit Dream bot use: `q` command.<br>\n",
         "\n",
-        "#Run from virtual env\n",
-        "\n",
+        "import gc\n",
+        "%cd /content/stable-diffusion/\n",
+        "%load_ext colabxterm\n",
         "%xterm\n",
         "gc.collect()"
       ],
@@ -220,18 +245,18 @@
     {
       "cell_type": "code",
       "source": [
-        "#@title 10. Show generated images\n",
-        "\n",
+        "#@title 10. Show the last 15 generated images\n",
+        "import gc\n",
         "import glob\n",
         "import matplotlib.pyplot as plt\n",
         "import matplotlib.image as mpimg\n",
         "%matplotlib inline\n",
         "\n",
         "images = []\n",
-        "for img_path in glob.glob('/content/stable-diffusion/outputs/img-samples/*.png'):\n",
+        "for img_path in sorted(glob.glob('/content/stable-diffusion/outputs/img-samples/*.png'), reverse=True):\n",
         "    images.append(mpimg.imread(img_path))\n",
         "\n",
-        "# Remove ticks and labels on x-axis and y-axis both\n",
+        "images = images[:15] \n",
         "\n",
         "plt.figure(figsize=(20,10))\n",
         "\n",
@@ -253,4 +278,4 @@
       "outputs": []
     }
   ]
-}
+}

VARIATIONS.md

@@ -108,6 +108,6 @@ the chosen two images. Here's the one I like best:
 
 <img src="static/variation_walkthru/000004.3747154981.png">
 
-As you can see, this is a very powerful too, which when combined with
+As you can see, this is a very powerful tool, which when combined with
 subprompt weighting, gives you great control over the content and
 quality of your generated images.

environment-mac.yaml

@@ -1,33 +1,29 @@
 name: ldm
 channels:
-  - pytorch-nightly
+  - pytorch
   - conda-forge
 dependencies:
-  - python==3.9.13
+  - python==3.10.5
   - pip==22.2.2
  
-  # pytorch-nightly, left unpinned
+  # pytorch left unpinned
   - pytorch
-  - torchmetrics
   - torchvision
 
   # I suggest to keep the other deps sorted for convenience.
-  # If you wish to upgrade to 3.10, try to run this:
+  # To determine what the latest versions should be, run:
   #
   # ```shell
-  # CONDA_CMD=conda
-  # sed -E 's/python==3.9.13/python==3.10.5/;s/ldm/ldm-3.10/;21,99s/- ([^=]+)==.+/- \1/' environment-mac.yaml > /tmp/environment-mac-updated.yml
-  # CONDA_SUBDIR=osx-arm64 $CONDA_CMD env create -f /tmp/environment-mac-updated.yml && $CONDA_CMD list -n ldm-3.10 | awk ' {print "  - " $1 "==" $2;} '
+  # sed -E 's/ldm/ldm-updated/;20,99s/- ([^=]+)==.+/- \1/' environment-mac.yaml > environment-mac-updated.yml
+  # CONDA_SUBDIR=osx-arm64 conda env create -f environment-mac-updated.yml && conda list -n ldm-updated | awk ' {print "  - " $1 "==" $2;} '
   # ```
-  #
-  # Unfortunately, as of 2022-08-31, this fails at the pip stage.
   - albumentations==1.2.1
   - coloredlogs==15.0.1
   - einops==0.4.1
   - grpcio==1.46.4
-  - humanfriendly
-  - imageio-ffmpeg==0.4.7
+  - humanfriendly==10.0
   - imageio==2.21.2
+  - imageio-ffmpeg==0.4.7
   - imgaug==0.4.0
   - kornia==0.6.7
   - mpmath==1.2.1
@@ -43,13 +39,11 @@ dependencies:
   - streamlit==1.12.2
   - sympy==1.10.1
   - tensorboard==2.9.0
-  - transformers==4.21.2
+  - torchmetrics==0.9.3
   - pip:
-    - invisible-watermark
-    - test-tube
-    - tokenizers
-    - torch-fidelity
-    - -e git+https://github.com/huggingface/diffusers.git@v0.2.4#egg=diffusers
+    - test-tube==0.7.5
+    - transformers==4.21.2
+    - torch-fidelity==0.3.0
     - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
     - -e git+https://github.com/openai/CLIP.git@main#egg=clip
     - -e git+https://github.com/Birch-san/k-diffusion.git@mps#egg=k_diffusion

ldm/dream/conditioning.py

@@ -0,0 +1,96 @@
+'''
+This module handles the generation of the conditioning tensors, including management of
+weighted subprompts.
+
+Useful function exports:
+
+get_uc_and_c()                  get the conditioned and unconditioned latent
+split_weighted_subpromopts()    split subprompts, normalize and weight them
+log_tokenization()              print out colour-coded tokens and warn if truncated
+
+'''
+import re
+import torch
+
+def get_uc_and_c(prompt, model, log_tokens=False, skip_normalize=False):
+    uc = model.get_learned_conditioning([''])
+
+    # get weighted sub-prompts
+    weighted_subprompts = split_weighted_subprompts(
+        prompt, skip_normalize
+    )
+
+    if len(weighted_subprompts) > 1:
+        # i dont know if this is correct.. but it works
+        c = torch.zeros_like(uc)
+        # normalize each "sub prompt" and add it
+        for subprompt, weight in weighted_subprompts:
+            log_tokenization(subprompt, model, log_tokens)
+            c = torch.add(
+                c,
+                model.get_learned_conditioning([subprompt]),
+                alpha=weight,
+            )
+    else:   # just standard 1 prompt
+        log_tokenization(prompt, model, log_tokens)
+        c = model.get_learned_conditioning([prompt])
+    return (uc, c)
+
+def split_weighted_subprompts(text, skip_normalize=False)->list:
+    """
+    grabs all text up to the first occurrence of ':'
+    uses the grabbed text as a sub-prompt, and takes the value following ':' as weight
+    if ':' has no value defined, defaults to 1.0
+    repeats until no text remaining
+    """
+    prompt_parser = re.compile("""
+            (?P<prompt>     # capture group for 'prompt'
+            (?:\\\:|[^:])+  # match one or more non ':' characters or escaped colons '\:'
+            )               # end 'prompt'
+            (?:             # non-capture group
+            :+              # match one or more ':' characters
+            (?P<weight>     # capture group for 'weight'
+            -?\d+(?:\.\d+)? # match positive or negative integer or decimal number
+            )?              # end weight capture group, make optional
+            \s*             # strip spaces after weight
+            |               # OR
+            $               # else, if no ':' then match end of line
+            )               # end non-capture group
+            """, re.VERBOSE)
+    parsed_prompts = [(match.group("prompt").replace("\\:", ":"), float(
+        match.group("weight") or 1)) for match in re.finditer(prompt_parser, text)]
+    if skip_normalize:
+        return parsed_prompts
+    weight_sum = sum(map(lambda x: x[1], parsed_prompts))
+    if weight_sum == 0:
+        print(
+            "Warning: Subprompt weights add up to zero. Discarding and using even weights instead.")
+        equal_weight = 1 / len(parsed_prompts)
+        return [(x[0], equal_weight) for x in parsed_prompts]
+    return [(x[0], x[1] / weight_sum) for x in parsed_prompts]
+        
+# shows how the prompt is tokenized
+# usually tokens have '</w>' to indicate end-of-word,
+# but for readability it has been replaced with ' '
+def log_tokenization(text, model, log=False):
+    if not log:
+        return
+    tokens    = model.cond_stage_model.tokenizer._tokenize(text)
+    tokenized = ""
+    discarded = ""
+    usedTokens = 0
+    totalTokens = len(tokens)
+    for i in range(0, totalTokens):
+        token = tokens[i].replace('</w>', ' ')
+        # alternate color
+        s = (usedTokens % 6) + 1
+        if i < model.cond_stage_model.max_length:
+            tokenized = tokenized + f"\x1b[0;3{s};40m{token}"
+            usedTokens += 1
+        else:  # over max token length
+            discarded = discarded + f"\x1b[0;3{s};40m{token}"
+        print(f"\n>> Tokens ({usedTokens}):\n{tokenized}\x1b[0m")
+        if discarded != "":
+            print(
+                f">> Tokens Discarded ({totalTokens-usedTokens}):\n{discarded}\x1b[0m"
+            )

ldm/dream/devices.py

@@ -1,4 +1,6 @@
 import torch
+from torch import autocast
+from contextlib import contextmanager, nullcontext
 
 def choose_torch_device() -> str:
     '''Convenience routine for guessing which GPU device to run model on'''
@@ -8,10 +10,11 @@ def choose_torch_device() -> str:
         return 'mps'
     return 'cpu'
 
-def choose_autocast_device(device) -> str:
+def choose_autocast_device(device):
     '''Returns an autocast compatible device from a torch device'''
     device_type = device.type # this returns 'mps' on M1
     # autocast only supports cuda or cpu
-    if device_type not in ('cuda','cpu'):
-        return 'cpu'
-    return device_type
+    if device_type in ('cuda','cpu'):
+        return device_type,autocast
+    else:
+        return 'cpu',nullcontext

ldm/dream/generator/__init__.py

@@ -0,0 +1,4 @@
+'''
+Initialization file for the ldm.dream.generator package
+'''
+from .base import Generator

ldm/dream/generator/base.py

@@ -0,0 +1,158 @@
+'''
+Base class for ldm.dream.generator.*
+including img2img, txt2img, and inpaint
+'''
+import torch
+import numpy as  np
+import random
+from tqdm import tqdm, trange
+from PIL               import Image
+from einops import rearrange, repeat
+from pytorch_lightning import seed_everything
+from ldm.dream.devices import choose_autocast_device
+
+downsampling = 8
+
+class Generator():
+    def __init__(self,model):
+        self.model               = model
+        self.seed                = None
+        self.latent_channels     = model.channels
+        self.downsampling_factor = downsampling   # BUG: should come from model or config
+        self.variation_amount    = 0
+        self.with_variations     = []
+
+    # this is going to be overridden in img2img.py, txt2img.py and inpaint.py
+    def get_make_image(self,prompt,**kwargs):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        """
+        raise NotImplementedError("image_iterator() must be implemented in a descendent class")
+
+    def set_variation(self, seed, variation_amount, with_variations):
+        self.seed             = seed
+        self.variation_amount = variation_amount
+        self.with_variations  = with_variations
+
+    def generate(self,prompt,init_image,width,height,iterations=1,seed=None,
+                 image_callback=None, step_callback=None,
+                 **kwargs):
+        device_type,scope   = choose_autocast_device(self.model.device)
+        make_image          = self.get_make_image(
+            prompt,
+            init_image    = init_image,
+            width         = width,
+            height        = height,
+            step_callback = step_callback,
+            **kwargs
+        )
+
+        results             = []
+        seed                = seed if seed else self.new_seed()
+        seed, initial_noise = self.generate_initial_noise(seed, width, height)
+        with scope(device_type), self.model.ema_scope():
+            for n in trange(iterations, desc='Generating'):
+                x_T = None
+                if self.variation_amount > 0:
+                    seed_everything(seed)
+                    target_noise = self.get_noise(width,height)
+                    x_T = self.slerp(self.variation_amount, initial_noise, target_noise)
+                elif initial_noise is not None:
+                    # i.e. we specified particular variations
+                    x_T = initial_noise
+                else:
+                    seed_everything(seed)
+                    if self.model.device.type == 'mps':
+                        x_T = self.get_noise(width,height)
+
+                # make_image will do the equivalent of get_noise itself
+                image = make_image(x_T)
+                results.append([image, seed])
+                if image_callback is not None:
+                    image_callback(image, seed)
+                seed = self.new_seed()
+        return results
+    
+    def sample_to_image(self,samples):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        """
+        x_samples = self.model.decode_first_stage(samples)
+        x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+        if len(x_samples) != 1:
+            raise Exception(
+                f'>> expected to get a single image, but got {len(x_samples)}')
+        x_sample = 255.0 * rearrange(
+            x_samples[0].cpu().numpy(), 'c h w -> h w c'
+        )
+        return Image.fromarray(x_sample.astype(np.uint8))
+
+    def generate_initial_noise(self, seed, width, height):
+        initial_noise = None
+        if self.variation_amount > 0 or len(self.with_variations) > 0:
+            # use fixed initial noise plus random noise per iteration
+            seed_everything(seed)
+            initial_noise = self.get_noise(width,height)
+            for v_seed, v_weight in self.with_variations:
+                seed = v_seed
+                seed_everything(seed)
+                next_noise = self.get_noise(width,height)
+                initial_noise = self.slerp(v_weight, initial_noise, next_noise)
+            if self.variation_amount > 0:
+                random.seed() # reset RNG to an actually random state, so we can get a random seed for variations
+                seed = random.randrange(0,np.iinfo(np.uint32).max)
+            return (seed, initial_noise)
+        else:
+            return (seed, None)
+
+    # returns a tensor filled with random numbers from a normal distribution
+    def get_noise(self,width,height):
+        """
+        Returns a tensor filled with random numbers, either form a normal distribution
+        (txt2img) or from the latent image (img2img, inpaint)
+        """
+        raise NotImplementedError("get_noise() must be implemented in a descendent class")
+    
+    def new_seed(self):
+        self.seed = random.randrange(0, np.iinfo(np.uint32).max)
+        return self.seed
+
+    def slerp(self, t, v0, v1, DOT_THRESHOLD=0.9995):
+        '''
+        Spherical linear interpolation
+        Args:
+            t (float/np.ndarray): Float value between 0.0 and 1.0
+            v0 (np.ndarray): Starting vector
+            v1 (np.ndarray): Final vector
+            DOT_THRESHOLD (float): Threshold for considering the two vectors as
+                                colineal. Not recommended to alter this.
+        Returns:
+            v2 (np.ndarray): Interpolation vector between v0 and v1
+        '''
+        inputs_are_torch = False
+        if not isinstance(v0, np.ndarray):
+            inputs_are_torch = True
+            v0 = v0.detach().cpu().numpy()
+        if not isinstance(v1, np.ndarray):
+            inputs_are_torch = True
+            v1 = v1.detach().cpu().numpy()
+
+        dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
+        if np.abs(dot) > DOT_THRESHOLD:
+            v2 = (1 - t) * v0 + t * v1
+        else:
+            theta_0 = np.arccos(dot)
+            sin_theta_0 = np.sin(theta_0)
+            theta_t = theta_0 * t
+            sin_theta_t = np.sin(theta_t)
+            s0 = np.sin(theta_0 - theta_t) / sin_theta_0
+            s1 = sin_theta_t / sin_theta_0
+            v2 = s0 * v0 + s1 * v1
+
+        if inputs_are_torch:
+            v2 = torch.from_numpy(v2).to(self.model.device)
+
+        return v2
+

ldm/dream/generator/img2img.py

@@ -0,0 +1,72 @@
+'''
+ldm.dream.generator.txt2img descends from ldm.dream.generator
+'''
+
+import torch
+import numpy as  np
+from ldm.dream.devices             import choose_autocast_device
+from ldm.dream.generator.base      import Generator
+from ldm.models.diffusion.ddim     import DDIMSampler
+
+class Img2Img(Generator):
+    def __init__(self,model):
+        super().__init__(model)
+        self.init_latent         = None    # by get_noise()
+    
+    @torch.no_grad()
+    def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
+                       conditioning,init_image,strength,step_callback=None,**kwargs):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it.
+        """
+
+        # PLMS sampler not supported yet, so ignore previous sampler
+        if not isinstance(sampler,DDIMSampler):
+            print(
+                f">> sampler '{sampler.__class__.__name__}' is not yet supported. Using DDIM sampler"
+            )
+            sampler = DDIMSampler(self.model, device=self.model.device)
+
+        sampler.make_schedule(
+            ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
+        )
+
+        device_type,scope   = choose_autocast_device(self.model.device)
+        with scope(device_type):
+            self.init_latent = self.model.get_first_stage_encoding(
+                self.model.encode_first_stage(init_image)
+            ) # move to latent space
+
+        t_enc = int(strength * steps)
+        uc, c   = conditioning
+
+        @torch.no_grad()
+        def make_image(x_T):
+            # encode (scaled latent)
+            z_enc = sampler.stochastic_encode(
+                self.init_latent,
+                torch.tensor([t_enc]).to(self.model.device),
+                noise=x_T
+            )
+            # decode it
+            samples = sampler.decode(
+                z_enc,
+                c,
+                t_enc,
+                img_callback = step_callback,
+                unconditional_guidance_scale=cfg_scale,
+                unconditional_conditioning=uc,
+            )
+            return self.sample_to_image(samples)
+
+        return make_image
+
+    def get_noise(self,width,height):
+        device      = self.model.device
+        init_latent = self.init_latent
+        assert init_latent is not None,'call to get_noise() when init_latent not set'
+        if device.type == 'mps':
+            return torch.randn_like(init_latent, device='cpu').to(device)
+        else:
+            return torch.randn_like(init_latent, device=device)

ldm/dream/generator/inpaint.py

@@ -0,0 +1,77 @@
+'''
+ldm.dream.generator.inpaint descends from ldm.dream.generator
+'''
+
+import torch
+import numpy as  np
+from einops import rearrange, repeat
+from ldm.dream.devices             import choose_autocast_device
+from ldm.dream.generator.img2img   import Img2Img
+from ldm.models.diffusion.ddim     import DDIMSampler
+
+class Inpaint(Img2Img):
+    def __init__(self,model):
+        self.init_latent = None
+        super().__init__(model)
+    
+    @torch.no_grad()
+    def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
+                       conditioning,init_image,mask_image,strength,
+                       step_callback=None,**kwargs):
+        """
+        Returns a function returning an image derived from the prompt and
+        the initial image + mask.  Return value depends on the seed at
+        the time you call it.  kwargs are 'init_latent' and 'strength'
+        """
+
+        mask_image = mask_image[0][0].unsqueeze(0).repeat(4,1,1).unsqueeze(0)
+        mask_image = repeat(mask_image, '1 ... -> b ...', b=1)
+
+        # PLMS sampler not supported yet, so ignore previous sampler
+        if not isinstance(sampler,DDIMSampler):
+            print(
+                f">> sampler '{sampler.__class__.__name__}' is not yet supported. Using DDIM sampler"
+            )
+            sampler = DDIMSampler(self.model, device=self.model.device)
+
+            sampler.make_schedule(
+                ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
+            )
+
+        device_type,scope   = choose_autocast_device(self.model.device)
+        with scope(device_type):
+            self.init_latent = self.model.get_first_stage_encoding(
+                self.model.encode_first_stage(init_image)
+            ) # move to latent space
+
+        t_enc   = int(strength * steps)
+        uc, c   = conditioning
+
+        print(f">> target t_enc is {t_enc} steps")
+
+        @torch.no_grad()
+        def make_image(x_T):
+            # encode (scaled latent)
+            z_enc = sampler.stochastic_encode(
+                self.init_latent,
+                torch.tensor([t_enc]).to(self.model.device),
+                noise=x_T
+            )
+                                       
+            # decode it
+            samples = sampler.decode(
+                z_enc,
+                c,
+                t_enc,
+                img_callback                 = step_callback,
+                unconditional_guidance_scale = cfg_scale,
+                unconditional_conditioning = uc,
+                mask                       = mask_image,
+                init_latent                = self.init_latent
+            )
+            return self.sample_to_image(samples)
+
+        return make_image
+
+
+

ldm/dream/generator/txt2img.py

@@ -0,0 +1,61 @@
+'''
+ldm.dream.generator.txt2img inherits from ldm.dream.generator
+'''
+
+import torch
+import numpy as  np
+from ldm.dream.generator.base import Generator
+
+class Txt2Img(Generator):
+    def __init__(self,model):
+        super().__init__(model)
+    
+    @torch.no_grad()
+    def get_make_image(self,prompt,sampler,steps,cfg_scale,ddim_eta,
+                       conditioning,width,height,step_callback=None,**kwargs):
+        """
+        Returns a function returning an image derived from the prompt and the initial image
+        Return value depends on the seed at the time you call it
+        kwargs are 'width' and 'height'
+        """
+        uc, c   = conditioning
+
+        @torch.no_grad()
+        def make_image(x_T):
+            shape = [
+                self.latent_channels,
+                height // self.downsampling_factor,
+                width  // self.downsampling_factor,
+            ]
+            samples, _ = sampler.sample(
+                batch_size                   = 1,
+                S                            = steps,
+                x_T                          = x_T,
+                conditioning                 = c,
+                shape                        = shape,
+                verbose                      = False,
+                unconditional_guidance_scale = cfg_scale,
+                unconditional_conditioning   = uc,
+                eta                          = ddim_eta,
+                img_callback                 = step_callback
+            )
+            return self.sample_to_image(samples)
+
+        return make_image
+
+
+    # returns a tensor filled with random numbers from a normal distribution
+    def get_noise(self,width,height):
+        device         = self.model.device
+        if device.type == 'mps':
+            return torch.randn([1,
+                                self.latent_channels,
+                                height // self.downsampling_factor,
+                                width  // self.downsampling_factor],
+                               device='cpu').to(device)
+        else:
+            return torch.randn([1,
+                                self.latent_channels,
+                                height // self.downsampling_factor,
+                                width  // self.downsampling_factor],
+                               device=device)

ldm/dream/pngwriter.py

@@ -59,6 +59,10 @@ class PromptFormatter:
         switches.append(f'-H{opt.height       or t2i.height}')
         switches.append(f'-C{opt.cfg_scale    or t2i.cfg_scale}')
         switches.append(f'-A{opt.sampler_name or t2i.sampler_name}')
+# to do: put model name into the t2i object
+#        switches.append(f'--model{t2i.model_name}')
+        if opt.seamless or t2i.seamless:
+            switches.append(f'--seamless')
         if opt.init_img:
             switches.append(f'-I{opt.init_img}')
         if opt.fit:

ldm/dream/readline.py

@@ -22,7 +22,7 @@ class Completer:
     def complete(self, text, state):
         buffer = readline.get_line_buffer()
 
-        if text.startswith(('-I', '--init_img')):
+        if text.startswith(('-I', '--init_img','-M','--init_mask')):
             return self._path_completions(text, state, ('.png','.jpg','.jpeg'))
 
         if buffer.strip().endswith('cd') or text.startswith(('.', '/')):
@@ -48,10 +48,15 @@ class Completer:
 
     def _path_completions(self, text, state, extensions):
         # get the path so far
+        # TODO: replace this mess with a regular expression match
         if text.startswith('-I'):
             path = text.replace('-I', '', 1).lstrip()
         elif text.startswith('--init_img='):
             path = text.replace('--init_img=', '', 1).lstrip()
+        elif text.startswith('--init_mask='):
+            path = text.replace('--init_mask=', '', 1).lstrip()
+        elif text.startswith('-M'):
+            path = text.replace('-M', '', 1).lstrip()
         else:
             path = text
 
@@ -94,6 +99,7 @@ if readline_available:
                 '--grid','-g',
                 '--individual','-i',
                 '--init_img','-I',
+                '--init_mask','-M',
                 '--strength','-f',
                 '--variants','-v',
                 '--outdir','-o',

ldm/dream/server.py

@@ -1,16 +1,65 @@
+import argparse
 import json
 import base64
 import mimetypes
 import os
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
-from ldm.dream.pngwriter import PngWriter
+from ldm.dream.pngwriter import PngWriter, PromptFormatter
 from threading import Event
 
+def build_opt(post_data, seed, gfpgan_model_exists):
+    opt = argparse.Namespace()
+    setattr(opt, 'prompt', post_data['prompt'])
+    setattr(opt, 'init_img', post_data['initimg'])
+    setattr(opt, 'strength', float(post_data['strength']))
+    setattr(opt, 'iterations', int(post_data['iterations']))
+    setattr(opt, 'steps', int(post_data['steps']))
+    setattr(opt, 'width', int(post_data['width']))
+    setattr(opt, 'height', int(post_data['height']))
+    setattr(opt, 'seamless', 'seamless' in post_data)
+    setattr(opt, 'fit', 'fit' in post_data)
+    setattr(opt, 'mask', 'mask' in post_data)
+    setattr(opt, 'invert_mask', 'invert_mask' in post_data)
+    setattr(opt, 'cfg_scale', float(post_data['cfg_scale']))
+    setattr(opt, 'sampler_name', post_data['sampler_name'])
+    setattr(opt, 'gfpgan_strength', float(post_data['gfpgan_strength']) if gfpgan_model_exists else 0)
+    setattr(opt, 'upscale', [int(post_data['upscale_level']), float(post_data['upscale_strength'])] if post_data['upscale_level'] != '' else None)
+    setattr(opt, 'progress_images', 'progress_images' in post_data)
+    setattr(opt, 'seed', None if int(post_data['seed']) == -1 else int(post_data['seed']))
+    setattr(opt, 'variation_amount', float(post_data['variation_amount']) if int(post_data['seed']) != -1 else 0)
+    setattr(opt, 'with_variations', [])
+
+    broken = False
+    if int(post_data['seed']) != -1 and post_data['with_variations'] != '':
+        for part in post_data['with_variations'].split(','):
+            seed_and_weight = part.split(':')
+            if len(seed_and_weight) != 2:
+                print(f'could not parse with_variation part "{part}"')
+                broken = True
+                break
+            try:
+                seed = int(seed_and_weight[0])
+                weight = float(seed_and_weight[1])
+            except ValueError:
+                print(f'could not parse with_variation part "{part}"')
+                broken = True
+                break
+            opt.with_variations.append([seed, weight])
+    
+    if broken:
+        raise CanceledException
+
+    if len(opt.with_variations) == 0:
+        opt.with_variations = None
+
+    return opt
+
 class CanceledException(Exception):
     pass
 
 class DreamServer(BaseHTTPRequestHandler):
     model = None
+    outdir = None
     canceled = Event()
 
     def do_GET(self):
@@ -30,6 +79,23 @@ class DreamServer(BaseHTTPRequestHandler):
                 'gfpgan_model_exists': gfpgan_model_exists
             }
             self.wfile.write(bytes("let config = " + json.dumps(config) + ";\n", "utf-8"))
+        elif self.path == "/run_log.json":
+            self.send_response(200)
+            self.send_header("Content-type", "application/json")
+            self.end_headers()
+            output = []
+            
+            log_file = os.path.join(self.outdir, "dream_web_log.txt")
+            if os.path.exists(log_file):
+                with open(log_file, "r") as log:
+                    for line in log:
+                        url, config = line.split(": {", maxsplit=1)
+                        config = json.loads("{" + config)
+                        config["url"] = url.lstrip(".")
+                        if os.path.exists(url):
+                            output.append(config)
+
+            self.wfile.write(bytes(json.dumps({"run_log": output}), "utf-8"))
         elif self.path == "/cancel":
             self.canceled.set()
             self.send_response(200)
@@ -63,34 +129,19 @@ class DreamServer(BaseHTTPRequestHandler):
 
         content_length = int(self.headers['Content-Length'])
         post_data = json.loads(self.rfile.read(content_length))
-        prompt = post_data['prompt']
-        initimg = post_data['initimg']
-        strength = float(post_data['strength'])
-        iterations = int(post_data['iterations'])
-        steps = int(post_data['steps'])
-        width = int(post_data['width'])
-        height = int(post_data['height'])
-        fit    = 'fit' in post_data
-        cfgscale = float(post_data['cfgscale'])
-        sampler_name  = post_data['sampler']
-        gfpgan_strength = float(post_data['gfpgan_strength']) if gfpgan_model_exists else 0
-        upscale_level    = post_data['upscale_level']
-        upscale_strength = post_data['upscale_strength']
-        upscale = [int(upscale_level),float(upscale_strength)] if upscale_level != '' else None
-        progress_images = 'progress_images' in post_data
-        seed = self.model.seed if int(post_data['seed']) == -1 else int(post_data['seed'])
+        opt = build_opt(post_data, self.model.seed, gfpgan_model_exists)
 
         self.canceled.clear()
-        print(f">> Request to generate with prompt: {prompt}")
+        print(f">> Request to generate with prompt: {opt.prompt}")
         # In order to handle upscaled images, the PngWriter needs to maintain state
         # across images generated by each call to prompt2img(), so we define it in
         # the outer scope of image_done()
         config = post_data.copy() # Shallow copy
-        config['initimg'] = ''
+        config['initimg'] = config.pop('initimg_name', '')
 
         images_generated = 0    # helps keep track of when upscaling is started
         images_upscaled = 0     # helps keep track of when upscaling is completed
-        pngwriter = PngWriter("./outputs/img-samples/")
+        pngwriter = PngWriter(self.outdir)
 
         prefix = pngwriter.unique_prefix()
         # if upscaling is requested, then this will be called twice, once when
@@ -99,11 +150,24 @@ class DreamServer(BaseHTTPRequestHandler):
         # entry should not be inserted into the image list.
         def image_done(image, seed, upscaled=False):
             name = f'{prefix}.{seed}.png'
-            path = pngwriter.save_image_and_prompt_to_png(image, f'{prompt} -S{seed}', name)
+            iter_opt = argparse.Namespace(**vars(opt)) # copy
+            if opt.variation_amount > 0:
+                this_variation = [[seed, opt.variation_amount]]
+                if opt.with_variations is None:
+                    iter_opt.with_variations = this_variation
+                else:
+                    iter_opt.with_variations = opt.with_variations + this_variation
+                iter_opt.variation_amount = 0
+            elif opt.with_variations is None:
+                iter_opt.seed = seed
+            normalized_prompt = PromptFormatter(self.model, iter_opt).normalize_prompt()
+            path = pngwriter.save_image_and_prompt_to_png(image, f'{normalized_prompt} -S{iter_opt.seed}', name)
 
+            if int(config['seed']) == -1:
+                config['seed'] = seed
             # Append post_data to log, but only once!
             if not upscaled:
-                with open("./outputs/img-samples/dream_web_log.txt", "a") as log:
+                with open(os.path.join(self.outdir, "dream_web_log.txt"), "a") as log:
                     log.write(f"{path}: {json.dumps(config)}\n")
 
                 self.wfile.write(bytes(json.dumps(
@@ -111,27 +175,27 @@ class DreamServer(BaseHTTPRequestHandler):
                 ) + '\n',"utf-8"))
 
             # control state of the "postprocessing..." message
-            upscaling_requested = upscale or gfpgan_strength>0
+            upscaling_requested = opt.upscale or opt.gfpgan_strength > 0
             nonlocal images_generated # NB: Is this bad python style? It is typical usage in a perl closure.
             nonlocal images_upscaled  # NB: Is this bad python style? It is typical usage in a perl closure.
             if upscaled:
                 images_upscaled += 1
             else:
-                images_generated +=1
+                images_generated += 1
             if upscaling_requested:
                 action = None
-                if images_generated >= iterations:
-                    if images_upscaled < iterations:
+                if images_generated >= opt.iterations:
+                    if images_upscaled < opt.iterations:
                         action = 'upscaling-started'
                     else:
                         action = 'upscaling-done'
                 if action:
-                    x = images_upscaled+1
+                    x = images_upscaled + 1
                     self.wfile.write(bytes(json.dumps(
-                        {'event':action,'processed_file_cnt':f'{x}/{iterations}'}
+                        {'event': action, 'processed_file_cnt': f'{x}/{opt.iterations}'}
                     ) + '\n',"utf-8"))
 
-        step_writer = PngWriter('./outputs/intermediates/')
+        step_writer = PngWriter(os.path.join(self.outdir, "intermediates"))
         step_index = 1
         def image_progress(sample, step):
             if self.canceled.is_set():
@@ -141,10 +205,10 @@ class DreamServer(BaseHTTPRequestHandler):
             # since rendering images is moderately expensive, only render every 5th image
             # and don't bother with the last one, since it'll render anyway
             nonlocal step_index
-            if progress_images and step % 5 == 0 and step < steps - 1:
-                image = self.model._sample_to_image(sample)
-                name = f'{prefix}.{seed}.{step_index}.png'
-                metadata = f'{prompt} -S{seed} [intermediate]'
+            if opt.progress_images and step % 5 == 0 and step < opt.steps - 1:
+                image = self.model.sample_to_image(sample)
+                name = f'{prefix}.{opt.seed}.{step_index}.png'
+                metadata = f'{opt.prompt} -S{opt.seed} [intermediate]'
                 path = step_writer.save_image_and_prompt_to_png(image, metadata, name)
                 step_index += 1
             self.wfile.write(bytes(json.dumps(
@@ -152,43 +216,20 @@ class DreamServer(BaseHTTPRequestHandler):
             ) + '\n',"utf-8"))
 
         try:
-            if initimg is None:
+            if opt.init_img is None:
                 # Run txt2img
-                self.model.prompt2image(prompt,
-                                        iterations=iterations,
-                                        cfg_scale = cfgscale,
-                                        width = width,
-                                        height = height,
-                                        seed = seed,
-                                        steps = steps,
-                                        gfpgan_strength = gfpgan_strength,
-                                        upscale         = upscale,
-                                        sampler_name    = sampler_name,
-                                        step_callback=image_progress,
-                                        image_callback=image_done)
+                self.model.prompt2image(**vars(opt), step_callback=image_progress, image_callback=image_done)
             else:
                 # Decode initimg as base64 to temp file
                 with open("./img2img-tmp.png", "wb") as f:
-                    initimg = initimg.split(",")[1] # Ignore mime type
+                    initimg = opt.init_img.split(",")[1] # Ignore mime type
                     f.write(base64.b64decode(initimg))
+                opt1 = argparse.Namespace(**vars(opt))
+                opt1.init_img = "./img2img-tmp.png"
 
                 try:
                     # Run img2img
-                    self.model.prompt2image(prompt,
-                                            init_img = "./img2img-tmp.png",
-                                            strength = strength,
-                                            iterations = iterations,
-                                            cfg_scale  = cfgscale,
-                                            seed       = seed,
-                                            steps      = steps,
-                                            sampler_name    = sampler_name,
-                                            width      = width,
-                                            height     = height,
-                                            fit        = fit,
-                                            gfpgan_strength=gfpgan_strength,
-                                            upscale         = upscale,
-                                            step_callback=image_progress,
-                                            image_callback=image_done)
+                    self.model.prompt2image(**vars(opt1), step_callback=image_progress, image_callback=image_done)
                 finally:
                     # Remove the temp file
                     os.remove("./img2img-tmp.png")

ldm/generate.py

@@ -0,0 +1,690 @@
+# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein)
+
+# Derived from source code carrying the following copyrights
+# Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
+# Copyright (c) 2022 Robin Rombach and Patrick Esser and contributors
+
+import torch
+import numpy as np
+import random
+import os
+import time
+import re
+import sys
+import traceback
+import transformers
+
+from omegaconf import OmegaConf
+from PIL import Image, ImageOps
+from torch import nn
+from pytorch_lightning import seed_everything
+
+from ldm.util                      import instantiate_from_config
+from ldm.models.diffusion.ddim     import DDIMSampler
+from ldm.models.diffusion.plms     import PLMSSampler
+from ldm.models.diffusion.ksampler import KSampler
+from ldm.dream.pngwriter           import PngWriter
+from ldm.dream.image_util          import InitImageResizer
+from ldm.dream.devices             import choose_torch_device
+from ldm.dream.conditioning        import get_uc_and_c
+
+"""Simplified text to image API for stable diffusion/latent diffusion
+
+Example Usage:
+
+from ldm.generate import Generate
+
+# Create an object with default values
+gr = Generate()
+
+# do the slow model initialization
+gr.load_model()
+
+# Do the fast inference & image generation. Any options passed here
+# override the default values assigned during class initialization
+# Will call load_model() if the model was not previously loaded and so
+# may be slow at first.
+# The method returns a list of images. Each row of the list is a sub-list of [filename,seed]
+results = gr.prompt2png(prompt     = "an astronaut riding a horse",
+                         outdir     = "./outputs/samples",
+                         iterations = 3)
+
+for row in results:
+    print(f'filename={row[0]}')
+    print(f'seed    ={row[1]}')
+
+# Same thing, but using an initial image.
+results = gr.prompt2png(prompt   = "an astronaut riding a horse",
+                         outdir   = "./outputs/,
+                         iterations = 3,
+                         init_img = "./sketches/horse+rider.png")
+
+for row in results:
+    print(f'filename={row[0]}')
+    print(f'seed    ={row[1]}')
+
+# Same thing, but we return a series of Image objects, which lets you manipulate them,
+# combine them, and save them under arbitrary names
+
+results = gr.prompt2image(prompt   = "an astronaut riding a horse"
+                           outdir   = "./outputs/")
+for row in results:
+    im   = row[0]
+    seed = row[1]
+    im.save(f'./outputs/samples/an_astronaut_riding_a_horse-{seed}.png')
+    im.thumbnail(100,100).save('./outputs/samples/astronaut_thumb.jpg')
+
+Note that the old txt2img() and img2img() calls are deprecated but will
+still work.
+
+The full list of arguments to Generate() are:
+gr = Generate(
+          weights     = path to model weights ('models/ldm/stable-diffusion-v1/model.ckpt')
+          config     = path to model configuraiton ('configs/stable-diffusion/v1-inference.yaml')
+          iterations  = <integer>     // how many times to run the sampling (1)
+          steps       = <integer>     // 50
+          seed        = <integer>     // current system time
+          sampler_name= ['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms']  // k_lms
+          grid        = <boolean>     // false
+          width       = <integer>     // image width, multiple of 64 (512)
+          height      = <integer>     // image height, multiple of 64 (512)
+          cfg_scale   = <float>       // condition-free guidance scale (7.5)
+          )
+
+"""
+
+
+class Generate:
+    """Generate class
+    Stores default values for multiple configuration items
+    """
+
+    def __init__(
+            self,
+            iterations            = 1,
+            steps                 = 50,
+            cfg_scale             = 7.5,
+            weights               = 'models/ldm/stable-diffusion-v1/model.ckpt',
+            config                = 'configs/stable-diffusion/v1-inference.yaml',
+            grid                  = False,
+            width                 = 512,
+            height                = 512,
+            sampler_name          = 'k_lms',
+            ddim_eta              = 0.0,  # deterministic
+            precision             = 'autocast',
+            full_precision        = False,
+            strength              = 0.75,  # default in scripts/img2img.py
+            seamless              = False,
+            embedding_path        = None,
+            device_type           = 'cuda',
+    ):
+        self.iterations               = iterations
+        self.width                    = width
+        self.height                   = height
+        self.steps                    = steps
+        self.cfg_scale                = cfg_scale
+        self.weights                  = weights
+        self.config                   = config
+        self.sampler_name             = sampler_name
+        self.grid                     = grid
+        self.ddim_eta                 = ddim_eta
+        self.precision                = precision
+        self.full_precision           = True if choose_torch_device() == 'mps' else full_precision
+        self.strength                 = strength
+        self.seamless                 = seamless
+        self.embedding_path           = embedding_path
+        self.device_type              = device_type
+        self.model                    = None     # empty for now
+        self.sampler                  = None
+        self.device                   = None
+        self.generators               = {}
+        self.base_generator           = None
+        self.seed                     = None
+
+        if device_type == 'cuda' and not torch.cuda.is_available():
+            device_type = choose_torch_device()
+            print(">> cuda not available, using device", device_type)
+        self.device = torch.device(device_type)
+
+        # for VRAM usage statistics
+        device_type          = choose_torch_device()
+        self.session_peakmem = torch.cuda.max_memory_allocated() if device_type == 'cuda' else None
+        transformers.logging.set_verbosity_error()
+
+    def prompt2png(self, prompt, outdir, **kwargs):
+        """
+        Takes a prompt and an output directory, writes out the requested number
+        of PNG files, and returns an array of [[filename,seed],[filename,seed]...]
+        Optional named arguments are the same as those passed to Generate and prompt2image()
+        """
+        results = self.prompt2image(prompt, **kwargs)
+        pngwriter = PngWriter(outdir)
+        prefix = pngwriter.unique_prefix()
+        outputs = []
+        for image, seed in results:
+            name = f'{prefix}.{seed}.png'
+            path = pngwriter.save_image_and_prompt_to_png(
+                image, f'{prompt} -S{seed}', name)
+            outputs.append([path, seed])
+        return outputs
+
+    def txt2img(self, prompt, **kwargs):
+        outdir = kwargs.pop('outdir', 'outputs/img-samples')
+        return self.prompt2png(prompt, outdir, **kwargs)
+
+    def img2img(self, prompt, **kwargs):
+        outdir = kwargs.pop('outdir', 'outputs/img-samples')
+        assert (
+            'init_img' in kwargs
+        ), 'call to img2img() must include the init_img argument'
+        return self.prompt2png(prompt, outdir, **kwargs)
+
+    def prompt2image(
+            self,
+            # these are common
+            prompt,
+            iterations     =    None,
+            steps          =    None,
+            seed           =    None,
+            cfg_scale      =    None,
+            ddim_eta       =    None,
+            skip_normalize =    False,
+            image_callback =    None,
+            step_callback  =    None,
+            width          =    None,
+            height         =    None,
+            sampler_name   =    None,
+            seamless       =    False,
+            log_tokenization=  False,
+            with_variations =   None,
+            variation_amount =  0.0,
+            # these are specific to img2img and inpaint
+            init_img       =    None,
+            init_mask      =    None,
+            fit            =    False,
+            strength       =    None,
+            # these are specific to GFPGAN/ESRGAN
+            gfpgan_strength=    0,
+            save_original  =    False,
+            upscale        =    None,
+            **args,
+    ):   # eat up additional cruft
+        """
+        ldm.prompt2image() is the common entry point for txt2img() and img2img()
+        It takes the following arguments:
+           prompt                          // prompt string (no default)
+           iterations                      // iterations (1); image count=iterations
+           steps                           // refinement steps per iteration
+           seed                            // seed for random number generator
+           width                           // width of image, in multiples of 64 (512)
+           height                          // height of image, in multiples of 64 (512)
+           cfg_scale                       // how strongly the prompt influences the image (7.5) (must be >1)
+           seamless                        // whether the generated image should tile
+           init_img                        // path to an initial image
+           strength                        // strength for noising/unnoising init_img. 0.0 preserves image exactly, 1.0 replaces it completely
+           gfpgan_strength                 // strength for GFPGAN. 0.0 preserves image exactly, 1.0 replaces it completely
+           ddim_eta                        // image randomness (eta=0.0 means the same seed always produces the same image)
+           step_callback                   // a function or method that will be called each step
+           image_callback                  // a function or method that will be called each time an image is generated
+           with_variations                 // a weighted list [(seed_1, weight_1), (seed_2, weight_2), ...] of variations which should be applied before doing any generation
+           variation_amount                // optional 0-1 value to slerp from -S noise to random noise (allows variations on an image)
+
+        To use the step callback, define a function that receives two arguments:
+        - Image GPU data
+        - The step number
+
+        To use the image callback, define a function of method that receives two arguments, an Image object
+        and the seed. You can then do whatever you like with the image, including converting it to
+        different formats and manipulating it. For example:
+
+            def process_image(image,seed):
+                image.save(f{'images/seed.png'})
+
+        The callback used by the prompt2png() can be found in ldm/dream_util.py. It contains code
+        to create the requested output directory, select a unique informative name for each image, and
+        write the prompt into the PNG metadata.
+        """
+        # TODO: convert this into a getattr() loop
+        steps                 = steps      or self.steps
+        width                 = width      or self.width
+        height                = height     or self.height
+        seamless              = seamless   or self.seamless
+        cfg_scale             = cfg_scale  or self.cfg_scale
+        ddim_eta              = ddim_eta   or self.ddim_eta
+        iterations            = iterations or self.iterations
+        strength              = strength   or self.strength
+        self.seed             = seed
+        self.log_tokenization = log_tokenization
+        with_variations = [] if with_variations is None else with_variations
+
+        model = (
+            self.load_model()
+        )  # will instantiate the model or return it from cache
+
+        for m in model.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                m.padding_mode = 'circular' if seamless else m._orig_padding_mode
+        
+        assert cfg_scale > 1.0, 'CFG_Scale (-C) must be >1.0'
+        assert (
+            0.0 < strength < 1.0
+        ), 'img2img and inpaint strength can only work with 0.0 < strength < 1.0'
+        assert (
+                0.0 <= variation_amount <= 1.0
+        ), '-v --variation_amount must be in [0.0, 1.0]'
+
+        # check this logic - doesn't look right
+        if len(with_variations) > 0 or variation_amount > 1.0:
+            assert seed is not None,\
+                'seed must be specified when using with_variations'
+            if variation_amount == 0.0:
+                assert iterations == 1,\
+                    'when using --with_variations, multiple iterations are only possible when using --variation_amount'
+            assert all(0 <= weight <= 1 for _, weight in with_variations),\
+                f'variation weights must be in [0.0, 1.0]: got {[weight for _, weight in with_variations]}'
+
+        width, height, _ = self._resolution_check(width, height, log=True)
+
+        if sampler_name and (sampler_name != self.sampler_name):
+            self.sampler_name = sampler_name
+            self._set_sampler()
+
+        tic = time.time()
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+        results          = list()
+        init_image       = None
+        mask_image       = None
+
+        try:
+            uc, c = get_uc_and_c(
+                prompt, model=self.model,
+                skip_normalize=skip_normalize,
+                log_tokens=self.log_tokenization
+            )
+
+            (init_image,mask_image) = self._make_images(init_img,init_mask, width, height, fit)
+            
+            if (init_image is not None) and (mask_image is not None):
+                generator = self._make_inpaint()
+            elif init_image is not None:
+                generator = self._make_img2img()
+            else:
+                generator = self._make_txt2img()
+
+            generator.set_variation(self.seed, variation_amount, with_variations)
+            results = generator.generate(
+                prompt,
+                iterations     = iterations,
+                seed           = self.seed,
+                sampler        = self.sampler,
+                steps          = steps,
+                cfg_scale      = cfg_scale,
+                conditioning   = (uc,c),
+                ddim_eta       = ddim_eta,
+                image_callback = image_callback,  # called after the final image is generated
+                step_callback  = step_callback,   # called after each intermediate image is generated
+                width          = width,
+                height         = height,
+                init_image     = init_image,      # notice that init_image is different from init_img
+                mask_image     = mask_image,
+                strength       = strength,
+            )
+
+            if upscale is not None or gfpgan_strength > 0:
+                self.upscale_and_reconstruct(results,
+                                             upscale        = upscale,
+                                             strength       = gfpgan_strength,
+                                             save_original  = save_original,
+                                             image_callback = image_callback)
+
+        except KeyboardInterrupt:
+            print('*interrupted*')
+            print(
+                '>> Partial results will be returned; if --grid was requested, nothing will be returned.'
+            )
+        except RuntimeError as e:
+            print(traceback.format_exc(), file=sys.stderr)
+            print('>> Could not generate image.')
+
+        toc = time.time()
+        print('>> Usage stats:')
+        print(
+            f'>>   {len(results)} image(s) generated in', '%4.2fs' % (toc - tic)
+        )
+        print(
+            f'>>   Max VRAM used for this generation:',
+            '%4.2fG' % (torch.cuda.max_memory_allocated() / 1e9),
+        )
+
+        if self.session_peakmem:
+            self.session_peakmem = max(
+                self.session_peakmem, torch.cuda.max_memory_allocated()
+            )
+            print(
+                f'>>   Max VRAM used since script start: ',
+                '%4.2fG' % (self.session_peakmem / 1e9),
+            )
+        return results
+
+    def _make_images(self, img_path, mask_path, width, height, fit=False):
+        init_image      = None
+        init_mask       = None
+        if not img_path:
+            return None,None
+
+        image        = self._load_img(img_path, width, height, fit=fit) # this returns an Image
+        init_image   = self._create_init_image(image)                   # this returns a torch tensor
+
+        if self._has_transparency(image) and not mask_path:      # if image has a transparent area and no mask was provided, then try to generate mask
+            print('>> Initial image has transparent areas. Will inpaint in these regions.')
+            if self._check_for_erasure(image):
+                print(
+                    '>> WARNING: Colors underneath the transparent region seem to have been erased.\n',
+                    '>>          Inpainting will be suboptimal. Please preserve the colors when making\n',
+                    '>>          a transparency mask, or provide mask explicitly using --init_mask (-M).'
+                )
+            init_mask = self._create_init_mask(image)                   # this returns a torch tensor
+
+        if mask_path:
+            mask_image  = self._load_img(mask_path, width, height, fit=fit) # this returns an Image
+            init_mask   = self._create_init_mask(mask_image)
+
+        return init_image,init_mask
+
+    def _make_img2img(self):
+        if not self.generators.get('img2img'):
+            from ldm.dream.generator.img2img import Img2Img
+            self.generators['img2img'] = Img2Img(self.model)
+        return self.generators['img2img']
+
+    def _make_txt2img(self):
+        if not self.generators.get('txt2img'):
+            from ldm.dream.generator.txt2img import Txt2Img
+            self.generators['txt2img'] = Txt2Img(self.model)
+        return self.generators['txt2img']
+
+    def _make_inpaint(self):
+        if not self.generators.get('inpaint'):
+            from ldm.dream.generator.inpaint import Inpaint
+            self.generators['inpaint'] = Inpaint(self.model)
+        return self.generators['inpaint']
+
+    def load_model(self):
+        """Load and initialize the model from configuration variables passed at object creation time"""
+        if self.model is None:
+            seed_everything(random.randrange(0, np.iinfo(np.uint32).max))
+            try:
+                config = OmegaConf.load(self.config)
+                model = self._load_model_from_config(config, self.weights)
+                if self.embedding_path is not None:
+                    model.embedding_manager.load(
+                        self.embedding_path, self.full_precision
+                    )
+                self.model = model.to(self.device)
+                # model.to doesn't change the cond_stage_model.device used to move the tokenizer output, so set it here
+                self.model.cond_stage_model.device = self.device
+            except AttributeError as e:
+                print(f'>> Error loading model. {str(e)}', file=sys.stderr)
+                print(traceback.format_exc(), file=sys.stderr)
+                raise SystemExit from e
+
+            self._set_sampler()
+
+            for m in self.model.modules():
+                if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
+                    m._orig_padding_mode = m.padding_mode
+
+        return self.model
+
+    def upscale_and_reconstruct(self,
+                                image_list,
+                                upscale       = None,
+                                strength      =  0.0,
+                                save_original = False,
+                                image_callback = None):
+        try:
+            if upscale is not None:
+                from ldm.gfpgan.gfpgan_tools import real_esrgan_upscale
+            if strength > 0:
+                from ldm.gfpgan.gfpgan_tools import run_gfpgan
+        except (ModuleNotFoundError, ImportError):
+            print(traceback.format_exc(), file=sys.stderr)
+            print('>> You may need to install the ESRGAN and/or GFPGAN modules')
+            return
+            
+        for r in image_list:
+            image, seed = r
+            try:
+                if upscale is not None:
+                    if len(upscale) < 2:
+                        upscale.append(0.75)
+                    image = real_esrgan_upscale(
+                        image,
+                        upscale[1],
+                        int(upscale[0]),
+                        seed,
+                    )
+                if strength > 0:
+                    image = run_gfpgan(
+                        image, strength, seed, 1
+                    )
+            except Exception as e:
+                print(
+                    f'>> Error running RealESRGAN or GFPGAN. Your image was not upscaled.\n{e}'
+                )
+
+            if image_callback is not None:
+                image_callback(image, seed, upscaled=True)
+            else:
+                r[0] = image
+
+    # to help WebGUI - front end to generator util function
+    def sample_to_image(self,samples):
+        return self._sample_to_image(samples)
+
+    def _sample_to_image(self,samples):
+        if not self.base_generator:
+            from ldm.dream.generator import Generator
+            self.base_generator = Generator(self.model)
+        return self.base_generator.sample_to_image(samples)
+
+    def _set_sampler(self):
+        msg = f'>> Setting Sampler to {self.sampler_name}'
+        if self.sampler_name == 'plms':
+            self.sampler = PLMSSampler(self.model, device=self.device)
+        elif self.sampler_name == 'ddim':
+            self.sampler = DDIMSampler(self.model, device=self.device)
+        elif self.sampler_name == 'k_dpm_2_a':
+            self.sampler = KSampler(
+                self.model, 'dpm_2_ancestral', device=self.device
+            )
+        elif self.sampler_name == 'k_dpm_2':
+            self.sampler = KSampler(self.model, 'dpm_2', device=self.device)
+        elif self.sampler_name == 'k_euler_a':
+            self.sampler = KSampler(
+                self.model, 'euler_ancestral', device=self.device
+            )
+        elif self.sampler_name == 'k_euler':
+            self.sampler = KSampler(self.model, 'euler', device=self.device)
+        elif self.sampler_name == 'k_heun':
+            self.sampler = KSampler(self.model, 'heun', device=self.device)
+        elif self.sampler_name == 'k_lms':
+            self.sampler = KSampler(self.model, 'lms', device=self.device)
+        else:
+            msg = f'>> Unsupported Sampler: {self.sampler_name}, Defaulting to plms'
+            self.sampler = PLMSSampler(self.model, device=self.device)
+
+        print(msg)
+
+    def _load_model_from_config(self, config, ckpt):
+        print(f'>> Loading model from {ckpt}')
+
+        # for usage statistics
+        device_type = choose_torch_device()
+        if device_type == 'cuda':
+            torch.cuda.reset_peak_memory_stats() 
+        tic = time.time()
+
+        # this does the work
+        pl_sd = torch.load(ckpt, map_location='cpu')
+        sd = pl_sd['state_dict']
+        model = instantiate_from_config(config.model)
+        m, u = model.load_state_dict(sd, strict=False)
+        model.to(self.device)
+        model.eval()
+
+        
+        if self.full_precision:
+            print(
+                '>> Using slower but more accurate full-precision math (--full_precision)'
+            )
+        else:
+            print(
+                '>> Using half precision math. Call with --full_precision to use more accurate but VRAM-intensive full precision.'
+            )
+            model.half()
+
+        # usage statistics
+        toc = time.time()
+        print(
+            f'>> Model loaded in', '%4.2fs' % (toc - tic)
+        )
+        if device_type == 'cuda':
+            print(
+                '>> Max VRAM used to load the model:',
+                '%4.2fG' % (torch.cuda.max_memory_allocated() / 1e9),
+                '\n>> Current VRAM usage:'
+                '%4.2fG' % (torch.cuda.memory_allocated() / 1e9),
+            )
+
+        return model
+
+    def _load_img(self, path, width, height, fit=False):
+        assert os.path.exists(path), f'>> {path}: File not found'
+
+        #        with Image.open(path) as img:
+        #            image = img.convert('RGBA')
+        image = Image.open(path)
+        print(
+            f'>> loaded input image of size {image.width}x{image.height} from {path}'
+        )
+        if fit:
+            image = self._fit_image(image,(width,height))
+        else:
+            image = self._squeeze_image(image)
+        return image
+
+    def _create_init_image(self,image):
+        image = image.convert('RGB')
+        # print(
+        #     f'>> DEBUG: writing the image to img.png'
+        # )
+        # image.save('img.png')
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image[None].transpose(0, 3, 1, 2)
+        image = torch.from_numpy(image)
+        image = 2.0 * image - 1.0 
+        return image.to(self.device)
+
+    def _create_init_mask(self, image):
+        # convert into a black/white mask
+        image = self._image_to_mask(image)
+        image = image.convert('RGB')
+        # BUG: We need to use the model's downsample factor rather than hardcoding "8"
+        from ldm.dream.generator.base import downsampling
+        image = image.resize((image.width//downsampling, image.height//downsampling), resample=Image.Resampling.LANCZOS)
+        # print(
+        #     f'>> DEBUG: writing the mask to mask.png'
+        #     )
+        # image.save('mask.png')
+        image = np.array(image)
+        image = image.astype(np.float32) / 255.0
+        image = image[None].transpose(0, 3, 1, 2)
+        image = torch.from_numpy(image)
+        return image.to(self.device)
+
+    # The mask is expected to have the region to be inpainted
+    # with alpha transparency. It converts it into a black/white
+    # image with the transparent part black.
+    def _image_to_mask(self, mask_image, invert=False) -> Image:
+        # Obtain the mask from the transparency channel
+        mask = Image.new(mode="L", size=mask_image.size, color=255)
+        mask.putdata(mask_image.getdata(band=3))
+        if invert:
+            mask = ImageOps.invert(mask)
+        return mask
+
+    def _has_transparency(self,image):
+        if image.info.get("transparency", None) is not None:
+            return True
+        if image.mode == "P":
+            transparent = image.info.get("transparency", -1)
+            for _, index in image.getcolors():
+                if index == transparent:
+                    return True
+        elif image.mode == "RGBA":
+            extrema = image.getextrema()
+            if extrema[3][0] < 255:
+                return True
+        return False
+
+    
+    def _check_for_erasure(self,image):
+        width, height = image.size
+        pixdata       = image.load()
+        colored       = 0
+        for y in range(height):
+            for x in range(width):
+                if pixdata[x, y][3] == 0:
+                    r, g, b, _ = pixdata[x, y]
+                    if (r, g, b) != (0, 0, 0) and \
+                       (r, g, b) != (255, 255, 255):
+                        colored += 1
+        return colored == 0
+
+    def _squeeze_image(self,image):
+        x,y,resize_needed = self._resolution_check(image.width,image.height)
+        if resize_needed:
+            return InitImageResizer(image).resize(x,y)
+        return image
+
+
+    def _fit_image(self,image,max_dimensions):
+        w,h = max_dimensions
+        print(
+            f'>> image will be resized to fit inside a box {w}x{h} in size.'
+        )
+        if image.width > image.height:
+            h   = None   # by setting h to none, we tell InitImageResizer to fit into the width and calculate height
+        elif image.height > image.width:
+            w   = None   # ditto for w
+        else:
+            pass
+        image = InitImageResizer(image).resize(w,h)   # note that InitImageResizer does the multiple of 64 truncation internally
+        print(
+            f'>> after adjusting image dimensions to be multiples of 64, init image is {image.width}x{image.height}'
+            )
+        return image
+
+    def _resolution_check(self, width, height, log=False):
+        resize_needed = False
+        w, h = map(
+            lambda x: x - x % 64, (width, height)
+        )  # resize to integer multiple of 64
+        if h != height or w != width:
+            if log:
+                print(
+                    f'>> Provided width and height must be multiples of 64. Auto-resizing to {w}x{h}'
+                )
+            height = h
+            width  = w
+            resize_needed = True
+
+        if (width * height) > (self.width * self.height):
+            print(">> This input is larger than your defaults. If you run out of memory, please use a smaller image.")
+
+        return width, height, resize_needed
+
+

ldm/gfpgan/gfpgan_tools.py

@@ -13,8 +13,8 @@ opt = arg_parser.parse_args()
 model_path = os.path.join(opt.gfpgan_dir, opt.gfpgan_model_path)
 gfpgan_model_exists = os.path.isfile(model_path)
 
-def _run_gfpgan(image, strength, prompt, seed, upsampler_scale=4):
-    print(f'>> GFPGAN - Restoring Faces: {prompt} : seed:{seed}')
+def run_gfpgan(image, strength, seed, upsampler_scale=4):
+    print(f'>> GFPGAN - Restoring Faces for image seed:{seed}')
     gfpgan = None
     with warnings.catch_warnings():
         warnings.filterwarnings('ignore', category=DeprecationWarning)
@@ -46,7 +46,7 @@ def _run_gfpgan(image, strength, prompt, seed, upsampler_scale=4):
 
     if gfpgan is None:
         print(
-            f'>> GFPGAN not initialized, it must be loaded via the --gfpgan argument'
+            f'>> GFPGAN not initialized. Their packages must be installed as siblings to the "stable-diffusion" folder, or set explicitly using the --gfpgan_dir option.'
         )
         return image
 
@@ -127,9 +127,9 @@ def _load_gfpgan_bg_upsampler(bg_upsampler, upsampler_scale, bg_tile=400):
     return bg_upsampler
 
 
-def real_esrgan_upscale(image, strength, upsampler_scale, prompt, seed):
+def real_esrgan_upscale(image, strength, upsampler_scale, seed):
     print(
-        f'>> Real-ESRGAN Upscaling: {prompt} : seed:{seed} : scale:{upsampler_scale}x'
+        f'>> Real-ESRGAN Upscaling seed:{seed} : scale:{upsampler_scale}x'
     )
 
     with warnings.catch_warnings():

ldm/models/diffusion/ddim.py

@@ -171,6 +171,7 @@ class DDIMSampler(object):
         )
         return samples, intermediates
 
+    # This routine gets called from img2img
     @torch.no_grad()
     def ddim_sampling(
         self,
@@ -270,6 +271,7 @@ class DDIMSampler(object):
 
         return img, intermediates
 
+    # This routine gets called from ddim_sampling() and decode()
     @torch.no_grad()
     def p_sample_ddim(
         self,
@@ -372,14 +374,16 @@ class DDIMSampler(object):
 
     @torch.no_grad()
     def decode(
-        self,
-        x_latent,
-        cond,
-        t_start,
-        img_callback=None,
-        unconditional_guidance_scale=1.0,
-        unconditional_conditioning=None,
-        use_original_steps=False,
+            self,
+            x_latent,
+            cond,
+            t_start,
+            img_callback=None,
+            unconditional_guidance_scale=1.0,
+            unconditional_conditioning=None,
+            use_original_steps=False,
+            init_latent       = None,
+            mask              = None,
     ):
 
         timesteps = (
@@ -395,6 +399,8 @@ class DDIMSampler(object):
 
         iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
         x_dec = x_latent
+        x0    = init_latent
+
         for i, step in enumerate(iterator):
             index = total_steps - i - 1
             ts = torch.full(
@@ -403,6 +409,14 @@ class DDIMSampler(object):
                 device=x_latent.device,
                 dtype=torch.long,
             )
+
+            if mask is not None:
+                assert x0 is not None
+                xdec_orig = self.model.q_sample(
+                    x0, ts
+                )  # TODO: deterministic forward pass?
+                x_dec = xdec_orig * mask + (1.0 - mask) * x_dec
+
             x_dec, _ = self.p_sample_ddim(
                 x_dec,
                 cond,
@@ -412,6 +426,7 @@ class DDIMSampler(object):
                 unconditional_guidance_scale=unconditional_guidance_scale,
                 unconditional_conditioning=unconditional_conditioning,
             )
+
             if img_callback:
                 img_callback(x_dec, i)
 

ldm/modules/attention.py

@@ -7,13 +7,14 @@ from einops import rearrange, repeat
 
 from ldm.modules.diffusionmodules.util import checkpoint
 
+import psutil
 
 def exists(val):
     return val is not None
 
 
 def uniq(arr):
-    return {el: True for el in arr}.keys()
+    return{el: True for el in arr}.keys()
 
 
 def default(val, d):
@@ -45,18 +46,19 @@ class GEGLU(nn.Module):
 
 
 class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
         super().__init__()
         inner_dim = int(dim * mult)
         dim_out = default(dim_out, dim)
-        project_in = (
-            nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
-            if not glu
-            else GEGLU(dim, inner_dim)
-        )
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            nn.GELU()
+        ) if not glu else GEGLU(dim, inner_dim)
 
         self.net = nn.Sequential(
-            project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
         )
 
     def forward(self, x):
@@ -73,9 +75,7 @@ def zero_module(module):
 
 
 def Normalize(in_channels):
-    return torch.nn.GroupNorm(
-        num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
-    )
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
 
 
 class LinearAttention(nn.Module):
@@ -83,28 +83,17 @@ class LinearAttention(nn.Module):
         super().__init__()
         self.heads = heads
         hidden_dim = dim_head * heads
-        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias = False)
         self.to_out = nn.Conv2d(hidden_dim, dim, 1)
 
     def forward(self, x):
         b, c, h, w = x.shape
         qkv = self.to_qkv(x)
-        q, k, v = rearrange(
-            qkv,
-            'b (qkv heads c) h w -> qkv b heads c (h w)',
-            heads=self.heads,
-            qkv=3,
-        )
-        k = k.softmax(dim=-1)
+        q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
+        k = k.softmax(dim=-1)  
         context = torch.einsum('bhdn,bhen->bhde', k, v)
         out = torch.einsum('bhde,bhdn->bhen', context, q)
-        out = rearrange(
-            out,
-            'b heads c (h w) -> b (heads c) h w',
-            heads=self.heads,
-            h=h,
-            w=w,
-        )
+        out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
         return self.to_out(out)
 
 
@@ -114,18 +103,26 @@ class SpatialSelfAttention(nn.Module):
         self.in_channels = in_channels
 
         self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.k = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.v = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
-        self.proj_out = torch.nn.Conv2d(
-            in_channels, in_channels, kernel_size=1, stride=1, padding=0
-        )
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
 
     def forward(self, x):
         h_ = x
@@ -135,12 +132,12 @@ class SpatialSelfAttention(nn.Module):
         v = self.v(h_)
 
         # compute attention
-        b, c, h, w = q.shape
+        b,c,h,w = q.shape
         q = rearrange(q, 'b c h w -> b (h w) c')
         k = rearrange(k, 'b c h w -> b c (h w)')
         w_ = torch.einsum('bij,bjk->bik', q, k)
 
-        w_ = w_ * (int(c) ** (-0.5))
+        w_ = w_ * (int(c)**(-0.5))
         w_ = torch.nn.functional.softmax(w_, dim=2)
 
         # attend to values
@@ -150,18 +147,16 @@ class SpatialSelfAttention(nn.Module):
         h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
         h_ = self.proj_out(h_)
 
-        return x + h_
+        return x+h_
 
 
 class CrossAttention(nn.Module):
-    def __init__(
-        self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0
-    ):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
         super().__init__()
         inner_dim = dim_head * heads
         context_dim = default(context_dim, query_dim)
 
-        self.scale = dim_head**-0.5
+        self.scale = dim_head ** -0.5
         self.heads = heads
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
@@ -169,69 +164,83 @@ class CrossAttention(nn.Module):
         self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
 
         self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+            nn.Linear(inner_dim, query_dim),
+            nn.Dropout(dropout)
         )
 
     def forward(self, x, context=None, mask=None):
         h = self.heads
 
-        q = self.to_q(x)
+        q_in = self.to_q(x)
         context = default(context, x)
-        k = self.to_k(context)
-        v = self.to_v(context)
+        k_in = self.to_k(context)
+        v_in = self.to_v(context)
+        device_type = 'mps' if x.device.type == 'mps' else 'cuda'
+        del context, x
 
-        q, k, v = map(
-            lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)
-        )
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
+        del q_in, k_in, v_in
 
-        sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)
 
-        if exists(mask):
-            mask = rearrange(mask, 'b ... -> b (...)')
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
-            sim.masked_fill_(~mask, max_neg_value)
+        if device_type == 'mps':
+            mem_free_total = psutil.virtual_memory().available
+        else:
+            stats = torch.cuda.memory_stats(q.device)
+            mem_active = stats['active_bytes.all.current']
+            mem_reserved = stats['reserved_bytes.all.current']
+            mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
+            mem_free_torch = mem_reserved - mem_active
+            mem_free_total = mem_free_cuda + mem_free_torch
 
-        # attention, what we cannot get enough of
-        attn = sim.softmax(dim=-1)
+        gb = 1024 ** 3
+        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * 4
+        mem_required = tensor_size * 2.5
+        steps = 1
 
-        out = einsum('b i j, b j d -> b i d', attn, v)
-        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
-        return self.to_out(out)
+        if mem_required > mem_free_total:
+            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
+            # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
+            #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
+
+        if steps > 64:
+            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
+            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
+                               f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
+
+        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
+        for i in range(0, q.shape[1], slice_size):
+            end = i + slice_size
+            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale
+
+            s2 = s1.softmax(dim=-1)
+            del s1
+
+            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
+            del s2
+
+        del q, k, v
+
+        r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
+        del r1
+
+        return self.to_out(r2)
 
 
 class BasicTransformerBlock(nn.Module):
-    def __init__(
-        self,
-        dim,
-        n_heads,
-        d_head,
-        dropout=0.0,
-        context_dim=None,
-        gated_ff=True,
-        checkpoint=True,
-    ):
+    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):
         super().__init__()
-        self.attn1 = CrossAttention(
-            query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout
-        )  # is a self-attention
+        self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)  # is a self-attention
         self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
-        self.attn2 = CrossAttention(
-            query_dim=dim,
-            context_dim=context_dim,
-            heads=n_heads,
-            dim_head=d_head,
-            dropout=dropout,
-        )  # is self-attn if context is none
+        self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
+                                    heads=n_heads, dim_head=d_head, dropout=dropout)  # is self-attn if context is none
         self.norm1 = nn.LayerNorm(dim)
         self.norm2 = nn.LayerNorm(dim)
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
     def forward(self, x, context=None):
-        return checkpoint(
-            self._forward, (x, context), self.parameters(), self.checkpoint
-        )
+        return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
 
     def _forward(self, x, context=None):
         x = x.contiguous() if x.device.type == 'mps' else x
@@ -249,43 +258,29 @@ class SpatialTransformer(nn.Module):
     Then apply standard transformer action.
     Finally, reshape to image
     """
-
-    def __init__(
-        self,
-        in_channels,
-        n_heads,
-        d_head,
-        depth=1,
-        dropout=0.0,
-        context_dim=None,
-    ):
+    def __init__(self, in_channels, n_heads, d_head,
+                 depth=1, dropout=0., context_dim=None):
         super().__init__()
         self.in_channels = in_channels
         inner_dim = n_heads * d_head
         self.norm = Normalize(in_channels)
 
-        self.proj_in = nn.Conv2d(
-            in_channels, inner_dim, kernel_size=1, stride=1, padding=0
-        )
+        self.proj_in = nn.Conv2d(in_channels,
+                                 inner_dim,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
 
         self.transformer_blocks = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    inner_dim,
-                    n_heads,
-                    d_head,
-                    dropout=dropout,
-                    context_dim=context_dim,
-                )
-                for d in range(depth)
-            ]
+            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
+                for d in range(depth)]
         )
 
-        self.proj_out = zero_module(
-            nn.Conv2d(
-                inner_dim, in_channels, kernel_size=1, stride=1, padding=0
-            )
-        )
+        self.proj_out = zero_module(nn.Conv2d(inner_dim,
+                                              in_channels,
+                                              kernel_size=1,
+                                              stride=1,
+                                              padding=0))
 
     def forward(self, x, context=None):
         # note: if no context is given, cross-attention defaults to self-attention

ldm/modules/diffusionmodules/util.py

@@ -81,7 +81,9 @@ def make_ddim_timesteps(
 
     # assert ddim_timesteps.shape[0] == num_ddim_timesteps
     # add one to get the final alpha values right (the ones from first scale to data during sampling)
-    steps_out = ddim_timesteps + 1
+#    steps_out = ddim_timesteps + 1
+    steps_out = ddim_timesteps
+
     if verbose:
         print(f'Selected timesteps for ddim sampler: {steps_out}')
     return steps_out

ldm/simplet2i.py

@@ -1,843 +1,13 @@
-# Copyright (c) 2022 Lincoln D. Stein (https://github.com/lstein)
+'''
+This module is provided for backward compatibility with the
+original (hasty) API.
 
-# Derived from source code carrying the following copyrights
-# Copyright (c) 2022 Machine Vision and Learning Group, LMU Munich
-# Copyright (c) 2022 Robin Rombach and Patrick Esser and contributors
+Please use ldm.generate instead.
+'''
 
-import torch
-import numpy as np
-import random
-import os
-import traceback
-from omegaconf import OmegaConf
-from PIL import Image
-from tqdm import tqdm, trange
-from itertools import islice
-from einops import rearrange, repeat
-from torchvision.utils import make_grid
-from pytorch_lightning import seed_everything
-from torch import autocast
-from contextlib import contextmanager, nullcontext
-import transformers
-import time
-import re
-import sys
+from ldm.generate import Generate
 
-from ldm.util                      import instantiate_from_config
-from ldm.models.diffusion.ddim     import DDIMSampler
-from ldm.models.diffusion.plms     import PLMSSampler
-from ldm.models.diffusion.ksampler import KSampler
-from ldm.dream.pngwriter           import PngWriter
-from ldm.dream.image_util          import InitImageResizer
-from ldm.dream.devices import choose_autocast_device, choose_torch_device
-
-"""Simplified text to image API for stable diffusion/latent diffusion
-
-Example Usage:
-
-from ldm.simplet2i import T2I
-
-# Create an object with default values
-t2i = T2I(model       = <path>        // models/ldm/stable-diffusion-v1/model.ckpt
-          config      = <path>        // configs/stable-diffusion/v1-inference.yaml
-          iterations  = <integer>     // how many times to run the sampling (1)
-          steps       = <integer>     // 50
-          seed        = <integer>     // current system time
-          sampler_name= ['ddim', 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler', 'k_heun', 'k_lms', 'plms']  // k_lms
-          grid        = <boolean>     // false
-          width       = <integer>     // image width, multiple of 64 (512)
-          height      = <integer>     // image height, multiple of 64 (512)
-          cfg_scale   = <float>       // unconditional guidance scale (7.5)
-          )
-
-# do the slow model initialization
-t2i.load_model()
-
-# Do the fast inference & image generation. Any options passed here
-# override the default values assigned during class initialization
-# Will call load_model() if the model was not previously loaded and so
-# may be slow at first.
-# The method returns a list of images. Each row of the list is a sub-list of [filename,seed]
-results = t2i.prompt2png(prompt     = "an astronaut riding a horse",
-                         outdir     = "./outputs/samples",
-                         iterations = 3)
-
-for row in results:
-    print(f'filename={row[0]}')
-    print(f'seed    ={row[1]}')
-
-# Same thing, but using an initial image.
-results = t2i.prompt2png(prompt   = "an astronaut riding a horse",
-                         outdir   = "./outputs/,
-                         iterations = 3,
-                         init_img = "./sketches/horse+rider.png")
-
-for row in results:
-    print(f'filename={row[0]}')
-    print(f'seed    ={row[1]}')
-
-# Same thing, but we return a series of Image objects, which lets you manipulate them,
-# combine them, and save them under arbitrary names
-
-results = t2i.prompt2image(prompt   = "an astronaut riding a horse"
-                           outdir   = "./outputs/")
-for row in results:
-    im   = row[0]
-    seed = row[1]
-    im.save(f'./outputs/samples/an_astronaut_riding_a_horse-{seed}.png')
-    im.thumbnail(100,100).save('./outputs/samples/astronaut_thumb.jpg')
-
-Note that the old txt2img() and img2img() calls are deprecated but will
-still work.
-"""
-
-
-class T2I:
-    """T2I class
-        Attributes
-        ----------
-        model
-        config
-        iterations
-        steps
-        seed
-        sampler_name
-        width
-        height
-        cfg_scale
-        latent_channels
-        downsampling_factor
-        precision
-        strength
-        embedding_path
-
-    The vast majority of these arguments default to reasonable values.
-"""
-
-    def __init__(
-            self,
-            iterations=1,
-            steps=50,
-            seed=None,
-            cfg_scale=7.5,
-            weights='models/ldm/stable-diffusion-v1/model.ckpt',
-            config='configs/stable-diffusion/v1-inference.yaml',
-            grid=False,
-            width=512,
-            height=512,
-            sampler_name='k_lms',
-            latent_channels=4,
-            downsampling_factor=8,
-            ddim_eta=0.0,  # deterministic
-            precision='autocast',
-            full_precision=False,
-            strength=0.75,  # default in scripts/img2img.py
-            embedding_path=None,
-            device_type = 'cuda',
-            # just to keep track of this parameter when regenerating prompt
-            # needs to be replaced when new configuration system implemented.
-            latent_diffusion_weights=False,
-    ):
-        self.iterations               = iterations
-        self.width                    = width
-        self.height                   = height
-        self.steps                    = steps
-        self.cfg_scale                = cfg_scale
-        self.weights                  = weights
-        self.config                   = config
-        self.sampler_name             = sampler_name
-        self.latent_channels          = latent_channels
-        self.downsampling_factor      = downsampling_factor
-        self.grid                     = grid
-        self.ddim_eta                 = ddim_eta
-        self.precision                = precision
-        self.full_precision           = True if choose_torch_device() == 'mps' else full_precision
-        self.strength                 = strength
-        self.embedding_path           = embedding_path
-        self.device_type              = device_type
-        self.model                    = None     # empty for now
-        self.sampler                  = None
-        self.device                   = None
-        self.latent_diffusion_weights = latent_diffusion_weights
-
-        if device_type == 'cuda' and not torch.cuda.is_available():
-            device_type = choose_torch_device()
-            print(">> cuda not available, using device", device_type)
-        self.device = torch.device(device_type)
-
-        # for VRAM usage statistics
-        device_type          = choose_torch_device()
-        self.session_peakmem = torch.cuda.max_memory_allocated() if device_type == 'cuda' else None
-
-        if seed is None:
-            self.seed = self._new_seed()
-        else:
-            self.seed = seed
-        transformers.logging.set_verbosity_error()
-
-    def prompt2png(self, prompt, outdir, **kwargs):
-        """
-        Takes a prompt and an output directory, writes out the requested number
-        of PNG files, and returns an array of [[filename,seed],[filename,seed]...]
-        Optional named arguments are the same as those passed to T2I and prompt2image()
-        """
-        results = self.prompt2image(prompt, **kwargs)
-        pngwriter = PngWriter(outdir)
-        prefix = pngwriter.unique_prefix()
-        outputs = []
-        for image, seed in results:
-            name = f'{prefix}.{seed}.png'
-            path = pngwriter.save_image_and_prompt_to_png(
-                image, f'{prompt} -S{seed}', name)
-            outputs.append([path, seed])
-        return outputs
-
-    def txt2img(self, prompt, **kwargs):
-        outdir = kwargs.pop('outdir', 'outputs/img-samples')
-        return self.prompt2png(prompt, outdir, **kwargs)
-
-    def img2img(self, prompt, **kwargs):
-        outdir = kwargs.pop('outdir', 'outputs/img-samples')
-        assert (
-            'init_img' in kwargs
-        ), 'call to img2img() must include the init_img argument'
-        return self.prompt2png(prompt, outdir, **kwargs)
-
-    def prompt2image(
-            self,
-            # these are common
-            prompt,
-            iterations     =    None,
-            steps          =    None,
-            seed           =    None,
-            cfg_scale      =    None,
-            ddim_eta       =    None,
-            skip_normalize =    False,
-            image_callback =    None,
-            step_callback  =    None,
-            width          =    None,
-            height         =    None,
-            # these are specific to img2img
-            init_img       =    None,
-            fit            =    False,
-            strength       =    None,
-            gfpgan_strength=    0,
-            save_original  =    False,
-            upscale        =    None,
-            sampler_name   =    None,
-            log_tokenization=  False,
-            with_variations =   None,
-            variation_amount =  0.0,
-            **args,
-    ):   # eat up additional cruft
-        """
-        ldm.prompt2image() is the common entry point for txt2img() and img2img()
-        It takes the following arguments:
-           prompt                          // prompt string (no default)
-           iterations                      // iterations (1); image count=iterations
-           steps                           // refinement steps per iteration
-           seed                            // seed for random number generator
-           width                           // width of image, in multiples of 64 (512)
-           height                          // height of image, in multiples of 64 (512)
-           cfg_scale                       // how strongly the prompt influences the image (7.5) (must be >1)
-           init_img                        // path to an initial image - its dimensions override width and height
-           strength                        // strength for noising/unnoising init_img. 0.0 preserves image exactly, 1.0 replaces it completely
-           gfpgan_strength                 // strength for GFPGAN. 0.0 preserves image exactly, 1.0 replaces it completely
-           ddim_eta                        // image randomness (eta=0.0 means the same seed always produces the same image)
-           step_callback                   // a function or method that will be called each step
-           image_callback                  // a function or method that will be called each time an image is generated
-           with_variations                 // a weighted list [(seed_1, weight_1), (seed_2, weight_2), ...] of variations which should be applied before doing any generation
-           variation_amount                // optional 0-1 value to slerp from -S noise to random noise (allows variations on an image)
-
-        To use the step callback, define a function that receives two arguments:
-        - Image GPU data
-        - The step number
-
-        To use the image callback, define a function of method that receives two arguments, an Image object
-        and the seed. You can then do whatever you like with the image, including converting it to
-        different formats and manipulating it. For example:
-
-            def process_image(image,seed):
-                image.save(f{'images/seed.png'})
-
-        The callback used by the prompt2png() can be found in ldm/dream_util.py. It contains code
-        to create the requested output directory, select a unique informative name for each image, and
-        write the prompt into the PNG metadata.
-        """
-        # TODO: convert this into a getattr() loop
-        steps                 = steps      or self.steps
-        width                 = width      or self.width
-        height                = height     or self.height
-        cfg_scale             = cfg_scale  or self.cfg_scale
-        ddim_eta              = ddim_eta   or self.ddim_eta
-        iterations            = iterations or self.iterations
-        strength              = strength   or self.strength
-        self.log_tokenization = log_tokenization
-        with_variations = [] if with_variations is None else with_variations
-
-        model = (
-            self.load_model()
-        )  # will instantiate the model or return it from cache
-        assert cfg_scale > 1.0, 'CFG_Scale (-C) must be >1.0'
-        assert (
-            0.0 <= strength <= 1.0
-        ), 'can only work with strength in [0.0, 1.0]'
-        assert (
-                0.0 <= variation_amount <= 1.0
-        ), '-v --variation_amount must be in [0.0, 1.0]'
-
-        if len(with_variations) > 0 or variation_amount > 1.0:
-            assert seed is not None,\
-                'seed must be specified when using with_variations'
-            if variation_amount == 0.0:
-                assert iterations == 1,\
-                    'when using --with_variations, multiple iterations are only possible when using --variation_amount'
-            assert all(0 <= weight <= 1 for _, weight in with_variations),\
-                f'variation weights must be in [0.0, 1.0]: got {[weight for _, weight in with_variations]}'
-
-        seed                  = seed       or self.seed
-        width, height, _ = self._resolution_check(width, height, log=True)
-
-        # TODO: - Check if this is still necessary to run on M1 devices.
-        #       - Move code into ldm.dream.devices to live alongside other
-        #         special-hardware casing code.
-        if self.precision == 'autocast' and torch.cuda.is_available():
-            scope = autocast
-        else:
-            scope = nullcontext
-
-        if sampler_name and (sampler_name != self.sampler_name):
-            self.sampler_name = sampler_name
-            self._set_sampler()
-
-        tic = time.time()
-        if torch.cuda.is_available():
-            torch.cuda.reset_peak_memory_stats()
-        results = list()
-
-        try:
-            if init_img:
-                assert os.path.exists(init_img), f'{init_img}: File not found'
-                init_image = self._load_img(init_img, width, height, fit).to(self.device)
-                with scope(self.device.type):
-                    init_latent = self.model.get_first_stage_encoding(
-                        self.model.encode_first_stage(init_image)
-                    ) # move to latent space
-
-                print(f' DEBUG: seed at make_image time ={seed}')
-                make_image = self._img2img(
-                    prompt,
-                    steps=steps,
-                    cfg_scale=cfg_scale,
-                    ddim_eta=ddim_eta,
-                    skip_normalize=skip_normalize,
-                    init_latent=init_latent,
-                    strength=strength,
-                    callback=step_callback,
-                )
-            else:
-                make_image = self._txt2img(
-                    prompt,
-                    steps=steps,
-                    cfg_scale=cfg_scale,
-                    ddim_eta=ddim_eta,
-                    skip_normalize=skip_normalize,
-                    width=width,
-                    height=height,
-                    callback=step_callback,
-                )
-
-            initial_noise = None
-            if variation_amount > 0 or len(with_variations) > 0:
-                # use fixed initial noise plus random noise per iteration
-                seed_everything(seed)
-                initial_noise = self._get_noise(init_img,width,height)
-                for v_seed, v_weight in with_variations:
-                    seed = v_seed
-                    seed_everything(seed)
-                    next_noise = self._get_noise(init_img,width,height)
-                    initial_noise = self.slerp(v_weight, initial_noise, next_noise)
-                if variation_amount > 0:
-                    random.seed() # reset RNG to an actually random state, so we can get a random seed for variations
-                    seed = random.randrange(0,np.iinfo(np.uint32).max)
-
-            device_type = choose_autocast_device(self.device)
-            with scope(device_type), self.model.ema_scope():
-                for n in trange(iterations, desc='Generating'):
-                    x_T = None
-                    if variation_amount > 0:
-                        seed_everything(seed)
-                        target_noise = self._get_noise(init_img,width,height)
-                        x_T = self.slerp(variation_amount, initial_noise, target_noise)
-                    elif initial_noise is not None:
-                        # i.e. we specified particular variations
-                        x_T = initial_noise
-                    else:
-                        seed_everything(seed)
-                        if self.device.type == 'mps':
-                            x_T = self._get_noise(init_img,width,height)
-                        # make_image will do the equivalent of get_noise itself
-                    print(f' DEBUG: seed at make_image() invocation time ={seed}')
-                    image = make_image(x_T)
-                    results.append([image, seed])
-                    if image_callback is not None:
-                        image_callback(image, seed)
-                    seed = self._new_seed()
-
-                if upscale is not None or gfpgan_strength > 0:
-                    for result in results:
-                        image, seed = result
-                        try:
-                            if upscale is not None:
-                                from ldm.gfpgan.gfpgan_tools import (
-                                    real_esrgan_upscale,
-                                )
-                                if len(upscale) < 2:
-                                    upscale.append(0.75)
-                                image = real_esrgan_upscale(
-                                    image,
-                                    upscale[1],
-                                    int(upscale[0]),
-                                    prompt,
-                                    seed,
-                                )
-                            if gfpgan_strength > 0:
-                                from ldm.gfpgan.gfpgan_tools import _run_gfpgan
-
-                                image = _run_gfpgan(
-                                    image, gfpgan_strength, prompt, seed, 1
-                                )
-                        except Exception as e:
-                            print(
-                                f'>> Error running RealESRGAN - Your image was not upscaled.\n{e}'
-                            )
-                        if image_callback is not None:
-                            if save_original:
-                                image_callback(image, seed)
-                            else:
-                                image_callback(image, seed, upscaled=True)
-                        else:  # no callback passed, so we simply replace old image with rescaled one
-                            result[0] = image
-
-        except KeyboardInterrupt:
-            print('*interrupted*')
-            print(
-                '>> Partial results will be returned; if --grid was requested, nothing will be returned.'
-            )
-        except RuntimeError as e:
-            print(traceback.format_exc(), file=sys.stderr)
-            print('>> Are you sure your system has an adequate NVIDIA GPU?')
-
-        toc = time.time()
-        print('>> Usage stats:')
-        print(
-            f'>>   {len(results)} image(s) generated in', '%4.2fs' % (toc - tic)
-        )
-        print(
-            f'>>   Max VRAM used for this generation:',
-            '%4.2fG' % (torch.cuda.max_memory_allocated() / 1e9),
-        )
-
-        if self.session_peakmem:
-            self.session_peakmem = max(
-                self.session_peakmem, torch.cuda.max_memory_allocated()
-            )
-            print(
-                f'>>   Max VRAM used since script start: ',
-                '%4.2fG' % (self.session_peakmem / 1e9),
-            )
-        return results
-
-    @torch.no_grad()
-    def _txt2img(
-        self,
-        prompt,
-        steps,
-        cfg_scale,
-        ddim_eta,
-        skip_normalize,
-        width,
-        height,
-        callback,
-    ):
-        """
-        Returns a function returning an image derived from the prompt and the initial image
-        Return value depends on the seed at the time you call it
-        """
-
-        sampler = self.sampler
-
-        def make_image(x_T):
-            uc, c = self._get_uc_and_c(prompt, skip_normalize)
-            shape = [
-                self.latent_channels,
-                height // self.downsampling_factor,
-                width // self.downsampling_factor,
-            ]
-            samples, _ = sampler.sample(
-                batch_size=1,
-                S=steps,
-                x_T=x_T,
-                conditioning=c,
-                shape=shape,
-                verbose=False,
-                unconditional_guidance_scale=cfg_scale,
-                unconditional_conditioning=uc,
-                eta=ddim_eta,
-                img_callback=callback
-            )
-            return self._sample_to_image(samples)
-        return make_image
-
-    @torch.no_grad()
-    def _img2img(
-            self,
-            prompt,
-            steps,
-            cfg_scale,
-            ddim_eta,
-            skip_normalize,
-            init_latent,
-            strength,
-            callback,  # Currently not implemented for img2img
-    ):
-        """
-        Returns a function returning an image derived from the prompt and the initial image
-        Return value depends on the seed at the time you call it
-        """
-
-        # PLMS sampler not supported yet, so ignore previous sampler
-        if self.sampler_name != 'ddim':
-            print(
-                f">> sampler '{self.sampler_name}' is not yet supported. Using DDIM sampler"
-            )
-            sampler = DDIMSampler(self.model, device=self.device)
-        else:
-            sampler = self.sampler
-
-        sampler.make_schedule(
-            ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False
-        )
-
-        t_enc = int(strength * steps)
-
-        def make_image(x_T):
-            uc, c = self._get_uc_and_c(prompt, skip_normalize)
-
-            # encode (scaled latent)
-            z_enc = sampler.stochastic_encode(
-                init_latent,
-                torch.tensor([t_enc]).to(self.device),
-                noise=x_T
-            )
-            # decode it
-            samples = sampler.decode(
-                z_enc,
-                c,
-                t_enc,
-                img_callback=callback,
-                unconditional_guidance_scale=cfg_scale,
-                unconditional_conditioning=uc,
-            )
-            return self._sample_to_image(samples)
-        return make_image
-
-    # TODO: does this actually need to run every loop? does anything in it vary by random seed?
-    def _get_uc_and_c(self, prompt, skip_normalize):
-
-        uc = self.model.get_learned_conditioning([''])
-
-        # get weighted sub-prompts
-        weighted_subprompts = T2I._split_weighted_subprompts(
-            prompt, skip_normalize)
-
-        if len(weighted_subprompts) > 1:
-            # i dont know if this is correct.. but it works
-            c = torch.zeros_like(uc)
-            # normalize each "sub prompt" and add it
-            for subprompt, weight in weighted_subprompts:
-                self._log_tokenization(subprompt)
-                c = torch.add(
-                    c,
-                    self.model.get_learned_conditioning([subprompt]),
-                    alpha=weight,
-                )
-        else:   # just standard 1 prompt
-            self._log_tokenization(prompt)
-            c = self.model.get_learned_conditioning([prompt])
-        return (uc, c)
-
-    def _sample_to_image(self, samples):
-        x_samples = self.model.decode_first_stage(samples)
-        x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
-        if len(x_samples) != 1:
-            raise Exception(
-                f'>> expected to get a single image, but got {len(x_samples)}')
-        x_sample = 255.0 * rearrange(
-            x_samples[0].cpu().numpy(), 'c h w -> h w c'
-        )
-        return Image.fromarray(x_sample.astype(np.uint8))
-
-    def _new_seed(self):
-        self.seed = random.randrange(0, np.iinfo(np.uint32).max)
-        return self.seed
-
-    def load_model(self):
-        """Load and initialize the model from configuration variables passed at object creation time"""
-        if self.model is None:
-            seed_everything(self.seed)
-            try:
-                config = OmegaConf.load(self.config)
-                model = self._load_model_from_config(config, self.weights)
-                if self.embedding_path is not None:
-                    model.embedding_manager.load(
-                        self.embedding_path, self.full_precision
-                    )
-                self.model = model.to(self.device)
-                # model.to doesn't change the cond_stage_model.device used to move the tokenizer output, so set it here
-                self.model.cond_stage_model.device = self.device
-            except AttributeError as e:
-                print(f'>> Error loading model. {str(e)}', file=sys.stderr)
-                print(traceback.format_exc(), file=sys.stderr)
-                raise SystemExit from e
-
-            self._set_sampler()
-
-        return self.model
-
-    # returns a tensor filled with random numbers from a normal distribution
-    def _get_noise(self,init_img,width,height):
-        if init_img:
-            if self.device.type == 'mps':
-                return torch.randn_like(init_latent, device='cpu').to(self.device)
-            else:
-                return torch.randn_like(init_latent, device=self.device)
-        else:
-            if self.device.type == 'mps':
-                return torch.randn([1,
-                                    self.latent_channels,
-                                    height // self.downsampling_factor,
-                                    width  // self.downsampling_factor],
-                                   device='cpu').to(self.device)
-            else:
-                return torch.randn([1,
-                                    self.latent_channels,
-                                    height // self.downsampling_factor,
-                                    width  // self.downsampling_factor],
-                                   device=self.device)
-
-    def _set_sampler(self):
-        msg = f'>> Setting Sampler to {self.sampler_name}'
-        if self.sampler_name == 'plms':
-            self.sampler = PLMSSampler(self.model, device=self.device)
-        elif self.sampler_name == 'ddim':
-            self.sampler = DDIMSampler(self.model, device=self.device)
-        elif self.sampler_name == 'k_dpm_2_a':
-            self.sampler = KSampler(
-                self.model, 'dpm_2_ancestral', device=self.device
-            )
-        elif self.sampler_name == 'k_dpm_2':
-            self.sampler = KSampler(self.model, 'dpm_2', device=self.device)
-        elif self.sampler_name == 'k_euler_a':
-            self.sampler = KSampler(
-                self.model, 'euler_ancestral', device=self.device
-            )
-        elif self.sampler_name == 'k_euler':
-            self.sampler = KSampler(self.model, 'euler', device=self.device)
-        elif self.sampler_name == 'k_heun':
-            self.sampler = KSampler(self.model, 'heun', device=self.device)
-        elif self.sampler_name == 'k_lms':
-            self.sampler = KSampler(self.model, 'lms', device=self.device)
-        else:
-            msg = f'>> Unsupported Sampler: {self.sampler_name}, Defaulting to plms'
-            self.sampler = PLMSSampler(self.model, device=self.device)
-
-        print(msg)
-
-    def _load_model_from_config(self, config, ckpt):
-        print(f'>> Loading model from {ckpt}')
-        pl_sd = torch.load(ckpt, map_location='cpu')
-        #        if "global_step" in pl_sd:
-        #            print(f"Global Step: {pl_sd['global_step']}")
-        sd = pl_sd['state_dict']
-        model = instantiate_from_config(config.model)
-        m, u = model.load_state_dict(sd, strict=False)
-        model.to(self.device)
-        model.eval()
-        if self.full_precision:
-            print(
-                'Using slower but more accurate full-precision math (--full_precision)'
-            )
-        else:
-            print(
-                '>> Using half precision math. Call with --full_precision to use more accurate but VRAM-intensive full precision.'
-            )
-            model.half()
-        return model
-
-    def _load_img(self, path, width, height, fit=False):
-        with Image.open(path) as img:
-            image = img.convert('RGB')
-        print(
-            f'>> loaded input image of size {image.width}x{image.height} from {path}'
-        )
-
-        # The logic here is:
-        # 1. If "fit" is true, then the image will be fit into the bounding box defined
-        #    by width and height. It will do this in a way that preserves the init image's
-        #    aspect ratio while preventing letterboxing. This means that if there is
-        #    leftover horizontal space after rescaling the image to fit in the bounding box,
-        #    the generated image's width will be reduced to the rescaled init image's width.
-        #    Similarly for the vertical space.
-        # 2. Otherwise, if "fit" is false, then the image will be scaled, preserving its
-        #    aspect ratio, to the nearest multiple of 64. Large images may generate an
-        #    unexpected OOM error.
-        if fit:
-            image = self._fit_image(image,(width,height))
-        else:
-            image = self._squeeze_image(image)
-        image = np.array(image).astype(np.float32) / 255.0
-        image = image[None].transpose(0, 3, 1, 2)
-        image = torch.from_numpy(image)
-        return 2.0 * image - 1.0
-
-    def _squeeze_image(self,image):
-        x,y,resize_needed = self._resolution_check(image.width,image.height)
-        if resize_needed:
-            return InitImageResizer(image).resize(x,y)
-        return image
-
-
-    def _fit_image(self,image,max_dimensions):
-        w,h = max_dimensions
-        print(
-            f'>> image will be resized to fit inside a box {w}x{h} in size.'
-        )
-        if image.width > image.height:
-            h   = None   # by setting h to none, we tell InitImageResizer to fit into the width and calculate height
-        elif image.height > image.width:
-            w   = None   # ditto for w
-        else:
-            pass
-        image = InitImageResizer(image).resize(w,h)   # note that InitImageResizer does the multiple of 64 truncation internally
-        print(
-            f'>> after adjusting image dimensions to be multiples of 64, init image is {image.width}x{image.height}'
-            )
-        return image
-
-
-    # TO DO: Move this and related weighted subprompt code into its own module.
-    def _split_weighted_subprompts(text, skip_normalize=False):
-        """
-        grabs all text up to the first occurrence of ':'
-        uses the grabbed text as a sub-prompt, and takes the value following ':' as weight
-        if ':' has no value defined, defaults to 1.0
-        repeats until no text remaining
-        """
-        prompt_parser = re.compile("""
-            (?P<prompt>     # capture group for 'prompt'
-            (?:\\\:|[^:])+  # match one or more non ':' characters or escaped colons '\:'
-            )               # end 'prompt'
-            (?:             # non-capture group
-            :+              # match one or more ':' characters
-            (?P<weight>     # capture group for 'weight'
-            -?\d+(?:\.\d+)? # match positive or negative integer or decimal number
-            )?              # end weight capture group, make optional
-            \s*             # strip spaces after weight
-            |               # OR
-            $               # else, if no ':' then match end of line
-            )               # end non-capture group
-        """, re.VERBOSE)
-        parsed_prompts = [(match.group("prompt").replace("\\:", ":"), float(
-            match.group("weight") or 1)) for match in re.finditer(prompt_parser, text)]
-        if skip_normalize:
-            return parsed_prompts
-        weight_sum = sum(map(lambda x: x[1], parsed_prompts))
-        if weight_sum == 0:
-            print(
-                "Warning: Subprompt weights add up to zero. Discarding and using even weights instead.")
-            equal_weight = 1 / len(parsed_prompts)
-            return [(x[0], equal_weight) for x in parsed_prompts]
-        return [(x[0], x[1] / weight_sum) for x in parsed_prompts]
-
-    # shows how the prompt is tokenized
-    # usually tokens have '</w>' to indicate end-of-word,
-    # but for readability it has been replaced with ' '
-    def _log_tokenization(self, text):
-        if not self.log_tokenization:
-            return
-        tokens = self.model.cond_stage_model.tokenizer._tokenize(text)
-        tokenized = ""
-        discarded = ""
-        usedTokens = 0
-        totalTokens = len(tokens)
-        for i in range(0, totalTokens):
-            token = tokens[i].replace('</w>', ' ')
-            # alternate color
-            s = (usedTokens % 6) + 1
-            if i < self.model.cond_stage_model.max_length:
-                tokenized = tokenized + f"\x1b[0;3{s};40m{token}"
-                usedTokens += 1
-            else:  # over max token length
-                discarded = discarded + f"\x1b[0;3{s};40m{token}"
-        print(f"\nTokens ({usedTokens}):\n{tokenized}\x1b[0m")
-        if discarded != "":
-            print(
-                f"Tokens Discarded ({totalTokens-usedTokens}):\n{discarded}\x1b[0m")
-
-    def _resolution_check(self, width, height, log=False):
-        resize_needed = False
-        w, h = map(
-            lambda x: x - x % 64, (width, height)
-        )  # resize to integer multiple of 64
-        if h != height or w != width:
-            if log:
-                print(
-                    f'>> Provided width and height must be multiples of 64. Auto-resizing to {w}x{h}'
-                )
-            height = h
-            width  = w
-            resize_needed = True
-
-        if (width * height) > (self.width * self.height):
-            print(">> This input is larger than your defaults. If you run out of memory, please use a smaller image.")
-
-        return width, height, resize_needed
-
-
-    def slerp(self, t, v0, v1, DOT_THRESHOLD=0.9995):
-        '''
-        Spherical linear interpolation
-        Args:
-            t (float/np.ndarray): Float value between 0.0 and 1.0
-            v0 (np.ndarray): Starting vector
-            v1 (np.ndarray): Final vector
-            DOT_THRESHOLD (float): Threshold for considering the two vectors as
-                                colineal. Not recommended to alter this.
-        Returns:
-            v2 (np.ndarray): Interpolation vector between v0 and v1
-        '''
-        inputs_are_torch = False
-        if not isinstance(v0, np.ndarray):
-            inputs_are_torch = True
-            v0 = v0.detach().cpu().numpy()
-        if not isinstance(v1, np.ndarray):
-            inputs_are_torch = True
-            v1 = v1.detach().cpu().numpy()
-
-        dot = np.sum(v0 * v1 / (np.linalg.norm(v0) * np.linalg.norm(v1)))
-        if np.abs(dot) > DOT_THRESHOLD:
-            v2 = (1 - t) * v0 + t * v1
-        else:
-            theta_0 = np.arccos(dot)
-            sin_theta_0 = np.sin(theta_0)
-            theta_t = theta_0 * t
-            sin_theta_t = np.sin(theta_t)
-            s0 = np.sin(theta_0 - theta_t) / sin_theta_0
-            s1 = sin_theta_t / sin_theta_0
-            v2 = s0 * v0 + s1 * v1
-
-        if inputs_are_torch:
-            v2 = torch.from_numpy(v2).to(self.device)
-
-        return v2
+class T2I(Generate):
+    def __init__(self,**kwargs):
+        print(f'>> The ldm.simplet2i module is deprecated. Use ldm.generate instead. It is a drop-in replacement.')
+        super().__init__(kwargs)

requirements-colab.txt

@@ -0,0 +1,26 @@
+albumentations==0.4.3
+clean-fid==0.1.29
+einops==0.3.0
+huggingface-hub==0.8.1
+imageio-ffmpeg==0.4.2
+imageio==2.9.0
+kornia==0.6.0
+numpy==1.21.6
+omegaconf==2.1.1
+opencv-python==4.6.0.66
+pillow==9.2.0
+pip>=22
+pudb==2019.2
+pytorch-lightning==1.4.2
+streamlit==1.12.0
+taming-transformers-rom1504==0.0.6
+test-tube>=0.7.5
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+torchtext==0.6.0
+transformers==4.19.2
+torch==1.12.1+cu113
+torchvision==0.13.1+cu113
+git+https://github.com/openai/CLIP.git@main#egg=clip
+git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
+-e .

requirements-lin.txt

@@ -0,0 +1,32 @@
+albumentations==0.4.3
+einops==0.3.0
+huggingface-hub==0.8.1
+imageio-ffmpeg==0.4.2
+imageio==2.9.0
+kornia==0.6.0
+# pip will resolve the version which matches torch
+numpy
+omegaconf==2.1.1
+opencv-python==4.6.0.66
+pillow==9.2.0
+pip>=22
+pudb==2019.2
+pytorch-lightning==1.4.2
+streamlit==1.12.0
+# "CompVis/taming-transformers" doesn't work
+# ldm\models\autoencoder.py", line 6, in <module>
+#    from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+# ModuleNotFoundError
+taming-transformers-rom1504==0.0.6
+test-tube>=0.7.5
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+transformers==4.19.2
+git+https://github.com/openai/CLIP.git@main#egg=clip
+git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
+# No CUDA in PyPi builds
+--extra-index-url https://download.pytorch.org/whl/cu113 --trusted-host https://download.pytorch.org
+torch==1.11.0
+# Same as numpy - let pip do its thing
+torchvision
+-e .

requirements-win.txt

@@ -0,0 +1,32 @@
+albumentations==0.4.3
+einops==0.3.0
+huggingface-hub==0.8.1
+imageio-ffmpeg==0.4.2
+imageio==2.9.0
+kornia==0.6.0
+# pip will resolve the version which matches torch
+numpy
+omegaconf==2.1.1
+opencv-python==4.6.0.66
+pillow==9.2.0
+pip>=22
+pudb==2019.2
+pytorch-lightning==1.4.2
+streamlit==1.12.0
+# "CompVis/taming-transformers" doesn't work
+# ldm\models\autoencoder.py", line 6, in <module>
+#    from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+# ModuleNotFoundError
+taming-transformers-rom1504==0.0.6
+test-tube>=0.7.5
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+transformers==4.19.2
+git+https://github.com/openai/CLIP.git@main#egg=clip
+git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
+# No CUDA in PyPi builds
+--extra-index-url https://download.pytorch.org/whl/cu113 --trusted-host https://download.pytorch.org
+torch==1.11.0
+# Same as numpy - let pip do its thing
+torchvision
+-e .

scripts/dream.py

@@ -40,7 +40,7 @@ def main():
     print('* Initializing, be patient...\n')
     sys.path.append('.')
     from pytorch_lightning import logging
-    from ldm.simplet2i import T2I
+    from ldm.generate import Generate
 
     # these two lines prevent a horrible warning message from appearing
     # when the frozen CLIP tokenizer is imported
@@ -52,18 +52,18 @@ def main():
     # defaults passed on the command line.
     # additional parameters will be added (or overriden) during
     # the user input loop
-    t2i = T2I(
-        width=width,
-        height=height,
-        sampler_name=opt.sampler_name,
-        weights=weights,
-        full_precision=opt.full_precision,
-        config=config,
-        grid  = opt.grid,
+    t2i = Generate(
+        width          = width,
+        height         = height,
+        sampler_name   = opt.sampler_name,
+        weights        = weights,
+        full_precision = opt.full_precision,
+        config         = config,
+        grid           = opt.grid,
         # this is solely for recreating the prompt
-        latent_diffusion_weights=opt.laion400m,
-        embedding_path=opt.embedding_path,
-        device_type=opt.device
+        seamless       = opt.seamless,
+        embedding_path = opt.embedding_path,
+        device_type    = opt.device
     )
 
     # make sure the output directory exists
@@ -87,12 +87,11 @@ def main():
             print(f'{e}. Aborting.')
             sys.exit(-1)
 
+    if opt.seamless:
+        print(">> changed to seamless tiling mode")
+
     # preload the model
-    tic = time.time()
     t2i.load_model()
-    print(
-        f'>> model loaded in', '%4.2fs' % (time.time() - tic)
-    )
 
     if not infile:
         print(
@@ -101,7 +100,7 @@ def main():
 
     cmd_parser = create_cmd_parser()
     if opt.web:
-        dream_server_loop(t2i, opt.host, opt.port)
+        dream_server_loop(t2i, opt.host, opt.port, opt.outdir)
     else:
         main_loop(t2i, opt.outdir, opt.prompt_as_dir, cmd_parser, infile)
 
@@ -312,7 +311,7 @@ def get_next_command(infile=None) -> str: #command string
         print(f'#{command}')
     return command
 
-def dream_server_loop(t2i, host, port):
+def dream_server_loop(t2i, host, port, outdir):
     print('\n* --web was specified, starting web server...')
     # Change working directory to the stable-diffusion directory
     os.chdir(
@@ -321,6 +320,7 @@ def dream_server_loop(t2i, host, port):
 
     # Start server
     DreamServer.model = t2i
+    DreamServer.outdir = outdir
     dream_server = ThreadingDreamServer((host, port))
     print(">> Started Stable Diffusion dream server!")
     if host == '0.0.0.0':
@@ -418,6 +418,11 @@ def create_argv_parser():
         default='outputs/img-samples',
         help='Directory to save generated images and a log of prompts and seeds. Default: outputs/img-samples',
     )
+    parser.add_argument(
+        '--seamless',
+        action='store_true',
+        help='Change the model to seamless tiling (circular) mode',
+    )
     parser.add_argument(
         '--embedding_path',
         type=str,
@@ -434,7 +439,7 @@ def create_argv_parser():
         '--gfpgan_bg_upsampler',
         type=str,
         default='realesrgan',
-        help='Background upsampler. Default: realesrgan. Options: realesrgan, none. Only used if --gfpgan is specified',
+        help='Background upsampler. Default: realesrgan. Options: realesrgan, none.',
 
     )
     parser.add_argument(
@@ -540,6 +545,11 @@ def create_cmd_parser():
         default=None,
         help='Directory to save generated images and a log of prompts and seeds',
     )
+    parser.add_argument(
+        '--seamless',
+        action='store_true',
+        help='Change the model to seamless tiling (circular) mode',
+    )
     parser.add_argument(
         '-i',
         '--individual',
@@ -552,6 +562,12 @@ def create_cmd_parser():
         type=str,
         help='Path to input image for img2img mode (supersedes width and height)',
     )
+    parser.add_argument(
+        '-M',
+        '--init_mask',
+        type=str,
+        help='Path to input mask for inpainting mode (supersedes width and height)',
+    )
     parser.add_argument(
         '-T',
         '-fit',

scripts/orig_scripts/txt2img.py

@@ -232,7 +232,12 @@ def main():
         print(f"reading prompts from {opt.from_file}")
         with open(opt.from_file, "r") as f:
             data = f.read().splitlines()
-            data = list(chunk(data, batch_size))
+            if (len(data) >= batch_size):
+                data = list(chunk(data, batch_size))
+            else:
+                while (len(data) < batch_size):
+                    data.append(data[-1])
+                data = [data]
 
     sample_path = os.path.join(outpath, "samples")
     os.makedirs(sample_path, exist_ok=True)
@@ -264,7 +269,7 @@ def main():
                             prompts = list(prompts)
                         c = model.get_learned_conditioning(prompts)
                         shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
-                        
+
                         if not opt.klms:
                             samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
                                                             conditioning=c,
@@ -284,7 +289,7 @@ def main():
                             model_wrap_cfg = CFGDenoiser(model_wrap)
                             extra_args = {'cond': c, 'uncond': uc, 'cond_scale': opt.scale}
                             samples_ddim = K.sampling.sample_lms(model_wrap_cfg, x, sigmas, extra_args=extra_args)
-                        
+
                         x_samples_ddim = model.decode_first_stage(samples_ddim)
                         x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
 

static/dream_web/index.css

@@ -1,11 +1,14 @@
 * {
     font-family: 'Arial';
+    font-size: 100%;
 }
-#header {
-    text-decoration: dotted underline;
+body {
+    font-size: 1em;
 }
-#search {
-    margin-top: 20vh;
+textarea {
+    font-size: 0.95em;
+}
+header, form, #progress-section {
     margin-left: auto;
     margin-right: auto;
     max-width: 1024px;
@@ -13,47 +16,79 @@
 }
 fieldset {
     border: none;
+    line-height: 2.2em;
+}
+select, input {
+    margin-right: 10px;
+    padding: 2px;
+}
+input[type=submit] {
+    background-color: #666;
+    color: white;
+}
+input[type=checkbox] {
+    margin-right: 0px;
+    width: 20px;
+    height: 20px;
+    vertical-align: middle;
+}
+input#seed {
+    margin-right: 0px;
 }
 div {
     padding: 10px 10px 10px 10px;
 }
-#fieldset-search {
+header {
+    margin-bottom: 16px;
+}
+header h1 {
+    margin-bottom: 0;
+    font-size: 2em;
+}
+#search-box {
     display: flex;
 }
-#scaling-inprocess-message{
+#scaling-inprocess-message {
     font-weight: bold;
     font-style: italic;
     display: none;
 }
 #prompt {
     flex-grow: 1;
-
-    border-radius: 20px 0px 0px 20px;
     padding: 5px 10px 5px 10px;
-    border: 1px solid black;
-    border-right: none;
+    border: 1px solid #999;
     outline: none;
 }
 #submit {
-    border-radius: 0px 20px 20px 0px;
     padding: 5px 10px 5px 10px;
-    border: 1px solid black;
+    border: 1px solid #999;
 }
-#reset-all {
+#reset-all, #remove-image {
+    margin-top: 12px;
+    font-size: 0.8em;
     background-color: pink;
+    border: 1px solid #999;
+    border-radius: 4px;
 }
 #results {
     text-align: center;
-//    max-width: 1024px;
     margin: auto;
     padding-top: 10px;
 }
-#results img {
-    cursor: pointer;
-    height: 30vh;
-    border-radius: 5px;
+#results figure {
+    display: inline-block;
     margin: 10px;
 }
+#results figcaption {
+    font-size: 0.8em;
+    padding: 3px;
+    color: #888;
+    cursor: pointer;
+}
+#results img {
+    border-radius: 5px;
+    object-fit: cover;
+}
 #fieldset-config {
     line-height:2em;
 }
@@ -63,8 +98,15 @@ input[type="number"] {
 #seed {
     width: 150px;
 }
-hr {
-//    width: 200px;
+button#reset-seed {
+    font-size: 1.7em;
+    background: #efefef;
+    border: 1px solid #999;
+    border-radius: 4px;
+    line-height: 0.8;
+    margin: 0 10px 0 0;
+    padding: 0 5px 3px;
+    vertical-align: middle;
 }
 label {
     white-space: nowrap;
@@ -83,6 +125,9 @@ label {
 #txt2img {
     background-color: #DCDCDC;
 }
+#variations {
+    background-color: #EEEEEE;
+}
 #img2img {
     background-color: #F5F5F5;
 }
@@ -92,6 +137,8 @@ label {
 #progress-section {
     background-color: #F5F5F5;
 }
-#about {
-    background-color: #DCDCDC;
+
+#no-results-message:not(:only-child) {
+    display: none;
 }
+

static/dream_web/index.html

@@ -2,110 +2,126 @@
   <head>
     <title>Stable Diffusion Dream Server</title>
     <meta charset="utf-8">
-    <link rel="icon" href="data:,">
+    <link rel="icon" type="image/x-icon" href="static/dream_web/favicon.ico" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
-
     <link rel="stylesheet" href="static/dream_web/index.css">
     <script src="config.js"></script>
     <script src="static/dream_web/index.js"></script>
   </head>
   <body>
-    <div id="search">
-      <h2 id="header">Stable Diffusion Dream Server</h2>
+    <header>
+      <h1>Stable Diffusion Dream Server</h1>
+      <div id="about">
+        For news and support for this web service, visit our <a href="http://github.com/lstein/stable-diffusion">GitHub site</a>
+      </div>
+    </header>
 
+    <main>
       <form id="generate-form" method="post" action="#">
-	<div id="txt2img">
-          <fieldset id="fieldset-search">
-            <input type="text" id="prompt" name="prompt">
+	<fieldset id="txt2img">
+          <div id="search-box">
+            <textarea rows="3" id="prompt" name="prompt"></textarea>
             <input type="submit" id="submit" value="Generate">
-          </fieldset>
-          <fieldset id="fieldset-config">
-            <label for="iterations">Images to generate:</label>
-            <input value="1" type="number" id="iterations" name="iterations" size="4">
-            <label for="steps">Steps:</label>
-            <input value="50" type="number" id="steps" name="steps">
-            <label for="cfgscale">Cfg Scale:</label>
-            <input value="7.5" type="number" id="cfgscale" name="cfgscale" step="any">
-            <label for="sampler">Sampler:</label>
-            <select id="sampler" name="sampler" value="k_lms">
-              <option value="ddim">DDIM</option>
-              <option value="plms">PLMS</option>
-              <option value="k_lms" selected>KLMS</option>
-              <option value="k_dpm_2">KDPM_2</option>
-              <option value="k_dpm_2_a">KDPM_2A</option>
-              <option value="k_euler">KEULER</option>
-	      <option value="k_euler_a">KEULER_A</option>
-              <option value="k_heun">KHEUN</option>
-            </select>
-            <br>
-            <label title="Set to multiple of 64" for="width">Width:</label>
-            <select id="width" name="width" value="512">
-              <option value="64">64</option> <option value="128">128</option>
-              <option value="192">192</option> <option value="256">256</option>
-              <option value="320">320</option> <option value="384">384</option>
-              <option value="448">448</option> <option value="512" selected>512</option>
-              <option value="576">576</option> <option value="640">640</option>
-              <option value="704">704</option> <option value="768">768</option>
-              <option value="832">832</option> <option value="896">896</option>
-              <option value="960">960</option> <option value="1024">1024</option>
-            </select>
-            <label title="Set to multiple of 64" for="height">Height:</label>
-            <select id="height" name="height" value="512">
-              <option value="64">64</option> <option value="128">128</option>
-              <option value="192">192</option> <option value="256">256</option>
-              <option value="320">320</option> <option value="384">384</option>
-              <option value="448">448</option> <option value="512" selected>512</option>
-              <option value="576">576</option> <option value="640">640</option>
-              <option value="704">704</option> <option value="768">768</option>
-              <option value="832">832</option> <option value="896">896</option>
-              <option value="960">960</option> <option value="1024">1024</option>
-            </select>
-            <label title="Set to -1 for random seed" for="seed">Seed:</label>
-            <input value="-1" type="number" id="seed" name="seed">
-            <button type="button" id="reset-seed">&olarr;</button>
-            <input type="checkbox" name="progress_images" id="progress_images">
-	    <label for="progress_images">Display in-progress images (slows down generation):</label>
-	    <button type="button" id="reset-all">Reset to Defaults</button>
-	</div>
-	<div id="img2img">
-          <label title="Upload an image to use img2img" for="initimg">Initial image:</label>
-          <input type="file" id="initimg" name="initimg" accept=".jpg, .jpeg, .png">
-	  <br>
-          <label for="strength">Img2Img Strength:</label>
-          <input value="0.75" type="number" id="strength" name="strength" step="0.01" min="0" max="1">
-          <input type="checkbox" id="fit" name="fit" checked>
-          <label title="Rescale image to fit within requested width and height" for="fit">Fit to width/height:</label>
-	</div>
-        <div id="gfpgan">
+	  </div>
+        </fieldset>
+        <fieldset id="fieldset-config">
+          <label for="iterations">Images to generate:</label>
+          <input value="1" type="number" id="iterations" name="iterations" size="4">
+          <label for="steps">Steps:</label>
+          <input value="50" type="number" id="steps" name="steps">
+          <label for="cfg_scale">Cfg Scale:</label>
+          <input value="7.5" type="number" id="cfg_scale" name="cfg_scale" step="any">
+          <label for="sampler_name">Sampler:</label>
+          <select id="sampler_name" name="sampler_name" value="k_lms">
+            <option value="ddim">DDIM</option>
+            <option value="plms">PLMS</option>
+            <option value="k_lms" selected>KLMS</option>
+            <option value="k_dpm_2">KDPM_2</option>
+            <option value="k_dpm_2_a">KDPM_2A</option>
+            <option value="k_euler">KEULER</option>
+	          <option value="k_euler_a">KEULER_A</option>
+            <option value="k_heun">KHEUN</option>
+          </select>
+          <input type="checkbox" name="seamless" id="seamless">
+	        <label for="seamless">Seamless circular tiling</label>
+          <br>
+          <label title="Set to multiple of 64" for="width">Width:</label>
+          <select id="width" name="width" value="512">
+            <option value="64">64</option> <option value="128">128</option>
+            <option value="192">192</option> <option value="256">256</option>
+            <option value="320">320</option> <option value="384">384</option>
+            <option value="448">448</option> <option value="512" selected>512</option>
+            <option value="576">576</option> <option value="640">640</option>
+            <option value="704">704</option> <option value="768">768</option>
+            <option value="832">832</option> <option value="896">896</option>
+            <option value="960">960</option> <option value="1024">1024</option>
+          </select>
+          <label title="Set to multiple of 64" for="height">Height:</label>
+          <select id="height" name="height" value="512">
+            <option value="64">64</option> <option value="128">128</option>
+            <option value="192">192</option> <option value="256">256</option>
+            <option value="320">320</option> <option value="384">384</option>
+            <option value="448">448</option> <option value="512" selected>512</option>
+            <option value="576">576</option> <option value="640">640</option>
+            <option value="704">704</option> <option value="768">768</option>
+            <option value="832">832</option> <option value="896">896</option>
+            <option value="960">960</option> <option value="1024">1024</option>
+          </select>
+          <label title="Set to -1 for random seed" for="seed">Seed:</label>
+          <input value="-1" type="number" id="seed" name="seed">
+          <button type="button" id="reset-seed">&olarr;</button>
+          <input type="checkbox" name="progress_images" id="progress_images">
+	  <label for="progress_images">Display in-progress images (slower)</label>
+	  <button type="button" id="reset-all">Reset to Defaults</button>
+	  <div id="variations">
+            <label title="If > 0, generates variations on the initial seed instead of random seeds per iteration. Must be between 0 and 1. Higher values will be more different." for="variation_amount">Variation amount (0 to disable):</label>
+            <input value="0" type="number" id="variation_amount" name="variation_amount" step="0.01" min="0" max="1">
+            <label title="list of variations to apply, in the format `seed:weight,seed:weight,..." for="with_variations">With variations (seed:weight,seed:weight,...):</label>
+            <input value="" type="text" id="with_variations" name="with_variations">
+	  </div>
+	  <div id="img2img">
+            <label title="Upload an image to use img2img" for="initimg">Initial image:</label>
+            <input type="file" id="initimg" name="initimg" accept=".jpg, .jpeg, .png">
+            <button type="button" id="remove-image">Remove Image</button>
+	    <br>
+	    <label for="strength">Img2Img Strength:</label>
+	    <input value="0.75" type="number" id="strength" name="strength" step="0.01" min="0" max="1">
+	    <input type="checkbox" id="fit" name="fit" checked>
+	    <label title="Rescale image to fit within requested width and height" for="fit">Fit to width/height:</label>
+	  </div>
+	</fieldset>
+        <fieldset id="gfpgan">
           <label title="Strength of the gfpgan (face fixing) algorithm." for="gfpgan_strength">GPFGAN Strength (0 to disable):</label>
           <input value="0.8" min="0" max="1" type="number" id="gfpgan_strength" name="gfpgan_strength" step="0.05">
           <label title="Upscaling to perform using ESRGAN." for="upscale_level">Upscaling Level</label>
           <select id="upscale_level" name="upscale_level" value="">
             <option value="" selected>None</option>
             <option value="2">2x</option>
+            <option value="3">3x</option>
             <option value="4">4x</option>
           </select>
           <label title="Strength of the esrgan (upscaling) algorithm." for="upscale_strength">Upscale Strength:</label>
           <input value="0.75" min="0" max="1" type="number" id="upscale_strength" name="upscale_strength" step="0.05">
-        </div>
         </fieldset>
       </form>
-      <div id="about">For news and support for this web service, visit our <a href="http://github.com/lstein/stable-diffusion">GitHub site</a></div>
       <br>
-      <div id="progress-section">
-        <progress id="progress-bar" value="0" max="1"></progress>
-        <span id="cancel-button" title="Cancel">&#10006;</span>
-        <br>
-        <img id="progress-image" src='data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>'></img>
-        <div id="scaling-inprocess-message">
-          <i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i>
+      <section id="progress-section">
+        <div id="progress-container">
+          <progress id="progress-bar" value="0" max="1"></progress>
+          <span id="cancel-button" title="Cancel">&#10006;</span>
+          <br>
+          <img id="progress-image" src='data:image/svg+xml,<svg xmlns="http://www.w3.org/2000/svg"/>'>
+          <div id="scaling-inprocess-message">
+            <i><span>Postprocessing...</span><span id="processing_cnt">1/3</span></i>
+          </div>
+        </div>
+      </section>
+
+      <div id="results">
+        <div id="no-results-message">
+          <i><p>No results...</p></i>
         </div>
       </div>
-    </div>
-    <div id="results">
-      <div id="no-results-message">
-        <i><p>No results...</p></i>
-      </div>
-    </div>
+    </main>
   </body>
 </html>

static/dream_web/index.js

@@ -8,20 +8,44 @@ function toBase64(file) {
 }
 
 function appendOutput(src, seed, config) {
-    let outputNode = document.createElement("img");
-    outputNode.src = src;
+    let outputNode = document.createElement("figure");
+    
+    let variations = config.with_variations;
+    if (config.variation_amount > 0) {
+        variations = (variations ? variations + ',' : '') + seed + ':' + config.variation_amount;
+    }
+    let baseseed = (config.with_variations || config.variation_amount > 0) ? config.seed : seed;
+    let altText = baseseed + ' | ' + (variations ? variations + ' | ' : '') + config.prompt;
 
-    let altText = seed.toString() + " | " + config.prompt;
-    outputNode.alt = altText;
-    outputNode.title = altText;
+    // img needs width and height for lazy loading to work
+    const figureContents = `
+        <a href="${src}" target="_blank">
+            <img src="${src}"
+                 alt="${altText}"
+                 title="${altText}"
+                 loading="lazy"
+                 width="256"
+                 height="256">
+        </a>
+        <figcaption>${seed}</figcaption>
+    `;
+
+    outputNode.innerHTML = figureContents;
+    let figcaption = outputNode.querySelector('figcaption');
 
     // Reload image config
-    outputNode.addEventListener('click', () => {
+    figcaption.addEventListener('click', () => {
         let form = document.querySelector("#generate-form");
         for (const [k, v] of new FormData(form)) {
+            if (k == 'initimg') { continue; }
             form.querySelector(`*[name=${k}]`).value = config[k];
         }
-        document.querySelector("#seed").value = seed;
+
+        document.querySelector("#seed").value = baseseed;
+        document.querySelector("#with_variations").value = variations || '';
+        if (document.querySelector("#variation_amount").value <= 0) {
+            document.querySelector("#variation_amount").value = 0.2;
+        }
 
         saveFields(document.querySelector("#generate-form"));
     });
@@ -59,6 +83,7 @@ async function generateSubmit(form) {
 
     // Convert file data to base64
     let formData = Object.fromEntries(new FormData(form));
+    formData.initimg_name = formData.initimg.name
     formData.initimg = formData.initimg.name !== '' ? await toBase64(formData.initimg) : null;
 
     let strength = formData.strength;
@@ -94,7 +119,6 @@ async function generateSubmit(form) {
 
                 if (data.event === 'result') {
                     noOutputs = false;
-                    document.querySelector("#no-results-message")?.remove();
                     appendOutput(data.url, data.seed, data.config);
                     progressEle.setAttribute('value', 0);
                     progressEle.setAttribute('max', totalSteps);
@@ -130,7 +154,25 @@ async function generateSubmit(form) {
     document.querySelector("#prompt").value = `Generating: "${prompt}"`;
 }
 
-window.onload = () => {
+async function fetchRunLog() {
+    try {
+        let response = await fetch('/run_log.json')
+        const data = await response.json();
+        for(let item of data.run_log) {
+            appendOutput(item.url, item.seed, item);
+        }
+    } catch (e) {
+        console.error(e);
+    }
+}
+
+window.onload = async () => {
+    document.querySelector("#prompt").addEventListener("keydown", (e) => {
+      if (e.key === "Enter" && !e.shiftKey) {
+        const form = e.target.form;
+        generateSubmit(form);
+      }
+    });
     document.querySelector("#generate-form").addEventListener('submit', (e) => {
         e.preventDefault();
         const form = e.target;
@@ -147,6 +189,9 @@ window.onload = () => {
     document.querySelector("#reset-all").addEventListener('click', (e) => {
         clearFields(e.target.form);
     });
+    document.querySelector("#remove-image").addEventListener('click', (e) => {
+        initimg.value=null;
+    });
     loadFields(document.querySelector("#generate-form"));
 
     document.querySelector('#cancel-button').addEventListener('click', () => {
@@ -154,8 +199,15 @@ window.onload = () => {
             console.error(e);
         });
     });
+    document.documentElement.addEventListener('keydown', (e) => {
+      if (e.key === "Escape")
+        fetch('/cancel').catch(err => {
+          console.error(err);
+        });
+    });
 
     if (!config.gfpgan_model_exists) {
         document.querySelector("#gfpgan").style.display = 'none';
     }
+    await fetchRunLog()
 };

tests/prompts.txt

@@ -0,0 +1 @@
+test trending on artstation -s 1 -S 1