feat: font detection using yuzumaker.FontDetection

koharu-ml/Cargo.toml

@@ -78,3 +78,7 @@ path = "bin/llm.rs"
 [[bin]]
 name = "manga-ocr"
 path = "bin/manga-ocr.rs"
+
+[[bin]]
+name = "font-detect"
+path = "bin/font-detect.rs"

koharu-ml/README.md

@@ -7,6 +7,7 @@ Model wrappers and CLI tools for the Koharu app. Each module lazily downloads it
 - `manga_ocr`: encoder/decoder OCR pipeline that reads cropped text regions.
 - `lama`: LaMa inpainting with tiled blending to remove text using a mask.
 - `llm`: quantized GGUF loader (Llama or Qwen2) using candle with chat-style prompting and generation controls.
+- `font_detect`: Candle ResNet50 that reproduces YuzuMarker.FontDetection (CJK font/style classifier).
 
 ## CLI tools
 ```bash
@@ -14,8 +15,34 @@ cargo run -p koharu-models --bin comic-text-detector -- --input page.png --outpu
 cargo run -p koharu-models --bin manga-ocr -- --input bubble.png
 cargo run -p koharu-models --bin lama -- --input page.png --mask mask.png --output filled.png
 cargo run -p koharu-models --bin llm -- --prompt "konnichiwa" --model vntl-llama3-8b-v2
+cargo run -p koharu-models --bin font-detect -- --input bubble.png --top-k 5 --model resnet50
 ```
 
+### Font detection weights
+
+The original checkpoints are published at [gyrojeff/YuzuMarker.FontDetection](https://huggingface.co/gyrojeff/YuzuMarker.FontDetection) in PyTorch Lightning format. Candle needs `safetensors`, so convert once and point the runtime to the file:
+
+```bash
+python scripts/convert_font_detection.py \
+  --checkpoint name=4x-epoch=84-step=1649340.ckpt \
+  --output ~/.cache/Koharu/models/yuzumarker-font-detection.safetensors
+```
+
+Set `KOHARU_FONT_DETECTION_WEIGHTS` to override the path if desired. The loader will look for the safetensors file in `~/.cache/Koharu/models/` by default.
+Supported backbones: `resnet18`, `resnet34`, `resnet50` (default), `resnet101`, `deepfont` (pads missing regression outputs).
+
+### Font labels (names)
+
+The original demo ships `font_demo_cache.bin` (Python pickle) that maps class ids to font paths. Convert it to JSON so Rust can read the names:
+
+```bash
+python scripts/convert_font_labels.py \
+  --input font_demo_cache.bin \
+  --output ~/.cache/Koharu/models/yuzumarker-font-labels.json
+```
+
+Set `KOHARU_FONT_DETECTION_LABELS` or pass `--labels` to the CLI to override the path. When labels are present, CLI output includes font names alongside ids.
+
 Feature `cuda` enables the CUDA execution provider for ONNX Runtime and candle; without it the models fall back to CPU.
 
 ## License

koharu-ml/bin/font-detect.rs

@@ -0,0 +1,74 @@
+use std::path::PathBuf;
+
+use anyhow::Result;
+use clap::Parser;
+use koharu_ml::font_detector::{FontDetector, ModelKind, TextDirection};
+
+#[derive(Parser, Debug)]
+#[command(
+    author,
+    version,
+    about = "Run YuzuMarker.FontDetection (Candle) on an image"
+)]
+struct Args {
+    /// Path to the input image.
+    #[arg(short, long)]
+    input: PathBuf,
+    /// Number of top font classes to return.
+    #[arg(short = 'k', long, default_value_t = 5)]
+    top_k: usize,
+    /// Force CPU even if GPU is available.
+    #[arg(long)]
+    cpu: bool,
+    /// Backbone architecture (must match the converted checkpoint).
+    #[arg(long, default_value = "resnet50", value_enum)]
+    model: ModelKind,
+}
+
+#[tokio::main]
+async fn main() -> Result<()> {
+    let args = Args::parse();
+    let detector = FontDetector::load_with_kind(args.cpu, args.model).await?;
+    let image = image::open(&args.input)?;
+    let start = std::time::Instant::now();
+    let result = detector.inference(&[image], args.top_k)?;
+    let Some(pred) = result.first() else {
+        return Ok(());
+    };
+    println!("Inference took: {:.2?}", start.elapsed());
+
+    println!("Top fonts:");
+    for (idx, prob) in &pred.top_fonts {
+        let name = pred.named_fonts.iter().find(|f| f.index == *idx);
+        if let Some(named) = name {
+            if let Some(language) = &named.language {
+                println!("  #{idx} ({} | lang={language}): {prob:.4}", named.name);
+            } else {
+                println!("  #{idx} ({}): {prob:.4}", named.name);
+            }
+        } else {
+            println!("  #{idx}: {prob:.4}");
+        }
+    }
+    println!(
+        "Direction: {:?}",
+        match pred.direction {
+            TextDirection::Horizontal => "horizontal",
+            TextDirection::Vertical => "vertical",
+        }
+    );
+    println!(
+        "Text color: rgb({},{},{})",
+        pred.text_color[0], pred.text_color[1], pred.text_color[2]
+    );
+    println!(
+        "Stroke color: rgb({},{},{}) width_px={:.2}",
+        pred.stroke_color[0], pred.stroke_color[1], pred.stroke_color[2], pred.stroke_width_px
+    );
+    println!(
+        "Font size (px): {:.2} | line height: {:.2} | angle: {:.1}°",
+        pred.font_size_px, pred.line_height, pred.angle_deg
+    );
+
+    Ok(())
+}

