Skip to content

Instantly share code, notes, and snippets.

@CoffeeVampir3
Created February 13, 2024 20:44
Show Gist options
  • Save CoffeeVampir3/7a7c5b78347aaee263c660ac8461b7b3 to your computer and use it in GitHub Desktop.
Save CoffeeVampir3/7a7c5b78347aaee263c660ac8461b7b3 to your computer and use it in GitHub Desktop.
moondream tagger
from transformers import AutoModelForCausalLM, CodeGenTokenizerFast as Tokenizer
from moondream import Moondream
from PIL import Image
from collections import Counter
import glob
import os
model = Moondream.from_pretrained("/home/blackroot/Desktop/moond/moondream1").to("cuda")
tokenizer = Tokenizer.from_pretrained("/home/blackroot/Desktop/moond/moondream1/tokenizer")
characters = [
"modeus",
"azazel",
"malina",
"lucifer",
"beelzebub",
"pandemonica",
"cerberus",
"millicent"
]
directory_path = '/home/blackroot/Desktop/moond/moondream1/assets'
pattern = f"{directory_path}/*.png"
for filepath in glob.glob(pattern):
#print(filepath)
image = Image.open(filepath)
enc_image = model.encode_image(image)
# Construct the text file name by replacing .png with .txt
txt_filename = os.path.splitext(filepath)[0] + '.txt'
txt_filepath = os.path.join(directory_path, txt_filename)
words = filepath.split('_')
matched_character = ""
for word in words:
word_cleaned = os.path.splitext(word)[0].lower()
if word_cleaned in characters:
matched_character = word
break
answer = model.answer_question(enc_image, f"Describe the character:", tokenizer)
answer = f"{matched_character} {answer}"
print(answer)
with open(txt_filepath, 'w') as txt_file:
txt_file.write(answer)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment