steered_diffusion_dataset.py

# Copyright (C) 2023-2024 Mitsubishi Electric Research Laboratories (MERL)
# Copyright (C) 2022-2023 Omri Avrahami

# SPDX-License-Identifier: AGPL-3.0-or-later
# SPDX-License-Identifier: MIT

# Code adapted  from https://github.com/omriav/blended-diffusion -- MIT License


import glob
import os
from pathlib import Path

import clip
import lpips
import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms.functional as F
import tqdm
import yaml
from numpy import random
from PIL import Image, ImageOps
from torchvision.transforms import functional as TF
from tqdm import tqdm

from guided_diffusion.guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults
from losses.ssim import SSIM

torch.autograd.set_detect_anomaly(True)

from losses.Full_loss import Full_loss


class ImageEditor:
    def __init__(self, args) -> None:
        self.args = args
        self.data = args["data"]
        self.params = args["params"]
        self.network = args["diffusion_network"]
        self.checkpoints = args["checkpoints"]
        self.Full_loss = Full_loss(args)
        count = 0
        out_path = os.path.join(self.params["results_dir"])
        if os.path.exists(out_path) == False:
            os.makedirs(out_path)
        self.data["output_path"] = out_path

        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(self.network)
        gpu_id = self.args["gpu_id"]
        self.device = torch.device(f"cuda:{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(self.checkpoints["ffhq"]))
        self.model.requires_grad_(False).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.image_size = (self.params["image_size"], self.params["image_size"])

    def edit_image(self):

        if self.params["use_ddim"]:
            self.init_ddim = self.diffusion.ddim_reverse_sample_loop(
                self.model,
                self.init_image_transfer,
            )["sample"]

        batch = self.params["batch_size"]

        images = glob.glob(os.path.join(self.data["data_fold"], "images", "*"))
        if self.params["cond"] == "inpaint":
            masks = glob.glob(os.path.join(self.data["data_fold"], "masks", "*"))
        else:
            masks = [None] * len(images)

        for img, mask in zip(images, masks):
            init_image = Image.open(img).convert("RGB")
            init_image = init_image.resize(self.image_size, Image.BICUBIC)

            init_image = TF.to_tensor(init_image).to(self.device).unsqueeze(0).mul(2).sub(1)
            self.init_image = init_image

            if self.params["cond"] == "inpaint":
                mask_image = Image.open(mask).convert("L")
                mask_image = mask_image.resize(self.image_size, Image.BICUBIC)
                mask_image = TF.to_tensor(mask_image).to(self.device).unsqueeze(0)
                mask_image = mask_image.repeat(1, 3, 1, 1)
            else:
                mask_image = None

            image_name = img.split("/")[-1].strip(".jpg")
            model_kwargs = {
                "cond": self.params["cond"],
                "mask_image": mask_image,
                "init_image": self.init_image,
                "num_iters": 1,
                "factor": self.params["scale_factor"],
            }

            shape = (
                batch,
                3,
                self.model_config["image_size"],
                self.model_config["image_size"],
            )

            model_kwargs["dest_fold"] = os.path.join("./results", model_kwargs["cond"])

            samples = self.diffusion.conditional_sample_loop_progressive(
                model=self.model,
                shape=shape,
                clip_denoised=False,
                model_kwargs=model_kwargs,
                noise=None,
                cond_fn=self.Full_loss,
                progress=True,
            )

            for count, sample in enumerate(samples):
                pred_image = sample["sample"]
                pred_image = pred_image.add(1).div(2).clamp(0, 1)
                degraded_image = sample["degraded"].add(1).div(2).clamp(0, 1)
                dest_fold = os.path.join(self.params["results_dir"], model_kwargs["cond"], image_name)
                if os.path.exists(dest_fold) == False:
                    os.makedirs(dest_fold)
                for j in range(pred_image.shape[0]):
                    degraded_pred = torch.cat([degraded_image[j], pred_image[j]], dim=2)
                    pred_image_pil = TF.to_pil_image(degraded_pred)
                    pred_path = os.path.join(dest_fold, str(j) + ".jpg")
                    pred_image_pil.save(pred_path)