Tune GPU duration based on model

app.py

@@ -10,11 +10,20 @@ models = [
     "stabilityai/stable-diffusion-xl-refiner-1.0",
     "timbrooks/instruct-pix2pix",
 ]
-default_model = "stabilityai/stable-diffusion-xl-refiner-1.0"
+DEFAULT_MODEL = "stabilityai/stable-diffusion-xl-refiner-1.0"
+
+
+@spaces.GPU
+def gpu(fn):
+    return fn()
 
 
-@logger.catch(reraise=True)
 @spaces.GPU(duration=180)
+def gpu_3min(fn):
+    return fn()
+
+
+@logger.catch(reraise=True)
 def generate(
     model: str,
     prompt: str,
@@ -49,12 +58,17 @@ def generate(
     additional_args = (
         {} if model == "timbrooks/instruct-pix2pix" else dict(strength=strength)
     )
-    images = pipe(
-        prompt=prompt,
-        image=init_image,
-        callback_on_step_end=progress_callback,
-        **additional_args,
-    ).images
+
+    gpu_runner = gpu_3min if model == "timbrooks/instruct-pix2pix" else gpu
+
+    images = gpu_runner(
+        lambda: pipe(
+            prompt=prompt,
+            image=init_image,
+            callback_on_step_end=progress_callback,
+            **additional_args,
+        ).images
+    )
     return images[0]
 
 
@@ -62,7 +76,7 @@ demo = gr.Interface(
     fn=generate,
     inputs=[
         gr.Dropdown(
-            label="Model", choices=models, value=default_model, allow_custom_value=True
+            label="Model", choices=models, value=DEFAULT_MODEL, allow_custom_value=True
         ),
         gr.Text(label="Prompt"),
         gr.Image(label="Init image", type="pil"),