Update app.py

app.py

@@ -5,23 +5,51 @@ from PIL import Image
 from io import BytesIO
 import base64
 import os
+import json
 
 # Initialize the Google Generative AI client with the API key from environment variables
-client = genai.Client(api_key=os.environ['GEMINI_API_KEY'])
+try:
+    api_key = os.environ['GEMINI_API_KEY']
+except KeyError:
+    raise ValueError("Please set the GEMINI_API_KEY environment variable.")
+client = genai.Client(api_key=api_key)
 
 def generate_item(tag):
-    # Generate text using Gemini LLM
-    prompt = f"Generate a short, engaging post about {tag} in the style of a TikTok caption."
+    """
+    Generate a single feed item consisting of text from Gemini LLM and an image from Imagen.
+    
+    Args:
+        tag (str): The tag to base the content on.
+    
+    Returns:
+        dict: A dictionary with 'text' (str) and 'image_base64' (str).
+    """
+    # Generate text using Gemini LLM with JSON output
+    prompt = f"""
+    Generate a short, engaging TikTok-style caption about {tag}. 
+    Return the response as a JSON object with a single key 'caption' containing the caption text.
+    Example: {{"caption": "Craving this yummy treat! 😍 #foodie"}}
+    Do not include additional commentary or options.
+    """
     text_response = client.models.generate_content(
         model='gemini-2.5-flash-preview-04-17',
         contents=[prompt]
     )
