Update app.py

app.py

@@ -141,188 +141,219 @@ def generate_item(user_input, ideas, generate_video=False, max_retries=3):
         user_input (str): The user's input concept or idea.
         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 for both image and video generation.
+        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).
     """
     video_base64 = None
+    max_total_attempts = 3  # Maximum total attempts for combined image and video generation cycles
     
-    # Retry loop for image generation
-    for attempt in range(max_retries):
-        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.
-        Return a JSON object with two keys:
-        - 'caption': A short, viral TikTok-style caption with hashtags that reflects "{user_input}".
-        - 'image_prompt': A detailed image prompt for generating a high-quality visual scene, ensuring the theme of "{user_input}" is central.
-        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": "Blindfolded Rubik's Cube MAGIC! 🤯 #rubiks", "image_prompt": "A close-up view of a person solving a Rubik's Cube blindfolded, in a dramatic 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 (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"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"
-
-        # Attempt to generate the image
-        try:
-            imagen = client.models.generate_images(
-                model='imagen-3.0-generate-002',
-                prompt=image_prompt,
-                config=types.GenerateImagesConfig(
-                    aspect_ratio="9:16",
-                    number_of_images=1
+    total_attempts = 0
+    while total_attempts < max_total_attempts:
+        total_attempts += 1
+        
+        # Step 1: Generate an image (retry up to max_retries times)
+        for image_attempt in range(max_retries):
+            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.
+            Return a JSON object with two keys:
+            - 'caption': A short, viral TikTok-style caption with hashtags that reflects "{user_input}".
+            - 'image_prompt': A detailed image prompt for generating a high-quality visual scene, ensuring the theme of "{user_input}" is central.
+            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": "Blindfolded Rubik's Cube MAGIC! 🤯 #rubiks", "image_prompt": "A close-up view of a person solving a Rubik's Cube blindfolded, in a dramatic 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 (image attempt {image_attempt + 1}, total attempt {total_attempts}): {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 (image attempt {image_attempt + 1}, total attempt {total_attempts}): {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"
+
+            # Attempt to generate the image
+            try:
+                imagen = client.models.generate_images(
+                    model='imagen-3.0-generate-002',
+                    prompt=image_prompt,
+                    config=types.GenerateImagesConfig(
+                        aspect_ratio="9:16",
+                        number_of_images=1
+                    )
                 )
-            )
-            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)
-                # Convert image to base64
-                buffered = BytesIO()
-                image.save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode()
-
-                # Generate video if enabled (with retries)
-                if generate_video:
-                    for video_attempt in range(max_retries):
-                        try:
-                            # Base video prompt
-                            video_prompt_base = 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, 
-                            using shallow focus on the main subject to emphasize details, in a realistic style with cinematic lighting.
-                            """
-                            # Modify the prompt slightly for each retry
-                            if video_attempt == 0:
-                                video_prompt = video_prompt_base
-                            elif video_attempt == 1:
-                                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 focusing on the subject, 
-                                with soft lighting and a realistic style.
-                                """
-                            else:
-                                video_prompt = f"""
-                                The user concept is "{user_input}". Based on this and a simplified scene: {image_prompt}, create a video.
-                                Use a static close-up shot of the subject in a realistic style.
-                                """
-
-                            print(f"Attempting video generation (attempt {video_attempt + 1}): {video_prompt}")
-                            operation = client.models.generate_videos(
-                                model="veo-2.0-generate-001",
-                                prompt=video_prompt,
-                                image=generated_image.image,
-                                config=types.GenerateVideosConfig(
-                                    aspect_ratio="9:16",
-                                    number_of_videos=1,
-                                    duration_seconds=8,
-                                    negative_prompt="blurry, low quality, text, letters"
-                                )
-                            )
-                            # Wait for video to generate
-                            while not operation.done:
-                                time.sleep(20)
-                                operation = client.operations.get(operation)
-                            
-                            # Enhanced error handling for video generation response
-                            if operation.error:
-                                raise ValueError(f"Video generation operation failed with error: {operation.error.message}")
-                            if operation.response is None:
-                                raise ValueError("Video generation operation failed: No response")
-                            if not hasattr(operation.response, 'generated_videos') or operation.response.generated_videos is None:
-                                raise ValueError("Video generation operation failed: No generated_videos in response")
-                            
-                            # Process the single generated video
-                            if len(operation.response.generated_videos) > 0:
-                                video = operation.response.generated_videos[0]
-                                if video is None or not hasattr(video, 'video'):
-                                    raise ValueError("Video is invalid or missing video data")
-                                fname = 'with_image_input.mp4'
-                                print(f"Generated video: {fname}")
-                                # Download the video and get the raw bytes
-                                video_data = client.files.download(file=video.video)
-                                # Ensure video_data is in bytes
-                                if isinstance(video_data, bytes):
-                                    video_bytes = video_data
-                                else:
-                                    # If video_data is a file-like object, read the bytes
-                                    video_buffer = BytesIO()
-                                    for chunk in video_data:
-                                        video_buffer.write(chunk)
-                                    video_bytes = video_buffer.getvalue()
-                                # Encode the video bytes as base64
-                                video_base64 = base64.b64encode(video_bytes).decode()
-                                break  # Success, exit the retry loop
-                            else:
-                                raise ValueError("No video was generated")
-                        except Exception as e:
-                            print(f"Error generating video (attempt {video_attempt + 1}): {e}")
-                            if video_attempt == max_retries - 1:
-                                print("Max retries reached for video generation. Proceeding without video.")
-                                video_base64 = None
-                            else:
-                                continue  # Retry with a modified prompt
-
-                return {
-                    'text': text,
-                    'image_base64': img_str,
-                    'video_base64': video_base64,
-                    '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')
+                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)
+                    # 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,
-                        'video_base64': None,
-                        'ideas': ideas
-                    }
-                # Retry with new ideas
-                ideas = generate_ideas(user_input)
-                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,
-                    'video_base64': None,
-                    'ideas': ideas
-                }
-            # Retry with new ideas
-            ideas = generate_ideas(user_input)
-            continue
+                    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, use a gray placeholder if max total attempts not reached
+                        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()
+                            return {
+                                'text': text,
+                                'image_base64': img_str,
+                                'video_base64': None,
+                                'ideas': ideas
+                            }
+                        # Otherwise, select a new idea and retry image generation in the next cycle
+                        continue
+            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
+                    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()
+                        return {
+                            'text': text,
+                            'image_base64': img_str,
+                            'video_base64': None,
+                            'ideas': ideas
+                        }
+                    # Otherwise, select a new idea and retry image generation in the next cycle
+                    continue
+
+        # Step 2: Generate video if enabled (with retries using the same image)
+        if generate_video:
+            max_video_retries_per_image = 1  # Try video generation twice per image
+            for video_attempt in range(max_video_retries_per_image):
+                try:
+                    # Base video prompt
+                    video_prompt_base = 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, 
+                    using shallow focus on the main subject to emphasize details, in a realistic style with cinematic lighting.
+                    """
+                    # Modify the prompt slightly for each retry
+                    if video_attempt == 0:
+                        video_prompt = video_prompt_base
+                    else:
+                        video_prompt = f"""
+                        The user concept is "{user_input}". Based on this and a simplified scene: {image_prompt}, create a video.
+                        Use a static close-up shot of the subject in a realistic style.
+                        """
+
+                    print(f"Attempting video generation (video attempt {video_attempt + 1}, total attempt {total_attempts}): {video_prompt}")
+                    operation = client.models.generate_videos(
+                        model="veo-2.0-generate-001",
+                        prompt=video_prompt,
+                        image=generated_image.image,
+                        config=types.GenerateVideosConfig(
+                            aspect_ratio="9:16",
+                            number_of_videos=1,
+                            duration_seconds=8,
+                            negative_prompt="blurry, low quality, text, letters"
+                        )
+                    )
+                    # Wait for video to generate
+                    while not operation.done:
+                        time.sleep(20)
+                        operation = client.operations.get(operation)
+                    
+                    # Enhanced error handling for video generation response
+                    if operation.error:
+                        raise ValueError(f"Video generation operation failed with error: {operation.error.message}")
+                    if operation.response is None:
+                        raise ValueError("Video generation operation failed: No response")
+                    if not hasattr(operation.response, 'generated_videos') or operation.response.generated_videos is None:
+                        raise ValueError("Video generation operation failed: No generated_videos in response")
+                    
+                    # Process the single generated video
+                    if len(operation.response.generated_videos) > 0:
+                        video = operation.response.generated_videos[0]
+                        if video is None or not hasattr(video, 'video'):
+                            raise ValueError("Video is invalid or missing video data")
+                        fname = 'with_image_input.mp4'
+                        print(f"Generated video: {fname}")
+                        # Download the video and get the raw bytes
+                        video_data = client.files.download(file=video.video)
+                        # Ensure video_data is in bytes
+                        if isinstance(video_data, bytes):
+                            video_bytes = video_data
+                        else:
+                            # If video_data is a file-like object, read the bytes
+                            video_buffer = BytesIO()
+                            for chunk in video_data:
+                                video_buffer.write(chunk)
+                            video_bytes = video_buffer.getvalue()
+                        # Encode the video bytes as base64
+                        video_base64 = base64.b64encode(video_bytes).decode()
+                        # Successfully generated video, return the result
+                        return {
+                            'text': text,
+                            'image_base64': img_str,
+                            'video_base64': video_base64,
+                            'ideas': ideas
+                        }
+                    else:
+                        raise ValueError("No video was generated")
+                except Exception as e:
+                    print(f"Error generating video (video attempt {video_attempt + 1}, total attempt {total_attempts}): {e}")
+                    if video_attempt == max_video_retries_per_image - 1:
+                        if total_attempts == max_total_attempts:
+                            print("Max total attempts reached. Proceeding without video.")
+                            video_base64 = None
+                            return {
+                                'text': text,
+                                'image_base64': img_str,
+                                'video_base64': video_base64,
+                                'ideas': ideas
+                            }
+                        # Video generation failed with this image, break to outer loop to try a new image
+                        print(f"Video generation failed after {max_video_retries_per_image} attempts with this image. Selecting a new idea and generating a new image.")
+                        break
+                    continue  # Retry video generation with the same image but a modified prompt
+        
+        # If video generation is not enabled or succeeded, return the result
+        return {
+            'text': text,
+            'image_base64': img_str,
+            'video_base64': video_base64,
+            'ideas': ideas
+        }
+
+    # If max total attempts reached without success, use a gray placeholder image
+    print("Max total attempts reached without successful image generation. Using 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': f"Amazing {user_input}! 🔥 #{user_input.replace(' ', '')}",
+        'image_base64': img_str,
+        'video_base64': None,
+        'ideas': ideas
+    }
 
 def start_feed(user_input, generate_video, current_index, feed_items):
     """