sample.py

import gc
import io
import math
import sys

from PIL import Image, ImageOps, ImageDraw
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 torchvision.ops import masks_to_boxes

from tqdm.notebook import tqdm

import numpy as np

from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults, classifier_defaults, \
    create_classifier, create_model, create_gaussian_diffusion

from omegaconf import OmegaConf
from ldm.util import instantiate_from_config

from einops import rearrange
from math import log2, sqrt

import argparse
import pickle

import os

from transformers import CLIPTokenizer, CLIPTextModel

# argument parsing

parser = argparse.ArgumentParser()

parser.add_argument('--model_path', type=str, default = 'diffusion.pt',
                   help='path to the diffusion model')

parser.add_argument('--kl_path', type=str, default = 'kl.pt',
                   help='path to the LDM first stage model')

parser.add_argument('--text', type = str, required = False, default = '',
                    help='your text prompt')

parser.add_argument('--classifier', type=str, default = '',
                   help='path to the classifier model')

parser.add_argument('--classifier_scale', type = int, required = False, default = 100,
                    help='amount of classifier guidance')

parser.add_argument('--edit', type = str, required = False,
                    help='path to the image you want to edit (either an image file or .npy containing a numpy array of the image embeddings)')

parser.add_argument('--outpaint', type = str, required = False, default = '',
                    help='options: expand (all directions), wider, taller, left, right, top, bottom')

parser.add_argument('--mask', type = str, required = False,
                    help='path to a mask image. white pixels = keep, black pixels = discard. width = image width/8, height = image height/8')

parser.add_argument('--negative', type = str, required = False, default = '',
                    help='negative text prompt')

parser.add_argument('--init_image', type=str, required = False, default = None,
                   help='init image to use')

parser.add_argument('--skip_timesteps', type=int, required = False, default = 0,
                   help='how many diffusion steps are gonna be skipped')

parser.add_argument('--prefix', type = str, required = False, default = '',
                    help='prefix for output files')

parser.add_argument('--num_batches', type = int, default = 1, required = False,
                    help='number of batches')

parser.add_argument('--batch_size', type = int, default = 1, required = False,
                    help='batch size')

parser.add_argument('--width', type = int, default = 512, required = False,
                    help='image size of output (multiple of 8)')

parser.add_argument('--height', type = int, default = 512, required = False,
                    help='image size of output (multiple of 8)')

parser.add_argument('--seed', type = int, default=-1, required = False,
                    help='random seed')

parser.add_argument('--guidance_scale', type = float, default = 7.0, required = False,
                    help='classifier-free guidance scale')

parser.add_argument('--steps', type = int, default = 0, required = False,
                    help='number of diffusion steps')

parser.add_argument('--cpu', dest='cpu', action='store_true')

parser.add_argument('--ddim', dest='ddim', action='store_true')

parser.add_argument('--ddpm', dest='ddpm', action='store_true')


def set_requires_grad(model, value):
    for param in model.parameters():
        param.requires_grad = value


