Update app.py

app.py

@@ -8,6 +8,9 @@ import pandas as pd
 from datasets import load_dataset
 from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
+from PIL import Image
+import requests
+from io import BytesIO
 
 # Load dataset
 dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:1000]')
@@ -76,20 +79,36 @@ model = CombinedModel()
 def predict(image):
     model.eval()
     with torch.no_grad():
-        image = transforms.ToTensor()(image).unsqueeze(0)
-        image = transforms.Resize((224, 224))(image)
+        image_tensor = transforms.ToTensor()(image).unsqueeze(0)
+        image_tensor = transforms.Resize((224, 224))(image_tensor)
         text_input = tokenizer("Sample prompt", return_tensors='pt', padding=True, truncation=True)
-        output = model(image, text_input)
+        output = model(image_tensor, text_input)
         _, indices = torch.topk(output, 5)
-        recommended_models = [dataset['Model'][i] for i in indices[0]]
+        
+        recommended_models = []
+        for i in indices[0]:
+            model_name = dataset['Model'][i]
+            image_url = dataset['url'][i]
+            response = requests.get(image_url)
+            model_image = Image.open(BytesIO(response.content))
+            recommended_models.append((model_name, model_image))
+
     return recommended_models
 
 # Set up Gradio interface
-interface = gr.Interface(fn=predict, 
-                         inputs=gr.Image(type="pil"), 
-                         outputs=gr.Textbox(label="Recommended Models"),
-                         title="AI Image Model Recommender",
-                         description="Upload an AI-generated image to receive model recommendations.")
+def display_predictions(image):
+    recommendations = predict(image)
+    model_names = [rec[0] for rec in recommendations]
+    model_images = [rec[1] for rec in recommendations]
+    return model_names, model_images
+
+interface = gr.Interface(
+    fn=display_predictions,
+    inputs=gr.Image(type="pil"), 
+    outputs=[gr.Textbox(label="Recommended Models"), gr.Image(type="pil", label="Model Images")],
+    title="AI Image Model Recommender",
+    description="Upload an AI-generated image to receive model recommendations."
+)
 
 # Launch the app
-interface.launch()
+interface.launch()