Update app.py

app.py

@@ -8,9 +8,12 @@ import pandas as pd
 from datasets import load_dataset
 from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
 
 # Load dataset
-dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
+dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:1000]')
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -72,7 +75,19 @@ class CombinedModel(nn.Module):
 # Instantiate model
 model = CombinedModel()
 
-# Define predict function
+def display_recommendations(image, recommended_models, distances):
+    fig, axes = plt.subplots(1, len(recommended_models), figsize=(16, 4))
+    fig.suptitle("Recommended Models")
+
+    for i, (model, distance) in enumerate(zip(recommended_models, distances)):
+        # Load and display the recommended model image
+        model_image = Image.open(io.BytesIO(dataset.get_example(model)['image']))
+        axes[i].imshow(model_image)
+        axes[i].axis('off')
+        axes[i].set_title(f"{model}\nDistance: {distance:.2f}")
+
+    return fig
+
 def predict(image):
     model.eval()
     with torch.no_grad():
@@ -80,16 +95,15 @@ def predict(image):
         image = transforms.Resize((224, 224))(image)
         text_input = tokenizer("Sample prompt", return_tensors='pt', padding=True, truncation=True)
         output = model(image, text_input)
-        _, indices = torch.topk(output, 5)
-        recommended_models = [dataset['Model'][i] for i in indices[0]]
-    return recommended_models
+        distances, indices = torch.topk(-output.squeeze(), 5)
+        recommended_models = [dataset['Model'][i] for i in indices]
+        distances = (-distances).tolist()
+    return display_recommendations(image, recommended_models, distances)
 
-# Set up Gradio interface
 interface = gr.Interface(fn=predict, 
-                         inputs=gr.inputs.Image(type="pil"), 
-                         outputs=gr.outputs.Textbox(label="Recommended Models"),
+                         inputs=gr.Image(type="pil"), 
+                         outputs=gr.Plot(),
                          title="AI Image Model Recommender",
                          description="Upload an AI-generated image to receive model recommendations.")
 
-# Launch the app
 interface.launch()