x

main.py

@@ -10,12 +10,14 @@ 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) ---
+# --- Load .env (for local development only; in Docker envs are already set) ---
 load_dotenv()
 
-# ---- Gemini SDK Setup ----
+# ---- Gemini SDK + Proxy Setup ----
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 
 # Strip whitespace (common issue with .env files)
@@ -25,11 +27,31 @@ if GEMINI_API_KEY:
 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"
+IDENTITY_MODEL = "gemini-2.5-pro"  # ✅ current stable multimodal 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
@@ -58,6 +80,164 @@ def find_first_image_part(resp):
     return None
 
 
+async def generate_tryon_image(user_part, outfit_part, prompt, previous_image_part=None):
+    """
+    Helper function to generate try-on image from Gemini.
+    Returns the raw image bytes.
+    """
+    contents = [prompt, user_part, outfit_part]
+    if previous_image_part:
+        contents.append(previous_image_part)
+        contents.append("Third image = previous incorrect attempt. Fix it by keeping the original face exactly.")
+    else:
+        contents.append("First image = user selfie. Second image = outfit reference.")
+    
+    print(f"[DEBUG] Generating try-on image with model: {IMAGE_MODEL}")
+    if previous_image_part:
+        print("[DEBUG] Correction mode: regenerating with face-lock instructions")
+    
+    try:
+        response = client.models.generate_content(
+            model=IMAGE_MODEL,
+            contents=contents,
+        )
+    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}")
+    
+    # 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",
+        )
+    
+    return raw
+
+
+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.
+    """
+    print(f"[DEBUG] Identity check using model: {IDENTITY_MODEL}")
+    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"}
@@ -105,29 +285,58 @@ async def try_on(
             # Try common key names for the UUID
             request_id = meta.get("request_id") or meta.get("id") or meta.get("uuid")
 
-            # Dynamic behavior based on bottom type
+            extra_sections = []
+
+            # Bottom: shorts
             bottom = meta.get("bottom") or {}
             bottom_type = (bottom.get("type") or "").lower()
-
             if bottom_type == "shorts":
-                # Special rule to force removal of long pants and show legs
-                extra_rules = """
-SPECIAL RULE FOR SHORTS (INSANELY STRICT, DO NOT IGNORE):
+                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.
+  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 (new, strict, merged 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.
 
@@ -137,9 +346,9 @@ You are performing a **strict outfit replacement operation** on the FIRST image
 - No face, body, or identity manipulation is permitted.
 
 -----------------------------------------------------
-🚨 ULTRA-HARD IDENTITY LOCK (ABSOLUTE RULES, DO NOT BREAK)
+🚨 ULTRA-HARD IDENTITY LOCK (ABSOLUTE, NON-NEGOTIABLE RULES)
 
-These identity rules are STRONGER than all other instructions:
+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:
@@ -147,19 +356,31 @@ These identity rules are STRONGER than all other instructions:
   • identity
   • skin tone
   • facial structure
-  • hair style and color
+  • 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.
 
-- 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.
+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**.
+- DO NOT TOUCH the face at all.
+- DO NOT TOUCH identity.
+- Only modify clothing regions.
+- Keep head/face exactly as the original.
+- Do not redraw the face at all.
+- 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
@@ -168,6 +389,7 @@ YOU MUST PROTECT THE FIRST PERSON'S IDENTITY and only adjust clothing areas.
 - 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
@@ -184,108 +406,32 @@ 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
+🎯 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)
+- Complete (ALL garments changed wherever indicated by the outfit and metadata)
 - Exact (copy the outfit items exactly)
 - Hyper-realistic
 - Seamlessly integrated
 - Physically correct
 
-Do NOT leave ANY part of the original outfit visible.
-
------------------------------------------------------
-👤 DO NOT CHANGE ANYTHING ABOUT THE USER (FIRST IMAGE)
-Strictly preserve:
-- Face and identity
-- 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)
+2) METADATA (DESCRIPTIVE SOURCE AND HARD CHECKLIST)
    The following metadata describes the outfit and items:
 
    {metadata_text}
 
    {extra_rules}
 
-   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, use it to refine, but do NOT invent new styles.
-
-3) COMBINED BEHAVIOR
-   - Outfit photo is the main visual truth.
-   - Metadata is the 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.
-
------------------------------------------------------
-🧵 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
@@ -293,17 +439,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.
@@ -330,121 +476,104 @@ 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}")
-    print(f"[DEBUG] Metadata received: {metadata[:200] if metadata else '(empty)'}")
-
+    print(f"\n[DEBUG] Metadata received: {metadata[:200] if metadata else '(empty)'}")
+    
+    # ============================================================
+    # STEP A: Generate initial try-on image (clothing only)
+    # ============================================================
+    print("\n[STEP A] Generating initial try-on image...")
     try:
