Spaces:
Runtime error
Runtime error
import os | |
import json | |
import random | |
from tqdm import tqdm | |
import numpy as np | |
from PIL import Image, ImageStat | |
import torch | |
from torch.utils.data import Dataset, DataLoader, IterableDataset, get_worker_info | |
from torchvision import transforms as T | |
### >>>>>>>> >>>>>>>> text related >>>>>>>> >>>>>>>> ### | |
class TokenizerWrapper(): | |
def __init__(self, tokenizer, is_train, proportion_empty_prompts, use_generic_prompts=False): | |
self.tokenizer = tokenizer | |
self.is_train = is_train | |
self.proportion_empty_prompts = proportion_empty_prompts | |
self.use_generic_prompts = use_generic_prompts | |
def __call__(self, prompts): | |
if isinstance(prompts, str): | |
prompts = [prompts] | |
captions = [] | |
for caption in prompts: | |
if random.random() < self.proportion_empty_prompts: | |
captions.append("") | |
else: | |
if self.use_generic_prompts: | |
captions.append("best quality, high quality") | |
elif isinstance(caption, str): | |
captions.append(caption) | |
elif isinstance(caption, (list, np.ndarray)): | |
# take a random caption if there are multiple | |
captions.append(random.choice(caption) if self.is_train else caption[0]) | |
else: | |
raise ValueError( | |
f"Caption column should contain either strings or lists of strings." | |
) | |
inputs = self.tokenizer( | |
captions, max_length=self.tokenizer.model_max_length, padding="max_length", | |
truncation=True, return_tensors="pt" | |
) | |
return inputs.input_ids | |
### >>>>>>>> >>>>>>>> image related >>>>>>>> >>>>>>>> ### | |
MONOCHROMATIC_MAX_VARIANCE = 0.3 | |
def is_monochromatic_image(pil_img): | |
v = ImageStat.Stat(pil_img.convert('RGB')).var | |
return sum(v)<MONOCHROMATIC_MAX_VARIANCE | |
def isnumeric(text): | |
return (''.join(filter(str.isalnum, text))).isnumeric() | |
class TextPromptDataset(IterableDataset): | |
''' | |
The dataset for (text embedding, noise, generated latent) triplets. | |
''' | |
def __init__(self, | |
data_root, | |
tokenizer = None, | |
transform = None, | |
rank = 0, | |
world_size = 1, | |
shuffle = True, | |
): | |
self.tokenizer = tokenizer | |
self.transform = transform | |
self.img_root = os.path.join(data_root, 'JPEGImages') | |
self.data_list = [] | |
print("#### Loading filename list...") | |
json_root = os.path.join(data_root, 'list') | |
json_list = [p for p in os.listdir(json_root) if p.startswith("shard") and p.endswith('.json')] | |
# duplicate several shards to make sure each process has the same number of shards | |
assert len(json_list) > world_size | |
duplicate = world_size - len(json_list)%world_size if len(json_list)%world_size>0 else 0 | |
json_list = json_list + json_list[:duplicate] | |
json_list = json_list[rank::world_size] | |
for json_file in tqdm(json_list): | |
shard_name = os.path.basename(json_file).split('.')[0] | |
with open(os.path.join(json_root, json_file)) as f: | |
key_text_pairs = json.load(f) | |
for pair in key_text_pairs: | |
self.data_list.append( [shard_name] + pair ) | |
print("#### All filename loaded...") | |
self.shuffle = shuffle | |
def __len__(self): | |
return len(self.data_list) | |
def __iter__(self): | |
worker_info = get_worker_info() | |
if worker_info is None: # single-process data loading, return the full iterator | |
data_list = self.data_list | |
else: | |
len_data = len(self.data_list) - len(self.data_list) % worker_info.num_workers | |
data_list = self.data_list[:len_data][worker_info.id :: worker_info.num_workers] | |
# print(worker_info.num_workers, worker_info.id, len(data_list)/len(self.data_list)) | |
if self.shuffle: | |
random.shuffle(data_list) | |
while True: | |
for idx in range(len(data_list)): | |
# try: | |
shard_name = data_list[idx][0] | |
data = {} | |
img_file = data_list[idx][1] | |
img = Image.open(os.path.join(self.img_root, shard_name, img_file+'.jpg')).convert("RGB") | |
if is_monochromatic_image(img): | |
continue | |
if self.transform is not None: | |
img = self.transform(img) | |
data['pixel_values'] = img | |
text = data_list[idx][2] | |
if self.tokenizer is not None: | |
if isinstance(self.tokenizer, list): | |
assert len(self.tokenizer)==2 | |
data['input_ids'] = self.tokenizer[0](text)[0] | |
data['input_ids_2'] = self.tokenizer[1](text)[0] | |
else: | |
data['input_ids'] = self.tokenizer(text)[0] | |
else: | |
data['input_ids'] = text | |
yield data | |
# except Exception as e: | |
# raise(e) | |
def collate_fn(self, examples): | |
pixel_values = torch.stack([example["pixel_values"] for example in examples]) | |
pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float() | |
if self.tokenizer is not None: | |
if isinstance(self.tokenizer, list): | |
assert len(self.tokenizer)==2 | |
input_ids = torch.stack([example["input_ids"] for example in examples]) | |
input_ids_2 = torch.stack([example["input_ids_2"] for example in examples]) | |
return {"pixel_values": pixel_values, "input_ids": input_ids, "input_ids_2": input_ids_2,} | |
else: | |
input_ids = torch.stack([example["input_ids"] for example in examples]) | |
return {"pixel_values": pixel_values, "input_ids": input_ids,} | |
else: | |
input_ids = [example["input_ids"] for example in examples] | |
return {"pixel_values": pixel_values, "input_ids": input_ids,} | |
def make_train_dataset( | |
train_data_path, | |
size = 512, | |
tokenizer=None, | |
cfg_drop_ratio=0, | |
rank=0, | |
world_size=1, | |
shuffle=True, | |
): | |
_image_transform = T.Compose([ | |
T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img), | |
T.Resize(size), | |
T.CenterCrop((size,size)), | |
T.ToTensor(), | |
T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) | |
]) | |
if tokenizer is not None: | |
if isinstance(tokenizer, list): | |
assert len(tokenizer)==2 | |
tokenizer_1 = TokenizerWrapper( | |
tokenizer[0], | |
is_train=True, | |
proportion_empty_prompts=cfg_drop_ratio, | |
use_generic_prompts=False, | |
) | |
tokenizer_2 = TokenizerWrapper( | |
tokenizer[1], | |
is_train=True, | |
proportion_empty_prompts=cfg_drop_ratio, | |
use_generic_prompts=False, | |
) | |
tokenizer = [tokenizer_1, tokenizer_2] | |
else: | |
tokenizer = TokenizerWrapper( | |
tokenizer, | |
is_train=True, | |
proportion_empty_prompts=cfg_drop_ratio, | |
use_generic_prompts=False, | |
) | |
train_dataset = TextPromptDataset( | |
data_root=train_data_path, | |
transform=_image_transform, | |
rank=rank, | |
world_size=world_size, | |
tokenizer=tokenizer, | |
shuffle=shuffle, | |
) | |
return train_dataset | |
### >>>>>>>> >>>>>>>> Test >>>>>>>> >>>>>>>> ### | |
if __name__ == "__main__": | |
from transformers import CLIPTextModel, CLIPTokenizer | |
tokenizer = CLIPTokenizer.from_pretrained( | |
"/mnt/bn/ic-research-aigc-editing/fast-diffusion-models/assets/public_models/StableDiffusion/stable-diffusion-v1-5", | |
subfolder="tokenizer" | |
) | |
train_dataset = make_train_dataset(tokenizer=tokenizer, rank=0, world_size=10) | |
loader = torch.utils.data.DataLoader( | |
train_dataset, batch_size=64, num_workers=0, | |
collate_fn=train_dataset.collect_fn if hasattr(train_dataset, 'collect_fn') else None, | |
) | |
for batch in loader: | |
pixel_values = batch["pixel_values"] | |
prompt_ids = batch['input_ids'] | |
from einops import rearrange | |
pixel_values = rearrange(pixel_values, 'b c h w -> b h w c') | |
for i in range(pixel_values.shape[0]): | |
import pdb; pdb.set_trace() | |
Image.fromarray(((pixel_values[i] + 1 )/2 * 255 ).numpy().astype(np.uint8)).save('tmp.png') | |
input_id = prompt_ids[i] | |
text = tokenizer.decode(input_id).split('<|startoftext|>')[-1].split('<|endoftext|>')[0] | |
print(text) | |
pass |