|
|
from modules import scripts |
|
|
import re |
|
|
|
|
|
from scripts.couple_mapping import empty_tensor, basic_mapping, advanced_mapping |
|
|
from scripts.couple_ui import couple_UI, validata_mapping, parse_mapping |
|
|
|
|
|
from scripts.attention_couple import AttentionCouple |
|
|
forgeAttentionCouple = AttentionCouple() |
|
|
|
|
|
VERSION = "1.3.0" |
|
|
|
|
|
|
|
|
class ForgeCouple(scripts.Script): |
|
|
|
|
|
def __init__(self): |
|
|
self.couples: list = None |
|
|
|
|
|
def title(self): |
|
|
return "Forge Couple" |
|
|
|
|
|
def show(self, is_img2img): |
|
|
return scripts.AlwaysVisible |
|
|
|
|
|
def ui(self, *args, **kwargs): |
|
|
return couple_UI(self, f"{self.title()} {VERSION}") |
|
|
|
|
|
def parse_networks(self, prompt: str) -> str: |
|
|
"""LoRAs are already parsed""" |
|
|
pattern = re.compile(r"<.*?>") |
|
|
cleaned = re.sub(pattern, "", prompt) |
|
|
|
|
|
return cleaned |
|
|
|
|
|
def after_extra_networks_activate( |
|
|
self, |
|
|
p, |
|
|
enable: bool, |
|
|
direction: str, |
|
|
background: str, |
|
|
separator: str, |
|
|
mode: str, |
|
|
mapping: list, |
|
|
*args, |
|
|
**kwargs, |
|
|
): |
|
|
if not enable: |
|
|
return |
|
|
|
|
|
if not separator.strip(): |
|
|
separator = "\n" |
|
|
|
|
|
couples = [] |
|
|
|
|
|
chunks = kwargs["prompts"][0].split(separator) |
|
|
for chunk in chunks: |
|
|
prompt = self.parse_networks(chunk).strip() |
|
|
|
|
|
if not prompt.strip(): |
|
|
|
|
|
continue |
|
|
|
|
|
couples.append(prompt) |
|
|
|
|
|
if len(couples) < (3 if background != "None" else 2): |
|
|
print("\n\n[Couple] Not Enough Lines in Prompt...\n\n") |
|
|
self.couples = None |
|
|
return |
|
|
|
|
|
if (mode == "Advanced") and (len(couples) != len(parse_mapping(mapping))): |
|
|
print("\n\n[Couple] Number of Couples and Mapping is not the same...\n\n") |
|
|
self.couples = None |
|
|
return |
|
|
|
|
|
if (mode == "Advanced") and not validata_mapping(mapping): |
|
|
self.couples = None |
|
|
return |
|
|
|
|
|
self.couples = couples |
|
|
|
|
|
def process_before_every_sampling( |
|
|
self, |
|
|
p, |
|
|
enable: bool, |
|
|
direction: str, |
|
|
background: str, |
|
|
separator: str, |
|
|
mode: str, |
|
|
mapping: list, |
|
|
*args, |
|
|
**kwargs, |
|
|
): |
|
|
|
|
|
if not enable or not self.couples: |
|
|
return |
|
|
|
|
|
|
|
|
p.extra_generation_params["forge_couple"] = True |
|
|
p.extra_generation_params["forge_couple_separator"] = "\n" if not separator.strip() else separator |
|
|
p.extra_generation_params["forge_couple_mode"] = mode |
|
|
if mode == "Basic": |
|
|
p.extra_generation_params["forge_couple_direction"] = direction |
|
|
p.extra_generation_params["forge_couple_background"] = background |
|
|
else: |
|
|
p.extra_generation_params["forge_couple_mapping"] = mapping |
|
|
|
|
|
|
|
|
|
|
|
unet = p.sd_model.forge_objects.unet |
|
|
|
|
|
WIDTH: int = p.width |
|
|
HEIGHT: int = p.height |
|
|
IS_HORIZONTAL: bool = direction == "Horizontal" |
|
|
|
|
|
LINE_COUNT: int = len(self.couples) |
|
|
TILE_COUNT: int = LINE_COUNT - (background != "None") |
|
|
|
|
|
if mode == "Basic": |
|
|
TILE_WEIGHT: float = 1.25 if background == "None" else 1.0 |
|
|
BG_WEIGHT: float = 0.0 if background == "None" else 0.5 |
|
|
|
|
|
TILE_SIZE: int = ( |
|
|
(WIDTH if IS_HORIZONTAL else HEIGHT) - 1 |
|
|
) // TILE_COUNT + 1 |
|
|
|
|
|
|
|
|
|
|
|
if mode == "Basic": |
|
|
ARGs = basic_mapping( |
|
|
p.sd_model, |
|
|
self.couples, |
|
|
WIDTH, |
|
|
HEIGHT, |
|
|
LINE_COUNT, |
|
|
IS_HORIZONTAL, |
|
|
background, |
|
|
TILE_SIZE, |
|
|
TILE_WEIGHT, |
|
|
BG_WEIGHT, |
|
|
) |
|
|
|
|
|
else: |
|
|
ARGs = advanced_mapping(p.sd_model, self.couples, WIDTH, HEIGHT, mapping) |
|
|
|
|
|
|
|
|
assert len(ARGs.keys()) // 2 == LINE_COUNT |
|
|
|
|
|
base_mask = empty_tensor(HEIGHT, WIDTH) |
|
|
patched_unet = forgeAttentionCouple.patch_unet(unet, base_mask, ARGs) |
|
|
p.sd_model.forge_objects.unet = patched_unet |
|
|
|
