Spaces:
Build error
Build error
| import spaces | |
| import gradio as gr | |
| import numpy as np | |
| import random | |
| import multiprocessing | |
| import subprocess | |
| import sys | |
| import time | |
| import signal | |
| import json | |
| import os | |
| import requests | |
| from loguru import logger | |
| from decouple import config | |
| from pathlib import Path | |
| from PIL import Image | |
| import io | |
| import torch | |
| URL="http://127.0.0.1" | |
| OUTPUT_DIR="Core/output" | |
| INPUT_DIR="Core/input" | |
| BACKEND_PATH="Core/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): | |
| initial_size = os.path.getsize(file_path) | |
| 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'))] | |
| 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(3) | |
| return latest_image | |
| return None | |
| 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, 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) | |
| return response.status_code == 200 | |
| except requests.RequestException: | |
| return False | |
| def generate_image(image): | |
| prefix_filename = str(random.randint(0, 999999)) | |
| 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') | |
| process = None | |
| new_port = str(random.randint(8123, 8200)) | |
| 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): | |
| 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}") | |
| try: | |
| return Image.open(latest_image) | |
| finally: | |
| delete_image_file(latest_image) | |
| delete_image_file(INPUT_DIR + '/input.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() | |
| if __name__ == "__main__": | |
| demo = gr.Interface( | |
| fn=generate_image, | |
| inputs=[gr.Image(image_mode='RGBA', type="numpy")], | |
| 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("finish") |