+h 8SSKrSSKrSSKJr SSKJr SSSSSS S S S S SSSSSS.rSSSS.rSrSr \ S:Xab\R"SS9r \ RS\SSS 9 \ RS!\SS"S 9 \ R5r\ "\R"\R$5 gg)#N) load_file)Kandinsky3UNetztime_embedding.linear_1ztime_embedding.linear_2conv_in conv_norm_outconv_out down_blocks up_blocksz"encoder_hid_proj.projection_linearz encoder_hid_proj.projection_normadd_time_conditionto_qto_kto_vzto_out.0z attentions.0)zto_time_embed.1zto_time_embed.3in_layerz out_layer.0z out_layer.2 down_samples up_samplesprojection_lin projection_lnfeature_poolingto_queryto_keyto_value output_layerself_attention_blockz resnets_in.*)z attentions.*z resnets_out.*)zresnet_attn_blocks.*.0zresnet_attn_blocks.*.1zresnet_attn_blocks.*.2c |0nUGH2nUn[R5HupEURXE5nM [R5HupgSn[R"USUS35(dM)U(aM2[ UR UR S5S5SR S5S5n [U[5(a XS-n USnOU n URS[U 55nURS[U 55nURXE5nS nM XX'GM5 U$) a6 Args: Convert the state dict of a U-Net model to match the key format expected by Kandinsky3UNet model. unet_model (torch.nn.Module): The original U-Net model. unet_kandi3_model (torch.nn.Module): The Kandinsky3UNet model to match keys with. Returns: OrderedDict: The converted state dictionary. Fz*.z.*.rr*T) MAPPINGitemsreplace DYNAMIC_MAPfnmatchintsplit isinstancetuplestr) unet_state_dictconverted_state_dictkeynew_keypattern new_pattern dyn_patterndyn_new_pattern has_matchedstarnew_stars p/home/james-whalen/.local/lib/python3.13/site-packages/diffusers/pipelines/kandinsky3/convert_kandinsky3_unet.pyconvert_state_dictr4$s-$+MMO Goog;G%4-8,=,=,? (KKw"[M(<==kk7==):):3)?)BCBGMMcRSTUVou55#b&99H&5a&8O#H%--c3t9=-55c3x=I !//'?" -@")8(<%-0 c[U5n0n[U5n[U5nURU5 UR U5 [ SU35 g)NzConverted model saved to )rr4rload_state_dictsave_pretrainedprint) model_path output_pathr(configr)unets r3mainr>KsW +OF.o> & !D-.% %k] 34r5__main__z4Convert U-Net PyTorch model to Kandinsky3UNet format) descriptionz --model_pathTz(Path to the original U-Net PyTorch model)typerequiredhelpz --output_pathz Path to save the converted model)argparser"safetensors.torchr diffusersrrr!r4r>__name__ArgumentParserparser add_argumentr' parse_argsargsr:r;r5r3rNs'$10"!:7+* &-1- $ N5" z  $ $1g hF S4Fpq cDGij    D$**+ r5