"""Execution-only CUDA adapter. Frozen routing, weights and FP32 math stay fixed.
CPU image transforms/postprocessing/order and native ORT banks remain unchanged.
"""
import types
import torch

def configure_fp32():
 assert torch.cuda.is_available(),'CUDA required for the GPU evaluation lane'
 torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False;torch.backends.cudnn.benchmark=False;torch.set_float32_matmul_precision('highest');torch.backends.cuda.enable_flash_sdp(False);torch.backends.cuda.enable_mem_efficient_sdp(False);torch.backends.cuda.enable_math_sdp(True)

def cpu_preprocess_cuda_forward(self,images,threshold=.4,return_features=False):
 from surya.common.rfdetr.predictor import IMAGENET_MEAN,IMAGENET_STD
 import torchvision.transforms.functional as TF
 if not images:return []
 tensors=[];sizes=[]
 for image in images:
  image=image.convert('RGB');sizes.append((image.height,image.width));t=TF.to_tensor(image);t=TF.normalize(t,IMAGENET_MEAN,IMAGENET_STD);t=TF.resize(t,(self.resolution,self.resolution));tensors.append(t)
 with torch.inference_mode():
  outputs=self.model(torch.stack(tensors).to('cuda'));features=self._feat.get('f')if return_features else None
  def cpu(x):
   if torch.is_tensor(x):return x.cpu()
   if isinstance(x,dict):return{k:cpu(v)for k,v in x.items()}
   if isinstance(x,list):return[cpu(v)for v in x]
   return x
  outputs=cpu(outputs);features=features.cpu()if features is not None else None;results=self.postprocess(outputs,target_sizes=torch.tensor(sizes));out=[]
  for i,res in enumerate(results):
   keep=res['scores']>threshold;record=dict(boxes=res['boxes'][keep],scores=res['scores'][keep],labels=res['labels'][keep])
   if return_features and features is not None:record['features']=features[i]
   out.append(record)
  return out

def accelerate(engine,reader):
 configure_fp32();reader.model.to('cuda').eval();reader.device=torch.device('cuda')
 assert all(p.dtype==torch.float32 for p in reader.model.parameters())
 engine.execution_devices=dict(layout_model='cpu:unchanged',RTL_student='cuda:FP32',layout_image_transforms='cpu:unchanged',layout_postprocess='cpu:unchanged',learned_reading_order='cpu:unchanged',native_ONNX_banks='CPUExecutionProvider',TF32=False)
 return engine