Update app.py

app.py

@@ -1,229 +1,210 @@
 import gradio as gr
 import torch
-from transformers import pipeline, set_seed
-import re
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, set_seed
 import random
+import re
 
-# 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
-    )
+# Load the model and tokenizer
+model_name = "gpt2"
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model = GPT2LMHeadModel.from_pretrained(model_name)
+
+# Add padding token
+tokenizer.pad_token = tokenizer.eos_token
 
-def clean_caption(text):
+def clean_caption(caption):
     """Clean and format the generated caption"""
-    # Remove the prompt part
-    lines = text.split('\n')
-    caption_lines = []
-    
-    for line in lines:
-        line = line.strip()
-        if line and not line.startswith('Photo:') and not line.startswith('Caption:'):
-            caption_lines.append(line)
-    
-    if caption_lines:
-        caption = caption_lines[0]
-        # Remove extra quotes or incomplete sentences
-        caption = re.sub(r'^["\']|["\']$', '', caption)
-        # Ensure it doesn't end mid-sentence
-        if caption and caption[-1] not in '.!?':
-            # Find the last complete sentence
-            sentences = re.split(r'[.!?]', caption)
-            if len(sentences) > 1:
-                caption = '.'.join(sentences[:-1]) + '.'
-            else:
-                caption = caption + '.'
-        return caption
-    return ""
+    # Remove any leading/trailing whitespace
+    caption = caption.strip()
+    
+    # Remove incomplete sentences at the end
+    sentences = caption.split('.')
+    if len(sentences) > 1:
+        caption = sentences[0] + '.'
+    
+    # Ensure it ends with proper punctuation
+    if not caption.endswith(('.', '!', '?')):
+        caption += '.'
+    
+    # Remove extra spaces
+    caption = re.sub(r'\s+', ' ', caption)
+    
+    # Capitalize first letter
+    if caption:
+        caption = caption[0].upper() + caption[1:]
+    
+    return caption
 
-def generate_instagram_caption(description):
-    """Generate an Instagram-style caption based on the photo description"""
-    if not description.strip():
-        return "Please enter a description of your photo!"
-    
-    # Set a random seed for variety
-    set_seed(random.randint(1, 1000))
-    
-    # Create Instagram-style prompts
-    prompts = [
-        f"Photo: {description}\nInstagram Caption: Living my best life",
-        f"Photo: {description}\nInstagram Caption: When life gives you",
-        f"Photo: {description}\nInstagram Caption: Just another day in paradise",
-        f"Photo: {description}\nInstagram Caption: Feeling grateful for",
-        f"Photo: {description}\nInstagram Caption: Making memories with",
-        f"Photo: {description}\nInstagram Caption: Can't get enough of",
-        f"Photo: {description}\nInstagram Caption: Sunday vibes",
-        f"Photo: {description}\nInstagram Caption: Good times and tan lines",
-        f"Photo: {description}\nInstagram Caption: Chasing sunsets and dreams",
-        f"Photo: {description}\nInstagram Caption: Life is better when you're"
+def generate_instagram_captions(description, num_captions=5):
+    """Generate Instagram-style captions based on description"""
+    
+    # Different prompt templates for variety
+    prompt_templates = [
+        f"Instagram caption about {description}:",
+        f"Perfect caption for {description}:",
+        f"Instagram post about {description}:",
+        f"Caption for a photo of {description}:",
+        f"Social media caption about {description}:"
     ]
     
-    # Randomly select a prompt style
-    prompt = random.choice(prompts)
+    captions = []
     
-    try:
+    for i in range(num_captions):
+        # Use different prompt templates for variety
+        prompt = prompt_templates[i % len(prompt_templates)]
+        
+        # Set different seeds for variety
+        set_seed(42 + i * 10)
+        
+        # Tokenize the prompt
+        inputs = tokenizer.encode(prompt, return_tensors="pt")
+        
         # Generate text
-        generated = generator(
-            prompt,
-            max_length=len(prompt.split()) + 25,  # Keep it short
-            num_return_sequences=1,
-            temperature=0.8,
-            do_sample=True,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            pad_token_id=50256
-        )
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=inputs.shape[1] + 50,
+                num_return_sequences=1,
+                temperature=0.8 + (i * 0.1),  # Vary temperature for diversity
+                do_sample=True,
+                top_p=0.9,
+                top_k=50,
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.2
+            )
+        
+        # Decode the generated text
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract and clean the caption
-        full_text = generated[0]['generated_text']
-        caption = clean_caption(full_text.replace(prompt, "").strip())
+        # Extract caption (remove the prompt)
+        caption = generated_text.replace(prompt, "").strip()
         
-        if not caption:
-            # Fallback captions if generation fails
-            fallback_captions = [
-                f"✨ {description} moments ✨",
-                f"Living for these {description} vibes 🌟",
-                f"Just me and my {description} mood 📸",
-                f"When {description} hits different 💫",
-                f"Grateful for {description} days like these 🙏",
-                f"Making memories with {description} energy ⭐",
-                f"Can't get enough of this {description} life 🌈",
-                f"Sunday vibes: {description} edition 🌅",
-                f"Chasing dreams and {description} scenes 🎨",
-                f"Life is better with {description} moments 💕"
-            ]
-            caption = random.choice(fallback_captions)
+        # Clean and format the caption
+        caption = clean_caption(caption)
         
