[content] /hkfs/home/project/hk-project-p0023960/tum_cte0515/Projects/jasmine/jasmine/sample.py (python)
    noise_level: float = 0.0
    max_noise_level: float = 0.7
    noise_buckets: int = 10
    # Tokenizer checkpoint
    tokenizer_dim: int = 512
    tokenizer_ffn_dim: int = 2048
    latent_patch_dim: int = 32
    num_patch_latents: int = 1024
    patch_size: int = 4
    tokenizer_num_blocks: int = 4
    tokenizer_num_heads: int = 8
    # LAM checkpoint
    lam_dim: int = 512
    lam_ffn_dim: int = 2048
    latent_action_dim: int = 32
    num_actions: int = 6
    lam_patch_size: int = 16
    lam_num_blocks: int = 4
    lam_num_heads: int = 8
    use_gt_actions: bool = False
    # Dynamics checkpoint
    dyna_type: str = "maskgit"
    dyna_dim: int = 512
    dyna_ffn_dim: int = 2048
    dyna_num_blocks: int = 6
    dyna_num_heads: int = 8
    param_dtype = jnp.float32
    dtype = jnp.bfloat16
    use_flash_attention: bool = True


args = tyro.cli(Args)

if __name__ == "__main__":
    """
    Dimension keys:
        B: batch size
        T: number of input (conditioning) frames
        N: number of patches per frame
        S: sequence length
        H: height
        W: width
        E: B * (S - 1)
    """
    jax.distributed.initialize()

    rng = jax.random.key(args.seed)

    # --- Load Genie checkpoint ---
    rngs = nnx.Rngs(rng)
    genie = Genie(
        # Tokenizer
        in_dim=args.image_channels,
        tokenizer_dim=args.tokenizer_dim,
        tokenizer_ffn_dim=args.tokenizer_ffn_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,
        lam_ffn_dim=args.lam_ffn_dim,
        latent_action_dim=args.latent_action_dim,
        num_actions=args.num_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=False,
        max_noise_level=args.max_noise_level,
        noise_buckets=args.noise_buckets,
        use_gt_actions=args.use_gt_actions,
        # Dynamics
        dyna_type=args.dyna_type,
        dyna_dim=args.dyna_dim,
        dyna_ffn_dim=args.dyna_ffn_dim,
        dyna_num_blocks=args.dyna_num_blocks,
        dyna_num_heads=args.dyna_num_heads,
        param_dtype=args.param_dtype,
        dtype=args.dtype,
        use_flash_attention=args.use_flash_attention,
        # FIXME (f.srambical): implement spatiotemporal KV caching and set decode=True
        decode=False,
        rngs=rngs,
    )

    # Need to delete lam decoder for checkpoint loading
    if not args.use_gt_actions:
        assert genie.lam is not None
        del genie.lam.decoder

    handler_registry = ocp.handlers.DefaultCheckpointHandlerRegistry()
    handler_registry.add(
        "model_state", ocp.args.PyTreeSave, ocp.handlers.PyTreeCheckpointHandler
    )
    handler_registry.add(
        "model_state", ocp.args.PyTreeRestore, ocp.handlers.PyTreeCheckpointHandler
    )
    checkpoint_options = ocp.CheckpointManagerOptions(
        step_format_fixed_length=6,
    )
    checkpoint_manager = ocp.CheckpointManager(
        args.checkpoint,
        options=checkpoint_options,
        handler_registry=handler_registry,
    )

    dummy_tx = optax.adamw(
        learning_rate=optax.linear_schedule(0.0001, 0.0001, 10000),
        b1=0.9,
        b2=0.9,
        weight_decay=1e-4,
        mu_dtype=args.dtype,
    )
    dummy_optimizer = nnx.ModelAndOptimizer(genie, dummy_tx)

    abstract_optimizer = nnx.eval_shape(lambda: dummy_optimizer)
    abstract_optimizer_state = nnx.state(abstract_optimizer)
    restored = checkpoint_manager.restore(
        checkpoint_manager.latest_step(),
        args=ocp.args.Composite(
            model_state=ocp.args.PyTreeRestore(abstract_optimizer_state),  # type: ignore
        ),
    )
    restored_optimizer_state = restored["model_state"]
    nnx.update(dummy_optimizer, restored_optimizer_state)

    # --- Define sampling function ---
    def _sampling_fn(model: Genie, batch: dict) -> jax.Array:
        """Runs Genie.sample with pre-defined generation hyper-parameters."""
        assert args.dyna_type in [
            "maskgit",
            "causal",
        ], f"Invalid dynamics type: {args.dyna_type}"
        frames, _ = model.sample(
            batch=batch,
            seq_len=args.seq_len,
            noise_level=args.noise_level,