Update app.py

app.py

@@ -3,85 +3,174 @@ import torch
 from transformers import pipeline, set_seed
 import re
 import random
+import os
 
-# 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 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
 
-def generate_short_caption(description):
-    """Generate a short, catchy caption (max 5 words)"""
-    if not description.strip():
-        return ""
+# 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
     
-    # Set random seed for variety
-    set_seed(random.randint(1, 10000))
+    words = caption.split()
     
-    # Short prompt templates for 5-word captions
-    prompts = [
-        f"Photo: {description}\nShort caption: Squad",
-        f"Photo: {description}\nCool caption: Good",
-        f"Photo: {description}\nCatchy caption: Living",
-        f"Photo: {description}\nFun caption: Making",
-        f"Photo: {description}\nTrendy caption: Caught"
+    # 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)
     ]
     
-    prompt = random.choice(prompts)
+    for pattern in invalid_patterns:
+        if re.search(pattern, caption, re.IGNORECASE):
+            return False
     
-    try:
-        generated = generator(
-            prompt,
-            max_length=len(prompt.split()) + 8,  # Very short
-            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 = full_text.replace(prompt, "").strip()
-        
-        # Clean and limit to 5 words
-        if caption:
-            words = caption.split()[:5]
-            if len(words) >= 2:
-                return " ".join(words)
+    # 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'
+    ]
     
-    except Exception as e:
-        print(f"Generation error: {e}")
+    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!"
+        return "Please enter a description of your photo! 📸"
     
-    # Cool short captions for different vibes
+    # Enhanced cool captions database with complete, meaningful phrases
     cool_captions = [
         "Squad goals right here",
         "Good vibes only today",
         "Living our best life",
         "Making memories with favorites",
         "Caught in the moment",
-        "Friends like these though",
+        "Friends like these matter",
         "Energy levels unmatched today",
         "Vibes on point always",
         "Good times roll forever",
@@ -92,7 +181,7 @@ def generate_instagram_captions(description):
         "Good company great memories",
         "Squad deeper than ocean",
         "Vibes immaculate as always",
-        "Energy infectious spread it",
+        "Energy infectious spread everywhere",
         "Good times never ending",
         "Making waves together always",
         "Caught smiling too much",
@@ -105,23 +194,54 @@ def generate_instagram_captions(description):
         "Squad tight bond stronger",
         "Good energy good people",
         "Making memories worth keeping",
-        "Caught being absolutely legendary"
+        "Caught being absolutely legendary",
+        "Weekend warriors in action",
+        "Sunshine mixed with chaos",
+        "Creating our own magic",
+        "Adventure awaits us always",
+        "Blessed with the best",
+        "Chaos and we prosper",
+        "Dreams do come true",
+        "Forever young at heart",
+        "Golden hour golden friends",
+        "Happiness looks like this",
+        "Peace love and yoga",
+        "Zen mode activated today",
+        "Morning ritual pure bliss",
+        "Mindful moments matter most",
+        "Grateful for this peace",
+        "Serenity found in stillness",
+        "Balance restored through practice",
+        "Inner peace outer strength",
+        "Breathe deeply live fully",
+        "Namaste all day long"
     ]
     
-    # Generate some AI captions
+    # Generate AI captions if model is available
     ai_captions = []
-    for i in range(3):
-        caption = generate_short_caption(description)
-        if caption:
-            ai_captions.append(caption)
+    if generator is not None:
+        attempts = 0
+        while len(ai_captions) < 2 and attempts < 10:  # Try to get at least 2 good 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 AI and predefined captions
     all_captions = ai_captions + cool_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)
+    
     # Select 5 random captions
-    selected_captions = random.sample(all_captions, min(5, len(all_captions)))
+    selected_captions = random.sample(unique_captions, min(5, len(unique_captions)))
     
-    # Ensure each caption is max 5 words
+    # Ensure each caption is max 5 words and properly formatted
     final_captions = []
     for caption in selected_captions:
         words = caption.split()[:5]
@@ -136,16 +256,18 @@ def generate_instagram_captions(description):
         if final_caption not in final_captions:
             final_captions.append(final_caption)
     
-    # Format as numbered list
-    result = ""
-    for i, caption in enumerate(final_captions[:5], 1):
-        result += f"{i}. {caption}\n"
+    # Format as numbered list with emojis
+    result = "✨ Your Instagram Captions:\n\n"
+    emojis = ["🎯", "⚡", "🌟", "🔥", "💫"]
+    
+    for i, caption in enumerate(final_captions[:5]):
+        result += f"{emojis[i]} {caption}\n"
     
     return result.strip()
 
