Update app.py

app.py

@@ -1,426 +1,298 @@
 import gradio as gr
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
-import random
+from transformers import pipeline, set_seed
 import re
-import os
-
-# Initialize model and tokenizer
-print("Loading GPT-2 model...")
-model_name = "gpt2"
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-model = GPT2LMHeadModel.from_pretrained(model_name)
+import random
 
-# Add padding token
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+# Set up the text generation pipeline
+try:
+    # Try to use GPU if available, otherwise fallback to CPU
+    device = 0 if torch.cuda.is_available() else -1
+    generator = pipeline(
+        "text-generation",
+        model="gpt2-medium",
+        device=device,
+        pad_token_id=50256
+    )
+except Exception as e:
+    print(f"Error loading model: {e}")
+    # Fallback to CPU with smaller model
+    generator = pipeline(
+        "text-generation",
+        model="gpt2",
+        device=-1,
+        pad_token_id=50256
+    )
 
-# Set model to evaluation mode for better performance
-model.eval()
+def create_short_caption(text, max_words=5):
+    """Ensure caption is max 5 words"""
+    words = text.split()
+    if len(words) <= max_words:
+        return text
+    return ' '.join(words[:max_words])
 
-def clean_caption(text):
-    """Clean and format generated caption"""
-    # Remove extra whitespace
-    text = re.sub(r'\s+', ' ', text.strip())
-    
-    # Split into sentences and take the first complete one
-    sentences = re.split(r'[.!?]+', text)
-    if sentences:
-        caption = sentences[0].strip()
-    else:
-        caption = text
-    
-    # Remove incomplete thoughts and weird artifacts
-    caption = re.sub(r'\b(the|and|or|but|so|if|when|where|who|what|why|how)\s*$', '', caption, flags=re.IGNORECASE)
+def generate_instagram_caption(description):
+    """Generate 5 short, catchy Instagram captions (max 5 words each)"""
+    if not description.strip():
+        return "Please enter a description of your photo!"
     
-    # Ensure it's not too short or too long
-    if len(caption) < 10:
-        return None
-    if len(caption) > 200:
-        caption = caption[:200]
+    # Set a random seed for variety
+    set_seed(random.randint(1, 10000))
     
-    # Capitalize first letter
-    if caption:
-        caption = caption[0].upper() + caption[1:]
+    # Cool and catchy caption templates (will be trimmed to 5 words)
+    cool_templates = [
+        "Squad goals right here",
+        "Living our best life",
+        "Good vibes only always",
+        "Making memories with favorites",
+        "This crew hits different",
+        "Unstoppable energy with besties",
+        "Main character energy activated",
+        "Plot twist we're legendary",
+        "Caught being absolutely iconic",
+        "Serving looks and laughs",
+        "That golden hour magic",
+        "Friendship level absolutely maximum",
+        "Current mood pure happiness",
+        "Creating chaos looking cute",
+        "Weekend warriors in action",
+        "Living proof dreams work",
+        "Grateful for these humans",
+        "When life gets good",
+        "Positive vibes attract everything",
+        "Making ordinary moments special"
+    ]
     
-    # Add punctuation if missing
-    if not caption.endswith(('.', '!', '?')):
-        caption += '.'
+    # Generate captions based on description context
+    captions = []
     
-    return caption
-
-def generate_instagram_caption(prompt, temperature=0.8, max_length=60):
-    """Generate a single Instagram caption"""
-    # Create input prompt
-    input_text = f"Instagram caption for {prompt}: "
+    # Try to generate some with AI
+    for i in range(3):
+        try:
+            prompt = f"Cool Instagram caption for {description}:"
+            generated = generator(
+                prompt,
+                max_length=len(prompt.split()) + 8,
+                num_return_sequences=1,
+                temperature=0.9,
+                do_sample=True,
+                top_p=0.95,
+                repetition_penalty=1.2,
+                pad_token_id=50256
+            )
+            
+            full_text = generated[0]['generated_text']
+            caption_text = full_text.replace(prompt, "").strip()
+            
+            # Clean and shorten
+            caption_text = caption_text.strip('"\'').strip()
+            if caption_text and len(caption_text) > 3:
+                short_caption = create_short_caption(caption_text, 5)
+                if short_caption not in captions:
+                    captions.append(short_caption)
+                    
+        except Exception as e:
+            print(f"Generation error: {e}")
+            continue
     
