Update app.py

app.py

@@ -25,7 +25,6 @@ def clean_response_text(response_text):
     Returns:
         str: The cleaned response text.
     """
-    # Remove ```json and ``` markers
     cleaned_text = response_text.strip()
     if cleaned_text.startswith("```json"):
         cleaned_text = cleaned_text[len("```json"):].strip()
@@ -59,7 +58,6 @@ def generate_ideas(tag):
         print(f"Raw response for ideas: {response.text}")  # Debugging
         if not response.text or response.text.isspace():
             raise ValueError("Empty response from API")
-        # Clean the response text
         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):
@@ -75,72 +73,92 @@ def generate_ideas(tag):
             f"An action-packed {tag} scene with dynamic colors"
         ]
 
-def generate_item(tag, ideas):
+def generate_item(tag, ideas, max_retries=3):
     """
-    Generate a single feed item using one of the ideas.
+    Generate a single feed item using one of the ideas, retrying if image generation fails.
     
     Args:
         tag (str): The tag to base the content on.
         ideas (list): List of ideas to choose from.
+        max_retries (int): Maximum number of retries if image generation fails.
     
     Returns:
-        dict: A dictionary with 'text' (str) and 'image_base64' (str).
+        dict: A dictionary with 'text' (str), 'image_base64' (str), and 'ideas' (list).
     """
-    selected_idea = random.choice(ideas)
-    prompt = f"""
-    Based on the idea "{selected_idea}", create content for a TikTok video about {tag}.
-    Return a JSON object with two keys:
-    - 'caption': A short, viral TikTok-style caption with hashtags.
-    - 'image_prompt': A detailed image prompt for generating a high-quality visual scene.
-    The image prompt should describe the scene vividly, specify a perspective and style, and ensure no text or letters are included.
-    Ensure the response is strictly in JSON format.
-    Example: {{"caption": "Neon vibes only! 🌌 #tech", "image_prompt": "A close-up view of a neon-lit gaming setup with RGB lights flashing, in a futuristic style, no text or letters"}}
-    """
-    try:
-        response = client.models.generate_content(
-            model='gemini-2.0-flash',
-            contents=[prompt],
-            config=types.GenerateContentConfig(temperature=1.2)
-        )
-        print(f"Raw response for item: {response.text}")  # Debugging
-        if not response.text or response.text.isspace():
-            raise ValueError("Empty response from API")
-        # Clean the response text
-        cleaned_text = clean_response_text(response.text)
-        response_json = json.loads(cleaned_text)
-        if 'caption' not in response_json or 'image_prompt' not in response_json:
-            raise ValueError("Invalid JSON format: 'caption' or 'image_prompt' key missing")
-        text = response_json['caption']
-        image_prompt = response_json['image_prompt']
-    except Exception as e:
-        print(f"Error generating item: {e}")
-        text = f"Obsessed with {tag}! 🔥 #{tag}"
-        image_prompt = f"A vivid scene of {selected_idea}, in a vibrant pop art style, no text or letters"
-
-    try:
-        image_response = client.models.generate_images(
-            model='imagen-3.0-generate-002',
-            prompt=image_prompt,
-            config=types.GenerateImagesConfig(
-                number_of_images=1,
-                aspect_ratio="9:16",
-                person_generation="DONT_ALLOW"
+    for attempt in range(max_retries):
+        selected_idea = random.choice(ideas)
+        prompt = f"""
+        Based on the idea "{selected_idea}", create content for a TikTok video about {tag}.
+        Return a JSON object with two keys:
+        - 'caption': A short, viral TikTok-style caption with hashtags.
+        - 'image_prompt': A detailed image prompt for generating a high-quality visual scene.
+        The image prompt should describe the scene vividly, specify a perspective and style, and ensure no text or letters are included.
+        Ensure the response is strictly in JSON format.
+        Example: {{"caption": "Neon vibes only! 🌌 #tech", "image_prompt": "A close-up view of a neon-lit gaming setup with RGB lights flashing, in a futuristic style, no text or letters"}}
+        """
+        try:
+            response = client.models.generate_content(
+                model='gemini-2.0-flash',
+                contents=[prompt],
+                config=types.GenerateContentConfig(temperature=1.2)
             )
-        )
-        if image_response.generated_images and len(image_response.generated_images) > 0:
-            generated_image = image_response.generated_images[0]
-            image = Image.open(BytesIO(generated_image.image.image_bytes))
-        else:
-            image = Image.new('RGB', (360, 640), color='gray')
-    except Exception as e:
-        print(f"Error generating image: {e}")
-        image = Image.new('RGB', (360, 640), color='gray')
-
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
+            print(f"Raw response for item (attempt {attempt + 1}): {response.text}")  # Debugging
+            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 'caption' not in response_json or 'image_prompt' not in response_json:
+                raise ValueError("Invalid JSON format: 'caption' or 'image_prompt' key missing")
+            text = response_json['caption']
+            image_prompt = response_json['image_prompt']
+        except Exception as e:
+            print(f"Error generating item (attempt {attempt + 1}): {e}")
+            text = f"Obsessed with {tag}! 🔥 #{tag}"
+            image_prompt = f"A vivid scene of {selected_idea}, in a vibrant pop art style, no text or letters"
 
-    return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+        # Attempt to generate the image
+        try:
+            image_response = client.models.generate_images(
+                model='imagen-3.0-generate-002',
+                prompt=image_prompt,
+                config=types.GenerateImagesConfig(
+                    number_of_images=1,
+                    aspect_ratio="9:16",
+                    person_generation="DONT_ALLOW"
+                )
+            )
+            if image_response.generated_images and len(image_response.generated_images) > 0:
+                generated_image = image_response.generated_images[0]
+                image = Image.open(BytesIO(generated_image.image.image_bytes))
+                # Successfully generated an image, proceed
+                buffered = BytesIO()
+                image.save(buffered, format="PNG")
+                img_str = base64.b64encode(buffered.getvalue()).decode()
+                return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+            else:
+                print(f"Image generation failed (attempt {attempt + 1}): No images returned")
+                if attempt == max_retries - 1:
+                    # Last attempt, 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()
+                    return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+                # Retry with new ideas
+                ideas = generate_ideas(tag)
+                continue
+        except Exception as e:
+            print(f"Error generating image (attempt {attempt + 1}): {e}")
+            if attempt == max_retries - 1:
+                # Last attempt, 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()
+                return {'text': text, 'image_base64': img_str, 'ideas': ideas}
+            # Retry with new ideas
+            ideas = generate_ideas(tag)
+            continue
 
 def start_feed(tag, current_index, feed_items):
     """