-# Instagram-inspired theme with custom CSS
+# Enhanced Instagram-inspired theme with custom CSS
 instagram_css = """
-/* Instagram color scheme */
+/* Instagram color scheme and modern design */
 :root {
     --ig-primary: #E4405F;
     --ig-secondary: #833AB4;
@@ -155,12 +277,14 @@ instagram_css = """
     --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 */
@@ -174,6 +298,7 @@ instagram_css = """
     font-weight: 900;
     margin-bottom: 0.5em;
     text-shadow: 0 2px 4px rgba(0,0,0,0.1);
+    padding: 1rem 0;
 }
 
 .subtitle {
@@ -182,6 +307,7 @@ instagram_css = """
     font-size: 1.2em;
     margin-bottom: 2em;
     text-shadow: 0 1px 2px rgba(0,0,0,0.3);
+    font-weight: 300;
 }
 
 /* Card styling */
@@ -193,6 +319,12 @@ instagram_css = """
     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 */
@@ -200,6 +332,8 @@ instagram_css = """
     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 {
@@ -218,6 +352,7 @@ instagram_css = """
     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 {
@@ -226,6 +361,10 @@ instagram_css = """
     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;
@@ -233,7 +372,8 @@ instagram_css = """
     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.6 !important;
+    line-height: 1.8 !important;
+    padding: 1rem !important;
 }
 
 /* Examples styling */
@@ -251,6 +391,7 @@ instagram_css = """
     color: white !important;
     margin: 0.25rem !important;
     transition: all 0.3s ease !important;
+    font-weight: 600 !important;
 }
 
 .gr-examples .gr-button:hover {
@@ -274,6 +415,16 @@ instagram_css = """
     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 {
@@ -284,6 +435,21 @@ instagram_css = """
         margin: 0.5rem !important;
         padding: 1.5rem !important;
     }
+    
+    .generate-btn {
+        font-size: 1.1em !important;
+        padding: 12px 25px !important;
+    }
+}
+
+/* Loading animation */
+@keyframes pulse {
+    0%, 100% { opacity: 1; }
+    50% { opacity: 0.5; }
+}
+
+.loading {
+    animation: pulse 2s infinite;
 }
 """
 
@@ -294,17 +460,18 @@ with gr.Blocks(
         secondary_hue="purple",
         neutral_hue="gray"
     ),
-    css=instagram_css
+    css=instagram_css,
+    title="InstaCap - AI Caption Generator"
 ) as demo:
     
     gr.HTML('<h1 class="main-header">📸 InstaCap – AI Caption Generator</h1>')
-    gr.HTML('<p class="subtitle">Create short, catchy Instagram captions in seconds! ✨</p>')
+    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, sunset vibes, coffee date...",
+                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"
@@ -318,25 +485,27 @@ with gr.Blocks(
         
         with gr.Column(scale=1, elem_classes="card"):
             caption_output = gr.Textbox(
-                label="🎯 Your 5 Short Captions (Max 5 Words Each)",
+                label="🎯 Your Instagram Captions",
                 lines=8,
-                max_lines=10,
+                max_lines=12,
                 interactive=False,
                 show_copy_button=True,
                 elem_classes="output-box"
             )
     
-    # Add examples with Instagram-style descriptions
+    # Add examples with diverse Instagram-style descriptions
     gr.Examples(
         examples=[
-            ["group photo with best friends"],
-            ["sunset beach vibes together"],
-            ["coffee date good times"],
-            ["squad night out fun"],
-            ["adventure hiking with crew"],
-            ["cozy indoor friend hangout"],
-            ["party celebration with everyone"],
-            ["weekend brunch with favorites"]
+            ["group photo with best friends at sunset"],
+            ["coffee date with my bestie"],
+            ["hiking adventure with the squad"],
+            ["beach day with good vibes"],
+            ["cozy movie night with favorites"],
+            ["birthday celebration with everyone"],
+            ["weekend brunch with the crew"],
+            ["road trip memories being made"],
+            ["concert night with amazing energy"],
+            ["graduation day with family pride"]
         ],
         inputs=description_input,
         outputs=caption_output,
@@ -362,15 +531,28 @@ with gr.Blocks(
     gr.HTML(
         """
         <div class="tips-section">
-            <p>💡 <strong>Pro Tips:</strong></p>
-            <p>🎯 Each caption is max 5 words for maximum impact</p>
-            <p>⚡ Perfect for group photos, adventures, and good vibes</p>
-            <p>📱 Copy and paste directly to Instagram</p>
-            <p>🔄 Generate again for new options</p>
+            <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()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )