Update image_fetcher.py

image_fetcher.py

@@ -1,73 +1,70 @@
-
 import os
-import time
+import asyncio
+import aiohttp
 import random
-import requests
 from PIL import Image
 from io import BytesIO
 from duckduckgo_search import DDGS
+import time
 
-# Create folders
-DIRS = {
-    "base": "images",
-    "temp": os.path.join("images", "tmp")
-}
-
-# Ensure directories exist
-path = os.path.join(os.getcwd(), 'images')
-os.makedirs(path, exist_ok=True)
+# Setup /tmp directory for Hugging Face
+IMAGE_DIR = "/tmp/images"
+os.makedirs(IMAGE_DIR, exist_ok=True)
 
-# Helper: Random delay + headers
+# Headers
 def get_headers():
     user_agents = [
         "Mozilla/5.0 (Windows NT 10.0; Win64; x64)",
         "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
         "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:90.0)"
     ]
-    return {'User-Agent': random.choice(user_agents)}
+    return {"User-Agent": random.choice(user_agents)}
 
-# Check if image meets requirements
-def is_valid_image(img):
-    width, height = img.size
-    ratio = round(width / height, 2)
-    return width >= 854 and height >= 480 and abs(ratio - (16 / 9)) <= 0.2  # relaxed ratio check
+# Image validation
+def is_valid_image(image):
+    width, height = image.size
+    aspect_ratio = round(width / height, 2)
+    return width >= 854 and height >= 480 and abs(aspect_ratio - (16 / 9)) <= 0.2
 
-# Try to download and validate image
-def validate_image(img_url, keyword):
+# Async fetch image
+async def fetch_image(session, url, name):
     try:
-        response = requests.get(img_url, headers=get_headers(), timeout=10)
-        img = Image.open(BytesIO(response.content)).convert("RGB")
-        if not is_valid_image(img):
-            print(f"Skipped (not valid): {img_url}")
-            return False
-        final_path = os.path.join("images", f"{keyword}.jpg")
-        img.save(final_path)
-        print(f"Saved: {final_path} | Size: {img.size}")
-        return True
-    except Exception as e:
-        print(f"Failed to validate {img_url}: {e}")
-        return False
+        async with session.get(url, timeout=10) as response:
+            content = await response.read()
+            image = Image.open(BytesIO(content)).convert("RGB")
 
-# Search and fetch images from DuckDuckGo
-def fetch_images_from_prompts(prompt_list):
-    for keyword in prompt_list:
-        safe_keyword = keyword.replace(" ", "_").lower()
-        print(f"\nSearching for: {keyword}")
-        found = False
+            if not is_valid_image(image):
+                return f"Skipped (invalid): {name}"
 
-        try:
-            with DDGS() as ddgs:
-                results = ddgs.images(keyword, max_results=20)
-                for result in results:
-                    if validate_image(result['image'], safe_keyword):
-                        found = True
-                        break
-                    time.sleep(random.uniform(1.5, 3.5))  # anti-bot delay
+            unique_name = f"{name}_{int(time.time() * 1000)}.jpg"
+            path = os.path.join(IMAGE_DIR, unique_name)
+            image.save(path)
+            return f"Saved: {unique_name}"
 
-        except Exception as e:
-            print(f"Rate limit or error: {e}. Sleeping 10 seconds...")
-            time.sleep(10)  # longer wait on rate limit
+    except Exception as e:
+        return f"Error: {name} | {e}"
 
-        if not found:
-            print(f"No suitable image found for: {keyword}")
+# Async search + download with DDGS inside thread
+async def search_and_download(session, prompt, sem):
+    async with sem:
+        name = prompt.replace(" ", "_").lower()
+        try:
+            loop = asyncio.get_event_loop()
+            results = await loop.run_in_executor(None, lambda: list(DDGS().images(prompt, max_results=15)))
+            for item in results:
+                url = item.get("image")
+                result = await fetch_image(session, url, name)
+                if "Saved" in result:
+                    return result
+            return f"No valid image for: {prompt}"
+        except Exception as e:
+            return f"Search failed for {prompt}: {e}"
 
+# Main runner
+async def main(prompts):
+    sem = asyncio.Semaphore(5)  # Limit concurrency
+    async with aiohttp.ClientSession(headers=get_headers()) as session:
+        tasks = [search_and_download(session, prompt, sem) for prompt in prompts]
+        results = await asyncio.gather(*tasks)
+        for res in results:
+            print(res)