app.py

import gradio as gr
import requests
import io
import random
import os
from PIL import Image
from deep_translator import GoogleTranslator
import json
from langdetect import detect

last_prompt = None
last_negative_prompt = None
last_translation = None
last_image_description = None

api_base = os.getenv("API_BASE")
mmodels = {
    "1": "openskyml/dalle-3-xl",
    "2": "playgroundai/playground-v2-1024px-aesthetic",
    "Openjourney 4": "prompthero/openjourney-v4",
    "AbsoluteReality 1.8.1": "digiplay/AbsoluteReality_v1.8.1",
    "Lyriel 1.6": "stablediffusionapi/lyrielv16",
    "Animagine XL 2.0": "Linaqruf/animagine-xl-2.0",
    "Counterfeit 2.5": "gsdf/Counterfeit-V2.5",
    "Realistic Vision 5.1": "stablediffusionapi/realistic-vision-v51",
    "Incursios 1.6": "digiplay/incursiosMemeDiffusion_v1.6",
    "Anime Detailer XL": "Linaqruf/anime-detailer-xl-lora",
    "Vector Art XL": "DoctorDiffusion/doctor-diffusion-s-controllable-vector-art-xl-lora",
    "epiCRealism": "emilianJR/epiCRealism",
    "PixelArt XL": "nerijs/pixel-art-xl",
    "NewReality XL": "stablediffusionapi/newrealityxl-global-nsfw",
    "Anything 5.0": "hogiahien/anything-v5-edited",
    "Disney": "goofyai/disney_style_xl",
    "CleanLinearMix": "digiplay/CleanLinearMix_nsfw",
    "Redmond SDXL": "artificialguybr/LogoRedmond-LogoLoraForSDXL-V2",
    "Arcane": "nitrosocke/Arcane-Diffusion"
}
def query(prompt, is_negative=False, model="1", gpt=True, steps=30, cfg_scale=7, sampler="DPM++ 2M Karras", seed=-1, strength=0.7, width=1024, height=1024):
    global last_prompt, last_negative_prompt, last_translation, last_image_description
    if prompt == "" or prompt == None:
        raise gr.Error("🙏 Опиши пожалуйста как должно выглядеть изображение")
        return None

    key = random.randint(0, 999)
    if gpt:
        if prompt != last_prompt:
            last_prompt = prompt
            payload = {
                "model": "gpt-3.5-turbo",
                "messages": [
                    {
                        "role": "system",
                        "content": "В этом чате тебе будут присылать prompt для нейросетей генераторов изображений. А ты должен будешь пожалуйста улучшить этот prompt. Если во входном сообщении будет что то NSFW, то придумай свой prompt, а не улучшай этот. Твоя задача только улучшить его и всё. Отвечай чисто (ТОЛЬКО улучшеный prompt, без лишнего). Prompt пиши ТОЛЬКО на английском языке пожалуйста.",
                    },
                    {
                        "role": "user",
                        "content": prompt,
                    }
                ],
                "max_tokens": 500,
            }    
    
         
            # API ключ для OpenAI
            api_key_oi = os.getenv("API_KEY_OPENAI")

            # Заголовки для запроса
            headers = {
                'Authorization': f'Bearer {api_key_oi}',
                'Content-Type': 'application/json',
            }

            # URL для запроса к API OpenAI
            url = "https://api.openai.com/v1/chat/completions"

            # Отправляем запрос в OpenAI
            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}")
                    raise gr.Info("🙄 Помогатор куда-то отошёл. Придётся без него...")
            else:
                # Если произошла ошибка, возвращаем сообщение об ошибке
                print(f"Error: {response.status_code} - {response.text}")
            last_image_description = prompt
        else:
            prompt = last_image_description
            print(f"Используем предыдущее описание изображения: {prompt}")
        
    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN"), os.getenv("HF_READ_TOKEN_2"), os.getenv("HF_READ_TOKEN_3"), os.getenv("HF_READ_TOKEN_4"), os.getenv("HF_READ_TOKEN_5")]) # it is free
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    if is_negative != last_negative_prompt:
        last_negative_prompt = is_negative
        language = detect(is_negative)
        
        if language != 'en':
            last_translation = GoogleTranslator(source=language, target='en').translate(is_negative)
            print(f'\033[1mГенерация {key} перевод negative:\033[0m {last_translation}')
    else:
        print(f"Используем предыдущий перевод для negative prompt: {last_translation}")

    if prompt != last_image_description:
        language2 = detect(prompt)
        
        if language2 != 'en':
            prompt = GoogleTranslator(source=language2, target='en').translate(prompt)
            last_image_description = prompt
            print(f'\033[1mГенерация {key} перевод negative:\033[0m {last_image_description})')
    else:
        print(f"Используем предыдущий перевод для negative prompt: {last_image_description}")

    prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
    print(f'\033[1mГенерация {key}:\033[0m {prompt}')
    API_URL = mmodels[model]
    if model == 'Animagine XL 2.0':
        prompt = f"Anime. {prompt}"
    if model == 'Anime Detailer XL':
        prompt = f"Anime. {prompt}"
    if model == 'Disney':
        prompt = f"Disney style. {prompt}"

    is_negative = f"{is_negative}, NSFW, porn, sex, 18+, vagina, sexy, naked"
    
    
    payload = {
        "inputs": prompt,
        "is_negative": is_negative,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "seed": seed if seed != -1 else random.randint(1, 1000000000),
        "strength": strength,
        "width": width,
        "height": height,
        "guidance_scale": cfg_scale,
        "num_inference_steps": steps,
        "resolution": f"{width} x {height}",
        "negative_prompt": is_negative
        }

    response = requests.post(f"{api_base}{API_URL}", headers=headers, json=payload, timeout=120)
    if response.status_code != 200:
        print(f"Ошибка: Не удалось получить изображение. Статус ответа: {response.status_code}")
        print(f"Содержимое ответа: {response.text}")
        if response.status_code == 503:
            raise gr.Error(f"😭 Прости... Но сервера перегружены...")
            return None
        if response.status_code == 500:
            raise gr.Error(f"🙄 Опять эти сервера сбоят, прости..")
            return None
        raise gr.Error(f"😫 Прости, произошла ошибочка.. Попробуй ещё раз позже.")
        return None
    
    try:
        image_bytes = response.content
        image = Image.open(io.BytesIO(image_bytes))
        print(f'\033[1mГенерация {key} завершена!\033[0m ({prompt})')
        return image
    except Exception as e:
        print(f"Ошибка при попытке открыть изображение: {e}")
        raise gr.Error(f"🥺 Ой, прости, я потерял фото...")
        return None

css = """
* {}
footer, svelte-17yhekk {visibility: hidden !important;}
"""

with gr.Blocks(css=css, theme="YTheme/KidsShool") as dalle:
    with gr.Row():
        with gr.Column():
            with gr.Tab("🏠 Базовые настройки"):
                with gr.Row():
                    with gr.Column(elem_id="prompt-container"):
                        with gr.Row():
                            text_prompt = gr.Textbox(label="🌅 Описание изображения", placeholder="Реалистичная кошка сидит на подоконнике...", lines=3, elem_id="prompt-text-input")
                

            with gr.Tab("⚙️ Расширенные настройки"):
                with gr.Row():
                    negative_prompt = gr.Textbox(label="🚫 Чего не должно быть на фото?", placeholder="Текст, кошка", value="Искажения, пиксельность, блюр, размыто, плохое качество, текст", lines=3, elem_id="negative-prompt-text-input")
                with gr.Row():
                    model = gr.Radio(label="Версия нейросети", choices=["1", "2"], value="1")
                with gr.Row():
                    gpt = gr.Checkbox(label="😊 Помогатор (Просить Помогатора улучшить описание?)", value=True)

            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, model], outputs=image_output, concurrency_limit=44)
dalle.launch(show_api=False, share=False)