Neural Networks
===============

Models
------

Conditional UNet
~~~~~~~~~~~~~~~~

.. currentmodule:: diffuse.neural_network.nn.condUNet

.. autoclass:: CondUNet2D
   :members:
   :show-inheritance:

VAE (Variational Autoencoder)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

.. currentmodule:: diffuse.neural_network.nn.sdVae

.. autoclass:: SDVae
   :members:
   :show-inheritance:

Network Parameters
~~~~~~~~~~~~~~~~~~

.. automodule:: diffuse.neural_network.nn.params
   :members:
   :undoc-members:
   :show-inheritance:

Building Blocks
---------------

Attention
~~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.attention
   :members:
   :undoc-members:
   :show-inheritance:

ResNet Block
~~~~~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.resnet_block
   :members:
   :undoc-members:
   :show-inheritance:

Time Embedding
~~~~~~~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.time_embedding
   :members:
   :undoc-members:
   :show-inheritance:

Timestep
~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.timestep
   :members:
   :undoc-members:
   :show-inheritance:

Encoder
~~~~~~~

.. automodule:: diffuse.neural_network.blocks.encoder
   :members:
   :undoc-members:
   :show-inheritance:

Decoder
~~~~~~~

.. automodule:: diffuse.neural_network.blocks.decoder
   :members:
   :undoc-members:
   :show-inheritance:

Downsample
~~~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.downsample
   :members:
   :undoc-members:
   :show-inheritance:

Upsample
~~~~~~~~

.. automodule:: diffuse.neural_network.blocks.upsample
   :members:
   :undoc-members:
   :show-inheritance: