feature/refactor/fix: better defaults. correct version in metadata

- feature: use different default steps and image sizes depending on sampler and model selceted
- fix: #110 use proper version in image metadata
- refactor: samplers all have their own class that inherits from ImageSampler

README.md

@@ -9,15 +9,13 @@ AI imagined images. Pythonic generation of stable diffusion images.
 
 "just works" on Linux and macOS(M1) (and maybe windows?).
 
-**🎉 Try Stable Diffusion v2 with alpha version of imaginairy:**
-
-`pip install imaginairy==6.0.0a0 --upgrade`
-
 ## Examples
 ```bash
 # on macOS, make sure rust is installed first
 >> pip install imaginairy
 >> imagine "a scenic landscape" "a photo of a dog" "photo of a fruit bowl" "portrait photo of a freckled woman"
+# Stable Diffusion 2.0
+>> imagine --model SD-2.0 "a forest"
 ```
 
 <details closed>
@@ -215,7 +213,7 @@ imagine_image_files(prompts, outdir="./my-art")
 ## Requirements
 - ~10 gb space for models to download
 - A decent computer with either a CUDA supported graphics card or M1 processor.
-- Python installed. Preferably Python 3.10.
+- Python installed. Preferably Python 3.10.  (not conda)
 - For macOS [rust](https://www.rust-lang.org/tools/install) and setuptools-rust must be installed to compile the `tokenizer` library.
 They can be installed via: `curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh` and `pip install setuptools-rust`
     
@@ -233,12 +231,16 @@ docker run -it --gpus all -v $HOME/.cache/huggingface:/root/.cache/huggingface -
 
 ## ChangeLog
 
-**6.0.0a**
+**6.1.0**
+- feature: use different default steps and image sizes depending on sampler and model selceted
+- fix: #110 use proper version in image metadata
+- refactor: samplers all have their own class that inherits from ImageSampler
 - feature: 🎉🎉🎉 Stable Diffusion 2.0
-  - Tested on MacOS and Linux
+  - `--model SD-2.0` to use (it makes worse images than 1.5 though...) 
+  - Tested on macOS and Linux
   - All samplers working for new 512x512 model
   - New inpainting model working
-  - 768x768 model working for DDIM sampler only
+  - 768x768 model working for DDIM sampler only (`--model SD-2.0 --sampler DDIM`)
 
 **5.1.0**
 - feature: add progress image callback
@@ -254,9 +256,9 @@ inpainting model will automatically be used for any image-masking task
 - feature: added `DPM++ 2S a` and `DPM++ 2M` samplers.
 - feature: improve progress image logging
 - fix: fix bug with `--show-work`. fixes #84
-- fix: add workaround for pytorch bug affecting MacOS users using the new `DPM++ 2S a` and `DPM++ 2M` samplers.
+- fix: add workaround for pytorch bug affecting macOS users using the new `DPM++ 2S a` and `DPM++ 2M` samplers.
 - fix: add workaround for pytorch mps bug affecting `k_dpm_fast` sampler. fixes #75
-- fix: larger image sizes now work on MacOS. fixes #8
+- fix: larger image sizes now work on macOS. fixes #8
 
 **4.1.0**
  - feature: allow dynamic switching between models/weights `--model SD-1.5` or `--model SD-1.4` or `--model path/my-custom-weights.ckpt`)

imaginairy/api.py

@@ -23,7 +23,8 @@ from imaginairy.log_utils import (
 )
 from imaginairy.model_manager import get_diffusion_model
 from imaginairy.safety import SafetyMode, create_safety_score
