Update app.py

app.py

@@ -30,6 +30,12 @@ def change_training_setup(training_type):
     elif training_type == "concept" :
         return 2000, 1000
 
+def swap_hardware(hf_token, hardware="cpu-basic"):
+    hardware_url = f"https://huggingface.co/spaces/{os.environ['SPACE_ID']}/hardware"
+    headers = { "authorization" : f"Bearer {hf_token}"}
+    body = {'flavor': hardware}
+    requests.post(hardware_url, json = body, headers=headers)
+    
 def train_dreambooth_blora_sdxl(instance_data_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps):
     
     script_filename = "train_dreambooth_b-lora_sdxl.py"  # Assuming it's in the same folder
@@ -104,6 +110,8 @@ def main(image_path, b_lora_trained_folder, instance_prompt, class_prompt, train
     train_dreambooth_blora_sdxl(local_dir, b_lora_trained_folder, instance_prompt, class_prompt, max_train_steps, checkpoint_steps)
     
     your_username = api.whoami(token=hf_token)["name"]
+
+    swap_hardware(hf_token, hardware="cpu-basic")
     
     return f"Done, your trained model has been stored in your models library: {your_username}/{b_lora_trained_folder}"