server.py 15.8 KB
Newer Older
pythongosssss's avatar
pythongosssss committed
1
2
3
4
import os
import sys
import asyncio
import nodes
5
import folder_paths
6
import execution
pythongosssss's avatar
pythongosssss committed
7
8
import uuid
import json
9
import glob
10
11
12
from PIL import Image
from io import BytesIO

pythongosssss's avatar
pythongosssss committed
13
14
15
16
17
18
19
20
21
22
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
23
import mimetypes
EllangoK's avatar
EllangoK committed
24
from comfy.cli_args import args
pythongosssss's avatar
pythongosssss committed
25

26
27
28
29
30
31
32
33

@web.middleware
async def cache_control(request: web.Request, handler):
    response: web.Response = await handler(request)
    if request.path.endswith('.js') or request.path.endswith('.css'):
        response.headers.setdefault('Cache-Control', 'no-cache')
    return response

EllangoK's avatar
EllangoK committed
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def create_cors_middleware(allowed_origin: str):
    @web.middleware
    async def cors_middleware(request: web.Request, handler):
        if request.method == "OPTIONS":
            # Pre-flight request. Reply successfully:
            response = web.Response()
        else:
            response = await handler(request)

        response.headers['Access-Control-Allow-Origin'] = allowed_origin
        response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS'
        response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'
        response.headers['Access-Control-Allow-Credentials'] = 'true'
        return response

    return cors_middleware
EllangoK's avatar
EllangoK committed
50

pythongosssss's avatar
pythongosssss committed
51
52
class PromptServer():
    def __init__(self, loop):
53
        PromptServer.instance = self
comfyanonymous's avatar
Style.  
comfyanonymous committed
54

pythongosssss's avatar
pythongosssss committed
55
56
        mimetypes.init(); 
        mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8'
pythongosssss's avatar
pythongosssss committed
57
58
59
60
        self.prompt_queue = None
        self.loop = loop
        self.messages = asyncio.Queue()
        self.number = 0
EllangoK's avatar
EllangoK committed
61
62

        middlewares = [cache_control]
63
64
        if args.enable_cors_header:
            middlewares.append(create_cors_middleware(args.enable_cors_header))
EllangoK's avatar
EllangoK committed
65
66

        self.app = web.Application(client_max_size=20971520, middlewares=middlewares)
pythongosssss's avatar
pythongosssss committed
67
68
        self.sockets = dict()
        self.web_root = os.path.join(os.path.dirname(
pythongosssss's avatar
pythongosssss committed
69
            os.path.realpath(__file__)), "web")
pythongosssss's avatar
pythongosssss committed
70
        routes = web.RouteTableDef()
71
        self.routes = routes
72
73
        self.last_node_id = None
        self.client_id = None
pythongosssss's avatar
pythongosssss committed
74
75
76
77
78

        @routes.get('/ws')
        async def websocket_handler(request):
            ws = web.WebSocketResponse()
            await ws.prepare(request)
79
80
81
82
83
            sid = request.rel_url.query.get('clientId', '')
            if sid:
                # Reusing existing session, remove old
                self.sockets.pop(sid, None)
            else:
84
                sid = uuid.uuid4().hex
85

pythongosssss's avatar
pythongosssss committed
86
            self.sockets[sid] = ws
87

pythongosssss's avatar
pythongosssss committed
88
89
90
            try:
                # Send initial state to the new client
                await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid)
91
92
93
94
                # 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
95
96
97
98
                async for msg in ws:
                    if msg.type == aiohttp.WSMsgType.ERROR:
                        print('ws connection closed with exception %s' % ws.exception())
            finally:
99
                self.sockets.pop(sid, None)
pythongosssss's avatar
pythongosssss committed
100
101
102
103
104
            return ws

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

106
107
        @routes.get("/embeddings")
        def get_embeddings(self):
108
            embeddings = folder_paths.get_filename_list("embeddings")
109
110
            return web.json_response(list(map(lambda a: os.path.splitext(a)[0].lower(), embeddings)))

111
112
113
114
115
        @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)))

116
117
118
119
120
121
122
123
124
125
126
127
        def get_dir_by_type(dir_type):
            if dir_type is None:
                type_dir = folder_paths.get_input_directory()
            elif dir_type == "input":
                type_dir = folder_paths.get_input_directory()
            elif dir_type == "temp":
                type_dir = folder_paths.get_temp_directory()
            elif dir_type == "output":
                type_dir = folder_paths.get_output_directory()

            return type_dir

comfyanonymous's avatar
comfyanonymous committed
128
        def image_upload(post, image_save_function=None):
