diffusion_models.models.SimpleUnet
==================================

.. py:module:: diffusion_models.models.SimpleUnet




Module Contents
---------------

.. py:class:: Block(in_ch, out_ch, time_emb_dim, up=False)



   .. py:attribute:: time_mlp


   .. py:attribute:: conv2


   .. py:attribute:: bnorm1


   .. py:attribute:: bnorm2


   .. py:attribute:: relu


   .. py:method:: forward(x, t)


.. py:class:: SinusoidalPositionEmbeddings(dim)



   .. py:attribute:: dim


   .. py:method:: forward(time)


.. py:class:: SimpleUnet(diffuser, image_channels)



   A Simplified variant of the Unet architecture used in DDPM.

   :param diffuser: A gaussian diffuser.
   :param image_channels: The number of image channels.


   .. py:attribute:: time_mlp


   .. py:attribute:: conv0


   .. py:attribute:: downs


   .. py:attribute:: ups


   .. py:attribute:: output


   .. py:method:: forward(x, timestep)

      Forward pass of the diffusion model.

      The forward pass of the diffusion model, predicting the noise at a single
      step.

      :param x: A batch of noisy images.
      :param timestep: The timesteps of each image in the batch.

      :returns: A tensor representing the noise predicted for each image.