Clean up

app.py

@@ -24,6 +24,7 @@ SHARED_UI_WARNING = f'''## Attention - Training doesn't work in this shared UI.
 <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>
 '''
 
+IS_SHARED_UI = SPACE_ID == ORIGINAL_SPACE_ID
 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:
@@ -57,7 +58,7 @@ pipe = InferencePipeline(HF_TOKEN)
 trainer = Trainer(HF_TOKEN)
 
 with gr.Blocks(css='style.css') as demo:
-    if SPACE_ID == ORIGINAL_SPACE_ID:
+    if IS_SHARED_UI:
         show_warning(SHARED_UI_WARNING)
     elif not torch.cuda.is_available():
         show_warning(CUDA_NOT_AVAILABLE_WARNING)
@@ -67,7 +68,7 @@ with gr.Blocks(css='style.css') as demo:
     gr.Markdown(TITLE)
     with gr.Tabs():
         with gr.TabItem('Train'):
-            create_training_demo(trainer, pipe)
+            create_training_demo(trainer, pipe, disable_training=IS_SHARED_UI)
         with gr.TabItem('Run'):
             create_inference_demo(pipe, HF_TOKEN)
         with gr.TabItem('Upload'):

app_training.py

@@ -13,7 +13,8 @@ from trainer import Trainer
 
 
 def create_training_demo(trainer: Trainer,
-                         pipe: InferencePipeline | None = None) -> gr.Blocks:
+                         pipe: InferencePipeline | None = None,
+                         disable_training: bool = False) -> gr.Blocks:
     def read_log() -> str:
         with open(trainer.log_file) as f:
             lines = f.readlines()
@@ -111,7 +112,8 @@ def create_training_demo(trainer: Trainer,
             value=False,
             interactive=bool(os.getenv('SPACE_ID')),
             visible=False)
-        run_button = gr.Button('Start Training')
+        run_button = gr.Button('Start Training',
+                               interactive=not disable_training)
 
         with gr.Box():
             gr.Text(label='Log',
@@ -119,7 +121,8 @@ def create_training_demo(trainer: Trainer,
                     lines=10,
                     max_lines=10,
                     every=1)
-            if not os.getenv('DISABLE_SYSTEM_MONITOR'):
+            if not disable_training and not os.getenv(
+                    'DISABLE_SYSTEM_MONITOR'):
                 with gr.Accordion(label='System info', open=False):
                     create_monitor_demo()
 

trainer.py

@@ -20,8 +20,6 @@ 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')
 
 
 class Trainer:
@@ -76,10 +74,6 @@ class Trainer:
         remove_gpu_after_training: bool,
         input_token: str,
     ) -> None:
-        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: