concept-ablation-diffusers/utils.py

import os
import random
import shutil
from io import BytesIO
from pathlib import Path

import numpy as np
import openai
import regex as re
import requests
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
from tqdm.auto import tqdm

from diffusers import DPMSolverMultistepScheduler

normalize = transforms.Normalize(
 mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225],
)
small_288 = transforms.Compose([
 transforms.Resize(288),
 transforms.ToTensor(),
 normalize,
])

def collate_fn(examples, with_prior_preservation):
 input_ids = [example["instance_prompt_ids"] for example in examples]
 input_anchor_ids = [example["instance_anchor_prompt_ids"]
 for example in examples]
 pixel_values = [example["instance_images"] for example in examples]
 mask = [example["mask"] for example in examples]
 # Concat class and instance examples for prior preservation.
 # We do this to avoid doing two forward passes.
 if with_prior_preservation:
 input_ids += [example["class_prompt_ids"] for example in examples]
 pixel_values += [example["class_images"] for example in examples]
 mask += [example["class_mask"] for example in examples]

 input_ids = torch.cat(input_ids, dim=0)
 input_anchor_ids = torch.cat(input_anchor_ids, dim=0)
 pixel_values = torch.stack(pixel_values)
 mask = torch.stack(mask)
 pixel_values = pixel_values.to(
 memory_format=torch.contiguous_format).float()
 mask = mask.to(memory_format=torch.contiguous_format).float()

 batch = {
 "input_ids": input_ids,
 "input_anchor_ids": input_anchor_ids,
 "pixel_values": pixel_values,
 "mask": mask.unsqueeze(1)
 }
 return batch


class PromptDataset(Dataset):
 "A simple dataset to prepare the prompts to generate class images on multiple GPUs."

 def __init__(self, prompt, num_samples):
 self.prompt = prompt
 self.num_samples = num_samples

 def __len__(self):
 return self.num_samples

 def __getitem__(self, index):
 example = {}
 example["prompt"] = self.prompt[index % len(self.prompt)]
 example["index"] = index
 return example


