SimpleUnet
==========

.. py:module:: SimpleUnet




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

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



   Initialize internal Module state, shared by both nn.Module and ScriptModule.


   .. 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)



   Initialize internal Module state, shared by both nn.Module and ScriptModule.


   .. 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:: image_channels


   .. py:attribute:: down_channels
      :value: (64, 128, 256, 512, 1024)



   .. py:attribute:: up_channels
      :value: (1024, 512, 256, 128, 64)



   .. py:attribute:: out_dim


   .. py:attribute:: time_emb_dim
      :value: 32



   .. 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.