import gradio as gr
from audiocraft.models import MusicGen
import torch

# Mampiasa 'small' mba ho haingana sy ho zakaky ny RAM
print("Loading model...")
model = MusicGen.get_pretrained('facebook/musicgen-small')

def generate_malagasy_loop(gadona_soratana, bpm, key, scale):
    # Zavamaneno 3 tianao
    instruments = "electric guitar, acoustic piano, saxophone"
    
    # Ny baiko ho an'ny AI
    prompt = f"A 15-second musical loop, {gadona_soratana} style from Madagascar, {instruments}, {bpm} BPM, {key} {scale}, high quality, seamless."
    
    # Famerana ny halavany ho 15 segondra
    model.set_generation_params(duration=15)
    
    # Famoronana ny feo
    wav = model.generate([prompt]) 
    
    # Fitahirizana vonjimaika
    path = "loop_2026.wav"
    from audiocraft.utils.export_utils import write_wav
    write_wav(path, wav[0].cpu(), model.sample_rate)
    
    return path

# CSS manafina ny soratra Gradio eo ambany
custom_css = "footer {visibility: hidden}"

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🎹 Malagasy AI Loop Maker 2026")
    gr.Markdown("Soraty ny gadona (Salegy, Tsapiky, sns) dia mamorona loop 15s misy Guitar, Piano, ary Sax ny AI.")
    
    with gr.Row():
        with gr.Column():
            gadona = gr.Textbox(label="Gadona tiana", value="Salegy Modern")
            bpm = gr.Slider(minimum=60, maximum=180, value=124, label="BPM")
            
            with gr.Row():
                key = gr.Dropdown(["C", "D", "E", "F", "G", "A", "B"], label="Clé", value="G")
                scale = gr.Radio(["Major", "Minor"], label="Loko", value="Major")
            
            btn = gr.Button("MAMORONA LOOP (15s)", variant="primary")
            
        with gr.Column():
            audio_out = gr.Audio(label="Ny Loop-nao")
            gr.Markdown("---")
            gr.Markdown("© 2026 - AI Music Lab Madagascar")

    btn.click(fn=generate_malagasy_loop, inputs=[gadona, bpm, key, scale], outputs=audio_out)

demo.launch()