diff --git a/yuisub/__main__.py b/yuisub/__main__.py index a5c4ece..8f5a9b4 100644 --- a/yuisub/__main__.py +++ b/yuisub/__main__.py @@ -19,8 +19,9 @@ parser.add_argument("-om", "--OPENAI_MODEL", type=str, help="Openai model name", required=True) parser.add_argument("-api", "--OPENAI_API_KEY", type=str, help="Openai API key", required=True) parser.add_argument("-url", "--OPENAI_BASE_URL", type=str, help="Openai base URL", required=True) -# bangumi url +# bangumi parser.add_argument("-bgm", "--BANGUMI_URL", type=str, help="Anime Bangumi URL", required=False) +parser.add_argument("-ac","--BANGUMI_ACCESS_TOKEN", type=str, help="Anime Bangumi Access Token", required=False) # whisper parser.add_argument("-d", "--TORCH_DEVICE", type=str, help="Pytorch device to use", required=False) parser.add_argument("-wm", "--WHISPER_MODEL", type=str, help="Whisper model to use", required=False) @@ -65,6 +66,7 @@ async def main() -> None: api_key=args.OPENAI_API_KEY, base_url=args.OPENAI_BASE_URL, bangumi_url=args.BANGUMI_URL, + bangumi_access_token=args.BANGUMI_ACCESS_TOKEN, ) sub_bilingual = await bilingual(sub_origin=sub, sub_zh=sub_zh) diff --git a/yuisub/bangumi.py b/yuisub/bangumi.py index 3942fd3..399ef2e 100644 --- a/yuisub/bangumi.py +++ b/yuisub/bangumi.py @@ -94,7 +94,7 @@ async def fetch_bangumi_data(client: httpx.AsyncClient, url: str) -> tuple[str, return response_info.json()["summary"], response_chars.json() -async def bangumi(url: Optional[str] = None) -> BGM: +async def bangumi(url: Optional[str] = None, token: Optional[str] = None) -> BGM: """ Get bangumi info and character list asynchronously @@ -120,6 +120,9 @@ async def bangumi(url: Optional[str] = None) -> BGM: "Connection": "keep-alive", } + if token: + headers["Authorization"] = f"Bearer {token}" + async with httpx.AsyncClient(headers=headers, timeout=30.0) as client: try: # 获取基本信息和角色列表 diff --git a/yuisub/sub.py b/yuisub/sub.py index 66508d4..41bacc3 100644 --- a/yuisub/sub.py +++ b/yuisub/sub.py @@ -82,6 +82,7 @@ async def translate( api_key: str, base_url: str, bangumi_url: Optional[str] = None, + bangumi_access_token: Optional[str] = None, styles: Optional[Dict[str, SSAStyle]] = None, ad: Optional[SSAEvent] = advertisement(), # noqa: B008 ) -> SSAFile: @@ -93,6 +94,7 @@ async def translate( :param api_key: llm api_key :param base_url: llm base_url :param bangumi_url: anime bangumi url + :param bangumi_access_token: anime bangumi access token :param styles: subtitle styles, default is PRESET_STYLES :param ad: add advertisement to subtitle, default is TensoRaws :return: @@ -101,7 +103,7 @@ async def translate( trans_list: List[str] = [s.text for s in sub] # get bangumi info asynchronously - bangumi_info = await bangumi(bangumi_url) if bangumi_url else None + bangumi_info = await bangumi(bangumi_url, bangumi_access_token) if bangumi_url else None # initialize summarizer summarizer = Summarizer(