app.py

import os
import time
import gradio as gr
from glm_huggingface import run_glm
from open_ai import run_openai
from  prompt_loader import get_prompt

from dotenv import load_dotenv
import sys

load_dotenv()

if os.getenv("OPENAI_API_KEY",False):
    status = "using OPENAI_API_KEY"
    run_llm=run_openai
else:
    status = "using glm from huggingface"
    run_llm=run_glm


PROMPT_STYLE = get_prompt("PROMPT_STYLE")
PROMPT_TEXT = get_prompt("PROMPT_TEXT")
PROMPT_TEXT2 = get_prompt("PROMPT_TEXT2") 
SYSPROMPT_STYLE = get_prompt("SYSPROMPT_STYLE")
SYSPROMPT_TEXT = get_prompt("SYSPROMPT_TEXT")
PROMT_GEN_LYRICS = get_prompt("PROMT_GEN_LYRICS")
SYSPROMT_GEN_LYRICS = get_prompt("SYSPROMT_GEN_LYRICS")
SYSPROMT_GEN_STYLE_FROM_LYRICS = get_prompt("SYSPROMT_GEN_STYLE_FROM_LYRICS")
PROMT_GEN_STYLE_FROM_LYRICS = get_prompt("PROMT_GEN_STYLE_FROM_LYRICS")

def get_style_lyrics(style: str, lyrics: str, gen_lyrics: bool, gen_style_from_lyrics: bool):
    # Detect current UI language from global TS variable
    current_lang = "RU" if TS == russian_strings else "EN"
    
    # Initialize timing and logging variables
    start_time = time.time()
    stage_times = {}
    errors = []

    
    # Helper function to get localized message
    def get_msg(key):
        return TS.get(key, english_strings.get(key, f"Missing key: {key}"))
    
    # Helper function to yield progress with timing
    def yield_progress(message_key, status="🔄", **kwargs):
        message = get_msg(message_key).format(**kwargs)
        elapsed = time.time() - start_time
        progress_msg = f"{status} {message} (⏱️ {elapsed:.2f}s)"
        print(f"[PROGRESS] {message}")
        return progress_msg

    log = ''
    yield ("", "", "🚀" + get_msg("log_starting"))
    
    gen_start = time.time()    
    # Stage 1: Lyrics Generation (if requested)
    if gen_lyrics:
        stage_start = time.time()
        yield ("", lyrics, log + " ⏳" + get_msg("log_lyrics_gen_start"))

        try:
            gen_lyrics_result = run_llm(f"{PROMT_GEN_LYRICS}\n<THEME>{lyrics}</THEME>", sys_prompt=SYSPROMT_GEN_LYRICS)
            if gen_lyrics_result:
                lyrics = gen_lyrics_result
            else:
                lyrics = ''            
        except Exception as e:
            stage_time = time.time() - stage_start
            stage_times['lyrics_generation'] = stage_time
            error_msg = f"ERROR infer llm: {e}"
            errors.append(error_msg)
            yield yield_progress("log_lyrics_gen_error", status="❌", error=str(e))
            print(error_msg)
            lyrics = ''
        stage_time = time.time() - stage_start
        stage_times['lyrics_generation'] = stage_time
        log += " "+yield_progress("log_lyrics_gen_complete", status=" 🖌️", time=stage_time)
        yield ("", lyrics, log + " ⏳" + get_msg("log_style_gen_start" if not gen_style_from_lyrics else "log_style_desc_start"))


    # Stage 2: Style Generation
    stage_start = time.time()
    if gen_style_from_lyrics:
        try:
            style_gen = run_llm(f"{PROMT_GEN_STYLE_FROM_LYRICS}\n<MUSIC_STYLE>{style}</MISIC_STYLE><SONG_TEXT><LYRICS>\n{lyrics}\n</LYRICS></SONG_TEXT>",
                                sys_prompt=SYSPROMT_GEN_STYLE_FROM_LYRICS)
            if not style_gen:
                style_gen = ''
        except Exception as e:
            error_msg = f"ERROR infer llm: {e}"
            errors.append(error_msg)
            yield yield_progress("log_style_gen_error", status="❌", error=str(e))
            print(error_msg)
            style_gen = ''
    else:
        try:
            style_gen = run_llm(f"Describe {style}", sys_prompt=SYSPROMPT_STYLE)
            if not style_gen:
                style_gen = ''
        except Exception as e:
            error_msg = f"ERROR infer llm: {e}"
            errors.append(error_msg)
            yield yield_progress("log_style_gen_error", status="❌", error=str(e))
            print(error_msg)
            style_gen = ''
    stage_time = time.time() - stage_start
    stage_times['style_generation'] = stage_time
    log += " "+ yield_progress("log_style_gen_complete", status=" 🎨", time=stage_time)
    yield (style_gen, lyrics, log + " ⏳" + get_msg("log_text_format_start"))


    # Stage 3: Text Formatting
    stage_start = time.time()
    try:
        text_glm = run_llm(f"{PROMPT_TEXT}\n{style_gen}\n#Text song:\n{lyrics}\n{PROMPT_TEXT2}", sys_prompt=SYSPROMPT_TEXT)
        if text_glm:
            text = text_glm
        else:
            text = ''
    except Exception as e:
        stage_time = time.time() - stage_start
        stage_times['text_formatting'] = stage_time
        error_msg = f"ERROR infer llm: {e}"
        errors.append(error_msg)
        yield yield_progress("log_text_format_error", status="❌", error=str(e))
        print(error_msg)
        text = ''

    stage_time = time.time() - stage_start
    # Final summary
    total_time = time.time() - gen_start
    log += " " + yield_progress("log_text_format_complete", status=" 🧾", time=stage_time)
    yield (style_gen, text, log)
   


isGenarateLyrics=False

# Define translatable strings
english_strings = {
    "title": "SUNO style generator and lyrics format",
    "performer_label": "Preferred Music Style:",
    "performer_placeholder": "Describe the desired music style in detail: you can blend different genres and influences, mention specific artists, instruments, rhythm, and overall mood. Specify the vocal type — male, female, high, raspy, smoky, operatic, or any other. Feel free to experiment — mix styles, evoke atmosphere, and express the emotions you want the music to convey.",
    "song_text_label": "Lyrics:",
    "song_text_placeholder": "Enter your song lyrics to format",
    "song_gen_label": "Idea for lyrics:",
    "song_gen_placeholder": "Share your inspiration — describe what you dream of singing about. Sketch out your song idea: emotions, events, places, characters — everything that matters. Tell the story you want to turn into music.",
    "language_english": "English",
    "language_russian": "Russian",
    "generate_button": "Generate",
    "style_label": "AI-Generated Style:",
    "song_result_label": "Formated song text:",
    "isGenerateLyrics": "generate lyrics (experemental)",
    "isGenerateStyle": "generate music style from lyrics",
    # Progress logging messages
    "log_starting": "Starting SUNO prompt generation...",
    "log_lyrics_gen_start": "Generating lyrics from idea...",
    "log_lyrics_gen_complete": "Lyrics generation completed in {time:.2f}s",
    "log_lyrics_gen_error": "Error generating lyrics: {error}",
    "log_style_gen_start": "Generating music style from lyrics...",
    "log_style_desc_start": "Generating music style description...",
    "log_style_gen_complete": "Style generation completed in {time:.2f}s",
    "log_style_gen_error": "Error generating style: {error}",
    "log_text_format_start": "Formatting song text...",
    "log_text_format_complete": "Text formatting completed in {time:.2f}s",
    "log_text_format_error": "Error formatting text: {error}",
    "log_summary_success": "✅ Generation completed successfully!",
    "log_summary_error": "❌ Generation completed with errors",
    "log_summary_total": "Total execution time: {time:.2f}s",
    "log_summary_stages": "Stage timings: {stages}",
    "generation_warning_msg": """**⚠️ Note:** the generation model (*GLM*) is not among the strongest ones and was chosen for its availability on *Hugging Face*.  
If the output looks off, simply try generating again — it usually fixes the issue.
"""
}

