SpongeML/main.py at main · spongedsc/SpongeML · 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
import socketio
import eventlet
import psutil
import textgenwui
import requests
from requests.exceptions import Timeout
import os

from transformers import pipeline
from apiCAI import send_message, new_chat

# create the Socket.IO server
sio = socketio.Server()
app = socketio.WSGIApp(sio)


@sio.event
def imgcaption(sid, data):
    try:
        if psutil.virtual_memory().available / 1048576 < 3000:
            print("Image recognition skipped due to low memory")
            return "", 200
        image_to_text = pipeline(
            "image-to-text", model="nlpconnect/vit-gpt2-image-captioning"
        )
        caption = image_to_text(data)
        return caption, 200
    except Exception as e:
        print(f"Error occurred during image recognition: {e}")
        return "", 500


@sio.event
def chat(sid, data: dict):
    try:
        if data.get("textgenwui") == True:
            msg = textgenwui.send_message(
                data["message"],
                os.getenv("TEXTGENWUI_CHARACTER"),
            )
        else:
            msg = send_message(send_message(data["message"]))
        return msg, 200
    except Exception as e:
        print(f"Error occurred during chat: {e}")
        return "", 500


@sio.event
def imagerecognitionenabled(sid, data):
    try:
        if psutil.virtual_memory().available / 1048576 < 3000:
            return False, 200
        else:
            return True, 200
    except Exception as e:
        print(f"Error occurred during chat: {e}")
        return "", 500


@sio.event
def localgenenabled(sid, data):
    try:
        response = requests.get(
            os.getenv("TEXTGENWUI_ENDPOINT").split("/v1/")[0], timeout=1
        )
        if response.status_code == 404:
            return True
        else:
            return False
    except Timeout:
        print("The request timed out")
        return False
    except Exception as e:
        print(f"An error occurred: {e}")
        return False


@sio.event
def newchat(sid, data):
    textgenwui.reset_history()
    return new_chat(), 200


@sio.event
def connect(sid, environ, auth):
    print("connect", sid)


@sio.event
def disconnect(sid):
    print("disconnect", sid)


if __name__ == "__main__":
    eventlet.wsgi.server(eventlet.listen(("", 6000)), app)