Apply pre-commit

app.py

@@ -38,7 +38,7 @@ You can use "T4 small/medium" to run this demo.
 </center>
 '''
 
-HF_TOKEN_NOT_SPECIFIED_WARNING = f'''The environment variable `HF_TOKEN` is not specified. Feel free to specify your Hugging Face token with write permission if you don't want to manually provide it for every run. 
+HF_TOKEN_NOT_SPECIFIED_WARNING = f'''The environment variable `HF_TOKEN` is not specified. Feel free to specify your Hugging Face token with write permission if you don't want to manually provide it for every run.
 <center>
 You can check and create your Hugging Face tokens <a href="https://huggingface.co/settings/tokens" target="_blank">here</a>.
 You can specify environment variables in the "Repository secrets" section of the {SETTINGS} tab.
@@ -63,9 +63,8 @@ with gr.Blocks(css='style.css') as demo:
         show_warning(SHARED_UI_WARNING)
     elif not torch.cuda.is_available():
         show_warning(CUDA_NOT_AVAILABLE_WARNING)
-    elif(not "T4" in GPU_DATA):
+    elif (not 'T4' in GPU_DATA):
         show_warning(INVALID_GPU_WARNING)
-    
 
     gr.Markdown(TITLE)
     with gr.Tabs():
@@ -78,8 +77,8 @@ with gr.Blocks(css='style.css') as demo:
             - You can use this tab to upload models later if you choose not to upload models in training time or if upload in training time failed.
             ''')
             create_upload_demo(HF_TOKEN)
-    
+
     if not HF_TOKEN:
         show_warning(HF_TOKEN_NOT_SPECIFIED_WARNING)
 
-demo.queue(max_size=1).launch(share=False)
+demo.queue(max_size=1).launch(share=False)

app_training.py

@@ -32,18 +32,23 @@ def create_training_demo(trainer: Trainer,
                 with gr.Box():
                     gr.Markdown('Training Parameters')
                     with gr.Row():
-                        base_model = gr.Text(label='Base Model',
-                                             value='CompVis/stable-diffusion-v1-4',
-                                             max_lines=1)
+                        base_model = gr.Text(
+                            label='Base Model',
+                            value='CompVis/stable-diffusion-v1-4',
+                            max_lines=1)
                         resolution = gr.Dropdown(choices=['512', '768'],
                                                  value='512',
                                                  label='Resolution',
                                                  visible=False)
-                    
-                    input_token = gr.Text(label="Hugging Face Write Token", placeholder="", visible=False if hf_token else True)
-                    with gr.Accordion("Advanced settings", open=False):
+
+                    input_token = gr.Text(label='Hugging Face Write Token',
+                                          placeholder='',
+                                          visible=False if hf_token else True)
+                    with gr.Accordion('Advanced settings', open=False):
                         num_training_steps = gr.Number(
-                            label='Number of Training Steps', value=300, precision=0)
+                            label='Number of Training Steps',
+                            value=300,
+                            precision=0)
                         learning_rate = gr.Number(label='Learning Rate',
                                                   value=0.000035)
                         gradient_accumulation = gr.Number(
@@ -57,33 +62,36 @@ def create_training_demo(trainer: Trainer,
                                          randomize=True,
                                          value=0)
                         fp16 = gr.Checkbox(label='FP16', value=True)
-                        use_8bit_adam = gr.Checkbox(label='Use 8bit Adam', value=False)
-                        checkpointing_steps = gr.Number(label='Checkpointing Steps',
-                                                        value=1000,
-                                                        precision=0)
-                        validation_epochs = gr.Number(label='Validation Epochs',
-                                                      value=100,
-                                                      precision=0)
+                        use_8bit_adam = gr.Checkbox(label='Use 8bit Adam',
+                                                    value=False)
+                        checkpointing_steps = gr.Number(
+                            label='Checkpointing Steps',
+                            value=1000,
+                            precision=0)
+                        validation_epochs = gr.Number(
+                            label='Validation Epochs', value=100, precision=0)
                     gr.Markdown('''
                         - The base model must be a Stable Diffusion model compatible with [diffusers](https://github.com/huggingface/diffusers) library.
                         - Expected time to train a model for 300 steps: ~20 minutes with T4
                         - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
                         ''')
-                    
+
         with gr.Row():
             with gr.Column():
                 gr.Markdown('Output Model')
                 output_model_name = gr.Text(label='Name of your model',
                                             placeholder='The surfer man',
                                             max_lines=1)
-                validation_prompt = gr.Text(label='Validation Prompt', placeholder='prompt to test the model, e.g: a dog is surfing')
+                validation_prompt = gr.Text(
+                    label='Validation Prompt',
+                    placeholder=
+                    'prompt to test the model, e.g: a dog is surfing')
             with gr.Column():
                 gr.Markdown('Upload Settings')
                 with gr.Row():
-                    upload_to_hub = gr.Checkbox(
-                        label='Upload model to Hub', value=True)
-                    use_private_repo = gr.Checkbox(label='Private',
-                                                   value=True)
+                    upload_to_hub = gr.Checkbox(label='Upload model to Hub',
+                                                value=True)
+                    use_private_repo = gr.Checkbox(label='Private', value=True)
                     delete_existing_repo = gr.Checkbox(
                         label='Delete existing repo of the same name',
                         value=False)
@@ -91,45 +99,32 @@ def create_training_demo(trainer: Trainer,
                         label='Upload to',
                         choices=[_.value for _ in UploadTarget],
                         value=UploadTarget.MODEL_LIBRARY.value)
-                    
+
         remove_gpu_after_training = gr.Checkbox(
             label='Remove GPU after training',
             value=False,
             interactive=bool(os.getenv('SPACE_ID')),
             visible=False)
         run_button = gr.Button('Start Training')
-    
+
         with gr.Box():
             gr.Markdown('Output message')
             output_message = gr.Markdown()
 
         if pipe is not None:
             run_button.click(fn=pipe.clear)
-        run_button.click(fn=trainer.run,
-                         inputs=[
-                             training_video,
-                             training_prompt,
-                             output_model_name,
-                             delete_existing_repo,
-                             validation_prompt,
-                             base_model,
-                             resolution,
-                             num_training_steps,
-                             learning_rate,
-                             gradient_accumulation,
-                             seed,
-                             fp16,
-                             use_8bit_adam,
-                             checkpointing_steps,
-                             validation_epochs,
-                             upload_to_hub,
-                             use_private_repo,
-                             delete_existing_repo,
-                             upload_to,
-                             remove_gpu_after_training,
-                             input_token
-                         ],
-                         outputs=output_message)
+        run_button.click(
+            fn=trainer.run,
+            inputs=[
+                training_video, training_prompt, output_model_name,
+                delete_existing_repo, validation_prompt, base_model,
+                resolution, num_training_steps, learning_rate,
+                gradient_accumulation, seed, fp16, use_8bit_adam,
+                checkpointing_steps, validation_epochs, upload_to_hub,
+                use_private_repo, delete_existing_repo, upload_to,
+                remove_gpu_after_training, input_token
+            ],
+            outputs=output_message)
     return demo
 
 

app_upload.py

@@ -70,7 +70,9 @@ def create_upload_demo(hf_token: str | None) -> gr.Blocks:
                                  choices=[_.value for _ in UploadTarget],
                                  value=UploadTarget.MODEL_LIBRARY.value)
             model_name = gr.Textbox(label='Model Name')
-            input_token = gr.Text(label="Hugging Face Write Token", placeholder="", visible=False if hf_token else True)
+            input_token = gr.Text(label='Hugging Face Write Token',
+                                  placeholder='',
+                                  visible=False if hf_token else True)
         upload_button = gr.Button('Upload')
         gr.Markdown(f'''
             - You can upload your trained model to your personal profile (i.e. https://huggingface.co/{{your_username}}/{{model_name}}) or to the public [Tune-A-Video Library](https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}) (i.e. https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}/{{model_name}}).

trainer.py

@@ -23,6 +23,7 @@ URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Vid
 ORIGINAL_SPACE_ID = 'Tune-A-Video-library/Tune-A-Video-Training-UI'
 SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
 
+
 class Trainer:
     def __init__(self, hf_token: str | None = None):
         self.hf_token = hf_token
@@ -73,7 +74,9 @@ class Trainer:
         input_token: str,
     ) -> str:
         if SPACE_ID == ORIGINAL_SPACE_ID:
-            raise gr.Error('This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU')
+            raise gr.Error(
+                'This Space does not work on this Shared UI. Duplicate the Space and attribute a GPU'
+            )
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
         if training_video is None:
@@ -97,7 +100,8 @@ class Trainer:
         output_dir.mkdir(parents=True)
 
         if upload_to_hub:
-            self.join_model_library_org(self.hf_token if self.hf_token else input_token)
+            self.join_model_library_org(
+                self.hf_token if self.hf_token else input_token)
 
         config = OmegaConf.load('Tune-A-Video/configs/man-surfing.yaml')
         config.pretrained_model_path = self.download_base_model(base_model)
@@ -154,7 +158,8 @@ class Trainer:
         if remove_gpu_after_training:
             space_id = os.getenv('SPACE_ID')
             if space_id:
-                api = HfApi(token=self.hf_token if self.hf_token else input_token)
+                api = HfApi(
+                    token=self.hf_token if self.hf_token else input_token)
                 api.request_space_hardware(repo_id=space_id,
                                            hardware='cpu-basic')
 

uploader.py

@@ -15,23 +15,23 @@ class Uploader:
                private: bool = True,
                delete_existing_repo: bool = False,
                input_token: str | None = None) -> str:
-        
+
         api = HfApi(token=self.hf_token if self.hf_token else input_token)
-        
+
         if not folder_path:
             raise ValueError
         if not repo_name:
             raise ValueError
         if not organization:
             organization = api.whoami()['name']
-        
+
         repo_id = f'{organization}/{repo_name}'
         if delete_existing_repo:
             try:
-                self.api.delete_repo(repo_id, repo_type=repo_type)
+                api.delete_repo(repo_id, repo_type=repo_type)
             except Exception:
                 pass
-        try:                
+        try:
             api.create_repo(repo_id, repo_type=repo_type, private=private)
             api.upload_folder(repo_id=repo_id,
                               folder_path=folder_path,