| import os | |
| def inject(assembler, chain_definition, chain_items): | |
| if not chain_items: | |
| return | |
| valid_images = [] | |
| for item in chain_items: | |
| if not item: | |
| continue | |
| img_path = item | |
| if isinstance(item, dict): | |
| img_path = item.get('image') or item.get('filename') or item.get('path') | |
| if img_path: | |
| valid_images.append(img_path) | |
| if not valid_images: | |
| return | |
| valid_images = valid_images[:2] | |
| pos_prompt_name = chain_definition.get('pos_prompt_node', 'pos_prompt') | |
| neg_prompt_name = chain_definition.get('neg_prompt_node', 'neg_prompt') | |
| vae_node_name = chain_definition.get('vae_node', 'vae_loader') | |
| if pos_prompt_name not in assembler.node_map: | |
| print(f"Warning: Positive prompt node '{pos_prompt_name}' not found for JoyAI Reference chain. Skipping.") | |
| return | |
| if vae_node_name not in assembler.node_map: | |
| print(f"Warning: VAE loader node '{vae_node_name}' not found for JoyAI Reference chain. Skipping.") | |
| return | |
| pos_prompt_id = assembler.node_map[pos_prompt_name] | |
| neg_prompt_id = assembler.node_map.get(neg_prompt_name) | |
| vae_node_id = assembler.node_map[vae_node_name] | |
| assembler.workflow[pos_prompt_id]['inputs']['vae'] = [vae_node_id, 0] | |
| if neg_prompt_id and neg_prompt_id in assembler.workflow: | |
| assembler.workflow[neg_prompt_id]['inputs']['vae'] = [vae_node_id, 0] | |
| for i, img_filename in enumerate(valid_images): | |
| load_id = assembler._get_unique_id() | |
| load_node = assembler._get_node_template("LoadImage") | |
| load_node['inputs']['image'] = img_filename | |
| load_node['_meta']['title'] = f"Load Reference Image {i+1}" | |
| assembler.workflow[load_id] = load_node | |
| scale_id = assembler._get_unique_id() | |
| scale_node = assembler._get_node_template("ImageScaleToTotalPixels") | |
| scale_node['inputs']['megapixels'] = 1.0 | |
| scale_node['inputs']['upscale_method'] = "nearest-exact" | |
| scale_node['inputs']['resolution_steps'] = 1 | |
| scale_node['inputs']['image'] = [load_id, 0] | |
| scale_node['_meta']['title'] = f"Scale Reference {i+1}" | |
| assembler.workflow[scale_id] = scale_node | |
| input_key = f"images.image{i}" | |
| assembler.workflow[pos_prompt_id]['inputs'][input_key] = [scale_id, 0] | |
| if neg_prompt_id and neg_prompt_id in assembler.workflow: | |
| assembler.workflow[neg_prompt_id]['inputs'][input_key] = [scale_id, 0] | |
| print(f"JoyAI Reference injector applied. Injected {len(valid_images)} reference images to JoyAI text encoding nodes.") | |