Update app.py

app.py

@@ -9,6 +9,7 @@ import logging
 import re
 import torch
 import gradio as gr
+import spaces
 
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
 
@@ -47,7 +48,8 @@ class PromptEnhancerV2:
         )
         self.processor = AutoProcessor.from_pretrained(models_root_path)
 
-    @torch.inference_mode()
+    # @torch.inference_mode()
+    @spaces.GPU
     def predict(
         self,
         prompt_cot,
@@ -111,7 +113,8 @@ class PromptEnhancerV2:
             print(f"✗ Re-prompting failed, so we are using the original prompt. Error: {e}")
 
         return prompt_cot
-    @torch.inference_mode()
+    # @torch.inference_mode()
+    @spaces.GPU
     def predict_stream(
         self,
         prompt_cot,