server.py 9.1 KB
Newer Older
pythongosssss's avatar
pythongosssss committed
1
2
3
4
import os
import sys
import asyncio
import nodes
5
import execution
pythongosssss's avatar
pythongosssss committed
6
7
import uuid
import json
8
import glob
pythongosssss's avatar
pythongosssss committed
9
10
11
12
13
14
15
16
17
18
try:
    import aiohttp
    from aiohttp import web
except ImportError:
    print("Module 'aiohttp' not installed. Please install it via:")
    print("pip install aiohttp")
    print("or")
    print("pip install -r requirements.txt")
    sys.exit()

comfyanonymous's avatar
Style.  
comfyanonymous committed
19
import mimetypes
pythongosssss's avatar
pythongosssss committed
20

pythongosssss's avatar
pythongosssss committed
21
22
class PromptServer():
    def __init__(self, loop):
pythongosssss's avatar
pythongosssss committed
23
24
        mimetypes.init(); 
        mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8'
pythongosssss's avatar
pythongosssss committed
25
26
27
28
        self.prompt_queue = None
        self.loop = loop
        self.messages = asyncio.Queue()
        self.number = 0
pythongosssss's avatar
pythongosssss committed
29
        self.app = web.Application(client_max_size=20971520)
pythongosssss's avatar
pythongosssss committed
30
31
        self.sockets = dict()
        self.web_root = os.path.join(os.path.dirname(
pythongosssss's avatar
pythongosssss committed
32
            os.path.realpath(__file__)), "web")
pythongosssss's avatar
pythongosssss committed
33
        routes = web.RouteTableDef()
34
35
        self.last_node_id = None
        self.client_id = None
pythongosssss's avatar
pythongosssss committed
36
37
38
39
40

        @routes.get('/ws')
        async def websocket_handler(request):
            ws = web.WebSocketResponse()
            await ws.prepare(request)
41
42
43
44
45
46
47
            sid = request.rel_url.query.get('clientId', '')
            if sid:
                # Reusing existing session, remove old
                self.sockets.pop(sid, None)
            else:
                sid = uuid.uuid4().hex      

pythongosssss's avatar
pythongosssss committed
48
            self.sockets[sid] = ws
49

pythongosssss's avatar
pythongosssss committed
50
51
52
            try:
                # Send initial state to the new client
                await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid)
53
54
55
56
                # On reconnect if we are the currently executing client send the current node
                if self.client_id == sid and self.last_node_id is not None:
                    await self.send("executing", { "node": self.last_node_id }, sid)
                    
pythongosssss's avatar
pythongosssss committed
57
58
59
60
                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.ERROR:
                        print('ws connection closed with exception %s' % ws.exception())
            finally:
61
                self.sockets.pop(sid, None)
pythongosssss's avatar
pythongosssss committed
62
63
64
65
66
            return ws

        @routes.get("/")
        async def get_root(request):
            return web.FileResponse(os.path.join(self.web_root, "index.html"))
67

68
69
70
71
72
        @routes.get("/extensions")
        async def get_extensions(request):
            files = glob.glob(os.path.join(self.web_root, 'extensions/**/*.js'), recursive=True)
            return web.json_response(list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)))

pythongosssss's avatar
pythongosssss committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
        @routes.post("/upload/image")
        async def upload_image(request):
            upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "uploads")

            if not os.path.exists(upload_dir):
                os.makedirs(upload_dir)
            
            post = await request.post()
            image = post.get("image")

            if image and image.file:
                filename = image.filename
                if not filename:
                    return web.Response(status=400)

                with open(os.path.join(upload_dir, filename), "wb") as f:
                    f.write(image.file.read())
                
                return web.json_response({"name" : filename})
            else:
                return web.Response(status=400)


pythongosssss's avatar
pythongosssss committed
96
97
98
        @routes.get("/view/{file}")
        async def view_image(request):
            if "file" in request.match_info:
pythongosssss's avatar
pythongosssss committed
99
100
101
102
103
                type = request.rel_url.query.get("type", "output")
                if type != "output" and type != "uploads":
                    return web.Response(status=400)

                output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type)
pythongosssss's avatar
pythongosssss committed
104
105
106
107
108
109
110
                file = request.match_info["file"]
                file = os.path.splitext(os.path.basename(file))[0] + ".png"
                file = os.path.join(output_dir, file)
                if os.path.isfile(file):
                    return web.FileResponse(file)
                
            return web.Response(status=404)
111

pythongosssss's avatar
pythongosssss committed
112
113
114
        @routes.get("/prompt")
        async def get_prompt(request):
            return web.json_response(self.get_queue_info())
115

