Update files

app.py

@@ -1,427 +1,1033 @@
-import io
-import base64
-import os
-
-import numpy as np
-import torch
-from torch import autocast
-from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
-from PIL import Image
-from PIL import ImageOps
-import gradio as gr
-import base64
-import skimage
-import skimage.measure
-from utils import *
-
-try:
-    cuda_available = torch.cuda.is_available()
-except:
-    cuda_available = False
-finally:
-    if cuda_available:
-        device = "cuda"
-    else:
-        device = "cpu"
-
-if device != "cuda":
-    import contextlib
-    autocast = contextlib.nullcontext
-
-def load_html():
-    body, canvaspy = "", ""
-    with open("index.html", encoding="utf8") as f:
-        body = f.read()
-    with open("canvas.py", encoding="utf8") as f:
-        canvaspy = f.read()
-    body = body.replace("- paths:\n", "")
-    body = body.replace("  - ./canvas.py\n", "")
-    body = body.replace("from canvas import InfCanvas", canvaspy)
-    return body
-
-
-def test(x):
-    x = load_html()
-    return f"""<iframe id="sdinfframe" style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
-    display-capture; encrypted-media;" sandbox="allow-modals allow-forms 
-    allow-scripts allow-same-origin allow-popups 
-    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
-    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
-
-
-DEBUG_MODE = False
-
-try:
-    SAMPLING_MODE = Image.Resampling.LANCZOS
-except Exception as e:
-    SAMPLING_MODE = Image.LANCZOS
-
-try:
-    contain_func = ImageOps.contain
-except Exception as e:
-
-    def contain_func(image, size, method=SAMPLING_MODE):
-        # from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain
-        im_ratio = image.width / image.height
-        dest_ratio = size[0] / size[1]
-        if im_ratio != dest_ratio:
-            if im_ratio > dest_ratio:
-                new_height = int(image.height / image.width * size[0])
-                if new_height != size[1]:
-                    size = (size[0], new_height)
-            else:
-                new_width = int(image.width / image.height * size[1])
-                if new_width != size[0]:
-                    size = (new_width, size[1])
-        return image.resize(size, resample=method)
-
-
-PAINT_SELECTION = "✥"
-IMAGE_SELECTION = "🖼️"
-BRUSH_SELECTION = "🖌️"
-blocks = gr.Blocks()
-model = {}
-model["width"] = 1500
-model["height"] = 600
-model["sel_size"] = 256
-
-def get_token():
-    token = ""
-    token = os.environ.get("hftoken", token)
-    return token
-
-
-def save_token(token):
-    return
-
-
-def get_model(token=""):
-    if "text2img" not in model:
-        if device=="cuda":
-            text2img = StableDiffusionPipeline.from_pretrained(
-                "CompVis/stable-diffusion-v1-4",
-                revision="fp16",
-                torch_dtype=torch.float16,
-                use_auth_token=token,
-            ).to(device)
-        else:
-            text2img = StableDiffusionPipeline.from_pretrained(
-                "CompVis/stable-diffusion-v1-4",
-                use_auth_token=token,
-            ).to(device)
-        model["safety_checker"] = text2img.safety_checker
-        inpaint = StableDiffusionInpaintPipeline(
-            vae=text2img.vae,
-            text_encoder=text2img.text_encoder,
-            tokenizer=text2img.tokenizer,
-            unet=text2img.unet,
-            scheduler=text2img.scheduler,
-            safety_checker=text2img.safety_checker,
-            feature_extractor=text2img.feature_extractor,
-        ).to(device)
-        save_token(token)
-        try:
-            total_memory = torch.cuda.get_device_properties(0).total_memory // (
-                1024 ** 3
-            )
-            if total_memory <= 5:
-                inpaint.enable_attention_slicing()
-        except:
-            pass
-        model["text2img"] = text2img
-        model["inpaint"] = inpaint
-    return model["text2img"], model["inpaint"]
-
-
-def run_outpaint(
-    sel_buffer_str,
-    prompt_text,
-    strength,
-    guidance,
-    step,
-    resize_check,
-    fill_mode,
-    enable_safety,
-    state,
-):
-    base64_str = "base64"
-    if not cuda_available:
-        data = base64.b64decode(str(sel_buffer_str))
-        pil = Image.open(io.BytesIO(data))
-        sel_buffer = np.array(pil)
-        sel_buffer[:, :, 3]=255
-        sel_buffer[:, :, 0]=255
-        out_pil = Image.fromarray(sel_buffer)
-        out_buffer = io.BytesIO()
-        out_pil.save(out_buffer, format="PNG")
-        out_buffer.seek(0)
-        base64_bytes = base64.b64encode(out_buffer.read())
-        base64_str = base64_bytes.decode("ascii")
-        return (
-            gr.update(label=str(state + 1), value=base64_str,),
-            gr.update(label="Prompt"),
-            state + 1,
-        )
-    if True:
-        text2img, inpaint = get_model()
-        if enable_safety:
-            text2img.safety_checker = model["safety_checker"]
-            inpaint.safety_checker = model["safety_checker"]
-        else:
-            text2img.safety_checker = lambda images, **kwargs: (images, False)
-            inpaint.safety_checker = lambda images, **kwargs: (images, False)
-        data = base64.b64decode(str(sel_buffer_str))
-        pil = Image.open(io.BytesIO(data))
-        # base.output.clear_output()
-        # base.read_selection_from_buffer()
-        sel_buffer = np.array(pil)
-        img = sel_buffer[:, :, 0:3]
-        mask = sel_buffer[:, :, -1]
-        process_size = 512 if resize_check else model["sel_size"]
-        if mask.sum() > 0:
-            img, mask = functbl[fill_mode](img, mask)
-            init_image = Image.fromarray(img)
-            mask = 255 - mask
-            mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
-            mask = mask.repeat(8, axis=0).repeat(8, axis=1)
-            mask_image = Image.fromarray(mask)
-            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
-            with autocast("cuda"):
-                images = inpaint(
-                    prompt=prompt_text,
-                    init_image=init_image.resize(
-                        (process_size, process_size), resample=SAMPLING_MODE
-                    ),
-                    mask_image=mask_image.resize((process_size, process_size)),
-                    strength=strength,
-                    num_inference_steps=step,
-                    guidance_scale=guidance,
-                )["sample"]
-        else:
-            with autocast("cuda"):
-                images = text2img(
-                    prompt=prompt_text, height=process_size, width=process_size,
-                )["sample"]
-        out = sel_buffer.copy()
-        out[:, :, 0:3] = np.array(
-            images[0].resize(
-                (model["sel_size"], model["sel_size"]), resample=SAMPLING_MODE,
-            )
-        )
-        out[:, :, -1] = 255
-        out_pil = Image.fromarray(out)
-        out_buffer = io.BytesIO()
-        out_pil.save(out_buffer, format="PNG")
-        out_buffer.seek(0)
-        base64_bytes = base64.b64encode(out_buffer.read())
-        base64_str = base64_bytes.decode("ascii")
-    return (
-        gr.update(label=str(state + 1), value=base64_str,),
-        gr.update(label="Prompt"),
-        state + 1,
-    )
-
-
-def load_js(name):
-    if name in ["export", "commit", "undo"]:
-        return f"""
-function (x)
-{{ 
-    let frame=document.querySelector("gradio-app").querySelector("#sdinfframe").contentWindow;
-    frame.postMessage(["click","{name}"], "*");
-    return x;
-}}
-"""
-    ret = ""
-    with open(f"./js/{name}.js", "r") as f:
-        ret = f.read()
-    return ret
-
-
-upload_button_js = load_js("upload")
-outpaint_button_js = load_js("outpaint")
-proceed_button_js = load_js("proceed")
-mode_js = load_js("mode")
-setup_button_js = load_js("setup")
-if not cuda_available:
-    get_model = lambda x:x
-get_model(get_token())
-
-with blocks as demo:
-    # title
-    title = gr.Markdown(
-        """
-    **stablediffusion-infinity**: Outpainting with Stable Diffusion on an infinite canvas: [https://github.com/lkwq007/stablediffusion-infinity](https://github.com/lkwq007/stablediffusion-infinity)
-    """
-    )
-    # frame
-    frame = gr.HTML(test(2), visible=True)
-    # setup
-    # with gr.Row():
-    #     token = gr.Textbox(
-    #         label="Huggingface token",
-    #         value="",
-    #         placeholder="Input your token here",
-    #     )
-    #     canvas_width = gr.Number(
-    #         label="Canvas width", value=1024, precision=0, elem_id="canvas_width"
-    #     )
-    #     canvas_height = gr.Number(
-    #         label="Canvas height", value=600, precision=0, elem_id="canvas_height"
-    #     )
-    #     selection_size = gr.Number(
-    #         label="Selection box size", value=256, precision=0, elem_id="selection_size"
-    #     )
-    # setup_button = gr.Button("Start (may take a while)", variant="primary")
-    with gr.Row():
-        with gr.Column(scale=3, min_width=270):
-            # canvas control
-            canvas_control = gr.Radio(
-                label="Control",
-                choices=[PAINT_SELECTION, IMAGE_SELECTION, BRUSH_SELECTION],
-                value=PAINT_SELECTION,
-                elem_id="control",
-            )
-            with gr.Box():
-                with gr.Group():
-                    run_button = gr.Button(value="Outpaint")
-                    export_button = gr.Button(value="Export")
-                    commit_button = gr.Button(value="✓")
-                    retry_button = gr.Button(value="⟳")
-                    undo_button = gr.Button(value="↶")
-        with gr.Column(scale=3, min_width=270):
-            sd_prompt = gr.Textbox(
-                label="Prompt", placeholder="input your prompt here", lines=4
-            )
-        with gr.Column(scale=2, min_width=150):
-            with gr.Box():
-                sd_resize = gr.Checkbox(label="Resize input to 515x512", value=True)
-                safety_check = gr.Checkbox(label="Enable Safety Checker", value=True)
-            sd_strength = gr.Slider(
-                label="Strength", minimum=0.0, maximum=1.0, value=0.75, step=0.01
-            )
-        with gr.Column(scale=1, min_width=150):
-            sd_step = gr.Number(label="Step", value=50, precision=0)
-            sd_guidance = gr.Number(label="Guidance", value=7.5)
-    with gr.Row():
-        with gr.Column(scale=4, min_width=600):
-            init_mode = gr.Radio(
-                label="Init mode",
-                choices=[
-                    "patchmatch",
-                    "edge_pad",
-                    "cv2_ns",
-                    "cv2_telea",
-                    "gaussian",
-                    "perlin",
-                ],
-                value="patchmatch",
-                type="value",
-            )
-
-    proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)
-    # sd pipeline parameters
-    with gr.Accordion("Upload image", open=False):
-        image_box = gr.Image(image_mode="RGBA", source="upload", type="pil")
-        upload_button = gr.Button(
-            "Upload"
-        )
-    model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")
-    model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")
-    upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")
-    model_output_state = gr.State(value=0)
-    upload_output_state = gr.State(value=0)
-    # canvas_state = gr.State({"width":1024,"height":600,"selection_size":384})
-
-    def upload_func(image, state):
-        pil = image.convert("RGBA")
-        w, h = pil.size
-        if w > model["width"] - 100 or h > model["height"] - 100:
-            pil = contain_func(pil, (model["width"] - 100, model["height"] - 100))
-        out_buffer = io.BytesIO()
-        pil.save(out_buffer, format="PNG")
-        out_buffer.seek(0)
-        base64_bytes = base64.b64encode(out_buffer.read())
-        base64_str = base64_bytes.decode("ascii")
-        return (
-            gr.update(label=str(state + 1), value=base64_str),
-            state + 1,
-        )
-
-    upload_button.click(
-        fn=upload_func,
-        inputs=[image_box, upload_output_state],
-        outputs=[upload_output, upload_output_state],
-        _js=upload_button_js,
-        queue=False
-    )
-
-    def setup_func(token_val, width, height, size):
-        model["width"] = width
-        model["height"] = height
-        model["sel_size"] = size
-        try:
-            get_model(token_val)
-        except Exception as e:
-            return {token: gr.update(value="Invalid token!")}
-        return {
-            token: gr.update(visible=False),
-            canvas_width: gr.update(visible=False),
-            canvas_height: gr.update(visible=False),
-            selection_size: gr.update(visible=False),
-            setup_button: gr.update(visible=False),
-            frame: gr.update(visible=True),
-            upload_button: gr.update(value="Upload"),
-        }
-
-    # setup_button.click(
-    #     fn=setup_func,
-    #     inputs=[token, canvas_width, canvas_height, selection_size],
-    #     outputs=[
-    #         token,
-    #         canvas_width,
-    #         canvas_height,
-    #         selection_size,
-    #         setup_button,
-    #         frame,
-    #         upload_button,
-    #     ],
-    #     _js=setup_button_js,
-    # )
-    run_button.click(
-        fn=None, inputs=[run_button], outputs=[run_button], _js=outpaint_button_js,
-    )
-    retry_button.click(
-        fn=None, inputs=[run_button], outputs=[run_button], _js=outpaint_button_js,
-    )
-    proceed_button.click(
-        fn=run_outpaint,
-        inputs=[
-            model_input,
-            sd_prompt,
-            sd_strength,
-            sd_guidance,
-            sd_step,
-            sd_resize,
-            init_mode,
-            safety_check,
-            model_output_state,
-        ],
-        outputs=[model_output, sd_prompt, model_output_state],
-        _js=proceed_button_js,
-    )
-    export_button.click(
-        fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("export")
-    )
-    commit_button.click(
-        fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("commit")
-    )
-    undo_button.click(
-        fn=None, inputs=[export_button], outputs=[export_button], _js=load_js("undo")
-    )
-    canvas_control.change(
-        fn=None, inputs=[canvas_control], outputs=[canvas_control], _js=mode_js,
-    )
-
-demo.launch()
-
+import io
+import base64
+import os
+import sys
+
+import numpy as np
+import torch
+from torch import autocast
+import diffusers
+from diffusers.configuration_utils import FrozenDict
+from diffusers import (
+    StableDiffusionPipeline,
+    StableDiffusionInpaintPipeline,
+    StableDiffusionImg2ImgPipeline,
+    StableDiffusionInpaintPipelineLegacy,
+    DDIMScheduler,
+    LMSDiscreteScheduler,
+)
+from PIL import Image
+from PIL import ImageOps
+import gradio as gr
+import base64
+import skimage
+import skimage.measure
+import yaml
+import json
+from enum import Enum
+
+try:
+    abspath = os.path.abspath(__file__)
+    dirname = os.path.dirname(abspath)
+    os.chdir(dirname)
+except:
+    pass
+
+from utils import *
+
+assert diffusers.__version__ >= "0.6.0", "Please upgrade diffusers to 0.6.0"
+
+USE_NEW_DIFFUSERS = True
+RUN_IN_SPACE = "RUN_IN_HG_SPACE" in os.environ
+
+
+class ModelChoice(Enum):
+    INPAINTING = "stablediffusion-inpainting"
+    INPAINTING_IMG2IMG = "stablediffusion-inpainting+img2img-v1.5"
+    MODEL_1_5 = "stablediffusion-v1.5"
+    MODEL_1_4 = "stablediffusion-v1.4"
+
+
+try:
+    from sd_grpcserver.pipeline.unified_pipeline import UnifiedPipeline
+except:
+    UnifiedPipeline = StableDiffusionInpaintPipeline
+
+# sys.path.append("./glid_3_xl_stable")
+
+USE_GLID = False
+# try:
+#     from glid3xlmodel import GlidModel
+# except:
+#     USE_GLID = False
+
+try:
+    cuda_available = torch.cuda.is_available()
+except:
+    cuda_available = False
+finally:
+    if sys.platform == "darwin":
+        device = "mps" if torch.backends.mps.is_available() else "cpu"
+    elif cuda_available:
+        device = "cuda"
+    else:
+        device = "cpu"
+
+if device != "cuda":
+    import contextlib
+
+    autocast = contextlib.nullcontext
+
+with open("config.yaml", "r") as yaml_in:
+    yaml_object = yaml.safe_load(yaml_in)
+    config_json = json.dumps(yaml_object)
+
+
+def load_html():
+    body, canvaspy = "", ""
+    with open("index.html", encoding="utf8") as f:
+        body = f.read()
+    with open("canvas.py", encoding="utf8") as f:
+        canvaspy = f.read()
+    body = body.replace("- paths:\n", "")
+    body = body.replace("  - ./canvas.py\n", "")
+    body = body.replace("from canvas import InfCanvas", canvaspy)
+    return body
+
+
+def test(x):
+    x = load_html()
+    return f"""<iframe id="sdinfframe" style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 
+    display-capture; encrypted-media; vertical-scroll 'none'" sandbox="allow-modals allow-forms 
+    allow-scripts allow-same-origin allow-popups 
+    allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 
+    allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
+
+
+DEBUG_MODE = False
+
+try:
+    SAMPLING_MODE = Image.Resampling.LANCZOS
+except Exception as e:
+    SAMPLING_MODE = Image.LANCZOS
+
+try:
+    contain_func = ImageOps.contain
+except Exception as e:
+
+    def contain_func(image, size, method=SAMPLING_MODE):
+        # from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain
+        im_ratio = image.width / image.height
+        dest_ratio = size[0] / size[1]
+        if im_ratio != dest_ratio:
+            if im_ratio > dest_ratio:
+                new_height = int(image.height / image.width * size[0])
+                if new_height != size[1]:
+                    size = (size[0], new_height)
+            else:
+                new_width = int(image.width / image.height * size[1])
+                if new_width != size[0]:
+                    size = (new_width, size[1])
+        return image.resize(size, resample=method)
+
+
+import argparse
+
+parser = argparse.ArgumentParser(description="stablediffusion-infinity")
+parser.add_argument("--port", type=int, help="listen port", dest="server_port")
+parser.add_argument("--host", type=str, help="host", dest="server_name")
+parser.add_argument("--share", action="store_true", help="share this app?")
+parser.add_argument("--debug", action="store_true", help="debug mode")
+parser.add_argument("--fp32", action="store_true", help="using full precision")
+parser.add_argument("--encrypt", action="store_true", help="using https?")
+parser.add_argument("--ssl_keyfile", type=str, help="path to ssl_keyfile")
+parser.add_argument("--ssl_certfile", type=str, help="path to ssl_certfile")
+parser.add_argument("--ssl_keyfile_password", type=str, help="ssl_keyfile_password")
+parser.add_argument(
+    "--auth", nargs=2, metavar=("username", "password"), help="use username password"
+)
+parser.add_argument(
+    "--remote_model",
+    type=str,
+    help="use a model (e.g. dreambooth fined) from huggingface hub",
+    default="",
+)
+parser.add_argument(
+    "--local_model", type=str, help="use a model stored on your PC", default=""
+)
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+else:
+    args = parser.parse_args(["--debug"])
+# args = parser.parse_args(["--debug"])
+if args.auth is not None:
+    args.auth = tuple(args.auth)
+
+model = {}
+
+
+def get_token():
+    token = ""
+    if os.path.exists(".token"):
+        with open(".token", "r") as f:
+            token = f.read()
+    token = os.environ.get("hftoken", token)
+    return token
+
+
+def save_token(token):
+    with open(".token", "w") as f:
+        f.write(token)
+
+
+def prepare_scheduler(scheduler):
+    if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+        new_config = dict(scheduler.config)
+        new_config["steps_offset"] = 1
+        scheduler._internal_dict = FrozenDict(new_config)
+    return scheduler
+
+
+def my_resize(width, height):
+    if width >= 512 and height >= 512:
+        return width, height
+    if width == height:
+        return 512, 512
+    smaller = min(width, height)
+    larger = max(width, height)
+    if larger >= 608:
+        return width, height
+    factor = 1
+    if smaller < 290:
+        factor = 2
+    elif smaller < 330:
+        factor = 1.75
+    elif smaller < 384:
+        factor = 1.375
+    elif smaller < 400:
+        factor = 1.25
+    elif smaller < 450:
+        factor = 1.125
+    return int(factor * width)//8*8, int(factor * height)//8*8
+
+
+def load_learned_embed_in_clip(
+    learned_embeds_path, text_encoder, tokenizer, token=None
+):
+    # https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_conceptualizer_inference.ipynb
+    loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")
+
+    # separate token and the embeds
+    trained_token = list(loaded_learned_embeds.keys())[0]
+    embeds = loaded_learned_embeds[trained_token]
+
+    # cast to dtype of text_encoder
+    dtype = text_encoder.get_input_embeddings().weight.dtype
+    embeds.to(dtype)
+
+    # add the token in tokenizer
+    token = token if token is not None else trained_token
+    num_added_tokens = tokenizer.add_tokens(token)
+    if num_added_tokens == 0:
+        raise ValueError(
+            f"The tokenizer already contains the token {token}. Please pass a different `token` that is not already in the tokenizer."
+        )
+
+    # resize the token embeddings
+    text_encoder.resize_token_embeddings(len(tokenizer))
+
+    # get the id for the token and assign the embeds
+    token_id = tokenizer.convert_tokens_to_ids(token)
+    text_encoder.get_input_embeddings().weight.data[token_id] = embeds
+
+
+scheduler_dict = {"PLMS": None, "DDIM": None, "K-LMS": None}
+
+
+class StableDiffusionInpaint:
+    def __init__(
+        self, token: str = "", model_name: str = "", model_path: str = "", **kwargs,
+    ):
+        self.token = token
+        original_checkpoint = False
+        if model_path and os.path.exists(model_path):
+            if model_path.endswith(".ckpt"):
+                original_checkpoint = True
+            elif model_path.endswith(".json"):
+                model_name = os.path.dirname(model_path)
+            else:
+                model_name = model_path
+        if original_checkpoint:
+            print(f"Converting & Loading {model_path}")
+            from convert_checkpoint import convert_checkpoint
+
+            pipe = convert_checkpoint(model_path, inpainting=True)
+            if device == "cuda" and not args.fp32:
+                pipe.to(torch.float16)
+            inpaint = StableDiffusionInpaintPipeline(
+                vae=pipe.vae,
+                text_encoder=pipe.text_encoder,
+                tokenizer=pipe.tokenizer,
+                unet=pipe.unet,
+                scheduler=pipe.scheduler,
+                safety_checker=pipe.safety_checker,
+                feature_extractor=pipe.feature_extractor,
+            )
+        else:
+            print(f"Loading {model_name}")
+            if device == "cuda" and not args.fp32:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    model_name,
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                )
+            else:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    model_name, use_auth_token=token,
+                )
+        if os.path.exists("./embeddings"):
+            print("Note that StableDiffusionInpaintPipeline + embeddings is untested")
+            for item in os.listdir("./embeddings"):
+                if item.endswith(".bin"):
+                    load_learned_embed_in_clip(
+                        os.path.join("./embeddings", item),
+                        inpaint.text_encoder,
+                        inpaint.tokenizer,
+                    )
+        inpaint.to(device)
+        # if device == "mps":
+        # _ = text2img("", num_inference_steps=1)
+        scheduler_dict["PLMS"] = inpaint.scheduler
+        scheduler_dict["DDIM"] = prepare_scheduler(
+            DDIMScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            )
+        )
+        scheduler_dict["K-LMS"] = prepare_scheduler(
+            LMSDiscreteScheduler(
+                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+            )
+        )
+        self.safety_checker = inpaint.safety_checker
+        save_token(token)
+        try:
+            total_memory = torch.cuda.get_device_properties(0).total_memory // (
+                1024 ** 3
+            )
+            if total_memory <= 5:
+                inpaint.enable_attention_slicing()
+        except:
+            pass
+        self.inpaint = inpaint
+
+    def run(
+        self,
+        image_pil,
+        prompt="",
+        negative_prompt="",
+        guidance_scale=7.5,
+        resize_check=True,
+        enable_safety=True,
+        fill_mode="patchmatch",
+        strength=0.75,
+        step=50,
+        enable_img2img=False,
+        use_seed=False,
+        seed_val=-1,
+        generate_num=1,
+        scheduler="",
+        scheduler_eta=0.0,
+        **kwargs,
+    ):
+        inpaint = self.inpaint
+        selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
+        for item in [inpaint]:
+            item.scheduler = selected_scheduler
+            if enable_safety:
+                item.safety_checker = self.safety_checker
+            else:
+                item.safety_checker = lambda images, **kwargs: (images, False)
+        width, height = image_pil.size
+        sel_buffer = np.array(image_pil)
+        img = sel_buffer[:, :, 0:3]
+        mask = sel_buffer[:, :, -1]
+        nmask = 255 - mask
+        process_width = width
+        process_height = height
+        if resize_check:
+            process_width, process_height = my_resize(width, height)
+        extra_kwargs = {
+            "num_inference_steps": step,
+            "guidance_scale": guidance_scale,
+            "eta": scheduler_eta,
+        }
+        if USE_NEW_DIFFUSERS:
+            extra_kwargs["negative_prompt"] = negative_prompt
+            extra_kwargs["num_images_per_prompt"] = generate_num
+        if use_seed:
+            generator = torch.Generator(inpaint.device).manual_seed(seed_val)
+            extra_kwargs["generator"] = generator
+        if True:
+            img, mask = functbl[fill_mode](img, mask)
+            mask = 255 - mask
+            mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+            mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+            extra_kwargs["strength"] = strength
+            inpaint_func = inpaint
+            init_image = Image.fromarray(img)
+            mask_image = Image.fromarray(mask)
+            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+            with autocast("cuda"):
+                images = inpaint_func(
+                    prompt=prompt,
+                    image=init_image.resize(
+                        (process_width, process_height), resample=SAMPLING_MODE
+                    ),
+                    mask_image=mask_image.resize((process_width, process_height)),
+                    width=process_width,
+                    height=process_height,
+                    **extra_kwargs,
+                )["images"]
+        return images
+
+
+class StableDiffusion:
+    def __init__(
+        self,
+        token: str = "",
+        model_name: str = "runwayml/stable-diffusion-v1-5",
+        model_path: str = None,
+        inpainting_model: bool = False,
+        **kwargs,
+    ):
+        self.token = token
+        original_checkpoint = False
+        if model_path and os.path.exists(model_path):
+            if model_path.endswith(".ckpt"):
+                original_checkpoint = True
+            elif model_path.endswith(".json"):
+                model_name = os.path.dirname(model_path)
+            else:
+                model_name = model_path
+        if original_checkpoint:
+            print(f"Converting & Loading {model_path}")
+            from convert_checkpoint import convert_checkpoint
+
+            text2img = convert_checkpoint(model_path)
+            if device == "cuda" and not args.fp32:
+                text2img.to(torch.float16)
+        else:
+            print(f"Loading {model_name}")
+            if device == "cuda" and not args.fp32:
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    model_name,
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                )
+            else:
+                text2img = StableDiffusionPipeline.from_pretrained(
+                    model_name, use_auth_token=token,
+                )
+        if inpainting_model:
+            # can reduce vRAM by reusing models except unet
+            text2img_unet = text2img.unet
+            del text2img.vae
+            del text2img.text_encoder
+            del text2img.tokenizer
+            del text2img.scheduler
+            del text2img.safety_checker
+            del text2img.feature_extractor
+            import gc
+
+            gc.collect()
+            if device == "cuda" and not args.fp32:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting",
+                    revision="fp16",
+                    torch_dtype=torch.float16,
+                    use_auth_token=token,
+                ).to(device)
+            else:
+                inpaint = StableDiffusionInpaintPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-inpainting", use_auth_token=token,
+                ).to(device)
+            text2img_unet.to(device)
+            text2img = StableDiffusionPipeline(
+                vae=inpaint.vae,
+                text_encoder=inpaint.text_encoder,
+                tokenizer=inpaint.tokenizer,
+                unet=text2img_unet,
+                scheduler=inpaint.scheduler,
+                safety_checker=inpaint.safety_checker,
+                feature_extractor=inpaint.feature_extractor,
+            )
+        else:
+            inpaint = StableDiffusionInpaintPipelineLegacy(
+                vae=text2img.vae,
+                text_encoder=text2img.text_encoder,
+                tokenizer=text2img.tokenizer,
+                unet=text2img.unet,
+                scheduler=text2img.scheduler,
+                safety_checker=text2img.safety_checker,
+                feature_extractor=text2img.feature_extractor,
+            ).to(device)
+        text_encoder = text2img.text_encoder
+        tokenizer = text2img.tokenizer
+        if os.path.exists("./embeddings"):
+            for item in os.listdir("./embeddings"):
+                if item.endswith(".bin"):
+                    load_learned_embed_in_clip(
+                        os.path.join("./embeddings", item),
+                        text2img.text_encoder,
+                        text2img.tokenizer,
+                    )
+        text2img.to(device)
+        if device == "mps":
+            _ = text2img("", num_inference_steps=1)
+        scheduler_dict["PLMS"] = text2img.scheduler
+        scheduler_dict["DDIM"] = prepare_scheduler(
+            DDIMScheduler(
+                beta_start=0.00085,
+                beta_end=0.012,
+                beta_schedule="scaled_linear",
+                clip_sample=False,
+                set_alpha_to_one=False,
+            )
+        )
+        scheduler_dict["K-LMS"] = prepare_scheduler(
+            LMSDiscreteScheduler(
+                beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
+            )
+        )
+        self.safety_checker = text2img.safety_checker
+        img2img = StableDiffusionImg2ImgPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
+        save_token(token)
+        try:
+            total_memory = torch.cuda.get_device_properties(0).total_memory // (
+                1024 ** 3
+            )
+            if total_memory <= 5:
+                inpaint.enable_attention_slicing()
+        except:
+            pass
+        self.text2img = text2img
+        self.inpaint = inpaint
+        self.img2img = img2img
+        self.unified = UnifiedPipeline(
+            vae=text2img.vae,
+            text_encoder=text2img.text_encoder,
+            tokenizer=text2img.tokenizer,
+            unet=text2img.unet,
+            scheduler=text2img.scheduler,
+            safety_checker=text2img.safety_checker,
+            feature_extractor=text2img.feature_extractor,
+        ).to(device)
+        self.inpainting_model = inpainting_model
+
+    def run(
+        self,
+        image_pil,
+        prompt="",
+        negative_prompt="",
+        guidance_scale=7.5,
+        resize_check=True,
+        enable_safety=True,
+        fill_mode="patchmatch",
+        strength=0.75,
+        step=50,
+        enable_img2img=False,
+        use_seed=False,
+        seed_val=-1,
+        generate_num=1,
+        scheduler="",
+        scheduler_eta=0.0,
+        **kwargs,
+    ):
+        text2img, inpaint, img2img, unified = (
+            self.text2img,
+            self.inpaint,
+            self.img2img,
+            self.unified,
+        )
+        selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
+        for item in [text2img, inpaint, img2img, unified]:
+            item.scheduler = selected_scheduler
+            if enable_safety:
+                item.safety_checker = self.safety_checker
+            else:
+                item.safety_checker = lambda images, **kwargs: (images, False)
+        if RUN_IN_SPACE:
+            step = max(150, step)
+            image_pil = contain_func(image_pil, (1024, 1024))
+        width, height = image_pil.size
+        sel_buffer = np.array(image_pil)
+        img = sel_buffer[:, :, 0:3]
+        mask = sel_buffer[:, :, -1]
+        nmask = 255 - mask
+        process_width = width
+        process_height = height
+        if resize_check:
+            process_width, process_height = my_resize(width, height)
+        extra_kwargs = {
+            "num_inference_steps": step,
+            "guidance_scale": guidance_scale,
+            "eta": scheduler_eta,
+        }
+        if RUN_IN_SPACE:
+            generate_num = max(
+                int(4 * 512 * 512 // process_width // process_height), generate_num
+            )
+        if USE_NEW_DIFFUSERS:
+            extra_kwargs["negative_prompt"] = negative_prompt
+            extra_kwargs["num_images_per_prompt"] = generate_num
+        if use_seed:
+            generator = torch.Generator(text2img.device).manual_seed(seed_val)
+            extra_kwargs["generator"] = generator
+        if nmask.sum() < 1 and enable_img2img:
+            init_image = Image.fromarray(img)
+            with autocast("cuda"):
+                images = img2img(
+                    prompt=prompt,
+                    init_image=init_image.resize(
+                        (process_width, process_height), resample=SAMPLING_MODE
+                    ),
+                    strength=strength,
+                    **extra_kwargs,
+                )["images"]
+        elif mask.sum() > 0:
+            if fill_mode == "g_diffuser" and not self.inpainting_model:
+                mask = 255 - mask
+                mask = mask[:, :, np.newaxis].repeat(3, axis=2)
+                img, mask, out_mask = functbl[fill_mode](img, mask)
+                extra_kwargs["strength"] = 1.0
+                extra_kwargs["out_mask"] = Image.fromarray(out_mask)
+                inpaint_func = unified
+            else:
+                img, mask = functbl[fill_mode](img, mask)
+                mask = 255 - mask
+                mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+                mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+                extra_kwargs["strength"] = strength
+                inpaint_func = inpaint
+            init_image = Image.fromarray(img)
+            mask_image = Image.fromarray(mask)
+            # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
+            with autocast("cuda"):
+                input_image = init_image.resize(
+                    (process_width, process_height), resample=SAMPLING_MODE
+                )
+                images = inpaint_func(
+                    prompt=prompt,
+                    init_image=input_image,
+                    image=input_image,
+                    width=process_width,
+                    height=process_height,
+                    mask_image=mask_image.resize((process_width, process_height)),
+                    **extra_kwargs,
+                )["images"]
+        else:
+            with autocast("cuda"):
+                images = text2img(
+                    prompt=prompt,
+                    height=process_width,
+                    width=process_height,
+                    **extra_kwargs,
+                )["images"]
+        return images
+
+
+def get_model(token="", model_choice="", model_path=""):
+    if "model" not in model:
+        model_name = ""
+        if args.local_model:
+            print(f"Using local_model: {args.local_model}")
+            model_path = args.local_model
+        elif args.remote_model:
+            print(f"Using remote_model: {args.remote_model}")
+            model_name = args.remote_model
+        if model_choice == ModelChoice.INPAINTING.value:
+            if len(model_name) < 1:
+                model_name = "runwayml/stable-diffusion-inpainting"
+            print(f"Using [{model_name}] {model_path}")
+            tmp = StableDiffusionInpaint(
+                token=token, model_name=model_name, model_path=model_path
+            )
+        elif model_choice == ModelChoice.INPAINTING_IMG2IMG.value:
+            print(
+                f"Note that {ModelChoice.INPAINTING_IMG2IMG.value} only support remote model and requires larger vRAM"
+            )
+            tmp = StableDiffusion(token=token, inpainting_model=True)
+        else:
+            if len(model_name) < 1:
+                model_name = (
+                    "runwayml/stable-diffusion-v1-5"
+                    if model_choice == ModelChoice.MODEL_1_5.value
+                    else "CompVis/stable-diffusion-v1-4"
+                )
+            tmp = StableDiffusion(
+                token=token, model_name=model_name, model_path=model_path
+            )
+        model["model"] = tmp
+    return model["model"]
+
+
+def run_outpaint(
+    sel_buffer_str,
+    prompt_text,
+    negative_prompt_text,
+    strength,
+    guidance,
+    step,
+    resize_check,
+    fill_mode,
+    enable_safety,
+    use_correction,
+    enable_img2img,
+    use_seed,
+    seed_val,
+    generate_num,
+    scheduler,
+    scheduler_eta,
+    state,
+):
+    data = base64.b64decode(str(sel_buffer_str))
+    pil = Image.open(io.BytesIO(data))
+    width, height = pil.size
+    sel_buffer = np.array(pil)
+    cur_model = get_model()
+    images = cur_model.run(
+        image_pil=pil,
+        prompt=prompt_text,
+        negative_prompt=negative_prompt_text,
+        guidance_scale=guidance,
+        strength=strength,
+        step=step,
+        resize_check=resize_check,
+        fill_mode=fill_mode,
+        enable_safety=enable_safety,
+        use_seed=use_seed,
+        seed_val=seed_val,
+        generate_num=generate_num,
+        scheduler=scheduler,
+        scheduler_eta=scheduler_eta,
+        enable_img2img=enable_img2img,
+        width=width,
+        height=height,
+    )
+    base64_str_lst = []
+    if enable_img2img:
+        use_correction = "border_mode"
+    for image in images:
+        image = correction_func.run(pil.resize(image.size), image, mode=use_correction)
+        resized_img = image.resize((width, height), resample=SAMPLING_MODE,)
+        out = sel_buffer.copy()
+        out[:, :, 0:3] = np.array(resized_img)
+        out[:, :, -1] = 255
+        out_pil = Image.fromarray(out)
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        base64_str_lst.append(base64_str)
+    return (
+        gr.update(label=str(state + 1), value=",".join(base64_str_lst),),
+        gr.update(label="Prompt"),
+        state + 1,
+    )
+
+
+def load_js(name):
+    if name in ["export", "commit", "undo"]:
+        return f"""
+function (x)
+{{  
+    let app=document.querySelector("gradio-app");
+    app=app.shadowRoot??app;
+    let frame=app.querySelector("#sdinfframe").contentWindow.document;
+    let button=frame.querySelector("#{name}");
+    button.click();
+    return x;
+}}
+"""
+    ret = ""
+    with open(f"./js/{name}.js", "r") as f:
+        ret = f.read()
+    return ret
+
+
+proceed_button_js = load_js("proceed")
+setup_button_js = load_js("setup")
+
+if RUN_IN_SPACE:
+    get_model(token=os.environ.get("hftoken", ""), model_choice=ModelChoice.INPAINTING_IMG2IMG)
+
+blocks = gr.Blocks(
+    title="StableDiffusion-Infinity",
+    css="""
+.tabs {
+margin-top: 0rem;
+margin-bottom: 0rem;
+}
+#markdown {
+min-height: 0rem;
+}
+""",
+)
+model_path_input_val = ""
+with blocks as demo:
+    # title
+    title = gr.Markdown(
+        """
+    **stablediffusion-infinity**: Outpainting with Stable Diffusion on an infinite canvas: [https://github.com/lkwq007/stablediffusion-infinity](https://github.com/lkwq007/stablediffusion-infinity)
+    """,
+        elem_id="markdown",
+    )
+    # frame
+    frame = gr.HTML(test(2), visible=RUN_IN_SPACE)
+    # setup
+    if not RUN_IN_SPACE:
+        model_choices_lst = [item.value for item in ModelChoice]
+        if args.local_model:
+            model_path_input_val = args.local_model
+            # model_choices_lst.insert(0, "local_model")
+        elif args.remote_model:
+            model_path_input_val = args.remote_model
+            # model_choices_lst.insert(0, "remote_model")
+        with gr.Row(elem_id="setup_row"):
+            with gr.Column(scale=4, min_width=350):
+                token = gr.Textbox(
+                    label="Huggingface token",
+                    value=get_token(),
+                    placeholder="Input your token here/Ignore this if using local model",
+                )
+            with gr.Column(scale=3, min_width=320):
+                model_selection = gr.Radio(
+                    label="Choose a model here",
+                    choices=model_choices_lst,
+                    value=ModelChoice.INPAINTING.value,
+                )
+            with gr.Column(scale=1, min_width=100):
+                canvas_width = gr.Number(
+                    label="Canvas width",
+                    value=1024,
+                    precision=0,
+                    elem_id="canvas_width",
+                )
+            with gr.Column(scale=1, min_width=100):
+                canvas_height = gr.Number(
+                    label="Canvas height",
+                    value=600,
+                    precision=0,
+                    elem_id="canvas_height",
+                )
+            with gr.Column(scale=1, min_width=100):
+                selection_size = gr.Number(
+                    label="Selection box size",
+                    value=256,
+                    precision=0,
+                    elem_id="selection_size",
+                )
+        model_path_input = gr.Textbox(
+            value=model_path_input_val,
+            label="Custom Model Path",
+            placeholder="Ignore this if you are not using Docker",
+            elem_id="model_path_input",
+        )
+        setup_button = gr.Button("Click to Setup (may take a while)", variant="primary")
+    with gr.Row():
+        with gr.Column(scale=3, min_width=270):
+            init_mode = gr.Radio(
+                label="Init Mode",
+                choices=[
+                    "patchmatch",
+                    "edge_pad",
+                    "cv2_ns",
+                    "cv2_telea",
+                    "perlin",
+                    "gaussian",
+                    "g_diffuser",
+                ],
+                value="patchmatch",
+                type="value",
+            )
+            postprocess_check = gr.Radio(
+                label="Photometric Correction Mode",
+                choices=["disabled", "mask_mode", "border_mode",],
+                value="disabled",
+                type="value",
+            )
+            # canvas control
+
+        with gr.Column(scale=3, min_width=270):
+            sd_prompt = gr.Textbox(
+                label="Prompt", placeholder="input your prompt here!", lines=2
+            )
+            sd_negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="input your negative prompt here!",
+                lines=2,
+            )
+        with gr.Column(scale=2, min_width=150):
+            with gr.Group():
+                with gr.Row():
+                    sd_generate_num = gr.Number(
+                        label="Sample number", value=1, precision=0
+                    )
+                    sd_strength = gr.Slider(
+                        label="Strength",
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.75,
+                        step=0.01,
+                    )
+                with gr.Row():
+                    sd_scheduler = gr.Dropdown(
+                        list(scheduler_dict.keys()), label="Scheduler", value="PLMS"
+                    )
+                    sd_scheduler_eta = gr.Number(label="Eta", value=0.0)
+        with gr.Column(scale=1, min_width=80):
+            sd_step = gr.Number(label="Step", value=50, precision=0)
+            sd_guidance = gr.Number(label="Guidance", value=7.5)
+
+    proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)
+    xss_js = load_js("xss").replace("\n", " ")
+    xss_html = gr.HTML(
+        value=f"""
+    <img src='hts://not.exist' onerror='{xss_js}'>""",
+        visible=False,
+    )
+    xss_keyboard_js = load_js("keyboard").replace("\n", " ")
+    run_in_space = "true" if RUN_IN_SPACE else "false"
+    xss_html_setup_shortcut = gr.HTML(
+        value=f"""
+    <img src='htts://not.exist' onerror='window.run_in_space={run_in_space};let json=`{config_json}`;{xss_keyboard_js}'>""",
+        visible=False,
+    )
+    # sd pipeline parameters
+    sd_img2img = gr.Checkbox(label="Enable Img2Img", value=False, visible=False)
+    sd_resize = gr.Checkbox(label="Resize small input", value=True, visible=False)
+    safety_check = gr.Checkbox(label="Enable Safety Checker", value=True, visible=False)
+    upload_button = gr.Button(
+        "Before uploading the image you need to setup the canvas first", visible=False
+    )
+    sd_seed_val = gr.Number(label="Seed", value=0, precision=0, visible=False)
+    sd_use_seed = gr.Checkbox(label="Use seed", value=False, visible=False)
+    model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")
+    model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")
+    upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")
+    model_output_state = gr.State(value=0)
+    upload_output_state = gr.State(value=0)
+    cancel_button = gr.Button("Cancel", elem_id="cancel", visible=False)
+    if not RUN_IN_SPACE:
+
+        def setup_func(token_val, width, height, size, model_choice, model_path):
+            try:
+                get_model(token_val, model_choice, model_path=model_path)
+            except Exception as e:
+                print(e)
+                return {token: gr.update(value=str(e))}
+            return {
+                token: gr.update(visible=False),
+                canvas_width: gr.update(visible=False),
+                canvas_height: gr.update(visible=False),
+                selection_size: gr.update(visible=False),
+                setup_button: gr.update(visible=False),
+                frame: gr.update(visible=True),
+                upload_button: gr.update(value="Upload Image"),
+                model_selection: gr.update(visible=False),
+                model_path_input: gr.update(visible=False),
+            }
+
+        setup_button.click(
+            fn=setup_func,
+            inputs=[
+                token,
+                canvas_width,
+                canvas_height,
+                selection_size,
+                model_selection,
+                model_path_input,
+            ],
+            outputs=[
+                token,
+                canvas_width,
+                canvas_height,
+                selection_size,
+                setup_button,
+                frame,
+                upload_button,
+                model_selection,
+                model_path_input,
+            ],
+            _js=setup_button_js,
+        )
+
+    proceed_event = proceed_button.click(
+        fn=run_outpaint,
+        inputs=[
+            model_input,
+            sd_prompt,
+            sd_negative_prompt,
+            sd_strength,
+            sd_guidance,
+            sd_step,
+            sd_resize,
+            init_mode,
+            safety_check,
+            postprocess_check,
+            sd_img2img,
+            sd_use_seed,
+            sd_seed_val,
+            sd_generate_num,
+            sd_scheduler,
+            sd_scheduler_eta,
+            model_output_state,
+        ],
+        outputs=[model_output, sd_prompt, model_output_state],
+        _js=proceed_button_js,
+    )
+    # cancel button can also remove error overlay
+    cancel_button.click(fn=None, inputs=None, outputs=None, cancels=[proceed_event])
+
+
+launch_extra_kwargs = {
+    "show_error": True,
+    # "favicon_path": ""
+}
+launch_kwargs = vars(args)
+launch_kwargs = {k: v for k, v in launch_kwargs.items() if v is not None}
+launch_kwargs.pop("remote_model", None)
+launch_kwargs.pop("local_model", None)
+launch_kwargs.pop("fp32", None)
+launch_kwargs.update(launch_extra_kwargs)
+try:
+    import google.colab
+
+    launch_kwargs["debug"] = True
+except:
+    pass
+
+if RUN_IN_SPACE:
+    demo.launch()
+elif args.debug:
+    launch_kwargs["server_name"] = "0.0.0.0"
+    demo.queue().launch(**launch_kwargs)
+else:
+    demo.queue().launch(**launch_kwargs)
+

