Update app.py

app.py

@@ -4,18 +4,18 @@ from transformers import pipeline
 from PIL import Image
 import io, os, traceback
 
-# Make sure Hugging Face cache is writable
+# Cache settings
 os.environ["HF_HOME"] = "/app/cache"
 os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
 os.environ["HF_HUB_CACHE"] = "/app/cache/hub"
 
 app = FastAPI()
 
-# Load SmolVLM with the pipeline API
+# Load SmolVLM as pure image-to-text
 pipe = pipeline(
     "image-to-text",
     model="HuggingFaceTB/SmolVLM-256M-Instruct",
-    device=-1  # CPU for free tier
+    device=-1  # CPU
 )
 
 @app.get("/")
@@ -45,16 +45,22 @@ async def predict_gender(file: UploadFile = File(...)):
         image_bytes = await file.read()
         image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
 
-        # Instruction for the model
-        prompt = "Is the person on this ID male or female?"
+        # Run the model (image only)
+        result = pipe(image, max_new_tokens=32)
+        caption = result[0]["generated_text"].strip()
 
-        # Run model (pipeline handles image + prompt via generate_kwargs)
-        result = pipe(image, generate_kwargs={"max_new_tokens": 32, "prompt": prompt})
+        # Very simple heuristic: check for male/female in caption
+        gender = "unknown"
+        lower_caption = caption.lower()
+        if "male" in lower_caption or "man" in lower_caption:
+            gender = "male"
+        elif "female" in lower_caption or "woman" in lower_caption:
+            gender = "female"
 
-        # Extract model output
-        answer = result[0]["generated_text"].strip()
-
-        return JSONResponse({"gender": answer})
+        return JSONResponse({
+            "caption": caption,
+            "gender": gender
+        })
 
     except Exception as e:
         traceback.print_exc()