516 lines
19 KiB
Python
516 lines
19 KiB
Python
import os
|
||
import json
|
||
from io import BytesIO
|
||
from pprint import pprint
|
||
|
||
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
|
||
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 google.genai import types as genai_types
|
||
from PIL import Image
|
||
import httpx # ✅ for proxy support
|
||
|
||
# --- Load .env (for local development only; in Docker envs are already set) ---
|
||
load_dotenv()
|
||
|
||
# ---- Gemini SDK + Proxy 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")
|
||
|
||
IMAGE_MODEL = "gemini-2.5-flash-image"
|
||
IDENTITY_MODEL = "gemini-1.5-flash" # text/vision model for identity check
|
||
|
||
AUTH_KEY = os.getenv("AIFIT_AUTH_KEY")
|
||
|
||
# 🔌 Proxy URL, e.g. "socks5://xray:10808"
|
||
PROXY_URL = os.getenv("AIFIT_PROXY_URL")
|
||
|
||
http_options = None
|
||
if PROXY_URL:
|
||
# Use httpx transports so ALL Gemini calls go through Xray
|
||
http_options = genai_types.HttpOptions(
|
||
client_args={
|
||
"transport": httpx.HTTPTransport(proxy=PROXY_URL),
|
||
},
|
||
async_client_args={
|
||
"transport": httpx.AsyncHTTPTransport(proxy=PROXY_URL),
|
||
},
|
||
)
|
||
|
||
client = genai.Client(
|
||
api_key=GEMINI_API_KEY,
|
||
http_options=http_options, # ✅ None locally, proxy in server
|
||
)
|
||
|
||
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
|
||
|
||
|
||
def faces_match_strict(original_bytes: bytes, edited_bytes: bytes) -> bool:
|
||
"""
|
||
Second Gemini call: compare ORIGINAL vs EDITED and decide
|
||
if identity/face has changed. Returns True if it's clearly
|
||
the same person, False otherwise.
|
||
"""
|
||
try:
|
||
original_part = Part.from_bytes(
|
||
data=original_bytes,
|
||
mime_type="image/jpeg",
|
||
)
|
||
edited_part = Part.from_bytes(
|
||
data=edited_bytes,
|
||
mime_type="image/png",
|
||
)
|
||
|
||
prompt = """
|
||
You are an identity consistency checker.
|
||
|
||
You receive TWO photos:
|
||
- FIRST: original_user (the real human photo from the user)
|
||
- SECOND: generated_tryon (the AI-edited try-on result)
|
||
|
||
Task:
|
||
1. Decide if the TWO images clearly show the SAME PERSON.
|
||
2. Focus on:
|
||
- overall facial structure and proportions
|
||
- jawline, chin, cheekbones
|
||
- nose shape
|
||
- eye / brow region
|
||
- beard/moustache presence, shape, and density
|
||
- hairline and hairstyle
|
||
3. Ignore clothing changes. We ONLY care about whether the face/identity changed.
|
||
|
||
Instructions:
|
||
- If the SECOND image looks like a different model, or the face/identity is noticeably altered
|
||
(different jaw, nose, beard, hairline, etc.), you MUST treat it as a DIFFERENT person.
|
||
- Do NOT be generous. If you are not sure, assume the identity has changed.
|
||
|
||
Output format:
|
||
Answer with EXACTLY ONE WORD:
|
||
- "OK" → clearly the same person, acceptable tiny rendering noise
|
||
- "CHANGE" → different person OR visibly altered identity
|
||
|
||
No explanations. No extra text.
|
||
"""
|
||
|
||
resp = client.models.generate_content(
|
||
model=IDENTITY_MODEL,
|
||
contents=[prompt, original_part, edited_part],
|
||
)
|
||
|
||
text = ""
|
||
if getattr(resp, "candidates", None):
|
||
cand = resp.candidates[0]
|
||
if hasattr(cand, "content") and getattr(cand.content, "parts", None):
|
||
for part in cand.content.parts:
|
||
if getattr(part, "text", None):
|
||
text += part.text
|
||
|
||
answer = (text or "").strip().upper()
|
||
print(f"[DEBUG] Identity check raw answer: {answer!r}")
|
||
return answer == "OK"
|
||
|
||
except Exception as e:
|
||
# If identity check fails for any reason, be SAFE and treat as mismatch
|
||
print(f"[ERROR] Identity check failed: {e!r}")
|
||
return False
|
||
|
||
|
||
@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(
|
||
request: Request,
|
||
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: '{"request_id": "uuid", "outer_layer": {...}, "bottom": {"type": "shorts"}}'
|
||
- any additional instructions
|
||
"""
|
||
|
||
# ---- AUTH CHECK (simple shared key) ----
|
||
if AUTH_KEY:
|
||
client_key = request.headers.get("X-AIFIT-Key")
|
||
if not client_key or client_key != AUTH_KEY:
|
||
raise HTTPException(status_code=401, detail="Invalid or missing authentication key")
|
||
|
||
# ---- 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 (JSON or plain text) + extract request_id + dynamic rules ----
|
||
metadata_text = ""
|
||
request_id = None
|
||
extra_rules = ""
|
||
|
||
if metadata:
|
||
try:
|
||
meta = json.loads(metadata)
|
||
metadata_text = json.dumps(meta, indent=2, ensure_ascii=False)
|
||
|
||
# Try common key names for the UUID
|
||
request_id = meta.get("request_id") or meta.get("id") or meta.get("uuid")
|
||
|
||
extra_sections = []
|
||
|
||
# Bottom: shorts
|
||
bottom = meta.get("bottom") or {}
|
||
bottom_type = (bottom.get("type") or "").lower()
|
||
if bottom_type == "shorts":
|
||
extra_sections.append(
|
||
"""
|
||
SPECIAL RULE FOR SHORTS (INSANELY STRICT, ZERO TOLERANCE – DO NOT IGNORE):
|
||
|
||
- Metadata says the bottom garment is SHORTS.
|
||
- You MUST COMPLETELY REMOVE ANY LONG PANTS OR TROUSERS from the FIRST image.
|
||
- ZERO pants are allowed to remain. No black fabric, no shadows shaped like pants,
|
||
no ghost textures. If any trace of the old pants is left, the result is WRONG and UNUSABLE.
|
||
- From the hem of the shorts down, the legs must be rendered as NATURAL, BARE LEGS
|
||
(with socks/shoes if present) with correct skin tone, shading, lighting, and anatomy.
|
||
- Only the shorts, socks and shoes may cover the legs.
|
||
- You are FORBIDDEN to stack shorts on top of pants. Shorts must fully replace pants.
|
||
- Shorts must match the SECOND image in color, length, silhouette and fit.
|
||
"""
|
||
)
|
||
|
||
# Top: t-shirt / tee / short-sleeve
|
||
top = meta.get("top") or meta.get("inner_layer") or meta.get("upper") or meta.get("outer_layer") or {}
|
||
top_type = (top.get("type") or "").lower()
|
||
if any(token in top_type for token in ["t-shirt", "tshirt", "tee", "t shirt"]):
|
||
extra_sections.append(
|
||
"""
|
||
SPECIAL RULE FOR T-SHIRT / TOP (EXTREMELY STRICT, ZERO TOLERANCE):
|
||
|
||
- Metadata says the upper garment is a T-SHIRT / SHORT-SLEEVE TOP.
|
||
- You MUST COMPLETELY REMOVE ANY ORIGINAL SHIRT, HOODIE, JACKET OR LONG SLEEVES
|
||
from the FIRST image in the areas covered by the new top.
|
||
- No ghost sleeves, no bits of old collar, no leftover cuffs. If any part of the old top
|
||
or sleeves is still visible, the result is WRONG and MUST BE TREATED AS FAILED.
|
||
- The new T-SHIRT / TOP must match the SECOND image in:
|
||
• color
|
||
• pattern
|
||
• sleeve length
|
||
• collar shape
|
||
• overall fit and silhouette.
|
||
- Do NOT layer the new t-shirt on top of the old garment. It MUST REPLACE the old garment.
|
||
"""
|
||
)
|
||
|
||
if extra_sections:
|
||
extra_rules = "\n".join(extra_sections)
|
||
|
||
except Exception:
|
||
metadata_text = metadata.strip()
|
||
|
||
# ---- Build Prompt (strict identity + clothing) ----
|
||
prompt = f"""
|
||
You are performing a **strict outfit replacement operation** on the FIRST image using the clothing and accessories from the SECOND image and the metadata provided.
|
||
|
||
⚠️ IMPORTANT SAFETY CONTEXT
|
||
- The individuals in the images are **not public figures**.
|
||
- The task is standard, allowed **fashion editing / virtual try-on** only.
|
||
- No face, body, or identity manipulation is permitted.
|
||
|
||
-----------------------------------------------------
|
||
🚨 ULTRA-HARD IDENTITY LOCK (ABSOLUTE, NON-NEGOTIABLE RULES)
|
||
|
||
These identity rules are STRONGER than all other instructions. If you cannot follow them, you MUST NOT output an image.
|
||
|
||
- The person in the FIRST image is the ONLY person you are allowed to show in the final image.
|
||
- You MUST keep the FIRST image person's:
|
||
• face
|
||
• identity
|
||
• skin tone
|
||
• facial structure
|
||
• jawline and chin shape
|
||
• nose shape
|
||
• lips and mouth shape
|
||
• facial hair (beard / moustache) style and density
|
||
• hair style and hairline
|
||
• body proportions
|
||
• pose
|
||
EXACTLY as they are.
|
||
|
||
ABSOLUTE FACE LOCK (READ-ONLY REGION):
|
||
|
||
- Treat the entire face, head, ears, and visible neck area of the FIRST image as a **READ-ONLY REGION**.
|
||
- You MUST NOT re-render, redraw, regenerate, beautify, smooth, reshape, de-age, re-gender,
|
||
or otherwise alter the face in any way.
|
||
- The face region in the final image must be **indistinguishable** from the original FIRST image,
|
||
except for tiny unavoidable differences at the clothing boundaries (e.g. collar touching neck).
|
||
|
||
If there is ANY conflict between clothing instructions and identity protection:
|
||
YOU MUST PROTECT THE FIRST PERSON'S IDENTITY and only adjust clothing areas.
|
||
If identity cannot be preserved, DO NOT generate an image.
|
||
|
||
-----------------------------------------------------
|
||
👤 SECOND IMAGE PERSON HANDLING — CLOTHES ONLY, NEVER THEIR IDENTITY
|
||
|
||
- The person in the SECOND image (outfit photo) is **NOT** allowed to appear in the final result.
|
||
- You must treat the SECOND person as a "clothing mannequin" ONLY.
|
||
- You are **STRICTLY FORBIDDEN** from copying:
|
||
• their face
|
||
• their moustache or beard
|
||
• their skin tone
|
||
• their hair
|
||
• their body shape
|
||
• their pose
|
||
• their tattoos or body details
|
||
|
||
You may ONLY use the SECOND image to extract:
|
||
• clothing items,
|
||
• shoes,
|
||
• accessories (including the glasses holder strap around the neck),
|
||
• their colors, patterns, materials, and shapes.
|
||
|
||
Under NO circumstances may the final image look like the SECOND person.
|
||
It must always clearly be the person from the FIRST image wearing those clothes.
|
||
|
||
-----------------------------------------------------
|
||
🎯 PRIMARY OBJECTIVE (ZERO TOLERANCE FOR MISTAKES)
|
||
|
||
Replace ALL clothing and accessories in the FIRST IMAGE with the outfit shown in the SECOND IMAGE, cross-verified with the metadata below.
|
||
|
||
The replacement must be:
|
||
- Complete (ALL garments changed wherever indicated by the outfit and metadata)
|
||
- Exact (copy the outfit items exactly)
|
||
- Hyper-realistic
|
||
- Seamlessly integrated
|
||
- Physically correct
|
||
|
||
-----------------------------------------------------
|
||
👕 OUTFIT SOURCES YOU MUST FOLLOW
|
||
|
||
1) SECOND IMAGE (OUTFIT PHOTO) – PRIMARY VISUAL SOURCE
|
||
- Carefully inspect ALL clothing and accessories.
|
||
- Transfer EVERY clothing item visible in the outfit photo.
|
||
|
||
2) METADATA (DESCRIPTIVE SOURCE AND HARD CHECKLIST)
|
||
The following metadata describes the outfit and items:
|
||
|
||
{metadata_text}
|
||
|
||
{extra_rules}
|
||
|
||
Treat metadata as a strict checklist: any item marked for change MUST be changed.
|
||
|
||
-----------------------------------------------------
|
||
📸 PHOTOREALISM REQUIREMENTS
|
||
|
||
- Maintain original exposure, contrast, and white balance.
|
||
- Preserve original depth of field (DOF) and perspective.
|
||
- Do NOT stylize, cartoonize, or apply filters.
|
||
- Avoid blurring, smoothing, or "AI look".
|
||
- The final render must look like a real photo from the same camera and setup.
|
||
|
||
-----------------------------------------------------
|
||
🚫 FORBIDDEN ACTIONS
|
||
|
||
- Do NOT alter the user's face, identity, or body.
|
||
- Do NOT change background, environment, camera angle, or lighting direction.
|
||
- Do NOT add or remove body parts or tattoos.
|
||
- Do NOT generate a different model.
|
||
- Do NOT copy the second person's face or body.
|
||
- Do NOT output multiple images or side-by-side comparisons.
|
||
- Do NOT output any text in the image.
|
||
- Do NOT ignore the outfit photo or metadata.
|
||
|
||
-----------------------------------------------------
|
||
✅ FINAL OUTPUT REQUIREMENT
|
||
|
||
Generate ONE single high-resolution, hyper-realistic edited image where:
|
||
- The person from the FIRST image is wearing the outfit from the SECOND image.
|
||
- All clothing and accessories are replaced according to the outfit photo and metadata.
|
||
- The person, pose, background, and scene remain identical except for clothes and accessories.
|
||
|
||
FIRST image = user_photo (to edit).
|
||
SECOND image = outfit_photo (outfit reference).
|
||
"""
|
||
|
||
# ---- 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 (image model) ----
|
||
print(f"\n[DEBUG] Using image model: {IMAGE_MODEL}")
|
||
print(f"[DEBUG] Metadata received: {metadata[:200] 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)
|
||
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."
|
||
),
|
||
)
|
||
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}")
|
||
|
||
text = getattr(part, "text", None)
|
||
if text:
|
||
print(f" TEXT: {text[:300]}")
|
||
|
||
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")
|
||
|
||
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):
|
||
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")
|
||
|
||
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
|
||
|
||
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 bytes
|
||
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",
|
||
)
|
||
|
||
# ---- SECOND CALL: IDENTITY VALIDATION ----
|
||
print("[DEBUG] Running identity consistency check...")
|
||
same_identity = faces_match_strict(user_bytes, raw)
|
||
print(f"[DEBUG] Identity check result: same_identity={same_identity}")
|
||
|
||
if not same_identity:
|
||
# Hard fail: we do NOT return the image if identity changed
|
||
raise HTTPException(
|
||
status_code=502,
|
||
detail="Model changed user identity; try-on result rejected by identity guard.",
|
||
)
|
||
|
||
headers = {}
|
||
if request_id:
|
||
headers["X-TryOn-Request-Id"] = str(request_id)
|
||
|
||
return Response(content=raw, media_type="image/png", headers=headers)
|