# from multiprocessing import process

# import torch 
# import gradio as gr 
# from PIL import Image
# from transformers import AutoModel, AutoImageProcessor
# import torch.nn.functional as f

# MODEL_NAME = "ht_hub:gaunernst/vit_tiny_patch8_112.arcface_ms1mv3"

# processor = AutoImageProcessor.from_pretrained(MODEL_NAME)

# model = AutoModel.from_pretrained(processor)

# model.eval()

# def get_embeddings(image):
#     image = image.convert("RGB")
#     inputs = processor(
#         images=image,
#         return_tensor= "pt"
#     )

#     with torch.no_grad():
#         outputs = model(**inputs)
#     embedding = outputs.last_hidden_state.mean(dim=1)
#     embedding = f.normalize(embedding,p=2,dim=1)
#     return embedding

# def compare_faces(image1, image2):
#     face1 = get_embeddings(image1)
#     face2 = get_embeddings(image2)

#     similarity = f.cosine_similarity(face1,face2).item()

#     if similarity >= 0.6:
#         result = "Face approved"
#     else:
#         result = "Face card declined"
#     return result,round(similarity,4)

# demo = gr.Interface(
#     fn=compare_faces,
#     inputs=[
#         gr.Image(type="pil",label="Upload Your first face"),
#         gr.Image(type="pil",label="Upload Your second face")],
#     outputs=[gr.Textbox(label="Result"),gr.Number(label="The Similarity")],
#     title = "Face recognition App",
#     description="This is a Face recognition app"
# )

# demo.launch()




import torch
import timm
import gradio as gr
from PIL import Image
import torch.nn.functional as F

MODEL_NAME = "hf_hub:gaunernst/vit_tiny_patch8_112.arcface_ms1mv3"

model = timm.create_model(MODEL_NAME, pretrained=True)
model.eval()

data_config = timm.data.resolve_model_data_config(model)
transform = timm.data.create_transform(**data_config, is_training=False)

def get_embeddings(image):
    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0)  # add batch dimension

    with torch.no_grad():
        embedding = model(input_tensor)

    embedding = F.normalize(embedding, p=2, dim=1)
    return embedding

def compare_faces(image1, image2):
    face1 = get_embeddings(image1)
    face2 = get_embeddings(image2)

    similarity = F.cosine_similarity(face1, face2).item()

    if similarity >= 0.6:
        result = "Face approved"
    else:
        result = "Face card declined"
    return result, round(similarity, 4)

demo = gr.Interface(
    fn=compare_faces,
    inputs=[
        gr.Image(type="pil", label="Upload Your first face"),
        gr.Image(type="pil", label="Upload Your second face")],
    outputs=[gr.Textbox(label="Result"), gr.Number(label="The Similarity")],
    title="Face recognition App",
    description="This is a Face recognition app"
)

demo.launch()