Update app.py

app.py

@@ -29,6 +29,12 @@ def load_models():
     
     return pipe
 
+@spaces.GPU(duration=60)
+def prompt_enhance(client, image_caption, text_caption):
+    combined_caption, enhanced_caption = client.predict(image_caption, text_caption)
+    return combined_caption, enhanced_caption
+    
+
 pipe = load_models()
 
 def truncate_caption_by_tokens(caption, max_tokens=256):
@@ -80,7 +86,7 @@ def run_pipeline(image_caption, text_caption, seed, num_inference_steps, guidanc
     
     if enable_enhancer:
         # combined_caption, enhanced_caption = generate_enhanced_caption(image_caption, text_caption)
-        combined_caption, enhanced_caption = client.predict(image_caption, text_caption)
+        combined_caption, enhanced_caption = prompt_enhance(client, image_caption, text_caption)
     else:
         enhanced_caption = combined_caption