done

main.py

@@ -3,7 +3,7 @@ import json
 from io import BytesIO
 from pprint import pprint
 
-from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
 from fastapi.responses import Response
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -28,6 +28,8 @@ if not GEMINI_API_KEY:
 client = genai.Client(api_key=GEMINI_API_KEY)
 IMAGE_MODEL = "gemini-2.5-flash-image"
 
+AUTH_KEY = os.getenv("AIFIT_AUTH_KEY")
+
 app = FastAPI()
 
 # (Optional) CORS if you call it from frontend directly
@@ -63,6 +65,7 @@ def health():
 
 @app.post("/api/try-on/", summary="Try-on (user_photo + outfit_photo + metadata)")
 async def try_on(
+    request: Request,
     user_photo: UploadFile = File(...),
     outfit_photo: UploadFile = File(...),
     metadata: str = Form("", description="JSON or text describing the outfit items"),
@@ -70,10 +73,16 @@ async def try_on(
     """
     metadata can be:
     - free text: 'black hoodie, patterned shirt, white shorts'
-    - JSON: '{"outer": "hoodie", "inner": "shirt", "bottom": "shorts"}'
+    - JSON: '{"request_id": "uuid", "outer_layer": {...}, "bottom": {"type": "shorts"}}'
     - any additional instructions
     """
 
+    # ---- AUTH CHECK (simple shared key) ----
+    if AUTH_KEY:
+        client_key = request.headers.get("X-AIFIT-Key")
+        if not client_key or client_key != AUTH_KEY:
+            raise HTTPException(status_code=401, detail="Invalid or missing authentication key")
+
     # ---- Read files ----
     user_bytes = await user_photo.read()
     outfit_bytes = await outfit_photo.read()
@@ -83,91 +92,248 @@ async def try_on(
     if not outfit_bytes:
         raise HTTPException(status_code=400, detail="outfit_photo is empty")
 
-    # ---- Parse metadata if JSON ----
+    # ---- Parse metadata (JSON or plain text) + extract request_id + dynamic rules ----
     metadata_text = ""
-    try:
+    request_id = None
-        meta = json.loads(metadata)
+    extra_rules = ""
-        metadata_text = json.dumps(meta, indent=2)
-    except:
-        metadata_text = metadata.strip()
 
-    # ---- Build Prompt ----
+    if metadata:
+        try:
+            meta = json.loads(metadata)
+            metadata_text = json.dumps(meta, indent=2, ensure_ascii=False)
+
+            # 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
+            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):
+
+- 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.
+- 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.
+"""
+        except Exception:
+            metadata_text = metadata.strip()
+
+    # ---- Build Prompt (new, strict, merged version) ----
     prompt = f"""
-EDIT TASK — FULL OUTFIT REPLACEMENT
+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.
 
-FIRST image = user body/mirror selfie
+⚠️ IMPORTANT SAFETY CONTEXT
-SECOND image = outfit reference
+- The individuals in the images are **not public figures**.
+- The task is standard, allowed **fashion editing / virtual try-on** only.
+- No face, body, or identity manipulation is permitted.
 
-Follow these rules:
+-----------------------------------------------------
+🚨 ULTRA-HARD IDENTITY LOCK (ABSOLUTE RULES, DO NOT BREAK)
 
-1. Edit ONLY the body/clothing of the FIRST image.
+These identity rules are STRONGER than all other instructions:
 
-2. Replace ALL clothing with the items visible in the SECOND image (outfit photo).
+- 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:
+  • face
+  • identity
+  • skin tone
+  • facial structure
+  • hair style and color
+  • body proportions
+  • pose
+  EXACTLY as they are.
 
-3. Keep the user's:
+- You are **FORBIDDEN** to:
-   - face
+  • change the face or body of the person in the FIRST image,
-   - skin
+  • beautify, smooth, reshape, de-age, re-gender, or re-style the face,
-   - hands
+  • swap the person with someone else,
-   - phone
+  • blend or mix the FIRST and SECOND person.
-   - pose
-   - environment
-   EXACTLY the same.
 
-4. OUTFIT PHOTO AND METADATA CONSIDERATION (CRITICAL):
+If there is ANY conflict between clothing instructions and identity protection:
-   
+YOU MUST PROTECT THE FIRST PERSON'S IDENTITY and only adjust clothing areas.
-   You MUST consider BOTH sources:
-   
-   A) SECOND IMAGE (OUTFIT PHOTO):
-      - Carefully examine ALL clothing items visible in the outfit photo
-      - Transfer EVERY piece of clothing from the outfit photo to the user photo
-      - Pay attention to the visual details: colors, patterns, textures, styles, fit, and layering
-      - Ensure all visible items (tops, bottoms, outerwear, accessories, etc.) are transferred
-   
-   B) METADATA:
-      The following metadata provides additional information about the outfit:
-      
-      Metadata:
-      {metadata_text}
-      
-      Use the metadata to:
-      - Identify and verify ALL items mentioned in the metadata are present in the outfit photo
-      - Understand the complete outfit composition and layering structure
-      - Ensure no items are missed - cross-reference metadata items with the outfit photo
-      - Apply any specific details from metadata (colors, patterns, materials, styles) that match the outfit photo
-      - If metadata mentions items not clearly visible in the photo, use metadata details to enhance accuracy
-   
-   C) COMBINE BOTH SOURCES:
-      - Use the outfit photo as the primary visual reference
-      - Use metadata to ensure completeness and accuracy
-      - Transfer ALL items that appear in the outfit photo AND are mentioned in metadata
-      - The final result should include every clothing item visible in the outfit photo, verified against the metadata
 
