Merge branch 'hwharrison-main' into main

This enables k_lms sampling (now the default)`:wq

README.md

@@ -1,4 +1,4 @@
-# Stable Diffusion
+# Stable Diffusion Dream Script
 
 This is a fork of CompVis/stable-diffusion, the wonderful open source
 text-to-image generator.
@@ -7,20 +7,24 @@ The original has been modified in several ways:
 
 ## Interactive command-line interface similar to the Discord bot
 
-There is now a command-line script, located in scripts/dream.py, which
+The *dream.py* script, located in scripts/dream.py, 
 provides an interactive interface to image generation similar to
 the "dream mothership" bot that Stable AI provided on its Discord
-server.  The advantage of this is that the lengthy model
-initialization only happens once. After that image generation is
-fast.
+server. Unlike the txt2img.py and img2img.py scripts provided in the
+original CompViz/stable-diffusion source code repository, the
+time-consuming initialization of the AI model
+initialization only happens once. After that image generation 
+from the command-line interface is very fast.
 
 The script uses the readline library to allow for in-line editing,
-command history (up and down arrows) and more.
+command history (up and down arrows), autocompletion, and more.
 
-Note that this has only been tested in the Linux environment!
+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.
 
 ~~~~
-(ldm) ~/stable-diffusion$ ./scripts/dream.py
+(ldm) ~/stable-diffusion$ python3 ./scripts/dream.py
 * Initializing, be patient...
 Loading model from models/ldm/text2img-large/model.ckpt
 LatentDiffusion: Running in eps-prediction mode
@@ -32,31 +36,192 @@ Loading Bert tokenizer from "models/bert"
 setting sampler to plms
 
 * Initialization done! Awaiting your command...
-dream> ashley judd riding a camel -n2
+dream> ashley judd riding a camel -n2 -s150
 Outputs:
-   outputs/txt2img-samples/00009.png: "ashley judd riding a camel" -n2 -S 416354203
-   outputs/txt2img-samples/00010.png: "ashley judd riding a camel" -n2 -S 1362479620
+   outputs/txt2img-samples/00009.png: "ashley judd riding a camel" -n2 -s150 -S 416354203
+   outputs/txt2img-samples/00010.png: "ashley judd riding a camel" -n2 -s150-S 1362479620
 
-dream> "your prompt here" -n6 -g
-    outputs/txt2img-samples/00041.png: "your prompt here" -n6 -g -S 2685670268
+dream> "there's a fly in my soup" -n6 -g
+    outputs/txt2img-samples/00041.png: "there's a fly in my soup" -n6 -g -S 2685670268
     seeds for individual rows: [2685670268, 1216708065, 2335773498, 822223658, 714542046, 3395302430]
 ~~~~
 
-Command-line arguments passed to the script allow you to change
-various defaults, and select between the mature stable-diffusion
-weights (512x512) and the older (256x256) latent diffusion weights
-(laion400m). From the dream> prompt, the arguments are (mostly)
+The dream> prompt's arguments are pretty much
 identical to those used in the Discord bot, except you don't need to
-type "!dream". Pass "-h" (or "--help") to list the arguments.
+type "!dream" (it doesn't hurt if you do). A significant change is that creation of individual images 
+is now the default
+unless --grid (-g) is given. For backward compatibility, the -i switch is recognized.
+For command-line help type -h (or --help) at the dream> prompt.
+
+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.
+
+## Image-to-Image
 
 This script also provides an img2img feature that lets you seed your
-creations with a drawing or photo.
+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:
 
-For command-line help, type -h (or --help) at the dream> prompt.
+~~~~
+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.
+
+## Changes
+
+- v1.01 (21 August 2022)
+* added k_lms sampling **Please run "conda update -f environment.yaml" to load the k_lms dependencies**
+* use half precision arithmetic by default, resulting in faster execution and lower memory requirements
+Pass argument --full_precision to dream.py to get slower but more accurate image generation
+
+
+## Installation
+
+### Linux/Mac
+
+1. You will need to install the following prerequisites if they are not already available. Use your
+operating system's preferred installer
+* Python (version 3.8 or higher)
+* git
+
+2. Install the Python Anaconda environment manager using pip3.
+```
+~$ pip3 install anaconda
+```
+After installing anaconda, you should log out of your system and log back in. If the installation
+worked, your command prompt will be prefixed by the name of the current anaconda environment, "(base)".
+
+3. Copy the stable-diffusion source code from GitHub:
+```
+(base) ~$ git clone https://github.com/lstein/stable-diffusion.git
+```
+This will create stable-diffusion folder where you will follow the rest of the steps.
+
+4. Enter the newly-created stable-diffusion folder. From this step forward make sure that you are working in the stable-diffusion directory!
+```
+(base) ~$ cd stable-diffusion
+(base) ~/stable-diffusion$
+```
+5. Use anaconda to copy necessary python packages, create a new python environment named "ldm",
+and activate the environment. 
+```
+(base) ~/stable-diffusion$ conda env create -f environment.yaml
+(base) ~/stable-diffusion$ conda activate ldm
+(ldm) ~/stable-diffusion$
+```
+After these steps, your command prompt will be prefixed by "(ldm)" as shown above.
+
+6. Load a couple of small machine-learning models required by stable diffusion:
+```
+(ldm) ~/stable-diffusion$ python3 scripts/preload_models.py
+```
+
+7. Now you need to install the weights for the stable diffusion model.
+
+For testing prior to the release of the real weights, you can use an older weight file that produces low-quality images. Create a directory within stable-diffusion named "models/ldm/text2img.large", and use the wget URL downloader tool to copy the weight file into it:
+```
+(ldm) ~/stable-diffusion$ mkdir -p models/ldm/text2img-large
+(ldm) ~/stable-diffusion$ wget -O models/ldm/text2img-large/model.ckpt https://ommer-lab.com/files/latent-diffusion/nitro/txt2img-f8-large/model.ckpt
+```
+For testing with the released weighs, you will do something similar, but with a directory named "models/ldm/stable-diffusion-v1"
+```
+(ldm) ~/stable-diffusion$ mkdir -p models/ldm/stable-diffusion-v1
+(ldm) ~/stable-diffusion$ wget -O models/ldm/stable-diffusion-v1/model.ckpt <ENTER URL HERE>
+```
+These weight files are ~5 GB in size, so downloading may take a while.
+
+8. Start generating images!
+```
+# for the pre-release weights use the -l or --liaon400m switch
+(ldm) ~/stable-diffusion$ python3 scripts/dream.py -l
+
+# for the post-release weights do not use the switch
+(ldm) ~/stable-diffusion$ python3 scripts/dream.py
+
+# for additional configuration switches and arguments, use -h or --help
+(ldm) ~/stable-diffusion$ python3 scripts/dream.py -h
+```
+9. Subsequently, to relaunch the script, be sure to run "conda activate ldm" (step 5, second command), enter the "stable-diffusion" 
+directory, and then launch the dream script (step 8). If you forget to activate the ldm environment, the script will fail with multiple ModuleNotFound errors.
+
+### Updating to newer versions of the script
+
+This distribution is changing rapidly. If you used the "git clone" method (step 5) to download the stable-diffusion directory, then to update to the latest and greatest version, launch the Anaconda window, enter "stable-diffusion", and type:
+```
+(ldm) ~/stable-diffusion$ git pull
+```
+This will bring your local copy into sync with the remote one.
+
+### Windows
+
+1. Install the most recent Python from here: https://www.python.org/downloads/windows/
+
+2. Install Anaconda3 (miniconda3 version) from here: https://docs.anaconda.com/anaconda/install/windows/
+
+3. Install Git from here: https://git-scm.com/download/win
+
+4. Launch Anaconda from the Windows Start menu. This will bring up a command window. Type all the remaining commands in this window.
+
+5. Run the command:
+```
+git clone https://github.com/lstein/stable-diffusion.git
+```
+This will create stable-diffusion folder where you will follow the rest of the steps.
+
+6. Enter the newly-created stable-diffusion folder. From this step forward make sure that you are working in the stable-diffusion directory!
+```
+cd stable-diffusion
+```
+
+7. Run the following two commands:
+```
+conda env create -f environment.yaml    (step 7a)
+conda activate ldm                      (step 7b)
+```
+This will install all python requirements and activate the "ldm" environment which sets PATH and other environment variables properly.
+
+8. Run the command:
+```
+python scripts\preload_models.py
+```
+This installs two machine learning models that stable diffusion requires.
+
+9. Now you need to install the weights for the big stable diffusion model.
+
+For testing prior to the release of the real weights, create a directory within stable-diffusion named "models\ldm\text2img.large".
+
+For testing with the released weights, create a directory within stable-diffusion named "models\ldm\stable-diffusion-v1".
+
+Then use a web browser to copy model.ckpt into the appropriate directory. For the text2img.large (pre-release) model, the weights are at https://ommer-lab.com/files/latent-diffusion/nitro/txt2img-f8-large/model.ckpt. Check back here later for the release URL.
+
+10. Start generating images!
+```
+# for the pre-release weights
+python scripts\dream.py -l
+
+# for the post-release weights
+python scripts\dream.py
+```
+11. Subsequently, to relaunch the script, first activate the Anaconda command window (step 4), run "conda activate ldm" (step 7b), and then launch the dream script (step 10).
+
+### Updating to newer versions of the script
+
+This distribution is changing rapidly. If you used the "git clone" method (step 5) to download the stable-diffusion directory, then to update to the latest and greatest version, launch the Anaconda window, enter "stable-diffusion", and type:
+```
+git pull
+```
+This will bring your local copy into sync with the remote one.
 
 ## Simplified API for text to image generation
 
-There is now a simplified API for text to image generation, which
+For programmers who wish to incorporate stable-diffusion into other
+products, this repository includes a simplified API for text to image generation, which
 lets you create images from a prompt in just three lines of code:
 
 ~~~~
@@ -68,6 +233,7 @@ outputs = model.text2image("a unicorn in manhattan")
 Outputs is a list of lists in the format [[filename1,seed1],[filename2,seed2]...]
 Please see ldm/simplet2i.py for more information.
 
+
 ## Workaround for machines with limited internet connectivity
 
 My development machine is a GPU node in a high-performance compute
@@ -103,14 +269,9 @@ time, copy over the file ldm/modules/encoders/modules.py from the
 CompVis/stable-diffusion repository. Or you can run preload_models.py
 on the target machine.
 
-## Minor fixes
+## Support
 
-I added the requirement for torchmetrics to environment.yaml.
-
-## Installation and support
-
-Follow the directions from the original README, which starts below, to
-configure the environment and install requirements. For support,
+For support,
 please use this repository's GitHub Issues tracking service. Feel free
 to send me an email if you use and like the script.
 
@@ -141,6 +302,7 @@ See [this section](#stable-diffusion-v1) below and the [model card](https://hugg
 
   
 ## Requirements
+
 A suitable [conda](https://conda.io/) environment named `ldm` can be created
 and activated with:
 
@@ -155,8 +317,7 @@ You can also update an existing [latent diffusion](https://github.com/CompVis/la
 conda install pytorch torchvision -c pytorch
 pip install transformers==4.19.2
 pip install -e .
-``` 
-
+```
 
 ## Stable Diffusion v1
 

environment.yaml

@@ -24,6 +24,8 @@ dependencies:
     - transformers==4.19.2
     - torchmetrics==0.6.0
     - kornia==0.6
-    - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+    - accelerate==0.12.0
     - -e git+https://github.com/openai/CLIP.git@main#egg=clip
+    - -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+    - -e git+https://github.com/lstein/k-diffusion.git@master#egg=k-diffusion
     - -e .

ldm/models/diffusion/ksampler.py

@@ -0,0 +1,74 @@
+'''wrapper around part of Karen Crownson's k-duffsion library, making it call compatible with other Samplers'''
+import k_diffusion as K
+import torch
+import torch.nn as nn
+import accelerate
+
+class CFGDenoiser(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.inner_model = model
+
+    def forward(self, x, sigma, uncond, cond, cond_scale):
+        x_in = torch.cat([x] * 2)
+        sigma_in = torch.cat([sigma] * 2)
+        cond_in = torch.cat([uncond, cond])
+        uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)
+        return uncond + (cond - uncond) * cond_scale
+
+class KSampler(object):
+    def __init__(self,model,schedule="lms", **kwargs):
+        super().__init__()
+        self.model        = K.external.CompVisDenoiser(model)
+        self.accelerator  = accelerate.Accelerator()
+        self.device       = self.accelerator.device
+        self.schedule = schedule
+
+        def forward(self, x, sigma, uncond, cond, cond_scale):
+            x_in = torch.cat([x] * 2)
+            sigma_in = torch.cat([sigma] * 2)
+            cond_in = torch.cat([uncond, cond])
+            uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)
+            return uncond + (cond - uncond) * cond_scale
+
+
+    # most of these arguments are ignored and are only present for compatibility with
+    # other samples
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               normals_sequence=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+               **kwargs
+               ):
+
+        sigmas = self.model.get_sigmas(S)
+        if x_T:
+            x = x_T
+        else:
+            x = torch.randn([batch_size, *shape], device=self.device) * sigmas[0] # for GPU draw
+        model_wrap_cfg = CFGDenoiser(self.model)
+        extra_args = {'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': unconditional_guidance_scale}
+        return (K.sampling.sample_lms(model_wrap_cfg, x, sigmas, extra_args=extra_args, disable=not self.accelerator.is_main_process),
+                None)
+
+    def gather(samples_ddim):
+        return self.accelerator.gather(samples_ddim)

ldm/simplet2i.py

@@ -11,7 +11,7 @@ t2i = T2I(outdir      = <path>        // outputs/txt2img-samples
           batch_size       = <integer>     // how many images to generate per sampling (1)
           steps       = <integer>     // 50
           seed        = <integer>     // current system time
-          sampler     = ['ddim','plms']  // ddim
+          sampler     = ['ddim','plms','klms']  // klms
           grid        = <boolean>     // false
           width       = <integer>     // image width, multiple of 64 (512)
           height      = <integer>     // image height, multiple of 64 (512)
@@ -62,8 +62,9 @@ import time
 import math
 
 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.ddim     import DDIMSampler
+from ldm.models.diffusion.plms     import PLMSSampler
+from ldm.models.diffusion.ksampler import KSampler
 
 class T2I:
     """T2I class
@@ -101,12 +102,13 @@ class T2I:
                  cfg_scale=7.5,
                  weights="models/ldm/stable-diffusion-v1/model.ckpt",
                  config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml",
-                 sampler="plms",
+                 sampler="klms",
                  latent_channels=4,
                  downsampling_factor=8,
                  ddim_eta=0.0,  # deterministic
                  fixed_code=False,
                  precision='autocast',
+                 full_precision=False,
                  strength=0.75 # default in scripts/img2img.py
     ):
         self.outdir     = outdir
@@ -125,6 +127,7 @@ class T2I:
         self.downsampling_factor = downsampling_factor
         self.ddim_eta            = ddim_eta
         self.precision           = precision
+        self.full_precision      = full_precision
         self.strength            = strength
         self.model      = None     # empty for now
         self.sampler    = None
@@ -256,6 +259,8 @@ class T2I:
 
         model = self.load_model()  # will instantiate the model or return it from cache
 
+        precision_scope = autocast if self.precision=="autocast" else nullcontext
+
         # grid and individual are mutually exclusive, with individual taking priority.
         # not necessary, but needed for compatability with dream bot
         if (grid is None):
@@ -279,7 +284,8 @@ class T2I:
         assert os.path.isfile(init_img)
         init_image = self._load_img(init_img).to(self.device)
         init_image = repeat(init_image, '1 ... -> b ...', b=batch_size)
-        init_latent = model.get_first_stage_encoding(model.encode_first_stage(init_image))  # move to latent space
+        with precision_scope("cuda"):
+            init_latent = model.get_first_stage_encoding(model.encode_first_stage(init_image))  # move to latent space
 
         sampler.make_schedule(ddim_num_steps=steps, ddim_eta=ddim_eta, verbose=False)
 
@@ -292,7 +298,6 @@ class T2I:
         t_enc = int(strength * steps)
         print(f"target t_enc is {t_enc} steps")
 
-        precision_scope = autocast if self.precision=="autocast" else nullcontext
         images = list()
         seeds  = list()
 
@@ -385,6 +390,9 @@ class T2I:
             elif self.sampler_name == 'ddim':
                 print("setting sampler to ddim")
                 self.sampler = DDIMSampler(self.model)
+            elif self.sampler_name == 'klms':
+                print("setting sampler to klms")
+                self.sampler = KSampler(self.model,'lms')
             else:
                 print(f"unsupported sampler {self.sampler_name}, defaulting to plms")
                 self.sampler = PLMSSampler(self.model)
@@ -401,6 +409,11 @@ class T2I:
         m, u = model.load_state_dict(sd, strict=False)
         model.cuda()
         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 full precision')
+            model.half()
         return model
 
     def _load_img(self,path):

scripts/dream.py

@@ -1,10 +1,16 @@
-#!/usr/bin/env python
-
-import readline
+#!/usr/bin/env python3
 import argparse
 import shlex
 import atexit
 import os
+import sys
+
+# readline unavailable on windows systems
+try:
+    import readline
+    readline_available = True
+except:
+    readline_available = False
 
 debugging = False
 
@@ -26,9 +32,11 @@ def main():
         weights = "models/ldm/stable-diffusion-v1/model.ckpt"
 
     # command line history will be stored in a file called "~/.dream_history"
-    setup_readline()
+    if readline_available:
+        setup_readline()
 
     print("* Initializing, be patient...\n")
+    sys.path.append('.')
     from pytorch_lightning import logging
     from ldm.simplet2i import T2I
 
@@ -42,6 +50,7 @@ def main():
               outdir=opt.outdir,
               sampler=opt.sampler,
               weights=weights,
+              full_precision=opt.full_precision,
               config=config)
 
     # make sure the output directory exists
@@ -54,7 +63,7 @@ def main():
     # preload the model
     if not debugging:
         t2i.load_model()
-    print("\n* Initialization done! Awaiting your command (-h for help)...")
+    print("\n* Initialization done! Awaiting your command (-h for help, q to quit)...")
 
     log_path   = os.path.join(opt.outdir,"dream_log.txt")
     with open(log_path,'a') as log:
@@ -62,17 +71,26 @@ def main():
         main_loop(t2i,cmd_parser,log)
         log.close()
 
+
 def main_loop(t2i,parser,log):
     ''' prompt/read/execute loop '''
-    while True:
+    done = False
+    
+    while not done:
         try:
             command = input("dream> ")
         except EOFError:
-            print("goodbye!")
+            done = True
             break
 
-        # rearrange the arguments to mimic how it works in the Dream bot.
         elements = shlex.split(command)
+        if elements[0]=='q':  # 
+            done = True
+            break
+        if elements[0].startswith('!dream'): # in case a stored prompt still contains the !dream command
+            elements.pop(0)
+            
+        # rearrange the arguments to mimic how it works in the Dream bot.
         switches = ['']
         switches_started = False
 
@@ -95,12 +113,19 @@ def main_loop(t2i,parser,log):
             print("Try again with a prompt!")
             continue
 
-        if opt.init_img is None:
-            results = t2i.txt2img(**vars(opt))
-        else:
-            results = t2i.img2img(**vars(opt))
-        print("Outputs:")
-        write_log_message(opt,switches,results,log)
+        try:
+            if opt.init_img is None:
+                results = t2i.txt2img(**vars(opt))
+            else:
+                results = t2i.img2img(**vars(opt))
+            print("Outputs:")
+            write_log_message(opt,switches,results,log)
+        except KeyboardInterrupt:
+            print('*interrupted*')
+            continue
+
+    print("goodbye!")
+
 
 def write_log_message(opt,switches,results,logfile):
     ''' logs the name of the output image, its prompt and seed to both the terminal and the log file '''
@@ -142,18 +167,22 @@ def create_argv_parser():
                         type=int,
                         default=1,
                         help="number of images to generate")
+    parser.add_argument('-F','--full_precision',
+                        dest='full_precision',
+                        action='store_true',
+                        help="use slower full precision math for calculations")
     parser.add_argument('-b','--batch_size',
                         type=int,
                         default=1,
                         help="number of images to produce per iteration (currently not working properly - producing too many images)")
     parser.add_argument('--sampler',
-                        choices=['plms','ddim'],
-                        default='plms',
-                        help="which sampler to use")
+                        choices=['plms','ddim', 'klms'],
+                        default='klms',
+                        help="which sampler to use (klms)")
     parser.add_argument('-o',
                         '--outdir',
                         type=str,
-                        default="outputs/txt2img-samples",
+                        default="outputs/img-samples",
                         help="directory in which to place generated images and a log of prompts and seeds")
     return parser
                         
@@ -174,80 +203,81 @@ def create_cmd_parser():
     parser.add_argument('-f','--strength',default=0.75,type=float,help="strength for noising/unnoising. 0.0 preserves image exactly, 1.0 replaces it completely")
     return parser
 
-def setup_readline():
-    readline.set_completer(Completer(['--steps','-s','--seed','-S','--iterations','-n','--batch_size','-b',
-                                      '--width','-W','--height','-H','--cfg_scale','-C','--grid','-g',
-                                      '--individual','-i','--init_img','-I','--strength','-f']).complete)
-    readline.set_completer_delims(" ")
-    readline.parse_and_bind('tab: complete')
-    load_history()
+if readline_available:
+    def setup_readline():
+        readline.set_completer(Completer(['--steps','-s','--seed','-S','--iterations','-n','--batch_size','-b',
+                                          '--width','-W','--height','-H','--cfg_scale','-C','--grid','-g',
+                                          '--individual','-i','--init_img','-I','--strength','-f']).complete)
+        readline.set_completer_delims(" ")
+        readline.parse_and_bind('tab: complete')
+        load_history()
 
-def load_history():
-    histfile = os.path.join(os.path.expanduser('~'),".dream_history")
-    try:
-        readline.read_history_file(histfile)
-        readline.set_history_length(1000)
-    except FileNotFoundError:
-        pass
-    atexit.register(readline.write_history_file,histfile)
+    def load_history():
+        histfile = os.path.join(os.path.expanduser('~'),".dream_history")
+        try:
+            readline.read_history_file(histfile)
+            readline.set_history_length(1000)
+        except FileNotFoundError:
+            pass
+        atexit.register(readline.write_history_file,histfile)
 
-class Completer():
-    def __init__(self,options):
-        self.options = sorted(options)
-        return
+    class Completer():
+        def __init__(self,options):
+            self.options = sorted(options)
+            return
 
-    def complete(self,text,state):
-        if text.startswith('-I') or text.startswith('--init_img'):
-            return self._image_completions(text,state)
-        
-        response = None
-        if state == 0:
-            # This is the first time for this text, so build a match list.
-            if text:
-                self.matches = [s 
-                                for s in self.options
-                                if s and s.startswith(text)]
+        def complete(self,text,state):
+            if text.startswith('-I') or text.startswith('--init_img'):
+                return self._image_completions(text,state)
+
+            response = None
+            if state == 0:
+                # This is the first time for this text, so build a match list.
+                if text:
+                    self.matches = [s 
+                                    for s in self.options
+                                    if s and s.startswith(text)]
+                else:
+                    self.matches = self.options[:]
+
+            # Return the state'th item from the match list,
+            # if we have that many.
+            try:
+                response = self.matches[state]
+            except IndexError:
+                response = None
+            return response
+
+        def _image_completions(self,text,state):
+            # get the path so far
+            if text.startswith('-I'):
+                path = text.replace('-I','',1).lstrip()
+            elif text.startswith('--init_img='):
+                path = text.replace('--init_img=','',1).lstrip()
+
+            matches  = list()
+
+            path = os.path.expanduser(path)
+            if len(path)==0:
+                matches.append(text+'./')
             else:
-                self.matches = self.options[:]
-        
-        # Return the state'th item from the match list,
-        # if we have that many.
-        try:
-            response = self.matches[state]
-        except IndexError:
-            response = None
-        return response
+                dir  = os.path.dirname(path)
+                dir_list = os.listdir(dir)
+                for n in dir_list:
+                    if n.startswith('.') and len(n)>1:
+                        continue
+                    full_path = os.path.join(dir,n)
+                    if full_path.startswith(path):
+                        if os.path.isdir(full_path):
+                            matches.append(os.path.join(os.path.dirname(text),n)+'/')
+                        elif n.endswith('.png'):
+                            matches.append(os.path.join(os.path.dirname(text),n))
 
-    def _image_completions(self,text,state):
-        # get the path so far
-        if text.startswith('-I'):
-            path = text.replace('-I','',1).lstrip()
-        elif text.startswith('--init_img='):
-            path = text.replace('--init_img=','',1).lstrip()
-
-        matches  = list()
-
-        path = os.path.expanduser(path)
-        if len(path)==0:
-            matches.append(text+'./')
-        else:
-            dir  = os.path.dirname(path)
-            dir_list = os.listdir(dir)
-            for n in dir_list:
-                if n.startswith('.') and len(n)>1:
-                    continue
-                full_path = os.path.join(dir,n)
-                if full_path.startswith(path):
-                    if os.path.isdir(full_path):
-                        matches.append(os.path.join(os.path.dirname(text),n)+'/')
-                    elif n.endswith('.png'):
-                        matches.append(os.path.join(os.path.dirname(text),n))
-                
-        try:
-            response = matches[state]
-        except IndexError:
-            response = None
-        return response
+            try:
+                response = matches[state]
+            except IndexError:
+                response = None
+            return response
         
 
 if __name__ == "__main__":

scripts/preload_models.py

@@ -1,5 +1,4 @@
-#!/usr/bin/env python
-
+#!/usr/bin/env python3
 # Before running stable-diffusion on an internet-isolated machine,
 # run this script from one with internet connectivity. The
 # two machines must share a common .cache directory.

scripts/txt2img.py

@@ -12,6 +12,10 @@ from pytorch_lightning import seed_everything
 from torch import autocast
 from contextlib import contextmanager, nullcontext
 
+import accelerate
+import k_diffusion as K
+import torch.nn as nn
+
 from ldm.util import instantiate_from_config
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
@@ -80,6 +84,11 @@ def main():
         action='store_true',
         help="use plms sampling",
     )
+    parser.add_argument(
+        "--klms",
+        action='store_true',
+        help="use klms sampling",
+    )
     parser.add_argument(
         "--laion400m",
         action='store_true',
@@ -190,6 +199,22 @@ def main():
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     model = model.to(device)
 
+    #for klms
+    model_wrap = K.external.CompVisDenoiser(model)
+    accelerator = accelerate.Accelerator()
+    device = accelerator.device
+    class CFGDenoiser(nn.Module):
+        def __init__(self, model):
+            super().__init__()
+            self.inner_model = model
+
+        def forward(self, x, sigma, uncond, cond, cond_scale):
+            x_in = torch.cat([x] * 2)
+            sigma_in = torch.cat([sigma] * 2)
+            cond_in = torch.cat([uncond, cond])
+            uncond, cond = self.inner_model(x_in, sigma_in, cond=cond_in).chunk(2)
+            return uncond + (cond - uncond) * cond_scale
+
     if opt.plms:
         sampler = PLMSSampler(model)
     else:
@@ -226,8 +251,8 @@ def main():
             with model.ema_scope():
                 tic = time.time()
                 all_samples = list()
-                for n in trange(opt.n_iter, desc="Sampling"):
-                    for prompts in tqdm(data, desc="data"):
+                for n in trange(opt.n_iter, desc="Sampling", disable =not accelerator.is_main_process):
+                    for prompts in tqdm(data, desc="data", disable =not accelerator.is_main_process):
                         uc = None
                         if opt.scale != 1.0:
                             uc = model.get_learned_conditioning(batch_size * [""])
@@ -235,18 +260,32 @@ def main():
                             prompts = list(prompts)
                         c = model.get_learned_conditioning(prompts)
                         shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
-                        samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
-                                                         conditioning=c,
-                                                         batch_size=opt.n_samples,
-                                                         shape=shape,
-                                                         verbose=False,
-                                                         unconditional_guidance_scale=opt.scale,
-                                                         unconditional_conditioning=uc,
-                                                         eta=opt.ddim_eta,
-                                                         x_T=start_code)
-
+                        
+                        if not opt.klms:
+                            samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
+                                                            conditioning=c,
+                                                            batch_size=opt.n_samples,
+                                                            shape=shape,
+                                                            verbose=False,
+                                                            unconditional_guidance_scale=opt.scale,
+                                                            unconditional_conditioning=uc,
+                                                            eta=opt.ddim_eta,
+                                                            x_T=start_code)
+                        else:
+                            sigmas = model_wrap.get_sigmas(opt.ddim_steps)
+                            if start_code:
+                                x = start_code
+                            else:
+                                x = torch.randn([opt.n_samples, *shape], device=device) * sigmas[0] # for GPU draw
+                            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, disable=not accelerator.is_main_process)
+                        
                         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)
+                        
+                        if opt.klms:
+                            x_sample = accelerator.gather(x_samples_ddim)
 
                         if not opt.skip_save:
                             for x_sample in x_samples_ddim: