FazBrowse GitHub Viewer
|
Trending
|
URL:
|
Home
Tools:
[Download Repo ZIP]
[View Raw Code]
[Original HTTPS Page]
diffusers/scripts/convert_wuerstchen.py at main · ssghost/diffusers · GitHub
ssghost
/
diffusers
Public
forked from
huggingface/diffusers
Notifications
You must be signed in to change notification settings
Fork
0
Star
0
Code
Pull requests
0
Actions
Projects
Security and quality
0
Insights
Additional navigation options
Code
Pull requests
Actions
Projects
Security and quality
Insights
Expand file tree
Breadcrumbs
diffusers
/
scripts
/
convert_wuerstchen.py
Copy path
More file actions
More file actions
Latest commit
History
History
History
115 lines (100 loc) · 4.78 KB
Breadcrumbs
diffusers
/
scripts
/
convert_wuerstchen.py
Copy path
File metadata and controls
115 lines (100 loc) · 4.78 KB
Raw
Copy raw file
Download raw file
Open symbols panel
Edit and raw actions
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# Run inside root directory of official source code: https://github.com/dome272/wuerstchen/
import
os
import
torch
from
transformers
import
AutoTokenizer
,
CLIPTextModel
from
vqgan
import
VQModel
from
diffusers
import
(
DDPMWuerstchenScheduler
,
WuerstchenCombinedPipeline
,
WuerstchenDecoderPipeline
,
WuerstchenPriorPipeline
,
)
from
diffusers
.
pipelines
.
wuerstchen
import
PaellaVQModel
,
WuerstchenDiffNeXt
,
WuerstchenPrior
model_path
=
"models/"
device
=
"cpu"
paella_vqmodel
=
VQModel
()
state_dict
=
torch
.
load
(
os
.
path
.
join
(
model_path
,
"vqgan_f4_v1_500k.pt"
),
map_location
=
device
)[
"state_dict"
]
paella_vqmodel
.
load_state_dict
(
state_dict
)
state_dict
[
"vquantizer.embedding.weight"
]
=
state_dict
[
"vquantizer.codebook.weight"
]
state_dict
.
pop
(
"vquantizer.codebook.weight"
)
vqmodel
=
PaellaVQModel
(
num_vq_embeddings
=
paella_vqmodel
.
codebook_size
,
latent_channels
=
paella_vqmodel
.
c_latent
)
vqmodel
.
load_state_dict
(
state_dict
)
# Clip Text encoder and tokenizer
text_encoder
=
CLIPTextModel
.
from_pretrained
(
"laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
)
tokenizer
=
AutoTokenizer
.
from_pretrained
(
"laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"
)
# Generator
gen_text_encoder
=
CLIPTextModel
.
from_pretrained
(
"laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
).
to
(
"cpu"
)
gen_tokenizer
=
AutoTokenizer
.
from_pretrained
(
"laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
)
orig_state_dict
=
torch
.
load
(
os
.
path
.
join
(
model_path
,
"model_v2_stage_b.pt"
),
map_location
=
device
)[
"state_dict"
]
state_dict
=
{}
for
key
in
orig_state_dict
.
keys
():
if
key
.
endswith
(
"in_proj_weight"
):
weights
=
orig_state_dict
[
key
].
chunk
(
3
,
0
)
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_q.weight"
)]
=
weights
[
0
]
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_k.weight"
)]
=
weights
[
1
]
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_v.weight"
)]
=
weights
[
2
]
elif
key
.
endswith
(
"in_proj_bias"
):
weights
=
orig_state_dict
[
key
].
chunk
(
3
,
0
)
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_q.bias"
)]
=
weights
[
0
]
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_k.bias"
)]
=
weights
[
1
]
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_v.bias"
)]
=
weights
[
2
]
elif
key
.
endswith
(
"out_proj.weight"
):
weights
=
orig_state_dict
[
key
]
state_dict
[
key
.
replace
(
"attn.out_proj.weight"
,
"to_out.0.weight"
)]
=
weights
elif
key
.
endswith
(
"out_proj.bias"
):
weights
=
orig_state_dict
[
key
]
state_dict
[
key
.
replace
(
"attn.out_proj.bias"
,
"to_out.0.bias"
)]
=
weights
else
:
state_dict
[
key
]
=
orig_state_dict
[
key
]
deocder
=
WuerstchenDiffNeXt
()
deocder
.
load_state_dict
(
state_dict
)
# Prior
orig_state_dict
=
torch
.
load
(
os
.
path
.
join
(
model_path
,
"model_v3_stage_c.pt"
),
map_location
=
device
)[
"ema_state_dict"
]
state_dict
=
{}
for
key
in
orig_state_dict
.
keys
():
if
key
.
endswith
(
"in_proj_weight"
):
weights
=
orig_state_dict
[
key
].
chunk
(
3
,
0
)
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_q.weight"
)]
=
weights
[
0
]
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_k.weight"
)]
=
weights
[
1
]
state_dict
[
key
.
replace
(
"attn.in_proj_weight"
,
"to_v.weight"
)]
=
weights
[
2
]
elif
key
.
endswith
(
"in_proj_bias"
):
weights
=
orig_state_dict
[
key
].
chunk
(
3
,
0
)
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_q.bias"
)]
=
weights
[
0
]
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_k.bias"
)]
=
weights
[
1
]
state_dict
[
key
.
replace
(
"attn.in_proj_bias"
,
"to_v.bias"
)]
=
weights
[
2
]
elif
key
.
endswith
(
"out_proj.weight"
):
weights
=
orig_state_dict
[
key
]
state_dict
[
key
.
replace
(
"attn.out_proj.weight"
,
"to_out.0.weight"
)]
=
weights
elif
key
.
endswith
(
"out_proj.bias"
):
weights
=
orig_state_dict
[
key
]
state_dict
[
key
.
replace
(
"attn.out_proj.bias"
,
"to_out.0.bias"
)]
=
weights
else
:
state_dict
[
key
]
=
orig_state_dict
[
key
]
prior_model
=
WuerstchenPrior
(
c_in
=
16
,
c
=
1536
,
c_cond
=
1280
,
c_r
=
64
,
depth
=
32
,
nhead
=
24
).
to
(
device
)
prior_model
.
load_state_dict
(
state_dict
)
# scheduler
scheduler
=
DDPMWuerstchenScheduler
()
# Prior pipeline
prior_pipeline
=
WuerstchenPriorPipeline
(
prior
=
prior_model
,
text_encoder
=
text_encoder
,
tokenizer
=
tokenizer
,
scheduler
=
scheduler
)
prior_pipeline
.
save_pretrained
(
"warp-ai/wuerstchen-prior"
)
decoder_pipeline
=
WuerstchenDecoderPipeline
(
text_encoder
=
gen_text_encoder
,
tokenizer
=
gen_tokenizer
,
vqgan
=
vqmodel
,
decoder
=
deocder
,
scheduler
=
scheduler
)
decoder_pipeline
.
save_pretrained
(
"warp-ai/wuerstchen"
)
# Wuerstchen pipeline
wuerstchen_pipeline
=
WuerstchenCombinedPipeline
(
# Decoder
text_encoder
=
gen_text_encoder
,
tokenizer
=
gen_tokenizer
,
decoder
=
deocder
,
scheduler
=
scheduler
,
vqgan
=
vqmodel
,
# Prior
prior_tokenizer
=
tokenizer
,
prior_text_encoder
=
text_encoder
,
prior
=
prior_model
,
prior_scheduler
=
scheduler
,
)
wuerstchen_pipeline
.
save_pretrained
(
"warp-ai/WuerstchenCombinedPipeline"
)
Back
|
FazBrowse Home
|
New Git URL