#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : AI.  @by PyCharm
# @File         : images
# @Time         : 2025/11/23 23:27
# @Author       : betterme
# @WeChat       : meutils
# @Software     : PyCharm
# @Description  : 

from meutils.pipe import *
from meutils.llm.openai_utils import to_openai_params
from meutils.schemas.image_types import ImageRequest, ImagesResponse
from openai import AsyncOpenAI


async def generate(request: ImageRequest, api_key: str = None, base_url: str = None):
    base_url = base_url or "https://api.aimlapi.com/v1"
    client = AsyncOpenAI(
        api_key=api_key,
        base_url=base_url,
    )
    request.prompt = request.cmdprompt

    extra_body = {"num_images": request.n}
    if request.model.startswith('google/gemini') and request.image_urls:
        request.model = f"{request.model.strip('-edit')}-edit"
        extra_body = {"image_urls": request.image_urls}

    data = to_openai_params(request)
    data["extra_body"] = {**data['extra_body'], **extra_body}

    any2log(bjson(data))

    response = await client.images.generate(**data)
    any2log(response)

    return response


if __name__ == "__main__":
    # main()
    api_key = "bbce99beae764486b59ed11daff6cc14"
    api_key = "a63443199c3e42ea90003e0261ccb246"

    data = {
        "prompt": "带个墨镜",
        "model": "google/gemini-3-pro-image-preview",
        "image_urls": [
            "https://s3.ffire.cc/files/jimeng.jpg",
        ],
        "aspect_ratio": "2:3"
    }

    request = ImageRequest(**data, image=data["image_urls"])

    arun(generate(request, api_key))
