Update main.py

main.py

@@ -11,41 +11,35 @@ from dotenv import load_dotenv
 
 load_dotenv()
 
-from models import AVAILABLE_MODELS, MODEL_ALIASES
+# Dummy imports for example
+from models import AVAILABLE_MODELS, MODEL_ALIASES  # Ensure these are defined properly
 
-# Require IMAGE_API_URL strictly from environment
+# Env variables
 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()
 
-
 def unix_id():
     return str(int(time.time() * 1000))
 
-
-# ==== Models ====
-
+# === Models ===
 @app.get("/v1/models")
 async def list_models():
     return {"object": "list", "data": AVAILABLE_MODELS}
 
-
-# ==== Chat Completion ====
+# === Chat Completion ===
 
 class Message(BaseModel):
     role: str
     content: str
 
-
 class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
 
-
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
     model_id = MODEL_ALIASES.get(request.model, request.model)
@@ -104,7 +98,6 @@ async def chat_completion(request: ChatRequest):
                             }
                             yield f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n"
         return StreamingResponse(event_stream(), media_type="text/event-stream")
-
     else:
         assistant_response = ""
         usage_info = {}
@@ -145,8 +138,7 @@ async def chat_completion(request: ChatRequest):
             }
         })
 
-
-# ==== Image Generation ====
+# === Image Generation ===
 
 class ImageGenerationRequest(BaseModel):
     prompt: str
@@ -155,14 +147,15 @@ class ImageGenerationRequest(BaseModel):
     user: Optional[str] = None
     model: Optional[str] = "default"
 
-
 @app.post("/v1/images/generations")
 async def generate_images(request: ImageGenerationRequest):
     results = []
 
     async with httpx.AsyncClient(timeout=60) as client:
         for _ in range(request.n):
-            if request.model == "gpt-image-1":
+            model = request.model or "default"
+
+            if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
                 headers = {
                     'Content-Type': 'application/json',
                     'User-Agent': 'Mozilla/5.0',
@@ -174,26 +167,23 @@ async def generate_images(request: ImageGenerationRequest):
 
                 payload = {
                     "prompt": request.prompt,
-                    "model": "gpt-image-1"
+                    "model": model
                 }
 
-                resp = await client.post(GPT_IMAGE_API_URL, headers=headers, json=payload)
+                resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
                 if resp.status_code != 200:
                     return JSONResponse(
                         status_code=502,
-                        content={"error": "GPT-Image generation failed", "details": resp.text}
+                        content={"error": f"{model} 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"})
+                    return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
 
-                # Reconstruct full data URI
-                full_image_data = f"data:image/png;base64,{b64_image}"
+                data_uri = 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')
@@ -207,15 +197,13 @@ async def generate_images(request: ImageGenerationRequest):
                     )
 
                 upload_data = upload_resp.json()
-                image_url = upload_data.get("url")
-
                 results.append({
-                    "url": image_url,
+                    "url": upload_data.get("url"),
                     "b64_json": b64_image,
-                    "data_uri": full_image_data
+                    "data_uri": data_uri,
+                    "model": model
                 })
             else:
-                # Default image generator
                 params = {
                     "prompt": request.prompt,
                     "aspect_ratio": request.aspect_ratio,
@@ -233,7 +221,8 @@ async def generate_images(request: ImageGenerationRequest):
                 results.append({
                     "url": data.get("image_link"),
                     "b64_json": data.get("base64_output"),
-                    "retries": data.get("attempt")
+                    "retries": data.get("attempt"),
+                    "model": "default"
                 })
 
     return {