chunyu
008b82055d
新应用aitool：POST /api/aitool/upscale（登录扣2积分、事务行锁防并发超扣，失败退款+failed记录）；产物存MEDIA_ROOT/aitool/upscale，返回可下载URL；provider默认本地Pillow，第三方网关走环境变量不落库。
71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
"""C-04 Provider 抽象层：本地 Pillow 实现 + HTTP 预留。"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import os
|
||
|
||
from PIL import Image
|
||
|
||
|
||
class BaseProvider:
|
||
name = 'base'
|
||
|
||
def upscale(self, raw: bytes, scale: int) -> bytes: # pragma: no cover
|
||
raise NotImplementedError
|
||
|
||
|
||
class PillowProvider(BaseProvider):
|
||
"""本地放大：LANCZOS 重采样。零 key、零外部调用，mock 链路即生产链路。"""
|
||
name = 'pillow'
|
||
MAX_INPUT_BYTES = 10 * 1024 * 1024
|
||
MAX_DIM = 4096
|
||
|
||
def upscale(self, raw: bytes, scale: int) -> bytes:
|
||
if scale not in (2, 4):
|
||
raise ValueError('仅支持 2x / 4x')
|
||
if len(raw) > self.MAX_INPUT_BYTES:
|
||
raise ValueError('图片过大（>10MB）')
|
||
img = Image.open(io.BytesIO(raw))
|
||
img.load()
|
||
if img.width > self.MAX_DIM or img.height > self.MAX_DIM:
|
||
raise ValueError('图片尺寸过大')
|
||
out = img.resize((img.width * scale, img.height * scale), Image.Resampling.LANCZOS)
|
||
buf = io.BytesIO()
|
||
fmt = (img.format or 'PNG').upper()
|
||
if fmt in ('JPEG', 'JPG'):
|
||
if out.mode in ('RGBA', 'P'):
|
||
out = out.convert('RGB')
|
||
out.save(buf, 'JPEG', quality=95, optimize=True)
|
||
return buf.getvalue()
|
||
out.save(buf, 'PNG', optimize=True)
|
||
return buf.getvalue()
|
||
|
||
|
||
class HTTPProvider(BaseProvider):
|
||
"""预留第三方网关：baseURL/key 全走环境变量。未配置时直接抛错（由视图转 503）。"""
|
||
name = 'http'
|
||
|
||
def __init__(self) -> None:
|
||
self.base = os.environ.get('AITOOL_API_BASE', '')
|
||
self.key = os.environ.get('AITOOL_API_KEY', '')
|
||
|
||
def upscale(self, raw: bytes, scale: int) -> bytes: # pragma: no cover
|
||
import urllib.request
|
||
if not self.base or not self.key:
|
||
raise RuntimeError('AITOOL_API_BASE / AITOOL_API_KEY 未配置')
|
||
req = urllib.request.Request(
|
||
f'{self.base.rstrip("/")}/upscale',
|
||
data=raw,
|
||
headers={'Authorization': f'Bearer {self.key}', 'X-Scale': str(scale)},
|
||
method='POST',
|
||
)
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
return resp.read()
|
||
|
||
|
||
def get_provider() -> BaseProvider:
|
||
which = os.environ.get('AITOOL_PROVIDER', 'pillow').lower()
|
||
if which == 'http':
|
||
return HTTPProvider()
|
||
return PillowProvider()
|