From 87f3c8a01ccee327283feb429e5e5077c3000c1e Mon Sep 17 00:00:00 2001 From: Anatoly Belikov Date: Tue, 9 Apr 2024 11:08:29 +0300 Subject: [PATCH] wrap queue processing in try except --- multigen/worker.py | 66 +++++++++++++++++++++++++--------------------- 1 file changed, 36 insertions(+), 30 deletions(-) diff --git a/multigen/worker.py b/multigen/worker.py index b29932f..7a82526 100755 --- a/multigen/worker.py +++ b/multigen/worker.py @@ -38,36 +38,42 @@ def _update(sess, job, gs): while not self._stop: while self.queue: # keep the job in the queue until complete - with self._lock: - data = self.queue[-1] - sess = self.sessions[data["session_id"]] - self.logger.info("GENERATING: " + str(data)) - if 'start_callback' in data: - data['start_callback']() + try: + with self._lock: + data = self.queue[-1] + sess = self.sessions[data["session_id"]] + self.logger.info("GENERATING: " + str(data)) + if 'start_callback' in data: + data['start_callback']() - pipe_name = sess.get('pipe', 'Prompt2ImPipe') - model_id = str(self.cwd/self.config["model_dir"]/self.models['base'][sess["model"]]['id']) - is_xl = self.models['base'][sess["model"]]['xl'] - pipe = self.get_pipeline(pipe_name, model_id, cnet=data.get('cnet', None), xl=is_xl) - class_name = str(pipe.__class__) - self.logger.debug(f'got pipeline {class_name}') + pipe_name = sess.get('pipe', 'Prompt2ImPipe') + model_id = str(self.cwd/self.config["model_dir"]/self.models['base'][sess["model"]]['id']) + is_xl = self.models['base'][sess["model"]]['xl'] + pipe = self.get_pipeline(pipe_name, model_id, cnet=data.get('cnet', None), xl=is_xl) + class_name = str(pipe.__class__) + self.logger.debug(f'got pipeline {class_name}') - images = data['images'] - if 'MaskedIm2ImPipe' in class_name: - pipe.setup(**data, original_image=str(images[0]), - image_painted=str(images[1])) - elif any([x in class_name for x in ('Im2ImPipe', 'Cond2ImPipe')]): - pipe.setup(**data, fimage=str(images[0])) - else: - pipe.setup(**data) - # TODO: add negative prompt to parameters - nprompt = "jpeg artifacts, blur, distortion, watermark, signature, extra fingers, fewer fingers, lowres, nude, bad hands, duplicate heads, bad anatomy, bad crop" - gs = GenSession(self.get_image_pathname(data["session_id"], None), - pipe, Cfgen(data["prompt"], nprompt)) - gs.gen_sess(add_count = data["count"], - callback = lambda: _update(sess, data, gs)) - if 'finish_callback' in data: - data['finish_callback']() - with self._lock: - self.queue.pop() + images = data['images'] + if 'MaskedIm2ImPipe' in class_name: + pipe.setup(**data, original_image=str(images[0]), + image_painted=str(images[1])) + elif any([x in class_name for x in ('Im2ImPipe', 'Cond2ImPipe')]): + pipe.setup(**data, fimage=str(images[0])) + else: + pipe.setup(**data) + # TODO: add negative prompt to parameters + nprompt = "jpeg artifacts, blur, distortion, watermark, signature, extra fingers, fewer fingers, lowres, nude, bad hands, duplicate heads, bad anatomy, bad crop" + gs = GenSession(self.get_image_pathname(data["session_id"], None), + pipe, Cfgen(data["prompt"], nprompt)) + gs.gen_sess(add_count = data["count"], + callback = lambda: _update(sess, data, gs)) + if 'finish_callback' in data: + data['finish_callback']() + except RuntimeError as e: + if 'finish_callback' in data: + self.logger.error("error in generation", exc_info=e) + data['finish_callback']("Can't generate image due to error") + finally: + with self._lock: + self.queue.pop() time.sleep(0.2)