Function sampling_main

sat/sample_video.py:191–424

(args, model_cls)

    print("!!!!!Setted seed to!!!!!!!", seed)

def sampling_main(args, model_cls):


    if isinstance(model_cls, type):
        model = get_model(args, model_cls)
    else:
        model = model_cls

    iteration = load_checkpoint(model, args)  # 20000
    model.eval()

    predictive_mode = False

    if args.input_type == "latent_dir":
        latent_files = get_all_file_from_dir(args.input_file)
        print("show latent_files\n", latent_files[:5])
        decode_latents(model, latent_files)
        return
    elif args.input_type == "dataset":
        data_class = get_obj_from_str(args.data_config["target"])
        create_dataset_function = partial(data_class.create_dataset_function, **args.data_config["params"])

        args.num_workers = 8  # * More workers
        train_data, val_data, test_data = make_loaders(args, create_dataset_function)
        data_iter = val_data  # * Checked, len(val_data) == 12146
        predictive_mode = True
    else:
        raise NotImplementedError
    
    # * Custom configs for sampling
    image_size = args.sampling_video_size  # * Remember set image size in your config
    n_prediction_round = args.n_prediction_round

    T, H, W, C, F = args.sampling_num_frames, image_size[0], image_size[1], args.latent_channels, 8
    num_samples = [1]
    
    force_uc_zero_embeddings = ["txt", "fut_traj", "with_human_drive_token"]


    APPLY_TRAJ = args.apply_traj  # * Default False
    SAVE_RECON = args.save_recon  # * Default True, Turn to false to speed up sampling
    SAVE_GT = args.save_gt
    
    CONCAT_RECON_FOR_DEMO = args.concat_recon_for_demo  # * Default False
    CONCAT_GT_FOR_DEMO = args.concat_gt_for_demo  # * Default False
    N_COND_FRAMES = args.n_cond_frames  # * Default 3

    device = model.device

    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
    out_root = os.path.join(repo_root, "outputs")

    if isinstance(args.base, list):
        cfg_path = args.base[0]
    else:
        cfg_path = args.base

    if "GROUP" not in cfg_path: