diffusion_models.diffusion_inference ==================================== .. py:module:: diffusion_models.diffusion_inference Module Contents --------------- .. py:class:: DiffusionInference(model, reverse_transforms = lambda x: x, image_shape = (3, 64), device = 'cuda') A diffusion inference framework. This is a simplified inference framework to easily start using your diffusion model. :param model: The trained diffusion model. :param reverse_transforms: A set of reverse transforms. :param image_shape: The shape of the image to produce. This is a tuple where the first value is the number of channels and the second is the size of the image. Images are expected to be square. :param device: The device to run the inference on. .. py:attribute:: image_channels The number of channels of the image. .. py:attribute:: image_size The size of the image. .. py:attribute:: model The trained diffusion model. .. py:attribute:: reverse_transforms The set of reverse transforms. .. py:attribute:: device The device to run the inference on. .. py:method:: generate(number_of_images, save_gif = False) Generate a batch of images. :param number_of_images: The number of images to generate. :param save_gif: Whether to save the generation process as a GIF. :returns: A PIL image of the generated images stacked together. .. py:method:: get_generator(number_of_images = 1) An image generator. This method is a generator that will generate a batch of images. At every call, the generator denoises the images by one more step until the image is fully generated. This can be particularly useful for running the image generation step by step or for a streaming API. :param number_of_images: The number of images the generator should generate.