Update app.py

app.py

@@ -61,14 +61,17 @@ def clean_response_text(response_text):
 def generate_ideas(user_input):
     """
     Generate a diverse set of ideas based on the user's input concept using the LLM.
+    Yields progress updates for the loading UI.
     
     Args:
-        user_input (str): The user's input concept or idea (e.g., "blindfolded Rubik's Cube challenge").
+        user_input (str): The user's input concept or idea.
     
-    Returns:
-        list: A list of ideas as strings.
+    Yields:
+        tuple: (progress_percentage, message) for UI updates.
+        list: Final list of ideas as strings.
     """
-    # Generate ideas using gemini-2.0-flash-lite
+    yield (10, f"Brainstorming epic ideas for {user_input}... 🌟")
+    
     prompt = f"""
     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.
@@ -102,10 +105,12 @@ def generate_ideas(user_input):
             raise ValueError("Invalid JSON format: 'ideas' key missing, not a list, or incorrect length")
         
         ideas = response_json['ideas']
+        yield (20, f"Ideas locked in for {user_input}! 🚀")
         return ideas
     
     except Exception as e:
         print(f"Error generating ideas: {e}")
+        yield (20, f"Oops, tweaking the plan for {user_input}... 🔧")
         return [
             f"A dramatic {user_input} scene with cinematic lighting",
             f"A close-up of {user_input} in a futuristic setting",
@@ -117,6 +122,7 @@ def generate_ideas(user_input):
 def generate_item(user_input, ideas, generate_video=False, max_retries=3):
     """
     Generate a single feed item (image and optionally one video) using one of the ideas.
+    Yields progress updates for the loading UI.
     
     Args:
         user_input (str): The user's input concept or idea.
@@ -124,8 +130,9 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
         generate_video (bool): Whether to generate a video from the image.
         max_retries (int): Maximum number of retries for image generation per cycle.
     
-    Returns:
-        dict: A dictionary with 'text' (str), 'image_base64' (str), 'video_base64' (str or None), and 'ideas' (list).
+    Yields:
+        tuple: (progress_percentage, message) for UI updates.
+        dict: Final dictionary with 'text', 'image_base64', 'video_base64', and 'ideas'.
     """
     video_base64 = None
     max_total_attempts = 3  # Maximum total attempts for combined image and video generation cycles
@@ -133,14 +140,16 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
     total_attempts = 0
     while total_attempts < max_total_attempts:
         total_attempts += 1
+        yield (20 + total_attempts * 10, f"Attempt {total_attempts} to craft your {user_input} masterpiece... 🎨")
         
         # Step 1: Generate an image (retry up to max_retries times)
-        generated_image = None  # Initialize to None
+        generated_image = None
         text = None
         img_str = None
         image_prompt = None
         
         for image_attempt in range(max_retries):
+            yield (20 + total_attempts * 10 + image_attempt * 5, f"Crafting a stunning image for {user_input}... 📸")
             selected_idea = random.choice(ideas)
             prompt = f"""
             The user has provided the concept: "{user_input}". Based on this concept and the specific idea "{selected_idea}", create content for a TikTok video.
@@ -176,6 +185,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
 
             # Attempt to generate the image
             try:
