File size: 4,277 Bytes
848c832 829c94d 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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
# -- coding: utf-8 --`
import argparse
import os
import random
import streamlit as st
from streamlit_drawable_canvas import st_canvas
import numpy as np
import cv2
from PIL import Image, ImageEnhance
import numpy as np
# engine
from stable_diffusion_engine import StableDiffusionEngine
# scheduler
from diffusers import PNDMScheduler
def run(engine):
with st.form(key="request"):
with st.sidebar:
prompt = st.text_area(label='Enter prompt')
with st.expander("Initial image"):
init_image = st.file_uploader("init_image", type=['jpg','png','jpeg'])
stroke_width = st.slider("stroke_width", 1, 100, 50)
stroke_color = st.color_picker("stroke_color", "#00FF00")
canvas_result = st_canvas(
fill_color="rgb(0, 0, 0)",
stroke_width = stroke_width,
stroke_color = stroke_color,
background_color = "#000000",
background_image = Image.open(init_image) if init_image else None,
height = 512,
width = 512,
drawing_mode = "freedraw",
key = "canvas"
)
if init_image is not None:
init_image = cv2.cvtColor(np.array(Image.open(init_image)), cv2.COLOR_RGB2BGR)
if canvas_result.image_data is not None:
mask = cv2.cvtColor(canvas_result.image_data, cv2.COLOR_BGRA2GRAY)
mask[mask > 0] = 255
else:
mask = None
num_inference_steps = st.select_slider(
label='num_inference_steps',
options=range(1, 150),
value=32
)
guidance_scale = st.select_slider(
label='guidance_scale',
options=range(1, 21),
value=7
)
strength = st.slider(
label='strength',
min_value = 0.0,
max_value = 1.0,
value = 0.5
)
seed = st.number_input(
label='seed',
min_value = 0,
max_value = 2 ** 31,
value = random.randint(0, 2 ** 31)
)
generate = st.form_submit_button(label = 'Generate')
if prompt:
np.random.seed(seed)
image = engine(
prompt = prompt,
init_image = init_image,
mask = mask,
strength = strength,
num_inference_steps = num_inference_steps,
guidance_scale = guidance_scale
)
st.image(Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)), width=512)
@st.cache(allow_output_mutation=True)
def load_engine(args):
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
)
return engine
if __name__ == "__main__":
parser = argparse.ArgumentParser()
# pipeline configure
parser.add_argument("--model", type=str, default="4eJIoBek/stable-diffusion-v1-4-openvino-fp32", help="model name")
# 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")
# tokenizer
parser.add_argument("--tokenizer", type=str, default="openai/clip-vit-large-patch14", help="tokenizer")
try:
args = parser.parse_args()
except SystemExit as e:
# This exception will be raised if --help or invalid command line arguments
# are used. Currently streamlit prevents the program from exiting normally
# so we have to do a hard exit.
os._exit(e.code)
engine = load_engine(args)
run(engine)
|