Upload 4 files

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.10
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["flask", "run", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,32 @@
+# server.py
+from flask import Flask, request, jsonify
+from image_handler import ImageHandler
+import os
+
+app = Flask(__name__)
+image_handler = ImageHandler()
+
+@app.route('/v2/images/generations', methods=['POST'])
+def generate_images():
+    try:
+        # Verify authorization
+        auth_header = request.headers.get('Authorization')
+        if not auth_header or not auth_header.startswith('Bearer '):
+            return jsonify({"error": "Invalid authorization"}), 401
+
+        # Get request data
+        data = request.get_json()
+        
+        # Validate required fields
+        if not data.get('prompt'):
+            return jsonify({"error": "prompt is required"}), 400
+
+        # Generate images
+        result = image_handler.generate_images(data)
+        return jsonify(result)
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == '__main__':
+    app.run(port=5000, debug=True)

image_handler.py

@@ -0,0 +1,86 @@
+# image_handler.py
+import time
+from typing import Dict, Any
+from datetime import datetime
+import requests
+import base64
+from requests.exceptions import RequestException
+import tempfile
+import os
+
+class ImageHandler:
+    def __init__(self):
+        self.available_models = ["flux", "flux-realism", "flux-3d", "flux-pro", 
+                               "flux-anime", "flux-cablyai", "turbo", "flux-dark"]
+        self.base_url = "https://image.pollinations.ai/prompt/{prompt}"
+        
+    def _get_size_dimensions(self, size: str) -> tuple:
+        size_map = {
+            "256x256": (256, 256),
+            "512x512": (512, 512),
+            "1024x1024": (1024, 1024),
+            "1792x1024": (1792, 1024),
+            "1024x1792": (1024, 1792)
+        }
+        return size_map.get(size, (768, 768))
+
+    def _save_temp_image(self, image_data: bytes) -> str:
+        """Save image to temporary file and return its URL."""
+        temp_dir = tempfile.gettempdir()
+        temp_file = os.path.join(temp_dir, f"image_{int(time.time())}.jpg")
+        with open(temp_file, 'wb') as f:
+            f.write(image_data)
+        return f"file://{temp_file}"
+
+    def generate_images(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        prompt = params.get("prompt")
+        n = min(params.get("n", 1), 10)
+        size = params.get("size", "1024x1024")
+        model = params.get("model", "flux-pro")
+        if model not in self.available_models:
+            model = "flux-pro"
+        response_format = params.get("response_format", "b64_json")
+        if response_format != "b64_json":
+            response_format = "bs64_json"
+        width, height = self._get_size_dimensions(size)
+        
+        headers = {
+            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
+            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:122.0) Gecko/20100101 Firefox/122.0",
+        }
+
+        images = []
+        for _ in range(n):
+            try:
+                url = self.base_url.format(prompt=prompt)
+                url += f"?width={width}&height={height}&model={model}"
+                
+                response = requests.get(url, headers=headers, timeout=60)
+                response.raise_for_status()
+                
+                # Get the image data
+                image_data = response.content
+                
+                if response_format == "b64_json":
+                    # Convert to base64
+                    image_b64 = base64.b64encode(image_data).decode('utf-8')
+                    images.append({"b64_json": image_b64})
+                else:
+                    # Two options for URL:
+                    # 1. Save locally and return file URL (good for testing)
+                    # 2. Return the pollinations URL (might expire)
+                    
+                    # Option 1: Local file
+                    temp_url = self._save_temp_image(image_data)
+                    images.append({"url": temp_url})
+                    
+                    # Option 2: Pollinations URL (uncomment if preferred)
+                    # images.append({"url": url})
+                    
+            except RequestException as e:
+                raise Exception(f"Image generation failed: {str(e)}")
+
+        return {
+            "created": int(datetime.now().timestamp()),
+            "data": images
+        }

requirements.txt

@@ -0,0 +1,5 @@
+fastapi==0.110.2
+Flask==3.0.3
+python-dotenv==1.0.1
+Requests==2.31.0
+uvicorn==0.29.0