Update app.py

app.py

@@ -9,6 +9,10 @@ import json
 import random
 import urllib.parse
 import time
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Initialize the Google Generative AI client with the API key from environment variables
 try:
@@ -79,6 +83,7 @@ def clean_response_text(response_text):
     """
     Clean the API response by removing Markdown code block markers.
     """
+    logging.debug("Cleaning response text")
     cleaned_text = response_text.strip()
     if cleaned_text.startswith("```json"):
         cleaned_text = cleaned_text[len("```json"):].strip()
@@ -90,6 +95,7 @@ def generate_ideas(user_input):
     """
     Generate a diverse set of ideas based on the user's input concept using the LLM.
     """
+    logging.debug(f"Generating ideas for input: {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.
@@ -106,15 +112,17 @@ def generate_ideas(user_input):
                 safety_settings=SAFETY_SETTINGS
             )
         )
+        logging.debug(f"Generate ideas response: {response.text}")
         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) or len(response_json['ideas']) != 5:
             raise ValueError("Invalid JSON format: 'ideas' key missing, not a list, or incorrect length")
+        logging.debug(f"Generated ideas: {response_json['ideas']}")
         return response_json['ideas']
     except Exception as e:
-        print(f"Error generating ideas: {e}")
+        logging.error(f"Error generating ideas: {e}")
         return [
             f"A dramatic {user_input} scene with cinematic lighting",
             f"A close-up of {user_input} in a futuristic setting",
@@ -128,12 +136,15 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
     Generate a single feed item with progress updates.
     Yields progress stage and message for UI updates.
     """
+    logging.debug("Starting generate_item")
     video_base64 = None
     max_total_attempts = 3
-    
+    timeout_seconds = 60  # Timeout for API calls
+
     total_attempts = 0
     while total_attempts < max_total_attempts:
         total_attempts += 1
+        logging.debug(f"Total attempt {total_attempts}")
         
         yield {"stage": "initializing", "message": random.choice(PROGRESS_STAGES["initializing"]), "progress": 10}
         
@@ -143,6 +154,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
         image_prompt = None
         
         for image_attempt in range(max_retries):
+            logging.debug(f"Image attempt {image_attempt + 1}")
             yield {"stage": "creating_image", "message": random.choice(PROGRESS_STAGES["creating_image"]), "progress": 30 + (image_attempt * 10)}
             
             selected_idea = random.choice(ideas)
@@ -155,6 +167,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
             Ensure the response is strictly in JSON format.
             """
             try:
+                start_time = time.time()
                 response = client.models.generate_content(
                     model='gemini-2.0-flash-lite',
                     contents=[prompt],
@@ -163,16 +176,20 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                         safety_settings=SAFETY_SETTINGS
                     )
                 )
+                if time.time() - start_time > timeout_seconds:
+                    raise TimeoutError("Image caption generation timed out")
+                logging.debug(f"Generate content response: {response.text}")
                 cleaned_text = clean_response_text(response.text)
                 response_json = json.loads(cleaned_text)
                 text = response_json['caption']
                 image_prompt = response_json['image_prompt']
             except Exception as e:
-                print(f"Error generating item: {e}")
+                logging.error(f"Error generating item: {e}")
                 text = f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}"
                 image_prompt = f"A vivid scene of {selected_idea} related to {user_input}, in a vibrant pop art style, no text or letters"
 
             try:
+                start_time = time.time()
                 imagen = client.models.generate_images(
                     model='imagen-3.0-generate-002',
                     prompt=image_prompt,
@@ -181,6 +198,8 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                         number_of_images=1
                     )
                 )
+                if time.time() - start_time > timeout_seconds:
+                    raise TimeoutError("Image generation timed out")
                 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))
@@ -190,42 +209,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()
+                    logging.debug("Image generated successfully")
                     break
                 else:
