app.py

from transformers import pipeline
import gradio as gr

# Initialize models suitable for free tier
try:
    # Try to load the primary model
    generator = pipeline('text-generation', model='distilgpt2')
    # Keep a backup model option ready
    backup_model_name = 'gpt2'
except Exception as e:
    print(f"Error loading primary model: {e}")
    # Fallback to an even smaller model if needed
    generator = pipeline('text-generation', model='gpt2')
    backup_model_name = None

def generate_lyrics(theme, genre, rhyme_scheme="AABB", max_length=150):
    """Generate song lyrics based on input parameters"""
    
    # A more structured prompt to guide the model towards a lyrical format
    prompt = f"""Task: Write lyrics for a {genre} song.
Theme: {theme}
Rhyme: {rhyme_scheme}
Style: Poetic, emotional, with verses and a chorus.

[Verse 1]"""
    
    try:
        # Increase temperature for more creativity and reduce repetition
        result = generator(prompt, 
                          max_length=max_length, 
                          num_return_sequences=1,
                          temperature=0.9,  # Increased from 0.7
                          top_p=0.92,       # Slightly increased
                          repetition_penalty=1.2)  # Added to reduce repetition
        
        # Extract the generated text
        lyrics = result[0]['generated_text']
        
        # Check if we got repetitive content
        if lyrics.count("[Verse]") > 3 or lyrics.count("Start with:") > 1:
            raise ValueError("Repetitive output detected")
            
    except Exception as e:
        print(f"Error with primary generation: {e}")
        
        # Try with a simpler prompt
        simple_prompt = f"A {genre} song about {theme}:\n\n[Verse 1]\n"
        
        try:
            # Try with different parameters
            result = generator(simple_prompt, 
                              max_length=max_length, 
                              num_return_sequences=1,
                              temperature=1.0,
                              top_p=0.95,
                              repetition_penalty=1.5)
            
            lyrics = result[0]['generated_text']
        except Exception as e2:
            print(f"Error with fallback generation: {e2}")
            return f"Sorry, I couldn't generate lyrics at this time. Please try again with a different theme or genre."
    
    # Better handling of output formatting
    # Find the first occurrence of [Verse] and keep everything after it
    verse_index = lyrics.find("[Verse]")
    
    if verse_index != -1:
        # Keep everything from [Verse] onwards
        formatted_lyrics = lyrics[verse_index:].strip()
    else:
        # If [Verse] not found, try to find the actual content
        lines = lyrics.split("\n")
        # Skip empty lines and prompt-related lines
        content_lines = []
        skip_keywords = ["Write a", "Start with:", "rhyme scheme", "using the", "Task:", "Theme:", "Rhyme:", "Style:", "Lyrics:"]
        
        for line in lines:
            # Skip empty lines and lines containing prompt keywords
            if not line.strip() or any(keyword in line for keyword in skip_keywords):
                continue
            content_lines.append(line)
        
        formatted_lyrics = "\n".join(content_lines).strip()
        
        # If we couldn't extract anything meaningful, use a generic message
        if not formatted_lyrics or len(formatted_lyrics.split()) < 5:
            return "I couldn't generate good lyrics with this theme. Please try a different theme or genre."
    
    # Check for repetitive content
    lines = formatted_lyrics.split("\n")
    unique_lines = set(line.strip() for line in lines if line.strip())
    
    # If we have very few unique lines compared to total lines, it's repetitive
    if len(lines) > 5 and len(unique_lines) < len(lines) / 2:
        return "The model generated repetitive content. Please try again with a different theme or genre."
    
    return formatted_lyrics

# Custom CSS for better appearance
css = """
.container {
    max-width: 900px;
    margin: auto;
    padding-top: 1.5rem;
}
.title {
    text-align: center;
    color: #7c3aed;
    margin-bottom: 1rem;
}
.subtitle {
    text-align: center;
    color: #6b7280;
    margin-bottom: 2rem;
}
.generate-btn {
    background-color: #7c3aed !important;
}
.footer {
    text-align: center;
    margin-top: 2rem;
    color: #6b7280;
    font-size: 0.875rem;
}
"""

# Example themes for the interface
examples = [
    ["lost in the forest", "folk", "AABB", 150],
    ["falling in love", "pop", "ABAB", 150],
    ["overcoming challenges", "rock", "AABB", 180],
    ["city life", "rap", "ABAB", 200],
]

# Create Gradio interface with a custom theme
with gr.Blocks(css=css, theme="soft") as demo:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 🎵 AI Song Lyric Generator", elem_classes="title")
        gr.Markdown("Create unique song lyrics with AI assistance", elem_classes="subtitle")
        
        with gr.Row():
            with gr.Column():
                theme_input = gr.Textbox(
                    label="Theme or Topic",
                    placeholder="Enter a theme (e.g., lost in the forest)",
                    info="What should your song be about?"
                )
                genre_input = gr.Dropdown(
                    ["pop", "rock", "folk", "rap"], 
                    label="Music Genre",
                    value="pop",
                    info="Select the musical style"
                )
                rhyme_input = gr.Dropdown(
                    ["AABB", "ABAB"], 
                    label="Rhyme Scheme",
                    value="AABB",
                    info="AABB: consecutive lines rhyme, ABAB: alternating lines rhyme"
                )
                length_slider = gr.Slider(
                    50, 200, value=150, 
                    label="Maximum Length",
                    info="Longer text = more content but may be less coherent"
                )
                generate_btn = gr.Button("✨ Generate Lyrics", elem_classes="generate-btn")
            
            with gr.Column():
                output = gr.Textbox(
                    label="Generated Lyrics", 
                    lines=12,
                    show_copy_button=True
                )
        
        gr.Examples(
            examples=examples,
            inputs=[theme_input, genre_input, rhyme_input, length_slider],
            outputs=output,
            fn=generate_lyrics,
            cache_examples=True,
        )
        
        gr.Markdown(
            """
            ### How to use
            1. Enter a theme or topic for your song
            2. Select a music genre
            3. Choose a rhyme scheme
            4. Adjust the length as needed
            5. Click "Generate Lyrics" and enjoy your AI-created song!
            
            ---
            
            *Note: This app uses a lightweight AI model (DistilGPT-2) optimized for Hugging Face Spaces.*
            """,
            elem_classes="footer"
        )
    
    generate_btn.click(
        generate_lyrics,
        inputs=[theme_input, genre_input, rhyme_input, length_slider],
        outputs=output
    )

# Launch the app
demo.launch()