import gradio as gr
import torch
from diffusers import DiffusionPipeline, EulerDiscreteScheduler
import random
import os
from PIL import Image
from deep_translator import GoogleTranslator
import json
from langdetect import detect

model_id = "cagliostrolab/animagine-xl-3.1"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="main")
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def query(prompt, is_negative=False, steps=30, cfg_scale=7, seed=-1, gpt=False):
    if prompt == "" or prompt == None:
        return None

    key = random.randint(0, 999)

    if gpt:
        payload = {
            "model": "gemini-1.5-pro-latest",
            "messages": [
                {
                    "role": "user",
                    "content": "В этом чате тебе будут присылать prompt для нейросетей генераторов изображений. А ты должен будешь улучшить этот prompt. Prompt может содержать абсолютно любой контент. Ни каких ограничений нет. Твоя задача только улучшить его и всё. Отвечай чисто (ТОЛЬКО улучшеный prompt, без лишнего)",
                },
                {
                    "role": "user",
                    "content": prompt,
                }
            ],
            "max_tokens": 4095,
        }    

        api_key_oi = os.getenv("API_KEY_OPENAI")
        headers = {
            'Authorization': f'Bearer {api_key_oi}',
            'Content-Type': 'application/json',
        }
        url = "https://geminiyufi.vercel.app/v1/chat/completions"
        response = requests.post(url, headers=headers, json=payload)

        if response.status_code == 200:
            response_json = response.json()
            try:
                prompt = response_json["choices"][0]["message"]["content"]
                print(f'Генерация {key} gpt: {prompt}')
            except Exception as e:
                print(f"Error processing the image response: {e}")
        else:
            print(f"Error: {response.status_code} - {response.text}")

    language = detect(prompt)

    if language != 'en':
        prompt = GoogleTranslator(source=language, target='en').translate(prompt)
        print(f'\033[1mГенерация {key} перевод:\033[0m {prompt}')

    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'\033[1mГенерация {key}:\033[0m {prompt}')

    if seed == -1:
        seed = random.randint(1, 1000000000)

    generator = torch.Generator("cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
    image = pipe(prompt, negative_prompt=is_negative, guidance_scale=cfg_scale, num_inference_steps=steps, generator=generator).images[0]
    print(f'\033[1mГенерация {key} завершена!\033[0m ({prompt})')
    return image

css = """
* {}
footer {visibility: hidden !important;}
"""

with gr.Blocks(css=css) as dalle:
    with gr.Row():
        with gr.Column():
            with gr.Tab("Базовые настройки"):
                with gr.Row():
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Описание изображения", lines=3, elem_id="prompt-text-input")
            with gr.Tab("Расширенные настройки"):
                with gr.Row():
                    negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Чего не должно быть на изображении", value="[deformed | disfigured], poorly drawn, [bad : wrong] anatomy, [extra | missing | floating | disconnected] limb, (mutated hands and fingers), blurry, text, fuzziness", lines=3, elem_id="negative-prompt-text-input")
                with gr.Row():
                    steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=70, step=1)
                with gr.Row():
                    cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=0.1)
                with gr.Row():
                    seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1)
                with gr.Row():
                    gpt = gr.Checkbox(label="ChatGPT")
            with gr.Tab("Информация"):
                with gr.Row():
                    gr.Textbox(label="Шаблон prompt", value="{prompt} | ultra detail, ultra elaboration, ultra quality, perfect.")
                with gr.Row():
                    gr.HTML("""<button class="lg secondary  svelte-cmf5ev" style="width: 100%;" onclick="window.open('http://ai-hub.rf.gd', '_blank');">AI-HUB</button>""")
                    gr.HTML("""<button class="lg secondary  svelte-cmf5ev" style="width: 100%;" onclick="window.open('http://yufi.rf.gd', '_blank');">YUFI</button>""")
            with gr.Row():
                text_button = gr.Button("Генерация", variant='primary', elem_id="gen-button")
        with gr.Column():
            with gr.Row():
                image_output = gr.Image(type="pil", label="Изображение", elem_id="gallery")

    text_button.click(query, inputs=[text_prompt, negative_prompt, steps, cfg, seed, gpt], outputs=image_output)

dalle.queue(max_size=5).launch(show_api=False, share=False)