Upload code/cube3d/inference/engine.py

code/cube3d/inference/engine.py

@@ -0,0 +1,924 @@
+from typing import Optional, Tuple
+
+import torch
+from peft import LoraConfig, get_peft_model, TaskType
+from tqdm import tqdm
+from transformers import CLIPTextModelWithProjection, CLIPTokenizerFast
+
+from cube3d.config import get_mapping_paths
+from cube3d.inference.logits_postprocesses import process_logits, process_logits_assembly
+from cube3d.inference.utils import load_config, load_model_weights, parse_structured, load_model_weights_adaption
+from cube3d.training.process_single_ldr import logits2ldr, load_mappings, logits2flatldrp, logits2flatldrpr, ids2flatldrpr
+from cube3d.model.autoencoder.one_d_autoencoder import OneDAutoEncoder
+from cube3d.model.gpt.dual_stream_roformer import DualStreamRoformer
+from cube3d.model.transformers.cache import Cache
+
+
+class Engine:
+    def __init__(
+        self,
+        config_path: str,
+        gpt_ckpt_path: str,
+        shape_ckpt_path: str,
+        save_gpt_ckpt_path: str,
+        device: torch.device,
+    ):
+        """
+        Initializes the inference engine with the given configuration and checkpoint paths.
+        Args:
+            config_path (str): Path to the configuration file.
+            gpt_ckpt_path (str): Path to the GPT model checkpoint file.
+            shape_ckpt_path (str): Path to the shape model checkpoint file.
+            device (torch.device): The device to run the models on (e.g., 'cpu' or 'cuda').
+        Attributes:
+            cfg (dict): Loaded configuration from the config file.
+            device (torch.device): The device to run the models on.
+            gpt_model (DualStreamRoformer): The GPT model initialized and loaded with weights.
+            shape_model (OneDAutoEncoder): The shape model initialized and loaded with weights.
+            text_model (CLIPTextModelWithProjection): The text model initialized from a pretrained model.
+            text_tokenizer (CLIPTokenizerFast): The tokenizer for the text model.
+            max_new_tokens (int): Maximum number of new tokens for the shape model.
+            min_id (int): Minimum ID for the shape model codes.
+            max_id (int): Maximum ID for the shape model codes.
+        """
+
+        self.cfg = load_config(config_path)
+        self.device = device
+
+        self.gpt_model = DualStreamRoformer(
+            parse_structured(DualStreamRoformer.Config, self.cfg.gpt_model)
+        )
+        # load_model_weights(
+        #     self.gpt_model,
+        #     gpt_ckpt_path,
+        # )
+        self.gpt_model = load_model_weights_adaption(
+            self.gpt_model, 
+            gpt_ckpt_path, 
+            save_gpt_ckpt_path
+        )
+        
+        self.gpt_model = self.gpt_model.eval().to(self.device)
+
+        self.shape_model = OneDAutoEncoder(
+            parse_structured(OneDAutoEncoder.Config, self.cfg.shape_model)
+        )
+        load_model_weights(
+            self.shape_model,
+            shape_ckpt_path,
+        )
+        self.shape_model = self.shape_model.eval().to(self.device)
+
+        # copy vq codebook to gpt
+        with torch.no_grad():
+            codebook = self.shape_model.bottleneck.block.get_codebook()
+            codebook = self.gpt_model.shape_proj(codebook).detach()
+        self.gpt_model.transformer.wte.weight.data[: codebook.shape[0]] = codebook
+        #import ipdb; ipdb.set_trace()
+        self.text_model = CLIPTextModelWithProjection.from_pretrained(
+            self.cfg.text_model_pretrained_model_name_or_path,
+            force_download=False,
+            device_map=self.device,
+        ).eval()
+        self.text_tokenizer = CLIPTokenizerFast.from_pretrained(
+            self.cfg.text_model_pretrained_model_name_or_path,
+            #force_download=False,
+        )
+
+        self.stride = 5
+        self.given = 0
+        self.max_new_tokens = 311*self.stride #self.shape_model.cfg.num_encoder_latents
+        self.min_id = 0
+        self.max_id = self.shape_model.cfg.num_codes
+
+    @torch.inference_mode()
+    def prepare_conditions_with_bbox(
+        self,
+        cond: torch.Tensor,
+        bounding_box_tensor: Optional[torch.Tensor] = None,
+    ):
+        """
+        Prepares condition embeddings by incorporating bounding box information.
+
+        Concatenates bounding box embeddings to the existing condition tensor if the model
+        supports bounding box projection. If no bounding box is provided, uses zero padding.
+
+        Args:
+            cond (torch.Tensor): The input condition embeddings tensor of shape (B, seq_len, dim).
+            bounding_box_xyz (Optional[torch.Tensor], optional): The size of the bounding box
+                as (x, y, z) dimensions represented as a tensor. If None, uses zero padding for
+                bounding box embeddings.
+
+        Returns:
+            torch.Tensor: The condition tensor with bounding box embeddings concatenated along
+                the sequence dimension if bounding box projection is supported, otherwise
+                returns the original condition tensor unchanged.
+        """
+        if not hasattr(self.gpt_model, "bbox_proj"):
+            return cond
+
+        if bounding_box_tensor is None:
+            B = cond.shape[0]
+            bounding_box_tensor = torch.zeros((B, 3), dtype=cond.dtype, device=self.device)
+
+        bbox_emb = self.gpt_model.bbox_proj(bounding_box_tensor).unsqueeze(dim=1)
+        cond = torch.cat([cond, bbox_emb], dim=1)
+        return cond
+
+    @torch.inference_mode()
+    def prepare_conditions_with_bboxs(
+        self,
+        cond: torch.Tensor,
+        bounding_box_tensor: Optional[torch.Tensor] = None,
+    ):
+        """
+        Prepares condition embeddings by incorporating bounding box information.
+
+        Concatenates bounding box embeddings to the existing condition tensor if the model
+        supports bounding box projection. If no bounding box is provided, uses zero padding.
+
+        Args:
+            cond (torch.Tensor): The input condition embeddings tensor of shape (B, seq_len, dim).
+            bounding_box_xyz (Optional[torch.Tensor], optional): The size of the bounding box
+                as (x, y, z) dimensions represented as a tensor. If None, uses zero padding for
+                bounding box embeddings.
+
+        Returns:
+            torch.Tensor: The condition tensor with bounding box embeddings concatenated along
+                the sequence dimension if bounding box projection is supported, otherwise
+                returns the original condition tensor unchanged.
+        """
+        if not hasattr(self.gpt_model, "bbox_proj"):
+            return cond
+
+        if bounding_box_tensor is None:
+            B = cond.shape[0]
+            bounding_box_tensor = torch.zeros((B, 3), dtype=cond.dtype, device=self.device)
+
+        bbox_emb = self.gpt_model.bbox_proj(bounding_box_tensor).unsqueeze(dim=1).expand(cond.shape[0], -1, -1)
+
+        cond = torch.cat([cond, bbox_emb], dim=1)
+        return cond
+
+    @torch.inference_mode()
+    def prepare_inputs(
+        self,
+        prompts: list[str],
+        guidance_scale: float,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Prepares the input embeddings for the model based on the provided prompts and guidance scale.
+        Args:
+            prompts (list[str]): A list of prompt strings to be encoded.
+            guidance_scale (float): A scaling factor for guidance. If greater than 0.0, additional processing is applied.
+            bounding_box_xyz (Optional[Tuple[float]], optional): The size of the bounding box for generation
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            tuple: A tuple containing:
+                - embed (torch.Tensor): The encoded input embeddings.
+                - cond (torch.Tensor): The condition embeddings, which may include unconditional embeddings if guidance_scale is greater than 0.0.
+        """
+
+        prompt_embeds = self.run_clip(prompts)
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            embed = self.encode_input(prompt_embeds, self.gpt_model.shape_bos_id)
+
+        if bounding_box_xyz is not None:
+            cond_bbox = torch.atleast_2d(torch.tensor(bounding_box_xyz)).to(self.device)
+            uncond_bbox = torch.zeros_like(cond_bbox).to(self.device)
+        else:
+            cond_bbox = None
+            uncond_bbox = None
+
+        cond = self.prepare_conditions_with_bbox(prompt_embeds, cond_bbox)
+        if guidance_scale > 0.0:
+            embed = torch.cat([embed, embed], dim=0)
+            uncond_embeds = self.run_clip([""] * len(prompts))
+            uncond = self.prepare_conditions_with_bbox(uncond_embeds, uncond_bbox)
+            cond = torch.cat([cond, uncond], dim=0)
+
+        return embed, cond
+
+    @torch.inference_mode()
+    def run_clip(self, text_inputs):
+        """
+        Processes the given text inputs using a text tokenizer and a text model, and returns the encoded text embeddings.
+        Args:
+            text_inputs (str or List[str]): The input text or list of texts to be processed.
+        Returns:
+            torch.Tensor: The encoded text embeddings.
+        """
+        text_inputs = self.text_tokenizer(
+            text_inputs,
+            max_length=self.text_tokenizer.model_max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )
+        with torch.no_grad():
+            text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
+            # use full precision for text encoder
+            with torch.autocast(device_type=self.device.type, enabled=False):
+                encoded = self.text_model(**text_inputs)
+            if self.gpt_model.cfg.use_pooled_text_embed:
+                embed = encoded.text_embeds.unsqueeze(1)  # [bs, 1, 512]
+            else:
+                embed = encoded.last_hidden_state  # [bs, 77, 512]
+        embed = self.gpt_model.encode_text(embed)
+
+        return embed
+
+    @torch.inference_mode()
+    def encode_input(self, inputs: torch.Tensor, bos: int):
+        """
+        Encodes the beginning of sequence (BOS) token for the given input tensor.
+        Args:
+            inputs (torch.Tensor): The input tensor containing sequences.
+            bos (int): The beginning of sequence token ID.
+        Returns:
+            torch.Tensor: The encoded BOS token embeddings.
+        """
+
+        b = inputs.shape[0]
+        bos_embed = self.gpt_model.encode_token(
+            torch.full(
+                (b, 1),
+                fill_value=bos,
+                dtype=torch.long,
+                device=self.device,
+            )
+        )
+        return bos_embed
+
+    @torch.inference_mode()
+    def run_gpt(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Generates text using a GPT model based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts to generate text from.
+            use_kv_cache (bool): Whether to use key-value caching for faster generation.
+            guidance_scale (float, optional): The scale for guidance during generation. Default is 3.0.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed (deterministic generation). Otherwise, smallest set of tokens with cumulative probability ≥ top_p are kept (stochastic generation).
+            bounding_box_xyz (Optional[Tuple[float]], optional): The size of the bounding box for generation
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            torch.Tensor: A tensor containing the generated token IDs.
+        """
+        embed, cond = self.prepare_inputs(prompts, guidance_scale, bounding_box_xyz)
+
+        output_ids = []
+#         import ipdb; ipdb.set_trace()
+        batch_size, input_seq_len, dim = embed.shape
+        max_seq_len = input_seq_len + self.max_new_tokens
+        embed_buffer = torch.zeros(
+            (batch_size, max_seq_len, dim), dtype=embed.dtype, device=embed.device
+        )
+        embed_buffer[:, :input_seq_len, :].copy_(embed)
+        cond_len = cond.shape[1]
+        kv_cache = None
+        if use_kv_cache:
+            kv_cache = self.gpt_model.init_kv_cache(
+                batch_size,
+                cond_len,
+                self.max_new_tokens + 1,  # +1 for the BOS token
+                torch.bfloat16,
+                embed.device,
+            )
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            for i in tqdm(range(self.max_new_tokens), desc=f"generating"):
+                curr_pos_id = torch.tensor([i], dtype=torch.long, device=embed.device)
+                logits = self.gpt_model(
+                    embed_buffer,
+                    cond,
+                    kv_cache=kv_cache,
+                    curr_pos_id=curr_pos_id if use_kv_cache else None,
+                    decode=(i > 0) if use_kv_cache else False,
+                )
+                if use_kv_cache:
+                    logits = logits[:, 0, ...]
+                else:
+                    logits = logits[:, i, ...]
+
+                logits = logits[..., self.min_id : self.max_id]
+
+                if guidance_scale > 0.0:
+                    logits, uncond_logits = logits.float().chunk(2, dim=0)
+                    # gamma = (
+                    #     guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                    # )
+
+                    logits = (1 + gamma) * logits - gamma * uncond_logits
+                next_id = process_logits(
+                    logits,
+                    top_p=top_p,
+                )
+                output_ids.append(next_id)
+                next_embed = self.gpt_model.encode_token(next_id)
+                if guidance_scale > 0.0:
+                    next_embed = torch.cat([next_embed, next_embed], dim=0)
+                embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+
+            #import ipdb; ipdb.set_trace()
+        return torch.cat(output_ids, dim=1)
+
+    @torch.inference_mode()
+    def run_apt(
+        self,
+        prompts: list[str],
+        inputs_ids: list[torch.Tensor],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Generates text using a GPT model based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts to generate text from.
+            use_kv_cache (bool): Whether to use key-value caching for faster generation.
+            guidance_scale (float, optional): The scale for guidance during generation. Default is 3.0.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed (deterministic generation). Otherwise, smallest set of tokens with cumulative probability ≥ top_p are kept (stochastic generation).
+            bounding_box_xyz (Optional[Tuple[float]], optional): The size of the bounding box for generation
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            torch.Tensor: A tensor containing the generated token IDs.
+        """
+        embed, _ = self.prepare_inputs(prompts, guidance_scale, bounding_box_xyz)
+
+        embed = embed[0][None]
+
+        output_ids = []
+        batch_size, input_seq_len, dim = embed.shape
+        max_seq_len = input_seq_len + self.max_new_tokens
+        embed_buffer = torch.zeros(
+            (batch_size, max_seq_len, dim), dtype=embed.dtype, device=embed.device
+        )
+        embed_buffer[:, :input_seq_len, :].copy_(embed)
+        #cond_len = cond.shape[1]
+        kv_cache = None
+        use_kv_cache = False
+        # if use_kv_cache:
+        #     import ipdb; ipdb.set_trace()
+        #     kv_cache = self.gpt_model.init_kv_cache(
+        #         batch_size,
+        #         cond_len,
+        #         self.max_new_tokens + 1,  # +1 for the BOS token
+        #         torch.bfloat16,
+        #         embed.device,
+        #     )
+
+
+        with torch.no_grad():
+            attention_mask = inputs_ids != -1
+        
+        cut_idx = (attention_mask == False)[:, :, 0].int().argmax(dim=1)
+        dat_id = inputs_ids[:,:,-6].long()
+        dat_id = torch.where(torch.arange(dat_id.shape[1], device=dat_id.device)[None,:] >= cut_idx[:,None], self.gpt_model.dat_num, dat_id)
+        
+        r_id = inputs_ids[:,:,0]
+        x_id = inputs_ids[:,:,-5]
+        y_id = inputs_ids[:,:,-4]
+        z_id = inputs_ids[:,:,-3]
+
+        r_id = torch.where(torch.arange(r_id.shape[1], device=r_id.device)[None,:] >= cut_idx[:,None], self.gpt_model.rot_num, r_id)
+        x_id = torch.where(torch.arange(x_id.shape[1], device=x_id.device)[None,:] >= cut_idx[:,None], self.gpt_model.x_num, x_id)
+        y_id = torch.where(torch.arange(y_id.shape[1], device=y_id.device)[None,:] >= cut_idx[:,None], self.gpt_model.y_num, y_id)
+        z_id = torch.where(torch.arange(z_id.shape[1], device=z_id.device)[None,:] >= cut_idx[:,None], self.gpt_model.z_num, z_id)
+
+
+        inputs_ids[:, :, 0] = r_id.clone()
+        inputs_ids[:, :, -6] = dat_id.clone()
+        inputs_ids[:, :, -5] = x_id.clone()
+        inputs_ids[:, :, -4] = y_id.clone()
+        inputs_ids[:, :, -3] = z_id.clone()
+
+        # xembeds_from_id = self.gpt_model.xte(x_id) 
+        # yembeds_from_id = self.gpt_model.yte(y_id) 
+        # zembeds_from_id = self.gpt_model.zte(z_id)    
+        
+        # embeds_from_id = torch.stack([yembeds_from_id, xembeds_from_id, zembeds_from_id], dim=2)  # [b, 310, 3, 1536]
+        # embeds_from_id = embeds_from_id.view(xembeds_from_id.shape[0], xembeds_from_id.shape[1] * 3, xembeds_from_id.shape[2])  # [b, 930, 1536]
+
+
+        inputs_embeds = self.gpt_model.dte(dat_id)   
+
+        inputs_embeds = self.prepare_conditions_with_bboxs(inputs_embeds, bounding_box_xyz.to(inputs_embeds.device))
+
+        inputs_embeds = embed #bos
+        
+        # #add bos
+        # place_holder = torch.ones_like(inputs_ids[:, 0, 0]).long()  # batch x 1
+        # bos_embed = self.gpt_model.encode_token(place_holder * self.gpt_model.shape_bos_id) #[1, 1536]
+        # embeds_from_id = torch.cat([bos_embed[:, None, :], embeds_from_id], dim=1)
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            for i in tqdm(range(self.max_new_tokens), desc=f"generating"):
+                curr_pos_id = torch.tensor([i], dtype=torch.long, device=embed.device)
+                #import ipdb; ipdb.set_trace()
+                logits = self.gpt_model(
+                    embed = embed_buffer,
+                    cond = inputs_embeds, #cond,
+                    kv_cache=kv_cache,
+                    curr_pos_id=curr_pos_id if use_kv_cache else None,
+                    decode=(i > 0) if use_kv_cache else False,
+                )
+                if use_kv_cache:
+                    logits = logits[:, 0, ...]
+                else:
+                    logits = logits[:, i, ...]
+
+                logits = logits[..., self.min_id : self.max_id]
+
+                # if guidance_scale > 0.0:
+                #     logits, uncond_logits = logits.float().chunk(2, dim=0)
+                #     # gamma = (
+                #     #     guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                #     # )
+                #     gamma = guidance_scale
+                #     logits = (1 + gamma) * logits - gamma * uncond_logits
+
+                next_id = process_logits_assembly(
+                    logits,
+                    top_p=0.9,
+                    pos_id=curr_pos_id,
+                    stride=self.stride
+                )
+
+                #next_embed = process_logits_assembly(logits) #self.gpt_model.encode_token(next_id)
+                #output_ids.append(next_id)
+                #output_ids.append(next_embed)
+                #import ipdb; ipdb.set_trace()
+                #next_embed = self.gpt_model.encode_token(next_id)
+                #next_embed = self.gpt_model.encode_embed(next_embed)
+
+                #output_ids.append(logits)
+
+                if curr_pos_id % self.stride == 0 and self.stride > 3:
+                    #next_embed = self.gpt_model.dte(next_id)
+                    if curr_pos_id<self.given*self.stride:
+                        next_embed = self.gpt_model.dte(dat_id[0, max(0, min(i//self.stride, 309))])[None]
+                    else:
+                        next_embed = self.gpt_model.dte(next_id)
+                elif curr_pos_id % self.stride == 1 and self.stride > 4 :
+                    if curr_pos_id<self.given*self.stride:
+                        next_embed = self.gpt_model.rte(r_id[0, max(0, min(i//self.stride, 309))])[None]
+                    else:
+                        next_embed = self.gpt_model.rte(next_id)
+                elif curr_pos_id % self.stride == (self.stride - 3):
+                    if curr_pos_id<self.given*self.stride:
+                        next_embed = self.gpt_model.yte(y_id[0, max(0, min(i//self.stride, 309))])[None]
+                    else:
+                        next_embed = self.gpt_model.yte(next_id)
+                elif curr_pos_id % self.stride == (self.stride - 2):
+                    if curr_pos_id<self.given*self.stride:
+                        next_embed = self.gpt_model.xte(x_id[0, max(0, min(i//self.stride, 309))])[None]
+                    else:
+                        next_embed = self.gpt_model.xte(next_id)
+                elif curr_pos_id % self.stride == (self.stride - 1):
+                    if curr_pos_id<self.given*self.stride:
+                        next_embed = self.gpt_model.zte(z_id[0, max(0, min(i//self.stride, 309))])[None]
+                    else:
+                        next_embed = self.gpt_model.zte(next_id)
+
+
+                output_ids.append(next_id)
+                # if guidance_scale > 0.0:
+                #     next_embed = torch.cat([next_embed, next_embed], dim=0)
+                embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+
+        #return torch.cat(output_ids, dim=1)
+        return torch.cat(output_ids, dim=0), inputs_ids
+
+    @torch.inference_mode()
+    def run_shape_decode(
+        self,
+        output_ids: torch.Tensor,
+        resolution_base: float = 8.0,
+        chunk_size: int = 100_000,
+    ):
+        """
+        Decodes the shape from the given output IDs and extracts the geometry.
+        Args:
+            output_ids (torch.Tensor): The tensor containing the output IDs.
+            resolution_base (float, optional): The base resolution for geometry extraction. Defaults to 8.43.
+            chunk_size (int, optional): The chunk size for processing. Defaults to 100,000.
+        Returns:
+            tuple: A tuple containing the vertices and faces of the mesh.
+        """
+        shape_ids = (
+            output_ids[:, : self.shape_model.cfg.num_encoder_latents, ...]
+            .clamp_(0, self.shape_model.cfg.num_codes - 1)
+            .view(-1, self.shape_model.cfg.num_encoder_latents)
+        )
+        latents = self.shape_model.decode_indices(shape_ids)
+        mesh_v_f, _ = self.shape_model.extract_geometry(
+            latents,
+            resolution_base=resolution_base,
+            chunk_size=chunk_size,
+            use_warp=True,
+        )
+        return mesh_v_f
+
+    @torch.inference_mode()
+    def t2s(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        resolution_base: float = 8.0,
+        chunk_size: int = 100_000,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Generates a 3D mesh from text prompts using a GPT model and shape decoder.
+        Args:
+            prompts (list[str]): A list of text prompts to guide the generation.
+            use_kv_cache (bool): Whether to use key-value caching for the GPT model.
+            guidance_scale (float, optional): The scale of guidance for the GPT model. Default is 3.0.
+            resolution_base (float, optional): The base resolution for the shape decoder. Default is 8.0.
+            chunk_size (int, optional): The chunk size for processing the shape decoding. Default is 100,000.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed (deterministic generation). Otherwise, smallest set of tokens with cumulative probability ≥ top_p are kept (stochastic generation).
+            bounding_box_xyz (Tuple[float] | None, optional): The size of the bounding box for the generated mesh
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            mesh_v_f: The generated 3D mesh vertices and faces.
+        """
+        #output_ids = self.run_gpt(
+        output_ids = self.run_apt(            
+            prompts, use_kv_cache, guidance_scale, top_p, bounding_box_xyz
+        )
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            mesh_v_f = self.run_shape_decode(output_ids, resolution_base, chunk_size)
+        return mesh_v_f
+
+    @torch.inference_mode()
+    def t2l(
+        self,
+        prompts: list[str],
+        inputs_ids: list[torch.Tensor],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        resolution_base: float = 8.0,
+        chunk_size: int = 100_000,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+        idx: int = 0
+    ):
+        """
+        Generates a ldr file from text prompts using a GPT model and ldr decoder.
+        Args:
+            prompts (list[str]): A list of text prompts to guide the generation.
+            use_kv_cache (bool): Whether to use key-value caching for the GPT model.
+            guidance_scale (float, optional): The scale of guidance for the GPT model. Default is 3.0.
+            resolution_base (float, optional): The base resolution for the shape decoder. Default is 8.0.
+            chunk_size (int, optional): The chunk size for processing the shape decoding. Default is 100,000.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed (deterministic generation). Otherwise, smallest set of tokens with cumulative probability ≥ top_p are kept (stochastic generation).
+            bounding_box_xyz (Tuple[float] | None, optional): The size of the bounding box for the generated mesh
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            ldr: The generated ldr file.
+        """
+        #output_ids = self.run_gpt(
+        output_ids, inputs_ids = self.run_apt(
+            prompts, inputs_ids, use_kv_cache, guidance_scale, top_p, bounding_box_xyz
+        )
+        #import ipdb; ipdb.set_trace()
+        # Use config-based paths (works in both local and HF Space environments)
+        forward_path, inverse_path = get_mapping_paths("subset_self")
+        label_mapping, label_inverse_mapping = load_mappings(forward_path, inverse_path)
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            #mesh_v_f = self.run_shape_decode(output_ids, resolution_base, chunk_size)
+            #ldr = logits2ldr(output_ids.cpu().detach().numpy(), label_inverse_mapping)
+            ldr = ids2flatldrpr(output_ids.float().cpu().detach().numpy(), inputs_ids[0].cpu().detach().numpy(), self.stride, self.given, output_file=f"outputs/sample/test_{self.given}_drp_{idx}_r512_uncond.ldr")
+        return ldr
+
+
+class EngineFast(Engine):
+    def __init__(
+        self,
+        config_path: str,
+        gpt_ckpt_path: str,
+        shape_ckpt_path: str,
+        device: torch.device,
+    ):
+        """
+        Initializes the inference engine with the given configuration and checkpoint paths.
+        Args:
+            config_path (str): Path to the configuration file.
+            gpt_ckpt_path (str): Path to the GPT checkpoint file.
+            shape_ckpt_path (str): Path to the shape checkpoint file.
+            device (torch.device): The device to run the inference on (e.g., CPU or CUDA).
+        """
+
+        assert (
+            device.type == "cuda"
+        ), "EngineFast is only supported on cuda devices, please use Engine on non-cuda devices"
+
+        super().__init__(config_path, gpt_ckpt_path, shape_ckpt_path, device)
+
+        # CUDA Graph params
+        self.graph = torch.cuda.CUDAGraph()
+        self.embed_buffer = torch.Tensor()
+        self.cond_buffer = torch.Tensor()
+        self.logits_buffer = torch.Tensor()
+        self.curr_pos_id = torch.tensor([0], dtype=torch.long, device=self.device)
+        self.kv_cache: list[Cache] = []
+
+        self._warmup_and_capture_graph()
+
+    def _warmup_and_capture_graph(self):
+        """
+        Warms up the model by running a series of forward passes and captures the CUDA graph for efficient execution.
+        This method performs the following steps:
+        1. Prepares the input embeddings and conditions using a warmup prompt.
+        2. Initializes buffers for embeddings and conditions.
+        3. Initializes the key-value cache for the GPT model.
+        4. Runs a series of warmup passes to prefill the model and generate logits.
+        5. Captures the CUDA graph for the model's forward pass to optimize future executions.
+        """
+
+        warmup_prompt = "A cube"
+        embed, cond = self.prepare_inputs([warmup_prompt], guidance_scale=3.0)
+
+        batch_size, input_seq_len, dim = embed.shape
+        max_seq_len = input_seq_len + self.max_new_tokens
+        self.embed_buffer = torch.zeros(
+            (batch_size, max_seq_len, dim), dtype=embed.dtype, device=self.device
+        )
+        self.embed_buffer[:, :input_seq_len, :].copy_(embed)
+
+        self.cond_buffer = torch.empty_like(cond)
+        self.cond_buffer.copy_(cond)
+        cond_len = self.cond_buffer.shape[1]
+
+        # Initialize kv_cache for the first time
+        self.kv_cache = self.gpt_model.init_kv_cache(
+            batch_size,
+            cond_len,
+            self.max_new_tokens + 1,  # +1 for the BOS token
+            #torch.bfloat16,
+            torch.float32,
+            self.device,
+        )
+
+        num_warmup_passes = 10
+
+        #with torch.autocast(self.device.type, dtype=torch.bfloat16):
+        self._set_curr_pos_id(0)
+        _ = self._prefill_and_return_logits()
+
+        for x in range(1, num_warmup_passes):
+            self._set_curr_pos_id(x)
+            self.logits_buffer = self.gpt_model(
+                embed=self.embed_buffer,
+                cond=self.cond_buffer,
+                kv_cache=self.kv_cache,
+                curr_pos_id=self.curr_pos_id,
+                decode=True,
+            )
+
+        side_stream = torch.cuda.Stream(device=self.device)
+        with torch.cuda.graph(self.graph, stream=side_stream):
+            #with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            self.logits_buffer = self.gpt_model(
+                embed=self.embed_buffer,
+                cond=self.cond_buffer,
+                kv_cache=self.kv_cache,
+                curr_pos_id=self.curr_pos_id,
+                decode=True,
+                #decode=False, #? should be false
+            )
+
+    def _reset_kv_cache(self):
+        """
+        Resets the key-value cache by setting all key and value states to zero.
+        This method iterates through each cache in the `kv_cache` attribute and
+        calls the `zero_()` method on both `key_states` and `value_states` to
+        reset them to their initial state.
+        """
+
+        for cache in self.kv_cache:
+            cache.key_states.zero_()
+            cache.value_states.zero_()
+
+    def _prefill_and_return_logits(self) -> torch.Tensor:
+        """
+        Prefills the model's key-value cache and returns the logits.
+        This method resets the key-value cache and then performs a forward pass
+        through the GPT model in eager mode to prefill the logits.
+        Returns:
+            torch.Tensor: The prefilled logits tensor with the first dimension removed.
+        """
+
+        self._reset_kv_cache()
+
+        # Prefill is always eager
+        prefill_logits = self.gpt_model(
+            embed=self.embed_buffer,
+            cond=self.cond_buffer,
+            kv_cache=self.kv_cache,
+            curr_pos_id=self.curr_pos_id,
+            decode=False,
+        )
+
+        return prefill_logits[:, 0, ...]
+
+    def _set_curr_pos_id(self, pos: int):
+        """
+        Set the current position ID.
+        This method updates the `curr_pos_id` attribute with the given position.
+        Args:
+            pos (int): The position ID to set.
+        """
+
+        self.curr_pos_id.copy_(
+            torch.tensor([pos], dtype=torch.long, device=self.device)
+        )
+
+    def run_gpt(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Runs the GPT model to generate text based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts for the GPT model. Only a single prompt is supported.
+            use_kv_cache (bool): Flag indicating whether to use key-value caching. (Currently not used)
+            guidance_scale (float, optional): The scale factor for guidance. Default is 3.0.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed. Otherwise, smallest
+                set of tokens with cumulative probability ≥ top_p are kept.
+            bounding_box_xyz (Tuple[float] | None, optional): The size of the bounding box for the generated mesh
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            torch.Tensor: A tensor containing the generated output token IDs.
+        Raises:
+            AssertionError: If the batch size is greater than 1.
+        """
+
+        embed, cond = self.prepare_inputs(prompts, guidance_scale, bounding_box_xyz)
+        assert len(prompts) == 1, "batch size > 1 not support for EngineFast"
+
+        batch_size, input_seq_len, _ = embed.shape
+        self.embed_buffer.zero_()
+        self.embed_buffer[:, :input_seq_len, :].copy_(embed)
+
+        assert self.cond_buffer.shape == cond.shape
+        self.cond_buffer.copy_(cond)
+
+        output_ids = torch.zeros(
+            (batch_size // 2, self.max_new_tokens), dtype=torch.int, device=self.device
+        )
+
+        with torch.autocast(self.device.type, dtype=torch.bfloat16):
+            self._set_curr_pos_id(0) #?
+
+            logits = self._prefill_and_return_logits()
+
+            logits = logits[..., self.min_id : self.max_id]
+            if guidance_scale > 0.0:
+                logits, uncond_logits = logits.float().chunk(2, dim=0)
+                gamma = guidance_scale
+                logits = (1 + gamma) * logits - gamma * uncond_logits
+            next_id = process_logits(logits, top_p=top_p)
+
+            output_ids[:, 0] = next_id.squeeze()
+            next_embed = self.gpt_model.encode_token(next_id)
+            next_embed = next_embed.repeat(2, 1, 1)
+            self.embed_buffer[:, input_seq_len, :].copy_(next_embed.squeeze(1))
+
+            for i in tqdm(range(1, self.max_new_tokens), desc=f"generating"):
+                self._set_curr_pos_id(i)
+                self.graph.replay()
+
+                logits = self.logits_buffer[:, 0, ...]
+
+                logits = logits[..., self.min_id : self.max_id]
+                if guidance_scale > 0.0:
+                    logits, uncond_logits = logits.float().chunk(2, dim=0)
+                    gamma = (
+                        guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                    )
+                    logits = (1 + gamma) * logits - gamma * uncond_logits
+                next_id = process_logits(logits, top_p=top_p)
+
+                output_ids[:, i] = next_id.squeeze()
+                next_embed = self.gpt_model.encode_token(next_id)
+                next_embed = next_embed.repeat(2, 1, 1)
+                self.embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+        
+        #import ipdb; ipdb.set_trace()
+        return output_ids # torch.Size([1, 1024])
+
+    def run_apt(
+        self,
+        prompts: list[str],
+        use_kv_cache: bool,
+        guidance_scale: float = 3.0,
+        top_p: float = None,
+        bounding_box_xyz: Optional[Tuple[float]] = None,
+    ):
+        """
+        Runs the GPT model to generate text based on the provided prompts.
+        Args:
+            prompts (list[str]): A list of input prompts for the GPT model. Only a single prompt is supported.
+            use_kv_cache (bool): Flag indicating whether to use key-value caching. (Currently not used)
+            guidance_scale (float, optional): The scale factor for guidance. Default is 3.0.
+            top_p (float, optional): The cumulative probability threshold for nucleus sampling.
+                If None, argmax selection is performed. Otherwise, smallest
+                set of tokens with cumulative probability ≥ top_p are kept.
+            bounding_box_xyz (Tuple[float] | None, optional): The size of the bounding box for the generated mesh
+                as (x, y, z) dimensions. Each value must be between 0 and 1.925. If None,
+                uses default bounding box sizing.
+        Returns:
+            torch.Tensor: A tensor containing the generated output token IDs.
+        Raises:
+            AssertionError: If the batch size is greater than 1.
+        """
+
+        embed, cond = self.prepare_inputs(prompts, guidance_scale, bounding_box_xyz)
+        assert len(prompts) == 1, "batch size > 1 not support for EngineFast"
+
+        batch_size, input_seq_len, _ = embed.shape
+        self.embed_buffer.zero_()
+        self.embed_buffer[:, :input_seq_len, :].copy_(embed)
+
+        assert self.cond_buffer.shape == cond.shape
+        self.cond_buffer.copy_(cond) #distinguishing between cond and embed
+
+        # output_ids = torch.zeros(
+        #     (batch_size // 2, self.max_new_tokens), dtype=torch.int, device=self.device
+        # )
+        output_ids = torch.zeros(
+            (batch_size // 2, self.max_new_tokens, embed.shape[2]), dtype=torch.int, device=self.device
+        )
+
+        #with torch.autocast(self.device.type, dtype=torch.bfloat16):
+        self._set_curr_pos_id(0)
+
+        logits = self._prefill_and_return_logits()
+
+        logits = logits[..., self.min_id : self.max_id]
+        if guidance_scale > 0.0:
+            logits, uncond_logits = logits.float().chunk(2, dim=0)
+            gamma = guidance_scale
+            logits = (1 + gamma) * logits - gamma * uncond_logits
+        #next_id = process_logits(logits, top_p=top_p)
+
+        #import ipdb; ipdb.set_trace()
+        #output_ids[:, 0] = next_id.squeeze()
+        next_embed = process_logits_assembly(logits, 0) #self.gpt_model.encode_token(next_id)
+        output_ids[:, 0] = next_embed
+        next_embed = next_embed.repeat(2, 1, 1)
+        self.embed_buffer[:, input_seq_len, :].copy_(next_embed.squeeze(1))
+
+        for i in tqdm(range(1, self.max_new_tokens), desc=f"generating"):
+            self._set_curr_pos_id(i) #position id, indicating the current token position for kv and mask
+            self.graph.replay()
+
+            logits = self.logits_buffer[:, 0, ...]
+
+            logits = logits[..., self.min_id : self.max_id]
+            if guidance_scale > 0.0:
+                logits, uncond_logits = logits.float().chunk(2, dim=0)
+                # gamma = (
+                #     guidance_scale * (self.max_new_tokens - i) / self.max_new_tokens
+                # )
+                gamma = guidance_scale
+                logits = (1 + gamma) * logits - gamma * uncond_logits
+            #next_id = process_logits(logits, top_p=top_p)
+
+            #output_ids[:, i] = next_id.squeeze()
+            
+            next_embed = process_logits_assembly(logits, i) #self.gpt_model.encode_token(next_id)
+            output_ids[:, i] = next_embed
+            next_embed = next_embed.repeat(2, 1, 1)
+            self.embed_buffer[:, i + input_seq_len, :].copy_(next_embed.squeeze(1))
+
+#         import ipdb; ipdb.set_trace()
+        return output_ids # torch.Size([1, 1024])