-    text = text_response.text.strip()
-
-    # Generate an image based on the text or tag
+    # Parse JSON response to extract the caption
+    try:
+        response_json = json.loads(text_response.text.strip())
+        text = response_json['caption']
+    except (json.JSONDecodeError, KeyError):
+        text = f"Wow, {tag} is amazing! 😍 #{tag}"  # Fallback caption
+
+    # Generate an image based on the tag, avoiding text
+    image_prompt = f"""
+    A vivid, high-quality visual scene representing {tag}, designed for a TikTok video. 
+    The image should be colorful and engaging, with no text or letters included.
+    """
     image_response = client.models.generate_images(
         model='imagen-3.0-generate-002',
-        prompt=text,  # Using the generated text as the prompt
+        prompt=image_prompt,
         config=types.GenerateImagesConfig(
             number_of_images=1,
             aspect_ratio="9:16",
@@ -34,8 +62,8 @@ def generate_item(tag):
         generated_image = image_response.generated_images[0]
         image = Image.open(BytesIO(generated_image.image.image_bytes))
     else:
-        # Fallback to a placeholder image if no images are generated
-        image = Image.new('RGB', (300, 533), color='gray')  # Size matches 9:16 aspect ratio
+        # Fallback to a placeholder image
+        image = Image.new('RGB', (300, 533), color='gray')  # 9:16 aspect ratio
 
     # Convert the image to base64
     buffered = BytesIO()
@@ -54,6 +82,8 @@ def start_feed(tag):
     Returns:
         tuple: (current_tag, feed_items, html_content)
     """
+    if not tag.strip():
+        tag = "trending"  # Default tag if empty
     item = generate_item(tag)
     feed_items = [item]
     html_content = generate_html(feed_items)
@@ -77,7 +107,7 @@ def load_more(current_tag, feed_items):
 
 def generate_html(feed_items):
     """
-    Generate an HTML string to display the feed items.
+    Generate an HTML string to display the feed items in a TikTok-like vertical layout.
     
     Args:
         feed_items (list): List of dictionaries containing 'text' and 'image_base64'.
@@ -85,41 +115,103 @@ def generate_html(feed_items):
     Returns:
         str: HTML string representing the feed.
     """
-    html_str = '<div style="max-height: 600px; overflow-y: auto; border: 1px solid #ccc; padding: 10px;">'
+    html_str = """
+    <div style="
+        display: flex;
+        flex-direction: column;
+        align-items: center;
+        max-width: 360px;
+        margin: 0 auto;
+        background-color: #000;
+        height: 640px;
+        overflow-y: auto;
+        scrollbar-width: none;
+        -ms-overflow-style: none;
+        border: 1px solid #333;
+        border-radius: 10px;
+    ">
+    """
+    # Hide scrollbar for a cleaner look
+    html_str += """
+    <style>
+        div::-webkit-scrollbar {
+            display: none;
+        }
+    </style>
+    """
     for item in feed_items:
         html_str += f"""
-        <div style="margin-bottom: 20px; border-bottom: 1px solid #eee; padding-bottom: 20px;">
-            <p style="font-size: 16px; margin-bottom: 10px;">{item['text']}</p>
-            <img src="data:image/png;base64,{item['image_base64']}" style="width: 100%; max-width: 300px; height: auto;">
+        <div style="
+            width: 100%;
+            height: 640px;
+            position: relative;
+            display: flex;
+            flex-direction: column;
+            justify-content: flex-end;
+            overflow: hidden;
+        ">
+            <img src="data:image/png;base64,{item['image_base64']}" style="
+                width: 100%;
+                height: 100%;
+                object-fit: cover;
+                position: absolute;
+                top: 0;
+                left: 0;
+                z-index: 1;
+            ">
+            <div style="
+                position: relative;
+                z-index: 2;
+                background: linear-gradient(to top, rgba(0,0,0,0.7), transparent);
+                padding: 20px;
+                color: white;
+                font-family: Arial, sans-serif;
+                font-size: 18px;
+                font-weight: bold;
+                text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
+            ">
+                {item['text']}
+            </div>
         </div>
         """
-    html_str += '</div>'
+    html_str += "</div>"
     return html_str
 
 # Define the Gradio interface
-with gr.Blocks(title="TikTok-Style Infinite Feed") as demo:
-    # Header
-    gr.Markdown("# TikTok-Style Infinite Feed Generator")
-    gr.Markdown("Enter a tag or select a suggested one to generate a scrollable feed of AI-generated content!")
-
-    # Input components
-    with gr.Row():
-        suggested_tags = gr.Dropdown(
-            choices=["technology", "nature", "art", "food"],
-            label="Suggested Tags",
-            value="nature"  # Default value
-        )
-        tag_input = gr.Textbox(label="Enter a Custom Tag", value="nature")
-
-    # Buttons
-    with gr.Row():
-        start_button = gr.Button("Start Feed")
-        load_more_button = gr.Button("Load More")
+with gr.Blocks(
+    css="""
+        body { background-color: #000; color: #fff; font-family: Arial, sans-serif; }
+        .gradio-container { max-width: 400px; margin: 0 auto; padding: 10px; }
+        input, select, button { border-radius: 5px; background-color: #222; color: #fff; border: 1px solid #444; }
+        button { background-color: #ff2d55; border: none; }
+        button:hover { background-color: #e0264b; }
+        .gr-button { width: 100%; margin-top: 10px; }
+        .gr-form { background-color: #111; padding: 15px; border-radius: 10px; }
+    """,
+    title="TikTok-Style Infinite Feed"
+) as demo:
+    # Input section
+    with gr.Column(elem_classes="gr-form"):
+        gr.Markdown("### Create Your TikTok Feed")
+        with gr.Row():
+            suggested_tags = gr.Dropdown(
+                choices=["food", "travel", "fashion", "tech"],
+                label="Pick a Tag",
+                value="food"
+            )
+            tag_input = gr.Textbox(
+                label="Or Enter a Custom Tag",
+                value="food",
+                placeholder="e.g., sushi, adventure"
+            )
+        with gr.Row():
+            start_button = gr.Button("Start Feed")
+            load_more_button = gr.Button("Load More")
 
     # Output display
-    feed_html = gr.HTML(label="Your Feed")
+    feed_html = gr.HTML()
 
-    # State variables to maintain feed and tag
+    # State variables
     current_tag = gr.State(value="")
     feed_items = gr.State(value=[])