Update app.py

app.py

@@ -5,6 +5,7 @@ import requests
 import random
 import re
 import tempfile
+import json
 
 # --- CORE FUNCTIONS ---
 
@@ -18,81 +19,101 @@ def get_client():
         return None, "❌ **HuggingFace Token Required**\n\n**Setup:**\n1. Go to Space Settings → Repository Secrets\n2. Add secret: Name=`HF_TOKEN`, Value=(your HF token)\n3. Get token: https://huggingface.co/settings/tokens\n4. Restart Space"
     return InferenceClient(token=token), None
 
-def generate_meme_text(idea: str):
-    """Generate meme text by trying a list of models using the correct chat API."""
+def generate_meme_content(idea: str):
+    """
+    NEW: In a single AI call, choose the best template AND generate the text.
+    Returns: template_name, top_text, bottom_text, error, model_used
+    """
     client, error = get_client()
-    if error: return None, None, error, None
+    if error: return None, None, None, error, None
 
-    MODELS_TO_TRY = ["mistralai/Mistral-7B-Instruct-v0.2", "HuggingFaceH4/zephyr-7b-beta", "google/flan-t5-large"]
+    # We "teach" the AI about our templates
+    template_descriptions = "\n".join([f"- {name}: {desc}" for name, desc in TEMPLATE_GUIDANCE.items()])
+
+    MODELS_TO_TRY = ["mistralai/Mistral-7B-Instruct-v0.2", "HuggingFaceH4/zephyr-7b-beta"]
     failed_models = []
 
     for model_id in MODELS_TO_TRY:
         try:
-            prompt = f"""Task: Create a funny, two-line meme caption for the idea: "{idea}".
+            prompt = f"""You are an AI expert in meme culture. Your task is to analyze an idea, choose the best meme template, and generate a funny caption.
+
+**1. Analyze the user's idea:** "{idea}"
+
+**2. Choose the single best meme template from this list:**
+{template_descriptions}
+
+**3. Generate a funny, two-line caption for the chosen template.**
 
-Format your response exactly like this, with no extra text:
-Top: [The first line of the meme]
-Bottom: [The second line, the punchline]"""
+**4. Format your response as a valid JSON object with three keys: "template", "top_text", "bottom_text".**
+
+Example Response:
+{{
+  "template": "Drake",
+  "top_text": "Manually selecting a meme template",
+  "bottom_text": "Letting the AI choose the template automatically"
+}}
+
+Your JSON response:"""
             
             messages = [{"role": "user", "content": prompt}]
-            
             response_stream = client.chat_completion(
-                messages, model=model_id, max_tokens=60, temperature=0.8, stream=False
+                messages, model=model_id, max_tokens=150, temperature=0.8, stream=False
             )
-            response = response_stream.choices[0].message.content
-            
-            top_text, bottom_text = "",""
-            top_match = re.search(r"Top:\s*(.*)", response, re.IGNORECASE)
-            bottom_match = re.search(r"Bottom:\s*(.*)", response, re.IGNORECASE)
+            response_text = response_stream.choices[0].message.content
 
-            if top_match: top_text = top_match.group(1).strip().strip('[]"')
-            if bottom_match: bottom_text = bottom_match.group(1).strip().strip('[]"')
-
-            if not top_text or not bottom_text:
-                lines = response.strip().split('\n')
-                if len(lines) >= 1: top_text = lines[0].replace("Top:", "").strip()
-                if len(lines) >= 2: bottom_text = lines[1].replace("Bottom:", "").strip()
+            # --- Robust JSON Parsing ---
+            # The model might sometimes add extra text or code blocks around the JSON
+            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+            if not json_match:
+                failed_models.append(f"{model_id.split('/')[-1]} (bad format)")
+                continue
             
-            if not top_text.strip() or not bottom_text.strip(): continue
+            parsed_json = json.loads(json_match.group(0))
+            template = parsed_json.get("template")
+            top_text = parsed_json.get("top_text")
+            bottom_text = parsed_json.get("bottom_text")
+
+            # Validate the response
+            if not all([template, top_text, bottom_text]) or template not in MEME_TEMPLATES:
+                failed_models.append(f"{model_id.split('/')[-1]} (invalid content)")
+                continue
 
