Remove GPU attribution if CUDA error

app.py

@@ -35,15 +35,8 @@ else:
     is_shared_ui = False
 is_gpu_associated = torch.cuda.is_available()
 
-css = '''
-    .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
-    .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
-    #component-4, #component-3, #component-10{min-height: 0}
-    .duplicate-button img{margin: 0}
-'''
-maximum_concepts = 3
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
-#Pre download the files
 if(is_gpu_associated):
     model_v1 = snapshot_download(repo_id="multimodalart/sd-fine-tunable")
     model_v2 = snapshot_download(repo_id="stabilityai/stable-diffusion-2-1", ignore_patterns=["*.ckpt", "*.safetensors"])
@@ -51,8 +44,25 @@ if(is_gpu_associated):
     safety_checker = snapshot_download(repo_id="multimodalart/sd-sc")
     model_to_load = model_v1
 
-#with zipfile.ZipFile("mix.zip", 'r') as zip_ref:
-#    zip_ref.extractall(".")
+def swap_base_model(selected_model):
+    if(is_gpu_associated):
+        global model_to_load
+        if(selected_model == "v1-5"):
+            model_to_load = model_v1
+        elif(selected_model == "v2-1-768"):
+            model_to_load = model_v2
+        else:
+            model_to_load = model_v2_512
+
+
+
+css = '''
+    .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
+    .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
+    #component-4, #component-3, #component-10{min-height: 0}
+    .duplicate-button img{margin: 0}
+'''
+maximum_concepts = 3
 
 def swap_text(option, base):
     resize_width = 768 if base == "v2-1-768" else 512
@@ -60,7 +70,7 @@ def swap_text(option, base):
     if(option == "object"):
         instance_prompt_example = "cttoy"
         freeze_for = 30
-        return [f"You are going to train `object`(s), upload 5-10 images of each object you are planning on training on from different angles/perspectives. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file/cat-toy.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}.", freeze_for, gr.update(visible=False)]
+        return [f"You are going to train `object`(s), upload 5-10 images of each object you are planning on training on from different angles/perspectives. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file=cat-toy.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}.", freeze_for, gr.update(visible=False)]
     elif(option == "person"):
        instance_prompt_example = "julcto"
        freeze_for = 70
@@ -70,27 +80,17 @@ def swap_text(option, base):
            prior_preservation_box_update = gr.update(visible=show_prior_preservation)
        else: 
            prior_preservation_box_update = gr.update(visible=show_prior_preservation, value=False)
-       return [f"You are going to train a `person`(s), upload 10-20 images of each person you are planning on training on from different angles/perspectives. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file/person.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}.", freeze_for, prior_preservation_box_update]
+       return [f"You are going to train a `person`(s), upload 10-20 images of each person you are planning on training on from different angles/perspectives. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file=person.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}.", freeze_for, prior_preservation_box_update]
     elif(option == "style"):
         instance_prompt_example = "trsldamrl"
         freeze_for = 10
-        return [f"You are going to train a `style`, upload 10-20 images of the style you are planning on training on. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file/trsl_style.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}", freeze_for, gr.update(visible=False)]
-
-def swap_base_model(selected_model):
-    if(is_gpu_associated):
-        global model_to_load
-        if(selected_model == "v1-5"):
-            model_to_load = model_v1
-        elif(selected_model == "v2-1-768"):
-            model_to_load = model_v2
-        else:
-            model_to_load = model_v2_512
+        return [f"You are going to train a `style`, upload 10-20 images of the style you are planning on training on. You can use services like <a style='text-decoration: underline' target='_blank' href='https://www.birme.net/?target_width={resize_width}&target_height={resize_width}'>birme</a> for smart cropping. {mandatory_liability}:", '''<img src="file=trsl_style.png" />''', f"You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `{instance_prompt_example}` here). Images will be automatically cropped to {resize_width}x{resize_width}", freeze_for, gr.update(visible=False)]
 
 def count_files(*inputs):
     file_counter = 0
     concept_counter = 0
     for i, input in enumerate(inputs):
-        if(i < maximum_concepts-1):
+        if(i < maximum_concepts):
             files = inputs[i]
             if(files):
                 concept_counter+=1
@@ -133,6 +133,9 @@ def update_steps(*files_list):
             file_counter+=len(files)
     return(gr.update(value=file_counter*200))
 
