Update gradio_app.py

gradio_app.py

@@ -2,10 +2,12 @@ import gradio as gr
 from test_lora import DanbooruTagTester
 import sys
 import io
+import spaces
 
 # Gradio's state management will hold the instance of our tester
 # This is better than a global variable as it's session-specific
 
+@spaces.GPU(duration=300) # Request GPU for model loading, with a 5-min timeout
 def load_model(model_path, base_model, use_4bit, progress=gr.Progress(track_tqdm=True)):
     """
     Loads the model and updates the UI.
@@ -41,7 +43,7 @@ def load_model(model_path, base_model, use_4bit, progress=gr.Progress(track_tqdm
     # Return the loaded model instance, the status message, and UI updates
     return tester, final_status, gr.update(interactive=success), gr.update(interactive=success), gr.update(interactive=success), gr.update(interactive=success), gr.update(interactive=success), gr.update(interactive=success)
 
-
+@spaces.GPU # Request GPU for generation
 def generate_tags(tester, prompt, max_new_tokens, temperature, top_k, top_p, do_sample):
     """
     Generates tags using the loaded model.