simplified-ui (#3)

- Simplified UI (17a686c02af144dbc69aabd538b946bb34cb3c4e)


Co-authored-by: Multimodal AI art <[email protected]>

app.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+from subprocess import getoutput
 
 import gradio as gr
 import torch
@@ -13,27 +14,31 @@ from app_upload import create_upload_demo
 from inference import InferencePipeline
 from trainer import Trainer
 
-TITLE = '# [Tune-A-Video](https://tuneavideo.github.io/) Training UI'
+TITLE = '# [Tune-A-Video](https://tuneavideo.github.io/) UI'
 
 ORIGINAL_SPACE_ID = 'Tune-A-Video-library/Tune-A-Video-Training-UI'
 SPACE_ID = os.getenv('SPACE_ID', ORIGINAL_SPACE_ID)
-SHARED_UI_WARNING = f'''# Attention - This Space doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU. (Please note that there seems to be an issue with training on the A10G GPU now. The model doesn't learn anything when trained on A10G. Training on T4 works perfectly fine and inference works fine on both.)
+GPU_DATA = getoutput('nvidia-smi')
+SHARED_UI_WARNING = f'''## Attention - Training doesn't work in this shared UI. You can duplicate and use it with a paid private T4 GPU.
 
-<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{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></center>
+<center><a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="margin-top:0;margin-bottom:0" 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></center>
 '''
 
 if os.getenv('SYSTEM') == 'spaces' and SPACE_ID != ORIGINAL_SPACE_ID:
     SETTINGS = f'<a href="https://huggingface.co/spaces/{SPACE_ID}/settings">Settings</a>'
 else:
     SETTINGS = 'Settings'
-CUDA_NOT_AVAILABLE_WARNING = f'''# Attention - Running on CPU.
+
+INVALID_GPU_WARNING = f'''## Attention - the specified GPU is invalid. Training may not work. Make sure you have selected a `T4 GPU` for this task.'''
+
+CUDA_NOT_AVAILABLE_WARNING = f'''## Attention - Running on CPU.
 <center>
 You can assign a GPU in the {SETTINGS} tab if you are running this on HF Spaces.
-You can use "T4 small/medium" or "A10G small/large" to run this demo.
+You can use "T4 small/medium" to run this demo.
 </center>
 '''
 
-HF_TOKEN_NOT_SPECIFIED_WARNING = f'''# Attention - The environment variable `HF_TOKEN` is not specified. Please specify your Hugging Face token with write permission as the value of it.
+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.
@@ -54,23 +59,27 @@ pipe = InferencePipeline(HF_TOKEN)
 trainer = Trainer(HF_TOKEN)
 
 with gr.Blocks(css='style.css') as demo:
-    if os.getenv('IS_SHARED_UI'):
+    if SPACE_ID == ORIGINAL_SPACE_ID:
         show_warning(SHARED_UI_WARNING)
-    if not torch.cuda.is_available():
+    elif not torch.cuda.is_available():
         show_warning(CUDA_NOT_AVAILABLE_WARNING)
-    if not HF_TOKEN:
-        show_warning(HF_TOKEN_NOT_SPECIFIED_WARNING)
+    elif(not "T4" in GPU_DATA):
+        show_warning(INVALID_GPU_WARNING)
+    
 
     gr.Markdown(TITLE)
     with gr.Tabs():
         with gr.TabItem('Train'):
             create_training_demo(trainer, pipe)
-        with gr.TabItem('Test'):
+        with gr.TabItem('Run'):
             create_inference_demo(pipe, HF_TOKEN)
         with gr.TabItem('Upload'):
             gr.Markdown('''
             - 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

@@ -13,6 +13,7 @@ from trainer import Trainer
 
 def create_training_demo(trainer: Trainer,
                          pipe: InferencePipeline | None = None) -> gr.Blocks:
+    hf_token = os.getenv('HF_TOKEN')
     with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
@@ -24,82 +25,80 @@ def create_training_demo(trainer: Trainer,
                         max_lines=1,
                         placeholder='A man is surfing')
                     gr.Markdown('''
-                        - Upload a video and write a prompt describing the video.
+                        - Upload a video and write a `Training Prompt` that describes the video.
                         ''')
+
+            with gr.Column():
                 with gr.Box():
-                    gr.Markdown('Output Model')
-                    output_model_name = gr.Text(label='Name of your model',
-                                                max_lines=1)
-                    delete_existing_model = gr.Checkbox(
-                        label='Delete existing model of the same name',
-                        value=False)
-                    validation_prompt = gr.Text(label='Validation Prompt')
-                with gr.Box():
-                    gr.Markdown('Upload Settings')
+                    gr.Markdown('Training Parameters')
                     with gr.Row():
-                        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)
+                        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):
+                        num_training_steps = gr.Number(
+                            label='Number of Training Steps', value=300, precision=0)
+                        learning_rate = gr.Number(label='Learning Rate',
+                                                  value=0.000035)
+                        gradient_accumulation = gr.Number(
+                            label='Number of Gradient Accumulation',
+                            value=1,
+                            precision=0)
+                        seed = gr.Slider(label='Seed',
+                                         minimum=0,
+                                         maximum=100000,
+                                         step=1,
+                                         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)
+                    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')
+            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)
+                    delete_existing_repo = gr.Checkbox(
+                        label='Delete existing repo of the same name',
+                        value=False)
                     upload_to = gr.Radio(
                         label='Upload to',
                         choices=[_.value for _ in UploadTarget],
                         value=UploadTarget.MODEL_LIBRARY.value)
