Update image_gen.py

image_gen.py

@@ -0,0 +1,209 @@
+# -----------------------
+# Image Generation
+# -----------------------
+
+
+import os
+import re
+import time
+import tempfile
+import requests
+import json
+from google import genai
+from google.genai import types
+import io
+import base64
+import numpy as np
+import cv2
+import logging
+import uuid
+import subprocess
+from pathlib import Path
+import urllib.parse
+import pandas as pd
+import plotly.graph_objects as go
+import matplotlib.pyplot as plt
+import base64
+import os
+import uuid
+import matplotlib
+import matplotlib.pyplot as plt
+from io import BytesIO
+import dataframe_image as dfi
+import uuid
+from PIL import ImageFont, ImageDraw, Image
+import seaborn as sns
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+
+def is_valid_png(file_path):
+    """Check if the PNG file at `file_path` is valid."""
+    try:
+        with open(file_path, "rb") as f:
+            # Read the first 8 bytes to check the PNG signature
+            header = f.read(8)
+            if header != b'\x89PNG\r\n\x1a\n':
+                return False
+            
+            # Attempt to open and verify the entire image
+            with Image.open(file_path) as img:
+                img.verify()  # Verify the file integrity
+            return True
+    except Exception as e:
+        print(f"Invalid PNG file at {file_path}: {e}")
+        return False
+
+
+def standardize_and_validate_image(file_path):
+    """Validate, standardize, and overwrite the image at `file_path`."""
+    try:
+        # Verify basic integrity
+        with Image.open(file_path) as img:
+            img.verify()
+        
+        # Reopen and convert to RGB
+        with Image.open(file_path) as img:
+            img = img.convert("RGB")  # Remove alpha channel if present
+            
+            # Save to a temporary BytesIO buffer first
+            buffer = io.BytesIO()
+            img.save(buffer, format="PNG")
+            buffer.seek(0)
+            
+            # Write the buffer to the file
+            with open(file_path, "wb") as f:
+                f.write(buffer.getvalue())
+            
+        return True
+    except Exception as e:
+        print(f"Failed to standardize/validate {file_path}: {e}")
+        return False
+
+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.
+    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).
+    Returns:
+        tuple: A tuple containing the generated PIL.Image and a Base64 string of the image.
+    """
+    try:
+        if model == "pollinations_turbo":
+            # URL-encode the prompt and add the query parameter to specify the model as "turbo"
+            prompt_encoded = urllib.parse.quote(prompt_text)
+            api_url = f"https://image.pollinations.ai/prompt/{prompt_encoded}?model=turbo"
+            response = requests.get(api_url)
+            if response.status_code != 200:
+                logger.error(f"Pollinations API error: {response.status_code}, {response.text}")
+                st.error(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")
+                    st.error("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")
+                st.error("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'")
+                return None, None
+                
+            except Exception as e:
+                logger.error(f"Gemini API error: {str(e)}")
+                st.error(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
+            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}"
+            payload = {"inputs": enhanced_prompt}
+            response = requests.post(api_url, headers=headers, json=payload)
+            if response.status_code != 200:
+                logger.error(f"Hugging Face API error: {response.status_code}, {response.text}")
+                st.error(f"Error from image generation API: {response.status_code}")
+                return None, None
+            image_bytes = response.content
+
+        # For HF and Pollinations models 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:
+        st.error(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