fix(api): get upscale params from request
This commit is contained in:
parent
120056f878
commit
1f0c19af04
|
@ -16,6 +16,7 @@ from .image import (
|
|||
)
|
||||
from .upscale import (
|
||||
upscale_resrgan,
|
||||
UpscaleParams,
|
||||
)
|
||||
from .utils import (
|
||||
safer_join,
|
||||
|
@ -75,7 +76,8 @@ def run_txt2img_pipeline(
|
|||
ctx: ServerContext,
|
||||
params: BaseParams,
|
||||
size: Size,
|
||||
output: str
|
||||
output: str,
|
||||
upscale: UpscaleParams
|
||||
):
|
||||
pipe = load_pipeline(OnnxStableDiffusionPipeline,
|
||||
params.model, params.provider, params.scheduler)
|
||||
|
@ -93,7 +95,9 @@ def run_txt2img_pipeline(
|
|||
negative_prompt=params.negative_prompt,
|
||||
num_inference_steps=params.steps,
|
||||
).images[0]
|
||||
image = upscale_resrgan(image, ctx.model_path)
|
||||
|
||||
if upscale.faces or upscale.scale > 1:
|
||||
image = upscale_resrgan(ctx, image, upscale)
|
||||
|
||||
dest = safer_join(ctx.output_path, output)
|
||||
image.save(dest)
|
||||
|
@ -105,8 +109,9 @@ def run_img2img_pipeline(
|
|||
ctx: ServerContext,
|
||||
params: BaseParams,
|
||||
output: str,
|
||||
upscale: UpscaleParams,
|
||||
source_image: Image,
|
||||
strength: float
|
||||
strength: float,
|
||||
):
|
||||
pipe = load_pipeline(OnnxStableDiffusionImg2ImgPipeline,
|
||||
params.model, params.provider, params.scheduler)
|
||||
|
@ -122,7 +127,9 @@ def run_img2img_pipeline(
|
|||
num_inference_steps=params.steps,
|
||||
strength=strength,
|
||||
).images[0]
|
||||
image = upscale_resrgan(image, ctx.model_path)
|
||||
|
||||
if upscale.faces or upscale.scale > 1:
|
||||
image = upscale_resrgan(ctx, image, upscale)
|
||||
|
||||
dest = safer_join(ctx.output_path, output)
|
||||
image.save(dest)
|
||||
|
@ -135,6 +142,7 @@ def run_inpaint_pipeline(
|
|||
params: BaseParams,
|
||||
size: Size,
|
||||
output: str,
|
||||
upscale: UpscaleParams,
|
||||
source_image: Image,
|
||||
mask_image: Image,
|
||||
expand: Border,
|
||||
|
@ -173,6 +181,9 @@ def run_inpaint_pipeline(
|
|||
width=size.width,
|
||||
).images[0]
|
||||
|
||||
if upscale.faces or upscale.scale > 1:
|
||||
image = upscale_resrgan(ctx, image, upscale)
|
||||
|
||||
dest = safer_join(ctx.output_path, output)
|
||||
image.save(dest)
|
||||
|
||||
|
|
|
@ -39,6 +39,9 @@ from .pipeline import (
|
|||
run_inpaint_pipeline,
|
||||
run_txt2img_pipeline,
|
||||
)
|
||||
from .upscale import (
|
||||
UpscaleParams,
|
||||
)
|
||||
from .utils import (
|
||||
get_and_clamp_float,
|
||||
get_and_clamp_int,
|
||||
|
@ -186,6 +189,13 @@ def border_from_request() -> Border:
|
|||
return Border(left, right, top, bottom)
|
||||
|
||||
|
||||
def upscale_from_request() -> UpscaleParams:
|
||||
denoise = get_and_clamp_float(request.args, 'denoise', 0.5, 1.0, 0.0)
|
||||
scale = get_and_clamp_int(request.args, 'scale', 1, 4, 1)
|
||||
faces = request.args.get('faces', 'false') == 'true'
|
||||
platform = 'onnx'
|
||||
return UpscaleParams(scale=scale, faces=faces, platform=platform, denoise=denoise)
|
||||
|
||||
def check_paths():
|
||||
if not path.exists(model_path):
|
||||
raise RuntimeError('model path must exist')
|
||||
|
@ -278,6 +288,7 @@ def img2img():
|
|||
source_image = Image.open(BytesIO(source_file.read())).convert('RGB')
|
||||
|
||||
params, size = pipeline_from_request()
|
||||
upscale = upscale_from_request()
|
||||
|
||||
strength = get_and_clamp_float(
|
||||
request.args,
|
||||
|
@ -294,7 +305,7 @@ def img2img():
|
|||
|
||||
source_image.thumbnail((size.width, size.height))
|
||||
executor.submit_stored(output, run_img2img_pipeline,
|
||||
context, params, output, source_image, strength)
|
||||
context, params, output, upscale, source_image, strength)
|
||||
|
||||
return jsonify({
|
||||
'output': output,
|
||||
|
@ -306,6 +317,7 @@ def img2img():
|
|||
@app.route('/api/txt2img', methods=['POST'])
|
||||
def txt2img():
|
||||
params, size = pipeline_from_request()
|
||||
upscale = upscale_from_request()
|
||||
|
||||
output = make_output_name(
|
||||
'txt2img',
|
||||
|
@ -314,7 +326,7 @@ def txt2img():
|
|||
print("txt2img output: %s" % (output))
|
||||
|
||||
executor.submit_stored(
|
||||
output, run_txt2img_pipeline, context, params, size, output)
|
||||
output, run_txt2img_pipeline, context, params, size, output, upscale)
|
||||
|
||||
return jsonify({
|
||||
'output': output,
|
||||
|
@ -333,6 +345,7 @@ def inpaint():
|
|||
|
||||
params, size = pipeline_from_request()
|
||||
expand = border_from_request()
|
||||
upscale = upscale_from_request()
|
||||
|
||||
mask_filter = get_from_map(request.args, 'filter', mask_filters, 'none')
|
||||
noise_source = get_from_map(
|
||||
|
@ -362,6 +375,7 @@ def inpaint():
|
|||
params,
|
||||
size,
|
||||
output,
|
||||
upscale,
|
||||
source_image,
|
||||
mask_image,
|
||||
expand,
|
||||
|
|
|
@ -129,6 +129,8 @@ def make_resrgan(ctx: ServerContext, params: UpscaleParams, tile=0):
|
|||
|
||||
|
||||
def upscale_resrgan(ctx: ServerContext, source_image: Image, params: UpscaleParams) -> Image:
|
||||
print('upscaling image with Real ESRGAN', params)
|
||||
|
||||
image = np.array(source_image)
|
||||
upsampler = make_resrgan(ctx.model_path)
|
||||
|
||||
|
@ -142,6 +144,8 @@ def upscale_resrgan(ctx: ServerContext, source_image: Image, params: UpscalePara
|
|||
|
||||
|
||||
def upscale_gfpgan(ctx: ServerContext, image, upsampler=None) -> Image:
|
||||
print('correcting faces with GFPGAN')
|
||||
|
||||
if upsampler is None:
|
||||
upsampler = make_resrgan(ctx.model_path, 512)
|
||||
|
||||
|
|
|
@ -17,8 +17,10 @@
|
|||
"CUDA",
|
||||
"ddim",
|
||||
"ddpm",
|
||||
"denoise",
|
||||
"directml",
|
||||
"ftfy",
|
||||
"gfpgan",
|
||||
"Heun",
|
||||
"huggingface",
|
||||
"Inpaint",
|
||||
|
@ -35,6 +37,7 @@
|
|||
"pndm",
|
||||
"pretrained",
|
||||
"protobuf",
|
||||
"resrgan",
|
||||
"runwayml",
|
||||
"scandir",
|
||||
"scipy",
|
||||
|
@ -42,6 +45,7 @@
|
|||
"spacy",
|
||||
"spinalcase",
|
||||
"stringcase",
|
||||
"upsampler",
|
||||
"venv",
|
||||
"virtualenv",
|
||||
"zustand"
|
||||
|
|
Loading…
Reference in New Issue