+                yield (40 + image_attempt * 5, f"Rendering your {user_input} vision... ✨")
                 imagen = client.models.generate_images(
                     model='imagen-3.0-generate-002',
                     prompt=image_prompt,
@@ -195,42 +205,43 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     buffered = BytesIO()
                     image.save(buffered, format="PNG")
                     img_str = base64.b64encode(buffered.getvalue()).decode()
+                    yield (50, f"Image for {user_input} is ready! 🎉")
                     break  # Successfully generated image, exit image retry loop
                 else:
                     print(f"Image generation failed (image attempt {image_attempt + 1}, total attempt {total_attempts}): No images returned")
                     if image_attempt == max_retries - 1:
-                        # Last image attempt in this cycle
+                        yield (50, f"Tweaking the image for {user_input}... 🔄")
                         if total_attempts == max_total_attempts:
                             # Max total attempts reached, 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()
+                            yield (60, f"Using a placeholder for {user_input}... 🖼️")
                             return {
                                 'text': text,
                                 'image_base64': img_str,
                                 'video_base64': None,
                                 'ideas': ideas
                             }
-                        # Otherwise, continue to next cycle
                         break  # Exit inner loop to retry with new idea
             except Exception as e:
                 print(f"Error generating image (image attempt {image_attempt + 1}, total attempt {total_attempts}): {e}")
                 if image_attempt == max_retries - 1:
-                    # Last image attempt in this cycle
+                    yield (50, f"Retrying image for {user_input}... 🔄")
                     if total_attempts == max_total_attempts:
                         # Max total attempts reached, 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()
+                        yield (60, f"Using a placeholder for {user_input}... 🖼️")
                         return {
                             'text': text,
                             'image_base64': img_str,
                             'video_base64': None,
                             'ideas': ideas
                         }
-                    # Otherwise, continue to next cycle
                     break  # Exit inner loop to retry with new idea
 
         # Step 2: Generate video if enabled (with fallback to text-to-video if image-to-video fails)
@@ -240,6 +251,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
 
             # First, try image-to-video generation (only once)
             try:
+                yield (60, f"Filming a viral video for {user_input}... 🎥")
                 video_prompt = f"""
                 The user concept is "{user_input}". Based on this and the scene: {image_prompt}, create a video.
                 Use a close-up shot with a slow dolly shot circling around the subject, 
@@ -309,6 +321,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     # Encode the video bytes as base64
                     video_base64 = base64.b64encode(video_bytes).decode()
                     video_generated = True
+                    yield (90, f"Video for {user_input} is a wrap! 🎬")
                     # Successfully generated video, return the result
                     return {
                         'text': text,
@@ -320,12 +333,14 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     raise ValueError("No video was generated")
             except Exception as e:
                 print(f"Error generating video (image-to-video, total attempt {total_attempts}): {e}")
+                yield (70, f"Switching to a new video approach for {user_input}... 🎞️")
                 print("Image-to-video generation failed. Falling back to text-to-video generation.")
 
             # If image-to-video generation failed, try text-to-video generation
             if not video_generated:
                 for video_attempt in range(max_video_retries_per_image):
                     try:
+                        yield (75 + video_attempt * 5, f"Trying a fresh video take for {user_input}... 📹")
                         # Use the same video prompt but without the image
                         video_prompt_base = f"""
                         The user concept is "{user_input}". Based on this and the scene: {image_prompt}, create a video.
@@ -403,6 +418,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                             # Encode the video bytes as base64
                             video_base64 = base64.b64encode(video_bytes).decode()
                             video_generated = True
+                            yield (90, f"Video for {user_input} is a wrap! 🎬")
                             # Successfully generated video, return the result
                             return {
                                 'text': text,
@@ -415,9 +431,11 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     except Exception as e:
                         print(f"Error generating video (text-to-video attempt {video_attempt + 1}, total attempt {total_attempts}): {e}")
                         if video_attempt == max_video_retries_per_image - 1:
+                            yield (85, f"Finalizing without video for {user_input}... 📌")
                             if total_attempts == max_total_attempts:
                                 print("Max total attempts reached. Proceeding without video.")
                                 video_base64 = None
+                                yield (95, f"Polishing your {user_input} masterpiece... ✨")
                                 return {
                                     'text': text,
                                     'image_base64': img_str,
@@ -431,6 +449,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
         
         # If video generation is not enabled or image generation failed, return the result
         if img_str is not None:
+            yield (95, f"Polishing your {user_input} masterpiece... ✨")
             return {
                 'text': text,
                 'image_base64': img_str,
@@ -441,10 +460,12 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
 
     # If max total attempts reached without success, use a gray placeholder image
     print("Max total attempts reached without successful image generation. Using placeholder.")
+    yield (95, f"Falling back to a placeholder for {user_input}... 🖼️")
     image = Image.new('RGB', (360, 640), color='gray')
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode()
+    yield (100, f"Ready to roll with {user_input}! 🚀")
     return {
         'text': f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}",
         'image_base64': img_str,
@@ -455,6 +476,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
 def start_feed(user_input, generate_video, current_index, feed_items):
     """
     Start or update the feed based on the user input.
+    Yields progress updates and final feed content.
     
     Args:
         user_input (str): The user's input concept or idea.
@@ -462,8 +484,9 @@ def start_feed(user_input, generate_video, current_index, feed_items):
         current_index (int): The current item index.
         feed_items (list): The current list of feed items.
     
-    Returns:
-        tuple: (current_user_input, current_index, feed_items, html_content, share_links, is_loading)
+    Yields:
+        tuple: (progress_percentage, message, html_content, is_loading) for UI updates.
+        tuple: (current_user_input, current_index, feed_items, html_content, share_links, is_loading) for final result.
     """
     if not user_input.strip():
         user_input = "trending"
@@ -471,14 +494,35 @@ def start_feed(user_input, generate_video, current_index, feed_items):
     # Update current_user_input with the new user_input
     current_user_input = user_input
     
-    # Set loading state
+    # Set initial loading state
     is_loading = True
-    html_content = generate_html([], False, 0, user_input, is_loading)
+    html_content = generate_html([], False, 0, user_input, is_loading, 0, f"Getting started with {user_input}... 🚀")
+    yield (0, f"Getting started with {user_input}... 🚀", html_content, is_loading)
     share_links = ""
 
     try:
-        ideas = generate_ideas(user_input)
-        item = generate_item(user_input, ideas, generate_video=generate_video)
+        # Generate ideas with progress updates
+        ideas_gen = generate_ideas(user_input)
+        ideas = None
+        for update in ideas_gen:
+            if isinstance(update, tuple):
+                progress, message = update
+                html_content = generate_html([], False, 0, user_input, is_loading, progress, message)
+                yield (progress, message, html_content, is_loading)
+            else:
+                ideas = update
+        
+        # Generate item with progress updates
+        item_gen = generate_item(user_input, ideas, generate_video=generate_video)
+        item = None
+        for update in item_gen:
+            if isinstance(update, tuple):
+                progress, message = update
+                html_content = generate_html([], False, 0, user_input, is_loading, progress, message)
+                yield (progress, message, html_content, is_loading)
+            else:
+                item = update
+        
         feed_items = [item]
         current_index = 0
         share_links = generate_share_links(
@@ -509,16 +553,19 @@ def start_feed(user_input, generate_video, current_index, feed_items):
         </div>
         """
         is_loading = False
+        yield (100, "Oops, something went wrong! 😅", html_content, is_loading)
         return current_user_input, 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, user_input, is_loading)
+    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, 100, "")
+    yield (100, f"Ready to roll with {user_input}! 🚀", html_content, is_loading)
     return current_user_input, current_index, feed_items, html_content, share_links, is_loading
 
 def load_next(user_input, generate_video, current_index, feed_items):
     """
     Load the next item in the feed.
+    Yields progress updates and final feed content.
     
     Args:
         user_input (str): The user's input concept or idea (updated from the textbox).
@@ -526,8 +573,9 @@ def load_next(user_input, generate_video, current_index, feed_items):
         current_index (int): The current item index.
         feed_items (list): The current list of feed items.
     
-    Returns:
-        tuple: (current_user_input, current_index, feed_items, html_content, share_links, is_loading)
+    Yields:
+        tuple: (progress_percentage, message, html_content, is_loading) for UI updates.
+        tuple: (current_user_input, current_index, feed_items, html_content, share_links, is_loading) for final result.
     """
     # Update current_user_input with the latest user_input from the textbox
     current_user_input = user_input if user_input.strip() else "trending"
@@ -536,17 +584,41 @@ def load_next(user_input, generate_video, current_index, feed_items):
     user_input = current_user_input
     
     is_loading = True
-    html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
+    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, 0, f"Loading next {user_input} vibe... 🚀")
+    yield (0, f"Loading next {user_input} vibe... 🚀", html_content, is_loading)
     share_links = ""
 
     try:
         if current_index + 1 < len(feed_items):
             current_index += 1
+            yield (50, f"Switching to the next {user_input} moment... 🔄", html_content, is_loading)
         else:
-            ideas = feed_items[-1]['ideas'] if feed_items else generate_ideas(user_input)
-            new_item = generate_item(user_input, ideas, generate_video=generate_video)
+            ideas = feed_items[-1]['ideas'] if feed_items else None
+            if not ideas:
+                # Generate new ideas if none exist
+                ideas_gen = generate_ideas(user_input)
+                for update in ideas_gen:
+                    if isinstance(update, tuple):
+                        progress, message = update
+                        html_content = generate_html(feed_items, False, current_index, user_input, is_loading, progress, message)
+                        yield (progress, message, html_content, is_loading)
+                    else:
+                        ideas = update
+            
+            # Generate new item
+            new_item_gen = generate_item(user_input, ideas, generate_video=generate_video)
+            new_item = None
+            for update in new_item_gen:
+                if isinstance(update, tuple):
+                    progress, message = update
+                    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, progress, message)
+                    yield (progress, message, html_content, is_loading)
+                else:
+                    new_item = update
+            
             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'],
@@ -573,10 +645,12 @@ def load_next(user_input, generate_video, current_index, feed_items):
         </div>
         """
         is_loading = False
+        yield (100, "Oops, something went wrong! 😅", html_content, is_loading)
         return current_user_input, current_index, feed_items, html_content, share_links, is_loading
 
     is_loading = False
-    html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
+    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, 100, "")
+    yield (100, f"Next {user_input} vibe is live! 🚀", html_content, is_loading)
     return current_user_input, current_index, feed_items, html_content, share_links, is_loading
 
 def load_previous(user_input, generate_video, current_index, feed_items):
@@ -597,7 +671,7 @@ def load_previous(user_input, generate_video, current_index, feed_items):
     
     if current_index > 0:
         current_index -= 1
-    html_content = generate_html(feed_items, False, current_index, user_input, False)
+    html_content = generate_html(feed_items, False, current_index, user_input, False, 100, "")
     share_links = generate_share_links(
         feed_items[current_index]['image_base64'],
         feed_items[current_index]['video_base64'],
@@ -730,7 +804,7 @@ def generate_share_links(image_base64, video_base64, caption):
                 justify-content: center;
                 margin-top: 10px;
             ">
-                <a href="https://www.youtube.com/upload?description={caption}" target="_blank" style="
+                <a href="https://studio.youtube.com/channel/UC/videos/upload?description={caption}" target="_blank" style="
                     background-color: #ff0000;
                     color: white;
                     padding: 8px 16px;
@@ -760,10 +834,9 @@ def generate_share_links(image_base64, video_base64, caption):
     </div>
     """.format(caption=encoded_caption)
 
-def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_input="", is_loading=False):
+def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_input="", is_loading=False, progress=0, message=""):
     """
-    Generate an HTML string to display the current feed item with click navigation.
-    Displays a video if available, otherwise falls back to the image.
+    Generate an HTML string to display the current feed item or loading state with progress bar.
     
     Args:
         feed_items (list): List of dictionaries containing 'text', 'image_base64', and 'video_base64'.
@@ -771,18 +844,12 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
         current_index (int): The index of the item to display.
         user_input (str): The user's input concept or idea for loading messages.
         is_loading (bool): Whether the feed is currently loading.
+        progress (float): Current progress percentage (0–100).
+        message (str): Current loading message to display.
     
     Returns:
-        str: HTML string representing the feed.
+        str: HTML string representing the feed or loading state.
     """
-    loading_messages = [
-        f"Cooking up a {user_input} masterpiece... 🍳",
-        f"Snapping a vibrant {user_input} moment... 📸",
-        f"Creating a {user_input} vibe that pops... ✨",
-        f"Getting that perfect {user_input} shot... 🎥",
-        f"Bringing {user_input} to life... 🌟"
-    ]
-
     if is_loading:
         return f"""
         <div id="feed-container" style="
@@ -807,7 +874,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
                 margin-bottom: 20px;
                 text-shadow: 1px 1px 2px rgba(0,0,0,0.5);
             ">
-                {loading_messages[0]}
+                {message}
             </div>
             <div style="
                 width: 80%;
@@ -816,30 +883,30 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
                 border-radius: 5px;
                 overflow: hidden;
             ">
-                <div style="
-                    width: 0%;
+                <div id="progress-bar" style="
+                    width: {progress}%;
                     height: 100%;
                     background: linear-gradient(to right, #ff2d55, #ff5e78);
-                    animation: loading 2s infinite;
+                    transition: width 0.5s ease-in-out;
                 "></div>
             </div>
+            <div style="
+                margin-top: 10px;
+                font-size: 14px;
+                color: #ccc;
+            ">
+                {int(progress)}% Complete
+            </div>
             <style>
-                @keyframes loading {{
-                    0% {{ width: 0%; }}
-                    50% {{ width: 100%; }}
-                    100% {{ width: 0%; }}
+                @keyframes pulse {{
+                    0% {{ opacity: 1; }}
+                    50% {{ opacity: 0.5; }}
+                    100% {{ opacity: 1; }}
                 }}
-            </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];
+                #loading-message {{
+                    animation: pulse 2s infinite;
                 }}
-                setInterval(rotateMessages, 2000);
-            </script>
+            </style>
         </div>
         """
 
@@ -992,14 +1059,40 @@ with gr.Blocks(
     user_input.submit(
         fn=start_feed,
         inputs=[user_input, generate_video_checkbox, current_index, feed_items],
-        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading]
+        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading],
+        _js="""
+        async (user_input, generate_video, current_index, feed_items) => {
+            const results = [];
+            for await (const update of gradioApp().fn_stream(start_feed, [user_input, generate_video, current_index, feed_items])) {
+                results.push(update);
+                if (update.length === 4) { // Progress update
+                    const [progress, message, html_content, is_loading] = update;
+                    gradioApp().update_component('feed_html', { value: html_content });
+                }
+            }
+            return results[results.length - 1]; // Return final result
+        }
+        """
     )
 
     # Handle magic button click to generate next item
     magic_button.click(
         fn=load_next,
         inputs=[user_input, generate_video_checkbox, current_index, feed_items],
-        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading]
+        outputs=[current_user_input, current_index, feed_items, feed_html, share_html, is_loading],
+        _js="""
+        async (user_input, generate_video, current_index, feed_items) => {
+            const results = [];
+            for await (const update of gradioApp().fn_stream(load_next, [user_input, generate_video, current_index, feed_items])) {
+                results.push(update);
+                if (update.length === 4) { // Progress update
+                    const [progress, message, html_content, is_loading] = update;
+                    gradioApp().update_component('feed_html', { value: html_content });
+                }
+            }
+            return results[results.length - 1]; // Return final result
+        }
+        """
     )
 
     # Hidden button for previous item navigation