-                    gr.Markdown(f'''
-                    - By default, trained models will be uploaded to [Tune-A-Video Library](https://huggingface.co/{MODEL_LIBRARY_ORG_NAME}) (see [this example model](https://huggingface.co/{SAMPLE_MODEL_REPO})).
-                    - You can also choose "Personal Profile", in which case, the model will be uploaded to https://huggingface.co/{{your_username}}/{{model_name}}.
-                    ''')
-
-            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)
-                    resolution = gr.Dropdown(choices=['512', '768'],
-                                             value='512',
-                                             label='Resolution',
-                                             visible=False)
-                num_training_steps = gr.Number(
-                    label='Number of Training Steps', value=300, precision=0)
-                learning_rate = gr.Number(label='Learning Rate',
-                                          value=0.000035)
-                gradient_accumulation = gr.Number(
-                    label='Number of Gradient Accumulation',
-                    value=1,
-                    precision=0)
-                seed = gr.Slider(label='Seed',
-                                 minimum=0,
-                                 maximum=100000,
-                                 step=1,
-                                 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)
-                gr.Markdown('''
-                    - The base model must be a model that is compatible with [diffusers](https://github.com/huggingface/diffusers) library.
-                    - It takes a few minutes to download the base model first.
-                    - Expected time to train a model for 300 steps: 20 minutes with T4, 8 minutes with A10G, (4 minutes with A100)
-                    - It takes a few minutes to upload your trained model.
-                    - You may want to try a small number of steps first, like 1, to see if everything works fine in your environment.
-                    - You can check the training status by pressing the "Open logs" button if you are running this on your Space.
-                    ''')
-
+                    
         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()
@@ -111,7 +110,7 @@ def create_training_demo(trainer: Trainer,
                              training_video,
                              training_prompt,
                              output_model_name,
-                             delete_existing_model,
+                             delete_existing_repo,
                              validation_prompt,
                              base_model,
                              resolution,
@@ -128,6 +127,7 @@ def create_training_demo(trainer: Trainer,
                              delete_existing_repo,
                              upload_to,
                              remove_gpu_after_training,
+                             input_token
                          ],
                          outputs=output_message)
     return demo

app_upload.py

@@ -20,6 +20,7 @@ class ModelUploader(Uploader):
         upload_to: str,
         private: bool,
         delete_existing_repo: bool,
+        input_token: str | None = None,
     ) -> str:
         if not folder_path:
             raise ValueError
@@ -38,7 +39,8 @@ class ModelUploader(Uploader):
                            repo_name,
                            organization=organization,
                            private=private,
