Enhance image generation and fetching with error handling and retry logic

image.py

@@ -26,12 +26,16 @@ def generate_image(title, summary):
         image_api_endpoint = "image.pollinations.ai"
         params = f"negative={quote(str(negative))}&seed=2342342340&width=1024&height=576&nologo=True&private=False&model=turbo&enhance=True"
         data_header = {"Content-Type": "application/json"}
-        prompt = quote(f"{title.strip()}: {extracted_summary.strip()}")
+        prompt = quote(f"[[(({title.strip()}))]]: {extracted_summary.strip()}")
         image_url = f"https://{image_api_endpoint}/prompt/{prompt}?{params}"
-        img_data = requests.get(image_url, headers=data_header, timeout=60).content
+        response = requests.get(image_url, headers=data_header)
+        if response.status_code != 200:
+            print(f"Error generating image: {response.text}")
+            return None
+        img_data = response.content
         if img_data:
-            base64_encoded = base64.b64encode(img_data).decode("utf-8")
-            return f"data:image/png;base64,{base64_encoded}"
+            data = base64.b64encode(img_data).decode("utf-8")
+            return f"data:image/png;base64,{data}"
         return None
     except Exception as e:
         print(f"An error occurred during image generation: {e}")
@@ -69,12 +73,22 @@ def fetch_image(title, summary, api_key):
     summary = r"{}".format(summary)
     image_url = "https://i.ibb.co/TBJqggw/Image-Not-Found.jpg"
     try:
-        data_uri = generate_image(title, summary)
+        data_uri = None
+        i = 1
+        while not data_uri and i <= 3:
+            print(f"Attempt {i} to fetch image")
+            data_uri = generate_image(title, summary)
+            i += 1
         if data_uri:
             base64_image = fix_base64_padding(data_uri.split(",")[1])
-            image_data = base64.b64decode(base64_image)
-            if verify_image(image_data):
-                image_url = upload_image(data_uri, api_key)
+            image_data = None
+            try:
+                image_data = base64.b64decode(base64_image, validate=True)
+            except Exception as e:
+                print(f"Invalid base64 data: {e}")
+            if image_data:
+                if verify_image(image_data):
+                    image_url = upload_image(data_uri, api_key)
         else:
             image_url = "https://i.ibb.co/TBJqggw/Image-Not-Found.jpg"
     except Exception as e:
@@ -84,3 +98,11 @@ def fetch_image(title, summary, api_key):
         if os.path.exists("image.png"):
             os.remove("image.png")
         return image_url
+
+
+if __name__ == "__main__":
+    title = "Accelerated cell-type-specific regulatory evolution of the Homo sapiens brain"
+    summary = "This study investigates the accelerated regulatory evolution of gene expression in human brain cell types compared to chimpanzees. It reveals significant differences in gene expression, particularly in excitatory and inhibitory neurons, highlighting the role of regulatory evolution in human cognitive and behavioral traits"
+    api_key = "aa38b04047587c609f5c7e22f9d840f0"
+    image_url = fetch_image(title, summary, api_key)
+    print(image_url)