"""o1key 临时素材上传工具。 只供后续生成请求本来就使用公网 URL 的节点调用。图片、视频和音频会通过 ``POST /v1/o1key/uploads`` 以 multipart/form-data 上传，并返回临时 HTTPS URL。 各节点已有的数量、尺寸、时长和体积校验仍由节点自身负责；本模块不新增统一 文件大小限制，也不替调用方缩放或转码素材。 """ import asyncio import io import json import os import re import uuid from typing import Optional import aiohttp from .config import get_api_key_or_raise, get_api_base_url from .video_task import check_interrupt, run_with_interrupt _UPLOAD_ENDPOINT = "/v1/o1key/uploads" _UPLOAD_RETRY_DELAYS = (2.0, 4.0, 8.0) _UPLOAD_TIMEOUT = aiohttp.ClientTimeout(total=180) def _redact_upload_text(value: object) -> str: text = str(value or "") text = re.sub(r"data:[^\s,;]+;base64,[A-Za-z0-9+/=_-]+", "", text) return re.sub(r"https?://[^\s\"'<>]+", "", text)[:500] def _upload_error_detail(text: str) -> str: try: payload = json.loads(text) except Exception: return _redact_upload_text(text) if isinstance(payload, dict): error = payload.get("error") or payload.get("message") or payload if isinstance(error, dict): error = error.get("message") or error.get("detail") or error return _redact_upload_text(error) return _redact_upload_text(payload) def _retry_after_seconds(value: Optional[str], fallback: float) -> float: try: return max(0.0, min(float(value), 120.0)) except (TypeError, ValueError): return fallback async def _upload_file( data: bytes, filename: str, content_type: str, base_url: Optional[str] = None, ) -> str: """上传一个附件并返回新版接口提供的临时 HTTPS URL。""" check_interrupt() api_key = get_api_key_or_raise() api_base_url = (base_url or get_api_base_url()).rstrip("/") upload_url = f"{api_base_url}{_UPLOAD_ENDPOINT}" last_error: Optional[BaseException] = None for attempt in range(len(_UPLOAD_RETRY_DELAYS) + 1): check_interrupt() form = aiohttp.FormData() form.add_field( "file", data, filename=filename, content_type=content_type, ) connector = aiohttp.TCPConnector(ssl=False) try: async with aiohttp.ClientSession(connector=connector) as session: async def _request(): async with session.post( upload_url, headers={"Authorization": f"Bearer {api_key}"}, data=form, timeout=_UPLOAD_TIMEOUT, ) as response: return response.status, await response.text(), dict(response.headers) status, text, response_headers = await run_with_interrupt(_request()) if status in (200, 201): try: payload = json.loads(text) except Exception: raise RuntimeError("临时素材上传响应不是有效 JSON。") from None public_url = str(payload.get("url") or "").strip() if not public_url.startswith("https://"): raise RuntimeError("临时素材上传响应缺少有效 HTTPS URL。") check_interrupt() print( f"[素材上传] {filename} 上传完成：{len(data) / 1024:.1f} KB" ) return public_url detail = _upload_error_detail(text) retryable = status == 429 or 500 <= status < 600 if not retryable or attempt >= len(_UPLOAD_RETRY_DELAYS): retry_after = response_headers.get("Retry-After") suffix = f"，Retry-After={retry_after}s" if retry_after else "" raise RuntimeError( f"临时素材上传失败 HTTP {status}{suffix}: {detail}" ) delay = _retry_after_seconds( response_headers.get("Retry-After"), _UPLOAD_RETRY_DELAYS[attempt], ) except (aiohttp.ClientError, asyncio.TimeoutError) as error: last_error = error if attempt >= len(_UPLOAD_RETRY_DELAYS): break delay = _UPLOAD_RETRY_DELAYS[attempt] print( f"[素材上传] 上传暂时失败，{delay:.0f}s 后重试 " f"({attempt + 1}/{len(_UPLOAD_RETRY_DELAYS)})..." ) await asyncio.sleep(delay) raise RuntimeError( f"临时素材上传失败（已重试 {len(_UPLOAD_RETRY_DELAYS)} 次）: {last_error}" ) from None async def upload_image(pil_image, base_url: Optional[str] = None) -> str: """ 接受 PIL Image 对象，编码为 PNG 上传并返回临时公网 URL。 """ check_interrupt() import io as _io buf = _io.BytesIO() pil_image.save(buf, format="PNG") data = buf.getvalue() filename = f"{uuid.uuid4()}.png" return await _upload_file(data, filename, "image/png", base_url=base_url) async def upload_video(video, base_url: Optional[str] = None) -> str: """ 接受 ComfyUI VIDEO 对象，上传并返回临时公网 URL。 支持 mp4 / mov 格式。 """ check_interrupt() if hasattr(video, "get_stream_source"): source = video.get_stream_source() elif isinstance(video, dict): source = ( video.get("video") or video.get("path") or video.get("file") or video.get("filename") or video.get("source_path") ) elif isinstance(video, str): source = video else: source = None for attr in ("source_path", "path", "video", "file", "filename"): if hasattr(video, attr): source = getattr(video, attr) break if isinstance(source, io.BytesIO): source.seek(0) data = source.read() ext = "mp4" else: video_path = source if not video_path or not os.path.isfile(video_path): raise ValueError(f"无法获取参考视频文件路径（当前路径：{video_path}）") ext = os.path.splitext(video_path)[1].lower().lstrip(".") if ext not in ("mp4", "mov"): raise ValueError(f"参考视频格式须为 mp4 或 mov，当前为 .{ext}") with open(video_path, "rb") as f: data = f.read() check_interrupt() content_type = "video/mp4" if ext == "mp4" else "video/quicktime" filename = f"{uuid.uuid4()}.{ext}" return await _upload_file(data, filename, content_type, base_url=base_url) async def upload_audio(audio, base_url: Optional[str] = None) -> str: """ 接受 ComfyUI AUDIO dict（waveform tensor + sample_rate）或音频文件路径。 AUDIO 对象编码为 WAV；文件路径保留原格式上传。 """ check_interrupt() if isinstance(audio, (str, os.PathLike)): audio_path = os.fspath(audio) if not os.path.isfile(audio_path): raise ValueError(f"无法获取参考音频文件路径（当前路径：{audio_path}）") ext = os.path.splitext(audio_path)[1].lower() content_types = { ".wav": "audio/wav", ".mp3": "audio/mpeg", ".m4a": "audio/mp4", ".aac": "audio/aac", ".flac": "audio/flac", ".ogg": "audio/ogg", } content_type = content_types.get(ext) if content_type is None: supported = "/".join(item.removeprefix(".") for item in content_types) raise ValueError(f"参考音频格式须为 {supported}，当前为 {ext or '无扩展名'}") with open(audio_path, "rb") as file: data = file.read() filename = f"{uuid.uuid4()}{ext}" return await _upload_file(data, filename, content_type, base_url=base_url) import struct import numpy as np waveform = audio["waveform"] # shape: [B, C, N] or [C, N] sample_rate = int(audio["sample_rate"]) if waveform.dim() == 3: waveform = waveform[0] wav_np = waveform.cpu().numpy() if wav_np.ndim == 2: wav_np = wav_np.mean(axis=0) wav_np = np.clip(wav_np, -1.0, 1.0) pcm = (wav_np * 32767).astype(np.int16) num_samples = len(pcm) num_channels = 1 bits_per_sample = 16 byte_rate = sample_rate * num_channels * bits_per_sample // 8 block_align = num_channels * bits_per_sample // 8 data_size = num_samples * block_align buf = io.BytesIO() buf.write(b"RIFF") buf.write(struct.pack("