Update main.py

main.py

@@ -10,6 +10,7 @@ from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionMod
 from transformers import CLIPTextModel, CLIPTokenizer
 from accelerate import Accelerator
 from tqdm.auto import tqdm
+from PIL import Image
 
 class CustomDataset(Dataset):
     def __init__(self, data_dir, prompt, tokenizer, size=512, center_crop=False):