Function create_model_and_diffusion_from_args

src/diffusion/script_util.py:7–19

Create model and diffusion from args.

(args)

def create_model_and_diffusion_from_args(args):
    """
    Create model and diffusion from args.
    """
    diffusion = create_gaussian_diffusion(**args_to_dict(args, diffusion_defaults().keys()))
    
    if type(args.channel_mult) is str:
        args.channel_mult = tuple(int(ch_mult) for ch_mult in args.channel_mult.split(","))
    if args.diff_net_type == "unet_small":
        model = TriplaneUNetModelSmall(**args_to_dict(args, diffusion_model_defaults().keys()))
    elif args.diff_net_type == "unet_raw":
        model = TriplaneUNetModelSmallRaw(**args_to_dict(args, diffusion_model_defaults().keys()))
    return model, diffusion


def create_gaussian_diffusion(