Added Dreambooth training script

1 files changed, 16 insertions, 0 deletions

diff --git a/data/dreambooth/prompt.py b/data/dreambooth/prompt.py
new file mode 100644
index 0000000..34f510d
--- /dev/null
+++ b/data/dreambooth/prompt.py
@@ -0,0 +1,16 @@
+from torch.utils.data import Dataset
+
+
+class PromptDataset(Dataset):
+    def __init__(self, prompt, num_samples):
+        self.prompt = prompt
+        self.num_samples = num_samples
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, index):
+        example = {}
+        example["prompt"] = self.prompt
+        example["index"] = index
+        return example