#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : AI.  @by PyCharm
# @File         : images
# @Time         : 2025/11/18 16:29
# @Author       : betterme
# @WeChat       : meutils
# @Software     : PyCharm
# @Description  : todo black-forest-labs/flux-1.1-pro-ultra
# https://replicate.com/black-forest-labs/flux-1.1-pro-ultra/api

from meutils.pipe import *
from meutils.schemas.image_types import ImageRequest, ImagesResponse
from meutils.apis.translator import deeplx

import replicate


async def generate(request: ImageRequest, api_key: str = None, base_url: str = None):
    api_key = api_key or os.getenv("REPLICATE_API_KEY")
    client = replicate.client.Client(api_token=api_key)

    payload = {
        "prompt": request.prompt,
        "aspect_ratio": request.aspect_ratio or "match_input_image",
        "output_format": request.output_format or "png",
    }
    if request.resolution:
        payload["resolution"] = request.resolution

    if request.model.startswith(("google/nano-banana",)):
        payload["image_input"] = request.image_urls

    elif request.model.startswith(("flux-kontext",)):
        payload["prompt"] = await deeplx.llm_translate(request.prompt)

        if request.model.endswith("-max"):
            request.model = "black-forest-labs/flux-kontext-max"
        else:
            request.model = "black-forest-labs/flux-kontext-pro"

        if len(request.image_urls) == 1:
            request.model = "black-forest-labs/flux-kontext-dev"  ##########
            payload["input_image"] = request.image_urls[0]

        elif len(request.image_urls) > 1:
            request.model = request.model.replace("black-forest-labs/flux-", "flux-kontext-apps/multi-image-")

            payload["input_image_1"] = request.image_urls[0]
            payload["input_image_2"] = request.image_urls[1]

    logger.debug(request.model)
    logger.debug(bjson(payload))

    output = await client.async_run(
        request.model,
        input=payload
    )

    return ImagesResponse(image=output.url)


#
# output = client.async_run(
#     "flux-kontext-apps/multi-image-kontext-pro",
#     input={
#         "prompt": "Put the woman next to the house",
#         "aspect_ratio": "match_input_image",
#         "input_image_1": "https://replicate.delivery/pbxt/N7gRAUNcVF6HarL0hdAQA2JYNMlJD52LP1wyaIWRUXWeHzqT/0_1-1.webp",
#         "input_image_2": "https://replicate.delivery/pbxt/N7gRAK5kbPwdsbOpqgyAIOFQX45U6suTlbL6ws2N74SnGFpo/test.jpg",
#         "output_format": "png",
#         "safety_tolerance": 2
#     }
# )

# To access the file URL:
# print(output.url)
# => "http://example.com"

if __name__ == '__main__':
    # image_request = ImageRequest(
    #     prompt="Put the woman next to the house",
    #     aspect_ratio="match_input_image",
    #     input_image_1="https://replicate.delivery/pbxt/N7gRAUNcVF6HarL0hdAQA2JYNMlJD52LP1wyaIWRUXWeHzqT/0_1-1.webp",
    #     input_image_2="https://replicate.delivery/pbxt/N7gRAK5kbPwdsbOpqgyAIOFQX45U6suTlbL6ws2N74SnGFpo/test.jpg",
    #     output_format="png",
    #     safety_tolerance=2
    # )

    # arun(output)
    model = "google/nano-banana-pro"
    model = "flux-kontext-apps/multi-image-kontext-pro"
    model = "flux-kontext-pro"
    model = "google/imagen-3"
    model = "google/imagen-4"
    model = "black-forest-labs/flux-1.1-pro-ultra"
    request = ImageRequest(
        model=model,
        # prompt="Put the woman next to the house",
        # prompt="一个裸体女人",
        prompt='带个墨镜',
        size="1x1",

        # aspect_ratio="match_input_image",
        # input_image_1="https://replicate.delivery/pbxt/N7gRAUNcVF6HarL0hdAQA2JYNMlJD52LP1wyaIWRUXWeHzqT/0_1-1.webp",
        # input_image_2="https://replicate.delivery/pbxt/N7gRAK5kbPwdsbOpqgyAIOFQX45U6suTlbL6ws2N74SnGFpo/test.jpg",
        # image="https://s3.ffire.cc/files/jimeng.jpg",
        # image="https://replicate.delivery/pbxt/N7gRAUNcVF6HarL0hdAQA2JYNMlJD52LP1wyaIWRUXWeHzqT/0_1-1.webp"
    )

    arun(generate(request))