ltdrdata's avatar
ltdrdata committed
129
            image = post.get("image")
130
            overwrite = post.get("overwrite")
ltdrdata's avatar
ltdrdata committed
131

comfyanonymous's avatar
comfyanonymous committed
132
133
            image_upload_type = post.get("type")
            upload_dir = get_dir_by_type(image_upload_type)
pythongosssss's avatar
pythongosssss committed
134
135
136
137
138
139

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

comfyanonymous's avatar
comfyanonymous committed
140
141
142
143
144
145
146
147
148
                subfolder = post.get("subfolder", "")
                full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder))

                if os.path.commonpath((upload_dir, os.path.abspath(full_output_folder))) != upload_dir:
                    return web.Response(status=400)

                if not os.path.exists(full_output_folder):
                    os.makedirs(full_output_folder)

pythongosssss's avatar
pythongosssss committed
149
                split = os.path.splitext(filename)
comfyanonymous's avatar
comfyanonymous committed
150
151
                filepath = os.path.join(full_output_folder, filename)

152
153
154
155
156
157
158
159
                if overwrite is not None and (overwrite == "true" or overwrite == "1"):
                    pass
                else:
                    i = 1
                    while os.path.exists(filepath):
                        filename = f"{split[0]} ({i}){split[1]}"
                        filepath = os.path.join(full_output_folder, filename)
                        i += 1
pythongosssss's avatar
pythongosssss committed
160

comfyanonymous's avatar
comfyanonymous committed
161
162
163
164
165
                if image_save_function is not None:
                    image_save_function(image, post, filepath)
                else:
                    with open(filepath, "wb") as f:
                        f.write(image.file.read())
pythongosssss's avatar
pythongosssss committed
166

comfyanonymous's avatar
comfyanonymous committed
167
                return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
pythongosssss's avatar
pythongosssss committed
168
169
170
            else:
                return web.Response(status=400)

comfyanonymous's avatar
comfyanonymous committed
171
172
173
174
175
        @routes.post("/upload/image")
        async def upload_image(request):
            post = await request.post()
            return image_upload(post)

176
177
178
179
        @routes.post("/upload/mask")
        async def upload_mask(request):
            post = await request.post()

comfyanonymous's avatar
comfyanonymous committed
180
181
            def image_save_function(image, post, filepath):
                original_pil = Image.open(post.get("original_image").file).convert('RGBA')
182
183
184
185
186
                mask_pil = Image.open(image.file).convert('RGBA')

                # alpha copy
                new_alpha = mask_pil.getchannel('A')
                original_pil.putalpha(new_alpha)
187
                original_pil.save(filepath, compress_level=4)
188

comfyanonymous's avatar
comfyanonymous committed
189
            return image_upload(post, image_save_function)
pythongosssss's avatar
pythongosssss committed
190

191
        @routes.get("/view")
pythongosssss's avatar
pythongosssss committed
192
        async def view_image(request):
m957ymj75urz's avatar
m957ymj75urz committed
193
            if "filename" in request.rel_url.query:
194
195
196
197
198
199
200
201
202
203
204
                filename = request.rel_url.query["filename"]
                filename,output_dir = folder_paths.annotated_filepath(filename)

                # validation for security: prevent accessing arbitrary path
                if filename[0] == '/' or '..' in filename:
                    return web.Response(status=400)

                if output_dir is None:
                    type = request.rel_url.query.get("type", "output")
                    output_dir = folder_paths.get_directory_by_type(type)

205
                if output_dir is None:
pythongosssss's avatar
pythongosssss committed
206
207
                    return web.Response(status=400)

208
                if "subfolder" in request.rel_url.query:
m957ymj75urz's avatar
m957ymj75urz committed
209
                    full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"])
210
                    if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:
m957ymj75urz's avatar
m957ymj75urz committed
211
212
                        return web.Response(status=403)
                    output_dir = full_output_dir
213

214
215
                filename = os.path.basename(filename)
                file = os.path.join(output_dir, filename)
m957ymj75urz's avatar
m957ymj75urz committed
216

