editing_diffusion/optimization/image_editor.py

import os
from pathlib import Path
from optimization.constants import ASSETS_DIR_NAME, RANKED_RESULTS_DIR

from utils.metrics_accumulator import MetricsAccumulator
from utils.video import save_video
from utils.fft_pytorch import HighFrequencyLoss

from numpy import random
from optimization.augmentations import ImageAugmentations

from PIL import Image
import torch
import torchvision
from torchvision import transforms
import torchvision.transforms.functional as F
from torchvision.transforms import functional as TF
from torch.nn.functional import mse_loss
from optimization.losses import range_loss, d_clip_loss
import lpips
import numpy as np

from CLIP import clip
from guided_diffusion.guided_diffusion.script_util import (
    create_model_and_diffusion,
    model_and_diffusion_defaults,
    create_classifier,
    classifier_defaults,
)
from utils.visualization import show_tensor_image, show_editied_masked_image
from utils.change_place import change_place, find_bbox

import pdb
import cv2

def create_classifier_ours():

    model = torchvision.models.resnet50()
    ckpt = torch.load('checkpoints/DRA_resnet50.pth')['model_state_dict']
    model.load_state_dict({k.replace('module.','').replace('last_linear','fc'):v for k,v in ckpt.items()})
    model = torch.nn.Sequential(*[torch.nn.Upsample(size=(256,256)), model])
    return model

class ImageEditor:
    def __init__(self, args) -> None:
        self.args = args
        os.makedirs(self.args.output_path, exist_ok=True)

        self.ranked_results_path = Path(os.path.join(self.args.output_path, RANKED_RESULTS_DIR))
        os.makedirs(self.ranked_results_path, exist_ok=True)

        if self.args.export_assets:
            self.assets_path = Path(os.path.join(self.args.output_path, ASSETS_DIR_NAME))
            os.makedirs(self.assets_path, exist_ok=True)
        if self.args.seed is not None:
            torch.manual_seed(self.args.seed)
            np.random.seed(self.args.seed)
            random.seed(self.args.seed)

        self.model_config = model_and_diffusion_defaults()
        self.model_config.update(
            {
                "attention_resolutions": "32, 16, 8",
                "class_cond": self.args.model_output_size == 512,
                "diffusion_steps": 1000,
                "rescale_timesteps": True,
                "timestep_respacing": self.args.timestep_respacing,
                "image_size": self.args.model_output_size,
                "learn_sigma": True,
                "noise_schedule": "linear",
                "num_channels": 256,
                "num_head_channels": 64,
                "num_res_blocks": 2,
                "resblock_updown": True,
                "use_fp16": True,
                "use_scale_shift_norm": True,
            }
        )

        self.classifier_config = classifier_defaults()
        self.classifier_config.update(
            {
                "image_size": self.args.model_output_size,
            }
        )

        # Load models
        self.device = torch.device(
            f"cuda:{self.args.gpu_id}" if torch.cuda.is_available() else "cpu"
        )
        print("Using device:", self.device)

        self.model, self.diffusion = create_model_and_diffusion(**self.model_config)
        self.model.load_state_dict(
            torch.load(
                "checkpoints/256x256_diffusion_uncond.pt"
                if self.args.model_output_size == 256
                else "checkpoints/512x512_diffusion.pt",
                map_location="cpu",
            )
        )
        # self.model.requires_grad_(False).eval().to(self.device)
        self.model.eval().to(self.device)
        for name, param in self.model.named_parameters():
            if "qkv" in name or "norm" in name or "proj" in name:
                param.requires_grad_()
        if self.model_config["use_fp16"]:
            self.model.convert_to_fp16()

        self.classifier = create_classifier(**self.classifier_config)
        self.classifier.load_state_dict(
            torch.load("checkpoints/256x256_classifier.pt", map_location="cpu")
        )
        # self.classifier.requires_grad_(False).eval().to(self.device)


        # self.classifier = create_classifier_ours()

        self.classifier.eval().to(self.device)
        if self.classifier_config["classifier_use_fp16"]:
            self.classifier.convert_to_fp16()

        self.clip_model = (
            clip.load("ViT-B/16", device=self.device, jit=False)[0].eval().requires_grad_(False)
        )
        self.clip_size = self.clip_model.visual.input_resolution
        self.clip_normalize = transforms.Normalize(
            mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]
        )
        self.to_tensor = transforms.ToTensor()
        self.lpips_model = lpips.LPIPS(net="vgg").to(self.device)

        self.image_augmentations = ImageAugmentations(self.clip_size, self.args.aug_num)
        self.metrics_accumulator = MetricsAccumulator()

        self.hf_loss = HighFrequencyLoss()


    def unscale_timestep(self, t):
        unscaled_timestep = (t * (self.diffusion.num_timesteps / 1000)).long()

        return unscaled_timestep


    def clip_loss(self, x_in, text_embed):
        clip_loss = torch.tensor(0)

        if self.mask is not None:
            masked_input = x_in * self.mask
        else:
            masked_input = x_in
        augmented_input = self.image_augmentations(masked_input).add(1).div(2) # shape: [N,C,H,W], range: [0,1]
        clip_in = self.clip_normalize(augmented_input)
        # pdb.set_trace()
        image_embeds = self.clip_model.encode_image(clip_in).float()
        dists = d_clip_loss(image_embeds, text_embed)

        # We want to sum over the averages
        for i in range(self.args.batch_size):
            # We want to average at the "augmentations level"
            clip_loss = clip_loss + dists[i :: self.args.batch_size].mean()

        return clip_loss

    def unaugmented_clip_distance(self, x, text_embed):
        x = F.resize(x, [self.clip_size, self.clip_size])
        image_embeds = self.clip_model.encode_image(x).float()
        dists = d_clip_loss(image_embeds, text_embed)

        return dists.item()

    def model_fn(self, x,t,y=None):
        return self.model(x, t, y if self.args.class_cond else None)

    def edit_image_by_prompt(self):
        if self.args.image_guide:
            img_guidance = Image.open(self.args.prompt).convert('RGB')
            img_guidance = img_guidance.resize((224,224), Image.LANCZOS)  # type: ignore
            img_guidance = self.clip_normalize(self.to_tensor(img_guidance).unsqueeze(0)).to(self.device)
            text_embed = self.clip_model.encode_image(img_guidance).float()

        else:
            text_embed = self.clip_model.encode_text(
                clip.tokenize(self.args.prompt).to(self.device)
            ).float()

        self.image_size = (self.model_config["image_size"], self.model_config["image_size"])
        self.init_image_pil = Image.open(self.args.init_image).convert("RGB")
        self.init_image_pil = self.init_image_pil.resize(self.image_size, Image.LANCZOS)  # type: ignore
        self.init_image = (
            TF.to_tensor(self.init_image_pil).to(self.device).unsqueeze(0).mul(2).sub(1)
        )
        self.init_image_pil_2 = Image.open(self.args.init_image_2).convert("RGB")
        if self.args.rotate_obj:
            # angle = random.randint(-45,45)
            angle = self.args.angle
            self.init_image_pil_2 = self.init_image_pil_2.rotate(angle)
        self.init_image_pil_2 = self.init_image_pil_2.resize(self.image_size, Image.LANCZOS)  # type: ignore
        self.init_image_2 = (
            TF.to_tensor(self.init_image_pil_2).to(self.device).unsqueeze(0).mul(2).sub(1)
        )

        '''
        # Init with the inpainting image
        self.init_image_pil_ = Image.open('output/ImageNet-S_val/bad_case_RN50/ILSVRC2012_val_00013212/ranked/08480_output_i_0_b_0.png').convert("RGB")
        self.init_image_pil_ = self.init_image_pil_.resize(self.image_size, Image.LANCZOS)  # type: ignore
        self.init_image_ = (
            TF.to_tensor(self.init_image_pil_).to(self.device).unsqueeze(0).mul(2).sub(1)
        )
        '''

        if self.args.export_assets:
            img_path = self.assets_path / Path(self.args.output_file)
            self.init_image_pil.save(img_path, quality=100)

        self.mask = torch.ones_like(self.init_image, device=self.device)
        self.mask_pil = None
        if self.args.mask is not None:
            self.mask_pil = Image.open(self.args.mask).convert("RGB")
            if self.args.rotate_obj:
                self.mask_pil = self.mask_pil.rotate(angle)
            if self.mask_pil.size != self.image_size:
                self.mask_pil = self.mask_pil.resize(self.image_size, Image.NEAREST)  # type: ignore
            if self.args.random_position:
                bbox = find_bbox(np.array(self.mask_pil))
                print(bbox)

            image_mask_pil_binarized = ((np.array(self.mask_pil) > 0.5) * 255).astype(np.uint8)
            # image_mask_pil_binarized = cv2.dilate(image_mask_pil_binarized, np.ones((50,50), np.uint8), iterations=1)
            if self.args.invert_mask:
                image_mask_pil_binarized = 255 - image_mask_pil_binarized
                self.mask_pil = TF.to_pil_image(image_mask_pil_binarized)
            self.mask = TF.to_tensor(Image.fromarray(image_mask_pil_binarized))
            self.mask = self.mask[0, ...].unsqueeze(0).unsqueeze(0).to(self.device)
            # self.mask[:] = 1

            if self.args.random_position:
                # print(self.init_image_2.shape, self.init_image_2.max(), self.init_image_2.min())
                # print(self.mask.shape, self.mask.max(), self.mask.min())
                # cv2.imwrite('tmp/init_before.jpg', np.transpose(((self.init_image_2+1)/2*255).cpu().numpy()[0], (1,2,0))[:,:,::-1])
                # cv2.imwrite('tmp/mask_before.jpg', (self.mask*255).cpu().numpy()[0][0])
                self.init_image_2, self.mask = change_place(self.init_image_2, self.mask, bbox, self.args.invert_mask)
                # cv2.imwrite('tmp/init_after.jpg', np.transpose(((self.init_image_2+1)/2*255).cpu().numpy()[0], (1,2,0))[:,:,::-1])
                # cv2.imwrite('tmp/mask_after.jpg', (self.mask*255).cpu().numpy()[0][0])

            if self.args.export_assets:
                mask_path = self.assets_path / Path(
                    self.args.output_file.replace(".png", "_mask.png")
                )
                self.mask_pil.save(mask_path, quality=100)

        def class_guided(x, y, t):
            assert y is not None
            with torch.enable_grad():
                x_in = x.detach().requires_grad_(True)
                # logits = self.classifier(x_in, t)
                logits = self.classifier(x_in)
                log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                selected = log_probs[range(len(logits)), y.view(-1)]
                loss = selected.sum()

                return -torch.autograd.grad(loss, x_in)[0] * self.args.classifier_scale

        def cond_fn(x, t, y=None):
            if self.args.prompt == "":
                return torch.zeros_like(x)
            # pdb.set_trace()
            with torch.enable_grad():
                x = x.detach().requires_grad_()

                t_unscale = self.unscale_timestep(t)

                '''
                out = self.diffusion.p_mean_variance(
                    self.model, x, t, clip_denoised=False, model_kwargs={"y": y}
                )
                '''
                out = self.diffusion.p_mean_variance(
                    self.model, x, t_unscale, clip_denoised=False, model_kwargs={"y": None}
                )

                fac = self.diffusion.sqrt_one_minus_alphas_cumprod[t_unscale[0].item()]
                # x_in = out["pred_xstart"] * fac + x * (1 - fac)
                x_in = out["pred_xstart"] # Revised by XX, 2022.07.14

                loss = torch.tensor(0)
                if self.args.classifier_scale != 0 and y is not None:
                    # gradient_class_guided = class_guided(x, y, t)
                    gradient_class_guided = class_guided(x_in, y, t)

                if self.args.background_complex != 0:
                    if self.args.hard:
                        loss = loss - self.args.background_complex*self.hf_loss((x_in+1.)/2.)
                    else:
                        loss = loss + self.args.background_complex*self.hf_loss((x_in+1.)/2.)

                if self.args.clip_guidance_lambda != 0:
                    clip_loss = self.clip_loss(x_in, text_embed) * self.args.clip_guidance_lambda
                    loss = loss + clip_loss
                    self.metrics_accumulator.update_metric("clip_loss", clip_loss.item())

                if self.args.range_lambda != 0:
                    r_loss = range_loss(out["pred_xstart"]).sum() * self.args.range_lambda
                    loss = loss + r_loss
                    self.metrics_accumulator.update_metric("range_loss", r_loss.item())

                if self.args.background_preservation_loss:
                    x_in = out["pred_xstart"] * fac + x * (1 - fac)
                    if self.mask is not None:
                        # masked_background = x_in * (1 - self.mask)
                        masked_background = x_in * self.mask # 2022.07.19
                    else:
                        masked_background = x_in

                    if self.args.lpips_sim_lambda:
                        '''
                        loss = (
                            loss
                            + self.lpips_model(masked_background, self.init_image).sum()
                            * self.args.lpips_sim_lambda
                        )
                        '''
                        # 2022.07.19
                        loss = (
                            loss
                            + self.lpips_model(masked_background, self.init_image*self.mask).sum()
                            * self.args.lpips_sim_lambda
                        )
                    if self.args.l2_sim_lambda:
                        '''
                        loss = (
                            loss
                            + mse_loss(masked_background, self.init_image) * self.args.l2_sim_lambda
                        )
                        '''
                        # 2022.07.19
                        loss = (
                            loss
                            + mse_loss(masked_background, self.init_image*self.mask) * self.args.l2_sim_lambda
                        )


                if self.args.classifier_scale != 0 and y is not None:
                    return -torch.autograd.grad(loss, x)[0] + gradient_class_guided
                else:
                    return -torch.autograd.grad(loss, x)[0]

        @torch.no_grad()
        def postprocess_fn(out, t):
            if self.args.coarse_to_fine:
                if t > 50:
                    kernel = 51
                elif t > 35:
                    kernel = 31
                else:
                    kernel = 0
                if kernel > 0:
                    max_pool = torch.nn.MaxPool2d(kernel_size=kernel, stride=1, padding=int((kernel-1)/2))
                    self.mask_d = 1 - self.mask
                    self.mask_d = max_pool(self.mask_d)
                    self.mask_d = 1 - self.mask_d
                else:
                    self.mask_d = self.mask
            else:
                self.mask_d = self.mask

            if self.mask is not None:
                background_stage_t = self.diffusion.q_sample(self.init_image_2, t[0])
                background_stage_t = torch.tile(
                    background_stage_t, dims=(self.args.batch_size, 1, 1, 1)
                )
                out["sample"] = out["sample"] * self.mask_d + background_stage_t * (1 - self.mask_d)

            return out

        save_image_interval = self.diffusion.num_timesteps // 5
        for iteration_number in range(self.args.iterations_num):
            print(f"Start iterations {iteration_number}")

            sample_func = (
                self.diffusion.ddim_sample_loop_progressive
                if self.args.ddim
                else self.diffusion.p_sample_loop_progressive
            )
            samples = sample_func(
                self.model_fn,
                (
                    self.args.batch_size,
                    3,
                    self.model_config["image_size"],
                    self.model_config["image_size"],
                ),
                clip_denoised=False,
                # model_kwargs={}
                # if self.args.model_output_size == 256
                # else {
                #     "y": torch.zeros([self.args.batch_size], device=self.device, dtype=torch.long)
                # },
                model_kwargs={}
                if self.args.classifier_scale == 0
                else {"y": self.args.y*torch.ones([self.args.batch_size], device=self.device, dtype=torch.long)},
                cond_fn=cond_fn,
                device=self.device,
                progress=True,
                skip_timesteps=self.args.skip_timesteps,
                init_image=self.init_image,
                # init_image=self.init_image_,
                postprocess_fn=None if self.args.local_clip_guided_diffusion else postprocess_fn,
                randomize_class=True if self.args.classifier_scale == 0 else False,
            )

            intermediate_samples = [[] for i in range(self.args.batch_size)]
            total_steps = self.diffusion.num_timesteps - self.args.skip_timesteps - 1
            for j, sample in enumerate(samples):
                should_save_image = j % save_image_interval == 0 or j == total_steps
                if should_save_image or self.args.save_video:
                    self.metrics_accumulator.print_average_metric()

                    for b in range(self.args.batch_size):
                        pred_image = sample["pred_xstart"][b]
                        visualization_path = Path(
                            os.path.join(self.args.output_path, self.args.output_file)
                        )
                        visualization_path = visualization_path.with_stem(
                            f"{visualization_path.stem}_i_{iteration_number}_b_{b}"
                        )
                        if (
                            self.mask is not None
                            and self.args.enforce_background
                            and j == total_steps
                            and not self.args.local_clip_guided_diffusion
                        ):
                            pred_image = (
                                self.init_image_2[0] * (1 - self.mask[0]) + pred_image * self.mask[0]
                            )
                        '''
                        if j == total_steps:
                            pdb.set_trace()
                        pred_image = (
                                self.init_image_2[0] * (1 - self.mask[0]) + pred_image * self.mask[0]
                            )
                        '''
                        pred_image = pred_image.add(1).div(2).clamp(0, 1)
                        pred_image_pil = TF.to_pil_image(pred_image)
                        masked_pred_image = self.mask * pred_image.unsqueeze(0)
                        final_distance = self.unaugmented_clip_distance(
                            masked_pred_image, text_embed
                        )
                        formatted_distance = f"{final_distance:.4f}"

                        if self.args.export_assets:
                            pred_path = self.assets_path / visualization_path.name
                            pred_image_pil.save(pred_path, quality=100)

                        if j == total_steps:
                            path_friendly_distance = formatted_distance.replace(".", "")
                            ranked_pred_path = self.ranked_results_path / (
                                path_friendly_distance + "_" + visualization_path.name
                            )
                            pred_image_pil.save(ranked_pred_path, quality=100)

                        intermediate_samples[b].append(pred_image_pil)
                        if should_save_image:
                            show_editied_masked_image(
                                title=self.args.prompt,
                                source_image=self.init_image_pil,
                                edited_image=pred_image_pil,
                                mask=self.mask_pil,
                                path=visualization_path,
                                distance=formatted_distance,
                            )

            if self.args.save_video:
                for b in range(self.args.batch_size):
                    video_name = self.args.output_file.replace(
                        ".png", f"_i_{iteration_number}_b_{b}.avi"
                    )
                    video_path = os.path.join(self.args.output_path, video_name)
                    save_video(intermediate_samples[b], video_path)

        visualize_size = (256,256)
        img_ori = cv2.imread(self.args.init_image_2) 
        img_ori = cv2.resize(img_ori, visualize_size)
        mask = cv2.imread(self.args.mask)
        mask = cv2.resize(mask, visualize_size)
        imgs = [img_ori, mask]
        for ii, img_name in enumerate(os.listdir(os.path.join(self.args.output_path, 'ranked'))):
            img_path = os.path.join(self.args.output_path, 'ranked', img_name)
            img = cv2.imread(img_path)
            img = cv2.resize(img, visualize_size)
            imgs.append(img)
            if ii >= 7:
                break

        img_whole = cv2.hconcat(imgs[2:])
        '''
        img_name = self.args.output_path.split('/')[-2]+'/'
        if self.args.coarse_to_fine:
            if self.args.clip_guidance_lambda == 0:
                prompt = 'coarse_to_fine_no_clip'
            else:
                prompt = 'coarse_to_fine'
        elif self.args.image_guide:
            prompt = 'image_guide'
        elif self.args.clip_guidance_lambda == 0:
            prompt = 'no_clip_guide'
        else:
            prompt = 'text_guide'
        '''

        cv2.imwrite(os.path.join(self.args.final_save_root, 'edited.png'), img_whole, [int(cv2.IMWRITE_PNG_COMPRESSION), 0])


    def reconstruct_image(self):
        init = Image.open(self.args.init_image).convert("RGB")
        init = init.resize(
            self.image_size,  # type: ignore
            Image.LANCZOS,
        )
        init = TF.to_tensor(init).to(self.device).unsqueeze(0).mul(2).sub(1)

        samples = self.diffusion.p_sample_loop_progressive(
            self.model,
            (1, 3, self.model_config["image_size"], self.model_config["image_size"],),
            clip_denoised=False,
            model_kwargs={}
            if self.args.model_output_size == 256
            else {"y": torch.zeros([self.args.batch_size], device=self.device, dtype=torch.long)},
            cond_fn=None,
            progress=True,
            skip_timesteps=self.args.skip_timesteps,
            init_image=init,
            randomize_class=True,
        )
        save_image_interval = self.diffusion.num_timesteps // 5
        max_iterations = self.diffusion.num_timesteps - self.args.skip_timesteps - 1

        for j, sample in enumerate(samples):
            if j % save_image_interval == 0 or j == max_iterations:
                print()
                filename = os.path.join(self.args.output_path, self.args.output_file)
                TF.to_pil_image(sample["pred_xstart"][0].add(1).div(2).clamp(0, 1)).save(filename)