app.py

import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import base64
import os
import json
import random
import urllib.parse
import time

# Initialize the Google Generative AI client with the API key from environment variables
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 clean_response_text(response_text):
    """
    Clean the API response by removing Markdown code block markers.
    
    Args:
        response_text (str): The raw response text from the API.
    
    Returns:
        str: The cleaned response text.
    """
    cleaned_text = response_text.strip()
    if cleaned_text.startswith("```json"):
        cleaned_text = cleaned_text[len("```json"):].strip()
    if cleaned_text.endswith("```"):
        cleaned_text = cleaned_text[:-len("```")].strip()
    return cleaned_text

def generate_ideas(tag):
    """
    Generate a diverse set of ideas related to the tag using the LLM.
    
    Args:
        tag (str): The tag to base the ideas on.
    
    Returns:
        list: A list of ideas as strings.
    """
    prompt = f"""
    Generate a list of 5 diverse and creative ideas related to {tag} that can be used for a TikTok video.
    Each idea should be a short sentence describing a specific scene or concept.
    Return the response as a JSON object with a single key 'ideas' containing a list of 5 ideas.
    Ensure the response is strictly in JSON format.
    Example: {{"ideas": ["A neon-lit gaming setup with RGB lights flashing", "A futuristic robot assembling a gadget"]}}
    """
    try:
        response = client.models.generate_content(
            model='gemini-2.0-flash',
            contents=[prompt],
            config=types.GenerateContentConfig(temperature=1.2)
        )
        print(f"Raw response for ideas: {response.text}")  # Debugging
        if not response.text or response.text.isspace():
            raise ValueError("Empty response from API")
        cleaned_text = clean_response_text(response.text)
        response_json = json.loads(cleaned_text)
        if 'ideas' not in response_json or not isinstance(response_json['ideas'], list):
            raise ValueError("Invalid JSON format: 'ideas' key missing or not a list")
        return response_json['ideas']
    except Exception as e:
        print(f"Error generating ideas: {e}")
        return [
            f"A vibrant {tag} scene at sunset",
            f"A close-up of {tag} with neon lights",
            f"A futuristic take on {tag} with holograms",
            f"A cozy {tag} moment with warm lighting",
            f"An action-packed {tag} scene with dynamic colors"
        ]

