Update app.py

app.py

@@ -8,6 +8,7 @@ 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:
@@ -74,18 +75,21 @@ def generate_ideas(tag):
             f"An action-packed {tag} scene with dynamic colors"
         ]
 
-def generate_item(tag, ideas, max_retries=3):
+def generate_item(tag, ideas, generate_video=False, max_retries=3):
     """
-    Generate a single feed item using one of the ideas, retrying if image generation fails.
+    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), and 'ideas' (list).
+        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"""
@@ -119,26 +123,61 @@ def generate_item(tag, ideas, max_retries=3):
 
         # Attempt to generate the image
         try:
-            # Pass parameters directly or adjust based on the correct API schema
-            image_response = client.models.generate_images(
+            imagen = client.models.generate_images(
                 model='imagen-3.0-generate-002',
                 prompt=image_prompt,
-                number_of_images=1,  # Pass directly instead of in config
-                aspect_ratio="9:16",
-                person_generation="DONT_ALLOW"
+                config=types.GenerateImagesConfig(
+                    aspect_ratio="9:16",
+                    number_of_images=1
+                )
             )
-            if image_response.generated_images and len(image_response.generated_images) > 0:
-                generated_image = image_response.generated_images[0]
+            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)
-                # Successfully generated an image, proceed
+                # Convert image to base64
                 buffered = BytesIO()
                 image.save(buffered, format="PNG")
                 img_str = base64.b64encode(buffered.getvalue()).decode()
-                return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+
+                # Generate video if enabled
+                if generate_video:
+                    try:
+                        operation = client.models.generate_videos(
+                            model="veo-2.0-generate-001",
+                            prompt=image_prompt,
+                            image=generated_image.image,
+                            config=types.GenerateVideosConfig(
+                                aspect_ratio="9:16",
+                                number_of_videos=2
+                            )
+                        )
+                        # 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:
@@ -147,7 +186,12 @@ def generate_item(tag, ideas, max_retries=3):
                     buffered = BytesIO()
                     image.save(buffered, format="PNG")
                     img_str = base64.b64encode(buffered.getvalue()).decode()
-                    return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+                    return {
+                        'text': text,
+                        'image_base64': img_str,
+                        'video_base64_list': [],
+                        'ideas': ideas
+                    }
                 # Retry with new ideas
                 ideas = generate_ideas(tag)
                 continue
@@ -159,17 +203,23 @@ def generate_item(tag, ideas, max_retries=3):
                 buffered = BytesIO()
                 image.save(buffered, format="PNG")
                 img_str = base64.b64encode(buffered.getvalue()).decode()
-                return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+                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, current_index, feed_items):
+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.
     
@@ -186,10 +236,14 @@ def start_feed(tag, current_index, feed_items):
 
     try:
         ideas = generate_ideas(tag)
-        item = generate_item(tag, ideas)
+        item = generate_item(tag, ideas, generate_video=generate_video)
         feed_items = [item]
         current_index = 0
-        share_links = generate_share_links(item['image_base64'], item['text'])
+        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 = []
@@ -220,12 +274,13 @@ def start_feed(tag, current_index, feed_items):
     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, current_index, feed_items):
+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.
     
@@ -241,10 +296,14 @@ def load_next(tag, current_index, feed_items):
             current_index += 1
         else:
             ideas = feed_items[-1]['ideas'] if feed_items else generate_ideas(tag)
-            new_item = generate_item(tag, ideas)
+            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]['text'])
+        share_links = generate_share_links(
+            feed_items[current_index]['image_base64'],
+            IRfeed_items[current_index]['video_base64_list'],
+            feed_items[current_index]['text']
+        )
     except Exception as e:
         print(f"Error in load_next: {e}")
         html_content = """
@@ -272,12 +331,13 @@ def load_next(tag, current_index, feed_items):
     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, current_index, feed_items):
+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.
     
@@ -287,16 +347,21 @@ def load_previous(tag, current_index, feed_items):
     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]['text'])
+    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, caption):
+def generate_share_links(image_base64, video_base64_list, caption):
     """
-    Generate share links for social media platforms with a download link.
+    Generate share links for social media platforms with download links for image and video.
     
     Args:
         image_base64 (str): The base64-encoded image data.
-        caption (str): The caption to share with the image.
+        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.
@@ -304,8 +369,33 @@ def generate_share_links(image_base64, caption):
     image_data_url = f"data:image/png;base64,{image_base64}"
     encoded_caption = urllib.parse.quote(caption)
 
-    # Generate share links with instructions to download the image
-    share_links = """
+    # 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;
@@ -315,14 +405,7 @@ def generate_share_links(image_base64, caption):
         color: white;
         font-family: Arial, sans-serif;
     ">
-        <p style="text-align: center;">To share, download the image and upload it to your platform:</p>
-        <a href="{image_url}" download="feed_item.png" style="
-            background-color: #4CAF50;
-            color: white;
-            padding: 5px 10px;
-            border-radius: 5px;
-            text-decoration: none;
-        ">Download Image</a>
+        {download_links}
         <div style="
             display: flex;
             flex-wrap: wrap;
@@ -367,8 +450,7 @@ def generate_share_links(image_base64, caption):
             ">Share on Pinterest</a>
         </div>
     </div>
-    """.format(image_url=image_data_url, caption=encoded_caption)
-
+    """
     return share_links
 
 def generate_html(feed_items, scroll_to_latest=False, current_index=0, tag="", is_loading=False):
@@ -542,7 +624,7 @@ 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; }
+        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; }
@@ -566,6 +648,10 @@ with gr.Blocks(
             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
@@ -576,14 +662,14 @@ with gr.Blocks(
     # Handle Enter keypress in the concept input
     tag_input.submit(
         fn=start_feed,
-        inputs=[tag_input, current_index, feed_items],
+        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, current_index, feed_items],
+        inputs=[current_tag, generate_video_checkbox, current_index, feed_items],
         outputs=[current_tag, current_index, feed_items, feed_html, share_html, is_loading]
     )
 
@@ -593,7 +679,7 @@ with gr.Blocks(
     # Handle click to go to previous item
     previous_button.click(
         fn=load_previous,
-        inputs=[current_tag, current_index, feed_items],
+        inputs=[current_tag, generate_video_checkbox, current_index, feed_items],
         outputs=[current_tag, current_index, feed_items, feed_html, share_html, is_loading]
     )