app.py

import gradio as gr
import torch
from transformers import pipeline, set_seed
import re
import random
import os

# Set up the text generation pipeline with better error handling
def initialize_model():
    """Initialize the text generation model with fallback options"""
    try:
        # Check if CUDA is available and model can fit in memory
        device = 0 if torch.cuda.is_available() else -1
        print(f"Using device: {'GPU' if device == 0 else 'CPU'}")
        
        # Try GPT2-medium first, fallback to GPT2 if memory issues
        try:
            generator = pipeline(
                "text-generation",
                model="gpt2-medium",
                device=device,
                pad_token_id=50256,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
            )
            print("Loaded GPT2-medium model successfully")
        except Exception as e:
            print(f"GPT2-medium failed, falling back to GPT2: {e}")
            generator = pipeline(
                "text-generation",
                model="gpt2",
                device=device,
                pad_token_id=50256
            )
            print("Loaded GPT2 model successfully")
        
        return generator
    except Exception as e:
        print(f"Error loading model: {e}")
        # Return None to use fallback captions only
        return None

# Initialize model at startup
generator = initialize_model()

def is_valid_caption(caption):
    """Check if caption is valid and makes sense"""
    if not caption or len(caption.strip()) < 3:
        return False
    
    words = caption.split()
    
    # Check for unwanted patterns
    invalid_patterns = [
        r'http[s]?://',  # URLs
        r'www\.',        # Website references
        r'\.com',        # Domain extensions
        r'\.org',
        r'\.net',
        r'\d{3,}',       # Long numbers
        r'[^\w\s]',      # Special characters (except basic punctuation)
    ]
    
    for pattern in invalid_patterns:
        if re.search(pattern, caption, re.IGNORECASE):
            return False
    
    # Check for incomplete sentences (common AI issues)
    incomplete_endings = [
        'i get', 'i am', 'i will', 'i have', 'i do', 'i can',
        'you get', 'you are', 'you will', 'you have', 'you do',
        'we get', 'we are', 'we will', 'we have', 'we do',
        'they get', 'they are', 'they will', 'they have',
        'there is a', 'there are', 'this is a', 'that is',
        'it is a', 'here is', 'when you', 'if you',
        'the best', 'the most', 'the only', 'the first'
    ]
    
    caption_lower = caption.lower()
    for ending in incomplete_endings:
        if caption_lower.endswith(ending):
            return False
    
    # Check for nonsensical word combinations
    nonsensical_words = [
        'lorem', 'ipsum', 'dolor', 'amet', 'consectetur',
        'adipiscing', 'elit', 'sed', 'eiusmod', 'tempor'
    ]
    
    for word in nonsensical_words:
        if word in caption_lower:
            return False
    
    # Must have at least 2 different words
    if len(set(words)) < 2:
        return False
    
    # Check for proper word structure (no single letters except 'I' and 'a')
    for word in words:
        if len(word) == 1 and word.lower() not in ['i', 'a']:
            return False
    
    return True

def generate_short_caption(description, max_retries=5):
    """Generate a short, catchy caption (max 5 words) with better quality control"""
    if not description.strip() or generator is None:
        return ""
    
    for attempt in range(max_retries):
        try:
            # Set random seed for variety
            set_seed(random.randint(1, 10000))
            
            # More specific prompt templates for complete captions
            prompts = [
                f"Complete Instagram caption about {description[:40]}:",
                f"Short social media caption: {description[:40]}. Caption:",
                f"Trendy Instagram post about {description[:40]}:",
                f"Cool caption for {description[:40]}:",
                f"Fun social media post: {description[:40]}. Text:"
            ]
            
            prompt = random.choice(prompts)
            
            generated = generator(
                prompt,
                max_length=len(prompt.split()) + 12,  # Allow more space for complete thoughts
                num_return_sequences=3,  # Generate multiple options
                temperature=0.7,  # Slightly lower temperature for more coherent text
                do_sample=True,
                top_p=0.85,
                repetition_penalty=1.4,
                pad_token_id=50256,
                truncation=True
            )
            
            # Try each generated option
            for result in generated:
                full_text = result['generated_text']
                caption = full_text.replace(prompt, "").strip()
                
                if caption:
                    # Clean the caption
                    caption = re.sub(r'[^\w\s\']', '', caption)  # Keep apostrophes
                    caption = ' '.join(caption.split())  # Normalize whitespace
                    
                    # Try different word lengths (3-5 words)
                    for max_words in [5, 4, 3]:
                        words = caption.split()[:max_words]
                        if len(words) >= 2:
                            test_caption = " ".join(words)
                            if is_valid_caption(test_caption):
                                return test_caption
        
        except Exception as e:
            print(f"Generation error (attempt {attempt + 1}): {e}")
            continue
    
    return ""

