polats Claude Opus 4.8 (1M context) commited on
Commit
acb2593
·
1 Parent(s): 2151ea1

Portraits: set correct Content-Type for cloud images (NIM returns JPEG)

Browse files

Sniff the image magic bytes so the response Content-Type matches the actual
format (FLUX/NIM returns JPEG, local Z-Image returns PNG) — keeps the cached
blob's type correct on the client.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Files changed (1) hide show
  1. app.py +11 -1
app.py CHANGED
@@ -508,6 +508,16 @@ def _hf_portrait(prompt, model="black-forest-labs/FLUX.1-schnell"):
508
  return data, None
509
 
510
 
 
 
 
 
 
 
 
 
 
 
511
  @fastapi_app.post("/portrait")
512
  async def portrait(request: Request):
513
  body = await request.json()
@@ -535,7 +545,7 @@ async def portrait(request: Request):
535
  return Response("no image provider (set NVIDIA_NIM_API_KEY / HF_TOKEN, or TINY_IMAGE_MODE=local)", status_code=503)
536
  if err:
537
  return Response(err, status_code=502)
538
- return Response(png, media_type="image/png", headers={"Cache-Control": "no-store"})
539
 
540
 
541
  @fastapi_app.get("/persona/status")
 
508
  return data, None
509
 
510
 
511
+ def _img_mime(data):
512
+ if data[:3] == b"\xff\xd8\xff":
513
+ return "image/jpeg" # NIM/FLUX returns JPEG
514
+ if data[:8] == b"\x89PNG\r\n\x1a\n":
515
+ return "image/png" # local Z-Image
516
+ if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
517
+ return "image/webp"
518
+ return "image/png"
519
+
520
+
521
  @fastapi_app.post("/portrait")
522
  async def portrait(request: Request):
523
  body = await request.json()
 
545
  return Response("no image provider (set NVIDIA_NIM_API_KEY / HF_TOKEN, or TINY_IMAGE_MODE=local)", status_code=503)
546
  if err:
547
  return Response(err, status_code=502)
548
+ return Response(png, media_type=_img_mime(png), headers={"Cache-Control": "no-store"})
549
 
550
 
551
  @fastapi_app.get("/persona/status")