File tree 1 file changed +8
-7
lines changed 1 file changed +8
-7
lines changed Original file line number Diff line number Diff line change 2
2
3
3
import numpy as np
4
4
import torch
5
- from transformers import AutoTokenizer , GemmaConfig , GemmaForCausalLM
5
+ from transformers import AutoTokenizer , Gemma2Config , Gemma2Model
6
6
7
7
from diffusers import (
8
8
AutoencoderKL ,
@@ -81,15 +81,16 @@ def get_dummy_components(self):
81
81
tokenizer = AutoTokenizer .from_pretrained ("hf-internal-testing/dummy-gemma" )
82
82
83
83
torch .manual_seed (0 )
84
- config = GemmaConfig (
85
- head_dim = 2 ,
84
+ config = Gemma2Config (
85
+ head_dim = 4 ,
86
86
hidden_size = 8 ,
87
- intermediate_size = 37 ,
88
- num_attention_heads = 4 ,
87
+ intermediate_size = 8 ,
88
+ num_attention_heads = 2 ,
89
89
num_hidden_layers = 2 ,
90
- num_key_value_heads = 4 ,
90
+ num_key_value_heads = 2 ,
91
+ sliding_window = 2 ,
91
92
)
92
- text_encoder = GemmaForCausalLM (config )
93
+ text_encoder = Gemma2Model (config )
93
94
94
95
components = {
95
96
"transformer" : transformer .eval (),
You can’t perform that action at this time.
0 commit comments