-    # Tokenize
-    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+    # Add cool predefined captions
+    random.shuffle(cool_templates)
+    for template in cool_templates:
+        if len(captions) >= 5:
+            break
+        short_caption = create_short_caption(template, 5)
+        if short_caption not in captions:
+            captions.append(short_caption)
     
-    # Generate with controlled parameters for speed
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs,
-            max_length=inputs.shape[1] + max_length,
-            num_return_sequences=1,
-            temperature=temperature,
-            do_sample=True,
-            top_p=0.9,
-            top_k=40,
-            pad_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.1,
-            no_repeat_ngram_size=2
-        )
+    # Context-specific captions based on description
+    context_captions = []
+    desc_lower = description.lower()
     
-    # Decode and clean
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    caption = generated_text.replace(input_text, "").strip()
+    if any(word in desc_lower for word in ['friend', 'group', 'squad', 'crew', 'team', 'gang']):
+        context_captions = [
+            "Squad goals achieved today",
+            "Friends who slay together",
+            "Main character energy activated",
+            "Chaos creators looking cute",
+            "Weekend warriors in action"
+        ]
+    elif any(word in desc_lower for word in ['sunset', 'beach', 'ocean', 'water']):
+        context_captions = [
+            "Chasing sunsets and dreams",
+            "Salt air good vibes",
+            "Beach therapy session complete",
+            "Ocean child energy activated",
+            "Waves and good times"
+        ]
+    elif any(word in desc_lower for word in ['coffee', 'morning', 'work', 'laptop']):
+        context_captions = [
+            "Morning hustle mode on",
+            "Coffee powered productivity session",
+            "Grinding with good vibes",
+            "Main character work energy",
+            "Caffeine and determination daily"
+        ]
+    elif any(word in desc_lower for word in ['food', 'dinner', 'eating', 'restaurant']):
+        context_captions = [
+            "Good food better company",
+            "Eating our way happy",
+            "Dinner dates hit different",
+            "Food coma incoming soon",
+            "Taste buds officially satisfied"
+        ]
+    elif any(word in desc_lower for word in ['adventure', 'hiking', 'mountain', 'nature']):
+        context_captions = [
+            "Adventure mode fully activated",
+            "Mountain therapy session complete",
+            "Nature calls we answer",
+            "Hiking boots happy feet",
+            "Views and good vibes"
+        ]
     
-    return clean_caption(caption)
-
-def add_instagram_flair(caption):
-    """Add Instagram-style elements to caption"""
-    if not caption:
-        return caption
+    # Add context-specific captions
+    for context_cap in context_captions:
+        if len(captions) >= 5:
+            break
+        short_caption = create_short_caption(context_cap, 5)
+        if short_caption not in captions:
+            captions.append(short_caption)
     
-    # Instagram-style emojis
-    emojis = ["✨", "🌟", "💫", "🔥", "💯", "🙌", "❤️", "🌈", "☀️", "🌙", "💕", "🎉", "🌸", "🦋", "💎", "🌊", "🏖️", "🌅", "📸", "💪"]
+    # Ensure we have exactly 5 captions
+    universal_captions = [
+        "Living our best life",
+        "Good vibes only always",
+        "Making memories count today",
+        "Plot twist we're iconic",
+        "Current mood pure happiness",
+        "Grateful for these moments",
+        "Creating magic happening now",
+        "Weekend warrior energy activated",
+        "Positive vibes attract everything",
+        "This crew hits different"
+    ]
     
-    # Add emoji with 40% chance
-    if random.random() < 0.4:
-        emoji = random.choice(emojis)
-        caption = f"{caption} {emoji}"
+    for universal_cap in universal_captions:
+        if len(captions) >= 5:
+            break
+        short_caption = create_short_caption(universal_cap, 5)
+        if short_caption not in captions:
+            captions.append(short_caption)
     
-    return caption
-
-def generate_captions(description):
-    """Generate 5 Instagram captions"""
-    if not description.strip():
-        return "Please enter a description to generate captions! 📝"
+    # Return exactly 5 captions
+    final_captions = captions[:5]
     
