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. 3. Keep the user's: - face - skin - hands - phone - pose - environment EXACTLY the same. 4. IF a metadata description is given, use it to understand layers: Metadata: {metadata_text} 5. Ensure: - realistic lighting - correct fabric physics - correct drape - realism - no blending of people - no copying faces from the second image 6. Output ONE final edited image. IMPORTANT: - FIRST = user_photo - SECOND = outfit_photo """ # ---- 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")