Fix

trainer.py

@@ -67,13 +67,12 @@ class Trainer:
         gradient_accumulation: int,
         fp16: bool,
         use_8bit_adam: bool,
-    ) -> tuple[dict, str]:
+    ) -> tuple[dict, list[pathlib.Path]]:
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
 
-        out_path = ''
         if self.is_running:
-            return gr.update(value=self.is_running_message), out_path
+            return gr.update(value=self.is_running_message), []
 
         if concept_images is None:
             raise gr.Error('You need to upload images.')
@@ -116,9 +115,7 @@ class Trainer:
 
         if res.returncode == 0:
             result_message = 'Training Completed!'
-            weight_path = self.output_dir / 'lora_weight.pt'
-            if weight_path.exists():
-                out_path = weight_path.as_posix()
         else:
             result_message = 'Training Failed!'
-        return gr.update(value=result_message), out_path
+        weight_paths = sorted(self.output_dir.glob('*.pt'))
+        return gr.update(value=result_message), weight_paths