pythongosssss's avatar
pythongosssss committed
217
                if os.path.isfile(file):
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
                    if 'channel' not in request.rel_url.query:
                        channel = 'rgba'
                    else:
                        channel = request.rel_url.query["channel"]

                    if channel == 'rgb':
                        with Image.open(file) as img:
                            if img.mode == "RGBA":
                                r, g, b, a = img.split()
                                new_img = Image.merge('RGB', (r, g, b))
                            else:
                                new_img = img.convert("RGB")

                            buffer = BytesIO()
                            new_img.save(buffer, format='PNG')
                            buffer.seek(0)

                            return web.Response(body=buffer.read(), content_type='image/png',
                                                headers={"Content-Disposition": f"filename=\"{filename}\""})

                    elif channel == 'a':
                        with Image.open(file) as img:
                            if img.mode == "RGBA":
                                _, _, _, a = img.split()
                            else:
                                a = Image.new('L', img.size, 255)

                            # alpha img
                            alpha_img = Image.new('RGBA', img.size)
                            alpha_img.putalpha(a)
                            alpha_buffer = BytesIO()
                            alpha_img.save(alpha_buffer, format='PNG')
                            alpha_buffer.seek(0)

                            return web.Response(body=alpha_buffer.read(), content_type='image/png',
                                                headers={"Content-Disposition": f"filename=\"{filename}\""})
                    else:
                        return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})

pythongosssss's avatar
pythongosssss committed
257
            return web.Response(status=404)
258

pythongosssss's avatar
pythongosssss committed
259
260
261
        @routes.get("/prompt")
        async def get_prompt(request):
            return web.json_response(self.get_queue_info())
262

pythongosssss's avatar
pythongosssss committed
263
264
265
266
267
268
269
270
        @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
271
                info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)
272
                info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']
City's avatar
City committed
273
274
                info['name'] = x
                info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[x] if x in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else x
pythongosssss's avatar
pythongosssss committed
275
276
277
278
279
280
                info['description'] = ''
                info['category'] = 'sd'
                if hasattr(obj_class, 'CATEGORY'):
                    info['category'] = obj_class.CATEGORY
                out[x] = info
            return web.json_response(out)
281

pythongosssss's avatar
pythongosssss committed
282
283
        @routes.get("/history")
        async def get_history(request):
284
285
            return web.json_response(self.prompt_queue.get_history())

pythongosssss's avatar
pythongosssss committed
286
287
288
289
290
291
292
        @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)
293

pythongosssss's avatar
pythongosssss committed
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
        @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"]
313
                valid = execution.validate_prompt(prompt)
pythongosssss's avatar
pythongosssss committed
314
315
316
317
318
319
320
                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]:
321
322
323
                    prompt_id = str(uuid.uuid4())
                    self.prompt_queue.put((number, prompt_id, prompt, extra_data, valid[2]))
                    return web.json_response({"prompt_id": prompt_id})
pythongosssss's avatar
pythongosssss committed
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
                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:
comfyanonymous's avatar
comfyanonymous committed
340
                    delete_func = lambda a: a[1] == id_to_delete
pythongosssss's avatar
pythongosssss committed
341
                    self.prompt_queue.delete_queue_item(delete_func)
comfyanonymous's avatar
comfyanonymous committed
342

pythongosssss's avatar
pythongosssss committed
343
            return web.Response(status=200)
pythongosssss's avatar
pythongosssss committed
344
345
346
347
348
349

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

pythongosssss's avatar
pythongosssss committed
350
351
352
353
354
        @routes.post("/history")
        async def post_history(request):
            json_data =  await request.json()
            if "clear" in json_data:
                if json_data["clear"]:
355
                    self.prompt_queue.wipe_history()
pythongosssss's avatar
pythongosssss committed
356
357
358
            if "delete" in json_data:
                to_delete = json_data['delete']
                for id_to_delete in to_delete:
359
360
                    self.prompt_queue.delete_history_item(id_to_delete)

pythongosssss's avatar
pythongosssss committed
361
            return web.Response(status=200)
362
363
364
        
    def add_routes(self):
        self.app.add_routes(self.routes)
pythongosssss's avatar
pythongosssss committed
365
        self.app.add_routes([
366
            web.static('/', self.web_root, follow_symlinks=True),
pythongosssss's avatar
pythongosssss committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
        ])

    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))
391

pythongosssss's avatar
pythongosssss committed
392
393
394
395
396
397
398
399
    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)

400
    async def start(self, address, port, verbose=True, call_on_start=None):
pythongosssss's avatar
pythongosssss committed
401
402
403
404
        runner = web.AppRunner(self.app)
        await runner.setup()
        site = web.TCPSite(runner, address, port)
        await site.start()
405

pythongosssss's avatar
pythongosssss committed
406
407
        if address == '':
            address = '0.0.0.0'
comfyanonymous's avatar
comfyanonymous committed
408
409
410
        if verbose:
            print("Starting server\n")
            print("To see the GUI go to: http://{}:{}".format(address, port))
411
412
413
        if call_on_start is not None:
            call_on_start(address, port)