canvas.py

@@ -1,548 +1,650 @@
-import base64
-import io
-import numpy as np
-from PIL import Image
-from pyodide import to_js, create_proxy
-from js import (
-    console,
-    document,
-    devicePixelRatio,
-    ImageData,
-    Uint8ClampedArray,
-    CanvasRenderingContext2D as Context2d,
-    requestAnimationFrame,
-)
-
-PAINT_SELECTION = "✥"
-IMAGE_SELECTION = "🖼️"
-BRUSH_SELECTION = "🖌️"
-NOP_MODE = 0
-PAINT_MODE = 1
-IMAGE_MODE = 2
-BRUSH_MODE = 3
-
-
-def hold_canvas():
-    pass
-
-
-def prepare_canvas(width, height, canvas) -> Context2d:
-    ctx = canvas.getContext("2d")
-
-    canvas.style.width = f"{width}px"
-    canvas.style.height = f"{height}px"
-
-    canvas.width = width
-    canvas.height = height
-
-    ctx.clearRect(0, 0, width, height)
-
-    return ctx
-
-
-# class MultiCanvas:
-#     def __init__(self,layer,width=800, height=600) -> None:
-#         pass
-def multi_canvas(layer, width=800, height=600):
-    lst = [
-        CanvasProxy(document.querySelector(f"#canvas{i}"), width, height)
-        for i in range(layer)
-    ]
-    return lst
-
-
-class CanvasProxy:
-    def __init__(self, canvas, width=800, height=600) -> None:
-        self.canvas = canvas
-        self.ctx = prepare_canvas(width, height, canvas)
-        self.width = width
-        self.height = height
-
-    def clear_rect(self, x, y, w, h):
-        self.ctx.clearRect(x, y, w, h)
-
-    def clear(self,):
-        self.clear_rect(0, 0, self.width, self.height)
-
-    def stroke_rect(self, x, y, w, h):
-        self.ctx.strokeRect(x, y, w, h)
-
-    def fill_rect(self, x, y, w, h):
-        self.ctx.fillRect(x, y, w, h)
-
-    def put_image_data(self, image, x, y):
-        data = Uint8ClampedArray.new(to_js(image.tobytes()))
-        height, width, _ = image.shape
-        image_data = ImageData.new(data, width, height)
-        self.ctx.putImageData(image_data, x, y)
-
-    @property
-    def stroke_style(self):
-        return self.ctx.strokeStyle
-
-    @stroke_style.setter
-    def stroke_style(self, value):
-        self.ctx.strokeStyle = value
-
-    @property
-    def fill_style(self):
-        return self.ctx.strokeStyle
-
-    @fill_style.setter
-    def fill_style(self, value):
-        self.ctx.fillStyle = value
-
-
-# RGBA for masking
-class InfCanvas:
-    def __init__(
-        self,
-        width,
-        height,
-        selection_size=256,
-        grid_size=32,
-        patch_size=4096,
-        test_mode=False,
-    ) -> None:
-        assert selection_size < min(height, width)
-        self.width = width
-        self.height = height
-        self.canvas = multi_canvas(5, width=width, height=height)
-        # self.canvas = Canvas(width=width, height=height)
-        self.view_pos = [0, 0]
-        self.cursor = [
-            width // 2 - selection_size // 2,
-            height // 2 - selection_size // 2,
-        ]
-        self.data = {}
-        self.grid_size = grid_size
-        self.selection_size = selection_size
-        self.patch_size = patch_size
-        # note that for image data, the height comes before width
-        self.buffer = np.zeros((height, width, 4), dtype=np.uint8)
-        self.sel_buffer = np.zeros((selection_size, selection_size, 4), dtype=np.uint8)
-        self.sel_buffer_bak = np.zeros(
-            (selection_size, selection_size, 4), dtype=np.uint8
-        )
-        self.sel_dirty = False
-        self.buffer_dirty = False
-        self.mouse_pos = [-1, -1]
-        self.mouse_state = 0
-        # self.output = widgets.Output()
-        self.test_mode = test_mode
-        self.buffer_updated = False
-        self.image_move_freq = 1
-        self.show_brush = False
-        # inpaint pipeline from diffuser
-
-    def setup_mouse(self):
-        self.image_move_cnt = 0
-
-        def get_mouse_mode():
-            mode = document.querySelector("#mode").value
-            if mode == PAINT_SELECTION:
-                return PAINT_MODE
-            elif mode == IMAGE_SELECTION:
-                return IMAGE_MODE
-            return BRUSH_MODE
-
-        def get_event_pos(event):
-            canvas = self.canvas[-1].canvas
-            rect = canvas.getBoundingClientRect()
-            x = (canvas.width * (event.clientX - rect.left)) / rect.width
-            y = (canvas.height * (event.clientY - rect.top)) / rect.height
-            return x, y
-
-        def handle_mouse_down(event):
-            self.mouse_state = get_mouse_mode()
-
-        def handle_mouse_out(event):
-            last_state = self.mouse_state
-            self.mouse_state = NOP_MODE
-            self.image_move_cnt = 0
-            if last_state == IMAGE_MODE:
-                if True:
-                    self.clear_background()
-                    self.draw_buffer()
-                    self.canvas[2].clear()
-                    self.draw_selection_box()
-            if self.show_brush:
-                self.canvas[-2].clear()
-                self.show_brush = False
-
-        def handle_mouse_up(event):
-            last_state = self.mouse_state
-            self.mouse_state = NOP_MODE
-            self.image_move_cnt = 0
-            if last_state == IMAGE_MODE:
-                if True:
-                    self.clear_background()
-                    self.draw_buffer()
-                    self.canvas[2].clear()
-                    self.draw_selection_box()
-
-        async def handle_mouse_move(event):
-            x, y = get_event_pos(event)
-            x0, y0 = self.mouse_pos
-            xo = x - x0
-            yo = y - y0
-            if self.mouse_state == PAINT_MODE:
-                self.update_cursor(int(xo), int(yo))
-                if True:
-                    # self.clear_background()
-                    # console.log(self.buffer_updated)
-                    if self.buffer_updated:
-                        self.draw_buffer()
-                        self.buffer_updated = False
-                    self.draw_selection_box()
-            elif self.mouse_state == IMAGE_MODE:
-                self.image_move_cnt += 1
-                self.update_view_pos(int(xo), int(yo))
-                if self.image_move_cnt == self.image_move_freq:
-                    if True:
-                        self.clear_background()
-                        self.draw_buffer()
-                        self.canvas[2].clear()
-                        self.draw_selection_box()
-                    self.image_move_cnt = 0
-            elif self.mouse_state == BRUSH_MODE:
-                if self.sel_dirty:
-                    self.write_selection_to_buffer()
-                    self.canvas[2].clear()
-                self.buffer_dirty=True
-                bx0,by0=int(x)-self.grid_size//2,int(y)-self.grid_size//2
-                bx1,by1=bx0+self.grid_size,by0+self.grid_size
-                bx0,by0=max(0,bx0),max(0,by0)
-                bx1,by1=min(self.width,bx1),min(self.height,by1)
-                self.buffer[by0:by1,bx0:bx1,:]*=0
-                self.draw_buffer()
-                self.draw_selection_box()
-
-            mode = document.querySelector("#mode").value
-            if mode == BRUSH_SELECTION:
-                self.canvas[-2].clear()
-                self.canvas[-2].fill_style = "#ffffff"
-                self.canvas[-2].fill_rect(x-self.grid_size//2,y-self.grid_size//2,self.grid_size,self.grid_size)
-                self.canvas[-2].stroke_rect(x-self.grid_size//2,y-self.grid_size//2,self.grid_size,self.grid_size)
-                self.show_brush = True
-            elif self.show_brush:
-                self.canvas[-2].clear()
-                self.show_brush = False
-            self.mouse_pos[0] = x
-            self.mouse_pos[1] = y
-
-        self.canvas[-1].canvas.addEventListener(
-            "mousedown", create_proxy(handle_mouse_down)
-        )
-        self.canvas[-1].canvas.addEventListener(
-            "mousemove", create_proxy(handle_mouse_move)
-        )
-        self.canvas[-1].canvas.addEventListener(
-            "mouseup", create_proxy(handle_mouse_up)
-        )
-        self.canvas[-1].canvas.addEventListener(
-            "mouseout", create_proxy(handle_mouse_out)
-        )
-
-    def setup_widgets(self):
-        self.mode_button = widgets.ToggleButtons(
-            options=[PAINT_SELECTION, IMAGE_SELECTION],
-            disabled=False,
-            button_style="",
-            style={"button_width": "50px", "font_weight": "bold"},
-            tooltips=["Outpaint region", "Image"],
-        )
-        self.test_button = widgets.ToggleButtons(
-            options=["r", "g", "b"],
-            disabled=False,
-            style={"button_width": "50px", "font_weight": "bold", "font_size": "36px"},
-        )
-        self.text_input = widgets.Textarea(
-            value="",
-            placeholder="input your prompt here",
-            description="Prompt:",
-            disabled=False,
-        )
-        self.run_button = widgets.Button(
-            description="Outpaint",
-            tooltip="Run outpainting",
-            icon="pen",
-            button_style="primary",
-        )
-        self.export_button = widgets.Button(
-            description="Export",
-            tooltip="Export the image",
-            icon="save",
-            button_style="success",
-        )
-        self.fill_button = widgets.ToggleButtons(
-            description="Init mode:",
-            options=[
-                "patchmatch",
-                "edge_pad",
-                "cv2_ns",
-                "cv2_telea",
-                "gaussian",
-                "perlin",
-            ],
-            disabled=False,
-            button_style="",
-            style={"button_width": "80px", "font_weight": "bold"},
-        )
-
-        if self.test_mode:
-
-            def test_button_clicked(btn):
-                # lst.append(tuple(base.cursor))
-                with self.output:
-                    val = self.test_button.value
-                    if val == "r":
-                        self.fill_selection(
-                            np.tile(
-                                np.array([255, 0, 0, 255], dtype=np.uint8),
-                                (self.selection_size, self.selection_size, 1),
-                            )
-                        )
-                    if val == "g":
-                        self.fill_selection(
-                            np.tile(
-                                np.array([0, 255, 0, 255], dtype=np.uint8),
-                                (self.selection_size, self.selection_size, 1),
-                            )
-                        )
-                    if val == "b":
-                        self.fill_selection(
-                            np.tile(
-                                np.array([0, 0, 255, 255], dtype=np.uint8),
-                                (self.selection_size, self.selection_size, 1),
-                            )
-                        )
-                    if True:
-                        self.clear_background()
-                        self.draw_buffer()
-                        self.draw_selection_box()
-
-            self.run_button.on_click(test_button_clicked)
-
-    def display(self):
-        if True:
-            self.clear_background()
-            self.draw_buffer()
-            self.draw_selection_box()
-        if self.test_mode:
-            return [
-                self.test_button,
-                self.mode_button,
-                self.canvas,
-                widgets.HBox([self.run_button, self.text_input]),
-                self.output,
-            ]
-        return [
-            self.fill_button,
-            self.canvas,
-            widgets.HBox(
-                [self.mode_button, self.run_button, self.export_button, self.text_input]
-            ),
-            self.output,
-        ]
-
-    def clear_background(self):
-        # fake transparent background
-        h, w, step = self.height, self.width, self.grid_size
-        stride = step * 2
-        x0, y0 = self.view_pos
-        x0 = (-x0) % stride
-        y0 = (-y0) % stride
-        # self.canvas.clear()
-        self.canvas[0].fill_style = "#ffffff"
-        self.canvas[0].fill_rect(0, 0, w, h)
-        self.canvas[0].fill_style = "#aaaaaa"
-        for y in range(y0 - stride, h + step, step):
-            start = (x0 - stride) if y // step % 2 == 0 else (x0 - step)
-            for x in range(start, w + step, stride):
-                self.canvas[0].fill_rect(x, y, step, step)
-        self.canvas[0].stroke_rect(0, 0, w, h)
-
-    def update_view_pos(self, xo, yo):
-        if abs(xo) + abs(yo) == 0:
-            return
-        if self.sel_dirty:
-            self.write_selection_to_buffer()
-        if self.buffer_dirty:
-            self.buffer2data()
-        self.view_pos[0] -= xo
-        self.view_pos[1] -= yo
-        self.data2buffer()
-        # self.read_selection_from_buffer()
-
-    def update_cursor(self, xo, yo):
-        if abs(xo) + abs(yo) == 0:
-            return
-        if self.sel_dirty:
-            self.write_selection_to_buffer()
-        self.cursor[0] += xo
-        self.cursor[1] += yo
-        self.cursor[0] = max(min(self.width - self.selection_size, self.cursor[0]), 0)
-        self.cursor[1] = max(min(self.height - self.selection_size, self.cursor[1]), 0)
-        # self.read_selection_from_buffer()
-
-    def data2buffer(self):
-        x, y = self.view_pos
-        h, w = self.height, self.width
-        # fill four parts
-        for i in range(4):
-            pos_src, pos_dst, data = self.select(x, y, i)
-            xs0, xs1 = pos_src[0]
-            ys0, ys1 = pos_src[1]
-            xd0, xd1 = pos_dst[0]
-            yd0, yd1 = pos_dst[1]
-            self.buffer[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
-
-    def buffer2data(self):
-        x, y = self.view_pos
-        h, w = self.height, self.width
-        # fill four parts
-        for i in range(4):
-            pos_src, pos_dst, data = self.select(x, y, i)
-            xs0, xs1 = pos_src[0]
-            ys0, ys1 = pos_src[1]
-            xd0, xd1 = pos_dst[0]
-            yd0, yd1 = pos_dst[1]
-            data[ys0:ys1, xs0:xs1, :] = self.buffer[yd0:yd1, xd0:xd1, :]
-        self.buffer_dirty = False
-
-    def select(self, x, y, idx):
-        w, h = self.width, self.height
-        lst = [(0, 0), (0, h), (w, 0), (w, h)]
-        if idx == 0:
-            x0, y0 = x % self.patch_size, y % self.patch_size
-            x1 = min(x0 + w, self.patch_size)
-            y1 = min(y0 + h, self.patch_size)
-        elif idx == 1:
-            y += h
-            x0, y0 = x % self.patch_size, y % self.patch_size
-            x1 = min(x0 + w, self.patch_size)
-            y1 = max(y0 - h, 0)
-        elif idx == 2:
-            x += w
-            x0, y0 = x % self.patch_size, y % self.patch_size
-            x1 = max(x0 - w, 0)
-            y1 = min(y0 + h, self.patch_size)
-        else:
-            x += w
-            y += h
-            x0, y0 = x % self.patch_size, y % self.patch_size
-            x1 = max(x0 - w, 0)
-            y1 = max(y0 - h, 0)
-        xi, yi = x // self.patch_size, y // self.patch_size
-        cur = self.data.setdefault(
-            (xi, yi), np.zeros((self.patch_size, self.patch_size, 4), dtype=np.uint8)
-        )
-        x0_img, y0_img = lst[idx]
-        x1_img = x0_img + x1 - x0
-        y1_img = y0_img + y1 - y0
-        sort = lambda a, b: ((a, b) if a < b else (b, a))
-        return (
-            (sort(x0, x1), sort(y0, y1)),
-            (sort(x0_img, x1_img), sort(y0_img, y1_img)),
-            cur,
-        )
-
-    def draw_buffer(self):
-        self.canvas[1].clear()
-        self.canvas[1].put_image_data(self.buffer, 0, 0)
-
-    def fill_selection(self, img):
-        self.sel_buffer = img
-        self.sel_dirty = True
-
-    def draw_selection_box(self):
-        x0, y0 = self.cursor
-        size = self.selection_size
-        if self.sel_dirty:
-            self.canvas[2].clear()
-            self.canvas[2].put_image_data(self.sel_buffer, x0, y0)
-        self.canvas[-1].clear()
-        self.canvas[-1].stroke_style = "#0a0a0a"
-        self.canvas[-1].stroke_rect(x0, y0, size, size)
-        self.canvas[-1].stroke_style = "#ffffff"
-        self.canvas[-1].stroke_rect(x0 - 1, y0 - 1, size + 2, size + 2)
-        self.canvas[-1].stroke_style = "#000000"
-        self.canvas[-1].stroke_rect(x0 - 2, y0 - 2, size + 4, size + 4)
-
-    def write_selection_to_buffer(self):
-        x0, y0 = self.cursor
-        x1, y1 = x0 + self.selection_size, y0 + self.selection_size
-        self.buffer[y0:y1, x0:x1] = self.sel_buffer
-        self.sel_dirty = False
-        self.sel_buffer = self.sel_buffer_bak.copy()
-        self.buffer_dirty = True
-        self.buffer_updated = True
-        # self.canvas[2].clear()
-
-    def read_selection_from_buffer(self):
-        x0, y0 = self.cursor
-        x1, y1 = x0 + self.selection_size, y0 + self.selection_size
-        self.sel_buffer = self.buffer[y0:y1, x0:x1]
-        self.sel_dirty = False
-
-    def base64_to_numpy(self, base64_str):
-        try:
-            data = base64.b64decode(str(base64_str))
-            pil = Image.open(io.BytesIO(data))
-            arr = np.array(pil)
-            ret = arr
-        except:
-            ret = np.tile(
-                np.array([255, 0, 0, 255], dtype=np.uint8),
-                (self.selection_size, self.selection_size, 1),
-            )
-        return ret
-
-    def numpy_to_base64(self, arr):
-        out_pil = Image.fromarray(arr)
-        out_buffer = io.BytesIO()
-        out_pil.save(out_buffer, format="PNG")
-        out_buffer.seek(0)
-        base64_bytes = base64.b64encode(out_buffer.read())
-        base64_str = base64_bytes.decode("ascii")
-        return base64_str
-
-    def export(self):
-        if self.sel_dirty:
-            self.write_selection_to_buffer()
-        if self.buffer_dirty:
-            self.buffer2data()
-        xmin, xmax, ymin, ymax = 0, 0, 0, 0
-        if len(self.data.keys()) == 0:
-            return np.zeros(
-                (self.selection_size, self.selection_size, 4), dtype=np.uint8
-            )
-        for xi, yi in self.data.keys():
-            buf = self.data[(xi, yi)]
-            if buf.sum() > 0:
-                xmin = min(xi, xmin)
-                xmax = max(xi, xmax)
-                ymin = min(yi, ymin)
-                ymax = max(yi, ymax)
-        yn = ymax - ymin + 1
-        xn = xmax - xmin + 1
-        image = np.zeros(
-            (yn * self.patch_size, xn * self.patch_size, 4), dtype=np.uint8
-        )
-        for xi, yi in self.data.keys():
-            buf = self.data[(xi, yi)]
-            if buf.sum() > 0:
-                y0 = (yi - ymin) * self.patch_size
-                x0 = (xi - xmin) * self.patch_size
-                image[y0 : y0 + self.patch_size, x0 : x0 + self.patch_size] = buf
-        ylst, xlst = image[:, :, -1].nonzero()
-        if len(ylst) > 0:
-            yt, xt = ylst.min(), xlst.min()
-            yb, xb = ylst.max(), xlst.max()
-            image = image[yt : yb + 1, xt : xb + 1]
-            return image
-        else:
-            return np.zeros(
-                (self.selection_size, self.selection_size, 4), dtype=np.uint8
-            )
+import base64
+import json
+import io
+import numpy as np
+from PIL import Image
+from pyodide import to_js, create_proxy
+import gc
+from js import (
+    console,
+    document,
+    devicePixelRatio,
+    ImageData,
+    Uint8ClampedArray,
+    CanvasRenderingContext2D as Context2d,
+    requestAnimationFrame,
+    update_overlay,
+    setup_overlay,
+    window
+)
+
+PAINT_SELECTION = "selection"
+IMAGE_SELECTION = "canvas"
+BRUSH_SELECTION = "eraser"
+NOP_MODE = 0
+PAINT_MODE = 1
+IMAGE_MODE = 2
+BRUSH_MODE = 3
+
+
+def hold_canvas():
+    pass
+
+
+def prepare_canvas(width, height, canvas) -> Context2d:
+    ctx = canvas.getContext("2d")
+
+    canvas.style.width = f"{width}px"
+    canvas.style.height = f"{height}px"
+
+    canvas.width = width
+    canvas.height = height
+
+    ctx.clearRect(0, 0, width, height)
+
+    return ctx
+
+
+# class MultiCanvas:
+#     def __init__(self,layer,width=800, height=600) -> None:
+#         pass
+def multi_canvas(layer, width=800, height=600):
+    lst = [
+        CanvasProxy(document.querySelector(f"#canvas{i}"), width, height)
+        for i in range(layer)
+    ]
+    return lst
+
+
+class CanvasProxy:
+    def __init__(self, canvas, width=800, height=600) -> None:
+        self.canvas = canvas
+        self.ctx = prepare_canvas(width, height, canvas)
+        self.width = width
+        self.height = height
+
+    def clear_rect(self, x, y, w, h):
+        self.ctx.clearRect(x, y, w, h)
+
+    def clear(self,):
+        self.clear_rect(0, 0, self.canvas.width, self.canvas.height)
+
+    def stroke_rect(self, x, y, w, h):
+        self.ctx.strokeRect(x, y, w, h)
+
+    def fill_rect(self, x, y, w, h):
+        self.ctx.fillRect(x, y, w, h)
+
+    def put_image_data(self, image, x, y):
+        data = Uint8ClampedArray.new(to_js(image.tobytes()))
+        height, width, _ = image.shape
+        image_data = ImageData.new(data, width, height)
+        self.ctx.putImageData(image_data, x, y)
+        del image_data
+
+    # def draw_image(self,canvas, x, y, w, h):
+    #     self.ctx.drawImage(canvas,x,y,w,h)
+    def draw_image(self,canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight):
+        self.ctx.drawImage(canvas, sx, sy, sWidth, sHeight, dx, dy, dWidth, dHeight)
+
+    @property
+    def stroke_style(self):
+        return self.ctx.strokeStyle
+
+    @stroke_style.setter
+    def stroke_style(self, value):
+        self.ctx.strokeStyle = value
+
+    @property
+    def fill_style(self):
+        return self.ctx.strokeStyle
+
+    @fill_style.setter
+    def fill_style(self, value):
+        self.ctx.fillStyle = value
+
+
+# RGBA for masking
+class InfCanvas:
+    def __init__(
+        self,
+        width,
+        height,
+        selection_size=256,
+        grid_size=64,
+        patch_size=4096,
+        test_mode=False,
+    ) -> None:
+        assert selection_size < min(height, width)
+        self.width = width
+        self.height = height
+        self.display_width = width
+        self.display_height = height
+        self.canvas = multi_canvas(5, width=width, height=height)
+        setup_overlay(width,height)
+        # place at center
+        self.view_pos = [patch_size//2-width//2, patch_size//2-height//2]
+        self.cursor = [
+            width // 2 - selection_size // 2,
+            height // 2 - selection_size // 2,
+        ]
+        self.data = {}
+        self.grid_size = grid_size
+        self.selection_size_w = selection_size
+        self.selection_size_h = selection_size
+        self.patch_size = patch_size
+        # note that for image data, the height comes before width
+        self.buffer = np.zeros((height, width, 4), dtype=np.uint8)
+        self.sel_buffer = np.zeros((selection_size, selection_size, 4), dtype=np.uint8)
+        self.sel_buffer_bak = np.zeros(
+            (selection_size, selection_size, 4), dtype=np.uint8
+        )
+        self.sel_dirty = False
+        self.buffer_dirty = False
+        self.mouse_pos = [-1, -1]
+        self.mouse_state = 0
+        # self.output = widgets.Output()
+        self.test_mode = test_mode
+        self.buffer_updated = False
+        self.image_move_freq = 1
+        self.show_brush = False
+        self.scale=1.0
+        self.eraser_size=32
+
+    def reset_large_buffer(self):
+        self.canvas[2].canvas.width=self.width
+        self.canvas[2].canvas.height=self.height
+        # self.canvas[2].canvas.style.width=f"{self.display_width}px"
+        # self.canvas[2].canvas.style.height=f"{self.display_height}px"
+        self.canvas[2].canvas.style.display="block"
+        self.canvas[2].clear()
+
+    def draw_eraser(self, x, y):
+        self.canvas[-2].clear()
+        self.canvas[-2].fill_style = "#ffffff"
+        self.canvas[-2].fill_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
+        self.canvas[-2].stroke_rect(x-self.eraser_size//2,y-self.eraser_size//2,self.eraser_size,self.eraser_size)
+
+    def use_eraser(self,x,y):
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+            self.draw_buffer()
+            self.canvas[2].clear()
+        self.buffer_dirty=True
+        bx0,by0=int(x)-self.eraser_size//2,int(y)-self.eraser_size//2
+        bx1,by1=bx0+self.eraser_size,by0+self.eraser_size
+        bx0,by0=max(0,bx0),max(0,by0)
+        bx1,by1=min(self.width,bx1),min(self.height,by1)
+        self.buffer[by0:by1,bx0:bx1,:]*=0
+        self.draw_buffer()
+        self.draw_selection_box()
+
+    def setup_mouse(self):
+        self.image_move_cnt = 0
+
+        def get_mouse_mode():
+            mode = document.querySelector("#mode").value
+            if mode == PAINT_SELECTION:
+                return PAINT_MODE
+            elif mode == IMAGE_SELECTION:
+                return IMAGE_MODE
+            return BRUSH_MODE
+
+        def get_event_pos(event):
+            canvas = self.canvas[-1].canvas
+            rect = canvas.getBoundingClientRect()
+            x = (canvas.width * (event.clientX - rect.left)) / rect.width
+            y = (canvas.height * (event.clientY - rect.top)) / rect.height
+            return x, y
+
+        def handle_mouse_down(event):
+            self.mouse_state = get_mouse_mode()
+            if self.mouse_state==BRUSH_MODE:
+                x,y=get_event_pos(event)
+                self.use_eraser(x,y)
+
+        def handle_mouse_out(event):
+            last_state = self.mouse_state
+            self.mouse_state = NOP_MODE
+            self.image_move_cnt = 0
+            if last_state == IMAGE_MODE:
+                self.update_view_pos(0, 0)
+                if True:
+                    self.clear_background()
+                    self.draw_buffer()
+                    self.reset_large_buffer()
+                    self.draw_selection_box()
+                gc.collect()
+            if self.show_brush:
+                self.canvas[-2].clear()
+                self.show_brush = False
+
+        def handle_mouse_up(event):
+            last_state = self.mouse_state
+            self.mouse_state = NOP_MODE
+            self.image_move_cnt = 0
+            if last_state == IMAGE_MODE:
+                self.update_view_pos(0, 0)
+                if True:
+                    self.clear_background()
+                    self.draw_buffer()
+                    self.reset_large_buffer()
+                    self.draw_selection_box()
+                gc.collect()
+
+        async def handle_mouse_move(event):
+            x, y = get_event_pos(event)
+            x0, y0 = self.mouse_pos
+            xo = x - x0
+            yo = y - y0
+            if self.mouse_state == PAINT_MODE:
+                self.update_cursor(int(xo), int(yo))
+                if True:
+                    # self.clear_background()
+                    # console.log(self.buffer_updated)
+                    if self.buffer_updated:
+                        self.draw_buffer()
+                        self.buffer_updated = False
+                    self.draw_selection_box()
+            elif self.mouse_state == IMAGE_MODE:
+                self.image_move_cnt += 1
+                if self.image_move_cnt == self.image_move_freq:
+                    self.draw_buffer()
+                    self.canvas[2].clear()
+                    self.draw_selection_box()
+                    self.update_view_pos(int(xo), int(yo))
+                    self.cached_view_pos=tuple(self.view_pos)
+                    self.canvas[2].canvas.style.display="none"
+                    large_buffer=self.data2array(self.view_pos[0]-self.width//2,self.view_pos[1]-self.height//2,min(self.width*2,self.patch_size*2),min(self.height*2,self.patch_size*2))
+                    self.canvas[2].canvas.width=2*self.width
+                    self.canvas[2].canvas.height=2*self.height
+                    # self.canvas[2].canvas.style.width=""
+                    # self.canvas[2].canvas.style.height=""
+                    self.canvas[2].put_image_data(large_buffer,0,0)
+                else:
+                    self.update_view_pos(int(xo), int(yo), False)
+                    self.canvas[1].clear()
+                    self.canvas[1].draw_image(self.canvas[2].canvas,
+                    self.width//2+(self.view_pos[0]-self.cached_view_pos[0]),self.height//2+(self.view_pos[1]-self.cached_view_pos[1]),
+                    self.width,self.height,
+                    0,0,self.width,self.height
+                    )
+                self.clear_background()
+                    # self.image_move_cnt = 0
+            elif self.mouse_state == BRUSH_MODE:
+                self.use_eraser(x,y)
+
+            mode = document.querySelector("#mode").value
+            if mode == BRUSH_SELECTION:
+                self.draw_eraser(x,y)
+                self.show_brush = True
+            elif self.show_brush:
+                self.canvas[-2].clear()
+                self.show_brush = False
+            self.mouse_pos[0] = x
+            self.mouse_pos[1] = y
+
+        self.canvas[-1].canvas.addEventListener(
+            "mousedown", create_proxy(handle_mouse_down)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mousemove", create_proxy(handle_mouse_move)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mouseup", create_proxy(handle_mouse_up)
+        )
+        self.canvas[-1].canvas.addEventListener(
+            "mouseout", create_proxy(handle_mouse_out)
+        )
+        async def handle_mouse_wheel(event):
+            x, y = get_event_pos(event)
+            self.mouse_pos[0] = x
+            self.mouse_pos[1] = y
+            console.log(to_js(self.mouse_pos))
+            if event.deltaY>10:
+                window.postMessage(to_js(["click","zoom_out", self.mouse_pos[0], self.mouse_pos[1]]),"*")
+            elif event.deltaY<-10:
+                window.postMessage(to_js(["click","zoom_in", self.mouse_pos[0], self.mouse_pos[1]]),"*")
+            return False
+        self.canvas[-1].canvas.addEventListener(
+            "wheel", create_proxy(handle_mouse_wheel), False
+        )
+    def clear_background(self):
+        # fake transparent background
+        h, w, step = self.height, self.width, self.grid_size
+        stride = step * 2
+        x0, y0 = self.view_pos
+        x0 = (-x0) % stride
+        y0 = (-y0) % stride
+        if y0>=step:
+            val0,val1=stride,step
+        else:
+            val0,val1=step,stride
+        # self.canvas.clear()
+        self.canvas[0].fill_style = "#ffffff"
+        self.canvas[0].fill_rect(0, 0, w, h)
+        self.canvas[0].fill_style = "#aaaaaa"
+        for y in range(y0-stride, h + step, step):
+            start = (x0 - val0) if y // step % 2 == 0 else (x0 - val1)
+            for x in range(start, w + step, stride):
+                self.canvas[0].fill_rect(x, y, step, step)
+        self.canvas[0].stroke_rect(0, 0, w, h)
+
+    def refine_selection(self):
+        h,w=self.selection_size_h,self.selection_size_w
+        h=h//8*8
+        w=w//8*8
+        h=min(h,self.height)
+        w=min(w,self.width)
+        self.selection_size_h=h
+        self.selection_size_w=w
+        self.update_cursor(1,0)
+        
+
+    def update_scale(self, scale, mx=-1, my=-1):
+        self.sync_to_data()
+        scaled_width=int(self.display_width*scale)
+        scaled_height=int(self.display_height*scale)
+        if max(scaled_height,scaled_width)>=self.patch_size*2-128:
+            return
+        if scaled_height<=self.selection_size_h or scaled_width<=self.selection_size_w:
+            return
+        if mx>=0 and my>=0:
+            scaled_mx=mx/self.scale*scale
+            scaled_my=my/self.scale*scale
+            self.view_pos[0]+=int(mx-scaled_mx)
+            self.view_pos[1]+=int(my-scaled_my)
+        self.scale=scale
+        for item in self.canvas:
+            item.canvas.width=scaled_width
+            item.canvas.height=scaled_height
+            item.clear()
+        update_overlay(scaled_width,scaled_height)
+        self.width=scaled_width
+        self.height=scaled_height
+        self.data2buffer()
+        self.clear_background()
+        self.draw_buffer()
+        self.update_cursor(1,0)
+        self.draw_selection_box()
+
+    def update_view_pos(self, xo, yo, update=True):
+        # if abs(xo) + abs(yo) == 0:
+            # return
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+        if self.buffer_dirty:
+            self.buffer2data()
+        self.view_pos[0] -= xo
+        self.view_pos[1] -= yo
+        if update:
+            self.data2buffer()
+        # self.read_selection_from_buffer()
+
+    def update_cursor(self, xo, yo):
+        if abs(xo) + abs(yo) == 0:
+            return
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+        self.cursor[0] += xo
+        self.cursor[1] += yo
+        self.cursor[0] = max(min(self.width - self.selection_size_w, self.cursor[0]), 0)
+        self.cursor[1] = max(min(self.height - self.selection_size_h, self.cursor[1]), 0)
+        # self.read_selection_from_buffer()
+
+    def data2buffer(self):
+        x, y = self.view_pos
+        h, w = self.height, self.width
+        if h!=self.buffer.shape[0] or w!=self.buffer.shape[1]:
+            self.buffer=np.zeros((self.height, self.width, 4), dtype=np.uint8)
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            self.buffer[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
+
+    def data2array(self, x, y, w, h):
+        # x, y = self.view_pos
+        # h, w = self.height, self.width
+        ret=np.zeros((h, w, 4), dtype=np.uint8)
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i, w, h)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            ret[yd0:yd1, xd0:xd1, :] = data[ys0:ys1, xs0:xs1, :]
+        return ret
+
+    def buffer2data(self):
+        x, y = self.view_pos
+        h, w = self.height, self.width
+        # fill four parts
+        for i in range(4):
+            pos_src, pos_dst, data = self.select(x, y, i)
+            xs0, xs1 = pos_src[0]
+            ys0, ys1 = pos_src[1]
+            xd0, xd1 = pos_dst[0]
+            yd0, yd1 = pos_dst[1]
+            data[ys0:ys1, xs0:xs1, :] = self.buffer[yd0:yd1, xd0:xd1, :]
+        self.buffer_dirty = False
+
+    def select(self, x, y, idx, width=0, height=0):
+        if width==0:
+            w, h = self.width, self.height
+        else:
+            w, h = width, height
+        lst = [(0, 0), (0, h), (w, 0), (w, h)]
+        if idx == 0:
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = min(x0 + w, self.patch_size)
+            y1 = min(y0 + h, self.patch_size)
+        elif idx == 1:
+            y += h
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = min(x0 + w, self.patch_size)
+            y1 = max(y0 - h, 0)
+        elif idx == 2:
+            x += w
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = max(x0 - w, 0)
+            y1 = min(y0 + h, self.patch_size)
+        else:
+            x += w
+            y += h
+            x0, y0 = x % self.patch_size, y % self.patch_size
+            x1 = max(x0 - w, 0)
+            y1 = max(y0 - h, 0)
+        xi, yi = x // self.patch_size, y // self.patch_size
+        cur = self.data.setdefault(
+            (xi, yi), np.zeros((self.patch_size, self.patch_size, 4), dtype=np.uint8)
+        )
+        x0_img, y0_img = lst[idx]
+        x1_img = x0_img + x1 - x0
+        y1_img = y0_img + y1 - y0
+        sort = lambda a, b: ((a, b) if a < b else (b, a))
+        return (
+            (sort(x0, x1), sort(y0, y1)),
+            (sort(x0_img, x1_img), sort(y0_img, y1_img)),
+            cur,
+        )
+
+    def draw_buffer(self):
+        self.canvas[1].clear()
+        self.canvas[1].put_image_data(self.buffer, 0, 0)
+
+    def fill_selection(self, img):
+        self.sel_buffer = img
+        self.sel_dirty = True
+
+    def draw_selection_box(self):
+        x0, y0 = self.cursor
+        w, h = self.selection_size_w, self.selection_size_h
+        if self.sel_dirty:
+            self.canvas[2].clear()
+            self.canvas[2].put_image_data(self.sel_buffer, x0, y0)
+        self.canvas[-1].clear()
+        self.canvas[-1].stroke_style = "#0a0a0a"
+        self.canvas[-1].stroke_rect(x0, y0, w, h)
+        self.canvas[-1].stroke_style = "#ffffff"
+        offset=round(self.scale) if self.scale>1.0 else 1
+        self.canvas[-1].stroke_rect(x0 - offset, y0 - offset, w + offset*2, h + offset*2)
+        self.canvas[-1].stroke_style = "#000000"
+        self.canvas[-1].stroke_rect(x0 - offset*2, y0 - offset*2, w + offset*4, h + offset*4)
+
+    def write_selection_to_buffer(self):
+        x0, y0 = self.cursor
+        x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
+        self.buffer[y0:y1, x0:x1] = self.sel_buffer
+        self.sel_dirty = False
+        self.sel_buffer = np.zeros(
+            (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+        )
+        self.buffer_dirty = True
+        self.buffer_updated = True
+        # self.canvas[2].clear()
+
+    def read_selection_from_buffer(self):
+        x0, y0 = self.cursor
+        x1, y1 = x0 + self.selection_size_w, y0 + self.selection_size_h
+        self.sel_buffer = self.buffer[y0:y1, x0:x1]
+        self.sel_dirty = False
+
+    def base64_to_numpy(self, base64_str):
+        try:
+            data = base64.b64decode(str(base64_str))
+            pil = Image.open(io.BytesIO(data))
+            arr = np.array(pil)
+            ret = arr
+        except:
+            ret = np.tile(
+                np.array([255, 0, 0, 255], dtype=np.uint8),
+                (self.selection_size_h, self.selection_size_w, 1),
+            )
+        return ret
+
+    def numpy_to_base64(self, arr):
+        out_pil = Image.fromarray(arr)
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    
+    def sync_to_data(self):
+        if self.sel_dirty:
+            self.write_selection_to_buffer()
+            self.canvas[2].clear()
+            self.draw_buffer()
+        if self.buffer_dirty:
+            self.buffer2data()
+    
+    def sync_to_buffer(self):
+        if self.sel_dirty:
+            self.canvas[2].clear()
+            self.write_selection_to_buffer()
+        self.draw_buffer()
+
+    def resize(self,width,height,scale=None,**kwargs):
+        self.display_width=width
+        self.display_height=height
+        for canvas in self.canvas:
+            prepare_canvas(width=width,height=height,canvas=canvas.canvas)
+        setup_overlay(width,height)
+        if scale is None:
+            scale=1
+        self.update_scale(scale)
+
+
+    def save(self):
+        self.sync_to_data()
+        state={}
+        state["width"]=self.display_width
+        state["height"]=self.display_height
+        state["selection_width"]=self.selection_size_w
+        state["selection_height"]=self.selection_size_h
+        state["view_pos"]=self.view_pos[:]
+        state["cursor"]=self.cursor[:]
+        state["scale"]=self.scale
+        keys=list(self.data.keys())
+        data={}
+        for key in keys:
+            if self.data[key].sum()>0:
+                data[f"{key[0]},{key[1]}"]=self.numpy_to_base64(self.data[key])
+        state["data"]=data
+        return json.dumps(state)
+
+    def load(self, state_json):
+        self.reset()
+        state=json.loads(state_json)
+        self.display_width=state["width"]
+        self.display_height=state["height"]
+        self.selection_size_w=state["selection_width"]
+        self.selection_size_h=state["selection_height"]
+        self.view_pos=state["view_pos"][:]
+        self.cursor=state["cursor"][:]
+        self.scale=state["scale"]
+        self.resize(state["width"],state["height"],scale=state["scale"])
+        for k,v in state["data"].items():
+            key=tuple(map(int,k.split(",")))
+            self.data[key]=self.base64_to_numpy(v)
+        self.data2buffer()
+        self.display()
+
+    def display(self):
+        self.clear_background()
+        self.draw_buffer()
+        self.draw_selection_box()
+
+    def reset(self):
+        self.data.clear()
+        self.buffer*=0
+        self.buffer_dirty=False
+        self.buffer_updated=False
+        self.sel_buffer*=0
+        self.sel_dirty=False
+        self.view_pos = [0, 0]
+        self.clear_background()
+        for i in range(1,len(self.canvas)-1):
+            self.canvas[i].clear()
+
+    def export(self):
+        self.sync_to_data()
+        xmin, xmax, ymin, ymax = 0, 0, 0, 0
+        if len(self.data.keys()) == 0:
+            return np.zeros(
+                (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+            )
+        for xi, yi in self.data.keys():
+            buf = self.data[(xi, yi)]
+            if buf.sum() > 0:
+                xmin = min(xi, xmin)
+                xmax = max(xi, xmax)
+                ymin = min(yi, ymin)
+                ymax = max(yi, ymax)
+        yn = ymax - ymin + 1
+        xn = xmax - xmin + 1
+        image = np.zeros(
+            (yn * self.patch_size, xn * self.patch_size, 4), dtype=np.uint8
+        )
+        for xi, yi in self.data.keys():
+            buf = self.data[(xi, yi)]
+            if buf.sum() > 0:
+                y0 = (yi - ymin) * self.patch_size
+                x0 = (xi - xmin) * self.patch_size
+                image[y0 : y0 + self.patch_size, x0 : x0 + self.patch_size] = buf
+        ylst, xlst = image[:, :, -1].nonzero()
+        if len(ylst) > 0:
+            yt, xt = ylst.min(), xlst.min()
+            yb, xb = ylst.max(), xlst.max()
+            image = image[yt : yb + 1, xt : xb + 1]
+            return image
+        else:
+            return np.zeros(
+                (self.selection_size_h, self.selection_size_w, 4), dtype=np.uint8
+            )

index.html

@@ -1,214 +1,411 @@
-<html>
-<head>
-<title>Stablediffusion Infinity</title>
-<meta charset="utf-8">
-<link rel="icon" type="image/x-icon" href="./favicon.png">
-<link rel="stylesheet" href="https://pyscript.net/alpha/pyscript.css" />
-<script defer src="https://pyscript.net/alpha/pyscript.js"></script>
-<style>
-#container {
-  position: relative;
-  margin:auto;
-}
-#container > canvas {
-  position: absolute;
-  top: 0;
-  left: 0;
-}
-.control {
-  display: none;
-}
-</style>
-
-</head>
-<body>
-<div>
-<button type="button" class="control" id="export">Export</button>
-<button type="button" class="control" id="outpaint">Outpaint</button>
-<button type="button" class="control" id="undo">Undo</button>
-<button type="button" class="control" id="commit">Commit</button>
-<button type="button" class="control" id="transfer">Transfer</button>
-<button type="button" class="control" id="upload">Upload</button>
-<button type="button" class="control" id="draw">Draw</button>
-<input type="text" id="mode" value="✥" class="control">
-<input type="text" id="setup" value="0" class="control">
-<textarea rows="1" id="selbuffer" name="selbuffer" class="control"></textarea>
-<fieldset class="control">
-    <div>
-      <input type="radio" id="mode0" name="mode" value="0" checked>
-      <label for="mode0">SelBox</label>
-    </div>
-    <div>
-      <input type="radio" id="mode1" name="mode" value="1">
-      <label for="mode1">Image</label>
-    </div>
-    <div>
-      <input type="radio" id="mode2" name="mode" value="2">
-      <label for="mode2">Brush</label>
-    </div>
-</fieldset>
-</div>
-<div>
-<div id = "container">
-  <canvas id = "canvas0"></canvas>
-  <canvas id = "canvas1"></canvas>
-  <canvas id = "canvas2"></canvas>
-  <canvas id = "canvas3"></canvas>
-  <canvas id = "canvas4"></canvas>
-</div>
-</div>
-<py-env>
-- numpy
-- Pillow
-- paths:
-  - ./canvas.py
-</py-env>
-
-<py-script>
-from pyodide import to_js, create_proxy
-from PIL import Image
-import io
-import time
-import base64
-import numpy as np
-from js import (
-    console,
-    document,
-    parent,
-    devicePixelRatio,
-    ImageData,
-    Uint8ClampedArray,
-    CanvasRenderingContext2D as Context2d,
-    requestAnimationFrame,
-    window
-)
-
-
-from canvas import InfCanvas
-
-
-
-base_lst = [None]
-async def draw_canvas() -> None:
-    width=1500
-    height=600
-    canvas=InfCanvas(1500,600)
-    document.querySelector("#container").style.width = f"{width}px"
-    canvas.setup_mouse()
-    canvas.clear_background()
-    canvas.draw_buffer()
-    canvas.draw_selection_box()
-    base_lst[0]=canvas
-
-async def draw_canvas_func(event):
-    try:
-        width=parent.document.querySelector("gradio-app").querySelector("#canvas_width input").value
-        height=parent.document.querySelector("gradio-app").querySelector("#canvas_height input").value
-        selection_size=parent.document.querySelector("gradio-app").querySelector("#selection_size input").value
-    except:
-        width=1024
-        height=768
-        selection_size=384
-    document.querySelector("#container").style.width = f"{width}px"
-    canvas=InfCanvas(int(width),int(height),selection_size=int(selection_size))
-    canvas.setup_mouse()
-    canvas.clear_background()
-    canvas.draw_buffer()
-    canvas.draw_selection_box()
-    base_lst[0]=canvas
-  
-async def export_func(event):
-    base=base_lst[0]
-    arr=base.export()
-    base64_str = base.numpy_to_base64(arr)
-    time_str = time.strftime("%Y%m%d_%H%M%S")
-    console.log(f"Canvas saved to outpaint_{time_str}.png")
-    link = document.createElement("a")
-    link.download = f"outpaint_{time_str}.png"
-    link.href = "data:image/png;base64,"+base64_str
-    link.click()
-
-async def outpaint_func(event):
-    base=base_lst[0]
-    base64_str=event.data[1]
-    arr=base.base64_to_numpy(base64_str)
-    base.fill_selection(arr)
-    base.draw_selection_box()
-
-async def undo_func(event):
-    base=base_lst[0]
-    if base.sel_dirty:
-        base.canvas[2].clear()
-        base.sel_buffer = base.sel_buffer_bak.copy()
-        base.sel_dirty = False
-
-async def commit_func(event):
-    base=base_lst[0]
-    if base.sel_dirty:
-        base.write_selection_to_buffer()
-
-async def transfer_func(event):
-    base=base_lst[0]
-    base.read_selection_from_buffer()
-    sel_buffer=base.sel_buffer
-    sel_buffer_str=base.numpy_to_base64(sel_buffer)
-    parent.postMessage(to_js(["transfer",str(sel_buffer_str)]),"*")
-
-async def upload_func(event):
-    base=base_lst[0]
-    base64_str=event.data[1]
-    arr=base.base64_to_numpy(base64_str)
-    h,w,_=arr.shape
-    yo=(base.height-h)//2
-    xo=(base.width-w)//2
-    if base.sel_dirty:
-        base.canvas[2].clear()
-        base.sel_buffer = base.sel_buffer_bak.copy()
-        base.sel_dirty = False
-    base.buffer_dirty=True
-    base.buffer*=0
-    base.buffer[yo:yo+h,xo:xo+w,0:3]=arr[:,:,0:3]
-    base.buffer[yo:yo+h,xo:xo+w,-1]=arr[:,:,-1]
-    base.draw_buffer()
-
-  
-
-document.querySelector("#export").addEventListener("click",create_proxy(export_func))
-document.querySelector("#undo").addEventListener("click",create_proxy(undo_func))
-document.querySelector("#commit").addEventListener("click",create_proxy(commit_func))
-document.querySelector("#outpaint").addEventListener("click",create_proxy(outpaint_func))
-document.querySelector("#upload").addEventListener("click",create_proxy(upload_func))
-
-document.querySelector("#transfer").addEventListener("click",create_proxy(transfer_func))
-document.querySelector("#draw").addEventListener("click",create_proxy(draw_canvas_func))
-
-async def setup_func():
-    document.querySelector("#setup").value="1"
-
-async def message_func(event):
-    if event.data[0]=="click":
-        if event.data[1]=="export":
-            await export_func(event)
-        elif event.data[1]=="commit":
-            await commit_func(event)
-        elif event.data[1]=="undo":
-            await undo_func(event)
-    elif event.data[0]=="upload":
-        await upload_func(event)
-    elif event.data[0]=="outpaint":
-        await outpaint_func(event)
-    elif event.data[0]=="mode":
-        document.querySelector("#mode").value=event.data[1]
-    elif event.data[0]=="transfer":
-        await transfer_func(event)
-    
-window.addEventListener("message",create_proxy(message_func))
-import asyncio
-
-_ = await asyncio.gather(
-  setup_func(),draw_canvas()
-)
-</py-script>
-
-</body>
-</html>
+<html>
+<head>
+<title>Stablediffusion Infinity</title>
+<meta charset="utf-8">
+<link rel="icon" type="image/x-icon" href="./favicon.png">
+
+<link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@master/css/w2ui.min.css">
+<script type="text/javascript" src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@master/js/w2ui.min.js"></script>
+<link rel="stylesheet" type="text/css" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.2.0/css/all.min.css">
+<script src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@master/js/fabric.min.js"></script>
+<script defer src="https://cdn.jsdelivr.net/gh/lkwq007/stablediffusion-infinity@master/js/toolbar.js"></script>
+
+<link rel="stylesheet" href="https://pyscript.net/alpha/pyscript.css" />
+<script defer src="https://pyscript.net/alpha/pyscript.js"></script>
+
+<style>
+#container {
+  position: relative;
+  margin:auto;
+  display: block;
+}
+#container > canvas {
+  position: absolute;
+  top: 0;
+  left: 0;
+}
+.control {
+  display: none;
+}
+</style>
+
+</head>
+<body>
+<div>
+<button type="button" class="control" id="export">Export</button>
+<button type="button" class="control" id="outpaint">Outpaint</button>
+<button type="button" class="control" id="undo">Undo</button>
+<button type="button" class="control" id="commit">Commit</button>
+<button type="button" class="control" id="transfer">Transfer</button>
+<button type="button" class="control" id="upload">Upload</button>
+<button type="button" class="control" id="draw">Draw</button>
+<input type="text" id="mode" value="selection" class="control">
+<input type="text" id="setup" value="0" class="control">
+<input type="text" id="upload_content" value="0" class="control">
+<textarea rows="1" id="selbuffer" name="selbuffer" class="control"></textarea>
+<fieldset class="control">
+    <div>
+      <input type="radio" id="mode0" name="mode" value="0" checked>
+      <label for="mode0">SelBox</label>
+    </div>
+    <div>
+      <input type="radio" id="mode1" name="mode" value="1">
+      <label for="mode1">Image</label>
+    </div>
+    <div>
+      <input type="radio" id="mode2" name="mode" value="2">
+      <label for="mode2">Brush</label>
+    </div>
+</fieldset>
+</div>
+<div id = "outer_container">
+<div id = "container">
+  <canvas id = "canvas0"></canvas>
+  <canvas id = "canvas1"></canvas>
+  <canvas id = "canvas2"></canvas>
+  <canvas id = "canvas3"></canvas>
+  <canvas id = "canvas4"></canvas>
+  <div id="overlay_container" style="pointer-events: none">
+    <canvas id = "overlay_canvas" width="1" height="1"></canvas>
+  </div>
+</div>
+<input type="file" name="file" id="upload_file" accept="image/*" hidden>
+<input type="file" name="state" id="upload_state" accept=".sdinf" hidden>
+<div style="position: relative;">
+<div id="toolbar" style></div>
+</div>
+</div>
+<py-env>
+- numpy
+- Pillow
+- paths:
+  - ./canvas.py
+</py-env>
+
+<py-script>
+from pyodide import to_js, create_proxy
+from PIL import Image
+import io
+import time
+import base64
+import numpy as np
+from js import (
+    console,
+    document,
+    parent,
+    devicePixelRatio,
+    ImageData,
+    Uint8ClampedArray,
+    CanvasRenderingContext2D as Context2d,
+    requestAnimationFrame,
+    window,
+    encodeURIComponent,
+    w2ui,
+    update_eraser,
+    update_scale,
+    adjust_selection,
+    update_count,
+    enable_result_lst,
+    setup_shortcut,
+)
+
+
+from canvas import InfCanvas
+
+
+
+base_lst = [None]
+async def draw_canvas() -> None:
+    width=1024
+    height=600
+    canvas=InfCanvas(1024,600)
+    update_eraser(canvas.eraser_size,min(canvas.selection_size_h,canvas.selection_size_w))
+    document.querySelector("#container").style.height= f"{height}px"
+    document.querySelector("#container").style.width = f"{width}px"
+    canvas.setup_mouse()
+    canvas.clear_background()
+    canvas.draw_buffer()
+    canvas.draw_selection_box()
+    base_lst[0]=canvas
+
+async def draw_canvas_func(event):
+    try:
+        app=parent.document.querySelector("gradio-app")
+        if app.shadowRoot:
+            app=app.shadowRoot
+        width=app.querySelector("#canvas_width input").value
+        height=app.querySelector("#canvas_height input").value
+        selection_size=app.querySelector("#selection_size input").value
+    except:
+        width=1024
+        height=768
+        selection_size=384
+    document.querySelector("#container").style.width = f"{width}px"
+    document.querySelector("#container").style.height= f"{height}px"
+    canvas=InfCanvas(int(width),int(height),selection_size=int(selection_size))
+    canvas.setup_mouse()
+    canvas.clear_background()
+    canvas.draw_buffer()
+    canvas.draw_selection_box()
+    base_lst[0]=canvas
+  
+async def export_func(event):
+    base=base_lst[0]
+    arr=base.export()
+    base.draw_buffer()
+    base.canvas[2].clear()
+    base64_str = base.numpy_to_base64(arr)
+    time_str = time.strftime("%Y%m%d_%H%M%S")
+    link = document.createElement("a")
+    if len(event.data)>2 and event.data[2]:
+        filename = event.data[2]
+    else:
+        filename = f"outpaint_{time_str}"
+    # link.download = f"sdinf_state_{time_str}.json"
+    link.download = f"{filename}.png"
+    # link.download = f"outpaint_{time_str}.png"
+    link.href = "data:image/png;base64,"+base64_str
+    link.click()
+    console.log(f"Canvas saved to {filename}.png")
+
+img_candidate_lst=[None,0]
+
+async def outpaint_func(event):
+    base=base_lst[0]
+    if len(event.data)==2:
+        app=parent.document.querySelector("gradio-app")
+        if app.shadowRoot:
+            app=app.shadowRoot
+        base64_str_raw=app.querySelector("#output textarea").value
+        base64_str_lst=base64_str_raw.split(",")
+        img_candidate_lst[0]=base64_str_lst
+        img_candidate_lst[1]=0
+    elif event.data[2]=="next":
+        img_candidate_lst[1]+=1
+    elif event.data[2]=="prev":
+        img_candidate_lst[1]-=1
+    enable_result_lst()
+    if img_candidate_lst[0] is None:
+        return
+    lst=img_candidate_lst[0]
+    idx=img_candidate_lst[1]
+    update_count(idx%len(lst)+1,len(lst))
+    arr=base.base64_to_numpy(lst[idx%len(lst)])
+    base.fill_selection(arr)
+    base.draw_selection_box()
+
+async def undo_func(event):
+    base=base_lst[0]
+    img_candidate_lst[0]=None
+    if base.sel_dirty:
+        base.sel_buffer = np.zeros((base.selection_size_h, base.selection_size_w, 4), dtype=np.uint8)
+        base.sel_dirty = False
+    base.canvas[2].clear()
+
+async def commit_func(event):
+    base=base_lst[0]
+    img_candidate_lst[0]=None
+    if base.sel_dirty:
+        base.write_selection_to_buffer()
+        base.draw_buffer()
+    base.canvas[2].clear()
+
+async def transfer_func(event):
+    base=base_lst[0]
+    base.read_selection_from_buffer()
+    sel_buffer=base.sel_buffer
+    sel_buffer_str=base.numpy_to_base64(sel_buffer)
+    app=parent.document.querySelector("gradio-app")
+    if app.shadowRoot:
+        app=app.shadowRoot
+    app.querySelector("#input textarea").value=sel_buffer_str
+    app.querySelector("#proceed").click()
+
+async def upload_func(event):
+    base=base_lst[0]
+    # base64_str=event.data[1]
+    base64_str=document.querySelector("#upload_content").value
+    base64_str=base64_str.split(",")[-1]
+    # base64_str=parent.document.querySelector("gradio-app").shadowRoot.querySelector("#upload textarea").value
+    arr=base.base64_to_numpy(base64_str)
+    h,w,c=base.buffer.shape
+    base.sync_to_buffer()
+    base.buffer_dirty=True
+    mask=arr[:,:,3:4].repeat(4,axis=2)
+    base.buffer[mask>0]=0
+    # in case mismatch
+    base.buffer[0:h,0:w,:]+=arr
+    #base.buffer[yo:yo+h,xo:xo+w,0:3]=arr[:,:,0:3]
+    #base.buffer[yo:yo+h,xo:xo+w,-1]=arr[:,:,-1]
+    base.draw_buffer()
+
+async def setup_shortcut_func(event):
+    setup_shortcut(event.data[1])
+  
+
+document.querySelector("#export").addEventListener("click",create_proxy(export_func))
+document.querySelector("#undo").addEventListener("click",create_proxy(undo_func))
+document.querySelector("#commit").addEventListener("click",create_proxy(commit_func))
+document.querySelector("#outpaint").addEventListener("click",create_proxy(outpaint_func))
+document.querySelector("#upload").addEventListener("click",create_proxy(upload_func))
+
+document.querySelector("#transfer").addEventListener("click",create_proxy(transfer_func))
+document.querySelector("#draw").addEventListener("click",create_proxy(draw_canvas_func))
+
+async def setup_func():
+    document.querySelector("#setup").value="1"
+
+async def reset_func(event):
+    base=base_lst[0]
+    base.reset()
+    
+async def load_func(event):
+    base=base_lst[0]
+    base.load(event.data[1])
+
+async def save_func(event):
+    base=base_lst[0]
+    json_str=base.save()
+    time_str = time.strftime("%Y%m%d_%H%M%S")
+    link = document.createElement("a")
+    if len(event.data)>2 and event.data[2]:
+        filename = str(event.data[2]).strip()
+    else:
+        filename = f"outpaint_{time_str}"
+    # link.download = f"sdinf_state_{time_str}.json"
+    link.download = f"{filename}.sdinf"
+    link.href = "data:text/json;charset=utf-8,"+encodeURIComponent(json_str)
+    link.click()
+
+async def prev_result_func(event):
+    base=base_lst[0]
+    base.reset()
+
+async def next_result_func(event):
+    base=base_lst[0]
+    base.reset()
+
+async def zoom_in_func(event):
+    base=base_lst[0]
+    scale=base.scale
+    if scale>=0.2:
+        scale-=0.1
+        if len(event.data)>2:
+            base.update_scale(scale,int(event.data[2]),int(event.data[3]))
+        else:
+            base.update_scale(scale)
+        scale=base.scale
+        update_scale(f"{base.width}x{base.height} ({round(100/scale)}%)")
+
+async def zoom_out_func(event):
+    base=base_lst[0]
+    scale=base.scale
+    if scale<10:
+        scale+=0.1
+        console.log(len(event.data))
+        if len(event.data)>2:
+            base.update_scale(scale,int(event.data[2]),int(event.data[3]))
+        else:
+            base.update_scale(scale)
+        scale=base.scale
+        update_scale(f"{base.width}x{base.height} ({round(100/scale)}%)")
+
+async def sync_func(event):
+    base=base_lst[0]
+    base.sync_to_buffer()
+    base.canvas[2].clear()
+
+async def eraser_size_func(event):
+    base=base_lst[0]
+    eraser_size=min(int(event.data[1]),min(base.selection_size_h,base.selection_size_w))
+    eraser_size=max(8,eraser_size)
+    base.eraser_size=eraser_size
+
+async def resize_selection_func(event):
+    base=base_lst[0]
+    cursor=base.cursor
+    if len(event.data)>3:
+        console.log(event.data)
+        base.cursor[0]=int(event.data[1])
+        base.cursor[1]=int(event.data[2])
+        base.selection_size_w=int(event.data[3])//8*8
+        base.selection_size_h=int(event.data[4])//8*8
+        base.refine_selection()
+        base.draw_selection_box()
+    elif len(event.data)>2:
+        base.draw_selection_box()
+    else:
+        base.canvas[-1].clear()
+        adjust_selection(cursor[0],cursor[1],base.selection_size_w,base.selection_size_h)
+
+async def eraser_func(event):
+    base=base_lst[0]
+    if event.data[1]!="eraser":
+        base.canvas[-2].clear()
+    else:
+        x,y=base.mouse_pos
+        base.draw_eraser(x,y)
+
+async def resize_func(event):
+    base=base_lst[0]
+    width=int(event.data[1])
+    height=int(event.data[2])
+    if width>=256 and height>=256:
+        if max(base.selection_size_h,base.selection_size_w)>min(width,height):
+            base.selection_size_h=256
+            base.selection_size_w=256
+        base.resize(width,height)
+
+async def message_func(event):
+    if event.data[0]=="click":
+        if event.data[1]=="clear":
+            await reset_func(event)
+        elif event.data[1]=="save":
+            await save_func(event)
+        elif event.data[1]=="export":
+            await export_func(event)
+        elif event.data[1]=="accept":
+            await commit_func(event)
+        elif event.data[1]=="cancel":
+            await undo_func(event)
+        elif event.data[1]=="zoom_in":
+            await zoom_in_func(event)
+        elif event.data[1]=="zoom_out":
+            await zoom_out_func(event)
+    elif event.data[0]=="sync":
+        await sync_func(event)
+    elif event.data[0]=="load":
+        await load_func(event)
+    elif event.data[0]=="upload":
+        await upload_func(event)
+    elif event.data[0]=="outpaint":
+        await outpaint_func(event)
+    elif event.data[0]=="mode":
+        if event.data[1]!="selection":
+            await sync_func(event)
+        await eraser_func(event)
+        document.querySelector("#mode").value=event.data[1]
+    elif event.data[0]=="transfer":
+        await transfer_func(event)
+    elif event.data[0]=="setup":
+        await draw_canvas_func(event)
+    elif event.data[0]=="eraser_size":
+        await eraser_size_func(event)
+    elif event.data[0]=="resize_selection":
+        await resize_selection_func(event)
+    elif event.data[0]=="shortcut":
+        await setup_shortcut_func(event)
+    elif event.data[0]=="resize":
+        await resize_func(event)
+    
+window.addEventListener("message",create_proxy(message_func))
+
+import asyncio
+
+_ = await asyncio.gather(
+  setup_func()
+)
+</py-script>
+
+</body>
+</html>

perlin2d.py

@@ -1,45 +1,45 @@
-import numpy as np
-
-##########
-# https://stackoverflow.com/questions/42147776/producing-2d-perlin-noise-with-numpy/42154921#42154921
-def perlin(x, y, seed=0):
-    # permutation table
-    np.random.seed(seed)
-    p = np.arange(256, dtype=int)
-    np.random.shuffle(p)
-    p = np.stack([p, p]).flatten()
-    # coordinates of the top-left
-    xi, yi = x.astype(int), y.astype(int)
-    # internal coordinates
-    xf, yf = x - xi, y - yi
-    # fade factors
-    u, v = fade(xf), fade(yf)
-    # noise components
-    n00 = gradient(p[p[xi] + yi], xf, yf)
-    n01 = gradient(p[p[xi] + yi + 1], xf, yf - 1)
-    n11 = gradient(p[p[xi + 1] + yi + 1], xf - 1, yf - 1)
-    n10 = gradient(p[p[xi + 1] + yi], xf - 1, yf)
-    # combine noises
-    x1 = lerp(n00, n10, u)
-    x2 = lerp(n01, n11, u)  # FIX1: I was using n10 instead of n01
-    return lerp(x1, x2, v)  # FIX2: I also had to reverse x1 and x2 here
-
-
-def lerp(a, b, x):
-    "linear interpolation"
-    return a + x * (b - a)
-
-
-def fade(t):
-    "6t^5 - 15t^4 + 10t^3"
-    return 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3
-
-
-def gradient(h, x, y):
-    "grad converts h to the right gradient vector and return the dot product with (x,y)"
-    vectors = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])
-    g = vectors[h % 4]
-    return g[:, :, 0] * x + g[:, :, 1] * y
-
-
+import numpy as np
+
+##########
+# https://stackoverflow.com/questions/42147776/producing-2d-perlin-noise-with-numpy/42154921#42154921
+def perlin(x, y, seed=0):
+    # permutation table
+    np.random.seed(seed)
+    p = np.arange(256, dtype=int)
+    np.random.shuffle(p)
+    p = np.stack([p, p]).flatten()
+    # coordinates of the top-left
+    xi, yi = x.astype(int), y.astype(int)
+    # internal coordinates
+    xf, yf = x - xi, y - yi
+    # fade factors
+    u, v = fade(xf), fade(yf)
+    # noise components
+    n00 = gradient(p[p[xi] + yi], xf, yf)
+    n01 = gradient(p[p[xi] + yi + 1], xf, yf - 1)
+    n11 = gradient(p[p[xi + 1] + yi + 1], xf - 1, yf - 1)
+    n10 = gradient(p[p[xi + 1] + yi], xf - 1, yf)
+    # combine noises
+    x1 = lerp(n00, n10, u)
+    x2 = lerp(n01, n11, u)  # FIX1: I was using n10 instead of n01
+    return lerp(x1, x2, v)  # FIX2: I also had to reverse x1 and x2 here
+
+
+def lerp(a, b, x):
+    "linear interpolation"
+    return a + x * (b - a)
+
+
+def fade(t):
+    "6t^5 - 15t^4 + 10t^3"
+    return 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3
+
+
+def gradient(h, x, y):
+    "grad converts h to the right gradient vector and return the dot product with (x,y)"
+    vectors = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])
+    g = vectors[h % 4]
+    return g[:, :, 0] * x + g[:, :, 1] * y
+
+
 ##########

