Fix locally and let it work with A10G

app.py

@@ -1,5 +1,17 @@
-import gradio as gr
+from subprocess import getoutput
 import os
+
+gpu_info = getoutput('nvidia-smi')
+if("A10G" in gpu_info):
+    which_gpu = "A10G"
+    os.system(f"pip install -q https://github.com/camenduru/stable-diffusion-webui-colab/releases/download/0.0.15/xformers-0.0.15.dev0+4c06c79.d20221205-cp38-cp38-linux_x86_64.whl")
+elif("T4" in gpu_info):
+    which_gpu = "T4"
+    os.system(f"pip install -q https://github.com/camenduru/stable-diffusion-webui-colab/releases/download/0.0.15/xformers-0.0.15.dev0+1515f77.d20221130-cp38-cp38-linux_x86_64.whl")
+else:
+    which_gpu = "CPU"
+    
+import gradio as gr
 from pathlib import Path
 import argparse
 import shutil
@@ -16,7 +28,6 @@ import gc
 from diffusers import StableDiffusionPipeline
 from huggingface_hub import snapshot_download, update_repo_visibility, HfApi
 
-
 is_spaces = True if "SPACE_ID" in os.environ else False
 if(is_spaces):
     is_shared_ui = True if "multimodalart/dreambooth-training" in os.environ['SPACE_ID'] else False
@@ -40,8 +51,8 @@ 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(".")
+#with zipfile.ZipFile("mix.zip", 'r') as zip_ref:
+#    zip_ref.extractall(".")
 
 def swap_text(option, base):
     resize_width = 768 if base == "v2-1-768" else 512
@@ -96,17 +107,19 @@ def count_files(*inputs):
             Training_Steps = 2400 #Avoid overfitting on person faces
     if(is_spaces):
         if(selected_model == "v1-5"):
-            its = 1.1
+            its = 1.1 if which_gpu == "T4" else 1.8
             if(experimental_faces):
                 its = 1
         elif(selected_model == "v2-1-512"):
-            its = 0.8
+            its = 0.8 if which_gpu == "T4" else 1.5
             if(experimental_faces):
                 its = 0.7
         elif(selected_model == "v2-1-768"):
-            its = 0.5
+            its = 0.48 if which_gpu == "T4" else 0.85
+       
+        gpu_price = 0.60 if which_gpu == "T4" else 1.10
         summary_sentence = f'''You are going to train {concept_counter} {type_of_thing}(s), with {file_counter} images for {Training_Steps} steps. The training should take around {round(Training_Steps/its, 2)} seconds, or {round((Training_Steps/its)/60, 2)} minutes.
-            The setup, compression and uploading the model can take up to 20 minutes.<br>As the T4-Small GPU costs US$0.60 for 1h, <span style="font-size: 120%"><b>the estimated cost for this training is below US${round((((Training_Steps/its)/3600)+0.3+0.1)*0.60, 2)}.</b></span><br><br>
+            The setup, compression and uploading the model can take up to 20 minutes.<br>As the {which_gpu}-Small GPU costs US${gpu_price} for 1h, <span style="font-size: 120%"><b>the estimated cost for this training is below US${round((((Training_Steps/its)/3600)+0.3+0.1)*gpu_price, 2)}.</b></span><br><br>
             If you check the box below the GPU attribution will automatically removed after training is done and the model is uploaded. If not, don't forget to come back here and swap the hardware back to CPU.<br><br>'''
     else:
         summary_sentence = f'''You are going to train {concept_counter} {type_of_thing}(s), with {file_counter} images for {Training_Steps} steps.<br><br>'''
@@ -141,7 +154,8 @@ def validate_model_upload(hf_token, model_name):
         except:
             raise gr.Error("You have inserted an invalid Hugging Face token")
         try:
-            update_repo_visibility(repo_id=os.environ['SPACE_ID'], private=True, token=hf_token, repo_type="space")
+            if(is_spaces):
+                update_repo_visibility(repo_id=os.environ['SPACE_ID'], private=True, token=hf_token, repo_type="space")
         except:
             raise gr.Error("Oops, you created a Hugging Face token with read permissions only. You need one with write permissions")
     else:
@@ -153,10 +167,14 @@ def train(*inputs):
     if is_shared_ui:
         raise gr.Error("This Space only works in duplicated instances")
     if not is_gpu_associated:
-        raise gr.Error("Please associate a T4 GPU for this Space")
+        raise gr.Error("Please associate a T4 or A10G GPU for this Space")
     hf_token = inputs[-5]
     model_name = inputs[-7]
-    remove_attribution_after = inputs[-6]
+    if(is_spaces):
+        remove_attribution_after = inputs[-6]
+    else:
+        remove_attribution_after = False
+    
     if(remove_attribution_after):
         validate_model_upload(hf_token, model_name)
     
@@ -469,7 +487,7 @@ with gr.Blocks(css=css) as demo:
             top_description = gr.HTML(f'''
                 <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 GPU for training. As each T4 costs US$0.60/h, it should cost < US$1 to train most models using default settings!&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>
+                <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" />
                 </div>
@@ -478,15 +496,15 @@ with gr.Blocks(css=css) as demo:
             if(is_gpu_associated):
                 top_description = gr.HTML(f'''
                         <div class="gr-prose" style="max-width: 80%">
-                        <h2>You have successfully associated a GPU to the Dreambooth Training Space 🎉</h2>
-                        <p>Certify that you got a T4. You can now train your model! You will be billed by the minute from when you activated the GPU until when it is turned it off.</p> 
+                        <h2>You have successfully associated a {which_gpu} GPU to the Dreambooth Training Space 🎉</h2>
+                        <p>You can now train your model! You will be billed by the minute from when you activated the GPU until when it is turned it off.</p> 
                         </div>
                 ''')
             else:
                 top_description = gr.HTML(f'''
                         <div class="gr-prose" style="max-width: 80%">
                         <h2>You have successfully duplicated the Dreambooth Training Space 🎉</h2>
-                        <p>There's only one step left before you can train your model: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4 GPU</b> to it (via the Settings tab)</a> and run the training below. Other GPUs are not compatible for now. You will be billed by the minute from when you activate the GPU until when it is turned it off.</p> 
+                        <p>There's only one step left before you can train your model: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4-small or A10G-small GPU</b> to it (via the Settings tab)</a> and run the training below. You will be billed by the minute from when you activate the GPU until when it is turned it off.</p> 
                         </div>
                 ''')
         else:
@@ -573,7 +591,7 @@ with gr.Blocks(css=css) as demo:
     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):
-        training_ongoing = gr.Markdown("## Oops, you haven't associated your T4 GPU to this Space. Visit the Settings tab, associate and try again.", visible=False)
+        training_ongoing = gr.Markdown("## Oops, you haven't associated your T4 or A10G GPU to this Space. Visit the Settings tab, associate and try again.", visible=False)
     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)
     

requirements.txt

@@ -14,5 +14,4 @@ triton==2.0.0.dev20220701
 bitsandbytes
 python-slugify
 requests
-tensorboard
-https://github.com/apolinario/xformers/releases/download/0.0.2/xformers-0.0.14.dev0-cp38-cp38-linux_x86_64.whl
+tensorboard