-from imaginairy.samplers.base import NoiseSchedule, get_sampler, noise_an_image
+from imaginairy.samplers import SAMPLER_LOOKUP
+from imaginairy.samplers.base import NoiseSchedule, noise_an_image
 from imaginairy.schema import ImaginePrompt, ImagineResult
 from imaginairy.utils import (
     fix_torch_group_norm,
@@ -169,8 +170,8 @@ def imagine(
                     prompt.height // downsampling_factor,
                     prompt.width // downsampling_factor,
                 ]
-
-                sampler = get_sampler(prompt.sampler_type, model)
+                SamplerCls = SAMPLER_LOOKUP[prompt.sampler_type.lower()]
+                sampler = SamplerCls(model)
                 mask = mask_image = mask_image_orig = mask_grayscale = None
                 t_enc = init_latent = init_latent_noised = None
                 if prompt.init_image:

imaginairy/cmds.py

@@ -4,11 +4,11 @@ import math
 import click
 from click_shell import shell
 
-from imaginairy import LazyLoadingImage, generate_caption
+from imaginairy import LazyLoadingImage, config, generate_caption
 from imaginairy.api import imagine_image_files
 from imaginairy.enhancers.prompt_expansion import expand_prompts
 from imaginairy.log_utils import configure_logging
-from imaginairy.samplers.base import SAMPLER_TYPE_OPTIONS
+from imaginairy.samplers import SAMPLER_TYPE_OPTIONS
 from imaginairy.schema import ImaginePrompt
 
 logger = logging.getLogger(__name__)
@@ -51,7 +51,7 @@ logger = logging.getLogger(__name__)
 @click.option(
     "-h",
     "--height",
-    default=512,
+    default=None,
     show_default=True,
     type=int,
     help="Image height. Should be multiple of 64.",
@@ -59,14 +59,14 @@ logger = logging.getLogger(__name__)
 @click.option(
     "-w",
     "--width",
-    default=512,
+    default=None,
     show_default=True,
     type=int,
     help="Image width. Should be multiple of 64.",
 )
 @click.option(
     "--steps",
-    default=15,
+    default=None,
     type=int,
     show_default=True,
     help="How many diffusion steps to run. More steps, more detail, but with diminishing returns.",
@@ -88,7 +88,7 @@ logger = logging.getLogger(__name__)
 @click.option(
     "--sampler-type",
     "--sampler",
-    default="k_dpmpp_2m",
+    default=config.DEFAULT_SAMPLER,
     show_default=True,
     type=click.Choice(SAMPLER_TYPE_OPTIONS),
     help="What sampling strategy to use.",
@@ -163,7 +163,8 @@ logger = logging.getLogger(__name__)
     "--model-weights-path",
     "--model",
     help="Model to use. Should be one of SD-1.4, SD-1.5, or a path to custom weights. Defaults to SD-1.5.",
-    default=None,
+    show_default=True,
+    default=config.DEFAULT_MODEL,
 )
 @click.option(
     "--prompt-library-path",

imaginairy/config.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+
+DEFAULT_MODEL = "SD-1.5"
+DEFAULT_SAMPLER = "k_dpmpp_2m"
+
+
+@dataclass
+class ModelConfig:
+    short_name: str
+    config_path: str
+    weights_url: str
+    default_image_size: int
+
+
+MODEL_CONFIGS = [
+    ModelConfig(
+        short_name="SD-1.4",
+        config_path="configs/stable-diffusion-v1.yaml",
+        weights_url="https://huggingface.co/bstddev/sd-v1-4/resolve/77221977fa8de8ab8f36fac0374c120bd5b53287/sd-v1-4.ckpt",
+        default_image_size=512,
+    ),
+    ModelConfig(
+        short_name="SD-1.5",
+        config_path="configs/stable-diffusion-v1.yaml",
+        weights_url="https://huggingface.co/acheong08/SD-V1-5-cloned/resolve/fc392f6bd4345b80fc2256fa8aded8766b6c629e/v1-5-pruned-emaonly.ckpt",
+        default_image_size=512,
+    ),
+    ModelConfig(
+        short_name="SD-1.5-inpaint",
+        config_path="configs/stable-diffusion-v1-inpaint.yaml",
+        weights_url="https://huggingface.co/julienacquaviva/inpainting/resolve/2155ff7fe38b55f4c0d99c2f1ab9b561f8311ca7/sd-v1-5-inpainting.ckpt",
+        default_image_size=512,
+    ),
+    ModelConfig(
+        short_name="SD-2.0",
+        config_path="configs/stable-diffusion-v2-inference.yaml",
+        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt",
+        default_image_size=512,
+    ),
+    ModelConfig(
+        short_name="SD-2.0-inpaint",
+        config_path="configs/stable-diffusion-v2-inpainting-inference.yaml",
+        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
+        default_image_size=512,
+    ),
+    ModelConfig(
+        short_name="SD-2.0-v",
+        config_path="configs/stable-diffusion-v2-inference-v.yaml",
+        weights_url="https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt",
+        default_image_size=768,
+    ),
+]
+
+MODEL_CONFIG_SHORTCUTS = {m.short_name: m for m in MODEL_CONFIGS}

imaginairy/model_manager.py

@@ -10,38 +10,12 @@ from transformers import cached_path
 from transformers.utils.hub import TRANSFORMERS_CACHE, HfFolder
 from transformers.utils.hub import url_to_filename as tf_url_to_filename
 
+from imaginairy import config as iconfig
 from imaginairy.paths import PKG_ROOT
 from imaginairy.utils import get_device, instantiate_from_config
 
 logger = logging.getLogger(__name__)
 
-MODEL_SHORTCUTS = {
-    "SD-1.4": (
-        "configs/stable-diffusion-v1.yaml",
-        "https://huggingface.co/bstddev/sd-v1-4/resolve/77221977fa8de8ab8f36fac0374c120bd5b53287/sd-v1-4.ckpt",
-    ),
-    "SD-1.5": (
-        "configs/stable-diffusion-v1.yaml",
-        "https://huggingface.co/acheong08/SD-V1-5-cloned/resolve/fc392f6bd4345b80fc2256fa8aded8766b6c629e/v1-5-pruned-emaonly.ckpt",
-    ),
-    "SD-1.5-inpaint": (
-        "configs/stable-diffusion-v1-inpaint.yaml",
-        "https://huggingface.co/julienacquaviva/inpainting/resolve/2155ff7fe38b55f4c0d99c2f1ab9b561f8311ca7/sd-v1-5-inpainting.ckpt",
-    ),
-    "SD-2.0": (
-        "configs/stable-diffusion-v2-inference.yaml",
-        "https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt",
-    ),
-    "SD-2.0-inpaint": (
-        "configs/stable-diffusion-v2-inpainting-inference.yaml",
-        "https://huggingface.co/stabilityai/stable-diffusion-2-inpainting/resolve/main/512-inpainting-ema.ckpt",
-    ),
-    "SD-2.0-v": (
-        "configs/stable-diffusion-v2-inference-v.yaml",
-        "https://huggingface.co/stabilityai/stable-diffusion-2/resolve/main/768-v-ema.ckpt",
-    ),
-}
-DEFAULT_MODEL = "SD-2.0"
 
 LOADED_MODELS = {}
 MOST_RECENTLY_LOADED_MODEL = None
@@ -126,7 +100,7 @@ def load_model_from_config(config, weights_location):
 
 
 def get_diffusion_model(
-    weights_location=DEFAULT_MODEL,
+    weights_location=iconfig.DEFAULT_MODEL,
     config_path="configs/stable-diffusion-v1.yaml",
     half_mode=None,
     for_inpainting=False,
@@ -146,13 +120,13 @@ def get_diffusion_model(
                 f"Failed to load inpainting model. Attempting to fall-back to standard model.   {str(e)}"
             )
             return _get_diffusion_model(
-                DEFAULT_MODEL, config_path, half_mode, for_inpainting=False
+                iconfig.DEFAULT_MODEL, config_path, half_mode, for_inpainting=False
             )
         raise e
 
 
 def _get_diffusion_model(
-    weights_location=DEFAULT_MODEL,
+    weights_location=iconfig.DEFAULT_MODEL,
     config_path="configs/stable-diffusion-v1.yaml",
     half_mode=None,
     for_inpainting=False,
@@ -164,11 +138,22 @@ def _get_diffusion_model(
     """
     global MOST_RECENTLY_LOADED_MODEL  # noqa
     if weights_location is None:
-        weights_location = DEFAULT_MODEL
-    if for_inpainting and f"{weights_location}-inpaint" in MODEL_SHORTCUTS:
-        config_path, weights_location = MODEL_SHORTCUTS[f"{weights_location}-inpaint"]
-    elif weights_location in MODEL_SHORTCUTS:
-        config_path, weights_location = MODEL_SHORTCUTS[weights_location]
+        weights_location = iconfig.DEFAULT_MODEL
+    if (
+        for_inpainting
+        and f"{weights_location}-inpaint" in iconfig.MODEL_CONFIG_SHORTCUTS
+    ):
+        model_config = iconfig.MODEL_CONFIG_SHORTCUTS[f"{weights_location}-inpaint"]
+        config_path, weights_location = (
+            model_config.config_path,
+            model_config.weights_url,
+        )
+    elif weights_location in iconfig.MODEL_CONFIG_SHORTCUTS:
+        model_config = iconfig.MODEL_CONFIG_SHORTCUTS[weights_location]
+        config_path, weights_location = (
+            model_config.config_path,
+            model_config.weights_url,
+        )
 
     key = (config_path, weights_location)
     if key not in LOADED_MODELS:
@@ -183,6 +168,13 @@ def _get_diffusion_model(
     return model
 
 
+def get_model_default_image_size(weights_location):
+    model_config = iconfig.MODEL_CONFIG_SHORTCUTS.get(weights_location, None)
+    if model_config:
+        return model_config.default_image_size
+    return 512
+
+
 def get_current_diffusion_model():
     return MOST_RECENTLY_LOADED_MODEL
 

imaginairy/samplers/__init__.py

@@ -0,0 +1,22 @@
+from imaginairy.samplers import kdiff
+from imaginairy.samplers.base import SamplerName  # noqa
+from imaginairy.samplers.ddim import DDIMSampler
+from imaginairy.samplers.plms import PLMSSampler
+
+SAMPLERS = [
+    PLMSSampler,
+    DDIMSampler,
+    kdiff.DPMFastSampler,
+    kdiff.DPMAdaptiveSampler,
+    kdiff.LMSSampler,
+    kdiff.DPM2Sampler,
+    kdiff.DPM2AncestralSampler,
+    kdiff.DPMPP2MSampler,
+    kdiff.DPMPP2SAncestralSampler,
+    kdiff.EulerSampler,
+    kdiff.EulerAncestralSampler,
+    kdiff.HeunSampler,
+]
+
+SAMPLER_LOOKUP = {sampler.short_name: sampler for sampler in SAMPLERS}
+SAMPLER_TYPE_OPTIONS = [sampler.short_name for sampler in SAMPLERS]

imaginairy/samplers/base.py

@@ -1,9 +1,9 @@
 # pylama:ignore=W0613
 import logging
+from abc import ABC
 
 import numpy as np
 import torch
-from torch import nn
 
 from imaginairy.log_utils import log_latent
 from imaginairy.modules.diffusion.util import (
@@ -15,100 +15,31 @@ from imaginairy.utils import get_device
 
 logger = logging.getLogger(__name__)
 
-SAMPLER_TYPE_OPTIONS = [
-    "plms",
-    "ddim",
-    "k_dpm_fast",
-    "k_dpm_adaptive",
-    "k_lms",
-    "k_dpm_2",
-    "k_dpm_2_a",
-    "k_dpmpp_2m",
-    "k_dpmpp_2s_a",
-    "k_euler",
-    "k_euler_a",
-    "k_heun",
-]
-
-_k_sampler_type_lookup = {
-    "k_dpm_fast": "dpm_fast",
-    "k_dpm_adaptive": "dpm_adaptive",
-    "k_dpm_2": "dpm_2",
-    "k_dpm_2_a": "dpm_2_ancestral",
-    "k_dpmpp_2m": "dpmpp_2m",
-    "k_dpmpp_2s_a": "dpmpp_2s_ancestral",
-    "k_euler": "euler",
-    "k_euler_a": "euler_ancestral",
-    "k_heun": "heun",
-    "k_lms": "lms",
-}
-
-
-def get_sampler(sampler_type, model):
-    from imaginairy.samplers.ddim import DDIMSampler  # noqa
-    from imaginairy.samplers.kdiff import KDiffusionSampler  # noqa
-    from imaginairy.samplers.plms import PLMSSampler  # noqa
-
-    sampler_type = sampler_type.lower()
-    if sampler_type == "plms":
-        return PLMSSampler(model)
-    if sampler_type == "ddim":
-        return DDIMSampler(model)
-    if sampler_type.startswith("k_"):
-        sampler_type = _k_sampler_type_lookup[sampler_type]
-        return KDiffusionSampler(model, sampler_type)
-    raise ValueError("invalid sampler_type")
-
-
-class CFGDenoiser(nn.Module):
-    """
-    Conditional forward guidance wrapper
-    """
 
-    def __init__(self, model):
-        super().__init__()
-        self.inner_model = model
-        self.device = get_device()
+class SamplerName:
+    PLMS = "plms"
+    DDIM = "ddim"
+    K_DPM_FAST = "k_dpm_fast"
+    K_DPM_ADAPTIVE = "k_dpm_adaptive"
+    K_LMS = "k_lms"
+    K_DPM_2 = "k_dpm_2"
+    K_DPM_2_ANCESTRAL = "k_dpm_2_a"
+    K_DPMPP_2M = "k_dpmpp_2m"
+    K_DPMPP_2S_ANCESTRAL = "k_dpmpp_2s_a"
+    K_EULER = "k_euler"
+    K_EULER_ANCESTRAL = "k_euler_a"
+    K_HEUN = "k_heun"
 
-    def forward(
-        self,
-        x,
-        sigma,
-        uncond,
-        cond,
-        cond_scale,
-        mask=None,
-        mask_noise=None,
-        orig_latent=None,
-    ):
-        def _wrapper(noisy_latent_in, time_encoding_in, conditioning_in):
-            return self.inner_model(
-                noisy_latent_in, time_encoding_in, cond=conditioning_in
-            )
-
-        if mask is not None:
-            assert orig_latent is not None
-            t = self.inner_model.sigma_to_t(sigma, quantize=True)
-            big_sigma = max(sigma, 1)
-            x = mask_blend(
-                noisy_latent=x,
-                orig_latent=orig_latent * big_sigma,
-                mask=mask,
-                mask_noise=mask_noise * big_sigma,
-                ts=t,
-                model=self.inner_model.inner_model,
-            )
-
-        noise_pred = get_noise_prediction(
-            denoise_func=_wrapper,
-            noisy_latent=x,
-            time_encoding=sigma,
-            neutral_conditioning=uncond,
-            positive_conditioning=cond,
-            signal_amplification=cond_scale,
-        )
 
-        return noise_pred
+class ImageSampler(ABC):
+    short_name: str
+    name: str
+    default_steps: int
+    default_size: int
+
+    def __init__(self, model):
+        self.model = model
+        self.device = get_device()
 
 
 def ensure_4_dim(t: torch.Tensor):

imaginairy/samplers/ddim.py

@@ -7,22 +7,28 @@ from tqdm import tqdm
 
 from imaginairy.log_utils import increment_step, log_latent
 from imaginairy.modules.diffusion.util import extract_into_tensor, noise_like
-from imaginairy.samplers.base import NoiseSchedule, get_noise_prediction, mask_blend
+from imaginairy.samplers.base import (
+    ImageSampler,
+    NoiseSchedule,
+    SamplerName,
+    get_noise_prediction,
+    mask_blend,
+)
 from imaginairy.utils import get_device
 
 logger = logging.getLogger(__name__)
 
 
-class DDIMSampler:
+class DDIMSampler(ImageSampler):
     """
     Denoising Diffusion Implicit Models
 
     https://arxiv.org/abs/2010.02502
     """
 
-    def __init__(self, model):
-        self.model = model
-        self.device = get_device()
+    short_name = SamplerName.DDIM
+    name = "Denoising Diffusion Implicit Models"
+    default_steps = 40
 
     @torch.no_grad()
     def sample(

imaginairy/samplers/kdiff.py

@@ -1,8 +1,16 @@
 # pylama:ignore=W0613
+from abc import ABC
+
 import torch
+from torch import nn
 
 from imaginairy.log_utils import increment_step, log_latent
-from imaginairy.samplers.base import CFGDenoiser
+from imaginairy.samplers.base import (
+    ImageSampler,
+    SamplerName,
+    get_noise_prediction,
+    mask_blend,
+)
 from imaginairy.utils import get_device
 from imaginairy.vendored.k_diffusion import sampling as k_sampling
 from imaginairy.vendored.k_diffusion.external import CompVisDenoiser
@@ -44,27 +52,12 @@ def sample_dpm_fast(model, x, sigmas, extra_args=None, disable=False, callback=N
     )
 
 
-class KDiffusionSampler:
-
-    sampler_lookup = {
-        "dpm_fast": sample_dpm_fast,
-        "dpm_adaptive": sample_dpm_adaptive,
-        "dpm_2": k_sampling.sample_dpm_2,
-        "dpm_2_ancestral": k_sampling.sample_dpm_2_ancestral,
-        "dpmpp_2m": k_sampling.sample_dpmpp_2m,
-        "dpmpp_2s_ancestral": k_sampling.sample_dpmpp_2s_ancestral,
-        "euler": k_sampling.sample_euler,
-        "euler_ancestral": k_sampling.sample_euler_ancestral,
-        "heun": k_sampling.sample_heun,
-        "lms": k_sampling.sample_lms,
-    }
+class KDiffusionSampler(ImageSampler, ABC):
+    sampler_func: callable
 
-    def __init__(self, model, sampler_name):
-        self.model = model
+    def __init__(self, model):
+        super().__init__(model)
         self.cv_denoiser = StandardCompVisDenoiser(model)
-        self.sampler_name = sampler_name
-        self.sampler_func = self.sampler_lookup[sampler_name]
-        self.device = get_device()
 
     def sample(
         self,
@@ -131,3 +124,124 @@ class KDiffusionSampler:
         )
 
         return samples
+
+
+class DPMFastSampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPM_FAST
+    name = "Diffusion probabilistic models - fast"
+    default_steps = 15
+    sampler_func = staticmethod(sample_dpm_fast)
+
+
+class DPMAdaptiveSampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPM_ADAPTIVE
+    name = "Diffusion probabilistic models - adaptive"
+    default_steps = 40
+    sampler_func = staticmethod(sample_dpm_adaptive)
+
+
+class DPM2Sampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPM_2
+    name = "Diffusion probabilistic models - 2"
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_dpm_2)
+
+
+class DPM2AncestralSampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPM_2_ANCESTRAL
+    name = "Diffusion probabilistic models - 2 ancestral"
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_dpm_2_ancestral)
+
+
+class DPMPP2MSampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPMPP_2M
+    name = "Diffusion probabilistic models - 2m"
+    default_steps = 15
+    sampler_func = staticmethod(k_sampling.sample_dpmpp_2m)
+
+
+class DPMPP2SAncestralSampler(KDiffusionSampler):
+    short_name = SamplerName.K_DPMPP_2S_ANCESTRAL
+    name = "Ancestral sampling with DPM-Solver++(2S) second-order steps."
+    default_steps = 15
+    sampler_func = staticmethod(k_sampling.sample_dpmpp_2s_ancestral)
+
+
+class EulerSampler(KDiffusionSampler):
+    short_name = SamplerName.K_EULER
+    name = "Algorithm 2 (Euler steps) from Karras et al. (2022)"
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_euler)
+
+
+class EulerAncestralSampler(KDiffusionSampler):
+    short_name = SamplerName.K_EULER_ANCESTRAL
+    name = "Euler ancestral"
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_euler_ancestral)
+
+
+class HeunSampler(KDiffusionSampler):
+    short_name = SamplerName.K_HEUN
+    name = "Algorithm 2 (Heun steps) from Karras et al. (2022)."
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_heun)
+
+
+class LMSSampler(KDiffusionSampler):
+    short_name = SamplerName.K_LMS
+    name = "LMS"
+    default_steps = 40
+    sampler_func = staticmethod(k_sampling.sample_lms)
+
+
+class CFGDenoiser(nn.Module):
+    """
+    Conditional forward guidance wrapper
+    """
+
+    def __init__(self, model):
+        super().__init__()
+        self.inner_model = model
+        self.device = get_device()
+
+    def forward(
+        self,
+        x,
+        sigma,
+        uncond,
+        cond,
+        cond_scale,
+        mask=None,
+        mask_noise=None,
+        orig_latent=None,
+    ):
+        def _wrapper(noisy_latent_in, time_encoding_in, conditioning_in):
+            return self.inner_model(
+                noisy_latent_in, time_encoding_in, cond=conditioning_in
+            )
+
+        if mask is not None:
+            assert orig_latent is not None
+            t = self.inner_model.sigma_to_t(sigma, quantize=True)
+            big_sigma = max(sigma, 1)
+            x = mask_blend(
+                noisy_latent=x,
+                orig_latent=orig_latent * big_sigma,
+                mask=mask,
+                mask_noise=mask_noise * big_sigma,
+                ts=t,
+                model=self.inner_model.inner_model,
+            )
+
+        noise_pred = get_noise_prediction(
+            denoise_func=_wrapper,
+            noisy_latent=x,
+            time_encoding=sigma,
+            neutral_conditioning=uncond,
+            positive_conditioning=cond,
+            signal_amplification=cond_scale,
+        )
+
+        return noise_pred

imaginairy/samplers/plms.py

@@ -7,13 +7,19 @@ from tqdm import tqdm
 
 from imaginairy.log_utils import increment_step, log_latent
 from imaginairy.modules.diffusion.util import extract_into_tensor, noise_like
-from imaginairy.samplers.base import NoiseSchedule, get_noise_prediction, mask_blend
+from imaginairy.samplers.base import (
+    ImageSampler,
+    NoiseSchedule,
+    SamplerName,
+    get_noise_prediction,
+    mask_blend,
+)
 from imaginairy.utils import get_device
 
 logger = logging.getLogger(__name__)
 
 
-class PLMSSampler:
+class PLMSSampler(ImageSampler):
     """
     probabilistic least-mean-squares
 
@@ -23,9 +29,9 @@ class PLMSSampler:
     https://github.com/luping-liu/PNDM
     """
 
-    def __init__(self, model):
-        self.model = model
-        self.device = get_device()
+    short_name = SamplerName.PLMS
+    name = "probabilistic least-mean-squares sampler"
+    default_steps = 40
 
     @torch.no_grad()
     def sample(

imaginairy/schema.py

@@ -11,7 +11,9 @@ from PIL import Image, ImageOps
 from urllib3.exceptions import LocationParseError
 from urllib3.util import parse_url
 
-from imaginairy.model_manager import DEFAULT_MODEL
+from imaginairy import config
+from imaginairy.model_manager import get_model_default_image_size
+from imaginairy.samplers import SAMPLER_LOOKUP
 from imaginairy.utils import get_device, get_hardware_description
 
 logger = logging.getLogger(__name__)
@@ -100,16 +102,16 @@ class ImaginePrompt:
         mask_mode=MaskMode.REPLACE,
         mask_modify_original=True,
         seed=None,
-        steps=50,
-        height=512,
-        width=512,
+        steps=None,
+        height=None,
+        width=None,
         upscale=False,
         fix_faces=False,
         fix_faces_fidelity=DEFAULT_FACE_FIDELITY,
-        sampler_type="plms",
+        sampler_type=config.DEFAULT_SAMPLER,
         conditioning=None,
         tile_mode=False,
-        model=DEFAULT_MODEL,
+        model=config.DEFAULT_MODEL,
     ):
         prompt = prompt if prompt is not None else ""
         fix_faces_fidelity = (
@@ -130,7 +132,7 @@ class ImaginePrompt:
         if mask_image is not None and mask_prompt is not None:
             raise ValueError("You can only set one of `mask_image` and `mask_prompt`")
         if model is None:
-            model = DEFAULT_MODEL
+            model = config.DEFAULT_MODEL
 
         self.init_image = init_image
         self.init_image_strength = init_image_strength
@@ -141,7 +143,7 @@ class ImaginePrompt:
         self.upscale = upscale
         self.fix_faces = fix_faces
         self.fix_faces_fidelity = fix_faces_fidelity
-        self.sampler_type = sampler_type
+        self.sampler_type = sampler_type.lower()
         self.conditioning = conditioning
         self.mask_prompt = mask_prompt
         self.mask_image = mask_image
@@ -150,6 +152,12 @@ class ImaginePrompt:
         self.tile_mode = tile_mode
         self.model = model
 
+        if self.height is None or self.width is None or self.steps is None:
+            SamplerCls = SAMPLER_LOOKUP[self.sampler_type]
+            self.steps = self.steps or SamplerCls.default_steps
+            self.width = self.width or get_model_default_image_size(self.model)
+            self.height = self.height or get_model_default_image_size(self.model)
+
     @property
     def prompt_text(self):
         if len(self.prompts) == 1:
@@ -249,7 +257,10 @@ class ImagineResult:
         exif[ExifCodes.ImageDescription] = self.prompt.prompt_description()
         exif[ExifCodes.UserComment] = json.dumps(self.metadata_dict())
         # help future web scrapes not ingest AI generated art
-        exif[ExifCodes.Software] = "Imaginairy / Stable Diffusion v1.4"
+        sd_version = self.prompt.model
+        if len(sd_version) > 20:
+            sd_version = "custom weights"
+        exif[ExifCodes.Software] = "Imaginairy / Stable Diffusion {sd_version}"
         exif[ExifCodes.DateTime] = self.created_at.isoformat(sep=" ")[:19]
         exif[ExifCodes.HostComputer] = f"{self.torch_backend}:{self.hardware_name}"
         return exif

setup.py

@@ -49,6 +49,7 @@ setup(
         "pytorch-lightning==1.4.2",
         "omegaconf==2.1.1",
         "open-clip-torch",
+        "requests",
         "einops==0.3.0",
         "timm>=0.4.12",  # for vendored blip
         "torchdiffeq",

tests/conftest.py

@@ -11,7 +11,7 @@ from urllib3 import HTTPConnectionPool
 
 from imaginairy import api
 from imaginairy.log_utils import suppress_annoying_logs_and_warnings
-from imaginairy.samplers.base import SAMPLER_TYPE_OPTIONS
+from imaginairy.samplers import SAMPLER_TYPE_OPTIONS
 from imaginairy.utils import (
     fix_torch_group_norm,
     fix_torch_nn_layer_norm,

tests/test_api.py

@@ -28,26 +28,40 @@ def test_imagine(sampler_type, filename_base_for_outputs):
     )
 
 
-@pytest.mark.skipif(get_device() == "cpu", reason="Too slow to run on CPU")
-def test_model_versions(filename_base_for_outputs):
+compare_prompts = [
+    "a photo of a bowl of fruit",
+    "a headshot photo of a happy couple smiling at the camera",
+    "a painting of a beautiful cloudy sunset at the beach",
+    "a photo of a dog",
+    "a photo of a handshake",
+    "a photo of an astronaut riding a horse on the moon. the earth visible in the background",
+]
+
+
+@pytest.mark.skipif(get_device() != "cuda", reason="Too slow to run on CPU or MPS")
+@pytest.mark.parametrize("model_version", ["SD-1.4", "SD-1.5", "SD-2.0", "SD-2.0-v"])
+def test_model_versions(filename_base_for_orig_outputs, model_version):
     """Test that we can switch between model versions"""
-    prompt_text = "a bowl of tropical fruit"
-    prompts = [
-        ImaginePrompt(
-            prompt_text, width=512, height=512, steps=20, seed=1, model="SD-1.5"
-        ),
-        ImaginePrompt(
-            prompt_text, width=512, height=512, steps=20, seed=1, model="SD-1.4"
-        ),
-    ]
+    prompts = []
+    for prompt_text in compare_prompts:
+        prompts.append(
+            ImaginePrompt(
+                prompt_text,
+                seed=1,
+                model=model_version,
+                sampler_type="ddim",
+                steps=30,
+            )
+        )
+
     threshold = 10000
 
     for i, result in enumerate(imagine(prompts)):
-        img_path = f"{filename_base_for_outputs}_{i}_{result.prompt.model}.png"
+        img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"
         result.img.save(img_path)
 
     for i, result in enumerate(imagine(prompts)):
-        img_path = f"{filename_base_for_outputs}_{i}_{result.prompt.model}.png"
+        img_path = f"{filename_base_for_orig_outputs}_{result.prompt.prompt_text}_{result.prompt.model}.png"
         assert_image_similar_to_expectation(
             result.img, img_path=img_path, threshold=threshold
         )

tests/test_enhancers.py

@@ -55,11 +55,12 @@ def test_clip_masking(filename_base_for_outputs):
         upscale=False,
         fix_faces=True,
         seed=42,
+        sampler_type="plms",
     )
 
     result = next(imagine(prompt))
     img_path = f"{filename_base_for_outputs}.png"
-    assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=300)
+    assert_image_similar_to_expectation(result.img, img_path=img_path, threshold=600)
 
 
 boolean_mask_test_cases = [