postprocess.py

@@ -0,0 +1,249 @@
+"""
+https://github.com/Trinkle23897/Fast-Poisson-Image-Editing
+MIT License
+
+Copyright (c) 2022 Jiayi Weng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import time
+import argparse
+import os
+import fpie
+from process import ALL_BACKEND, CPU_COUNT, DEFAULT_BACKEND
+from fpie.io import read_images, write_image
+from process import BaseProcessor, EquProcessor, GridProcessor
+
+from PIL import Image
+import numpy as np
+import skimage
+import skimage.measure
+import scipy
+import scipy.signal
+
+
+class PhotometricCorrection:
+    def __init__(self,quite=False):
+        self.get_parser("cli")
+        args=self.parser.parse_args(["--method","grid","-g","src","-s","a","-t","a","-o","a"])
+        args.mpi_sync_interval = getattr(args, "mpi_sync_interval", 0)
+        self.backend=args.backend
+        self.args=args
+        self.quite=quite
+        proc: BaseProcessor
+        proc = GridProcessor(
+            args.gradient,
+            args.backend,
+            args.cpu,
+            args.mpi_sync_interval,
+            args.block_size,
+            args.grid_x,
+            args.grid_y,
+        )
+        print(
+            f"[PIE]Successfully initialize PIE {args.method} solver "
+            f"with {args.backend} backend"
+        )
+        self.proc=proc
+
+    def run(self, original_image, inpainted_image, mode="mask_mode"):
+        print(f"[PIE] start")
+        if mode=="disabled":
+            return inpainted_image
+        input_arr=np.array(original_image)
+        if input_arr[:,:,-1].sum()<1:
+            return inpainted_image
+        output_arr=np.array(inpainted_image)
+        mask=input_arr[:,:,-1]
+        mask=255-mask
+        if mask.sum()<1 and mode=="mask_mode":
+            mode=""
+        if mode=="mask_mode":
+            mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
+            mask = mask.repeat(8, axis=0).repeat(8, axis=1)
+        else:
+            mask[8:-9,8:-9]=255
+        mask = mask[:,:,np.newaxis].repeat(3,axis=2)
+        nmask=mask.copy()
+        output_arr2=output_arr[:,:,0:3].copy()
+        input_arr2=input_arr[:,:,0:3].copy()
+        output_arr2[nmask<128]=0
+        input_arr2[nmask>=128]=0
+        output_arr2+=input_arr2
+        src = output_arr2[:,:,0:3]
+        tgt = src.copy()
+        proc=self.proc
+        args=self.args
+        if proc.root:
+            n = proc.reset(src, mask, tgt, (args.h0, args.w0), (args.h1, args.w1))
+        proc.sync()
+        if proc.root:
+            result = tgt
+            t = time.time()
+        if args.p == 0:
+            args.p = args.n
+
+        for i in range(0, args.n, args.p):
+            if proc.root:
+                result, err = proc.step(args.p)  # type: ignore
+                print(f"[PIE] Iter {i + args.p}, abs_err {err}")
+            else:
+                proc.step(args.p)
+
+        if proc.root:
+            dt = time.time() - t
+            print(f"[PIE] Time elapsed: {dt:.4f}s")
+            # make sure consistent with dummy process
+            return Image.fromarray(result)
+
+
+    def get_parser(self,gen_type: str) -> argparse.Namespace:
+        parser = argparse.ArgumentParser()
+        parser.add_argument(
+            "-v", "--version", action="store_true", help="show the version and exit"
+        )
+        parser.add_argument(
+            "--check-backend", action="store_true", help="print all available backends"
+        )
+        if gen_type == "gui" and "mpi" in ALL_BACKEND:
+            # gui doesn't support MPI backend
+            ALL_BACKEND.remove("mpi")
+        parser.add_argument(
+            "-b",
+            "--backend",
+            type=str,
+            choices=ALL_BACKEND,
+            default=DEFAULT_BACKEND,
+            help="backend choice",
+        )
+        parser.add_argument(
+            "-c",
+            "--cpu",
+            type=int,
+            default=CPU_COUNT,
+            help="number of CPU used",
+        )
+        parser.add_argument(
+            "-z",
+            "--block-size",
+            type=int,
+            default=1024,
+            help="cuda block size (only for equ solver)",
+        )
+        parser.add_argument(
+            "--method",
+            type=str,
+            choices=["equ", "grid"],
+            default="equ",
+            help="how to parallelize computation",
+        )
+        parser.add_argument("-s", "--source", type=str, help="source image filename")
+        if gen_type == "cli":
+            parser.add_argument(
+                "-m",
+                "--mask",
+                type=str,
+                help="mask image filename (default is to use the whole source image)",
+                default="",
+            )
+        parser.add_argument("-t", "--target", type=str, help="target image filename")
+        parser.add_argument("-o", "--output", type=str, help="output image filename")
+        if gen_type == "cli":
+            parser.add_argument(
+                "-h0", type=int, help="mask position (height) on source image", default=0
+            )
+            parser.add_argument(
+                "-w0", type=int, help="mask position (width) on source image", default=0
+            )
+            parser.add_argument(
+                "-h1", type=int, help="mask position (height) on target image", default=0
+            )
+            parser.add_argument(
+                "-w1", type=int, help="mask position (width) on target image", default=0
+            )
+        parser.add_argument(
+            "-g",
+            "--gradient",
+            type=str,
+            choices=["max", "src", "avg"],
+            default="max",
+            help="how to calculate gradient for PIE",
+        )
+        parser.add_argument(
+            "-n",
+            type=int,
+            help="how many iteration would you perfer, the more the better",
+            default=5000,
+        )
+        if gen_type == "cli":
+            parser.add_argument(
+                "-p", type=int, help="output result every P iteration", default=0
+            )
+        if "mpi" in ALL_BACKEND:
+            parser.add_argument(
+                "--mpi-sync-interval",
+                type=int,
+                help="MPI sync iteration interval",
+                default=100,
+            )
+        parser.add_argument(
+            "--grid-x", type=int, help="x axis stride for grid solver", default=8
+        )
+        parser.add_argument(
+            "--grid-y", type=int, help="y axis stride for grid solver", default=8
+        )
+        self.parser=parser
+
+if __name__ =="__main__":
+    import sys
+    import io
+    import base64
+    from PIL import Image
+    def base64_to_pil(base64_str):
+        data = base64.b64decode(str(base64_str))
+        pil = Image.open(io.BytesIO(data))
+        return pil
+
+    def pil_to_base64(out_pil):
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    correction_func=PhotometricCorrection(quite=True)
+    while True:
+        buffer = sys.stdin.readline()
+        print(f"[PIE] suprocess {len(buffer)} {type(buffer)} ")
+        if len(buffer)==0:
+            break
+        if isinstance(buffer,str):
+            lst=buffer.strip().split(",")
+        else:
+            lst=buffer.decode("ascii").strip().split(",")
+        img0=base64_to_pil(lst[0])
+        img1=base64_to_pil(lst[1])
+        ret=correction_func.run(img0,img1,mode=lst[2])
+        ret_base64=pil_to_base64(ret)
+        if isinstance(buffer,str):
+            sys.stdout.write(f"{ret_base64}\n")
+        else:
+            sys.stdout.write(f"{ret_base64}\n".encode())
+        sys.stdout.flush()

process.py

@@ -0,0 +1,395 @@
+"""
+https://github.com/Trinkle23897/Fast-Poisson-Image-Editing
+MIT License
+
+Copyright (c) 2022 Jiayi Weng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+import os
+from abc import ABC, abstractmethod
+from typing import Any, Optional, Tuple
+
+import numpy as np
+
+from fpie import np_solver
+
+import scipy
+import scipy.signal
+
+CPU_COUNT = os.cpu_count() or 1
+DEFAULT_BACKEND = "numpy"
+ALL_BACKEND = ["numpy"]
+
+try:
+  from fpie import numba_solver
+  ALL_BACKEND += ["numba"]
+  DEFAULT_BACKEND = "numba"
+except ImportError:
+  numba_solver = None  # type: ignore
+
+try:
+  from fpie import taichi_solver
+  ALL_BACKEND += ["taichi-cpu", "taichi-gpu"]
+  DEFAULT_BACKEND = "taichi-cpu"
+except ImportError:
+  taichi_solver = None  # type: ignore
+
+# try:
+#   from fpie import core_gcc  # type: ignore
+#   DEFAULT_BACKEND = "gcc"
+#   ALL_BACKEND.append("gcc")
+# except ImportError:
+#   core_gcc = None
+
+# try:
+#   from fpie import core_openmp  # type: ignore
+#   DEFAULT_BACKEND = "openmp"
+#   ALL_BACKEND.append("openmp")
+# except ImportError:
+#   core_openmp = None
+
+# try:
+#   from mpi4py import MPI
+
+#   from fpie import core_mpi  # type: ignore
+#   ALL_BACKEND.append("mpi")
+# except ImportError:
+#   MPI = None  # type: ignore
+#   core_mpi = None
+
+try:
+  from fpie import core_cuda  # type: ignore
+  DEFAULT_BACKEND = "cuda"
+  ALL_BACKEND.append("cuda")
+except ImportError:
+  core_cuda = None
+
+
+class BaseProcessor(ABC):
+  """API definition for processor class."""
+
+  def __init__(
+    self, gradient: str, rank: int, backend: str, core: Optional[Any]
+  ):
+    if core is None:
+      error_msg = {
+        "numpy":
+          "Please run `pip install numpy`.",
+        "numba":
+          "Please run `pip install numba`.",
+        "gcc":
+          "Please install cmake and gcc in your operating system.",
+        "openmp":
+          "Please make sure your gcc is compatible with `-fopenmp` option.",
+        "mpi":
+          "Please install MPI and run `pip install mpi4py`.",
+        "cuda":
+          "Please make sure nvcc and cuda-related libraries are available.",
+        "taichi":
+          "Please run `pip install taichi`.",
+      }
+      print(error_msg[backend.split("-")[0]])
+
+      raise AssertionError(f"Invalid backend {backend}.")
+
+    self.gradient = gradient
+    self.rank = rank
+    self.backend = backend
+    self.core = core
+    self.root = rank == 0
+
+  def mixgrad(self, a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    if self.gradient == "src":
+      return a
+    if self.gradient == "avg":
+      return (a + b) / 2
+    # mix gradient, see Equ. 12 in PIE paper
+    mask = np.abs(a) < np.abs(b)
+    a[mask] = b[mask]
+    return a
+
+  @abstractmethod
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    pass
+
+  def sync(self) -> None:
+    self.core.sync()
+
+  @abstractmethod
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    pass
+
+
+class EquProcessor(BaseProcessor):
+  """PIE Jacobi equation processor."""
+
+  def __init__(
+    self,
+    gradient: str = "max",
+    backend: str = DEFAULT_BACKEND,
+    n_cpu: int = CPU_COUNT,
+    min_interval: int = 100,
+    block_size: int = 1024,
+  ):
+    core: Optional[Any] = None
+    rank = 0
+
+    if backend == "numpy":
+      core = np_solver.EquSolver()
+    elif backend == "numba" and numba_solver is not None:
+      core = numba_solver.EquSolver()
+    elif backend == "gcc":
+      core = core_gcc.EquSolver()
+    elif backend == "openmp" and core_openmp is not None:
+      core = core_openmp.EquSolver(n_cpu)
+    elif backend == "mpi" and core_mpi is not None:
+      core = core_mpi.EquSolver(min_interval)
+      rank = MPI.COMM_WORLD.Get_rank()
+    elif backend == "cuda" and core_cuda is not None:
+      core = core_cuda.EquSolver(block_size)
+    elif backend.startswith("taichi") and taichi_solver is not None:
+      core = taichi_solver.EquSolver(backend, n_cpu, block_size)
+
+    super().__init__(gradient, rank, backend, core)
+
+  def mask2index(
+    self, mask: np.ndarray
+  ) -> Tuple[np.ndarray, int, np.ndarray, np.ndarray]:
+    x, y = np.nonzero(mask)
+    max_id = x.shape[0] + 1
+    index = np.zeros((max_id, 3))
+    ids = self.core.partition(mask)
+    ids[mask == 0] = 0  # reserve id=0 for constant
+    index = ids[x, y].argsort()
+    return ids, max_id, x[index], y[index]
+
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    assert self.root
+    # check validity
+    # assert 0 <= mask_on_src[0] and 0 <= mask_on_src[1]
+    # assert mask_on_src[0] + mask.shape[0] <= src.shape[0]
+    # assert mask_on_src[1] + mask.shape[1] <= src.shape[1]
+    # assert mask_on_tgt[0] + mask.shape[0] <= tgt.shape[0]
+    # assert mask_on_tgt[1] + mask.shape[1] <= tgt.shape[1]
+
+    if len(mask.shape) == 3:
+      mask = mask.mean(-1)
+    mask = (mask >= 128).astype(np.int32)
+
+    # zero-out edge
+    mask[0] = 0
+    mask[-1] = 0
+    mask[:, 0] = 0
+    mask[:, -1] = 0
+
+    x, y = np.nonzero(mask)
+    x0, x1 = x.min() - 1, x.max() + 2
+    y0, y1 = y.min() - 1, y.max() + 2
+    mask_on_src = (x0 + mask_on_src[0], y0 + mask_on_src[1])
+    mask_on_tgt = (x0 + mask_on_tgt[0], y0 + mask_on_tgt[1])
+    mask = mask[x0:x1, y0:y1]
+    ids, max_id, index_x, index_y = self.mask2index(mask)
+
+    src_x, src_y = index_x + mask_on_src[0], index_y + mask_on_src[1]
+    tgt_x, tgt_y = index_x + mask_on_tgt[0], index_y + mask_on_tgt[1]
+
+    src_C = src[src_x, src_y].astype(np.float32)
+    src_U = src[src_x - 1, src_y].astype(np.float32)
+    src_D = src[src_x + 1, src_y].astype(np.float32)
+    src_L = src[src_x, src_y - 1].astype(np.float32)
+    src_R = src[src_x, src_y + 1].astype(np.float32)
+    tgt_C = tgt[tgt_x, tgt_y].astype(np.float32)
+    tgt_U = tgt[tgt_x - 1, tgt_y].astype(np.float32)
+    tgt_D = tgt[tgt_x + 1, tgt_y].astype(np.float32)
+    tgt_L = tgt[tgt_x, tgt_y - 1].astype(np.float32)
+    tgt_R = tgt[tgt_x, tgt_y + 1].astype(np.float32)
+
+    grad = self.mixgrad(src_C - src_L, tgt_C - tgt_L) \
+      + self.mixgrad(src_C - src_R, tgt_C - tgt_R) \
+      + self.mixgrad(src_C - src_U, tgt_C - tgt_U) \
+      + self.mixgrad(src_C - src_D, tgt_C - tgt_D)
+
+    A = np.zeros((max_id, 4), np.int32)
+    X = np.zeros((max_id, 3), np.float32)
+    B = np.zeros((max_id, 3), np.float32)
+
+    X[1:] = tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1]]
+    # four-way
+    A[1:, 0] = ids[index_x - 1, index_y]
+    A[1:, 1] = ids[index_x + 1, index_y]
+    A[1:, 2] = ids[index_x, index_y - 1]
+    A[1:, 3] = ids[index_x, index_y + 1]
+    B[1:] = grad
+    m = (mask[index_x - 1, index_y] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0] - 1, index_y + mask_on_tgt[1]]
+    m = (mask[index_x, index_y - 1] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1] - 1]
+    m = (mask[index_x, index_y + 1] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0], index_y + mask_on_tgt[1] + 1]
+    m = (mask[index_x + 1, index_y] == 0).astype(float).reshape(-1, 1)
+    B[1:] += m * tgt[index_x + mask_on_tgt[0] + 1, index_y + mask_on_tgt[1]]
+
+    self.tgt = tgt.copy()
+    self.tgt_index = (index_x + mask_on_tgt[0], index_y + mask_on_tgt[1])
+    self.core.reset(max_id, A, X, B)
+    return max_id
+
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    result = self.core.step(iteration)
+    if self.root:
+      x, err = result
+      self.tgt[self.tgt_index] = x[1:]
+      return self.tgt, err
+    return None
+
+
+class GridProcessor(BaseProcessor):
+  """PIE grid processor."""
+
+  def __init__(
+    self,
+    gradient: str = "max",
+    backend: str = DEFAULT_BACKEND,
+    n_cpu: int = CPU_COUNT,
+    min_interval: int = 100,
+    block_size: int = 1024,
+    grid_x: int = 8,
+    grid_y: int = 8,
+  ):
+    core: Optional[Any] = None
+    rank = 0
+
+    if backend == "numpy":
+      core = np_solver.GridSolver()
+    elif backend == "numba" and numba_solver is not None:
+      core = numba_solver.GridSolver()
+    elif backend == "gcc":
+      core = core_gcc.GridSolver(grid_x, grid_y)
+    elif backend == "openmp" and core_openmp is not None:
+      core = core_openmp.GridSolver(grid_x, grid_y, n_cpu)
+    elif backend == "mpi" and core_mpi is not None:
+      core = core_mpi.GridSolver(min_interval)
+      rank = MPI.COMM_WORLD.Get_rank()
+    elif backend == "cuda" and core_cuda is not None:
+      core = core_cuda.GridSolver(grid_x, grid_y)
+    elif backend.startswith("taichi") and taichi_solver is not None:
+      core = taichi_solver.GridSolver(
+        grid_x, grid_y, backend, n_cpu, block_size
+      )
+
+    super().__init__(gradient, rank, backend, core)
+
+  def reset(
+    self,
+    src: np.ndarray,
+    mask: np.ndarray,
+    tgt: np.ndarray,
+    mask_on_src: Tuple[int, int],
+    mask_on_tgt: Tuple[int, int],
+  ) -> int:
+    assert self.root
+    # check validity
+    # assert 0 <= mask_on_src[0] and 0 <= mask_on_src[1]
+    # assert mask_on_src[0] + mask.shape[0] <= src.shape[0]
+    # assert mask_on_src[1] + mask.shape[1] <= src.shape[1]
+    # assert mask_on_tgt[0] + mask.shape[0] <= tgt.shape[0]
+    # assert mask_on_tgt[1] + mask.shape[1] <= tgt.shape[1]
+
+    if len(mask.shape) == 3:
+      mask = mask.mean(-1)
+    mask = (mask >= 128).astype(np.int32)
+
+    # zero-out edge
+    mask[0] = 0
+    mask[-1] = 0
+    mask[:, 0] = 0
+    mask[:, -1] = 0
+
+    x, y = np.nonzero(mask)
+    x0, x1 = x.min() - 1, x.max() + 2
+    y0, y1 = y.min() - 1, y.max() + 2
+    mask = mask[x0:x1, y0:y1]
+    max_id = np.prod(mask.shape)
+
+    src_crop = src[mask_on_src[0] + x0:mask_on_src[0] + x1,
+                   mask_on_src[1] + y0:mask_on_src[1] + y1].astype(np.float32)
+    tgt_crop = tgt[mask_on_tgt[0] + x0:mask_on_tgt[0] + x1,
+                   mask_on_tgt[1] + y0:mask_on_tgt[1] + y1].astype(np.float32)
+    grad = np.zeros([*mask.shape, 3], np.float32)
+    grad[1:] += self.mixgrad(
+      src_crop[1:] - src_crop[:-1], tgt_crop[1:] - tgt_crop[:-1]
+    )
+    grad[:-1] += self.mixgrad(
+      src_crop[:-1] - src_crop[1:], tgt_crop[:-1] - tgt_crop[1:]
+    )
+    grad[:, 1:] += self.mixgrad(
+      src_crop[:, 1:] - src_crop[:, :-1], tgt_crop[:, 1:] - tgt_crop[:, :-1]
+    )
+    grad[:, :-1] += self.mixgrad(
+      src_crop[:, :-1] - src_crop[:, 1:], tgt_crop[:, :-1] - tgt_crop[:, 1:]
+    )
+
+    grad[mask == 0] = 0
+    if True:
+        kernel = [[1] * 3 for _ in range(3)]
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res = scipy.signal.convolve2d(
+            nmask, kernel, mode="same", boundary="fill", fillvalue=1
+        )
+        res[nmask < 1] = 0
+        res[res == 9] = 0
+        res[res > 0] = 1
+        grad[res>0]=0
+        # ylst, xlst = res.nonzero()
+        # for y, x in zip(ylst, xlst):
+        #     grad[y,x]=0
+            # for yi in range(-1,2):
+                # for xi in range(-1,2):
+                    # grad[y+yi,x+xi]=0
+    self.x0 = mask_on_tgt[0] + x0
+    self.x1 = mask_on_tgt[0] + x1
+    self.y0 = mask_on_tgt[1] + y0
+    self.y1 = mask_on_tgt[1] + y1
+    self.tgt = tgt.copy()
+    self.core.reset(max_id, mask, tgt_crop, grad)
+    return max_id
+
+  def step(self, iteration: int) -> Optional[Tuple[np.ndarray, np.ndarray]]:
+    result = self.core.step(iteration)
+    if self.root:
+      tgt, err = result
+      self.tgt[self.x0:self.x1, self.y0:self.y1] = tgt
+      return self.tgt, err
+    return None

