novel-downloader 1.1.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

novel_downloader/core/requesters/base_async_session.py

@@ -0,0 +1,299 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+novel_downloader.core.requesters.base_async_session
+---------------------------------------------------
+
+This module defines the BaseAsyncSession class, which provides asynchronous
+HTTP request capabilities using aiohttp. It maintains a persistent
+client session and supports retries, headers, timeout configurations,
+cookie handling, and defines abstract methods for subclasses.
+"""
+
+import abc
+import asyncio
+import time
+from typing import Any, Dict, Optional, Union
+
+import aiohttp
+from aiohttp import ClientResponse, ClientSession, ClientTimeout, TCPConnector
+
+from novel_downloader.config.models import RequesterConfig
+from novel_downloader.core.interfaces import AsyncRequesterProtocol
+from novel_downloader.utils.constants import DEFAULT_USER_HEADERS
+
+
+class RateLimiter:
+    """
+    Simple async token-bucket rate limiter: ensures no more than rate_per_sec
+    requests are started per second, across all coroutines.
+    """
+
+    def __init__(self, rate_per_sec: float):
+        self._interval = 1.0 / rate_per_sec
+        self._lock = asyncio.Lock()
+        self._last = time.monotonic()
+
+    async def wait(self) -> None:
+        async with self._lock:
+            now = time.monotonic()
+            elapsed = now - self._last
+            delay = self._interval - elapsed
+            if delay > 0:
+                await asyncio.sleep(delay)
+            self._last = time.monotonic()
+
+
+class BaseAsyncSession(AsyncRequesterProtocol, abc.ABC):
+    """
+    BaseAsyncSession wraps basic HTTP operations using aiohttp.ClientSession,
+    supporting retry logic, timeout, persistent connections, and cookie management.
+
+    Attributes:
+        _session (ClientSession): The persistent aiohttp client session.
+        _timeout (float): Timeout for each request in seconds.
+        _retry_times (int): Number of retry attempts on failure.
+        _retry_interval (float): Delay (in seconds) between retries.
+        _headers (Dict[str, str]): Default HTTP headers to send.
+        _cookies (Dict[str, str]): Optional cookie jar for the session.
+    """
+
+    def _init_session(
+        self,
+        config: RequesterConfig,
+        cookies: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """
+        Initialize the async session with configuration.
+
+        :param config: Configuration object for session behavior
+                       (timeouts, retries, headers, etc.)
+        :param cookies: Optional initial cookies to set on the session.
+        """
+        self._config = config
+        self._timeout = config.timeout
+        self._retry_times = config.retry_times
+        self._retry_interval = config.retry_interval
+        self._cookies = cookies or {}
+        self._headers = DEFAULT_USER_HEADERS.copy()
+        self._session: Optional[ClientSession] = None
+        self._rate_limiter: Optional[RateLimiter] = None
+
+    async def _setup(self) -> None:
+        """
+        Set up the aiohttp.ClientSession with timeout, connector, headers, and cookies.
+        """
+        max_rps = getattr(self._config, "max_rps", None)
+        if max_rps is not None:
+            self._rate_limiter = RateLimiter(max_rps)
+
+        timeout = ClientTimeout(total=self._timeout)
+        connector = TCPConnector(
+            limit_per_host=getattr(self._config, "max_connections", 10)
+        )
+        self._session = ClientSession(
+            timeout=timeout,
+            connector=connector,
+            headers=self._headers,
+            cookies=self._cookies,
+        )
+
+    async def login(self, max_retries: int = 3, manual_login: bool = False) -> bool:
+        """
+        Attempt to log in asynchronously.
+        Override in subclasses that require authentication.
+
+        :returns: True if login succeeded, False otherwise.
+        """
+        raise NotImplementedError(
+            "Login is not supported by this session type. "
+            "Override login() in your subclass to enable it."
+        )
+
+    @abc.abstractmethod
+    async def get_book_info(
+        self, book_id: str, wait_time: Optional[float] = None
+    ) -> str:
+        """
+        Fetch the raw HTML (or JSON) of the book info page asynchronously.
+
+        :param book_id: The book identifier.
+        :param wait_time: Base number of seconds to wait before returning content.
+        :return: The page content as a string.
+        """
+        ...
+
+    @abc.abstractmethod
+    async def get_book_chapter(
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
+    ) -> str:
+        """
+        Fetch the raw HTML (or JSON) of a single chapter asynchronously.
+
+        :param book_id: The book identifier.
+        :param chapter_id: The chapter identifier.
+        :param wait_time: Base number of seconds to wait before returning content.
+        :return: The chapter content as a string.
+        """
+        ...
+
+    async def get_bookcase(self, wait_time: Optional[float] = None) -> str:
+        """
+        Optional: Retrieve the HTML content of the authenticated user's bookcase page.
+        Subclasses that support user login/bookcase should override this.
+
+        :param wait_time: Base number of seconds to wait before returning content.
+        :return: The HTML of the bookcase page.
+        """
+        raise NotImplementedError(
+            "Bookcase fetching is not supported by this session type. "
+            "Override get_bookcase() in your subclass to enable it."
+        )
+
+    async def fetch(self, url: str, **kwargs: Any) -> str:
+        """
+        Fetch the content from the given URL asynchronously, with retry support.
+
+        :param url: The target URL to fetch.
+        :param kwargs: Additional keyword arguments to pass to `session.get`.
+        :return: The response body as text.
+        :raises: aiohttp.ClientError on final failure.
+        """
+        if self._session is None:
+            await self._setup()
+        if self._session is None:
+            raise RuntimeError("Session not initialized after setup")
+
+        if self._rate_limiter:
+            await self._rate_limiter.wait()
+
+        for attempt in range(self._retry_times + 1):
+            try:
+                async with self._session.get(url, **kwargs) as resp:
+                    resp.raise_for_status()
+                    text: str = await resp.text()
+                    return text
+            except aiohttp.ClientError:
+                if attempt < self._retry_times:
+                    await asyncio.sleep(self._retry_interval)
+                    continue
+                raise
+
+        raise RuntimeError("Unreachable code reached in fetch()")
+
+    async def get(
+        self, url: str, params: Optional[Dict[str, Any]] = None, **kwargs: Any
+    ) -> ClientResponse:
+        """
+        Send an HTTP GET request asynchronously.
+
+        :param url: The target URL.
+        :param params: Query parameters to include in the request.
+        :param kwargs: Additional args passed to session.get().
+        :return: aiohttp.ClientResponse object.
+        :raises RuntimeError: If the session is not initialized.
+        """
+        if self._session is None:
+            await self._setup()
+        if self._session is None:
+            raise RuntimeError("Session not initialized after setup")
+
+        if self._rate_limiter:
+            await self._rate_limiter.wait()
+        return await self._session.get(url, params=params, **kwargs)
+
+    async def post(
+        self,
+        url: str,
+        data: Optional[Union[Dict[str, Any], bytes]] = None,
+        json: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> ClientResponse:
+        """
+        Send an HTTP POST request asynchronously.
+
+        :param url: The target URL.
+        :param data: Form data to include in the request body.
+        :param json: JSON body to include in the request.
+        :param kwargs: Additional args passed to session.post().
+        :return: aiohttp.ClientResponse object.
+        :raises RuntimeError: If the session is not initialized.
+        """
+        if self._session is None:
+            await self._setup()
+        if self._session is None:
+            raise RuntimeError("Session not initialized after setup")
+
+        if self._rate_limiter:
+            await self._rate_limiter.wait()
+        return await self._session.post(url, data=data, json=json, **kwargs)
+
+    @property
+    def session(self) -> ClientSession:
+        """
+        Return the active aiohttp.ClientSession.
+
+        :raises RuntimeError: If the session is uninitialized.
+        """
+        if self._session is None:
+            raise RuntimeError("Session is not initialized or has been shut down.")
+        return self._session
+
+    @property
+    def timeout(self) -> float:
+        """Return the default timeout setting."""
+        return self._timeout
+
+    @property
+    def retry_times(self) -> int:
+        """Return the maximum number of retry attempts."""
+        return self._retry_times
+
+    @property
+    def retry_interval(self) -> float:
+        """Return the base interval (in seconds) between retries."""
+        return self._retry_interval
+
+    async def update_cookies(
+        self, cookies: Dict[str, str], overwrite: bool = True
+    ) -> None:
+        """
+        Update cookies for the current session and internal cache.
+
+        :param cookies: New cookies to merge.
+        :param overwrite: If True, replace existing; else, only set missing.
+        """
+        # update internal cache
+        if overwrite:
+            self._cookies.update({str(k): str(v) for k, v in cookies.items()})
+        else:
+            for k, v in cookies.items():
+                self._cookies.setdefault(str(k), str(v))
+
+        # apply to live session
+        if self._session:
+            self._session.cookie_jar.update_cookies(self._cookies)
+
+    async def shutdown(self) -> None:
+        """
+        Shutdown and clean up the session. Closes connection pool.
+        """
+        if self._session:
+            await self._session.close()
+            self._session = None
+
+    def __getstate__(self) -> Dict[str, Any]:
+        """
+        Prepare object state for serialization: remove unpickleable session.
+        """
+        state = self.__dict__.copy()
+        state.pop("_session", None)
+        state.pop("_rate_limiter", None)
+        return state
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        """
+        Restore object state. Session will be lazily reinitialized on next request.
+        """
+        self.__dict__.update(state)
+        self._session = None

novel_downloader/core/requesters/base_browser.py

@@ -111,7 +111,7 @@ class BaseBrowser(RequesterProtocol, abc.ABC):
         )
 
     @abc.abstractmethod
-    def get_book_info(self, book_id: str, wait_time: Optional[int] = None) -> str:
+    def get_book_info(self, book_id: str, wait_time: Optional[float] = None) -> str:
         """
         Fetch the raw HTML (or JSON) of the book info page.
 
