Update app.py

app.py

@@ -135,92 +135,17 @@ 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)
+def get_duration(prompt, image, duration):
+    return int(duration)
 
-@spaces.GPU(duration=230)
-def generate_image_230(prompt, image):
+@spaces.GPU(duration=get_duration)
+def generate_image(prompt, image, duration):
   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_wrapper,
+        fn=generate_image,
         inputs=[
         "text",
         gr.Image(image_mode='RGBA', type="numpy"),