-5. Ensure:
+-----------------------------------------------------
-   - realistic lighting
+👤 SECOND IMAGE PERSON HANDLING — CLOTHES ONLY, NEVER THEIR IDENTITY
-   - correct fabric physics
-   - correct drape
-   - realism
-   - no blending of people
-   - no copying faces from the second image
-   - all items from both the outfit photo and metadata are properly integrated into the user photo
 
-6. Output ONE final edited image with ALL items from the outfit photo transferred to the user, ensuring all metadata items are also included.
+- The person in the SECOND image (outfit photo) is **NOT** allowed to appear in the final result.
+- You must treat the SECOND person as a "clothing mannequin" ONLY.
+- You are **STRICTLY FORBIDDEN** from copying:
+  • their face
+  • their skin tone
+  • their hair
+  • their body shape
+  • their pose
+  • their tattoos or body details
 
-IMPORTANT:
+You may ONLY use the SECOND image to extract:
-- FIRST = user_photo
+  • clothing items,
-- SECOND = outfit_photo
+  • shoes,
-- You must use BOTH the outfit photo (visual reference) AND metadata (descriptive reference)
+  • accessories (including the glasses holder strap around the neck),
-- ALL items visible in the outfit photo AND mentioned in metadata must appear in the final image
+  • their colors, patterns, materials, and shapes.
+
+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
+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)
+- 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
+
+2) METADATA (DESCRIPTIVE SOURCE)
+   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.
+
+-----------------------------------------------------
+📸 PHOTOREALISM REQUIREMENTS
+
+- 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”.
+- 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 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 output multiple images or side-by-side comparisons.
+- Do NOT output any text in the image.
+- Do NOT ignore the outfit photo or metadata.
+
+-----------------------------------------------------
+✅ FINAL OUTPUT REQUIREMENT
+
+Generate ONE single high-resolution, hyper-realistic edited image where:
+- The person from the FIRST image is wearing the outfit from the SECOND image.
+- 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.
+
+FIRST image = user_photo (to edit).
+SECOND image = outfit_photo (outfit reference).
 """
 
     # ---- Wrap image files ----
-    user_part = Part.from_bytes(data=user_bytes, mime_type=user_photo.content_type or "image/jpeg")
+    user_part = Part.from_bytes(
-    outfit_part = Part.from_bytes(data=outfit_bytes, mime_type=outfit_photo.content_type or "image/jpeg")
+        data=user_bytes,
+        mime_type=user_photo.content_type or "image/jpeg",
+    )
+    outfit_part = Part.from_bytes(
+        data=outfit_bytes,
+        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[:100] if metadata else '(empty)'}")
+    print(f"[DEBUG] Metadata received: {metadata[:200] if metadata else '(empty)'}")
-    
+
     try:
         response = client.models.generate_content(
             model=IMAGE_MODEL,
@@ -180,16 +346,14 @@ IMPORTANT:
         )
     except Exception as e:
         error_msg = str(e)
-        # Check if it's an API key error
         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."
-                )
+                ),
             )
-        # Log the full error for debugging
         import traceback
         print(f"[ERROR] Gemini API error: {error_msg}")
         print(traceback.format_exc())
@@ -205,27 +369,24 @@ IMPORTANT:
         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):
             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}")
-                
+
-                # Check for text
                 text = getattr(part, "text", None)
                 if text:
                     print(f"    TEXT: {text[:300]}")
-                
+
-                # Check for inline_data
                 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")
-                
+
-                # Check for blob
                 if hasattr(part, "blob"):
                     blob = part.blob
                     if blob:
@@ -237,33 +398,30 @@ IMPORTANT:
     # ---- Extract output image ----
     img_part = find_first_image_part(response)
     if not img_part or not getattr(img_part, "data", None):
-        # Try to extract any text message from Gemini
         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")
-            
+
-            # Extract text explanation from Gemini
             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
-            
+
-            # Check safety ratings
             if hasattr(cand, "safety_ratings"):
                 safety_ratings = cand.safety_ratings
                 if safety_ratings:
                     blocked = [
-                        f"{getattr(r, 'category', 'unknown')}" 
+                        f"{getattr(r, 'category', 'unknown')}"
-                        for r in safety_ratings 
+                        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}",
@@ -281,8 +439,12 @@ IMPORTANT:
         Image.open(BytesIO(raw))
     except Exception as e:
         raise HTTPException(
-            status_code=500, 
+            status_code=500,
-            detail=f"Returned image is invalid: {e}. Image size: {len(raw)} bytes"
+            detail=f"Returned image is invalid: {e}. Image size: {len(raw)} bytes",
         )
 
-    return Response(content=raw, media_type="image/png")
+    headers = {}
+    if request_id:
+        headers["X-TryOn-Request-Id"] = str(request_id)
+
+    return Response(content=raw, media_type="image/png", headers=headers)