@@ -265,7 +283,7 @@ def load_previous(tag, current_index, feed_items):
 
 def generate_html(feed_items, scroll_to_latest=False, current_index=0, tag="", is_loading=False):
     """
-    Generate an HTML string to display the current feed item with click navigation.
+    Generate an HTML string to display the current feed item with click navigation and a next button.
     
     Args:
         feed_items (list): List of dictionaries containing 'text' and 'image_base64'.
@@ -411,6 +429,20 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, tag="", i
             ">
                 {text}
             </div>
+            <button id="next-button" style="
+                position: absolute;
+                bottom: 80px;
+                right: 20px;
+                background-color: #ff2d55;
+                color: white;
+                border: none;
+                border-radius: 50%;
+                width: 50px;
+                height: 50px;
+                font-size: 24px;
+                cursor: pointer;
+                z-index: 3;
+            " onclick="document.getElementById('next-button').click()">✨</button>
         </div>
     </div>
     <script>
@@ -421,8 +453,6 @@ def generate_html(feed_items, scroll_to_latest=False, current_index=0, tag="", i
             const width = rect.width;
             if (clickX > width * 0.75) {{
                 document.getElementById('previous-button').click();
-            }} else {{
-                document.getElementById('next-button').click();
             }}
         }}
     </script>
@@ -508,5 +538,5 @@ with gr.Blocks(
         outputs=[current_tag, current_index, feed_items, feed_html, is_loading]
     )
 
-# Launch the app (removed share=True since we're on Hugging Face Spaces)
+# Launch the app
 demo.launch()