def generate_item(tag, ideas, generate_video=False, max_retries=3):
    """
    Generate a single feed item (image and optionally video) using one of the ideas.
    
    Args:
        tag (str): The tag to base the content on.
        ideas (list): List of ideas to choose from.
        generate_video (bool): Whether to generate a video from the image.
        max_retries (int): Maximum number of retries if image generation fails.
    
    Returns:
        dict: A dictionary with 'text' (str), 'image_base64' (str), 'video_base64_list' (list of str), and 'ideas' (list).
    """
    video_base64_list = []
    
    for attempt in range(max_retries):
        selected_idea = random.choice(ideas)
        prompt = f"""
        Based on the idea "{selected_idea}", create content for a TikTok video about {tag}.
        Return a JSON object with two keys:
        - 'caption': A short, viral TikTok-style caption with hashtags.
        - 'image_prompt': A detailed image prompt for generating a high-quality visual scene.
        The image prompt should describe the scene vividly, specify a perspective and style, and ensure no text or letters are included.
        Ensure the response is strictly in JSON format.
        Example: {{"caption": "Neon vibes only! 🌌 #tech", "image_prompt": "A close-up view of a neon-lit gaming setup with RGB lights flashing, in a futuristic style, no text or letters"}}
        """
        try:
            response = client.models.generate_content(
                model='gemini-2.0-flash',
                contents=[prompt],
                config=types.GenerateContentConfig(temperature=1.2)
            )
            print(f"Raw response for item (attempt {attempt + 1}): {response.text}")  # Debugging
            if not response.text or response.text.isspace():
                raise ValueError("Empty response from API")
            cleaned_text = clean_response_text(response.text)
            response_json = json.loads(cleaned_text)
            if 'caption' not in response_json or 'image_prompt' not in response_json:
                raise ValueError("Invalid JSON format: 'caption' or 'image_prompt' key missing")
            text = response_json['caption']
            image_prompt = response_json['image_prompt']
        except Exception as e:
            print(f"Error generating item (attempt {attempt + 1}): {e}")
            text = f"Obsessed with {tag}! 🔥 #{tag}"
            image_prompt = f"A vivid scene of {selected_idea}, in a vibrant pop art style, no text or letters"

        # Attempt to generate the image
        try:
            imagen = client.models.generate_images(
                model='imagen-3.0-generate-002',
                prompt=image_prompt,
                config=types.GenerateImagesConfig(
                    aspect_ratio="9:16",
                    number_of_images=1
                )
            )
            if imagen.generated_images and len(imagen.generated_images) > 0:
                generated_image = imagen.generated_images[0]
                image = Image.open(BytesIO(generated_image.image.image_bytes))
                # Ensure the image matches the desired aspect ratio (9:16 = 0.5625)
                target_width = 360
                target_height = int(target_width / 9 * 16)  # 9:16 aspect ratio
                image = image.resize((target_width, target_height), Image.LANCZOS)
                # Convert image to base64
                buffered = BytesIO()
                image.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode()

                # Generate video if enabled
                if generate_video:
                    try:
                        # Enhance the image prompt for video generation
                        video_prompt = f"""
                        {image_prompt} Create a close-up shot with a slow dolly shot circling around the subject, 
                        using shallow focus on the main subject to emphasize details, in a realistic style with cinematic lighting.
                        """
                        operation = client.models.generate_videos(
                            model="veo-2.0-generate-001",
                            prompt=video_prompt,
                            image=generated_image.image,
                            config=types.GenerateVideosConfig(
                                aspect_ratio="9:16",
                                number_of_videos=2,
                                duration_seconds=8,
                                negative_prompt="blurry, low quality, text, letters",
                            )
                        )
                        # Wait for videos to generate
                        while not operation.done:
                            time.sleep(20)
                            operation = client.operations.get(operation)
                        
                        for n, video in enumerate(operation.response.generated_videos):
                            fname = f'with_image_input{n}.mp4'
                            print(f"Generated video: {fname}")
                            client.files.download(file=video.video)
                            video_buffer = BytesIO()
                            video.video.save(video_buffer)
                            video_base64 = base64.b64encode(video_buffer.getvalue()).decode()
                            video_base64_list.append(video_base64)
                    except Exception as e:
                        print(f"Error generating video: {e}")
                        video_base64_list = []  # Proceed without video if generation fails

                return {
                    'text': text,
                    'image_base64': img_str,
                    'video_base64_list': video_base64_list,
                    'ideas': ideas
                }
            else:
                print(f"Image generation failed (attempt {attempt + 1}): No images returned")
                if attempt == max_retries - 1:
                    # Last attempt, use a gray placeholder
                    image = Image.new('RGB', (360, 640), color='gray')
                    buffered = BytesIO()
                    image.save(buffered, format="PNG")
                    img_str = base64.b64encode(buffered.getvalue()).decode()
                    return {
                        'text': text,
                        'image_base64': img_str,
                        'video_base64_list': [],
                        'ideas': ideas
                    }
                # Retry with new ideas
                ideas = generate_ideas(tag)
                continue
        except Exception as e:
            print(f"Error generating image (attempt {attempt + 1}): {e}")
            if attempt == max_retries - 1:
                # Last attempt, use a gray placeholder
                image = Image.new('RGB', (360, 640), color='gray')
                buffered = BytesIO()
                image.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode()
                return {
                    'text': text,
                    'image_base64': img_str,
                    'video_base64_list': [],
                    'ideas': ideas
                }
            # Retry with new ideas
            ideas = generate_ideas(tag)
            continue

