updated README with info on weighted partial prompts

Fix usage of simplified API in readme

README.md

@@ -17,7 +17,11 @@ 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), autocompletion, and more.
+command history (up and down arrows), autocompletion, and more. To help
+keep track of which prompts generated which images, the script writes a
+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
@@ -38,24 +42,31 @@ setting sampler to plms
 * Initialization done! Awaiting your command...
 dream> ashley judd riding a camel -n2 -s150
 Outputs:
-   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
+   outputs/img-samples/00009.png: "ashley judd riding a camel" -n2 -s150 -S 416354203
+   outputs/img-samples/00010.png: "ashley judd riding a camel" -n2 -s150 -S 1362479620
 
 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
+    outputs/img-samples/00011.png: "there's a fly in my soup" -n6 -g -S 2685670268
     seeds for individual rows: [2685670268, 1216708065, 2335773498, 822223658, 714542046, 3395302430]
+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
+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
 ~~~~
 
-The dream> prompt's arguments are pretty much
-identical to those used in the Discord bot, except you don't need to
-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 dream> prompt's arguments are pretty much identical to those used
+in the Discord bot, except you don't need to 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.
+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
 
@@ -73,21 +84,76 @@ The --init_img (-I) option gives the path to the seed picture. --strength (-f) c
 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.
 
+## Weighted Prompts
+
+You may weight different sections of the prompt to tell the sampler to attach different levels of
+priority to them, by adding :(number) to the end of the section you wish to up- or downweight.
+For example consider this prompt:
+
+~~~~
+    tabby cat:0.25 white duck:0.75 hybrid
+~~~~
+
+This will tell the sampler to invest 25% of its effort on the tabby
+cat aspect of the image and 75% on the white duck aspect
+(surprisingly, this example actually works). The prompt weights can
+use any combination of integers and floating point numbers, and they
+do not need to add up to 1. A practical example of using this type of
+weighting is described here:
+https://www.reddit.com/r/StableDiffusion/comments/wvb7q7/using_prompt_weights_to_tweak_an_image_with/
+
 ## 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
