x

main.py

@@ -80,6 +80,93 @@ 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
@@ -281,6 +368,11 @@ These identity rules are STRONGER than all other instructions. If you cannot fol
 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,
@@ -384,133 +476,104 @@ SECOND image = outfit_photo (outfit reference).
         mime_type=outfit_photo.content_type or "image/jpeg",
     )
 
-    # ---- Call Gemini (image model) ----
+    print(f"\n[DEBUG] Metadata received: {metadata[:200] if metadata else '(empty)'}")
-    print(f"\n[DEBUG] Using image model: {IMAGE_MODEL}")
+    
-    print(f"[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(
+        raw_v1 = await generate_tryon_image(user_part, outfit_part, prompt)
-            model=IMAGE_MODEL,
+        print("[STEP A] ✅ Initial try-on image generated")
-            contents=[
+    except HTTPException:
-                prompt,
+        raise
-                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(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 ----
+🚫 CRITICAL CORRECTION INSTRUCTIONS:
-    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):
+DO NOT TOUCH THE USER'S FACE AT ALL
-            print(f"[DEBUG] Parts in first candidate ({len(cand.content.parts)} total):")
+DO NOT REDRAW THE FACE
-            for i, part in enumerate(cand.content.parts):
+DO NOT CHANGE IDENTITY, HAIRLINE, BEARD, OR ANY FACIAL FEATURES
-                part_type = type(part).__name__
+DO NOT CHANGE HEAD SHAPE OR NECK
-                print(f"  Part {i}: type={part_type}")
+DO NOT BEAUTIFY, SMOOTH, OR RESHAPE THE FACE
+DO NOT MODIFY SKIN TONE ON THE FACE
 
-                text = getattr(part, "text", None)
+YOU MUST KEEP THE FIRST IMAGE FACE EXACTLY, PIXEL BY PIXEL IF POSSIBLE.
-                if text:
-                    print(f"    TEXT: {text[:300]}")
 
-                if hasattr(part, "inline_data"):
+ONLY modify clothing areas using the outfit reference from the SECOND image.
-                    inline = part.inline_data
+Keep the background, lighting, pose, and identity EXACTLY as in the original user photo (FIRST image).
-                    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"):
+The face, head, ears, and neck region from the FIRST image are READ-ONLY.
-                    blob = part.blob
+Only clothing below the neck and on the body should be changed.
-                    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 ----
+{metadata_text}
-    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:
+{extra_rules}
-            cand = response.candidates[0]
+"""
-            finish_reason = getattr(cand, "finish_reason", "unknown")
+    
-
+    # Create a Part from the previous incorrect image for reference
-            if hasattr(cand, "content") and getattr(cand.content, "parts", None):
+    previous_image_part = Part.from_bytes(
-                for part in cand.content.parts:
+        data=raw_v1,
-                    txt = getattr(part, "text", None)
+        mime_type="image/png",
-                    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))
+        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:
-        raise HTTPException(
+        import traceback
-            status_code=500,
+        print(f"[ERROR] Step C failed: {e}")
-            detail=f"Returned image is invalid: {e}. Image size: {len(raw)} bytes",
+        print(traceback.format_exc())
-        )
-
-    # ---- 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.",
+            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)