From d39dad28782b6fd4380a94a1fdcc8d49f77dddd8 Mon Sep 17 00:00:00 2001 From: arena Date: Fri, 6 Jan 2023 11:38:31 +0800 Subject: [PATCH] allow extracting EMA UNET --- components/UNET-v1-EMA.txt | 686 +++++++++++++++++++++++++++++++++++++ scripts/toolkit_gui.py | 23 +- toolkit.py | 60 +++- 3 files changed, 752 insertions(+), 17 deletions(-) create mode 100644 components/UNET-v1-EMA.txt diff --git a/components/UNET-v1-EMA.txt b/components/UNET-v1-EMA.txt new file mode 100644 index 0000000..10b4928 --- /dev/null +++ b/components/UNET-v1-EMA.txt @@ -0,0 +1,686 @@ +input_blocks00bias [320] +input_blocks00weight [320,4,3,3] +input_blocks100emb_layers1bias [1280] +input_blocks100emb_layers1weight [1280,1280] +input_blocks100in_layers0bias [1280] +input_blocks100in_layers0weight [1280] +input_blocks100in_layers2bias [1280] +input_blocks100in_layers2weight [1280,1280,3,3] +input_blocks100out_layers0bias [1280] +input_blocks100out_layers0weight [1280] +input_blocks100out_layers3bias [1280] +input_blocks100out_layers3weight [1280,1280,3,3] +input_blocks10emb_layers1bias [320] +input_blocks10emb_layers1weight [320,1280] +input_blocks10in_layers0bias [320] +input_blocks10in_layers0weight [320] +input_blocks10in_layers2bias [320] +input_blocks10in_layers2weight [320,320,3,3] +input_blocks10out_layers0bias [320] +input_blocks10out_layers0weight [320] +input_blocks10out_layers3bias [320] +input_blocks10out_layers3weight [320,320,3,3] +input_blocks110emb_layers1bias [1280] +input_blocks110emb_layers1weight [1280,1280] +input_blocks110in_layers0bias [1280] +input_blocks110in_layers0weight [1280] +input_blocks110in_layers2bias [1280] +input_blocks110in_layers2weight [1280,1280,3,3] +input_blocks110out_layers0bias [1280] +input_blocks110out_layers0weight [1280] +input_blocks110out_layers3bias [1280] +input_blocks110out_layers3weight [1280,1280,3,3] +input_blocks11normbias [320] +input_blocks11normweight [320] +input_blocks11proj_inbias [320] +input_blocks11proj_inweight [320,320,1,1] +input_blocks11proj_outbias [320] +input_blocks11proj_outweight [320,320,1,1] +input_blocks11transformer_blocks0attn1to_kweight [320,320] +input_blocks11transformer_blocks0attn1to_out0bias [320] +input_blocks11transformer_blocks0attn1to_out0weight [320,320] +input_blocks11transformer_blocks0attn1to_qweight [320,320] +input_blocks11transformer_blocks0attn1to_vweight [320,320] +input_blocks11transformer_blocks0attn2to_kweight [320,768] +input_blocks11transformer_blocks0attn2to_out0bias [320] +input_blocks11transformer_blocks0attn2to_out0weight [320,320] +input_blocks11transformer_blocks0attn2to_qweight [320,320] +input_blocks11transformer_blocks0attn2to_vweight [320,768] +input_blocks11transformer_blocks0ffnet0projbias [2560] +input_blocks11transformer_blocks0ffnet0projweight [2560,320] +input_blocks11transformer_blocks0ffnet2bias [320] +input_blocks11transformer_blocks0ffnet2weight [320,1280] +input_blocks11transformer_blocks0norm1bias [320] +input_blocks11transformer_blocks0norm1weight [320] +input_blocks11transformer_blocks0norm2bias [320] +input_blocks11transformer_blocks0norm2weight [320] +input_blocks11transformer_blocks0norm3bias [320] +input_blocks11transformer_blocks0norm3weight [320] +input_blocks20emb_layers1bias [320] +input_blocks20emb_layers1weight [320,1280] +input_blocks20in_layers0bias [320] +input_blocks20in_layers0weight [320] +input_blocks20in_layers2bias [320] +input_blocks20in_layers2weight [320,320,3,3] +input_blocks20out_layers0bias [320] +input_blocks20out_layers0weight [320] +input_blocks20out_layers3bias [320] +input_blocks20out_layers3weight [320,320,3,3] +input_blocks21normbias [320] +input_blocks21normweight [320] +input_blocks21proj_inbias [320] +input_blocks21proj_inweight [320,320,1,1] +input_blocks21proj_outbias [320] +input_blocks21proj_outweight [320,320,1,1] +input_blocks21transformer_blocks0attn1to_kweight [320,320] +input_blocks21transformer_blocks0attn1to_out0bias [320] +input_blocks21transformer_blocks0attn1to_out0weight [320,320] +input_blocks21transformer_blocks0attn1to_qweight [320,320] +input_blocks21transformer_blocks0attn1to_vweight [320,320] +input_blocks21transformer_blocks0attn2to_kweight [320,768] +input_blocks21transformer_blocks0attn2to_out0bias [320] +input_blocks21transformer_blocks0attn2to_out0weight [320,320] +input_blocks21transformer_blocks0attn2to_qweight [320,320] +input_blocks21transformer_blocks0attn2to_vweight [320,768] +input_blocks21transformer_blocks0ffnet0projbias [2560] +input_blocks21transformer_blocks0ffnet0projweight [2560,320] +input_blocks21transformer_blocks0ffnet2bias [320] +input_blocks21transformer_blocks0ffnet2weight [320,1280] +input_blocks21transformer_blocks0norm1bias [320] +input_blocks21transformer_blocks0norm1weight [320] +input_blocks21transformer_blocks0norm2bias [320] +input_blocks21transformer_blocks0norm2weight [320] +input_blocks21transformer_blocks0norm3bias [320] +input_blocks21transformer_blocks0norm3weight [320] +input_blocks30opbias [320] +input_blocks30opweight [320,320,3,3] +input_blocks40emb_layers1bias [640] +input_blocks40emb_layers1weight [640,1280] +input_blocks40in_layers0bias [320] +input_blocks40in_layers0weight [320] +input_blocks40in_layers2bias [640] +input_blocks40in_layers2weight [640,320,3,3] +input_blocks40out_layers0bias [640] +input_blocks40out_layers0weight [640] +input_blocks40out_layers3bias [640] +input_blocks40out_layers3weight [640,640,3,3] +input_blocks40skip_connectionbias [640] +input_blocks40skip_connectionweight [640,320,1,1] +input_blocks41normbias [640] +input_blocks41normweight [640] +input_blocks41proj_inbias [640] +input_blocks41proj_inweight [640,640,1,1] +input_blocks41proj_outbias [640] +input_blocks41proj_outweight [640,640,1,1] +input_blocks41transformer_blocks0attn1to_kweight [640,640] +input_blocks41transformer_blocks0attn1to_out0bias [640] +input_blocks41transformer_blocks0attn1to_out0weight [640,640] +input_blocks41transformer_blocks0attn1to_qweight [640,640] +input_blocks41transformer_blocks0attn1to_vweight [640,640] +input_blocks41transformer_blocks0attn2to_kweight [640,768] +input_blocks41transformer_blocks0attn2to_out0bias [640] +input_blocks41transformer_blocks0attn2to_out0weight [640,640] +input_blocks41transformer_blocks0attn2to_qweight [640,640] +input_blocks41transformer_blocks0attn2to_vweight [640,768] +input_blocks41transformer_blocks0ffnet0projbias [5120] +input_blocks41transformer_blocks0ffnet0projweight [5120,640] +input_blocks41transformer_blocks0ffnet2bias [640] +input_blocks41transformer_blocks0ffnet2weight [640,2560] +input_blocks41transformer_blocks0norm1bias [640] +input_blocks41transformer_blocks0norm1weight [640] +input_blocks41transformer_blocks0norm2bias [640] +input_blocks41transformer_blocks0norm2weight [640] +input_blocks41transformer_blocks0norm3bias [640] +input_blocks41transformer_blocks0norm3weight [640] +input_blocks50emb_layers1bias [640] +input_blocks50emb_layers1weight [640,1280] +input_blocks50in_layers0bias [640] +input_blocks50in_layers0weight [640] +input_blocks50in_layers2bias [640] +input_blocks50in_layers2weight [640,640,3,3] +input_blocks50out_layers0bias [640] +input_blocks50out_layers0weight [640] +input_blocks50out_layers3bias [640] +input_blocks50out_layers3weight [640,640,3,3] +input_blocks51normbias [640] +input_blocks51normweight [640] +input_blocks51proj_inbias [640] +input_blocks51proj_inweight [640,640,1,1] +input_blocks51proj_outbias [640] +input_blocks51proj_outweight [640,640,1,1] +input_blocks51transformer_blocks0attn1to_kweight [640,640] +input_blocks51transformer_blocks0attn1to_out0bias [640] +input_blocks51transformer_blocks0attn1to_out0weight [640,640] +input_blocks51transformer_blocks0attn1to_qweight [640,640] +input_blocks51transformer_blocks0attn1to_vweight [640,640] +input_blocks51transformer_blocks0attn2to_kweight [640,768] +input_blocks51transformer_blocks0attn2to_out0bias [640] +input_blocks51transformer_blocks0attn2to_out0weight [640,640] +input_blocks51transformer_blocks0attn2to_qweight [640,640] +input_blocks51transformer_blocks0attn2to_vweight [640,768] +input_blocks51transformer_blocks0ffnet0projbias [5120] +input_blocks51transformer_blocks0ffnet0projweight [5120,640] +input_blocks51transformer_blocks0ffnet2bias [640] +input_blocks51transformer_blocks0ffnet2weight [640,2560] +input_blocks51transformer_blocks0norm1bias [640] +input_blocks51transformer_blocks0norm1weight [640] +input_blocks51transformer_blocks0norm2bias [640] +input_blocks51transformer_blocks0norm2weight [640] +input_blocks51transformer_blocks0norm3bias [640] +input_blocks51transformer_blocks0norm3weight [640] +input_blocks60opbias [640] +input_blocks60opweight [640,640,3,3] +input_blocks70emb_layers1bias [1280] +input_blocks70emb_layers1weight [1280,1280] +input_blocks70in_layers0bias [640] +input_blocks70in_layers0weight [640] +input_blocks70in_layers2bias [1280] +input_blocks70in_layers2weight [1280,640,3,3] +input_blocks70out_layers0bias [1280] +input_blocks70out_layers0weight [1280] +input_blocks70out_layers3bias [1280] +input_blocks70out_layers3weight [1280,1280,3,3] +input_blocks70skip_connectionbias [1280] +input_blocks70skip_connectionweight [1280,640,1,1] +input_blocks71normbias [1280] +input_blocks71normweight [1280] +input_blocks71proj_inbias [1280] +input_blocks71proj_inweight [1280,1280,1,1] +input_blocks71proj_outbias [1280] +input_blocks71proj_outweight [1280,1280,1,1] +input_blocks71transformer_blocks0attn1to_kweight [1280,1280] +input_blocks71transformer_blocks0attn1to_out0bias [1280] +input_blocks71transformer_blocks0attn1to_out0weight [1280,1280] +input_blocks71transformer_blocks0attn1to_qweight [1280,1280] +input_blocks71transformer_blocks0attn1to_vweight [1280,1280] +input_blocks71transformer_blocks0attn2to_kweight [1280,768] +input_blocks71transformer_blocks0attn2to_out0bias [1280] +input_blocks71transformer_blocks0attn2to_out0weight [1280,1280] +input_blocks71transformer_blocks0attn2to_qweight [1280,1280] +input_blocks71transformer_blocks0attn2to_vweight [1280,768] +input_blocks71transformer_blocks0ffnet0projbias [10240] +input_blocks71transformer_blocks0ffnet0projweight [10240,1280] +input_blocks71transformer_blocks0ffnet2bias [1280] +input_blocks71transformer_blocks0ffnet2weight [1280,5120] +input_blocks71transformer_blocks0norm1bias [1280] +input_blocks71transformer_blocks0norm1weight [1280] +input_blocks71transformer_blocks0norm2bias [1280] +input_blocks71transformer_blocks0norm2weight [1280] +input_blocks71transformer_blocks0norm3bias [1280] +input_blocks71transformer_blocks0norm3weight [1280] +input_blocks80emb_layers1bias [1280] +input_blocks80emb_layers1weight [1280,1280] +input_blocks80in_layers0bias [1280] +input_blocks80in_layers0weight [1280] +input_blocks80in_layers2bias [1280] +input_blocks80in_layers2weight [1280,1280,3,3] +input_blocks80out_layers0bias [1280] +input_blocks80out_layers0weight [1280] +input_blocks80out_layers3bias [1280] +input_blocks80out_layers3weight [1280,1280,3,3] +input_blocks81normbias [1280] +input_blocks81normweight [1280] +input_blocks81proj_inbias [1280] +input_blocks81proj_inweight [1280,1280,1,1] +input_blocks81proj_outbias [1280] +input_blocks81proj_outweight [1280,1280,1,1] +input_blocks81transformer_blocks0attn1to_kweight [1280,1280] +input_blocks81transformer_blocks0attn1to_out0bias [1280] +input_blocks81transformer_blocks0attn1to_out0weight [1280,1280] +input_blocks81transformer_blocks0attn1to_qweight [1280,1280] +input_blocks81transformer_blocks0attn1to_vweight [1280,1280] +input_blocks81transformer_blocks0attn2to_kweight [1280,768] +input_blocks81transformer_blocks0attn2to_out0bias [1280] +input_blocks81transformer_blocks0attn2to_out0weight [1280,1280] +input_blocks81transformer_blocks0attn2to_qweight [1280,1280] +input_blocks81transformer_blocks0attn2to_vweight [1280,768] +input_blocks81transformer_blocks0ffnet0projbias [10240] +input_blocks81transformer_blocks0ffnet0projweight [10240,1280] +input_blocks81transformer_blocks0ffnet2bias [1280] +input_blocks81transformer_blocks0ffnet2weight [1280,5120] +input_blocks81transformer_blocks0norm1bias [1280] +input_blocks81transformer_blocks0norm1weight [1280] +input_blocks81transformer_blocks0norm2bias [1280] +input_blocks81transformer_blocks0norm2weight [1280] +input_blocks81transformer_blocks0norm3bias [1280] +input_blocks81transformer_blocks0norm3weight [1280] +input_blocks90opbias [1280] +input_blocks90opweight [1280,1280,3,3] +middle_block0emb_layers1bias [1280] +middle_block0emb_layers1weight [1280,1280] +middle_block0in_layers0bias [1280] +middle_block0in_layers0weight [1280] +middle_block0in_layers2bias [1280] +middle_block0in_layers2weight [1280,1280,3,3] +middle_block0out_layers0bias [1280] +middle_block0out_layers0weight [1280] +middle_block0out_layers3bias [1280] +middle_block0out_layers3weight [1280,1280,3,3] +middle_block1normbias [1280] +middle_block1normweight [1280] +middle_block1proj_inbias [1280] +middle_block1proj_inweight [1280,1280,1,1] +middle_block1proj_outbias [1280] +middle_block1proj_outweight [1280,1280,1,1] +middle_block1transformer_blocks0attn1to_kweight [1280,1280] +middle_block1transformer_blocks0attn1to_out0bias [1280] +middle_block1transformer_blocks0attn1to_out0weight [1280,1280] +middle_block1transformer_blocks0attn1to_qweight [1280,1280] +middle_block1transformer_blocks0attn1to_vweight [1280,1280] +middle_block1transformer_blocks0attn2to_kweight [1280,768] +middle_block1transformer_blocks0attn2to_out0bias [1280] +middle_block1transformer_blocks0attn2to_out0weight [1280,1280] +middle_block1transformer_blocks0attn2to_qweight [1280,1280] +middle_block1transformer_blocks0attn2to_vweight [1280,768] +middle_block1transformer_blocks0ffnet0projbias [10240] +middle_block1transformer_blocks0ffnet0projweight [10240,1280] +middle_block1transformer_blocks0ffnet2bias [1280] +middle_block1transformer_blocks0ffnet2weight [1280,5120] +middle_block1transformer_blocks0norm1bias [1280] +middle_block1transformer_blocks0norm1weight [1280] +middle_block1transformer_blocks0norm2bias [1280] +middle_block1transformer_blocks0norm2weight [1280] +middle_block1transformer_blocks0norm3bias [1280] +middle_block1transformer_blocks0norm3weight [1280] +middle_block2emb_layers1bias [1280] +middle_block2emb_layers1weight [1280,1280] +middle_block2in_layers0bias [1280] +middle_block2in_layers0weight [1280] +middle_block2in_layers2bias [1280] +middle_block2in_layers2weight [1280,1280,3,3] +middle_block2out_layers0bias [1280] +middle_block2out_layers0weight [1280] +middle_block2out_layers3bias [1280] +middle_block2out_layers3weight [1280,1280,3,3] +out0bias [320] +out0weight [320] +out2bias [4] +out2weight [4,320,3,3] +output_blocks00emb_layers1bias [1280] +output_blocks00emb_layers1weight [1280,1280] +output_blocks00in_layers0bias [2560] +output_blocks00in_layers0weight [2560] +output_blocks00in_layers2bias [1280] +output_blocks00in_layers2weight [1280,2560,3,3] +output_blocks00out_layers0bias [1280] +output_blocks00out_layers0weight [1280] +output_blocks00out_layers3bias [1280] +output_blocks00out_layers3weight [1280,1280,3,3] +output_blocks00skip_connectionbias [1280] +output_blocks00skip_connectionweight [1280,2560,1,1] +output_blocks100emb_layers1bias [320] +output_blocks100emb_layers1weight [320,1280] +output_blocks100in_layers0bias [640] +output_blocks100in_layers0weight [640] +output_blocks100in_layers2bias [320] +output_blocks100in_layers2weight [320,640,3,3] +output_blocks100out_layers0bias [320] +output_blocks100out_layers0weight [320] +output_blocks100out_layers3bias [320] +output_blocks100out_layers3weight [320,320,3,3] +output_blocks100skip_connectionbias [320] +output_blocks100skip_connectionweight [320,640,1,1] +output_blocks101normbias [320] +output_blocks101normweight [320] +output_blocks101proj_inbias [320] +output_blocks101proj_inweight [320,320,1,1] +output_blocks101proj_outbias [320] +output_blocks101proj_outweight [320,320,1,1] +output_blocks101transformer_blocks0attn1to_kweight [320,320] +output_blocks101transformer_blocks0attn1to_out0bias [320] +output_blocks101transformer_blocks0attn1to_out0weight [320,320] +output_blocks101transformer_blocks0attn1to_qweight [320,320] +output_blocks101transformer_blocks0attn1to_vweight [320,320] +output_blocks101transformer_blocks0attn2to_kweight [320,768] +output_blocks101transformer_blocks0attn2to_out0bias [320] +output_blocks101transformer_blocks0attn2to_out0weight [320,320] +output_blocks101transformer_blocks0attn2to_qweight [320,320] +output_blocks101transformer_blocks0attn2to_vweight [320,768] +output_blocks101transformer_blocks0ffnet0projbias [2560] +output_blocks101transformer_blocks0ffnet0projweight [2560,320] +output_blocks101transformer_blocks0ffnet2bias [320] +output_blocks101transformer_blocks0ffnet2weight [320,1280] +output_blocks101transformer_blocks0norm1bias [320] +output_blocks101transformer_blocks0norm1weight [320] +output_blocks101transformer_blocks0norm2bias [320] +output_blocks101transformer_blocks0norm2weight [320] +output_blocks101transformer_blocks0norm3bias [320] +output_blocks101transformer_blocks0norm3weight [320] +output_blocks10emb_layers1bias [1280] +output_blocks10emb_layers1weight [1280,1280] +output_blocks10in_layers0bias [2560] +output_blocks10in_layers0weight [2560] +output_blocks10in_layers2bias [1280] +output_blocks10in_layers2weight [1280,2560,3,3] +output_blocks10out_layers0bias [1280] +output_blocks10out_layers0weight [1280] +output_blocks10out_layers3bias [1280] +output_blocks10out_layers3weight [1280,1280,3,3] +output_blocks10skip_connectionbias [1280] +output_blocks10skip_connectionweight [1280,2560,1,1] +output_blocks110emb_layers1bias [320] +output_blocks110emb_layers1weight [320,1280] +output_blocks110in_layers0bias [640] +output_blocks110in_layers0weight [640] +output_blocks110in_layers2bias [320] +output_blocks110in_layers2weight [320,640,3,3] +output_blocks110out_layers0bias [320] +output_blocks110out_layers0weight [320] +output_blocks110out_layers3bias [320] +output_blocks110out_layers3weight [320,320,3,3] +output_blocks110skip_connectionbias [320] +output_blocks110skip_connectionweight [320,640,1,1] +output_blocks111normbias [320] +output_blocks111normweight [320] +output_blocks111proj_inbias [320] +output_blocks111proj_inweight [320,320,1,1] +output_blocks111proj_outbias [320] +output_blocks111proj_outweight [320,320,1,1] +output_blocks111transformer_blocks0attn1to_kweight [320,320] +output_blocks111transformer_blocks0attn1to_out0bias [320] +output_blocks111transformer_blocks0attn1to_out0weight [320,320] +output_blocks111transformer_blocks0attn1to_qweight [320,320] +output_blocks111transformer_blocks0attn1to_vweight [320,320] +output_blocks111transformer_blocks0attn2to_kweight [320,768] +output_blocks111transformer_blocks0attn2to_out0bias [320] +output_blocks111transformer_blocks0attn2to_out0weight [320,320] +output_blocks111transformer_blocks0attn2to_qweight [320,320] +output_blocks111transformer_blocks0attn2to_vweight [320,768] +output_blocks111transformer_blocks0ffnet0projbias [2560] +output_blocks111transformer_blocks0ffnet0projweight [2560,320] +output_blocks111transformer_blocks0ffnet2bias [320] +output_blocks111transformer_blocks0ffnet2weight [320,1280] +output_blocks111transformer_blocks0norm1bias [320] +output_blocks111transformer_blocks0norm1weight [320] +output_blocks111transformer_blocks0norm2bias [320] +output_blocks111transformer_blocks0norm2weight [320] +output_blocks111transformer_blocks0norm3bias [320] +output_blocks111transformer_blocks0norm3weight [320] +output_blocks20emb_layers1bias [1280] +output_blocks20emb_layers1weight [1280,1280] +output_blocks20in_layers0bias [2560] +output_blocks20in_layers0weight [2560] +output_blocks20in_layers2bias [1280] +output_blocks20in_layers2weight [1280,2560,3,3] +output_blocks20out_layers0bias [1280] +output_blocks20out_layers0weight [1280] +output_blocks20out_layers3bias [1280] +output_blocks20out_layers3weight [1280,1280,3,3] +output_blocks20skip_connectionbias [1280] +output_blocks20skip_connectionweight [1280,2560,1,1] +output_blocks21convbias [1280] +output_blocks21convweight [1280,1280,3,3] +output_blocks30emb_layers1bias [1280] +output_blocks30emb_layers1weight [1280,1280] +output_blocks30in_layers0bias [2560] +output_blocks30in_layers0weight [2560] +output_blocks30in_layers2bias [1280] +output_blocks30in_layers2weight [1280,2560,3,3] +output_blocks30out_layers0bias [1280] +output_blocks30out_layers0weight [1280] +output_blocks30out_layers3bias [1280] +output_blocks30out_layers3weight [1280,1280,3,3] +output_blocks30skip_connectionbias [1280] +output_blocks30skip_connectionweight [1280,2560,1,1] +output_blocks31normbias [1280] +output_blocks31normweight [1280] +output_blocks31proj_inbias [1280] +output_blocks31proj_inweight [1280,1280,1,1] +output_blocks31proj_outbias [1280] +output_blocks31proj_outweight [1280,1280,1,1] +output_blocks31transformer_blocks0attn1to_kweight [1280,1280] +output_blocks31transformer_blocks0attn1to_out0bias [1280] +output_blocks31transformer_blocks0attn1to_out0weight [1280,1280] +output_blocks31transformer_blocks0attn1to_qweight [1280,1280] +output_blocks31transformer_blocks0attn1to_vweight [1280,1280] +output_blocks31transformer_blocks0attn2to_kweight [1280,768] +output_blocks31transformer_blocks0attn2to_out0bias [1280] +output_blocks31transformer_blocks0attn2to_out0weight [1280,1280] +output_blocks31transformer_blocks0attn2to_qweight [1280,1280] +output_blocks31transformer_blocks0attn2to_vweight [1280,768] +output_blocks31transformer_blocks0ffnet0projbias [10240] +output_blocks31transformer_blocks0ffnet0projweight [10240,1280] +output_blocks31transformer_blocks0ffnet2bias [1280] +output_blocks31transformer_blocks0ffnet2weight [1280,5120] +output_blocks31transformer_blocks0norm1bias [1280] +output_blocks31transformer_blocks0norm1weight [1280] +output_blocks31transformer_blocks0norm2bias [1280] +output_blocks31transformer_blocks0norm2weight [1280] +output_blocks31transformer_blocks0norm3bias [1280] +output_blocks31transformer_blocks0norm3weight [1280] +output_blocks40emb_layers1bias [1280] +output_blocks40emb_layers1weight [1280,1280] +output_blocks40in_layers0bias [2560] +output_blocks40in_layers0weight [2560] +output_blocks40in_layers2bias [1280] +output_blocks40in_layers2weight [1280,2560,3,3] +output_blocks40out_layers0bias [1280] +output_blocks40out_layers0weight [1280] +output_blocks40out_layers3bias [1280] +output_blocks40out_layers3weight [1280,1280,3,3] +output_blocks40skip_connectionbias [1280] +output_blocks40skip_connectionweight [1280,2560,1,1] +output_blocks41normbias [1280] +output_blocks41normweight [1280] +output_blocks41proj_inbias [1280] +output_blocks41proj_inweight [1280,1280,1,1] +output_blocks41proj_outbias [1280] +output_blocks41proj_outweight [1280,1280,1,1] +output_blocks41transformer_blocks0attn1to_kweight [1280,1280] +output_blocks41transformer_blocks0attn1to_out0bias [1280] +output_blocks41transformer_blocks0attn1to_out0weight [1280,1280] +output_blocks41transformer_blocks0attn1to_qweight [1280,1280] +output_blocks41transformer_blocks0attn1to_vweight [1280,1280] +output_blocks41transformer_blocks0attn2to_kweight [1280,768] +output_blocks41transformer_blocks0attn2to_out0bias [1280] +output_blocks41transformer_blocks0attn2to_out0weight [1280,1280] +output_blocks41transformer_blocks0attn2to_qweight [1280,1280] +output_blocks41transformer_blocks0attn2to_vweight [1280,768] +output_blocks41transformer_blocks0ffnet0projbias [10240] +output_blocks41transformer_blocks0ffnet0projweight [10240,1280] +output_blocks41transformer_blocks0ffnet2bias [1280] +output_blocks41transformer_blocks0ffnet2weight [1280,5120] +output_blocks41transformer_blocks0norm1bias [1280] +output_blocks41transformer_blocks0norm1weight [1280] +output_blocks41transformer_blocks0norm2bias [1280] +output_blocks41transformer_blocks0norm2weight [1280] +output_blocks41transformer_blocks0norm3bias [1280] +output_blocks41transformer_blocks0norm3weight [1280] +output_blocks50emb_layers1bias [1280] +output_blocks50emb_layers1weight [1280,1280] +output_blocks50in_layers0bias [1920] +output_blocks50in_layers0weight [1920] +output_blocks50in_layers2bias [1280] +output_blocks50in_layers2weight [1280,1920,3,3] +output_blocks50out_layers0bias [1280] +output_blocks50out_layers0weight [1280] +output_blocks50out_layers3bias [1280] +output_blocks50out_layers3weight [1280,1280,3,3] +output_blocks50skip_connectionbias [1280] +output_blocks50skip_connectionweight [1280,1920,1,1] +output_blocks51normbias [1280] +output_blocks51normweight [1280] +output_blocks51proj_inbias [1280] +output_blocks51proj_inweight [1280,1280,1,1] +output_blocks51proj_outbias [1280] +output_blocks51proj_outweight [1280,1280,1,1] +output_blocks51transformer_blocks0attn1to_kweight [1280,1280] +output_blocks51transformer_blocks0attn1to_out0bias [1280] +output_blocks51transformer_blocks0attn1to_out0weight [1280,1280] +output_blocks51transformer_blocks0attn1to_qweight [1280,1280] +output_blocks51transformer_blocks0attn1to_vweight [1280,1280] +output_blocks51transformer_blocks0attn2to_kweight [1280,768] +output_blocks51transformer_blocks0attn2to_out0bias [1280] +output_blocks51transformer_blocks0attn2to_out0weight [1280,1280] +output_blocks51transformer_blocks0attn2to_qweight [1280,1280] +output_blocks51transformer_blocks0attn2to_vweight [1280,768] +output_blocks51transformer_blocks0ffnet0projbias [10240] +output_blocks51transformer_blocks0ffnet0projweight [10240,1280] +output_blocks51transformer_blocks0ffnet2bias [1280] +output_blocks51transformer_blocks0ffnet2weight [1280,5120] +output_blocks51transformer_blocks0norm1bias [1280] +output_blocks51transformer_blocks0norm1weight [1280] +output_blocks51transformer_blocks0norm2bias [1280] +output_blocks51transformer_blocks0norm2weight [1280] +output_blocks51transformer_blocks0norm3bias [1280] +output_blocks51transformer_blocks0norm3weight [1280] +output_blocks52convbias [1280] +output_blocks52convweight [1280,1280,3,3] +output_blocks60emb_layers1bias [640] +output_blocks60emb_layers1weight [640,1280] +output_blocks60in_layers0bias [1920] +output_blocks60in_layers0weight [1920] +output_blocks60in_layers2bias [640] +output_blocks60in_layers2weight [640,1920,3,3] +output_blocks60out_layers0bias [640] +output_blocks60out_layers0weight [640] +output_blocks60out_layers3bias [640] +output_blocks60out_layers3weight [640,640,3,3] +output_blocks60skip_connectionbias [640] +output_blocks60skip_connectionweight [640,1920,1,1] +output_blocks61normbias [640] +output_blocks61normweight [640] +output_blocks61proj_inbias [640] +output_blocks61proj_inweight [640,640,1,1] +output_blocks61proj_outbias [640] +output_blocks61proj_outweight [640,640,1,1] +output_blocks61transformer_blocks0attn1to_kweight [640,640] +output_blocks61transformer_blocks0attn1to_out0bias [640] +output_blocks61transformer_blocks0attn1to_out0weight [640,640] +output_blocks61transformer_blocks0attn1to_qweight [640,640] +output_blocks61transformer_blocks0attn1to_vweight [640,640] +output_blocks61transformer_blocks0attn2to_kweight [640,768] +output_blocks61transformer_blocks0attn2to_out0bias [640] +output_blocks61transformer_blocks0attn2to_out0weight [640,640] +output_blocks61transformer_blocks0attn2to_qweight [640,640] +output_blocks61transformer_blocks0attn2to_vweight [640,768] +output_blocks61transformer_blocks0ffnet0projbias [5120] +output_blocks61transformer_blocks0ffnet0projweight [5120,640] +output_blocks61transformer_blocks0ffnet2bias [640] +output_blocks61transformer_blocks0ffnet2weight [640,2560] +output_blocks61transformer_blocks0norm1bias [640] +output_blocks61transformer_blocks0norm1weight [640] +output_blocks61transformer_blocks0norm2bias [640] +output_blocks61transformer_blocks0norm2weight [640] +output_blocks61transformer_blocks0norm3bias [640] +output_blocks61transformer_blocks0norm3weight [640] +output_blocks70emb_layers1bias [640] +output_blocks70emb_layers1weight [640,1280] +output_blocks70in_layers0bias [1280] +output_blocks70in_layers0weight [1280] +output_blocks70in_layers2bias [640] +output_blocks70in_layers2weight [640,1280,3,3] +output_blocks70out_layers0bias [640] +output_blocks70out_layers0weight [640] +output_blocks70out_layers3bias [640] +output_blocks70out_layers3weight [640,640,3,3] +output_blocks70skip_connectionbias [640] +output_blocks70skip_connectionweight [640,1280,1,1] +output_blocks71normbias [640] +output_blocks71normweight [640] +output_blocks71proj_inbias [640] +output_blocks71proj_inweight [640,640,1,1] +output_blocks71proj_outbias [640] +output_blocks71proj_outweight [640,640,1,1] +output_blocks71transformer_blocks0attn1to_kweight [640,640] +output_blocks71transformer_blocks0attn1to_out0bias [640] +output_blocks71transformer_blocks0attn1to_out0weight [640,640] +output_blocks71transformer_blocks0attn1to_qweight [640,640] +output_blocks71transformer_blocks0attn1to_vweight [640,640] +output_blocks71transformer_blocks0attn2to_kweight [640,768] +output_blocks71transformer_blocks0attn2to_out0bias [640] +output_blocks71transformer_blocks0attn2to_out0weight [640,640] +output_blocks71transformer_blocks0attn2to_qweight [640,640] +output_blocks71transformer_blocks0attn2to_vweight [640,768] +output_blocks71transformer_blocks0ffnet0projbias [5120] +output_blocks71transformer_blocks0ffnet0projweight [5120,640] +output_blocks71transformer_blocks0ffnet2bias [640] +output_blocks71transformer_blocks0ffnet2weight [640,2560] +output_blocks71transformer_blocks0norm1bias [640] +output_blocks71transformer_blocks0norm1weight [640] +output_blocks71transformer_blocks0norm2bias [640] +output_blocks71transformer_blocks0norm2weight [640] +output_blocks71transformer_blocks0norm3bias [640] +output_blocks71transformer_blocks0norm3weight [640] +output_blocks80emb_layers1bias [640] +output_blocks80emb_layers1weight [640,1280] +output_blocks80in_layers0bias [960] +output_blocks80in_layers0weight [960] +output_blocks80in_layers2bias [640] +output_blocks80in_layers2weight [640,960,3,3] +output_blocks80out_layers0bias [640] +output_blocks80out_layers0weight [640] +output_blocks80out_layers3bias [640] +output_blocks80out_layers3weight [640,640,3,3] +output_blocks80skip_connectionbias [640] +output_blocks80skip_connectionweight [640,960,1,1] +output_blocks81normbias [640] +output_blocks81normweight [640] +output_blocks81proj_inbias [640] +output_blocks81proj_inweight [640,640,1,1] +output_blocks81proj_outbias [640] +output_blocks81proj_outweight [640,640,1,1] +output_blocks81transformer_blocks0attn1to_kweight [640,640] +output_blocks81transformer_blocks0attn1to_out0bias [640] +output_blocks81transformer_blocks0attn1to_out0weight [640,640] +output_blocks81transformer_blocks0attn1to_qweight [640,640] +output_blocks81transformer_blocks0attn1to_vweight [640,640] +output_blocks81transformer_blocks0attn2to_kweight [640,768] +output_blocks81transformer_blocks0attn2to_out0bias [640] +output_blocks81transformer_blocks0attn2to_out0weight [640,640] +output_blocks81transformer_blocks0attn2to_qweight [640,640] +output_blocks81transformer_blocks0attn2to_vweight [640,768] +output_blocks81transformer_blocks0ffnet0projbias [5120] +output_blocks81transformer_blocks0ffnet0projweight [5120,640] +output_blocks81transformer_blocks0ffnet2bias [640] +output_blocks81transformer_blocks0ffnet2weight [640,2560] +output_blocks81transformer_blocks0norm1bias [640] +output_blocks81transformer_blocks0norm1weight [640] +output_blocks81transformer_blocks0norm2bias [640] +output_blocks81transformer_blocks0norm2weight [640] +output_blocks81transformer_blocks0norm3bias [640] +output_blocks81transformer_blocks0norm3weight [640] +output_blocks82convbias [640] +output_blocks82convweight [640,640,3,3] +output_blocks90emb_layers1bias [320] +output_blocks90emb_layers1weight [320,1280] +output_blocks90in_layers0bias [960] +output_blocks90in_layers0weight [960] +output_blocks90in_layers2bias [320] +output_blocks90in_layers2weight [320,960,3,3] +output_blocks90out_layers0bias [320] +output_blocks90out_layers0weight [320] +output_blocks90out_layers3bias [320] +output_blocks90out_layers3weight [320,320,3,3] +output_blocks90skip_connectionbias [320] +output_blocks90skip_connectionweight [320,960,1,1] +output_blocks91normbias [320] +output_blocks91normweight [320] +output_blocks91proj_inbias [320] +output_blocks91proj_inweight [320,320,1,1] +output_blocks91proj_outbias [320] +output_blocks91proj_outweight [320,320,1,1] +output_blocks91transformer_blocks0attn1to_kweight [320,320] +output_blocks91transformer_blocks0attn1to_out0bias [320] +output_blocks91transformer_blocks0attn1to_out0weight [320,320] +output_blocks91transformer_blocks0attn1to_qweight [320,320] +output_blocks91transformer_blocks0attn1to_vweight [320,320] +output_blocks91transformer_blocks0attn2to_kweight [320,768] +output_blocks91transformer_blocks0attn2to_out0bias [320] +output_blocks91transformer_blocks0attn2to_out0weight [320,320] +output_blocks91transformer_blocks0attn2to_qweight [320,320] +output_blocks91transformer_blocks0attn2to_vweight [320,768] +output_blocks91transformer_blocks0ffnet0projbias [2560] +output_blocks91transformer_blocks0ffnet0projweight [2560,320] +output_blocks91transformer_blocks0ffnet2bias [320] +output_blocks91transformer_blocks0ffnet2weight [320,1280] +output_blocks91transformer_blocks0norm1bias [320] +output_blocks91transformer_blocks0norm1weight [320] +output_blocks91transformer_blocks0norm2bias [320] +output_blocks91transformer_blocks0norm2weight [320] +output_blocks91transformer_blocks0norm3bias [320] +output_blocks91transformer_blocks0norm3weight [320] +time_embed0bias [1280] +time_embed0weight [1280,320] +time_embed2bias [1280] +time_embed2weight [1280,1280] \ No newline at end of file diff --git a/scripts/toolkit_gui.py b/scripts/toolkit_gui.py index 54accf5..04dc98d 100644 --- a/scripts/toolkit_gui.py +++ b/scripts/toolkit_gui.py @@ -19,6 +19,7 @@ VAE_PATH = os.path.join(ROOT_PATH, "VAE") MODEL_EXT = [".ckpt", ".pt", ".safetensors"] COMPONENT_EXT = { "UNET-v1": ".unet.pt", + "EMA-UNET-v1": ".unet.pt", "UNET-v2": ".unet-v2.pt", "UNET-v2-Depth": ".unet-v2-d.pt", "VAE-v1": ".vae.pt", @@ -122,7 +123,7 @@ def do_basic_report(details: ToolkitModel, dont_half, keep_ema): out = [f"Model is **{get_size(d.z_total)}**."] if len(d.a_potential) > 1: - out += [f"**Multiple model types identified: {', '.join(d.a_potential)}.**"] + out += [f"Multiple model types identified: **{', '.join(d.a_potential)}**."] out += [f"Model type **{d.a_type}** will be used."] else: out += [f"Model type identified as **{d.a_type}**."] @@ -388,6 +389,7 @@ def do_load(source, precision): error = f"Cannot find {source}!" else: model, _ = load(filename) + fix_model(model, fix_clip=False) loaded = do_analysis(model) loaded.filename = filename if loaded: @@ -424,9 +426,9 @@ def do_select(drop_arch, drop_class, drop_comp): if not loaded: return [gr.update(choices=[], value="") for _ in range(3)] + [gr.update(value="")] - arch_list = sorted(loaded.a_found.keys()) + arch_list = sorted(loaded.a_potential) if not drop_arch in arch_list: - drop_arch = arch_list[0] + drop_arch = loaded.a_type arch = loaded.a_found[drop_arch] @@ -536,15 +538,16 @@ def do_export(drop_arch, drop_class, drop_comp, export_name): prefix = COMPONENTS[comp]["prefix"] extract_component(model, comp, prefixed) - - print(len(model), len(loaded.model)) for k in model: kk = prefix + k if prefixed else k model[k] = loaded.model[kk] + if "EMA" in comp: + fix_ema(model) + folder = COMPONENT_PATH - if "VAE" in loaded.a_type: + if "VAE" in comp: folder = VAE_PATH filename = os.path.join(folder, export_name) @@ -566,6 +569,7 @@ def do_import(drop_arch, drop_class, drop_comp, import_drop, precision): if not error: filename = find_source(import_drop) model, _ = load(filename) + fix_model(model, fix_clip=False) found, _ = inspect_model(model, all=True) if not found or not model: error = "### ERROR: Imported model could not be identified!\n----" @@ -592,10 +596,11 @@ def do_import(drop_arch, drop_class, drop_comp, import_drop, precision): names = [gr.update(), gr.update()] if not error: - # delete the other conflicitng components + # delete the other conflicting components delete_class(loaded.model, drop_arch, drop_class) - - # replace the component + + extract_component(model, choosen) + replace_component(loaded.model, drop_arch, model, choosen) # update analysis diff --git a/toolkit.py b/toolkit.py index f3fa7ef..27d7200 100644 --- a/toolkit.py +++ b/toolkit.py @@ -31,6 +31,11 @@ COMPONENTS = { "source": "UNET-v1-SD.txt", "prefix": "model.diffusion_model." }, + "UNET-v1-EMA": { + "keys": {}, + "source": "UNET-v1-EMA.txt", + "prefix": "model_ema.diffusion_model" + }, "UNET-v1-Inpainting": { "keys": {}, "source": "UNET-v1-Inpainting.txt", @@ -80,6 +85,7 @@ COMPONENTS = { COMPONENT_CLASS = { "UNET-v1-SD": "UNET-v1", + "UNET-v1-EMA": "EMA-UNET-v1", "UNET-v1-Inpainting": "UNET-v1", "UNET-v2-SD": "UNET-v2", "UNET-v2-Depth": "UNET-v2-Depth", @@ -170,6 +176,12 @@ ARCHITECTURES = { "required": [], "prefixed": True }, + "EMA-v1": { + "classes": ["EMA-UNET-v1"], + "optional": OPTIONAL, + "required": [], + "prefixed": True + } } def tensor_size(t): @@ -211,7 +223,6 @@ def get_keys_size(model, keys): z += tensor_size(model[k]) return z - class FakeTensor(): def __init__(self, shape): self.shape = shape @@ -362,7 +373,7 @@ def get_allowed_keys(arch, allowed_classes=None): allowed.update(comp_keys) return allowed -def fix_model(model): +def fix_model(model, fix_clip=False): # fix NAI nonsense nai_keys = { 'cond_stage_model.transformer.embeddings.': 'cond_stage_model.transformer.text_model.embeddings.', @@ -382,10 +393,32 @@ def fix_model(model): # fix merging nonsense i = "cond_stage_model.transformer.text_model.embeddings.position_ids" if i in model: - model[i] = model[i].to(torch.int64) + if fix_clip: + # actually fix the ids + model[i] = torch.Tensor([list(range(77))]).to(torch.int64) + else: + # ensure fp16 looks the same as fp32 + model[i] = model[i].to(torch.int64) return renamed +def fix_ema(model): + # turns UNET-v1-EMA into UNET-v1-SD + # but only when in component form (unprefixed) + + # example keys + # EMA = model_ema.diffusion_modeloutput_blocks91transformer_blocks0norm3weight + # SD = model.diffusion_model.output_blocks9.1.transformer_blocks.0.norm3.weight + + normal = COMPONENTS["UNET-v1-SD"]["keys"] + for k, _ in normal: + kk = k.replace(".", "") + if kk in model: + model[k] = model[kk] + del model[kk] + else: + print(kk) + def compute_metric(model, arch=None): def tensor_metric(t): t = t.to(torch.float16).to(torch.float32) @@ -462,8 +495,6 @@ def load(file): metadata[k] = model[k] model = model['state_dict'] - fix_model(model) - return model, metadata def save(model, metadata, file): @@ -497,7 +528,7 @@ def extract_component(model, component, prefixed=None): allowed = allowed.union(COMPONENTS[component]["keys"]) if prefixed != False: allowed = allowed.union(get_prefixed_keys(component)) - + for k in list(model.keys()): z = tensor_shape(model[k]) if (k, z) in allowed: @@ -519,7 +550,6 @@ def replace_component(target, target_arch, source, source_component): # find out if we should prefix the component is_prefixed = ARCHITECTURES[target_arch]["prefixed"] - print(len(list(source.keys())), len(component_keys)) for k in list(source.keys()): src_z = tensor_shape(source[k]) src_k = k[len(prefix):] if k.startswith(prefix) else k @@ -538,6 +568,7 @@ def delete_class(model, model_arch, component_class): for k in component_keys: if k in keys: del model[k[0]] + keys.remove(k) def log(model, file): keys = [] @@ -547,4 +578,17 @@ def log(model, file): keys.sort() out = "\n".join(keys) with open(file, "w") as f: - f.write(out) \ No newline at end of file + f.write(out) + +if __name__ == '__main__': + r = "/run/media/pul/ssd/stable-diffusion-webui/models/Stable-diffusion/Anything-V3.0.ckpt" + + a, _ = load(r) + + print(list(a.keys())) + + load_components("components") + + extract_component(a, "UNET-v1-EMA") + + print(list(a.keys())) \ No newline at end of file