Update app.py

app.py

@@ -154,34 +154,38 @@ async def v1_chat(request: Request, authorization: str = Header(None)):
 # ---------------------------------------------------------------------
 @api.post("/v1/images/generations")
 async def v1_images_generations(request: Request, authorization: str = Header(None)):
-    import base64
-
+    """
+    OpenAI-compatible image generation endpoint → Bytez DALL-E 3
+    """
     check_key(authorization)
 
     if not BYTEZ_AUTH:
         raise HTTPException(status_code=500, detail="Server BYTEZ_API_KEY not configured")
 
-    # ---------------- parse request ----------------
+    # -------- Parse request --------
     try:
         payload = await request.json()
     except json.JSONDecodeError:
         raise HTTPException(status_code=400, detail="Invalid JSON")
 
-    prompt = payload.get("prompt") or payload.get("text")
-    if not prompt or not str(prompt).strip():
+    prompt = payload.get("prompt")
+    if not prompt:
         raise HTTPException(status_code=400, detail="Field 'prompt' is required")
 
     n = int(payload.get("n", 1))
-    response_format = payload.get("response_format", "b64_json")  # ignored; we always return b64_json
+    response_format = payload.get("response_format", "url")  # url | b64_json
 
-    # ---------------- Bytez request ----------------
-    bytez_payload = { "text": prompt }
+    # -------- Bytez request format --------
+    bytez_payload = {
+        "text": prompt
+    }
 
     headers = {
-        "Authorization": BYTEZ_AUTH,
+        "Authorization": f"Bearer {BYTEZ_AUTH}",
         "Content-Type": "application/json",
     }
 
+    # -------- Call Bytez --------
     async with httpx.AsyncClient(timeout=200) as client:
         try:
             resp = await client.post(BYTEZ_IMAGE_URL, json=bytez_payload, headers=headers)
@@ -193,76 +197,64 @@ async def v1_images_generations(request: Request, authorization: str = Header(No
             except:
                 detail = e.response.text
             raise HTTPException(status_code=e.response.status_code, detail={"upstream_error": detail})
-        except Exception as e:
-            raise HTTPException(status_code=502, detail=f"Bytez unreachable: {str(e)}")
 
     try:
         bytez_data = resp.json()
-    except json.JSONDecodeError:
+    except:
         raise HTTPException(status_code=502, detail="Bytez returned invalid JSON")
 
-    # ---------------- normalize Bytez → list of images ----------------
-    images = []
-
-    if isinstance(bytez_data, dict) and "images" in bytez_data:
-        images = bytez_data["images"]
-
-    elif isinstance(bytez_data, dict) and "data" in bytez_data:
-        for item in bytez_data["data"]:
-            images.append(item.get("url") or item.get("b64_json") or str(item))
-
-    elif isinstance(bytez_data, str):
-        images = [bytez_data]
-
-    else:
-        images = [str(bytez_data)]
-
-    if not images:
-        raise HTTPException(status_code=500, detail="No images returned from Bytez")
-
-    images = images[:n]
-
-    # ---------------- ALWAYS RETURN BASE64 ----------------
+    # -----------------------------------------------------------------
+    # -------- Extract URL (bytez_data.output) --------
+    # -----------------------------------------------------------------
+    url_output = None
+    if isinstance(bytez_data, dict):
+        url_output = bytez_data.get("output")
+
+    # -----------------------------------------------------------------
+    # -------- Extract Base64 (provider.generatedImages[].imageBytes) --
+    # -----------------------------------------------------------------
+    b64_list = []
+    try:
+        generated = bytez_data["provider"]["generatedImages"]
+        for g in generated:
+            img = g.get("image", {})
+            if "imageBytes" in img:
+                b64_list.append(img["imageBytes"])
+    except Exception:
+        pass  # ignore if provider/genImages is missing
+
+    # -----------------------------------------------------------------
+    # -------- Build OpenAI-compatible response ------------------------
+    # -----------------------------------------------------------------
     openai_data = []
+    count = max(len(b64_list), 1)
+    count = min(count, n)
+
+    for i in range(count):
+        # --- Base64 Output ---
+        if response_format == "b64_json":
+            if i < len(b64_list):
+                openai_data.append({"b64_json": b64_list[i]})
+            else:
+                # fallback: turn URL into base64-json payload (OpenAI accepts raw string)
+                openai_data.append({"b64_json": url_output})
+        else:
+            # --- URL Output ---
+            if url_output:
+                openai_data.append({"url": url_output})
+            elif i < len(b64_list):  
+                # fallback: convert base64 into a data URL
+                openai_data.append({"url": f"data:image/png;base64,{b64_list[i]}"})
+            else:
+                openai_data.append({"url": ""})
+
+    result = {
+        "created": int(time.time()),
+        "data": openai_data,
+    }
 
-    async def fetch_and_b64(url: str) -> str:
-        async with httpx.AsyncClient(timeout=60) as client:
-            r = await client.get(url)
-            r.raise_for_status()
-            return base64.b64encode(r.content).decode("utf-8")
-
-    for img in images:
-
-        # Case 1: data:image/...;base64,...
-        if isinstance(img, str) and img.startswith("data:image"):
-            b64 = img.split("base64,", 1)[1]
-            openai_data.append({"b64_json": b64})
-            continue
-
-        # Case 2: pure base64
-        if isinstance(img, str) and not img.startswith("http"):
-            openai_data.append({"b64_json": img})
-            continue
+    return JSONResponse(result, headers={"Access-Control-Allow-Origin": "*"})
 
-        # Case 3: URL → fetch + convert
-        if isinstance(img, str) and img.startswith("http"):
-            try:
-                b64 = await fetch_and_b64(img)
-                openai_data.append({"b64_json": b64})
-            except Exception as e:
-                raise HTTPException(
-                    status_code=502,
-                    detail=f"Failed to fetch image from URL returned by Bytez: {str(e)}"
-                )
-
-    # ---------------- return clean OpenAI-style result ----------------
-    return JSONResponse(
-        {
-            "created": int(time.time()),
-            "data": openai_data
-        },
-        headers={"Access-Control-Allow-Origin": "*"}
-    )