change a alot

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(
+            raise HTTPException(status_code=401, detail="Invalid or missing authentication key")
-                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 = (
+            request_id = meta.get("request_id") or meta.get("id") or meta.get("uuid")
-                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 = (
+            top = meta.get("top") or meta.get("inner_layer") or meta.get("upper") or meta.get("outer_layer") or {}
-                meta.get("top")
-                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:
+   - Transfer EVERY clothing item visible in the outfit photo.
-     • 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
 
 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**:
+   Treat metadata as a strict checklist: any item marked for change MUST be changed.
-   - 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.
 
 -----------------------------------------------------
 📸 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 ----
+    # ---- Call Gemini (image model) ----
-    print(f"\n[DEBUG] Using model: {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 (
+        if "API key" in error_msg or "INVALID_ARGUMENT" in error_msg or "API_KEY" in error_msg:
-            "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(
+            print(f"[DEBUG] Parts in first candidate ({len(cand.content.parts)} total):")
-                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(
+                        print(f"    INLINE_DATA: mime={mime}, size={data_size} bytes")
-                            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)

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
+