Ma/sora2api
1
0
Fork 0
mirror of https://github.com/TheSmallHanCat/sora2api.git synced 2026-02-15 19:14:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 为token新增图片、视频开关

Browse Source
This commit is contained in:
TheSmallHanCat
2025-11-09 09:57:10 +08:00
parent 8758447b34
commit 7269e3fa79
8 changed files with 114 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/services/load_balancer.py
View File

@@ -19,8 +19,8 @@ class LoadBalancer:
Select a token using random load balancing
Args:
for_image_generation: If True, only select tokens that are not locked for image generation
for_video_generation: If True, filter out tokens with Sora2 quota exhausted (sora2_cooldown_until not expired) and tokens that don't support Sora2
for_image_generation: If True, only select tokens that are not locked for image generation and have image_enabled=True
for_video_generation: If True, filter out tokens with Sora2 quota exhausted (sora2_cooldown_until not expired), tokens that don't support Sora2, and tokens with video_enabled=False
Returns:
Selected token or None if no available tokens
@@ -35,6 +35,10 @@ class LoadBalancer:
from datetime import datetime
available_tokens = []
for token in active_tokens:
# Skip tokens that don't have video enabled
if not token.video_enabled:
continue
# Skip tokens that don't support Sora2
if not token.sora2_supported:
continue
@@ -57,10 +61,14 @@ class LoadBalancer:
active_tokens = available_tokens
# If for image generation, filter out locked tokens
# If for image generation, filter out locked tokens and tokens without image enabled
if for_image_generation:
available_tokens = []
for token in active_tokens:
# Skip tokens that don't have image enabled
if not token.image_enabled:
continue
if not await self.token_lock.is_locked(token.id):
available_tokens.append(token)

19
src/services/token_manager.py
View File

@@ -494,7 +494,9 @@ class TokenManager:
st: Optional[str] = None,
rt: Optional[str] = None,
remark: Optional[str] = None,
update_if_exists: bool = False) -> Token:
update_if_exists: bool = False,
image_enabled: bool = True,
video_enabled: bool = True) -> Token:
"""Add a new Access Token to database
Args:
@@ -503,6 +505,8 @@ class TokenManager:
rt: Refresh Token (optional)
remark: Remark (optional)
update_if_exists: If True, update existing token instead of raising error
image_enabled: Enable image generation (default: True)
video_enabled: Enable video generation (default: True)
Returns:
Token object
@@ -634,7 +638,9 @@ class TokenManager:
sora2_invite_code=sora2_invite_code,
sora2_redeemed_count=sora2_redeemed_count,
sora2_total_count=sora2_total_count,
sora2_remaining_count=sora2_remaining_count
sora2_remaining_count=sora2_remaining_count,
image_enabled=image_enabled,
video_enabled=video_enabled
)
# Save to database
@@ -704,8 +710,10 @@ class TokenManager:
token: Optional[str] = None,
st: Optional[str] = None,
rt: Optional[str] = None,
remark: Optional[str] = None):
"""Update token (AT, ST, RT, remark)"""
remark: Optional[str] = None,
image_enabled: Optional[bool] = None,
video_enabled: Optional[bool] = None):
"""Update token (AT, ST, RT, remark, image_enabled, video_enabled)"""
# If token (AT) is updated, decode JWT to get new expiry time
expiry_time = None
if token:
@@ -715,7 +723,8 @@ class TokenManager:
except Exception:
pass # If JWT decode fails, keep expiry_time as None
await self.db.update_token(token_id, token=token, st=st, rt=rt, remark=remark, expiry_time=expiry_time)
await self.db.update_token(token_id, token=token, st=st, rt=rt, remark=remark, expiry_time=expiry_time,
image_enabled=image_enabled, video_enabled=video_enabled)
async def get_active_tokens(self) -> List[Token]:
"""Get all active tokens (not cooled down)"""