Update app.py

app.py

@@ -1,6 +1,6 @@
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
+import spaces
 import gradio as gr
 import torch
 import base64
@@ -89,6 +89,7 @@ class MolmoCaptionProcessor:
             device_map='auto'
         )
 
+    @spaces.GPU(duration=120)
     def generate_response(self, image: Image.Image, msg: str) -> str:
         inputs = self.processor.process(
             images=[image],
@@ -218,6 +219,8 @@ def update_prompt(images: List[Image.Image], task_type: str, custom_msg: str = N
     except Exception as e:
         return "", f"Error generating captions: {str(e)}"
 
+
+@spaces.GPU(duration=120)
 def generate_image(images: List[Image.Image], prompt: str, negative_prompt: str, num_inference_steps: int, guidance_scale: float, 
                    denoise_mask: List[str], task_type: str, azimuth: str, elevation: str, distance: str, focal_length: float,
                    height: int = 1024, width: int = 1024, scale_factor: float = 1.0, scale_watershed: float = 1.0,