@@ -123,7 +123,7 @@ class BaseBrowser(RequesterProtocol, abc.ABC):
 
     @abc.abstractmethod
     def get_book_chapter(
-        self, book_id: str, chapter_id: str, wait_time: Optional[int] = None
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
     ) -> str:
         """
         Fetch the raw HTML (or JSON) of a single chapter.
@@ -135,7 +135,7 @@ class BaseBrowser(RequesterProtocol, abc.ABC):
         """
         ...
 
-    def get_bookcase(self, wait_time: Optional[int] = None) -> str:
+    def get_bookcase(self, wait_time: Optional[float] = None) -> str:
         """
         Optional: Retrieve the HTML content of the authenticated user's bookcase page.
 

novel_downloader/core/requesters/base_session.py

@@ -28,7 +28,7 @@ class BaseSession(RequesterProtocol, abc.ABC):
 
     Attributes:
         _session (requests.Session): The persistent HTTP session.
-        _timeout (int): Timeout for each request in seconds.
+        _timeout (float): Timeout for each request in seconds.
     """
 
     def _init_session(
@@ -81,7 +81,7 @@ class BaseSession(RequesterProtocol, abc.ABC):
         )
 
     @abc.abstractmethod
-    def get_book_info(self, book_id: str, wait_time: Optional[int] = None) -> str:
+    def get_book_info(self, book_id: str, wait_time: Optional[float] = None) -> str:
         """
         Fetch the raw HTML (or JSON) of the book info page.
 
@@ -93,7 +93,7 @@ class BaseSession(RequesterProtocol, abc.ABC):
 
     @abc.abstractmethod
     def get_book_chapter(
-        self, book_id: str, chapter_id: str, wait_time: Optional[int] = None
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
     ) -> str:
         """
         Fetch the raw HTML (or JSON) of a single chapter.
