Add support for image generation