def generate_instagram_captions(description):
    """Generate 5 short, catchy Instagram captions (max 5 words each)"""
    if not description.strip():
        return "Please enter a description of your photo! 📸"
    
    # Keyword-based caption training for specific contexts
    keyword_captions = {
        # Beach/Sea related
        "beach": ["Beach vibes only"],
        "sea": ["Sea you later"],
        "ocean": ["Ocean child wild heart"],
        "sand": ["Sandy toes salty kisses"],
        "sun": ["Chasing golden hour"],
        "sunset": ["Golden hour magic"],
        "swimming": ["Making waves today"],
        "surfing": ["Riding the waves"],
        
        # Food related
        "food": ["Good food good mood"],
        "coffee": ["But first coffee"],
        "pizza": ["Pizza my heart"],
        "brunch": ["Brunch squad assembled"],
        
        # Fitness/Health
        "gym": ["Sweat now shine later"],
        "workout": ["No pain no gain"],
        "yoga": ["Peace love and yoga"],
        "running": ["Run like you stole"],
        
        # Travel/Adventure
        "travel": ["Wanderlust and city dust"],
        "adventure": ["Adventure is out there"],
        "hiking": ["Take only pictures"],
        "camping": ["Into the wild"],
        
        # Party/Celebration
        "party": ["Party like its"],
        "birthday": ["Another year wiser"],
        "celebration": ["Cheers to good times"],
        
        # Friends/Social
        "friends": ["Squad goals right here"],
        "squad": ["Squad goals achieved"],
        "bestie": ["Partner in crime"],
        
        # Fashion/Style
        "outfit": ["Outfit on point"],
        "fashion": ["Fashion is art"],
        "style": ["Style is eternal"],
        
        # Love/Romance
        "love": ["Love is everything"],
        "couple": ["Better together always"],
        "date": ["Date night vibes"],
        
        # Nature/Outdoors
        "nature": ["Nature is medicine"],
        "flowers": ["Bloom where planted"],
        "garden": ["Garden therapy session"],
        
        # Work/Success
        "work": ["Work hard play harder"],
        "success": ["Success tastes sweet"],
        "business": ["Business mode activated"]
    }
    
    # Find keyword-specific captions
    description_lower = description.lower()
    keyword_specific_captions = []
    
    for keyword, captions in keyword_captions.items():
        if keyword in description_lower:
            keyword_specific_captions.extend(captions)
    
    # Generate AI captions if model is available
    ai_captions = []
    if generator is not None:
        attempts = 0
        while len(ai_captions) < 5 and attempts < 15:  # Try to get 5 AI captions
            caption = generate_short_caption(description)
            if caption and caption not in ai_captions and is_valid_caption(caption):
                ai_captions.append(caption)
            attempts += 1
    
    # Combine keyword-specific and AI captions
    # Prioritize keyword-specific captions if found
    if keyword_specific_captions:
        all_captions = keyword_specific_captions + ai_captions
    else:
        all_captions = ai_captions
    
    # Remove duplicates while preserving order
    seen = set()
    unique_captions = []
    for caption in all_captions:
        if caption not in seen:
            seen.add(caption)
            unique_captions.append(caption)
    
    # If we don't have enough captions, create basic fallbacks
    if len(unique_captions) < 5:
        basic_fallbacks = [
            "Picture perfect moment",
            "Living my best life",
            "Good vibes only",
            "Making memories today",
            "Life is beautiful"
        ]
        
        for fallback in basic_fallbacks:
            if fallback not in unique_captions:
                unique_captions.append(fallback)
                if len(unique_captions) >= 5:
                    break
    
    # Select 5 captions
    selected_captions = unique_captions[:5]
    
    # Ensure each caption is max 5 words and properly formatted
    final_captions = []
    for caption in selected_captions:
        words = caption.split()[:5]
        if len(words) >= 2:
            final_captions.append(" ".join(words))
    
    # Format as simple numbered list
    result = ""
    emojis = ["🎯", "⚡", "🌟", "🔥", "💫"]
    
    for i, caption in enumerate(final_captions[:5]):
        result += f"{emojis[i]} {caption}\n"
    
    return result.strip()

