Update app.py

app.py

@@ -8,12 +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]')
+dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -79,36 +76,20 @@ model = CombinedModel()
 def predict(image):
     model.eval()
     with torch.no_grad():
-        image_tensor = transforms.ToTensor()(image).unsqueeze(0)
-        image_tensor = transforms.Resize((224, 224))(image_tensor)
+        image = transforms.ToTensor()(image).unsqueeze(0)
+        image = transforms.Resize((224, 224))(image)
         text_input = tokenizer("Sample prompt", return_tensors='pt', padding=True, truncation=True)
-        output = model(image_tensor, text_input)
+        output = model(image, text_input)
         _, indices = torch.topk(output, 5)
-        
-        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))
-
+        recommended_models = [dataset['Model'][i] for i in indices[0]]
     return recommended_models
 
 # Set up Gradio interface
-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."
-)
+interface = gr.Interface(fn=predict, 
+                         inputs=gr.inputs.Image(type="pil"), 
+                         outputs=gr.outputs.Textbox(label="Recommended Models"),
+                         title="AI Image Model Recommender",
+                         description="Upload an AI-generated image to receive model recommendations.")
 
 # Launch the app
-interface.launch()
+interface.launch()