+* v1.06 (23 August 2022)
+   * Added weighted prompt support contributed by [xraxra](https://github.com/xraxra)
+   * Example of using weighted prompts to tweak a demonic figure contributed by [bmaltais](https://github.com/bmaltais)
+
+* v1.05 (22 August 2022 - after the drop)
+   * Filenames now use the following formats:
+       000010.95183149.png      -- Two files produced by the same command (e.g. -n2),
+       000010.26742632.png      -- distinguished by a different seed.
+
+       000011.455191342.01.png  -- Two files produced by the same command using
+       000011.455191342.02.png  -- a batch size>1 (e.g. -b2). They have the same seed.
+
+       000011.4160627868.grid#1-4.png  -- a grid of four images (-g); the whole grid can
+                                          be regenerated with the indicated key
+
+    * It should no longer be possible for one image to overwrite another
+    * You can use the "cd" and "pwd" commands at the dream> prompt to set and retrieve
+      the path of the output directory.
+     
+* v1.04 (22 August 2022 - after the drop)
+   * Updated README to reflect installation of the released weights.
+   * Suppressed very noisy and inconsequential warning when loading the frozen CLIP
+   tokenizer.
+
+* v1.03 (22 August 2022)
+   * The original txt2img and img2img scripts from the CompViz repository have been moved into
+     a subfolder named "orig_scripts", to reduce confusion.
+   
+* v1.02 (21 August 2022)
+    * A copy of the prompt and all of its switches and options is now stored in the corresponding
+    image in a tEXt metadata field named "Dream". You can read the prompt using scripts/images2prompt.py,
+    or an image editor that allows you to explore the full metadata.
+        **Please run "conda env update -f environment.yaml" to load the k_lms dependencies!!**
+
+* v1.01 (21 August 2022)
+    * added k_lms sampling. 
+      **Please run "conda env 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
 
+There are separate installation walkthroughs for [Linux/Mac](#linuxmac) and [Windows](#windows).
+
 ### 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)
+* Python (version 3.8.5 recommended; higher may work)
 * git
 
 2. Install the Python Anaconda environment manager using pip3.
@@ -122,19 +188,26 @@ After these steps, your command prompt will be prefixed by "(ldm)" as shown abov
 (ldm) ~/stable-diffusion$ python3 scripts/preload_models.py
 ```
 
+Note that this step is necessary because I modified the original
+just-in-time model loading scheme to allow the script to work on GPU
+machines that are not internet connected. See [Workaround for machines with limited internet connectivity](#workaround-for-machines-with-limited-internet-connectivity)
+
 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"
+For running with the released weights, you will first need to set up an acount with Hugging Face (https://huggingface.co).
+Use your credentials to log in, and then point your browser at https://huggingface.co/CompVis/stable-diffusion-v-1-4-original.
+You may be asked to sign a license agreement at this point.
+
+Click on "Files and versions" near the top of the page, and then click on the file named "sd-v1-4.ckpt". You'll be taken
+to a page that prompts you to click the "download" link. Save the file somewhere safe on your local machine.
+
+Now run the following commands from within the stable-diffusion directory. This will create a symbolic
+link from the stable-diffusion model.ckpt file, to the true location of the sd-v1-4.ckpt file.
+   
 ```
 (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>
+(ldm) ~/stable-diffusion$ ln -sf /path/to/sd-v1-4.ckpt models/ldm/stable-diffusion-v1/model.ckpt
 ```
-These weight files are ~5 GB in size, so downloading may take a while.
 
 8. Start generating images!
 ```
@@ -150,7 +223,7 @@ These weight files are ~5 GB in size, so downloading may take a while.
 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
+#### 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:
 ```
@@ -160,7 +233,8 @@ 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/
+1. Install Python version 3.8.5 from here: https://www.python.org/downloads/windows/
+   (note that several users have reported that later versions do not work properly)
 
 2. Install Anaconda3 (miniconda3 version) from here: https://docs.anaconda.com/anaconda/install/windows/
 
@@ -190,15 +264,37 @@ This will install all python requirements and activate the "ldm" environment whi
 ```
 python scripts\preload_models.py
 ```
-This installs two machine learning models that stable diffusion requires.
+
+This installs several machine learning models that stable diffusion
+requires. (Note that this step is required. I created it because some people
+are using GPU systems that are behind a firewall and the models can't be
+downloaded just-in-time)
 
 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 running with the released weights, you will first need to set up
+an acount with Hugging Face (https://huggingface.co).  Use your
+credentials to log in, and then point your browser at
+https://huggingface.co/CompVis/stable-diffusion-v-1-4-original.  You
+may be asked to sign a license agreement at this point.
 
-For testing with the released weights, create a directory within stable-diffusion named "models\ldm\stable-diffusion-v1".
+Click on "Files and versions" near the top of the page, and then click
+on the file named "sd-v1-4.ckpt". You'll be taken to a page that
+prompts you to click the "download" link. Now save the file somewhere
+safe on your local machine.  The weight file is >4 GB in size, so
+downloading may take a while.
 
-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.
+Now run the following commands from **within the stable-diffusion
+directory** to copy the weights file to the right place:
+   
+```
+mkdir -p models/ldm/stable-diffusion-v1
+copy C:\path\to\sd-v1-4.ckpt models\ldm\stable-diffusion-v1\model.ckpt
+```
+Please replace "C:\path\to\sd-v1.4.ckpt" with the correct path to wherever
+you stashed this file. If you prefer not to copy or move the .ckpt file, 
+you may instead create a shortcut to it from within
+"models\ldm\stable-diffusion-v1\".
 
 10. Start generating images!
 ```
@@ -208,9 +304,9 @@ 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).
+11. Subsequently, to relaunch the script, first activate the Anaconda command window (step 4), enter the stable-diffusion directory (step 6, "cd \path\to\stable-diffusion"), run "conda activate ldm" (step 7b), and then launch the dream script (step 10).
 
-### Updating to newer versions of the script
+#### 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:
 ```
@@ -227,7 +323,7 @@ lets you create images from a prompt in just three lines of code:
 ~~~~
 from ldm.simplet2i import T2I
 model   = T2I()
-outputs = model.text2image("a unicorn in manhattan")
+outputs = model.txt2img("a unicorn in manhattan")
 ~~~~
 
 Outputs is a list of lists in the format [[filename1,seed1],[filename2,seed2]...]
@@ -275,7 +371,9 @@ 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.
 
-*Author:* Lincoln D. Stein <lincoln.stein@gmail.com>
+*Original Author:* Lincoln D. Stein <lincoln.stein@gmail.com>
+
+*Contributions by:* [Peter Kowalczyk](https://github.com/slix), [Henry Harrison](https://github.com/hwharrison), [xraxra](https://github.com/xraxra), and [bmaltais](https://github.com/bmaltais)
 
 # Original README from CompViz/stable-diffusion
 *Stable Diffusion was made possible thanks to a collaboration with [Stability AI](https://stability.ai/) and [Runway](https://runwayml.com/) and builds upon our previous work:*

ldm/modules/encoders/modules.py

@@ -146,8 +146,8 @@ class FrozenCLIPEmbedder(AbstractEncoder):
     """Uses the CLIP transformer encoder for text (from Hugging Face)"""
     def __init__(self, version="openai/clip-vit-large-patch14", device="cuda", max_length=77):
         super().__init__()
-        self.tokenizer = CLIPTokenizer.from_pretrained(version)
-        self.transformer = CLIPTextModel.from_pretrained(version)
+        self.tokenizer = CLIPTokenizer.from_pretrained(version,local_files_only=True)
+        self.transformer = CLIPTextModel.from_pretrained(version,local_files_only=True)
         self.device = device
         self.max_length = max_length
         self.freeze()

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','klms']  // klms
+          sampler_name= ['ddim','plms','klms']  // klms
           grid        = <boolean>     // false
           width       = <integer>     // image width, multiple of 64 (512)
           height      = <integer>     // image height, multiple of 64 (512)
@@ -60,6 +60,7 @@ from torch import autocast
 from contextlib import contextmanager, nullcontext
 import time
 import math
+import re
 
 from ldm.util import instantiate_from_config
 from ldm.models.diffusion.ddim     import DDIMSampler
@@ -77,7 +78,7 @@ class T2I:
     batch_size
     steps
     seed
-    sampler
+    sampler_name
     grid
     individual
     width
@@ -88,6 +89,8 @@ class T2I:
     downsampling_factor
     precision
     strength
+
+The vast majority of these arguments default to reasonable values.
 """
     def __init__(self,
                  outdir="outputs/txt2img-samples",
@@ -101,15 +104,16 @@ class T2I:
                  seed=None,
                  cfg_scale=7.5,
                  weights="models/ldm/stable-diffusion-v1/model.ckpt",
-                 config = "configs/latent-diffusion/txt2img-1p4B-eval.yaml",
-                 sampler="klms",
+                 config = "configs/stable-diffusion/v1-inference.yaml",
+                 sampler_name="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
+                 strength=0.75, # default in scripts/img2img.py
+                 latent_diffusion_weights=False  # just to keep track of this parameter when regenerating prompt
     ):
         self.outdir     = outdir
         self.batch_size      = batch_size
@@ -119,9 +123,9 @@ class T2I:
         self.grid       = grid
         self.steps      = steps
         self.cfg_scale  = cfg_scale
-        self.weights   = weights
+        self.weights    = weights
         self.config     = config
-        self.sampler_name  = sampler
+        self.sampler_name  = sampler_name
         self.fixed_code    = fixed_code
         self.latent_channels     = latent_channels
         self.downsampling_factor = downsampling_factor
@@ -131,6 +135,7 @@ class T2I:
         self.strength            = strength
         self.model      = None     # empty for now
         self.sampler    = None
+        self.latent_diffusion_weights=latent_diffusion_weights
         if seed is None:
             self.seed = self._new_seed()
         else:
@@ -138,7 +143,7 @@ class T2I:
 
     def txt2img(self,prompt,outdir=None,batch_size=None,iterations=None,
                 steps=None,seed=None,grid=None,individual=None,width=None,height=None,
-                cfg_scale=None,ddim_eta=None,strength=None,init_img=None):
+                cfg_scale=None,ddim_eta=None,strength=None,init_img=None,skip_normalize=False):
         """
         Generate an image from the prompt, writing iteration images into the outdir
         The output is a list of lists in the format: [[filename1,seed1], [filename2,seed2],...]
@@ -167,7 +172,6 @@ class T2I:
 
         # make directories and establish names for the output files
         os.makedirs(outdir, exist_ok=True)
-        base_count = len(os.listdir(outdir))-1
 
         start_code = None
         if self.fixed_code:
@@ -181,65 +185,90 @@ class T2I:
         sampler         = self.sampler
         images = list()
         seeds  = list()
-
+        filename = None
+        image_count = 0
         tic    = time.time()
-        
-        with torch.no_grad():
-            with precision_scope("cuda"):
-                with model.ema_scope():
-                    all_samples = list()
-                    for n in trange(iterations, desc="Sampling"):
-                        seed_everything(seed)
-                        for prompts in tqdm(data, desc="data", dynamic_ncols=True):
-                            uc = None
-                            if cfg_scale != 1.0:
-                                uc = model.get_learned_conditioning(batch_size * [""])
-                            if isinstance(prompts, tuple):
-                                prompts = list(prompts)
-                            c = model.get_learned_conditioning(prompts)
-                            shape = [self.latent_channels, height // self.downsampling_factor, width // self.downsampling_factor]
-                            samples_ddim, _ = sampler.sample(S=steps,
-                                                             conditioning=c,
-                                                             batch_size=batch_size,
-                                                             shape=shape,
-                                                             verbose=False,
-                                                             unconditional_guidance_scale=cfg_scale,
-                                                             unconditional_conditioning=uc,
-                                                             eta=ddim_eta,
-                                                             x_T=start_code)
 
-                            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)
+        # Gawd. Too many levels of indent here. Need to refactor into smaller routines!
+        try:
+            with torch.no_grad():
+                with precision_scope("cuda"):
+                    with model.ema_scope():
+                        all_samples = list()
+                        for n in trange(iterations, desc="Sampling"):
+                            seed_everything(seed)
+                            for prompts in tqdm(data, desc="data", dynamic_ncols=True):
+                                uc = None
+                                if cfg_scale != 1.0:
+                                    uc = model.get_learned_conditioning(batch_size * [""])
+                                if isinstance(prompts, tuple):
+                                    prompts = list(prompts)
 
-                            if not grid:
-                                for x_sample in x_samples_ddim:
-                                    x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                                    filename = os.path.join(outdir, f"{base_count:05}.png")
-                                    Image.fromarray(x_sample.astype(np.uint8)).save(filename)
-                                    images.append([filename,seed])
-                                    base_count += 1
-                            else:
-                                all_samples.append(x_samples_ddim)
-                                seeds.append(seed)
+                                # weighted sub-prompts
+                                subprompts,weights = T2I._split_weighted_subprompts(prompts[0])
+                                if len(subprompts) > 1:
+                                    # i dont know if this is correct.. but it works
+                                    c = torch.zeros_like(uc)
+                                    # get total weight for normalizing
+                                    totalWeight = sum(weights)
+                                    # normalize each "sub prompt" and add it
+                                    for i in range(0,len(subprompts)):
+                                        weight = weights[i]
+                                        if not skip_normalize:
+                                            weight = weight / totalWeight
+                                        c = torch.add(c,model.get_learned_conditioning(subprompts[i]), alpha=weight)
+                                else: # just standard 1 prompt
+                                    c = model.get_learned_conditioning(prompts)
 
-                        seed = self._new_seed()
- 
-                    if grid:
-                        images = self._make_grid(samples=all_samples,
-                                                 seeds=seeds,
-                                                 batch_size=batch_size,
-                                                 iterations=iterations,
-                                                 outdir=outdir)
+                                shape = [self.latent_channels, height // self.downsampling_factor, width // self.downsampling_factor]
+                                samples_ddim, _ = sampler.sample(S=steps,
+                                                                 conditioning=c,
+                                                                 batch_size=batch_size,
+                                                                 shape=shape,
+                                                                 verbose=False,
+                                                                 unconditional_guidance_scale=cfg_scale,
+                                                                 unconditional_conditioning=uc,
+                                                                 eta=ddim_eta,
+                                                                 x_T=start_code)
+
+                                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 not grid:
+                                    for x_sample in x_samples_ddim:
+                                        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
+                                        filename = self._unique_filename(outdir,previousname=filename,
+                                                                         seed=seed,isbatch=(batch_size>1))
+                                        assert not os.path.exists(filename)
+                                        Image.fromarray(x_sample.astype(np.uint8)).save(filename)
+                                        images.append([filename,seed])
+                                else:
+                                    all_samples.append(x_samples_ddim)
+                                    seeds.append(seed)
+
+                            image_count += 1
+                            seed = self._new_seed()
+                        if grid:
+                            images = self._make_grid(samples=all_samples,
+                                                     seeds=seeds,
+                                                     batch_size=batch_size,
+                                                     iterations=iterations,
+                                                     outdir=outdir)
+        except KeyboardInterrupt:
+            print('*interrupted*')
+            print('Partial results will be returned; if --grid was requested, nothing will be returned.')
+        except RuntimeError as e:
+            print(str(e))
 
         toc = time.time()
-        print(f'{batch_size * iterations} images generated in',"%4.2fs"% (toc-tic))
+        print(f'{image_count} images generated in',"%4.2fs"% (toc-tic))
 
         return images
         
     # There is lots of shared code between this and txt2img and should be refactored.
     def img2img(self,prompt,outdir=None,init_img=None,batch_size=None,iterations=None,
                 steps=None,seed=None,grid=None,individual=None,width=None,height=None,
-                cfg_scale=None,ddim_eta=None,strength=None):
+                cfg_scale=None,ddim_eta=None,strength=None,skip_normalize=False):
         """
         Generate an image from the prompt and the initial image, writing iteration images into the outdir
         The output is a list of lists in the format: [[filename1,seed1], [filename2,seed2],...]
@@ -279,7 +308,6 @@ class T2I:
 
         # make directories and establish names for the output files
         os.makedirs(outdir, exist_ok=True)
-        base_count = len(os.listdir(outdir))-1
 
         assert os.path.isfile(init_img)
         init_image = self._load_img(init_img).to(self.device)
@@ -300,60 +328,83 @@ class T2I:
 
         images = list()
         seeds  = list()
-
+        filename = None
+        image_count = 0 # actual number of iterations performed
         tic    = time.time()
-        
-        with torch.no_grad():
-            with precision_scope("cuda"):
-                with model.ema_scope():
-                    all_samples = list()
-                    for n in trange(iterations, desc="Sampling"):
-                        seed_everything(seed)
-                        for prompts in tqdm(data, desc="data", dynamic_ncols=True):
-                            uc = None
-                            if cfg_scale != 1.0:
-                                uc = model.get_learned_conditioning(batch_size * [""])
-                            if isinstance(prompts, tuple):
-                                prompts = list(prompts)
-                            c = model.get_learned_conditioning(prompts)
 
-                            # encode (scaled latent)
-                            z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(self.device))
-                            # decode it
-                            samples = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=cfg_scale,
-                                                     unconditional_conditioning=uc,)
+        # Gawd. Too many levels of indent here. Need to refactor into smaller routines!
+        try:
+            with torch.no_grad():
+                with precision_scope("cuda"):
+                    with model.ema_scope():
+                        all_samples = list()
+                        for n in trange(iterations, desc="Sampling"):
+                            seed_everything(seed)
+                            for prompts in tqdm(data, desc="data", dynamic_ncols=True):
+                                uc = None
+                                if cfg_scale != 1.0:
+                                    uc = model.get_learned_conditioning(batch_size * [""])
+                                if isinstance(prompts, tuple):
+                                    prompts = list(prompts)
 
-                            x_samples = model.decode_first_stage(samples)
-                            x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+                                # weighted sub-prompts
+                                subprompts,weights = T2I._split_weighted_subprompts(prompts[0])
+                                if len(subprompts) > 1:
+                                    # i dont know if this is correct.. but it works
+                                    c = torch.zeros_like(uc)
+                                    # get total weight for normalizing
+                                    totalWeight = sum(weights)
+                                    # normalize each "sub prompt" and add it
+                                    for i in range(0,len(subprompts)):
+                                        weight = weights[i]
+                                        if not skip_normalize:
+                                            weight = weight / totalWeight
+                                        c = torch.add(c,model.get_learned_conditioning(subprompts[i]), alpha=weight)
+                                else: # just standard 1 prompt
+                                    c = model.get_learned_conditioning(prompts)
 
-                            if not grid:
-                                for x_sample in x_samples:
-                                    x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
-                                    filename = os.path.join(outdir, f"{base_count:05}.png")
-                                    Image.fromarray(x_sample.astype(np.uint8)).save(filename)
-                                    images.append([filename,seed])
-                                    base_count += 1
-                            else:
-                                all_samples.append(x_samples)
-                                seeds.append(seed)
+                                # encode (scaled latent)
+                                z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(self.device))
+                                # decode it
+                                samples = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=cfg_scale,
+                                                         unconditional_conditioning=uc,)
 
-                        seed = self._new_seed()
+                                x_samples = model.decode_first_stage(samples)
+                                x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
 
-                    if grid:
-                        images = self._make_grid(samples=all_samples,
-                                                 seeds=seeds,
-                                                 batch_size=batch_size,
-                                                 iterations=iterations,
-                                                 outdir=outdir)
+                                if not grid:
+                                    for x_sample in x_samples:
+                                        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
+                                        filename = self._unique_filename(outdir,previousname=filename,
+                                                                         seed=seed,isbatch=(batch_size>1))
+                                        assert not os.path.exists(filename)
+                                        Image.fromarray(x_sample.astype(np.uint8)).save(filename)
+                                        images.append([filename,seed])
+                                else:
+                                    all_samples.append(x_samples)
+                                    seeds.append(seed)
+                            image_count +=1
+                            seed = self._new_seed()
+                        if grid:
+                            images = self._make_grid(samples=all_samples,
+                                                     seeds=seeds,
+                                                     batch_size=batch_size,
+                                                     iterations=iterations,
+                                                     outdir=outdir)
+
+        except KeyboardInterrupt:
+            print('*interrupted*')
+            print('Partial results will be returned; if --grid was requested, nothing will be returned.')
+        except RuntimeError as e:
+            print(str(e))
 
         toc = time.time()
-        print(f'{batch_size * iterations} images generated in',"%4.2fs"% (toc-tic))
+        print(f'{image_count} images generated in',"%4.2fs"% (toc-tic))
 
         return images
 
     def _make_grid(self,samples,seeds,batch_size,iterations,outdir):
         images = list()
-        base_count = len(os.listdir(outdir))-1
         n_rows = batch_size if batch_size>1 else int(math.sqrt(batch_size * iterations))
         # save as grid
         grid = torch.stack(samples, 0)
@@ -362,7 +413,7 @@ class T2I:
 
         # to image
         grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
-        filename = os.path.join(outdir, f"{base_count:05}.png")
+        filename = self._unique_filename(outdir,seed=seeds[0],grid_count=batch_size*iterations)
         Image.fromarray(grid.astype(np.uint8)).save(filename)
         for s in seeds:
             images.append([filename,s])
@@ -412,7 +463,7 @@ class T2I:
         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')
+            print('Using half precision math. Call with --full_precision to use slower but more accurate full precision.')
             model.half()
         return model
 
@@ -426,3 +477,85 @@ class T2I:
         image = image[None].transpose(0, 3, 1, 2)
         image = torch.from_numpy(image)
         return 2.*image - 1.
+
+    def _unique_filename(self,outdir,previousname=None,seed=0,isbatch=False,grid_count=None):
+        revision = 1
+
+        if previousname is None:
+            # count up until we find an unfilled slot
+            dir_list  = [a.split('.',1)[0] for a in os.listdir(outdir)]
+            uniques   = dict.fromkeys(dir_list,True)
+            basecount = 1
+            while f'{basecount:06}' in uniques:
+                basecount += 1
+            if grid_count is not None:
+                grid_label = f'grid#1-{grid_count}'
+                filename = f'{basecount:06}.{seed}.{grid_label}.png'
+            elif isbatch:
+                filename = f'{basecount:06}.{seed}.01.png'
+            else:
+                filename = f'{basecount:06}.{seed}.png'
+            
+            return os.path.join(outdir,filename)
+
+        else:
+            previousname = os.path.basename(previousname)
+            x = re.match('^(\d+)\..*\.png',previousname)
+            if not x:
+                return self._unique_filename(outdir,previousname,seed)
+
+            basecount = int(x.groups()[0])
+            series = 0 
+            finished = False
+            while not finished:
+                series += 1
+                filename = f'{basecount:06}.{seed}.png'
+                if isbatch or os.path.exists(os.path.join(outdir,filename)):
+                    filename = f'{basecount:06}.{seed}.{series:02}.png'
+                finished = not os.path.exists(os.path.join(outdir,filename))
+            return os.path.join(outdir,filename)
+
+    def _split_weighted_subprompts(text):
+        """
+        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
+        """
+        remaining = len(text)
+        prompts = []
+        weights = []
+        while remaining > 0:
+            if ":" in text:
+                idx = text.index(":") # first occurrence from start
+                # grab up to index as sub-prompt
+                prompt = text[:idx]
+                remaining -= idx
+                # remove from main text
+                text = text[idx+1:]
+                # find value for weight 
+                if " " in text:
+                    idx = text.index(" ") # first occurence
+                else: # no space, read to end
+                    idx = len(text)
+                if idx != 0:
+                    try:
+                        weight = float(text[:idx])
+                    except: # couldn't treat as float
+                        print(f"Warning: '{text[:idx]}' is not a value, are you missing a space?")
+                        weight = 1.0
+                else: # no value found
+                    weight = 1.0
+                # remove from main text
+                remaining -= idx
+                text = text[idx+1:]
+                # append the sub-prompt and its weight
+                prompts.append(prompt)
+                weights.append(weight)
+            else: # no : found
+                if len(text) > 0: # there is still text though
+                    # take remainder as weight 1
+                    prompts.append(text)
+                    weights.append(1.0)
+                remaining = 0
+        return prompts, weights

scripts/dream.py

@@ -4,6 +4,7 @@ import shlex
 import atexit
 import os
 import sys
+from PIL import Image,PngImagePlugin
 
 # readline unavailable on windows systems
 try:
@@ -39,7 +40,11 @@ def main():
     sys.path.append('.')
     from pytorch_lightning import logging
     from ldm.simplet2i import T2I
-
+    # these two lines prevent a horrible warning message from appearing
+    # when the frozen CLIP tokenizer is imported
+    import transformers
+    transformers.logging.set_verbosity_error()
+    
     # creating a simple text2image object with a handful of
     # defaults passed on the command line.
     # additional parameters will be added (or overriden) during
@@ -48,10 +53,12 @@ def main():
               height=height,
               batch_size=opt.batch_size,
               outdir=opt.outdir,
-              sampler=opt.sampler,
+              sampler_name=opt.sampler_name,
               weights=weights,
               full_precision=opt.full_precision,
-              config=config)
+              config=config,
+              latent_diffusion_weights=opt.laion400m # this is solely for recreating the prompt
+    )
 
     # make sure the output directory exists
     if not os.path.exists(opt.outdir):
@@ -60,33 +67,71 @@ def main():
     # gets rid of annoying messages about random seed
     logging.getLogger("pytorch_lightning").setLevel(logging.ERROR)
 
+    infile = None
+    try:
+        if opt.infile is not None:
+            infile = open(opt.infile,'r')
+    except FileNotFoundError as e:
+        print(e)
+        exit(-1)
+
     # preload the model
     if not debugging:
         t2i.load_model()
-    print("\n* Initialization done! Awaiting your command (-h for help, q to quit)...")
+    print("\n* Initialization done! Awaiting your command (-h for help, 'q' to quit, 'cd' to change output dir, 'pwd' to print output dir)...")
 
-    log_path   = os.path.join(opt.outdir,"dream_log.txt")
+    log_path   = os.path.join(opt.outdir,'dream_log.txt')
     with open(log_path,'a') as log:
         cmd_parser = create_cmd_parser()
-        main_loop(t2i,cmd_parser,log)
+        main_loop(t2i,cmd_parser,log,infile)
         log.close()
+    if infile:
+        infile.close()
 
 
-def main_loop(t2i,parser,log):
+def main_loop(t2i,parser,log,infile):
     ''' prompt/read/execute loop '''
     done = False
     
     while not done:
         try:
-            command = input("dream> ")
+            command = infile.readline() if infile else input("dream> ") 
         except EOFError:
             done = True
             break
 
-        elements = shlex.split(command)
-        if elements[0]=='q':  # 
+        if infile and len(command)==0:
             done = True
             break
+
+        if command.startswith(('#','//')):
+            continue
+
+        try:
+            elements = shlex.split(command)
+        except ValueError as e:
+            print(str(e))
+            continue
+        
+        if len(elements)==0:
+            continue
+
+        if elements[0]=='q':
+            done = True
+            break
+
+        if elements[0]=='cd' and len(elements)>1:
+            if os.path.exists(elements[1]):
+                print(f"setting image output directory to {elements[1]}")
+                t2i.outdir=elements[1]
+            else:
+                print(f"directory {elements[1]} does not exist")
+            continue
+
+        if elements[0]=='pwd':
+            print(f"current output directory is {t2i.outdir}")
+            continue
+        
         if elements[0].startswith('!dream'): # in case a stored prompt still contains the !dream command
             elements.pop(0)
             
@@ -113,48 +158,82 @@ def main_loop(t2i,parser,log):
             print("Try again with a prompt!")
             continue
 
-        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
+        if opt.init_img is None:
+            results = t2i.txt2img(**vars(opt))
+        else:
+            results = t2i.img2img(**vars(opt))
+        print("Outputs:")
+        write_log_message(t2i,opt,results,log)
+            
 
     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 '''
-    if opt.grid:
-        _output_for_grid(switches,results,logfile)
-    else:
-        _output_for_individual(switches,results,logfile)
+def write_log_message(t2i,opt,results,logfile):
+    ''' logs the name of the output image, its prompt and seed to the terminal, log file, and a Dream text chunk in the PNG metadata '''
+    switches = _reconstruct_switches(t2i,opt)
+    prompt_str = ' '.join(switches)
+
+    # when multiple images are produced in batch, then we keep track of where each starts
+    last_seed  = None
+    img_num    = 1
+    batch_size = opt.batch_size or t2i.batch_size
+    seenit     = {}
+
+    seeds = [a[1] for a in results]
+    if batch_size > 1:
+        seeds = f"(seeds for each batch row: {seeds})"
+    else:
+        seeds = f"(seeds for individual images: {seeds})"
 
-def _output_for_individual(switches,results,logfile):
     for r in results:
-        log_message = " ".join(['   ',str(r[0])+':',
-                                f'"{switches[0]}"',
-                                *switches[1:],f'-S {r[1]}'])
+        seed = r[1]
+        log_message = (f'{r[0]}: {prompt_str} -S{seed}')
+
+        if batch_size > 1:
+            if seed != last_seed:
+                img_num = 1
+                log_message += f' # (batch image {img_num} of {batch_size})'
+            else:
+                img_num += 1
+                log_message += f' # (batch image {img_num} of {batch_size})'
+            last_seed = seed
         print(log_message)
         logfile.write(log_message+"\n")
         logfile.flush()
+        if r[0] not in seenit:
+            seenit[r[0]] = True
+            try:
+                if opt.grid:
+                    _write_prompt_to_png(r[0],f'{prompt_str} -g -S{seed} {seeds}')
+                else:
+                    _write_prompt_to_png(r[0],f'{prompt_str} -S{seed}')
+            except FileNotFoundError:
+                print(f"Could not open file '{r[0]}' for reading")
 
-def _output_for_grid(switches,results,logfile):
-    first_seed = results[0][1]
-    log_message = " ".join(['   ',str(results[0][0])+':',
-                            f'"{switches[0]}"',
-                            *switches[1:],f'-S {results[0][1]}'])
-    print(log_message)
-    logfile.write(log_message+"\n")
-    all_seeds   = [row[1] for row in results]
-    log_message = f'    seeds for individual rows: {all_seeds}'
-    print(log_message)
-    logfile.write(log_message+"\n")
+def _reconstruct_switches(t2i,opt):
+    '''Normalize the prompt and switches'''
+    switches = list()
+    switches.append(f'"{opt.prompt}"')
+    switches.append(f'-s{opt.steps        or t2i.steps}')
+    switches.append(f'-b{opt.batch_size   or t2i.batch_size}')
+    switches.append(f'-W{opt.width        or t2i.width}')
+    switches.append(f'-H{opt.height       or t2i.height}')
+    switches.append(f'-C{opt.cfg_scale    or t2i.cfg_scale}')
+    if opt.init_img:
+        switches.append(f'-I{opt.init_img}')
+    if opt.strength and opt.init_img is not None:
+        switches.append(f'-f{opt.strength or t2i.strength}')
+    if t2i.full_precision:
+        switches.append('-F')
+    return switches
 
+def _write_prompt_to_png(path,prompt):
+    info = PngImagePlugin.PngInfo()
+    info.add_text("Dream",prompt)
+    im = Image.open(path)
+    im.save(path,"PNG",pnginfo=info)
+    
 def create_argv_parser():
     parser = argparse.ArgumentParser(description="Parse script's command line args")
     parser.add_argument("--laion400m",
@@ -162,7 +241,11 @@ def create_argv_parser():
                         "-l",
                         dest='laion400m',
                         action='store_true',
-                        help="fallback to the latent diffusion (LAION4400M) weights and config")
+                        help="fallback to the latent diffusion (laion400m) weights and config")
+    parser.add_argument("--from_file",
+                        dest='infile',
+                        type=str,
+                        help="if specified, load prompts from this file")
     parser.add_argument('-n','--iterations',
                         type=int,
                         default=1,
@@ -174,11 +257,12 @@ def create_argv_parser():
     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',
+                        help="number of images to produce per iteration (faster, but doesn't generate individual seeds")
+    parser.add_argument('--sampler','-m',
+                        dest="sampler_name",
                         choices=['plms','ddim', 'klms'],
                         default='klms',
-                        help="which sampler to use (klms)")
+                        help="which sampler to use (klms) - can only be set on command line")
     parser.add_argument('-o',
                         '--outdir',
                         type=str,
@@ -193,7 +277,7 @@ def create_cmd_parser():
     parser.add_argument('-s','--steps',type=int,help="number of steps")
     parser.add_argument('-S','--seed',type=int,help="image seed")
     parser.add_argument('-n','--iterations',type=int,default=1,help="number of samplings to perform")
-    parser.add_argument('-b','--batch_size',type=int,default=1,help="number of images to produce per sampling (currently broken)")
+    parser.add_argument('-b','--batch_size',type=int,default=1,help="number of images to produce per sampling")
     parser.add_argument('-W','--width',type=int,help="image width, multiple of 64")
     parser.add_argument('-H','--height',type=int,help="image height, multiple of 64")
     parser.add_argument('-C','--cfg_scale',default=7.5,type=float,help="prompt configuration scale")
@@ -201,11 +285,13 @@ def create_cmd_parser():
     parser.add_argument('-i','--individual',action='store_true',help="generate individual files (default)")
     parser.add_argument('-I','--init_img',type=str,help="path to input image (supersedes width and height)")
     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")
+    parser.add_argument('-x','--skip_normalize',action='store_true',help="skip subprompt weight normalization")
     return parser
 
 if readline_available:
     def setup_readline():
-        readline.set_completer(Completer(['--steps','-s','--seed','-S','--iterations','-n','--batch_size','-b',
+        readline.set_completer(Completer(['cd','pwd',
+                                          '--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(" ")
@@ -227,8 +313,13 @@ if readline_available:
             return
 
         def complete(self,text,state):
-            if text.startswith('-I') or text.startswith('--init_img'):
-                return self._image_completions(text,state)
+            buffer = readline.get_line_buffer()
+            
+            if text.startswith(('-I','--init_img')):
+                return self._path_completions(text,state,('.png'))
+
+            if buffer.strip().endswith('cd') or text.startswith(('.','/')):
+                return self._path_completions(text,state,())
 
             response = None
             if state == 0:
@@ -248,12 +339,14 @@ if readline_available:
                 response = None
             return response
 
-        def _image_completions(self,text,state):
+        def _path_completions(self,text,state,extensions):
             # 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()
+            else:
+                path = text
 
             matches  = list()
 
@@ -270,7 +363,7 @@ if readline_available:
                     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'):
+                        elif n.endswith(extensions):
                             matches.append(os.path.join(os.path.dirname(text),n))
 
             try:
@@ -278,7 +371,6 @@ if readline_available:
             except IndexError:
                 response = None
             return response
-        
 
 if __name__ == "__main__":
     main()

scripts/images2prompt.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+'''This script reads the "Dream" Stable Diffusion prompt embedded in files generated by dream.py'''
+
+import sys
+from PIL import Image,PngImagePlugin
+
+if len(sys.argv) < 2:
+    print("Usage: file2prompt.py <file1.png> <file2.png> <file3.png>...")
+    print("This script opens up the indicated dream.py-generated PNG file(s) and prints out the prompt used to generate them.")
+    exit(-1)
+
+filenames = sys.argv[1:]
+for f in filenames:
+    try:
+        im = Image.open(f)
+        try:
+            prompt = im.text['Dream']
+        except KeyError:
+            prompt = ''
+        print(f'{f}: {prompt}')
+    except FileNotFoundError:
+        sys.stderr.write(f'{f} not found\n')
+        continue
+    except PermissionError:
+        sys.stderr.write(f'{f} could not be opened due to inadequate permissions\n')
+        continue
+        
+
+
+

scripts/preload_models.py

@@ -2,6 +2,10 @@
 # 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.
+import sys
+import transformers
+
+transformers.logging.set_verbosity_error()
 
 # this will preload the Bert tokenizer fles
 print("preloading bert tokenizer...")
@@ -10,7 +14,19 @@ tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
 print("...success")
 
 # this will download requirements for Kornia
-print("preloading Kornia requirements...")
+print("preloading Kornia requirements (ignore the warnings)...")
 import kornia
 print("...success")
 
+# doesn't work - probably wrong logger
+# logging.getLogger('transformers.tokenization_utils').setLevel(logging.ERROR)
+version='openai/clip-vit-large-patch14'
+
+print('preloading CLIP model (Ignore the warnings)...')
+sys.stdout.flush()
+import clip
+from transformers import CLIPTokenizer, CLIPTextModel
+tokenizer  =CLIPTokenizer.from_pretrained(version)
+transformer=CLIPTextModel.from_pretrained(version)
+print('\n\n...success')
+