Update app.py

app.py

@@ -33,21 +33,39 @@ device = torch.cuda.get_device_name(torch.cuda.current_device())
 print(device)
 
 
-def get_latest_image(folder):
-    files = os.listdir(folder)
-    image_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
-    image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder, x)))
-    latest_image = os.path.join(folder, image_files[-1]) if image_files else None
+def wait_for_image_with_prefix(folder, prefix):
     def is_file_ready(file_path):
         initial_size = os.path.getsize(file_path)
         time.sleep(0.5)
         return initial_size == os.path.getsize(file_path)
 
-    if latest_image:
-        if is_file_ready(latest_image):
-            return latest_image
-    return latest_image
+    while True:
+        files = os.listdir(folder)
+        image_files = [f for f in files if f.lower().startswith(prefix.lower()) and 
+                       f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        
+        if image_files:
+            # Sort by modification time to get the latest file
+            image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder, x)), reverse=True)
+            latest_image = os.path.join(folder, image_files[0])
+            
+            if is_file_ready(latest_image):
+                # Wait a bit more to ensure the file is completely written
+                time.sleep(1)
+                return latest_image
+        
+        # If no matching file found, wait before checking again
+        time.sleep(1)
 
+def delete_image_file(file_path):
+    try:
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            logger.debug(f"file {file_path} deleted")
+        else:
+            logger.debug(f"file {file_path} is not exist")
+    except Exception as e:
+        logger.debug(f"error {file_path}: {str(e)}")
 
 def start_queue(prompt_workflow):
     p = {"prompt": prompt_workflow}
@@ -55,9 +73,9 @@ def start_queue(prompt_workflow):
     requests.post(URL, data=data)
 
 
-def check_server_ready():
+def check_server_ready(new_port):
     try:
-        response = requests.get(f"http://127.0.0.1:8188/history/123", timeout=5)
+        response = requests.get(f"http://127.0.0.1:{new_port}/history/123", timeout=5)
         return response.status_code == 200
     except requests.RequestException:
         return False
@@ -67,20 +85,21 @@ def check_server_ready():
 def generate_image(prompt, image):
     prompt = json.loads(prompt)
     
-    previous_image = get_latest_image(OUTPUT_DIR)
+    previous_image = None
 
     image = Image.fromarray(image)
     image.save(INPUT_DIR+'/input.png', format='PNG')
 
-
+    new_port = random.randint(8188, 8288)
+    prefix_filename = random.randint(0, 999999)
     # Запускаем скрипт как подпроцесс
-    process = subprocess.Popen([sys.executable, COMF_PATH, "--listen", "127.0.0.1"])
+    process = subprocess.Popen([sys.executable, COMF_PATH, "--listen", "127.0.0.1", "--port", f"{new_port}"])
     logger.debug(f'Subprocess started with PID: {process.pid}')
     latest_image = None
     try:
         # Ожидание запуска сервера
         for _ in range(20):  # Максимум 20 секунд ожидания
-            if check_server_ready():  # Вам нужно реализовать эту функцию
+            if check_server_ready(new_port):  # Вам нужно реализовать эту функцию
                 break
             time.sleep(1)
         else:
@@ -92,9 +111,14 @@ def generate_image(prompt, image):
         timeout = 220  # Максимальное время ожидания в секундах
         start_time = time.time()
         while time.time() - start_time < timeout:
-            latest_image = get_latest_image(OUTPUT_DIR)
+            latest_image = wait_for_image_with_prefix(OUTPUT_DIR, prefix_filename)
             if latest_image != previous_image:
-                return latest_image
+                logger.debug(f"file is: {latest_image}")
+                with open(latest_image, 'rb') as f:
+                    photo = file.read()
+                logger.debug(f"file bytes size: {len(photo)}")
+                delete_image_file(latest_image)
+                return io.BytesIO(photo)
             time.sleep(1)
 
         raise TimeoutError("New image was not generated in time")
@@ -109,8 +133,8 @@ def generate_image(prompt, image):
             try:
                 process.wait(timeout=5)
             except subprocess.TimeoutExpired:
-                process.kill()
-
+                process.kill()        
+        
         #logger.error("No new image was generated")
         return latest_image