diff --git a/backend/invoke_ai_web_server.py b/backend/invoke_ai_web_server.py index 36494f720c..0c29807878 100644 --- a/backend/invoke_ai_web_server.py +++ b/backend/invoke_ai_web_server.py @@ -22,8 +22,8 @@ from ldm.invoke.pngwriter import PngWriter, retrieve_metadata from ldm.invoke.prompt_parser import split_weighted_subprompts from backend.modules.parameters import parameters_to_command -from backend.modules.get_outpainting_generation_mode import ( - get_outpainting_generation_mode, +from backend.modules.get_canvas_generation_mode import ( + get_canvas_generation_mode, ) # Loading Arguments @@ -606,7 +606,10 @@ class InvokeAIWebServer: """ Prepare for generation based on generation_mode """ - if generation_parameters["generation_mode"] == "outpainting": + if generation_parameters["generation_mode"] in [ + "outpainting", + "inpainting", + ]: """ generation_parameters["init_img"] is a base64 image generation_parameters["init_mask"] is a base64 image @@ -617,46 +620,60 @@ class InvokeAIWebServer: truncated_outpaint_image_b64 = generation_parameters["init_img"][:64] truncated_outpaint_mask_b64 = generation_parameters["init_mask"][:64] - outpaint_image = dataURL_to_image( - generation_parameters["init_img"] - ).convert("RGBA") + init_img_url = generation_parameters["init_img"] + + original_bounding_box = generation_parameters["bounding_box"].copy() + + if generation_parameters["generation_mode"] == "outpainting": + initial_image = dataURL_to_image( + generation_parameters["init_img"] + ).convert("RGBA") + + """ + The outpaint image and mask are pre-cropped by the UI, so the bounding box we pass + to the generator should be: + { + "x": 0, + "y": 0, + "width": original_bounding_box["width"], + "height": original_bounding_box["height"] + } + """ + + generation_parameters["bounding_box"]["x"] = 0 + generation_parameters["bounding_box"]["y"] = 0 + elif generation_parameters["generation_mode"] == "inpainting": + init_img_path = self.get_image_path_from_url(init_img_url) + initial_image = Image.open(init_img_path) + + """ + For inpainting, only the mask is pre-cropped by the UI, so we need to crop out a copy + of the region of the image to be inpainted to match the size of the mask image. + """ + initial_image = copy_image_from_bounding_box( + initial_image, **generation_parameters["bounding_box"] + ) # Convert mask dataURL to an image and convert to greyscale - outpaint_mask = dataURL_to_image( + mask_image = dataURL_to_image( generation_parameters["init_mask"] ).convert("L") - actual_generation_mode = get_outpainting_generation_mode( - outpaint_image, outpaint_mask + actual_generation_mode = get_canvas_generation_mode( + initial_image, mask_image ) - """ - The outpaint image and mask are pre-cropped by the UI, so the bounding box we pass - to the generator should be: - { - "x": 0, - "y": 0, - "width": original_bounding_box["width"], - "height": original_bounding_box["height"] - } - - Save the original bounding box, we need to give it back to the UI when finished, - because the UI needs to know where to put the inpainted image on the canvas. - """ - original_bounding_box = generation_parameters["bounding_box"].copy() - - generation_parameters["bounding_box"]["x"] = 0 - generation_parameters["bounding_box"]["y"] = 0 + print(initial_image, mask_image) """ Apply the mask to the init image, creating a "mask" image with transparency where inpainting should occur. This is the kind of mask that prompt2image() needs. """ - alpha_mask = outpaint_image.copy() - alpha_mask.putalpha(outpaint_mask) + alpha_mask = initial_image.copy() + alpha_mask.putalpha(mask_image) - generation_parameters["init_img"] = outpaint_image + generation_parameters["init_img"] = initial_image generation_parameters["init_mask"] = alpha_mask # Remove the unneeded parameters for whichever mode we are doing @@ -691,48 +708,6 @@ class InvokeAIWebServer: generation_parameters.pop("tile_size", None) generation_parameters.pop("force_outpaint", None) - elif generation_parameters["generation_mode"] == "inpainting": - """ - generation_parameters["init_img"] is a url - generation_parameters["init_mask"] is a base64 image - - So we need to convert each into a PIL Image. - """ - truncated_outpaint_image_b64 = generation_parameters["init_img"][:64] - truncated_outpaint_mask_b64 = generation_parameters["init_mask"][:64] - - init_img_url = generation_parameters["init_img"] - init_mask_url = generation_parameters["init_mask"] - - init_img_path = self.get_image_path_from_url(init_img_url) - - original_image = Image.open(init_img_path) - - rgba_image = original_image.convert("RGBA") - # copy a region from it which we will inpaint - cropped_init_image = copy_image_from_bounding_box( - rgba_image, **generation_parameters["bounding_box"] - ) - - original_bounding_box = generation_parameters["bounding_box"].copy() - - generation_parameters["init_img"] = cropped_init_image - - # Convert mask dataURL to an image and convert to greyscale - mask_image = dataURL_to_image( - generation_parameters["init_mask"] - ).convert("L") - - """ - Apply the mask to the init image, creating a "mask" image with - transparency where inpainting should occur. This is the kind of - mask that prompt2image() needs. - """ - alpha_mask = cropped_init_image.copy() - alpha_mask.putalpha(mask_image) - - generation_parameters["init_mask"] = alpha_mask - elif generation_parameters["generation_mode"] == "img2img": init_img_url = generation_parameters["init_img"] init_img_path = self.get_image_path_from_url(init_img_url) diff --git a/backend/modules/get_outpainting_generation_mode.py b/backend/modules/get_canvas_generation_mode.py similarity index 88% rename from backend/modules/get_outpainting_generation_mode.py rename to backend/modules/get_canvas_generation_mode.py index d21e231671..764cf474ff 100644 --- a/backend/modules/get_outpainting_generation_mode.py +++ b/backend/modules/get_canvas_generation_mode.py @@ -21,7 +21,7 @@ def check_for_any_transparency(img: Union[ImageType, str]) -> bool: return False -def get_outpainting_generation_mode( +def get_canvas_generation_mode( init_img: Union[ImageType, str], init_mask: Union[ImageType, str] ) -> Literal["txt2img", "outpainting", "inpainting", "img2img",]: if type(init_img) is str: @@ -80,36 +80,36 @@ def main(): print( "OPAQUE IMAGE, NO MASK, expect img2img, got ", - get_outpainting_generation_mode(init_img_opaque, init_mask_no_mask), + get_canvas_generation_mode(init_img_opaque, init_mask_no_mask), ) print( "IMAGE WITH TRANSPARENCY, NO MASK, expect outpainting, got ", - get_outpainting_generation_mode( + get_canvas_generation_mode( init_img_partial_transparency, init_mask_no_mask ), ) print( "FULLY TRANSPARENT IMAGE NO MASK, expect txt2img, got ", - get_outpainting_generation_mode(init_img_full_transparency, init_mask_no_mask), + get_canvas_generation_mode(init_img_full_transparency, init_mask_no_mask), ) print( "OPAQUE IMAGE, WITH MASK, expect inpainting, got ", - get_outpainting_generation_mode(init_img_opaque, init_mask_has_mask), + get_canvas_generation_mode(init_img_opaque, init_mask_has_mask), ) print( "IMAGE WITH TRANSPARENCY, WITH MASK, expect outpainting, got ", - get_outpainting_generation_mode( + get_canvas_generation_mode( init_img_partial_transparency, init_mask_has_mask ), ) print( "FULLY TRANSPARENT IMAGE WITH MASK, expect txt2img, got ", - get_outpainting_generation_mode(init_img_full_transparency, init_mask_has_mask), + get_canvas_generation_mode(init_img_full_transparency, init_mask_has_mask), )