前端更新, inference.py更新

1) 前端支持剪贴板粘贴图片.
2) 前端支持模型配置.
3) 修改了inference.py的接口.
4) 删除了不必要的文件
This commit is contained in:
三洋三洋
2024-04-17 09:12:07 +00:00
parent 66d4902871
commit 3cebc2eb2a
11 changed files with 181 additions and 105 deletions

View File

@@ -14,7 +14,7 @@ def inference(
model: TexTeller,
tokenizer: RobertaTokenizerFast,
imgs_path: Union[List[str], List[np.ndarray]],
use_cuda: bool,
inf_mode: str = 'cpu',
num_beams: int = 1,
) -> List[str]:
model.eval()
@@ -26,9 +26,8 @@ def inference(
imgs = inference_transform(imgs)
pixel_values = torch.stack(imgs)
if use_cuda:
model = model.to('cuda')
pixel_values = pixel_values.to('cuda')
model = model.to(inf_mode)
pixel_values = pixel_values.to(inf_mode)
generate_config = GenerationConfig(
max_new_tokens=MAX_TOKEN_SIZE,