diff --git a/main.py b/main.py index 9dc3835..4c9a2a3 100644 --- a/main.py +++ b/main.py @@ -28,6 +28,7 @@ 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") @@ -79,6 +80,76 @@ def find_first_image_part(resp): 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"} @@ -102,9 +173,7 @@ async def try_on( 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" - ) + raise HTTPException(status_code=401, detail="Invalid or missing authentication key") # ---- Read files ---- user_bytes = await user_photo.read() @@ -126,17 +195,13 @@ async def try_on( 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") - ) + request_id = meta.get("request_id") or meta.get("id") or meta.get("uuid") - # ---- Dynamic behavior based on bottom/top type ---- extra_sections = [] # Bottom: shorts bottom = meta.get("bottom") or {} bottom_type = (bottom.get("type") or "").lower() - if bottom_type == "shorts": extra_sections.append( """ @@ -155,18 +220,9 @@ SPECIAL RULE FOR SHORTS (INSANELY STRICT, ZERO TOLERANCE – DO NOT IGNORE): ) # Top: t-shirt / tee / short-sleeve - top = ( - meta.get("top") - or meta.get("upper") - or meta.get("outer_layer") - or {} - ) + 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"] - ): + 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): @@ -192,7 +248,7 @@ SPECIAL RULE FOR T-SHIRT / TOP (EXTREMELY STRICT, ZERO TOLERANCE): except Exception: metadata_text = metadata.strip() - # ---- Build Prompt (new, ultra strict version) ---- + # ---- 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. @@ -217,7 +273,6 @@ These identity rules are STRONGER than all other instructions. If you cannot fol • lips and mouth shape • facial hair (beard / moustache) style and density • hair style and hairline - • ears • body proportions • pose EXACTLY as they are. @@ -227,33 +282,9 @@ 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. -- You MUST NOT change: - • beard thickness or shape - • moustache presence or shape - • hairline or hairstyle - • skin texture or tone on the face - • teeth alignment or smile shape - 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). -IDENTITY FAILURE CONDITION: - -- If completing the clothing replacement would require you to change the user's face, head, - or identity in any way, you MUST NOT generate an image. -- In that situation, you must respond with TEXT ONLY explaining that identity rules prevented - you from generating an image. -- You are not allowed to "approximate" the face. Either: - • keep the original face exactly, OR - • do not output an image at all. - -- You are **FORBIDDEN** to: - • change the face or body of the person in the FIRST image, - • beautify, smooth, reshape, de-age, re-gender, or re-style the face, - • swap the person with someone else, - • blend or mix the FIRST and SECOND person, - • recreate a new face that only "resembles" the original, - • copy the second person's moustache, beard, jawline, nose, or hairstyle. - 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. @@ -293,47 +324,12 @@ The replacement must be: - Seamlessly integrated - Physically correct -If you leave **ANY** part of the original outfit visible in an area that should be replaced, -the result is considered **WRONG and FAILED**. - ------------------------------------------------------ -👤 DO NOT CHANGE ANYTHING ABOUT THE USER (FIRST IMAGE) -Strictly preserve: -- Face and identity (see ABSOLUTE FACE LOCK above) -- Skin tone -- Hair -- Expression -- Body proportions -- Pose -- Hands -- Phone -- Non-clothing jewelry -- Background and environment -- Lighting and shadows -- Camera angle, framing, and composition -- Depth of field and bokeh - -The edited image must look like the same person photographed in the same moment, but wearing the uploaded outfit. - ----------------------------------------------------- 👕 OUTFIT SOURCES YOU MUST FOLLOW -You MUST use BOTH sources: - 1) SECOND IMAGE (OUTFIT PHOTO) – PRIMARY VISUAL SOURCE - Carefully inspect ALL clothing and accessories. - - Transfer EVERY clothing item visible in the outfit photo to the user: - • Outer layers (jackets, hoodies, coats) - • Inner layers (shirts, t-shirts, tops) - • Bottoms (pants, shorts, skirts) - • Footwear - • Accessories (hats, glasses, glasses holder strap around the neck, necklaces, etc.) - - Copy exactly: - • Colors and hues (no tinting) - • Patterns and their scale (no warping or stretching) - • Textures, weave, and material type - • Stitching, seams, hems, cuffs, collars, zippers, buttons, pockets, labels, prints, logos, trims - • Shape, silhouette, and fabric density + - Transfer EVERY clothing item visible in the outfit photo. 2) METADATA (DESCRIPTIVE SOURCE AND HARD CHECKLIST) The following metadata describes the outfit and items: @@ -342,57 +338,7 @@ You MUST use BOTH sources: {extra_rules} - You MUST treat the metadata as a **STRICT CHECKLIST**: - - Every clothing item described in metadata that differs from the FIRST image MUST be updated in the final image. - - Do NOT skip or ignore ANY listed item (top, bottom, shoes, accessories, straps, etc.). - - If something is mentioned in metadata for replacement and you leave the original visible, - the result is **INCORRECT**. - - Use metadata to: - • Confirm all clothing layers and components. - • Ensure no item mentioned in metadata is forgotten. - • Validate colors, patterns, and materials. - • Resolve ambiguous details from the outfit photo. - - If metadata mentions items not clearly visible, prefer the metadata description, - but do NOT invent a style that contradicts the outfit photo. - -3) COMBINED BEHAVIOR - - Outfit photo is the main visual truth. - - Metadata is the strict checklist and description to avoid missing items. - - The final result MUST: - • Include ALL clothing and accessories visible in the outfit photo. - • Include ALL relevant items mentioned in metadata that are consistent with the outfit photo. - • Remove the original garments completely wherever replacement is implied. - ------------------------------------------------------ -🧵 STRICT CLOTHING REPLACEMENT RULES - -1. REMOVE original clothing completely from the FIRST image before placing replacements. - - No underlayers or ghost fabrics. - - No visible seams, colors, or textures from the old outfit. - -2. APPLY the new outfit exactly as shown in the SECOND image: - - Perfect color match (no hue or saturation shift). - - Original pattern scale and orientation. - - Correct sleeve length, collar height, garment length, and silhouette. - - Only scale garments enough to fit the body while preserving pattern proportions. - -3. NATURAL, PHYSICS-AWARE FIT: - - Fabric drapes according to gravity and the model’s pose. - - Folds and creases look realistic but do NOT distort patterns or design. - -4. PERFECT INTEGRATION: - - Lighting and reflections must match the original scene. - - Contact shadows and occlusion must be consistent. - - No haloing, floating edges, or mismatched shadow direction. - -5. ACCESSORIES: - - Apply all visible accessories from the outfit photo. - - If metadata mentions a glasses holder strap around the neck, it must appear naturally placed - according to anatomy and perspective. - -6. NO NEW DESIGN ELEMENTS: - - Do NOT invent new zippers, seams, logos, or prints. - - Do NOT remove real logos or prints that exist on the uploaded outfit unless clearly obstructed. + Treat metadata as a strict checklist: any item marked for change MUST be changed. ----------------------------------------------------- 📸 PHOTOREALISM REQUIREMENTS @@ -400,17 +346,17 @@ You MUST use BOTH sources: - 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”. +- 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 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 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. @@ -423,9 +369,6 @@ Generate ONE single high-resolution, hyper-realistic edited image where: - 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. -If you cannot satisfy the IDENTITY LOCK and ABSOLUTE FACE LOCK rules, -you MUST NOT generate an image and should respond with TEXT ONLY explaining why. - FIRST image = user_photo (to edit). SECOND image = outfit_photo (outfit reference). """ @@ -440,8 +383,8 @@ SECOND image = outfit_photo (outfit reference). mime_type=outfit_photo.content_type or "image/jpeg", ) - # ---- Call Gemini ---- - print(f"\n[DEBUG] Using model: {IMAGE_MODEL}") + # ---- 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: @@ -456,11 +399,7 @@ SECOND image = outfit_photo (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 - ): + if "API key" in error_msg or "INVALID_ARGUMENT" in error_msg or "API_KEY" in error_msg: raise HTTPException( status_code=401, detail=( @@ -469,7 +408,6 @@ SECOND image = outfit_photo (outfit reference). ), ) 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}") @@ -486,9 +424,7 @@ SECOND image = outfit_photo (outfit reference). 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):" - ) + 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}") @@ -502,9 +438,7 @@ SECOND image = outfit_photo (outfit reference). 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" - ) + print(f" INLINE_DATA: mime={mime}, size={data_size} bytes") if hasattr(part, "blob"): blob = part.blob @@ -553,7 +487,7 @@ SECOND image = outfit_photo (outfit reference). raw = img_part.data - # Ensure valid image + # Ensure valid image bytes try: Image.open(BytesIO(raw)) except Exception as e: @@ -562,6 +496,18 @@ SECOND image = outfit_photo (outfit reference). 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) diff --git a/prompts/identity_check_prompt.txt b/prompts/identity_check_prompt.txt new file mode 100644 index 0000000..e9a5bb2 --- /dev/null +++ b/prompts/identity_check_prompt.txt @@ -0,0 +1,30 @@ +You are an identity-protection validator. + +You receive two images: + +1. original_user: the original human photo from the user (this identity must NEVER change) + +2. generated_tryon: the try-on output image generated by another model + +Your job: +- Compare faces +- Compare facial structure, beard, hairline, skin tone, nose shape, jaw shape +- Check if the generated image kept the SAME identity +- Detect if the AI replaced the face with a model +- Detect if the AI adjusted or beautified the face too much +- Detect if body proportions changed unnaturally + +Return STRICT JSON ONLY: + +{ + "identity_match_score": 0–100, + "identity_changed": true/false, + "reason": "string explaining mismatch", + "safe_to_use": true/false +} + +Rules: +- If identity_match_score < 85 → identity_changed must be true +- If anything seems suspicious → safe_to_use must be false +- No extra text, no explanations outside the JSON +