stablediffusion-wgpu 0.1.2

Stable diffusion burn with wgpu
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

use super::GroupNorm;
use crate::model::load::*;

use std::error::Error;

use burn::tensor::{backend::Backend, Tensor};

pub fn load_group_norm<B: Backend>(
    path: &str,
    device: &B::Device,
) -> Result<GroupNorm<B>, Box<dyn Error>> {
    let n_group = load_usize::<B>("n_group", path, device)?.into();
    let n_channel = load_usize::<B>("n_channel", path, device)?.into();
    let eps = load_f32::<B>("eps", path, device)?.into();

    let gamma = load_tensor::<B, 1>("weight", path, device)
        .ok()
        .unwrap_or_else(|| Tensor::ones_device([n_channel], device))
        .into();
    let beta = load_tensor::<B, 1>("bias", path, device)
        .ok()
        .unwrap_or_else(|| Tensor::zeros_device([n_channel], device))
        .into();

    Ok(GroupNorm {
        n_group,
        n_channel,
        gamma,
        beta,
        eps,
    })
}