472 lines
17 KiB
Python
472 lines
17 KiB
Python
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"
|
||
|
||
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
|
||
|
||
|
||
@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")
|
||
|
||
# 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"""
|
||
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 RULES, DO NOT BREAK)
|
||
|
||
These identity rules are STRONGER than all other instructions:
|
||
|
||
- 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.
|
||
|
||
- You are **FORBIDDEN** to:
|
||
• change the face or body of the person in the FIRST image,
|
||
• beautify, smooth, reshape, de-age, re-gender, or re-style the face,
|
||
• swap the person with someone else,
|
||
• blend or mix the FIRST and SECOND person.
|
||
|
||
If there is ANY conflict between clothing instructions and identity protection:
|
||
YOU MUST PROTECT THE FIRST PERSON'S IDENTITY and only adjust clothing areas.
|
||
|
||
-----------------------------------------------------
|
||
👤 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 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
|
||
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",
|
||
)
|
||
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[:200] if metadata else '(empty)'}")
|
||
|
||
try:
|
||
response = client.models.generate_content(
|
||
model=IMAGE_MODEL,
|
||
contents=[
|
||
prompt,
|
||
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(traceback.format_exc())
|
||
raise HTTPException(status_code=502, detail=f"Gemini error: {error_msg}")
|
||
|
||
# ---- 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)}")
|
||
|
||
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}")
|
||
|
||
text = getattr(part, "text", None)
|
||
if text:
|
||
print(f" TEXT: {text[:300]}")
|
||
|
||
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")
|
||
|
||
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")
|
||
|
||
# ---- 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
|
||
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",
|
||
)
|
||
|
||
headers = {}
|
||
if request_id:
|
||
headers["X-TryOn-Request-Id"] = str(request_id)
|
||
|
||
return Response(content=raw, media_type="image/png", headers=headers)
|