refactor: move commands to entrypoints

imaginairy/cmds.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+
+import click
+
+from imaginairy.imagine import ImaginePrompt, imagine as imagine_f
+
+
+@click.command()
+@click.argument(
+    "prompt_texts", default=None, help="text to render to an image", 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")
+@click.option(
+    "-h",
+    "--height",
+    default=512,
+    type=int,
+    help="image height. should be multiple of 64",
+)
+@click.option(
+    "-w", "--width", default=512, type=int, help="image width. should be multiple of 64"
+)
+@click.option(
+    "--steps",
+    default=50,
+    type=int,
+    help="How many diffusion steps to run. More steps, more detail, but with diminishing returns",
+)
+@click.option(
+    "--seed",
+    default=None,
+    type=int,
+    help="What seed to use for randomness. Allows reproducible image renders",
+)
+@click.option(
+    "--prompt-strength",
+    default=7.5,
+    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("--ddim-eta", default=0.0, type=float)
+def imagine_cmd(
+    prompt_texts,
+    outdir,
+    repeats,
+    height,
+    width,
+    steps,
+    seed,
+    prompt_strength,
+    sampler_type,
+    ddim_eta,
+):
+    prompts = []
+    for _ in range(repeats):
+        for prompt_text in prompt_texts:
+            prompt = ImaginePrompt(
+                prompt_text,
+                seed=seed,
+                sampler_type=sampler_type,
+                steps=steps,
+                height=height,
+                width=width,
+                prompt_strength=prompt_strength,
+                upscale=True,
+                fix_faces=True,
+            )
+            prompts.append(prompt)
+
+    imagine_f(
+        prompts,
+        outdir=outdir,
+        ddim_eta=ddim_eta,
+    )
+
+
+if __name__ == "__main__":
+    imagine_cmd()

imaginairy/imagine.py

@@ -42,134 +42,6 @@ def load_model_from_config(config, ckpt, verbose=False):
     return model
 
 
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--prompt", type=str, nargs="?", default=None, help="the prompt to render"
-    )
-    parser.add_argument(
-        "--outdir",
-        type=str,
-        nargs="?",
-        help="dir to write results to",
-        default="outputs/txt2img-samples",
-    )
-    parser.add_argument(
-        "--skip_save",
-        action="store_true",
-        help="do not save individual samples. For speed measurements.",
-    )
-    parser.add_argument(
-        "--steps",
-        type=int,
-        default=50,
-        help="number of sampling steps",
-    )
-    parser.add_argument(
-        "--plms", action="store_true", help="use plms sampling", default=True
-    )
-    parser.add_argument(
-        "--fixed_code",
-        action="store_true",
-        help="if enabled, uses the same starting code across samples ",
-    )
-    parser.add_argument(
-        "--ddim_eta",
-        type=float,
-        default=0.0,
-        help="ddim eta (eta=0.0 corresponds to deterministic sampling",
-    )
-    parser.add_argument(
-        "--n",
-        type=int,
-        default=1,
-        help="how many images",
-    )
-    parser.add_argument(
-        "--H",
-        type=int,
-        default=512,
-        help="image height, in pixel space",
-    )
-    parser.add_argument(
-        "--W",
-        type=int,
-        default=512,
-        help="image width, in pixel space",
-    )
-    parser.add_argument(
-        "--C",
-        type=int,
-        default=4,
-        help="latent channels",
-    )
-    parser.add_argument(
-        "--f",
-        type=int,
-        default=8,
-        help="downsampling factor",
-    )
-    parser.add_argument(
-        "--scale",
-        type=float,
-        default=7.5,
-        help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
-    )
-    parser.add_argument(
-        "--config",
-        type=str,
-        default="configs/stable-diffusion/v1-inference.yaml",
-        help="path to config which constructs model",
-    )
-    parser.add_argument(
-        "--ckpt",
-        type=str,
-        default="data/sd-v1-4.ckpt",
-        help="path to checkpoint of model",
-    )
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=None,
-        help="the seed (for reproducible sampling)",
-    )
-    parser.add_argument(
-        "--precision",
-        type=str,
-        help="evaluate at this precision",
-        choices=["full", "autocast"],
-        default="autocast",
-    )
-    opt = parser.parse_args()
-    if opt.plms:
-        sampler = "PLMS"
-    else:
-        sampler = "DDIM"
-    prompt = ImaginePrompt(
-        opt.prompt,
-        seed=opt.seed,
-        sampler_type=sampler,
-        steps=opt.steps,
-        height=opt.H,
-        width=opt.W,
-        upscale=True,
-        fix_faces=True,
-    )
-    imagine(
-        [prompt],
-        config=opt.config,
-        ckpt=opt.ckpt,
-        outdir=opt.outdir,
-        fixed_code=opt.fixed_code,
-        latent_channels=opt.C,
-        precision=opt.precision,
-        downsampling_factor=opt.f,
-        skip_save=opt.skip_save,
-        ddim_eta=opt.ddim_eta,
-    )
-
-
 class WeightedPrompt:
     def __init__(self, text, weight=1):
         self.text = text

setup.py

@@ -5,16 +5,19 @@ setup(
     version='0.0.1',
     description='AI imagined images.',
     packages=find_packages("imaginairy"),
+    entry_points={
+        'console_scripts': ['imagine=imaginairy.cmds:imagine_cmd'],
+    },
     install_requires=[
         'click',
         'torch',
         'numpy',
         'tqdm',
-        "albumentations==0.4.3",
+        # "albumentations==0.4.3",
         "diffusers",
         # opencv-python==4.1.2.30
         "pudb==2019.2",
-        "invisible-watermark",
+        # "invisible-watermark",
         "imageio==2.9.0",
         "imageio-ffmpeg==0.4.2",
         "pytorch-lightning==1.4.2",