russian_strings = {
    "title": "SUNO генератор стиля музыки и форматирования структуры песни",
    "performer_label": "Желаемый музыкальный стиль:",
    "performer_placeholder": "Опишите стиль музыки: можно сочетать разные жанры и направления, конкретных исполнителей, инструменты, ритм и настроение, тип вокала — мужской, женский, высокий, надрывный, прокуренный, оперный или любой другой. Передавайте атмосферу и эмоции, которые хотите услышать.",
    "song_text_label": "Текст песни:",
    "song_text_placeholder": "Введите текст песни для форматирования",
    "song_gen_label": "Опишите идею для текста песни:",
    "song_gen_placeholder": "Поделитесь вдохновением — опишите, о чём вы мечтаете спеть. Набросайте идею для песни: эмоции, события, место, герои — всё, что важно. Расскажите историю, которую хотите превратить в музыку.",
    "language_english": "Английский: ",
    "language_russian": "Русский: ",
    "generate_button": "Сгенерировать",
    "style_label": "AI-сгенерированный музыкальный стиль: ",
    "song_result_label": "Отформатированный текст песни: ",
    "isGenerateLyrics": "придумать текст песни (экспериментальная)",
    "isGenerateStyle": "придумать стиль музыки по тексту песни",
    # Progress logging messages
    "log_starting": "Начинаем генерацию SUNO промптов...",
    "log_lyrics_gen_start": "Генерируем текст песни из идеи...",
    "log_lyrics_gen_complete": "Генерация текста завершена за {time:.2f}с",
    "log_lyrics_gen_error": "Ошибка генерации текста: {error}",
    "log_style_gen_start": "Генерируем музыкальный стиль из текста...",
    "log_style_desc_start": "Генерируем описание музыкального стиля...",
    "log_style_gen_complete": "Генерация стиля завершена за {time:.2f}с",
    "log_style_gen_error": "Ошибка генерации стиля: {error}",
    "log_text_format_start": "Форматируем текст песни...",
    "log_text_format_complete": "Форматирование текста завершено за {time:.2f}с",
    "log_text_format_error": "Ошибка форматирования текста: {error}",
    "log_summary_success": "✅ Генерация успешно завершена!",
    "log_summary_error": "❌ Генерация завершена с ошибками",
    "log_summary_total": "Общее время выполнения: {time:.2f}с",
    "log_summary_stages": "Время этапов: {stages}",
    "generation_warning_msg": """
**⚠️ Примечание:** используется не самая мощная модель генерации (*GLM*), выбранная из-за доступности на *Hugging Face*.  
Иногда результат может содержать ошибки — чаще всего помогает просто повторная генерация."""    
}

lang = "EN"

if lang=="EN":
    TS=english_strings
else:
    TS=russian_strings

def process_lang(selected_lang, isGenerateLyrics, isGenerateStyle):
    global TS
    lang=selected_lang
    if selected_lang == "RU":
        TS=russian_strings
        message="Вы выбрали русский"
        isVisible=True
    elif selected_lang == "EN":
        TS=english_strings
        message="You selected English"
    else:
        message=""
        isVisible=False
    ret = [gr.update(label=TS["performer_label"],placeholder=TS["performer_placeholder"]),
            gr.update(label=TS["song_text_label" if not isGenerateLyrics else "song_gen_label"],
                    placeholder=TS["song_text_placeholder" if not isGenerateLyrics else "song_gen_placeholder"]),
            gr.update(label=TS["style_label"]),
            gr.update(label=TS["song_result_label"]),
            gr.update(value=TS["generate_button"]),
            gr.update(label=TS["isGenerateLyrics"]),
            gr.update(label=TS["isGenerateStyle"]),
            gr.update(value=TS["generation_warning_msg"])
        ]
    return message, *ret