class CustomDiffusionDataset(Dataset):
 """
 A dataset to prepare the instance and class images with the prompts for fine-tuning the model.
 It pre-processes the images and the tokenizes prompts.
 """

 def __init__(
 self,
 concepts_list,
 concept_type,
 tokenizer,
 size=512,
 center_crop=False,
 with_prior_preservation=False,
 num_class_images=200,
 hflip=False,
 aug=True,
 ):
 self.size = size
 self.center_crop = center_crop
 self.tokenizer = tokenizer
 self.interpolation = Image.BILINEAR
 self.aug = aug
 self.concept_type = concept_type

 self.instance_images_path = []
 self.class_images_path = []
 self.with_prior_preservation = with_prior_preservation
 for concept in concepts_list:
 with open(concept["instance_data_dir"], "r") as f:
 inst_images_path = f.read().splitlines()
 with open(concept["instance_prompt"], "r") as f:
 inst_prompt = f.read().splitlines()
 inst_img_path = [(x, y, concept['caption_target'])
 for (x, y) in zip(inst_images_path, inst_prompt)]
 self.instance_images_path.extend(inst_img_path)

 if with_prior_preservation:
 class_data_root = Path(concept["class_data_dir"])
 if os.path.isdir(class_data_root):
 class_images_path = list(class_data_root.iterdir())
 class_prompt = [concept["class_prompt"]
 for _ in range(len(class_images_path))]
 else:
 with open(class_data_root, "r") as f:
 class_images_path = f.read().splitlines()
 with open(concept["class_prompt"], "r") as f:
 class_prompt = f.read().splitlines()

 class_img_path = [(x, y) for (x, y) in zip(
 class_images_path, class_prompt)]
 self.class_images_path.extend(
 class_img_path[:num_class_images])

 random.shuffle(self.instance_images_path)
 self.num_instance_images = len(self.instance_images_path)
 self.num_class_images = len(self.class_images_path)
 self._length = max(self.num_class_images, self.num_instance_images)
 self.flip = transforms.RandomHorizontalFlip(0.5 * hflip)

 self.image_transforms = transforms.Compose(
 [
 self.flip,
 transforms.Resize(
 size, interpolation=transforms.InterpolationMode.BILINEAR),
 transforms.CenterCrop(
 size) if center_crop else transforms.RandomCrop(size),
 transforms.ToTensor(),
 transforms.Normalize([0.5], [0.5]),
 ]
 )

 def __len__(self):
 return self._length

 def preprocess(self, image, scale, resample):
 outer, inner = self.size, scale
 if scale > self.size:
 outer, inner = scale, self.size
 top, left = np.random.randint(
 0, outer - inner + 1), np.random.randint(0, outer - inner + 1)
 image = image.resize((scale, scale), resample=resample)
 image = np.array(image).astype(np.uint8)
 image = (image / 127.5 - 1.0).astype(np.float32)
 instance_image = np.zeros((self.size, self.size, 3), dtype=np.float32)
 mask = np.zeros((self.size // 8, self.size // 8))
 if scale > self.size:
 instance_image = image[top: top + inner, left: left + inner, :]
 mask = np.ones((self.size // 8, self.size // 8))
 else:
 instance_image[top: top + inner, left: left + inner, :] = image
 mask[top // 8 + 1: (top + scale) // 8 - 1, left //
 8 + 1: (left + scale) // 8 - 1] = 1.
 return instance_image, mask

 def __getprompt__(self, instance_prompt, instance_target):
 if self.concept_type == 'style':
 r = np.random.choice([0, 1, 2])
 instance_prompt = f'{instance_prompt}, in the style of {instance_target}' if r == 0 else f'in {instance_target}\'s style, {instance_prompt}' if r == 1 else f'in {instance_target}\'s style, {instance_prompt}'
 elif self.concept_type == 'object':
 anchor, target = instance_target.split('+')
 instance_prompt = instance_prompt.replace(anchor, target)
 elif self.concept_type == 'memorization':
 instance_prompt = instance_target.split('+')[1]
 return instance_prompt

 def __getitem__(self, index):
 example = {}
 instance_image, instance_prompt, instance_target = self.instance_images_path[
 index % self.num_instance_images]
 instance_image = Image.open(instance_image)
 if not instance_image.mode == "RGB":
 instance_image = instance_image.convert("RGB")
 instance_image = self.flip(instance_image)
 # modify instance prompt according to the concept_type to include target concept
 # multiple style/object fine-tuning
 if ';' in instance_target:
 instance_target = instance_target.split(';')
 instance_target = instance_target[index % len(instance_target)]

 instance_anchor_prompt = instance_prompt
 instance_prompt = self.__getprompt__(instance_prompt, instance_target)
 # apply resize augmentation and create a valid image region mask
 random_scale = self.size
 if self.aug:
 random_scale = np.random.randint(self.size // 3, self.size + 1) if np.random.uniform(
 ) < 0.66 else np.random.randint(int(1.2 * self.size), int(1.4 * self.size))
 instance_image, mask = self.preprocess(
 instance_image, random_scale, self.interpolation)

 if random_scale < 0.6 * self.size:
 instance_prompt = np.random.choice(
 ["a far away ", "very small "]) + instance_prompt
 elif random_scale > self.size:
 instance_prompt = np.random.choice(
 ["zoomed in ", "close up "]) + instance_prompt

 example["instance_images"] = torch.from_numpy(
 instance_image).permute(2, 0, 1)
 example["mask"] = torch.from_numpy(mask)

 example["instance_prompt_ids"] = self.tokenizer(
 instance_prompt,
 truncation=True,
 padding="max_length",
 max_length=self.tokenizer.model_max_length,
 return_tensors="pt",
 ).input_ids
 example["instance_anchor_prompt_ids"] = self.tokenizer(
 instance_anchor_prompt,
 truncation=True,
 padding="max_length",
 max_length=self.tokenizer.model_max_length,
 return_tensors="pt",
 ).input_ids

 if self.with_prior_preservation:
 class_image, class_prompt = self.class_images_path[index %
 self.num_class_images]
 class_image = Image.open(class_image)
 if not class_image.mode == "RGB":
 class_image = class_image.convert("RGB")
 example["class_images"] = self.image_transforms(class_image)
 example["class_mask"] = torch.ones_like(example["mask"])
 example["class_prompt_ids"] = self.tokenizer(
 class_prompt,
 truncation=True,
 padding="max_length",
 max_length=self.tokenizer.model_max_length,
 return_tensors="pt",
 ).input_ids

 return example


def isimage(path):
 if 'png' in path.lower() or 'jpg' in path.lower() or 'jpeg' in path.lower():
 return True


def filter(folder, impath, outpath=None, unfiltered_path=None, threshold=0.15,
 image_threshold=0.5, anchor_size=10, target_size=3, return_score=False):
 model = torch.jit.load(
 "./assets/sscd_imagenet_mixup.torchscript.pt")
 if isinstance(folder, list):
 image_paths = folder
 image_captions = ["None" for _ in range(len(image_paths))]
 elif Path(folder / 'images.txt').exists():
 with open(f'{folder}/images.txt', "r") as f:
 image_paths = f.read().splitlines()
 with open(f'{folder}/caption.txt', "r") as f:
 image_captions = f.read().splitlines()
 else:
 image_paths = [os.path.join(str(folder), file_path)
 for file_path in os.listdir(folder) if isimage(file_path)]
 image_captions = ["None" for _ in range(len(image_paths))]

 batch = small_288(Image.open(impath).convert('RGB')).unsqueeze(0)
 embedding_target = model(batch)[0, :]

 filtered_paths = []
 filtered_captions = []
 unfiltered_paths = []
 unfiltered_captions = []
 count_dict = {}
 for im, c in zip(image_paths, image_captions):
 if c not in count_dict:
 count_dict[c] = 0
 if isinstance(folder, list):
 batch = small_288(im).unsqueeze(0)
 else:
 batch = small_288(Image.open(im).convert('RGB')).unsqueeze(0)
 embedding = model(batch)[0, :]

 diff_sscd = (embedding * embedding_target).sum()

 if diff_sscd <= image_threshold:
 filtered_paths.append(im)
 filtered_captions.append(c)
 count_dict[c] += 1
 else:
 unfiltered_paths.append(im)
 unfiltered_captions.append(c)

 # only return score
 if return_score:
 score = len(unfiltered_paths) / \
 (len(unfiltered_paths)+len(filtered_paths))
 return score

 os.makedirs(outpath, exist_ok=True)
 os.makedirs(f'{outpath}/samples', exist_ok=True)
 with open(f'{outpath}/caption.txt', 'w') as f:
 for each in filtered_captions:
 f.write(each.strip() + '\n')

 with open(f'{outpath}/images.txt', 'w') as f:
 for each in filtered_paths:
 f.write(each.strip() + '\n')
 imbase = Path(each).name
 shutil.copy(each, f'{outpath}/samples/{imbase}')

 print('++++++++++++++++++++++++++++++++++++++++++++++++')
 print('+ Filter Summary +')
 print(f'+ Remained images: {len(filtered_paths)}')
 print(f'+ Filtered images: {len(unfiltered_paths)}')
 print('++++++++++++++++++++++++++++++++++++++++++++++++')

 sorted_list = sorted(list(count_dict.items()),
 key=lambda x: x[1], reverse=True)
 anchor_prompts = [c[0] for c in sorted_list[:anchor_size]]
 target_prompts = [c[0] for c in sorted_list[-target_size:]]
 return anchor_prompts, target_prompts, len(filtered_paths)


def getanchorprompts(pipeline, accelerator, class_prompt, concept_type, class_images_dir, api_key, num_class_images=200, mem_impath=None):
 openai.api_key = api_key
 class_prompt_collection = []
 caption_target = []
 if concept_type == 'object':
 messages = [{"role": "system", "content": "You can describe any image via text and provide captions for wide variety of images that is possible to generate."}]
 messages = [{"role": "user", "content": f"Generate {num_class_images} captions for images containing a {class_prompt}. The caption should also contain the word \"{class_prompt}\" "}]
 while True:
 completion = openai.ChatCompletion.create(
 model="gpt-3.5-turbo",
 messages=messages
 )
 class_prompt_collection += [x for x in completion.choices[0].message.content.lower(
 ).split('\n') if class_prompt in x]
 messages.append(
 {"role": "assistant", "content": completion.choices[0].message.content})
 messages.append(
 {"role": "user", "content": f"Generate {num_class_images-len(class_prompt_collection)} more captions"})
 if len(class_prompt_collection) >= num_class_images:
 break
 class_prompt_collection = clean_prompt(class_prompt_collection)[
 :num_class_images]

 elif concept_type == 'memorization':
 pipeline.scheduler = DPMSolverMultistepScheduler.from_config(
 pipeline.scheduler.config)
 num_prompts_firstpass = 5
 num_prompts_secondpass = 2
 threshold = 0.3
 # Generate num_prompts_firstpass paraphrases which generate different content at least 1-threshold % of the times.
 os.makedirs(class_images_dir / 'temp/', exist_ok=True)
 class_prompt_collection_counter = []
 caption_target = []
 prev_captions = []
 messages = [{"role": "user", "content": f"Generate {4*num_prompts_firstpass} different paraphrase of the caption: {class_prompt}. Preserve the meaning when paraphrasing."}]
 while True:
 completion = openai.ChatCompletion.create(
 model="gpt-3.5-turbo",
 messages=messages
 )
 # print(completion.choices[0].message.content.lower().split('\n'))
 class_prompt_collection_ = [x.strip(
 ) for x in completion.choices[0].message.content.lower().split('\n') if x.strip() != '']
 class_prompt_collection_ = clean_prompt(class_prompt_collection_)
 # print(class_prompt_collection_)
 for prompt in tqdm(
 class_prompt_collection_, desc="Generating anchor and target prompts ", disable=not accelerator.is_local_main_process
 ):
 print(f'Prompt: {prompt}')
 images = pipeline([prompt]*10, num_inference_steps=25,).images

 score = filter(images, mem_impath, return_score=True)
 print(f'Memorization rate: {score}')
 if score <= threshold and prompt not in class_prompt_collection and len(class_prompt_collection) < num_prompts_firstpass:
 class_prompt_collection += [prompt]
 class_prompt_collection_counter += [score]
 elif score >= 0.6 and prompt not in caption_target and len(caption_target) < 2:
 caption_target += [prompt]
 if len(class_prompt_collection) >= num_prompts_firstpass and len(caption_target) >= 2:
 break

 if len(class_prompt_collection) >= num_prompts_firstpass:
 break
 # print("prompts till now", class_prompt_collection, caption_target)
 # print("prompts till now", len(
 # class_prompt_collection), len(caption_target))
 prev_captions += class_prompt_collection_
 prev_captions_ = ','.join(prev_captions[-40:])

 messages = [
 {"role": "user", "content": f"Generate {4*(num_prompts_firstpass- len(class_prompt_collection))} different paraphrase of the caption: {class_prompt}. Preserve the meaning the most when paraphrasing. Also make sure that the new captions are different from the following captions: {prev_captions_[:4000]}"}]

 # Generate more paraphrases using the captions we retrieved above.
 for prompt in class_prompt_collection[:num_prompts_firstpass]:
 completion = openai.ChatCompletion.create(
 model="gpt-3.5-turbo",
 messages=[
 {"role": "user", "content": f"Generate {num_prompts_secondpass} different paraphrases of: {prompt}. "}]

 )
 class_prompt_collection += clean_prompt(
 [x.strip() for x in completion.choices[0].message.content.lower().split('\n') if x.strip() != ''])

 for prompt in tqdm(class_prompt_collection[num_prompts_firstpass:], desc="Memorization rate for final prompts"):
 images = pipeline([prompt]*10, num_inference_steps=25,).images

 class_prompt_collection_counter += [
 filter(images, mem_impath, return_score=True)]

 # select least ten and most memorized text prompts to be selected as anchor and target prompts.
 class_prompt_collection = sorted(
 zip(class_prompt_collection, class_prompt_collection_counter), key=lambda x: x[1])
 caption_target += [x for (x, y) in class_prompt_collection if y >= 0.6]
 class_prompt_collection = [
 x for (x, y) in class_prompt_collection if y <= threshold][:10]
 print("Anchor prompts:", class_prompt_collection)
 print("Target prompts:", caption_target)
 return class_prompt_collection, ';*+'.join(caption_target)


def clean_prompt(class_prompt_collection):
 class_prompt_collection = [re.sub(
 r"[0-9]+", lambda num: '' * len(num.group(0)), prompt) for prompt in class_prompt_collection]
 class_prompt_collection = [re.sub(
 r"^\.+", lambda dots: '' * len(dots.group(0)), prompt) for prompt in class_prompt_collection]
 class_prompt_collection = [x.strip() for x in class_prompt_collection]
 class_prompt_collection = [x.replace('"', '') for x in class_prompt_collection]
 return class_prompt_collection


def safe_dir(dir):
 if not dir.exists():
 dir.mkdir()
 return dir