def start_feed(tag, generate_video, current_index, feed_items):
    """
    Start or update the feed based on the tag.
    
    Args:
        tag (str): The tag to generate content for.
        generate_video (bool): Whether to generate a video.
        current_index (int): The current item index.
        feed_items (list): The current list of feed items.
    
    Returns:
        tuple: (current_tag, current_index, feed_items, html_content, share_links, is_loading)
    """
    if not tag.strip():
        tag = "trending"
    
    # Set loading state
    is_loading = True
    html_content = generate_html([], False, 0, tag, is_loading)
    share_links = ""

    try:
        ideas = generate_ideas(tag)
        item = generate_item(tag, ideas, generate_video=generate_video)
        feed_items = [item]
        current_index = 0
        share_links = generate_share_links(
            item['image_base64'],
            item['video_base64_list'],
            item['text']
        )
    except Exception as e:
        print(f"Error in start_feed: {e}")
        feed_items = []
        current_index = 0
        html_content = """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Error generating content. Please try again!</p>
        </div>
        """
        is_loading = False
        return tag, current_index, feed_items, html_content, share_links, is_loading

    # Set loading state to False and update UI
    is_loading = False
    html_content = generate_html(feed_items, False, current_index, tag, is_loading)
    return tag, current_index, feed_items, html_content, share_links, is_loading

def load_next(tag, generate_video, current_index, feed_items):
    """
    Load the next item in the feed.
    
    Args:
        tag (str): The tag to generate content for.
        generate_video (bool): Whether to generate a video.
        current_index (int): The current item index.
        feed_items (list): The current list of feed items.
    
    Returns:
        tuple: (current_tag, current_index, feed_items, html_content, share_links, is_loading)
    """
    is_loading = True
    html_content = generate_html(feed_items, False, current_index, tag, is_loading)
    share_links = ""

    try:
        if current_index + 1 < len(feed_items):
            current_index += 1
        else:
            ideas = feed_items[-1]['ideas'] if feed_items else generate_ideas(tag)
            new_item = generate_item(tag, ideas, generate_video=generate_video)
            feed_items.append(new_item)
            current_index = len(feed_items) - 1
        share_links = generate_share_links(
            feed_items[current_index]['image_base64'],
            feed_items[current_index]['video_base64_list'],  # Fixed typo: IRfeed_items -> feed_items
            feed_items[current_index]['text']
        )
    except Exception as e:
        print(f"Error in load_next: {e}")
        html_content = """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Error generating content. Please try again!</p>
        </div>
        """
        is_loading = False
        return tag, current_index, feed_items, html_content, share_links, is_loading

    is_loading = False
    html_content = generate_html(feed_items, False, current_index, tag, is_loading)
    return tag, current_index, feed_items, html_content, share_links, is_loading

def load_previous(tag, generate_video, current_index, feed_items):
    """
    Load the previous item in the feed.
    
    Args:
        tag (str): The tag to generate content for.
        generate_video (bool): Whether to generate a video (not used here).
        current_index (int): The current item index.
        feed_items (list): The current list of feed items.
    
    Returns:
        tuple: (current_tag, current_index, feed_items, html_content, share_links, is_loading)
    """
    if current_index > 0:
        current_index -= 1
    html_content = generate_html(feed_items, False, current_index, tag, False)
    share_links = generate_share_links(
        feed_items[current_index]['image_base64'],
        feed_items[current_index]['video_base64_list'],
        feed_items[current_index]['text']
    )
    return tag, current_index, feed_items, html_content, share_links, False

def generate_share_links(image_base64, video_base64_list, caption):
    """
    Generate share links for social media platforms with download links for image and video.
    
    Args:
        image_base64 (str): The base64-encoded image data.
        video_base64_list (list): List of base64-encoded video data.
        caption (str): The caption to share.
    
    Returns:
        str: HTML string with share links and download instructions.
    """
    image_data_url = f"data:image/png;base64,{image_base64}"
    encoded_caption = urllib.parse.quote(caption)

    # Generate download links for image and videos
    download_links = f"""
        <p style="text-align: center;">Download and attach the image/video to share:</p>
        <a href="{image_data_url}" download="feed_item.png" style="
            background-color: #4CAF50;
            color: white;
            padding: 5px 10px;
            border-radius: 5px;
            text-decoration: none;
            margin: 5px;
        ">Download Image</a>
    """
    for i, video_base64 in enumerate(video_base64_list):
        video_data_url = f"data:video/mp4;base64,{video_base64}"
        download_links += f"""
            <a href="{video_data_url}" download="feed_video_{i}.mp4" style="
                background-color: #4CAF50;
                color: white;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
                margin: 5px;
            ">Download Video {i+1}</a>
        """

    # Generate share links using only the caption
    share_links = f"""
    <div style="
        display: flex;
        flex-direction: column;
        align-items: center;
        gap: 10px;
        margin-top: 10px;
        color: white;
        font-family: Arial, sans-serif;
    ">
        {download_links}
        <div style="
            display: flex;
            flex-wrap: wrap;
            justify-content: center;
            gap: 10px;
            margin-top: 10px;
        ">
            <a href="https://www.tiktok.com/upload?caption={caption}" target="_blank" style="
                background-color: #00f2ea;
                color: #000;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
            ">Share on TikTok</a>
            <a href="https://www.instagram.com/?caption={caption}" target="_blank" style="
                background-color: #e1306c;
                color: white;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
            ">Share on Instagram</a>
            <a href="https://www.facebook.com/sharer/sharer.php?quote={caption}" target="_blank" style="
                background-color: #4267b2;
                color: white;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
            ">Share on Facebook</a>
            <a href="https://twitter.com/intent/tweet?text={caption}" target="_blank" style="
                background-color: #1da1f2;
                color: white;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
            ">Share on X</a>
            <a href="https://pinterest.com/pin/create/button/?description={caption}" target="_blank" style="
                background-color: #bd081c;
                color: white;
                padding: 5px 10px;
                border-radius: 5px;
                text-decoration: none;
            ">Share on Pinterest</a>
        </div>
    </div>
    """
    return share_links

