Update backend to the new Inpainting model

app.py

@@ -1,427 +1,422 @@
-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 numpy as np
+import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
+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 = DiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-inpainting",
+            use_auth_token=token,
+        ).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,
+                    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()
+