def do_load(args):
    device = torch.device('cuda:0' if (torch.cuda.is_available() and not args.cpu) else 'cpu')

    model_state_dict = torch.load(args.model_path, map_location='cpu')

    model_config = model_and_diffusion_defaults()
    model_config.update({
        'attention_resolutions': '32,16,8',
        'class_cond': False,
        'diffusion_steps': 1000,
        'rescale_timesteps': True,
        'timestep_respacing': '50',
        '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': None,
        'image_condition': True if model_state_dict['input_blocks.0.0.weight'].shape[1] == 8 else False,
        'super_res_condition': True if 'external_block.0.0.weight' in model_state_dict else False,
    })

    if args.cpu:
        model_config['use_fp16'] = False

    model = create_model(
        image_size=model_config["image_size"],
        num_channels=model_config["num_channels"],
        num_res_blocks=model_config["num_res_blocks"],
        channel_mult=model_config["channel_mult"],
        learn_sigma=model_config["learn_sigma"],
        class_cond=model_config["class_cond"],
        use_checkpoint=model_config["use_checkpoint"],
        attention_resolutions=model_config["attention_resolutions"],
        num_heads=model_config["num_heads"],
        num_head_channels=model_config["num_head_channels"],
        num_heads_upsample=model_config["num_heads_upsample"],
        use_scale_shift_norm=model_config["use_scale_shift_norm"],
        dropout=model_config["dropout"],
        resblock_updown=model_config["resblock_updown"],
        use_fp16=model_config["use_fp16"],
        use_spatial_transformer=model_config["use_spatial_transformer"],
        context_dim=model_config["context_dim"],
        clip_embed_dim=model_config["clip_embed_dim"],
        image_condition=model_config["image_condition"],
        super_res_condition=model_config["super_res_condition"],
    )
    model.load_state_dict(model_state_dict, strict=True)
    model.requires_grad_(False).eval().to(device)

    if model_config['use_fp16']:
        model.convert_to_fp16()
    else:
        model.convert_to_fp32()

    classifier = None

    def set_requires_grad(model, value):
        for param in model.parameters():
            param.requires_grad = value

    # load classifier
    if args.classifier:
        classifier_config = classifier_defaults()
        classifier_config['classifier_width'] = 128
        classifier_config['classifier_depth'] = 4
        classifier_config['classifier_attention_resolutions'] = '64,32,16,8'
        classifier = create_classifier(**classifier_config)
        classifier.load_state_dict(
            torch.load(args.classifier, map_location="cpu")
        )
        classifier.to(device)
        classifier.convert_to_fp16()
        classifier.eval()

    # vae
    kl_config = OmegaConf.load('kl.yaml')
    kl_sd = torch.load(args.kl_path, map_location="cpu")

    ldm = instantiate_from_config(kl_config.model)
    ldm.load_state_dict(kl_sd, strict=True)

    ldm.to(device)
    ldm.eval()
    ldm.requires_grad_(False)
    set_requires_grad(ldm, False)

    # clip
    clip_version = 'openai/clip-vit-large-patch14'
    clip_tokenizer = CLIPTokenizer.from_pretrained(clip_version)
    clip_transformer = CLIPTextModel.from_pretrained(clip_version)
    clip_transformer.eval().requires_grad_(False).to(device)

    return model, model_config, device, ldm, classifier, clip_tokenizer, clip_transformer


