Added gradio app

app.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import math
+import random
+
+import gradio as gr
+import torch
+from PIL import Image, ImageOps
+from run_edit import run_model
+from cool_models import make_models
+
+help_text = """"""
+
+example_instructions = [
+    "Make it a picasso painting",
+    "as if it were by modigliani",
+    "convert to a bronze statue",
+    "Turn it into an anime.",
+    "have it look like a graphic novel",
+    "make him gain weight",
+    "what would he look like bald?",
+    "Have him smile",
+    "Put him in a cocktail party.",
+    "move him at the beach.",
+    "add dramatic lighting",
+    "Convert to black and white",
+    "What if it were snowing?",
+    "Give him a leather jacket",
+    "Turn him into a cyborg!",
+    "make him wear a beanie",
+]
+
+model_id = "timbrooks/instruct-pix2pix"
+
+def main():
+    # pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None).to("cuda")
+    segmodel, model, diffusion, ldm, bert, clip_model, model_params = make_models()
+
+    def generate(
+        input_image: Image.Image,
+        from_text: str,
+        instruction: str,
+        negative_prompt: str,
+        randomize_seed: bool,
+        seed: int,
+        guidance_scale: float,
+        clip_guidance_scale: float,
+        cutn: int,
+        l2_sim_lambda: float
+    ):
+        seed = random.randint(0, 100000) if randomize_seed else seed
+
+        if instruction == "":
+            return [seed, input_image]
+
+        generator = torch.manual_seed(seed)
+
+        edited_image_1 = run_model(
+            segmodel, model, diffusion, ldm, bert, clip_model, model_params,
+            from_text, instruction, negative_prompt, input_image.convert('RGB'), seed, guidance_scale, clip_guidance_scale, cutn, l2_sim_lambda
+        )
+
+        # edited_image = input_image
+        return [seed, edited_image_1]
+
+    def reset():
+        return [
+            "Randomize Seed", 1371, None, 5.0,
+            150, 16, 10000
+        ]
+
+    with gr.Blocks() as demo:
+        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">
+   RDM: Region-Aware Diffusion for Zero-shot Text-driven Image Editing
+</h1>
+<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
+<br/>
+<a href="https://huggingface.co/spaces/timbrooks/instruct-pix2pix?duplicate=true">
+<img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
+<p/>""")
+        with gr.Row():
+            with gr.Column(scale=1, min_width=100):
+                generate_button = gr.Button("Generate")
+            # with gr.Column(scale=1, min_width=100):
+            #     load_button = gr.Button("Load Example")
+            with gr.Column(scale=1, min_width=100):
+                reset_button = gr.Button("Reset")
+            with gr.Column(scale=3):
+                from_text = gr.Textbox(lines=1, label="From Text", interactive=True)
+                instruction = gr.Textbox(lines=1, label="Edit Instruction", interactive=True)
+                negative_prompt = gr.Textbox(lines=1, label="Negative Prompt", interactive=True)
+
+        with gr.Row():
+            input_image = gr.Image(label="Input Image", type="pil", interactive=True)
+            edited_image_1 = gr.Image(label=f"Edited Image", type="pil", interactive=False)
+            # edited_image_2 = gr.Image(label=f"Edited Image", type="pil", interactive=False)
+            input_image.style(height=512, width=512)
+            edited_image_1.style(height=512, width=512)
+            # edited_image_2.style(height=512, width=512)
+
+        with gr.Row():
+            # steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)
+            seed = gr.Number(value=1371, precision=0, label="Seed", interactive=True)
+            guidance_scale = gr.Number(value=5.0, precision=1, label="Guidance Scale", interactive=True)
+            clip_guidance_scale = gr.Number(value=150, precision=1, label="Clip Guidance Scale", interactive=True)
+            cutn = gr.Number(value=16, precision=1, label="Number of Cuts", interactive=True)
+            l2_sim_lambda = gr.Number(value=10000, precision=1, label="L2 similarity to original image")
+
+            randomize_seed = gr.Radio(
+                ["Fix Seed", "Randomize Seed"],
+                value="Randomize Seed",
+                type="index",
+                show_label=False,
+                interactive=True,
+            )
+            # use_ddim = gr.Checkbox(label="Use 50-step DDIM?", value=True)
+            # use_ddpm = gr.Checkbox(label="Use 50-step DDPM?", value=True)
+        
+        gr.Markdown(help_text)
+
+        generate_button.click(
+            fn=generate,
+            inputs=[
+                input_image,
+                from_text,
+                instruction,
+                negative_prompt,
+                randomize_seed,
+                seed,
+                guidance_scale,
+                clip_guidance_scale,
+                cutn,
+                l2_sim_lambda
+            ],
+            outputs=[seed, edited_image_1],
+        )
+        reset_button.click(
+            fn=reset,
+            inputs=[],
+            outputs=[
+                randomize_seed, seed, edited_image_1, guidance_scale,
+                clip_guidance_scale, cutn, l2_sim_lambda
+            ],
+        )
+
+    demo.queue(concurrency_count=1)
+    demo.launch(share=False, server_name="0.0.0.0")
+
+
+if __name__ == "__main__":
+    main()

cool_models.py

@@ -0,0 +1,132 @@
+import torch
+from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults
+import lpips
+import clip
+
+
+from encoders.modules import BERTEmbedder
+from models.clipseg import CLIPDensePredT
+
+from huggingface_hub import hf_hub_download
+
+STEPS = 100
+USE_DDPM = False
+USE_DDIM = False
+USE_CPU = False
+BERT_PATH = "./weights/bert.pt"
+KL_PATH = "./weights/kl-f8.pt"
+INPAINT_PATH = "./weights/inpaint.pt"
+CLIP_SEG_PATH = './weights/rd64-uni.pth'
+CLIP_GUIDANCE = False
+
+def make_models():
+    segmodel = CLIPDensePredT(version='ViT-B/16', reduce_dim=64)
+    segmodel.eval()
+
+    # non-strict, because we only stored decoder weights (not CLIP weights)
+    segmodel.load_state_dict(torch.load(CLIP_SEG_PATH, map_location=torch.device('cpu')), strict=False)
+    # segmodel.save_pretrained("./weights/hf_clipseg")
+
+    device = torch.device('cuda:0' if (torch.cuda.is_available() and not USE_CPU) else 'cpu')
+    print('Using device:', device)
+
+    hf_inpaint_path = hf_hub_download("alvanlii/rdm_inpaint", "inpaint.pt")
+    model_state_dict = torch.load(hf_inpaint_path, map_location='cpu')
+
+    # print(
+    #     'hey',
+    #     'clip_proj.weight' in model_state_dict, # True
+    #     model_state_dict['input_blocks.0.0.weight'].shape[1] == 8, # True
+    #     'external_block.0.0.weight' in model_state_dict # False
+    # )
+
+    model_params = {
+        'attention_resolutions': '32,16,8',
+        'class_cond': False,
+        'diffusion_steps': 1000,
+        'rescale_timesteps': True,
+        'timestep_respacing': STEPS,  # Modify this value to decrease the number of
+                                    # timesteps.
+        'image_size': 32,
+        'learn_sigma': False,
+        'noise_schedule': 'linear',
+        'num_channels': 320,
+        'num_heads': 8,
+        'num_res_blocks': 2,
+        'resblock_updown': False,
+        'use_fp16': False,
+        'use_scale_shift_norm': False,
+        'clip_embed_dim': 768,
+        'image_condition': True,
+        'super_res_condition': False,
+    }
+
+    if USE_DDPM:
+        model_params['timestep_respacing'] = '1000'
+    if USE_DDIM:
+        if STEPS:
+            model_params['timestep_respacing'] = 'ddim'+str(STEPS)
+        else:
+            model_params['timestep_respacing'] = 'ddim50'
+    elif STEPS:
+        model_params['timestep_respacing'] = str(STEPS)
+
+    model_config = model_and_diffusion_defaults()
+    model_config.update(model_params)
+
+    if USE_CPU:
+        model_config['use_fp16'] = False
+
+
+    model, diffusion = create_model_and_diffusion(**model_config)
+
+    # model.from_pretrained("alvanlii/rdm_inpaint")
+    model.load_state_dict(model_state_dict, strict=False)
+    # model.save_pretrained("./weights/hf_inpaint")
+
+    model.requires_grad_(CLIP_GUIDANCE).eval().to(device)
+
+    if model_config['use_fp16']:
+        model.convert_to_fp16()
+    else:
+        model.convert_to_fp32()
+
+    def set_requires_grad(model, value):
+        for param in model.parameters():
+            param.requires_grad = value
+
+
+    lpips_model = lpips.LPIPS(net="vgg").to(device)
+    hf_kl_path = hf_hub_download("alvanlii/rdm_kl", "kl-f8.pt")
+
+    # kl_model_url = hf_hub_url("alvanlii/rdm_kl", "kl-f8.pt")
+    # kl_cache_path = cached_download(kl_model_url, cache_dir=".")
+
+    ldm = torch.load(hf_kl_path, map_location="cpu")
+    
+    # torch.save(ldm, "./weights/hf_ldm")
+    ldm.to(device)
+    ldm.eval()
+    ldm.requires_grad_(CLIP_GUIDANCE)
+    set_requires_grad(ldm, CLIP_GUIDANCE)
+
+    bert = BERTEmbedder(1280, 32)
+    hf_bert_path = hf_hub_download("alvanlii/rdm_bert", 'bert.pt')
+    # bert = BERTEmbedder.from_pretrained("alvanlii/rdm_bert")
+    sd = torch.load(hf_bert_path, map_location="cpu")
+    bert.load_state_dict(sd)
+    # bert.save_pretrained("./weights/hf_bert")
+
+    bert.to(device)
+    bert.half().eval()
+    set_requires_grad(bert, False)
+
+
+    clip_model, clip_preprocess = clip.load('ViT-L/14', device=device, jit=False)
+    clip_model.eval().requires_grad_(False)
+
+    return segmodel, model, diffusion, ldm, bert, clip_model, model_params
+
+
+if __name__ == "__main__":
+    make_models()

requirements.txt

@@ -0,0 +1,13 @@
+einops==0.6.0
+lpips
+gradio
+opencv-python
+--extra-index-url https://download.pytorch.org/whl/cu116
+torch 
+--extra-index-url https://download.pytorch.org/whl/cu116
+torchvision 
+transformers
+pytorch-lightning
+git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
+git+https://github.com/openai/CLIP.git@main#egg=clip
+git+https://github.com/alvanli/RDM-Region-Aware-Diffusion-Model.git@main#egg=guided_diffusion

run_edit.py

@@ -0,0 +1,288 @@
+import gc
+import os
+import io
+import math
+import sys
+import tempfile
+
+from PIL import Image, ImageOps
+import requests
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torchvision import transforms
+from torchvision.transforms import functional as TF
+from tqdm.notebook import tqdm
+
+import numpy as np
+
+from math import log2, sqrt
+
+import argparse
+import pickle
+
+
+
+
+################################### mask_fusion ######################################
+from util.metrics_accumulator import MetricsAccumulator
+metrics_accumulator = MetricsAccumulator()
+
+from pathlib import Path
+from PIL import Image
+################################### mask_fusion ######################################
+
+import clip
+import lpips
+from torch.nn.functional import mse_loss
+
+################################### CLIPseg ######################################
+from  torchvision import utils as vutils
+import cv2  
+
+################################### CLIPseg ######################################
+
+def str2bool(x):
+    return x.lower() in ('true')
+    
+USE_CPU = False
+device = torch.device('cuda:0' if (torch.cuda.is_available() and not USE_CPU) else 'cpu')
+
+
+def fetch(url_or_path):
+    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):
+        r = requests.get(url_or_path)
+        r.raise_for_status()
+        fd = io.BytesIO()
+        fd.write(r.content)
+        fd.seek(0)
+        return fd
+    return open(url_or_path, 'rb')
+
+
+class MakeCutouts(nn.Module):
+    def __init__(self, cut_size, cutn, cut_pow=1.):
+        super().__init__()
+
+        self.cut_size = cut_size
+        self.cutn = cutn
+        self.cut_pow = cut_pow
+
+    def forward(self, input):
+        sideY, sideX = input.shape[2:4]
+        max_size = min(sideX, sideY)
+        min_size = min(sideX, sideY, self.cut_size)
+        cutouts = []
+        for _ in range(self.cutn):
+            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)
+            offsetx = torch.randint(0, sideX - size + 1, ())
+            offsety = torch.randint(0, sideY - size + 1, ())
+            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]
+            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))
+        return torch.cat(cutouts)
+
+def spherical_dist_loss(x, y):
+    x = F.normalize(x, dim=-1)
+    y = F.normalize(y, dim=-1)
+    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)
+
+
+def do_run(
+    arg_seed, arg_text, arg_batch_size, arg_num_batches, arg_negative, arg_cutn, arg_edit, arg_height, arg_width,
+    arg_edit_y, arg_edit_x, arg_edit_width, arg_edit_height, mask, arg_guidance_scale, arg_background_preservation_loss,
+    arg_lpips_sim_lambda, arg_l2_sim_lambda, arg_ddpm, arg_ddim, arg_enforce_background, arg_clip_guidance_scale,
+    arg_clip_guidance, model_params, model, diffusion, ldm, bert, clip_model
+):
+    normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])
+
+    if arg_seed >= 0:
+        torch.manual_seed(arg_seed)
+
+    text_emb = bert.encode([arg_text] * arg_batch_size).to(device).float()
+    text_blank = bert.encode([arg_negative] * arg_batch_size).to(device).float()
+
+    text = clip.tokenize([arg_text] * arg_batch_size, truncate=True).to(device)
+    text_clip_blank = clip.tokenize([arg_negative] * arg_batch_size, truncate=True).to(device)
+
+
+
+    text_emb_clip = clip_model.encode_text(text)
+    text_emb_clip_blank = clip_model.encode_text(text_clip_blank)
+    make_cutouts = MakeCutouts(clip_model.visual.input_resolution, arg_cutn)
+    text_emb_norm = text_emb_clip[0] / text_emb_clip[0].norm(dim=-1, keepdim=True)
+    image_embed = None
+
+    if arg_edit:
+        w = arg_edit_width if arg_edit_width else arg_width
+        h = arg_edit_height if arg_edit_height else arg_height
+
+        arg_edit = arg_edit.convert('RGB')
+        input_image_pil = arg_edit
+
+        init_image_pil = input_image_pil.resize((arg_height, arg_width), Image.Resampling.LANCZOS)
+
+        input_image_pil = ImageOps.fit(input_image_pil, (w, h))
+
+        im = transforms.ToTensor()(input_image_pil).unsqueeze(0).to(device)
+
+        init_image = (TF.to_tensor(init_image_pil).to(device).unsqueeze(0).mul(2).sub(1))
+
+        im = 2*im-1
+        im = ldm.encode(im).sample()
+
+        y = arg_edit_y//8
+        x = arg_edit_x//8
+
+        input_image = torch.zeros(1, 4, arg_height//8, arg_width//8, device=device)
+
+        ycrop = y + im.shape[2] - input_image.shape[2]
+        xcrop = x + im.shape[3] - input_image.shape[3]
+
+        ycrop = ycrop if ycrop > 0 else 0
+        xcrop = xcrop if xcrop > 0 else 0
+
+        input_image[0,:,y if y >=0 else 0:y+im.shape[2],x if x >=0 else 0:x+im.shape[3]] = im[:,:,0 if y > 0 else -y:im.shape[2]-ycrop,0 if x > 0 else -x:im.shape[3]-xcrop]
+
+        input_image_pil = ldm.decode(input_image)
+        input_image_pil = TF.to_pil_image(input_image_pil.squeeze(0).add(1).div(2).clamp(0, 1))
+
+        input_image *= 0.18215
+
+        new_mask = TF.resize(mask.unsqueeze(0).unsqueeze(0).to(device), (arg_width//8, arg_height//8))
+
+        mask1 = (new_mask > 0.5)
+        mask1 = mask1.float()
+
+        input_image *= mask1
+
+        image_embed = torch.cat(arg_batch_size*2*[input_image], dim=0).float()
+    elif model_params['image_condition']:
+        # using inpaint model but no image is provided
+        image_embed = torch.zeros(arg_batch_size*2, 4, arg_height//8, arg_width//8, device=device)
+
+    kwargs = {
+        "context": torch.cat([text_emb, text_blank], dim=0).float(),
+        "clip_embed": torch.cat([text_emb_clip, text_emb_clip_blank], dim=0).float() if model_params['clip_embed_dim'] else None,
+        "image_embed": image_embed
+    }
+
+    # Create a classifier-free guidance sampling function
+    def model_fn(x_t, ts, **kwargs):
+        half = x_t[: len(x_t) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = model(combined, ts, **kwargs)
+        eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + arg_guidance_scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+    cur_t = None
+
+    @torch.no_grad()
+    def postprocess_fn(out, t):
+        if mask is not None:
+            background_stage_t = diffusion.q_sample(init_image, t[0])
+            background_stage_t = torch.tile(
+                background_stage_t, dims=(arg_batch_size, 1, 1, 1)
+            )
+            out["sample"] = out["sample"] * mask + background_stage_t * (1 - mask)
+        return out
+
+    # if arg_ddpm:
+    #     sample_fn = diffusion.p_sample_loop_progressive
+    # elif arg_ddim:
+    #     sample_fn = diffusion.ddim_sample_loop_progressive
+    # else:
+    sample_fn = diffusion.plms_sample_loop_progressive
+
+    def save_sample(i, sample):
+        out_ims = []
+        for k, image in enumerate(sample['pred_xstart'][:arg_batch_size]):
+            image /= 0.18215
+            im = image.unsqueeze(0)
+            out = ldm.decode(im)
+            metrics_accumulator.print_average_metric()
+
+            for b in range(arg_batch_size):
+                pred_image = sample["pred_xstart"][b]
+
+                if arg_enforce_background:
+                    new_mask = TF.resize(mask.unsqueeze(0).unsqueeze(0).to(device), (arg_width, arg_height))
+                    pred_image = (
+                        init_image[0] * new_mask[0] + out * (1 - new_mask[0])
+                    )
+
+                pred_image_pil = TF.to_pil_image(pred_image.squeeze(0).add(1).div(2).clamp(0, 1))
+                out_ims.append(pred_image_pil)
+        return out_ims
+            
+
+    all_saved_ims = []
+    for i in range(arg_num_batches):
+        cur_t = diffusion.num_timesteps - 1
+
+        samples = sample_fn(
+            model_fn,
+            (arg_batch_size*2, 4, int(arg_height//8), int(arg_width//8)),
+            clip_denoised=False,
+            model_kwargs=kwargs,
+            cond_fn=None,
+            device=device,
+            progress=True,
+        )
+    
+        for j, sample in enumerate(samples):
+            cur_t -= 1
+            if j % 5 == 0 and j != diffusion.num_timesteps - 1:
+                all_saved_ims += save_sample(i, sample)
+        all_saved_ims += save_sample(i, sample)
+
+    return all_saved_ims
+
+def run_model(
+    segmodel, model, diffusion, ldm, bert, clip_model, model_params,
+    from_text, instruction, negative_prompt, original_img, seed, guidance_scale, clip_guidance_scale, cutn, l2_sim_lambda
+):
+    input_image = original_img
+
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        transforms.Resize((256, 256)),
+    ])
+    img = transform(input_image).unsqueeze(0)
+
+    with torch.no_grad():
+        preds = segmodel(img.repeat(1,1,1,1), from_text)[0]
+
+    mask = torch.sigmoid(preds[0][0])
+    image = (mask.detach().cpu().numpy() * 255).astype(np.uint8) # cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)  
+    ret, thresh = cv2.threshold(image, 100, 255, cv2.THRESH_TRUNC, image) 
+    timg = np.array(thresh)
+    x, y = timg.shape
+    for row in range(x):
+        for col in range(y):
+            if (timg[row][col]) == 100:
+                timg[row][col] = 255
+            if (timg[row][col]) < 100:
+                timg[row][col] = 0
+    
+    fulltensor = torch.full_like(mask, fill_value=255)
+    bgtensor = fulltensor-timg
+    mask = bgtensor / 255.0
+
+    gc.collect()
+    use_ddim = False
+    use_ddpm = False
+    all_saved_ims = do_run(
+        seed, instruction, 1, 1, negative_prompt, cutn, input_image, 256, 256, 
+        0, 0, 0, 0, mask, guidance_scale, True,
+        1000, l2_sim_lambda, use_ddpm, use_ddim, True, clip_guidance_scale, False,
+        model_params, model, diffusion, ldm, bert, clip_model
+    )
+        
+    return all_saved_ims[-1]
+
+
+

weights/rd64-uni.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13845f6cee4d54ca46f62ee19dd354822094a26e0efccc64e606be93d6a7e26f
+size 4306645