fix 7

app.py

@@ -1,16 +1,33 @@
 import gradio as gr
-from fastai.vision.all import *
+from transformers import AutoModelForImageClassification, AutoTokenizer
+from PIL import Image
 
+# Load Hugging Face model and tokenizer
+model_name = 'ImageDifferentiator'  # Replace with the specific model name
+model = AutoModelForImageClassification.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
+# Get class labels from the model configuration
+labels = model.config.id2label
 
-learn = load_learner('ImageDifferentiator.pkl')
-
-labels = learn.dls.vocab
 def predict(img):
-    
-    img = PILImage.create(img)
-    pred, pred_idx, probs = learn.predict(img)
+    # Tokenize and preprocess the image
+    inputs = tokenizer(img, return_tensors="pt")
+
+    # Make prediction using the Hugging Face model
+    outputs = model(**inputs)
+    logits = outputs.logits
+
+    # Get class probabilities
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0].tolist()
+
+    # Create result dictionary
     return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
-iface = gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(512, 512)), outputs=gr.outputs.Label(num_top_classes=3))
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(shape=(512, 512)),
+    outputs=gr.outputs.Label(num_top_classes=3)
+)
+
 iface.launch(share=True)