Update app.py

app.py

@@ -11,8 +11,12 @@ from sklearn.preprocessing import LabelEncoder
 # Load dataset
 dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:10000]')
 
-# Text preprocessing function
+# Text preprocessing function with None handling
 def preprocess_text(text, max_length=100):
+    # Handle None or empty text
+    if text is None or not isinstance(text, str):
+        text = ""
+    
     # Convert text to lowercase and split into words
     words = text.lower().split()
     # Truncate or pad to max_length
@@ -29,32 +33,55 @@ class CustomDataset(Dataset):
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
         ])
+        
+        # Filter out None values from Model column
+        valid_indices = [i for i, model in enumerate(dataset['Model']) if model is not None]
+        self.valid_dataset = dataset.select(valid_indices)
+        
         self.label_encoder = LabelEncoder()
-        self.labels = self.label_encoder.fit_transform(dataset['Model'])
+        self.labels = self.label_encoder.fit_transform(self.valid_dataset['Model'])
         
         # Create vocabulary from all prompts
         self.vocab = set()
-        for item in dataset['prompt']:
-            self.vocab.update(preprocess_text(item))
+        for item in self.valid_dataset['prompt']:
+            try:
+                self.vocab.update(preprocess_text(item))
+            except Exception as e:
+                print(f"Error processing prompt: {e}")
+                continue
+        
+        # Remove empty string from vocabulary if present
+        self.vocab.discard('')
         self.vocab = list(self.vocab)
         self.word_to_idx = {word: idx for idx, word in enumerate(self.vocab)}
         
     def __len__(self):
-        return len(self.dataset)
+        return len(self.valid_dataset)
     
     def text_to_vector(self, text):
-        words = preprocess_text(text)
-        vector = torch.zeros(len(self.vocab))
-        for word in words:
-            if word in self.word_to_idx:
-                vector[self.word_to_idx[word]] += 1
-        return vector
+        try:
+            words = preprocess_text(text)
+            vector = torch.zeros(len(self.vocab))
+            for word in words:
+                if word in self.word_to_idx:
+                    vector[self.word_to_idx[word]] += 1
+            return vector
+        except Exception as e:
+            print(f"Error converting text to vector: {e}")
+            return torch.zeros(len(self.vocab))
     
     def __getitem__(self, idx):
-        image = self.transform(self.dataset[idx]['image'])
-        text_vector = self.text_to_vector(self.dataset[idx]['prompt'])
-        label = self.labels[idx]
-        return image, text_vector, label
+        try:
+            image = self.transform(self.valid_dataset[idx]['image'])
+            text_vector = self.text_to_vector(self.valid_dataset[idx]['prompt'])
+            label = self.labels[idx]
+            return image, text_vector, label
+        except Exception as e:
+            print(f"Error getting item at index {idx}: {e}")
+            # Return zero tensors as fallback
+            return (torch.zeros((3, 224, 224)), 
+                   torch.zeros(len(self.vocab)), 
+                   0)
 
 # Define CNN for image processing
 class ImageModel(nn.Module):
@@ -85,11 +112,11 @@ class TextMLP(nn.Module):
 
 # Combined model
 class CombinedModel(nn.Module):
-    def __init__(self, vocab_size):
+    def __init__(self, vocab_size, num_classes):
         super(CombinedModel, self).__init__()
         self.image_model = ImageModel()
         self.text_model = TextMLP(vocab_size)
-        self.fc = nn.Linear(1024, len(dataset['Model'].unique()))
+        self.fc = nn.Linear(1024, num_classes)
         
     def forward(self, image, text):
         image_features = self.image_model(image)
@@ -97,9 +124,15 @@ class CombinedModel(nn.Module):
         combined = torch.cat((image_features, text_features), dim=1)
         return self.fc(combined)
 
-# Create dataset instance and model
+# Create dataset instance
+print("Creating dataset...")
 custom_dataset = CustomDataset(dataset)
-model = CombinedModel(len(custom_dataset.vocab))
+print(f"Vocabulary size: {len(custom_dataset.vocab)}")
+print(f"Number of valid samples: {len(custom_dataset)}")
+
+# Create model
+num_classes = len(custom_dataset.label_encoder.classes_)
+model = CombinedModel(len(custom_dataset.vocab), num_classes)
 
 def get_recommendations(image):
     model.eval()
@@ -111,7 +144,7 @@ def get_recommendations(image):
         ])
         image_tensor = transform(image).unsqueeze(0)
         
-        # Create dummy text vector (since we're only doing image-based recommendations)
+        # Create dummy text vector
         dummy_text = torch.zeros((1, len(custom_dataset.vocab)))
         
         # Get model output
@@ -121,9 +154,13 @@ def get_recommendations(image):
         # Get recommended images and their information
         recommendations = []
         for idx in indices[0]:
-            recommended_image = dataset[idx.item()]['image']
-            model_name = dataset[idx.item()]['Model']
-            recommendations.append((recommended_image, f"{model_name}"))
+            try:
+                recommended_image = custom_dataset.valid_dataset[idx.item()]['image']
+                model_name = custom_dataset.valid_dataset[idx.item()]['Model']
+                recommendations.append((recommended_image, f"{model_name}"))
+            except Exception as e:
+                print(f"Error getting recommendation for index {idx}: {e}")
+                continue
         
     return recommendations
 
@@ -137,4 +174,5 @@ interface = gr.Interface(
 )
 
 # Launch the app
-interface.launch()
+if __name__ == "__main__":
+    interface.launch()