Update app.py

app.py

@@ -3,6 +3,7 @@ from huggingface_hub import InferenceClient
 import os
 import requests
 import random
+import re
 
 # --- CORE FUNCTIONS ---
 
@@ -20,52 +21,45 @@ def generate_meme_text(idea: str, model: str):
         return None, None, error
     
     try:
-        # A robust prompt that works well with smaller models
-        prompt = f"""You are a meme generator. Create funny text for a meme about "{idea}".
-
-Respond ONLY in this exact format:
-TOP: [Top text, max 8 words]
-BOTTOM: [Bottom text, max 8 words]
-
-Example for "code works on first try":
-TOP: My code works
-BOTTOM: And I don't know why
-
-Now, generate for "{idea}":
-TOP:"""
+        # A simple, direct prompt for older models
+        prompt = f"<|prompter|>Create funny meme text for: \"{idea}\". Respond ONLY in the format:\nTOP: [Top text]\nBOTTOM: [Bottom text]<|endoftext|><|assistant|>"
 
         # Using text_generation which is compatible with these models
         response = client.text_generation(
             prompt,
             model=model,
             max_new_tokens=100,
-            temperature=0.8,
+            temperature=0.9,
             return_full_text=False,
             do_sample=True,
+            stop_sequences=["<|", "\n\n"],
         )
         
         # --- Robust Parsing Logic ---
-        full_response = "TOP:" + response
-        lines = full_response.strip().split('\n')
-        
         top_text = ""
         bottom_text = ""
         
-        for line in lines:
-            line = line.strip()
-            if line.startswith("TOP:") and not top_text:
-                top_text = line.replace("TOP:", "").strip().strip('[]"\'.,')
-            elif line.startswith("BOTTOM:") and not bottom_text:
-                bottom_text = line.replace("BOTTOM:", "").strip().strip('[]"\'.,')
+        # Use regex to find TOP: and BOTTOM:
+        top_match = re.search(r"TOP:\s*(.*)", response, re.IGNORECASE)
+        bottom_match = re.search(r"BOTTOM:\s*(.*)", response, re.IGNORECASE)
+
+        if top_match:
+            top_text = top_match.group(1).strip().strip('[]"\'.,')
+        if bottom_match:
+            bottom_text = bottom_match.group(1).strip().strip('[]"\'.,')
         
         # Fallback if AI gives a weird response
         if not top_text or not bottom_text:
