joy: artist-from-folder flag
Browse files
joy
CHANGED
@@ -319,7 +319,6 @@ class JoyCaptionModel:
|
|
319 |
Tuple[str, float]: The generated caption and its entropy.
|
320 |
"""
|
321 |
torch.cuda.empty_cache()
|
322 |
-
logging.info(f"Prompt: {prompt_str}")
|
323 |
|
324 |
pixel_values = self._preprocess_image(input_image)
|
325 |
|
@@ -687,6 +686,11 @@ def main():
|
|
687 |
"Optionally specify the number of tags to use."
|
688 |
),
|
689 |
)
|
|
|
|
|
|
|
|
|
|
|
690 |
parser.add_argument(
|
691 |
"--random-tags",
|
692 |
type=int,
|
@@ -745,14 +749,22 @@ def main():
|
|
745 |
args.caption_type, args.caption_length
|
746 |
)
|
747 |
|
|
|
|
|
|
|
|
|
748 |
if args.feed_from_tags is not None:
|
749 |
prompt = prompt_from_tags(args, image_path, tagset_normalizer, prompt)
|
|
|
|
|
750 |
|
751 |
if args.dry_run:
|
752 |
logging.info(
|
753 |
f"Dry run: Skipping caption generation for {image_path} with prompt:\n\t{prompt}"
|
754 |
)
|
755 |
continue
|
|
|
|
|
756 |
|
757 |
caption = joy_caption_model.generate_valid_caption(input_image, prompt)
|
758 |
|
@@ -1029,5 +1041,12 @@ def find_tag_file(image_path):
|
|
1029 |
return None
|
1030 |
|
1031 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1032 |
if __name__ == "__main__":
|
1033 |
main()
|
|
|
319 |
Tuple[str, float]: The generated caption and its entropy.
|
320 |
"""
|
321 |
torch.cuda.empty_cache()
|
|
|
322 |
|
323 |
pixel_values = self._preprocess_image(input_image)
|
324 |
|
|
|
686 |
"Optionally specify the number of tags to use."
|
687 |
),
|
688 |
)
|
689 |
+
parser.add_argument(
|
690 |
+
"--artist-from-folder",
|
691 |
+
action="store_true",
|
692 |
+
help="Get the artist name from the parent folder",
|
693 |
+
)
|
694 |
parser.add_argument(
|
695 |
"--random-tags",
|
696 |
type=int,
|
|
|
749 |
args.caption_type, args.caption_length
|
750 |
)
|
751 |
|
752 |
+
if args.feed_from_tags is not None and args.artist_from_folder:
|
753 |
+
raise ValueError(
|
754 |
+
"feed-from-tags and artist-from-folder can't be used together"
|
755 |
+
)
|
756 |
if args.feed_from_tags is not None:
|
757 |
prompt = prompt_from_tags(args, image_path, tagset_normalizer, prompt)
|
758 |
+
elif args.artist_from_folder:
|
759 |
+
prompt = prompt_from_folder(prompt, image_path.resolve())
|
760 |
|
761 |
if args.dry_run:
|
762 |
logging.info(
|
763 |
f"Dry run: Skipping caption generation for {image_path} with prompt:\n\t{prompt}"
|
764 |
)
|
765 |
continue
|
766 |
+
else:
|
767 |
+
logging.info(f"Prompt for {image_path}:\n\t{prompt}")
|
768 |
|
769 |
caption = joy_caption_model.generate_valid_caption(input_image, prompt)
|
770 |
|
|
|
1041 |
return None
|
1042 |
|
1043 |
|
1044 |
+
def prompt_from_folder(prompt, path):
|
1045 |
+
artist = (
|
1046 |
+
path.parent.name.replace("_", " ").replace("-", " ").replace(".", " ").title()
|
1047 |
+
)
|
1048 |
+
return prompt.replace("image", f"image by {artist}")
|
1049 |
+
|
1050 |
+
|
1051 |
if __name__ == "__main__":
|
1052 |
main()
|