Initial commit

app.py

@@ -0,0 +1,736 @@
+import gradio as gr
+
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from collections import Counter
+import math
+
+from gradio import processing_utils
+from typing import Optional
+
+import warnings
+
+from datetime import datetime
+
+import torch
+from PIL import Image
+import matplotlib.pyplot as plt
+from diffusers import StableDiffusionInpaintPipeline
+from accelerate.utils import set_seed
+
+class Instance:
+    def __init__(self, capacity = 2):
+        self.model_type = 'base'
+        self.loaded_model_list = {}
+        self.counter = Counter()
+        self.global_counter = Counter()
+        self.capacity = capacity
+
+        self.loaded_model = None
+
+    def _log(self, model_type, batch_size, instruction, phrase_list):
+        self.counter[model_type] += 1
+        self.global_counter[model_type] += 1
+        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        print('[{}] Current: {}, All: {}. Samples: {}, prompt: {}, phrases: {}'.format(
+            current_time, dict(self.counter), dict(self.global_counter), batch_size, instruction, phrase_list
+        ))
+
+    def get_model(self):
+        if self.loaded_model is None:
+            self.loaded_model = self.load_model()
+        return self.loaded_model
+    
+    def load_model(self, model_id='j-min/IterInpaint-CLEVR'):
+        pipe = StableDiffusionInpaintPipeline.from_pretrained(model_id)
+        pipe = pipe.to('cuda')
+        def dummy(images, **kwargs):
+            return images, False
+        pipe.safety_checker = dummy
+        print("Disabled safety checker")
+
+        print("Loaded model")
+        return pipe
+
+instance = Instance()
+
+# from ldm.viz_utils import plot_results, fig2img, show_images
+from gen_utils import encode_from_custom_annotation, iterinpaint_sample_diffusers
+
+class ImageMask(gr.components.Image):
+    """
+    Sets: source="canvas", tool="sketch"
+    """
+
+    is_template = True
+
+    def __init__(self, **kwargs):
+        super().__init__(source="upload", tool="sketch", interactive=True, **kwargs)
+
+    def preprocess(self, x):
+        if x is None:
+            return x
+        if self.tool == "sketch" and self.source in ["upload", "webcam"] and type(x) != dict:
+            decode_image = processing_utils.decode_base64_to_image(x)
+            width, height = decode_image.size
+            mask = np.zeros((height, width, 4), dtype=np.uint8)
+            mask[..., -1] = 255
+            mask = self.postprocess(mask)
+            x = {'image': x, 'mask': mask}
+        return super().preprocess(x)
+
+
+class Blocks(gr.Blocks):
+
+    def __init__(
+        self,
+        theme: str = "default",
+        analytics_enabled: Optional[bool] = None,
+        mode: str = "blocks",
+        title: str = "Gradio",
+        css: Optional[str] = None,
+        **kwargs,
+    ):
+
+        self.extra_configs = {
+            'thumbnail': kwargs.pop('thumbnail', ''),
+            'url': kwargs.pop('url', 'https://gradio.app/'),
+            'creator': kwargs.pop('creator', '@teamGradio'),
+        }
+
+        super(Blocks, self).__init__(
+            theme, analytics_enabled, mode, title, css, **kwargs)
+        warnings.filterwarnings("ignore")
+
+    def get_config_file(self):
+        config = super(Blocks, self).get_config_file()
+
+        for k, v in self.extra_configs.items():
+            config[k] = v
+
+        return config
+
+def draw_box(boxes=[], texts=[], img=None):
+    if len(boxes) == 0 and img is None:
+        return None
+
+    if img is None:
+        img = Image.new('RGB', (512, 512), (255, 255, 255))
+    colors = ["red", "olive", "blue", "green", "orange", "brown", "cyan", "purple"]
+    draw = ImageDraw.Draw(img)
+    font = ImageFont.truetype("DejaVuSansMono.ttf", size=20)
+    for bid, box in enumerate(boxes):
+        draw.rectangle([box[0], box[1], box[2], box[3]], outline=colors[bid % len(colors)], width=4)
+        anno_text = texts[bid]
+        draw.rectangle([box[0], box[3] - int(font.size * 1.2), box[0] + int((len(anno_text) + 0.8) * font.size * 0.6), box[3]], outline=colors[bid % len(colors)], fill=colors[bid % len(colors)], width=4)
+        draw.text([box[0] + int(font.size * 0.2), box[3] - int(font.size*1.2)], anno_text, font=font, fill=(255,255,255))
+    return img
+
+def get_concat(ims):
+    if len(ims) == 1:
+        n_col = 1
+    else:
+        n_col = 2
+    n_row = math.ceil(len(ims) / 2)
+    dst = Image.new('RGB', (ims[0].width * n_col, ims[0].height * n_row), color="white")
+    for i, im in enumerate(ims):
+        row_id = i // n_col
+        col_id = i % n_col
+        dst.paste(im, (im.width * col_id, im.height * row_id))
+    return dst
+
+
+def inference(language_instruction, grounding_texts, boxes, guidance_scale):
+
+    # custom_annotations = [
+    #     {'x': 19,
+    #      'y': 61,
+    #      'width': 158,
+    #      'height': 169,
+    #      'label': 'blue metal cube'},
+    #     {'x': 183,
+    #      'y': 94,
+    #      'width': 103,
+    #      'height': 109,
+    #      'label': 'brown rubber sphere'},
+    # ]
+
+    # # boxes - normalized -> unnormalized
+    # boxes = np.array(boxes) * 512
+
+    custom_annotations = []
+    for i in range(len(boxes)):
+        box = boxes[i]
+        custom_annotations.append({'x': box[0],
+                                   'y': box[1],
+                                   'width': box[2] - box[0],
+                                   'height': box[3] - box[1],
+                                   'label': grounding_texts[i]})
+    # # 1) convert xywh to xyxy
+    # # 2) normalize coordinates
+    scene = encode_from_custom_annotation(custom_annotations, size=512)
+     
+    print(scene['box_captions'])
+    print(scene['boxes_normalized'])
+
+    pipe = instance.get_model()
+
+    out = iterinpaint_sample_diffusers(
+        pipe, scene, paste=True, verbose=True, size=512, guidance_scale=guidance_scale)
+    
+    final_image = out['generated_images'][-1].copy()
+    
+    # Create Generation GIF
+    prompts = out['prompts']
+
+    fps = 4
+
+    def create_gif_source_images(images, prompts):
+        """Create source images for gif
+        Each frame consists of a intermediate image with a prompt as title.
+        Don't change size of the original images.
+        """
+
+        step_images = []
+        font = ImageFont.truetype("DejaVuSansMono.ttf", size=20)
+        for i, img in enumerate(images):
+            draw = ImageDraw.Draw(img)
+            draw.text((0, 0), prompts[i], (255, 255, 255), font=font)
+            step_images.append(img)
+        return step_images
+    
+    import imageio
+
+    step_images = create_gif_source_images(out['generated_images'], prompts)
+    print("Number of frames in GIF: {}".format(len(step_images)))
+    # create temp path
+    import tempfile
+    import os
+    gif_save_path = os.path.join(tempfile.gettempdir(), 'gen.gif')
+
+    # create gif
+    imageio.mimsave(gif_save_path, step_images, fps=fps)
+    print('GIF saved to {}'.format(gif_save_path))
+
+    out_images = [
+        final_image,
+        gif_save_path
+    ]
+    
+    return out_images
+    
+def generate(task, language_instruction, grounding_texts, sketch_pad,
+             alpha_sample, guidance_scale, batch_size,
+             fix_seed, rand_seed, use_actual_mask, append_grounding, style_cond_image,
+             state):
+    if 'boxes' not in state:
+        state['boxes'] = []
+
+    boxes = state['boxes']
+    grounding_texts = [x.strip() for x in grounding_texts.split(';')]
+    # assert len(boxes) == len(grounding_texts)
+    if len(boxes) != len(grounding_texts):
+        if len(boxes) < len(grounding_texts):
+            raise ValueError("""The number of boxes should be equal to the number of grounding objects.
+Number of boxes drawn: {}, number of grounding tokens: {}.
+Please draw boxes accordingly on the sketch pad.""".format(len(boxes), len(grounding_texts)))
+        grounding_texts = grounding_texts + [""] * (len(boxes) - len(grounding_texts))
+
+    # # normalize boxes
+    # boxes = (np.asarray(boxes) / 512).tolist()
+
+    print('input boxes: ', boxes)
+    print('input grounding_texts: ', grounding_texts)
+    print('input language instruction: ', language_instruction)
+
+    if fix_seed:
+        set_seed(rand_seed)
+        print('seed set to: ', rand_seed)
+
+    gen_image, gen_animation = inference(
+        language_instruction, grounding_texts, boxes,
+        guidance_scale=guidance_scale,
+    )
+
+    # for idx, gen_image in enumerate(gen_images):
+
+    #     if task == 'Grounded Inpainting' and state.get('inpaint_hw', None):
+    #         hw = min(*state['original_image'].shape[:2])
+    #         gen_image = sized_center_fill(state['original_image'].copy(), np.array(gen_image.resize((hw, hw))), hw, hw)
+    #         gen_image = Image.fromarray(gen_image)
+        
+    #     gen_images[idx] = gen_image
+
+    # blank_samples = batch_size % 2 if batch_size > 1 else 0
+    # gen_images = [gr.Image.update(value=x, visible=True) for i,x in enumerate(gen_images)] \
+    #                 + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
+    #                 + [gr.Image.update(value=None, visible=False) for _ in range(4 - batch_size - blank_samples)]
+
+    # gen_images = [gr.Image.update(value=x, visible=True) for i,x in enumerate(gen_images)] \
+    #                 + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
+
+    gen_images = [
+        gr.Image.update(value=gen_image, visible=True),
+        gr.Image.update(value=gen_animation, visible=True)
+    ]
+
+    return gen_images + [state]
+
+
+def binarize(x):
+    return (x != 0).astype('uint8') * 255
+
+def sized_center_crop(img, cropx, cropy):
+    y, x = img.shape[:2]
+    startx = x // 2 - (cropx // 2)
+    starty = y // 2 - (cropy // 2)    
+    return img[starty:starty+cropy, startx:startx+cropx]
+
+def sized_center_fill(img, fill, cropx, cropy):
+    y, x = img.shape[:2]
+    startx = x // 2 - (cropx // 2)
+    starty = y // 2 - (cropy // 2)    
+    img[starty:starty+cropy, startx:startx+cropx] = fill
+    return img
+
+def sized_center_mask(img, cropx, cropy):
+    y, x = img.shape[:2]
+    startx = x // 2 - (cropx // 2)
+    starty = y // 2 - (cropy // 2)    
+    center_region = img[starty:starty+cropy, startx:startx+cropx].copy()
+    img = (img * 0.2).astype('uint8')
+    img[starty:starty+cropy, startx:startx+cropx] = center_region
+    return img
+
+def center_crop(img, HW=None, tgt_size=(512, 512)):
+    if HW is None:
+        H, W = img.shape[:2]
+        HW = min(H, W)
+    img = sized_center_crop(img, HW, HW)
+    img = Image.fromarray(img)
+    img = img.resize(tgt_size)
+    return np.array(img)
+
+def draw(task, input, grounding_texts, new_image_trigger, state):
+    if type(input) == dict:
+        image = input['image']
+        mask = input['mask']
+    else:
+        mask = input
+
+    if mask.ndim == 3:
+        mask = mask[..., 0]
+
+    image_scale = 1.0
+
+    # resize trigger
+    if task == "Grounded Inpainting":
+        mask_cond = mask.sum() == 0
+        # size_cond = mask.shape != (512, 512)
+        if mask_cond and 'original_image' not in state:
+            image = Image.fromarray(image)
+            width, height = image.size
+            scale = 600 / min(width, height)
+            image = image.resize((int(width * scale), int(height * scale)))
+            state['original_image'] = np.array(image).copy()
+            image_scale = float(height / width)
+            return [None, new_image_trigger + 1, image_scale, state]
+        else:
+            original_image = state['original_image']
+            H, W = original_image.shape[:2]
+            image_scale = float(H / W)
+
+    mask = binarize(mask)
+    if mask.shape != (512, 512):
+        # assert False, "should not receive any non- 512x512 masks."
+        if 'original_image' in state and state['original_image'].shape[:2] == mask.shape:
+            mask = center_crop(mask, state['inpaint_hw'])
+            image = center_crop(state['original_image'], state['inpaint_hw'])
+        else:
+            mask = np.zeros((512, 512), dtype=np.uint8)
+    # mask = center_crop(mask)
+    mask = binarize(mask)
+
+    if type(mask) != np.ndarray:
+        mask = np.array(mask)
+
+    if mask.sum() == 0 and task != "Grounded Inpainting":
+        state = {}
+
+    if task != 'Grounded Inpainting':
+        image = None
+    else:
+        image = Image.fromarray(image)
+
+    if 'boxes' not in state:
+        state['boxes'] = []
+
+    if 'masks' not in state or len(state['masks']) == 0:
+        state['masks'] = []
+        last_mask = np.zeros_like(mask)
+    else:
+        last_mask = state['masks'][-1]
+
+    if type(mask) == np.ndarray and mask.size > 1:
+        diff_mask = mask - last_mask
+    else:
+        diff_mask = np.zeros([])
+
+    if diff_mask.sum() > 0:
+        x1x2 = np.where(diff_mask.max(0) != 0)[0]
+        y1y2 = np.where(diff_mask.max(1) != 0)[0]
+        y1, y2 = y1y2.min(), y1y2.max()
+        x1, x2 = x1x2.min(), x1x2.max()
+
+        if (x2 - x1 > 5) and (y2 - y1 > 5):
+            state['masks'].append(mask.copy())
+            state['boxes'].append((x1, y1, x2, y2))
+
+    grounding_texts = [x.strip() for x in grounding_texts.split(';')]
+    grounding_texts = [x for x in grounding_texts if len(x) > 0]
+    if len(grounding_texts) < len(state['boxes']):
+        grounding_texts += [f'Obj. {bid+1}' for bid in range(len(grounding_texts), len(state['boxes']))]
+
+    box_image = draw_box(state['boxes'], grounding_texts, image)
+
+    if box_image is not None and state.get('inpaint_hw', None):
+        inpaint_hw = state['inpaint_hw']
+        box_image_resize = np.array(box_image.resize((inpaint_hw, inpaint_hw)))
+        original_image = state['original_image'].copy()
+        box_image = sized_center_fill(original_image, box_image_resize, inpaint_hw, inpaint_hw)
+
+    return [box_image, new_image_trigger, image_scale, state]
+
+def clear(task, sketch_pad_trigger, batch_size, state, switch_task=False):
+    if task != 'Grounded Inpainting':
+        sketch_pad_trigger = sketch_pad_trigger + 1
+    blank_samples = batch_size % 2 if batch_size > 1 else 0
+    # out_images = [gr.Image.update(value=None, visible=True) for i in range(batch_size)] \
+    #                 + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
+    #                 + [gr.Image.update(value=None, visible=False) for _ in range(4 - batch_size - blank_samples)]
+
+    out_images = [gr.Image.update(value=None, visible=True) for i in range(1)] \
+                    + [gr.Image.update(value=None, visible=True) for _ in range(1)]
+    state = {}
+    return [None, sketch_pad_trigger, None, 1.0] + out_images + [state]
+
+css = """
+#img2img_image, #img2img_image > .fixed-height, #img2img_image > .fixed-height > div, #img2img_image > .fixed-height > div > img
+{
+    height: var(--height) !important;
+    max-height: var(--height) !important;
+    min-height: var(--height) !important;
+}
+#paper-info a {
+    color:#008AD7;
+    text-decoration: none;
+}
+#paper-info a:hover {
+    cursor: pointer;
+    text-decoration: none;
+}
+"""
+
+rescale_js = """
+function(x) {
+    const root = document.querySelector('gradio-app').shadowRoot || document.querySelector('gradio-app');
+    let image_scale = parseFloat(root.querySelector('#image_scale input').value) || 1.0;
+    const image_width = root.querySelector('#img2img_image').clientWidth;
+    const target_height = parseInt(image_width * image_scale);
+    document.body.style.setProperty('--height', `${target_height}px`);
+    root.querySelectorAll('button.justify-center.rounded')[0].style.display='none';
+    root.querySelectorAll('button.justify-center.rounded')[1].style.display='none';
+    return x;
+}
+"""
+
+with Blocks(
+    # css=css,
+    analytics_enabled=False,
+    title="IterInpaint demo",
+) as main:
+    description = """
+    <p style="text-align: center; font-weight: bold;">
+        <span style="font-size: 28px">IterInpaint CLEVR Demo</span>
+        <br>
+        <span style="font-size: 18px" id="paper-info">
+            [<a href="https://layoutbench.github.io" target="_blank">Project Page</a>]
+            [<a href="https://arxiv.org/abs/2304.06671" target="_blank">Paper</a>]
+            [<a href="https://github.com/j-min/IterInpaint" target="_blank">GitHub</a>]
+        </span>
+    </p>
+    <span style="font-size: 14px">
+    <b>IterInpaint</b> is a new baseline for layout-guided image generation. 
+    Unlike previous methods that generate all objects in a single step, IterInpaint decomposes image generation process into multiple steps and uses an inpainting model to update regions step-by-step.
+    </span>
+    <br>
+    <br>
+    <span style="font-size: 18px" id="instruction">
+        Instructions:
+    </span>
+    <p>
+        (1) &#9000;&#65039; Enter the object names in <em> Region Captions</em>
+        <br>
+        (2) &#128433;&#65039; Draw their corresponding bounding boxes one by one using <em> Sketch Pad</em> -- the parsed boxes will be displayed automatically.
+        <br>
+        For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. <a href="https://huggingface.co/spaces/j-min/iterinpaint-CLEVR?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a>
+    </p>
+    """
+    gr.HTML(description)
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            sketch_pad_trigger = gr.Number(value=0, visible=False)
+            sketch_pad_resize_trigger = gr.Number(value=0, visible=False)
+            init_white_trigger = gr.Number(value=0, visible=False)
+            image_scale = gr.Number(
+                value=0, elem_id="image_scale", visible=False)
+            new_image_trigger = gr.Number(value=0, visible=False)
+
+            # task = gr.Radio(
+            #     choices=["Grounded Generation", 'Grounded Inpainting'],
+            #     type="value",
+            #     value="Grounded Generation",
+            #     label="Task",
+            # )
+            task = gr.State("Grounded Generation")
+            # language_instruction = gr.Textbox(
+            #     label="Language instruction",
+            # )
+            language_instruction = gr.State("")
+
+            grounding_instruction = gr.Textbox(
+                label="""
+                Region Captions (Separated by semicolon)
+                e.g., "blue metal cube; red rubber cylinder"
+                """,
+            )
+            with gr.Row():
+                sketch_pad = ImageMask(
+                    label="Draw bounding boxes", elem_id="img2img_image")
+                out_imagebox = gr.Image(type="pil", label="Parsed Layout")
+            with gr.Row():
+                clear_btn = gr.Button(value='Clear')
+                gen_btn = gr.Button(value='Generate')
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Column():
+                    # alpha_sample = gr.Slider(
+                    #     minimum=0, maximum=1.0, step=0.1, value=0.3, label="Scheduled Sampling (τ)")
+                    alpha_sample = gr.State(0.3)
+                    guidance_scale = gr.Slider(
+                        minimum=0, maximum=50, step=0.5, value=4.0, label="Guidance Scale")
+                    # batch_size = gr.Slider(
+                    #     minimum=1, maximum=4, step=1, value=2, label="Number of Samples")
+                    # batch_size = gr.Slider(
+                    #     minimum=1, maximum=1, step=1, value=1, label="Number of Samples")
+                    batch_size = gr.State(1)
+                    # append_grounding = gr.Checkbox(
+                    #     value=True, label="Append grounding instructions to the caption")
+                    append_grounding = gr.State(False)
+                    # use_actual_mask = gr.Checkbox(
+                    #     value=False, label="Use actual mask for inpainting", visible=False)
+                    use_actual_mask = gr.State(False)
+                    with gr.Row():
+                        # fix_seed = gr.Checkbox(value=True, label="Fixed seed")
+                        fix_seed = gr.State(True)
+                        rand_seed = gr.Slider(
+                            minimum=0, maximum=1000, step=1, value=0, label="Seed")
+                    with gr.Row():
+                        # use_style_cond = gr.Checkbox(
+                        #     value=False, label="Enable Style Condition")
+                        # style_cond_image = gr.Image(
+                        #     type="pil", label="Style Condition", visible=False, interactive=True)
+                        use_style_cond = gr.State(False)
+                        style_cond_image = gr.State(None)
+        with gr.Column(scale=3):
+            gr.HTML(
+                '<span style="font-size: 20px; font-weight: bold">Generated Image</span>')
+            # with gr.Row():
+            out_gen_1 = gr.Image(
+                type="pil", visible=True, show_label=False)
+            gr.HTML(
+                '<span style="font-size: 20px; font-weight: bold">Step-by-Step Animation</span>')
+            out_gen_2 = gr.Image(
+                type="pil", visible=True, show_label=False)
+            # with gr.Row():
+            #     out_gen_3 = gr.Image(
+            #         type="pil", visible=False, show_label=False)
+            #     out_gen_4 = gr.Image(
+            #         type="pil", visible=False, show_label=False)
+
+        state = gr.State({})
+
+        class Controller:
+            def __init__(self):
+                self.calls = 0
+                self.tracks = 0
+                self.resizes = 0
+                self.scales = 0
+
+            def init_white(self, init_white_trigger):
+                self.calls += 1
+                return np.ones((512, 512), dtype='uint8') * 255, 1.0, init_white_trigger+1
+
+            # def change_n_samples(self, n_samples):
+                # blank_samples = n_samples % 2 if n_samples > 1 else 0
+                # return [gr.Image.update(visible=True) for _ in range(n_samples + blank_samples)] \
+                #     + [gr.Image.update(visible=False)
+                #        for _ in range(4 - n_samples - blank_samples)]
+
+            def resize_centercrop(self, state):
+                self.resizes += 1
+                image = state['original_image'].copy()
+                inpaint_hw = int(0.9 * min(*image.shape[:2]))
+                state['inpaint_hw'] = inpaint_hw
+                image_cc = center_crop(image, inpaint_hw)
+                # print(f'resize triggered {self.resizes}', image.shape, '->', image_cc.shape)
+                return image_cc, state
+
+            def resize_masked(self, state):
+                self.resizes += 1
+                image = state['original_image'].copy()
+                inpaint_hw = int(0.9 * min(*image.shape[:2]))
+                state['inpaint_hw'] = inpaint_hw
+                image_mask = sized_center_mask(image, inpaint_hw, inpaint_hw)
+                state['masked_image'] = image_mask.copy()
+                # print(f'mask triggered {self.resizes}')
+                return image_mask, state
+
+            def switch_task_hide_cond(self, task):
+                cond = False
+                if task == "Grounded Generation":
+                    cond = True
+
+                return gr.Checkbox.update(visible=cond, value=False), gr.Image.update(value=None, visible=False), gr.Slider.update(visible=cond), gr.Checkbox.update(visible=(not cond), value=False)
+
+        controller = Controller()
+        main.load(
+            lambda x: x+1,
+            inputs=sketch_pad_trigger,
+            outputs=sketch_pad_trigger,
+            queue=False)
+        sketch_pad.edit(
+            draw,
+            inputs=[task, sketch_pad, grounding_instruction,
+                    sketch_pad_resize_trigger, state],
+            outputs=[out_imagebox, sketch_pad_resize_trigger,
+                     image_scale, state],
+            queue=False,
+        )
+        grounding_instruction.change(
+            draw,
+            inputs=[task, sketch_pad, grounding_instruction,
+                    sketch_pad_resize_trigger, state],
+            outputs=[out_imagebox, sketch_pad_resize_trigger,
+                     image_scale, state],
+            queue=False,
+        )
+        clear_btn.click(
+            clear,
+            inputs=[task, sketch_pad_trigger, batch_size, state],
+            outputs=[sketch_pad, sketch_pad_trigger, out_imagebox,
+                    #  image_scale, out_gen_1, out_gen_2, out_gen_3, out_gen_4, state],
+                     image_scale, out_gen_1, out_gen_2, state],
+            queue=False)
+        # task.change(
+        #     partial(clear, switch_task=True),
+        #     inputs=[task, sketch_pad_trigger, batch_size, state],
+        #     outputs=[sketch_pad, sketch_pad_trigger, out_imagebox,
+        #              image_scale, out_gen_1, out_gen_2, out_gen_3, out_gen_4, state],
+        #     queue=False)
+        sketch_pad_trigger.change(
+            controller.init_white,
+            inputs=[init_white_trigger],
+            outputs=[sketch_pad, image_scale, init_white_trigger],
+            queue=False)
+        sketch_pad_resize_trigger.change(
+            controller.resize_masked,
+            inputs=[state],
+            outputs=[sketch_pad, state],
+            queue=False)
+        # batch_size.change(
+        #     controller.change_n_samples,
+        #     inputs=[batch_size],
+        #     outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4],
+        #     queue=False)
+        gen_btn.click(
+            generate,
+            inputs=[
+                task, language_instruction, grounding_instruction, sketch_pad,
+                alpha_sample, guidance_scale, batch_size,
+                fix_seed, rand_seed,
+                use_actual_mask,
+                append_grounding, style_cond_image,
+                state,
+            ],
+            # outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4, state],
+            outputs=[out_gen_1, out_gen_2, state],
+            queue=True
+        )
+        sketch_pad_resize_trigger.change(
+            None,
+            None,
+            sketch_pad_resize_trigger,
+            _js=rescale_js,
+            queue=False)
+        init_white_trigger.change(
+            None,
+            None,
+            init_white_trigger,
+            _js=rescale_js,
+            queue=False)
+        # use_style_cond.change(
+        #     lambda cond: gr.Image.update(visible=cond),
+        #     use_style_cond,
+        #     style_cond_image,
+        #     queue=False)
+        # task.change(
+        #     controller.switch_task_hide_cond,
+        #     inputs=task,
+        #     outputs=[use_style_cond, style_cond_image,
+        #              alpha_sample, use_actual_mask],
+        #     queue=False)
+
+    with gr.Column():
+        gr.Examples(
+            examples=[
+                [
+                    "images/blank.png",
+                    "blue metal cube",
+                ],
+                [
+                    "images/blank.png",
+                    "green metal cube; red metal sphere; brown rubber cube",
+                ],
+                [
+                    "images/blank.png",
+                    "blue metal cube; brown rubber sphere; gray metal sphere; yellow rubber cylinder; gray metal cylinder; cyan rubber sphere; green rubber cube; red metal cylinder",
+                ]
+            ],
+            inputs=[
+                sketch_pad,
+                grounding_instruction
+            ],
+            outputs=None,
+            fn=None,
+            cache_examples=False,
+        )
+
+    # https://huggingface.co/spaces/gligen/demo
+    # add hyperlink
+    thank_desc = """
+    Thanks
+    <a href="https://huggingface.co/spaces/gligen/demo" target="_blank">GLIGEN demo</a>, for providing bounding box parsing module.
+    """
+    gr.HTML(thank_desc)
+
+main.queue(concurrency_count=1, api_open=False)
+# main.launch(share=False, show_api=False, show_error=True)
+main.launch(
+    server_name="0.0.0.0",
+    share=True,
+    # server_port=7864,
+    show_api=False, show_error=True
+    )