def do_run(args, model, model_config, device, ldm, classifier, clip_tokenizer, clip_transformer):
    if args.edit and not args.mask:
        from PyQt5.QtWidgets import QApplication, QMainWindow
        from PyQt5.QtGui import QPainter, QPen
        from PyQt5.QtCore import Qt, QPoint, QRect, QBuffer
        import PyQt5.QtGui as QtGui

        class Draw(QMainWindow):

            def __init__(self, width, height, im):
                super().__init__()
                self.drawing = False
                self.lastPoint = QPoint()

                self.qim = QtGui.QImage(im.tobytes("raw", "RGB"), im.width, im.height, QtGui.QImage.Format_RGB888)
                self.image = QtGui.QPixmap.fromImage(self.qim)

                canvas = QtGui.QImage(im.width, im.height, QtGui.QImage.Format_ARGB32)
                self.canvas = QtGui.QPixmap.fromImage(canvas)
                self.canvas.fill(Qt.transparent)

                self.setGeometry(0, 0, im.width, im.height)
                self.resize(self.image.width(), self.image.height())
                self.show()

            def paintEvent(self, event):
                painter = QPainter(self)
                painter.drawPixmap(QRect(0, 0, self.image.width(), self.image.height()), self.image)
                painter.drawPixmap(QRect(0, 0, self.canvas.width(), self.canvas.height()), self.canvas)

            def mousePressEvent(self, event):
                if event.button() == Qt.LeftButton:
                    self.drawing = True
                    self.lastPoint = event.pos()

            def mouseMoveEvent(self, event):
                if event.buttons() and Qt.LeftButton and self.drawing:
                    painter = QPainter(self.canvas)
                    painter.setPen(
                        QPen(Qt.red, (self.width() + self.height()) / 20, Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin))
                    painter.drawLine(self.lastPoint, event.pos())
                    self.lastPoint = event.pos()
                    self.update()

            def mouseReleaseEvent(self, event):
                if event.button == Qt.LeftButton:
                    self.drawing = False

            def getCanvas(self):
                image = self.canvas.toImage()
                buffer = QBuffer()
                buffer.open(QBuffer.ReadWrite)
                image.save(buffer, "PNG")
                pil_im = Image.open(io.BytesIO(buffer.data()))
                return pil_im

            def resizeEvent(self, event):
                self.image = QtGui.QPixmap.fromImage(self.qim)
                self.image = self.image.scaled(self.width(), self.height())

                canvas = QtGui.QImage(self.width(), self.height(), QtGui.QImage.Format_ARGB32)
                self.canvas = QtGui.QPixmap.fromImage(canvas)
                self.canvas.fill(Qt.transparent)

    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')

    print('Using device:', device)

    if args.ddpm:
        model_config['timestep_respacing'] = '1000'
    if args.ddim:
        if args.steps:
            model_config['timestep_respacing'] = 'ddim'+str(args.steps)
        else:
            model_config['timestep_respacing'] = 'ddim250'
    elif args.steps:
        model_config['timestep_respacing'] = str(args.steps)

    # Create output folders
    os.makedirs("output", exist_ok = True)
    os.makedirs("output_npy", exist_ok = True)

    # Load models
    diffusion = create_gaussian_diffusion(
        steps=model_config['diffusion_steps'],
        learn_sigma=model_config['learn_sigma'],
        noise_schedule=model_config['noise_schedule'],
        use_kl=model_config['use_kl'],
        predict_xstart=model_config['predict_xstart'],
        rescale_timesteps=model_config['rescale_timesteps'],
        rescale_learned_sigmas=model_config['rescale_learned_sigmas'],
        timestep_respacing=model_config['timestep_respacing'],
    )

    if args.seed >= 0:
        torch.manual_seed(args.seed)

    # clip context
    text = clip_tokenizer([args.text]*args.batch_size, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
    text_blank = clip_tokenizer([args.negative]*args.batch_size, truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
    text_tokens = text["input_ids"].to(device)
    text_blank_tokens = text_blank["input_ids"].to(device)

    text_emb = clip_transformer(input_ids=text_tokens).last_hidden_state
    text_emb_blank = clip_transformer(input_ids=text_blank_tokens).last_hidden_state

    image_embed = None

    # image context
    if args.edit:
        if args.edit.endswith('.npy'):
            with open(args.edit, 'rb') as f:
                im = np.load(f)
                im = torch.from_numpy(im).unsqueeze(0).to(device)
        else:
            input_image_pil = Image.open(fetch(args.edit)).convert('RGB')

            im = transforms.ToTensor()(input_image_pil).unsqueeze(0).to(device)
            im = 2*im-1
            im = ldm.encode(im).sample()

        if im.shape[3] < 64:
            im2 = torch.zeros(1,4,im.shape[2],64)
            x = (64-im.shape[3])//2
            im2[:,:,:,x:x+im.shape[3]] = im
            im = im2

        if im.shape[2] < 64:
            im2 = torch.zeros(1,4,64,im.shape[3])
            y = (64-im.shape[2])//2
            im2[:,:,y:y+im.shape[2],:] = im
            im = im2


        if args.outpaint == 'expand':
            input_image = torch.zeros(1, 4, im.shape[2]+64, im.shape[3]+64, device=device)
            input_image[:,:,32:32+im.shape[2],32:32+im.shape[3]] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2]+64, im.shape[3]+64, device=device, dtype=torch.bool)
            input_image_mask[:,:,32:32+im.shape[2],32:32+im.shape[3]] = True
        elif args.outpaint == 'wider':
            input_image = torch.zeros(1, 4, im.shape[2], im.shape[3]+64, device=device)
            input_image[:,:,:,32:32+im.shape[3]] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2], im.shape[3]+64, device=device, dtype=torch.bool)
            input_image_mask[:,:,:,32:32+im.shape[3]] = True
        elif args.outpaint == 'taller':
            input_image = torch.zeros(1, 4, im.shape[2]+64, im.shape[3], device=device)
            input_image[:,:,32:32+im.shape[2],:] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2]+64, im.shape[3], device=device, dtype=torch.bool)
            input_image_mask[:,:,32:32+im.shape[2],:] = True
        elif args.outpaint == 'left':
            input_image = torch.zeros(1, 4, im.shape[2], im.shape[3]+32, device=device)
            input_image[:,:,:,32:32+im.shape[3]] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2], im.shape[3]+32, device=device, dtype=torch.bool)
            input_image_mask[:,:,:,32:32+im.shape[3]] = True
        elif args.outpaint == 'right':
            input_image = torch.zeros(1, 4, im.shape[2], im.shape[3]+32, device=device)
            input_image[:,:,:,0:im.shape[3]] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2], im.shape[3]+32, device=device, dtype=torch.bool)
            input_image_mask[:,:,:,0:im.shape[3]] = True
        elif args.outpaint == 'top':
            input_image = torch.zeros(1, 4, im.shape[2]+32, im.shape[3], device=device)
            input_image[:,:,32:32+im.shape[2],:] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2]+32, im.shape[3], device=device, dtype=torch.bool)
            input_image_mask[:,:,32:32+im.shape[2],:] = True
        elif args.outpaint == 'bottom':
            input_image = torch.zeros(1, 4, im.shape[2]+32, im.shape[3], device=device)
            input_image[:,:,0:im.shape[2],:] = im
            input_image_mask = torch.zeros(1, 1, im.shape[2]+32, im.shape[3], device=device, dtype=torch.bool)
            input_image_mask[:,:,0:im.shape[2],:] = True
        else:
            input_image = im
            input_image_mask = torch.ones(1,1,im.shape[2], im.shape[3], device=device, dtype=torch.bool)

        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

        if args.mask:
            mask_image = Image.open(fetch(args.mask)).convert('L')
            mask_image = mask_image.resize((input_image.shape[3],input_image.shape[2]), Image.ANTIALIAS)
            mask = transforms.ToTensor()(mask_image).unsqueeze(0).to(device)
        else:
            print('draw the area for inpainting, then close the window')
            app = QApplication(sys.argv)
            d = Draw(args.width, args.height, input_image_pil)
            app.exec_()
            mask_image = d.getCanvas().convert('L').point( lambda p: 255 if p < 1 else 0 )
            mask_image.save('mask.png')
            mask_image = mask_image.resize((input_image.shape[3],input_image.shape[2]), Image.ANTIALIAS)
            mask = transforms.ToTensor()(mask_image).unsqueeze(0).to(device)

        mask1 = (mask > 0.5)
        input_image_mask *= mask1

        #mask1 = mask1.float()
        #input_image *= mask1


        image_embed = torch.cat(args.batch_size*2*[input_image], dim=0).float()

    elif model_config['image_condition']:
        # using inpaint model but no image is provided
        image_embed = torch.zeros(args.batch_size*2, 4, args.height//8, args.width//8, device=device)

    kwargs = {
        "context": torch.cat([text_emb, text_emb_blank], dim=0).float(),
        "clip_embed": 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 + args.guidance_scale * (cond_eps - uncond_eps)
        eps = torch.cat([half_eps, half_eps], dim=0)
        return torch.cat([eps, rest], dim=1)

    cond_fn = None

    if args.classifier:
        def cond_fn(x, t, context=None, clip_embed=None, image_embed=None):
            with torch.enable_grad():
                x_in = x[:x.shape[0]//2].detach().requires_grad_(True)
                logits = classifier(x_in, t)
                log_probs = F.log_softmax(logits, dim=-1)
                selected = log_probs[range(len(logits)), torch.ones(x_in.shape[0], dtype=torch.long)]
                return torch.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale

    cur_t = None

    if args.ddpm:
        sample_fn = diffusion.ddpm_sample_loop_progressive
    elif args.ddim:
        sample_fn = diffusion.ddim_sample_loop_progressive
    else:
        sample_fn = diffusion.plms_sample_loop_progressive

    def save_sample(i, samples, square=None):
        for k, image in enumerate(samples):
            image_scaled = image/0.18215
            im = image_scaled.unsqueeze(0)
            out = ldm.decode(im)

            npy_filename = f'output_npy/{args.prefix}{i * args.batch_size + k:05}.npy'
            with open(npy_filename, 'wb') as outfile:
                np.save(outfile, image_scaled.detach().cpu().numpy())

            out = TF.to_pil_image(out.squeeze(0).add(1).div(2).clamp(0, 1))

            if square is not None:
                outdraw = ImageDraw.Draw(out)
                outdraw.rectangle([(square[0]*8, square[1]*8),(square[0]*8+512, square[1]*8+512)], fill=None, outline ="red")

            filename = f'output/{args.prefix}{i * args.batch_size + k:05}.png'
            out.save(filename)

    if args.init_image:
        init = Image.open(args.init_image).convert('RGB')
        init = init.resize((int(args.width),  int(args.height)), Image.LANCZOS)
        init = TF.to_tensor(init).to(device).unsqueeze(0).clamp(0,1)
        h = ldm.encode(init * 2 - 1).sample() *  0.18215
        init = torch.cat(args.batch_size*2*[h], dim=0)
    else:
        init = None

    overlap = 32

    if args.edit:
        for i in range(args.num_batches):
            output = input_image.detach().clone()
            output *= input_image_mask.repeat(1, 4, 1, 1).float()

            mask = input_image_mask.detach().clone()

            box = masks_to_boxes(~mask.squeeze(0))[0]

            x0 = int(box[0])
            y0 = int(box[1])
            x1 = int(box[2] + 1)
            y1 = int(box[3] + 1)

            x_num = math.ceil(((x1-x0)-overlap)/(64-overlap))
            y_num = math.ceil(((y1-y0)-overlap)/(64-overlap))

            if x_num < 1:
                x_num = 1
            if y_num < 1:
                y_num = 1

            for y in range(y_num):
                for x in range(x_num):
                    offsetx = x0 + x*(64-overlap)
                    offsety = y0 + y*(64-overlap)

                    if offsetx + 64 > x1:
                        offsetx = x1 - 64
                    if offsetx < 0:
                        offsetx = 0

                    if offsety + 64 > y1:
                        offsety = y1 - 64
                    if offsety < 0:
                        offsety = 0

                    patch_input = output[:,:, offsety:offsety+64, offsetx:offsetx+64]
                    patch_mask = mask[:,:, offsety:offsety+64, offsetx:offsetx+64]

                    if not torch.any(~patch_mask):
                        # region does not require any inpainting
                        output[:,:, offsety:offsety+64, offsetx:offsetx+64] = patch_input
                        continue

                    mask[:,:, offsety:offsety+64, offsetx:offsetx+64] = True

                    patch_init = None

                    if args.skip_timesteps > 0:
                        patch_init = input_image[:,:, offsety:offsety+64, offsetx:offsetx+64]
                        patch_init = torch.cat([patch_init, patch_init], dim=0)

                    skip_timesteps = args.skip_timesteps

                    if not torch.any(patch_mask):
                        # region has no input image, cannot use init
                        patch_init = None
                        skip_timesteps = 0

                    patch_kwargs = {
                        "context": kwargs["context"],
                        "clip_embed": None,
                        "image_embed": torch.cat([patch_input, patch_input], dim=0)
                    }

                    cur_t = diffusion.num_timesteps - 1

                    samples = sample_fn(
                        model_fn,
                        (2, 4, 64, 64),
                        clip_denoised=False,
                        model_kwargs=patch_kwargs,
                        cond_fn=cond_fn,
                        device=device,
                        progress=True,
                        init_image=patch_init,
                        skip_timesteps=skip_timesteps,
                    )

                    for j, sample in enumerate(samples):
                        cur_t -= 1
                        output[0,:, offsety:offsety+64, offsetx:offsetx+64] = sample['pred_xstart'][0]
                        if j % 25 == 0:
                            save_sample(i, output, square=(offsetx, offsety))

                    save_sample(i, output)

    else:
        for i in range(args.num_batches):
            cur_t = diffusion.num_timesteps - 1

            samples = sample_fn(
                model_fn,
                (args.batch_size*2, 4, int(args.height/8), int(args.width/8)),
                clip_denoised=False,
                model_kwargs=kwargs,
                cond_fn=cond_fn,
                device=device,
                progress=True,
                init_image=init,
                skip_timesteps=args.skip_timesteps,
            )

            for j, sample in enumerate(samples):
                cur_t -= 1
                if j % 20 == 0:
                    save_sample(i, sample['pred_xstart'][:args.batch_size])

            save_sample(i, sample['pred_xstart'][:args.batch_size])

if __name__ == '__main__':
    args = parser.parse_args()
    params = do_load(args)
    gc.collect()
    do_run(args, *params)