Update app.py

app.py

@@ -9,8 +9,10 @@ from datasets import load_dataset
 from torch.utils.data import DataLoader, Dataset
 from sklearn.preprocessing import LabelEncoder
 
-# Load dataset
+# Load dataset and filter out null/none values
 dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
+# Filter out entries where Model is None or empty
+dataset = dataset.filter(lambda example: example['Model'] is not None and example['Model'].strip() != '')
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
@@ -77,47 +79,31 @@ class CombinedModel(nn.Module):
 # Instantiate model
 model = CombinedModel()
 
-def get_recommendations(image):
+# Define predict function
+def predict(image):
     model.eval()
     with torch.no_grad():
-        # Process image
-        transform = transforms.Compose([
-            transforms.Resize((224, 224)),
-            transforms.ToTensor()
-        ])
-        image_tensor = transform(image).unsqueeze(0)
-        
-        # Process text
+        image = transforms.ToTensor()(image).unsqueeze(0)
+        image = transforms.Resize((224, 224))(image)
         text_input = tokenizer(
             "Sample prompt",
             return_tensors='pt',
             padding=True,
             truncation=True
         )
-        
-        # Get predictions
-        output = model(image_tensor, text_input)
-        scores, indices = torch.topk(output, 5)
-        
-        # Prepare gallery output
-        recommendations = []
-        for idx, score in zip(indices[0], scores[0]):
-            sample_data = dataset[int(idx)]
-            recommendations.append({
-                'image': sample_data['image'],
-                'label': f"Model: {sample_data['Model']}\nScore: {score:.2f}"
-            })
-        
-        return recommendations
+        output = model(image, text_input)
+        _, indices = torch.topk(output, 5)
+        recommended_models = [dataset['Model'][i] for i in indices[0]]
+    return recommended_models
 
-# Gradio interface
+# Set up Gradio interface
 interface = gr.Interface(
-    fn=get_recommendations,
+    fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Gallery(label="Recommended Images"),
-    title="Image Recommendation System",
-    description="Upload an image and get similar images with their model names and distances."
+    outputs=gr.Textbox(label="Recommended Models"),
+    title="AI Image Model Recommender",
+    description="Upload an AI-generated image to receive model recommendations."
 )
 
-if __name__ == "__main__":
-    interface.launch()
+# Launch the app
+interface.launch()