Update app.py

app.py

@@ -45,11 +45,19 @@ def generate_ideas(user_input):
         list: A list of ideas as strings.
     """
     prompt = f"""
-    The user has provided the concept: "{user_input}". Based on this concept, generate a list of 5 diverse and creative ideas for a TikTok video that directly relate to "{user_input}".
-    Each idea should be a short sentence describing a specific scene or concept, ensuring the core theme of "{user_input}" is central to each idea.
+    The user has provided the concept: "{user_input}". You must generate 5 diverse and creative ideas for a TikTok video that are directly and explicitly related to "{user_input}".
+    Each idea must clearly incorporate and focus on the core theme of "{user_input}" without deviating into unrelated topics.
+    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 Rubik's Cube challenge in the dark", "A blindfolded speedcubing competition with dramatic music"]}}
+    Example for "blindfolded Rubik's Cube challenge":
+    {{"ideas": [
+        "A blindfolded speedcubing competition with dramatic music",
+        "A close-up of a person solving a Rubik's Cube blindfolded under a spotlight",
+        "A time-lapse of a blindfolded Rubik's Cube solve with colorful lighting",
+        "A blindfolded Rubik's Cube challenge in a futuristic setting",
+        "A split-screen of two people racing to solve a Rubik's Cube blindfolded"
+    ]}}
     """
     try:
         response = client.models.generate_content(
@@ -62,9 +70,22 @@ def generate_ideas(user_input):
             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']
+        if 'ideas' not in response_json or not isinstance(response_json['ideas'], list) or len(response_json['ideas']) != 5:
+            raise ValueError("Invalid JSON format: 'ideas' key missing, not a list, or incorrect length")
+        
+        # Validate that ideas are related to user_input
+        ideas = response_json['ideas']
+        for idea in ideas:
+            if user_input.lower() not in idea.lower():
+                print(f"Warning: Idea '{idea}' does not seem related to '{user_input}'. Falling back to default ideas.")
+                return [
+                    f"A dramatic {user_input} scene with cinematic lighting",
+                    f"A close-up of {user_input} in a futuristic setting",
+                    f"A high-energy {user_input} moment with vibrant colors",
+                    f"A serene {user_input} scene with soft focus",
+                    f"An action-packed {user_input} challenge with dynamic angles"
+                ]
+        return ideas
     except Exception as e:
         print(f"Error generating ideas: {e}")
         return [
@@ -161,7 +182,6 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                                 number_of_videos=2,
                                 duration_seconds=8,
                                 negative_prompt="blurry, low quality, text, letters"
-                                # Removed enhance_prompt=True since it's the default
                             )
                         )
                         # Wait for videos to generate
@@ -183,10 +203,19 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                                 raise ValueError(f"Video {n} is invalid or missing video data")
                             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()
+                            # Download the video and get the raw bytes
+                            video_data = client.files.download(file=video.video)
+                            # Ensure video_data is in bytes
+                            if isinstance(video_data, bytes):
+                                video_bytes = video_data
+                            else:
+                                # If video_data is a file-like object, read the bytes
+                                video_buffer = BytesIO()
+                                for chunk in video_data:
+                                    video_buffer.write(chunk)
+                                video_bytes = video_buffer.getvalue()
+                            # Encode the video bytes as base64
+                            video_base64 = base64.b64encode(video_bytes).decode()
                             video_base64_list.append(video_base64)
                     except Exception as e:
                         print(f"Error generating video: {e}")
@@ -476,9 +505,10 @@ def generate_share_links(image_base64, video_base64_list, caption):
 def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_input="", is_loading=False):
     """
     Generate an HTML string to display the current feed item with click navigation.
+    Displays a video if available, otherwise falls back to the image.
     
     Args:
-        feed_items (list): List of dictionaries containing 'text' and 'image_base64'.
+        feed_items (list): List of dictionaries containing 'text', 'image_base64', and 'video_base64_list'.
         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.
         user_input (str): The user's input concept or idea for loading messages.
@@ -576,7 +606,37 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
         """
 
     item = feed_items[current_index]
-    html_str = """
+    # Check if there are videos to display; if so, show the first video
+    if item['video_base64_list']:
+        media_element = f"""
+            <video id="feed-video" controls style="
+                width: 100%;
+                height: 100%;
+                object-fit: cover;
+                position: absolute;
+                top: 0;
+                left: 0;
+                z-index: 1;
+            ">
+                <source src="data:video/mp4;base64,{item['video_base64_list'][0]}" type="video/mp4">
+                Your browser does not support the video tag.
+            </video>
+        """
+    else:
+        # Fallback to image if no video is available
+        media_element = f"""
+            <img id="feed-image" 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;
+            ">
+        """
+
+    html_str = f"""
     <div id="feed-container" style="
         display: flex;
         flex-direction: column;
@@ -599,15 +659,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
             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;
-            ">
+            {media_element}
             <div style="
                 position: relative;
                 z-index: 2;
@@ -619,14 +671,14 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
                 font-weight: bold;
                 text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
             ">
-                {text}
+                {item['text']}
             </div>
         </div>
     </div>
     <script>
         function handleClick(event) {{
-            const image = document.getElementById('feed-image');
-            const rect = image.getBoundingClientRect();
+            const media = document.getElementById('feed-video') || document.getElementById('feed-image');
+            const rect = media.getBoundingClientRect();
             const clickX = event.clientX - rect.left;
             const width = rect.width;
             if (clickX > width * 0.75) {{
@@ -635,8 +687,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
         }}
     </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