def generate_html(feed_items, scroll_to_latest=False, current_index=0, tag="", is_loading=False):
    """
    Generate an HTML string to display the current feed item with click navigation.
    
    Args:
        feed_items (list): List of dictionaries containing 'text' and 'image_base64'.
        scroll_to_latest (bool): Whether to auto-scroll to the latest item (not used here).
        current_index (int): The index of the item to display.
        tag (str): The current tag for loading messages.
        is_loading (bool): Whether the feed is currently loading.
    
    Returns:
        str: HTML string representing the feed.
    """
    loading_messages = [
        f"Cooking up a {tag} masterpiece... 🍳",
        f"Snapping a vibrant {tag} moment... 📸",
        f"Creating a {tag} vibe that pops... ✨",
        f"Getting that perfect {tag} shot... 🎥",
        f"Bringing {tag} to life... 🌟"
    ]

    if is_loading:
        return f"""
        <div id="feed-container" style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
            position: relative;
        ">
            <div id="loading-message" style="
                font-size: 18px;
                font-weight: bold;
                text-align: center;
                margin-bottom: 20px;
                text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
            ">
                {loading_messages[0]}
            </div>
            <div style="
                width: 80%;
                height: 10px;
                background-color: #333;
                border-radius: 5px;
                overflow: hidden;
            ">
                <div style="
                    width: 0%;
                    height: 100%;
                    background: linear-gradient(to right, #ff2d55, #ff5e78);
                    animation: loading 2s infinite;
                "></div>
            </div>
            <style>
                @keyframes loading {{
                    0% {{ width: 0%; }}
                    50% {{ width: 100%; }}
                    100% {{ width: 0%; }}
                }}
            </style>
            <script>
                const messages = {json.dumps(loading_messages)};
                let currentMessageIndex = 0;
                const messageElement = document.getElementById('loading-message');
                function rotateMessages() {{
                    currentMessageIndex = (currentMessageIndex + 1) % messages.length;
                    messageElement.textContent = messages[currentMessageIndex];
                }}
                setInterval(rotateMessages, 2000);
            </script>
        </div>
        """

    if not feed_items or current_index >= len(feed_items):
        return """
        <div style="
            display: flex;
            flex-direction: column;
            align-items: center;
            justify-content: center;
            max-width: 360px;
            margin: 0 auto;
            background-color: #000;
            height: 640px;
            border: 1px solid #333;
            border-radius: 10px;
            color: white;
            font-family: Arial, sans-serif;
        ">
            <p>Enter a concept or idea to start your feed!</p>
        </div>
        """

    item = feed_items[current_index]
    html_str = """
    <div id="feed-container" style="
        display: flex;
        flex-direction: column;
        align-items: center;
        max-width: 360px;
        margin: 0 auto;
        background-color: #000;
        height: 640px;
        border: 1px solid #333;
        border-radius: 10px;
        position: relative;
    ">
        <div class="feed-item" style="
            width: 100%;
            height: 100%;
            position: relative;
            display: flex;
            flex-direction: column;
            justify-content: flex-end;
            overflow: hidden;
            cursor: pointer;
        " onclick="handleClick(event)">
            <img id="feed-image" src="data:image/png;base64,{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);
            ">
                {text}
            </div>
        </div>
    </div>
    <script>
        function handleClick(event) {{
            const image = document.getElementById('feed-image');
            const rect = image.getBoundingClientRect();
            const clickX = event.clientX - rect.left;
            const width = rect.width;
            if (clickX > width * 0.75) {{
                document.getElementById('previous-button').click();
            }}
        }}
    </script>
    <button id="previous-button" style="display: none;" onclick="document.getElementById('previous-button').click()"></button>
    """.format(image_base64=item['image_base64'], text=item['text'])

    return html_str

# Define the Gradio interface
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, .gr-checkbox { 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="Create Your Feed"
) as demo:
    # State variables
    current_tag = gr.State(value="")
    current_index = gr.State(value=0)
    feed_items = gr.State(value=[])
    is_loading = gr.State(value=False)
    share_links = gr.State(value="")

    # Input section
    with gr.Column(elem_classes="gr-form"):
        gr.Markdown("### Create Your Feed")
        tag_input = gr.Textbox(
            label="Enter Concept or Ideas",
            value="",
            placeholder="e.g., sushi adventure, neon tech",
            submit_btn=False
        )
        generate_video_checkbox = gr.Checkbox(
            label="Generate Video (may take longer)",
            value=False
        )
        magic_button = gr.Button("✨ Generate Next Item", elem_classes="gr-button")

    # Output display
    feed_html = gr.HTML()
    share_html = gr.HTML(label="Share this item:")

    # Event handlers
    # Handle Enter keypress in the concept input
    tag_input.submit(
        fn=start_feed,
        inputs=[tag_input, generate_video_checkbox, current_index, feed_items],
        outputs=[current_tag, current_index, feed_items, feed_html, share_html, is_loading]
    )

    # Handle magic button click to generate next item
    magic_button.click(
        fn=load_next,
        inputs=[current_tag, generate_video_checkbox, current_index, feed_items],
        outputs=[current_tag, current_index, feed_items, feed_html, share_html, is_loading]
    )

    # Hidden button for previous item navigation
    previous_button = gr.Button("Previous", elem_id="previous-button", visible=False)

    # Handle click to go to previous item
    previous_button.click(
        fn=load_previous,
        inputs=[current_tag, generate_video_checkbox, current_index, feed_items],
        outputs=[current_tag, current_index, feed_items, feed_html, share_html, is_loading]
    )

# Launch the app
demo.launch()