-            return top_text.strip(), bottom_text.strip(), None, model_id
+            # SUCCESS!
+            return template, top_text, bottom_text, None, model_id
 
         except Exception as e:
             error_msg = str(e).lower()
-            if "404" in error_msg or "503" in error_msg or "is currently loading" in error_msg:
+            if "404" in error_msg or "503" in error_msg or "is currently loading" in error_msg or "invalid json" in error_msg:
                 failed_models.append(model_id.split('/')[-1])
                 continue
             else:
-                return None, None, f"❌ **AI Error:** {str(e)[:250]}", model_id
+                return None, None, None, f"❌ **AI Error:** {str(e)[:250]}", model_id
 
-    return None, None, f"❌ **All AI Models Are Offline**\n\n**Models Tried:** {', '.join(failed_models)}", None
+    return None, None, None, f"❌ **All AI Models Are Offline or Failing**\n\n**Models Tried:** {', '.join(failed_models)}", None
 
 
-def create_meme(idea: str, template: str):
-    """Main function to generate the complete meme."""
+def create_meme(idea: str):
+    """Main function to generate the complete meme, now without a template input."""
     if not idea or len(idea.strip()) < 3:
         return None, "❌ Please enter a meme idea (at least 3 characters)!"
 
-    # --- FINAL FIX: Check for ImgFlip credentials ---
+    # Check for ImgFlip credentials
     imgflip_user = os.environ.get("IMGFLIP_USERNAME")
     imgflip_pass = os.environ.get("IMGFLIP_PASSWORD")
-
     if not imgflip_user or not imgflip_pass:
-        return None, "❌ **ImgFlip Credentials Required**\n\n**Setup:**\n1. Create a free account at Imgflip.com\n2. Go to Space Settings → Repository Secrets\n3. Add `IMGFLIP_USERNAME` and `IMGFLIP_PASSWORD` secrets\n4. Restart Space."
+        return None, "❌ **ImgFlip Credentials Required in Secrets**"
 
-    top, bottom, error, model_used = generate_meme_text(idea)
-    if error: return None, error
+    # Call the new intelligent function
+    template_name, top, bottom, error, model_used = generate_meme_content(idea)
+    if error:
+        return None, error
 
-    template_id = MEME_TEMPLATES.get(template, "181913649")
+    template_id = MEME_TEMPLATES.get(template_name)
     url = "https://api.imgflip.com/caption_image"
     
-    # --- FINAL FIX: Use your credentials from the secrets ---
     payload = {
-        'template_id': template_id,
-        'username': imgflip_user,
-        'password': imgflip_pass,
-        'text0': top,
-        'text1': bottom
+        'template_id': template_id, 'username': imgflip_user,
+        'password': imgflip_pass, 'text0': top, 'text1': bottom
     }
     
     try:
@@ -110,9 +131,10 @@ def create_meme(idea: str, template: str):
                 temp_path = tmpfile.name
             
             status_message = (f"✅ **Success!**\n\n"
+                            f"🧠 **AI Chose:** {template_name}\n"
                             f"📝 **Top Text:** {top}\n"
                             f"📝 **Bottom Text:** {bottom}\n\n"
-                            f"🤖 **AI Model Used:** {model_used.split('/')[-1] if model_used else 'N/A'}")
+                            f"🤖 **Model Used:** {model_used.split('/')[-1] if model_used else 'N/A'}")
             return temp_path, status_message
         else:
             return None, f"❌ **ImgFlip API Error:** {data.get('error_message', 'Unknown error')}"
@@ -122,53 +144,67 @@ def create_meme(idea: str, template: str):
     except Exception as e:
         return None, f"❌ **An unexpected error occurred:** {str(e)}"
 
-# --- CONFIGURATION & UI (No changes needed) ---
-# [UI code remains the same]
+# --- CONFIGURATION & UI ---
 
