Function load_state_dict

diffportrait360_release/code/inference.py:18–23

(model, ckpt_path, strict=True, map_location="cpu")

from torchvision.utils import save_image

def load_state_dict(model, ckpt_path, strict=True, map_location="cpu"):
    print(f"Loading model state dict from {ckpt_path} ...")
    state_dict = load_from_pretrain(ckpt_path, map_location=map_location)
    state_dict = state_dict.get('state_dict', state_dict)
    model.load_state_dict(state_dict, strict=strict)
    del state_dict  
def get_cond_control(args, batch_data, control_type, device, model=None,batch_size=None, train=True, seg_model=None):
    # Single-control
    control_type = copy.deepcopy(control_type)[0]