cymic/Waifu_Diffusion_Webui

1# We need this so Python doesn't complain about the unknown StableDiffusionProcessing-typehint at runtime2from __future__ import annotations3 4import csv5import os6import os.path7import typing8import collections.abc as abc9import tempfile10import shutil11 12if typing.TYPE_CHECKING:13    # Only import this when code is being type-checked, it doesn't have any effect at runtime14    from .processing import StableDiffusionProcessing15 16 17class PromptStyle(typing.NamedTuple):18    name: str19    prompt: str20    negative_prompt: str21 22 23def merge_prompts(style_prompt: str, prompt: str) -> str:24    if "{prompt}" in style_prompt:25        res = style_prompt.replace("{prompt}", prompt)26    else:27        parts = filter(None, (prompt.strip(), style_prompt.strip()))28        res = ", ".join(parts)29 30    return res31 32 33def apply_styles_to_prompt(prompt, styles):34    for style in styles:35        prompt = merge_prompts(style, prompt)36 37    return prompt38 39 40class StyleDatabase:41    def __init__(self, path: str):42        self.no_style = PromptStyle("None", "", "")43        self.styles = {"None": self.no_style}44 45        if not os.path.exists(path):46            return47 48        with open(path, "r", encoding="utf8", newline='') as file:49            reader = csv.DictReader(file)50            for row in reader:51                # Support loading old CSV format with "name, text"-columns52                prompt = row["prompt"] if "prompt" in row else row["text"]53                negative_prompt = row.get("negative_prompt", "")54                self.styles[row["name"]] = PromptStyle(row["name"], prompt, negative_prompt)55 56    def get_style_prompts(self, styles):57        return [self.styles.get(x, self.no_style).prompt for x in styles]58 59    def get_negative_style_prompts(self, styles):60        return [self.styles.get(x, self.no_style).negative_prompt for x in styles]61 62    def apply_styles_to_prompt(self, prompt, styles):63        return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).prompt for x in styles])64 65    def apply_negative_styles_to_prompt(self, prompt, styles):66        return apply_styles_to_prompt(prompt, [self.styles.get(x, self.no_style).negative_prompt for x in styles])67 68    def apply_styles(self, p: StableDiffusionProcessing) -> None:69        if isinstance(p.prompt, list):70            p.prompt = [self.apply_styles_to_prompt(prompt, p.styles) for prompt in p.prompt]71        else:72            p.prompt = self.apply_styles_to_prompt(p.prompt, p.styles)73 74        if isinstance(p.negative_prompt, list):75            p.negative_prompt = [self.apply_negative_styles_to_prompt(prompt, p.styles) for prompt in p.negative_prompt]76        else:77            p.negative_prompt = self.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)78 79    def save_styles(self, path: str) -> None:80        # Write to temporary file first, so we don't nuke the file if something goes wrong81        fd, temp_path = tempfile.mkstemp(".csv")82        with os.fdopen(fd, "w", encoding="utf8", newline='') as file:83            # _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,84            # and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()85            writer = csv.DictWriter(file, fieldnames=PromptStyle._fields)86            writer.writeheader()87            writer.writerows(style._asdict() for k,     style in self.styles.items())88 89        # Always keep a backup file around90        if os.path.exists(path):91            shutil.move(path, path + ".bak")92        shutil.move(temp_path, path)93