with gr.Blocks(title="Suno PromptGen", fill_height=True,
               theme=gr.themes.Default(primary_hue=gr.themes.colors.sky, secondary_hue=gr.themes.colors.indigo),
               analytics_enabled=False, css="footer{display:none !important}") as demo:
    logo_info = demo.serve_static_file("splash.webp")
    logo_url = logo_info["url"]

    with gr.Row(elem_id="header", equal_height=True):
        with gr.Column(scale=90, min_width=200):
            gr.HTML(f"""
            <div style='display: flex; align-items: center; gap: 16px;'>
                <img src='{logo_url}' alt='AiCave logo'
                    width='100' height='100'
                    style='background: transparent; display: block;' />
                <div style='line-height: 1.2;'>
                    <h3 style='margin: 0; font-weight: 600;'>SUNO style generator and lyrics format</h3>
                    <p style='margin: 4px 0 0 0; font-size: 0.9em;'>by <a href='https://boosty.to/aicave/donate' style='color: #60a5fa;'>AiCave</a></p>
                </div>
            </div>
            """)
        with gr.Column(scale=5):
            radio_lang = gr.Radio(choices = ["RU",  "EN"],  show_label = False, container = False, type = "value", value=lang)
        with gr.Column(scale=5):
            generate_button = gr.Button(TS["generate_button"], variant="primary", size="lg")
    with gr.Row():
        with gr.Column():
            name_input = gr.Textbox(label=TS["performer_label"],lines=5, max_lines=10,
                                    placeholder=TS["performer_placeholder"], autofocus=True)
            isGenerateStyle = gr.Checkbox(value=False, label=TS["isGenerateStyle"], show_label=True)
            gr.Markdown("")
            text_input = gr.Textbox(label=TS["song_text_label"], lines=23, max_lines=23,
                                    placeholder=TS["song_text_placeholder"])

            isGenerateLyrics = gr.Checkbox(value=False, label=TS["isGenerateLyrics"], show_label=True)

        with gr.Column(show_progress=False):
            style_output = gr.Textbox(label=TS["style_label"],lines=7, max_lines=10,
                                      show_copy_button=True,)
            gr.Markdown("")
            song_output = gr.Textbox(label=TS["song_result_label"], lines=25, max_lines=25,show_copy_button=True)

    with gr.Row(variant="default"):
            log_text = gr.Textbox(value=status, container=False, lines=1, max_lines=3, label="log")
    warning_msg = gr.Markdown(value=TS["generation_warning_msg"])

    generate_button.click(
        fn=get_style_lyrics,
        inputs=[name_input, text_input, isGenerateLyrics, isGenerateStyle],
        outputs=[style_output, song_output, log_text],
        api_name="GenSuno",
        show_api=True
    )

    radio_lang.change(process_lang, inputs=[radio_lang, isGenerateLyrics, isGenerateStyle],
                outputs=[log_text,name_input,text_input,style_output,song_output,generate_button,
                        isGenerateLyrics, isGenerateStyle, warning_msg],
                api_name=False, show_api=False)

    isGenerateLyrics.change(process_lang, inputs=[radio_lang, isGenerateLyrics, isGenerateStyle],
                outputs=[log_text,name_input,text_input,style_output,song_output,generate_button,
                        isGenerateLyrics, isGenerateStyle, warning_msg],
                api_name=False, show_api=False)
    isGenerateStyle.change(process_lang, inputs=[radio_lang, isGenerateLyrics, isGenerateStyle],
                outputs=[log_text,name_input,text_input,style_output,song_output,generate_button,
                        isGenerateLyrics, isGenerateStyle, warning_msg],
                api_name=False, show_api=False)
    
demo.css = """
#header {
    align-items: center;
    padding: 0px 0px;
    border-bottom: None;
}
"""

if __name__ == "__main__":
    demo.launch(quiet=False, share=False, server_name="0.0.0.0", server_port=7860, ssr_mode=False)