|
| 1 | +""" |
| 2 | +generate.py - the Image-Generator without the questions, for other programs |
| 3 | +(DREAM paints her dreams with it). Same SDXL-Turbo pipelines as |
| 4 | +image_generator.py, driven by arguments instead of input(). |
| 5 | +
|
| 6 | + python scripts/generate.py --prompt "a red fox in snow" --out fox.jpg |
| 7 | + python scripts/generate.py --prompt "as a watercolour" --init in.jpg --out out.jpg |
| 8 | + python scripts/generate.py --job job.json # many images, one model load |
| 9 | + python scripts/generate.py --download # just fetch the model |
| 10 | +
|
| 11 | +A job file: |
| 12 | + {"prompt": "...", "keyframe": "key.jpg", text-to-image (optional) |
| 13 | + "frames": [["in1.jpg", "out1.jpg"], ...], image-to-image, each in -> out |
| 14 | + "strength": 0.5, "size": 512, "seed": 1234} |
| 15 | +
|
| 16 | +Each image is written as soon as it's done, so a job that gets killed part |
| 17 | +way still leaves what it finished. --model picks another turbo model, e.g. |
| 18 | +stabilityai/sd-turbo (smaller: fits a 4 GB GPU and 8 GB of RAM). |
| 19 | +""" |
| 20 | + |
| 21 | +import argparse |
| 22 | +import json |
| 23 | +import os |
| 24 | +import sys |
| 25 | + |
| 26 | +DEFAULT_MODEL = "stabilityai/sdxl-turbo" |
| 27 | + |
| 28 | + |
| 29 | +def load_pipelines(model, device): |
| 30 | + import torch |
| 31 | + from diffusers import AutoPipelineForImage2Image, AutoPipelineForText2Image |
| 32 | + |
| 33 | + if device == "auto": |
| 34 | + device = "cuda" if torch.cuda.is_available() else "cpu" |
| 35 | + dtype = torch.float16 if device == "cuda" else torch.float32 |
| 36 | + t2i = AutoPipelineForText2Image.from_pretrained(model, torch_dtype=dtype, variant="fp16") |
| 37 | + if device == "cuda": |
| 38 | + t2i.enable_model_cpu_offload() # only the part that's working sits on the GPU |
| 39 | + t2i.enable_attention_slicing() |
| 40 | + else: |
| 41 | + t2i.to("cpu") |
| 42 | + t2i.set_progress_bar_config(disable=True) |
| 43 | + i2i = AutoPipelineForImage2Image.from_pipe(t2i) |
| 44 | + return t2i, i2i |
| 45 | + |
| 46 | + |
| 47 | +def generator_for(seed): |
| 48 | + import torch |
| 49 | + return None if seed is None else torch.Generator("cpu").manual_seed(int(seed)) |
| 50 | + |
| 51 | + |
| 52 | +def text_to_image(t2i, prompt, size, seed, steps=1): |
| 53 | + return t2i(prompt=prompt, num_inference_steps=steps, guidance_scale=0.0, |
| 54 | + width=size, height=size, generator=generator_for(seed)).images[0] |
| 55 | + |
| 56 | + |
| 57 | +def image_to_image(i2i, prompt, init_path, size, strength, seed, steps=2): |
| 58 | + from diffusers.utils import load_image |
| 59 | + init = load_image(init_path).convert("RGB").resize((size, size)) |
| 60 | + # turbo models need steps * strength >= 1 |
| 61 | + steps = max(steps, int(1 / max(strength, 0.05)) + 1) |
| 62 | + return i2i(prompt, image=init, num_inference_steps=steps, strength=strength, |
| 63 | + guidance_scale=0.0, generator=generator_for(seed)).images[0] |
| 64 | + |
| 65 | + |
| 66 | +def save(img, path): |
| 67 | + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) |
| 68 | + img.save(path, quality=92) |
| 69 | + print(f"wrote {path}", flush=True) |
| 70 | + |
| 71 | + |
| 72 | +def main(): |
| 73 | + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| 74 | + ap.add_argument("--prompt") |
| 75 | + ap.add_argument("--out") |
| 76 | + ap.add_argument("--init", help="image-to-image: the picture to start from") |
| 77 | + ap.add_argument("--strength", type=float, default=0.5) |
| 78 | + ap.add_argument("--size", type=int, default=512) |
| 79 | + ap.add_argument("--seed", type=int) |
| 80 | + ap.add_argument("--job", help="JSON job file (see above)") |
| 81 | + ap.add_argument("--model", default=DEFAULT_MODEL) |
| 82 | + ap.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"]) |
| 83 | + ap.add_argument("--download", action="store_true", help="fetch the model and exit") |
| 84 | + args = ap.parse_args() |
| 85 | + |
| 86 | + if args.download: |
| 87 | + from diffusers import AutoPipelineForText2Image |
| 88 | + AutoPipelineForText2Image.from_pretrained(args.model, variant="fp16") |
| 89 | + print(f"{args.model} is ready.") |
| 90 | + return 0 |
| 91 | + |
| 92 | + if args.job: |
| 93 | + with open(args.job, encoding="utf-8") as f: |
| 94 | + job = json.load(f) |
| 95 | + elif args.prompt and args.out: |
| 96 | + job = {"prompt": args.prompt, "strength": args.strength, "size": args.size, "seed": args.seed} |
| 97 | + if args.init: |
| 98 | + job["frames"] = [[args.init, args.out]] |
| 99 | + else: |
| 100 | + job["keyframe"] = args.out |
| 101 | + else: |
| 102 | + ap.error("give --prompt and --out, or --job") |
| 103 | + |
| 104 | + t2i, i2i = load_pipelines(args.model, args.device) |
| 105 | + prompt, size, seed = job["prompt"], int(job.get("size", 512)), job.get("seed") |
| 106 | + strength = float(job.get("strength", 0.5)) |
| 107 | + if job.get("keyframe"): |
| 108 | + save(text_to_image(t2i, prompt, size, seed), job["keyframe"]) |
| 109 | + for src, dst in job.get("frames", []): |
| 110 | + save(image_to_image(i2i, prompt, src, size, strength, seed), dst) |
| 111 | + return 0 |
| 112 | + |
| 113 | + |
| 114 | +if __name__ == "__main__": |
| 115 | + sys.exit(main()) |
0 commit comments