diff --git a/main.py b/main.py index f918ed5..e4f4519 100644 --- a/main.py +++ b/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] Using image 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 bytes +{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: - 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 + import traceback + print(f"[ERROR] Step C failed: {e}") + print(traceback.format_exc()) 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)