Update

trainer.py

@@ -32,8 +32,8 @@ class Trainer:
         self.is_running = False
         self.is_running_message = 'Another training is in progress.'
 
-        self.instance_data_dir = pathlib.Path('training_data')
         self.output_dir = pathlib.Path('results')
+        self.instance_data_dir = self.output_dir / 'training_data'
 
     def check_if_running(self) -> dict:
         if self.is_running:
@@ -42,11 +42,10 @@ class Trainer:
             return gr.update(value='No training is running.')
 
     def cleanup_dirs(self) -> None:
-        shutil.rmtree(self.instance_data_dir, ignore_errors=True)
         shutil.rmtree(self.output_dir, ignore_errors=True)
 
     def prepare_dataset(self, concept_images: list, resolution: int) -> None:
-        self.instance_data_dir.mkdir()
+        self.instance_data_dir.mkdir(parents=True)
         for i, temp_path in enumerate(concept_images):
             image = PIL.Image.open(temp_path.name)
             image = pad_image(image)
@@ -96,6 +95,9 @@ class Trainer:
           --lr_warmup_steps=0 \
           --max_train_steps={n_steps}
         '''
+        with open(self.output_dir / 'train.sh', 'w') as f:
+            command_s = ' '.join(command.split())
+            f.write(command_s)
         res = subprocess.run(shlex.split(command))
         self.is_running = False