-            words = idea.split()
-            mid_point = max(1, len(words) // 2)
-            top_text = ' '.join(words[:mid_point])
-            bottom_text = ' '.join(words[mid_point:])
-        
-        # Ensure text is not empty
+            lines = response.strip().split('\n')
+            if len(lines) >= 2:
+                top_text = lines[0].replace("TOP:", "").strip()
+                bottom_text = lines[1].replace("BOTTOM:", "").strip()
+            else: # Final fallback
+                words = idea.split()
+                mid_point = max(1, len(words) // 2)
+                top_text = ' '.join(words[:mid_point])
+                bottom_text = ' '.join(words[mid_point:])
+
         if not top_text.strip(): top_text = "Top Text"
         if not bottom_text.strip(): bottom_text = "Bottom Text"
         
@@ -76,7 +70,9 @@ TOP:"""
         if "rate limit" in error_msg:
             return None, None, "❌ **Rate Limit Exceeded**\n\nWait 60 seconds and try again. Free tier has limits."
         elif "503" in error_msg:
-            return None, None, f"❌ **Model is Loading**\n\nWait 30 seconds for '{model}' to wake up and try again."
+            return None, None, f"❌ **Model is Loading**\n\nWait 30 seconds for the model to wake up and try again."
+        elif "404" in error_msg:
+             return None, None, f"❌ **Model Not Found**\n\nThe model '{model}' is not available on the free tier right now. This can change frequently."
         else:
             return None, None, f"❌ **AI Error:** {error_msg[:200]}"
 
@@ -94,7 +90,6 @@ def create_meme(idea: str, template: str, model: str):
     template_id = MEME_TEMPLATES.get(template, "181913649") # Default to Drake
     url = "https://api.imgflip.com/caption_image"
     
-    # Using public credentials for ImgFlip API
     payload = {
         'template_id': template_id,
         'username': 'imgflip_huggingface',
@@ -105,13 +100,12 @@ def create_meme(idea: str, template: str, model: str):
     
     try:
         response = requests.post(url, data=payload, timeout=15)
-        response.raise_for_status() # Raise error for bad responses
+        response.raise_for_status()
         data = response.json()
         
         if data.get('success'):
             meme_url = data['data']['url']
             
-            # Download image locally to display in Gradio
             img_response = requests.get(meme_url, timeout=15)
             if img_response.status_code == 200:
                 temp_path = f"/tmp/meme_{random.randint(1000, 9999)}.jpg"
@@ -138,11 +132,11 @@ def create_meme(idea: str, template: str, model: str):
 
 # --- CONFIGURATION & UI ---
 
-# ✅ WORKING MODELS YOU PROVIDED
+# ✅ STABLE & RELIABLE FREE MODELS
 MODELS = {
-    "Mistral 7B (Fast & Reliable)": "mistralai/Mistral-7B-Instruct-v0.3",
-    "Zephyr 7B": "HuggingFaceH4/zephyr-7b-beta",
-    "Phi-3 Mini": "microsoft/Phi-3-mini-4k-instruct"
+    "OpenAssistant Pythia 12B (Stable)": "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+    "Zephyr 7B (Good Alternative)": "HuggingFaceH4/zephyr-7b-beta",
+    "Mistral 7B v0.2": "mistralai/Mistral-7B-Instruct-v0.2",
 }
 
 # Popular meme templates
@@ -159,10 +153,9 @@ MEME_TEMPLATES = {
 
 # Example prompts
 examples = [
-    ["When you finally fix the bug at 3 AM", "Success Kid", "Mistral 7B (Fast & Reliable)"],
-    ["Junior dev vs Senior dev looking at same error", "Drake", "Mistral 7B (Fast & Reliable)"],
-    ["My code in development vs in production", "Distracted Boyfriend", "Zephyr 7B"],
-    ["Trying to explain AI to my parents", "Expanding Brain", "Phi-3 Mini"],
+    ["When you finally fix the bug at 3 AM", "Success Kid", "OpenAssistant Pythia 12B (Stable)"],
+    ["Junior dev vs Senior dev looking at same error", "Drake", "Zephyr 7B (Good Alternative)"],
+    ["My code in development vs in production", "Distracted Boyfriend", "Mistral 7B v0.2"],
 ]
 
 # Gradio UI
@@ -172,7 +165,7 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Meme Generator") as demo:
     <div style='text-align: center; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                 padding: 30px; border-radius: 15px; color: white; margin-bottom: 20px;'>
         <h1>🥸 AI Meme Generator</h1>
-        <h3>Powered by Actually Working HuggingFace Models 🤗</h3>
+        <h3>Powered by Stable Free HuggingFace Models 🤗</h3>
     </div>
     """)
     
@@ -180,32 +173,22 @@ with gr.Blocks(theme=gr.themes.Soft(), title="AI Meme Generator") as demo:
     
     with gr.Row():
         with gr.Column(scale=2):
-            idea_input = gr.Textbox(
-                label="🎨 Your Meme Idea",
-                placeholder="Example: When the client says 'just one small change'...",
-                lines=2,
-            )
-            
+            idea_input = gr.Textbox(label="🎨 Your Meme Idea", placeholder="Example: When the client says 'just one small change'...", lines=2)
             with gr.Row():
-                template_dropdown = gr.Dropdown(
-                    choices=list(MEME_TEMPLATES.keys()), value="Drake", label="🖼️ Meme Template"
-                )
-                model_dropdown = gr.Dropdown(
-                    choices=list(MODELS.keys()), value="Mistral 7B (Fast & Reliable)", label="🤖 AI Model"
-                )
-            
+                template_dropdown = gr.Dropdown(choices=list(MEME_TEMPLATES.keys()), value="Drake", label="🖼️ Meme Template")
+                model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="OpenAssistant Pythia 12B (Stable)", label="🤖 AI Model (Free Tier)")
             generate_button = gr.Button("🚀 Generate Meme", variant="primary", size="lg")
         
         with gr.Column(scale=1):
             gr.Markdown("""
             ### 📖 How to Use
             
-            1. **Enter Your Idea:** Describe the meme.
-            2. **Pick Template:** Choose a meme format.
-            3. **Choose AI:** Mistral is fast and reliable.
+            1. **Enter Idea:** Describe the meme.
+            2. **Pick Template:** Choose a format.
+            3. **Choose AI:** OpenAssistant is very stable.
             4. **Click Generate!**
             
-            **Setup (First Time Only):**
+            **Setup (First Time):**
             - Get [HF Token](https://huggingface.co/settings/tokens)
             - Go to Settings → Secrets
             - Add `HF_TOKEN`