-                           delete_existing_repo=delete_existing_repo)
+                           delete_existing_repo=delete_existing_repo,
+                           input_token=input_token)
 
 
 def load_local_model_list() -> dict:
@@ -68,6 +70,7 @@ 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)
         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}}).
@@ -86,6 +89,7 @@ def create_upload_demo(hf_token: str | None) -> gr.Blocks:
                                 upload_to,
                                 use_private_repo,
                                 delete_existing_repo,
+                                input_token,
                             ],
                             outputs=output_message)
 

trainer.py

@@ -20,12 +20,12 @@ from utils import save_model_card
 sys.path.append('Tune-A-Video')
 
 URL_TO_JOIN_MODEL_LIBRARY_ORG = 'https://huggingface.co/organizations/Tune-A-Video-library/share/YjTcaNJmKyeHFpMBioHhzBcTzCYddVErEk'
-
+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
-        self.api = HfApi(token=hf_token)
         self.model_uploader = ModelUploader(hf_token)
 
         self.checkpoint_dir = pathlib.Path('checkpoints')
@@ -42,10 +42,10 @@ class Trainer:
                            cwd=org_dir)
         return model_dir.as_posix()
 
-    def join_model_library_org(self) -> None:
+    def join_model_library_org(self, token: str) -> None:
         subprocess.run(
             shlex.split(
-                f'curl -X POST -H "Authorization: Bearer {self.hf_token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
+                f'curl -X POST -H "Authorization: Bearer {token}" -H "Content-Type: application/json" {URL_TO_JOIN_MODEL_LIBRARY_ORG}'
             ))
 
     def run(
@@ -70,7 +70,10 @@ class Trainer:
         delete_existing_repo: bool,
         upload_to: str,
         remove_gpu_after_training: bool,
+        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')
         if not torch.cuda.is_available():
             raise gr.Error('CUDA is not available.')
         if training_video is None:
@@ -94,7 +97,7 @@ class Trainer:
         output_dir.mkdir(parents=True)
 
         if upload_to_hub:
-            self.join_model_library_org()
+            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)
@@ -143,14 +146,16 @@ class Trainer:
                 repo_name=output_model_name,
                 upload_to=upload_to,
                 private=use_private_repo,
-                delete_existing_repo=delete_existing_repo)
+                delete_existing_repo=delete_existing_repo,
+                input_token=input_token)
             print(upload_message)
             message = message + '\n' + upload_message
 
         if remove_gpu_after_training:
             space_id = os.getenv('SPACE_ID')
             if space_id:
-                self.api.request_space_hardware(repo_id=space_id,
-                                                hardware='cpu-basic')
+                api = HfApi(token=self.hf_token if self.hf_token else input_token)
+                api.request_space_hardware(repo_id=space_id,
+                                           hardware='cpu-basic')
 
         return message

uploader.py

@@ -5,10 +5,7 @@ from huggingface_hub import HfApi
 
 class Uploader:
     def __init__(self, hf_token: str | None):
-        self.api = HfApi(token=hf_token)
-
-    def get_username(self) -> str:
-        return self.api.whoami()['name']
+        self.hf_token = hf_token
 
     def upload(self,
                folder_path: str,
@@ -16,25 +13,30 @@ class Uploader:
                organization: str = '',
                repo_type: str = 'model',
                private: bool = True,
-               delete_existing_repo: bool = False) -> str:
+               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 = self.get_username()
+            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)
             except Exception:
                 pass
-        try:
-            self.api.create_repo(repo_id, repo_type=repo_type, private=private)
-            self.api.upload_folder(repo_id=repo_id,
-                                   folder_path=folder_path,
-                                   path_in_repo='.',
-                                   repo_type=repo_type)
+        try:                
+            api.create_repo(repo_id, repo_type=repo_type, private=private)
+            api.upload_folder(repo_id=repo_id,
+                              folder_path=folder_path,
+                              path_in_repo='.',
+                              repo_type=repo_type)
             url = f'https://huggingface.co/{repo_id}'
             message = f'Your model was successfully uploaded to <a href="{url}" target="_blank">{url}</a>.'
         except Exception as e: