Update image_gen.py

image_gen.py

@@ -1,8 +1,6 @@
 # -----------------------
 # Image Generation
 # -----------------------
-
-
 import os
 import re
 import time
@@ -87,15 +85,20 @@ def standardize_and_validate_image(file_path):
 
 def generate_image(prompt_text, style, model="hf"):
     """
-    Generate an image from a text prompt using either Hugging Face's, Pollinations Turbo's,
-    or Google's Gemini API.
+    Generate an image from a text prompt using one of the following:
+      - Hugging Face's FLUX.1-dev
+      - Pollinations Turbo
+      - Google's Gemini
+      - Pexels API (for a real photo instead of AI)
+    
     Args:
-        prompt_text (str): The text prompt for image generation.
-        style (str or None): The style of the image (used for HF and Gemini models).
-        model (str): Which model to use ("hf" for Hugging Face, "pollinations_turbo" for Pollinations Turbo,
-                     or "gemini" for Google's Gemini).
+        prompt_text (str): The text prompt for image generation or search.
+        style (str or None): The style of the image (used for HF and Gemini models, ignored for Pexels).
+        model (str): Which model to use 
+                     ("hf", "pollinations_turbo", "gemini", or "pexels").
+    
     Returns:
-        tuple: A tuple containing the generated PIL.Image and a Base64 string of the image.
+        tuple: (PIL.Image, base64_string) or (None, None) on error.
     """
     try:
         if model == "pollinations_turbo":
@@ -108,60 +111,104 @@ def generate_image(prompt_text, style, model="hf"):
                 print(f"Error from image generation API: {response.status_code}")
                 return None, None
             image_bytes = response.content
-            
+
         elif model == "gemini":
             # For Google's Gemini model
             try:
-                
-                # Get API key from environment variable
                 g_api_key = os.getenv("GEMINI")
                 if not g_api_key:
                     logger.error("GEMINI_API_KEY not found in environment variables")
                     print("Google Gemini API key is missing. Please set the GEMINI_API_KEY environment variable.")
                     return None, None
-                    
+
                 # Initialize Gemini client
                 client = genai.Client(api_key=g_api_key)
-                
+
                 # Enhance prompt with style
                 enhanced_prompt = f"image of {prompt_text} in {style} style, high quality, detailed illustration"
-                
+
                 # Generate content
                 response = client.models.generate_content(
                     model="models/gemini-2.0-flash-exp",
                     contents=enhanced_prompt,
                     config=types.GenerateContentConfig(response_modalities=['Text', 'Image'])
                 )
-                
+
                 # Extract image from response
                 for part in response.candidates[0].content.parts:
                     if part.inline_data is not None:
                         image = Image.open(BytesIO(part.inline_data.data))
-                        
                         # Convert to base64 string
                         buffered = io.BytesIO()
                         image.save(buffered, format="JPEG")
                         img_str = base64.b64encode(buffered.getvalue()).decode()
-                        
                         return image, img_str
-                
+
                 # If no image was found in the response
                 logger.error("No image was found in the Gemini API response")
                 print("Gemini API didn't return an image")
                 return None, None
-                
+
             except ImportError:
                 logger.error("Google Gemini libraries not installed")
-                st.error("Google Gemini libraries not installed. Install with 'pip install google-genai'")
+                print("Google Gemini libraries not installed. Install with 'pip install google-genai'")
                 return None, None
-                
+
             except Exception as e:
                 logger.error(f"Gemini API error: {str(e)}")
                 print(f"Error from Gemini image generation: {str(e)}")
                 return None, None
-                
-        else:  # Default to Hugging Face model
-            # For Hugging Face model, include style details in the prompt
+
+        elif model == "pexels":
+            # ---------- NEW BRANCH FOR PEXELS -----------
+            pexels_api_key = os.getenv("PEXELS_API_KEY")
+            if not pexels_api_key:
+                logger.error("PEXELS_API_KEY not found in environment variables")
+                print("Pexels API key is missing. Please set the PEXELS_API_KEY environment variable.")
+                return None, None
+
+            # Call Pexels search endpoint
+            # e.g. GET https://api.pexels.com/v1/search?query={prompt_text}&per_page=1
+            
+            search_url = "https://api.pexels.com/v1/search"
+            headers_pexels = {
+                "Authorization": pexels_api_key
+            }
+            params = {
+                "query": prompt_text,
+                "per_page": 1
+            }
+            response = requests.get(search_url, headers=headers_pexels, params=params)
+            if response.status_code != 200:
+                logger.error(f"Pexels API error: {response.status_code}, {response.text}")
+                print(f"Error from Pexels API: {response.status_code}")
+                return None, None
+
+            data = response.json()
+            photos = data.get("photos", [])
+            if not photos:
+                logger.error("No photos found for the given prompt on Pexels")
+                print("No photos found on Pexels for this prompt.")
+                return None, None
+
+            # Take the first photo
+            photo = photos[0]
+            # We can pick "src" => "original" or "large2x", etc.
+            image_url = photo["src"].get("large2x") or photo["src"].get("original")
+            if not image_url:
+                logger.error("No suitable image URL found in Pexels photo object")
+                return None, None
+
+            # Download the image
+            img_resp = requests.get(image_url)
+            if img_resp.status_code != 200:
+                logger.error(f"Failed to download Pexels image from {image_url}")
+                return None, None
+
+            image_bytes = img_resp.content
+
+        else:
+            # Default to Hugging Face model
             enhanced_prompt = f"{prompt_text} in {style} style, high quality, detailed illustration"
             model_id = "black-forest-labs/FLUX.1-dev"
             api_url = f"https://api-inference.huggingface.co/models/{model_id}"
@@ -173,43 +220,21 @@ def generate_image(prompt_text, style, model="hf"):
                 return None, None
             image_bytes = response.content
 
-        # For HF and Pollinations models that return image bytes
+        # For HF, Pollinations, or Pexels that return image bytes
         if model != "gemini":
             image = Image.open(io.BytesIO(image_bytes))
             buffered = io.BytesIO()
             image.save(buffered, format="JPEG")
             img_str = base64.b64encode(buffered.getvalue()).decode()
             return image, img_str
-            
+
     except Exception as e:
         print(f"Error generating image: {e}")
         logger.error(f"Image generation error: {str(e)}")
-        
+
     # Return a placeholder image in case of failure
     return Image.new('RGB', (1024, 1024), color=(200,200,200)), None
 
-def generate_image_with_retry(prompt_text, style, model="hf", max_retries=3):
-    """
-    Attempt to generate an image using generate_image, retrying up to max_retries if needed.
-    Args:
-        prompt_text (str): The text prompt for image generation.
-        style (str or None): The style of the image (ignored for Pollinations Turbo).
-        model (str): Which model to use ("hf" or "pollinations_turbo").
-        max_retries (int): Maximum number of retries.
-    Returns:
-        tuple: The generated image and its Base64 string.
-    """
-    for attempt in range(max_retries):
-        try:
-            if attempt > 0:
-                time.sleep(2 ** attempt)
-            return generate_image(prompt_text, style, model=model)
-        except Exception as e:
-            logger.error(f"Attempt {attempt+1} failed: {e}")
-            if attempt == max_retries - 1:
-                raise
-    return None, None
-
 # edit image function
 def edit_section_image(image_url: str, gemini_prompt: str):
     """