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