Update model.py

model.py

@@ -31,4 +31,37 @@ class AgeGenderViTModel(ViTPreTrainedModel):
         gender_output = self.gender_head(pooled_output)
         
         logits = torch.cat([age_output, gender_output], dim=1)
-        return {"logits": logits}
+        return {"logits": logits}
+
+    # Add this to the END of your model.py file
+
+    def predict_age_gender(image_path):
+        """
+        Simple one-liner function for age-gender prediction
+        
+        Args:
+            image_path: Path to image file or URL
+            
+        Returns:
+            Dictionary with age, gender, confidence
+        """
+        from transformers import pipeline
+        
+        classifier = pipeline("image-classification", model="abhilash88/age-gender-prediction", trust_remote_code=True)
+        raw = classifier(image_path)
+        result = raw[0]  # Get first result
+        
+        return {
+            'age': result['age'], 
+            'gender': result['gender'],
+            'confidence': result['gender_confidence'],
+            'summary': f"{result['age']} years, {result['gender']} ({result['gender_confidence']:.1%} confidence)"
+        }
+    
+    
+    def simple_predict(image_path):
+        """
+        Even simpler - just returns a string
+        """
+        result = predict_age_gender(image_path)
+        return result['summary']