-    try:
-        captions = []
-        max_attempts = 15  # Try up to 15 times to get 5 good captions
-        
-        # Different prompt variations for diversity
-        prompt_variations = [
-            description,
-            f"aesthetic {description}",
-            f"vibes {description}",
-            f"mood {description}",
-            f"feeling {description}"
-        ]
-        
-        for i in range(max_attempts):
-            if len(captions) >= 5:
-                break
-            
-            # Use different prompts and temperatures for variety
-            prompt = prompt_variations[i % len(prompt_variations)]
-            temperature = 0.7 + (i % 3) * 0.1  # Vary between 0.7-0.9
-            
-            caption = generate_instagram_caption(prompt, temperature)
-            
-            if caption and caption not in captions:
-                # Add Instagram flair
-                caption = add_instagram_flair(caption)
-                captions.append(caption)
-        
-        # Fill with backup captions if needed
-        while len(captions) < 5:
-            backup_captions = [
-                f"Living my best {description} life ✨",
-                f"Good vibes and {description} energy 🌟",
-                f"Making memories with {description} moments 💫",
-                f"Embracing the {description} aesthetic 🌈",
-                f"Forever grateful for {description} days 💕"
-            ]
-            for backup in backup_captions:
-                if backup not in captions and len(captions) < 5:
-                    captions.append(backup)
-        
-        # Format as numbered list
-        formatted_output = []
-        for i, caption in enumerate(captions[:5], 1):
-            formatted_output.append(f"{i}. {caption}")
-        
-        return "\n\n".join(formatted_output)
+    # Format as numbered list
+    result = ""
+    for i, caption in enumerate(final_captions, 1):
+        result += f"{i}. {caption}\n"
     
-    except Exception as e:
-        return f"Oops! Something went wrong: {str(e)} 😅\nTry a different description!"
-
-# Custom CSS for Instagram-style design
-css = """
-/* Import Google Fonts */
-@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;500;600;700&display=swap');
-
-/* Main gradient background */
-.gradio-container {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
-    font-family: 'Poppins', sans-serif !important;
-}
-
-/* Container styling */
-.contain {
-    background: rgba(255, 255, 255, 0.1) !important;
-    backdrop-filter: blur(10px) !important;
-    border-radius: 20px !important;
-    border: 1px solid rgba(255, 255, 255, 0.2) !important;
-    box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37) !important;
-}
-
-/* Title styling */
-h1 {
-    background: linear-gradient(45deg, #ff6b6b, #4ecdc4, #45b7d1, #f9ca24) !important;
-    -webkit-background-clip: text !important;
-    -webkit-text-fill-color: transparent !important;
-    background-clip: text !important;
-    text-align: center !important;
-    font-weight: 700 !important;
-    font-size: 2.5rem !important;
-    margin-bottom: 1rem !important;
-    text-shadow: 0 4px 8px rgba(0,0,0,0.1) !important;
-}
-
-/* Input styling */
-.gr-textbox {
-    background: rgba(255, 255, 255, 0.9) !important;
-    border-radius: 25px !important;
-    border: 2px solid transparent !important;
-    background-clip: padding-box !important;
-    box-shadow: 0 4px 16px rgba(0,0,0,0.1) !important;
-    font-size: 16px !important;
-    padding: 15px 20px !important;
-    transition: all 0.3s ease !important;
-}
-
-.gr-textbox:focus {
-    border: 2px solid #4ecdc4 !important;
-    box-shadow: 0 4px 20px rgba(78, 205, 196, 0.3) !important;
-    transform: translateY(-2px) !important;
-}
-
-/* Button styling */
-.gr-button {
-    background: linear-gradient(45deg, #ff6b6b, #4ecdc4) !important;
-    border: none !important;
-    border-radius: 25px !important;
-    color: white !important;
-    font-weight: 600 !important;
-    font-size: 18px !important;
-    padding: 15px 30px !important;
-    box-shadow: 0 6px 20px rgba(255, 107, 107, 0.3) !important;
-    transition: all 0.3s ease !important;
-    text-transform: uppercase !important;
-    letter-spacing: 1px !important;
-}
-
-.gr-button:hover {
-    transform: translateY(-3px) !important;
-    box-shadow: 0 8px 25px rgba(255, 107, 107, 0.4) !important;
-}
-
-/* Output styling */
-.gr-textbox[data-testid="textbox"] {
-    background: rgba(255, 255, 255, 0.95) !important;
-    border-radius: 20px !important;
-    border: 1px solid rgba(255, 255, 255, 0.3) !important;
-    font-size: 16px !important;
-    line-height: 1.6 !important;
-    padding: 20px !important;
-}
-
-/* Examples styling */
-.gr-examples {
-    background: rgba(255, 255, 255, 0.1) !important;
-    border-radius: 15px !important;
-    padding: 20px !important;
-    backdrop-filter: blur(5px) !important;
-}
-
-.gr-examples .gr-button {
-    background: rgba(255, 255, 255, 0.2) !important;
-    border: 1px solid rgba(255, 255, 255, 0.3) !important;
-    border-radius: 20px !important;
-    color: white !important;
-    transition: all 0.3s ease !important;
-}
-
-.gr-examples .gr-button:hover {
-    background: rgba(255, 255, 255, 0.3) !important;
-    transform: translateY(-1px) !important;
-}
-
-/* Labels */
-.gr-label {
-    color: white !important;
-    font-weight: 500 !important;
-    font-size: 16px !important;
-    margin-bottom: 8px !important;
-}
-
-/* Markdown styling */
-.gr-markdown {
-    color: white !important;
-    font-size: 16px !important;
-    line-height: 1.6 !important;
-}
-
-.gr-markdown h3 {
-    color: #4ecdc4 !important;
-    font-weight: 600 !important;
-    margin-top: 20px !important;
-}
-
-/* Copy button */
-.gr-copy-button {
-    background: linear-gradient(45deg, #45b7d1, #f9ca24) !important;
-    border: none !important;
-    border-radius: 10px !important;
-    color: white !important;
-    font-weight: 500 !important;
-}
-
-/* Scrollbar styling */
-::-webkit-scrollbar {
-    width: 8px;
-}
-
-::-webkit-scrollbar-track {
-    background: rgba(255, 255, 255, 0.1);
-    border-radius: 10px;
-}
-
-::-webkit-scrollbar-thumb {
-    background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
-    border-radius: 10px;
-}
-
-::-webkit-scrollbar-thumb:hover {
-    background: linear-gradient(45deg, #ff5252, #26a69a);
-}
-
-/* Animation */
-@keyframes float {
-    0% { transform: translateY(0px); }
-    50% { transform: translateY(-10px); }
-    100% { transform: translateY(0px); }
-}
-
-.gr-button {
-    animation: float 3s ease-in-out infinite;
-}
-"""
+    return result.strip()
 
 # Create the Gradio interface
