k4d3 commited on
Commit
6f8fd26
·
1 Parent(s): 6f44ddc

joy: artist-from-folder flag

Browse files
Files changed (1) hide show
  1. joy +20 -1
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()