pagesift_ocr/
device.rs

1//! 设备与模型的选择：跑在哪（[`Device`]）、实际跑在哪（[`DeviceStatus`]）、用哪档模型（[`select_models`]）。
2
3use pagesift_gpu::GpuInfo;
4
5use crate::models::{MEDIUM_DET, MEDIUM_REC, Model, SMALL_DET, SMALL_REC};
6
7/// 想用的设备。配置 `[gpu] mode` 的对应：`off` → `Cpu`，`cuda` → `Gpu { index: 0 }`，
8/// `cuda:N` → `Gpu { index: N }`；`auto` 也是 `Gpu { index: 0 }`，打不开时退回 CPU。
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Device {
11    Cpu,
12    Gpu { index: usize },
13}
14
15/// 实际用的设备，以及想用 GPU 却退回 CPU 的原因。
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct DeviceStatus {
18    /// 调用方要的设备。
19    pub requested: Device,
20    /// 实际用的设备。
21    pub actual: Device,
22    /// `actual` 是 GPU 时它的情况（显存空闲量是打开时的值）。
23    pub gpu: Option<GpuInfo>,
24    /// 永久退回 CPU 的原因；没有退回（包括本来就要 CPU）时为 None。个别页、行显存不够只算那一页的事，
25    /// 不记在这里，见 [`PageRun::fallback`](crate::pipeline::PageRun::fallback)。
26    pub fallback: Option<String>,
27}
28
29impl DeviceStatus {
30    /// 是否想用 GPU 却退回了 CPU（对应警告 `gpu_fallback`）。
31    pub fn fell_back(&self) -> bool {
32        self.fallback.is_some()
33    }
34
35    /// 不想用 GPU。
36    pub(crate) fn cpu() -> DeviceStatus {
37        DeviceStatus {
38            requested: Device::Cpu,
39            actual: Device::Cpu,
40            gpu: None,
41            fallback: None,
42        }
43    }
44
45    /// GPU 打不开，退回 CPU。
46    pub(crate) fn fallback(index: usize, why: String) -> DeviceStatus {
47        DeviceStatus {
48            requested: Device::Gpu { index },
49            actual: Device::Cpu,
50            gpu: None,
51            fallback: Some(why),
52        }
53    }
54
55    /// GPU 打开了。
56    pub(crate) fn on_gpu(index: usize, info: GpuInfo) -> DeviceStatus {
57        let device = Device::Gpu { index };
58        DeviceStatus {
59            requested: device,
60            actual: device,
61            gpu: Some(info),
62            fallback: None,
63        }
64    }
65}
66
67/// 配置 `[ocr] quality`。
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69pub enum Quality {
70    /// 有 GPU 且显存放得下 medium 就用 medium，否则 small。
71    Auto,
72    /// small。
73    Fast,
74    /// medium，不看设备（CPU 上也用，只是慢）。
75    Accurate,
76}
77
78/// 选中的模型。
79#[derive(Debug, Clone, Copy, PartialEq, Eq)]
80pub struct ModelChoice {
81    pub det: Model,
82    pub rec: Model,
83    /// “small” 或 “medium”。
84    pub tier: &'static str,
85}
86
87impl ModelChoice {
88    fn small() -> ModelChoice {
89        ModelChoice {
90            det: SMALL_DET,
91            rec: SMALL_REC,
92            tier: "small",
93        }
94    }
95
96    fn medium() -> ModelChoice {
97        ModelChoice {
98            det: MEDIUM_DET,
99            rec: MEDIUM_REC,
100            tier: "medium",
101        }
102    }
103}
104
105/// medium 在显存里至少要占多少字节（权重加一页的工作内存）。
106pub fn medium_vram_bytes() -> u64 {
107    // [推导] B5 上实测的“真正需要”峰值 2288 MiB（benchmarks/history/1.0/g3，权重 + 在用张量，并行 1 页）。
108    // 这是批大小按显存放开时的值，显存少时批会变小，所以这个门槛偏保守
109    2288 << 20
110}
111
112/// 按质量档和设备选模型。`Quality::Auto` 用 [`pagesift_gpu::probe`] 查空闲显存；查不到（没有 GPU、
113/// `PAGESIFT_GPU=off`）时选 small。
114pub fn select_models(quality: Quality, device: Device) -> ModelChoice {
115    match (quality, device) {
116        (Quality::Fast, _) | (Quality::Auto, Device::Cpu) => ModelChoice::small(),
117        (Quality::Accurate, _) => ModelChoice::medium(),
118        (Quality::Auto, Device::Gpu { index }) => match pagesift_gpu::probe(index) {
119            Ok(info) if info.memory_free >= medium_vram_bytes() => ModelChoice::medium(),
120            _ => ModelChoice::small(),
121        },
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn 按质量档选模型() {
131        let g = Device::Gpu { index: 0 };
132        assert_eq!(select_models(Quality::Fast, g).tier, "small");
133        assert_eq!(select_models(Quality::Auto, Device::Cpu).tier, "small");
134        assert_eq!(select_models(Quality::Accurate, Device::Cpu).tier, "medium");
135    }
136
137    #[test]
138    fn 不存在的显卡退回_cpu_并带原因() {
139        let e = pagesift_gpu::probe(9999).unwrap_err();
140        let s = DeviceStatus::fallback(9999, e);
141        assert_eq!(s.actual, Device::Cpu);
142        assert!(s.fell_back());
143        assert_eq!(s.requested, Device::Gpu { index: 9999 });
144        assert_eq!(
145            select_models(Quality::Auto, Device::Gpu { index: 9999 }).tier,
146            "small"
147        );
148    }
149}
pagesift_ocr/device.rs

pagesift_ocr/
device.rs
