feature: cleaned up logging

- cleans up all the logging. hide most of it
 - create better readme. show example images
 - save metadata into image

.gitignore

@@ -13,4 +13,5 @@ downloads
 build
 dist
 **/*.ckpt
-**/*.egg-info
+**/*.egg-info
+tests/test_output

README.md

@@ -2,18 +2,83 @@
 
 AI imagined images.
 
+```bash
+>> pip install imaginairy
 
+>> imagine "a scenic landscape" "a photo of a dog" "photo of a fruit bowl" "portrait photo of a freckled woman"
+```
+<img src="assets/000019_786355545_PLMS50_PS7.5_a_scenic_landscape.jpg" width="256" height="256">
+<img src="assets/000032_337692011_PLMS40_PS7.5_a_photo_of_a_dog.jpg" width="256" height="256">
+<img src="assets/000056_293284644_PLMS40_PS7.5_photo_of_a_bowl_of_fruit.jpg" width="256" height="256">
+<img src="assets/000078_260972468_PLMS40_PS7.5_portrait_photo_of_a_freckled_woman.jpg" width="256" height="256">
 
+# Features
+ 
+ - It makes images from text descriptions!
+ - Generate images either in code or from command line.
+ - It just works (if you have the right hardware)
+ - Noisy logs are gone (which was surprisingly hard to accomplish)
+ - WeightedPrompts let you smash together separate prompts ()
 
-# Models
+# How To
 
+```python
+from imaginairy import imagine_images, imagine_image_files, ImaginePrompt, WeightedPrompt
+
+prompts = [
+    ImaginePrompt("a scenic landscape", seed=1),
+    ImaginePrompt("a bowl of fruit"),
+    ImaginePrompt([
+       WeightedPrompt("cat", weight=1),
+       WeightedPrompt("dog", weight=1),
+    ])
+]
+for result in imagine_images(prompts):
+    # do something
+    result.save("my_image.jpg")
+    
+# or
+
+imagine_image_files(prompts, outdir="./my-art")
+
+```
+
+# Requirements
+
+- Computer with CUDA supported graphics card. ~10 gb video ram
+OR
+- Apple M1 computer
+
+# Improvements from CompVis
+ - img2img actually does # of steps you specify
+
+# Models Used
+ - CLIP
  - LDM - Latent Diffusion
- - Stable Diffusion
+ - Stable Diffusion - https://github.com/CompVis/stable-diffusion
 
 # Todo
- - add tests
+ - add safety model - https://github.com/CompVis/stable-diffusion/blob/main/scripts/txt2img.py#L21-L28
  - add docs
- - remove yaml config
  - deploy to pypi
- - add image describe feature
+ - add tests
+ - set up ci (test/lint/format)
+ - remove yaml config
+ - performance optimizations https://github.com/huggingface/diffusers/blob/main/docs/source/optimization/fp16.mdx 
+ - Interface improvements
+   - init-image at command line
+   - prompt expansion?
+   - webserver interface (low priority, this is a library)
+ - Image Generation Features
+   - image describe feature
+   - outpainting
+   - inpainting
+   - face improvements
+   - upscaling
+   - cross-attention control: 
+     - https://github.com/bloc97/CrossAttentionControl/blob/main/CrossAttention_Release_NoImages.ipynb
+   - tiling
+   - output show-work videos
+   - zooming videos? a la disco diffusion
+
  

imaginairy/__init__.py

@@ -0,0 +1,6 @@
+import os
+
+os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+from api import imagine_images, imagine_image_files
+from schema import ImaginePrompt, ImagineResult, WeightedPrompt

imaginairy/api.py

@@ -37,20 +37,20 @@ logger = logging.getLogger(__name__)
 
 
 def load_model_from_config(config):
-    ckpt_path = cached_path(
-        "https://www.googleapis.com/storage/v1/b/aai-blog-files/o/sd-v1-4.ckpt?alt=media"
-    )
-    logger.info(f"Loading model from {ckpt_path}")
+    url = "https://www.googleapis.com/storage/v1/b/aai-blog-files/o/sd-v1-4.ckpt?alt=media"
+    ckpt_path = cached_path(url)
+    logger.info(f"Loading model onto {get_device()} backend...")
+    logger.debug(f"Loading model from {ckpt_path}")
     pl_sd = torch.load(ckpt_path, map_location="cpu")
     if "global_step" in pl_sd:
-        logger.info(f"Global Step: {pl_sd['global_step']}")
+        logger.debug(f"Global Step: {pl_sd['global_step']}")
     sd = pl_sd["state_dict"]
     model = instantiate_from_config(config.model)
     m, u = model.load_state_dict(sd, strict=False)
     if len(m) > 0:
-        logger.info(f"missing keys: {m}")
+        logger.debug(f"missing keys: {m}")
     if len(u) > 0:
-        logger.info(f"unexpected keys: {u}")
+        logger.debug(f"unexpected keys: {u}")
 
     model.to(get_device())
     model.eval()
@@ -88,13 +88,17 @@ def imagine_image_files(
     downsampling_factor=8,
     precision="autocast",
     ddim_eta=0.0,
-    record_steps=False,
+    record_step_images=False,
+    output_file_extension="jpg",
 ):
     big_path = os.path.join(outdir, "upscaled")
     os.makedirs(outdir, exist_ok=True)
     os.makedirs(big_path, exist_ok=True)
     base_count = len(os.listdir(outdir))
     step_count = 0
+    output_file_extension = output_file_extension.lower()
+    if output_file_extension not in {"jpg", "png"}:
+        raise ValueError("Must output a png or jpg")
 
     def _record_steps(samples, i, model, prompt):
         nonlocal step_count
@@ -110,7 +114,7 @@ def imagine_image_files(
                 os.path.join(steps_path, filename)
             )
 
-    img_callback = _record_steps if record_steps else None
+    img_callback = _record_steps if record_step_images else None
     for result in imagine_images(
         prompts,
         latent_channels=latent_channels,
@@ -120,16 +124,15 @@ def imagine_image_files(
         img_callback=img_callback,
     ):
         prompt = result.prompt
-        img = result.img
         basefilename = f"{base_count:06}_{prompt.seed}_{prompt.sampler_type}{prompt.steps}_PS{prompt.prompt_strength}_{prompt_normalized(prompt.prompt_text)}"
         filepath = os.path.join(outdir, f"{basefilename}.jpg")
 
-        img.save(filepath)
+        result.save(filepath)
+        logger.info(f"    🖼  saved to: {filepath}")
         if prompt.upscale:
-            enlarge_realesrgan2x(
-                filepath,
-                os.path.join(big_path, basefilename) + ".jpg",
-            )
+            bigfilepath = (os.path.join(big_path, basefilename) + ".jpg",)
+            enlarge_realesrgan2x(filepath, bigfilepath)
+            logger.info(f"    upscaled 🖼  saved to: {filepath}")
         base_count += 1
 
 
@@ -146,9 +149,14 @@ def imagine_images(
     prompts = [prompts] if isinstance(prompts, ImaginePrompt) else prompts
     _img_callback = None
 
-    precision_scope = autocast if precision == "autocast" else nullcontext
-    with (torch.no_grad(), precision_scope("cuda"), fix_torch_nn_layer_norm()):
+    precision_scope = (
+        autocast
+        if precision == "autocast" and get_device() in ("cuda", "cpu")
+        else nullcontext
+    )
+    with (torch.no_grad(), precision_scope(get_device()), fix_torch_nn_layer_norm()):
         for prompt in prompts:
+            logger.info(f"Generating {prompt.prompt_description()}")
             seed_everything(prompt.seed)
             uc = None
             if prompt.prompt_strength != 1.0:

imaginairy/cmd_wrap.py

@@ -0,0 +1,64 @@
+# only builtin imports allowed at this point since we want to modify
+# the environment and code before it's loaded
+import importlib.abc
+import importlib.util
+import logging.config
+import os
+import site
+import sys
+import warnings
+
+# tells pytorch to allow MPS usage (for Mac M1 compatibility)
+os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+
+def disable_transformers_logging():
+    """
+    Disable `transformers` package custom logging.
+
+    I can't believe it came to this. I tried like four other approaches first
+
+    Loads up the source code from the transformers file and turns it into a module.
+    We then modify the module.  Every other approach (import hooks, custom import function)
+    loaded the module before it could be modified.
+    """
+    t_logging_path = f"{site.getsitepackages()[0]}/transformers/utils/logging.py"
+    with open(t_logging_path, "r", encoding="utf-8") as f:
+        src_code = f.read()
+
+    spec = importlib.util.spec_from_loader("transformers.utils.logging", loader=None)
+    module = importlib.util.module_from_spec(spec)
+
+    exec(src_code, module.__dict__)
+    module.get_logger = logging.getLogger
+    sys.modules["transformers.utils.logging"] = module
+
+
+def disable_pytorch_lighting_custom_logging():
+    from pytorch_lightning import _logger
+
+    _logger.setLevel(logging.NOTSET)
+
+
+def filter_torch_warnings():
+    warnings.filterwarnings(
+        "ignore",
+        category=UserWarning,
+        message=r"The operator .*?is not currently supported.*",
+    )
+
+
+def setup_env():
+    disable_transformers_logging()
+    disable_pytorch_lighting_custom_logging()
+    filter_torch_warnings()
+
+
+setup_env()
+
+
+from imaginairy.cmds import imagine_cmd  # noqa
+
+# imagine_cmd = disable_transformers_logging_mess()(imagine_cmd)
+if __name__ == "__main__":
+    imagine_cmd()

imaginairy/cmds.py

@@ -1,19 +1,56 @@
-#!/usr/bin/env python
-import os
-
-os.putenv("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+import logging.config
 
 import click
+from imaginairy.imagine import load_model
+
+logger = logging.getLogger(__name__)
+
+
+def configure_logging(level="INFO"):
+    fmt = "%(message)s"
+    if level == "DEBUG":
+        fmt = "%(asctime)s [%(levelname)s] %(name)s:%(lineno)d: %(message)s"
 
-from imaginairy.imagine import imagine_image_files
-from imaginairy.schema import ImaginePrompt
+    LOGGING_CONFIG = {
+        "version": 1,
+        "disable_existing_loggers": True,
+        "formatters": {
+            "standard": {"format": fmt},
+        },
+        "handlers": {
+            "default": {
+                "level": "INFO",
+                "formatter": "standard",
+                "class": "logging.StreamHandler",
+                "stream": "ext://sys.stdout",  # Default is stderr
+            },
+        },
+        "loggers": {
+            "": {  # root logger
+                "handlers": ["default"],
+                "level": "WARNING",
+                "propagate": False,
+            },
+            "imaginairy": {"handlers": ["default"], "level": level, "propagate": False},
+            "transformers.modeling_utils": {
+                "handlers": ["default"],
+                "level": "ERROR",
+                "propagate": False,
+            },
+        },
+    }
+    logging.config.dictConfig(LOGGING_CONFIG)
 
 
 @click.command()
-@click.argument("prompt_texts", default=None, nargs=-1)
+@click.argument("prompt_texts", nargs=-1)
 @click.option("--outdir", default="./outputs", help="where to write results to")
 @click.option(
-    "-r", "--repeats", default=1, type=int, help="How many times to repeat the renders"
+    "-r",
+    "--repeats",
+    default=1,
+    type=int,
+    help="How many times to repeat the renders. If you provide two prompts and --repeat=3 then six images will be generated",
 )
 @click.option(
     "-h",
@@ -27,8 +64,9 @@ from imaginairy.schema import ImaginePrompt
 )
 @click.option(
     "--steps",
-    default=50,
+    default=40,
     type=int,
+    show_default=True,
     help="How many diffusion steps to run. More steps, more detail, but with diminishing returns",
 )
 @click.option(
@@ -40,10 +78,29 @@ from imaginairy.schema import ImaginePrompt
 @click.option(
     "--prompt-strength",
     default=7.5,
+    show_default=True,
     help="How closely to follow the prompt. Image looks unnatural at higher values",
 )
-@click.option("--sampler-type", default="PLMS", help="What sampling strategy to use")
+@click.option(
+    "--sampler-type",
+    default="PLMS",
+    type=click.Choice(["PLMS", "DDIM"]),
+    help="What sampling strategy to use",
+)
 @click.option("--ddim-eta", default=0.0, type=float)
+@click.option(
+    "--log-level",
+    default="INFO",
+    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR"]),
+    help="What level of logs to show.",
+)
+@click.option(
+    "--show-work",
+    default=["none"],
+    type=click.Choice(["none", "images", "video"]),
+    multiple=True,
+    help="Make a video showing the image being created",
+)
 def imagine_cmd(
     prompt_texts,
     outdir,
@@ -55,9 +112,21 @@ def imagine_cmd(
     prompt_strength,
     sampler_type,
     ddim_eta,
+    log_level,
+    show_work,
 ):
     """Render an image"""
+    configure_logging(log_level)
+    from imaginairy.imagine import imagine_image_files
+    from imaginairy.schema import ImaginePrompt
+
+    total_image_count = len(prompt_texts) * repeats
+    logger.info(
+        f"🤖🧠 received {len(prompt_texts)} prompt(s) and will repeat them {repeats} times to create {total_image_count} images."
+    )
+
     prompts = []
+    load_model()
     for _ in range(repeats):
         for prompt_text in prompt_texts:
             prompt = ImaginePrompt(
@@ -77,6 +146,7 @@ def imagine_cmd(
         prompts,
         outdir=outdir,
         ddim_eta=ddim_eta,
+        record_step_images="images" in show_work,
     )
 
 

imaginairy/models/diffusion/ddpm.py

@@ -77,7 +77,7 @@ class DDPM(pl.LightningModule):
             "x0",
         ], 'currently only supporting "eps" and "x0"'
         self.parameterization = parameterization
-        logger.info(
+        logger.debug(
             f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode"
         )
         self.cond_stage_model = None
@@ -309,10 +309,10 @@ class LatentDiffusion(DDPM):
     def instantiate_cond_stage(self, config):
         if not self.cond_stage_trainable:
             if config == "__is_first_stage__":
-                logger.info("Using first stage also as cond stage.")
+                logger.debug("Using first stage also as cond stage.")
                 self.cond_stage_model = self.first_stage_model
             elif config == "__is_unconditional__":
-                logger.info(
+                logger.debug(
                     f"Training {self.__class__.__name__} as an unconditional model."
                 )
                 self.cond_stage_model = None

imaginairy/models/diffusion/plms.py

@@ -129,7 +129,7 @@ class PLMSSampler(object):
         # sampling
         C, H, W = shape
         size = (batch_size, C, H, W)
-        logger.info(f"Data shape for PLMS sampling is {size}")
+        logger.debug(f"Data shape for PLMS sampling is {size}")
 
         samples, intermediates = self.plms_sampling(
             conditioning,
@@ -202,9 +202,9 @@ class PLMSSampler(object):
             else np.flip(timesteps)
         )
         total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
-        logger.info(f"Running PLMS Sampling with {total_steps} timesteps")
+        logger.debug(f"Running PLMS Sampling with {total_steps} timesteps")
 
-        iterator = tqdm(time_range, desc="PLMS Sampler", total=total_steps)
+        iterator = tqdm(time_range, desc="    PLMS Sampler", total=total_steps)
         old_eps = []
 
         for i, step in enumerate(iterator):

imaginairy/modules/diffusionmodules/model.py

@@ -196,7 +196,7 @@ class AttnBlock(nn.Module):
 
 def make_attn(in_channels, attn_type="vanilla"):
     assert attn_type in ["vanilla", "linear", "none"], f"attn_type {attn_type} unknown"
-    logger.info(
+    logger.debug(
         f"making attention of type '{attn_type}' with {in_channels} in_channels"
     )
     if attn_type == "vanilla":
@@ -361,7 +361,7 @@ class Decoder(nn.Module):
         block_in = ch * ch_mult[self.num_resolutions - 1]
         curr_res = resolution // 2 ** (self.num_resolutions - 1)
         self.z_shape = (1, z_channels, curr_res, curr_res)
-        logger.info(
+        logger.debug(
             f"Working with z of shape {self.z_shape} = {np.prod(self.z_shape)} dimensions."
         )
 
@@ -516,7 +516,7 @@ class Upsampler(nn.Module):
         assert out_size >= in_size
         num_blocks = int(np.log2(out_size // in_size)) + 1
         factor_up = 1.0 + (out_size % in_size)
-        logger.info(
+        logger.debug(
             f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}"
         )
         self.rescaler = LatentRescaler(

imaginairy/schema.py

@@ -1,7 +1,12 @@
 import hashlib
+import json
 import random
+from datetime import datetime, timezone
 
 import numpy
+from PIL.Image import Exif
+
+from imaginairy.utils import get_device, get_device_name
 
 
 class WeightedPrompt:
@@ -17,35 +22,33 @@ class ImaginePrompt:
     def __init__(
         self,
         prompt=None,
-        seed=None,
         prompt_strength=7.5,
-        sampler_type="PLMS",
         init_image=None,
         init_image_strength=0.3,
+        seed=None,
         steps=50,
         height=512,
         width=512,
         upscale=False,
         fix_faces=False,
-        parts=None,
+        sampler_type="PLMS",
     ):
         prompt = prompt if prompt is not None else "a scenic landscape"
         if isinstance(prompt, str):
             self.prompts = [WeightedPrompt(prompt, 1)]
         else:
             self.prompts = prompt
+        self.prompts.sort(key=lambda p: p.weight, reverse=True)
+        self.prompt_strength = prompt_strength
         self.init_image = init_image
         self.init_image_strength = init_image_strength
-        self.prompts.sort(key=lambda p: p.weight, reverse=True)
         self.seed = random.randint(1, 1_000_000_000) if seed is None else seed
-        self.prompt_strength = prompt_strength
-        self.sampler_type = sampler_type
         self.steps = steps
         self.height = height
         self.width = width
         self.upscale = upscale
         self.fix_faces = fix_faces
-        self.parts = parts or {}
+        self.sampler_type = sampler_type
 
     @property
     def prompt_text(self):
@@ -53,11 +56,46 @@ class ImaginePrompt:
             return self.prompts[0].text
         return "|".join(str(p) for p in self.prompts)
 
+    def prompt_description(self):
+        return (
+            f'🖼  : "{self.prompt_text}" {self.width}x{self.height}px '
+            f"seed:{self.seed} prompt-strength:{self.prompt_strength} steps:{self.steps} sampler-type:{self.sampler_type}"
+        )
+
+    def as_dict(self):
+        prompts = [(p.weight, p.text) for p in self.prompts]
+        return {
+            "software": "imaginairy",
+            "prompts": prompts,
+            "prompt_strength": self.prompt_strength,
+            "init_image": self.init_image,
+            "init_image_strength": self.init_image_strength,
+            "seed": self.seed,
+            "steps": self.steps,
+            "height": self.height,
+            "width": self.width,
+            "upscale": self.upscale,
+            "fix_faces": self.fix_faces,
+            "sampler_type": self.sampler_type,
+        }
+
+
+class ExifCodes:
+    """https://www.awaresystems.be/imaging/tiff/tifftags/baseline.html"""
+    ImageDescription = 0x010E
+    Software = 0x0131
+    DateTime = 0x0132
+    HostComputer = 0x013C
+    UserComment = 0x9286
+
 
 class ImagineResult:
-    def __init__(self, img, prompt):
+    def __init__(self, img, prompt: ImaginePrompt):
         self.img = img
         self.prompt = prompt
+        self.created_at = datetime.utcnow().replace(tzinfo=timezone.utc)
+        self.torch_backend = get_device()
+        self.hardware_name = get_device_name(get_device())
 
     def cv2_img(self):
         open_cv_image = numpy.array(self.img)
@@ -68,3 +106,18 @@ class ImagineResult:
 
     def md5(self):
         return hashlib.md5(self.img.tobytes()).hexdigest()
+
+    def metadata_dict(self):
+        return {
+            "prompt": self.prompt.as_dict(),
+        }
+
+    def save(self, save_path):
+        exif = Exif()
+        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"
+        exif[ExifCodes.DateTime] = self.created_at.isoformat(sep=" ")[:19]
+        exif[ExifCodes.HostComputer] = f"{self.torch_backend}:{self.hardware_name}"
+        self.img.save(save_path, exif=exif)

imaginairy/utils.py

@@ -1,6 +1,6 @@
-import os
 import importlib
 import logging
+import platform
 from contextlib import contextmanager
 from functools import lru_cache
 from typing import List, Optional
@@ -21,9 +21,16 @@ def get_device():
         return "cpu"
 
 
+@lru_cache()
+def get_device_name(device_type):
+    if device_type == "cuda":
+        return torch.cuda.get_device_name(0)
+    return platform.processor()
+
+
 def log_params(model):
     total_params = sum(p.numel() for p in model.parameters())
-    logger.info(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
+    logger.debug(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
 
 
 def instantiate_from_config(config):

setup.py

@@ -6,7 +6,7 @@ setup(
     description="AI imagined images.",
     packages=find_packages(include=("imaginairy", "imaginairy.*")),
     entry_points={
-        "console_scripts": ["imagine=imaginairy.cmds:imagine_cmd"],
+        "console_scripts": ["imagine=imaginairy.cmd_wrap:imagine_cmd"],
     },
     package_data={"imaginairy": ["configs/*.yaml"]},
     install_requires=[

tests/test_imagine.py

@@ -1,5 +1,5 @@
-from imaginairy.imagine import imagine_images, imagine_image_files
-from imaginairy.schema import ImaginePrompt, WeightedPrompt
+from imaginairy.api import imagine_images, imagine_image_files
+from imaginairy.schema import ImaginePrompt
 from . import TESTS_FOLDER
 
 
@@ -24,43 +24,18 @@ def test_img_to_img():
         sampler_type="DDIM",
     )
     out_folder = f"{TESTS_FOLDER}/test_output"
-    out_folder = "/home/bryce/Mounts/drennanfiles/art/tests"
     imagine_image_files(prompt, outdir=out_folder)
 
 
 def test_img_to_file():
     prompt = ImaginePrompt(
-        [
-            WeightedPrompt(
-                "an old growth forest, diffuse light poking through the canopy. high-resolution, nature photography, nat geo photo"
-            )
-        ],
-        # init_image=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg",
-        init_image_strength=0.5,
+        "an old growth forest, diffuse light poking through the canopy. high-resolution, nature photography, nat geo photo",
         width=512 + 64,
         height=512 - 64,
         steps=50,
-        # seed=2,
+        seed=2,
         sampler_type="PLMS",
         upscale=True,
     )
     out_folder = f"{TESTS_FOLDER}/test_output"
-    out_folder = "/home/bryce/Mounts/drennanfiles/art/tests"
     imagine_image_files(prompt, outdir=out_folder)
-
-
-def test_img_conditioning():
-    prompt = ImaginePrompt(
-        "photo",
-        init_image=f"{TESTS_FOLDER}/data/beach_at_sainte_adresse.jpg",
-        init_image_strength=0.5,
-        width=512 + 64,
-        height=512 - 64,
-        steps=50,
-        # seed=2,
-        sampler_type="PLMS",
-        upscale=True,
-    )
-    out_folder = f"{TESTS_FOLDER}/test_output"
-    out_folder = "/home/bryce/Mounts/drennanfiles/art/tests"
-    imagine_image_files(prompt, outdir=out_folder, record_steps=True)