Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import random
@@ -16,22 +17,18 @@ from decouple import config
 from pathlib import Path
 from PIL import Image
 import io
-
-URL="http://127.0.0.1"
-OUTPUT_DIR = config('OUTPUT_DIR')
-INPUT_DIR = config('INPUT_DIR')
-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())}")
-print(torch.version.cuda)
-device = torch.cuda.get_device_name(torch.cuda.current_device())
-print(device)
+URL="http://127.0.0.1"
+OUTPUT_DIR="ComfyUI/output"
+INPUT_DIR="ComfyUI/input"
+BACKEND_PATH="ComfyUI/main.py"
+
 
+def read_prompt_from_file(filename):
+    with open(filename, 'r', encoding='utf-8') as file:
+        return json.load(file)
 
 def wait_for_image_with_prefix(folder, prefix):
     def is_file_ready(file_path):
@@ -39,7 +36,6 @@ def wait_for_image_with_prefix(folder, prefix):
         time.sleep(1)
         return initial_size == os.path.getsize(file_path)
 
-    
     files = os.listdir(folder)
     image_files = [f for f in files if f.lower().startswith(prefix.lower()) and
                    f.lower().endswith(('.png', '.jpg', '.jpeg'))]
@@ -54,10 +50,8 @@ def wait_for_image_with_prefix(folder, prefix):
             time.sleep(3)
             return latest_image
 
-    
     return None
 
-
 def delete_image_file(file_path):
     try:
         if os.path.exists(file_path):
@@ -68,13 +62,11 @@ def delete_image_file(file_path):
     except Exception as e:
         logger.debug(f"error {file_path}: {str(e)}")
 
-
 def start_queue(prompt_workflow, port):
     p = {"prompt": prompt_workflow}
     data = json.dumps(p).encode('utf-8')
     requests.post(f"{URL}:{port}/prompt", data=data)
 
-
 def check_server_ready(port):
     try:
         response = requests.get(f"{URL}:{port}/history/123", timeout=5)
@@ -82,13 +74,16 @@ def check_server_ready(port):
     except requests.RequestException:
         return False
 
-
-
-@spaces.GPU(duration=175)
+@spaces.GPU(duration=170)
 def generate_image(prompt, image):
     prefix_filename = str(random.randint(0, 999999))
-    prompt = prompt.replace('ComfyUI', prefix_filename)
-    prompt = json.loads(prompt)
+    try:
+        prompt = prompt.replace('ComfyUI', prefix_filename)
+        prompt = json.loads(prompt)
+    except:
+        prompt = read_prompt_from_file('config.json')
+        prompt = json.dumps(prompt, ensure_ascii=False).replace('output_image', prefix_filename)
+        prompt = json.loads(prompt)
 
     image = Image.fromarray(image)
     image.save(INPUT_DIR + '/input.png', format='PNG')
@@ -96,13 +91,11 @@ def generate_image(prompt, image):
     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])
+    try:        
+        process = subprocess.Popen([sys.executable, BACKEND_PATH, "--listen", "127.0.0.1", "--port", new_port])
         logger.debug(f'Subprocess started with PID: {process.pid}')
-
-        # Ожидание запуска сервера
-        for _ in range(40):  # Максимум 20 секунд ожидания
+      
+        for _ in range(40): 
             if check_server_ready(new_port):
                 break
             time.sleep(1)
@@ -110,9 +103,8 @@ def generate_image(prompt, image):
             raise TimeoutError("Server did not start in time")
 
         start_queue(prompt, new_port)
-
-        # Ожидание нового изображения
-        timeout = 240  # Максимальное время ожидания в секундах
+        
+        timeout = 240
         start_time = time.time()
         while time.time() - start_time < timeout:
             latest_image = wait_for_image_with_prefix(OUTPUT_DIR, prefix_filename)
@@ -141,17 +133,18 @@ def generate_image(prompt, image):
                 logger.debug("process.kill()")
                 process.kill()
 
-
-
 if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image, inputs=[
+    demo = gr.Interface(
+        fn=generate_image,
+        inputs=[
         "text",
         gr.Image(image_mode='RGBA', type="numpy")
     ],
-                        outputs=[
-                            gr.Image(type="numpy", image_mode='RGBA')
-                        ])
+        outputs=[gr.Image(type="numpy", image_mode='RGBA')],
+        title="Image Upscaler",
+        description="BEST UPSCALER EVER!!!!!"
+    )
     demo.launch(debug=True)
     logger.debug('demo.launch()')
 
-    logger.info("Основной скрипт завершил работу.")
+    logger.info("finish")