+def visualise_progress_bar():
+    return gr.update(visible=True)
+
 def pad_image(image):
     w, h = image.size
     if w == h:
@@ -163,7 +166,34 @@ def validate_model_upload(hf_token, model_name):
     if(model_name == ""):
         raise gr.Error("Please fill in your model's name")
 
-def train(*inputs):
+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 swap_sleep_time(hf_token,sleep_time):
+    sleep_time_url = f"https://huggingface.co/api/spaces/{os.environ['SPACE_ID']}/sleeptime"
+    headers = { "authorization" : f"Bearer {hf_token}"}
+    body = {'seconds':sleep_time}
+    requests.post(sleep_time_url,json=body,headers=headers)
+
+def get_sleep_time(hf_token):
+    sleep_time_url = f"https://huggingface.co/api/spaces/{os.environ['SPACE_ID']}"
+    headers = { "authorization" : f"Bearer {hf_token}"}
+    response = requests.get(sleep_time_url,headers=headers)
+    return response.json()['runtime']['gcTimeout']
+
+def write_to_community(title, description,hf_token): 
+    from huggingface_hub import HfApi
+    api = HfApi()
+    api.create_discussion(repo_id=os.environ['SPACE_ID'], title=title, description=description,repo_type="space", token=hf_token)
+
+def train(progress=gr.Progress(track_tqdm=True), *inputs):
+    which_model = inputs[-10]
+    if(which_model == ""):
+        raise gr.Error("You forgot to select a base model to use")
+        
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
     if not is_gpu_associated:
@@ -171,6 +201,9 @@ def train(*inputs):
     hf_token = inputs[-5]
     model_name = inputs[-7]
     if(is_spaces):
+        sleep_time = get_sleep_time(hf_token)
+        if sleep_time:
+            swap_sleep_time(hf_token, -1)
         remove_attribution_after = inputs[-6]
     else:
         remove_attribution_after = False
@@ -191,7 +224,6 @@ def train(*inputs):
     if os.path.exists("model.ckpt"): os.remove("model.ckpt")
     if os.path.exists("hastrained.success"): os.remove("hastrained.success")
     file_counter = 0
-    which_model = inputs[-10]
     resolution = 512 if which_model != "v2-1-768" else 768
     for i, input in enumerate(inputs):
         if(i < maximum_concepts-1):
@@ -261,38 +293,22 @@ def train(*inputs):
         print("Starting single training...")
         lock_file = open("intraining.lock", "w")
         lock_file.close()
-        run_training(args_general)
-    else:
-        args_general = argparse.Namespace(
-            image_captions_filename = True,
-            train_text_encoder = True if stptxt > 0 else False,
-            stop_text_encoder_training = stptxt,
-            save_n_steps = 0,
-            pretrained_model_name_or_path = model_to_load,
-            instance_data_dir="instance_images",
-            class_data_dir="Mix",
-            output_dir="output_model",
-            with_prior_preservation=True,
-            prior_loss_weight=1.0,
-            instance_prompt="",
-            seed=42,
-            resolution=resolution,
-            mixed_precision="fp16",
-            train_batch_size=1,
-            gradient_accumulation_steps=1,
-            use_8bit_adam=True,
-            learning_rate=2e-6,
-            lr_scheduler="polynomial",
-            lr_warmup_steps = 0,
-            max_train_steps=Training_Steps,
-            num_class_images=200,     
-            gradient_checkpointing=gradient_checkpointing,
-            cache_latents=cache_latents,
-        )
-        print("Starting multi-training...")
-        lock_file = open("intraining.lock", "w")
-        lock_file.close()
-        run_training(args_general)
+        try:
+            run_training(args_general)
+        except Exception as e:
+            if(is_spaces):
+                title="There was an error on during your training"
+                description=f'''
+                Unfortunately there was an error during training your {model_name} model. 
+                Please check it out below. Feel free to report this issue to [Dreambooth Training](https://huggingface.co/spaces/multimodalart/dreambooth-training): 
+                ```
+                {str(e)}
+                ```
+                '''
+                swap_hardware(hf_token, "cpu-basic")
+                write_to_community(title,description,hf_token)
+            
+            
     gc.collect()
     torch.cuda.empty_cache()
     if(which_model == "v1-5"):
@@ -302,6 +318,7 @@ def train(*inputs):
         shutil.copy(f"model_index.json", "output_model/model_index.json")
     
     if(not remove_attribution_after):