-                    if image_attempt == max_retries - 1:
-                        if total_attempts == max_total_attempts:
-                            image = Image.new('RGB', (360, 640), color='gray')
-                            buffered = BytesIO()
-                            image.save(buffered, format="PNG")
-                            img_str = base64.b64encode(buffered.getvalue()).decode()
-                            yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
-                            return {
-                                'text': text,
-                                'image_base64': img_str,
-                                'video_base64': None,
-                                'ideas': ideas
-                            }
-                        break
-                    yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}
-            except Exception as e:
-                print(f"Error generating image: {e}")
-                if image_attempt == max_retries - 1:
-                    if total_attempts == max_total_attempts:
+                    logging.warning("No images generated")
+                    if image_attempt == max_retries - 1 and total_attempts == max_total_attempts:
                         image = Image.new('RGB', (360, 640), color='gray')
                         buffered = BytesIO()
                         image.save(buffered, format="PNG")
                         img_str = base64.b64encode(buffered.getvalue()).decode()
                         yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+                        logging.debug("Returning with placeholder image")
                         return {
                             'text': text,
                             'image_base64': img_str,
                             'video_base64': None,
                             'ideas': ideas
                         }
-                    break
+                    yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}
+            except Exception as e:
+                logging.error(f"Error generating image: {e}")
+                if image_attempt == max_retries - 1 and total_attempts == max_total_attempts:
+                    image = Image.new('RGB', (360, 640), color='gray')
+                    buffered = BytesIO()
+                    image.save(buffered, format="PNG")
+                    img_str = base64.b64encode(buffered.getvalue()).decode()
+                    yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+                    logging.debug("Returning with placeholder image")
+                    return {
+                        'text': text,
+                        'image_base64': img_str,
+                        'video_base64': None,
+                        'ideas': ideas
+                    }
                 yield {"stage": "retrying_image", "message": random.choice(PROGRESS_STAGES["retrying_image"]), "progress": 40 + (image_attempt * 10)}
 
         if generate_video and generated_image is not None:
+            logging.debug("Attempting video generation")
             yield {"stage": "generating_video", "message": random.choice(PROGRESS_STAGES["generating_video"]), "progress": 70}
             try:
                 video_prompt = f"""
@@ -233,6 +253,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                 Use 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.
                 """
+                start_time = time.time()
                 operation = client.models.generate_videos(
                     model="veo-2.0-generate-001",
                     prompt=video_prompt,
@@ -245,7 +266,9 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     )
                 )
                 while not operation.done:
-                    time.sleep(20)
+                    if time.time() - start_time > timeout_seconds:
+                        raise TimeoutError("Video generation timed out")
+                    time.sleep(5)  # Reduced polling interval
                     operation = client.operations.get(operation)
                 if operation.error:
                     raise ValueError(f"Video generation failed: {operation.error.message}")
@@ -257,6 +280,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                             video_bytes = video_data if isinstance(video_data, bytes) else BytesIO(video_data).getvalue()
                             video_base64 = base64.b64encode(video_bytes).decode()
                             yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+                            logging.debug("Video generated successfully")
                             return {
                                 'text': text,
                                 'image_base64': img_str,
@@ -267,9 +291,11 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                 else:
                     raise ValueError("Video generation operation failed: No generated_videos in response")
             except Exception as e:
-                print(f"Error generating video: {e}")
+                logging.error(f"Error generating video: {e}")
+                logging.debug("Falling back to text-to-video")
                 yield {"stage": "generating_video", "message": random.choice(PROGRESS_STAGES["generating_video"]), "progress": 80}
                 try:
+                    start_time = time.time()
                     operation = client.models.generate_videos(
                         model="veo-2.0-generate-001",
                         prompt=video_prompt,
@@ -281,7 +307,9 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                         )
                     )
                     while not operation.done:
-                        time.sleep(20)
+                        if time.time() - start_time > timeout_seconds:
+                            raise TimeoutError("Text-to-video generation timed out")
+                        time.sleep(5)
                         operation = client.operations.get(operation)
                     if operation.error:
                         raise ValueError(f"Video generation failed: {operation.error.message}")
