Update app.py

app.py

@@ -3,6 +3,7 @@ import spaces
 from image_gen_aux import UpscaleWithModel
 from image_gen_aux.utils import load_image
 from fastapi.middleware.cors import CORSMiddleware
+import math
 
 MODELS = {
     "4xNomosWebPhotoRealPLKSR": "Phips/4xNomosWebPhoto_RealPLKSR",
@@ -12,7 +13,18 @@ MODELS = {
 }
 
 
-@spaces.GPU
+def get_duration(
+    image, model_selection
+):
+    height, width, _ = image.shape
+    megapixel = math.ceil((width * height) / 1_000_000)
+
+    if model_selection in ["4xNomosWebPhotoRealPLKSR", "4xRealESRGAN"]:
+        return int(megapixel * 10)
+    else:
+        return int(megapixel * 30)
+
+@spaces.GPU(duration = get_duration)
 def upscale_image(image, model_selection):
     original = load_image(image)