Refactor

app_upload.py

@@ -13,9 +13,14 @@ from utils import find_exp_dirs
 
 
 class LoRAModelUploader(Uploader):
-    def upload_lora_model(self, folder_path: str, repo_name: str,
-                          upload_to: str, private: bool,
-                          delete_existing_repo: bool) -> str:
+    def upload_lora_model(
+        self,
+        folder_path: str,
+        repo_name: str,
+        upload_to: str,
+        private: bool,
+        delete_existing_repo: bool,
+    ) -> str:
         if not repo_name:
             repo_name = pathlib.Path(folder_path).name
         repo_name = slugify.slugify(repo_name)

trainer.py

@@ -13,7 +13,7 @@ import slugify
 import torch
 from huggingface_hub import HfApi
 
-from constants import UploadTarget
+from app_upload import LoRAModelUploader
 
 
 def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
@@ -32,8 +32,8 @@ def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
 
 class Trainer:
     def __init__(self, hf_token: str | None = None):
-        self.hf_token = hf_token
         self.api = HfApi(token=hf_token)
+        self.model_uploader = LoRAModelUploader(hf_token)
 
     def prepare_dataset(self, instance_images: list, resolution: int,
                         instance_data_dir: pathlib.Path) -> None:
@@ -91,8 +91,7 @@ class Trainer:
         output_dir = repo_dir / 'experiments' / output_model_name
         if overwrite_existing_model or upload_to_hub:
             shutil.rmtree(output_dir, ignore_errors=True)
-        if not upload_to_hub:
-            output_dir.mkdir(parents=True)
+        output_dir.mkdir(parents=True)
 
         instance_data_dir = repo_dir / 'training_data' / output_model_name
         self.prepare_dataset(instance_images, resolution, instance_data_dir)
@@ -121,16 +120,23 @@ class Trainer:
             command += ' --use_8bit_adam'
         if use_wandb:
             command += ' --report_to wandb'
-        if upload_to_hub:
-            command += f' --push_to_hub --hub_token {self.hf_token}'
-            if use_private_repo:
-                command += ' --private_repo'
-            if delete_existing_repo:
-                command += ' --delete_existing_repo'
-            if upload_to == UploadTarget.LORA_LIBRARY.value:
-                command += ' --upload_to_lora_library'
 
+        with open(output_dir / 'train.sh', 'w') as f:
+            command_s = ' '.join(command.split())
+            f.write(command_s)
         subprocess.run(shlex.split(command))
+        message = 'Training completed!'
+        print(message)
+
+        if upload_to_hub:
+            upload_message = self.model_uploader.upload_lora_model(
+                folder_path=output_dir.as_posix(),
+                repo_name=output_model_name,
+                upload_to=upload_to,
+                private=use_private_repo,
+                delete_existing_repo=delete_existing_repo)
+            print(upload_message)
+            message = message + '\n' + upload_message
 
         if remove_gpu_after_training:
             space_id = os.getenv('SPACE_ID')
@@ -138,8 +144,4 @@ class Trainer:
                 self.api.request_space_hardware(repo_id=space_id,
                                                 hardware='cpu-basic')
 
-        with open(output_dir / 'train.sh', 'w') as f:
-            command_s = ' '.join(command.split())
-            f.write(command_s)
-
-        return 'Training completed!'
+        return message

uploader.py

@@ -35,5 +35,4 @@ class Uploader:
             message = f'Your model was successfully uploaded to <a href="{url}" target="_blank">{url}</a>.'
         except Exception as e:
             message = str(e)
-        print(message)
         return message