File size: 1,564 Bytes
59da1c6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import os
import pickle
import imageio
import matplotlib.pyplot as plt
from qdhf_things import run_qdhf, many_pictures
EXAMPLE_PROMPTS = [
'an image of a cat on the sofa',
'an image of a bear in a national park',
'a photo of an astronaut riding a horse on mars',
'a drawing of a tree behind a fence',
'a painting of a sunset over the ocean',
'a sketch of a racoon sitting on a mushroom',
'a picture of a dragon flying over a castle',
'a photo of a robot playing the guitar',
]
if __name__ == '__main__':
print('Hello! I am a script!')
for i, example_prompt in enumerate(EXAMPLE_PROMPTS):
# Initialize list to store images for GIF
images = []
# Run QDHF
for archive, plt in run_qdhf(example_prompt):
# Save current plot to a temporary file
temp_filename = f'./examples/temp_plot_{i}.png'
plt.savefig(temp_filename)
plt.close()
# Read the saved image and append to images list
images.append(imageio.imread(temp_filename))
os.remove(temp_filename)
# Create a GIF from the images
gif_filename = f'./examples/archive_{i}.gif'
imageio.mimsave(gif_filename, images, duration=0.5) # Adjust duration as needed
# Save archive with pickle
pickle.dump(archive, open(f'./examples/archive_{i}.pkl', 'wb'))
# Save the final archive plot
plt = many_pictures(archive, example_prompt)
plt.savefig(f'./examples/archive_pics_{i}.png')
plt.close()
|