from __future__ import annotations import torch try: from .trt_torch import TorchTensorRTEngine except ImportError: from trt_torch import TorchTensorRTEngine class SSFlowTRT(torch.nn.Module): resolution=16; in_channels=8; in_channels_skl=4; out_channels=8; out_channels_skl=4 z_is_global=False; z_skl_is_global=False; global_token_num=1024; global_token_num_skl=1024 def __init__(self,plan): super().__init__(); self.engine=TorchTensorRTEngine(plan); self.device=torch.device('cuda'); self._dummy=torch.nn.Parameter(torch.empty(0,device='cuda'),requires_grad=False) def forward(self,x,x_skl,t,cond,**kwargs): o=self.engine.run({'x':x,'x_skl':x_skl,'timestep':t,'cond':cond});return o['velocity'],o['velocity_skl'] def to(self,*a,**k):return self def eval(self):return self class SSDecoderTRT(torch.nn.Module): def __init__(self,plan):super().__init__();self.engine=TorchTensorRTEngine(plan);self.device=torch.device('cuda');self._dummy=torch.nn.Parameter(torch.empty(0,device='cuda'),requires_grad=False) def forward(self,z,z_skl): o=self.engine.run({'z':z,'z_skl':z_skl});return o['occupancy'],o['occupancy_skl'] def to(self,*a,**k):return self def eval(self):return self class SkinDecoderTRT(torch.nn.Module): skin_feat_channels=4 def __init__(self,plan): super().__init__();self.engine=TorchTensorRTEngine(plan);self.device=torch.device('cuda');self._dummy=torch.nn.Parameter(torch.empty(0,device='cuda'),requires_grad=False) def _cast_for_binding(self,name,t): dt=self.engine._torch_dtype(self.engine.engine.get_tensor_dtype(name)) return t.to(device='cuda',dtype=dt).contiguous() def forward(self,vertex_features,joint_features,parents): e=self.engine;ctx=e.context vf=self._cast_for_binding('vertex_features',vertex_features) jf=self._cast_for_binding('joint_features',joint_features) pp=self._cast_for_binding('parents',parents) feeds={'vertex_features':vf,'joint_features':jf,'parents':pp} for n,t in feeds.items(): ctx.set_input_shape(n,tuple(t.shape));ctx.set_tensor_address(n,int(t.data_ptr())) if hasattr(ctx,'infer_shapes'): try:ctx.infer_shapes() except Exception:pass out_name='skin_weights';dt=e._torch_dtype(e.engine.get_tensor_dtype(out_name)) shape=tuple(ctx.get_tensor_shape(out_name)) if any(int(x)<0 for x in shape):shape=(vf.shape[0],vf.shape[1],jf.shape[1]) y=torch.empty(shape,device='cuda',dtype=dt);ctx.set_tensor_address(out_name,int(y.data_ptr())) ok=ctx.execute_async_v3(stream_handle=torch.cuda.current_stream().cuda_stream) if not ok:raise RuntimeError(f'TensorRT skin decoder execute failed; vf={tuple(vf.shape)}/{vf.dtype} jf={tuple(jf.shape)}/{jf.dtype} parents={tuple(pp.shape)}/{pp.dtype} out={shape}/{dt}') return y def to(self,*a,**k):return self def eval(self):return self class SLatFlowCoreTRTAdapter(torch.nn.Module): def __init__(self,shell_model,plan,max_geo=16384,max_skl=8192,fallback_factory=None): super().__init__();self.native=shell_model;self.engine=TorchTensorRTEngine(plan);self.max_geo=max_geo;self.max_skl=max_skl;self.fallback_factory=fallback_factory;self._fallback_model=None for name in ['resolution','in_channels','in_channels_vert_skin','in_channels_skl','out_channels','out_channels_vert_skin','out_channels_skl','dtype','pe_mode','share_mod','z_is_global','z_skl_is_global','global_token_num','global_token_num_skl','predict_x0','t_scale','t_eps','use_joint_num_cond']: if hasattr(shell_model,name):setattr(self,name,getattr(shell_model,name)) @property def device(self):return self.native.device def _fallback(self,x,x_skl,t,cond,**kwargs): if self._fallback_model is None: if self.fallback_factory is None:raise RuntimeError('SLat TRT outside profile and no fallback') self._fallback_model=self.fallback_factory() return self._fallback_model(x,x_skl,t,cond,**kwargs) def forward(self,x,x_skl,t,cond,joints_num=None,**kwargs): m=self.native;n_geo=int(x.feats.shape[0]);n_skl=int(x_skl.feats.shape[0]) if x.shape[0]!=1 or x_skl.shape[0]!=1 or n_geo>self.max_geo or n_skl>self.max_skl:return self._fallback(x,x_skl,t,cond,joints_num=joints_num,**kwargs) cond=cond.type(m.dtype);feats,feats_skin=x.feats[:,:m.in_channels],x.feats[:,m.in_channels:];x_geo,x_skin=x.replace(feats),x.replace(feats_skin) if getattr(m,'predict_x0',False):xt_skin,xt_skl=feats_skin.clone(),x_skl.feats.clone() joint_vs=joint_skl=None if getattr(m,'use_joint_num_cond',False): joint_vs,joint_skl=m._get_joint_num_emb(joints_num,x.shape[0],x.device);joint_vs=joint_vs.type(m.dtype);joint_skl=joint_skl.type(m.dtype) d={'':x_geo,'_vert_skin':x_skin,'_skl':x_skl};ce={'':None,'_vert_skin':joint_vs,'_skl':joint_skl} for postfix in ['', '_vert_skin','_skl']: d[postfix],d[f't_emb{postfix}'],d[f't_mod{postfix}'],d[f'skips{postfix}']=m.forward_stage(d[postfix],t,postfix,stage='in',cond_emb=ce[postfix]) try: o=self.engine.run({'geo':d[''].feats.unsqueeze(0).contiguous(),'skin':d['_vert_skin'].feats.unsqueeze(0).contiguous(),'skl':d['_skl'].feats.unsqueeze(0).contiguous(),'mod_geo':d['t_mod'].contiguous(),'mod_skin':d['t_mod_vert_skin'].contiguous(),'mod_skl':d['t_mod_skl'].contiguous(),'cond':cond.contiguous()}) d['']=d[''].replace(o['geo_out'][0]);d['_vert_skin']=d['_vert_skin'].replace(o['skin_out'][0]);d['_skl']=d['_skl'].replace(o['skl_out'][0]) except Exception as exc: print(f'[companion-runtime] SLat TRT execution fallback: {exc}',flush=True);return self._fallback(x,x_skl,t,cond,joints_num=joints_num,**kwargs) for postfix in ['', '_vert_skin','_skl']: d[postfix]=m.forward_stage(d[postfix],t,postfix,stage='out',t_emb=d[f't_emb{postfix}'],skips=d[f'skips{postfix}'],original_dtype=x.dtype) if getattr(m,'predict_x0',False): tn=t/m.t_scale;factor=(1/tn.clamp_min(m.t_eps))[:,None] d['_vert_skin']=d['_vert_skin'].replace((d['_vert_skin'].feats-xt_skin)*factor[d['_vert_skin'].coords[:,0]]) d['_skl']=d['_skl'].replace((d['_skl'].feats-xt_skl)*factor[d['_skl'].coords[:,0]]) return x.replace(torch.cat([d[''].feats,d['_vert_skin'].feats],dim=1)),x_skl.replace(d['_skl'].feats) def to(self,*a,**k):self.native.to(*a,**k);return self def eval(self):self.native.eval();return self