- Cross-platform reproducibility (`--rng cuda`, consistent with the`stable-diffusion-webui GPU RNG`)

- Embedds generation parameters into png output as webui-compatible text string

- Supported platforms

1.0 corresponds to full destruction of information in init image

-H, --height H image height, in pixel space (default: 512)

-W, --width W image width, in pixel space (default: 512)

--sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2}

--sampling-method {euler, euler_a, heun, dpm++2m, dpm++2mv2, lcm}

sampling method (default: "euler_a")

--steps STEPS number of sample steps (default: 20)

--rng {std_default, cuda} RNG (default: cuda)

-[stable-diffusion-stability-ai](https://github.com/Stability-AI/stablediffusion)

-[stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui)

-[k-diffusion](https://github.com/crowsonkb/k-diffusion)

-[latent-consistency-model](https://github.com/luosiallen/latent-consistency-model)

"dpm2",

"dpm++2s_a",

"dpm++2m",

"dpm++2mv2"};

"dpm++2mv2",

"lcm",

};

constchar* schedule_str[] = {

"default",

"discrete",

"karras"};

};

structOption {

int n_threads = -1;

printf(" 1.0 corresponds to full destruction of information in init image\n");

printf(" -H, --height H image height, in pixel space (default: 512)\n");

printf(" -W, --width W image width, in pixel space (default: 512)\n");

printf(" --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2}\n");

printf(" --sampling-method {euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, lcm}\n");

printf(" sampling method (default:\"euler_a\")\n");

printf(" --steps STEPS number of sample steps (default: 20)\n");

printf(" --rng {std_default, cuda} RNG (default: cuda)\n");

importos.pathasosp

fromsafetensors.torchimportload_file,save_file

unet_conversion_map= [

("time_embed.0.weight","time_embedding.linear_1.weight"),

("time_embed.0.bias","time_embedding.linear_1.bias"),

("time_embed.2.weight","time_embedding.linear_2.weight"),

("time_embed.2.bias","time_embedding.linear_2.bias"),

("input_blocks.0.0.weight","conv_in.weight"),

("input_blocks.0.0.bias","conv_in.bias"),

("out.0.weight","conv_norm_out.weight"),

("out.0.bias","conv_norm_out.bias"),

("out.2.weight","conv_out.weight"),

("out.2.bias","conv_out.bias"),

]

unet_conversion_map_resnet= [

("in_layers.0","norm1"),

("in_layers.2","conv1"),

("out_layers.0","norm2"),

("out_layers.3","conv2"),

("emb_layers.1","time_emb_proj"),

("skip_connection","conv_shortcut"),

]

unet_conversion_map_layer= []

foriinrange(4):

forjinrange(2):

unet_conversion_map_layer.append((sd_down_res_prefix,hf_down_res_prefix))

ifi<3:

unet_conversion_map_layer.append((sd_down_atn_prefix,hf_down_atn_prefix))

forjinrange(3):

unet_conversion_map_layer.append((sd_up_res_prefix,hf_up_res_prefix))

ifi>0:

unet_conversion_map_layer.append((sd_up_atn_prefix,hf_up_atn_prefix))

ifi<3:

unet_conversion_map_layer.append((sd_downsample_prefix,hf_downsample_prefix))

unet_conversion_map_layer.append((sd_upsample_prefix,hf_upsample_prefix))

unet_conversion_map_layer.append((sd_mid_atn_prefix,hf_mid_atn_prefix))

forjinrange(2):

unet_conversion_map_layer.append((sd_mid_res_prefix,hf_mid_res_prefix))

defconvert_unet_state_dict(unet_state_dict):

mapping= {k:kforkinunet_state_dict.keys()}

forsd_name,hf_nameinunet_conversion_map:

mapping[hf_name]=sd_name

fork,vinmapping.items():

if"resnets"ink:

forsd_part,hf_partinunet_conversion_map_resnet:

v=v.replace(hf_part,sd_part)

mapping[k]=v

fork,vinmapping.items():

forsd_part,hf_partinunet_conversion_map_layer:

v=v.replace(hf_part,sd_part)

mapping[k]=v

new_state_dict= {v:unet_state_dict[k]fork,vinmapping.items()}

vae_conversion_map= [

("nin_shortcut","conv_shortcut"),

("norm_out","conv_norm_out"),

("mid.attn_1.","mid_block.attentions.0."),

]

foriinrange(4):

forjinrange(2):

vae_conversion_map.append((sd_down_prefix,hf_down_prefix))

ifi<3:

vae_conversion_map.append((sd_downsample_prefix,hf_downsample_prefix))

vae_conversion_map.append((sd_upsample_prefix,hf_upsample_prefix))

forjinrange(3):

vae_conversion_map.append((sd_up_prefix,hf_up_prefix))

foriinrange(2):

vae_conversion_map.append((sd_mid_res_prefix,hf_mid_res_prefix))

vae_conversion_map_attn= [

("norm.","group_norm."),

("q.","query."),

("k.","key."),

("v.","value."),

("proj_out.","proj_attn."),

]

defreshape_weight_for_sd(w):

returnw.reshape(*w.shape,1,1)

defconvert_vae_state_dict(vae_state_dict):

mapping= {k:kforkinvae_state_dict.keys()}

fork,vinmapping.items():

forsd_part,hf_partinvae_conversion_map:

v=v.replace(hf_part,sd_part)

mapping[k]=v

fork,vinmapping.items():

if"attentions"ink:

forsd_part,hf_partinvae_conversion_map_attn:

v=v.replace(hf_part,sd_part)

mapping[k]=v

new_state_dict= {v:vae_state_dict[k]fork,vinmapping.items()}

weights_to_convert= ["q","k","v","proj_out"]

fork,vinnew_state_dict.items():

forweight_nameinweights_to_convert:

iff"mid.attn_1.{weight_name}.weight"ink:

print(f"Reshaping{k} for SD format")

new_state_dict[k]=reshape_weight_for_sd(v)

textenc_conversion_lst= [

("resblocks.","text_model.encoder.layers."),

("ln_1","layer_norm1"),

("ln_2","layer_norm2"),

(".c_fc.",".fc1."),

(".c_proj.",".fc2."),

(".attn",".self_attn"),

("ln_final.","transformer.text_model.final_layer_norm."),

("token_embedding.weight","transformer.text_model.embeddings.token_embedding.weight"),

("positional_embedding","transformer.text_model.embeddings.position_embedding.weight"),

]

protected= {re.escape(x[1]):x[0]forxintextenc_conversion_lst}

textenc_pattern=re.compile("|".join(protected.keys()))

code2idx= {"q":0,"k":1,"v":2}

defconvert_text_enc_state_dict_v20(text_enc_dict):

new_state_dict= {}

capture_qkv_weight= {}

capture_qkv_bias= {}

fork,vintext_enc_dict.items():

if (

k.endswith(".self_attn.q_proj.weight")

ork.endswith(".self_attn.k_proj.weight")

ork.endswith(".self_attn.v_proj.weight")

):

k_pre=k[:-len(".q_proj.weight")]

k_code=k[-len("q_proj.weight")]

ifk_prenotincapture_qkv_weight:

capture_qkv_weight[k_pre]= [None,None,None]

capture_qkv_weight[k_pre][code2idx[k_code]]=v

if (

k.endswith(".self_attn.q_proj.bias")

ork.endswith(".self_attn.k_proj.bias")

ork.endswith(".self_attn.v_proj.bias")

):

k_pre=k[:-len(".q_proj.bias")]

k_code=k[-len("q_proj.bias")]

ifk_prenotincapture_qkv_bias:

capture_qkv_bias[k_pre]= [None,None,None]

capture_qkv_bias[k_pre][code2idx[k_code]]=v

relabelled_key=textenc_pattern.sub(lambdam:protected[re.escape(m.group(0))],k)

new_state_dict[relabelled_key]=v

fork_pre,tensorsincapture_qkv_weight.items():

ifNoneintensors:

raiseException("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")

relabelled_key=textenc_pattern.sub(lambdam:protected[re.escape(m.group(0))],k_pre)

new_state_dict[relabelled_key+".in_proj_weight"]=torch.cat(tensors)

fork_pre,tensorsincapture_qkv_bias.items():

ifNoneintensors:

raiseException("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")

relabelled_key=textenc_pattern.sub(lambdam:protected[re.escape(m.group(0))],k_pre)

new_state_dict[relabelled_key+".in_proj_bias"]=torch.cat(tensors)

defconvert_text_enc_state_dict(text_enc_dict):

if__name__=="__main__":

parser=argparse.ArgumentParser()

parser.add_argument("--model_path",default=None,type=str,required=True,help="Path to the model to convert.")

parser.add_argument("--checkpoint_path",default=None,type=str,required=True,help="Path to the output model.")

parser.add_argument("--half",action="store_true",help="Save weights in half precision.")

parser.add_argument(

"--use_safetensors",action="store_true",help="Save weights use safetensors, default is ckpt."

)

args=parser.parse_args()

assertargs.model_pathisnotNone,"Must provide a model path!"

assertargs.checkpoint_pathisnotNone,"Must provide a checkpoint path!"

unet_path=osp.join(args.model_path,"unet","diffusion_pytorch_model.safetensors")

vae_path=osp.join(args.model_path,"vae","diffusion_pytorch_model.safetensors")

text_enc_path=osp.join(args.model_path,"text_encoder","model.safetensors")

ifosp.exists(unet_path):

unet_state_dict=load_file(unet_path,device="cpu")

else:

unet_path=osp.join(args.model_path,"unet","diffusion_pytorch_model.bin")

unet_state_dict=torch.load(unet_path,map_location="cpu")

ifosp.exists(vae_path):

vae_state_dict=load_file(vae_path,device="cpu")

else:

vae_path=osp.join(args.model_path,"vae","diffusion_pytorch_model.bin")

vae_state_dict=torch.load(vae_path,map_location="cpu")

ifosp.exists(text_enc_path):

text_enc_dict=load_file(text_enc_path,device="cpu")

else:

text_enc_path=osp.join(args.model_path,"text_encoder","pytorch_model.bin")

text_enc_dict=torch.load(text_enc_path,map_location="cpu")

unet_state_dict=convert_unet_state_dict(unet_state_dict)

unet_state_dict= {"model.diffusion_model."+k:vfork,vinunet_state_dict.items()}

vae_state_dict=convert_vae_state_dict(vae_state_dict)

vae_state_dict= {"first_stage_model."+k:vfork,vinvae_state_dict.items()}

ifis_v20_model:

text_enc_dict= {"transformer."+k:vfork,vintext_enc_dict.items()}

text_enc_dict=convert_text_enc_state_dict_v20(text_enc_dict)

text_enc_dict= {"cond_stage_model.model."+k:vfork,vintext_enc_dict.items()}

else:

text_enc_dict=convert_text_enc_state_dict(text_enc_dict)

text_enc_dict= {"cond_stage_model.transformer."+k:vfork,vintext_enc_dict.items()}

state_dict= {**unet_state_dict,**vae_state_dict,**text_enc_dict}

ifargs.half:

state_dict= {k:v.half()fork,vinstate_dict.items()}

ifargs.use_safetensors:

save_file(state_dict,args.checkpoint_path)

else:

state_dict= {"state_dict":state_dict}

torch.save(state_dict,args.checkpoint_path)