| 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 |
|
|
