File size: 2,619 Bytes
e0b7ec8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | 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.")
|