# Simplified Instagram-inspired theme with custom CSS
instagram_css = """
/* Instagram color scheme and modern design */
:root {
    --ig-primary: #E4405F;
    --ig-secondary: #833AB4;
    --ig-tertiary: #F77737;
    --ig-quaternary: #FCAF45;
    --ig-blue: #405DE6;
    --ig-bg: #FAFAFA;
    --ig-text: #262626;
    --ig-border: #DBDBDB;
    --ig-light: #FFFFFF;
}

/* Main container styling */
.gradio-container {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    min-height: 100vh;
    font-family: 'Instagram Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

/* Header styling */
.main-header {
    text-align: center;
    background: linear-gradient(45deg, var(--ig-primary), var(--ig-secondary));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    color: transparent;
    font-size: 3em;
    font-weight: 400;
    margin-bottom: 0.5em;
    padding: 1rem 0;
    background-color: rgba(255, 255, 255, 0.9);
    border-radius: 15px;
    margin: 1rem;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1);
}

.subtitle {
    text-align: center;
    color: white;
    font-size: 1.2em;
    margin-bottom: 2em;
    text-shadow: 0 1px 2px rgba(0,0,0,0.3);
    font-weight: 300;
}

/* Card styling */
.card {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 20px !important;
    box-shadow: 0 8px 32px rgba(0,0,0,0.1) !important;
    backdrop-filter: blur(10px) !important;
    border: 1px solid rgba(255,255,255,0.2) !important;
    padding: 2rem !important;
    margin: 1rem !important;
    transition: all 0.3s ease !important;
}

.card:hover {
    transform: translateY(-5px) !important;
    box-shadow: 0 12px 40px rgba(0,0,0,0.15) !important;
}

/* Input styling */
.gr-textbox {
    border-radius: 15px !important;
    border: 2px solid var(--ig-border) !important;
    transition: all 0.3s ease !important;
    font-size: 1.1em !important;
    padding: 1rem !important;
}

.gr-textbox:focus {
    border-color: var(--ig-primary) !important;
    box-shadow: 0 0 0 3px rgba(228, 64, 95, 0.1) !important;
}

/* Button styling */
.generate-btn {
    background: linear-gradient(45deg, var(--ig-primary), var(--ig-secondary)) !important;
    border: none !important;
    color: white !important;
    font-weight: bold !important;
    font-size: 1.2em !important;
    padding: 15px 30px !important;
    border-radius: 25px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(228, 64, 95, 0.3) !important;
    cursor: pointer !important;
}

.generate-btn:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 8px 25px rgba(228, 64, 95, 0.4) !important;
    background: linear-gradient(45deg, var(--ig-secondary), var(--ig-primary)) !important;
}

.generate-btn:active {
    transform: translateY(-1px) !important;
}

/* Output styling */
.output-box {
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 15px !important;
    border: 2px solid var(--ig-border) !important;
    font-family: 'Instagram Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
    font-size: 1.1em !important;
    line-height: 1.8 !important;
    padding: 1rem !important;
}

/* Tips section */
.tips-section {
    background: rgba(255, 255, 255, 0.9) !important;
    border-radius: 15px !important;
    padding: 1.5rem !important;
    margin-top: 2rem !important;
    text-align: center !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1) !important;
}

.tips-section p {
    color: var(--ig-text) !important;
    font-size: 1em !important;
    margin: 0.5rem 0 !important;
}

/* Footer styling */
.footer {
    text-align: center;
    color: white;
    font-size: 0.9em;
    margin-top: 2rem;
    padding: 1rem;
    opacity: 0.8;
}

/* Mobile responsive */
@media (max-width: 768px) {
    .main-header {
        font-size: 2.2em !important;
    }
    
    .card {
        margin: 0.5rem !important;
        padding: 1.5rem !important;
    }
    
    .generate-btn {
        font-size: 1.1em !important;
        padding: 12px 25px !important;
    }
}
"""

# Create the Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="red",
        secondary_hue="purple",
        neutral_hue="gray"
    ),
    css=instagram_css,
    title="InstaCap - AI Caption Generator"
) as demo:
    
    # Header
    gr.HTML('<h1 class="main-header">InstaCap – AI Caption Generator</h1>')
    gr.HTML('<p class="subtitle">Create short, catchy Instagram captions in seconds! Perfect for your social media posts ✨</p>')
    
    with gr.Row():
        with gr.Column(scale=1, elem_classes="card"):
            description_input = gr.Textbox(
                label="📝 Describe your photo",
                placeholder="E.g., group photo with friends at sunset, coffee date with bestie, hiking adventure with squad...",
                lines=3,
                max_lines=5,
                elem_classes="gr-textbox"
            )
            
            generate_btn = gr.Button(
                "✨ Generate Captions",
                variant="primary",
                elem_classes="generate-btn"
            )
        
        with gr.Column(scale=1, elem_classes="card"):
            caption_output = gr.Textbox(
                label="🎯 Your Instagram Captions",
                lines=8,
                max_lines=12,
                interactive=False,
                show_copy_button=True,
                elem_classes="output-box"
            )
    
    # Connect the button to the function
    generate_btn.click(
        fn=generate_instagram_captions,
        inputs=description_input,
        outputs=caption_output
    )
    
    # Also trigger on Enter key
    description_input.submit(
        fn=generate_instagram_captions,
        inputs=description_input,
        outputs=caption_output
    )
    
    gr.HTML(
        """
        <div class="tips-section">
            <p>💡 <strong>Pro Tips for Better Captions:</strong></p>
            <p>🎯 Each caption is maximum 5 words for maximum impact</p>
            <p>⚡ Be specific in your description for better AI-generated captions</p>
            <p>📱 Copy and paste directly to Instagram, TikTok, or any social platform</p>
            <p>🔄 Click generate again for fresh new options</p>
            <p>✨ Mix AI-generated captions with trending phrases</p>
        </div>
        """
    )
    
    gr.HTML(
        """
        <div class="footer">
            <p>Made with ❤️ for social media creators • Perfect for Instagram, TikTok, and more!</p>
        </div>
        """
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )