Update app.py

app.py

@@ -3,9 +3,13 @@ import re
 import zipfile
 import torch
 import gradio as gr
+
+print('hello', gr.__version__)
+
+import numpy as np
 import time
 from transformers import CLIPTextModel, CLIPTokenizer
-from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel, DiffusionPipeline, LCMScheduler
+from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel, DiffusionPipeline
 from tqdm import tqdm
 from PIL import Image
 from PIL import Image, ImageDraw, ImageFont
@@ -26,29 +30,9 @@ if not os.path.exists('images2'):
 # os.system('nvidia-smi')
 os.system('ls')
 
-#### import m1
-from fastchat.model import load_model, get_conversation_template
-from transformers import AutoTokenizer, AutoModelForCausalLM
-m1_model_path = 'JingyeChen22/textdiffuser2_layout_planner'
-# m1_model, m1_tokenizer = load_model(
-#     m1_model_path,
-#     'cuda',
-#     1,
-#     None,
-#     False,
-#     False,
-#     revision="main",
-#     debug=False,
-# )
-
-m1_tokenizer = AutoTokenizer.from_pretrained(m1_model_path, use_fast=False)
-m1_model = AutoModelForCausalLM.from_pretrained(
-    m1_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
-).cuda()
-
 #### import diffusion models
 text_encoder = CLIPTextModel.from_pretrained(
-    'JingyeChen22/textdiffuser2-full-ft', subfolder="text_encoder"
+    'JingyeChen22/textdiffuser2-full-ft-inpainting', subfolder="text_encoder"
 ).cuda().half()
 tokenizer = CLIPTokenizer.from_pretrained(
     'runwayml/stable-diffusion-v1-5', subfolder="tokenizer"
@@ -69,50 +53,49 @@ print('***************')
 
 vae = AutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="vae").half().cuda()
 unet = UNet2DConditionModel.from_pretrained(
-    'JingyeChen22/textdiffuser2-full-ft', subfolder="unet"
+    'JingyeChen22/textdiffuser2-full-ft-inpainting', subfolder="unet"
 ).half().cuda()
 text_encoder.resize_token_embeddings(len(tokenizer))
 
+global_dict = {}
+#### for interactive
+# stack = []
+# state = 0   
+font = ImageFont.truetype("./Arial.ttf", 20)
 
-#### load lcm components
-model_id = "lambdalabs/sd-pokemon-diffusers"
-lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
-pipe = DiffusionPipeline.from_pretrained(model_id, unet=copy.deepcopy(unet), tokenizer=tokenizer, text_encoder=copy.deepcopy(text_encoder), torch_dtype=torch.float16)
-pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.load_lora_weights(lcm_lora_id)
-pipe.to(device="cuda")
+def skip_fun(i, t, guest_id):
+    global_dict[guest_id]['state'] = 0
+    # global state
+    # state = 0
 
 
-#### for interactive
-stack = []
-state = 0   
-font = ImageFont.truetype("./Arial.ttf", 32)
-
-def skip_fun(i, t):
-    global state
-    state = 0
-
-
-def exe_undo(i, t):
-    global stack
-    global state
-    state = 0
-    stack = []
-    image = Image.open(f'./gray256.jpg')
-    print('stack', stack)
-    return image
+def exe_undo(i, orig_i, t, guest_id):
+
+    global_dict[guest_id]['stack'] = []
+    global_dict[guest_id]['state'] = 0
+
+    return copy.deepcopy(orig_i)
+
+
+def exe_redo(i, orig_i, t, guest_id):
 
+    print('redo ',orig_i)
 
-def exe_redo(i, t):
-    global state 
-    state = 0
+    if type(orig_i) == str:
+        orig_i = Image.open(orig_i)
 
-    if len(stack) > 0:
-        stack.pop()
-    image = Image.open(f'./gray256.jpg')
+    # global state 
+    # state = 0
+    global_dict[guest_id]['state'] = 0
+
+    if len(global_dict[guest_id]['stack']) > 0:
+        global_dict[guest_id]['stack'].pop()
+
+    image = copy.deepcopy(orig_i)
+    
     draw = ImageDraw.Draw(image)
 
-    for items in stack:
+    for items in global_dict[guest_id]['stack']:
         # print('now', items)
         text_position, t = items
         if len(text_position) == 2:
@@ -133,57 +116,194 @@ def exe_redo(i, t):
             draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
             draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
 
-    print('stack', stack)
+    print('stack', global_dict[guest_id]['stack'])
     return image
 
-def get_pixels(i, t, evt: gr.SelectData):
-    global state
+def get_pixels(i, orig_i, radio, t, guest_id, evt: gr.SelectData):  
 
-    text_position = evt.index
+    print('hi1 ', i)
+    print('hi2 ', orig_i)
 
-    if state == 0:
-        stack.append(
-            (text_position, t)
-        )
-        print(text_position, stack)
-        state = 1
+    width, height = Image.open(i).size
+
+    # register
+    if guest_id == '-1': # register for the first time
+        seed = str(int(time.time()))
+        global_dict[str(seed)] = {
+            'state': 0,
+            'stack': [],
+            'image_id': [list(Image.open(i).resize((512,512)).getdata())] # an image has been recorded
+        }
+        guest_id = str(seed)
+    else:
+        seed = guest_id
+
+    if type(i) == str:
+        i = Image.open(i)
+        i = i.resize((512,512))
+
+    images = global_dict[str(seed)]['image_id']
+    flag = False
+    for image in images:
+        if image == list(i.getdata()):
+            print('find it')
+            flag = True
+            break
+    
+    if not flag:
+        global_dict[str(seed)]['image_id'] = [list(i.getdata())]
+        global_dict[str(seed)]['stack'] = []
+        global_dict[str(seed)]['state'] = 0
+        orig_i = i
     else:
-        
-        (_, t) = stack.pop()
-        x, y = _
-        stack.append(
-            ((x,y,text_position[0],text_position[1]), t)
-        )
-        state = 0
 
+        if orig_i is not None: 
+            orig_i = Image.open(orig_i)
+            orig_i = orig_i.resize((512,512))
+        else: 
+            orig_i = i 
+            global_dict[guest_id]['stack'] = []
+            global_dict[guest_id]['state'] = 0
 
-    image = Image.open(f'./gray256.jpg')
-    draw = ImageDraw.Draw(image)
+    text_position = evt.index
 
-    for items in stack:
-        # print('now', items)
-        text_position, t = items
-        if len(text_position) == 2:
-            x, y = text_position
-            text_color = (255, 0, 0)  
-            draw.text((x+2, y), t, font=font, fill=text_color)
-            r = 4
-            leftUpPoint = (x-r, y-r)
-            rightDownPoint = (x+r, y+r)
-            draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
-        elif len(text_position) == 4:
-            x0, y0, x1, y1 = text_position
-            text_color = (255, 0, 0)  
-            draw.text((x0+2, y0), t, font=font, fill=text_color)
-            r = 4
-            leftUpPoint = (x0-r, y0-r)
-            rightDownPoint = (x0+r, y0+r)
-            draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
-            draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
+    print('hello ', text_position)
 
-    print('stack', stack)
+    if radio == 'Two Points':
 