+# Maps template names to their ImgFlip IDs
 MEME_TEMPLATES = {
     "Drake": "181913649", "Distracted Boyfriend": "112126428", "Two Buttons": "87743020",
     "Expanding Brain": "93895088", "Success Kid": "61544", "Batman Slapping Robin": "438680",
     "Change My Mind": "129242436", "Woman Yelling at a Cat": "188390779", "Surprised Pikachu": "155067746",
 }
 
+# NEW: Guidance for the AI on what each template means
+TEMPLATE_GUIDANCE = {
+    "Drake": "Represents choosing one thing (good) over another (bad). Good for showing preference.",
+    "Distracted Boyfriend": "Represents being tempted by something new while neglecting something you already have.",
+    "Two Buttons": "Represents a difficult choice, a dilemma, or inner conflict.",
+    "Expanding Brain": "Shows increasing levels of enlightenment or absurdity on a topic.",
+    "Success Kid": "Represents a small victory, unexpected success, or relief.",
+    "Batman Slapping Robin": "Represents a sharp rebuke or correction of a silly idea.",
+    "Change My Mind": "For presenting a controversial opinion that you are confident about.",
+    "Woman Yelling at a Cat": "Represents a misunderstanding, with one side angry and the other confused.",
+    "Surprised Pikachu": "Represents feigned surprise at an obvious outcome.",
+}
+
 examples = [
     ["When you fix a bug you don't understand"], ["My plans for the weekend vs. what I actually do"],
     ["Saying you'll just have one slice of pizza"], ["Me pretending to be productive in a Zoom meeting"],
-    ["Checking my bank account after a long weekend"], ["Trying to assemble IKEA furniture with the instructions"],
 ]
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), title="AI Meme Generator") as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), title="Intelligent AI Meme Generator") as demo:
     
     gr.HTML("""
     <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>Auto-Healing AI that ALWAYS Finds a Working Model 🤗</h3>
+        <h1>🧠 Intelligent AI Meme Generator</h1>
+        <h3>Enter an idea and let the AI choose the best meme template for you!</h3>
     </div>
     """)
 
     with gr.Row():
         with gr.Column(scale=2):
-            idea_input = gr.Textbox(label="🎨 Your Meme Idea", placeholder="Example: When the CI/CD pipeline finally passes...", lines=3)
-            template_dropdown = gr.Dropdown(choices=list(MEME_TEMPLATES.keys()), value="Drake", label="🖼️ Meme Template")
+            idea_input = gr.Textbox(label="🎨 Your Meme Idea", placeholder="Example: When the CI/CD pipeline finally passes...", lines=4)
             generate_button = gr.Button("🚀 Generate Meme", variant="primary", size="lg")
+            
+            with gr.Accordion("How does this work?", open=False):
+                gr.Markdown("""
+                This app is **intelligent and resilient**. When you click generate:
+                1.  The AI analyzes your idea and chooses the best meme template from its knowledge base.
+                2.  It then generates a funny caption for that specific template.
+                3.  It automatically finds a working AI model from a list of free options.
+                """)
         
         with gr.Column(scale=1):
             output_image = gr.Image(label="🖼️ Your Generated Meme", type="filepath", show_download_button=True)
 
-    output_status = gr.Textbox(label="📊 Status & Details", lines=4, show_copy_button=True)
+    output_status = gr.Textbox(label="📊 Status & Details", lines=5, show_copy_button=True)
     
-    with gr.Accordion("How does this work?", open=False):
-        gr.Markdown("""
-        This app is **resilient**. It has a list of reliable free models from Hugging Face and automatically tries them until it finds one that is online and working.
-        If your `HF_TOKEN` is set up correctly in your Space Secrets, this app should always work.
-        """)
-
     gr.Examples(examples=examples, inputs=[idea_input], label="💡 Meme Ideas to Try")
     
+    # UPDATED: The click function no longer needs the template dropdown
     generate_button.click(
         fn=create_meme,
-        inputs=[idea_input, template_dropdown],
+        inputs=[idea_input],
         outputs=[output_image, output_status]
     )