Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import numpy as np
 import random
@@ -16,22 +17,19 @@ from decouple import config
 from pathlib import Path
 from PIL import Image
 import io
+import torch
 
-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
+URL="http://127.0.0.1"
+OUTPUT_DIR="ComfyUI/output"
+INPUT_DIR="ComfyUI/input"
+BACKEND_PATH="ComfyUI/main.py"
 
-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)
 
+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 +37,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 +51,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 +63,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)
@@ -83,12 +76,15 @@ def check_server_ready(port):
         return False
 
 
-
-@spaces.GPU(duration=190)
 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 +92,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 +104,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)
@@ -143,15 +136,102 @@ def generate_image(prompt, image):
 
 
 
+@spaces.GPU(duration=50)
+def generate_image_50(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=70)
+def generate_image_70(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=90)
+def generate_image_90(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=110)
+def generate_image_110(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=130)
+def generate_image_130(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=150)
+def generate_image_150(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=170)
+def generate_image_170(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=190) 
+def generate_image_190(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=200)
+def generate_image_200(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=210)
+def generate_image_210(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=220)
+def generate_image_220(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=230)
+def generate_image_230(prompt, image):
+  return generate_image(prompt, image)
+
+@spaces.GPU(duration=240)
+def generate_image_240(prompt, image):
+  return generate_image(prompt, image)
+
+
+def generate_image_wrapper(prompt, image, duration):
+    if duration == "50":
+        return generate_image_50(prompt, image)
+    elif duration == "70":
+        return generate_image_70(prompt, image)
+    elif duration == "90":
+        return generate_image_90(prompt, image)
+    elif duration == "110":
+        return generate_image_110(prompt, image)
+    elif duration == "130":
+        return generate_image_130(prompt, image)
+    elif duration == "150":
+        return generate_image_150(prompt, image)
+    elif duration == "170":
+        return generate_image_170(prompt, image)
+    elif duration == "190":
+        return generate_image_190(prompt, image)
+    elif duration == "200":
+        return generate_image_200(prompt, image)
+    elif duration == "210":
+        return generate_image_210(prompt, image)
+    elif duration == "220":
+        return generate_image_220(prompt, image)
+    elif duration == "230":
+        return generate_image_230(prompt, image)
+    elif duration == "240":
+        return generate_image_240(prompt, image)
+    else:
+        return generate_image_170(prompt, image)
+
 if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image, inputs=[
+    demo = gr.Interface(
+        fn=generate_image_wrapper,
+        inputs=[
         "text",
-        gr.Image(image_mode='RGBA', type="numpy")
+        gr.Image(image_mode='RGBA', type="numpy"),
+        "text"
     ],
-                        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")