@@ -105,7 +105,7 @@ class BaseSession(RequesterProtocol, abc.ABC):
         """
         ...
 
-    def get_bookcase(self, wait_time: Optional[int] = None) -> str:
+    def get_bookcase(self, wait_time: Optional[float] = None) -> str:
         """
         Optional: Retrieve the HTML content of the authenticated user's bookcase page.
 
@@ -171,7 +171,7 @@ class BaseSession(RequesterProtocol, abc.ABC):
         return self._session
 
     @property
-    def timeout(self) -> int:
+    def timeout(self) -> float:
         """Return the default timeout setting."""
         return self._timeout
 

novel_downloader/core/requesters/common_requester/__init__.py

@@ -9,6 +9,10 @@ request operations to novel websites. It serves as a unified access
 point to import `CommonSession` without exposing lower-level modules.
 """
 
+from .common_async_session import CommonAsyncSession
 from .common_session import CommonSession
 
-__all__ = ["CommonSession"]
+__all__ = [
+    "CommonAsyncSession",
+    "CommonSession",
+]

novel_downloader/core/requesters/common_requester/common_async_session.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+novel_downloader.core.requesters.common_requester.common_async_session
+----------------------------------------------------------------------
+
+This module defines a `CommonAsyncSession` class for handling HTTP requests
+to common novel sites **asynchronously**. It provides methods to retrieve
+raw book info pages and chapter contents using a flexible URL templating
+system defined by a site profile, with retry logic and random delays.
+"""
+
+import asyncio
+import random
+from typing import Dict, Optional
+
+from novel_downloader.config import RequesterConfig, SiteProfile
+from novel_downloader.core.requesters.base_async_session import BaseAsyncSession
+
+
+class CommonAsyncSession(BaseAsyncSession):
+    """
+    A common async session for handling site-specific HTTP requests.
+
+    :ivar _site: The unique identifier or name of the site.
+    :ivar _profile: Metadata and URL templates related to the site.
+    """
+
+    def __init__(
+        self,
+        config: RequesterConfig,
+        site: str,
+        profile: SiteProfile,
+        cookies: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """
+        Initialize a CommonAsyncSession instance.
+
+        :param config: The RequesterConfig instance containing settings.
+        :param site: The identifier or domain of the target site.
+        :param profile: The site's metadata and URL templates.
+        :param cookies: Optional cookies to preload into the session.
+        """
+        self._init_session(config=config, cookies=cookies)
+        self._site = site
+        self._profile = profile
+
+    async def get_book_info(
+        self, book_id: str, wait_time: Optional[float] = None
+    ) -> str:
+        """
+        Fetch the raw HTML of the book info page asynchronously.
+
+        Relies on BaseAsyncSession.fetch for retry logic, then sleeps with jitter.
+
+        :param book_id:   The book identifier.
+        :param wait_time: Base seconds to sleep (with 0.5-1.5x random factor).
+        :return:          The page content as a string.
+        """
+        url = self.book_info_url.format(book_id=book_id)
+        html = await self.fetch(url)
+        base = wait_time if wait_time is not None else self._config.wait_time
+        await asyncio.sleep(base * random.uniform(0.5, 1.5))
+        return html
+
+    async def get_book_chapter(
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
+    ) -> str:
+        """
+        Fetch the raw HTML of a single chapter asynchronously.
+
+        Relies on BaseAsyncSession.fetch for retry logic, then sleeps with jitter.
+
+        :param book_id:    The book identifier.
+        :param chapter_id: The chapter identifier.
+        :param wait_time:  Base seconds to sleep (with 0.5-1.5x random factor).
+        :return:           The chapter content as a string.
+        """
+        url = self.chapter_url.format(book_id=book_id, chapter_id=chapter_id)
+        html = await self.fetch(url)
+        base = wait_time if wait_time is not None else self._config.wait_time
+        await asyncio.sleep(base * random.uniform(0.5, 1.5))
+        return html
+
+    @property
+    def site(self) -> str:
+        """Return the site name."""
+        return self._site
+
+    @property
+    def book_info_url(self) -> str:
+        """Return the URL template for fetching book info."""
+        return self._profile["book_info_url"]
+
+    @property
+    def chapter_url(self) -> str:
+        """Return the URL template for fetching chapter content."""
+        return self._profile["chapter_url"]

novel_downloader/core/requesters/common_requester/common_session.py

@@ -47,7 +47,7 @@ class CommonSession(BaseSession):
         self._site = site
         self._profile = profile
 
-    def get_book_info(self, book_id: str, wait_time: Optional[int] = None) -> str:
+    def get_book_info(self, book_id: str, wait_time: Optional[float] = None) -> str:
         """
         Fetch the raw HTML (or JSON) of the book info page.
 
@@ -75,7 +75,7 @@ class CommonSession(BaseSession):
         raise RuntimeError("Unexpected error: get_book_info failed without returning")
 
     def get_book_chapter(
-        self, book_id: str, chapter_id: str, wait_time: Optional[int] = None
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
     ) -> str:
         """
         Fetch the raw HTML (or JSON) of a single chapter.

novel_downloader/core/requesters/qidian_requester/qidian_broswer.py

@@ -266,7 +266,7 @@ class QidianBrowser(BaseBrowser):
         """
         return self.QIDIAN_BOOKCASE_URL
 
-    def get_book_info(self, book_id: str, wait_time: Optional[int] = None) -> str:
+    def get_book_info(self, book_id: str, wait_time: Optional[float] = None) -> str:
         """
         Retrieve the HTML of a Qidian book info page.
 
@@ -311,7 +311,7 @@ class QidianBrowser(BaseBrowser):
             time.sleep(pause)
 
     def get_book_chapter(
-        self, book_id: str, chapter_id: str, wait_time: Optional[int] = None
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
     ) -> str:
         """
         Retrieve the HTML content of a specific chapter.
@@ -347,7 +347,7 @@ class QidianBrowser(BaseBrowser):
             logger.warning("[fetch] Error fetching chapter from '%s': %s", url, e)
             return ""
 
-    def get_bookcase(self, wait_time: Optional[int] = None) -> str:
+    def get_bookcase(self, wait_time: Optional[float] = None) -> str:
         """
         Retrieve the HTML content of the logged‑in user's Qidian bookcase page.
 

novel_downloader/core/requesters/qidian_requester/qidian_session.py

@@ -70,7 +70,7 @@ class QidianSession(BaseSession):
         3. Updates both the live ``requests.Session`` and the internal cache;
         4. Delegates the actual request to ``super().get``.
         """
-        if self._session is None:  # defensive – mirrors BaseSession check
+        if self._session is None:  # defensive - mirrors BaseSession check
             raise RuntimeError("Session is not initialized or has been shut down.")
 
         # ---- 1. refresh token cookie --------------------------------------
@@ -108,7 +108,7 @@ class QidianSession(BaseSession):
         self.get("https://www.qidian.com")
         return True
 
-    def get_book_info(self, book_id: str, wait_time: Optional[int] = None) -> str:
+    def get_book_info(self, book_id: str, wait_time: Optional[float] = None) -> str:
         """
         Fetch the raw HTML of the book info page.
 
@@ -140,7 +140,7 @@ class QidianSession(BaseSession):
         raise RuntimeError("Unexpected fall-through in get_book_info")
 
     def get_book_chapter(
-        self, book_id: str, chapter_id: str, wait_time: Optional[int] = None
+        self, book_id: str, chapter_id: str, wait_time: Optional[float] = None
     ) -> str:
         """
         Fetch the HTML of a single chapter.
@@ -174,7 +174,7 @@ class QidianSession(BaseSession):
 
         raise RuntimeError("Unexpected fall-through in get_book_chapter")
 
-    def get_bookcase(self, wait_time: Optional[int] = None) -> str:
+    def get_bookcase(self, wait_time: Optional[float] = None) -> str:
         """
         Retrieve the user's *bookcase* page.
 

novel_downloader/resources/config/settings.yaml

@@ -1,9 +1,10 @@
 # 网络请求层设置
 requests:
-  wait_time: 5                      # 每次请求等待时间 (秒)
+  wait_time: 5.0                    # 每次请求等待时间 (秒)
   retry_times: 3                    # 请求失败重试次数
-  retry_interval: 5
-  timeout: 30                       # 页面加载超时时间 (秒)
+  retry_interval: 5.0
+  timeout: 30.0                     # 页面加载超时时间 (秒)
+  max_rps: null                     # 最大请求速率 (requests per second), 为 null 则不限制
   # DrissionPage 专用设置
   headless: false                   # 是否以无头模式启动浏览器
   user_data_folder: ""              # 浏览器用户数据目录: 为空则使用默认目录
@@ -14,15 +15,29 @@ requests:
 
 # 全局通用设置
 general:
-  request_interval: 5               # 同一本书各章节请求间隔 (秒)
+  request_interval: 5.0             # 同一本书各章节请求间隔 (秒)
   raw_data_dir: "./raw_data"        # 原始章节 HTML/JSON 存放目录
   output_dir: "./downloads"         # 最终输出文件存放目录
   cache_dir: "./novel_cache"        # 本地缓存目录 (字体 / 图片等)
-  max_threads: 4                    # 最大并发下载线程数 (未实现)
+  download_workers: 4               # 并发下载线程数
+  parser_workers: 4                 # 并发解析线程数
+  use_process_pool: false           # 是否使用多进程池来处理任务
   skip_existing: true               # 是否跳过已存在章节
   debug:
     save_html: false                # 是否将抓取到的原始 HTML 保留到磁盘
     log_level: "INFO"               # 日志级别: DEBUG, INFO, WARNING, ERROR
+  font_ocr:
+    decode_font: false              # 是否尝试本地解码混淆字体
+    use_freq: false                 # 是否使用频率分析
+    ocr_version: "v2.0"             # "v1.0" / "v2.0"
+    use_ocr: true                   # 是否使用 OCR 辅助识别文本
+    use_vec: false                  # 是否使用 Vector 辅助识别文本
+    save_font_debug: false          # 是否保存字体解码调试数据
+    batch_size: 32
+    gpu_mem: 500                    # GPU 显存限制 (MB)
+    gpu_id: null                    # 使用哪个 GPU
+    ocr_weight: 0.6
+    vec_weight: 0.4
 
 # 各站点的特定配置
 sites:
@@ -35,15 +50,6 @@ sites:
       - "0000000000"
     mode: "browser"                 # browser / session
     login_required: true            # 是否需要登录才能访问
-    decode_font: false              # 是否尝试本地解码混淆字体
-    use_freq: false                 # 是否使用频率分析
-    ocr_version: "v2.0"             # "v1.0" / "v2.0"
-    use_ocr: true                   # 是否使用 OCR 辅助识别文本
-    use_vec: false                  # 是否使用 Vector 辅助识别文本
-    save_font_debug: false          # 是否保存字体解码调试数据
-    batch_size: 32
-    ocr_weight: 0.6
-    vec_weight: 0.4
     #
   sample_site:
     book_ids:

novel_downloader/utils/crypto_utils.py

@@ -108,17 +108,17 @@ def patch_qd_payload_token(
     if not key:
         key = _get_key()
 
-    # Step 1 – decrypt --------------------------------------------------
+    # Step 1 - decrypt --------------------------------------------------
     decrypted_json: str = rc4_crypt(key, enc_token, mode="decrypt")
     payload: Dict[str, Any] = json.loads(decrypted_json)
 
-    # Step 2 – rebuild timing fields -----------------------------------
+    # Step 2 - rebuild timing fields -----------------------------------
     loadts = int(time.time() * 1000)  # ms since epoch
     # Simulate the JS duration: N(600, 150)  pushed into [300, 1000]
     duration = max(300, min(1000, int(random.normalvariate(600, 150))))
     timestamp = loadts + duration
 
-    # Step 3 – recalculate ------------------------------------
+    # Step 3 - recalculate ------------------------------------
     fp_key = _d("ZmluZ2VycHJpbnQ=")
     ab_key = _d("YWJub3JtYWw=")
     ck_key = _d("Y2hlY2tzdW0=")
@@ -138,7 +138,7 @@ def patch_qd_payload_token(
         ck_key: ck_val,
     }
 
-    # Step 4 – encrypt and return --------------------------------------
+    # Step 4 - encrypt and return --------------------------------------
     return rc4_crypt(
         key, json.dumps(new_payload, separators=(",", ":")), mode="encrypt"
     )

novel_downloader/utils/fontocr/ocr_v2.py

@@ -221,6 +221,8 @@ class FontOCRV2:
         use_ocr: bool = True,
         use_vec: bool = False,
         batch_size: int = 32,
+        gpu_mem: int = 500,
+        gpu_id: Optional[int] = None,
         ocr_weight: float = 0.6,
         vec_weight: float = 0.4,
         ocr_version: str = "v1.0",
@@ -232,6 +234,8 @@ class FontOCRV2:
         self.use_ocr = use_ocr
         self.use_vec = use_vec
         self.batch_size = batch_size
+        self.gpu_mem = gpu_mem
+        self.gpu_id = gpu_id
         self.ocr_weight = ocr_weight
         self.vec_weight = vec_weight
         self.ocr_version = ocr_version
@@ -279,6 +283,8 @@ class FontOCRV2:
             rec_batch_num=self.batch_size,
             use_space_char=False,
             use_gpu=gpu_available,
+            gpu_mem=self.gpu_mem,
+            gpu_id=self.gpu_id,
         )
 
     def _load_char_freq_db(self) -> bool: