Update main.py

main.py

@@ -5,6 +5,7 @@ from typing import List, Optional
 import time
 import json
 import os
+import base64
 import httpx
 from dotenv import load_dotenv
 
@@ -14,6 +15,9 @@ from models import AVAILABLE_MODELS, MODEL_ALIASES
 
 # Require IMAGE_API_URL strictly from environment
 IMAGE_API_URL = os.environ["IMAGE_API_URL"]
+GPT_IMAGE_API_URL = "https://www.chatwithmono.xyz/api/image"
+SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
+SNAPZION_API_KEY = os.environ["SNAP"]
 
 app = FastAPI()
 
@@ -149,6 +153,7 @@ class ImageGenerationRequest(BaseModel):
     aspect_ratio: Optional[str] = "1:1"
     n: Optional[int] = 1
     user: Optional[str] = None
+    model: Optional[str] = "default"
 
 
 @app.post("/v1/images/generations")
@@ -157,26 +162,79 @@ async def generate_images(request: ImageGenerationRequest):
 
     async with httpx.AsyncClient(timeout=60) as client:
         for _ in range(request.n):
-            params = {
-                "prompt": request.prompt,
-                "aspect_ratio": request.aspect_ratio,
-                "link": "typegpt.net"
-            }
-
-            resp = await client.get(IMAGE_API_URL, params=params)
-
-            if resp.status_code != 200:
-                return JSONResponse(
-                    status_code=502,
-                    content={"error": "Image generation failed", "details": resp.text}
-                )
-
-            data = resp.json()
-            results.append({
-                "url": data.get("image_link"),
-                "b64_json": data.get("base64_output"),
-                "retries": data.get("attempt")
-            })
+            if request.model == "gpt-image-1":
+                headers = {
+                    'Content-Type': 'application/json',
+                    'User-Agent': 'Mozilla/5.0',
+                    'Referer': 'https://www.chatwithmono.xyz/',
+                    'sec-ch-ua-platform': '"Windows"',
+                    'sec-ch-ua': '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
+                    'sec-ch-ua-mobile': '?0',
+                }
+
+                payload = {
+                    "prompt": request.prompt,
+                    "model": "gpt-image-1"
+                }
+
+                resp = await client.post(GPT_IMAGE_API_URL, headers=headers, json=payload)
+                if resp.status_code != 200:
+                    return JSONResponse(
+                        status_code=502,
+                        content={"error": "GPT-Image generation failed", "details": resp.text}
+                    )
+
+                data = resp.json()
+                b64_image = data.get("image")
+
+                if not b64_image:
+                    return JSONResponse(status_code=502, content={"error": "No base64 image in response"})
+
+                # Reconstruct full data URI
+                full_image_data = f"data:image/png;base64,{b64_image}"
+
+                # Upload to Snapzion
+                upload_headers = {"Authorization": SNAPZION_API_KEY}
+                upload_files = {
+                    'file': ('image.png', base64.b64decode(b64_image), 'image/png')
+                }
+
+                upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
+                if upload_resp.status_code != 200:
+                    return JSONResponse(
+                        status_code=502,
+                        content={"error": "Upload failed", "details": upload_resp.text}
+                    )
+
+                upload_data = upload_resp.json()
+                image_url = upload_data.get("url")
+
+                results.append({
+                    "url": image_url,
+                    "b64_json": b64_image,
+                    "data_uri": full_image_data
+                })
+            else:
+                # Default image generator
+                params = {
+                    "prompt": request.prompt,
+                    "aspect_ratio": request.aspect_ratio,
+                    "link": "typegpt.net"
+                }
+
+                resp = await client.get(IMAGE_API_URL, params=params)
+                if resp.status_code != 200:
+                    return JSONResponse(
+                        status_code=502,
+                        content={"error": "Image generation failed", "details": resp.text}
+                    )
+
+                data = resp.json()
+                results.append({
+                    "url": data.get("image_link"),
+                    "b64_json": data.get("base64_output"),
+                    "retries": data.get("attempt")
+                })
 
     return {
         "created": int(time.time()),