Added class data support

app.py

@@ -77,6 +77,7 @@ def create_training_demo(trainer: Trainer, pipe: InferencePipeline) -> gr.Blocks
             with gr.Box():
                 gr.Markdown("Training Data")
                 concept_images = gr.Files(label="Images for your concept")
+                class_images = gr.Files(label="Class images")
                 concept_prompt = gr.Textbox(label="Concept Prompt", max_lines=1)
                 gr.Markdown(
                     """
@@ -202,6 +203,7 @@ def create_training_demo(trainer: Trainer, pipe: InferencePipeline) -> gr.Blocks
                 num_training_steps,
                 concept_images,
                 concept_prompt,
+                class_images,
                 learning_rate,
                 gradient_accumulation,
                 fp16,

trainer.py

@@ -31,6 +31,7 @@ class Trainer:
         self.is_running_message = "Another training is in progress."
 
         self.output_dir = pathlib.Path("results")
+        self.class_dir = self.output_dir / "class_data"
         self.instance_data_dir = self.output_dir / "training_data"
 
     def check_if_running(self) -> dict:
@@ -52,6 +53,13 @@ class Trainer:
             out_path = self.instance_data_dir / f"{i:03d}.jpg"
             image.save(out_path, format="JPEG", quality=100)
 
+    def copy_class_data(self, class_images: list) -> None:
+        self.class_dir.mkdir(parents=True)
+        for i, temp_path in enumerate(class_images):
+            image = PIL.Image.open(temp_path.name)
+            out_path = self.class_dir / f"{i:03d}.jpg"
+            image.save(out_path, format="JPEG", quality=100)
+
     def run(
         self,
         base_model: str,
@@ -59,6 +67,7 @@ class Trainer:
         n_steps: int,
         concept_images: list | None,
         concept_prompt: str,
+        class_images: list | None,
         learning_rate: float,
         gradient_accumulation: int,
         fp16: bool,
@@ -93,6 +102,7 @@ class Trainer:
 
         self.cleanup_dirs()
         self.prepare_dataset(concept_images, resolution)
+        self.copy_class_data(class_images)
 
         command = f"""
         accelerate launch train_dreambooth.py \
@@ -116,7 +126,7 @@ class Trainer:
             command += f""" --with_prior_preservation \
                 --prior_loss_weight={prior_loss_weight} \
                 --class_prompt="{class_prompt}" \
-                --class_data_dir={self.output_dir / 'class_data'}
+                --class_data_dir={self.class_dir}
                 """
 
         command += f""" --use_lora \