code: def __init__(
      self,
      data_root,
      width,
      height,
      repeats,
      flip_p=0.5,
      placeholder_token="*",
      model=None,
      cond_model=None,
      device=None,
      template_file=None,
      include_cond=False,
      batch_size=1,
      gradient_step=1,
      shuffle_tags=False,
      tag_drop_out=0,
      latent_sampling_method="once",
      varsize=False,
      use_weight=False,
  ):
      re_word = (
          re.compile(shared.opts.dataset_filename_word_regex)
          if shared.opts.dataset_filename_word_regex
          else None
      )
  
      self.placeholder_token = placeholder_token
  
      self.flip = transforms.RandomHorizontalFlip(p=flip_p)
  
      self.dataset = []
  
      with open(template_file, "r") as file:
          lines = [x.strip() for x in file.readlines()]
  
      self.lines = lines
  
      assert data_root, "dataset directory not specified"
      assert os.path.isdir(data_root), "Dataset directory doesn't exist"
      assert os.listdir(data_root), "Dataset directory is empty"
  
      self.image_paths = [
          os.path.join(data_root, file_path)
          for file_path in os.listdir(data_root)
      ]
  
      self.shuffle_tags = shuffle_tags
      self.tag_drop_out = tag_drop_out
      groups = defaultdict(list)
  
      print("Preparing dataset...")
      for path in tqdm.tqdm(self.image_paths):
          alpha_channel = None
          if shared.state.interrupted:
              raise Exception("interrupted")
          try:
              image = Image.open(path)
              # Currently does not work for single color transparency
              # We would need to read image.info['transparency'] for that
              if use_weight and "A" in image.getbands():
                  alpha_channel = image.getchannel("A")
              image = image.convert("RGB")
              if not varsize:
                  image = image.resize(
                      (width, height), PIL.Image.BICUBIC
                  )
          except Exception:
              continue
  
          text_filename = f"{os.path.splitext(path)[0]}.txt"
          filename = os.path.basename(path)
  
          if os.path.exists(text_filename):
              with open(text_filename, "r", encoding="utf8") as file:
                  filename_text = file.read()
          else:
              filename_text = os.path.splitext(filename)[0]
              filename_text = re.sub(
                  re_numbers_at_start, "", filename_text
              )
              if re_word:
                  tokens = re_word.findall(filename_text)
                  filename_text = (
                      shared.opts.dataset_filename_join_string or ""
                  ).join(tokens)
  
          npimage = np.array(image).astype(np.uint8)
          npimage = (npimage / 127.5 - 1.0).astype(np.float32)
  
          torchdata = (
              torch.from_numpy(npimage)
              .permute(2, 0, 1)
              .to(device=device, dtype=torch.float32)
          )
          latent_sample = None
  
          with devices.autocast():
              latent_dist = model.encode_first_stage(
                  torchdata.unsqueeze(dim=0)
              )
  
          # Perform latent sampling, even for random sampling.
          # We need the sample dimensions for the weights
          if latent_sampling_method == "deterministic":
              if isinstance(latent_dist, DiagonalGaussianDistribution):
                  # Works only for DiagonalGaussianDistribution
                  latent_dist.std = 0
              else:
                  latent_sampling_method = "once"
          latent_sample = (
              model.get_first_stage_encoding(latent_dist)
              .squeeze()
              .to(devices.cpu)
          )
  
          if use_weight and alpha_channel is not None:
              channels, *latent_size = latent_sample.shape
              weight_img = alpha_channel.resize(latent_size)
              npweight = np.array(weight_img).astype(np.float32)
              # Repeat for every channel in the latent sample
              weight = torch.tensor([npweight] * channels).reshape(
                  [channels] + latent_size
              )
              # Normalize the weight to a minimum of 0 and a mean of 1, that way the loss will be comparable to default.
              weight -= weight.min()
              weight /= weight.mean()
          elif use_weight:
              # If an image does not have a alpha channel, add a ones weight map anyway so we can stack it later
              weight = torch.ones(latent_sample.shape)
          else:
              weight = None
  
          if latent_sampling_method == "random":
              entry = DatasetEntry(
                  filename=path,
                  filename_text=filename_text,
                  latent_dist=latent_dist,
                  weight=weight,
              )
          else:
              entry = DatasetEntry(
                  filename=path,
                  filename_text=filename_text,
                  latent_sample=latent_sample,
                  weight=weight,
              )
  
          if not (self.tag_drop_out != 0 or self.shuffle_tags):
              entry.cond_text = self.create_text(filename_text)
  
          if include_cond and not (
              self.tag_drop_out != 0 or self.shuffle_tags
          ):
              with devices.autocast():
                  entry.cond = (
                      cond_model([entry.cond_text])
                      .to(devices.cpu)
                      .squeeze(0)
                  )
          groups[image.size].append(len(self.dataset))
          self.dataset.append(entry)
          del torchdata
          del latent_dist
          del latent_sample
          del weight
  
      self.length = len(self.dataset)
      self.groups = list(groups.values())
      assert (
          self.length > 0
      ), "No images have been found in the dataset."
      self.batch_size = min(batch_size, self.length)
      self.gradient_step = min(
          gradient_step, self.length // self.batch_size
      )
      self.latent_sampling_method = latent_sampling_method
  
      if len(groups) > 1:
          print("Buckets:")
          for (w, h), ids in sorted(groups.items(), key=lambda x: x[0]):
              print(f"  {w}x{h}: {len(ids)}")
          print()
  
imports: import os
  import numpy
  import PIL
  import torch
  from PIL import Image
  from torch.utils.data import Dataset, DataLoader, Sampler
  from torchvision import transforms
  from collections import defaultdict
  from random import shuffle, choices
  import random
  import tqdm
  from modules import devices, shared
  import re
  from ldm.modules.distributions.distributions import DiagonalGaussianDistribution
label_desc: Use image node_id 1 to create a new PersonalizedBase object from inherited base classes: Dataset with example: obj = PersonalizedBase(data_root, width, height, repeats, flip_p, placeholder_token, model, cond_model, device, template_file, include_cond, batch_size, gradient_step, shuffle_tags, tag_drop_out, latent_sampling_method, varsize, use_weight)