diff --git a/.gitignore b/.gitignore
index a4affa9..0a61b0f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,4 +13,5 @@ downloads
build
dist
**/*.ckpt
-**/*.egg-info
\ No newline at end of file
+**/*.egg-info
+tests/test_output
\ No newline at end of file
diff --git a/README.md b/README.md
index 0a53249..a4e138a 100644
--- a/README.md
+++ b/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"
+```
+
+
+
+
+# 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
+
\ No newline at end of file
diff --git a/assets/000019_786355545_PLMS50_PS7.5_a_scenic_landscape.jpg b/assets/000019_786355545_PLMS50_PS7.5_a_scenic_landscape.jpg
new file mode 100644
index 0000000..8987a20
Binary files /dev/null and b/assets/000019_786355545_PLMS50_PS7.5_a_scenic_landscape.jpg differ
diff --git a/assets/000032_337692011_PLMS40_PS7.5_a_photo_of_a_dog.jpg b/assets/000032_337692011_PLMS40_PS7.5_a_photo_of_a_dog.jpg
new file mode 100644
index 0000000..b8b4ca5
Binary files /dev/null and b/assets/000032_337692011_PLMS40_PS7.5_a_photo_of_a_dog.jpg differ
diff --git a/assets/000056_293284644_PLMS40_PS7.5_photo_of_a_bowl_of_fruit.jpg b/assets/000056_293284644_PLMS40_PS7.5_photo_of_a_bowl_of_fruit.jpg
new file mode 100644
index 0000000..37ed94f
Binary files /dev/null and b/assets/000056_293284644_PLMS40_PS7.5_photo_of_a_bowl_of_fruit.jpg differ
diff --git a/assets/000078_260972468_PLMS40_PS7.5_portrait_photo_of_a_freckled_woman.jpg b/assets/000078_260972468_PLMS40_PS7.5_portrait_photo_of_a_freckled_woman.jpg
new file mode 100644
index 0000000..5af579c
Binary files /dev/null and b/assets/000078_260972468_PLMS40_PS7.5_portrait_photo_of_a_freckled_woman.jpg differ
diff --git a/imaginairy/__init__.py b/imaginairy/__init__.py
index e69de29..4b2fdbb 100644
--- a/imaginairy/__init__.py
+++ b/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
diff --git a/imaginairy/imagine.py b/imaginairy/api.py
similarity index 87%
rename from imaginairy/imagine.py
rename to imaginairy/api.py
index 68b2c80..962efac 100755
--- a/imaginairy/imagine.py
+++ b/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:
diff --git a/imaginairy/cmd_wrap.py b/imaginairy/cmd_wrap.py
new file mode 100644
index 0000000..2f67e9b
--- /dev/null
+++ b/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()
diff --git a/imaginairy/cmds.py b/imaginairy/cmds.py
index 4109bce..fde4200 100644
--- a/imaginairy/cmds.py
+++ b/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,
)
diff --git a/imaginairy/models/diffusion/ddpm.py b/imaginairy/models/diffusion/ddpm.py
index db0b49e..985fab4 100644
--- a/imaginairy/models/diffusion/ddpm.py
+++ b/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
diff --git a/imaginairy/models/diffusion/plms.py b/imaginairy/models/diffusion/plms.py
index 6506c53..0ff6841 100644
--- a/imaginairy/models/diffusion/plms.py
+++ b/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):
diff --git a/imaginairy/modules/diffusionmodules/model.py b/imaginairy/modules/diffusionmodules/model.py
index 5a8a464..6fbab45 100644
--- a/imaginairy/modules/diffusionmodules/model.py
+++ b/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(
diff --git a/imaginairy/schema.py b/imaginairy/schema.py
index fb1e767..cc61f84 100644
--- a/imaginairy/schema.py
+++ b/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)
diff --git a/imaginairy/utils.py b/imaginairy/utils.py
index 068c8e0..8c084fa 100644
--- a/imaginairy/utils.py
+++ b/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):
diff --git a/setup.py b/setup.py
index b99b502..218b2cb 100644
--- a/setup.py
+++ b/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=[
diff --git a/tests/data/beach_at_sainte_adresse.jpg b/tests/data/beach_at_sainte_adresse.jpg
new file mode 100644
index 0000000..d74cd44
Binary files /dev/null and b/tests/data/beach_at_sainte_adresse.jpg differ
diff --git a/tests/data/girl_with_a_pearl_earring.jpg b/tests/data/girl_with_a_pearl_earring.jpg
new file mode 100644
index 0000000..2925183
Binary files /dev/null and b/tests/data/girl_with_a_pearl_earring.jpg differ
diff --git a/tests/test_imagine.py b/tests/test_imagine.py
index 6754229..e6ce7f5 100644
--- a/tests/test_imagine.py
+++ b/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)