x
This commit is contained in:
parent
0f00b8ffbd
commit
60f86db416
285
main.py
285
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)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user