Support training text encoder

app.py

@@ -83,6 +83,10 @@ 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)
+                train_text_encoder = gr.Checkbox(label='Train Text Encoder',
+                                                 value=False)
+                learning_rate_text = gr.Number(
+                    label='Learning Rate for Text Encoder', value=0.00005)
                 gradient_accumulation = gr.Number(
                     label='Number of Gradient Accumulation',
                     value=1,
@@ -113,6 +117,8 @@ def create_training_demo(trainer: Trainer,
                              concept_prompt,
                              num_training_steps,
                              learning_rate,
+                             train_text_encoder,
+                             learning_rate_text,
                              gradient_accumulation,
                              fp16,
                              use_8bit_adam,
@@ -136,6 +142,7 @@ def create_training_demo(trainer: Trainer,
 def find_weight_files() -> list[str]:
     curr_dir = pathlib.Path(__file__).parent
     paths = sorted(curr_dir.rglob('*.pt'))
+    paths = [path for path in paths if not path.stem.endswith('.text_encoder')]
     return [path.relative_to(curr_dir).as_posix() for path in paths]
 
 
@@ -165,6 +172,11 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
                                   maximum=2,
                                   step=0.05,
                                   value=1)
+                alpha_for_text = gr.Slider(label='Alpha for Text Encoder',
+                                           minimum=0,
+                                           maximum=2,
+                                           step=0.05,
+                                           value=1)
                 seed = gr.Slider(label='Seed',
                                  minimum=0,
                                  maximum=100000,
@@ -185,7 +197,8 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
                 run_button = gr.Button('Generate')
 
                 gr.Markdown('''
-                - The pretrained models are trained with the concept prompt "style of sks".
+                - The pretrained models for "disney", "illust" and "pop" are trained with the concept prompt "style of sks".
+                - The pretrained model for "kiriko" is trained with the concept prompt "game character bnha". For this model, the text encoder is also trained.
                 ''')
             with gr.Column():
                 result = gr.Image(label='Result')
@@ -199,6 +212,7 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
                           lora_weight_name,
                           prompt,
                           alpha,
+                          alpha_for_text,
                           seed,
                           num_steps,
                           guidance_scale,
@@ -211,6 +225,7 @@ def create_inference_demo(pipe: InferencePipeline) -> gr.Blocks:
                              lora_weight_name,
                              prompt,
                              alpha,
+                             alpha_for_text,
                              seed,
                              num_steps,
                              guidance_scale,

inference.py

@@ -32,6 +32,14 @@ class InferencePipeline:
         curr_dir = pathlib.Path(__file__).parent
         return curr_dir / name
 
+    @staticmethod
+    def get_lora_text_encoder_weight_path(path: pathlib.Path) -> str:
+        parent_dir = path.parent
+        stem = path.stem
+        text_encoder_filename = f'{stem}.text_encoder.pt'
+        path = parent_dir / text_encoder_filename
+        return path.as_posix() if path.exists() else ''
+
     def load_pipe(self, model_id: str, lora_filename: str) -> None:
         weight_path = self.get_lora_weight_path(lora_filename)
         if weight_path == self.weight_path:
@@ -47,6 +55,16 @@ class InferencePipeline:
             pipe = pipe.to(self.device)
 
         monkeypatch_lora(pipe.unet, lora_weight)
+
+        lora_text_encoder_weight_path = self.get_lora_text_encoder_weight_path(
+            weight_path)
+        if lora_text_encoder_weight_path:
+            lora_text_encoder_weight = torch.load(
+                lora_text_encoder_weight_path, map_location=self.device)
+            monkeypatch_lora(pipe.text_encoder,
+                             lora_text_encoder_weight,
+                             target_replace_module=['CLIPAttention'])
+
         self.pipe = pipe
 
     def run(
@@ -55,6 +73,7 @@ class InferencePipeline:
         lora_weight_name: str,
         prompt: str,
         alpha: float,
+        alpha_for_text: float,
         seed: int,
         n_steps: int,
         guidance_scale: float,
@@ -66,6 +85,7 @@ class InferencePipeline:
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         tune_lora_scale(self.pipe.unet, alpha)  # type: ignore
+        tune_lora_scale(self.pipe.text_encoder, alpha_for_text)  # type: ignore
         out = self.pipe(prompt,
                         num_inference_steps=n_steps,
                         guidance_scale=guidance_scale,

lora

@@ -1 +1 @@
-Subproject commit ba349e56e23e92e3b128c7c67ae58d3067540daa
+Subproject commit 26787a09bff4ebcb08f0ad4e848b67bce4389a7a

trainer.py

@@ -54,10 +54,20 @@ 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, gradient_accumulation: int, fp16: bool,
-            use_8bit_adam: bool) -> 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,
+        train_text_encoder: bool,
+        learning_rate_text: 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.')
 
@@ -93,6 +103,8 @@ class Trainer:
             command += ' --mixed_precision fp16 '
         if use_8bit_adam:
             command += ' --use_8bit_adam'
+        if train_text_encoder:
+            command += f' --train_text_encoder --learning_rate_text={learning_rate_text} --color_jitter'
 
         with open(self.output_dir / 'train.sh', 'w') as f:
             command_s = ' '.join(command.split())