Kleinhe/SemanticBoost

1from argparse import Namespace2import torch3from motion.dataset.recover_joints import recover_from_ric4from motion.model.cfg_sampler import ClassifierFreeSampleModel5from motion.model_util import create_model_and_diffusion, load_model_wo_clip6import os7import numpy as np8from motion.dataset.recover_smr import *9import json10from motion.double_take import double_take11 12class Predictor(object):13    def __init__(self, **kargs):14        self.path = kargs["path"]15        self.handshake_size = 2016        self.blend_size = 1017 18        args = Namespace()19        with open(self.path["config"], 'r') as f:20            params1 = json.load(f)21        for key, value in params1.items():22            setattr(args, key, value)23 24 25        mode = kargs.get("mode", "cadm")26        if mode == "cadm":27            args.arch = "refined_decoder"28            args.encode_full = 229            args.txt_tokens = 130            args.model_path = self.path["cadm"]31            args.rep = "smr"32        elif mode == "cadm-augment":33            args.arch = "refined_decoder"34            args.encode_full = 235            args.txt_tokens = 136            args.model_path = self.path["cadm-augment"]37            args.rep = "smr"38        elif mode == "mdm":39            args.arch = "trans_enc"40            args.encode_full = 041            args.txt_tokens = 042            args.model_path = self.path["mdm"]43            args.rep = "t2m"44 45        self.skip_steps = kargs.get("skip_steps", 0)46        self.device = kargs.get("device", "cpu")47        self.args = args48        self.rep = args.rep49        self.num_frames = args.num_frames50        self.condition = kargs.get("condition", "text")51        if self.condition == "uncond":52            self.args.guidance_param = 053 54        if self.rep == "t2m":55            extension = ""56        elif self.rep == "smr":57            extension = "_smr"58 59        self.mean = torch.from_numpy(np.load(os.path.join(self.path["dataset_dir"], 'Mean{}.npy'.format(extension)))).to(self.device)60        self.std = torch.from_numpy(np.load(os.path.join(self.path["dataset_dir"], 'Std{}.npy'.format(extension)))).to(self.device)61 62        print(f"Loading checkpoints from...")63        self.model, self.diffusion = create_model_and_diffusion(args, args.control_signal, self.path)64        state_dict = torch.load(self.args.model_path, map_location='cpu')65        try:66            if self.args.ema:67                print("EMA Checkpoints Loading.")68                load_model_wo_clip(self.model, state_dict["ema"])69            else:70                print("Normal Checkpoints Loading.")71                load_model_wo_clip(self.model, state_dict["model"])72        except:73            load_model_wo_clip(self.model, state_dict)74 75        if self.args.guidance_param != 1 and not self.args.unconstrained:76            self.model = ClassifierFreeSampleModel(self.model)   # wrapping model with the classifier-free sampler77        self.model.to(self.device)78        self.model.eval()  # disable random masking79 80    def predict(self,prompt, num_repetitions=1, path=None):81        double_split = prompt.split("|")82        if len(double_split) > 1:83            print("sample mode - double_take long motion")84            sample, step_sizes = double_take(prompt, path, num_repetitions, self.model, self.diffusion, self.handshake_size, 85                                    self.blend_size, self.num_frames, self.args.guidance_param, self.device)86            87            sample = sample.permute(0, 2, 3, 1).float() 88            sample = sample * self.std + self.mean   89            if self.rep == "t2m":90                sample = recover_from_ric(sample, 22)    91                sample = sample.view(-1, *sample.shape[2:]).permute(0, 2, 3, 1) 92            elif self.rep == "smr":93                sample = sample.permute(0, 2, 3, 1)       94        else:95            nframes = prompt.split(",")[0]96            try:97                nframes = int(nframes)98                prompt = prompt.split(",")[1::]99                prompt = ",".join(prompt)100            except:101                nframes = self.num_frames102            103            model_kwargs = {'y':{'text': str(prompt), 'lengths':nframes}}104            if self.args.guidance_param != 1:105                model_kwargs['y']['scale'] = torch.ones(num_repetitions, device=self.device) * self.args.guidance_param106 107            sample_fn = self.diffusion.p_sample_loop108            sample = sample_fn(109                self.model,110                (num_repetitions, self.model.njoints, self.model.nfeats, nframes),111                clip_denoised=False,112                model_kwargs=model_kwargs,113                skip_timesteps=self.skip_steps,  # 0 is the default value - i.e. don't skip any step114                init_image=None,115                progress=True,116                dump_steps=None,117                noise=None,118                const_noise=False119            )120            sample = sample["output"]121            sample = sample.permute(0, 2, 3, 1).float() 122            sample = sample * self.std + self.mean            123 124            if self.rep == "t2m":125                sample = recover_from_ric(sample, 22)    126                sample = sample.view(-1, *sample.shape[2:]).permute(0, 2, 3, 1) 127            elif self.rep == "smr":128                sample = sample.permute(0, 2, 3, 1) 129 130        all_motions = sample.permute(0, 3, 1, 2)131        return all_motions