encode.py

import random
import time
from typing import Optional, Union, Tuple, List, Callable, Dict
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler, AutoencoderKL, UNet2DConditionModel,DDIMInverseScheduler
import torch.nn.functional as nnf
import numpy as np
import abc
from scipy import  stats
import shutil
from albumentations import Compose, RandomBrightnessContrast, HorizontalFlip, FancyPCA, HueSaturationValue, OneOf, \
    ToGray, ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur, Resize, RandomRotate90, Flip, \
    RGBShift, Sharpen, CenterCrop
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
from albumentations import Compose, RandomBrightnessContrast, HorizontalFlip, FancyPCA, HueSaturationValue, OneOf, \
    ToGray, ShiftScaleRotate, ImageCompression, PadIfNeeded, GaussNoise, GaussianBlur, Resize, RandomRotate90, Flip, \
    RGBShift, Sharpen, CenterCrop,RandomCrop,OpticalDistortion,GridDistortion
from diffusers.image_processor import VaeImageProcessor
from torch.optim.adam import Adam
from PIL import Image
from PIL import Image, ImageOps
import cv2
from tqdm import tqdm
import argparse
import os
from transformers import CLIPTokenizer, CLIPTextModel
from utils import *
from DDIM_watermark import PAI
from dataset.watermark_dataloader import PAI_Watermark_dataset
from einops import rearrange, repeat
from scipy.stats import norm

M_inject = 0.1
MY_TOKEN = ''
height = 512
width = 512
LOW_RESOURCE = False
GUIDANCE_SCALE = 7.5
MAX_NUM_WORDS = 77

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--wm_steps', type=int, default=5, help='add watermark steps')
    parser.add_argument('--start_wm_step', type=int, default=1, help='start to add watermark')
    parser.add_argument('--dataset', default="", type=str, help="dataset txt path")  #
    parser.add_argument('--save_path', type=str, default='./output', help='text prompt of the public key')
    parser.add_argument('--num_steps', type=int, default=50, help='sampling step of DDIM')
    parser.add_argument("--gpu_num", default="0", type=str,help="gpu number")
    parser.add_argument("--batch_size",type=int,default=6,help="batch size")
    args = parser.parse_args()
    device_ids = list(map(int, args.gpu_num.split(',')))
    device = torch.device("cuda", device_ids[0])

    name = "stable-diffusion-2-1"

    data_loader         = DataLoader(PAI_Watermark_dataset(args.dataset), shuffle=False, batch_size=args.batch_size)
    batch_size          = args.batch_size
    wm_steps            = args.wm_steps
    start_wm_step       = args.start_wm_step
    pi                  = PAI(name, 
                              NUM_DDIM_STEPS = args.num_steps, 
                              start_wm_step  = start_wm_step, 
                              GUIDANCE_SCALE = 7.5,
                              device         = device,
                              M_INJECT       = M_inject)

    record = []
    record_dict = {}
    save_wm_path = f"{args.save_path}/wm-img"
    os.makedirs(save_wm_path, exist_ok=True)
    shape       = (pi.unet.in_channels, height // 8, width // 8)
    key         = torch.randn(shape).to(device)
    key         = repeat(key, 'c h w -> b c h w', b=1)
    for epoch, prompts in enumerate(data_loader):

        init_latents_noises = []
        public_times        = []
        for _ in prompts:
            public_time = int(time.time())
            np.random.seed(public_time)
            U1                 = np.random.uniform(0, 1, 
                                (1, pi.unet.in_channels, 
                                 height // 8, 
                                 width // 8))
            U2                 = norm.cdf(key.cpu().numpy())
            init_latents_noise = np.sqrt(-2 * np.log(U1)) * np.cos(2 * np.pi * U2)
            init_latents_noise = torch.Tensor(init_latents_noise).to(device)
            init_latents_noises.append(init_latents_noise)
            public_times.append(public_time)
        init_latents_noise           = torch.cat(init_latents_noises)
        '''add watermark'''
        pi.wm_steps                  = wm_steps
        pi.key                       = key.clone()
        pi.update_step()
        latent_img_wm, _ = \
        pi.invert(prompts, init_latents_noise.clone(), is_forward=False)
        
        image_wm = pi.latent2image(latent_img_wm)
        for i, prompt in enumerate(prompts):
            file_name    = str(time.time())
            wm_path      = f'{save_wm_path}/{file_name}.png'
            if latent_img_wm.shape[0] == 1:
                cv2.imwrite(wm_path, cv2.cvtColor(image_wm, cv2.COLOR_RGB2BGR))
            else:
                cv2.imwrite(wm_path, cv2.cvtColor(image_wm[i], cv2.COLOR_RGB2BGR))

            record_dict[f'{file_name}.png'] = {
                'watermarked_img_path': wm_path,
                'key'                 : key[i].cpu().numpy(),
                'prompt'              : prompt,
                'time'                : public_times[i],
            }
            record.append({
                'watermarked_img_path': wm_path,
                'key'                 : key[i].cpu().numpy(),
                'prompt'              : prompt,
                'time'                : public_times[i],
            })

    np.save(f"{args.save_path}/watermark.npy", [f"{args.save_path}/watermark.npy", 
                                                record,
                                                record_dict])