import gradio as gr
import requests
import json
import os
import time

def generate_promt(description):
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {os.getenv("API_KEY")}'
    }

    payload = {
        'messages': [{'role': 'system', 'content': f'Сгенерируй пожалуйста promt для нейросети генерирующая изображения, соенерируй по описанию \"{description}\". Должно быть описание изображения (prompt), а не просьба, приказ. Нейросеть не знает не каких персонажей, построек, городов, по этому если там будет что нибудь из этого, то опиши это отлично, понятно, идеально, в подробностях и деталях (у персонажей например подробно описывай внешность, у городов как они ввглядят, какие здания, постройки, структура, вид). В конце всегда добавляй разные: 8k, ультра детализация, ультра проработка, ультра реалистичность, anime, красиво и т.д., но что бв подходило по смыслу. Напиши ТОЛЬКО promt, без дополнительного текста, кавычек в начале и в конце.'}],
        'max_tokens': 5000,
        'model': os.getenv("MODEL")
    }

    response = requests.post(os.getenv("BASE_URL"), headers=headers, json=payload)
    data = json.loads(response.text)

    if 'choices' in data and len(data['choices']) > 0:
        command = data['choices'][0]['message']['content'].strip()
        return command
    elif 'error' in data:
        error_message = data['error']['message']
        print(f'Ошибка: {error_message}')
        time.sleep(1)
    else:
        print(f'Не удалось сгенерировать promt. {data}')
        time.sleep(1)

iface = gr.Interface(fn=generate_promt, inputs=[
    gr.Textbox(label="Описание")
], outputs=gr.Textbox(label="Promt"), title="Генератор prompt для нейросетей")
iface.launch()