pythongosssss's avatar
pythongosssss committed
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
        @routes.get("/object_info")
        async def get_object_info(request):
            out = {}
            for x in nodes.NODE_CLASS_MAPPINGS:
                obj_class = nodes.NODE_CLASS_MAPPINGS[x]
                info = {}
                info['input'] = obj_class.INPUT_TYPES()
                info['output'] = obj_class.RETURN_TYPES
                info['name'] = x #TODO
                info['description'] = ''
                info['category'] = 'sd'
                if hasattr(obj_class, 'CATEGORY'):
                    info['category'] = obj_class.CATEGORY
                out[x] = info
            return web.json_response(out)
131

pythongosssss's avatar
pythongosssss committed
132
133
        @routes.get("/history")
        async def get_history(request):
134
135
            return web.json_response(self.prompt_queue.get_history())

pythongosssss's avatar
pythongosssss committed
136
137
138
139
140
141
142
        @routes.get("/queue")
        async def get_queue(request):
            queue_info = {}
            current_queue = self.prompt_queue.get_current_queue()
            queue_info['queue_running'] = current_queue[0]
            queue_info['queue_pending'] = current_queue[1]
            return web.json_response(queue_info)
143

pythongosssss's avatar
pythongosssss committed
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
        @routes.post("/prompt")
        async def post_prompt(request):
            print("got prompt")
            resp_code = 200
            out_string = ""
            json_data =  await request.json()

            if "number" in json_data:
                number = float(json_data['number'])
            else:
                number = self.number
                if "front" in json_data:
                    if json_data['front']:
                        number = -number

                self.number += 1

            if "prompt" in json_data:
                prompt = json_data["prompt"]
163
                valid = execution.validate_prompt(prompt)
pythongosssss's avatar
pythongosssss committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
                extra_data = {}
                if "extra_data" in json_data:
                    extra_data = json_data["extra_data"]

                if "client_id" in json_data:
                    extra_data["client_id"] = json_data["client_id"]
                if valid[0]:
                    self.prompt_queue.put((number, id(prompt), prompt, extra_data))
                else:
                    resp_code = 400
                    out_string = valid[1]
                    print("invalid prompt:", valid[1])

            return web.Response(body=out_string, status=resp_code)
        
        @routes.post("/queue")
        async def post_queue(request):
            json_data =  await request.json()
            if "clear" in json_data:
                if json_data["clear"]:
                    self.prompt_queue.wipe_queue()
            if "delete" in json_data:
                to_delete = json_data['delete']
                for id_to_delete in to_delete:
                    delete_func = lambda a: a[1] == int(id_to_delete)
                    self.prompt_queue.delete_queue_item(delete_func)
                    
            return web.Response(status=200)
pythongosssss's avatar
pythongosssss committed
192
193
194
195
196
197

        @routes.post("/interrupt")
        async def post_interrupt(request):
            nodes.interrupt_processing()
            return web.Response(status=200)

pythongosssss's avatar
pythongosssss committed
198
199
200
201
202
        @routes.post("/history")
        async def post_history(request):
            json_data =  await request.json()
            if "clear" in json_data:
                if json_data["clear"]:
203
                    self.prompt_queue.wipe_history()
pythongosssss's avatar
pythongosssss committed
204
205
206
            if "delete" in json_data:
                to_delete = json_data['delete']
                for id_to_delete in to_delete:
207
208
                    self.prompt_queue.delete_history_item(id_to_delete)

pythongosssss's avatar
pythongosssss committed
209
            return web.Response(status=200)
210

pythongosssss's avatar
pythongosssss committed
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        self.app.add_routes(routes)
        self.app.add_routes([
            web.static('/', self.web_root),
        ])

    def get_queue_info(self):
        prompt_info = {}
        exec_info = {}
        exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()
        prompt_info['exec_info'] = exec_info
        return prompt_info

    async def send(self, event, data, sid=None):
        message = {"type": event, "data": data}
       
        if isinstance(message, str) == False:
            message = json.dumps(message)

        if sid is None:
            for ws in self.sockets.values():
                await ws.send_str(message)
        elif sid in self.sockets:
            await self.sockets[sid].send_str(message)

    def send_sync(self, event, data, sid=None):
        self.loop.call_soon_threadsafe(
            self.messages.put_nowait, (event, data, sid))
238

pythongosssss's avatar
pythongosssss committed
239
240
241
242
243
244
245
246
    def queue_updated(self):
        self.send_sync("status", { "status": self.get_queue_info() })

    async def publish_loop(self):
        while True:
            msg = await self.messages.get()
            await self.send(*msg)

comfyanonymous's avatar
comfyanonymous committed
247
    async def start(self, address, port, verbose=True):
pythongosssss's avatar
pythongosssss committed
248
249
250
251
        runner = web.AppRunner(self.app)
        await runner.setup()
        site = web.TCPSite(runner, address, port)
        await site.start()
252

pythongosssss's avatar
pythongosssss committed
253
254
        if address == '':
            address = '0.0.0.0'
comfyanonymous's avatar
comfyanonymous committed
255
256
257
        if verbose:
            print("Starting server\n")
            print("To see the GUI go to: http://{}:{}".format(address, port))