    def sample_hf(
        self,
        phone_ids,  # [B, T]
        prompt_ids,  # [8, B, T]
        first_stage_ids,  # [B, T]
        top_k=50,
        top_p=1,
        temperature=1.1,
        first_stage_ids_gt=None,  # [Q, B, T]
        first_stage_ids_gt_end_layer=None,  # 2 to 8
    ):
        """
        phone_ids: [B, T]
        prompt_ids: [8, B, T]
        first_stage_ids: [B, T] result from first quant layer. Should be continuation of prompt_ids
        """
        phone_mask = torch.ones_like(phone_ids, dtype=torch.long)

        assert prompt_ids.shape[-1] >= 5, "prompt_ids should have at least 5 tokens"
        target_ids = torch.cat(
            [prompt_ids, first_stage_ids.expand(prompt_ids.shape[0], -1, -1)], dim=-1
        )
        target_mask = torch.ones_like(target_ids[0], dtype=torch.long)

        if first_stage_ids_gt is not None:
            target_ids[
                :first_stage_ids_gt_end_layer, :, -first_stage_ids_gt.shape[-1] :
            ] = first_stage_ids_gt[:first_stage_ids_gt_end_layer]

        gen_len = first_stage_ids.shape[-1]

        start_qnt_layer = 1
        if first_stage_ids_gt_end_layer is not None:
            start_qnt_layer = first_stage_ids_gt_end_layer
        for qnt_level in range(start_qnt_layer, 8):
            out = self.forward(
                phone_ids=phone_ids,
                phone_mask=phone_mask,
                target_ids=target_ids,
                target_mask=target_mask,
                target_quantization_layer=qnt_level,
                prompt_len=prompt_ids.shape[-1],
            )
            logits = out.logits
            gen_tokens = torch.argmax(logits, dim=-1).reshape(-1)[
                -gen_len:
            ]  # [T], generated tokens in this level

            # overwrite the target_ids with the generated tokens
            target_ids[qnt_level, :, -gen_len:] = gen_tokens

        return target_ids[:, :, -gen_len:]
phone_ids: [B, T]
prompt_ids: [8, B, T]
first_stage_ids: [B, T] result from first quant layer. Should be continuation of prompt_ids
models/tts/valle_v2/valle_nar.py