Update

app.py

@@ -73,7 +73,6 @@ def create_training_demo(trainer: Trainer,
                 gr.Markdown('Training Data')
                 concept_images = gr.Files(label='Images for your concept')
                 concept_prompt = gr.Textbox(label='Concept Prompt',
-                                            value='sks',
                                             max_lines=1)
                 gr.Markdown('''
                     - Upload images of the style you are planning on training on.
@@ -84,8 +83,14 @@ def create_training_demo(trainer: Trainer,
                 num_training_steps = gr.Number(
                     label='Number of Training Steps', value=1000, precision=0)
                 learning_rate = gr.Number(label='Learning Rate', value=0.0001)
+                gradient_accumulation = gr.Number(
+                    label='Number of Gradient Accumulation',
+                    value=1,
+                    precision=0)
+                fp16 = gr.Checkbox(label='FP16', value=True)
+                use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=True)
                 gr.Markdown('''
-                    - It will take about 15-20 minutes to train for 1000 steps with a T4 GPU.
+                    - It will take about 8 minutes to train for 1000 steps with a T4 GPU.
                     - You may want to try a small number of steps first, like 1, to see if everything works fine in your environment.
                     ''')
 
@@ -108,6 +113,9 @@ def create_training_demo(trainer: Trainer,
                              concept_prompt,
                              num_training_steps,
                              learning_rate,
+                             gradient_accumulation,
+                             fp16,
+                             use_8bit_adam,
                          ],
                          outputs=[
                              training_status,
@@ -175,6 +183,10 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
                                                value=7)
 
                 run_button = gr.Button('Generate')
+
+                gr.Markdown('''
+                - The pretrained models are trained with the concept prompt "style of sks".
+                ''')
             with gr.Column():
                 result = gr.Image(label='Result')
 

trainer.py

@@ -54,15 +54,10 @@ class Trainer:
             out_path = self.instance_data_dir / f'{i:03d}.jpg'
             image.save(out_path, format='JPEG', quality=100)
 
-    def run(
-        self,
-        base_model: str,
-        resolution_s: str,
-        concept_images: list | None,
-        concept_prompt: str,
-        n_steps: int,
-        learning_rate: float,
-    ) -> tuple[dict, str]:
+    def run(self, base_model: str, resolution_s: str,
+            concept_images: list | None, concept_prompt: str, n_steps: int,
+            learning_rate: float, gradient_accumulation: int, fp16: bool,
+            use_8bit_adam: bool) -> tuple[dict, str]:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
 
@@ -80,24 +75,30 @@ class Trainer:
         self.cleanup_dirs()
         self.prepare_dataset(concept_images, resolution)
 
-        self.is_running = True
         command = f'''
         accelerate launch lora/train_lora_dreambooth.py \
           --pretrained_model_name_or_path={base_model}  \
           --instance_data_dir={self.instance_data_dir} \
           --output_dir={self.output_dir} \
-          --instance_prompt="style of {concept_prompt}" \
+          --instance_prompt="{concept_prompt}" \
           --resolution={resolution} \
           --train_batch_size=1 \
-          --gradient_accumulation_steps=1 \
+          --gradient_accumulation_steps={gradient_accumulation} \
           --learning_rate={learning_rate} \
           --lr_scheduler=constant \
           --lr_warmup_steps=0 \
           --max_train_steps={n_steps}
         '''
+        if fp16:
+            command += ' --mixed_precision fp16 '
+        if use_8bit_adam:
+            command += ' --use_8bit_adam'
+
         with open(self.output_dir / 'train.sh', 'w') as f:
             command_s = ' '.join(command.split())
             f.write(command_s)
+
+        self.is_running = True
         res = subprocess.run(shlex.split(command))
         self.is_running = False