File size: 3,142 Bytes
848c832 ff1df96 848c832 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
# -- coding: utf-8 --`
import argparse
import os
# engine
from stable_diffusion_engine import StableDiffusionEngine
# scheduler
from diffusers import LMSDiscreteScheduler, PNDMScheduler
# utils
import cv2
import numpy as np
def main(args):
if args.seed is not None:
np.random.seed(args.seed)
if args.init_image is None:
scheduler = LMSDiscreteScheduler(
beta_start=args.beta_start,
beta_end=args.beta_end,
beta_schedule=args.beta_schedule,
tensor_format="np"
)
else:
scheduler = PNDMScheduler(
beta_start=args.beta_start,
beta_end=args.beta_end,
beta_schedule=args.beta_schedule,
skip_prk_steps = True,
tensor_format="np"
)
engine = StableDiffusionEngine(
model = args.model,
scheduler = scheduler,
tokenizer = args.tokenizer
)
image = engine(
prompt = args.prompt,
init_image = None if args.init_image is None else cv2.imread(args.init_image),
mask = None if args.mask is None else cv2.imread(args.mask, 0),
strength = args.strength,
num_inference_steps = args.num_inference_steps,
guidance_scale = args.guidance_scale,
eta = args.eta
)
cv2.imwrite(args.output, image)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# pipeline configure
parser.add_argument("--model", type=str, default="bes-dev/stable-diffusion-v1-4-openvino", help="model name")
# randomizer params
parser.add_argument("--seed", type=int, default=None, help="random seed for generating consistent images per prompt")
# scheduler params
parser.add_argument("--beta-start", type=float, default=0.00085, help="LMSDiscreteScheduler::beta_start")
parser.add_argument("--beta-end", type=float, default=0.012, help="LMSDiscreteScheduler::beta_end")
parser.add_argument("--beta-schedule", type=str, default="scaled_linear", help="LMSDiscreteScheduler::beta_schedule")
# diffusion params
parser.add_argument("--num-inference-steps", type=int, default=32, help="num inference steps")
parser.add_argument("--guidance-scale", type=float, default=7.5, help="guidance scale")
parser.add_argument("--eta", type=float, default=0.0, help="eta")
# tokenizer
parser.add_argument("--tokenizer", type=str, default="openai/clip-vit-large-patch14", help="tokenizer")
# prompt
parser.add_argument("--prompt", type=str, default="Street-art painting of Emilia Clarke in style of Banksy, photorealism", help="prompt")
# img2img params
parser.add_argument("--init-image", type=str, default=None, help="path to initial image")
parser.add_argument("--strength", type=float, default=0.5, help="how strong the initial image should be noised [0.0, 1.0]")
# inpainting
parser.add_argument("--mask", type=str, default=None, help="mask of the region to inpaint on the initial image")
# output name
parser.add_argument("--output", type=str, default="output.png", help="output image name")
args = parser.parse_args()
main(args)
|