Update app.py

app.py

@@ -154,26 +154,10 @@ 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)):
-    """
-    OpenAI-compatible image generation endpoint that forwards to:
-      https://api.bytez.com/models/v2/openai/dall-e-3
-
-    It accepts the usual OpenAI body:
-      {
-        "model": "dall-e-3",           # REQUIRED by clients, but ignored
-        "prompt": "text",              # REQUIRED
-        "n": 1,
-        "size": "1024x1024",
-        "response_format": "url" | "b64_json",
-        "quality": "standard" | "hd",
-        "style": "vivid" | "natural"
-      }
-    and converts it to the Bytez format from your curl:
-      { "text": "<prompt>" }
-    """
+    import base64
+
     check_key(authorization)
 
-    # Use SAME key as chat unless you *really* have a separate one
     if not BYTEZ_AUTH:
         raise HTTPException(status_code=500, detail="Server BYTEZ_API_KEY not configured")
 
@@ -187,54 +171,28 @@ async def v1_images_generations(request: Request, authorization: str = Header(No
     if not prompt or not str(prompt).strip():
         raise HTTPException(status_code=400, detail="Field 'prompt' is required")
 
-    # These are there only to satisfy OpenAI-compatible clients
-    _model = payload.get("model", "dall-e-3")  # ignored
     n = int(payload.get("n", 1))
-    response_format = payload.get("response_format", "b64_json")  # "url" | "b64_json"
-    size = payload.get("size", "1024x1024")
-    quality = payload.get("quality")
-    style = payload.get("style")
-
-    # ---------------- build Bytez request ----------------
-    # Minimal payload that we KNOW works from your curl:
-    bytez_payload = {
-        "text": prompt,
-    }
+    response_format = payload.get("response_format", "b64_json")  # ignored; we always return b64_json
 
-    # Only add extras if Bytez actually supports them.
-    # Comment these out if you’re unsure; or keep them if Bytez docs say so.
-    # bytez_payload["num_outputs"] = n
-    # bytez_payload["size"] = size
-    # if quality in ["standard", "hd"]:
-    #     bytez_payload["quality"] = quality
-    # if style in ["vivid", "natural"]:
-    #     bytez_payload["style"] = style
+    # ---------------- Bytez request ----------------
+    bytez_payload = { "text": prompt }
 
     headers = {
-        "Authorization": BYTEZ_AUTH,         # << important
+        "Authorization": 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,
-            )
+            resp = await client.post(BYTEZ_IMAGE_URL, json=bytez_payload, headers=headers)
             resp.raise_for_status()
         except httpx.HTTPStatusError as e:
-            # Surface Bytez error message directly to client
             detail = None
             try:
                 detail = e.response.json()
-            except Exception:
+            except:
                 detail = e.response.text
-            raise HTTPException(
-                status_code=e.response.status_code,
-                detail={"upstream_error": detail},
-            )
+            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)}")
 
@@ -243,69 +201,68 @@ async def v1_images_generations(request: Request, authorization: str = Header(No
     except json.JSONDecodeError:
         raise HTTPException(status_code=502, detail="Bytez returned invalid JSON")
 
-    # ---------------- map Bytez → OpenAI image response ----------------
-    # We don't know the exact shape, so we handle several possibilities.
-
+    # ---------------- normalize Bytez → list of images ----------------
     images = []
 
-    # Case 1: { "images": ["url_or_b64", ...] }
-    if isinstance(bytez_data, dict) and "images" in bytez_data and isinstance(bytez_data["images"], list):
+    if isinstance(bytez_data, dict) and "images" in bytez_data:
         images = bytez_data["images"]
 
-    # Case 2: { "data": [ { "url": "..." } , ... ] }
-    elif isinstance(bytez_data, dict) and "data" in bytez_data and isinstance(bytez_data["data"], list):
+    elif isinstance(bytez_data, dict) and "data" in bytez_data:
         for item in bytez_data["data"]:
-            if "url" in item:
-                images.append(item["url"])
-            elif "b64_json" in item:
-                images.append(item["b64_json"])
-            else:
-                images.append(str(item))
-
-    # Case 3: single string
+            images.append(item.get("url") or item.get("b64_json") or str(item))
+
     elif isinstance(bytez_data, str):
         images = [bytez_data]
 
-    # Fallback: treat whole thing as one string
     else:
         images = [str(bytez_data)]
 
     if not images:
-        raise HTTPException(status_code=500, detail={"error": "No images returned from Bytez", "raw": bytez_data})
+        raise HTTPException(status_code=500, detail="No images returned from Bytez")
 
-    # truncate to n images
     images = images[:n]
 
+    # ---------------- ALWAYS RETURN BASE64 ----------------
     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: Already base64
-    if isinstance(img, str) and img.startswith("data:image"):
-        b64_part = img.split("base64,", 1)[1]
-        openai_data.append({"b64_json": b64_part})
-        continue
-
-    # Case 2: Looks like base64 (no URL)
-    if isinstance(img, str) and not img.startswith("http"):
-        openai_data.append({"b64_json": img})
-        continue
-
-    # Case 3: It is a URL → fetch + convert to base64
-    if isinstance(img, str) and img.startswith("http"):
-        try:
-            b64 = await fetch_and_b64(img)
+    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})
-        except Exception as e:
-            raise HTTPException(
-                status_code=502,
-                detail=f"Failed to fetch image from URL returned by Bytez: {str(e)}"
-            )
+            continue
+
+        # Case 2: pure base64
+        if isinstance(img, str) and not img.startswith("http"):
+            openai_data.append({"b64_json": img})
+            continue
+
+        # 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": "*"}
+    )