Update app.py

app.py

@@ -84,76 +84,84 @@ def check_server_ready(port):
 
 
 
-@spaces.GPU(duration=172)
-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):  # Максимум 20 секунд ожидания
-            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 = 240  # Максимальное время ожидания в секундах
-        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}")
+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()")
                 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()")
-            try:
-                logger.debug("process.wait(timeout=5)")
-                process.wait(timeout=5)
-            except subprocess.TimeoutExpired:
-                logger.debug("process.kill()")
-                process.kill()
+                    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)
 
 
 
 if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image,
+    demo = gr.Interface(fn=generate_image_with_allocation,
                         inputs=[
                             "text",
                             gr.Image(image_mode='RGBA', type="numpy"),
-                            gr.Image(image_mode='RGBA', type="numpy")
+                            gr.Image(image_mode='RGBA', type="numpy"),
+                            "text"
                         ],
                         outputs=[
                             gr.Image(type="numpy", image_mode='RGBA')