Spaces:
Runtime error
Runtime error
File size: 4,491 Bytes
c6919c4 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
# We need this so Python doesn't complain about the unknown StableDiffusionProcessing-typehint at runtime
from __future__ import annotations
import csv
import os
import os.path
import typing
import collections.abc as abc
import tempfile
import shutil
#if typing.TYPE_CHECKING:
# # Only import this when code is being type-checked, it doesn't have any effect at runtime
# from .processing import StableDiffusionProcessing
class PromptTransformation(typing.NamedTuple):
name: str
regex: str
replacement: str
flags: str
long_description: str
def merge_prompts(transformation_prompt: str, prompt: str) -> str:
if "{prompt}" in transformation_prompt:
res = transformation_prompt.replace("{prompt}", prompt)
else:
parts = filter(None, (prompt.strip(), transformation_prompt.strip()))
res = ", ".join(parts)
return res
import re
def apply_rule_to_prompt(rule, text):
regex = rule.regex
replacement = rule.replacement
flags = 0
if 'MULTILINE' in rule.flags:
flags |= re.MULTILINE
return re.sub(regex, replacement, text, flags=flags)
def apply_transformations_to_prompt(prompt, transformations):
for transformation in transformations:
#print(f"Applying transformation {transformation.name}, regex: {transformation.regex}, replacement: {transformation.replacement}, flags: {transformation.flags}")
#print(f"Before: {prompt}")
prompt = apply_rule_to_prompt(transformation, prompt)
return prompt
class TransformationDatabase:
def __init__(self, path: str, user_path: str):
self.no_transformation = PromptTransformation("None", "", "", "","")
self.transformations = {}
self.path = path
self.user_path = user_path
#print(path)
self.reload()
def reload(self):
self.transformations.clear()
if not os.path.exists(self.path):
print(f"Can't find transformations at {self.path}")
else:
with open(self.path, "r", encoding="utf-8-sig", newline='') as file:
reader = csv.DictReader(file)
for row in reader:
#print(f"row: {row}")
# Support loading old CSV format with "name, text"-columns
regex = row["regex"]
replacement = row.get("replacement", "")
flags = row.get("flags", "")
long_description = row.get("long_description", "")
self.transformations[row["name"]] = PromptTransformation(row["name"], regex, replacement, flags, long_description)
if not os.path.exists(self.user_path):
print(f"Can't find transformations at {self.user_path}")
else:
with open(self.user_path, "r", encoding="utf-8-sig", newline='') as file:
reader = csv.DictReader(file)
for row in reader:
#print(f"row: {row}")
# Support loading old CSV format with "name, text"-columns
regex = row["regex"]
replacement = row.get("replacement", "")
flags = row.get("flags", "")
self.transformations[row["name"]] = PromptTransformation(row["name"], regex, replacement, flags, long_description)
#def get_transformation_prompts(self, transformations):
# return [self.transformations.get(x, self.no_transformation).prompt for x in transformations]
def apply_transformations_to_prompt(self, prompt, transformations):
return apply_transformations_to_prompt(prompt, [self.transformations.get(x, self.no_transformation) for x in transformations])
def save_transformations(self, path: str) -> None:
# Always keep a backup file around
if os.path.exists(path):
shutil.copy(path, path + ".bak")
fd = os.open(path, os.O_RDWR|os.O_CREAT)
with os.fdopen(fd, "w", encoding="utf-8-sig", newline='') as file:
# _fields is actually part of the public API: typing.NamedTuple is a replacement for collections.NamedTuple,
# and collections.NamedTuple has explicit documentation for accessing _fields. Same goes for _asdict()
writer = csv.DictWriter(file, fieldnames=PromptTransformation._fields)
writer.writeheader()
writer.writerows(transformation._asdict() for k, transformation in self.transformations.items())
|