atomium-backend-python/checker.py at master · markjkaem/atomium-backend-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

from openai_client import openai_client
import requests
import os
from dotenv import load_dotenv
import time
from restart_program import restart_program

load_dotenv()

client = openai_client()

def count_images_in_folder(folder_path):
     if not os.path.exists(folder_path):
        return 0

     # Define a set of valid image extensions
     image_extensions = {'.jpg', '.jpeg', '.png', '.webp'}

     # Use a list comprehension to count files with valid image extensions
     image_count = sum(1 for filename in os.listdir(folder_path)
                      if os.path.splitext(filename)[1].lower() in image_extensions)

     return image_count

def check_openai():
    engines = client.models.list()
    if engines:
        # If the call is successful, print the engine list
        return "LIVE"

    else:
    # If there's an error, print the error message
        return "OFFLINE"

def check_prodia():

    url = "https://api.prodia.com/v1/sd/models"

    headers = {
        "accept": "application/json",
        "X-Prodia-Key": "1ac454f7-67d7-4950-823f-6b4f8369d975"
    }

    response = requests.get(url, headers=headers)

    if response.text:
        # If the call is successful, print the engine list
        return "LIVE"

    else:
    # If there's an error, print the error message
        return "OFFLINE"

def check_imagefolders(channel):

    amount = count_images_in_folder("channels/" + channel + "/images/")
    # Check if the number of images is greater than or equal to 30
    if amount >= 30:
       return True, amount
    else:
       return False, amount

def check_imagefolders_short(channel):

    amount = count_images_in_folder("channels/" + channel + "/images-short/")
    # Check if the number of images is greater than or equal to 30
    if amount >= 30:
       return True, amount
    else:
       return False, amount


def middleware(channel):
    isMissing, amount = check_imagefolders(channel)
    print(isMissing, amount)
    if isMissing == False:
        print("\n" + channel.capitalize().replace("_", " ") + "has not enough images: \033[91m" + str(amount) + "\033[0m images.\n")
        time.sleep(1)
        print("\n\033[93mPlease wait... Closing application...\033[0m\n")
        exit()
    else:
        print("")


def checker():
    print("Checking endpoints and folders...")
    openai_status = check_openai()
    if openai_status == "LIVE":
        print("Openai: " + "\033[32m" + openai_status + "\033[0m")
    else:
        print("Openai: " + "\033[91m" + openai_status + "\033[0m")
    prodia_status = check_prodia()
    if prodia_status == "LIVE":
        print("Prodia: " + "\033[32m" + prodia_status + "\033[0m")
    else:
        print("Prodia: " + "\033[91m" + prodia_status + "\033[0m")
    folder1_status, folder1_amount = check_imagefolders("spiritual_mutation")
    if folder1_status == True:
        print("SpiritualMutation: " + "\033[32m" + str(folder1_amount) + "\033[0m" + " images (min. 30)")
    else:
        print("SpiritualMutation: " + "\033[91m" + str(folder1_amount) + "\033[0m" + " images (min. 30)" )

    folder2_status, folder2_amount = check_imagefolders("the_universe_decoded")
    if folder2_status == True:
        print("TheUniverseDecoded: " + "\033[32m" + str(folder2_amount) + "\033[0m" + " images (min. 30)")
    else:
        print("TheUniverseDecoded: " + "\033[91m" + str(folder2_amount) + "\033[0m" + " images (min. 30)")

    folder1_status, folder1_amount = check_imagefolders_short("spiritual_mutation")
    if folder1_status == True:
        print("SpiritualMutation: " + "\033[32m" + str(folder1_amount) + "\033[0m" + " short images (min. 10)")
    else:
        print("SpiritualMutation: " + "\033[91m" + str(folder1_amount) + "\033[0m" + " short images (min. 10)" )

    folder2_status, folder2_amount = check_imagefolders_short("the_universe_decoded")
    if folder2_status == True:
        print("TheUniverseDecoded: " + "\033[32m" + str(folder2_amount) + "\033[0m" + " short images (min. 10)")
    else:
        print("TheUniverseDecoded: " + "\033[91m" + str(folder2_amount) + "\033[0m" + " short images (min. 10)")
    time.sleep(2)