from datasets import load_dataset
import torch
import config
from model import load_image

from diffusers.utils.loading_utils import load_image as dl_to_pil

def my_collate(batch):
    try:
      img = [item['image'] for item in batch]
      img = torch.cat([load_image(pil_image=i, image_file=None) for i in img])
      text = ['''<|user|><img><IMG_CONTEXT></img><|end|><|assistant|>'''+item['prompt'] for item in batch]
    except Exception as e:
      print(e)
      return None
    return {'image':img, 'text':text}

ds = load_dataset("stylebreeder/stylebreeder", split='2M_sample', streaming=True).shuffle(seed=7, buffer_size=1)
dataloader = torch.utils.data.DataLoader(ds, num_workers=32, collate_fn=my_collate, batch_size=config.batch_size)