gen_utils.py

@@ -0,0 +1,208 @@
+import torch
+from PIL import Image
+from PIL import ImageDraw
+
+def encode_scene(obj_list, H=320, W=320, src_bbox_format='xywh', tgt_bbox_format='xyxy'):
+    """Encode scene into text and bounding boxes
+    Args:
+        obj_list: list of dicts
+            Each dict has keys:
+                
+                'color': str
+                'material': str
+                'shape': str
+                or 
+                'caption': str
+
+                and
+
+                'bbox': list of 4 floats (unnormalized)
+                    [x0, y0, x1, y1] or [x0, y0, w, h]
+    """
+    box_captions = []
+    for obj in obj_list:
+        if 'caption' in obj:
+            box_caption = obj['caption']
+        else:
+            box_caption = f"{obj['color']} {obj['material']} {obj['shape']}"
+        box_captions += [box_caption]
+    
+    assert src_bbox_format in ['xywh', 'xyxy'], f"src_bbox_format must be 'xywh' or 'xyxy', not {src_bbox_format}"
+    assert tgt_bbox_format in ['xywh', 'xyxy'], f"tgt_bbox_format must be 'xywh' or 'xyxy', not {tgt_bbox_format}"
+
+    boxes_unnormalized = []
+    boxes_normalized = []
+    for obj in obj_list:
+        if src_bbox_format == 'xywh':
+            x0, y0, w, h = obj['bbox']
+            x1 = x0 + w
+            y1 = y0 + h
+        elif src_bbox_format == 'xyxy':
+            x0, y0, x1, y1 = obj['bbox']
+            w = x1 - x0
+            h = y1 - y0
+        assert x1 > x0, f"x1={x1} <= x0={x0}"
+        assert y1 > y0, f"y1={y1} <= y0={y0}"
+        assert x1 <= W, f"x1={x1} > W={W}"
+        assert y1 <= H, f"y1={y1} > H={H}"
+
+        if tgt_bbox_format == 'xywh':
+            bbox_unnormalized = [x0, y0, w, h]
+            bbox_normalized = [x0 / W, y0 / H, w / W, h / H]
+
+        elif tgt_bbox_format == 'xyxy':
+            bbox_unnormalized = [x0, y0, x1, y1]
+            bbox_normalized = [x0 / W, y0 / H, x1 / W, y1 / H]
+            
+        boxes_unnormalized += [bbox_unnormalized]
+        boxes_normalized += [bbox_normalized]
+
+    assert len(box_captions) == len(boxes_normalized), f"len(box_captions)={len(box_captions)} != len(boxes_normalized)={len(boxes_normalized)}"
+        
+    
+    out = {}
+    out['box_captions'] = box_captions
+    out['boxes_normalized'] = boxes_normalized
+    out['boxes_unnormalized'] = boxes_unnormalized
+        
+    return out
+
+def encode_from_custom_annotation(custom_annotations, size=512):
+    #     custom_annotations = [
+    #     {'x': 83, 'y': 335, 'width': 70, 'height': 69, 'label': 'blue metal cube'},
+    #     {'x': 162, 'y': 302, 'width': 110, 'height': 138, 'label': 'blue metal cube'},
+    #     {'x': 274, 'y': 250, 'width': 191, 'height': 234, 'label': 'blue metal cube'},
+    #     {'x': 14, 'y': 18, 'width': 155, 'height': 205, 'label': 'blue metal cube'},
+    #     {'x': 175, 'y': 79, 'width': 106, 'height': 119, 'label': 'blue metal cube'},
+    #     {'x': 288, 'y': 111, 'width': 69, 'height': 63, 'label': 'blue metal cube'}
+    # ]
+    H, W = size, size
+
+    objects = []
+    for j in range(len(custom_annotations)):
+        xyxy = [
+            custom_annotations[j]['x'],
+            custom_annotations[j]['y'],
+            custom_annotations[j]['x'] + custom_annotations[j]['width'],
+            custom_annotations[j]['y'] + custom_annotations[j]['height']]
+        objects.append({
+            'caption': custom_annotations[j]['label'],
+            'bbox': xyxy,
+        })
+
+    out = encode_scene(objects, H=H, W=W,
+                       src_bbox_format='xyxy', tgt_bbox_format='xyxy')
+
+    return out
+
+
+
+#### Below are for HF diffusers
+
+def iterinpaint_sample_diffusers(pipe, datum, paste=True, verbose=False, guidance_scale=4.0, size=512, background_instruction='Add gray background'):
+    d = datum
+
+    d['unnormalized_boxes'] = d['boxes_unnormalized']
+    
+    n_total_boxes = len(d['unnormalized_boxes'])
+
+    context_imgs = []
+    mask_imgs = []
+    # masked_imgs = []
+    generated_images = []
+    prompts = []
+
+    context_img = Image.new('RGB', (size, size))
+    # context_draw = ImageDraw.Draw(context_img)
+    if verbose:
+        print('Initiailzed context image')
+
+    background_mask_img = Image.new('L', (size, size))
+    background_mask_draw = ImageDraw.Draw(background_mask_img)
+    background_mask_draw.rectangle([(0, 0), background_mask_img.size], fill=255)
+
+    for i in range(n_total_boxes):
+        if verbose:
+            print('Iter: ', i+1, 'total: ', n_total_boxes)
+
+        target_caption = d['box_captions'][i]
+        if verbose:
+            print('Drawing ', target_caption)
+
+        mask_img = Image.new('L', context_img.size)
+        mask_draw = ImageDraw.Draw(mask_img)
+        mask_draw.rectangle([(0, 0), mask_img.size], fill=0)
+
+        box = d['unnormalized_boxes'][i]
+        if type(box) == list:
+            box = torch.tensor(box) 
+        mask_draw.rectangle(box.long().tolist(), fill=255)
+        background_mask_draw.rectangle(box.long().tolist(), fill=0)
+
+        mask_imgs.append(mask_img.copy())
+
+        
+        prompt = f"Add {d['box_captions'][i]}"
+
+        if verbose:
+            print('prompt:', prompt)
+        prompts += [prompt]
+
+        context_imgs.append(context_img.copy())
+
+        generated_image = pipe(
+            prompt,
+            context_img,
+            mask_img,
+            guidance_scale=guidance_scale).images[0]
+        
+        if paste:
+            # context_img.paste(generated_image.crop(box.long().tolist()), box.long().tolist())
+            
+
+            src_box = box.long().tolist()
+
+            # x1 -> x1 + 1
+            # y1 -> y1 + 1
+            paste_box = box.long().tolist()
+            paste_box[0] -= 1
+            paste_box[1] -= 1
+            paste_box[2] += 1
+            paste_box[3] += 1
+
+            box_w = paste_box[2] - paste_box[0]
+            box_h = paste_box[3] - paste_box[1]
+
+            context_img.paste(generated_image.crop(src_box).resize((box_w, box_h)), paste_box)
+            generated_images.append(context_img.copy())
+        else:
+            context_img = generated_image
+            generated_images.append(context_img.copy())
+
+    if verbose:
+        print('Fill background')
+
+    mask_img = background_mask_img
+
+    mask_imgs.append(mask_img)
+
+    prompt = background_instruction
+
+    if verbose:
+        print('prompt:', prompt)
+    prompts += [prompt]
+
+    generated_image = pipe(
+        prompt,
+        context_img,
+        mask_img,
+        guidance_scale=guidance_scale).images[0]
+
+    generated_images.append(generated_image)
+    
+    return {
+        'context_imgs': context_imgs,
+        'mask_imgs': mask_imgs,
+        'prompts': prompts,
+        'generated_images': generated_images,
+    }

requirements.txt

@@ -0,0 +1,32 @@
+
+transformers
+diffusers
+accelerate
+opencv-python
+timm
+einops
+datasets
+rouge_score
+omegaconf
+ftfy
+pycocotools
+pycocoevalcap
+albumentations
+pudb
+imageio
+imageio-ffmpeg
+clean-fid
+h5py
+pillow
+setuptools
+pytorch-lightning==1.5.9
+tensorboardX==2.4.1
+test-tube>=0.7.5
+streamlit>=0.73.1
+torch-fidelity==0.3.0
+torchmetrics==0.6.0
+kornia==0.6
+chardet
+cchardet
+taming-transformers-rom1504
+git+https://github.com/openai/CLIP.git