Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import random
@@ -24,7 +25,7 @@ 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 +85,6 @@ def check_server_ready(port):
 
 
 
-@spaces.GPU(duration=170)
 def generate_image(prompt, image, image2):
     prefix_filename = str(random.randint(0, 999999))
     prompt = prompt.replace('ComfyUI', prefix_filename)
@@ -105,7 +105,7 @@ def generate_image(prompt, image, image2):
         logger.debug(f'Subprocess started with PID: {process.pid}')
 
         # Ожидание запуска сервера
-        for _ in range(30):  # Максимум 20 секунд ожидания
+        for _ in range(30):  # Максимум 30 секунд ожидания
             if check_server_ready(new_port):
                 break
             time.sleep(1)
@@ -115,7 +115,7 @@ def generate_image(prompt, image, image2):
         start_queue(prompt, new_port)
 
         # Ожидание нового изображения
-        timeout = 240  # Максимальное время ожидания в секундах
+        timeout = 360  # Максимальное время ожидания в секундах
         start_time = time.time()
         while time.time() - start_time < timeout:
             latest_image = wait_for_image_with_prefix(OUTPUT_DIR, prefix_filename)
@@ -135,7 +135,7 @@ def generate_image(prompt, image, image2):
     except Exception as e:
         logger.error(f"Error in generate_image: {e}")
 
-    finally:        
+    finally:
         if process and process.poll() is None:
             process.terminate()
             logger.debug("process.terminate()")
@@ -145,15 +145,49 @@ def generate_image(prompt, image, image2):
             except subprocess.TimeoutExpired:
                 logger.debug("process.kill()")
                 process.kill()
+return generate_image(prompt, image, image2)
 
+@spaces.GPU(duration=130)
+def generate_image_130(prompt, image, image2):
+  return generate_image(prompt, image, image2)
 
+@spaces.GPU(duration=150)
+def generate_image_150(prompt, image, image2):
+  return generate_image(prompt, image, image2)
+
+@spaces.GPU(duration=170)
+def generate_image_170(prompt, image, image2):
+  return generate_image(prompt, image, image2)
+
+@spaces.GPU(duration=190) 
+def generate_image_190(prompt, image, image2):
+  return generate_image(prompt, image, image2)
+
+@spaces.GPU(duration=200)
+def generate_image_200(prompt, image, image2):
+  return generate_image(prompt, image, image2)
+
+def generate_image_wrapper(prompt, image, image2, duration):
+  if duration == "130":
+    return generate_image_130(prompt, image, image2)
+  elif duration == "150":
+    return generate_image_150(prompt, image, image2)
+  elif duration == "170":
+    return generate_image_170(prompt, image, image2)
+  elif duration == "190":
+    return generate_image_190(prompt, image, image2)
+  elif duration == "200":
+    return generate_image_200(prompt, image, image2)
+  else:
+    return generate_image_170(prompt, image, image2)
 
 if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image,
+    demo = gr.Interface(fn=generate_image_wrapper,
                         inputs=[
                             "text",
                             gr.Image(image_mode='RGBA', type="numpy"),
-                            gr.Image(image_mode='RGBA', type="numpy")
+                            gr.Image(image_mode='RGBA', type="numpy"),
+                            "text", label="GPU acclocation in secs")
                         ],
                         outputs=[
                             gr.Image(type="numpy", image_mode='RGBA')