import base64 import json import logging import mimetypes import os from datetime import datetime import httpx from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from volcenginesdkarkruntime import AsyncArk from app.config import settings from app.models.image_engine import ImageEngine from app.models.generation_record import GenerationRecord from app.services.log_config import is_enabled, LOG_DIR, LOG_DATE_FORMAT, encrypt_data from app.services.error_codes import extract_error_message logger = logging.getLogger("videogen") def _log_image_request(engine, record_id: str, request_data: dict): """Log image generation request to log/AiModel/YYYY-MM-DD.log""" if not is_enabled(): return try: os.makedirs(LOG_DIR, exist_ok=True) today = datetime.now().strftime(LOG_DATE_FORMAT) log_file = os.path.join(LOG_DIR, f"{today}.log") request_str = json.dumps(request_data, ensure_ascii=False) request_encrypted = encrypt_data(request_data) entry = { "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "type": "image_gen_request", "engine": engine.name, "model": engine.model_name, "record_id": record_id, "request": request_encrypted, "request_length": len(request_str), } with open(log_file, "a", encoding="utf-8") as f: f.write(json.dumps(entry, ensure_ascii=False) + "\n") except Exception: pass def _log_image_response(record_id: str, response_data: dict, error: str | None = None): """Log image generation response to log/AiModel/YYYY-MM-DD.log""" if not is_enabled(): return try: os.makedirs(LOG_DIR, exist_ok=True) today = datetime.now().strftime(LOG_DATE_FORMAT) log_file = os.path.join(LOG_DIR, f"{today}.log") response_encrypted = encrypt_data(response_data) if response_data else "" entry = { "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "type": "image_gen_response", "record_id": record_id, "response": response_encrypted, "error": error, } with open(log_file, "a", encoding="utf-8") as f: f.write(json.dumps(entry, ensure_ascii=False) + "\n") except Exception: pass async def get_active_image_engine(db: AsyncSession) -> ImageEngine: """Get the active image engine with highest priority.""" result = await db.execute( select(ImageEngine) .where(ImageEngine.is_active == True) .order_by(ImageEngine.priority.desc()) .limit(1) ) engine = result.scalar_one_or_none() if not engine: raise ValueError("没有可用的图片引擎，请联系管理员配置") return engine def _resolve_url(url: str) -> str: """Convert local path to base64 data URI, pass through remote URLs.""" # if url.startswith("http"): # return url # file_path = os.path.join(settings.UPLOAD_LOCAL_PATH, url.replace("/uploads/", "")) # if not os.path.exists(file_path): # settings.BASE_URL + mime = mimetypes.guess_type(file_path)[0] or "application/octet-stream" # with open(file_path, "rb") as f: # b64 = base64.b64encode(f.read()).decode() # return f"data:{mime};base64,{b64}" if url.startswith(("http://", "https://", "data:")): return url return f"{settings.BASE_URL.rstrip('/')}/{url.lstrip('/')}" def submit_image_task( db, engine: ImageEngine, record: GenerationRecord, ) -> str: """Submit an image generation task via Ark SDK. Returns image_url.""" from volcenginesdkarkruntime import Ark client = Ark( base_url=engine.api_base, api_key=engine.api_key, timeout=300, ) prompt = record.optimized_prompt image_urls = [] if record.media_references: try: refs = json.loads(record.media_references) for ref in refs: ref_type = ref.get("type") ref_url = ref.get("url", "") if ref_type == "image" and ref_url: resolved = _resolve_url(ref_url) image_urls.append(resolved) except (json.JSONDecodeError, TypeError): pass request_payload = { "model": engine.model_name, "prompt": prompt, "size": record.image_size or engine.default_size, "sequential_image_generation": "disabled", "output_format": "png", "response_format": "url", "watermark": False, } if image_urls: request_payload["image"] = image_urls _log_image_request(engine, record.id, request_payload) try: result = client.images.generate( model=engine.model_name, prompt=prompt, size=record.image_size or engine.default_size, output_format="png", response_format="url", watermark=False, image=image_urls if image_urls else None, ) image_url = result.data[0].url response_data = { "model": result.model, "created": result.created, "data": [{"url": item.url, "size": item.size} for item in result.data] if result.data else [], "usage": { "generated_images": result.usage.generated_images if hasattr(result.usage, 'generated_images') else 0, "output_tokens": result.usage.output_tokens if hasattr(result.usage, 'output_tokens') else 0, "total_tokens": result.usage.total_tokens if hasattr(result.usage, 'total_tokens') else 0, } } except httpx.TimeoutException: error_msg = "图片生成超时，请稍后重试" logger.error(f"Image generation timeout for record {record.id}") _log_image_response(record.id, {}, error_msg) raise TimeoutError(error_msg) except Exception as e: error_msg = str(e) logger.error(f"Image generation failed for record {record.id}: {error_msg}") _log_image_response(record.id, {}, error_msg) raise finally: client.close() return { "image_url": image_url, "image_tokens": getattr(result.usage, "total_tokens", 0), "response_data": json.dumps(response_data, ensure_ascii=False, default=str), "error": str(result.error) if result.error else "", } async def poll_image_task_status(engine: ImageEngine, task_id: str) -> dict: """Query image task status via Ark SDK. Returns {status, image_url, response_data}.""" client = AsyncArk( base_url=engine.api_base, api_key=engine.api_key, ) result = await client.image_generation.tasks.get(task_id=task_id) await client.close() response_dict = { "id": result.id, "model": result.model, "status": result.status, "created_at": result.created_at, "updated_at": result.updated_at, } image_url = None image_tokens = 0 if result.status == "succeeded" and result.content: image_url = getattr(result.content, "image_url", None) response_dict["image_url"] = image_url response_dict["ratio"] = getattr(result, "ratio", None) response_dict["size"] = getattr(result, "size", None) usage = getattr(result, "usage", None) if usage: response_dict["usage"] = { "input_tokens": getattr(usage, "input_tokens", 0), "output_tokens": getattr(usage, "output_tokens", 0), "total_tokens": getattr(usage, "total_tokens", 0), } image_tokens = getattr(usage, "total_tokens", 0) elif result.status == "failed": response_dict["error"] = str(getattr(result, "error", "图片生成失败")) return { "status": result.status, "image_url": image_url, "image_tokens": image_tokens, "response_data": json.dumps(response_dict, ensure_ascii=False, default=str), "error": response_dict.get("error"), } async def download_image(image_url: str, dest_path: str) -> str: """Download image to local storage.""" os.makedirs(os.path.dirname(dest_path), exist_ok=True) async with httpx.AsyncClient(timeout=300) as client: async with client.stream("GET", image_url) as response: response.raise_for_status() with open(dest_path, "wb") as f: async for chunk in response.aiter_bytes(chunk_size=8192): f.write(chunk) return dest_path