| from dataclasses import dataclass |
| from torch import Tensor |
| from typing import Dict, Optional, List, Tuple |
|
|
| import io |
| import os |
| import torch |
|
|
| from ..rig_package.info.asset import Asset |
| from ..model.tokenrig import TokenRig |
|
|
| PORT = 59875 |
| SERVER = f"http://localhost:{PORT}" |
| TMP_CKPT_DIR = "./tmp_ckpt" |
|
|
| BPY_PORT = 59876 |
| BPY_SERVER = f"http://localhost:{BPY_PORT}" |
|
|
| @dataclass |
| class TensorPacket: |
| """make sure stays on cpu""" |
| validate: bool=False |
| know_skeleton: bool=False |
| learned_mesh_cond: Optional[Tensor]=None |
| cond_latents: Optional[Tensor]=None |
| mesh_cond: Optional[Tensor]=None |
| vertices: Optional[Tensor]=None |
| assets: Optional[List[Asset]]=None |
| output_ids: Optional[Tensor]=None |
| start_embed_list: Optional[List[Tensor]]=None |
| start_tokens_list: Optional[List[List[int]]]=None |
|
|
| def to_device(self, device): |
| if self.learned_mesh_cond is not None: |
| self.learned_mesh_cond = self.learned_mesh_cond.to(device) |
| if self.cond_latents is not None: |
| self.cond_latents = self.cond_latents.to(device) |
| if self.mesh_cond is not None: |
| self.mesh_cond = self.mesh_cond.to(device) |
| if self.vertices is not None: |
| self.vertices = self.vertices.to(device) |
| if self.output_ids is not None: |
| self.output_ids = self.output_ids.to(device) |
| if self.start_embed_list is not None: |
| self.start_embed_list = [x.to(device) for x in self.start_embed_list] |
|
|
| @property |
| def B(self): |
| assert self.learned_mesh_cond is not None |
| return self.learned_mesh_cond.shape[0] |
|
|
| def to_bytes(self): |
| return object_to_bytes(self) |
|
|
| @classmethod |
| def from_bytes(cls, bytes) -> 'TensorPacket': |
| return bytes_to_object(bytes) |
|
|
|
|
| def object_to_bytes(t): |
| buffer = io.BytesIO() |
| torch.save(t, buffer) |
| return buffer.getvalue() |
|
|
| def bytes_to_object(b, map_location=None): |
| return torch.load(io.BytesIO(b), weights_only=False, map_location=map_location) |
|
|
| def get_model( |
| ckpt_path: str, |
| hf_path: Optional[str]=None, |
| device='cuda', |
| ) -> TokenRig: |
| model = TokenRig.load_from_system_checkpoint(checkpoint_path=ckpt_path) |
| if hf_path is not None: |
| from transformers import AutoModel |
| a = AutoModel.from_pretrained( |
| hf_path, |
| local_files_only=True, |
| _attn_implementation="flash_attention_2", |
| torch_dtype=torch.bfloat16, |
| ) |
| model.transformer.model.load_state_dict(a.state_dict()) |
|
|
| model = model.to(device) |
| return model |
|
|
