Update app.py

app.py

@@ -85,84 +85,109 @@ def check_server_ready(port):
 
 
 
-def generate_image_with_allocation(prompt, image, image2, allocation_time):
-    if allocation_time is None:
-        allocation_time = 170
-    else:
-        allocation_time = int(allocation_time)
-
-    @spaces.GPU(duration=allocation_time)
-    def generate_image(prompt, image, image2):
-        prefix_filename = str(random.randint(0, 999999))
-        prompt = prompt.replace('ComfyUI', prefix_filename)
-        prompt = json.loads(prompt)
-
-        image = Image.fromarray(image)
-        image.save(INPUT_DIR + '/input.png', format='PNG')
-        if image2 is not None:
-            image2 = Image.fromarray(image2)
-            image2.save(INPUT_DIR + '/input2.png', format='PNG')
-
-        process = None
-        new_port = str(random.randint(8123, 8200))
-
-        try:
-            # Запускаем скрипт как подпроцесс
-            process = subprocess.Popen([sys.executable, COMF_PATH, "--listen", "127.0.0.1", "--port", new_port])
-            logger.debug(f'Subprocess started with PID: {process.pid}')
-
-            # Ожидание запуска сервера
-            for _ in range(30):  # Максимум 30 секунд ожидания
-                if check_server_ready(new_port):
-                    break
-                time.sleep(1)
-            else:
-                raise TimeoutError("Server did not start in time")
-
-            start_queue(prompt, new_port)
-
-            # Ожидание нового изображения
-            timeout = 360  # Максимальное время ожидания в секундах
-            start_time = time.time()
-            while time.time() - start_time < timeout:
-                latest_image = wait_for_image_with_prefix(OUTPUT_DIR, prefix_filename)
-                if latest_image:
-                    logger.debug(f"file is: {latest_image}")
-                    try:
-                        return Image.open(latest_image)
-                    finally:
-                        delete_image_file(latest_image)
-                        delete_image_file(INPUT_DIR + '/input.png')
-                        if image2 is not None:
-                            delete_image_file(INPUT_DIR + '/input2.png')
-                time.sleep(1)
-
-            raise TimeoutError("New image was not generated in time")
-
-        except Exception as e:
-            logger.error(f"Error in generate_image: {e}")
-
-        finally:
-            if process and process.poll() is None:
-                process.terminate()
-                logger.debug("process.terminate()")
+def generate_image(prompt, image, image2):
+    prefix_filename = str(random.randint(0, 999999))
+    prompt = prompt.replace('ComfyUI', prefix_filename)
+    prompt = json.loads(prompt)
+
+    image = Image.fromarray(image)
+    image.save(INPUT_DIR + '/input.png', format='PNG')
+    if image2 is not None:
+        image2 = Image.fromarray(image2)
+        image2.save(INPUT_DIR + '/input2.png', format='PNG')
+
+    process = None
+    new_port = str(random.randint(8123, 8200))
+
+    try:
+        # Запускаем скрипт как подпроцесс
+        process = subprocess.Popen([sys.executable, COMF_PATH, "--listen", "127.0.0.1", "--port", new_port])
+        logger.debug(f'Subprocess started with PID: {process.pid}')
+
+        # Ожидание запуска сервера
+        for _ in range(30):  # Максимум 30 секунд ожидания
+            if check_server_ready(new_port):
+                break
+            time.sleep(1)
+        else:
+            raise TimeoutError("Server did not start in time")
+
+        start_queue(prompt, new_port)
+
+        # Ожидание нового изображения
+        timeout = 360  # Максимальное время ожидания в секундах
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            latest_image = wait_for_image_with_prefix(OUTPUT_DIR, prefix_filename)
+            if latest_image:
+                logger.debug(f"file is: {latest_image}")
                 try:
-                    logger.debug("process.wait(timeout=5)")
-                    process.wait(timeout=5)
-                except subprocess.TimeoutExpired:
-                    logger.debug("process.kill()")
-                    process.kill()
-    return generate_image(prompt, image, image2)
+                    return Image.open(latest_image)
+                finally:
+                    delete_image_file(latest_image)
+                    delete_image_file(INPUT_DIR + '/input.png')
+                    if image2 is not None:
+                        delete_image_file(INPUT_DIR + '/input2.png')
+            time.sleep(1)
 
+        raise TimeoutError("New image was not generated in time")
 
+    except Exception as e:
+        logger.error(f"Error in generate_image: {e}")
+
+    finally:
+        if process and process.poll() is None:
+            process.terminate()
+            logger.debug("process.terminate()")
+            try:
+                logger.debug("process.wait(timeout=5)")
+                process.wait(timeout=5)
+            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_with_allocation,
+    demo = gr.Interface(fn=generate_image_wrapper,
                         inputs=[
                             "text",
                             gr.Image(image_mode='RGBA', type="numpy"),
                             gr.Image(image_mode='RGBA', type="numpy"),
-                            "text"
+                            gr.Radio(["130", "150", "170", "190", "200"], label="GPU acclocation in secs")
                         ],
                         outputs=[
                             gr.Image(type="numpy", image_mode='RGBA')