Create run_edit.py

run_edit.py

@@ -0,0 +1,287 @@
+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]
+
+