-    return image
+        if global_dict[guest_id]['state'] == 0:
+            global_dict[guest_id]['stack'].append(
+                (text_position, t)
+            )
+            print(text_position, global_dict[guest_id]['stack'])
+            global_dict[guest_id]['state'] = 1
+        else:
+            
+            (_, t) = global_dict[guest_id]['stack'].pop()
+            x, y = _
+            global_dict[guest_id]['stack'].append(
+                ((x,y,text_position[0],text_position[1]), t)
+            )
+            global_dict[guest_id]['state'] = 0
+
+        image = copy.deepcopy(orig_i) 
+        draw = ImageDraw.Draw(image)
+
+        for items in global_dict[guest_id]['stack']:
+            text_position, t = items
+            if len(text_position) == 2:
+                x, y = text_position
+
+                x = int(512 * x / width)
+                y = int(512 * y / height)
+
+                text_color = (255, 0, 0)  
+                draw.text((x+2, y), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x-r, y-r)
+                rightDownPoint = (x+r, y+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+            elif len(text_position) == 4:
+                x0, y0, x1, y1 = text_position
+
+                x0 = int(512 * x0 / width)
+                x1 = int(512 * x1 / width)
+                y0 = int(512 * y0 / height)
+                y1 = int(512 * y1 / height)
+
+                text_color = (255, 0, 0)  
+                draw.text((x0+2, y0), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x0-r, y0-r)
+                rightDownPoint = (x0+r, y0+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+                draw.rectangle((x0,y0,x1,y1), outline=(255, 0, 0) )
+
+    elif radio == 'Four Points':
+
+        if global_dict[guest_id]['state'] == 0:
+            global_dict[guest_id]['stack'].append(
+                (text_position, t)
+            )
+            print(text_position, global_dict[guest_id]['stack'])
+            global_dict[guest_id]['state'] = 1
+        elif global_dict[guest_id]['state'] == 1:
+            (_, t) = global_dict[guest_id]['stack'].pop()
+            x, y = _
+            global_dict[guest_id]['stack'].append(
+                ((x,y,text_position[0],text_position[1]), t)
+            )
+            global_dict[guest_id]['state'] = 2
+        elif global_dict[guest_id]['state'] == 2:
+            (_, t) = global_dict[guest_id]['stack'].pop()
+            x0, y0, x1, y1 = _
+            global_dict[guest_id]['stack'].append(
+                ((x0, y0, x1, y1,text_position[0],text_position[1]), t)
+            )
+            global_dict[guest_id]['state'] = 3
+        elif global_dict[guest_id]['state'] == 3:
+            (_, t) = global_dict[guest_id]['stack'].pop()
+            x0, y0, x1, y1, x2, y2 = _
+            global_dict[guest_id]['stack'].append(
+                ((x0, y0, x1, y1, x2, y2,text_position[0],text_position[1]), t)
+            )
+            global_dict[guest_id]['state'] = 0
+
+        image = copy.deepcopy(orig_i) 
+        draw = ImageDraw.Draw(image)
+
+        for items in global_dict[guest_id]['stack']:
+            text_position, t = items
+            if len(text_position) == 2:
+                x, y = text_position
+
+                x = int(512 * x / width)
+                y = int(512 * y / height)
+
+                text_color = (255, 0, 0)  
+                draw.text((x+2, y), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x-r, y-r)
+                rightDownPoint = (x+r, y+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+            elif len(text_position) == 4:
+                x0, y0, x1, y1 = text_position
+                text_color = (255, 0, 0)  
+                draw.text((x0+2, y0), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x0-r, y0-r)
+                rightDownPoint = (x0+r, y0+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+                draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
+            elif len(text_position) == 6:
+                x0, y0, x1, y1, x2, y2 = text_position
+                text_color = (255, 0, 0)  
+                draw.text((x0+2, y0), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x0-r, y0-r)
+                rightDownPoint = (x0+r, y0+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+                draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
+                draw.line(((x1,y1),(x2,y2)), fill=(255, 0, 0) )
+            elif len(text_position) == 8:
+                x0, y0, x1, y1, x2, y2, x3, y3 = text_position
+                text_color = (255, 0, 0)  
+                draw.text((x0+2, y0), t, font=font, fill=text_color)
+                r = 4
+                leftUpPoint = (x0-r, y0-r)
+                rightDownPoint = (x0+r, y0+r)
+                draw.ellipse((leftUpPoint,rightDownPoint), fill='red')
+                draw.line(((x0,y0),(x1,y1)), fill=(255, 0, 0) )
+                draw.line(((x1,y1),(x2,y2)), fill=(255, 0, 0) )
+                draw.line(((x2,y2),(x3,y3)), fill=(255, 0, 0) )
+                draw.line(((x3,y3),(x0,y0)), fill=(255, 0, 0) )
+
+
+    print('stack', global_dict[guest_id]['stack'])
+
+    global_dict[str(seed)]['image_id'].append(list(image.getdata()))
+
+    return image, orig_i, seed
 
 
 font_layout = ImageFont.truetype('./Arial.ttf', 16)
@@ -208,11 +328,30 @@ def get_layout_image(ocrs):
     return blank
 
 
+def to_tensor(image):
+    if isinstance(image, Image.Image):  
+        image = np.array(image)
+    elif not isinstance(image, np.ndarray):  
+        raise TypeError("Error")
+
+    image = image.astype(np.float32) / 255.0
+    image = np.transpose(image, (2, 0, 1))
+    tensor = torch.from_numpy(image)
+
+    return tensor
 
-def text_to_image(prompt,keywords,positive_prompt,radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural):
+def test_fn(x,y):
+    print('hello')
 
-    global stack
-    global state
+def text_to_image(guest_id, i, orig_i, prompt,keywords,positive_prompt,radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural):
+
+    # print(type(i))
+    # exit(0)
+
+    print(f'[info] Prompt: {prompt} | Keywords: {keywords} | Radio: {radio} | Steps: {slider_step} | Guidance: {slider_guidance} | Natural: {slider_natural}')
+
+    # global stack
+    # global state
 
     if len(positive_prompt.strip()) != 0:
         prompt += positive_prompt
@@ -227,7 +366,7 @@ def text_to_image(prompt,keywords,positive_prompt,radio,slider_step,slider_guida
             prompt = tokenizer.encode(user_prompt)
             layout_image = None
         else:
-            if len(stack) == 0:
+            if guest_id not in global_dict or len(global_dict[guest_id]['stack']) == 0:
 
                 if len(keywords.strip()) == 0:
                     template = f'Given a prompt that will be used to generate an image, plan the layout of visual text for the image. The size of the image is 128x128. Therefore, all properties of the positions should not exceed 128, including the coordinates of top, left, right, and bottom. All keywords are included in the caption. You dont need to specify the details of font styles. At each line, the format should be keyword left, top, right, bottom. So let us begin. Prompt: {user_prompt}'
@@ -308,19 +447,25 @@ def text_to_image(prompt,keywords,positive_prompt,radio,slider_step,slider_guida
                 composed_prompt = tokenizer.decode(prompt)
             
             else:
-                user_prompt += ' <|endoftext|>'
+                user_prompt += ' <|endoftext|><|startoftext|>'
                 layout_image = None
-                
-                for items in stack:
+
+                image_mask = Image.new('L', (512,512), 0)
+                draw = ImageDraw.Draw(image_mask)
+
+                for items in global_dict[guest_id]['stack']:
                     position, text = items
 
+                    # feature_mask
+                    # masked_feature
                     
                     if len(position) == 2:
                         x, y = position
                         x = x // 4
                         y = y // 4
                         text_str = ' '.join([f'[{c}]' for c in list(text)])
-                        user_prompt += f'<|startoftext|> l{x} t{y} {text_str} <|endoftext|>'
+                        user_prompt += f' l{x} t{y} {text_str} <|endoftext|>'
+
                     elif len(position) == 4:
                         x0, y0, x1, y1 = position
                         x0 = x0 // 4
@@ -328,9 +473,32 @@ def text_to_image(prompt,keywords,positive_prompt,radio,slider_step,slider_guida
                         x1 = x1 // 4
                         y1 = y1 // 4
                         text_str = ' '.join([f'[{c}]' for c in list(text)])
-                        user_prompt += f'<|startoftext|> l{x0} t{y0} r{x1} b{y1} {text_str} <|endoftext|>'
+                        user_prompt += f' l{x0} t{y0} r{x1} b{y1} {text_str} <|endoftext|>'
+
+                        draw.rectangle((x0*4, y0*4, x1*4, y1*4), fill=1)
+                        print('prompt ', user_prompt)
+
+                    elif len(position) == 8: # four points
+                        x0, y0, x1, y1, x2, y2, x3, y3 = position
+                        draw.polygon([(x0, y0), (x1, y1), (x2, y2), (x3, y3)], fill=1)
+                        x0 = x0 // 4
+                        y0 = y0 // 4
+                        x1 = x1 // 4
+                        y1 = y1 // 4
+                        x2 = x2 // 4
+                        y2 = y2 // 4
+                        x3 = x3 // 4
+                        y3 = y3 // 4
+                        xmin = min(x0, x1, x2, x3)
+                        ymin = min(y0, y1, y2, y3)
+                        xmax = max(x0, x1, x2, x3)
+                        ymax = max(y0, y1, y2, y3)
+                        text_str = ' '.join([f'[{c}]' for c in list(text)])
+                        user_prompt += f' l{xmin} t{ymin} r{xmax} b{ymax} {text_str} <|endoftext|>'
+
+                        print('prompt ', user_prompt)
+
 
-                    # composed_prompt = user_prompt
                     prompt = tokenizer.encode(user_prompt)
                     composed_prompt = tokenizer.decode(prompt)
 
@@ -338,70 +506,67 @@ def text_to_image(prompt,keywords,positive_prompt,radio,slider_step,slider_guida
         while len(prompt) < 77: 
             prompt.append(tokenizer.pad_token_id) 
 
-        if radio == 'TextDiffuser-2':
-            
-            prompts_cond = prompt
-            prompts_nocond = [tokenizer.pad_token_id]*77
-
-            prompts_cond = [prompts_cond] * slider_batch
-            prompts_nocond = [prompts_nocond] * slider_batch
-
-            prompts_cond = torch.Tensor(prompts_cond).long().cuda()
-            prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
-
-            scheduler = DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="scheduler") 
-            scheduler.set_timesteps(slider_step) 
-            noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda").half()
-            input = noise
-
-            encoder_hidden_states_cond = text_encoder(prompts_cond)[0].half()
-            encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0].half()
-
-
-            for t in tqdm(scheduler.timesteps):
-                with torch.no_grad():  # classifier free guidance
-                    noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:slider_batch]).sample # b, 4, 64, 64
-                    noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:slider_batch]).sample # b, 4, 64, 64
-                    noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64     
-                    input = scheduler.step(noisy_residual, t, input).prev_sample
-                    del noise_pred_cond
-                    del noise_pred_uncond
-
-                    torch.cuda.empty_cache()
-
-            # decode
-            input = 1 / vae.config.scaling_factor * input 
-            images = vae.decode(input, return_dict=False)[0] 
-            width, height = 512, 512
-            results = []
-            new_image = Image.new('RGB', (2*width, 2*height))
-            for index, image in enumerate(images.cpu().float()):
-                image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
-                image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
-                image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
-                results.append(image)
-                row = index // 2
-                col = index % 2
-                new_image.paste(image, (col*width, row*height))
-            # os.system('nvidia-smi')
-            torch.cuda.empty_cache()
-            # os.system('nvidia-smi')
-            return tuple(results),  composed_prompt, layout_image
-        
-        elif radio == 'TextDiffuser-2-LCM':
-            generator = torch.Generator(device=pipe.device).manual_seed(random.randint(0,1000))
-            image = pipe(
-                prompt=user_prompt,
-                generator=generator,
-                # negative_prompt=negative_prompt,
-                num_inference_steps=slider_step,
-                guidance_scale=1,
-                # num_images_per_prompt=slider_batch,
-            ).images
-            # os.system('nvidia-smi')
-            torch.cuda.empty_cache()
-            # os.system('nvidia-smi')
-            return tuple(image), composed_prompt, layout_image
+        prompts_cond = prompt
+        prompts_nocond = [tokenizer.pad_token_id]*77
+
+        prompts_cond = [prompts_cond] * slider_batch
+        prompts_nocond = [prompts_nocond] * slider_batch
+
+        prompts_cond = torch.Tensor(prompts_cond).long().cuda()
+        prompts_nocond = torch.Tensor(prompts_nocond).long().cuda()
+
+        scheduler = DDPMScheduler.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder="scheduler") 
+        scheduler.set_timesteps(slider_step) 
+        noise = torch.randn((slider_batch, 4, 64, 64)).to("cuda").half()
+        input = noise
+
+        encoder_hidden_states_cond = text_encoder(prompts_cond)[0].half()
+        encoder_hidden_states_nocond = text_encoder(prompts_nocond)[0].half()
+
+        image_mask = torch.Tensor(np.array(image_mask)).float().half().cuda()
+        image_mask = image_mask.unsqueeze(0).unsqueeze(0).repeat(slider_batch, 1, 1, 1)
+
+        image = Image.open(orig_i).resize((512,512))
+        image_tensor = to_tensor(image).unsqueeze(0).cuda().sub_(0.5).div_(0.5)   
+        print(f'image_tensor.shape {image_tensor.shape}')
+        masked_image = image_tensor * (1-image_mask)
+        masked_feature = vae.encode(masked_image.half()).latent_dist.sample() 
+        masked_feature = masked_feature * vae.config.scaling_factor
+        masked_feature = masked_feature.half()
+        print(f'masked_feature.shape {masked_feature.shape}')
+
+        feature_mask = torch.nn.functional.interpolate(image_mask, size=(64,64), mode='nearest').cuda()
+
+        for t in tqdm(scheduler.timesteps):
+            with torch.no_grad():  # classifier free guidance
+
+                noise_pred_cond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_cond[:slider_batch],feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
+                noise_pred_uncond = unet(sample=input, timestep=t, encoder_hidden_states=encoder_hidden_states_nocond[:slider_batch],feature_mask=feature_mask, masked_feature=masked_feature).sample # b, 4, 64, 64
+                noisy_residual = noise_pred_uncond + slider_guidance * (noise_pred_cond - noise_pred_uncond) # b, 4, 64, 64     
+                input = scheduler.step(noisy_residual, t, input).prev_sample
+                del noise_pred_cond
+                del noise_pred_uncond
+
+                torch.cuda.empty_cache()
+
+        # decode
+        input = 1 / vae.config.scaling_factor * input 
+        images = vae.decode(input, return_dict=False)[0] 
+        width, height = 512, 512
+        results = []
+        new_image = Image.new('RGB', (2*width, 2*height))
+        for index, image in enumerate(images.cpu().float()):
+            image = (image / 2 + 0.5).clamp(0, 1).unsqueeze(0)
+            image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
+            image = Image.fromarray((image * 255).round().astype("uint8")).convert('RGB')
+            results.append(image)
+            row = index // 2
+            col = index % 2
+            new_image.paste(image, (col*width, row*height))
+        # os.system('nvidia-smi')
+        torch.cuda.empty_cache()
+        # os.system('nvidia-smi')
+        return tuple(results), composed_prompt
         
 with gr.Blocks() as demo:
 
@@ -411,6 +576,9 @@ with gr.Blocks() as demo:
         <h2 style="font-weight: 900; font-size: 2.3rem; margin: 0rem">
             TextDiffuser-2: Unleashing the Power of Language Models for Text Rendering
         </h2>
+        <h2 style="font-weight: 900; font-size: 1.3rem; margin: 0rem">
+            (Demo for <b>Text Inpainting</b> 🖼️🖌️)
+        </h2>
         <h2 style="font-weight: 460; font-size: 1.1rem; margin: 0rem">
             <a href="https://jingyechen.github.io/">Jingye Chen</a>, <a href="https://hypjudy.github.io/website/">Yupan Huang</a>, <a href="https://scholar.google.com/citations?user=0LTZGhUAAAAJ&hl=en">Tengchao Lv</a>, <a href="https://www.microsoft.com/en-us/research/people/lecu/">Lei Cui</a>, <a href="https://cqf.io/">Qifeng Chen</a>, <a href="https://thegenerality.com/">Furu Wei</a>
         </h2>      
@@ -420,96 +588,87 @@ with gr.Blocks() as demo:
         <h3 style="font-weight: 450; font-size: 1rem; margin: 0rem"> 
         [<a href="https://arxiv.org/abs/2311.16465" style="color:blue;">arXiv</a>] 
         [<a href="https://github.com/microsoft/unilm/tree/master/textdiffuser-2" style="color:blue;">Code</a>]
+        [<a href="https://jingyechen.github.io/textdiffuser2/" style="color:blue;">Project Page</a>]
+        [<a href="https://discord.gg/q7eHPupu" style="color:purple;">Discord</a>]
         </h3> 
         <h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
-        We propose <b>TextDiffuser-2</b>, aiming at unleashing the power of language models for text rendering. Specifically, we <b>tame a language model into a layout planner</b> to transform user prompt into a layout using the caption-OCR pairs. The language model demonstrates flexibility and automation by inferring keywords from user prompts or incorporating user-specified keywords to determine their positions. Secondly, we <b>leverage the language model in the diffusion model as the layout encoder</b> to represent the position and content of text at the line level. This approach enables diffusion models to generate text images with broader diversity.
+        TextDiffuser-2 leverages language models to enhance text rendering, achieving greater flexibility. Different from text editing, the text inpainting task aims to add or modify text guided by users, ensuring that the inpainted text has a reasonable style (i.e., no need to match the style of the original text during modification exactly) and is coherent with backgrounds. TextDiffuser-2 offers an <b>improved user experience</b>. Specifically, users only need to type the text they wish to inpaint into the provided input box and then select key points on the Canvas.
         </h2>
         <h2 style="text-align: left; font-weight: 450; font-size: 1rem; margin-top: 0.5rem; margin-bottom: 0.5rem">
-        👀 <b>Tips for using this demo</b>: <b>(1)</b> Please carefully read the disclaimer in the below. Current verison can only support English. <b>(2)</b> The specification of keywords is optional. If provided, the language model will do its best to plan layouts using the given keywords. <b>(3)</b> If a template is given, the layout planner (M1) is not used. <b>(4)</b> Three operations, including redo, undo, and skip are provided. When using skip, only the left-top point of a keyword will be recorded, resulting in more diversity but sometimes decreasing the accuracy. <b>(5)</b> The layout planner can produce different layouts. You can increase the temperature to enhance the diversity. ✨ <b>(6)</b> We also provide the experimental demo combining <b>TextDiffuser-2</b> and <b>LCM</b>. The inference is fast using less sampling steps, although the precision in text rendering might decrease.
+        👀 <b>Tips for using this demo</b>: <b>(1)</b> Please carefully read the disclaimer in the below. Current verison can only support English. <b>(2)</b> The <b>prompt is optional</b>. If provided, the generated image may be more accurate. <b>(3)</b> Redo is used to cancel the last keyword, and undo is used to clear all keywords. <b>(4)</b> Current version only supports input image with resolution 512x512. <b>(5)</b> You can use either two points or four points to specify the text box. Using four points can better represent the perspective boxes. <b>(6)</b> Leave "Text to be inpaintd" empty can function as the text removal task. <b>(7)</b> Classifier-free guidance is set to a small value (e.g. 1) in default. It is noticed that a larger cfg may result in chromatic aberration against the background. <b>(8)</b> You can inpaint many text regions at one time. <b>(9)</b> Thanks for reading these tips, shall we start now?
         </h2>
-        <style>
-            .scaled-image {
-                transform: scale(1);
-            }
-        </style>
-        
-        <img src="https://i.ibb.co/56JVg5j/architecture.jpg" alt="textdiffuser-2" class="scaled-image">
+        <img src="https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/inpainting_blank.jpg" alt="textdiffuser-2">
         </div>
         """)
 
-    with gr.Tab("Text-to-Image"):
+    with gr.Tab("Text Inpainting"):
         with gr.Row():
-            with gr.Column(scale=1):
-                prompt = gr.Textbox(label="Prompt. You can let language model automatically identify keywords, or provide them below", placeholder="A beautiful city skyline stamp of Shanghai")
-                keywords = gr.Textbox(label="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...)", placeholder="keyword1/keyword2")
-                positive_prompt = gr.Textbox(label="(Optional) Positive prompt", value=", digital art, very detailed, fantasy, high definition, cinematic light, dnd, trending on artstation")
-
-                with gr.Accordion("(Optional) Template - Click to paint", open=False):
-                    with gr.Row():
-                        with gr.Column(scale=1):
-                            i = gr.Image(label="Canvas", type='filepath', value=f'./gray256.jpg', height=256, width=256)
-                        with gr.Column(scale=1):
-                            t = gr.Textbox(label="Keyword", value='input_keyword')
-                            redo = gr.Button(value='Redo - Cancel the last keyword') 
-                            undo = gr.Button(value='Undo - Clear the canvas') 
-                            skip_button = gr.Button(value='Skip - Operate the next keyword') 
-
-                i.select(get_pixels,[i,t],[i])
-                redo.click(exe_redo, [i,t],[i])
-                undo.click(exe_undo, [i,t],[i])
-                skip_button.click(skip_fun, [i,t])
-
-                radio = gr.Radio(["TextDiffuser-2", "TextDiffuser-2-LCM"], label="Choice of models", value="TextDiffuser-2")
-                slider_natural = gr.Checkbox(label="Natural image generation", value=False, info="The text position and content info will not be incorporated.")
-                slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser-2. You may decease the step to 4 when using LCM.")
-                slider_guidance = gr.Slider(minimum=1, maximum=13, value=7.5, step=0.5, label="Scale of classifier-free guidance", info="The scale of cfg and is set to 7.5 in default. When using LCM, cfg is set to 1.")
-                slider_batch = gr.Slider(minimum=1, maximum=4, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
-                slider_temperature = gr.Slider(minimum=0.1, maximum=2, value=1.4, step=0.1, label="Temperature", info="Control the diversity of layout planner. Higher value indicates more diversity.")
+            with gr.Column():
+
+                keywords = gr.Textbox(label="(Optional) Keywords. Should be seperated by / (e.g., keyword1/keyword2/...)", placeholder="keyword1/keyword2", visible=False)
+                positive_prompt = gr.Textbox(label="(Optional) Positive prompt", value="", visible=False)
+
+                i = gr.Image(label="Image", type='filepath', value='https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example11.jpg')
+                orig_i = gr.Image(label="Placeholder", type='filepath', height=512, width=512, visible=False)
+
+                radio = gr.Radio(["Two Points", "Four Points"], label="Number of points to represent the text box.", value="Two Points", visible=True)
+
+                with gr.Row():
+                    t = gr.Textbox(label="Text to be inpainted", value='Test')
+                    prompt = gr.Textbox(label="(Optional) Prompt.")
+                with gr.Row():
+                    redo = gr.Button(value='Redo - Cancel the last keyword') 
+                    undo = gr.Button(value='Undo - Clear the canvas') 
+                # skip_button = gr.Button(value='Skip - Operate the next keyword') 
+
+                slider_natural = gr.Checkbox(label="Natural image generation", value=False, info="The text position and content info will not be incorporated.", visible=False)
+                slider_step = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Sampling step", info="The sampling step for TextDiffuser-2.")
+                slider_guidance = gr.Slider(minimum=1, maximum=13, value=1, step=0.5, label="Scale of classifier-free guidance", info="The scale of cfg and is set to 1 in default. Smaller cfg produce stable results.")
+                slider_batch = gr.Slider(minimum=1, maximum=6, value=4, step=1, label="Batch size", info="The number of images to be sampled.")
+                slider_temperature = gr.Slider(minimum=0.1, maximum=2, value=1.4, step=0.1, label="Temperature", info="Control the diversity of layout planner. Higher value indicates more diversity.", visible=False)
                 # slider_seed = gr.Slider(minimum=1, maximum=10000, label="Seed", randomize=True)
                 button = gr.Button("Generate")
+
+                guest_id_box = gr.Textbox(label="guest_id", value=f"-1", visible=False)
+                i.select(get_pixels,[i,orig_i,radio,t,guest_id_box],[i,orig_i,guest_id_box])
+                redo.click(exe_redo, [i,orig_i,t,guest_id_box],[i])
+                undo.click(exe_undo, [i,orig_i,t,guest_id_box],[i])
+                # skip_button.click(skip_fun, [i,t,guest_id_box])
+
                             
-            with gr.Column(scale=1):
-                output = gr.Gallery(label='Generated image')
+            with gr.Column():
+                output = gr.Gallery(label='Generated image', rows=2, height=768)
 
-                with gr.Accordion("Intermediate results", open=False):
+                with gr.Accordion("Intermediate results", open=False, visible=False):
                     gr.Markdown("Composed prompt")
                     composed_prompt = gr.Textbox(label='')
-                    gr.Markdown("Layout visualization")
-                    layout = gr.Image(height=256, width=256)
+                    # gr.Markdown("Layout visualization")
+                    # layout = gr.Image(height=256, width=256)
 
 
-        button.click(text_to_image, inputs=[prompt,keywords,positive_prompt, radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural], outputs=[output, composed_prompt, layout])
+        button.click(text_to_image, inputs=[guest_id_box, i, orig_i, prompt,keywords,positive_prompt, radio,slider_step,slider_guidance,slider_batch,slider_temperature,slider_natural], outputs=[output, composed_prompt])
 
-        gr.Markdown("## Prompt Examples")
-        gr.Examples(
+        gr.Markdown("## Image Examples")
+        template = None
+        gr.Examples(   
             [
-                ["A beautiful city skyline stamp of Shanghai", "", False],
-                ["A logo of superman", "", False],
-                ["A pencil sketch of a tree with the title nothing to tree here", "", False],
-                ["handwritten signature of peter", "", False],
-                ["Delicate greeting card of happy birthday to xyz", "", False],
-                ["Book cover of good morning baby ", "", False],
-                ["The handwritten words Hello World displayed on a wall in a neon light effect", "", False],
-                ["Logo of winter in artistic font, made by snowflake", "", False],
-                ["A book cover named summer vibe", "", False],
-                ["Newspaper with the title Love Story", "", False],
-                ["A logo for the company EcoGrow, where the letters look like plants", "EcoGrow", False],
-                ["A poster titled 'Quails of North America', showing different kinds of quails.", "Quails/of/North/America", False],
-                ["A detailed portrait of a fox guardian with a shield with Kung Fu written on it, by victo ngai and justin gerard, digital art, realistic painting", "kung/fu", False],
-                ["A stamp of breath of the wild", "breath/of/the/wild", False],
-                ["Poster of the incoming movie Transformers", "Transformers", False],
-                ["Some apples are on a table", "", True],
-                ["a hotdog with mustard and other toppings on it", "", True],
-                ["a bathroom that has a slanted ceiling and a large bath tub", "", True],
-                ["a man holding a tennis racquet on a tennis court", "", True],
-                ["hamburger with bacon, lettuce, tomato and cheese| promotional image| hyperquality| products shot| full - color| extreme render| mouthwatering", "", True],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example1.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example2.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example3.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example4.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example5.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example7.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example8.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example11.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example12.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example13.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example14.jpg"],
+                ["https://raw.githubusercontent.com/JingyeChen/jingyechen.github.io/master/textdiffuser2/static/images/example15.jpg"],
             ],
             [
-                prompt,
-                keywords,
-                slider_natural
+                i
             ],
-            examples_per_page=20
+            examples_per_page=25,
         )
 
     gr.HTML(