Update app.py

app.py

@@ -14,8 +14,9 @@ dataset = load_dataset('thefcraft/civitai-stable-diffusion-337k', split='train[:
 
 # Preprocess text data
 tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
 class CustomDataset(Dataset):
-    def init(self, dataset):
+    def __init__(self, dataset):
         self.dataset = dataset
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
@@ -23,48 +24,54 @@ class CustomDataset(Dataset):
         ])
         self.label_encoder = LabelEncoder()
         self.labels = self.label_encoder.fit_transform(dataset['Model'])
-    def len(self):
+
+    def __len__(self):
         return len(self.dataset)
-    def getitem(self, idx):
+
+    def __getitem__(self, idx):
         image = self.transform(self.dataset[idx]['image'])
         text = tokenizer(self.dataset[idx]['prompt'], padding='max_length', truncation=True, return_tensors='pt')
         label = self.labels[idx]
         return image, text, label
-        
+
 # Define CNN for image processing
 class ImageModel(nn.Module):
-    def init(self):
-        super(ImageModel, self).init()
+    def __init__(self):
+        super(ImageModel, self).__init__()
         self.model = models.resnet18(pretrained=True)
         self.model.fc = nn.Linear(self.model.fc.in_features, 512)
+
     def forward(self, x):
         return self.model(x)
-        
+
 # Define MLP for text processing
 class TextModel(nn.Module):
-    def init(self):
-        super(TextModel, self).init()
+    def __init__(self):
+        super(TextModel, self).__init__()
         self.bert = BertModel.from_pretrained('bert-base-uncased')
         self.fc = nn.Linear(768, 512)
+
     def forward(self, x):
-        output = self.bert(x)
+        output = self.bert(**x)
         return self.fc(output.pooler_output)
-        
+
 # Combined model
 class CombinedModel(nn.Module):
-    def init(self):
-        super(CombinedModel, self).init()
+    def __init__(self):
+        super(CombinedModel, self).__init__()
         self.image_model = ImageModel()
         self.text_model = TextModel()
         self.fc = nn.Linear(1024, len(dataset['Model']))
+
     def forward(self, image, text):
         image_features = self.image_model(image)
         text_features = self.text_model(text)
         combined = torch.cat((image_features, text_features), dim=1)
         return self.fc(combined)
-        
+
 # Instantiate model
 model = CombinedModel()
+
 # Define predict function
 def predict(image):
     model.eval()
@@ -72,16 +79,17 @@ def predict(image):
         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, textinput)
-        , indices = torch.topk(output, 5)
+        output = model(image, text_input)
+        _, indices = torch.topk(output, 5)
         recommended_models = [dataset['Model'][i] for i in indices[0]]
     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.")
+
 # Launch the app
 interface.launch()