gpt4free/g4f/Provider/deprecated/ARTA.py

from __future__ import annotations

import os
import time
import json
import random
from pathlib import Path
from aiohttp import ClientSession, ClientResponse
import asyncio

from ...typing import AsyncResult, Messages
from ...providers.response import ImageResponse, Reasoning
from ...errors import ResponseError, ModelNotFoundError
from ...cookies import get_cookies_dir
from ..base_provider import AsyncGeneratorProvider, ProviderModelMixin
from ..helper import format_media_prompt
from ... import debug

class ARTA(AsyncGeneratorProvider, ProviderModelMixin):
    url = "https://ai-arta.com"
    auth_url = "https://www.googleapis.com/identitytoolkit/v3/relyingparty/signupNewUser?key=AIzaSyB3-71wG0fIt0shj0ee4fvx1shcjJHGrrQ"
    token_refresh_url = "https://securetoken.googleapis.com/v1/token?key=AIzaSyB3-71wG0fIt0shj0ee4fvx1shcjJHGrrQ"
    image_generation_url = "https://img-gen-prod.ai-arta.com/api/v1/text2image"
    status_check_url = "https://img-gen-prod.ai-arta.com/api/v1/text2image/{record_id}/status"

    working = False # Take down request

    default_model = "flux"
    default_image_model = default_model
    model_aliases = {
        "anything-xl": "Anything-xl",
        "high-gpt4o": "High GPT4o",
        "on-limbs-black": "On limbs black",
        "f-dev": "F Dev",
        "flux-dev": "F Dev", # Added
        "sdxl-1.0": "SDXL 1.0", # Added
        "old-school": "Old School",
        "vincent-van-gogh": "Vincent Van Gogh",
        "cor-epica-xl": "Cor-epica-xl",
        "professional": "Professional",
        "cheyenne-xl": "Cheyenne-xl",
        "chicano": "Chicano",
        "sdxl-l": "SDXL L", # Added
        "black-ink": "Black Ink",
        "juggernaut-xl": "Juggernaut-xl",
        "cinematic-art": "Cinematic Art",
        "dreamshaper-xl": "Dreamshaper-xl",
        "fantasy-art": "Fantasy Art",
        "neo-traditional": "Neo-traditional",
        "realistic-stock-xl": "Realistic-stock-xl",
        "flame-design": "Flame design",
        "japanese-2": "Japanese_2",
        "medieval": "Medieval",
        "surrealism": "Surrealism",
        "dotwork": "Dotwork",
        "graffiti": "Graffiti",
        "revanimated": "RevAnimated",
        "on-limbs-color": "On limbs color",
        "old-school-colored": "Old school colored",
        "gpt4o-ghibli": "GPT4o Ghibli",
        "low-poly": "Low Poly",
        "gpt4o": "GPT4o",
        "gpt-image": ["GPT4o", "High GPT4o", "GPT4o Ghibli"],
        "no-style": "No Style",
        "anime": "Anime",
        "tattoo": "tattoo",
        "embroidery-tattoo": "Embroidery tattoo",
        "mini-tattoo": "Mini tattoo",
        "realistic-tattoo": "Realistic tattoo",
        "playground-xl": "Playground-xl",
        "Watercolor": "Watercolor",
        "f-pro": "F Pro",
        "flux-pro": "F Pro", # Added
        "kawaii": "Kawaii",
        "photographic": "Photographic",
        "katayama-mix-xl": "Katayama-mix-xl",
        "death-metal": "Death metal",
        "new-school": "New School",
        "pony-xl": "Pony-xl",
        "anima-pencil-xl": "Anima-pencil-xl",
        default_image_model: "Flux", # Added
        "biomech": "Biomech",
        "yamers-realistic-xl": "Yamers-realistic-xl",
        "trash-polka": "Trash Polka",
        "red-and-black": "Red and Black",
    }
    image_models = list(model_aliases.keys())
    models = image_models

    @classmethod
    def get_model(cls, model: str) -> str:
        """Get the internal model name from the user-provided model name."""
        if not model:
            return cls.model_aliases[cls.default_model]
        
        # Always check aliases first to get the proper API name
        if model in cls.model_aliases:
            alias = cls.model_aliases[model]
            # If the alias is a list, randomly select one of the options
            if isinstance(alias, list):
                selected_model = random.choice(alias)
                debug.log(f"ARTA: Selected model '{selected_model}' from alias '{model}'")
                return selected_model
            debug.log(f"ARTA: Using model '{alias}' for alias '{model}'")
            return alias
        
        # If not in aliases, check if it's a direct API model name
        api_model_names = [v for v in cls.model_aliases.values() if isinstance(v, str)]
        if model in api_model_names:
            return model
        
        raise ModelNotFoundError(f"Model {model} not found")



    @classmethod
    def get_auth_file(cls):
        path = Path(get_cookies_dir())
        path.mkdir(exist_ok=True)
        filename = f"auth_{cls.__name__}.json"
        return path / filename

    @classmethod
    async def create_token(cls, path: Path, proxy: str | None = None):
        async with ClientSession() as session:
            # Step 1: Generate Authentication Token
            auth_payload = {"clientType": "CLIENT_TYPE_ANDROID"}
            async with session.post(cls.auth_url, json=auth_payload, proxy=proxy) as auth_response:
                await raise_error(f"Failed to obtain authentication token", auth_response)
                auth_data = await auth_response.json()
                auth_token = auth_data.get("idToken")
                #refresh_token = auth_data.get("refreshToken")
                if not auth_token:
                    raise ResponseError("Failed to obtain authentication token.")
                json.dump(auth_data, path.open("w"))
                return auth_data

    @classmethod
    async def refresh_token(cls, refresh_token: str, proxy: str = None) -> tuple[str, str]:
        async with ClientSession() as session:
            payload = {
                "grant_type": "refresh_token",
                "refresh_token": refresh_token,
            }
            async with session.post(cls.token_refresh_url, data=payload, proxy=proxy) as response:
                await raise_error(f"Failed to refresh token", response)
                response_data = await response.json()
                return response_data.get("id_token"), response_data.get("refresh_token")

    @classmethod
    async def read_and_refresh_token(cls, proxy: str | None = None) -> str:
        path = cls.get_auth_file()
        if path.is_file():
            auth_data = json.load(path.open("rb"))
            diff = time.time() - os.path.getmtime(path)
            expiresIn = int(auth_data.get("expiresIn"))
            if diff < expiresIn:
                if diff > expiresIn / 2:
                    auth_data["idToken"], auth_data["refreshToken"] = await cls.refresh_token(auth_data.get("refreshToken"), proxy)
                    json.dump(auth_data, path.open("w"))
                return auth_data
        return await cls.create_token(path, proxy)

    @classmethod
    async def create_async_generator(
        cls,
        model: str,
        messages: Messages,
        proxy: str = None,
        prompt: str = None,
        negative_prompt: str = "blurry, deformed hands, ugly",
        n: int = 1,
        guidance_scale: int = 7,
        num_inference_steps: int = 30,
        aspect_ratio: str = None,
        seed: int = None,
        **kwargs
    ) -> AsyncResult:
        model = cls.get_model(model)
        prompt = format_media_prompt(messages, prompt)

        # Generate a random seed if not provided
        if seed is None:
            seed = random.randint(9999, 99999999)  # Common range for random seeds

        # Step 1: Get Authentication Token
        auth_data = await cls.read_and_refresh_token(proxy)
        auth_token = auth_data.get("idToken")

        async with ClientSession() as session:
            # Step 2: Generate Images
            # Create a form data structure as the API might expect form data instead of JSON
            form_data = {
                "prompt": prompt,
                "negative_prompt": negative_prompt,
                "style": model,
                "images_num": str(n),
                "cfg_scale": str(guidance_scale),
                "steps": str(num_inference_steps),
                "aspect_ratio": "1:1" if aspect_ratio is None else aspect_ratio,
                "seed": str(seed),
            }

            headers = {
                "Authorization": auth_token,
                # No Content-Type header for multipart/form-data, aiohttp sets it automatically
            }

            # Try with form data instead of JSON
            async with session.post(cls.image_generation_url, data=form_data, headers=headers, proxy=proxy) as image_response:
                await raise_error(f"Failed to initiate image generation", image_response)
                image_data = await image_response.json()
                record_id = image_data.get("record_id")
                if not record_id:
                    raise ResponseError(f"Failed to initiate image generation: {image_data}")

            # Step 3: Check Generation Status
            status_url = cls.status_check_url.format(record_id=record_id)
            start_time = time.time()
            last_status = None
            while True:
                async with session.get(status_url, headers=headers, proxy=proxy) as status_response:
                    await raise_error(f"Failed to check image generation status", status_response)
                    status_data = await status_response.json()
                    status = status_data.get("status")

                    if status == "DONE":
                        image_urls = [image["url"] for image in status_data.get("response", [])]
                        duration = time.time() - start_time
                        yield Reasoning(label="Generated", status=f"{n} image in {duration:.2f}s" if n == 1 else f"{n} images in {duration:.2f}s")
                        yield ImageResponse(urls=image_urls, alt=prompt)
                        return
                    elif status in ("IN_QUEUE", "IN_PROGRESS"):
                        if last_status != status:
                            last_status = status
                            if status == "IN_QUEUE":
                                yield Reasoning(label="Waiting")
                            else:
                                yield Reasoning(label="Generating")
                        await asyncio.sleep(2)  # Poll every 2 seconds
                    else:
                        raise ResponseError(f"Image generation failed with status: {status}")

async def raise_error(message: str, response: ClientResponse):
    if response.ok:
        return
    error_text = await response.text()
    content_type = response.headers.get('Content-Type', 'unknown')
    raise ResponseError(f"{message}. Content-Type: {content_type}, Response: {error_text}")