Update app.py

app.py

@@ -388,48 +388,43 @@ def start_feed(user_input, generate_video, current_index, feed_items):
         ideas = generate_ideas(user_input)
         
         item_generator = generate_item(user_input, ideas, generate_video=generate_video)
-        for progress in item_generator:
-            if time.time() - start_time > timeout_seconds:
-                logging.error("Generation timed out")
-                raise TimeoutError("Feed generation timed out")
-            logging.debug(f"Progress update: {progress}")
-            if isinstance(progress, dict) and "stage" in progress:
-                html_content = generate_html([], False, 0, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
-                yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
-            else:
-                logging.debug("Received final item from generator")
-                item = progress
+        try:
+            while True:
+                if time.time() - start_time > timeout_seconds:
+                    logging.error("Generation timed out")
+                    raise TimeoutError("Feed generation timed out")
+                progress = next(item_generator)
+                logging.debug(f"Progress update: {progress}")
+                if isinstance(progress, dict) and "stage" in progress:
+                    html_content = generate_html([], False, 0, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
+                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
+                else:
+                    logging.debug(f"Received final item: {progress}")
+                    item = progress
+                    feed_items = [item]
+                    current_index = 0
+                    share_links = generate_share_links(item['image_base64'], item['video_base64'], item['text'])
+                    is_loading = False
+                    html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
+                    logging.debug("Feed generation complete")
+                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
+                    return
+        except StopIteration as e:
+            # Handle the generator's return value
+            if e.value is not None:
+                logging.debug(f"Generator returned final item: {e.value}")
+                item = e.value
                 feed_items = [item]
                 current_index = 0
                 share_links = generate_share_links(item['image_base64'], item['video_base64'], item['text'])
                 is_loading = False
                 html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
-                logging.debug("Feed generation complete")
+                logging.debug("Feed generation complete via StopIteration")
                 yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                 return
-        logging.warning("Generator exhausted without returning final item")
-        # Fallback if generator doesn't return
-        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>Generation incomplete. Please try again!</p>
-        </div>
-        """
-        is_loading = False
-        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
-        return
+            else:
+                logging.warning("Generator returned without a final item")
+                raise ValueError("Generator did not return a valid item")
     except Exception as e:
         logging.error(f"Error in start_feed: {e}")
         html_content = """
@@ -488,17 +483,36 @@ def load_next(user_input, generate_video, current_index, feed_items):
             yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
             
             item_generator = generate_item(user_input, ideas, generate_video=generate_video)
-            for progress in item_generator:
-                if time.time() - start_time > timeout_seconds:
-                    logging.error("Generation timed out")
-                    raise TimeoutError("Feed generation timed out")
-                logging.debug(f"Progress update: {progress}")
-                if isinstance(progress, dict) and "stage" in progress:
-                    html_content = generate_html(feed_items, False, current_index, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
-                    yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
-                else:
-                    logging.debug("Received final item from generator")
-                    new_item = progress
+            try:
+                while True:
+                    if time.time() - start_time > timeout_seconds:
+                        logging.error("Generation timed out")
+                        raise TimeoutError("Feed generation timed out")
+                    progress = next(item_generator)
+                    logging.debug(f"Progress update: {progress}")
+                    if isinstance(progress, dict) and "stage" in progress:
+                        html_content = generate_html(feed_items, False, current_index, user_input, is_loading, progress["stage"], progress["message"], progress["progress"])
+                        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
+                    else:
+                        logging.debug(f"Received final item: {progress}")
+                        new_item = progress
+                        feed_items.append(new_item)
+                        current_index = len(feed_items) - 1
+                        is_loading = False
+                        html_content = generate_html(feed_items, False, current_index, user_input, is_loading)
+                        share_links = generate_share_links(
+                            feed_items[current_index]['image_base64'],
+                            feed_items[current_index]['video_base64'],
+                            feed_items[current_index]['text']
+                        )
+                        logging.debug("New feed item generated")
+                        yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
+                        return
+            except StopIteration as e:
+                # Handle the generator's return value
+                if e.value is not None:
+                    logging.debug(f"Generator returned final item: {e.value}")
+                    new_item = e.value
                     feed_items.append(new_item)
                     current_index = len(feed_items) - 1
                     is_loading = False
@@ -508,32 +522,12 @@ def load_next(user_input, generate_video, current_index, feed_items):
                         feed_items[current_index]['video_base64'],
                         feed_items[current_index]['text']
                     )
-                    logging.debug("New feed item generated")
+                    logging.debug("New feed item generated via StopIteration")
                     yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                     return
-            logging.warning("Generator exhausted without returning final item")
-            # Fallback if generator doesn't return
-            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>Generation incomplete. Please try again!</p>
-            </div>
-            """
-            is_loading = False
-            yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
-            return
+                else:
+                    logging.warning("Generator returned without a final item")
+                    raise ValueError("Generator did not return a valid item")
     except Exception as e:
         logging.error(f"Error in load_next: {e}")
         html_content = """
@@ -602,15 +596,15 @@ def generate_share_links(image_base64, video_base64, caption):
         video_data_url = f"data:video/mp4;base64,{video_base64}"
         download_links += f"""
             <a href="{video_data_url}" download="feed_video.mp4" style="
-            background-color: #4CAF50;
-            color: white;
-            padding: 8px 16px;
-            border-radius: 5px;
-            text-decoration: none;
-            font-size: 14px;
-            font-weight: bold;
-            transition: background-color 0.3s;
-        " onmouseover="this.style.backgroundColor='#45a049'" onmouseout="this.style.backgroundColor='#4CAF50'">Download Video</a>
+                background-color: #4CAF50;
+                color: white;
+                padding: 8px 16px;
+                border-radius: 5px;
+                text-decoration: none;
+                font-size: 14px;
+                font-weight: bold;
+                transition: background-color 0.3s;
+            " onmouseover="this.style.backgroundColor='#45a049'" onmouseout="this.style.backgroundColor='#4CAF50'">Download Video</a>
         """
     download_links += "</div>"
 
@@ -924,7 +918,7 @@ with gr.Blocks(
             label="Generate Video (may take longer)",
             value=False
         )
-        magic_button = gr.Button("✨Create✨", elem_classes="gr-button")
+        magic_button = gr.Button("✨ Generate Next Item", elem_classes="gr-button")
 
     feed_html = gr.HTML()
     share_html = gr.HTML(label="Share this item:")