FastApi-GeminiApi/main.py
2025-12-01 15:17:48 +03:30

580 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import os
import json
from io import BytesIO
from pprint import pprint
from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
from fastapi.responses import Response
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 only; in Docker envs are already set) ---
load_dotenv()
# ---- Gemini SDK + Proxy Setup ----
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
# Strip whitespace (common issue with .env files)
if GEMINI_API_KEY:
GEMINI_API_KEY = GEMINI_API_KEY.strip()
if not GEMINI_API_KEY:
raise RuntimeError("GEMINI_API_KEY not set")
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
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # tighten in production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
def find_first_image_part(resp):
"""Return the actual output image part."""
if not getattr(resp, "candidates", None):
return None
for cand in resp.candidates:
content = getattr(cand, "content", None)
if not content or not getattr(content, "parts", None):
continue
for p in content.parts:
if hasattr(p, "inline_data") and getattr(p.inline_data, "data", None):
return p.inline_data
if hasattr(p, "blob") and getattr(p.blob, "data", None):
return p.blob
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"}
@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"),
):
"""
metadata can be:
- free text: 'black hoodie, patterned shirt, white 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()
if not user_bytes:
raise HTTPException(status_code=400, detail="user_photo is empty")
if not outfit_bytes:
raise HTTPException(status_code=400, detail="outfit_photo is empty")
# ---- Parse metadata (JSON or plain text) + extract request_id + dynamic rules ----
metadata_text = ""
request_id = None
extra_rules = ""
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")
extra_sections = []
# Bottom: shorts
bottom = meta.get("bottom") or {}
bottom_type = (bottom.get("type") or "").lower()
if bottom_type == "shorts":
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 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 (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.
⚠️ IMPORTANT SAFETY CONTEXT
- 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.
-----------------------------------------------------
🚨 ULTRA-HARD IDENTITY LOCK (ABSOLUTE, NON-NEGOTIABLE RULES)
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:
• face
• identity
• skin tone
• facial structure
• 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.
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
- 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 moustache or beard
• their skin tone
• their hair
• their body shape
• their pose
• their tattoos or body details
You may ONLY use the SECOND image to extract:
• clothing items,
• shoes,
• accessories (including the glasses holder strap around the neck),
• 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 (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 wherever indicated by the outfit and metadata)
- Exact (copy the outfit items exactly)
- Hyper-realistic
- Seamlessly integrated
- Physically correct
-----------------------------------------------------
👕 OUTFIT SOURCES YOU MUST FOLLOW
1) SECOND IMAGE (OUTFIT PHOTO) PRIMARY VISUAL SOURCE
- Carefully inspect ALL clothing and accessories.
- Transfer EVERY clothing item visible in the outfit photo.
2) METADATA (DESCRIPTIVE SOURCE AND HARD CHECKLIST)
The following metadata describes the outfit and items:
{metadata_text}
{extra_rules}
Treat metadata as a strict checklist: any item marked for change MUST be changed.
-----------------------------------------------------
📸 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",
)
outfit_part = Part.from_bytes(
data=outfit_bytes,
mime_type=outfit_photo.content_type or "image/jpeg",
)
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:
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:
import traceback
print(f"[ERROR] Step A failed: {e}")
print(traceback.format_exc())
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.
🚫 CRITICAL CORRECTION INSTRUCTIONS:
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
YOU MUST KEEP THE FIRST IMAGE FACE EXACTLY, PIXEL BY PIXEL IF POSSIBLE.
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).
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.
{metadata_text}
{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:
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=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_v2, media_type="image/png", headers=headers)