package fetch
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
)
// ModelFile describes one pinned OCR model asset on Hugging Face.
type ModelFile struct {
Name string
URL string
SHA256 string
Size int64
}
// Pinned model files for the Linux OCR pipeline:
// - Meiki text detector (D-FINE): detects horizontal and vertical text lines
// - Meiki text recognizer: recognized horizontal dialogue / subtitles / UI
// - Manga-OCR (kha-white base ONNX): recognized vertical speech bubbles and whole images
var PinnedModels = []ModelFile{
{
Name: "meiki.text.detect.v0.1.960x544.onnx",
URL: "https://huggingface.co/rtr46/meiki.text.detect.v0/resolve/a9cffa4f60cbf72ddb87edf19c6f98a01cd042e6/meiki.text.detect.v0.1.960x544.onnx",
SHA256: "40b6a016667745cae7d3055929ae3b8b1e7716aac795f5904cd3c2c7c3b8404b",
Size: 14503825,
},
{
Name: "meiki.text.rec.v0.960x32.onnx",
URL: "https://huggingface.co/rtr46/meiki.txt.recognition.v0/resolve/a28cf5874dc2438ebb1c86336be26bcec51e3375/meiki.text.rec.v0.960x32.onnx",
SHA256: "3e96bc772fbee9717e536a6353032bb944c3382dd2f6960ef4890decda43b000",
Size: 18593254,
},
{
Name: "encoder_model.onnx",
URL: "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/encoder_model.onnx",
SHA256: "15fa8155fe9bc1a7d25d9bb353debaa4def033d0174e907dbd2dd6d995def85f",
Size: 343454249,
},
{
Name: "decoder_model.onnx",
URL: "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/decoder_model.onnx",
SHA256: "ef7765261e9d1cdc34d89356986c2bbc2a082897f753a89605ae80fdfa61f5e8",
Size: 117480262,
},
{
Name: "vocab.txt",
URL: "https://huggingface.co/mayocream/manga-ocr-onnx/resolve/24b12778d85800835e2ca409236de281b8ab7b9f/vocab.txt",
SHA256: "5cb5c5586d98a2f331d9f8828e4586479b0611bfba5d8c3b6dadffc84d6a36a3",
Size: 30216,
},
}
// VerifyFile returns true if path exists, matches expectedSize, and has expectedSHA256.
func VerifyFile(path string, expectedSHA string, expectedSize int64) bool {
info, err := os.Stat(path)
if err != nil || info.Size() != expectedSize {
return false
}
f, err := os.Open(path)
if err != nil {
return false
}
defer f.Close()
h := sha256.New()
if _, err := io.Copy(h, f); err != nil {
return false
}
return hex.EncodeToString(h.Sum(nil)) == expectedSHA
}
// EnsureModels verifies all required model files in dir, downloading any missing
// or corrupted files atomically into dir.
func EnsureModels(ctx context.Context, dir string) error {
if err := os.MkdirAll(dir, 0755); err != nil {
return fmt.Errorf("create models dir %s: %w", dir, err)
}
for _, m := range PinnedModels {
dest := filepath.Join(dir, m.Name)
if VerifyFile(dest, m.SHA256, m.Size) {
continue
}
if err := downloadFile(ctx, dir, m); err != nil {
return fmt.Errorf("download %s: %w", m.Name, err)
}
}
return nil
}
func downloadFile(ctx context.Context, dir string, m ModelFile) error {
tmpFile, err := os.CreateTemp(dir, ".tmp-"+m.Name+"-*")
if err != nil {
return fmt.Errorf("create temp file: %w", err)
}
tmpPath := tmpFile.Name()
defer func() {
tmpFile.Close()
os.Remove(tmpPath)
}()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, m.URL, nil)
if err != nil {
return fmt.Errorf("create request: %w", err)
}
req.Header.Set("User-Agent", "chirami-ocr/1.0")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return fmt.Errorf("http get: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("unexpected status %s from %s", resp.Status, m.URL)
}
h := sha256.New()
tee := io.TeeReader(resp.Body, h)
n, err := io.Copy(tmpFile, tee)
if err != nil {
return fmt.Errorf("copying body: %w", err)
}
if n != m.Size {
return fmt.Errorf("expected %d bytes, got %d", m.Size, n)
}
actualSHA := hex.EncodeToString(h.Sum(nil))
if actualSHA != m.SHA256 {
return fmt.Errorf("sha256 mismatch: got %s, want %s", actualSHA, m.SHA256)
}
if err := tmpFile.Sync(); err != nil {
return fmt.Errorf("sync temp file: %w", err)
}
if err := tmpFile.Close(); err != nil {
return fmt.Errorf("close temp file: %w", err)
}
destPath := filepath.Join(dir, m.Name)
if err := os.Rename(tmpPath, destPath); err != nil {
return fmt.Errorf("atomic rename: %w", err)
}
return nil
}
// Missing returns any model files from PinnedModels that are not present and valid in dir.
func Missing(dir string) []ModelFile {
var missing []ModelFile
for _, m := range PinnedModels {
dest := filepath.Join(dir, m.Name)
if !VerifyFile(dest, m.SHA256, m.Size) {
missing = append(missing, m)
}
}
return missing
}
// Absent returns the model files that are missing from dir or have the wrong
// size, without hashing them.
func Absent(dir string) []ModelFile {
var absent []ModelFile
for _, m := range PinnedModels {
if st, err := os.Stat(filepath.Join(dir, m.Name)); err != nil || st.Size() != m.Size {
absent = append(absent, m)
}
}
return absent
}