p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] /home/hk-project-p0023960/tum_cte0515/Projects/tmp/jafar/sample.py (python)
def _sampling_wrapper(module, batch):
    return module.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) 
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        params,
        batch
    )
    return generated_vid
[content] /home/hk-project-p0023960/tum_cte0515/Projects/tmp/jafar/genie.py (python)
def _sampling_wrapper(module, batch):
    return module.sample(batch, args.seq_len, args.maskgit_steps, args.temperature, args.sample_argmax)

# --- Define autoregressive sampling loop ---
def _autoreg_sample(rng, video_batch, action_batch):
    vid = video_batch[:, : args.start_frame + 1]
    sampling_fn = jax.jit(nn.apply(_sampling_wrapper, genie)) 
    rng, _rng = jax.random.split(rng)
    batch = dict(videos=vid, latent_actions=action_batch, rng=_rng)
    generated_vid = sampling_fn(
        params,
        batch
    )
    return generated_vid
[content] /home/hk-project-p0023960/tum_cte0515/Projects/tmp/jafar/genie.py (python)

    @nn.compact
    def sample(
        self,
        batch: Dict[str, Any],
        steps: int = 25,
        temperature: int = 1,
        sample_argmax: bool = False,
    ) -> Any:
        # --- Encode videos and actions ---
        tokenizer_out = self.tokenizer.vq_encode(batch["videos"], training=False)
        token_idxs = tokenizer_out["indices"]
        new_frame_idxs = jnp.zeros_like(token_idxs)[:, 0]
        action_tokens = self.lam.vq.get_codes(batch["latent_actions"])

        # --- Initialize MaskGIT ---
        init_mask = jnp.ones_like(token_idxs, dtype=bool)[:, 0]
        init_carry = (
            batch["rng"],
            new_frame_idxs,
            init_mask,
            token_idxs,
            action_tokens,
        )
        MaskGITLoop = nn.scan(
            MaskGITStep,
            variable_broadcast="params",
            split_rngs={"params": False},
            in_axes=0,
            out_axes=0,
            length=steps,
        )

        # --- Run MaskGIT loop ---
        loop_fn = MaskGITLoop(
            dynamics=self.dynamics,
            tokenizer=self.tokenizer,
            temperature=temperature,
            sample_argmax=sample_argmax,
            steps=steps,
        )
        final_carry, _ = loop_fn(init_carry, jnp.arange(steps))
        new_frame_idxs = final_carry[1]
        new_frame_pixels = self.tokenizer.decode(
            jnp.expand_dims(new_frame_idxs, 1),
            video_hw=batch["videos"].shape[2:4],
        )
        return new_frame_pixels
[content] /home/hk-project-p0023960/tum_cte0515/Projects/tmp/jafar/genie.py (python)


class MaskGITStep(nn.Module):
    dynamics: nn.Module
    tokenizer: nn.Module
    temperature: float
    sample_argmax: bool
    steps: int

    @nn.compact
    def __call__(self, carry, x):
        rng, token_idxs, mask, action_tokens = carry
        step = x
        N = token_idxs.shape[2]

        # --- Construct + encode video ---
        vid_embed = self.dynamics.patch_embed(token_idxs) # (B, S, N, D)
        mask_token = self.dynamics.mask_token  # (1, 1, 1, D,)
        mask_expanded = mask[..., None] # (B, S, N, 1) 
        vid_embed = jnp.where(mask_expanded, mask_token, vid_embed)

        # --- Predict transition ---
        act_embed = self.dynamics.action_up(action_tokens)
        vid_embed += jnp.pad(act_embed, ((0, 0), (1, 0), (0, 0), (0, 0)))
        unmasked_ratio = jnp.cos(jnp.pi * (step + 1) / (self.steps * 2))
        step_temp = self.temperature * (1.0 - unmasked_ratio)
        final_logits = self.dynamics.dynamics(vid_embed) / step_temp

        # --- Sample new tokens for final frame ---
        if self.sample_argmax:
            sampled_token_idxs = jnp.argmax(final_logits, axis=-1)
        else:
            rng, _rng = jax.random.split(rng)
            sampled_token_idxs = jnp.where(
                step == self.steps - 1,
                jnp.argmax(final_logits, axis=-1),
                jax.random.categorical(_rng, final_logits),
            )
        gather_fn = jax.vmap(jax.vmap(jax.vmap(lambda x, y: x[y])))
        final_token_probs = gather_fn(jax.nn.softmax(final_logits), sampled_token_idxs)
        final_token_probs += ~mask
        # Update masked tokens only
        token_idxs = jnp.where(mask, sampled_token_idxs, token_idxs)

        # --- Update mask ---
        num_unmasked_tokens = jnp.round(N * (1.0 - unmasked_ratio)).astype(int)
        idx_mask = jnp.arange(final_token_probs.shape[-1]) > num_unmasked_tokens
        sorted_idxs = jnp.argsort(final_token_probs, axis=-1, descending=True)
        mask_update_fn = jax.vmap(lambda msk, ids: msk.at[ids].set(idx_mask))
        new_mask = mask_update_fn(mask, sorted_idxs)

        new_carry = (rng, token_idxs, new_mask, action_tokens)
        return new_carry, None