Download api/image_tasks.py from zlcy/abc: direct link, hf CLI and curl.
- Browser
- Download file 18.7 kB
-
https://huggingface.co/spaces/zlcy/abc/resolve/main/api/image_tasks.py
- Command line
-
hf download hf://spaces/zlcy/abc/api/image_tasks.py
-
curl -L -o image_tasks.py https://huggingface.co/spaces/zlcy/abc/resolve/main/api/image_tasks.py
18.7 kB
| from __future__ import annotations | |
| from decimal import Decimal, ROUND_HALF_UP | |
| from fastapi import APIRouter, Header, HTTPException, Query, Request | |
| from fastapi.concurrency import run_in_threadpool | |
| from pydantic import BaseModel, Field | |
| from api.image_inputs import parse_image_edit_request, read_image_sources | |
| from api.support import require_identity, resolve_image_base_url | |
| from services.content_filter import check_request | |
| from services.config import config | |
| from services.auth_service import DAILY_GIFT_ACCOUNT_POOL_MODEL_ID, auth_service | |
| from services.image_task_service import image_task_service | |
| from services.log_service import LoggedCall | |
| from services.pool_pricing_service import pool_pricing_service | |
| from services.share_evidence_service import share_evidence_service | |
| from services.third_party_image_service import third_party_image_service | |
| def _credit_amount(value: object) -> int | float: | |
| try: | |
| amount = max(Decimal("0"), Decimal(str(value or 0))).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP) | |
| except Exception: | |
| amount = Decimal("0") | |
| return int(amount) if amount == amount.to_integral_value() else float(amount) | |
| def _third_party_activity_quote(identity: dict[str, object], *, group_id: str, base_unit_price: object) -> dict[str, object]: | |
| pricing_config = config.normal_user_third_party_image_pricing | |
| if str(identity.get("role") or "").strip().lower() != "normal" or not bool(pricing_config.get("enabled")): | |
| base = _credit_amount(base_unit_price) | |
| return { | |
| "applies": False, | |
| "target_type": "third_party_group", | |
| "third_party_group_id": group_id, | |
| "base_unit_price": base, | |
| "pay_percentage": 100, | |
| "unit_price": base, | |
| "quantity": 1, | |
| "estimated_total": base, | |
| "rule_id": "", | |
| "rule_name": "", | |
| } | |
| return pool_pricing_service.quote( | |
| identity, | |
| quantity=1, | |
| pricing_config=pricing_config, | |
| target_type="third_party_group", | |
| third_party_group_id=group_id, | |
| base_unit_price=base_unit_price, | |
| ) | |
| class ImageGenerationTaskRequest(BaseModel): | |
| client_task_id: str = Field(..., min_length=1) | |
| prompt: str = Field(..., min_length=1) | |
| model: str = "gpt-image-2" | |
| size: str | None = None | |
| quality: str | None = None | |
| resolution: str | None = None | |
| image_resolution: str | None = None | |
| third_party_image_group_id: str | None = None | |
| account_pool_strategy: str = "own_first" | |
| memory: bool = False | |
| memory_conversation_id: str = "" | |
| memory_reset: bool = False | |
| class ResumePollRequest(BaseModel): | |
| extra_timeout_secs: float | None = Field(default=None, ge=5.0, le=120.0) | |
| class ShareEvidenceRequest(BaseModel): | |
| conversation_id: str = Field(..., min_length=1) | |
| def _parse_task_ids(value: str) -> list[str]: | |
| return [item.strip() for item in value.split(",") if item.strip()] | |
| def _require_owned_task(identity: dict[str, object], conversation_id: str) -> dict[str, object]: | |
| """校验这条上游会话确实属于当前用户,返回内部任务字典。 | |
| 没有这层校验的话,任何人拿到一个别人的 conversation_id 就能把它变成公开分享链接。 | |
| ⚠️ 生图任务记录可能已被清理,但官网取证台账里还有这条会话 → 也算「属于你」。 | |
| 否则用户点「查看官网」时只会看到一句像越权警告的 404,而真实原因是 | |
| 「这条会话已过保留期,程序已按策略把官网对话删了」。 | |
| """ | |
| task = image_task_service.find_owned_conversation(identity, conversation_id) | |
| if task is None: | |
| raise HTTPException(status_code=404, detail={ | |
| "error": "这条会话的官网记录已不存在:可能已过保留期被自动清理,或不属于当前账号", | |
| "reason": "conversation_not_found", | |
| }) | |
| return task | |
| def _internal_model_limit_kwargs(identity: dict[str, object], model: object) -> dict[str, object]: | |
| if str(identity.get("role") or "").strip().lower() != "user": | |
| return {} | |
| model_id = str(model or "gpt-image-2").strip() or "gpt-image-2" | |
| return { | |
| "internal_model_limit_user_id": str(identity.get("id") or "").strip(), | |
| "internal_model_limit_model_id": model_id, | |
| "internal_model_limit_amount": 1, | |
| } | |
| def _credit_kwargs(identity: dict[str, object], model: object) -> dict[str, object]: | |
| if str(identity.get("role") or "").strip().lower() != "normal": | |
| return _internal_model_limit_kwargs(identity, model) | |
| pricing = pool_pricing_service.quote(identity, 1) | |
| return { | |
| "credit_user_id": str(identity.get("id") or "").strip(), | |
| "credit_amount": pricing["estimated_total"], | |
| "credit_model_id": DAILY_GIFT_ACCOUNT_POOL_MODEL_ID, | |
| "credit_pricing": pricing, | |
| } | |
| def _third_party_wallet_remaining(identity: dict[str, object], group_id: str, model_id: str) -> int | float: | |
| try: | |
| summary = auth_service.credit_summary(identity) | |
| except Exception: | |
| return 0 | |
| wallets = summary.get("third_party_image_wallets") if isinstance(summary, dict) else [] | |
| wallet_key = f"{group_id}:{model_id}" | |
| wallet_items = wallets if isinstance(wallets, list) else [] | |
| balances: list[int | float] = [] | |
| for wallet in wallet_items: | |
| if not isinstance(wallet, dict): | |
| continue | |
| if str(wallet.get("id") or "") == wallet_key or ( | |
| str(wallet.get("group_id") or "") == group_id | |
| and str(wallet.get("model_id") or "").strip() == model_id | |
| ): | |
| try: | |
| balances.append(_credit_amount(wallet.get("remaining"))) | |
| except (TypeError, ValueError): | |
| balances.append(0) | |
| return max(balances) if balances else 0 | |
| def _image_api_billing_kwargs( | |
| identity: dict[str, object], | |
| *, | |
| model: object, | |
| quality: object = "", | |
| resolution: object = "", | |
| image_resolution: object = "", | |
| third_party_image_group_id: object = "", | |
| ) -> dict[str, object]: | |
| resolution = str(resolution or image_resolution or "").strip().lower() | |
| group_id = str(third_party_image_group_id or "").strip() | |
| explicit_third_party = resolution in {"1k", "2k", "4k"} or bool(group_id) | |
| route = None | |
| if not explicit_third_party: | |
| find_named_route = getattr(third_party_image_service, "find_named_route", None) | |
| route = find_named_route(model, advance=False) if callable(find_named_route) else None | |
| if not explicit_third_party and route is None: | |
| return _credit_kwargs(identity, model) | |
| require_4k = resolution == "4k" | |
| if route is None and str(identity.get("role") or "").strip().lower() == "normal": | |
| find_wallet_route = getattr(third_party_image_service, "find_wallet_route_for_identity", None) | |
| if callable(find_wallet_route): | |
| route = find_wallet_route( | |
| identity, | |
| model, | |
| group_id=group_id, | |
| require_4k=require_4k, | |
| resolution=resolution, | |
| advance=False, | |
| ) | |
| if route is None: | |
| route = third_party_image_service.find_route( | |
| model, | |
| group_id=group_id, | |
| require_4k=require_4k, | |
| resolution=resolution or "1k", | |
| advance=False, | |
| ) | |
| if route is None: | |
| raise ValueError("No enabled image API model route") | |
| try: | |
| third_party_image_service.assert_access(identity, route) | |
| except PermissionError as exc: | |
| raise PermissionError(str(exc)) from exc | |
| cost = third_party_image_service.cost_for_route(route, quality, 1) | |
| common = { | |
| "resolution": str(route.model.get("resolution") or resolution or "1k"), | |
| "third_party_image_group_id": str(route.group.get("id") or ""), | |
| } | |
| route_model_id = str(route.model.get("id") or "") | |
| if str(identity.get("role") or "").strip().lower() != "normal": | |
| return { | |
| **common, | |
| **_internal_model_limit_kwargs(identity, route_model_id), | |
| } | |
| billing_mode = str(route.group.get("billing_mode") or "") | |
| route_group_id = str(route.group.get("id") or "") | |
| if billing_mode == "general_multiplier" and _third_party_wallet_remaining(identity, route_group_id, route_model_id) > 0: | |
| # Keep existing model-specific dedicated balances usable even if an | |
| # older/stale group setting still says general-only. | |
| billing_mode = "both" | |
| wallet_pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=1) | |
| wallet_cost = _credit_amount(wallet_pricing.get("estimated_total", 1)) | |
| if billing_mode == "dedicated_wallet" or ( | |
| billing_mode == "both" and cost > 0 and _third_party_wallet_remaining(identity, route_group_id, route_model_id) >= wallet_cost | |
| ): | |
| wallet_kwargs: dict[str, object] = { | |
| **common, | |
| "third_party_wallet_user_id": str(identity.get("id") or "").strip(), | |
| "third_party_wallet_group_id": route_group_id, | |
| "third_party_wallet_model_id": route_model_id, | |
| "third_party_wallet_amount": wallet_cost, | |
| "third_party_wallet_cost_per_image": _credit_amount(wallet_pricing.get("unit_price", 1)), | |
| "credit_pricing": wallet_pricing, | |
| } | |
| if billing_mode == "both" and cost > 0: | |
| fallback_pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=cost) | |
| wallet_kwargs["third_party_wallet_fallback_credit_user_id"] = str(identity.get("id") or "").strip() | |
| wallet_kwargs["third_party_wallet_fallback_credit_amount"] = fallback_pricing["estimated_total"] | |
| wallet_kwargs["third_party_wallet_fallback_credit_model_id"] = f"{route_group_id}:{route_model_id}" | |
| wallet_kwargs["third_party_wallet_fallback_credit_pricing"] = fallback_pricing | |
| return wallet_kwargs | |
| pricing = _third_party_activity_quote(identity, group_id=route_group_id, base_unit_price=cost) | |
| return { | |
| **common, | |
| "credit_user_id": str(identity.get("id") or "").strip(), | |
| "credit_amount": pricing["estimated_total"], | |
| "credit_model_id": f"{route_group_id}:{route_model_id}", | |
| "credit_pricing": pricing, | |
| } | |
| def _generation_billing_kwargs(identity: dict[str, object], body: ImageGenerationTaskRequest) -> dict[str, object]: | |
| return _image_api_billing_kwargs( | |
| identity, | |
| model=body.model, | |
| quality=body.quality, | |
| resolution=body.resolution, | |
| image_resolution=body.image_resolution, | |
| third_party_image_group_id=body.third_party_image_group_id, | |
| ) | |
| def _effective_account_pool_strategy(identity: dict[str, object], value: object) -> str: | |
| if identity.get("role") == "admin" or bool(identity.get("account_pool_enabled")): | |
| return str(value or "own_first") | |
| return "admin_only" | |
| async def filter_or_log(call: LoggedCall, text: str) -> None: | |
| try: | |
| await run_in_threadpool(check_request, text) | |
| except HTTPException as exc: | |
| call.log("调用失败", status="failed", error=str(exc.detail)) | |
| raise | |
| def create_router() -> APIRouter: | |
| router = APIRouter() | |
| async def list_image_tasks( | |
| ids: str = Query(default=""), | |
| authorization: str | None = Header(default=None), | |
| ): | |
| identity = require_identity(authorization) | |
| return await run_in_threadpool(image_task_service.list_tasks, identity, _parse_task_ids(ids)) | |
| async def create_generation_task( | |
| body: ImageGenerationTaskRequest, | |
| request: Request, | |
| authorization: str | None = Header(default=None), | |
| ): | |
| identity = require_identity(authorization) | |
| await filter_or_log(LoggedCall(identity, "/api/image-tasks/generations", body.model, "文生图任务", request_text=body.prompt), body.prompt) | |
| try: | |
| return await run_in_threadpool( | |
| image_task_service.submit_generation, | |
| identity, | |
| client_task_id=body.client_task_id, | |
| prompt=body.prompt, | |
| model=body.model, | |
| size=body.size, | |
| quality=str(body.quality or ""), | |
| base_url=resolve_image_base_url(request), | |
| memory_enabled=body.memory, | |
| memory_conversation_id=body.memory_conversation_id, | |
| memory_reset=body.memory_reset, | |
| account_pool_strategy=_effective_account_pool_strategy(identity, body.account_pool_strategy), | |
| **_generation_billing_kwargs(identity, body), | |
| ) | |
| except PermissionError as exc: | |
| raise HTTPException(status_code=403, detail={"error": str(exc)}) from exc | |
| except ValueError as exc: | |
| error_text = str(exc) | |
| status_code = 402 if "次数不足" in error_text or "not enough" in error_text.lower() or "balance" in error_text.lower() else 400 | |
| raise HTTPException(status_code=status_code, detail={"error": str(exc)}) from exc | |
| async def create_edit_task( | |
| request: Request, | |
| authorization: str | None = Header(default=None), | |
| ): | |
| identity = require_identity(authorization) | |
| payload, image_sources, mask_sources = await parse_image_edit_request(request) | |
| client_task_id = str(payload.get("client_task_id") or "").strip() | |
| if not client_task_id: | |
| raise HTTPException(status_code=400, detail={"error": "client_task_id is required"}) | |
| prompt = str(payload["prompt"]) | |
| model = str(payload["model"]) | |
| await filter_or_log(LoggedCall(identity, "/api/image-tasks/edits", model, "图生图任务", request_text=prompt), prompt) | |
| images = await read_image_sources(image_sources) | |
| masks = await read_image_sources(mask_sources) if mask_sources else None | |
| try: | |
| return await run_in_threadpool( | |
| image_task_service.submit_edit, | |
| identity, | |
| client_task_id=client_task_id, | |
| prompt=prompt, | |
| model=model, | |
| size=payload["size"], | |
| base_url=resolve_image_base_url(request), | |
| images=images, | |
| masks=masks, | |
| memory_enabled=bool(payload.get("memory")), | |
| memory_conversation_id=str(payload.get("memory_conversation_id") or ""), | |
| memory_reset=bool(payload.get("memory_reset")), | |
| account_pool_strategy=_effective_account_pool_strategy(identity, payload.get("account_pool_strategy") or "own_first"), | |
| quality=str(payload.get("quality") or ""), | |
| **_image_api_billing_kwargs( | |
| identity, | |
| model=model, | |
| quality=payload.get("quality"), | |
| resolution=payload.get("resolution"), | |
| image_resolution=payload.get("image_resolution"), | |
| third_party_image_group_id=payload.get("third_party_image_group_id"), | |
| ), | |
| ) | |
| except PermissionError as exc: | |
| raise HTTPException(status_code=403, detail={"error": str(exc)}) from exc | |
| except ValueError as exc: | |
| status_code = 402 if "次数不足" in str(exc) else 400 | |
| raise HTTPException(status_code=status_code, detail={"error": str(exc)}) from exc | |
| async def resume_image_poll( | |
| task_id: str, | |
| body: ResumePollRequest, | |
| authorization: str | None = Header(default=None), | |
| ): | |
| identity = require_identity(authorization) | |
| try: | |
| return await run_in_threadpool( | |
| image_task_service.resume_poll, | |
| identity, | |
| task_id, | |
| body.extra_timeout_secs or config.image_timeout_retry_secs, | |
| ) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc | |
| async def share_evidence_status( | |
| conversation_id: str = Query(..., min_length=1), | |
| authorization: str | None = Header(default=None), | |
| ): | |
| """查询某条会话的「官网来源」取证状态(链接、剩余有效秒数)。""" | |
| identity = require_identity(authorization) | |
| _require_owned_task(identity, conversation_id) | |
| return share_evidence_service.describe(conversation_id) | |
| async def create_share_evidence( | |
| body: ShareEvidenceRequest, | |
| authorization: str | None = Header(default=None), | |
| ): | |
| """把一条会话取成为 ChatGPT 官方匿名分享链接(幂等)。""" | |
| identity = require_identity(authorization) | |
| if not config.image_share_evidence_enabled: | |
| raise HTTPException(status_code=403, detail={"error": "后台未开启「官网来源取证」,无法生成分享链接"}) | |
| task = _require_owned_task(identity, body.conversation_id) | |
| try: | |
| await run_in_threadpool( | |
| share_evidence_service.create_share, | |
| conversation_id=body.conversation_id, | |
| owner_id=str(identity.get("id") or ""), | |
| account_email=str(task.get("account_email") or ""), | |
| ) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc | |
| except RuntimeError as exc: | |
| raise HTTPException(status_code=502, detail={"error": str(exc)}) from exc | |
| return share_evidence_service.describe(body.conversation_id) | |
| async def revoke_share_evidence( | |
| conversation_id: str = Query(..., min_length=1), | |
| authorization: str | None = Header(default=None), | |
| ): | |
| """提前撤销分享链接并回收会话。""" | |
| identity = require_identity(authorization) | |
| _require_owned_task(identity, conversation_id) | |
| try: | |
| await run_in_threadpool(share_evidence_service.revoke, conversation_id) | |
| except ValueError as exc: | |
| raise HTTPException(status_code=400, detail={"error": str(exc)}) from exc | |
| except RuntimeError as exc: | |
| raise HTTPException(status_code=502, detail={"error": str(exc)}) from exc | |
| return share_evidence_service.describe(conversation_id) | |
| return router | |