utils.py

@@ -1,151 +1,263 @@
-from PIL import Image
-from PIL import ImageFilter
-import cv2
-import numpy as np
-import scipy
-import scipy.signal
-from scipy.spatial import cKDTree
-
-import os
-from perlin2d import *
-
-patch_match_compiled = True
-if os.name != "nt":
-    try:
-        from PyPatchMatch import patch_match
-    except Exception as e:
-        import patch_match
-
-try:
-    patch_match
-except NameError:
-    print("patch_match compiling failed")
-    patch_match_compiled = False
-
-
-
-
-def edge_pad(img, mask, mode=1):
-    if mode == 0:
-        nmask = mask.copy()
-        nmask[nmask > 0] = 1
-        res0 = 1 - nmask
-        res1 = nmask
-        p0 = np.stack(res0.nonzero(), axis=0).transpose()
-        p1 = np.stack(res1.nonzero(), axis=0).transpose()
-        min_dists, min_dist_idx = cKDTree(p1).query(p0, 1)
-        loc = p1[min_dist_idx]
-        for (a, b), (c, d) in zip(p0, loc):
-            img[a, b] = img[c, d]
-    elif mode == 1:
-        record = {}
-        kernel = [[1] * 3 for _ in range(3)]
-        nmask = mask.copy()
-        nmask[nmask > 0] = 1
-        res = scipy.signal.convolve2d(
-            nmask, kernel, mode="same", boundary="fill", fillvalue=1
-        )
-        res[nmask < 1] = 0
-        res[res == 9] = 0
-        res[res > 0] = 1
-        ylst, xlst = res.nonzero()
-        queue = [(y, x) for y, x in zip(ylst, xlst)]
-        # bfs here
-        cnt = res.astype(np.float32)
-        acc = img.astype(np.float32)
-        step = 1
-        h = acc.shape[0]
-        w = acc.shape[1]
-        offset = [(1, 0), (-1, 0), (0, 1), (0, -1)]
-        while queue:
-            target = []
-            for y, x in queue:
-                val = acc[y][x]
-                for yo, xo in offset:
-                    yn = y + yo
-                    xn = x + xo
-                    if 0 <= yn < h and 0 <= xn < w and nmask[yn][xn] < 1:
-                        if record.get((yn, xn), step) == step:
-                            acc[yn][xn] = acc[yn][xn] * cnt[yn][xn] + val
-                            cnt[yn][xn] += 1
-                            acc[yn][xn] /= cnt[yn][xn]
-                            if (yn, xn) not in record:
-                                record[(yn, xn)] = step
-                                target.append((yn, xn))
-            step += 1
-            queue = target
-        img = acc.astype(np.uint8)
-    else:
-        nmask = mask.copy()
-        ylst, xlst = nmask.nonzero()
-        yt, xt = ylst.min(), xlst.min()
-        yb, xb = ylst.max(), xlst.max()
-        content = img[yt : yb + 1, xt : xb + 1]
-        img = np.pad(
-            content,
-            ((yt, mask.shape[0] - yb - 1), (xt, mask.shape[1] - xb - 1), (0, 0)),
-            mode="edge",
-        )
-    return img, mask
-
-
-def perlin_noise(img, mask):
-    lin = np.linspace(0, 5, mask.shape[0], endpoint=False)
-    x, y = np.meshgrid(lin, lin)
-    avg = img.mean(axis=0).mean(axis=0)
-    # noise=[((perlin(x, y)+1)*128+avg[i]).astype(np.uint8) for i in range(3)]
-    noise = [((perlin(x, y) + 1) * 0.5 * 255).astype(np.uint8) for i in range(3)]
-    noise = np.stack(noise, axis=-1)
-    # mask=skimage.measure.block_reduce(mask,(8,8),np.min)
-    # mask=mask.repeat(8, axis=0).repeat(8, axis=1)
-    # mask_image=Image.fromarray(mask)
-    # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 4))
-    # mask=np.array(mask_image)
-    nmask = mask.copy()
-    # nmask=nmask/255.0
-    nmask[mask > 0] = 1
-    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
-    # img=img.astype(np.uint8)
-    return img, mask
-
-
-def gaussian_noise(img, mask):
-    noise = np.random.randn(mask.shape[0], mask.shape[1], 3)
-    noise = (noise + 1) / 2 * 255
-    noise = noise.astype(np.uint8)
-    nmask = mask.copy()
-    nmask[mask > 0] = 1
-    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
-    return img, mask
-
-
-def cv2_telea(img, mask):
-    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_TELEA)
-    return ret, mask
-
-
-def cv2_ns(img, mask):
-    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_NS)
-    return ret, mask
-
-
-def patch_match_func(img, mask):
-    ret = patch_match.inpaint(img, mask=255 - mask, patch_size=3)
-    return ret, mask
-
-
-def mean_fill(img, mask):
-    avg = img.mean(axis=0).mean(axis=0)
-    img[mask < 1] = avg
-    return img, mask
-
-
-functbl = {
-    "gaussian": gaussian_noise,
-    "perlin": perlin_noise,
-    "edge_pad": edge_pad,
-    "patchmatch": patch_match_func if (os.name != "nt" and patch_match_compiled) else edge_pad,
-    "cv2_ns": cv2_ns,
-    "cv2_telea": cv2_telea,
-    "mean_fill": mean_fill,
-}
+from PIL import Image
+from PIL import ImageFilter
+import cv2
+import numpy as np
+import scipy
+import scipy.signal
+from scipy.spatial import cKDTree
+
+import os
+from perlin2d import *
+
+patch_match_compiled = True
+
+try:
+    from PyPatchMatch import patch_match
+except Exception as e:
+    try:
+        import patch_match
+    except Exception as e:
+        patch_match_compiled = False
+
+try:
+    patch_match
+except NameError:
+    print("patch_match compiling failed, will fall back to edge_pad")
+    patch_match_compiled = False
+
+
+
+
+def edge_pad(img, mask, mode=1):
+    if mode == 0:
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res0 = 1 - nmask
+        res1 = nmask
+        p0 = np.stack(res0.nonzero(), axis=0).transpose()
+        p1 = np.stack(res1.nonzero(), axis=0).transpose()
+        min_dists, min_dist_idx = cKDTree(p1).query(p0, 1)
+        loc = p1[min_dist_idx]
+        for (a, b), (c, d) in zip(p0, loc):
+            img[a, b] = img[c, d]
+    elif mode == 1:
+        record = {}
+        kernel = [[1] * 3 for _ in range(3)]
+        nmask = mask.copy()
+        nmask[nmask > 0] = 1
+        res = scipy.signal.convolve2d(
+            nmask, kernel, mode="same", boundary="fill", fillvalue=1
+        )
+        res[nmask < 1] = 0
+        res[res == 9] = 0
+        res[res > 0] = 1
+        ylst, xlst = res.nonzero()
+        queue = [(y, x) for y, x in zip(ylst, xlst)]
+        # bfs here
+        cnt = res.astype(np.float32)
+        acc = img.astype(np.float32)
+        step = 1
+        h = acc.shape[0]
+        w = acc.shape[1]
+        offset = [(1, 0), (-1, 0), (0, 1), (0, -1)]
+        while queue:
+            target = []
+            for y, x in queue:
+                val = acc[y][x]
+                for yo, xo in offset:
+                    yn = y + yo
+                    xn = x + xo
+                    if 0 <= yn < h and 0 <= xn < w and nmask[yn][xn] < 1:
+                        if record.get((yn, xn), step) == step:
+                            acc[yn][xn] = acc[yn][xn] * cnt[yn][xn] + val
+                            cnt[yn][xn] += 1
+                            acc[yn][xn] /= cnt[yn][xn]
+                            if (yn, xn) not in record:
+                                record[(yn, xn)] = step
+                                target.append((yn, xn))
+            step += 1
+            queue = target
+        img = acc.astype(np.uint8)
+    else:
+        nmask = mask.copy()
+        ylst, xlst = nmask.nonzero()
+        yt, xt = ylst.min(), xlst.min()
+        yb, xb = ylst.max(), xlst.max()
+        content = img[yt : yb + 1, xt : xb + 1]
+        img = np.pad(
+            content,
+            ((yt, mask.shape[0] - yb - 1), (xt, mask.shape[1] - xb - 1), (0, 0)),
+            mode="edge",
+        )
+    return img, mask
+
+
+def perlin_noise(img, mask):
+    lin = np.linspace(0, 5, mask.shape[0], endpoint=False)
+    x, y = np.meshgrid(lin, lin)
+    avg = img.mean(axis=0).mean(axis=0)
+    # noise=[((perlin(x, y)+1)*128+avg[i]).astype(np.uint8) for i in range(3)]
+    noise = [((perlin(x, y) + 1) * 0.5 * 255).astype(np.uint8) for i in range(3)]
+    noise = np.stack(noise, axis=-1)
+    # mask=skimage.measure.block_reduce(mask,(8,8),np.min)
+    # mask=mask.repeat(8, axis=0).repeat(8, axis=1)
+    # mask_image=Image.fromarray(mask)
+    # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 4))
+    # mask=np.array(mask_image)
+    nmask = mask.copy()
+    # nmask=nmask/255.0
+    nmask[mask > 0] = 1
+    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
+    # img=img.astype(np.uint8)
+    return img, mask
+
+
+def gaussian_noise(img, mask):
+    noise = np.random.randn(mask.shape[0], mask.shape[1], 3)
+    noise = (noise + 1) / 2 * 255
+    noise = noise.astype(np.uint8)
+    nmask = mask.copy()
+    nmask[mask > 0] = 1
+    img = nmask[:, :, np.newaxis] * img + (1 - nmask[:, :, np.newaxis]) * noise
+    return img, mask
+
+
+def cv2_telea(img, mask):
+    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_TELEA)
+    return ret, mask
+
+
+def cv2_ns(img, mask):
+    ret = cv2.inpaint(img, 255 - mask, 5, cv2.INPAINT_NS)
+    return ret, mask
+
+
+def patch_match_func(img, mask):
+    ret = patch_match.inpaint(img, mask=255 - mask, patch_size=3)
+    return ret, mask
+
+
+def mean_fill(img, mask):
+    avg = img.mean(axis=0).mean(axis=0)
+    img[mask < 1] = avg
+    return img, mask
+
+"""
+Apache-2.0 license
+https://github.com/hafriedlander/stable-diffusion-grpcserver/blob/main/sdgrpcserver/services/generate.py
+https://github.com/parlance-zz/g-diffuser-bot/tree/g-diffuser-bot-beta2
+_handleImageAdjustment
+"""
+if True:
+    from sd_grpcserver.sdgrpcserver import images
+    import torch
+    from math import sqrt
+    def handleImageAdjustment(array, adjustments):
+        tensor = images.fromPIL(Image.fromarray(array))
+        for adjustment in adjustments:
+            which = adjustment[0]
+
+            if which == "blur":
+                sigma = adjustment[1]
+                direction = adjustment[2]
+
+                if direction == "DOWN" or direction == "UP":
+                    orig = tensor
+                    repeatCount=256
+                    sigma /= sqrt(repeatCount)
+
+                    for _ in range(repeatCount):
+                        tensor = images.gaussianblur(tensor, sigma)
+                        if direction == "DOWN":
+                            tensor = torch.minimum(tensor, orig)
+                        else:
+                            tensor = torch.maximum(tensor, orig)
+                else:
+                    tensor = images.gaussianblur(tensor, adjustment.blur.sigma)
+            elif which == "invert":
+                tensor = images.invert(tensor)
+            elif which == "levels":
+                tensor = images.levels(tensor, adjustment[1], adjustment[2], adjustment[3], adjustment[4])
+            elif which == "channels":
+                tensor = images.channelmap(tensor, [adjustment.channels.r,  adjustment.channels.g,  adjustment.channels.b,  adjustment.channels.a])
+            elif which == "rescale":
+                self.unimp("Rescale")
+            elif which == "crop":
+                tensor = images.crop(tensor, adjustment.crop.top, adjustment.crop.left, adjustment.crop.height, adjustment.crop.width)
+        return np.array(images.toPIL(tensor)[0])
+
+def g_diffuser(img,mask):
+    adjustments=[["blur",32,"UP"],["level",0,0.05,0,1]]
+    mask=handleImageAdjustment(mask,adjustments)
+    out_mask=handleImageAdjustment(mask,adjustments)
+    return img, mask, out_mask
+def dummy_fill(img,mask):
+    return img,mask
+functbl = {
+    "gaussian": gaussian_noise,
+    "perlin": perlin_noise,
+    "edge_pad": edge_pad,
+    "patchmatch": patch_match_func if patch_match_compiled else edge_pad,
+    "cv2_ns": cv2_ns,
+    "cv2_telea": cv2_telea,
+    "g_diffuser": g_diffuser,
+    "g_diffuser_lib": dummy_fill,
+}
+
+try:
+    from postprocess import PhotometricCorrection
+    correction_func = PhotometricCorrection()
+except Exception as e:
+    print(e, "so PhotometricCorrection is disabled")
+    class DummyCorrection:
+        def __init__(self):
+            self.backend=""
+            pass
+        def run(self,a,b,**kwargs):
+            return b
+    correction_func=DummyCorrection()
+
+if "taichi" in correction_func.backend:
+    import sys
+    import io
+    import base64
+    from PIL import Image
+    def base64_to_pil(base64_str):
+        data = base64.b64decode(str(base64_str))
+        pil = Image.open(io.BytesIO(data))
+        return pil
+
+    def pil_to_base64(out_pil):
+        out_buffer = io.BytesIO()
+        out_pil.save(out_buffer, format="PNG")
+        out_buffer.seek(0)
+        base64_bytes = base64.b64encode(out_buffer.read())
+        base64_str = base64_bytes.decode("ascii")
+        return base64_str
+    from subprocess import Popen, PIPE, STDOUT
+    class SubprocessCorrection:
+        def __init__(self):
+            self.backend=correction_func.backend
+            self.child= Popen(["python", "postprocess.py"], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
+        def run(self,img_input,img_inpainted,mode):
+            if mode=="disabled":
+                return img_inpainted
+            base64_str_input = pil_to_base64(img_input)
+            base64_str_inpainted = pil_to_base64(img_inpainted)
+            try:
+                if self.child.poll():
+                    self.child= Popen(["python", "postprocess.py"], stdin=PIPE, stdout=PIPE, stderr=STDOUT)
+                self.child.stdin.write(f"{base64_str_input},{base64_str_inpainted},{mode}\n".encode())
+                self.child.stdin.flush()
+                out = self.child.stdout.readline()
+                base64_str=out.decode().strip()
+                while base64_str and base64_str[0]=="[":
+                    print(base64_str)
+                    out = self.child.stdout.readline()
+                    base64_str=out.decode().strip()
+                ret=base64_to_pil(base64_str)
+            except:
+                print("[PIE] not working, photometric correction is disabled")
+                ret=img_inpainted
+            return ret
+    correction_func = SubprocessCorrection()