from pathlib import Path

AnimeHeads

from deep_t2i.model import Anime_Export
from deep_t2i.inference_anime_heads import pred_and_show
model_dir = Path('../models/large_anime_heads/gan')
model = Anime_Export.from_pretrained(model_dir/'gan_export.pt')
hairs = ['orange hair', 'white hair', 'aqua hair', 'gray hair','green hair', 'red hair',
         'purple hair', 'pink hair','blue hair', 'black hair', 'brown hair', 'blonde hair']
eyes = ['black eyes', 'orange eyes', 'purple eyes', 'pink eyes', 'yellow eyes', 'aqua eyes', 
        'green eyes', 'brown eyes', 'red eyes', 'blue eyes']
pred_and_show(model, f'{hairs[0]} {eyes[4]}')

Birds

from deep_t2i.model import Birds_Export
from deep_t2i.inference_birds import pred_and_show
model_dir = Path('../models/large_birds/gan')
model = Birds_Export.from_pretrained(model_dir/'gan_export.pt')
examples = [
    "this bird has a green crown, black wings and a yellow belly", 
    "the bird has a yellow crown and a black eyering that is round", 
    "a small red bird", 
    "a small orange bird",
    "this bird has wings that are red and has a white belly",
    "this goofy looking bird has a bright red beak and dull white and gray plumage.",
    "this bird is white in color, with a yellow beak.",
    "this bird is red and black in color and has a sharp black beak",
    "a small black bird",
    "a small yellow bird",
]
pred_and_show(model, examples[0])
pred_and_show(model, examples[1])
pred_and_show(model, examples[2])
pred_and_show(model, examples[3])
pred_and_show(model, examples[4])
pred_and_show(model, examples[5])
pred_and_show(model, examples[6])
pred_and_show(model, examples[7])
pred_and_show(model, examples[8])
pred_and_show(model, examples[9])