Function model_run

app.py:178–291

(prompts, choose_model, infer_mode, seed, n_samples, camera_args=None)

def model_run(prompts, choose_model, infer_mode, seed, n_samples, camera_args=None):
    traj_list = get_traj_list()
    camera_dict = get_camera_dict()

    RT = process_camera(camera_dict, camera_args).reshape(-1,12)
    traj_flow = process_traj(traj_list).transpose(3,0,1,2)

    if choose_model == BASE_MODEL[0]:
        model = model_v1
        noise_shape = [1, 4, 16, 32, 32]
    else:
        model = model_v2
        noise_shape = [1, 4, 16, 40, 64]
    unconditional_guidance_scale = 7.5
    unconditional_guidance_scale_temporal = None

    ddim_steps= 50
    ddim_eta=1.0
    cond_T=800

    if n_samples < 1:
        n_samples = 1
    if n_samples > 4:
        n_samples = 4

    seed_everything(seed)

    if infer_mode == MODE[0]:
        camera_poses = RT
        camera_poses = torch.tensor(camera_poses).float()
        camera_poses = camera_poses.unsqueeze(0)
        trajs = None
        if torch.cuda.is_available():
            camera_poses = camera_poses.cuda()
    elif infer_mode == MODE[1]:
        trajs = traj_flow
        trajs = torch.tensor(trajs).float()
        trajs = trajs.unsqueeze(0)
        camera_poses = None
        if torch.cuda.is_available():
            trajs = trajs.cuda()
    else:
        camera_poses = RT
        trajs = traj_flow
        camera_poses = torch.tensor(camera_poses).float()
        trajs = torch.tensor(trajs).float()
        camera_poses = camera_poses.unsqueeze(0)
        trajs = trajs.unsqueeze(0)
        if torch.cuda.is_available():
            camera_poses = camera_poses.cuda()
            trajs = trajs.cuda()


    ddim_sampler = DDIMSampler(model)
    batch_size = noise_shape[0]
    ## get condition embeddings (support single prompt only)
    if isinstance(prompts, str):
        prompts = [prompts]