-        response = client.models.generate_content(
-            model=IMAGE_MODEL,
-            contents=[
-                prompt,
-                user_part,
-                outfit_part,
-                "First image = user selfie. Second image = outfit reference.",
-            ],
-        )
+        raw_v1 = await generate_tryon_image(user_part, outfit_part, prompt)
+        print("[STEP A] ✅ Initial try-on image generated")
+    except HTTPException:
+        raise
     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(f"[ERROR] Step A failed: {e}")
         print(traceback.format_exc())
-        raise HTTPException(status_code=502, detail=f"Gemini error: {error_msg}")
+        raise HTTPException(status_code=502, detail=f"Failed to generate initial try-on: {str(e)}")
+    
+    # ============================================================
+    # STEP B: Identity validation check
+    # ============================================================
+    print("\n[STEP B] Running identity consistency check on initial result...")
+    same_identity = faces_match_strict(user_bytes, raw_v1)
+    print(f"[STEP B] Identity check result: same_identity={same_identity}")
+    
+    if same_identity:
+        # ✅ Identity preserved - return the image
+        print("[STEP B] ✅ Identity validated - returning initial result")
+        headers = {}
+        if request_id:
+            headers["X-TryOn-Request-Id"] = str(request_id)
+        return Response(content=raw_v1, media_type="image/png", headers=headers)
+    
+    # ============================================================
+    # STEP C: Identity changed - regenerate with HARD face lock
+    # ============================================================
+    print("\n[STEP C] ❌ Identity changed detected - regenerating with face-lock correction...")
+    
+    correction_prompt = f"""
+You changed the user's identity in the previous image generation attempt.
+THIS IS ABSOLUTELY FORBIDDEN.
 
-    # ---- 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)}")
+🚫 CRITICAL CORRECTION INSTRUCTIONS:
 
-        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}")
+DO NOT TOUCH THE USER'S FACE AT ALL
+DO NOT REDRAW THE FACE
+DO NOT CHANGE IDENTITY, HAIRLINE, BEARD, OR ANY FACIAL FEATURES
+DO NOT CHANGE HEAD SHAPE OR NECK
+DO NOT BEAUTIFY, SMOOTH, OR RESHAPE THE FACE
+DO NOT MODIFY SKIN TONE ON THE FACE
 
-                text = getattr(part, "text", None)
-                if text:
-                    print(f"    TEXT: {text[:300]}")
+YOU MUST KEEP THE FIRST IMAGE FACE EXACTLY, PIXEL BY PIXEL IF POSSIBLE.
 
-                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")
+ONLY modify clothing areas using the outfit reference from the SECOND image.
+Keep the background, lighting, pose, and identity EXACTLY as in the original user photo (FIRST image).
 
-                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")
+The face, head, ears, and neck region from the FIRST image are READ-ONLY.
+Only clothing below the neck and on the body should be changed.
 
-    # ---- 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."
+{metadata_text}
 
-        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
+{extra_rules}
+"""
+    
+    # Create a Part from the previous incorrect image for reference
+    previous_image_part = Part.from_bytes(
+        data=raw_v1,
+        mime_type="image/png",
+    )
+    
     try:
-        Image.open(BytesIO(raw))
+        raw_v2 = await generate_tryon_image(user_part, outfit_part, correction_prompt, previous_image_part)
+        print("[STEP C] ✅ Correction image generated")
+    except HTTPException:
+        raise
     except Exception as e:
+        import traceback
+        print(f"[ERROR] Step C failed: {e}")
+        print(traceback.format_exc())
         raise HTTPException(
-            status_code=500,
-            detail=f"Returned image is invalid: {e}. Image size: {len(raw)} bytes",
+            status_code=502,
+            detail="Failed to generate corrected try-on after identity mismatch.",
         )
-
+    
+    # ============================================================
+    # STEP D: Re-validate corrected version (optional check, always return)
+    # ============================================================
+    print("\n[STEP D] Running identity consistency check on corrected result...")
+    same_identity_v2 = faces_match_strict(user_bytes, raw_v2)
+    print(f"[STEP D] Identity check result: same_identity={same_identity_v2}")
+    
+    if same_identity_v2:
+        print("[STEP D] ✅ Identity validated after correction")
+    else:
+        print("[STEP D] ⚠️ Identity check failed, but returning corrected image anyway")
+    
+    # Always return the corrected image, regardless of identity check result
     headers = {}
     if request_id:
         headers["X-TryOn-Request-Id"] = str(request_id)
-
-    return Response(content=raw, media_type="image/png", headers=headers)
+    
+    return Response(content=raw_v2, media_type="image/png", headers=headers)

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
+

requirements.txt

@@ -4,4 +4,4 @@ google-genai
 Pillow
 python-multipart
 python-dotenv
-
+httpx[socks]