hysts commited on
Commit
f689e96
1 Parent(s): 94efb76

Let users jointhe LoRA Library organization

Browse files
Files changed (2) hide show
  1. app_training.py +1 -2
  2. trainer.py +12 -0
app_training.py CHANGED
@@ -48,8 +48,7 @@ def create_training_demo(trainer: Trainer,
48
  choices=[_.value for _ in UploadTarget],
49
  value=UploadTarget.LORA_LIBRARY.value)
50
  gr.Markdown('''
51
- - By default, trained models will be uploaded to [LoRA Library](https://huggingface.co/lora-library) (See [this example model](https://huggingface.co/lora-library/lora-dreambooth-sample-dog)).
52
- Note that it will fail if you are not a member of the organization. So please join the org first. In the case uploading failed, you can use the "Upload" tab to upload your model later.
53
  - You can also choose "Personal Profile", in which case, the model will be uploaded to https://huggingface.co/{your_username}/{model_name}.
54
  ''')
55
 
 
48
  choices=[_.value for _ in UploadTarget],
49
  value=UploadTarget.LORA_LIBRARY.value)
50
  gr.Markdown('''
51
+ - By default, trained models will be uploaded to [LoRA Library](https://huggingface.co/lora-library) (see [this example model](https://huggingface.co/lora-library/lora-dreambooth-sample-dog)).
 
52
  - You can also choose "Personal Profile", in which case, the model will be uploaded to https://huggingface.co/{your_username}/{model_name}.
53
  ''')
54
 
trainer.py CHANGED
@@ -16,6 +16,8 @@ from huggingface_hub import HfApi
16
  from app_upload import LoRAModelUploader
17
  from utils import save_model_card
18
 
 
 
19
 
20
  def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
21
  w, h = image.size
@@ -33,6 +35,7 @@ def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
33
 
34
  class Trainer:
35
  def __init__(self, hf_token: str | None = None):
 
36
  self.api = HfApi(token=hf_token)
37
  self.model_uploader = LoRAModelUploader(hf_token)
38
 
@@ -48,6 +51,12 @@ class Trainer:
48
  out_path = instance_data_dir / f'{i:03d}.jpg'
49
  image.save(out_path, format='JPEG', quality=100)
50
 
 
 
 
 
 
 
51
  def run(
52
  self,
53
  instance_images: list | None,
@@ -97,6 +106,9 @@ class Trainer:
97
  instance_data_dir = repo_dir / 'training_data' / output_model_name
98
  self.prepare_dataset(instance_images, resolution, instance_data_dir)
99
 
 
 
 
100
  command = f'''
101
  accelerate launch train_dreambooth_lora.py \
102
  --pretrained_model_name_or_path={base_model} \
 
16
  from app_upload import LoRAModelUploader
17
  from utils import save_model_card
18
 
19
+ URL_TO_JOIN_LORA_LIBRARY_ORG = 'https://huggingface.co/organizations/lora-library/share/hjetHAcKjnPHXhHfbeEcqnBqmhgilFfpOL'
20
+
21
 
22
  def pad_image(image: PIL.Image.Image) -> PIL.Image.Image:
23
  w, h = image.size
 
35
 
36
  class Trainer:
37
  def __init__(self, hf_token: str | None = None):
38
+ self.hf_token = hf_token
39
  self.api = HfApi(token=hf_token)
40
  self.model_uploader = LoRAModelUploader(hf_token)
41
 
 
51
  out_path = instance_data_dir / f'{i:03d}.jpg'
52
  image.save(out_path, format='JPEG', quality=100)
53
 
54
+ def join_lora_library_org(self) -> None:
55
+ subprocess.run(
56
+ shlex.split(
57
+ f'curl -X POST -H "Authorization: Bearer {self.hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_LORA_LIBRARY_ORG}'
58
+ ))
59
+
60
  def run(
61
  self,
62
  instance_images: list | None,
 
106
  instance_data_dir = repo_dir / 'training_data' / output_model_name
107
  self.prepare_dataset(instance_images, resolution, instance_data_dir)
108
 
109
+ if upload_to_hub:
110
+ self.join_lora_library_org()
111
+
112
  command = f'''
113
  accelerate launch train_dreambooth_lora.py \
114
  --pretrained_model_name_or_path={base_model} \