Add files via upload

Add initial project files
2023-08-04 14:32:47 -04:00
parent 1aed8b655a
commit e4145441eb
31 changed files with 266571 additions and 0 deletions
--- a/src/model/unet/mod.rs
+++ b/src/model/unet/mod.rs
@@ -0,0 +1,757 @@
+pub mod load;
+
+use burn::{
+    config::Config, 
+    module::{Module, Param},
+    nn::{self, PaddingConfig2d, GELU, conv::{Conv2d, Conv2dConfig}},
+    tensor::{
+        backend::Backend,
+        activation::softmax, 
+        module::embedding, 
+        Tensor,
+        Distribution, 
+        Int, 
+    },
+};
+
+use super::silu::*;
+use super::groupnorm::*;
+use crate::helper::to_float;
+
+use super::attention::qkv_attention;
+
+
+fn timestep_embedding<B: Backend>(timesteps: Tensor<B, 1, Int>, dim: usize, max_period: usize) -> Tensor<B, 2> {
+    let half = dim / 2;
+    let freqs = ( to_float(Tensor::arange_device(0..half, &timesteps.device())) * (-(max_period as f64).ln() / half as f64 ) ).exp();
+    let args = to_float(timesteps) * freqs;
+    Tensor::cat(vec![args.clone().cos(), args.sin()], 0).unsqueeze()
+}
+
+
+#[derive(Config)]
+pub struct UNetConfig {}
+
+impl UNetConfig {
+    pub fn init<B: Backend>(&self) -> UNet<B> {
+        let lin1_time_embed = nn::LinearConfig::new(320, 1280).init();
+        let silu_time_embed = SILU::new();
+        let lin2_time_embed = nn::LinearConfig::new(1280, 1280).init();
+
+        let input_blocks = UNetInputBlocks {
+            conv: Conv2dConfig::new([4, 320], [3, 3]).with_padding(PaddingConfig2d::Explicit(1, 1)).init(),
+            rt1: ResTransformerConfig::new(320, 1280, 320, 768, 8).init(),
+            rt2: ResTransformerConfig::new(320, 1280, 320, 768, 8).init(),
+            d1: DownsampleConfig::new(320).init(),
+            rt3: ResTransformerConfig::new(320, 1280, 640, 768, 8).init(),
+            rt4: ResTransformerConfig::new(640, 1280, 640, 768, 8).init(),
+            d2: DownsampleConfig::new(640).init(),
+            rt5: ResTransformerConfig::new(640, 1280, 1280, 768, 8).init(),
+            rt6: ResTransformerConfig::new(1280, 1280, 1280, 768, 8).init(),
+            d3: DownsampleConfig::new(1280).init(),
+            r1: ResBlockConfig::new(1280, 1280, 1280).init(),
+            r2: ResBlockConfig::new(1280, 1280, 1280).init(),
+        };
+        
+        let middle_block = ResTransformerResConfig::new(1280, 1280, 1280, 768, 8).init();
+
+        let output_blocks = UNetOutputBlocks {
+            r1: ResBlockConfig::new(2560, 1280, 1280).init(),
+            r2: ResBlockConfig::new(2560, 1280, 1280).init(),
+            ru: ResUpSampleConfig::new(2560, 1280, 1280).init(),
+            rt1: ResTransformerConfig::new(2560, 1280, 1280, 768, 8).init(),
+            rt2: ResTransformerConfig::new(2560, 1280, 1280, 768, 8).init(),
+            rtu1: ResTransformerUpsampleConfig::new(1920, 1280, 1280, 768, 8).init(),
+            rt3: ResTransformerConfig::new(1920, 1280, 640, 768, 8).init(),
+            rt4: ResTransformerConfig::new(1280, 1280, 640, 768, 8).init(),
+            rtu2: ResTransformerUpsampleConfig::new(960, 1280, 640, 768, 8).init(),
+            rt5: ResTransformerConfig::new(960, 1280, 320, 768, 8).init(),
+            rt6: ResTransformerConfig::new(640, 1280, 320, 768, 8).init(),
+            rt7: ResTransformerConfig::new(640, 1280, 320, 768, 8).init(),
+        };
+
+        let norm_out = GroupNormConfig::new(32, 320).init();
+        let silu_out = SILU::new();
+        let conv_out = Conv2dConfig::new([320, 4], [3, 3]).with_padding(PaddingConfig2d::Explicit(1, 1)).init();
+
+        UNet {
+            lin1_time_embed, 
+            silu_time_embed, 
+            lin2_time_embed, 
+            input_blocks, 
+            middle_block, 
+            output_blocks, 
+            norm_out, 
+            silu_out, 
+            conv_out, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct UNet<B: Backend> {
+    lin1_time_embed: nn::Linear<B>, 
+    silu_time_embed: SILU, 
+    lin2_time_embed: nn::Linear<B>, 
+    input_blocks: UNetInputBlocks<B>, 
+    middle_block: ResTransformerRes<B>, 
+    output_blocks: UNetOutputBlocks<B>, 
+    norm_out: GroupNorm<B>, 
+    silu_out: SILU, 
+    conv_out: Conv2d<B>, 
+}
+
+impl<B: Backend> UNet<B> {
+    pub fn forward(&self, x: Tensor<B, 4>, timesteps: Tensor<B, 1, Int>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let t_emb = timestep_embedding(timesteps, 320, 10000);
+        let emb = self.lin1_time_embed.forward(t_emb);
+        let emb = self.silu_time_embed.forward(emb);
+        let emb = self.lin2_time_embed.forward(emb);
+
+        let mut saved_inputs = Vec::new();
+        let mut x = x;
+
+        // input blocks
+        for block in self.input_blocks.as_array() {
+            println!("{:?}", x.clone().flatten::<1>(0, 3).slice([0..100]).into_data());
+            x = block.forward(x, emb.clone(), context.clone());
+            saved_inputs.push(x.clone())
+        }
+
+        // middle block
+        x = self.middle_block.forward(x, emb.clone(), context.clone());
+
+        // output blocks
+        for block in self.output_blocks.as_array() {
+            x = Tensor::cat(vec![x, saved_inputs.pop().unwrap()], 1);
+            x = block.forward(x, emb.clone(), context.clone());
+        }
+
+        let x = self.norm_out.forward(x);
+        let x = self.silu_out.forward(x);
+        let x = self.conv_out.forward(x);
+        x
+    }
+}
+
+
+
+#[derive(Module, Debug)]
+pub struct UNetInputBlocks<B: Backend> {
+    conv: Conv2d<B>, 
+    rt1: ResTransformer<B>, 
+    rt2: ResTransformer<B>, 
+    d1: Downsample<B>, 
+    rt3: ResTransformer<B>, 
+    rt4: ResTransformer<B>, 
+    d2: Downsample<B>, 
+    rt5: ResTransformer<B>, 
+    rt6: ResTransformer<B>, 
+    d3: Downsample<B>, 
+    r1: ResBlock<B>, 
+    r2: ResBlock<B>, 
+}
+
+impl<B: Backend> UNetInputBlocks<B> {
+    fn as_array(&self) -> [&dyn UNetBlock<B>; 12] {
+        [
+            &self.conv, 
+            &self.rt1, 
+            &self.rt2, 
+            &self.d1, 
+            &self.rt3, 
+            &self.rt4,
+            &self.d2, 
+            &self.rt5, 
+            &self.rt6, 
+            &self.d3, 
+            &self.r1, 
+            &self.r2, 
+        ]
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct UNetOutputBlocks<B: Backend> {
+    r1: ResBlock<B>,
+    r2: ResBlock<B>,
+    ru: ResUpSample<B>,
+    rt1: ResTransformer<B>,
+    rt2: ResTransformer<B>,
+    rtu1: ResTransformerUpsample<B>, 
+    rt3: ResTransformer<B>,
+    rt4: ResTransformer<B>,
+    rtu2: ResTransformerUpsample<B>, 
+    rt5: ResTransformer<B>, 
+    rt6: ResTransformer<B>, 
+    rt7: ResTransformer<B>, 
+}
+
+impl<B: Backend> UNetOutputBlocks<B> {
+    fn as_array(&self) -> [&dyn UNetBlock<B>; 12] {
+        [
+            &self.r1,
+            &self.r2,
+            &self.ru,
+            &self.rt1,
+            &self.rt2,
+            &self.rtu1,
+            &self.rt3,
+            &self.rt4,
+            &self.rtu2,
+            &self.rt5,
+            &self.rt6,
+            &self.rt7,
+        ]
+    }
+}
+
+
+
+
+
+trait UNetBlock<B: Backend> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4>;
+}
+
+#[derive(Config)]
+pub struct ResTransformerConfig {
+    n_channels_in: usize, 
+    n_channels_embed: usize, 
+    n_channels_out: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl ResTransformerConfig {
+    fn init<B: Backend>(&self) -> ResTransformer<B> {
+        let res = ResBlockConfig::new(self.n_channels_in, self.n_channels_embed, self.n_channels_out).init();
+        let transformer = SpatialTransformerConfig::new(self.n_channels_out, self.n_context_state, self.n_head).init();
+
+        ResTransformer {
+            res, 
+            transformer, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct ResTransformer<B: Backend> {
+    res: ResBlock<B>, 
+    transformer: SpatialTransformer<B>, 
+}
+
+impl<B: Backend> UNetBlock<B> for ResTransformer<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let x = self.res.forward(x, emb);
+        let x = self.transformer.forward(x, context);
+        x
+    }
+}
+
+#[derive(Config)]
+pub struct ResUpSampleConfig {
+    n_channels_in: usize, 
+    n_channels_embed: usize, 
+    n_channels_out: usize, 
+}
+
+impl ResUpSampleConfig {
+    fn init<B: Backend>(&self) -> ResUpSample<B> {
+        let res = ResBlockConfig::new(self.n_channels_in, self.n_channels_embed, self.n_channels_out).init();
+        let upsample = UpsampleConfig::new(self.n_channels_out).init();
+
+        ResUpSample {
+            res, 
+            upsample, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct ResUpSample<B: Backend> {
+    res: ResBlock<B>, 
+    upsample: Upsample<B>, 
+}
+
+impl<B: Backend> UNetBlock<B> for ResUpSample<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let x = self.res.forward(x, emb);
+        let x = self.upsample.forward(x);
+        x
+    }
+}
+
+#[derive(Config)]
+pub struct ResTransformerUpsampleConfig {
+    n_channels_in: usize, 
+    n_channels_embed: usize, 
+    n_channels_out: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl ResTransformerUpsampleConfig {
+    fn init<B: Backend>(&self) -> ResTransformerUpsample<B> {
+        let res = ResBlockConfig::new(self.n_channels_in, self.n_channels_embed, self.n_channels_out).init();
+        let transformer = SpatialTransformerConfig::new(self.n_channels_out, self.n_context_state, self.n_head).init();
+        let upsample = UpsampleConfig::new(self.n_channels_out).init();
+
+        ResTransformerUpsample {
+            res, 
+            transformer, 
+            upsample, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct ResTransformerUpsample<B: Backend> {
+    res: ResBlock<B>, 
+    transformer: SpatialTransformer<B>, 
+    upsample: Upsample<B>, 
+}
+
+impl<B: Backend> UNetBlock<B> for ResTransformerUpsample<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let x = self.res.forward(x, emb);
+        let x = self.transformer.forward(x, context);
+        let x = self.upsample.forward(x);
+        x
+    }
+}
+
+#[derive(Config)]
+pub struct ResTransformerResConfig {
+    n_channels_in: usize, 
+    n_channels_embed: usize, 
+    n_channels_out: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl ResTransformerResConfig {
+    fn init<B: Backend>(&self) -> ResTransformerRes<B> {
+        let res1 = ResBlockConfig::new(self.n_channels_in, self.n_channels_embed, self.n_channels_out).init();
+        let transformer = SpatialTransformerConfig::new(self.n_channels_out, self.n_context_state, self.n_head).init();
+        let res2 = ResBlockConfig::new(self.n_channels_in, self.n_channels_embed, self.n_channels_out).init();
+
+        ResTransformerRes {
+            res1, 
+            transformer, 
+            res2, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct ResTransformerRes<B: Backend> {
+    res1: ResBlock<B>, 
+    transformer: SpatialTransformer<B>, 
+    res2: ResBlock<B>, 
+}
+
+impl<B: Backend> UNetBlock<B> for ResTransformerRes<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let x = self.res1.forward(x, emb.clone());
+        let x = self.transformer.forward(x, context);
+        let x = self.res2.forward(x, emb);
+        x
+    }
+}
+
+
+
+#[derive(Config)]
+pub struct UpsampleConfig {
+    n_channels: usize, 
+}
+
+impl UpsampleConfig {
+    fn init<B: Backend>(&self) -> Upsample<B> {
+        let conv = Conv2dConfig::new([self.n_channels, self.n_channels], [3, 3])
+            .with_stride([2, 2])
+            .with_padding(PaddingConfig2d::Explicit(1, 1))
+            .init();
+
+        Upsample {
+            conv, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Upsample<B: Backend> {
+    conv: Conv2d<B>, 
+}
+
+impl<B: Backend> Upsample<B> {
+    fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
+        let [n_batch, n_channel, height, width] = x.dims();
+        let x = x
+                .reshape([n_batch, n_channel, height, 1, width, 1])
+                .repeat(3, 2)
+                .repeat(5, 2)
+                .reshape([n_batch, n_channel, 2 * height, 2 * width]);
+        self.conv.forward(x)
+    }
+}
+
+impl<B: Backend> UNetBlock<B> for Upsample<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        self.forward(x)
+    }
+}
+
+#[derive(Config)]
+pub struct DownsampleConfig {
+    n_channels: usize, 
+}
+
+impl DownsampleConfig {
+    fn init<B: Backend>(&self) -> Conv2d<B> {
+        Conv2dConfig::new([self.n_channels, self.n_channels], [3, 3])
+            .with_stride([2, 2])
+            .with_padding(PaddingConfig2d::Explicit(1, 1))
+            .init()
+    }
+}
+
+type Downsample<B> = Conv2d<B>;
+
+impl<B: Backend> UNetBlock<B> for Conv2d<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        self.forward(x)
+    }
+}
+
+
+
+
+#[derive(Config)]
+pub struct SpatialTransformerConfig {
+    n_channels: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl SpatialTransformerConfig {
+    fn init<B: Backend>(&self) -> SpatialTransformer<B> {
+        let norm = GroupNormConfig::new(32, self.n_channels).init();
+        let proj_in = Conv2dConfig::new([self.n_channels, self.n_channels], [1, 1]).init();
+        let transformer = TransformerBlockConfig::new(self.n_channels, self.n_context_state, self.n_head).init();
+        let proj_out = Conv2dConfig::new([self.n_channels, self.n_channels], [1, 1]).init();
+
+        SpatialTransformer {
+            norm, 
+            proj_in, 
+            transformer, 
+            proj_out, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct SpatialTransformer<B: Backend> {
+    norm: GroupNorm<B>, 
+    proj_in: Conv2d<B>, 
+    transformer: TransformerBlock<B>, 
+    proj_out: Conv2d<B>, 
+}
+
+impl<B: Backend> SpatialTransformer<B> {
+    fn forward(&self, x: Tensor<B, 4>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        let [n_batch, n_channel, height, width] = x.dims();
+
+        let x_in = x.clone();
+
+        let x = self.norm.forward(x);
+        let x = self.proj_in.forward(x);
+        let x = x.reshape([n_batch, n_channel, height * width]).swap_dims(1, 2);
+
+        let x = self.transformer.forward(x, context)
+            .swap_dims(1, 2)
+            .reshape([n_batch, n_channel, height, width]);
+
+        x_in + self.proj_out.forward(x)
+    }
+}
+
+
+
+
+
+
+
+
+#[derive(Config)]
+pub struct TransformerBlockConfig {
+    n_state: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl TransformerBlockConfig {
+    fn init<B: Backend>(&self) -> TransformerBlock<B> {
+        let norm1 = nn::LayerNormConfig::new(self.n_state).init();
+        let attn1 = MultiHeadAttentionConfig::new(self.n_state, self.n_context_state, self.n_head).init();
+        let norm2 = nn::LayerNormConfig::new(self.n_state).init();
+        let attn2 = MultiHeadAttentionConfig::new(self.n_state, self.n_context_state, self.n_head).init();
+        let norm3 = nn::LayerNormConfig::new(self.n_state).init();
+        let mlp = MLPConfig::new(self.n_state, 4).init();
+
+        TransformerBlock {
+            norm1, 
+            attn1, 
+            norm2, 
+            attn2, 
+            norm3, 
+            mlp, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct TransformerBlock<B: Backend> {
+    norm1: nn::LayerNorm<B>, 
+    attn1: MultiHeadAttention<B>, 
+    norm2: nn::LayerNorm<B>, 
+    attn2: MultiHeadAttention<B>, 
+    norm3: nn::LayerNorm<B>, 
+    mlp: MLP<B>, 
+}
+
+impl<B: Backend> TransformerBlock<B> {
+    fn forward(&self, x: Tensor<B, 3>, context: Tensor<B, 3>) -> Tensor<B, 3> {
+        let x = x.clone() + self.attn1.forward( self.norm1.forward(x), None);
+        let x = x.clone() + self.attn2.forward( self.norm2.forward(x), Some(context));
+        x.clone() + self.mlp.forward( self.norm3.forward(x) )
+    }
+}
+
+
+#[derive(Config)]
+pub struct MLPConfig {
+    n_state: usize, 
+    mult: usize, 
+}
+
+impl MLPConfig {
+    pub fn init<B: Backend>(&self) -> MLP<B> {
+        let n_state_hidden = self.n_state * self.mult;
+        let geglu = GEGLUConfig::new(self.n_state, n_state_hidden).init();
+        let lin = nn::LinearConfig::new(n_state_hidden, self.n_state).init();
+
+        MLP {
+            geglu, 
+            lin, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct MLP<B: Backend> {
+    geglu: GEGLU<B>, 
+    lin: nn::Linear<B>, 
+}
+
+impl<B: Backend> MLP<B> {
+    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
+        self.lin.forward( self.geglu.forward(x) )
+    }
+}
+
+
+#[derive(Config)]
+pub struct GEGLUConfig {
+    n_state_in: usize, 
+    n_state_out: usize, 
+}
+
+impl GEGLUConfig {
+    fn init<B: Backend>(&self) -> GEGLU<B> {
+        let proj = nn::LinearConfig::new(self.n_state_in, 2 * self.n_state_out).init();
+        let gelu = GELU::new();
+
+        GEGLU {
+            proj, 
+            gelu, 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct GEGLU<B: Backend> {
+    proj: nn::Linear<B>, 
+    gelu: GELU, 
+}
+
+impl<B: Backend> GEGLU<B> {
+    fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
+        let projected = self.proj.forward(x);
+        let [n_batch, n_ctx, n_state] = projected.dims();
+
+        let n_state_out = n_state / 2;
+
+        let x = projected.clone().slice([0..n_batch, 0..n_ctx, 0..n_state_out]);
+        let gate = projected.slice([0..n_batch, 0..n_ctx, n_state_out..n_state]);
+
+        x * self.gelu.forward(gate)
+    }
+}
+
+
+
+
+
+#[derive(Config)]
+pub struct MultiHeadAttentionConfig {
+    n_state: usize, 
+    n_context_state: usize, 
+    n_head: usize, 
+}
+
+impl MultiHeadAttentionConfig {
+    fn init<B: Backend>(&self) -> MultiHeadAttention<B> {
+        assert!(self.n_state % self.n_head == 0, "State size {} must be a multiple of head size {}", self.n_state, self.n_head);
+
+        let n_head = self.n_head;
+        let query = nn::LinearConfig::new(self.n_state, self.n_state).with_bias(false).init();
+        let key = nn::LinearConfig::new(self.n_context_state, self.n_state).with_bias(false).init();
+        let value = nn::LinearConfig::new(self.n_context_state, self.n_state).with_bias(false).init();
+        let out = nn::LinearConfig::new(self.n_state, self.n_state).init();
+
+        MultiHeadAttention { 
+            n_head, 
+            query, 
+            key, 
+            value, 
+            out 
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct MultiHeadAttention<B: Backend> {
+    n_head: usize, 
+    query: nn::Linear<B>, 
+    key: nn::Linear<B>, 
+    value: nn::Linear<B>, 
+    out: nn::Linear<B>, 
+}
+
+impl<B: Backend> MultiHeadAttention<B> {
+    pub fn forward(&self, x: Tensor<B, 3>, context: Option<Tensor<B, 3>>) -> Tensor<B, 3> {
+        let xa = context.unwrap_or_else(|| x.clone());
+
+        let q = self.query.forward(x);
+        let k = self.key.forward(xa.clone());
+        let v = self.value.forward(xa);
+
+        let wv = qkv_attention(q, k, v, None, self.n_head);
+
+        self.out.forward(wv)
+    }
+}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+#[derive(Config)]
+pub struct ResBlockConfig {
+    n_channels_in: usize, 
+    n_channels_embed: usize, 
+    n_channels_out: usize, 
+}
+
+
+impl ResBlockConfig {
+    fn init<B: Backend>(&self) -> ResBlock<B> {
+        let norm_in = GroupNormConfig::new(32, self.n_channels_in).init();
+        let silu_in = SILU::new();
+        let conv_in = Conv2dConfig::new([self.n_channels_in, self.n_channels_out], [3, 3]).with_padding(PaddingConfig2d::Explicit(1, 1)).init();
+
+        let silu_embed = SILU::new();
+        let lin_embed = nn::LinearConfig::new(self.n_channels_embed, self.n_channels_out).init();
+
+        let norm_out = GroupNormConfig::new(32, self.n_channels_out).init();
+        let silu_out = SILU::new();
+        let conv_out = Conv2dConfig::new([self.n_channels_out, self.n_channels_out], [3, 3]).with_padding(PaddingConfig2d::Explicit(1, 1)).init();
+
+        let skip_connection = if self.n_channels_in != self.n_channels_out {
+            Some( Conv2dConfig::new([self.n_channels_in, self.n_channels_out], [1, 1]).init() )
+        } else {
+            None
+        };
+
+        ResBlock {
+            norm_in, 
+            silu_in, 
+            conv_in, 
+            silu_embed, 
+            lin_embed, 
+            norm_out, 
+            silu_out, 
+            conv_out, 
+            skip_connection, 
+        }
+    }
+}
+
+
+#[derive(Module, Debug)]
+pub struct ResBlock<B: Backend> {
+    norm_in: GroupNorm<B>, 
+    silu_in: SILU, 
+    conv_in: Conv2d<B>, 
+    silu_embed: SILU, 
+    lin_embed: nn::Linear<B>, 
+    norm_out: GroupNorm<B>, 
+    silu_out: SILU, 
+    conv_out: Conv2d<B>, 
+    skip_connection: Option<Conv2d<B>>, 
+}
+
+impl<B: Backend> ResBlock<B> {
+    fn forward(&self, x: Tensor<B, 4>, embed: Tensor<B, 2>) -> Tensor<B, 4> {
+        let h = self.norm_in.forward(x.clone());
+        let h = self.silu_in.forward(h);
+        let h = self.conv_in.forward(h);
+
+        let embed_out = self.silu_embed.forward(embed);
+        let embed_out = self.lin_embed.forward(embed_out);
+        
+        let [n_batch_embed, n_state_embed] = embed_out.dims();
+        let h = h + embed_out.reshape([n_batch_embed, n_state_embed, 1, 1]);
+
+        let h = self.norm_out.forward(h);
+        let h = self.silu_out.forward(h);
+        let h = self.conv_out.forward(h);
+
+        if let Some(skipc) = self.skip_connection.as_ref() {
+            skipc.forward(x) + h
+        } else {
+            x + h
+        }
+    }
+}
+
+impl<B: Backend> UNetBlock<B> for ResBlock<B> {
+    fn forward(&self, x: Tensor<B, 4>, emb: Tensor<B, 2>, context: Tensor<B, 3>) -> Tensor<B, 4> {
+        self.forward(x, emb)
+    }
+}
+
+