koharu-ml/src/font_detector/mod.rs

@@ -0,0 +1,256 @@
+use std::{fs, path::PathBuf};
+
+use anyhow::{Context, Result};
+use candle_core::{DType, Device, IndexOp, Tensor};
+use candle_nn::{
+    VarBuilder,
+    ops::{sigmoid, softmax},
+};
+use image::{DynamicImage, GenericImageView, imageops::FilterType};
+use serde::{Deserialize, Serialize};
+
+use crate::{define_models, device};
+
+mod models;
+pub use models::ModelKind;
+
+const FONT_COUNT: usize = 6_150;
+const REGRESSION_START: usize = FONT_COUNT + 2;
+const REGRESSION_DIM: usize = 10;
+
+define_models! {
+    FontWeights => ("fffonion/yuzumarker-font-detection", "yuzumarker-font-detection.safetensors"),
+    FontNames => ("fffonion/yuzumarker-font-detection", "font-labels-ex.json"),
+}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
+pub enum TextDirection {
+    Horizontal,
+    Vertical,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct NamedFontPrediction {
+    pub index: usize,
+    pub name: String,
+    pub language: Option<String>,
+    pub probability: f32,
+    pub serif: bool,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FontPrediction {
+    pub top_fonts: Vec<(usize, f32)>,
+    pub named_fonts: Vec<NamedFontPrediction>,
+    pub direction: TextDirection,
+    pub text_color: [u8; 3],
+    pub stroke_color: [u8; 3],
+    pub font_size_px: f32,
+    pub stroke_width_px: f32,
+    pub line_height: f32,
+    pub angle_deg: f32,
+}
+
+// implement Default for FontPrediction
+impl Default for FontPrediction {
+    fn default() -> Self {
+        Self {
+            top_fonts: Vec::new(),
+            named_fonts: Vec::new(),
+            direction: TextDirection::Horizontal,
+            text_color: [0, 0, 0],
+            stroke_color: [0, 0, 0],
+            font_size_px: 0.0,
+            stroke_width_px: 0.0,
+            line_height: 1.0,
+            angle_deg: 0.0,
+        }
+    }
+}
+
+pub struct FontDetector {
+    model: models::Model,
+    labels: FontLabels,
+    device: Device,
+}
+
+impl FontDetector {
+    pub async fn load(use_cpu: bool) -> Result<Self> {
+        Self::load_with_kind(use_cpu, ModelKind::default()).await
+    }
+
+    pub async fn load_with_kind(use_cpu: bool, kind: ModelKind) -> Result<Self> {
+        let device = device(use_cpu)?;
+        let weights = Manifest::FontWeights.get().await?;
+        let vb = unsafe {
+            VarBuilder::from_mmaped_safetensors(&[weights], DType::F32, &device)?
+                .pp("model._orig_mod.model")
+        };
+        let model = models::Model::load(vb, kind)?;
+        let labels = FontLabels::load().await?;
+
+        Ok(Self {
+            model,
+            device,
+            labels,
+        })
+    }
+
+    pub fn inference(&self, images: &[DynamicImage], top_k: usize) -> Result<Vec<FontPrediction>> {
+        if images.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        let mut processed = Vec::with_capacity(images.len());
+        let mut original_sizes = Vec::with_capacity(images.len());
+        let input_size = self.model.input_size();
+        for image in images {
+            let (w, _h) = image.dimensions();
+            original_sizes.push(w);
+            processed.push(preprocess_image(image, input_size, &self.device)?);
+        }
+        let batch = Tensor::stack(&processed, 0)?;
+        let logits = self.model.forward(&batch, false)?;
+
+        let mut predictions = Vec::with_capacity(images.len());
+        for (index, width) in original_sizes.into_iter().enumerate() {
+            let example = logits.i(index)?;
+            let font_logits = example.narrow(0, 0, FONT_COUNT)?;
+            let font_probs = softmax(&font_logits, 0)?;
+            let font_probs_vec: Vec<f32> = font_probs.to_vec1()?;
+            let mut ranked: Vec<(usize, f32)> = font_probs_vec.into_iter().enumerate().collect();
+            ranked.sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+            ranked.truncate(top_k.min(FONT_COUNT));
+
+            let named_fonts = ranked
+                .iter()
+                .filter_map(|(idx, prob)| {
+                    self.labels.entry(*idx).map(|label| NamedFontPrediction {
+                        index: *idx,
+                        name: label.name.clone(),
+                        language: label.language.clone(),
+                        probability: *prob,
+                        serif: label.serif,
+                    })
+                })
+                .collect();
+
+            let direction_logits = example.narrow(0, FONT_COUNT, 2)?;
+            let direction_vec: Vec<f32> = direction_logits.to_vec1()?;
+            let direction = if direction_vec.len() == 2 && direction_vec[1] > direction_vec[0] {
+                TextDirection::Vertical
+            } else {
+                TextDirection::Horizontal
+            };
+
+            let regression = example.narrow(0, REGRESSION_START, REGRESSION_DIM)?;
+            // Regression head is trained on normalized values; bring logits into [0, 1].
+            let regression = sigmoid(&regression)?;
+            let mut regression: Vec<f32> = regression.to_vec1()?;
+            regression.resize(REGRESSION_DIM, 0.0);
+            let clamp01 = |v: f32| v.clamp(0.0, 1.0);
+            let text_color = [
+                (clamp01(regression[0]) * 255.0).round() as u8,
+                (clamp01(regression[1]) * 255.0).round() as u8,
+                (clamp01(regression[2]) * 255.0).round() as u8,
+            ];
+            let font_size_px = clamp01(regression[3]) * width as f32;
+            let stroke_width_px = clamp01(regression[4]) * width as f32;
+            let stroke_color = [
+                (clamp01(regression[5]) * 255.0).round() as u8,
+                (clamp01(regression[6]) * 255.0).round() as u8,
+                (clamp01(regression[7]) * 255.0).round() as u8,
+            ];
+            let line_spacing_px = clamp01(regression[8]) * width as f32;
+            let line_height = if font_size_px > 0.0 {
+                1.0 + line_spacing_px / font_size_px
+            } else {
+                1.2
+            };
+            let angle_deg = (regression[9] - 0.5) * 180.0;
+
+            predictions.push(FontPrediction {
+                top_fonts: ranked,
+                named_fonts,
+                direction,
+                text_color,
+                stroke_color,
+                font_size_px,
+                stroke_width_px,
+                line_height,
+                angle_deg,
+            });
+        }
+
+        Ok(predictions)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct FontLabel {
+    pub name: String,
+    pub language: Option<String>,
+    pub serif: bool,
+}
+
+#[derive(Debug, Clone)]
+pub struct FontLabels {
+    labels: Vec<FontLabel>,
+}
+
+impl FontLabels {
+    pub async fn load() -> Result<Self> {
+        let path = Manifest::FontNames.get().await?;
+        Self::from_path(&path)
+    }
+
+    pub fn from_path(path: &PathBuf) -> Result<Self> {
+        let data = fs::read_to_string(&path)
+            .with_context(|| format!("Failed to read labels file {}", path.display()))?;
+        let entries: Vec<FontLabelEntry> = serde_json::from_str(&data)
+            .with_context(|| format!("Failed to parse labels file {}", path.display()))?;
+        let mut labels = Vec::with_capacity(entries.len());
+        for entry in entries {
+            labels.push(FontLabel {
+                name: entry.path,
+                language: entry.language,
+                serif: entry.serif,
+            });
+        }
+        Ok(Self { labels })
+    }
+
+    pub fn entry(&self, idx: usize) -> Option<&FontLabel> {
+        self.labels.get(idx)
+    }
+
+    pub fn name(&self, idx: usize) -> Option<&str> {
+        self.entry(idx).map(|label| label.name.as_str())
+    }
+
+    pub fn language(&self, idx: usize) -> Option<&str> {
+        self.entry(idx).and_then(|label| label.language.as_deref())
+    }
+}
+
+#[derive(serde::Deserialize)]
+struct FontLabelEntry {
+    path: String,
+    language: Option<String>,
+    serif: bool,
+}
+
+fn preprocess_image(image: &DynamicImage, target: usize, device: &Device) -> Result<Tensor> {
+    let resized = image.resize_exact(target as u32, target as u32, FilterType::CatmullRom);
+    let data = resized.to_rgb8().into_raw();
+    let tensor = Tensor::from_vec(
+        data,
+        (target, target, 3),
+        &Device::Cpu,
+    )?
+    .to_dtype(DType::F32)?
+    .permute((2, 0, 1))? // (3, H, W)
+    * (1.0 / 255.0);
+    let tensor = tensor?;
+    Ok(tensor.to_device(device)?)
+}

koharu-ml/src/font_detector/models.rs

@@ -0,0 +1,524 @@
+use anyhow::Result;
+use candle_core::{DType, Module, ModuleT, Tensor};
+use candle_nn::{BatchNorm, Conv2d, Conv2dConfig, Linear, VarBuilder};
+use clap::ValueEnum;
+
+use super::{FONT_COUNT, REGRESSION_DIM};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
+#[value(rename_all = "kebab-case")]
+pub enum ModelKind {
+    Resnet18,
+    Resnet34,
+    Resnet50,
+    Resnet101,
+    Deepfont,
+}
+
+impl Default for ModelKind {
+    fn default() -> Self {
+        ModelKind::Resnet50
+    }
+}
+
+pub struct Model {
+    kind: ModelKind,
+    inner: ModelImpl,
+}
+
+enum ModelImpl {
+    ResNet(ResNet),
+    DeepFont(DeepFont),
+}
+
+impl Model {
+    pub fn load(vb: VarBuilder, kind: ModelKind) -> Result<Self> {
+        let model = match kind {
+            ModelKind::Resnet18 => ModelImpl::ResNet(ResNet::load_basic(vb, [2, 2, 2, 2], 1)?),
+            ModelKind::Resnet34 => ModelImpl::ResNet(ResNet::load_basic(vb, [3, 4, 6, 3], 1)?),
+            ModelKind::Resnet50 => ModelImpl::ResNet(ResNet::load_bottleneck(vb, [3, 4, 6, 3], 4)?),
+            ModelKind::Resnet101 => {
+                ModelImpl::ResNet(ResNet::load_bottleneck(vb, [3, 4, 23, 3], 4)?)
+            }
+            ModelKind::Deepfont => ModelImpl::DeepFont(DeepFont::load(vb)?),
+        };
+        Ok(Self { kind, inner: model })
+    }
+
+    pub fn input_size(&self) -> usize {
+        match self.kind {
+            ModelKind::Deepfont => 105,
+            _ => 512,
+        }
+    }
+
+    pub fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        let logits = match &self.inner {
+            ModelImpl::ResNet(m) => m.forward(xs, train)?,
+            ModelImpl::DeepFont(m) => m.forward(xs, train)?,
+        };
+
+        let (_, dim) = logits.dims2()?;
+        if dim == FONT_COUNT + REGRESSION_DIM + 2 {
+            return Ok(logits);
+        }
+
+        // For models that only output font logits (e.g., DeepFont), pad zeros for direction/regression.
+        if dim == FONT_COUNT {
+            let device = logits.device();
+            let zeros = Tensor::zeros((logits.dim(0)?, REGRESSION_DIM + 2), DType::F32, device)?;
+            return Tensor::cat(&[logits, zeros], 1);
+        }
+
+        Err(candle_core::Error::Msg(format!(
+            "Unexpected output dimension from backbone: got {}, expected {}",
+            dim,
+            FONT_COUNT + REGRESSION_DIM + 2
+        )))
+    }
+}
+
+#[derive(Clone)]
+struct BasicBlock {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
+    downsample: Option<(Conv2d, BatchNorm)>,
+}
+
+impl BasicBlock {
+    fn load(vb: VarBuilder, in_channels: usize, planes: usize, stride: usize) -> Result<Self> {
+        let conv1 = Conv2d::new(
+            vb.pp("conv1").get((planes, in_channels, 3, 3), "weight")?,
+            None,
+            Conv2dConfig {
+                stride,
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let bn1 = load_batch_norm(&vb.pp("bn1"), planes)?;
+        let conv2 = Conv2d::new(
+            vb.pp("conv2").get((planes, planes, 3, 3), "weight")?,
+            None,
+            Conv2dConfig {
+                stride: 1,
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let bn2 = load_batch_norm(&vb.pp("bn2"), planes)?;
+
+        let downsample = if stride != 1 || in_channels != planes {
+            let conv = Conv2d::new(
+                vb.pp("downsample.0")
+                    .get((planes, in_channels, 1, 1), "weight")?,
+                None,
+                Conv2dConfig {
+                    stride,
+                    ..Default::default()
+                },
+            );
+            let bn = load_batch_norm(&vb.pp("downsample.1"), planes)?;
+            Some((conv, bn))
+        } else {
+            None
+        };
+
+        Ok(Self {
+            conv1,
+            bn1,
+            conv2,
+            bn2,
+            downsample,
+        })
+    }
+
+    fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        let mut out = self.conv1.forward(xs)?;
+        out = self.bn1.forward_t(&out, train)?;
+        out = out.relu()?;
+
+        out = self.conv2.forward(&out)?;
+        out = self.bn2.forward_t(&out, train)?;
+
+        let residual = if let Some((conv, bn)) = &self.downsample {
+            let mut y = conv.forward(xs)?;
+            y = bn.forward_t(&y, train)?;
+            y
+        } else {
+            xs.clone()
+        };
+
+        (out + residual)?.relu()
+    }
+}
+
+#[derive(Clone)]
+struct Bottleneck {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
+    conv3: Conv2d,
+    bn3: BatchNorm,
+    downsample: Option<(Conv2d, BatchNorm)>,
+}
+
+impl Bottleneck {
+    fn load(
+        vb: VarBuilder,
+        in_channels: usize,
+        planes: usize,
+        stride: usize,
+        expansion: usize,
+    ) -> Result<Self> {
+        let conv1 = Conv2d::new(
+            vb.pp("conv1").get((planes, in_channels, 1, 1), "weight")?,
+            None,
+            Conv2dConfig::default(),
+        );
+        let bn1 = load_batch_norm(&vb.pp("bn1"), planes)?;
+        let conv2 = Conv2d::new(
+            vb.pp("conv2").get((planes, planes, 3, 3), "weight")?,
+            None,
+            Conv2dConfig {
+                stride,
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let bn2 = load_batch_norm(&vb.pp("bn2"), planes)?;
+        let conv3 = Conv2d::new(
+            vb.pp("conv3")
+                .get((planes * expansion, planes, 1, 1), "weight")?,
+            None,
+            Conv2dConfig::default(),
+        );
+        let bn3 = load_batch_norm(&vb.pp("bn3"), planes * expansion)?;
+
+        let downsample = if in_channels != planes * expansion || stride != 1 {
+            let conv = Conv2d::new(
+                vb.pp("downsample.0")
+                    .get((planes * expansion, in_channels, 1, 1), "weight")?,
+                None,
+                Conv2dConfig {
+                    stride,
+                    ..Default::default()
+                },
+            );
+            let bn = load_batch_norm(&vb.pp("downsample.1"), planes * expansion)?;
+            Some((conv, bn))
+        } else {
+            None
+        };
+
+        Ok(Self {
+            conv1,
+            bn1,
+            conv2,
+            bn2,
+            conv3,
+            bn3,
+            downsample,
+        })
+    }
+
+    fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        let mut out = self.conv1.forward(xs)?;
+        out = self.bn1.forward_t(&out, train)?;
+        out = out.relu()?;
+
+        out = self.conv2.forward(&out)?;
+        out = self.bn2.forward_t(&out, train)?;
+        out = out.relu()?;
+
+        out = self.conv3.forward(&out)?;
+        out = self.bn3.forward_t(&out, train)?;
+
+        let residual = if let Some((conv, bn)) = &self.downsample {
+            let mut y = conv.forward(xs)?;
+            y = bn.forward_t(&y, train)?;
+            y
+        } else {
+            xs.clone()
+        };
+
+        (out + residual)?.relu()
+    }
+}
+
+struct ResNet {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    layer1: Vec<ResBlock>,
+    layer2: Vec<ResBlock>,
+    layer3: Vec<ResBlock>,
+    layer4: Vec<ResBlock>,
+    fc: Linear,
+}
+
+enum ResBlock {
+    Basic(BasicBlock),
+    Bottleneck(Bottleneck),
+}
+
+impl ResNet {
+    fn load_basic(vb: VarBuilder, layers: [usize; 4], expansion: usize) -> Result<Self> {
+        Self::load_impl(vb, layers, BlockKind::Basic, expansion)
+    }
+
+    fn load_bottleneck(vb: VarBuilder, layers: [usize; 4], expansion: usize) -> Result<Self> {
+        Self::load_impl(vb, layers, BlockKind::Bottleneck, expansion)
+    }
+
+    fn load_impl(
+        vb: VarBuilder,
+        layers: [usize; 4],
+        block: BlockKind,
+        expansion: usize,
+    ) -> Result<Self> {
+        let conv1 = Conv2d::new(
+            vb.pp("conv1").get((64, 3, 7, 7), "weight")?,
+            None,
+            Conv2dConfig {
+                stride: 2,
+                padding: 3,
+                ..Default::default()
+            },
+        );
+        let bn1 = load_batch_norm(&vb.pp("bn1"), 64)?;
+
+        let (layer1, c1) =
+            Self::make_layer(vb.pp("layer1"), 64, 64, layers[0], 1, block, expansion)?;
+        let (layer2, c2) =
+            Self::make_layer(vb.pp("layer2"), c1, 128, layers[1], 2, block, expansion)?;
+        let (layer3, c3) =
+            Self::make_layer(vb.pp("layer3"), c2, 256, layers[2], 2, block, expansion)?;
+        let (layer4, c4) =
+            Self::make_layer(vb.pp("layer4"), c3, 512, layers[3], 2, block, expansion)?;
+
+        let fc = Linear::new(
+            vb.pp("fc")
+                .get((FONT_COUNT + REGRESSION_DIM + 2, c4), "weight")?,
+            Some(vb.pp("fc").get(FONT_COUNT + REGRESSION_DIM + 2, "bias")?),
+        );
+
+        Ok(Self {
+            conv1,
+            bn1,
+            layer1,
+            layer2,
+            layer3,
+            layer4,
+            fc,
+        })
+    }
+
+    fn make_layer(
+        vb: VarBuilder,
+        in_channels: usize,
+        planes: usize,
+        blocks: usize,
+        stride: usize,
+        block_kind: BlockKind,
+        expansion: usize,
+    ) -> Result<(Vec<ResBlock>, usize)> {
+        let mut layers = Vec::with_capacity(blocks);
+        let first = match block_kind {
+            BlockKind::Basic => {
+                ResBlock::Basic(BasicBlock::load(vb.pp("0"), in_channels, planes, stride)?)
+            }
+            BlockKind::Bottleneck => ResBlock::Bottleneck(Bottleneck::load(
+                vb.pp("0"),
+                in_channels,
+                planes,
+                stride,
+                expansion,
+            )?),
+        };
+        layers.push(first);
+        let current_channels = planes * expansion;
+        for idx in 1..blocks {
+            let block_vb = vb.pp(idx.to_string());
+            let block = match block_kind {
+                BlockKind::Basic => {
+                    ResBlock::Basic(BasicBlock::load(block_vb, current_channels, planes, 1)?)
+                }
+                BlockKind::Bottleneck => ResBlock::Bottleneck(Bottleneck::load(
+                    block_vb,
+                    current_channels,
+                    planes,
+                    1,
+                    expansion,
+                )?),
+            };
+            layers.push(block);
+        }
+        Ok((layers, current_channels))
+    }
+
+    fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        let mut x = self.conv1.forward(xs)?;
+        x = self.bn1.forward_t(&x, train)?;
+        x = x.relu()?;
+        x = x.max_pool2d_with_stride(3, 2)?;
+
+        for b in &self.layer1 {
+            x = b.forward(&x, train)?;
+        }
+        for b in &self.layer2 {
+            x = b.forward(&x, train)?;
+        }
+        for b in &self.layer3 {
+            x = b.forward(&x, train)?;
+        }
+        for b in &self.layer4 {
+            x = b.forward(&x, train)?;
+        }
+
+        let (_, c, h, w) = x.dims4()?;
+        let mut x = x.sum_keepdim(2)?;
+        x = x.sum_keepdim(3)?;
+        x = (x / ((h * w) as f64))?.reshape((xs.dim(0)?, c))?;
+        self.fc.forward(&x)
+    }
+}
+
+impl ResBlock {
+    fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        match self {
+            ResBlock::Basic(b) => b.forward(xs, train),
+            ResBlock::Bottleneck(b) => b.forward(xs, train),
+        }
+    }
+}
+
+#[derive(Clone, Copy)]
+enum BlockKind {
+    Basic,
+    Bottleneck,
+}
+
+struct DeepFont {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
+    conv3: Conv2d,
+    conv4: Conv2d,
+    conv5: Conv2d,
+    fc1: Linear,
+    fc2: Linear,
+    fc3: Linear,
+}
+
+impl DeepFont {
+    fn load(vb: VarBuilder) -> Result<Self> {
+        let conv1 = Conv2d::new(
+            vb.pp("0").get((64, 3, 11, 11), "weight")?,
+            Some(vb.pp("0").get(64, "bias")?),
+            Conv2dConfig {
+                stride: 2,
+                ..Default::default()
+            },
+        );
+        let bn1 = load_batch_norm(&vb.pp("1"), 64)?;
+        let conv2 = Conv2d::new(
+            vb.pp("4").get((128, 64, 3, 3), "weight")?,
+            Some(vb.pp("4").get(128, "bias")?),
+            Conv2dConfig {
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let bn2 = load_batch_norm(&vb.pp("5"), 128)?;
+        let conv3 = Conv2d::new(
+            vb.pp("8").get((256, 128, 3, 3), "weight")?,
+            Some(vb.pp("8").get(256, "bias")?),
+            Conv2dConfig {
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let conv4 = Conv2d::new(
+            vb.pp("9").get((256, 256, 3, 3), "weight")?,
+            Some(vb.pp("9").get(256, "bias")?),
+            Conv2dConfig {
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let conv5 = Conv2d::new(
+            vb.pp("10").get((256, 256, 3, 3), "weight")?,
+            Some(vb.pp("10").get(256, "bias")?),
+            Conv2dConfig {
+                padding: 1,
+                ..Default::default()
+            },
+        );
+        let fc1 = Linear::new(
+            vb.pp("14").get((4096, 256 * 12 * 12), "weight")?,
+            Some(vb.pp("14").get(4096, "bias")?),
+        );
+        let fc2 = Linear::new(
+            vb.pp("16").get((4096, 4096), "weight")?,
+            Some(vb.pp("16").get(4096, "bias")?),
+        );
+        let fc3 = Linear::new(
+            vb.pp("18").get((FONT_COUNT, 4096), "weight")?,
+            Some(vb.pp("18").get(FONT_COUNT, "bias")?),
+        );
+
+        Ok(Self {
+            conv1,
+            bn1,
+            conv2,
+            bn2,
+            conv3,
+            conv4,
+            conv5,
+            fc1,
+            fc2,
+            fc3,
+        })
+    }
+
+    fn forward(&self, xs: &Tensor, train: bool) -> candle_core::Result<Tensor> {
+        let mut x = self.conv1.forward(xs)?;
+        x = self.bn1.forward_t(&x, train)?;
+        x = x.relu()?;
+        x = x.max_pool2d_with_stride(2, 2)?;
+
+        x = self.conv2.forward(&x)?;
+        x = self.bn2.forward_t(&x, train)?;
+        x = x.relu()?;
+        x = x.max_pool2d_with_stride(2, 2)?;
+
+        x = self.conv3.forward(&x)?;
+        x = x.relu()?;
+        x = self.conv4.forward(&x)?;
+        x = x.relu()?;
+        x = self.conv5.forward(&x)?;
+        x = x.relu()?;
+
+        x = x.flatten(1, x.rank() - 1)?;
+        x = self.fc1.forward(&x)?;
+        x = x.relu()?;
+        x = self.fc2.forward(&x)?;
+        x = x.relu()?;
+        self.fc3.forward(&x)
+    }
+}
+
+fn load_batch_norm(vb: &VarBuilder, channels: usize) -> Result<BatchNorm> {
+    Ok(BatchNorm::new(
+        channels,
+        vb.get(channels, "running_mean")?,
+        vb.get(channels, "running_var")?,
+        vb.get(channels, "weight")?,
+        vb.get(channels, "bias")?,
+        1e-5,
+    )?)
+}

koharu-ml/src/lib.rs

@@ -1,6 +1,7 @@
 mod hf_hub;
 
 pub mod comic_text_detector;
+pub mod font_detector;
 pub mod lama;
 pub mod llm;
 pub mod manga_ocr;

koharu/src/command.rs

@@ -123,10 +123,36 @@ pub async fn detect(
         .get_mut(index)
         .ok_or_else(|| anyhow::anyhow!("Document not found"))?;
 
-    let (text_blocks, segment) = model.detect(&document.image).await?;
+    let (text_blocks, segment) = model.detect_dialog(&document.image).await?;
     document.text_blocks = text_blocks;
     document.segment = Some(segment);
 
+    // detect fonts for each text block
+    if !document.text_blocks.is_empty() {
+        let images: Vec<image::DynamicImage> = document
+            .text_blocks
+            .iter()
+            .map(|block| {
+                let sub_image = document.image.crop_imm(
+                    block.x as u32,
+                    block.y as u32,
+                    block.width as u32,
+                    block.height as u32,
+                );
+                sub_image
+            })
+            .collect();
+        let font_predictions = model.detect_fonts(&images, 1).await?;
+        for (block, prediction) in document
+            .text_blocks
+            .iter_mut()
+            .zip(font_predictions.into_iter())
+        {
+            tracing::info!("Detected font for block {:?}: {:?}", block.text, prediction);
+            block.font_info = Some(prediction);
+        }
+    }
+
     Ok(document.clone())
 }
 

koharu/src/ml.rs

@@ -1,6 +1,7 @@
 use anyhow::Result;
 use image::DynamicImage;
 use koharu_ml::comic_text_detector::{self, ComicTextDetector};
+use koharu_ml::font_detector::{self, FontDetector};
 use koharu_ml::lama::{self, Lama};
 use koharu_ml::manga_ocr::{self, MangaOcr};
 
@@ -8,25 +9,27 @@ use crate::image::SerializableDynamicImage;
 use crate::state::TextBlock;
 
 pub struct Model {
-    detector: ComicTextDetector,
+    dialog_detector: ComicTextDetector,
     ocr: MangaOcr,
     lama: Lama,
+    font_detector: FontDetector,
 }
 
 impl Model {
     pub async fn new(use_cpu: bool) -> Result<Self> {
         Ok(Self {
-            detector: ComicTextDetector::load(use_cpu).await?,
+            dialog_detector: ComicTextDetector::load(use_cpu).await?,
             ocr: MangaOcr::load(use_cpu).await?,
             lama: Lama::load(use_cpu).await?,
+            font_detector: FontDetector::load(use_cpu).await?,
         })
     }
 
-    pub async fn detect(
+    pub async fn detect_dialog(
         &self,
         image: &SerializableDynamicImage,
     ) -> Result<(Vec<TextBlock>, SerializableDynamicImage)> {
-        let (bboxes, segment) = self.detector.inference(image)?;
+        let (bboxes, segment) = self.dialog_detector.inference(image)?;
 
         let mut text_blocks: Vec<TextBlock> = bboxes
             .into_iter()
@@ -91,12 +94,36 @@ impl Model {
 
         Ok(result.into())
     }
+
+    pub async fn detect_font(
+        &self,
+        image: &DynamicImage,
+        top_k: usize,
+    ) -> Result<font_detector::FontPrediction> {
+        let mut results = self
+            .detect_fonts(std::slice::from_ref(image), top_k)
+            .await?;
+        Ok(results.pop().unwrap_or_default())
+    }
+
+    pub async fn detect_fonts(
+        &self,
+        images: &[DynamicImage],
+        top_k: usize,
+    ) -> Result<Vec<font_detector::FontPrediction>> {
+        if images.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        self.font_detector.inference(images, top_k)
+    }
 }
 
 pub async fn prefetch() -> Result<()> {
     comic_text_detector::prefetch().await?;
     manga_ocr::prefetch().await?;
     lama::prefetch().await?;
+    font_detector::prefetch().await?;
 
     Ok(())
 }

koharu/src/state.rs

@@ -1,6 +1,7 @@
 use std::{path::PathBuf, sync::Arc};
 
 use image::GenericImageView;
+use koharu_ml::font_detector::FontPrediction;
 use koharu_renderer::types::Color;
 use serde::{Deserialize, Serialize};
 use tokio::sync::RwLock;
@@ -17,6 +18,7 @@ pub struct TextBlock {
     pub text: Option<String>,
     pub translation: Option<String>,
     pub style: Option<TextStyle>,
+    pub font_info: Option<FontPrediction>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

scripts/convert_font_detection.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""
+Convert YuzuMarker.FontDetection checkpoints (.ckpt) to safetensors for Candle.
+
+Example:
+    python scripts/convert_font_detection.py \
+        --checkpoint name=4x-epoch=84-step=1649340.ckpt
+"""
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+import torch
+from safetensors.torch import save_file
+
+
+DEFAULT_CKPT = "name=4x-epoch=84-step=1649340.ckpt"
+REPO_ID = "gyrojeff/YuzuMarker.FontDetection"
+
+
+def parse_args() -> argparse.Namespace:
+    cache_dir = (
+        Path.home() / ".cache" / "Koharu" / "models" / "yuzumarker-font-detection.safetensors"
+    )
+    parser = argparse.ArgumentParser(description="Convert YuzuMarker.FontDetection checkpoint.")
+    parser.add_argument(
+        "-c",
+        "--checkpoint",
+        default=DEFAULT_CKPT,
+        help=f"Checkpoint filename from {REPO_ID} (default: {DEFAULT_CKPT})",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=cache_dir,
+        help=f"Output safetensors path (default: {cache_dir})",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Downloading {args.checkpoint} from {REPO_ID} ...")
+    ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=args.checkpoint)
+    print(f"Loaded checkpoint at {ckpt_path}")
+
+    state = torch.load(ckpt_path, map_location="cpu")
+    if "state_dict" not in state:
+        raise RuntimeError("Unexpected checkpoint format: missing state_dict")
+    state_dict = state["state_dict"]
+    print(f"Saving {len(state_dict)} tensors to {args.output}")
+    save_file(state_dict, str(args.output))
+    print("Done.")
+
+
+if __name__ == "__main__":
+    main()

scripts/convert_font_labels.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""
+Convert font_demo_cache.bin (from YuzuMarker.FontDetection) to a JSON list that Rust can read.
+
+Usage:
+    python scripts/convert_font_labels.py \
+        --input font_demo_cache.bin \
+        --output yuzumarker-font-labels.json
+"""
+
+import argparse
+import json
+import sys
+import types
+from pathlib import Path
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Convert font_demo_cache.bin to JSON")
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=Path,
+        default=Path("font_demo_cache.bin"),
+        help="Input pickle file (font_demo_cache.bin from the original repo)",
+    )
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=Path,
+        default=Path("yuzumarker-font-labels.json"),
+        help="Output JSON path (default: yuzumarker-font-labels.json)",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if not args.input.exists():
+        sys.exit(f"Input file not found: {args.input}")
+
+    # Stub module/classes so pickle can load without the original code
+    font_dataset_mod = types.ModuleType("font_dataset")
+    font_mod = types.ModuleType("font_dataset.font")
+
+    class DSFont:
+        def __init__(self, path=None, language=None):
+            self.path = path
+            self.language = language
+
+    font_mod.DSFont = DSFont
+    sys.modules["font_dataset"] = font_dataset_mod
+    sys.modules["font_dataset.font"] = font_mod
+    font_dataset_mod.font = font_mod
+
+    import pickle  # noqa: E402
+
+    with open(args.input, "rb") as f:
+        data = pickle.load(f)
+
+    entries = []
+    for item in data:
+        path = getattr(item, "path", None)
+        language = getattr(item, "language", None)
+        if path is None:
+            continue
+        entries.append({"path": path, "language": language})
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with open(args.output, "w", encoding="utf-8") as f:
+        json.dump(entries, f, ensure_ascii=False, indent=2)
+
+    print(f"Wrote {len(entries)} labels to {args.output}")
+
+
+if __name__ == "__main__":
+    main()