"""
Author: Zane
Date: 2024-02-25 21:27:18
Description: 
LastEditors: Zane zanekwok73@gmail.com
LastEditTime: 2024-02-26 00:05:36
"""

import gradio as gr
from gradio.themes.base import Base
from typing import Iterable
import random
from gradio.themes.utils import colors, fonts, sizes
import requests
import io
from PIL import Image
import os


API_URL = os.environ.get("MODEL_API_URL")


def query(payload):
    auth_hf_api_token = os.environ.get("AUTH_HF_API_TOKEN")
    authorization = "Bearer " + auth_hf_api_token
    headers = {"Authorization": authorization}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content


def genImage(prompt, negative_prompt, steps, cfg_scale, width, height, seed):
    image_bytes = query(
        {
            "inputs": prompt,
            "negative_prompt": negative_prompt,
            "steps": steps,
            "cfg_scale": cfg_scale,
            "sampler": "DPM++ 2M Karras",
            "width": width,
            "height": height,
            "seed": seed
        }
    )
    image = Image.open(io.BytesIO(image_bytes))
    return image


class Seafoam(Base):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.blue,
        secondary_hue: colors.Color | str = colors.blue,
        neutral_hue: colors.Color | str = colors.gray,
        spacing_size: sizes.Size | str = sizes.spacing_md,
        radius_size: sizes.Size | str = sizes.radius_md,
        text_size: sizes.Size | str = sizes.text_md,
        font: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Quicksand"),
            "ui-sans-serif",
            "sans-serif",
        ),
        font_mono: fonts.Font
        | str
        | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"),
            "ui-monospace",
            "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue,
            secondary_hue=secondary_hue,
            neutral_hue=neutral_hue,
            spacing_size=spacing_size,
            radius_size=radius_size,
            text_size=text_size,
            font=font,
            font_mono=font_mono,
        )

theme = Seafoam(radius_size=gr.themes.sizes.radius_sm).set(body_background_fill="white")


with gr.Blocks(theme=theme) as demo:
    gr.HTML(
        """
        <div style="text-align: center; margin: 0 auto;">
            <div style="display: inline-flex; align-items: center;gap: 0.8rem;font-size: 1.75rem;">
                <h1 style="font-weight: 900; margin-bottom: 7px;margin-top:5px">
                  AI Cartoon Generator
                </h1> 
            </div>
            <p>AI Character Generator creator by <a href="https://aicartoongenerator.net/" title="AI Cartoon Generator">https://aicartoongenerator.net/</a><p>
        </div>
        """
    )
    with gr.Row():
        with gr.Column():    
            prompt = gr.Textbox(label="Prompt", placeholder="a cute cat, 8k",value="a cute cat, 8k", show_label=True, lines=5)
            with gr.Row():
                with gr.Accordion("Additionals inputs", open=False):
                    with gr.Column(scale=1):
                        negative_prompt = gr.Textbox(label="Negative Prompt", value="text, blurry", placeholder="What you don't want to see in the image", show_label=True, lines=1)
                    with gr.Column(scale=1):
                        steps = gr.Slider(label="Sampling Steps", minimum=1, maximum=30, value=25, step=1)
                        cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7, step=1)
                        seed = gr.Number(label="Seed", value=-1)
                    with gr.Column(scale=1):
                        width = gr.Slider(label="↔️ Width", minimum=1024, maximum=1024, value=1024, step=8)
                        height = gr.Slider(label="↕️ Height", minimum=1024, maximum=1024, value=1024, step=8)
            text_button = gr.Button("Generate", variant='primary')
        with gr.Column():    
            image_output = gr.Image(type="pil", label="Output Image",height="auto",show_share_button=False)

        text_button.click(genImage, inputs=[prompt, negative_prompt, steps, cfg_scale, width, height, seed], concurrency_limit=100,outputs=image_output)
    with gr.Row():
        gr.Examples(
            examples=[
                "It’s spring, a super cute little girl IP sitting on a huge flower as if dreaming, sunny day, chibi, 3D,natural lighting, full body portralt, 8k best quality, super detail, super detail, Ultra HD",
                "A high-detail Pixar-style, especially Wreck-It Ralph style woman cartoon character with long brown hair and black eyes with big smile in a peaceful and happy scene, set in an environment with natural and clear weather, enhanced with bright sunshine.",
                "A high-detail Pixar-style, especially Wreck-It Ralph style woman cartoon character with long brown hair and black eyes with big smile In a joyful Christmas scene, set in a snowy environment during the night, where the twinkling lights create an atmosphere full of hope and festive cheer."
                ],
            inputs=prompt

        )


demo.queue(api_open=False).launch(show_api=False, show_error=False)