Update cog_sdxl_dataset_and_utils.py

cog_sdxl_dataset_and_utils.py

@@ -33,6 +33,11 @@ def prepare_mask(mask: PIL.Image.Image, width: int = 512, height: int = 512) ->
     return torch.from_numpy(np.expand_dims(arr, 0)).unsqueeze(0)
 
 
+class TokenEmbeddingsHandler:
+    def __init__(self, text_encoders, tokenizers):
+        self.text_encoders = text_encoders
+        self.tokenizers = tokenizers
+
 class PreprocessedDataset(Dataset):
     def __init__(
         self,
@@ -175,3 +180,4 @@ def load_models(pretrained_model_name_or_path, revision, device, weight_dtype):
     unet.to(device, dtype=weight_dtype)
 
     return tokenizer_1, tokenizer_2, noise_scheduler, text_encoder_1, text_encoder_2, vae, unet
+