-with gr.Blocks(css=css, theme=gr.themes.Soft(), title="📸 Instacap – AI Instagram Caption Generator") as app:
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css="""
+    .main-header {
+        text-align: center;
+        background: linear-gradient(45deg, #ff6b6b, #4ecdc4, #45b7d1);
+        -webkit-background-clip: text;
+        -webkit-text-fill-color: transparent;
+        font-size: 2.5em;
+        font-weight: bold;
+        margin-bottom: 0.5em;
+    }
+    .subtitle {
+        text-align: center;
+        color: #666;
+        font-size: 1.1em;
+        margin-bottom: 2em;
+    }
+    .generate-btn {
+        background: linear-gradient(45deg, #ff6b6b, #4ecdc4) !important;
+        border: none !important;
+        color: white !important;
+        font-weight: bold !important;
+        font-size: 1.1em !important;
+        padding: 12px 24px !important;
+        border-radius: 25px !important;
+        transition: all 0.3s ease !important;
+    }
+    .generate-btn:hover {
+        transform: translateY(-2px) !important;
+        box-shadow: 0 5px 15px rgba(0,0,0,0.2) !important;
+    }
+    """
+) as demo:
     
-    gr.Markdown(
-        """
-        # 📸 Instacap – AI Instagram Caption Generator
-        
-        Transform your photo descriptions into viral Instagram captions! ✨
-        
-        **Powered by GPT-2 AI** • **Instant Generation** • **100% Unique Captions**
-        """
-    )
+    gr.HTML('<h1 class="main-header">📸 Instacap – AI Instagram Caption Generator</h1>')
+    gr.HTML('<p class="subtitle">Transform your photo descriptions into short, catchy Instagram captions!</p>')
     
     with gr.Row():
         with gr.Column(scale=1):
             description_input = gr.Textbox(
-                label="🎯 Describe Your Photo or Vibe",
-                placeholder="beach sunset, coffee morning, city lights...",
+                label="Describe your photo",
+                placeholder="E.g., group photo with friends, sunset beach picnic, coffee shop vibes",
                 lines=3,
-                max_lines=4
+                max_lines=5
             )
             
             generate_btn = gr.Button(
-                "✨ Generate Captions",
+                "Generate Caption ✨",
                 variant="primary",
-                size="lg"
-            )
-            
-            gr.Markdown(
-                """
-                ### 💡 **Pro Tips:**
-                - Be specific: "golden hour beach" vs "beach"
-                - Include emotions: "cozy rainy day vibes"
-                - Mention key elements: "coffee, books, morning light"
-                """
+                elem_classes="generate-btn"
             )
         
         with gr.Column(scale=1):
-            output_captions = gr.Textbox(
-                label="🔥 Your Instagram Captions",
-                lines=12,
-                max_lines=15,
-                show_copy_button=True,
-                placeholder="Your amazing captions will appear here... ✨"
+            caption_output = gr.Textbox(
+                label="Your 5 Short & Catchy Captions (Max 5 words each)",
+                lines=8,
+                max_lines=10,
+                interactive=False,
+                show_copy_button=True
             )
     
-    # Examples section
-    gr.Markdown("### 🌟 **Try These Trending Vibes:**")
-    
-    examples = [
-        "golden hour beach sunset",
-        "cozy coffee shop morning",
-        "city lights night photography",
-        "aesthetic workout motivation",
-        "minimalist home decor",
-        "travel adventure mountains",
-        "friend group birthday party",
-        "food photography brunch"
-    ]
-    
+    # Add some example inputs
     gr.Examples(
-        examples=examples,
+        examples=[
+            ["group photo with friends laughing"],
+            ["sunset beach picnic, cozy vibe"],
+            ["coffee shop laptop work, productive morning"],
+            ["hiking mountain trail, adventure time"],
+            ["friends dinner party, good food"],
+            ["morning yoga session, peaceful zen"],
+            ["city skyline night view, urban lights"],
+            ["weekend brunch with squad, happy vibes"]
+        ],
         inputs=description_input,
-        outputs=output_captions,
-        fn=generate_captions,
+        outputs=caption_output,
+        fn=generate_instagram_caption,
         cache_examples=False
     )
     
-    # Event handlers
+    # Connect the button to the function
     generate_btn.click(
-        fn=generate_captions,
+        fn=generate_instagram_caption,
         inputs=description_input,
-        outputs=output_captions
+        outputs=caption_output
     )
     
+    # Also trigger on Enter key
     description_input.submit(
-        fn=generate_captions,
+        fn=generate_instagram_caption,
         inputs=description_input,
-        outputs=output_captions
+        outputs=caption_output
     )
     
-    gr.Markdown(
+    gr.HTML(
         """
-        ---
-        
-        ### 🚀 **What Makes Instacap Special:**
-        
-        **🤖 AI-Powered:** Advanced GPT-2 model trained on millions of social media posts
-        
-        **⚡ Lightning Fast:** Optimized for instant caption generation
-        
-        **🎨 Instagram-Ready:** Perfect length, tone, and style for maximum engagement
-        
-        **🔄 Always Unique:** Get different captions every time you generate
-        
-        **📱 Mobile Friendly:** Works perfectly on all devices
-        
-        ---
-        
-        *Built with ❤️ for content creators, influencers, and anyone who wants to level up their Instagram game!*
+        <div style="text-align: center; margin-top: 2em; color: #888; font-size: 0.9em;">
+            <p>💡 Tip: Be specific for better context-matched captions!</p>
+            <p>🎨 Get 5 short, catchy captions (max 5 words each)</p>
+            <p>📋 Perfect for group photos, adventures, and good vibes!</p>
+        </div>
         """
     )
 
 # Launch the app
 if __name__ == "__main__":
-    app.launch(
-        share=True,
-        server_name="0.0.0.0",
-        server_port=7860,
-        favicon_path=None,
-        show_error=True
-    )
+    demo.launch()