huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-21 12:20:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use the module trait in stable-diffusion. (#817)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-11 20:40:07 +01:00
committed by GitHub
parent 59e63d690c
commit c5a058b169
7 changed files with 30 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-transformers/src/models/stable_diffusion/attention.rs
View File

@ -17,7 +17,7 @@ impl GeGlu {
}
}
impl GeGlu {
impl Module for GeGlu {
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let _enter = self.span.enter();
let hidden_states_and_gate = self.proj.forward(xs)?.chunk(2, D::Minus1)?;
@ -53,7 +53,7 @@ impl FeedForward {
}
}
impl FeedForward {
impl Module for FeedForward {
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let _enter = self.span.enter();
let xs = self.project_in.forward(xs)?;
@ -501,8 +501,10 @@ impl AttentionBlock {
xs.reshape((batch, t, self.num_heads, h_times_d / self.num_heads))?
.transpose(1, 2)
}
}
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
impl Module for AttentionBlock {
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let _enter = self.span.enter();
let in_dtype = xs.dtype();
let residual = xs;