-        # Add some popular hashtags occasionally
-        if random.random() < 0.3:  # 30% chance to add hashtags
-            hashtags = ["#vibes", "#mood", "#life", "#grateful", "#memories", "#blessed", "#goodtimes", "#livingmybestlife"]
-            selected_tags = random.sample(hashtags, min(2, len(hashtags)))
-            caption += " " + " ".join(selected_tags)
+        # Add some Instagram-style elements randomly
+        instagram_elements = ["✨", "🌟", "💫", "🔥", "💯", "🙌", "❤️", "🌈", "☀️", "🌙"]
+        if random.random() < 0.3:  # 30% chance to add an emoji
+            caption = caption.rstrip('.!?') + f" {random.choice(instagram_elements)}"
         
-        return caption
+        captions.append(caption)
+    
+    return captions
+
+def generate_captions_interface(description):
+    """Interface function for Gradio"""
+    if not description.strip():
+        return "Please enter a description to generate captions!"
+    
+    try:
+        captions = generate_instagram_captions(description.strip())
         
+        # Format captions as numbered list
+        formatted_captions = []
+        for i, caption in enumerate(captions, 1):
+            formatted_captions.append(f"{i}. {caption}")
+        
+        return "\n\n".join(formatted_captions)
+    
     except Exception as e:
-        print(f"Generation error: {e}")
-        return f"✨ {description} vibes ✨ #mood #life"
+        return f"Error generating captions: {str(e)}"
 
 # Create the Gradio interface
 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.HTML('<h1 class="main-header">📸 Instacap – AI Instagram Caption Generator</h1>')
-    gr.HTML('<p class="subtitle">Transform your photo descriptions into engaging Instagram captions!</p>')
+    theme=gr.themes.Soft(
+        primary_hue="blue",
+        secondary_hue="purple"
+    ),
+    title="📸 Instacap – AI Instagram Caption Generator"
+) as app:
+    
+    gr.Markdown(
+        """
+        # 📸 Instacap – AI Instagram Caption Generator
+        
+        Transform your photo descriptions into catchy Instagram captions! 
+        Simply describe your photo or the vibe you're going for, and get 5 unique, engaging captions.
+        
+        **Examples:** "beach sunset", "coffee and books", "city nightlife", "cozy home vibes"
+        """
+    )
     
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column(scale=2):
             description_input = gr.Textbox(
-                label="Describe your photo",
-                placeholder="E.g., sunset beach picnic, cozy vibe",
-                lines=3,
-                max_lines=5
+                label="Photo Description or Vibe",
+                placeholder="e.g., beach, sea & sun",
+                lines=2,
+                max_lines=3
             )
             
             generate_btn = gr.Button(
-                "Generate Caption ✨",
+                "Generate Captions",
                 variant="primary",
-                elem_classes="generate-btn"
+                size="lg"
             )
         
-        with gr.Column(scale=1):
-            caption_output = gr.Textbox(
-                label="Your Instagram Caption",
-                lines=4,
-                max_lines=6,
-                interactive=False,
+        with gr.Column(scale=3):
+            output_captions = gr.Textbox(
+                label="Generated Instagram Captions",
+                lines=10,
+                max_lines=15,
                 show_copy_button=True
             )
     
-    # Add some example inputs
+    # Examples section
+    gr.Markdown("### 💡 Try these examples:")
+    
+    examples = [
+        "beach sunset vibes",
+        "coffee morning routine",
+        "city lights at night",
+        "cozy reading corner",
+        "workout motivation",
+        "food photography",
+        "travel adventure",
+        "friendship moments"
+    ]
+    
     gr.Examples(
-        examples=[
-            ["sunset beach picnic, cozy vibe"],
-            ["coffee shop laptop work, productive morning"],
-            ["hiking mountain trail, adventure time"],
-            ["cozy rainy day reading, warm blanket"],
-            ["friends dinner party, laughter and good food"],
-            ["morning yoga session, peaceful zen"],
-            ["city skyline night view, urban lights"],
-            ["homemade cooking, delicious comfort food"]
-        ],
+        examples=examples,
         inputs=description_input,
-        outputs=caption_output,
-        fn=generate_instagram_caption,
+        outputs=output_captions,
+        fn=generate_captions_interface,
         cache_examples=False
     )
     
-    # Connect the button to the function
+    # Set up the event handlers
     generate_btn.click(
-        fn=generate_instagram_caption,
+        fn=generate_captions_interface,
         inputs=description_input,
-        outputs=caption_output
+        outputs=output_captions
     )
     
-    # Also trigger on Enter key
     description_input.submit(
-        fn=generate_instagram_caption,
+        fn=generate_captions_interface,
         inputs=description_input,
-        outputs=caption_output
+        outputs=output_captions
     )
     
-    gr.HTML(
+    gr.Markdown(
         """
-        <div style="text-align: center; margin-top: 2em; color: #888; font-size: 0.9em;">
-            <p>💡 Tip: Be specific with your descriptions for better captions!</p>
-            <p>🎨 Each generation creates a unique caption with different vibes</p>
-        </div>
+        ---
+        
+        **Tips for better captions:**
+        - Be specific with your description
+        - Include emotions or moods
+        - Mention key visual elements
+        - Try different variations of the same concept
+        
+        *Generated captions are AI-created and may need editing for your specific needs.*
         """
     )
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    app.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )