289 lines
10 KiB
Python
289 lines
10 KiB
Python
import os
|
|
import json
|
|
from io import BytesIO
|
|
from pprint import pprint
|
|
|
|
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
|
|
from fastapi.responses import Response
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
from dotenv import load_dotenv
|
|
from google import genai
|
|
from google.genai.types import Part
|
|
from PIL import Image
|
|
|
|
# --- Load .env (for local development) ---
|
|
load_dotenv()
|
|
|
|
# ---- Gemini SDK Setup ----
|
|
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
|
|
|
# Strip whitespace (common issue with .env files)
|
|
if GEMINI_API_KEY:
|
|
GEMINI_API_KEY = GEMINI_API_KEY.strip()
|
|
|
|
if not GEMINI_API_KEY:
|
|
raise RuntimeError("GEMINI_API_KEY not set")
|
|
|
|
client = genai.Client(api_key=GEMINI_API_KEY)
|
|
IMAGE_MODEL = "gemini-2.5-flash-image"
|
|
|
|
app = FastAPI()
|
|
|
|
# (Optional) CORS if you call it from frontend directly
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"], # tighten in production
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
def find_first_image_part(resp):
|
|
"""Return the actual output image part."""
|
|
if not getattr(resp, "candidates", None):
|
|
return None
|
|
for cand in resp.candidates:
|
|
content = getattr(cand, "content", None)
|
|
if not content or not getattr(content, "parts", None):
|
|
continue
|
|
for p in content.parts:
|
|
if hasattr(p, "inline_data") and getattr(p.inline_data, "data", None):
|
|
return p.inline_data
|
|
if hasattr(p, "blob") and getattr(p.blob, "data", None):
|
|
return p.blob
|
|
return None
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/api/try-on/", summary="Try-on (user_photo + outfit_photo + metadata)")
|
|
async def try_on(
|
|
user_photo: UploadFile = File(...),
|
|
outfit_photo: UploadFile = File(...),
|
|
metadata: str = Form("", description="JSON or text describing the outfit items"),
|
|
):
|
|
"""
|
|
metadata can be:
|
|
- free text: 'black hoodie, patterned shirt, white shorts'
|
|
- JSON: '{"outer": "hoodie", "inner": "shirt", "bottom": "shorts"}'
|
|
- any additional instructions
|
|
"""
|
|
|
|
# ---- Read files ----
|
|
user_bytes = await user_photo.read()
|
|
outfit_bytes = await outfit_photo.read()
|
|
|
|
if not user_bytes:
|
|
raise HTTPException(status_code=400, detail="user_photo is empty")
|
|
if not outfit_bytes:
|
|
raise HTTPException(status_code=400, detail="outfit_photo is empty")
|
|
|
|
# ---- Parse metadata if JSON ----
|
|
metadata_text = ""
|
|
try:
|
|
meta = json.loads(metadata)
|
|
metadata_text = json.dumps(meta, indent=2)
|
|
except:
|
|
metadata_text = metadata.strip()
|
|
|
|
# ---- Build Prompt ----
|
|
prompt = f"""
|
|
EDIT TASK — FULL OUTFIT REPLACEMENT
|
|
|
|
FIRST image = user body/mirror selfie
|
|
SECOND image = outfit reference
|
|
|
|
Follow these rules:
|
|
|
|
1. Edit ONLY the body/clothing of the FIRST image.
|
|
|
|
2. Replace ALL clothing with the items visible in the SECOND image (outfit photo).
|
|
|
|
3. Keep the user's:
|
|
- face
|
|
- skin
|
|
- hands
|
|
- phone
|
|
- pose
|
|
- environment
|
|
EXACTLY the same.
|
|
|
|
4. OUTFIT PHOTO AND METADATA CONSIDERATION (CRITICAL):
|
|
|
|
You MUST consider BOTH sources:
|
|
|
|
A) SECOND IMAGE (OUTFIT PHOTO):
|
|
- Carefully examine ALL clothing items visible in the outfit photo
|
|
- Transfer EVERY piece of clothing from the outfit photo to the user photo
|
|
- Pay attention to the visual details: colors, patterns, textures, styles, fit, and layering
|
|
- Ensure all visible items (tops, bottoms, outerwear, accessories, etc.) are transferred
|
|
|
|
B) METADATA:
|
|
The following metadata provides additional information about the outfit:
|
|
|
|
Metadata:
|
|
{metadata_text}
|
|
|
|
Use the metadata to:
|
|
- Identify and verify ALL items mentioned in the metadata are present in the outfit photo
|
|
- Understand the complete outfit composition and layering structure
|
|
- Ensure no items are missed - cross-reference metadata items with the outfit photo
|
|
- Apply any specific details from metadata (colors, patterns, materials, styles) that match the outfit photo
|
|
- If metadata mentions items not clearly visible in the photo, use metadata details to enhance accuracy
|
|
|
|
C) COMBINE BOTH SOURCES:
|
|
- Use the outfit photo as the primary visual reference
|
|
- Use metadata to ensure completeness and accuracy
|
|
- Transfer ALL items that appear in the outfit photo AND are mentioned in metadata
|
|
- The final result should include every clothing item visible in the outfit photo, verified against the metadata
|
|
|
|
5. Ensure:
|
|
- realistic lighting
|
|
- correct fabric physics
|
|
- correct drape
|
|
- realism
|
|
- no blending of people
|
|
- no copying faces from the second image
|
|
- all items from both the outfit photo and metadata are properly integrated into the user photo
|
|
|
|
6. Output ONE final edited image with ALL items from the outfit photo transferred to the user, ensuring all metadata items are also included.
|
|
|
|
IMPORTANT:
|
|
- FIRST = user_photo
|
|
- SECOND = outfit_photo
|
|
- You must use BOTH the outfit photo (visual reference) AND metadata (descriptive reference)
|
|
- ALL items visible in the outfit photo AND mentioned in metadata must appear in the final image
|
|
"""
|
|
|
|
# ---- Wrap image files ----
|
|
user_part = Part.from_bytes(data=user_bytes, mime_type=user_photo.content_type or "image/jpeg")
|
|
outfit_part = Part.from_bytes(data=outfit_bytes, mime_type=outfit_photo.content_type or "image/jpeg")
|
|
|
|
# ---- Call Gemini ----
|
|
print(f"\n[DEBUG] Using model: {IMAGE_MODEL}")
|
|
print(f"[DEBUG] Metadata received: {metadata[:100] if metadata else '(empty)'}")
|
|
|
|
try:
|
|
response = client.models.generate_content(
|
|
model=IMAGE_MODEL,
|
|
contents=[
|
|
prompt,
|
|
user_part,
|
|
outfit_part,
|
|
"First image = user selfie. Second image = outfit reference.",
|
|
],
|
|
)
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
# Check if it's an API key error
|
|
if "API key" in error_msg or "INVALID_ARGUMENT" in error_msg or "API_KEY" in error_msg:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail=(
|
|
f"Gemini API key error: {error_msg}. "
|
|
"Please verify your GEMINI_API_KEY in the .env file is correct."
|
|
)
|
|
)
|
|
# Log the full error for debugging
|
|
import traceback
|
|
print(f"[ERROR] Gemini API error: {error_msg}")
|
|
print(traceback.format_exc())
|
|
raise HTTPException(status_code=502, detail=f"Gemini error: {error_msg}")
|
|
|
|
# ---- Debug: Inspect response ----
|
|
print("\n[DEBUG] === Gemini Response Analysis ===")
|
|
if not getattr(response, "candidates", None):
|
|
print("[DEBUG] No candidates at all. Full response:")
|
|
pprint(response)
|
|
else:
|
|
cand = response.candidates[0]
|
|
finish_reason = getattr(cand, "finish_reason", None)
|
|
print(f"[DEBUG] Finish reason: {finish_reason}")
|
|
print(f"[DEBUG] Safety ratings: {getattr(cand, 'safety_ratings', None)}")
|
|
|
|
if hasattr(cand, "content") and getattr(cand.content, "parts", None):
|
|
print(f"[DEBUG] Parts in first candidate ({len(cand.content.parts)} total):")
|
|
for i, part in enumerate(cand.content.parts):
|
|
part_type = type(part).__name__
|
|
print(f" Part {i}: type={part_type}")
|
|
|
|
# Check for text
|
|
text = getattr(part, "text", None)
|
|
if text:
|
|
print(f" TEXT: {text[:300]}")
|
|
|
|
# Check for inline_data
|
|
if hasattr(part, "inline_data"):
|
|
inline = part.inline_data
|
|
if inline:
|
|
data_size = len(getattr(inline, "data", b""))
|
|
mime = getattr(inline, "mime_type", "unknown")
|
|
print(f" INLINE_DATA: mime={mime}, size={data_size} bytes")
|
|
|
|
# Check for blob
|
|
if hasattr(part, "blob"):
|
|
blob = part.blob
|
|
if blob:
|
|
data_size = len(getattr(blob, "data", b""))
|
|
mime = getattr(blob, "mime_type", "unknown")
|
|
print(f" BLOB: mime={mime}, size={data_size} bytes")
|
|
print("[DEBUG] === End Response Analysis ===\n")
|
|
|
|
# ---- Extract output image ----
|
|
img_part = find_first_image_part(response)
|
|
if not img_part or not getattr(img_part, "data", None):
|
|
# Try to extract any text message from Gemini
|
|
message = "No image returned from Gemini."
|
|
|
|
if getattr(response, "candidates", None) and len(response.candidates) > 0:
|
|
cand = response.candidates[0]
|
|
finish_reason = getattr(cand, "finish_reason", "unknown")
|
|
|
|
# Extract text explanation from Gemini
|
|
if hasattr(cand, "content") and getattr(cand.content, "parts", None):
|
|
for part in cand.content.parts:
|
|
txt = getattr(part, "text", None)
|
|
if txt:
|
|
message += f" Gemini says: {txt[:500]}"
|
|
break
|
|
|
|
# Check safety ratings
|
|
if hasattr(cand, "safety_ratings"):
|
|
safety_ratings = cand.safety_ratings
|
|
if safety_ratings:
|
|
blocked = [
|
|
f"{getattr(r, 'category', 'unknown')}"
|
|
for r in safety_ratings
|
|
if hasattr(r, "blocked") and getattr(r, "blocked", False)
|
|
]
|
|
if blocked:
|
|
message += f" Safety blocked: {', '.join(blocked)}"
|
|
|
|
raise HTTPException(
|
|
status_code=502,
|
|
detail=f"{message} Finish reason: {finish_reason}",
|
|
)
|
|
else:
|
|
raise HTTPException(
|
|
status_code=502,
|
|
detail="No image returned from Gemini. No candidates in response.",
|
|
)
|
|
|
|
raw = img_part.data
|
|
|
|
# Ensure valid image
|
|
try:
|
|
Image.open(BytesIO(raw))
|
|
except Exception as e:
|
|
raise HTTPException(
|
|
status_code=500,
|
|
detail=f"Returned image is invalid: {e}. Image size: {len(raw)} bytes"
|
|
)
|
|
|
|
return Response(content=raw, media_type="image/png")
|