Add CLI script

README.md

@@ -107,7 +107,7 @@ video_path = pipeline.walk(
 )
 ```
 
-#### Run the App Locally
+### Run the App Locally
 
 ```python
 from stable_diffusion_videos import StableDiffusionWalkPipeline, Interface
@@ -123,6 +123,13 @@ interface = Interface(pipeline)
 interface.launch()
 ```
 
+### CLI
+
+The script `scripts/make_video.py` also provides a CLI. Example:
+```bash
+python scripts/make_video.py --prompts "a cat" "a dog" --fps 10
+```
+
 ## Credits
 
 This work built off of [a script](https://gist.github.com/karpathy/00103b0037c5aaea32fe1da1af553355

scripts/make_video.py

@@ -0,0 +1,154 @@
+import argparse
+import random
+
+import torch
+import yaml
+
+from diffusers import DPMSolverMultistepScheduler
+from stable_diffusion_videos import StableDiffusionWalkPipeline
+
+
+def init_arg_parser():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument('--checkpoint_id',
+                        default="stabilityai/stable-diffusion-2-1",
+                        help="checkpoint id on huggingface")
+    parser.add_argument('--prompts', nargs='+',
+                        help='sequence of prompts')
+    parser.add_argument('--seeds', type=int, nargs='+',
+                        help='seed for each prompt')
+    parser.add_argument('--num_interpolation_steps', type=int, nargs='+',
+                        help='number of steps between each image')
+    parser.add_argument('--output_dir', default="dreams",
+                        help='output directory')
+    parser.add_argument('--name',
+                        help='output sub-directory')
+    parser.add_argument('--fps', type=int, default=10,
+                        help='frames per second')
+    parser.add_argument('--guidance_scale', type=float, default=7.5,
+                        help='diffusion guidance scale')
+    parser.add_argument('--num_inference_steps', type=int, default=50,
+                        help='number of diffusion inference steps')
+    parser.add_argument('--height', type=int, default=512,
+                        help='output image height')
+    parser.add_argument('--width', type=int, default=512,
+                        help='output image width')
+    parser.add_argument('--upsample', action='store_true',
+                        help='upscale x4 using Real-ESRGAN')
+    parser.add_argument('--batch_size', type=int, default=1,
+                        help='batch size')
+    parser.add_argument('--audio_filepath',
+                        help='path to audio file')
+    parser.add_argument('--audio_offsets', type=int, nargs='+',
+                        help='audio offset for each prompt')
+    parser.add_argument('--negative_prompt',
+                        help='negative prompt (one for all images)')
+
+    parser.add_argument('--cfg',
+                        help='yaml config file (overwrites other options)')
+
+    return parser
+
+
+def parse_args(parser):
+    args = parser.parse_args()
+
+    # read config file
+    if args.cfg is not None:
+        with open(args.cfg) as f:
+            cfg = yaml.safe_load(f)
+        for key, val in cfg.items():
+            if hasattr(args, key):
+                setattr(args, key, val)
+            else:
+                raise ValueError(f'bad field in config file: {key}')
+
+    # check for prompts
+    if args.prompts is None:
+        raise ValueError('no prompt provided')
+    if args.seeds is None:
+        args.seeds = [random.getrandbits(16) for _ in args.prompts]
+
+    # check audio arguments
+    if args.audio_filepath is not None and args.audio_offsets is None:
+        raise ValueError('must provide audio_offsets when providing '
+                         'audio_filepath')
+    if args.audio_offsets is not None and args.audio_filepath is None:
+        raise ValueError('must provide audio_filepath when providing '
+                         'audio_offsets')
+
+    # check lengths
+    if args.audio_offsets is not None:
+        if not len(args.prompts) == len(args.seeds) == len(args.audio_offsets):
+            raise ValueError('prompts, seeds and audio_offsets must have same '
+                             f'length, got lengths {len(args.prompts)}, '
+                             f'{len(args.seeds)} and '
+                             f'{len(args.audio_offsets)} respectively')
+    else:
+        if not len(args.prompts) == len(args.seeds):
+            raise ValueError('prompts and seeds must have same length, got '
+                             f'lengths {len(args.prompts)} and '
+                             f'{len(args.seeds)} respectively')
+
+    # set num_interpolation_steps
+    if args.audio_offsets is not None \
+            and args.num_interpolation_steps is not None:
+        raise ValueError('cannot provide both audio_offsets and '
+                         'num_interpolation_steps')
+    elif args.audio_offsets is not None:
+        args.num_interpolation_steps = [
+            (b-a)*args.fps for a, b in zip(
+                args.audio_offsets, args.audio_offsets[1:]
+            )
+        ]
+    elif args.num_interpolation_steps is not None \
+            and not len(args.num_interpolation_steps) == len(args.prompts)-1:
+        raise ValueError('num_interpolation_steps must have length '
+                         f'len(prompts)-1, got '
+                         f'{len(args.num_interpolation_steps)} != '
+                         f'{len(args.prompts)-1}')
+    else:
+        args.num_interpolation_steps = args.fps*10  # 10 second video
+
+    return args
+
+
+def main():
+    parser = init_arg_parser()
+    args = parse_args(parser)
+
+    pipe = StableDiffusionWalkPipeline.from_pretrained(
+        args.checkpoint_id,
+        torch_dtype=torch.float16,
+        revision="fp16",
+        feature_extractor=None,
+        safety_checker=None,
+    ).to("cuda")
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+        pipe.scheduler.config
+    )
+
+    pipe.walk(
+        prompts=args.prompts,
+        seeds=args.seeds,
+        num_interpolation_steps=args.num_interpolation_steps,
+        output_dir=args.output_dir,
+        name=args.name,
+        fps=args.fps,
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=args.height,
+        width=args.width,
+        upsample=args.upsample,
+        batch_size=args.batch_size,
+        audio_filepath=args.audio_filepath,
+        audio_start_sec=None if args.audio_offsets is None else args.audio_offsets[0],
+        negative_prompt=args.negative_prompt,
+    )
+
+
+if __name__ == '__main__':
+    main()