Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import random
@@ -24,7 +25,6 @@ COMF_PATH = config('COMF_PATH')
 
 import torch
 
-import spaces
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
@@ -84,7 +84,7 @@ def check_server_ready(port):
 
 
 
-@spaces.GPU(duration=170)
+
 def generate_image(prompt, image):
     prefix_filename = str(random.randint(0, 999999))
     prompt = prompt.replace('ComfyUI', prefix_filename)
@@ -143,10 +143,96 @@ def generate_image(prompt, image):
 
 
 
+
+@spaces.GPU(duration=50)
+def generate_image_50(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=70)
+def generate_image_70(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=90)
+def generate_image_90(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=110)
+def generate_image_110(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=130)
+def generate_image_130(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=150)
+def generate_image_150(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=170)
+def generate_image_170(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=190) 
+def generate_image_190(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=200)
+def generate_image_200(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=210)
+def generate_image_210(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=220)
+def generate_image_220(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=230)
+def generate_image_230(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=240)
+def generate_image_240(prompt, image):
+  return generate_image(prompt, image)
+
+
+def generate_image_wrapper(prompt, image, duration):
+    if duration == "50":
+        return generate_image_50(prompt, image)
+    elif duration == "70":
+        return generate_image_70(prompt, image)
+    elif duration == "90":
+        return generate_image_90(prompt, image)
+    elif duration == "110":
+        return generate_image_110(prompt, image)
+    elif duration == "130":
+        return generate_image_130(prompt, image)
+    elif duration == "150":
+        return generate_image_150(prompt, image)
+    elif duration == "170":
+        return generate_image_170(prompt, image)
+    elif duration == "190":
+        return generate_image_190(prompt, image)
+    elif duration == "200":
+        return generate_image_200(prompt, image)
+    elif duration == "210":
+        return generate_image_210(prompt, image)
+    elif duration == "220":
+        return generate_image_220(prompt, image)
+    elif duration == "230":
+        return generate_image_230(prompt, image)
+    elif duration == "240":
+        return generate_image_240(prompt, image)
+    else:
+        return generate_image_170(prompt, image)
+
+
 if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image, inputs=[
+    demo = gr.Interface(fn=generate_image_wrapper, inputs=[
         "text",
-        gr.Image(image_mode='RGBA', type="numpy")
+        gr.Image(image_mode='RGBA', type="numpy"),
+        "text"
     ],
                         outputs=[
                             gr.Image(type="numpy", image_mode='RGBA')