@@ -293,6 +321,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                                 video_bytes = video_data if isinstance(video_data, bytes) else BytesIO(video_data).getvalue()
                                 video_base64 = base64.b64encode(video_bytes).decode()
                                 yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+                                logging.debug("Text-to-video generated successfully")
                                 return {
                                     'text': text,
                                     'image_base64': img_str,
@@ -303,9 +332,10 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
                     else:
                         raise ValueError("Video generation operation failed: No generated_videos in response")
                 except Exception as e:
-                    print(f"Error generating text-to-video: {e}")
+                    logging.error(f"Error generating text-to-video: {e}")
                     if total_attempts == max_total_attempts:
                         yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+                        logging.debug("Returning without video")
                         return {
                             'text': text,
                             'image_base64': img_str,
@@ -315,6 +345,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
         
         if img_str is not None:
             yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+            logging.debug("Returning with image only")
             return {
                 'text': text,
                 'image_base64': img_str,
@@ -327,6 +358,7 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
     image.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode()
     yield {"stage": "finalizing", "message": random.choice(PROGRESS_STAGES["finalizing"]), "progress": 90}
+    logging.debug("Returning with placeholder image")
     return {
         'text': f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}",
         'image_base64': img_str,
@@ -338,6 +370,7 @@ def start_feed(user_input, generate_video, current_index, feed_items):
     """
     Start or update the feed based on the user input with progress updates.
     """
+    logging.debug("Starting start_feed")
     if not user_input.strip():
         user_input = "trending"
     current_user_input = user_input
@@ -354,6 +387,7 @@ def start_feed(user_input, generate_video, current_index, feed_items):
         
         item_generator = generate_item(user_input, ideas, generate_video=generate_video)
         for progress in 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
@@ -364,10 +398,11 @@ def start_feed(user_input, generate_video, current_index, feed_items):
                 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 Exception as e:
-        print(f"Error in start_feed: {e}")
+        logging.error(f"Error in start_feed: {e}")
         html_content = """
         <div style="
             display: flex;
@@ -394,6 +429,7 @@ def load_next(user_input, generate_video, current_index, feed_items):
     """
     Load the next item in the feed with progress updates.
     """
+    logging.debug("Starting load_next")
     current_user_input = user_input if user_input.strip() else "trending"
     user_input = current_user_input
     is_loading = True
@@ -412,6 +448,7 @@ def load_next(user_input, generate_video, current_index, feed_items):
                 feed_items[current_index]['video_base64'],
                 feed_items[current_index]['text']
             )
+            logging.debug("Loaded existing feed item")
             yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
             return
         else:
@@ -421,6 +458,7 @@ def load_next(user_input, generate_video, current_index, feed_items):
             
             item_generator = generate_item(user_input, ideas, generate_video=generate_video)
             for progress in 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
@@ -435,10 +473,11 @@ 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")
                     yield current_user_input, current_index, feed_items, html_content, share_links, is_loading
                     return
     except Exception as e:
-        print(f"Error in load_next: {e}")
+        logging.error(f"Error in load_next: {e}")
         html_content = """
         <div style="
             display: flex;
@@ -465,6 +504,7 @@ def load_previous(user_input, generate_video, current_index, feed_items):
     """
     Load the previous item in the feed.
     """
+    logging.debug("Loading previous item")
     current_user_input = user_input if user_input.strip() else "trending"
     
     if current_index > 0:
@@ -475,12 +515,14 @@ def load_previous(user_input, generate_video, current_index, feed_items):
         feed_items[current_index]['video_base64'],
         feed_items[current_index]['text']
     )
+    logging.debug("Previous item loaded")
     return current_user_input, current_index, feed_items, html_content, share_links, False
 
 def generate_share_links(image_base64, video_base64, caption):
     """
     Generate share links for social media platforms with download links for image and video.
     """
+    logging.debug("Generating share links")
     image_data_url = f"data:image/png;base64,{image_base64}"
     encoded_caption = urllib.parse.quote(caption)
 
@@ -623,6 +665,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
     """
     Generate an HTML string to display the current feed item or a loading screen with a progress bar.
     """
+    logging.debug(f"Generating HTML, is_loading: {is_loading}, progress_stage: {progress_stage}")
     if is_loading:
         progress_percent = max(0, min(100, progress_percent))
         return f"""
@@ -692,6 +735,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
         """
 
     if not feed_items or current_index >= len(feed_items):
+        logging.debug("No feed items to display")
         return """
         <div style="
             display: flex;
@@ -788,6 +832,7 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, user_inpu
         }}
     </script>
     """
+    logging.debug("Feed item HTML generated")
     return html_str
 
 # Define the Gradio interface
@@ -821,7 +866,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("✨Create✨", elem_classes="gr-button")
 
     feed_html = gr.HTML()
     share_html = gr.HTML(label="Share this item:")