gpt4free

from __future__ import annotations

import json
import uuid

from ...typing import AsyncResult, Messages
from ...providers.response import (
    ImageResponse,
    ImagePreview,
    JsonConversation,
    Reasoning,
)
from ...requests import StreamSession
from ...image import use_aspect_ratio
from ...errors import ResponseError
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..helper import format_media_prompt
from .utils import get_zerogpu_token
from .raise_for_status import raise_for_status


class BlackForestLabs_Flux1Dev(AsyncGeneratorProvider, ProviderModelMixin):
    label = "BlackForestLabs Flux-1-Dev"
    url = "https://black-forest-labs-flux-1-dev.hf.space"
    space = "black-forest-labs/FLUX.1-dev"
    referer = f"{url}/?__theme=light"

    working = True

    default_model = "black-forest-labs-flux-1-dev"
    default_image_model = default_model
    model_aliases = {"flux-dev": default_image_model, "flux": default_image_model}
    image_models = list(model_aliases.keys())
    models = image_models

    @classmethod
    def run(
        cls,
        method: str,
        session: StreamSession,
        conversation: JsonConversation,
        data: list = None,
    ):
        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "x-zerogpu-token": conversation.zerogpu_token,
            "x-zerogpu-uuid": conversation.zerogpu_uuid,
            "referer": cls.referer,
        }
        if method == "post":
            return session.post(
                f"{cls.url}/gradio_api/queue/join?__theme=light",
                **{
                    "headers": {k: v for k, v in headers.items() if v is not None},
                    "json": {
                        "data": data,
                        "event_data": None,
                        "fn_index": 2,
                        "trigger_id": 4,
                        "session_hash": conversation.session_hash,
                    },
                },
            )
        return session.get(
            f"{cls.url}/gradio_api/queue/data?session_hash={conversation.session_hash}",
            **{
                "headers": {
                    "accept": "text/event-stream",
                    "content-type": "application/json",
                    "referer": cls.referer,
                }
            },
        )

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        prompt: str = None,
        proxy: str = None,
        aspect_ratio: str = "1:1",
        width: int = None,
        height: int = None,
        guidance_scale: float = 3.5,
        num_inference_steps: int = 28,
        seed: int = 0,
        randomize_seed: bool = True,
        cookies: dict = None,
        api_key: str = None,
        zerogpu_uuid: str = "[object Object]",
        **kwargs,
    ) -> AsyncResult:
        async with StreamSession(impersonate="chrome", proxy=proxy) as session:
            prompt = format_media_prompt(messages, prompt)
            data = use_aspect_ratio({"width": width, "height": height}, aspect_ratio)
            data = [
                prompt,
                seed,
                randomize_seed,
                data.get("width"),
                data.get("height"),
                guidance_scale,
                num_inference_steps,
            ]
            conversation = JsonConversation(
                zerogpu_token=api_key,
                zerogpu_uuid=zerogpu_uuid,
                session_hash=uuid.uuid4().hex,
            )
            if conversation.zerogpu_token is None:
                (
                    conversation.zerogpu_uuid,
                    conversation.zerogpu_token,
                ) = await get_zerogpu_token(cls.space, session, conversation, cookies)
            async with cls.run(f"post", session, conversation, data) as response:
                await raise_for_status(response)
                assert (await response.json()).get("event_id")
                async with cls.run("get", session, conversation) as event_response:
                    await raise_for_status(event_response)
                    async for chunk in event_response.iter_lines():
                        if chunk.startswith(b"data: "):
                            try:
                                json_data = json.loads(chunk[6:])
                                if json_data is None:
                                    continue
                                if json_data.get("msg") == "log":
                                    yield Reasoning(status=json_data["log"])

                                if json_data.get("msg") == "progress":
                                    if "progress_data" in json_data:
                                        if json_data["progress_data"]:
                                            progress = json_data["progress_data"][0]
                                            yield Reasoning(
                                                status=f"{progress['desc']} {progress['index']}/{progress['length']}"
                                            )
                                        else:
                                            yield Reasoning(status=f"Generating")

                                elif json_data.get("msg") == "process_generating":
                                    for item in json_data["output"]["data"][0]:
                                        if isinstance(item, dict) and "url" in item:
                                            yield ImagePreview(item["url"], prompt)
                                        elif (
                                            isinstance(item, list)
                                            and len(item) > 2
                                            and "url" in item[1]
                                        ):
                                            yield ImagePreview(item[2], prompt)

                                elif json_data.get("msg") == "process_completed":
                                    if (
                                        "output" in json_data
                                        and "error" in json_data["output"]
                                    ):
                                        json_data["output"]["error"] = json_data[
                                            "output"
                                        ]["error"].split(" <a ")[0]
                                        raise ResponseError(
                                            json_data["output"]["error"]
                                        )
                                    if (
                                        "output" in json_data
                                        and "data" in json_data["output"]
                                    ):
                                        yield Reasoning(status="")
                                        if len(json_data["output"]["data"]) > 0:
                                            yield ImageResponse(
                                                json_data["output"]["data"][0]["url"],
                                                prompt,
                                            )
                                    break
                            except (json.JSONDecodeError, KeyError, TypeError) as e:
                                raise RuntimeError(
                                    f"Failed to parse message: {chunk.decode(errors='replace')}",
                                    e,
                                )