Update app.py

app.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+from subprocess import getoutput
 
 import gradio as gr
 import torch
@@ -17,7 +18,8 @@ TITLE = '# [Tune-A-Video](https://tuneavideo.github.io/) Training 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 - This Space 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>
 '''
@@ -26,20 +28,25 @@ 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.
 </center>
 '''
 
+
+
 HF_TOKEN = os.getenv('HF_TOKEN')
 
 
@@ -56,6 +63,8 @@ trainer = Trainer(HF_TOKEN)
 with gr.Blocks(css='style.css') as demo:
     if SPACE_ID != ORIGINAL_SPACE_ID:
         show_warning(SHARED_UI_WARNING)
+    elif(not "T4" in GPU_DATA):
+        show_warning(INVALID_GPU_WARNING)
     elif not torch.cuda.is_available():
         show_warning(CUDA_NOT_AVAILABLE_WARNING)
 
@@ -74,4 +83,4 @@ with gr.Blocks(css='style.css') as demo:
     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)