add from_image_json caption plugin to add all metadata from image json to prompt

This commit is contained in:
Victor Hall 2024-06-21 15:27:19 -04:00
parent de651dc6fb
commit b1b73838f0
1 changed files with 31 additions and 7 deletions

View File

@ -228,14 +228,40 @@ class TitleAndTagsFromImageJson(PromptIdentityBase):
logging.debug(f" {self.key}: prompt after: {prompt}")
return prompt
class VogueRunwayImageJson(PromptIdentityBase):
class AllMetadataFromImageJson(PromptIdentityBase):
def __init__(self, args:Namespace=None):
super().__init__(key="vogue_runway_from_image_json",
super().__init__(key="from_image_json",
description="Adds title and tags hint from metadata.json (in the samefolder as the image) to the prompt",
fn=self._title_and_tags_from_metadata_json,
args=args)
def try_get_kvps(self, metadata, keys:list):
def _title_and_tags_from_metadata_json(self, args:Namespace) -> str:
prompt = args.prompt
logging.debug(f" {self.key}: prompt before: {prompt}")
image_path = args.image_path
current_dir = os.path.dirname(image_path)
image_path_base = os.path.basename(image_path)
image_path_without_extension = os.path.splitext(image_path_base)[0]
candidate_json_path = os.path.join(current_dir, f"{image_path_without_extension}.json")
if os.path.exists(candidate_json_path):
with open(candidate_json_path, "r") as f:
metadata = json.load(f)
hint = json.dumps(metadata)
prompt = self._add_hint_to_prompt(hint, prompt)
logging.debug(f" {self.key}: prompt after: {prompt}")
return prompt
class VogueRunwayImageJson(PromptIdentityBase):
def __init__(self, args:Namespace=None):
super().__init__(key="vogue_runway_from_image_json",
description="Adds designer, season, category, year, and tags from [image_path_without_extension].json to the prompt",
fn=self._title_and_tags_from_metadata_json,
args=args)
def _try_get_kvps(self, metadata, keys:list):
values = []
for key in keys:
val = metadata.get(key, "")
@ -261,10 +287,8 @@ class VogueRunwayImageJson(PromptIdentityBase):
with open(candidate_json_path, "r") as f:
metadata = json.load(f)
keys = ["designer","season","category","year"]
hint = ""
hint = self.try_get_kvps(metadata, keys)
hint = self._try_get_kvps(metadata, ["designer","season","category","year"])
tags = metadata.get("tags", [])
tags = tags.split(",") if isinstance(tags, str) else tags # can be csv or list