[content] /fast/home/franz.srambical/jafar/sample.py (python)
    args=ocp.args.Composite(
        model_state=ocp.args.StandardRestore(abstract_train_state),
    ),
)
restored_train_state = restored["model_state"]
params = restored_train_state.params

inference_genie = Genie(
    # Tokenizer
    in_dim=args.image_channels,
    tokenizer_dim=args.tokenizer_dim,
    latent_patch_dim=args.latent_patch_dim,
    num_patch_latents=args.num_patch_latents,
    patch_size=args.patch_size,
    tokenizer_num_blocks=args.tokenizer_num_blocks,
    tokenizer_num_heads=args.tokenizer_num_heads,
    # LAM
    lam_dim=args.lam_dim,
    latent_action_dim=args.latent_action_dim,
    num_latent_actions=args.num_latent_actions,
    lam_patch_size=args.lam_patch_size,
    lam_num_blocks=args.lam_num_blocks,
    lam_num_heads=args.lam_num_heads,
    lam_co_train=args.lam_co_train,
    # Dynamics
    dyna_dim=args.dyna_dim,
    dyna_num_blocks=args.dyna_num_blocks,
    dyna_num_heads=args.dyna_num_heads,
    use_maskgit=False,
    param_dtype=args.param_dtype,
    dtype=args.dtype,
    use_flash_attention=args.use_flash_attention,
    is_training=False,
)

rng, _rng = jax.random.split(rng)
inference_variables = inference_genie.init(_rng, dummy_inputs)

cache_struct, _ = flax.core.pop(inference_variables, 'cache')
variables_for_sampling = freeze({
    'params': params,
    'cache': cache_struct
})

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    """
    A clean wrapper to call the JIT-compiled sample_causal method from the Genie module.
    """
    context_frames = video_batch[:, : args.start_frame + 1]
    rng, sample_rng = jax.random.split(rng)
    batch = {
        'videos': context_frames,
        'latent_actions': action_batch,
        'rng': sample_rng,
    }

    @jax.jit
    def sampling_fn(variables, input_batch):
        return inference_genie.apply(
            variables,
            input_batch,
            args.seq_len,
            temperature=args.temperature,
            sample_argmax=args.sample_argmax,
            method=Genie.sample_causal,
        )

    generated_vid = sampling_fn(
        {'params': variables_for_sampling["params"]["params"], 'cache': variables_for_sampling["cache"]},
        batch,
    )
    return generated_vid

def _get_dataloader_iterator():
    array_record_files = [
        os.path.join(args.data_dir, x)
        for x in os.listdir(args.data_dir)
        if x.endswith(".array_record")
    ]
    grain_dataloader = get_dataloader(
        array_record_files,
        args.seq_len,
        # NOTE: We deliberately pass the global batch size
        # The dataloader shards the dataset across all processes
        args.batch_size,
        *image_shape,
        num_workers=0,
        prefetch_buffer_size=1,
        seed=args.seed,
    )
    initial_state = grain_dataloader._create_initial_state()
    grain_iterator = grain.DataLoaderIterator(grain_dataloader, initial_state)
    return grain_iterator

# --- Get video + latent actions ---
grain_iterator = _get_dataloader_iterator()
video_batch = next(grain_iterator)
# video_batch = np.load("overfit_dir/single_sample_corner.npy")

video_batch = video_batch.astype(args.dtype) / 255.0
# Get latent actions for all videos in the batch
batch = dict(videos=video_batch[:,:args.seq_len])
action_batch = inference_genie.apply(variables_for_sampling["params"], batch, False, method=Genie.vq_encode)
action_batch = action_batch.reshape(video_batch.shape[0], args.seq_len - 1, 1)

# --- Sample + evaluate video ---
print("autoreg sampling...")
vid = _autoreg_sample(rng, video_batch, action_batch)
print("autoreg sampling done. calculating ssim and saving video")
gt = video_batch[:, : vid.shape[1]].clip(0, 1).reshape(-1, *video_batch.shape[2:])
recon = vid.clip(0, 1).reshape(-1, *vid.shape[2:])
ssim = pix.ssim(gt[:, args.start_frame + 1 :], recon[:, args.start_frame + 1 :]).mean()