File tree Expand file tree Collapse file tree 2 files changed +5
-4
lines changed Expand file tree Collapse file tree 2 files changed +5
-4
lines changed Original file line number Diff line number Diff line change @@ -90,8 +90,12 @@ def first_stage_model_decode_wrap(z):
90
90
sd_model .conditioner .register_forward_pre_hook (send_me_to_gpu )
91
91
elif is_sd2 :
92
92
sd_model .cond_stage_model .model .register_forward_pre_hook (send_me_to_gpu )
93
+ sd_model .cond_stage_model .model .token_embedding .register_forward_pre_hook (send_me_to_gpu )
94
+ parents [sd_model .cond_stage_model .model ] = sd_model .cond_stage_model
95
+ parents [sd_model .cond_stage_model .model .token_embedding ] = sd_model .cond_stage_model
93
96
else :
94
97
sd_model .cond_stage_model .transformer .register_forward_pre_hook (send_me_to_gpu )
98
+ parents [sd_model .cond_stage_model .transformer ] = sd_model .cond_stage_model
95
99
96
100
sd_model .first_stage_model .register_forward_pre_hook (send_me_to_gpu )
97
101
sd_model .first_stage_model .encode = first_stage_model_encode_wrap
@@ -101,9 +105,6 @@ def first_stage_model_decode_wrap(z):
101
105
if sd_model .embedder :
102
106
sd_model .embedder .register_forward_pre_hook (send_me_to_gpu )
103
107
104
- if hasattr (sd_model , 'cond_stage_model' ):
105
- parents [sd_model .cond_stage_model .transformer ] = sd_model .cond_stage_model
106
-
107
108
if use_medvram :
108
109
sd_model .model .register_forward_pre_hook (send_me_to_gpu )
109
110
else :
Original file line number Diff line number Diff line change @@ -32,7 +32,7 @@ def encode_with_transformers(self, tokens):
32
32
def encode_embedding_init_text (self , init_text , nvpt ):
33
33
ids = tokenizer .encode (init_text )
34
34
ids = torch .asarray ([ids ], device = devices .device , dtype = torch .int )
35
- embedded = self .wrapped .model .token_embedding .wrapped (ids . to ( self . wrapped . model . token_embedding . wrapped . weight . device ) ).squeeze (0 )
35
+ embedded = self .wrapped .model .token_embedding .wrapped (ids ).squeeze (0 )
36
36
37
37
return embedded
38
38
You can’t perform that action at this time.
0 commit comments