import os import threading import traceback from aiohttp import web import impact import folder_paths import torchvision import impact.core as core import impact.impact_pack as impact_pack from impact.utils import to_tensor from segment_anything import SamPredictor, sam_model_registry import numpy as np import nodes from PIL import Image import io import impact.wildcards as wildcards import comfy from io import BytesIO import random from server import PromptServer import logging sam_predictor = None default_sam_model_name = os.path.join(impact_pack.model_path, "sams", "sam_vit_b_01ec64.pth") sam_lock = threading.Condition() last_prepare_data = None def async_prepare_sam(image_dir, model_name, filename): with sam_lock: global sam_predictor if 'vit_h' in model_name: model_kind = 'vit_h' elif 'vit_l' in model_name: model_kind = 'vit_l' else: model_kind = 'vit_b' sam_model = sam_model_registry[model_kind](checkpoint=model_name) sam_predictor = SamPredictor(sam_model) image_path = os.path.join(image_dir, filename) image = nodes.LoadImage().load_image(image_path)[0] image = np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8) if impact.config.get_config()['sam_editor_cpu']: device = 'cpu' else: device = comfy.model_management.get_torch_device() sam_predictor.model.to(device=device) sam_predictor.set_image(image, "RGB") sam_predictor.model.cpu() @PromptServer.instance.routes.post("/sam/prepare") async def sam_prepare(request): global sam_predictor global last_prepare_data data = await request.json() with sam_lock: if last_prepare_data is not None and last_prepare_data == data: # already loaded: skip -- prevent redundant loading return web.Response(status=200) last_prepare_data = data model_name = 'sam_vit_b_01ec64.pth' if data['sam_model_name'] == 'auto': model_name = impact.config.get_config()['sam_editor_model'] model_path = folder_paths.get_full_path("sams", model_name) if model_path is None: logging.error(f"[Impact Pack] The '{model_name}' model file cannot be found in any sams model path.") return web.Response(status=400) logging.info(f"[Impact Pack] Loading SAM model '{model_path}'") filename, image_dir = folder_paths.annotated_filepath(data["filename"]) if image_dir is None: typ = data['type'] if data['type'] != '' else 'output' image_dir = folder_paths.get_directory_by_type(typ) if data['subfolder'] is not None and data['subfolder'] != '': image_dir += f"/{data['subfolder']}" if image_dir is None: return web.Response(status=400) thread = threading.Thread(target=async_prepare_sam, args=(image_dir, model_path, filename,)) thread.start() logging.info("[Impact Pack] SAM model loaded. ") return web.Response(status=200) @PromptServer.instance.routes.post("/sam/release") async def release_sam(request): global sam_predictor with sam_lock: del sam_predictor sam_predictor = None logging.info("[Impact Pack]: unloading SAM model") @PromptServer.instance.routes.post("/sam/detect") async def sam_detect(request): global sam_predictor with sam_lock: if sam_predictor is not None: if impact.config.get_config()['sam_editor_cpu']: device = 'cpu' else: device = comfy.model_management.get_torch_device() sam_predictor.model.to(device=device) try: data = await request.json() positive_points = data['positive_points'] negative_points = data['negative_points'] threshold = data['threshold'] points = [] plabs = [] for p in positive_points: points.append(p) plabs.append(1) for p in negative_points: points.append(p) plabs.append(0) detected_masks = core.sam_predict(sam_predictor, points, plabs, None, threshold) mask = core.combine_masks2(detected_masks) if mask is None: return web.Response(status=400) image = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) i = 255. * image.cpu().numpy() img = Image.fromarray(np.clip(i[0], 0, 255).astype(np.uint8)) img_buffer = io.BytesIO() img.save(img_buffer, format='png') headers = {'Content-Type': 'image/png'} finally: sam_predictor.model.to(device="cpu") return web.Response(body=img_buffer.getvalue(), headers=headers) else: return web.Response(status=400) @PromptServer.instance.routes.get("/impact/wildcards/refresh") async def wildcards_refresh(request): impact.wildcards.wildcard_load() return web.Response(status=200) @PromptServer.instance.routes.get("/impact/wildcards/list") async def wildcards_list(request): data = {'data': impact.wildcards.get_wildcard_list()} return web.json_response(data) @PromptServer.instance.routes.post("/impact/wildcards") async def populate_wildcards(request): data = await request.json() populated = wildcards.process(data['text'], data.get('seed', None)) return web.json_response({"text": populated}) segs_picker_map = {} @PromptServer.instance.routes.get("/impact/segs/picker/count") async def segs_picker_count(request): node_id = request.rel_url.query.get('id', '') if node_id in segs_picker_map: res = len(segs_picker_map[node_id]) return web.Response(status=200, text=str(res)) return web.Response(status=400) @PromptServer.instance.routes.get("/impact/segs/picker/view") async def segs_picker(request): node_id = request.rel_url.query.get('id', '') idx = int(request.rel_url.query.get('idx', '')) if node_id in segs_picker_map and idx < len(segs_picker_map[node_id]): img = to_tensor(segs_picker_map[node_id][idx]).permute(0, 3, 1, 2).squeeze(0) pil = torchvision.transforms.ToPILImage('RGB')(img) image_bytes = BytesIO() pil.save(image_bytes, format="PNG") image_bytes.seek(0) return web.Response(status=200, body=image_bytes, content_type='image/png', headers={"Content-Disposition": f"filename={node_id}{idx}.png"}) return web.Response(status=400) @PromptServer.instance.routes.get("/view/validate") async def view_validate(request): if "filename" in request.rel_url.query: filename = request.rel_url.query["filename"] subfolder = request.rel_url.query["subfolder"] filename, base_dir = folder_paths.annotated_filepath(filename) if filename == '' or filename[0] == '/' or '..' in filename: return web.Response(status=400) if base_dir is None: base_dir = folder_paths.get_input_directory() file = os.path.join(base_dir, subfolder, filename) if os.path.isfile(file): return web.Response(status=200) return web.Response(status=400) @PromptServer.instance.routes.get("/impact/validate/pb_id_image") async def view_validate(request): if "id" in request.rel_url.query: pb_id = request.rel_url.query["id"] if pb_id not in core.preview_bridge_image_id_map: return web.Response(status=400) file = core.preview_bridge_image_id_map[pb_id] if os.path.isfile(file): return web.Response(status=200) return web.Response(status=400) @PromptServer.instance.routes.get("/impact/set/pb_id_image") async def set_previewbridge_image(request): try: if "filename" in request.rel_url.query: node_id = request.rel_url.query["node_id"] filename = request.rel_url.query["filename"] path_type = request.rel_url.query["type"] subfolder = request.rel_url.query["subfolder"] filename, output_dir = folder_paths.annotated_filepath(filename) if filename == '' or filename[0] == '/' or '..' in filename: return web.Response(status=400) if output_dir is None: if path_type == 'input': output_dir = folder_paths.get_input_directory() elif path_type == 'output': output_dir = folder_paths.get_output_directory() else: output_dir = folder_paths.get_temp_directory() file = os.path.join(output_dir, subfolder, filename) item = { 'filename': filename, 'type': path_type, 'subfolder': subfolder, } pb_id = core.set_previewbridge_image(node_id, file, item) return web.Response(status=200, text=pb_id) except Exception: traceback.print_exc() return web.Response(status=400) @PromptServer.instance.routes.get("/impact/get/pb_id_image") async def get_previewbridge_image(request): if "id" in request.rel_url.query: pb_id = request.rel_url.query["id"] if pb_id in core.preview_bridge_image_id_map: _, path_item = core.preview_bridge_image_id_map[pb_id] return web.json_response(path_item) return web.Response(status=400) @PromptServer.instance.routes.get("/impact/view/pb_id_image") async def view_previewbridge_image(request): if "id" in request.rel_url.query: pb_id = request.rel_url.query["id"] if pb_id in core.preview_bridge_image_id_map: file = core.preview_bridge_image_id_map[pb_id] with Image.open(file) as img: filename = os.path.basename(file) return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""}) return web.Response(status=400) def onprompt_for_switch(json_data): inversed_switch_info = {} onprompt_switch_info = {} onprompt_cond_branch_info = {} disabled_switch = set() for k, v in json_data['prompt'].items(): if 'class_type' not in v: continue cls = v['class_type'] if cls == 'ImpactInversedSwitch': # if 'sel_mode' is 'select_on_prompt' if 'sel_mode' in v['inputs'] and v['inputs']['sel_mode'] and 'select' in v['inputs']: select_input = v['inputs']['select'] # if 'select' is converted input if isinstance(select_input, list) and len(select_input) == 2: input_node = json_data['prompt'][select_input[0]] if input_node['class_type'] == 'ImpactInt' and 'inputs' in input_node and 'value' in input_node['inputs']: inversed_switch_info[k] = input_node['inputs']['value'] else: logging.warning(f"\n##### ##### #####\n[Impact Pack] {cls}: For the 'select' operation, only 'select_index' of the 'ImpactInversedSwitch', which is not an input, or 'ImpactInt' and 'Primitive' are allowed as inputs if 'select_on_prompt' is selected.\n##### ##### #####\n") else: inversed_switch_info[k] = select_input elif cls in ['ImpactSwitch', 'LatentSwitch', 'SEGSSwitch', 'ImpactMakeImageList']: # if 'sel_mode' is 'select_on_prompt' if 'sel_mode' in v['inputs'] and v['inputs']['sel_mode'] and 'select' in v['inputs']: select_input = v['inputs']['select'] # if 'select' is converted input if isinstance(select_input, list) and len(select_input) == 2: input_node = json_data['prompt'][select_input[0]] if input_node['class_type'] == 'ImpactInt' and 'inputs' in input_node and 'value' in input_node['inputs']: onprompt_switch_info[k] = input_node['inputs']['value'] if input_node['class_type'] == 'ImpactSwitch' and 'inputs' in input_node and 'select' in input_node['inputs']: if isinstance(input_node['inputs']['select'], int): onprompt_switch_info[k] = input_node['inputs']['select'] else: logging.warning(f"\n##### ##### #####\n[Impact Pack] {cls}: For the 'select' operation, only 'select_index' of the 'ImpactSwitch', which is not an input, or 'ImpactInt' and 'Primitive' are allowed as inputs if 'select_on_prompt' is selected.\n##### ##### #####\n") else: onprompt_switch_info[k] = select_input if k in onprompt_switch_info and f'input{onprompt_switch_info[k]}' not in v['inputs']: # disconnect output disabled_switch.add(k) elif cls == 'ImpactConditionalBranchSelMode': if 'sel_mode' in v['inputs'] and v['inputs']['sel_mode'] and 'cond' in v['inputs']: cond_input = v['inputs']['cond'] if isinstance(cond_input, list) and len(cond_input) == 2: input_node = json_data['prompt'][cond_input[0]] if (input_node['class_type'] == 'ImpactValueReceiver' and 'inputs' in input_node and 'value' in input_node['inputs'] and 'typ' in input_node['inputs']): if 'BOOLEAN' == input_node['inputs']['typ']: try: onprompt_cond_branch_info[k] = input_node['inputs']['value'].lower() == "true" except: pass else: onprompt_cond_branch_info[k] = cond_input for k, v in json_data['prompt'].items(): disable_targets = set() for kk, vv in v['inputs'].items(): if isinstance(vv, list) and len(vv) == 2: if vv[0] in inversed_switch_info: if vv[1] + 1 != inversed_switch_info[vv[0]]: disable_targets.add(kk) else: del inversed_switch_info[k] if vv[0] in disabled_switch: disable_targets.add(kk) if k in onprompt_switch_info: selected_slot_name = f"input{onprompt_switch_info[k]}" for kk, vv in v['inputs'].items(): if kk != selected_slot_name and kk.startswith('input'): disable_targets.add(kk) if k in onprompt_cond_branch_info: selected_slot_name = "tt_value" if onprompt_cond_branch_info[k] else "ff_value" for kk, vv in v['inputs'].items(): if kk in ['tt_value', 'ff_value'] and kk != selected_slot_name: disable_targets.add(kk) for kk in disable_targets: del v['inputs'][kk] # inversed_switch - select out of range for target in inversed_switch_info.keys(): del json_data['prompt'][target]['inputs']['input'] def onprompt_for_pickers(json_data): detected_pickers = set() for k, v in json_data['prompt'].items(): if 'class_type' not in v: continue cls = v['class_type'] if cls == 'ImpactSEGSPicker': detected_pickers.add(k) # garbage collection keys_to_remove = [key for key in segs_picker_map if key not in detected_pickers] for key in keys_to_remove: del segs_picker_map[key] def gc_preview_bridge_cache(json_data): prompt_keys = json_data['prompt'].keys() for key in list(core.preview_bridge_cache.keys()): if key not in prompt_keys: # print(f"key deleted [PB]: {key}") del core.preview_bridge_cache[key] for key in list(core.preview_bridge_last_mask_cache.keys()): if key not in prompt_keys: # print(f"key deleted [PB_last_mask]: {key}") del core.preview_bridge_last_mask_cache[key] def workflow_imagereceiver_update(json_data): prompt = json_data['prompt'] for v in prompt.values(): if 'class_type' in v and v['class_type'] == 'ImageReceiver': if v['inputs']['save_to_workflow']: v['inputs']['image'] = "#DATA" def regional_sampler_seed_update(json_data): prompt = json_data['prompt'] for k, v in prompt.items(): if 'class_type' in v and v['class_type'] == 'RegionalSampler': seed_2nd_mode = v['inputs']['seed_2nd_mode'] new_seed = None if seed_2nd_mode == 'increment': new_seed = v['inputs']['seed_2nd']+1 if new_seed > 1125899906842624: new_seed = 0 elif seed_2nd_mode == 'decrement': new_seed = v['inputs']['seed_2nd']-1 if new_seed < 0: new_seed = 1125899906842624 elif seed_2nd_mode == 'randomize': new_seed = random.randint(0, 1125899906842624) if new_seed is not None: PromptServer.instance.send_sync("impact-node-feedback", {"node_id": k, "widget_name": "seed_2nd", "type": "INT", "value": new_seed}) def onprompt_populate_wildcards(json_data): prompt = json_data['prompt'] updated_widget_values = {} for k, v in prompt.items(): if 'class_type' in v and (v['class_type'] == 'ImpactWildcardEncode' or v['class_type'] == 'ImpactWildcardProcessor'): inputs = v['inputs'] # legacy adapter if isinstance(inputs['mode'], bool): if inputs['mode']: new_mode = 'populate' else: new_mode = 'fixed' inputs['mode'] = new_mode if inputs['mode'] == 'populate' and isinstance(inputs['populated_text'], str): if isinstance(inputs['seed'], list): try: input_node = prompt[inputs['seed'][0]] if input_node['class_type'] == 'ImpactInt': input_seed = int(input_node['inputs']['value']) if not isinstance(input_seed, int): continue if input_node['class_type'] == 'Seed (rgthree)': input_seed = int(input_node['inputs']['seed']) if not isinstance(input_seed, int): continue else: logging.info(f"[Impact Pack] Only `ImpactInt`, `Seed (rgthree)` and `Primitive` Node are allowed as the seed for '{v['class_type']}'. It will be ignored. ") continue except: continue else: input_seed = int(inputs['seed']) inputs['populated_text'] = wildcards.process(inputs['wildcard_text'], input_seed) inputs['mode'] = 'reproduce' PromptServer.instance.send_sync("impact-node-feedback", {"node_id": k, "widget_name": "populated_text", "type": "STRING", "value": inputs['populated_text']}) updated_widget_values[k] = inputs['populated_text'] if inputs['mode'] == 'reproduce': PromptServer.instance.send_sync("impact-node-feedback", {"node_id": k, "widget_name": "mode", "type": "STRING", "value": 'populate'}) if 'extra_data' in json_data and 'extra_pnginfo' in json_data['extra_data']: for node in json_data['extra_data']['extra_pnginfo']['workflow']['nodes']: key = str(node['id']) if key in updated_widget_values: node['widgets_values'][1] = updated_widget_values[key] node['widgets_values'][2] = 'reproduce' def onprompt_for_remote(json_data): prompt = json_data['prompt'] for v in prompt.values(): if 'class_type' in v: cls = v['class_type'] if cls == 'ImpactRemoteBoolean' or cls == 'ImpactRemoteInt': inputs = v['inputs'] node_id = str(inputs['node_id']) if node_id not in prompt: continue target_inputs = prompt[node_id]['inputs'] widget_name = inputs['widget_name'] if widget_name in target_inputs: widget_type = None if cls == 'ImpactRemoteBoolean' and isinstance(target_inputs[widget_name], bool): widget_type = 'BOOLEAN' elif cls == 'ImpactRemoteInt' and (isinstance(target_inputs[widget_name], int) or isinstance(target_inputs[widget_name], float)): widget_type = 'INT' if widget_type is None: break target_inputs[widget_name] = inputs['value'] PromptServer.instance.send_sync("impact-node-feedback", {"node_id": node_id, "widget_name": widget_name, "type": widget_type, "value": inputs['value']}) def onprompt(json_data): try: onprompt_for_remote(json_data) # NOTE: top priority onprompt_for_switch(json_data) onprompt_for_pickers(json_data) onprompt_populate_wildcards(json_data) gc_preview_bridge_cache(json_data) workflow_imagereceiver_update(json_data) regional_sampler_seed_update(json_data) core.current_prompt = json_data except Exception as e: logging.warning(f"[Impact Pack] ComfyUI-Impact-Pack: Error on prompt - several features will not work.\n{e}") return json_data PromptServer.instance.add_on_prompt_handler(onprompt)