diff --git a/api/onnx_web/serve.py b/api/onnx_web/serve.py index e9cef04d..2248423f 100644 --- a/api/onnx_web/serve.py +++ b/api/onnx_web/serve.py @@ -103,22 +103,30 @@ def get_latents_from_seed(seed: int, width: int, height: int) -> np.ndarray: # based on https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/master/scripts/outpainting_mk_2.py#L175-L232 -def expand_image(source_image: Image, dims: Tuple[int, int, int, int]): +def expand_image(source_image: Image, mask_image: Image, dims: Tuple[int, int, int, int]): (left, right, top, bottom) = dims full_width = left + source_image.width + right full_height = top + source_image.height + bottom - full_image = Image.new('RGB', (full_width, full_height), 'white') - full_image.paste(source_image, (left, top)) + full_source = Image.new('RGB', (full_width, full_height), 'white') + full_source.paste(source_image, (left, top)) - # draw = ImageDraw.Draw(full_image) - # draw.rectangle((0, 0, left, full_height), 'white') - # draw.rectangle((0, 0, full_width, top), 'white') - # draw.rectangle((full_width - right, 0, full_width, full_height), 'white') - # draw.rectangle((0, full_height - bottom, full_width, full_height), 'white') + full_mask = Image.new('RGB', (full_width, full_height), 'white') + full_mask.paste(mask_image, (left, top)) - return full_image + full_noise = Image.effect_noise((full_width, full_height), 2) + + for x in range(full_source.width): + for y in range(full_source.height): + mask_color = full_mask.getpixel((x, y)) + noise_color = full_noise.getpixel((x, y)) + source_color = full_source.getpixel((x, y)) + + if mask_color[0] > 0: + full_source.putpixel((x, y), source_color + (mask_color * noise_color)) + + return (full_source, full_mask, (full_width, full_height)) def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, scheduler): @@ -288,9 +296,8 @@ def run_inpaint_pipeline(model, provider, scheduler, prompt, negative_prompt, cf latents = get_latents_from_seed(seed, width, height) rng = np.random.RandomState(seed) - full_mask = expand_image(mask_image, (256, 256, 256, 256)) - full_source = expand_image(source_image, (256, 256, 256, 256)) - # full_source.thumbnail((512, 512)) + extra = 256 + (full_source, full_mask, full_dims) = expand_image(source_image, mask_image, (extra, extra, extra, extra)) image = pipe( prompt,