adding train method dropdown

app.py

@@ -230,12 +230,13 @@ class Demo:
             self.iterations_input,
             self.lr_input,
             self.attributes_input,
-            self.is_person
+            self.is_person,
+            self.train_method_input
         ],
         outputs=[self.train_button,  self.train_status, self.download, self.model_dropdown]
         )
 
-    def train(self, target_concept,positive_prompt, negative_prompt, rank, iterations_input, lr_input, attributes_input, is_person, pbar = gr.Progress(track_tqdm=True)):
+    def train(self, target_concept,positive_prompt, negative_prompt, rank, iterations_input, lr_input, attributes_input, is_person, train_method_input, pbar = gr.Progress(track_tqdm=True)):
         iterations_input = min(int(iterations_input),1000)
         if attributes_input == '':
             attributes_input = None
@@ -257,13 +258,13 @@ class Demo:
             attributes = 'white, black, asian, hispanic, indian, male, female'
         
         self.training = True
-        train_xl(target=target_concept, positive=positive_prompt, negative=negative_prompt, lr=lr_input, iterations=iterations_input, config_file='trainscripts/textsliders/data/config-xl.yaml', rank=int(rank), device=self.device, attributes=attributes, save_name=save_name)
+        train_xl(target=target_concept, positive=positive_prompt, negative=negative_prompt, lr=lr_input, iterations=iterations_input, config_file='trainscripts/textsliders/data/config-xl.yaml', rank=int(rank), train_method=train_method_input, device=self.device, attributes=attributes, save_name=save_name)
         self.training = False
 
         torch.cuda.empty_cache()
-        model_map['Custom Slider'] = f'models/{save_name}'
+        model_map[save_name.replace('.pt','')] = f'models/{save_name}'
         
-        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom slider in the "Test" tab'), f'models/{save_name}', gr.update(choices=list(model_map.keys()), value='Custom Slider')]
+        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom slider in the "Test" tab'), f'models/{save_name}', gr.update(choices=list(model_map.keys()), value=save_name.replace('.pt',''))]
 
     
     def inference(self, prompt, seed, start_noise, scale, model_name, pbar = gr.Progress(track_tqdm=True)):

trainscripts/textsliders/data/config-xl.yaml

@@ -7,7 +7,7 @@ network:
   type: "c3lier" # or "c3lier" or "lierla"
   rank: 4
   alpha: 1.0
-  training_method: "xattn"
+  training_method: "noxattn"
 train:
   precision: "bfloat16"
   noise_scheduler: "ddim" # or "ddpm", "lms", "euler_a"

trainscripts/textsliders/demotrain.py

@@ -411,7 +411,7 @@ def train(
 #     train(config, prompts, device)
 
 
-def train_xl(target, positive, negative, lr, iterations, config_file, rank, device, attributes,save_name):
+def train_xl(target, positive, negative, lr, iterations, config_file, rank, train_method, device, attributes,save_name):
     
     config = config_util.load_config_from_yaml(config_file)
     randn = torch.randint(1, 10000000, (1,)).item()
@@ -427,6 +427,7 @@ def train_xl(target, positive, negative, lr, iterations, config_file, rank, devi
         attributes = []
     config.network.alpha = 1.0
     config.network.rank = int(rank)
+    config.network.training_method = train_method
     
 #     config.save.path += f'/{config.save.name}'