+        swap_sleep_time(hf_token, sleep_time)
         print("Archiving model file...")
         with tarfile.open("diffusers_model.tar", "w") as tar:
             tar.add("output_model", arcname=os.path.basename("output_model"))
@@ -310,6 +327,7 @@ def train(*inputs):
         trained_file.close()
         print("Training completed!")
         return [
+            gr.update(visible=False), #progress_bar
             gr.update(visible=True, value=["diffusers_model.tar"]), #result
             gr.update(visible=True), #try_your_model
             gr.update(visible=True), #push_to_hub
@@ -320,10 +338,7 @@ def train(*inputs):
     else:
         where_to_upload = inputs[-8]
         push(model_name, where_to_upload, hf_token, which_model, True)
-        hardware_url = f"https://huggingface.co/spaces/{os.environ['SPACE_ID']}/hardware"
-        headers = { "authorization" : f"Bearer {hf_token}"}
-        body = {'flavor': 'cpu-basic'}
-        requests.post(hardware_url, json = body, headers=headers)
+        swap_hardware(hf_token, "cpu-basic")
 
 pipe_is_set = False
 def generate(prompt, steps):
@@ -338,7 +353,7 @@ def generate(prompt, steps):
         
     image = pipe(prompt, num_inference_steps=steps).images[0]  
     return(image)
-    
+
 def push(model_name, where_to_upload, hf_token, which_model, comes_from_automated=False):
     validate_model_upload(hf_token, model_name)
     if(not os.path.exists("model.ckpt")):
@@ -425,7 +440,10 @@ Sample pictures of:
             extra_message = "Don't forget to remove the GPU attribution after you play with it."
         else:
             extra_message = "The GPU has been removed automatically as requested, and you can try the model via the model page"
-        api.create_discussion(repo_id=os.environ['SPACE_ID'], title=f"Your model {model_name} has finished trained from the Dreambooth Train Spaces!", description=f"Your model has been successfully uploaded to: https://huggingface.co/{model_id}. {extra_message}",repo_type="space", token=hf_token)
+        title=f"Your model {model_name} has finished trained from the Dreambooth Train Spaces!"
+        description=f"Your model has been successfully uploaded to: https://huggingface.co/{model_id}. {extra_message}"
+        write_to_community(title, description, hf_token)
+        #api.create_discussion(repo_id=os.environ['SPACE_ID'], title=f"Your model {model_name} has finished trained from the Dreambooth Train Spaces!", description=f"Your model has been successfully uploaded to: https://huggingface.co/{model_id}. {extra_message}",repo_type="space", token=hf_token)
     print("Model uploaded successfully!")
     return [gr.update(visible=True, value=f"Successfully uploaded your model. Access it [here](https://huggingface.co/{model_id})"), gr.update(visible=True, value=["diffusers_model.tar", "model.ckpt"])]
 
@@ -488,8 +506,8 @@ with gr.Blocks(css=css) as demo:
                 <div class="gr-prose" style="max-width: 80%">
                 <h2>Attention - This Space doesn't work in this shared UI</h2>
                 <p>For it to work, you can either run locally or duplicate the Space and run it on your own profile using a (paid) private T4-small or A10G-small GPU for training. A T4 costs US$0.60/h, so it should cost < US$1 to train most models using default settings with it!&nbsp;&nbsp;<a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p>
-                <img class="instruction" src="file/duplicate.png"> 
-                <img class="arrow" src="file/arrow.png" />
+                <img class="instruction" src="file=duplicate.png"> 
+                <img class="arrow" src="file=arrow.png" />
                 </div>
             ''')
         elif(is_spaces):
@@ -519,14 +537,15 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Row() as what_are_you_training:
         type_of_thing = gr.Dropdown(label="What would you like to train?", choices=["object", "person", "style"], value="object", interactive=True)
-        base_model_to_use = gr.Dropdown(label="Which base model would you like to use?", choices=["v1-5", "v2-1-512", "v2-1-768"], value="v1-5", interactive=True)
-    
+        with gr.Column():
+            base_model_to_use = gr.Dropdown(label="Which base model would you like to use?", choices=["v1-5", "v2-1-512", "v2-1-768"], value="v1-5", interactive=True)
+
     #Very hacky approach to emulate dynamically created Gradio components   
     with gr.Row() as upload_your_concept:
         with gr.Column():
             thing_description = gr.Markdown("You are going to train an `object`, please upload 5-10 images of the object you are planning on training on from different angles/perspectives. You must have the right to do so and you are liable for the images you use, example")
             thing_experimental = gr.Checkbox(label="Improve faces (prior preservation) - can take longer training but can improve faces", visible=False, value=False)
-            thing_image_example = gr.HTML('''<img src="file/cat-toy.png" />''')
+            thing_image_example = gr.HTML('''<img src="file=cat-toy.png" />''')
             things_naming = gr.Markdown("You should name your concept with a unique made up word that has low chance of the model already knowing it (e.g.: `cttoy` here). Images will be automatically cropped to 512x512.")
             
         with gr.Column():
@@ -588,6 +607,7 @@ with gr.Blocks(css=css) as demo:
         training_summary_token = gr.Textbox(label="Hugging Face Write Token", type="password", visible=True)
         
     train_btn = gr.Button("Start Training")
+    progress_bar = gr.Textbox(visible=False)
     if(is_shared_ui):
         training_ongoing = gr.Markdown("## This Space only works in duplicated instances. Please duplicate it and try again!", visible=False)
     elif(not is_gpu_associated):
@@ -595,6 +615,7 @@ with gr.Blocks(css=css) as demo:
     else:
         training_ongoing = gr.Markdown("## Training is ongoing ⌛... You can close this tab if you like or just wait. If you did not check the `Remove GPU After training`, you can come back here to try your model and upload it after training. Don't forget to remove the GPU attribution after you are done. ", visible=False)
     
+    
     #Post-training UI
     completed_training = gr.Markdown('''# ✅ Training completed. 
     ### Don't forget to remove the GPU attribution after you are done trying and uploading your model''', visible=False)
@@ -624,9 +645,10 @@ with gr.Blocks(css=css) as demo:
     type_of_thing.change(fn=swap_text, inputs=[type_of_thing, base_model_to_use], outputs=[thing_description, thing_image_example, things_naming, perc_txt_encoder, thing_experimental], queue=False, show_progress=False)
     
     #Swap the base model
+    
     base_model_to_use.change(fn=swap_text, inputs=[type_of_thing, base_model_to_use], outputs=[thing_description, thing_image_example, things_naming, perc_txt_encoder, thing_experimental], queue=False, show_progress=False)
+    #base_model_to_use.change(fn=visualise_progress_bar, inputs=[], outputs=progress_bar)
     base_model_to_use.change(fn=swap_base_model, inputs=base_model_to_use, outputs=[])
-
     #Update the summary box below the UI according to how many images are uploaded and whether users are using custom settings or not 
     for file in file_collection:
         #file.change(fn=update_steps,inputs=file_collection, outputs=steps)
@@ -641,10 +663,12 @@ with gr.Blocks(css=css) as demo:
     if(is_spaces):
         training_summary_checkbox.change(fn=checkbox_swap, inputs=training_summary_checkbox, outputs=[training_summary_token_message, training_summary_token, training_summary_model_name, training_summary_where_to_upload],queue=False, show_progress=False)
     #Add a message for while it is in training
-    train_btn.click(lambda:gr.update(visible=True), inputs=None, outputs=training_ongoing)
+    
+    #train_btn.click(lambda:gr.update(visible=True), inputs=None, outputs=training_ongoing)
     
     #The main train function
-    train_btn.click(fn=train, inputs=is_visible+concept_collection+file_collection+[base_model_to_use]+[thing_experimental]+[training_summary_where_to_upload]+[training_summary_model_name]+[training_summary_checkbox]+[training_summary_token]+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[result, try_your_model, push_to_hub, convert_button, training_ongoing, completed_training], queue=False)
+    train_btn.click(lambda:gr.update(visible=True), inputs=[], outputs=progress_bar)
+    train_btn.click(fn=train, inputs=is_visible+concept_collection+file_collection+[base_model_to_use]+[thing_experimental]+[training_summary_where_to_upload]+[training_summary_model_name]+[training_summary_checkbox]+[training_summary_token]+[type_of_thing]+[steps]+[perc_txt_encoder]+[swap_auto_calculated], outputs=[progress_bar, result, try_your_model, push_to_hub, convert_button, training_ongoing, completed_training], queue=False)
     
     #Button to generate an image from your trained model after training
     generate_button.click(fn=generate, inputs=[prompt, inference_steps], outputs=result_image, queue=False)