add seed to output filename, apply pep8 to server
This commit is contained in:
parent
d8b6d7fc15
commit
d93d4659fa
16
api/serve.py
16
api/serve.py
|
@ -41,9 +41,11 @@ scheduler_list = {
|
||||||
'dpm-multi': DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder="scheduler"),
|
'dpm-multi': DPMSolverMultistepScheduler.from_pretrained(model_path, subfolder="scheduler"),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
def get_and_clamp(args, key, default_value, max_value, min_value=1):
|
def get_and_clamp(args, key, default_value, max_value, min_value=1):
|
||||||
return min(max(int(args.get(key, default_value)), min_value), max_value)
|
return min(max(int(args.get(key, default_value)), min_value), max_value)
|
||||||
|
|
||||||
|
|
||||||
def get_from_map(args, key, values, default):
|
def get_from_map(args, key, values, default):
|
||||||
selected = args.get(key, default)
|
selected = args.get(key, default)
|
||||||
if selected in values:
|
if selected in values:
|
||||||
|
@ -51,6 +53,7 @@ def get_from_map(args, key, values, default):
|
||||||
else:
|
else:
|
||||||
return values[default]
|
return values[default]
|
||||||
|
|
||||||
|
|
||||||
def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray:
|
def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray:
|
||||||
# 1 is batch size
|
# 1 is batch size
|
||||||
latents_shape = (1, 4, height // 8, width // 8)
|
latents_shape = (1, 4, height // 8, width // 8)
|
||||||
|
@ -59,6 +62,7 @@ def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray:
|
||||||
image_latents = rng.standard_normal(latents_shape).astype(np.float32)
|
image_latents = rng.standard_normal(latents_shape).astype(np.float32)
|
||||||
return image_latents
|
return image_latents
|
||||||
|
|
||||||
|
|
||||||
# setup
|
# setup
|
||||||
if not path.exists(model_path):
|
if not path.exists(model_path):
|
||||||
raise RuntimeError('model path must exist')
|
raise RuntimeError('model path must exist')
|
||||||
|
@ -69,16 +73,20 @@ if not path.exists(output_path):
|
||||||
app = Flask(__name__)
|
app = Flask(__name__)
|
||||||
|
|
||||||
# routes
|
# routes
|
||||||
|
|
||||||
|
|
||||||
@app.route('/')
|
@app.route('/')
|
||||||
def hello():
|
def hello():
|
||||||
return 'Hello, %s' % (__name__)
|
return 'Hello, %s' % (__name__)
|
||||||
|
|
||||||
|
|
||||||
@app.route('/txt2img')
|
@app.route('/txt2img')
|
||||||
def txt2img():
|
def txt2img():
|
||||||
user = request.remote_addr
|
user = request.remote_addr
|
||||||
|
|
||||||
prompt = request.args.get('prompt', default_prompt)
|
prompt = request.args.get('prompt', default_prompt)
|
||||||
scheduler = get_from_map(request.args, 'scheduler', scheduler_list, 'euler-a')
|
scheduler = get_from_map(request.args, 'scheduler',
|
||||||
|
scheduler_list, 'euler-a')
|
||||||
cfg = get_and_clamp(request.args, 'cfg', default_cfg, max_cfg, 0)
|
cfg = get_and_clamp(request.args, 'cfg', default_cfg, max_cfg, 0)
|
||||||
steps = get_and_clamp(request.args, 'steps', default_steps, max_steps)
|
steps = get_and_clamp(request.args, 'steps', default_steps, max_steps)
|
||||||
height = get_and_clamp(request.args, 'height', default_height, max_height)
|
height = get_and_clamp(request.args, 'height', default_height, max_height)
|
||||||
|
@ -90,7 +98,8 @@ def txt2img():
|
||||||
|
|
||||||
latents = get_latents_from_seed(seed, width, height)
|
latents = get_latents_from_seed(seed, width, height)
|
||||||
|
|
||||||
print("txt2img from %s: %s/%s, %sx%s, %s, %s" % (user, cfg, steps, width, height, seed, prompt))
|
print("txt2img from %s: %s/%s, %sx%s, %s, %s" %
|
||||||
|
(user, cfg, steps, width, height, seed, prompt))
|
||||||
|
|
||||||
pipe = OnnxStableDiffusionPipeline.from_pretrained(
|
pipe = OnnxStableDiffusionPipeline.from_pretrained(
|
||||||
model_path,
|
model_path,
|
||||||
|
@ -107,7 +116,8 @@ def txt2img():
|
||||||
latents=latents
|
latents=latents
|
||||||
).images[0]
|
).images[0]
|
||||||
|
|
||||||
output = '%s/txt2img_%s.png' % (output_path, spinalcase(prompt[0:64]))
|
output = '%s/txt2img_%s_%s.png' % (output_path,
|
||||||
|
seed, spinalcase(prompt[0:64]))
|
||||||
print("txt2img output: %s" % (output))
|
print("txt2img output: %s" % (output))
|
||||||
image.save(output)
|
image.save(output)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue